core: Add connection-level concurrency limiting to DNS-over-QUIC (#8213)

* core: Add connection-level concurrency limiting to DNS-over-QUIC

This PR adds max connections to prevent unbounded connection goroutine growth
for DoQ

Signed-off-by: Yong Tang <yong.tang.github@outlook.com>

* Add test to cover change

Signed-off-by: Yong Tang <yong.tang.github@outlook.com>

---------

Signed-off-by: Yong Tang <yong.tang.github@outlook.com>

core/dnsserver/config.go

@@ -74,6 +74,9 @@ type Config struct {
 	// This is nil if not specified, allowing for a default to be used.
 	MaxQUICStreams *int
 
+	// MaxQUICConnections is the maximum number of concurrent connections.
+	MaxQUICConnections *int
+
 	// MaxQUICWorkerPoolSize defines the size of the worker pool for processing QUIC streams.
 	// This is nil if not specified, allowing for a default to be used.
 	MaxQUICWorkerPoolSize *int

core/dnsserver/server_quic.go

@@ -43,6 +43,9 @@ const (
 
 	// DefaultQUICStreamWorkers is the default number of workers for processing QUIC streams.
 	DefaultQUICStreamWorkers = 1024
+
+	// DefaultQUICMaxConnections is the default maximum number of concurrent connections.
+	DefaultQUICMaxConnections = 200
 )
 
 // ServerQUIC represents an instance of a DNS-over-QUIC server.
@@ -54,6 +57,8 @@ type ServerQUIC struct {
 	quicListener      *quic.Listener
 	maxStreams        int
 	streamProcessPool chan struct{}
+	maxConnections    int
+	connSem           chan struct{}
 }
 
 // NewServerQUIC returns a new CoreDNS QUIC server and compiles all plugin in to it.
@@ -93,6 +98,15 @@ func NewServerQUIC(addr string, group []*Config) (*ServerQUIC, error) {
 		// Enable 0-RTT by default for all connections on the server-side.
 		Allow0RTT: true,
 	}
+	maxConnections := DefaultQUICMaxConnections
+	if len(group) > 0 && group[0] != nil && group[0].MaxQUICConnections != nil {
+		maxConnections = *group[0].MaxQUICConnections
+	}
+
+	var connSem chan struct{}
+	if maxConnections > 0 {
+		connSem = make(chan struct{}, maxConnections)
+	}
 
 	return &ServerQUIC{
 		Server:            s,
@@ -100,6 +114,8 @@ func NewServerQUIC(addr string, group []*Config) (*ServerQUIC, error) {
 		quicConfig:        quicConfig,
 		maxStreams:        maxStreams,
 		streamProcessPool: make(chan struct{}, streamProcessPoolSize),
+		maxConnections:    maxConnections,
+		connSem:           connSem,
 	}, nil
 }
 
@@ -133,8 +149,21 @@ func (s *ServerQUIC) ServeQUIC() error {
 			s.closeQUICConn(conn, DoQCodeInternalError)
 			return err
 		}
+		if s.connSem == nil {
+			go s.serveQUICConnection(conn)
+			continue
+		}
 
-		go s.serveQUICConnection(conn)
+		select {
+		case s.connSem <- struct{}{}:
+			go func(c *quic.Conn) {
+				defer func() { <-s.connSem }()
+				s.serveQUICConnection(c)
+			}(conn)
+
+		default:
+			_ = conn.CloseWithError(0, "too many connections")
+		}
 	}
 }
 

core/dnsserver/server_quic_test.go

@@ -10,6 +10,7 @@ import (
 	"crypto/x509/pkix"
 	"errors"
 	"math/big"
+	"strings"
 	"testing"
 	"time"
 
@@ -997,3 +998,124 @@ func TestServerQUIC_ServeQUIC_ConnectionStateExposesSNI(t *testing.T) {
 		t.Fatal("did not receive connection state from plugin")
 	}
 }
