From a4d61a4be4201637a78d88312df5a17f31b823c8 Mon Sep 17 00:00:00 2001
From: Simone Basso <bassosimone@gmail.com>
Date: Fri, 25 Jun 2021 17:58:42 +0200
Subject: [PATCH] fix(netxlite): close quic packetconn (#407)

Noticed when working on https://github.com/ooni/probe/issues/1505.

Justification for this diff:

1. [DialEarlyContext calls dialContext with the last argument set to false](https://github.com/lucas-clemente/quic-go/blob/v0.21.1/client.go#L153);

2. [the semantics of the last argument is whether we own the connection](https://github.com/lucas-clemente/quic-go/blob/v0.21.1/client.go#L187);

3. [this value is propagated to the client data structure](https://github.com/lucas-clemente/quic-go/blob/v0.21.1/client.go#L269);

4. [client.dial](https://github.com/lucas-clemente/quic-go/blob/v0.21.1/client.go#L302) runs the session in a background goroutine and only destroys the `packetHandlers` when the connection is owned;

5. [packetHandlerMap.Destroy](https://github.com/lucas-clemente/quic-go/blob/v0.21.1/packet_handler_map.go#L293) closes the underlying PacketConn.

6. also, the documentation clearly states that when you use `DialEarlyContext` you can use the same packet conn multiple times, so it does not take ownership.
---
 internal/engine/netx/quicdialer/system.go |  2 +-
 internal/netxlite/quic.go                 | 25 +++++++++++++++++++++--
 internal/netxlite/quic_test.go            | 21 +++++++++++++++++++
 3 files changed, 45 insertions(+), 3 deletions(-)

diff --git a/internal/engine/netx/quicdialer/system.go b/internal/engine/netx/quicdialer/system.go
index 392378f..3507460 100644
--- a/internal/engine/netx/quicdialer/system.go
+++ b/internal/engine/netx/quicdialer/system.go
@@ -11,7 +11,7 @@ import (
 
 // QUICListener listens for QUIC connections.
 type QUICListener interface {
-	// Listen creates a new listening net.PacketConn.
+	// Listen creates a new listening PacketConn.
 	Listen(addr *net.UDPAddr) (net.PacketConn, error)
 }
 
diff --git a/internal/netxlite/quic.go b/internal/netxlite/quic.go
index 5ae1749..d02e5ef 100644
--- a/internal/netxlite/quic.go
+++ b/internal/netxlite/quic.go
@@ -30,7 +30,7 @@ type QUICDialer interface {
 
 // QUICListener listens for QUIC connections.
 type QUICListener interface {
-	// Listen creates a new listening net.PacketConn.
+	// Listen creates a new listening PacketConn.
 	Listen(addr *net.UDPAddr) (net.PacketConn, error)
 }
 
@@ -76,6 +76,27 @@ func (d *QUICDialerQUICGo) DialContext(ctx context.Context, network string,
 		return nil, err
 	}
 	udpAddr := &net.UDPAddr{IP: ip, Port: port, Zone: ""}
-	return quic.DialEarlyContext(
+	sess, err := quic.DialEarlyContext(
 		ctx, pconn, udpAddr, address, tlsConfig, quicConfig)
+	if err != nil {
+		return nil, err
+	}
+	return &quicSessionOwnsConn{EarlySession: sess, conn: pconn}, nil
+}
+
+// quicSessionOwnsConn ensures that we close the PacketConn.
+type quicSessionOwnsConn struct {
+	// EarlySession is the embedded early session
+	quic.EarlySession
+
+	// conn is the connection we own
+	conn net.PacketConn
+}
+
+// CloseWithError implements quic.EarlySession.CloseWithError.
+func (sess *quicSessionOwnsConn) CloseWithError(
+	code quic.ApplicationErrorCode, reason string) error {
+	err := sess.EarlySession.CloseWithError(code, reason)
+	sess.conn.Close()
+	return err
 }
diff --git a/internal/netxlite/quic_test.go b/internal/netxlite/quic_test.go
index 4d6e2b2..873a279 100644
--- a/internal/netxlite/quic_test.go
+++ b/internal/netxlite/quic_test.go
@@ -94,6 +94,26 @@ func TestQUICDialerQUICGoCannotListen(t *testing.T) {
 	}
 }
 
+func TestQUICDialerCannotPerformHandshake(t *testing.T) {
+	tlsConfig := &tls.Config{
+		NextProtos: []string{"h3"},
+		ServerName: "dns.google",
+	}
+	systemdialer := QUICDialerQUICGo{
+		QUICListener: &QUICListenerStdlib{},
+	}
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel() // fail immediately
+	sess, err := systemdialer.DialContext(
+		ctx, "udp", "8.8.8.8:443", tlsConfig, &quic.Config{})
+	if !errors.Is(err, context.Canceled) {
+		t.Fatal("not the error we expected", err)
+	}
+	if sess != nil {
+		log.Fatal("expected nil session here")
+	}
+}
+
 func TestQUICDialerWorksAsIntended(t *testing.T) {
 	tlsConfig := &tls.Config{
 		NextProtos: []string{"h3"},
@@ -108,6 +128,7 @@ func TestQUICDialerWorksAsIntended(t *testing.T) {
 	if err != nil {
 		t.Fatal("not the error we expected", err)
 	}
+	<-sess.HandshakeComplete().Done()
 	if err := sess.CloseWithError(0, ""); err != nil {
 		log.Fatal(err)
 	}