fix: close SSH sessions bottom-up if top-down fails

spikecurtis · spikecurtis · commit d9dc47ae8aa8 · 2024-09-16T10:20:47.000+04:00
diff --git a/cli/ssh.go b/cli/ssh.go
@@ -37,6 +37,7 @@ import (
 	"github.com/coder/coder/v2/codersdk/workspacesdk"
 	"github.com/coder/coder/v2/cryptorand"
 	"github.com/coder/coder/v2/pty"
+	"github.com/coder/quartz"
 	"github.com/coder/retry"
 	"github.com/coder/serpent"
 )
@@ -48,6 +49,7 @@ const (
 var (
 	workspacePollInterval   = time.Minute
 	autostopNotifyCountdown = []time.Duration{30 * time.Minute}
+	gracefulShutdownTimeout = 5 * time.Second
 )
 
 func (r *RootCmd) ssh() *serpent.Command {
@@ -250,7 +252,16 @@ func (r *RootCmd) ssh() *serpent.Command {
 			if err != nil {
 				return xerrors.Errorf("dial agent: %w", err)
 			}
-			if err = stack.push("agent conn", conn); err != nil {
+			if err = stack.push(
+				"agent conn",
+				// We set a long TCP timeout on SSH connections, which means if the underlying
+				// network fails, the SSH layer can hang for a really long time trying to send a
+				// shutdown message for any remote forwards (https://github.com/golang/go/issues/69484)
+				// Normally, we want to tear stuff down top to bottom, but if we get stuck doing it
+				// that way, this timeoutCloser will trip and close the underlying connection,
+				// bottom-up.
+				newTimeoutCloser(ctx, logger, gracefulShutdownTimeout, conn, quartz.NewReal()),
+			); err != nil {
 				return err
 			}
 			conn.AwaitReachable(ctx)
@@ -1085,3 +1096,49 @@ func getUsageAppName(usageApp string) codersdk.UsageAppName {
 
 	return codersdk.UsageAppNameSSH
 }
+
+type timeoutCloser struct {
+	target      io.Closer
+	closeCalled chan struct{}
+
+	// for testing
+	clock quartz.Clock
+}
+
+func newTimeoutCloser(
+	ctx context.Context, logger slog.Logger, timeout time.Duration, target io.Closer, clock quartz.Clock,
+) *timeoutCloser {
+	b := &timeoutCloser{
+		target:      target,
+		closeCalled: make(chan struct{}),
+		clock:       clock,
+	}
+	go b.waitForCtxOrClose(ctx, logger, timeout)
+	return b
+}
+
+func (t *timeoutCloser) waitForCtxOrClose(ctx context.Context, logger slog.Logger, timeout time.Duration) {
+	select {
+	case <-t.closeCalled:
+		return
+	case <-ctx.Done():
+	}
+	tmr := t.clock.NewTimer(timeout, "timeoutCloser", "waitForCtxOrClose")
+	defer tmr.Stop()
+	select {
+	case <-t.closeCalled:
+		return
+	case <-tmr.C:
+		logger.Warn(ctx, "timed out waiting for graceful shutdown")
+		err := t.target.Close()
+		if err != nil {
+			logger.Debug(ctx, "error closing target", slog.Error(err))
+		}
+	}
+}
+
+// Close should only be called at most once, e.g. in the closerStack
+func (t *timeoutCloser) Close() error {
+	close(t.closeCalled)
+	return t.target.Close()
+}
diff --git a/cli/ssh_internal_test.go b/cli/ssh_internal_test.go
@@ -12,6 +12,7 @@ import (
 
 	"cdr.dev/slog"
 	"cdr.dev/slog/sloggers/slogtest"
+	"github.com/coder/quartz"
 
 	"github.com/coder/coder/v2/codersdk"
 	"github.com/coder/coder/v2/testutil"
@@ -192,3 +193,39 @@ func (c *asyncCloser) Close() error {
 		return nil
 	}
 }
+
+func TestTimeoutCloser_Close(t *testing.T) {
+	t.Parallel()
+	ctx := testutil.Context(t, testutil.WaitShort)
+	logger := slogtest.Make(t, nil).Leveled(slog.LevelDebug)
+	mClock := quartz.NewMock(t)
+	closes := new([]*fakeCloser)
+	fc0 := &fakeCloser{closes: closes}
+	uut := newTimeoutCloser(ctx, logger, time.Second, fc0, mClock)
+	err := uut.Close()
+	require.NoError(t, err)
+	require.Equal(t, []*fakeCloser{fc0}, *closes, "should close fc0")
+}
+
+func TestTimeoutCloser_Timeout(t *testing.T) {
+	t.Parallel()
+	testCtx := testutil.Context(t, testutil.WaitShort)
+	ctx, cancel := context.WithCancel(testCtx)
+	defer cancel()
+	logger := slogtest.Make(t, nil).Leveled(slog.LevelDebug)
+	mClock := quartz.NewMock(t)
+	trap := mClock.Trap().NewTimer("timeoutCloser", "waitForCtxOrClose")
+	defer trap.Close()
+	ac := &asyncCloser{
+		t:        t,
+		ctx:      testCtx,
+		complete: make(chan struct{}),
+		started:  make(chan struct{}),
+	}
+	_ = newTimeoutCloser(ctx, logger, time.Second, ac, mClock)
+	cancel()
+	trap.MustWait(testCtx).Release()
+	mClock.Advance(time.Second).MustWait(testCtx)
+	testutil.RequireRecvCtx(testCtx, t, ac.started)
+	close(ac.complete)
+}
