Skip to content

Commit d4691d9

Browse files
kylecarbskylecarbs
committed
Merge branch 'main' into gitsshwin
2 parents 6fc556f + fccd4fa commit d4691d9

24 files changed

+719
-79
lines changed

‎.github/workflows/coder.yaml

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -172,14 +172,14 @@ jobs:
172172
key: ${{ runner.os }}-go-mod-${{ hashFiles('**/go.sum') }}
173173

174174
- name: Install goreleaser
175-
uses: jaxxstorm/action-install-gh-release@v1.4.0
175+
uses: jaxxstorm/action-install-gh-release@v1.5.0
176176
env:
177177
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
178178
with:
179179
repo: gotestyourself/gotestsum
180180
tag: v1.7.0
181181

182-
- uses: hashicorp/setup-terraform@v1
182+
- uses: hashicorp/setup-terraform@v2
183183
with:
184184
terraform_version: 1.1.2
185185
terraform_wrapper: false
@@ -241,14 +241,14 @@ jobs:
241241
key: ${{ runner.os }}-go-mod-${{ hashFiles('**/go.sum') }}
242242

243243
- name: Install goreleaser
244-
uses: jaxxstorm/action-install-gh-release@v1.4.0
244+
uses: jaxxstorm/action-install-gh-release@v1.5.0
245245
env:
246246
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
247247
with:
248248
repo: gotestyourself/gotestsum
249249
tag: v1.7.0
250250

251-
- uses: hashicorp/setup-terraform@v1
251+
- uses: hashicorp/setup-terraform@v2
252252
with:
253253
terraform_version: 1.1.2
254254
terraform_wrapper: false
@@ -449,7 +449,7 @@ jobs:
449449
with:
450450
go-version: "~1.18"
451451

452-
- uses: hashicorp/setup-terraform@v1
452+
- uses: hashicorp/setup-terraform@v2
453453
with:
454454
terraform_version: 1.1.2
455455
terraform_wrapper: false

‎.vscode/settings.json

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616
"gographviz",
1717
"goleak",
1818
"gossh",
19+
"gsyslog",
1920
"hashicorp",
2021
"hclsyntax",
2122
"httpmw",

‎agent/agent.go

