Skip to content

Commit db9df4b

Browse files
johnstcnjohnstcn
committed
feat(agent): write up reconnectingpty server to container exec
1 parent 304007b commit db9df4b

File tree

8 files changed

+149
-11
lines changed

8 files changed

+149
-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)

‎site/src/pages/TerminalPage/TerminalPage.tsx

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,7 @@ const TerminalPage: FC = () => {
5555
// a round-trip, and must be a UUIDv4.
5656
const reconnectionToken = searchParams.get("reconnect") ?? uuidv4();
5757
const command = searchParams.get("command") || undefined;
58+
const containerName = searchParams.get("container") || undefined;
5859
// The workspace name is in the format:
5960
// <workspace name>[.<agent name>]
6061
const workspaceNameParts = params.workspace?.split(".");
@@ -232,6 +233,7 @@ const TerminalPage: FC = () => {
232233
reconnectionToken,
233234
workspaceAgent.id,
234235
command,
236+
containerName,
235237
terminal.rows,
236238
terminal.cols,
237239
)
@@ -253,6 +255,7 @@ const TerminalPage: FC = () => {
253255
JSON.stringify({
254256
height: terminal.rows,
255257
width: terminal.cols,
258+
container: containerName,
256259
}),
257260
),
258261
);

‎site/src/utils/terminal.ts

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,13 +5,20 @@ export const terminalWebsocketUrl = async (
55
reconnect: string,
66
agentId: string,
77
command: string | undefined,
8+
containerName: string | undefined,
89
height: number,
910
width: number,
1011
): Promise<string> => {
1112
const query = new URLSearchParams({ reconnect });
1213
if (command) {
1314
query.set("command", command);
1415
}
16+
if (containerName) {
17+
query.set("container", containerName);
18+
}
19+
if (command) {
20+
query.set("command", command);
21+
}
1522
query.set("height", height.toString());
1623
query.set("width", width.toString());
1724

0 commit comments

Comments
 (0)
pFad - Phonifier reborn

Pfad - The Proxy pFad of © 2024 Garber Painting. All rights reserved.

Note: This service is not intended for secure transactions such as banking, social media, email, or purchasing. Use at your own risk. We assume no liability whatsoever for broken pages.


Alternative Proxies:

Alternative Proxy

pFad Proxy

pFad v3 Proxy

pFad v4 Proxy