diff --git a/internal/servers/controller/handlers/workers/worker_service.go b/internal/servers/controller/handlers/workers/worker_service.go index aea0bb1d9..98399425b 100644 --- a/internal/servers/controller/handlers/workers/worker_service.go +++ b/internal/servers/controller/handlers/workers/worker_service.go @@ -425,12 +425,12 @@ func (ws *workerServiceServer) AuthorizeConnection(ctx context.Context, req *pbs func (ws *workerServiceServer) ConnectConnection(ctx context.Context, req *pbs.ConnectConnectionRequest) (*pbs.ConnectConnectionResponse, error) { const op = "workers.(workerServiceServer).ConnectConnection" - sessRepo, err := ws.sessionRepoFn() + connRepo, err := ws.connectionRepoFn() if err != nil { return nil, status.Errorf(codes.Internal, "error getting session repo: %v", err) } - connectionInfo, connStates, err := sessRepo.ConnectConnection(ctx, session.ConnectWith{ + connectionInfo, connStates, err := connRepo.ConnectConnection(ctx, session.ConnectWith{ ConnectionId: req.GetConnectionId(), ClientTcpAddress: req.GetClientTcpAddress(), ClientTcpPort: req.GetClientTcpPort(), diff --git a/internal/servers/controller/session_cleanup_job_test.go b/internal/servers/controller/session_cleanup_job_test.go index 612594169..721389267 100644 --- a/internal/servers/controller/session_cleanup_job_test.go +++ b/internal/servers/controller/session_cleanup_job_test.go @@ -72,7 +72,7 @@ func TestSessionConnectionCleanupJob(t *testing.T) { // This is just to ensure we have a spread when we test it out. for i, connId := range connIds { if i%2 == 0 { - _, cs, err := sessionRepo.ConnectConnection(ctx, session.ConnectWith{ + _, cs, err := connectionRepo.ConnectConnection(ctx, session.ConnectWith{ ConnectionId: connId, ClientTcpAddress: "127.0.0.1", ClientTcpPort: 22, diff --git a/internal/session/repository_connection.go b/internal/session/repository_connection.go index 40440b1d4..e98e82924 100644 --- a/internal/session/repository_connection.go +++ b/internal/session/repository_connection.go @@ -184,6 +184,62 @@ func (r *ConnectionRepository) ListConnectionsBySessionId(ctx context.Context, s return connections, nil } +// ConnectConnection updates a connection in the repo with a state of "connected". +func (r *ConnectionRepository) ConnectConnection(ctx context.Context, c ConnectWith) (*Connection, []*ConnectionState, error) { + const op = "session.(ConnectionRepository).ConnectConnection" + // ConnectWith.validate will check all the fields... + if err := c.validate(); err != nil { + return nil, nil, errors.Wrap(ctx, err, op) + } + var connection Connection + var connectionStates []*ConnectionState + _, err := r.writer.DoTx( + ctx, + db.StdRetryCnt, + db.ExpBackoff{}, + func(reader db.Reader, w db.Writer) error { + connection = AllocConnection() + connection.PublicId = c.ConnectionId + connection.ClientTcpAddress = c.ClientTcpAddress + connection.ClientTcpPort = c.ClientTcpPort + connection.EndpointTcpAddress = c.EndpointTcpAddress + connection.EndpointTcpPort = c.EndpointTcpPort + connection.UserClientIp = c.UserClientIp + fieldMask := []string{ + "ClientTcpAddress", + "ClientTcpPort", + "EndpointTcpAddress", + "EndpointTcpPort", + "UserClientIp", + } + rowsUpdated, err := w.Update(ctx, &connection, fieldMask, nil) + if err != nil { + return errors.Wrap(ctx, err, op) + } + if err == nil && rowsUpdated > 1 { + // return err, which will result in a rollback of the update + return errors.New(ctx, errors.MultipleRecords, op, "more than 1 resource would have been updated") + } + newState, err := NewConnectionState(connection.PublicId, StatusConnected) + if err != nil { + return errors.Wrap(ctx, err, op) + } + if err := w.Create(ctx, newState); err != nil { + return errors.Wrap(ctx, err, op) + } + connectionStates, err = fetchConnectionStates(ctx, reader, c.ConnectionId, db.WithOrder("start_time desc")) + if err != nil { + return errors.Wrap(ctx, err, op) + } + return nil + }, + ) + if err != nil { + return nil, nil, errors.Wrap(ctx, err, op) + } + return &connection, connectionStates, nil +} + // DeleteConnection will delete a connection from the repository. func (r *ConnectionRepository) DeleteConnection(ctx context.Context, publicId string, _ ...Option) (int, error) { const op = "session.(ConnectionRepository).DeleteConnection" diff --git a/internal/session/repository_connection_test.go b/internal/session/repository_connection_test.go index 82fd57e96..529bbf9ba 100644 --- a/internal/session/repository_connection_test.go +++ b/internal/session/repository_connection_test.go @@ -133,6 +133,131 @@ func TestRepository_ListConnection(t *testing.T) { }) } +func TestRepository_ConnectConnection(t *testing.T) { + t.Parallel() + ctx := context.Background() + conn, _ := db.TestSetup(t, "postgres") + rw := db.New(conn) + wrapper := db.TestWrapper(t) + iamRepo := iam.TestRepo(t, conn, wrapper) + kms := kms.TestKms(t, conn, wrapper) + repo, err := NewRepository(rw, rw, kms) + connRepo, err := NewConnectionRepository(ctx, rw, rw, kms) + require.NoError(t, err) + + setupFn := func() ConnectWith { + 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) + c := TestConnection(t, conn, s.PublicId, "127.0.0.1", 22, "127.0.0.1", 2222, "127.0.0.1") + return ConnectWith{ + ConnectionId: c.PublicId, + ClientTcpAddress: "127.0.0.1", + ClientTcpPort: 22, + EndpointTcpAddress: "127.0.0.1", + EndpointTcpPort: 2222, + UserClientIp: "127.0.0.1", + } + } + tests := []struct { + name string + connectWith ConnectWith + wantErr bool + wantIsError errors.Code + }{ + { + name: "valid", + connectWith: setupFn(), + }, + { + name: "empty-SessionId", + connectWith: func() ConnectWith { + cw := setupFn() + cw.ConnectionId = "" + return cw + }(), + wantErr: true, + wantIsError: errors.InvalidParameter, + }, + { + name: "empty-ClientTcpAddress", + connectWith: func() ConnectWith { + cw := setupFn() + cw.ClientTcpAddress = "" + return cw + }(), + wantErr: true, + wantIsError: errors.InvalidParameter, + }, + { + name: "empty-ClientTcpPort", + connectWith: func() ConnectWith { + cw := setupFn() + cw.ClientTcpPort = 0 + return cw + }(), + wantErr: true, + wantIsError: errors.InvalidParameter, + }, + { + name: "empty-EndpointTcpAddress", + connectWith: func() ConnectWith { + cw := setupFn() + cw.EndpointTcpAddress = "" + return cw + }(), + wantErr: true, + wantIsError: errors.InvalidParameter, + }, + { + name: "empty-EndpointTcpPort", + connectWith: func() ConnectWith { + cw := setupFn() + cw.EndpointTcpPort = 0 + return cw + }(), + wantErr: true, + wantIsError: errors.InvalidParameter, + }, + { + name: "empty-UserClientIp", + connectWith: func() ConnectWith { + cw := setupFn() + cw.UserClientIp = "" + return cw + }(), + wantErr: true, + wantIsError: errors.InvalidParameter, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + assert, require := assert.New(t), require.New(t) + + c, cs, err := connRepo.ConnectConnection(context.Background(), tt.connectWith) + if tt.wantErr { + require.Error(err) + assert.Truef(errors.Match(errors.T(tt.wantIsError), err), "unexpected error %s", err.Error()) + return + } + require.NoError(err) + require.NotNil(c) + require.NotNil(cs) + assert.Equal(StatusConnected, cs[0].Status) + gotConn, _, err := connRepo.LookupConnection(context.Background(), c.PublicId) + require.NoError(err) + assert.Equal(tt.connectWith.ClientTcpAddress, gotConn.ClientTcpAddress) + assert.Equal(tt.connectWith.ClientTcpPort, gotConn.ClientTcpPort) + assert.Equal(tt.connectWith.ClientTcpAddress, gotConn.ClientTcpAddress) + assert.Equal(tt.connectWith.EndpointTcpAddress, gotConn.EndpointTcpAddress) + assert.Equal(tt.connectWith.EndpointTcpPort, gotConn.EndpointTcpPort) + assert.Equal(tt.connectWith.UserClientIp, gotConn.UserClientIp) + }) + } +} + func TestRepository_DeleteConnection(t *testing.T) { t.Parallel() ctx := context.Background() @@ -261,7 +386,7 @@ func TestRepository_CloseDeadConnectionsOnWorker(t *testing.T) { // This is just to ensure we have a spread when we test it out. for i, connId := range connIds { if i%2 == 0 { - _, cs, err := repo.ConnectConnection(ctx, ConnectWith{ + _, cs, err := connRepo.ConnectConnection(ctx, ConnectWith{ ConnectionId: connId, ClientTcpAddress: "127.0.0.1", ClientTcpPort: 22, @@ -406,7 +531,7 @@ func TestRepository_CloseConnectionsForDeadWorkers(t *testing.T) { return s }() { if i%3 == 0 { - _, cs, err := repo.ConnectConnection(ctx, ConnectWith{ + _, cs, err := connRepo.ConnectConnection(ctx, ConnectWith{ ConnectionId: connId, ClientTcpAddress: "127.0.0.1", ClientTcpPort: 22, @@ -662,7 +787,7 @@ func TestRepository_ShouldCloseConnectionsOnWorker(t *testing.T) { // Mark half of the connections connected, close the other half. for i, connId := range connIds { if i%2 == 0 { - _, cs, err := repo.ConnectConnection(ctx, ConnectWith{ + _, cs, err := connRepo.ConnectConnection(ctx, ConnectWith{ ConnectionId: connId, ClientTcpAddress: "127.0.0.1", ClientTcpPort: 22, diff --git a/internal/session/repository_session.go b/internal/session/repository_session.go index 861e2c5b0..5fcccf912 100644 --- a/internal/session/repository_session.go +++ b/internal/session/repository_session.go @@ -402,62 +402,6 @@ func (r *Repository) sessionAuthzSummary(ctx context.Context, sessionId string) return info, nil } -// ConnectConnection updates a connection in the repo with a state of "connected". -func (r *Repository) ConnectConnection(ctx context.Context, c ConnectWith) (*Connection, []*ConnectionState, error) { - const op = "session.(ConnectionRepository).ConnectConnection" - // ConnectWith.validate will check all the fields... - if err := c.validate(); err != nil { - return nil, nil, errors.Wrap(ctx, err, op) - } - var connection Connection - var connectionStates []*ConnectionState - _, err := r.writer.DoTx( - ctx, - db.StdRetryCnt, - db.ExpBackoff{}, - func(reader db.Reader, w db.Writer) error { - connection = AllocConnection() - connection.PublicId = c.ConnectionId - connection.ClientTcpAddress = c.ClientTcpAddress - connection.ClientTcpPort = c.ClientTcpPort - connection.EndpointTcpAddress = c.EndpointTcpAddress - connection.EndpointTcpPort = c.EndpointTcpPort - connection.UserClientIp = c.UserClientIp - fieldMask := []string{ - "ClientTcpAddress", - "ClientTcpPort", - "EndpointTcpAddress", - "EndpointTcpPort", - "UserClientIp", - } - rowsUpdated, err := w.Update(ctx, &connection, fieldMask, nil) - if err != nil { - return errors.Wrap(ctx, err, op) - } - if err == nil && rowsUpdated > 1 { - // return err, which will result in a rollback of the update - return errors.New(ctx, errors.MultipleRecords, op, "more than 1 resource would have been updated") - } - newState, err := NewConnectionState(connection.PublicId, StatusConnected) - if err != nil { - return errors.Wrap(ctx, err, op) - } - if err := w.Create(ctx, newState); err != nil { - return errors.Wrap(ctx, err, op) - } - connectionStates, err = fetchConnectionStates(ctx, reader, c.ConnectionId, db.WithOrder("start_time desc")) - if err != nil { - return errors.Wrap(ctx, err, op) - } - return nil - }, - ) - if err != nil { - return nil, nil, errors.Wrap(ctx, err, op) - } - return &connection, connectionStates, nil -} - // CloseConnectionResp is just a wrapper for the response from CloseConnections. // It wraps the connection and its states for each connection closed. type CloseConnectionResp struct { diff --git a/internal/session/repository_session_test.go b/internal/session/repository_session_test.go index b2fbca45f..243e568f0 100644 --- a/internal/session/repository_session_test.go +++ b/internal/session/repository_session_test.go @@ -541,131 +541,6 @@ func TestRepository_updateState(t *testing.T) { } } -func TestRepository_ConnectConnection(t *testing.T) { - t.Parallel() - ctx := context.Background() - conn, _ := db.TestSetup(t, "postgres") - rw := db.New(conn) - wrapper := db.TestWrapper(t) - iamRepo := iam.TestRepo(t, conn, wrapper) - kms := kms.TestKms(t, conn, wrapper) - repo, err := NewRepository(rw, rw, kms) - connRepo, err := NewConnectionRepository(ctx, rw, rw, kms) - require.NoError(t, err) - - setupFn := func() ConnectWith { - 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) - c := TestConnection(t, conn, s.PublicId, "127.0.0.1", 22, "127.0.0.1", 2222, "127.0.0.1") - return ConnectWith{ - ConnectionId: c.PublicId, - ClientTcpAddress: "127.0.0.1", - ClientTcpPort: 22, - EndpointTcpAddress: "127.0.0.1", - EndpointTcpPort: 2222, - UserClientIp: "127.0.0.1", - } - } - tests := []struct { - name string - connectWith ConnectWith - wantErr bool - wantIsError errors.Code - }{ - { - name: "valid", - connectWith: setupFn(), - }, - { - name: "empty-SessionId", - connectWith: func() ConnectWith { - cw := setupFn() - cw.ConnectionId = "" - return cw - }(), - wantErr: true, - wantIsError: errors.InvalidParameter, - }, - { - name: "empty-ClientTcpAddress", - connectWith: func() ConnectWith { - cw := setupFn() - cw.ClientTcpAddress = "" - return cw - }(), - wantErr: true, - wantIsError: errors.InvalidParameter, - }, - { - name: "empty-ClientTcpPort", - connectWith: func() ConnectWith { - cw := setupFn() - cw.ClientTcpPort = 0 - return cw - }(), - wantErr: true, - wantIsError: errors.InvalidParameter, - }, - { - name: "empty-EndpointTcpAddress", - connectWith: func() ConnectWith { - cw := setupFn() - cw.EndpointTcpAddress = "" - return cw - }(), - wantErr: true, - wantIsError: errors.InvalidParameter, - }, - { - name: "empty-EndpointTcpPort", - connectWith: func() ConnectWith { - cw := setupFn() - cw.EndpointTcpPort = 0 - return cw - }(), - wantErr: true, - wantIsError: errors.InvalidParameter, - }, - { - name: "empty-UserClientIp", - connectWith: func() ConnectWith { - cw := setupFn() - cw.UserClientIp = "" - return cw - }(), - wantErr: true, - wantIsError: errors.InvalidParameter, - }, - } - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - assert, require := assert.New(t), require.New(t) - - c, cs, err := repo.ConnectConnection(context.Background(), tt.connectWith) - if tt.wantErr { - require.Error(err) - assert.Truef(errors.Match(errors.T(tt.wantIsError), err), "unexpected error %s", err.Error()) - return - } - require.NoError(err) - require.NotNil(c) - require.NotNil(cs) - assert.Equal(StatusConnected, cs[0].Status) - gotConn, _, err := connRepo.LookupConnection(context.Background(), c.PublicId) - require.NoError(err) - assert.Equal(tt.connectWith.ClientTcpAddress, gotConn.ClientTcpAddress) - assert.Equal(tt.connectWith.ClientTcpPort, gotConn.ClientTcpPort) - assert.Equal(tt.connectWith.ClientTcpAddress, gotConn.ClientTcpAddress) - assert.Equal(tt.connectWith.EndpointTcpAddress, gotConn.EndpointTcpAddress) - assert.Equal(tt.connectWith.EndpointTcpPort, gotConn.EndpointTcpPort) - assert.Equal(tt.connectWith.UserClientIp, gotConn.UserClientIp) - }) - } -} - func TestRepository_TerminateCompletedSessions(t *testing.T) { t.Parallel() ctx := context.Background()