Refactor Common AEAD Code in Ciphersuite Package (#789)

Signed-off-by: Adriano Sela Aviles <adriano.selaviles@gmail.com>

Author: adrianosela

pkg/crypto/ciphersuite/README.md

@@ -70,13 +70,13 @@ go test -tags=bench -bench=BenchmarkCBCEncrypt -benchmem
 go test -tags=bench -bench=BenchmarkCBCDecrypt -benchmem
 ```
 
-- All cipgers, with 1KB payloads only
+- All ciphers, with 1KB payloads only
 
 ```bash
 go test -tags=bench -bench=/1KB -benchmem
 ```
 
-- All cipgers, with 16B payloads only
+- All ciphers, with 16B payloads only
 
 ```bash
 go test -tags=bench -bench=/16B -benchmem

pkg/crypto/ciphersuite/bench_helpers_test.go

@@ -5,7 +5,109 @@
 
 package ciphersuite
 
-import "testing"
+import (
+	"testing"
+
+	"github.com/pion/dtls/v3/pkg/protocol"
+	"github.com/pion/dtls/v3/pkg/protocol/recordlayer"
+)
+
+type testCipher interface {
+	Encrypt(pkt *recordlayer.RecordLayer, raw []byte) ([]byte, error)
+	Decrypt(header recordlayer.Header, in []byte) ([]byte, error)
+}
+
+// benchmarkEncrypt benchmarks a cipher's encryption with various payload sizes.
+func benchmarkEncrypt(b *testing.B, cipher testCipher) {
+	b.Helper()
+
+	payloadSizes := []int{16, 64, 128, 256, 512, 800, 1024, 1200, 1500, 4096, 8192}
+
+	for _, size := range payloadSizes {
+		b.Run(formatSize(b, size), func(b *testing.B) {
+			hdr := recordlayer.Header{
+				ContentType:    protocol.ContentTypeApplicationData,
+				Version:        protocol.Version1_2,
+				Epoch:          1,
+				SequenceNumber: 12345,
+			}
+
+			headerRaw, err := hdr.Marshal()
+			if err != nil {
+				b.Fatal(err)
+			}
+
+			payload := make([]byte, size)
+			raw := make([]byte, len(headerRaw)+len(payload))
+			copy(raw, headerRaw)
+			copy(raw[len(headerRaw):], payload)
+
+			pkt := &recordlayer.RecordLayer{Header: hdr}
+
+			b.ReportAllocs()
+			b.SetBytes(int64(size))
+			b.ResetTimer()
+
+			for i := 0; i < b.N; i++ {
+				rawCopy := make([]byte, len(raw))
+				copy(rawCopy, raw)
+
+				_, err := cipher.Encrypt(pkt, rawCopy)
+				if err != nil {
+					b.Fatal(err)
+				}
+			}
+		})
+	}
+}
+
+// benchmarkDecrypt benchmarks a cipher's decryption with various payload sizes.
+func benchmarkDecrypt(b *testing.B, cipher testCipher) {
+	b.Helper()
+
+	payloadSizes := []int{16, 64, 256, 512, 1024, 1500}
+
+	for _, size := range payloadSizes {
+		b.Run(formatSize(b, size), func(b *testing.B) {
+			hdr := recordlayer.Header{
+				ContentType:    protocol.ContentTypeApplicationData,
+				Version:        protocol.Version1_2,
+				Epoch:          1,
+				SequenceNumber: 12345,
+			}
+
+			headerRaw, err := hdr.Marshal()
+			if err != nil {
+				b.Fatal(err)
+			}
+
+			payload := make([]byte, size)
+			raw := make([]byte, len(headerRaw)+len(payload))
+			copy(raw, headerRaw)
+			copy(raw[len(headerRaw):], payload)
+
+			pkt := &recordlayer.RecordLayer{Header: hdr}
+			encrypted, err := cipher.Encrypt(pkt, raw)
+			if err != nil {
+				b.Fatal(err)
+			}
+
+			b.ReportAllocs()
+			b.SetBytes(int64(size))
+			b.ResetTimer()
+
+			for i := 0; i < b.N; i++ {
+				encCopy := make([]byte, len(encrypted))
+				copy(encCopy, encrypted)
+
+				_, err := cipher.Decrypt(hdr, encCopy)
+				if err != nil {
+					b.Fatal(err)
+				}
+			}
+		})
+	}
+}
 
 func formatSize(b *testing.B, size int) string {
 	b.Helper()

pkg/crypto/ciphersuite/cbc_bench_test.go

@@ -8,9 +8,6 @@ package ciphersuite
 import (
 	"crypto/sha256"
 	"testing"
-
-	"github.com/pion/dtls/v3/pkg/protocol"
-	"github.com/pion/dtls/v3/pkg/protocol/recordlayer"
 )
 
 // BenchmarkCBCEncrypt benchmarks CBC encryption with various payload sizes.
@@ -26,45 +23,7 @@ func BenchmarkCBCEncrypt(b *testing.B) {
 		b.Fatal(err)
 	}
 
-	payloadSizes := []int{16, 64, 256, 512, 1024, 1500}
-
-	// nolint:dupl
-	for _, size := range payloadSizes {
-		b.Run(formatSize(b, size), func(b *testing.B) {
-			hdr := recordlayer.Header{
-				ContentType:    protocol.ContentTypeApplicationData,
-				Version:        protocol.Version1_2,
-				Epoch:          1,
-				SequenceNumber: 12345,
-			}
-
-			headerRaw, err := hdr.Marshal()
-			if err != nil {
-				b.Fatal(err)
-			}
-
-			payload := make([]byte, size)
-			raw := make([]byte, len(headerRaw)+len(payload))
-			copy(raw, headerRaw)
-			copy(raw[len(headerRaw):], payload)
-
-			pkt := &recordlayer.RecordLayer{Header: hdr}
-
-			b.ReportAllocs()
-			b.SetBytes(int64(size))
-			b.ResetTimer()
-
-			for i := 0; i < b.N; i++ {
-				rawCopy := make([]byte, len(raw))
-				copy(rawCopy, raw)
-
-				_, err := cbcCipher.Encrypt(pkt, rawCopy)
-				if err != nil {
-					b.Fatal(err)
-				}
-			}
-		})
-	}
+	benchmarkEncrypt(b, cbcCipher)
 }
 
 // BenchmarkCBCDecrypt benchmarks CBC decryption with various payload sizes.
@@ -80,47 +39,5 @@ func BenchmarkCBCDecrypt(b *testing.B) {
 		b.Fatal(err)
 	}
 
-	payloadSizes := []int{16, 64, 256, 512, 1024, 1500}
-
-	// nolint:dupl
-	for _, size := range payloadSizes {
-		b.Run(formatSize(b, size), func(b *testing.B) {
-			hdr := recordlayer.Header{
-				ContentType:    protocol.ContentTypeApplicationData,
-				Version:        protocol.Version1_2,
-				Epoch:          1,
-				SequenceNumber: 12345,
-			}
-
-			headerRaw, err := hdr.Marshal()
-			if err != nil {
-				b.Fatal(err)
-			}
-
-			payload := make([]byte, size)
-			raw := make([]byte, len(headerRaw)+len(payload))
-			copy(raw, headerRaw)
-			copy(raw[len(headerRaw):], payload)
-
-			pkt := &recordlayer.RecordLayer{Header: hdr}
-			encrypted, err := cbcCipher.Encrypt(pkt, raw)
-			if err != nil {
-				b.Fatal(err)
-			}
-
-			b.ReportAllocs()
-			b.SetBytes(int64(size))
-			b.ResetTimer()
-
-			for i := 0; i < b.N; i++ {
-				encCopy := make([]byte, len(encrypted))
-				copy(encCopy, encrypted)
-
-				_, err := cbcCipher.Decrypt(hdr, encCopy)
-				if err != nil {
-					b.Fatal(err)
-				}
-			}
-		})
-	}
+	benchmarkDecrypt(b, cbcCipher)
 }

