Optimize session auth tag generation

I profiled my application and crypto/hmac.New was taking 5.71% of the
total CPU time. Since the session key does not change, you can call
crypto/hmac.Reset() instead of making a new object. This reuses previous
work done for the key.

With a separate generateAuthTag function for RTP/RTCP, it's now possible
to pass in the ROC instead of appending it to the buffer. This primarily
helps DecryptRTP, as it was otherwise doing an extra append just for the
ROC.

name                 old time/op    new time/op    delta
EncryptRTP-8           2.20µs ± 5%    2.00µs ± 4%   -9.00%
EncryptRTPInPlace-8    2.20µs ± 1%    1.97µs ± 5%  -10.52%
DecryptRTP-8           2.13µs ± 3%    1.79µs ± 4%  -16.18%

name                 old alloc/op   new alloc/op   delta
EncryptRTP-8           1.23kB ± 0%    0.79kB ± 0%  -36.04%
EncryptRTPInPlace-8    1.10kB ± 0%    0.66kB ± 0%  -40.22%
DecryptRTP-8           1.17kB ± 0%    0.69kB ± 0%  -40.75%

name                 old allocs/op  new allocs/op  delta
EncryptRTP-8             11.0 ± 0%       7.0 ± 0%  -36.36%
EncryptRTPInPlace-8      10.0 ± 0%       6.0 ± 0%  -40.00%
DecryptRTP-8             12.0 ± 0%       7.0 ± 0%  -41.67%

context.go

@@ -5,8 +5,8 @@ import (
 	"crypto/cipher"
 	"crypto/hmac"
 	"crypto/sha1" // #nosec
-	"crypto/subtle"
 	"encoding/binary"
+	"hash"
 
 	"github.com/pkg/errors"
 )
