From 1d8c47e4ed760efe9f00c233eab569d5a6ab3caa Mon Sep 17 00:00:00 2001
From: Aaron LI <aly@aaronly.me>
Date: Thu, 23 Nov 2023 17:56:45 +0800
Subject: [PATCH 2/3] crypto: Refactor the chacha20poly1305 code to be more
 flexible

Introduce the 'chacha20poly1305_ctx' struct to hold the context, and
implement the _init()/_update()/_final() functions as the building
blocks to perform encryption/decryption in a more generic way.

The main intention is to help implement the in-place encryption and
decryption of data in an mbuf chain.  That would reduce the unnecessary
memory allocations and data copies in packet manipulations, as needed by
the in-kernel WireGuard.  This API will be done in a later commit.

Rewrite the original chacha20poly1305_{encrypt,decrypt}() functions
using the new blocks.
---
 sys/crypto/chachapoly.c | 238 ++++++++++++++++++++++++++++------------
 sys/crypto/chachapoly.h |  12 ++
 2 files changed, 179 insertions(+), 71 deletions(-)

diff --git a/sys/crypto/chachapoly.c b/sys/crypto/chachapoly.c
index 13ea6bf0871..34b0e468f7e 100644
--- a/sys/crypto/chachapoly.c
+++ b/sys/crypto/chachapoly.c
@@ -24,57 +24,182 @@
 #include <crypto/poly1305/poly1305.h>
 
 
+struct chacha20poly1305_ctx {
+	struct chacha_ctx	chacha;
+	poly1305_state		poly;
+	uint8_t			tag[CHACHA20POLY1305_AUTHTAG_SIZE];
+	size_t			ad_len;
+	size_t			data_len;
+	int			flags;
+#define F_MODE_ENCRYPTION	0x001	/* encryption operation */
+#define F_INITIALIZED		0x002	/* context initialized */
+#define F_AD_DONE		0x004	/* no more additional data */
+};
+
 static const uint8_t pad0[16] = { 0 };
 
 
