cmd/tailscale/cli: Add clear subcommand to servekevin/add_services_flag_to_tailscale_serve_and_enrich_output

This commit adds a subcommand for serve to remove all configs for a service.
This include all proxies added for the service in the serve config's services section.

Updates tailscale/corp#22954

Signed-off-by: KevinLiang10 <37811973+KevinLiang10@users.noreply.github.com>

2 files changed, 101 insertions, 5 deletions

diff --git a/cmd/tailscale/cli/serve_legacy.go b/cmd/tailscale/cli/serve_legacy.go
index 1789ad93e..eac17af6f 100644
--- a/cmd/tailscale/cli/serve_legacy.go
+++ b/cmd/tailscale/cli/serve_legacy.go
@@ -754,6 +754,46 @@ func (e *serveEnv) runServeReset(ctx context.Context, args []string) error {
 	return e.lc.SetServeConfig(ctx, sc)
 }
 
+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]
+	err := tailcfg.ServiceName(svc).Validate()
+	if err != nil {
+		return fmt.Errorf("failed to parse service name: %w", err)
+	}
+	return e.removeServiceFromPrefs(ctx, svc)
+}
+
+func (e *serveEnv) runServeClear(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]
+	err := tailcfg.ServiceName(svc).Validate()
+	if err != nil {
+		return fmt.Errorf("failed to parse service name: %w", err)
+	}
+	sc, err := e.lc.GetServeConfig(ctx)
+	if err != nil {
+		return fmt.Errorf("error getting serve config: %w", err)
+	}
+	if _, ok := sc.Services[tailcfg.ServiceName(svc)]; !ok {
+		return fmt.Errorf("service %q not found in serve config", svc)
+	}
+	delete(sc.Services, tailcfg.ServiceName(svc))
+	return e.lc.SetServeConfig(ctx, sc)
+}
+
 // parseServePort parses a port number from a string and returns it as a
 // uint16. It returns an error if the port number is invalid or zero.
 func parseServePort(s string) (uint16, error) {
diff --git a/cmd/tailscale/cli/serve_v2.go b/cmd/tailscale/cli/serve_v2.go
index 83277891c..0fd147b52 100644
--- a/cmd/tailscale/cli/serve_v2.go
+++ b/cmd/tailscale/cli/serve_v2.go
@@ -177,6 +177,23 @@ 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 <service>", 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" +
+					"<service> should be a service name (e.g., svc:my-service).",
+				Exec: e.runServeDrain,
+			},
+			{
+				Name:       "clear",
+				ShortUsage: fmt.Sprintf("tailscale %s clear <service>", info.Name),
+				ShortHelp:  "Remove all config for a service",
+				LongHelp:   "Remove all handlers configured for the specified service.",
+				Exec:       e.runServeClear,
+			},
 		},
 	}
 }
@@ -275,11 +292,6 @@ func (e *serveEnv) runServeCombined(subcmd serveMode) execFunc {
 			return fmt.Errorf("error getting serve config: %w", err)
 		}
 
-		prefs, err := e.lc.GetPrefs(ctx)
-		if err != nil {
-			return fmt.Errorf("error getting prefs: %w", err)
-		}
-
 		// nil if no config
 		if sc == nil {
 			sc = new(ipn.ServeConfig)
@@ -400,6 +412,50 @@ func (e *serveEnv) runServeCombined(subcmd serveMode) execFunc {
 	}
 }
 
+func (e *serveEnv) addServiceToPrefs(ctx context.Context, serviceName string) error {
+	prefs, err := e.lc.GetPrefs(ctx)
+	if err != nil {
+		return fmt.Errorf("error getting prefs: %w", err)
+	}
+	advertisedServices := prefs.AdvertiseServices
+	if !slices.Contains(advertisedServices, serviceName) {
+		advertisedServices = append(advertisedServices, serviceName)
+	}
+	_, err = e.lc.EditPrefs(ctx, &ipn.MaskedPrefs{
+		AdvertiseServicesSet: true,
+		Prefs: ipn.Prefs{
+			AdvertiseServices: advertisedServices,
+		},
+	})
+	return err
+}
+
+func (e *serveEnv) removeServiceFromPrefs(ctx context.Context, serviceName string) 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
+	}
+	var advertisedServices []string
+	for _, svc := range prefs.AdvertiseServices {
+		if svc != serviceName {
+			advertisedServices = append(advertisedServices, svc)
+		}
+	}
+	if len(advertisedServices) == len(prefs.AdvertiseServices) {
+		return fmt.Errorf("service %q was not advertised", serviceName)
+	}
+	_, err = e.lc.EditPrefs(ctx, &ipn.MaskedPrefs{
+		AdvertiseServicesSet: true,
+		Prefs: ipn.Prefs{
+			AdvertiseServices: advertisedServices,
+		},
+	})
+	return err
+}
+
 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.