package agent_test

import (

"bufio"

"bytes"

"context"

"encoding/json"

"errors"

"fmt"

"io"

"net"

"net/http"

"net/http/httptest"

"net/netip"

"os"

"os/exec"

"os/user"

"path"

"path/filepath"

"regexp"

"runtime"

"slices"

"strconv"

"strings"

"testing"

"time"

"github.com/bramvdbogaerde/go-scp"

"github.com/google/uuid"

"github.com/ory/dockertest/v3"

"github.com/ory/dockertest/v3/docker"

"github.com/pion/udp"

"github.com/pkg/sftp"

"github.com/prometheus/client_golang/prometheus"

promgo "github.com/prometheus/client_model/go"

"github.com/spf13/afero"

"github.com/stretchr/testify/assert"

"github.com/stretchr/testify/require"

"go.uber.org/goleak"

"golang.org/x/crypto/ssh"

"golang.org/x/xerrors"

"tailscale.com/net/speedtest"

"tailscale.com/tailcfg"

"cdr.dev/slog/v3"

"cdr.dev/slog/v3/sloggers/slogtest"

"github.com/coder/coder/v2/agent"

"github.com/coder/coder/v2/agent/agentcontainers"

"github.com/coder/coder/v2/agent/agentssh"

"github.com/coder/coder/v2/agent/agenttest"

"github.com/coder/coder/v2/agent/proto"

"github.com/coder/coder/v2/agent/usershell"

"github.com/coder/coder/v2/codersdk"

"github.com/coder/coder/v2/codersdk/agentsdk"

"github.com/coder/coder/v2/codersdk/workspacesdk"

"github.com/coder/coder/v2/cryptorand"

"github.com/coder/coder/v2/tailnet"

"github.com/coder/coder/v2/tailnet/tailnettest"

"github.com/coder/coder/v2/testutil"

"github.com/coder/coder/v2/testutil/expecter"

"github.com/coder/quartz"

)

func TestMain(m *testing.M) {

if os.Getenv("CODER_TEST_RUN_SUB_AGENT_MAIN") == "1" {

// If we're running as a subagent, we don't want to run the main tests.

// Instead, we just run the subagent tests.

exit := runSubAgentMain()

os.Exit(exit)

}

goleak.VerifyTestMain(m, testutil.GoleakOptions...)

}

var sshPorts = []uint16{workspacesdk.AgentSSHPort, workspacesdk.AgentStandardSSHPort}

// TestAgent_CloseWhileStarting is a regression test for https://github.com/coder/coder/issues/17328

func TestAgent_ImmediateClose(t *testing.T) {

t.Parallel()

ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)

defer cancel()

logger := slogtest.Make(t, &slogtest.Options{

// Agent can drop errors when shutting down, and some, like the

// fasthttplistener connection closed error, are unexported.

IgnoreErrors: true,

}).Leveled(slog.LevelDebug)

manifest := agentsdk.Manifest{

AgentID: uuid.New(),

AgentName: "test-agent",

WorkspaceName: "test-workspace",

WorkspaceID: uuid.New(),

}

coordinator := tailnet.NewCoordinator(logger)

t.Cleanup(func() {

_ = coordinator.Close()

})

statsCh := make(chan *proto.Stats, 50)

fs := afero.NewMemMapFs()

client := agenttest.NewClient(t, logger.Named("agenttest"), manifest.AgentID, manifest, statsCh, coordinator)

t.Cleanup(client.Close)

options := agent.Options{

Client: client,

Filesystem: fs,

Logger: logger.Named("agent"),

ReconnectingPTYTimeout: 0,

EnvironmentVariables: map[string]string{},

}

agentUnderTest := agent.New(options)

t.Cleanup(func() {

_ = agentUnderTest.Close()

})

// wait until the agent has connected and is starting to find races in the startup code

_ = testutil.TryReceive(ctx, t, client.GetStartup())

t.Log("Closing Agent")

err := agentUnderTest.Close()

require.NoError(t, err)

}

// NOTE(Cian): I noticed that these tests would fail when my default shell was zsh.

// Writing "exit 0" to stdin before closing fixed the issue for me.

