cmd: read environment variables from generic function

Signed-off-by: Morten Linderud <[email protected]>

Author: Foxboron

cmd/ssh-tpm-add/main.go

@@ -50,7 +50,7 @@ func main() {
 	flag.StringVar(&user, "user", "", "remote ssh user")
 	flag.Parse()
 
-	socket := os.Getenv("SSH_TPM_AUTH_SOCK")
+	socket := utils.EnvSocketPath("")
 	if socket == "" {
 		fmt.Println("Can't find any ssh-tpm-agent socket.")
 		os.Exit(1)

cmd/ssh-tpm-agent/main.go

@@ -9,7 +9,6 @@ import (
 	"net"
 	"os"
 	"os/signal"
-	"path"
 	"path/filepath"
 	"syscall"
 
@@ -112,22 +111,9 @@ func main() {
 		noCache                          bool
 	)
 
-	envSocketPath := func() string {
-		// Find a default socket name from ssh-tpm-agent.service
-		if val, ok := os.LookupEnv("SSH_TPM_AUTH_SOCK"); ok && socketPath == "" {
-			return val
-		}
-
-		dir := os.Getenv("XDG_RUNTIME_DIR")
-		if dir == "" {
-			dir = "/var/tmp"
-		}
-		return path.Join(dir, "ssh-tpm-agent.sock")
-	}()
-
 	var sockets SocketSet
 
-	flag.StringVar(&socketPath, "l", envSocketPath, "path of the UNIX socket to listen on")
+	flag.StringVar(&socketPath, "l", func(s string) string { return utils.EnvSocketPath(s) }(socketPath), "path of the UNIX socket to listen on")
 	flag.Var(&sockets, "A", "fallback ssh-agent sockets")
 	flag.BoolVar(&swtpmFlag, "swtpm", false, "use swtpm instead of actual tpm")
 	flag.BoolVar(&printSocketFlag, "print-socket", false, "print path of UNIX socket to stdout")

utils/tpm.go

@@ -80,3 +80,16 @@ func TPM(f bool) (transport.TPMCloser, error) {
 	}
 	return tpm, nil
 }
+
+func EnvSocketPath(socketPath string) string {
+	// Find a default socket name from ssh-tpm-agent.service
+	if val, ok := os.LookupEnv("SSH_TPM_AUTH_SOCK"); ok && socketPath == "" {
+		return val
+	}
+
+	dir := os.Getenv("XDG_RUNTIME_DIR")
+	if dir == "" {
+		dir = "/var/tmp"
+	}
+	return path.Join(dir, "ssh-tpm-agent.sock")
+}