mirror of
https://github.com/coder/coder.git
synced 2026-10-16 14:48:56 +08:00
test: add memory guard and heap profiles for CI timeouts Add a memory guard to `cli`, `agent`, and `enterprise/cli` to diagnose recurring macOS `test-go-pg` timeouts. Recent logs show runaway memory usage and heavy swapping, but the allocation source remains unknown. - Check Go runtime memory once per second, with per-process CI limits of 5 GiB on macOS and 6 GiB on Linux and Windows. Disabled when `CODER_TEST_MEMORY_LIMIT_MIB` is unset. - On exceeding the limit, save a heap profile, attempt a second profile after forced GC with a five-second wait limit, then panic with memory usage, profile paths, and all goroutine stacks. - Upload heap profiles from unsuccessful CI jobs on every OS, retaining them for seven days. - Test guard lifecycle, environment parsing, panic diagnostics, profile validity, and post-GC allocation attribution. Validate with race tests, vet, lint, and a manual macOS test. Refs https://linear.app/codercom/issue/PLAT-251/flake-test-go-pg-macos-job-timeoutcancelled Refs https://github.com/coder/internal/issues/1365
4802 lines
150 KiB
Go
4802 lines
150 KiB
Go
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"
|
||
"sync"
|
||
"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"
|
||
tailnetproto "github.com/coder/coder/v2/tailnet/proto"
|
||
"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)
|
||
}
|
||
// Report multi-second pauses of the whole process, see coder/internal#1365.
|
||
testutil.StartStallDetector(5 * time.Second)
|
||
// Fail with a heap profile before a runaway allocation pushes the runner into swap.
|
||
testutil.StartMemoryGuardFromEnv()
|
||
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)
|
||
}
|
||
|
||
// eofClient forces the agent's RPC connection to fail with io.EOF so runLoop
|
||
// takes the disconnect path.
|
||
type eofClient struct {
|
||
*agenttest.Client
|
||
}
|
||
|
||
func (*eofClient) ConnectRPC214WithRole(context.Context, string) (
|
||
proto.DRPCAgentClient214, tailnetproto.DRPCTailnetClient28, error,
|
||
) {
|
||
return nil, nil, io.EOF
|
||
}
|
||
|
||
// flushTestSink is a concurrency-safe slog.Sink that records entry messages.
|
||
type flushTestSink struct {
|
||
mu sync.Mutex
|
||
messages []string
|
||
}
|
||
|
||
func (s *flushTestSink) LogEntry(_ context.Context, e slog.SinkEntry) {
|
||
s.mu.Lock()
|
||
defer s.mu.Unlock()
|
||
s.messages = append(s.messages, e.Message)
|
||
}
|
||
|
||
func (*flushTestSink) Sync() {}
|
||
|
||
func (s *flushTestSink) contains(msg string) bool {
|
||
s.mu.Lock()
|
||
defer s.mu.Unlock()
|
||
return slices.Contains(s.messages, msg)
|
||
}
|
||
|
||
// TestAgent_FlushFlightRecorderOnDisconnect verifies the agent emits its recorded
|
||
// debug history when it loses the connection to coderd, so the detail leading
|
||
// up to the disconnect is available even though the agent runs at Info.
|
||
func TestAgent_FlushFlightRecorderOnDisconnect(t *testing.T) {
|
||
t.Parallel()
|
||
|
||
ctx := testutil.Context(t, testutil.WaitShort)
|
||
|
||
sink := &flushTestSink{}
|
||
logger := slog.Make(sink).Leveled(slog.LevelInfo).FlightRecorder(1000)
|
||
|
||
// The debug entry is held by the flight recorder, not emitted, during normal
|
||
// operation at Info.
|
||
const sentinel = "recorded debug sentinel"
|
||
logger.Debug(ctx, sentinel)
|
||
require.False(t, sink.contains(sentinel), "debug entry should be recorded, not emitted")
|
||
|
||
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)
|
||
baseClient := agenttest.NewClient(t, logger.Named("agenttest"), manifest.AgentID, manifest, statsCh, coordinator)
|
||
t.Cleanup(baseClient.Close)
|
||
|
||
agentUnderTest := agent.New(agent.Options{
|
||
Client: &eofClient{Client: baseClient},
|
||
Filesystem: afero.NewMemMapFs(),
|
||
Logger: logger.Named("agent"),
|
||
})
|
||
t.Cleanup(func() {
|
||
_ = agentUnderTest.Close()
|
||
})
|
||
|
||
// Losing the connection to coderd flushes the recorded debug history to the
|
||
// sink.
|
||
require.Eventually(t, func() bool {
|
||
return sink.contains(sentinel)
|
||
}, testutil.WaitShort, testutil.IntervalFast)
|
||
}
|
||
|
||
// 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, agentClient, 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")
|
||
|
||
assertConnectionReport(t, agentClient,
|
||
proto.ConnectEvent{Method: proto.Connection_METHOD_SSH},
|
||
proto.DisconnectEvent{Reason: string(codersdk.DisconnectReasonGraceful)},
|
||
)
|
||
})
|
||
}
|
||
|
||
// 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.ConnectEvent{Method: proto.Connection_METHOD_SSH, AppName: string(codersdk.AppFamilyVSCode)},
|
||
proto.DisconnectEvent{Reason: string(codersdk.DisconnectReasonGraceful)},
|
||
)
|
||
})
|
||
|
||
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.ConnectEvent{Method: proto.Connection_METHOD_SSH, AppName: string(codersdk.AppFamilyJetBrains)},
|
||
proto.DisconnectEvent{Reason: "normal close"},
|
||
)
|
||
})
|
||
}
|
||
|
||
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"
|
||
opts.EnvInfo = sessionEnvInfo{}
|
||
})
|
||
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, func(_ *agenttest.Client, opts *agent.Options) {
|
||
opts.EnvInfo = sessionEnvInfo{}
|
||
})
|
||
|
||
// 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))
|
||
})
|
||
}
|
||
}
|
||
|
||
// Environment assertions need a shell that does not run host startup files.
|
||
// The default shell can run those files before it executes the SSH command,
|
||
// even when that command invokes another shell.
|
||
type sessionEnvInfo struct {
|
||
usershell.SystemEnvInfo
|
||
}
|
||
|
||
func (sessionEnvInfo) Shell(string) (string, error) {
|
||
if runtime.GOOS == "windows" {
|
||
return exec.LookPath("cmd.exe")
|
||
}
|
||
return "/bin/sh", nil
|
||
}
|
||
|
||
func (sessionEnvInfo) ModifyCommand(name string, args ...string) (string, []string) {
|
||
if runtime.GOOS == "windows" {
|
||
// Disable cmd.exe AutoRun commands from the host registry.
|
||
args = append([]string{"/d"}, args...)
|
||
}
|
||
return name, args
|
||
}
|
||
|
||
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 := "printf '%s\\n' \"$" + envName + "\""
|
||
if runtime.GOOS == "windows" {
|
||
command = `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() })
|
||
|
||
err = session.Wait()
|
||
stopClose()
|
||
require.NoError(t, err)
|
||
|
||
require.Contains(t, stdout.String(), wantServiceBanner, "should show service banner")
|
||
require.NotContains(t, stdout.String(), wantNotMOTD, "should not show motd")
|
||
})
|
||
}
|
||
|
||
func TestAgent_Session_TTY_FastCommandHasOutput(t *testing.T) {
|
||
t.Parallel()
|
||
// This test is here to prevent regressions where quickly executing
|
||
// commands (with TTY) don't sync their output to the SSH session.
|
||
//
|
||
// See: https://github.com/coder/coder/issues/6656
|
||
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()
|
||
|
||
var stdout bytes.Buffer
|
||
// NOTE(mafredri): Increase iterations to increase chance of failure,
|
||
// assuming bug is present. Limiting GOMAXPROCS further
|
||
// increases the chance of failure.
|
||
// Using 1000 iterations is basically a guaranteed failure (but let's
|
||
// not increase test times needlessly).
|
||
// Limit GOMAXPROCS (e.g. `export GOMAXPROCS=1`) to further increase
|
||
// chance of failure. Also -race helps.
|
||
for i := 0; i < 5; i++ {
|
||
func() {
|
||
stdout.Reset()
|
||
|
||
session, err := sshClient.NewSession()
|
||
require.NoError(t, err)
|
||
defer session.Close()
|
||
err = session.RequestPty("xterm", 128, 128, ssh.TerminalModes{})
|
||
require.NoError(t, err)
|
||
|
||
session.Stdout = &stdout
|
||
err = session.Start("echo wazzup")
|
||
require.NoError(t, err)
|
||
|
||
err = session.Wait()
|
||
require.NoError(t, err)
|
||
require.Contains(t, stdout.String(), "wazzup", "should output greeting")
|
||
}()
|
||
}
|
||
}
|
||
|
||
func TestAgent_Session_TTY_HugeOutputIsNotLost(t *testing.T) {
|
||
t.Parallel()
|
||
|
||
// This test is here to prevent regressions where a command (with or
|
||
// without) a large amount of output would not be fully copied to the
|
||
// SSH session. On unix systems, this was fixed by duplicating the file
|
||
// descriptor of the PTY master and using it for copying the output.
|
||
//
|
||
// See: https://github.com/coder/coder/issues/6656
|
||
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()
|
||
|
||
var stdout bytes.Buffer
|
||
// NOTE(mafredri): Increase iterations to increase chance of failure,
|
||
// assuming bug is present.
|
||
// Using 10 iterations is basically a guaranteed failure (but let's
|
||
// not increase test times needlessly). Run with -race and do not
|
||
// limit parallelism (`export GOMAXPROCS=10`) to increase the chance
|
||
// of failure.
|
||
for i := 0; i < 1; i++ {
|
||
func() {
|
||
stdout.Reset()
|
||
|
||
session, err := sshClient.NewSession()
|
||
require.NoError(t, err)
|
||
defer session.Close()
|
||
err = session.RequestPty("xterm", 128, 128, ssh.TerminalModes{})
|
||
require.NoError(t, err)
|
||
|
||
session.Stdout = &stdout
|
||
want := strings.Repeat("wazzup", 1024+1) // ~6KB, +1 because 1024 is a common buffer size.
|
||
err = session.Start("echo " + want)
|
||
require.NoError(t, err)
|
||
|
||
err = session.Wait()
|
||
require.NoError(t, err)
|
||
require.Contains(t, stdout.String(), want, "should output entire greeting")
|
||
}()
|
||
}
|
||
}
|
||
|
||
func TestAgent_TCPLocalForwarding(t *testing.T) {
|
||
t.Parallel()
|
||
ctx := testutil.Context(t, testutil.WaitLong)
|
||
|
||
rl, err := net.Listen("tcp", "127.0.0.1:0")
|
||
require.NoError(t, err)
|
||
defer rl.Close()
|
||
tcpAddr, valid := rl.Addr().(*net.TCPAddr)
|
||
require.True(t, valid)
|
||
remotePort := tcpAddr.Port
|
||
go echoOnce(t, rl)
|
||
|
||
sshClient := setupAgentSSHClient(ctx, t)
|
||
|
||
conn, err := sshClient.Dial("tcp", fmt.Sprintf("127.0.0.1:%d", remotePort))
|
||
require.NoError(t, err)
|
||
defer conn.Close()
|
||
requireEcho(t, conn)
|
||
}
|
||
|
||
func TestAgent_TCPRemoteForwarding(t *testing.T) {
|
||
t.Parallel()
|
||
ctx := testutil.Context(t, testutil.WaitLong)
|
||
sshClient := setupAgentSSHClient(ctx, t)
|
||
|
||
localhost := netip.MustParseAddr("127.0.0.1")
|
||
var randomPort uint16
|
||
var ll net.Listener
|
||
var err error
|
||
for {
|
||
randomPort = testutil.RandomPortNoListen(t)
|
||
addr := net.TCPAddrFromAddrPort(netip.AddrPortFrom(localhost, randomPort))
|
||
ll, err = sshClient.ListenTCP(addr)
|
||
if err != nil {
|
||
t.Logf("error remote forwarding: %s", err.Error())
|
||
select {
|
||
case <-ctx.Done():
|
||
t.Fatal("timed out getting random listener")
|
||
default:
|
||
continue
|
||
}
|
||
}
|
||
break
|
||
}
|
||
defer ll.Close()
|
||
go echoOnce(t, ll)
|
||
|
||
conn, err := net.Dial("tcp", fmt.Sprintf("127.0.0.1:%d", randomPort))
|
||
require.NoError(t, err)
|
||
defer conn.Close()
|
||
requireEcho(t, conn)
|
||
}
|
||
|
||
func TestAgent_TCPLocalForwardingBlocked(t *testing.T) {
|
||
t.Parallel()
|
||
ctx := testutil.Context(t, testutil.WaitLong)
|
||
|
||
rl, err := net.Listen("tcp", "127.0.0.1:0")
|
||
require.NoError(t, err)
|
||
defer rl.Close()
|
||
tcpAddr, valid := rl.Addr().(*net.TCPAddr)
|
||
require.True(t, valid)
|
||
remotePort := tcpAddr.Port
|
||
|
||
//nolint:dogsled
|
||
agentConn, _, _, _, _ := setupAgent(t, agentsdk.Manifest{}, 0, func(_ *agenttest.Client, o *agent.Options) {
|
||
o.BlockLocalPortForwarding = true
|
||
})
|
||
sshClient, err := agentConn.SSHClient(ctx)
|
||
require.NoError(t, err)
|
||
defer sshClient.Close()
|
||
|
||
_, err = sshClient.Dial("tcp", fmt.Sprintf("127.0.0.1:%d", remotePort))
|
||
require.ErrorContains(t, err, "administratively prohibited")
|
||
}
|
||
|
||
func TestAgent_TCPRemoteForwardingBlocked(t *testing.T) {
|
||
t.Parallel()
|
||
ctx := testutil.Context(t, testutil.WaitLong)
|
||
|
||
//nolint:dogsled
|
||
agentConn, _, _, _, _ := setupAgent(t, agentsdk.Manifest{}, 0, func(_ *agenttest.Client, o *agent.Options) {
|
||
o.BlockReversePortForwarding = true
|
||
})
|
||
sshClient, err := agentConn.SSHClient(ctx)
|
||
require.NoError(t, err)
|
||
defer sshClient.Close()
|
||
|
||
localhost := netip.MustParseAddr("127.0.0.1")
|
||
randomPort := testutil.RandomPortNoListen(t)
|
||
addr := net.TCPAddrFromAddrPort(netip.AddrPortFrom(localhost, randomPort))
|
||
_, err = sshClient.ListenTCP(addr)
|
||
require.ErrorContains(t, err, "tcpip-forward request denied by peer")
|
||
}
|
||
|
||
func TestAgent_UnixLocalForwardingBlocked(t *testing.T) {
|
||
t.Parallel()
|
||
if runtime.GOOS == "windows" {
|
||
t.Skip("unix domain sockets are not fully supported on Windows")
|
||
}
|
||
ctx := testutil.Context(t, testutil.WaitLong)
|
||
tmpdir := testutil.TempDirUnixSocket(t)
|
||
remoteSocketPath := filepath.Join(tmpdir, "remote-socket")
|
||
|
||
l, err := net.Listen("unix", remoteSocketPath)
|
||
require.NoError(t, err)
|
||
defer l.Close()
|
||
|
||
//nolint:dogsled
|
||
agentConn, _, _, _, _ := setupAgent(t, agentsdk.Manifest{}, 0, func(_ *agenttest.Client, o *agent.Options) {
|
||
o.BlockLocalPortForwarding = true
|
||
})
|
||
sshClient, err := agentConn.SSHClient(ctx)
|
||
require.NoError(t, err)
|
||
defer sshClient.Close()
|
||
|
||
_, err = sshClient.Dial("unix", remoteSocketPath)
|
||
require.ErrorContains(t, err, "administratively prohibited")
|
||
}
|
||
|
||
func TestAgent_UnixRemoteForwardingBlocked(t *testing.T) {
|
||
t.Parallel()
|
||
if runtime.GOOS == "windows" {
|
||
t.Skip("unix domain sockets are not fully supported on Windows")
|
||
}
|
||
ctx := testutil.Context(t, testutil.WaitLong)
|
||
tmpdir := testutil.TempDirUnixSocket(t)
|
||
remoteSocketPath := filepath.Join(tmpdir, "remote-socket")
|
||
|
||
//nolint:dogsled
|
||
agentConn, _, _, _, _ := setupAgent(t, agentsdk.Manifest{}, 0, func(_ *agenttest.Client, o *agent.Options) {
|
||
o.BlockReversePortForwarding = true
|
||
})
|
||
sshClient, err := agentConn.SSHClient(ctx)
|
||
require.NoError(t, err)
|
||
defer sshClient.Close()
|
||
|
||
_, err = sshClient.ListenUnix(remoteSocketPath)
|
||
require.ErrorContains(t, err, "streamlocal-forward@openssh.com request denied by peer")
|
||
}
|
||
|
||
// TestAgent_LocalBlockedDoesNotAffectReverse verifies that blocking
|
||
// local port forwarding does not prevent reverse port forwarding from
|
||
// working. A field-name transposition at any plumbing hop would cause
|
||
// both directions to be blocked when only one flag is set.
|
||
func TestAgent_LocalBlockedDoesNotAffectReverse(t *testing.T) {
|
||
t.Parallel()
|
||
ctx := testutil.Context(t, testutil.WaitLong)
|
||
|
||
//nolint:dogsled
|
||
agentConn, _, _, _, _ := setupAgent(t, agentsdk.Manifest{}, 0, func(_ *agenttest.Client, o *agent.Options) {
|
||
o.BlockLocalPortForwarding = true
|
||
})
|
||
sshClient, err := agentConn.SSHClient(ctx)
|
||
require.NoError(t, err)
|
||
defer sshClient.Close()
|
||
|
||
// Reverse forwarding must still work.
|
||
localhost := netip.MustParseAddr("127.0.0.1")
|
||
var ll net.Listener
|
||
for {
|
||
randomPort := testutil.RandomPortNoListen(t)
|
||
addr := net.TCPAddrFromAddrPort(netip.AddrPortFrom(localhost, randomPort))
|
||
ll, err = sshClient.ListenTCP(addr)
|
||
if err != nil {
|
||
t.Logf("error remote forwarding: %s", err.Error())
|
||
select {
|
||
case <-ctx.Done():
|
||
t.Fatal("timed out getting random listener")
|
||
default:
|
||
continue
|
||
}
|
||
}
|
||
break
|
||
}
|
||
_ = ll.Close()
|
||
}
|
||
|
||
// TestAgent_ReverseBlockedDoesNotAffectLocal verifies that blocking
|
||
// reverse port forwarding does not prevent local port forwarding from
|
||
// working.
|
||
func TestAgent_ReverseBlockedDoesNotAffectLocal(t *testing.T) {
|
||
t.Parallel()
|
||
ctx := testutil.Context(t, testutil.WaitLong)
|
||
|
||
rl, err := net.Listen("tcp", "127.0.0.1:0")
|
||
require.NoError(t, err)
|
||
defer rl.Close()
|
||
tcpAddr, valid := rl.Addr().(*net.TCPAddr)
|
||
require.True(t, valid)
|
||
remotePort := tcpAddr.Port
|
||
go echoOnce(t, rl)
|
||
|
||
//nolint:dogsled
|
||
agentConn, _, _, _, _ := setupAgent(t, agentsdk.Manifest{}, 0, func(_ *agenttest.Client, o *agent.Options) {
|
||
o.BlockReversePortForwarding = true
|
||
})
|
||
sshClient, err := agentConn.SSHClient(ctx)
|
||
require.NoError(t, err)
|
||
defer sshClient.Close()
|
||
|
||
// Local forwarding must still work.
|
||
conn, err := sshClient.Dial("tcp", fmt.Sprintf("127.0.0.1:%d", remotePort))
|
||
require.NoError(t, err)
|
||
defer conn.Close()
|
||
requireEcho(t, conn)
|
||
}
|
||
|
||
func TestAgent_UnixLocalForwarding(t *testing.T) {
|
||
t.Parallel()
|
||
if runtime.GOOS == "windows" {
|
||
t.Skip("unix domain sockets are not fully supported on Windows")
|
||
}
|
||
ctx := testutil.Context(t, testutil.WaitLong)
|
||
tmpdir := testutil.TempDirUnixSocket(t)
|
||
remoteSocketPath := filepath.Join(tmpdir, "remote-socket")
|
||
|
||
l, err := net.Listen("unix", remoteSocketPath)
|
||
require.NoError(t, err)
|
||
defer l.Close()
|
||
go echoOnce(t, l)
|
||
|
||
sshClient := setupAgentSSHClient(ctx, t)
|
||
|
||
conn, err := sshClient.Dial("unix", remoteSocketPath)
|
||
require.NoError(t, err)
|
||
defer conn.Close()
|
||
_, err = conn.Write([]byte("test"))
|
||
require.NoError(t, err)
|
||
b := make([]byte, 4)
|
||
_, err = conn.Read(b)
|
||
require.NoError(t, err)
|
||
require.Equal(t, "test", string(b))
|
||
_ = conn.Close()
|
||
}
|
||
|
||
func TestAgent_UnixRemoteForwarding(t *testing.T) {
|
||
t.Parallel()
|
||
if runtime.GOOS == "windows" {
|
||
t.Skip("unix domain sockets are not fully supported on Windows")
|
||
}
|
||
|
||
tmpdir := testutil.TempDirUnixSocket(t)
|
||
remoteSocketPath := filepath.Join(tmpdir, "remote-socket")
|
||
|
||
ctx := testutil.Context(t, testutil.WaitLong)
|
||
sshClient := setupAgentSSHClient(ctx, t)
|
||
|
||
l, err := sshClient.ListenUnix(remoteSocketPath)
|
||
require.NoError(t, err)
|
||
defer l.Close()
|
||
go echoOnce(t, l)
|
||
|
||
conn, err := net.Dial("unix", remoteSocketPath)
|
||
require.NoError(t, err)
|
||
defer conn.Close()
|
||
requireEcho(t, conn)
|
||
}
|
||
|
||
func TestAgent_SFTP(t *testing.T) {
|
||
t.Parallel()
|
||
|
||
t.Run("DefaultWorkingDirectory", func(t *testing.T) {
|
||
t.Parallel()
|
||
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
|
||
defer cancel()
|
||
u, err := user.Current()
|
||
require.NoError(t, err, "get current user")
|
||
home := u.HomeDir
|
||
if runtime.GOOS == "windows" {
|
||
home = "/" + strings.ReplaceAll(home, "\\", "/")
|
||
}
|
||
//nolint:dogsled
|
||
conn, agentClient, _, _, _ := setupAgent(t, agentsdk.Manifest{}, 0)
|
||
sshClient, err := conn.SSHClient(ctx)
|
||
require.NoError(t, err)
|
||
defer sshClient.Close()
|
||
client, err := sftp.NewClient(sshClient)
|
||
require.NoError(t, err)
|
||
defer client.Close()
|
||
wd, err := client.Getwd()
|
||
require.NoError(t, err, "get working directory")
|
||
require.Equal(t, home, wd, "working directory should be user home")
|
||
tempFile := filepath.Join(t.TempDir(), "sftp")
|
||
// SFTP only accepts unix-y paths.
|
||
remoteFile := filepath.ToSlash(tempFile)
|
||
if !path.IsAbs(remoteFile) {
|
||
// On Windows, e.g. "/C:/Users/...".
|
||
remoteFile = path.Join("/", remoteFile)
|
||
}
|
||
file, err := client.Create(remoteFile)
|
||
require.NoError(t, err)
|
||
err = file.Close()
|
||
require.NoError(t, err)
|
||
_, err = os.Stat(tempFile)
|
||
require.NoError(t, err)
|
||
|
||
// Close the client to trigger disconnect event.
|
||
_ = client.Close()
|
||
assertConnectionReport(t, agentClient,
|
||
proto.ConnectEvent{Method: proto.Connection_METHOD_SSH},
|
||
proto.DisconnectEvent{},
|
||
)
|
||
})
|
||
|
||
t.Run("CustomWorkingDirectory", func(t *testing.T) {
|
||
t.Parallel()
|
||
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
|
||
defer cancel()
|
||
|
||
// Create a custom directory for the agent to use.
|
||
customDir := t.TempDir()
|
||
expectedDir := customDir
|
||
if runtime.GOOS == "windows" {
|
||
expectedDir = "/" + strings.ReplaceAll(customDir, "\\", "/")
|
||
}
|
||
|
||
conn, agentClient, _, fs, _ := setupAgent(t, agentsdk.Manifest{
|
||
Directory: customDir,
|
||
}, 0)
|
||
// The agent stats the working directory against its filesystem, so
|
||
// the directory must exist there for it to be honored.
|
||
require.NoError(t, fs.MkdirAll(customDir, 0o700))
|
||
sshClient, err := conn.SSHClient(ctx)
|
||
require.NoError(t, err)
|
||
defer sshClient.Close()
|
||
client, err := sftp.NewClient(sshClient)
|
||
require.NoError(t, err)
|
||
defer client.Close()
|
||
wd, err := client.Getwd()
|
||
require.NoError(t, err, "get working directory")
|
||
require.Equal(t, expectedDir, wd, "working directory should be custom directory")
|
||
|
||
// Close the client to trigger disconnect event.
|
||
_ = client.Close()
|
||
assertConnectionReport(t, agentClient,
|
||
proto.ConnectEvent{Method: proto.Connection_METHOD_SSH},
|
||
proto.DisconnectEvent{},
|
||
)
|
||
})
|
||
|
||
t.Run("MissingWorkingDirectory", func(t *testing.T) {
|
||
t.Parallel()
|
||
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
|
||
defer cancel()
|
||
home, err := os.UserHomeDir()
|
||
require.NoError(t, err, "get home dir")
|
||
if runtime.GOOS == "windows" {
|
||
home = "/" + strings.ReplaceAll(home, "\\", "/")
|
||
}
|
||
|
||
// A configured directory that does not exist on the agent's
|
||
// filesystem must fall back to the home directory.
|
||
missingDir := filepath.Join(t.TempDir(), "does-not-exist")
|
||
//nolint:dogsled
|
||
conn, _, _, _, _ := setupAgent(t, agentsdk.Manifest{
|
||
Directory: missingDir,
|
||
}, 0)
|
||
sshClient, err := conn.SSHClient(ctx)
|
||
require.NoError(t, err)
|
||
defer sshClient.Close()
|
||
client, err := sftp.NewClient(sshClient)
|
||
require.NoError(t, err)
|
||
defer client.Close()
|
||
wd, err := client.Getwd()
|
||
require.NoError(t, err, "get working directory")
|
||
require.Equal(t, home, wd, "working directory should fall back to user home")
|
||
})
|
||
}
|
||
|
||
func TestAgent_SCP(t *testing.T) {
|
||
t.Parallel()
|
||
|
||
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
|
||
defer cancel()
|
||
|
||
//nolint:dogsled
|
||
conn, agentClient, _, _, _ := setupAgent(t, agentsdk.Manifest{}, 0)
|
||
sshClient, err := conn.SSHClient(ctx)
|
||
require.NoError(t, err)
|
||
defer sshClient.Close()
|
||
scpClient, err := scp.NewClientBySSH(sshClient)
|
||
require.NoError(t, err)
|
||
defer scpClient.Close()
|
||
tempFile := filepath.Join(t.TempDir(), "scp")
|
||
content := "hello world"
|
||
err = scpClient.CopyFile(context.Background(), strings.NewReader(content), tempFile, "0755")
|
||
require.NoError(t, err)
|
||
_, err = os.Stat(tempFile)
|
||
require.NoError(t, err)
|
||
|
||
// Close the client to trigger disconnect event.
|
||
scpClient.Close()
|
||
assertConnectionReport(t, agentClient,
|
||
proto.ConnectEvent{Method: proto.Connection_METHOD_SSH},
|
||
proto.DisconnectEvent{Reason: string(codersdk.DisconnectReasonGraceful)},
|
||
)
|
||
}
|
||
|
||
func TestAgent_FileTransferBlocked(t *testing.T) {
|
||
t.Parallel()
|
||
|
||
assertFileTransferBlocked := func(t *testing.T, errorMessage string) {
|
||
// NOTE: Checking content of the error message is flaky. Most likely there is a race condition, which results
|
||
// in stopping the client in different phases, and returning different errors:
|
||
// - client read the full error message: File transfer has been disabled.
|
||
// - client's stream was terminated before reading the error message: EOF
|
||
// - client just read the error code (Windows): Process exited with status 65
|
||
isErr := strings.Contains(errorMessage, agentssh.BlockedFileTransferErrorMessage) ||
|
||
strings.Contains(errorMessage, "EOF") ||
|
||
strings.Contains(errorMessage, "Process exited with status 65")
|
||
require.True(t, isErr, "Message: "+errorMessage)
|
||
}
|
||
|
||
t.Run("SFTP", func(t *testing.T) {
|
||
t.Parallel()
|
||
|
||
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
|
||
defer cancel()
|
||
|
||
//nolint:dogsled
|
||
conn, agentClient, _, _, _ := setupAgent(t, agentsdk.Manifest{}, 0, func(_ *agenttest.Client, o *agent.Options) {
|
||
o.BlockFileTransfer = true
|
||
})
|
||
sshClient, err := conn.SSHClient(ctx)
|
||
require.NoError(t, err)
|
||
defer sshClient.Close()
|
||
_, err = sftp.NewClient(sshClient)
|
||
require.Error(t, err)
|
||
assertFileTransferBlocked(t, err.Error())
|
||
|
||
assertConnectionReport(t, agentClient,
|
||
proto.ConnectEvent{Method: proto.Connection_METHOD_SSH},
|
||
proto.DisconnectEvent{
|
||
Code: agentssh.BlockedFileTransferErrorCode,
|
||
Reason: "file transfer blocked",
|
||
},
|
||
)
|
||
})
|
||
|
||
t.Run("SCP with go-scp package", func(t *testing.T) {
|
||
t.Parallel()
|
||
|
||
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
|
||
defer cancel()
|
||
|
||
//nolint:dogsled
|
||
conn, agentClient, _, _, _ := setupAgent(t, agentsdk.Manifest{}, 0, func(_ *agenttest.Client, o *agent.Options) {
|
||
o.BlockFileTransfer = true
|
||
})
|
||
sshClient, err := conn.SSHClient(ctx)
|
||
require.NoError(t, err)
|
||
defer sshClient.Close()
|
||
scpClient, err := scp.NewClientBySSH(sshClient)
|
||
require.NoError(t, err)
|
||
defer scpClient.Close()
|
||
tempFile := filepath.Join(t.TempDir(), "scp")
|
||
err = scpClient.CopyFile(context.Background(), strings.NewReader("hello world"), tempFile, "0755")
|
||
require.Error(t, err)
|
||
assertFileTransferBlocked(t, err.Error())
|
||
|
||
assertConnectionReport(t, agentClient,
|
||
proto.ConnectEvent{Method: proto.Connection_METHOD_SSH},
|
||
proto.DisconnectEvent{
|
||
Code: agentssh.BlockedFileTransferErrorCode,
|
||
Reason: "file transfer blocked",
|
||
},
|
||
)
|
||
})
|
||
|
||
t.Run("Forbidden commands", func(t *testing.T) {
|
||
t.Parallel()
|
||
|
||
for _, c := range agentssh.BlockedFileTransferCommands {
|
||
t.Run(c, func(t *testing.T) {
|
||
t.Parallel()
|
||
|
||
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
|
||
defer cancel()
|
||
|
||
//nolint:dogsled
|
||
conn, agentClient, _, _, _ := setupAgent(t, agentsdk.Manifest{}, 0, func(_ *agenttest.Client, o *agent.Options) {
|
||
o.BlockFileTransfer = true
|
||
})
|
||
sshClient, err := conn.SSHClient(ctx)
|
||
require.NoError(t, err)
|
||
defer sshClient.Close()
|
||
|
||
session, err := sshClient.NewSession()
|
||
require.NoError(t, err)
|
||
defer session.Close()
|
||
|
||
stdout, err := session.StdoutPipe()
|
||
require.NoError(t, err)
|
||
|
||
//nolint:govet // we don't need `c := c` in Go 1.22
|
||
err = session.Start(c)
|
||
require.NoError(t, err)
|
||
defer session.Close()
|
||
|
||
msg, err := io.ReadAll(stdout)
|
||
require.NoError(t, err)
|
||
assertFileTransferBlocked(t, string(msg))
|
||
|
||
assertConnectionReport(t, agentClient,
|
||
proto.ConnectEvent{Method: proto.Connection_METHOD_SSH},
|
||
proto.DisconnectEvent{
|
||
Code: agentssh.BlockedFileTransferErrorCode,
|
||
Reason: "file transfer blocked",
|
||
},
|
||
)
|
||
})
|
||
}
|
||
})
|
||
}
|
||
|
||
func TestAgent_EnvironmentVariables(t *testing.T) {
|
||
t.Parallel()
|
||
key := "EXAMPLE"
|
||
value := "value"
|
||
session := setupSSHSession(t, agentsdk.Manifest{
|
||
EnvironmentVariables: map[string]string{
|
||
key: value,
|
||
},
|
||
}, codersdk.ServiceBannerConfig{}, nil)
|
||
command := "sh -c 'echo $" + key + "'"
|
||
if runtime.GOOS == "windows" {
|
||
command = "cmd.exe /c echo %" + key + "%"
|
||
}
|
||
output, err := session.Output(command)
|
||
require.NoError(t, err)
|
||
require.Equal(t, value, strings.TrimSpace(string(output)))
|
||
}
|
||
|
||
func TestAgent_EnvironmentVariableExpansion(t *testing.T) {
|
||
t.Parallel()
|
||
key := "EXAMPLE"
|
||
session := setupSSHSession(t, agentsdk.Manifest{
|
||
EnvironmentVariables: map[string]string{
|
||
key: "$SOMETHINGNOTSET",
|
||
},
|
||
}, codersdk.ServiceBannerConfig{}, nil)
|
||
command := "sh -c 'echo $" + key + "'"
|
||
if runtime.GOOS == "windows" {
|
||
command = "cmd.exe /c echo %" + key + "%"
|
||
}
|
||
output, err := session.Output(command)
|
||
require.NoError(t, err)
|
||
expect := ""
|
||
if runtime.GOOS == "windows" {
|
||
expect = "%EXAMPLE%"
|
||
}
|
||
// Output should be empty, because the variable is not set!
|
||
require.Equal(t, expect, strings.TrimSpace(string(output)))
|
||
}
|
||
|
||
func TestAgent_CoderEnvVars(t *testing.T) {
|
||
t.Parallel()
|
||
|
||
for _, key := range []string{"CODER", "CODER_WORKSPACE_NAME", "CODER_WORKSPACE_OWNER_NAME", "CODER_WORKSPACE_AGENT_NAME"} {
|
||
t.Run(key, func(t *testing.T) {
|
||
t.Parallel()
|
||
|
||
session := setupSSHSession(t, agentsdk.Manifest{}, codersdk.ServiceBannerConfig{}, nil)
|
||
command := "sh -c 'echo $" + key + "'"
|
||
if runtime.GOOS == "windows" {
|
||
command = "cmd.exe /c echo %" + key + "%"
|
||
}
|
||
output, err := session.Output(command)
|
||
require.NoError(t, err)
|
||
require.NotEmpty(t, strings.TrimSpace(string(output)))
|
||
})
|
||
}
|
||
}
|
||
|
||
func TestAgent_SSHConnectionEnvVars(t *testing.T) {
|
||
t.Parallel()
|
||
|
||
// Note: the SSH_TTY environment variable should only be set for TTYs.
|
||
// For some reason this test produces a TTY locally and a non-TTY in CI
|
||
// so we don't test for the absence of SSH_TTY.
|
||
for _, key := range []string{"SSH_CONNECTION", "SSH_CLIENT"} {
|
||
t.Run(key, func(t *testing.T) {
|
||
t.Parallel()
|
||
|
||
session := setupSSHSession(t, agentsdk.Manifest{}, codersdk.ServiceBannerConfig{}, nil)
|
||
command := "sh -c 'echo $" + key + "'"
|
||
if runtime.GOOS == "windows" {
|
||
command = "cmd.exe /c echo %" + key + "%"
|
||
}
|
||
output, err := session.Output(command)
|
||
require.NoError(t, err)
|
||
require.NotEmpty(t, strings.TrimSpace(string(output)))
|
||
})
|
||
}
|
||
}
|
||
|
||
func TestAgent_SSHConnectionLoginVars(t *testing.T) {
|
||
t.Parallel()
|
||
|
||
envInfo := usershell.SystemEnvInfo{}
|
||
u, err := envInfo.User()
|
||
require.NoError(t, err, "get current user")
|
||
shell, err := envInfo.Shell(u.Username)
|
||
require.NoError(t, err, "get current shell")
|
||
|
||
tests := []struct {
|
||
key string
|
||
want string
|
||
}{
|
||
{
|
||
key: "USER",
|
||
want: u.Username,
|
||
},
|
||
{
|
||
key: "LOGNAME",
|
||
want: u.Username,
|
||
},
|
||
{
|
||
key: "SHELL",
|
||
want: shell,
|
||
},
|
||
}
|
||
for _, tt := range tests {
|
||
t.Run(tt.key, func(t *testing.T) {
|
||
t.Parallel()
|
||
|
||
session := setupSSHSession(t, agentsdk.Manifest{}, codersdk.ServiceBannerConfig{}, nil)
|
||
command := "sh -c 'echo $" + tt.key + "'"
|
||
if runtime.GOOS == "windows" {
|
||
command = "cmd.exe /c echo %" + tt.key + "%"
|
||
}
|
||
output, err := session.Output(command)
|
||
require.NoError(t, err)
|
||
require.Equal(t, tt.want, strings.TrimSpace(string(output)))
|
||
})
|
||
}
|
||
}
|
||
|
||
// TestAgent_SSHEnvInfoShell verifies that an agent.Options.EnvInfo whose
|
||
// Shell() reports a custom shell is piped through to the SSH session, so the
|
||
// session command runs under that shell instead of the host default.
|
||
func TestAgent_SSHEnvInfoShell(t *testing.T) {
|
||
t.Parallel()
|
||
if runtime.GOOS == "windows" {
|
||
t.Skip("the fake shell is a POSIX script")
|
||
}
|
||
|
||
// A fake shell that ignores its arguments and prints a sentinel. The
|
||
// sentinel only appears in the session output if the injected Shell() was
|
||
// honored. Otherwise the command's own output ("should-not-run") appears.
|
||
const marker = "injected-shell-was-used"
|
||
shellPath := filepath.Join(t.TempDir(), "fakeshell")
|
||
//nolint:gosec // Executable test shell with test-controlled content.
|
||
err := os.WriteFile(shellPath, []byte("#!/bin/sh\necho "+marker+"\n"), 0o700)
|
||
require.NoError(t, err)
|
||
|
||
session := setupSSHSession(t, agentsdk.Manifest{}, codersdk.ServiceBannerConfig{}, nil, func(_ *agenttest.Client, o *agent.Options) {
|
||
o.EnvInfo = shellOverrideEnvInfo{shell: shellPath}
|
||
})
|
||
|
||
output, err := session.Output("echo should-not-run")
|
||
require.NoError(t, err)
|
||
require.Contains(t, string(output), marker)
|
||
require.NotContains(t, string(output), "should-not-run")
|
||
}
|
||
|
||
// shellOverrideEnvInfo is a usershell.EnvInfoer that delegates to the system
|
||
// implementation but reports a custom shell.
|
||
type shellOverrideEnvInfo struct {
|
||
usershell.SystemEnvInfo
|
||
shell string
|
||
}
|
||
|
||
func (e shellOverrideEnvInfo) Shell(string) (string, error) { return e.shell, nil }
|
||
|
||
func TestAgent_Metadata(t *testing.T) {
|
||
t.Parallel()
|
||
|
||
echoHello := "echo 'hello'"
|
||
|
||
t.Run("Once", func(t *testing.T) {
|
||
t.Parallel()
|
||
|
||
//nolint:dogsled
|
||
_, client, _, _, _ := setupAgent(t, agentsdk.Manifest{
|
||
Metadata: []codersdk.WorkspaceAgentMetadataDescription{
|
||
{
|
||
Key: "greeting1",
|
||
Interval: 0,
|
||
Script: echoHello,
|
||
},
|
||
{
|
||
Key: "greeting2",
|
||
Interval: 1,
|
||
Script: echoHello,
|
||
},
|
||
},
|
||
}, 0, func(_ *agenttest.Client, opts *agent.Options) {
|
||
opts.ReportMetadataInterval = testutil.IntervalFast
|
||
})
|
||
|
||
var gotMd map[string]agentsdk.Metadata
|
||
require.Eventually(t, func() bool {
|
||
gotMd = client.GetMetadata()
|
||
return len(gotMd) == 2
|
||
}, testutil.WaitShort, testutil.IntervalFast/2)
|
||
|
||
collectedAt1 := gotMd["greeting1"].CollectedAt
|
||
collectedAt2 := gotMd["greeting2"].CollectedAt
|
||
|
||
require.Eventually(t, func() bool {
|
||
gotMd = client.GetMetadata()
|
||
if len(gotMd) != 2 {
|
||
panic("unexpected number of metadata")
|
||
}
|
||
return !gotMd["greeting2"].CollectedAt.Equal(collectedAt2)
|
||
}, testutil.WaitShort, testutil.IntervalFast/2)
|
||
|
||
require.Equal(t, gotMd["greeting1"].CollectedAt, collectedAt1, "metadata should not be collected again")
|
||
})
|
||
|
||
t.Run("Many", func(t *testing.T) {
|
||
t.Parallel()
|
||
//nolint:dogsled
|
||
_, client, _, _, _ := setupAgent(t, agentsdk.Manifest{
|
||
Metadata: []codersdk.WorkspaceAgentMetadataDescription{
|
||
{
|
||
Key: "greeting",
|
||
Interval: 1,
|
||
Timeout: 100,
|
||
Script: echoHello,
|
||
},
|
||
},
|
||
}, 0, func(_ *agenttest.Client, opts *agent.Options) {
|
||
opts.ReportMetadataInterval = testutil.IntervalFast
|
||
})
|
||
|
||
var gotMd map[string]agentsdk.Metadata
|
||
require.Eventually(t, func() bool {
|
||
gotMd = client.GetMetadata()
|
||
return len(gotMd) == 1
|
||
}, testutil.WaitShort, testutil.IntervalFast/2)
|
||
|
||
collectedAt1 := gotMd["greeting"].CollectedAt
|
||
require.Equal(t, "hello", strings.TrimSpace(gotMd["greeting"].Value))
|
||
|
||
if !assert.Eventually(t, func() bool {
|
||
gotMd = client.GetMetadata()
|
||
return gotMd["greeting"].CollectedAt.After(collectedAt1)
|
||
}, testutil.WaitShort, testutil.IntervalFast/2) {
|
||
t.Fatalf("expected metadata to be collected again")
|
||
}
|
||
})
|
||
}
|
||
|
||
func TestAgentMetadata_Timing(t *testing.T) {
|
||
if runtime.GOOS == "windows" {
|
||
// Shell scripting in Windows is a pain, and we have already tested
|
||
// that the OS logic works in the simpler tests.
|
||
t.SkipNow()
|
||
}
|
||
testutil.SkipIfNotTiming(t)
|
||
t.Parallel()
|
||
|
||
dir := t.TempDir()
|
||
|
||
const reportInterval = 2
|
||
const intervalUnit = 100 * time.Millisecond
|
||
var (
|
||
greetingPath = filepath.Join(dir, "greeting")
|
||
script = "echo hello | tee -a " + greetingPath
|
||
)
|
||
//nolint:dogsled
|
||
_, client, _, _, _ := setupAgent(t, agentsdk.Manifest{
|
||
Metadata: []codersdk.WorkspaceAgentMetadataDescription{
|
||
{
|
||
Key: "greeting",
|
||
Interval: reportInterval,
|
||
Script: script,
|
||
},
|
||
{
|
||
Key: "bad",
|
||
Interval: reportInterval,
|
||
Script: "exit 1",
|
||
},
|
||
},
|
||
}, 0, func(_ *agenttest.Client, opts *agent.Options) {
|
||
opts.ReportMetadataInterval = intervalUnit
|
||
})
|
||
|
||
require.Eventually(t, func() bool {
|
||
return len(client.GetMetadata()) == 2
|
||
}, testutil.WaitShort, testutil.IntervalMedium)
|
||
|
||
for start := time.Now(); time.Since(start) < testutil.WaitMedium; time.Sleep(testutil.IntervalMedium) {
|
||
md := client.GetMetadata()
|
||
require.Len(t, md, 2, "got: %+v", md)
|
||
|
||
require.Equal(t, "hello\n", md["greeting"].Value)
|
||
require.Equal(t, "run cmd: exit status 1", md["bad"].Error)
|
||
|
||
greetingByt, err := os.ReadFile(greetingPath)
|
||
require.NoError(t, err)
|
||
|
||
var (
|
||
numGreetings = bytes.Count(greetingByt, []byte("hello"))
|
||
idealNumGreetings = time.Since(start) / (reportInterval * intervalUnit)
|
||
// We allow a 50% error margin because the report loop may backlog
|
||
// in CI and other toasters. In production, there is no hard
|
||
// guarantee on timing either, and the frontend gives similar
|
||
// wiggle room to the staleness of the value.
|
||
upperBound = int(idealNumGreetings) + 1
|
||
lowerBound = (int(idealNumGreetings) / 2)
|
||
)
|
||
|
||
if idealNumGreetings < 50 {
|
||
// There is an insufficient sample size.
|
||
continue
|
||
}
|
||
|
||
t.Logf("numGreetings: %d, idealNumGreetings: %d", numGreetings, idealNumGreetings)
|
||
// The report loop may slow down on load, but it should never, ever
|
||
// speed up.
|
||
if numGreetings > upperBound {
|
||
t.Fatalf("too many greetings: %d > %d in %v", numGreetings, upperBound, time.Since(start))
|
||
} else if numGreetings < lowerBound {
|
||
t.Fatalf("too few greetings: %d < %d", numGreetings, lowerBound)
|
||
}
|
||
}
|
||
}
|
||
|
||
func TestAgent_Lifecycle(t *testing.T) {
|
||
t.Parallel()
|
||
|
||
t.Run("StartTimeout", func(t *testing.T) {
|
||
t.Parallel()
|
||
|
||
_, client, _, _, _ := setupAgent(t, agentsdk.Manifest{
|
||
Scripts: []codersdk.WorkspaceAgentScript{{
|
||
Script: "sleep 3",
|
||
Timeout: time.Millisecond,
|
||
RunOnStart: true,
|
||
}},
|
||
}, 0)
|
||
|
||
want := []codersdk.WorkspaceAgentLifecycle{
|
||
codersdk.WorkspaceAgentLifecycleStarting,
|
||
codersdk.WorkspaceAgentLifecycleStartTimeout,
|
||
}
|
||
|
||
var got []codersdk.WorkspaceAgentLifecycle
|
||
assert.Eventually(t, func() bool {
|
||
got = client.GetLifecycleStates()
|
||
return slices.Contains(got, want[len(want)-1])
|
||
}, testutil.WaitShort, testutil.IntervalMedium)
|
||
|
||
require.Equal(t, want, got[:len(want)])
|
||
})
|
||
|
||
t.Run("StartError", func(t *testing.T) {
|
||
t.Parallel()
|
||
|
||
_, client, _, _, _ := setupAgent(t, agentsdk.Manifest{
|
||
Scripts: []codersdk.WorkspaceAgentScript{{
|
||
Script: "false",
|
||
Timeout: 30 * time.Second,
|
||
RunOnStart: true,
|
||
}},
|
||
}, 0)
|
||
|
||
want := []codersdk.WorkspaceAgentLifecycle{
|
||
codersdk.WorkspaceAgentLifecycleStarting,
|
||
codersdk.WorkspaceAgentLifecycleStartError,
|
||
}
|
||
|
||
var got []codersdk.WorkspaceAgentLifecycle
|
||
assert.Eventually(t, func() bool {
|
||
got = client.GetLifecycleStates()
|
||
return slices.Contains(got, want[len(want)-1])
|
||
}, testutil.WaitShort, testutil.IntervalMedium)
|
||
|
||
require.Equal(t, want, got[:len(want)])
|
||
})
|
||
|
||
t.Run("Ready", func(t *testing.T) {
|
||
t.Parallel()
|
||
|
||
_, client, _, _, _ := setupAgent(t, agentsdk.Manifest{
|
||
Scripts: []codersdk.WorkspaceAgentScript{{
|
||
Script: "echo foo",
|
||
Timeout: 30 * time.Second,
|
||
RunOnStart: true,
|
||
}},
|
||
}, 0)
|
||
|
||
want := []codersdk.WorkspaceAgentLifecycle{
|
||
codersdk.WorkspaceAgentLifecycleStarting,
|
||
codersdk.WorkspaceAgentLifecycleReady,
|
||
}
|
||
|
||
var got []codersdk.WorkspaceAgentLifecycle
|
||
assert.Eventually(t, func() bool {
|
||
got = client.GetLifecycleStates()
|
||
return len(got) > 0 && got[len(got)-1] == want[len(want)-1]
|
||
}, testutil.WaitShort, testutil.IntervalMedium)
|
||
|
||
require.Equal(t, want, got)
|
||
})
|
||
|
||
t.Run("ShuttingDown", func(t *testing.T) {
|
||
t.Parallel()
|
||
|
||
_, client, _, _, closer := setupAgent(t, agentsdk.Manifest{
|
||
Scripts: []codersdk.WorkspaceAgentScript{{
|
||
Script: "sleep 3",
|
||
Timeout: 30 * time.Second,
|
||
RunOnStop: true,
|
||
}},
|
||
}, 0)
|
||
|
||
assert.Eventually(t, func() bool {
|
||
return slices.Contains(client.GetLifecycleStates(), codersdk.WorkspaceAgentLifecycleReady)
|
||
}, testutil.WaitShort, testutil.IntervalMedium)
|
||
|
||
// Start close asynchronously so that we an inspect the state.
|
||
done := make(chan struct{})
|
||
go func() {
|
||
defer close(done)
|
||
err := closer.Close()
|
||
assert.NoError(t, err)
|
||
}()
|
||
t.Cleanup(func() {
|
||
<-done
|
||
})
|
||
|
||
want := []codersdk.WorkspaceAgentLifecycle{
|
||
codersdk.WorkspaceAgentLifecycleStarting,
|
||
codersdk.WorkspaceAgentLifecycleReady,
|
||
codersdk.WorkspaceAgentLifecycleShuttingDown,
|
||
}
|
||
|
||
var got []codersdk.WorkspaceAgentLifecycle
|
||
assert.Eventually(t, func() bool {
|
||
got = client.GetLifecycleStates()
|
||
return slices.Contains(got, want[len(want)-1])
|
||
}, testutil.WaitShort, testutil.IntervalMedium)
|
||
|
||
require.Equal(t, want, got[:len(want)])
|
||
})
|
||
|
||
t.Run("ShutdownTimeout", func(t *testing.T) {
|
||
t.Parallel()
|
||
|
||
_, client, _, _, closer := setupAgent(t, agentsdk.Manifest{
|
||
Scripts: []codersdk.WorkspaceAgentScript{{
|
||
Script: "sleep 3",
|
||
Timeout: time.Millisecond,
|
||
RunOnStop: true,
|
||
}},
|
||
}, 0)
|
||
|
||
assert.Eventually(t, func() bool {
|
||
return slices.Contains(client.GetLifecycleStates(), codersdk.WorkspaceAgentLifecycleReady)
|
||
}, testutil.WaitShort, testutil.IntervalMedium)
|
||
|
||
// Start close asynchronously so that we an inspect the state.
|
||
done := make(chan struct{})
|
||
go func() {
|
||
defer close(done)
|
||
err := closer.Close()
|
||
assert.NoError(t, err)
|
||
}()
|
||
t.Cleanup(func() {
|
||
<-done
|
||
})
|
||
|
||
want := []codersdk.WorkspaceAgentLifecycle{
|
||
codersdk.WorkspaceAgentLifecycleStarting,
|
||
codersdk.WorkspaceAgentLifecycleReady,
|
||
codersdk.WorkspaceAgentLifecycleShuttingDown,
|
||
codersdk.WorkspaceAgentLifecycleShutdownTimeout,
|
||
}
|
||
|
||
var got []codersdk.WorkspaceAgentLifecycle
|
||
assert.Eventually(t, func() bool {
|
||
got = client.GetLifecycleStates()
|
||
return slices.Contains(got, want[len(want)-1])
|
||
}, testutil.WaitShort, testutil.IntervalMedium)
|
||
|
||
require.Equal(t, want, got[:len(want)])
|
||
})
|
||
|
||
t.Run("ShutdownError", func(t *testing.T) {
|
||
t.Parallel()
|
||
|
||
_, client, _, _, closer := setupAgent(t, agentsdk.Manifest{
|
||
Scripts: []codersdk.WorkspaceAgentScript{{
|
||
Script: "false",
|
||
Timeout: 30 * time.Second,
|
||
RunOnStop: true,
|
||
}},
|
||
}, 0)
|
||
|
||
assert.Eventually(t, func() bool {
|
||
return slices.Contains(client.GetLifecycleStates(), codersdk.WorkspaceAgentLifecycleReady)
|
||
}, testutil.WaitShort, testutil.IntervalMedium)
|
||
|
||
// Start close asynchronously so that we an inspect the state.
|
||
done := make(chan struct{})
|
||
go func() {
|
||
defer close(done)
|
||
err := closer.Close()
|
||
assert.NoError(t, err)
|
||
}()
|
||
t.Cleanup(func() {
|
||
<-done
|
||
})
|
||
|
||
want := []codersdk.WorkspaceAgentLifecycle{
|
||
codersdk.WorkspaceAgentLifecycleStarting,
|
||
codersdk.WorkspaceAgentLifecycleReady,
|
||
codersdk.WorkspaceAgentLifecycleShuttingDown,
|
||
codersdk.WorkspaceAgentLifecycleShutdownError,
|
||
}
|
||
|
||
var got []codersdk.WorkspaceAgentLifecycle
|
||
assert.Eventually(t, func() bool {
|
||
got = client.GetLifecycleStates()
|
||
return slices.Contains(got, want[len(want)-1])
|
||
}, testutil.WaitShort, testutil.IntervalMedium)
|
||
|
||
require.Equal(t, want, got[:len(want)])
|
||
})
|
||
|
||
t.Run("ShutdownScriptOnce", func(t *testing.T) {
|
||
t.Parallel()
|
||
logger := testutil.Logger(t)
|
||
ctx := testutil.Context(t, testutil.WaitMedium)
|
||
expected := "this-is-shutdown"
|
||
derpMap, _ := tailnettest.RunDERPAndSTUN(t)
|
||
statsCh := make(chan *proto.Stats, 50)
|
||
|
||
client := agenttest.NewClient(t,
|
||
logger,
|
||
uuid.New(),
|
||
agentsdk.Manifest{
|
||
DERPMap: derpMap,
|
||
Scripts: []codersdk.WorkspaceAgentScript{{
|
||
ID: uuid.New(),
|
||
LogPath: "coder-startup-script.log",
|
||
Script: "echo 1",
|
||
RunOnStart: true,
|
||
}, {
|
||
ID: uuid.New(),
|
||
LogPath: "coder-shutdown-script.log",
|
||
Script: "echo " + expected,
|
||
RunOnStop: true,
|
||
}},
|
||
},
|
||
statsCh,
|
||
tailnet.NewCoordinator(logger),
|
||
)
|
||
defer client.Close()
|
||
|
||
fs := afero.NewMemMapFs()
|
||
agent := agent.New(agent.Options{
|
||
Client: client,
|
||
Logger: logger.Named("agent"),
|
||
Filesystem: fs,
|
||
})
|
||
|
||
// agent.Close() loads the shutdown script from the agent metadata.
|
||
// The metadata is populated just before execution of the startup script, so it's mandatory to wait
|
||
// until the startup starts.
|
||
require.Eventually(t, func() bool {
|
||
outputPath := filepath.Join(os.TempDir(), "coder-startup-script.log")
|
||
content, err := afero.ReadFile(fs, outputPath)
|
||
if err != nil {
|
||
t.Logf("read file %q: %s", outputPath, err)
|
||
return false
|
||
}
|
||
return len(content) > 0 // something is in the startup log file
|
||
}, testutil.WaitShort, testutil.IntervalMedium)
|
||
|
||
// In order to avoid shutting down the agent before it is fully started and triggering
|
||
// errors, we'll wait until the agent is fully up. It's a bit hokey, but among the last things the agent starts
|
||
// is the stats reporting, so getting a stats report is a good indication the agent is fully up.
|
||
_ = testutil.TryReceive(ctx, t, statsCh)
|
||
|
||
err := agent.Close()
|
||
require.NoError(t, err, "agent should be closed successfully")
|
||
|
||
outputPath := filepath.Join(os.TempDir(), "coder-shutdown-script.log")
|
||
logFirstRead, err := afero.ReadFile(fs, outputPath)
|
||
require.NoError(t, err, "log file should be present")
|
||
require.Equal(t, expected, string(bytes.TrimSpace(logFirstRead)))
|
||
|
||
// Make sure that script can't be executed twice.
|
||
err = agent.Close()
|
||
require.NoError(t, err, "don't need to close the agent twice, no effect")
|
||
|
||
logSecondRead, err := afero.ReadFile(fs, outputPath)
|
||
require.NoError(t, err, "log file should be present")
|
||
require.Equal(t, string(bytes.TrimSpace(logFirstRead)), string(bytes.TrimSpace(logSecondRead)))
|
||
})
|
||
}
|
||
|
||
func TestAgent_Startup(t *testing.T) {
|
||
t.Parallel()
|
||
|
||
t.Run("EmptyDirectory", func(t *testing.T) {
|
||
t.Parallel()
|
||
ctx := testutil.Context(t, testutil.WaitShort)
|
||
|
||
_, client, _, _, _ := setupAgent(t, agentsdk.Manifest{
|
||
Directory: "",
|
||
}, 0)
|
||
startup := testutil.TryReceive(ctx, t, client.GetStartup())
|
||
require.Equal(t, "", startup.GetExpandedDirectory())
|
||
})
|
||
|
||
t.Run("HomeDirectory", func(t *testing.T) {
|
||
t.Parallel()
|
||
ctx := testutil.Context(t, testutil.WaitShort)
|
||
|
||
_, client, _, _, _ := setupAgent(t, agentsdk.Manifest{
|
||
Directory: "~",
|
||
}, 0)
|
||
startup := testutil.TryReceive(ctx, t, client.GetStartup())
|
||
homeDir, err := os.UserHomeDir()
|
||
require.NoError(t, err)
|
||
require.Equal(t, homeDir, startup.GetExpandedDirectory())
|
||
})
|
||
|
||
t.Run("NotAbsoluteDirectory", func(t *testing.T) {
|
||
t.Parallel()
|
||
ctx := testutil.Context(t, testutil.WaitShort)
|
||
|
||
_, client, _, _, _ := setupAgent(t, agentsdk.Manifest{
|
||
Directory: "coder/coder",
|
||
}, 0)
|
||
startup := testutil.TryReceive(ctx, t, client.GetStartup())
|
||
homeDir, err := os.UserHomeDir()
|
||
require.NoError(t, err)
|
||
require.Equal(t, filepath.Join(homeDir, "coder/coder"), startup.GetExpandedDirectory())
|
||
})
|
||
|
||
t.Run("HomeEnvironmentVariable", func(t *testing.T) {
|
||
t.Parallel()
|
||
ctx := testutil.Context(t, testutil.WaitShort)
|
||
|
||
_, client, _, _, _ := setupAgent(t, agentsdk.Manifest{
|
||
Directory: "$HOME",
|
||
}, 0)
|
||
startup := testutil.TryReceive(ctx, t, client.GetStartup())
|
||
homeDir, err := os.UserHomeDir()
|
||
require.NoError(t, err)
|
||
require.Equal(t, homeDir, startup.GetExpandedDirectory())
|
||
})
|
||
}
|
||
|
||
//nolint:paralleltest // This test sets an environment variable.
|
||
func TestAgent_ReconnectingPTY(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.")
|
||
}
|
||
|
||
backends := []string{"Buffered", "Screen"}
|
||
|
||
_, err := exec.LookPath("screen")
|
||
hasScreen := err == nil
|
||
|
||
tmuxPath, err := exec.LookPath("tmux")
|
||
hasTmux := err == nil
|
||
|
||
// Make sure UTF-8 works even with locale variables set to C.
|
||
t.Setenv("LANG", "C")
|
||
t.Setenv("LC_CTYPE", "C")
|
||
t.Setenv("LC_ALL", "")
|
||
|
||
for _, backendType := range backends {
|
||
t.Run(backendType, func(t *testing.T) {
|
||
if backendType == "Screen" {
|
||
if runtime.GOOS != "linux" {
|
||
t.Skipf("`screen` is not supported on %s", runtime.GOOS)
|
||
} else if !hasScreen {
|
||
t.Skip("`screen` not found")
|
||
}
|
||
} else if hasScreen && runtime.GOOS == "linux" {
|
||
// Set up a PATH that does not have screen in it.
|
||
bashPath, err := exec.LookPath("bash")
|
||
require.NoError(t, err)
|
||
dir, err := os.MkdirTemp("/tmp", "coder-test-reconnecting-pty-PATH")
|
||
require.NoError(t, err, "create temp dir for reconnecting pty PATH")
|
||
err = os.Symlink(bashPath, filepath.Join(dir, "bash"))
|
||
require.NoError(t, err, "symlink bash into reconnecting pty PATH")
|
||
t.Setenv("PATH", dir)
|
||
}
|
||
|
||
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
|
||
defer cancel()
|
||
|
||
//nolint:dogsled
|
||
conn, agentClient, _, _, _ := setupAgent(t, agentsdk.Manifest{}, 0)
|
||
idConnectionReport := uuid.New()
|
||
id := uuid.New()
|
||
|
||
// Test that the connection is reported. This must be tested in the
|
||
// first connection because we care about verifying all of these.
|
||
netConn0, err := conn.ReconnectingPTY(ctx, idConnectionReport, 80, 80, "bash --norc")
|
||
require.NoError(t, err)
|
||
_ = netConn0.Close()
|
||
assertConnectionReport(t, agentClient,
|
||
proto.ConnectEvent{Method: proto.Connection_METHOD_RECONNECTING_PTY},
|
||
proto.DisconnectEvent{Reason: string(codersdk.DisconnectReasonGraceful)},
|
||
)
|
||
|
||
// --norc disables executing .bashrc, which is often used to customize the bash prompt
|
||
netConn1, err := conn.ReconnectingPTY(ctx, id, 80, 80, "bash --norc")
|
||
require.NoError(t, err)
|
||
defer netConn1.Close()
|
||
tr1 := testutil.NewTerminalReader(t, netConn1)
|
||
|
||
// A second simultaneous connection.
|
||
netConn2, err := conn.ReconnectingPTY(ctx, id, 80, 80, "bash --norc")
|
||
require.NoError(t, err)
|
||
defer netConn2.Close()
|
||
tr2 := testutil.NewTerminalReader(t, netConn2)
|
||
|
||
matchPrompt := func(line string) bool {
|
||
return strings.Contains(line, "$ ") || strings.Contains(line, "# ")
|
||
}
|
||
matchEchoCommand := func(line string) bool {
|
||
return strings.Contains(line, "echo test")
|
||
}
|
||
matchEchoOutput := func(line string) bool {
|
||
return strings.Contains(line, "test") && !strings.Contains(line, "echo")
|
||
}
|
||
matchExitCommand := func(line string) bool {
|
||
return strings.Contains(line, "exit")
|
||
}
|
||
matchExitOutput := func(line string) bool {
|
||
return strings.Contains(line, "exit") || strings.Contains(line, "logout")
|
||
}
|
||
|
||
// Wait for the prompt before writing commands. If the command
|
||
// arrives before the prompt is written, screen will sometimes put
|
||
// the command output on the same line as the command and the test
|
||
// will flake.
|
||
require.NoError(t, tr1.ReadUntil(ctx, matchPrompt), "find prompt")
|
||
require.NoError(t, tr2.ReadUntil(ctx, matchPrompt), "find prompt")
|
||
|
||
data, err := json.Marshal(workspacesdk.ReconnectingPTYRequest{
|
||
Data: "printf '%s\\n' \"$TERM\"\r",
|
||
})
|
||
require.NoError(t, err)
|
||
_, err = netConn1.Write(data)
|
||
require.NoError(t, err)
|
||
require.NoError(t, tr1.ReadUntilString(ctx, "xterm-256color"), "find TERM output")
|
||
require.NoError(t, tr2.ReadUntilString(ctx, "xterm-256color"), "find TERM output")
|
||
require.NoError(t, tr1.ReadUntil(ctx, matchPrompt), "find prompt")
|
||
require.NoError(t, tr2.ReadUntil(ctx, matchPrompt), "find prompt")
|
||
|
||
data, err = json.Marshal(workspacesdk.ReconnectingPTYRequest{
|
||
Data: "echo test\r",
|
||
})
|
||
require.NoError(t, err)
|
||
_, err = netConn1.Write(data)
|
||
require.NoError(t, err)
|
||
|
||
// Once for typing the command...
|
||
require.NoError(t, tr1.ReadUntil(ctx, matchEchoCommand), "find echo command")
|
||
// And another time for the actual output.
|
||
require.NoError(t, tr1.ReadUntil(ctx, matchEchoOutput), "find echo output")
|
||
|
||
// Same for the other connection.
|
||
require.NoError(t, tr2.ReadUntil(ctx, matchEchoCommand), "find echo command")
|
||
require.NoError(t, tr2.ReadUntil(ctx, matchEchoOutput), "find echo output")
|
||
|
||
_ = netConn1.Close()
|
||
_ = netConn2.Close()
|
||
netConn3, err := conn.ReconnectingPTY(ctx, id, 80, 80, "bash --norc")
|
||
require.NoError(t, err)
|
||
defer netConn3.Close()
|
||
tr3 := testutil.NewTerminalReader(t, netConn3)
|
||
|
||
// Same output again!
|
||
require.NoError(t, tr3.ReadUntil(ctx, matchEchoCommand), "find echo command")
|
||
require.NoError(t, tr3.ReadUntil(ctx, matchEchoOutput), "find echo output")
|
||
|
||
// Exit should cause the connection to close.
|
||
data, err = json.Marshal(workspacesdk.ReconnectingPTYRequest{
|
||
Data: "exit\r",
|
||
})
|
||
require.NoError(t, err)
|
||
_, err = netConn3.Write(data)
|
||
require.NoError(t, err)
|
||
|
||
// Once for the input and again for the output.
|
||
require.NoError(t, tr3.ReadUntil(ctx, matchExitCommand), "find exit command")
|
||
require.NoError(t, tr3.ReadUntil(ctx, matchExitOutput), "find exit output")
|
||
|
||
// Wait for the connection to close.
|
||
require.ErrorIs(t, tr3.ReadUntil(ctx, nil), io.EOF)
|
||
|
||
// Try a non-shell command. It should output then immediately exit.
|
||
netConn4, err := conn.ReconnectingPTY(ctx, uuid.New(), 80, 80, "echo test")
|
||
require.NoError(t, err)
|
||
defer netConn4.Close()
|
||
|
||
tr4 := testutil.NewTerminalReader(t, netConn4)
|
||
require.NoError(t, tr4.ReadUntil(ctx, matchEchoOutput), "find echo output")
|
||
require.ErrorIs(t, tr4.ReadUntil(ctx, nil), io.EOF)
|
||
|
||
// Ensure that UTF-8 is supported. Avoid the terminal emulator because it
|
||
// does not appear to support UTF-8, just make sure the bytes that come
|
||
// back have the character in it.
|
||
netConn5, err := conn.ReconnectingPTY(ctx, uuid.New(), 80, 80, "echo ❯")
|
||
require.NoError(t, err)
|
||
defer netConn5.Close()
|
||
|
||
bytes, err := io.ReadAll(netConn5)
|
||
require.NoError(t, err)
|
||
require.Contains(t, string(bytes), "❯")
|
||
|
||
if !hasTmux {
|
||
t.Log("`tmux` not found, skipping tmux glyph regression")
|
||
} else {
|
||
glyphs := "⚠╭╮╰╯•›│─█▓░▄❯✔╌"
|
||
tmuxSocket := "coder-test-" + strings.ReplaceAll(uuid.NewString(), "-", "")
|
||
t.Cleanup(func() {
|
||
_ = exec.Command(tmuxPath, "-L", tmuxSocket, "kill-server").Run()
|
||
})
|
||
// Keep the pane alive with a shell builtin until the read loop sees
|
||
// the glyphs, otherwise tmux can restore the alternate screen first.
|
||
command := fmt.Sprintf(
|
||
"%s -L %s new-session %q",
|
||
strconv.Quote(tmuxPath),
|
||
tmuxSocket,
|
||
fmt.Sprintf("printf '%%s\\n' '%s'; read _", glyphs),
|
||
)
|
||
netConn6, err := conn.ReconnectingPTY(ctx, uuid.New(), 80, 80, command)
|
||
require.NoError(t, err)
|
||
defer netConn6.Close()
|
||
|
||
var output strings.Builder
|
||
buffer := make([]byte, 1024)
|
||
deadline := time.Now().Add(testutil.WaitMedium)
|
||
for !strings.Contains(output.String(), glyphs) {
|
||
if time.Now().After(deadline) {
|
||
require.Contains(t, output.String(), glyphs)
|
||
}
|
||
require.NoError(t, netConn6.SetReadDeadline(time.Now().Add(testutil.IntervalMedium)))
|
||
read, err := netConn6.Read(buffer)
|
||
if read > 0 {
|
||
_, _ = output.Write(buffer[:read])
|
||
}
|
||
var netErr net.Error
|
||
if errors.As(err, &netErr) && netErr.Timeout() {
|
||
continue
|
||
}
|
||
require.NoError(t, err)
|
||
}
|
||
}
|
||
})
|
||
}
|
||
}
|
||
|
||
// This tests end-to-end functionality of connecting to a running container
|
||
// and executing a command. It creates a real Docker container and runs a
|
||
// command. As such, it does not run by default in CI.
|
||
// You can run it manually as follows:
|
||
//
|
||
// CODER_TEST_USE_DOCKER=1 go test -count=1 ./agent -run TestAgent_ReconnectingPTYContainer
|
||
func TestAgent_ReconnectingPTYContainer(t *testing.T) {
|
||
t.Parallel()
|
||
if os.Getenv("CODER_TEST_USE_DOCKER") != "1" {
|
||
t.Skip("Set CODER_TEST_USE_DOCKER=1 to run this test")
|
||
}
|
||
if _, err := exec.LookPath("devcontainer"); err != nil {
|
||
t.Skip("This test requires the devcontainer CLI: npm install -g @devcontainers/cli")
|
||
}
|
||
|
||
pool, err := dockertest.NewPool("")
|
||
require.NoError(t, err, "Could not connect to docker")
|
||
ct, err := pool.RunWithOptions(&dockertest.RunOptions{
|
||
Repository: "busybox",
|
||
Tag: "latest",
|
||
Cmd: []string{"sleep", "infnity"},
|
||
}, func(config *docker.HostConfig) {
|
||
config.AutoRemove = true
|
||
config.RestartPolicy = docker.RestartPolicy{Name: "no"}
|
||
})
|
||
require.NoError(t, err, "Could not start container")
|
||
defer func() {
|
||
err := pool.Purge(ct)
|
||
require.NoError(t, err, "Could not stop container")
|
||
}()
|
||
// Wait for container to start
|
||
require.Eventually(t, func() bool {
|
||
ct, ok := pool.ContainerByName(ct.Container.Name)
|
||
return ok && ct.Container.State.Running
|
||
}, testutil.WaitShort, testutil.IntervalSlow, "Container did not start in time")
|
||
|
||
// nolint: dogsled
|
||
conn, _, _, _, _ := setupAgent(t, agentsdk.Manifest{}, 0, func(_ *agenttest.Client, o *agent.Options) {
|
||
o.Devcontainers = true
|
||
o.DevcontainerAPIOptions = append(o.DevcontainerAPIOptions,
|
||
agentcontainers.WithContainerLabelIncludeFilter("this.label.does.not.exist.ignore.devcontainers", "true"),
|
||
)
|
||
})
|
||
ctx := testutil.Context(t, testutil.WaitLong)
|
||
ac, err := conn.ReconnectingPTY(ctx, uuid.New(), 80, 80, "/bin/sh", func(arp *workspacesdk.AgentReconnectingPTYInit) {
|
||
arp.Container = ct.Container.ID
|
||
})
|
||
require.NoError(t, err, "failed to create ReconnectingPTY")
|
||
defer ac.Close()
|
||
tr := testutil.NewTerminalReader(t, ac)
|
||
|
||
require.NoError(t, tr.ReadUntil(ctx, func(line string) bool {
|
||
return strings.Contains(line, "#") || strings.Contains(line, "$")
|
||
}), "find prompt")
|
||
|
||
require.NoError(t, json.NewEncoder(ac).Encode(workspacesdk.ReconnectingPTYRequest{
|
||
Data: "hostname\r",
|
||
}), "write hostname")
|
||
require.NoError(t, tr.ReadUntil(ctx, func(line string) bool {
|
||
return strings.Contains(line, "hostname")
|
||
}), "find hostname command")
|
||
|
||
require.NoError(t, tr.ReadUntil(ctx, func(line string) bool {
|
||
return strings.Contains(line, ct.Container.Config.Hostname)
|
||
}), "find hostname output")
|
||
require.NoError(t, json.NewEncoder(ac).Encode(workspacesdk.ReconnectingPTYRequest{
|
||
Data: "exit\r",
|
||
}), "write exit command")
|
||
|
||
// Wait for the connection to close.
|
||
require.ErrorIs(t, tr.ReadUntil(ctx, nil), io.EOF)
|
||
}
|
||
|
||
type subAgentRequestPayload struct {
|
||
Token string `json:"token"`
|
||
Directory string `json:"directory"`
|
||
}
|
||
|
||
// runSubAgentMain is the main function for the sub-agent that connects
|
||
// to the control plane. It reads the CODER_AGENT_URL and
|
||
// CODER_AGENT_TOKEN environment variables, sends the token, and exits
|
||
// with a status code based on the response.
|
||
func runSubAgentMain() int {
|
||
url := os.Getenv("CODER_AGENT_URL")
|
||
token := os.Getenv("CODER_AGENT_TOKEN")
|
||
if url == "" || token == "" {
|
||
_, _ = fmt.Fprintln(os.Stderr, "CODER_AGENT_URL and CODER_AGENT_TOKEN must be set")
|
||
return 10
|
||
}
|
||
|
||
dir, err := os.Getwd()
|
||
if err != nil {
|
||
_, _ = fmt.Fprintf(os.Stderr, "failed to get current working directory: %v\n", err)
|
||
return 1
|
||
}
|
||
payload := subAgentRequestPayload{
|
||
Token: token,
|
||
Directory: dir,
|
||
}
|
||
b, err := json.Marshal(payload)
|
||
if err != nil {
|
||
_, _ = fmt.Fprintf(os.Stderr, "failed to marshal payload: %v\n", err)
|
||
return 1
|
||
}
|
||
|
||
req, err := http.NewRequest("POST", url, bytes.NewReader(b))
|
||
if err != nil {
|
||
_, _ = fmt.Fprintf(os.Stderr, "failed to create request: %v\n", err)
|
||
return 1
|
||
}
|
||
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
|
||
defer cancel()
|
||
req = req.WithContext(ctx)
|
||
client := &http.Client{}
|
||
resp, err := client.Do(req)
|
||
if err != nil {
|
||
_, _ = fmt.Fprintf(os.Stderr, "agent connection failed: %v\n", err)
|
||
return 11
|
||
}
|
||
defer resp.Body.Close()
|
||
if resp.StatusCode != http.StatusOK {
|
||
_, _ = fmt.Fprintf(os.Stderr, "agent exiting with non-zero exit code %d\n", resp.StatusCode)
|
||
return 12
|
||
}
|
||
_, _ = fmt.Println("sub-agent connected successfully")
|
||
return 0
|
||
}
|
||
|
||
// This tests end-to-end functionality of auto-starting a devcontainer.
|
||
// It runs "devcontainer up" which creates a real Docker container. As
|
||
// such, it does not run by default in CI.
|
||
//
|
||
// You can run it manually as follows:
|
||
//
|
||
// CODER_TEST_USE_DOCKER=1 go test -count=1 ./agent -run TestAgent_DevcontainerAutostart
|
||
//
|
||
//nolint:paralleltest // This test sets an environment variable.
|
||
func TestAgent_DevcontainerAutostart(t *testing.T) {
|
||
if os.Getenv("CODER_TEST_USE_DOCKER") != "1" {
|
||
t.Skip("Set CODER_TEST_USE_DOCKER=1 to run this test")
|
||
}
|
||
if _, err := exec.LookPath("devcontainer"); err != nil {
|
||
t.Skip("This test requires the devcontainer CLI: npm install -g @devcontainers/cli")
|
||
}
|
||
|
||
// This HTTP handler handles requests from runSubAgentMain which
|
||
// acts as a fake sub-agent. We want to verify that the sub-agent
|
||
// connects and sends its token. We use a channel to signal
|
||
// that the sub-agent has connected successfully and then we wait
|
||
// until we receive another signal to return from the handler. This
|
||
// keeps the agent "alive" for as long as we want.
|
||
subAgentConnected := make(chan subAgentRequestPayload, 1)
|
||
subAgentReady := make(chan struct{}, 1)
|
||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||
if r.Method == http.MethodGet && strings.HasPrefix(r.URL.Path, "/api/v2/workspaceagents/me/") {
|
||
return
|
||
}
|
||
|
||
t.Logf("Sub-agent request received: %s %s", r.Method, r.URL.Path)
|
||
|
||
if r.Method != http.MethodPost {
|
||
http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
|
||
return
|
||
}
|
||
|
||
// Read the token from the request body.
|
||
var payload subAgentRequestPayload
|
||
if err := json.NewDecoder(r.Body).Decode(&payload); err != nil {
|
||
http.Error(w, "Failed to read token", http.StatusBadRequest)
|
||
t.Logf("Failed to read token: %v", err)
|
||
return
|
||
}
|
||
defer r.Body.Close()
|
||
|
||
t.Logf("Sub-agent request payload received: %+v", payload)
|
||
|
||
// Signal that the sub-agent has connected successfully.
|
||
select {
|
||
case <-t.Context().Done():
|
||
t.Logf("Test context done, not processing sub-agent request")
|
||
return
|
||
case subAgentConnected <- payload:
|
||
}
|
||
|
||
// Wait for the signal to return from the handler.
|
||
select {
|
||
case <-t.Context().Done():
|
||
t.Logf("Test context done, not waiting for sub-agent ready")
|
||
return
|
||
case <-subAgentReady:
|
||
}
|
||
|
||
w.WriteHeader(http.StatusOK)
|
||
}))
|
||
defer srv.Close()
|
||
|
||
pool, err := dockertest.NewPool("")
|
||
require.NoError(t, err, "Could not connect to docker")
|
||
|
||
// Prepare temporary devcontainer for test (mywork).
|
||
devcontainerID := uuid.New()
|
||
tmpdir := t.TempDir()
|
||
t.Setenv("HOME", tmpdir)
|
||
tempWorkspaceFolder := filepath.Join(tmpdir, "mywork")
|
||
unexpandedWorkspaceFolder := filepath.Join("~", "mywork")
|
||
t.Logf("Workspace folder: %s", tempWorkspaceFolder)
|
||
t.Logf("Unexpanded workspace folder: %s", unexpandedWorkspaceFolder)
|
||
devcontainerPath := filepath.Join(tempWorkspaceFolder, ".devcontainer")
|
||
err = os.MkdirAll(devcontainerPath, 0o755)
|
||
require.NoError(t, err, "create devcontainer directory")
|
||
devcontainerFile := filepath.Join(devcontainerPath, "devcontainer.json")
|
||
err = os.WriteFile(devcontainerFile, []byte(`{
|
||
"name": "mywork",
|
||
"image": "ubuntu:latest",
|
||
"cmd": ["sleep", "infinity"],
|
||
"runArgs": ["--network=host", "--label=`+agentcontainers.DevcontainerIsTestRunLabel+`=true"]
|
||
}`), 0o600)
|
||
require.NoError(t, err, "write devcontainer.json")
|
||
|
||
manifest := agentsdk.Manifest{
|
||
// Set up pre-conditions for auto-starting a devcontainer, the script
|
||
// is expected to be prepared by the provisioner normally.
|
||
Devcontainers: []codersdk.WorkspaceAgentDevcontainer{
|
||
{
|
||
ID: devcontainerID,
|
||
Name: "test",
|
||
// Use an unexpanded path to test the expansion.
|
||
WorkspaceFolder: unexpandedWorkspaceFolder,
|
||
},
|
||
},
|
||
Scripts: []codersdk.WorkspaceAgentScript{
|
||
{
|
||
ID: devcontainerID,
|
||
LogSourceID: agentsdk.ExternalLogSourceID,
|
||
RunOnStart: true,
|
||
Script: "echo this-will-be-replaced",
|
||
DisplayName: "Dev Container (test)",
|
||
},
|
||
},
|
||
}
|
||
mClock := quartz.NewMock(t)
|
||
mClock.Set(time.Now())
|
||
tickerFuncTrap := mClock.Trap().TickerFunc("agentcontainers")
|
||
|
||
//nolint:dogsled
|
||
_, agentClient, _, _, _ := setupAgent(t, manifest, 0, func(_ *agenttest.Client, o *agent.Options) {
|
||
o.Devcontainers = true
|
||
o.DevcontainerAPIOptions = append(
|
||
o.DevcontainerAPIOptions,
|
||
// Only match this specific dev container.
|
||
agentcontainers.WithClock(mClock),
|
||
agentcontainers.WithContainerLabelIncludeFilter("devcontainer.local_folder", tempWorkspaceFolder),
|
||
agentcontainers.WithContainerLabelIncludeFilter(agentcontainers.DevcontainerIsTestRunLabel, "true"),
|
||
agentcontainers.WithSubAgentURL(srv.URL),
|
||
// The agent will copy "itself", but in the case of this test, the
|
||
// agent is actually this test binary. So we'll tell the test binary
|
||
// to execute the sub-agent main function via this env.
|
||
agentcontainers.WithSubAgentEnv("CODER_TEST_RUN_SUB_AGENT_MAIN=1"),
|
||
)
|
||
})
|
||
|
||
t.Logf("Waiting for container with label: devcontainer.local_folder=%s", tempWorkspaceFolder)
|
||
|
||
var container docker.APIContainers
|
||
require.Eventually(t, func() bool {
|
||
containers, err := pool.Client.ListContainers(docker.ListContainersOptions{All: true})
|
||
if err != nil {
|
||
t.Logf("Error listing containers: %v", err)
|
||
return false
|
||
}
|
||
|
||
for _, c := range containers {
|
||
t.Logf("Found container: %s with labels: %v", c.ID[:12], c.Labels)
|
||
if labelValue, ok := c.Labels["devcontainer.local_folder"]; ok {
|
||
if labelValue == tempWorkspaceFolder {
|
||
t.Logf("Found matching container: %s", c.ID[:12])
|
||
container = c
|
||
return true
|
||
}
|
||
}
|
||
}
|
||
|
||
return false
|
||
}, testutil.WaitSuperLong, testutil.IntervalMedium, "no container with workspace folder label found")
|
||
defer func() {
|
||
// We can't rely on pool here because the container is not
|
||
// managed by it (it is managed by @devcontainer/cli).
|
||
err := pool.Client.RemoveContainer(docker.RemoveContainerOptions{
|
||
ID: container.ID,
|
||
RemoveVolumes: true,
|
||
Force: true,
|
||
})
|
||
assert.NoError(t, err, "remove container")
|
||
}()
|
||
|
||
containerInfo, err := pool.Client.InspectContainer(container.ID)
|
||
require.NoError(t, err, "inspect container")
|
||
t.Logf("Container state: status: %v", containerInfo.State.Status)
|
||
require.True(t, containerInfo.State.Running, "container should be running")
|
||
|
||
ctx := testutil.Context(t, testutil.WaitLong)
|
||
|
||
// Ensure the container update routine runs.
|
||
tickerFuncTrap.MustWait(ctx).MustRelease(ctx)
|
||
tickerFuncTrap.Close()
|
||
|
||
// Since the agent does RefreshContainers, and the ticker function
|
||
// is set to skip instead of queue, we must advance the clock
|
||
// multiple times to ensure that the sub-agent is created.
|
||
var subAgents []*proto.SubAgent
|
||
for {
|
||
_, next := mClock.AdvanceNext()
|
||
next.MustWait(ctx)
|
||
|
||
// Verify that a subagent was created.
|
||
subAgents = agentClient.GetSubAgents()
|
||
if len(subAgents) > 0 {
|
||
t.Logf("Found sub-agents: %d", len(subAgents))
|
||
break
|
||
}
|
||
}
|
||
require.Len(t, subAgents, 1, "expected one sub agent")
|
||
|
||
subAgent := subAgents[0]
|
||
subAgentID, err := uuid.FromBytes(subAgent.GetId())
|
||
require.NoError(t, err, "failed to parse sub-agent ID")
|
||
t.Logf("Connecting to sub-agent: %s (ID: %s)", subAgent.Name, subAgentID)
|
||
|
||
gotDir, err := agentClient.GetSubAgentDirectory(subAgentID)
|
||
require.NoError(t, err, "failed to get sub-agent directory")
|
||
require.Equal(t, "/workspaces/mywork", gotDir, "sub-agent directory should match")
|
||
|
||
subAgentToken, err := uuid.FromBytes(subAgent.GetAuthToken())
|
||
require.NoError(t, err, "failed to parse sub-agent token")
|
||
|
||
payload := testutil.RequireReceive(ctx, t, subAgentConnected)
|
||
require.Equal(t, subAgentToken.String(), payload.Token, "sub-agent token should match")
|
||
require.Equal(t, "/workspaces/mywork", payload.Directory, "sub-agent directory should match")
|
||
|
||
// Allow the subagent to exit.
|
||
close(subAgentReady)
|
||
}
|
||
|
||
// TestAgent_DevcontainerRecreate tests that RecreateDevcontainer
|
||
// recreates a devcontainer and emits logs.
|
||
//
|
||
// This tests end-to-end functionality of auto-starting a devcontainer.
|
||
// It runs "devcontainer up" which creates a real Docker container. As
|
||
// such, it does not run by default in CI.
|
||
//
|
||
// You can run it manually as follows:
|
||
//
|
||
// CODER_TEST_USE_DOCKER=1 go test -count=1 ./agent -run TestAgent_DevcontainerRecreate
|
||
func TestAgent_DevcontainerRecreate(t *testing.T) {
|
||
if os.Getenv("CODER_TEST_USE_DOCKER") != "1" {
|
||
t.Skip("Set CODER_TEST_USE_DOCKER=1 to run this test")
|
||
}
|
||
t.Parallel()
|
||
|
||
pool, err := dockertest.NewPool("")
|
||
require.NoError(t, err, "Could not connect to docker")
|
||
|
||
// Prepare temporary devcontainer for test (mywork).
|
||
devcontainerID := uuid.New()
|
||
devcontainerLogSourceID := uuid.New()
|
||
workspaceFolder := filepath.Join(t.TempDir(), "mywork")
|
||
t.Logf("Workspace folder: %s", workspaceFolder)
|
||
devcontainerPath := filepath.Join(workspaceFolder, ".devcontainer")
|
||
err = os.MkdirAll(devcontainerPath, 0o755)
|
||
require.NoError(t, err, "create devcontainer directory")
|
||
devcontainerFile := filepath.Join(devcontainerPath, "devcontainer.json")
|
||
err = os.WriteFile(devcontainerFile, []byte(`{
|
||
"name": "mywork",
|
||
"image": "busybox:latest",
|
||
"cmd": ["sleep", "infinity"],
|
||
"runArgs": ["--label=`+agentcontainers.DevcontainerIsTestRunLabel+`=true"]
|
||
}`), 0o600)
|
||
require.NoError(t, err, "write devcontainer.json")
|
||
|
||
manifest := agentsdk.Manifest{
|
||
// Set up pre-conditions for auto-starting a devcontainer, the
|
||
// script is used to extract the log source ID.
|
||
Devcontainers: []codersdk.WorkspaceAgentDevcontainer{
|
||
{
|
||
ID: devcontainerID,
|
||
Name: "test",
|
||
WorkspaceFolder: workspaceFolder,
|
||
},
|
||
},
|
||
Scripts: []codersdk.WorkspaceAgentScript{
|
||
{
|
||
ID: devcontainerID,
|
||
LogSourceID: devcontainerLogSourceID,
|
||
},
|
||
},
|
||
}
|
||
|
||
//nolint:dogsled
|
||
conn, client, _, _, _ := setupAgent(t, manifest, 0, func(_ *agenttest.Client, o *agent.Options) {
|
||
o.Devcontainers = true
|
||
o.DevcontainerAPIOptions = append(o.DevcontainerAPIOptions,
|
||
agentcontainers.WithContainerLabelIncludeFilter("devcontainer.local_folder", workspaceFolder),
|
||
agentcontainers.WithContainerLabelIncludeFilter(agentcontainers.DevcontainerIsTestRunLabel, "true"),
|
||
)
|
||
})
|
||
|
||
ctx := testutil.Context(t, testutil.WaitLong)
|
||
|
||
// We enabled autostart for the devcontainer, so ready is a good
|
||
// indication that the devcontainer is up and running. Importantly,
|
||
// this also means that the devcontainer startup is no longer
|
||
// producing logs that may interfere with the recreate logs.
|
||
testutil.Eventually(ctx, t, func(context.Context) bool {
|
||
states := client.GetLifecycleStates()
|
||
return slices.Contains(states, codersdk.WorkspaceAgentLifecycleReady)
|
||
}, testutil.IntervalMedium, "devcontainer not ready")
|
||
|
||
t.Logf("Looking for container with label: devcontainer.local_folder=%s", workspaceFolder)
|
||
|
||
var container codersdk.WorkspaceAgentContainer
|
||
testutil.Eventually(ctx, t, func(context.Context) bool {
|
||
resp, err := conn.ListContainers(ctx)
|
||
if err != nil {
|
||
t.Logf("Error listing containers: %v", err)
|
||
return false
|
||
}
|
||
for _, c := range resp.Containers {
|
||
t.Logf("Found container: %s with labels: %v", c.ID[:12], c.Labels)
|
||
if v, ok := c.Labels["devcontainer.local_folder"]; ok && v == workspaceFolder {
|
||
t.Logf("Found matching container: %s", c.ID[:12])
|
||
container = c
|
||
return true
|
||
}
|
||
}
|
||
return false
|
||
}, testutil.IntervalMedium, "no container with workspace folder label found")
|
||
defer func(container codersdk.WorkspaceAgentContainer) {
|
||
// We can't rely on pool here because the container is not
|
||
// managed by it (it is managed by @devcontainer/cli).
|
||
err := pool.Client.RemoveContainer(docker.RemoveContainerOptions{
|
||
ID: container.ID,
|
||
RemoveVolumes: true,
|
||
Force: true,
|
||
})
|
||
assert.Error(t, err, "container should be removed by recreate")
|
||
}(container)
|
||
|
||
ctx = testutil.Context(t, testutil.WaitLong) // Reset context.
|
||
|
||
// Capture logs via ScriptLogger.
|
||
logsCh := make(chan *proto.BatchCreateLogsRequest, 1)
|
||
client.SetLogsChannel(logsCh)
|
||
|
||
// Invoke recreate to trigger the destruction and recreation of the
|
||
// devcontainer, we do it in a goroutine so we can process logs
|
||
// concurrently.
|
||
go func(container codersdk.WorkspaceAgentContainer) {
|
||
_, err := conn.RecreateDevcontainer(ctx, devcontainerID.String())
|
||
assert.NoError(t, err, "recreate devcontainer should succeed")
|
||
}(container)
|
||
|
||
t.Logf("Checking recreate logs for outcome...")
|
||
|
||
// Wait for the logs to be emitted, the @devcontainer/cli up command
|
||
// will emit a log with the outcome at the end suggesting we did
|
||
// receive all the logs.
|
||
waitForOutcomeLoop:
|
||
for {
|
||
batch := testutil.RequireReceive(ctx, t, logsCh)
|
||
|
||
if bytes.Equal(batch.LogSourceId, devcontainerLogSourceID[:]) {
|
||
for _, log := range batch.Logs {
|
||
t.Logf("Received log: %s", log.Output)
|
||
if strings.Contains(log.Output, "\"outcome\"") {
|
||
break waitForOutcomeLoop
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
t.Logf("Checking there's a new container with label: devcontainer.local_folder=%s", workspaceFolder)
|
||
|
||
// Make sure the container exists and isn't the same as the old one.
|
||
testutil.Eventually(ctx, t, func(context.Context) bool {
|
||
resp, err := conn.ListContainers(ctx)
|
||
if err != nil {
|
||
t.Logf("Error listing containers: %v", err)
|
||
return false
|
||
}
|
||
for _, c := range resp.Containers {
|
||
t.Logf("Found container: %s with labels: %v", c.ID[:12], c.Labels)
|
||
if v, ok := c.Labels["devcontainer.local_folder"]; ok && v == workspaceFolder {
|
||
if c.ID == container.ID {
|
||
t.Logf("Found same container: %s", c.ID[:12])
|
||
return false
|
||
}
|
||
t.Logf("Found new container: %s", c.ID[:12])
|
||
container = c
|
||
return true
|
||
}
|
||
}
|
||
return false
|
||
}, testutil.IntervalMedium, "new devcontainer not found")
|
||
defer func(container codersdk.WorkspaceAgentContainer) {
|
||
// We can't rely on pool here because the container is not
|
||
// managed by it (it is managed by @devcontainer/cli).
|
||
err := pool.Client.RemoveContainer(docker.RemoveContainerOptions{
|
||
ID: container.ID,
|
||
RemoveVolumes: true,
|
||
Force: true,
|
||
})
|
||
assert.NoError(t, err, "remove container")
|
||
}(container)
|
||
}
|
||
|
||
func TestAgent_DevcontainersDisabledForSubAgent(t *testing.T) {
|
||
t.Parallel()
|
||
|
||
// Create a manifest with a ParentID to make this a sub agent.
|
||
manifest := agentsdk.Manifest{
|
||
AgentID: uuid.New(),
|
||
ParentID: uuid.New(),
|
||
}
|
||
|
||
// Setup the agent with devcontainers enabled initially.
|
||
//nolint:dogsled
|
||
conn, _, _, _, _ := setupAgent(t, manifest, 0, func(_ *agenttest.Client, o *agent.Options) {
|
||
o.Devcontainers = true
|
||
})
|
||
|
||
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitMedium)
|
||
defer cancel()
|
||
|
||
var err error
|
||
// setupAgent only waits for tailnet reachability, not for the HTTP API
|
||
// listener to serve the expected sub-agent rejection response.
|
||
testutil.Eventually(ctx, t, func(ctx context.Context) bool {
|
||
_, err = conn.ListContainers(ctx)
|
||
if err != nil {
|
||
t.Logf("Error listing containers: %v", err)
|
||
}
|
||
return err != nil && strings.Contains(err.Error(), "Dev Container feature not supported.")
|
||
}, testutil.IntervalFast, "containers endpoint should reject devcontainers inside sub agents")
|
||
require.Error(t, err)
|
||
require.Contains(t, err.Error(), "Dev Container feature not supported.")
|
||
require.Contains(t, err.Error(), "Dev Container integration inside other Dev Containers is explicitly not supported.")
|
||
}
|
||
|
||
// TestAgent_DevcontainerPrebuildClaim tests that we correctly handle
|
||
// the claiming process for running devcontainers.
|
||
//
|
||
// You can run it manually as follows:
|
||
//
|
||
// CODER_TEST_USE_DOCKER=1 go test -count=1 ./agent -run TestAgent_DevcontainerPrebuildClaim
|
||
//
|
||
//nolint:paralleltest // This test sets an environment variable.
|
||
func TestAgent_DevcontainerPrebuildClaim(t *testing.T) {
|
||
if os.Getenv("CODER_TEST_USE_DOCKER") != "1" {
|
||
t.Skip("Set CODER_TEST_USE_DOCKER=1 to run this test")
|
||
}
|
||
if _, err := exec.LookPath("devcontainer"); err != nil {
|
||
t.Skip("This test requires the devcontainer CLI: npm install -g @devcontainers/cli")
|
||
}
|
||
|
||
pool, err := dockertest.NewPool("")
|
||
require.NoError(t, err, "Could not connect to docker")
|
||
|
||
var (
|
||
ctx = testutil.Context(t, testutil.WaitShort)
|
||
|
||
devcontainerID = uuid.New()
|
||
devcontainerLogSourceID = uuid.New()
|
||
|
||
workspaceFolder = filepath.Join(t.TempDir(), "project")
|
||
devcontainerPath = filepath.Join(workspaceFolder, ".devcontainer")
|
||
devcontainerConfig = filepath.Join(devcontainerPath, "devcontainer.json")
|
||
)
|
||
|
||
// Given: A devcontainer project.
|
||
t.Logf("Workspace folder: %s", workspaceFolder)
|
||
|
||
err = os.MkdirAll(devcontainerPath, 0o755)
|
||
require.NoError(t, err, "create dev container directory")
|
||
|
||
// Given: This devcontainer project specifies an app that uses the owner name and workspace name.
|
||
err = os.WriteFile(devcontainerConfig, []byte(`{
|
||
"name": "project",
|
||
"image": "busybox:latest",
|
||
"cmd": ["sleep", "infinity"],
|
||
"runArgs": ["--label=`+agentcontainers.DevcontainerIsTestRunLabel+`=true"],
|
||
"customizations": {
|
||
"coder": {
|
||
"apps": [{
|
||
"slug": "zed",
|
||
"url": "zed://ssh/${localEnv:CODER_WORKSPACE_AGENT_NAME}.${localEnv:CODER_WORKSPACE_NAME}.${localEnv:CODER_WORKSPACE_OWNER_NAME}.coder${containerWorkspaceFolder}"
|
||
}]
|
||
}
|
||
}
|
||
}`), 0o600)
|
||
require.NoError(t, err, "write devcontainer config")
|
||
|
||
// Given: A manifest with a prebuild username and workspace name.
|
||
manifest := agentsdk.Manifest{
|
||
OwnerName: "prebuilds",
|
||
WorkspaceName: "prebuilds-xyz-123",
|
||
|
||
Devcontainers: []codersdk.WorkspaceAgentDevcontainer{
|
||
{ID: devcontainerID, Name: "test", WorkspaceFolder: workspaceFolder},
|
||
},
|
||
Scripts: []codersdk.WorkspaceAgentScript{
|
||
{ID: devcontainerID, LogSourceID: devcontainerLogSourceID},
|
||
},
|
||
}
|
||
|
||
// When: We create an agent with devcontainers enabled.
|
||
//nolint:dogsled
|
||
conn, client, _, _, _ := setupAgent(t, manifest, 0, func(_ *agenttest.Client, o *agent.Options) {
|
||
o.Devcontainers = true
|
||
o.DevcontainerAPIOptions = append(o.DevcontainerAPIOptions,
|
||
agentcontainers.WithContainerLabelIncludeFilter(agentcontainers.DevcontainerLocalFolderLabel, workspaceFolder),
|
||
agentcontainers.WithContainerLabelIncludeFilter(agentcontainers.DevcontainerIsTestRunLabel, "true"),
|
||
)
|
||
})
|
||
|
||
testutil.Eventually(ctx, t, func(ctx context.Context) bool {
|
||
return slices.Contains(client.GetLifecycleStates(), codersdk.WorkspaceAgentLifecycleReady)
|
||
}, testutil.IntervalMedium, "agent not ready")
|
||
|
||
var dcPrebuild codersdk.WorkspaceAgentDevcontainer
|
||
testutil.Eventually(ctx, t, func(ctx context.Context) bool {
|
||
resp, err := conn.ListContainers(ctx)
|
||
require.NoError(t, err)
|
||
|
||
for _, dc := range resp.Devcontainers {
|
||
if dc.Container == nil {
|
||
continue
|
||
}
|
||
|
||
v, ok := dc.Container.Labels[agentcontainers.DevcontainerLocalFolderLabel]
|
||
if ok && v == workspaceFolder {
|
||
dcPrebuild = dc
|
||
return true
|
||
}
|
||
}
|
||
|
||
return false
|
||
}, testutil.IntervalMedium, "devcontainer not found")
|
||
defer func() {
|
||
pool.Client.RemoveContainer(docker.RemoveContainerOptions{
|
||
ID: dcPrebuild.Container.ID,
|
||
RemoveVolumes: true,
|
||
Force: true,
|
||
})
|
||
}()
|
||
|
||
// Then: We expect a sub agent to have been created.
|
||
subAgents := client.GetSubAgents()
|
||
require.Len(t, subAgents, 1)
|
||
|
||
subAgent := subAgents[0]
|
||
subAgentID, err := uuid.FromBytes(subAgent.GetId())
|
||
require.NoError(t, err)
|
||
|
||
// And: We expect there to be 1 app.
|
||
subAgentApps, err := client.GetSubAgentApps(subAgentID)
|
||
require.NoError(t, err)
|
||
require.Len(t, subAgentApps, 1)
|
||
|
||
// And: This app should contain the prebuild workspace name and owner name.
|
||
subAgentApp := subAgentApps[0]
|
||
require.Equal(t, "zed://ssh/project.prebuilds-xyz-123.prebuilds.coder/workspaces/project", subAgentApp.GetUrl())
|
||
|
||
// Given: We close the client and connection
|
||
client.Close()
|
||
conn.Close()
|
||
|
||
// Given: A new manifest with a regular user owner name and workspace name.
|
||
manifest = agentsdk.Manifest{
|
||
OwnerName: "user",
|
||
WorkspaceName: "user-workspace",
|
||
|
||
Devcontainers: []codersdk.WorkspaceAgentDevcontainer{
|
||
{ID: devcontainerID, Name: "test", WorkspaceFolder: workspaceFolder},
|
||
},
|
||
Scripts: []codersdk.WorkspaceAgentScript{
|
||
{ID: devcontainerID, LogSourceID: devcontainerLogSourceID},
|
||
},
|
||
}
|
||
|
||
// When: We create an agent with devcontainers enabled.
|
||
//nolint:dogsled
|
||
conn, client, _, _, _ = setupAgent(t, manifest, 0, func(_ *agenttest.Client, o *agent.Options) {
|
||
o.Devcontainers = true
|
||
o.DevcontainerAPIOptions = append(o.DevcontainerAPIOptions,
|
||
agentcontainers.WithContainerLabelIncludeFilter(agentcontainers.DevcontainerLocalFolderLabel, workspaceFolder),
|
||
agentcontainers.WithContainerLabelIncludeFilter(agentcontainers.DevcontainerIsTestRunLabel, "true"),
|
||
)
|
||
})
|
||
|
||
testutil.Eventually(ctx, t, func(ctx context.Context) bool {
|
||
return slices.Contains(client.GetLifecycleStates(), codersdk.WorkspaceAgentLifecycleReady)
|
||
}, testutil.IntervalMedium, "agent not ready")
|
||
|
||
var dcClaimed codersdk.WorkspaceAgentDevcontainer
|
||
testutil.Eventually(ctx, t, func(ctx context.Context) bool {
|
||
resp, err := conn.ListContainers(ctx)
|
||
require.NoError(t, err)
|
||
|
||
for _, dc := range resp.Devcontainers {
|
||
if dc.Container == nil {
|
||
continue
|
||
}
|
||
|
||
v, ok := dc.Container.Labels[agentcontainers.DevcontainerLocalFolderLabel]
|
||
if ok && v == workspaceFolder {
|
||
dcClaimed = dc
|
||
return true
|
||
}
|
||
}
|
||
|
||
return false
|
||
}, testutil.IntervalMedium, "devcontainer not found")
|
||
defer func() {
|
||
if dcClaimed.Container.ID != dcPrebuild.Container.ID {
|
||
pool.Client.RemoveContainer(docker.RemoveContainerOptions{
|
||
ID: dcClaimed.Container.ID,
|
||
RemoveVolumes: true,
|
||
Force: true,
|
||
})
|
||
}
|
||
}()
|
||
|
||
// Then: We expect the claimed devcontainer and prebuild devcontainer
|
||
// to be using the same underlying container.
|
||
require.Equal(t, dcPrebuild.Container.ID, dcClaimed.Container.ID)
|
||
|
||
// And: We expect there to be a sub agent created.
|
||
subAgents = client.GetSubAgents()
|
||
require.Len(t, subAgents, 1)
|
||
|
||
subAgent = subAgents[0]
|
||
subAgentID, err = uuid.FromBytes(subAgent.GetId())
|
||
require.NoError(t, err)
|
||
|
||
// And: We expect there to be an app.
|
||
subAgentApps, err = client.GetSubAgentApps(subAgentID)
|
||
require.NoError(t, err)
|
||
require.Len(t, subAgentApps, 1)
|
||
|
||
// And: We expect this app to have the user's owner name and workspace name.
|
||
subAgentApp = subAgentApps[0]
|
||
require.Equal(t, "zed://ssh/project.user-workspace.user.coder/workspaces/project", subAgentApp.GetUrl())
|
||
}
|
||
|
||
func TestAgent_Dial(t *testing.T) {
|
||
t.Parallel()
|
||
|
||
cases := []struct {
|
||
name string
|
||
setup func(t testing.TB) net.Listener
|
||
}{
|
||
{
|
||
name: "TCP",
|
||
setup: func(t testing.TB) net.Listener {
|
||
l, err := net.Listen("tcp", "127.0.0.1:0")
|
||
require.NoError(t, err, "create TCP listener")
|
||
return l
|
||
},
|
||
},
|
||
{
|
||
name: "UDP",
|
||
setup: func(t testing.TB) net.Listener {
|
||
addr := net.UDPAddr{
|
||
IP: net.ParseIP("127.0.0.1"),
|
||
Port: 0,
|
||
}
|
||
l, err := udp.Listen("udp", &addr)
|
||
require.NoError(t, err, "create UDP listener")
|
||
return l
|
||
},
|
||
},
|
||
}
|
||
|
||
for _, c := range cases {
|
||
t.Run(c.name, func(t *testing.T) {
|
||
t.Parallel()
|
||
|
||
// The purpose of this test is to ensure that a client can dial a
|
||
// listener in the workspace over tailnet.
|
||
//
|
||
// The OS sometimes drops packets if the system can't keep up with
|
||
// them. For TCP packets, it's typically fine due to
|
||
// retransmissions, but for UDP packets, it can fail this test.
|
||
//
|
||
// The OS gets involved for the Wireguard traffic (either via DERP
|
||
// or direct UDP), and also for the traffic between the agent and
|
||
// the listener in the "workspace".
|
||
//
|
||
// To avoid this, we'll retry this test up to 3 times.
|
||
//nolint:gocritic // This test is flaky due to uncontrollable OS packet drops under heavy load.
|
||
testutil.RunRetry(t, 3, func(t testing.TB) {
|
||
ctx := testutil.Context(t, testutil.WaitLong)
|
||
|
||
l := c.setup(t)
|
||
done := make(chan struct{})
|
||
defer func() {
|
||
l.Close()
|
||
<-done
|
||
}()
|
||
|
||
go func() {
|
||
defer close(done)
|
||
for range 2 {
|
||
c, err := l.Accept()
|
||
if assert.NoError(t, err, "accept connection") {
|
||
testAccept(ctx, t, c)
|
||
_ = c.Close()
|
||
}
|
||
}
|
||
}()
|
||
|
||
agentID := uuid.UUID{0, 0, 0, 0, 0, 1, 2, 3, 4, 5, 6, 7, 8}
|
||
//nolint:dogsled
|
||
agentConn, _, _, _, _ := setupAgent(t, agentsdk.Manifest{
|
||
AgentID: agentID,
|
||
}, 0)
|
||
require.True(t, agentConn.AwaitReachable(ctx))
|
||
conn, err := agentConn.DialContext(ctx, l.Addr().Network(), l.Addr().String())
|
||
require.NoError(t, err)
|
||
testDial(ctx, t, conn)
|
||
err = conn.Close()
|
||
require.NoError(t, err)
|
||
|
||
// also connect via the CoderServicePrefix, to test that we can reach the agent on this
|
||
// IP. This will be required for CoderVPN.
|
||
_, rawPort, _ := net.SplitHostPort(l.Addr().String())
|
||
port, _ := strconv.ParseUint(rawPort, 10, 16)
|
||
ipp := netip.AddrPortFrom(tailnet.CoderServicePrefix.AddrFromUUID(agentID), uint16(port))
|
||
|
||
switch l.Addr().Network() {
|
||
case "tcp":
|
||
conn, err = agentConn.TailnetConn().DialContextTCP(ctx, ipp)
|
||
case "udp":
|
||
conn, err = agentConn.TailnetConn().DialContextUDP(ctx, ipp)
|
||
default:
|
||
t.Fatalf("unknown network: %s", l.Addr().Network())
|
||
}
|
||
require.NoError(t, err)
|
||
testDial(ctx, t, conn)
|
||
err = conn.Close()
|
||
require.NoError(t, err)
|
||
})
|
||
})
|
||
}
|
||
}
|
||
|
||
// TestAgent_UpdatedDERP checks that agents can handle their DERP map being
|
||
// updated, and that clients can also handle it.
|
||
func TestAgent_UpdatedDERP(t *testing.T) {
|
||
t.Parallel()
|
||
|
||
logger := testutil.Logger(t)
|
||
|
||
originalDerpMap, _ := tailnettest.RunDERPAndSTUN(t)
|
||
require.NotNil(t, originalDerpMap)
|
||
|
||
coordinator := tailnet.NewCoordinator(logger)
|
||
// use t.Cleanup so the coordinator closing doesn't deadlock with in-memory
|
||
// coordination
|
||
t.Cleanup(func() {
|
||
_ = coordinator.Close()
|
||
})
|
||
agentID := uuid.New()
|
||
statsCh := make(chan *proto.Stats, 50)
|
||
fs := afero.NewMemMapFs()
|
||
client := agenttest.NewClient(t,
|
||
logger.Named("agent"),
|
||
agentID,
|
||
agentsdk.Manifest{
|
||
DERPMap: originalDerpMap,
|
||
// Force DERP.
|
||
DisableDirectConnections: true,
|
||
},
|
||
statsCh,
|
||
coordinator,
|
||
)
|
||
t.Cleanup(func() {
|
||
t.Log("closing client")
|
||
client.Close()
|
||
})
|
||
uut := agent.New(agent.Options{
|
||
Client: client,
|
||
Filesystem: fs,
|
||
Logger: logger.Named("agent"),
|
||
ReconnectingPTYTimeout: time.Minute,
|
||
})
|
||
t.Cleanup(func() {
|
||
t.Log("closing agent")
|
||
_ = uut.Close()
|
||
})
|
||
|
||
// Setup a client connection.
|
||
newClientConn := func(derpMap *tailcfg.DERPMap, name string) workspacesdk.AgentConn {
|
||
conn, err := tailnet.NewConn(&tailnet.Options{
|
||
Addresses: []netip.Prefix{tailnet.TailscaleServicePrefix.RandomPrefix()},
|
||
DERPMap: derpMap,
|
||
Logger: logger.Named(name),
|
||
})
|
||
require.NoError(t, err)
|
||
t.Cleanup(func() {
|
||
t.Logf("closing conn %s", name)
|
||
_ = conn.Close()
|
||
})
|
||
testCtx, testCtxCancel := context.WithCancel(context.Background())
|
||
t.Cleanup(testCtxCancel)
|
||
clientID := uuid.New()
|
||
ctrl := tailnet.NewTunnelSrcCoordController(logger, conn)
|
||
ctrl.AddDestination(agentID)
|
||
auth := tailnet.ClientCoordinateeAuth{AgentID: agentID}
|
||
coordination := ctrl.New(tailnet.NewInMemoryCoordinatorClient(logger, clientID, auth, coordinator))
|
||
t.Cleanup(func() {
|
||
t.Logf("closing coordination %s", name)
|
||
cctx, ccancel := context.WithTimeout(testCtx, testutil.WaitShort)
|
||
defer ccancel()
|
||
err := coordination.Close(cctx)
|
||
if err != nil {
|
||
t.Logf("error closing in-memory coordination: %s", err.Error())
|
||
}
|
||
t.Logf("closed coordination %s", name)
|
||
})
|
||
// Force DERP.
|
||
conn.SetBlockEndpoints(true)
|
||
|
||
sdkConn := workspacesdk.NewAgentConn(conn, workspacesdk.AgentConnOptions{
|
||
AgentID: agentID,
|
||
CloseFunc: func() error { return workspacesdk.ErrSkipClose },
|
||
})
|
||
t.Cleanup(func() {
|
||
t.Logf("closing sdkConn %s", name)
|
||
_ = sdkConn.Close()
|
||
})
|
||
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
|
||
defer cancel()
|
||
if !sdkConn.AwaitReachable(ctx) {
|
||
t.Fatal("agent not reachable")
|
||
}
|
||
|
||
return sdkConn
|
||
}
|
||
conn1 := newClientConn(originalDerpMap, "client1")
|
||
|
||
// Change the DERP map.
|
||
newDerpMap, _ := tailnettest.RunDERPAndSTUN(t)
|
||
require.NotNil(t, newDerpMap)
|
||
|
||
// Change the region ID.
|
||
newDerpMap.Regions[2] = newDerpMap.Regions[1]
|
||
delete(newDerpMap.Regions, 1)
|
||
newDerpMap.Regions[2].RegionID = 2
|
||
for _, node := range newDerpMap.Regions[2].Nodes {
|
||
node.RegionID = 2
|
||
}
|
||
|
||
// Push a new DERP map to the agent.
|
||
err := client.PushDERPMapUpdate(newDerpMap)
|
||
require.NoError(t, err)
|
||
t.Log("pushed DERPMap update to agent")
|
||
|
||
require.Eventually(t, func() bool {
|
||
conn := uut.TailnetConn()
|
||
if conn == nil {
|
||
return false
|
||
}
|
||
regionIDs := conn.DERPMap().RegionIDs()
|
||
preferredDERP := conn.Node().PreferredDERP
|
||
t.Logf("agent Conn DERPMap with regionIDs %v, PreferredDERP %d", regionIDs, preferredDERP)
|
||
return len(regionIDs) == 1 && regionIDs[0] == 2 && preferredDERP == 2
|
||
}, testutil.WaitLong, testutil.IntervalFast)
|
||
t.Log("agent got the new DERPMap")
|
||
|
||
// Connect from a second client and make sure it uses the new DERP map.
|
||
conn2 := newClientConn(newDerpMap, "client2")
|
||
require.Equal(t, []int{2}, conn2.TailnetConn().DERPMap().RegionIDs())
|
||
t.Log("conn2 got the new DERPMap")
|
||
|
||
// If the first client gets a DERP map update, it should be able to
|
||
// reconnect just fine.
|
||
conn1.TailnetConn().SetDERPMap(newDerpMap)
|
||
require.Equal(t, []int{2}, conn1.TailnetConn().DERPMap().RegionIDs())
|
||
t.Log("set the new DERPMap on conn1")
|
||
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
|
||
defer cancel()
|
||
require.True(t, conn1.AwaitReachable(ctx))
|
||
t.Log("conn1 reached agent with new DERP")
|
||
}
|
||
|
||
func TestAgent_Speedtest(t *testing.T) {
|
||
t.Parallel()
|
||
t.Skip("This test is relatively flakey because of Tailscale's speedtest code...")
|
||
logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}).Leveled(slog.LevelDebug)
|
||
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
|
||
defer cancel()
|
||
derpMap, _ := tailnettest.RunDERPAndSTUN(t)
|
||
//nolint:dogsled
|
||
conn, _, _, _, _ := setupAgent(t, agentsdk.Manifest{
|
||
DERPMap: derpMap,
|
||
}, 0, func(client *agenttest.Client, options *agent.Options) {
|
||
options.Logger = logger.Named("agent")
|
||
})
|
||
defer conn.Close()
|
||
res, err := conn.Speedtest(ctx, speedtest.Upload, 250*time.Millisecond)
|
||
require.NoError(t, err)
|
||
t.Logf("%.2f MBits/s", res[len(res)-1].MBitsPerSecond())
|
||
}
|
||
|
||
func TestAgent_Reconnect(t *testing.T) {
|
||
t.Parallel()
|
||
ctx := testutil.Context(t, testutil.WaitLong)
|
||
logger := testutil.Logger(t)
|
||
// After the agent is disconnected from a coordinator, it's supposed
|
||
// to reconnect!
|
||
fCoordinator := tailnettest.NewFakeCoordinator()
|
||
|
||
agentID := uuid.New()
|
||
statsCh := make(chan *proto.Stats, 50)
|
||
derpMap, _ := tailnettest.RunDERPAndSTUN(t)
|
||
client := agenttest.NewClient(t,
|
||
logger,
|
||
agentID,
|
||
agentsdk.Manifest{
|
||
DERPMap: derpMap,
|
||
Directory: "/test/workspace",
|
||
},
|
||
statsCh,
|
||
fCoordinator,
|
||
)
|
||
defer client.Close()
|
||
|
||
closer := agent.New(agent.Options{
|
||
Client: client,
|
||
Logger: logger.Named("agent"),
|
||
})
|
||
defer closer.Close()
|
||
|
||
// Each iteration forces the agent to reconnect by closing
|
||
// the current coordinate call while the tracked HTTP server
|
||
// goroutine (from connection 1's createTailnet) is still
|
||
// alive, widening the race window.
|
||
const reconnections = 5
|
||
for i := range reconnections {
|
||
call := testutil.RequireReceive(ctx, t, fCoordinator.CoordinateCalls)
|
||
require.Equal(t, i+1, client.GetNumRefreshTokenCalls())
|
||
close(call.Resps) // hang up — triggers reconnect
|
||
}
|
||
// Verify final reconnect succeeds.
|
||
testutil.RequireReceive(ctx, t, fCoordinator.CoordinateCalls)
|
||
require.Equal(t, reconnections+1, client.GetNumRefreshTokenCalls())
|
||
closer.Close()
|
||
}
|
||
|
||
func TestAgent_ReconnectNoLifecycleReemit(t *testing.T) {
|
||
t.Parallel()
|
||
ctx := testutil.Context(t, testutil.WaitLong)
|
||
logger := testutil.Logger(t)
|
||
|
||
fCoordinator := tailnettest.NewFakeCoordinator()
|
||
agentID := uuid.New()
|
||
statsCh := make(chan *proto.Stats, 50)
|
||
derpMap, _ := tailnettest.RunDERPAndSTUN(t)
|
||
|
||
client := agenttest.NewClient(t,
|
||
logger,
|
||
agentID,
|
||
agentsdk.Manifest{
|
||
DERPMap: derpMap,
|
||
Scripts: []codersdk.WorkspaceAgentScript{{
|
||
Script: "echo hello",
|
||
Timeout: 30 * time.Second,
|
||
RunOnStart: true,
|
||
}},
|
||
},
|
||
statsCh,
|
||
fCoordinator,
|
||
)
|
||
defer client.Close()
|
||
|
||
closer := agent.New(agent.Options{
|
||
Client: client,
|
||
Logger: logger.Named("agent"),
|
||
})
|
||
defer closer.Close()
|
||
|
||
// Wait for the agent to reach Ready state.
|
||
require.Eventually(t, func() bool {
|
||
return slices.Contains(client.GetLifecycleStates(), codersdk.WorkspaceAgentLifecycleReady)
|
||
}, testutil.WaitShort, testutil.IntervalFast)
|
||
|
||
statesBefore := slices.Clone(client.GetLifecycleStates())
|
||
|
||
// Disconnect by closing the coordinator response channel.
|
||
call1 := testutil.RequireReceive(ctx, t, fCoordinator.CoordinateCalls)
|
||
close(call1.Resps)
|
||
|
||
// Wait for reconnect.
|
||
testutil.RequireReceive(ctx, t, fCoordinator.CoordinateCalls)
|
||
|
||
// Wait for a stats report as a deterministic steady-state proof.
|
||
testutil.RequireReceive(ctx, t, statsCh)
|
||
|
||
statesAfter := client.GetLifecycleStates()
|
||
require.Equal(t, statesBefore, statesAfter,
|
||
"lifecycle states should not be re-reported after reconnect")
|
||
|
||
closer.Close()
|
||
}
|
||
|
||
func TestAgent_WriteVSCodeConfigs(t *testing.T) {
|
||
t.Parallel()
|
||
logger := testutil.Logger(t)
|
||
coordinator := tailnet.NewCoordinator(logger)
|
||
defer coordinator.Close()
|
||
|
||
client := agenttest.NewClient(t,
|
||
logger,
|
||
uuid.New(),
|
||
agentsdk.Manifest{
|
||
GitAuthConfigs: 1,
|
||
DERPMap: &tailcfg.DERPMap{},
|
||
},
|
||
make(chan *proto.Stats, 50),
|
||
coordinator,
|
||
)
|
||
defer client.Close()
|
||
filesystem := afero.NewMemMapFs()
|
||
closer := agent.New(agent.Options{
|
||
Client: client,
|
||
Logger: logger.Named("agent"),
|
||
Filesystem: filesystem,
|
||
})
|
||
defer closer.Close()
|
||
|
||
home, err := os.UserHomeDir()
|
||
require.NoError(t, err)
|
||
name := filepath.Join(home, ".vscode-server", "data", "Machine", "settings.json")
|
||
require.Eventually(t, func() bool {
|
||
_, err := filesystem.Stat(name)
|
||
return err == nil
|
||
}, testutil.WaitShort, testutil.IntervalFast)
|
||
}
|
||
|
||
func TestAgent_DebugServer(t *testing.T) {
|
||
t.Parallel()
|
||
|
||
logDir := t.TempDir()
|
||
logPath := filepath.Join(logDir, "coder-agent.log")
|
||
randLogStr, err := cryptorand.String(32)
|
||
require.NoError(t, err)
|
||
require.NoError(t, os.WriteFile(logPath, []byte(randLogStr), 0o600))
|
||
newRotatedLogPath := filepath.Join(logDir, "coder-agent-2026-05-17T20-00-00.000.log")
|
||
oldRotatedLogPath := filepath.Join(logDir, "coder-agent-2026-05-17T19-00-00.000.log")
|
||
require.NoError(t, os.WriteFile(newRotatedLogPath, []byte("new rotated log"), 0o600))
|
||
require.NoError(t, os.WriteFile(oldRotatedLogPath, []byte("old rotated log"), 0o600))
|
||
now := time.Now()
|
||
newRotatedModTime := now.Add(-time.Minute)
|
||
oldRotatedModTime := now.Add(-48 * time.Hour)
|
||
require.NoError(t, os.Chtimes(newRotatedLogPath, newRotatedModTime, newRotatedModTime))
|
||
require.NoError(t, os.Chtimes(oldRotatedLogPath, oldRotatedModTime, oldRotatedModTime))
|
||
derpMap, _ := tailnettest.RunDERPAndSTUN(t)
|
||
//nolint:dogsled
|
||
conn, _, _, _, agnt := setupAgentWithSecrets(t, agentsdk.Manifest{
|
||
DERPMap: derpMap,
|
||
EnvironmentVariables: map[string]string{
|
||
"AWS_SECRET_ACCESS_KEY": "env-value-should-be-redacted-67890",
|
||
"EMPTY_VAR": "",
|
||
},
|
||
}, []agentsdk.WorkspaceSecret{
|
||
{EnvName: "DEBUG_SECRET", Value: []byte("super-secret-value-12345")},
|
||
}, 0, func(c *agenttest.Client, o *agent.Options) {
|
||
o.LogDir = logDir
|
||
})
|
||
|
||
awaitReachableCtx := testutil.Context(t, testutil.WaitLong)
|
||
ok := conn.AwaitReachable(awaitReachableCtx)
|
||
require.True(t, ok)
|
||
_ = conn.Close()
|
||
|
||
srv := httptest.NewServer(agnt.HTTPDebug())
|
||
t.Cleanup(srv.Close)
|
||
|
||
t.Run("MagicsockDebug", func(t *testing.T) {
|
||
t.Parallel()
|
||
|
||
ctx := testutil.Context(t, testutil.WaitLong)
|
||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, srv.URL+"/debug/magicsock", nil)
|
||
require.NoError(t, err)
|
||
|
||
res, err := srv.Client().Do(req)
|
||
require.NoError(t, err)
|
||
defer res.Body.Close()
|
||
require.Equal(t, http.StatusOK, res.StatusCode)
|
||
|
||
resBody, err := io.ReadAll(res.Body)
|
||
require.NoError(t, err)
|
||
require.Contains(t, string(resBody), "<h1>magicsock</h1>")
|
||
})
|
||
|
||
t.Run("MagicsockDebugLogging", func(t *testing.T) {
|
||
t.Parallel()
|
||
|
||
t.Run("Enable", func(t *testing.T) {
|
||
t.Parallel()
|
||
|
||
ctx := testutil.Context(t, testutil.WaitLong)
|
||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, srv.URL+"/debug/magicsock/debug-logging/t", nil)
|
||
require.NoError(t, err)
|
||
|
||
res, err := srv.Client().Do(req)
|
||
require.NoError(t, err)
|
||
defer res.Body.Close()
|
||
require.Equal(t, http.StatusOK, res.StatusCode)
|
||
|
||
resBody, err := io.ReadAll(res.Body)
|
||
require.NoError(t, err)
|
||
require.Contains(t, string(resBody), "updated magicsock debug logging to true")
|
||
})
|
||
|
||
t.Run("Disable", func(t *testing.T) {
|
||
t.Parallel()
|
||
|
||
ctx := testutil.Context(t, testutil.WaitLong)
|
||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, srv.URL+"/debug/magicsock/debug-logging/0", nil)
|
||
require.NoError(t, err)
|
||
|
||
res, err := srv.Client().Do(req)
|
||
require.NoError(t, err)
|
||
defer res.Body.Close()
|
||
require.Equal(t, http.StatusOK, res.StatusCode)
|
||
|
||
resBody, err := io.ReadAll(res.Body)
|
||
require.NoError(t, err)
|
||
require.Contains(t, string(resBody), "updated magicsock debug logging to false")
|
||
})
|
||
|
||
t.Run("Invalid", func(t *testing.T) {
|
||
t.Parallel()
|
||
|
||
ctx := testutil.Context(t, testutil.WaitLong)
|
||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, srv.URL+"/debug/magicsock/debug-logging/blah", nil)
|
||
require.NoError(t, err)
|
||
|
||
res, err := srv.Client().Do(req)
|
||
require.NoError(t, err)
|
||
defer res.Body.Close()
|
||
require.Equal(t, http.StatusBadRequest, res.StatusCode)
|
||
|
||
resBody, err := io.ReadAll(res.Body)
|
||
require.NoError(t, err)
|
||
require.Contains(t, string(resBody), `invalid state "blah", must be a boolean`)
|
||
})
|
||
})
|
||
|
||
t.Run("Manifest", func(t *testing.T) {
|
||
t.Parallel()
|
||
|
||
ctx := testutil.Context(t, testutil.WaitLong)
|
||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, srv.URL+"/debug/manifest", nil)
|
||
require.NoError(t, err)
|
||
|
||
res, err := srv.Client().Do(req)
|
||
require.NoError(t, err)
|
||
defer res.Body.Close()
|
||
require.Equal(t, http.StatusOK, res.StatusCode)
|
||
|
||
var v agentsdk.Manifest
|
||
require.NoError(t, json.NewDecoder(res.Body).Decode(&v))
|
||
require.NotNil(t, v)
|
||
})
|
||
|
||
t.Run("ManifestSecretsStripped", func(t *testing.T) {
|
||
t.Parallel()
|
||
|
||
ctx := testutil.Context(t, testutil.WaitLong)
|
||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, srv.URL+"/debug/manifest", nil)
|
||
require.NoError(t, err)
|
||
|
||
res, err := srv.Client().Do(req)
|
||
require.NoError(t, err)
|
||
defer res.Body.Close()
|
||
require.Equal(t, http.StatusOK, res.StatusCode)
|
||
|
||
body, err := io.ReadAll(res.Body)
|
||
require.NoError(t, err)
|
||
|
||
// The response must not contain the secret value.
|
||
require.NotContains(t, string(body), "super-secret-value-12345")
|
||
|
||
// Confirm we can decode as a Manifest. The SDK type
|
||
// intentionally has no Secrets field, so there is nothing
|
||
// to leak through JSON encoding.
|
||
var v agentsdk.Manifest
|
||
require.NoError(t, json.Unmarshal(body, &v))
|
||
})
|
||
|
||
t.Run("ManifestEnvVarValuesRedacted", func(t *testing.T) {
|
||
t.Parallel()
|
||
|
||
ctx := testutil.Context(t, testutil.WaitLong)
|
||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, srv.URL+"/debug/manifest", nil)
|
||
require.NoError(t, err)
|
||
|
||
res, err := srv.Client().Do(req)
|
||
require.NoError(t, err)
|
||
defer res.Body.Close()
|
||
require.Equal(t, http.StatusOK, res.StatusCode)
|
||
|
||
body, err := io.ReadAll(res.Body)
|
||
require.NoError(t, err)
|
||
|
||
require.NotContains(t, string(body), "env-value-should-be-redacted-67890")
|
||
|
||
var v agentsdk.Manifest
|
||
require.NoError(t, json.Unmarshal(body, &v))
|
||
|
||
require.Contains(t, v.EnvironmentVariables, "AWS_SECRET_ACCESS_KEY")
|
||
require.Equal(t, "***REDACTED***", v.EnvironmentVariables["AWS_SECRET_ACCESS_KEY"])
|
||
|
||
// Empty values carry no secret and are preserved as empty.
|
||
require.Contains(t, v.EnvironmentVariables, "EMPTY_VAR")
|
||
require.Equal(t, "", v.EnvironmentVariables["EMPTY_VAR"])
|
||
})
|
||
|
||
t.Run("Logs", func(t *testing.T) {
|
||
t.Parallel()
|
||
|
||
ctx := testutil.Context(t, testutil.WaitLong)
|
||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, srv.URL+"/debug/logs", nil)
|
||
require.NoError(t, err)
|
||
|
||
res, err := srv.Client().Do(req)
|
||
require.NoError(t, err)
|
||
require.Equal(t, http.StatusOK, res.StatusCode)
|
||
defer res.Body.Close()
|
||
resBody, err := io.ReadAll(res.Body)
|
||
require.NoError(t, err)
|
||
require.NotEmpty(t, string(resBody))
|
||
require.Contains(t, string(resBody), randLogStr)
|
||
require.NotContains(t, string(resBody), "new rotated log")
|
||
})
|
||
|
||
t.Run("LogsIncludeActiveOnlyWithAfter", func(t *testing.T) {
|
||
t.Parallel()
|
||
|
||
ctx := testutil.Context(t, testutil.WaitLong)
|
||
url := srv.URL + "/debug/logs?after=" + newRotatedModTime.Add(time.Minute).UTC().Format(time.RFC3339Nano)
|
||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
|
||
require.NoError(t, err)
|
||
|
||
res, err := srv.Client().Do(req)
|
||
require.NoError(t, err)
|
||
require.Equal(t, http.StatusOK, res.StatusCode)
|
||
defer res.Body.Close()
|
||
resBody, err := io.ReadAll(res.Body)
|
||
require.NoError(t, err)
|
||
body := string(resBody)
|
||
require.Contains(t, body, randLogStr)
|
||
require.Contains(t, body, "coder-agent.log")
|
||
require.NotContains(t, body, "new rotated log")
|
||
require.NotContains(t, body, "old rotated log")
|
||
})
|
||
|
||
t.Run("LogsIncludeRotatedWithAfter", func(t *testing.T) {
|
||
t.Parallel()
|
||
|
||
ctx := testutil.Context(t, testutil.WaitLong)
|
||
url := srv.URL + "/debug/logs?after=" + newRotatedModTime.Add(-time.Minute).UTC().Format(time.RFC3339Nano)
|
||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
|
||
require.NoError(t, err)
|
||
|
||
res, err := srv.Client().Do(req)
|
||
require.NoError(t, err)
|
||
require.Equal(t, http.StatusOK, res.StatusCode)
|
||
defer res.Body.Close()
|
||
resBody, err := io.ReadAll(res.Body)
|
||
require.NoError(t, err)
|
||
body := string(resBody)
|
||
require.Contains(t, body, randLogStr)
|
||
require.Contains(t, body, "coder-agent.log")
|
||
require.Contains(t, body, "coder-agent-2026-05-17T20-00-00.000.log")
|
||
require.Contains(t, body, "new rotated log")
|
||
require.NotContains(t, body, "old rotated log")
|
||
})
|
||
|
||
t.Run("LogsIncludeRotatedWithOlderAfter", func(t *testing.T) {
|
||
t.Parallel()
|
||
|
||
ctx := testutil.Context(t, testutil.WaitLong)
|
||
url := srv.URL + "/debug/logs?after=" + oldRotatedModTime.Add(-time.Minute).UTC().Format(time.RFC3339Nano)
|
||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
|
||
require.NoError(t, err)
|
||
|
||
res, err := srv.Client().Do(req)
|
||
require.NoError(t, err)
|
||
require.Equal(t, http.StatusOK, res.StatusCode)
|
||
defer res.Body.Close()
|
||
resBody, err := io.ReadAll(res.Body)
|
||
require.NoError(t, err)
|
||
body := string(resBody)
|
||
require.Contains(t, body, randLogStr)
|
||
require.Contains(t, body, "new rotated log")
|
||
require.Contains(t, body, "old rotated log")
|
||
require.Less(t, strings.Index(body, randLogStr), strings.Index(body, "new rotated log"))
|
||
require.Less(t, strings.Index(body, "new rotated log"), strings.Index(body, "old rotated log"))
|
||
})
|
||
|
||
t.Run("LogsInvalidAfter", func(t *testing.T) {
|
||
t.Parallel()
|
||
|
||
ctx := testutil.Context(t, testutil.WaitLong)
|
||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, srv.URL+"/debug/logs?after=nope", nil)
|
||
require.NoError(t, err)
|
||
|
||
res, err := srv.Client().Do(req)
|
||
require.NoError(t, err)
|
||
defer res.Body.Close()
|
||
require.Equal(t, http.StatusBadRequest, res.StatusCode)
|
||
})
|
||
}
|
||
|
||
func TestAgent_ScriptLogging(t *testing.T) {
|
||
if runtime.GOOS == "windows" {
|
||
t.Skip("bash scripts only")
|
||
}
|
||
t.Parallel()
|
||
ctx := testutil.Context(t, testutil.WaitMedium)
|
||
|
||
derpMap, _ := tailnettest.RunDERPAndSTUN(t)
|
||
logsCh := make(chan *proto.BatchCreateLogsRequest, 100)
|
||
lsStart := uuid.UUID{0x11}
|
||
lsStop := uuid.UUID{0x22}
|
||
//nolint:dogsled
|
||
_, _, _, _, agnt := setupAgent(
|
||
t,
|
||
agentsdk.Manifest{
|
||
DERPMap: derpMap,
|
||
Scripts: []codersdk.WorkspaceAgentScript{
|
||
{
|
||
LogSourceID: lsStart,
|
||
RunOnStart: true,
|
||
Script: `#!/bin/sh
|
||
i=0
|
||
while [ $i -ne 5 ]
|
||
do
|
||
i=$(($i+1))
|
||
echo "start $i"
|
||
done
|
||
`,
|
||
},
|
||
{
|
||
LogSourceID: lsStop,
|
||
RunOnStop: true,
|
||
Script: `#!/bin/sh
|
||
i=0
|
||
while [ $i -ne 3000 ]
|
||
do
|
||
i=$(($i+1))
|
||
echo "stop $i"
|
||
done
|
||
`, // send a lot of stop logs to make sure we don't truncate shutdown logs before closing the API conn
|
||
},
|
||
},
|
||
},
|
||
0,
|
||
func(cl *agenttest.Client, _ *agent.Options) {
|
||
cl.SetLogsChannel(logsCh)
|
||
},
|
||
)
|
||
|
||
n := 1
|
||
for n <= 5 {
|
||
logs := testutil.TryReceive(ctx, t, logsCh)
|
||
require.NotNil(t, logs)
|
||
for _, l := range logs.GetLogs() {
|
||
require.Equal(t, fmt.Sprintf("start %d", n), l.GetOutput())
|
||
n++
|
||
}
|
||
}
|
||
|
||
err := agnt.Close()
|
||
require.NoError(t, err)
|
||
|
||
n = 1
|
||
for n <= 3000 {
|
||
logs := testutil.TryReceive(ctx, t, logsCh)
|
||
require.NotNil(t, logs)
|
||
for _, l := range logs.GetLogs() {
|
||
require.Equal(t, fmt.Sprintf("stop %d", n), l.GetOutput())
|
||
n++
|
||
}
|
||
t.Logf("got %d stop logs", n-1)
|
||
}
|
||
}
|
||
|
||
// setupAgentSSHClient creates an agent, dials it, and sets up an ssh.Client for it
|
||
func setupAgentSSHClient(ctx context.Context, t *testing.T) *ssh.Client {
|
||
//nolint: dogsled
|
||
agentConn, _, _, _, _ := setupAgent(t, agentsdk.Manifest{}, 0)
|
||
sshClient, err := agentConn.SSHClient(ctx)
|
||
require.NoError(t, err)
|
||
t.Cleanup(func() { sshClient.Close() })
|
||
return sshClient
|
||
}
|
||
|
||
func setupSSHSession(
|
||
t *testing.T,
|
||
manifest agentsdk.Manifest,
|
||
banner codersdk.BannerConfig,
|
||
prepareFS func(fs afero.Fs),
|
||
opts ...func(*agenttest.Client, *agent.Options),
|
||
) *ssh.Session {
|
||
return setupSSHSessionOnPort(t, manifest, banner, prepareFS, workspacesdk.AgentSSHPort, opts...)
|
||
}
|
||
|
||
func setupSSHSessionOnPort(
|
||
t *testing.T,
|
||
manifest agentsdk.Manifest,
|
||
banner codersdk.BannerConfig,
|
||
prepareFS func(fs afero.Fs),
|
||
port uint16,
|
||
opts ...func(*agenttest.Client, *agent.Options),
|
||
) *ssh.Session {
|
||
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
|
||
defer cancel()
|
||
opts = append(opts, func(c *agenttest.Client, o *agent.Options) {
|
||
c.SetAnnouncementBannersFunc(func() ([]codersdk.BannerConfig, error) {
|
||
return []codersdk.BannerConfig{banner}, nil
|
||
})
|
||
})
|
||
//nolint:dogsled
|
||
conn, _, _, fs, _ := setupAgent(t, manifest, 0, opts...)
|
||
if prepareFS != nil {
|
||
prepareFS(fs)
|
||
}
|
||
sshClient, err := conn.SSHClientOnPort(ctx, port)
|
||
require.NoError(t, err)
|
||
t.Cleanup(func() {
|
||
_ = sshClient.Close()
|
||
})
|
||
session, err := sshClient.NewSession()
|
||
require.NoError(t, err)
|
||
t.Cleanup(func() {
|
||
_ = session.Close()
|
||
})
|
||
return session
|
||
}
|
||
|
||
func setupAgent(t testing.TB, metadata agentsdk.Manifest, ptyTimeout time.Duration, opts ...func(*agenttest.Client, *agent.Options)) (
|
||
workspacesdk.AgentConn,
|
||
*agenttest.Client,
|
||
<-chan *proto.Stats,
|
||
afero.Fs,
|
||
agent.Agent,
|
||
) {
|
||
return setupAgentWithSecrets(t, metadata, nil, ptyTimeout, opts...)
|
||
}
|
||
|
||
// setupAgentWithSecrets is like setupAgent but also injects user
|
||
// secrets into the agent's proto manifest. Separate from setupAgent
|
||
// because agentsdk.Manifest intentionally does not carry secrets; see
|
||
// the Manifest doc comment in codersdk/agentsdk.
|
||
func setupAgentWithSecrets(t testing.TB, metadata agentsdk.Manifest, secrets []agentsdk.WorkspaceSecret, ptyTimeout time.Duration, opts ...func(*agenttest.Client, *agent.Options)) (
|
||
workspacesdk.AgentConn,
|
||
*agenttest.Client,
|
||
<-chan *proto.Stats,
|
||
afero.Fs,
|
||
agent.Agent,
|
||
) {
|
||
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)
|
||
if metadata.DERPMap == nil {
|
||
metadata.DERPMap, _ = tailnettest.RunDERPAndSTUN(t)
|
||
}
|
||
if metadata.AgentID == uuid.Nil {
|
||
metadata.AgentID = uuid.New()
|
||
}
|
||
if metadata.AgentName == "" {
|
||
metadata.AgentName = "test-agent"
|
||
}
|
||
if metadata.WorkspaceName == "" {
|
||
metadata.WorkspaceName = "test-workspace"
|
||
}
|
||
if metadata.OwnerName == "" {
|
||
metadata.OwnerName = "test-user"
|
||
}
|
||
if metadata.WorkspaceID == uuid.Nil {
|
||
metadata.WorkspaceID = uuid.New()
|
||
}
|
||
coordinator := tailnet.NewCoordinator(logger)
|
||
t.Cleanup(func() {
|
||
_ = coordinator.Close()
|
||
})
|
||
statsCh := make(chan *proto.Stats, 50)
|
||
fs := afero.NewMemMapFs()
|
||
c := agenttest.NewClientWithSecrets(t, logger.Named("agenttest"), metadata.AgentID, metadata, secrets, statsCh, coordinator)
|
||
t.Cleanup(c.Close)
|
||
|
||
options := agent.Options{
|
||
Client: c,
|
||
Filesystem: fs,
|
||
Logger: logger.Named("agent"),
|
||
ReconnectingPTYTimeout: ptyTimeout,
|
||
EnvironmentVariables: map[string]string{},
|
||
StatsReportInterval: agenttest.StatsInterval,
|
||
}
|
||
|
||
for _, opt := range opts {
|
||
opt(c, &options)
|
||
}
|
||
|
||
agnt := agent.New(options)
|
||
t.Cleanup(func() {
|
||
_ = agnt.Close()
|
||
})
|
||
conn, err := tailnet.NewConn(&tailnet.Options{
|
||
Addresses: []netip.Prefix{netip.PrefixFrom(tailnet.TailscaleServicePrefix.RandomAddr(), 128)},
|
||
DERPMap: metadata.DERPMap,
|
||
Logger: logger.Named("client"),
|
||
})
|
||
require.NoError(t, err)
|
||
t.Cleanup(func() {
|
||
_ = conn.Close()
|
||
})
|
||
testCtx, testCtxCancel := context.WithCancel(context.Background())
|
||
t.Cleanup(testCtxCancel)
|
||
clientID := uuid.New()
|
||
ctrl := tailnet.NewTunnelSrcCoordController(logger, conn)
|
||
ctrl.AddDestination(metadata.AgentID)
|
||
auth := tailnet.ClientCoordinateeAuth{AgentID: metadata.AgentID}
|
||
coordination := ctrl.New(tailnet.NewInMemoryCoordinatorClient(
|
||
logger, clientID, auth, coordinator))
|
||
t.Cleanup(func() {
|
||
cctx, ccancel := context.WithTimeout(testCtx, testutil.WaitShort)
|
||
defer ccancel()
|
||
err := coordination.Close(cctx)
|
||
if err != nil {
|
||
t.Logf("error closing in-mem coordination: %s", err.Error())
|
||
}
|
||
})
|
||
agentConn := workspacesdk.NewAgentConn(conn, workspacesdk.AgentConnOptions{
|
||
AgentID: metadata.AgentID,
|
||
})
|
||
t.Cleanup(func() {
|
||
_ = agentConn.Close()
|
||
})
|
||
// Ideally we wouldn't wait too long here, but sometimes the the
|
||
// networking needs more time to resolve itself.
|
||
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
|
||
defer cancel()
|
||
if !agentConn.AwaitReachable(ctx) {
|
||
t.Fatal("agent not reachable")
|
||
}
|
||
return agentConn, c, statsCh, fs, agnt
|
||
}
|
||
|
||
var dialTestPayload = []byte("dean-was-here123")
|
||
|
||
func testDial(ctx context.Context, t testing.TB, c net.Conn) {
|
||
t.Helper()
|
||
|
||
if deadline, ok := ctx.Deadline(); ok {
|
||
err := c.SetDeadline(deadline)
|
||
assert.NoError(t, err)
|
||
defer func() {
|
||
err := c.SetDeadline(time.Time{})
|
||
assert.NoError(t, err)
|
||
}()
|
||
}
|
||
|
||
assertWritePayload(t, c, dialTestPayload)
|
||
assertReadPayload(t, c, dialTestPayload)
|
||
}
|
||
|
||
func testAccept(ctx context.Context, t testing.TB, c net.Conn) {
|
||
t.Helper()
|
||
defer c.Close()
|
||
|
||
if deadline, ok := ctx.Deadline(); ok {
|
||
err := c.SetDeadline(deadline)
|
||
assert.NoError(t, err)
|
||
defer func() {
|
||
err := c.SetDeadline(time.Time{})
|
||
assert.NoError(t, err)
|
||
}()
|
||
}
|
||
|
||
assertReadPayload(t, c, dialTestPayload)
|
||
assertWritePayload(t, c, dialTestPayload)
|
||
}
|
||
|
||
func assertReadPayload(t testing.TB, r io.Reader, payload []byte) {
|
||
t.Helper()
|
||
b := make([]byte, len(payload)+16)
|
||
n, err := r.Read(b)
|
||
assert.NoError(t, err, "read payload")
|
||
assert.Equal(t, len(payload), n, "read payload length does not match")
|
||
assert.Equal(t, payload, b[:n])
|
||
}
|
||
|
||
func assertWritePayload(t testing.TB, w io.Writer, payload []byte) {
|
||
t.Helper()
|
||
n, err := w.Write(payload)
|
||
assert.NoError(t, err, "write payload")
|
||
assert.Equal(t, len(payload), n, "written payload length does not match")
|
||
}
|
||
|
||
func testSessionOutput(t *testing.T, session *ssh.Session, expected, unexpected []string, expectedRe *regexp.Regexp) {
|
||
t.Helper()
|
||
logger := testutil.Logger(t)
|
||
|
||
err := session.RequestPty("xterm", 128, 128, ssh.TerminalModes{})
|
||
require.NoError(t, err)
|
||
|
||
var stdout bytes.Buffer
|
||
session.Stdout = &stdout
|
||
stdin := testutil.NewWriterAttachedToSSHSession(t, logger.Named("sshin"), session)
|
||
err = session.Shell()
|
||
require.NoError(t, err)
|
||
|
||
stdin.WriteLine("exit 0")
|
||
|
||
waitErr := make(chan error, 1)
|
||
go func() {
|
||
waitErr <- session.Wait()
|
||
}()
|
||
select {
|
||
case err = <-waitErr:
|
||
require.NoError(t, err)
|
||
case <-time.After(testutil.WaitLong):
|
||
require.Fail(t, "timed out waiting for session to exit")
|
||
}
|
||
|
||
for _, unexpected := range unexpected {
|
||
require.NotContains(t, stdout.String(), unexpected, "should not show output")
|
||
}
|
||
for _, expect := range expected {
|
||
require.Contains(t, stdout.String(), expect, "should show output")
|
||
}
|
||
if expectedRe != nil {
|
||
require.Regexp(t, expectedRe, stdout.String())
|
||
}
|
||
}
|
||
|
||
func TestAgent_Metrics_SSH(t *testing.T) {
|
||
t.Parallel()
|
||
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
|
||
defer cancel()
|
||
|
||
registry := prometheus.NewRegistry()
|
||
|
||
//nolint:dogsled
|
||
conn, _, _, _, _ := setupAgent(t, agentsdk.Manifest{}, 0, func(_ *agenttest.Client, o *agent.Options) {
|
||
o.PrometheusRegistry = registry
|
||
})
|
||
|
||
sshClient, err := conn.SSHClient(ctx)
|
||
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)
|
||
|
||
expected := []struct {
|
||
Name string
|
||
Type proto.Stats_Metric_Type
|
||
CheckFn func(float64) error
|
||
Labels []*proto.Stats_Metric_Label
|
||
}{
|
||
{
|
||
Name: "agent_reconnecting_pty_connections_total",
|
||
Type: proto.Stats_Metric_COUNTER,
|
||
CheckFn: func(v float64) error {
|
||
if v == 0 {
|
||
return nil
|
||
}
|
||
return xerrors.Errorf("expected 0, got %f", v)
|
||
},
|
||
},
|
||
{
|
||
Name: "agent_sessions_total",
|
||
Type: proto.Stats_Metric_COUNTER,
|
||
CheckFn: func(v float64) error {
|
||
if v == 1 {
|
||
return nil
|
||
}
|
||
return xerrors.Errorf("expected 1, got %f", v)
|
||
},
|
||
Labels: []*proto.Stats_Metric_Label{
|
||
{
|
||
Name: "magic_type",
|
||
Value: "ssh",
|
||
},
|
||
{
|
||
Name: "pty",
|
||
Value: "no",
|
||
},
|
||
},
|
||
},
|
||
{
|
||
Name: "agent_ssh_server_failed_connections_total",
|
||
Type: proto.Stats_Metric_COUNTER,
|
||
CheckFn: func(v float64) error {
|
||
if v == 0 {
|
||
return nil
|
||
}
|
||
return xerrors.Errorf("expected 0, got %f", v)
|
||
},
|
||
},
|
||
{
|
||
Name: "agent_ssh_server_sftp_connections_total",
|
||
Type: proto.Stats_Metric_COUNTER,
|
||
CheckFn: func(v float64) error {
|
||
if v == 0 {
|
||
return nil
|
||
}
|
||
return xerrors.Errorf("expected 0, got %f", v)
|
||
},
|
||
},
|
||
{
|
||
Name: "agent_ssh_server_sftp_server_errors_total",
|
||
Type: proto.Stats_Metric_COUNTER,
|
||
CheckFn: func(v float64) error {
|
||
if v == 0 {
|
||
return nil
|
||
}
|
||
return xerrors.Errorf("expected 0, got %f", v)
|
||
},
|
||
},
|
||
{
|
||
Name: "coderd_agentstats_currently_reachable_peers",
|
||
Type: proto.Stats_Metric_GAUGE,
|
||
CheckFn: func(float64) error {
|
||
// We can't reliably ping a peer here, and networking is out of
|
||
// scope of this test, so we just test that the metric exists
|
||
// with the correct labels.
|
||
return nil
|
||
},
|
||
Labels: []*proto.Stats_Metric_Label{
|
||
{
|
||
Name: "connection_type",
|
||
Value: "derp",
|
||
},
|
||
},
|
||
},
|
||
{
|
||
Name: "coderd_agentstats_currently_reachable_peers",
|
||
Type: proto.Stats_Metric_GAUGE,
|
||
CheckFn: func(float64) error {
|
||
return nil
|
||
},
|
||
Labels: []*proto.Stats_Metric_Label{
|
||
{
|
||
Name: "connection_type",
|
||
Value: "p2p",
|
||
},
|
||
},
|
||
},
|
||
{
|
||
Name: "coderd_agentstats_startup_script_seconds",
|
||
Type: proto.Stats_Metric_GAUGE,
|
||
CheckFn: func(f float64) error {
|
||
if f >= 0 {
|
||
return nil
|
||
}
|
||
return xerrors.Errorf("expected >= 0, got %f", f)
|
||
},
|
||
Labels: []*proto.Stats_Metric_Label{
|
||
{
|
||
Name: "success",
|
||
Value: "true",
|
||
},
|
||
},
|
||
},
|
||
}
|
||
|
||
var actual []*promgo.MetricFamily
|
||
assert.Eventually(t, func() bool {
|
||
actual, err = registry.Gather()
|
||
if err != nil {
|
||
return false
|
||
}
|
||
count := 0
|
||
for _, m := range actual {
|
||
count += len(m.GetMetric())
|
||
}
|
||
return count == len(expected)
|
||
}, testutil.WaitLong, testutil.IntervalFast)
|
||
|
||
i := 0
|
||
for _, mf := range actual {
|
||
for _, m := range mf.GetMetric() {
|
||
assert.Equal(t, expected[i].Name, mf.GetName())
|
||
assert.Equal(t, expected[i].Type.String(), mf.GetType().String())
|
||
switch expected[i].Type {
|
||
case proto.Stats_Metric_GAUGE:
|
||
assert.NoError(t, expected[i].CheckFn(m.GetGauge().GetValue()), "check fn for %s failed", expected[i].Name)
|
||
case proto.Stats_Metric_COUNTER:
|
||
assert.NoError(t, expected[i].CheckFn(m.GetCounter().GetValue()), "check fn for %s failed", expected[i].Name)
|
||
}
|
||
for j, lbl := range expected[i].Labels {
|
||
assert.Equal(t, m.GetLabel()[j], &promgo.LabelPair{
|
||
Name: &lbl.Name,
|
||
Value: &lbl.Value,
|
||
})
|
||
}
|
||
i++
|
||
}
|
||
}
|
||
|
||
_, 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")
|
||
}
|
||
|
||
// echoOnce accepts a single connection, reads 4 bytes and echos them back
|
||
func echoOnce(t *testing.T, ll net.Listener) {
|
||
t.Helper()
|
||
conn, err := ll.Accept()
|
||
if err != nil {
|
||
return
|
||
}
|
||
defer conn.Close()
|
||
b := make([]byte, 4)
|
||
_, err = conn.Read(b)
|
||
if !assert.NoError(t, err) {
|
||
return
|
||
}
|
||
_, err = conn.Write(b)
|
||
if !assert.NoError(t, err) {
|
||
return
|
||
}
|
||
}
|
||
|
||
// requireEcho sends 4 bytes and requires the read response to match what was sent.
|
||
func requireEcho(t *testing.T, conn net.Conn) {
|
||
t.Helper()
|
||
_, err := conn.Write([]byte("test"))
|
||
require.NoError(t, err)
|
||
b := make([]byte, 4)
|
||
_, err = conn.Read(b)
|
||
require.NoError(t, err)
|
||
require.Equal(t, "test", string(b))
|
||
}
|
||
|
||
// assertConnectionReport expects the provided connect and disconnect events
|
||
// have been reported.
|
||
func assertConnectionReport(t testing.TB, agentClient *agenttest.Client,
|
||
connect proto.ConnectEvent, disconnect proto.DisconnectEvent,
|
||
) {
|
||
t.Helper()
|
||
|
||
var reports []*proto.ReportConnectionRequest
|
||
if !assert.Eventually(t, func() bool {
|
||
reports = agentClient.GetConnectionReports()
|
||
return len(reports) >= 2
|
||
}, testutil.WaitMedium, testutil.IntervalFast, "waiting for 2 connection reports or more; got %d", len(reports)) {
|
||
return
|
||
}
|
||
|
||
assert.Len(t, reports, 2, "want 2 connection reports")
|
||
|
||
assert.Equal(t, proto.Connection_CONNECT, reports[0].GetConnection().GetAction(), "first report should be connect")
|
||
assert.Equal(t, proto.Connection_DISCONNECT, reports[1].GetConnection().GetAction(), "second report should be disconnect")
|
||
assert.Equal(t, connect.Method, reports[0].GetConnection().GetConnectionMethod(), "connect method should be %s", connect.Method)
|
||
assert.Equal(t, connect.Method, reports[1].GetConnection().GetConnectionMethod(), "disconnect method should be %s", connect.Method)
|
||
assert.Equal(t, connect.AppName, reports[0].GetConnection().GetAppName(), "connect app name should be %s", connect.AppName)
|
||
assert.Equal(t, connect.AppName, reports[1].GetConnection().GetAppName(), "disconnect app name should be %s", connect.AppName)
|
||
assert.Equal(t, reports[0].GetConnection().GetId(), reports[1].GetConnection().GetId(), "disconnect id should match connect id")
|
||
t1 := reports[0].GetConnection().GetTimestamp().AsTime()
|
||
t2 := reports[1].GetConnection().GetTimestamp().AsTime()
|
||
assert.True(t, t1.Before(t2) || t1.Equal(t2), "connect timestamp should be before or equal to disconnect timestamp")
|
||
assert.NotEmpty(t, reports[0].GetConnection().GetIp(), "connect ip should not be empty")
|
||
assert.NotEmpty(t, reports[1].GetConnection().GetIp(), "disconnect ip should not be empty")
|
||
assert.Equal(t, 0, int(reports[0].GetConnection().GetStatusCode()), "connect status code should be 0")
|
||
assert.Equal(t, disconnect.Code, int(reports[1].GetConnection().GetStatusCode()), "disconnect status code should be %d", disconnect.Code)
|
||
assert.Equal(t, "", reports[0].GetConnection().GetReason(), "connect reason should be empty")
|
||
if disconnect.Reason != "" {
|
||
assert.Contains(t, reports[1].GetConnection().GetReason(), disconnect.Reason, "disconnect reason should contain %s", disconnect.Reason)
|
||
} else {
|
||
assert.Equal(t, "", reports[1].GetConnection().GetReason(), "disconnect reason should be empty")
|
||
}
|
||
assert.Equal(t, connect.ClientSessionID, reports[0].GetConnection().GetClientSessionId(), "connect client session id should be %s", connect.ClientSessionID)
|
||
assert.Equal(t, connect.ClientSessionID, reports[1].GetConnection().GetClientSessionId(), "disconnect client session id should be %s", connect.ClientSessionID)
|
||
}
|
||
|
||
func TestAgent_ToolCall(t *testing.T) {
|
||
t.Skip("flaky until PLAT-717: the agent netstack drops a SYN that reuses a source port it holds in TIME_WAIT")
|
||
t.Parallel()
|
||
//nolint:dogsled
|
||
conn, _, _, agentFS, _ := setupAgent(t, agentsdk.Manifest{}, 0)
|
||
conn.SetExtraHeaders(http.Header{workspacesdk.CoderChatIDHeader: {uuid.NewString()}})
|
||
|
||
// Every call uses the tool call context, as chatd does, so a method
|
||
// that wrongly sent the tool call ID would get a saved response.
|
||
tests := []struct {
|
||
name string
|
||
fs afero.Fs
|
||
act func(ctx context.Context, path string) (any, error)
|
||
saved func(workspacesdk.CancelToolCallResponse) (any, error)
|
||
}{
|
||
{
|
||
name: "StartProcess",
|
||
fs: afero.NewOsFs(),
|
||
act: func(ctx context.Context, path string) (any, error) {
|
||
resp, err := conn.StartProcess(ctx, workspacesdk.StartProcessRequest{
|
||
Command: "printf run > " + filepath.Base(path),
|
||
WorkDir: filepath.Dir(path),
|
||
})
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
out, err := conn.ProcessOutput(ctx, resp.ID, &workspacesdk.ProcessOutputOptions{Wait: true})
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
if out.ExitCode == nil {
|
||
return nil, xerrors.Errorf("process output reports no exit: %+v", out)
|
||
}
|
||
return resp, nil
|
||
},
|
||
saved: func(r workspacesdk.CancelToolCallResponse) (any, error) { return r.StartProcessResult() },
|
||
},
|
||
{
|
||
name: "WriteFile",
|
||
fs: agentFS,
|
||
act: func(ctx context.Context, path string) (any, error) {
|
||
return nil, conn.WriteFile(ctx, path, strings.NewReader("run"))
|
||
},
|
||
saved: func(r workspacesdk.CancelToolCallResponse) (any, error) { return nil, r.WriteFileResult() },
|
||
},
|
||
{
|
||
name: "EditFiles",
|
||
fs: agentFS,
|
||
act: func(ctx context.Context, path string) (any, error) {
|
||
return conn.EditFiles(ctx, workspacesdk.FileEditRequest{
|
||
Files: []workspacesdk.FileEdits{{Path: path, Edits: []workspacesdk.FileEdit{{OldText: "reset", NewText: "run"}}}},
|
||
IncludeDiff: true,
|
||
})
|
||
},
|
||
saved: func(r workspacesdk.CancelToolCallResponse) (any, error) { return r.EditFilesResult() },
|
||
},
|
||
}
|
||
for _, tc := range tests {
|
||
t.Run(tc.name, func(t *testing.T) {
|
||
t.Parallel()
|
||
ctx := testutil.Context(t, testutil.WaitLong)
|
||
path := filepath.Join(t.TempDir(), "tool-call.txt")
|
||
reset := func() { require.NoError(t, afero.WriteFile(tc.fs, path, []byte("reset"), 0o600)) }
|
||
requireContent := func(want string) {
|
||
t.Helper()
|
||
content, err := afero.ReadFile(tc.fs, path)
|
||
require.NoError(t, err)
|
||
require.Equal(t, want, string(content))
|
||
}
|
||
id := uuid.New()
|
||
toolCtx := workspacesdk.WithToolCallID(ctx, id)
|
||
|
||
reset()
|
||
first, err := tc.act(toolCtx, path)
|
||
require.NoError(t, err)
|
||
requireContent("run")
|
||
reset()
|
||
repeat, err := tc.act(toolCtx, path)
|
||
require.NoError(t, err)
|
||
require.Equal(t, first, repeat)
|
||
requireContent("reset")
|
||
|
||
canceled, err := conn.CancelToolCall(toolCtx, id)
|
||
require.NoError(t, err)
|
||
require.True(t, canceled.Received)
|
||
saved, err := tc.saved(canceled)
|
||
require.NoError(t, err)
|
||
require.Equal(t, first, saved)
|
||
|
||
_, err = tc.act(toolCtx, path)
|
||
var sdkErr *codersdk.Error
|
||
require.ErrorAs(t, err, &sdkErr)
|
||
require.Equal(t, http.StatusConflict, sdkErr.StatusCode())
|
||
})
|
||
}
|
||
}
|
||
|
||
// TestAgent_ToolCallOtherChat checks that two chats using the same tool
|
||
// call ID get separate processes and cancels.
|
||
func TestAgent_ToolCallOtherChat(t *testing.T) {
|
||
t.Parallel()
|
||
ctx := testutil.Context(t, testutil.WaitLong)
|
||
agentID := uuid.New()
|
||
//nolint:dogsled
|
||
connA, _, _, _, _ := setupAgent(t, agentsdk.Manifest{AgentID: agentID}, 0)
|
||
connA.SetExtraHeaders(http.Header{workspacesdk.CoderChatIDHeader: {uuid.NewString()}})
|
||
// connB shares connA's tailnet connection, which setupAgent closes.
|
||
connB := workspacesdk.NewAgentConn(connA.TailnetConn(), workspacesdk.AgentConnOptions{AgentID: agentID})
|
||
connB.SetExtraHeaders(http.Header{workspacesdk.CoderChatIDHeader: {uuid.NewString()}})
|
||
|
||
id := uuid.New()
|
||
toolCtx := workspacesdk.WithToolCallID(ctx, id)
|
||
_, err := connA.StartProcess(toolCtx, workspacesdk.StartProcessRequest{Command: "sleep 300"})
|
||
require.NoError(t, err)
|
||
t.Cleanup(func() {
|
||
_, err := connA.CancelToolCall(workspacesdk.WithToolCallID(context.Background(), id), id)
|
||
assert.NoError(t, err)
|
||
})
|
||
|
||
_, err = connB.StartProcess(toolCtx, workspacesdk.StartProcessRequest{Command: "printf b"})
|
||
require.NoError(t, err)
|
||
outB, err := connB.ProcessOutput(toolCtx, id.String(), &workspacesdk.ProcessOutputOptions{Wait: true})
|
||
require.NoError(t, err)
|
||
require.Equal(t, "printf b", outB.Command)
|
||
require.Equal(t, "b", outB.Output)
|
||
|
||
_, err = connB.CancelToolCall(toolCtx, id)
|
||
require.NoError(t, err)
|
||
// CancelToolCall marks the process canceled before returning, so a
|
||
// wrong cancel shows here without waiting.
|
||
outA, err := connA.ProcessOutput(toolCtx, id.String(), nil)
|
||
require.NoError(t, err)
|
||
require.Equal(t, "sleep 300", outA.Command)
|
||
require.False(t, outA.Canceled)
|
||
}
|
||
|
||
func TestAgent_ProcessOutputTimeoutFromStartProcess(t *testing.T) {
|
||
t.Parallel()
|
||
//nolint:dogsled
|
||
conn, _, _, _, _ := setupAgent(t, agentsdk.Manifest{}, 0)
|
||
conn.SetExtraHeaders(http.Header{workspacesdk.CoderChatIDHeader: {uuid.NewString()}})
|
||
|
||
// A 1 ms timeout has passed before any wait uses it and one hour cannot
|
||
// pass during the test, so no case depends on timing.
|
||
tests := []struct {
|
||
name string
|
||
timeoutMs int64
|
||
opts workspacesdk.ProcessOutputOptions
|
||
cancel bool
|
||
want workspacesdk.ProcessOutputResponse
|
||
}{
|
||
{
|
||
name: "TimeoutFromStartProcess",
|
||
timeoutMs: 1,
|
||
opts: workspacesdk.ProcessOutputOptions{Wait: true, TimeoutFromStartProcess: true},
|
||
want: workspacesdk.ProcessOutputResponse{Running: true, TimedOut: true},
|
||
},
|
||
{
|
||
name: "PlainWait",
|
||
timeoutMs: 1,
|
||
opts: workspacesdk.ProcessOutputOptions{Wait: true},
|
||
cancel: true,
|
||
want: workspacesdk.ProcessOutputResponse{Canceled: true},
|
||
},
|
||
{
|
||
name: "BeforeTimeout",
|
||
timeoutMs: time.Hour.Milliseconds(),
|
||
want: workspacesdk.ProcessOutputResponse{Running: true},
|
||
},
|
||
}
|
||
for _, tc := range tests {
|
||
t.Run(tc.name, func(t *testing.T) {
|
||
t.Parallel()
|
||
ctx := testutil.Context(t, testutil.WaitLong)
|
||
id := uuid.New()
|
||
toolCtx := workspacesdk.WithToolCallID(ctx, id)
|
||
_, err := conn.StartProcess(toolCtx, workspacesdk.StartProcessRequest{Command: "sleep 300", TimeoutMs: tc.timeoutMs})
|
||
require.NoError(t, err)
|
||
// Cancel the tool call so the agent does not wait for its process
|
||
// on close.
|
||
t.Cleanup(func() {
|
||
_, err := conn.CancelToolCall(workspacesdk.WithToolCallID(context.Background(), id), id)
|
||
assert.NoError(t, err)
|
||
})
|
||
|
||
type result struct {
|
||
resp workspacesdk.ProcessOutputResponse
|
||
err error
|
||
}
|
||
waited := make(chan result, 1)
|
||
go func() {
|
||
resp, err := conn.ProcessOutput(toolCtx, id.String(), &tc.opts)
|
||
waited <- result{resp, err}
|
||
}()
|
||
if tc.cancel {
|
||
_, err := conn.CancelToolCall(toolCtx, id)
|
||
require.NoError(t, err)
|
||
}
|
||
|
||
got := testutil.RequireReceive(ctx, t, waited)
|
||
require.NoError(t, got.err)
|
||
require.Equal(t, tc.want.Running, got.resp.Running)
|
||
require.Equal(t, tc.want.TimedOut, got.resp.TimedOut)
|
||
require.Equal(t, tc.want.Canceled, got.resp.Canceled)
|
||
})
|
||
}
|
||
}
|