feat(auth): add PAT auth, scoped REST/MCP access, CSRF, and tenant isolation

Introduce Personal Access Tokens (nzp_*) as a stateless auth path alongside
JWT, gated per-endpoint by a scope middleware (nezha:{resource}:{verb}) with
fail-closed empty-scope defaults and a server-id whitelist. Self-management
endpoints (profile, api-tokens, oauth2 bind, refresh-token) explicitly reject
PATs to block privilege-escalation chains. A revoke registry tears down active
long-lived connections (terminal, fm, ws, transfer, mcp) the moment a PAT is
deleted, with a tombstone closing the revoke->register race.

Add an MCP endpoint that proxies tool calls (exec, fs read/write/delete,
transfer) to agents over gRPC, guarded by origin/DNS-rebinding checks, a
per-token rate limiter, audit logging, and a kill switch. Serialize all
sends through the IOStream wrapper to honour grpc-go's concurrency contract.

Add CSRF double-submit protection on unsafe cookie-authenticated methods,
exempting authenticated PAT requests by context identity (not a forgeable
Authorization header). Apply visibility/whitelist filtering consistently
across list, get-by-id, and mutate paths to enforce tenant isolation.

Migrate legacy mcp:* scopes: rewrite read/exec to nezha:* equivalents and
drop dangerous write/delete/wildcard grants.

