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