From 967603abcce5864212dbe22e0e7a3df9e25313ed Mon Sep 17 00:00:00 2001
From: Gusted <postmaster@gusted.xyz>
Date: Wed, 25 Dec 2024 00:06:55 +0100
Subject: [PATCH] feat: harden keying implementation

Harden the current checks in place, I doubt these will ever hit (you can
prove easily by reading the current source code this cannot happen) but
just in case a new Go version does something weird or something else
goes catastrophicly wrong, this should add an extra defense-in-depth
layer.

`n != aeadKeySize` will panic a nil error, don't think it's needed to
add more logic to this, a nil error is enough to indicate that that
condition failed (given the other condition is `err != nil`).

Also move constant integers to being `const`, this helps reducing the
amount of instructions being done for the extra check.
---
 modules/keying/keying.go | 15 +++++++++------
 1 file changed, 9 insertions(+), 6 deletions(-)

diff --git a/modules/keying/keying.go b/modules/keying/keying.go
index 6082a91e8f..0b161b39c5 100644
--- a/modules/keying/keying.go
+++ b/modules/keying/keying.go
@@ -28,13 +28,16 @@ var (
 	// The hash used for HKDF.
 	hash = sha256.New
 	// The AEAD used for encryption/decryption.
-	aead          = chacha20poly1305.NewX
-	aeadKeySize   = chacha20poly1305.KeySize
-	aeadNonceSize = chacha20poly1305.NonceSizeX
+	aead = chacha20poly1305.NewX
 	// The pseudorandom key generated by HKDF-Extract.
 	prk []byte
 )
 
+const (
+	aeadKeySize   = chacha20poly1305.KeySize
+	aeadNonceSize = chacha20poly1305.NonceSizeX
+)
+
 // Set the main IKM for this module.
 func Init(ikm []byte) {
 	// Salt is intentionally left empty, it's not useful to Forgejo's use case.
@@ -55,7 +58,7 @@ var (
 // Derive *the* key for a given context, this is a deterministic function.
 // The same key will be provided for the same context.
 func DeriveKey(context Context) *Key {
-	if len(prk) == 0 {
+	if len(prk) != sha256.Size {
 		panic("keying: not initialized")
 	}
 
@@ -63,7 +66,7 @@ func DeriveKey(context Context) *Key {
 
 	key := make([]byte, aeadKeySize)
 	// This should never return an error, but if it does, panic.
-	if _, err := r.Read(key); err != nil {
+	if n, err := r.Read(key); err != nil || n != aeadKeySize {
 		panic(err)
 	}
 
@@ -92,7 +95,7 @@ func (k *Key) Encrypt(plaintext, additionalData []byte) []byte {
 
 	// Generate a random nonce.
 	nonce := make([]byte, aeadNonceSize)
-	if _, err := rand.Read(nonce); err != nil {
+	if n, err := rand.Read(nonce); err != nil || n != aeadNonceSize {
 		panic(err)
 	}