Pull request 340: Fix QUIC server

Squashed commit of the following: commit ce1a921 Author: Eugene Burkov <E.Burkov@AdGuard.COM> Date: Mon Apr 8 17:43:36 2024 +0300 proxy: imp tests commit a5863ce Author: Eugene Burkov <E.Burkov@AdGuard.COM> Date: Mon Apr 8 17:27:51 2024 +0300 proxy: imp code commit 3e0bad3 Author: Eugene Burkov <E.Burkov@AdGuard.COM> Date: Mon Apr 8 17:19:25 2024 +0300 proxy: fix quic listeners
AdguardTeam · Apr 8, 2024 · 480eb52 · 480eb52
1 parent 0e2cfca
commit 480eb52
Show file tree

Hide file tree

Showing 3 changed files with 62 additions and 28 deletions.
diff --git a/proxy/proxy.go b/proxy/proxy.go
@@ -129,6 +129,10 @@ type Proxy struct {
 	// quicListen are the listened QUIC connections.
 	quicListen []*quic.EarlyListener
 
+	// quicConns are UDP connections for all listened QUIC connections.  Those
+	// should be closed on shutdown, since *quic.EarlyListener doesn't close it.
+	quicConns []*net.UDPConn
+
 	// httpsListen are the listened HTTPS connections.
 	httpsListen []net.Listener
 
@@ -411,6 +415,9 @@ func (p *Proxy) Shutdown(_ context.Context) (err error) {
 	errs = closeAll(errs, p.quicListen...)
 	p.quicListen = nil
 
+	errs = closeAll(errs, p.quicConns...)
+	p.quicConns = nil
+
 	errs = closeAll(errs, p.dnsCryptUDPListen...)
 	p.dnsCryptUDPListen = nil
 

diff --git a/proxy/server_quic.go b/proxy/server_quic.go
@@ -64,13 +64,15 @@ const (
 // createQUICListeners creates QUIC listeners for the DoQ server.
 func (p *Proxy) createQUICListeners() error {
 	for _, a := range p.QUICListenAddr {
-		log.Info("Creating a QUIC listener")
+		log.Info("creating listener quic://%s", a)
 
 		conn, err := net.ListenUDP(bootstrap.NetworkUDP, a)
 		if err != nil {
 			return fmt.Errorf("listening to %s: %w", a, err)
 		}
 
+		p.quicConns = append(p.quicConns, conn)
+
 		v := newQUICAddrValidator(quicAddrValidatorCacheSize, quicAddrValidatorCacheTTL)
 		transport := &quic.Transport{
 			Conn:                conn,
@@ -88,7 +90,8 @@ func (p *Proxy) createQUICListeners() error {
 		}
 
 		p.quicListen = append(p.quicListen, quicListen)
-		log.Info("Listening to quic://%s", quicListen.Addr())
+
+		log.Info("listening quic://%s", quicListen.Addr())
 	}
 	return nil
 }

diff --git a/proxy/server_quic_test.go b/proxy/server_quic_test.go
@@ -18,22 +18,6 @@ import (
 
 func TestQuicProxy(t *testing.T) {
 	serverConfig, caPem := newTLSConfig(t)
-	dnsProxy := mustNew(t, &Config{
-		TLSListenAddr:          []*net.TCPAddr{net.TCPAddrFromAddrPort(localhostAnyPort)},
-		HTTPSListenAddr:        []*net.TCPAddr{net.TCPAddrFromAddrPort(localhostAnyPort)},
-		QUICListenAddr:         []*net.UDPAddr{net.UDPAddrFromAddrPort(localhostAnyPort)},
-		TLSConfig:              serverConfig,
-		UpstreamConfig:         newTestUpstreamConfig(t, defaultTimeout, testDefaultUpstreamAddr),
-		TrustedProxies:         defaultTrustedProxies,
-		RatelimitSubnetLenIPv4: 24,
-		RatelimitSubnetLenIPv6: 64,
-	})
-
-	// Start listening.
-	ctx := context.Background()
-	err := dnsProxy.Start(ctx)
-	require.NoError(t, err)
-	testutil.CleanupAndRequireSuccess(t, func() (err error) { return dnsProxy.Shutdown(ctx) })
 
 	roots := x509.NewCertPool()
 	roots.AppendCertsFromPEM(caPem)
@@ -43,23 +27,63 @@ func TestQuicProxy(t *testing.T) {
 		NextProtos: append([]string{NextProtoDQ}, compatProtoDQ...),
 	}
 
-	// Create a DNS-over-QUIC client connection.
-	addr := dnsProxy.Addr(ProtoQUIC)
+	conf := &Config{
+		QUICListenAddr:         []*net.UDPAddr{net.UDPAddrFromAddrPort(localhostAnyPort)},
+		TLSConfig:              serverConfig,
+		UpstreamConfig:         newTestUpstreamConfig(t, defaultTimeout, testDefaultUpstreamAddr),
+		TrustedProxies:         defaultTrustedProxies,
+		RatelimitSubnetLenIPv4: 24,
+		RatelimitSubnetLenIPv6: 64,
+	}
 
-	// Open a QUIC connection.
-	conn, err := quic.DialAddrEarly(context.Background(), addr.String(), tlsConfig, nil)
-	require.NoError(t, err)
-	testutil.CleanupAndRequireSuccess(t, func() (err error) {
-		return conn.CloseWithError(DoQCodeNoError, "")
+	var addr *net.UDPAddr
+	t.Run("run", func(t *testing.T) {
+		dnsProxy := mustNew(t, conf)
+
+		ctx := context.Background()
+		err := dnsProxy.Start(ctx)
+		require.NoError(t, err)
+		testutil.CleanupAndRequireSuccess(t, func() (err error) { return dnsProxy.Shutdown(ctx) })
+
+		addr = testutil.RequireTypeAssert[*net.UDPAddr](t, dnsProxy.Addr(ProtoQUIC))
+
+		conn, err := quic.DialAddrEarly(context.Background(), addr.String(), tlsConfig, nil)
+		require.NoError(t, err)
+		testutil.CleanupAndRequireSuccess(t, func() (err error) {
+			return conn.CloseWithError(DoQCodeNoError, "")
+		})
+
+		for range 10 {
+			sendTestQUICMessage(t, conn, DoQv1)
+
+			// Send a message encoded for a draft version as well.
+			sendTestQUICMessage(t, conn, DoQv1Draft)
+		}
 	})
+	require.False(t, t.Failed())
+
+	conf.QUICListenAddr = []*net.UDPAddr{addr}
+	conf.UpstreamConfig = newTestUpstreamConfig(t, defaultTimeout, testDefaultUpstreamAddr)
+
+	t.Run("rerun", func(t *testing.T) {
+		dnsProxy := mustNew(t, conf)
+
+		ctx := context.Background()
+		err := dnsProxy.Start(ctx)
+		require.NoError(t, err)
+		testutil.CleanupAndRequireSuccess(t, func() (err error) { return dnsProxy.Shutdown(ctx) })
+
+		conn, err := quic.DialAddrEarly(context.Background(), addr.String(), tlsConfig, nil)
+		require.NoError(t, err)
+		testutil.CleanupAndRequireSuccess(t, func() (err error) {
+			return conn.CloseWithError(DoQCodeNoError, "")
+		})
 
-	// Send several test messages.
-	for range 10 {
 		sendTestQUICMessage(t, conn, DoQv1)
 
 		// Send a message encoded for a draft version as well.
 		sendTestQUICMessage(t, conn, DoQv1Draft)
-	}
+	})
 }
 
 func TestQuicProxy_largePackets(t *testing.T) {