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

refact(session connection): remove session connection state table (#4617)

Browse Source 
* refact(session connection): remove session connection state table
This commit is contained in:
Irena Rindos authored and GitHub committed 2024-07-23 14:19:21 -04:00
1 parent eb6e1b558d
commit 2191aa1d03
27 files changed
+543 -729
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

internal/session/connection.go
+4 -1
View File
@@ -13,7 +13,7 @@ import (
)
const (
defaultConnectionTableName = "session_connection"
defaultConnectionTableName = "session_connection_with_status_view" // "session_connection"
)
// Connection contains information about session's connection to a target
@@ -44,6 +44,8 @@ type Connection struct {
UpdateTime *timestamp.Timestamp `json:"update_time,omitempty" gorm:"default:current_timestamp"`
// Version of the connection
Version uint32 `json:"version,omitempty" gorm:"default:null"`
// Status is a field derived from connected_time_range
Status string `json:"status,omitempty" gorm:"default:null"`
tableName string `gorm:"-"`
}
@@ -94,6 +96,7 @@ func (c *Connection) Clone() any {
BytesDown: c.BytesDown,
ClosedReason: c.ClosedReason,
Version: c.Version,
Status: c.Status,
}
if c.CreateTime != nil {
clone.CreateTime = &timestamp.Timestamp{
internal/session/connection_state.go
+10 -129
View File
@@ -4,20 +4,9 @@
package session
import (
"context"
"github.com/hashicorp/boundary/internal/db"
"github.com/hashicorp/boundary/internal/db/timestamp"
"github.com/hashicorp/boundary/internal/errors"
"google.golang.org/protobuf/types/known/timestamppb"
workerpbs "github.com/hashicorp/boundary/internal/gen/controller/servers/services"
)
const (
defaultConnectionStateTableName = "session_connection_state"
)
// ConnectionStatus of the connection's state
type ConnectionStatus string
@@ -60,122 +49,14 @@ func ConnectionStatusFromProtoVal(s workerpbs.CONNECTIONSTATUS) ConnectionStatus
return StatusUnspecified
}
// ConnectionState of the state of the connection
type ConnectionState struct {
// ConnectionId is used to access the state via an API
ConnectionId string `json:"public_id,omitempty" gorm:"primary_key"`
// status of the connection
Status ConnectionStatus `protobuf:"bytes,20,opt,name=status,proto3" json:"status,omitempty" gorm:"column:state"`
// PreviousEndTime from the RDBMS
PreviousEndTime *timestamp.Timestamp `json:"previous_end_time,omitempty" gorm:"default:current_timestamp"`
// StartTime from the RDBMS
StartTime *timestamp.Timestamp `json:"start_time,omitempty" gorm:"default:current_timestamp;primary_key"`
// EndTime from the RDBMS
EndTime *timestamp.Timestamp `json:"end_time,omitempty" gorm:"default:current_timestamp"`
tableName string `gorm:"-"`
}
var (
_ Cloneable = (*ConnectionState)(nil)
_ db.VetForWriter = (*ConnectionState)(nil)
)
// NewConnectionState creates a new in memory connection state. No options
// are currently supported.
func NewConnectionState(ctx context.Context, connectionId string, state ConnectionStatus, _ ...Option) (*ConnectionState, error) {
const op = "session.NewConnectionState"
s := ConnectionState{
ConnectionId: connectionId,
Status: state,
}
if err := s.validate(ctx); err != nil {
return nil, errors.Wrap(ctx, err, op)
}
return &s, nil
}
// allocConnectionState will allocate a connection State
func allocConnectionState() ConnectionState {
return ConnectionState{}
}
// Clone creates a clone of the State
func (s *ConnectionState) Clone() any {
clone := &ConnectionState{
ConnectionId: s.ConnectionId,
Status: s.Status,
}
if s.PreviousEndTime != nil {
clone.PreviousEndTime = &timestamp.Timestamp{
Timestamp: &timestamppb.Timestamp{
Seconds: s.PreviousEndTime.Timestamp.Seconds,
Nanos: s.PreviousEndTime.Timestamp.Nanos,
},
}
}
if s.StartTime != nil {
clone.StartTime = &timestamp.Timestamp{
Timestamp: &timestamppb.Timestamp{
Seconds: s.StartTime.Timestamp.Seconds,
Nanos: s.StartTime.Timestamp.Nanos,
},
}
}
if s.EndTime != nil {
clone.EndTime = &timestamp.Timestamp{
Timestamp: &timestamppb.Timestamp{
Seconds: s.EndTime.Timestamp.Seconds,
Nanos: s.EndTime.Timestamp.Nanos,
},
}
}
return clone
}
// VetForWrite implements db.VetForWrite() interface and validates the state
// before it's written.
func (s *ConnectionState) VetForWrite(ctx context.Context, _ db.Reader, _ db.OpType, _ ...db.Option) error {
const op = "session.(ConnectionState).VetForWrite"
if err := s.validate(ctx); err != nil {
return errors.Wrap(ctx, err, op)
}
return nil
}
// TableName returns the tablename to override the default gorm table name
func (s *ConnectionState) TableName() string {
if s.tableName != "" {
return s.tableName
}
return defaultConnectionStateTableName
}
// 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 *ConnectionState) SetTableName(n string) {
s.tableName = n
}
// validate checks the session state
func (s *ConnectionState) validate(ctx context.Context) error {
const op = "session.(ConnectionState).validate"
if s.Status == "" {
return errors.New(ctx, errors.InvalidParameter, op, "missing status")
}
if s.ConnectionId == "" {
return errors.New(ctx, errors.InvalidParameter, op, "missing connection id")
}
if s.StartTime != nil {
return errors.New(ctx, errors.InvalidParameter, op, "start time is not settable")
}
if s.EndTime != nil {
return errors.New(ctx, errors.InvalidParameter, op, "end time is not settable")
}
if s.PreviousEndTime != nil {
return errors.New(ctx, errors.InvalidParameter, op, "previous end time is not settable")
}
return nil
func ConnectionStatusFromString(s string) ConnectionStatus {
switch s {
case "authorized":
return StatusAuthorized
case "connected":
return StatusConnected
case "closed":
return StatusClosed
}
return StatusUnspecified
}
internal/session/connection_state_test.go
-217
View File
@@ -1,217 +0,0 @@
// Copyright (c) HashiCorp, Inc.
// SPDX-License-Identifier: BUSL-1.1
package session
import (
"context"
"testing"
"github.com/hashicorp/boundary/internal/db"
"github.com/hashicorp/boundary/internal/errors"
"github.com/hashicorp/boundary/internal/iam"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestConnectionState_Create(t *testing.T) {
t.Parallel()
ctx := context.Background()
conn, _ := db.TestSetup(t, "postgres")
wrapper := db.TestWrapper(t)
iamRepo := iam.TestRepo(t, conn, wrapper)
session := TestDefaultSession(t, conn, wrapper, iamRepo)
connection := TestConnection(t, conn, session.PublicId, "127.0.0.1", 443, "127.0.0.1", 4443, "127.0.0.1")
type args struct {
connectionId string
status ConnectionStatus
}
tests := []struct {
name string
args args
want *ConnectionState
wantErr bool
wantIsErr errors.Code
create bool
wantCreateErr bool
}{
{
name: "valid",
args: args{
connectionId: connection.PublicId,
status: StatusClosed,
},
want: &ConnectionState{
ConnectionId: connection.PublicId,
Status: StatusClosed,
},
create: true,
},
{
name: "empty-connectionId",
args: args{
status: StatusClosed,
},
wantErr: true,
wantIsErr: errors.InvalidParameter,
},
{
name: "empty-status",
args: args{
connectionId: connection.PublicId,
},
wantErr: true,
wantIsErr: errors.InvalidParameter,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
assert, require := assert.New(t), require.New(t)
got, err := NewConnectionState(ctx, tt.args.connectionId, tt.args.status)
if tt.wantErr {
require.Error(err)
assert.True(errors.Match(errors.T(tt.wantIsErr), err))
return
}
require.NoError(err)
assert.Equal(tt.want, got)
if tt.create {
err = db.New(conn).Create(ctx, got)
if tt.wantCreateErr {
assert.Error(err)
return
} else {
assert.NoError(err)
}
}
})
}
}
func TestConnectionState_Delete(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)
s := TestDefaultSession(t, conn, wrapper, iamRepo)
c := TestConnection(t, conn, s.PublicId, "127.0.0.1", 22, "127.0.0.1", 2222, "127.0.0.1")
c2 := TestConnection(t, conn, s.PublicId, "127.0.0.1", 22, "127.0.0.1", 2222, "127.0.0.1")
tests := []struct {
name string
state *ConnectionState
deleteConnectionStateId string
wantRowsDeleted int
wantErr bool
wantErrMsg string
}{
{
name: "valid",
state: TestConnectionState(t, conn, c.PublicId, StatusClosed),
wantErr: false,
wantRowsDeleted: 1,
},
{
name: "bad-id",
state: TestConnectionState(t, conn, c2.PublicId, StatusClosed),
deleteConnectionStateId: func() string {
id, err := db.NewPublicId(ctx, ConnectionStatePrefix)
require.NoError(t, err)
return id
}(),
wantErr: false,
wantRowsDeleted: 0,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
assert, require := assert.New(t), require.New(t)
var initialState ConnectionState
err := rw.LookupWhere(context.Background(), &initialState, "connection_id = ? and state = ?", []any{tt.state.ConnectionId, tt.state.Status})
require.NoError(err)
deleteState := allocConnectionState()
if tt.deleteConnectionStateId != "" {
deleteState.ConnectionId = tt.deleteConnectionStateId
} else {
deleteState.ConnectionId = tt.state.ConnectionId
}
deleteState.StartTime = initialState.StartTime
deletedRows, err := rw.Delete(ctx, &deleteState)
if tt.wantErr {
require.Error(err)
return
}
require.NoError(err)
if tt.wantRowsDeleted == 0 {
assert.Equal(tt.wantRowsDeleted, deletedRows)
return
}
assert.Equal(tt.wantRowsDeleted, deletedRows)
foundState := allocConnectionState()
err = rw.LookupWhere(ctx, &foundState, "connection_id = ? and start_time = ?", []any{tt.state.ConnectionId, initialState.StartTime})
require.Error(err)
assert.True(errors.IsNotFoundError(err))
})
}
}
func TestConnectionState_Clone(t *testing.T) {
t.Parallel()
conn, _ := db.TestSetup(t, "postgres")
wrapper := db.TestWrapper(t)
iamRepo := iam.TestRepo(t, conn, wrapper)
t.Run("valid", func(t *testing.T) {
assert := assert.New(t)
s := TestDefaultSession(t, conn, wrapper, iamRepo)
c := TestConnection(t, conn, s.PublicId, "127.0.0.1", 22, "127.0.0.1", 2222, "127.0.0.1")
state := TestConnectionState(t, conn, c.PublicId, StatusConnected)
cp := state.Clone()
assert.Equal(cp.(*ConnectionState), state)
})
t.Run("not-equal", func(t *testing.T) {
assert := assert.New(t)
s := TestDefaultSession(t, conn, wrapper, iamRepo)
c := TestConnection(t, conn, s.PublicId, "127.0.0.1", 22, "127.0.0.1", 2222, "127.0.0.1")
state := TestConnectionState(t, conn, c.PublicId, StatusConnected)
state2 := TestConnectionState(t, conn, c.PublicId, StatusConnected)
cp := state.Clone()
assert.NotEqual(cp.(*ConnectionState), state2)
})
}
func TestConnectionState_SetTableName(t *testing.T) {
t.Parallel()
defaultTableName := defaultConnectionStateTableName
tests := []struct {
name string
setNameTo string
want string
}{
{
name: "new-name",
setNameTo: "new-name",
want: "new-name",
},
{
name: "reset to default",
setNameTo: "",
want: defaultTableName,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
assert, require := assert.New(t), require.New(t)
def := allocConnectionState()
require.Equal(defaultTableName, def.TableName())
s := allocConnectionState()
s.SetTableName(tt.setNameTo)
assert.Equal(tt.want, s.TableName())
})
}
}
internal/session/immutable_fields_test.go
-94
View File
@@ -9,7 +9,6 @@ import (
"github.com/hashicorp/boundary/internal/db"
"github.com/hashicorp/boundary/internal/db/timestamp"
"github.com/hashicorp/boundary/internal/errors"
"github.com/hashicorp/boundary/internal/iam"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
@@ -236,96 +235,3 @@ func TestConnection_ImmutableFields(t *testing.T) {
})
}
}
func TestConnectionState_ImmutableFields(t *testing.T) {
t.Parallel()
conn, _ := db.TestSetup(t, "postgres")
rw := db.New(conn)
wrapper := db.TestWrapper(t)
iamRepo := iam.TestRepo(t, conn, wrapper)
ts := timestamp.Timestamp{Timestamp: &timestamppb.Timestamp{Seconds: 0, Nanos: 0}}
_, _ = iam.TestScopes(t, iam.TestRepo(t, conn, wrapper))
session := TestDefaultSession(t, conn, wrapper, iamRepo)
connection := TestConnection(t, conn, session.PublicId, "127.0.0.1", 22, "127.0.0.1", 2222, "127.0.0.1")
state := TestConnectionState(t, conn, connection.PublicId, StatusConnected)
var new ConnectionState
err := rw.LookupWhere(context.Background(), &new, "connection_id = ? and state = ?", []any{state.ConnectionId, state.Status})
require.NoError(t, err)
tests := []struct {
name string
update *ConnectionState
fieldMask []string
wantErrMatch *errors.Template
wantErrContains string
}{
{
name: "session_id",
update: func() *ConnectionState {
s := new.Clone().(*ConnectionState)
s.ConnectionId = "sc_thisIsNotAValidId"
return s
}(),
fieldMask: []string{"PublicId"},
},
{
name: "status",
update: func() *ConnectionState {
s := new.Clone().(*ConnectionState)
s.Status = "closed"
return s
}(),
fieldMask: []string{"Status"},
wantErrMatch: errors.T(errors.NotSpecificIntegrity),
wantErrContains: "immutable column",
},
{
name: "start time",
update: func() *ConnectionState {
s := new.Clone().(*ConnectionState)
s.StartTime = &ts
return s
}(),
fieldMask: []string{"StartTime"},
wantErrMatch: errors.T(errors.InvalidFieldMask),
wantErrContains: "parameter violation",
},
{
name: "previous_end_time",
update: func() *ConnectionState {
s := new.Clone().(*ConnectionState)
s.PreviousEndTime = &ts
return s
}(),
fieldMask: []string{"PreviousEndTime"},
wantErrMatch: errors.T(errors.NotSpecificIntegrity),
wantErrContains: "immutable column",
},
}
for _, tt := range tests {
tt := tt
t.Run(tt.name, func(t *testing.T) {
assert, require := assert.New(t), require.New(t)
orig := new.Clone()
err := rw.LookupWhere(context.Background(), orig, "connection_id = ? and start_time = ?", []any{new.ConnectionId, new.StartTime})
require.NoError(err)
rowsUpdated, err := rw.Update(context.Background(), tt.update, tt.fieldMask, nil, db.WithSkipVetForWrite(true))
require.Error(err)
assert.Equal(0, rowsUpdated)
if tt.wantErrMatch != nil {
assert.Truef(errors.Match(tt.wantErrMatch, err), "wanted error %s and got: %s", tt.wantErrMatch.Code, err.Error())
}
if tt.wantErrContains != "" {
assert.Contains(err.Error(), tt.wantErrContains)
}
after := new.Clone()
err = rw.LookupWhere(context.Background(), after, "connection_id = ? and start_time = ?", []any{new.ConnectionId, new.StartTime})
require.NoError(err)
assert.Equal(orig.(*ConnectionState), after)
})
}
}
internal/session/job_session_cleanup_test.go
+33 -72
View File
@@ -65,10 +65,9 @@ func TestSessionConnectionCleanupJob(t *testing.T) {
sess := TestDefaultSession(t, conn, wrapper, iamRepo, WithDbOpts(db.WithSkipVetForWrite(true)))
sess, _, err = sessionRepo.ActivateSession(ctx, sess.GetPublicId(), sess.Version, []byte("foo"))
require.NoError(err)
c, cs, _, err := AuthorizeConnection(ctx, sessionRepo, connectionRepo, sess.GetPublicId(), serverId)
c, _, err := AuthorizeConnection(ctx, sessionRepo, connectionRepo, sess.GetPublicId(), serverId)
require.NoError(err)
require.Len(cs, 1)
require.Equal(StatusAuthorized, cs[0].Status)
require.Equal(StatusAuthorized, ConnectionStatusFromString(c.Status))
connIds = append(connIds, c.GetPublicId())
if i%2 == 0 {
connIdsByWorker[worker2.PublicId] = append(connIdsByWorker[worker2.PublicId], c.GetPublicId())
@@ -81,7 +80,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 := connectionRepo.ConnectConnection(ctx, ConnectWith{
cc, err := connectionRepo.ConnectConnection(ctx, ConnectWith{
ConnectionId: connId,
ClientTcpAddress: "127.0.0.1",
ClientTcpPort: 22,
@@ -90,18 +89,7 @@ func TestSessionConnectionCleanupJob(t *testing.T) {
UserClientIp: "127.0.0.1",
})
require.NoError(err)
require.Len(cs, 2)
var foundAuthorized, foundConnected bool
for _, status := range cs {
if status.Status == StatusAuthorized {
foundAuthorized = true
}
if status.Status == StatusConnected {
foundConnected = true
}
}
require.True(foundAuthorized)
require.True(foundConnected)
require.Equal(StatusConnected, ConnectionStatusFromString(cc.Status))
}
}
@@ -130,14 +118,11 @@ func TestSessionConnectionCleanupJob(t *testing.T) {
require.True(ok)
require.Len(connIds, 6)
for _, connId := range connIds {
_, states, err := connectionRepo.LookupConnection(ctx, connId)
conn, err := connectionRepo.LookupConnection(ctx, connId)
require.NoError(err)
var foundClosed bool
for _, state := range states {
if state.Status == StatusClosed {
foundClosed = true
break
}
if ConnectionStatusFromString(conn.Status) == StatusClosed {
foundClosed = true
}
assert.Equal(closed, foundClosed)
}
@@ -239,10 +224,9 @@ func TestCloseConnectionsForDeadWorkers(t *testing.T) {
sess := TestDefaultSession(t, conn, wrapper, iamRepo, WithDbOpts(db.WithSkipVetForWrite(true)))
sess, _, err = repo.ActivateSession(ctx, sess.GetPublicId(), sess.Version, []byte("foo"))
require.NoError(err)
c, cs, err := connRepo.AuthorizeConnection(ctx, sess.GetPublicId(), serverId)
c, err := connRepo.AuthorizeConnection(ctx, sess.GetPublicId(), serverId)
require.NoError(err)
require.Len(cs, 1)
require.Equal(StatusAuthorized, cs[0].Status)
require.Equal(StatusAuthorized, ConnectionStatusFromString(c.Status))
if i%3 == 0 {
worker1ConnIds = append(worker1ConnIds, c.GetPublicId())
} else if i%3 == 1 {
@@ -263,7 +247,7 @@ func TestCloseConnectionsForDeadWorkers(t *testing.T) {
return s
}() {
if i%3 == 0 {
_, cs, err := connRepo.ConnectConnection(ctx, ConnectWith{
cc, err := connRepo.ConnectConnection(ctx, ConnectWith{
ConnectionId: connId,
ClientTcpAddress: "127.0.0.1",
ClientTcpPort: 22,
@@ -272,18 +256,7 @@ func TestCloseConnectionsForDeadWorkers(t *testing.T) {
UserClientIp: "127.0.0.1",
})
require.NoError(err)
require.Len(cs, 2)
var foundAuthorized, foundConnected bool
for _, status := range cs {
if status.Status == StatusAuthorized {
foundAuthorized = true
}
if status.Status == StatusConnected {
foundConnected = true
}
}
require.True(foundAuthorized)
require.True(foundConnected)
require.Equal(StatusConnected, ConnectionStatusFromString(cc.Status))
} else if i%3 == 1 {
resp, err := connRepo.closeConnections(ctx, []CloseWith{
{
@@ -293,19 +266,8 @@ func TestCloseConnectionsForDeadWorkers(t *testing.T) {
})
require.NoError(err)
require.Len(resp, 1)
cs := resp[0].ConnectionStates
require.Len(cs, 2)
var foundAuthorized, foundClosed bool
for _, status := range cs {
if status.Status == StatusAuthorized {
foundAuthorized = true
}
if status.Status == StatusClosed {
foundClosed = true
}
}
require.True(foundAuthorized)
require.True(foundClosed)
cs := resp[0].ConnectionState
require.Equal(StatusClosed, cs)
}
}
@@ -344,9 +306,9 @@ func TestCloseConnectionsForDeadWorkers(t *testing.T) {
expected = StatusAuthorized
}
_, states, err := connRepo.LookupConnection(ctx, connId)
conn, err := connRepo.LookupConnection(ctx, connId)
require.NoError(err)
require.Equal(expected, states[0].Status, "expected latest status for %q (index %d) to be %v", connId, i, expected)
require.Equal(expected, ConnectionStatusFromString(conn.Status), "expected latest status for %q (index %d) to be %v", connId, i, expected)
}
}
@@ -480,10 +442,9 @@ func TestCloseWorkerlessConnections(t *testing.T) {
sess, _, err = repo.ActivateSession(ctx, sess.GetPublicId(), sess.Version, []byte("foo"))
require.NoError(err)
conn, cs, err := connRepo.AuthorizeConnection(ctx, sess.GetPublicId(), workerId)
conn, err := connRepo.AuthorizeConnection(ctx, sess.GetPublicId(), workerId)
require.NoError(err)
require.Len(cs, 1)
require.Equal(StatusAuthorized, cs[0].Status)
require.Equal(StatusAuthorized, ConnectionStatusFromString(conn.Status))
return conn
}
@@ -513,21 +474,21 @@ func TestCloseWorkerlessConnections(t *testing.T) {
}})
require.NoError(err)
_, st, err := connRepo.LookupConnection(ctx, dActiveConn.GetPublicId())
con, err := connRepo.LookupConnection(ctx, dActiveConn.GetPublicId())
require.NoError(err)
require.Equal(StatusAuthorized, st[0].Status)
require.Equal(StatusAuthorized, ConnectionStatusFromString(con.Status))
_, st, err = connRepo.LookupConnection(ctx, dClosedConn.GetPublicId())
con, err = connRepo.LookupConnection(ctx, dClosedConn.GetPublicId())
require.NoError(err)
require.Equal(StatusClosed, st[0].Status)
require.Equal(StatusClosed, ConnectionStatusFromString(con.Status))
_, st, err = connRepo.LookupConnection(ctx, activeConn.GetPublicId())
con, err = connRepo.LookupConnection(ctx, activeConn.GetPublicId())
require.NoError(err)
require.Equal(StatusAuthorized, st[0].Status)
require.Equal(StatusAuthorized, ConnectionStatusFromString(con.Status))
_, st, err = connRepo.LookupConnection(ctx, closedConn.GetPublicId())
con, err = connRepo.LookupConnection(ctx, closedConn.GetPublicId())
require.NoError(err)
require.Equal(StatusClosed, st[0].Status)
require.Equal(StatusClosed, ConnectionStatusFromString(con.Status))
// Run the job
numClosed, err := job.closeWorkerlessConnections(ctx)
@@ -535,19 +496,19 @@ func TestCloseWorkerlessConnections(t *testing.T) {
assert.Equal(t, 1, numClosed)
// This is the only one that the job should have actually closed.
_, st, err = connRepo.LookupConnection(ctx, dActiveConn.GetPublicId())
con, err = connRepo.LookupConnection(ctx, dActiveConn.GetPublicId())
require.NoError(err)
require.Equal(StatusClosed, st[0].Status)
require.Equal(StatusClosed, ConnectionStatusFromString(con.Status))
_, st, err = connRepo.LookupConnection(ctx, dClosedConn.GetPublicId())
con, err = connRepo.LookupConnection(ctx, dClosedConn.GetPublicId())
require.NoError(err)
require.Equal(StatusClosed, st[0].Status)
require.Equal(StatusClosed, ConnectionStatusFromString(con.Status))
_, st, err = connRepo.LookupConnection(ctx, activeConn.GetPublicId())
con, err = connRepo.LookupConnection(ctx, activeConn.GetPublicId())
require.NoError(err)
require.Equal(StatusAuthorized, st[0].Status)
require.Equal(StatusAuthorized, ConnectionStatusFromString(con.Status))
_, st, err = connRepo.LookupConnection(ctx, closedConn.GetPublicId())
con, err = connRepo.LookupConnection(ctx, closedConn.GetPublicId())
require.NoError(err)
require.Equal(StatusClosed, st[0].Status)
require.Equal(StatusClosed, ConnectionStatusFromString(con.Status))
}
internal/session/query.go
+26 -35
View File
@@ -124,6 +124,14 @@ from
session_connection_limit, session_connection_count;
`
// connectConnection sets the connected time range to (now, infinity) to
// indicate the connection is connected.
connectConnection = `
update session_connection
set connected_time_range=tstzrange(now(),'infinity')
where public_id=@public_id
`
terminateSessionIfPossible = `
-- is terminate_session_id in a canceling state
with session_version as (
@@ -197,17 +205,11 @@ from
select
session_id
from
session_connection
where public_id in (
select
connection_id
from
session_connection_state
where
state != 'closed' and
end_time is null
)
)
session_connection
where
upper(connected_time_range) > now() or
connected_time_range is null
)
`
// termSessionUpdate is one stmt that terminates sessions for the following
@@ -271,21 +273,12 @@ where
)
) and
-- make sure there are no existing connections
us.public_id not in (
select
session_id
from
session_connection
where public_id in (
select
connection_id
from
session_connection_state
where
state != 'closed' and
end_time is null
)
);
us.public_id not in (
select session_id
from session_connection
where upper(connected_time_range) > now()
or connected_time_range is null
);
`
// closeConnectionsForDeadServersCte finds connections that are:
@@ -336,22 +329,20 @@ where
and closed_reason is null
returning public_id;
`
orphanedConnectionsCte = `
-- Find connections that are not closed so we can reference those IDs
with
unclosed_connections as (
select connection_id
from session_connection_state
select public_id
from session_connection
where
-- It's the current state
end_time is null
-- Current state isn't closed state
and state in ('authorized', 'connected')
-- It's not closed
upper(connected_time_range) > now() or
connected_time_range is null
-- It's not in limbo between when it moved into this state and when
-- it started being reported by the worker, which is roughly every
-- 2-3 seconds
and start_time < wt_sub_seconds_from_now(@worker_state_delay_seconds)
and update_time < wt_sub_seconds_from_now(@worker_state_delay_seconds)
),
connections_to_close as (
select public_id
@@ -360,7 +351,7 @@ with
-- Related to the worker that just reported to us
worker_id = @worker_id
-- Only unclosed ones
and public_id in (select connection_id from unclosed_connections)
and public_id in (select public_id from unclosed_connections)
-- These are connection IDs that just got reported to us by the given
-- worker, so they should not be considered closed.
%s
internal/session/repository_connection.go
+30 -55
View File
@@ -134,19 +134,18 @@ func (r *ConnectionRepository) updateBytesUpBytesDown(ctx context.Context, conns
// If authorization is success, it creates/stores a new connection in the repo
// and returns it, along with its states. If the authorization fails, it
// an error with Code InvalidSessionState.
func (r *ConnectionRepository) AuthorizeConnection(ctx context.Context, sessionId, workerId string) (*Connection, []*ConnectionState, error) {
func (r *ConnectionRepository) AuthorizeConnection(ctx context.Context, sessionId, workerId string) (*Connection, error) {
const op = "session.(ConnectionRepository).AuthorizeConnection"
if sessionId == "" {
return nil, nil, errors.Wrap(ctx, status.Error(codes.FailedPrecondition, "missing session id"), op, errors.WithCode(errors.InvalidParameter))
return nil, errors.Wrap(ctx, status.Error(codes.FailedPrecondition, "missing session id"), op, errors.WithCode(errors.InvalidParameter))
}
connectionId, err := newConnectionId(ctx)
if err != nil {
return nil, nil, errors.Wrap(ctx, err, op)
return nil, errors.Wrap(ctx, err, op)
}
connection := AllocConnection()
connection.PublicId = connectionId
var connectionStates []*ConnectionState
_, err = r.writer.DoTx(
ctx,
db.StdRetryCnt,
@@ -166,31 +165,26 @@ func (r *ConnectionRepository) AuthorizeConnection(ctx context.Context, sessionI
if err := reader.LookupById(ctx, &connection); err != nil {
return errors.Wrap(ctx, err, op, errors.WithMsg(fmt.Sprintf("failed for session %s", sessionId)))
}
connectionStates, err = fetchConnectionStates(ctx, reader, 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 nil, errors.Wrap(ctx, err, op)
}
return &connection, connectionStates, nil
return &connection, nil
}
// 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.
// with its state. If the connection is not found, it will return nil, nil.
// No options are currently supported.
func (r *ConnectionRepository) LookupConnection(ctx context.Context, connectionId string, _ ...Option) (*Connection, []*ConnectionState, error) {
func (r *ConnectionRepository) LookupConnection(ctx context.Context, connectionId string, _ ...Option) (*Connection, error) {
const op = "session.(ConnectionRepository).LookupConnection"
if connectionId == "" {
return nil, nil, errors.New(ctx, errors.InvalidParameter, op, "missing connectionId id")
return nil, errors.New(ctx, errors.InvalidParameter, op, "missing connectionId id")
}
connection := AllocConnection()
connection.PublicId = connectionId
var states []*ConnectionState
_, err := r.writer.DoTx(
ctx,
db.StdRetryCnt,
@@ -199,20 +193,16 @@ func (r *ConnectionRepository) LookupConnection(ctx context.Context, connectionI
if err := read.LookupById(ctx, &connection); err != nil {
return errors.Wrap(ctx, err, op, errors.WithMsg(fmt.Sprintf("failed for %s", connectionId)))
}
var err error
if states, err = fetchConnectionStates(ctx, read, connectionId, db.WithOrder("start_time desc")); err != nil {
return errors.Wrap(ctx, err, op)
}
return nil
},
)
if err != nil {
if errors.IsNotFoundError(err) {
return nil, nil, nil
return nil, nil
}
return nil, nil, errors.Wrap(ctx, err, op)
return nil, errors.Wrap(ctx, err, op)
}
return &connection, states, nil
return &connection, nil
}
// ListConnectionsBySessionId will list connections by session ID. Supports the
@@ -231,14 +221,13 @@ func (r *ConnectionRepository) ListConnectionsBySessionId(ctx context.Context, s
}
// ConnectConnection updates a connection in the repo with a state of "connected".
func (r *ConnectionRepository) ConnectConnection(ctx context.Context, c ConnectWith) (*Connection, []*ConnectionState, error) {
func (r *ConnectionRepository) ConnectConnection(ctx context.Context, c ConnectWith) (*Connection, error) {
const op = "session.(ConnectionRepository).ConnectConnection"
// ConnectWith.validate will check all the fields...
if err := c.validate(ctx); err != nil {
return nil, nil, errors.Wrap(ctx, err, op)
return nil, errors.Wrap(ctx, err, op)
}
var connection Connection
var connectionStates []*ConnectionState
_, err := r.writer.DoTx(
ctx,
db.StdRetryCnt,
@@ -266,31 +255,33 @@ func (r *ConnectionRepository) ConnectConnection(ctx context.Context, c ConnectW
// 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(ctx, connection.PublicId, StatusConnected)
// Set the lower bound of the connected_time_range to indicate the connection is connected
rowsUpdated, err = w.Exec(ctx, connectConnection, []any{
sql.Named("public_id", c.ConnectionId),
})
if err != nil {
return errors.Wrap(ctx, err, op)
}
if err := w.Create(ctx, newState); err != nil {
return errors.Wrap(ctx, err, op)
if rowsUpdated != 1 {
return errors.Wrap(ctx, err, op, errors.WithMsg(fmt.Sprintf("unable to connect connection %s", c.ConnectionId)))
}
connectionStates, err = fetchConnectionStates(ctx, reader, c.ConnectionId, db.WithOrder("start_time desc"))
if err != nil {
return errors.Wrap(ctx, err, op)
if err := reader.LookupById(ctx, &connection); err != nil {
return errors.Wrap(ctx, err, op, errors.WithMsg(fmt.Sprintf("failed for connection %s", c.ConnectionId)))
}
return nil
},
)
if err != nil {
return nil, nil, errors.Wrap(ctx, err, op)
return nil, errors.Wrap(ctx, err, op)
}
return &connection, connectionStates, nil
return &connection, nil
}
// closeConnectionResp is just a wrapper for the response from CloseConnections.
// It wraps the connection and its states for each connection closed.
// It wraps the connection and its state for each connection closed.
type closeConnectionResp struct {
Connection *Connection
ConnectionStates []*ConnectionState
Connection *Connection
ConnectionState ConnectionStatus
}
// closeConnections set's a connection's state to "closed" in the repo. It's
@@ -318,8 +309,8 @@ func (r *ConnectionRepository) closeConnections(ctx context.Context, closeWith [
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.
// updating the ClosedReason will trigger the session_connection to set the
// upper limit of connection_time_range to indicate the connection is closed.
rowsUpdated, err := w.Update(
ctx,
&updateConnection,
@@ -332,13 +323,9 @@ func (r *ConnectionRepository) closeConnections(ctx context.Context, closeWith [
if rowsUpdated != 1 {
return errors.New(ctx, errors.MultipleRecords, op, fmt.Sprintf("%d would have been updated for connection %s", rowsUpdated, cw.ConnectionId))
}
states, err := fetchConnectionStates(ctx, reader, cw.ConnectionId, db.WithOrder("start_time desc"))
if err != nil {
return errors.Wrap(ctx, err, op)
}
resp = append(resp, closeConnectionResp{
Connection: &updateConnection,
ConnectionStates: states,
Connection: &updateConnection,
ConnectionState: StatusClosed,
})
}
@@ -441,15 +428,3 @@ func (r *ConnectionRepository) closeOrphanedConnections(ctx context.Context, wor
}
return orphanedConns, nil
}
func fetchConnectionStates(ctx context.Context, r db.Reader, connectionId string, opt ...db.Option) ([]*ConnectionState, error) {
const op = "session.fetchConnectionStates"
var states []*ConnectionState
if err := r.SearchWhere(ctx, &states, "connection_id = ?", []any{connectionId}, opt...); err != nil {
return nil, errors.Wrap(ctx, err, op)
}
if len(states) == 0 {
return nil, nil
}
return states, nil
}
internal/session/repository_connection_test.go
+107 -25
View File
@@ -240,7 +240,7 @@ func TestRepository_ConnectConnection(t *testing.T) {
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)
c, 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())
@@ -248,9 +248,8 @@ func TestRepository_ConnectConnection(t *testing.T) {
}
require.NoError(err)
require.NotNil(c)
require.NotNil(cs)
assert.Equal(StatusConnected, cs[0].Status)
gotConn, _, err := connRepo.LookupConnection(context.Background(), c.PublicId)
assert.Equal(StatusConnected, ConnectionStatusFromString(c.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)
@@ -336,7 +335,7 @@ func TestRepository_DeleteConnection(t *testing.T) {
}
assert.NoError(err)
assert.Equal(tt.wantRowsDeleted, deletedRows)
found, _, err := connRepo.LookupConnection(context.Background(), tt.args.connection.PublicId)
found, err := connRepo.LookupConnection(context.Background(), tt.args.connection.PublicId)
assert.NoError(err)
assert.Nil(found)
@@ -378,10 +377,9 @@ func TestRepository_orphanedConnections(t *testing.T) {
sess := TestDefaultSession(t, conn, wrapper, iamRepo, WithDbOpts(db.WithSkipVetForWrite(true)))
sess, _, err = repo.ActivateSession(ctx, sess.GetPublicId(), sess.Version, []byte("foo"))
require.NoError(err)
c, cs, err := connRepo.AuthorizeConnection(ctx, sess.GetPublicId(), serverId)
c, err := connRepo.AuthorizeConnection(ctx, sess.GetPublicId(), serverId)
require.NoError(err)
require.Len(cs, 1)
require.Equal(StatusAuthorized, cs[0].Status)
require.Equal(StatusAuthorized, ConnectionStatusFromString(c.Status))
connIds = append(connIds, c.GetPublicId())
if i%2 == 0 {
worker2ConnIds = append(worker2ConnIds, c.GetPublicId())
@@ -394,7 +392,7 @@ func TestRepository_orphanedConnections(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 := connRepo.ConnectConnection(ctx, ConnectWith{
cc, err := connRepo.ConnectConnection(ctx, ConnectWith{
ConnectionId: connId,
ClientTcpAddress: "127.0.0.1",
ClientTcpPort: 22,
@@ -403,18 +401,7 @@ func TestRepository_orphanedConnections(t *testing.T) {
UserClientIp: "127.0.0.1",
})
require.NoError(err)
require.Len(cs, 2)
var foundAuthorized, foundConnected bool
for _, status := range cs {
if status.Status == StatusAuthorized {
foundAuthorized = true
}
if status.Status == StatusConnected {
foundConnected = true
}
}
require.True(foundAuthorized)
require.True(foundConnected)
require.Equal(StatusConnected, ConnectionStatusFromString(cc.Status))
}
}
@@ -517,8 +504,8 @@ func TestRepository_CloseConnections(t *testing.T) {
assert.Equal(len(tt.closeWith), len(resp))
for _, r := range resp {
require.NotNil(r.Connection)
require.NotNil(r.ConnectionStates)
assert.Equal(StatusClosed, r.ConnectionStates[0].Status)
require.NotNil(r.ConnectionState)
assert.Equal(StatusClosed, r.ConnectionState)
}
})
}
@@ -561,7 +548,7 @@ func TestUpdateBytesUpDown(t *testing.T) {
// Assert that the bytes up and down values have been persisted.
for i := 0; i < len(conns); i++ {
c, _, err := connRepo.LookupConnection(ctx, conns[i].GetPublicId())
c, err := connRepo.LookupConnection(ctx, conns[i].GetPublicId())
require.NoError(t, err)
require.Equal(t, conns[i].BytesUp, c.BytesUp)
@@ -604,10 +591,105 @@ func TestUpdateBytesUpDown(t *testing.T) {
// BytesUp and BytesDown values should be set to the old ones.
for i := 0; i < len(conns); i++ {
c, _, err := connRepo.LookupConnection(ctx, conns[i].GetPublicId())
c, err := connRepo.LookupConnection(ctx, conns[i].GetPublicId())
require.NoError(t, err)
require.Equal(t, conns[i].BytesUp, c.BytesUp)
require.Equal(t, conns[i].BytesDown, c.BytesDown)
}
}
func TestRepository_StateTransitions(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(ctx, rw, rw, kms)
require.NoError(t, err)
connRepo, err := NewConnectionRepository(ctx, rw, rw, kms)
require.NoError(t, err)
s := TestDefaultSession(t, conn, wrapper, iamRepo)
tofu := TestTofu(t)
_, _, err = repo.ActivateSession(context.Background(), s.PublicId, s.Version, tofu)
require.NoError(t, err)
// First connection will transition authorized -> connected -> closed
c := TestConnection(t, conn, s.PublicId, "127.0.0.1", 22, "127.0.0.1", 2222, "127.0.0.1")
cw := ConnectWith{
ConnectionId: c.PublicId,
ClientTcpAddress: "127.0.0.1",
ClientTcpPort: 22,
EndpointTcpAddress: "127.0.0.1",
EndpointTcpPort: 2222,
UserClientIp: "127.0.0.1",
}
gotConn, err := connRepo.LookupConnection(context.Background(), c.PublicId)
require.NoError(t, err)
require.NotNil(t, gotConn)
require.Equal(t, StatusAuthorized, ConnectionStatusFromString(gotConn.Status))
_, err = connRepo.ConnectConnection(context.Background(), cw)
require.NoError(t, err)
gotConn, err = connRepo.LookupConnection(context.Background(), c.PublicId)
require.NoError(t, err)
require.NotNil(t, gotConn)
require.Equal(t, StatusConnected, ConnectionStatusFromString(gotConn.Status))
// Attempt to connect again, expect failure
_, err = connRepo.ConnectConnection(context.Background(), cw)
require.Error(t, err)
require.Contains(t, err.Error(), "Invalid state transition from connected")
closeWith := CloseWith{
ConnectionId: c.PublicId,
ClosedReason: ConnectionClosedByUser,
}
resp, err := connRepo.closeConnections(context.Background(), []CloseWith{closeWith})
require.NoError(t, err)
require.NotNil(t, resp)
require.Equal(t, StatusClosed, resp[0].ConnectionState)
// Second connection will transition from authorized -> closed
c2 := TestConnection(t, conn, s.PublicId, "127.0.0.1", 22, "127.0.0.1", 2222, "127.0.0.1")
gotConn, err = connRepo.LookupConnection(context.Background(), c2.PublicId)
require.NoError(t, err)
require.NotNil(t, gotConn)
require.Equal(t, StatusAuthorized, ConnectionStatusFromString(gotConn.Status))
closeWith2 := CloseWith{
ConnectionId: c2.PublicId,
ClosedReason: ConnectionClosedByUser,
}
resp, err = connRepo.closeConnections(context.Background(), []CloseWith{closeWith2})
require.NoError(t, err)
require.NotNil(t, resp)
require.Equal(t, StatusClosed, resp[0].ConnectionState)
gotConn, err = connRepo.LookupConnection(context.Background(), c2.PublicId)
require.NoError(t, err)
require.NotNil(t, gotConn)
require.Equal(t, StatusClosed, ConnectionStatusFromString(gotConn.Status))
// Now try to connect it while closed and ensure it can't transition to connected
cw2 := ConnectWith{
ConnectionId: c2.PublicId,
ClientTcpAddress: "127.0.0.1",
ClientTcpPort: 22,
EndpointTcpAddress: "127.0.0.1",
EndpointTcpPort: 2222,
UserClientIp: "127.0.0.1",
}
_, err = connRepo.ConnectConnection(context.Background(), cw2)
require.Error(t, err)
require.Contains(t, err.Error(), "Invalid state transition from closed")
gotConn, err = connRepo.LookupConnection(context.Background(), c2.PublicId)
require.NoError(t, err)
require.NotNil(t, gotConn)
require.Equal(t, StatusClosed, ConnectionStatusFromString(gotConn.Status))
}
internal/session/repository_session_test.go
+4 -8
View File
@@ -1149,12 +1149,10 @@ func TestRepository_TerminateCompletedSessions(t *testing.T) {
conn, err := connRepo.ListConnectionsBySessionId(context.Background(), found.PublicId)
require.NoError(err)
for _, sc := range conn {
c, cs, err := connRepo.LookupConnection(context.Background(), sc.PublicId)
c, err := connRepo.LookupConnection(context.Background(), sc.PublicId)
require.NoError(err)
assert.NotEmpty(c.ClosedReason)
for _, s := range cs {
t.Logf("%s session %s connection state %s at %s", found.PublicId, s.ConnectionId, s.Status, s.EndTime)
}
t.Logf("%s session connection state %s", found.PublicId, c.Status)
}
} else {
t.Logf("not terminated %s has a connection limit of %d", found.PublicId, found.ConnectionLimit)
@@ -1162,11 +1160,9 @@ func TestRepository_TerminateCompletedSessions(t *testing.T) {
conn, err := connRepo.ListConnectionsBySessionId(context.Background(), found.PublicId)
require.NoError(err)
for _, sc := range conn {
cs, err := fetchConnectionStates(context.Background(), rw, sc.PublicId)
c, err := connRepo.LookupConnection(context.Background(), sc.PublicId)
require.NoError(err)
for _, s := range cs {
t.Logf("%s session %s connection state %s at %s", found.PublicId, s.ConnectionId, s.Status, s.EndTime)
}
t.Logf("%s session connection state %s", found.PublicId, c.Status)
}
}
}
internal/session/service_authorize_connection.go
+5 -5
View File
@@ -17,17 +17,17 @@ import (
// If any of these criteria is not met, it returns an error with Code InvalidSessionState.
func AuthorizeConnection(ctx context.Context, sessionRepoFn *Repository, connectionRepoFn *ConnectionRepository,
sessionId, workerId string, opt ...Option,
) (*Connection, []*ConnectionState, *AuthzSummary, error) {
) (*Connection, *AuthzSummary, error) {
const op = "session.AuthorizeConnection"
connection, connectionStates, err := connectionRepoFn.AuthorizeConnection(ctx, sessionId, workerId)
connection, err := connectionRepoFn.AuthorizeConnection(ctx, sessionId, workerId)
if err != nil {
return nil, nil, nil, errors.Wrap(ctx, err, op)
return nil, nil, errors.Wrap(ctx, err, op)
}
authzSummary, err := sessionRepoFn.sessionAuthzSummary(ctx, sessionId)
if err != nil {
return nil, nil, nil, errors.Wrap(ctx, err, op)
return nil, nil, errors.Wrap(ctx, err, op)
}
return connection, connectionStates, authzSummary, nil
return connection, authzSummary, nil
}
internal/session/service_authorize_connection_test.go
+3 -3
View File
@@ -138,7 +138,7 @@ func TestService_AuthorizeConnection(t *testing.T) {
t.Run(tt.name, func(t *testing.T) {
assert, require := assert.New(t), require.New(t)
c, cs, authzInfo, err := AuthorizeConnection(context.Background(), repo, connRepo, tt.session.PublicId, testServer)
c, authzInfo, err := AuthorizeConnection(context.Background(), repo, connRepo, tt.session.PublicId, testServer)
if tt.wantErr {
require.Error(err)
// TODO (jimlambrt 9/2020): add in tests for errorsIs once we
@@ -150,8 +150,8 @@ func TestService_AuthorizeConnection(t *testing.T) {
}
require.NoError(err)
require.NotNil(c)
require.NotNil(cs)
assert.Equal(StatusAuthorized, cs[0].Status)
require.NotNil(c.Status)
assert.Equal(StatusAuthorized, ConnectionStatusFromString(c.Status))
assert.True(authzInfo.ExpirationTime.GetTimestamp().AsTime().Sub(tt.wantAuthzInfo.ExpirationTime.GetTimestamp().AsTime()) < 10*time.Millisecond)
tt.wantAuthzInfo.ExpirationTime = authzInfo.ExpirationTime
internal/session/service_close_connections_test.go
+2 -2
View File
@@ -131,8 +131,8 @@ func TestServiceCloseConnections(t *testing.T) {
for _, r := range resp {
require.NotNil(r.Connection)
require.NotNil(r.ConnectionStates)
assert.Equal(StatusClosed, r.ConnectionStates[0].Status)
require.NotNil(r.ConnectionState)
assert.Equal(StatusClosed, r.ConnectionState)
}
// Ensure session is in the state we want- terminated if all conns closed, else active
internal/session/service_worker_status_report_test.go
+12 -14
View File
@@ -94,7 +94,7 @@ func TestWorkerStatusReport(t *testing.T) {
require.NoError(t, err)
require.NoError(t, err)
_, _, err = connRepo.AuthorizeConnection(ctx, sess.PublicId, worker.PublicId)
_, err = connRepo.AuthorizeConnection(ctx, sess.PublicId, worker.PublicId)
require.NoError(t, err)
_, err = repo.CancelSession(ctx, sess.PublicId, sess.Version)
@@ -126,7 +126,7 @@ func TestWorkerStatusReport(t *testing.T) {
require.NoError(t, err)
require.NoError(t, err)
connection, _, err := connRepo.AuthorizeConnection(ctx, sess.PublicId, worker.PublicId)
connection, err := connRepo.AuthorizeConnection(ctx, sess.PublicId, worker.PublicId)
require.NoError(t, err)
return testCase{
worker: worker,
@@ -160,7 +160,7 @@ func TestWorkerStatusReport(t *testing.T) {
tofu := session.TestTofu(t)
sess, _, err = repo.ActivateSession(ctx, sess.PublicId, sess.Version, tofu)
require.NoError(t, err)
connection, _, err := connRepo.AuthorizeConnection(ctx, sess.PublicId, worker.PublicId)
connection, err := connRepo.AuthorizeConnection(ctx, sess.PublicId, worker.PublicId)
require.NoError(t, err)
_, err = repo.CancelSession(ctx, sess.PublicId, sess.Version)
require.NoError(t, err)
@@ -224,7 +224,7 @@ func TestWorkerStatusReport(t *testing.T) {
tofu := session.TestTofu(t)
sess, _, err = repo.ActivateSession(ctx, sess.PublicId, sess.Version, tofu)
require.NoError(t, err)
connection, _, err := connRepo.AuthorizeConnection(ctx, sess.PublicId, worker.PublicId)
connection, err := connRepo.AuthorizeConnection(ctx, sess.PublicId, worker.PublicId)
require.NoError(t, err)
_, err = repo.CancelSession(ctx, sess.PublicId, sess.Version)
require.NoError(t, err)
@@ -242,7 +242,7 @@ func TestWorkerStatusReport(t *testing.T) {
tofu2 := session.TestTofu(t)
sess2, _, err = repo.ActivateSession(ctx, sess2.PublicId, sess2.Version, tofu2)
require.NoError(t, err)
connection2, _, err := connRepo.AuthorizeConnection(ctx, sess2.PublicId, worker.PublicId)
connection2, err := connRepo.AuthorizeConnection(ctx, sess2.PublicId, worker.PublicId)
require.NoError(t, err)
_, err = repo.CancelSession(ctx, sess2.PublicId, sess2.Version)
require.NoError(t, err)
@@ -295,7 +295,7 @@ func TestWorkerStatusReport(t *testing.T) {
tofu := session.TestTofu(t)
sess, _, err = repo.ActivateSession(ctx, sess.PublicId, sess.Version, tofu)
require.NoError(t, err)
connection, _, err := connRepo.AuthorizeConnection(ctx, sess.PublicId, worker.PublicId)
connection, err := connRepo.AuthorizeConnection(ctx, sess.PublicId, worker.PublicId)
require.NoError(t, err)
sess2 := session.TestSession(t, conn, wrapper, session.ComposedOf{
@@ -311,7 +311,7 @@ func TestWorkerStatusReport(t *testing.T) {
tofu2 := session.TestTofu(t)
sess2, _, err = repo.ActivateSession(ctx, sess2.PublicId, sess2.Version, tofu2)
require.NoError(t, err)
connection2, _, err := connRepo.AuthorizeConnection(ctx, sess2.PublicId, worker.PublicId)
connection2, err := connRepo.AuthorizeConnection(ctx, sess2.PublicId, worker.PublicId)
require.NoError(t, err)
require.NotEqual(t, connection.PublicId, connection2.PublicId)
@@ -348,7 +348,7 @@ func TestWorkerStatusReport(t *testing.T) {
tofu := session.TestTofu(t)
sess, _, err = repo.ActivateSession(ctx, sess.PublicId, sess.Version, tofu)
require.NoError(t, err)
connection, _, err := connRepo.AuthorizeConnection(ctx, sess.PublicId, worker.PublicId)
connection, err := connRepo.AuthorizeConnection(ctx, sess.PublicId, worker.PublicId)
require.NoError(t, err)
_, err = repo.CancelSession(ctx, sess.PublicId, sess.Version)
require.NoError(t, err)
@@ -366,9 +366,9 @@ func TestWorkerStatusReport(t *testing.T) {
tofu2 := session.TestTofu(t)
sess2, _, err = repo.ActivateSession(ctx, sess2.PublicId, sess2.Version, tofu2)
require.NoError(t, err)
connection2, _, err := connRepo.AuthorizeConnection(ctx, sess2.PublicId, worker.PublicId)
connection2, err := connRepo.AuthorizeConnection(ctx, sess2.PublicId, worker.PublicId)
require.NoError(t, err)
connection3, _, err := connRepo.AuthorizeConnection(ctx, sess2.PublicId, worker.PublicId)
connection3, err := connRepo.AuthorizeConnection(ctx, sess2.PublicId, worker.PublicId)
require.NoError(t, err)
_, err = repo.CancelSession(ctx, sess2.PublicId, sess2.Version)
require.NoError(t, err)
@@ -417,12 +417,10 @@ func TestWorkerStatusReport(t *testing.T) {
require.NoError(err)
assert.ElementsMatch(tc.want, got)
for _, dc := range tc.orphanedConnections {
gotConn, states, err := connRepo.LookupConnection(ctx, dc)
gotConn, err := connRepo.LookupConnection(ctx, dc)
require.NoError(err)
assert.Equal(session.ConnectionSystemError, session.ClosedReason(gotConn.ClosedReason))
assert.Equal(2, len(states))
assert.Nil(states[0].EndTime)
assert.Equal(session.StatusClosed, states[0].Status)
assert.Equal(session.StatusClosed, session.ConnectionStatusFromString(gotConn.Status))
}
})
}
internal/session/testing.go
-17
View File
@@ -41,26 +41,9 @@ func TestConnection(t testing.TB, conn *db.DB, sessionId, clientTcpAddr string,
err = rw.Create(ctx, c)
require.NoError(err)
connectedState, err := NewConnectionState(ctx, c.PublicId, StatusConnected)
require.NoError(err)
err = rw.Create(ctx, connectedState)
require.NoError(err)
return c
}
// TestConnectionState creates a test connection state for the connectionId in the repository.
func TestConnectionState(t testing.TB, conn *db.DB, connectionId string, state ConnectionStatus) *ConnectionState {
t.Helper()
ctx := context.Background()
require := require.New(t)
rw := db.New(conn)
s, err := NewConnectionState(ctx, connectionId, state)
require.NoError(err)
err = rw.Create(context.Background(), s)
require.NoError(err)
return s
}
// TestState creates a test state for the sessionId in the repository.
func TestState(t testing.TB, conn *db.DB, sessionId string, state Status) *State {
t.Helper()
internal/session/testing_test.go
-23
View File
@@ -50,29 +50,6 @@ func Test_TestConnection(t *testing.T) {
require.NotNil(c)
}
func Test_TestConnectionState(t *testing.T) {
assert, require := assert.New(t), require.New(t)
conn, _ := db.TestSetup(t, "postgres")
wrapper := db.TestWrapper(t)
iamRepo := iam.TestRepo(t, conn, wrapper)
s := TestDefaultSession(t, conn, wrapper, iamRepo)
require.NotNil(s)
assert.NotEmpty(s.PublicId)
c := TestConnection(t, conn, s.PublicId, "0.0.0.0", 22, "0.0.0.0", 2222, "127.0.0.1")
require.NotNil(c)
assert.NotEmpty(c.PublicId)
cs := TestConnectionState(t, conn, c.PublicId, StatusClosed)
require.NotNil(cs)
rw := db.New(conn)
var initialState ConnectionState
err := rw.LookupWhere(context.Background(), &initialState, "connection_id = ? and state = ?", []any{cs.ConnectionId, cs.Status})
require.NoError(err)
assert.NotEmpty(initialState.StartTime)
}
func Test_TestWorker(t *testing.T) {
require := require.New(t)
conn, _ := db.TestSetup(t, "postgres")