Cyrene_v2/internal/services/gateway/ws_role_gate_test.go

package gateway

import (
	"context"
	"encoding/json"
	"errors"
	"net"
	"sync"
	"testing"
	"time"

	"github.com/gorilla/websocket"

	"git.yeij.top/AskaEth/Cyrene_v2/internal/auth"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/history"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/loop"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/model"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/orchestration"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/task"
)

// This file locks design/11 §7.2 — the two remaining places where a `viewer`
// (the read-only monitoring role) used to reach beyond its lane through
// /ws/chat, both found while implementing the §7.1 downlink gate:
//
//	A. the handshake history replay (ws_history.go) carries the main session's
//	   message bodies, so it is owner/admin only;
//	B. every inbound (client→server) JSON frame (ws_route.go) is denied by
//	   default for a non owner/admin connection, so a viewer can follow the
//	   session state machine but can never send an instruction.
//
// Together with the §7.1 downlink gate they pin the combined contract:
// `viewer` = may connect, may read its 12 read-only REST monitoring endpoints,
// may receive the 7 state-machine signals — and gets and sends nothing else over
// WS. The connections and tokens are real (httptest server + signed JWTs), so a
// mistake in the routing tables cannot hide behind a fake.
const (
	// wsRoleGateReadWait bounds a read that is expected to produce a frame.
	// gorilla corrupts a connection after a read timeout, so every connection
	// here performs at most one timeout read, always as its last read.
	wsRoleGateReadWait = 3 * time.Second
	// wsRoleGateNoFrameWait is the window in which a frame that must not exist
	// would have arrived (a directed reply is queued synchronously by the
	// handler, so an idle window this long is conclusive).
	wsRoleGateNoFrameWait = 250 * time.Millisecond
	// wsRoleGateSettle lets the side effects of an accepted frame land before they
	// are asserted to be absent.
	wsRoleGateSettle = 200 * time.Millisecond
	// wsRoleGateWriteWait bounds a control-frame write (ping / close).
	wsRoleGateWriteWait = 2 * time.Second
)

// countingRunner is the Chatter seam with a call counter: design/11 §7.2 B has to
// prove that a viewer's `chat` frame never reaches the model, so "was the runner
// driven" must be observable rather than inferred from a broadcast.
type countingRunner struct {
	mu    sync.Mutex
	calls int
}

func (r *countingRunner) Run(context.Context, string, []model.Message) (*loop.Result, error) {
	r.mu.Lock()
	r.calls++
	r.mu.Unlock()
	return &loop.Result{Text: "reply", Steps: 1, StopReason: "stop"}, nil
}

func (r *countingRunner) count() int {
	r.mu.Lock()
	defer r.mu.Unlock()
	return r.calls
}

// replayHistoryStore serves the handshake replay with a fixed transcript while
// inheriting fakeHistory's Append recording, so one store both makes the replay
// deterministic and reports what an accepted `chat` frame persisted.
type replayHistoryStore struct {
	*fakeHistory
	mu   sync.Mutex
	msgs []history.Message
}

func newReplayHistoryStore(msgs ...history.Message) *replayHistoryStore {
	return &replayHistoryStore{fakeHistory: &fakeHistory{}, msgs: msgs}
}

func (s *replayHistoryStore) Recent(context.Context, string, int) ([]history.Message, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	return append([]history.Message(nil), s.msgs...), nil
}

// historyFrameWire is the handshake replay frame decoded for assertions.
type historyFrameWire struct {
	Type      string `json:"type"`
	SessionID string `json:"sessionId"`
	Data      struct {
		Messages []struct {
			Role     string `json:"role"`
			Content  string `json:"content"`
			Thinking string `json:"thinking"`
		} `json:"messages"`
	} `json:"data"`
}

// errorFrameWire is any {type:"error"} frame decoded for assertions.
type errorFrameWire struct {
	Type string `json:"type"`
	Data struct {
		Code    string `json:"code"`
		Message string `json:"message"`
	} `json:"data"`
}

// roleGateEnv wires a gateway for these tests: a counting runner (so a leaked
// inbound frame is visible as a model call), a store that replays a transcript and
// records persistence, and the sub-agent/task seams so a leaked subagent_ask or
// task_output_query would produce a real, assertable content frame instead of a
// handler-level error.
type roleGateEnv struct {
	srv  *Server
	url  string
	jwt  *auth.JWTService
	run  *countingRunner
	hist *replayHistoryStore
}