func TestAgent_Stats_SSH(t *testing.T) {

t.Parallel()

for _, port := range sshPorts {

t.Run(fmt.Sprintf("(:%d)", port), func(t *testing.T) {

t.Parallel()

ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)

defer cancel()

//nolint:dogsled

conn, _, stats, _, _ := setupAgent(t, agentsdk.Manifest{}, 0)

sshClient, err := conn.SSHClientOnPort(ctx, port)

require.NoError(t, err)

defer sshClient.Close()

session, err := sshClient.NewSession()

require.NoError(t, err)

defer session.Close()

stdin, err := session.StdinPipe()

require.NoError(t, err)

err = session.Shell()

require.NoError(t, err)

// Generate SSH traffic so the connstats window sees the session.

_, err = stdin.Write([]byte("echo test\n"))

require.NoError(t, err)

assertSSHStats(t, stats)

_, err = stdin.Write([]byte("exit 0\n"))

require.NoError(t, err, "writing exit to stdin")

_ = stdin.Close()

err = session.Wait()

require.NoError(t, err, "waiting for session to exit")

})

}

// Regression test for CODAGT-517: the barrier blocks reportLoop's

// initial UpdateStats, so on unfixed code the connstats callback is

// never installed and handshake traffic is lost. On fixed code the

// callback is installed at creation, so traffic is captured.

t.Run("StatsCallbackRace", func(t *testing.T) {

t.Parallel()

ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)

defer cancel()

barrier := make(chan struct{})

//nolint:dogsled

conn, _, stats, _, _ := setupAgent(t, agentsdk.Manifest{}, 0,

func(c *agenttest.Client, _ *agent.Options) {

c.SetUpdateStatsOverride(func(

ctx context.Context,

req *proto.UpdateStatsRequest,

next func(context.Context, *proto.UpdateStatsRequest) (*proto.UpdateStatsResponse, error),

) (*proto.UpdateStatsResponse, error) {

if req.Stats == nil {

select {

case <-ctx.Done():

return nil, ctx.Err()

case <-barrier:

}

}

return next(ctx, req)

})

},

)

// Connect SSH while the barrier holds reportLoop blocked.

sshClient, err := conn.SSHClientOnPort(ctx, workspacesdk.AgentStandardSSHPort)

require.NoError(t, err)

defer sshClient.Close()

session, err := sshClient.NewSession()

require.NoError(t, err)

defer session.Close()

stdin, err := session.StdinPipe()

require.NoError(t, err)

err = session.Shell()

require.NoError(t, err)

// Shell must be idle so the only traffic is the SSH handshake.

close(barrier)

assertSSHStats(t, stats)

_, err = stdin.Write([]byte("exit 0\n"))

require.NoError(t, err, "writing exit to stdin")

_ = stdin.Close()

err = session.Wait()

require.NoError(t, err, "waiting for session to exit")

})

}

// assertSSHStats waits for ConnectionCount, RxBytes, TxBytes, and

// SessionCountSsh to be nonzero on the stats channel.

func assertSSHStats(t *testing.T, stats <-chan *proto.Stats) {

t.Helper()

var connectionCountSeen, rxBytesSeen, txBytesSeen, sessionCountSSHSeen bool

require.Eventuallyf(t, func() bool {

s, ok := <-stats

if !ok {

return false

}

t.Logf("got stats: ConnectionCount=%d, RxBytes=%d, TxBytes=%d, SessionCountSsh=%d",

s.ConnectionCount, s.RxBytes, s.TxBytes, s.SessionCounts["ssh"])

if s.ConnectionCount > 0 {

connectionCountSeen = true

}

if s.RxBytes > 0 {

rxBytesSeen = true

}

if s.TxBytes > 0 {

txBytesSeen = true

}

if s.SessionCounts["ssh"] == 1 {

sessionCountSSHSeen = true

}

return connectionCountSeen && rxBytesSeen && txBytesSeen && sessionCountSSHSeen

}, testutil.WaitLong, testutil.IntervalFast,

"never saw all SSH stats",

)

}

func TestAgent_Stats_ReconnectingPTY(t *testing.T) {

t.Parallel()

ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)

defer cancel()

//nolint:dogsled

conn, _, stats, _, _ := setupAgent(t, agentsdk.Manifest{}, 0)

ptyConn, err := conn.ReconnectingPTY(ctx, uuid.New(), 128, 128, "bash")

require.NoError(t, err)

defer ptyConn.Close()

data, err := json.Marshal(workspacesdk.ReconnectingPTYRequest{

Data: "echo test\r\n",

})

require.NoError(t, err)

_, err = ptyConn.Write(data)

require.NoError(t, err)

var s *proto.Stats

// We are looking for four different stats to be reported. They might not all

