new domain functions for the session repo (#368)

internal/db/migrations/postgres.gen.go

@@ -3508,6 +3508,7 @@ begin;
   as $$
   begin
     if new.termination_reason is not null then
+      -- check to see if there are any open connections.
       perform from
         session_connection sc,
         session_connection_state scs
@@ -3518,6 +3519,18 @@ begin;
       if found then 
         raise 'session %s has existing open connections', new.public_id;
       end if;
+      
+      -- check to see if there's a terminated state already, before inserting a
+      -- new one.
+      perform from
+        session_state ss
+      where
+        ss.session_id = new.public_id and 
+        ss.state = 'terminated';
+      if found then 
+        return new;
+      end if;
+
       insert into session_state (session_id, state)
       values
         (new.public_id, 'terminated');
@@ -3738,8 +3751,8 @@ begin;
     ('system error');
 
   -- A session connection is one connection proxied by a worker from a client to
-  -- a backend for a session. The client initiates the connection to the worker
-  -- and the worker initiates the connection to the backend.
+  -- a endpoint for a session. The client initiates the connection to the worker
+  -- and the worker initiates the connection to the endpoint.
   -- A session can have zero or more session connections.
   create table session_connection (
     public_id wt_public_id primary key,
@@ -3758,26 +3771,26 @@ begin;
         and
         client_tcp_port <= 65535
       ),
-    -- the backend_tcp_address is the network address of the backend which the
+    -- the endpoint_tcp_address is the network address of the endpoint which the
     -- worker initiated the connection to, for the user
-    backend_tcp_address inet not null,
-    -- the backend_tcp_port is the network port at the address of the backend the
+    endpoint_tcp_address inet not null,
+    -- the endpoint_tcp_port is the network port at the address of the endpoint the
     -- worker proxied a connection to, for the user
-    backend_tcp_port integer not null
+    endpoint_tcp_port integer not null
       check(
-        backend_tcp_port > 0
+        endpoint_tcp_port > 0
         and
-        backend_tcp_port <= 65535
+        endpoint_tcp_port <= 65535
       ),
     -- the total number of bytes received by the worker from the client and sent
-    -- to the backend for this connection
+    -- to the endpoint for this connection
     bytes_up bigint -- can be null
       check (
         bytes_up is null
         or
         bytes_up >= 0
       ),
-    -- the total number of bytes received by the worker from the backend and sent
+    -- the total number of bytes received by the worker from the endpoint and sent
     -- to the client for this connection
     bytes_down bigint -- can be null
       check (
@@ -3798,7 +3811,7 @@ begin;
     immutable_columns
   before
   update on session_connection
-    for each row execute procedure immutable_columns('public_id', 'session_id', 'client_tcp_address', 'client_tcp_port', 'backend_tcp_address', 'backend_tcp_port', 'create_time');
+    for each row execute procedure immutable_columns('public_id', 'session_id', 'client_tcp_address', 'client_tcp_port', 'endpoint_tcp_address', 'endpoint_tcp_port', 'create_time');
 
   create trigger 
     update_version_column 
@@ -3845,11 +3858,20 @@ begin;
   as $$
   begin
     if new.closed_reason is not null then
-      insert into session_connection_state (connection_id, state)
-      values
-        (new.public_id, 'closed');
+      -- 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;
-      return new;
+    end if;
+    return new;
   end;
   $$ language plpgsql;
 

internal/db/migrations/postgres/50_session.up.sql

@@ -232,6 +232,7 @@ begin;
   as $$
   begin
     if new.termination_reason is not null then
+      -- check to see if there are any open connections.
       perform from
         session_connection sc,
         session_connection_state scs
@@ -242,6 +243,18 @@ begin;
       if found then 
         raise 'session %s has existing open connections', new.public_id;
       end if;
+      
+      -- check to see if there's a terminated state already, before inserting a
+      -- new one.
+      perform from
+        session_state ss
+      where
+        ss.session_id = new.public_id and 
+        ss.state = 'terminated';
+      if found then 
+        return new;
+      end if;
+
       insert into session_state (session_id, state)
       values
         (new.public_id, 'terminated');

internal/db/migrations/postgres/51_connection.up.sql

@@ -75,8 +75,8 @@ begin;
     ('system error');
 
   -- A session connection is one connection proxied by a worker from a client to
-  -- a backend for a session. The client initiates the connection to the worker
-  -- and the worker initiates the connection to the backend.
+  -- a endpoint for a session. The client initiates the connection to the worker
+  -- and the worker initiates the connection to the endpoint.
   -- A session can have zero or more session connections.
   create table session_connection (
     public_id wt_public_id primary key,
@@ -95,26 +95,26 @@ begin;
         and
         client_tcp_port <= 65535
       ),
-    -- the backend_tcp_address is the network address of the backend which the
+    -- the endpoint_tcp_address is the network address of the endpoint which the
     -- worker initiated the connection to, for the user
-    backend_tcp_address inet not null,
-    -- the backend_tcp_port is the network port at the address of the backend the
+    endpoint_tcp_address inet not null,
+    -- the endpoint_tcp_port is the network port at the address of the endpoint the
     -- worker proxied a connection to, for the user
-    backend_tcp_port integer not null
+    endpoint_tcp_port integer not null
       check(
-        backend_tcp_port > 0
+        endpoint_tcp_port > 0
         and
-        backend_tcp_port <= 65535
+        endpoint_tcp_port <= 65535
       ),
     -- the total number of bytes received by the worker from the client and sent
-    -- to the backend for this connection
+    -- to the endpoint for this connection
     bytes_up bigint -- can be null
       check (
         bytes_up is null
         or
         bytes_up >= 0
       ),
-    -- the total number of bytes received by the worker from the backend and sent
+    -- the total number of bytes received by the worker from the endpoint and sent
     -- to the client for this connection
     bytes_down bigint -- can be null
       check (
@@ -135,7 +135,7 @@ begin;
     immutable_columns
   before
   update on session_connection
-    for each row execute procedure immutable_columns('public_id', 'session_id', 'client_tcp_address', 'client_tcp_port', 'backend_tcp_address', 'backend_tcp_port', 'create_time');
+    for each row execute procedure immutable_columns('public_id', 'session_id', 'client_tcp_address', 'client_tcp_port', 'endpoint_tcp_address', 'endpoint_tcp_port', 'create_time');
 
   create trigger 
     update_version_column 
@@ -182,11 +182,20 @@ begin;
   as $$
   begin
     if new.closed_reason is not null then
-      insert into session_connection_state (connection_id, state)
-      values
-        (new.public_id, 'closed');
+      -- 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;
-      return new;
+    end if;
+    return new;
   end;
   $$ language plpgsql;
 

internal/session/connection.go

@@ -23,10 +23,10 @@ type Connection struct {
 	ClientTcpAddress string `json:"client_tcp_address,omitempty" gorm:"default:null"`
 	// ClientTcpPort of the connection
 	ClientTcpPort uint32 `json:"client_tcp_port,omitempty" gorm:"default:null"`
-	// BackendTcpAddress of the connection
-	BackendTcpAddress string `json:"backend_tcp_address,omitempty" gorm:"default:null"`
-	// BackendTcpPort of the connection
-	BackendTcpPort uint32 `json:"backend_tcp_port,omitempty" gorm:"default:null"`
+	// EndpointTcpAddress of the connection
+	EndpointTcpAddress string `json:"endpoint_tcp_address,omitempty" gorm:"default:null"`
+	// EndpointTcpPort of the connection
+	EndpointTcpPort uint32 `json:"endpoint_tcp_port,omitempty" gorm:"default:null"`
 	// BytesUp of the connection
 	BytesUp uint64 `json:"bytes_up,omitempty" gorm:"default:null"`
 	// BytesDown of the connection
@@ -52,13 +52,13 @@ var _ db.VetForWriter = (*Connection)(nil)
 
 // New creates a new in memory session.  No options
 // are currently supported.
-func NewConnection(sessionID, clientTcpAddress string, clientTcpPort uint32, backendTcpAddr string, backendTcpPort uint32, opt ...Option) (*Connection, error) {
+func NewConnection(sessionID, clientTcpAddress string, clientTcpPort uint32, endpointTcpAddr string, endpointTcpPort uint32, opt ...Option) (*Connection, error) {
 	c := Connection{
-		SessionId:         sessionID,
-		ClientTcpAddress:  clientTcpAddress,
-		ClientTcpPort:     clientTcpPort,
-		BackendTcpAddress: backendTcpAddr,
-		BackendTcpPort:    backendTcpPort,
+		SessionId:          sessionID,
+		ClientTcpAddress:   clientTcpAddress,
+		ClientTcpPort:      clientTcpPort,
+		EndpointTcpAddress: endpointTcpAddr,
+		EndpointTcpPort:    endpointTcpPort,
 	}
 	if err := c.validateNewConnection("new connection:"); err != nil {
 		return nil, err
@@ -74,16 +74,16 @@ func AllocConnection() Connection {
 // Clone creates a clone of the Session
 func (c *Connection) Clone() interface{} {
 	clone := &Connection{
-		PublicId:          c.PublicId,
-		SessionId:         c.SessionId,
-		ClientTcpAddress:  c.ClientTcpAddress,
-		ClientTcpPort:     c.ClientTcpPort,
-		BackendTcpAddress: c.BackendTcpAddress,
-		BackendTcpPort:    c.BackendTcpPort,
-		BytesUp:           c.BytesUp,
-		BytesDown:         c.BytesDown,
-		ClosedReason:      c.ClosedReason,
-		Version:           c.Version,
+		PublicId:           c.PublicId,
+		SessionId:          c.SessionId,
+		ClientTcpAddress:   c.ClientTcpAddress,
+		ClientTcpPort:      c.ClientTcpPort,
+		EndpointTcpAddress: c.EndpointTcpAddress,
+		EndpointTcpPort:    c.EndpointTcpPort,
+		BytesUp:            c.BytesUp,
+		BytesDown:          c.BytesDown,
+		ClosedReason:       c.ClosedReason,
+		Version:            c.Version,
 	}
 	if c.CreateTime != nil {
 		clone.CreateTime = &timestamp.Timestamp{
@@ -126,10 +126,10 @@ func (c *Connection) VetForWrite(ctx context.Context, r db.Reader, opType db.OpT
 			return fmt.Errorf("connection vet for write: client address is immutable: %w", db.ErrInvalidParameter)
 		case contains(opts.WithFieldMaskPaths, "ClientTcpPort"):
 			return fmt.Errorf("connection vet for write: client port is immutable: %w", db.ErrInvalidParameter)
-		case contains(opts.WithFieldMaskPaths, "BackendTcpAddress"):
-			return fmt.Errorf("connection vet for write: backend address is immutable: %w", db.ErrInvalidParameter)
-		case contains(opts.WithFieldMaskPaths, "BackendTcpPort"):
-			return fmt.Errorf("connection vet for write: backend port is immutable: %w", db.ErrInvalidParameter)
+		case contains(opts.WithFieldMaskPaths, "EndpointTcpAddress"):
+			return fmt.Errorf("connection vet for write: endpoint address is immutable: %w", db.ErrInvalidParameter)
+		case contains(opts.WithFieldMaskPaths, "EndpointTcpPort"):
+			return fmt.Errorf("connection vet for write: endpoint port is immutable: %w", db.ErrInvalidParameter)
 		case contains(opts.WithFieldMaskPaths, "CreateTime"):
 			return fmt.Errorf("connection vet for write: create time is immutable: %w", db.ErrInvalidParameter)
 		case contains(opts.WithFieldMaskPaths, "UpdateTime"):
@@ -169,11 +169,11 @@ func (c *Connection) validateNewConnection(errorPrefix string) error {
 	if c.ClientTcpPort == 0 {
 		return fmt.Errorf("%s missing client port: %w", errorPrefix, db.ErrInvalidParameter)
 	}
-	if c.BackendTcpAddress == "" {
-		return fmt.Errorf("%s missing backend address: %w", errorPrefix, db.ErrInvalidParameter)
+	if c.EndpointTcpAddress == "" {
+		return fmt.Errorf("%s missing endpoint address: %w", errorPrefix, db.ErrInvalidParameter)
 	}
-	if c.BackendTcpPort == 0 {
-		return fmt.Errorf("%s missing backend port: %w", errorPrefix, db.ErrInvalidParameter)
+	if c.EndpointTcpPort == 0 {
+		return fmt.Errorf("%s missing endpoint port: %w", errorPrefix, db.ErrInvalidParameter)
 	}
 	return nil
 }

internal/session/connection_closed_with.go

@@ -0,0 +1,31 @@
+package session
+
+import (
+	"fmt"
+
+	"github.com/hashicorp/boundary/internal/db"
+)
+
+// CloseWith defines the boundary data that is saved in the repo when the
+// worker closes a connection between the client and the endpoint.
+type CloseWith struct {
+	ConnectionId      string
+	ConnectionVersion uint32
+	BytesUp           uint64
+	BytesDown         uint64
+	ClosedReason      ClosedReason
+}
+
+func (c CloseWith) validate() error {
+	if c.ConnectionId == "" {
+		return fmt.Errorf("missing connection id: %w", db.ErrInvalidParameter)
+	}
+	if c.ConnectionVersion == 0 {
+		return fmt.Errorf("missing connection version: %w", db.ErrInvalidParameter)
+	}
+	if c.ClosedReason.String() == "" {
+		return fmt.Errorf("missing closed reason: %w", db.ErrInvalidParameter)
+	}
+	// 0 is valid for BytesUp and BytesDown
+	return nil
+}

internal/session/connection_closed_with_test.go

@@ -0,0 +1,83 @@
+package session
+
+import (
+	"testing"
+
+	"github.com/hashicorp/boundary/internal/db"
+	"github.com/hashicorp/boundary/internal/iam"
+)
+
+func TestClosedWith_validate(t *testing.T) {
+	conn, _ := db.TestSetup(t, "postgres")
+	wrapper := db.TestWrapper(t)
+	iamRepo := iam.TestRepo(t, conn, wrapper)
+	session := TestDefaultSession(t, conn, wrapper, iamRepo)
+	sessionConnection := TestConnection(t, conn, session.PublicId, "127.0.0.1", 22, "127.0.0.1", 2222)
+	type fields struct {
+		ConnectionId      string
+		ConnectionVersion uint32
+		BytesUp           uint64
+		BytesDown         uint64
+		ClosedReason      ClosedReason
+	}
+	tests := []struct {
+		name    string
+		fields  fields
+		wantErr bool
+	}{
+		{
+			name: "valid",
+			fields: fields{
+				ConnectionId:      sessionConnection.PublicId,
+				ConnectionVersion: sessionConnection.Version,
+				BytesUp:           1,
+				BytesDown:         2,
+				ClosedReason:      ConnectionClosedByUser,
+			},
+		},
+		{
+			name: "missing-ConnectionId",
+			fields: fields{
+				ConnectionVersion: sessionConnection.Version,
+				BytesUp:           1,
+				BytesDown:         2,
+				ClosedReason:      ConnectionClosedByUser,
+			},
+			wantErr: true,
+		},
+		{
+			name: "missing-ConnectionVersion",
+			fields: fields{
+				ConnectionId: sessionConnection.PublicId,
+				BytesUp:      1,
+				BytesDown:    2,
+				ClosedReason: ConnectionClosedByUser,
+			},
+			wantErr: true,
+		},
+		{
+			name: "missing-ClosedReason",
+			fields: fields{
+				ConnectionId:      sessionConnection.PublicId,
+				ConnectionVersion: sessionConnection.Version,
+				BytesUp:           1,
+				BytesDown:         2,
+			},
+			wantErr: true,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			c := CloseWith{
+				ConnectionId:      tt.fields.ConnectionId,
+				ConnectionVersion: tt.fields.ConnectionVersion,
+				BytesUp:           tt.fields.BytesUp,
+				BytesDown:         tt.fields.BytesDown,
+				ClosedReason:      tt.fields.ClosedReason,
+			}
+			if err := c.validate(); (err != nil) != tt.wantErr {
+				t.Errorf("ClosedWith.validate() error = %v, wantErr %v", err, tt.wantErr)
+			}
+		})
+	}
+}

internal/session/connection_test.go

@@ -19,11 +19,11 @@ func TestConnection_Create(t *testing.T) {
 	s := TestDefaultSession(t, conn, wrapper, iamRepo)
 
 	type args struct {
-		sessionId         string
-		clientTcpAddress  string
-		clientTcpPort     uint32
-		backendTcpAddress string
-		backendTcpPort    uint32
+		sessionId          string
+		clientTcpAddress   string
+		clientTcpPort      uint32
+		endpointTcpAddress string
+		endpointTcpPort    uint32
 	}
 	tests := []struct {
 		name          string
@@ -37,28 +37,28 @@ func TestConnection_Create(t *testing.T) {
 		{
 			name: "valid",
 			args: args{
-				sessionId:         s.PublicId,
-				clientTcpAddress:  "127.0.0.1",
-				clientTcpPort:     22,
-				backendTcpAddress: "127.0.0.1",
-				backendTcpPort:    2222,
+				sessionId:          s.PublicId,
+				clientTcpAddress:   "127.0.0.1",
+				clientTcpPort:      22,
+				endpointTcpAddress: "127.0.0.1",
+				endpointTcpPort:    2222,
 			},
 			want: &Connection{
-				SessionId:         s.PublicId,
-				ClientTcpAddress:  "127.0.0.1",
-				ClientTcpPort:     22,
-				BackendTcpAddress: "127.0.0.1",
-				BackendTcpPort:    2222,
+				SessionId:          s.PublicId,
+				ClientTcpAddress:   "127.0.0.1",
+				ClientTcpPort:      22,
+				EndpointTcpAddress: "127.0.0.1",
+				EndpointTcpPort:    2222,
 			},
 			create: true,
 		},
 		{
 			name: "empty-session-id",
 			args: args{
-				clientTcpAddress:  "127.0.0.1",
-				clientTcpPort:     22,
-				backendTcpAddress: "127.0.0.1",
-				backendTcpPort:    2222,
+				clientTcpAddress:   "127.0.0.1",
+				clientTcpPort:      22,
+				endpointTcpAddress: "127.0.0.1",
+				endpointTcpPort:    2222,
 			},
 			wantErr:   true,
 			wantIsErr: db.ErrInvalidParameter,
@@ -66,10 +66,10 @@ func TestConnection_Create(t *testing.T) {
 		{
 			name: "empty-client-address",
 			args: args{
-				sessionId:         s.PublicId,
-				clientTcpPort:     22,
-				backendTcpAddress: "127.0.0.1",
-				backendTcpPort:    2222,
+				sessionId:          s.PublicId,
+				clientTcpPort:      22,
+				endpointTcpAddress: "127.0.0.1",
+				endpointTcpPort:    2222,
 			},
 			wantErr:   true,
 			wantIsErr: db.ErrInvalidParameter,
@@ -77,32 +77,32 @@ func TestConnection_Create(t *testing.T) {
 		{
 			name: "empty-client-port",
 			args: args{
-				sessionId:         s.PublicId,
-				clientTcpAddress:  "localhost",
-				backendTcpAddress: "127.0.0.1",
-				backendTcpPort:    2222,
+				sessionId:          s.PublicId,
+				clientTcpAddress:   "localhost",
+				endpointTcpAddress: "127.0.0.1",
+				endpointTcpPort:    2222,
 			},
 			wantErr:   true,
 			wantIsErr: db.ErrInvalidParameter,
 		},
 		{
-			name: "empty-backend-address",
+			name: "empty-endpoint-address",
 			args: args{
 				sessionId:        s.PublicId,
 				clientTcpAddress: "localhost",
 				clientTcpPort:    22,
-				backendTcpPort:   2222,
+				endpointTcpPort:  2222,
 			},
 			wantErr:   true,
 			wantIsErr: db.ErrInvalidParameter,
 		},
 		{
-			name: "empty-backend-port",
+			name: "empty-endpoint-port",
 			args: args{
-				sessionId:         s.PublicId,
-				clientTcpAddress:  "localhost",
-				clientTcpPort:     22,
-				backendTcpAddress: "127.0.0.1",
+				sessionId:          s.PublicId,
+				clientTcpAddress:   "localhost",
+				clientTcpPort:      22,
+				endpointTcpAddress: "127.0.0.1",
 			},
 			wantErr:   true,
 			wantIsErr: db.ErrInvalidParameter,
@@ -115,8 +115,8 @@ func TestConnection_Create(t *testing.T) {
 				tt.args.sessionId,
 				tt.args.clientTcpAddress,
 				tt.args.clientTcpPort,
-				tt.args.backendTcpAddress,
-				tt.args.backendTcpPort,
+				tt.args.endpointTcpAddress,
+				tt.args.endpointTcpPort,
 			)
 			if tt.wantErr {
 				require.Error(err)

internal/session/error.go

@@ -0,0 +1,19 @@
+package session
+
+import "errors"
+
+// Errors returned from this package may be tested against these errors
+// with errors.Is.
+var (
+	// ErrInvalidStateForOperation indicates the session's state does not allow the
+	// operation.
+	ErrInvalidStateForOperation = errors.New("state is invalid for operation")
+
+	// ErrSessionNotPending indicates that a session cannot be activated
+	// because it's not in a pending state.
+	ErrSessionNotPending = errors.New("session is not in a pending state")
+
+	// ErrOpenConnection indicates that a session can not be terminated because
+	// it has open connections.
+	ErrOpenConnection = errors.New("session has open connections")
+)

internal/session/immutable_fields_test.go

@@ -213,22 +213,22 @@ func TestConnection_ImmutableFields(t *testing.T) {
 			fieldMask: []string{"ClientTcpPort"},
 		},
 		{
-			name: "backend_tcp_address",
+			name: "endpoint_tcp_address",
 			update: func() *Connection {
 				c := new.Clone().(*Connection)
-				c.BackendTcpAddress = "0.0.0.0"
+				c.EndpointTcpAddress = "0.0.0.0"
 				return c
 			}(),
-			fieldMask: []string{"BackendTcpAddress"},
+			fieldMask: []string{"EndpointTcpAddress"},
 		},
 		{
-			name: "backend_tcp_port",
+			name: "endpoint_tcp_port",
 			update: func() *Connection {
 				c := new.Clone().(*Connection)
-				c.BackendTcpPort = 443
+				c.EndpointTcpPort = 443
 				return c
 			}(),
-			fieldMask: []string{"BackendTcpPort"},
+			fieldMask: []string{"EndpointTcpPort"},
 		},
 		{
 			name: "create time",

internal/session/query.go

@@ -16,5 +16,82 @@ with not_active as (
 		s.public_id not in(select session_id from session_state where session_id = $1 and state = 'active') 
 )
 select * from not_active;
+`
+
+	terminateSessionCte = `
+insert into session_state
+with terminated as (
+	select s.public_id, 'terminated' as state
+	from 
+		session s
+	where 
+		s.version = $2  and
+		s.public_id in (
+			-- sessions without any connections
+			select s.public_id 
+			from 
+				session s
+			left join session_connection sc on sc.session_id = s.public_id 
+			where 
+				sc.session_id is null
+				and s.public_id = $1
+			union
+			-- sessions where all connections are closed
+			select s.public_id 
+			from
+				session s,
+				session_connection c,
+				session_connection_state cs
+			where
+				s.public_id = c.session_id and
+				c.public_id = cs.connection_id and
+				cs.state = 'closed' and 
+				s.public_id = $1
+		) 
+)
+select * from terminated;
+`
+	createConnectionCte = `
+insert into session_connection (
+	session_id, 
+	public_id, 
+	client_tcp_address, 
+	client_tcp_port,
+	endpoint_tcp_address,
+	endpoint_tcp_port
+)
+with active_session as ( 
+	select 
+		$1 as session_id,
+		$2 as public_id,
+		$3::inet as client_tcp_address,
+		$4::int as client_tcp_port,
+		$5::inet as endpoint_tcp_address,
+		$6::int as endpoint_tcp_port
+	from
+		session s
+	where
+		-- check that there's a state of active
+		s.public_id in (
+			select 
+				ss.session_id 
+			from 
+				session_state ss
+			where 
+				ss.session_id = $1 and 
+				ss.state = 'active'
+		) and 
+		-- check that there are no cancelling or terminated states
+		s.public_id not in(
+			select 
+				ss.session_id 
+			from 
+				session_state ss 
+			where
+				ss.session_id = $1 and 
+				ss.state in('cancelling', 'terminated') 			
+		)
+)
+select * from active_session;
 `
 )

internal/session/repository_connection.go

@@ -4,80 +4,10 @@ import (
 	"context"
 	"errors"
 	"fmt"
-	"strings"
 
 	"github.com/hashicorp/boundary/internal/db"
-	dbcommon "github.com/hashicorp/boundary/internal/db/common"
 )
 
-// CreateConnection inserts into the repository and returns the new Connection with
-// its State of "Connected".  The following fields must be empty when creating a
-// session: PublicId, BytesUp, BytesDown, ClosedReason, Version, CreateTime,
-// UpdateTime.  No options are currently supported.
-func (r *Repository) CreateConnection(ctx context.Context, newConnection *Connection, opt ...Option) (*Connection, *ConnectionState, error) {
-	if newConnection == nil {
-		return nil, nil, fmt.Errorf("create connection: missing connection: %w", db.ErrInvalidParameter)
-	}
-	if newConnection.PublicId != "" {
-		return nil, nil, fmt.Errorf("create connection: public id is not empty: %w", db.ErrInvalidParameter)
-	}
-	if newConnection.BytesUp != 0 {
-		return nil, nil, fmt.Errorf("create connection: bytes down is not empty: %w", db.ErrInvalidParameter)
-	}
-	if newConnection.BytesDown != 0 {
-		return nil, nil, fmt.Errorf("create connection: bytes up is not empty: %w", db.ErrInvalidParameter)
-	}
-	if newConnection.ClosedReason != "" {
-		return nil, nil, fmt.Errorf("create connection: closed reason is not empty: %w", db.ErrInvalidParameter)
-	}
-	if newConnection.Version != 0 {
-		return nil, nil, fmt.Errorf("create connection: version is not empty: %w", db.ErrInvalidParameter)
-	}
-	if newConnection.CreateTime != nil {
-		return nil, nil, fmt.Errorf("create connection: create time is not empty: %w", db.ErrInvalidParameter)
-	}
-	if newConnection.UpdateTime != nil {
-		return nil, nil, fmt.Errorf("create connection: update time is not empty: %w", db.ErrInvalidParameter)
-	}
-
-	id, err := newConnectionId()
-	if err != nil {
-		return nil, nil, fmt.Errorf("create connection: %w", err)
-	}
-	newConnection.PublicId = id
-
-	var returnedConnection *Connection
-	var returnedState *ConnectionState
-	_, err = r.writer.DoTx(
-		ctx,
-		db.StdRetryCnt,
-		db.ExpBackoff{},
-		func(read db.Reader, w db.Writer) error {
-			returnedConnection = newConnection.Clone().(*Connection)
-			if err = w.Create(ctx, returnedConnection); err != nil {
-				return err
-			}
-			var foundStates []*ConnectionState
-			// trigger will create new "Connected" state
-			if foundStates, err = fetchConnectionStates(ctx, read, returnedConnection.PublicId); err != nil {
-				return err
-			}
-			if len(foundStates) != 1 {
-				return fmt.Errorf("%d states found for new connection %s", len(foundStates), returnedConnection.PublicId)
-			}
-			returnedState = foundStates[0]
-			if returnedState.Status != StatusConnected.String() {
-				return fmt.Errorf("new connection %s state is not valid: %s", returnedConnection.PublicId, returnedState.Status)
-			}
-			return nil
-		},
-	)
-	if err != nil {
-		return nil, nil, fmt.Errorf("create connection: %w", err)
-	}
-	return returnedConnection, returnedState, err
-}
-
 // 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.
@@ -158,133 +88,6 @@ func (r *Repository) DeleteConnection(ctx context.Context, publicId string, opt
 	return rowsDeleted, nil
 }
 
-// UpdateConnection updates the repository entry for the connection, using the
-// fieldMaskPaths.  Only BytesUp, BytesDown, and ClosedReason are mutable and
-// will be set to NULL if set to a zero value and included in the fieldMaskPaths.
-func (r *Repository) UpdateConnection(ctx context.Context, connection *Connection, version uint32, fieldMaskPaths []string, opt ...Option) (*Connection, []*ConnectionState, int, error) {
-	if connection == nil {
-		return nil, nil, db.NoRowsAffected, fmt.Errorf("update connection: missing connection %w", db.ErrInvalidParameter)
-	}
-	if connection.PublicId == "" {
-		return nil, nil, db.NoRowsAffected, fmt.Errorf("update connection: missing connection public id %w", db.ErrInvalidParameter)
-	}
-	for _, f := range fieldMaskPaths {
-		switch {
-		case strings.EqualFold("BytesUp", f):
-		case strings.EqualFold("BytesDown", f):
-		case strings.EqualFold("ClosedReason", f):
-		default:
-			return nil, nil, db.NoRowsAffected, fmt.Errorf("update connection: field: %s: %w", f, db.ErrInvalidFieldMask)
-		}
-	}
-	var dbMask, nullFields []string
-	dbMask, nullFields = dbcommon.BuildUpdatePaths(
-		map[string]interface{}{
-			"BytesUp":      connection.BytesUp,
-			"BytesDown":    connection.BytesDown,
-			"ClosedReason": connection.ClosedReason,
-		},
-		fieldMaskPaths,
-	)
-	if len(dbMask) == 0 && len(nullFields) == 0 {
-		return nil, nil, db.NoRowsAffected, fmt.Errorf("update connection: %w", db.ErrEmptyFieldMask)
-	}
-
-	var c *Connection
-	var states []*ConnectionState
-	var rowsUpdated int
-	_, err := r.writer.DoTx(
-		ctx,
-		db.StdRetryCnt,
-		db.ExpBackoff{},
-		func(reader db.Reader, w db.Writer) error {
-			var err error
-			c = connection.Clone().(*Connection)
-			rowsUpdated, err = w.Update(
-				ctx,
-				c,
-				dbMask,
-				nullFields,
-			)
-			if err != nil {
-				return err
-			}
-			if err == nil && rowsUpdated > 1 {
-				// return err, which will result in a rollback of the update
-				return errors.New("error more than 1 connection would have been updated ")
-			}
-			states, err = fetchConnectionStates(ctx, reader, c.PublicId, db.WithOrder("start_time desc"))
-			if err != nil {
-				return err
-			}
-			return nil
-		},
-	)
-	if err != nil {
-		return nil, nil, db.NoRowsAffected, fmt.Errorf("update connection: %w for %s", err, connection.PublicId)
-	}
-	return c, states, rowsUpdated, err
-}
-
-// UpdateConnectionState will update the connection's state using the connection id and its
-// version.  No options are currently supported.
-func (r *Repository) UpdateConnectionState(ctx context.Context, connectionId string, connectionVersion uint32, s ConnectionStatus, opt ...Option) (*Connection, []*ConnectionState, error) {
-	if connectionId == "" {
-		return nil, nil, fmt.Errorf("update connection state: missing session id %w", db.ErrInvalidParameter)
-	}
-	if connectionVersion == 0 {
-		return nil, nil, fmt.Errorf("update connection state: version cannot be zero: %w", db.ErrInvalidParameter)
-	}
-	if s == "" {
-		return nil, nil, fmt.Errorf("update connection state: missing connection status: %w", db.ErrInvalidParameter)
-	}
-
-	newState, err := NewConnectionState(connectionId, s)
-	if err != nil {
-		return nil, nil, fmt.Errorf("update connection state: %w", err)
-	}
-	sessionConnection, _, err := r.LookupConnection(ctx, connectionId)
-	if err != nil {
-		return nil, nil, fmt.Errorf("update connection state: %w", err)
-	}
-	if sessionConnection == nil {
-		return nil, nil, fmt.Errorf("update connection state: unable to look up connection for %s: %w", connectionId, err)
-	}
-
-	updatedConnection := AllocConnection()
-	var returnedStates []*ConnectionState
-	_, err = r.writer.DoTx(
-		ctx,
-		db.StdRetryCnt,
-		db.ExpBackoff{},
-		func(reader db.Reader, w db.Writer) error {
-			// We need to update the connection version as that's the aggregate
-			updatedConnection.PublicId = connectionId
-			updatedConnection.Version = uint32(connectionVersion) + 1
-			rowsUpdated, err := w.Update(ctx, &updatedConnection, []string{"Version"}, nil, db.WithVersion(&connectionVersion))
-			if err != nil {
-				return fmt.Errorf("unable to update connection version: %w", err)
-			}
-			if rowsUpdated != 1 {
-				return fmt.Errorf("updated connection and %d rows updated", rowsUpdated)
-			}
-			if err := w.Create(ctx, newState); err != nil {
-				return fmt.Errorf("unable to add new state: %w", err)
-			}
-
-			returnedStates, err = fetchConnectionStates(ctx, reader, connectionId, db.WithOrder("start_time desc"))
-			if err != nil {
-				return err
-			}
-			return nil
-		},
-	)
-	if err != nil {
-		return nil, nil, fmt.Errorf("update connection state: error creating new state: %w", err)
-	}
-	return &updatedConnection, returnedStates, nil
-}
-
 func fetchConnectionStates(ctx context.Context, r db.Reader, connectionId string, opt ...db.Option) ([]*ConnectionState, error) {
 	var states []*ConnectionState
 	if err := r.SearchWhere(ctx, &states, "connection_id = ?", []interface{}{connectionId}, opt...); err != nil {

internal/session/repository_connection_test.go

@@ -8,7 +8,6 @@ import (
 
 	"github.com/golang/protobuf/ptypes"
 	"github.com/hashicorp/boundary/internal/db"
-	dbassert "github.com/hashicorp/boundary/internal/db/assert"
 	"github.com/hashicorp/boundary/internal/iam"
 	"github.com/hashicorp/boundary/internal/kms"
 	"github.com/hashicorp/boundary/internal/oplog"
@@ -130,431 +129,6 @@ func TestRepository_ListConnection(t *testing.T) {
 	})
 }
 
-func TestRepository_CreateConnection(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)
-	session := TestDefaultSession(t, conn, wrapper, iamRepo)
-
-	type args struct {
-		connection *Connection
-	}
-	tests := []struct {
-		name        string
-		args        args
-		wantErr     bool
-		wantIsError error
-	}{
-		{
-			name: "valid",
-			args: args{
-				connection: func() *Connection {
-					c, err := NewConnection(
-						session.PublicId,
-						"127.0.0.1",
-						22,
-						"127.0.0.1",
-						2222,
-					)
-					require.NoError(t, err)
-					return c
-				}(),
-			},
-			wantErr: false,
-		},
-		{
-			name: "empty-session-id",
-			args: args{
-				connection: &Connection{
-					ClientTcpAddress:  "127.0.0.1",
-					ClientTcpPort:     22,
-					BackendTcpAddress: "127.0.0.1",
-					BackendTcpPort:    2222,
-				},
-			},
-			wantErr:     true,
-			wantIsError: db.ErrInvalidParameter,
-		},
-		{
-			name: "empty-client-address",
-			args: args{
-				connection: &Connection{
-					SessionId:         session.PublicId,
-					ClientTcpPort:     22,
-					BackendTcpAddress: "127.0.0.1",
-					BackendTcpPort:    2222,
-				},
-			},
-			wantErr:     true,
-			wantIsError: db.ErrInvalidParameter,
-		},
-		{
-			name: "empty-client-port",
-			args: args{
-				connection: &Connection{
-					SessionId:         session.PublicId,
-					ClientTcpAddress:  "127.0.0.1",
-					BackendTcpAddress: "127.0.0.1",
-					BackendTcpPort:    2222,
-				},
-			},
-			wantErr:     true,
-			wantIsError: db.ErrInvalidParameter,
-		},
-		{
-			name: "empty-backend-address",
-			args: args{
-				connection: &Connection{
-					SessionId:        session.PublicId,
-					ClientTcpAddress: "127.0.0.1",
-					ClientTcpPort:    22,
-					BackendTcpPort:   2222,
-				},
-			},
-			wantErr:     true,
-			wantIsError: db.ErrInvalidParameter,
-		},
-		{
-			name: "empty-backend-port",
-			args: args{
-				connection: &Connection{
-					SessionId:         session.PublicId,
-					ClientTcpAddress:  "127.0.0.1",
-					ClientTcpPort:     22,
-					BackendTcpAddress: "127.0.0.1",
-				},
-			},
-			wantErr:     true,
-			wantIsError: db.ErrInvalidParameter,
-		},
-	}
-	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
-			assert, require := assert.New(t), require.New(t)
-			connection, st, err := repo.CreateConnection(context.Background(), tt.args.connection)
-			if tt.wantErr {
-				assert.Error(err)
-				assert.Nil(connection)
-				assert.Nil(st)
-				if tt.wantIsError != nil {
-					assert.True(errors.Is(err, tt.wantIsError))
-				}
-				return
-			}
-			require.NoError(err)
-			assert.NotNil(connection.CreateTime)
-			assert.NotNil(st.StartTime)
-			assert.Equal(st.Status, StatusConnected.String())
-			found, foundStates, err := repo.LookupConnection(context.Background(), connection.PublicId)
-			assert.NoError(err)
-			assert.Equal(found, connection)
-
-			err = db.TestVerifyOplog(t, rw, connection.PublicId, db.WithOperation(oplog.OpType_OP_TYPE_CREATE), db.WithCreateNotBefore(10*time.Second))
-			assert.Error(err)
-
-			require.Equal(1, len(foundStates))
-			assert.Equal(foundStates[0].Status, StatusConnected.String())
-		})
-	}
-}
-
-func TestRepository_UpdateConnectionState(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)
-	session := TestDefaultSession(t, conn, wrapper, iamRepo)
-
-	tests := []struct {
-		name                      string
-		connection                *Connection
-		newStatus                 ConnectionStatus
-		overrideConnectionId      *string
-		overrideConnectionVersion *uint32
-		wantStateCnt              int
-		wantErr                   bool
-		wantIsError               error
-	}{
-		{
-			name:         "closed",
-			connection:   TestConnection(t, conn, session.PublicId, "0.0.0.0", 22, "0.0.0.0", 2222),
-			newStatus:    StatusClosed,
-			wantStateCnt: 2,
-			wantErr:      false,
-		},
-		{
-			name:       "bad-version",
-			connection: TestConnection(t, conn, session.PublicId, "0.0.0.0", 22, "0.0.0.0", 2222),
-			newStatus:  StatusClosed,
-			overrideConnectionVersion: func() *uint32 {
-				v := uint32(22)
-				return &v
-			}(),
-			wantErr: true,
-		},
-		{
-			name:       "empty-version",
-			connection: TestConnection(t, conn, session.PublicId, "0.0.0.0", 22, "0.0.0.0", 2222),
-			newStatus:  StatusClosed,
-			overrideConnectionVersion: func() *uint32 {
-				v := uint32(0)
-				return &v
-			}(),
-			wantErr:     true,
-			wantIsError: db.ErrInvalidParameter,
-		},
-		{
-			name:       "bad-connectionId",
-			connection: TestConnection(t, conn, session.PublicId, "0.0.0.0", 22, "0.0.0.0", 2222),
-			newStatus:  StatusClosed,
-			overrideConnectionId: func() *string {
-				s := "sc_thisIsNotValid"
-				return &s
-			}(),
-			wantErr: true,
-		},
-		{
-			name:       "empty-connectionId",
-			connection: TestConnection(t, conn, session.PublicId, "0.0.0.0", 22, "0.0.0.0", 2222),
-			newStatus:  StatusClosed,
-			overrideConnectionId: func() *string {
-				s := ""
-				return &s
-			}(),
-			wantErr:     true,
-			wantIsError: db.ErrInvalidParameter,
-		},
-	}
-	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
-			assert, require := assert.New(t), require.New(t)
-			var id string
-			var version uint32
-			switch {
-			case tt.overrideConnectionId != nil:
-				id = *tt.overrideConnectionId
-			default:
-				id = tt.connection.PublicId
-			}
-			switch {
-			case tt.overrideConnectionVersion != nil:
-				version = *tt.overrideConnectionVersion
-			default:
-				version = tt.connection.Version
-			}
-
-			s, ss, err := repo.UpdateConnectionState(context.Background(), id, version, tt.newStatus)
-			if tt.wantErr {
-				require.Error(err)
-				if tt.wantIsError != nil {
-					assert.Truef(errors.Is(err, tt.wantIsError), "unexpected error %s", err.Error())
-				}
-				return
-			}
-			require.NoError(err)
-			require.NotNil(s)
-			require.NotNil(ss)
-			assert.Equal(tt.wantStateCnt, len(ss))
-			assert.Equal(tt.newStatus.String(), ss[0].Status)
-		})
-	}
-}
-
-func TestRepository_UpdateConnection(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)
-	session := TestDefaultSession(t, conn, wrapper, iamRepo)
-
-	type args struct {
-		closedReason      ClosedReason
-		bytesUp           uint64
-		bytesDown         uint64
-		fieldMaskPaths    []string
-		opt               []Option
-		publicId          *string // not updateable - db.ErrInvalidFieldMask
-		sessionId         string  // not updateable - db.ErrInvalidFieldMask
-		clientTcpAddress  string  // not updateable - db.ErrInvalidFieldMask
-		clientTcpPort     uint32  // not updateable - db.ErrInvalidFieldMask
-		backendTcpAddress string  // not updateable - db.ErrInvalidFieldMask
-		backendTcpPort    uint32  // not updateable - db.ErrInvalidFieldMask
-	}
-	tests := []struct {
-		name           string
-		args           args
-		wantRowsUpdate int
-		wantErr        bool
-		wantIsError    error
-	}{
-		{
-			name: "valid",
-			args: args{
-				closedReason:   ConnectionClosedByUser,
-				bytesUp:        uint64(111),
-				bytesDown:      uint64(1),
-				fieldMaskPaths: []string{"ClosedReason", "BytesUp", "BytesDown"},
-			},
-			wantErr:        false,
-			wantRowsUpdate: 1,
-		},
-		{
-			name: "publicId",
-			args: args{
-				publicId: func() *string {
-					id, err := newConnectionId()
-					require.NoError(t, err)
-					return &id
-				}(),
-				fieldMaskPaths: []string{"PublicId"},
-			},
-			wantErr:        true,
-			wantRowsUpdate: 0,
-			wantIsError:    db.ErrInvalidFieldMask,
-		},
-		{
-			name: "sessionId",
-			args: args{
-				sessionId: func() string {
-					id, err := newId()
-					require.NoError(t, err)
-					return id
-				}(),
-				fieldMaskPaths: []string{"SessionId"},
-			},
-			wantErr:        true,
-			wantRowsUpdate: 0,
-			wantIsError:    db.ErrInvalidFieldMask,
-		},
-		{
-			name: "clientTcpAddress",
-			args: args{
-				clientTcpAddress: "127.0.0.1",
-				fieldMaskPaths:   []string{"ClientTcpAddress"},
-			},
-			wantErr:        true,
-			wantRowsUpdate: 0,
-			wantIsError:    db.ErrInvalidFieldMask,
-		},
-		{
-			name: "clientTcpPort",
-			args: args{
-				clientTcpPort:  443,
-				fieldMaskPaths: []string{"ClientTcpPort"},
-			},
-			wantErr:        true,
-			wantRowsUpdate: 0,
-			wantIsError:    db.ErrInvalidFieldMask,
-		},
-		{
-			name: "backendTcpAddress",
-			args: args{
-				backendTcpAddress: "127.0.0.1",
-				fieldMaskPaths:    []string{"BackendTcpAddress"},
-			},
-			wantErr:        true,
-			wantRowsUpdate: 0,
-			wantIsError:    db.ErrInvalidFieldMask,
-		},
-		{
-			name: "backendTcpPort",
-			args: args{
-				backendTcpPort: 4443,
-				fieldMaskPaths: []string{"BackendTcpPort"},
-			},
-			wantErr:        true,
-			wantRowsUpdate: 0,
-			wantIsError:    db.ErrInvalidFieldMask,
-		},
-	}
-	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
-			assert, require := assert.New(t), require.New(t)
-
-			c := TestConnection(t, conn, session.PublicId, "0.0.0.0", 22, "127.0.0.1", 2222)
-
-			updateConnection := AllocConnection()
-			updateConnection.PublicId = c.PublicId
-			if tt.args.publicId != nil {
-				updateConnection.PublicId = *tt.args.publicId
-			}
-			updateConnection.BytesUp = tt.args.bytesUp
-			updateConnection.BytesDown = tt.args.bytesDown
-			updateConnection.ClosedReason = tt.args.closedReason.String()
-			updateConnection.Version = c.Version
-			afterUpdate, afterUpdateState, updatedRows, err := repo.UpdateConnection(context.Background(), &updateConnection, updateConnection.Version, tt.args.fieldMaskPaths, tt.args.opt...)
-
-			if tt.wantErr {
-				require.Error(err)
-				if tt.wantIsError != nil {
-					assert.Truef(errors.Is(err, tt.wantIsError), "unexpected error: %s", err.Error())
-				}
-				assert.Nil(afterUpdate)
-				assert.Nil(afterUpdateState)
-				assert.Equal(0, updatedRows)
-				err = db.TestVerifyOplog(t, rw, c.PublicId, db.WithOperation(oplog.OpType_OP_TYPE_UPDATE), db.WithCreateNotBefore(10*time.Second))
-				assert.Error(err)
-				assert.True(errors.Is(db.ErrRecordNotFound, err))
-				return
-			}
-			require.NoError(err)
-			assert.Equal(tt.wantRowsUpdate, updatedRows)
-			require.NotNil(afterUpdate)
-			require.NotNil(afterUpdateState)
-			switch tt.name {
-			case "valid-no-op":
-				assert.Equal(c.UpdateTime, afterUpdate.UpdateTime)
-			default:
-				assert.NotEqual(c.UpdateTime, afterUpdate.UpdateTime)
-			}
-			found, foundStates, err := repo.LookupConnection(context.Background(), c.PublicId)
-			require.NoError(err)
-			assert.Equal(afterUpdate, found)
-			dbassrt := dbassert.New(t, rw)
-			if tt.args.bytesUp == 0 {
-				dbassrt.IsNull(found, "BytesUp")
-			}
-			if tt.args.bytesDown == 0 {
-				dbassrt.IsNull(found, "BytesDown")
-			}
-			if tt.args.closedReason == "" {
-				dbassrt.IsNull(found, "ClosedReason")
-			}
-			assert.Equal(tt.args.closedReason.String(), found.ClosedReason)
-			assert.Equal(tt.args.bytesUp, found.BytesUp)
-			assert.Equal(tt.args.bytesDown, found.BytesDown)
-
-			err = db.TestVerifyOplog(t, rw, c.PublicId, db.WithOperation(oplog.OpType_OP_TYPE_UPDATE), db.WithCreateNotBefore(10*time.Second))
-			assert.Error(err)
-
-			switch {
-			case tt.args.closedReason != "":
-				require.Equal(2, len(foundStates))
-				// sorted by StartTime desc
-				assert.Equal(StatusClosed.String(), foundStates[0].Status)
-				assert.Equal(StatusConnected.String(), foundStates[1].Status)
-			default:
-				require.Equal(1, len(foundStates))
-				assert.Equal(StatusConnected.String(), foundStates[0].Status)
-			}
-		})
-	}
-
-}
-
 func TestRepository_DeleteConnection(t *testing.T) {
 	t.Parallel()
 	conn, _ := db.TestSetup(t, "postgres")

internal/session/repository_session.go

@@ -8,10 +8,8 @@ import (
 	"strings"
 
 	"github.com/hashicorp/boundary/internal/db"
-	dbcommon "github.com/hashicorp/boundary/internal/db/common"
 	"github.com/hashicorp/boundary/internal/kms"
 	wrapping "github.com/hashicorp/go-kms-wrapping"
-	"github.com/hashicorp/vault/sdk/helper/strutil"
 )
 
 // CreateSession inserts into the repository and returns the new Session with
@@ -210,90 +208,61 @@ func (r *Repository) DeleteSession(ctx context.Context, publicId string, opt ...
 	return rowsDeleted, nil
 }
 
-// UpdateSession updates the repository entry for the session, using the
-// fieldMaskPaths.  Only TerminationReason, ServerId and ServerType a muttable
-// and will be set to NULL if set to a zero value and included in the
-// fieldMaskPaths. Returned States are ordered by start time descending.
-func (r *Repository) UpdateSession(ctx context.Context, session *Session, version uint32, fieldMaskPaths []string, opt ...Option) (*Session, []*State, int, error) {
-	if session == nil {
-		return nil, nil, db.NoRowsAffected, fmt.Errorf("update session: missing session %w", db.ErrInvalidParameter)
+// CancelSession sets a session's state to "cancelling" in the repo.  It's
+// called when the user cancels a session and the controller wants to update the
+// session state to "cancelling" for the given reason, so the workers can get
+// the "cancelling signal" during their next status heartbeat.
+func (r *Repository) CancelSession(ctx context.Context, sessionId string, sessionVersion uint32) (*Session, []*State, error) {
+	if sessionId == "" {
+		return nil, nil, fmt.Errorf("cancel session: missing session id: %w", db.ErrInvalidParameter)
 	}
-	if session.PublicId == "" {
-		return nil, nil, db.NoRowsAffected, fmt.Errorf("update session: missing session public id %w", db.ErrInvalidParameter)
+	if sessionVersion == 0 {
+		return nil, nil, fmt.Errorf("cancel session: missing session version: %w", db.ErrInvalidParameter)
 	}
-	if session.CtTofuToken != nil {
-		return nil, nil, db.NoRowsAffected, fmt.Errorf("update session: ct must be empty: %w", db.ErrInvalidParameter)
+	s, ss, err := r.updateState(ctx, sessionId, sessionVersion, StatusCancelling)
+	if err != nil {
+		return nil, nil, fmt.Errorf("cancel session: %w", err)
+	}
+	return s, ss, nil
+}
+
+// TerminateSession sets a session's state to "terminated" in the repo.  It's
+// called by the worker when the session has been terminated or by a controller
+// when all of a session's workers have stopped sending heartbeat status for a
+// period of time.  Sessions cannot be terminated which still have connections
+// that are not closed.
+func (r *Repository) TerminateSession(ctx context.Context, sessionId string, sessionVersion uint32, reason TerminationReason) (*Session, []*State, error) {
+	if sessionId == "" {
+		return nil, nil, fmt.Errorf("terminate session: missing session id: %w", db.ErrInvalidParameter)
+	}
+	if sessionVersion == 0 {
+		return nil, nil, fmt.Errorf("terminate session: version cannot be zero: %w", db.ErrInvalidParameter)
 	}
 
-	translatedFieldMasks := make([]string, 0, len(fieldMaskPaths))
-	for _, f := range fieldMaskPaths {
-		switch {
-		case strings.EqualFold("TerminationReason", f):
-			translatedFieldMasks = append(translatedFieldMasks, f)
-		case strings.EqualFold("ServerId", f):
-			translatedFieldMasks = append(translatedFieldMasks, f)
-		case strings.EqualFold("ServerType", f):
-			translatedFieldMasks = append(translatedFieldMasks, f)
-		case strings.EqualFold("TofuToken", f):
-			translatedFieldMasks = append(translatedFieldMasks, "CtTofuToken")
-		default:
-			return nil, nil, db.NoRowsAffected, fmt.Errorf("update session: field: %s: %w", f, db.ErrInvalidFieldMask)
-		}
-	}
-
-	updateSession := session.Clone().(*Session)
-	if strutil.StrListContains(translatedFieldMasks, "CtTofuToken") && len(updateSession.TofuToken) != 0 {
-		if err := r.reader.LookupById(ctx, updateSession); err != nil {
-			return nil, nil, db.NoRowsAffected, fmt.Errorf("update session: %w for %s", err, updateSession.PublicId)
-		}
-		databaseWrapper, err := r.kms.GetWrapper(ctx, updateSession.ScopeId, kms.KeyPurposeDatabase)
-		if err != nil {
-			return nil, nil, db.NoRowsAffected, fmt.Errorf("update session: unable to get database wrapper: %w", err)
-		}
-		if err := updateSession.encrypt(ctx, databaseWrapper); err != nil {
-			return nil, nil, db.NoRowsAffected, fmt.Errorf("create session: %w", err)
-		}
-	}
-
-	var dbMask, nullFields []string
-	dbMask, nullFields = dbcommon.BuildUpdatePaths(
-		map[string]interface{}{
-			"TerminationReason": updateSession.TerminationReason,
-			"ServerId":          updateSession.ServerId,
-			"ServerType":        updateSession.ServerType,
-			"CtTofuToken":       updateSession.CtTofuToken,
-		},
-		translatedFieldMasks,
-	)
-	if len(dbMask) == 0 && len(nullFields) == 0 {
-		return nil, nil, db.NoRowsAffected, fmt.Errorf("update session: %w", db.ErrEmptyFieldMask)
-	}
-
-	var s *Session
-	var states []*State
-	var rowsUpdated int
+	updatedSession := AllocSession()
+	updatedSession.PublicId = sessionId
+	updatedSession.TerminationReason = reason.String()
+	var returnedStates []*State
 	_, err := r.writer.DoTx(
 		ctx,
 		db.StdRetryCnt,
 		db.ExpBackoff{},
 		func(reader db.Reader, w db.Writer) error {
-			var err error
-			s = updateSession.Clone().(*Session)
-
-			rowsUpdated, err = w.Update(
-				ctx,
-				s,
-				dbMask,
-				nullFields,
-			)
+			rowsAffected, err := w.Exec(terminateSessionCte, []interface{}{sessionId, sessionVersion})
 			if err != nil {
-				return err
+				return fmt.Errorf("unable to terminate session %s: %w", sessionId, err)
 			}
-			if err == nil && rowsUpdated > 1 {
-				// return err, which will result in a rollback of the update
-				return errors.New("error more than 1 session would have been updated ")
+			if rowsAffected == 0 {
+				return fmt.Errorf("unable to terminate session %s", sessionId)
 			}
-			states, err = fetchStates(ctx, reader, s.PublicId, db.WithOrder("start_time desc"))
+			rowsUpdated, err := w.Update(ctx, &updatedSession, []string{"TerminationReason"}, nil, db.WithVersion(&sessionVersion))
+			if err != nil {
+				return fmt.Errorf("update session: failed %w for %s", err, sessionId)
+			}
+			if rowsUpdated != 1 {
+				return fmt.Errorf("update to session %s would have updated %d session", updatedSession.PublicId, rowsUpdated)
+			}
+			returnedStates, err = fetchStates(ctx, reader, sessionId, db.WithOrder("start_time desc"))
 			if err != nil {
 				return err
 			}
@@ -301,24 +270,138 @@ func (r *Repository) UpdateSession(ctx context.Context, session *Session, versio
 		},
 	)
 	if err != nil {
-		return nil, nil, db.NoRowsAffected, fmt.Errorf("update session: %w for %s", err, session.PublicId)
+		return nil, nil, fmt.Errorf("terminate session: %w", err)
 	}
-	if len(s.CtTofuToken) == 0 {
-		s.CtTofuToken = nil
+	return &updatedSession, returnedStates, nil
+}
+
+// ConnectSession creates a connection in the repo with a state of "connected".
+// Returns an ErrCancelledOrTerminatedSession error if a connection cannot be made
+// because the session was cancelled or terminated.
+func (r *Repository) ConnectSession(ctx context.Context, c ConnectWith) (*Connection, []*ConnectionState, error) {
+	// ConnectWith.validate will check all the fields...
+	if err := c.validate(); err != nil {
+		return nil, nil, fmt.Errorf("connect session: %w", err)
 	}
-	return s, states, rowsUpdated, err
+	connectionId, err := newConnectionId()
+	if err != nil {
+		return nil, nil, fmt.Errorf("connect session: %w", err)
+	}
+
+	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(createConnectionCte, []interface{}{c.SessionId, connectionId, c.ClientTcpAddress, c.ClientTcpPort, c.EndpointTcpAddress, c.EndpointTcpPort})
+			if err != nil {
+				return fmt.Errorf("unable to connect session %s: %w", c.SessionId, err)
+			}
+			if rowsAffected == 0 {
+				return fmt.Errorf("session %s is not active: %w", c.SessionId, ErrInvalidStateForOperation)
+			}
+			if err := reader.LookupById(ctx, &connection); err != nil {
+				return fmt.Errorf("lookup session: failed %w for %s", err, c.SessionId)
+			}
+			connectionStates, err = fetchConnectionStates(ctx, reader, connectionId, db.WithOrder("start_time desc"))
+			if err != nil {
+				return err
+			}
+			return nil
+		},
+	)
+	if err != nil {
+		return nil, nil, fmt.Errorf("connect session: %w", err)
+	}
+	return &connection, connectionStates, nil
+}
+
+// CloseConnectionRep 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, opt ...Option) ([]CloseConnectionResp, error) {
+	if len(closeWith) == 0 {
+		return nil, fmt.Errorf("close connections: missing connections to close: %w", db.ErrInvalidParameter)
+	}
+	for _, cw := range closeWith {
+		if err := cw.validate(); err != nil {
+			return nil, fmt.Errorf("close connections: %s was invalid: %w", cw.ConnectionId, err)
+		}
+	}
+	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,
+					db.WithVersion(&cw.ConnectionVersion),
+				)
+				if err != nil {
+					return fmt.Errorf("unable to update connection %s: %w", cw.ConnectionId, err)
+				}
+				if rowsUpdated != 1 {
+					return fmt.Errorf("%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 err
+				}
+				resp = append(resp, CloseConnectionResp{
+					Connection:       &updateConnection,
+					ConnectionStates: states,
+				})
+
+			}
+			return nil
+		},
+	)
+	if err != nil {
+		return nil, fmt.Errorf("close connections: %w", err)
+	}
+	return resp, nil
 }
 
 // ActivateSession will activate the session and is called by a worker after
-// authenticating the session. States are ordered by start time descending.
-func (r *Repository) ActivateSession(ctx context.Context, sessionId string, sessionVersion uint32, tofuToken []byte) (*Session, []*State, error) {
+// authenticating the session. The session must be in a "pending" state to be
+// activated. States are ordered by start time descending. Returns an
+// ErrSessionNotPending error if a connection cannot be made because the session
+// was cancelled or terminated.
+func (r *Repository) ActivateSession(ctx context.Context, sessionId string, sessionVersion uint32, serverId, serverType string, tofuToken []byte) (*Session, []*State, error) {
 	if sessionId == "" {
-		return nil, nil, fmt.Errorf("activate session state: missing session id %w", db.ErrInvalidParameter)
+		return nil, nil, fmt.Errorf("activate session: missing session id %w", db.ErrInvalidParameter)
 	}
 	if sessionVersion == 0 {
-		return nil, nil, fmt.Errorf("activate session state: version cannot be zero: %w", db.ErrInvalidParameter)
+		return nil, nil, fmt.Errorf("activate session: version cannot be zero: %w", db.ErrInvalidParameter)
+	}
+	if serverId == "" {
+		return nil, nil, fmt.Errorf("activate session: missing server id: %w", db.ErrInvalidParameter)
+	}
+	if serverType == "" {
+		return nil, nil, fmt.Errorf("activate session: missing server type: %w", db.ErrInvalidParameter)
 	}
-
 	updatedSession := AllocSession()
 	updatedSession.PublicId = sessionId
 	var returnedStates []*State
@@ -332,19 +415,20 @@ func (r *Repository) ActivateSession(ctx context.Context, sessionId string, sess
 				return fmt.Errorf("unable to activate session %s: %w", sessionId, err)
 			}
 			if rowsAffected == 0 {
-				return fmt.Errorf("unable to activate session %s", sessionId)
+				return fmt.Errorf("unable to activate session %s: %w", sessionId, ErrSessionNotPending)
 			}
 			foundSession := AllocSession()
 			foundSession.PublicId = sessionId
-			if err := r.reader.LookupById(ctx, &foundSession); err != nil {
+			if err := reader.LookupById(ctx, &foundSession); err != nil {
 				return fmt.Errorf("lookup session: failed %w for %s", err, sessionId)
 			}
 			databaseWrapper, err := r.kms.GetWrapper(ctx, foundSession.ScopeId, kms.KeyPurposeDatabase)
 			if err != nil {
 				return fmt.Errorf("unable to get database wrapper: %w", err)
 			}
-
 			updatedSession.TofuToken = tofuToken
+			updatedSession.ServerId = serverId
+			updatedSession.ServerType = serverType
 			if err := updatedSession.encrypt(ctx, databaseWrapper); err != nil {
 				return err
 			}
@@ -370,10 +454,10 @@ func (r *Repository) ActivateSession(ctx context.Context, sessionId string, sess
 	return &updatedSession, returnedStates, nil
 }
 
-// UpdateState will update the session's state using the session id and its
+// updateState will update the session's state using the session id and its
 // version.  States are ordered by start time descending. No options are
 // currently supported.
-func (r *Repository) UpdateState(ctx context.Context, sessionId string, sessionVersion uint32, s Status, opt ...Option) (*Session, []*State, error) {
+func (r *Repository) updateState(ctx context.Context, sessionId string, sessionVersion uint32, s Status, opt ...Option) (*Session, []*State, error) {
 	if sessionId == "" {
 		return nil, nil, fmt.Errorf("update session state: missing session id %w", db.ErrInvalidParameter)
 	}

internal/session/repository_session_test.go

@@ -7,17 +7,10 @@ import (
 	"time"
 
 	"github.com/golang/protobuf/ptypes"
-	"github.com/hashicorp/boundary/internal/auth/password"
-	"github.com/hashicorp/boundary/internal/authtoken"
 	"github.com/hashicorp/boundary/internal/db"
-	dbassert "github.com/hashicorp/boundary/internal/db/assert"
-	"github.com/hashicorp/boundary/internal/host/static"
 	"github.com/hashicorp/boundary/internal/iam"
 	"github.com/hashicorp/boundary/internal/kms"
 	"github.com/hashicorp/boundary/internal/oplog"
-	"github.com/hashicorp/boundary/internal/servers"
-	"github.com/hashicorp/boundary/internal/target"
-	"github.com/hashicorp/go-uuid"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
@@ -304,7 +297,7 @@ func TestRepository_UpdateState(t *testing.T) {
 		{
 			name:         "cancelling",
 			session:      TestDefaultSession(t, conn, wrapper, iamRepo),
-			newStatus:    StatusCanceling,
+			newStatus:    StatusCancelling,
 			wantStateCnt: 2,
 			wantErr:      false,
 		},
@@ -322,7 +315,7 @@ func TestRepository_UpdateState(t *testing.T) {
 		{
 			name:      "bad-version",
 			session:   TestDefaultSession(t, conn, wrapper, iamRepo),
-			newStatus: StatusCanceling,
+			newStatus: StatusCancelling,
 			overrideSessionVersion: func() *uint32 {
 				v := uint32(22)
 				return &v
@@ -332,7 +325,7 @@ func TestRepository_UpdateState(t *testing.T) {
 		{
 			name:      "empty-version",
 			session:   TestDefaultSession(t, conn, wrapper, iamRepo),
-			newStatus: StatusCanceling,
+			newStatus: StatusCancelling,
 			overrideSessionVersion: func() *uint32 {
 				v := uint32(0)
 				return &v
@@ -343,7 +336,7 @@ func TestRepository_UpdateState(t *testing.T) {
 		{
 			name:      "bad-sessionId",
 			session:   TestDefaultSession(t, conn, wrapper, iamRepo),
-			newStatus: StatusCanceling,
+			newStatus: StatusCancelling,
 			overrideSessionId: func() *string {
 				s := "s_thisIsNotValid"
 				return &s
@@ -353,7 +346,7 @@ func TestRepository_UpdateState(t *testing.T) {
 		{
 			name:      "empty-session",
 			session:   TestDefaultSession(t, conn, wrapper, iamRepo),
-			newStatus: StatusCanceling,
+			newStatus: StatusCancelling,
 			overrideSessionId: func() *string {
 				s := ""
 				return &s
@@ -380,7 +373,7 @@ func TestRepository_UpdateState(t *testing.T) {
 				version = tt.session.Version
 			}
 
-			s, ss, err := repo.UpdateState(context.Background(), id, version, tt.newStatus)
+			s, ss, err := repo.updateState(context.Background(), id, version, tt.newStatus)
 			if tt.wantErr {
 				require.Error(err)
 				if tt.wantIsError != nil {
@@ -397,7 +390,7 @@ func TestRepository_UpdateState(t *testing.T) {
 	}
 }
 
-func TestRepository_ActivateState(t *testing.T) {
+func TestRepository_ConnectSession(t *testing.T) {
 	t.Parallel()
 	conn, _ := db.TestSetup(t, "postgres")
 	rw := db.New(conn)
@@ -407,6 +400,394 @@ func TestRepository_ActivateState(t *testing.T) {
 	repo, err := NewRepository(rw, rw, kms)
 	require.NoError(t, err)
 
+	setupFn := func() ConnectWith {
+		s := TestDefaultSession(t, conn, wrapper, iamRepo)
+		srv := TestWorker(t, conn, wrapper)
+		tofu := TestTofu(t)
+		_, _, err := repo.ActivateSession(context.Background(), s.PublicId, s.Version, srv.PrivateId, srv.Type, tofu)
+		require.NoError(t, err)
+		_ = TestConnection(t, conn, s.PublicId, "127.0.0.1", 22, "127.0.0.1", 2222)
+		return ConnectWith{
+			SessionId:          s.PublicId,
+			ClientTcpAddress:   "127.0.0.1",
+			ClientTcpPort:      22,
+			EndpointTcpAddress: "127.0.0.1",
+			EndpointTcpPort:    2222,
+		}
+	}
+	tests := []struct {
+		name        string
+		connectWith ConnectWith
+		wantErr     bool
+		wantIsError error
+	}{
+		{
+			name:        "valid",
+			connectWith: setupFn(),
+		},
+		{
+			name: "empty-SessionId",
+			connectWith: func() ConnectWith {
+				cw := setupFn()
+				cw.SessionId = ""
+				return cw
+			}(),
+			wantErr:     true,
+			wantIsError: db.ErrInvalidParameter,
+		},
+		{
+			name: "empty-ClientTcpAddress",
+			connectWith: func() ConnectWith {
+				cw := setupFn()
+				cw.ClientTcpAddress = ""
+				return cw
+			}(),
+			wantErr:     true,
+			wantIsError: db.ErrInvalidParameter,
+		},
+		{
+			name: "empty-ClientTcpPort",
+			connectWith: func() ConnectWith {
+				cw := setupFn()
+				cw.ClientTcpPort = 0
+				return cw
+			}(),
+			wantErr:     true,
+			wantIsError: db.ErrInvalidParameter,
+		},
+		{
+			name: "empty-EndpointTcpAddress",
+			connectWith: func() ConnectWith {
+				cw := setupFn()
+				cw.EndpointTcpAddress = ""
+				return cw
+			}(),
+			wantErr:     true,
+			wantIsError: db.ErrInvalidParameter,
+		},
+		{
+			name: "empty-EndpointTcpPort",
+			connectWith: func() ConnectWith {
+				cw := setupFn()
+				cw.EndpointTcpPort = 0
+				return cw
+			}(),
+			wantErr:     true,
+			wantIsError: db.ErrInvalidParameter,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			assert, require := assert.New(t), require.New(t)
+
+			c, cs, err := repo.ConnectSession(context.Background(), tt.connectWith)
+			if tt.wantErr {
+				require.Error(err)
+				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(StatusConnected.String(), cs[0].Status)
+		})
+	}
+}
+
+func TestRepository_TerminateSession(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() *Session {
+		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)
+		return s
+	}
+	tests := []struct {
+		name        string
+		session     *Session
+		reason      TerminationReason
+		wantErr     bool
+		wantIsError error
+	}{
+		{
+			name:    "valid-active-session",
+			session: setupFn(),
+			reason:  ClosedByUser,
+		},
+		{
+			name:    "valid-pending-session",
+			session: TestDefaultSession(t, conn, wrapper, iamRepo),
+			reason:  ClosedByUser,
+		},
+		{
+			name: "empty-session-id",
+			session: func() *Session {
+				s := setupFn()
+				s.PublicId = ""
+				return s
+			}(),
+			reason:      ClosedByUser,
+			wantErr:     true,
+			wantIsError: db.ErrInvalidParameter,
+		},
+		{
+			name: "empty-session-version",
+			session: func() *Session {
+				s := setupFn()
+				s.Version = 0
+				return s
+			}(),
+			reason:      ClosedByUser,
+			wantErr:     true,
+			wantIsError: db.ErrInvalidParameter,
+		},
+		{
+			name: "open-connection",
+			session: func() *Session {
+				s := setupFn()
+				_ = TestConnection(t, conn, s.PublicId, "127.0.0.1", 22, "127.0.0.1", 222)
+				return s
+			}(),
+			reason:  ClosedByUser,
+			wantErr: true,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			assert, require := assert.New(t), require.New(t)
+			s, ss, err := repo.TerminateSession(context.Background(), tt.session.PublicId, tt.session.Version, tt.reason)
+			if tt.wantErr {
+				require.Error(err)
+				if tt.wantIsError != nil {
+					assert.Truef(errors.Is(err, tt.wantIsError), "unexpected error %s", err.Error())
+				}
+				return
+			}
+			require.NoError(err)
+			assert.Equal(tt.reason.String(), s.TerminationReason)
+			assert.Equal(StatusTerminated.String(), ss[0].Status)
+		})
+	}
+}
+
+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)
+			require.NoError(t, err)
+			cw = append(cw, CloseWith{
+				ConnectionId:      c.PublicId,
+				ConnectionVersion: c.Version,
+				BytesUp:           1,
+				BytesDown:         2,
+				ClosedReason:      ConnectionClosedByUser,
+			})
+		}
+		return cw
+	}
+	tests := []struct {
+		name        string
+		closeWith   []CloseWith
+		reason      TerminationReason
+		wantErr     bool
+		wantIsError error
+	}{
+		{
+			name:      "valid",
+			closeWith: setupFn(2),
+			reason:    ClosedByUser,
+		},
+		{
+			name:        "empty-closed-with",
+			closeWith:   []CloseWith{},
+			reason:      ClosedByUser,
+			wantErr:     true,
+			wantIsError: db.ErrInvalidParameter,
+		},
+		{
+			name: "missing-ConnectionId",
+			closeWith: func() []CloseWith {
+				cw := setupFn(2)
+				cw[1].ConnectionId = ""
+				return cw
+
+			}(),
+			reason:      ClosedByUser,
+			wantErr:     true,
+			wantIsError: db.ErrInvalidParameter,
+		},
+		{
+			name: "ConnectionVersion",
+			closeWith: func() []CloseWith {
+				cw := setupFn(2)
+				cw[1].ConnectionVersion = 0
+				return cw
+
+			}(),
+			reason:      ClosedByUser,
+			wantErr:     true,
+			wantIsError: db.ErrInvalidParameter,
+		},
+	}
+	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)
+				if tt.wantIsError != nil {
+					assert.Truef(errors.Is(err, tt.wantIsError), "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.String(), r.ConnectionStates[0].Status)
+			}
+		})
+	}
+}
+
+func TestRepository_CancelSession(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() *Session {
+		session := TestDefaultSession(t, conn, wrapper, iamRepo)
+		_ = TestConnection(t, conn, session.PublicId, "127.0.0.1", 22, "127.0.0.1", 2222)
+		return session
+	}
+	tests := []struct {
+		name                   string
+		session                *Session
+		overrideSessionId      *string
+		overrideSessionVersion *uint32
+		wantErr                bool
+		wantIsError            error
+	}{
+		{
+			name:    "valid",
+			session: setupFn(),
+		},
+		{
+			name:    "bad-session-id",
+			session: setupFn(),
+			overrideSessionId: func() *string {
+				id, err := newId()
+				require.NoError(t, err)
+				return &id
+			}(),
+			wantErr: true,
+		},
+		{
+			name:    "missing-session-id",
+			session: setupFn(),
+			overrideSessionId: func() *string {
+				id := ""
+				return &id
+			}(),
+			wantErr:     true,
+			wantIsError: db.ErrInvalidParameter,
+		},
+		{
+			name:    "bad-version-id",
+			session: setupFn(),
+			overrideSessionVersion: func() *uint32 {
+				v := uint32(101)
+				return &v
+			}(),
+			wantErr: true,
+		},
+		{
+			name:    "missing-version-id",
+			session: setupFn(),
+			overrideSessionVersion: func() *uint32 {
+				v := uint32(0)
+				return &v
+			}(),
+			wantErr:     true,
+			wantIsError: db.ErrInvalidParameter,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			assert, require := assert.New(t), require.New(t)
+			var id string
+			var version uint32
+			switch {
+			case tt.overrideSessionId != nil:
+				id = *tt.overrideSessionId
+			default:
+				id = tt.session.PublicId
+			}
+			switch {
+			case tt.overrideSessionVersion != nil:
+				version = *tt.overrideSessionVersion
+			default:
+				version = tt.session.Version
+			}
+
+			s, ss, err := repo.CancelSession(context.Background(), id, version)
+			if tt.wantErr {
+				require.Error(err)
+				if tt.wantIsError != nil {
+					assert.Truef(errors.Is(err, tt.wantIsError), "unexpected error %s", err.Error())
+				}
+				return
+			}
+			require.NoError(err)
+			require.NotNil(s)
+			require.NotNil(ss)
+			assert.Equal(StatusCancelling.String(), ss[0].Status)
+		})
+	}
+}
+func TestRepository_ActivateSession(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)
+	worker := TestWorker(t, conn, wrapper)
+
 	tofu := TestTofu(t)
 	tests := []struct {
 		name                   string
@@ -425,7 +806,7 @@ func TestRepository_ActivateState(t *testing.T) {
 			name: "already-active",
 			session: func() *Session {
 				s := TestDefaultSession(t, conn, wrapper, iamRepo)
-				activeSession, _, err := repo.ActivateSession(context.Background(), s.PublicId, s.Version, tofu)
+				activeSession, _, err := repo.ActivateSession(context.Background(), s.PublicId, s.Version, worker.PrivateId, worker.Type, tofu)
 				require.NoError(t, err)
 				return activeSession
 			}(),
@@ -488,7 +869,7 @@ func TestRepository_ActivateState(t *testing.T) {
 			default:
 				version = tt.session.Version
 			}
-			s, ss, err := repo.ActivateSession(context.Background(), id, version, tofu)
+			s, ss, err := repo.ActivateSession(context.Background(), id, version, worker.PrivateId, worker.Type, tofu)
 			if tt.wantErr {
 				require.Error(err)
 				if tt.wantIsError != nil {
@@ -506,272 +887,21 @@ func TestRepository_ActivateState(t *testing.T) {
 		t.Run("already active", func(t *testing.T) {
 			assert, require := assert.New(t), require.New(t)
 			session := TestDefaultSession(t, conn, wrapper, iamRepo)
-			s, ss, err := repo.ActivateSession(context.Background(), session.PublicId, 1, tofu)
+			s, ss, err := repo.ActivateSession(context.Background(), session.PublicId, 1, worker.PrivateId, worker.Type, tofu)
 			require.NoError(err)
 			require.NotNil(s)
 			require.NotNil(ss)
 			assert.Equal(2, len(ss))
 			assert.Equal(StatusActive.String(), ss[0].Status)
 
-			_, _, err = repo.ActivateSession(context.Background(), session.PublicId, 1, tofu)
+			_, _, err = repo.ActivateSession(context.Background(), session.PublicId, 1, worker.PrivateId, worker.Type, tofu)
 			require.Error(err)
 
-			_, _, err = repo.ActivateSession(context.Background(), session.PublicId, 2, tofu)
+			_, _, err = repo.ActivateSession(context.Background(), session.PublicId, 2, worker.PrivateId, worker.Type, tofu)
 			require.Error(err)
 		})
 	}
 }
-func TestRepository_UpdateSession(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)
-	serversRepo, err := servers.NewRepository(rw, rw, kms)
-	require.NoError(t, err)
-
-	newServerFunc := func() string {
-		id, err := uuid.GenerateUUID()
-		require.NoError(t, err)
-		worker := &servers.Server{
-			Name:        "test-session-worker-" + id,
-			Type:        servers.ServerTypeWorker.String(),
-			Description: "Test Session Worker",
-			Address:     "127.0.0.1",
-		}
-		_, _, err = serversRepo.UpsertServer(context.Background(), worker)
-		require.NoError(t, err)
-		return worker.Name
-	}
-
-	type args struct {
-		terminationReason TerminationReason
-		serverId          string
-		serverType        string
-		tofu              []byte
-		fieldMaskPaths    []string
-		opt               []Option
-		publicId          *string // not updateable - db.ErrInvalidFieldMask
-		userId            string  // not updateable - db.ErrInvalidFieldMask
-		hostId            string  // not updateable - db.ErrInvalidFieldMask
-		targetId          string  // not updateable - db.ErrInvalidFieldMask
-		hostSetId         string  // not updateable - db.ErrInvalidFieldMask
-		authTokenId       string  // not updateable - db.ErrInvalidFieldMask
-		scopeId           string  // not updateable - db.ErrInvalidFieldMask
-	}
-	tests := []struct {
-		name           string
-		args           args
-		wantRowsUpdate int
-		wantErr        bool
-		wantIsError    error
-	}{
-		{
-			name: "valid",
-			args: args{
-				terminationReason: Terminated,
-				serverId:          newServerFunc(),
-				serverType:        servers.ServerTypeWorker.String(),
-				tofu:              TestTofu(t),
-				fieldMaskPaths:    []string{"TerminationReason", "ServerId", "ServerType", "TofuToken"},
-			},
-			wantErr:        false,
-			wantRowsUpdate: 1,
-		},
-		{
-			name: "publicId",
-			args: args{
-				publicId: func() *string {
-					id, err := newId()
-					require.NoError(t, err)
-					return &id
-				}(),
-				fieldMaskPaths: []string{"PublicId"},
-			},
-			wantErr:        true,
-			wantRowsUpdate: 0,
-			wantIsError:    db.ErrInvalidFieldMask,
-		},
-		{
-			name: "userId",
-			args: args{
-				userId: func() string {
-					org, _ := iam.TestScopes(t, iamRepo)
-					u := iam.TestUser(t, iamRepo, org.PublicId)
-					return u.PublicId
-				}(),
-				fieldMaskPaths: []string{"UserId"},
-			},
-			wantErr:        true,
-			wantRowsUpdate: 0,
-			wantIsError:    db.ErrInvalidFieldMask,
-		},
-		{
-			name: "hostId",
-			args: args{
-				hostId: func() string {
-					_, proj := iam.TestScopes(t, iamRepo)
-					cats := static.TestCatalogs(t, conn, proj.PublicId, 1)
-					hosts := static.TestHosts(t, conn, cats[0].PublicId, 1)
-					return hosts[0].PublicId
-				}(),
-				fieldMaskPaths: []string{"HostId"},
-			},
-			wantErr:        true,
-			wantRowsUpdate: 0,
-			wantIsError:    db.ErrInvalidFieldMask,
-		},
-		{
-			name: "targetId",
-			args: args{
-				targetId: func() string {
-					_, proj := iam.TestScopes(t, iamRepo)
-					tcpTarget := target.TestTcpTarget(t, conn, proj.PublicId, "test target")
-					return tcpTarget.PublicId
-				}(),
-				fieldMaskPaths: []string{"TargetId"},
-			},
-			wantErr:        true,
-			wantRowsUpdate: 0,
-			wantIsError:    db.ErrInvalidFieldMask,
-		},
-		{
-			name: "hostSetId",
-			args: args{
-				hostSetId: func() string {
-					_, proj := iam.TestScopes(t, iamRepo)
-					cats := static.TestCatalogs(t, conn, proj.PublicId, 1)
-					sets := static.TestSets(t, conn, cats[0].PublicId, 1)
-					return sets[0].PublicId
-				}(),
-				fieldMaskPaths: []string{"HostSetId"},
-			},
-			wantErr:        true,
-			wantRowsUpdate: 0,
-			wantIsError:    db.ErrInvalidFieldMask,
-		},
-		{
-			name: "AuthTokenId",
-			args: args{
-				authTokenId: func() string {
-					ctx := context.Background()
-					org, _ := iam.TestScopes(t, iamRepo)
-					authMethod := password.TestAuthMethods(t, conn, org.PublicId, 1)[0]
-					acct := password.TestAccounts(t, conn, authMethod.GetPublicId(), 1)[0]
-					user, err := iamRepo.LookupUserWithLogin(ctx, acct.GetPublicId(), iam.WithAutoVivify(true))
-					require.NoError(t, err)
-
-					authTokenRepo, err := authtoken.NewRepository(rw, rw, kms)
-					require.NoError(t, err)
-					at, err := authTokenRepo.CreateAuthToken(ctx, user, acct.GetPublicId())
-					require.NoError(t, err)
-					return at.PublicId
-				}(),
-				fieldMaskPaths: []string{"AuthTokenId"},
-			},
-			wantErr:        true,
-			wantRowsUpdate: 0,
-			wantIsError:    db.ErrInvalidFieldMask,
-		},
-		{
-			name: "ScopeId",
-			args: args{
-				scopeId: func() string {
-					_, proj := iam.TestScopes(t, iamRepo)
-					return proj.PublicId
-				}(),
-				fieldMaskPaths: []string{"ScopeId"},
-			},
-			wantErr:        true,
-			wantRowsUpdate: 0,
-			wantIsError:    db.ErrInvalidFieldMask,
-		},
-	}
-	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
-			assert, require := assert.New(t), require.New(t)
-
-			composedOf := TestSessionParams(t, conn, wrapper, iamRepo)
-			s := TestSession(t, conn, wrapper, composedOf)
-
-			updateSession := AllocSession()
-			updateSession.PublicId = s.PublicId
-			if tt.args.publicId != nil {
-				updateSession.PublicId = *tt.args.publicId
-			}
-			updateSession.ServerId = tt.args.serverId
-			updateSession.ServerType = tt.args.serverType
-			updateSession.TerminationReason = tt.args.terminationReason.String()
-			if tt.args.tofu != nil {
-				updateSession.TofuToken = make([]byte, len(tt.args.tofu))
-				copy(updateSession.TofuToken, tt.args.tofu)
-			}
-			updateSession.Version = s.Version
-			afterUpdateSession, afterUpdateState, updatedRows, err := repo.UpdateSession(context.Background(), &updateSession, updateSession.Version, tt.args.fieldMaskPaths, tt.args.opt...)
-
-			if tt.wantErr {
-				require.Error(err)
-				if tt.wantIsError != nil {
-					assert.Truef(errors.Is(err, tt.wantIsError), "unexpected error: %s", err.Error())
-				}
-				assert.Nil(afterUpdateSession)
-				assert.Nil(afterUpdateState)
-				assert.Equal(0, updatedRows)
-				err = db.TestVerifyOplog(t, rw, s.PublicId, db.WithOperation(oplog.OpType_OP_TYPE_UPDATE), db.WithCreateNotBefore(10*time.Second))
-				assert.Error(err)
-				assert.True(errors.Is(db.ErrRecordNotFound, err))
-				return
-			}
-			require.NoError(err)
-			assert.Equal(tt.wantRowsUpdate, updatedRows)
-			require.NotNil(afterUpdateSession)
-			require.NotNil(afterUpdateState)
-			switch tt.name {
-			case "valid-no-op":
-				assert.Equal(s.UpdateTime, afterUpdateSession.UpdateTime)
-			default:
-				assert.NotEqual(s.UpdateTime, afterUpdateSession.UpdateTime)
-			}
-			foundSession, foundStates, err := repo.LookupSession(context.Background(), s.PublicId)
-			require.NoError(err)
-			assert.Equal(afterUpdateSession, foundSession)
-			dbassrt := dbassert.New(t, rw)
-			if tt.args.serverId == "" {
-				dbassrt.IsNull(foundSession, "ServerId")
-			}
-			if tt.args.serverType == "" {
-				dbassrt.IsNull(foundSession, "ServerType")
-			}
-			if tt.args.tofu == nil {
-				dbassrt.IsNull(foundSession, "CtTofuToken")
-				dbassrt.IsNull(foundSession, "KeyId")
-			} else {
-				dbassrt.NotNull(foundSession, "CtTofuToken")
-				dbassrt.NotNull(foundSession, "KeyId")
-			}
-			assert.Equal(tt.args.terminationReason.String(), foundSession.TerminationReason)
-			assert.Equal(tt.args.serverId, foundSession.ServerId)
-			assert.Equal(tt.args.serverType, foundSession.ServerType)
-
-			err = db.TestVerifyOplog(t, rw, s.PublicId, db.WithOperation(oplog.OpType_OP_TYPE_UPDATE), db.WithCreateNotBefore(10*time.Second))
-			assert.Error(err)
-
-			switch {
-			case tt.args.terminationReason != "":
-				require.Equal(2, len(foundStates))
-				assert.Equal(StatusTerminated.String(), foundStates[0].Status)
-				assert.Equal(StatusPending.String(), foundStates[1].Status)
-			default:
-				require.Equal(1, len(foundStates))
-				assert.Equal(StatusPending.String(), foundStates[0].Status)
-			}
-		})
-	}
-
-}
 
 func TestRepository_DeleteSession(t *testing.T) {
 	t.Parallel()

internal/session/session.go

@@ -209,7 +209,7 @@ func (s *Session) VetForWrite(ctx context.Context, r db.Reader, opType db.OpType
 			return fmt.Errorf("session vet for write: update time is immutable: %w", db.ErrInvalidParameter)
 		case contains(opts.WithFieldMaskPaths, "TerminationReason"):
 			if _, err := convertToReason(s.TerminationReason); err != nil {
-				return fmt.Errorf("session vet for write: %w", db.ErrInvalidParameter)
+				return fmt.Errorf("session vet for write: termination reason '%s' is invalid: %w", s.TerminationReason, db.ErrInvalidParameter)
 			}
 		}
 	}

internal/session/session_connect_with.go

@@ -0,0 +1,36 @@
+package session
+
+import (
+	"fmt"
+
+	"github.com/hashicorp/boundary/internal/db"
+)
+
+// ConnectWith defines the boundary data that is saved in the repo when the
+// worker has established a connection between the client and the endpoint.
+type ConnectWith struct {
+	SessionId          string
+	ClientTcpAddress   string
+	ClientTcpPort      uint32
+	EndpointTcpAddress string
+	EndpointTcpPort    uint32
+}
+
+func (c ConnectWith) validate() error {
+	if c.SessionId == "" {
+		return fmt.Errorf("missing session id: %w", db.ErrInvalidParameter)
+	}
+	if c.ClientTcpAddress == "" {
+		return fmt.Errorf("missing client tcp address: %w", db.ErrInvalidParameter)
+	}
+	if c.ClientTcpPort == 0 {
+		return fmt.Errorf("missing client ctp port: %w", db.ErrInvalidParameter)
+	}
+	if c.EndpointTcpAddress == "" {
+		return fmt.Errorf("missing endpoint tcp address: %w", db.ErrInvalidParameter)
+	}
+	if c.EndpointTcpPort == 0 {
+		return fmt.Errorf("missing endpoint ctp port: %w", db.ErrInvalidParameter)
+	}
+	return nil
+}

internal/session/session_connect_with_test.go

@@ -0,0 +1,100 @@
+package session
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestConnectWith_validate(t *testing.T) {
+	id, err := newId()
+	require.NoError(t, err)
+
+	type fields struct {
+		SessionId          string
+		ClientTcpAddress   string
+		ClientTcpPort      uint32
+		EndpointTcpAddress string
+		EndpointTcpPort    uint32
+	}
+	tests := []struct {
+		name    string
+		fields  fields
+		wantErr bool
+	}{
+		{
+			name: "valid",
+			fields: fields{
+				SessionId:          id,
+				ClientTcpAddress:   "0.0.0.1",
+				ClientTcpPort:      22,
+				EndpointTcpAddress: "0.0.0.1",
+				EndpointTcpPort:    2222,
+			},
+		},
+		{
+			name: "missing-SessionId",
+			fields: fields{
+				ClientTcpAddress:   "0.0.0.1",
+				ClientTcpPort:      22,
+				EndpointTcpAddress: "0.0.0.1",
+				EndpointTcpPort:    2222,
+			},
+			wantErr: true,
+		},
+		{
+			name: "missing-ClientTcpAddress",
+			fields: fields{
+				SessionId:          id,
+				ClientTcpPort:      22,
+				EndpointTcpAddress: "0.0.0.1",
+				EndpointTcpPort:    2222,
+			},
+			wantErr: true,
+		},
+		{
+			name: "missing-ClientTcpPort",
+			fields: fields{
+				SessionId:          id,
+				ClientTcpAddress:   "0.0.0.1",
+				EndpointTcpAddress: "0.0.0.1",
+				EndpointTcpPort:    2222,
+			},
+			wantErr: true,
+		},
+		{
+			name: "missing-EndpointTcpAddress",
+			fields: fields{
+				SessionId:        id,
+				ClientTcpAddress: "0.0.0.1",
+				ClientTcpPort:    22,
+				EndpointTcpPort:  2222,
+			},
+			wantErr: true,
+		},
+		{
+			name: "missing-EndpointTcpPort",
+			fields: fields{
+				SessionId:          id,
+				ClientTcpAddress:   "0.0.0.1",
+				ClientTcpPort:      22,
+				EndpointTcpAddress: "0.0.0.1",
+			},
+			wantErr: true,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			c := ConnectWith{
+				SessionId:          tt.fields.SessionId,
+				ClientTcpAddress:   tt.fields.ClientTcpAddress,
+				ClientTcpPort:      tt.fields.ClientTcpPort,
+				EndpointTcpAddress: tt.fields.EndpointTcpAddress,
+				EndpointTcpPort:    tt.fields.EndpointTcpPort,
+			}
+			if err := c.validate(); (err != nil) != tt.wantErr {
+				t.Errorf("ConnectWith.validate() error = %v, wantErr %v", err, tt.wantErr)
+			}
+		})
+	}
+}

internal/session/state.go

@@ -19,7 +19,7 @@ type Status string
 const (
 	StatusPending    Status = "pending"
 	StatusActive     Status = "active"
-	StatusCanceling  Status = "canceling"
+	StatusCancelling Status = "canceling"
 	StatusTerminated Status = "terminated"
 )
 

internal/session/testing.go

@@ -14,19 +14,21 @@ import (
 	"github.com/hashicorp/boundary/internal/host/static"
 	"github.com/hashicorp/boundary/internal/iam"
 	"github.com/hashicorp/boundary/internal/kms"
+	"github.com/hashicorp/boundary/internal/servers"
 	"github.com/hashicorp/boundary/internal/target"
 	wrapping "github.com/hashicorp/go-kms-wrapping"
+	"github.com/hashicorp/go-uuid"
 	"github.com/hashicorp/vault/sdk/helper/base62"
 	"github.com/jinzhu/gorm"
 	"github.com/stretchr/testify/require"
 )
 
 // TestConnection creates a test connection for the sessionId in the repository.
-func TestConnection(t *testing.T, conn *gorm.DB, sessionId, clientTcpAddr string, clientTcpPort uint32, backendTcpAddr string, backendTcpPort uint32) *Connection {
+func TestConnection(t *testing.T, conn *gorm.DB, sessionId, clientTcpAddr string, clientTcpPort uint32, endpointTcpAddr string, endpointTcpPort uint32) *Connection {
 	t.Helper()
 	require := require.New(t)
 	rw := db.New(conn)
-	c, err := NewConnection(sessionId, clientTcpAddr, clientTcpPort, backendTcpAddr, backendTcpPort)
+	c, err := NewConnection(sessionId, clientTcpAddr, clientTcpPort, endpointTcpAddr, endpointTcpPort)
 	require.NoError(err)
 	id, err := newConnectionId()
 	require.NoError(err)
@@ -143,6 +145,26 @@ func TestTofu(t *testing.T) []byte {
 	return []byte(tofu)
 }
 
+func TestWorker(t *testing.T, conn *gorm.DB, wrapper wrapping.Wrapper) *servers.Server {
+	t.Helper()
+	rw := db.New(conn)
+	kms := kms.TestKms(t, conn, wrapper)
+	serversRepo, err := servers.NewRepository(rw, rw, kms)
+	require.NoError(t, err)
+
+	id, err := uuid.GenerateUUID()
+	require.NoError(t, err)
+	worker := &servers.Server{
+		Name:        "test-session-worker-" + id,
+		Type:        servers.ServerTypeWorker.String(),
+		Description: "Test Session Worker",
+		Address:     "127.0.0.1",
+	}
+	_, _, err = serversRepo.UpsertServer(context.Background(), worker)
+	require.NoError(t, err)
+	return worker
+}
+
 // TestCert is a temporary test func that intentionally doesn't take testing.T
 // as a parameter.  It's currently used in controller.jobTestingHandler() and
 // should be deprecated once that function is refactored to use sessions properly.

internal/session/testing_test.go

@@ -11,7 +11,6 @@ import (
 )
 
 func Test_TestSession(t *testing.T) {
-	t.Helper()
 	assert, require := assert.New(t), require.New(t)
 	conn, _ := db.TestSetup(t, "postgres")
 	wrapper := db.TestWrapper(t)
@@ -22,7 +21,6 @@ func Test_TestSession(t *testing.T) {
 }
 
 func Test_TestState(t *testing.T) {
-	t.Helper()
 	assert, require := assert.New(t), require.New(t)
 	conn, _ := db.TestSetup(t, "postgres")
 	wrapper := db.TestWrapper(t)
@@ -36,7 +34,6 @@ func Test_TestState(t *testing.T) {
 }
 
 func Test_TestConnection(t *testing.T) {
-	t.Helper()
 	assert, require := assert.New(t), require.New(t)
 	conn, _ := db.TestSetup(t, "postgres")
 	wrapper := db.TestWrapper(t)
@@ -50,7 +47,6 @@ func Test_TestConnection(t *testing.T) {
 }
 
 func Test_TestConnectionState(t *testing.T) {
-	t.Helper()
 	assert, require := assert.New(t), require.New(t)
 	conn, _ := db.TestSetup(t, "postgres")
 	wrapper := db.TestWrapper(t)
@@ -73,6 +69,13 @@ func Test_TestConnectionState(t *testing.T) {
 	assert.NotEmpty(initialState.StartTime)
 }
 
+func Test_TestWorker(t *testing.T) {
+	require := require.New(t)
+	conn, _ := db.TestSetup(t, "postgres")
+	wrapper := db.TestWrapper(t)
+	w := TestWorker(t, conn, wrapper)
+	require.NotNil(w)
+}
 func Test_TestCert(t *testing.T) {
 	t.Helper()
 	assert, require := assert.New(t), require.New(t)