Skip to content

Commit d96fc39

Browse files
DevelopmentCatsDevelopmentCats
authored
Merge branch 'main' into feat-gemini-tmux-icon
2 parents 218c957 + 812d72c commit d96fc39

File tree

9 files changed

+322
-55
lines changed

9 files changed

+322
-55
lines changed

‎coderd/util/strings/strings.go

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,12 @@ package strings
22

33
import (
44
"fmt"
5+
"strconv"
56
"strings"
7+
"unicode"
8+
9+
"github.com/acarl005/stripansi"
10+
"github.com/microcosm-cc/bluemonday"
611
)
712

813
// JoinWithConjunction joins a slice of strings with commas except for the last
@@ -28,3 +33,38 @@ func Truncate(s string, n int) string {
2833
}
2934
return s[:n]
3035
}
36+
37+
var bmPolicy = bluemonday.StrictPolicy()
38+
39+
// UISanitize sanitizes a string for display in the UI.
40+
// The following transformations are applied, in order:
41+
// - HTML tags are removed using bluemonday's strict policy.
42+
// - ANSI escape codes are stripped using stripansi.
43+
// - Consecutive backslashes are replaced with a single backslash.
44+
// - Non-printable characters are removed.
45+
// - Whitespace characters are replaced with spaces.
46+
// - Multiple spaces are collapsed into a single space.
47+
// - Leading and trailing whitespace is trimmed.
48+
func UISanitize(in string) string {
49+
if unq, err := strconv.Unquote(`"` + in + `"`); err == nil {
50+
in = unq
51+
}
52+
in = bmPolicy.Sanitize(in)
53+
in = stripansi.Strip(in)
54+
var b strings.Builder
55+
var spaceSeen bool
56+
for _, r := range in {
57+
if unicode.IsSpace(r) {
58+
if !spaceSeen {
59+
_, _ = b.WriteRune(' ')
60+
spaceSeen = true
61+
}
62+
continue
63+
}
64+
spaceSeen = false
65+
if unicode.IsPrint(r) {
66+
_, _ = b.WriteRune(r)
67+
}
68+
}
69+
return strings.TrimSpace(b.String())
70+
}

‎coderd/util/strings/strings_test.go

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ package strings_test
33
import (
44
"testing"
55

6+
"github.com/stretchr/testify/assert"
67
"github.com/stretchr/testify/require"
78

89
"github.com/coder/coder/v2/coderd/util/strings"
@@ -37,3 +38,41 @@ func TestTruncate(t *testing.T) {
3738
})
3839
}
3940
}
41+
42+
func TestUISanitize(t *testing.T) {
43+
t.Parallel()
44+
45+
for _, tt := range []struct {
46+
s string
47+
expected string
48+
}{
49+
{"normal text", "normal text"},
50+
{"\tfoo \r\\nbar ", "foo bar"},
51+
{"通常のテキスト", "通常のテキスト"},
52+
{"foo\nbar", "foo bar"},
53+
{"foo\tbar", "foo bar"},
54+
{"foo\rbar", "foo bar"},
55+
{"foo\x00bar", "foobar"},
56+
{"\u202Eabc", "abc"},
57+
{"\u200Bzero width", "zero width"},
58+
{"foo\x1b[31mred\x1b[0mbar", "fooredbar"},
59+
{"foo\u0008bar", "foobar"},
60+
{"foo\x07bar", "foobar"},
61+
{"foo\uFEFFbar", "foobar"},
62+
{"<a href='javascript:alert(1)'>link</a>", "link"},
63+
{"<style>body{display:none}</style>", ""},
64+
{"<html>HTML</html>", "HTML"},
65+
{"<br>line break", "line break"},
66+
{"<link rel='stylesheet' href='evil.css'>", ""},
67+
{"<img src=1 onerror=alert(1)>", ""},
68+
{"<!-- comment -->visible", "visible"},
69+
{"<script>alert('xss')</script>", ""},
70+
{"<iframe src='evil.com'></iframe>", ""},
71+
} {
72+
t.Run(tt.expected, func(t *testing.T) {
73+
t.Parallel()
74+
actual := strings.UISanitize(tt.s)
75+
assert.Equal(t, tt.expected, actual)
76+
})
77+
}
78+
}