// arrive at the same time, so we loop until we've seen them all.

var connectionCountSeen, rxBytesSeen, txBytesSeen, sessionCountReconnectingPTYSeen bool

require.Eventuallyf(t, func() bool {

var ok bool

s, ok = <-stats

if !ok {

return false

}

if s.ConnectionCount > 0 {

connectionCountSeen = true

}

if s.RxBytes > 0 {

rxBytesSeen = true

}

if s.TxBytes > 0 {

txBytesSeen = true

}

if s.SessionCounts["reconnecting_pty"] == 1 {

sessionCountReconnectingPTYSeen = true

}

return connectionCountSeen && rxBytesSeen && txBytesSeen && sessionCountReconnectingPTYSeen

}, testutil.WaitLong, testutil.IntervalFast,

"never saw all stats: %+v, saw connectionCount: %t, rxBytes: %t, txBytes: %t, sessionCountReconnectingPTY: %t",

s, connectionCountSeen, rxBytesSeen, txBytesSeen, sessionCountReconnectingPTYSeen,

)

}

func TestAgent_Stats_Magic(t *testing.T) {

t.Parallel()

t.Run("StripsEnvironmentVariable", func(t *testing.T) {

t.Parallel()

ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)

defer cancel()

//nolint:dogsled

conn, _, _, _, _ := setupAgent(t, agentsdk.Manifest{}, 0)

sshClient, err := conn.SSHClient(ctx)

require.NoError(t, err)

defer sshClient.Close()

session, err := sshClient.NewSession()

require.NoError(t, err)

session.Setenv(agentssh.AppNameEnvironmentVariable, string(codersdk.AppFamilyVSCode))

defer session.Close()

command := "sh -c 'echo $" + agentssh.AppNameEnvironmentVariable + "'"

expected := ""

if runtime.GOOS == "windows" {

expected = "%" + agentssh.AppNameEnvironmentVariable + "%"

command = "cmd.exe /c echo " + expected

}

output, err := session.Output(command)

require.NoError(t, err)

require.Equal(t, expected, strings.TrimSpace(string(output)))

})

t.Run("TracksVSCode", func(t *testing.T) {

t.Parallel()

if runtime.GOOS == "windows" {

t.Skip("Sleeping for infinity doesn't work on Windows")

}

ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)

defer cancel()

//nolint:dogsled

conn, agentClient, stats, _, _ := setupAgent(t, agentsdk.Manifest{}, 0)

sshClient, err := conn.SSHClient(ctx)

require.NoError(t, err)

defer sshClient.Close()

session, err := sshClient.NewSession()

require.NoError(t, err)

session.Setenv(agentssh.AppNameEnvironmentVariable, string(codersdk.AppFamilyVSCode))

defer session.Close()

stdin, err := session.StdinPipe()

require.NoError(t, err)

err = session.Shell()

require.NoError(t, err)

require.Eventuallyf(t, func() bool {

s, ok := <-stats

t.Logf("got stats: ok=%t, ConnectionCount=%d, RxBytes=%d, TxBytes=%d, SessionCounts[vscode]=%d, ConnectionMedianLatencyMS=%f",

ok, s.ConnectionCount, s.RxBytes, s.TxBytes, s.SessionCounts["vscode"], s.ConnectionMedianLatencyMs)

return ok &&

// Ensure that the connection didn't count as a "normal" SSH session.

// This was a special one, so it should be labeled specially in the stats!

s.SessionCounts["vscode"] == 1 &&

// Ensure that connection latency is being counted!

// If it isn't, it's set to -1.

s.ConnectionMedianLatencyMs >= 0

}, testutil.WaitLong, testutil.IntervalFast,

"never saw stats",

)

_, err = stdin.Write([]byte("exit 0\n"))

require.NoError(t, err, "writing exit to stdin")

_ = stdin.Close()

err = session.Wait()

require.NoError(t, err)

assertConnectionReport(t, agentClient, proto.Connection_VSCODE, 0, "")

})

