CyberCoveLLC/threshold
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refactor(connections): Move connect connect to the Connection repo

Browse Source
This commit is contained in:
irenarindos committed 2022-03-15 12:11:31 -04:00
1 parent ac9fb8016f
commit 3482f91a17
6 files changed
+187 -187
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

internal/servers/controller/handlers/workers/worker_service.go
+2 -2
View File
@@ -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(),
internal/servers/controller/session_cleanup_job_test.go
+1 -1
View File
@@ -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,
internal/session/repository_connection.go
+56
View File
@@ -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"
internal/session/repository_connection_test.go
+128 -3
View File
@@ -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,
internal/session/repository_session.go
-56
View File
@@ -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 {
internal/session/repository_session_test.go
-125
View File
@@ -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()