fix: serialize tunnel lifecycle operations

Author: gitlayzer

cmd/cleanup.go

@@ -21,22 +21,17 @@ var cleanupCmd = &cobra.Command{
 			return fmt.Errorf("--all cannot be used with a specific tunnel id")
 		}
 		if len(args) > 0 {
-			sess, err := findSession(args[0])
+			eligible, deleteFailed, err := cleanupTunnelWithLock(cmd, args[0], cleanupAll)
 			if err != nil {
-				return err
-			}
-			if !cleanupAll && !sessionCleanupEligible(*sess, time.Minute) {
-				return fmt.Errorf("tunnel %s is not stopped, expired, stale, or error; refusing cleanup without --all", sess.TunnelID)
-			}
-			ctx, cancel := context.WithTimeout(cmd.Context(), 30*time.Second)
-			defer cancel()
-			if err := cleanupSessionResources(ctx, *sess); err != nil {
-				return fmt.Errorf("cleanup tunnel %s: %w", sess.TunnelID, err)
+				if deleteFailed {
+					return err
+				}
+				return fmt.Errorf("cleanup tunnel %s: %w", args[0], err)
 			}
-			if err := session.Delete(sess.TunnelID); err != nil {
-				return fmt.Errorf("delete local session %s: %w", sess.TunnelID, err)
+			if !eligible {
+				return fmt.Errorf("tunnel %s is not stopped, expired, stale, or error; refusing cleanup without --all", args[0])
 			}
-			fmt.Printf("Cleanup complete. Removed tunnel %s and its remote resources.\n", sess.TunnelID)
+			fmt.Printf("Cleanup complete. Removed tunnel %s and its remote resources.\n", args[0])
 			return nil
 		}
 		sessions, err := session.List()
@@ -48,17 +43,15 @@ var cleanupCmd = &cobra.Command{
 			removed := 0
 			failed := 0
 			for _, sess := range sessions {
-				ctx, cancel := context.WithTimeout(cmd.Context(), 30*time.Second)
-				if err := cleanupSessionResources(ctx, sess); err != nil {
-					cancel()
+				_, deleteFailed, err := cleanupTunnelWithLock(cmd, sess.TunnelID, true)
+				if err != nil {
+					if deleteFailed {
+						return err
+					}
 					failed++
 					fmt.Fprintf(cmd.ErrOrStderr(), "[!] Skipped tunnel %s: %v\n", sess.TunnelID, err)
 					continue
 				}
-				cancel()
-				if err := session.Delete(sess.TunnelID); err != nil {
-					return fmt.Errorf("delete local session %s: %w", sess.TunnelID, err)
-				}
 				removed++
 			}
 
@@ -73,14 +66,11 @@ var cleanupCmd = &cobra.Command{
 		skipped := 0
 		failed := 0
 		for _, sess := range sessions {
-			if !sessionCleanupEligible(sess, time.Minute) {
-				skipped++
-				continue
-			}
-
-			ctx, cancel := context.WithTimeout(cmd.Context(), 30*time.Second)
-			if err := cleanupSessionResources(ctx, sess); err != nil {
-				cancel()
+			eligible, deleteFailed, err := cleanupTunnelWithLock(cmd, sess.TunnelID, false)
+			if err != nil {
+				if deleteFailed {
+					return err
+				}
 				failed++
 				if errors.Is(err, errMissingSessionKubeconfig) {
 					fmt.Fprintf(cmd.ErrOrStderr(), "[!] Skipped cleanup-eligible tunnel %s: %v\n", sess.TunnelID, err)
@@ -89,9 +79,9 @@ var cleanupCmd = &cobra.Command{
 				fmt.Fprintf(cmd.ErrOrStderr(), "[!] Failed to clean up cleanup-eligible tunnel %s: %v\n", sess.TunnelID, err)
 				continue
 			}
-			cancel()
-			if err := session.Delete(sess.TunnelID); err != nil {
-				return fmt.Errorf("delete local session %s: %w", sess.TunnelID, err)
+			if !eligible {
+				skipped++
+				continue
 			}
 			cleaned++
 		}
@@ -108,3 +98,29 @@ func init() {
 	rootCmd.AddCommand(cleanupCmd)
 	cleanupCmd.Flags().BoolVar(&cleanupAll, "all", false, "Force delete all locally tracked Sealtun tunnel resources and remove matching local session records")
 }
+
+func cleanupTunnelWithLock(cmd *cobra.Command, tunnelID string, force bool) (bool, bool, error) {
+	eligible := false
+	deleteFailed := false
+	err := withTunnelOperationLock(tunnelID, func() error {
+		sess, err := findSession(tunnelID)
+		if err != nil {
+			return err
+		}
+		if !force && !sessionCleanupEligible(*sess, time.Minute) {
+			return nil
+		}
+		eligible = true
+		ctx, cancel := context.WithTimeout(cmd.Context(), 30*time.Second)
+		defer cancel()
+		if err := cleanupSessionResources(ctx, *sess); err != nil {
+			return err
+		}
+		if err := session.Delete(sess.TunnelID); err != nil {
+			deleteFailed = true
+			return fmt.Errorf("delete local session %s: %w", sess.TunnelID, err)
+		}
+		return nil
+	})
+	return eligible, deleteFailed, err
+}

cmd/cleanup_test.go

@@ -2,6 +2,8 @@ package cmd
 
 import (
 	"context"
+	"errors"
+	"fmt"
 	"strings"
 	"testing"
 	"time"
@@ -77,3 +79,37 @@ func TestCleanupSpecificTunnelRejectsActiveSession(t *testing.T) {
 		t.Fatalf("expected active cleanup refusal, got %v", err)
 	}
 }
+
+func TestCleanupWaitsForTunnelOperationLock(t *testing.T) {
+	t.Setenv("SEALTUN_HOME", t.TempDir())
+	if err := session.Save(session.TunnelSession{
+		TunnelID:        "cleanlocked",
+		ConnectionState: session.ConnectionStateStopped,
+		CreatedAt:       time.Now().Format(time.RFC3339),
+	}); err != nil {
+		t.Fatalf("save session: %v", err)
+	}
+
+	releaseLock := holdTunnelOperationLock(t, "cleanlocked")
+	defer releaseLock()
+	previousCleanup := cleanupSessionResources
+	cleanupCalled := make(chan struct{}, 1)
+	want := fmt.Errorf("stop after lock")
+	cleanupSessionResources = func(context.Context, session.TunnelSession) error {
+		cleanupCalled <- struct{}{}
+		return want
+	}
+	t.Cleanup(func() { cleanupSessionResources = previousCleanup })
+
+	done := make(chan error, 1)
+	go func() {
+		cmd := *cleanupCmd
+		cmd.SetContext(context.Background())
+		done <- cmd.RunE(&cmd, []string{"cleanlocked"})
+	}()
+	assertOperationBlocked(t, cleanupCalled)
+	releaseLock()
+	if err := <-done; !errors.Is(err, want) {
+		t.Fatalf("cleanup error = %v, want %v", err, want)
+	}
+}

cmd/connect.go

@@ -6,7 +6,6 @@ import (
 	"fmt"
 	"os"
 	"os/signal"
-	"syscall"
 
 	"github.com/labring/sealtun/pkg/clusterconnect"
 	"github.com/spf13/cobra"
@@ -37,6 +36,9 @@ var (
 	connectSaveState          = clusterconnect.SaveState
 	connectRemoveState        = clusterconnect.RemoveState
 	connectAcquireRuntimeLock = clusterconnect.AcquireRuntimeLock
+	connectNotifyContext      = signal.NotifyContext
+	connectStopStateProcess   = clusterconnect.StopStateProcess
+	connectCleanupState       = clusterconnect.CleanupTransparentState
 )
 
 var connectCmd = &cobra.Command{
@@ -99,6 +101,10 @@ func runConnectCheck(cmd *cobra.Command, opts connectOptions) error {
 	if err != nil {
 		return err
 	}
+	return runConnectCheckWithEnvironment(cmd, opts, env)
+}
+
+func runConnectCheckWithEnvironment(cmd *cobra.Command, opts connectOptions, env connectPreflighter) error {
 	ctx := connectCommandContext(cmd)
 	preflight, err := env.Preflight(ctx, clusterconnect.Options{
 		Mode:      opts.Mode,
@@ -107,16 +113,12 @@ func runConnectCheck(cmd *cobra.Command, opts connectOptions) error {
 	if opts.JSON {
 		enc := json.NewEncoder(cmd.OutOrStdout())
 		enc.SetIndent("", "  ")
-		_ = enc.Encode(preflight)
-		if preflight != nil {
-			return nil
+		if encodeErr := enc.Encode(preflight); encodeErr != nil {
+			return encodeErr
 		}
 		return err
 	}
 	printConnectPreflight(cmd, preflight)
-	if preflight != nil {
-		return nil
-	}
 	return err
 }
 
@@ -183,7 +185,7 @@ func runConnectWithEnvironment(cmd *cobra.Command, opts connectOptions, env conn
 	}
 	defer connectRemoveState()
 
-	connectCtx, stopSignals := signal.NotifyContext(ctx, os.Interrupt, syscall.SIGTERM)
+	connectCtx, stopSignals := connectNotifyContext(ctx, signalCleanupSignals()...)
 	defer stopSignals()
 	err = server.RunPlan(connectCtx, plan)
 	if err == nil || err == context.Canceled {
@@ -254,17 +256,17 @@ func runDisconnect(cmd *cobra.Command) error {
 		}
 		return err
 	}
-	stopErr := clusterconnect.StopStateProcess(*state)
+	stopErr := connectStopStateProcess(*state)
+	if stopErr != nil {
+		return stopErr
+	}
 	plan := &clusterconnect.TransparentPlan{
 		Namespace: state.Namespace,
 		Listen:    state.Listen,
 		Rules:     state.Rules,
 		Hosts:     state.Hosts,
 	}
-	cleanupErr := clusterconnect.CleanupTransparentState(plan)
-	if stopErr != nil {
-		return stopErr
-	}
+	cleanupErr := connectCleanupState(plan)
 	if cleanupErr != nil {
 		return cleanupErr
 	}

cmd/connect_test.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"errors"
 	"io"
+	"os"
 	"reflect"
 	"testing"
 
@@ -131,6 +132,96 @@ func TestRunConnectRejectsConcurrentRuntime(t *testing.T) {
 	}
 }
 
+func TestRunConnectCheckReturnsPreflightErrorWithPayload(t *testing.T) {
+	want := errors.New("no compatible connect mode")
+	for _, asJSON := range []bool{false, true} {
+		t.Run(map[bool]string{false: "text", true: "json"}[asJSON], func(t *testing.T) {
+			env := &fakeConnectEnv{
+				preflight: &clusterconnect.Preflight{Mode: clusterconnect.ModeAuto},
+				err:       want,
+			}
+			err := runConnectCheckWithEnvironment(newTestConnectCommand(), connectOptions{JSON: asJSON}, env)
+			if !errors.Is(err, want) {
+				t.Fatalf("runConnectCheckWithEnvironment error = %v, want %v", err, want)
+			}
+		})
+	}
+}
+
+func TestRunConnectUsesCleanupSignalSet(t *testing.T) {
+	stubConnectRuntimeLock(t)
+	previousNotify := connectNotifyContext
+	previousSave := connectSaveState
+	previousRemove := connectRemoveState
+	t.Cleanup(func() {
+		connectNotifyContext = previousNotify
+		connectSaveState = previousSave
+		connectRemoveState = previousRemove
+	})
+
+	var gotSignals []os.Signal
+	connectNotifyContext = func(parent context.Context, signals ...os.Signal) (context.Context, context.CancelFunc) {
+		gotSignals = append([]os.Signal(nil), signals...)
+		return context.WithCancel(parent)
+	}
+	connectSaveState = func(clusterconnect.State) error { return nil }
+	connectRemoveState = func() error { return nil }
+
+	events := []string{}
+	server := &fakeConnectServer{
+		plan:   &clusterconnect.TransparentPlan{},
+		events: &events,
+	}
+	env := &fakeConnectEnv{preflight: &clusterconnect.Preflight{SelectedMode: clusterconnect.ModeTun}}
+	if err := runConnectWithEnvironment(newTestConnectCommand(), connectOptions{}, env, func(clusterconnect.TransparentOptions) connectPlanRunner {
+		return server
+	}); err != nil {
+		t.Fatalf("runConnectWithEnvironment returned error: %v", err)
+	}
+	if want := signalCleanupSignals(); !reflect.DeepEqual(gotSignals, want) {
+		t.Fatalf("cleanup signals = %v, want %v", gotSignals, want)
+	}
+}
+
+func TestRunDisconnectDoesNotCleanupWhenStopFails(t *testing.T) {
+	t.Setenv("SEALTUN_HOME", t.TempDir())
+	if err := clusterconnect.SaveState(clusterconnect.State{Mode: clusterconnect.ModeTun}); err != nil {
+		t.Fatalf("save connect state: %v", err)
+	}
+
+	previousStop := connectStopStateProcess
+	previousCleanup := connectCleanupState
+	previousRemove := connectRemoveState
+	t.Cleanup(func() {
+		connectStopStateProcess = previousStop
+		connectCleanupState = previousCleanup
+		connectRemoveState = previousRemove
+	})
+	want := errors.New("connect process did not stop")
+	connectStopStateProcess = func(clusterconnect.State) error { return want }
+	cleanupCalled := false
+	connectCleanupState = func(*clusterconnect.TransparentPlan) error {
+		cleanupCalled = true
+		return nil
+	}
+	removeCalled := false
+	connectRemoveState = func() error {
+		removeCalled = true
+		return nil
+	}
+
+	err := runDisconnect(newTestConnectCommand())
+	if !errors.Is(err, want) {
+		t.Fatalf("runDisconnect error = %v, want %v", err, want)
+	}
+	if cleanupCalled {
+		t.Fatal("transparent state was cleaned before the connect process stopped")
+	}
+	if removeCalled {
+		t.Fatal("connect state was removed after the process failed to stop")
+	}
+}
+
 func stubConnectRuntimeLock(t *testing.T) {
 	t.Helper()
 	previous := connectAcquireRuntimeLock