-void
-chacha20poly1305_encrypt(uint8_t *dst, const uint8_t *src, size_t src_len,
-			 const uint8_t *ad, size_t ad_len,
-			 const uint8_t nonce[CHACHA20POLY1305_NONCE_SIZE],
-			 const uint8_t key[CHACHA20POLY1305_KEY_SIZE])
+static void
+_chacha20poly1305_init(struct chacha20poly1305_ctx *ctx, bool enc_mode,
+		       const uint8_t nonce[CHACHA20POLY1305_NONCE_SIZE],
+		       const uint8_t key[CHACHA20POLY1305_KEY_SIZE])
 {
-	struct chacha_ctx chacha;
-	poly1305_state poly;
-	uint64_t lens[2];
 	uint8_t poly_key[CHACHA20POLY1305_KEY_SIZE];
 	uint8_t nonce64[CHACHA_NONCELEN], counter64[CHACHA_CTRLEN];
-	size_t pad_len;
 
+	bzero(ctx, sizeof(*ctx));
 	bzero(poly_key, sizeof(poly_key));
 	bzero(nonce64, sizeof(nonce64));
 	bzero(counter64, sizeof(counter64));
 
 	/*
 	 * The original ChaCha uses a 64-bit nonce and a 64-bit counter,
-	 * while the IETF version modified it to be a 96-bit nonce and
+	 * but the IETF version is modified to use a 96-bit nonce and
 	 * a 32-bit counter. (RFC 8439)
 	 */
 	memcpy(counter64 + 4, nonce, 4);
 	memcpy(nonce64, nonce + 4, 8);
 
 	/* Generate the Poly1305 one-time key. */
-	chacha_keysetup(&chacha, key, CHACHA20POLY1305_KEY_SIZE * 8);
-	chacha_ivsetup(&chacha, nonce64, counter64);
-	chacha_encrypt_bytes(&chacha, poly_key, poly_key, sizeof(poly_key));
+	chacha_keysetup(&ctx->chacha, key, CHACHA20POLY1305_KEY_SIZE * 8);
+	chacha_ivsetup(&ctx->chacha, nonce64, counter64);
+	chacha_encrypt_bytes(&ctx->chacha, poly_key, poly_key,
+			     sizeof(poly_key));
+
+	/* Initialize the authenticator. */
+	poly1305_init(&ctx->poly, poly_key);
+
+	ctx->flags |= (enc_mode ? F_MODE_ENCRYPTION : 0);
+	ctx->flags |= F_INITIALIZED;
+}
+
+/*
+ * Process the AD (additional data) and cipher data (i.e., plaintext or
+ * ciphertext).
+ *
+ * If the output buffer <out> is NULL, then the input data are regarded
+ * as AAD.  The AAD can be provided in multiple pieces, but must be done
+ * before any cipher data.
+ *
+ * NOTE: The cipher data must be complete blocks.  This requirement is
+ *       placed to help easily implement the in-place encryption/
+ *       decryption for data in non-contiguous buffers (e.g., mbuf chain).
+ */
+static void
+_chacha20poly1305_update(struct chacha20poly1305_ctx *ctx,
+			 uint8_t *out, const uint8_t *in, size_t in_len)
+{
+	KKASSERT((ctx->flags & F_INITIALIZED) != 0);
+
+	if (out == NULL) {
+		/* Additional data */
+		KKASSERT((ctx->flags & F_AD_DONE) == 0);
+		poly1305_update(&ctx->poly, in, in_len);
+		ctx->ad_len += in_len;
+		return;
+	}
+
+	/* Cipher data: must be complete blocks. */
+	KKASSERT(in_len % CHACHA_BLOCKLEN == 0);
+
+	if ((ctx->flags & F_AD_DONE) == 0) {
+		poly1305_update(&ctx->poly, pad0,
+				(0x10 - (ctx->ad_len & 0xF)));
+		ctx->flags |= F_AD_DONE;
+	}
+
+	/* Swap the Poly1305/ChaCha order to support in-place operation. */
+	if (ctx->flags & F_MODE_ENCRYPTION) {
+		chacha_encrypt_bytes(&ctx->chacha, in, out, in_len);
+		poly1305_update(&ctx->poly, out, in_len);
+	} else {
+		poly1305_update(&ctx->poly, in, in_len);
+		chacha_encrypt_bytes(&ctx->chacha, in, out, in_len);
+	}
+
+	ctx->data_len += in_len;
+}
+
+/*
+ * Process the last cipher block, which can be empty, complete or incomplete.
+ *
+ * In the encryption mode, the tag is calculated and is available via
+ * 'ctx->tag'.  Let the caller get the tag and append to the output, because
+ * the allocated tag space may not directly follow the output buffer; e.g.,
+ * the tag is allocated on a separate mbuf.
+ *
+ * In the decryption mode, the caller must set the expected tag in
+ * 'ctx->tag' before calling this function.  The tag will be verified and
+ * a boolean is returned to indicate whether the tag matches.
+ */
+static bool
+_chacha20poly1305_final(struct chacha20poly1305_ctx *ctx,
+			uint8_t *out, const uint8_t *in, size_t in_len)
+{
+	uint64_t lens[2];
+	uint8_t tag[CHACHA20POLY1305_AUTHTAG_SIZE];
+	bool ret;
+
+	KKASSERT((ctx->flags & F_INITIALIZED) != 0);
+	KKASSERT(in_len <= CHACHA_BLOCKLEN);
+	KKASSERT(out != NULL);
 
-	/* Encrypt the plaintext. */
-	chacha_encrypt_bytes(&chacha, src, dst, src_len);
+	if (ctx->ad_len > 0 && (ctx->flags & F_AD_DONE) == 0) {
+		KKASSERT(ctx->data_len == 0);
+		poly1305_update(&ctx->poly, pad0,
+				(0x10 - (ctx->ad_len & 0xF)));
+		ctx->flags |= F_AD_DONE;
+	}
+
+	if (ctx->flags & F_MODE_ENCRYPTION) {
+		chacha_encrypt_bytes(&ctx->chacha, in, out, in_len);
+		poly1305_update(&ctx->poly, out, in_len);
+	} else {
+		poly1305_update(&ctx->poly, in, in_len);
+		chacha_encrypt_bytes(&ctx->chacha, in, out, in_len);
+	}
+	poly1305_update(&ctx->poly, pad0, (0x10 - (in_len & 0xF)));
+
+	ctx->data_len += in_len;
 
 	/* Calculate the tag. */
-	poly1305_init(&poly, poly_key);
-	poly1305_update(&poly, ad, ad_len);
-	if ((pad_len = 0x10 - (ad_len & 0xF)) > 0)
-		poly1305_update(&poly, pad0, pad_len);
-	poly1305_update(&poly, dst, src_len);
-	if ((pad_len = 0x10 - (src_len & 0xF)) > 0)
-		poly1305_update(&poly, pad0, pad_len);
-	lens[0] = htole64(ad_len);
-	lens[1] = htole64(src_len);
-	poly1305_update(&poly, (uint8_t *)lens, sizeof(lens));
-	poly1305_finish(&poly, dst + src_len);
-
-	explicit_bzero(&chacha, sizeof(chacha));
-	explicit_bzero(poly_key, sizeof(poly_key));
+	lens[0] = htole64(ctx->ad_len);
+	lens[1] = htole64(ctx->data_len);
+	poly1305_update(&ctx->poly, (uint8_t *)lens, sizeof(lens));
+	poly1305_finish(&ctx->poly, tag);
+
+	if (ctx->flags & F_MODE_ENCRYPTION) {
+		ret = true;
+		explicit_bzero(ctx, sizeof(*ctx));
+		memcpy(ctx->tag, tag, sizeof(tag));
+	} else {
+		ret = (timingsafe_bcmp(ctx->tag, tag, sizeof(tag)) == 0);
+		explicit_bzero(ctx, sizeof(*ctx));
+	}
+
+	return (ret);
+}
+
+/*-------------------------------------------------------------------*/
+
+void
+chacha20poly1305_encrypt(uint8_t *dst, const uint8_t *src, size_t src_len,
+			 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;
+	size_t len;
+
+	_chacha20poly1305_init(&ctx, true, nonce, key);
+
+	_chacha20poly1305_update(&ctx, NULL, ad, ad_len);
+
+	len = rounddown2(src_len, CHACHA_BLOCKLEN);
+	_chacha20poly1305_update(&ctx, dst, src, len);
+
+	_chacha20poly1305_final(&ctx, dst + len, src + len, src_len - len);
+	memcpy(dst + src_len, ctx.tag, sizeof(ctx.tag));
 }
 
 bool
