Create session manager for workers to use to interact with sessions (#2235)

internal/daemon/worker/auth_rotation_test.go

@@ -110,9 +110,8 @@ func TestRotationTicking(t *testing.T) {
 		// are valid; if not, we won't establish a new connection and rotation
 		// will fail
 		w.Worker().Resolver().UpdateState(resolver.State{Addresses: []resolver.Address{}})
-		cc, err := w.Worker().GrpcClientConn()
-		require.NoError(err)
-		require.NoError(cc.Close())
+		require.NotNil(w.Worker().GrpcClientConn)
+		require.NoError(w.Worker().GrpcClientConn.Close())
 		require.NoError(w.Worker().StartControllerConnections())
 	}
 }

internal/daemon/worker/controller_connection.go

@@ -197,10 +197,9 @@ func (w *Worker) createClientConn(addr string) error {
 	if err != nil {
 		return fmt.Errorf("error dialing controller for worker auth: %w", err)
 	}
-	w.grpcClientConn.Store(cc)
+	w.GrpcClientConn = cc
 
 	w.controllerStatusConn.Store(pbs.NewServerCoordinationServiceClient(cc))
-	w.controllerSessionConn.Store(pbs.NewSessionServiceClient(cc))
 	w.controllerMultihopConn.Store(multihop.NewMultihopServiceClient(cc))
 	return nil
 }

internal/daemon/worker/handler.go

@@ -20,6 +20,7 @@ import (
 	"github.com/hashicorp/boundary/internal/observability/event"
 	"github.com/hashicorp/boundary/internal/proxy"
 	"github.com/hashicorp/go-secure-stdlib/listenerutil"
+	"google.golang.org/protobuf/types/known/timestamppb"
 	"nhooyr.io/websocket"
 	"nhooyr.io/websocket/wspb"
 )
