Make agent channel setup lazy.

Changes agent channel setup behavior to be consistent
openssh by having servers lazily request agent channels
when they are needed, rather than immediately starting a
single connection-wide channel as soon as forwarding is
requested.  Fixes an issue introduced in #3613 which
caused openssh clients to hang on exit due to persistent
agent channel.

integration/helpers.go

@@ -1406,8 +1406,11 @@ func createAgent(me *user.User, privateKeyByte []byte, certificateBytes []byte)
 	}
 
 	// create a (unstarted) agent and add the key to it
-	teleAgent := teleagent.NewServer()
-	teleAgent.Add(agentKey)
+	keyring := agent.NewKeyring()
+	keyring.Add(agentKey)
+	teleAgent := teleagent.NewServer(func() (teleagent.Agent, error) {
+		return teleagent.WrapAgent(keyring), nil
+	})
 
 	// start the SSH agent
 	err = teleAgent.ListenUnixSocket(sockPath, uid, gid, 0600)

lib/srv/ctx.go

@@ -27,7 +27,6 @@ import (
 	"time"
 
 	"golang.org/x/crypto/ssh"
-	"golang.org/x/crypto/ssh/agent"
 
 	"github.com/gravitational/teleport"
 	"github.com/gravitational/teleport/lib/auth"
@@ -38,6 +37,7 @@ import (
 	"github.com/gravitational/teleport/lib/services"
 	rsession "github.com/gravitational/teleport/lib/session"
 	"github.com/gravitational/teleport/lib/sshutils"
+	"github.com/gravitational/teleport/lib/teleagent"
 	"github.com/gravitational/teleport/lib/utils"
 
 	"github.com/gravitational/trace"
@@ -293,7 +293,7 @@ func NewServerContext(ccx *sshutils.ConnectionContext, srv Server, identityConte
 		return nil, trace.Wrap(err)
 	}
 
-	cancelContext, cancel := context.WithCancel(context.TODO())
+	cancelContext, cancel := context.WithCancel(ccx)
 
 	ctx := &ServerContext{
 		Parent:            ccx,
@@ -448,24 +448,15 @@ func (c *ServerContext) AddCloser(closer io.Closer) {
 	c.closers = append(c.closers, closer)
 }
 
-// GetAgent returns a agent.Agent which represents the capabilities of an SSH agent,
+// GetAgent returns a teleagent.Agent which represents the capabilities of an SSH agent,
 // or nil if no agent is available in this context.
-func (c *ServerContext) GetAgent() agent.Agent {
+func (c *ServerContext) GetAgent() teleagent.Agent {
 	if c.Parent == nil {
 		return nil
 	}
 	return c.Parent.GetAgent()
 }
 
-// GetAgentChannel returns the channel over which communication with the agent occurs,
-// or nil if no agent is available in this context.
-func (c *ServerContext) GetAgentChannel() ssh.Channel {
-	if c.Parent == nil {
-		return nil
-	}
-	return c.Parent.GetAgentChannel()
-}
-
 // GetTerm returns a Terminal.
 func (c *ServerContext) GetTerm() Terminal {
 	c.RLock()

lib/srv/forward/sshserver.go

@@ -439,7 +439,7 @@ func (s *Server) Serve() {
 	}
 	s.sconn = sconn
 
-	s.connectionContext = sshutils.NewConnectionContext(s.serverConn, s.sconn)
+	s.connectionContext = sshutils.NewConnectionContext(s.closeContext, s.serverConn, s.sconn)
 
 	// Take connection and extract identity information for the user from it.
 	s.identityContext, err = s.authHandlers.CreateIdentityContext(sconn)

lib/srv/regular/proxy.go

@@ -27,14 +27,14 @@ import (
 	"sync"
 
 	"golang.org/x/crypto/ssh"
-	"golang.org/x/crypto/ssh/agent"
 
 	"github.com/gravitational/teleport"
 	"github.com/gravitational/teleport/lib/defaults"
 	"github.com/gravitational/teleport/lib/reversetunnel"
 	"github.com/gravitational/teleport/lib/services"
 	"github.com/gravitational/teleport/lib/srv"
 	"github.com/gravitational/teleport/lib/sshutils"
+	"github.com/gravitational/teleport/lib/teleagent"
 	"github.com/gravitational/teleport/lib/utils"
 
 	"github.com/gravitational/trace"
@@ -46,12 +46,11 @@ import (
 // remote hosts to a proxy client (AKA port mapping)
 type proxySubsys struct {
 	proxySubsysConfig
-	log          *logrus.Entry
-	closeC       chan struct{}
-	error        error
-	closeOnce    sync.Once
-	agent        agent.Agent
-	agentChannel ssh.Channel
+	log       *logrus.Entry
+	closeC    chan struct{}
+	error     error
+	closeOnce sync.Once
+	agent     teleagent.Agent
 }
 
 // parseProxySubsys looks at the requested subsystem name and returns a fully configured
@@ -172,9 +171,8 @@ func newProxySubsys(cfg proxySubsysConfig) (*proxySubsys, error) {
 			trace.Component:       teleport.ComponentSubsystemProxy,
 			trace.ComponentFields: map[string]string{},
 		}),
-		closeC:       make(chan struct{}),
-		agent:        cfg.ctx.GetAgent(),
-		agentChannel: cfg.ctx.GetAgentChannel(),
+		closeC: make(chan struct{}),
+		agent:  cfg.ctx.GetAgent(),
 	}, nil
 }
 
@@ -449,6 +447,9 @@ func (t *proxySubsys) close(err error) {
 	t.closeOnce.Do(func() {
 		t.error = err
 		close(t.closeC)
+		if t.agent != nil {
+			t.agent.Close()
+		}
 	})
 }
 

lib/srv/regular/sshserver.go

@@ -34,7 +34,6 @@ import (
 	"time"
 
 	"golang.org/x/crypto/ssh"
-	"golang.org/x/crypto/ssh/agent"
 
 	"github.com/gravitational/teleport"
 	"github.com/gravitational/teleport/lib/auth"
@@ -775,8 +774,9 @@ func (s *Server) serveAgent(ctx *srv.ServerContext) error {
 		return trace.ConvertSystemError(err)
 	}
 
-	// start an agent on a unix socket
-	agentServer := &teleagent.AgentServer{Agent: ctx.Parent.GetAgent()}
+	// start an agent server on a unix socket.  each incoming connection
+	// will result in a separate agent request.
+	agentServer := teleagent.NewServer(ctx.Parent.StartAgentChannel)
 	err = agentServer.ListenUnixSocket(socketPath, uid, gid, 0600)
 	if err != nil {
 		return trace.Wrap(err)
@@ -788,7 +788,7 @@ func (s *Server) serveAgent(ctx *srv.ServerContext) error {
 	// ensure that SSHAuthSock and SSHAgentPID are imported into
 	// the current child context.
 	ctx.ImportParentEnv()
-	ctx.Debugf("Opened agent channel for Teleport user %v and socket %v.", ctx.Identity.TeleportUser, socketPath)
+	ctx.Debugf("Starting agent server for Teleport user %v and socket %v.", ctx.Identity.TeleportUser, socketPath)
 	go agentServer.Serve()
 
 	return nil
@@ -1184,14 +1184,7 @@ func (s *Server) handleAgentForwardNode(req *ssh.Request, ctx *srv.ServerContext
 		return trace.Wrap(err)
 	}
 
-	// open a channel to the client where the client will serve an agent
-	authChannel, _, err := ctx.Conn.OpenChannel(sshutils.AuthAgentRequest, nil)
-	if err != nil {
-		return trace.Wrap(err)
-	}
-
-	// save the agent in the context so it can be used later
-	ctx.Parent.SetAgent(agent.NewClient(authChannel), authChannel)
+	ctx.Parent.SetForwardAgent(true)
 
 	// serve an agent on a unix socket on this node
 	err = s.serveAgent(ctx)
@@ -1220,16 +1213,7 @@ func (s *Server) handleAgentForwardProxy(req *ssh.Request, ctx *srv.ServerContex
 		return trace.Wrap(err)
 	}
 
-	// Open a channel to the client where the client will serve an agent.
-	authChannel, _, err := ctx.Conn.OpenChannel(sshutils.AuthAgentRequest, nil)
-	if err != nil {
-		return trace.Wrap(err)
-	}
-
-	// Save the agent so it can be used when making a proxy subsystem request
-	// later. It will also be used when building a remote connection to the
-	// target node.
-	ctx.Parent.SetAgent(agent.NewClient(authChannel), authChannel)
+	ctx.Parent.SetForwardAgent(true)
 
 	return nil
 }

lib/sshutils/ctx.go

@@ -17,10 +17,13 @@ limitations under the License.
 package sshutils
 
 import (
+	"context"
 	"io"
 	"net"
 	"sync"
 
+	"github.com/gravitational/teleport/lib/teleagent"
+
 	"golang.org/x/crypto/ssh"
 	"golang.org/x/crypto/ssh/agent"
 
@@ -29,6 +32,7 @@ import (
 
 // ConnectionContext manages connection-level state.
 type ConnectionContext struct {
+	context.Context
 	// NetConn is the base connection object.
 	NetConn net.Conn
 
@@ -42,24 +46,63 @@ type ConnectionContext struct {
 	// set for all channels.
 	env map[string]string
 
-	// agent is a client to remote SSH agent.
-	agent agent.Agent
+	// forwardAgent indicates that agent forwarding has
+	// been requested for this connection.
+	forwardAgent bool
 
-	// agentCh is SSH channel using SSH agent protocol.
-	agentChannel ssh.Channel
 	// closers is a list of io.Closer that will be called when session closes
 	// this is handy as sometimes client closes session, in this case resources
 	// will be properly closed and deallocated, otherwise they could be kept hanging.
 	closers []io.Closer
+
+	// closed indicates that closers have been run.
+	closed bool
+
+	// cancel cancels the context.Context scope associated with this ConnectionContext.
+	cancel context.CancelFunc
 }
 
 // NewConnectionContext creates a new ConnectionContext instance.
-func NewConnectionContext(nconn net.Conn, sconn *ssh.ServerConn) *ConnectionContext {
+func NewConnectionContext(ctx context.Context, nconn net.Conn, sconn *ssh.ServerConn) *ConnectionContext {
+	ctx, cancel := context.WithCancel(ctx)
 	return &ConnectionContext{
+		Context:    ctx,
 		NetConn:    nconn,
 		ServerConn: sconn,
 		env:        make(map[string]string),
+		cancel:     cancel,
+	}
+}
+
+// agentChannel implements the extended teleteleagent.Agent interface,
+// allowing the underlying ssh.Channel to be closed when the agent
+// is no longer needed.
+type agentChannel struct {
+	agent.Agent
+	ch ssh.Channel
+}
+
+func (a *agentChannel) Close() error {
+	return a.ch.Close()
+}
+
+// StartAgentChannel sets up a new agent forwarding channel against this connection.  The channel
+// is automatically closed when either ConnectionContext, or the supplied context.Context
+// gets canceled.
+func (c *ConnectionContext) StartAgentChannel() (teleagent.Agent, error) {
+	// refuse to start an agent if forwardAgent has not yet been set.
+	if !c.getForwardAgent() {
+		return nil, trace.AccessDenied("agent forwarding not requested or not authorized")
+	}
+	// open a agent channel to client
+	ch, _, err := c.ServerConn.OpenChannel(AuthAgentRequest, nil)
+	if err != nil {
+		return nil, trace.Wrap(err)
 	}
+	return &agentChannel{
+		Agent: agent.NewClient(ch),
+		ch:    ch,
+	}, nil
 }
 
 // SetEnv sets a environment variable within this context.
@@ -87,38 +130,44 @@ func (c *ConnectionContext) ExportEnv(m map[string]string) {
 	}
 }
 
-// GetAgent returns a agent.Agent which represents the capabilities of an SSH agent,
-// or nil if no agent is available in this context.
-func (c *ConnectionContext) GetAgent() agent.Agent {
-	c.mu.RLock()
-	defer c.mu.RUnlock()
-	return c.agent
+// GetAgent returns a teleagent.Agent which represents the capabilities of an SSH agent,
+// or nil if no agent is available in this context.  This is a legacy method used in
+// parts of Teleport which still expect a single connection-level agent initialized once
+// when the client initially requests forwarding. Prefer using the
+// StartAgentChannel method directly where possible.
+func (c *ConnectionContext) GetAgent() teleagent.Agent {
+	if c.getForwardAgent() {
+		return teleagent.NewLazy(c.StartAgentChannel)
+	}
+	return nil
 }
 
-// GetAgentChannel returns the channel over which communication with the agent occurs,
-// or nil if no agent is available in this context.
-func (c *ConnectionContext) GetAgentChannel() ssh.Channel {
-	c.mu.RLock()
-	defer c.mu.RUnlock()
-	return c.agentChannel
+// SetForwardAgent configures this context to support agent forwarding.
+// Must not be set until agent forwarding is explicitly requested.
+func (c *ConnectionContext) SetForwardAgent(forwardAgent bool) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	c.forwardAgent = forwardAgent
 }
 
-// SetAgent sets the agent and channel over which communication with the agent occurs.
-func (c *ConnectionContext) SetAgent(a agent.Agent, channel ssh.Channel) {
+// getForwardAgent loads the forwardAgent flag with lock.
+func (c *ConnectionContext) getForwardAgent() bool {
 	c.mu.Lock()
 	defer c.mu.Unlock()
-	if c.agentChannel != nil {
-		c.agentChannel.Close()
-	}
-	c.agentChannel = channel
-	c.agent = a
+	return c.forwardAgent
 }
 
 // AddCloser adds any closer in ctx that will be called
 // when the underlying connection is closed.
 func (c *ConnectionContext) AddCloser(closer io.Closer) {
 	c.mu.Lock()
 	defer c.mu.Unlock()
+	// if context was already closed, run the closer immediately
+	// in the background.
+	if c.closed {
+		go closer.Close()
+		return
+	}
 	c.closers = append(c.closers, closer)
 }
 
@@ -131,17 +180,17 @@ func (c *ConnectionContext) takeClosers() []io.Closer {
 
 	closers := c.closers
 	c.closers = nil
-	if c.agentChannel != nil {
-		closers = append(closers, c.agentChannel)
-		c.agentChannel = nil
-	}
+	c.closed = true
+
 	return closers
 }
 
 // Close closes associated resources (e.g. agent channel).
 func (c *ConnectionContext) Close() error {
 	var errs []error
 
+	c.cancel()
+
 	closers := c.takeClosers()
 
 	for _, cl := range closers {

lib/sshutils/server.go

@@ -437,7 +437,11 @@ func (s *Server) HandleConnection(conn net.Conn) {
 	defer keepAliveTick.Stop()
 	keepAlivePayload := [8]byte{0}
 
-	ccx := NewConnectionContext(wconn, sconn)
+	// NOTE: we deliberately don't use s.closeContext here because the server's
+	// closeContext field is used to trigger starvation on cancellation by halting
+	// the acceptance of new connections; it is not intended to halt in-progress
+	// connection handling, and is therefore orthogonal to the role of ConnectionContext.
+	ccx := NewConnectionContext(context.Background(), wconn, sconn)
 	defer ccx.Close()
 
 	for {
@@ -466,6 +470,8 @@ func (s *Server) HandleConnection(conn net.Conn) {
 			if err != nil {
 				log.Errorf("Failed sending keepalive request: %v", err)
 			}
+		case <-ccx.Done():
+			log.Debugf("Connection context canceled: %v -> %v", conn.RemoteAddr(), conn.LocalAddr())
 		}
 	}
 }