1 files changed, 52 insertions, 25 deletions

diff --git a/third_party/rust/neqo-transport/src/path.rs b/third_party/rust/neqo-transport/src/path.rs
index 50e458ff36..0e4c82b1ca 100644
--- a/third_party/rust/neqo-transport/src/path.rs
+++ b/third_party/rust/neqo-transport/src/path.rs
@@ -22,6 +22,7 @@ use crate::{
     ackrate::{AckRate, PeerAckDelay},
     cc::CongestionControlAlgorithm,
     cid::{ConnectionId, ConnectionIdRef, ConnectionIdStore, RemoteConnectionIdEntry},
+    ecn::{EcnCount, EcnInfo},
     frame::{FRAME_TYPE_PATH_CHALLENGE, FRAME_TYPE_PATH_RESPONSE, FRAME_TYPE_RETIRE_CONNECTION_ID},
     packet::PacketBuilder,
     recovery::RecoveryToken,
@@ -145,15 +146,8 @@ impl Paths {
             })
     }
 
-    /// Get a reference to the primary path.  This will assert if there is no primary
-    /// path, which happens at a server prior to receiving a valid Initial packet
-    /// from a client.  So be careful using this method.
-    pub fn primary(&self) -> PathRef {
-        self.primary_fallible().unwrap()
-    }
-
-    /// Get a reference to the primary path.  Use this prior to handshake completion.
-    pub fn primary_fallible(&self) -> Option<PathRef> {
+    /// Get a reference to the primary path, if one exists.
+    pub fn primary(&self) -> Option<PathRef> {
         self.primary.clone()
     }
 
@@ -242,6 +236,11 @@ impl Paths {
     /// Returns `true` if the path was migrated.
     pub fn migrate(&mut self, path: &PathRef, force: bool, now: Instant) -> bool {
         debug_assert!(!self.is_temporary(path));
+        let baseline = self.primary().map_or_else(
+            || EcnInfo::default().baseline(),
+            |p| p.borrow().ecn_info.baseline(),
+        );
+        path.borrow_mut().set_ecn_baseline(baseline);
         if force || path.borrow().is_valid() {
             path.borrow_mut().set_valid(now);
             mem::drop(self.select_primary(path));
@@ -307,7 +306,6 @@ impl Paths {
     /// Set the identified path to be primary.
     /// This panics if `make_permanent` hasn't been called.
     pub fn handle_migration(&mut self, path: &PathRef, remote: SocketAddr, now: Instant) {
-        qtrace!([self.primary().borrow()], "handle_migration");
         // The update here needs to match the checks in `Path::received_on`.
         // Here, we update the remote port number to match the source port on the
         // datagram that was received.  This ensures that we send subsequent
@@ -425,10 +423,10 @@ impl Paths {
             stats.retire_connection_id += 1;
         }
 
-        // Write out any ACK_FREQUENCY frames.
-        self.primary()
-            .borrow_mut()
-            .write_cc_frames(builder, tokens, stats);
+        if let Some(path) = self.primary() {
+            // Write out any ACK_FREQUENCY frames.
+            path.borrow_mut().write_cc_frames(builder, tokens, stats);
+        }
     }
 
     pub fn lost_retire_cid(&mut self, lost: u64) {
@@ -440,11 +438,15 @@ impl Paths {
     }
 
     pub fn lost_ack_frequency(&mut self, lost: &AckRate) {
-        self.primary().borrow_mut().lost_ack_frequency(lost);
+        if let Some(path) = self.primary() {
+            path.borrow_mut().lost_ack_frequency(lost);
+        }
     }
 
     pub fn acked_ack_frequency(&mut self, acked: &AckRate) {
-        self.primary().borrow_mut().acked_ack_frequency(acked);
+        if let Some(path) = self.primary() {
+            path.borrow_mut().acked_ack_frequency(acked);
+        }
     }
 
     /// Get an estimate of the RTT on the primary path.
@@ -454,7 +456,7 @@ impl Paths {
         // make a new RTT esimate and interrogate that.
         // That is more expensive, but it should be rare and breaking encapsulation
         // is worse, especially as this is only used in tests.
-        self.primary_fallible()
+        self.primary()
             .map_or(RttEstimate::default().estimate(), |p| {
                 p.borrow().rtt().estimate()
             })
@@ -532,8 +534,6 @@ pub struct Path {
     rtt: RttEstimate,
     /// A packet sender for the path, which includes congestion control and a pacer.
     sender: PacketSender,
-    /// The DSCP/ECN marking to use for outgoing packets on this path.
-    tos: IpTos,
     /// The IP TTL to use for outgoing packets on this path.
     ttl: u8,
 
@@ -543,7 +543,8 @@ pub struct Path {
     received_bytes: usize,
     /// The number of bytes sent on this path.
     sent_bytes: usize,
-
+    /// The ECN-related state for this path (see RFC9000, Section 13.4 and Appendix A.4)
+    ecn_info: EcnInfo,
     /// For logging of events.
     qlog: NeqoQlog,
 }
@@ -572,14 +573,23 @@ impl Path {
             challenge: None,
             rtt: RttEstimate::default(),
             sender,
-            tos: IpTos::default(), // TODO: Default to Ect0 when ECN is supported.
-            ttl: 64,               // This is the default TTL on many OSes.
+            ttl: 64, // This is the default TTL on many OSes.
             received_bytes: 0,
             sent_bytes: 0,
+            ecn_info: EcnInfo::default(),
             qlog,
         }
     }
 
+    pub fn set_ecn_baseline(&mut self, baseline: EcnCount) {
+        self.ecn_info.set_baseline(baseline);
+    }
+
+    /// Return the DSCP/ECN marking to use for outgoing packets on this path.
+    pub fn tos(&self) -> IpTos {
+        self.ecn_info.ecn_mark().into()
+    }
+
     /// Whether this path is the primary or current path for the connection.
     pub fn is_primary(&self) -> bool {
         self.primary
@@ -695,8 +705,9 @@ impl Path {
     }
 
     /// Make a datagram.
-    pub fn datagram<V: Into<Vec<u8>>>(&self, payload: V) -> Datagram {
-        Datagram::new(self.local, self.remote, self.tos, Some(self.ttl), payload)
+    pub fn datagram<V: Into<Vec<u8>>>(&mut self, payload: V) -> Datagram {
+        self.ecn_info.on_packet_sent();
+        Datagram::new(self.local, self.remote, self.tos(), Some(self.ttl), payload)
     }
 
     /// Get local address as `SocketAddr`
@@ -959,8 +970,24 @@ impl Path {
     }
 
     /// Record packets as acknowledged with the sender.
-    pub fn on_packets_acked(&mut self, acked_pkts: &[SentPacket], now: Instant) {
+    pub fn on_packets_acked(
+        &mut self,
+        acked_pkts: &[SentPacket],
+        ack_ecn: Option<EcnCount>,
+        now: Instant,
+    ) {
         debug_assert!(self.is_primary());
+
+        let ecn_ce_received = self.ecn_info.on_packets_acked(acked_pkts, ack_ecn);
+        if ecn_ce_received {
+            let cwnd_reduced = self
+                .sender
+                .on_ecn_ce_received(acked_pkts.first().expect("must be there"));
+            if cwnd_reduced {
+                self.rtt.update_ack_delay(self.sender.cwnd(), self.mtu());
+            }
+        }
+
         self.sender.on_packets_acked(acked_pkts, &self.rtt, now);
     }