From d4f0b962ed23488f31dc8683430981da6f69d8e8 Mon Sep 17 00:00:00 2001
From: Sean DuBois <sean@siobud.com>
Date: Fri, 19 May 2023 15:12:24 -0400
Subject: [PATCH] Fix incorrect AES_CM_PRF logic

The PRF internally always uses AES-128. Before it was incorrectly
written to use a block the size of the SRTP Master Key. This would cause
AEAD_AES_256_GCM to use the incorrect cipher key. See [0] for logic.

[0] https://datatracker.ietf.org/doc/html/rfc6188#section-7.2
---
 key_derivation.go                |  6 +++---
 srtp_cipher_aead_aes_gcm_test.go | 20 ++++++++++----------
 2 files changed, 13 insertions(+), 13 deletions(-)

diff --git a/key_derivation.go b/key_derivation.go
index f0677d7..05f0f29 100644
--- a/key_derivation.go
+++ b/key_derivation.go
@@ -22,7 +22,7 @@ func aesCmKeyDerivation(label byte, masterKey, masterSalt []byte, indexOverKdr i
 	nMasterKey := len(masterKey)
 	nMasterSalt := len(masterSalt)
 
-	prfIn := make([]byte, nMasterKey)
+	prfIn := make([]byte, 16)
 	copy(prfIn[:nMasterSalt], masterSalt)
 
 	prfIn[7] ^= label
@@ -35,8 +35,8 @@ func aesCmKeyDerivation(label byte, masterKey, masterSalt []byte, indexOverKdr i
 
 	out := make([]byte, ((outLen+nMasterKey)/nMasterKey)*nMasterKey)
 	var i uint16
-	for n := 0; n < outLen; n += nMasterKey {
-		binary.BigEndian.PutUint16(prfIn[nMasterKey-2:], i)
+	for n := 0; n < outLen; n += block.BlockSize() {
+		binary.BigEndian.PutUint16(prfIn[len(prfIn)-2:], i)
 		block.Encrypt(out[n:n+nMasterKey], prfIn)
 		i++
 	}
diff --git a/srtp_cipher_aead_aes_gcm_test.go b/srtp_cipher_aead_aes_gcm_test.go
index a6fc4db..90fd88c 100644
--- a/srtp_cipher_aead_aes_gcm_test.go
+++ b/srtp_cipher_aead_aes_gcm_test.go
@@ -94,12 +94,12 @@ func TestSrtpCipherAedAes256Gcm(t *testing.T) {
 		0xab, 0xab, 0xab, 0xab,
 	}
 	encryptedRTPPacket := []byte{
-		0x80, 0x0f, 0x12, 0x34, 0xde, 0xca, 0xfb, 0xad,
-		0xca, 0xfe, 0xba, 0xbe, 0x0b, 0x16, 0x5c, 0x30,
-		0xca, 0xa3, 0xae, 0xce, 0xc6, 0x18, 0x45, 0x92,
-		0x2e, 0x74, 0xb9, 0x7f, 0xb, 0x2b, 0x50, 0x03,
-		0x7a, 0x6c, 0x86, 0x8a, 0xa7, 0xf4, 0x39, 0xfd,
-		0xbc, 0x0e, 0x11, 0x67,
+		0x80, 0xf, 0x12, 0x34, 0xde, 0xca, 0xfb, 0xad,
+		0xca, 0xfe, 0xba, 0xbe, 0xaf, 0x49, 0x96, 0x8f,
+		0x7e, 0x9c, 0x43, 0xf8, 0x01, 0xdd, 0x0c, 0x84,
+		0x8b, 0x1e, 0xc9, 0xb0, 0x29, 0xcd, 0xf8, 0x5c,
+		0xb7, 0x9a, 0x2f, 0x95, 0x60, 0xd4, 0x69, 0x75,
+		0x98, 0x50, 0x77, 0x25,
 	}
 	decryptedRtcpPacket := []byte{
 		0x81, 0xc8, 0x00, 0x0b, 0xca, 0xfe, 0xba, 0xbe,
@@ -108,10 +108,10 @@ func TestSrtpCipherAedAes256Gcm(t *testing.T) {
 	}
 	encryptedRtcpPacket := []byte{
 		0x81, 0xc8, 0x00, 0x0b, 0xca, 0xfe, 0xba, 0xbe,
-		0xe8, 0x0e, 0x69, 0x88, 0x59, 0x1b, 0xaf, 0xc8,
-		0x28, 0x33, 0x5c, 0x29, 0x0a, 0x0f, 0xa9, 0x18,
-		0xf2, 0x84, 0xf2, 0x90, 0xa3, 0xaa, 0x4b, 0xe5,
-		0x35, 0xa4, 0x28, 0xc6, 0xa0, 0xd7, 0x1e, 0xef,
+		0x98, 0x22, 0xba, 0x22, 0x96, 0x1c, 0x31, 0x48,
+		0xe7, 0xb7, 0xec, 0x4f, 0x09, 0xf4, 0x26, 0xdc,
+		0xf6, 0xb5, 0x9a, 0x75, 0xad, 0xec, 0x74, 0xfd,
+		0xb9, 0x51, 0xb6, 0x66, 0x84, 0x24, 0xd4, 0xe2,
 		0x80, 0x00, 0x00, 0x01,
 	}