Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit 3dc994a

Browse files
johnstcnjohnstcn
committed
feat(cli): add capability for SSH command to connect to a running container
1 parent cccdf1e commit 3dc994a

File tree

4 files changed

+167
-12
lines changed

4 files changed

+167
-12
lines changed

‎agent/agent.go

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -307,6 +307,8 @@ func (a *agent) init() {
307307

308308
return a.reportConnection(id, connectionType, ip)
309309
},
310+
311+
ExperimentalContainersEnabled: a.experimentalDevcontainersEnabled,
310312
})
311313
if err != nil {
312314
panic(err)

‎agent/agentssh/agentssh.go

Lines changed: 48 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@ import (
2929

3030
"cdr.dev/slog"
3131

32+
"github.com/coder/coder/v2/agent/agentcontainers"
3233
"github.com/coder/coder/v2/agent/agentexec"
3334
"github.com/coder/coder/v2/agent/agentrsa"
3435
"github.com/coder/coder/v2/agent/usershell"
@@ -104,6 +105,9 @@ type Config struct {
104105
BlockFileTransfer bool
105106
// ReportConnection.
106107
ReportConnection reportConnectionFunc
108+
// Experimental: allow connecting to running containers if
109+
// CODER_AGENT_DEVCONTAINERS_ENABLE=true.
110+
ExperimentalContainersEnabled bool
107111
}
108112

109113
type Server struct {
@@ -324,6 +328,22 @@ func (s *sessionCloseTracker) Close() error {
324328
return s.Session.Close()
325329
}
326330

331+
func extractContainerInfo(env []string) (container, containerUser string, filteredEnv []string) {
332+
for _, kv := range env {
333+
if strings.HasPrefix(kv, "CODER_CONTAINER=") {
334+
container = strings.TrimPrefix(kv, "CODER_CONTAINER=")
335+
}
336+
337+
if strings.HasPrefix(kv, "CODER_CONTAINER_USER=") {
338+
containerUser = strings.TrimPrefix(kv, "CODER_CONTAINER_USER=")
339+
}
340+
}
341+
342+
return container, containerUser, slices.DeleteFunc(env, func(kv string) bool {
343+
return strings.HasPrefix(kv, "CODER_CONTAINER=") || strings.HasPrefix(kv, "CODER_CONTAINER_USER=")
344+
})
345+
}
346+
327347
func (s *Server) sessionHandler(session ssh.Session) {
328348
ctx := session.Context()
329349
id := uuid.New()
@@ -353,6 +373,7 @@ func (s *Server) sessionHandler(session ssh.Session) {
353373
defer s.trackSession(session, false)
354374

355375
reportSession := true
376+
356377
switch magicType {
357378
case MagicSessionTypeVSCode:
358379
s.connCountVSCode.Add(1)
@@ -395,9 +416,19 @@ func (s *Server) sessionHandler(session ssh.Session) {
395416
return
396417
}
397418

419+
container, containerUser, env := extractContainerInfo(env)
420+
s.logger.Debug(ctx, "container info",
421+
slog.F("container", container),
422+
slog.F("container_user", containerUser),
423+
)
424+
398425
switch ss := session.Subsystem(); ss {
399426
case "":
400427
case "sftp":
428+
if s.config.ExperimentalContainersEnabled && container != "" {
429+
closeCause("sftp not yet supported with containers")
430+
return
431+
}
401432
err := s.sftpHandler(logger, session)
402433
if err != nil {
403434
closeCause(err.Error())
@@ -422,7 +453,7 @@ func (s *Server) sessionHandler(session ssh.Session) {
422453
env = append(env, fmt.Sprintf("DISPLAY=localhost:%d.%d", display, x11.ScreenNumber))
423454
}
424455

425-
err := s.sessionStart(logger, session, env, magicType)
456+
err := s.sessionStart(logger, session, env, magicType, container, containerUser)
426457
var exitError *exec.ExitError
427458
if xerrors.As(err, &exitError) {
428459
code := exitError.ExitCode()
@@ -495,30 +526,35 @@ func (s *Server) fileTransferBlocked(session ssh.Session) bool {
495526
return false
496527
}
497528

498-
func (s *Server) sessionStart(logger slog.Logger, session ssh.Session, env []string, magicType MagicSessionType) (retErr error) {
529+
func (s *Server) sessionStart(logger slog.Logger, session ssh.Session, env []string, magicType MagicSessionType, container, containerUser string) (retErr error) {
499530
ctx := session.Context()
500531

501532
magicTypeLabel := magicTypeMetricLabel(magicType)
502533
sshPty, windowSize, isPty := session.Pty()
534+
ptyLabel := "no"
535+
if isPty {
536+
ptyLabel = "yes"
537+
}
503538

504-
cmd, err := s.CreateCommand(ctx, session.RawCommand(), env, nil)
505-
if err != nil {
506-
ptyLabel := "no"
507-
if isPty {
508-
ptyLabel = "yes"
539+
// plumb in envinfoer here to modify command for container exec?
540+
var ei usershell.EnvInfoer
541+
var err error
542+
if s.config.ExperimentalContainersEnabled && container != "" {
543+
ei, err = agentcontainers.EnvInfo(ctx, s.Execer, container, containerUser)
544+
if err != nil {
545+
s.metrics.sessionErrors.WithLabelValues(magicTypeLabel, ptyLabel, "container_env_info").Add(1)
546+
return err
509547
}
548+
}
549+
cmd, err := s.CreateCommand(ctx, session.RawCommand(), env, ei)
550+
if err != nil {
510551
s.metrics.sessionErrors.WithLabelValues(magicTypeLabel, ptyLabel, "create_command").Add(1)
511552
return err
512553
}
513554

514555
if ssh.AgentRequested(session) {
515556
l, err := ssh.NewAgentListener()
516557
if err != nil {
517-
ptyLabel := "no"
518-
if isPty {
519-
ptyLabel = "yes"
520-
}
521-
522558
s.metrics.sessionErrors.WithLabelValues(magicTypeLabel, ptyLabel, "listener").Add(1)
523559
return xerrors.Errorf("new agent listener: %w", err)
524560
}

‎cli/ssh.go

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -76,6 +76,9 @@ func (r *RootCmd) ssh() *serpent.Command {
7676
appearanceConfig codersdk.AppearanceConfig
7777
networkInfoDir string
7878
networkInfoInterval time.Duration
79+
80+
container string
81+
containerUser string
7982
)
8083
client := new(codersdk.Client)
8184
cmd := &serpent.Command{
@@ -282,6 +285,23 @@ func (r *RootCmd) ssh() *serpent.Command {
282285
}
283286
conn.AwaitReachable(ctx)
284287

288+
if container != "" {
289+
cts, err := client.WorkspaceAgentListContainers(ctx, workspaceAgent.ID, nil)
290+
if err != nil {
291+
return xerrors.Errorf("list containers: %w", err)
292+
}
293+
var found bool
294+
for _, c := range cts.Containers {
295+
if c.FriendlyName == container || c.ID == container {
296+
found = true
297+
break
298+
}
299+
}
300+
if !found {
301+
return xerrors.Errorf("container not found: %q", container)
302+
}
303+
}
304+
285305
stopPolling := tryPollWorkspaceAutostop(ctx, client, workspace)
286306
defer stopPolling()
287307

@@ -454,6 +474,17 @@ func (r *RootCmd) ssh() *serpent.Command {
454474
}
455475
}
456476

477+
if container != "" {
478+
for k, v := range map[string]string{
479+
"CODER_CONTAINER": container,
480+
"CODER_CONTAINER_USER": containerUser,
481+
} {
482+
if err := sshSession.Setenv(k, v); err != nil {
483+
return xerrors.Errorf("setenv: %w", err)
484+
}
485+
}
486+
}
487+
457488
err = sshSession.RequestPty("xterm-256color", 128, 128, gossh.TerminalModes{})
458489
if err != nil {
459490
return xerrors.Errorf("request pty: %w", err)
@@ -594,6 +625,20 @@ func (r *RootCmd) ssh() *serpent.Command {
594625
Default: "5s",
595626
Value: serpent.DurationOf(&networkInfoInterval),
596627
},
628+
{
629+
Flag: "container",
630+
FlagShorthand: "c",
631+
Description: "Specifies a container inside the workspace to connect to.",
632+
Value: serpent.StringOf(&container),
633+
Hidden: true, // Hidden until this features is at least in beta.
634+
},
635+
{
636+
Flag: "container-user",
637+
FlagShorthand: "u",
638+
Description: "When connecting to a container, specifies the user to connect as.",
639+
Value: serpent.StringOf(&containerUser),
640+
Hidden: true, // Hidden until this features is at least in beta.
641+
},
597642
sshDisableAutostartOption(serpent.BoolOf(&disableAutostart)),
598643
}
599644
return cmd

‎cli/ssh_test.go

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,8 @@ import (
2424
"time"
2525

2626
"github.com/google/uuid"
27+
"github.com/ory/dockertest/v3"
28+
"github.com/ory/dockertest/v3/docker"
2729
"github.com/spf13/afero"
2830
"github.com/stretchr/testify/assert"
2931
"github.com/stretchr/testify/require"
@@ -1924,6 +1926,76 @@ Expire-Date: 0
19241926
<-cmdDone
19251927
}
19261928

1929+
func TestSSH_Container(t *testing.T) {
1930+
t.Parallel()
1931+
if runtime.GOOS != "linux" {
1932+
t.Skip("Skipping test on non-Linux platform")
1933+
}
1934+
1935+
t.Run("OK", func(t *testing.T) {
1936+
t.Parallel()
1937+
1938+
client, workspace, agentToken := setupWorkspaceForAgent(t)
1939+
ctx := testutil.Context(t, testutil.WaitLong)
1940+
pool, err := dockertest.NewPool("")
1941+
require.NoError(t, err, "Could not connect to docker")
1942+
ct, err := pool.RunWithOptions(&dockertest.RunOptions{
1943+
Repository: "busybox",
1944+
Tag: "latest",
1945+
Cmd: []string{"sleep", "infnity"},
1946+
}, func(config *docker.HostConfig) {
1947+
config.AutoRemove = true
1948+
config.RestartPolicy = docker.RestartPolicy{Name: "no"}
1949+
})
1950+
require.NoError(t, err, "Could not start container")
1951+
// Wait for container to start
1952+
require.Eventually(t, func() bool {
1953+
ct, ok := pool.ContainerByName(ct.Container.Name)
1954+
return ok && ct.Container.State.Running
1955+
}, testutil.WaitShort, testutil.IntervalSlow, "Container did not start in time")
1956+
t.Cleanup(func() {
1957+
err := pool.Purge(ct)
1958+
require.NoError(t, err, "Could not stop container")
1959+
})
1960+
1961+
_ = agenttest.New(t, client.URL, agentToken, func(o *agent.Options) {
1962+
o.ExperimentalContainersEnabled = true
1963+
})
1964+
_ = coderdtest.NewWorkspaceAgentWaiter(t, client, workspace.ID).Wait()
1965+
1966+
inv, root := clitest.New(t, "ssh", workspace.Name, "-c", ct.Container.ID)
1967+
clitest.SetupConfig(t, client, root)
1968+
ptty := ptytest.New(t).Attach(inv)
1969+
1970+
cmdDone := tGo(t, func() {
1971+
err := inv.WithContext(ctx).Run()
1972+
assert.NoError(t, err)
1973+
})
1974+
1975+
ptty.ExpectMatch(" #")
1976+
ptty.WriteLine("hostname")
1977+
ptty.ExpectMatch(ct.Container.Config.Hostname)
1978+
ptty.WriteLine("exit")
1979+
<-cmdDone
1980+
})
1981+
1982+
t.Run("NotFound", func(t *testing.T) {
1983+
t.Parallel()
1984+
1985+
ctx := testutil.Context(t, testutil.WaitShort)
1986+
client, workspace, agentToken := setupWorkspaceForAgent(t)
1987+
_ = agenttest.New(t, client.URL, agentToken, func(o *agent.Options) {
1988+
o.ExperimentalContainersEnabled = true
1989+
})
1990+
_ = coderdtest.NewWorkspaceAgentWaiter(t, client, workspace.ID).Wait()
1991+
1992+
inv, root := clitest.New(t, "ssh", workspace.Name, "-c", uuid.NewString())
1993+
clitest.SetupConfig(t, client, root)
1994+
err := inv.WithContext(ctx).Run()
1995+
require.ErrorContains(t, err, "container not found:")
1996+
})
1997+
}
1998+
19271999
// tGoContext runs fn in a goroutine passing a context that will be
19282000
// canceled on test completion and wait until fn has finished executing.
19292001
// Done and cancel are returned for optionally waiting until completion

0 commit comments

Comments
 (0)