Launch live browser opener as invoking sudo user

2 files changed, 180 insertions, 2 deletions

diff --git a/internal/flamegraph/liveserver.go b/internal/flamegraph/liveserver.go
index de65ee3..4cc5629 100644
--- a/internal/flamegraph/liveserver.go
+++ b/internal/flamegraph/liveserver.go
@@ -6,8 +6,13 @@ import (
 	"errors"
 	"fmt"
 	"net/http"
+	"os"
 	"os/exec"
+	"os/user"
+	"path/filepath"
+	"strconv"
 	"strings"
+	"syscall"
 	"time"
 )
 
@@ -19,6 +24,7 @@ var liveServerTimeouts = serverTimeouts{
 
 type LiveServerOptions struct {
 	OpenCommand string
+	WarningCb   func(message string)
 }
 
 var openBrowserURLFn = openBrowserURL
@@ -39,7 +45,7 @@ func ServeLiveWithOptions(ctx context.Context, lt *LiveTrie, interval time.Durat
 		url := fmt.Sprintf("http://%s:%d/", hostname, port)
 		fmt.Printf("Live flamegraph available at %s\n", url)
 		if err := maybeOpenLiveBrowser(url, options); err != nil {
-			fmt.Printf("Live flamegraph browser auto-open failed: %v\n", err)
+			notifyLiveWarning(options.WarningCb, fmt.Sprintf("Live flamegraph browser auto-open failed: %v", err))
 		}
 	})
 }
@@ -57,13 +63,110 @@ func openBrowserURL(url, openCommand string) error {
 		return err
 	}
 	cmd := exec.Command(parts[0], parts[1:]...)
+	applySudoInvokerContext(cmd)
 	if err := cmd.Start(); err != nil {
 		return err
 	}
-	go func() { _ = cmd.Wait() }()
+
+	waitCh := make(chan error, 1)
+	go func() { waitCh <- cmd.Wait() }()
+
+	select {
+	case waitErr := <-waitCh:
+		if waitErr != nil {
+			return fmt.Errorf("browser command exited early: %w", waitErr)
+		}
+	case <-time.After(750 * time.Millisecond):
+	}
 	return nil
 }
 
