support 1password agent #51

IdentityAgent
trzsz · Nov 11, 2023 · faef20f · faef20f
1 parent 577ecdb
commit faef20f
Show file tree

Hide file tree

Showing 4 changed files with 120 additions and 38 deletions.
diff --git a/tssh/agent_unix.go b/tssh/agent_unix.go
@@ -30,9 +30,11 @@ package tssh
 import (
 	"net"
 	"os"
+	"strings"
 	"sync"
 	"time"
 
+	"golang.org/x/crypto/ssh"
 	"golang.org/x/crypto/ssh/agent"
 )
 
@@ -42,32 +44,43 @@ var (
 	agentClient agent.ExtendedAgent
 )
 
-func getAgentClient() agent.ExtendedAgent {
+func getAgentAddr(args *sshArgs) string {
+	if addr := getOptionConfig(args, "IdentityAgent"); addr != "" {
+		if strings.ToLower(addr) == "none" {
+			return ""
+		}
+		return addr
+	}
+	return os.Getenv("SSH_AUTH_SOCK")
+}
+
+func getAgentClient(args *sshArgs) agent.ExtendedAgent {
 	agentOnce.Do(func() {
-		sock := os.Getenv("SSH_AUTH_SOCK")
-		if sock == "" {
-			debug("no ssh agent environment variable SSH_AUTH_SOCK")
+		addr := resolveHomeDir(getAgentAddr(args))
+		if addr == "" {
+			debug("ssh agent unix socket is not set")
 			return
 		}
 
 		var err error
-		agentConn, err = net.DialTimeout("unix", sock, time.Second)
+		agentConn, err = net.DialTimeout("unix", addr, time.Second)
 		if err != nil {
-			debug("dial ssh agent unix socket [%s] failed: %v", sock, err)
+			debug("dial ssh agent unix socket [%s] failed: %v", addr, err)
 			return
 		}
 
 		agentClient = agent.NewClient(agentConn)
-		debug("new ssh agent client [%s] success", sock)
-	})
+		debug("new ssh agent client [%s] success", addr)
 
+		cleanupAfterLogined = append(cleanupAfterLogined, func() {
+			agentConn.Close()
+			agentConn = nil
+			agentClient = nil
+		})
+	})
 	return agentClient
 }
 