@@ -83,56 +208,27 @@ chacha20poly1305_decrypt(uint8_t *dst, const uint8_t *src, size_t src_len,
 			 const uint8_t nonce[CHACHA20POLY1305_NONCE_SIZE],
 			 const uint8_t key[CHACHA20POLY1305_KEY_SIZE])
 {
-	struct chacha_ctx chacha;
-	poly1305_state poly;
-	uint64_t lens[2];
-	uint8_t poly_key[CHACHA20POLY1305_KEY_SIZE];
-	uint8_t tag[CHACHA20POLY1305_AUTHTAG_SIZE];
-	uint8_t nonce64[CHACHA_NONCELEN], counter64[CHACHA_CTRLEN];
-	size_t pad_len, ct_len;
-	int ret;
+	struct chacha20poly1305_ctx ctx;
+	size_t data_len, len;
 
-	if (src_len < CHACHA20POLY1305_AUTHTAG_SIZE)
+	if (src_len < sizeof(ctx.tag))
 		return (false);
 
-	bzero(poly_key, sizeof(poly_key));
-	bzero(nonce64, sizeof(nonce64));
-	bzero(counter64, sizeof(counter64));
-
-	memcpy(counter64 + 4, nonce, 4);
-	memcpy(nonce64, nonce + 4, 8);
+	_chacha20poly1305_init(&ctx, false, nonce, key);
 
-	/* Generate the Poly1305 one-time key. */
-	chacha_keysetup(&chacha, key, CHACHA20POLY1305_KEY_SIZE * 8);
-	chacha_ivsetup(&chacha, nonce64, counter64);
-	chacha_encrypt_bytes(&chacha, poly_key, poly_key, sizeof(poly_key));
+	_chacha20poly1305_update(&ctx, NULL, ad, ad_len);
 
-	/* Calculate the tag. */
-	poly1305_init(&poly, poly_key);
-	poly1305_update(&poly, ad, ad_len);
-	if ((pad_len = 0x10 - (ad_len & 0xF)) > 0)
-		poly1305_update(&poly, pad0, pad_len);
-	ct_len = src_len - CHACHA20POLY1305_AUTHTAG_SIZE;
-	poly1305_update(&poly, src, ct_len);
-	if ((pad_len = 0x10 - (ct_len & 0xF)) > 0)
-		poly1305_update(&poly, pad0, pad_len);
-	lens[0] = htole64(ad_len);
-	lens[1] = htole64(ct_len);
-	poly1305_update(&poly, (uint8_t *)lens, sizeof(lens));
-	poly1305_finish(&poly, tag);
-
-	/* Compare the tag. */
-	ret = timingsafe_bcmp(tag, src + ct_len, sizeof(tag));
-	if (ret == 0) {
-		/* Decrypt the ciphertext. */
-		chacha_encrypt_bytes(&chacha, src, dst, ct_len);
-	}
+	data_len = src_len - sizeof(ctx.tag);
+	len = rounddown2(data_len, CHACHA_BLOCKLEN);
+	_chacha20poly1305_update(&ctx, dst, src, len);
 
-	explicit_bzero(&chacha, sizeof(chacha));
-	explicit_bzero(poly_key, sizeof(poly_key));
-	return (ret == 0);
+	memcpy(ctx.tag, src + data_len, sizeof(ctx.tag));
+	return _chacha20poly1305_final(&ctx, dst + len, src + len,
+				       data_len - len);
 }
 