Co-authored-by: cloudcode <cloudcode@users.noreply.github.com>
This commit is contained in:
naiba
2026-05-30 15:56:44 +00:00
co-authored by cloudcode
parent 58c1ea7b0b
commit ab25662ddd
153 changed files with 16974 additions and 244 deletions
+146 -37
View File
@@ -1,24 +1,39 @@
package rpc
import (
"context"
"errors"
"io"
"sync"
"sync/atomic"
"time"
"github.com/nezhahq/nezha/service/singleton"
)
// StreamPurpose tags every IOStream with the feature that opened it so
// admin actions can drop only the relevant subset. Existing call sites
// (terminal / fm / NAT / server-transfer) keep PurposeLegacy and the
// previous semantics; only the new MCP fs.transfer path uses
// PurposeMCPTransfer, which is what EnableMCP=false revokes.
type StreamPurpose uint8
const (
PurposeLegacy StreamPurpose = iota
PurposeMCPTransfer
)
type ioStreamContext struct {
creatorUserID uint64
targetServerID uint64
purpose StreamPurpose
userIo io.ReadWriteCloser
agentIo io.ReadWriteCloser
userIoConnectCh chan struct{}
agentIoConnectCh chan struct{}
userIoChOnce sync.Once
agentIoChOnce sync.Once
revokedCh chan struct{}
revokedOnce sync.Once
}
type bp struct {
@@ -34,14 +49,20 @@ var bufPool = sync.Pool{
}
func (s *NezhaHandler) CreateStream(streamId string, creatorUserID uint64, targetServerID uint64) {
s.CreateStreamWithPurpose(streamId, creatorUserID, targetServerID, PurposeLegacy)
}
func (s *NezhaHandler) CreateStreamWithPurpose(streamId string, creatorUserID uint64, targetServerID uint64, purpose StreamPurpose) {
s.ioStreamMutex.Lock()
defer s.ioStreamMutex.Unlock()
s.ioStreams[streamId] = &ioStreamContext{
creatorUserID: creatorUserID,
targetServerID: targetServerID,
purpose: purpose,
userIoConnectCh: make(chan struct{}),
agentIoConnectCh: make(chan struct{}),
revokedCh: make(chan struct{}),
}
}
@@ -64,6 +85,42 @@ func (s *NezhaHandler) IsStreamAuthorizedForAgent(streamId string, agentServerID
return ctx.targetServerID != 0 && ctx.targetServerID == agentServerID
}
// WaitForAgent 阻塞等待 agent 端通过 IOStream 接入并完成 AgentConnected。
// dashboard 把 MCP 大文件传输的 task 派给 agent 后,需要等 agent dial 回来
// 才能开始 Read/Write,这里以 timeout 内的轻量轮询暴露给 controller。
//
// 同时返回 agent 端流(io.ReadWriteCloser)以便 controller 调 io.CopyN 转发
// HTTP bodyok=false 表示超时或流已被关闭。
func (s *NezhaHandler) WaitForAgent(ctx context.Context, streamId string, timeout time.Duration) (io.ReadWriteCloser, bool) {
deadline := time.NewTimer(timeout)
defer deadline.Stop()
for {
s.ioStreamMutex.RLock()
sc, ok := s.ioStreams[streamId]
if ok && sc.agentIo != nil {
s.ioStreamMutex.RUnlock()
return sc.agentIo, true
}
s.ioStreamMutex.RUnlock()
if !ok {
return nil, false
}
select {
case <-ctx.Done():
return nil, false
case <-deadline.C:
return nil, false
case <-sc.revokedCh:
return nil, false
case <-sc.agentIoConnectCh:
s.ioStreamMutex.RLock()
ag := sc.agentIo
s.ioStreamMutex.RUnlock()
return ag, ag != nil
}
}
}
// IsStreamAuthorizedForUser checks whether the requesting user may attach to
// the stream. A stream is reachable only by its creator or by an admin; any
// other authenticated user must be rejected. Unknown streams are always
@@ -108,6 +165,23 @@ func (s *NezhaHandler) StreamOwnership(streamId string) (uint64, bool) {
return ctx.creatorUserID, true
}
// StreamTarget returns the server ID the stream was opened against and
// whether the stream is still tracked. Callers MUST pass this through the
// requesting PAT's CanAccessServer check before allowing attachment —
// IsStreamAuthorizedForUser only knows about creator/admin, so without this
// dual gate an admin's server-limited PAT can hijack any stream by knowing
// the streamId.
func (s *NezhaHandler) StreamTarget(streamId string) (uint64, bool) {
s.ioStreamMutex.RLock()
defer s.ioStreamMutex.RUnlock()
ctx, ok := s.ioStreams[streamId]
if !ok {
return 0, false
}
return ctx.targetServerID, true
}
func (s *NezhaHandler) GetStream(streamId string) (*ioStreamContext, error) {
s.ioStreamMutex.RLock()
defer s.ioStreamMutex.RUnlock()
@@ -148,6 +222,37 @@ func (s *NezhaHandler) RevokeStreamsForServer(serverID uint64) {
}
}
// RevokeStreamsForPurpose tears down every IOStream tagged with the given
// purpose. Used as the IOStream half of the MCP kill switch: when the
// admin flips EnableMCP=false, any in-flight fs.transfer / fs.upload /
// fs.download must drop immediately rather than wait out the 5min IO
// timeout. Returns the number of streams revoked so the caller can log
// the blast radius.
func (s *NezhaHandler) RevokeStreamsForPurpose(purpose StreamPurpose) int {
s.ioStreamMutex.Lock()
defer s.ioStreamMutex.Unlock()
revoked := 0
for streamId, ctx := range s.ioStreams {
if ctx.purpose != purpose {
continue
}
ctx.revokedOnce.Do(func() {
if ctx.revokedCh != nil {
close(ctx.revokedCh)
}
})
if ctx.userIo != nil {
ctx.userIo.Close()
}
if ctx.agentIo != nil {
ctx.agentIo.Close()
}
delete(s.ioStreams, streamId)
revoked++
}
return revoked
}
func (s *NezhaHandler) CloseStream(streamId string) error {
s.ioStreamMutex.Lock()
defer s.ioStreamMutex.Unlock()
@@ -167,34 +272,50 @@ func (s *NezhaHandler) CloseStream(streamId string) error {
// UserConnected publishes the user-side IO under ioStreamMutex so concurrent
// Revoke* / WaitForAgent / StartStream see a consistent stream view.
// Without the lock, the bare assignment to stream.userIo races with the
// revoker's lock-protected read and triggers go-race.
func (s *NezhaHandler) UserConnected(streamId string, userIo io.ReadWriteCloser) error {
stream, err := s.GetStream(streamId)
if err != nil {
return err
s.ioStreamMutex.Lock()
stream, ok := s.ioStreams[streamId]
if !ok {
s.ioStreamMutex.Unlock()
return errors.New("stream not found")
}
stream.userIo = userIo
s.ioStreamMutex.Unlock()
stream.userIoChOnce.Do(func() {
close(stream.userIoConnectCh)
})
return nil
}
// AgentConnected is the agent-side dual of UserConnected. Same locking
// rationale.
func (s *NezhaHandler) AgentConnected(streamId string, agentIo io.ReadWriteCloser) error {
stream, err := s.GetStream(streamId)
if err != nil {
return err
s.ioStreamMutex.Lock()
stream, ok := s.ioStreams[streamId]
if !ok {
s.ioStreamMutex.Unlock()
return errors.New("stream not found")
}
stream.agentIo = agentIo
s.ioStreamMutex.Unlock()
stream.agentIoChOnce.Do(func() {
close(stream.agentIoConnectCh)
})
return nil
}
// streamEndpoints returns the user/agent IO under ioStreamMutex so callers
// never read the interface fields while UserConnected/AgentConnected write them.
func (s *NezhaHandler) streamEndpoints(stream *ioStreamContext) (userIo, agentIo io.ReadWriteCloser) {
s.ioStreamMutex.RLock()
defer s.ioStreamMutex.RUnlock()
return stream.userIo, stream.agentIo
}
func (s *NezhaHandler) StartStream(streamId string, timeout time.Duration) error {
stream, err := s.GetStream(streamId)
if err != nil {
@@ -202,62 +323,50 @@ func (s *NezhaHandler) StartStream(streamId string, timeout time.Duration) error
}
timeoutTimer := time.NewTimer(timeout)
defer timeoutTimer.Stop()
LOOP:
for {
select {
case <-stream.userIoConnectCh:
if stream.agentIo != nil {
timeoutTimer.Stop()
if _, agentIo := s.streamEndpoints(stream); agentIo != nil {
break LOOP
}
case <-stream.agentIoConnectCh:
if stream.userIo != nil {
timeoutTimer.Stop()
if userIo, _ := s.streamEndpoints(stream); userIo != nil {
break LOOP
}
case <-time.After(timeout):
case <-timeoutTimer.C:
break LOOP
}
time.Sleep(time.Millisecond * 500)
}
if stream.userIo == nil && stream.agentIo == nil {
userIo, agentIo := s.streamEndpoints(stream)
if userIo == nil && agentIo == nil {
return singleton.Localizer.ErrorT("timeout: no connection established")
}
if stream.userIo == nil {
if userIo == nil {
return singleton.Localizer.ErrorT("timeout: user connection not established")
}
if stream.agentIo == nil {
if agentIo == nil {
return singleton.Localizer.ErrorT("timeout: agent connection not established")
}
isDone := new(atomic.Bool)
endCh := make(chan struct{})
errCh := make(chan error, 2)
go func() {
bp := bufPool.Get().(*bp)
defer bufPool.Put(bp)
_, innerErr := io.CopyBuffer(stream.userIo, stream.agentIo, bp.buf)
if innerErr != nil {
err = innerErr
}
if isDone.CompareAndSwap(false, true) {
close(endCh)
}
_, innerErr := io.CopyBuffer(userIo, agentIo, bp.buf)
errCh <- innerErr
}()
go func() {
bp := bufPool.Get().(*bp)
defer bufPool.Put(bp)
_, innerErr := io.CopyBuffer(stream.agentIo, stream.userIo, bp.buf)
if innerErr != nil {
err = innerErr
}
if isDone.CompareAndSwap(false, true) {
close(endCh)
}
_, innerErr := io.CopyBuffer(agentIo, userIo, bp.buf)
errCh <- innerErr
}()
<-endCh
return err
return <-errCh
}
+92
View File
@@ -0,0 +1,92 @@
package rpc
import (
"io"
"sync"
"testing"
"time"
)
// nopRWC is a minimal io.ReadWriteCloser used for race tests; Close is a
// no-op so the racer goroutines do not panic on shared state.
type nopRWC struct{}
func (nopRWC) Read(p []byte) (int, error) { return 0, io.EOF }
func (nopRWC) Write(p []byte) (int, error) { return len(p), nil }
func (nopRWC) Close() error { return nil }
// H10 regression: UserConnected/AgentConnected mutate stream.userIo /
// stream.agentIo without holding ioStreamMutex, while WaitForAgent /
// RevokeStreamsForPurpose / RevokeStreamsForServer read & close the same
// fields under the lock. The go race detector catches it deterministically
// under -race; without the fix this test fails.
func TestIOStream_AgentConnectedIsRaceFreeUnderLock(t *testing.T) {
h := NewNezhaHandler()
const streamId = "race-test"
h.CreateStream(streamId, 1, 1)
t.Cleanup(func() { _ = h.CloseStream(streamId) })
var wg sync.WaitGroup
wg.Add(3)
go func() {
defer wg.Done()
// repeatedly attach an agent
for i := 0; i < 200; i++ {
_ = h.AgentConnected(streamId, nopRWC{})
}
}()
go func() {
defer wg.Done()
// concurrently attach a user
for i := 0; i < 200; i++ {
_ = h.UserConnected(streamId, nopRWC{})
}
}()
go func() {
defer wg.Done()
// Revoker takes the write lock and reads the same userIo/agentIo
// fields the unsynchronised writers above are setting. Use a real
// targetServerID (1) so RevokeStreamsForServer actually inspects
// the entry's userIo/agentIo before deleting.
for i := 0; i < 200; i++ {
h.RevokeStreamsForServer(1)
h.CreateStream(streamId, 1, 1)
}
}()
wg.Wait()
}
// StartStream reads stream.userIo/agentIo while it waits for both endpoints.
// Those reads must be lock-protected against the concurrent writes done by
// UserConnected/AgentConnected; otherwise -race flags the data race on the
// interface fields.
func TestIOStream_StartStreamReadsAreRaceFree(t *testing.T) {
h := NewNezhaHandler()
const streamId = "startstream-race"
h.CreateStream(streamId, 2, 2)
t.Cleanup(func() { _ = h.CloseStream(streamId) })
var wg sync.WaitGroup
wg.Add(3)
go func() {
defer wg.Done()
_ = h.StartStream(streamId, 50*time.Millisecond)
}()
go func() {
defer wg.Done()
time.Sleep(5 * time.Millisecond)
_ = h.AgentConnected(streamId, nopRWC{})
}()
go func() {
defer wg.Done()
time.Sleep(5 * time.Millisecond)
_ = h.UserConnected(streamId, nopRWC{})
}()
wg.Wait()
}
@@ -0,0 +1,55 @@
package rpc
import (
"sync"
"sync/atomic"
"testing"
pb "github.com/nezhahq/nezha/proto"
)
// updateConfig has no serialization, so two concurrent admin PATCH /setting
// requests can both flip EnableMCP true->false and both invoke
// CancelAllMCPInflight concurrently. The sweep must close each entry's cancel
// channel at most once; a non-atomic check-then-close double-closes the same
// channel and panics, crashing the dashboard.
func TestCancelAllMCPInflight_ConcurrentSweepsDoNotDoubleClose(t *testing.T) {
mcpInflight.Range(func(key, _ any) bool {
mcpInflight.Delete(key)
return true
})
t.Cleanup(func() {
mcpInflight.Range(func(key, _ any) bool {
mcpInflight.Delete(key)
return true
})
})
const entries = 256
for i := 0; i < entries; i++ {
mcpInflight.Store(uint64(i+1), &mcpInflightEntry{
serverID: uint64(i + 1),
result: make(chan *pb.TaskResult, 1),
cancel: make(chan struct{}),
cancelled: new(atomic.Bool),
})
}
const sweepers = 8
var wg sync.WaitGroup
wg.Add(sweepers)
for i := 0; i < sweepers; i++ {
go func() {
defer wg.Done()
// A double-close inside CancelAllMCPInflight panics here and
// fails the test (panic in a goroutine aborts the test binary).
CancelAllMCPInflight()
}()
}
wg.Wait()
mcpInflight.Range(func(key, _ any) bool {
t.Fatalf("inflight entry %v survived the sweep", key)
return false
})
}
+35
View File
@@ -0,0 +1,35 @@
package rpc
import (
"context"
"testing"
"time"
"github.com/nezhahq/nezha/model"
)
// H9 regression: CallAgent must consult mcpKillSwitchObserved before any
// side-effects. Without this gate, mcpEndpoint's EnableMCP read and
// CancelAllMCPInflight race against a fresh CallAgent that registers
// AFTER the cancel sweep, surviving the disabled state.
func TestCallAgent_RefusesWhenKillSwitchObserved(t *testing.T) {
prevCheck := mcpKillSwitchObserver()
SetMCPKillSwitchObserver(func() bool { return true })
t.Cleanup(func() { SetMCPKillSwitchObserver(prevCheck) })
_, err := CallAgent(context.Background(), 1, model.TaskTypeExec, struct{}{}, 50*time.Millisecond)
if err != ErrMCPDisabled {
t.Fatalf("CallAgent must short-circuit to ErrMCPDisabled when kill switch is observed, got %v", err)
}
}
// The default hook must keep production behaviour disarmed so tests and
// unconfigured deployments do not short-circuit CallAgent.
func TestCallAgent_KillSwitchHookDefaultsToDisarmed(t *testing.T) {
if mcpKillSwitchObserver() == nil {
t.Fatal("mcpKillSwitchObserver must always return a non-nil probe so dashboard can wire it")
}
if mcpKillSwitchObserver()() {
t.Fatal("default hook must return false so unconfigured dashboards / tests don't accidentally short-circuit CallAgent")
}
}
@@ -0,0 +1,80 @@
package rpc
import (
"context"
"errors"
"sync"
"testing"
"time"
"github.com/nezhahq/nezha/model"
)
// Registration-after-sweep race (review issue #1): CancelAllMCPInflight only
// cancels entries already present in the inflight map. A CallAgent that passes
// the upfront kill-switch check but has not yet Store()d its entry is invisible
// to the sweep, so without a post-registration re-check it goes on to SendTask
// a fresh exec/fs task to the agent AFTER EnableMCP=false.
//
// This test drives the worst-case interleaving deterministically:
// 1. CallAgent passes the upfront observer check (observer still false).
// 2. The operator flips the observer to "disabled" and runs the cancel sweep
// while CallAgent is paused between the check and Store.
// 3. CallAgent resumes; it MUST observe the kill switch on the post-Store
// re-check and return ErrMCPDisabled WITHOUT sending the task.
//
// With the race present, CallAgent sends the task and blocks until timeout
// (ErrAgentTimeout) — the agent received a fresh task past the kill switch.
func TestCallAgent_KillSwitchBeatsRegistrationAfterSweep(t *testing.T) {
const target uint64 = 7401
stream := newFakeStream()
cleanup := installFakeServer(t, target, stream)
defer cleanup()
var killed bool
var mu sync.Mutex
prev := mcpKillSwitchObserver()
// Observer returns the operator-controlled flag. CallAgent reads it both
// before and (with the fix) after registering the inflight entry.
SetMCPKillSwitchObserver(func() bool {
mu.Lock()
defer mu.Unlock()
return killed
})
t.Cleanup(func() { SetMCPKillSwitchObserver(prev) })
// Fail loudly if the agent ever receives a task: that means a fresh call
// leaked past the kill switch.
leaked := make(chan *struct{}, 1)
go func() {
select {
case <-stream.sent:
leaked <- nil
case <-time.After(2 * time.Second):
}
}()
// Arrange the interleaving: hook fires once CallAgent is about to register,
// flipping the kill switch and running the cancel sweep so the not-yet-Stored
// entry is missed by the sweep.
testKillSwitchAfterUpfrontCheck = func() {
mu.Lock()
killed = true
mu.Unlock()
CancelAllMCPInflight()
}
t.Cleanup(func() { testKillSwitchAfterUpfrontCheck = nil })
_, err := CallAgent(context.Background(), target, model.TaskTypeExec,
model.ExecRequest{Cmd: "x"}, 1*time.Second)
if !errors.Is(err, ErrMCPDisabled) {
t.Fatalf("CallAgent must return ErrMCPDisabled when the kill switch fires during registration; got %v", err)
}
select {
case <-leaked:
t.Fatal("a fresh MCP task leaked to the agent past the kill switch")
default:
}
}
+351
View File
@@ -0,0 +1,351 @@
package rpc
import (
"context"
"encoding/json"
"errors"
"log"
"sync"
"sync/atomic"
"time"
"github.com/nezhahq/nezha/model"
pb "github.com/nezhahq/nezha/proto"
"github.com/nezhahq/nezha/service/singleton"
)
// MCP 的"调用-响应"模式复用了 RequestTask 双向流:
// - dashboard 发 Task(带新分配的 taskID + JSON params
// - agent 执行后回 TaskResult(同 taskID + JSON result
// - RequestTask 接收循环把这种 TaskType 识别后路由到 inflight 等待方
//
// 不污染 model.Server 字段:用本包内的全局 inflight 表按 taskID 关联,
// 跨 server 共享单一命名空间。
var (
mcpTaskIDCounter atomic.Uint64
mcpInflight sync.Map // key: uint64 (taskID), value: chan *pb.TaskResult
)
// ErrMCPDisabled 是 CallAgent 在 MCP kill switch 被触发时返回的哨兵错误。
// 与 ErrAgentTimeout / ErrAgentOffline 平级,便于 controller 把它映射到
// MCPOutcomeForbidden 之类的审计 code 而不是误报 agent 故障。
var ErrMCPDisabled = errors.New("MCP is disabled by the dashboard administrator")
// mcpKillSwitchObserved is a process-level hook the dashboard wires to
// singleton.Conf.EnableMCP. CallAgent consults it before any side-effects so
// the entry-check / cancel-sweep / registration race cannot leak a fresh
// call past EnableMCP=false. Defaults to "disarmed" so tests and headless
// builds are unaffected.
//
// Stored behind atomic.Pointer because SetMCPKillSwitchObserver (startup +
// tests) and CallAgent (any RPC goroutine) touch it concurrently; a plain
// func variable is a data race under -race.
var mcpKillSwitchObserved atomic.Pointer[func() bool]
// disarmedKillSwitch is the default probe: never trips the kill switch.
var disarmedKillSwitch = func() bool { return false }
// testKillSwitchAfterUpfrontCheck, when non-nil, runs inside CallAgent between
// the upfront kill-switch check and the inflight registration. Production
// leaves it nil; tests use it to drive the registration-after-sweep race
// deterministically. Guarded by the same atomic.Pointer for race-freedom.
var testKillSwitchAfterUpfrontCheck func()
// SetMCPKillSwitchObserver installs the kill-switch probe the dashboard
// owns. Idempotent; the dashboard wires it at startup. Passing nil
// restores the default disarmed hook (used by tests to undo overrides).
func SetMCPKillSwitchObserver(fn func() bool) {
if fn == nil {
mcpKillSwitchObserved.Store(&disarmedKillSwitch)
return
}
mcpKillSwitchObserved.Store(&fn)
}
// mcpKillSwitchObserver returns the currently installed probe, never nil.
func mcpKillSwitchObserver() func() bool {
if p := mcpKillSwitchObserved.Load(); p != nil {
return *p
}
return disarmedKillSwitch
}
// allocateMCPTaskID 分配下一个 MCP 用的 task ID。
// 取 1<<32 起步以与可能存在的 cron/transfer 等已有 ID 空间错开(cron.id 由
// DB 自增,常量级,不会触及 1<<32)。
func allocateMCPTaskID() uint64 {
const base uint64 = 1 << 32
v := mcpTaskIDCounter.Add(1)
return base + v
}
// CallAgent 给 serverID 对应的 agent 发一条 MCP-RPC 风格的 Task,并阻塞等待 TaskResult 回包。
//
// taskType 必须是 model.IsMCPRPCResult 返回 true 的类型;params 会被 JSON 编码进 Task.Data。
// 超时由调用方控制;触发超时后从 inflight 表移除等待 slot(晚到的回包会被丢弃)。
//
// 错误语义:
// - server 未在线 / 未连接 task stream → ErrAgentOffline
// - 超时 → ctx.Err 或 ErrAgentTimeout
// - agent 回包 successful=false → 把 result.Data 当错误字符串返回
// - CancelAllMCPInflight 期间被中断 → ErrMCPDisabled
// - 任何 send 失败、序列化失败 → 原始 error
//
// 返回的 raw JSON 是 agent 端 TaskResult.Data 的原文。
func CallAgent(ctx context.Context, serverID uint64, taskType uint64, params any, timeout time.Duration) (json.RawMessage, error) {
if !model.IsMCPRPCResult(taskType) {
return nil, errors.New("CallAgent: task type is not registered as MCP RPC")
}
killSwitch := mcpKillSwitchObserver()
if killSwitch() {
return nil, ErrMCPDisabled
}
server, _ := singleton.ServerShared.Get(serverID)
if server == nil {
return nil, ErrAgentOffline
}
if server.GetTaskStream() == nil {
return nil, ErrAgentOffline
}
body, err := json.Marshal(params)
if err != nil {
return nil, err
}
taskID := allocateMCPTaskID()
resultCh := make(chan *pb.TaskResult, 1)
cancelCh := make(chan struct{})
entry := &mcpInflightEntry{
serverID: serverID,
result: resultCh,
cancel: cancelCh,
cancelled: new(atomic.Bool),
}
if hook := testKillSwitchAfterUpfrontCheck; hook != nil {
hook()
}
mcpInflight.Store(taskID, entry)
defer mcpInflight.Delete(taskID)
// Close the registration-after-sweep window: a kill switch that fired
// between the upfront check and this Store is invisible to
// CancelAllMCPInflight (our entry was not in the map yet). Because the
// operator sets EnableMCP=false BEFORE running the sweep, re-reading the
// observer here after Store guarantees we either see it disabled, or the
// sweep saw our now-registered entry and flipped entry.cancelled.
if killSwitch() || entry.cancelled.Load() {
return nil, ErrMCPDisabled
}
if err := server.SendTask(&pb.Task{
Id: taskID,
Type: taskType,
Data: string(body),
}); err != nil {
if errors.Is(err, model.ErrTaskStreamOffline) {
return nil, ErrAgentOffline
}
return nil, err
}
waitCtx := ctx
var cancel context.CancelFunc
if timeout > 0 {
waitCtx, cancel = context.WithTimeout(ctx, timeout)
defer cancel()
}
select {
case res := <-resultCh:
// Cancel must beat a late agent reply: Go select picks a random
// ready case, so if CancelAllMCPInflight closed cancelCh after the
// agent already filled resultCh we could still surface success.
// Re-check the cancel flag and prefer ErrMCPDisabled, matching the
// contract documented above ("CancelAllMCPInflight 期间被中断 →
// ErrMCPDisabled") and what TestUpdateConfig_DisablingMCPInvokesKillSwitch
// expects.
if entry.cancelled.Load() {
return nil, ErrMCPDisabled
}
if res == nil {
return nil, errors.New("agent returned nil result")
}
if !res.GetSuccessful() {
if res.GetData() != "" {
return nil, errors.New(res.GetData())
}
return nil, errors.New("agent returned unsuccessful result")
}
return json.RawMessage(res.GetData()), nil
case <-cancelCh:
return nil, ErrMCPDisabled
case <-waitCtx.Done():
if errors.Is(waitCtx.Err(), context.DeadlineExceeded) {
return nil, ErrAgentTimeout
}
return nil, waitCtx.Err()
}
}
// mcpInflightEntry binds an in-flight MCP call to its target serverID and
// pairs the result channel with a per-call cancel channel so the kill switch
// can break out of CallAgent without leaving the result channel dangling for
// the next late agent reply. The serverID is the authoritative reporter
// identity check at delivery time — without it deliverMCPResult would route
// purely by attacker-controlled TaskResult.Id (same bug class as commit
// 02129f1 in the cron path).
//
// cancelled flips to true the instant CancelAllMCPInflight observes the
// entry. Every code path that could complete the call — the CallAgent
// select on resultCh, deliverMCPResult, deliverMCPResultFromReporter —
// MUST consult it before treating an agent reply as authoritative, otherwise
// a TaskResult delivered concurrently with the kill switch can win Go's
// random select tiebreak and surface success after EnableMCP=false.
type mcpInflightEntry struct {
serverID uint64
result chan *pb.TaskResult
cancel chan struct{}
cancelled *atomic.Bool
closeOnce sync.Once
}
// closeCancel closes the entry's cancel channel exactly once. Concurrent
// CancelAllMCPInflight sweeps (two admin PATCH /setting requests both
// disabling MCP) would otherwise race a non-atomic check-then-close and
// panic on the second close.
func (e *mcpInflightEntry) closeCancel() {
e.closeOnce.Do(func() { close(e.cancel) })
}
// CancelAllMCPInflight closes every in-flight CallAgent so they return
// ErrMCPDisabled immediately. Used by the EnableMCP=false transition: by
// itself the inflight table holds the dashboard goroutine hostage until
// the agent replies (or the per-call timeout fires, up to ~305s for
// server.exec). Returns the number of calls cancelled for audit.
//
// Implementation notes:
// - Set the cancelled flag BEFORE closing cancelCh so any goroutine that
// already woke on resultCh observes it on the post-select re-check.
// - Delete the entry from mcpInflight immediately. Late agent replies via
// deliverMCPResult* would otherwise still find it (their own cancelled
// check covers concurrent delete, but evicting eagerly keeps the table
// small under repeated kill switch / re-enable cycles).
func CancelAllMCPInflight() int {
cancelled := 0
mcpInflight.Range(func(key, value any) bool {
entry, ok := value.(*mcpInflightEntry)
if !ok {
return true
}
if entry.cancelled != nil {
entry.cancelled.Store(true)
}
entry.closeCancel()
mcpInflight.Delete(key)
cancelled++
return true
})
return cancelled
}
// DeliverMCPResultForTest 暴露 deliverMCPResult 给跨包测试用:这是显式的
// "信任路径 / 不做 reporter 校验"入口,专给不关心来源的旧测试用。
// 安全敏感测试请用 DeliverMCPResultFromReporterForTest 并传入真实 reporterID。
func DeliverMCPResultForTest(res *pb.TaskResult) { deliverMCPResult(res) }
// DeliverMCPResultFromReporterForTest 暴露带 reporter 校验的投递入口给跨包
// 测试用,与生产 RequestTask 路径同语义:reporterID 必须等于 inflight 条目
// 登记的目标 serverID 才会投递。reporterID == 0 视为 "未知 reporter" 并被
// 拒绝;要绕过 reporter 校验请改用 DeliverMCPResultForTest。
func DeliverMCPResultFromReporterForTest(res *pb.TaskResult, reporterID uint64) {
deliverMCPResultFromReporter(res, reporterID)
}
// inflightServerIDForTest 返回某个 taskID 当前挂载的目标 serverID。用于安全
// 回归测试断言 inflight 条目确实把目标 server 绑进了路由表。
// 未找到时返回 (0, false)。
func inflightServerIDForTest(taskID uint64) (uint64, bool) {
v, ok := mcpInflight.Load(taskID)
if !ok {
return 0, false
}
entry, ok := v.(*mcpInflightEntry)
if !ok {
return 0, false
}
return entry.serverID, true
}
// deliverMCPResult 把 RequestTask 收到的 MCP-RPC TaskResult 路由到等待方。
// 找不到等待 slot(已超时被移除)则丢弃。
//
// 此变体不做 reporter 校验,仅用于不关心 reporter 的内部/测试路径。生产
// RequestTask 接收循环必须走 deliverMCPResultFromReporter,把 stream 上
// 已认证的 clientID 作为 reporter 传入。
func deliverMCPResult(res *pb.TaskResult) {
if res == nil {
return
}
v, ok := mcpInflight.Load(res.GetId())
if !ok {
return
}
entry, ok := v.(*mcpInflightEntry)
if !ok {
return
}
if entry.cancelled != nil && entry.cancelled.Load() {
return
}
select {
case entry.result <- res:
default:
}
}
// deliverMCPResultFromReporter 是生产路径的入口:要求 reporterID 与 inflight
// 条目登记的目标 serverID 一致才投递;否则丢弃并打日志。reporterID == 0
// 视为“未知 reporter”,安全起见也丢弃。
//
// 这条校验是必要的:mcpInflight 用全局递增 taskID 做键,跨 server 共享
// 单一命名空间;如果不在投递时核对上报 agent 是 CallAgent 的目标 server
// 任何已认证的恶意/失陷 agent 都能用猜到的 taskID 抢答其他 server 的
// MCP 调用(resultCh 容量 1,先到者覆盖真正回包)——和 commit 02129f1
// 在 cron 路径修过的攻击面同类。
func deliverMCPResultFromReporter(res *pb.TaskResult, reporterID uint64) {
if res == nil {
return
}
v, ok := mcpInflight.Load(res.GetId())
if !ok {
return
}
entry, ok := v.(*mcpInflightEntry)
if !ok {
return
}
if reporterID == 0 || entry.serverID != reporterID {
log.Printf("NEZHA>> MCP result ignored: taskID=%d targetServerID=%d reporterID=%d",
res.GetId(), entry.serverID, reporterID)
return
}
if entry.cancelled != nil && entry.cancelled.Load() {
return
}
select {
case entry.result <- res:
default:
}
}
// 错误类型
var (
ErrAgentOffline = errors.New("agent offline or task stream not connected")
ErrAgentTimeout = errors.New("agent did not respond within timeout")
)
+61
View File
@@ -0,0 +1,61 @@
package rpc
import (
"strings"
"sync/atomic"
"testing"
pb "github.com/nezhahq/nezha/proto"
)
// 把"测试 helper 的注释与运行时语义"钉成测试,避免 helper 文档骗读者:
//
// 1. DeliverMCPResultForTest 是显式的"信任路径 / 不做 reporter 校验"入口。
// 2. DeliverMCPResultFromReporterForTest 是带 reporter 校验的入口;
// reporterID == 0 视为"未知 reporter",必须被拒绝,不能像旧注释暗示的
// 那样当作"未知/不校验"放行。
//
// 这条契约决定了任何安全敏感的跨包测试调用方式:要绕过 reporter,
// 必须用 DeliverMCPResultForTest,而不是 reporterID=0 通过 reporter 入口。
func TestDeliverMCPResultFromReporterForTest_ZeroReporterIDIsRejected(t *testing.T) {
taskID := allocateMCPTaskID()
resultCh := make(chan *pb.TaskResult, 1)
cancelCh := make(chan struct{})
mcpInflight.Store(taskID, &mcpInflightEntry{
serverID: 7,
result: resultCh,
cancel: cancelCh,
cancelled: new(atomic.Bool),
})
t.Cleanup(func() { mcpInflight.Delete(taskID) })
DeliverMCPResultFromReporterForTest(&pb.TaskResult{Id: taskID, Data: "x", Successful: true}, 0)
select {
case <-resultCh:
t.Fatalf("reporterID==0 must be rejected by the reporter-checked helper; expected no delivery")
default:
}
}
// 同时把"测试 helper 自身的文档约束"钉到代码里:注释必须明确说出
// "reporterID == 0 视为未知 reporter 并被拒绝",否则未来维护者很容易看着
// "不校验"的旧措辞写出绕过 reporter 的安全敏感测试。
func TestDeliverMCPResultFromReporterForTest_DocStatesZeroIsRejected(t *testing.T) {
src := mustReadFile(t, "mcp_rpc.go")
if !strings.Contains(src, "DeliverMCPResultFromReporterForTest") {
t.Fatalf("expected helper to live in mcp_rpc.go")
}
// 提取 helper 上方的注释块:从 helper 名字往上找到第一段连续的 // 行。
idx := strings.Index(src, "func DeliverMCPResultFromReporterForTest(")
if idx < 0 {
t.Fatalf("helper not found in source")
}
prefix := src[:idx]
if !strings.Contains(prefix, "reporterID == 0") {
t.Fatalf("doc must mention reporterID == 0 contract explicitly")
}
if strings.Contains(prefix, "不校验") {
t.Fatalf("doc still claims reporterID==0 is 不校验; this contradicts deliverMCPResultFromReporter which drops it")
}
}
@@ -0,0 +1,95 @@
package rpc
import (
"context"
"errors"
"testing"
"time"
"github.com/nezhahq/nezha/model"
pb "github.com/nezhahq/nezha/proto"
)
// Kill switch must beat a late agent reply. Without the cancelled-flag
// re-check in CallAgent, the following sequence surfaces success after
// EnableMCP=false:
//
// t0 agent puts TaskResult into resultCh (capacity 1, non-blocking)
// t1 admin flips EnableMCP=false → CancelAllMCPInflight closes cancelCh
// t2 CallAgent's select sees BOTH cases ready; Go picks one at random;
// if it picks resultCh, the call returns the agent's payload even
// though the operator's kill switch fired.
//
// The fix is to mark the entry cancelled BEFORE closing cancelCh and have
// the resultCh branch re-check that flag. This test pins the contract by
// driving the worst-case ordering: result is delivered FIRST, then the
// kill switch fires, then CallAgent observes both. With the race in place
// this would flake (random select); with the fix it always returns
// ErrMCPDisabled.
func TestCallAgent_KillSwitchBeatsConcurrentLateResult(t *testing.T) {
const target uint64 = 7301
stream := newFakeStream()
cleanup := installFakeServer(t, target, stream)
defer cleanup()
delivered := make(chan struct{})
go func() {
sent := <-stream.sent
deliverMCPResultFromReporter(&pb.TaskResult{
Id: sent.GetId(),
Type: model.TaskTypeExec,
Successful: true,
Data: `{"exit_code":0,"stdout":"should-not-surface"}`,
}, target)
CancelAllMCPInflight()
close(delivered)
}()
_, err := CallAgent(context.Background(), target, model.TaskTypeExec,
model.ExecRequest{Cmd: "x"}, 2*time.Second)
<-delivered
if !errors.Is(err, ErrMCPDisabled) {
t.Fatalf("kill switch must win the race with a late agent reply; want ErrMCPDisabled, got %v", err)
}
}
// CancelAllMCPInflight must eagerly evict entries so a stale TaskResult
// that arrives after the kill switch cannot still land in resultCh.
// Without the cancelled flag this entry would still be reachable through
// deliverMCPResultFromReporter; the flag guarantees the late delivery is
// silently dropped even if the caller has not returned yet.
func TestCancelAllMCPInflight_LaterResultIsSwallowed(t *testing.T) {
const target uint64 = 7302
stream := newFakeStream()
cleanup := installFakeServer(t, target, stream)
defer cleanup()
taskIDCh := make(chan uint64, 1)
go func() {
sent := <-stream.sent
taskIDCh <- sent.GetId()
}()
resultCh := make(chan error, 1)
go func() {
_, err := CallAgent(context.Background(), target, model.TaskTypeExec,
model.ExecRequest{Cmd: "x"}, 5*time.Second)
resultCh <- err
}()
taskID := <-taskIDCh
CancelAllMCPInflight()
if err := <-resultCh; !errors.Is(err, ErrMCPDisabled) {
t.Fatalf("CallAgent must return ErrMCPDisabled after kill switch; got %v", err)
}
deliverMCPResultFromReporter(&pb.TaskResult{
Id: taskID,
Type: model.TaskTypeExec,
Successful: true,
Data: `{"exit_code":0}`,
}, target)
}
+149
View File
@@ -0,0 +1,149 @@
package rpc
import (
"context"
"encoding/json"
"errors"
"testing"
"time"
"github.com/nezhahq/nezha/model"
pb "github.com/nezhahq/nezha/proto"
)
// These tests pin the security invariant that an MCP TaskResult delivered
// back through RequestTask must come from the SAME agent the CallAgent was
// targeted at. The receive loop in service/rpc/nezha.go has the authenticated
// clientID in scope; deliverMCPResult must consume it and reject mismatches.
//
// Why the invariant matters: mcpInflight is keyed by a globally increasing
// counter (allocateMCPTaskID) and the lookup table is shared across servers.
// Without binding the inflight entry to the target serverID and verifying it
// against the reporter clientID, any compromised agent A can race a forged
// TaskResult for server B's CallAgent (resultCh capacity is 1; first reply
// wins, real reply is dropped). The same class of attack motivated the cron
// path's CanReportCronResult and the transfer path's pending.ID == result.Id
// check in this very file's RequestTask switch.
// TestDeliverMCPResult_RejectsForeignReporter is the security regression: a
// reporter that is NOT the call target must not be able to deliver into
// another server's inflight slot, even with a correctly-guessed taskID.
func TestDeliverMCPResult_RejectsForeignReporter(t *testing.T) {
const (
targetServerID uint64 = 6101
foreignAgentID uint64 = 6102
)
stream := newFakeStream()
cleanup := installFakeServer(t, targetServerID, stream)
defer cleanup()
captured := make(chan uint64, 1)
go func() {
sent := <-stream.sent
// Foreign agent racing a forged TaskResult with the right taskID.
DeliverMCPResultFromReporterForTest(&pb.TaskResult{
Id: sent.GetId(),
Type: model.TaskTypeExec,
Successful: true,
Data: `{"exit_code":0,"stdout":"forged"}`,
}, foreignAgentID)
captured <- sent.GetId()
}()
_, err := CallAgent(context.Background(), targetServerID, model.TaskTypeExec,
model.ExecRequest{Cmd: "x"}, 200*time.Millisecond)
if !errors.Is(err, ErrAgentTimeout) {
t.Fatalf("forged result from foreign reporter must NOT deliver; want ErrAgentTimeout, got %v", err)
}
select {
case <-captured:
case <-time.After(time.Second):
t.Fatalf("test stream never observed the dispatched task")
}
}
// TestDeliverMCPResult_AcceptsMatchingReporter is the green companion: when
// the reporter clientID matches the inflight target, the result must still
// route correctly (we are not breaking the happy path).
func TestDeliverMCPResult_AcceptsMatchingReporter(t *testing.T) {
const targetServerID uint64 = 6103
stream := newFakeStream()
cleanup := installFakeServer(t, targetServerID, stream)
defer cleanup()
want := model.ExecResult{ExitCode: 0, Stdout: "ok"}
payload, _ := json.Marshal(want)
go func() {
sent := <-stream.sent
DeliverMCPResultFromReporterForTest(&pb.TaskResult{
Id: sent.GetId(),
Type: model.TaskTypeExec,
Successful: true,
Data: string(payload),
}, targetServerID)
}()
raw, err := CallAgent(context.Background(), targetServerID, model.TaskTypeExec,
model.ExecRequest{Cmd: "x"}, 2*time.Second)
if err != nil {
t.Fatalf("matching reporter must deliver, got %v", err)
}
var got model.ExecResult
if err := json.Unmarshal(raw, &got); err != nil {
t.Fatalf("bad result json: %v", err)
}
if got.Stdout != "ok" {
t.Fatalf("payload not propagated, got %+v", got)
}
}
// TestDeliverMCPResult_InflightEntryBoundToServerID locks in the structural
// requirement that the inflight table records the target serverID. Without
// this binding deliverMCPResult cannot perform the reporter check above.
// Probing via reflection avoids exporting mcpInflight just for tests.
func TestDeliverMCPResult_InflightEntryBoundToServerID(t *testing.T) {
const targetServerID uint64 = 6104
stream := newFakeStream()
cleanup := installFakeServer(t, targetServerID, stream)
defer cleanup()
gotEntry := make(chan struct {
taskID uint64
serverID uint64
found bool
}, 1)
go func() {
sent := <-stream.sent
taskID := sent.GetId()
serverID, ok := inflightServerIDForTest(taskID)
gotEntry <- struct {
taskID uint64
serverID uint64
found bool
}{taskID, serverID, ok}
// Unblock CallAgent so the inflight slot is cleaned up.
DeliverMCPResultFromReporterForTest(&pb.TaskResult{
Id: taskID,
Type: model.TaskTypeExec,
Successful: true,
Data: "{}",
}, targetServerID)
}()
_, err := CallAgent(context.Background(), targetServerID, model.TaskTypeExec,
model.ExecRequest{Cmd: "x"}, 2*time.Second)
if err != nil {
t.Fatalf("unexpected CallAgent error: %v", err)
}
probe := <-gotEntry
if !probe.found {
t.Fatalf("inflight entry for taskID=%d not found while CallAgent was blocking", probe.taskID)
}
if probe.serverID != targetServerID {
t.Fatalf("inflight entry must carry target serverID=%d, got %d", targetServerID, probe.serverID)
}
}
+165
View File
@@ -0,0 +1,165 @@
package rpc
import (
"context"
"encoding/json"
"errors"
"sync/atomic"
"testing"
"time"
"google.golang.org/grpc/metadata"
"github.com/nezhahq/nezha/model"
pb "github.com/nezhahq/nezha/proto"
"github.com/nezhahq/nezha/service/singleton"
)
type fakeTaskStream struct {
sent chan *pb.Task
delay time.Duration
}
func newFakeStream() *fakeTaskStream {
return &fakeTaskStream{sent: make(chan *pb.Task, 4)}
}
func (s *fakeTaskStream) Send(t *pb.Task) error { s.sent <- t; return nil }
func (s *fakeTaskStream) Recv() (*pb.TaskResult, error) { return nil, context.Canceled }
func (s *fakeTaskStream) SetHeader(metadata.MD) error { return nil }
func (s *fakeTaskStream) SendHeader(metadata.MD) error { return nil }
func (s *fakeTaskStream) SetTrailer(metadata.MD) {}
func (s *fakeTaskStream) Context() context.Context { return context.Background() }
func (s *fakeTaskStream) SendMsg(any) error { return nil }
func (s *fakeTaskStream) RecvMsg(any) error { return context.Canceled }
func installFakeServer(t *testing.T, id uint64, stream pb.NezhaService_RequestTaskServer) func() {
t.Helper()
original := singleton.ServerShared
sc := singleton.NewEmptyServerClassForTest()
srv := &model.Server{}
srv.ID = id
srv.SetTaskStream(stream)
sc.InsertForTest(srv)
singleton.ServerShared = sc
return func() { singleton.ServerShared = original }
}
func TestCallAgent_RejectsNonMCPType(t *testing.T) {
_, err := CallAgent(context.Background(), 1, model.TaskTypeCommand, struct{}{}, time.Second)
if err == nil {
t.Fatalf("expected error for non-MCP type")
}
}
func TestCallAgent_OfflineWhenNoStream(t *testing.T) {
original := singleton.ServerShared
sc := singleton.NewEmptyServerClassForTest()
srv := &model.Server{}
srv.ID = 7
sc.InsertForTest(srv)
singleton.ServerShared = sc
t.Cleanup(func() { singleton.ServerShared = original })
_, err := CallAgent(context.Background(), 7, model.TaskTypeExec, struct{}{}, time.Second)
if !errors.Is(err, ErrAgentOffline) {
t.Fatalf("expected ErrAgentOffline, got %v", err)
}
}
func TestCallAgent_HappyPath_DelivlersResultByTaskID(t *testing.T) {
stream := newFakeStream()
cleanup := installFakeServer(t, 42, stream)
defer cleanup()
resultPayload, _ := json.Marshal(model.ExecResult{ExitCode: 0, Stdout: "hello"})
var captured atomic.Uint64
done := make(chan struct{})
go func() {
sent := <-stream.sent
captured.Store(sent.GetId())
deliverMCPResult(&pb.TaskResult{
Id: sent.GetId(),
Type: model.TaskTypeExec,
Data: string(resultPayload),
Successful: true,
})
close(done)
}()
raw, err := CallAgent(context.Background(), 42, model.TaskTypeExec, model.ExecRequest{Cmd: "x"}, 5*time.Second)
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
<-done
if captured.Load() == 0 {
t.Fatalf("task id never captured")
}
var got model.ExecResult
if err := json.Unmarshal(raw, &got); err != nil {
t.Fatalf("bad result json: %v", err)
}
if got.Stdout != "hello" {
t.Fatalf("payload not propagated, got %+v", got)
}
}
func TestCallAgent_Timeout(t *testing.T) {
stream := newFakeStream()
cleanup := installFakeServer(t, 43, stream)
defer cleanup()
go func() { <-stream.sent }()
_, err := CallAgent(context.Background(), 43, model.TaskTypeFsRead, model.FsReadRequest{Path: "/x"}, 50*time.Millisecond)
if !errors.Is(err, ErrAgentTimeout) {
t.Fatalf("expected ErrAgentTimeout, got %v", err)
}
}
func TestCallAgent_LateResultIsDropped(t *testing.T) {
stream := newFakeStream()
cleanup := installFakeServer(t, 44, stream)
defer cleanup()
var taskID uint64
got := make(chan struct{})
go func() {
sent := <-stream.sent
taskID = sent.GetId()
close(got)
}()
_, err := CallAgent(context.Background(), 44, model.TaskTypeFsDelete, model.FsDeleteRequest{Path: "/x"}, 50*time.Millisecond)
if !errors.Is(err, ErrAgentTimeout) {
t.Fatalf("expected timeout")
}
<-got
deliverMCPResult(&pb.TaskResult{Id: taskID, Type: model.TaskTypeFsDelete, Successful: true, Data: "{}"})
if _, ok := mcpInflight.Load(taskID); ok {
t.Fatalf("inflight entry must be cleaned up after timeout")
}
}
func TestCallAgent_UnsuccessfulIsError(t *testing.T) {
stream := newFakeStream()
cleanup := installFakeServer(t, 45, stream)
defer cleanup()
go func() {
sent := <-stream.sent
deliverMCPResult(&pb.TaskResult{
Id: sent.GetId(),
Type: sent.GetType(),
Successful: false,
Data: "agent says nope",
})
}()
_, err := CallAgent(context.Background(), 45, model.TaskTypeFsWrite, model.FsWriteRequest{Path: "/x", Content: "y"}, time.Second)
if err == nil || err.Error() != "agent says nope" {
t.Fatalf("expected agent error message, got %v", err)
}
}
+66 -13
View File
@@ -37,6 +37,32 @@ func NewNezhaHandler() *NezhaHandler {
}
}
// attachRequestTaskStream resolves the server for clientID and publishes the
// task stream. It mirrors the !ok || server == nil guard the other RPC entry
// points use: the server can be deleted between CheckRequestTask and this
// lookup, in which case Get returns a nil *Server and SetTaskStream would
// panic.
func attachRequestTaskStream(clientID uint64, stream pb.NezhaService_RequestTaskServer) (*model.Server, bool) {
server, ok := singleton.ServerShared.Get(clientID)
if !ok || server == nil {
return nil, false
}
server.SetTaskStream(stream)
return server, true
}
// clearRequestTaskStream detaches the dropped stream from whichever *Server is
// currently published for clientID. Edit and transfer rotation publish a new
// *Server that adopts the same stream holder, so cleanup must target the live
// map entry; the captured server is only the fallback for a removed entry.
func clearRequestTaskStream(clientID uint64, captured *model.Server, stream pb.NezhaService_RequestTaskServer) {
if current, ok := singleton.ServerShared.Get(clientID); ok && current != nil {
current.ClearTaskStreamIfCurrent(stream)
return
}
captured.ClearTaskStreamIfCurrent(stream)
}
func (s *NezhaHandler) RequestTask(stream pb.NezhaService_RequestTaskServer) error {
var clientID uint64
var err error
@@ -44,9 +70,11 @@ func (s *NezhaHandler) RequestTask(stream pb.NezhaService_RequestTaskServer) err
return err
}
server, _ := singleton.ServerShared.Get(clientID)
server.SetTaskStream(stream)
defer server.ClearTaskStreamIfCurrent(stream)
server, ok := attachRequestTaskStream(clientID, stream)
if !ok {
return nil
}
defer clearRequestTaskStream(clientID, server, stream)
// If a transfer is mid-flight for this server, the agent has just brought
// up a fresh bidi stream — this is the moment to (re)deliver the
// ApplyConfig task carrying the new owner's AgentSecret. Pushes from
@@ -114,6 +142,10 @@ func (s *NezhaHandler) RequestTask(stream pb.NezhaService_RequestTaskServer) err
log.Printf("NEZHA>> ServerTransfer MarkFailed(%d) failed: %v", result.GetId(), err)
}
default:
if model.IsMCPRPCResult(result.GetType()) {
deliverMCPResultFromReporter(result, clientID)
continue
}
if model.IsServiceSentinelNeeded(result.GetType()) {
singleton.ServiceSentinelShared.Dispatch(singleton.ReportData{
Data: result,
@@ -265,24 +297,45 @@ func (s *NezhaHandler) IOStream(stream pb.NezhaService_IOStreamServer) error {
return fmt.Errorf("stream not authorized for agent")
}
go func() {
for {
if err := stream.Send(&pb.IOStreamData{Data: []byte{}}); err != nil {
log.Printf("NEZHA>> IOStream keepAlive error: %v\n", err)
return
}
time.Sleep(time.Second * 30)
}
}()
if _, err := s.GetStream(streamId); err != nil {
return err
}
iw := grpcx.NewIOStreamWrapper(stream)
// Keepalive MUST go through the wrapper so it shares the same sendMu as
// MCP fs.transfer / terminal / fm Writers. Calling stream.Send directly
// here used to race those Writers — grpc-go forbids concurrent SendMsg
// on the same stream. The wrapper's sendMu is the dashboard-side dual
// of agent/cmd/agent/mcp_fs_transfer.go's serialIOStreamSender.
keepaliveDone := make(chan struct{})
go func() {
defer close(keepaliveDone)
ticker := time.NewTicker(time.Second * 30)
defer ticker.Stop()
for {
select {
case <-iw.Context().Done():
return
case <-iw.Done():
// 业务侧(CloseStream / RevokeStreamsForPurpose)调过
// iw.Close()。即便底层 gRPC stream context 尚未取消,也
// 必须立刻收手——否则要再等一整个 30s tickhandler 在
// iw.Wait() 之后又得多等一拍 keepaliveDone 才能返回。
return
case <-ticker.C:
if err := iw.SendKeepalive(); err != nil {
log.Printf("NEZHA>> IOStream keepAlive error: %v\n", err)
return
}
}
}
}()
if err := s.AgentConnected(streamId, iw); err != nil {
return err
}
iw.Wait()
<-keepaliveDone
return nil
}
@@ -0,0 +1,25 @@
package rpc
import (
"testing"
"github.com/nezhahq/nezha/model"
"github.com/nezhahq/nezha/service/singleton"
)
func TestAttachRequestTaskStream_MissingServerDoesNotPanic(t *testing.T) {
reporter := requestTaskSecurityServer(7, 200, "eeeeeeee-eeee-eeee-eeee-eeeeeeeeeeee")
setupRequestTaskSecurityFixture(t, []*model.Server{reporter}, nil, map[uint64]model.UserInfo{
200: {Role: model.RoleMember},
}, map[string]uint64{"reporter-secret": 200})
singleton.ServerShared.Delete([]uint64{reporter.ID})
srv, ok := attachRequestTaskStream(reporter.ID, nil)
if ok {
t.Fatal("attach must report not-ok when the server was deleted between auth and lookup")
}
if srv != nil {
t.Fatalf("attach must return a nil server for a deleted id, got %#v", srv)
}
}
@@ -0,0 +1,46 @@
package rpc
import (
"context"
"errors"
"testing"
"github.com/nezhahq/nezha/model"
"github.com/nezhahq/nezha/service/singleton"
)
// When a server is edited mid-session, updateServer swaps a new *Server into
// ServerShared that adopts the live stream holder. The agent's RequestTask
// cleanup must detach the stream from whichever *Server is currently published,
// not the stale object captured when the stream attached — otherwise the new
// object keeps reporting the agent as online on a dead stream.
func TestRequestTaskCleanupDetachesStreamFromCurrentServerAfterEdit(t *testing.T) {
reporter := requestTaskSecurityServer(7, 200, "ffffffff-ffff-ffff-ffff-ffffffffffff")
setupRequestTaskSecurityFixture(t, []*model.Server{reporter}, nil, map[uint64]model.UserInfo{
200: {Role: model.RoleMember},
}, map[string]uint64{"reporter-secret": 200})
old, ok := singleton.ServerShared.Get(reporter.ID)
if !ok {
t.Fatalf("server %d not found", reporter.ID)
}
stream := requestTaskSecurityAuthedStream("reporter-secret", reporter.UUID)
stream.onRecv = func() {
edited := &model.Server{Common: model.Common{ID: old.ID, UserID: old.UserID}, UUID: old.UUID, Name: "edited"}
edited.CopyFromRunningServer(old)
singleton.ServerShared.Update(edited, "")
}
if err := NewNezhaHandler().RequestTask(stream); !errors.Is(err, context.Canceled) {
t.Fatalf("expected RequestTask to finish after Recv error, got %v", err)
}
current, ok := singleton.ServerShared.Get(reporter.ID)
if !ok {
t.Fatalf("server %d not found after edit", reporter.ID)
}
if got := current.GetTaskStream(); got != nil {
t.Fatalf("edited server must report offline after the agent stream dropped, got %T", got)
}
}
+20
View File
@@ -0,0 +1,20 @@
package rpc
import (
"os"
"path/filepath"
"testing"
)
func mustReadFile(t *testing.T, name string) string {
t.Helper()
wd, err := os.Getwd()
if err != nil {
t.Fatalf("getwd: %v", err)
}
b, err := os.ReadFile(filepath.Join(wd, name))
if err != nil {
t.Fatalf("read %s: %v", name, err)
}
return string(b)
}
+56
View File
@@ -0,0 +1,56 @@
package rpc
import (
"context"
"testing"
"time"
)
// WaitForAgent 必须在 stream 被 RevokeStreamsForPurpose 强制下线时立刻返回,
// 否则 EnableMCP=false 的 kill switch 不能真的“立即”切断那些还卡在
// “等待 agent attach” 阶段的 transfer 请求 —— 它们会一直等到 timeout
// (生产路径上是 30 秒)。
//
// 期望行为:调用 RevokeStreamsForPurpose 后,WaitForAgent 在远小于 timeout
// 的时间内返回 (nil, false)。
func TestWaitForAgent_RevokeWakesUpWaiter(t *testing.T) {
h := NewNezhaHandler()
const streamID = "kill-switch-wait"
h.CreateStreamWithPurpose(streamID, 0, 7, PurposeMCPTransfer)
done := make(chan struct {
io any
ok bool
dur time.Duration
}, 1)
start := time.Now()
go func() {
// 给一个明显大于 revoke 触发延时的 timeout;如果 revoke 没唤醒,
// WaitForAgent 会一直等到这里,下面的 assertion 就会失败。
stream, ok := h.WaitForAgent(context.Background(), streamID, 5*time.Second)
done <- struct {
io any
ok bool
dur time.Duration
}{stream, ok, time.Since(start)}
}()
// 让 WaitForAgent 真的进入 select 等待,再触发 kill switch。
time.Sleep(50 * time.Millisecond)
if revoked := h.RevokeStreamsForPurpose(PurposeMCPTransfer); revoked != 1 {
t.Fatalf("expected to revoke exactly 1 MCP stream, got %d", revoked)
}
select {
case res := <-done:
if res.ok {
t.Fatalf("WaitForAgent must return ok=false after revoke; got ok=true")
}
if res.dur > time.Second {
t.Fatalf("WaitForAgent did not wake up promptly after revoke (took %s); kill switch is not immediate", res.dur)
}
case <-time.After(2 * time.Second):
t.Fatalf("WaitForAgent never returned after revoke; kill switch did not wake the waiter")
}
}
+11 -6
View File
@@ -264,13 +264,12 @@ func CronTrigger(cr *model.Cron, triggerServer ...uint64) func() {
if !cronCanSendToServer(cr, s) {
return
}
stream := s.GetTaskStream()
if stream != nil {
if s.GetTaskStream() != nil {
cronShared := CronShared
if cronShared != nil {
cronShared.reserveAlertTriggerCronResult(cr.ID, s.ID)
}
if err := stream.Send(&pb.Task{
if err := s.SendTask(&pb.Task{
Id: cr.ID,
Data: cr.Command,
Type: model.TaskTypeCommand,
@@ -287,7 +286,13 @@ func CronTrigger(cr *model.Cron, triggerServer ...uint64) func() {
return
}
for _, s := range ServerShared.Range {
// 先在锁内快照 server 列表再逐个 SendTaskServerShared.Range 会在整个
// 回调期间持 listMu.RLock,而 SendTask 走阻塞 gRPC,一个卡死的 agent
// 会让需要写锁的 server 编辑/删除被拖死。GetList 克隆后即释放锁。
for _, s := range ServerShared.GetList() {
if s == nil {
continue
}
if !cronCanSendToServer(cr, s) {
continue
}
@@ -297,8 +302,8 @@ func CronTrigger(cr *model.Cron, triggerServer ...uint64) func() {
if cr.Cover == model.CronCoverIgnoreAll && !crIgnoreMap[s.ID] {
continue
}
if stream := s.GetTaskStream(); stream != nil {
stream.Send(&pb.Task{
if s.GetTaskStream() != nil {
_ = s.SendTask(&pb.Task{
Id: cr.ID,
Data: cr.Command,
Type: model.TaskTypeCommand,
+35 -2
View File
@@ -38,9 +38,39 @@ func NewServerClass() *ServerClass {
}
sc.sortList()
model.OwnerServerIDsLookup = sc.ownerServerIDs
model.AllServerIDsLookup = sc.allServerIDs
model.OwnerIsAdminLookup = ownerIsAdmin
return sc
}
func (c *ServerClass) ownerServerIDs(ownerUID uint64) []uint64 {
var ids []uint64
c.Range(func(id uint64, s *model.Server) bool {
if s != nil && s.GetUserID() == ownerUID {
ids = append(ids, id)
}
return true
})
return ids
}
func (c *ServerClass) allServerIDs() []uint64 {
var ids []uint64
c.Range(func(id uint64, s *model.Server) bool {
if s != nil {
ids = append(ids, id)
}
return true
})
return ids
}
func ownerIsAdmin(ownerUID uint64) bool {
return userIsAdmin(ownerUID)
}
func (c *ServerClass) Update(s *model.Server, uuid string) {
c.listMu.Lock()
@@ -64,8 +94,11 @@ func (c *ServerClass) Delete(idList []uint64) {
c.listMu.Lock()
for _, id := range idList {
serverUUID := c.list[id].UUID
delete(c.uuidToID, serverUUID)
s, ok := c.list[id]
if !ok {
continue
}
delete(c.uuidToID, s.UUID)
delete(c.list, id)
}
@@ -0,0 +1,32 @@
package singleton
import (
"testing"
"github.com/nezhahq/nezha/model"
)
func TestServerClassDeleteMissingIDNoPanic(t *testing.T) {
c := &ServerClass{
class: class[uint64, *model.Server]{
list: map[uint64]*model.Server{
1: {Common: model.Common{ID: 1}, UUID: "uuid-1"},
},
},
uuidToID: map[string]uint64{"uuid-1": 1},
}
c.Delete([]uint64{999999})
if _, ok := c.list[1]; !ok {
t.Fatalf("existing server 1 must remain after deleting a non-existent id")
}
c.Delete([]uint64{1, 424242})
if _, ok := c.list[1]; ok {
t.Fatalf("server 1 should be removed")
}
if _, ok := c.uuidToID["uuid-1"]; ok {
t.Fatalf("uuid mapping for server 1 should be removed")
}
}
+8 -15
View File
@@ -232,20 +232,6 @@ func NewServerTransferClass() *ServerTransferClass {
}
c.pending[t.ServerID] = &t
}
for i := range pending {
t := pending[i]
// Skip ghost rows whose server has been deleted out from under
// the transfer (e.g. before OnServersDeleted existed, or because
// the row predates this branch). Loading them would resurrect a
// HasPending state that no longer corresponds to a real server
// and the timeout sweeper would log errors every 30s without
// being able to settle the row.
if s, ok := ServerShared.Get(t.ServerID); !ok || s == nil {
log.Printf("NEZHA>> ServerTransferClass: ignoring pending transfer %d for missing server %d (likely a leftover from before OnServersDeleted was wired)", t.ID, t.ServerID)
continue
}
c.pending[t.ServerID] = &t
}
var reverted []model.ServerTransfer
// acked_at IS NULL is non-negotiable: MarkRevertDelivered persists
@@ -927,7 +913,14 @@ func (c *ServerTransferClass) sendApplyConfigTask(s *model.Server, stream pb.Nez
// Keep Send synchronous under the per-server lock. A goroutine+timeout cannot
// cancel grpc.ServerStream.Send; returning early would let a stale new-secret
// ApplyConfig complete after a cancel/fail revert and overwrite the rollback.
if err := stream.Send(task); err != nil {
//
// Route through Server.SendTask so the holder-scoped send mutex is
// honoured: cron / MCP CallAgent / MCP fs.transfer dispatch on the same
// gRPC stream and would otherwise race grpc-go's one-SendMsg-per-stream
// invariant. The captured stream argument is still passed to
// ClearTaskStreamIfCurrent so a reconnect mid-Send cannot wipe a newer
// published stream when Send fails on the stale one.
if err := s.SendTask(task); err != nil {
log.Printf("NEZHA>> ServerTransfer ApplyConfig send failed: serverID=%d transferID=%d: %v", s.ID, task.Id, err)
s.ClearTaskStreamIfCurrent(stream)
return err
+11 -1
View File
@@ -94,11 +94,21 @@ func InitDBFromPath(path string) error {
model.Notification{}, model.AlertRule{}, model.Service{}, model.NotificationGroupNotification{},
model.Cron{}, model.Transfer{}, model.ServerGroupServer{},
model.NAT{}, model.DDNSProfile{}, model.NotificationGroupNotification{},
model.WAF{}, model.Oauth2Bind{}, model.ServerTransfer{}, model.JWTSession{})
model.WAF{}, model.Oauth2Bind{}, model.ServerTransfer{}, model.JWTSession{},
model.APIToken{}, model.MCPAuditLog{})
if err != nil {
return err
}
// 旧 mcp:* scope 与 nezha:* 并行了一段时间,HasScope 通过别名让 mcp:fs:write
// 静默扩到 REST nezha:server:write。统一命名后这里把残留旧 scope 一次性
// 归一化(或在仅剩危险旧 scope 时整张 PAT 删除),保证运行时不再依赖别名。
if rewritten, deleted, mErr := model.MigrateLegacyMCPScopes(DB); mErr != nil {
log.Printf("NEZHA>> MigrateLegacyMCPScopes failed: %v", mErr)
} else if rewritten > 0 || deleted > 0 {
log.Printf("NEZHA>> Migrated legacy mcp:* api token scopes: rewritten=%d deleted=%d", rewritten, deleted)
}
return nil
}
+65
View File
@@ -0,0 +1,65 @@
package singleton
import "github.com/nezhahq/nezha/model"
// NewEmptyServerClassForTest 构造一个不依赖 DB 的空 ServerClass,仅用于单测。
// 生产路径请用 NewServerClass。
func NewEmptyServerClassForTest() *ServerClass {
sc := &ServerClass{
class: class[uint64, *model.Server]{
list: make(map[uint64]*model.Server),
},
uuidToID: make(map[string]uint64),
}
model.OwnerServerIDsLookup = sc.ownerServerIDs
model.AllServerIDsLookup = sc.allServerIDs
model.OwnerIsAdminLookup = ownerIsAdmin
return sc
}
// InsertForTest 把一个 server 直接塞进内存表与排序快照,跳过 DB & InitServer 逻辑。
// 调用方需保证 server.ID 已经设置。
func (c *ServerClass) InsertForTest(s *model.Server) {
c.listMu.Lock()
c.list[s.ID] = s
if s.UUID != "" {
c.uuidToID[s.UUID] = s.ID
}
c.listMu.Unlock()
c.sortList()
}
// NewEmptyDDNSClassForTest 构造一个不依赖 DB 的空 DDNSClass,仅用于单测。
func NewEmptyDDNSClassForTest() *DDNSClass {
return &DDNSClass{
class: class[uint64, *model.DDNSProfile]{
list: make(map[uint64]*model.DDNSProfile),
},
}
}
// InsertForTest 把一个 DDNS profile 直接塞进内存表,跳过 DB。
func (c *DDNSClass) InsertForTest(p *model.DDNSProfile) {
c.listMu.Lock()
c.list[p.ID] = p
c.listMu.Unlock()
}
// NewEmptyNotificationClassForTest 构造空 NotificationClass。
func NewEmptyNotificationClassForTest() *NotificationClass {
return &NotificationClass{
class: class[uint64, *model.Notification]{
list: make(map[uint64]*model.Notification),
},
groupToIDList: make(map[uint64]map[uint64]*model.Notification),
idToGroupList: make(map[uint64]map[uint64]struct{}),
groupList: make(map[uint64]string),
}
}
// InsertForTest 把一个 Notification 直接塞进内存表。
func (c *NotificationClass) InsertForTest(n *model.Notification) {
c.listMu.Lock()
c.list[n.ID] = n
c.listMu.Unlock()
}