pkg/crypto/ciphersuite/ccm.go

@@ -5,13 +5,8 @@ package ciphersuite
 
 import (
 	"crypto/aes"
-	"crypto/rand"
-	"encoding/binary"
-	"fmt"
-	"sync"
 
 	"github.com/pion/dtls/v3/pkg/crypto/ccm"
-	"github.com/pion/dtls/v3/pkg/protocol"
 	"github.com/pion/dtls/v3/pkg/protocol/recordlayer"
 )
 
@@ -27,12 +22,7 @@ const (
 
 // CCM Provides an API to Encrypt/Decrypt DTLS 1.2 Packets.
 type CCM struct {
-	localCCM, remoteCCM         ccm.CCM
-	localWriteIV, remoteWriteIV []byte
-	tagLen                      CCMTagLen
-
-	// buffer pool for nonces.
-	nonceBufferPool sync.Pool
+	aead *aead
 }
 
 // NewCCM creates a DTLS GCM Cipher.
@@ -56,82 +46,23 @@ func NewCCM(tagLen CCMTagLen, localKey, localWriteIV, remoteKey, remoteWriteIV [
 	}
 
 	return &CCM{
-		localCCM:      localCCM,
-		localWriteIV:  localWriteIV,
-		remoteCCM:     remoteCCM,
-		remoteWriteIV: remoteWriteIV,
-		tagLen:        tagLen,
-
-		nonceBufferPool: sync.Pool{
-			New: func() any {
-				b := make([]byte, ccmNonceLength)
-				return &b // nolint:nlreturn
-			},
-		},
+		aead: newAEAD(
+			localCCM,
+			localWriteIV,
+			remoteCCM,
+			remoteWriteIV,
+			ccmNonceLength,
+			int(tagLen),
+		),
 	}, nil
 }
 
 // Encrypt encrypt a DTLS RecordLayer message.
 func (c *CCM) Encrypt(pkt *recordlayer.RecordLayer, raw []byte) ([]byte, error) {
-	payload := raw[pkt.Header.Size():]
-	raw = raw[:pkt.Header.Size()]
-
-	noncePtr := c.nonceBufferPool.Get().(*[]byte) // nolint:forcetypeassert
-	nonce := *noncePtr
-	defer c.nonceBufferPool.Put(noncePtr)
-
-	copy(nonce[:4], c.localWriteIV[:4])
-	if _, err := rand.Read(nonce[4:]); err != nil {
-		return nil, err
-	}
-
-	var additionalData []byte
-	if pkt.Header.ContentType == protocol.ContentTypeConnectionID {
-		additionalData = generateAEADAdditionalDataCID(&pkt.Header, len(payload))
-	} else {
-		additionalData = generateAEADAdditionalData(&pkt.Header, len(payload))
-	}
-
-	finalSize := len(raw) + 8 + len(payload) + int(c.tagLen)
-	result := make([]byte, finalSize)
-	copy(result, raw)
-	copy(result[len(raw):], nonce[4:])
-
-	c.localCCM.Seal(result[len(raw)+8:len(raw)+8], nonce, payload, additionalData)
-
-	// Update recordLayer size to include explicit nonce
-	binary.BigEndian.PutUint16(result[pkt.Header.Size()-2:], uint16(len(result)-pkt.Header.Size())) //nolint:gosec //G115
-
-	return result, nil
+	return c.aead.encrypt(pkt, raw)
 }
 
 // Decrypt decrypts a DTLS RecordLayer message.
 func (c *CCM) Decrypt(header recordlayer.Header, in []byte) ([]byte, error) {
-	if err := header.Unmarshal(in); err != nil {
-		return nil, err
-	}
-	switch {
-	case header.ContentType == protocol.ContentTypeChangeCipherSpec:
-		// Nothing to encrypt with ChangeCipherSpec
-		return in, nil
-	case len(in) <= (8 + header.Size()):
-		return nil, errNotEnoughRoomForNonce
-	}
-
-	nonce := append(append([]byte{}, c.remoteWriteIV[:4]...), in[header.Size():header.Size()+8]...)
-	out := in[header.Size()+8:]
-
-	var additionalData []byte
-	if header.ContentType == protocol.ContentTypeConnectionID {
-		additionalData = generateAEADAdditionalDataCID(&header, len(out)-int(c.tagLen))
-	} else {
-		additionalData = generateAEADAdditionalData(&header, len(out)-int(c.tagLen))
-	}
-	var err error
-	out, err = c.remoteCCM.Open(out[:0], nonce, out, additionalData)
-	if err != nil {
-		return nil, fmt.Errorf("%w: %v", errDecryptPacket, err) //nolint:errorlint
-	}
-
-	return append(in[:header.Size()], out...), nil
+	return c.aead.decrypt(header, in)
 }