refactor(session connections): Refactor connection closure

internal/db/schema/migrations/oss/postgres/0/51_connection.up.sql

@@ -175,6 +175,7 @@ begin;
   after insert on session_connection
     for each row execute procedure insert_new_connection_state();
 
+-- Replaced in 27/01_disable_terminate_session.up.sql
   -- update_connection_state_on_closed_reason() is used in an update trigger on the
   -- session_connection table.  it will insert a state of "closed" in
   -- session_connection_state for the closed session connection. 
@@ -284,6 +285,7 @@ begin;
   create trigger insert_session_connection_state before insert on session_connection_state
     for each row execute procedure insert_session_connection_state();
 
+-- Removed in 27/01_disable_terminate_session.up.sql
 -- terminate_session_if_possible takes a session id and terminates the session
 -- if the following conditions are met:
 --    * the session is expired and all its connections are closed.

internal/db/schema/migrations/oss/postgres/0/69_wh_session_facts.up.sql

@@ -146,6 +146,7 @@ begin;
     for each row
     execute function wh_insert_session_connection();
 
+  -- Updated in 27/01_disable_terminate_session.up.sql
   -- wh_update_session_connection returns an after update trigger for the
   -- session_connection table which updates a row in
   -- wh_session_connection_accumulating_fact for the session connection.

internal/db/schema/migrations/oss/postgres/16/05_wh_credential_dimension.up.sql

@@ -8,6 +8,7 @@ begin;
   alter table wh_session_connection_accumulating_fact
     alter column credential_group_key drop default;
 
+  -- Updated in 27/01_disable_terminate_session.up.sql
   -- replaces function from 15/01_wh_rename_key_columns.up.sql
   drop trigger wh_insert_session_connection on session_connection;
   drop function wh_insert_session_connection;

internal/db/schema/migrations/oss/postgres/27/01_disable_terminate_session.up.sql

@@ -0,0 +1,36 @@
+begin;
+-- Replaces function from 0/51_connection.up.sql
+-- Remove call to terminate_session_if_possible
+drop trigger update_connection_state_on_closed_reason on session_connection;
+drop function update_connection_state_on_closed_reason();
+create function
+    update_connection_state_on_closed_reason()
+    returns trigger
+as $$
+    begin
+        if new.closed_reason is not null then
+            -- check to see if there's a closed state already, before inserting a new one.
+            perform from
+                session_connection_state cs
+            where
+                    cs.connection_id = new.public_id and
+                    cs.state = 'closed';
+            if not found then
+                insert into session_connection_state (connection_id, state)
+                values
+                    (new.public_id, 'closed');
+            end if;
+        end if;
+    return new;
+    end;
+$$ language plpgsql;
+
+create trigger
+    update_connection_state_on_closed_reason
+    after update of closed_reason on session_connection
+    for each row execute procedure update_connection_state_on_closed_reason();
+
+-- Remove function, defined in 0/51_connection.up.sql
+drop function terminate_session_if_possible;
+
+commit;

internal/db/schema/migrations/oss/postgres/27/02_wh_session_facts.up.sql