func newRoleGateEnv(t *testing.T, transcript []history.Message) *roleGateEnv {
	t.Helper()
	run := &countingRunner{}
	hist := newReplayHistoryStore(transcript...)
	srv, ts, jwt := buildEnv(t, cfgDefault(), Deps{
		Loop:    run,
		History: hist,
		Scheduler: &stubAuditor{res: &orchestration.AuditResult{
			AgentID: "sub-1",
			Mode:    orchestration.AuditFull,
			Goal:    "子代理的私密目标",
		}},
		Tasks: &stubTaskOutputs{ev: task.OutputEvent{
			ID:     "task-1",
			Output: task.Output{Stdout: "任务的私密输出"},
		}},
	})
	return &roleGateEnv{srv: srv, url: ts.URL, jwt: jwt, run: run, hist: hist}
}

func (e *roleGateEnv) token(t *testing.T, role auth.Role) string {
	t.Helper()
	return signAny(t, e.jwt, "acct", role)
}

// TestWSHandshakeReplayPerRole locks design/11 §7.2 A: the handshake history
// replay holds the main session's message bodies, so only owner/admin connections
// receive it. A leaked replay is always the FIRST frame on a connection — it is
// queued during the upgrade, before the read pump starts and before any reply —
// so reading the viewer's first frame turns "no replay" into an exact assertion
// instead of a guessed idle window.
func TestWSHandshakeReplayPerRole(t *testing.T) {
	t.Parallel()
	env := newRoleGateEnv(t, []history.Message{
		{ID: "m1", SessionID: history.MainSessionID, Role: history.RoleUser, Content: "主人的私密提问", Thinking: "不该泄漏的推理", CreatedAt: 1},
		{ID: "m2", SessionID: history.MainSessionID, Role: history.RoleAssistant, Content: "角色的私密回答", CreatedAt: 2},
	})

	owner := dialWS(t, env.url, env.token(t, auth.RoleOwner))
	defer owner.Close()
	admin := dialWS(t, env.url, env.token(t, auth.RoleAdmin))
	defer admin.Close()
	viewer := dialWS(t, env.url, env.token(t, auth.RoleViewer))
	defer viewer.Close()
	waitConnCount(t, env.srv, 3)

	for _, tc := range []struct {
		role auth.Role
		conn *websocket.Conn
	}{{auth.RoleOwner, owner}, {auth.RoleAdmin, admin}} {
		typ, raw := readWire(t, tc.conn, string(tc.role)+" handshake")
		if typ != "history" {
			t.Fatalf("%s first frame = %q, want the history replay", tc.role, typ)
		}
		var hf historyFrameWire
		if err := json.Unmarshal(raw, &hf); err != nil {
			t.Fatalf("%s decode history frame %s: %v", tc.role, raw, err)
		}
		if hf.SessionID != history.MainSessionID {
			t.Fatalf("%s history sessionId = %q, want %q", tc.role, hf.SessionID, history.MainSessionID)
		}
		if len(hf.Data.Messages) != 2 {
			t.Fatalf("%s replayed %d messages, want 2 (%+v)", tc.role, len(hf.Data.Messages), hf.Data.Messages)
		}
		if got := hf.Data.Messages[0].Content; got != "主人的私密提问" {
			t.Fatalf("%s replayed content = %q, want the stored user message", tc.role, got)
		}
		for i, m := range hf.Data.Messages {
			if m.Thinking != "" {
				t.Fatalf("%s replayed message[%d] carries thinking %q; it must be stripped", tc.role, i, m.Thinking)
			}
		}
	}

	// The viewer is probed so that a frame is guaranteed to come back (the §7.2 B
	// refusal); if the replay leaked it would be ahead of that reply.
	if err := viewer.WriteJSON(wsInbound{Type: "chat", Data: json.RawMessage(`{"sessionId":"web","text":"viewer"}`)}); err != nil {
		t.Fatalf("viewer write chat: %v", err)
	}
	if typ, raw := readWire(t, viewer, "viewer handshake"); typ == "history" {
		t.Fatalf("viewer received the handshake replay %s; design/11 §7.2 A allows owner/admin only", raw)
	}
}

