Skip to content

chore(agent/agentssh): extract EnvInfoer #16603

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 4 commits into from
Feb 19, 2025
Merged

chore(agent/agentssh): extract EnvInfoer #16603

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
7b28cfa
chore(agent/agentssh): extract CreateCommandDeps
johnstcn Feb 18, 2025
251d1d2
make fmt
johnstcn Feb 18, 2025
1259518
EnvInfoer is so Go
johnstcn Feb 18, 2025
c5f45e4
add test for custom implementation
johnstcn Feb 18, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion agent/agent.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -340,7 +340,7 @@ func (a *agent) collectMetadata(ctx context.Context, md codersdk.WorkspaceAgentM
// if it can guarantee the clocks are synchronized.
CollectedAt: now,
}
cmdPty, err := a.sshServer.CreateCommand(ctx, md.Script, nil)
cmdPty, err := a.sshServer.CreateCommand(ctx, md.Script, nil, nil)
if err != nil {
result.Error = fmt.Sprintf("create cmd: %+v", err)
return result
Expand Down
2 changes: 1 addition & 1 deletion agent/agentscripts/agentscripts.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -283,7 +283,7 @@ func (r *Runner) run(ctx context.Context, script codersdk.WorkspaceAgentScript,
cmdCtx, ctxCancel = context.WithTimeout(ctx, script.Timeout)
defer ctxCancel()
}
cmdPty, err := r.SSHServer.CreateCommand(cmdCtx, script.Script, nil)
cmdPty, err := r.SSHServer.CreateCommand(cmdCtx, script.Script, nil, nil)
if err != nil {
return xerrors.Errorf("%s script: create command: %w", logPath, err)
}
Expand Down
58 changes: 52 additions & 6 deletions agent/agentssh/agentssh.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -409,7 +409,7 @@ func (s *Server) sessionStart(logger slog.Logger, session ssh.Session, extraEnv
magicTypeLabel := magicTypeMetricLabel(magicType)
sshPty, windowSize, isPty := session.Pty()

cmd, err := s.CreateCommand(ctx, session.RawCommand(), env)
cmd, err := s.CreateCommand(ctx, session.RawCommand(), env, nil)
if err != nil {
ptyLabel := "no"
if isPty {
Expand Down Expand Up @@ -670,17 +670,63 @@ func (s *Server) sftpHandler(logger slog.Logger, session ssh.Session) {
_ = session.Exit(1)
}

// EnvInfoer encapsulates external information required by CreateCommand.
type EnvInfoer interface {
// CurrentUser returns the current user.
CurrentUser() (*user.User, error)
// Environ returns the environment variables of the current process.
Environ() []string
// UserHomeDir returns the home directory of the current user.
UserHomeDir() (string, error)
// UserShell returns the shell of the given user.
UserShell(username string) (string, error)
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I feel like all these methods are related to the environment, figuring out the user, env, home and shell, so EnvInfo sounds like a pretty good name to me. Deps is a bit vague IMO and makes me think of dependency injection. 😄

We could also simplify this interface somewhat. All we really need are:

  • User()
  • Environ()

We'd just make sure the returned user has the correct home directory and shell populated. Not sure if that'd be worth the lift, though.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

EnvInfoer it is!

I'd prefer to do the bare minimum here, but we can consolidate + simplify down the line!

}

type systemEnvInfoer struct{}

var defaultEnvInfoer EnvInfoer = &systemEnvInfoer{}

// DefaultEnvInfoer returns a default implementation of
// EnvInfoer. This reads information using the default Go
// implementations.
func DefaultEnvInfoer() EnvInfoer {
return defaultEnvInfoer
}

func (systemEnvInfoer) CurrentUser() (*user.User, error) {
return user.Current()
}

func (systemEnvInfoer) Environ() []string {
return os.Environ()
}

func (systemEnvInfoer) UserHomeDir() (string, error) {
return userHomeDir()
}

func (systemEnvInfoer) UserShell(username string) (string, error) {
return usershell.Get(username)
}

// CreateCommand processes raw command input with OpenSSH-like behavior.
// If the script provided is empty, it will default to the users shell.
// This injects environment variables specified by the user at launch too.
func (s *Server) CreateCommand(ctx context.Context, script string, env []string) (*pty.Cmd, error) {
currentUser, err := user.Current()
// The final argument is an interface that allows the caller to provide
// alternative implementations for the dependencies of CreateCommand.
// This is useful when creating a command to be run in a separate environment
// (for example, a Docker container). Pass in nil to use the default.
func (s *Server) CreateCommand(ctx context.Context, script string, env []string, deps EnvInfoer) (*pty.Cmd, error) {
if deps == nil {
deps = DefaultEnvInfoer()
}
currentUser, err := deps.CurrentUser()
if err != nil {
return nil, xerrors.Errorf("get current user: %w", err)
}
username := currentUser.Username

shell, err := usershell.Get(username)
shell, err := deps.UserShell(username)
if err != nil {
return nil, xerrors.Errorf("get user shell: %w", err)
}
Expand Down Expand Up @@ -736,13 +782,13 @@ func (s *Server) CreateCommand(ctx context.Context, script string, env []string)
_, err = os.Stat(cmd.Dir)
if cmd.Dir == "" || err != nil {
// Default to user home if a directory is not set.
homedir, err := userHomeDir()
homedir, err := deps.UserHomeDir()
if err != nil {
return nil, xerrors.Errorf("get home dir: %w", err)
}
cmd.Dir = homedir
}
cmd.Env = append(os.Environ(), env...)
cmd.Env = append(deps.Environ(), env...)
cmd.Env = append(cmd.Env, fmt.Sprintf("USER=%s", username))

// Set SSH connection environment variables (these are also set by OpenSSH
Expand Down
38 changes: 36 additions & 2 deletions agent/agentssh/agentssh_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ import (
"context"
"fmt"
"net"
"os/user"
"runtime"
"strings"
"sync"
Expand Down Expand Up @@ -87,7 +88,7 @@ func TestNewServer_ExecuteShebang(t *testing.T) {
t.Run("Basic", func(t *testing.T) {
t.Parallel()
cmd, err := s.CreateCommand(ctx, `#!/bin/bash
echo test`, nil)
echo test`, nil, nil)
require.NoError(t, err)
output, err := cmd.AsExec().CombinedOutput()
require.NoError(t, err)
Expand All @@ -96,12 +97,45 @@ func TestNewServer_ExecuteShebang(t *testing.T) {
t.Run("Args", func(t *testing.T) {
t.Parallel()
cmd, err := s.CreateCommand(ctx, `#!/usr/bin/env bash
echo test`, nil)
echo test`, nil, nil)
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Perhaps we should add a test for a custom implementation as well?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can do!

require.NoError(t, err)
output, err := cmd.AsExec().CombinedOutput()
require.NoError(t, err)
require.Equal(t, "test\n", string(output))
})
t.Run("CustomEnvInfoer", func(t *testing.T) {
t.Parallel()
ei := &fakeEnvInfoer{
CurrentUserFn: func() (u *user.User, err error) {
return nil, assert.AnError
},
}
_, err := s.CreateCommand(ctx, `whatever`, nil, ei)
require.ErrorIs(t, err, assert.AnError)
})
}

type fakeEnvInfoer struct {
CurrentUserFn func() (*user.User, error)
EnvironFn func() []string
UserHomeDirFn func() (string, error)
UserShellFn func(string) (string, error)
}

func (f *fakeEnvInfoer) CurrentUser() (u *user.User, err error) {
return f.CurrentUserFn()
}

func (f *fakeEnvInfoer) Environ() []string {
return f.EnvironFn()
}

func (f *fakeEnvInfoer) UserHomeDir() (string, error) {
return f.UserHomeDirFn()
}

func (f *fakeEnvInfoer) UserShell(u string) (string, error) {
return f.UserShellFn(u)
}

func TestNewServer_CloseActiveConnections(t *testing.T) {
Expand Down
2 changes: 1 addition & 1 deletion agent/reconnectingpty/server.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -159,7 +159,7 @@ func (s *Server) handleConn(ctx context.Context, logger slog.Logger, conn net.Co
}()

// Empty command will default to the users shell!
cmd, err := s.commandCreator.CreateCommand(ctx, msg.Command, nil)
cmd, err := s.commandCreator.CreateCommand(ctx, msg.Command, nil, nil)
if err != nil {
s.errorsTotal.WithLabelValues("create_command").Add(1)
return xerrors.Errorf("create command: %w", err)
Expand Down
Loading