Fix PID lifecycle race and improve CLI routing for diagnostic commands

5 files changed, 202 insertions, 51 deletions

diff --git a/Makefile b/Makefile
index 3dcd788..dc57dfa 100644
--- a/Makefile
+++ b/Makefile
@@ -19,16 +19,16 @@ FUZZ_TIME ?= 30s
 all: $(BINARY)
 
 # Build the Go binary
+# Use a temporary build file to ensure we don't leave a partially-built or stale binary
 $(BINARY): $(LAUNCHER_BIN)
-	CGO_ENABLED=$(CGO_ENABLED) go build $(GO_BUILD_FLAGS) -o $(BINARY) cmd/wg-wrap/main.go
+	CGO_ENABLED=$(CGO_ENABLED) go build $(GO_BUILD_FLAGS) -o $(BINARY).tmp cmd/wg-wrap/main.go && mv $(BINARY).tmp $(BINARY)
 
 # Build the embedded C launcher binary
 $(LAUNCHER_BIN): $(LAUNCHER_SRC)
 	$(CC) $(CFLAGS) $(LAUNCHER_SRC) -o $(LAUNCHER_BIN)
 
 # Run tests
-test: clean
-	$(MAKE) $(BINARY)
+test: $(BINARY)
 	@echo "Running tests with WG_WRAP_BIN=$(shell pwd)/$(BINARY)"
 	WG_WRAP_BIN=$(shell pwd)/$(BINARY) go test -v -race ./...
 
diff --git a/internal/cli/cli.go b/internal/cli/cli.go
index aa4268a..13a4a6b 100644
--- a/internal/cli/cli.go
+++ b/internal/cli/cli.go
@@ -21,7 +21,7 @@ func NewApp(args []string) *App {
 	return &App{Args: args}
 }
 
