From f01b1f8baccd29292f222fc1d7360a0755a27e59 Mon Sep 17 00:00:00 2001
From: Ivan <ivan@quad4.io>
Date: Sat, 27 Sep 2025 04:40:35 -0500
Subject: [PATCH] Update Decrypt method in Identity to validate token structure
 and HMAC. Update extraction logic for ephemeral public key, ciphertext, and
 MAC, ensuring proper error handling for token size and HMAC validation.

---
 pkg/identity/identity.go | 16 +++++++++++++---
 1 file changed, 13 insertions(+), 3 deletions(-)

diff --git a/pkg/identity/identity.go b/pkg/identity/identity.go
index 8685443..7edc929 100644
--- a/pkg/identity/identity.go
+++ b/pkg/identity/identity.go
@@ -307,9 +307,14 @@ func (i *Identity) Decrypt(ciphertextToken []byte, ratchets [][]byte, enforceRat
 		return nil, errors.New("decryption failed because the token size was invalid")
 	}
 
-	// Extract peer public key and ciphertext
-	peerPubBytes := ciphertextToken[:KEYSIZE/8/2]
-	ciphertext := ciphertextToken[KEYSIZE/8/2:]
+	// Extract components: ephemeralPubKey(32) + ciphertext + mac(32)
+	if len(ciphertextToken) < 32+32+32 { // minimum sizes
+		return nil, errors.New("token too short")
+	}
+
+	peerPubBytes := ciphertextToken[:32]
+	ciphertext := ciphertextToken[32 : len(ciphertextToken)-32]
+	mac := ciphertextToken[len(ciphertextToken)-32:]
 
 	// Try decryption with ratchets first if provided
 	if len(ratchets) > 0 {
@@ -343,6 +348,11 @@ func (i *Identity) Decrypt(ciphertextToken []byte, ratchets [][]byte, enforceRat
 		return nil, fmt.Errorf("failed to derive key: %v", err)
 	}
 
+	// Validate HMAC
+	if !cryptography.ValidateHMAC(derivedKey, append(peerPubBytes, ciphertext...), mac) {
+		return nil, errors.New("invalid HMAC")
+	}
+
 	// Create AES cipher
 	block, err := aes.NewCipher(derivedKey)
 	if err != nil {