From c3772f61c7bff98ee70bc2e2086b0209f20f0f22 Mon Sep 17 00:00:00 2001
From: Marten Seemann <martenseemann@gmail.com>
Date: Mon, 24 Oct 2022 12:35:37 +0100
Subject: [PATCH] expose quic-go's connection tracing ID on the Session.Context
 (#35)

---
 session.go      | 3 ++-
 session_test.go | 7 ++++++-
 2 files changed, 8 insertions(+), 2 deletions(-)

diff --git a/session.go b/session.go
index 656f4e2..937fce5 100644
--- a/session.go
+++ b/session.go
@@ -74,7 +74,8 @@ type Session struct {
 }
 
 func newSession(sessionID sessionID, qconn http3.StreamCreator, requestStr quic.Stream) *Session {
-	ctx, ctxCancel := context.WithCancel(context.Background())
+	tracingID := qconn.Context().Value(quic.ConnectionTracingKey).(uint64)
+	ctx, ctxCancel := context.WithCancel(context.WithValue(context.Background(), quic.ConnectionTracingKey, tracingID))
 	c := &Session{
 		sessionID:       sessionID,
 		qconn:           qconn,
diff --git a/session_test.go b/session_test.go
index 7ceed8f..fbd2086 100644
--- a/session_test.go
+++ b/session_test.go
@@ -1,6 +1,7 @@
 package webtransport
 
 import (
+	"context"
 	"io"
 	"testing"
 
@@ -40,6 +41,7 @@ func TestCloseStreamsOnClose(t *testing.T) {
 	defer ctrl.Finish()
 
 	mockSess := NewMockStreamCreator(ctrl)
+	mockSess.EXPECT().Context().Return(context.WithValue(context.Background(), quic.ConnectionTracingKey, uint64(1337)))
 	sess := newSession(42, mockSess, newMockRequestStream(ctrl))
 
 	str := NewMockStream(ctrl)
@@ -63,7 +65,10 @@ func TestAddStreamAfterSessionClose(t *testing.T) {
 	ctrl := gomock.NewController(t)
 	defer ctrl.Finish()
 
-	sess := newSession(42, NewMockStreamCreator(ctrl), newMockRequestStream(ctrl))
+	mockSess := NewMockStreamCreator(ctrl)
+	mockSess.EXPECT().Context().Return(context.WithValue(context.Background(), quic.ConnectionTracingKey, uint64(1337)))
+
+	sess := newSession(42, mockSess, newMockRequestStream(ctrl))
 	require.NoError(t, sess.Close())
 
 	str := NewMockStream(ctrl)