Repository: quic-go/quic-go
Pull request #3829 (2023-05-14 05:24:29)

wire: save ECN counts on the ACK frame
I'm not sure why we haven't done this before. We're not using the counts, but the fields are there on the frame, and at the very least, we might correctly log them.


Patch:
diff --git a/.github/workflows/unit.yml b/.github/workflows/unit.yml
index 1b678555..eb4ce725 100644
--- a/.github/workflows/unit.yml
+++ b/.github/workflows/unit.yml
@@ -44,5 +44,5 @@ jobs:
       - name: Upload coverage to Codecov
         uses: codecov/codecov-action@v3
         with:
-          files: coverage.txt coverage-root.txt
+          files: coverage.txt,coverage-root.txt
           env_vars: OS=${{ matrix.os }}, GO=${{ matrix.go }}
diff --git a/http3/roundtrip.go b/http3/roundtrip.go
index 53885e1e..eef93c28 100644
--- a/http3/roundtrip.go
+++ b/http3/roundtrip.go
@@ -10,6 +10,7 @@ import (
 	"net/http"
 	"strings"
 	"sync"
+	"sync/atomic"
 
 	"golang.org/x/net/http/httpguts"
 
@@ -25,6 +26,11 @@ type roundTripCloser interface {
 	io.Closer
 }
 
+type roundTripCloserWithCount struct {
+	roundTripCloser
+	useCount atomic.Int64
+}
+
 // RoundTripper implements the http.RoundTripper interface
 type RoundTripper struct {
 	mutex sync.Mutex
@@ -82,7 +88,7 @@ type RoundTripper struct {
 	MaxResponseHeaderBytes int64
 
 	newClient func(hostname string, tlsConf *tls.Config, opts *roundTripperOpts, conf *quic.Config, dialer dialFunc) (roundTripCloser, error) // so we can mock it in tests
-	clients   map[string]roundTripCloser
+	clients   map[string]*roundTripCloserWithCount
 	udpConn   *net.UDPConn
 }
 
@@ -143,6 +149,7 @@ func (r *RoundTripper) RoundTripOpt(req *http.Request, opt RoundTripOpt) (*http.
 	if err != nil {
 		return nil, err
 	}
+	defer cl.useCount.Add(-1)
 	rsp, err := cl.RoundTripOpt(req, opt)
 	if err != nil {
 		r.removeClient(hostname)
@@ -160,12 +167,12 @@ func (r *RoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
 	return r.RoundTripOpt(req, RoundTripOpt{})
 }
 
-func (r *RoundTripper) getClient(hostname string, onlyCached bool) (rtc roundTripCloser, isReused bool, err error) {
+func (r *RoundTripper) getClient(hostname string, onlyCached bool) (rtc *roundTripCloserWithCount, isReused bool, err error) {
 	r.mutex.Lock()
 	defer r.mutex.Unlock()
 
 	if r.clients == nil {
-		r.clients = make(map[string]roundTripCloser)
+		r.clients = make(map[string]*roundTripCloserWithCount)
 	}
 
 	client, ok := r.clients[hostname]
@@ -188,7 +195,7 @@ func (r *RoundTripper) getClient(hostname string, onlyCached bool) (rtc roundTri
 			}
 			dial = r.makeDialer()
 		}
-		client, err = newCl(
+		c, err := newCl(
 			hostname,
 			r.TLSClientConfig,
 			&roundTripperOpts{
@@ -204,10 +211,12 @@ func (r *RoundTripper) getClient(hostname string, onlyCached bool) (rtc roundTri
 		if err != nil {
 			return nil, false, err
 		}
+		client = &roundTripCloserWithCount{roundTripCloser: c}
 		r.clients[hostname] = client
 	} else if client.HandshakeComplete() {
 		isReused = true
 	}
+	client.useCount.Add(1)
 	return client, isReused, nil
 }
 
@@ -276,3 +285,14 @@ func (r *RoundTripper) makeDialer() func(ctx context.Context, addr string, tlsCf
 		return quicDialer(ctx, r.udpConn, udpAddr, tlsCfg, cfg)
 	}
 }
+
+func (r *RoundTripper) CloseIdleConnections() {
+	r.mutex.Lock()
+	defer r.mutex.Unlock()
+	for hostname, client := range r.clients {
+		if client.useCount.Load() == 0 {
+			client.Close()
+			delete(r.clients, hostname)
+		}
+	}
+}
diff --git a/internal/ackhandler/packet.go b/internal/ackhandler/packet.go
index 394ee40a..849eea69 100644
--- a/internal/ackhandler/packet.go
+++ b/internal/ackhandler/packet.go
@@ -9,12 +9,12 @@ import (
 
 // A Packet is a packet
 type Packet struct {
+	SendTime        time.Time
 	PacketNumber    protocol.PacketNumber
 	Frames          []*Frame
 	LargestAcked    protocol.PacketNumber // InvalidPacketNumber if the packet doesn't contain an ACK
 	Length          protocol.ByteCount
 	EncryptionLevel protocol.EncryptionLevel
-	SendTime        time.Time
 
 	IsPathMTUProbePacket bool // We don't report the loss of Path MTU probe packets to the congestion controller.
 
diff --git a/internal/utils/rtt_stats.go b/internal/utils/rtt_stats.go
index 527539e1..2cd9a191 100644
--- a/internal/utils/rtt_stats.go
+++ b/internal/utils/rtt_stats.go
@@ -103,8 +103,12 @@ func (r *RTTStats) SetMaxAckDelay(mad time.Duration) {
 // SetInitialRTT sets the initial RTT.
 // It is used during the 0-RTT handshake when restoring the RTT stats from the session state.
 func (r *RTTStats) SetInitialRTT(t time.Duration) {
+	// On the server side, by the time we get to process the session ticket,
+	// we might already have obtained an RTT measurement.
+	// This can happen if we received the ClientHello in multiple pieces, and one of those pieces was lost.
+	// Discard the restored value. A fresh measurement is always better.
 	if r.hasMeasurement {
-		panic("initial RTT set after first measurement")
+		return
 	}
 	r.smoothedRTT = t
 	r.latestRTT = t
diff --git a/internal/wire/ack_frame.go b/internal/wire/ack_frame.go
index f145c8b4..2f816c78 100644
--- a/internal/wire/ack_frame.go
+++ b/internal/wire/ack_frame.go
@@ -90,13 +90,22 @@ func parseAckFrame(r *bytes.Reader, typ uint64, ackDelayExponent uint8, _ protoc
 		return nil, errInvalidAckRanges
 	}
 
-	// parse (and skip) the ECN section
 	if ecn {
-		for i := 0; i < 3; i++ {
-			if _, err := quicvarint.Read(r); err != nil {
-				return nil, err
-			}
+		ect0, err := quicvarint.Read(r)
+		if err != nil {
+			return nil, err
+		}
+		frame.ECT0 = ect0
+		ect1, err := quicvarint.Read(r)
+		if err != nil {
+			return nil, err
+		}
+		frame.ECT1 = ect1
+		ecnce, err := quicvarint.Read(r)
+		if err != nil {
+			return nil, err
 		}
+		frame.ECNCE = ecnce
 	}
 
 	return frame, nil
diff --git a/internal/wire/header.go b/internal/wire/header.go
index 6e8d4f9f..37f48cce 100644
--- a/internal/wire/header.go
+++ b/internal/wire/header.go
@@ -13,8 +13,6 @@ import (
 )
 
 // ParseConnectionID parses the destination connection ID of a packet.
-// It uses the data slice for the connection ID.
-// That means that the connection ID must not be used after the packet buffer is released.
 func ParseConnectionID(data []byte, shortHeaderConnIDLen int) (protocol.ConnectionID, error) {
 	if len(data) == 0 {
 		return protocol.ConnectionID{}, io.EOF
diff --git a/transport.go b/transport.go
index 55bc7c48..c62ee388 100644
--- a/transport.go
+++ b/transport.go
@@ -66,7 +66,8 @@ type Transport struct {
 
 	conn rawConn
 
-	closeQueue chan closePacket
+	closeQueue          chan closePacket
+	statelessResetQueue chan *receivedPacket
 
 	listening   chan struct{} // is closed when listen returns
 	closed      bool
@@ -182,6 +183,7 @@ func (t *Transport) init(isServer bool) error {
 		t.listening = make(chan struct{})
 
 		t.closeQueue = make(chan closePacket, 4)
+		t.statelessResetQueue = make(chan *receivedPacket, 4)
 
 		if t.ConnectionIDGenerator != nil {
 			t.connIDGenerator = t.ConnectionIDGenerator
@@ -196,7 +198,7 @@ func (t *Transport) init(isServer bool) error {
 		}
 
 		go t.listen(conn)
-		go t.runCloseQueue()
+		go t.runSendQueue()
 	})
 	return t.initErr
 }
@@ -210,13 +212,15 @@ func (t *Transport) enqueueClosePacket(p closePacket) {
 	}
 }
 
-func (t *Transport) runCloseQueue() {
+func (t *Transport) runSendQueue() {
 	for {
 		select {
 		case <-t.listening:
 			return
 		case p := <-t.closeQueue:
 			t.conn.WritePacket(p.payload, p.addr, p.info.OOB())
+		case p := <-t.statelessResetQueue:
+			t.sendStatelessReset(p)
 		}
 	}
 }
@@ -340,7 +344,7 @@ func (t *Transport) handlePacket(p *receivedPacket) {
 		return
 	}
 	if !wire.IsLongHeaderPacket(p.data[0]) {
-		go t.maybeSendStatelessReset(p, connID)
+		t.maybeSendStatelessReset(p)
 		return
 	}
 
@@ -353,14 +357,33 @@ func (t *Transport) handlePacket(p *receivedPacket) {
 	t.server.handlePacket(p)
 }
 
-func (t *Transport) maybeSendStatelessReset(p *receivedPacket, connID protocol.ConnectionID) {
-	defer p.buffer.Release()
+func (t *Transport) maybeSendStatelessReset(p *receivedPacket) {
 	if t.StatelessResetKey == nil {
+		p.buffer.Release()
 		return
 	}
+
 	// Don't send a stateless reset in response to very small packets.
 	// This includes packets that could be stateless resets.
 	if len(p.data) <= protocol.MinStatelessResetSize {
+		p.buffer.Release()
+		return
+	}
+
+	select {
+	case t.statelessResetQueue <- p:
+	default:
+		// it's fine to not send a stateless reset when we're busy
+		p.buffer.Release()
+	}
+}
+
+func (t *Transport) sendStatelessReset(p *receivedPacket) {
+	defer p.buffer.Release()
+
+	connID, err := wire.ParseConnectionID(p.data, t.connIDLen)
+	if err != nil {
+		t.logger.Errorf("error parsing connection ID on packet from %s: %s", p.remoteAddr, err)
 		return
 	}
 	token := t.handlerMap.GetStatelessResetToken(connID)
@@ -370,7 +393,7 @@ func (t *Transport) maybeSendStatelessReset(p *receivedPacket, connID protocol.C
 	data[0] = (data[0] & 0x7f) | 0x40
 	data = append(data, token[:]...)
 	if _, err := t.conn.WritePacket(data, p.remoteAddr, p.info.OOB()); err != nil {
-		t.logger.Debugf("Error sending Stateless Reset: %s", err)
+		t.logger.Debugf("Error sending Stateless Reset to %s: %s", p.remoteAddr, err)
 	}
 }
 


Test patch:
diff --git a/http3/roundtrip_test.go b/http3/roundtrip_test.go
index 14388b5f..8e859de5 100644
--- a/http3/roundtrip_test.go
+++ b/http3/roundtrip_test.go
@@ -8,6 +8,7 @@ import (
 	"io"
 	"net"
 	"net/http"
+	"sync/atomic"
 	"time"
 
 	"github.com/quic-go/quic-go"
@@ -304,10 +305,10 @@ var _ = Describe("RoundTripper", func() {
 
 	Context("closing", func() {
 		It("closes", func() {
-			rt.clients = make(map[string]roundTripCloser)
+			rt.clients = make(map[string]*roundTripCloserWithCount)
 			cl := NewMockRoundTripCloser(mockCtrl)
 			cl.EXPECT().Close()
-			rt.clients["foo.bar"] = cl
+			rt.clients["foo.bar"] = &roundTripCloserWithCount{cl, atomic.Int64{}}
 			err := rt.Close()
 			Expect(err).ToNot(HaveOccurred())
 			Expect(len(rt.clients)).To(BeZero())
@@ -319,6 +320,53 @@ var _ = Describe("RoundTripper", func() {
 			Expect(err).ToNot(HaveOccurred())
 			Expect(len(rt.clients)).To(BeZero())
 		})
+
+		It("closes idle connections", func() {
+			Expect(len(rt.clients)).To(Equal(0))
+			req1, err := http.NewRequest("GET", "https://site1.com", nil)
+			Expect(err).ToNot(HaveOccurred())
+			req2, err := http.NewRequest("GET", "https://site2.com", nil)
+			Expect(err).ToNot(HaveOccurred())
+			Expect(req1.Host).ToNot(Equal(req2.Host))
+			ctx1, cancel1 := context.WithCancel(context.Background())
+			ctx2, cancel2 := context.WithCancel(context.Background())
+			req1 = req1.WithContext(ctx1)
+			req2 = req2.WithContext(ctx2)
+			roundTripCalled := make(chan struct{})
+			reqFinished := make(chan struct{})
+			rt.newClient = func(hostname string, tlsConf *tls.Config, opts *roundTripperOpts, conf *quic.Config, dialer dialFunc) (roundTripCloser, error) {
+				cl := NewMockRoundTripCloser(mockCtrl)
+				cl.EXPECT().Close()
+				cl.EXPECT().RoundTripOpt(gomock.Any(), gomock.Any()).DoAndReturn(func(r *http.Request, _ RoundTripOpt) (*http.Response, error) {
+					roundTripCalled <- struct{}{}
+					<-r.Context().Done()
+					return nil, nil
+				})
+				return cl, nil
+			}
+			go func() {
+				rt.RoundTrip(req1)
+				reqFinished <- struct{}{}
+			}()
+			go func() {
+				rt.RoundTrip(req2)
+				reqFinished <- struct{}{}
+			}()
+			<-roundTripCalled
+			<-roundTripCalled
+			// Both two requests are started.
+			Expect(len(rt.clients)).To(Equal(2))
+			cancel1()
+			<-reqFinished
+			// req1 is finished
+			rt.CloseIdleConnections()
+			Expect(len(rt.clients)).To(Equal(1))
+			cancel2()
+			<-reqFinished
+			// all requests are finished
+			rt.CloseIdleConnections()
+			Expect(len(rt.clients)).To(Equal(0))
+		})
 	})
 
 	Context("reusing udpconn", func() {
diff --git a/internal/utils/rtt_stats_test.go b/internal/utils/rtt_stats_test.go
index 29892c6f..eae31926 100644
--- a/internal/utils/rtt_stats_test.go
+++ b/internal/utils/rtt_stats_test.go
@@ -154,4 +154,14 @@ var _ = Describe("RTT stats", func() {
 		Expect(rttStats.SmoothedRTT()).To(Equal(200 * time.Millisecond))
 		Expect(rttStats.MeanDeviation()).To(Equal(100 * time.Millisecond))
 	})
+
+	It("doesn't restore the RTT if we already have a measurement", func() {
+		const rtt = 10 * time.Millisecond
+		rttStats.UpdateRTT(rtt, 0, time.Now())
+		Expect(rttStats.LatestRTT()).To(Equal(rtt))
+		Expect(rttStats.SmoothedRTT()).To(Equal(rtt))
+		rttStats.SetInitialRTT(time.Minute)
+		Expect(rttStats.LatestRTT()).To(Equal(rtt))
+		Expect(rttStats.SmoothedRTT()).To(Equal(rtt))
+	})
 })
diff --git a/internal/wire/ack_frame_test.go b/internal/wire/ack_frame_test.go
index 49d11374..f80335c0 100644
--- a/internal/wire/ack_frame_test.go
+++ b/internal/wire/ack_frame_test.go
@@ -170,6 +170,9 @@ var _ = Describe("ACK Frame (for IETF QUIC)", func() {
 				Expect(frame.LargestAcked()).To(Equal(protocol.PacketNumber(100)))
 				Expect(frame.LowestAcked()).To(Equal(protocol.PacketNumber(90)))
 				Expect(frame.HasMissingRanges()).To(BeFalse())
+				Expect(frame.ECT0).To(BeEquivalentTo(0x42))
+				Expect(frame.ECT1).To(BeEquivalentTo(0x12345))
+				Expect(frame.ECNCE).To(BeEquivalentTo(0x12345678))
 				Expect(b.Len()).To(BeZero())
 			})