Add some more agent tests.

stun-proto/src/agent.rs

@@ -871,6 +871,7 @@ pub(crate) mod tests {
             .remote_addr(remote_addr)
             .build();
         let msg = Message::new_request(BINDING);
+        let transaction_id = msg.transaction_id();
         let transmit = agent.send(msg, remote_addr).unwrap();
         let now = Instant::now();
         assert_eq!(transmit.transport, TransportType::Udp);
@@ -881,11 +882,100 @@ pub(crate) mod tests {
         let resp_data = response.to_bytes();
         let ret = agent.handle_incoming_data(&resp_data, remote_addr).unwrap();
         assert!(matches!(ret[0], HandleStunReply::StunResponse(_, _)));
+        assert!(agent.request_transaction(transaction_id).is_none());
+        assert!(agent.mut_request_transaction(transaction_id).is_none());
 
         let ret = agent.poll(now);
         assert!(matches!(ret, StunAgentPollRet::WaitUntil(_)));
     }
 
+    #[test]
+    fn indication_with_invalid_response() {
+        init();
+        let local_addr = "127.0.0.1:2000".parse().unwrap();
+        let remote_addr = "127.0.0.1:1000".parse().unwrap();
+        let mut agent = StunAgent::builder(TransportType::Udp, local_addr)
+            .remote_addr(remote_addr)
+            .build();
+        let transaction_id = Message::generate_transaction();
+        let msg = Message::new(
+            MessageType::from_class_method(MessageClass::Indication, BINDING),
+            transaction_id,
+        );
+        let transmit = agent.send(msg, remote_addr).unwrap();
+        assert_eq!(transmit.transport, TransportType::Udp);
+        assert_eq!(transmit.from, local_addr);
+        assert_eq!(transmit.to, remote_addr);
+        let _indication = Message::from_bytes(&transmit.data).unwrap();
+        assert!(agent.request_transaction(transaction_id).is_none());
+        assert!(agent.mut_request_transaction(transaction_id).is_none());
+        // you should definitely never do this ;). Indications should never get replies.
+        let response = Message::new(
+            MessageType::from_class_method(MessageClass::Error, BINDING),
+            transaction_id,
+        );
+        let resp_data = response.to_bytes();
+        // response without a request is dropped.
+        let ret = agent.handle_incoming_data(&resp_data, remote_addr).unwrap();
+        assert!(ret.is_empty());
+    }
+
+    #[test]
+    fn request_with_credentials() {
+        init();
+        let local_addr = "10.0.0.1:12345".parse().unwrap();
+        let remote_addr = "10.0.0.2:3478".parse().unwrap();
+
+        let mut agent = StunAgent::builder(TransportType::Udp, local_addr).build();
+        let local_credentials = ShortTermCredentials::new(String::from("local_password"));
+        let remote_credentials = ShortTermCredentials::new(String::from("remote_password"));
+        agent.set_local_credentials(local_credentials.clone().into());
+        agent.set_remote_credentials(remote_credentials.clone().into());
+
+        // unvalidated peer data should be dropped
+        let data = vec![20; 4];
+        let replies = agent.handle_incoming_data(&data, remote_addr).unwrap();
+        assert!(replies.is_empty());
+
+        let mut msg = Message::new_request(BINDING);
+        let transaction_id = msg.transaction_id();
+        msg.add_message_integrity(&local_credentials.clone().into(), IntegrityAlgorithm::Sha1)
+            .unwrap();
+        let transmit = agent.send(msg, remote_addr).unwrap();
+
+        let request = Message::from_bytes(&transmit.data).unwrap();
+
+        let mut response = Message::new_success(&request);
+        response
+            .add_attribute(XorMappedAddress::new(
+                transmit.from,
+                request.transaction_id(),
+            ))
+            .unwrap();
+        response
+            .add_message_integrity(&remote_credentials.into(), IntegrityAlgorithm::Sha1)
+            .unwrap();
+
+        let data = response.to_bytes();
+        let to = transmit.to;
+        let mut reply = agent.handle_incoming_data(&data, to).unwrap();
+        let HandleStunReply::StunResponse(request, response) = reply.remove(0) else {
+            unreachable!();
+        };
+
+        assert_eq!(request.transaction_id(), transaction_id);
+        assert_eq!(response.transaction_id(), transaction_id);
+        assert!(agent.request_transaction(transaction_id).is_none());
+        assert!(agent.mut_request_transaction(transaction_id).is_none());
+
+        let data = vec![20; 4];
+        let mut replies = agent.handle_incoming_data(&data, remote_addr).unwrap();
+        let HandleStunReply::Data(received) = replies.remove(0) else {
+            unreachable!();
+        };
+        assert_eq!(data, received);
+    }
+
     #[test]
     fn request_unanswered() {
         init();
@@ -895,6 +985,7 @@ pub(crate) mod tests {
             .remote_addr(remote_addr)
             .build();
         let msg = Message::new_request(BINDING);
+        let transaction_id = msg.transaction_id();
         agent.send(msg, remote_addr).unwrap();
         let mut now = Instant::now();
         loop {
@@ -907,6 +998,13 @@ pub(crate) mod tests {
                 _ => unreachable!(),
             }
         }
+        assert!(agent.request_transaction(transaction_id).is_none());
+        assert!(agent.mut_request_transaction(transaction_id).is_none());
+
+        // unvalidated peer data should be dropped
+        let data = vec![20; 4];
+        let replies = agent.handle_incoming_data(&data, remote_addr).unwrap();
+        assert!(replies.is_empty());
     }
 
     #[test]
@@ -917,7 +1015,13 @@ pub(crate) mod tests {
 
         let mut agent = StunAgent::builder(TransportType::Udp, local_addr).build();
 
+        // unvalidated peer data should be dropped
+        let data = vec![20; 4];
+        let replies = agent.handle_incoming_data(&data, remote_addr).unwrap();
+        assert!(replies.is_empty());
+
         let msg = Message::new_request(BINDING);
+        let transaction_id = msg.transaction_id();
         let transmit = agent.send(msg, remote_addr).unwrap();
 
         let request = Message::from_bytes(&transmit.data).unwrap();
@@ -935,12 +1039,21 @@ pub(crate) mod tests {
         let reply = agent.handle_incoming_data(&data, to).unwrap();
 
         assert!(matches!(reply[0], HandleStunReply::StunResponse(_, _)));
+        assert!(agent.request_transaction(transaction_id).is_none());
+        assert!(agent.mut_request_transaction(transaction_id).is_none());
 
         let data = vec![42; 8];
         let transmit = agent.send_data(&data, remote_addr);
         assert_eq!(transmit.data(), &data);
         assert_eq!(transmit.from, local_addr);
         assert_eq!(transmit.to, remote_addr);
+
+        let data = vec![20; 4];
+        let mut replies = agent.handle_incoming_data(&data, remote_addr).unwrap();
+        let HandleStunReply::Data(received) = replies.remove(0) else {
+            unreachable!();
+        };
+        assert_eq!(data, received);
     }
 
     #[test]
@@ -956,6 +1069,7 @@ pub(crate) mod tests {
         agent.set_remote_credentials(remote_credentials.clone().into());
 
         let mut msg = Message::new_request(BINDING);
+        let transaction_id = msg.transaction_id();
         msg.add_message_integrity(&local_credentials.into(), IntegrityAlgorithm::Sha1)
             .unwrap();
         let transmit = agent.send(msg, remote_addr).unwrap();
@@ -975,6 +1089,13 @@ pub(crate) mod tests {
         let reply = agent.handle_incoming_data(&data, to).unwrap();
         // reply is ignored as it does not have credentials
         assert!(reply.is_empty());
+        assert!(agent.request_transaction(transaction_id).is_some());
+        assert!(agent.mut_request_transaction(transaction_id).is_some());
+
+        // unvalidated peer data should be dropped
+        let data = vec![20; 4];
+        let replies = agent.handle_incoming_data(&data, remote_addr).unwrap();
+        assert!(replies.is_empty());
     }
 
     #[test]
@@ -1013,6 +1134,61 @@ pub(crate) mod tests {
         let reply = agent.handle_incoming_data(&data, to).unwrap();
         // reply is ignored as it does not have credentials
         assert!(reply.is_empty());
+
+        // unvalidated peer data should be dropped
+        let data = vec![20; 4];
+        let replies = agent.handle_incoming_data(&data, remote_addr).unwrap();
+        assert!(replies.is_empty());
+    }
+
+    #[test]
+    fn duplicate_response_ignored() {
+        init();
+        let local_addr = "10.0.0.1:12345".parse().unwrap();
+        let remote_addr = "10.0.0.2:3478".parse().unwrap();
+
+        let mut agent = StunAgent::builder(TransportType::Udp, local_addr).build();
+
+        // unvalidated peer data should be dropped
+        let data = vec![20; 4];
+        let replies = agent.handle_incoming_data(&data, remote_addr).unwrap();
+        assert!(replies.is_empty());
+
+        let msg = Message::new_request(BINDING);
+        let transmit = agent.send(msg, remote_addr).unwrap();
+
+        let request = Message::from_bytes(&transmit.data).unwrap();
+
+        let mut response = Message::new_success(&request);
+        response
+            .add_attribute(XorMappedAddress::new(
+                transmit.from,
+                request.transaction_id(),
+            ))
+            .unwrap();
+
+        let data = response.to_bytes();
+        let to = transmit.to;
+        let reply = agent.handle_incoming_data(&data, to).unwrap();
+
+        assert!(matches!(reply[0], HandleStunReply::StunResponse(_, _)));
+
+        let data = vec![42; 8];
+        let transmit = agent.send_data(&data, remote_addr);
+        assert_eq!(transmit.data(), &data);
+        assert_eq!(transmit.from, local_addr);
+        assert_eq!(transmit.to, remote_addr);
+
+        let data = vec![20; 4];
+        let mut replies = agent.handle_incoming_data(&data, remote_addr).unwrap();
+        let HandleStunReply::Data(received) = replies.remove(0) else {
+            unreachable!();
+        };
+        assert_eq!(data, received);
+
+        let data = response.to_bytes();
+        let reply = agent.handle_incoming_data(&data, to).unwrap();
+        assert!(reply.is_empty());
     }
 
     #[test]
@@ -1041,6 +1217,57 @@ pub(crate) mod tests {
 
         let ret = agent.poll(now);
         assert!(matches!(ret, StunAgentPollRet::WaitUntil(_)));
+
+        let data = vec![42; 8];
+        let transmit = agent.send_data(&data, remote_addr);
+        assert_eq!(&transmit.data()[2..], &data);
+        assert_eq!(transmit.from, local_addr);
+        assert_eq!(transmit.to, remote_addr);
+
+        let data = vec![0, 2, 4, 8];
+        let mut replies = agent.handle_incoming_data(&data, remote_addr).unwrap();
+        let HandleStunReply::Data(received) = replies.remove(0) else {
+            unreachable!();
+        };
+        assert_eq!(&data[2..], received);
+    }
+
+    #[test]
+    fn tcp_data_before_request() {
+        init();
+        let local_addr = "127.0.0.1:2000".parse().unwrap();
+        let remote_addr = "127.0.0.1:1000".parse().unwrap();
+        let mut agent = StunAgent::builder(TransportType::Tcp, local_addr)
+            .remote_addr(remote_addr)
+            .build();
+        let data = [0, 2, 42, 42];
+
+        assert!(matches!(
+            agent.handle_incoming_data(&data, remote_addr),
+            Err(StunError::ProtocolViolation)
+        ));
+    }
+
+    #[test]
+    fn tcp_split_recv() {
+        init();
+        let local_addr = "127.0.0.1:2000".parse().unwrap();
+        let remote_addr = "127.0.0.1:1000".parse().unwrap();
+        let mut agent = StunAgent::builder(TransportType::Tcp, local_addr)
+            .remote_addr(remote_addr)
+            .build();
+        let msg = Message::new_request(BINDING);
+
+        let msg_data = msg.to_bytes();
+        let mut data = Vec::with_capacity(msg_data.len() + 2);
+        data.resize(2, 0);
+        BigEndian::write_u16(&mut data[..2], msg_data.len() as u16);
+        data.extend(msg_data);
+
+        let ret = agent.handle_incoming_data(&data[..8], remote_addr).unwrap();
+        assert!(ret.is_empty());
+        let ret = agent.handle_incoming_data(&data[8..], remote_addr).unwrap();
+        assert!(matches!(ret[0], HandleStunReply::IncomingStun(_)));
     }
 
     #[test]
@@ -1066,6 +1293,40 @@ pub(crate) mod tests {
             unreachable!();
         };
         assert_eq!(request.transaction_id(), transaction_id);
+        assert!(agent.request_transaction(transaction_id).is_none());
+        assert!(agent.mut_request_transaction(transaction_id).is_none());
+    }
+
+    #[test]
+    fn request_cancel_send() {
+        let local_addr = "10.0.0.1:12345".parse().unwrap();
+        let remote_addr = "10.0.0.2:3478".parse().unwrap();
+
+        let mut agent = StunAgent::builder(TransportType::Udp, local_addr).build();
+
+        let msg = Message::new_request(BINDING);
+        let transaction_id = msg.transaction_id();
+        let _transmit = agent.send(msg, remote_addr).unwrap();
+
+        let mut request = agent.mut_request_transaction(transaction_id).unwrap();
+        assert_eq!(request.request().transaction_id(), transaction_id);
+        assert_eq!(request.agent().local_addr(), local_addr);
+        assert_eq!(request.mut_agent().local_addr(), local_addr);
+        assert_eq!(request.peer_address(), remote_addr);
+        request.cancel_retransmissions();
+
+        let mut now = Instant::now();
+        loop {
+            match agent.poll(now) {
+                StunAgentPollRet::WaitUntil(new_now) => {
+                    now = new_now;
+                }
+                StunAgentPollRet::TransactionCancelled(_) => break,
+                _ => unreachable!(),
+            }
+        }
+        assert!(agent.request_transaction(transaction_id).is_none());
+        assert!(agent.mut_request_transaction(transaction_id).is_none());
     }
 
     #[test]
@@ -1103,5 +1364,49 @@ pub(crate) mod tests {
         };
         assert_eq!(request.transaction_id(), transaction_id);
         assert_eq!(response.transaction_id(), transaction_id);
+
+        let data = vec![20; 4];
+        let mut replies = agent.handle_incoming_data(&data, remote_addr).unwrap();
+        let HandleStunReply::Data(received) = replies.remove(0) else {
+            unreachable!();
+        };
+        assert_eq!(data, received);
+    }
+
+    #[test]
+    fn incoming_request() {
+        let local_addr = "10.0.0.1:12345".parse().unwrap();
+        let remote_addr = "10.0.0.2:3478".parse().unwrap();
+
+        let mut agent = StunAgent::builder(TransportType::Udp, local_addr).build();
+
+        let msg = Message::new_request(BINDING);
+        let data = msg.to_bytes();
+        let HandleStunReply::IncomingStun(request) = agent
+            .handle_incoming_data(&data, remote_addr)
+            .unwrap()
+            .remove(0)
+        else {
+            unreachable!()
+        };
+        assert_eq!(msg.transaction_id(), request.transaction_id());
+    }
+
+    #[test]
+    fn data_access() {
+        let array = [0, 1, 2, 3];
+        let borrowed_data = Data::from(array.as_slice());
+        assert_eq!(array.as_slice(), &*borrowed_data);
+        let owned_data = borrowed_data.into_owned();
+        assert_eq!(array.as_slice(), &*owned_data);
+        let Data::Owned(owned) = owned_data else {
+            unreachable!();
+        };
+        let owned = DataOwned::take(owned);
+        assert_eq!(array.as_slice(), &*owned);
+        let data = Data::from(owned);
+        assert_eq!(array.as_slice(), &*data);
+        let borrowed = DataSlice::from(&*data);
+        assert_eq!(array.as_slice(), &*borrowed);
     }
 }