This commit adds some rudimentary functionality for cleaning up connections on controller failure, or more specifically, a failure to communicate with the controllers. In this scenario, when a status cannot be sent for any reason for longer than 30 seconds, the worker will close all connections that are normally not cleaned up (so for sessions that currently are *not* in the progress of being canceled or terminated). Status requests now time out after 10 seconds if they are not made successfully. Session status is not touched during this process, but natural cleanup or connectionless sessions will still occur. This should be okay, as the session should be added on re-connect if need be. A few other changes are made to cancelConnections: first, error logging has been moved to the function so that we can handle errors requesting cancellation directly within the function; this entails simply just locally marking these connections as closed with a warning. Second, there was a minor logging bug where closed connections were not properly being noted in the last trace log, this has been fixed. Additional changes: * Circle: Set acceptance test no_output_timeout to 30m
This commit is contained in:
1 parent
2b4155d01e
commit
9f9acce972
10 files changed
+1008
-45
No files matched your search
@@ -4,6 +4,12 @@ Canonical reference for changes, improvements, and bugfixes for Boundary.
|
||||
|
||||
## Next
|
||||
|
||||
### New and Improved
|
||||
|
||||
* The worker will now close any existing proxy connections it is handling when
|
||||
it cannot make a status request to the worker. The timeout for this behavior
|
||||
is 15 seconds.
|
||||
|
||||
### Bug Fixes
|
||||
|
||||
* scheduler: removes a Postgres check constraint, on the length of the controller name,
|
||||
|
||||
@@ -19,6 +19,7 @@ require (
|
||||
github.com/grpc-ecosystem/grpc-gateway/v2 v2.3.0
|
||||
github.com/hashicorp/boundary/api v0.0.11
|
||||
github.com/hashicorp/boundary/sdk v0.0.4
|
||||
github.com/hashicorp/dawdle v0.4.0
|
||||
github.com/hashicorp/cap v0.1.1
|
||||
github.com/hashicorp/dbassert v0.0.0-20200930125617-6218396928df
|
||||
github.com/hashicorp/errwrap v1.1.0
|
||||
|
||||
@@ -447,6 +447,8 @@ github.com/hashicorp/consul/api v1.1.0/go.mod h1:VmuI/Lkw1nC05EYQWNKwWGbkg+FbDBt
|
||||
github.com/hashicorp/consul/api v1.3.0/go.mod h1:MmDNSzIMUjNpY/mQ398R4bk2FnqQLoPndWW5VkKPlCE=
|
||||
github.com/hashicorp/consul/sdk v0.1.1/go.mod h1:VKf9jXwCTEY1QZP2MOLRhb5i/I/ssyNV1vwHyQBF0x8=
|
||||
github.com/hashicorp/consul/sdk v0.3.0/go.mod h1:VKf9jXwCTEY1QZP2MOLRhb5i/I/ssyNV1vwHyQBF0x8=
|
||||
github.com/hashicorp/dawdle v0.4.0 h1:nAE+aCoAmVulgCPPPrp6oGOykjPhxEe4VUNms/mTKeA=
|
||||
github.com/hashicorp/dawdle v0.4.0/go.mod h1:gOBY8rFXpW4Rw39up1FIacw/1pG/fO779LfchCY3KUw=
|
||||
github.com/hashicorp/dbassert v0.0.0-20200930125617-6218396928df h1:Y5++qCYLnJsz036XXKjawzg09aCvWe5Pq2n2DhbgyrA=
|
||||
github.com/hashicorp/dbassert v0.0.0-20200930125617-6218396928df/go.mod h1:+B5eZq7vXqN4Gxkgjm4ub1bwG6OTcYzFL0r+bMvWorg=
|
||||
github.com/hashicorp/errwrap v1.0.0/go.mod h1:YH+1FKiLXxHSkmPseP+kNlulaMuP3n2brvKWEqk/Jc4=
|
||||
|
||||
@@ -367,6 +367,10 @@ type TestControllerOpts struct {
|
||||
|
||||
// The logger to use, or one will be created
|
||||
Logger hclog.Logger
|
||||
|
||||
// A cluster address for overriding the advertised controller listener
|
||||
// (overrides address provided in config, if any)
|
||||
PublicClusterAddr string
|
||||
}
|
||||
|
||||
func NewTestController(t *testing.T, opts *TestControllerOpts) *TestController {
|
||||
@@ -503,6 +507,11 @@ func NewTestController(t *testing.T, opts *TestControllerOpts) *TestController {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
// Set cluster address if we supplied one (overrides one in config)
|
||||
if opts.PublicClusterAddr != "" {
|
||||
opts.Config.Controller.PublicClusterAddr = opts.PublicClusterAddr
|
||||
}
|
||||
|
||||
if opts.DatabaseUrl != "" {
|
||||
tc.b.DatabaseUrl = opts.DatabaseUrl
|
||||
if _, err := schema.MigrateStore(ctx, "postgres", tc.b.DatabaseUrl); err != nil {
|
||||
@@ -600,6 +609,7 @@ func (tc *TestController) AddClusterControllerMember(t *testing.T, opts *TestCon
|
||||
DefaultPassword: tc.b.DevPassword,
|
||||
DisableKmsKeyCreation: true,
|
||||
DisableAuthMethodCreation: true,
|
||||
PublicClusterAddr: opts.PublicClusterAddr,
|
||||
}
|
||||
if opts.Logger != nil {
|
||||
nextOpts.Logger = opts.Logger
|
||||
|
||||
@@ -153,12 +153,7 @@ func (w *Worker) handleProxy() http.HandlerFunc {
|
||||
}
|
||||
|
||||
defer func() {
|
||||
connectionId := ci.id
|
||||
if err := w.closeConnections(r.Context(), map[string]string{
|
||||
connectionId: si.id,
|
||||
}); err != nil {
|
||||
w.logger.Error("error marking connection closed", "error", err, "connection_id", connectionId)
|
||||
}
|
||||
w.closeConnections(r.Context(), map[string]string{ci.id: si.id})
|
||||
}()
|
||||
|
||||
si.Lock()
|
||||
|
||||
@@ -16,7 +16,8 @@ import (
|
||||
)
|
||||
|
||||
const (
|
||||
validateSessionTimeout = 90 * time.Second
|
||||
validateSessionTimeout = 90 * time.Second
|
||||
errMakeSessionCloseInfoNilCloseInfo = "nil closeInfo supplied to makeSessionCloseInfo, this is a bug, please report it"
|
||||
)
|
||||
|
||||
type connInfo struct {
|
||||
@@ -251,49 +252,132 @@ func (w *Worker) closeConnection(ctx context.Context, req *pbs.CloseConnectionRe
|
||||
return resp, nil
|
||||
}
|
||||
|
||||
func (w *Worker) closeConnections(ctx context.Context, closeMap map[string]string) error {
|
||||
w.logger.Trace("marking connections as closed", "session_and_connection_ids", fmt.Sprintf("%#v", closeMap))
|
||||
// closeConnections is a helper worker function that sends connection
|
||||
// close requests to the controller, and sets close times within the
|
||||
// worker. It is called during the worker status loop and on
|
||||
// connection exit on the proxy.
|
||||
//
|
||||
// closeInfo is a map of connections mapped to their individual
|
||||
// session.
|
||||
func (w *Worker) closeConnections(ctx context.Context, closeInfo map[string]string) {
|
||||
if closeInfo == nil {
|
||||
// This should not happen, but it's a no-op if it does. Just
|
||||
// return.
|
||||
return
|
||||
}
|
||||
|
||||
closeData := make([]*pbs.CloseConnectionRequestData, 0, len(closeMap))
|
||||
for connId := range closeMap {
|
||||
w.logger.Trace("marking connections as closed", "session_and_connection_ids", fmt.Sprintf("%#v", closeInfo))
|
||||
response, err := w.closeConnection(ctx, w.makeCloseConnectionRequest(closeInfo))
|
||||
if err != nil {
|
||||
w.logger.Error("error marking connections closed", "error", err)
|
||||
w.logger.Warn(
|
||||
"error contacting controller, connections will be closed only on worker",
|
||||
"session_and_connection_ids", fmt.Sprintf("%#v", closeInfo),
|
||||
)
|
||||
}
|
||||
|
||||
// Mark connections as closed
|
||||
closedIds, errs := w.setCloseTimeForResponse(w.makeSessionCloseInfo(closeInfo, response))
|
||||
if len(errs) > 0 {
|
||||
for _, err := range errs {
|
||||
w.logger.Error("error marking connection closed in state", "err", err)
|
||||
}
|
||||
}
|
||||
|
||||
w.logger.Trace("connections successfully marked closed", "connection_ids", closedIds)
|
||||
}
|
||||
|
||||
// makeCloseConnectionRequest creates a CloseConnectionRequest for
|
||||
// use with closing connections.
|
||||
//
|
||||
// closeInfo is a map, indexed by connection ID, to the individual
|
||||
// sessions IDs that those connections belong to. The values are
|
||||
// ignored; the parameter is expected as such just for convenience of
|
||||
// its caller.
|
||||
func (w *Worker) makeCloseConnectionRequest(closeInfo map[string]string) *pbs.CloseConnectionRequest {
|
||||
closeData := make([]*pbs.CloseConnectionRequestData, 0, len(closeInfo))
|
||||
for connId := range closeInfo {
|
||||
closeData = append(closeData, &pbs.CloseConnectionRequestData{
|
||||
ConnectionId: connId,
|
||||
Reason: session.UnknownReason.String(),
|
||||
})
|
||||
}
|
||||
closeInfo := &pbs.CloseConnectionRequest{
|
||||
|
||||
return &pbs.CloseConnectionRequest{
|
||||
CloseRequestData: closeData,
|
||||
}
|
||||
}
|
||||
|
||||
connStatus, err := w.closeConnection(ctx, closeInfo)
|
||||
if err != nil {
|
||||
return err
|
||||
// makeSessionCloseInfo takes the response from CloseConnections and
|
||||
// our original closeInfo map and makes a map of slices, indexed by
|
||||
// session ID, of all of the connection responses. This allows us to
|
||||
// easily lock on session once for all connections in
|
||||
// setCloseTimeForResponse.
|
||||
func (w *Worker) makeSessionCloseInfo(
|
||||
closeInfo map[string]string,
|
||||
response *pbs.CloseConnectionResponse,
|
||||
) map[string][]*pbs.CloseConnectionResponseData {
|
||||
if closeInfo == nil {
|
||||
// Should never happen, panic if it does. Results will be
|
||||
// undefined.
|
||||
panic(errMakeSessionCloseInfoNilCloseInfo)
|
||||
}
|
||||
closedIds := make([]string, 0, len(connStatus.GetCloseResponseData()))
|
||||
|
||||
// Here we build a reverse map from closeMap, that is, session ID to
|
||||
// connection IDs, for more efficient locking
|
||||
revMap := make(map[string][]*pbs.CloseConnectionResponseData)
|
||||
for _, v := range connStatus.GetCloseResponseData() {
|
||||
revMap[closeMap[v.GetConnectionId()]] = append(revMap[closeMap[v.GetConnectionId()]], v)
|
||||
result := make(map[string][]*pbs.CloseConnectionResponseData)
|
||||
for _, v := range response.GetCloseResponseData() {
|
||||
result[closeInfo[v.GetConnectionId()]] = append(result[closeInfo[v.GetConnectionId()]], v)
|
||||
}
|
||||
for k, v := range revMap {
|
||||
siRaw, ok := w.sessionInfoMap.Load(k)
|
||||
|
||||
return result
|
||||
}
|
||||
|
||||
// setCloseTimeForResponse iterates a CloseConnectionResponse and
|
||||
// sets the close time for any connection found to be closed to the
|
||||
// current time.
|
||||
//
|
||||
// sessionCloseInfo can be derived from the closeInfo supplied to
|
||||
// makeCloseConnectionRequest through reverseCloseInfo, which creates
|
||||
// a session ID to connection ID mapping.
|
||||
//
|
||||
// A non-zero error count does not necessarily mean the operation
|
||||
// failed, as some connections may have been marked as closed. The
|
||||
// actual list of connection IDs closed is returned as the first
|
||||
// return value.
|
||||
func (w *Worker) setCloseTimeForResponse(sessionCloseInfo map[string][]*pbs.CloseConnectionResponseData) ([]string, []error) {
|
||||
closedIds := make([]string, 0)
|
||||
var errors []error
|
||||
for sessionId, responses := range sessionCloseInfo {
|
||||
siRaw, ok := w.sessionInfoMap.Load(sessionId)
|
||||
if !ok {
|
||||
w.logger.Warn("could not find session ID in info map after closing connections", "session_id", k)
|
||||
errors = append(errors, fmt.Errorf("could not find session ID %q in local state after closing connections", sessionId))
|
||||
continue
|
||||
}
|
||||
|
||||
si := siRaw.(*sessionInfo)
|
||||
si.Lock()
|
||||
for _, connResult := range v {
|
||||
ci := si.connInfoMap[connResult.GetConnectionId()]
|
||||
ci.status = connResult.GetStatus()
|
||||
|
||||
for _, response := range responses {
|
||||
ci, ok := si.connInfoMap[response.GetConnectionId()]
|
||||
if !ok {
|
||||
errors = append(errors,
|
||||
fmt.Errorf(
|
||||
"could not find connection ID %q for session ID %q in local state after closing connections",
|
||||
response.GetConnectionId(),
|
||||
sessionId,
|
||||
),
|
||||
)
|
||||
continue
|
||||
}
|
||||
|
||||
ci.status = response.GetStatus()
|
||||
if ci.status == pbs.CONNECTIONSTATUS_CONNECTIONSTATUS_CLOSED {
|
||||
ci.closeTime = time.Now()
|
||||
closedIds = append(closedIds, ci.id)
|
||||
}
|
||||
}
|
||||
|
||||
si.Unlock()
|
||||
}
|
||||
w.logger.Trace("connections successfully marked closed", "connection_ids", closedIds)
|
||||
return nil
|
||||
|
||||
return closedIds, errors
|
||||
}
|
||||
@@ -0,0 +1,249 @@
|
||||
package worker
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
pbs "github.com/hashicorp/boundary/internal/gen/controller/servers/services"
|
||||
"github.com/hashicorp/boundary/internal/session"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestWorkerMakeCloseConnectionRequest(t *testing.T) {
|
||||
require := require.New(t)
|
||||
in := map[string]string{"foo": "one", "bar": "two"}
|
||||
expected := &pbs.CloseConnectionRequest{
|
||||
CloseRequestData: []*pbs.CloseConnectionRequestData{
|
||||
{ConnectionId: "foo", Reason: session.UnknownReason.String()},
|
||||
{ConnectionId: "bar", Reason: session.UnknownReason.String()},
|
||||
},
|
||||
}
|
||||
actual := new(Worker).makeCloseConnectionRequest(in)
|
||||
require.ElementsMatch(expected.GetCloseRequestData(), actual.GetCloseRequestData())
|
||||
}
|
||||
|
||||
func TestMakeSessionCloseInfo(t *testing.T) {
|
||||
require := require.New(t)
|
||||
closeInfo := map[string]string{"foo": "one", "bar": "two"}
|
||||
response := &pbs.CloseConnectionResponse{
|
||||
CloseResponseData: []*pbs.CloseConnectionResponseData{
|
||||
{ConnectionId: "foo", Status: pbs.CONNECTIONSTATUS_CONNECTIONSTATUS_CLOSED},
|
||||
{ConnectionId: "bar", Status: pbs.CONNECTIONSTATUS_CONNECTIONSTATUS_CLOSED},
|
||||
},
|
||||
}
|
||||
expected := map[string][]*pbs.CloseConnectionResponseData{
|
||||
"one": []*pbs.CloseConnectionResponseData{
|
||||
{ConnectionId: "foo", Status: pbs.CONNECTIONSTATUS_CONNECTIONSTATUS_CLOSED},
|
||||
},
|
||||
"two": []*pbs.CloseConnectionResponseData{
|
||||
{ConnectionId: "bar", Status: pbs.CONNECTIONSTATUS_CONNECTIONSTATUS_CLOSED},
|
||||
},
|
||||
}
|
||||
actual := new(Worker).makeSessionCloseInfo(closeInfo, response)
|
||||
require.Equal(expected, actual)
|
||||
}
|
||||
|
||||
func TestMakeSessionCloseInfoPanicIfCloseInfoNil(t *testing.T) {
|
||||
require := require.New(t)
|
||||
require.Panics(func() {
|
||||
new(Worker).makeSessionCloseInfo(nil, nil)
|
||||
})
|
||||
}
|
||||
|
||||
func TestMakeSessionCloseInfoEmpty(t *testing.T) {
|
||||
require := require.New(t)
|
||||
require.Equal(
|
||||
make(map[string][]*pbs.CloseConnectionResponseData),
|
||||
new(Worker).makeSessionCloseInfo(make(map[string]string), nil),
|
||||
)
|
||||
}
|
||||
|
||||
func TestWorkerSetCloseTimeForResponse(t *testing.T) {
|
||||
cases := []struct {
|
||||
name string
|
||||
sessionCloseInfo map[string][]*pbs.CloseConnectionResponseData
|
||||
sessionInfoMap func() *sync.Map
|
||||
expected []string
|
||||
expectedClosed map[string]struct{}
|
||||
expectedErr []error
|
||||
}{
|
||||
{
|
||||
name: "basic",
|
||||
sessionCloseInfo: map[string][]*pbs.CloseConnectionResponseData{
|
||||
"one": []*pbs.CloseConnectionResponseData{
|
||||
{ConnectionId: "foo", Status: pbs.CONNECTIONSTATUS_CONNECTIONSTATUS_CLOSED},
|
||||
},
|
||||
"two": []*pbs.CloseConnectionResponseData{
|
||||
{ConnectionId: "bar", Status: pbs.CONNECTIONSTATUS_CONNECTIONSTATUS_CLOSED},
|
||||
},
|
||||
},
|
||||
sessionInfoMap: func() *sync.Map {
|
||||
m := new(sync.Map)
|
||||
m.Store("one", &sessionInfo{
|
||||
id: "one",
|
||||
connInfoMap: map[string]*connInfo{
|
||||
"foo": &connInfo{id: "foo"},
|
||||
},
|
||||
})
|
||||
m.Store("two", &sessionInfo{
|
||||
id: "two",
|
||||
connInfoMap: map[string]*connInfo{
|
||||
"bar": &connInfo{id: "bar"},
|
||||
},
|
||||
})
|
||||
m.Store("three", &sessionInfo{
|
||||
id: "three",
|
||||
connInfoMap: map[string]*connInfo{
|
||||
"baz": &connInfo{id: "baz"},
|
||||
},
|
||||
})
|
||||
|
||||
return m
|
||||
},
|
||||
expected: []string{"foo", "bar"},
|
||||
expectedClosed: map[string]struct{}{
|
||||
"foo": struct{}{},
|
||||
"bar": struct{}{},
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "not closed",
|
||||
sessionCloseInfo: map[string][]*pbs.CloseConnectionResponseData{
|
||||
"one": []*pbs.CloseConnectionResponseData{
|
||||
{ConnectionId: "foo", Status: pbs.CONNECTIONSTATUS_CONNECTIONSTATUS_CLOSED},
|
||||
},
|
||||
"two": []*pbs.CloseConnectionResponseData{
|
||||
{ConnectionId: "bar", Status: pbs.CONNECTIONSTATUS_CONNECTIONSTATUS_CONNECTED},
|
||||
},
|
||||
},
|
||||
sessionInfoMap: func() *sync.Map {
|
||||
m := new(sync.Map)
|
||||
m.Store("one", &sessionInfo{
|
||||
id: "one",
|
||||
connInfoMap: map[string]*connInfo{
|
||||
"foo": &connInfo{id: "foo"},
|
||||
},
|
||||
})
|
||||
m.Store("two", &sessionInfo{
|
||||
id: "two",
|
||||
connInfoMap: map[string]*connInfo{
|
||||
"bar": &connInfo{id: "bar"},
|
||||
},
|
||||
})
|
||||
|
||||
return m
|
||||
},
|
||||
expected: []string{"foo"},
|
||||
expectedClosed: map[string]struct{}{
|
||||
"foo": struct{}{},
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "missing session",
|
||||
sessionCloseInfo: map[string][]*pbs.CloseConnectionResponseData{
|
||||
"one": []*pbs.CloseConnectionResponseData{
|
||||
{ConnectionId: "foo", Status: pbs.CONNECTIONSTATUS_CONNECTIONSTATUS_CLOSED},
|
||||
},
|
||||
"two": []*pbs.CloseConnectionResponseData{
|
||||
{ConnectionId: "bar", Status: pbs.CONNECTIONSTATUS_CONNECTIONSTATUS_CLOSED},
|
||||
},
|
||||
},
|
||||
sessionInfoMap: func() *sync.Map {
|
||||
m := new(sync.Map)
|
||||
m.Store("one", &sessionInfo{
|
||||
id: "one",
|
||||
connInfoMap: map[string]*connInfo{
|
||||
"foo": &connInfo{id: "foo"},
|
||||
},
|
||||
})
|
||||
|
||||
return m
|
||||
},
|
||||
expected: []string{"foo"},
|
||||
expectedClosed: map[string]struct{}{
|
||||
"foo": struct{}{},
|
||||
},
|
||||
expectedErr: []error{
|
||||
errors.New(`could not find session ID "two" in local state after closing connections`),
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "missing connection",
|
||||
sessionCloseInfo: map[string][]*pbs.CloseConnectionResponseData{
|
||||
"one": []*pbs.CloseConnectionResponseData{
|
||||
{ConnectionId: "foo", Status: pbs.CONNECTIONSTATUS_CONNECTIONSTATUS_CLOSED},
|
||||
},
|
||||
"two": []*pbs.CloseConnectionResponseData{
|
||||
{ConnectionId: "bar", Status: pbs.CONNECTIONSTATUS_CONNECTIONSTATUS_CLOSED},
|
||||
},
|
||||
},
|
||||
sessionInfoMap: func() *sync.Map {
|
||||
m := new(sync.Map)
|
||||
m.Store("one", &sessionInfo{
|
||||
id: "one",
|
||||
connInfoMap: map[string]*connInfo{
|
||||
"foo": &connInfo{id: "foo"},
|
||||
},
|
||||
})
|
||||
m.Store("two", &sessionInfo{id: "two"})
|
||||
|
||||
return m
|
||||
},
|
||||
expected: []string{"foo"},
|
||||
expectedClosed: map[string]struct{}{
|
||||
"foo": struct{}{},
|
||||
},
|
||||
expectedErr: []error{
|
||||
errors.New(`could not find connection ID "bar" for session ID "two" in local state after closing connections`),
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "empty",
|
||||
sessionCloseInfo: make(map[string][]*pbs.CloseConnectionResponseData),
|
||||
sessionInfoMap: func() *sync.Map {
|
||||
m := new(sync.Map)
|
||||
m.Store("one", &sessionInfo{
|
||||
id: "one",
|
||||
connInfoMap: map[string]*connInfo{
|
||||
"foo": &connInfo{id: "foo"},
|
||||
},
|
||||
})
|
||||
|
||||
return m
|
||||
},
|
||||
expected: []string{},
|
||||
expectedClosed: map[string]struct{}{},
|
||||
},
|
||||
}
|
||||
|
||||
for _, tc := range cases {
|
||||
tc := tc
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
require := require.New(t)
|
||||
w := &Worker{
|
||||
sessionInfoMap: tc.sessionInfoMap(),
|
||||
}
|
||||
actual, actualErr := w.setCloseTimeForResponse(tc.sessionCloseInfo)
|
||||
|
||||
// Assert all close times were set
|
||||
w.sessionInfoMap.Range(func(key, value interface{}) bool {
|
||||
t.Helper()
|
||||
for _, ci := range value.(*sessionInfo).connInfoMap {
|
||||
if _, ok := tc.expectedClosed[ci.id]; ok {
|
||||
require.NotEqual(time.Time{}, ci.closeTime)
|
||||
} else {
|
||||
require.Equal(time.Time{}, ci.closeTime)
|
||||
}
|
||||
}
|
||||
|
||||
return true
|
||||
})
|
||||
|
||||
// Assert return values
|
||||
require.ElementsMatch(tc.expected, actual)
|
||||
require.ElementsMatch(tc.expectedErr, actualErr)
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -3,6 +3,8 @@ package worker
|
||||
import (
|
||||
"context"
|
||||
"math/rand"
|
||||
"os"
|
||||
"strconv"
|
||||
"time"
|
||||
|
||||
pbs "github.com/hashicorp/boundary/internal/gen/controller/servers/services"
|
||||
@@ -13,9 +15,48 @@ import (
|
||||
|
||||
// In the future we could make this configurable
|
||||
const (
|
||||
statusInterval = 2 * time.Second
|
||||
statusInterval = 2 * time.Second
|
||||
statusTimeout = 10 * time.Second
|
||||
defaultStatusGracePeriod = 30 * time.Second
|
||||
statusGracePeriodEnvVar = "BOUNDARY_STATUS_GRACE_PERIOD"
|
||||
)
|
||||
|
||||
// statusGracePeriod returns the status grace period setting for this
|
||||
// worker, in seconds.
|
||||
//
|
||||
// The grace period is the length of time we allow connections to run
|
||||
// on a worker in the event of an error sending status updates. The
|
||||
// period is defined the length of time since the last successful
|
||||
// update.
|
||||
//
|
||||
// The setting is derived from one of the following:
|
||||
//
|
||||
// * BOUNDARY_STATUS_GRACE_PERIOD, if defined, can be set to an
|
||||
// integer value to define the setting.
|
||||
// * If this is missing, the default (30 seconds) is used.
|
||||
//
|
||||
func (w *Worker) statusGracePeriod() time.Duration {
|
||||
if v := os.Getenv(statusGracePeriodEnvVar); v != "" {
|
||||
n, err := strconv.Atoi(v)
|
||||
if err != nil {
|
||||
w.logger.Error("could not read setting for BOUNDARY_STATUS_GRACE_PERIOD, using default",
|
||||
"err", err,
|
||||
"value", v,
|
||||
)
|
||||
return defaultStatusGracePeriod
|
||||
}
|
||||
|
||||
if n < 1 {
|
||||
w.logger.Error("invalid setting for BOUNDARY_STATUS_GRACE_PERIOD, using default", "value", v)
|
||||
return defaultStatusGracePeriod
|
||||
}
|
||||
|
||||
return time.Second * time.Duration(n)
|
||||
}
|
||||
|
||||
return defaultStatusGracePeriod
|
||||
}
|
||||
|
||||
type LastStatusInformation struct {
|
||||
*pbs.StatusResponse
|
||||
StatusTime time.Time
|
||||
@@ -52,8 +93,8 @@ func (w *Worker) startStatusTicking(cancelCtx context.Context) {
|
||||
}()
|
||||
}
|
||||
|
||||
// LastStatusSuccess is used in tests (it's exported for tests in other
|
||||
// packages) to verify the last time we sent status
|
||||
// LastStatusSuccess reports the last time we sent a successful
|
||||
// status request.
|
||||
func (w *Worker) LastStatusSuccess() *LastStatusInformation {
|
||||
return w.lastStatusSuccess.Load().(*LastStatusInformation)
|
||||
}
|
||||
@@ -102,7 +143,9 @@ func (w *Worker) sendWorkerStatus(cancelCtx context.Context) {
|
||||
if w.updateTags.Load() {
|
||||
tags = w.tags.Load().(map[string]*servers.TagValues)
|
||||
}
|
||||
result, err := client.Status(cancelCtx, &pbs.StatusRequest{
|
||||
statusCtx, statusCancel := context.WithTimeout(cancelCtx, statusTimeout)
|
||||
defer statusCancel()
|
||||
result, err := client.Status(statusCtx, &pbs.StatusRequest{
|
||||
Jobs: activeJobs,
|
||||
Worker: &servers.Server{
|
||||
PrivateId: w.conf.RawConfig.Worker.Name,
|
||||
@@ -115,6 +158,25 @@ func (w *Worker) sendWorkerStatus(cancelCtx context.Context) {
|
||||
})
|
||||
if err != nil {
|
||||
w.logger.Error("error making status request to controller", "error", err)
|
||||
// Check for last successful status. Ignore nil last status, this probably
|
||||
// means that we've never connected to a controller, and as such probably
|
||||
// don't have any sessions to worry about anyway.
|
||||
//
|
||||
// If a length of time has passed since we've been able to communicate, we
|
||||
// want to start terminating all sessions as a "break glass" kind of
|
||||
// scenario, as there will be no way we can really tell if these
|
||||
// connections should continue to exist.
|
||||
|
||||
if isPastGrace, lastStatusTime, gracePeriod := w.isPastGrace(); isPastGrace {
|
||||
w.logger.Warn("status error grace period has expired, canceling all sessions on worker",
|
||||
"last_status_time", lastStatusTime.String(),
|
||||
"grace_period", gracePeriod,
|
||||
)
|
||||
|
||||
// Run a "cleanup" for all sessions that will not be caught by
|
||||
// our standard cleanup routine.
|
||||
w.cleanupConnections(cancelCtx, true)
|
||||
}
|
||||
} else {
|
||||
w.logger.Trace("successfully sent status to controller")
|
||||
w.updateTags.Store(false)
|
||||
@@ -154,17 +216,28 @@ func (w *Worker) sendWorkerStatus(cancelCtx context.Context) {
|
||||
}
|
||||
}
|
||||
|
||||
// Cleanup: Run through current jobs. Cancel connections for any
|
||||
// canceling session or any session that is expired. Clear out
|
||||
// sessions that are canceled or expired with all connections
|
||||
// marked as closed. Close any that aren't marked as such.
|
||||
// Standard cleanup: Run through current jobs. Cancel connections
|
||||
// for any canceling session or any session that is expired.
|
||||
w.cleanupConnections(cancelCtx, false)
|
||||
}
|
||||
|
||||
// cleanupConnections walks all sessions and shuts down connections.
|
||||
// Additionally, sessions without connections are cleaned up from the
|
||||
// local worker's state.
|
||||
//
|
||||
// Use ignoreSessionState to ignore the state checks, this closes all
|
||||
// connections, regardless of whether or not the session is still
|
||||
// active.
|
||||
func (w *Worker) cleanupConnections(cancelCtx context.Context, ignoreSessionState bool) {
|
||||
closeInfo := make(map[string]string)
|
||||
cleanSessionIds := make([]string, 0)
|
||||
w.sessionInfoMap.Range(func(key, value interface{}) bool {
|
||||
si := value.(*sessionInfo)
|
||||
si.Lock()
|
||||
defer si.Unlock()
|
||||
switch {
|
||||
case si.status == pbs.SESSIONSTATUS_SESSIONSTATUS_CANCELING,
|
||||
case ignoreSessionState,
|
||||
si.status == pbs.SESSIONSTATUS_SESSIONSTATUS_CANCELING,
|
||||
si.status == pbs.SESSIONSTATUS_SESSIONSTATUS_TERMINATED,
|
||||
time.Until(si.lookupSessionResponse.Expiration.AsTime()) < 0:
|
||||
var toClose int
|
||||
@@ -184,16 +257,17 @@ func (w *Worker) sendWorkerStatus(cancelCtx context.Context) {
|
||||
cleanSessionIds = append(cleanSessionIds, si.id)
|
||||
}
|
||||
}
|
||||
si.Unlock()
|
||||
|
||||
return true
|
||||
})
|
||||
|
||||
// Note that we won't clean these from the info map until the
|
||||
// next time we run this function
|
||||
if len(closeInfo) > 0 {
|
||||
if err := w.closeConnections(cancelCtx, closeInfo); err != nil {
|
||||
w.logger.Error("error marking connections closed", "error", err)
|
||||
}
|
||||
// Call out to a helper to send the connection close requests to the
|
||||
// controller, and set the close time. This functionality is shared with
|
||||
// post-close functionality in the proxy handler.
|
||||
w.closeConnections(cancelCtx, closeInfo)
|
||||
}
|
||||
|
||||
// Forget sessions where the session is expired/canceled and all
|
||||
@@ -202,3 +276,19 @@ func (w *Worker) sendWorkerStatus(cancelCtx context.Context) {
|
||||
w.sessionInfoMap.Delete(v)
|
||||
}
|
||||
}
|
||||
|
||||
func (w *Worker) lastSuccessfulStatusTime() time.Time {
|
||||
lastStatus := w.LastStatusSuccess()
|
||||
if lastStatus == nil {
|
||||
return w.workerStartTime
|
||||
}
|
||||
|
||||
return lastStatus.StatusTime
|
||||
}
|
||||
|
||||
func (w *Worker) isPastGrace() (bool, time.Time, time.Duration) {
|
||||
t := w.lastSuccessfulStatusTime()
|
||||
u := w.statusGracePeriod()
|
||||
v := time.Since(t)
|
||||
return v > u, t, u
|
||||
}
|
||||
@@ -28,6 +28,7 @@ type Worker struct {
|
||||
started *ua.Bool
|
||||
|
||||
controllerStatusConn *atomic.Value
|
||||
workerStartTime time.Time
|
||||
lastStatusSuccess *atomic.Value
|
||||
|
||||
controllerResolver *atomic.Value
|
||||
@@ -60,6 +61,10 @@ func New(conf *Config) (*Worker, error) {
|
||||
w.lastStatusSuccess.Store((*LastStatusInformation)(nil))
|
||||
w.controllerResolver.Store((*manual.Resolver)(nil))
|
||||
|
||||
if conf.RawConfig.Worker == nil {
|
||||
conf.RawConfig.Worker = new(config.Worker)
|
||||
}
|
||||
|
||||
w.ParseAndStoreTags(conf.RawConfig.Worker.Tags)
|
||||
|
||||
if conf.SecureRandomReader == nil {
|
||||
@@ -67,9 +72,6 @@ func New(conf *Config) (*Worker, error) {
|
||||
}
|
||||
|
||||
var err error
|
||||
if conf.RawConfig.Worker == nil {
|
||||
conf.RawConfig.Worker = new(config.Worker)
|
||||
}
|
||||
if conf.RawConfig.Worker.Name == "" {
|
||||
if conf.RawConfig.Worker.Name, err = base62.Random(10); err != nil {
|
||||
return nil, fmt.Errorf("error auto-generating worker name: %w", err)
|
||||
@@ -115,6 +117,7 @@ func (w *Worker) Start() error {
|
||||
}
|
||||
|
||||
w.startStatusTicking(w.baseContext)
|
||||
w.workerStartTime = time.Now()
|
||||
w.started.Store(true)
|
||||
|
||||
return nil
|
||||
|
||||
@@ -0,0 +1,523 @@
|
||||
package cluster
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/ed25519"
|
||||
"crypto/tls"
|
||||
"crypto/x509"
|
||||
"encoding/binary"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net"
|
||||
"net/http"
|
||||
"strconv"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/hashicorp/boundary/api/targets"
|
||||
"github.com/hashicorp/boundary/globals"
|
||||
"github.com/hashicorp/boundary/internal/cmd/config"
|
||||
targetspb "github.com/hashicorp/boundary/internal/gen/controller/api/resources/targets"
|
||||
"github.com/hashicorp/boundary/internal/proxy"
|
||||
"github.com/hashicorp/boundary/internal/servers/controller"
|
||||
"github.com/hashicorp/boundary/internal/servers/worker"
|
||||
"github.com/hashicorp/dawdle"
|
||||
"github.com/hashicorp/go-cleanhttp"
|
||||
"github.com/hashicorp/go-hclog"
|
||||
"github.com/hashicorp/vault/sdk/helper/base62"
|
||||
"github.com/mr-tron/base58"
|
||||
"github.com/stretchr/testify/require"
|
||||
"google.golang.org/protobuf/proto"
|
||||
"nhooyr.io/websocket"
|
||||
"nhooyr.io/websocket/wspb"
|
||||
)
|
||||
|
||||
const testSendRecvSendMax = 60
|
||||
|
||||
func TestWorkerSessionCleanup(t *testing.T) {
|
||||
require := require.New(t)
|
||||
logger := hclog.New(&hclog.LoggerOptions{
|
||||
Level: hclog.Trace,
|
||||
})
|
||||
|
||||
conf, err := config.DevController()
|
||||
require.NoError(err)
|
||||
|
||||
pl, err := net.Listen("tcp", "localhost:0")
|
||||
require.NoError(err)
|
||||
c1 := controller.NewTestController(t, &controller.TestControllerOpts{
|
||||
Config: conf,
|
||||
InitialResourcesSuffix: "1234567890",
|
||||
Logger: logger.Named("c1"),
|
||||
PublicClusterAddr: pl.Addr().String(),
|
||||
})
|
||||
defer c1.Shutdown()
|
||||
|
||||
expectWorkers(t, c1)
|
||||
|
||||
// Wire up the testing proxies
|
||||
require.Len(c1.ClusterAddrs(), 1)
|
||||
proxy, err := dawdle.NewProxy("tcp", "", c1.ClusterAddrs()[0],
|
||||
dawdle.WithListener(pl),
|
||||
dawdle.WithRbufSize(512),
|
||||
dawdle.WithWbufSize(512),
|
||||
)
|
||||
require.NoError(err)
|
||||
defer proxy.Close()
|
||||
require.NotEmpty(t, proxy.ListenerAddr())
|
||||
|
||||
w1 := worker.NewTestWorker(t, &worker.TestWorkerOpts{
|
||||
WorkerAuthKms: c1.Config().WorkerAuthKms,
|
||||
InitialControllers: []string{proxy.ListenerAddr()},
|
||||
Logger: logger.Named("w1"),
|
||||
})
|
||||
defer w1.Shutdown()
|
||||
|
||||
time.Sleep(10 * time.Second)
|
||||
expectWorkers(t, c1, w1)
|
||||
|
||||
// Use an independent context for test things that take a context so
|
||||
// that we aren't tied to any timeouts in the controller, etc. This
|
||||
// can interfere with some of the test operations.
|
||||
ctx := context.Background()
|
||||
|
||||
// Connect target
|
||||
client := c1.Client()
|
||||
client.SetToken(c1.Token().Token)
|
||||
tcl := targets.NewClient(client)
|
||||
tgt, err := tcl.Read(ctx, "ttcp_1234567890")
|
||||
require.NoError(err)
|
||||
require.NotNil(tgt)
|
||||
|
||||
// Create test server, update default port on target
|
||||
ts := newTestTcpServer(t, logger)
|
||||
require.NotNil(t, ts)
|
||||
defer ts.Close()
|
||||
tgt, err = tcl.Update(ctx, tgt.Item.Id, tgt.Item.Version, targets.WithTcpTargetDefaultPort(ts.Port()))
|
||||
require.NoError(err)
|
||||
require.NotNil(tgt)
|
||||
|
||||
// Authorize and connect
|
||||
sess := newTestSession(ctx, t, tcl, "ttcp_1234567890")
|
||||
sConn := sess.Connect(ctx, t, logger)
|
||||
|
||||
// Run initial send/receive test, make sure things are working
|
||||
sConn.TestSendRecvAll(t)
|
||||
|
||||
// Kill the link
|
||||
proxy.Pause()
|
||||
|
||||
// Run again, ensure connection is dead
|
||||
sConn.TestSendRecvFail(t)
|
||||
|
||||
// Assert we have no connections left (should be default behavior)
|
||||
sess.TestNoConnectionsLeft(t)
|
||||
|
||||
// Resume the connection, and reconnect.
|
||||
proxy.Resume()
|
||||
time.Sleep(time.Second * 10) // Sleep to wait for worker to report back as healthy
|
||||
sess = newTestSession(ctx, t, tcl, "ttcp_1234567890") // re-assign, other connection will close in t.Cleanup()
|
||||
sConn = sess.Connect(ctx, t, logger)
|
||||
sConn.TestSendRecvAll(t)
|
||||
}
|
||||
|
||||
func TestWorkerSessionCleanupMultiController(t *testing.T) {
|
||||
require := require.New(t)
|
||||
logger := hclog.New(&hclog.LoggerOptions{
|
||||
Level: hclog.Trace,
|
||||
})
|
||||
|
||||
// ******************
|
||||
// ** Controller 1 **
|
||||
// ******************
|
||||
conf1, err := config.DevController()
|
||||
require.NoError(err)
|
||||
|
||||
pl1, err := net.Listen("tcp", "localhost:0")
|
||||
require.NoError(err)
|
||||
c1 := controller.NewTestController(t, &controller.TestControllerOpts{
|
||||
Config: conf1,
|
||||
InitialResourcesSuffix: "1234567890",
|
||||
Logger: logger.Named("c1"),
|
||||
PublicClusterAddr: pl1.Addr().String(),
|
||||
})
|
||||
defer c1.Shutdown()
|
||||
|
||||
// ******************
|
||||
// ** Controller 2 **
|
||||
// ******************
|
||||
pl2, err := net.Listen("tcp", "localhost:0")
|
||||
require.NoError(err)
|
||||
c2 := c1.AddClusterControllerMember(t, &controller.TestControllerOpts{
|
||||
Logger: c1.Config().Logger.ResetNamed("c2"),
|
||||
PublicClusterAddr: pl2.Addr().String(),
|
||||
})
|
||||
defer c2.Shutdown()
|
||||
expectWorkers(t, c1)
|
||||
expectWorkers(t, c2)
|
||||
|
||||
// *************
|
||||
// ** Proxy 1 **
|
||||
// *************
|
||||
require.Len(c1.ClusterAddrs(), 1)
|
||||
p1, err := dawdle.NewProxy("tcp", "", c1.ClusterAddrs()[0],
|
||||
dawdle.WithListener(pl1),
|
||||
dawdle.WithRbufSize(512),
|
||||
dawdle.WithWbufSize(512),
|
||||
)
|
||||
require.NoError(err)
|
||||
defer p1.Close()
|
||||
require.NotEmpty(t, p1.ListenerAddr())
|
||||
|
||||
// *************
|
||||
// ** Proxy 2 **
|
||||
// *************
|
||||
require.Len(c2.ClusterAddrs(), 1)
|
||||
p2, err := dawdle.NewProxy("tcp", "", c2.ClusterAddrs()[0],
|
||||
dawdle.WithListener(pl2),
|
||||
dawdle.WithRbufSize(512),
|
||||
dawdle.WithWbufSize(512),
|
||||
)
|
||||
require.NoError(err)
|
||||
defer p2.Close()
|
||||
require.NotEmpty(t, p2.ListenerAddr())
|
||||
|
||||
// ************
|
||||
// ** Worker **
|
||||
// ************
|
||||
w1 := worker.NewTestWorker(t, &worker.TestWorkerOpts{
|
||||
WorkerAuthKms: c1.Config().WorkerAuthKms,
|
||||
InitialControllers: []string{p1.ListenerAddr(), p2.ListenerAddr()},
|
||||
Logger: logger.Named("w1"),
|
||||
})
|
||||
defer w1.Shutdown()
|
||||
|
||||
time.Sleep(10 * time.Second)
|
||||
expectWorkers(t, c1, w1)
|
||||
expectWorkers(t, c2, w1)
|
||||
|
||||
// Use an independent context for test things that take a context so
|
||||
// that we aren't tied to any timeouts in the controller, etc. This
|
||||
// can interfere with some of the test operations.
|
||||
ctx := context.Background()
|
||||
|
||||
// Connect target
|
||||
client := c1.Client()
|
||||
client.SetToken(c1.Token().Token)
|
||||
tcl := targets.NewClient(client)
|
||||
tgt, err := tcl.Read(ctx, "ttcp_1234567890")
|
||||
require.NoError(err)
|
||||
require.NotNil(tgt)
|
||||
|
||||
// Create test server, update default port on target
|
||||
ts := newTestTcpServer(t, logger)
|
||||
require.NotNil(ts)
|
||||
defer ts.Close()
|
||||
tgt, err = tcl.Update(ctx, tgt.Item.Id, tgt.Item.Version, targets.WithTcpTargetDefaultPort(ts.Port()))
|
||||
require.NoError(err)
|
||||
require.NotNil(tgt)
|
||||
|
||||
// Authorize and connect
|
||||
sess := newTestSession(ctx, t, tcl, "ttcp_1234567890")
|
||||
sConn := sess.Connect(ctx, t, logger)
|
||||
|
||||
// Run initial send/receive test, make sure things are working
|
||||
sConn.TestSendRecvAll(t)
|
||||
|
||||
// Kill connection to first controller, and run test again, should
|
||||
// pass, deferring to other controller.
|
||||
p1.Pause()
|
||||
sConn.TestSendRecvAll(t)
|
||||
|
||||
// Resume first controller, pause second. This one should work too.
|
||||
p1.Resume()
|
||||
p2.Pause()
|
||||
sConn.TestSendRecvAll(t)
|
||||
|
||||
// Kill the first controller connection again. This one should fail
|
||||
// due to lack of any connection.
|
||||
p1.Pause()
|
||||
sConn.TestSendRecvFail(t)
|
||||
|
||||
// Finally resume both, try again. Should behave as per normal.
|
||||
p1.Resume()
|
||||
p2.Resume()
|
||||
time.Sleep(time.Second * 10) // Sleep to wait for worker to report back as healthy
|
||||
sess = newTestSession(ctx, t, tcl, "ttcp_1234567890") // re-assign, other connection will close in t.Cleanup()
|
||||
sConn = sess.Connect(ctx, t, logger)
|
||||
sConn.TestSendRecvAll(t)
|
||||
}
|
||||
|
||||
// testSession represents an authorized session.
|
||||
type testSession struct {
|
||||
workerAddr string
|
||||
transport *http.Transport
|
||||
tofuToken string
|
||||
connectionsLeft int32
|
||||
}
|
||||
|
||||
// newTestSession authorizes a session and creates all of the data
|
||||
// necessary to initialize.
|
||||
func newTestSession(
|
||||
ctx context.Context,
|
||||
t *testing.T,
|
||||
tcl *targets.Client,
|
||||
targetId string,
|
||||
) *testSession {
|
||||
t.Helper()
|
||||
require := require.New(t)
|
||||
sar, err := tcl.AuthorizeSession(ctx, "ttcp_1234567890")
|
||||
require.NoError(err)
|
||||
require.NotNil(sar)
|
||||
|
||||
s := new(testSession)
|
||||
authzString := sar.GetItem().(*targets.SessionAuthorization).AuthorizationToken
|
||||
marshaled, err := base58.FastBase58Decoding(authzString)
|
||||
require.NoError(err)
|
||||
require.NotZero(marshaled)
|
||||
|
||||
sessionAuthzData := new(targetspb.SessionAuthorizationData)
|
||||
err = proto.Unmarshal(marshaled, sessionAuthzData)
|
||||
require.NoError(err)
|
||||
require.NotZero(sessionAuthzData.GetWorkerInfo())
|
||||
|
||||
s.workerAddr = sessionAuthzData.GetWorkerInfo()[0].GetAddress()
|
||||
|
||||
parsedCert, err := x509.ParseCertificate(sessionAuthzData.Certificate)
|
||||
require.NoError(err)
|
||||
require.Len(parsedCert.DNSNames, 1)
|
||||
|
||||
certPool := x509.NewCertPool()
|
||||
certPool.AddCert(parsedCert)
|
||||
tlsConf := &tls.Config{
|
||||
Certificates: []tls.Certificate{
|
||||
{
|
||||
Certificate: [][]byte{sessionAuthzData.Certificate},
|
||||
PrivateKey: ed25519.PrivateKey(sessionAuthzData.PrivateKey),
|
||||
Leaf: parsedCert,
|
||||
},
|
||||
},
|
||||
RootCAs: certPool,
|
||||
ServerName: parsedCert.DNSNames[0],
|
||||
MinVersion: tls.VersionTLS13,
|
||||
}
|
||||
|
||||
s.transport = cleanhttp.DefaultTransport()
|
||||
s.transport.DisableKeepAlives = false
|
||||
s.transport.TLSClientConfig = tlsConf
|
||||
s.transport.IdleConnTimeout = 0
|
||||
|
||||
return s
|
||||
}
|
||||
|
||||
// connect returns a connected websocket for the stored session,
|
||||
// connecting to the stored workerAddr with the configured transport.
|
||||
//
|
||||
// The returned (wrapped) net.Conn should be ready for communication.
|
||||
func (s *testSession) connect(ctx context.Context, t *testing.T) net.Conn {
|
||||
t.Helper()
|
||||
require := require.New(t)
|
||||
conn, resp, err := websocket.Dial(
|
||||
ctx,
|
||||
fmt.Sprintf("wss://%s/v1/proxy", s.workerAddr),
|
||||
&websocket.DialOptions{
|
||||
HTTPClient: &http.Client{
|
||||
Transport: s.transport,
|
||||
},
|
||||
Subprotocols: []string{globals.TcpProxyV1},
|
||||
},
|
||||
)
|
||||
require.NoError(err)
|
||||
require.NotNil(conn)
|
||||
require.NotNil(resp)
|
||||
require.Equal(resp.Header.Get("Sec-WebSocket-Protocol"), globals.TcpProxyV1)
|
||||
|
||||
// Send the handshake.
|
||||
if s.tofuToken == "" {
|
||||
s.tofuToken, err = base62.Random(20)
|
||||
}
|
||||
|
||||
require.NoError(err)
|
||||
handshake := proxy.ClientHandshake{TofuToken: s.tofuToken}
|
||||
err = wspb.Write(ctx, conn, &handshake)
|
||||
require.NoError(err)
|
||||
|
||||
// Receive/check the handshake
|
||||
var handshakeResult proxy.HandshakeResult
|
||||
err = wspb.Read(ctx, conn, &handshakeResult)
|
||||
require.NoError(err)
|
||||
|
||||
// This is just a cursory check to make sure that the handshake is
|
||||
// populated. We could check connections remaining too, but that
|
||||
// could legitimately be a trivial (zero) value.
|
||||
require.NotNil(handshakeResult.GetExpiration())
|
||||
s.connectionsLeft = handshakeResult.GetConnectionsLeft()
|
||||
|
||||
return websocket.NetConn(ctx, conn, websocket.MessageBinary)
|
||||
}
|
||||
|
||||
// TestNoConnectionsLeft asserts that there are no connections left.
|
||||
func (s *testSession) TestNoConnectionsLeft(t *testing.T) {
|
||||
t.Helper()
|
||||
require.Zero(t, s.connectionsLeft)
|
||||
}
|
||||
|
||||
// testSessionConnection abstracts a connected session.
|
||||
type testSessionConnection struct {
|
||||
conn net.Conn
|
||||
logger hclog.Logger
|
||||
}
|
||||
|
||||
// Connect returns a testSessionConnection for a testSession. Check
|
||||
// the unexported connect method for the lower-level details.
|
||||
func (s *testSession) Connect(
|
||||
ctx context.Context,
|
||||
t *testing.T, // Just to add cleanup
|
||||
logger hclog.Logger,
|
||||
) *testSessionConnection {
|
||||
t.Helper()
|
||||
require := require.New(t)
|
||||
conn := s.connect(ctx, t)
|
||||
require.NotNil(conn)
|
||||
t.Cleanup(func() {
|
||||
conn.Close()
|
||||
})
|
||||
require.NotNil(conn)
|
||||
|
||||
return &testSessionConnection{
|
||||
conn: conn,
|
||||
logger: logger,
|
||||
}
|
||||
}
|
||||
|
||||
// testSendRecv runs a basic send/receive test over the returned
|
||||
// connection, and returns whether or not all "pings" made it
|
||||
// through.
|
||||
//
|
||||
// The test is a simple sequence number, ticking up every second to
|
||||
// max. The passed in conn is expected to copy whatever it is
|
||||
// received.
|
||||
func (c *testSessionConnection) testSendRecv(t *testing.T) bool {
|
||||
t.Helper()
|
||||
require := require.New(t)
|
||||
for i := uint32(0); i < testSendRecvSendMax; i++ {
|
||||
// Shuttle over the sequence number as base64.
|
||||
err := binary.Write(c.conn, binary.LittleEndian, i)
|
||||
if err != nil {
|
||||
c.logger.Debug("received error during write", "err", err)
|
||||
if errors.Is(err, net.ErrClosed) ||
|
||||
errors.Is(err, io.EOF) ||
|
||||
errors.Is(err, websocket.CloseError{Code: websocket.StatusPolicyViolation, Reason: "timed out"}) {
|
||||
return false
|
||||
}
|
||||
|
||||
require.FailNow(err.Error())
|
||||
}
|
||||
|
||||
// Read it back
|
||||
var j uint32
|
||||
err = binary.Read(c.conn, binary.LittleEndian, &j)
|
||||
if err != nil {
|
||||
c.logger.Debug("received error during read", "err", err)
|
||||
if errors.Is(err, net.ErrClosed) ||
|
||||
errors.Is(err, io.EOF) ||
|
||||
errors.Is(err, websocket.CloseError{Code: websocket.StatusPolicyViolation, Reason: "timed out"}) {
|
||||
return false
|
||||
}
|
||||
|
||||
require.FailNow(err.Error())
|
||||
}
|
||||
|
||||
require.Equal(j, i)
|
||||
|
||||
// Sleep 1s
|
||||
time.Sleep(time.Second)
|
||||
}
|
||||
|
||||
return true
|
||||
}
|
||||
|
||||
// TestSendRecvAll asserts that we were able to send/recv all pings
|
||||
// over the test connection.
|
||||
func (c *testSessionConnection) TestSendRecvAll(t *testing.T) {
|
||||
t.Helper()
|
||||
require.True(t, c.testSendRecv(t))
|
||||
}
|
||||
|
||||
// TestSendRecvFail asserts that we were able to send/recv all pings
|
||||
// over the test connection.
|
||||
func (c *testSessionConnection) TestSendRecvFail(t *testing.T) {
|
||||
t.Helper()
|
||||
require.False(t, c.testSendRecv(t))
|
||||
}
|
||||
|
||||
type testTcpServer struct {
|
||||
logger hclog.Logger
|
||||
ln net.Listener
|
||||
conns map[string]net.Conn
|
||||
}
|
||||
|
||||
func (ts *testTcpServer) Port() uint32 {
|
||||
_, portS, err := net.SplitHostPort(ts.ln.Addr().String())
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
|
||||
if portS == "" {
|
||||
panic("empty port in what should be TCP listener")
|
||||
}
|
||||
|
||||
port, err := strconv.Atoi(portS)
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
|
||||
if port < 1 {
|
||||
panic("zero or negative port in what should be a TCP listener")
|
||||
}
|
||||
|
||||
return uint32(port)
|
||||
}
|
||||
|
||||
func (ts *testTcpServer) Close() {
|
||||
ts.ln.Close()
|
||||
for _, conn := range ts.conns {
|
||||
conn.Close()
|
||||
}
|
||||
}
|
||||
|
||||
func (ts *testTcpServer) run() {
|
||||
for {
|
||||
conn, err := ts.ln.Accept()
|
||||
if err != nil {
|
||||
if !errors.Is(err, net.ErrClosed) {
|
||||
ts.logger.Error("Accept() error in testTcpServer", "err", err)
|
||||
}
|
||||
|
||||
return
|
||||
}
|
||||
|
||||
ts.conns[conn.RemoteAddr().String()] = conn
|
||||
|
||||
go func(c net.Conn) {
|
||||
io.Copy(c, c)
|
||||
c.Close()
|
||||
}(conn)
|
||||
}
|
||||
}
|
||||
|
||||
func newTestTcpServer(t *testing.T, logger hclog.Logger) *testTcpServer {
|
||||
t.Helper()
|
||||
require := require.New(t)
|
||||
ts := &testTcpServer{
|
||||
logger: logger,
|
||||
conns: make(map[string]net.Conn),
|
||||
}
|
||||
var err error
|
||||
ts.ln, err = net.Listen("tcp", ":0")
|
||||
require.NoError(err)
|
||||
|
||||
go ts.run()
|
||||
return ts
|
||||
}
|
||||
Reference in new issue
Block a user