diff --git a/server/cmd/api/api/process.go b/server/cmd/api/api/process.go index 38a458f4..b9b6ebb8 100644 --- a/server/cmd/api/api/process.go +++ b/server/cmd/api/api/process.go @@ -30,6 +30,11 @@ import ( openapi_types "github.com/oapi-codegen/runtime/types" ) +// processOutputDrainTimeout bounds how long a spawned process's exit is held +// back while its output readers drain the pipes. A child the process started +// in the background can keep the pipes open after the process itself exits. +const processOutputDrainTimeout = 2 * time.Second + type processHandle struct { id openapi_types.UUID pid int @@ -298,19 +303,12 @@ func (s *ApiService) ProcessSpawn(ctx context.Context, request oapi.ProcessSpawn stdin = ptyFile isTTY = true } else { - stdout, err = cmd.StdoutPipe() - if err != nil { - return oapi.ProcessSpawn500JSONResponse{InternalErrorJSONResponse: oapi.InternalErrorJSONResponse{Message: "failed to open stdout"}}, nil - } - stderr, err = cmd.StderrPipe() - if err != nil { - return oapi.ProcessSpawn500JSONResponse{InternalErrorJSONResponse: oapi.InternalErrorJSONResponse{Message: "failed to open stderr"}}, nil - } stdin, err = cmd.StdinPipe() if err != nil { return oapi.ProcessSpawn500JSONResponse{InternalErrorJSONResponse: oapi.InternalErrorJSONResponse{Message: "failed to open stdin"}}, nil } - if err := cmd.Start(); err != nil { + stdout, stderr, err = startWithOutputPipes(cmd) + if err != nil { if isUserCmdError(err) { return oapi.ProcessSpawn400JSONResponse{BadRequestErrorJSONResponse: oapi.BadRequestErrorJSONResponse{Message: err.Error()}}, nil } @@ -354,8 +352,11 @@ func (s *ApiService) ProcessSpawn(ctx context.Context, request oapi.ProcessSpawn // Reader goroutines // In PTY mode, do NOT read from the PTY here to avoid competing with the /attach endpoint. // In non‑PTY mode, stdout and stderr are separate pipes, so we run two readers and tag chunks accordingly. + var readers sync.WaitGroup if !isTTY { + readers.Add(2) go func() { + defer readers.Done() reader := bufio.NewReader(stdout) buf := make([]byte, 4096) for { @@ -371,6 +372,7 @@ func (s *ApiService) ProcessSpawn(ctx context.Context, request oapi.ProcessSpawn } }() go func() { + defer readers.Done() reader := bufio.NewReader(stderr) buf := make([]byte, 4096) for { @@ -405,6 +407,17 @@ func (s *ApiService) ProcessSpawn(ctx context.Context, request oapi.ProcessSpawn } } h.setExited(code) + // Let the readers drain what the process wrote before it exited, so + // the exit event follows all of its output. + readersDone := make(chan struct{}) + go func() { + readers.Wait() + close(readersDone) + }() + select { + case <-readersDone: + case <-time.After(processOutputDrainTimeout): + } // Ensure all related FDs are closed to avoid leaking descriptors. // In PTY mode, close the PTY master; in non-PTY mode, close individual pipes. if h.isTTY { @@ -431,6 +444,11 @@ func (s *ApiService) ProcessSpawn(ctx context.Context, request oapi.ProcessSpawn } } + // Closing the pipes above ends any reader still reading them. A reader + // blocked on a full outCh finishes once a stream takes its event, just + // as the exit event below would have waited for one. + <-readersDone + // Send exit event evt := oapi.ProcessStreamEventEvent("exit") h.outCh <- oapi.ProcessStreamEvent{Event: &evt, ExitCode: &code} @@ -457,6 +475,36 @@ func (s *ApiService) ProcessSpawn(ctx context.Context, request oapi.ProcessSpawn }, nil } +// startWithOutputPipes starts cmd with stdout and stderr connected to pipes +// whose read ends the caller owns and must close. cmd.StdoutPipe is not used +// because cmd.Wait closes its read end as soon as the process exits, which +// discards output the readers have not consumed yet. +func startWithOutputPipes(cmd *exec.Cmd) (stdout, stderr *os.File, err error) { + stdoutR, stdoutW, err := os.Pipe() + if err != nil { + return nil, nil, err + } + stderrR, stderrW, err := os.Pipe() + if err != nil { + stdoutR.Close() + stdoutW.Close() + return nil, nil, err + } + cmd.Stdout = stdoutW + cmd.Stderr = stderrW + err = cmd.Start() + // The child holds its own copies of the write ends; closing ours lets the + // readers see EOF once the child, and anything it started, closes them. + stdoutW.Close() + stderrW.Close() + if err != nil { + stdoutR.Close() + stderrR.Close() + return nil, nil, err + } + return stdoutR, stderrR, nil +} + // Send signal to process // (POST /process/{process_id}/kill) func (s *ApiService) ProcessKill(ctx context.Context, request oapi.ProcessKillRequestObject) (oapi.ProcessKillResponseObject, error) { @@ -581,27 +629,42 @@ func (s *ApiService) ProcessStdoutStream(ctx context.Context, request oapi.Proce pr, pw := io.Pipe() go func() { defer pw.Close() + // writeEvent writes evt as an SSE frame: data: \n\n + writeEvent := func(evt oapi.ProcessStreamEvent) error { + var buf bytes.Buffer + if err := json.NewEncoder(&buf).Encode(evt); err != nil { + log.Error("failed to marshal event", "err", err) + return err + } + line := bytes.TrimRight(buf.Bytes(), "\n") + if _, err := pw.Write([]byte("data: ")); err != nil { + return err + } + if _, err := pw.Write(line); err != nil { + return err + } + _, err := pw.Write([]byte("\n\n")) + return err + } for { select { case evt := <-h.outCh: - // Write SSE: data: \n\n - var buf bytes.Buffer - if err := json.NewEncoder(&buf).Encode(evt); err != nil { - log.Error("failed to marshal event", "err", err) - return - } - line := bytes.TrimRight(buf.Bytes(), "\n") - if _, err := pw.Write([]byte("data: ")); err != nil { - return - } - if _, err := pw.Write(line); err != nil { - return - } - if _, err := pw.Write([]byte("\n\n")); err != nil { + if writeEvent(evt) != nil { return } case <-h.doneCh: - return + // The exit event and any output before it were queued before + // doneCh closed; deliver them before ending the stream. + for { + select { + case evt := <-h.outCh: + if writeEvent(evt) != nil { + return + } + default: + return + } + } } } }() diff --git a/server/cmd/api/api/process_test.go b/server/cmd/api/api/process_test.go index 741734d7..06ff8226 100644 --- a/server/cmd/api/api/process_test.go +++ b/server/cmd/api/api/process_test.go @@ -9,6 +9,7 @@ import ( "os/user" "strconv" "strings" + "syscall" "testing" "time" @@ -271,3 +272,138 @@ func TestBuildCmd_AsRootAndAsUserConflict(t *testing.T) { _, err := buildCmd(body) require.Error(t, err, "expected error when both as_root and as_user are set") } + +// spawnShForTest spawns `sh -c script` and returns its process ID. +func spawnShForTest(t *testing.T, svc *ApiService, script string) openapi_types.UUID { + t.Helper() + args := []string{"-c", script} + resp, err := svc.ProcessSpawn(context.Background(), oapi.ProcessSpawnRequestObject{Body: &oapi.ProcessSpawnRequest{Command: "sh", Args: &args}}) + require.NoError(t, err, "ProcessSpawn error") + s200, ok := resp.(oapi.ProcessSpawn200JSONResponse) + require.True(t, ok, "unexpected spawn resp type: %T", resp) + require.NotNil(t, s200.ProcessId, "missing ProcessId in spawn resp") + return *s200.ProcessId +} + +// readProcessStream reads a process's SSE stream until it ends and returns +// its stdout, its stderr and every event in the order received. +func readProcessStream(t *testing.T, svc *ApiService, id openapi_types.UUID) (string, string, []oapi.ProcessStreamEvent) { + t.Helper() + resp, err := svc.ProcessStdoutStream(context.Background(), oapi.ProcessStdoutStreamRequestObject{ProcessId: id}) + require.NoError(t, err, "StdoutStream error") + st200, ok := resp.(oapi.ProcessStdoutStream200TexteventStreamResponse) + require.True(t, ok, "unexpected stream resp: %T", resp) + + type result struct { + stdout, stderr strings.Builder + events []oapi.ProcessStreamEvent + err error + } + done := make(chan *result, 1) + go func() { + r := &result{} + scanner := bufio.NewScanner(st200.Body) + for scanner.Scan() { + line := scanner.Text() + if !strings.HasPrefix(line, "data: ") { + continue + } + var evt oapi.ProcessStreamEvent + if r.err = json.Unmarshal([]byte(strings.TrimPrefix(line, "data: ")), &evt); r.err != nil { + break + } + r.events = append(r.events, evt) + if evt.DataB64 != nil { + b, _ := base64.StdEncoding.DecodeString(*evt.DataB64) + if evt.Stream != nil && *evt.Stream == "stderr" { + r.stderr.Write(b) + } else { + r.stdout.Write(b) + } + } + } + if r.err == nil { + r.err = scanner.Err() + } + done <- r + }() + select { + case r := <-done: + require.NoError(t, r.err, "read SSE stream") + return r.stdout.String(), r.stderr.String(), r.events + case <-time.After(10 * time.Second): + st200.Body.(io.Closer).Close() + t.Fatal("stream did not end within 10s") + return "", "", nil + } +} + +// requireExitLast checks that the stream ended with exactly one exit event. +func requireExitLast(t *testing.T, events []oapi.ProcessStreamEvent, wantCode int) { + t.Helper() + require.NotEmpty(t, events, "stream delivered no events") + last := events[len(events)-1] + require.True(t, last.Event != nil && *last.Event == "exit", "last event is not exit: %+v", last) + require.NotNil(t, last.ExitCode, "exit event has no exit code") + require.Equal(t, wantCode, *last.ExitCode, "exit code mismatch") + for _, evt := range events[:len(events)-1] { + require.Nil(t, evt.Event, "unexpected event before exit: %+v", evt) + } +} + +func TestProcessStdoutStream_AfterExitDeliversQueuedEvents(t *testing.T) { + t.Parallel() + svc := &ApiService{procs: make(map[string]*processHandle), stz: scaletozero.NewNoopController()} + + // Opening the stream only after the process has exited is the common case + // for short commands. Everything queued for it must still be delivered. + for i := 0; i < 200; i++ { + id := spawnShForTest(t, svc, "printf one; printf two 1>&2; exit 7") + svc.procMu.RLock() + h := svc.procs[id.String()] + svc.procMu.RUnlock() + <-h.doneCh + + stdout, stderr, events := readProcessStream(t, svc, id) + require.Equal(t, "one", stdout, "run %d: stdout mismatch", i) + require.Equal(t, "two", stderr, "run %d: stderr mismatch", i) + requireExitLast(t, events, 7) + } +} + +func TestProcessStdoutStream_DeliversAllOutputBeforeExit(t *testing.T) { + t.Parallel() + svc := &ApiService{procs: make(map[string]*processHandle), stz: scaletozero.NewNoopController()} + + // The process exits right after writing more than a pipe buffer, so output + // is still unread in the pipe when it exits. + want := strings.Repeat("kernel\n", 50000) + for i := 0; i < 5; i++ { + id := spawnShForTest(t, svc, "yes kernel | head -n 50000; exit 3") + stdout, stderr, events := readProcessStream(t, svc, id) + require.Equal(t, len(want), len(stdout), "run %d: stdout length mismatch", i) + require.Equal(t, want, stdout, "run %d: stdout mismatch", i) + require.Empty(t, stderr, "run %d: unexpected stderr", i) + requireExitLast(t, events, 3) + } +} + +func TestProcessSpawn_BackgroundChildDoesNotHoldExit(t *testing.T) { + t.Parallel() + svc := &ApiService{procs: make(map[string]*processHandle), stz: scaletozero.NewNoopController()} + + // The shell exits at once, but the background sleep keeps the output + // pipes open. The stream must still end with the shell's exit. + id := spawnShForTest(t, svc, "sleep 30 & echo $!") + stdout, _, events := readProcessStream(t, svc, id) + pid, err := strconv.Atoi(strings.TrimSpace(stdout)) + require.NoError(t, err, "background pid not printed: %q", stdout) + t.Cleanup(func() { _ = syscall.Kill(pid, syscall.SIGKILL) }) + requireExitLast(t, events, 0) + + statusResp, err := svc.ProcessStatus(context.Background(), oapi.ProcessStatusRequestObject{ProcessId: id}) + require.NoError(t, err, "ProcessStatus error") + sr, ok := statusResp.(oapi.ProcessStatus200JSONResponse) + require.True(t, ok, "unexpected status resp: %T", statusResp) + require.Equal(t, oapi.ProcessStatusState("exited"), *sr.State) +}