t.Run("TracksJetBrains", func(t *testing.T) {

t.Parallel()

if runtime.GOOS != "linux" {

t.Skip("JetBrains tracking is only supported on Linux")

}

ctx := testutil.Context(t, testutil.WaitLong)

// JetBrains tracking works by looking at the process name listening on the

// forwarded port. If the process's command line includes the magic string

// we are looking for, then we assume it is a JetBrains editor. So when we

// connect to the port we must ensure the process includes that magic string

// to fool the agent into thinking this is JetBrains. To do this we need to

// spawn an external process (in this case a simple echo server) so we can

// control the process name. The -D here is just to mimic how Java options

// are set but is not necessary as the agent looks only for the magic

// string itself anywhere in the command.

_, b, _, ok := runtime.Caller(0)

require.True(t, ok)

dir := filepath.Join(filepath.Dir(b), "../scripts/echoserver/main.go")

echoServerCmd := exec.Command("go", "run", dir,

"-D", agentssh.MagicProcessCmdlineJetBrains)

stdout, err := echoServerCmd.StdoutPipe()

require.NoError(t, err)

err = echoServerCmd.Start()

require.NoError(t, err)

defer echoServerCmd.Process.Kill()

// The echo server prints its port as the first line.

sc := bufio.NewScanner(stdout)

sc.Scan()

remotePort := sc.Text()

//nolint:dogsled

conn, agentClient, stats, _, _ := setupAgent(t, agentsdk.Manifest{}, 0)

sshClient, err := conn.SSHClient(ctx)

require.NoError(t, err)

tunneledConn, err := sshClient.Dial("tcp", fmt.Sprintf("127.0.0.1:%s", remotePort))

require.NoError(t, err)

t.Cleanup(func() {

// always close on failure of test

_ = conn.Close()

_ = tunneledConn.Close()

})

require.Eventuallyf(t, func() bool {

s, ok := <-stats

t.Logf("got stats with conn open: ok=%t, ConnectionCount=%d, SessionCountJetBrains=%d",

ok, s.ConnectionCount, s.SessionCounts["jetbrains"])

return ok && s.SessionCounts["jetbrains"] == 1

}, testutil.WaitLong, testutil.IntervalFast,

"never saw stats with conn open",

)

// Kill the server and connection after checking for the echo.

requireEcho(t, tunneledConn)

_ = echoServerCmd.Process.Kill()

_ = tunneledConn.Close()

require.Eventuallyf(t, func() bool {

s, ok := <-stats

t.Logf("got stats after disconnect %t, %d",

ok, s.SessionCounts["jetbrains"])

return ok &&

s.SessionCounts["jetbrains"] == 0

}, testutil.WaitLong, testutil.IntervalFast,

"never saw stats after conn closes",

)

assertConnectionReport(t, agentClient, proto.Connection_JETBRAINS, 0, "")

})

}

func TestAgent_SessionExec(t *testing.T) {

t.Parallel()

for _, port := range sshPorts {

t.Run(fmt.Sprintf("(:%d)", port), func(t *testing.T) {

t.Parallel()

session := setupSSHSessionOnPort(t, agentsdk.Manifest{}, codersdk.ServiceBannerConfig{}, nil, port)

command := "echo test"

if runtime.GOOS == "windows" {

command = "cmd.exe /c echo test"

}

output, err := session.Output(command)

require.NoError(t, err)

require.Equal(t, "test", strings.TrimSpace(string(output)))

})

}

}

func TestAgent_Session_EnvironmentVariables(t *testing.T) {

t.Parallel()

tmpdir := t.TempDir()

// Defined by the coder script runner, hardcoded here since we don't

// have a reference to it.

scriptBinDir := filepath.Join(tmpdir, "coder-script-data", "bin")

manifest := agentsdk.Manifest{

EnvironmentVariables: map[string]string{

"MY_MANIFEST": "true",

"MY_OVERRIDE": "false",

"MY_SESSION_MANIFEST": "false",

},

}

ctx := testutil.Context(t, testutil.WaitLong)

//nolint:dogsled

conn, _, _, _, _ := setupAgent(t, manifest, 0, func(_ *agenttest.Client, opts *agent.Options) {

opts.ScriptDataDir = tmpdir

opts.EnvironmentVariables["MY_OVERRIDE"] = "true"

})

sshClient, err := conn.SSHClient(ctx)

require.NoError(t, err)

t.Cleanup(func() { _ = sshClient.Close() })

for envName, want := range map[string]string{

"CODER": "true", // From the agent.

"MY_MANIFEST": "true", // From the manifest.

"MY_OVERRIDE": "true", // From the agent environment variables option, overrides manifest.

"MY_SESSION_MANIFEST": "false", // From the manifest, overrides session env.

"MY_SESSION": "true", // From the session.

"PATH": scriptBinDir,

} {

t.Run(envName, func(t *testing.T) {

t.Parallel()

got := sessionEnvValue(t, sshClient, envName, map[string]string{

"MY_SESSION": "true",

"MY_SESSION_MANIFEST": "true",

})

if envName == "PATH" {

require.Contains(t, filepath.SplitList(got), want)

return

}

require.Equal(t, want, got)

})

}

}

