From f780504d74e5bebc0d298fc45c708417a6c6ba29 Mon Sep 17 00:00:00 2001
From: Lars Eggert <lars@eggert.org>
Date: Fri, 24 May 2024 12:10:30 +0300
Subject: [PATCH] fix: Use correct ECN mark when sending datagram

EcnInfo::on_packet_sent may update the ECN state and can hence change it.

Also throw in some unrelated fixes to make the ECN log output more
quiet, and to test/test.sh.
---
 neqo-transport/src/ecn.rs  | 2 +-
 neqo-transport/src/path.rs | 6 +++++-
 test/test.sh               | 7 ++++---
 3 files changed, 10 insertions(+), 5 deletions(-)

diff --git a/neqo-transport/src/ecn.rs b/neqo-transport/src/ecn.rs
index 84eee9b3f8..10cf98163c 100644
--- a/neqo-transport/src/ecn.rs
+++ b/neqo-transport/src/ecn.rs
@@ -207,7 +207,7 @@ impl EcnInfo {
         } else if ecn_diff[IpTosEcn::Ect1] > 0 {
             qwarn!("ECN validation failed, ACK counted ECT(1) marks that were never sent");
             self.state = EcnValidationState::Failed;
-        } else {
+        } else if self.state != EcnValidationState::Capable {
             qinfo!("ECN validation succeeded, path is capable");
             self.state = EcnValidationState::Capable;
         }
diff --git a/neqo-transport/src/path.rs b/neqo-transport/src/path.rs
index 037a04408a..2aecd209b1 100644
--- a/neqo-transport/src/path.rs
+++ b/neqo-transport/src/path.rs
@@ -706,8 +706,12 @@ impl Path {
 
     /// Make a datagram.
     pub fn datagram<V: Into<Vec<u8>>>(&mut self, payload: V) -> Datagram {
+        // Make sure to use the TOS value from before calling EcnInfo::on_packet_sent, which may
+        // update the ECN state and can hence change it - this packet should still be sent
+        // with the current value.
+        let tos = self.tos();
         self.ecn_info.on_packet_sent();
-        Datagram::new(self.local, self.remote, self.tos(), Some(self.ttl), payload)
+        Datagram::new(self.local, self.remote, tos, Some(self.ttl), payload)
     }
 
     /// Get local address as `SocketAddr`
diff --git a/test/test.sh b/test/test.sh
index 2c22e9b6c8..195f8d7297 100755
--- a/test/test.sh
+++ b/test/test.sh
@@ -10,7 +10,6 @@
 
 set -e
 tmp=$(mktemp -d)
-trap 'rm -rf "$tmp"' EXIT
 
 cargo build --bin neqo-client --bin neqo-server
 
@@ -29,12 +28,14 @@ server="SSLKEYLOGFILE=$tmp/test.tlskey ./target/debug/neqo-server $flags $addr:$
 
 tcpdump -U -i "$iface" -w "$tmp/test.pcap" host $addr and port $port >/dev/null 2>&1 &
 tcpdump_pid=$!
+trap 'rm -rf "$tmp"; kill -USR2 $tcpdump_pid' EXIT
 
 tmux -CC \
         set-option -g default-shell "$(which bash)" \; \
-        new-session "$client && kill -USR2 $tcpdump_pid && touch $tmp/done" \; \
+        new-session "$client; kill -USR2 $tcpdump_pid; touch $tmp/done" \; \
         split-window -h "$server" \; \
         split-window -v -f "\
-                until [ -e $tmp/done ]; do sleep 1; done && \
+                until [ -e $tmp/done ]; do sleep 1; done; \
+                echo $tmp; ls -l $tmp; echo; \
                 tshark -r $tmp/test.pcap -o tls.keylog_file:$tmp/test.tlskey" \; \
         set remain-on-exit on