lib/client: reopen an agent connection for each request

lib/client/keyagent.go

@@ -150,11 +150,17 @@ func NewLocalAgent(conf LocalAgentConfig) (a *LocalKeyAgent, err error) {
 
 	if shouldAddKeysToAgent(conf.KeysOption) {
 		a.log.DebugContext(context.Background(), "Connecting to the system agent")
+		// Check whether the system agent is reachable
+		// and close the connection right away
 		systemAgent, err := sshagent.NewSystemAgentClient()
 		if err != nil {
 			a.log.WarnContext(context.Background(), "Unable to connect to system agent", "error", err)
 		} else {
-			a.systemAgent = systemAgent
+			if err := systemAgent.Close(); err != nil {
+				a.log.DebugContext(context.Background(), "Failed to close connection to system agent", "error", err)
+			}
+			// Build a client that connects to the system agent for each request.
+			a.systemAgent = sshagent.NewSingleRequestClient(sshagent.NewSystemAgentClient)
 		}
 	} else {
 		log.DebugContext(context.Background(), "Skipping connection to the local ssh-agent.")

lib/client/keyagent_test.go

@@ -28,6 +28,7 @@ import (
 	"path/filepath"
 	"runtime"
 	"sync"
+	"sync/atomic"
 	"testing"
 	"time"
 
@@ -63,6 +64,10 @@ type KeyAgentTestSuite struct {
 	clusterName string
 	tlsca       *tlsca.CertAuthority
 	tlscaCert   authclient.TrustedCerts
+
+	// openSystemAgentConns is the number of client connections currently
+	// open to the agent served on $SSH_AUTH_SOCK.
+	openSystemAgentConns *atomic.Int32
 }
 
 type keyAgentTestSuiteFunc func(opt *keyAgentTestSuiteOpt)
@@ -96,14 +101,15 @@ func makeSuite(t *testing.T, opts ...keyAgentTestSuiteFunc) *KeyAgentTestSuite {
 		o(&settings)
 	}
 
-	err := startDebugAgent(t)
+	openSystemAgentConns, err := startDebugAgent(t)
 	require.NoError(t, err)
 
 	s := &KeyAgentTestSuite{
-		keyDir:      t.TempDir(),
-		username:    "foo",
-		hostname:    settings.hostname,
-		clusterName: settings.clusterName,
+		keyDir:               t.TempDir(),
+		username:             "foo",
+		hostname:             settings.hostname,
+		clusterName:          settings.clusterName,
+		openSystemAgentConns: openSystemAgentConns,
 	}
 
 	pemBytes, ok := fixtures.PEMBytes["rsa"]
@@ -183,6 +189,40 @@ func TestAddKey(t *testing.T) {
 
 }
 
+// TestSystemAgentConnections ensures that a LocalKeyAgent does not hold
+// connections to the system agent open.
+func TestSystemAgentConnections(t *testing.T) {
+	s := makeSuite(t)
+
+	requireNoOpenConns := func(t *testing.T) {
+		t.Helper()
+		require.EventuallyWithT(t, func(t *assert.CollectT) {
+			assert.Zero(t, s.openSystemAgentConns.Load())
+		}, 5*time.Second, 10*time.Millisecond, "connections to the system agent were leaked")
+	}
+
+	// Creating many key agents must not accumulate connections to the system agent.
+	var keyAgents []*LocalKeyAgent
+	for range 10 {
+		keyAgent := s.newKeyAgent(t)
+		require.NotNil(t, keyAgent.systemAgent, "expected the key agent to use the system agent")
+		keyAgents = append(keyAgents, keyAgent)
+	}
+	requireNoOpenConns(t)
+
+	for _, keyAgent := range keyAgents {
+		require.NoError(t, keyAgent.AddKeyRing(s.keyRing))
+		requireNoOpenConns(t)
+
+		_, err := keyAgent.Signers()
+		require.NoError(t, err)
+		requireNoOpenConns(t)
+
+		require.NoError(t, keyAgent.UnloadKeyRing(s.keyRing.KeyRingIndex))
+		requireNoOpenConns(t)
+	}
+}
+
 // TestLoadKey ensures correct loading of a key into an agent. This test
 // checks the following:
 //   - Loading a key multiple times overwrites the same key.
@@ -812,7 +852,10 @@ func (s *KeyAgentTestSuite) makeKeyRing(t *testing.T, username, proxyHost string
 	}
 }
 
-func startDebugAgent(t *testing.T) error {
+// startDebugAgent serves an in-memory keyring over $SSH_AUTH_SOCK, mimicking
+// the system agent. It returns an atomic integer tracking the number of client
+// connections that are currently open to that agent.
+func startDebugAgent(t *testing.T) (openConns *atomic.Int32, err error) {
 	// Create own tmp dir instead of using t.TmpDir
 	// because net.Listen("unix", path) has dir path length limitation
 	tempDir, err := os.MkdirTemp("", "teleport-test")
@@ -824,18 +867,18 @@ func startDebugAgent(t *testing.T) error {
 	socketpath := filepath.Join(tempDir, "agent.sock")
 	listener, err := net.Listen("unix", socketpath)
 	if err != nil {
-		return trace.Wrap(err)
+		return nil, trace.Wrap(err)
 	}
 
 	systemAgent := agent.NewKeyring()
 	t.Setenv(teleport.SSHAuthSock, socketpath)
 
+	openConns = new(atomic.Int32)
 	startedC := make(chan struct{})
 	doneC := make(chan struct{})
+
 	var wg sync.WaitGroup
-	wg.Add(2)
-	go func() {
-		defer wg.Done()
+	wg.Go(func() {
 		// agent is listening and environment variable is set, unblock now
 		close(startedC)
 		for {
@@ -846,24 +889,22 @@ func startDebugAgent(t *testing.T) error {
 				}
 				return
 			}
-			wg.Add(2)
-			go func() {
+			openConns.Add(1)
+			wg.Go(func() {
 				agent.ServeAgent(systemAgent, conn)
-				wg.Done()
-			}()
-			go func() {
+				openConns.Add(-1)
+			})
+			wg.Go(func() {
 				<-doneC
 				conn.Close()
-				wg.Done()
-			}()
+			})
 		}
-	}()
+	})
 
-	go func() {
+	wg.Go(func() {
 		<-doneC
 		listener.Close()
-		wg.Done()
-	}()
+	})
 
 	t.Cleanup(func() {
 		close(doneC)
@@ -872,7 +913,7 @@ func startDebugAgent(t *testing.T) error {
 
 	// block until agent is started
 	<-startedC
-	return nil
+	return openConns, nil
 }
 
 func (s *KeyAgentTestSuite) newKeyAgent(t *testing.T) *LocalKeyAgent {

lib/sshagent/client.go

@@ -17,6 +17,7 @@
 package sshagent
 
 import (
+	"bytes"
 	"context"
 	"errors"
 	"io"
@@ -89,6 +90,162 @@ func (c *client) Close() error {
 	return trace.Wrap(err)
 }
 
+// NewSingleRequestClient returns a client that opens a new connection for each request
+// and closes it once the request completes.
+//
+// The returned agent is safe for concurrent use as long as getClient is.
+func NewSingleRequestClient(getClient ClientGetter) agent.ExtendedAgent {
+	return singleRequestClient{getClient: getClient}
+}
+
+type singleRequestClient struct {
+	getClient ClientGetter
+}
+
+// withSingleRequestClient opens a new agent connection, runs fn against it,
+// and closes the connection.
+func withSingleRequestClient[T any](getClient ClientGetter, fn func(Client) (T, error)) (T, error) {
+	agentClient, err := getClient()
+	if err != nil {
+		var zero T
+		return zero, trace.Wrap(err)
+	}
+	defer agentClient.Close()
+
+	out, err := fn(agentClient)
+	return out, trace.Wrap(err)
+}
+
+// doWithSingleRequestClient is [withSingleRequestClient] for requests with no return value.
+func doWithSingleRequestClient(getClient ClientGetter, fn func(Client) error) error {
+	_, err := withSingleRequestClient(getClient, func(agentClient Client) (struct{}, error) {
+		return struct{}{}, fn(agentClient)
+	})
+	return trace.Wrap(err)
+}
+
+func (s singleRequestClient) List() ([]*agent.Key, error) {
+	return withSingleRequestClient(s.getClient, func(agentClient Client) ([]*agent.Key, error) {
+		return agentClient.List()
+	})
+}
+
+func (s singleRequestClient) Sign(key ssh.PublicKey, data []byte) (*ssh.Signature, error) {
+	return withSingleRequestClient(s.getClient, func(agentClient Client) (*ssh.Signature, error) {
+		return agentClient.Sign(key, data)
+	})
+}
+
+func (s singleRequestClient) SignWithFlags(key ssh.PublicKey, data []byte, flags agent.SignatureFlags) (*ssh.Signature, error) {
+	return withSingleRequestClient(s.getClient, func(agentClient Client) (*ssh.Signature, error) {
+		return agentClient.SignWithFlags(key, data, flags)
+	})
+}
+
+func (s singleRequestClient) Add(key agent.AddedKey) error {
+	return doWithSingleRequestClient(s.getClient, func(agentClient Client) error {
+		return agentClient.Add(key)
+	})
+}
+
+func (s singleRequestClient) Remove(key ssh.PublicKey) error {
+	return doWithSingleRequestClient(s.getClient, func(agentClient Client) error {
+		return agentClient.Remove(key)
+	})
+}
+
+func (s singleRequestClient) RemoveAll() error {
+	return doWithSingleRequestClient(s.getClient, func(agentClient Client) error {
+		return agentClient.RemoveAll()
+	})
+}
+
+func (s singleRequestClient) Lock(passphrase []byte) error {
+	return doWithSingleRequestClient(s.getClient, func(agentClient Client) error {
+		return agentClient.Lock(passphrase)
+	})
+}
+
+func (s singleRequestClient) Unlock(passphrase []byte) error {
+	return doWithSingleRequestClient(s.getClient, func(agentClient Client) error {
+		return agentClient.Unlock(passphrase)
+	})
+}
+
+func (s singleRequestClient) Extension(extensionType string, contents []byte) ([]byte, error) {
+	return withSingleRequestClient(s.getClient, func(agentClient Client) ([]byte, error) {
+		return agentClient.Extension(extensionType, contents)
+	})
+}
+
+// Signers returns signers for all the keys currently known to the agent.
+// The signers do not hold a connection open, they connect to the agent on demand
+// for each signature.
+func (s singleRequestClient) Signers() ([]ssh.Signer, error) {
+	keys, err := s.List()
+	if err != nil {
+		return nil, trace.Wrap(err)
+	}
+
+	signers := make([]ssh.Signer, 0, len(keys))
+	for _, key := range keys {
+		signers = append(signers, singleRequestSigner{getClient: s.getClient, pub: key})
+	}
+	return signers, nil
+}
+
+// singleRequestSigner is a signer for a key held by an agent that
+// opens a new connection for each signature request.
+type singleRequestSigner struct {
+	getClient ClientGetter
+	pub       ssh.PublicKey
+}
+
+var _ ssh.AlgorithmSigner = singleRequestSigner{}
+
+func (s singleRequestSigner) PublicKey() ssh.PublicKey {
+	return s.pub
+}
+
+func (s singleRequestSigner) Sign(rand io.Reader, data []byte) (*ssh.Signature, error) {
+	// Note: the agent has its own entropy source, so the rand argument is ignored.
+	return withSingleRequestClient(s.getClient, func(agentClient Client) (*ssh.Signature, error) {
+		return agentClient.Sign(s.pub, data)
+	})
+}
+
+func (s singleRequestSigner) SignWithAlgorithm(rand io.Reader, data []byte, algorithm string) (*ssh.Signature, error) {
+	return withSingleRequestClient(s.getClient, func(agentClient Client) (*ssh.Signature, error) {
+		signer, err := agentSigner(agentClient, s.pub)
+		if err != nil {
+			return nil, trace.Wrap(err)
+		}
+		return signer.SignWithAlgorithm(rand, data, algorithm)
+	})
+}
+
+// agentSigner returns the agent's own signer for the given public key.
+func agentSigner(agentClient Client, pub ssh.PublicKey) (ssh.AlgorithmSigner, error) {
+	signers, err := agentClient.Signers()
+	if err != nil {
+		return nil, trace.Wrap(err)
+	}
+
+	pubBytes := pub.Marshal()
+	for _, signer := range signers {
+		if !bytes.Equal(signer.PublicKey().Marshal(), pubBytes) {
+			continue
+		}
+		algorithmSigner, ok := signer.(ssh.AlgorithmSigner)
+		if !ok {
+			return nil, trace.NotImplemented("agent signer of type %T does not support signing with a specific algorithm", signer)
+		}
+		return algorithmSigner, nil
+	}
+
+	return nil, trace.NotFound("agent no longer holds the requested %v key", pub.Type())
+}
+
 const channelType = "auth-agent@openssh.com"
 
 // ServeChannelRequests routes agent channel requests to a new agent