fix: correctly unmap IPv4-mapped local addresses (#2454)

Author: camshaft

quic/s2n-quic-core/src/inet/ip.rs

@@ -35,6 +35,14 @@ impl IpAddress {
         }
     }
 
+    /// Returns `true` if the two addresses are equal from a network perspective.
+    ///
+    /// This will unmap IPv4-mapped addresses to IpV4 tagged enum values
+    #[inline]
+    pub fn unmapped_eq(&self, other: &Self) -> bool {
+        self.unmap() == other.unmap()
+    }
+
     /// Converts the IP address into IPv6 if it is IPv4, otherwise the address is unchanged
     #[inline]
     #[must_use]
@@ -167,6 +175,14 @@ impl SocketAddress {
             Self::IpV6(addr) => addr.unmap(),
         }
     }
+
+    /// Returns `true` if the two addresses are equal from a network perspective.
+    ///
+    /// This will unmap IPv4-mapped addresses to IpV4 tagged enum values
+    #[inline]
+    pub fn unmapped_eq(&self, other: &Self) -> bool {
+        self.unmap() == other.unmap()
+    }
 }
 
 impl Default for SocketAddress {

quic/s2n-quic-core/src/packet/interceptor.rs

@@ -5,6 +5,7 @@ use crate::{
     event::api::{SocketAddress, Subject},
     havoc,
     packet::number::{PacketNumber, PacketNumberSpace},
+    path::{LocalAddress, RemoteAddress},
     time::Timestamp,
     varint::VarInt,
 };
@@ -45,9 +46,15 @@ pub trait Interceptor: 'static + Send {
     }
 
     #[inline(always)]
-    fn intercept_rx_remote_port(&mut self, subject: &Subject, port: &mut u16) {
+    fn intercept_rx_local_address(&mut self, subject: &Subject, addr: &mut LocalAddress) {
         let _ = subject;
-        let _ = port;
+        let _ = addr;
+    }
+
+    #[inline(always)]
+    fn intercept_rx_remote_address(&mut self, subject: &Subject, addr: &mut RemoteAddress) {
+        let _ = subject;
+        let _ = addr;
     }
 
     #[inline(always)]
@@ -116,9 +123,15 @@ where
     }
 
     #[inline(always)]