-func (a *App) Run() error {
+func (a *App) Route() error {
 	// 1. Validate arguments for null bytes to prevent exec failures in the C launcher
 	for i, arg := range a.Args {
 		for j := 0; j < len(arg); j++ {
@@ -31,10 +31,34 @@ func (a *App) Run() error {
 		}
 	}
 
-	// Handle the internal diagnostic commands first
+	// Handle the internal diagnostic commands that should run on the HOST
+	if len(a.Args) > 1 {
+		switch a.Args[1] {
+		case "show-config":
+			return a.showConfig()
+		}
+	}
+
+	// Handle subcommands first (profile list, import, configure, delete, stop)
+	if len(a.Args) > 1 && a.Args[1] == "profile" {
+		return a.handleProfileCmd()
+	}
+
+	// If we reach here, we are either wrapping a process or running a command
+	// that requires isolation (e.g., test-ns, test-args).
+	return a.Run()
+}
+
+func (a *App) Run() error {
+	// Handle the internal diagnostic commands that require ISOLATION
 	if len(a.Args) > 1 {
 		switch a.Args[1] {
 		case "test-ns":
+			if !namespace.IsIsolated() {
+				if err := namespace.Bootstrap(); err != nil {
+					return fmt.Errorf("bootstrap failed: %w", err)
+				}
+			}
 			ok, msg := namespace.VerifyIsolation()
 			if !ok {
 				return fmt.Errorf("isolation check failed: %s", msg)
@@ -42,16 +66,15 @@ func (a *App) Run() error {
 			fmt.Println("Isolation Verified: OK")
 			return nil
 		case "test-args":
+			if !namespace.IsIsolated() {
+				if err := namespace.Bootstrap(); err != nil {
+					return fmt.Errorf("bootstrap failed: %w", err)
+				}
+			}
 			return namespace.VerifyArguments(a.Args)
 		}
 	}
 
-	// Handle subcommands first (profile list, import, configure, delete, stop)
-	if len(a.Args) > 1 && a.Args[1] == "profile" {
-		return a.handleProfileCmd()
-	}
-	// ...
-
 	cfg := &config.Config{}
 
 	fs := flag.NewFlagSet("wg-wrap", flag.ExitOnError)
@@ -194,3 +217,39 @@ func (a *App) handleProfileCmd() error {
 		return fmt.Errorf("unknown profile subcommand: %s", subCmd)
 	}
 }
+
+func (a *App) showConfig() error {
+	cfg := &config.Config{}
+	fs := flag.NewFlagSet("wg-wrap", flag.ExitOnError)
+	fs.StringVar(&cfg.Profile, "profile", "default", "WireGuard profile to use")
+	fs.StringVar(&cfg.DNSServer, "dns-server", "", "Override DNS server to use")
+
+	// Parse the arguments that follow 'show-config'
+	if len(a.Args) > 2 {
+		_ = fs.Parse(a.Args[2:])
+	}
+
+	// Determine runtime base directory
+	runtimeBase := a.RuntimeBaseDir
+	if runtimeBase == "" {
+		runtimeBase = os.Getenv("XDG_RUNTIME_DIR")
+		if runtimeBase == "" {
+			runtimeBase = fmt.Sprintf("/run/user/%d", os.Getuid())
+		}
+	}
+
+	// Resolve paths
+	profilePath := namespace.GetProfileNamespacePath(runtimeBase, cfg.Profile)
+	pidsPath := namespace.GetPidsDirPath(runtimeBase, cfg.Profile)
+
+	fmt.Printf("Configuration:\n")
+	fmt.Printf("  Profile:       %s\n", cfg.Profile)
+	fmt.Printf("  DNS Server:    %s\n", cfg.DNSServer)
+	fmt.Printf("  Runtime Base:  %s\n", runtimeBase)
+	fmt.Printf("  Profile Path:  %s\n", profilePath)
+	fmt.Printf("  PIDs Path:     %s\n", pidsPath)
+	fmt.Printf("  Isolated:      %v\n", namespace.IsIsolated())
+	fmt.Printf("  UID:           %d\n", os.Getuid())
+
+	return nil
+}
diff --git a/internal/namespace/launcher_src/launcher.c b/internal/namespace/launcher_src/launcher.c
index 63dd6ff..4311430 100644
--- a/internal/namespace/launcher_src/launcher.c
+++ b/internal/namespace/launcher_src/launcher.c
@@ -65,12 +65,15 @@ int main(int argc, char **argv) {
     // as the first element of the argv array.
     // Therefore, argv[0] is the path to the binary we want to execute.
     if (argv[0] == NULL) {
-        fprintf(stderr, "No target binary provided in argv[0]\n");
+        fprintf(stderr, "No target binary provided in argv[0]\\n");
         return 1;
     }
 
-    // Use execv instead of execvp to avoid PATH search issues
-    // since we are providing an absolute path from Go.
+    // Prepare a new argv for the target command.
+    // We want the target binary to see itself as argv[0].
+    // The current argv is [target_binary, arg1, arg2, ...].
+    // execv expects argv[0] to be the filename, and the rest as arguments.
+    // This is already the case here, but let's be explicit.
     if (execv(argv[0], argv) == -1) {
         perror("execv failed");
         return 1;
diff --git a/tests/e2e/config_test.go b/tests/e2e/config_test.go
new file mode 100644
index 0000000..83cfc15
--- /dev/null
+++ b/tests/e2e/config_test.go
@@ -0,0 +1,58 @@
+package e2e
+
+import (
+	"fmt"
+	"os"
+	"os/exec"
+	"strings"
+	"testing"
+)
+
+func TestConfigPropagation(t *testing.T) {
+	binaryPath, err := GetBinaryPath()
+	if err != nil {
+		t.Skipf("Skipping test: %v", err)
+	}
+
+	tmpRuntimeDir := t.TempDir()
+	profile := "config-test-vpn"
+
+	// Test 1: Non-isolated configuration
+	cmd := exec.Command(binaryPath, "show-config", "--profile", profile)
+	cmd.Env = append(os.Environ(), fmt.Sprintf("XDG_RUNTIME_DIR=%s", tmpRuntimeDir))
+	out, err := cmd.CombinedOutput()
+	if err != nil {
+		t.Fatalf("show-config failed: %v\nOutput: %s", err, string(out))
+	}
+
+	output := string(out)
+	expectedBase := tmpRuntimeDir
+	expectedPids := fmt.Sprintf("%s/profiles/%s/pids", tmpRuntimeDir, profile)
+
+	if !strings.Contains(output, fmt.Sprintf("Runtime Base:  %s", expectedBase)) {
+		t.Errorf("Expected Runtime Base %s in output: %s", expectedBase, output)
+	}
+	if !strings.Contains(output, fmt.Sprintf("PIDs Path:     %s", expectedPids)) {
+		t.Errorf("Expected PIDs Path %s in output: %s", expectedPids, output)
+	}
+
+	// Test 2: Configuration after bootstrap (Isolated)
+	// We use 'test-ns' as a way to run a command that we know is isolated.
+	// Actually, we can just run 'show-config' but the current 'Route' 
+	// handles 'show-config' BEFORE the bootstrap. 
+	// To test isolated config, we can't use 'show-config' because it's a diagnostic 
+	// command designed to run outside isolation.
+	
+	// To verify what an isolated process sees, we can use a target command 
+	// that prints the environment.
+	cmdIsolated := exec.Command(binaryPath, "--profile", profile, "--", "sh", "-c", "echo $XDG_RUNTIME_DIR")
+	cmdIsolated.Env = append(os.Environ(), fmt.Sprintf("XDG_RUNTIME_DIR=%s", tmpRuntimeDir))
+	outIso, err := cmdIsolated.CombinedOutput()
+	if err != nil {
+		t.Fatalf("Isolated command failed: %v\nOutput: %s", err, string(outIso))
+	}
+
+	if !strings.Contains(string(outIso), expectedBase) {
+		t.Errorf("Expected isolated process to see XDG_RUNTIME_DIR=%s, got: %s", expectedBase, string(outIso))
+	}
+}
diff --git a/tests/e2e/lifecycle_test.go b/tests/e2e/lifecycle_test.go
index baf9f56..649dbc0 100644
--- a/tests/e2e/lifecycle_test.go
+++ b/tests/e2e/lifecycle_test.go
@@ -9,6 +9,34 @@ import (
 	"time"
 )
 
+func waitForPids(t *testing.T, pidsDir string, expectedCount int) {
+	timeout := time.After(2 * time.Second)
+	tick := time.NewTicker(50 * time.Millisecond)
+	defer tick.Stop()
+
+	for {
+		select {
+		case <-timeout:
+			files, err := os.ReadDir(pidsDir)
+			if err != nil {
+				t.Fatalf("Failed to read pids dir during timeout: %v", err)
+			}
+			t.Fatalf("Timed out waiting for %d PID files, got %d", expectedCount, len(files))
+		case <-tick.C:
+			files, err := os.ReadDir(pidsDir)
+			if err != nil {
+				if os.IsNotExist(err) {
+					continue
+				}
+				t.Fatalf("Failed to read pids dir: %v", err)
+			}
+			if len(files) == expectedCount {
+				return
+			}
+		}
+	}
+}
+
 func TestNamespaceLifecycleAutomation(t *testing.T) {
 	// 1. Setup Environment
 	binaryPath, err := GetBinaryPath()
@@ -32,17 +60,8 @@ func TestNamespaceLifecycleAutomation(t *testing.T) {
 			t.Fatalf("Failed to start cmd1: %v", err)
 		}
 
-		// Allow a moment for the bootstrap loop to complete and register the PID
-		time.Sleep(500 * time.Millisecond)
-
-		// Verify PID file exists
-		files, err := os.ReadDir(pidsDir)
-		if err != nil {
-			t.Fatalf("Failed to read pids dir: %v", err)
-		}
-		if len(files) != 1 {
-			t.Errorf("Expected 1 PID file, got %d", len(files))
-		}
+		// Verify PID file exists using polling
+		waitForPids(t, pidsDir, 1)
 
 		// Start a second process using the same profile
 		cmd2 := exec.Command(binaryPath, "--profile", profile, "--", "sleep", "0.1")
@@ -50,43 +69,55 @@ func TestNamespaceLifecycleAutomation(t *testing.T) {
 		if err := cmd2.Start(); err != nil {
 			t.Fatalf("Failed to start cmd2: %v", err)
 		}
-		time.Sleep(500 * time.Millisecond)
-
-		files, err = os.ReadDir(pidsDir)
-		if err != nil {
-			t.Fatalf("Failed to read pids dir: %v", err)
-		}
-		if len(files) != 2 {
-			t.Errorf("Expected 2 PID files, got %d", len(files))
-		}
+		waitForPids(t, pidsDir, 2)
 
 		// Wait for first process to exit naturally (triggering defer)
 		if err := cmd1.Wait(); err != nil {
 			t.Fatalf("cmd1 failed: %v", err)
 		}
-		time.Sleep(500 * time.Millisecond)
-
-		files, err = os.ReadDir(pidsDir)
-		if err != nil {
-			t.Fatalf("Failed to read pids dir: %v", err)
-		}
-		if len(files) != 1 {
-			t.Errorf("Expected 1 PID file after first exit, got %d", len(files))
+		
+		// Poll for the count to drop back to 1
+		timeout := time.After(2 * time.Second)
+		found := false
+		for !found {
+			select {
+			case <-timeout:
+				t.Fatalf("Timed out waiting for first process to unregister")
+			default:
+				files, err := os.ReadDir(pidsDir)
+				if err == nil && len(files) == 1 {
+					found = true
+					break
+				}
+				time.Sleep(50 * time.Millisecond)
+			}
 		}
 
 		// Wait for second process to exit naturally
 		if err := cmd2.Wait(); err != nil {
 			t.Fatalf("cmd2 failed: %v", err)
 		}
-		time.Sleep(500 * time.Millisecond)
-
-		// Verify a clean state
-		files, err = os.ReadDir(pidsDir)
-		if err != nil && !os.IsNotExist(err) {
-			t.Fatalf("Failed to read pids dir: %v", err)
-		}
-		if err == nil && len(files) != 0 {
-			t.Errorf("Expected 0 PID files after all exits, got %d", len(files))
+		
+		// Verify a clean state (expect 0 files)
+		timeout = time.After(2 * time.Second)
+		found = false
+		for !found {
+			select {
+			case <-timeout:
+				files, _ := os.ReadDir(pidsDir)
+				t.Fatalf("Expected 0 PID files after all exits, got %d", len(files))
+			default:
+				files, err := os.ReadDir(pidsDir)
+				if err != nil && os.IsNotExist(err) {
+					found = true
+					break
+				}
+				if err == nil && len(files) == 0 {
+					found = true
+					break
+				}
+				time.Sleep(50 * time.Millisecond)
+			}
 		}
 	})
 }