func TestAgent_Session_SecretInjection(t *testing.T) {

t.Parallel()

manifest := agentsdk.Manifest{

EnvironmentVariables: map[string]string{

"SHOULD_BE_OVERRIDDEN": "manifest-value",

},

}

secrets := []agentsdk.WorkspaceSecret{

{EnvName: "MY_SECRET_ENV", Value: []byte("env-secret-value")},

{FilePath: "/tmp/secret-file", Value: []byte("file-secret-content")},

{EnvName: "BOTH_ENV", FilePath: "/tmp/both-file", Value: []byte("both-value")},

{EnvName: "SHOULD_BE_OVERRIDDEN", Value: []byte("secret-wins")},

}

ctx := testutil.Context(t, testutil.WaitLong)

//nolint:dogsled

conn, _, _, fs, _ := setupAgentWithSecrets(t, manifest, secrets, 0)

// Verify file injection via the agent's filesystem.

content, err := afero.ReadFile(fs, "/tmp/secret-file")

require.NoError(t, err)

require.Equal(t, "file-secret-content", string(content))

content, err = afero.ReadFile(fs, "/tmp/both-file")

require.NoError(t, err)

require.Equal(t, "both-value", string(content))

sshClient, err := conn.SSHClient(ctx)

require.NoError(t, err)

t.Cleanup(func() { _ = sshClient.Close() })

for envName, want := range map[string]string{

"MY_SECRET_ENV": "env-secret-value",

"BOTH_ENV": "both-value",

"SHOULD_BE_OVERRIDDEN": "secret-wins",

} {

t.Run(envName, func(t *testing.T) {

t.Parallel()

require.Equal(t, want, sessionEnvValue(t, sshClient, envName, nil))

})

}

}

func sessionEnvValue(t *testing.T, sshClient *ssh.Client, envName string, sessionEnv map[string]string) string {

t.Helper()

ctx := testutil.Context(t, testutil.WaitLong)

session, err := sshClient.NewSession()

require.NoError(t, err)

defer session.Close()

go func() {

<-ctx.Done()

_ = session.Close()

}()

for name, value := range sessionEnv {

require.NoError(t, session.Setenv(name, value))

}

stderr := &bytes.Buffer{}

session.Stderr = stderr

command := "sh -c 'echo $" + envName + "'"

if runtime.GOOS == "windows" {

command = `cmd.exe /c echo %` + envName + `%`

}

out, err := session.Output(command)

if err != nil && ctx.Err() != nil {

t.Fatalf("SSH session deadline expired: %v; output: %q, stderr: %q", ctx.Err(), out, stderr.String())

}

require.NoError(t, err, "output: %q, stderr: %q", out, stderr.String())

return strings.TrimSpace(string(out))

}

func TestAgent_StartupScript_SecretInjection(t *testing.T) {

t.Parallel()

if runtime.GOOS == "windows" {

t.Skip("startup script test uses sh syntax")

}

tmpDir := t.TempDir()

secretFilePath := filepath.Join(tmpDir, "secret-file")

envProofPath := filepath.Join(tmpDir, "env-proof")

fileProofPath := filepath.Join(tmpDir, "file-proof")

// The startup script reads the secret env var and the secret file,

// writing both to proof files so we can verify they were available

// at script execution time.

script := fmt.Sprintf(

"echo \"$MY_STARTUP_SECRET\" > %s && cat %s > %s",

envProofPath, secretFilePath, fileProofPath,

)

manifest := agentsdk.Manifest{

Scripts: []codersdk.WorkspaceAgentScript{{

Script: script,

Timeout: 30 * time.Second,

RunOnStart: true,

}},

}

secrets := []agentsdk.WorkspaceSecret{

{EnvName: "MY_STARTUP_SECRET", Value: []byte("startup-env-value")},

{FilePath: secretFilePath, Value: []byte("startup-file-content")},

}

// Use the real OS filesystem so that both writeSecretFiles and

// the startup script operate on the same filesystem.

//nolint:dogsled

_, client, _, _, _ := setupAgentWithSecrets(t, manifest, secrets, 0, func(_ *agenttest.Client, opts *agent.Options) {

opts.Filesystem = afero.NewOsFs()

})

// Wait for the startup script to complete.

var got []codersdk.WorkspaceAgentLifecycle

assert.Eventually(t, func() bool {

got = client.GetLifecycleStates()

return len(got) > 0 && got[len(got)-1] == codersdk.WorkspaceAgentLifecycleReady

}, testutil.WaitLong, testutil.IntervalMedium)

require.Contains(t, got, codersdk.WorkspaceAgentLifecycleReady, "agent never reached ready")

// Verify the startup script could read the secret env var.

envProof, err := os.ReadFile(envProofPath)

require.NoError(t, err)

require.Equal(t, "startup-env-value", strings.TrimSpace(string(envProof)))

// Verify the startup script could read the secret file.

fileProof, err := os.ReadFile(fileProofPath)

require.NoError(t, err)

require.Equal(t, "startup-file-content", string(fileProof))

}