-    fn intercept_rx_remote_port(&mut self, subject: &Subject, port: &mut u16) {
-        self.0.intercept_rx_remote_port(subject, port);
-        self.1.intercept_rx_remote_port(subject, port);
+    fn intercept_rx_local_address(&mut self, subject: &Subject, addr: &mut LocalAddress) {
+        self.0.intercept_rx_local_address(subject, addr);
+        self.1.intercept_rx_local_address(subject, addr);
+    }
+
+    #[inline(always)]
+    fn intercept_rx_remote_address(&mut self, subject: &Subject, addr: &mut RemoteAddress) {
+        self.0.intercept_rx_remote_address(subject, addr);
+        self.1.intercept_rx_remote_address(subject, addr);
     }
 
     #[inline(always)]
@@ -188,8 +201,10 @@ where
     R: 'static + Send + havoc::Random,
 {
     #[inline]
-    fn intercept_rx_remote_port(&mut self, _subject: &Subject, port: &mut u16) {
-        *port = self.port.havoc_u16(&mut self.random, *port);
+    fn intercept_rx_remote_address(&mut self, _subject: &Subject, addr: &mut RemoteAddress) {
+        let port = addr.port();
+        let port = self.port.havoc_u16(&mut self.random, port);
+        addr.set_port(port);
     }
 
     #[inline]
@@ -229,9 +244,16 @@ where
 
 impl<T: Interceptor> Interceptor for Option<T> {
     #[inline]
-    fn intercept_rx_remote_port(&mut self, subject: &Subject, port: &mut u16) {
+    fn intercept_rx_local_address(&mut self, subject: &Subject, addr: &mut LocalAddress) {
+        if let Some(inner) = self.as_mut() {
+            inner.intercept_rx_local_address(subject, addr)
+        }
+    }
+
+    #[inline]
+    fn intercept_rx_remote_address(&mut self, subject: &Subject, addr: &mut RemoteAddress) {
         if let Some(inner) = self.as_mut() {
-            inner.intercept_rx_remote_port(subject, port)
+            inner.intercept_rx_remote_address(subject, addr)
         }
     }
 

quic/s2n-quic-core/src/path/mod.rs

@@ -59,17 +59,20 @@ pub trait Handle: 'static + Copy + Send + fmt::Debug {
     /// Returns the remote address for the given handle
     fn remote_address(&self) -> RemoteAddress;
 
-    /// Updates the remote port to the given value
-    fn set_remote_port(&mut self, port: u16);
+    /// Updates the remote address to the given value
+    fn set_remote_address(&mut self, addr: RemoteAddress);
 
     /// Returns the local address for the given handle
     fn local_address(&self) -> LocalAddress;
 
+    /// Updates the local address to the given value
+    fn set_local_address(&mut self, addr: LocalAddress);
+
     /// Returns `true` if the two handles are equal from a network perspective
     ///
     /// This function is used to determine if a connection has migrated to another
     /// path.
-    fn eq(&self, other: &Self) -> bool;
+    fn unmapped_eq(&self, other: &Self) -> bool;
 
     /// Returns `true` if the two handles are strictly equal to each other, i.e.
     /// byte-for-byte.
@@ -95,6 +98,16 @@ macro_rules! impl_addr {
         #[cfg_attr(kani, derive(kani::Arbitrary))]
         pub struct $name(pub SocketAddress);
 
+        impl $name {
+            /// Returns `true` if the two addresses are equal from a network perspective.
+            ///
+            /// This will unmap IPv4-mapped addresses to IpV4 tagged enum values
+            #[inline]
+            pub fn unmapped_eq(&self, other: &Self) -> bool {
+                self.0.unmapped_eq(&other.0)
+            }
+        }
+
         impl From<event::api::SocketAddress<'_>> for $name {
             #[inline]
             fn from(value: event::api::SocketAddress<'_>) -> Self {
@@ -162,8 +175,8 @@ impl Handle for RemoteAddress {
     }
 
     #[inline]
-    fn set_remote_port(&mut self, port: u16) {
-        self.0.set_port(port)
+    fn set_remote_address(&mut self, addr: RemoteAddress) {
+        *self = addr;
     }
 
     #[inline]
@@ -172,8 +185,13 @@ impl Handle for RemoteAddress {
     }
 
     #[inline]
-    fn eq(&self, other: &Self) -> bool {
-        PartialEq::eq(&self.unmap(), &other.unmap())
+    fn set_local_address(&mut self, _addr: LocalAddress) {
+        // nothing to update
+    }
+
+    #[inline]
+    fn unmapped_eq(&self, other: &Self) -> bool {
+        Self::unmapped_eq(self, other)
     }
 
     #[inline]
@@ -218,8 +236,8 @@ impl Handle for Tuple {
     }
 
     #[inline]
-    fn set_remote_port(&mut self, port: u16) {
-        self.remote_address.set_port(port)
+    fn set_remote_address(&mut self, addr: RemoteAddress) {
+        self.remote_address = addr;
     }
 
     #[inline]
@@ -228,9 +246,14 @@ impl Handle for Tuple {
     }
 
     #[inline]
-    fn eq(&self, other: &Self) -> bool {
-        PartialEq::eq(&self.local_address.unmap(), &other.local_address.unmap())
-            && Handle::eq(&self.remote_address, &other.remote_address)
+    fn set_local_address(&mut self, addr: LocalAddress) {
+        self.local_address = addr;
+    }
+
+    #[inline]
+    fn unmapped_eq(&self, other: &Self) -> bool {
+        self.local_address.unmapped_eq(&other.local_address)
+            && self.remote_address.unmapped_eq(&other.remote_address)
     }
 
     #[inline]

quic/s2n-quic-core/src/xdp/encoder.rs

@@ -409,7 +409,7 @@ mod tests {
 
             header.path.swap();
 
-            assert!(Handle::eq(&header.path, &message.path));
+            assert!(header.path.unmapped_eq(&message.path));
             assert_eq!(header.ecn, message.ecn);
             assert_eq!(payload.into_less_safe_slice(), &message.payload);
         });

quic/s2n-quic-core/src/xdp/path.rs

@@ -97,8 +97,9 @@ impl Handle for Tuple {
     }
 
     #[inline]
-    fn set_remote_port(&mut self, port: u16) {
-        self.remote_address.port = port;
+    fn set_remote_address(&mut self, addr: path::RemoteAddress) {
+        self.remote_address.ip = addr.ip();
+        self.remote_address.port = addr.port();
     }
 
     #[inline]
@@ -107,7 +108,13 @@ impl Handle for Tuple {
     }
 
     #[inline]
-    fn eq(&self, other: &Self) -> bool {
+    fn set_local_address(&mut self, addr: path::LocalAddress) {
+        self.local_address.ip = addr.ip();
+        self.local_address.port = addr.port();
+    }
+
+    #[inline]
+    fn unmapped_eq(&self, other: &Self) -> bool {
         // TODO only compare everything if the other is all filled out
         PartialEq::eq(&self.local_address.unmap(), &other.local_address.unmap())
             && PartialEq::eq(&self.remote_address.unmap(), &other.remote_address.unmap())

quic/s2n-quic-platform/src/message/msg/handle.rs

@@ -5,7 +5,8 @@ use super::ext::Ext as _;
 use crate::{features, message::cmsg::Encoder};
 use libc::msghdr;
 use s2n_quic_core::{
-    inet::{AncillaryData, SocketAddressV4},
+    ensure,
+    inet::{AncillaryData, SocketAddressV4, Unspecified},
     path::{self, LocalAddress, RemoteAddress},
 };
 
@@ -54,8 +55,8 @@ impl path::Handle for Handle {
     }
 
     #[inline]
-    fn set_remote_port(&mut self, port: u16) {
-        self.remote_address.0.set_port(port);
+    fn set_remote_address(&mut self, addr: RemoteAddress) {
+        self.remote_address = addr;
     }
 
     #[inline]
@@ -64,15 +65,41 @@ impl path::Handle for Handle {
     }
 
     #[inline]
-    fn eq(&self, other: &Self) -> bool {
-        let mut eq = true;
+    fn set_local_address(&mut self, addr: LocalAddress) {
+        self.local_address = addr;
+    }
+
+    #[inline]
+    fn unmapped_eq(&self, other: &Self) -> bool {
+        ensure!(
+            self.remote_address.unmapped_eq(&other.remote_address),
+            false
+        );
 
         // only compare local addresses if the OS returns them
-        if features::pktinfo::IS_SUPPORTED {
-            eq &= self.local_address.eq(&other.local_address);
+        ensure!(features::pktinfo::IS_SUPPORTED, true);
+
+        // Make sure to only compare the fields if they're both set
+        //
+        // This avoids cases where we don't have the full context for the local address and find it
+        // out with a later packet.
+        if !self.local_address.ip().is_unspecified() && !other.local_address.ip().is_unspecified() {
+            ensure!(
+                self.local_address
+                    .ip()
+                    .unmapped_eq(&other.local_address.ip()),
+                false
+            );
         }
 
-        eq && path::Handle::eq(&self.remote_address, &other.remote_address)
+        if self.local_address.port() > 0 && other.local_address.port() > 0 {
+            ensure!(
+                self.local_address.port() == other.local_address.port(),
+                false
+            );
+        }
+
+        true
     }
 
     #[inline]
@@ -92,3 +119,50 @@ impl path::Handle for Handle {
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use crate::message::msg::Handle;
+    use s2n_quic_core::{
+        inet::{IpAddress, IpV4Address},
+        path::{Handle as _, LocalAddress},
+    };
+
+    /// Checks that unmapped_eq is correct independent of argument ordering
+    fn reflexive_check(a: Handle, b: Handle) {
+        assert!(a.unmapped_eq(&b));
+        assert!(b.unmapped_eq(&a));
+    }
+
+    #[test]
+    fn unmapped_eq_test() {
+        // All of these values should be considered equivalent for local addresses
+        let ips: &[IpAddress] = &[
+            // if we have an unspecified IP address then don't consider it for equality
+            IpV4Address::new([0, 0, 0, 0]).into(),
+            // a regular IPv4 IP should match the IPv4-mapped into IPv6
+            IpV4Address::new([1, 1, 1, 1]).into(),
+            IpV4Address::new([1, 1, 1, 1]).to_ipv6_mapped().into(),
+        ];
+        let ports = [0u16, 4440];
+
+        for ip_a in ips {
+            for ip_b in ips {
+                for port_a in ports {
+                    for port_b in ports {
+                        reflexive_check(
+                            Handle {
+                                remote_address: Default::default(),
+                                local_address: LocalAddress::from(ip_a.with_port(port_a)),
+                            },
+                            Handle {
+                                remote_address: Default::default(),
+                                local_address: LocalAddress::from(ip_b.with_port(port_b)),
+                            },
+                        );
+                    }
+                }
+            }
+        }
+    }
+}

quic/s2n-quic-qns/src/intercept.rs

@@ -7,8 +7,11 @@ use s2n_codec::encoder::scatter;
 use s2n_quic_core::{
     event::api::Subject,
     havoc::{self, Strategy as _, *},
-    packet,
-    packet::interceptor::{DecoderBufferMut, Havoc},
+    packet::{
+        self,
+        interceptor::{DecoderBufferMut, Havoc},
+    },
+    path::RemoteAddress,
 };
 use structopt::StructOpt;
 
@@ -113,10 +116,10 @@ impl Interceptor {
 
 impl packet::interceptor::Interceptor for Interceptor {
     #[inline]
-    fn intercept_rx_remote_port(&mut self, subject: &Subject, port: &mut u16) {
+    fn intercept_rx_remote_address(&mut self, subject: &Subject, addr: &mut RemoteAddress) {
         if self.port {
             self.strategy_for(subject)
-                .intercept_rx_remote_port(subject, port)
+                .intercept_rx_remote_address(subject, addr)
         }
     }