refactor: improve resource management and cleanup patterns

- Simplify namespace bootstrapping by introducing `prepareLauncher` helper
- Implement a cleanup stack in `StartTunnel` to ensure orderly resource release on error
- Streamline temporary file and mount lifecycles in `ConfigureResolvConf` and `BlockHostServices`
- Ensure `Tunnel.Close()` also closes the underlying TUN device
- Reduce redundant manual cleanup calls using defer-based error handling

2 files changed, 63 insertions, 116 deletions

diff --git a/internal/namespace/namespace.go b/internal/namespace/namespace.go
index 0f2618b..6f56a84 100644
--- a/internal/namespace/namespace.go
+++ b/internal/namespace/namespace.go
@@ -84,7 +84,6 @@ func Bootstrap() (err error) {
 	}()
 
 	// 0. Validate current arguments for null bytes before proceeding.
-	// If any argument contains a null byte, syscall.Exec will fail with 'invalid argument'.
 	for i, arg := range os.Args {
 		for j := 0; j < len(arg); j++ {
 			if arg[j] == 0 {
@@ -98,57 +97,22 @@ func Bootstrap() (err error) {
 		return fmt.Errorf("failed to get executable path: %w", err)
 	}
 
-	// 1. Create a secure temporary file for the launcher binary.
-	// os.CreateTemp ensures a unique, unpredictable filename and restrictive permissions.
-	tmpFile, err := os.CreateTemp("", "wg-wrap-launcher-")
-	if err != nil {
-		return fmt.Errorf("failed to create temp launcher file: %w", err)
-	}
-	launcherPath := tmpFile.Name()
-
-	// 2. Write the embedded launcher binary to the temp file.
-	if _, err := tmpFile.Write(launcherBytes); err != nil {
-		_ = tmpFile.Close()
-		_ = os.Remove(launcherPath)
-		return fmt.Errorf("failed to write launcher binary: %w", err)
-	}
-
-	// Ensure the binary is executable (0700)
-	if err := tmpFile.Chmod(0700); err != nil {
-		_ = tmpFile.Close()
-		_ = os.Remove(launcherPath)
-		return fmt.Errorf("failed to set launcher permissions: %w", err)
-	}
-
-	// 2b. Open a read-only fd of the launcher to exec
-	execFd, err := syscall.Open(launcherPath, syscall.O_RDONLY, 0)
+	execFd, launcherPath, err := prepareLauncher()
 	if err != nil {
-		_ = tmpFile.Close()
-		_ = os.Remove(launcherPath)
-		return fmt.Errorf("failed to open launcher for exec: %w", err)
+		return err
 	}
 	fdsToClose = append(fdsToClose, execFd)
+	_ = os.Remove(launcherPath) // Unlink early; fd remains valid
 
-	// Close the write file descriptor (to avoid ETXTBSY)
-	_ = tmpFile.Close()
-
-	// Unlink the file from disk (makes it invisible and ensures it is deleted on exit)
-	_ = os.Remove(launcherPath)
-
-	// Clear close-on-exec so it remains open across syscall.Exec
+	// Clear close-on-exec
 	if flags, err := unix.FcntlInt(uintptr(execFd), unix.F_GETFD, 0); err == nil {
 		_, _ = unix.FcntlInt(uintptr(execFd), unix.F_SETFD, flags&^unix.FD_CLOEXEC)
 	}
 
 	// 3. Prepare arguments for the launcher.
-	// The launcher expects: launcher <command_to_run> [args...]
 	args := []string{self}
 	args = append(args, os.Args[1:]...)
 
-	// 4. Replace the current process with the launcher.
-	// We must check for null bytes in the arguments here because syscall.Exec
-	// (which calls execve) will return 'invalid argument' (EINVAL) if any
-	// string in the argv array contains a null byte.
 	for i, arg := range args {
 		for j := 0; j < len(arg); j++ {
 			if arg[j] == 0 {
@@ -164,7 +128,7 @@ func Bootstrap() (err error) {
 	}
 	fdsToClose = append(fdsToClose, hostNetFd)
 
-	// Clear close-on-exec so it remains open across syscall.Exec
+	// Clear close-on-exec
 	if flags, err := unix.FcntlInt(uintptr(hostNetFd), unix.F_GETFD, 0); err == nil {
 		_, _ = unix.FcntlInt(uintptr(hostNetFd), unix.F_SETFD, flags&^unix.FD_CLOEXEC)
 	}
@@ -203,12 +167,13 @@ func BootstrapJoin(targetPid int) (err error) {
 
 	var fdsToClose []int
 	defer func() {
-		for _, fd := range fdsToClose {
-			_ = syscall.Close(fd)
+		if err != nil {
+			for _, fd := range fdsToClose {
+				_ = syscall.Close(fd)
+			}
 		}
 	}()
 
-	// Validate current arguments for null bytes before proceeding.
 	for i, arg := range os.Args {
 		for j := 0; j < len(arg); j++ {
 			if arg[j] == 0 {
@@ -222,48 +187,17 @@ func BootstrapJoin(targetPid int) (err error) {
 		return fmt.Errorf("failed to get executable path: %w", err)
 	}
 
-	// 1. Create a secure temporary file for the launcher binary.
-	tmpFile, err := os.CreateTemp("", "wg-wrap-launcher-")
+	execFd, launcherPath, err := prepareLauncher()
 	if err != nil {
-		return fmt.Errorf("failed to create temp launcher file: %w", err)
-	}
-	launcherPath := tmpFile.Name()
-
-	// 2. Write the embedded launcher binary to the temp file.
-	if _, err := tmpFile.Write(launcherBytes); err != nil {
-		_ = tmpFile.Close()
-		_ = os.Remove(launcherPath)
-		return fmt.Errorf("failed to write launcher binary: %w", err)
-	}
-
-	// Ensure the binary is executable (0700)
-	if err := tmpFile.Chmod(0700); err != nil {
-		_ = tmpFile.Close()
-		_ = os.Remove(launcherPath)
-		return fmt.Errorf("failed to set launcher permissions: %w", err)
-	}
-
-	// 2b. Open a read-only fd of the launcher to exec
-	execFd, err := syscall.Open(launcherPath, syscall.O_RDONLY, 0)
-	if err != nil {
-		_ = tmpFile.Close()
-		_ = os.Remove(launcherPath)
-		return fmt.Errorf("failed to open launcher for exec: %w", err)
+		return err
 	}
 	fdsToClose = append(fdsToClose, execFd)
-
-	// Close the write file descriptor (to avoid ETXTBSY)
-	_ = tmpFile.Close()
-
-	// Unlink the file from disk (makes it invisible and ensures it is deleted on exit)
 	_ = os.Remove(launcherPath)
 
-	// Clear close-on-exec so it remains open across syscall.Exec
 	if flags, err := unix.FcntlInt(uintptr(execFd), unix.F_GETFD, 0); err == nil {
 		_, _ = unix.FcntlInt(uintptr(execFd), unix.F_SETFD, flags&^unix.FD_CLOEXEC)
 	}
 
-	// 3. Prepare arguments for the launcher.
 	args := []string{self}
 	args = append(args, os.Args[1:]...)
 
@@ -275,8 +209,6 @@ func BootstrapJoin(targetPid int) (err error) {
 		}
 	}
 
-	// Set environment variables to tell the C launcher to join,
-	// and to tell the second wg-wrap instance that we are in a joined session.
 	env := append(os.Environ(),
 		fmt.Sprintf("WG_WRAP_JOIN_PID=%d", targetPid),
 		"WG_WRAP_JOINED=1",
@@ -289,3 +221,34 @@ func BootstrapJoin(targetPid int) (err error) {
 
 	return nil
 }
+
+func prepareLauncher() (int, string, error) {
+	tmpFile, err := os.CreateTemp("", "wg-wrap-launcher-")
+	if err != nil {
+		return 0, "", fmt.Errorf("failed to create temp launcher file: %w", err)
+	}
+	launcherPath := tmpFile.Name()
+
+	defer func() {
+		if err != nil {
+			_ = tmpFile.Close()
+			_ = os.Remove(launcherPath)
+		}
+	}()
+
+	if _, err = tmpFile.Write(launcherBytes); err != nil {
+		return 0, "", fmt.Errorf("failed to write launcher binary: %w", err)
+	}
+
+	if err = tmpFile.Chmod(0700); err != nil {
+		return 0, "", fmt.Errorf("failed to set launcher permissions: %w", err)
+	}
+
+	execFd, err := syscall.Open(launcherPath, syscall.O_RDONLY, 0)
+	if err != nil {
+		return 0, "", fmt.Errorf("failed to open launcher for exec: %w", err)
+	}
+
+	_ = tmpFile.Close()
+	return execFd, launcherPath, nil
+}
diff --git a/internal/wireguard/wireguard.go b/internal/wireguard/wireguard.go
index 3f17392..3a2bfa3 100644
--- a/internal/wireguard/wireguard.go
+++ b/internal/wireguard/wireguard.go
@@ -28,21 +28,24 @@ type Tunnel struct {
 }
 
 // StartTunnel creates a TUN device, launches wireguard-go over it, and configures IPs/routes.
-func StartTunnel(cfg *wgconf.Config, dnsServer string) (*Tunnel, error) {
+func StartTunnel(cfg *wgconf.Config, dnsServer string) (t *Tunnel, err error) {
+	var cleanups []func()
+	defer func() {
+		if err != nil {
+			for i := len(cleanups) - 1; i >= 0; i-- {
+				cleanups[i]()
+			}
+		}
+	}()
+
 	// 1. Create the TUN device inside the current (isolated) namespace
-	// We use the default name 'tun0'
 	tunName := "tun0"
 	mtu := 1420
 
-	// Ensure the mount namespace is private to prevent mount propagation to the host.
-	// This is critical for the bind-mount of /etc/resolv.conf to work in rootless environments.
 	if err := unix.Mount("", "/", "", unix.MS_REC|unix.MS_PRIVATE, ""); err != nil {
-		// We log this as a warning because some environments might not allow this,
-		// but we can still try to proceed.
 		fmt.Printf("warning: failed to make mount namespace private: %v\n", err)
 	}
 
-	// Block host services (D-Bus, nscd) to prevent name resolution leak bypasses
 	if err := BlockHostServices(); err != nil {
 		fmt.Printf("warning: failed to block host services: %v\n", err)
 	}
@@ -51,12 +54,12 @@ func StartTunnel(cfg *wgconf.Config, dnsServer string) (*Tunnel, error) {
 	if err != nil {
 		return nil, fmt.Errorf("failed to create TUN device %s: %w", tunName, err)
 	}
+	cleanups = append(cleanups, func() { tunDev.Close() })
 
 	// 2. Instantiate the userspace WireGuard device
 	logger := device.NewLogger(device.LogLevelSilent, "[wg-wrap] ")
 	var bind conn.Bind
 
-	// Check if a pre-opened host UDP socket file descriptor was passed first (Approach A - FD Passing)
 	if hostSocketFdStr := os.Getenv("WG_WRAP_HOST_SOCKET_FD"); hostSocketFdStr != "" {
 		if fd, err := strconv.Atoi(hostSocketFdStr); err == nil && fd > 0 {
 			if fdBind, err := NewFDBind(fd); err == nil {
@@ -65,7 +68,6 @@ func StartTunnel(cfg *wgconf.Config, dnsServer string) (*Tunnel, error) {
 		}
 	}
 
-	// Fallback to NewHostBind or standard Bind if no host socket was passed
 	if bind == nil {
 		bind = conn.NewDefaultBind()
 		if hostNetNSFdStr := os.Getenv("WG_WRAP_HOST_NETNS_FD"); hostNetNSFdStr != "" {
@@ -76,38 +78,28 @@ func StartTunnel(cfg *wgconf.Config, dnsServer string) (*Tunnel, error) {
 	}
 
 	wgDev := device.NewDevice(tunDev, bind, logger)
+	cleanups = append(cleanups, func() { wgDev.Close() })
 
 	// 3. Formulate the UAPI configuration string to configure peers/keys
 	uapiConf, err := buildUAPIConfig(cfg)
 	if err != nil {
-		wgDev.Close()
 		return nil, fmt.Errorf("failed to build UAPI config: %w", err)
 	}
 
-	// Apply configuration via UAPI (IpcSet)
 	if err := wgDev.IpcSet(uapiConf); err != nil {
-		wgDev.Close()
 		return nil, fmt.Errorf("failed to configure WireGuard device: %w", err)
 	}
 
-	// Enable device
 	if err := wgDev.Up(); err != nil {
-		wgDev.Close()
 		return nil, fmt.Errorf("failed to bring up WireGuard device: %w", err)
 	}
 
 	// 4. Configure network interface using standard Linux network commands (iproute2)
-	// Since we are mapped to root (UID 0) inside our isolated network namespace,
-	// we have complete control over local network interfaces without affecting the host.
 	if err := configureInterface(tunName, cfg.Address, mtu); err != nil {
-		wgDev.Close()
 		return nil, fmt.Errorf("failed to configure network interface %s: %w", tunName, err)
 	}
 
-	// Configure DNS resolver inside the namespace
 	if err := ConfigureResolvConf(dnsServer); err != nil {
-		// We treat DNS failure as a warning rather than a fatal error to allow
-		// the tunnel to function even if /etc/resolv.conf is read-only.
 		fmt.Printf("warning: failed to configure DNS resolver: %v\n", err)
 	}
 
@@ -235,33 +227,26 @@ func ConfigureResolvConf(dns string) error {
 		return nil
 	}
 
-	// To avoid modifying the host's /etc/resolv.conf, we use the private mount namespace.
 	tmpFile, err := os.CreateTemp("", "resolvconf")
 	if err != nil {
 		return fmt.Errorf("failed to create temp resolv.conf: %w", err)
 	}
-	defer func() { _ = tmpFile.Close() }()
+	launcherPath := tmpFile.Name()
+	defer func() {
+		_ = tmpFile.Close()
+		_ = os.Remove(launcherPath)
+	}()
 
 	content := fmt.Sprintf("nameserver %s\n", dns)
 	if _, err := tmpFile.WriteString(content); err != nil {
 		return fmt.Errorf("failed to write to temp resolv.conf: %w", err)
 	}
 
-	// 1. Bind-mount the temp file over /etc/resolv.conf
-	if err := unix.Mount(tmpFile.Name(), "/etc/resolv.conf", "", unix.MS_BIND, ""); err != nil {
-		_ = os.Remove(tmpFile.Name())
-		return fmt.Errorf("failed to bind-mount %s to /etc/resolv.conf: %w", tmpFile.Name(), err)
+	if err := unix.Mount(launcherPath, "/etc/resolv.conf", "", unix.MS_BIND, ""); err != nil {
+		return fmt.Errorf("failed to bind-mount %s to /etc/resolv.conf: %w", launcherPath, err)
 	}
 
-	// Unlink the temporary source file. Since /etc/resolv.conf is a bind mount,
-	// the kernel will keep the inode alive, but the file is removed from /tmp.
-	_ = os.Remove(tmpFile.Name())
-
-	// 2. Make the mount private to ensure it doesn't propagate back to the host
-	// and to satisfy kernel requirements for mount transitions in some environments.
-	// We do this by applying MS_PRIVATE in a separate mount call.
 	if err := unix.Mount("", "/etc/resolv.conf", "", unix.MS_PRIVATE, ""); err != nil {
-		// If MS_PRIVATE fails, we can log a warning but proceed since / is already private
 		fmt.Printf("warning: failed to make /etc/resolv.conf mount private: %v\n", err)
 	}
 
@@ -276,7 +261,7 @@ func BlockHostServices() error {
 	if err != nil {
 		return fmt.Errorf("failed to create temp dir: %w", err)
 	}
-	defer func() { _ = os.Remove(tmpDir) }()
+	defer os.RemoveAll(tmpDir)
 
 	tmpFile, err := os.CreateTemp("", "wg-wrap-block-file-")
 	if err != nil {
@@ -284,9 +269,8 @@ func BlockHostServices() error {
 	}
 	tmpFileName := tmpFile.Name()
 	_ = tmpFile.Close()
-	defer func() { _ = os.Remove(tmpFileName) }()
+	defer os.Remove(tmpFileName)
 
-	// Specific socket files and directories to block
 	pathsToBlock := []string{
 		"/run/dbus/system_bus_socket",
 		"/run/systemd/resolve/io.systemd.Resolve",