func TestAgent_GitSSH(t *testing.T) {

t.Parallel()

session := setupSSHSession(t, agentsdk.Manifest{}, codersdk.ServiceBannerConfig{}, nil)

command := "sh -c 'echo $GIT_SSH_COMMAND'"

if runtime.GOOS == "windows" {

command = "cmd.exe /c echo %GIT_SSH_COMMAND%"

}

output, err := session.Output(command)

require.NoError(t, err)

require.True(t, strings.HasSuffix(strings.TrimSpace(string(output)), "gitssh --"))

}

func TestAgent_SessionTTYShell(t *testing.T) {

t.Parallel()

if runtime.GOOS == "windows" {

// This might be our implementation, or ConPTY itself.

// It's difficult to find extensive tests for it, so

// it seems like it could be either.

t.Skip("ConPTY appears to be inconsistent on Windows.")

}

for _, port := range sshPorts {

t.Run(fmt.Sprintf("(%d)", port), func(t *testing.T) {

t.Parallel()

ctx := testutil.Context(t, testutil.WaitMedium)

logger := testutil.Logger(t)

session := setupSSHSessionOnPort(t, agentsdk.Manifest{}, codersdk.ServiceBannerConfig{}, nil, port)

command := "sh"

if runtime.GOOS == "windows" {

command = "cmd.exe"

}

err := session.RequestPty("xterm", 128, 128, ssh.TerminalModes{})

require.NoError(t, err)

stdout := expecter.NewAttachedToSSHSession(t, session)

stdin := testutil.NewWriterAttachedToSSHSession(t, logger.Named("sshin"), session)

err = session.Start(command)

require.NoError(t, err)

_ = stdout.Peek(ctx, 1) // wait for the prompt

stdin.WriteLine("echo test")

stdout.ExpectMatch(ctx, "test")

stdin.WriteLine("exit")

err = session.Wait()

require.NoError(t, err)

})

}

}

func TestAgent_SessionTTYExitCode(t *testing.T) {

t.Parallel()

session := setupSSHSession(t, agentsdk.Manifest{}, codersdk.ServiceBannerConfig{}, nil)

command := "areallynotrealcommand"

err := session.RequestPty("xterm", 128, 128, ssh.TerminalModes{})

require.NoError(t, err)

err = session.Start(command)

require.NoError(t, err)

err = session.Wait()

exitErr := &ssh.ExitError{}

require.True(t, xerrors.As(err, &exitErr))

if runtime.GOOS == "windows" {

assert.Equal(t, 1, exitErr.ExitStatus())

} else {

assert.Equal(t, 127, exitErr.ExitStatus())

}

}

