1 use super::*; 2 use crate::error::Result; 3 use crate::protection_profile::*; 4 5 use bytes::{Bytes, BytesMut}; 6 use std::{collections::HashMap, sync::Arc}; 7 use tokio::{ 8 net::UdpSocket, 9 sync::{mpsc, Mutex}, 10 }; 11 12 async fn build_session_srtp_pair() -> Result<(Session, Session)> { 13 let ua = UdpSocket::bind("127.0.0.1:0").await?; 14 let ub = UdpSocket::bind("127.0.0.1:0").await?; 15 16 ua.connect(ub.local_addr()?).await?; 17 ub.connect(ua.local_addr()?).await?; 18 19 let ca = Config { 20 profile: ProtectionProfile::Aes128CmHmacSha1_80, 21 keys: SessionKeys { 22 local_master_key: vec![ 23 0xE1, 0xF9, 0x7A, 0x0D, 0x3E, 0x01, 0x8B, 0xE0, 0xD6, 0x4F, 0xA3, 0x2C, 0x06, 0xDE, 24 0x41, 0x39, 25 ], 26 local_master_salt: vec![ 27 0x0E, 0xC6, 0x75, 0xAD, 0x49, 0x8A, 0xFE, 0xEB, 0xB6, 0x96, 0x0B, 0x3A, 0xAB, 0xE6, 28 ], 29 remote_master_key: vec![ 30 0xE1, 0xF9, 0x7A, 0x0D, 0x3E, 0x01, 0x8B, 0xE0, 0xD6, 0x4F, 0xA3, 0x2C, 0x06, 0xDE, 31 0x41, 0x39, 32 ], 33 remote_master_salt: vec![ 34 0x0E, 0xC6, 0x75, 0xAD, 0x49, 0x8A, 0xFE, 0xEB, 0xB6, 0x96, 0x0B, 0x3A, 0xAB, 0xE6, 35 ], 36 }, 37 38 local_rtp_options: None, 39 remote_rtp_options: None, 40 41 local_rtcp_options: None, 42 remote_rtcp_options: None, 43 }; 44 45 let cb = Config { 46 profile: ProtectionProfile::Aes128CmHmacSha1_80, 47 keys: SessionKeys { 48 local_master_key: vec![ 49 0xE1, 0xF9, 0x7A, 0x0D, 0x3E, 0x01, 0x8B, 0xE0, 0xD6, 0x4F, 0xA3, 0x2C, 0x06, 0xDE, 50 0x41, 0x39, 51 ], 52 local_master_salt: vec![ 53 0x0E, 0xC6, 0x75, 0xAD, 0x49, 0x8A, 0xFE, 0xEB, 0xB6, 0x96, 0x0B, 0x3A, 0xAB, 0xE6, 54 ], 55 remote_master_key: vec![ 56 0xE1, 0xF9, 0x7A, 0x0D, 0x3E, 0x01, 0x8B, 0xE0, 0xD6, 0x4F, 0xA3, 0x2C, 0x06, 0xDE, 57 0x41, 0x39, 58 ], 59 remote_master_salt: vec![ 60 0x0E, 0xC6, 0x75, 0xAD, 0x49, 0x8A, 0xFE, 0xEB, 0xB6, 0x96, 0x0B, 0x3A, 0xAB, 0xE6, 61 ], 62 }, 63 64 local_rtp_options: None, 65 remote_rtp_options: None, 66 67 local_rtcp_options: None, 68 remote_rtcp_options: None, 69 }; 70 71 let sa = Session::new(Arc::new(ua), ca, true).await?; 72 let sb = Session::new(Arc::new(ub), cb, true).await?; 73 74 Ok((sa, sb)) 75 } 76 77 const TEST_SSRC: u32 = 5000; 78 const RTP_HEADER_SIZE: usize = 12; 79 80 #[tokio::test] 81 async fn test_session_srtp_accept() -> Result<()> { 82 let test_payload = Bytes::from_static(&[0x00, 0x01, 0x03, 0x04]); 83 let mut read_buffer = BytesMut::with_capacity(RTP_HEADER_SIZE + test_payload.len()); 84 read_buffer.resize(RTP_HEADER_SIZE + test_payload.len(), 0u8); 85 let (sa, sb) = build_session_srtp_pair().await?; 86 87 let packet = rtp::packet::Packet { 88 header: rtp::header::Header { 89 ssrc: TEST_SSRC, 90 ..Default::default() 91 }, 92 payload: test_payload.clone(), 93 }; 94 sa.write_rtp(&packet).await?; 95 96 let read_stream = sb.accept().await?; 97 let ssrc = read_stream.get_ssrc(); 98 assert_eq!( 99 ssrc, TEST_SSRC, 100 "SSRC mismatch during accept exp({}) actual({})", 101 TEST_SSRC, ssrc 102 ); 103 104 read_stream.read(&mut read_buffer).await?; 105 106 assert_eq!( 107 &test_payload[..], 108 &read_buffer[RTP_HEADER_SIZE..], 109 "Sent buffer does not match the one received exp({:?}) actual({:?})", 110 &test_payload[..], 111 &read_buffer[RTP_HEADER_SIZE..] 112 ); 113 114 sa.close().await?; 115 sb.close().await?; 116 117 Ok(()) 118 } 119 120 #[tokio::test] 121 async fn test_session_srtp_listen() -> Result<()> { 122 let test_payload = Bytes::from_static(&[0x00, 0x01, 0x03, 0x04]); 123 let mut read_buffer = BytesMut::with_capacity(RTP_HEADER_SIZE + test_payload.len()); 124 read_buffer.resize(RTP_HEADER_SIZE + test_payload.len(), 0u8); 125 let (sa, sb) = build_session_srtp_pair().await?; 126 127 let packet = rtp::packet::Packet { 128 header: rtp::header::Header { 129 ssrc: TEST_SSRC, 130 ..Default::default() 131 }, 132 payload: test_payload.clone(), 133 }; 134 135 let read_stream = sb.open(TEST_SSRC).await; 136 137 sa.write_rtp(&packet).await?; 138 139 read_stream.read(&mut read_buffer).await?; 140 141 assert_eq!( 142 &test_payload[..], 143 &read_buffer[RTP_HEADER_SIZE..], 144 "Sent buffer does not match the one received exp({:?}) actual({:?})", 145 &test_payload[..], 146 &read_buffer[RTP_HEADER_SIZE..] 147 ); 148 149 sa.close().await?; 150 sb.close().await?; 151 152 Ok(()) 153 } 154 155 #[tokio::test] 156 async fn test_session_srtp_multi_ssrc() -> Result<()> { 157 let ssrcs = vec![5000, 5001, 5002]; 158 let test_payload = Bytes::from_static(&[0x00, 0x01, 0x03, 0x04]); 159 let mut read_buffer = BytesMut::with_capacity(RTP_HEADER_SIZE + test_payload.len()); 160 read_buffer.resize(RTP_HEADER_SIZE + test_payload.len(), 0u8); 161 let (sa, sb) = build_session_srtp_pair().await?; 162 163 let mut read_streams = HashMap::new(); 164 for ssrc in &ssrcs { 165 let read_stream = sb.open(*ssrc).await; 166 read_streams.insert(*ssrc, read_stream); 167 } 168 169 for ssrc in &ssrcs { 170 let packet = rtp::packet::Packet { 171 header: rtp::header::Header { 172 ssrc: *ssrc, 173 ..Default::default() 174 }, 175 payload: test_payload.clone(), 176 }; 177 sa.write_rtp(&packet).await?; 178 179 if let Some(read_stream) = read_streams.get_mut(ssrc) { 180 read_stream.read(&mut read_buffer).await?; 181 182 assert_eq!( 183 &test_payload[..], 184 &read_buffer[RTP_HEADER_SIZE..], 185 "Sent buffer does not match the one received exp({:?}) actual({:?})", 186 &test_payload[..], 187 &read_buffer[RTP_HEADER_SIZE..] 188 ); 189 } else { 190 panic!("ssrc {} not found", *ssrc); 191 } 192 } 193 194 sa.close().await?; 195 sb.close().await?; 196 197 Ok(()) 198 } 199 200 fn encrypt_srtp(context: &mut Context, pkt: &rtp::packet::Packet) -> Result<Bytes> { 201 let decrypted = pkt.marshal()?; 202 let encrypted = context.encrypt_rtp(&decrypted)?; 203 Ok(encrypted) 204 } 205 206 async fn payload_srtp( 207 read_stream: &Arc<Stream>, 208 header_size: usize, 209 expected_payload: &[u8], 210 ) -> Result<u16> { 211 let mut read_buffer = BytesMut::with_capacity(header_size + expected_payload.len()); 212 read_buffer.resize(header_size + expected_payload.len(), 0u8); 213 214 let (n, hdr) = read_stream.read_rtp(&mut read_buffer).await?; 215 216 assert_eq!( 217 expected_payload, 218 &read_buffer[header_size..n], 219 "Sent buffer does not match the one received exp({:?}) actual({:?})", 220 expected_payload, 221 &read_buffer[header_size..n] 222 ); 223 224 Ok(hdr.sequence_number) 225 } 226 227 #[tokio::test] 228 async fn test_session_srtp_replay_protection() -> Result<()> { 229 let test_payload = Bytes::from_static(&[0x00, 0x01, 0x03, 0x04]); 230 231 let (sa, sb) = build_session_srtp_pair().await?; 232 233 let read_stream = sb.open(TEST_SSRC).await; 234 235 // Generate test packets 236 let mut packets = vec![]; 237 let mut expected_sequence_number = vec![]; 238 { 239 let mut local_context = sa.local_context.lock().await; 240 let mut i = 0xFFF0u16; 241 while i != 0x10 { 242 expected_sequence_number.push(i); 243 244 let packet = rtp::packet::Packet { 245 header: rtp::header::Header { 246 ssrc: TEST_SSRC, 247 sequence_number: i, 248 ..Default::default() 249 }, 250 payload: test_payload.clone(), 251 }; 252 253 let encrypted = encrypt_srtp(&mut local_context, &packet)?; 254 255 packets.push(encrypted); 256 257 if i == 0xFFFF { 258 i = 0; 259 } else { 260 i += 1; 261 } 262 } 263 } 264 265 let (done_tx, mut done_rx) = mpsc::channel::<()>(1); 266 267 let received_sequence_number = Arc::new(Mutex::new(vec![])); 268 let cloned_received_sequence_number = Arc::clone(&received_sequence_number); 269 let count = expected_sequence_number.len(); 270 271 tokio::spawn(async move { 272 let mut i = 0; 273 while i < count { 274 let seq = payload_srtp(&read_stream, RTP_HEADER_SIZE, &test_payload) 275 .await 276 .unwrap(); 277 let mut r = cloned_received_sequence_number.lock().await; 278 r.push(seq); 279 280 i += 1; 281 } 282 283 drop(done_tx); 284 }); 285 286 // Write with replay attack 287 for packet in &packets { 288 sa.udp_tx.send(packet).await?; 289 290 // Immediately replay 291 sa.udp_tx.send(packet).await?; 292 } 293 for packet in &packets { 294 // Delayed replay 295 sa.udp_tx.send(packet).await?; 296 } 297 298 done_rx.recv().await; 299 300 sa.close().await?; 301 sb.close().await?; 302 303 { 304 let received_sequence_number = received_sequence_number.lock().await; 305 assert_eq!(&received_sequence_number[..], &expected_sequence_number[..]); 306 } 307 308 Ok(()) 309 } 310