Buriburizaemon/nezha_domains
1
0
Fork 0
mirror of https://github.com/Buriburizaem0n/nezha_domains.git synced 2026-09-19 09:40:12 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(mcp): harden dashboard dispatch lifecycle

Browse Source 
Co-authored-by: naiba/CloudCode <hi+cloudcode@nai.ba>
This commit is contained in:
naiba
2026-07-20 04:32:16 +00:00
co-authored by naiba/CloudCode
parent 2640b86d3a
commit 3f6de265a0
48 changed files with 1332 additions and 265 deletions
Download Patch File Download Diff File Expand all files Collapse all files
service/rpc/mcp_kill_switch_registration_race_test.go
+8 -7
View File
@@ -17,11 +17,11 @@ import (
// 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.
// 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.
@@ -58,13 +58,14 @@ func TestCallAgent_KillSwitchBeatsRegistrationAfterSweep(t *testing.T) {
// 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() {
hook := func() {
mu.Lock()
killed = true
mu.Unlock()
CancelAllMCPInflight()
}
t.Cleanup(func() { testKillSwitchAfterUpfrontCheck = nil })
testKillSwitchAfterUpfrontCheck.Store(&hook)
t.Cleanup(func() { testKillSwitchAfterUpfrontCheck.Store(nil) })
_, err := CallAgent(context.Background(), target, model.TaskTypeExec,
model.ExecRequest{Cmd: "x"}, 1*time.Second)
service/rpc/mcp_rpc.go
+50 -17
View File
@@ -49,8 +49,13 @@ 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()
// deterministically.
var testKillSwitchAfterUpfrontCheck atomic.Pointer[func()]
var (
testMCPResultBeforeCancellationCheck atomic.Pointer[func()]
testMCPResultAfterCancellationCheck atomic.Pointer[func()]
)
// SetMCPKillSwitchObserver installs the kill-switch probe the dashboard
// owns. Idempotent; the dashboard wires it at startup. Passing nil
@@ -126,8 +131,8 @@ func CallAgent(ctx context.Context, serverID uint64, taskType uint64, params any
cancelled: new(atomic.Bool),
}
if hook := testKillSwitchAfterUpfrontCheck; hook != nil {
hook()
if hook := testKillSwitchAfterUpfrontCheck.Load(); hook != nil {
(*hook)()
}
mcpInflight.Store(taskID, entry)
@@ -153,6 +158,7 @@ func CallAgent(ctx context.Context, serverID uint64, taskType uint64, params any
}
return nil, err
}
notifyMCPTaskDispatched(serverID, taskID, taskType)
waitCtx := ctx
var cancel context.CancelFunc
@@ -163,6 +169,9 @@ func CallAgent(ctx context.Context, serverID uint64, taskType uint64, params any
select {
case res := <-resultCh:
if hook := testMCPResultBeforeCancellationCheck.Load(); hook != nil {
(*hook)()
}
// 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.
@@ -170,9 +179,13 @@ func CallAgent(ctx context.Context, serverID uint64, taskType uint64, params any
// contract documented above ("CancelAllMCPInflight 期间被中断 →
// ErrMCPDisabled") and what TestUpdateConfig_DisablingMCPInvokesKillSwitch
// expects.
if entry.cancelled.Load() {
if !entry.claimResult() {
return nil, ErrMCPDisabled
}
notifyMCPTaskResultAccepted(entry.serverID, res.GetId(), res.GetType())
if hook := testMCPResultAfterCancellationCheck.Load(); hook != nil {
(*hook)()
}
if res == nil {
return nil, errors.New("agent returned nil result")
}
@@ -201,20 +214,39 @@ func CallAgent(ctx context.Context, serverID uint64, taskType uint64, params any
// 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.
// cancelled flips to true when CancelAllMCPInflight wins the entry lock before
// the result is claimed. 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.
type mcpInflightEntry struct {
serverID uint64
result chan *pb.TaskResult
cancel chan struct{}
cancelled *atomic.Bool
mu sync.Mutex
claimed bool
closeOnce sync.Once
}
func (e *mcpInflightEntry) claimResult() bool {
e.mu.Lock()
defer e.mu.Unlock()
if e.cancelled.Load() {
return false
}
e.claimed = true
return true
}
func (e *mcpInflightEntry) cancelCall() {
e.mu.Lock()
if !e.claimed {
e.cancelled.Store(true)
}
e.mu.Unlock()
e.closeCancel()
}
// 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
@@ -243,10 +275,7 @@ func CancelAllMCPInflight() int {
if !ok {
return true
}
if entry.cancelled != nil {
entry.cancelled.Store(true)
}
entry.closeCancel()
entry.cancelCall()
mcpInflight.Delete(key)
cancelled++
return true
@@ -300,7 +329,9 @@ func deliverMCPResult(res *pb.TaskResult) {
if !ok {
return
}
if entry.cancelled != nil && entry.cancelled.Load() {
entry.mu.Lock()
defer entry.mu.Unlock()
if entry.cancelled.Load() {
return
}
select {
@@ -335,7 +366,9 @@ func deliverMCPResultFromReporter(res *pb.TaskResult, reporterID uint64) {
res.GetId(), entry.serverID, reporterID)
return
}
if entry.cancelled != nil && entry.cancelled.Load() {
entry.mu.Lock()
defer entry.mu.Unlock()
if entry.cancelled.Load() {
return
}
select {
service/rpc/mcp_rpc_kill_switch_race_test.go
+64 -3
View File
@@ -3,6 +3,7 @@ package rpc
import (
"context"
"errors"
"sync/atomic"
"testing"
"time"
@@ -33,6 +34,17 @@ func TestCallAgent_KillSwitchBeatsConcurrentLateResult(t *testing.T) {
cleanup := installFakeServer(t, target, stream)
defer cleanup()
resultSelected := make(chan struct{})
resumeResult := make(chan struct{})
var resultHook atomic.Pointer[func()]
hook := func() {
close(resultSelected)
<-resumeResult
}
resultHook.Store(&hook)
testMCPResultBeforeCancellationCheck.Store(resultHook.Load())
t.Cleanup(func() { testMCPResultBeforeCancellationCheck.Store(nil) })
delivered := make(chan struct{})
go func() {
sent := <-stream.sent
@@ -42,18 +54,67 @@ func TestCallAgent_KillSwitchBeatsConcurrentLateResult(t *testing.T) {
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)
errCh := make(chan error, 1)
go func() {
_, err := CallAgent(context.Background(), target, model.TaskTypeExec,
model.ExecRequest{Cmd: "x"}, 2*time.Second)
errCh <- err
}()
<-resultSelected
CancelAllMCPInflight()
close(resumeResult)
<-delivered
err := <-errCh
if !errors.Is(err, ErrMCPDisabled) {
t.Fatalf("kill switch must win the race with a late agent reply; want ErrMCPDisabled, got %v", err)
}
}
func TestCallAgent_ResultBeforeKillSwitchReturnsSuccess(t *testing.T) {
const target uint64 = 7303
stream := newFakeStream()
cleanup := installFakeServer(t, target, stream)
defer cleanup()
resultClaimed := make(chan struct{})
resumeResult := make(chan struct{})
var resultHook atomic.Pointer[func()]
hook := func() {
close(resultClaimed)
<-resumeResult
}
resultHook.Store(&hook)
testMCPResultAfterCancellationCheck.Store(resultHook.Load())
t.Cleanup(func() { testMCPResultAfterCancellationCheck.Store(nil) })
go func() {
sent := <-stream.sent
deliverMCPResultFromReporter(&pb.TaskResult{
Id: sent.GetId(),
Type: model.TaskTypeExec,
Successful: true,
Data: `{"exit_code":0}`,
}, target)
}()
errCh := make(chan error, 1)
go func() {
_, err := CallAgent(context.Background(), target, model.TaskTypeExec,
model.ExecRequest{Cmd: "x"}, 2*time.Second)
errCh <- err
}()
<-resultClaimed
CancelAllMCPInflight()
close(resumeResult)
if err := <-errCh; err != nil {
t.Fatalf("result claimed before kill switch must succeed, 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
service/rpc/mcp_rpc_test.go
+15 -8
View File
@@ -18,20 +18,27 @@ import (
type fakeTaskStream struct {
sent chan *pb.Task
delay time.Duration
err error
}
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 (s *fakeTaskStream) Send(t *pb.Task) error {
if s.err != nil {
return s.err
}
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()
service/singleton/nat_reserved_host_test.go
+6 -6
View File
@@ -120,12 +120,12 @@ func TestIsReservedDashboardHostCollapsesEquivalentForms(t *testing.T) {
})
reserved := []string{
"panel.example.com.", // trailing dot, no port
"panel.example.com.:8008", // trailing dot with port
"PANEL.EXAMPLE.COM.", // trailing dot, mixed case
"[0:0:0:0:0:0:0:1]:8008", // IPv6 expanded form of ::1
"::1", // IPv6 compressed, bare
"[::1]", // IPv6 compressed, bracketed
"panel.example.com.", // trailing dot, no port
"panel.example.com.:8008", // trailing dot with port
"PANEL.EXAMPLE.COM.", // trailing dot, mixed case
"[0:0:0:0:0:0:0:1]:8008", // IPv6 expanded form of ::1
"::1", // IPv6 compressed, bare
"[::1]", // IPv6 compressed, bracketed
}
for _, d := range reserved {
if !IsReservedDashboardHost(d) {