From 1d695c1d1f6e0573ea66cd9ad7f1031c574037b2 Mon Sep 17 00:00:00 2001
From: Aaron LI <aly@aaronly.me>
Date: Fri, 24 Nov 2023 13:39:56 +0800
Subject: [PATCH 3/3] crypto: Implement chacha20poly1305 in-place process for
 mbuf chains

Implement the chacha20poly1305_{encrypt,decrypt}_mbuf() functions that
performs in-place encryption and decryption for data in an mbuf chain.

The in-kernel WireGuard will use these two functions to encrypt/decrypt
packets.
---
 sys/crypto/chachapoly.c | 110 ++++++++++++++++++++++++++++++++++++++++
 sys/crypto/chachapoly.h |  11 ++++
 2 files changed, 121 insertions(+)

diff --git a/sys/crypto/chachapoly.c b/sys/crypto/chachapoly.c
index 34b0e468f7e..2e925db1d4f 100644
--- a/sys/crypto/chachapoly.c
+++ b/sys/crypto/chachapoly.c
@@ -18,6 +18,7 @@
 #include <sys/param.h>
 #include <sys/systm.h>
 #include <sys/endian.h>
+#include <sys/mbuf.h>
 
 #include <crypto/chachapoly.h>
 #include <crypto/chacha20/chacha.h>
@@ -37,6 +38,7 @@ struct chacha20poly1305_ctx {
 };
 
 static const uint8_t pad0[16] = { 0 };
+static const uint8_t blank_tag[CHACHA20POLY1305_AUTHTAG_SIZE] = { 0 };
 
 
 static void
@@ -276,3 +278,111 @@ xchacha20poly1305_decrypt(uint8_t *dst, const uint8_t *src, size_t src_len,
 	explicit_bzero(derived_key, sizeof(derived_key));
 	return (ret);
 }
+
+/*-------------------------------------------------------------------*/
+
+static int
+_chacha20poly1305_mbuf(struct chacha20poly1305_ctx *ctx,
+		       struct mbuf *m, int len)
+{
+	uint8_t block[CHACHA_BLOCKLEN], *p;
+	int off, remainder, blen, ret;
+
+	off = 0;
+	remainder = len;
+
+	while (remainder >= sizeof(block)) {
+		if (m->m_len != off && m->m_len < off + sizeof(block)) {
+			/*
+			 * Insufficient data at the end of an mbuf;
+			 * have to do some copying.
+			 */
+			m_copydata(m, off, sizeof(block), block);
+			_chacha20poly1305_update(ctx, block, block,
+						 sizeof(block));
+			m_copyback(m, off, sizeof(block), block);
+
+			/* Advance pointer. */
+			m = m_getptr(m, off + sizeof(block), &off);
+			if (m == NULL)
+				return (EINVAL);
+
+			remainder -= sizeof(block);
+			if (remainder == 0)
+				break;
+		}
+
+		if (m->m_len == off) {
+			off = 0;
+			m = m->m_next;
+			/* Skip possibly empty mbufs. */
+			while (m != NULL && m->m_len == 0)
+				m = m->m_next;
+			if (m == NULL) /* invalid because remainder > 0 */
+				return (EINVAL);
+		}
+
+		p = mtod(m, uint8_t *) + off;
+		blen = rounddown2(m->m_len - off, sizeof(block));
+		_chacha20poly1305_update(ctx, p, p, blen);
+
+		off += blen;
+		remainder -= blen;
+	}
+
+	m_copydata(m, off, remainder, block); /* remainder may be 0 */
+	if (ctx->flags & F_MODE_ENCRYPTION) {
+		ret = _chacha20poly1305_final(ctx, block, block, remainder) ?
+		      0 : EBADMSG;
+		m_copyback(m, off + remainder, sizeof(ctx->tag), ctx->tag);
+	} else {
+		m_copydata(m, off + remainder, sizeof(ctx->tag), ctx->tag);
+		ret = _chacha20poly1305_final(ctx, block, block, remainder) ?
+		      0 : EBADMSG;
+	}
+	m_copyback(m, off, remainder, block);
+
+	return (ret);
+}
+
+int
+chacha20poly1305_encrypt_mbuf(struct mbuf *m, const uint8_t *ad, size_t ad_len,
+			      const uint8_t nonce[CHACHA20POLY1305_NONCE_SIZE],
+			      const uint8_t key[CHACHA20POLY1305_KEY_SIZE])
+{
+	struct chacha20poly1305_ctx ctx;
+	int len;
+
+	if (!m_append(m, sizeof(blank_tag), blank_tag))
+		return (ENOMEM);
+
+	_chacha20poly1305_init(&ctx, true, nonce, key);
+
+	_chacha20poly1305_update(&ctx, NULL, ad, ad_len);
+
+	len = m->m_pkthdr.len - sizeof(blank_tag);
+	return _chacha20poly1305_mbuf(&ctx, m, len);
+}
+
+int
+chacha20poly1305_decrypt_mbuf(struct mbuf *m, const uint8_t *ad, size_t ad_len,
+			      const uint8_t nonce[CHACHA20POLY1305_NONCE_SIZE],
+			      const uint8_t key[CHACHA20POLY1305_KEY_SIZE])
+{
+	struct chacha20poly1305_ctx ctx;
+	int len, ret;
+
+	len = m->m_pkthdr.len - sizeof(ctx.tag);
+	if (len < 0)
+		return (EBADMSG);
+
+	_chacha20poly1305_init(&ctx, false, nonce, key);
+
+	_chacha20poly1305_update(&ctx, NULL, ad, ad_len);
+
+	ret = _chacha20poly1305_mbuf(&ctx, m, len);
+	if (ret == 0)
+		m_adj(m, -(int)sizeof(ctx.tag));
+
+	return (ret);
+}
diff --git a/sys/crypto/chachapoly.h b/sys/crypto/chachapoly.h
index 1900d185bbc..3ebf6a77b7e 100644
--- a/sys/crypto/chachapoly.h
+++ b/sys/crypto/chachapoly.h
@@ -53,4 +53,15 @@ bool xchacha20poly1305_decrypt(uint8_t *, const uint8_t *, size_t,
 			       const uint8_t[XCHACHA20POLY1305_NONCE_SIZE],
 			       const uint8_t[CHACHA20POLY1305_KEY_SIZE]);
 
+/*
+ * Perform in-place encryption/decryption for data in an mbuf chain.
+ */
+struct mbuf;
+int chacha20poly1305_encrypt_mbuf(struct mbuf *, const uint8_t *, size_t,
+				  const uint8_t[CHACHA20POLY1305_NONCE_SIZE],
+				  const uint8_t[CHACHA20POLY1305_KEY_SIZE]);
+int chacha20poly1305_decrypt_mbuf(struct mbuf *, const uint8_t *, size_t,
+				  const uint8_t[CHACHA20POLY1305_NONCE_SIZE],
+				  const uint8_t[CHACHA20POLY1305_KEY_SIZE]);
+
 #endif	/* _CHACHAPOLY_H_ */
-- 
2.40.1

