No more dns globals, proper cleanup on shutdown (#1667)

dns_server_test.go

@@ -1,19 +1,31 @@
 package nebula
 
 import (
+	"context"
+	"io"
+	"net"
 	"net/netip"
+	"strconv"
 	"testing"
+	"time"
 
 	"github.com/miekg/dns"
 	"github.com/sirupsen/logrus"
 	"github.com/slackhq/nebula/config"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
 func TestParsequery(t *testing.T) {
 	l := logrus.New()
 	hostMap := &HostMap{}
-	ds := newDnsRecords(l, &CertState{}, hostMap)
+	ds := &dnsServer{
+		l:       l,
+		dnsMap4: make(map[string]netip.Addr),
+		dnsMap6: make(map[string]netip.Addr),
+		hostMap: hostMap,
+	}
+	ds.enabled.Store(true)
 	addrs := []netip.Addr{
 		netip.MustParseAddr("1.2.3.4"),
 		netip.MustParseAddr("1.2.3.5"),
@@ -71,3 +83,209 @@ func Test_getDnsServerAddr(t *testing.T) {
 	}
 	assert.Equal(t, "[::]:1", getDnsServerAddr(c))
 }
+
+func newTestDnsServer(t *testing.T) (*dnsServer, *config.C) {
+	t.Helper()
+	l := logrus.New()
+	l.Out = io.Discard
+	ds := &dnsServer{
+		l:       l,
+		ctx:     context.Background(),
+		dnsMap4: make(map[string]netip.Addr),
+		dnsMap6: make(map[string]netip.Addr),
+		hostMap: &HostMap{},
+	}
+	ds.mux = dns.NewServeMux()
+	ds.mux.HandleFunc(".", ds.handleDnsRequest)
+	return ds, config.NewC(l)
+}
+
+func setDnsConfig(c *config.C, host string, port string, amLighthouse, serveDns bool) {
+	c.Settings["lighthouse"] = map[string]any{
+		"am_lighthouse": amLighthouse,
+		"serve_dns":     serveDns,
+		"dns": map[string]any{
+			"host": host,
+			"port": port,
+		},
+	}
+}
+
+func TestDnsServer_reload_initial_disabled(t *testing.T) {
+	ds, c := newTestDnsServer(t)
+	setDnsConfig(c, "127.0.0.1", "0", true, false)
+
+	require.NoError(t, ds.reload(c, true))
+	assert.False(t, ds.enabled.Load())
+	assert.Equal(t, "127.0.0.1:0", ds.addr)
+	assert.Nil(t, ds.server)
+}
+
+func TestDnsServer_reload_initial_enabled(t *testing.T) {
+	ds, c := newTestDnsServer(t)
+	setDnsConfig(c, "127.0.0.1", "0", true, true)
+
+	require.NoError(t, ds.reload(c, true))
+	assert.True(t, ds.enabled.Load())
+	assert.Equal(t, "127.0.0.1:0", ds.addr)
+	// initial never starts a runner; that's Control.Start's job
+	assert.Nil(t, ds.server)
+}
+
+func TestDnsServer_reload_initial_serveDnsWithoutLighthouse(t *testing.T) {
+	ds, c := newTestDnsServer(t)
+	setDnsConfig(c, "127.0.0.1", "0", false, true)
+
+	require.NoError(t, ds.reload(c, true))
+	// Wants DNS but isn't a lighthouse: gated off, no runner.
+	assert.False(t, ds.enabled.Load())
+}
+
+func TestDnsServer_reload_sameAddr_noOp(t *testing.T) {
+	ds, c := newTestDnsServer(t)
+	setDnsConfig(c, "127.0.0.1", "0", true, true)
+
+	require.NoError(t, ds.reload(c, true))
+	// No server running yet, no addr change. Reload should not spawn anything.
+	require.NoError(t, ds.reload(c, false))
+	assert.True(t, ds.enabled.Load())
+	assert.Nil(t, ds.server)
+}
+
+func TestDnsServer_StartStop_lifecycle(t *testing.T) {
+	// Bind to a real (random) UDP port so we exercise the actual
+	// ListenAndServe + Shutdown plumbing including the started-chan race fix.
+	port := freeUDPPort(t)
+
+	ds, c := newTestDnsServer(t)
+	setDnsConfig(c, "127.0.0.1", port, true, true)
+	require.NoError(t, ds.reload(c, true))
+
+	done := make(chan struct{})
+	go func() {
+		ds.Start()
+		close(done)
+	}()
+
+	waitFor(t, func() bool {
+		ds.serverMu.Lock()
+		started := ds.started
+		ds.serverMu.Unlock()
+		if started == nil {
+			return false
+		}
+		select {
+		case <-started:
+			return true
+		default:
+			return false
+		}
+	})
+
+	ds.Stop()
+	select {
+	case <-done:
+	case <-time.After(5 * time.Second):
+		t.Fatal("Start did not return after Stop")
+	}
+}
+
+func TestDnsServer_Stop_beforeBind_doesNotHang(t *testing.T) {
+	// Stop called immediately after Start should not deadlock even if bind
+	// hasn't completed yet. This exercises the started-chan close-on-bind-fail
+	// path: by binding to an obviously bad port (privileged) we get a fast
+	// bind error before NotifyStartedFunc fires.
+	ds, c := newTestDnsServer(t)
+	// Use a port that should fail to bind (negative would be invalid, use a
+	// host that won't resolve to ensure listenUDP fails quickly).
+	setDnsConfig(c, "256.256.256.256", "53", true, true)
+	require.NoError(t, ds.reload(c, true))
+
+	done := make(chan struct{})
+	go func() {
+		ds.Start()
+		close(done)
+	}()
+
+	// Give Start a moment to attempt the bind and fail.
+	select {
+	case <-done:
+		// Bind failed and Start returned; Stop should be a no-op.
+	case <-time.After(time.Second):
+		t.Fatal("Start did not return after a bad bind")
+	}
+
+	stopped := make(chan struct{})
+	go func() {
+		ds.Stop()
+		close(stopped)
+	}()
+	select {
+	case <-stopped:
+	case <-time.After(time.Second):
+		t.Fatal("Stop hung after a failed bind")
+	}
+}
+
+func TestDnsServer_reload_disable_stopsRunningServer(t *testing.T) {
+	port := freeUDPPort(t)
+	ds, c := newTestDnsServer(t)
+	setDnsConfig(c, "127.0.0.1", port, true, true)
+	require.NoError(t, ds.reload(c, true))
+
+	startReturned := make(chan struct{})
+	go func() {
+		ds.Start()
+		close(startReturned)
+	}()
+	waitForBind(t, ds)
+
+	// Toggle serve_dns off; reload should shut the running server down.
+	setDnsConfig(c, "127.0.0.1", port, true, false)
+	require.NoError(t, ds.reload(c, false))
+	select {
+	case <-startReturned:
+	case <-time.After(5 * time.Second):
+		t.Fatal("Start did not return after reload disabled DNS")
+	}
+	assert.False(t, ds.enabled.Load())
+}
+
+func freeUDPPort(t *testing.T) string {
+	t.Helper()
+	conn, err := net.ListenPacket("udp", "127.0.0.1:0")
+	require.NoError(t, err)
+	port := conn.LocalAddr().(*net.UDPAddr).Port
+	require.NoError(t, conn.Close())
+	return strconv.Itoa(port)
+}
+
+func waitForBind(t *testing.T, ds *dnsServer) {
+	t.Helper()
+	waitFor(t, func() bool {
+		ds.serverMu.Lock()
+		started := ds.started
+		ds.serverMu.Unlock()
+		if started == nil {
+			return false
+		}
+		select {
+		case <-started:
+			return true
+		default:
+			return false
+		}
+	})
+}
+
+func waitFor(t *testing.T, cond func() bool) {
+	t.Helper()
+	deadline := time.Now().Add(5 * time.Second)
+	for time.Now().Before(deadline) {
+		if cond() {
+			return
+		}
+		time.Sleep(5 * time.Millisecond)
+	}
+	t.Fatal("timed out waiting for condition")
+}