@@ -0,0 +1,114 @@
+begin;
+-- Updating definition from 16/05_wh_credential_dimension.up.sql
+-- Remove call to wh_rollup_connections(new.session_id) from function
+drop trigger wh_insert_session_connection on session_connection;
+drop function wh_insert_session_connection();
+create function wh_insert_session_connection()
+    returns trigger
+as $$
+declare
+new_row wh_session_connection_accumulating_fact%rowtype;
+begin
+with
+    authorized_timestamp (date_dim_key, time_dim_key, ts) as (
+        select wh_date_key(start_time), wh_time_key(start_time), start_time
+        from session_connection_state
+        where connection_id = new.public_id
+          and state = 'authorized'
+    ),
+    session_dimension (host_dim_key, user_dim_key, credential_group_dim_key) as (
+        select host_key, user_key, credential_group_key
+        from wh_session_accumulating_fact
+        where session_id = new.session_id
+    )
+insert into wh_session_connection_accumulating_fact (
+        connection_id,
+        session_id,
+        host_key,
+        user_key,
+        credential_group_key,
+        connection_authorized_date_key,
+        connection_authorized_time_key,
+        connection_authorized_time,
+        client_tcp_address,
+        client_tcp_port_number,
+        endpoint_tcp_address,
+        endpoint_tcp_port_number,
+        bytes_up,
+        bytes_down
+    )
+select new.public_id,
+       new.session_id,
+       session_dimension.host_dim_key,
+       session_dimension.user_dim_key,
+       session_dimension.credential_group_dim_key,
+       authorized_timestamp.date_dim_key,
+       authorized_timestamp.time_dim_key,
+       authorized_timestamp.ts,
+       new.client_tcp_address,
+       new.client_tcp_port,
+       new.endpoint_tcp_address,
+       new.endpoint_tcp_port,
+       new.bytes_up,
+       new.bytes_down
+from authorized_timestamp,
+     session_dimension
+         returning * into strict new_row;
+return null;
+end;
+$$ language plpgsql;
+
+create trigger wh_insert_session_connection
+    after insert on session_connection
+    for each row
+    execute function wh_insert_session_connection();
+
+-- Updating definition from 0/69_wh_session_facts.up.sql
+-- Remove call to wh_rollup_connections(new.session_id) from function
+drop trigger wh_update_session_connection on session_connection;
+drop function wh_update_session_connection;
+create function wh_update_session_connection()
+    returns trigger
+as $$
+declare
+updated_row wh_session_connection_accumulating_fact%rowtype;
+begin
+update wh_session_connection_accumulating_fact
+set client_tcp_address       = new.client_tcp_address,
+    client_tcp_port_number   = new.client_tcp_port,
+    endpoint_tcp_address     = new.endpoint_tcp_address,
+    endpoint_tcp_port_number = new.endpoint_tcp_port,
+    bytes_up                 = new.bytes_up,
+    bytes_down               = new.bytes_down
+where connection_id = new.public_id
+    returning * into strict updated_row;
+return null;
+end;
+$$ language plpgsql;
+
+create trigger wh_update_session_connection
+    after update on session_connection
+    for each row
+    execute function wh_update_session_connection();
+
+create function
+    wh_session_rollup()
+    returns trigger
+as $$
+begin
+    if new.termination_reason is not null then
+        -- Rollup will fail if no connections were made for a session
+        if exists (select from session_connection where session_id = new.public_id) then
+            perform wh_rollup_connections(new.public_id);
+        end if;
+    end if;
+return null;
+end;
+$$ language plpgsql;
+
+create trigger
+    wh_rollup_connections_on_session_termination
+    after update of termination_reason on session
+    for each row execute procedure wh_session_rollup();
+
+commit;

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

@@ -469,12 +469,17 @@ func (ws *workerServiceServer) CloseConnection(ctx context.Context, req *pbs.Clo
 			ClosedReason: session.ClosedReason(v.GetReason()),
 		})
 	}
+	connRepo, err := ws.connectionRepoFn()
+	if err != nil {
+		return nil, status.Errorf(codes.Internal, "error getting connection repo: %v", err)
+	}
+
 	sessRepo, err := ws.sessionRepoFn()
 	if err != nil {
 		return nil, status.Errorf(codes.Internal, "error getting session repo: %v", err)
 	}
 
-	closeInfos, err := sessRepo.CloseConnections(ctx, closeWiths)
+	closeInfos, err := session.CloseConnections(ctx, sessRepo, connRepo, closeWiths)
 	if err != nil {
 		return nil, err
 	}

internal/session/query.go

