From a84465ed37f8eda7c08bc0533bb1398ef7813918 Mon Sep 17 00:00:00 2001
From: Josh Deprez <jd@buildkite.com>
Date: Tue, 17 Sep 2024 15:45:36 +1000
Subject: [PATCH] Add osutil.UserHomeDir

---
 clicommand/agent_start.go       |  2 +-
 internal/job/knownhosts.go      |  3 +-
 internal/osutil/homedir.go      | 12 ++++++++
 internal/osutil/homedir_test.go | 50 +++++++++++++++++++++++++++++++++
 4 files changed, 65 insertions(+), 2 deletions(-)
 create mode 100644 internal/osutil/homedir.go
 create mode 100644 internal/osutil/homedir_test.go

diff --git a/clicommand/agent_start.go b/clicommand/agent_start.go
index 70fcaec883..157bb2a22c 100644
--- a/clicommand/agent_start.go
+++ b/clicommand/agent_start.go
@@ -1383,7 +1383,7 @@ func agentLifecycleHook(hookName string, log logger.Logger, cfg AgentStartConfig
 }
 
 func defaultSocketsPath() string {
-	home, err := os.UserHomeDir()
+	home, err := osutil.UserHomeDir()
 	if err != nil {
 		return filepath.Join(os.TempDir(), "buildkite-sockets")
 	}
diff --git a/internal/job/knownhosts.go b/internal/job/knownhosts.go
index 03ab7e4de7..4b6f1c5043 100644
--- a/internal/job/knownhosts.go
+++ b/internal/job/knownhosts.go
@@ -10,6 +10,7 @@ import (
 	"time"
 
 	"github.com/buildkite/agent/v3/internal/job/shell"
+	"github.com/buildkite/agent/v3/internal/osutil"
 	"golang.org/x/crypto/ssh/knownhosts"
 )
 
@@ -19,7 +20,7 @@ type knownHosts struct {
 }
 
 func findKnownHosts(sh *shell.Shell) (*knownHosts, error) {
-	userHomePath, err := os.UserHomeDir()
+	userHomePath, err := osutil.UserHomeDir()
 	if err != nil {
 		return nil, fmt.Errorf("Could not find the current users home directory (%s)", err)
 	}
diff --git a/internal/osutil/homedir.go b/internal/osutil/homedir.go
new file mode 100644
index 0000000000..e56cae9739
--- /dev/null
+++ b/internal/osutil/homedir.go
@@ -0,0 +1,12 @@
+package osutil
+
+import "os"
+
+// UserHomeDir is similar to os.UserHomeDir, but prefers $HOME when available
+// over other options (such as USERPROFILE on Windows).
+func UserHomeDir() (string, error) {
+	if h := os.Getenv("HOME"); h != "" {
+		return h, nil
+	}
+	return os.UserHomeDir()
+}
diff --git a/internal/osutil/homedir_test.go b/internal/osutil/homedir_test.go
new file mode 100644
index 0000000000..7fa2f9450c
--- /dev/null
+++ b/internal/osutil/homedir_test.go
@@ -0,0 +1,50 @@
+package osutil
+
+import (
+	"os"
+	"runtime"
+	"testing"
+)
+
+func TestUserHomeDir(t *testing.T) {
+	// not parallel because it messes with env vars
+	origHome := os.Getenv("HOME")
+	origUserProfile := os.Getenv("USERPROFILE")
+	t.Cleanup(func() {
+		os.Setenv("HOME", origHome)
+		os.Setenv("USERPROFILE", origUserProfile)
+	})
+
+	type testCase struct {
+		home, userProfile, want string
+	}
+
+	tests := []testCase{
+		{
+			// Prefer $HOME on all platforms
+			home:        "home",
+			userProfile: "userProfile",
+			want:        "home",
+		},
+	}
+	if runtime.GOOS == "windows" {
+		// Windows can use %USERPROFILE% as a treat when $HOME is unavailable
+		tests = append(tests, testCase{
+			home:        "",
+			userProfile: "userProfile",
+			want:        "userProfile",
+		})
+	}
+
+	for _, test := range tests {
+		os.Setenv("HOME", test.home)
+		os.Setenv("USERPROFILE", test.userProfile)
+		got, err := UserHomeDir()
+		if err != nil {
+			t.Errorf("HOME=%q USERPROFILE=%q UserHomeDir() error = %v", test.home, test.userProfile, err)
+		}
+		if got != test.want {
+			t.Errorf("HOME=%q USERPROFILE=%q UserHomeDir() = %q, want %q", test.home, test.userProfile, got, test.want)
+		}
+	}
+}