-func closeAgentClient() {
-	if agentConn != nil {
-		agentConn.Close()
-		agentConn = nil
-	}
-	agentClient = nil
+func forwardToRemote(client *ssh.Client, addr string) error {
+	return agent.ForwardToRemote(client, addr)
 }
diff --git a/tssh/agent_windows.go b/tssh/agent_windows.go
@@ -26,11 +26,15 @@ SOFTWARE.
 */
 
 import (
+	"fmt"
+	"io"
 	"os"
+	"strings"
 	"sync"
 	"time"
 
 	"github.com/trzsz/npipe"
+	"golang.org/x/crypto/ssh"
 	"golang.org/x/crypto/ssh/agent"
 )
 
@@ -40,36 +44,94 @@ var (
 	agentClient agent.ExtendedAgent
 )
 
-func getAgentClient() agent.ExtendedAgent {
-	agentOnce.Do(func() {
-		name := `\\.\pipe\openssh-ssh-agent`
-		if sock := os.Getenv("SSH_AUTH_SOCK"); sock != "" {
-			name = sock
+func getAgentAddr(args *sshArgs) string {
+	if addr := getOptionConfig(args, "IdentityAgent"); addr != "" {
+		if strings.ToLower(addr) == "none" {
+			return ""
 		}
+		return addr
+	}
+	if addr := os.Getenv("SSH_AUTH_SOCK"); addr != "" {
+		return addr
+	}
+	if addr := `\\.\pipe\openssh-ssh-agent`; isFileExist(addr) {
+		return addr
+	}
+	return ""
+}
 
-		if !isFileExist(name) {
-			debug("ssh agent named pipe [%s] does not exist", name)
+func getAgentClient(args *sshArgs) agent.ExtendedAgent {
+	agentOnce.Do(func() {
+		addr := getAgentAddr(args)
+		if addr == "" {
+			debug("ssh agent named pipe is not set")
 			return
 		}
 
 		var err error
-		agentConn, err = npipe.DialTimeout(name, time.Second)
+		agentConn, err = npipe.DialTimeout(addr, time.Second)
 		if err != nil {
-			debug("dial ssh agent named pipe [%s] failed: %v", name, err)
+			debug("dial ssh agent named pipe [%s] failed: %v", addr, err)
 			return
 		}
 
 		agentClient = agent.NewClient(agentConn)
-		debug("new ssh agent client [%s] success", name)
-	})
+		debug("new ssh agent client [%s] success", addr)
 
+		cleanupAfterLogined = append(cleanupAfterLogined, func() {
+			agentConn.Close()
+			agentConn = nil
+			agentClient = nil
+		})
+	})
 	return agentClient
 }
 
-func closeAgentClient() {
-	if agentConn != nil {
-		agentConn.Close()
-		agentConn = nil
+const channelType = "[email protected]"
+
+func forwardToRemote(client *ssh.Client, addr string) error {
+	channels := client.HandleChannelOpen(channelType)
+	if channels == nil {
+		return fmt.Errorf("agent: already have handler for %s", channelType)
 	}
-	agentClient = nil
+	conn, err := npipe.DialTimeout(addr, time.Second)
+	if err != nil {
+		return err
+	}
+	conn.Close()
+
+	go func() {
+		for ch := range channels {
+			channel, reqs, err := ch.Accept()
+			if err != nil {
+				continue
+			}
+			go ssh.DiscardRequests(reqs)
+			go forwardNamedPipe(channel, addr)
+		}
+	}()
+	return nil
+}
+
+func forwardNamedPipe(channel ssh.Channel, addr string) {
+	conn, err := npipe.DialTimeout(addr, time.Second)
+	if err != nil {
+		return
+	}
+
+	var wg sync.WaitGroup
+	wg.Add(2)
+	go func() {
+		io.Copy(conn, channel)
+		wg.Done()
+	}()
+	go func() {
+		io.Copy(channel, conn)
+		channel.CloseWrite()
+		wg.Done()
+	}()
+
+	wg.Wait()
+	conn.Close()
+	channel.Close()
 }
diff --git a/tssh/ctrl_unix.go b/tssh/ctrl_unix.go
@@ -191,10 +191,13 @@ func startControlMaster(args *sshArgs) {
 		return
 	}
 
-	cmdArgs := []string{"-a", "-T", "-oClearAllForwardings=yes", "-oRemoteCommand=none", "-oConnectTimeout=5"}
+	cmdArgs := []string{"-T", "-oClearAllForwardings=yes", "-oRemoteCommand=none", "-oConnectTimeout=5"}
 	if args.Debug {
 		cmdArgs = append(cmdArgs, "-v")
 	}
+	if !args.NoForwardAgent && args.ForwardAgent {
+		cmdArgs = append(cmdArgs, "-A")
+	}
 	if args.LoginName != "" {
 		cmdArgs = append(cmdArgs, "-l", args.LoginName)
 	}
@@ -219,6 +222,8 @@ func startControlMaster(args *sshArgs) {
 			cmdArgs = append(cmdArgs, fmt.Sprintf("-oControlPath=%s", value))
 		case "controlpersist":
 			cmdArgs = append(cmdArgs, fmt.Sprintf("-oControlPersist=%s", value))
+		case "forwardagent":
+			cmdArgs = append(cmdArgs, fmt.Sprintf("-oForwardAgent=%s", value))
 		}
 	}
 

diff --git a/tssh/login.go b/tssh/login.go
@@ -529,7 +529,7 @@ func getPublicKeysAuthMethod(args *sshArgs) ssh.AuthMethod {
 		}
 	}
 
-	if agentClient := getAgentClient(); agentClient != nil {
+	if agentClient := getAgentClient(args); agentClient != nil {
 		signers, err := agentClient.Signers()
 		if err != nil {
 			warning("get ssh agent signers failed: %v", err)
@@ -838,15 +838,15 @@ func keepAlive(client *ssh.Client, args *sshArgs) {
 
 func sshAgentForward(args *sshArgs, client *ssh.Client, session *ssh.Session) {
 	if args.NoForwardAgent || !args.ForwardAgent && strings.ToLower(getOptionConfig(args, "ForwardAgent")) != "yes" {
-		closeAgentClient()
 		return
 	}
-	agentClient := getAgentClient()
-	if agentClient == nil {
+	addr := resolveHomeDir(getAgentAddr(args))
+	if addr == "" {
+		warning("forward agent but the socket address is not set")
 		return
 	}
-	if err := agent.ForwardToAgent(client, agentClient); err != nil {
-		warning("forward to agent failed: %v", err)
+	if err := forwardToRemote(client, addr); err != nil {
+		warning("forward to agent [%s] failed: %v", addr, err)
 		return
 	}
 	if err := agent.RequestAgentForwarding(session); err != nil {
@@ -915,7 +915,9 @@ func sshLogin(args *sshArgs, tty bool) (client *ssh.Client, session *ssh.Session
 	session.Stderr = os.Stderr
 
 	// ssh agent forward
-	sshAgentForward(args, client, session)
+	if !control {
+		sshAgentForward(args, client, session)
+	}
 
 	// not terminal or not tty
 	if !isTerminal || !tty {