feat(gateway): 网关（REST /api/v1 + WS Hub /ws/chat + admin /api/v1/admin/* + JWT/RBAC + 审计，单端口 :8080）
This commit is contained in:
Aska committed
1 parent
ff34b62079
commit
05e23ea590
7 files changed
+1535
No files matched your search
@@ -0,0 +1,383 @@
|
||||
package gateway
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"log/slog"
|
||||
"net/http"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/gorilla/websocket"
|
||||
|
||||
"git.yeij.top/AskaEth/Cyrene_v2/internal/auth"
|
||||
"git.yeij.top/AskaEth/Cyrene_v2/internal/bus"
|
||||
apperr "git.yeij.top/AskaEth/Cyrene_v2/internal/errors"
|
||||
"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/model"
|
||||
)
|
||||
|
||||
// WebSocket tuning constants. The ping/pong cadence keeps intermediary proxies
|
||||
// from dropping idle connections.
|
||||
const (
|
||||
wsEventsPattern = "ws.>"
|
||||
wsPrefix = "ws."
|
||||
wsChatSubject = "ws.chat"
|
||||
wsReadLimit = 4096
|
||||
wsWriteWait = 10 * time.Second
|
||||
wsPongWait = 60 * time.Second
|
||||
wsPingPeriod = 54 * time.Second
|
||||
wsSendBuffer = 32
|
||||
)
|
||||
|
||||
// wsEvent is the outbound (server→client) frame.
|
||||
type wsEvent struct {
|
||||
Type string `json:"type"`
|
||||
Data any `json:"data,omitempty"`
|
||||
}
|
||||
|
||||
// wsInbound is the client→server frame. Data carries the per-type payload.
|
||||
type wsInbound struct {
|
||||
Type string `json:"type"`
|
||||
ReqID string `json:"reqId,omitempty"`
|
||||
Data json.RawMessage `json:"data,omitempty"`
|
||||
}
|
||||
|
||||
// wsConn wraps one WebSocket connection. Writes are serialized through the
|
||||
// buffered send channel drained by a single writePump, so the gorilla conn is
|
||||
// only ever touched by one goroutine at a time.
|
||||
type wsConn struct {
|
||||
conn *websocket.Conn
|
||||
send chan []byte
|
||||
done chan struct{}
|
||||
h *hub
|
||||
claims *auth.Claims
|
||||
once sync.Once
|
||||
}
|
||||
|
||||
// hub owns the connection set and relays bus events to every client.
|
||||
type hub struct {
|
||||
mu sync.RWMutex
|
||||
conns map[*wsConn]struct{}
|
||||
bus bus.Bus
|
||||
runner Chatter
|
||||
log *slog.Logger
|
||||
upgrader websocket.Upgrader
|
||||
origins map[string]bool
|
||||
allowAll bool
|
||||
cancel context.CancelFunc
|
||||
}
|
||||
|
||||
// newHub builds a hub bound to the server's dependencies.
|
||||
func newHub(s *Server) *hub {
|
||||
origins, allowAll := s.allowedOrigins()
|
||||
return &hub{
|
||||
conns: make(map[*wsConn]struct{}),
|
||||
bus: s.bus,
|
||||
runner: s.runner,
|
||||
log: s.log,
|
||||
origins: origins,
|
||||
allowAll: allowAll,
|
||||
upgrader: websocket.Upgrader{
|
||||
ReadBufferSize: 1024,
|
||||
WriteBufferSize: 1024,
|
||||
CheckOrigin: originChecker(origins, allowAll),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// start subscribes the hub to the bus event stream. Publishing to ws.> arrives
|
||||
// here and is fanned out to clients; nothing happens when no bus is wired.
|
||||
func (h *hub) start() {
|
||||
if h.bus == nil {
|
||||
return
|
||||
}
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
h.cancel = cancel
|
||||
if _, err := h.bus.Subscribe(ctx, wsEventsPattern, h.relay); err != nil {
|
||||
h.log.Warn("bus subscribe", "pattern", wsEventsPattern, "err", err)
|
||||
}
|
||||
}
|
||||
|
||||
// close cancels the bus subscription and shuts every live connection down.
|
||||
func (h *hub) close() {
|
||||
if h.cancel != nil {
|
||||
h.cancel()
|
||||
}
|
||||
h.mu.Lock()
|
||||
conns := make([]*wsConn, 0, len(h.conns))
|
||||
for c := range h.conns {
|
||||
conns = append(conns, c)
|
||||
}
|
||||
h.conns = make(map[*wsConn]struct{})
|
||||
h.mu.Unlock()
|
||||
for _, c := range conns {
|
||||
c.shutdown()
|
||||
}
|
||||
}
|
||||
|
||||
// register adds a connection to the set.
|
||||
func (h *hub) register(c *wsConn) {
|
||||
h.mu.Lock()
|
||||
h.conns[c] = struct{}{}
|
||||
h.mu.Unlock()
|
||||
}
|
||||
|
||||
// unregister removes a connection from the set.
|
||||
func (h *hub) unregister(c *wsConn) {
|
||||
h.mu.Lock()
|
||||
if _, ok := h.conns[c]; ok {
|
||||
delete(h.conns, c)
|
||||
}
|
||||
h.mu.Unlock()
|
||||
}
|
||||
|
||||
// connCount reports the number of live connections (observability/tests).
|
||||
func (h *hub) connCount() int {
|
||||
h.mu.RLock()
|
||||
defer h.mu.RUnlock()
|
||||
return len(h.conns)
|
||||
}
|
||||
|
||||
// broadcast fans an event out to every registered connection without blocking on
|
||||
// any single (possibly slow) client: a full send buffer drops the frame.
|
||||
func (h *hub) broadcast(ev wsEvent) {
|
||||
b, err := json.Marshal(ev)
|
||||
if err != nil {
|
||||
h.log.Warn("marshal broadcast", "err", err)
|
||||
return
|
||||
}
|
||||
h.mu.RLock()
|
||||
defer h.mu.RUnlock()
|
||||
for c := range h.conns {
|
||||
c.enqueue(b)
|
||||
}
|
||||
}
|
||||
|
||||
// relay is the bus handler for the ws.> pattern. The ws.chat subject is a
|
||||
// streaming trigger (run the loop and emit chunks); every other subject becomes
|
||||
// an outbound event of its design type.
|
||||
func (h *hub) relay(ctx context.Context, e bus.Event) error {
|
||||
if e.Subject == wsChatSubject {
|
||||
go h.streamChat(ctx, e.Payload)
|
||||
return nil
|
||||
}
|
||||
h.broadcast(wsEvent{Type: eventType(e.Subject), Data: e.Payload})
|
||||
return nil
|
||||
}
|
||||
|
||||
// streamChat runs one agent turn for a chat trigger and broadcasts the result as
|
||||
// stream_chunk/stream_end. It runs on its own goroutine so a slow model call does
|
||||
// not stall the bus subscription worker.
|
||||
func (h *hub) streamChat(ctx context.Context, payload any) {
|
||||
var req chatStreamRequest
|
||||
if b, err := json.Marshal(payload); err == nil {
|
||||
_ = json.Unmarshal(b, &req)
|
||||
}
|
||||
if h.runner == nil || strings.TrimSpace(req.Text) == "" {
|
||||
h.broadcast(wsEvent{Type: "stream_end", Data: map[string]any{"stopReason": "idle"}})
|
||||
return
|
||||
}
|
||||
if req.SessionID == "" {
|
||||
req.SessionID = "web"
|
||||
}
|
||||
res, err := h.runner.Run(ctx, req.SessionID, []model.Message{model.TextMessage(model.RoleUser, req.Text)})
|
||||
if err != nil {
|
||||
h.log.Error("ws chat run", "err", err, "session", req.SessionID)
|
||||
h.broadcast(wsEvent{Type: "stream_end", Data: map[string]any{"stopReason": "error", "error": err.Error()}})
|
||||
return
|
||||
}
|
||||
h.broadcast(wsEvent{Type: "stream_chunk", Data: map[string]any{"delta": res.Text}})
|
||||
h.broadcast(wsEvent{Type: "stream_end", Data: map[string]any{"stopReason": res.StopReason, "steps": res.Steps}})
|
||||
}
|
||||
|
||||
// chatStreamRequest is the decoded ws.chat payload.
|
||||
type chatStreamRequest struct {
|
||||
SessionID string `json:"sessionId"`
|
||||
Text string `json:"text"`
|
||||
}
|
||||
|
||||
// handleWS upgrades /ws/chat after authenticating the JWT from the query token or
|
||||
// the Authorization header. A failed upgrade writes a 401 {error:...} before the
|
||||
// handshake, so a browser/IM client sees a clean reject.
|
||||
func (s *Server) handleWS(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Method != http.MethodGet {
|
||||
writeError(w, apperr.Wrap(apperr.ErrBadRequest, "bad_request", "websocket requires GET"))
|
||||
return
|
||||
}
|
||||
claims, err := s.auth.Verify(wsToken(r))
|
||||
if err != nil {
|
||||
s.log.Debug("ws auth", "err", err)
|
||||
writeError(w, apperr.ErrUnauthorized)
|
||||
return
|
||||
}
|
||||
conn, err := s.hub.upgrader.Upgrade(w, r, nil)
|
||||
if err != nil {
|
||||
s.log.Debug("ws upgrade", "err", err)
|
||||
return
|
||||
}
|
||||
c := &wsConn{
|
||||
conn: conn,
|
||||
send: make(chan []byte, wsSendBuffer),
|
||||
done: make(chan struct{}),
|
||||
h: s.hub,
|
||||
claims: claims,
|
||||
}
|
||||
s.hub.register(c)
|
||||
go c.writePump()
|
||||
go c.readPump()
|
||||
}
|
||||
|
||||
// readPump reads client frames until the connection closes, routing each by type.
|
||||
func (c *wsConn) readPump() {
|
||||
defer func() {
|
||||
c.h.unregister(c)
|
||||
c.shutdown()
|
||||
}()
|
||||
c.conn.SetReadLimit(wsReadLimit)
|
||||
_ = c.conn.SetReadDeadline(time.Now().Add(wsPongWait))
|
||||
c.conn.SetPongHandler(func(string) error {
|
||||
return c.conn.SetReadDeadline(time.Now().Add(wsPongWait))
|
||||
})
|
||||
for {
|
||||
_, msg, err := c.conn.ReadMessage()
|
||||
if err != nil {
|
||||
if websocket.IsUnexpectedCloseError(err, websocket.CloseNormalClosure, websocket.CloseGoingAway) {
|
||||
c.h.log.Debug("ws read", "err", err)
|
||||
}
|
||||
return
|
||||
}
|
||||
var in wsInbound
|
||||
if err := json.Unmarshal(msg, &in); err != nil {
|
||||
c.enqueueJSON(wsEvent{Type: "error", Data: map[string]any{"message": "invalid json"}})
|
||||
continue
|
||||
}
|
||||
c.route(in)
|
||||
}
|
||||
}
|
||||
|
||||
// writePump drains the send channel (and sends keepalive pings) until the
|
||||
// connection closes. It is the only writer on the gorilla conn.
|
||||
func (c *wsConn) writePump() {
|
||||
ticker := time.NewTicker(wsPingPeriod)
|
||||
defer func() {
|
||||
ticker.Stop()
|
||||
_ = c.conn.Close()
|
||||
}()
|
||||
for {
|
||||
select {
|
||||
case <-c.done:
|
||||
return
|
||||
case msg := <-c.send:
|
||||
_ = c.conn.SetWriteDeadline(time.Now().Add(wsWriteWait))
|
||||
if err := c.conn.WriteMessage(websocket.TextMessage, msg); err != nil {
|
||||
return
|
||||
}
|
||||
case <-ticker.C:
|
||||
_ = c.conn.SetWriteDeadline(time.Now().Add(wsWriteWait))
|
||||
if err := c.conn.WriteMessage(websocket.PingMessage, nil); err != nil {
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// route dispatches one inbound frame by its type.
|
||||
func (c *wsConn) route(in wsInbound) {
|
||||
switch in.Type {
|
||||
case "mode_switch":
|
||||
c.publishPublic("ws.mode_switch", rawData(in.Data))
|
||||
case "chat":
|
||||
c.publishPublic(wsChatSubject, rawData(in.Data))
|
||||
case "subagent_ask":
|
||||
c.enqueueJSON(wsEvent{Type: "subagent_event", Data: map[string]any{"state": "pending"}})
|
||||
case "task_output_query":
|
||||
c.enqueueJSON(wsEvent{Type: "task_output", Data: map[string]any{}})
|
||||
default:
|
||||
c.enqueueJSON(wsEvent{Type: "error", Data: map[string]any{"message": "unknown event type"}})
|
||||
}
|
||||
}
|
||||
|
||||
// publishPublic publishes a bus event without any per-connection reply.
|
||||
func (c *wsConn) publishPublic(subject string, payload any) {
|
||||
if c.h.bus == nil {
|
||||
return
|
||||
}
|
||||
_ = c.h.bus.Publish(context.Background(), subject, payload)
|
||||
}
|
||||
|
||||
// shutdown closes the connection exactly once.
|
||||
func (c *wsConn) shutdown() {
|
||||
c.once.Do(func() {
|
||||
close(c.done)
|
||||
_ = c.conn.Close()
|
||||
})
|
||||
}
|
||||
|
||||
// enqueue queues a serialized frame, dropping it when a slow client's buffer is
|
||||
// full so the producer never blocks.
|
||||
func (c *wsConn) enqueue(b []byte) {
|
||||
select {
|
||||
case c.send <- b:
|
||||
default:
|
||||
}
|
||||
}
|
||||
|
||||
// enqueueJSON serializes ev and queues it.
|
||||
func (c *wsConn) enqueueJSON(ev wsEvent) {
|
||||
b, err := json.Marshal(ev)
|
||||
if err != nil {
|
||||
c.h.log.Warn("marshal ws event", "err", err)
|
||||
return
|
||||
}
|
||||
c.enqueue(b)
|
||||
}
|
||||
|
||||
// publishChatStream publishes a ws.chat streaming trigger (REST stream mode).
|
||||
func (s *Server) publishChatStream(ctx context.Context, sessionID, text string) {
|
||||
if s.bus == nil {
|
||||
return
|
||||
}
|
||||
_ = s.bus.Publish(ctx, wsChatSubject, map[string]any{"sessionId": sessionID, "text": text})
|
||||
}
|
||||
|
||||
// wsToken extracts the JWT from the ?token= query or the Authorization header.
|
||||
func wsToken(r *http.Request) string {
|
||||
if t := r.URL.Query().Get("token"); t != "" {
|
||||
return t
|
||||
}
|
||||
const prefix = "Bearer "
|
||||
h := r.Header.Get("Authorization")
|
||||
if strings.HasPrefix(h, prefix) {
|
||||
return strings.TrimSpace(h[len(prefix):])
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
// eventType maps a bus subject to its outbound event type, dropping the ws. prefix.
|
||||
func eventType(subject string) string {
|
||||
return strings.TrimPrefix(subject, wsPrefix)
|
||||
}
|
||||
|
||||
// rawData converts a possibly-empty raw payload into a value safe to publish.
|
||||
func rawData(d json.RawMessage) any {
|
||||
if len(d) == 0 {
|
||||
return nil
|
||||
}
|
||||
return d
|
||||
}
|
||||
|
||||
// originChecker allows same-origin (empty) and whitelisted WebSocket origins,
|
||||
// mirroring the CORS allow list. A "*" default allows all.
|
||||
func originChecker(allowed map[string]bool, allowAll bool) func(*http.Request) bool {
|
||||
return func(r *http.Request) bool {
|
||||
origin := r.Header.Get("Origin")
|
||||
if origin == "" {
|
||||
return true
|
||||
}
|
||||
if allowAll {
|
||||
return true
|
||||
}
|
||||
return allowed[origin]
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user