fix(session connections): Create session connection repository

internal/cmd/commands/server/worker_shutdown_reload_test.go

@@ -163,7 +163,7 @@ func TestServer_ShutdownWorker(t *testing.T) {
 
 	// Connection should fail, and the session should be closed on the DB.
 	sConn.TestSendRecvFail(t)
-	sess.ExpectConnectionStateOnController(ctx, t, controllerCmd.controller.SessionRepoFn, session.StatusClosed)
+	sess.ExpectConnectionStateOnController(ctx, t, controllerCmd.controller.ConnectionRepoFn, session.StatusClosed)
 
 	// We're done! Shutdown the controller, and that's it.
 	close(controllerCmd.ShutdownCh)

internal/db/schema/migrations/oss/postgres_24_01_test.go

@@ -123,6 +123,7 @@ func TestMigrations_SessionState(t *testing.T) {
 	require.Equal(want, state)
 
 	sessionRepo, err := session.NewRepository(rw, rw, kmsCache)
+	connectionRepo, err := session.NewConnectionRepository(ctx, rw, rw, kmsCache)
 	require.NoError(err)
 
 	// Ensure session is cancelled
@@ -138,7 +139,7 @@ func TestMigrations_SessionState(t *testing.T) {
 	require.Equal([]string{"canceled"}, sessionTermReason)
 
 	// Ensure connection is also cancelled
-	connections, err := sessionRepo.ListConnectionsBySessionId(ctx, repoSessionId)
+	connections, err := connectionRepo.ListConnectionsBySessionId(ctx, repoSessionId)
 	require.NoError(err)
 
 	var connTermReason []string

internal/servers/controller/common/common.go

@@ -24,5 +24,6 @@ type (
 	PluginHostRepoFactory      func() (*pluginhost.Repository, error)
 	HostPluginRepoFactory      func() (*hostplugin.Repository, error)
 	SessionRepoFactory         func() (*session.Repository, error)
+	ConnectionRepoFactory      func() (*session.ConnectionRepository, error)
 	TargetRepoFactory          func() (*target.Repository, error)
 )

internal/servers/controller/controller.go

@@ -68,6 +68,7 @@ type Controller struct {
 	PasswordAuthRepoFn    common.PasswordAuthRepoFactory
 	ServersRepoFn         common.ServersRepoFactory
 	SessionRepoFn         common.SessionRepoFactory
+	ConnectionRepoFn      common.ConnectionRepoFactory
 	StaticHostRepoFn      common.StaticRepoFactory
 	PluginHostRepoFn      common.PluginHostRepoFactory
 	HostPluginRepoFn      common.HostPluginRepoFactory
@@ -235,7 +236,9 @@ func New(ctx context.Context, conf *Config) (*Controller, error) {
 	c.SessionRepoFn = func() (*session.Repository, error) {
 		return session.NewRepository(dbase, dbase, c.kms)
 	}
-
+	c.ConnectionRepoFn = func() (*session.ConnectionRepository, error) {
+		return session.NewConnectionRepository(ctx, dbase, dbase, c.kms)
+	}
 	return c, nil
 }
 
@@ -290,21 +293,21 @@ func (c *Controller) registerJobs() error {
 		return err
 	}
 
-	if err := c.registerSessionCleanupJob(); err != nil {
+	if err := c.registerSessionConnectionCleanupJob(); err != nil {
 		return err
 	}
 
 	return nil
 }
 
-// registerSessionCleanupJob is a helper method to abstract
-// registering the session cleanup job specifically.
-func (c *Controller) registerSessionCleanupJob() error {
-	sessionCleanupJob, err := newSessionCleanupJob(c.SessionRepoFn, int(c.conf.StatusGracePeriodDuration.Seconds()))
+// registerSessionConnectionCleanupJob is a helper method to abstract
+// registering the session connection cleanup job specifically.
+func (c *Controller) registerSessionConnectionCleanupJob() error {
+	sessionConnectionCleanupJob, err := newSessionConnectionCleanupJob(c.ConnectionRepoFn, int(c.conf.StatusGracePeriodDuration.Seconds()))
 	if err != nil {
 		return fmt.Errorf("error creating session cleanup job: %w", err)
 	}
-	if err = c.scheduler.RegisterJob(c.baseContext, sessionCleanupJob); err != nil {
+	if err = c.scheduler.RegisterJob(c.baseContext, sessionConnectionCleanupJob); err != nil {
 		return fmt.Errorf("error registering session cleanup job: %w", err)
 	}
 

internal/servers/controller/handlers/targets/tcp/target_service_test.go

@@ -2628,6 +2628,9 @@ func TestAuthorizeSession(t *testing.T) {
 	sessionRepoFn := func() (*session.Repository, error) {
 		return session.NewRepository(rw, rw, kms)
 	}
+	connectionRepoFn := func() (*session.ConnectionRepository, error) {
+		return session.NewConnectionRepository(ctx, rw, rw, kms)
+	}
 	staticHostRepoFn := func() (*static.Repository, error) {
 		return static.NewRepository(rw, rw, kms)
 	}
@@ -2758,7 +2761,7 @@ func TestAuthorizeSession(t *testing.T) {
 			require.NoError(t, err)
 
 			// Tell our DB that there is a worker ready to serve the data
-			workerService := workers.NewWorkerServiceServer(serversRepoFn, sessionRepoFn, &sync.Map{}, kms)
+			workerService := workers.NewWorkerServiceServer(serversRepoFn, sessionRepoFn, connectionRepoFn, &sync.Map{}, kms)
 			_, err = workerService.Status(ctx, &spbs.StatusRequest{
 				Worker: &spb.Server{
 					PrivateId: "testworker",
@@ -2874,6 +2877,9 @@ func TestAuthorizeSessionTypedCredentials(t *testing.T) {
 	sessionRepoFn := func() (*session.Repository, error) {
 		return session.NewRepository(rw, rw, kms)
 	}
+	connectionRepoFn := func() (*session.ConnectionRepository, error) {
+		return session.NewConnectionRepository(ctx, rw, rw, kms)
+	}
 	staticHostRepoFn := func() (*static.Repository, error) {
 		return static.NewRepository(rw, rw, kms)
 	}
@@ -3037,7 +3043,7 @@ func TestAuthorizeSessionTypedCredentials(t *testing.T) {
 			require.NoError(t, err)
 
 			// Tell our DB that there is a worker ready to serve the data
-			workerService := workers.NewWorkerServiceServer(serversRepoFn, sessionRepoFn, &sync.Map{}, kms)
+			workerService := workers.NewWorkerServiceServer(serversRepoFn, sessionRepoFn, connectionRepoFn, &sync.Map{}, kms)
 			_, err = workerService.Status(ctx, &spbs.StatusRequest{
 				Worker: &spb.Server{
 					PrivateId: "testworker",
@@ -3126,6 +3132,9 @@ func TestAuthorizeSession_Errors(t *testing.T) {
 	sessionRepoFn := func() (*session.Repository, error) {
 		return session.NewRepository(rw, rw, kms)
 	}
+	connectionRepoFn := func() (*session.ConnectionRepository, error) {
+		return session.NewConnectionRepository(ctx, rw, rw, kms)
+	}
 	staticHostRepoFn := func() (*static.Repository, error) {
 		return static.NewRepository(rw, rw, kms)
 	}
@@ -3165,7 +3174,7 @@ func TestAuthorizeSession_Errors(t *testing.T) {
 	store := vault.TestCredentialStore(t, conn, wrapper, proj.GetPublicId(), v.Addr, tok, sec.Auth.Accessor)
 
 	workerExists := func(tar target.Target) (version uint32) {
-		workerService := workers.NewWorkerServiceServer(serversRepoFn, sessionRepoFn, &sync.Map{}, kms)
+		workerService := workers.NewWorkerServiceServer(serversRepoFn, sessionRepoFn, connectionRepoFn, &sync.Map{}, kms)
 		_, err := workerService.Status(context.Background(), &spbs.StatusRequest{
 			Worker: &spb.Server{
 				PrivateId: "testworker",

internal/servers/controller/handlers/workers/worker_service.go

@@ -26,22 +26,25 @@ type workerServiceServer struct {
 	pbs.UnimplementedServerCoordinationServiceServer
 	pbs.UnimplementedSessionServiceServer
 
-	serversRepoFn common.ServersRepoFactory
-	sessionRepoFn common.SessionRepoFactory
-	updateTimes   *sync.Map
-	kms           *kms.Kms
+	serversRepoFn    common.ServersRepoFactory
+	sessionRepoFn    common.SessionRepoFactory
+	connectionRepoFn common.ConnectionRepoFactory
+	updateTimes      *sync.Map
+	kms              *kms.Kms
 }
 
 func NewWorkerServiceServer(
 	serversRepoFn common.ServersRepoFactory,
 	sessionRepoFn common.SessionRepoFactory,
+	connectionRepoFn common.ConnectionRepoFactory,
 	updateTimes *sync.Map,
 	kms *kms.Kms) *workerServiceServer {
 	return &workerServiceServer{
-		serversRepoFn: serversRepoFn,
-		sessionRepoFn: sessionRepoFn,
-		updateTimes:   updateTimes,
-		kms:           kms,
+		serversRepoFn:    serversRepoFn,
+		sessionRepoFn:    sessionRepoFn,
+		connectionRepoFn: connectionRepoFn,
+		updateTimes:      updateTimes,
+		kms:              kms,
 	}
 }
 
@@ -61,6 +64,7 @@ func (ws *workerServiceServer) Status(ctx context.Context, req *pbs.StatusReques
 		return &pbs.StatusResponse{}, status.Errorf(codes.Internal, "Error acquiring repo to store worker status: %v", err)
 	}
 	sessRepo, err := ws.sessionRepoFn()
+	connectionRepo, err := ws.connectionRepoFn()
 	if err != nil {
 		event.WriteError(ctx, op, err, event.WithInfoMsg("error getting sessions repo"))
 		return &pbs.StatusResponse{}, status.Errorf(codes.Internal, "Error acquiring repo to query session status: %v", err)
@@ -173,7 +177,7 @@ func (ws *workerServiceServer) Status(ctx context.Context, req *pbs.StatusReques
 	// additive to it, we don't add sessions that have already been
 	// added there as canceling sessions already closes the
 	// connections.
-	shouldCloseConnections, err := sessRepo.ShouldCloseConnectionsOnWorker(ctx, reportedOpenConns, requestedSessionCancelIds)
+	shouldCloseConnections, err := connectionRepo.ShouldCloseConnectionsOnWorker(ctx, reportedOpenConns, requestedSessionCancelIds)
 	if err != nil {
 		return nil, status.Errorf(codes.Internal, "Error fetching connections that should be closed: %v", err)
 	}
@@ -203,7 +207,7 @@ func (ws *workerServiceServer) Status(ctx context.Context, req *pbs.StatusReques
 	}
 
 	// Run our controller-side cleanup function.
-	closedConns, err := sessRepo.CloseDeadConnectionsForWorker(ctx, req.Worker.PrivateId, reportedOpenConns)
+	closedConns, err := connectionRepo.CloseDeadConnectionsForWorker(ctx, req.Worker.PrivateId, reportedOpenConns)
 	if err != nil {
 		return nil, status.Errorf(codes.Internal, "Error closing dead conns for worker %s: %v", req.Worker.PrivateId, err)
 	}
@@ -386,12 +390,17 @@ func (ws *workerServiceServer) ActivateSession(ctx context.Context, req *pbs.Act
 
 func (ws *workerServiceServer) AuthorizeConnection(ctx context.Context, req *pbs.AuthorizeConnectionRequest) (*pbs.AuthorizeConnectionResponse, error) {
 	const op = "workers.(workerServiceServer).AuthorizeConnection"
-	sessRepo, err := ws.sessionRepoFn()
+	connectionRepo, err := ws.connectionRepoFn()
 	if err != nil {
 		return nil, status.Errorf(codes.Internal, "error getting session repo: %v", err)
 	}
 
-	connectionInfo, connStates, authzSummary, err := sessRepo.AuthorizeConnection(ctx, req.GetSessionId(), req.GetWorkerId())
+	sessionRepo, err := ws.sessionRepoFn()
+	if err != nil {
+		return nil, status.Errorf(codes.Internal, "error getting session repo: %v", err)
+	}
+
+	connectionInfo, connStates, authzSummary, err := session.AuthorizeConnection(ctx, sessionRepo, connectionRepo, req.GetSessionId(), req.GetWorkerId())
 	if err != nil {
 		return nil, err
 	}

internal/servers/controller/handlers/workers/worker_service_test.go

@@ -37,6 +37,9 @@ func TestLookupSession(t *testing.T) {
 	sessionRepoFn := func() (*session.Repository, error) {
 		return session.NewRepository(rw, rw, kms)
 	}
+	connectionRepoFn := func() (*session.ConnectionRepository, error) {
+		return session.NewConnectionRepository(ctx, rw, rw, kms)
+	}
 
 	at := authtoken.TestAuthToken(t, conn, kms, org.GetPublicId())
 	uId := at.GetIamUserId()
@@ -97,7 +100,7 @@ func TestLookupSession(t *testing.T) {
 	err = repo.AddSessionCredentials(ctx, egressSess.ScopeId, egressSess.GetPublicId(), workerCreds)
 	require.NoError(t, err)
 
-	s := workers.NewWorkerServiceServer(serversRepoFn, sessionRepoFn, new(sync.Map), kms)
+	s := workers.NewWorkerServiceServer(serversRepoFn, sessionRepoFn, connectionRepoFn, new(sync.Map), kms)
 	require.NotNil(t, s)
 
 	cases := []struct {

internal/servers/controller/listeners.go

@@ -123,7 +123,8 @@ func (c *Controller) startListeners(ctx context.Context) error {
 				),
 			),
 		)
-		workerService := workers.NewWorkerServiceServer(c.ServersRepoFn, c.SessionRepoFn, c.workerStatusUpdateTimes, c.kms)
+		workerService := workers.NewWorkerServiceServer(c.ServersRepoFn, c.SessionRepoFn, c.ConnectionRepoFn,
+			c.workerStatusUpdateTimes, c.kms)
 		pbs.RegisterServerCoordinationServiceServer(workerServer, workerService)
 		pbs.RegisterSessionServiceServer(workerServer, workerService)
 

internal/servers/controller/session_cleanup_job.go

@@ -13,7 +13,7 @@ import (
 	"github.com/hashicorp/boundary/internal/session"
 )
 
-// sessionCleanupJob defines a periodic job that monitors workers for
+// sessionConnectionCleanupJob defines a periodic job that monitors workers for
 // loss of connection and terminates connections on workers that have
 // not sent a heartbeat in a significant period of time.
 //
@@ -22,8 +22,8 @@ import (
 // worker, or the event of a synchronization issue between the two,
 // the controller will win out and order that the connections be
 // closed on the worker.
-type sessionCleanupJob struct {
-	sessionRepoFn common.SessionRepoFactory
+type sessionConnectionCleanupJob struct {
+	connectionRepoFn common.ConnectionRepoFactory
 
 	// The amount of time to give disconnected workers before marking
 	// their connections as closed.
@@ -33,47 +33,47 @@ type sessionCleanupJob struct {
 	totalClosed int
 }
 
-// newSessionCleanupJob instantiates the session cleanup job.
-func newSessionCleanupJob(
-	sessionRepoFn common.SessionRepoFactory,
+// newSessionConnectionCleanupJob instantiates the session cleanup job.
+func newSessionConnectionCleanupJob(
+	connectionRepoFn common.ConnectionRepoFactory,
 	gracePeriod int,
-) (*sessionCleanupJob, error) {
-	const op = "controller.newNewSessionCleanupJob"
+) (*sessionConnectionCleanupJob, error) {
+	const op = "controller.newNewSessionConnectionCleanupJob"
 	switch {
-	case sessionRepoFn == nil:
-		return nil, errors.NewDeprecated(errors.InvalidParameter, op, "missing sessionRepoFn")
+	case connectionRepoFn == nil:
+		return nil, errors.NewDeprecated(errors.InvalidParameter, op, "missing connectionRepoFn")
 	case gracePeriod < session.DeadWorkerConnCloseMinGrace:
 		return nil, errors.NewDeprecated(
 			errors.InvalidParameter, op, fmt.Sprintf("invalid gracePeriod, must be greater than %d", session.DeadWorkerConnCloseMinGrace))
 	}
 
-	return &sessionCleanupJob{
-		sessionRepoFn: sessionRepoFn,
-		gracePeriod:   gracePeriod,
+	return &sessionConnectionCleanupJob{
+		connectionRepoFn: connectionRepoFn,
+		gracePeriod:      gracePeriod,
 	}, nil
 }
 
 // Name returns a short, unique name for the job.
-func (j *sessionCleanupJob) Name() string { return "session_cleanup" }
+func (j *sessionConnectionCleanupJob) Name() string { return "session_cleanup" }
 
 // Description returns the description for the job.
-func (j *sessionCleanupJob) Description() string {
+func (j *sessionConnectionCleanupJob) Description() string {
 	return "Clean up session connections from disconnected workers"
 }
 
 // NextRunIn returns the next run time after a job is completed.
 //
-// The next run time is defined for sessionCleanupJob as one second.
+// The next run time is defined for sessionConnectionCleanupJob as one second.
 // This is because the job should run continuously to terminate
 // connections as soon as a worker has not reported in for a long
 // enough time. Only one job will ever run at once, so there is no
 // reason why it cannot run again immediately.
-func (j *sessionCleanupJob) NextRunIn(_ context.Context) (time.Duration, error) {
+func (j *sessionConnectionCleanupJob) NextRunIn(_ context.Context) (time.Duration, error) {
 	return time.Second, nil
 }
 
 // Status returns the status of the running job.
-func (j *sessionCleanupJob) Status() scheduler.JobStatus {
+func (j *sessionConnectionCleanupJob) Status() scheduler.JobStatus {
 	return scheduler.JobStatus{
 		Completed: j.totalClosed,
 		Total:     j.totalClosed,
@@ -81,18 +81,18 @@ func (j *sessionCleanupJob) Status() scheduler.JobStatus {
 }
 
 // Run executes the job.
-func (j *sessionCleanupJob) Run(ctx context.Context) error {
-	const op = "controller.(sessionCleanupJob).Run"
+func (j *sessionConnectionCleanupJob) Run(ctx context.Context) error {
+	const op = "controller.(sessionConnectionCleanupJob).Run"
 	j.totalClosed = 0
 
 	// Load repos.
-	sessionRepo, err := j.sessionRepoFn()
+	connectionRepo, err := j.connectionRepoFn()
 	if err != nil {
 		return errors.Wrap(ctx, err, op, errors.WithMsg("error getting session repo"))
 	}
 
 	// Run the atomic dead worker cleanup job.
-	results, err := sessionRepo.CloseConnectionsForDeadWorkers(ctx, j.gracePeriod)
+	results, err := connectionRepo.CloseConnectionsForDeadWorkers(ctx, j.gracePeriod)
 	if err != nil {
 		return errors.Wrap(ctx, err, op)
 	}

internal/servers/controller/session_cleanup_job_test.go

@@ -18,13 +18,14 @@ import (
 )
 
 // assert the interface
-var _ = scheduler.Job(new(sessionCleanupJob))
+var _ = scheduler.Job(new(sessionConnectionCleanupJob))
 
 // This test has been largely adapted from
 // TestRepository_CloseDeadConnectionsOnWorker in
 // internal/session/repository_connection_test.go.
-func TestSessionCleanupJob(t *testing.T) {
+func TestSessionConnectionCleanupJob(t *testing.T) {
 	t.Parallel()
+	ctx := context.Background()
 	require, assert := require.New(t), assert.New(t)
 	conn, _ := db.TestSetup(t, "postgres")
 	rw := db.New(conn)
@@ -34,8 +35,9 @@ func TestSessionCleanupJob(t *testing.T) {
 	serversRepo, err := servers.NewRepository(rw, rw, kms)
 	require.NoError(err)
 	sessionRepo, err := session.NewRepository(rw, rw, kms)
+	connectionRepo, err := session.NewConnectionRepository(ctx, rw, rw, kms)
 	require.NoError(err)
-	ctx := context.Background()
+
 	numConns := 12
 
 	// Create two "workers". One will remain untouched while the other "goes
@@ -54,7 +56,7 @@ func TestSessionCleanupJob(t *testing.T) {
 		sess := session.TestDefaultSession(t, conn, wrapper, iamRepo, session.WithServerId(serverId), session.WithDbOpts(db.WithSkipVetForWrite(true)))
 		sess, _, err = sessionRepo.ActivateSession(ctx, sess.GetPublicId(), sess.Version, serverId, "worker", []byte("foo"))
 		require.NoError(err)
-		c, cs, _, err := sessionRepo.AuthorizeConnection(ctx, sess.GetPublicId(), serverId)
+		c, cs, _, err := session.AuthorizeConnection(ctx, sessionRepo, connectionRepo, sess.GetPublicId(), serverId)
 		require.NoError(err)
 		require.Len(cs, 1)
 		require.Equal(session.StatusAuthorized, cs[0].Status)
@@ -95,8 +97,8 @@ func TestSessionCleanupJob(t *testing.T) {
 	}
 
 	// Create the job.
-	job, err := newSessionCleanupJob(
-		func() (*session.Repository, error) { return sessionRepo, nil },
+	job, err := newSessionConnectionCleanupJob(
+		func() (*session.ConnectionRepository, error) { return connectionRepo, nil },
 		session.DeadWorkerConnCloseMinGrace,
 	)
 	require.NoError(err)
@@ -119,7 +121,7 @@ func TestSessionCleanupJob(t *testing.T) {
 		require.True(ok)
 		require.Len(connIds, 6)
 		for _, connId := range connIds {
-			_, states, err := sessionRepo.LookupConnection(ctx, connId, nil)
+			_, states, err := connectionRepo.LookupConnection(ctx, connId, nil)
 			require.NoError(err)
 			var foundClosed bool
 			for _, state := range states {
@@ -138,22 +140,22 @@ func TestSessionCleanupJob(t *testing.T) {
 	assertConnections(worker1.PrivateId, false)
 }
 
-func TestSessionCleanupJobNewJobErr(t *testing.T) {
+func TestSessionConnectionCleanupJobNewJobErr(t *testing.T) {
 	t.Parallel()
 	ctx := context.TODO()
-	const op = "controller.newNewSessionCleanupJob"
+	const op = "controller.newNewSessionConnectionCleanupJob"
 	require := require.New(t)
 
-	job, err := newSessionCleanupJob(nil, 0)
+	job, err := newSessionConnectionCleanupJob(nil, 0)
 	require.Equal(err, errors.E(
 		ctx,
 		errors.WithCode(errors.InvalidParameter),
 		errors.WithOp(op),
-		errors.WithMsg("missing sessionRepoFn"),
+		errors.WithMsg("missing connectionRepoFn"),
 	))
 	require.Nil(job)
 
-	job, err = newSessionCleanupJob(func() (*session.Repository, error) { return nil, nil }, 0)
+	job, err = newSessionConnectionCleanupJob(func() (*session.ConnectionRepository, error) { return nil, nil }, 0)
 	require.Equal(err, errors.E(
 		ctx,
 		errors.WithCode(errors.InvalidParameter),

internal/servers/controller/testing.go

@@ -10,6 +10,8 @@ import (
 	"testing"
 	"time"
 
+	"github.com/hashicorp/boundary/internal/session"
+
 	grpc_middleware "github.com/grpc-ecosystem/go-grpc-middleware"
 	"github.com/hashicorp/boundary/api"
 	"github.com/hashicorp/boundary/api/authmethods"
@@ -116,6 +118,14 @@ func (tc *TestController) ServersRepo() *servers.Repository {
 	return repo
 }
 
+func (tc *TestController) ConnectionsRepo() *session.ConnectionRepository {
+	repo, err := tc.c.ConnectionRepoFn()
+	if err != nil {
+		tc.t.Fatal(err)
+	}
+	return repo
+}
+
 func (tc *TestController) Cancel() {
 	tc.cancel()
 }

internal/session/query.go

@@ -58,43 +58,44 @@ where
 	) 
 `
 	authorizeConnectionCte = `
-insert into session_connection (
-	session_id, 
-	public_id,
-	server_id
-)
-with active_session as ( 
+with connections_available as (
 	select 
-		@session_id as session_id,
+		s.public_id
+ 	from 
+		session s 
+	where 
+		s.public_id = @session_id and 
+ 		(s.connection_limit = -1 or 
+		s.connection_limit > (select count(*) from session_connection sc where sc.session_id = @session_id ))
+),
+unexpired_session as (
+	select 
+		s.public_id 
+	from 
+		session s 
+	where 
+		s.public_id in (select * from  connections_available) and 
+		s.expiration_time > now() 
+),
+active_session as (
+	select 
+		ss.session_id as session_id,
 		@public_id as public_id,
 		@worker_id as server_id
-	from
-		session s
-	where
-		-- check that the session hasn't expired.
-		s.expiration_time > now() and
-		-- check that there are still connections available. connection_limit of -1 equals unlimited connections
-		(
-			s.connection_limit = -1
-				or 
-			s.connection_limit > (select count(*) from session_connection sc where sc.session_id = @session_id)
-		) and
-		-- check that there's a state of active
-		s.public_id in (
-			select 
-				ss.session_id 
-			from 
-				session_state ss
-			where 
-				ss.session_id = @session_id and 
-				ss.state = 'active' and
-				-- if there's no end_time, then this is the current state.
-				ss.end_time is null 
-		) 
+	from 
+		session_state ss 
+	where 
+		ss.session_id in (select * from unexpired_session) and 
+		ss.state = 'active' and
+		ss.end_time is null 
+) 
+insert into session_connection (
+  	session_id, 
+ 	public_id, 
+	server_id
 )
 select * from active_session;
 `
-
 	remainingConnectionsCte = `
 with
 session_connection_count(current_connection_count) as (

internal/session/repository_connection.go

@@ -7,11 +7,128 @@ import (
 	"strings"
 	"time"
 
+	"github.com/hashicorp/boundary/internal/kms"
+
 	"github.com/hashicorp/boundary/internal/db"
 	"github.com/hashicorp/boundary/internal/errors"
 	"github.com/hashicorp/boundary/internal/servers"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
 )
 
+// ConnectionRepository is the session connection database repository.
+type ConnectionRepository struct {
+	reader db.Reader
+	writer db.Writer
+	kms    *kms.Kms
+
+	// defaultLimit provides a default for limiting the number of results returned from the repo
+	defaultLimit int
+}
+
+// NewConnectionRepository creates a new session Connection Repository. Supports the options: WithLimit
+// which sets a default limit on results returned by repo operations.
+func NewConnectionRepository(ctx context.Context, r db.Reader, w db.Writer, kms *kms.Kms, opt ...Option) (*ConnectionRepository, error) {
+	const op = "sessionConnection.NewRepository"
+	if r == nil {
+		return nil, errors.New(ctx, errors.InvalidParameter, op, "nil reader")
+	}
+	if w == nil {
+		return nil, errors.New(ctx, errors.InvalidParameter, op, "nil writer")
+	}
+	if kms == nil {
+		return nil, errors.New(ctx, errors.InvalidParameter, op, "nil kms")
+	}
+	opts := getOpts(opt...)
+	if opts.withLimit == 0 {
+		// zero signals the boundary defaults should be used.
+		opts.withLimit = db.DefaultLimit
+	}
+	return &ConnectionRepository{
+		reader:       r,
+		writer:       w,
+		kms:          kms,
+		defaultLimit: opts.withLimit,
+	}, nil
+}
+
+// list will return a listing of resources and honor the WithLimit option or the
+// repo defaultLimit.  Supports WithOrder option.
+func (r *ConnectionRepository) list(ctx context.Context, resources interface{}, where string, args []interface{}, opt ...Option) error {
+	const op = "session.(ConnectionRepository).list"
+	opts := getOpts(opt...)
+	limit := r.defaultLimit
+	var dbOpts []db.Option
+	if opts.withLimit != 0 {
+		// non-zero signals an override of the default limit for the repo.
+		limit = opts.withLimit
+	}
+	dbOpts = append(dbOpts, db.WithLimit(limit))
+	switch opts.withOrderByCreateTime {
+	case db.AscendingOrderBy:
+		dbOpts = append(dbOpts, db.WithOrder("create_time asc"))
+	case db.DescendingOrderBy:
+		dbOpts = append(dbOpts, db.WithOrder("create_time"))
+	}
+	if err := r.reader.SearchWhere(ctx, resources, where, args, dbOpts...); err != nil {
+		return errors.Wrap(ctx, err, op)
+	}
+	return nil
+}
+
+// AuthorizeConnection will check to see if a connection is allowed.  Currently,
+// that authorization checks:
+// * the hasn't expired based on the session.Expiration
+// * number of connections already created is less than session.ConnectionLimit
+// If authorization is success, it creates/stores a new connection in the repo
+// and returns it, along with its states.  If the authorization fails, it
+// an error with Code InvalidSessionState.
+func (r *ConnectionRepository) AuthorizeConnection(ctx context.Context, sessionId, workerId string) (*Connection, []*ConnectionState, error) {
+	const op = "session.(ConnectionRepository).AuthorizeConnection"
+	if sessionId == "" {
+		return nil, nil, errors.Wrap(ctx, status.Error(codes.FailedPrecondition, "missing session id"), op, errors.WithCode(errors.InvalidParameter))
+	}
+	connectionId, err := newConnectionId()
+	if err != nil {
+		return nil, nil, errors.Wrap(ctx, err, op)
+	}
+
+	connection := AllocConnection()
+	connection.PublicId = connectionId
+	var connectionStates []*ConnectionState
+	_, err = r.writer.DoTx(
+		ctx,
+		db.StdRetryCnt,
+		db.ExpBackoff{},
+		func(reader db.Reader, w db.Writer) error {
+			rowsAffected, err := w.Exec(ctx, authorizeConnectionCte, []interface{}{
+				sql.Named("session_id", sessionId),
+				sql.Named("public_id", connectionId),
+				sql.Named("worker_id", workerId),
+			})
+			if err != nil {
+				return errors.Wrap(ctx, err, op, errors.WithMsg(fmt.Sprintf("unable to authorize connection %s", sessionId)))
+			}
+			if rowsAffected == 0 {
+				return errors.Wrap(ctx, status.Errorf(codes.PermissionDenied, "session %s is not authorized (not active, expired or connection limit reached)", sessionId), op, errors.WithCode(errors.InvalidSessionState))
+			}
+			if err := reader.LookupById(ctx, &connection); err != nil {
+				return errors.Wrap(ctx, err, op, errors.WithMsg(fmt.Sprintf("failed for session %s", sessionId)))
+			}
+			connectionStates, err = fetchConnectionStates(ctx, reader, connectionId, db.WithOrder("start_time desc"))
+			if err != nil {
+				return errors.Wrap(ctx, err, op)
+			}
+			return nil
+		},
+	)
+	if err != nil {
+		return nil, nil, errors.Wrap(ctx, err, op)
+	}
+
+	return &connection, connectionStates, nil
+}
+
 // deadWorkerConnCloseMinGrace is the minimum allowable setting for
 // the CloseConnectionsForDeadWorkers method. This is synced with
 // the default server liveness setting.
@@ -20,8 +137,8 @@ var DeadWorkerConnCloseMinGrace = int(servers.DefaultLiveness.Seconds())
 // LookupConnection will look up a connection in the repository and return the connection
 // with its states. If the connection is not found, it will return nil, nil, nil.
 // No options are currently supported.
-func (r *Repository) LookupConnection(ctx context.Context, connectionId string, _ ...Option) (*Connection, []*ConnectionState, error) {
-	const op = "session.(Repository).LookupConnection"
+func (r *ConnectionRepository) LookupConnection(ctx context.Context, connectionId string, _ ...Option) (*Connection, []*ConnectionState, error) {
+	const op = "session.(ConnectionRepository).LookupConnection"
 	if connectionId == "" {
 		return nil, nil, errors.New(ctx, errors.InvalidParameter, op, "missing connectionId id")
 	}
@@ -54,8 +171,8 @@ func (r *Repository) LookupConnection(ctx context.Context, connectionId string,
 
 // ListConnectionsBySessionId will list connections by session ID. Supports the
 // WithLimit and WithOrder options.
-func (r *Repository) ListConnectionsBySessionId(ctx context.Context, sessionId string, opt ...Option) ([]*Connection, error) {
-	const op = "session.(Repository).ListConnectionsBySessionId"
+func (r *ConnectionRepository) ListConnectionsBySessionId(ctx context.Context, sessionId string, opt ...Option) ([]*Connection, error) {
+	const op = "session.(ConnectionRepository).ListConnectionsBySessionId"
 	if sessionId == "" {
 		return nil, errors.New(ctx, errors.InvalidParameter, op, "no session ID supplied")
 	}
@@ -68,8 +185,8 @@ func (r *Repository) ListConnectionsBySessionId(ctx context.Context, sessionId s
 }
 
 // DeleteConnection will delete a connection from the repository.
-func (r *Repository) DeleteConnection(ctx context.Context, publicId string, _ ...Option) (int, error) {
-	const op = "session.(Repository).DeleteConnection"
+func (r *ConnectionRepository) DeleteConnection(ctx context.Context, publicId string, _ ...Option) (int, error) {
+	const op = "session.(ConnectionRepository).DeleteConnection"
 	if publicId == "" {
 		return db.NoRowsAffected, errors.New(ctx, errors.InvalidParameter, op, "missing public id")
 	}
@@ -115,8 +232,8 @@ func (r *Repository) DeleteConnection(ctx context.Context, publicId string, _ ..
 // uses a NOT IN clause to ensure these are excluded. It is not an error for
 // this to be empty as the worker could claim no connections; in that case all
 // connections will immediately transition to closed.
-func (r *Repository) CloseDeadConnectionsForWorker(ctx context.Context, serverId string, foundConns []string) (int, error) {
-	const op = "session.(Repository).CloseDeadConnectionsForWorker"
+func (r *ConnectionRepository) CloseDeadConnectionsForWorker(ctx context.Context, serverId string, foundConns []string) (int, error) {
+	const op = "session.(ConnectionRepository).CloseDeadConnectionsForWorker"
 	if serverId == "" {
 		return db.NoRowsAffected, errors.New(ctx, errors.InvalidParameter, op, "missing server id")
 	}
@@ -166,8 +283,8 @@ type CloseConnectionsForDeadWorkersResult struct {
 // sending status updates to the controller(s).
 //
 // The only input to the method is the grace period, in seconds.
-func (r *Repository) CloseConnectionsForDeadWorkers(ctx context.Context, gracePeriod int) ([]CloseConnectionsForDeadWorkersResult, error) {
-	const op = "session.(Repository).CloseConnectionsForDeadWorkers"
+func (r *ConnectionRepository) CloseConnectionsForDeadWorkers(ctx context.Context, gracePeriod int) ([]CloseConnectionsForDeadWorkersResult, error) {
+	const op = "session.(ConnectionRepository).CloseConnectionsForDeadWorkers"
 	if gracePeriod < DeadWorkerConnCloseMinGrace {
 		return nil, errors.New(ctx,
 			errors.InvalidParameter, op, fmt.Sprintf("gracePeriod must be at least %d seconds", DeadWorkerConnCloseMinGrace))
@@ -214,8 +331,8 @@ func (r *Repository) CloseConnectionsForDeadWorkers(ctx context.Context, gracePe
 // them again would be redundant.
 //
 // The returned map[string][]string is indexed by session ID.
-func (r *Repository) ShouldCloseConnectionsOnWorker(ctx context.Context, foundConns, filterSessions []string) (map[string][]string, error) {
-	const op = "session.(Repository).ShouldCloseConnectionsOnWorker"
+func (r *ConnectionRepository) ShouldCloseConnectionsOnWorker(ctx context.Context, foundConns, filterSessions []string) (map[string][]string, error) {
+	const op = "session.(ConnectionRepository).ShouldCloseConnectionsOnWorker"
 	if len(foundConns) < 1 {
 		return nil, nil // nothing to do
 	}

internal/session/repository_connection_test.go

@@ -19,6 +19,7 @@ import (
 
 func TestRepository_ListConnection(t *testing.T) {
 	t.Parallel()
+	ctx := context.Background()
 	conn, _ := db.TestSetup(t, "postgres")
 	const testLimit = 10
 	wrapper := db.TestWrapper(t)
@@ -26,6 +27,7 @@ func TestRepository_ListConnection(t *testing.T) {
 	rw := db.New(conn)
 	kms := kms.TestKms(t, conn, wrapper)
 	repo, err := NewRepository(rw, rw, kms, WithLimit(testLimit))
+	connRepo, err := NewConnectionRepository(ctx, rw, rw, kms, WithLimit(testLimit))
 	require.NoError(t, err)
 	session := TestDefaultSession(t, conn, wrapper, iamRepo)
 
@@ -96,7 +98,7 @@ func TestRepository_ListConnection(t *testing.T) {
 				testConnections = append(testConnections, c)
 			}
 			assert.Equal(tt.createCnt, len(testConnections))
-			got, err := repo.ListConnectionsBySessionId(context.Background(), tt.args.searchForSessionId, tt.args.opt...)
+			got, err := connRepo.ListConnectionsBySessionId(context.Background(), tt.args.searchForSessionId, tt.args.opt...)
 			if tt.wantErr {
 				require.Error(err)
 				return
@@ -119,7 +121,7 @@ func TestRepository_ListConnection(t *testing.T) {
 				"127.0.0.1",
 			)
 		}
-		got, err := repo.ListConnectionsBySessionId(context.Background(), session.PublicId, WithOrderByCreateTime(db.AscendingOrderBy))
+		got, err := connRepo.ListConnectionsBySessionId(context.Background(), session.PublicId, WithOrderByCreateTime(db.AscendingOrderBy))
 		require.NoError(err)
 		assert.Equal(wantCnt, len(got))
 
@@ -133,12 +135,13 @@ func TestRepository_ListConnection(t *testing.T) {
 
 func TestRepository_DeleteConnection(t *testing.T) {
 	t.Parallel()
+	ctx := context.Background()
 	conn, _ := db.TestSetup(t, "postgres")
 	rw := db.New(conn)
 	wrapper := db.TestWrapper(t)
 	iamRepo := iam.TestRepo(t, conn, wrapper)
 	kms := kms.TestKms(t, conn, wrapper)
-	repo, err := NewRepository(rw, rw, kms)
+	connRepo, err := NewConnectionRepository(ctx, rw, rw, kms)
 	require.NoError(t, err)
 	session := TestDefaultSession(t, conn, wrapper, iamRepo)
 
@@ -171,7 +174,7 @@ func TestRepository_DeleteConnection(t *testing.T) {
 			},
 			wantRowsDeleted: 0,
 			wantErr:         true,
-			wantErrMsg:      "session.(Repository).DeleteConnection: missing public id: parameter violation: error #100",
+			wantErrMsg:      "session.(ConnectionRepository).DeleteConnection: missing public id: parameter violation: error #100",
 		},
 		{
 			name: "not-found",
@@ -192,7 +195,7 @@ func TestRepository_DeleteConnection(t *testing.T) {
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 			assert := assert.New(t)
-			deletedRows, err := repo.DeleteConnection(context.Background(), tt.args.connection.PublicId, tt.args.opt...)
+			deletedRows, err := connRepo.DeleteConnection(context.Background(), tt.args.connection.PublicId, tt.args.opt...)
 			if tt.wantErr {
 				assert.Error(err)
 				assert.Equal(0, deletedRows)
@@ -204,7 +207,7 @@ func TestRepository_DeleteConnection(t *testing.T) {
 			}
 			assert.NoError(err)
 			assert.Equal(tt.wantRowsDeleted, deletedRows)
-			found, _, err := repo.LookupConnection(context.Background(), tt.args.connection.PublicId)
+			found, _, err := connRepo.LookupConnection(context.Background(), tt.args.connection.PublicId)
 			assert.NoError(err)
 			assert.Nil(found)
 
@@ -216,6 +219,7 @@ func TestRepository_DeleteConnection(t *testing.T) {
 
 func TestRepository_CloseDeadConnectionsOnWorker(t *testing.T) {
 	t.Parallel()
+	ctx := context.Background()
 	require, assert := require.New(t), assert.New(t)
 	conn, _ := db.TestSetup(t, "postgres")
 	rw := db.New(conn)
@@ -223,8 +227,8 @@ func TestRepository_CloseDeadConnectionsOnWorker(t *testing.T) {
 	iamRepo := iam.TestRepo(t, conn, wrapper)
 	kms := kms.TestKms(t, conn, wrapper)
 	repo, err := NewRepository(rw, rw, kms)
+	connRepo, err := NewConnectionRepository(ctx, rw, rw, kms)
 	require.NoError(err)
-	ctx := context.Background()
 	numConns := 12
 
 	// Create two "workers". One will remain untouched while the other "goes
@@ -243,7 +247,7 @@ func TestRepository_CloseDeadConnectionsOnWorker(t *testing.T) {
 		sess := TestDefaultSession(t, conn, wrapper, iamRepo, WithServerId(serverId), WithDbOpts(db.WithSkipVetForWrite(true)))
 		sess, _, err = repo.ActivateSession(ctx, sess.GetPublicId(), sess.Version, serverId, "worker", []byte("foo"))
 		require.NoError(err)
-		c, cs, _, err := repo.AuthorizeConnection(ctx, sess.GetPublicId(), serverId)
+		c, cs, err := connRepo.AuthorizeConnection(ctx, sess.GetPublicId(), serverId)
 		require.NoError(err)
 		require.Len(cs, 1)
 		require.Equal(StatusAuthorized, cs[0].Status)
@@ -289,7 +293,7 @@ func TestRepository_CloseDeadConnectionsOnWorker(t *testing.T) {
 	// all connection IDs for worker 1 should be showing as non-closed, and
 	// the ones for worker 2 not advertised should be closed.
 	shouldStayOpen := worker2ConnIds[0:2]
-	count, err := repo.CloseDeadConnectionsForWorker(ctx, worker2.GetPrivateId(), shouldStayOpen)
+	count, err := connRepo.CloseDeadConnectionsForWorker(ctx, worker2.GetPrivateId(), shouldStayOpen)
 	require.NoError(err)
 	assert.Equal(4, count)
 
@@ -300,7 +304,7 @@ func TestRepository_CloseDeadConnectionsOnWorker(t *testing.T) {
 	var conns []*Connection
 	require.NoError(repo.list(ctx, &conns, "", nil))
 	for _, conn := range conns {
-		_, states, err := repo.LookupConnection(ctx, conn.PublicId)
+		_, states, err := connRepo.LookupConnection(ctx, conn.PublicId)
 		require.NoError(err)
 		var foundClosed bool
 		for _, state := range states {
@@ -315,7 +319,7 @@ func TestRepository_CloseDeadConnectionsOnWorker(t *testing.T) {
 	// Now, advertise none of the connection IDs for worker 2. This is mainly to
 	// test that handling the case where we do not include IDs works properly as
 	// it changes the where clause.
-	count, err = repo.CloseDeadConnectionsForWorker(ctx, worker1.GetPrivateId(), nil)
+	count, err = connRepo.CloseDeadConnectionsForWorker(ctx, worker1.GetPrivateId(), nil)
 	require.NoError(err)
 	assert.Equal(6, count)
 
@@ -323,7 +327,7 @@ func TestRepository_CloseDeadConnectionsOnWorker(t *testing.T) {
 	conns = nil
 	require.NoError(repo.list(ctx, &conns, "", nil))
 	for _, conn := range conns {
-		_, states, err := repo.LookupConnection(ctx, conn.PublicId)
+		_, states, err := connRepo.LookupConnection(ctx, conn.PublicId)
 		require.NoError(err)
 		var foundClosed bool
 		for _, state := range states {
@@ -338,6 +342,7 @@ func TestRepository_CloseDeadConnectionsOnWorker(t *testing.T) {
 
 func TestRepository_CloseConnectionsForDeadWorkers(t *testing.T) {
 	t.Parallel()
+	ctx := context.Background()
 	require := require.New(t)
 	conn, _ := db.TestSetup(t, "postgres")
 	rw := db.New(conn)
@@ -345,10 +350,10 @@ func TestRepository_CloseConnectionsForDeadWorkers(t *testing.T) {
 	iamRepo := iam.TestRepo(t, conn, wrapper)
 	kms := kms.TestKms(t, conn, wrapper)
 	repo, err := NewRepository(rw, rw, kms)
+	connRepo, err := NewConnectionRepository(ctx, rw, rw, kms)
 	require.NoError(err)
 	serversRepo, err := servers.NewRepository(rw, rw, kms)
 	require.NoError(err)
-	ctx := context.Background()
 
 	// connection count = 6 * states(authorized, connected, closed = 3) * servers_with_open_connections(3)
 	numConns := 54
@@ -377,7 +382,7 @@ func TestRepository_CloseConnectionsForDeadWorkers(t *testing.T) {
 		sess := TestDefaultSession(t, conn, wrapper, iamRepo, WithServerId(serverId), WithDbOpts(db.WithSkipVetForWrite(true)))
 		sess, _, err = repo.ActivateSession(ctx, sess.GetPublicId(), sess.Version, serverId, "worker", []byte("foo"))
 		require.NoError(err)
-		c, cs, _, err := repo.AuthorizeConnection(ctx, sess.GetPublicId(), serverId)
+		c, cs, err := connRepo.AuthorizeConnection(ctx, sess.GetPublicId(), serverId)
 		require.NoError(err)
 		require.Len(cs, 1)
 		require.Equal(StatusAuthorized, cs[0].Status)
@@ -496,7 +501,7 @@ func TestRepository_CloseConnectionsForDeadWorkers(t *testing.T) {
 				expected = StatusAuthorized
 			}
 
-			_, states, err := repo.LookupConnection(ctx, connId)
+			_, states, err := connRepo.LookupConnection(ctx, connId)
 			require.NoError(err)
 			require.Equal(expected, states[0].Status, "expected latest status for %q (index %d) to be %v", connId, i, expected)
 		}
@@ -515,10 +520,10 @@ func TestRepository_CloseConnectionsForDeadWorkers(t *testing.T) {
 	// Now try some scenarios.
 	{
 		// First, test the error/validation case.
-		result, err := repo.CloseConnectionsForDeadWorkers(ctx, 0)
+		result, err := connRepo.CloseConnectionsForDeadWorkers(ctx, 0)
 		require.Equal(err, errors.E(ctx,
 			errors.WithCode(errors.InvalidParameter),
-			errors.WithOp("session.(Repository).CloseConnectionsForDeadWorkers"),
+			errors.WithOp("session.(ConnectionRepository).CloseConnectionsForDeadWorkers"),
 			errors.WithMsg(fmt.Sprintf("gracePeriod must be at least %d seconds", DeadWorkerConnCloseMinGrace)),
 		))
 		require.Nil(result)
@@ -531,7 +536,7 @@ func TestRepository_CloseConnectionsForDeadWorkers(t *testing.T) {
 		worker3 = updateServer(t, worker3)
 		updateServer(t, worker4) // no re-assignment here because we never reference the server again
 
-		result, err := repo.CloseConnectionsForDeadWorkers(ctx, DeadWorkerConnCloseMinGrace)
+		result, err := connRepo.CloseConnectionsForDeadWorkers(ctx, DeadWorkerConnCloseMinGrace)
 		require.NoError(err)
 		require.Empty(result)
 		// Expect appropriate split connection state on worker1
@@ -551,7 +556,7 @@ func TestRepository_CloseConnectionsForDeadWorkers(t *testing.T) {
 		worker2 = updateServer(t, worker2)
 		worker3 = updateServer(t, worker3)
 
-		result, err := repo.CloseConnectionsForDeadWorkers(ctx, DeadWorkerConnCloseMinGrace)
+		result, err := connRepo.CloseConnectionsForDeadWorkers(ctx, DeadWorkerConnCloseMinGrace)
 		require.NoError(err)
 		require.Empty(result)
 		// Expect appropriate split connection state on worker1
@@ -569,7 +574,7 @@ func TestRepository_CloseConnectionsForDeadWorkers(t *testing.T) {
 		worker2 = updateServer(t, worker2)
 		worker3 = updateServer(t, worker3)
 
-		result, err := repo.CloseConnectionsForDeadWorkers(ctx, DeadWorkerConnCloseMinGrace)
+		result, err := connRepo.CloseConnectionsForDeadWorkers(ctx, DeadWorkerConnCloseMinGrace)
 		require.NoError(err)
 		// Assert that we have one result with the appropriate ID and
 		// number of connections closed. Due to how things are
@@ -594,7 +599,7 @@ func TestRepository_CloseConnectionsForDeadWorkers(t *testing.T) {
 		// appropriate message from the next two servers acted on.
 		time.Sleep(time.Second * time.Duration(DeadWorkerConnCloseMinGrace))
 
-		result, err := repo.CloseConnectionsForDeadWorkers(ctx, DeadWorkerConnCloseMinGrace)
+		result, err := connRepo.CloseConnectionsForDeadWorkers(ctx, DeadWorkerConnCloseMinGrace)
 		require.NoError(err)
 		// Assert that we have one result with the appropriate ID and number of connections closed.
 		require.Equal([]CloseConnectionsForDeadWorkersResult{
@@ -620,6 +625,7 @@ func TestRepository_CloseConnectionsForDeadWorkers(t *testing.T) {
 
 func TestRepository_ShouldCloseConnectionsOnWorker(t *testing.T) {
 	t.Parallel()
+	ctx := context.Background()
 	require := require.New(t)
 	conn, _ := db.TestSetup(t, "postgres")
 	rw := db.New(conn)
@@ -627,8 +633,8 @@ func TestRepository_ShouldCloseConnectionsOnWorker(t *testing.T) {
 	iamRepo := iam.TestRepo(t, conn, wrapper)
 	kms := kms.TestKms(t, conn, wrapper)
 	repo, err := NewRepository(rw, rw, kms)
+	connRepo, err := NewConnectionRepository(ctx, rw, rw, kms)
 	require.NoError(err)
-	ctx := context.Background()
 	numConns := 12
 
 	// Create a worker, we only need one here as our query is dependent
@@ -644,7 +650,7 @@ func TestRepository_ShouldCloseConnectionsOnWorker(t *testing.T) {
 		sessionId := sess.GetPublicId()
 		sess, _, err = repo.ActivateSession(ctx, sessionId, sess.Version, serverId, "worker", []byte("foo"))
 		require.NoError(err)
-		c, cs, _, err := repo.AuthorizeConnection(ctx, sess.GetPublicId(), serverId)
+		c, cs, err := connRepo.AuthorizeConnection(ctx, sess.GetPublicId(), serverId)
 		require.NoError(err)
 		require.Len(cs, 1)
 		require.Equal(StatusAuthorized, cs[0].Status)
@@ -709,7 +715,7 @@ func TestRepository_ShouldCloseConnectionsOnWorker(t *testing.T) {
 	// Now we try some scenarios.
 	{
 		// First test an empty set.
-		result, err := repo.ShouldCloseConnectionsOnWorker(ctx, nil, nil)
+		result, err := connRepo.ShouldCloseConnectionsOnWorker(ctx, nil, nil)
 		require.NoError(err)
 		require.Zero(result, "should be empty when no connections are supplied")
 	}
@@ -724,14 +730,14 @@ func TestRepository_ShouldCloseConnectionsOnWorker(t *testing.T) {
 		expectedSessionConnIds := make(map[string][]string)
 		for sessionId, connIds := range sessionConnIds {
 			for _, connId := range connIds {
-				if testIsConnectionClosed(ctx, t, repo, connId) {
+				if testIsConnectionClosed(ctx, t, connRepo, connId) {
 					expectedSessionConnIds[sessionId] = append(expectedSessionConnIds[sessionId], connId)
 				}
 			}
 		}
 
 		// Send query, use all connections w/o a filter on sessions.
-		actualSessionConnIds, err := repo.ShouldCloseConnectionsOnWorker(ctx, connIds, nil)
+		actualSessionConnIds, err := connRepo.ShouldCloseConnectionsOnWorker(ctx, connIds, nil)
 		require.NoError(err)
 		require.Equal(expectedSessionConnIds, actualSessionConnIds)
 	}
@@ -744,7 +750,7 @@ func TestRepository_ShouldCloseConnectionsOnWorker(t *testing.T) {
 		var filterSession bool
 		for sessionId, connIds := range sessionConnIds {
 			for _, connId := range connIds {
-				if testIsConnectionClosed(ctx, t, repo, connId) {
+				if testIsConnectionClosed(ctx, t, connRepo, connId) {
 					if !filterSession {
 						expectedSessionConnIds[sessionId] = append(expectedSessionConnIds[sessionId], connId)
 					} else {
@@ -760,13 +766,13 @@ func TestRepository_ShouldCloseConnectionsOnWorker(t *testing.T) {
 		}
 
 		// Send query with the session filter.
-		actualSessionConnIds, err := repo.ShouldCloseConnectionsOnWorker(ctx, connIds, filterSessionIds)
+		actualSessionConnIds, err := connRepo.ShouldCloseConnectionsOnWorker(ctx, connIds, filterSessionIds)
 		require.NoError(err)
 		require.Equal(expectedSessionConnIds, actualSessionConnIds)
 	}
 }
 
-func testIsConnectionClosed(ctx context.Context, t *testing.T, repo *Repository, connId string) bool {
+func testIsConnectionClosed(ctx context.Context, t *testing.T, repo *ConnectionRepository, connId string) bool {
 	require := require.New(t)
 	_, states, err := repo.LookupConnection(ctx, connId)
 	require.NoError(err)

internal/session/repository_session.go

@@ -13,8 +13,6 @@ import (
 	"github.com/hashicorp/boundary/internal/errors"
 	"github.com/hashicorp/boundary/internal/kms"
 	wrapping "github.com/hashicorp/go-kms-wrapping"
-	"google.golang.org/grpc/codes"
-	"google.golang.org/grpc/status"
 )
 
 // CreateSession inserts into the repository and returns the new Session with
@@ -138,7 +136,7 @@ func (r *Repository) CreateSession(ctx context.Context, sessionWrapper wrapping.
 // with its states.  Returned States are ordered by start time descending.  If the
 // session is not found, it will return nil, nil, nil. No options are currently
 // supported.
-func (r *Repository) LookupSession(ctx context.Context, sessionId string, _ ...Option) (*Session, *ConnectionAuthzSummary, error) {
+func (r *Repository) LookupSession(ctx context.Context, sessionId string, _ ...Option) (*Session, *AuthzSummary, error) {
 	const op = "session.(Repository).LookupSession"
 	if sessionId == "" {
 		return nil, nil, errors.New(ctx, errors.InvalidParameter, op, "missing session id")
@@ -377,69 +375,13 @@ func (r *Repository) TerminateCompletedSessions(ctx context.Context) (int, error
 	return rowsAffected, nil
 }
 
-// AuthorizeConnection will check to see if a connection is allowed.  Currently,
-// that authorization checks:
-// * the hasn't expired based on the session.Expiration
-// * number of connections already created is less than session.ConnectionLimit
-// If authorization is success, it creates/stores a new connection in the repo
-// and returns it, along with its states.  If the authorization fails, it
-// an error with Code InvalidSessionState.
-func (r *Repository) AuthorizeConnection(ctx context.Context, sessionId, workerId string) (*Connection, []*ConnectionState, *ConnectionAuthzSummary, error) {
-	const op = "session.(Repository).AuthorizeConnection"
-	if sessionId == "" {
-		return nil, nil, nil, errors.Wrap(ctx, status.Error(codes.FailedPrecondition, "missing session id"), op, errors.WithCode(errors.InvalidParameter))
-	}
-	connectionId, err := newConnectionId()
-	if err != nil {
-		return nil, nil, nil, errors.Wrap(ctx, err, op)
-	}
-
-	connection := AllocConnection()
-	connection.PublicId = connectionId
-	var connectionStates []*ConnectionState
-	_, err = r.writer.DoTx(
-		ctx,
-		db.StdRetryCnt,
-		db.ExpBackoff{},
-		func(reader db.Reader, w db.Writer) error {
-			rowsAffected, err := w.Exec(ctx, authorizeConnectionCte, []interface{}{
-				sql.Named("session_id", sessionId),
-				sql.Named("public_id", connectionId),
-				sql.Named("worker_id", workerId),
-			})
-			if err != nil {
-				return errors.Wrap(ctx, err, op, errors.WithMsg(fmt.Sprintf("unable to authorize connection %s", sessionId)))
-			}
-			if rowsAffected == 0 {
-				return errors.Wrap(ctx, status.Errorf(codes.PermissionDenied, "session %s is not authorized (not active, expired or connection limit reached)", sessionId), op, errors.WithCode(errors.InvalidSessionState))
-			}
-			if err := reader.LookupById(ctx, &connection); err != nil {
-				return errors.Wrap(ctx, err, op, errors.WithMsg(fmt.Sprintf("failed for session %s", sessionId)))
-			}
-			connectionStates, err = fetchConnectionStates(ctx, reader, connectionId, db.WithOrder("start_time desc"))
-			if err != nil {
-				return errors.Wrap(ctx, err, op)
-			}
-			return nil
-		},
-	)
-	if err != nil {
-		return nil, nil, nil, errors.Wrap(ctx, err, op)
-	}
-	authzSummary, err := r.sessionAuthzSummary(ctx, connection.SessionId)
-	if err != nil {
-		return nil, nil, nil, errors.Wrap(ctx, err, op)
-	}
-	return &connection, connectionStates, authzSummary, nil
-}
-
-type ConnectionAuthzSummary struct {
+type AuthzSummary struct {
 	ExpirationTime         *timestamp.Timestamp
 	ConnectionLimit        int32
 	CurrentConnectionCount uint32
 }
 
-func (r *Repository) sessionAuthzSummary(ctx context.Context, sessionId string) (*ConnectionAuthzSummary, error) {
+func (r *Repository) sessionAuthzSummary(ctx context.Context, sessionId string) (*AuthzSummary, error) {
 	const op = "session.(Repository).sessionAuthzSummary"
 	rows, err := r.reader.Query(ctx, remainingConnectionsCte, []interface{}{sql.Named("session_id", sessionId)})
 	if err != nil {
@@ -447,12 +389,12 @@ func (r *Repository) sessionAuthzSummary(ctx context.Context, sessionId string)
 	}
 	defer rows.Close()
 
-	var info *ConnectionAuthzSummary
+	var info *AuthzSummary
 	for rows.Next() {
 		if info != nil {
 			return nil, errors.New(ctx, errors.MultipleRecords, op, "query returned more than one row")
 		}
-		info = &ConnectionAuthzSummary{}
+		info = &AuthzSummary{}
 		if err := r.reader.ScanRows(ctx, rows, info); err != nil {
 			return nil, errors.Wrap(ctx, err, op, errors.WithMsg("scan row failed"))
 		}
@@ -462,7 +404,7 @@ func (r *Repository) sessionAuthzSummary(ctx context.Context, sessionId string)
 
 // ConnectConnection updates a connection in the repo with a state of "connected".
 func (r *Repository) ConnectConnection(ctx context.Context, c ConnectWith) (*Connection, []*ConnectionState, error) {
-	const op = "session.(Repository).ConnectConnection"
+	const op = "session.(ConnectionRepository).ConnectConnection"
 	// ConnectWith.validate will check all the fields...
 	if err := c.validate(); err != nil {
 		return nil, nil, errors.Wrap(ctx, err, op)
@@ -527,7 +469,7 @@ type CloseConnectionResp struct {
 // called by a worker after it's closed a connection between the client and the
 // endpoint
 func (r *Repository) CloseConnections(ctx context.Context, closeWith []CloseWith, _ ...Option) ([]CloseConnectionResp, error) {
-	const op = "session.(Repository).CloseConnections"
+	const op = "session.(ConnectionRepository).CloseConnections"
 	if len(closeWith) == 0 {
 		return nil, errors.New(ctx, errors.InvalidParameter, op, "missing connections")
 	}

internal/session/repository_session_test.go

@@ -541,108 +541,16 @@ func TestRepository_updateState(t *testing.T) {
 	}
 }
 
-func TestRepository_AuthorizeConnect(t *testing.T) {
-	t.Parallel()
-	conn, _ := db.TestSetup(t, "postgres")
-	rw := db.New(conn)
-	wrapper := db.TestWrapper(t)
-	iamRepo := iam.TestRepo(t, conn, wrapper)
-	kms := kms.TestKms(t, conn, wrapper)
-	repo, err := NewRepository(rw, rw, kms)
-	require.NoError(t, err)
-
-	var testServer string
-	setupFn := func(exp *timestamp.Timestamp) *Session {
-		composedOf := TestSessionParams(t, conn, wrapper, iamRepo)
-		if exp != nil {
-			composedOf.ExpirationTime = exp
-		}
-		s := TestSession(t, conn, wrapper, composedOf)
-		srv := TestWorker(t, conn, wrapper)
-		testServer = srv.PrivateId
-		tofu := TestTofu(t)
-		_, _, err := repo.ActivateSession(context.Background(), s.PublicId, s.Version, srv.PrivateId, srv.Type, tofu)
-		require.NoError(t, err)
-		return s
-	}
-	testSession := setupFn(nil)
-
-	tests := []struct {
-		name          string
-		session       *Session
-		wantErr       bool
-		wantIsError   error
-		wantAuthzInfo ConnectionAuthzSummary
-	}{
-		{
-			name:    "valid",
-			session: testSession,
-			wantAuthzInfo: ConnectionAuthzSummary{
-				ConnectionLimit:        1,
-				CurrentConnectionCount: 1,
-				ExpirationTime:         testSession.ExpirationTime,
-			},
-		},
-		{
-			name: "empty-sessionId",
-			session: func() *Session {
-				s := AllocSession()
-				return &s
-			}(),
-			wantErr: true,
-		},
-		{
-			name: "exceeded-connection-limit",
-			session: func() *Session {
-				session := setupFn(nil)
-				_ = TestConnection(t, conn, session.PublicId, "127.0.0.1", 22, "127.0.0.1", 2222, "127.0.0.1")
-				return session
-			}(),
-			wantErr: true,
-		},
-		{
-			name:    "expired-session",
-			session: setupFn(&timestamp.Timestamp{Timestamp: timestamppb.Now()}),
-			wantErr: true,
-		},
-	}
-	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
-			assert, require := assert.New(t), require.New(t)
-
-			c, cs, authzInfo, err := repo.AuthorizeConnection(context.Background(), tt.session.PublicId, testServer)
-			if tt.wantErr {
-				require.Error(err)
-				// TODO (jimlambrt 9/2020): add in tests for errorsIs once we
-				// remove the grpc errors from the repo.
-				// if tt.wantIsError != nil {
-				// 	assert.Truef(errors.Is(err, tt.wantIsError), "unexpected error %s", err.Error())
-				// }
-				return
-			}
-			require.NoError(err)
-			require.NotNil(c)
-			require.NotNil(cs)
-			assert.Equal(StatusAuthorized, cs[0].Status)
-
-			assert.True(authzInfo.ExpirationTime.GetTimestamp().AsTime().Sub(tt.wantAuthzInfo.ExpirationTime.GetTimestamp().AsTime()) < 10*time.Millisecond)
-			tt.wantAuthzInfo.ExpirationTime = authzInfo.ExpirationTime
-
-			assert.Equal(tt.wantAuthzInfo.ExpirationTime, authzInfo.ExpirationTime)
-			assert.Equal(tt.wantAuthzInfo.ConnectionLimit, authzInfo.ConnectionLimit)
-			assert.Equal(tt.wantAuthzInfo.CurrentConnectionCount, authzInfo.CurrentConnectionCount)
-		})
-	}
-}
-
 func TestRepository_ConnectConnection(t *testing.T) {
 	t.Parallel()
+	ctx := context.Background()
 	conn, _ := db.TestSetup(t, "postgres")
 	rw := db.New(conn)
 	wrapper := db.TestWrapper(t)
 	iamRepo := iam.TestRepo(t, conn, wrapper)
 	kms := kms.TestKms(t, conn, wrapper)
 	repo, err := NewRepository(rw, rw, kms)
+	connRepo, err := NewConnectionRepository(ctx, rw, rw, kms)
 	require.NoError(t, err)
 
 	setupFn := func() ConnectWith {
@@ -746,7 +654,7 @@ func TestRepository_ConnectConnection(t *testing.T) {
 			require.NotNil(c)
 			require.NotNil(cs)
 			assert.Equal(StatusConnected, cs[0].Status)
-			gotConn, _, err := repo.LookupConnection(context.Background(), c.PublicId)
+			gotConn, _, err := connRepo.LookupConnection(context.Background(), c.PublicId)
 			require.NoError(err)
 			assert.Equal(tt.connectWith.ClientTcpAddress, gotConn.ClientTcpAddress)
 			assert.Equal(tt.connectWith.ClientTcpPort, gotConn.ClientTcpPort)
@@ -760,12 +668,14 @@ func TestRepository_ConnectConnection(t *testing.T) {
 
 func TestRepository_TerminateCompletedSessions(t *testing.T) {
 	t.Parallel()
+	ctx := context.Background()
 	conn, _ := db.TestSetup(t, "postgres")
 	rw := db.New(conn)
 	wrapper := db.TestWrapper(t)
 	iamRepo := iam.TestRepo(t, conn, wrapper)
 	kms := kms.TestKms(t, conn, wrapper)
 	repo, err := NewRepository(rw, rw, kms)
+	connRepo, err := NewConnectionRepository(ctx, rw, rw, kms)
 	require.NoError(t, err)
 
 	setupFn := func(limit int32, expireIn time.Duration, leaveOpen bool) *Session {
@@ -967,10 +877,10 @@ func TestRepository_TerminateCompletedSessions(t *testing.T) {
 					}
 					assert.Equal(args.wantTermed[found.PublicId].String(), found.TerminationReason)
 					t.Logf("terminated %s has a connection limit of %d", found.PublicId, found.ConnectionLimit)
-					conn, err := repo.ListConnectionsBySessionId(context.Background(), found.PublicId)
+					conn, err := connRepo.ListConnectionsBySessionId(context.Background(), found.PublicId)
 					require.NoError(err)
 					for _, sc := range conn {
-						c, cs, err := repo.LookupConnection(context.Background(), sc.PublicId)
+						c, cs, err := connRepo.LookupConnection(context.Background(), sc.PublicId)
 						require.NoError(err)
 						assert.NotEmpty(c.ClosedReason)
 						for _, s := range cs {
@@ -980,7 +890,7 @@ func TestRepository_TerminateCompletedSessions(t *testing.T) {
 				} else {
 					t.Logf("not terminated %s has a connection limit of %d", found.PublicId, found.ConnectionLimit)
 					assert.Equal("", found.TerminationReason)
-					conn, err := repo.ListConnectionsBySessionId(context.Background(), found.PublicId)
+					conn, err := connRepo.ListConnectionsBySessionId(context.Background(), found.PublicId)
 					require.NoError(err)
 					for _, sc := range conn {
 						cs, err := fetchConnectionStates(context.Background(), rw, sc.PublicId)

internal/session/service_authorize_connection.go

@@ -0,0 +1,29 @@
+package session
+
+import (
+	"context"
+
+	"github.com/hashicorp/boundary/internal/errors"
+)
+
+// AuthorizeConnection is a domain service function that will create a Connection
+// for a session if the following criteria are met:
+// * The session is active.
+// * The session is not expired.
+// * The session has not reached its connection limit or has a connection limit of -1.
+// If any of these criteria is not met, it returns an error with Code InvalidSessionState.
+func AuthorizeConnection(ctx context.Context, sessionRepoFn *Repository, connectionRepoFn *ConnectionRepository,
+	sessionId, workerId string, opt ...Option) (*Connection, []*ConnectionState, *AuthzSummary, error) {
+	const op = "session.AuthorizeConnection"
+
+	connection, connectionStates, err := connectionRepoFn.AuthorizeConnection(ctx, sessionId, workerId)
+	if err != nil {
+		return nil, nil, nil, errors.Wrap(ctx, err, op)
+	}
+
+	authzSummary, err := sessionRepoFn.sessionAuthzSummary(ctx, sessionId)
+	if err != nil {
+		return nil, nil, nil, errors.Wrap(ctx, err, op)
+	}
+	return connection, connectionStates, authzSummary, nil
+}

internal/session/service_authorize_connection_test.go

@@ -0,0 +1,112 @@
+package session
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"google.golang.org/protobuf/types/known/timestamppb"
+
+	"github.com/hashicorp/boundary/internal/db"
+	"github.com/hashicorp/boundary/internal/db/timestamp"
+	"github.com/hashicorp/boundary/internal/iam"
+	"github.com/hashicorp/boundary/internal/kms"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestService_AuthorizeConnection(t *testing.T) {
+	t.Parallel()
+	ctx := context.Background()
+	conn, _ := db.TestSetup(t, "postgres")
+	rw := db.New(conn)
+	wrapper := db.TestWrapper(t)
+	iamRepo := iam.TestRepo(t, conn, wrapper)
+	kms := kms.TestKms(t, conn, wrapper)
+	repo, err := NewRepository(rw, rw, kms)
+	connRepo, err := NewConnectionRepository(ctx, rw, rw, kms)
+	require.NoError(t, err)
+
+	var testServer string
+	setupFn := func(exp *timestamp.Timestamp) *Session {
+		composedOf := TestSessionParams(t, conn, wrapper, iamRepo)
+		if exp != nil {
+			composedOf.ExpirationTime = exp
+		}
+		s := TestSession(t, conn, wrapper, composedOf)
+		srv := TestWorker(t, conn, wrapper)
+		testServer = srv.PrivateId
+		tofu := TestTofu(t)
+		_, _, err := repo.ActivateSession(context.Background(), s.PublicId, s.Version, srv.PrivateId, srv.Type, tofu)
+		require.NoError(t, err)
+		return s
+	}
+	testSession := setupFn(nil)
+
+	tests := []struct {
+		name          string
+		session       *Session
+		wantErr       bool
+		wantIsError   error
+		wantAuthzInfo AuthzSummary
+	}{
+		{
+			name:    "valid",
+			session: testSession,
+			wantAuthzInfo: AuthzSummary{
+				ConnectionLimit:        1,
+				CurrentConnectionCount: 1,
+				ExpirationTime:         testSession.ExpirationTime,
+			},
+		},
+		{
+			name: "empty-sessionId",
+			session: func() *Session {
+				s := AllocSession()
+				return &s
+			}(),
+			wantErr: true,
+		},
+		{
+			name: "exceeded-connection-limit",
+			session: func() *Session {
+				session := setupFn(nil)
+				_ = TestConnection(t, conn, session.PublicId, "127.0.0.1", 22, "127.0.0.1", 2222, "127.0.0.1")
+				return session
+			}(),
+			wantErr: true,
+		},
+		{
+			name:    "expired-session",
+			session: setupFn(&timestamp.Timestamp{Timestamp: timestamppb.Now()}),
+			wantErr: true,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			assert, require := assert.New(t), require.New(t)
+
+			c, cs, authzInfo, err := AuthorizeConnection(context.Background(), repo, connRepo, tt.session.PublicId, testServer)
+			if tt.wantErr {
+				require.Error(err)
+				// TODO (jimlambrt 9/2020): add in tests for errorsIs once we
+				// remove the grpc errors from the repo.
+				// if tt.wantIsError != nil {
+				// 	assert.Truef(errors.Is(err, tt.wantIsError), "unexpected error %s", err.Error())
+				// }
+				return
+			}
+			require.NoError(err)
+			require.NotNil(c)
+			require.NotNil(cs)
+			assert.Equal(StatusAuthorized, cs[0].Status)
+
+			assert.True(authzInfo.ExpirationTime.GetTimestamp().AsTime().Sub(tt.wantAuthzInfo.ExpirationTime.GetTimestamp().AsTime()) < 10*time.Millisecond)
+			tt.wantAuthzInfo.ExpirationTime = authzInfo.ExpirationTime
+
+			assert.Equal(tt.wantAuthzInfo.ExpirationTime, authzInfo.ExpirationTime)
+			assert.Equal(tt.wantAuthzInfo.ConnectionLimit, authzInfo.ConnectionLimit)
+			assert.Equal(tt.wantAuthzInfo.CurrentConnectionCount, authzInfo.CurrentConnectionCount)
+		})
+	}
+}

internal/tests/cluster/session_cleanup_test.go

@@ -170,18 +170,18 @@ func testWorkerSessionCleanupSingle(burdenCase timeoutBurdenType) func(t *testin
 		case timeoutBurdenTypeWorker:
 			// Wait on worker, then check controller
 			sess.ExpectConnectionStateOnWorker(ctx, t, w1, session.StatusClosed)
-			sess.ExpectConnectionStateOnController(ctx, t, c1.Controller().SessionRepoFn, session.StatusConnected)
+			sess.ExpectConnectionStateOnController(ctx, t, c1.Controller().ConnectionRepoFn, session.StatusConnected)
 
 		case timeoutBurdenTypeController:
 			// Wait on controller, then check worker
-			sess.ExpectConnectionStateOnController(ctx, t, c1.Controller().SessionRepoFn, session.StatusClosed)
+			sess.ExpectConnectionStateOnController(ctx, t, c1.Controller().ConnectionRepoFn, session.StatusClosed)
 			sess.ExpectConnectionStateOnWorker(ctx, t, w1, session.StatusConnected)
 
 		default:
 			// Should be closed on both worker and controller. Wait on
 			// worker then check controller.
 			sess.ExpectConnectionStateOnWorker(ctx, t, w1, session.StatusClosed)
-			sess.ExpectConnectionStateOnController(ctx, t, c1.Controller().SessionRepoFn, session.StatusClosed)
+			sess.ExpectConnectionStateOnController(ctx, t, c1.Controller().ConnectionRepoFn, session.StatusClosed)
 		}
 
 		// Run send/receive test again to check expected connection-level
@@ -211,7 +211,7 @@ func testWorkerSessionCleanupSingle(burdenCase timeoutBurdenType) func(t *testin
 			// a connection status, ensure that our old session's
 			// connections are actually closed now that the worker is
 			// properly reporting in again.
-			sess.ExpectConnectionStateOnController(ctx, t, c1.Controller().SessionRepoFn, session.StatusClosed)
+			sess.ExpectConnectionStateOnController(ctx, t, c1.Controller().ConnectionRepoFn, session.StatusClosed)
 
 		case timeoutBurdenTypeController:
 			// If we are expecting the controller to be the source of
@@ -371,18 +371,18 @@ func testWorkerSessionCleanupMulti(burdenCase timeoutBurdenType) func(t *testing
 		case timeoutBurdenTypeWorker:
 			// Wait on worker, then check controller
 			sess.ExpectConnectionStateOnWorker(ctx, t, w1, session.StatusClosed)
-			sess.ExpectConnectionStateOnController(ctx, t, c1.Controller().SessionRepoFn, session.StatusConnected)
+			sess.ExpectConnectionStateOnController(ctx, t, c1.Controller().ConnectionRepoFn, session.StatusConnected)
 
 		case timeoutBurdenTypeController:
 			// Wait on controller, then check worker
-			sess.ExpectConnectionStateOnController(ctx, t, c1.Controller().SessionRepoFn, session.StatusClosed)
+			sess.ExpectConnectionStateOnController(ctx, t, c1.Controller().ConnectionRepoFn, session.StatusClosed)
 			sess.ExpectConnectionStateOnWorker(ctx, t, w1, session.StatusConnected)
 
 		default:
 			// Should be closed on both worker and controller. Wait on
 			// worker then check controller.
 			sess.ExpectConnectionStateOnWorker(ctx, t, w1, session.StatusClosed)
-			sess.ExpectConnectionStateOnController(ctx, t, c1.Controller().SessionRepoFn, session.StatusClosed)
+			sess.ExpectConnectionStateOnController(ctx, t, c1.Controller().ConnectionRepoFn, session.StatusClosed)
 		}
 
 		// Run send/receive test again to check expected connection-level
@@ -413,7 +413,7 @@ func testWorkerSessionCleanupMulti(burdenCase timeoutBurdenType) func(t *testing
 			// a connection status, ensure that our old session's
 			// connections are actually closed now that the worker is
 			// properly reporting in again.
-			sess.ExpectConnectionStateOnController(ctx, t, c1.Controller().SessionRepoFn, session.StatusClosed)
+			sess.ExpectConnectionStateOnController(ctx, t, c1.Controller().ConnectionRepoFn, session.StatusClosed)
 
 		case timeoutBurdenTypeController:
 			// If we are expecting the controller to be the source of

internal/tests/helper/testing_helper.go

@@ -168,7 +168,7 @@ func (s *TestSession) connect(ctx context.Context, t *testing.T) net.Conn {
 func (s *TestSession) ExpectConnectionStateOnController(
 	ctx context.Context,
 	t *testing.T,
-	sessionRepoFn common.SessionRepoFactory,
+	connectionRepoFn common.ConnectionRepoFactory,
 	expectState session.ConnectionStatus,
 ) {
 	t.Helper()
@@ -181,10 +181,10 @@ func (s *TestSession) ExpectConnectionStateOnController(
 	// This is just for initialization of the actual state set.
 	const sessionStatusUnknown session.ConnectionStatus = "unknown"
 
-	sessionRepo, err := sessionRepoFn()
+	connectionRepo, err := connectionRepoFn()
 	require.NoError(err)
 
-	conns, err := sessionRepo.ListConnectionsBySessionId(ctx, s.sessionId)
+	conns, err := connectionRepo.ListConnectionsBySessionId(ctx, s.sessionId)
 	require.NoError(err)
 	// To avoid misleading passing tests, we require this test be used
 	// with sessions with connections..
@@ -208,7 +208,7 @@ func (s *TestSession) ExpectConnectionStateOnController(
 		}
 
 		for i, conn := range conns {
-			_, states, err := sessionRepo.LookupConnection(ctx, conn.PublicId, nil)
+			_, states, err := connectionRepo.LookupConnection(ctx, conn.PublicId, nil)
 			require.NoError(err)
 			// Look at the first state in the returned list, which will
 			// be the most recent state.