Close connections on worker on shutdown, add force shutdown (#1354)

This adds functionality to ensure that connections are closed on a
worker when the server is shut down or reloaded.

Connections are closed, and then the shutdown process waits for the next
report to the controller to go through, waiting the grace period that
was added in #1330. The rest of shutdown is then proceeded with.
Listener shutdown has been moved to the start of the process to ensure
that new connections can't be made while we are shutting down.

Additionally, functionality has been added to force shut down servers
after a second interrupt has been received after the first one. This
takes the form of code in both the server and dev mode code paths that
just performs an os.Exit if another shutdown is received during the
normal shutdown process.

Additionally, the commit also includes some refactoring to tests so that
we could use some of the new E2E testing code to test session
connections during shutdown and reload.

Additionally, this fixes a panic in worker startup when there are no
controllers defined. Looks like the shutdown call was attempting to call
into the controller, not the worker.

Finally, this also removes the extra shutdown within the outer worker
startup.

Shutdown is actually handled within the inner StartWorker(), attempting
shutdown on errors was actually leading to potential use of an
uninitialized worker, which is leading to panics in testing.

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 {

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 {

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())

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

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")
+	}
+}

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())

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
 }
 

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
-}

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
+}