ssh-tpm-agent: unroll two functions used once

Signed-off-by: Morten Linderud <morten@linderud.pw>
Foxboron · Aug 2, 2023 · ed9794c · ed9794c
1 parent 6092d3f
commit ed9794c
Show file tree

Hide file tree

Showing 2 changed files with 42 additions and 49 deletions.
diff --git a/agent/agent.go b/agent/agent.go
@@ -10,20 +10,17 @@ import (
 	"log"
 	"net"
 	"os"
-	"os/signal"
 	"path"
 	"path/filepath"
 	"strings"
 	"sync"
-	"syscall"
 	"time"
 
 	"github.com/foxboron/ssh-tpm-agent/key"
 	"github.com/foxboron/ssh-tpm-agent/signer"
 	"github.com/google/go-tpm/tpm2/transport"
 	"golang.org/x/crypto/ssh"
 	"golang.org/x/crypto/ssh/agent"
-	"golang.org/x/term"
 )
 
 var ErrOperationUnsupported = errors.New("operation unsupported")
@@ -250,37 +247,3 @@ func NewAgent(socketPath string, tpmFetch func() transport.TPMCloser, pin func(*
 	go a.serve()
 	return a
 }
-
-func execAgent(socketPath string, tpmFetch func() transport.TPMCloser, pin func(*key.Key) ([]byte, error)) *Agent {
-	os.Remove(socketPath)
-	if err := os.MkdirAll(filepath.Dir(socketPath), 0777); err != nil {
-		log.Fatalln("Failed to create UNIX socket folder:", err)
-	}
-	log.Printf("Listening on %v\n", socketPath);
-	a := NewAgent(socketPath, tpmFetch, pin)
-
-	c := make(chan os.Signal, 1)
-	signal.Notify(c, syscall.SIGHUP)
-	go func() {
-		for range c {
-			a.Stop()
-		}
-	}()
-
-	return a
-}
-
-func RunAgent(socketPath string, tpmFetch func() transport.TPMCloser, pin func(*key.Key) ([]byte, error)) {
-	if term.IsTerminal(int(os.Stdin.Fd())) {
-		log.Println("Warning: ssh-tpm-agent is meant to run as a background daemon.")
-		log.Println("Running multiple instances is likely to lead to conflicts.")
-		log.Println("Consider using a systemd service.")
-	}
-
-	a := execAgent(socketPath, tpmFetch, pin)
-
-	//TODO: Maybe we should allow people to not auto-load keys
-	a.LoadKeys()
-
-	a.Wait()
-}
diff --git a/cmd/ssh-tpm-agent/main.go b/cmd/ssh-tpm-agent/main.go
@@ -77,18 +77,48 @@ func main() {
 		os.Exit(0)
 	}
 
-	tpmFetch := func() (tpm transport.TPMCloser) {
-		// the agent will close the TPM after this is called
-		tpm, err := utils.GetTPM(swtpmFlag)
-		if err != nil {
-			log.Fatal(err)
-		}
-		return tpm
+	if term.IsTerminal(int(os.Stdin.Fd())) {
+		log.Println("Warning: ssh-tpm-agent is meant to run as a background daemon.")
+		log.Println("Running multiple instances is likely to lead to conflicts.")
+		log.Println("Consider using a systemd service.")
 	}
-	pin := func(key *key.Key) ([]byte, error) {
-		keyHash := sha256.Sum256(key.Public.Bytes())
-		keyInfo := fmt.Sprintf("ssh-tpm-agent/%x", keyHash)
-		return pinentry.GetPinentry(keyInfo)
+
+	os.Remove(socketPath)
+	if err := os.MkdirAll(filepath.Dir(socketPath), 0777); err != nil {
+		log.Fatalln("Failed to create UNIX socket folder:", err)
 	}
-	agent.RunAgent(socketPath, tpmFetch, pin)
+	log.Printf("Listening on %v\n", socketPath)
+
+	a := agent.NewAgent(socketPath,
+		// TPM Callback
+		func() (tpm transport.TPMCloser) {
+			// the agent will close the TPM after this is called
+			tpm, err := utils.GetTPM(swtpmFlag)
+			if err != nil {
+				log.Fatal(err)
+			}
+			return tpm
+		},
+
+		// PIN Callback
+		func(key *key.Key) ([]byte, error) {
+			keyHash := sha256.Sum256(key.Public.Bytes())
+			keyInfo := fmt.Sprintf("ssh-tpm-agent/%x", keyHash)
+			return pinentry.GetPinentry(keyInfo)
+		},
+	)
+
+	// Signal handling
+	c := make(chan os.Signal, 1)
+	signal.Notify(c, syscall.SIGHUP)
+	go func() {
+		for range c {
+			a.Stop()
+		}
+	}()
+
+	//TODO: Maybe we should allow people to not auto-load keys
+	a.LoadKeys()
+
+	a.Wait()
 }