@@ -129,6 +129,92 @@ where
 %s
 `
 
+	terminateSessionIfPossible = `
+    -- is terminate_session_id in a canceling state
+    with session_version as (
+		select 
+			version
+		from 
+			session
+		where public_id = @public_id
+	),
+    canceling_session(session_id) as
+    (
+      select 
+        session_id
+      from
+        session_state ss
+      where 
+        ss.session_id = @public_id and
+        ss.state = 'canceling' and 
+        ss.end_time is null
+    )
+    update session us
+      set version = version +1,
+	  termination_reason = 
+      case 
+        -- timed out sessions
+        when now() > us.expiration_time then 'timed out'
+        -- canceling sessions
+        when us.public_id in(
+          select 
+            session_id 
+          from 
+            canceling_session cs 
+          where
+            us.public_id = cs.session_id
+          ) then 'canceled' 
+        -- default: session connection limit reached.
+        else 'connection limit'
+      end
+    where
+      -- limit update to just the terminating_session_id
+      us.public_id = @public_id and
+  	  us.version = (select * from session_version) and
+      termination_reason is null and
+      -- session expired or connection limit reached
+      (
+        -- expired sessions...
+        now() > us.expiration_time or 
+        -- connection limit reached...
+        (
+          -- handle unlimited connections...
+          connection_limit != -1 and
+          (
+            select count (*) 
+              from session_connection sc 
+            where 
+              sc.session_id = us.public_id
+          ) >= connection_limit
+        ) or 
+        -- canceled sessions
+        us.public_id in (
+          select 
+            session_id
+          from
+            canceling_session cs
+          where 
+            us.public_id = cs.session_id 
+        )
+      ) and 
+      -- make sure there are no existing connections
+      us.public_id not in (
+        select 
+          session_id 
+        from 
+            session_connection
+          where public_id in (
+          select 
+            connection_id
+          from 
+            session_connection_state
+          where 
+            state != 'closed' and
+            end_time is null
+        )
+    )
+`
+
 	// termSessionUpdate is one stmt that terminates sessions for the following
 	// reasons:
 	//	* sessions that are expired and all their connections are closed.

internal/session/repository_connection.go

@@ -240,6 +240,71 @@ func (r *ConnectionRepository) ConnectConnection(ctx context.Context, c ConnectW
 	return &connection, connectionStates, nil
 }
 
+// closeConnectionResp is just a wrapper for the response from CloseConnections.
+// It wraps the connection and its states for each connection closed.
+type closeConnectionResp struct {
+	Connection       *Connection
+	ConnectionStates []*ConnectionState
+}
+
+// closeConnections set's a connection's state to "closed" in the repo.  It's
+// called by a worker after it's closed a connection between the client and the
+// endpoint
+func (r *ConnectionRepository) closeConnections(ctx context.Context, closeWith []CloseWith, _ ...Option) ([]closeConnectionResp, error) {
+	const op = "session.(ConnectionRepository).closeConnections"
+	if len(closeWith) == 0 {
+		return nil, errors.New(ctx, errors.InvalidParameter, op, "missing connections")
+	}
+	for _, cw := range closeWith {
+		if err := cw.validate(); err != nil {
+			return nil, errors.Wrap(ctx, err, op, errors.WithMsg(fmt.Sprintf("%s was invalid", cw.ConnectionId)))
+		}
+	}
+	var resp []closeConnectionResp
+	_, err := r.writer.DoTx(
+		ctx,
+		db.StdRetryCnt,
+		db.ExpBackoff{},
+		func(reader db.Reader, w db.Writer) error {
+			for _, cw := range closeWith {
+				updateConnection := AllocConnection()
+				updateConnection.PublicId = cw.ConnectionId
+				updateConnection.BytesUp = cw.BytesUp
+				updateConnection.BytesDown = cw.BytesDown
+				updateConnection.ClosedReason = cw.ClosedReason.String()
+				// updating the ClosedReason will trigger an insert into the
+				// session_connection_state with a state of closed.
+				rowsUpdated, err := w.Update(
+					ctx,
+					&updateConnection,
+					[]string{"BytesUp", "BytesDown", "ClosedReason"},
+					nil,
+				)
+				if err != nil {
+					return errors.Wrap(ctx, err, op, errors.WithMsg(fmt.Sprintf("unable to update connection %s", cw.ConnectionId)))
+				}
+				if rowsUpdated != 1 {
+					return errors.New(ctx, errors.MultipleRecords, op, fmt.Sprintf("%d would have been updated for connection %s", rowsUpdated, cw.ConnectionId))
+				}
+				states, err := fetchConnectionStates(ctx, reader, cw.ConnectionId, db.WithOrder("start_time desc"))
+				if err != nil {
+					return errors.Wrap(ctx, err, op)
+				}
+				resp = append(resp, closeConnectionResp{
+					Connection:       &updateConnection,
+					ConnectionStates: states,
+				})
+
+			}
+			return nil
+		},
+	)
+	if err != nil {
+		return nil, errors.Wrap(ctx, err, op)
+	}
+	return resp, nil
+}
+
 // DeleteConnection will delete a connection from the repository.
 func (r *ConnectionRepository) DeleteConnection(ctx context.Context, publicId string, _ ...Option) (int, error) {
 	const op = "session.(ConnectionRepository).DeleteConnection"

internal/session/repository_connection_test.go

@@ -553,7 +553,7 @@ func TestRepository_CloseConnectionsForDeadWorkers(t *testing.T) {
 			require.True(foundAuthorized)
 			require.True(foundConnected)
 		} else if i%3 == 1 {
-			resp, err := repo.CloseConnections(ctx, []CloseWith{
+			resp, err := connRepo.closeConnections(ctx, []CloseWith{
 				{
 					ConnectionId: connId,
 					ClosedReason: ConnectionCanceled,
@@ -809,7 +809,7 @@ func TestRepository_ShouldCloseConnectionsOnWorker(t *testing.T) {
 			require.True(foundAuthorized)
 			require.True(foundConnected)
 		} else {
-			resp, err := repo.CloseConnections(ctx, []CloseWith{
+			resp, err := connRepo.closeConnections(ctx, []CloseWith{
 				{
 					ConnectionId: connId,
 					ClosedReason: ConnectionCanceled,
@@ -897,6 +897,88 @@ func TestRepository_ShouldCloseConnectionsOnWorker(t *testing.T) {
 	}
 }
 
+func TestRepository_CloseConnections(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(cnt int) []CloseWith {
+		s := TestDefaultSession(t, conn, wrapper, iamRepo)
+		srv := TestWorker(t, conn, wrapper)
+		tofu := TestTofu(t)
+		s, _, err := repo.ActivateSession(context.Background(), s.PublicId, s.Version, srv.PrivateId, srv.Type, tofu)
+		require.NoError(t, err)
+		cw := make([]CloseWith, 0, cnt)
+		for i := 0; i < cnt; i++ {
+			c := TestConnection(t, conn, s.PublicId, "127.0.0.1", 22, "127.0.0.1", 2222, "127.0.0.1")
+			require.NoError(t, err)
+			cw = append(cw, CloseWith{
+				ConnectionId: c.PublicId,
+				BytesUp:      1,
+				BytesDown:    2,
+				ClosedReason: ConnectionClosedByUser,
+			})
+		}
+		return cw
+	}
+	tests := []struct {
+		name        string
+		closeWith   []CloseWith
+		reason      TerminationReason
+		wantErr     bool
+		wantIsError errors.Code
+	}{
+		{
+			name:      "valid",
+			closeWith: setupFn(2),
+			reason:    ClosedByUser,
+		},
+		{
+			name:        "empty-closed-with",
+			closeWith:   []CloseWith{},
+			reason:      ClosedByUser,
+			wantErr:     true,
+			wantIsError: errors.InvalidParameter,
+		},
+		{
+			name: "missing-ConnectionId",
+			closeWith: func() []CloseWith {
+				cw := setupFn(2)
+				cw[1].ConnectionId = ""
+				return cw
+			}(),
+			reason:      ClosedByUser,
+			wantErr:     true,
+			wantIsError: errors.InvalidParameter,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			assert, require := assert.New(t), require.New(t)
+			resp, err := connRepo.closeConnections(context.Background(), tt.closeWith)
+			if tt.wantErr {
+				require.Error(err)
+				assert.Truef(errors.Match(errors.T(tt.wantIsError), err), "unexpected error %s", err.Error())
+				return
+			}
+			require.NoError(err)
+			assert.Equal(len(tt.closeWith), len(resp))
+			for _, r := range resp {
+				require.NotNil(r.Connection)
+				require.NotNil(r.ConnectionStates)
+				assert.Equal(StatusClosed, r.ConnectionStates[0].Status)
+			}
+		})
+	}
+}
+
 func testIsConnectionClosed(ctx context.Context, t *testing.T, repo *ConnectionRepository, connId string) bool {
 	require := require.New(t)
 	_, states, err := repo.LookupConnection(ctx, connId)

internal/session/repository_session.go

@@ -375,6 +375,35 @@ func (r *Repository) TerminateCompletedSessions(ctx context.Context) (int, error
 	return rowsAffected, nil
 }
 
+// terminateSessionIfPossible is called on connection close and will attempt to close the connection's
+// session if the following conditions are met:
+//  * sessions that have exhausted their connection limit and all their connections are closed.
+//	* sessions that are expired and all their connections are closed.
+//	* sessions that are canceling and all their connections are closed
+func (r *Repository) terminateSessionIfPossible(ctx context.Context, sessionId string) (int, error) {
+	const op = "session.(Repository).terminateSessionIfPossible"
+	rowsAffected := 0
+
+	_, err := r.writer.DoTx(
+		ctx,
+		db.StdRetryCnt,
+		db.ExpBackoff{},
+		func(reader db.Reader, w db.Writer) error {
+			var err error
+			rowsAffected, err = w.Exec(ctx, terminateSessionIfPossible,
+				[]interface{}{sql.Named("public_id", sessionId)})
+			if err != nil {
+				return errors.Wrap(ctx, err, op)
+			}
+			return nil
+		},
+	)
+	if err != nil {
+		return db.NoRowsAffected, errors.Wrap(ctx, err, op)
+	}
+	return rowsAffected, nil
+}
+
 type AuthzSummary struct {
 	ExpirationTime         *timestamp.Timestamp
 	ConnectionLimit        int32
@@ -402,71 +431,6 @@ func (r *Repository) sessionAuthzSummary(ctx context.Context, sessionId string)
 	return info, nil
 }
 
-// CloseConnectionResp is just a wrapper for the response from CloseConnections.
-// It wraps the connection and its states for each connection closed.
-type CloseConnectionResp struct {
-	Connection       *Connection
-	ConnectionStates []*ConnectionState
-}
-
-// CloseConnections set's a connection's state to "closed" in the repo.  It's
-// 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.(ConnectionRepository).CloseConnections"
-	if len(closeWith) == 0 {
-		return nil, errors.New(ctx, errors.InvalidParameter, op, "missing connections")
-	}
-	for _, cw := range closeWith {
-		if err := cw.validate(); err != nil {
-			return nil, errors.Wrap(ctx, err, op, errors.WithMsg(fmt.Sprintf("%s was invalid", cw.ConnectionId)))
-		}
-	}
-	var resp []CloseConnectionResp
-	_, err := r.writer.DoTx(
-		ctx,
-		db.StdRetryCnt,
-		db.ExpBackoff{},
-		func(reader db.Reader, w db.Writer) error {
-			for _, cw := range closeWith {
-				updateConnection := AllocConnection()
-				updateConnection.PublicId = cw.ConnectionId
-				updateConnection.BytesUp = cw.BytesUp
-				updateConnection.BytesDown = cw.BytesDown
-				updateConnection.ClosedReason = cw.ClosedReason.String()
-				// updating the ClosedReason will trigger an insert into the
-				// session_connection_state with a state of closed.
-				rowsUpdated, err := w.Update(
-					ctx,
-					&updateConnection,
-					[]string{"BytesUp", "BytesDown", "ClosedReason"},
-					nil,
-				)
-				if err != nil {
-					return errors.Wrap(ctx, err, op, errors.WithMsg(fmt.Sprintf("unable to update connection %s", cw.ConnectionId)))
-				}
-				if rowsUpdated != 1 {
-					return errors.New(ctx, errors.MultipleRecords, op, fmt.Sprintf("%d would have been updated for connection %s", rowsUpdated, cw.ConnectionId))
-				}
-				states, err := fetchConnectionStates(ctx, reader, cw.ConnectionId, db.WithOrder("start_time desc"))
-				if err != nil {
-					return errors.Wrap(ctx, err, op)
-				}
-				resp = append(resp, CloseConnectionResp{
-					Connection:       &updateConnection,
-					ConnectionStates: states,
-				})
-
-			}
-			return nil
-		},
-	)
-	if err != nil {
-		return nil, errors.Wrap(ctx, err, op)
-	}
-	return resp, nil
-}
-
 // ActivateSession will activate the session and is called by a worker after
 // authenticating the session. The session must be in a "pending" state to be
 // activated. States are ordered by start time descending. Returns an

internal/session/repository_session_test.go

@@ -573,7 +573,7 @@ func TestRepository_TerminateCompletedSessions(t *testing.T) {
 				BytesDown:    1,
 				ClosedReason: ConnectionClosedByUser,
 			}
-			_, err = repo.CloseConnections(context.Background(), []CloseWith{cw})
+			_, err = connRepo.closeConnections(context.Background(), []CloseWith{cw})
 			require.NoError(t, err)
 		}
 		return s
@@ -781,94 +781,16 @@ func TestRepository_TerminateCompletedSessions(t *testing.T) {
 	}
 }
 
-func TestRepository_CloseConnections(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)
-
-	setupFn := func(cnt int) []CloseWith {
-		s := TestDefaultSession(t, conn, wrapper, iamRepo)
-		srv := TestWorker(t, conn, wrapper)
-		tofu := TestTofu(t)
-		s, _, err := repo.ActivateSession(context.Background(), s.PublicId, s.Version, srv.PrivateId, srv.Type, tofu)
-		require.NoError(t, err)
-		cw := make([]CloseWith, 0, cnt)
-		for i := 0; i < cnt; i++ {
-			c := TestConnection(t, conn, s.PublicId, "127.0.0.1", 22, "127.0.0.1", 2222, "127.0.0.1")
-			require.NoError(t, err)
-			cw = append(cw, CloseWith{
-				ConnectionId: c.PublicId,
-				BytesUp:      1,
-				BytesDown:    2,
-				ClosedReason: ConnectionClosedByUser,
-			})
-		}
-		return cw
-	}
-	tests := []struct {
-		name        string
-		closeWith   []CloseWith
-		reason      TerminationReason
-		wantErr     bool
-		wantIsError errors.Code
-	}{
-		{
-			name:      "valid",
-			closeWith: setupFn(2),
-			reason:    ClosedByUser,
-		},
-		{
-			name:        "empty-closed-with",
-			closeWith:   []CloseWith{},
-			reason:      ClosedByUser,
-			wantErr:     true,
-			wantIsError: errors.InvalidParameter,
-		},
-		{
-			name: "missing-ConnectionId",
-			closeWith: func() []CloseWith {
-				cw := setupFn(2)
-				cw[1].ConnectionId = ""
-				return cw
-			}(),
-			reason:      ClosedByUser,
-			wantErr:     true,
-			wantIsError: errors.InvalidParameter,
-		},
-	}
-	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
-			assert, require := assert.New(t), require.New(t)
-			resp, err := repo.CloseConnections(context.Background(), tt.closeWith)
-			if tt.wantErr {
-				require.Error(err)
-				assert.Truef(errors.Match(errors.T(tt.wantIsError), err), "unexpected error %s", err.Error())
-				return
-			}
-			require.NoError(err)
-			assert.Equal(len(tt.closeWith), len(resp))
-			for _, r := range resp {
-				require.NotNil(r.Connection)
-				require.NotNil(r.ConnectionStates)
-				assert.Equal(StatusClosed, r.ConnectionStates[0].Status)
-			}
-		})
-	}
-}
-
 func TestRepository_CancelSession(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() *Session {
 		session := TestDefaultSession(t, conn, wrapper, iamRepo)
@@ -900,11 +822,13 @@ func TestRepository_CancelSession(t *testing.T) {
 					BytesDown:    1,
 					ClosedReason: ConnectionClosedByUser,
 				}
-				_, err = repo.CloseConnections(context.Background(), []CloseWith{cw})
+				_, err = CloseConnections(ctx, repo, connRepo, []CloseWith{cw})
 				require.NoError(t, err)
-				s, _, err := repo.LookupSession(context.Background(), session.PublicId)
+				s, _, err := repo.LookupSession(ctx, session.PublicId)
 				require.NoError(t, err)
 				assert.Equal(t, StatusTerminated, s.States[0].Status)
+				// The two transactions to cancel connections and terminate the session will result in version being 2, not 1
+				session.Version = s.Version
 				return session
 			}(),
 			wantStatus: StatusTerminated,
@@ -970,7 +894,7 @@ func TestRepository_CancelSession(t *testing.T) {
 			default:
 				version = tt.session.Version
 			}
-			s, err := repo.CancelSession(context.Background(), id, version)
+			s, err := repo.CancelSession(ctx, id, version)
 			if tt.wantErr {
 				require.Error(err)
 				assert.Truef(errors.Match(errors.T(tt.wantIsError), err), "unexpected error %s", err.Error())

internal/session/service_close_connections.go

@@ -0,0 +1,31 @@
+package session
+
+import (
+	"context"
+
+	"github.com/hashicorp/boundary/internal/errors"
+)
+
+// CloseConnections is a domain service function that:
+// * closes requested connections
+// * uses the sessionId of the connection to see if the session meets conditions for termination
+func CloseConnections(ctx context.Context, sessionRepoFn *Repository, connectionRepoFn *ConnectionRepository,
+	closeWiths []CloseWith) ([]closeConnectionResp, error) {
+	const op = "session.AuthorizeConnection"
+
+	closeInfos, err := connectionRepoFn.closeConnections(ctx, closeWiths)
+	if err != nil {
+		return nil, errors.Wrap(ctx, err, op)
+	}
+
+	// Attempt to terminate only once per sessionId
+	sessionIdsProcessed := make(map[string]bool)
+	for _, c := range closeInfos {
+		if !sessionIdsProcessed[c.Connection.SessionId] {
+			sessionRepoFn.terminateSessionIfPossible(ctx, c.Connection.SessionId)
+			sessionIdsProcessed[c.Connection.SessionId] = true
+		}
+	}
+
+	return closeInfos, nil
+}

internal/session/service_close_connections_test.go

@@ -0,0 +1,97 @@
+package session
+
+import (
+	"context"
+	"testing"
+
+	"github.com/hashicorp/boundary/internal/db"
+	"github.com/hashicorp/boundary/internal/iam"
+	"github.com/hashicorp/boundary/internal/kms"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestServiceCloseConnections(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)
+
+	type sessionAndCloseWiths struct {
+		session   *Session
+		closeWith []CloseWith
+	}
+
+	setupFn := func(cnt int, addtlConn int) sessionAndCloseWiths {
+		s := TestDefaultSession(t, conn, wrapper, iamRepo)
+		srv := TestWorker(t, conn, wrapper)
+		tofu := TestTofu(t)
+		s, _, err = repo.ActivateSession(context.Background(), s.PublicId, s.Version, srv.PrivateId, srv.Type, tofu)
+		require.NoError(t, err)
+
+		require.NoError(t, err)
+		cw := make([]CloseWith, 0, cnt)
+		for i := 0; i < cnt; i++ {
+			c := TestConnection(t, conn, s.PublicId, "127.0.0.1", 22, "127.0.0.1", 2222, "127.0.0.1")
+			require.NoError(t, err)
+			cw = append(cw, CloseWith{
+				ConnectionId: c.PublicId,
+				BytesUp:      1,
+				BytesDown:    2,
+				ClosedReason: ConnectionClosedByUser,
+			})
+		}
+
+		for i := 0; i < addtlConn; i++ {
+			TestConnection(t, conn, s.PublicId, "127.0.0.1", 22, "127.0.0.1", 2222, "127.0.0.1")
+			require.NoError(t, err)
+		}
+		return sessionAndCloseWiths{s, cw}
+	}
+
+	tests := []struct {
+		name              string
+		sessionCW         sessionAndCloseWiths
+		wantClosedSession bool
+	}{
+		{
+			name:              "close-multiple-connections-and-session",
+			sessionCW:         setupFn(4, 0),
+			wantClosedSession: true,
+		},
+		{
+			name:              "close-subset-of-connections",
+			sessionCW:         setupFn(2, 1),
+			wantClosedSession: false,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			assert, require := assert.New(t), require.New(t)
+
+			resp, err := CloseConnections(ctx, repo, connRepo, tt.sessionCW.closeWith)
+			require.NoError(err)
+
+			for _, r := range resp {
+				require.NotNil(r.Connection)
+				require.NotNil(r.ConnectionStates)
+				assert.Equal(StatusClosed, r.ConnectionStates[0].Status)
+			}
+
+			// Ensure session is in the state we want- terminated if all conns closed, else active
+			ses, _, err := repo.LookupSession(ctx, tt.sessionCW.session.PublicId)
+			require.NoError(err)
+			if tt.wantClosedSession {
+				assert.Equal(StatusTerminated, ses.States[0].Status)
+			} else {
+				assert.Equal(StatusActive, ses.States[0].Status)
+			}
+		})
+	}
+}