Function TestNewServer_CloseActiveConnections

agent/agentssh/agentssh_test.go:150–256

(t *testing.T)

}

func TestNewServer_CloseActiveConnections(t *testing.T) {
	t.Parallel()

	prepare := func(ctx context.Context, t *testing.T) (*agentssh.Server, func()) {
		t.Helper()
		logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}).Leveled(slog.LevelDebug)
		s, err := agentssh.NewServer(ctx, logger, prometheus.NewRegistry(), afero.NewMemMapFs(), agentexec.DefaultExecer, nil)
		require.NoError(t, err)
		t.Cleanup(func() {
			_ = s.Close()
		})
		err = s.UpdateHostSigner(42)
		assert.NoError(t, err)

		ln, err := net.Listen("tcp", "127.0.0.1:0")
		require.NoError(t, err)

		waitConns := make([]chan struct{}, 4)

		var wg sync.WaitGroup
		wg.Add(1 + len(waitConns))

		go func() {
			defer wg.Done()
			err := s.Serve(ln)
			assert.Error(t, err) // Server is closed.
		}()

		for i := 0; i < len(waitConns); i++ {
			waitConns[i] = make(chan struct{})
			go func(ch chan struct{}) {
				defer wg.Done()
				c := sshClient(t, ln.Addr().String())
				sess, err := c.NewSession()
				assert.NoError(t, err)
				pty := ptytest.New(t)
				sess.Stdin = pty.Input()
				sess.Stdout = pty.Output()
				sess.Stderr = pty.Output()

				// Every other session will request a PTY.
				if i%2 == 0 {
					err = sess.RequestPty("xterm", 80, 80, nil)
					assert.NoError(t, err)
				}
				// The 60 seconds here is intended to be longer than the
				// test. The shutdown should propagate.
				if runtime.GOOS == "windows" {
					// Best effort to at least partially test this in Windows.
					err = sess.Start("echo start\"ed\" && sleep 60")
				} else {
					err = sess.Start("/bin/bash -c 'trap \"sleep 60\" SIGTERM; echo start\"ed\"; sleep 60'")
				}
				assert.NoError(t, err)

				// Allow the session to settle (i.e. reach echo).
				pty.ExpectMatchContext(ctx, "started")
				// Sleep a bit to ensure the sleep has started.