client: Check agent connection first

1 files changed, 12 insertions, 5 deletions

diff --git a/cmd/client/client.go b/cmd/client/client.go
index 7799b87..a404c65 100644
--- a/cmd/client/client.go
+++ b/cmd/client/client.go
@@ -120,16 +120,18 @@ func getCertificateFromCA(ctx context.Context, oauthToken string, certRequest []
         return cert, nil
 }
 
-func addCertificateToAgent(private any, cert *ssh.Certificate) error {
+func connectToAgent() (agent.ExtendedAgent, error) {
         socket := os.Getenv("SSH_AUTH_SOCK")
         conn, err := net.Dial("unix", socket)
         if err != nil {
-                return err
+                return nil, err
         }
 
-        agentConn := agent.NewClient(conn)
+        return agent.NewClient(conn), nil
+}
 
-        return agentConn.Add(agent.AddedKey{
+func addCertificateToAgent(conn agent.ExtendedAgent, private any, cert *ssh.Certificate) error {
+        return conn.Add(agent.AddedKey{
                 PrivateKey:   private,
                 Certificate:  cert,
                 LifetimeSecs: 10,
@@ -191,6 +193,11 @@ func clientMain(cfg app.Config, host, port, username string) {
         ctx, cancel := context.WithCancel(context.Background())
         defer cancel()
 
+        agentConn, err := connectToAgent()
+        if err != nil {
+                log.Fatalf("Error connecting to agent, is it started?")
+        }
+
         oauthToken, err := fetchOauthToken(ctx, clientId, cfg.ClientHost)
         if err != nil {
                 log.Fatalf("Error fetching oauth token: %s", err)
@@ -206,7 +213,7 @@ func clientMain(cfg app.Config, host, port, username string) {
                 log.Fatalf("Error fetching certificate: %s", err)
         }
 
-        if err := addCertificateToAgent(privateKey, certificate); err != nil {
+        if err := addCertificateToAgent(agentConn, privateKey, certificate); err != nil {
                 log.Fatalf("Error adding certificate to agent: %s", err)
         }