diff --git a/internal/cmd/commands/dev/dev.go b/internal/cmd/commands/dev/dev.go index a4b663f9a..c373849bf 100644 --- a/internal/cmd/commands/dev/dev.go +++ b/internal/cmd/commands/dev/dev.go @@ -5,8 +5,11 @@ import ( "errors" "fmt" "net" + "os" + "os/signal" "runtime" "strings" + "syscall" "github.com/hashicorp/boundary/internal/auth/oidc" "github.com/hashicorp/boundary/internal/auth/password" @@ -626,7 +629,23 @@ func (c *Command) Run(args []string) int { for !shutdownTriggered { select { case <-c.ShutdownCh: - c.UI.Output("==> Boundary dev environment shutdown triggered") + c.UI.Output("==> Boundary dev environment shutdown triggered, interrupt again to force") + + // Add a force-shutdown goroutine to consume another interrupt + abortForceShutdownCh := make(chan struct{}) + defer close(abortForceShutdownCh) + go func() { + shutdownCh := make(chan os.Signal, 4) + signal.Notify(shutdownCh, os.Interrupt, syscall.SIGTERM) + select { + case <-shutdownCh: + c.UI.Error("Second interrupt received, forcing shutdown") + os.Exit(base.CommandUserError) + + case <-abortForceShutdownCh: + // No-op, we just use this to shut down the goroutine + } + }() if !c.flagControllerOnly { if err := c.worker.Shutdown(false); err != nil { diff --git a/internal/cmd/commands/server/listener_reload_test.go b/internal/cmd/commands/server/listener_reload_test.go index 4b7af2658..2bb6ab03b 100644 --- a/internal/cmd/commands/server/listener_reload_test.go +++ b/internal/cmd/commands/server/listener_reload_test.go @@ -92,7 +92,12 @@ func TestServer_ReloadListener(t *testing.T) { controllerKey, workerAuthKey, recoveryKey := config.DevKeyGeneration() - cmd := testServerCommand(t, controllerKey) + cmd := testServerCommand(t, testServerCommandOpts{ + CreateDevDatabase: true, + ControllerKey: controllerKey, + UseDevAuthMethod: true, + UseDevTarget: true, + }) defer func() { if cmd.DevDatabaseCleanupFunc != nil { diff --git a/internal/cmd/commands/server/server.go b/internal/cmd/commands/server/server.go index 7c9a8b3c1..b23205e52 100644 --- a/internal/cmd/commands/server/server.go +++ b/internal/cmd/commands/server/server.go @@ -5,8 +5,11 @@ import ( stderrors "errors" "fmt" "net" + "os" + "os/signal" "runtime" "strings" + "syscall" "time" "github.com/hashicorp/boundary/globals" @@ -459,9 +462,6 @@ func (c *Command) Run(args []string) int { if c.Config.Worker != nil { if err := c.StartWorker(); err != nil { c.UI.Error(err.Error()) - if err := c.controller.Shutdown(false); err != nil { - c.UI.Error(fmt.Errorf("Error with controller shutdown: %w", err).Error()) - } return base.CommandCliError } } @@ -592,14 +592,31 @@ func (c *Command) WaitForInterrupt() int { for !shutdownTriggered { select { case <-c.ShutdownCh: - c.UI.Output("==> Boundary server shutdown triggered") + c.UI.Output("==> Boundary server shutdown triggered, interrupt again to force") + // Add a force-shutdown goroutine to consume another interrupt + abortForceShutdownCh := make(chan struct{}) + defer close(abortForceShutdownCh) + go func() { + shutdownCh := make(chan os.Signal, 4) + signal.Notify(shutdownCh, os.Interrupt, syscall.SIGTERM) + select { + case <-shutdownCh: + c.UI.Error("Second interrupt received, forcing shutdown") + os.Exit(base.CommandUserError) + case <-abortForceShutdownCh: + // No-op, we just use this to shut down the goroutine + } + }() + + // Do worker shutdown if c.Config.Worker != nil { if err := c.worker.Shutdown(false); err != nil { c.UI.Error(fmt.Errorf("Error shutting down worker: %w", err).Error()) } } + // Do controller shutdown if c.Config.Controller != nil { if err := c.controller.Shutdown(c.Config.Worker != nil); err != nil { c.UI.Error(fmt.Errorf("Error shutting down controller: %w", err).Error()) diff --git a/internal/cmd/commands/server/server_test.go b/internal/cmd/commands/server/server_test.go index 5a7db61c6..0b4f81ce7 100644 --- a/internal/cmd/commands/server/server_test.go +++ b/internal/cmd/commands/server/server_test.go @@ -6,10 +6,16 @@ import ( "github.com/hashicorp/boundary/internal/cmd/base" "github.com/hashicorp/boundary/internal/cmd/config" + "github.com/hashicorp/boundary/internal/servers/controller" "github.com/mitchellh/cli" "github.com/stretchr/testify/require" ) +// We try to pull these from TestController, but targets et al are +// computed off of a suffix instead of having constants. Just define +// for now here, this can be tweaked later if need be. +const defaultTestTargetId = "ttcp_1234567890" + const rootKmsConfig = ` kms "aead" { purpose = "root" @@ -18,7 +24,21 @@ kms "aead" { key_id = "global_root" }` -func testServerCommand(t *testing.T, controllerKey string) *Command { +type testServerCommandOpts struct { + // Whether or not to create the dev database + CreateDevDatabase bool + + // The controller key used in dev database creation + ControllerKey string + + // Use the well-known dev mode auth method id (1234567890) + UseDevAuthMethod bool + + // Use the well-known dev mode target method id (1234567890) + UseDevTarget bool +} + +func testServerCommand(t *testing.T, opts testServerCommandOpts) *Command { require := require.New(t) t.Helper() cmd := &Command{ @@ -31,17 +51,29 @@ func testServerCommand(t *testing.T, controllerKey string) *Command { require.NoError(cmd.SetupLogging("trace", "", "", "")) require.NoError(cmd.SetupEventing(cmd.Logger, cmd.StderrLock, "test-server-command")) - kmsHcl := fmt.Sprintf(rootKmsConfig, controllerKey) - parsedKmsConfig, err := config.Parse(kmsHcl) - require.NoError(err) - require.NoError(cmd.SetupKMSes(cmd.UI, parsedKmsConfig)) - - err = cmd.CreateDevDatabase(cmd.Context, base.WithContainerImage("postgres"), base.WithSkipOidcAuthMethodCreation()) - if err != nil { - if cmd.DevDatabaseCleanupFunc != nil { - require.NoError(cmd.DevDatabaseCleanupFunc()) - } + if opts.CreateDevDatabase { + kmsHcl := fmt.Sprintf(rootKmsConfig, opts.ControllerKey) + parsedKmsConfig, err := config.Parse(kmsHcl) require.NoError(err) + require.NoError(cmd.SetupKMSes(cmd.UI, parsedKmsConfig)) + + if opts.UseDevAuthMethod { + cmd.Server.DevPasswordAuthMethodId = controller.DefaultTestPasswordAuthMethodId + cmd.Server.DevLoginName = controller.DefaultTestLoginName + cmd.Server.DevPassword = controller.DefaultTestPassword + } + + if opts.UseDevTarget { + cmd.Server.DevTargetId = defaultTestTargetId + } + + err = cmd.CreateDevDatabase(cmd.Context, base.WithContainerImage("postgres"), base.WithSkipOidcAuthMethodCreation()) + if err != nil { + if cmd.DevDatabaseCleanupFunc != nil { + require.NoError(cmd.DevDatabaseCleanupFunc()) + } + require.NoError(err) + } } return cmd diff --git a/internal/cmd/commands/server/worker_shutdown_reload_test.go b/internal/cmd/commands/server/worker_shutdown_reload_test.go new file mode 100644 index 000000000..ee6a8af18 --- /dev/null +++ b/internal/cmd/commands/server/worker_shutdown_reload_test.go @@ -0,0 +1,216 @@ +// +build !hsm + +package server + +import ( + "context" + "encoding/json" + "fmt" + "testing" + "time" + + "github.com/hashicorp/boundary/api" + "github.com/hashicorp/boundary/api/authmethods" + "github.com/hashicorp/boundary/api/authtokens" + "github.com/hashicorp/boundary/api/targets" + "github.com/hashicorp/boundary/internal/cmd/config" + "github.com/hashicorp/boundary/internal/session" + "github.com/hashicorp/boundary/internal/tests/helper" + "github.com/hashicorp/go-hclog" + "github.com/mitchellh/cli" + "github.com/stretchr/testify/require" + "go.uber.org/atomic" +) + +const shutdownReloadApiAddr = "http://127.0.0.1:9203" + +const shutdownReloadControllerConfig = ` +disable_mlock = true + +controller { + name = "dev-controller" + description = "A default controller created in dev mode" + database { + url = "%s" + } +} + +kms "aead" { + purpose = "root" + aead_type = "aes-gcm" + key = "%s" + key_id = "global_root" +} + +kms "aead" { + purpose = "worker-auth" + aead_type = "aes-gcm" + key = "%s" + key_id = "global_worker-auth" +} + +listener "tcp" { + purpose = "api" + address = "127.0.0.1:9203" + tls_disable = true + cors_enabled = true + cors_allowed_origins = ["*"] +} + +listener "tcp" { + purpose = "cluster" + address = "127.0.0.1:9204" +} +` + +const shutdownReloadWorkerConfig = ` +disable_mlock = true + +worker { + name = "dev-worker" + description = "A default worker created in dev mode" + controllers = ["127.0.0.1:9204"] +} + +kms "aead" { + purpose = "worker-auth" + aead_type = "aes-gcm" + key = "%s" + key_id = "global_worker-auth" +} + +listener "tcp" { + purpose = "proxy" + address = "127.0.0.1:9205" +} +` + +func TestServer_ShutdownWorker(t *testing.T) { + logger := hclog.New(&hclog.LoggerOptions{ + Name: t.Name(), + Level: hclog.Trace, + }) + require := require.New(t) + controllerKey, workerAuthKey, _ := config.DevKeyGeneration() + + // Start the controller + controllerCmd := testServerCommand(t, testServerCommandOpts{ + CreateDevDatabase: true, + ControllerKey: controllerKey, + UseDevAuthMethod: true, + UseDevTarget: true, + }) + controllerDoneCh := make(chan struct{}) + defer func() { + if controllerCmd.DevDatabaseCleanupFunc != nil { + require.NoError(controllerCmd.DevDatabaseCleanupFunc()) + } + }() + controllerCmd.presetConfig = atomic.NewString(fmt.Sprintf(shutdownReloadControllerConfig, controllerCmd.DatabaseUrl, controllerKey, workerAuthKey)) + + go func() { + defer close(controllerDoneCh) + if code := controllerCmd.Run(nil); code != 0 { + output := controllerCmd.UI.(*cli.MockUi).ErrorWriter.String() + controllerCmd.UI.(*cli.MockUi).OutputWriter.String() + require.FailNow(output, "command exited with non-zero error code") + } + }() + + waitCh(t, controllerCmd.startedCh) + + // Start the worker + workerCmd := testServerCommand(t, testServerCommandOpts{}) + workerDoneCh := make(chan struct{}) + workerCmd.presetConfig = atomic.NewString(fmt.Sprintf(shutdownReloadWorkerConfig, workerAuthKey)) + + go func() { + defer close(workerDoneCh) + if code := workerCmd.Run(nil); code != 0 { + output := workerCmd.UI.(*cli.MockUi).ErrorWriter.String() + workerCmd.UI.(*cli.MockUi).OutputWriter.String() + require.FailNow(output, "command exited with non-zero error code") + } + }() + + waitCh(t, workerCmd.startedCh) + + // Set up the target + ctx := context.Background() + client := buildClient(t, shutdownReloadApiAddr) + setAuthToken(ctx, t, client) + + 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 := helper.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 := helper.NewTestSession(ctx, t, logger, tcl, "ttcp_1234567890") + sConn := sess.Connect(ctx, t) + + // Run initial send/receive test, make sure things are working + logger.Debug("running initial send/recv test") + sConn.TestSendRecvAll(t) + + // Now, shut the worker down. + close(workerCmd.ShutdownCh) + waitCh(t, workerDoneCh) + + // Connection should fail, and the session should be closed on the DB. + sConn.TestSendRecvFail(t) + sess.ExpectConnectionStateOnController(ctx, t, controllerCmd.controller.SessionRepoFn, session.StatusClosed) + + // We're done! Shutdown the controller, and that's it. + close(controllerCmd.ShutdownCh) + waitCh(t, controllerDoneCh) +} + +// largely copied from controller/testing.go +func buildClient(t *testing.T, addr string) *api.Client { + require := require.New(t) + client, err := api.NewClient(nil) + require.NoError(err) + err = client.SetAddr(addr) + require.NoError(err) + // Because this is using the real lib it can pick up from stored locations + // like the system keychain. Explicitly clear the token to ensure we + // understand the client state at the start of each test. + client.SetToken("") + + return client +} + +func setAuthToken(ctx context.Context, t *testing.T, client *api.Client) { + require := require.New(t) + result, err := authmethods.NewClient(client).Authenticate( + ctx, + "ampw_1234567890", + "login", + map[string]interface{}{ + "login_name": "admin", + "password": "passpass", + }, + ) + require.NoError(err) + token := new(authtokens.AuthToken) + err = json.Unmarshal(result.GetRawAttributes(), token) + require.NoError(err) + + client.SetToken(token.Token) +} + +func waitCh(t *testing.T, c chan struct{}) { + require := require.New(t) + select { + case <-c: + case <-time.After(15 * time.Second): + require.FailNow("timeout") + } +} diff --git a/internal/cmd/commands/server/worker_tags_reload_test.go b/internal/cmd/commands/server/worker_tags_reload_test.go index da20bfc0f..8b53b2fc7 100644 --- a/internal/cmd/commands/server/worker_tags_reload_test.go +++ b/internal/cmd/commands/server/worker_tags_reload_test.go @@ -106,7 +106,12 @@ func TestServer_ReloadWorkerTags(t *testing.T) { controllerKey, workerAuthKey, recoveryKey := config.DevKeyGeneration() - cmd := testServerCommand(t, controllerKey) + cmd := testServerCommand(t, testServerCommandOpts{ + CreateDevDatabase: true, + ControllerKey: controllerKey, + UseDevAuthMethod: true, + UseDevTarget: true, + }) defer func() { if cmd.DevDatabaseCleanupFunc != nil { require.NoError(cmd.DevDatabaseCleanupFunc()) diff --git a/internal/servers/worker/worker.go b/internal/servers/worker/worker.go index 57e23ae5d..c354413d9 100644 --- a/internal/servers/worker/worker.go +++ b/internal/servers/worker/worker.go @@ -139,20 +139,58 @@ func (w *Worker) Shutdown(skipListeners bool) error { event.WriteSysEvent(context.TODO(), op, "already shut down, skipping") return nil } + + // Stop listeners first to prevent new connections to the + // controller. + w.logger.Debug("beginning shutdown") defer w.started.Store(false) w.Resolver().UpdateState(resolver.State{Addresses: []resolver.Address{}}) w.baseCancel() if !skipListeners { + w.logger.Debug("stopping listeners") if err := w.stopListeners(); err != nil { return fmt.Errorf("error stopping worker listeners: %w", err) } } + + // Shut down all connections. + w.logger.Debug("shutting down all connections due to worker shutdown or reload") + w.cleanupConnections(w.baseContext, true) + + // Wait for next status request to succeed. Don't wait too long; + // wrap the base context in a timeout equal to our status grace + // period. + w.logger.Debug("waiting for next status report to controller") + waitStatusStart := time.Now() + nextStatusCtx, nextStatusCancel := context.WithTimeout(w.baseContext, w.conf.StatusGracePeriodDuration) + defer nextStatusCancel() + for { + if err := nextStatusCtx.Err(); err != nil { + w.logger.Error("error waiting for next status report to controller", "err", err) + break + } + + if w.lastSuccessfulStatusTime().Sub(waitStatusStart) > 0 { + break + } + + time.Sleep(time.Second) + } + + // Proceed with remainder of shutdown. + w.logger.Debug("canceling base context and shutting down connection to controller") + w.baseCancel() + w.Resolver().UpdateState(resolver.State{Addresses: []resolver.Address{}}) + + w.started.Store(false) w.tickerWg.Wait() if w.conf.Eventer != nil { if err := w.conf.Eventer.FlushNodes(context.Background()); err != nil { return fmt.Errorf("error flushing worker eventer nodes: %w", err) } } + + w.logger.Debug("shutdown successful") return nil } diff --git a/internal/tests/cluster/session_cleanup_test.go b/internal/tests/cluster/session_cleanup_test.go index d2caed451..2ba8ce0d9 100644 --- a/internal/tests/cluster/session_cleanup_test.go +++ b/internal/tests/cluster/session_cleanup_test.go @@ -2,58 +2,20 @@ package cluster import ( "context" - "crypto/ed25519" - "crypto/tls" - "crypto/x509" - "encoding/binary" - "errors" - "fmt" - "io" "net" - "net/http" - "reflect" - "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/observability/event" - "github.com/hashicorp/boundary/internal/proxy" "github.com/hashicorp/boundary/internal/servers/controller" "github.com/hashicorp/boundary/internal/servers/worker" "github.com/hashicorp/boundary/internal/session" + "github.com/hashicorp/boundary/internal/tests/helper" "github.com/hashicorp/dawdle" - "github.com/hashicorp/go-cleanhttp" "github.com/hashicorp/go-hclog" - "github.com/hashicorp/go-secure-stdlib/base62" - "github.com/mr-tron/base58" - "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" - "google.golang.org/protobuf/proto" - "nhooyr.io/websocket" - "nhooyr.io/websocket/wspb" -) - -const ( - defaultGracePeriod = time.Second * 15 - expectConnectionStateOnControllerTimeout = time.Minute * 2 - expectConnectionStateOnWorkerTimeout = defaultGracePeriod * 2 - - // This is the interval that we check states on in the worker. It - // needs to be particularly granular to ensure that we allow for - // adequate time to catch any edge cases where the connection state - // disappears from the worker before we can update the state. - // - // As of this writing, the (hardcoded) status request interval on - // the worker is 2 seconds, and a session is removed from the state - // when it has no connections left on the *next* pass. A one second - // interval would only mean two chances to check with a high - // possibility of skew; while it seems okay I'm still not 100% - // comfortable with this little resolution. - expectConnectionStateOnWorkerInterval = time.Millisecond * 100 ) // timeoutBurdenType details our "burden cases" for the session @@ -91,18 +53,18 @@ var timeoutBurdenCases = []timeoutBurdenType{timeoutBurdenTypeDefault, timeoutBu func controllerGracePeriod(ty timeoutBurdenType) time.Duration { if ty == timeoutBurdenTypeWorker { - return defaultGracePeriod * 10 + return helper.DefaultGracePeriod * 10 } - return defaultGracePeriod + return helper.DefaultGracePeriod } func workerGracePeriod(ty timeoutBurdenType) time.Duration { if ty == timeoutBurdenTypeController { - return defaultGracePeriod * 10 + return helper.DefaultGracePeriod * 10 } - return defaultGracePeriod + return helper.DefaultGracePeriod } // TestSessionCleanup is the main test for session cleanup, and @@ -184,7 +146,7 @@ func testWorkerSessionCleanupSingle(burdenCase timeoutBurdenType) func(t *testin require.NotNil(tgt) // Create test server, update default port on target - ts := newTestTcpServer(t, logger) + ts := helper.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()), targets.WithSessionConnectionLimit(-1)) @@ -192,7 +154,7 @@ func testWorkerSessionCleanupSingle(burdenCase timeoutBurdenType) func(t *testin require.NotNil(tgt) // Authorize and connect - sess := newTestSession(ctx, t, logger, tcl, "ttcp_1234567890") + sess := helper.NewTestSession(ctx, t, logger, tcl, "ttcp_1234567890") sConn := sess.Connect(ctx, t) // Run initial send/receive test, make sure things are working @@ -208,18 +170,18 @@ func testWorkerSessionCleanupSingle(burdenCase timeoutBurdenType) func(t *testin case timeoutBurdenTypeWorker: // Wait on worker, then check controller sess.ExpectConnectionStateOnWorker(ctx, t, w1, session.StatusClosed) - sess.ExpectConnectionStateOnController(ctx, t, c1, session.StatusConnected) + sess.ExpectConnectionStateOnController(ctx, t, c1.Controller().SessionRepoFn, session.StatusConnected) case timeoutBurdenTypeController: // Wait on controller, then check worker - sess.ExpectConnectionStateOnController(ctx, t, c1, session.StatusClosed) + sess.ExpectConnectionStateOnController(ctx, t, c1.Controller().SessionRepoFn, session.StatusClosed) sess.ExpectConnectionStateOnWorker(ctx, t, w1, session.StatusConnected) default: // Should be closed on both worker and controller. Wait on // worker then check controller. sess.ExpectConnectionStateOnWorker(ctx, t, w1, session.StatusClosed) - sess.ExpectConnectionStateOnController(ctx, t, c1, session.StatusClosed) + sess.ExpectConnectionStateOnController(ctx, t, c1.Controller().SessionRepoFn, session.StatusClosed) } // Run send/receive test again to check expected connection-level @@ -249,7 +211,7 @@ func testWorkerSessionCleanupSingle(burdenCase timeoutBurdenType) func(t *testin // a connection status, ensure that our old session's // connections are actually closed now that the worker is // properly reporting in again. - sess.ExpectConnectionStateOnController(ctx, t, c1, session.StatusClosed) + sess.ExpectConnectionStateOnController(ctx, t, c1.Controller().SessionRepoFn, session.StatusClosed) case timeoutBurdenTypeController: // If we are expecting the controller to be the source of @@ -260,7 +222,7 @@ func testWorkerSessionCleanupSingle(burdenCase timeoutBurdenType) func(t *testin // Proceed with new connection test event.WriteSysEvent(ctx, op, "connecting to new session after resuming controller/worker link") - sess = newTestSession(ctx, t, logger, tcl, "ttcp_1234567890") // re-assign, other connection will close in t.Cleanup() + sess = helper.NewTestSession(ctx, t, logger, tcl, "ttcp_1234567890") // re-assign, other connection will close in t.Cleanup() sConn = sess.Connect(ctx, t) sConn.TestSendRecvAll(t) } @@ -367,7 +329,7 @@ func testWorkerSessionCleanupMulti(burdenCase timeoutBurdenType) func(t *testing require.NotNil(tgt) // Create test server, update default port on target - ts := newTestTcpServer(t, logger) + ts := helper.NewTestTcpServer(t, logger) require.NotNil(ts) defer ts.Close() tgt, err = tcl.Update(ctx, tgt.Item.Id, tgt.Item.Version, targets.WithTcpTargetDefaultPort(ts.Port()), targets.WithSessionConnectionLimit(-1)) @@ -375,7 +337,7 @@ func testWorkerSessionCleanupMulti(burdenCase timeoutBurdenType) func(t *testing require.NotNil(tgt) // Authorize and connect - sess := newTestSession(ctx, t, logger, tcl, "ttcp_1234567890") + sess := helper.NewTestSession(ctx, t, logger, tcl, "ttcp_1234567890") sConn := sess.Connect(ctx, t) // Run initial send/receive test, make sure things are working @@ -409,18 +371,18 @@ func testWorkerSessionCleanupMulti(burdenCase timeoutBurdenType) func(t *testing case timeoutBurdenTypeWorker: // Wait on worker, then check controller sess.ExpectConnectionStateOnWorker(ctx, t, w1, session.StatusClosed) - sess.ExpectConnectionStateOnController(ctx, t, c1, session.StatusConnected) + sess.ExpectConnectionStateOnController(ctx, t, c1.Controller().SessionRepoFn, session.StatusConnected) case timeoutBurdenTypeController: // Wait on controller, then check worker - sess.ExpectConnectionStateOnController(ctx, t, c1, session.StatusClosed) + sess.ExpectConnectionStateOnController(ctx, t, c1.Controller().SessionRepoFn, session.StatusClosed) sess.ExpectConnectionStateOnWorker(ctx, t, w1, session.StatusConnected) default: // Should be closed on both worker and controller. Wait on // worker then check controller. sess.ExpectConnectionStateOnWorker(ctx, t, w1, session.StatusClosed) - sess.ExpectConnectionStateOnController(ctx, t, c1, session.StatusClosed) + sess.ExpectConnectionStateOnController(ctx, t, c1.Controller().SessionRepoFn, session.StatusClosed) } // Run send/receive test again to check expected connection-level @@ -451,7 +413,7 @@ func testWorkerSessionCleanupMulti(burdenCase timeoutBurdenType) func(t *testing // a connection status, ensure that our old session's // connections are actually closed now that the worker is // properly reporting in again. - sess.ExpectConnectionStateOnController(ctx, t, c1, session.StatusClosed) + sess.ExpectConnectionStateOnController(ctx, t, c1.Controller().SessionRepoFn, session.StatusClosed) case timeoutBurdenTypeController: // If we are expecting the controller to be the source of @@ -462,447 +424,8 @@ func testWorkerSessionCleanupMulti(burdenCase timeoutBurdenType) func(t *testing // Proceed with new connection test event.WriteSysEvent(ctx, op, "connecting to new session after resuming controller/worker link") - sess = newTestSession(ctx, t, logger, tcl, "ttcp_1234567890") // re-assign, other connection will close in t.Cleanup() + sess = helper.NewTestSession(ctx, t, logger, tcl, "ttcp_1234567890") // re-assign, other connection will close in t.Cleanup() sConn = sess.Connect(ctx, t) sConn.TestSendRecvAll(t) } } - -// testSession represents an authorized session. -type testSession struct { - sessionId string - workerAddr string - transport *http.Transport - tofuToken string - connectionsLeft int32 - logger hclog.Logger -} - -// newTestSession authorizes a session and creates all of the data -// necessary to initialize. -func newTestSession( - ctx context.Context, - t *testing.T, - logger hclog.Logger, - 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 := &testSession{ - sessionId: sar.Item.SessionId, - logger: logger, - } - 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) -} - -// ExpectConnectionStateOnController waits until all connections in a -// session have transitioned to a particular state on the controller. -func (s *testSession) ExpectConnectionStateOnController( - ctx context.Context, - t *testing.T, - tc *controller.TestController, - expectState session.ConnectionStatus, -) { - const op = "cluster.(testSession).ExpectConnectionStateOnController" - t.Helper() - require := require.New(t) - assert := assert.New(t) - - ctx, cancel := context.WithTimeout(ctx, expectConnectionStateOnControllerTimeout) - defer cancel() - - // This is just for initialization of the actual state set. - const sessionStatusUnknown session.ConnectionStatus = "unknown" - - // Get all connections for the session on the controller. - sessionRepo, err := tc.Controller().SessionRepoFn() - require.NoError(err) - - conns, err := sessionRepo.ListConnectionsBySessionId(ctx, s.sessionId) - require.NoError(err) - // To avoid misleading passing tests, we require this test be used - // with sessions with connections.. - require.Greater(len(conns), 0, "should have at least one connection") - - // Make a set of states, 1 per connection - actualStates := make([]session.ConnectionStatus, len(conns)) - for i := range actualStates { - actualStates[i] = sessionStatusUnknown - } - - // Make expect set for comparison - expectStates := make([]session.ConnectionStatus, len(conns)) - for i := range expectStates { - expectStates[i] = expectState - } - - for { - if ctx.Err() != nil { - break - } - - for i, conn := range conns { - _, states, err := sessionRepo.LookupConnection(ctx, conn.PublicId, nil) - require.NoError(err) - // Look at the first state in the returned list, which will - // be the most recent state. - actualStates[i] = states[0].Status - } - - if reflect.DeepEqual(expectStates, actualStates) { - break - } - - time.Sleep(time.Second) - } - - // "non-fatal" assert here, so that we can surface both timeouts - // and invalid state - assert.NoError(ctx.Err()) - - // Assert - require.Equal(expectStates, actualStates) - event.WriteSysEvent(ctx, op, "successfully asserted all connection states on controller", "expected_states", expectStates, "actual_states", actualStates) -} - -// ExpectConnectionStateOnWorker waits until all connections in a -// session have transitioned to a particular state on the worker. -func (s *testSession) ExpectConnectionStateOnWorker( - ctx context.Context, - t *testing.T, - tw *worker.TestWorker, - expectState session.ConnectionStatus, -) { - const op = "cluster.(testSession).ExpectConnectionStateOnWorker" - t.Helper() - require := require.New(t) - assert := assert.New(t) - - ctx, cancel := context.WithTimeout(ctx, expectConnectionStateOnWorkerTimeout) - defer cancel() - - // This is just for initialization of the actual state set. - const sessionStatusUnknown session.ConnectionStatus = "unknown" - - connIds := s.testWorkerConnectionIds(t, tw) - // Make a set of states, 1 per connection - actualStates := make(map[string]session.ConnectionStatus, len(connIds)) - for _, id := range connIds { - actualStates[id] = sessionStatusUnknown - } - - // Make expect set for comparison - expectStates := make(map[string]session.ConnectionStatus, len(connIds)) - for _, id := range connIds { - expectStates[id] = expectState - } - - for { - if ctx.Err() != nil { - break - } - - for _, conn := range s.testWorkerConnectionInfo(t, tw) { - actualStates[conn.Id] = session.ConnectionStatusFromProtoVal(conn.Status) - } - - if reflect.DeepEqual(expectStates, actualStates) { - break - } - - time.Sleep(expectConnectionStateOnWorkerInterval) - } - - // "non-fatal" assert here, so that we can surface both timeouts - // and invalid state - assert.NoError(ctx.Err()) - - // Assert - require.Equal(expectStates, actualStates) - event.WriteSysEvent(ctx, op, "successfully asserted all connection states on worker", "expected_states", expectStates, "actual_states", actualStates) -} - -func (s *testSession) testWorkerConnectionInfo(t *testing.T, tw *worker.TestWorker) map[string]worker.TestConnectionInfo { - t.Helper() - require := require.New(t) - si, ok := tw.LookupSession(s.sessionId) - // This is always an error if the session has been removed from the - // local state. - require.True(ok) - - // Likewise, we require the helper to always be used with - // connections. - require.Greater(len(si.Connections), 0, "should have at least one connection") - - return si.Connections -} - -func (s *testSession) testWorkerConnectionIds(t *testing.T, tw *worker.TestWorker) []string { - t.Helper() - conns := s.testWorkerConnectionInfo(t, tw) - result := make([]string, 0, len(conns)) - for _, conn := range conns { - result = append(result, conn.Id) - } - - return result -} - -// 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 -) *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: s.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 { - const op = "cluster.(testSessionConnection).testSendRecv" - ctx := context.TODO() - t.Helper() - require := require.New(t) - - // This is a fairly arbitrary value, as the send/recv is - // instantaneous. The main key here is just to make sure that we do - // it a reasonable amount of times to know the connection is - // stable. - const testSendRecvSendMax = 100 - 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 { - event.WriteSysEvent(ctx, op, "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 { - event.WriteSysEvent(ctx, op, "received error during read", "err", err, "num_successfully_sent", i) - 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) - } - - event.WriteSysEvent(ctx, op, "finished send/recv successfully", "num_successfully_sent", testSendRecvSendMax) - return true -} - -// TestSendRecvAll asserts that we were able to send/recv all pings -// over the test connection. -func (c *testSessionConnection) TestSendRecvAll(t *testing.T) { - const op = "cluster.(testSessionConnection).TestSendRecvAll" - t.Helper() - require.True(t, c.testSendRecv(t)) - event.WriteSysEvent(context.TODO(), op, "successfully asserted send/recv as passing") -} - -// TestSendRecvFail asserts that we were able to send/recv all pings -// over the test connection. -func (c *testSessionConnection) TestSendRecvFail(t *testing.T) { - const op = "cluster.(testSessionConnection).TestSendRecvFail" - t.Helper() - require.False(t, c.testSendRecv(t)) - event.WriteSysEvent(context.TODO(), op, "successfully asserted send/recv as failing") -} - -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() { - const op = "cluster.(testTcpServer).run" - for { - conn, err := ts.ln.Accept() - if err != nil { - if !errors.Is(err, net.ErrClosed) { - event.WriteError(context.TODO(), op, err, event.WithInfoMsg("Accept() error in testTcpServer")) - } - - 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 -} diff --git a/internal/tests/helper/testing_helper.go b/internal/tests/helper/testing_helper.go new file mode 100644 index 000000000..95ff0808a --- /dev/null +++ b/internal/tests/helper/testing_helper.go @@ -0,0 +1,490 @@ +package helper + +import ( + "context" + "crypto/ed25519" + "crypto/tls" + "crypto/x509" + "encoding/binary" + "errors" + "fmt" + "io" + "net" + "net/http" + "reflect" + "strconv" + "testing" + "time" + + "github.com/hashicorp/boundary/api/targets" + "github.com/hashicorp/boundary/globals" + targetspb "github.com/hashicorp/boundary/internal/gen/controller/api/resources/targets" + "github.com/hashicorp/boundary/internal/proxy" + "github.com/hashicorp/boundary/internal/servers/controller/common" + "github.com/hashicorp/boundary/internal/servers/worker" + "github.com/hashicorp/boundary/internal/session" + "github.com/hashicorp/go-cleanhttp" + "github.com/hashicorp/go-hclog" + "github.com/hashicorp/go-secure-stdlib/base62" + "github.com/mr-tron/base58" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "google.golang.org/protobuf/proto" + "nhooyr.io/websocket" + "nhooyr.io/websocket/wspb" +) + +const ( + DefaultGracePeriod = time.Second * 15 + expectConnectionStateOnControllerTimeout = time.Minute * 2 + expectConnectionStateOnWorkerTimeout = DefaultGracePeriod * 2 + + // This is the interval that we check states on in the worker. It + // needs to be particularly granular to ensure that we allow for + // adequate time to catch any edge cases where the connection state + // disappears from the worker before we can update the state. + // + // As of this writing, the (hardcoded) status request interval on + // the worker is 2 seconds, and a session is removed from the state + // when it has no connections left on the *next* pass. A one second + // interval would only mean two chances to check with a high + // possibility of skew; while it seems okay I'm still not 100% + // comfortable with this little resolution. + expectConnectionStateOnWorkerInterval = time.Millisecond * 100 +) + +// TestSession represents an authorized session. +type TestSession struct { + sessionId string + workerAddr string + transport *http.Transport + tofuToken string + connectionsLeft int32 + logger hclog.Logger +} + +// NewTestSession authorizes a session and creates all of the data +// necessary to initialize. +func NewTestSession( + ctx context.Context, + t *testing.T, + logger hclog.Logger, + 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 := &TestSession{ + sessionId: sar.Item.SessionId, + logger: logger, + } + 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) +} + +// ExpectConnectionStateOnController waits until all connections in a +// session have transitioned to a particular state on the controller. +func (s *TestSession) ExpectConnectionStateOnController( + ctx context.Context, + t *testing.T, + sessionRepoFn common.SessionRepoFactory, + expectState session.ConnectionStatus, +) { + t.Helper() + require := require.New(t) + assert := assert.New(t) + + ctx, cancel := context.WithTimeout(ctx, expectConnectionStateOnControllerTimeout) + defer cancel() + + // This is just for initialization of the actual state set. + const sessionStatusUnknown session.ConnectionStatus = "unknown" + + sessionRepo, err := sessionRepoFn() + require.NoError(err) + + conns, err := sessionRepo.ListConnectionsBySessionId(ctx, s.sessionId) + require.NoError(err) + // To avoid misleading passing tests, we require this test be used + // with sessions with connections.. + require.Greater(len(conns), 0, "should have at least one connection") + + // Make a set of states, 1 per connection + actualStates := make([]session.ConnectionStatus, len(conns)) + for i := range actualStates { + actualStates[i] = sessionStatusUnknown + } + + // Make expect set for comparison + expectStates := make([]session.ConnectionStatus, len(conns)) + for i := range expectStates { + expectStates[i] = expectState + } + + for { + if ctx.Err() != nil { + break + } + + for i, conn := range conns { + _, states, err := sessionRepo.LookupConnection(ctx, conn.PublicId, nil) + require.NoError(err) + // Look at the first state in the returned list, which will + // be the most recent state. + actualStates[i] = states[0].Status + } + + if reflect.DeepEqual(expectStates, actualStates) { + break + } + + time.Sleep(time.Second) + } + + // "non-fatal" assert here, so that we can surface both timeouts + // and invalid state + assert.NoError(ctx.Err()) + + // Assert + require.Equal(expectStates, actualStates) + s.logger.Debug("successfully asserted all connection states on controller", "expected_states", expectStates, "actual_states", actualStates) +} + +// ExpectConnectionStateOnWorker waits until all connections in a +// session have transitioned to a particular state on the worker. +func (s *TestSession) ExpectConnectionStateOnWorker( + ctx context.Context, + t *testing.T, + tw *worker.TestWorker, + expectState session.ConnectionStatus, +) { + t.Helper() + require := require.New(t) + assert := assert.New(t) + + ctx, cancel := context.WithTimeout(ctx, expectConnectionStateOnWorkerTimeout) + defer cancel() + + // This is just for initialization of the actual state set. + const sessionStatusUnknown session.ConnectionStatus = "unknown" + + connIds := s.testWorkerConnectionIds(t, tw) + // Make a set of states, 1 per connection + actualStates := make(map[string]session.ConnectionStatus, len(connIds)) + for _, id := range connIds { + actualStates[id] = sessionStatusUnknown + } + + // Make expect set for comparison + expectStates := make(map[string]session.ConnectionStatus, len(connIds)) + for _, id := range connIds { + expectStates[id] = expectState + } + + for { + if ctx.Err() != nil { + break + } + + for _, conn := range s.testWorkerConnectionInfo(t, tw) { + actualStates[conn.Id] = session.ConnectionStatusFromProtoVal(conn.Status) + } + + if reflect.DeepEqual(expectStates, actualStates) { + break + } + + time.Sleep(expectConnectionStateOnWorkerInterval) + } + + // "non-fatal" assert here, so that we can surface both timeouts + // and invalid state + assert.NoError(ctx.Err()) + + // Assert + require.Equal(expectStates, actualStates) + s.logger.Debug("successfully asserted all connection states on worker", "expected_states", expectStates, "actual_states", actualStates) +} + +func (s *TestSession) testWorkerConnectionInfo(t *testing.T, tw *worker.TestWorker) map[string]worker.TestConnectionInfo { + t.Helper() + require := require.New(t) + si, ok := tw.LookupSession(s.sessionId) + // This is always an error if the session has been removed from the + // local state. + require.True(ok) + + // Likewise, we require the helper to always be used with + // connections. + require.Greater(len(si.Connections), 0, "should have at least one connection") + + return si.Connections +} + +func (s *TestSession) testWorkerConnectionIds(t *testing.T, tw *worker.TestWorker) []string { + t.Helper() + conns := s.testWorkerConnectionInfo(t, tw) + result := make([]string, 0, len(conns)) + for _, conn := range conns { + result = append(result, conn.Id) + } + + return result +} + +// 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 +) *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: s.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) + + // This is a fairly arbitrary value, as the send/recv is + // instantaneous. The main key here is just to make sure that we do + // it a reasonable amount of times to know the connection is + // stable. + const testSendRecvSendMax = 100 + 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, "num_successfully_sent", i) + 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) + } + + c.logger.Debug("finished send/recv successfully", "num_successfully_sent", testSendRecvSendMax) + 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)) + c.logger.Debug("successfully asserted send/recv as passing") +} + +// 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)) + c.logger.Debug("successfully asserted send/recv as failing") +} + +// TestTcpServer defines a simple testing TCP server. Data is simply +// copied from its input to its output. +type TestTcpServer struct { + logger hclog.Logger + ln net.Listener + conns map[string]net.Conn +} + +// Port returns the port number for the test server listener. +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) +} + +// Close closes the listener and all active connections. +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) + } +} + +// NewTestTcpServer creates a TestTcpServer and runs it. +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 +}