cleanup(quic): wait for handshake completion in netxlite (#729)

See https://github.com/ooni/probe/issues/2097

internal/netxlite/quic.go

@@ -55,9 +55,12 @@ func NewQUICDialerWithResolver(listener model.QUICListener,
 		Dialer: &quicDialerResolver{
 			Dialer: &quicDialerLogger{
 				Dialer: &quicDialerErrWrapper{
-					QUICDialer: &quicDialerQUICGo{
-						QUICListener: listener,
-					}},
+					QUICDialer: &quicDialerHandshakeCompleter{
+						Dialer: &quicDialerQUICGo{
+							QUICListener: listener,
+						},
+					},
+				},
 				Logger:          logger,
 				operationSuffix: "_address",
 			},
@@ -164,6 +167,33 @@ func (d *quicDialerQUICGo) CloseIdleConnections() {
 	// nothing to do
 }
 
+// quicDialerHandshakeCompleter ensures we complete the handshake.
+type quicDialerHandshakeCompleter struct {
+	Dialer model.QUICDialer
+}
+
+// DialContext implements model.QUICDialer.DialContext.
+func (d *quicDialerHandshakeCompleter) DialContext(
+	ctx context.Context, network, address string,
+	tlsConfig *tls.Config, quicConfig *quic.Config) (quic.EarlyConnection, error) {
+	conn, err := d.Dialer.DialContext(ctx, network, address, tlsConfig, quicConfig)
+	if err != nil {
+		return nil, err
+	}
+	select {
+	case <-conn.HandshakeComplete().Done():
+		return conn, nil
+	case <-ctx.Done():
+		conn.CloseWithError(0, "") // we own the conn
+		return nil, ctx.Err()
+	}
+}
+
+// CloseIdleConnections implements model.QUICDialer.CloseIdleConnections.
+func (d *quicDialerHandshakeCompleter) CloseIdleConnections() {
+	d.Dialer.CloseIdleConnections()
+}
+
 // quicConnectionOwnsConn ensures that we close the UDPLikeConn.
 type quicConnectionOwnsConn struct {
 	// EarlyConnection is the embedded early connection

internal/netxlite/quic_test.go

@@ -38,7 +38,8 @@ func TestNewQUICDialer(t *testing.T) {
 		t.Fatal("invalid logger")
 	}
 	errWrapper := logger.Dialer.(*quicDialerErrWrapper)
-	base := errWrapper.QUICDialer.(*quicDialerQUICGo)
+	handshakeCompleter := errWrapper.QUICDialer.(*quicDialerHandshakeCompleter)
+	base := handshakeCompleter.Dialer.(*quicDialerQUICGo)
 	if base.QUICListener != ql {
 		t.Fatal("invalid quic listener")
 	}
@@ -227,6 +228,107 @@ func TestQUICDialerQUICGo(t *testing.T) {
 	})
 }
 
+func TestQUICDialerHandshakeCompleter(t *testing.T) {
+	t.Run("DialContext", func(t *testing.T) {
+		t.Run("in case of failure", func(t *testing.T) {
+			expected := errors.New("mocked error")
+			d := &quicDialerHandshakeCompleter{
+				Dialer: &mocks.QUICDialer{
+					MockDialContext: func(ctx context.Context, network, address string,
+						tlsConfig *tls.Config, quicConfig *quic.Config) (quic.EarlyConnection, error) {
+						return nil, expected
+					},
+				},
+			}
+			ctx := context.Background()
+			conn, err := d.DialContext(ctx, "udp", "8.8.8.8:443", &tls.Config{}, &quic.Config{})
+			if !errors.Is(err, expected) {
+				t.Fatal("unexpected err", err)
+			}
+			if conn != nil {
+				t.Fatal("expected nil conn")
+			}
+		})
+
+		t.Run("in case of context cancellation", func(t *testing.T) {
+			handshakeCtx, handshakeCancel := context.WithCancel(context.Background())
+			defer handshakeCancel()
+			ctx, cancel := context.WithCancel(context.Background())
+			var called bool
+			expected := &mocks.QUICEarlyConnection{
+				MockHandshakeComplete: func() context.Context {
+					cancel()
+					return handshakeCtx
+				},
+				MockCloseWithError: func(code quic.ApplicationErrorCode, reason string) error {
+					called = true
+					return nil
+				},
+			}
+			d := &quicDialerHandshakeCompleter{
+				Dialer: &mocks.QUICDialer{
+					MockDialContext: func(ctx context.Context, network, address string,
+						tlsConfig *tls.Config, quicConfig *quic.Config) (quic.EarlyConnection, error) {
+						return expected, nil
+					},
+				},
+			}
+			conn, err := d.DialContext(ctx, "udp", "8.8.8.8:443", &tls.Config{}, &quic.Config{})
+			if !errors.Is(err, context.Canceled) {
+				t.Fatal("unexpected err", err)
+			}
+			if conn != nil {
+				t.Fatal("expected nil conn")
+			}
+			if !called {
+				t.Fatal("not called")
+			}
+		})
+
+		t.Run("in case of success", func(t *testing.T) {
+			handshakeCtx, handshakeCancel := context.WithCancel(context.Background())
+			defer handshakeCancel()
+			expected := &mocks.QUICEarlyConnection{
+				MockHandshakeComplete: func() context.Context {
+					handshakeCancel()
+					return handshakeCtx
+				},
+			}
+			d := &quicDialerHandshakeCompleter{
+				Dialer: &mocks.QUICDialer{
+					MockDialContext: func(ctx context.Context, network, address string,
+						tlsConfig *tls.Config, quicConfig *quic.Config) (quic.EarlyConnection, error) {
+						return expected, nil
+					},
+				},
+			}
+			conn, err := d.DialContext(
+				context.Background(), "udp", "8.8.8.8:443", &tls.Config{}, &quic.Config{})
+			if err != nil {
+				t.Fatal(err)
+			}
+			if conn == nil {
+				t.Fatal("expected non-nil conn")
+			}
+		})
+	})
+
+	t.Run("CloseIdleConnections", func(t *testing.T) {
+		var forDialer bool
+		d := &quicDialerHandshakeCompleter{
+			Dialer: &mocks.QUICDialer{
+				MockCloseIdleConnections: func() {
+					forDialer = true
+				},
+			},
+		}
+		d.CloseIdleConnections()
+		if !forDialer {
+			t.Fatal("not called")
+		}
+	})
+}
+
 func TestQUICDialerResolver(t *testing.T) {
 	t.Run("CloseIdleConnections", func(t *testing.T) {
 		var (