// TestWSInboundDeniedForViewer locks design/11 §7.2 B: a non owner/admin
// connection may not send any in-app instruction. Every rejected frame gets a
// connection-directed {type:"error", data:{code:"forbidden"}} and nothing else —
// no content frame, no broadcast, no model call, no persisted message.
func TestWSInboundDeniedForViewer(t *testing.T) {
	t.Parallel()
	env := newRoleGateEnv(t, nil)

	// A second viewer connection observes that the refusals are directed at the
	// sender: a broadcast refusal would land here too.
	observer := dialWS(t, env.url, env.token(t, auth.RoleViewer))
	defer observer.Close()

	cases := []struct {
		name  string
		frame wsInbound
	}{
		{"chat", wsInbound{Type: "chat", Data: json.RawMessage(`{"sessionId":"web","text":"替主人发一条"}`)}},
		{"interrupt", wsInbound{Type: "interrupt", Data: json.RawMessage(`{"sessionId":"web"}`)}},
		{"subagent_ask", wsInbound{Type: "subagent_ask", Data: json.RawMessage(`{"id":"sub-1","mode":"full"}`)}},
		{"task_output_query", wsInbound{Type: "task_output_query", Data: json.RawMessage(`{"taskId":"task-1"}`)}},
		{"voice_input", wsInbound{Type: "voice_input", Data: json.RawMessage(`{"sessionId":"web","text":"你好"}`)}},
		{"mode_switch", wsInbound{Type: "mode_switch", Data: json.RawMessage(`{"mode":"chat"}`)}},
	}
	for _, tc := range cases {
		tc := tc
		t.Run(tc.name, func(t *testing.T) {
			conn := dialWS(t, env.url, env.token(t, auth.RoleViewer))
			defer conn.Close()
			if err := conn.WriteJSON(tc.frame); err != nil {
				t.Fatalf("write %s: %v", tc.name, err)
			}
			typ, raw := readWire(t, conn, tc.name)
			if typ != "error" {
				t.Fatalf("%s reply = %q (%s), want a directed error frame", tc.name, typ, raw)
			}
			var ef errorFrameWire
			if err := json.Unmarshal(raw, &ef); err != nil {
				t.Fatalf("%s decode reply %s: %v", tc.name, raw, err)
			}
			if ef.Data.Code != wsForbiddenCode {
				t.Fatalf("%s error code = %q, want %q (frame=%s)", tc.name, ef.Data.Code, wsForbiddenCode, raw)
			}
			// For subagent_ask / task_output_query a leak would be a directed
			// subagent_event / task_output reply carrying a sub-agent's context or
			// a task's output, so assert the refusal is the whole answer.
			assertNoFrame(t, conn, tc.name+" follow-up")
		})
	}

	// The `chat` case is the one that used to drive the model and persist the
	// message; give its (refused) side effects a moment to not happen.
	time.Sleep(wsRoleGateSettle)
	calls, records := env.run.count(), env.hist.records()
	if calls != 0 || len(records) != 0 {
		t.Fatalf("viewer frames reached the runner %d time(s) and persisted %d message(s) %+v; a viewer may do neither", calls, len(records), records)
	}
	assertNoFrame(t, observer, "observer viewer (a refusal must not be broadcast)")
}

// TestWSInboundChatDrivesOwnerTurn is the owner half of design/11 §7.2 B: the very
// frame a viewer is refused (`chat`) still drives the model and persists the turn
// for an owner, so the gate is proven to be a role gate rather than a broken chat
// path.
func TestWSInboundChatDrivesOwnerTurn(t *testing.T) {
	t.Parallel()
	env := newRoleGateEnv(t, nil)
	owner := dialWS(t, env.url, env.token(t, auth.RoleOwner))
	defer owner.Close()
	waitConnCount(t, env.srv, 1)

	if err := owner.WriteJSON(wsInbound{Type: "chat", Data: json.RawMessage(`{"sessionId":"web","text":"主人问一句"}`)}); err != nil {
		t.Fatalf("write chat: %v", err)
	}
	waitUntil(t, "the owner's chat to drive the runner", func() bool { return env.run.count() == 1 })
	waitUntil(t, "the owner's turn to be persisted", func() bool { return len(env.hist.records()) == 2 })

	got := env.hist.records()
	if got[0].Role != history.RoleUser || got[1].Role != history.RoleAssistant {
		t.Fatalf("persisted roles = %q, %q, want user then assistant", got[0].Role, got[1].Role)
	}
	if got[1].Content != "reply" {
		t.Fatalf("persisted reply = %q, want the runner's text", got[1].Content)
	}
}

