Skip to content

Commit 36707f9

Browse files
johnstcnjohnstcn
committed
feat(agent): write up reconnectingpty server to container exec
1 parent a2d4b99 commit 36707f9

File tree

6 files changed

+139
-11
lines changed

6 files changed

+139
-11
lines changed

‎agent/agent_test.go

Lines changed: 70 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -25,24 +25,28 @@ import (
2525
"testing"
2626
"time"
2727

28+
"go.uber.org/goleak"
29+
"tailscale.com/net/speedtest"
30+
"tailscale.com/tailcfg"
31+
2832
"github.com/bramvdbogaerde/go-scp"
2933
"github.com/google/uuid"
34+
"github.com/ory/dockertest/v3"
35+
"github.com/ory/dockertest/v3/docker"
3036
"github.com/pion/udp"
3137
"github.com/pkg/sftp"
3238
"github.com/prometheus/client_golang/prometheus"
3339
promgo "github.com/prometheus/client_model/go"
3440
"github.com/spf13/afero"
3541
"github.com/stretchr/testify/assert"
3642
"github.com/stretchr/testify/require"
37-
"go.uber.org/goleak"
3843
"golang.org/x/crypto/ssh"
3944
"golang.org/x/exp/slices"
4045
"golang.org/x/xerrors"
41-
"tailscale.com/net/speedtest"
42-
"tailscale.com/tailcfg"
4346

4447
"cdr.dev/slog"
4548
"cdr.dev/slog/sloggers/slogtest"
49+
4650
"github.com/coder/coder/v2/agent"
4751
"github.com/coder/coder/v2/agent/agentssh"
4852
"github.com/coder/coder/v2/agent/agenttest"
@@ -1761,6 +1765,69 @@ func TestAgent_ReconnectingPTY(t *testing.T) {
17611765
}
17621766
}
17631767

1768+
// This tests end-to-end functionality of connecting to a running container
1769+
// and executing a command. It creates a real Docker container and runs a
1770+
// command. As such, it does not run by default in CI.
1771+
// You can run it manually as follows:
1772+
//
1773+
// CODER_TEST_USE_DOCKER=1 go test -count=1 ./agent -run TestAgent_ReconnectingPTYContainer
1774+
func TestAgent_ReconnectingPTYContainer(t *testing.T) {
1775+
t.Parallel()
1776+
if ctud, ok := os.LookupEnv("CODER_TEST_USE_DOCKER"); !ok || ctud != "1" {
1777+
t.Skip("Set CODER_TEST_USE_DOCKER=1 to run this test")
1778+
}
1779+
1780+
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
1781+
defer cancel()
1782+
1783+
pool, err := dockertest.NewPool("")
1784+
require.NoError(t, err, "Could not connect to docker")
1785+
ct, err := pool.RunWithOptions(&dockertest.RunOptions{
1786+
Repository: "busybox",
1787+
Tag: "latest",
1788+
Cmd: []string{"sleep", "infnity"},
1789+
}, func(config *docker.HostConfig) {
1790+
config.AutoRemove = true
1791+
config.RestartPolicy = docker.RestartPolicy{Name: "no"}
1792+
})
1793+
require.NoError(t, err, "Could not start container")
1794+
// Wait for container to start
1795+
require.Eventually(t, func() bool {
1796+
ct, ok := pool.ContainerByName(ct.Container.Name)
1797+
return ok && ct.Container.State.Running
1798+
}, testutil.WaitShort, testutil.IntervalSlow, "Container did not start in time")
1799+
1800+
// nolint: dogsled
1801+
conn, _, _, _, _ := setupAgent(t, agentsdk.Manifest{}, 0)
1802+
ac, err := conn.ReconnectingPTY(ctx, uuid.New(), 80, 80, "/bin/sh", func(arp *workspacesdk.AgentReconnectingPTYInit) {
1803+
arp.Container = ct.Container.ID
1804+
})
1805+
require.NoError(t, err, "failed to create ReconnectingPTY")
1806+
defer ac.Close()
1807+
tr := testutil.NewTerminalReader(t, ac)
1808+
1809+
require.NoError(t, tr.ReadUntil(ctx, func(line string) bool {
1810+
return strings.Contains(line, "#") || strings.Contains(line, "$")
1811+
}), "find prompt")
1812+
1813+
require.NoError(t, json.NewEncoder(ac).Encode(workspacesdk.ReconnectingPTYRequest{
1814+
Data: "hostname\r",
1815+
}), "write hostname")
1816+
require.NoError(t, tr.ReadUntil(ctx, func(line string) bool {
1817+
return strings.Contains(line, "hostname")
1818+
}), "find hostname command")
1819+
1820+
require.NoError(t, tr.ReadUntil(ctx, func(line string) bool {
1821+
return strings.Contains(line, ct.Container.Config.Hostname)
1822+
}), "find hostname output")
1823+
require.NoError(t, json.NewEncoder(ac).Encode(workspacesdk.ReconnectingPTYRequest{
1824+
Data: "exit\r",
1825+
}), "write exit command")
1826+
1827+
// Wait for the connection to close.
1828+
require.ErrorIs(t, tr.ReadUntil(ctx, nil), io.EOF)
1829+
}
1830+
17641831
func TestAgent_Dial(t *testing.T) {
17651832
t.Parallel()
17661833

‎agent/agentssh/agentssh.go

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -708,6 +708,8 @@ type EnvInfoer interface {
708708
UserHomeDir() (string, error)
709709
// UserShell returns the shell of the given user.
710710
UserShell(username string) (string, error)
711+
// ModifyCommand modifies the command and arguments before execution.
712+
ModifyCommand(name string, args ...string) (string, []string)
711713
}
712714

713715
type systemEnvInfoer struct{}
@@ -737,6 +739,10 @@ func (systemEnvInfoer) UserShell(username string) (string, error) {
737739
return usershell.Get(username)
738740
}
739741

742+
func (systemEnvInfoer) ModifyCommand(name string, args ...string) (string, []string) {
743+
return name, args
744+
}
745+
740746
// CreateCommand processes raw command input with OpenSSH-like behavior.
741747
// If the script provided is empty, it will default to the users shell.
742748
// This injects environment variables specified by the user at launch too.
@@ -802,7 +808,13 @@ func (s *Server) CreateCommand(ctx context.Context, script string, env []string,
802808
}
803809
}
804810