@@ -30,12 +31,12 @@ type HandlerProperties struct {
 
 // Handler returns a http.Handler for the API. This can be used on
 // its own to mount the Worker API within another web server.
-func (w *Worker) handler(props HandlerProperties) (http.Handler, error) {
+func (w *Worker) handler(props HandlerProperties, sm session.Manager) (http.Handler, error) {
 	const op = "worker.(Worker).handler"
 	// Create the muxer to handle the actual endpoints
 	mux := http.NewServeMux()
 
-	h, err := w.handleProxy(props.ListenerConfig)
+	h, err := w.handleProxy(props.ListenerConfig, sm)
 	if err != nil {
 		return nil, fmt.Errorf("%s: %w", op, err)
 	}
@@ -46,7 +47,7 @@ func (w *Worker) handler(props HandlerProperties) (http.Handler, error) {
 	return metricHandler, nil
 }
 
-func (w *Worker) handleProxy(listenerCfg *listenerutil.ListenerConfig) (http.HandlerFunc, error) {
+func (w *Worker) handleProxy(listenerCfg *listenerutil.ListenerConfig, sessionManager session.Manager) (http.HandlerFunc, error) {
 	const op = "worker.(Worker).handleProxy"
 	if listenerCfg == nil {
 		return nil, fmt.Errorf("%s: missing listener config", op)
@@ -98,21 +99,12 @@ func (w *Worker) handleProxy(listenerCfg *listenerutil.ListenerConfig) (http.Han
 			wr.WriteHeader(http.StatusInternalServerError)
 		}
 
-		siRaw, valid := w.sessionInfoMap.Load(sessionId)
-		if !valid {
-			event.WriteError(ctx, op, errors.New("session not found in info map"), event.WithInfo("session_id", sessionId))
+		sess := sessionManager.Get(sessionId)
+		if sess == nil {
+			event.WriteError(ctx, op, errors.New("session not found locally"), event.WithInfo("session_id", sessionId))
 			wr.WriteHeader(http.StatusInternalServerError)
 			return
 		}
-		si := siRaw.(*session.Info)
-		si.RLock()
-		expiration := si.LookupSessionResponse.GetExpiration()
-		tofuToken := si.LookupSessionResponse.GetTofuToken()
-		version := si.LookupSessionResponse.GetVersion()
-		endpoint := si.LookupSessionResponse.GetEndpoint()
-		credentials := si.LookupSessionResponse.GetCredentials()
-		sessStatus := si.Status
-		si.RUnlock()
 
 		opts := &websocket.AcceptOptions{
 			Subprotocols: []string{globals.TcpProxyV1},
@@ -126,18 +118,9 @@ func (w *Worker) handleProxy(listenerCfg *listenerutil.ListenerConfig) (http.Han
 		// Later calls will cause this to noop if they return a different status
 		defer conn.Close(websocket.StatusNormalClosure, "done")
 
-		connCtx, connCancel := context.WithDeadline(ctx, expiration.AsTime())
+		connCtx, connCancel := context.WithDeadline(ctx, sess.GetExpiration())
 		defer connCancel()
 
-		sessClient, err := w.ControllerSessionConn()
-		if err != nil {
-			event.WriteError(ctx, op, err)
-			if err = conn.Close(websocket.StatusInternalError, "unable to get controller session client"); err != nil {
-				event.WriteError(ctx, op, err, event.WithInfoMsg("error closing client connection"))
-			}
-			return
-		}
-
 		var handshake proxy.ClientHandshake
 		if err := wspb.Read(connCtx, conn, &handshake); err != nil {
 			event.WriteError(ctx, op, err, event.WithInfoMsg("error reading handshake from client"))
@@ -155,8 +138,7 @@ func (w *Worker) handleProxy(listenerCfg *listenerutil.ListenerConfig) (http.Han
 		}
 
 		if handshake.Command == proxy.HANDSHAKECOMMAND_HANDSHAKECOMMAND_SESSION_CANCEL {
-			_, err := session.Cancel(ctx, sessClient, sessionId)
-			if err != nil {
+			if err := sess.RequestCancel(ctx); err != nil {
 				event.WriteError(ctx, op, err, event.WithInfoMsg("unable to cancel session"))
 				if err = conn.Close(websocket.StatusInternalError, "unable to cancel session"); err != nil && !errors.Is(err, io.EOF) {
 					event.WriteError(ctx, op, err, event.WithInfoMsg("error closing client connection"))
@@ -169,8 +151,8 @@ func (w *Worker) handleProxy(listenerCfg *listenerutil.ListenerConfig) (http.Han
 			return
 		}
 
-		if tofuToken != "" {
-			if tofuToken != handshake.GetTofuToken() {
+		if sess.GetTofuToken() != "" {
+			if sess.GetTofuToken() != handshake.GetTofuToken() {
 				event.WriteError(ctx, op, errors.New("WARNING: mismatched tofu token"), event.WithInfo("session_id", sessionId))
 				if err = conn.Close(websocket.StatusPolicyViolation, "tofu token not allowed"); err != nil {
 					event.WriteError(ctx, op, err, event.WithInfoMsg("error closing client connection"))
@@ -178,7 +160,7 @@ func (w *Worker) handleProxy(listenerCfg *listenerutil.ListenerConfig) (http.Han
 				return
 			}
 		} else {
-			if sessStatus != pbs.SESSIONSTATUS_SESSIONSTATUS_PENDING {
+			if sess.GetStatus() != pbs.SESSIONSTATUS_SESSIONSTATUS_PENDING {
 				event.WriteError(ctx, op, errors.New("no tofu token but not in correct session state"), event.WithInfo("session_id", sessionId))
 				if err = conn.Close(websocket.StatusInternalError, "refusing to activate session"); err != nil {
 					event.WriteError(ctx, op, err, event.WithInfoMsg("error closing client connection"))
@@ -186,7 +168,7 @@ func (w *Worker) handleProxy(listenerCfg *listenerutil.ListenerConfig) (http.Han
 				return
 			}
 			if handshake.Command == proxy.HANDSHAKECOMMAND_HANDSHAKECOMMAND_UNSPECIFIED {
-				sessStatus, err = session.Activate(ctx, sessClient, sessionId, handshake.GetTofuToken(), version)
+				err = sess.RequestActivate(ctx, handshake.GetTofuToken())
 				if err != nil {
 					event.WriteError(ctx, op, err, event.WithInfoMsg("unable to validate session"))
 					if err = conn.Close(websocket.StatusInternalError, "unable to activate session"); err != nil {
@@ -198,10 +180,10 @@ func (w *Worker) handleProxy(listenerCfg *listenerutil.ListenerConfig) (http.Han
 			}
 		}
 
-		// Verify the protocol has a supported proxy before calling AuthorizeConnection
-		endpointUrl, err := url.Parse(endpoint)
+		// Verify the protocol has a supported proxy before calling RequestAuthorizeConnection
+		endpointUrl, err := url.Parse(sess.GetEndpoint())
 		if err != nil {
-			event.WriteError(ctx, op, err, event.WithInfoMsg("worker failed to parse target endpoint", "endpoint", endpoint))
+			event.WriteError(ctx, op, err, event.WithInfoMsg("worker failed to parse target endpoint", "endpoint", sess.GetEndpoint()))
 			if err = conn.Close(websocket.StatusProtocolError, "unsupported-protocol"); err != nil {
 				event.WriteError(ctx, op, err, event.WithInfoMsg("error closing client connection"))
 			}
@@ -225,9 +207,9 @@ func (w *Worker) handleProxy(listenerCfg *listenerutil.ListenerConfig) (http.Han
 		}
 		workerId := w.LastStatusSuccess().WorkerId
 
-		var ci *session.ConnInfo
+		var ci session.ConnInfo
 		var connsLeft int32
-		ci, connsLeft, err = session.AuthorizeConnection(ctx, sessClient, workerId, sessionId)
+		ci, connsLeft, err = sess.RequestAuthorizeConnection(ctx, workerId, connCancel)
 		if err != nil {
 			event.WriteError(ctx, op, err, event.WithInfoMsg("unable to authorize connection"))
 			if err = conn.Close(websocket.StatusInternalError, "unable to authorize connection"); err != nil {
@@ -237,22 +219,14 @@ func (w *Worker) handleProxy(listenerCfg *listenerutil.ListenerConfig) (http.Han
 		}
 		event.WriteSysEvent(ctx, op, "connection successfully authorized", "session_id", sessionId, "connection_id", ci.Id)
 		defer func() {
-			if session.CloseConnections(ctx, sessClient, w.sessionInfoMap, map[string]string{ci.Id: si.Id}) {
+			if sessionManager.RequestCloseConnections(ctx, map[string]string{ci.Id: sess.GetId()}) {
 				event.WriteSysEvent(ctx, op, "connection closed", "session_id", sessionId, "connection_id", ci.Id)
 			}
 		}()
 
-		si.Lock()
-		ci.ConnCtx = connCtx
-		ci.ConnCancel = connCancel
-		si.ConnInfoMap[ci.Id] = ci
-		si.Status = sessStatus
-		connectionLimit := si.LookupSessionResponse.GetConnectionLimit()
-		si.Unlock()
-
 		handshakeResult := &proxy.HandshakeResult{
-			Expiration:      expiration,
-			ConnectionLimit: connectionLimit,
+			Expiration:      timestamppb.New(sess.GetExpiration()),
+			ConnectionLimit: sess.GetConnectionLimit(),
 			ConnectionsLeft: connsLeft,
 		}
 		if err := wspb.Write(connCtx, conn, handshakeResult); err != nil {
@@ -267,9 +241,8 @@ func (w *Worker) handleProxy(listenerCfg *listenerutil.ListenerConfig) (http.Han
 			UserClientIp:   net.ParseIP(userClientIp),
 			ClientAddress:  clientAddr,
 			ClientConn:     conn,
-			RemoteEndpoint: endpoint,
-			SessionClient:  sessClient,
-			SessionInfo:    si,
+			RemoteEndpoint: sess.GetEndpoint(),
+			Session:        sess,
 			ConnectionId:   ci.Id,
 		}
 
@@ -281,13 +254,14 @@ func (w *Worker) handleProxy(listenerCfg *listenerutil.ListenerConfig) (http.Han
 			return
 		}
 
+		credentials := sess.GetCredentials()
 		var proxyOpts []proxyHandlers.Option
 		if len(credentials) > 0 {
 			proxyOpts = append(proxyOpts, proxyHandlers.WithEgressCredentials(credentials))
 		}
 
 		if err = handleProxyFn(connCtx, conf, proxyOpts...); err != nil {
-			event.WriteError(ctx, op, err, event.WithInfoMsg("error handling proxy", "session_id", sessionId, "endpoint", endpoint))
+			event.WriteError(ctx, op, err, event.WithInfoMsg("error handling proxy", "session_id", sessionId, "endpoint", sess.GetEndpoint()))
 			if err = conn.Close(websocket.StatusInternalError, "unable to establish proxy"); err != nil {
 				event.WriteError(ctx, op, err, event.WithInfoMsg("error closing client connection"))
 			}

internal/daemon/worker/listeners.go

@@ -14,6 +14,7 @@ import (
 
 	"github.com/hashicorp/boundary/internal/cmd/base"
 	"github.com/hashicorp/boundary/internal/daemon/cluster/handlers"
+	"github.com/hashicorp/boundary/internal/daemon/worker/session"
 	pbs "github.com/hashicorp/boundary/internal/gen/controller/servers/services"
 	"github.com/hashicorp/boundary/internal/observability/event"
 	"github.com/hashicorp/go-multierror"
@@ -26,7 +27,7 @@ import (
 	"google.golang.org/grpc"
 )
 
-func (w *Worker) startListeners() error {
+func (w *Worker) startListeners(sm session.Manager) error {
 	const op = "worker.(Worker).startListeners"
 
 	e := event.SysEventer()
@@ -41,7 +42,7 @@ func (w *Worker) startListeners() error {
 		return fmt.Errorf("%s: nil proxy listener", op)
 	}
 
-	workerServer, err := w.configureForWorker(w.proxyListener, logger)
+	workerServer, err := w.configureForWorker(w.proxyListener, logger, sm)
 	if err != nil {
 		return fmt.Errorf("%s: failed to configure for worker: %w", op, err)
 	}
@@ -51,10 +52,9 @@ func (w *Worker) startListeners() error {
 	return nil
 }
 
-func (w *Worker) configureForWorker(ln *base.ServerListener, logger *log.Logger) (func(), error) {
+func (w *Worker) configureForWorker(ln *base.ServerListener, logger *log.Logger, sessionManager session.Manager) (func(), error) {
 	const op = "worker.configureForWorker"
-
-	handler, err := w.handler(HandlerProperties{ListenerConfig: ln.Config})
+	handler, err := w.handler(HandlerProperties{ListenerConfig: ln.Config}, sessionManager)
 	if err != nil {
 		return nil, err
 	}
@@ -126,7 +126,7 @@ func (w *Worker) configureForWorker(ln *base.ServerListener, logger *log.Logger)
 			Storage:      w.WorkerAuthStorage,
 			BaseListener: ln.ProxyListener,
 			BaseTlsConfiguration: &tls.Config{
-				GetConfigForClient: w.getSessionTls,
+				GetConfigForClient: w.getSessionTls(sessionManager),
 			},
 			FetchCredsFunc:                 fetchCredsFn,
 			GenerateServerCertificatesFunc: generateServerCertificatesFn,
@@ -150,7 +150,7 @@ func (w *Worker) configureForWorker(ln *base.ServerListener, logger *log.Logger)
 		return nil, fmt.Errorf("%s: error creating multihop service handler: %w", op, err)
 	}
 	multihop.RegisterMultihopServiceServer(downstreamServer, multihopService)
-	statusSessionService := NewWorkerProxyServiceServer(w.controllerStatusConn, w.controllerSessionConn)
+	statusSessionService := NewWorkerProxyServiceServer(w.GrpcClientConn, w.controllerStatusConn)
 	pbs.RegisterServerCoordinationServiceServer(downstreamServer, statusSessionService)
 	pbs.RegisterSessionServiceServer(downstreamServer, statusSessionService)
 

internal/daemon/worker/listeners_test.go

@@ -9,6 +9,8 @@ import (
 	"testing"
 
 	"github.com/hashicorp/boundary/internal/cmd/base"
+	"github.com/hashicorp/boundary/internal/daemon/worker/session"
+	pbs "github.com/hashicorp/boundary/internal/gen/controller/servers/services"
 	"github.com/hashicorp/go-secure-stdlib/listenerutil"
 	"github.com/stretchr/testify/require"
 )
@@ -76,7 +78,9 @@ func TestStartListeners(t *testing.T) {
 			require.NoError(t, err)
 			w.baseContext = context.Background()
 
-			err = w.startListeners()
+			manager, err := session.NewManager(pbs.NewSessionServiceClient(w.GrpcClientConn))
+			require.NoError(t, err)
+			err = w.startListeners(manager)
 			if tt.expErr {
 				require.EqualError(t, err, tt.expErrMsg)
 				return

internal/daemon/worker/proxy/proxy.go

@@ -7,7 +7,6 @@ import (
 	"sync"
 
 	"github.com/hashicorp/boundary/internal/daemon/worker/session"
-	pbs "github.com/hashicorp/boundary/internal/gen/controller/servers/services"
 	"nhooyr.io/websocket"
 )
 
@@ -24,9 +23,8 @@ type Config struct {
 	ClientConn     *websocket.Conn
 	RemoteEndpoint string
 
-	SessionClient pbs.SessionServiceClient
-	SessionInfo   *session.Info
-	ConnectionId  string
+	Session      session.Session
+	ConnectionId string
 }
 
 // Validate checks that the provided config is valid. If invalid, an error is returned
@@ -39,10 +37,8 @@ func (c Config) Validate() error {
 		return errors.New("missing client connection")
 	case c.RemoteEndpoint == "":
 		return errors.New("missing remote endpoint")
-	case c.SessionClient == nil:
-		return errors.New("missing session client")
-	case c.SessionInfo == nil:
-		return errors.New("missing session info")
+	case c.Session == nil:
+		return errors.New("missing session")
 	case c.ConnectionId == "":
 		return errors.New("missing connection id")
 	default:

internal/daemon/worker/proxy/proxy_test.go

@@ -6,12 +6,15 @@ import (
 	"testing"
 
 	"github.com/hashicorp/boundary/internal/daemon/worker/session"
-	pbs "github.com/hashicorp/boundary/internal/gen/controller/servers/services"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"nhooyr.io/websocket"
 )
 
+type testSession struct {
+	session.Session
+}
+
 func TestConfigValidate(t *testing.T) {
 	t.Parallel()
 
@@ -20,10 +23,7 @@ func TestConfigValidate(t *testing.T) {
 		IP:   net.ParseIP("127.0.0.1"),
 		Port: 50000,
 	}
-	sessClient := pbs.NewMockSessionServiceClient()
-	si := &session.Info{
-		Id: "one",
-	}
+	si := testSession{}
 	conn := &websocket.Conn{}
 
 	tests := []struct {
@@ -37,8 +37,7 @@ func TestConfigValidate(t *testing.T) {
 			conf: Config{
 				ClientConn:     conn,
 				RemoteEndpoint: "tcp://remote",
-				SessionClient:  sessClient,
-				SessionInfo:    si,
+				Session:        si,
 				ConnectionId:   "connection-id",
 			},
 			wantErr:    true,
@@ -49,8 +48,7 @@ func TestConfigValidate(t *testing.T) {
 			conf: Config{
 				ClientAddress:  clientAddr,
 				RemoteEndpoint: "tcp://remote",
-				SessionClient:  sessClient,
-				SessionInfo:    si,
+				Session:        si,
 				ConnectionId:   "connection-id",
 			},
 			wantErr:    true,
@@ -61,36 +59,22 @@ func TestConfigValidate(t *testing.T) {
 			conf: Config{
 				ClientAddress: clientAddr,
 				ClientConn:    conn,
-				SessionClient: sessClient,
-				SessionInfo:   si,
+				Session:       si,
 				ConnectionId:  "connection-id",
 			},
 			wantErr:    true,
 			wantErrMsg: "missing remote endpoint",
 		},
 		{
-			name: "missing-session-client",
+			name: "missing-session",
 			conf: Config{
 				ClientAddress:  clientAddr,
 				ClientConn:     conn,
 				RemoteEndpoint: "tcp://remote",
-				SessionInfo:    si,
 				ConnectionId:   "connection-id",
 			},
 			wantErr:    true,
-			wantErrMsg: "missing session client",
-		},
-		{
-			name: "missing-session-info",
-			conf: Config{
-				ClientAddress:  clientAddr,
-				ClientConn:     conn,
-				RemoteEndpoint: "tcp://remote",
-				SessionClient:  sessClient,
-				ConnectionId:   "connection-id",
-			},
-			wantErr:    true,
-			wantErrMsg: "missing session info",
+			wantErrMsg: "missing session",
 		},
 		{
 			name: "missing-connection-id",
@@ -98,8 +82,7 @@ func TestConfigValidate(t *testing.T) {
 				ClientAddress:  clientAddr,
 				ClientConn:     conn,
 				RemoteEndpoint: "tcp://remote",
-				SessionClient:  sessClient,
-				SessionInfo:    si,
+				Session:        si,
 			},
 			wantErr:    true,
 			wantErrMsg: "missing connection id",
@@ -110,8 +93,7 @@ func TestConfigValidate(t *testing.T) {
 				ClientAddress:  clientAddr,
 				ClientConn:     conn,
 				RemoteEndpoint: "tcp://remote",
-				SessionClient:  sessClient,
-				SessionInfo:    si,
+				Session:        si,
 				ConnectionId:   "connection-id",
 			},
 		},

internal/daemon/worker/proxy/tcp/tcp.go

@@ -9,7 +9,6 @@ import (
 	"sync"
 
 	"github.com/hashicorp/boundary/internal/daemon/worker/proxy"
-	"github.com/hashicorp/boundary/internal/daemon/worker/session"
 	pbs "github.com/hashicorp/boundary/internal/gen/controller/servers/services"
 	"nhooyr.io/websocket"
 )
@@ -56,16 +55,10 @@ func handleProxy(ctx context.Context, conf proxy.Config, _ ...proxy.Option) erro
 		UserClientIp:       conf.UserClientIp.String(),
 	}
 
-	connStatus, err := session.ConnectConnection(ctx, conf.SessionClient, connectionInfo)
-	if err != nil {
+	if err := conf.Session.RequestConnectConnection(ctx, connectionInfo); err != nil {
 		return fmt.Errorf("error marking connection as connected: %w", err)
 	}
 
-	// Update connection info to set connection status
-	conf.SessionInfo.Lock()
-	conf.SessionInfo.ConnInfoMap[conf.ConnectionId].Status = connStatus
-	conf.SessionInfo.Unlock()
-
 	// Get a wrapped net.Conn so we can use io.Copy
 	netConn := websocket.NetConn(ctx, conn, websocket.MessageBinary)
 

internal/daemon/worker/proxy/tcp/tcp_test.go

@@ -2,9 +2,14 @@ package tcp
 
 import (
 	"context"
+	"crypto/ed25519"
+	"crypto/rand"
+	"crypto/x509"
 	"fmt"
+	"math/big"
 	"net"
 	"testing"
+	"time"
 
 	"github.com/hashicorp/boundary/internal/daemon/worker/proxy"
 	"github.com/hashicorp/boundary/internal/daemon/worker/session"
@@ -13,6 +18,7 @@ import (
 	"github.com/hashicorp/boundary/sdk/testutil"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	"google.golang.org/protobuf/types/known/timestamppb"
 	"nhooyr.io/websocket"
 )
 
@@ -49,24 +55,41 @@ func TestHandleTcpProxyV1(t *testing.T) {
 		Port: 50000,
 	}
 	sessClient := pbs.NewMockSessionServiceClient()
-	si := &session.Info{
-		Id: "one",
-		LookupSessionResponse: &pbs.LookupSessionResponse{
+	sessClient.LookupSessionFn = func(_ context.Context, request *pbs.LookupSessionRequest) (*pbs.LookupSessionResponse, error) {
+		cert, _, _ := createTestCert(t)
+		return &pbs.LookupSessionResponse{
 			Authorization: &targets.SessionAuthorizationData{
-				SessionId: "mock-session",
+				SessionId:   request.GetSessionId(),
+				Certificate: cert,
 			},
-		},
-		ConnInfoMap: map[string]*session.ConnInfo{
-			"mock-connection": {},
-		},
+			Expiration: timestamppb.New(time.Now().Add(time.Hour)),
+		}, nil
 	}
+	sessClient.AuthorizeConnectionFn = func(_ context.Context, req *pbs.AuthorizeConnectionRequest) (*pbs.AuthorizeConnectionResponse, error) {
+		return &pbs.AuthorizeConnectionResponse{
+			ConnectionId:    "mock-connection",
+			Status:          pbs.CONNECTIONSTATUS_CONNECTIONSTATUS_AUTHORIZED,
+			ConnectionsLeft: -1,
+		}, nil
+	}
+	sessClient.ConnectConnectionFn = func(_ context.Context, _ *pbs.ConnectConnectionRequest) (*pbs.ConnectConnectionResponse, error) {
+		return &pbs.ConnectConnectionResponse{
+			Status: pbs.CONNECTIONSTATUS_CONNECTIONSTATUS_CONNECTED,
+		}, nil
+	}
+	manager, err := session.NewManager(sessClient)
+	require.NoError(err)
+	s, err := manager.LoadLocalSession(ctx, "one", "workerid")
+	require.NoError(err)
+	_, connCancelFn := context.WithCancel(context.Background())
+	_, _, err = s.RequestAuthorizeConnection(ctx, "workerid", connCancelFn)
+	require.NoError(err)
 
 	conf := proxy.Config{
 		ClientAddress:  clientAddr,
 		ClientConn:     proxyConn,
 		RemoteEndpoint: fmt.Sprintf("tcp://localhost:%d", port),
-		SessionClient:  sessClient,
-		SessionInfo:    si,
+		Session:        s,
 		ConnectionId:   "mock-connection",
 		UserClientIp:   net.ParseIP("127.0.0.1"),
 	}
@@ -109,3 +132,25 @@ func TestHandleTcpProxyV1(t *testing.T) {
 
 	cancelCtx()
 }
+
+func createTestCert(t *testing.T) ([]byte, ed25519.PublicKey, ed25519.PrivateKey) {
+	pub, priv, err := ed25519.GenerateKey(rand.Reader)
+	require.NoError(t, err)
+
+	template := &x509.Certificate{
+		ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth, x509.ExtKeyUsageClientAuth},
+		KeyUsage:              x509.KeyUsageDigitalSignature | x509.KeyUsageKeyEncipherment | x509.KeyUsageKeyAgreement | x509.KeyUsageCertSign,
+		SerialNumber:          big.NewInt(0),
+		NotBefore:             time.Now().Add(-30 * time.Second),
+		NotAfter:              time.Now().Add(5 * time.Minute),
+		BasicConstraintsValid: true,
+		IsCA:                  true,
+		IPAddresses:           []net.IP{net.ParseIP("127.0.0.1")},
+		DNSNames:              []string{"/tmp/boundary-opslistener-test0.sock", "/tmp/boundary-opslistener-test1.sock"},
+	}
+
+	certBytes, err := x509.CreateCertificate(rand.Reader, template, template, pub, priv)
+	require.NoError(t, err)
+
+	return certBytes, pub, priv
+}

internal/daemon/worker/session/manager.go

@@ -0,0 +1,139 @@
+package session
+
+import (
+	"context"
+	"fmt"
+	"reflect"
+	"sync"
+
+	"github.com/hashicorp/boundary/internal/errors"
+	pbs "github.com/hashicorp/boundary/internal/gen/controller/servers/services"
+)
+
+// Manager stores session information locally and exposes ways to operate on the
+// set of sessions locally in batch.
+// This is thread-safe.
+type Manager interface {
+	Get(string) Session
+
+	// ForEachLocalSession calls the provided function with each Session object.
+	// If the provided function ever returns false the iteration stops.
+	// If changes to the sessions in the manager happen concurrently, this function
+	// does not guarantee that the key will or will not be provided in the
+	// iteration.
+	ForEachLocalSession(func(Session) bool)
+
+	// LoadLocalSession looks up from the source of truth the session information,
+	// validates it is valid, and then refreshes the local manager's data.
+	// On a local worker, the only value of a Session that we care about that can
+	// be modified is the Status.  Because of that, if LoadLocalSession is called on
+	// a Session that is already in the manager's data, only the Status is updated.
+	LoadLocalSession(ctx context.Context, id string, workerId string) (Session, error)
+
+	// DeleteLocalSession removes all sessions with the provided id from the
+	// local manager.  If ids are passed in which do not exist in the manager
+	// no error is returned.
+	DeleteLocalSession([]string)
+
+	// RequestCloseConnections sends connection close requests to the controller,
+	// and sets close times within the worker. It should be called during the worker
+	// status loop and on connection exit on the proxy.
+	//
+	// The boolean indicates whether the function was successful, e.g. had any
+	// errors. Individual events will be sent for the errors if there are any.
+	//
+	// closeInfo is a map of connection ids mapped to their individual session ids.
+	RequestCloseConnections(context.Context, map[string]string) bool
+}
+
+type manager struct {
+	controllerSessionConn pbs.SessionServiceClient
+	sessionMap            sync.Map
+}
+
+var _ Manager = (*manager)(nil)
+
+// NewManager returns a *Manager which uses the provided ServiceServiceClient to
+// perform actions on Sessions and Connections on the Controller.
+func NewManager(client pbs.SessionServiceClient) (*manager, error) {
+	if isNil(client) {
+		return nil, fmt.Errorf("SessionServiceClient is nil")
+	}
+	return &manager{
+		controllerSessionConn: client,
+	}, nil
+}
+
+func (m *manager) Get(id string) Session {
+	if s, ok := m.sessionMap.Load(id); ok {
+		return s.(*sess)
+	}
+	return nil
+}
+
+func (m *manager) ForEachLocalSession(f func(Session) bool) {
+	// TODO: Periodically clean this up. We can't rely on things in here but
+	// not in cancellation because they could be on the way to being
+	// established. However, since cert lifetimes are short, we can simply range
+	// through and remove values that are expired.
+	m.sessionMap.Range(func(_, value any) bool {
+		s, ok := value.(Session)
+		if !ok {
+			return false
+		}
+		return f(s)
+	})
+}
+
+func (m *manager) LoadLocalSession(ctx context.Context, id string, workerId string) (Session, error) {
+	const op = "session.(*manager).LoadLocalSession"
+	switch {
+	case id == "":
+		return nil, errors.New(ctx, errors.InvalidParameter, op, "id is not set")
+	case workerId == "":
+		return nil, errors.New(ctx, errors.InvalidParameter, op, "workerId is not set")
+	}
+
+	resp, err := m.controllerSessionConn.LookupSession(ctx, &pbs.LookupSessionRequest{
+		SessionId: id,
+		WorkerId:  workerId,
+	})
+	if err != nil {
+		return nil, err
+	}
+
+	s, err := newSess(m.controllerSessionConn, resp)
+	if err != nil {
+		return nil, err
+	}
+
+	actualSessRaw, loaded := m.sessionMap.LoadOrStore(s.GetId(), s)
+	if !loaded {
+		return s, nil
+	}
+	// Update the response to the latest
+	actualSess := actualSessRaw.(*sess)
+	actualSess.ApplyLocalStatus(s.GetStatus())
+	return actualSess, nil
+}
+
+func (m *manager) DeleteLocalSession(sessIds []string) {
+	for _, s := range sessIds {
+		m.sessionMap.Delete(s)
+	}
+}
+
+func (m *manager) RequestCloseConnections(ctx context.Context, closeInfo map[string]string) bool {
+	return closeConnections(ctx, m.controllerSessionConn, m, closeInfo)
+}
+
+func isNil(i interface{}) bool {
+	if i == nil {
+		return true
+	}
+	switch reflect.TypeOf(i).Kind() {
+	case reflect.Ptr, reflect.Map, reflect.Array, reflect.Chan, reflect.Slice:
+		return reflect.ValueOf(i).IsNil()
+	}
+	return false
+}

internal/daemon/worker/session/manager_test.go

@@ -0,0 +1,474 @@
+package session
+
+import (
+	"context"
+	"crypto/ed25519"
+	"crypto/rand"
+	"crypto/x509"
+	"errors"
+	"fmt"
+	"math/big"
+	"net"
+	"sync"
+	"testing"
+	"time"
+
+	pbs "github.com/hashicorp/boundary/internal/gen/controller/servers/services"
+	"github.com/hashicorp/boundary/sdk/pbs/controller/api/resources/targets"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"google.golang.org/protobuf/types/known/timestamppb"
+)
+
+func TestManager_Get(t *testing.T) {
+	mockSessionClient := pbs.NewMockSessionServiceClient()
+	mockSessionClient.LookupSessionFn = func(context.Context, *pbs.LookupSessionRequest) (*pbs.LookupSessionResponse, error) {
+		return &pbs.LookupSessionResponse{
+			Authorization: &targets.SessionAuthorizationData{
+				SessionId:   "foo",
+				Certificate: createTestCert(t),
+			},
+			Version:    1,
+			Expiration: timestamppb.New(time.Now().Add(time.Hour)),
+			Status:     pbs.SESSIONSTATUS_SESSIONSTATUS_PENDING,
+		}, nil
+	}
+	manager, err := NewManager(mockSessionClient)
+	require.NoError(t, err)
+	_, err = manager.LoadLocalSession(context.Background(), "foo", "worker id")
+	require.NoError(t, err)
+
+	assert.NotNil(t, manager.Get("foo"))
+	assert.Nil(t, manager.Get("unknown"))
+}
+
+func TestManager_DeleteLocalSession(t *testing.T) {
+	mockSessionClient := pbs.NewMockSessionServiceClient()
+	mockSessionClient.LookupSessionFn = func(_ context.Context, req *pbs.LookupSessionRequest) (*pbs.LookupSessionResponse, error) {
+		return &pbs.LookupSessionResponse{
+			Authorization: &targets.SessionAuthorizationData{
+				SessionId:   req.GetSessionId(),
+				Certificate: createTestCert(t),
+			},
+			Version:    1,
+			Expiration: timestamppb.New(time.Now().Add(time.Hour)),
+			Status:     pbs.SESSIONSTATUS_SESSIONSTATUS_PENDING,
+		}, nil
+	}
+	manager, err := NewManager(mockSessionClient)
+	require.NoError(t, err)
+	_, err = manager.LoadLocalSession(context.Background(), "foo", "worker id")
+	require.NoError(t, err)
+
+	require.NotNil(t, manager.Get("foo"))
+	manager.DeleteLocalSession([]string{"foo"})
+	assert.Nil(t, manager.Get("foo"))
+	// A second call to DeleteLocalSession is ok.
+	manager.DeleteLocalSession([]string{"foo"})
+	assert.Nil(t, manager.Get("foo"))
+}
+
+func TestManager_RequestCloseConnections(t *testing.T) {
+	ctx := context.Background()
+	mockSessionClient := pbs.NewMockSessionServiceClient()
+
+	manager, err := NewManager(mockSessionClient)
+	require.NoError(t, err)
+	assert.False(t, manager.RequestCloseConnections(ctx, nil))
+	assert.False(t, manager.RequestCloseConnections(ctx, map[string]string{}))
+
+	mockSessionClient.LookupSessionFn = func(_ context.Context, req *pbs.LookupSessionRequest) (*pbs.LookupSessionResponse, error) {
+		return &pbs.LookupSessionResponse{
+			Authorization: &targets.SessionAuthorizationData{
+				SessionId:   req.GetSessionId(),
+				Certificate: createTestCert(t),
+			},
+			Version:    1,
+			Expiration: timestamppb.New(time.Now().Add(time.Hour)),
+			Status:     pbs.SESSIONSTATUS_SESSIONSTATUS_PENDING,
+		}, nil
+	}
+	session1, err := manager.LoadLocalSession(ctx, "sess1", "worker id")
+	require.NoError(t, err)
+	session2, err := manager.LoadLocalSession(ctx, "sess2", "worker id")
+	require.NoError(t, err)
+	session3, err := manager.LoadLocalSession(ctx, "sess3", "worker id")
+	require.NoError(t, err)
+
+	mockSessionClient.CloseConnectionFn = func(context.Context, *pbs.CloseConnectionRequest) (*pbs.CloseConnectionResponse, error) {
+		return nil, errors.New("error")
+	}
+	assert.False(t, manager.RequestCloseConnections(ctx, map[string]string{"connection id": session1.GetId()}))
+	mockSessionClient.CloseConnectionFn = func(_ context.Context, req *pbs.CloseConnectionRequest) (*pbs.CloseConnectionResponse, error) {
+		var data []*pbs.CloseConnectionResponseData
+		for _, r := range req.GetCloseRequestData() {
+			data = append(data, &pbs.CloseConnectionResponseData{
+				ConnectionId: r.GetConnectionId(),
+				Status:       pbs.CONNECTIONSTATUS_CONNECTIONSTATUS_CLOSED,
+			})
+		}
+		return &pbs.CloseConnectionResponse{CloseResponseData: data}, nil
+	}
+	// There is no connection information yet for this connection.
+	assert.False(t, manager.RequestCloseConnections(ctx, map[string]string{
+		"random_connection_id": session1.GetId(),
+	}))
+
+	// Load the connection info into the local storage
+	mockSessionClient.AuthorizeConnectionFn = func(_ context.Context, req *pbs.AuthorizeConnectionRequest) (*pbs.AuthorizeConnectionResponse, error) {
+		return &pbs.AuthorizeConnectionResponse{
+			ConnectionId:    fmt.Sprintf("connection_%s", req.GetSessionId()),
+			Status:          pbs.CONNECTIONSTATUS_CONNECTIONSTATUS_AUTHORIZED,
+			ConnectionsLeft: -1,
+		}, nil
+	}
+	_, cancelFn := context.WithCancel(ctx)
+	c1, _, err := session1.RequestAuthorizeConnection(ctx, "worker id", cancelFn)
+	require.NoError(t, err)
+
+	assert.True(t, manager.RequestCloseConnections(ctx, map[string]string{
+		c1.Id: session1.GetId(),
+	}))
+
+	c2, _, err := session2.RequestAuthorizeConnection(ctx, "worker id", cancelFn)
+	require.NoError(t, err)
+	c3, _, err := session3.RequestAuthorizeConnection(ctx, "worker id", cancelFn)
+	require.NoError(t, err)
+	assert.True(t, manager.RequestCloseConnections(ctx, map[string]string{
+		c2.Id: session2.GetId(),
+		c3.Id: session3.GetId(),
+	}))
+}
+
+func TestManager_LoadLocalSession(t *testing.T) {
+	mockSessionClient := pbs.NewMockSessionServiceClient()
+	errorCases := []struct {
+		name               string
+		inId               string
+		inWorkerId         string
+		controllerResponse func() (*pbs.LookupSessionResponse, error)
+		wantError          string
+	}{
+		{
+			name:       "no id",
+			inWorkerId: "worker id",
+			wantError:  "id is not set",
+		},
+		{
+			name:      "no worker id",
+			inId:      "id",
+			wantError: "workerId is not set",
+		},
+		{
+			name:       "controller error",
+			inId:       "id",
+			inWorkerId: "worker id",
+			controllerResponse: func() (*pbs.LookupSessionResponse, error) {
+				return nil, errors.New("some error")
+			},
+			wantError: "some error",
+		},
+		{
+			name:       "cant parse cert",
+			inId:       "id",
+			inWorkerId: "worker id",
+			controllerResponse: func() (*pbs.LookupSessionResponse, error) {
+				return &pbs.LookupSessionResponse{
+					Authorization: &targets.SessionAuthorizationData{
+						SessionId:   "foo",
+						Certificate: []byte("unparseable"),
+					},
+					Version:    1,
+					Expiration: timestamppb.New(time.Now().Add(time.Hour)),
+					Status:     pbs.SESSIONSTATUS_SESSIONSTATUS_PENDING,
+				}, nil
+			},
+			wantError: "error parsing session certificate",
+		},
+		{
+			name:       "expired session",
+			inId:       "id",
+			inWorkerId: "worker id",
+			controllerResponse: func() (*pbs.LookupSessionResponse, error) {
+				return &pbs.LookupSessionResponse{
+					Authorization: &targets.SessionAuthorizationData{
+						SessionId:   "foo",
+						Certificate: createTestCert(t),
+					},
+					Version:    1,
+					Expiration: timestamppb.New(time.Now().Add(-time.Hour)),
+					Status:     pbs.SESSIONSTATUS_SESSIONSTATUS_PENDING,
+				}, nil
+			},
+			wantError: "session is expired",
+		},
+	}
+	for _, tc := range errorCases {
+		t.Run(tc.name, func(t *testing.T) {
+			mockSessionClient.LookupSessionFn = func(context.Context, *pbs.LookupSessionRequest) (*pbs.LookupSessionResponse, error) {
+				return tc.controllerResponse()
+			}
+			manager, err := NewManager(mockSessionClient)
+			require.NoError(t, err)
+			s, err := manager.LoadLocalSession(context.Background(), tc.inId, tc.inWorkerId)
+			require.Error(t, err)
+			assert.Nil(t, s)
+
+			assert.ErrorContains(t, err, tc.wantError)
+		})
+	}
+
+	t.Run("success", func(t *testing.T) {
+		expirationTime := time.Now().Add(time.Hour).UTC()
+		mockSessionClient.LookupSessionFn = func(context.Context, *pbs.LookupSessionRequest) (*pbs.LookupSessionResponse, error) {
+			return &pbs.LookupSessionResponse{
+				Authorization: &targets.SessionAuthorizationData{
+					SessionId:   "foo",
+					Certificate: createTestCert(t),
+				},
+				ConnectionLimit: -1,
+				Version:         1,
+				Expiration:      timestamppb.New(expirationTime),
+				Status:          pbs.SESSIONSTATUS_SESSIONSTATUS_PENDING,
+			}, nil
+		}
+		manager, err := NewManager(mockSessionClient)
+		require.NoError(t, err)
+		s, err := manager.LoadLocalSession(context.Background(), "foo", "worker id")
+		require.NoError(t, err)
+		assert.Equal(t, "foo", s.GetId())
+		assert.Equal(t, pbs.SESSIONSTATUS_SESSIONSTATUS_PENDING, s.GetStatus())
+		assert.Empty(t, s.GetLocalConnections())
+		assert.Equal(t, int32(-1), s.GetConnectionLimit())
+		assert.Equal(t, expirationTime, s.GetExpiration())
+	})
+}
+
+func createTestCert(t *testing.T) []byte {
+	pub, priv, err := ed25519.GenerateKey(rand.Reader)
+	require.NoError(t, err)
+
+	template := &x509.Certificate{
+		ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth, x509.ExtKeyUsageClientAuth},
+		KeyUsage:              x509.KeyUsageDigitalSignature | x509.KeyUsageKeyEncipherment | x509.KeyUsageKeyAgreement | x509.KeyUsageCertSign,
+		SerialNumber:          big.NewInt(0),
+		NotBefore:             time.Now().Add(-30 * time.Second),
+		NotAfter:              time.Now().Add(5 * time.Minute),
+		BasicConstraintsValid: true,
+		IsCA:                  true,
+		IPAddresses:           []net.IP{net.ParseIP("127.0.0.1")},
+		DNSNames:              []string{"/tmp/boundary-opslistener-test0.sock", "/tmp/boundary-opslistener-test1.sock"},
+	}
+
+	certBytes, err := x509.CreateCertificate(rand.Reader, template, template, pub, priv)
+	require.NoError(t, err)
+
+	return certBytes
+}
+
+func TestWorkerSetCloseTimeForResponse(t *testing.T) {
+	cases := []struct {
+		name             string
+		sessionCloseInfo map[string][]*pbs.CloseConnectionResponseData
+		sessionInfoMap   func() sync.Map
+		expected         []string
+		expectedClosed   map[string]struct{}
+		expectedErr      []error
+	}{
+		{
+			name: "basic",
+			sessionCloseInfo: map[string][]*pbs.CloseConnectionResponseData{
+				"one": {
+					{ConnectionId: "foo", Status: pbs.CONNECTIONSTATUS_CONNECTIONSTATUS_CLOSED},
+				},
+				"two": {
+					{ConnectionId: "bar", Status: pbs.CONNECTIONSTATUS_CONNECTIONSTATUS_CLOSED},
+				},
+			},
+			sessionInfoMap: func() sync.Map {
+				var m sync.Map
+				m.Store("one", &sess{
+					resp: &pbs.LookupSessionResponse{Authorization: &targets.SessionAuthorizationData{
+						SessionId: "one",
+					}},
+					connInfoMap: map[string]*ConnInfo{
+						"foo": {Id: "foo"},
+					},
+				})
+				m.Store("two", &sess{
+					resp: &pbs.LookupSessionResponse{Authorization: &targets.SessionAuthorizationData{
+						SessionId: "two",
+					}},
+					connInfoMap: map[string]*ConnInfo{
+						"bar": {Id: "bar"},
+					},
+				})
+				m.Store("three", &sess{
+					resp: &pbs.LookupSessionResponse{Authorization: &targets.SessionAuthorizationData{
+						SessionId: "three",
+					}},
+					connInfoMap: map[string]*ConnInfo{
+						"baz": {Id: "baz"},
+					},
+				})
+
+				return m
+			},
+			expected: []string{"foo", "bar"},
+			expectedClosed: map[string]struct{}{
+				"foo": {},
+				"bar": {},
+			},
+		},
+		{
+			name: "not closed",
+			sessionCloseInfo: map[string][]*pbs.CloseConnectionResponseData{
+				"one": {
+					{ConnectionId: "foo", Status: pbs.CONNECTIONSTATUS_CONNECTIONSTATUS_CLOSED},
+				},
+				"two": {
+					{ConnectionId: "bar", Status: pbs.CONNECTIONSTATUS_CONNECTIONSTATUS_CONNECTED},
+				},
+			},
+			sessionInfoMap: func() sync.Map {
+				var m sync.Map
+				m.Store("one", &sess{
+					resp: &pbs.LookupSessionResponse{Authorization: &targets.SessionAuthorizationData{
+						SessionId: "one",
+					}},
+					connInfoMap: map[string]*ConnInfo{
+						"foo": {Id: "foo"},
+					},
+				})
+				m.Store("two", &sess{
+					resp: &pbs.LookupSessionResponse{Authorization: &targets.SessionAuthorizationData{
+						SessionId: "two",
+					}},
+					connInfoMap: map[string]*ConnInfo{
+						"bar": {Id: "bar"},
+					},
+				})
+
+				return m
+			},
+			expected: []string{"foo"},
+			expectedClosed: map[string]struct{}{
+				"foo": {},
+			},
+		},
+		{
+			name: "missing session",
+			sessionCloseInfo: map[string][]*pbs.CloseConnectionResponseData{
+				"one": {
+					{ConnectionId: "foo", Status: pbs.CONNECTIONSTATUS_CONNECTIONSTATUS_CLOSED},
+				},
+				"two": {
+					{ConnectionId: "bar", Status: pbs.CONNECTIONSTATUS_CONNECTIONSTATUS_CLOSED},
+				},
+			},
+			sessionInfoMap: func() sync.Map {
+				var m sync.Map
+				m.Store("one", &sess{
+					resp: &pbs.LookupSessionResponse{Authorization: &targets.SessionAuthorizationData{
+						SessionId: "one",
+					}},
+					connInfoMap: map[string]*ConnInfo{
+						"foo": {Id: "foo"},
+					},
+				})
+
+				return m
+			},
+			expected: []string{"foo"},
+			expectedClosed: map[string]struct{}{
+				"foo": {},
+			},
+			expectedErr: []error{
+				errors.New(`could not find session ID "two" in local state after closing connections`),
+			},
+		},
+		{
+			name: "missing connection",
+			sessionCloseInfo: map[string][]*pbs.CloseConnectionResponseData{
+				"one": {
+					{ConnectionId: "foo", Status: pbs.CONNECTIONSTATUS_CONNECTIONSTATUS_CLOSED},
+				},
+				"two": {
+					{ConnectionId: "bar", Status: pbs.CONNECTIONSTATUS_CONNECTIONSTATUS_CLOSED},
+				},
+			},
+			sessionInfoMap: func() sync.Map {
+				var m sync.Map
+				m.Store("one", &sess{
+					resp: &pbs.LookupSessionResponse{Authorization: &targets.SessionAuthorizationData{
+						SessionId: "one",
+					}},
+					connInfoMap: map[string]*ConnInfo{
+						"foo": {Id: "foo"},
+					},
+				})
+				m.Store("two", &sess{
+					resp: &pbs.LookupSessionResponse{Authorization: &targets.SessionAuthorizationData{
+						SessionId: "two",
+					}},
+				})
+
+				return m
+			},
+			expected: []string{"foo"},
+			expectedClosed: map[string]struct{}{
+				"foo": {},
+			},
+			expectedErr: []error{
+				errors.New(`could not find connection ID "bar" for session ID "two" in local state`),
+			},
+		},
+		{
+			name:             "empty",
+			sessionCloseInfo: make(map[string][]*pbs.CloseConnectionResponseData),
+			sessionInfoMap: func() sync.Map {
+				var m sync.Map
+				m.Store("one", &sess{
+					resp: &pbs.LookupSessionResponse{Authorization: &targets.SessionAuthorizationData{
+						SessionId: "one",
+					}},
+					connInfoMap: map[string]*ConnInfo{
+						"foo": {Id: "foo"},
+					},
+				})
+
+				return m
+			},
+			expected:       []string{},
+			expectedClosed: map[string]struct{}{},
+		},
+	}
+
+	for _, tc := range cases {
+		tc := tc
+		t.Run(tc.name, func(t *testing.T) {
+			require := require.New(t)
+			manager := &manager{}
+			manager.sessionMap = tc.sessionInfoMap()
+			actual, actualErr := setCloseTimeForResponse(manager, tc.sessionCloseInfo)
+
+			// Assert all close times were set
+			manager.ForEachLocalSession(func(value Session) bool {
+				t.Helper()
+				for _, ci := range value.GetLocalConnections() {
+					if _, ok := tc.expectedClosed[ci.Id]; ok {
+						require.NotEqual(time.Time{}, ci.CloseTime)
+					} else {
+						require.Equal(time.Time{}, ci.CloseTime)
+					}
+				}
+
+				return true
+			})
+
+			// Assert return values
+			require.ElementsMatch(tc.expected, actual)
+			require.ElementsMatch(tc.expectedErr, actualErr)
+		})
+	}
+}

internal/daemon/worker/session/session.go

@@ -2,7 +2,7 @@ package session
 
 import (
 	"context"
-	"crypto/tls"
+	"crypto/x509"
 	"errors"
 	"fmt"
 	"sync"
@@ -22,26 +22,291 @@ var errMakeSessionCloseInfoNilCloseInfo = errors.New("nil closeInfo supplied to
 
 // ConnInfo defines the information about a connection attached to a session
 type ConnInfo struct {
-	Id         string
-	ConnCtx    context.Context
-	ConnCancel context.CancelFunc
-	Status     pbs.CONNECTIONSTATUS
-	CloseTime  time.Time
+	Id     string
+	Status pbs.CONNECTIONSTATUS
+
+	// The context.CancelFunc for the proxy connection.  Calling this function
+	// closes the proxy connection.
+	connCtxCancelFunc context.CancelFunc
+
+	// The time the controller has successfully reported that this connection is
+	// closed.
+	CloseTime time.Time
 }
 
-// Info defines the information about a session
-type Info struct {
-	sync.RWMutex
-	Id                    string
-	SessionTls            *tls.Config
-	Status                pbs.SESSIONSTATUS
-	LookupSessionResponse *pbs.LookupSessionResponse
-	ConnInfoMap           map[string]*ConnInfo
+// Session is the local representation of a session.  After initial loading
+// the only values that will change will be the status (readable from
+// GetStatus()) and the Connections (GetLocalConnections()).
+type Session interface {
+	// ApplyLocalConnectionStatus set's a connection's status to the one provided.
+	// If there is no connection with the provided id, an error is returned.
+	ApplyLocalConnectionStatus(connId string, status pbs.CONNECTIONSTATUS) error
+
+	// ApplyLocalStatus updates the given session with the status provided by
+	// the SessionJobInfo.  It returns an error if any of the connections
+	// in the SessionJobInfo are not present, however, it still applies the
+	// status change to the session and the connections which are present.
+	ApplyLocalStatus(st pbs.SESSIONSTATUS)
+
+	GetStatus() pbs.SESSIONSTATUS
+	// GetLocalConnections returns the connections this session is handling.
+	GetLocalConnections() map[string]ConnInfo
+
+	GetTofuToken() string
+	GetConnectionLimit() int32
+	GetEndpoint() string
+	GetCredentials() []*pbs.Credential
+	GetExpiration() time.Time
+	GetCertificate() *x509.Certificate
+	GetPrivateKey() []byte
+	GetId() string
+
+	// CancelOpenLocalConnections closes the local connections in this session
+	//based on the connection's state by calling the connections context cancel
+	// function.
+	//
+	// The returned slice are connection ids that were closed.
+	CancelOpenLocalConnections() []string
+
+	// CancelAllLocalConnections close connections regardless of connection's state
+	// by calling the connection context's CancelFunc.
+	//
+	// The returned slice is the connection ids which were closed.
+	CancelAllLocalConnections() []string
+
+	// RequestCancel sends session cancellation request to the controller.  If there is no
+	// error the local session's status is updated with the result of the cancel
+	// request
+	RequestCancel(ctx context.Context) error
+
+	// RequestActivate sends session activation request to the controller.  The Session's
+	// status is then updated with the result of the call.  After a successful
+	// call to RequestActivate, subsequent calls will fail.
+	RequestActivate(ctx context.Context, tofu string) error
+
+	// RequestAuthorizeConnection sends an AuthorizeConnection request to
+	// the controller.
+	// It is called by the worker handler after a connection has been received by
+	// the worker, and the session has been validated.
+	// The passed in context.CancelFunc is used to terminate any ongoing local proxy
+	// connections.
+	// The connection status is then viewable in this session's GetLocalConnections() call.
+	RequestAuthorizeConnection(ctx context.Context, workerId string, connCancel context.CancelFunc) (ConnInfo, int32, error)
+
+	// RequestConnectConnection sends a RequestConnectConnection request to the controller. It
+	// should only be called by the worker handler after a connection has been
+	// authorized.  The local connection's status is updated with the result of the
+	// call.
+	RequestConnectConnection(ctx context.Context, info *pbs.ConnectConnectionRequest) error
 }
 
-// Activate is a helper worker function that sends session activation request to the
+type sess struct {
+	lock        sync.RWMutex
+	client      pbs.SessionServiceClient
+	connInfoMap map[string]*ConnInfo
+	resp        *pbs.LookupSessionResponse
+	status      pbs.SESSIONSTATUS
+	cert        *x509.Certificate
+}
+
+func newSess(client pbs.SessionServiceClient, resp *pbs.LookupSessionResponse) (*sess, error) {
+	switch {
+	case isNil(client):
+		return nil, errors.New("SessionServiceClient is nil")
+	case resp.GetExpiration().AsTime().Before(time.Now()):
+		return nil, fmt.Errorf("session is expired")
+	}
+
+	parsedCert, err := x509.ParseCertificate(resp.GetAuthorization().GetCertificate())
+	if err != nil {
+		return nil, fmt.Errorf("error parsing session certificate: %w", err)
+	}
+
+	s := &sess{
+		client:      client,
+		connInfoMap: make(map[string]*ConnInfo),
+		resp:        resp,
+		status:      resp.GetStatus(),
+		cert:        parsedCert,
+	}
+	return s, nil
+}
+
+// ApplyLocalConnectionStatus Satisfies the Session interface
+func (s *sess) ApplyLocalConnectionStatus(connId string, status pbs.CONNECTIONSTATUS) error {
+	s.lock.Lock()
+	defer s.lock.Unlock()
+	// Update connection status if there are any connections in
+	// the request.
+	connInfo, ok := s.connInfoMap[connId]
+	if !ok {
+		return fmt.Errorf("could not find connection ID %q for session ID %q in local state",
+			connId,
+			s.GetId())
+	}
+
+	connInfo.Status = status
+	if connInfo.Status == pbs.CONNECTIONSTATUS_CONNECTIONSTATUS_CLOSED {
+		connInfo.CloseTime = time.Now()
+	}
+	return nil
+}
+
+func (s *sess) ApplyLocalStatus(st pbs.SESSIONSTATUS) {
+	s.lock.Lock()
+	defer s.lock.Unlock()
+	s.status = st
+}
+
+func (s *sess) GetLocalConnections() map[string]ConnInfo {
+	s.lock.RLock()
+	defer s.lock.RUnlock()
+	res := make(map[string]ConnInfo, len(s.connInfoMap))
+	// Returning the s.connInfoMap directly wouldn't be thread safe.
+	for k, v := range s.connInfoMap {
+		res[k] = ConnInfo{
+			Id:        v.Id,
+			Status:    v.Status,
+			CloseTime: v.CloseTime,
+		}
+	}
+	return res
+}
+
+func (s *sess) GetTofuToken() string {
+	return s.resp.GetTofuToken()
+}
+
+func (s *sess) GetConnectionLimit() int32 {
+	return s.resp.GetConnectionLimit()
+}
+
+func (s *sess) GetEndpoint() string {
+	return s.resp.GetEndpoint()
+}
+
+func (s *sess) GetCredentials() []*pbs.Credential {
+	return s.resp.GetCredentials()
+}
+
+func (s *sess) GetStatus() pbs.SESSIONSTATUS {
+	s.lock.RLock()
+	defer s.lock.RUnlock()
+	return s.status
+}
+
+func (s *sess) GetExpiration() time.Time {
+	return s.resp.GetExpiration().AsTime()
+}
+
+func (s *sess) GetCertificate() *x509.Certificate {
+	return s.cert
+}
+
+func (s *sess) GetPrivateKey() []byte {
+	return s.resp.GetAuthorization().GetPrivateKey()
+}
+
+func (s *sess) GetId() string {
+	return s.resp.GetAuthorization().GetSessionId()
+}
+
+func (s *sess) RequestCancel(ctx context.Context) error {
+	st, err := cancel(ctx, s.client, s.GetId())
+	if err != nil {
+		return err
+	}
+	s.ApplyLocalStatus(st)
+	return nil
+}
+
+func (s *sess) RequestActivate(ctx context.Context, tofu string) error {
+	st, err := activate(ctx, s.client, s.GetId(), tofu, s.resp.GetVersion())
+	if err != nil {
+		return err
+	}
+	s.ApplyLocalStatus(st)
+	return nil
+}
+
+func (s *sess) RequestAuthorizeConnection(ctx context.Context, workerId string, connCancel context.CancelFunc) (ConnInfo, int32, error) {
+	switch {
+	case connCancel == nil:
+		return ConnInfo{}, 0, errors.New("the provided context.CancelFunc was nil")
+	case workerId == "":
+		return ConnInfo{}, 0, errors.New("worker id is empty")
+	}
+	ci, connsLeft, err := authorizeConnection(ctx, s.client, workerId, s.GetId())
+	if err != nil {
+		return ConnInfo{}, connsLeft, err
+	}
+	ci.connCtxCancelFunc = connCancel
+	s.lock.Lock()
+	defer s.lock.Unlock()
+	s.connInfoMap[ci.Id] = ci
+	return ConnInfo{
+		Id:        ci.Id,
+		Status:    ci.Status,
+		CloseTime: ci.CloseTime,
+	}, connsLeft, err
+}
+
+func (s *sess) RequestConnectConnection(ctx context.Context, info *pbs.ConnectConnectionRequest) error {
+	st, err := connectConnection(ctx, s.client, info)
+	if err != nil {
+		return err
+	}
+	s.ApplyLocalConnectionStatus(info.GetConnectionId(), st)
+	return nil
+}
+
+// CancelOpenLocalConnections closes the local connections in this session
+//based on the connection's state by calling the connections context cancel
+// function.
+//
+// The returned slice are connection ids that were closed.
+func (s *sess) CancelOpenLocalConnections() []string {
+	var closedIds []string
+	for k, v := range s.connInfoMap {
+		if v.Status != pbs.CONNECTIONSTATUS_CONNECTIONSTATUS_CLOSED {
+			continue
+		}
+		v.connCtxCancelFunc()
+
+		// CloseTime is set when the controller has reported that the connection
+		// is closed.  The worker should only request a connection be marked
+		// closed after it has already been cancelled.
+		if v.CloseTime.IsZero() {
+			closedIds = append(closedIds, k)
+		}
+	}
+
+	return closedIds
+}
+
+// CancelAllLocalConnections close connections regardless of connection's state
+// by calling the connection context's CancelFunc.
+//
+// The returned slice is the connection ids which were closed.
+func (s *sess) CancelAllLocalConnections() []string {
+	var closedIds []string
+	for k, v := range s.connInfoMap {
+		v.connCtxCancelFunc()
+
+		// CloseTime is set when the controller has reported that the connection
+		// is closed.  The worker should only request a connection be marked
+		// closed after it has already been cancelled.
+		if v.CloseTime.IsZero() {
+			closedIds = append(closedIds, k)
+		}
+	}
+
+	return closedIds
+}
+
+// activate is a helper worker function that sends session activation request to the
 // controller.
-func Activate(ctx context.Context, sessClient pbs.SessionServiceClient, sessionId, tofuToken string, version uint32) (pbs.SESSIONSTATUS, error) {
+func activate(ctx context.Context, sessClient pbs.SessionServiceClient, sessionId, tofuToken string, version uint32) (pbs.SESSIONSTATUS, error) {
 	resp, err := sessClient.ActivateSession(ctx, &pbs.ActivateSessionRequest{
 		SessionId: sessionId,
 		TofuToken: tofuToken,
@@ -53,9 +318,9 @@ func Activate(ctx context.Context, sessClient pbs.SessionServiceClient, sessionI
 	return resp.GetStatus(), nil
 }
 
-// Cancel is a helper worker function that sends session cancellation request to the
+// cancel is a helper worker function that sends session cancellation request to the
 // controller.
-func Cancel(ctx context.Context, sessClient pbs.SessionServiceClient, sessionId string) (pbs.SESSIONSTATUS, error) {
+func cancel(ctx context.Context, sessClient pbs.SessionServiceClient, sessionId string) (pbs.SESSIONSTATUS, error) {
 	resp, err := sessClient.CancelSession(ctx, &pbs.CancelSessionRequest{
 		SessionId: sessionId,
 	})
@@ -65,10 +330,10 @@ func Cancel(ctx context.Context, sessClient pbs.SessionServiceClient, sessionId
 	return resp.GetStatus(), nil
 }
 
-// AuthorizeConnection is a helper worker function that sends connection
+// authorizeConnection is a helper worker function that sends connection
 // authorization request to the controller. It is called by the worker handler after a
 // connection has been received by the worker, and the session has been validated.
-func AuthorizeConnection(ctx context.Context, sessClient pbs.SessionServiceClient, workerId, sessionId string) (*ConnInfo, int32, error) {
+func authorizeConnection(ctx context.Context, sessClient pbs.SessionServiceClient, workerId, sessionId string) (*ConnInfo, int32, error) {
 	resp, err := sessClient.AuthorizeConnection(ctx, &pbs.AuthorizeConnectionRequest{
 		SessionId: sessionId,
 		WorkerId:  workerId,
@@ -83,10 +348,10 @@ func AuthorizeConnection(ctx context.Context, sessClient pbs.SessionServiceClien
 	}, resp.GetConnectionsLeft(), nil
 }
 
-// ConnectConnection is a helper worker function that sends connection
+// connectConnection is a helper worker function that sends connection
 // connect request to the controller. It is called by the worker handler after a
 // connection has been authorized.
-func ConnectConnection(ctx context.Context, sessClient pbs.SessionServiceClient, req *pbs.ConnectConnectionRequest) (pbs.CONNECTIONSTATUS, error) {
+func connectConnection(ctx context.Context, sessClient pbs.SessionServiceClient, req *pbs.ConnectConnectionRequest) (pbs.CONNECTIONSTATUS, error) {
 	resp, err := sessClient.ConnectConnection(ctx, req)
 	if err != nil {
 		return pbs.CONNECTIONSTATUS_CONNECTIONSTATUS_UNSPECIFIED, err
@@ -99,6 +364,8 @@ func ConnectConnection(ctx context.Context, sessClient pbs.SessionServiceClient,
 	return resp.GetStatus(), nil
 }
 
+// TODO: Move these to manager.go.  This is kept here for now simply to make it easier to see the diff.
+
 func closeConnection(ctx context.Context, sessClient pbs.SessionServiceClient, req *pbs.CloseConnectionRequest) (*pbs.CloseConnectionResponse, error) {
 	const op = "session.closeConnection"
 	resp, err := sessClient.CloseConnection(ctx, req)
@@ -112,7 +379,7 @@ func closeConnection(ctx context.Context, sessClient pbs.SessionServiceClient, r
 	return resp, nil
 }
 
-// CloseConnections is a helper worker function that sends connection close
+// closeConnections is a helper worker function that sends connection close
 // requests to the controller, and sets close times within the worker. It is
 // called during the worker status loop and on connection exit on the proxy.
 //
@@ -120,9 +387,9 @@ func closeConnection(ctx context.Context, sessClient pbs.SessionServiceClient, r
 // errors. Individual events will be sent for the errors if there are any.
 //
 // closeInfo is a map of connections mapped to their individual session.
-func CloseConnections(ctx context.Context, sessClient pbs.SessionServiceClient, sessionInfo *sync.Map, closeInfo map[string]string) bool {
-	const op = "session.CloseConnections"
-	if closeInfo == nil {
+func closeConnections(ctx context.Context, sessClient pbs.SessionServiceClient, sManager Manager, closeInfo map[string]string) bool {
+	const op = "session.closeConnections"
+	if len(closeInfo) == 0 {
 		// This should not happen, but it's a no-op if it does. Just
 		// return.
 		return false
@@ -160,7 +427,7 @@ func CloseConnections(ctx context.Context, sessClient pbs.SessionServiceClient,
 	}
 
 	// Mark connections as closed
-	_, errs := setCloseTimeForResponse(sessionInfo, sessionCloseInfo)
+	_, errs := setCloseTimeForResponse(sManager, sessionCloseInfo)
 	if len(errs) > 0 {
 		for _, err := range errs {
 			event.WriteError(ctx, op, err, event.WithInfoMsg("error marking connection closed in state"))
@@ -192,7 +459,7 @@ func makeCloseConnectionRequest(closeInfo map[string]string) *pbs.CloseConnectio
 	}
 }
 
-// makeSessionCloseInfo takes the response from CloseConnections and
+// makeSessionCloseInfo takes the response from closeConnections and
 // our original closeInfo map and makes a map of slices, indexed by
 // session ID, of all of the connection responses. This allows us to
 // easily lock on session once for all connections in
@@ -245,41 +512,27 @@ func makeFakeSessionCloseInfo(
 // failed, as some connections may have been marked as closed. The
 // actual list of connection IDs closed is returned as the first
 // return value.
-func setCloseTimeForResponse(sessionInfo *sync.Map, sessionCloseInfo map[string][]*pbs.CloseConnectionResponseData) ([]string, []error) {
+func setCloseTimeForResponse(sManager Manager, sessionCloseInfo map[string][]*pbs.CloseConnectionResponseData) ([]string, []error) {
 	closedIds := make([]string, 0)
-	var errors []error
+	var result []error
 	for sessionId, responses := range sessionCloseInfo {
-		siRaw, ok := sessionInfo.Load(sessionId)
-		if !ok {
-			errors = append(errors, fmt.Errorf("could not find session ID %q in local state after closing connections", sessionId))
+		si := sManager.Get(sessionId)
+		if si == nil {
+			result = append(result, fmt.Errorf("could not find session ID %q in local state after closing connections", sessionId))
 			continue
 		}
-
-		si := siRaw.(*Info)
-		si.Lock()
-
+		connStatus := make(map[string]pbs.CONNECTIONSTATUS, len(responses))
 		for _, response := range responses {
-			ci, ok := si.ConnInfoMap[response.GetConnectionId()]
-			if !ok {
-				errors = append(errors,
-					fmt.Errorf(
-						"could not find connection ID %q for session ID %q in local state after closing connections",
-						response.GetConnectionId(),
-						sessionId,
-					),
-				)
+			if err := si.ApplyLocalConnectionStatus(response.GetConnectionId(), response.GetStatus()); err != nil {
+				result = append(result, err)
 				continue
 			}
-
-			ci.Status = response.GetStatus()
-			if ci.Status == pbs.CONNECTIONSTATUS_CONNECTIONSTATUS_CLOSED {
-				ci.CloseTime = time.Now()
-				closedIds = append(closedIds, ci.Id)
+			connStatus[response.GetConnectionId()] = response.GetStatus()
+			if response.GetStatus() == pbs.CONNECTIONSTATUS_CONNECTIONSTATUS_CLOSED {
+				closedIds = append(closedIds, response.GetConnectionId())
 			}
 		}
-
-		si.Unlock()
 	}
 
-	return closedIds, errors
+	return closedIds, result
 }

internal/daemon/worker/session/session_test.go

@@ -1,16 +1,179 @@
 package session
 
 import (
-	"errors"
-	"sync"
+	"context"
+	"fmt"
 	"testing"
 	"time"
 
 	pbs "github.com/hashicorp/boundary/internal/gen/controller/servers/services"
 	"github.com/hashicorp/boundary/internal/session"
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
 
+func TestSession_ApplyLocalStatus(t *testing.T) {
+	sess := &sess{
+		status: pbs.SESSIONSTATUS_SESSIONSTATUS_PENDING,
+	}
+	for _, s := range []pbs.SESSIONSTATUS{
+		pbs.SESSIONSTATUS_SESSIONSTATUS_ACTIVE,
+		pbs.SESSIONSTATUS_SESSIONSTATUS_PENDING,
+		pbs.SESSIONSTATUS_SESSIONSTATUS_CANCELING,
+		pbs.SESSIONSTATUS_SESSIONSTATUS_TERMINATED,
+	} {
+		sess.ApplyLocalStatus(s)
+		assert.Equal(t, s, sess.GetStatus())
+	}
+}
+
+func TestSession_CancelAllLocalConnections(t *testing.T) {
+	var closedContextCalled []string
+	cancelFn := func(id string) context.CancelFunc {
+		return func() {
+			closedContextCalled = append(closedContextCalled, id)
+		}
+	}
+	connInfo := map[string]*ConnInfo{
+		"1": {
+			Id:                "1",
+			connCtxCancelFunc: cancelFn("1"),
+		},
+		"2": {
+			Id:                "2",
+			connCtxCancelFunc: cancelFn("2"),
+		},
+		"3": {
+			Id:                "3",
+			connCtxCancelFunc: cancelFn("3"),
+			CloseTime:         time.Now(),
+		},
+	}
+	sess := &sess{
+		connInfoMap: connInfo,
+	}
+	assert.ElementsMatch(t, sess.CancelAllLocalConnections(), []string{"1", "2"})
+	// We can call the cancel context multiple times, even if it was marked
+	// closed previously.
+	assert.ElementsMatch(t, closedContextCalled, []string{"1", "2", "3"})
+}
+
+func TestSession_CancelOpenLocalConnections(t *testing.T) {
+	var closedContextCalled []string
+	cancelFn := func(id string) context.CancelFunc {
+		return func() {
+			closedContextCalled = append(closedContextCalled, id)
+		}
+	}
+	connInfo := map[string]*ConnInfo{
+		"1": {
+			Id:                "1",
+			connCtxCancelFunc: cancelFn("1"),
+			Status:            pbs.CONNECTIONSTATUS_CONNECTIONSTATUS_CLOSED,
+		},
+		"2": {
+			Id:                "2",
+			connCtxCancelFunc: cancelFn("2"),
+			Status:            pbs.CONNECTIONSTATUS_CONNECTIONSTATUS_CLOSED,
+			CloseTime:         time.Now(),
+		},
+		"3": {
+			Id:                "3",
+			connCtxCancelFunc: cancelFn("3"),
+			Status:            pbs.CONNECTIONSTATUS_CONNECTIONSTATUS_CONNECTED,
+		},
+	}
+	sess := &sess{
+		connInfoMap: connInfo,
+	}
+	assert.ElementsMatch(t, sess.CancelOpenLocalConnections(), []string{"1"})
+	// We call the cancel context multiple times, even if it was marked
+	// closed previously like connection 2 was (by setting CloseTime)
+	assert.ElementsMatch(t, closedContextCalled, []string{"1", "2"})
+}
+
+func TestSession_RequestActivate(t *testing.T) {
+	mockClient := pbs.NewMockSessionServiceClient()
+	mockClient.ActivateSessionFn = func(context.Context, *pbs.ActivateSessionRequest) (*pbs.ActivateSessionResponse, error) {
+		return nil, fmt.Errorf("test error")
+	}
+	sess := &sess{
+		client: mockClient,
+		resp: &pbs.LookupSessionResponse{
+			TofuToken:       "tofu",
+			Version:         1,
+			Status:          pbs.SESSIONSTATUS_SESSIONSTATUS_PENDING,
+			ConnectionLimit: -1,
+		},
+		status: pbs.SESSIONSTATUS_SESSIONSTATUS_PENDING,
+	}
+	assert.Error(t, sess.RequestActivate(context.Background(), "tofu"))
+
+	mockClient.ActivateSessionFn = func(context.Context, *pbs.ActivateSessionRequest) (*pbs.ActivateSessionResponse, error) {
+		return &pbs.ActivateSessionResponse{Status: pbs.SESSIONSTATUS_SESSIONSTATUS_ACTIVE}, nil
+	}
+	assert.NoError(t, sess.RequestActivate(context.Background(), "tofu"))
+	assert.Equal(t, pbs.SESSIONSTATUS_SESSIONSTATUS_ACTIVE, sess.GetStatus())
+}
+
+func TestSession_RequestCancel(t *testing.T) {
+	mockClient := pbs.NewMockSessionServiceClient()
+	mockClient.CancelSessionFn = func(context.Context, *pbs.CancelSessionRequest) (*pbs.CancelSessionResponse, error) {
+		return nil, fmt.Errorf("test error")
+	}
+	sess := &sess{
+		client: mockClient,
+		resp: &pbs.LookupSessionResponse{
+			TofuToken:       "tofu",
+			Version:         1,
+			Status:          pbs.SESSIONSTATUS_SESSIONSTATUS_PENDING,
+			ConnectionLimit: -1,
+		},
+		status: pbs.SESSIONSTATUS_SESSIONSTATUS_PENDING,
+	}
+	assert.Error(t, sess.RequestCancel(context.Background()))
+
+	mockClient.CancelSessionFn = func(context.Context, *pbs.CancelSessionRequest) (*pbs.CancelSessionResponse, error) {
+		return &pbs.CancelSessionResponse{Status: pbs.SESSIONSTATUS_SESSIONSTATUS_CANCELING}, nil
+	}
+	assert.NoError(t, sess.RequestCancel(context.Background()))
+	assert.Equal(t, pbs.SESSIONSTATUS_SESSIONSTATUS_CANCELING, sess.GetStatus())
+}
+
+func TestSession_RequestAuthorizeConnection(t *testing.T) {
+	mockClient := pbs.NewMockSessionServiceClient()
+	mockClient.AuthorizeConnectionFn = func(ctx context.Context, request *pbs.AuthorizeConnectionRequest) (*pbs.AuthorizeConnectionResponse, error) {
+		return nil, fmt.Errorf("test error")
+	}
+	sess := &sess{
+		client:      mockClient,
+		connInfoMap: make(map[string]*ConnInfo),
+		resp: &pbs.LookupSessionResponse{
+			TofuToken:       "tofu",
+			Version:         1,
+			Status:          pbs.SESSIONSTATUS_SESSIONSTATUS_PENDING,
+			ConnectionLimit: -1,
+		},
+		status: pbs.SESSIONSTATUS_SESSIONSTATUS_PENDING,
+	}
+	_, cancel := context.WithCancel(context.Background())
+	conn, _, err := sess.RequestAuthorizeConnection(context.Background(), "workerid", cancel)
+	require.Error(t, err)
+	assert.Equal(t, ConnInfo{}, conn)
+
+	mockClient.AuthorizeConnectionFn = func(ctx context.Context, request *pbs.AuthorizeConnectionRequest) (*pbs.AuthorizeConnectionResponse, error) {
+		return &pbs.AuthorizeConnectionResponse{
+			ConnectionId:    "conn1",
+			Status:          pbs.CONNECTIONSTATUS_CONNECTIONSTATUS_AUTHORIZED,
+			ConnectionsLeft: -1,
+		}, nil
+	}
+	conn, left, err := sess.RequestAuthorizeConnection(context.Background(), "workerid", cancel)
+	require.NoError(t, err)
+	assert.Equal(t, ConnInfo{Id: "conn1", Status: pbs.CONNECTIONSTATUS_CONNECTIONSTATUS_AUTHORIZED}, conn)
+	assert.Equal(t, int32(-1), left)
+}
+
 func TestWorkerMakeCloseConnectionRequest(t *testing.T) {
 	require := require.New(t)
 	in := map[string]string{"foo": "one", "bar": "two"}
@@ -95,189 +258,3 @@ func TestMakeFakeSessionCloseInfoEmpty(t *testing.T) {
 		actual,
 	)
 }
-
-func TestWorkerSetCloseTimeForResponse(t *testing.T) {
-	cases := []struct {
-		name             string
-		sessionCloseInfo map[string][]*pbs.CloseConnectionResponseData
-		sessionInfoMap   func() *sync.Map
-		expected         []string
-		expectedClosed   map[string]struct{}
-		expectedErr      []error
-	}{
-		{
-			name: "basic",
-			sessionCloseInfo: map[string][]*pbs.CloseConnectionResponseData{
-				"one": {
-					{ConnectionId: "foo", Status: pbs.CONNECTIONSTATUS_CONNECTIONSTATUS_CLOSED},
-				},
-				"two": {
-					{ConnectionId: "bar", Status: pbs.CONNECTIONSTATUS_CONNECTIONSTATUS_CLOSED},
-				},
-			},
-			sessionInfoMap: func() *sync.Map {
-				m := new(sync.Map)
-				m.Store("one", &Info{
-					Id: "one",
-					ConnInfoMap: map[string]*ConnInfo{
-						"foo": {Id: "foo"},
-					},
-				})
-				m.Store("two", &Info{
-					Id: "two",
-					ConnInfoMap: map[string]*ConnInfo{
-						"bar": {Id: "bar"},
-					},
-				})
-				m.Store("three", &Info{
-					Id: "three",
-					ConnInfoMap: map[string]*ConnInfo{
-						"baz": {Id: "baz"},
-					},
-				})
-
-				return m
-			},
-			expected: []string{"foo", "bar"},
-			expectedClosed: map[string]struct{}{
-				"foo": {},
-				"bar": {},
-			},
-		},
-		{
-			name: "not closed",
-			sessionCloseInfo: map[string][]*pbs.CloseConnectionResponseData{
-				"one": {
-					{ConnectionId: "foo", Status: pbs.CONNECTIONSTATUS_CONNECTIONSTATUS_CLOSED},
-				},
-				"two": {
-					{ConnectionId: "bar", Status: pbs.CONNECTIONSTATUS_CONNECTIONSTATUS_CONNECTED},
-				},
-			},
-			sessionInfoMap: func() *sync.Map {
-				m := new(sync.Map)
-				m.Store("one", &Info{
-					Id: "one",
-					ConnInfoMap: map[string]*ConnInfo{
-						"foo": {Id: "foo"},
-					},
-				})
-				m.Store("two", &Info{
-					Id: "two",
-					ConnInfoMap: map[string]*ConnInfo{
-						"bar": {Id: "bar"},
-					},
-				})
-
-				return m
-			},
-			expected: []string{"foo"},
-			expectedClosed: map[string]struct{}{
-				"foo": {},
-			},
-		},
-		{
-			name: "missing session",
-			sessionCloseInfo: map[string][]*pbs.CloseConnectionResponseData{
-				"one": {
-					{ConnectionId: "foo", Status: pbs.CONNECTIONSTATUS_CONNECTIONSTATUS_CLOSED},
-				},
-				"two": {
-					{ConnectionId: "bar", Status: pbs.CONNECTIONSTATUS_CONNECTIONSTATUS_CLOSED},
-				},
-			},
-			sessionInfoMap: func() *sync.Map {
-				m := new(sync.Map)
-				m.Store("one", &Info{
-					Id: "one",
-					ConnInfoMap: map[string]*ConnInfo{
-						"foo": {Id: "foo"},
-					},
-				})
-
-				return m
-			},
-			expected: []string{"foo"},
-			expectedClosed: map[string]struct{}{
-				"foo": {},
-			},
-			expectedErr: []error{
-				errors.New(`could not find session ID "two" in local state after closing connections`),
-			},
-		},
-		{
-			name: "missing connection",
-			sessionCloseInfo: map[string][]*pbs.CloseConnectionResponseData{
-				"one": {
-					{ConnectionId: "foo", Status: pbs.CONNECTIONSTATUS_CONNECTIONSTATUS_CLOSED},
-				},
-				"two": {
-					{ConnectionId: "bar", Status: pbs.CONNECTIONSTATUS_CONNECTIONSTATUS_CLOSED},
-				},
-			},
-			sessionInfoMap: func() *sync.Map {
-				m := new(sync.Map)
-				m.Store("one", &Info{
-					Id: "one",
-					ConnInfoMap: map[string]*ConnInfo{
-						"foo": {Id: "foo"},
-					},
-				})
-				m.Store("two", &Info{Id: "two"})
-
-				return m
-			},
-			expected: []string{"foo"},
-			expectedClosed: map[string]struct{}{
-				"foo": {},
-			},
-			expectedErr: []error{
-				errors.New(`could not find connection ID "bar" for session ID "two" in local state after closing connections`),
-			},
-		},
-		{
-			name:             "empty",
-			sessionCloseInfo: make(map[string][]*pbs.CloseConnectionResponseData),
-			sessionInfoMap: func() *sync.Map {
-				m := new(sync.Map)
-				m.Store("one", &Info{
-					Id: "one",
-					ConnInfoMap: map[string]*ConnInfo{
-						"foo": {Id: "foo"},
-					},
-				})
-
-				return m
-			},
-			expected:       []string{},
-			expectedClosed: map[string]struct{}{},
-		},
-	}
-
-	for _, tc := range cases {
-		tc := tc
-		t.Run(tc.name, func(t *testing.T) {
-			require := require.New(t)
-			sessionInfoMap := tc.sessionInfoMap()
-			actual, actualErr := setCloseTimeForResponse(sessionInfoMap, tc.sessionCloseInfo)
-
-			// Assert all close times were set
-			sessionInfoMap.Range(func(key, value interface{}) bool {
-				t.Helper()
-				for _, ci := range value.(*Info).ConnInfoMap {
-					if _, ok := tc.expectedClosed[ci.Id]; ok {
-						require.NotEqual(time.Time{}, ci.CloseTime)
-					} else {
-						require.Equal(time.Time{}, ci.CloseTime)
-					}
-				}
-
-				return true
-			})
-
-			// Assert return values
-			require.ElementsMatch(tc.expected, actual)
-			require.ElementsMatch(tc.expectedErr, actualErr)
-		})
-	}
-}

internal/daemon/worker/status.go

@@ -23,7 +23,7 @@ type LastStatusInformation struct {
 	LastCalculatedUpstreams []string
 }
 
-func (w *Worker) startStatusTicking(cancelCtx context.Context) {
+func (w *Worker) startStatusTicking(cancelCtx context.Context, sessionManager session.Manager) {
 	const op = "worker.(Worker).startStatusTicking"
 	r := rand.New(rand.NewSource(time.Now().UnixNano()))
 	// This function exists to desynchronize calls to controllers from
@@ -47,7 +47,7 @@ func (w *Worker) startStatusTicking(cancelCtx context.Context) {
 			return
 
 		case <-timer.C:
-			w.sendWorkerStatus(cancelCtx)
+			w.sendWorkerStatus(cancelCtx, sessionManager)
 			timer.Reset(getRandomInterval())
 		}
 	}
@@ -91,7 +91,7 @@ func (w *Worker) WaitForNextSuccessfulStatusUpdate() error {
 	return nil
 }
 
-func (w *Worker) sendWorkerStatus(cancelCtx context.Context) {
+func (w *Worker) sendWorkerStatus(cancelCtx context.Context, sessionManager session.Manager) {
 	const op = "worker.(Worker).sendWorkerStatus"
 
 	// If we've never managed to successfully authenticate then we won't have
@@ -107,20 +107,17 @@ func (w *Worker) sendWorkerStatus(cancelCtx context.Context) {
 	var activeJobs []*pbs.JobStatus
 
 	// Range over known sessions and collect info
-	w.sessionInfoMap.Range(func(key, value interface{}) bool {
+	sessionManager.ForEachLocalSession(func(s session.Session) bool {
 		var jobInfo pbs.SessionJobInfo
-		sessionId := key.(string)
-		si := value.(*session.Info)
-		si.RLock()
-		status := si.Status
-		connections := make([]*pbs.Connection, 0, len(si.ConnInfoMap))
-		for k, v := range si.ConnInfoMap {
+		status := s.GetStatus()
+		sessionId := s.GetId()
+		connections := make([]*pbs.Connection, 0, len(s.GetLocalConnections()))
+		for k, v := range s.GetLocalConnections() {
 			connections = append(connections, &pbs.Connection{
 				ConnectionId: k,
 				Status:       v.Status,
 			})
 		}
-		si.RUnlock()
 		jobInfo.SessionId = sessionId
 		activeJobs = append(activeJobs, &pbs.JobStatus{
 			Job: &pbs.Job{
@@ -176,7 +173,6 @@ func (w *Worker) sendWorkerStatus(cancelCtx context.Context) {
 		// want to start terminating all connections as a "break glass" kind of
 		// scenario, as there will be no way we can really tell if these
 		// connections should continue to exist.
-
 		if isPastGrace, lastStatusTime, gracePeriod := w.isPastGrace(); isPastGrace {
 			event.WriteError(statusCtx, op,
 				errors.New("status error grace period has expired, canceling all sessions on worker"),
@@ -185,14 +181,9 @@ func (w *Worker) sendWorkerStatus(cancelCtx context.Context) {
 
 			// Cancel connections if grace period has expired. These Connections will be closed in the
 			// database on the next successful status report, or via the Controller’s dead Worker cleanup connections job.
-			w.sessionInfoMap.Range(func(key, value interface{}) bool {
-				si := value.(*session.Info)
-				si.Lock()
-				defer si.Unlock()
-
-				closedIds := w.cancelConnections(si.ConnInfoMap, true)
-				for _, connId := range closedIds {
-					event.WriteSysEvent(cancelCtx, op, "terminated connection due to status grace period expiration", "session_id", si.Id, "connection_id", connId)
+			sessionManager.ForEachLocalSession(func(s session.Session) bool {
+				for _, connId := range s.CancelAllLocalConnections() {
+					event.WriteSysEvent(cancelCtx, op, "terminated connection due to status grace period expiration", "session_id", s.GetId(), "connection_id", connId)
 				}
 				return true
 			})
@@ -231,28 +222,20 @@ func (w *Worker) sendWorkerStatus(cancelCtx context.Context) {
 				case pbs.JOBTYPE_JOBTYPE_SESSION:
 					sessInfo := request.GetJob().GetSessionInfo()
 					sessionId := sessInfo.GetSessionId()
-					siRaw, ok := w.sessionInfoMap.Load(sessionId)
-					if !ok {
+					si := sessionManager.Get(sessionId)
+					if si == nil {
 						event.WriteError(statusCtx, op, errors.New("session change requested but could not find local information for it"), event.WithInfo("session_id", sessionId))
 						continue
 					}
-					si := siRaw.(*session.Info)
-					si.Lock()
-					si.Status = sessInfo.GetStatus()
+					si.ApplyLocalStatus(sessInfo.GetStatus())
+
 					// Update connection state if there are any connections in
 					// the request.
 					for _, conn := range sessInfo.GetConnections() {
-						connId := conn.GetConnectionId()
-						connInfo, ok := si.ConnInfoMap[connId]
-						if !ok {
-							event.WriteError(statusCtx, op, errors.New("connection change requested but could not find local information for it"), event.WithInfo("connection_id", connId))
-							continue
+						if err := si.ApplyLocalConnectionStatus(conn.GetConnectionId(), conn.GetStatus()); err != nil {
+							event.WriteError(statusCtx, op, err, event.WithInfo("connection_id", conn.GetConnectionId()))
 						}
-
-						connInfo.Status = conn.GetStatus()
 					}
-
-					si.Unlock()
 				}
 			}
 		}
@@ -260,53 +243,48 @@ func (w *Worker) sendWorkerStatus(cancelCtx context.Context) {
 
 	// Standard cleanup: Run through current jobs. Cancel connections
 	// for any canceling session or any session that is expired.
-	w.cleanupConnections(cancelCtx, false)
+	w.cleanupConnections(cancelCtx, false, sessionManager)
 }
 
-// cleanupConnections walks all sessions and shuts down connections.
+// cleanupConnections walks all sessions and shuts down all proxy connections.
+// After the local connections are terminated, they are requested to be marked
+// close on the controller.
 // Additionally, sessions without connections are cleaned up from the
 // local worker's state.
 //
 // Use ignoreSessionState to ignore the state checks, this closes all
-// connections, regardless of whether or not the session is still
-// active.
-func (w *Worker) cleanupConnections(cancelCtx context.Context, ignoreSessionState bool) {
+// connections, regardless of whether or not the session is still active.
+func (w *Worker) cleanupConnections(cancelCtx context.Context, ignoreSessionState bool, sessionManager session.Manager) {
 	const op = "worker.(Worker).cleanupConnections"
 	closeInfo := make(map[string]string)
 	cleanSessionIds := make([]string, 0)
-	w.sessionInfoMap.Range(func(key, value interface{}) bool {
-		si := value.(*session.Info)
-		si.Lock()
-		defer si.Unlock()
+	sessionManager.ForEachLocalSession(func(s session.Session) bool {
 		switch {
 		case ignoreSessionState,
-			si.Status == pbs.SESSIONSTATUS_SESSIONSTATUS_CANCELING,
-			si.Status == pbs.SESSIONSTATUS_SESSIONSTATUS_TERMINATED,
-			time.Until(si.LookupSessionResponse.Expiration.AsTime()) < 0:
+			s.GetStatus() == pbs.SESSIONSTATUS_SESSIONSTATUS_CANCELING,
+			s.GetStatus() == pbs.SESSIONSTATUS_SESSIONSTATUS_TERMINATED,
+			time.Until(s.GetExpiration()) < 0:
 			// Cancel connections without regard to individual connection
 			// state.
-			closedIds := w.cancelConnections(si.ConnInfoMap, true)
+			closedIds := s.CancelAllLocalConnections()
 			for _, connId := range closedIds {
-				closeInfo[connId] = si.Id
-				event.WriteSysEvent(cancelCtx, op, "terminated connection due to cancellation or expiration", "session_id", si.Id, "connection_id", connId)
+				closeInfo[connId] = s.GetId()
+				event.WriteSysEvent(cancelCtx, op, "terminated connection due to cancellation or expiration", "session_id", s.GetId(), "connection_id", connId)
 			}
 
-			// CloseTime is marked by CloseConnections iff the
-			// status is returned for that connection as closed. If
-			// the session is no longer valid and all connections
+			// If the session is no longer valid and all connections
 			// are marked closed, clean up the session.
 			if len(closedIds) == 0 {
-				cleanSessionIds = append(cleanSessionIds, si.Id)
+				cleanSessionIds = append(cleanSessionIds, s.GetId())
 			}
 
 		default:
 			// Cancel connections *with* regard to individual connection
 			// state (ie: only ones that the controller has requested be
 			// terminated).
-			closedIds := w.cancelConnections(si.ConnInfoMap, false)
-			for _, connId := range closedIds {
-				closeInfo[connId] = si.Id
-				event.WriteSysEvent(cancelCtx, op, "terminated connection due to cancellation or expiration", "session_id", si.Id, "connection_id", connId)
+			for _, connId := range s.CancelOpenLocalConnections() {
+				closeInfo[connId] = s.GetId()
+				event.WriteSysEvent(cancelCtx, op, "terminated connection due to cancellation or expiration", "session_id", s.GetId(), "connection_id", connId)
 			}
 		}
 
@@ -319,42 +297,11 @@ func (w *Worker) cleanupConnections(cancelCtx context.Context, ignoreSessionStat
 		// Call out to a helper to send the connection close requests to the
 		// controller, and set the close time. This functionality is shared with
 		// post-close functionality in the proxy handler.
-		sessClient, err := w.ControllerSessionConn()
-		if err != nil {
-			event.WriteError(cancelCtx, op, err, event.WithInfo("failed to create controller session client, connections won't be cleaned up"))
-		} else {
-			_ = session.CloseConnections(cancelCtx, sessClient, w.sessionInfoMap, closeInfo)
-		}
+		_ = sessionManager.RequestCloseConnections(cancelCtx, closeInfo)
 	}
-
 	// Forget sessions where the session is expired/canceled and all
 	// connections are canceled and marked closed
-	for _, v := range cleanSessionIds {
-		w.sessionInfoMap.Delete(v)
-	}
-}
-
-// cancelConnections is run by cleanupConnections to iterate over a
-// session's ConnInfoMap and close connections based on the
-// connection's state (or regardless if ignoreConnectionState is
-// set).
-//
-// The returned map and slice are the maps of connection -> session,
-// and sessions to completely remove from local state, respectively.
-func (w *Worker) cancelConnections(connInfoMap map[string]*session.ConnInfo, ignoreConnectionState bool) []string {
-	var closedIds []string
-	for k, v := range connInfoMap {
-		if v.CloseTime.IsZero() {
-			if !ignoreConnectionState && v.Status != pbs.CONNECTIONSTATUS_CONNECTIONSTATUS_CLOSED {
-				continue
-			}
-
-			v.ConnCancel()
-			closedIds = append(closedIds, k)
-		}
-	}
-
-	return closedIds
+	sessionManager.DeleteLocalSession(cleanSessionIds)
 }
 
 func (w *Worker) lastSuccessfulStatusTime() time.Time {

internal/daemon/worker/testing.go

@@ -11,7 +11,6 @@ import (
 
 	"github.com/hashicorp/boundary/internal/cmd/base"
 	"github.com/hashicorp/boundary/internal/cmd/config"
-	"github.com/hashicorp/boundary/internal/daemon/worker/session"
 	"github.com/hashicorp/boundary/internal/db"
 	pbs "github.com/hashicorp/boundary/internal/gen/controller/servers/services"
 	"github.com/hashicorp/boundary/internal/observability/event"
@@ -116,17 +115,13 @@ type TestConnectionInfo struct {
 // See TestSessionInfo for details on how to use this info.
 func (tw *TestWorker) LookupSession(id string) (TestSessionInfo, bool) {
 	var result TestSessionInfo
-	raw, ok := tw.w.sessionInfoMap.Load(id)
-	if !ok {
-		return result, false
+	sess := tw.w.sessionManager.Get(id)
+	if sess == nil {
+		return TestSessionInfo{}, false
 	}
 
-	sess := raw.(*session.Info)
-	sess.RLock()
-	defer sess.RUnlock()
-
 	conns := make(map[string]TestConnectionInfo)
-	for _, conn := range sess.ConnInfoMap {
+	for _, conn := range sess.GetLocalConnections() {
 		conns[conn.Id] = TestConnectionInfo{
 			Id:        conn.Id,
 			Status:    conn.Status,
@@ -134,8 +129,8 @@ func (tw *TestWorker) LookupSession(id string) (TestSessionInfo, bool) {
 		}
 	}
 
-	result.Id = sess.Id
-	result.Status = sess.Status
+	result.Id = sess.GetId()
+	result.Status = sess.GetStatus()
 	result.Connections = conns
 
 	return result, true

internal/daemon/worker/testing_test.go

@@ -1,40 +1,72 @@
 package worker
 
 import (
-	"sync"
+	"context"
+	"crypto/ed25519"
+	"crypto/rand"
+	"crypto/x509"
+	"math/big"
+	"net"
 	"testing"
 	"time"
 
 	"github.com/hashicorp/boundary/internal/daemon/worker/session"
 	"github.com/hashicorp/boundary/internal/db"
 	pbs "github.com/hashicorp/boundary/internal/gen/controller/servers/services"
+	"github.com/hashicorp/boundary/sdk/pbs/controller/api/resources/targets"
 	wrapping "github.com/hashicorp/go-kms-wrapping/v2"
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	"google.golang.org/protobuf/types/known/timestamppb"
+
+	_ "github.com/hashicorp/boundary/internal/daemon/controller/handlers/targets/tcp"
 )
 
 func TestTestWorkerLookupSession(t *testing.T) {
 	require := require.New(t)
-	// This loads the golang reference time, see those docs for more details. We
-	// just use this as a stable non-zero time.
-	refTime, err := time.Parse(time.RFC3339, "2006-01-02T15:04:05+07:00")
+	ctx := context.Background()
+
+	mockSessionClient := pbs.NewMockSessionServiceClient()
+	manager, err := session.NewManager(mockSessionClient)
 	require.NoError(err)
+	mockSessionClient.LookupSessionFn = func(_ context.Context, request *pbs.LookupSessionRequest) (*pbs.LookupSessionResponse, error) {
+		cert, _, _ := createTestCert(t)
+		return &pbs.LookupSessionResponse{
+			Authorization: &targets.SessionAuthorizationData{
+				SessionId:   request.GetSessionId(),
+				Certificate: cert,
+			},
+			Version:    1,
+			TofuToken:  "tofu",
+			Expiration: timestamppb.New(time.Now().Add(time.Hour)),
+		}, nil
+	}
+	s, err := manager.LoadLocalSession(ctx, "foo", "worker id")
+	require.NoError(err)
+	mockSessionClient.ActivateSessionFn = func(_ context.Context, _ *pbs.ActivateSessionRequest) (*pbs.ActivateSessionResponse, error) {
+		return &pbs.ActivateSessionResponse{Status: pbs.SESSIONSTATUS_SESSIONSTATUS_ACTIVE}, nil
+	}
+	require.NoError(s.RequestActivate(ctx, "tofu"))
+
+	mockSessionClient.AuthorizeConnectionFn = func(_ context.Context, req *pbs.AuthorizeConnectionRequest) (*pbs.AuthorizeConnectionResponse, error) {
+		return &pbs.AuthorizeConnectionResponse{
+			ConnectionId:    "one",
+			Status:          pbs.CONNECTIONSTATUS_CONNECTIONSTATUS_AUTHORIZED,
+			ConnectionsLeft: -1,
+		}, nil
+	}
+	_, cancelFn := context.WithCancel(context.Background())
+	_, _, err = s.RequestAuthorizeConnection(ctx, "worker id", cancelFn)
+	require.NoError(err)
+	require.NoError(s.ApplyLocalConnectionStatus("one", pbs.CONNECTIONSTATUS_CONNECTIONSTATUS_CLOSED))
 
 	tw := new(TestWorker)
 	tw.w = &Worker{
-		sessionInfoMap: new(sync.Map),
+		sessionManager: manager,
 	}
-	tw.w.sessionInfoMap.Store("foo", &session.Info{
-		Id:     "foo",
-		Status: pbs.SESSIONSTATUS_SESSIONSTATUS_ACTIVE,
-		ConnInfoMap: map[string]*session.ConnInfo{
-			"one": {
-				Id:        "one",
-				Status:    pbs.CONNECTIONSTATUS_CONNECTIONSTATUS_CLOSED,
-				CloseTime: refTime,
-			},
-		},
-	})
 
+	closeTime := s.GetLocalConnections()["one"].CloseTime
+	assert.NotZero(t, closeTime)
 	expected := TestSessionInfo{
 		Id:     "foo",
 		Status: pbs.SESSIONSTATUS_SESSIONSTATUS_ACTIVE,
@@ -42,7 +74,7 @@ func TestTestWorkerLookupSession(t *testing.T) {
 			"one": {
 				Id:        "one",
 				Status:    pbs.CONNECTIONSTATUS_CONNECTIONSTATUS_CLOSED,
-				CloseTime: refTime,
+				CloseTime: closeTime,
 			},
 		},
 	}
@@ -54,10 +86,7 @@ func TestTestWorkerLookupSession(t *testing.T) {
 
 func TestTestWorkerLookupSessionMissing(t *testing.T) {
 	require := require.New(t)
-	tw := new(TestWorker)
-	tw.w = &Worker{
-		sessionInfoMap: new(sync.Map),
-	}
+	tw := NewTestWorker(t, nil)
 	actual, ok := tw.LookupSession("missing")
 	require.False(ok)
 	require.Equal(TestSessionInfo{}, actual)
@@ -87,3 +116,25 @@ func TestTestWorker_WorkerAuthStorageKms(t *testing.T) {
 		})
 	}
 }
+
+func createTestCert(t *testing.T) ([]byte, ed25519.PublicKey, ed25519.PrivateKey) {
+	pub, priv, err := ed25519.GenerateKey(rand.Reader)
+	require.NoError(t, err)
+
+	template := &x509.Certificate{
+		ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth, x509.ExtKeyUsageClientAuth},
+		KeyUsage:              x509.KeyUsageDigitalSignature | x509.KeyUsageKeyEncipherment | x509.KeyUsageKeyAgreement | x509.KeyUsageCertSign,
+		SerialNumber:          big.NewInt(0),
+		NotBefore:             time.Now().Add(-30 * time.Second),
+		NotAfter:              time.Now().Add(5 * time.Minute),
+		BasicConstraintsValid: true,
+		IsCA:                  true,
+		IPAddresses:           []net.IP{net.ParseIP("127.0.0.1")},
+		DNSNames:              []string{"/tmp/boundary-opslistener-test0.sock", "/tmp/boundary-opslistener-test1.sock"},
+	}
+
+	certBytes, err := x509.CreateCertificate(rand.Reader, template, template, pub, priv)
+	require.NoError(t, err)
+
+	return certBytes, pub, priv
+}

internal/daemon/worker/worker.go

@@ -52,7 +52,12 @@ type Worker struct {
 
 	tickerWg sync.WaitGroup
 
-	grpcClientConn *atomic.Value
+	// grpc.ClientConns are thread safe.
+	// See https://github.com/grpc/grpc-go/blob/master/Documentation/concurrency.md#clients
+	// This is exported for tests.
+	GrpcClientConn *grpc.ClientConn
+
+	sessionManager session.Manager
 
 	controllerStatusConn *atomic.Value
 	everAuthenticated    *ua.Bool
@@ -61,9 +66,6 @@ type Worker struct {
 
 	controllerResolver *atomic.Value
 
-	controllerSessionConn *atomic.Value
-	sessionInfoMap        *sync.Map
-
 	controllerMultihopConn *atomic.Value
 
 	proxyListener *base.ServerListener
@@ -100,13 +102,10 @@ func New(conf *Config) (*Worker, error) {
 		conf:                   conf,
 		logger:                 conf.Logger.Named("worker"),
 		started:                ua.NewBool(false),
-		grpcClientConn:         new(atomic.Value),
 		controllerStatusConn:   new(atomic.Value),
 		everAuthenticated:      ua.NewBool(false),
 		lastStatusSuccess:      new(atomic.Value),
 		controllerResolver:     new(atomic.Value),
-		controllerSessionConn:  new(atomic.Value),
-		sessionInfoMap:         new(sync.Map),
 		controllerMultihopConn: new(atomic.Value),
 		tags:                   new(atomic.Value),
 		updateTags:             ua.NewBool(false),
@@ -115,7 +114,9 @@ func New(conf *Config) (*Worker, error) {
 	}
 
 	w.lastStatusSuccess.Store((*LastStatusInformation)(nil))
-	w.controllerResolver.Store((*manual.Resolver)(nil))
+	scheme := strconv.FormatInt(time.Now().UnixNano(), 36)
+	controllerResolver := manual.NewBuilderWithScheme(scheme)
+	w.controllerResolver.Store(controllerResolver)
 
 	if conf.RawConfig.Worker == nil {
 		conf.RawConfig.Worker = new(config.Worker)
@@ -177,10 +178,6 @@ func (w *Worker) Start() error {
 		return nil
 	}
 
-	scheme := strconv.FormatInt(time.Now().UnixNano(), 36)
-	controllerResolver := manual.NewBuilderWithScheme(scheme)
-	w.controllerResolver.Store(controllerResolver)
-
 	if w.conf.WorkerAuthKms == nil || w.conf.DevUsePkiForUpstream {
 		// In this section, we look for existing worker credentials. The two
 		// variables below store whether to create new credentials and whether
@@ -294,13 +291,20 @@ func (w *Worker) Start() error {
 		w.WorkerAuthCurrentKeyId.Store(currentKeyId)
 	}
 
-	if err := w.startListeners(); err != nil {
-		return fmt.Errorf("error starting worker listeners: %w", err)
-	}
 	if err := w.StartControllerConnections(); err != nil {
 		return fmt.Errorf("error making controller connections: %w", err)
 	}
 
+	var err error
+	w.sessionManager, err = session.NewManager(pbs.NewSessionServiceClient(w.GrpcClientConn))
+	if err != nil {
+		return fmt.Errorf("error creating session manager: %w", err)
+	}
+
+	if err := w.startListeners(w.sessionManager); err != nil {
+		return fmt.Errorf("error starting worker listeners: %w", err)
+	}
+
 	// Rather than deal with some of the potential error conditions for Add on
 	// the waitgroup vs. Done (in case a function exits immediately), we will
 	// always start rotation and simply exit early if we're using KMS, and
@@ -308,7 +312,7 @@ func (w *Worker) Start() error {
 	w.tickerWg.Add(2)
 	go func() {
 		defer w.tickerWg.Done()
-		w.startStatusTicking(w.baseContext)
+		w.startStatusTicking(w.baseContext, w.sessionManager)
 	}()
 	go func() {
 		defer w.tickerWg.Done()
@@ -343,7 +347,7 @@ func (w *Worker) Shutdown() error {
 	}
 
 	// Shut down all connections.
-	w.cleanupConnections(w.baseContext, true)
+	w.cleanupConnections(w.baseContext, true, w.sessionManager)
 
 	// Wait for next status request to succeed. Don't wait too long;
 	// wrap the base context in a timeout equal to our status grace
@@ -405,40 +409,6 @@ func (w *Worker) ParseAndStoreTags(incoming map[string][]string) {
 	w.updateTags.Store(true)
 }
 
-// GrpcClientConn returns the underlying client connection
-func (w *Worker) GrpcClientConn() (*grpc.ClientConn, error) {
-	rawConn := w.grpcClientConn.Load()
-	if rawConn == nil {
-		return nil, errors.New("unable to load grpc client conn")
-	}
-	cc, ok := rawConn.(*grpc.ClientConn)
-	if !ok {
-		return nil, fmt.Errorf("invalid grpc client connection %T", rawConn)
-	}
-	if cc == nil {
-		return nil, fmt.Errorf("client connection is nil")
-	}
-
-	return cc, nil
-}
-
-// ControllerSessionConn returns the underlying session service client
-func (w *Worker) ControllerSessionConn() (pbs.SessionServiceClient, error) {
-	rawConn := w.controllerSessionConn.Load()
-	if rawConn == nil {
-		return nil, errors.New("unable to load controller session service connection")
-	}
-	sessClient, ok := rawConn.(pbs.SessionServiceClient)
-	if !ok {
-		return nil, fmt.Errorf("invalid session client %T", rawConn)
-	}
-	if sessClient == nil {
-		return nil, fmt.Errorf("session client is nil")
-	}
-
-	return sessClient, nil
-}
-
 // ControllerServerCoordinationServiceConn returns the underlying server coordination service client
 func (w *Worker) ControllerServerCoordinationServiceConn() (pbs.ServerCoordinationServiceClient, error) {
 	rawConn := w.controllerStatusConn.Load()
@@ -473,125 +443,88 @@ func (w *Worker) ControllerMultihopConn() (multihop.MultihopServiceClient, error
 	return multihopClient, nil
 }
 
-func (w *Worker) getSessionTls(hello *tls.ClientHelloInfo) (*tls.Config, error) {
+func (w *Worker) getSessionTls(sessionManager session.Manager) func(hello *tls.ClientHelloInfo) (*tls.Config, error) {
 	const op = "worker.(Worker).getSessionTls"
-	ctx := w.baseContext
-	var sessionId string
-	switch {
-	case strings.HasPrefix(hello.ServerName, globals.SessionPrefix):
-		sessionId = hello.ServerName
-	default:
-		for _, proto := range hello.SupportedProtos {
-			if strings.HasPrefix(proto, globals.SessionPrefix) {
-				sessionId = proto
-				break
+	return func(hello *tls.ClientHelloInfo) (*tls.Config, error) {
+		ctx := w.baseContext
+		var sessionId string
+		switch {
+		case strings.HasPrefix(hello.ServerName, globals.SessionPrefix):
+			sessionId = hello.ServerName
+		default:
+			for _, proto := range hello.SupportedProtos {
+				if strings.HasPrefix(proto, globals.SessionPrefix) {
+					sessionId = proto
+					break
+				}
 			}
 		}
-	}
 
-	if sessionId == "" {
-		event.WriteSysEvent(ctx, op, "session_id not found in either SNI or ALPN protos", "server_name", hello.ServerName)
-		return nil, fmt.Errorf("could not find session ID in SNI or ALPN protos")
-	}
-
-	conn, err := w.ControllerSessionConn()
-	if err != nil {
-		event.WriteError(ctx, op, err, event.WithInfo("failed to create controller session client"))
-	}
-
-	lastSuccess := w.LastStatusSuccess()
-	if lastSuccess == nil {
-		event.WriteSysEvent(ctx, op, "no last status information found at session acceptance time")
-		return nil, fmt.Errorf("no last status information found at session acceptance time")
-	}
-
-	timeoutContext, cancel := context.WithTimeout(w.baseContext, session.ValidateSessionTimeout)
-	defer cancel()
-
-	resp, err := conn.LookupSession(timeoutContext, &pbs.LookupSessionRequest{
-		SessionId: sessionId,
-		WorkerId:  lastSuccess.WorkerId,
-	})
-	if err != nil {
-		return nil, fmt.Errorf("error validating session: %w", err)
-	}
-
-	if resp.GetExpiration().AsTime().Before(time.Now()) {
-		return nil, fmt.Errorf("session is expired")
-	}
-
-	parsedCert, err := x509.ParseCertificate(resp.GetAuthorization().Certificate)
-	if err != nil {
-		return nil, fmt.Errorf("error parsing session certificate: %w", err)
-	}
-
-	certPool := x509.NewCertPool()
-	certPool.AddCert(parsedCert)
-
-	tlsConf := &tls.Config{
-		Certificates: []tls.Certificate{
-			{
-				Certificate: [][]byte{resp.GetAuthorization().Certificate},
-				PrivateKey:  ed25519.PrivateKey(resp.GetAuthorization().PrivateKey),
-				Leaf:        parsedCert,
-			},
-		},
-		NextProtos: []string{"http/1.1"},
-		MinVersion: tls.VersionTLS13,
-
-		// These two are set this way so we can make use of VerifyConnection,
-		// which we set on this TLS config below. We are not skipping
-		// verification!
-		ClientAuth:         tls.RequireAnyClientCert,
-		InsecureSkipVerify: true,
-	}
-
-	// We disable normal DNS SAN behavior as we don't rely on DNS or IP
-	// addresses for security and want to avoid issues with including localhost
-	// etc.
-	verifyOpts := x509.VerifyOptions{
-		DNSName: sessionId,
-		Roots:   certPool,
-		KeyUsages: []x509.ExtKeyUsage{
-			x509.ExtKeyUsageClientAuth,
-			x509.ExtKeyUsageServerAuth,
-		},
-	}
-	if w.TestOverrideX509VerifyCertPool != nil {
-		verifyOpts.Roots = w.TestOverrideX509VerifyCertPool
-	}
-	if w.TestOverrideX509VerifyDnsName != "" {
-		verifyOpts.DNSName = w.TestOverrideX509VerifyDnsName
-	}
-	tlsConf.VerifyConnection = func(cs tls.ConnectionState) error {
-		// Go will not run this without at least one peer certificate, but
-		// doesn't hurt to check
-		if len(cs.PeerCertificates) == 0 {
-			return errors.New("no peer certificates provided")
+		if sessionId == "" {
+			event.WriteSysEvent(ctx, op, "session_id not found in either SNI or ALPN protos", "server_name", hello.ServerName)
+			return nil, fmt.Errorf("could not find session ID in SNI or ALPN protos")
 		}
-		_, err := cs.PeerCertificates[0].Verify(verifyOpts)
-		return err
-	}
 
-	si := &session.Info{
-		Id:                    resp.GetAuthorization().GetSessionId(),
-		SessionTls:            tlsConf,
-		LookupSessionResponse: resp,
-		Status:                resp.GetStatus(),
-		ConnInfoMap:           make(map[string]*session.ConnInfo),
-	}
-	// TODO: Periodically clean this up. We can't rely on things in here but
-	// not in cancellation because they could be on the way to being
-	// established. However, since cert lifetimes are short, we can simply range
-	// through and remove values that are expired.
-	actualSiRaw, loaded := w.sessionInfoMap.LoadOrStore(sessionId, si)
-	if loaded {
-		// Update the response to the latest
-		actualSi := actualSiRaw.(*session.Info)
-		actualSi.Lock()
-		actualSi.LookupSessionResponse = resp
-		actualSi.Unlock()
-	}
+		lastSuccess := w.LastStatusSuccess()
+		if lastSuccess == nil {
+			event.WriteSysEvent(ctx, op, "no last status information found at session acceptance time")
+			return nil, fmt.Errorf("no last status information found at session acceptance time")
+		}
 
-	return tlsConf, nil
+		timeoutContext, cancel := context.WithTimeout(w.baseContext, session.ValidateSessionTimeout)
+		defer cancel()
+		sess, err := sessionManager.LoadLocalSession(timeoutContext, sessionId, lastSuccess.GetWorkerId())
+		if err != nil {
+			return nil, fmt.Errorf("error refreshing session: %w", err)
+		}
+
+		certPool := x509.NewCertPool()
+		certPool.AddCert(sess.GetCertificate())
+
+		tlsConf := &tls.Config{
+			Certificates: []tls.Certificate{
+				{
+					Certificate: [][]byte{sess.GetCertificate().Raw},
+					PrivateKey:  ed25519.PrivateKey(sess.GetPrivateKey()),
+					Leaf:        sess.GetCertificate(),
+				},
+			},
+			NextProtos: []string{"http/1.1"},
+			MinVersion: tls.VersionTLS13,
+
+			// These two are set this way so we can make use of VerifyConnection,
+			// which we set on this TLS config below. We are not skipping
+			// verification!
+			ClientAuth:         tls.RequireAnyClientCert,
+			InsecureSkipVerify: true,
+		}
+
+		// We disable normal DNS SAN behavior as we don't rely on DNS or IP
+		// addresses for security and want to avoid issues with including localhost
+		// etc.
+		verifyOpts := x509.VerifyOptions{
+			DNSName: sessionId,
+			Roots:   certPool,
+			KeyUsages: []x509.ExtKeyUsage{
+				x509.ExtKeyUsageClientAuth,
+				x509.ExtKeyUsageServerAuth,
+			},
+		}
+		if w.TestOverrideX509VerifyCertPool != nil {
+			verifyOpts.Roots = w.TestOverrideX509VerifyCertPool
+		}
+		if w.TestOverrideX509VerifyDnsName != "" {
+			verifyOpts.DNSName = w.TestOverrideX509VerifyDnsName
+		}
+		tlsConf.VerifyConnection = func(cs tls.ConnectionState) error {
+			// Go will not run this without at least one peer certificate, but
+			// doesn't hurt to check
+			if len(cs.PeerCertificates) == 0 {
+				return errors.New("no peer certificates provided")
+			}
+			_, err := cs.PeerCertificates[0].Verify(verifyOpts)
+			return err
+		}
+		return tlsConf, nil
+	}
 }

internal/daemon/worker/worker_proxy_service.go

@@ -5,6 +5,7 @@ import (
 	"sync/atomic"
 
 	pbs "github.com/hashicorp/boundary/internal/gen/controller/servers/services"
+	"google.golang.org/grpc"
 )
 
 type workerProxyServiceServer struct {
@@ -12,16 +13,16 @@ type workerProxyServiceServer struct {
 	pbs.UnimplementedSessionServiceServer
 
 	scsClient *atomic.Value
-	ssClient  *atomic.Value
+	ssClient  pbs.SessionServiceClient
 }
 
 func NewWorkerProxyServiceServer(
+	cc *grpc.ClientConn,
 	scsClient *atomic.Value,
-	ssClient *atomic.Value,
 ) *workerProxyServiceServer {
 	return &workerProxyServiceServer{
 		scsClient: scsClient,
-		ssClient:  ssClient,
+		ssClient:  pbs.NewSessionServiceClient(cc),
 	}
 }
 
@@ -43,25 +44,25 @@ func (ws *workerProxyServiceServer) Status(ctx context.Context, req *pbs.StatusR
 }
 
 func (ws *workerProxyServiceServer) LookupSession(ctx context.Context, req *pbs.LookupSessionRequest) (*pbs.LookupSessionResponse, error) {
-	return ws.ssClient.Load().(pbs.SessionServiceClient).LookupSession(ctx, req)
+	return ws.ssClient.LookupSession(ctx, req)
 }
 
 func (ws *workerProxyServiceServer) CancelSession(ctx context.Context, req *pbs.CancelSessionRequest) (*pbs.CancelSessionResponse, error) {
-	return ws.ssClient.Load().(pbs.SessionServiceClient).CancelSession(ctx, req)
+	return ws.ssClient.CancelSession(ctx, req)
 }
 
 func (ws *workerProxyServiceServer) ActivateSession(ctx context.Context, req *pbs.ActivateSessionRequest) (*pbs.ActivateSessionResponse, error) {
-	return ws.ssClient.Load().(pbs.SessionServiceClient).ActivateSession(ctx, req)
+	return ws.ssClient.ActivateSession(ctx, req)
 }
 
 func (ws *workerProxyServiceServer) AuthorizeConnection(ctx context.Context, req *pbs.AuthorizeConnectionRequest) (*pbs.AuthorizeConnectionResponse, error) {
-	return ws.ssClient.Load().(pbs.SessionServiceClient).AuthorizeConnection(ctx, req)
+	return ws.ssClient.AuthorizeConnection(ctx, req)
 }
 
 func (ws *workerProxyServiceServer) ConnectConnection(ctx context.Context, req *pbs.ConnectConnectionRequest) (*pbs.ConnectConnectionResponse, error) {
-	return ws.ssClient.Load().(pbs.SessionServiceClient).ConnectConnection(ctx, req)
+	return ws.ssClient.ConnectConnection(ctx, req)
 }
 
 func (ws *workerProxyServiceServer) CloseConnection(ctx context.Context, req *pbs.CloseConnectionRequest) (*pbs.CloseConnectionResponse, error) {
-	return ws.ssClient.Load().(pbs.SessionServiceClient).CloseConnection(ctx, req)
+	return ws.ssClient.CloseConnection(ctx, req)
 }

internal/gen/controller/servers/services/testing.go

@@ -6,36 +6,60 @@ import (
 	"google.golang.org/grpc"
 )
 
-// TODO (lcr 08/2021): implement mockSessionServiceClient outside of primitive support for
-// ConnectConnection
-type mockSessionServiceClient struct{}
+type mockSessionServiceClient struct {
+	LookupSessionFn       func(context.Context, *LookupSessionRequest) (*LookupSessionResponse, error)
+	ActivateSessionFn     func(context.Context, *ActivateSessionRequest) (*ActivateSessionResponse, error)
+	CancelSessionFn       func(context.Context, *CancelSessionRequest) (*CancelSessionResponse, error)
+	AuthorizeConnectionFn func(context.Context, *AuthorizeConnectionRequest) (*AuthorizeConnectionResponse, error)
+	ConnectConnectionFn   func(context.Context, *ConnectConnectionRequest) (*ConnectConnectionResponse, error)
+	CloseConnectionFn     func(context.Context, *CloseConnectionRequest) (*CloseConnectionResponse, error)
+}
 
-func NewMockSessionServiceClient() SessionServiceClient {
+// NewMockSessionServiceClient returns a mock SessionServiceClient which allows
+// the mocking out specific calls to a SessionService by assigning values to
+// the respective mock client member variables.
+func NewMockSessionServiceClient() *mockSessionServiceClient {
 	return &mockSessionServiceClient{}
 }
 
-func (c *mockSessionServiceClient) LookupSession(_ context.Context, _ *LookupSessionRequest, _ ...grpc.CallOption) (*LookupSessionResponse, error) {
+func (c *mockSessionServiceClient) LookupSession(ctx context.Context, req *LookupSessionRequest, _ ...grpc.CallOption) (*LookupSessionResponse, error) {
+	if c.LookupSessionFn != nil {
+		return c.LookupSessionFn(ctx, req)
+	}
 	panic("not implemented")
 }
 
-func (c *mockSessionServiceClient) ActivateSession(_ context.Context, _ *ActivateSessionRequest, _ ...grpc.CallOption) (*ActivateSessionResponse, error) {
+func (c *mockSessionServiceClient) ActivateSession(ctx context.Context, req *ActivateSessionRequest, _ ...grpc.CallOption) (*ActivateSessionResponse, error) {
+	if c.ActivateSessionFn != nil {
+		return c.ActivateSessionFn(ctx, req)
+	}
 	panic("not implemented")
 }
 
-func (c *mockSessionServiceClient) CancelSession(_ context.Context, _ *CancelSessionRequest, _ ...grpc.CallOption) (*CancelSessionResponse, error) {
+func (c *mockSessionServiceClient) CancelSession(ctx context.Context, req *CancelSessionRequest, _ ...grpc.CallOption) (*CancelSessionResponse, error) {
+	if c.CancelSessionFn != nil {
+		return c.CancelSessionFn(ctx, req)
+	}
 	panic("not implemented")
 }
 
-func (c *mockSessionServiceClient) AuthorizeConnection(_ context.Context, _ *AuthorizeConnectionRequest, _ ...grpc.CallOption) (*AuthorizeConnectionResponse, error) {
+func (c *mockSessionServiceClient) AuthorizeConnection(ctx context.Context, req *AuthorizeConnectionRequest, _ ...grpc.CallOption) (*AuthorizeConnectionResponse, error) {
+	if c.AuthorizeConnectionFn != nil {
+		return c.AuthorizeConnectionFn(ctx, req)
+	}
 	panic("not implemented")
 }
 
-func (c *mockSessionServiceClient) ConnectConnection(_ context.Context, _ *ConnectConnectionRequest, _ ...grpc.CallOption) (*ConnectConnectionResponse, error) {
-	return &ConnectConnectionResponse{
-		Status: CONNECTIONSTATUS_CONNECTIONSTATUS_CONNECTED,
-	}, nil
-}
-
-func (c *mockSessionServiceClient) CloseConnection(_ context.Context, _ *CloseConnectionRequest, _ ...grpc.CallOption) (*CloseConnectionResponse, error) {
+func (c *mockSessionServiceClient) ConnectConnection(ctx context.Context, req *ConnectConnectionRequest, _ ...grpc.CallOption) (*ConnectConnectionResponse, error) {
+	if c.ConnectConnectionFn != nil {
+		return c.ConnectConnectionFn(ctx, req)
+	}
+	panic("not implemented")
+}
+
+func (c *mockSessionServiceClient) CloseConnection(ctx context.Context, req *CloseConnectionRequest, _ ...grpc.CallOption) (*CloseConnectionResponse, error) {
+	if c.CloseConnectionFn != nil {
+		return c.CloseConnectionFn(ctx, req)
+	}
 	panic("not implemented")
 }