func TestAgent_Session_TTY_MOTD(t *testing.T) {

t.Parallel()

if runtime.GOOS == "windows" {

// This might be our implementation, or ConPTY itself.

// It's difficult to find extensive tests for it, so

// it seems like it could be either.

t.Skip("ConPTY appears to be inconsistent on Windows.")

}

u, err := user.Current()

require.NoError(t, err, "get current user")

name := filepath.Join(u.HomeDir, "motd")

wantMOTD := "Welcome to your Coder workspace!"

wantServiceBanner := "Service banner text goes here"

tests := []struct {

name string

manifest agentsdk.Manifest

banner codersdk.ServiceBannerConfig

expected []string

unexpected []string

expectedRe *regexp.Regexp

}{

{

name: "WithoutServiceBanner",

manifest: agentsdk.Manifest{MOTDFile: name},

banner: codersdk.ServiceBannerConfig{},

expected: []string{wantMOTD},

unexpected: []string{wantServiceBanner},

},

{

name: "WithServiceBanner",

manifest: agentsdk.Manifest{MOTDFile: name},

banner: codersdk.ServiceBannerConfig{

Enabled: true,

Message: wantServiceBanner,

},

expected: []string{wantMOTD, wantServiceBanner},

},

{

name: "ServiceBannerDisabled",

manifest: agentsdk.Manifest{MOTDFile: name},

banner: codersdk.ServiceBannerConfig{

Enabled: false,

Message: wantServiceBanner,

},

expected: []string{wantMOTD},

unexpected: []string{wantServiceBanner},

},

{

name: "ServiceBannerOnly",

manifest: agentsdk.Manifest{},

banner: codersdk.ServiceBannerConfig{

Enabled: true,

Message: wantServiceBanner,

},

expected: []string{wantServiceBanner},

unexpected: []string{wantMOTD},

},

{

name: "None",

manifest: agentsdk.Manifest{},

banner: codersdk.ServiceBannerConfig{},

unexpected: []string{wantServiceBanner, wantMOTD},

},

{

name: "CarriageReturns",

manifest: agentsdk.Manifest{},

banner: codersdk.ServiceBannerConfig{

Enabled: true,

Message: "service\n\nbanner\nhere",

},

expected: []string{"service\r\n\r\nbanner\r\nhere\r\n\r\n"},

unexpected: []string{},

},

{

name: "Trim",

// Enable motd since it will be printed after the banner,

// this ensures that we can test for an exact mount of

// newlines.

manifest: agentsdk.Manifest{

MOTDFile: name,

},

banner: codersdk.ServiceBannerConfig{

Enabled: true,

Message: "\n\n\n\n\n\nbanner\n\n\n\n\n\n",

},

expectedRe: regexp.MustCompile(`([^\n\r]|^)banner\r\n\r\n[^\r\n]`),

},

}

for _, test := range tests {

t.Run(test.name, func(t *testing.T) {

t.Parallel()

session := setupSSHSession(t, test.manifest, test.banner, func(fs afero.Fs) {

err := fs.MkdirAll(filepath.Dir(name), 0o700)

require.NoError(t, err)

err = afero.WriteFile(fs, name, []byte(wantMOTD), 0o600)

require.NoError(t, err)

})

testSessionOutput(t, session, test.expected, test.unexpected, test.expectedRe)

})

}

}

//nolint:tparallel // Sub tests need to run sequentially.

func TestAgent_Session_TTY_MOTD_Update(t *testing.T) {

t.Parallel()

if runtime.GOOS == "windows" {

// This might be our implementation, or ConPTY itself.

// It's difficult to find extensive tests for it, so

// it seems like it could be either.

t.Skip("ConPTY appears to be inconsistent on Windows.")

}

// Only the banner updates dynamically; the MOTD file does not.

wantServiceBanner := "Service banner text goes here"

tests := []struct {

banner codersdk.ServiceBannerConfig

expected []string

unexpected []string

}{

{

banner: codersdk.ServiceBannerConfig{},

expected: []string{},

unexpected: []string{wantServiceBanner},

},

{

banner: codersdk.ServiceBannerConfig{

Enabled: true,

Message: wantServiceBanner,

},

expected: []string{wantServiceBanner},

},

{

banner: codersdk.ServiceBannerConfig{

Enabled: false,

Message: wantServiceBanner,

},

expected: []string{},

unexpected: []string{wantServiceBanner},

},

{

banner: codersdk.ServiceBannerConfig{

Enabled: true,

Message: wantServiceBanner,

},

expected: []string{wantServiceBanner},

unexpected: []string{},

},

{

banner: codersdk.ServiceBannerConfig{},

unexpected: []string{wantServiceBanner},

},

}

setSBInterval := func(_ *agenttest.Client, opts *agent.Options) {

opts.ServiceBannerRefreshInterval = testutil.IntervalFast

}

//nolint:dogsled // Allow the blank identifiers.

conn, client, _, _, _ := setupAgent(t, agentsdk.Manifest{}, 0, setSBInterval)

ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)

defer cancel()

//nolint:paralleltest // These tests need to swap the banner func.

