Add tracking of worker used for injection (#3204)

internal/daemon/controller/handlers/targets/target_service.go

@@ -922,6 +922,9 @@ func (s Service) AuthorizeSession(ctx context.Context, req *pbs.AuthorizeSession
 		DynamicCredentials:  dynCreds,
 		StaticCredentials:   staticCreds,
 	}
+	if protoWorker != nil {
+		sessionComposition.ProtocolWorkerId = protoWorker.GetPublicId()
+	}
 	sess, err := session.New(sessionComposition)
 	if err != nil {
 		return nil, err

internal/db/schema/migrations/oss/postgres/60/02_sessions.up.sql

@@ -176,6 +176,7 @@ begin;
     for each row execute procedure insert_session();
 
   -- Replaces view from 59/01_target_ingress_egress_worker_filters.up.sql
+  -- Replaced in 69/02_session_worker_protocol.up.sql
   create view session_list as
   select
     s.public_id,

internal/db/schema/migrations/oss/postgres/69/02_session_worker_protocol.up.sql

@@ -0,0 +1,67 @@
+-- Copyright (c) HashiCorp, Inc.
+-- SPDX-License-Identifier: MPL-2.0
+
+begin;
+
+  create table session_worker_protocol (
+    session_id wt_public_id
+      constraint session_fkey
+        references session(public_id)
+        on delete cascade
+        on update cascade,
+    worker_id wt_public_id
+      constraint server_worker_fkey
+        references server_worker(public_id)
+        on delete cascade
+        on update cascade
+  );
+
+  drop view session_list;
+  -- Replaces view from 60/02_sessions.up.sql to add swp.worker_id
+  create view session_list as
+  select
+    s.public_id,
+    s.user_id,
+    shsh.host_id,
+    s.target_id,
+    shsh.host_set_id,
+    s.auth_token_id,
+    s.project_id,
+    s.certificate,
+    s.certificate_private_key,
+    s.expiration_time,
+    s.connection_limit,
+    s.tofu_token,
+    s.key_id,
+    s.termination_reason,
+    s.version,
+    s.create_time,
+    s.update_time,
+    s.endpoint,
+    s.worker_filter,
+    s.egress_worker_filter,
+    s.ingress_worker_filter,
+    swp.worker_id,
+    ss.state,
+    ss.previous_end_time,
+    ss.start_time,
+    ss.end_time,
+    sc.public_id as connection_id,
+    sc.client_tcp_address,
+    sc.client_tcp_port,
+    sc.endpoint_tcp_address,
+    sc.endpoint_tcp_port,
+    sc.bytes_up,
+    sc.bytes_down,
+    sc.closed_reason
+  from session s
+    join session_state ss on
+      s.public_id = ss.session_id
+    left join session_connection sc on
+      s.public_id = sc.session_id
+    left join session_host_set_host shsh on
+      s.public_id = shsh.session_id
+    left join session_worker_protocol swp on
+      s.public_id = swp.session_id;
+
+commit;

internal/session/repository_session.go

@@ -112,6 +112,17 @@ func (r *Repository) CreateSession(ctx context.Context, sessionWrapper wrapping.
 				}
 			}
 
+			if newSession.ProtocolWorkerId != "" {
+				swp, err := NewSessionWorkerProtocol(ctx, newSession.PublicId, newSession.ProtocolWorkerId)
+				if err != nil {
+					return errors.Wrap(ctx, err, op)
+				}
+				if err = w.Create(ctx, swp); err != nil {
+					return errors.Wrap(ctx, err, op)
+				}
+				returnedSession.ProtocolWorkerId = swp.WorkerId
+			}
+
 			for _, cred := range newSession.DynamicCredentials {
 				cred.SessionId = newSession.PublicId
 			}
@@ -225,6 +236,12 @@ func (r *Repository) LookupSession(ctx context.Context, sessionId string, opt ..
 			session.HostSetId = sessionHostSetHost.HostSetId
 			session.HostId = sessionHostSetHost.HostId
 
+			sessionWorkerProtocol := AllocSessionWorkerProtocol()
+			if err := read.LookupWhere(ctx, sessionWorkerProtocol, "session_id = ?", []any{sessionId}); err != nil && !errors.IsNotFoundError(err) {
+				return errors.Wrap(ctx, err, op)
+			}
+			session.ProtocolWorkerId = sessionWorkerProtocol.WorkerId
+
 			connections, err := fetchConnections(ctx, read, sessionId, db.WithOrder("create_time desc"))
 			if err != nil {
 				return errors.Wrap(ctx, err, op)

internal/session/repository_session_test.go

@@ -25,6 +25,7 @@ import (
 	"github.com/hashicorp/boundary/internal/kms"
 	"github.com/hashicorp/boundary/internal/oplog"
 	"github.com/hashicorp/boundary/internal/perms"
+	"github.com/hashicorp/boundary/internal/server"
 	"github.com/hashicorp/boundary/internal/target"
 	"github.com/hashicorp/boundary/internal/target/tcp"
 	tcpStore "github.com/hashicorp/boundary/internal/target/tcp/store"
@@ -282,6 +283,7 @@ func TestRepository_CreateSession(t *testing.T) {
 	ctx := context.Background()
 	repo, err := NewRepository(ctx, rw, rw, kmsCache)
 	require.NoError(t, err)
+	worker := server.TestKmsWorker(t, conn, wrapper)
 
 	workerAddresses := []string{"1.2.3.4"}
 	type args struct {
@@ -318,6 +320,30 @@ func TestRepository_CreateSession(t *testing.T) {
 			},
 			wantErr: false,
 		},
+		{
+			name: "valid-with-protocol-worker",
+			args: args{
+				composedOf: func() ComposedOf {
+					c := TestSessionParams(t, conn, wrapper, iamRepo)
+					c.ProtocolWorkerId = worker.PublicId
+					return c
+				}(),
+				workerAddresses: workerAddresses,
+			},
+		},
+		{
+			name: "invalid-protocol-worker",
+			args: args{
+				composedOf: func() ComposedOf {
+					c := TestSessionParams(t, conn, wrapper, iamRepo)
+					c.ProtocolWorkerId = "something"
+					return c
+				}(),
+				workerAddresses: workerAddresses,
+			},
+			wantErr:     true,
+			wantIsError: errors.CheckConstraint,
+		},
 		{
 			name: "empty-host-source-endpoint",
 			args: args{
@@ -445,6 +471,7 @@ func TestRepository_CreateSession(t *testing.T) {
 				ConnectionLimit:    tt.args.composedOf.ConnectionLimit,
 				DynamicCredentials: tt.args.composedOf.DynamicCredentials,
 				StaticCredentials:  tt.args.composedOf.StaticCredentials,
+				ProtocolWorkerId:   tt.args.composedOf.ProtocolWorkerId,
 			}
 			ses, err := repo.CreateSession(context.Background(), wrapper, s, tt.args.workerAddresses)
 

internal/session/session.go

@@ -63,6 +63,8 @@ type ComposedOf struct {
 	// StaticCredentials are static credentials that will be retrieved
 	// for the session. StaticCredentials optional.
 	StaticCredentials []*StaticCredential
+	// Which worker is performing protocol-related tasks
+	ProtocolWorkerId string
 }
 
 // Session contains information about a user's session with a target
@@ -129,6 +131,9 @@ type Session struct {
 	// HostId of the session
 	HostId string `gorm:"-"`
 
+	// ProtocolWorkerId of the session
+	ProtocolWorkerId string `gorm:"-"`
+
 	// Connections for the session are for read only and are ignored during write operations
 	Connections []*Connection `gorm:"-"`
 
@@ -171,6 +176,7 @@ func New(c ComposedOf, _ ...Option) (*Session, error) {
 		IngressWorkerFilter: c.IngressWorkerFilter,
 		DynamicCredentials:  c.DynamicCredentials,
 		StaticCredentials:   c.StaticCredentials,
+		ProtocolWorkerId:    c.ProtocolWorkerId,
 	}
 	if err := s.validateNewSession(); err != nil {
 		return nil, errors.WrapDeprecated(err, op)
@@ -201,6 +207,7 @@ func (s *Session) Clone() any {
 		EgressWorkerFilter:  s.EgressWorkerFilter,
 		IngressWorkerFilter: s.IngressWorkerFilter,
 		KeyId:               s.KeyId,
+		ProtocolWorkerId:    s.ProtocolWorkerId,
 	}
 	if len(s.States) > 0 {
 		clone.States = make([]*State, 0, len(s.States))
@@ -322,6 +329,8 @@ func (s *Session) VetForWrite(ctx context.Context, _ db.Reader, opType db.OpType
 			return errors.New(ctx, errors.InvalidParameter, op, "dynamic credentials are immutable")
 		case contains(opts.WithFieldMaskPaths, "StaticCredentials"):
 			return errors.New(ctx, errors.InvalidParameter, op, "static credentials are immutable")
+		case contains(opts.WithFieldMaskPaths, "ProtocolWorkerId"):
+			return errors.New(ctx, errors.InvalidParameter, op, "protocol worker id is immutable")
 		case contains(opts.WithFieldMaskPaths, "TerminationReason"):
 			if _, err := convertToReason(s.TerminationReason); err != nil {
 				return errors.Wrap(ctx, err, op)
@@ -376,7 +385,8 @@ func (s *Session) validateNewSession() error {
 	if s.CtTofuToken != nil {
 		return errors.NewDeprecated(errors.InvalidParameter, op, "ct must be empty")
 	}
-	// It is okay for the worker filter to be empty, so it is not checked here.
+	// It is okay for the worker filter and protocol worker ID to be empty, so
+	// they are not checked here.
 	return nil
 }
 
@@ -548,6 +558,7 @@ type sessionListView struct {
 	Endpoint                string               `json:"-" gorm:"default:null"`
 	ConnectionLimit         int32                `json:"connection_limit,omitempty" gorm:"default:null"`
 	KeyId                   string               `json:"key_id,omitempty" gorm:"default:null"`
+	ProtocolWorkerId        string               `json:"protocol_worker_id,omitempty" gorm:"default:null"`
 
 	// State fields
 	Status          string               `json:"state,omitempty" gorm:"column:state"`

internal/session/session_worker_protocol.go

@@ -0,0 +1,71 @@
+// Copyright (c) HashiCorp, Inc.
+// SPDX-License-Identifier: MPL-2.0
+
+package session
+
+import (
+	"context"
+
+	"github.com/hashicorp/boundary/internal/errors"
+)
+
+const (
+	defaultSessionWorkerProtocolTableName = "session_worker_protocol"
+)
+
+// SessionWorkerProtocol contains information about a chosen protocol-aware
+// worker for a session
+type SessionWorkerProtocol struct {
+	// SessionId of the session
+	SessionId string `json:"session_id,omitempty" gorm:"primary_key"`
+	// WorkerId chosen for protocol tasks
+	WorkerId string `json:"worker_id,omitempty" gorm:"primary_key"`
+
+	tableName string `gorm:"-"`
+}
+
+// NewSessionWorkerProtocol creates a new in-memory session to protocol worker
+// association
+func NewSessionWorkerProtocol(ctx context.Context, sessionId, workerId string) (*SessionWorkerProtocol, error) {
+	const op = "session.NewSessionWorkerProtocol"
+	if sessionId == "" {
+		return nil, errors.New(ctx, errors.InvalidParameter, op, "missing session id")
+	}
+	if workerId == "" {
+		return nil, errors.New(ctx, errors.InvalidParameter, op, "missing worker id")
+	}
+	swp := &SessionWorkerProtocol{
+		SessionId: sessionId,
+		WorkerId:  workerId,
+	}
+	return swp, nil
+}
+
+// TableName returns the tablename to override the default gorm table name
+func (s *SessionWorkerProtocol) TableName() string {
+	if s.tableName != "" {
+		return s.tableName
+	}
+	return defaultSessionWorkerProtocolTableName
+}
+
+// SetTableName sets the tablename and satisfies the ReplayableMessage
+// interface. If the caller attempts to set the name to "" the name will be
+// reset to the default name.
+func (s *SessionWorkerProtocol) SetTableName(n string) {
+	s.tableName = n
+}
+
+// AllocSessionWorkerProtocol will allocate a SessionHostSetHost
+func AllocSessionWorkerProtocol() *SessionWorkerProtocol {
+	return &SessionWorkerProtocol{}
+}
+
+// Clone creates a clone of the SessionWorkerProtocol
+func (s *SessionWorkerProtocol) Clone() any {
+	clone := &SessionWorkerProtocol{
+		SessionId: s.SessionId,
+		WorkerId:  s.WorkerId,
+	}
+	return clone
+}