Lines changed: 87 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -11,10 +11,14 @@ import (
1111
"os"
1212
"os/exec"
1313
"os/user"
14+
"runtime"
1415
"strings"
1516
"sync"
1617
"time"
1718

19+
gsyslog "github.com/hashicorp/go-syslog"
20+
"go.uber.org/atomic"
21+
1822
"cdr.dev/slog"
1923
"github.com/coder/coder/agent/usershell"
2024
"github.com/coder/coder/peer"
@@ -30,10 +34,11 @@ import (
3034
)
3135

3236
type Options struct {
33-
Logger slog.Logger
37+
EnvironmentVariables map[string]string
38+
StartupScript string
3439
}
3540

36-
type Dialer func(ctx context.Context, logger slog.Logger) (*peerbroker.Listener, error)
41+
type Dialer func(ctx context.Context, logger slog.Logger) (*Options, *peerbroker.Listener, error)
3742

3843
func New(dialer Dialer, logger slog.Logger) io.Closer {
3944
ctx, cancelFunc := context.WithCancel(context.Background())
@@ -56,16 +61,21 @@ type agent struct {
5661
closeMutex sync.Mutex
5762
closed chan struct{}
5863

59-
sshServer *ssh.Server
64+
// Environment variables sent by Coder to inject for shell sessions.
65+
// This is atomic because values can change after reconnect.
66+
envVars atomic.Value
67+
startupScript atomic.Bool
68+
sshServer *ssh.Server
6069
}
6170

6271
func (a *agent) run(ctx context.Context) {
72+
var options *Options
6373
var peerListener *peerbroker.Listener
6474
var err error
6575
// An exponential back-off occurs when the connection is failing to dial.
6676
// This is to prevent server spam in case of a coderd outage.
6777
for retrier := retry.New(50*time.Millisecond, 10*time.Second); retrier.Wait(ctx); {
68-
peerListener, err = a.dialer(ctx, a.logger)
78+
options, peerListener, err = a.dialer(ctx, a.logger)
6979
if err != nil {
7080
if errors.Is(err, context.Canceled) {
7181
return
@@ -84,6 +94,20 @@ func (a *agent) run(ctx context.Context) {
8494
return
8595
default:
8696
}
97+
a.envVars.Store(options.EnvironmentVariables)
98+
99+
if a.startupScript.CAS(false, true) {
100+
// The startup script has not ran yet!
101+
go func() {
102+
err := a.runStartupScript(ctx, options.StartupScript)
103+
if errors.Is(err, context.Canceled) {
104+
return
105+
}
106+
if err != nil {
107+
a.logger.Warn(ctx, "agent script failed", slog.Error(err))
108+
}
109+
}()
110+
}
87111

88112
for {
89113
conn, err := peerListener.Accept()
@@ -102,6 +126,48 @@ func (a *agent) run(ctx context.Context) {
102126
}
103127
}
104128

129+
func (*agent) runStartupScript(ctx context.Context, script string) error {
130+
if script == "" {
131+
return nil
132+
}
133+
currentUser, err := user.Current()
134+
if err != nil {
135+
return xerrors.Errorf("get current user: %w", err)
136+
}
137+
username := currentUser.Username
138+
139+
shell, err := usershell.Get(username)
140+
if err != nil {
141+
return xerrors.Errorf("get user shell: %w", err)
142+
}
143+
144+
var writer io.WriteCloser
145+
// Attempt to use the syslog to write startup information.
146+
writer, err = gsyslog.NewLogger(gsyslog.LOG_INFO, "USER", "coder-startup-script")
147+
if err != nil {
148+
// If the syslog isn't supported or cannot be created, use a text file in temp.
149+
writer, err = os.CreateTemp("", "coder-startup-script.txt")
150+
if err != nil {
151+
return xerrors.Errorf("open startup script log file: %w", err)
152+
}
153+
}
154+
defer func() {
155+
_ = writer.Close()
156+
}()
157+
caller := "-c"
158+
if runtime.GOOS == "windows" {
159+
caller = "/c"
160+
}
161+
cmd := exec.CommandContext(ctx, shell, caller, script)
162+
cmd.Stdout = writer
163+
cmd.Stderr = writer
164+
err = cmd.Run()
165+
if err != nil {
166+
return xerrors.Errorf("run: %w", err)
167+
}
168+
return nil
169+
}
170+
105171
func (a *agent) handlePeerConn(ctx context.Context, conn *peer.Conn) {
106172
go func() {
107173
select {
@@ -231,8 +297,24 @@ func (a *agent) handleSSHSession(session ssh.Session) error {
231297

232298
// OpenSSH executes all commands with the users current shell.
233299
// We replicate that behavior for IDE support.
234-
cmd := exec.CommandContext(session.Context(), shell, "-c", command)
300+
caller := "-c"
301+
if runtime.GOOS == "windows" {
302+
caller = "/c"
303+
}
304+
cmd := exec.CommandContext(session.Context(), shell, caller, command)
235305
cmd.Env = append(os.Environ(), session.Environ()...)
306+
307+
// Load environment variables passed via the agent.
308+
envVars := a.envVars.Load()
309+
if envVars != nil {
310+
envVarMap, ok := envVars.(map[string]string)
311+
if ok {
312+
for key, value := range envVarMap {
313+
cmd.Env = append(cmd.Env, fmt.Sprintf("%s=%s", key, value))
314+
}
315+
}
316+
}
317+
236318
executablePath, err := os.Executable()
237319
if err != nil {
238320
return xerrors.Errorf("getting os executable: %w", err)

‎agent/agent_test.go

Lines changed: 73 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -12,12 +12,15 @@ import (
1212
"strconv"
1313
"strings"
1414
"testing"
15+
"time"
1516

1617
"github.com/pion/webrtc/v3"
1718
"github.com/pkg/sftp"
1819
"github.com/stretchr/testify/require"
1920
"go.uber.org/goleak"
2021
"golang.org/x/crypto/ssh"
22+
"golang.org/x/text/encoding/unicode"
23+
"golang.org/x/text/transform"
2124

2225
"cdr.dev/slog"
2326
"cdr.dev/slog/sloggers/slogtest"
@@ -37,7 +40,7 @@ func TestAgent(t *testing.T) {
3740
t.Parallel()
3841
t.Run("SessionExec", func(t *testing.T) {
3942
t.Parallel()
40-
session := setupSSHSession(t)
43+
session := setupSSHSession(t, nil)
4144

4245
command := "echo test"
4346
if runtime.GOOS == "windows" {
@@ -50,7 +53,7 @@ func TestAgent(t *testing.T) {
5053

5154
t.Run("GitSSH", func(t *testing.T) {
5255
t.Parallel()
53-
session := setupSSHSession(t)
56+
session := setupSSHSession(t, nil)
5457
command := "sh -c 'echo $GIT_SSH_COMMAND'"
5558
if runtime.GOOS == "windows" {
5659
command = "cmd.exe /c echo %GIT_SSH_COMMAND%"
@@ -62,7 +65,13 @@ func TestAgent(t *testing.T) {
6265

6366
t.Run("SessionTTY", func(t *testing.T) {
6467
t.Parallel()
65-
session := setupSSHSession(t)
68+
if runtime.GOOS == "windows" {
69+
// This might be our implementation, or ConPTY itself.
70+
// It's difficult to find extensive tests for it, so
71+
// it seems like it could be either.
72+
t.Skip("ConPTY appears to be inconsistent on Windows.")
73+
}
74+
session := setupSSHSession(t, nil)
6675
command := "bash"
6776
if runtime.GOOS == "windows" {
6877
command = "cmd.exe"
@@ -76,6 +85,11 @@ func TestAgent(t *testing.T) {
7685
session.Stdin = ptty.Input()
7786
err = session.Start(command)
7887
require.NoError(t, err)
88+
caret := "$"
89+
if runtime.GOOS == "windows" {
90+
caret = ">"
91+
}
92+
ptty.ExpectMatch(caret)
7993
ptty.WriteLine("echo test")
8094
ptty.ExpectMatch("test")
8195
ptty.WriteLine("exit")
@@ -117,7 +131,7 @@ func TestAgent(t *testing.T) {
117131

118132
t.Run("SFTP", func(t *testing.T) {
119133
t.Parallel()
120-
sshClient, err := setupAgent(t).SSHClient()
134+
sshClient, err := setupAgent(t, nil).SSHClient()
121135
require.NoError(t, err)
122136
client, err := sftp.NewClient(sshClient)
123137
require.NoError(t, err)
@@ -129,10 +143,55 @@ func TestAgent(t *testing.T) {
129143
_, err = os.Stat(tempFile)
130144
require.NoError(t, err)
131145
})
146+
147+
t.Run("EnvironmentVariables", func(t *testing.T) {
148+
t.Parallel()
149+
key := "EXAMPLE"
150+
value := "value"
151+
session := setupSSHSession(t, &agent.Options{
152+
EnvironmentVariables: map[string]string{
153+
key: value,
154+
},
155+
})
156+
command := "sh -c 'echo $" + key + "'"
157+
if runtime.GOOS == "windows" {
158+
command = "cmd.exe /c echo %" + key + "%"
159+
}
160+
output, err := session.Output(command)
161+
require.NoError(t, err)
162+
require.Equal(t, value, strings.TrimSpace(string(output)))
163+
})
164+
165+
t.Run("StartupScript", func(t *testing.T) {
166+
t.Parallel()
167+
tempPath := filepath.Join(os.TempDir(), "content.txt")
168+
content := "somethingnice"
169+
setupAgent(t, &agent.Options{
170+
StartupScript: "echo " + content + " > " + tempPath,
171+
})
172+
var gotContent string
173+
require.Eventually(t, func() bool {
174+
content, err := os.ReadFile(tempPath)
175+
if err != nil {
176+
return false
177+
}
178+
if len(content) == 0 {
179+
return false
180+
}
181+
if runtime.GOOS == "windows" {
182+
// Windows uses UTF16! 🪟🪟🪟
183+
content, _, err = transform.Bytes(unicode.UTF16(unicode.LittleEndian, unicode.UseBOM).NewDecoder(), content)
184+
require.NoError(t, err)
185+
}
186+
gotContent = string(content)
187+
return true
188+
}, 15*time.Second, 100*time.Millisecond)
189+
require.Equal(t, content, strings.TrimSpace(gotContent))
190+
})
132191
}
133192

134193
func setupSSHCommand(t *testing.T, beforeArgs []string, afterArgs []string) *exec.Cmd {
135-
agentConn := setupAgent(t)
194+
agentConn := setupAgent(t, nil)
136195
listener, err := net.Listen("tcp", "127.0.0.1:0")
137196
require.NoError(t, err)
138197
go func() {
@@ -160,18 +219,22 @@ func setupSSHCommand(t *testing.T, beforeArgs []string, afterArgs []string) *exe
160219
return exec.Command("ssh", args...)
161220
}
162221

163-
func setupSSHSession(t *testing.T) *ssh.Session {
164-
sshClient, err := setupAgent(t).SSHClient()
222+
func setupSSHSession(t *testing.T, options *agent.Options) *ssh.Session {
223+
sshClient, err := setupAgent(t, options).SSHClient()
165224
require.NoError(t, err)
166225
session, err := sshClient.NewSession()
167226
require.NoError(t, err)
168227
return session
169228
}
170229

171-
func setupAgent(t *testing.T) *agent.Conn {
230+
func setupAgent(t *testing.T, options *agent.Options) *agent.Conn {
231+
if options == nil {
232+
options = &agent.Options{}
233+
}
172234
client, server := provisionersdk.TransportPipe()
173-
closer := agent.New(func(ctx context.Context, logger slog.Logger) (*peerbroker.Listener, error) {
174-
return peerbroker.Listen(server, nil)
235+
closer := agent.New(func(ctx context.Context, logger slog.Logger) (*agent.Options, *peerbroker.Listener, error) {
236+
listener, err := peerbroker.Listen(server, nil)
237+
return options, listener, err
175238
}, slogtest.Make(t, nil).Leveled(slog.LevelDebug))
176239
t.Cleanup(func() {
177240
_ = client.Close()

‎cli/cliui/resources_test.go

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ func TestWorkspaceResources(t *testing.T) {
1717
t.Run("SingleAgentSSH", func(t *testing.T) {
1818
t.Parallel()
1919
ptty := ptytest.New(t)
20+
done := make(chan struct{})
2021
go func() {
2122
err := cliui.WorkspaceResources(ptty.Output(), []codersdk.WorkspaceResource{{
2223
Type: "google_compute_instance",
@@ -32,14 +33,17 @@ func TestWorkspaceResources(t *testing.T) {
3233
WorkspaceName: "example",
3334
})
3435
require.NoError(t, err)
36+
close(done)
3537
}()
3638
ptty.ExpectMatch("coder ssh example")
39+
<-done
3740
})
3841

3942
t.Run("MultipleStates", func(t *testing.T) {
4043
t.Parallel()
4144
ptty := ptytest.New(t)
4245
disconnected := database.Now().Add(-4 * time.Second)
46+
done := make(chan struct{})
4347
go func() {
4448
err := cliui.WorkspaceResources(ptty.Output(), []codersdk.WorkspaceResource{{
4549
Transition: database.WorkspaceTransitionStart,
@@ -82,9 +86,11 @@ func TestWorkspaceResources(t *testing.T) {
8286
HideAccess: false,
8387
})
8488
require.NoError(t, err)
89+
close(done)
8590
}()
8691
ptty.ExpectMatch("google_compute_disk.root")
8792
ptty.ExpectMatch("google_compute_instance.dev")
8893
ptty.ExpectMatch("coder ssh dev.postgres")
94+
<-done
8995
})
9096
}

0 commit comments

Comments
 (0)