optimizations

Author: alexlapa

srtp/Cargo.toml

@@ -27,7 +27,7 @@ rtcp = { version = "0.9.0", path = "../rtcp" }
 byteorder = "1"
 bytes = "1"
 thiserror = "1.0"
-hmac = { version = "0.12.1", features = ["std", "reset"] }
+hmac = { version = "0.12.1", features = ["std"] }
 sha1 = "0.10.5"
 ctr = "0.8.0"
 aes = "0.7.5"

srtp/benches/srtp_bench.rs

@@ -3,7 +3,7 @@ use criterion::{criterion_group, criterion_main, Criterion};
 use util::Marshal;
 use webrtc_srtp::{context::Context, protection_profile::ProtectionProfile};
 
-fn benchmark_buffer(c: &mut Criterion) {
+fn benchmark_encrypt_rtp_aes_128_cm_hmac_sha1(c: &mut Criterion) {
     let mut ctx = Context::new(
         &vec![
             96, 180, 31, 4, 119, 137, 128, 252, 75, 194, 252, 44, 63, 56, 61, 55,
@@ -16,20 +16,18 @@ fn benchmark_buffer(c: &mut Criterion) {
     .unwrap();
 
     let mut pld = BytesMut::new();
-
-    for i in 0..1000 {
+    for i in 0..1200 {
         pld.extend_from_slice(&[i as u8]);
     }
 
-    let mut count = 1;
-
     c.bench_function("Benchmark context ", |b| {
+        let mut seq = 1;
         b.iter_batched(
             || {
                 let pkt = rtp::packet::Packet {
                     header: rtp::header::Header {
-                        sequence_number: count,
-                        timestamp: count.into(),
+                        sequence_number: seq,
+                        timestamp: seq.into(),
                         extension_profile: 48862,
                         marker: true,
                         padding: false,
@@ -39,7 +37,7 @@ fn benchmark_buffer(c: &mut Criterion) {
                     },
                     payload: pld.clone().into(),
                 };
-                count += 1;
+                seq += 1;
                 pkt.marshal().unwrap()
             },
             |pkt_raw| {
@@ -50,5 +48,5 @@ fn benchmark_buffer(c: &mut Criterion) {
     });
 }
 
-criterion_group!(benches, benchmark_buffer);
+criterion_group!(benches, benchmark_encrypt_rtp_aes_128_cm_hmac_sha1);
 criterion_main!(benches);

srtp/src/cipher/cipher_aead_aes_gcm.rs

@@ -33,18 +33,19 @@ impl Cipher for CipherAeadAesGcm {
         roc: u32,
     ) -> Result<Bytes> {
         // Grow the given buffer to fit the output.
+        let header_len = header.marshal_size();
         let mut writer =
-            BytesMut::with_capacity(header.marshal_size() + payload.len() + self.auth_tag_len());
+            BytesMut::with_capacity(payload.len() + self.auth_tag_len());
 
-        let data = header.marshal()?;
-        writer.extend(data);
+        // Copy header unencrypted.
+        writer.extend_from_slice(&payload[..header_len]);
 
         let nonce = self.rtp_initialization_vector(header, roc);
 
         let encrypted = self.srtp_cipher.encrypt(
             Nonce::from_slice(&nonce),
             Payload {
-                msg: payload,
+                msg: &payload[header_len..],
                 aad: &writer,
             },
         )?;

srtp/src/cipher/cipher_aes_cm_hmac_sha1.rs

@@ -121,19 +121,15 @@ impl CipherAesCmHmacSha1 {
     /// - 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
-    fn generate_srtp_auth_tag(&mut self, buf: &[u8], roc: u32) -> Vec<u8> {
-        self.srtp_session_auth.reset();
+    fn generate_srtp_auth_tag(&self, buf: &[u8], roc: u32) -> [u8; 20] {
+        let mut signer = self.srtp_session_auth.clone();
 
-        self.srtp_session_auth.update(buf);
+        signer.update(buf);
 
         // For SRTP only, we need to hash the rollover counter as well.
-        self.srtp_session_auth.update(&roc.to_be_bytes());
+        signer.update(&roc.to_be_bytes());
 
-        let result = self.srtp_session_auth.clone().finalize();
-        let code_bytes = result.into_bytes();
-
-        // Truncate the hash to the first AUTH_TAG_SIZE bytes.
-        code_bytes[0..self.auth_tag_len()].to_vec()
+        signer.finalize().into_bytes().into()
     }
 
     /// https://tools.ietf.org/html/rfc3711#section-4.2
@@ -147,13 +143,12 @@ impl CipherAesCmHmacSha1 {
     /// - 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
-    fn generate_srtcp_auth_tag(&mut self, buf: &[u8]) -> Vec<u8> {
-        self.srtcp_session_auth.reset();
+    fn generate_srtcp_auth_tag(&self, buf: &[u8]) -> Vec<u8> {
+        let mut signer = self.srtcp_session_auth.clone();
 
-        self.srtcp_session_auth.update(buf);
+        signer.update(buf);
 
-        let result = self.srtcp_session_auth.clone().finalize();
-        let code_bytes = result.into_bytes();
+        let code_bytes = signer.finalize().into_bytes();
 
         // Truncate the hash to the first AUTH_TAG_SIZE bytes.
         code_bytes[0..self.auth_tag_len()].to_vec()
@@ -179,26 +174,26 @@ impl Cipher for CipherAesCmHmacSha1 {
     ) -> Result<Bytes> {
         let header_len = header.marshal_size();
         let mut writer =
-            BytesMut::with_capacity(header_len + payload.len() + self.auth_tag_len());
+            Vec::with_capacity(payload.len() + self.auth_tag_len());
 
         // Copy the header unencrypted.
-        writer.extend(header.marshal());
+        writer.extend_from_slice(&payload[..header_len]);
+
         // Encrypt the payload
         let nonce = generate_counter(
             header.sequence_number,
             roc,
             header.ssrc,
             &self.srtp_session_salt,
         );
-
-        writer.put_bytes(0, payload.len());
+        writer.resize(payload.len(), 0);
         self.ctx.encrypt_init(None, None, Some(&nonce)).unwrap();
-        let count = self.ctx.cipher_update(&payload, Some(&mut writer[header_len..])).unwrap();
+        let count = self.ctx.cipher_update(&payload[header_len..], Some(&mut writer[header_len..])).unwrap();
         self.ctx.cipher_final(&mut writer[count..]).unwrap();
 
-        // Generate the auth tag.
-        let auth_tag = self.generate_srtp_auth_tag(&writer, roc);
-        writer.extend(auth_tag);
+        // Generate and write the auth tag.
+        let auth_tag = &self.generate_srtp_auth_tag(&writer, roc)[..self.auth_tag_len()];
+        writer.extend_from_slice(auth_tag);
 
         Ok(Bytes::from(writer))
     }
@@ -220,11 +215,11 @@ impl Cipher for CipherAesCmHmacSha1 {
         let cipher_text = &encrypted[..encrypted.len() - self.auth_tag_len()];
 
         // Generate the auth tag we expect to see from the ciphertext.
-        let expected_tag = self.generate_srtp_auth_tag(cipher_text, roc);
+        let expected_tag = &self.generate_srtp_auth_tag(cipher_text, roc)[..self.auth_tag_len()];
 
         // See if the auth tag actually matches.
         // We use a constant time comparison to prevent timing attacks.
-        if actual_tag.ct_eq(&expected_tag).unwrap_u8() != 1 {
+        if actual_tag.ct_eq(expected_tag).unwrap_u8() != 1 {
             return Err(Error::RtpFailedToVerifyAuthTag);
         }
 

srtp/src/context/mod.rs

@@ -170,14 +170,13 @@ impl Context {
         self.srtp_ssrc_states.entry(ssrc).or_insert(s)
     }
 
-    fn get_srtcp_ssrc_state(&mut self, ssrc: u32) -> Option<&mut SrtcpSsrcState> {
+    fn get_srtcp_ssrc_state(&mut self, ssrc: u32) -> &mut SrtcpSsrcState {
         let s = SrtcpSsrcState {
             ssrc,
             replay_detector: Some((self.new_srtcp_replay_detector)()),
             ..Default::default()
         };
-        self.srtcp_ssrc_states.entry(ssrc).or_insert(s);
-        self.srtcp_ssrc_states.get_mut(&ssrc)
+        self.srtcp_ssrc_states.entry(ssrc).or_insert(s)
     }
 
     /// roc returns SRTP rollover counter value of specified SSRC.
@@ -197,8 +196,6 @@ impl Context {
 
     /// set_index sets SRTCP index value of specified SSRC.
     fn set_index(&mut self, ssrc: u32, index: usize) {
-        if let Some(s) = self.get_srtcp_ssrc_state(ssrc) {
-            s.srtcp_index = index;
-        }
+        self.get_srtcp_ssrc_state(ssrc).srtcp_index = index;
     }
 }

srtp/src/context/srtcp.rs

@@ -13,26 +13,16 @@ impl Context {
         let index = self.cipher.get_rtcp_index(encrypted);
         let ssrc = u32::from_be_bytes([encrypted[4], encrypted[5], encrypted[6], encrypted[7]]);
 
-        {
-            if let Some(state) = self.get_srtcp_ssrc_state(ssrc) {
-                if let Some(replay_detector) = &mut state.replay_detector {
-                    if !replay_detector.check(index as u64) {
-                        return Err(Error::SrtcpSsrcDuplicated(ssrc, index));
-                    }
-                }
-            } else {
-                return Err(Error::SsrcMissingFromSrtcp(ssrc));
+        if let Some(replay_detector) = &mut self.get_srtcp_ssrc_state(ssrc).replay_detector {
+            if !replay_detector.check(index as u64) {
+                return Err(Error::SrtcpSsrcDuplicated(ssrc, index));
             }
         }
 
         let dst = self.cipher.decrypt_rtcp(encrypted, index, ssrc)?;
 
-        {
-            if let Some(state) = self.get_srtcp_ssrc_state(ssrc) {
-                if let Some(replay_detector) = &mut state.replay_detector {
-                    replay_detector.accept();
-                }
-            }
+        if let Some(replay_detector) = &mut self.get_srtcp_ssrc_state(ssrc).replay_detector {
+            replay_detector.accept();
         }
 
         Ok(dst)
@@ -46,18 +36,14 @@ impl Context {
 
         let ssrc = u32::from_be_bytes([decrypted[4], decrypted[5], decrypted[6], decrypted[7]]);
 
-        let index;
-        {
-            if let Some(state) = self.get_srtcp_ssrc_state(ssrc) {
-                state.srtcp_index += 1;
-                if state.srtcp_index > MAX_SRTCP_INDEX {
-                    state.srtcp_index = 0;
-                }
-                index = state.srtcp_index;
-            } else {
-                return Err(Error::SsrcMissingFromSrtcp(ssrc));
+        let index = {
+            let state = self.get_srtcp_ssrc_state(ssrc);
+            state.srtcp_index += 1;
+            if state.srtcp_index > MAX_SRTCP_INDEX {
+                state.srtcp_index = 0;
             }
-        }
+            state.srtcp_index
+        };
 
         self.cipher.encrypt_rtcp(decrypted, index, ssrc)
     }

srtp/src/context/srtp.rs

@@ -10,8 +10,7 @@ impl Context {
         encrypted: &[u8],
         header: &rtp::header::Header,
     ) -> Result<Bytes> {
-        let roc;
-        {
+        let roc = {
             let state = self.get_srtp_ssrc_state(header.ssrc);
             if let Some(replay_detector) = &mut state.replay_detector {
                 if !replay_detector.check(header.sequence_number as u64) {
@@ -22,8 +21,8 @@ impl Context {
                 }
             }
 
-            roc = state.next_rollover_count(header.sequence_number);
-        }
+            state.next_rollover_count(header.sequence_number)
+        };
 
         let dst = self.cipher.decrypt_rtp(encrypted, header, roc)?;
         {
@@ -46,14 +45,14 @@ impl Context {
 
     pub fn encrypt_rtp_with_header(
         &mut self,
-        plaintext: &[u8],
+        payload: &[u8],
         header: &rtp::header::Header,
     ) -> Result<Bytes> {
         let roc = self.get_srtp_ssrc_state(header.ssrc).next_rollover_count(header.sequence_number);
 
         let dst = self
             .cipher
-            .encrypt_rtp(&plaintext[header.marshal_size()..], header, roc)?;
+            .encrypt_rtp(&payload, header, roc)?;
 
         self.get_srtp_ssrc_state(header.ssrc).update_rollover_count(header.sequence_number);
 

srtp/src/error.rs

@@ -49,8 +49,6 @@ pub enum Error {
     SrtpSaltLength(usize, usize),
     #[error("SyntaxError: {0}")]
     ExtMapParse(String),
-    #[error("ssrc {0} not exist in srtp_ssrc_state")]
-    SsrcMissingFromSrtp(u32),
     #[error("srtp ssrc={0} index={1}: duplicated")]
     SrtpSsrcDuplicated(u32, u16),
     #[error("srtcp ssrc={0} index={1}: duplicated")]