From 33a4d6ef05c82e84c4b19b496bc9388999c1192c Mon Sep 17 00:00:00 2001
From: Troels Liebe Bentsen <tlb@connectedcars.io>
Date: Mon, 17 Oct 2022 23:35:56 +0200
Subject: [PATCH] Allow setting auth sock location

---
 cmd/authwrapper/main.go  | 29 ++++++++++++++++++++++++++++-
 cmd/authwrapper/setup.go |  3 +++
 cmd/authwrapper/utils.go |  9 ++-------
 sshagent/sshagent.go     | 23 +----------------------
 4 files changed, 34 insertions(+), 30 deletions(-)

diff --git a/cmd/authwrapper/main.go b/cmd/authwrapper/main.go
index 27aad47..e242b56 100644
--- a/cmd/authwrapper/main.go
+++ b/cmd/authwrapper/main.go
@@ -2,11 +2,14 @@ package main
 
 import (
 	"fmt"
+	"io/ioutil"
 	"log"
+	"math/rand"
 	"os"
 	"strings"
 	"time"
 
+	"github.com/connectedcars/auth-wrapper/sshagent"
 	"golang.org/x/crypto/ssh"
 )
 
@@ -67,7 +70,21 @@ func main() {
 		}
 	}
 
-	exitCode, err := runCommandWithSSHAgent(agent, config.Command, config.Args)
+	var sshAuthSock string
+	if config.SSHAgentSocketPath != "" {
+		sshAuthSock = config.SSHAgentSocketPath
+	} else {
+		// Generate random filename
+		dir, err := ioutil.TempDir(os.TempDir(), "")
+		if err != nil {
+			log.Fatal(err)
+		}
+		sshAuthSock = dir + "/" + generateRandomString(8) + ".sock"
+	}
+
+	sshagent.StartSSHAgentServer(agent, sshAuthSock)
+
+	exitCode, err := runCommandWithSSHAgent(sshAuthSock, config.Command, config.Args)
 	if err != nil {
 		log.Fatalf("runCommandWithSSHAgent: %v", err)
 	}
@@ -77,3 +94,13 @@ func main() {
 	}
 	os.Exit(exitCode)
 }
+
+const letterBytes = "abcdefghijklmnopqrstuvwxyz"
+
+func generateRandomString(n int) string {
+	b := make([]byte, n)
+	for i := range b {
+		b[i] = letterBytes[rand.Intn(len(letterBytes))]
+	}
+	return string(b)
+}
diff --git a/cmd/authwrapper/setup.go b/cmd/authwrapper/setup.go
index 4704b48..64c5174 100644
--- a/cmd/authwrapper/setup.go
+++ b/cmd/authwrapper/setup.go
@@ -30,6 +30,7 @@ type Config struct {
 	SSHCaAuthorizedKeysPath string
 	SSHSigningServerAddress string
 	SSHAgentSocket          string
+	SSHAgentSocketPath      string
 	AuthWrapperQuiet        bool
 }
 
@@ -55,6 +56,7 @@ func parseEnvironment() (*Config, error) {
 		SSHCaAuthorizedKeysPath: os.Getenv("SSH_CA_AUTHORIZED_KEYS_PATH"),
 		SSHSigningServerAddress: os.Getenv("SSH_SIGNING_SERVER_LISTEN_ADDRESS"),
 		SSHAgentSocket:          os.Getenv("SSH_AUTH_SOCK"),
+		SSHAgentSocketPath:      os.Getenv("SSH_AUTH_SOCK_PATH"),
 		AuthWrapperQuiet:        isAuthWrapperQuiet,
 	}
 	os.Unsetenv("WRAP_COMMAND")
@@ -65,6 +67,7 @@ func parseEnvironment() (*Config, error) {
 	os.Unsetenv("SSH_CA_KEY_PASSWORD")
 	os.Unsetenv("SSH_SIGNING_SERVER_LISTEN_ADDRESS")
 	os.Unsetenv("SSH_AUTH_SOCK")
+	os.Unsetenv("SSH_AUTH_SOCK_PATH")
 	os.Unsetenv("AUTH_WRAPPER_QUIET")
 
 	if *principalsFlag != "" {
diff --git a/cmd/authwrapper/utils.go b/cmd/authwrapper/utils.go
index 3cce7a2..0e39f91 100644
--- a/cmd/authwrapper/utils.go
+++ b/cmd/authwrapper/utils.go
@@ -17,18 +17,13 @@ import (
 
 	"github.com/connectedcars/auth-wrapper/kms/google"
 	"github.com/connectedcars/auth-wrapper/server"
-	"github.com/connectedcars/auth-wrapper/sshagent"
+
 	"golang.org/x/crypto/ssh"
-	"golang.org/x/crypto/ssh/agent"
 )
 
 var httpClient = &http.Client{Timeout: 10 * time.Second}
 
-func runCommandWithSSHAgent(agent agent.ExtendedAgent, command string, args []string) (exitCode int, err error) {
-	sshAuthSock, err := sshagent.StartSSHAgentServer(agent)
-	if err != nil {
-		return 255, fmt.Errorf("Failed to start ssh agent server: %v", err)
-	}
+func runCommandWithSSHAgent(sshAuthSock string, command string, args []string) (exitCode int, err error) {
 
 	os.Setenv("SSH_AUTH_SOCK", sshAuthSock)
 
diff --git a/sshagent/sshagent.go b/sshagent/sshagent.go
index 82e151d..407bbec 100644
--- a/sshagent/sshagent.go
+++ b/sshagent/sshagent.go
@@ -2,9 +2,7 @@ package sshagent
 
 import (
 	"fmt"
-	"io/ioutil"
 	"log"
-	"math/rand"
 	"net"
 	"os"
 	"strings"
@@ -14,14 +12,7 @@ import (
 )
 
 // StartSSHAgentServer start an SSH Agent server and loads the given private key
-func StartSSHAgentServer(sshAgent agent.Agent) (sshAuthSock string, error error) {
-	// Generate random filename
-	dir, err := ioutil.TempDir(os.TempDir(), "")
-	if err != nil {
-		log.Fatal(err)
-	}
-	sshAuthSock = dir + "/" + generateRandomString(8) + ".sock"
-
+func StartSSHAgentServer(sshAgent agent.Agent, sshAuthSock string) {
 	go func() {
 		// Open SSH agent socket
 		if err := os.RemoveAll(sshAuthSock); err != nil {
@@ -42,8 +33,6 @@ func StartSSHAgentServer(sshAgent agent.Agent) (sshAuthSock string, error error)
 			go agent.ServeAgent(sshAgent, conn)
 		}
 	}()
-
-	return sshAuthSock, err
 }
 
 // ConnectSSHAgent connects to a SSH agent socket and returns a agent.ExtendedAgent
@@ -55,16 +44,6 @@ func ConnectSSHAgent(socket string) (agent.ExtendedAgent, error) {
 	return agent.NewClient(conn), nil
 }
 
-const letterBytes = "abcdefghijklmnopqrstuvwxyz"
-
-func generateRandomString(n int) string {
-	b := make([]byte, n)
-	for i := range b {
-		b[i] = letterBytes[rand.Intn(len(letterBytes))]
-	}
-	return string(b)
-}
-
 // ParsePrivateSSHKey parses a private key
 func ParsePrivateSSHKey(privateKeyBytes []byte, passphrase string) (interface{}, error) {
 	var err error