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)
|
||||
}
|
||||
Reference in New Issue
Block a user