805-
cmd := s.Execer.PTYCommandContext(ctx, name, args...)
811+
// Modify command prior to execution. This will usually be a no-op, but not always.
812+
modifiedName, modifiedArgs := deps.ModifyCommand(name, args...)
813+
s.logger.Info(ctx, "modified command",
814+
slog.F("before", append([]string{name}, args...)),
815+
slog.F("after", append([]string{modifiedName}, modifiedArgs...)),
816+
)
817+
cmd := s.Execer.PTYCommandContext(ctx, modifiedName, modifiedArgs...)
806818
cmd.Dir = s.config.WorkingDirectory()
807819

808820
// If the metadata directory doesn't exist, we run the command

‎agent/agentssh/agentssh_test.go

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -140,6 +140,10 @@ func (f *fakeEnvInfoer) UserShell(u string) (string, error) {
140140
return f.UserShellFn(u)
141141
}
142142

143+
func (*fakeEnvInfoer) ModifyCommand(cmd string, args ...string) (string, []string) {
144+
return cmd, args
145+
}
146+
143147
func TestNewServer_CloseActiveConnections(t *testing.T) {
144148
t.Parallel()
145149

‎agent/reconnectingpty/server.go

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@ import (
1414
"golang.org/x/xerrors"
1515

1616
"cdr.dev/slog"
17+
"github.com/coder/coder/v2/agent/agentcontainers"
1718
"github.com/coder/coder/v2/agent/agentssh"
1819
"github.com/coder/coder/v2/codersdk/workspacesdk"
1920
)
@@ -116,7 +117,7 @@ func (s *Server) handleConn(ctx context.Context, logger slog.Logger, conn net.Co
116117
}
117118

118119
connectionID := uuid.NewString()
119-
connLogger := logger.With(slog.F("message_id", msg.ID), slog.F("connection_id", connectionID))
120+
connLogger := logger.With(slog.F("message_id", msg.ID), slog.F("connection_id", connectionID), slog.F("container", msg.Container), slog.F("container_user", msg.ContainerUser))
120121
connLogger.Debug(ctx, "starting handler")
121122

122123
defer func() {
@@ -158,8 +159,17 @@ func (s *Server) handleConn(ctx context.Context, logger slog.Logger, conn net.Co
158159
}
159160
}()
160161

162+
var ei agentssh.EnvInfoer
163+
if msg.Container != "" {
164+
dei, err := agentcontainers.EnvInfo(ctx, s.commandCreator.Execer, msg.Container, msg.ContainerUser)
165+
if err != nil {
166+
return xerrors.Errorf("get container env info: %w", err)
167+
}
168+
ei = dei
169+
s.logger.Info(ctx, "got container env info", slog.F("container", msg.Container))
170+
}
161171
// Empty command will default to the users shell!
162-
cmd, err := s.commandCreator.CreateCommand(ctx, msg.Command, nil, nil)
172+
cmd, err := s.commandCreator.CreateCommand(ctx, msg.Command, nil, ei)
163173
if err != nil {
164174
s.errorsTotal.WithLabelValues("create_command").Add(1)
165175
return xerrors.Errorf("create command: %w", err)

‎codersdk/workspacesdk/agentconn.go

Lines changed: 25 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -93,6 +93,24 @@ type AgentReconnectingPTYInit struct {
9393
Height uint16
9494
Width uint16
9595
Command string
96+
// Container, if set, will attempt to exec into a running container visible to the agent.
97+
// This should be a unique container ID (implementation-dependent).
98+
Container string
99+
// ContainerUser, if set, will set the target user when execing into a container.
100+
// This can be a username or UID, depending on the underlying implementation.
101+
// This is ignored if Container is not set.
102+
ContainerUser string
103+
}
104+
105+
// AgentReconnectingPTYInitOption is a functional option for AgentReconnectingPTYInit.
106+
type AgentReconnectingPTYInitOption func(*AgentReconnectingPTYInit)
107+
108+
// AgentReconnectingPTYInitWithContainer sets the container and container user for the reconnecting PTY session.
109+
func AgentReconnectingPTYInitWithContainer(container, containerUser string) AgentReconnectingPTYInitOption {
110+
return func(init *AgentReconnectingPTYInit) {
111+
init.Container = container
112+
init.ContainerUser = containerUser
113+
}
96114
}
97115

98116
// ReconnectingPTYRequest is sent from the client to the server
@@ -107,7 +125,7 @@ type ReconnectingPTYRequest struct {
107125
// ReconnectingPTY spawns a new reconnecting terminal session.
108126
// `ReconnectingPTYRequest` should be JSON marshaled and written to the returned net.Conn.
109127
// Raw terminal output will be read from the returned net.Conn.
110-
func (c *AgentConn) ReconnectingPTY(ctx context.Context, id uuid.UUID, height, width uint16, command string) (net.Conn, error) {
128+
func (c *AgentConn) ReconnectingPTY(ctx context.Context, id uuid.UUID, height, width uint16, command string, initOpts ...AgentReconnectingPTYInitOption) (net.Conn, error) {
111129
ctx, span := tracing.StartSpan(ctx)
112130
defer span.End()
113131

@@ -119,12 +137,16 @@ func (c *AgentConn) ReconnectingPTY(ctx context.Context, id uuid.UUID, height, w
119137
if err != nil {
120138
return nil, err
121139
}
122-
data, err := json.Marshal(AgentReconnectingPTYInit{
140+
rptyInit := AgentReconnectingPTYInit{
123141
ID: id,
124142
Height: height,
125143
Width: width,
126144
Command: command,
127-
})
145+
}
146+
for _, o := range initOpts {
147+
o(&rptyInit)
148+
}
149+
data, err := json.Marshal(rptyInit)
128150
if err != nil {
129151
_ = conn.Close()
130152
return nil, err

‎codersdk/workspacesdk/workspacesdk.go

Lines changed: 15 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -12,12 +12,14 @@ import (
1212
"strconv"
1313
"strings"
1414

15-
"github.com/google/uuid"
16-
"golang.org/x/xerrors"
1715
"tailscale.com/tailcfg"
1816
"tailscale.com/wgengine/capture"
1917

18+
"github.com/google/uuid"
19+
"golang.org/x/xerrors"
20+
2021
"cdr.dev/slog"
22+
2123
"github.com/coder/coder/v2/codersdk"
2224
"github.com/coder/coder/v2/tailnet"
2325
"github.com/coder/coder/v2/tailnet/proto"
@@ -305,6 +307,11 @@ type WorkspaceAgentReconnectingPTYOpts struct {
305307
// issue-reconnecting-pty-signed-token endpoint. If set, the session token
306308
// on the client will not be sent.
307309
SignedToken string
310+
311+
// Container, if set, will attempt to exec into a running container visible to the agent.
312+
// This should be a unique container ID (implementation-dependent).
313+
Container string
314+
ContainerUser string
308315
}
309316

310317
// AgentReconnectingPTY spawns a PTY that reconnects using the token provided.
@@ -320,6 +327,12 @@ func (c *Client) AgentReconnectingPTY(ctx context.Context, opts WorkspaceAgentRe
320327
q.Set("width", strconv.Itoa(int(opts.Width)))
321328
q.Set("height", strconv.Itoa(int(opts.Height)))
322329
q.Set("command", opts.Command)
330+
if opts.Container != "" {
331+
q.Set("container", opts.Container)
332+
}
333+
if opts.ContainerUser != "" {
334+
q.Set("container_user", opts.ContainerUser)
335+
}
323336
// If we're using a signed token, set the query parameter.
324337
if opts.SignedToken != "" {
325338
q.Set(codersdk.SignedAppTokenQueryParameter, opts.SignedToken)

0 commit comments

Comments
 (0)