diff --git a/core/dnsserver/config.go b/core/dnsserver/config.go index 494c3d3ae..f17dd4570 100644 --- a/core/dnsserver/config.go +++ b/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 diff --git a/core/dnsserver/server_quic.go b/core/dnsserver/server_quic.go index 9c89334b9..a2357bcfa 100644 --- a/core/dnsserver/server_quic.go +++ b/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") + } } } diff --git a/core/dnsserver/server_quic_test.go b/core/dnsserver/server_quic_test.go index bcbb269ed..d15164358 100644 --- a/core/dnsserver/server_quic_test.go +++ b/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) + } +}