Add -max-udp-size option

Summary:
Add a -max-udp-size command-line argument to fbdns.  This allows us to limit the size of UDP responses, as [recommended in a forthcoming IETF standard](https://datatracker.ietf.org/doc/draft-ietf-dnsop-avoid-fragmentation/).

This argument does not affect the Don't Fragment setting.

Reviewed By: deathowl

Differential Revision: D48613667

fbshipit-source-id: c5c8b9e7e5e51b9e0061cfb54f31ab21b942b49f

dnsrocks/cmd/dnsrocks/dnsrocks.go

@@ -81,6 +81,7 @@ func main() {
 
 	// DNS Server config
 	cliflags.IntVar(&serverConfig.Port, "port", 8053, "port to run on")
+	cliflags.IntVar(&serverConfig.MaxUDPSize, "max-udp-size", 0, "Maximum UDP response size (default: none)")
 	cliflags.BoolVar(&serverConfig.TCP, "tcp", true, "Whether or not to also listen on TCP.")
 	cliflags.IntVar(&serverConfig.MaxTCPQueries, "tcp-max-queries", -1, "Maximum number of queries handled on a single TCP connection before closing the socket. This also applies for TLS. (unlimited if -1).")
 	// Idle Timeout default is based on miekg/dns original default: https://fburl.com/t0tmjp2c

dnsrocks/fbserver/config.go

@@ -30,6 +30,7 @@ import (
 type ServerConfig struct {
 	IPAns          ipAns
 	Port           int
+	MaxUDPSize     int
 	TCP            bool
 	TLS            bool
 	ReusePort      int

dnsrocks/fbserver/server.go

@@ -26,6 +26,7 @@ import (
 	"time"
 
 	"github.com/coredns/coredns/plugin"
+	"github.com/coredns/coredns/plugin/bufsize"
 	"github.com/golang/glog"
 	"github.com/miekg/dns"
 	"golang.org/x/sys/unix"
@@ -305,6 +306,18 @@ func (srv *Server) Start() (err error) {
 				srv.conf.DNSSECConfig.Keys)
 		}
 
+		if srv.conf.MaxUDPSize >= dns.MinMsgSize && srv.conf.MaxUDPSize <= dns.MaxMsgSize {
+			glog.Infof("Limiting UDP response size to %d", srv.conf.MaxUDPSize)
+			handler.defaultHandler = bufsize.Bufsize{
+				Size: srv.conf.MaxUDPSize,
+				Next: handler.defaultHandler,
+			}
+		} else if srv.conf.MaxUDPSize > 0 {
+			glog.Warningf("Ignoring non-compliant -max-udp-size: %d", srv.conf.MaxUDPSize)
+		} else {
+			glog.Infof("Max UDP size not set")
+		}
+
 		addr := joinAddress(ip, srv.conf.Port)
 
 		for i := 0; i < numListeners; i++ {

dnsrocks/fbserver/server_test.go

@@ -15,6 +15,7 @@ package fbserver
 
 import (
 	"crypto/tls"
+	"fmt"
 	"io"
 	"os"
 	"testing"
@@ -378,3 +379,96 @@ func TestMultipleQueryoverTCP(t *testing.T) {
 		require.Equal(t, m.Question[0].Qtype, r.Question[0].Qtype, "Mismatching question type")
 	}
 }
+
+func TestMaxUDPSize(t *testing.T) {
+	testCases := []struct {
+		clientMax uint16
+		serverMax int
+		truncated bool
+	}{
+		// No EDNS0 (clientMax is unset)
+		{
+			truncated: true,
+		},
+		{
+			serverMax: 2048,
+			truncated: true,
+		},
+		// One or both is too small.
+		{
+			clientMax: 512,
+			// serverMax is unset
+			truncated: true,
+		},
+		{
+			clientMax: 512,
+			serverMax: 512,
+			truncated: true,
+		},
+		{
+			clientMax: 2048,
+			serverMax: 512,
+			truncated: true,
+		},
+		{
+			clientMax: 512,
+			serverMax: 2048,
+			truncated: true,
+		},
+		// Both are big enough.
+		{
+			clientMax: 2048,
+			serverMax: 2048,
+			truncated: false,
+		},
+		{
+			clientMax: 2048,
+			serverMax: 8192,
+			truncated: false,
+		},
+		// Invalid server max sizes
+		{
+			clientMax: 2048,
+			serverMax: 511, // Too small, ignored.
+			truncated: false,
+		},
+		{
+			clientMax: 2048,
+			serverMax: 65536, // Too large, ignored.
+			truncated: false,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(fmt.Sprintf("%d/%d", tc.clientMax, tc.serverMax), func(t *testing.T) {
+			config := makeTestServerConfig(true, false)
+			config.MaxUDPSize = tc.serverMax
+			portMap, srv := makeTestServer(t, config)
+			defer srv.Shutdown()
+
+			c := new(dns.Client)
+			query := new(dns.Msg)
+			qname := "lotofns.example.org."
+			query.SetQuestion(qname, dns.TypeNS)
+			if tc.clientMax > 0 {
+				query.SetEdns0(tc.clientMax, false)
+			}
+			udpResponse, _, err := c.Exchange(query, portMap["udp"])
+			require.Nil(t, err)
+
+			if tc.clientMax > 0 {
+				require.NotNil(t, udpResponse.IsEdns0())
+			} else {
+				require.Nil(t, udpResponse.IsEdns0())
+			}
+
+			require.Equal(t, tc.truncated, udpResponse.Truncated)
+
+			c.Net = "tcp"
+			tcpResponse, _, err := c.Exchange(query, portMap["tcp"])
+			require.Nil(t, err)
+			require.NotNil(t, tcpResponse)
+			require.False(t, tcpResponse.Truncated)
+		})
+	}
+}