Properly encrypt RTP packets with padding

Author: sirzooro

session_srtp.go

@@ -150,8 +150,20 @@ func (s *SessionSRTP) writeRTP(header *rtp.Header, payload []byte) (int, error)
 	ibuf := bufferpool.Get()
 	defer bufferpool.Put(ibuf)
 
+	buf := ibuf.([]byte)                                                      // nolint:forcetypeassert
+	headerLen, marshalSize := rtp.HeaderAndPacketMarshalSize(header, payload) // nolint:staticcheck
+	if len(buf) < marshalSize+20 {
+		// The buffer is too small, so we need to allocate a new one. Add 20 bytes for auth tag like
+		// for bufferpool above.
+		buf = make([]byte, marshalSize+20)
+	}
+	_, err := rtp.MarshalPacketTo(buf, header, payload) // nolint:staticcheck
+	if err != nil {
+		return 0, err
+	}
+
 	s.session.localContextMutex.Lock()
-	encrypted, err := s.localContext.encryptRTP(ibuf.([]byte), header, payload) //nolint:forcetypeassert
+	encrypted, err := s.localContext.encryptRTP(buf, header, headerLen, buf[:marshalSize])
 	s.session.localContextMutex.Unlock()
 
 	if err != nil {

session_srtp_test.go

@@ -363,3 +363,51 @@ func encryptSRTP(context *Context, pkt *rtp.Packet) ([]byte, error) {
 
 	return encrypted, nil
 }
+
+func TestSessionSRTPPacketWithPadding(t *testing.T) {
+	lim := test.TimeOut(time.Second * 5)
+	defer lim.Stop()
+
+	report := test.CheckRoutines(t)
+	defer report()
+
+	const (
+		testSSRC      = 5000
+		rtpHeaderSize = 12
+		paddingSize   = 5
+		authTagLen    = 10 // For AES_CM_128_HMAC_SHA1_80, the auth tag length is 10 bytes.
+	)
+	testPayload := []byte{0x00, 0x01, 0x03, 0x04}
+	readBuffer := make([]byte, rtpHeaderSize+paddingSize+len(testPayload))
+	aSession, bSession := buildSessionSRTPPair(t)
+
+	aWriteStream, err := aSession.OpenWriteStream()
+	assert.NoError(t, err)
+
+	writeBytes, err := aWriteStream.WriteRTP(&rtp.Header{SSRC: testSSRC, Padding: true, PaddingSize: paddingSize},
+		append([]byte{}, testPayload...))
+	assert.NoError(t, err)
+	assert.Equalf(t, rtpHeaderSize+paddingSize+len(testPayload)+authTagLen, writeBytes,
+		"WriteRTP should return the size of the packet including padding, exp(%v) actual(%v)",
+		rtpHeaderSize+paddingSize+len(testPayload)+authTagLen, writeBytes)
+
+	bReadStream, ssrc, err := bSession.AcceptStream()
+	assert.NoError(t, err)
+	assert.Equalf(t, uint32(testSSRC), ssrc, "SSRC mismatch during accept exp(%v) actual(%v)", testSSRC, ssrc)
+
+	readBytes, err := bReadStream.Read(readBuffer)
+	assert.NoError(t, err)
+	assert.Equal(t, rtpHeaderSize+paddingSize+len(testPayload), readBytes,
+		"Read should return the size of the packet including padding, exp(%v) actual(%v)",
+		rtpHeaderSize+paddingSize+len(testPayload), readBytes)
+
+	var rtpPacket rtp.Packet
+	err = rtpPacket.Unmarshal(readBuffer[:readBytes])
+	assert.NoError(t, err)
+	assert.Equal(t, rtpPacket.Padding, true)
+	assert.Equal(t, rtpPacket.PaddingSize, byte(paddingSize))
+	assert.Equal(t, rtpPacket.Payload, testPayload)
+
+	assert.NoError(t, aSession.Close())
+	assert.NoError(t, bSession.Close())
+}

srtp.go

@@ -113,13 +113,14 @@ func (c *Context) EncryptRTP(dst []byte, plaintext []byte, header *rtp.Header) (
 		return nil, err
 	}
 
-	return c.encryptRTP(dst, header, plaintext[headerLen:])
+	return c.encryptRTP(dst, header, headerLen, plaintext)
 }
 
 // encryptRTP marshals and encrypts an RTP packet, writing to the dst buffer provided.
 // If the dst buffer does not have the capacity, a new one will be allocated and returned.
 // Similar to above but faster because it can avoid unmarshaling the header and marshaling the payload.