‎coderd/workspaceagents.go

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,7 @@ import (
4141
"github.com/coder/coder/v2/coderd/rbac/policy"
4242
"github.com/coder/coder/v2/coderd/telemetry"
4343
maputil "github.com/coder/coder/v2/coderd/util/maps"
44+
strutil "github.com/coder/coder/v2/coderd/util/strings"
4445
"github.com/coder/coder/v2/coderd/wspubsub"
4546
"github.com/coder/coder/v2/codersdk"
4647
"github.com/coder/coder/v2/codersdk/agentsdk"
@@ -383,6 +384,9 @@ func (api *API) patchWorkspaceAgentAppStatus(rw http.ResponseWriter, r *http.Req
383384
return
384385
}
385386

387+
// Treat the message as untrusted input.
388+
cleaned := strutil.UISanitize(req.Message)
389+
386390
// nolint:gocritic // This is a system restricted operation.
387391
_, err = api.Database.InsertWorkspaceAppStatus(dbauthz.AsSystemRestricted(ctx), database.InsertWorkspaceAppStatusParams{
388392
ID: uuid.New(),
@@ -391,7 +395,7 @@ func (api *API) patchWorkspaceAgentAppStatus(rw http.ResponseWriter, r *http.Req
391395
AgentID: workspaceAgent.ID,
392396
AppID: app.ID,
393397
State: database.WorkspaceAppStatusState(req.State),
394-
Message: req.Message,
398+
Message: cleaned,
395399
Uri: sql.NullString{
396400
String: req.URI,
397401
Valid: req.URI != "",

‎codersdk/toolsdk/bash.go

Lines changed: 47 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -21,9 +21,10 @@ import (
2121
)
2222

2323
type WorkspaceBashArgs struct {
24-
Workspace string `json:"workspace"`
25-
Command string `json:"command"`
26-
TimeoutMs int `json:"timeout_ms,omitempty"`
24+
Workspace string `json:"workspace"`
25+
Command string `json:"command"`
26+
TimeoutMs int `json:"timeout_ms,omitempty"`
27+
Background bool `json:"background,omitempty"`
2728
}
2829

2930
type WorkspaceBashResult struct {
@@ -50,9 +51,13 @@ The workspace parameter supports various formats:
5051
The timeout_ms parameter specifies the command timeout in milliseconds (defaults to 60000ms, maximum of 300000ms).
5152
If the command times out, all output captured up to that point is returned with a cancellation message.
5253
54+
For background commands (background: true), output is captured until the timeout is reached, then the command
55+
continues running in the background. The captured output is returned as the result.
56+
5357
Examples:
5458
- workspace: "my-workspace", command: "ls -la"
5559
- workspace: "john/dev-env", command: "git status", timeout_ms: 30000
60+
- workspace: "my-workspace", command: "npm run dev", background: true, timeout_ms: 10000
5661
- workspace: "my-workspace.main", command: "docker ps"`,
5762
Schema: aisdk.Schema{
5863
Properties: map[string]any{
@@ -70,6 +75,10 @@ Examples:
7075
"default": 60000,
7176
"minimum": 1,
7277
},
78+
"background": map[string]any{
79+
"type": "boolean",
80+
"description": "Whether to run the command in the background. Output is captured until timeout, then the command continues running in the background.",
81+
},
7382
},
7483
Required: []string{"workspace", "command"},
7584
},
@@ -137,23 +146,35 @@ Examples:
137146

138147
// Set default timeout if not specified (60 seconds)
139148
timeoutMs := args.TimeoutMs
149+
defaultTimeoutMs := 60000
140150
if timeoutMs <= 0 {
141-
timeoutMs = 60000
151+
timeoutMs = defaultTimeoutMs
152+
}
153+
command := args.Command
154+
if args.Background {
155+
// For background commands, use nohup directly to ensure they survive SSH session
156+
// termination. This captures output normally but allows the process to continue
157+
// running even after the SSH connection closes.
158+
command = fmt.Sprintf("nohup %s </dev/null 2>&1", args.Command)
142159
}
143160

144-
// Create context with timeout
145-
ctx, cancel = context.WithTimeout(ctx, time.Duration(timeoutMs)*time.Millisecond)
146-
defer cancel()
161+
// Create context with command timeout (replace the broader MCP timeout)
162+
commandCtx, commandCancel := context.WithTimeout(ctx, time.Duration(timeoutMs)*time.Millisecond)
163+
defer commandCancel()
147164

148165
// Execute command with timeout handling
149-
output, err := executeCommandWithTimeout(ctx, session, args.Command)
166+
output, err := executeCommandWithTimeout(commandCtx, session, command)
150167
outputStr := strings.TrimSpace(string(output))
151168

152169
// Handle command execution results
153170
if err != nil {
154171
// Check if the command timed out
155-
if errors.Is(context.Cause(ctx), context.DeadlineExceeded) {
156-
outputStr += "\nCommand canceled due to timeout"
172+
if errors.Is(context.Cause(commandCtx), context.DeadlineExceeded) {
173+
if args.Background {
174+
outputStr += "\nCommand continues running in background"
175+
} else {
176+
outputStr += "\nCommand canceled due to timeout"
177+
}
157178
return WorkspaceBashResult{
158179
Output: outputStr,
159180
ExitCode: 124,
@@ -387,21 +408,27 @@ func executeCommandWithTimeout(ctx context.Context, session *gossh.Session, comm
387408
return safeWriter.Bytes(), err
388409
case <-ctx.Done():
389410
// Context was canceled (timeout or other cancellation)
390-
// Close the session to stop the command
391-
_ = session.Close()
411+
// Close the session to stop the command, but handle errors gracefully
412+
closeErr := session.Close()
392413

393-
// Give a brief moment to collect any remaining output
394-
timer := time.NewTimer(50 * time.Millisecond)
414+
// Give a brief moment to collect any remaining output and for goroutines to finish
415+
timer := time.NewTimer(100 * time.Millisecond)
395416
defer timer.Stop()
396417

397418
select {
398419
case <-timer.C:
399420
// Timer expired, return what we have
421+
break
400422
case err := <-done:
401423
// Command finished during grace period
402-
return safeWriter.Bytes(), err
424+
if closeErr == nil {
425+
return safeWriter.Bytes(), err
426+
}
427+
// If session close failed, prioritize the context error
428+
break
403429
}
404430

431+
// Return the collected output with the context error
405432
return safeWriter.Bytes(), context.Cause(ctx)
406433
}
407434
}
@@ -421,5 +448,9 @@ func (sw *syncWriter) Write(p []byte) (n int, err error) {
421448
func (sw *syncWriter) Bytes() []byte {
422449
sw.mu.Lock()
423450
defer sw.mu.Unlock()
424-
return sw.w.Bytes()
451+
// Return a copy to prevent race conditions with the underlying buffer
452+
b := sw.w.Bytes()
453+
result := make([]byte, len(b))
454+
copy(result, b)
455+
return result
425456
}

0 commit comments

Comments
 (0)