+
+func TestServerQUICServeQUICDefaultMaxConnections(t *testing.T) {
+	const maxConnections = 200
+
+	config := testConfig("quic", echoPlugin{})
+	config.TLSConfig = mustMakeQUICServerTLSConfig(t)
+
+	server, err := NewServerQUIC(transport.QUIC+"://127.0.0.1:0", []*Config{config})
+	if err != nil {
+		t.Fatalf("NewServerQUIC() failed: %v", err)
+	}
+
+	pc, err := server.ListenPacket()
+	if err != nil {
+		t.Fatalf("ListenPacket() failed: %v", err)
+	}
+	defer pc.Close()
+
+	serveErrCh := make(chan error, 1)
+	go func() {
+		serveErrCh <- server.ServeQUIC()
+	}()
+
+	defer func() {
+		_ = server.Stop()
+		select {
+		case <-serveErrCh:
+		case <-time.After(2 * time.Second):
+		}
+	}()
+
+	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
+	defer cancel()
+
+	query := func(conn *quic.Conn) error {
+		stream, err := conn.OpenStreamSync(ctx)
+		if err != nil {
+			return err
+		}
+
+		q := new(dns.Msg)
+		q.SetQuestion("example.com.", dns.TypeA)
+		q.Id = 0
+		wire, err := q.Pack()
+		if err != nil {
+			return err
+		}
+		if _, err := stream.Write(AddPrefix(wire)); err != nil {
+			return err
+		}
+		if err := stream.Close(); err != nil {
+			return err
+		}
+		_, err = readDOQMessage(stream)
+		return err
+	}
+
+	connections := make([]*quic.Conn, 0, maxConnections)
+	defer func() {
+		for _, conn := range connections {
+			_ = conn.CloseWithError(DoQCodeNoError, "")
+		}
+	}()
+
+	addr := pc.LocalAddr().String()
+	for i := range maxConnections {
+		conn, err := quic.DialAddr(ctx, addr, mustMakeQUICClientTLSConfig(), &quic.Config{})
+		if err != nil {
+			t.Fatalf("quic.DialAddr() for connection %d failed: %v", i+1, err)
+		}
+		connections = append(connections, conn)
+
+		// Receiving a response proves ServeQUIC accepted the connection and
+		// started serveQUICConnection, which keeps its connection slot held.
+		if err := query(conn); err != nil {
+			t.Fatalf("query on connection %d failed: %v", i+1, err)
+		}
+	}
+
+	overflow, err := quic.DialAddr(ctx, addr, mustMakeQUICClientTLSConfig(), &quic.Config{})
+	if err == nil {
+		err = query(overflow)
+	}
+	if err == nil {
+		_ = overflow.CloseWithError(DoQCodeNoError, "")
+		t.Fatalf("connection %d was served; want it rejected by the default connection limit", maxConnections+1)
+	}
+
+	if overflow != nil {
+		select {
+		case <-overflow.Context().Done():
+			err = context.Cause(overflow.Context())
+		case <-time.After(2 * time.Second):
+		}
+		_ = overflow.CloseWithError(DoQCodeNoError, "")
+	}
+	if err == nil || !strings.Contains(err.Error(), "too many connections") {
+		t.Fatalf("connection %d rejection error = %v, want %q", maxConnections+1, err, "too many connections")
+	}
+
+	// Releasing one accepted connection must release its semaphore slot.
+	_ = connections[0].CloseWithError(DoQCodeNoError, "")
+	connections = connections[1:]
+
+	deadline := time.Now().Add(5 * time.Second)
+	for {
+		replacement, dialErr := quic.DialAddr(ctx, addr, mustMakeQUICClientTLSConfig(), &quic.Config{})
+		if dialErr == nil {
+			if queryErr := query(replacement); queryErr == nil {
+				connections = append(connections, replacement)
+				break
+			}
+			_ = replacement.CloseWithError(DoQCodeNoError, "")
+		}
+
+		if time.Now().After(deadline) {
+			t.Fatal("replacement connection was not served after a connection slot was released")
+		}
+		time.Sleep(10 * time.Millisecond)
+	}
+}