CyberCoveLLC/threshold
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
1f550722a2aadace5ffbf13fcb343465d38ed830
threshold/internal/session/repository_connection_test.go
T
Todd 472d7d520a Remove the server_id from session table, change it to worker_id on session_connection (#2070) 
* Remove the server_id from session table, change it to worker_id on session_connection

This is step 1 in updating the worker table.  The worker is becoming a public resource so it must have
a public id.  We drop the server_id column on the session_connection table and add the worker_id column
so we can use the wt_public_key type.  Since db migrations are applied while the controller is offline,
it should be safe to assume all session connections are not in a running state, thus allowing us to null
out all the worker ids in the session_connection table.

For now we don't have very good enforcement of the wt_public_id constraints, since the id is currently
defined in the worker config, but this will soon change and the worker's id will be generated by the
controller.
2022-05-16 15:14:12 -07:00

521 lines
15 KiB
Go
Raw Blame History

package session
import (
"context"
"fmt"
"testing"
"time"
"github.com/hashicorp/boundary/internal/db"
"github.com/hashicorp/boundary/internal/errors"
"github.com/hashicorp/boundary/internal/iam"
"github.com/hashicorp/boundary/internal/kms"
"github.com/hashicorp/boundary/internal/oplog"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestRepository_ListConnection(t *testing.T) {
t.Parallel()
ctx := context.Background()
conn, _ := db.TestSetup(t, "postgres")
const testLimit = 10
wrapper := db.TestWrapper(t)
iamRepo := iam.TestRepo(t, conn, wrapper)
rw := db.New(conn)
kms := kms.TestKms(t, conn, wrapper)
repo, err := NewRepository(rw, rw, kms, WithLimit(testLimit))
require.NoError(t, err)
connRepo, err := NewConnectionRepository(ctx, rw, rw, kms, WithLimit(testLimit))
require.NoError(t, err)
session := TestDefaultSession(t, conn, wrapper, iamRepo)
type args struct {
searchForSessionId string
opt []Option
}
tests := []struct {
name string
createCnt int
args args
wantCnt int
wantErr bool
}{
{
name: "no-limit",
createCnt: repo.defaultLimit + 1,
args: args{
searchForSessionId: session.PublicId,
opt: []Option{WithLimit(-1)},
},
wantCnt: repo.defaultLimit + 1,
wantErr: false,
},
{
name: "default-limit",
createCnt: repo.defaultLimit + 1,
args: args{
searchForSessionId: session.PublicId,
},
wantCnt: repo.defaultLimit,
wantErr: false,
},
{
name: "custom-limit",
createCnt: repo.defaultLimit + 1,
args: args{
searchForSessionId: session.PublicId,
opt: []Option{WithLimit(3)},
},
wantCnt: 3,
wantErr: false,
},
{
name: "bad-session-id",
createCnt: repo.defaultLimit + 1,
args: args{
searchForSessionId: "s_thisIsNotValid",
},
wantCnt: 0,
wantErr: false,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
assert, require := assert.New(t), require.New(t)
db.TestDeleteWhere(t, conn, func() interface{} { i := AllocConnection(); return &i }(), "1=1")
testConnections := []*Connection{}
for i := 0; i < tt.createCnt; i++ {
c := TestConnection(t, conn,
session.PublicId,
"127.0.0.1",
22,
"127.0.0.1",
2222,
"127.0.0.1",
)
testConnections = append(testConnections, c)
}
assert.Equal(tt.createCnt, len(testConnections))
got, err := connRepo.ListConnectionsBySessionId(context.Background(), tt.args.searchForSessionId, tt.args.opt...)
if tt.wantErr {
require.Error(err)
return
}
require.NoError(err)
assert.Equal(tt.wantCnt, len(got))
})
}
t.Run("withOrder", func(t *testing.T) {
assert, require := assert.New(t), require.New(t)
db.TestDeleteWhere(t, conn, func() interface{} { i := AllocConnection(); return &i }(), "1=1")
wantCnt := 5
for i := 0; i < wantCnt; i++ {
_ = TestConnection(t, conn,
session.PublicId,
"127.0.0.1",
22,
"127.0.0.1",
2222,
"127.0.0.1",
)
}
got, err := connRepo.ListConnectionsBySessionId(context.Background(), session.PublicId, WithOrderByCreateTime(db.AscendingOrderBy))
require.NoError(err)
assert.Equal(wantCnt, len(got))
for i := 0; i < len(got)-1; i++ {
first := got[i].CreateTime.Timestamp.AsTime()
second := got[i+1].CreateTime.Timestamp.AsTime()
assert.True(first.Before(second))
}
})
}
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)
require.NoError(t, err)
connRepo, err := NewConnectionRepository(ctx, rw, rw, kms)
require.NoError(t, err)
setupFn := func() ConnectWith {
s := TestDefaultSession(t, conn, wrapper, iamRepo)
tofu := TestTofu(t)
_, _, err := repo.ActivateSession(context.Background(), s.PublicId, s.Version, 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()
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)
connRepo, err := NewConnectionRepository(ctx, rw, rw, kms)
require.NoError(t, err)
session := TestDefaultSession(t, conn, wrapper, iamRepo)
type args struct {
connection *Connection
opt []Option
}
tests := []struct {
name string
args args
wantRowsDeleted int
wantErr bool
wantErrMsg string
}{
{
name: "valid",
args: args{
connection: TestConnection(t, conn, session.PublicId, "127.0.0.1", 22, "127.0.0.1", 2222, "127.0.0.1"),
},
wantRowsDeleted: 1,
wantErr: false,
},
{
name: "no-public-id",
args: args{
connection: func() *Connection {
c := AllocConnection()
return &c
}(),
},
wantRowsDeleted: 0,
wantErr: true,
wantErrMsg: "session.(ConnectionRepository).DeleteConnection: missing public id: parameter violation: error #100",
},
{
name: "not-found",
args: args{
connection: func() *Connection {
c := AllocConnection()
id, err := newConnectionId()
require.NoError(t, err)
c.PublicId = id
return &c
}(),
},
wantRowsDeleted: 0,
wantErr: true,
wantErrMsg: "db.LookupById: record not found",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
assert := assert.New(t)
deletedRows, err := connRepo.DeleteConnection(context.Background(), tt.args.connection.PublicId, tt.args.opt...)
if tt.wantErr {
assert.Error(err)
assert.Equal(0, deletedRows)
assert.Contains(err.Error(), tt.wantErrMsg)
err = db.TestVerifyOplog(t, rw, tt.args.connection.PublicId, db.WithOperation(oplog.OpType_OP_TYPE_DELETE), db.WithCreateNotBefore(10*time.Second))
assert.Error(err)
assert.True(errors.IsNotFoundError(err))
return
}
assert.NoError(err)
assert.Equal(tt.wantRowsDeleted, deletedRows)
found, _, err := connRepo.LookupConnection(context.Background(), tt.args.connection.PublicId)
assert.NoError(err)
assert.Nil(found)
err = db.TestVerifyOplog(t, rw, tt.args.connection.PublicId, db.WithOperation(oplog.OpType_OP_TYPE_DELETE), db.WithCreateNotBefore(10*time.Second))
assert.Error(err)
})
}
}
func TestRepository_orphanedConnections(t *testing.T) {
t.Parallel()
ctx := context.Background()
require, assert := require.New(t), assert.New(t)
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(err)
connRepo, err := NewConnectionRepository(ctx, rw, rw, kms, WithWorkerStateDelay(0))
require.NoError(err)
numConns := 12
// Create two "workers". One will remain untouched while the other "goes
// away and comes back" (worker 2).
worker1 := TestWorker(t, conn, wrapper)
worker2 := TestWorker(t, conn, wrapper)
// Create a few sessions on each, activate, and authorize a connection
var connIds []string
var worker1ConnIds []string
var worker2ConnIds []string
for i := 0; i < numConns; i++ {
serverId := worker1.PublicId
if i%2 == 0 {
serverId = worker2.PublicId
}
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)
require.NoError(err)
require.Len(cs, 1)
require.Equal(StatusAuthorized, cs[0].Status)
connIds = append(connIds, c.GetPublicId())
if i%2 == 0 {
worker2ConnIds = append(worker2ConnIds, c.GetPublicId())
} else {
worker1ConnIds = append(worker1ConnIds, c.GetPublicId())
}
}
// Mark half of the connections connected and leave the others authorized.
// 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{
ConnectionId: connId,
ClientTcpAddress: "127.0.0.1",
ClientTcpPort: 22,
EndpointTcpAddress: "127.0.0.1",
EndpointTcpPort: 22,
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)
}
}
// Now, advertise only some of the connection IDs for worker 2. After,
// all connection IDs for worker 1 should be showing as non-closed, and
// the ones for worker 2 not advertised should be closed.
shouldStayOpen := worker2ConnIds[0:2]
found, err := connRepo.closeOrphanedConnections(ctx, worker2.GetPublicId(), shouldStayOpen)
require.NoError(err)
fmt.Printf("shouldstate: %v\nfound: %v\n", shouldStayOpen, found)
require.Equal(4, len(found))
// For the ones we didn't specify, we expect those to now be closed. We
// expect all others to be open.
shouldBeFound := worker2ConnIds[2:]
assert.ElementsMatch(found, shouldBeFound)
// Now, advertise none of the connection IDs for worker 2. This is mainly to
// test that handling the case where we do not include IDs works properly as
// it changes the where clause.
found, err = connRepo.closeOrphanedConnections(ctx, worker1.GetPublicId(), nil)
require.NoError(err)
assert.Equal(6, len(found))
assert.ElementsMatch(found, worker1ConnIds)
}
func TestRepository_CloseConnections(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)
require.NoError(t, err)
connRepo, err := NewConnectionRepository(ctx, rw, rw, kms)
require.NoError(t, err)
setupFn := func(cnt int) []CloseWith {
s := TestDefaultSession(t, conn, wrapper, iamRepo)
tofu := TestTofu(t)
s, _, err := repo.ActivateSession(context.Background(), s.PublicId, s.Version, 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, "127.0.0.1")
require.NoError(t, err)
cw = append(cw, CloseWith{
ConnectionId: c.PublicId,
BytesUp: 1,
BytesDown: 2,
ClosedReason: ConnectionClosedByUser,
})
}
return cw
}
tests := []struct {
name string
closeWith []CloseWith
reason TerminationReason
wantErr bool
wantIsError errors.Code
}{
{
name: "valid",
closeWith: setupFn(2),
reason: ClosedByUser,
},
{
name: "empty-closed-with",
closeWith: []CloseWith{},
reason: ClosedByUser,
wantErr: true,
wantIsError: errors.InvalidParameter,
},
{
name: "missing-ConnectionId",
closeWith: func() []CloseWith {
cw := setupFn(2)
cw[1].ConnectionId = ""
return cw
}(),
reason: ClosedByUser,
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)
resp, err := connRepo.closeConnections(context.Background(), tt.closeWith)
if tt.wantErr {
require.Error(err)
assert.Truef(errors.Match(errors.T(tt.wantIsError), err), "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, r.ConnectionStates[0].Status)
}
})
}
}
Reference in new issue View Git Blame Copy Permalink