You've already forked agentic-coding-workflow
TTY
This commit is contained in:
@@ -0,0 +1,220 @@
|
||||
package terminal
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"strconv"
|
||||
"time"
|
||||
|
||||
"github.com/go-chi/chi/v5"
|
||||
"github.com/google/uuid"
|
||||
|
||||
ws "github.com/coder/websocket"
|
||||
|
||||
"github.com/yan1h/agent-coding-workflow/internal/audit"
|
||||
"github.com/yan1h/agent-coding-workflow/internal/infra/errs"
|
||||
httpx "github.com/yan1h/agent-coding-workflow/internal/transport/http"
|
||||
"github.com/yan1h/agent-coding-workflow/internal/transport/http/middleware"
|
||||
)
|
||||
|
||||
// readBufSize 是从 shell 读输出的缓冲块大小。
|
||||
const readBufSize = 32 * 1024
|
||||
|
||||
// wsReadLimit 放宽单条消息上限以容纳较大的粘贴输入（默认 32KB 对粘贴偏小）。
|
||||
const wsReadLimit = 1 << 20
|
||||
|
||||
// AdminLookup 解析用户是否为管理员（与 run/acp 模块同一模式）。
|
||||
type AdminLookup interface {
|
||||
IsAdmin(ctx context.Context, userID uuid.UUID) (bool, error)
|
||||
}
|
||||
|
||||
// WSConfig 控制 WebSocket 心跳。
|
||||
type WSConfig struct {
|
||||
PingInterval time.Duration
|
||||
PongTimeout time.Duration
|
||||
}
|
||||
|
||||
// Handler 暴露仅限管理员的终端 WS 端点。
|
||||
type Handler struct {
|
||||
sup *Supervisor
|
||||
resolver middleware.SessionResolver
|
||||
adminLookup AdminLookup
|
||||
audit audit.Recorder
|
||||
shell string
|
||||
cfg WSConfig
|
||||
}
|
||||
|
||||
// NewHandler 构造 terminal.Handler。
|
||||
func NewHandler(sup *Supervisor, resolver middleware.SessionResolver, al AdminLookup, rec audit.Recorder, shell string, cfg WSConfig) *Handler {
|
||||
return &Handler{sup: sup, resolver: resolver, adminLookup: al, audit: rec, shell: shell, cfg: cfg}
|
||||
}
|
||||
|
||||
// Mount 注册终端路由（仅 admin）。WS 鉴权经 middleware.Auth（浏览器原生 WebSocket
|
||||
// 无法带自定义 header，token 走 ?token= 查询参数，沿用 run/acp 的 WS 约定）。
|
||||
func (h *Handler) Mount(r chi.Router) {
|
||||
r.Route("/api/v1/admin/terminal", func(r chi.Router) {
|
||||
r.Use(middleware.Auth(h.resolver, middleware.AuthOptions{}))
|
||||
r.Get("/stream", h.stream)
|
||||
})
|
||||
}
|
||||
|
||||
// requireAdmin 从请求解析用户并要求其为管理员，否则返回 forbidden。
|
||||
func (h *Handler) requireAdmin(r *http.Request) (uuid.UUID, error) {
|
||||
uid, ok := middleware.UserIDFromContext(r.Context())
|
||||
if !ok {
|
||||
return uuid.Nil, errs.New(errs.CodeUnauthorized, "not authenticated")
|
||||
}
|
||||
admin, err := h.adminLookup.IsAdmin(r.Context(), uid)
|
||||
if err != nil {
|
||||
return uuid.Nil, err
|
||||
}
|
||||
if !admin {
|
||||
return uuid.Nil, errs.New(errs.CodeForbidden, "admin only")
|
||||
}
|
||||
return uid, nil
|
||||
}
|
||||
|
||||
// resizeMsg 是客户端经文本帧上报的窗口尺寸控制消息。
|
||||
type resizeMsg struct {
|
||||
Type string `json:"type"`
|
||||
Rows uint16 `json:"rows"`
|
||||
Cols uint16 `json:"cols"`
|
||||
}
|
||||
|
||||
// stream 升级为 WebSocket，开一个 shell 会话并双向桥接：
|
||||
// - server→client：shell 输出作为二进制帧推送
|
||||
// - client→server：二进制帧 = stdin 原始字节；文本帧 = JSON 控制（resize）
|
||||
//
|
||||
// 二进制/文本帧天然区分 stdin 与控制消息，无需额外前缀字节。
|
||||
func (h *Handler) stream(w http.ResponseWriter, r *http.Request) {
|
||||
uid, err := h.requireAdmin(r)
|
||||
if err != nil {
|
||||
writeErr(w, r, err)
|
||||
return
|
||||
}
|
||||
rows := uint16(parseUint(r.URL.Query().Get("rows"), uint64(defaultRows)))
|
||||
cols := uint16(parseUint(r.URL.Query().Get("cols"), uint64(defaultCols)))
|
||||
|
||||
conn, aerr := ws.Accept(w, r, &ws.AcceptOptions{})
|
||||
if aerr != nil {
|
||||
return
|
||||
}
|
||||
defer conn.Close(ws.StatusInternalError, "internal")
|
||||
conn.SetReadLimit(wsReadLimit)
|
||||
|
||||
sess, oerr := h.sup.Open(OpenParams{Rows: rows, Cols: cols})
|
||||
if oerr != nil {
|
||||
conn.Close(ws.StatusTryAgainLater, "cannot open terminal session")
|
||||
return
|
||||
}
|
||||
|
||||
h.record(r.Context(), r.RemoteAddr, uid, "admin.terminal.opened", sess.ID)
|
||||
defer func() {
|
||||
// 连接已断时 r.Context() 多半已取消，审计与终止改用独立 context。
|
||||
closeCtx, cancel := context.WithTimeout(context.Background(), h.cfg.PongTimeout)
|
||||
defer cancel()
|
||||
h.sup.Close(closeCtx, sess.ID)
|
||||
h.record(closeCtx, r.RemoteAddr, uid, "admin.terminal.closed", sess.ID)
|
||||
}()
|
||||
|
||||
ctx, cancel := context.WithCancel(r.Context())
|
||||
defer cancel()
|
||||
|
||||
// shell 进程自行退出（如用户输入 exit）→ 取消，触发主循环退出。
|
||||
go func() {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
case <-sess.Done():
|
||||
cancel()
|
||||
}
|
||||
}()
|
||||
|
||||
// read pump：shell 输出 → WS 二进制帧。
|
||||
go func() {
|
||||
buf := make([]byte, readBufSize)
|
||||
for {
|
||||
n, rerr := sess.Read(buf)
|
||||
if n > 0 {
|
||||
if werr := conn.Write(ctx, ws.MessageBinary, buf[:n]); werr != nil {
|
||||
cancel()
|
||||
return
|
||||
}
|
||||
}
|
||||
if rerr != nil {
|
||||
cancel()
|
||||
return
|
||||
}
|
||||
}
|
||||
}()
|
||||
|
||||
// 心跳。
|
||||
pingTicker := time.NewTicker(h.cfg.PingInterval)
|
||||
defer pingTicker.Stop()
|
||||
go func() {
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case <-pingTicker.C:
|
||||
pingCtx, pcancel := context.WithTimeout(ctx, h.cfg.PongTimeout)
|
||||
perr := conn.Ping(pingCtx)
|
||||
pcancel()
|
||||
if perr != nil {
|
||||
cancel()
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}()
|
||||
|
||||
// write pump（主循环）：WS → shell stdin / resize。
|
||||
for {
|
||||
typ, data, rerr := conn.Read(ctx)
|
||||
if rerr != nil {
|
||||
return
|
||||
}
|
||||
switch typ {
|
||||
case ws.MessageBinary:
|
||||
if _, werr := sess.Write(data); werr != nil {
|
||||
return
|
||||
}
|
||||
case ws.MessageText:
|
||||
var m resizeMsg
|
||||
if json.Unmarshal(data, &m) == nil && m.Type == "resize" {
|
||||
_ = sess.Resize(m.Rows, m.Cols)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// record 写一条审计日志，吞错（审计失败不应影响终端会话）。
|
||||
func (h *Handler) record(ctx context.Context, remoteAddr string, uid uuid.UUID, action string, sessID uuid.UUID) {
|
||||
if h.audit == nil {
|
||||
return
|
||||
}
|
||||
u := uid
|
||||
_ = h.audit.Record(ctx, audit.Entry{
|
||||
UserID: &u,
|
||||
Action: action,
|
||||
TargetType: "terminal_session",
|
||||
TargetID: sessID.String(),
|
||||
IP: remoteAddr,
|
||||
Metadata: map[string]any{"shell": h.shell},
|
||||
})
|
||||
}
|
||||
|
||||
func parseUint(s string, def uint64) uint64 {
|
||||
if s == "" {
|
||||
return def
|
||||
}
|
||||
n, err := strconv.ParseUint(s, 10, 16)
|
||||
if err != nil {
|
||||
return def
|
||||
}
|
||||
return n
|
||||
}
|
||||
|
||||
func writeErr(w http.ResponseWriter, r *http.Request, err error) {
|
||||
httpx.WriteError(w, middleware.RequestIDFromContext(r.Context()), err)
|
||||
}
|
||||
@@ -0,0 +1,73 @@
|
||||
package terminal
|
||||
|
||||
import (
|
||||
"io"
|
||||
"os"
|
||||
"os/exec"
|
||||
"strings"
|
||||
|
||||
procgrp "github.com/yan1h/agent-coding-workflow/internal/infra/proc"
|
||||
)
|
||||
|
||||
// 终端窗口的兜底尺寸（前端尚未上报 resize 前使用）。
|
||||
const (
|
||||
defaultRows uint16 = 24
|
||||
defaultCols uint16 = 80
|
||||
)
|
||||
|
||||
// ptySession 持有平台相关的 shell 进程与 IO。resizeFn 由 startPTY 注入：
|
||||
// Unix 调 pty.Setsize，Windows 为 nil（降级实现无 resize）。这样 ptySession
|
||||
// 类型与生命周期方法可保持平台无关，仅 startPTY 按 build tag 特化。
|
||||
type ptySession struct {
|
||||
cmd *exec.Cmd
|
||||
group procgrp.Group
|
||||
rw io.ReadWriteCloser // PTY master（Unix）或 stdio 桥接（Windows）
|
||||
resizeFn func(rows, cols uint16) error
|
||||
}
|
||||
|
||||
func (p *ptySession) resize(rows, cols uint16) error {
|
||||
if p.resizeFn == nil {
|
||||
return nil
|
||||
}
|
||||
return p.resizeFn(rows, cols)
|
||||
}
|
||||
|
||||
// terminate 是 Open 失败路径的尽力清理（kill 进程、释放句柄、关闭 IO）。
|
||||
// 正常关闭走 Supervisor.Close → group.TerminateGroup 的整树终止。
|
||||
func (p *ptySession) terminate() {
|
||||
if p.cmd.Process != nil {
|
||||
_ = p.cmd.Process.Kill()
|
||||
}
|
||||
p.group.Close()
|
||||
_ = p.rw.Close()
|
||||
}
|
||||
|
||||
// buildEnv 从宿主环境按白名单过滤出传给 shell 的环境变量，并补上 TERM 以便
|
||||
// 终端程序正确渲染。白名单确保 APP_MASTER_KEY、DB DSN、git 凭据等敏感变量
|
||||
// 不会透传给交互式 shell。
|
||||
func buildEnv(whitelist []string) []string {
|
||||
allow := make(map[string]bool, len(whitelist))
|
||||
for _, k := range whitelist {
|
||||
allow[k] = true
|
||||
}
|
||||
out := make([]string, 0, len(whitelist)+1)
|
||||
hasTerm := false
|
||||
for _, kv := range os.Environ() {
|
||||
i := strings.IndexByte(kv, '=')
|
||||
if i < 0 {
|
||||
continue
|
||||
}
|
||||
name := kv[:i]
|
||||
if !allow[name] {
|
||||
continue
|
||||
}
|
||||
if name == "TERM" {
|
||||
hasTerm = true
|
||||
}
|
||||
out = append(out, kv)
|
||||
}
|
||||
if !hasTerm {
|
||||
out = append(out, "TERM=xterm-256color")
|
||||
}
|
||||
return out
|
||||
}
|
||||
@@ -0,0 +1,46 @@
|
||||
//go:build !windows
|
||||
|
||||
package terminal
|
||||
|
||||
import (
|
||||
"os/exec"
|
||||
|
||||
"github.com/creack/pty"
|
||||
|
||||
procgrp "github.com/yan1h/agent-coding-workflow/internal/infra/proc"
|
||||
)
|
||||
|
||||
// startPTY 在真 PTY 中拉起 shell（Unix）。
|
||||
//
|
||||
// 关于进程组：pty.StartWithSize 内部会设置 SysProcAttr.Setsid，使 shell 成为新
|
||||
// 会话的首进程，其进程组 pgid == pid。因此这里**不**调用 group.Prepare（它设的
|
||||
// Setpgid 与 Setsid 同时设置在 Go exec 下是非法/冗余的）；group.Adopt 直接记录
|
||||
// pgid = cmd.Process.Pid，TerminateGroup 向 -pgid 发信号即可整树终止 shell 及其
|
||||
// 派生的全部子孙进程。
|
||||
func startPTY(shell string, env []string, rows, cols uint16) (*ptySession, error) {
|
||||
cmd := exec.Command(shell)
|
||||
cmd.Env = env
|
||||
|
||||
ptmx, err := pty.StartWithSize(cmd, &pty.Winsize{Rows: rows, Cols: cols})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
group := procgrp.NewGroup()
|
||||
if err := group.Adopt(cmd); err != nil {
|
||||
_ = ptmx.Close()
|
||||
if cmd.Process != nil {
|
||||
_ = cmd.Process.Kill()
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return &ptySession{
|
||||
cmd: cmd,
|
||||
group: group,
|
||||
rw: ptmx,
|
||||
resizeFn: func(r, c uint16) error {
|
||||
return pty.Setsize(ptmx, &pty.Winsize{Rows: r, Cols: c})
|
||||
},
|
||||
}, nil
|
||||
}
|
||||
@@ -0,0 +1,75 @@
|
||||
//go:build windows
|
||||
|
||||
package terminal
|
||||
|
||||
import (
|
||||
"io"
|
||||
"os"
|
||||
"os/exec"
|
||||
|
||||
procgrp "github.com/yan1h/agent-coding-workflow/internal/infra/proc"
|
||||
)
|
||||
|
||||
// stdioBridge 把 shell 的 stdin（写端）与合并后的 stdout/stderr（读端）合成一个
|
||||
// io.ReadWriteCloser，供降级实现使用。
|
||||
type stdioBridge struct {
|
||||
stdin io.WriteCloser
|
||||
stdout io.ReadCloser
|
||||
}
|
||||
|
||||
func (b *stdioBridge) Read(p []byte) (int, error) { return b.stdout.Read(p) }
|
||||
func (b *stdioBridge) Write(p []byte) (int, error) { return b.stdin.Write(p) }
|
||||
func (b *stdioBridge) Close() error {
|
||||
_ = b.stdin.Close()
|
||||
return b.stdout.Close()
|
||||
}
|
||||
|
||||
// startPTY 是 Windows 降级实现：用普通管道桥接 stdin/stdout，无 PTY、无 resize，
|
||||
// 仅供本地开发调试，不追求生产可用（无行编辑/光标控制/作业控制）。
|
||||
//
|
||||
// stdout/stderr 合并：手动建 os.Pipe，把写端同时交给子进程的 Stdout/Stderr，父进程
|
||||
// Start 后立即关闭自己持有的写端；子进程退出后读端 pr 收到 EOF，read pump 据此感知结束。
|
||||
func startPTY(shell string, env []string, rows, cols uint16) (*ptySession, error) {
|
||||
cmd := exec.Command(shell)
|
||||
cmd.Env = env
|
||||
|
||||
pr, pw, err := os.Pipe()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
cmd.Stdout = pw
|
||||
cmd.Stderr = pw
|
||||
|
||||
stdin, err := cmd.StdinPipe()
|
||||
if err != nil {
|
||||
_ = pr.Close()
|
||||
_ = pw.Close()
|
||||
return nil, err
|
||||
}
|
||||
|
||||
group := procgrp.NewGroup()
|
||||
group.Prepare(cmd) // Windows: no-op
|
||||
if err := cmd.Start(); err != nil {
|
||||
_ = pr.Close()
|
||||
_ = pw.Close()
|
||||
_ = stdin.Close()
|
||||
return nil, err
|
||||
}
|
||||
_ = pw.Close() // 父进程放掉写端；子进程仍持有，退出后 pr 读到 EOF
|
||||
|
||||
if err := group.Adopt(cmd); err != nil {
|
||||
if cmd.Process != nil {
|
||||
_ = cmd.Process.Kill()
|
||||
}
|
||||
_ = pr.Close()
|
||||
_ = stdin.Close()
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return &ptySession{
|
||||
cmd: cmd,
|
||||
group: group,
|
||||
rw: &stdioBridge{stdin: stdin, stdout: pr},
|
||||
resizeFn: nil, // 降级实现不支持 resize
|
||||
}, nil
|
||||
}
|
||||
@@ -0,0 +1,185 @@
|
||||
// Package terminal 提供一个仅限管理员的网页端交互式终端：在 server 宿主进程
|
||||
// （部署时即 runtime 容器）内拉起一个 shell，经 WebSocket 双向桥接 stdin/stdout，
|
||||
// 主要用途是运维（例如容器内 npm i -g 安装 ACP 客户端 CLI）。
|
||||
//
|
||||
// 会话纯内存、不落库：开即用、断即清，server 重启即全部销毁（沿用 run profiles
|
||||
// 「运行态不持久化」的取舍）。进程树终止复用 internal/infra/proc，断连时整树清理，
|
||||
// 避免 shell 派生的子孙进程（npm→node…）残留为孤儿。
|
||||
//
|
||||
// 跨平台：Unix 走真 PTY（github.com/creack/pty）；Windows 走降级的 stdio 管道
|
||||
// 实现（无 PTY、无 resize），仅供本地开发调试，不追求生产可用。
|
||||
package terminal
|
||||
|
||||
import (
|
||||
"context"
|
||||
"log/slog"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
// SupervisorConfig 是从 config.TerminalConfig 派生的 supervisor 子集。
|
||||
type SupervisorConfig struct {
|
||||
Shell string // 启动的 shell，如 /bin/bash
|
||||
MaxSessions int // 同时存在的会话上限
|
||||
KillGrace time.Duration // 终止整组时 SIGTERM→SIGKILL 的宽限
|
||||
// EnvWhitelist 是允许从宿主进程透传给 shell 的环境变量名白名单。
|
||||
// APP_MASTER_KEY、DB DSN、git 凭据等敏感变量绝不应出现在此名单中。
|
||||
EnvWhitelist []string
|
||||
}
|
||||
|
||||
// ErrTooManySessions 在达到 MaxSessions 上限时返回，由 Handler 翻译为 errs code。
|
||||
var ErrTooManySessions = errSentinel("terminal: too many active sessions")
|
||||
|
||||
type errSentinel string
|
||||
|
||||
func (e errSentinel) Error() string { return string(e) }
|
||||
|
||||
// Session 是单个终端会话的运行时句柄。pty 持有平台相关的 shell 进程与 IO，
|
||||
// done 在进程退出（monitor）时关闭。
|
||||
type Session struct {
|
||||
ID uuid.UUID
|
||||
pty *ptySession
|
||||
done chan struct{} // monitor 在 shell 退出后 close
|
||||
}
|
||||
|
||||
// Read 从 shell 读输出（PTY master / 管道读端）。
|
||||
func (s *Session) Read(p []byte) (int, error) { return s.pty.rw.Read(p) }
|
||||
|
||||
// Write 向 shell 写输入（PTY master / stdin 管道）。
|
||||
func (s *Session) Write(p []byte) (int, error) { return s.pty.rw.Write(p) }
|
||||
|
||||
// Resize 调整终端窗口尺寸（Windows 降级实现为 no-op）。
|
||||
func (s *Session) Resize(rows, cols uint16) error { return s.pty.resize(rows, cols) }
|
||||
|
||||
// Done 在 shell 进程退出后被关闭，供调用方感知进程自行结束（如用户输入 exit）。
|
||||
func (s *Session) Done() <-chan struct{} { return s.done }
|
||||
|
||||
// Supervisor 是终端会话总管。被动结构：Open 时起 monitor goroutine，无主循环。
|
||||
// 并发骨架照搬 run.RunSupervisor 已被测试验证的不变量（sessions map 写锁仅在
|
||||
// Open 插入 / monitor 删除时持有；done 严格 close-once，monitor 唯一持有）。
|
||||
type Supervisor struct {
|
||||
mu sync.RWMutex
|
||||
sessions map[uuid.UUID]*Session
|
||||
|
||||
cfg SupervisorConfig
|
||||
log *slog.Logger
|
||||
}
|
||||
|
||||
// NewSupervisor 构造空 Supervisor。
|
||||
func NewSupervisor(cfg SupervisorConfig, log *slog.Logger) *Supervisor {
|
||||
if log == nil {
|
||||
log = slog.Default()
|
||||
}
|
||||
return &Supervisor{
|
||||
sessions: map[uuid.UUID]*Session{},
|
||||
cfg: cfg,
|
||||
log: log,
|
||||
}
|
||||
}
|
||||
|
||||
// OpenParams 描述一次终端会话的初始参数。
|
||||
type OpenParams struct {
|
||||
Rows uint16
|
||||
Cols uint16
|
||||
}
|
||||
|
||||
// Open 启动一个 shell 会话并托管。达到 MaxSessions 上限时返回 ErrTooManySessions。
|
||||
func (s *Supervisor) Open(p OpenParams) (*Session, error) {
|
||||
s.mu.RLock()
|
||||
full := len(s.sessions) >= s.cfg.MaxSessions
|
||||
s.mu.RUnlock()
|
||||
if full {
|
||||
return nil, ErrTooManySessions
|
||||
}
|
||||
|
||||
rows, cols := p.Rows, p.Cols
|
||||
if rows == 0 {
|
||||
rows = defaultRows
|
||||
}
|
||||
if cols == 0 {
|
||||
cols = defaultCols
|
||||
}
|
||||
|
||||
ps, err := startPTY(s.cfg.Shell, buildEnv(s.cfg.EnvWhitelist), rows, cols)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
id, err := uuid.NewRandom()
|
||||
if err != nil {
|
||||
ps.terminate()
|
||||
return nil, err
|
||||
}
|
||||
sess := &Session{ID: id, pty: ps, done: make(chan struct{})}
|
||||
|
||||
s.mu.Lock()
|
||||
s.sessions[id] = sess
|
||||
s.mu.Unlock()
|
||||
|
||||
go s.monitor(sess)
|
||||
return sess, nil
|
||||
}
|
||||
|
||||
// monitor 阻塞等待 shell 退出后回收：关闭 done、移除映射、释放平台句柄。
|
||||
func (s *Supervisor) monitor(sess *Session) {
|
||||
_ = sess.pty.cmd.Wait() // reap，避免僵尸进程
|
||||
close(sess.done)
|
||||
s.mu.Lock()
|
||||
delete(s.sessions, sess.ID)
|
||||
s.mu.Unlock()
|
||||
sess.pty.group.Close()
|
||||
}
|
||||
|
||||
// Close 终止指定会话（整树）。幂等：会话已不存在时直接返回。阻塞直到 done 关闭
|
||||
// 或宽限耗尽。
|
||||
func (s *Supervisor) Close(ctx context.Context, id uuid.UUID) {
|
||||
s.mu.RLock()
|
||||
sess, ok := s.sessions[id]
|
||||
s.mu.RUnlock()
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
sess.pty.group.TerminateGroup(sess.done, s.cfg.KillGrace)
|
||||
select {
|
||||
case <-sess.done:
|
||||
case <-ctx.Done():
|
||||
}
|
||||
}
|
||||
|
||||
// Count 返回当前活跃会话数。
|
||||
func (s *Supervisor) Count() int {
|
||||
s.mu.RLock()
|
||||
defer s.mu.RUnlock()
|
||||
return len(s.sessions)
|
||||
}
|
||||
|
||||
// CloseAll 在服务关停时终止全部会话。
|
||||
func (s *Supervisor) CloseAll(ctx context.Context) {
|
||||
s.mu.RLock()
|
||||
ids := make([]uuid.UUID, 0, len(s.sessions))
|
||||
for id := range s.sessions {
|
||||
ids = append(ids, id)
|
||||
}
|
||||
s.mu.RUnlock()
|
||||
|
||||
var wg sync.WaitGroup
|
||||
for _, id := range ids {
|
||||
wg.Add(1)
|
||||
go func(sid uuid.UUID) {
|
||||
defer wg.Done()
|
||||
s.Close(ctx, sid)
|
||||
}(id)
|
||||
}
|
||||
|
||||
done := make(chan struct{})
|
||||
go func() {
|
||||
wg.Wait()
|
||||
close(done)
|
||||
}()
|
||||
select {
|
||||
case <-done:
|
||||
case <-ctx.Done():
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,178 @@
|
||||
package terminal
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"runtime"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
// testShell 返回各平台一个可交互读 stdin 的 shell。
|
||||
func testShell() string {
|
||||
if runtime.GOOS == "windows" {
|
||||
return "cmd"
|
||||
}
|
||||
return "/bin/sh"
|
||||
}
|
||||
|
||||
func newTestSup(t *testing.T, max int) *Supervisor {
|
||||
t.Helper()
|
||||
sup := NewSupervisor(SupervisorConfig{
|
||||
Shell: testShell(),
|
||||
MaxSessions: max,
|
||||
KillGrace: 2 * time.Second,
|
||||
EnvWhitelist: []string{
|
||||
"PATH", "HOME", "TERM",
|
||||
"SystemRoot", "TEMP", "TMP", "USERPROFILE", "ComSpec", "PATHEXT",
|
||||
},
|
||||
}, nil)
|
||||
t.Cleanup(func() {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
sup.CloseAll(ctx)
|
||||
})
|
||||
return sup
|
||||
}
|
||||
|
||||
// readUntil 在独立 goroutine 中持续读会话输出，直到累计内容包含 marker 或超时。
|
||||
// 阻塞的 Read goroutine 会在会话被关闭（defer/Cleanup）后随 Read 报错而退出。
|
||||
func readUntil(sess *Session, marker string, d time.Duration) string {
|
||||
out := make(chan []byte, 16)
|
||||
go func() {
|
||||
buf := make([]byte, 4096)
|
||||
for {
|
||||
n, err := sess.Read(buf)
|
||||
if n > 0 {
|
||||
b := make([]byte, n)
|
||||
copy(b, buf[:n])
|
||||
out <- b
|
||||
}
|
||||
if err != nil {
|
||||
close(out)
|
||||
return
|
||||
}
|
||||
}
|
||||
}()
|
||||
deadline := time.After(d)
|
||||
var acc strings.Builder
|
||||
for {
|
||||
select {
|
||||
case b, ok := <-out:
|
||||
if !ok {
|
||||
return acc.String()
|
||||
}
|
||||
acc.Write(b)
|
||||
if strings.Contains(acc.String(), marker) {
|
||||
return acc.String()
|
||||
}
|
||||
case <-deadline:
|
||||
return acc.String()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestOpen_ReadsShellOutput 验证开会话后向 shell 写命令能读回其输出。
|
||||
// Unix 走真 PTY；Windows 走降级管道实现。
|
||||
func TestOpen_ReadsShellOutput(t *testing.T) {
|
||||
sup := newTestSup(t, 5)
|
||||
sess, err := sup.Open(OpenParams{})
|
||||
if err != nil {
|
||||
t.Fatalf("open: %v", err)
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
defer sup.Close(ctx, sess.ID)
|
||||
|
||||
const marker = "ACWMARK42"
|
||||
cmd := "echo " + marker + "\n"
|
||||
if runtime.GOOS == "windows" {
|
||||
cmd = "echo " + marker + "\r\n"
|
||||
}
|
||||
if _, err := sess.Write([]byte(cmd)); err != nil {
|
||||
t.Fatalf("write: %v", err)
|
||||
}
|
||||
|
||||
got := readUntil(sess, marker, 5*time.Second)
|
||||
if !strings.Contains(got, marker) {
|
||||
t.Fatalf("marker %q not found in shell output; got %q", marker, got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestOpen_MaxSessions 验证达到上限后再开会话返回 ErrTooManySessions。
|
||||
func TestOpen_MaxSessions(t *testing.T) {
|
||||
sup := newTestSup(t, 1)
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
|
||||
s1, err := sup.Open(OpenParams{})
|
||||
if err != nil {
|
||||
t.Fatalf("first open: %v", err)
|
||||
}
|
||||
defer sup.Close(ctx, s1.ID)
|
||||
|
||||
if _, err := sup.Open(OpenParams{}); !errors.Is(err, ErrTooManySessions) {
|
||||
t.Fatalf("expected ErrTooManySessions, got %v", err)
|
||||
}
|
||||
if got := sup.Count(); got != 1 {
|
||||
t.Fatalf("expected 1 active session, got %d", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestClose_TerminatesSession 验证 Close 终止 shell（done 关闭）并从映射移除。
|
||||
func TestClose_TerminatesSession(t *testing.T) {
|
||||
sup := newTestSup(t, 5)
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
|
||||
sess, err := sup.Open(OpenParams{})
|
||||
if err != nil {
|
||||
t.Fatalf("open: %v", err)
|
||||
}
|
||||
if got := sup.Count(); got != 1 {
|
||||
t.Fatalf("expected 1 active session, got %d", got)
|
||||
}
|
||||
|
||||
sup.Close(ctx, sess.ID)
|
||||
|
||||
select {
|
||||
case <-sess.Done():
|
||||
case <-time.After(8 * time.Second):
|
||||
t.Fatal("session not terminated within 8s after Close")
|
||||
}
|
||||
|
||||
// monitor 在 done 关闭后异步移除映射，轮询直至归零。
|
||||
deadline := time.Now().Add(3 * time.Second)
|
||||
for sup.Count() != 0 {
|
||||
if time.Now().After(deadline) {
|
||||
t.Fatalf("session not removed from supervisor; count=%d", sup.Count())
|
||||
}
|
||||
time.Sleep(20 * time.Millisecond)
|
||||
}
|
||||
}
|
||||
|
||||
// TestCloseAll_ClosesEverySession 验证 CloseAll 终止全部会话。
|
||||
func TestCloseAll_ClosesEverySession(t *testing.T) {
|
||||
sup := newTestSup(t, 5)
|
||||
for i := 0; i < 3; i++ {
|
||||
if _, err := sup.Open(OpenParams{}); err != nil {
|
||||
t.Fatalf("open %d: %v", i, err)
|
||||
}
|
||||
}
|
||||
if got := sup.Count(); got != 3 {
|
||||
t.Fatalf("expected 3 active sessions, got %d", got)
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 8*time.Second)
|
||||
defer cancel()
|
||||
sup.CloseAll(ctx)
|
||||
|
||||
deadline := time.Now().Add(3 * time.Second)
|
||||
for sup.Count() != 0 {
|
||||
if time.Now().After(deadline) {
|
||||
t.Fatalf("sessions not all closed; count=%d", sup.Count())
|
||||
}
|
||||
time.Sleep(20 * time.Millisecond)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user