-func (c *Context) encryptRTP(dst []byte, header *rtp.Header, payload []byte) (ciphertext []byte, err error) {
+func (c *Context) encryptRTP(dst []byte, header *rtp.Header, headerLen int, plaintext []byte,
+) (ciphertext []byte, err error) {
 	s := c.getSRTPSSRCState(header.SSRC)
 	roc, diff, ovf := s.nextRolloverCount(header.SequenceNumber)
 	if ovf {
@@ -136,7 +137,7 @@ func (c *Context) encryptRTP(dst []byte, header *rtp.Header, payload []byte) (ci
 		rocInPacket = true
 	}
 
-	return c.cipher.encryptRTP(dst, header, payload, roc, rocInPacket)
+	return c.cipher.encryptRTP(dst, header, headerLen, plaintext, roc, rocInPacket)
 }
 
 func (c *Context) hasROCInPacket(header *rtp.Header, authTagLen int) (bool, int) {

srtp_cipher.go

@@ -17,7 +17,7 @@ type srtpCipher interface {
 	AEADAuthTagLen() (int, error)
 	getRTCPIndex([]byte) uint32
 
-	encryptRTP([]byte, *rtp.Header, []byte, uint32, bool) ([]byte, error)
+	encryptRTP([]byte, *rtp.Header, int, []byte, uint32, bool) ([]byte, error)
 	encryptRTCP([]byte, []byte, uint32, uint32) ([]byte, error)
 
 	decryptRTP([]byte, []byte, *rtp.Header, int, uint32, bool) ([]byte, error)

srtp_cipher_aead_aes_gcm.go

@@ -91,10 +91,14 @@ func newSrtpCipherAeadAesGcm(
 func (s *srtpCipherAeadAesGcm) encryptRTP(
 	dst []byte,
 	header *rtp.Header,
-	payload []byte,
+	headerLen int,
+	plaintext []byte,
 	roc uint32,
 	rocInAuthTag bool,
 ) (ciphertext []byte, err error) {
+	payload := plaintext[headerLen:]
+	payloadLen := len(payload)
+
 	// Grow the given buffer to fit the output.
 	authTagLen, err := s.AEADAuthTagLen()
 	if err != nil {
@@ -106,18 +110,21 @@ func (s *srtpCipherAeadAesGcm) encryptRTP(
 		dstLen += 4
 	}
 	dst = growBufferSize(dst, dstLen)
+	sameBuffer := isSameBuffer(dst, plaintext)
 
-	n, err := header.MarshalTo(dst)
-	if err != nil {
-		return nil, err
+	// Copy the header unencrypted.
+	if !sameBuffer {
+		copy(dst, plaintext[:headerLen])
 	}
 
 	iv := s.rtpInitializationVector(header, roc)
 	if s.srtpEncrypted {
-		s.srtpCipher.Seal(dst[n:n], iv[:], payload, dst[:n])
+		s.srtpCipher.Seal(dst[headerLen:headerLen], iv[:], payload, dst[:headerLen])
 	} else {
-		clearLen := n + len(payload)
-		copy(dst[n:], payload)
+		clearLen := headerLen + payloadLen
+		if !sameBuffer {
+			copy(dst[headerLen:], payload)
+		}
 		s.srtpCipher.Seal(dst[clearLen:clearLen], iv[:], nil, dst[:clearLen])
 	}
 

srtp_cipher_aes_cm_hmac_sha1.go

@@ -104,33 +104,37 @@ func newSrtpCipherAesCmHmacSha1(
 func (s *srtpCipherAesCmHmacSha1) encryptRTP(
 	dst []byte,
 	header *rtp.Header,
-	payload []byte,
+	headerLen int,
+	plaintext []byte,
 	roc uint32,
 	rocInAuthTag bool,
 ) (ciphertext []byte, err error) {
+	payload := plaintext[headerLen:]
+	payloadLen := len(payload)
+
 	// Grow the given buffer to fit the output.
 	authTagLen, err := s.AuthTagRTPLen()
 	if err != nil {
 		return nil, err
 	}
-	dst = growBufferSize(dst, header.MarshalSize()+len(payload)+len(s.mki)+authTagLen)
+	dst = growBufferSize(dst, headerLen+payloadLen+len(s.mki)+authTagLen)
+	sameBuffer := isSameBuffer(dst, plaintext)
 
 	// Copy the header unencrypted.
-	n, err := header.MarshalTo(dst)
-	if err != nil {
-		return nil, err
+	if !sameBuffer {
+		copy(dst, plaintext[:headerLen])
 	}
 
 	// Encrypt the payload
 	if s.srtpEncrypted {
 		counter := generateCounter(header.SequenceNumber, roc, header.SSRC, s.srtpSessionSalt)
-		if err = xorBytesCTR(s.srtpBlock, counter[:], dst[n:], payload); err != nil {
+		if err = xorBytesCTR(s.srtpBlock, counter[:], dst[headerLen:], payload); err != nil {
 			return nil, err
 		}
-	} else {
-		copy(dst[n:], payload)
+	} else if !sameBuffer {
+		copy(dst[headerLen:], payload)
 	}
-	n += len(payload)
+	n := headerLen + payloadLen
 
 	// Generate the auth tag.
 	authTag, err := s.generateSrtpAuthTag(dst[:n], roc, rocInAuthTag)

srtp_cipher_test.go

@@ -591,6 +591,18 @@ func TestSrtpCipher(t *testing.T) {
 					assert.NoError(t, err)
 					assert.Equal(t, testCase.encryptedRTPPacket, actualEncrypted)
 				})
+
+				t.Run("Same buffer", func(t *testing.T) {
+					buffer := make([]byte, 0, 1000)
+					src, dst := buffer, buffer
+					src = append(src, testCase.decryptedRTPPacket...)
+					assert.True(t, isSameBuffer(dst, src))
+
+					actualEncrypted, err := ctx.EncryptRTP(dst, src, nil)
+					assert.NoError(t, err)
+					assert.Equal(t, testCase.encryptedRTPPacket, actualEncrypted)
+					assert.True(t, isSameBuffer(actualEncrypted, src))
+				})
 			})
 
 			t.Run("Decrypt RTP", func(t *testing.T) {

util.go

@@ -3,9 +3,12 @@
 
 package srtp
 
-import "bytes"
+import (
+	"bytes"
+	"unsafe"
+)
 
-// Grow the buffer size to the given number of bytes.
+// growBufferSize grows the buffer size to the given number of bytes.
 func growBufferSize(buf []byte, size int) []byte {
 	if size <= cap(buf) {
 		return buf[:size]
@@ -17,7 +20,7 @@ func growBufferSize(buf []byte, size int) []byte {
 	return buf2
 }
 
-// Check if buffers match, if not allocate a new buffer and return it.
+// allocateIfMismatch checks if buffers match, if not allocates a new buffer and returns it.
 func allocateIfMismatch(dst, src []byte) []byte {
 	if dst == nil {
 		dst = make([]byte, len(src))
@@ -35,3 +38,22 @@ func allocateIfMismatch(dst, src []byte) []byte {
 
 	return dst
 }
+
+// isSameBuffer returns true if slices a and b share the same underlying buffer.
+func isSameBuffer(a, b []byte) bool {
+	// If both are nil, they are technically the same (no buffer)
+	if a == nil && b == nil {
+		return true
+	}
+
+	// If either is nil, or both have 0 capacity, they can't share backing buffer
+	if cap(a) == 0 || cap(b) == 0 {
+		return false
+	}
+
+	// Create a slice of length 1 from each if possible
+	aPtr := unsafe.Pointer(&a[:1][0]) // nolint:gosec
+	bPtr := unsafe.Pointer(&b[:1][0]) // nolint:gosec
+
+	return aPtr == bPtr
+}