Implement automatic namespace lifecycle cleanup with last-man-out reference counting

1 files changed, 60 insertions, 9 deletions

diff --git a/internal/cli/cli.go b/internal/cli/cli.go
index eba7f68..f88a623 100644
--- a/internal/cli/cli.go
+++ b/internal/cli/cli.go
@@ -3,14 +3,18 @@ package cli
 import (
 	"flag"
 	"fmt"
+	"os"
+	"os/exec"
 
 	"git.theodohertyfamily.com/tools/wg-wrap/internal/config"
 	"git.theodohertyfamily.com/tools/wg-wrap/internal/namespace"
 )
 
+
 type App struct {
-	Args      []string
-	ConfigDir string // Optional override for profile storage location
+	Args           []string
+	ConfigDir      string // Optional override for profile storage location
+	RuntimeBaseDir string // Optional override for namespace/PID tracking
 }
 
 func NewApp(args []string) *App {
@@ -88,15 +92,62 @@ func (a *App) Run() error {
 		cfg.Profile = "default"
 	}
 
-	profilesDir := a.ConfigDir
-	if profilesDir == "" {
-		profilesDir = config.GetDefaultProfilesDir()
+	if namespace.IsIsolated() {
+		// Inject runtime base dir if provided
+		if a.RuntimeBaseDir != "" {
+			namespace.SetRuntimeBaseDir(a.RuntimeBaseDir)
+		}
+		return a.ExecuteCommand(cfg)
+	}
+
+	// If we are not isolated, we bootstrap.
+	// The Bootstrap process will replace this process and restart it.
+	if err := namespace.Bootstrap(); err != nil {
+		return fmt.Errorf("bootstrap failed: %w", err)
+	}
+
+	// This point is never reached because Bootstrap uses syscall.Exec
+	return nil
+}
+
+// ExecuteCommand handles the isolated execution of the target application.
+// This is called after the bootstrap loop has successfully isolated the process.
+func (a *App) ExecuteCommand(cfg *config.Config) error {
+	if !namespace.IsIsolated() {
+		return fmt.Errorf("ExecuteCommand called without namespace isolation")
+	}
+
+	// 1. Prepare the namespace
+	namespace.PruneStalePids(cfg.Profile)
+	if err := namespace.RegisterProcess(cfg.Profile); err != nil {
+		return fmt.Errorf("failed to register process: %w", err)
+	}
+
+	// Ensure we unregister and check for cleanup on exit
+	defer func() {
+		namespace.UnregisterProcess(cfg.Profile)
+		if last, err := namespace.IsLastProcess(cfg.Profile); err == nil && last {
+			fmt.Printf("Last process exiting. Cleaning up profile %s...\n", cfg.Profile)
+			// Here we would call namespace.UnpinNamespace(cfg.Profile) 
+			// and terminate the userspace WG process.
+		}
+	}()
+
+	// 2. VPN Setup (Stubbed)
+	fmt.Printf("Initializing WireGuard tunnel for profile %s...\n", cfg.Profile)
+	// TODO: Integrate with internal/wireguard to set up TUN and WG-Go
+
+	// 3. Execute the target command
+	cmd := exec.Command(cfg.Command[0], cfg.Command[1:]...)
+	cmd.Stdin = os.Stdin
+	cmd.Stdout = os.Stdout
+	cmd.Stderr = os.Stderr
+	cmd.Env = os.Environ()
+
+	if err := cmd.Run(); err != nil {
+		return fmt.Errorf("command execution failed: %w", err)
 	}
 
-	fmt.Printf("Profile: %s\n", cfg.Profile)
-	fmt.Printf("Profiles Directory: %s\n", profilesDir)
-	fmt.Printf("DNS Server: %s\n", cfg.DNSServer)
-	fmt.Printf("Command: %v\n", cfg.Command)
 	return nil
 }