diff --git a/cmd/tailscale/cli/serve_v2.go b/cmd/tailscale/cli/serve_v2.go index 15de0609c..6fa1a1b08 100644 --- a/cmd/tailscale/cli/serve_v2.go +++ b/cmd/tailscale/cli/serve_v2.go @@ -203,6 +203,16 @@ func newServeV2Command(e *serveEnv, subcmd serveMode) *ffcli.Command { Exec: e.runServeReset, FlagSet: e.newFlags("serve-reset", nil), }, + { + Name: "drain", + ShortUsage: fmt.Sprintf("tailscale %s drain ", info.Name), + ShortHelp: "Drain a service from the current node", + LongHelp: "Make the current node no longer accept new connections for the specified service.\n" + + "Existing connections will continue to work until they are closed, but no new connections will be accepted.\n" + + "Use this command to gracefully remove a service from the current node without disrupting existing connections.\n" + + " should be a service name (e.g., svc:my-service).", + Exec: e.runServeDrain, + }, }, } } @@ -443,6 +453,44 @@ func (e *serveEnv) addServiceToPrefs(ctx context.Context, serviceName string) er return err } +func (e *serveEnv) removeServiceFromPrefs(ctx context.Context, serviceName tailcfg.ServiceName) error { + prefs, err := e.lc.GetPrefs(ctx) + if err != nil { + return fmt.Errorf("error getting prefs: %w", err) + } + if len(prefs.AdvertiseServices) == 0 { + return nil // nothing to remove + } + initialLen := len(prefs.AdvertiseServices) + prefs.AdvertiseServices = slices.DeleteFunc(prefs.AdvertiseServices, func(s string) bool { return s == serviceName.String() }) + if initialLen == len(prefs.AdvertiseServices) { + return nil // serviceName not advertised + } + _, err = e.lc.EditPrefs(ctx, &ipn.MaskedPrefs{ + AdvertiseServicesSet: true, + Prefs: ipn.Prefs{ + AdvertiseServices: prefs.AdvertiseServices, + }, + }) + return err +} + +func (e *serveEnv) runServeDrain(ctx context.Context, args []string) error { + if len(args) == 0 { + return errHelp + } + if len(args) != 1 { + fmt.Fprintf(Stderr, "error: invalid number of arguments\n\n") + return errHelp + } + svc := args[0] + svcName := tailcfg.ServiceName(svc) + if err := svcName.Validate(); err != nil { + return fmt.Errorf("invalid service name: %s", err) + } + return e.removeServiceFromPrefs(ctx, svcName) +} + const backgroundExistsMsg = "background configuration already exists, use `tailscale %s --%s=%d off` to remove the existing configuration" // validateConfig checks if the serve config is valid to serve the type wanted on the port. diff --git a/cmd/tailscale/cli/serve_v2_test.go b/cmd/tailscale/cli/serve_v2_test.go index b3e7ea773..2ba0b3f84 100644 --- a/cmd/tailscale/cli/serve_v2_test.go +++ b/cmd/tailscale/cli/serve_v2_test.go @@ -1212,6 +1212,54 @@ func TestAddServiceToPrefs(t *testing.T) { } +func TestRemoveServiceFromPrefs(t *testing.T) { + tests := []struct { + name string + svcName tailcfg.ServiceName + startServices []string + expected []string + }{ + { + name: "remove service from empty prefs", + svcName: "svc:foo", + expected: []string{}, + }, + { + name: "remove existing service from prefs", + svcName: "svc:foo", + startServices: []string{"svc:foo"}, + expected: []string{}, + }, + { + name: "remove service not in prefs", + svcName: "svc:bar", + startServices: []string{"svc:foo"}, + expected: []string{"svc:foo"}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + lc := &fakeLocalServeClient{} + ctx := t.Context() + lc.EditPrefs(ctx, &ipn.MaskedPrefs{ + AdvertiseServicesSet: true, + Prefs: ipn.Prefs{ + AdvertiseServices: tt.startServices, + }, + }) + e := &serveEnv{lc: lc, bg: bgBoolFlag{true, false}} + err := e.removeServiceFromPrefs(ctx, tt.svcName) + if err != nil { + t.Fatalf("removeServiceFromPrefs(%q) returned unexpected error: %v", tt.svcName, err) + } + if !slices.Equal(lc.prefs.AdvertiseServices, tt.expected) { + t.Errorf("removeServiceFromPrefs(%q) = %v, want %v", tt.svcName, lc.prefs.AdvertiseServices, tt.expected) + } + }) + } +} + func TestMessageForPort(t *testing.T) { svcIPMap := tailcfg.ServiceIPMappings{ "svc:foo": []netip.Addr{