From 98b45ef12cdb4cdda20d3fa4169219054e900c88 Mon Sep 17 00:00:00 2001
From: Maisem Ali <maisem@tailscale.com>
Date: Mon, 14 Mar 2022 13:26:06 -0700
Subject: [PATCH] ssh/tailssh: add support for agent forwarding.

Updates #3802

Signed-off-by: Maisem Ali <maisem@tailscale.com>
---
 ssh/tailssh/incubator.go |  4 +++
 ssh/tailssh/tailssh.go   | 60 ++++++++++++++++++++++++++++++++++++----
 tailcfg/tailcfg.go       |  4 +++
 3 files changed, 63 insertions(+), 5 deletions(-)

diff --git a/ssh/tailssh/incubator.go b/ssh/tailssh/incubator.go
index 3c24ba38b..dc073883f 100644
--- a/ssh/tailssh/incubator.go
+++ b/ssh/tailssh/incubator.go
@@ -186,6 +186,10 @@ func (ss *sshSession) launchProcess(ctx context.Context) error {
 
 	ss.cmd = cmd
 
+	if ss.agentListener != nil {
+		cmd.Env = append(cmd.Env, fmt.Sprintf("SSH_AUTH_SOCK=%s", ss.agentListener.Addr()))
+	}
+
 	ptyReq, winCh, isPty := ss.Pty()
 	if !isPty {
 		ss.logf("starting non-pty command: %+v", cmd.Args)
diff --git a/ssh/tailssh/tailssh.go b/ssh/tailssh/tailssh.go
index f8b769f12..2e6efaa78 100644
--- a/ssh/tailssh/tailssh.go
+++ b/ssh/tailssh/tailssh.go
@@ -20,6 +20,8 @@
 	"os"
 	"os/exec"
 	"os/user"
+	"path/filepath"
+	"strconv"
 	"strings"
 	"sync"
 	"time"
@@ -257,11 +259,12 @@ type sshSession struct {
 	sharedID string // ID that's shared with control
 	logf     logger.Logf
 
-	ctx       *sshContext // implements context.Context
-	srv       *server
-	connInfo  *sshConnInfo
-	action    *tailcfg.SSHAction
-	localUser *user.User
+	ctx           *sshContext // implements context.Context
+	srv           *server
+	connInfo      *sshConnInfo
+	action        *tailcfg.SSHAction
+	localUser     *user.User
+	agentListener net.Listener // non-nil if agent-forwarding requested+allowed
 
 	// initialized by launchProcess:
 	cmd    *exec.Cmd
@@ -385,6 +388,47 @@ func (srv *server) endSession(ss *sshSession) {
 
 var errSessionDone = errors.New("session is done")
 
+// handleSSHAgentForwarding starts a Unix socket listener and in the background
+// forwards agent connections between the listenr and the ssh.Session.
+// On success, it assigns ss.agentListener.
+func (ss *sshSession) handleSSHAgentForwarding(s ssh.Session, lu *user.User) error {
+	if !ssh.AgentRequested(ss) || !ss.action.AllowAgentForwarding {
+		return nil
+	}
+	ss.logf("ssh: agent forwarding requested")
+	ln, err := ssh.NewAgentListener()
+	if err != nil {
+		return err
+	}
+	defer func() {
+		if err != nil && ln != nil {
+			ln.Close()
+		}
+	}()
+
+	uid, err := strconv.ParseUint(lu.Uid, 10, 32)
+	if err != nil {
+		return err
+	}
+	gid, err := strconv.ParseUint(lu.Gid, 10, 32)
+	if err != nil {
+		return err
+	}
+	socket := ln.Addr().String()
+	dir := filepath.Dir(socket)
+	// Make sure the socket is accessible by the user.
+	if err := os.Chown(socket, int(uid), int(gid)); err != nil {
+		return err
+	}
+	if err := os.Chmod(dir, 0755); err != nil {
+		return err
+	}
+
+	go ssh.ForwardAgentConnections(ln, s)
+	ss.agentListener = ln
+	return nil
+}
+
 // run is the entrypoint for a newly accepted SSH session.
 //
 // When ctx is done, the session is forcefully terminated. If its Err
@@ -424,6 +468,12 @@ func (ss *sshSession) run() {
 	// See https://github.com/tailscale/tailscale/issues/4146
 	ss.DisablePTYEmulation()
 
+	if err := ss.handleSSHAgentForwarding(ss, lu); err != nil {
+		logf("ssh: agent forwarding failed: %v", err)
+	} else if ss.agentListener != nil {
+		// TODO(maisem/bradfitz): add a way to close all session resources
+		defer ss.agentListener.Close()
+	}
 	err := ss.launchProcess(ss.ctx)
 	if err != nil {
 		logf("start failed: %v", err.Error())
diff --git a/tailcfg/tailcfg.go b/tailcfg/tailcfg.go
index 5305dea4f..594e057b5 100644
--- a/tailcfg/tailcfg.go
+++ b/tailcfg/tailcfg.go
@@ -1615,6 +1615,10 @@ type SSHAction struct {
 	// before being forcefully terminated.
 	SesssionDuration time.Duration `json:"sessionDuration,omitempty"`
 
+	// AllowAgentForwarding, if true, allows accepted connections to forward
+	// the ssh agent if requested.
+	AllowAgentForwarding bool `json:"allowAgentForwarding,omitempty"`
+
 	// HoldAndDelegate, if non-empty, is a URL that serves an
 	// outcome verdict.  The connection will be accepted and will
 	// block until the provided long-polling URL serves a new