+func notifyLiveWarning(warningCb func(string), message string) {
+	if message == "" {
+		return
+	}
+	if warningCb != nil {
+		warningCb(message)
+		return
+	}
+	fmt.Println(message)
+}
+
+func applySudoInvokerContext(cmd *exec.Cmd) {
+	applySudoInvokerContextWithEnv(cmd, os.Geteuid(), os.Environ())
+}
+
+func applySudoInvokerContextWithEnv(cmd *exec.Cmd, euid int, env []string) {
+	if cmd == nil || euid != 0 {
+		return
+	}
+
+	sudoUIDStr, okUID := lookupEnvValue(env, "SUDO_UID")
+	sudoGIDStr, okGID := lookupEnvValue(env, "SUDO_GID")
+	if !okUID || !okGID {
+		return
+	}
+
+	uid, errUID := strconv.ParseUint(strings.TrimSpace(sudoUIDStr), 10, 32)
+	gid, errGID := strconv.ParseUint(strings.TrimSpace(sudoGIDStr), 10, 32)
+	if errUID != nil || errGID != nil {
+		return
+	}
+
+	cmd.SysProcAttr = &syscall.SysProcAttr{
+		Credential: &syscall.Credential{
+			Uid: uint32(uid),
+			Gid: uint32(gid),
+		},
+	}
+
+	launchEnv := append([]string(nil), env...)
+	if sudoUser, ok := lookupEnvValue(env, "SUDO_USER"); ok && strings.TrimSpace(sudoUser) != "" {
+		launchEnv = upsertEnvValue(launchEnv, "USER", sudoUser)
+		launchEnv = upsertEnvValue(launchEnv, "LOGNAME", sudoUser)
+	}
+
+	if sudoUser, err := user.LookupId(strconv.FormatUint(uid, 10)); err == nil && strings.TrimSpace(sudoUser.HomeDir) != "" {
+		launchEnv = upsertEnvValue(launchEnv, "HOME", sudoUser.HomeDir)
+		if _, ok := lookupEnvValue(launchEnv, "XAUTHORITY"); !ok {
+			xauth := filepath.Join(sudoUser.HomeDir, ".Xauthority")
+			if info, statErr := os.Stat(xauth); statErr == nil && !info.IsDir() {
+				launchEnv = upsertEnvValue(launchEnv, "XAUTHORITY", xauth)
+			}
+		}
+	}
+
+	if _, ok := lookupEnvValue(launchEnv, "XDG_RUNTIME_DIR"); !ok {
+		runtimeDir := fmt.Sprintf("/run/user/%d", uid)
+		if info, statErr := os.Stat(runtimeDir); statErr == nil && info.IsDir() {
+			launchEnv = upsertEnvValue(launchEnv, "XDG_RUNTIME_DIR", runtimeDir)
+		}
+	}
+
+	cmd.Env = launchEnv
+}
+
+func lookupEnvValue(env []string, key string) (string, bool) {
+	prefix := key + "="
+	for _, entry := range env {
+		if strings.HasPrefix(entry, prefix) {
+			return strings.TrimPrefix(entry, prefix), true
+		}
+	}
+	return "", false
+}
+
+func upsertEnvValue(env []string, key, value string) []string {
+	prefix := key + "="
+	for i := range env {
+		if strings.HasPrefix(env[i], prefix) {
+			env[i] = prefix + value
+			return env
+		}
+	}
+	return append(env, prefix+value)
+}
+
 func browserOpenCommandParts(openCommand, url string) ([]string, error) {
 	parts := strings.Fields(strings.TrimSpace(openCommand))
 	if len(parts) == 0 {
diff --git a/internal/flamegraph/liveserver_open_test.go b/internal/flamegraph/liveserver_open_test.go
index 0a42a61..aa9340a 100644
--- a/internal/flamegraph/liveserver_open_test.go
+++ b/internal/flamegraph/liveserver_open_test.go
@@ -2,6 +2,9 @@ package flamegraph
 
 import (
 	"errors"
+	"os"
+	"os/exec"
+	"strconv"
 	"testing"
 )
 
@@ -102,3 +105,75 @@ func TestMaybeOpenLiveBrowserPropagatesOpenError(t *testing.T) {
 		t.Fatalf("expected launch failed error, got %v", err)
 	}
 }
+
+func TestOpenBrowserURLReturnsErrorWhenCommandExitsNonZero(t *testing.T) {
+	err := openBrowserURL("http://localhost:1234/", "false")
+	if err == nil {
+		t.Fatalf("expected non-nil error")
+	}
+}
+
+func TestOpenBrowserURLReturnsNilWhenCommandExitsZero(t *testing.T) {
+	err := openBrowserURL("http://localhost:1234/", "true")
+	if err != nil {
+		t.Fatalf("expected nil error, got %v", err)
+	}
+}
+
+func TestApplySudoInvokerContextWithEnvSetsCredential(t *testing.T) {
+	cmd := exec.Command("echo")
+	uid := os.Getuid()
+	gid := os.Getgid()
+	env := []string{
+		"SUDO_UID=" + strconv.Itoa(uid),
+		"SUDO_GID=" + strconv.Itoa(gid),
+		"SUDO_USER=tester",
+		"HOME=/root",
+	}
+
+	applySudoInvokerContextWithEnv(cmd, 0, env)
+
+	if cmd.SysProcAttr == nil || cmd.SysProcAttr.Credential == nil {
+		t.Fatalf("expected process credentials to be configured")
+	}
+	if got := cmd.SysProcAttr.Credential.Uid; got != uint32(uid) {
+		t.Fatalf("credential uid = %d, want %d", got, uint32(uid))
+	}
+	if got := cmd.SysProcAttr.Credential.Gid; got != uint32(gid) {
+		t.Fatalf("credential gid = %d, want %d", got, uint32(gid))
+	}
+	if got, ok := lookupEnvValue(cmd.Env, "USER"); !ok || got != "tester" {
+		t.Fatalf("USER env = %q (ok=%v), want %q", got, ok, "tester")
+	}
+	if got, ok := lookupEnvValue(cmd.Env, "LOGNAME"); !ok || got != "tester" {
+		t.Fatalf("LOGNAME env = %q (ok=%v), want %q", got, ok, "tester")
+	}
+}
+
+func TestApplySudoInvokerContextWithEnvSkipsWhenNotRoot(t *testing.T) {
+	cmd := exec.Command("echo")
+	env := []string{
+		"SUDO_UID=1000",
+		"SUDO_GID=1000",
+		"SUDO_USER=tester",
+	}
+
+	applySudoInvokerContextWithEnv(cmd, 1000, env)
+
+	if cmd.SysProcAttr != nil {
+		t.Fatalf("expected credentials to remain nil for non-root euid")
+	}
+	if cmd.Env != nil {
+		t.Fatalf("expected environment to remain nil for non-root euid")
+	}
+}
+
+func TestNotifyLiveWarningUsesCallback(t *testing.T) {
+	var got string
+	notifyLiveWarning(func(message string) {
+		got = message
+	}, "open failed")
+	if got != "open failed" {
+		t.Fatalf("warning callback got %q, want %q", got, "open failed")
+	}
+}