tsnet: ensure funnel listener cleans up after itself when closed

Previously the funnel listener would leave artifacts in the serve
config. This caused weird out-of-sync effects like the admin panel
showing that funnel was enabled for a node, but the node rejecting
packets because the listener was closed.

This change resolves these synchronization issues by ensuring that
funnel listeners clean up the serve config when closed.

See also:
https://github.com/tailscale/tailscale/commit/e109cf9fdd405153a8d8c0ec52a87d7c8ce8689b

Updates #cleanup
Signed-off-by: Harry Harpham <harry@tailscale.com>

tsnet/tsnet.go

@@ -1228,12 +1228,26 @@ func (s *Server) ListenFunnel(network, addr string, opts ...FunnelOption) (net.L
 	}
 	domain := st.CertDomains[0]
 	hp := ipn.HostPort(domain + ":" + portStr)
+	var cleanupOnClose func() error
 	if !srvConfig.AllowFunnel[hp] {
 		mak.Set(&srvConfig.AllowFunnel, hp, true)
 		srvConfig.AllowFunnel[hp] = true
 		if err := lc.SetServeConfig(ctx, srvConfig); err != nil {
 			return nil, err
 		}
+		cleanupOnClose = func() error {
+			sc, err := lc.GetServeConfig(ctx)
+			if err != nil {
+				return fmt.Errorf("cleaning config changes: %w", err)
+			}
+			if sc.AllowFunnel != nil {
+				delete(sc.AllowFunnel, hp)
+			}
+			if err := lc.SetServeConfig(ctx, sc); err != nil {
+				return fmt.Errorf("cleaning config changes: %w", err)
+			}
+			return nil
+		}
 	}
 
 	// Start a funnel listener.
@@ -1241,6 +1255,7 @@ func (s *Server) ListenFunnel(network, addr string, opts ...FunnelOption) (net.L
 	if err != nil {
 		return nil, err
 	}
+	ln = &cleanupListener{Listener: ln, cleanup: cleanupOnClose}
 	return tls.NewListener(ln, tlsConfig), nil
 }
 
@@ -1449,3 +1464,30 @@ type addr struct{ ln *listener }
 
 func (a addr) Network() string { return a.ln.keys[0].network }
 func (a addr) String() string  { return a.ln.addr }
+
+// cleanupListener wraps a net.Listener with a function to be run on Close.
+type cleanupListener struct {
+	net.Listener
+	cleanup     func() error
+	cleanupOnce sync.Once
+}
+
+func (cl *cleanupListener) Close() error {
+	var cleanupErr error
+	cl.cleanupOnce.Do(func() {
+		if cl.cleanup != nil {
+			cleanupErr = cl.cleanup()
+		}
+	})
+	closeErr := cl.Listener.Close()
+	switch {
+	case closeErr != nil && cleanupErr != nil:
+		return fmt.Errorf("%w; also: %w", closeErr, cleanupErr)
+	case closeErr != nil:
+		return closeErr
+	case cleanupErr != nil:
+		return cleanupErr
+	default:
+		return nil
+	}
+}

tsnet/tsnet_test.go

@@ -13,6 +13,7 @@ import (
 	"crypto/tls"
 	"crypto/x509"
 	"crypto/x509/pkix"
+	"encoding/json"
 	"errors"
 	"flag"
 	"fmt"
@@ -33,6 +34,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/google/go-cmp/cmp"
 	dto "github.com/prometheus/client_model/go"
 	"github.com/prometheus/common/expfmt"
 	"golang.org/x/net/proxy"
@@ -741,6 +743,105 @@ func TestFunnel(t *testing.T) {
 	}
 }
 
