Add support for ssh-agent

Signed-off-by: yuguorui <[email protected]>

Author: yuguorui

main.go

@@ -45,6 +45,7 @@ import (
 
 	"github.com/dsnet/golib/jsonfmt"
 	"golang.org/x/crypto/ssh"
+	"golang.org/x/crypto/ssh/agent"
 	"golang.org/x/crypto/ssh/knownhosts"
 )
 
@@ -116,6 +117,23 @@ type KeepAliveConfig struct {
 	CountMax uint
 }
 
+func setupSshAgent() ssh.AuthMethod {
+	// ssh-agent(1) provides a UNIX socket at $SSH_AUTH_SOCK.
+	socket := os.Getenv("SSH_AUTH_SOCK")
+	if len(socket) == 0 {
+		return nil
+	}
+
+	conn, err := net.Dial("unix", socket)
+	if err != nil {
+		log.Printf("Failed to open SSH_AUTH_SOCK %s: %v\n", socket, err)
+		return nil
+	}
+
+	agentClient := agent.NewClient(conn)
+	return ssh.PublicKeysCallback(agentClient.Signers)
+}
+
 func loadConfig(conf string) (tunns []tunnel, logger *log.Logger, closer func() error) {
 	var logBuf bytes.Buffer
 	logger = log.New(io.MultiWriter(os.Stderr, &logBuf), "", log.Ldate|log.Ltime|log.Lshortfile)
@@ -171,11 +189,10 @@ func loadConfig(conf string) (tunns []tunnel, logger *log.Logger, closer func()
 		closer = f.Close
 	}
 
+	var auth []ssh.AuthMethod
+
 	// Parse all of the private keys.
 	var keys []ssh.Signer
-	if len(config.KeyFiles) == 0 {
-		logger.Fatal("no private keys specified")
-	}
 	for _, kf := range config.KeyFiles {
 		b, err := ioutil.ReadFile(kf)
 		if err != nil {
@@ -187,7 +204,19 @@ func loadConfig(conf string) (tunns []tunnel, logger *log.Logger, closer func()
 		}
 		keys = append(keys, k)
 	}
-	auth := []ssh.AuthMethod{ssh.PublicKeys(keys...)}
+	if len(keys) > 0 {
+		auth = append(auth, ssh.PublicKeys(keys...))
+	}
+
+	// Setup ssh-agent(1)
+	agent := setupSshAgent()
+	if agent != nil {
+		auth = append(auth, agent)
+	}
+
+	if len(auth) == 0 {
+		logger.Panic("no private keys and ssh-agent usable")
+	}
 
 	// Parse all of the host public keys.
 	if len(config.KnownHostFiles) == 0 {