for _, port := range sshPorts {

sshClient, err := conn.SSHClientOnPort(ctx, port)

require.NoError(t, err)

t.Cleanup(func() {

_ = sshClient.Close()

})

for i, test := range tests {

t.Run(fmt.Sprintf("(:%d)/%d", port, i), func(t *testing.T) {

// Set new banner func and wait for the agent to call it to update the

// banner. We wait for two calls to ensure the value has been stored:

// the second call can only begin after the first iteration of

// fetchServiceBannerLoop completes (call + store), so after

// receiving two signals at least one store has happened.

ready := make(chan struct{}, 2)

client.SetAnnouncementBannersFunc(func() ([]codersdk.BannerConfig, error) {

select {

case ready <- struct{}{}:

default:

}

return []codersdk.BannerConfig{test.banner}, nil

})

testutil.TryReceive(ctx, t, ready)

testutil.TryReceive(ctx, t, ready)

session, err := sshClient.NewSession()

require.NoError(t, err)

t.Cleanup(func() {

_ = session.Close()

})

testSessionOutput(t, session, test.expected, test.unexpected, nil)

})

}

}

}

//nolint:paralleltest // This test sets an environment variable.

func TestAgent_Session_TTY_QuietLogin(t *testing.T) {

if runtime.GOOS == "windows" {

// This might be our implementation, or ConPTY itself.

// It's difficult to find extensive tests for it, so

// it seems like it could be either.

t.Skip("ConPTY appears to be inconsistent on Windows.")

}

wantNotMOTD := "Welcome to your Coder workspace!"

wantServiceBanner := "Service banner text goes here"

u, err := user.Current()

require.NoError(t, err, "get current user")

motdPath := filepath.Join(u.HomeDir, "motd")

hushloginPath := filepath.Join(u.HomeDir, ".hushlogin")

// Neither banner nor MOTD should show if not a login shell.

t.Run("NotLogin", func(t *testing.T) {

session := setupSSHSession(t, agentsdk.Manifest{

MOTDFile: motdPath,

}, codersdk.ServiceBannerConfig{

Enabled: true,

Message: wantServiceBanner,

}, func(fs afero.Fs) {

err := afero.WriteFile(fs, motdPath, []byte(wantNotMOTD), 0o600)

require.NoError(t, err, "write motd file")

})

err = session.RequestPty("xterm", 128, 128, ssh.TerminalModes{})

require.NoError(t, err)

wantEcho := "foobar"

command := "echo " + wantEcho

output, err := session.Output(command)

require.NoError(t, err)

require.Contains(t, string(output), wantEcho, "should show echo")

require.NotContains(t, string(output), wantNotMOTD, "should not show motd")

require.NotContains(t, string(output), wantServiceBanner, "should not show service banner")

})

// Only the MOTD should be silenced when hushlogin is present.

t.Run("Hushlogin", func(t *testing.T) {

// The agent writes banners and the MOTD before starting the

// login shell. The working theory is that, under CI contention,

// the host login shell can stall during startup and fail to

// process the PTY exit command, leaving session.Wait blocked.

// Use a self-terminating test shell so unrelated host-shell

// behavior cannot block the quiet-login assertions.

shellPath := filepath.Join(t.TempDir(), "shell")

//nolint:gosec // Executable test shell with test-controlled content.

err := os.WriteFile(shellPath, []byte("#!/bin/sh\nexit 0\n"), 0o700)

require.NoError(t, err, "write test shell")

session := setupSSHSession(t, agentsdk.Manifest{

MOTDFile: motdPath,

}, codersdk.BannerConfig{

Enabled: true,

Message: wantServiceBanner,

}, func(fs afero.Fs) {

err := afero.WriteFile(fs, motdPath, []byte(wantNotMOTD), 0o600)

require.NoError(t, err, "write motd file")

// Place an empty .hushlogin in the user's home so the agent's

// isQuietLogin lookup succeeds and showMOTD is skipped.

err = afero.WriteFile(fs, hushloginPath, []byte{}, 0o600)

require.NoError(t, err, "write hushlogin file")

}, func(_ *agenttest.Client, opts *agent.Options) {

opts.EnvInfo = shellOverrideEnvInfo{shell: shellPath}

})

err = session.RequestPty("xterm", 128, 128, ssh.TerminalModes{})

require.NoError(t, err)

stdout := testutil.NewWaitBuffer()

session.Stdout = stdout

require.NoError(t, session.Shell())

ctx := testutil.Context(t, testutil.WaitShort)

stopClose := context.AfterFunc(ctx, func() { _ = session.Close() })