udp: split AnnounceResponse into V4 and V6 versions

aquatic_udp_protocol/src/common.rs

@@ -1,4 +1,4 @@
-use std::net::IpAddr;
+use std::net::{Ipv4Addr, Ipv6Addr};
 
 #[derive(PartialEq, Eq, Hash, Clone, Copy, Debug)]
 pub enum IpVersion {
@@ -37,8 +37,14 @@ pub struct PeerId(pub [u8; 20]);
 pub struct PeerKey(pub u32);
 
 #[derive(Hash, PartialEq, Eq, Clone, Debug)]
-pub struct ResponsePeer {
-    pub ip_address: IpAddr,
+pub struct ResponsePeerIpv4 {
+    pub ip_address: Ipv4Addr,
+    pub port: Port,
+}
+
+#[derive(Hash, PartialEq, Eq, Clone, Debug)]
+pub struct ResponsePeerIpv6 {
+    pub ip_address: Ipv6Addr,
     pub port: Port,
 }
 
@@ -80,11 +86,21 @@ impl quickcheck::Arbitrary for PeerId {
 }
 
 #[cfg(test)]
-impl quickcheck::Arbitrary for ResponsePeer {
+impl quickcheck::Arbitrary for ResponsePeerIpv4 {
     fn arbitrary(g: &mut quickcheck::Gen) -> Self {
         Self {
-            ip_address: ::std::net::IpAddr::arbitrary(g),
-            port: Port(u16::arbitrary(g)),
+            ip_address: quickcheck::Arbitrary::arbitrary(g),
+            port: Port(u16::arbitrary(g).into()),
+        }
+    }
+}
+
+#[cfg(test)]
+impl quickcheck::Arbitrary for ResponsePeerIpv6 {
+    fn arbitrary(g: &mut quickcheck::Gen) -> Self {
+        Self {
+            ip_address: quickcheck::Arbitrary::arbitrary(g),
+            port: Port(u16::arbitrary(g).into()),
         }
     }
 }

aquatic_udp_protocol/src/response.rs

@@ -1,7 +1,7 @@
 use std::borrow::Cow;
 use std::convert::TryInto;
 use std::io::{self, Cursor, Write};
-use std::net::{IpAddr, Ipv4Addr, Ipv6Addr};
+use std::net::{Ipv4Addr, Ipv6Addr};
 
 use byteorder::{NetworkEndian, ReadBytesExt, WriteBytesExt};
 
@@ -21,12 +21,21 @@ pub struct ConnectResponse {
 }
 
 #[derive(PartialEq, Eq, Clone, Debug)]
-pub struct AnnounceResponse {
+pub struct AnnounceResponseIpv4 {
     pub transaction_id: TransactionId,
     pub announce_interval: AnnounceInterval,
     pub leechers: NumberOfPeers,
     pub seeders: NumberOfPeers,
-    pub peers: Vec<ResponsePeer>,
+    pub peers: Vec<ResponsePeerIpv4>,
+}
+
+#[derive(PartialEq, Eq, Clone, Debug)]
+pub struct AnnounceResponseIpv6 {
+    pub transaction_id: TransactionId,
+    pub announce_interval: AnnounceInterval,
+    pub leechers: NumberOfPeers,
+    pub seeders: NumberOfPeers,
+    pub peers: Vec<ResponsePeerIpv6>,
 }
 
 #[derive(PartialEq, Eq, Clone, Debug)]
@@ -44,7 +53,8 @@ pub struct ErrorResponse {
 #[derive(PartialEq, Eq, Clone, Debug)]
 pub enum Response {
     Connect(ConnectResponse),
-    Announce(AnnounceResponse),
+    AnnounceIpv4(AnnounceResponseIpv4),
+    AnnounceIpv6(AnnounceResponseIpv6),
     Scrape(ScrapeResponse),
     Error(ErrorResponse),
 }
@@ -55,9 +65,15 @@ impl From<ConnectResponse> for Response {
     }
 }
 
-impl From<AnnounceResponse> for Response {
-    fn from(r: AnnounceResponse) -> Self {
-        Self::Announce(r)
+impl From<AnnounceResponseIpv4> for Response {
+    fn from(r: AnnounceResponseIpv4) -> Self {
+        Self::AnnounceIpv4(r)
+    }
+}
+
+impl From<AnnounceResponseIpv6> for Response {
+    fn from(r: AnnounceResponseIpv6) -> Self {
+        Self::AnnounceIpv6(r)
     }
 }
 
@@ -81,42 +97,23 @@ impl Response {
     /// addresses. Clients seem not to support it very well, but due to a lack
     /// of alternative solutions, it is implemented here.
     #[inline]
-    pub fn write(self, bytes: &mut impl Write, ip_version: IpVersion) -> Result<(), io::Error> {
+    pub fn write(self, bytes: &mut impl Write) -> Result<(), io::Error> {
         match self {
             Response::Connect(r) => {
                 bytes.write_i32::<NetworkEndian>(0)?;
                 bytes.write_i32::<NetworkEndian>(r.transaction_id.0)?;
                 bytes.write_i64::<NetworkEndian>(r.connection_id.0)?;
             }
-            Response::Announce(r) => {
-                if ip_version == IpVersion::IPv4 {
-                    bytes.write_i32::<NetworkEndian>(1)?;
-                    bytes.write_i32::<NetworkEndian>(r.transaction_id.0)?;
-                    bytes.write_i32::<NetworkEndian>(r.announce_interval.0)?;
-                    bytes.write_i32::<NetworkEndian>(r.leechers.0)?;
-                    bytes.write_i32::<NetworkEndian>(r.seeders.0)?;
+            Response::AnnounceIpv4(r) => {
+                bytes.write_i32::<NetworkEndian>(1)?;
+                bytes.write_i32::<NetworkEndian>(r.transaction_id.0)?;
+                bytes.write_i32::<NetworkEndian>(r.announce_interval.0)?;
+                bytes.write_i32::<NetworkEndian>(r.leechers.0)?;
+                bytes.write_i32::<NetworkEndian>(r.seeders.0)?;
 
-                    // Silently ignore peers with wrong IP version
-                    for peer in r.peers {
-                        if let IpAddr::V4(ip) = peer.ip_address {
-                            bytes.write_all(&ip.octets())?;
-                            bytes.write_u16::<NetworkEndian>(peer.port.0)?;
-                        }
-                    }
-                } else {
-                    bytes.write_i32::<NetworkEndian>(4)?;
-                    bytes.write_i32::<NetworkEndian>(r.transaction_id.0)?;
-                    bytes.write_i32::<NetworkEndian>(r.announce_interval.0)?;
-                    bytes.write_i32::<NetworkEndian>(r.leechers.0)?;
-                    bytes.write_i32::<NetworkEndian>(r.seeders.0)?;
-
-                    // Silently ignore peers with wrong IP version
-                    for peer in r.peers {
-                        if let IpAddr::V6(ip) = peer.ip_address {
-                            bytes.write_all(&ip.octets())?;
-                            bytes.write_u16::<NetworkEndian>(peer.port.0)?;
-                        }
-                    }
+                for peer in r.peers {
+                    bytes.write_all(&peer.ip_address.octets())?;
+                    bytes.write_u16::<NetworkEndian>(peer.port.0)?;
                 }
             }
             Response::Scrape(r) => {
@@ -135,6 +132,18 @@ impl Response {
 
                 bytes.write_all(r.message.as_bytes())?;
             }
+            Response::AnnounceIpv6(r) => {
+                bytes.write_i32::<NetworkEndian>(4)?;
+                bytes.write_i32::<NetworkEndian>(r.transaction_id.0)?;
+                bytes.write_i32::<NetworkEndian>(r.announce_interval.0)?;
+                bytes.write_i32::<NetworkEndian>(r.leechers.0)?;
+                bytes.write_i32::<NetworkEndian>(r.seeders.0)?;
+
+                for peer in r.peers {
+                    bytes.write_all(&peer.ip_address.octets())?;
+                    bytes.write_u16::<NetworkEndian>(peer.port.0)?;
+                }
+            }
         }
 
         Ok(())
@@ -171,17 +180,17 @@ impl Response {
                     .chunks_exact(6)
                     .map(|chunk| {
                         let ip_bytes: [u8; 4] = (&chunk[..4]).try_into().unwrap();
-                        let ip_address = IpAddr::V4(Ipv4Addr::from(ip_bytes));
+                        let ip_address = Ipv4Addr::from(ip_bytes);
                         let port = (&chunk[4..]).read_u16::<NetworkEndian>().unwrap();
 
-                        ResponsePeer {
+                        ResponsePeerIpv4 {
                             ip_address,
                             port: Port(port),
                         }
                     })
                     .collect();
 
-                Ok((AnnounceResponse {
+                Ok((AnnounceResponseIpv4 {
                     transaction_id: TransactionId(transaction_id),
                     announce_interval: AnnounceInterval(announce_interval),
                     leechers: NumberOfPeers(leechers),
@@ -244,17 +253,17 @@ impl Response {
                     .chunks_exact(18)
                     .map(|chunk| {
                         let ip_bytes: [u8; 16] = (&chunk[..16]).try_into().unwrap();
-                        let ip_address = IpAddr::V6(Ipv6Addr::from(ip_bytes));
+                        let ip_address = Ipv6Addr::from(ip_bytes);
                         let port = (&chunk[16..]).read_u16::<NetworkEndian>().unwrap();
 
-                        ResponsePeer {
+                        ResponsePeerIpv6 {
                             ip_address,
                             port: Port(port),
                         }
                     })
                     .collect();
 
-                Ok((AnnounceResponse {
+                Ok((AnnounceResponseIpv6 {
                     transaction_id: TransactionId(transaction_id),
                     announce_interval: AnnounceInterval(announce_interval),
                     leechers: NumberOfPeers(leechers),
@@ -297,10 +306,26 @@ mod tests {
         }
     }
 
-    impl quickcheck::Arbitrary for AnnounceResponse {
+    impl quickcheck::Arbitrary for AnnounceResponseIpv4 {
         fn arbitrary(g: &mut quickcheck::Gen) -> Self {
             let peers = (0..u8::arbitrary(g))
-                .map(|_| ResponsePeer::arbitrary(g))
+                .map(|_| ResponsePeerIpv4::arbitrary(g))
+                .collect();
+
+            Self {
+                transaction_id: TransactionId(i32::arbitrary(g)),
+                announce_interval: AnnounceInterval(i32::arbitrary(g)),
+                leechers: NumberOfPeers(i32::arbitrary(g)),
+                seeders: NumberOfPeers(i32::arbitrary(g)),
+                peers,
+            }
+        }
+    }
+
+    impl quickcheck::Arbitrary for AnnounceResponseIpv6 {
+        fn arbitrary(g: &mut quickcheck::Gen) -> Self {
+            let peers = (0..u8::arbitrary(g))
+                .map(|_| ResponsePeerIpv6::arbitrary(g))
                 .collect();
 
             Self {
@@ -326,10 +351,10 @@ mod tests {
         }
     }
 
-    fn same_after_conversion(response: Response, ip_version: IpVersion) -> bool {
+    fn same_after_conversion(response: Response) -> bool {
         let mut buf = Vec::new();
 
-        response.clone().write(&mut buf, ip_version).unwrap();
+        response.clone().write(&mut buf).unwrap();
         let r2 = Response::from_bytes(&buf[..]).unwrap();
 
         let success = response == r2;
@@ -343,24 +368,21 @@ mod tests {
 
     #[quickcheck]
     fn test_connect_response_convert_identity(response: ConnectResponse) -> bool {
-        same_after_conversion(response.into(), IpVersion::IPv4)
+        same_after_conversion(response.into())
     }
 
     #[quickcheck]
-    fn test_announce_response_convert_identity(data: (AnnounceResponse, IpVersion)) -> bool {
-        let mut r = data.0;
+    fn test_announce_response_ipv4_convert_identity(response: AnnounceResponseIpv4) -> bool {
+        same_after_conversion(response.into())
+    }
 
-        if data.1 == IpVersion::IPv4 {
-            r.peers.retain(|peer| peer.ip_address.is_ipv4());
-        } else {
-            r.peers.retain(|peer| peer.ip_address.is_ipv6());
-        }
-
-        same_after_conversion(r.into(), data.1)
+    #[quickcheck]
+    fn test_announce_response_ipv6_convert_identity(response: AnnounceResponseIpv6) -> bool {
+        same_after_conversion(response.into())
     }
 
     #[quickcheck]
     fn test_scrape_response_convert_identity(response: ScrapeResponse) -> bool {
-        same_after_conversion(response.into(), IpVersion::IPv4)
+        same_after_conversion(response.into())
     }
 }