Add E2EE example, style changes, more tests (#583)

livekit · Dec 19, 2024 · 6eb9e8f · 6eb9e8f
1 parent 218ce74
commit 6eb9e8f
Show file tree

Hide file tree

Showing 3 changed files with 323 additions and 56 deletions.
diff --git a/encryption.go b/encryption.go
@@ -22,11 +22,14 @@ const (
 	unencrypted_audio_bytes  = 1
 )
 
-var ErrIncorrectKeyLength = errors.New("incorrect key length for encryption/decryption")
-var ErrUnableGenerateIV = errors.New("unable to generate iv for encryption")
-var ErrIncorrectIVLength = errors.New("incorrect iv length")
-var ErrIncorrectSecretLength = errors.New("input secret provided to derivation function cannot be empty or nil")
-var ErrIncorrectSaltLength = errors.New("input salt provided to derivation function cannot be empty or nil")
+var (
+	ErrIncorrectKeyLength    = errors.New("incorrect key length for encryption/decryption")
+	ErrUnableGenerateIV      = errors.New("unable to generate iv for encryption")
+	ErrIncorrectIVLength     = errors.New("incorrect iv length")
+	ErrIncorrectSecretLength = errors.New("input secret provided to derivation function cannot be empty or nil")
+	ErrIncorrectSaltLength   = errors.New("input salt provided to derivation function cannot be empty or nil")
+	ErrBlockCipherRequired   = errors.New("input block cipher cannot be nil")
+)
 
 func DeriveKeyFromString(password string) ([]byte, error) {
 	return DeriveKeyFromStringCustomSalt(password, LIVEKIT_SDK_SALT)
@@ -77,6 +80,34 @@ func DeriveKeyFromBytesCustomSalt(secret []byte, salt string) ([]byte, error) {
 }
 
 // Take audio sample (body of RTP) encrypted by LiveKit client SDK, extract IV and decrypt using provided key
+// If sample matches sifTrailer, it's considered to be a non-encrypted Server Injected Frame and nil is returned
+// Use DecryptGCMAudioSampleCustomCipher with cached aes cipher block for better (30%) performance
+func DecryptGCMAudioSample(sample, key, sifTrailer []byte) ([]byte, error) {
+
+	if len(key) != 16 {
+		return nil, ErrIncorrectKeyLength
+	}
+
+	if sifTrailer != nil && len(sample) >= len(sifTrailer) {
+		possibleTrailer := sample[len(sample)-len(sifTrailer):]
+		if bytes.Equal(possibleTrailer, sifTrailer) {
+			// this is unencrypted Server Injected Frame (SIF) that should be dropped
+			return nil, nil
+		}
+
+	}
+
+	cipherBlock, err := aes.NewCipher(key)
+	if err != nil {
+		return nil, err
+	}
+
+	return DecryptGCMAudioSampleCustomCipher(sample, sifTrailer, cipherBlock)
+
+}
+
+// Take audio sample (body of RTP) encrypted by LiveKit client SDK, extract IV and decrypt using provided cipherBlock
+// If sample matches sifTrailer, it's considered to be a non-encrypted Server Injected Frame and nil is returned
 // Encrypted sample format based on livekit client sdk
 // ---------+-------------------------+---------+----
 // payload  |IV...(length = IV_LENGTH)|IV_LENGTH|KID|
@@ -86,10 +117,14 @@ func DeriveKeyFromBytesCustomSalt(secret []byte, salt string) ([]byte, error) {
 // IV - variable bytes (equal to IV_LENGTH bytes)
 // IV_LENGTH - 1 byte
 // KID (Key ID) - 1 byte - ignored here, key is provided as parameter to function
-func DecryptGCMAudioSample(sample, key, sifTrailer []byte) ([]byte, error) {
+func DecryptGCMAudioSampleCustomCipher(sample, sifTrailer []byte, cipherBlock cipher.Block) ([]byte, error) {
 
-	if len(key) != 16 {
-		return nil, ErrIncorrectKeyLength
+	if cipherBlock == nil {
+		return nil, ErrBlockCipherRequired
+	}
+
+	if sample == nil {
+		return nil, nil
 	}
 
 	if sifTrailer != nil && len(sample) >= len(sifTrailer) {
@@ -120,18 +155,11 @@ func DecryptGCMAudioSample(sample, key, sifTrailer []byte) ([]byte, error) {
 	cipherText := make([]byte, cipherTextLength)
 	copy(cipherText, sample[cipherTextStart:cipherTextStart+cipherTextLength])
 
-	// setup AES
-	aesCipher, err := aes.NewCipher(key)
+	aesGCM, err := cipher.NewGCMWithNonceSize(cipherBlock, ivLength) // standard Nonce size is 12 bytes, but since it MAY be different in the sample, we use the one from the sample
 	if err != nil {
 		return nil, err
 	}
 
-	aesGCM, err := cipher.NewGCMWithNonceSize(aesCipher, ivLength) // standard Nonce size is 12 bytes, but since it MAY be different in the sample, we use the one from the sample
-	if err != nil {
-		return nil, err
-	}
-
-	// fmt.Println("**** DECRYPTION BEGIN ********")
 	plainText, err := aesGCM.Open(nil, iv, cipherText, frameHeader)
 	if err != nil {
 		return nil, err
@@ -147,6 +175,23 @@ func DecryptGCMAudioSample(sample, key, sifTrailer []byte) ([]byte, error) {
 }
 
 // Take audio sample (body of RTP) and encrypts it using AES-GCM 128bit with provided key
+// Use EncryptGCMAudioSampleCustomCipher with cached aes cipher block for better (20%) performance
+func EncryptGCMAudioSample(sample, key []byte, kid uint8) ([]byte, error) {
+
+	if len(key) != 16 {
+		return nil, ErrIncorrectKeyLength
+	}
+
+	cipherBlock, err := aes.NewCipher(key)
+	if err != nil {
+		return nil, err
+	}
+
+	return EncryptGCMAudioSampleCustomCipher(sample, kid, cipherBlock)
+
+}
+
+// Take audio sample (body of RTP) and encrypts it using AES-GCM 128bit with provided cipher block
 // Encrypted sample format based on livekit client sdk
 // ---------+-------------------------+---------+----
 // payload  |IV...(length = IV_LENGTH)|IV_LENGTH|KID|
@@ -156,10 +201,14 @@ func DecryptGCMAudioSample(sample, key, sifTrailer []byte) ([]byte, error) {
 // IV - variable bytes (equal to IV_LENGTH bytes) - 12 random bytes
 // IV_LENGTH - 1 byte - 12 bytes fixed
 // KID (Key ID) - 1 byte - taken from "kid" parameter
-func EncryptGCMAudioSample(sample, key []byte, kid uint8) ([]byte, error) {
+func EncryptGCMAudioSampleCustomCipher(sample []byte, kid uint8, cipherBlock cipher.Block) ([]byte, error) {
 
-	if len(key) != 16 {
-		return nil, ErrIncorrectKeyLength
+	if cipherBlock == nil {
+		return nil, ErrBlockCipherRequired
+	}
+
+	if sample == nil {
+		return nil, nil
 	}
 
 	// variable naming is kept close to LiveKit client SDK decrypt function
@@ -179,13 +228,7 @@ func EncryptGCMAudioSample(sample, key []byte, kid uint8) ([]byte, error) {
 	plainText := make([]byte, plainTextLength)
 	copy(plainText, sample[plainTextStart:plainTextStart+plainTextLength])
 
-	// setup AES
-	aesCipher, err := aes.NewCipher(key)
-	if err != nil {
-		return nil, err
-	}
-
-	aesGCM, err := cipher.NewGCMWithNonceSize(aesCipher, LIVEKIT_IV_LENGTH) // standard Nonce size is 12 bytes, but using one from defined constant (which matches Javascript SDK)
+	aesGCM, err := cipher.NewGCMWithNonceSize(cipherBlock, LIVEKIT_IV_LENGTH) // standard Nonce size is 12 bytes, but using one from defined constant (which matches Javascript SDK)
 	if err != nil {
 		return nil, err
 	}

diff --git a/encryption_test.go b/encryption_test.go
@@ -1,25 +1,31 @@
 package lksdk
 
 import (
+	"crypto/aes"
 	"testing"
 
-	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
-var opusEncryptedFrame = []byte{120, 145, 24, 159, 76, 65, 130, 48, 144, 249, 17, 112, 134, 78, 250, 129, 171, 194, 16, 173, 73, 196, 5, 152, 69, 225, 28, 210, 196, 241, 226, 139, 231, 172, 51, 38, 139, 179, 245, 182, 170, 8, 122, 117, 98, 144, 123, 95, 73, 89, 119, 39, 205, 20, 191, 55, 121, 59, 239, 192, 85, 224, 228, 143, 10, 113, 195, 223, 118, 42, 2, 32, 22, 17, 77, 227, 109, 160, 245, 202, 189, 63, 162, 164, 5, 241, 24, 151, 45, 42, 165, 131, 171, 243, 141, 53, 35, 131, 141, 52, 253, 188, 12, 0}
-var opusDecryptedFrame = []byte{120, 11, 109, 82, 113, 132, 189, 156, 220, 173, 30, 109, 87, 54, 173, 99, 26, 126, 166, 37, 127, 234, 110, 211, 230, 152, 181, 235, 197, 19, 140, 230, 179, 35, 131, 132, 29, 192, 97, 247, 108, 53, 183, 214, 77, 181, 173, 206, 175, 7, 228, 145, 93, 155, 155, 142, 14, 27, 111, 64, 96, 196, 229, 189, 142, 59, 149, 169, 99, 225, 216, 85, 186, 182}
-var opusSilenceFrame = []byte{0xf8, 0xff, 0xfe, 0x00, 0x00, 0x00, 0x00, 0x00,
-	0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
-	0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
-	0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
-	0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
-	0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
-	0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
-	0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
-	0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
-	0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00}
-var sifTrailer = []byte{50, 86, 10, 220, 108, 185, 57, 211}
-var testPassphrase = "12345"
+var (
+	opusEncryptedFrame = []byte{120, 145, 24, 159, 76, 65, 130, 48, 144, 249, 17, 112, 134, 78, 250, 129, 171, 194, 16, 173, 73, 196, 5, 152, 69, 225, 28, 210, 196, 241, 226, 139, 231, 172, 51, 38, 139, 179, 245, 182, 170, 8, 122, 117, 98, 144, 123, 95, 73, 89, 119, 39, 205, 20, 191, 55, 121, 59, 239, 192, 85, 224, 228, 143, 10, 113, 195, 223, 118, 42, 2, 32, 22, 17, 77, 227, 109, 160, 245, 202, 189, 63, 162, 164, 5, 241, 24, 151, 45, 42, 165, 131, 171, 243, 141, 53, 35, 131, 141, 52, 253, 188, 12, 0}
+	opusDecryptedFrame = []byte{120, 11, 109, 82, 113, 132, 189, 156, 220, 173, 30, 109, 87, 54, 173, 99, 26, 126, 166, 37, 127, 234, 110, 211, 230, 152, 181, 235, 197, 19, 140, 230, 179, 35, 131, 132, 29, 192, 97, 247, 108, 53, 183, 214, 77, 181, 173, 206, 175, 7, 228, 145, 93, 155, 155, 142, 14, 27, 111, 64, 96, 196, 229, 189, 142, 59, 149, 169, 99, 225, 216, 85, 186, 182}
+	opusSilenceFrame   = []byte{
+		0xf8, 0xff, 0xfe, 0x00, 0x00, 0x00, 0x00, 0x00,
+		0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+		0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+		0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+		0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+		0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+		0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+		0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+		0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+		0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+	}
+	sifTrailer         = []byte{50, 86, 10, 220, 108, 185, 57, 211}
+	testPassphrase     = "12345"
+	keyIncorrectLength = []byte{1, 2, 3}
+)
 
 func TestDeriveKeyFromString(t *testing.T) {
 
@@ -28,8 +34,8 @@ func TestDeriveKeyFromString(t *testing.T) {
 	key, err := DeriveKeyFromString(password)
 	expectedKey := []byte{15, 94, 198, 66, 93, 211, 116, 46, 55, 97, 232, 121, 189, 233, 224, 22}
 
-	assert.Nil(t, err)
-	assert.Equal(t, key, expectedKey)
+	require.Nil(t, err)
+	require.Equal(t, key, expectedKey)
 }
 
 func TestDeriveKeyFromBytes(t *testing.T) {
@@ -38,43 +44,101 @@ func TestDeriveKeyFromBytes(t *testing.T) {
 	expectedKey := []byte{129, 224, 93, 62, 17, 203, 99, 136, 101, 35, 149, 128, 189, 152, 251, 76}
 
 	key, err := DeriveKeyFromBytes(inputSecret)
-	assert.Nil(t, err)
-	assert.Equal(t, expectedKey, key)
+	require.Nil(t, err)
+	require.Equal(t, expectedKey, key)
 
 }
 
 func TestDecryptAudioSample(t *testing.T) {
 
 	key, err := DeriveKeyFromString(testPassphrase)
-	assert.Nil(t, err)
+	require.Nil(t, err)
+
+	_, err = DecryptGCMAudioSample(opusEncryptedFrame, keyIncorrectLength, sifTrailer)
+	require.ErrorIs(t, err, ErrIncorrectKeyLength)
 
-	decryptedFrame, err := DecryptGCMAudioSample(opusEncryptedFrame, key, sifTrailer)
+	decryptedFrame, err := DecryptGCMAudioSample(nil, key, sifTrailer)
+	require.Nil(t, err)
+	require.Nil(t, decryptedFrame)
 
-	assert.Nil(t, err)
-	assert.Equal(t, opusDecryptedFrame, decryptedFrame)
+	decryptedFrame, err = DecryptGCMAudioSample(opusEncryptedFrame, key, sifTrailer)
+
+	require.Nil(t, err)
+	require.Equal(t, opusDecryptedFrame, decryptedFrame)
 
 	var sifFrame []byte
 	sifFrame = append(sifFrame, opusSilenceFrame...)
 	sifFrame = append(sifFrame, sifTrailer...)
 
 	decryptedFrame, err = DecryptGCMAudioSample(sifFrame, key, sifTrailer)
-	assert.Nil(t, err)
-	assert.Nil(t, decryptedFrame)
+	require.Nil(t, err)
+	require.Nil(t, decryptedFrame)
 
 }
 
 func TestEncryptAudioSample(t *testing.T) {
 
 	key, err := DeriveKeyFromString(testPassphrase)
-	assert.Nil(t, err)
+	require.Nil(t, err)
+
+	_, err = EncryptGCMAudioSample(opusDecryptedFrame, keyIncorrectLength, 0)
+	require.ErrorIs(t, err, ErrIncorrectKeyLength)
 
-	encryptedFrame, err := EncryptGCMAudioSample(opusDecryptedFrame, key, 0)
+	encryptedFrame, err := EncryptGCMAudioSample(nil, key, 0)
+	require.Nil(t, err)
+	require.Nil(t, encryptedFrame)
 
-	assert.Nil(t, err)
+	encryptedFrame, err = EncryptGCMAudioSample(opusDecryptedFrame, key, 0)
+
+	require.Nil(t, err)
 
 	// IV is generated randomly so to verify we decrypt and make sure that we got the expected plain text frame
 	decryptedFrame, err := DecryptGCMAudioSample(encryptedFrame, key, sifTrailer)
-	assert.Nil(t, err)
-	assert.Equal(t, opusDecryptedFrame, decryptedFrame)
+	require.Nil(t, err)
+	require.Equal(t, opusDecryptedFrame, decryptedFrame)
+
+}
+
+func BenchmarkDecryptAudioNewCipher(b *testing.B) {
+
+	key, _ := DeriveKeyFromString(testPassphrase)
+	for i := 0; i < b.N; i++ {
+		cipherBlock, _ := aes.NewCipher(key)
+		DecryptGCMAudioSampleCustomCipher(opusEncryptedFrame, sifTrailer, cipherBlock)
+
+	}
+
+}
+
+func BenchmarkDecryptAudioCachedCipher(b *testing.B) {
+
+	key, _ := DeriveKeyFromString(testPassphrase)
+	cipherBlock, _ := aes.NewCipher(key)
+	for i := 0; i < b.N; i++ {
+		DecryptGCMAudioSampleCustomCipher(opusEncryptedFrame, sifTrailer, cipherBlock)
+
+	}
+
+}
+
+func BenchmarkEncryptAudioCachedCipher(b *testing.B) {
+
+	key, _ := DeriveKeyFromString(testPassphrase)
+	cipherBlock, _ := aes.NewCipher(key)
+	for i := 0; i < b.N; i++ {
+		EncryptGCMAudioSampleCustomCipher(opusDecryptedFrame, 0, cipherBlock)
+
+	}
+
+}
+
+func BenchmarkEncryptAudioNewCipher(b *testing.B) {
+
+	key, _ := DeriveKeyFromString(testPassphrase)
+	for i := 0; i < b.N; i++ {
+		cipherBlock, _ := aes.NewCipher(key)
+		EncryptGCMAudioSampleCustomCipher(opusDecryptedFrame, 0, cipherBlock)
+
+	}
 
 }