+/*-------------------------------------------------------------------*/
+
 /*
  * XChaCha: eXtended-nonce ChaCha and AEAD_XChaCha20_Poly1305
  * RFC draft: https://datatracker.ietf.org/doc/html/draft-irtf-cfrg-xchacha
diff --git a/sys/crypto/chachapoly.h b/sys/crypto/chachapoly.h
index e7f1d8ff982..1900d185bbc 100644
--- a/sys/crypto/chachapoly.h
+++ b/sys/crypto/chachapoly.h
@@ -23,6 +23,12 @@
 #define CHACHA20POLY1305_NONCE_SIZE	12
 #define XCHACHA20POLY1305_NONCE_SIZE	24
 
+/*
+ * ChaCha20-Poly1305 AEAD cipher (RFC 8439)
+ *
+ * NOTE: Support in-place encryption/decryption; i.e., the output buffer
+ *       points to the same location as the input.
+ */
 void chacha20poly1305_encrypt(uint8_t *, const uint8_t *, size_t,
 			      const uint8_t *, size_t,
 			      const uint8_t[CHACHA20POLY1305_NONCE_SIZE],
@@ -32,6 +38,12 @@ bool chacha20poly1305_decrypt(uint8_t *, const uint8_t *, size_t,
 			      const uint8_t[CHACHA20POLY1305_NONCE_SIZE],
 			      const uint8_t[CHACHA20POLY1305_KEY_SIZE]);
 
+/*
+ * XChaCha20-Poly1305 AEAD cipher
+ * (extended nonce size from 96 bits to 192 bits)
+ *
+ * NOTE: Support in-place encryption/decryption, as above.
+ */
 void xchacha20poly1305_encrypt(uint8_t *, const uint8_t *, size_t,
 			       const uint8_t *, size_t,
 			       const uint8_t[XCHACHA20POLY1305_NONCE_SIZE],
-- 
2.40.1