// TestWSInboundGateKeepsProtocolFrames proves the inbound gate cannot kill a
// connection (design/11 §7.2 B 「协议级帧不受影响」): after a refusal the viewer
// still completes a ping/pong round trip and a normal close, and the server
// deregisters it.
func TestWSInboundGateKeepsProtocolFrames(t *testing.T) {
	t.Parallel()
	env := newRoleGateEnv(t, nil)
	viewer := dialWS(t, env.url, env.token(t, auth.RoleViewer))
	waitConnCount(t, env.srv, 1)

	if err := viewer.WriteJSON(wsInbound{Type: "chat", Data: json.RawMessage(`{"sessionId":"web","text":"x"}`)}); err != nil {
		t.Fatalf("write chat: %v", err)
	}
	if typ, raw := readWire(t, viewer, "probe"); typ != "error" {
		t.Fatalf("probe reply = %q (%s), want the forbidden error frame", typ, raw)
	}

	// A second refusal round trip proves the read pump survived the first: the
	// gate returns before any handler, so the connection is untouched.
	if err := viewer.WriteJSON(wsInbound{Type: "interrupt", Data: json.RawMessage(`{"sessionId":"web"}`)}); err != nil {
		t.Fatalf("write interrupt after refusal: %v", err)
	}
	if typ, raw := readWire(t, viewer, "second probe"); typ != "error" {
		t.Fatalf("second probe reply = %q (%s), want the forbidden error frame", typ, raw)
	}

	// A client ping is answered by the server's default ping handler. The pong is
	// consumed by the client read below, so the handler is what reports it. This
	// is the connection's single timeout read (gorilla corrupts a connection after
	// a read deadline), hence it comes after every assertion on the read side.
	ponged := make(chan struct{}, 1)
	viewer.SetPongHandler(func(string) error {
		select {
		case ponged <- struct{}{}:
		default:
		}
		return nil
	})
	if err := viewer.WriteControl(websocket.PingMessage, nil, time.Now().Add(wsRoleGateWriteWait)); err != nil {
		t.Fatalf("write ping: %v", err)
	}
	_ = viewer.SetReadDeadline(time.Now().Add(wsRoleGateReadWait))
	if _, msg, err := viewer.ReadMessage(); err == nil {
		t.Fatalf("data frame %s while waiting for the pong", msg)
	}
	select {
	case <-ponged:
	default:
		t.Fatal("the server never answered the ping; the gate broke the protocol path")
	}

	// A normal close still deregisters the connection.
	if err := viewer.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""), time.Now().Add(wsRoleGateWriteWait)); err != nil {
		t.Fatalf("write close: %v", err)
	}
	waitConnCount(t, env.srv, 0)
}

// readWire reads one frame, returning its type plus the raw bytes (so a caller can
// decode the type-specific payload). It fails the test when no frame arrives
// before the deadline.
func readWire(t *testing.T, conn *websocket.Conn, what string) (string, []byte) {
	t.Helper()
	_ = conn.SetReadDeadline(time.Now().Add(wsRoleGateReadWait))
	_, msg, err := conn.ReadMessage()
	if err != nil {
		t.Fatalf("%s: read frame: %v", what, err)
	}
	var env struct {
		Type string `json:"type"`
	}
	if err := json.Unmarshal(msg, &env); err != nil {
		t.Fatalf("%s: decode frame %s: %v", what, msg, err)
	}
	return env.Type, msg
}

// assertNoFrame asserts that nothing further arrives on conn within a short
// window. It must be the connection's last read: gorilla leaves a connection
// corrupt after a read deadline expires.
func assertNoFrame(t *testing.T, conn *websocket.Conn, what string) {
	t.Helper()
	_ = conn.SetReadDeadline(time.Now().Add(wsRoleGateNoFrameWait))
	_, msg, err := conn.ReadMessage()
	if err == nil {
		t.Fatalf("%s: unexpected frame %s", what, msg)
	}
	var netErr net.Error
	if !errors.As(err, &netErr) || !netErr.Timeout() {
		t.Fatalf("%s: read failed with %v, want a timeout (no frame expected)", what, err)
	}
}

// waitUntil polls cond until it holds or the deadline passes, so a test can assert
// on an asynchronous turn (run → persist) without a fixed sleep that could pass
// vacuously.
func waitUntil(t *testing.T, what string, cond func() bool) {
	t.Helper()
	deadline := time.Now().Add(wsRoleGateReadWait)
	for time.Now().Before(deadline) {
		if cond() {
			return
		}
		time.Sleep(2 * time.Millisecond)
	}
	t.Fatalf("timed out waiting for %s", what)
}