+// TestFunnelClose ensures that the listener returned by ListenFunnel cleans up
+// after itself when closed. Specifically, changes made to the serve config
+// should be cleared.
+func TestFunnelClose(t *testing.T) {
+	marshalServeConfig := func(t *testing.T, sc ipn.ServeConfigView) string {
+		t.Helper()
+		return string(must.Get(json.MarshalIndent(sc, "", "\t")))
+	}
+
+	t.Run("simple", func(t *testing.T) {
+		controlURL, _ := startControl(t)
+		s, _, _ := startServer(t, t.Context(), controlURL, "s")
+
+		before := s.lb.ServeConfig()
+
+		ln := must.Get(s.ListenFunnel("tcp", ":443"))
+		ln.Close()
+
+		after := s.lb.ServeConfig()
+		if diff := cmp.Diff(marshalServeConfig(t, after), marshalServeConfig(t, before)); diff != "" {
+			t.Fatalf("expected serve config to be unchanged after close (-got, +want):\n%s", diff)
+		}
+	})
+
+	// Closing the listener shouldn't clear out config that predates it.
+	t.Run("no_clobbering", func(t *testing.T) {
+		controlURL, _ := startControl(t)
+		s, _, _ := startServer(t, t.Context(), controlURL, "s")
+
+		// To obtain config the listener might want to clobber, we:
+		//  - run a listener
+		//  - grab the config
+		//  - close the listener (clearing config)
+		ln := must.Get(s.ListenFunnel("tcp", ":443"))
+		before := s.lb.ServeConfig()
+		ln.Close()
+
+		// Now we manually write the config to the local backend (it should have
+		// been cleared), run the listener again, and close it again.
+		must.Do(s.lb.SetServeConfig(before.AsStruct(), ""))
+		ln = must.Get(s.ListenFunnel("tcp", ":443"))
+		ln.Close()
+
+		// The config should not have been cleared this time since it predated
+		// the most recent run.
+		after := s.lb.ServeConfig()
+		if diff := cmp.Diff(marshalServeConfig(t, after), marshalServeConfig(t, before)); diff != "" {
+			t.Fatalf("expected existing config to remain intact (-got, +want):\n%s", diff)
+		}
+	})
+
+	// Closing one listener shouldn't affect config for another listener.
+	t.Run("two_listeners", func(t *testing.T) {
+		controlURL, _ := startControl(t)
+		s, _, _ := startServer(t, t.Context(), controlURL, "s1")
+
+		// Start a listener on 443.
+		ln1 := must.Get(s.ListenFunnel("tcp", ":443"))
+		defer ln1.Close()
+
+		// Save the serve config for this original listener.
+		before := s.lb.ServeConfig()
+
+		// Now start and close a new listener on a different port.
+		ln2 := must.Get(s.ListenFunnel("tcp", ":8080"))
+		ln2.Close()
+
+		// The serve config for the original listener should be intact.
+		after := s.lb.ServeConfig()
+		if diff := cmp.Diff(marshalServeConfig(t, after), marshalServeConfig(t, before)); diff != "" {
+			t.Fatalf("expected existing config to remain intact (-got, +want):\n%s", diff)
+		}
+	})
+
+	// It should be possible to close a listener and free system resources even
+	// when the Server has been closed (or the listener should be automatically
+	// closed).
+	t.Run("after_server_close", func(t *testing.T) {
+		controlURL, _ := startControl(t)
+		s, _, _ := startServer(t, t.Context(), controlURL, "s")
+
+		ln := must.Get(s.ListenFunnel("tcp", ":443"))
+
+		// Close the server, then close the listener.
+		must.Do(s.Close())
+		// We don't care whether we get an error from the listener closing.
+		ln.Close()
+
+		// The listener should immediately return an error indicating closure.
+		_, err := ln.Accept()
+		// Looking for a string in the error sucks, but it's supposed to stay
+		// consistent:
+		// https://github.com/golang/go/blob/108b333d510c1f60877ac917375d7931791acfe6/src/internal/poll/fd.go#L20-L24
+		if err == nil || !strings.Contains(err.Error(), "use of closed network connection") {
+			t.Fatal("expected listener to be closed, got:", err)
+		}
+	})
+}
+
 func TestListenerClose(t *testing.T) {
 	tstest.Shard(t)
 	ctx := context.Background()