Add option to specify DERP servers

cmd/wush/cp.go

@@ -1,6 +1,7 @@
 package main
 
 import (
+	"encoding/json"
 	"errors"
 	"fmt"
 	"io"
@@ -19,6 +20,7 @@ import (
 	"github.com/coder/wush/tsserver"
 	"github.com/schollz/progressbar/v3"
 	"tailscale.com/net/netns"
+	"tailscale.com/tailcfg"
 	"tailscale.com/types/ptr"
 )
 
@@ -65,13 +67,10 @@ func initAuth(authFlag *string, ca *overlay.ClientAuth) serpent.MiddlewareFunc {
 	}
 }
 
-func sendOverlayMW(opts *sendOverlayOpts, send **overlay.Send, logger *slog.Logger, logf *func(str string, args ...any)) serpent.MiddlewareFunc {
+func sendOverlayMW(opts *sendOverlayOpts, send **overlay.Send, logger *slog.Logger, dm *tailcfg.DERPMap, logf *func(str string, args ...any)) serpent.MiddlewareFunc {
 	return func(next serpent.HandlerFunc) serpent.HandlerFunc {
 		return func(i *serpent.Invocation) error {
-			dm, err := tsserver.DERPMapTailscale(i.Context())
-			if err != nil {
-				return err
-			}
+			var err error
 
 			newSend := overlay.NewSendOverlay(logger, dm)
 			newSend.Auth = opts.clientAuth
@@ -90,6 +89,30 @@ func sendOverlayMW(opts *sendOverlayOpts, send **overlay.Send, logger *slog.Logg
 	}
 }
 
+func derpMap(fi *string, dm *tailcfg.DERPMap) serpent.MiddlewareFunc {
+	return func(next serpent.HandlerFunc) serpent.HandlerFunc {
+		return func(i *serpent.Invocation) error {
+			if *fi == "" {
+				_dm, err := tsserver.DERPMapTailscale(i.Context())
+				if err != nil {
+					return fmt.Errorf("request derpmap from tailscale: %w", err)
+				}
+				*dm = *_dm
+			} else {
+				data, err := os.ReadFile(*fi)
+				if err != nil {
+					return fmt.Errorf("read derp config file: %w", err)
+				}
+				if err := json.Unmarshal(data, dm); err != nil {
+					return fmt.Errorf("unmarshal derp config: %w", err)
+				}
+			}
+
+			return next(i)
+		}
+	}
+}
+
 type sendOverlayOpts struct {
 	authKey          string
 	clientAuth       overlay.ClientAuth
@@ -99,10 +122,12 @@ type sendOverlayOpts struct {
 
 func cpCmd() *serpent.Command {
 	var (
-		verbose bool
-		logger  = new(slog.Logger)
-		logf    = func(str string, args ...any) {}
+		verbose   bool
+		derpmapFi string
+		logger    = new(slog.Logger)
+		logf      = func(str string, args ...any) {}
 
+		dm          = new(tailcfg.DERPMap)
 		overlayOpts = new(sendOverlayOpts)
 		send        = new(overlay.Send)
 	)
@@ -119,12 +144,13 @@ func cpCmd() *serpent.Command {
 			serpent.RequireNArgs(1),
 			initLogger(&verbose, ptr.To(false), logger, &logf),
 			initAuth(&overlayOpts.authKey, &overlayOpts.clientAuth),
-			sendOverlayMW(overlayOpts, &send, logger, &logf),
+			derpMap(&derpmapFi, dm),
+			sendOverlayMW(overlayOpts, &send, logger, dm, &logf),
 		),
 		Handler: func(inv *serpent.Invocation) error {
 			ctx := inv.Context()
 
-			s, err := tsserver.NewServer(ctx, logger, send)
+			s, err := tsserver.NewServer(ctx, logger, send, dm)
 			if err != nil {
 				return err
 			}
@@ -215,6 +241,12 @@ func cpCmd() *serpent.Command {
 				Default:     "",
 				Value:       serpent.StringOf(&overlayOpts.authKey),
 			},
+			{
+				Flag:        "derp-config-file",
+				Description: "File which specifies the DERP config to use. In the structure of https://pkg.go.dev/tailscale.com@v1.74.1/tailcfg#DERPMap. By default, https://controlplane.tailscale.com/derpmap/default is used.",
+				Default:     "",
+				Value:       serpent.StringOf(&derpmapFi),
+			},
 			{
 				Flag:    "stun-ip-override",
 				Default: "",

cmd/wush/portforward.go

@@ -15,6 +15,7 @@ import (
 
 	"golang.org/x/xerrors"
 	"tailscale.com/net/netns"
+	"tailscale.com/tailcfg"
 	"tailscale.com/tsnet"
 	"tailscale.com/types/ptr"
 
@@ -27,10 +28,12 @@ import (
 
 func portForwardCmd() *serpent.Command {
 	var (
-		verbose bool
-		logger  = new(slog.Logger)
-		logf    = func(str string, args ...any) {}
+		verbose   bool
+		derpmapFi string
+		logger    = new(slog.Logger)
+		logf      = func(str string, args ...any) {}
 
+		dm          = new(tailcfg.DERPMap)
 		overlayOpts = new(sendOverlayOpts)
 		send        = new(overlay.Send)
 		tcpForwards []string // <port>:<port>
@@ -64,7 +67,8 @@ func portForwardCmd() *serpent.Command {
 		Middleware: serpent.Chain(
 			initLogger(&verbose, ptr.To(false), logger, &logf),
 			initAuth(&overlayOpts.authKey, &overlayOpts.clientAuth),
-			sendOverlayMW(overlayOpts, &send, logger, &logf),
+			derpMap(nil, dm),
+			sendOverlayMW(overlayOpts, &send, logger, dm, &logf),
 		),
 		Handler: func(inv *serpent.Invocation) error {
 			ctx, cancel := context.WithCancel(inv.Context())
@@ -78,7 +82,7 @@ func portForwardCmd() *serpent.Command {
 				return errors.New("no port-forwards requested")
 			}
 
-			s, err := tsserver.NewServer(ctx, logger, send)
+			s, err := tsserver.NewServer(ctx, logger, send, dm)
 			if err != nil {
 				return err
 			}
@@ -97,8 +101,6 @@ func portForwardCmd() *serpent.Command {
 			if err != nil {
 				return err
 			}
-			ts.Logf = func(string, ...any) {}
-			ts.UserLogf = func(string, ...any) {}
 
 			logf("Bringing WireGuard up..")
 			ts.Up(ctx)
@@ -179,6 +181,12 @@ func portForwardCmd() *serpent.Command {
 				Default:     "",
 				Value:       serpent.StringOf(&overlayOpts.authKey),
 			},
+			{
+				Flag:        "derp-config-file",
+				Description: "File which specifies the DERP config to use. In the structure of https://pkg.go.dev/tailscale.com@v1.74.1/tailcfg#DERPMap.",
+				Default:     "",
+				Value:       serpent.StringOf(&derpmapFi),
+			},
 			{
 				Flag:    "stun-ip-override",
 				Default: "",

cmd/wush/serve.go

@@ -20,6 +20,7 @@ import (
 	"golang.org/x/xerrors"
 	"tailscale.com/ipn/store"
 	"tailscale.com/net/netns"
+	"tailscale.com/tailcfg"
 	"tailscale.com/tsnet"
 
 	cslog "cdr.dev/slog"
@@ -38,12 +39,18 @@ func serveCmd() *serpent.Command {
 		verbose     bool
 		enabled     = []string{}
 		disabled    = []string{}
+		derpmapFi   string
+
+		dm = new(tailcfg.DERPMap)
 	)
 	return &serpent.Command{
 		Use:     "serve",
 		Aliases: []string{"host"},
 		Short:   "Run the wush server.",
 		Long:    "Runs the wush server. Allows other wush CLIs to connect to this computer.",
+		Middleware: serpent.Chain(
+			derpMap(&derpmapFi, dm),
+		),
 		Handler: func(inv *serpent.Invocation) error {
 			ctx := inv.Context()
 			var logSink io.Writer = io.Discard
@@ -54,12 +61,9 @@ func serveCmd() *serpent.Command {
 			hlog := func(format string, args ...any) {
 				fmt.Fprintf(inv.Stderr, format+"\n", args...)
 			}
-			dm, err := tsserver.DERPMapTailscale(ctx)
-			if err != nil {
-				return err
-			}
 			r := overlay.NewReceiveOverlay(logger, hlog, dm)
 
+			var err error
 			switch overlayType {
 			case "derp":
 				err = r.PickDERPHome(ctx)
@@ -89,7 +93,7 @@ func serveCmd() *serpent.Command {
 				hlog("The auth key has been printed to stdout")
 			}
 
-			s, err := tsserver.NewServer(ctx, logger, r)
+			s, err := tsserver.NewServer(ctx, logger, r, dm)
 			if err != nil {
 				return err
 			}
@@ -209,6 +213,12 @@ func serveCmd() *serpent.Command {
 				Default:     "",
 				Value:       serpent.EnumArrayOf(&disabled, "ssh", "cp", "port-forward"),
 			},
+			{
+				Flag:        "derp-config-file",
+				Description: "File which specifies the DERP config to use. In the structure of https://pkg.go.dev/tailscale.com@v1.74.1/tailcfg#DERPMap.",
+				Default:     "",
+				Value:       serpent.StringOf(&derpmapFi),
+			},
 		},
 	}
 }

cmd/wush/ssh.go

@@ -21,10 +21,13 @@ import (
 
 func sshCmd() *serpent.Command {
 	var (
-		verbose     bool
-		quiet       bool
-		logger      = new(slog.Logger)
-		logf        = func(str string, args ...any) {}
+		verbose   bool
+		quiet     bool
+		derpmapFi string
+		logger    = new(slog.Logger)
+		logf      = func(str string, args ...any) {}
+
+		dm          = new(tailcfg.DERPMap)
 		overlayOpts = new(sendOverlayOpts)
 		send        = new(overlay.Send)
 	)
@@ -37,12 +40,13 @@ func sshCmd() *serpent.Command {
 		Middleware: serpent.Chain(
 			initLogger(&verbose, &quiet, logger, &logf),
 			initAuth(&overlayOpts.authKey, &overlayOpts.clientAuth),
-			sendOverlayMW(overlayOpts, &send, logger, &logf),
+			derpMap(&derpmapFi, dm),
+			sendOverlayMW(overlayOpts, &send, logger, dm, &logf),
 		),
 		Handler: func(inv *serpent.Invocation) error {
 			ctx := inv.Context()
 
-			s, err := tsserver.NewServer(ctx, logger, send)
+			s, err := tsserver.NewServer(ctx, logger, send, dm)
 			if err != nil {
 				return err
 			}
@@ -93,6 +97,12 @@ func sshCmd() *serpent.Command {
 				Default:     "",
 				Value:       serpent.StringOf(&overlayOpts.authKey),
 			},
+			{
+				Flag:        "derp-config-file",
+				Description: "File which specifies the DERP config to use. In the structure of https://pkg.go.dev/tailscale.com@v1.74.1/tailcfg#DERPMap.",
+				Default:     "",
+				Value:       serpent.StringOf(&derpmapFi),
+			},
 			{
 				Flag:    "stun-ip-override",
 				Default: "",

tsserver/server.go

@@ -64,12 +64,7 @@ func DERPMapTailscale(ctx context.Context) (*tailcfg.DERPMap, error) {
 	return dm, nil
 }
 
-func NewServer(ctx context.Context, logger *slog.Logger, ov overlay.Overlay) (*server, error) {
-	dm, err := DERPMapTailscale(ctx)
-	if err != nil {
-		return nil, err
-	}
-
+func NewServer(ctx context.Context, logger *slog.Logger, ov overlay.Overlay, dm *tailcfg.DERPMap) (*server, error) {
 	s := &server{
 		logger:          logger,
 		noisePrivateKey: key.NewMachine(),