diff --git a/internal/db/schema/migrations/oss/postgres/0/51_connection.up.sql b/internal/db/schema/migrations/oss/postgres/0/51_connection.up.sql index ef7a7e9bf..dd6a5e697 100644 --- a/internal/db/schema/migrations/oss/postgres/0/51_connection.up.sql +++ b/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. diff --git a/internal/db/schema/migrations/oss/postgres/0/69_wh_session_facts.up.sql b/internal/db/schema/migrations/oss/postgres/0/69_wh_session_facts.up.sql index a2df68b35..a8c5b8db2 100644 --- a/internal/db/schema/migrations/oss/postgres/0/69_wh_session_facts.up.sql +++ b/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. diff --git a/internal/db/schema/migrations/oss/postgres/16/05_wh_credential_dimension.up.sql b/internal/db/schema/migrations/oss/postgres/16/05_wh_credential_dimension.up.sql index 4627b57a0..eed53efee 100644 --- a/internal/db/schema/migrations/oss/postgres/16/05_wh_credential_dimension.up.sql +++ b/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; diff --git a/internal/db/schema/migrations/oss/postgres/27/01_disable_terminate_session.up.sql b/internal/db/schema/migrations/oss/postgres/27/01_disable_terminate_session.up.sql new file mode 100644 index 000000000..c1b8db55c --- /dev/null +++ b/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; \ No newline at end of file diff --git a/internal/db/schema/migrations/oss/postgres/27/02_wh_session_facts.up.sql b/internal/db/schema/migrations/oss/postgres/27/02_wh_session_facts.up.sql new file mode 100644 index 000000000..8f781c93d --- /dev/null +++ b/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; \ No newline at end of file diff --git a/internal/servers/controller/handlers/workers/worker_service.go b/internal/servers/controller/handlers/workers/worker_service.go index 98399425b..6307b27ae 100644 --- a/internal/servers/controller/handlers/workers/worker_service.go +++ b/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 } diff --git a/internal/session/query.go b/internal/session/query.go index a32af4a09..2e85ad435 100644 --- a/internal/session/query.go +++ b/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. diff --git a/internal/session/repository_connection.go b/internal/session/repository_connection.go index e98e82924..3639da422 100644 --- a/internal/session/repository_connection.go +++ b/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" diff --git a/internal/session/repository_connection_test.go b/internal/session/repository_connection_test.go index 529bbf9ba..d31ea1524 100644 --- a/internal/session/repository_connection_test.go +++ b/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) diff --git a/internal/session/repository_session.go b/internal/session/repository_session.go index 5fcccf912..4ecc5c6f7 100644 --- a/internal/session/repository_session.go +++ b/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 diff --git a/internal/session/repository_session_test.go b/internal/session/repository_session_test.go index 243e568f0..516c76281 100644 --- a/internal/session/repository_session_test.go +++ b/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()) diff --git a/internal/session/service_close_connections.go b/internal/session/service_close_connections.go new file mode 100644 index 000000000..34971722a --- /dev/null +++ b/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 +} diff --git a/internal/session/service_close_connections_test.go b/internal/session/service_close_connections_test.go new file mode 100644 index 000000000..8eb5d3fa6 --- /dev/null +++ b/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) + } + }) + } +}