From edd7a9ddc9138715bc69385c3be4bd30e2d1bfdf Mon Sep 17 00:00:00 2001
From: lordwelch <timmy@narnian.us>
Date: Mon, 4 May 2020 04:06:11 -0700
Subject: [PATCH] Improve the agent signal handling and error handling

There is probably a better way to make the list of signals to catch

I also noticed that the agent kept dying on me when the CA was not setup
correctly with an EOF error from the ssh client
---
 cmd/sshrimp-agent/main.go | 33 +++++++++++++++++++++++++--------
 cmd/sshrimp-agent/unix.go | 10 ++++++++++
 2 files changed, 35 insertions(+), 8 deletions(-)
 create mode 100644 cmd/sshrimp-agent/unix.go

diff --git a/cmd/sshrimp-agent/main.go b/cmd/sshrimp-agent/main.go
index 6c8a14e..ce534fb 100644
--- a/cmd/sshrimp-agent/main.go
+++ b/cmd/sshrimp-agent/main.go
@@ -1,9 +1,12 @@
 package main
 
 import (
+	"crypto"
 	"crypto/rand"
 	"crypto/rsa"
+	"errors"
 	"fmt"
+	"io"
 	"net"
 	"os"
 	"os/signal"
@@ -18,6 +21,8 @@ import (
 	"golang.org/x/crypto/ssh/agent"
 )
 
+var sigs = []os.Signal{os.Kill, os.Interrupt}
+
 var cli struct {
 	Config string `kong:"arg,type='string',help='sshrimp config file (default: ${config_file} or ${env_var_name} environment variable)',default='${config_file}',env='SSHRIMP_CONFIG'"`
 }
@@ -38,15 +43,26 @@ func main() {
 }
 
 func launchAgent(c *config.SSHrimp, ctx *kong.Context) error {
+	var (
+		err        error
+		listener   net.Listener
+		privateKey crypto.Signer
+		signer     ssh.Signer
+	)
 
-	if _, err := os.Stat(c.Agent.Socket); err == nil {
-		return fmt.Errorf("socket %s already exists", c.Agent.Socket)
+	if _, err = os.Stat(c.Agent.Socket); err == nil {
+		conn, sockErr := net.Dial("unix", c.Agent.Socket)
+		if sockErr == nil { // socket is accepting connections
+			conn.Close()
+			return fmt.Errorf("socket %s already exists", c.Agent.Socket)
+		}
+		os.Remove(c.Agent.Socket) // socket is not accepting connections, assuming safe to remove
 	}
 
 	// This affects all files created for the process. Since this is a sensitive
 	// socket, only allow the current user to write to the socket.
 	syscall.Umask(0077)
-	listener, err := net.Listen("unix", c.Agent.Socket)
+	listener, err = net.Listen("unix", c.Agent.Socket)
 	if err != nil {
 		return err
 	}
@@ -55,11 +71,11 @@ func launchAgent(c *config.SSHrimp, ctx *kong.Context) error {
 	ctx.Printf("listening on %s", c.Agent.Socket)
 
 	// Generate a new SSH private/public key pair
-	privateKey, err := rsa.GenerateKey(rand.Reader, 2048)
+	privateKey, err = rsa.GenerateKey(rand.Reader, 2048)
 	if err != nil {
 		return err
 	}
-	signer, err := ssh.NewSignerFromKey(privateKey)
+	signer, err = ssh.NewSignerFromKey(privateKey)
 	if err != nil {
 		return err
 	}
@@ -69,7 +85,7 @@ func launchAgent(c *config.SSHrimp, ctx *kong.Context) error {
 
 	// Listen for signals so that we can close the listener and exit nicely
 	osSignals := make(chan os.Signal)
-	signal.Notify(osSignals, os.Interrupt, os.Kill)
+	signal.Notify(osSignals, sigs...)
 	go func() {
 		_ = <-osSignals
 		listener.Close()
@@ -77,7 +93,8 @@ func launchAgent(c *config.SSHrimp, ctx *kong.Context) error {
 
 	// Accept connections and serve the agent
 	for {
-		conn, err := listener.Accept()
+		var conn net.Conn
+		conn, err = listener.Accept()
 		if err != nil {
 			if strings.Contains(err.Error(), "use of closed network connection") {
 				// Occurs if the user interrupts the agent with a ctrl-c signal
@@ -85,7 +102,7 @@ func launchAgent(c *config.SSHrimp, ctx *kong.Context) error {
 			}
 			return err
 		}
-		if err := agent.ServeAgent(sshrimpAgent, conn); err != nil {
+		if err = agent.ServeAgent(sshrimpAgent, conn); err != nil && !errors.Is(err, io.EOF) {
 			return err
 		}
 	}
diff --git a/cmd/sshrimp-agent/unix.go b/cmd/sshrimp-agent/unix.go
new file mode 100644
index 0000000..5a47e69
--- /dev/null
+++ b/cmd/sshrimp-agent/unix.go
@@ -0,0 +1,10 @@
+// +build darwin linux
+package main
+
+import (
+	"syscall"
+)
+
+func init() {
+	sigs = append(sigs, syscall.SIGTERM, syscall.SIGHUP)
+}