Repository: quic-go/quic-go
Pull request #4401 (2024-04-01 22:58:27)

http3: pass tracing ID instead of quic.Connection to stream hijackers
~Depends on #4400.~

The stream hijackers only need to be able to associate the stream with the underlying QUIC connection. They are not supposed to call any functions on the `quic.Connection`. As such, the better API is to just pass them a unique identifier.

This fixes a problem in webtransport-go introduced by #4389, where we're using a HTTP/3 connection, which, although satisfying the `quic.Connection` interface, led to a type confusing.

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/quic-go/quic-go/pull/4401?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=quic-go) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Project coverage is 85.32%. Comparing base [(`32f32df`)](https://app.codecov.io/gh/quic-go/quic-go/commit/32f32df0cb7d418ab1b675a77c12bb64fa473369?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=quic-go) to head [(`0d62a15`)](https://app.codecov.io/gh/quic-go/quic-go/pull/4401?dropdown=coverage&src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=quic-go).


<details><summary>Additional details and impacted files</summary>


```diff
@@                    Coverage Diff                    @@
##           connection-tracing-id    #4401      +/-   ##
=========================================================
+ Coverage                  84.90%   85.32%   +0.42%     
=========================================================
  Files                        152      149       -3     
  Lines                      14342    12816    -1526     
=========================================================
- Hits                       12177    10935    -1242     
+ Misses                      1666     1384     -282     
+ Partials                     499      497       -2     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/quic-go/quic-go/pull/4401?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=quic-go).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=quic-go).


Patch:
diff --git a/http3/client.go b/http3/client.go
index 1008eeea..1a0a891e 100644
--- a/http3/client.go
+++ b/http3/client.go
@@ -49,8 +49,8 @@ type roundTripperOpts struct {
 	EnableDatagram     bool
 	MaxHeaderBytes     int64
 	AdditionalSettings map[uint64]uint64
-	StreamHijacker     func(FrameType, quic.Connection, quic.Stream, error) (hijacked bool, err error)
-	UniStreamHijacker  func(StreamType, quic.Connection, quic.ReceiveStream, error) (hijacked bool)
+	StreamHijacker     func(FrameType, quic.ConnectionTracingID, quic.Stream, error) (hijacked bool, err error)
+	UniStreamHijacker  func(StreamType, quic.ConnectionTracingID, quic.ReceiveStream, error) (hijacked bool)
 }
 
 // client is a HTTP3 client doing requests
@@ -183,7 +183,8 @@ func (c *client) handleBidirectionalStreams(conn quic.EarlyConnection) {
 		}
 		go func(str quic.Stream) {
 			_, err := parseNextFrame(str, func(ft FrameType, e error) (processed bool, err error) {
-				return c.opts.StreamHijacker(ft, conn, str, e)
+				id := conn.Context().Value(quic.ConnectionTracingKey).(quic.ConnectionTracingID)
+				return c.opts.StreamHijacker(ft, id, str, e)
 			})
 			if err == errHijacked {
 				return
diff --git a/http3/conn.go b/http3/conn.go
index cae1f2ce..e154f17c 100644
--- a/http3/conn.go
+++ b/http3/conn.go
@@ -17,7 +17,7 @@ type connection struct {
 	logger      utils.Logger
 
 	enableDatagrams   bool
-	uniStreamHijacker func(StreamType, quic.Connection, quic.ReceiveStream, error) (hijacked bool)
+	uniStreamHijacker func(StreamType, quic.ConnectionTracingID, quic.ReceiveStream, error) (hijacked bool)
 
 	settings         *Settings
 	receivedSettings chan struct{}
@@ -26,7 +26,7 @@ type connection struct {
 func newConnection(
 	quicConn quic.Connection,
 	enableDatagrams bool,
-	uniStreamHijacker func(StreamType, quic.Connection, quic.ReceiveStream, error) (hijacked bool),
+	uniStreamHijacker func(StreamType, quic.ConnectionTracingID, quic.ReceiveStream, error) (hijacked bool),
 	perspective protocol.Perspective,
 	logger utils.Logger,
 ) *connection {
@@ -57,7 +57,8 @@ func (c *connection) HandleUnidirectionalStreams() {
 		go func(str quic.ReceiveStream) {
 			streamType, err := quicvarint.Read(quicvarint.NewReader(str))
 			if err != nil {
-				if c.uniStreamHijacker != nil && c.uniStreamHijacker(StreamType(streamType), c.Connection, str, err) {
+				id := c.Connection.Context().Value(quic.ConnectionTracingKey).(quic.ConnectionTracingID)
+				if c.uniStreamHijacker != nil && c.uniStreamHijacker(StreamType(streamType), id, str, err) {
 					return
 				}
 				c.logger.Debugf("reading stream type on stream %d failed: %s", str.StreamID(), err)
@@ -89,8 +90,15 @@ func (c *connection) HandleUnidirectionalStreams() {
 				}
 				return
 			default:
-				if c.uniStreamHijacker != nil && c.uniStreamHijacker(StreamType(streamType), c.Connection, str, nil) {
-					return
+				if c.uniStreamHijacker != nil {
+					if c.uniStreamHijacker(
+						StreamType(streamType),
+						c.Connection.Context().Value(quic.ConnectionTracingKey).(quic.ConnectionTracingID),
+						str,
+						nil,
+					) {
+						return
+					}
 				}
 				str.CancelRead(quic.StreamErrorCode(ErrCodeStreamCreationError))
 				return
diff --git a/http3/roundtrip.go b/http3/roundtrip.go
index 8dedaa21..ecbb6974 100644
--- a/http3/roundtrip.go
+++ b/http3/roundtrip.go
@@ -89,12 +89,12 @@ type RoundTripper struct {
 	// Callers can either ignore the frame and return control of the stream back to HTTP/3
 	// (by returning hijacked false).
 	// Alternatively, callers can take over the QUIC stream (by returning hijacked true).
-	StreamHijacker func(FrameType, quic.Connection, quic.Stream, error) (hijacked bool, err error)
+	StreamHijacker func(FrameType, quic.ConnectionTracingID, quic.Stream, error) (hijacked bool, err error)
 
 	// When set, this callback is called for unknown unidirectional stream of unknown stream type.
 	// If parsing the stream type fails, the error is passed to the callback.
 	// In that case, the stream type will not be set.
-	UniStreamHijacker func(StreamType, quic.Connection, quic.ReceiveStream, error) (hijacked bool)
+	UniStreamHijacker func(StreamType, quic.ConnectionTracingID, quic.ReceiveStream, error) (hijacked bool)
 
 	// Dial specifies an optional dial function for creating QUIC
 	// connections for requests.
diff --git a/http3/server.go b/http3/server.go
index 538d2637..b1449e81 100644
--- a/http3/server.go
+++ b/http3/server.go
@@ -202,12 +202,12 @@ type Server struct {
 	// Callers can either ignore the frame and return control of the stream back to HTTP/3
 	// (by returning hijacked false).
 	// Alternatively, callers can take over the QUIC stream (by returning hijacked true).
-	StreamHijacker func(FrameType, quic.Connection, quic.Stream, error) (hijacked bool, err error)
+	StreamHijacker func(FrameType, quic.ConnectionTracingID, quic.Stream, error) (hijacked bool, err error)
 
 	// UniStreamHijacker, when set, is called for unknown unidirectional stream of unknown stream type.
 	// If parsing the stream type fails, the error is passed to the callback.
 	// In that case, the stream type will not be set.
-	UniStreamHijacker func(StreamType, quic.Connection, quic.ReceiveStream, error) (hijacked bool)
+	UniStreamHijacker func(StreamType, quic.ConnectionTracingID, quic.ReceiveStream, error) (hijacked bool)
 
 	// ConnContext optionally specifies a function that modifies
 	// the context used for a new connection c. The provided ctx
@@ -512,7 +512,14 @@ func (s *Server) maxHeaderBytes() uint64 {
 func (s *Server) handleRequest(conn *connection, str quic.Stream, decoder *qpack.Decoder, onFrameError func()) requestError {
 	var ufh unknownFrameHandlerFunc
 	if s.StreamHijacker != nil {
-		ufh = func(ft FrameType, e error) (processed bool, err error) { return s.StreamHijacker(ft, conn, str, e) }
+		ufh = func(ft FrameType, e error) (processed bool, err error) {
+			return s.StreamHijacker(
+				ft,
+				conn.Context().Value(quic.ConnectionTracingKey).(quic.ConnectionTracingID),
+				str,
+				e,
+			)
+		}
 	}
 	frame, err := parseNextFrame(str, ufh)
 	if err != nil {


Test patch:
diff --git a/http3/client_test.go b/http3/client_test.go
index 1155701d..2a0ef451 100644
--- a/http3/client_test.go
+++ b/http3/client_test.go
@@ -237,9 +237,11 @@ var _ = Describe("Client", func() {
 		})
 
 		It("hijacks a bidirectional stream of unknown frame type", func() {
+			id := quic.ConnectionTracingID(1234)
 			frameTypeChan := make(chan FrameType, 1)
-			cl.opts.StreamHijacker = func(ft FrameType, c quic.Connection, s quic.Stream, e error) (hijacked bool, err error) {
+			cl.opts.StreamHijacker = func(ft FrameType, connTracingID quic.ConnectionTracingID, _ quic.Stream, e error) (hijacked bool, err error) {
 				Expect(e).ToNot(HaveOccurred())
+				Expect(connTracingID).To(Equal(id))
 				frameTypeChan <- ft
 				return true, nil
 			}
@@ -252,6 +254,8 @@ var _ = Describe("Client", func() {
 				<-testDone
 				return nil, errors.New("test done")
 			})
+			ctx := context.WithValue(context.Background(), quic.ConnectionTracingKey, id)
+			conn.EXPECT().Context().Return(ctx).AnyTimes()
 			_, err := cl.RoundTripOpt(request, RoundTripOpt{})
 			Expect(err).To(MatchError("done"))
 			Eventually(frameTypeChan).Should(Receive(BeEquivalentTo(0x41)))
@@ -260,7 +264,7 @@ var _ = Describe("Client", func() {
 
 		It("closes the connection when hijacker didn't hijack a bidirectional stream", func() {
 			frameTypeChan := make(chan FrameType, 1)
-			cl.opts.StreamHijacker = func(ft FrameType, c quic.Connection, s quic.Stream, e error) (hijacked bool, err error) {
+			cl.opts.StreamHijacker = func(ft FrameType, _ quic.ConnectionTracingID, _ quic.Stream, e error) (hijacked bool, err error) {
 				Expect(e).ToNot(HaveOccurred())
 				frameTypeChan <- ft
 				return false, nil
@@ -274,6 +278,8 @@ var _ = Describe("Client", func() {
 				<-testDone
 				return nil, errors.New("test done")
 			})
+			ctx := context.WithValue(context.Background(), quic.ConnectionTracingKey, quic.ConnectionTracingID(1234))
+			conn.EXPECT().Context().Return(ctx).AnyTimes()
 			conn.EXPECT().CloseWithError(quic.ApplicationErrorCode(ErrCodeFrameUnexpected), gomock.Any()).Return(nil).AnyTimes()
 			_, err := cl.RoundTripOpt(request, RoundTripOpt{})
 			Expect(err).To(MatchError("done"))
@@ -282,7 +288,7 @@ var _ = Describe("Client", func() {
 
 		It("closes the connection when hijacker returned error", func() {
 			frameTypeChan := make(chan FrameType, 1)
-			cl.opts.StreamHijacker = func(ft FrameType, c quic.Connection, s quic.Stream, e error) (hijacked bool, err error) {
+			cl.opts.StreamHijacker = func(ft FrameType, _ quic.ConnectionTracingID, _ quic.Stream, e error) (hijacked bool, err error) {
 				Expect(e).ToNot(HaveOccurred())
 				frameTypeChan <- ft
 				return false, errors.New("error in hijacker")
@@ -296,6 +302,8 @@ var _ = Describe("Client", func() {
 				<-testDone
 				return nil, errors.New("test done")
 			})
+			ctx := context.WithValue(context.Background(), quic.ConnectionTracingKey, quic.ConnectionTracingID(1234))
+			conn.EXPECT().Context().Return(ctx).AnyTimes()
 			conn.EXPECT().CloseWithError(quic.ApplicationErrorCode(ErrCodeFrameUnexpected), gomock.Any()).Return(nil).AnyTimes()
 			_, err := cl.RoundTripOpt(request, RoundTripOpt{})
 			Expect(err).To(MatchError("done"))
@@ -306,7 +314,7 @@ var _ = Describe("Client", func() {
 			testErr := errors.New("test error")
 			unknownStr := mockquic.NewMockStream(mockCtrl)
 			done := make(chan struct{})
-			cl.opts.StreamHijacker = func(ft FrameType, c quic.Connection, str quic.Stream, e error) (hijacked bool, err error) {
+			cl.opts.StreamHijacker = func(ft FrameType, _ quic.ConnectionTracingID, str quic.Stream, e error) (hijacked bool, err error) {
 				defer close(done)
 				Expect(e).To(MatchError(testErr))
 				Expect(ft).To(BeZero())
@@ -320,6 +328,8 @@ var _ = Describe("Client", func() {
 				<-testDone
 				return nil, errors.New("test done")
 			})
+			ctx := context.WithValue(context.Background(), quic.ConnectionTracingKey, quic.ConnectionTracingID(1234))
+			conn.EXPECT().Context().Return(ctx).AnyTimes()
 			conn.EXPECT().CloseWithError(quic.ApplicationErrorCode(ErrCodeFrameUnexpected), gomock.Any()).Return(nil).AnyTimes()
 			_, err := cl.RoundTripOpt(request, RoundTripOpt{})
 			Expect(err).To(MatchError("done"))
@@ -363,8 +373,10 @@ var _ = Describe("Client", func() {
 		})
 
 		It("hijacks an unidirectional stream of unknown stream type", func() {
+			id := quic.ConnectionTracingID(100)
 			streamTypeChan := make(chan StreamType, 1)
-			cl.opts.UniStreamHijacker = func(st StreamType, _ quic.Connection, _ quic.ReceiveStream, err error) bool {
+			cl.opts.UniStreamHijacker = func(st StreamType, connTracingID quic.ConnectionTracingID, _ quic.ReceiveStream, err error) bool {
+				Expect(connTracingID).To(Equal(id))
 				Expect(err).ToNot(HaveOccurred())
 				streamTypeChan <- st
 				return true
@@ -380,6 +392,8 @@ var _ = Describe("Client", func() {
 				<-testDone
 				return nil, errors.New("test done")
 			})
+			ctx := context.WithValue(context.Background(), quic.ConnectionTracingKey, id)
+			conn.EXPECT().Context().Return(ctx).AnyTimes()
 			_, err := cl.RoundTripOpt(req, RoundTripOpt{})
 			Expect(err).To(MatchError("done"))
 			Eventually(streamTypeChan).Should(Receive(BeEquivalentTo(0x54)))
@@ -390,7 +404,7 @@ var _ = Describe("Client", func() {
 			testErr := errors.New("test error")
 			done := make(chan struct{})
 			unknownStr := mockquic.NewMockStream(mockCtrl)
-			cl.opts.UniStreamHijacker = func(st StreamType, _ quic.Connection, str quic.ReceiveStream, err error) bool {
+			cl.opts.UniStreamHijacker = func(st StreamType, _ quic.ConnectionTracingID, str quic.ReceiveStream, err error) bool {
 				defer close(done)
 				Expect(st).To(BeZero())
 				Expect(str).To(Equal(unknownStr))
@@ -404,6 +418,8 @@ var _ = Describe("Client", func() {
 				<-testDone
 				return nil, errors.New("test done")
 			})
+			ctx := context.WithValue(context.Background(), quic.ConnectionTracingKey, quic.ConnectionTracingID(1234))
+			conn.EXPECT().Context().Return(ctx).AnyTimes()
 			_, err := cl.RoundTripOpt(req, RoundTripOpt{})
 			Expect(err).To(MatchError("done"))
 			Eventually(done).Should(BeClosed())
@@ -412,7 +428,7 @@ var _ = Describe("Client", func() {
 
 		It("cancels reading when hijacker didn't hijack an unidirectional stream", func() {
 			streamTypeChan := make(chan StreamType, 1)
-			cl.opts.UniStreamHijacker = func(st StreamType, _ quic.Connection, _ quic.ReceiveStream, err error) bool {
+			cl.opts.UniStreamHijacker = func(st StreamType, _ quic.ConnectionTracingID, _ quic.ReceiveStream, err error) bool {
 				Expect(err).ToNot(HaveOccurred())
 				streamTypeChan <- st
 				return false
@@ -429,6 +445,8 @@ var _ = Describe("Client", func() {
 				<-testDone
 				return nil, errors.New("test done")
 			})
+			ctx := context.WithValue(context.Background(), quic.ConnectionTracingKey, quic.ConnectionTracingID(1234))
+			conn.EXPECT().Context().Return(ctx).AnyTimes()
 			_, err := cl.RoundTripOpt(req, RoundTripOpt{})
 			Expect(err).To(MatchError("done"))
 			Eventually(streamTypeChan).Should(Receive(BeEquivalentTo(0x54)))
diff --git a/http3/server_test.go b/http3/server_test.go
index 789770f1..10f4f2b5 100644
--- a/http3/server_test.go
+++ b/http3/server_test.go
@@ -315,9 +315,12 @@ var _ = Describe("Server", func() {
 			AfterEach(func() { testDone <- struct{}{} })
 
 			It("hijacks a bidirectional stream of unknown frame type", func() {
+				id := quic.ConnectionTracingID(1337)
 				frameTypeChan := make(chan FrameType, 1)
-				s.StreamHijacker = func(ft FrameType, c quic.Connection, s quic.Stream, e error) (hijacked bool, err error) {
+				s.StreamHijacker = func(ft FrameType, connTracingID quic.ConnectionTracingID, _ quic.Stream, e error) (hijacked bool, err error) {
+					defer GinkgoRecover()
 					Expect(e).ToNot(HaveOccurred())
+					Expect(connTracingID).To(Equal(id))
 					frameTypeChan <- ft
 					return true, nil
 				}
@@ -331,6 +334,8 @@ var _ = Describe("Server", func() {
 					<-testDone
 					return nil, errors.New("test done")
 				})
+				ctx := context.WithValue(context.Background(), quic.ConnectionTracingKey, id)
+				conn.EXPECT().Context().Return(ctx).AnyTimes()
 				s.handleConn(conn)
 				Eventually(frameTypeChan).Should(Receive(BeEquivalentTo(0x41)))
 				time.Sleep(scaleDuration(20 * time.Millisecond)) // don't EXPECT any calls to conn.CloseWithError
@@ -338,7 +343,7 @@ var _ = Describe("Server", func() {
 
 			It("cancels writing when hijacker didn't hijack a bidirectional stream", func() {
 				frameTypeChan := make(chan FrameType, 1)
-				s.StreamHijacker = func(ft FrameType, c quic.Connection, s quic.Stream, e error) (hijacked bool, err error) {
+				s.StreamHijacker = func(ft FrameType, _ quic.ConnectionTracingID, _ quic.Stream, e error) (hijacked bool, err error) {
 					Expect(e).ToNot(HaveOccurred())
 					frameTypeChan <- ft
 					return false, nil
@@ -354,6 +359,8 @@ var _ = Describe("Server", func() {
 					<-testDone
 					return nil, errors.New("test done")
 				})
+				ctx := context.WithValue(context.Background(), quic.ConnectionTracingKey, quic.ConnectionTracingID(1234))
+				conn.EXPECT().Context().Return(ctx).AnyTimes()
 				s.handleConn(conn)
 				Eventually(frameTypeChan).Should(Receive(BeEquivalentTo(0x41)))
 				time.Sleep(scaleDuration(20 * time.Millisecond)) // don't EXPECT any calls to conn.CloseWithError
@@ -361,7 +368,7 @@ var _ = Describe("Server", func() {
 
 			It("cancels writing when hijacker returned error", func() {
 				frameTypeChan := make(chan FrameType, 1)
-				s.StreamHijacker = func(ft FrameType, c quic.Connection, s quic.Stream, e error) (hijacked bool, err error) {
+				s.StreamHijacker = func(ft FrameType, _ quic.ConnectionTracingID, _ quic.Stream, e error) (hijacked bool, err error) {
 					Expect(e).ToNot(HaveOccurred())
 					frameTypeChan <- ft
 					return false, errors.New("error in hijacker")
@@ -377,6 +384,8 @@ var _ = Describe("Server", func() {
 					<-testDone
 					return nil, errors.New("test done")
 				})
+				ctx := context.WithValue(context.Background(), quic.ConnectionTracingKey, quic.ConnectionTracingID(1234))
+				conn.EXPECT().Context().Return(ctx).AnyTimes()
 				s.handleConn(conn)
 				Eventually(frameTypeChan).Should(Receive(BeEquivalentTo(0x41)))
 				time.Sleep(scaleDuration(20 * time.Millisecond)) // don't EXPECT any calls to conn.CloseWithError
@@ -386,7 +395,7 @@ var _ = Describe("Server", func() {
 				testErr := errors.New("test error")
 				done := make(chan struct{})
 				unknownStr := mockquic.NewMockStream(mockCtrl)
-				s.StreamHijacker = func(ft FrameType, _ quic.Connection, str quic.Stream, err error) (bool, error) {
+				s.StreamHijacker = func(ft FrameType, _ quic.ConnectionTracingID, str quic.Stream, err error) (bool, error) {
 					defer close(done)
 					Expect(ft).To(BeZero())
 					Expect(str).To(Equal(unknownStr))
@@ -401,6 +410,8 @@ var _ = Describe("Server", func() {
 					<-testDone
 					return nil, errors.New("test done")
 				})
+				ctx := context.WithValue(context.Background(), quic.ConnectionTracingKey, quic.ConnectionTracingID(1234))
+				conn.EXPECT().Context().Return(ctx).AnyTimes()
 				s.handleConn(conn)
 				Eventually(done).Should(BeClosed())
 				time.Sleep(scaleDuration(20 * time.Millisecond)) // don't EXPECT any calls to conn.CloseWithError
@@ -425,9 +436,11 @@ var _ = Describe("Server", func() {
 			AfterEach(func() { testDone <- struct{}{} })
 
 			It("hijacks an unidirectional stream of unknown stream type", func() {
+				id := quic.ConnectionTracingID(42)
 				streamTypeChan := make(chan StreamType, 1)
-				s.UniStreamHijacker = func(st StreamType, _ quic.Connection, _ quic.ReceiveStream, err error) bool {
+				s.UniStreamHijacker = func(st StreamType, connTracingID quic.ConnectionTracingID, _ quic.ReceiveStream, err error) bool {
 					Expect(err).ToNot(HaveOccurred())
+					Expect(connTracingID).To(Equal(id))
 					streamTypeChan <- st
 					return true
 				}
@@ -442,6 +455,8 @@ var _ = Describe("Server", func() {
 					<-testDone
 					return nil, errors.New("test done")
 				})
+				ctx := context.WithValue(context.Background(), quic.ConnectionTracingKey, id)
+				conn.EXPECT().Context().Return(ctx).AnyTimes()
 				s.handleConn(conn)
 				Eventually(streamTypeChan).Should(Receive(BeEquivalentTo(0x54)))
 				time.Sleep(scaleDuration(20 * time.Millisecond)) // don't EXPECT any calls to conn.CloseWithError
@@ -451,7 +466,7 @@ var _ = Describe("Server", func() {
 				testErr := errors.New("test error")
 				done := make(chan struct{})
 				unknownStr := mockquic.NewMockStream(mockCtrl)
-				s.UniStreamHijacker = func(st StreamType, _ quic.Connection, str quic.ReceiveStream, err error) bool {
+				s.UniStreamHijacker = func(st StreamType, _ quic.ConnectionTracingID, str quic.ReceiveStream, err error) bool {
 					defer close(done)
 					Expect(st).To(BeZero())
 					Expect(str).To(Equal(unknownStr))
@@ -465,6 +480,8 @@ var _ = Describe("Server", func() {
 					<-testDone
 					return nil, errors.New("test done")
 				})
+				ctx := context.WithValue(context.Background(), quic.ConnectionTracingKey, quic.ConnectionTracingID(1234))
+				conn.EXPECT().Context().Return(ctx).AnyTimes()
 				s.handleConn(conn)
 				Eventually(done).Should(BeClosed())
 				time.Sleep(scaleDuration(20 * time.Millisecond)) // don't EXPECT any calls to conn.CloseWithError
@@ -472,7 +489,7 @@ var _ = Describe("Server", func() {
 
 			It("cancels reading when hijacker didn't hijack an unidirectional stream", func() {
 				streamTypeChan := make(chan StreamType, 1)
-				s.UniStreamHijacker = func(st StreamType, _ quic.Connection, _ quic.ReceiveStream, err error) bool {
+				s.UniStreamHijacker = func(st StreamType, _ quic.ConnectionTracingID, _ quic.ReceiveStream, err error) bool {
 					Expect(err).ToNot(HaveOccurred())
 					streamTypeChan <- st
 					return false
@@ -490,6 +507,8 @@ var _ = Describe("Server", func() {
 					<-testDone
 					return nil, errors.New("test done")
 				})
+				ctx := context.WithValue(context.Background(), quic.ConnectionTracingKey, quic.ConnectionTracingID(1234))
+				conn.EXPECT().Context().Return(ctx).AnyTimes()
 				s.handleConn(conn)
 				Eventually(streamTypeChan).Should(Receive(BeEquivalentTo(0x54)))
 				time.Sleep(scaleDuration(20 * time.Millisecond)) // don't EXPECT any calls to conn.CloseWithError