@@ -56,11 +56,13 @@ type Context struct {
 	ssrcStates         map[uint32]*ssrcState
 	srtpSessionKey     []byte
 	srtpSessionSalt    []byte
+	srtpSessionAuth    hash.Hash
 	srtpSessionAuthTag []byte
 	srtpBlock          cipher.Block
 
 	srtcpSessionKey     []byte
 	srtcpSessionSalt    []byte
+	srtcpSessionAuth    hash.Hash
 	srtcpSessionAuthTag []byte
 	srtcpIndex          uint32
 	srtcpBlock          cipher.Block
@@ -90,6 +92,8 @@ func CreateContext(masterKey, masterSalt []byte, profile ProtectionProfile) (c *
 		return nil, err
 	}
 
+	c.srtpSessionAuth = hmac.New(sha1.New, c.srtpSessionAuthTag)
+
 	if c.srtcpSessionKey, err = c.generateSessionKey(labelSRTCPEncryption); err != nil {
 		return nil, err
 	} else if c.srtcpSessionSalt, err = c.generateSessionSalt(labelSRTCPSalt); err != nil {
@@ -100,6 +104,8 @@ func CreateContext(masterKey, masterSalt []byte, profile ProtectionProfile) (c *
 		return nil, err
 	}
 
+	c.srtcpSessionAuth = hmac.New(sha1.New, c.srtcpSessionAuthTag)
+
 	return c, nil
 }
 
@@ -151,6 +157,7 @@ func (c *Context) generateSessionSalt(label byte) ([]byte, error) {
 	block.Encrypt(sessionSalt, sessionSalt)
 	return sessionSalt[0:saltLen], nil
 }
+
 func (c *Context) generateSessionAuthTag(label byte) ([]byte, error) {
 	// https://tools.ietf.org/html/rfc3711#appendix-B.3
 	// We now show how the auth key is generated.  The input block for AES-
@@ -199,7 +206,7 @@ func (c *Context) generateCounter(sequenceNumber uint16, rolloverCounter uint32,
 	return counter
 }
 
-func (c *Context) generateAuthTag(buf, sessionAuthTag []byte) ([]byte, error) {
+func (c *Context) generateSrtpAuthTag(buf []byte, roc uint32) ([]byte, error) {
 	// https://tools.ietf.org/html/rfc3711#section-4.2
 	// In the case of SRTP, M SHALL consist of the Authenticated
 	// Portion of the packet (as specified in Figure 1) concatenated with
@@ -214,27 +221,44 @@ func (c *Context) generateAuthTag(buf, sessionAuthTag []byte) ([]byte, error) {
 	// - Authenticated portion of the packet is everything BEFORE MKI
 	// - k_a is the session message authentication key
 	// - n_tag is the bit-length of the output authentication tag
-	// - ROC is already added by caller (to allow RTP + RTCP support)
-	mac := hmac.New(sha1.New, sessionAuthTag)
+	c.srtpSessionAuth.Reset()
 
-	if _, err := mac.Write(buf); err != nil {
+	if _, err := c.srtpSessionAuth.Write(buf); err != nil {
 		return nil, err
 	}
 
-	return mac.Sum(nil)[0:10], nil
-}
+	// For SRTP only, we need to hash the rollover counter as well.
+	rocRaw := [4]byte{}
+	binary.BigEndian.PutUint32(rocRaw[:], roc)
 
-func (c *Context) verifyAuthTag(buf, actualAuthTag []byte) (bool, error) {
-	expectedAuthTag, err := c.generateAuthTag(buf, c.srtpSessionAuthTag)
+	_, err := c.srtpSessionAuth.Write(rocRaw[:])
 	if err != nil {
-		return false, err
+		return nil, err
+	}
+
+	// Truncate the hash to the first 10 bytes.
+	return c.srtpSessionAuth.Sum(nil)[0:10], nil
+}
+
+func (c *Context) generateSrtcpAuthTag(buf []byte) ([]byte, error) {
+	// https://tools.ietf.org/html/rfc3711#section-4.2
+	//
+	// The pre-defined authentication transform for SRTP is HMAC-SHA1
+	// [RFC2104].  With HMAC-SHA1, the SRTP_PREFIX_LENGTH (Figure 3) SHALL
+	// be 0.  For SRTP (respectively SRTCP), the HMAC SHALL be applied to
+	// the session authentication key and M as specified above, i.e.,
+	// HMAC(k_a, M).  The HMAC output SHALL then be truncated to the n_tag
+	// left-most bits.
+	// - Authenticated portion of the packet is everything BEFORE MKI
+	// - k_a is the session message authentication key
+	// - n_tag is the bit-length of the output authentication tag
+	c.srtcpSessionAuth.Reset()
+
+	if _, err := c.srtcpSessionAuth.Write(buf); err != nil {
+		return nil, err
 	}
 
-	// We use a constant time comparison to prevent timing attacks.
-	// If we compared each byte and broke out early, like bytes.Equal,
-	// then an attacker could measure the time it takes for the auth comparison to fail.
-	// The longer it takes, the more correct the forged auth tag.
-	return subtle.ConstantTimeCompare(actualAuthTag, expectedAuthTag) == 1, nil
+	return c.srtcpSessionAuth.Sum(nil)[0:10], nil
 }
 
 // https://tools.ietf.org/html/rfc3550#appendix-A.1

srtcp.go

@@ -62,7 +62,7 @@ func (c *Context) encryptRTCP(dst, decrypted []byte) ([]byte, error) {
 	binary.BigEndian.PutUint32(out[len(out)-4:], c.srtcpIndex)
 	out[len(out)-4] |= 0x80
 
-	authTag, err := c.generateAuthTag(out, c.srtcpSessionAuthTag)
+	authTag, err := c.generateSrtcpAuthTag(out)
 	if err != nil {
 		return nil, err
 	}

srtp.go

@@ -2,33 +2,43 @@ package srtp
 
 import (
 	"crypto/cipher"
-	"encoding/binary"
+	"crypto/subtle"
 
 	"github.com/pions/rtp"
 	"github.com/pkg/errors"
 )
 
-func (c *Context) decryptRTP(dst, encrypted []byte, header *rtp.Header) ([]byte, error) {
-	dst = allocateIfMismatch(dst, encrypted)
+func (c *Context) decryptRTP(dst []byte, ciphertext []byte, header *rtp.Header) ([]byte, error) {
+	dst = growBufferSize(dst, len(ciphertext)-authTagSize)
 
 	s := c.getSSRCState(header.SSRC)
 	c.updateRolloverCount(header.SequenceNumber, s)
 
-	pktWithROC := append(append([]byte{}, dst[:len(dst)-authTagSize]...), make([]byte, 4)...)
-	binary.BigEndian.PutUint32(pktWithROC[len(pktWithROC)-4:], s.rolloverCounter)
+	// Split the auth tag and the cipher text into two parts.
+	actualTag := ciphertext[len(ciphertext)-authTagSize:]
+	ciphertext = ciphertext[:len(ciphertext)-authTagSize]
 
-	actualAuthTag := dst[len(dst)-authTagSize:]
-	verified, err := c.verifyAuthTag(pktWithROC, actualAuthTag)
+	// Generate the auth tag we expect to see from the ciphertext.
+	expectedTag, err := c.generateSrtpAuthTag(ciphertext, s.rolloverCounter)
 	if err != nil {
 		return nil, err
-	} else if !verified {
-		return nil, errors.Errorf("Failed to verify auth tag")
 	}
 
-	stream := cipher.NewCTR(c.srtpBlock, c.generateCounter(header.SequenceNumber, s.rolloverCounter, s.ssrc, c.srtpSessionSalt))
-	stream.XORKeyStream(dst[header.PayloadOffset:], dst[header.PayloadOffset:])
+	// See if the auth tag actually matches.
+	// We use a constant time comparison to prevent timing attacks.
+	if subtle.ConstantTimeCompare(actualTag, expectedTag) != 1 {
+		return nil, errors.Errorf("failed to verify auth tag")
+	}
+
+	// Write the plaintext header to the destination buffer.
+	copy(dst, ciphertext[:header.PayloadOffset])
+
+	// Decrypt the ciphertext for the payload.
+	counter := c.generateCounter(header.SequenceNumber, s.rolloverCounter, s.ssrc, c.srtpSessionSalt)
+	stream := cipher.NewCTR(c.srtpBlock, counter)
+	stream.XORKeyStream(dst[header.PayloadOffset:], ciphertext[header.PayloadOffset:])
 
-	return dst[:len(dst)-authTagSize], nil
+	return dst, nil
 }
 
 // DecryptRTP decrypts a RTP packet with an encrypted payload
@@ -84,20 +94,12 @@ func (c *Context) EncryptRTP(dst []byte, plaintext []byte, header *rtp.Header) (
 	stream.XORKeyStream(dst[offset:], plaintext[headerSize:])
 	offset += len(plaintext) - headerSize
 
-	// Write the rollover counter just for computing the auth tag.
-	// We will overwrite it immediately afterwards.
-	binary.BigEndian.PutUint32(dst[offset:], s.rolloverCounter)
-	offset += 4
-
 	// Generate the auth tag.
-	authTag, err := c.generateAuthTag(dst[:offset], c.srtpSessionAuthTag)
+	authTag, err := c.generateSrtpAuthTag(dst[:offset], s.rolloverCounter)
 	if err != nil {
 		return nil, err
 	}
 
-	// Pop off the rollover counter.
-	offset -= 4
-
 	// Write the auth tag to the dest.
 	copy(dst[offset:], authTag)
 

srtp_test.go

@@ -331,3 +331,34 @@ func BenchmarkEncryptRTPInPlace(b *testing.B) {
 		}
 	}
 }
+
+func BenchmarkDecryptRTP(b *testing.B) {
+	sequenceNumber := uint16(5000)
+	encrypted := []byte{0x6d, 0xd3, 0x7e, 0xd5, 0x99, 0xb7, 0x2d, 0x28, 0xb1, 0xf3, 0xa1, 0xf0, 0xc, 0xfb, 0xfd, 0x8}
+
+	encryptedPkt := &rtp.Packet{
+		Payload: encrypted,
+		Header: rtp.Header{
+			SequenceNumber: sequenceNumber,
+		},
+	}
+
+	encryptedRaw, err := encryptedPkt.Marshal()
+	if err != nil {
+		b.Fatal(err)
+	}
+
+	context, err := buildTestContext()
+	if err != nil {
+		b.Fatal(err)
+	}
+
+	b.ResetTimer()
+
+	for i := 0; i < b.N; i++ {
+		_, err := context.DecryptRTP(nil, encryptedRaw, nil)
+		if err != nil {
+			b.Fatal(err)
+		}
+	}
+}