From 15cb3db8eb2838ba9ae1079fbce2485b0bd167e4 Mon Sep 17 00:00:00 2001 From: chaojixinren Date: Wed, 8 Jul 2026 18:25:33 +0800 Subject: [PATCH] feat(sdk): add PTY (interactive terminal) support to Go SDK Add sandbox.Pty() mirroring the Python/Node pty API: Create/Connect streaming RPCs plus Kill/SendStdin/Resize selector RPCs over envd's Connect-JSON protocol. PtyHandle streams raw output, exposes Wait for the exit code (surfacing envd end errors such as "signal: killed"), and supports Disconnect to detach and later reattach via Connect. Includes a client-side idle timeout, exit-code-from-status parsing, and httptest-based unit tests. Closes #811 Signed-off-by: chaojixinren --- sdk/go/README.md | 57 +++- sdk/go/pty.go | 694 +++++++++++++++++++++++++++++++++++++++ sdk/go/pty_test.go | 803 +++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 1553 insertions(+), 1 deletion(-) create mode 100644 sdk/go/pty.go create mode 100644 sdk/go/pty_test.go diff --git a/sdk/go/README.md b/sdk/go/README.md index 130c397f2..c10369ced 100644 --- a/sdk/go/README.md +++ b/sdk/go/README.md @@ -1,6 +1,6 @@ # cubesandbox Go SDK -Go SDK for [CubeSandbox](https://github.com/TencentCloud/CubeSandbox). It matches the current Python SDK surface: sandbox lifecycle, code execution, commands, filesystem operations (read, write, list, stat, exists, remove, rename, mkdir, watch), snapshots, clone, rollback, and L7 egress policy. +Go SDK for [CubeSandbox](https://github.com/TencentCloud/CubeSandbox). It matches the current Python SDK surface: sandbox lifecycle, code execution, commands, PTY (interactive terminal), filesystem operations (read, write, list, stat, exists, remove, rename, mkdir, watch), snapshots, clone, rollback, and L7 egress policy. ## Install @@ -71,6 +71,61 @@ fmt.Println(result.Stdout, result.Stderr, result.ExitCode) `Commands.Run` starts `/bin/bash -l -c ` through envd's `process.Process/Start` API and returns stdout, stderr, and the `EndEvent` exit code. Callers are still responsible for treating untrusted shell input carefully. +## PTY (interactive terminal) + +`sb.Pty()` opens a real pseudo-terminal for interactive programs that need a TTY — shells/REPLs, full-screen tools (`vim`, `top`), or agent-driven terminals. It mirrors the Python/Node `sandbox.pty` surface. + +```go +pty := sb.Pty() + +// Start an interactive login shell (80x24). Optionally set opts.User / opts.Cwd. +handle, err := pty.Create(ctx, cubesandbox.PtySize{Rows: 24, Cols: 80}, cubesandbox.PtyCreateOptions{}) +if err != nil { + panic(err) +} + +// Drive it: send input, resize the window. +_ = handle.SendStdin(ctx, []byte("echo hi && stty size\n")) +_ = handle.Resize(ctx, cubesandbox.PtySize{Rows: 40, Cols: 120}) + +// Consume output with Wait's callback OR by ranging handle.Output() — not both, +// they share one stream. Wait blocks until the shell exits (or you call +// handle.Kill / handle.Disconnect) and returns the exit code. +code, err := handle.Wait(func(chunk []byte) { + os.Stdout.Write(chunk) +}) +fmt.Println("pty exited with", code, err) +``` + +Reattach to a still-running PTY from elsewhere with `pty.Connect(ctx, pid, ...)`, or control one by PID without a handle: + +```go +handle, _ := pty.Create(ctx, cubesandbox.PtySize{Rows: 24, Cols: 80}, cubesandbox.PtyCreateOptions{}) +handle.Disconnect() // detach without killing; the shell keeps running + +again, _ := pty.Connect(ctx, handle.PID(), cubesandbox.PtyConnectOptions{}) +_ = pty.SendStdin(ctx, again.PID(), []byte("ls\n")) +killed, _ := pty.Kill(ctx, again.PID()) // false if the PID already exited +_ = killed +``` + +| Method | Description | +|---|---| +| `Pty.Create(ctx, size, opts)` | Start `/bin/bash -i -l` with a PTY; seeds `TERM`/`LANG`/`LC_ALL` (overridable via `opts.Envs`). Streaming `process.Process/Start`. | +| `Pty.Connect(ctx, pid, opts)` | Reattach to a running PTY. Streaming `process.Process/Connect`. | +| `Pty.Kill(ctx, pid)` | `SIGKILL` a PTY; returns `false` (not an error) if the PID was not found. | +| `Pty.SendStdin(ctx, pid, data)` | Write bytes to the PTY master (`SendInput`). | +| `Pty.Resize(ctx, pid, size)` | Resize the window (`Update`). | +| `handle.Output()` | Channel of raw output chunks; closed when the stream ends. | +| `handle.Wait(onData)` | Block until exit, return the exit code; surfaces envd errors (e.g. `signal: killed`). | +| `handle.Disconnect()` | Stop receiving output without killing the PTY. | +| `handle.PID()` / `ExitCode()` / `ErrorMessage()` | PTY process ID; exit code (returns `0, false` until known); and envd end error. | +| `handle.Kill` / `SendStdin` / `Resize` | Per-handle shortcuts that target this PTY's PID. | + +Consume output via **either** `Output()` **or** `Wait(onData)`, not both — they share one stream. + +`PtyCreateOptions.Timeout` / `PtyConnectOptions.Timeout` (default 60s, `<= 0` uses the default) is both sent to envd as `Connect-Timeout-Ms` and enforced client-side as an idle abort that resets on every received frame; on expiry `Wait` returns an "idle" timeout error. + ## Files ```go diff --git a/sdk/go/pty.go b/sdk/go/pty.go new file mode 100644 index 000000000..0ada5a6dd --- /dev/null +++ b/sdk/go/pty.go @@ -0,0 +1,694 @@ +// Copyright (c) 2026 Tencent Inc. +// SPDX-License-Identifier: Apache-2.0 + +package cubesandbox + +import ( + "bytes" + "context" + "encoding/base64" + "encoding/binary" + "encoding/json" + "fmt" + "io" + "net/http" + "regexp" + "strconv" + "strings" + "sync" + "time" +) + +// defaultPtyTimeout mirrors the Python/Node SDKs' 60s default. It is applied two +// ways, matching them: sent to envd as Connect-Timeout-Ms (a server-side +// deadline) and used as a client-side idle abort that resets on every frame +// received. A timeout <= 0 uses this default. +const defaultPtyTimeout = 60 * time.Second + +// signalSIGKILL is the Connect-JSON Signal enum name. The wire format uses the +// string name rather than the protobuf integer. +const signalSIGKILL = "SIGNAL_SIGKILL" + +// PtySize describes a pseudo-terminal window size. +type PtySize struct { + Rows int + Cols int +} + +// PtyCreateOptions configures Pty.Create. +type PtyCreateOptions struct { + // User authenticates the envd process call (Basic auth). Empty defaults to + // "root" to match the Python/Node SDKs. + User string + // Cwd is the working directory for the shell. Empty uses envd's default. + Cwd string + // Envs are extra environment variables. TERM/LANG/LC_ALL are seeded with + // sensible interactive defaults unless overridden here. + Envs map[string]string + // Timeout is the server-side deadline for the stream (Connect-Timeout-Ms). + // Zero uses defaultPtyTimeout. + Timeout time.Duration +} + +// PtyConnectOptions configures Pty.Connect. +type PtyConnectOptions struct { + // Timeout is the server-side deadline for the stream (Connect-Timeout-Ms). + // Zero uses defaultPtyTimeout. + Timeout time.Duration +} + +// Pty is the entry point for interacting with pseudo-terminals in a sandbox. +// +// It mirrors the E2B / Python / Node "sandbox.pty" namespace: Create starts a +// new interactive shell, Connect reattaches to an existing one, and +// Kill/SendStdin/Resize control a PTY by PID without holding a PtyHandle. +type Pty struct { + sandbox *Sandbox +} + +// Pty returns the PTY namespace for this sandbox. +func (s *Sandbox) Pty() *Pty { + return &Pty{sandbox: s} +} + +// Create starts a new PTY running an interactive login bash shell ("/bin/bash +// -i -l") sized to size. The returned PtyHandle streams raw terminal output +// until the process exits or the caller disconnects. +func (p *Pty) Create(ctx context.Context, size PtySize, opts PtyCreateOptions) (*PtyHandle, error) { + if p == nil || p.sandbox == nil { + return nil, fmt.Errorf("pty is not attached to a sandbox") + } + + envs := map[string]string{} + for k, v := range opts.Envs { + envs[k] = v + } + setDefaultEnv(envs, "TERM", "xterm-256color") + setDefaultEnv(envs, "LANG", "C.UTF-8") + setDefaultEnv(envs, "LC_ALL", "C.UTF-8") + + user := opts.User + if user == "" { + user = defaultEnvdUser + } + timeout := opts.Timeout + if timeout <= 0 { + timeout = defaultPtyTimeout + } + + payload := ptyStartRequest{ + Process: ptyProcessConfig{ + Cmd: "/bin/bash", + Args: []string{"-i", "-l"}, + Envs: envs, + Cwd: opts.Cwd, + }, + PTY: ptyConfig{Size: ptySizeWire{Rows: size.Rows, Cols: size.Cols}}, + } + return p.openStream(ctx, "Start", payload, user, timeout) +} + +// Connect reattaches to an already-running PTY identified by pid, returning a +// fresh PtyHandle that streams its output. The PTY is not affected if a +// previous handle was disconnected. +func (p *Pty) Connect(ctx context.Context, pid int, opts PtyConnectOptions) (*PtyHandle, error) { + if p == nil || p.sandbox == nil { + return nil, fmt.Errorf("pty is not attached to a sandbox") + } + timeout := opts.Timeout + if timeout <= 0 { + timeout = defaultPtyTimeout + } + payload := ptySelectorRequest{Process: ptyProcessSelector{PID: pid}} + return p.openStream(ctx, "Connect", payload, "", timeout) +} + +// Kill sends SIGKILL to the PTY process identified by pid. It reports false +// (not an error) when the PID could not be found — e.g. it already exited. +func (p *Pty) Kill(ctx context.Context, pid int) (bool, error) { + if p == nil || p.sandbox == nil { + return false, fmt.Errorf("pty is not attached to a sandbox") + } + return p.unary(ctx, "SendSignal", ptySignalRequest{ + Process: ptyProcessSelector{PID: pid}, + Signal: signalSIGKILL, + }, true) +} + +// SendStdin writes data to the master side of the PTY identified by pid. +func (p *Pty) SendStdin(ctx context.Context, pid int, data []byte) error { + if p == nil || p.sandbox == nil { + return fmt.Errorf("pty is not attached to a sandbox") + } + _, err := p.unary(ctx, "SendInput", ptyInputRequest{ + Process: ptyProcessSelector{PID: pid}, + Input: ptyInput{PTY: base64.StdEncoding.EncodeToString(data)}, + }, false) + return err +} + +// Resize changes the window size of the PTY identified by pid. +func (p *Pty) Resize(ctx context.Context, pid int, size PtySize) error { + if p == nil || p.sandbox == nil { + return fmt.Errorf("pty is not attached to a sandbox") + } + _, err := p.unary(ctx, "Update", ptyUpdateRequest{ + Process: ptyProcessSelector{PID: pid}, + PTY: ptyConfig{Size: ptySizeWire{Rows: size.Rows, Cols: size.Cols}}, + }, false) + return err +} + +// PtyHandle is a handle to a running PTY. Consume its output either by ranging +// over Output() or by calling Wait with an on-data callback (use one, not +// both). Kill/SendStdin/Resize operate on this handle's PID. +type PtyHandle struct { + pid int + pty *Pty + output chan []byte + done chan struct{} + ctx context.Context + control *ptyStreamControl + body io.ReadCloser + once sync.Once + + mu sync.Mutex + exitCode *int + errMsg string + exited bool + readErr error +} + +// PID returns the PTY process ID. +func (h *PtyHandle) PID() int { return h.pid } + +// Output returns the channel of raw PTY output chunks. It is closed when the +// stream ends (process exit, disconnect, or error). +func (h *PtyHandle) Output() <-chan []byte { return h.output } + +// ExitCode returns the PTY's exit code once known. The second result is false +// while the process is still running or if envd never reported one. +func (h *PtyHandle) ExitCode() (int, bool) { + h.mu.Lock() + defer h.mu.Unlock() + if h.exitCode == nil { + return 0, false + } + return *h.exitCode, true +} + +// ErrorMessage returns the error envd reported for the PTY (e.g. "signal: +// killed"), or an empty string if none. +func (h *PtyHandle) ErrorMessage() string { + h.mu.Lock() + defer h.mu.Unlock() + return h.errMsg +} + +// Kill sends SIGKILL to this PTY. See Pty.Kill. +func (h *PtyHandle) Kill(ctx context.Context) (bool, error) { + return h.pty.Kill(ctx, h.pid) +} + +// SendStdin writes data to this PTY's master side. +func (h *PtyHandle) SendStdin(ctx context.Context, data []byte) error { + return h.pty.SendStdin(ctx, h.pid, data) +} + +// Resize changes this PTY's window size. +func (h *PtyHandle) Resize(ctx context.Context, size PtySize) error { + return h.pty.Resize(ctx, h.pid, size) +} + +// Disconnect stops receiving output without killing the PTY. The process keeps +// running inside the sandbox and can be reattached via Pty.Connect. +func (h *PtyHandle) Disconnect() error { + h.control.disconnect() + h.closeBody() + return nil +} + +// closeBody closes the response body exactly once, whether the stream ended +// naturally (readLoop) or was torn down early (Disconnect / openStream errors). +// io.ReadCloser does not guarantee a tolerant double-close, so we gate it here. +func (h *PtyHandle) closeBody() { + h.once.Do(func() { h.body.Close() }) +} + +// Wait blocks until the PTY exits and returns its exit code. onData, if +// non-nil, is invoked with each output chunk as it arrives. It returns an error +// if the stream ended without an end event or if envd reported a PTY error +// (e.g. the process was killed). +func (h *PtyHandle) Wait(onData func([]byte)) (int, error) { + for chunk := range h.output { + if onData != nil { + onData(chunk) + } + } + <-h.done + + h.mu.Lock() + defer h.mu.Unlock() + if h.readErr != nil { + return 0, h.readErr + } + if !h.exited { + return 0, fmt.Errorf("PTY stream ended without an end event") + } + code := 0 + if h.exitCode != nil { + code = *h.exitCode + } + if h.errMsg != "" { + return code, fmt.Errorf("PTY exited with error: %s", h.errMsg) + } + return code, nil +} + +func (h *PtyHandle) readLoop() { + defer close(h.output) + defer close(h.done) + defer h.closeBody() + defer h.control.clearIdle() + + for { + ev, eos, err := readPtyEvent(h.body) + if err != nil { + // A cancelled context (Disconnect / idle timeout) is an expected + // stop; setReadErr suppresses or reclassifies it accordingly. + if err != io.EOF && err != io.ErrUnexpectedEOF { + h.setReadErr(err) + } + return + } + h.control.reset() + if eos { + return + } + if ev == nil { + continue + } + if ev.Data != nil && ev.Data.PTY != "" { + raw, decErr := base64.StdEncoding.DecodeString(ev.Data.PTY) + if decErr != nil { + h.setReadErr(fmt.Errorf("decode pty output: %w", decErr)) + return + } + select { + case h.output <- raw: + case <-h.ctx.Done(): + return + } + } + if ev.End != nil { + h.recordEnd(ev.End) + } + } +} + +func (h *PtyHandle) recordEnd(end *processEndEvent) { + h.mu.Lock() + defer h.mu.Unlock() + if code, ok := end.exitCode(); ok { + h.exitCode = &code + } else if code, ok := exitCodeFromStatus(end.Status); ok { + h.exitCode = &code + } else if end.Exited { + zero := 0 + h.exitCode = &zero + } + if end.Error != "" { + h.errMsg = end.Error + } + h.exited = true +} + +func (h *PtyHandle) setReadErr(err error) { + // Classify the read error while holding h.mu so the disconnect / idle-timeout + // checks are atomic with the write: otherwise the idle timer could fire in + // the window between the check and the lock, and the descriptive timeout + // message would be lost to the raw I/O error. No lock-ordering inversion — + // the control methods only take c.mu, never h.mu. + h.mu.Lock() + defer h.mu.Unlock() + if h.readErr != nil { + return + } + // A user-initiated Disconnect is a clean stop, not a failure. + if h.control.isDisconnected() { + return + } + if h.control.idleFired() { + h.readErr = fmt.Errorf("PTY stream timed out after %s of inactivity", h.control.timeout) + return + } + h.readErr = err +} + +// openStream opens a streaming Connect RPC (Start/Connect), reads frames until +// the start event to learn the PID, then hands the still-open body to a +// background read loop. +func (p *Pty) openStream(ctx context.Context, method string, payload any, user string, timeout time.Duration) (*PtyHandle, error) { + s := p.sandbox + if err := s.ensureClient(); err != nil { + return nil, err + } + + raw, err := json.Marshal(payload) + if err != nil { + return nil, err + } + + streamCtx, cancel := context.WithCancel(ctx) + control := newPtyStreamControl(timeout, cancel) + + req, err := s.newEnvdRequest(streamCtx, http.MethodPost, "/process.Process/"+method, nil, encodeConnectEnvelope(raw)) + if err != nil { + control.disconnect() + return nil, err + } + req.Header.Set("Content-Type", connectContentType) + req.Header.Set("Connect-Protocol-Version", connectProtocolVersion) + req.Header.Set("Connect-Content-Encoding", "identity") + if user != "" { + req.Header.Set("Authorization", basicAuthUser(user)) + } + setConnectTimeout(req, timeout) + + resp, err := s.client.dataHTTP.Do(req) + if err != nil { + control.disconnect() + if control.idleFired() { + return nil, fmt.Errorf("%s timed out after %s", method, timeout) + } + return nil, err + } + if resp.StatusCode >= http.StatusBadRequest { + defer resp.Body.Close() + control.disconnect() + return nil, apiErrorFromResponse(resp) + } + + pid, err := readPtyStartPID(resp.Body, method, control) + if err != nil { + resp.Body.Close() + control.disconnect() + if control.idleFired() { + return nil, fmt.Errorf("%s timed out after %s", method, timeout) + } + return nil, err + } + + h := &PtyHandle{ + pid: pid, + pty: p, + output: make(chan []byte, 64), + done: make(chan struct{}), + ctx: streamCtx, + control: control, + body: resp.Body, + } + go h.readLoop() + return h, nil +} + +// unary sends a unary Connect-JSON request (plain application/json, no 5-byte +// envelope) and reports success. When allowNotFound is set, an HTTP 404 or a +// Connect "not_found" body yields (false, nil) instead of an error. +func (p *Pty) unary(ctx context.Context, method string, payload any, allowNotFound bool) (bool, error) { + s := p.sandbox + if err := s.ensureClient(); err != nil { + return false, err + } + + raw, err := json.Marshal(payload) + if err != nil { + return false, err + } + req, err := s.newEnvdRequest(ctx, http.MethodPost, "/process.Process/"+method, nil, bytes.NewReader(raw)) + if err != nil { + return false, err + } + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Connect-Protocol-Version", connectProtocolVersion) + + resp, err := s.client.dataHTTP.Do(req) + if err != nil { + return false, err + } + defer resp.Body.Close() + + if resp.StatusCode >= http.StatusBadRequest { + body, _ := io.ReadAll(io.LimitReader(resp.Body, 64<<10)) + if allowNotFound && isConnectNotFound(resp.StatusCode, body) { + return false, nil + } + return false, apiErrorFromStatus(resp.StatusCode, fmt.Sprintf("%s failed: %s", method, extractErrorMessage(body, resp.StatusCode))) + } + return true, nil +} + +func readPtyStartPID(r io.Reader, method string, control *ptyStreamControl) (int, error) { + for { + ev, eos, err := readPtyEvent(r) + if err != nil { + return 0, err + } + control.reset() + if eos { + return 0, fmt.Errorf("%s: stream closed before start event", method) + } + if ev != nil && ev.Start != nil { + return ev.Start.PID, nil + } + // Skip keepalive / non-start events while waiting for the PID. + } +} + +// ptyStreamControl drives the client-side idle timeout and manual disconnect +// for a streaming PTY, mirroring the Node SDK's StreamControl. A single reusable +// timer guards the stream; each received frame stamps a "last activity" time +// instead of touching the timer, so the hot path is allocation-free. When the +// timer fires it re-arms itself for the remaining window if a frame arrived in +// the meantime, otherwise it aborts the request (via cancel) and records +// idleFired so callers can surface a timeout error. +type ptyStreamControl struct { + timeout time.Duration + cancel context.CancelFunc + + mu sync.Mutex + timer *time.Timer + last time.Time + fired bool + disc bool +} + +func newPtyStreamControl(timeout time.Duration, cancel context.CancelFunc) *ptyStreamControl { + c := &ptyStreamControl{timeout: timeout, cancel: cancel, last: time.Now()} + if timeout > 0 { + c.mu.Lock() + c.timer = time.AfterFunc(timeout, c.onFire) + c.mu.Unlock() + } + return c +} + +// onFire runs when the single idle timer elapses. If a frame arrived within the +// timeout window (recorded by reset), it re-arms for the remaining time rather +// than aborting. This makes the timeout race-free: a just-received frame can +// never be lost to a timer that fired a moment earlier, because the timer owns +// the decision under c.mu and re-checks the activity timestamp before firing. +func (c *ptyStreamControl) onFire() { + c.mu.Lock() + if c.disc || c.timer == nil { + c.mu.Unlock() + return + } + if idle := time.Since(c.last); idle < c.timeout { + c.timer.Reset(c.timeout - idle) + c.mu.Unlock() + return + } + c.fired = true + c.mu.Unlock() + c.cancel() +} + +// reset records that a frame just arrived. It only stamps the activity time; +// the timer is re-armed lazily by onFire, keeping this O(1) and allocation-free. +func (c *ptyStreamControl) reset() { + if c.timeout <= 0 { + return + } + c.mu.Lock() + c.last = time.Now() + c.mu.Unlock() +} + +// clearIdle stops the idle timer without aborting the request. +func (c *ptyStreamControl) clearIdle() { + c.mu.Lock() + defer c.mu.Unlock() + if c.timer != nil { + c.timer.Stop() + c.timer = nil + } +} + +// disconnect marks the stream as intentionally closed and aborts the request. +func (c *ptyStreamControl) disconnect() { + c.mu.Lock() + already := c.disc + c.disc = true + if c.timer != nil { + c.timer.Stop() + c.timer = nil + } + c.mu.Unlock() + if !already { + c.cancel() + } +} + +func (c *ptyStreamControl) idleFired() bool { + c.mu.Lock() + defer c.mu.Unlock() + return c.fired +} + +func (c *ptyStreamControl) isDisconnected() bool { + c.mu.Lock() + defer c.mu.Unlock() + return c.disc +} + +var ( + reExitStatus = regexp.MustCompile(`(?:exit status|exited with code)\s+(-?\d+)`) + reSignalStatus = regexp.MustCompile(`(?:signal|terminated by signal)\s+(\d+)`) +) + +// exitCodeFromStatus best-effort parses an exit code from envd's free-form +// end-event "status" string, mirroring the Python/Node SDKs. Signals map to +// 128+signal, matching shell conventions. +func exitCodeFromStatus(status string) (int, bool) { + if status == "" { + return 0, false + } + if m := reExitStatus.FindStringSubmatch(status); m != nil { + if n, err := strconv.Atoi(m[1]); err == nil { + return n, true + } + } + if m := reSignalStatus.FindStringSubmatch(status); m != nil { + if n, err := strconv.Atoi(m[1]); err == nil { + return 128 + n, true + } + } + if status == "exited" { + return 0, true + } + return 0, false +} + +// readPtyEvent reads one Connect frame and decodes its ProcessEvent. eos is +// true when an end-stream trailer was seen; a trailer carrying an error is +// returned as err with eos true. +func readPtyEvent(r io.Reader) (event *processEvent, eos bool, err error) { + flags, payload, err := readConnectEnvelope(r) + if err != nil { + return nil, false, err + } + if flags&connectCompressedFlag != 0 { + return nil, false, fmt.Errorf("unsupported compressed Connect stream message") + } + if flags&connectEndStreamFlag != 0 { + if err := parseConnectEndStream(payload); err != nil { + return nil, true, err + } + return nil, true, nil + } + + var response processStartResponse + if err := json.Unmarshal(payload, &response); err != nil { + return nil, false, fmt.Errorf("decode pty event: %w", err) + } + return response.Event, false, nil +} + +func encodeConnectEnvelope(payload []byte) *bytes.Buffer { + var buf bytes.Buffer + var header [5]byte + binary.BigEndian.PutUint32(header[1:], uint32(len(payload))) + buf.Write(header[:]) + buf.Write(payload) + return &buf +} + +func isConnectNotFound(status int, body []byte) bool { + if status == http.StatusNotFound { + return true + } + var parsed struct { + Code string `json:"code"` + } + if json.Unmarshal(body, &parsed) == nil && strings.EqualFold(parsed.Code, "not_found") { + return true + } + return false +} + +func setDefaultEnv(m map[string]string, key, value string) { + if _, ok := m[key]; !ok { + m[key] = value + } +} + +// --- Connect-JSON request bodies -------------------------------------- + +type ptyStartRequest struct { + Process ptyProcessConfig `json:"process"` + PTY ptyConfig `json:"pty"` +} + +type ptyProcessConfig struct { + Cmd string `json:"cmd"` + Args []string `json:"args"` + Envs map[string]string `json:"envs"` + Cwd string `json:"cwd,omitempty"` +} + +type ptyConfig struct { + Size ptySizeWire `json:"size"` +} + +type ptySizeWire struct { + Rows int `json:"rows"` + Cols int `json:"cols"` +} + +type ptyProcessSelector struct { + PID int `json:"pid"` +} + +type ptySelectorRequest struct { + Process ptyProcessSelector `json:"process"` +} + +type ptySignalRequest struct { + Process ptyProcessSelector `json:"process"` + Signal string `json:"signal"` +} + +type ptyInput struct { + PTY string `json:"pty"` +} + +type ptyInputRequest struct { + Process ptyProcessSelector `json:"process"` + Input ptyInput `json:"input"` +} + +type ptyUpdateRequest struct { + Process ptyProcessSelector `json:"process"` + PTY ptyConfig `json:"pty"` +} diff --git a/sdk/go/pty_test.go b/sdk/go/pty_test.go new file mode 100644 index 000000000..c53b23e78 --- /dev/null +++ b/sdk/go/pty_test.go @@ -0,0 +1,803 @@ +// Copyright (c) 2026 Tencent Inc. +// SPDX-License-Identifier: Apache-2.0 + +package cubesandbox + +import ( + "bytes" + "context" + "encoding/base64" + "encoding/json" + "errors" + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" +) + +func newPtyTestSandbox(t *testing.T, server *httptest.Server) *Sandbox { + t.Helper() + host, port := serverHostPort(t, server.URL) + client := NewClient(Config{ + ProxyNodeIP: host, + ProxyPortHTTP: port, + SandboxDomain: "cube.test", + RequestTimeout: 5 * time.Second, + }) + return &Sandbox{client: client, SandboxID: "sb-pty", EnvdAccessToken: "tok"} +} + +// ptyDataFrame builds a Connect data frame carrying base64-encoded PTY output. +func ptyDataFrame(text string) []byte { + encoded := base64.StdEncoding.EncodeToString([]byte(text)) + return connectEnvelope(0, `{"event":{"data":{"pty":"`+encoded+`"}}}`) +} + +// decodeEnvelopeJSON strips the 5-byte Connect envelope header and unmarshals +// the JSON body into v. +func decodeEnvelopeJSON(t *testing.T, body []byte, v any) { + t.Helper() + if len(body) < 5 { + t.Fatalf("body too short for Connect envelope: %d bytes", len(body)) + } + if err := json.Unmarshal(body[5:], v); err != nil { + t.Fatalf("decode enveloped body: %v", err) + } +} + +func TestPtyCreateStreamsAndWaits(t *testing.T) { + var gotPath, gotCT, gotAuth, gotToken string + var gotBody []byte + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotPath = r.URL.Path + gotCT = r.Header.Get("Content-Type") + gotAuth = r.Header.Get("Authorization") + gotToken = r.Header.Get("X-Access-Token") + gotBody, _ = io.ReadAll(r.Body) + + w.Header().Set("Content-Type", connectContentType) + w.Write(connectEnvelope(0, `{"event":{"start":{"pid":4321}}}`)) + w.Write(ptyDataFrame("hello ")) + w.Write(ptyDataFrame("world")) + w.Write(connectEnvelope(0, `{"event":{"end":{"exitCode":0,"exited":true}}}`)) + w.Write(connectEnvelope(connectEndStreamFlag, `{}`)) + })) + defer server.Close() + + sb := newPtyTestSandbox(t, server) + handle, err := sb.Pty().Create(context.Background(), PtySize{Rows: 24, Cols: 80}, PtyCreateOptions{}) + if err != nil { + t.Fatalf("Create: %v", err) + } + if handle.PID() != 4321 { + t.Fatalf("PID=%d, want 4321", handle.PID()) + } + + var buf bytes.Buffer + code, err := handle.Wait(func(chunk []byte) { buf.Write(chunk) }) + if err != nil { + t.Fatalf("Wait: %v", err) + } + if code != 0 { + t.Fatalf("exit code=%d, want 0", code) + } + if buf.String() != "hello world" { + t.Fatalf("output=%q, want %q", buf.String(), "hello world") + } + if got, ok := handle.ExitCode(); !ok || got != 0 { + t.Fatalf("ExitCode=(%d,%v), want (0,true)", got, ok) + } + + if gotPath != "/process.Process/Start" { + t.Fatalf("path=%q", gotPath) + } + if gotCT != connectContentType { + t.Fatalf("content-type=%q, want %q", gotCT, connectContentType) + } + if gotAuth != basicAuthUser("root") { + t.Fatalf("Authorization=%q, want root basic auth", gotAuth) + } + if gotToken != "tok" { + t.Fatalf("X-Access-Token=%q, want tok", gotToken) + } + + var req ptyStartRequest + decodeEnvelopeJSON(t, gotBody, &req) + if req.Process.Cmd != "/bin/bash" { + t.Fatalf("cmd=%q", req.Process.Cmd) + } + if strings.Join(req.Process.Args, " ") != "-i -l" { + t.Fatalf("args=%v, want [-i -l]", req.Process.Args) + } + if req.PTY.Size.Rows != 24 || req.PTY.Size.Cols != 80 { + t.Fatalf("size=%+v, want {24 80}", req.PTY.Size) + } + if req.Process.Envs["TERM"] != "xterm-256color" { + t.Fatalf("TERM=%q, want xterm-256color", req.Process.Envs["TERM"]) + } + if req.Process.Envs["LANG"] != "C.UTF-8" || req.Process.Envs["LC_ALL"] != "C.UTF-8" { + t.Fatalf("locale envs=%v", req.Process.Envs) + } +} + +func TestPtyCreateUserEnvsCwd(t *testing.T) { + var gotAuth string + var gotBody []byte + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotAuth = r.Header.Get("Authorization") + gotBody, _ = io.ReadAll(r.Body) + w.Header().Set("Content-Type", connectContentType) + w.Write(connectEnvelope(0, `{"event":{"start":{"pid":10}}}`)) + w.Write(connectEnvelope(0, `{"event":{"end":{"exitCode":0,"exited":true}}}`)) + w.Write(connectEnvelope(connectEndStreamFlag, `{}`)) + })) + defer server.Close() + + sb := newPtyTestSandbox(t, server) + handle, err := sb.Pty().Create(context.Background(), PtySize{Rows: 10, Cols: 40}, PtyCreateOptions{ + User: "app", + Cwd: "/work", + Envs: map[string]string{"TERM": "vt100", "FOO": "bar"}, + }) + if err != nil { + t.Fatalf("Create: %v", err) + } + if _, err := handle.Wait(nil); err != nil { + t.Fatalf("Wait: %v", err) + } + + if gotAuth != basicAuthUser("app") { + t.Fatalf("Authorization=%q, want app basic auth", gotAuth) + } + var req ptyStartRequest + decodeEnvelopeJSON(t, gotBody, &req) + if req.Process.Cwd != "/work" { + t.Fatalf("cwd=%q", req.Process.Cwd) + } + if req.Process.Envs["TERM"] != "vt100" { + t.Fatalf("TERM override lost: %q", req.Process.Envs["TERM"]) + } + if req.Process.Envs["FOO"] != "bar" { + t.Fatalf("FOO=%q", req.Process.Envs["FOO"]) + } + if req.Process.Envs["LANG"] != "C.UTF-8" { + t.Fatalf("LANG default lost: %q", req.Process.Envs["LANG"]) + } +} + +func TestPtyWaitSurfacesEndError(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", connectContentType) + w.Write(connectEnvelope(0, `{"event":{"start":{"pid":7}}}`)) + w.Write(connectEnvelope(0, `{"event":{"end":{"error":"signal: killed"}}}`)) + w.Write(connectEnvelope(connectEndStreamFlag, `{}`)) + })) + defer server.Close() + + sb := newPtyTestSandbox(t, server) + handle, err := sb.Pty().Create(context.Background(), PtySize{Rows: 24, Cols: 80}, PtyCreateOptions{}) + if err != nil { + t.Fatalf("Create: %v", err) + } + _, err = handle.Wait(nil) + if err == nil || !strings.Contains(err.Error(), "signal: killed") { + t.Fatalf("Wait error=%v, want it to mention 'signal: killed'", err) + } + if handle.ErrorMessage() != "signal: killed" { + t.Fatalf("ErrorMessage=%q", handle.ErrorMessage()) + } +} + +func TestPtyConnectNoAuthHeader(t *testing.T) { + var gotPath, gotAuth string + var gotBody []byte + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotPath = r.URL.Path + gotAuth = r.Header.Get("Authorization") + gotBody, _ = io.ReadAll(r.Body) + w.Header().Set("Content-Type", connectContentType) + w.Write(connectEnvelope(0, `{"event":{"start":{"pid":99}}}`)) + w.Write(connectEnvelope(0, `{"event":{"end":{"exitCode":0,"exited":true}}}`)) + w.Write(connectEnvelope(connectEndStreamFlag, `{}`)) + })) + defer server.Close() + + sb := newPtyTestSandbox(t, server) + handle, err := sb.Pty().Connect(context.Background(), 99, PtyConnectOptions{}) + if err != nil { + t.Fatalf("Connect: %v", err) + } + if handle.PID() != 99 { + t.Fatalf("PID=%d, want 99", handle.PID()) + } + if _, err := handle.Wait(nil); err != nil { + t.Fatalf("Wait: %v", err) + } + if gotPath != "/process.Process/Connect" { + t.Fatalf("path=%q", gotPath) + } + if gotAuth != "" { + t.Fatalf("Connect should not send Authorization, got %q", gotAuth) + } + var req ptySelectorRequest + decodeEnvelopeJSON(t, gotBody, &req) + if req.Process.PID != 99 { + t.Fatalf("selector pid=%d, want 99", req.Process.PID) + } +} + +func TestPtyStartHTTPError(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + http.Error(w, `{"message":"boom"}`, http.StatusInternalServerError) + })) + defer server.Close() + + sb := newPtyTestSandbox(t, server) + _, err := sb.Pty().Create(context.Background(), PtySize{Rows: 1, Cols: 1}, PtyCreateOptions{}) + var apiErr *APIError + if err == nil { + t.Fatal("Create returned nil error") + } + if !errors.As(err, &apiErr) || apiErr.StatusCode != http.StatusInternalServerError { + t.Fatalf("Create error=%v, want APIError 500", err) + } +} + +func TestPtyKill(t *testing.T) { + t.Run("success", func(t *testing.T) { + var gotPath, gotCT string + var gotBody []byte + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotPath = r.URL.Path + gotCT = r.Header.Get("Content-Type") + gotBody, _ = io.ReadAll(r.Body) + w.WriteHeader(http.StatusOK) + w.Write([]byte(`{}`)) + })) + defer server.Close() + + sb := newPtyTestSandbox(t, server) + killed, err := sb.Pty().Kill(context.Background(), 42) + if err != nil { + t.Fatalf("Kill: %v", err) + } + if !killed { + t.Fatal("Kill returned false, want true") + } + if gotPath != "/process.Process/SendSignal" { + t.Fatalf("path=%q", gotPath) + } + if gotCT != "application/json" { + t.Fatalf("content-type=%q, want application/json", gotCT) + } + var req ptySignalRequest + if err := json.Unmarshal(gotBody, &req); err != nil { + t.Fatalf("decode body: %v", err) + } + if req.Process.PID != 42 || req.Signal != signalSIGKILL { + t.Fatalf("signal req=%+v", req) + } + }) + + t.Run("http 404 is not found", func(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + http.Error(w, `{"message":"no such process"}`, http.StatusNotFound) + })) + defer server.Close() + + sb := newPtyTestSandbox(t, server) + killed, err := sb.Pty().Kill(context.Background(), 42) + if err != nil { + t.Fatalf("Kill returned error on 404: %v", err) + } + if killed { + t.Fatal("Kill returned true on 404, want false") + } + }) + + t.Run("connect not_found body is not found", func(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusBadRequest) + w.Write([]byte(`{"code":"not_found","message":"gone"}`)) + })) + defer server.Close() + + sb := newPtyTestSandbox(t, server) + killed, err := sb.Pty().Kill(context.Background(), 42) + if err != nil { + t.Fatalf("Kill returned error on not_found body: %v", err) + } + if killed { + t.Fatal("Kill returned true on not_found body, want false") + } + }) + + t.Run("other error propagates", func(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + http.Error(w, `{"message":"boom"}`, http.StatusInternalServerError) + })) + defer server.Close() + + sb := newPtyTestSandbox(t, server) + if _, err := sb.Pty().Kill(context.Background(), 42); err == nil { + t.Fatal("Kill returned nil error on 500") + } + }) +} + +func TestPtySendStdin(t *testing.T) { + var gotPath, gotCT string + var gotBody []byte + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotPath = r.URL.Path + gotCT = r.Header.Get("Content-Type") + gotBody, _ = io.ReadAll(r.Body) + w.WriteHeader(http.StatusOK) + })) + defer server.Close() + + sb := newPtyTestSandbox(t, server) + if err := sb.Pty().SendStdin(context.Background(), 5, []byte("ls -la\n")); err != nil { + t.Fatalf("SendStdin: %v", err) + } + if gotPath != "/process.Process/SendInput" { + t.Fatalf("path=%q", gotPath) + } + if gotCT != "application/json" { + t.Fatalf("content-type=%q", gotCT) + } + var req ptyInputRequest + if err := json.Unmarshal(gotBody, &req); err != nil { + t.Fatalf("decode body: %v", err) + } + if req.Process.PID != 5 { + t.Fatalf("pid=%d, want 5", req.Process.PID) + } + decoded, err := base64.StdEncoding.DecodeString(req.Input.PTY) + if err != nil { + t.Fatalf("input.pty not base64: %v", err) + } + if string(decoded) != "ls -la\n" { + t.Fatalf("input=%q, want %q", decoded, "ls -la\n") + } +} + +func TestPtyResize(t *testing.T) { + var gotPath string + var gotBody []byte + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotPath = r.URL.Path + gotBody, _ = io.ReadAll(r.Body) + w.WriteHeader(http.StatusOK) + })) + defer server.Close() + + sb := newPtyTestSandbox(t, server) + if err := sb.Pty().Resize(context.Background(), 8, PtySize{Rows: 30, Cols: 120}); err != nil { + t.Fatalf("Resize: %v", err) + } + if gotPath != "/process.Process/Update" { + t.Fatalf("path=%q", gotPath) + } + var req ptyUpdateRequest + if err := json.Unmarshal(gotBody, &req); err != nil { + t.Fatalf("decode body: %v", err) + } + if req.Process.PID != 8 || req.PTY.Size.Rows != 30 || req.PTY.Size.Cols != 120 { + t.Fatalf("update req=%+v", req) + } +} + +func TestPtyHandleDelegatesToHandlePID(t *testing.T) { + var stdinPID int + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/process.Process/Start": + w.Header().Set("Content-Type", connectContentType) + w.Write(connectEnvelope(0, `{"event":{"start":{"pid":314}}}`)) + w.Write(connectEnvelope(0, `{"event":{"end":{"exitCode":0,"exited":true}}}`)) + w.Write(connectEnvelope(connectEndStreamFlag, `{}`)) + case "/process.Process/SendInput": + body, _ := io.ReadAll(r.Body) + var req ptyInputRequest + json.Unmarshal(body, &req) + stdinPID = req.Process.PID + w.WriteHeader(http.StatusOK) + default: + w.WriteHeader(http.StatusOK) + } + })) + defer server.Close() + + sb := newPtyTestSandbox(t, server) + handle, err := sb.Pty().Create(context.Background(), PtySize{Rows: 24, Cols: 80}, PtyCreateOptions{}) + if err != nil { + t.Fatalf("Create: %v", err) + } + if err := handle.SendStdin(context.Background(), []byte("x")); err != nil { + t.Fatalf("handle.SendStdin: %v", err) + } + if stdinPID != handle.PID() { + t.Fatalf("SendInput pid=%d, want handle pid %d", stdinPID, handle.PID()) + } + if _, err := handle.Wait(nil); err != nil { + t.Fatalf("Wait: %v", err) + } +} + +func TestPtyIdleTimeout(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", connectContentType) + w.Write(connectEnvelope(0, `{"event":{"start":{"pid":5}}}`)) + if f, ok := w.(http.Flusher); ok { + f.Flush() + } + // Go quiet: send no further frames so the client-side idle timer fires. + select { + case <-r.Context().Done(): + case <-time.After(3 * time.Second): + } + })) + defer server.Close() + + sb := newPtyTestSandbox(t, server) + handle, err := sb.Pty().Create(context.Background(), PtySize{Rows: 24, Cols: 80}, PtyCreateOptions{ + Timeout: 150 * time.Millisecond, + }) + if err != nil { + t.Fatalf("Create: %v", err) + } + + done := make(chan struct{}) + var werr error + go func() { + _, werr = handle.Wait(nil) + close(done) + }() + select { + case <-done: + case <-time.After(3 * time.Second): + t.Fatal("Wait did not return after idle timeout") + } + if werr == nil || !strings.Contains(werr.Error(), "timed out") { + t.Fatalf("Wait error=%v, want an idle timeout error", werr) + } +} + +func TestPtyDisconnectKeepsStreamOpen(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", connectContentType) + w.Write(connectEnvelope(0, `{"event":{"start":{"pid":55}}}`)) + if f, ok := w.(http.Flusher); ok { + f.Flush() + } + select { + case <-r.Context().Done(): + case <-time.After(2 * time.Second): + } + })) + defer server.Close() + + sb := newPtyTestSandbox(t, server) + handle, err := sb.Pty().Create(context.Background(), PtySize{Rows: 24, Cols: 80}, PtyCreateOptions{}) + if err != nil { + t.Fatalf("Create: %v", err) + } + if handle.PID() != 55 { + t.Fatalf("PID=%d, want 55", handle.PID()) + } + if err := handle.Disconnect(); err != nil { + t.Fatalf("Disconnect: %v", err) + } + // Second Disconnect must be a no-op, not a panic on double-close. + if err := handle.Disconnect(); err != nil { + t.Fatalf("second Disconnect: %v", err) + } + + done := make(chan struct{}) + var werr error + go func() { + _, werr = handle.Wait(nil) + close(done) + }() + select { + case <-done: + case <-time.After(3 * time.Second): + t.Fatal("Wait did not return after Disconnect") + } + if werr == nil || !strings.Contains(werr.Error(), "without an end event") { + t.Fatalf("Wait after Disconnect error=%v, want 'without an end event'", werr) + } +} + +func TestExitCodeFromStatus(t *testing.T) { + // -1 is accepted on purpose: the regex mirrors the Python/Node SDKs + // (`-?\d+`) so the three SDKs parse envd's free-form status identically. + cases := []struct { + name string + status string + wantCode int + wantOK bool + }{ + {"empty", "", 0, false}, + {"exit status", "exit status 3", 3, true}, + {"exited with code", "exited with code 7", 7, true}, + {"signal", "signal 9", 137, true}, + {"terminated by signal", "terminated by signal 15", 143, true}, + {"plain exited", "exited", 0, true}, + {"negative parity", "exit status -1", -1, true}, + {"unrecognized", "still running", 0, false}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + got, ok := exitCodeFromStatus(tc.status) + if got != tc.wantCode || ok != tc.wantOK { + t.Fatalf("exitCodeFromStatus(%q)=(%d,%v), want (%d,%v)", tc.status, got, ok, tc.wantCode, tc.wantOK) + } + }) + } +} + +// TestPtyRecordEndFallbacks exercises recordEnd's non-exitCode branches: exit +// code parsed from the free-form "status" string, and the "exited" flag with no +// code (defaults to 0). +func TestPtyRecordEndFallbacks(t *testing.T) { + t.Run("status string", func(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", connectContentType) + w.Write(connectEnvelope(0, `{"event":{"start":{"pid":7}}}`)) + w.Write(connectEnvelope(0, `{"event":{"end":{"status":"exit status 5","exited":true}}}`)) + w.Write(connectEnvelope(connectEndStreamFlag, `{}`)) + })) + defer server.Close() + + sb := newPtyTestSandbox(t, server) + handle, err := sb.Pty().Create(context.Background(), PtySize{Rows: 24, Cols: 80}, PtyCreateOptions{}) + if err != nil { + t.Fatalf("Create: %v", err) + } + code, err := handle.Wait(nil) + if err != nil { + t.Fatalf("Wait: %v", err) + } + if code != 5 { + t.Fatalf("exit code=%d, want 5 (parsed from status)", code) + } + }) + + t.Run("exited flag only", func(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", connectContentType) + w.Write(connectEnvelope(0, `{"event":{"start":{"pid":7}}}`)) + w.Write(connectEnvelope(0, `{"event":{"end":{"exited":true}}}`)) + w.Write(connectEnvelope(connectEndStreamFlag, `{}`)) + })) + defer server.Close() + + sb := newPtyTestSandbox(t, server) + handle, err := sb.Pty().Create(context.Background(), PtySize{Rows: 24, Cols: 80}, PtyCreateOptions{}) + if err != nil { + t.Fatalf("Create: %v", err) + } + code, err := handle.Wait(nil) + if err != nil { + t.Fatalf("Wait: %v", err) + } + if got, ok := handle.ExitCode(); !ok || got != 0 || code != 0 { + t.Fatalf("ExitCode=(%d,%v), Wait code=%d, want 0/true/0", got, ok, code) + } + }) +} + +// TestPtyOutputChannelStreaming consumes output by ranging Output() directly, +// the alternative to Wait's callback. +func TestPtyOutputChannelStreaming(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", connectContentType) + w.Write(connectEnvelope(0, `{"event":{"start":{"pid":8}}}`)) + w.Write(ptyDataFrame("hello ")) + w.Write(ptyDataFrame("world")) + w.Write(connectEnvelope(0, `{"event":{"end":{"exitCode":0,"exited":true}}}`)) + w.Write(connectEnvelope(connectEndStreamFlag, `{}`)) + })) + defer server.Close() + + sb := newPtyTestSandbox(t, server) + handle, err := sb.Pty().Create(context.Background(), PtySize{Rows: 24, Cols: 80}, PtyCreateOptions{}) + if err != nil { + t.Fatalf("Create: %v", err) + } + var buf bytes.Buffer + for chunk := range handle.Output() { + buf.Write(chunk) + } + if buf.String() != "hello world" { + t.Fatalf("output=%q, want %q", buf.String(), "hello world") + } +} + +func TestPtyCompressedFrameError(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", connectContentType) + w.Write(connectEnvelope(0, `{"event":{"start":{"pid":9}}}`)) + w.Write(connectEnvelope(connectCompressedFlag, `{}`)) + })) + defer server.Close() + + sb := newPtyTestSandbox(t, server) + handle, err := sb.Pty().Create(context.Background(), PtySize{Rows: 24, Cols: 80}, PtyCreateOptions{}) + if err != nil { + t.Fatalf("Create: %v", err) + } + _, err = handle.Wait(nil) + if err == nil || !strings.Contains(err.Error(), "compressed") { + t.Fatalf("Wait error=%v, want a compressed-frame error", err) + } +} + +func TestPtyEndStreamTrailerError(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", connectContentType) + w.Write(connectEnvelope(0, `{"event":{"start":{"pid":10}}}`)) + w.Write(connectEnvelope(connectEndStreamFlag, `{"error":{"code":"internal","message":"boom"}}`)) + })) + defer server.Close() + + sb := newPtyTestSandbox(t, server) + handle, err := sb.Pty().Create(context.Background(), PtySize{Rows: 24, Cols: 80}, PtyCreateOptions{}) + if err != nil { + t.Fatalf("Create: %v", err) + } + _, err = handle.Wait(nil) + if err == nil || !strings.Contains(err.Error(), "boom") { + t.Fatalf("Wait error=%v, want the trailer error 'boom'", err) + } +} + +func TestPtyBase64DecodeError(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", connectContentType) + w.Write(connectEnvelope(0, `{"event":{"start":{"pid":11}}}`)) + w.Write(connectEnvelope(0, `{"event":{"data":{"pty":"@@not-base64@@"}}}`)) + })) + defer server.Close() + + sb := newPtyTestSandbox(t, server) + handle, err := sb.Pty().Create(context.Background(), PtySize{Rows: 24, Cols: 80}, PtyCreateOptions{}) + if err != nil { + t.Fatalf("Create: %v", err) + } + _, err = handle.Wait(nil) + if err == nil || !strings.Contains(err.Error(), "decode pty output") { + t.Fatalf("Wait error=%v, want a base64 decode error", err) + } +} + +func TestPtyStartStreamClosedBeforeStart(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", connectContentType) + // End the stream immediately, before any start event. + w.Write(connectEnvelope(connectEndStreamFlag, `{}`)) + })) + defer server.Close() + + sb := newPtyTestSandbox(t, server) + _, err := sb.Pty().Create(context.Background(), PtySize{Rows: 1, Cols: 1}, PtyCreateOptions{}) + if err == nil || !strings.Contains(err.Error(), "stream closed before start event") { + t.Fatalf("Create error=%v, want 'stream closed before start event'", err) + } +} + +func TestPtyConnectHTTPError(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + http.Error(w, `{"message":"nope"}`, http.StatusInternalServerError) + })) + defer server.Close() + + sb := newPtyTestSandbox(t, server) + _, err := sb.Pty().Connect(context.Background(), 1, PtyConnectOptions{}) + var apiErr *APIError + if err == nil || !errors.As(err, &apiErr) || apiErr.StatusCode != http.StatusInternalServerError { + t.Fatalf("Connect error=%v, want APIError 500", err) + } +} + +func TestPtySendStdinAndResizeErrors(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + http.Error(w, `{"message":"boom"}`, http.StatusInternalServerError) + })) + defer server.Close() + + sb := newPtyTestSandbox(t, server) + if err := sb.Pty().SendStdin(context.Background(), 5, []byte("x")); err == nil { + t.Fatal("SendStdin returned nil error on 500") + } + if err := sb.Pty().Resize(context.Background(), 5, PtySize{Rows: 1, Cols: 1}); err == nil { + t.Fatal("Resize returned nil error on 500") + } +} + +// TestPtyHandleKillResizeDelegation verifies the per-handle Kill and Resize +// shortcuts target the handle's own PID (SendStdin is covered separately). +func TestPtyHandleKillResizeDelegation(t *testing.T) { + var killPID, resizePID int + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body, _ := io.ReadAll(r.Body) + switch r.URL.Path { + case "/process.Process/Start": + w.Header().Set("Content-Type", connectContentType) + w.Write(connectEnvelope(0, `{"event":{"start":{"pid":271}}}`)) + w.Write(connectEnvelope(0, `{"event":{"end":{"exitCode":0,"exited":true}}}`)) + w.Write(connectEnvelope(connectEndStreamFlag, `{}`)) + case "/process.Process/SendSignal": + var req ptySignalRequest + json.Unmarshal(body, &req) + killPID = req.Process.PID + w.Write([]byte(`{}`)) + case "/process.Process/Update": + var req ptyUpdateRequest + json.Unmarshal(body, &req) + resizePID = req.Process.PID + w.WriteHeader(http.StatusOK) + default: + w.WriteHeader(http.StatusOK) + } + })) + defer server.Close() + + sb := newPtyTestSandbox(t, server) + handle, err := sb.Pty().Create(context.Background(), PtySize{Rows: 24, Cols: 80}, PtyCreateOptions{}) + if err != nil { + t.Fatalf("Create: %v", err) + } + if _, err := handle.Kill(context.Background()); err != nil { + t.Fatalf("handle.Kill: %v", err) + } + if err := handle.Resize(context.Background(), PtySize{Rows: 40, Cols: 120}); err != nil { + t.Fatalf("handle.Resize: %v", err) + } + if killPID != handle.PID() || resizePID != handle.PID() { + t.Fatalf("delegated PIDs kill=%d resize=%d, want %d", killPID, resizePID, handle.PID()) + } + if _, err := handle.Wait(nil); err != nil { + t.Fatalf("Wait: %v", err) + } +} + +// TestPtyContextCancellation verifies cancelling the caller's context tears the +// stream down and unblocks Wait. +func TestPtyContextCancellation(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", connectContentType) + w.Write(connectEnvelope(0, `{"event":{"start":{"pid":12}}}`)) + if f, ok := w.(http.Flusher); ok { + f.Flush() + } + select { + case <-r.Context().Done(): + case <-time.After(3 * time.Second): + } + })) + defer server.Close() + + sb := newPtyTestSandbox(t, server) + ctx, cancel := context.WithCancel(context.Background()) + handle, err := sb.Pty().Create(ctx, PtySize{Rows: 24, Cols: 80}, PtyCreateOptions{}) + if err != nil { + t.Fatalf("Create: %v", err) + } + + done := make(chan struct{}) + go func() { + handle.Wait(nil) + close(done) + }() + cancel() + select { + case <-done: + case <-time.After(3 * time.Second): + t.Fatal("Wait did not return after context cancellation") + } +}