correctly handle fresh contexts returned from ConnContext

integrationtests/self/handshake_test.go

@@ -213,6 +213,37 @@ var _ = Describe("Handshake tests", func() {
  checkContext(tlsGetCertificateContextChan, false)
  })
 
+ It("correctly handles a fresh context returned from ConnContext", func() {
+ conn, err := net.ListenUDP("udp", &net.UDPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 0})
+ Expect(err).ToNot(HaveOccurred())
+ defer conn.Close()
+ tr := &quic.Transport{
+ Conn: conn,
+ ConnContext: func(ctx context.Context) context.Context { return context.Background() },
+ }
+ server, err := tr.Listen(getTLSConfig(), getQuicConfig(nil))
+ Expect(err).ToNot(HaveOccurred())
+ done := make(chan struct{})
+ go func() {
+ defer GinkgoRecover()
+ defer close(done)
+ conn, err := server.Accept(context.Background())
+ if err != nil {
+ return
+ }
+ Eventually(conn.Context().Done).Should(BeClosed())
+ }()
+
+ c, err := quic.DialAddr(
+ context.Background(),
+ fmt.Sprintf("localhost:%d", server.Addr().(*net.UDPAddr).Port),
+ getTLSClientConfig(),
+ getQuicConfig(nil),
+ )
+ Expect(err).ToNot(HaveOccurred())
+ c.CloseWithError(1337, "bye")
+ })
+
  It("uses the context everywhere, on the client side", func() {
  tlsServerConf := getTLSConfig()
  tlsServerConf.ClientAuth = tls.RequestClientCert

server.go

@@ -636,12 +636,24 @@ func (s *baseServer) handleInitialImpl(p receivedPacket, hdr *wire.Header) error
  }
 
  var conn quicConn
- ctx, cancel := context.WithCancelCause(context.Background())
+ var cancel context.CancelCauseFunc
+ ctx, cancel1 := context.WithCancelCause(context.Background())
  if s.connContext != nil {
  ctx = s.connContext(ctx)
  if ctx == nil {
  panic("quic: ConnContext returned nil")
  }
+ // There's no guarantee that the application returns a context
+ // that's derived from the context we passed into ConnContext.
+ // We need to make sure that both contexts are cancelled.
+ var cancel2 context.CancelCauseFunc
+ ctx, cancel2 = context.WithCancelCause(ctx)
+ cancel = func(cause error) {
+ cancel1(cause)
+ cancel2(cause)
+ }
+ } else {
+ cancel = cancel1
  }
  ctx = context.WithValue(ctx, ConnectionTracingKey, nextConnTracingID())
  var tracer *logging.ConnectionTracer