mirror of
https://github.com/Buriburizaem0n/nezha_domains.git
synced 2026-09-19 09:40:12 +00:00
fix(mcp): harden dashboard dispatch lifecycle
Co-authored-by: naiba/CloudCode <hi+cloudcode@nai.ba>
This commit is contained in:
@@ -33,6 +33,8 @@ func setupAlertRuleFanoutFixture(t *testing.T) {
|
||||
|
||||
db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
|
||||
require.NoError(t, err)
|
||||
sqlDB, err := db.DB()
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, db.AutoMigrate(&model.Server{}, &model.AlertRule{}, &model.Cron{}, &model.User{}))
|
||||
|
||||
singleton.DB = db
|
||||
@@ -50,6 +52,8 @@ func setupAlertRuleFanoutFixture(t *testing.T) {
|
||||
singleton.CronShared = singleton.NewCronClass()
|
||||
|
||||
t.Cleanup(func() {
|
||||
singleton.CronShared.Close()
|
||||
_ = sqlDB.Close()
|
||||
singleton.DB = originalDB
|
||||
singleton.Cache = originalCache
|
||||
singleton.Loc = originalLoc
|
||||
|
||||
@@ -17,9 +17,10 @@ import (
|
||||
)
|
||||
|
||||
const (
|
||||
apiTokenSecretLength = 32 // 明文 token 随机部分长度(hex 编码前)
|
||||
apiTokenCtxKey = "nz_api_token" // #nosec G101 -- gin context key name, not a credential
|
||||
apiTokenLastUsedCtxKey = "nz_api_token_used_marker" // #nosec G101 -- gin context key name, not a credential
|
||||
apiTokenSecretLength = 32 // 明文 token 随机部分长度(hex 编码前)
|
||||
apiTokenCtxKey = "nz_api_token" // #nosec G101 -- gin context key name, not a credential
|
||||
apiTokenLastUsedCtxKey = "nz_api_token_used_marker" // #nosec G101 -- gin context key name, not a credential
|
||||
apiTokenReadOnlyCtxKey = "nz_api_token_read_only" // #nosec G101 -- gin context key name, not a credential
|
||||
apiTokenAuthSchemePrefix = "Bearer "
|
||||
)
|
||||
|
||||
@@ -204,6 +205,7 @@ func deleteAPIToken(c *gin.Context) (any, error) {
|
||||
// 命中但 token 无效:直接 401 并 abort,不再走到 JWT。
|
||||
func apiTokenAuthMiddleware() gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
readOnly := apiTokenAuthReadOnly(c)
|
||||
raw := strings.TrimSpace(c.GetHeader("Authorization"))
|
||||
if raw == "" {
|
||||
return
|
||||
@@ -223,7 +225,9 @@ func apiTokenAuthMiddleware() gin.HandlerFunc {
|
||||
err := singleton.DB.Where("token_hash = ?", model.HashAPIToken(plaintext)).First(&tok).Error
|
||||
if err != nil {
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
model.BlockIP(singleton.DB, realIP, model.WAFBlockReasonTypeBruteForceToken, model.BlockIDToken)
|
||||
if !readOnly {
|
||||
model.BlockIP(singleton.DB, realIP, model.WAFBlockReasonTypeBruteForceToken, model.BlockIDToken)
|
||||
}
|
||||
abortAPITokenUnauthorized(c, "invalid api token")
|
||||
return
|
||||
}
|
||||
@@ -232,19 +236,25 @@ func apiTokenAuthMiddleware() gin.HandlerFunc {
|
||||
}
|
||||
now := time.Now()
|
||||
if tok.IsExpired(now) {
|
||||
model.BlockIP(singleton.DB, realIP, model.WAFBlockReasonTypeBruteForceToken, model.BlockIDToken)
|
||||
if !readOnly {
|
||||
model.BlockIP(singleton.DB, realIP, model.WAFBlockReasonTypeBruteForceToken, model.BlockIDToken)
|
||||
}
|
||||
abortAPITokenUnauthorized(c, "api token expired")
|
||||
return
|
||||
}
|
||||
|
||||
var user model.User
|
||||
if err := singleton.DB.First(&user, tok.UserID).Error; err != nil {
|
||||
model.BlockIP(singleton.DB, realIP, model.WAFBlockReasonTypeBruteForceToken, model.BlockIDToken)
|
||||
if !readOnly {
|
||||
model.BlockIP(singleton.DB, realIP, model.WAFBlockReasonTypeBruteForceToken, model.BlockIDToken)
|
||||
}
|
||||
abortAPITokenUnauthorized(c, "owner of api token not found")
|
||||
return
|
||||
}
|
||||
|
||||
model.UnblockIP(singleton.DB, realIP, model.BlockIDToken)
|
||||
if !readOnly {
|
||||
model.UnblockIP(singleton.DB, realIP, model.BlockIDToken)
|
||||
}
|
||||
|
||||
c.Set(model.CtxKeyAuthorizedUser, &user)
|
||||
c.Set(apiTokenCtxKey, &tok)
|
||||
@@ -253,7 +263,7 @@ func apiTokenAuthMiddleware() gin.HandlerFunc {
|
||||
// last_used 同步更新:开销极低(一行 UPDATE),异步路径在
|
||||
// 多连接 sqlite 测试场景下会和测试 teardown 形成竞态,并把
|
||||
// `last_used_*` 写丢到不可见的 :memory: 实例。生产路径上等价。
|
||||
if v, ok := c.Get(apiTokenLastUsedCtxKey); !ok || v != true {
|
||||
if !readOnly && apiTokenLastUsedOnce(c) {
|
||||
c.Set(apiTokenLastUsedCtxKey, true)
|
||||
_ = singleton.DB.Model(&model.APIToken{}).
|
||||
Where("id = ?", tok.ID).
|
||||
|
||||
@@ -0,0 +1,19 @@
|
||||
package controller
|
||||
|
||||
import "github.com/gin-gonic/gin"
|
||||
|
||||
func suppressAPITokenAuthWrites(c *gin.Context) { c.Set(apiTokenReadOnlyCtxKey, true) }
|
||||
|
||||
func apiTokenAuthReadOnly(c *gin.Context) bool {
|
||||
value, ok := c.Get(apiTokenReadOnlyCtxKey)
|
||||
return ok && value == true
|
||||
}
|
||||
|
||||
func apiTokenLastUsedOnce(c *gin.Context) bool {
|
||||
value, seen := c.Get(apiTokenLastUsedCtxKey)
|
||||
if seen && value == true {
|
||||
return false
|
||||
}
|
||||
c.Set(apiTokenLastUsedCtxKey, true)
|
||||
return true
|
||||
}
|
||||
@@ -15,9 +15,9 @@ import (
|
||||
|
||||
// restScopeMiddleware 的"空 scope"实际行为:PAT 调用方一律 403。
|
||||
// 这条测试把注释与实现的契约对齐:
|
||||
// 1. 实际行为:PAT + scope="" → 403。
|
||||
// 2. 文档约束:源码注释必须明确说出"空 scope 对 PAT 仍被拒绝",
|
||||
// 不能再保留"空字符串 = 放行"这种与实现相反的旧措辞。
|
||||
// 1. 实际行为:PAT + scope="" → 403。
|
||||
// 2. 文档约束:源码注释必须明确说出"空 scope 对 PAT 仍被拒绝",
|
||||
// 不能再保留"空字符串 = 放行"这种与实现相反的旧措辞。
|
||||
func TestRestScopeMiddleware_EmptyScopeRejectsPAT(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
r := gin.New()
|
||||
|
||||
@@ -23,9 +23,12 @@ func setupAPITokenTest(t *testing.T) func() {
|
||||
originalDB := singleton.DB
|
||||
db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
|
||||
require.NoError(t, err)
|
||||
sqlDB, err := db.DB()
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, db.AutoMigrate(&model.User{}, &model.APIToken{}, &model.Server{}))
|
||||
singleton.DB = db
|
||||
return func() {
|
||||
_ = sqlDB.Close()
|
||||
singleton.DB = originalDB
|
||||
}
|
||||
}
|
||||
|
||||
@@ -27,9 +27,9 @@ func callBatchMoveWithPAT(t *testing.T, callerID uint64, role model.Role, tok *m
|
||||
r.ServeHTTP(w, req)
|
||||
|
||||
var resp struct {
|
||||
Success bool `json:"success"`
|
||||
Error string `json:"error"`
|
||||
Data []model.BatchMoveServerResult `json:"data"`
|
||||
Success bool `json:"success"`
|
||||
Error string `json:"error"`
|
||||
Data []model.BatchMoveServerResult `json:"data"`
|
||||
}
|
||||
require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp))
|
||||
return resp.Data, resp.Success, resp.Error
|
||||
|
||||
@@ -46,13 +46,14 @@ func setupCronDispatchPATFixture(t *testing.T) {
|
||||
|
||||
db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
|
||||
require.NoError(t, err)
|
||||
sqlDB, err := db.DB()
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, db.AutoMigrate(&model.Cron{}, &model.Server{}, &model.User{}, &model.NotificationGroup{}, &model.Notification{}))
|
||||
|
||||
singleton.DB = db
|
||||
singleton.Loc = time.UTC
|
||||
singleton.Cache = cache.New(time.Minute, time.Minute)
|
||||
singleton.Localizer = i18n.NewLocalizer("en_US", "nezha", "translations", i18n.Translations)
|
||||
singleton.CronShared = singleton.NewCronClass()
|
||||
// CronTrigger 的 fan-out 路径会在 server 没接入 task-stream 时调
|
||||
// NotificationShared.SendNotification 上报「离线」;这里给出一个空的
|
||||
// notification class,避免 nil deref。本测试不验证通知内容。
|
||||
@@ -66,12 +67,16 @@ func setupCronDispatchPATFixture(t *testing.T) {
|
||||
sc.InsertForTest(s)
|
||||
}
|
||||
singleton.ServerShared = sc
|
||||
singleton.CronShared = singleton.NewCronClass()
|
||||
|
||||
singleton.UserLock.Lock()
|
||||
singleton.UserInfoMap = map[uint64]model.UserInfo{100: {Role: model.RoleMember}}
|
||||
singleton.UserLock.Unlock()
|
||||
|
||||
t.Cleanup(func() {
|
||||
// Test-owned cron jobs must be joined before restoring process-global singleton dependencies.
|
||||
singleton.CronShared.Close()
|
||||
_ = sqlDB.Close()
|
||||
singleton.DB = originalDB
|
||||
singleton.Cache = originalCache
|
||||
singleton.Loc = originalLoc
|
||||
|
||||
@@ -30,14 +30,16 @@ func setupCronManualTriggerFixture(t *testing.T) {
|
||||
originalUserInfo := singleton.UserInfoMap
|
||||
db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
|
||||
require.NoError(t, err)
|
||||
sqlDB, err := db.DB()
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, db.AutoMigrate(&model.Cron{}, &model.Server{}, &model.User{}))
|
||||
|
||||
singleton.DB = db
|
||||
singleton.Loc = time.UTC
|
||||
singleton.Cache = cache.New(time.Minute, time.Minute)
|
||||
singleton.Localizer = i18n.NewLocalizer("en_US", "nezha", "translations", i18n.Translations)
|
||||
singleton.CronShared = singleton.NewCronClass()
|
||||
singleton.ServerShared = singleton.NewServerClass()
|
||||
singleton.CronShared = singleton.NewCronClass()
|
||||
singleton.UserLock.Lock()
|
||||
singleton.UserInfoMap = map[uint64]model.UserInfo{100: {Role: model.RoleMember}}
|
||||
singleton.UserLock.Unlock()
|
||||
@@ -53,6 +55,8 @@ func setupCronManualTriggerFixture(t *testing.T) {
|
||||
singleton.CronShared.Update(cr)
|
||||
|
||||
t.Cleanup(func() {
|
||||
singleton.CronShared.Close()
|
||||
_ = sqlDB.Close()
|
||||
singleton.DB = originalDB
|
||||
singleton.Cache = originalCache
|
||||
singleton.Loc = originalLoc
|
||||
|
||||
@@ -33,14 +33,16 @@ func setupCronPATWhitelistFixture(t *testing.T) (cronID7, cronID8 uint64) {
|
||||
|
||||
db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
|
||||
require.NoError(t, err)
|
||||
sqlDB, err := db.DB()
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, db.AutoMigrate(&model.Cron{}, &model.Server{}, &model.User{}))
|
||||
|
||||
singleton.DB = db
|
||||
singleton.Loc = time.UTC
|
||||
singleton.Cache = cache.New(time.Minute, time.Minute)
|
||||
singleton.Localizer = i18n.NewLocalizer("en_US", "nezha", "translations", i18n.Translations)
|
||||
singleton.CronShared = singleton.NewCronClass()
|
||||
singleton.ServerShared = singleton.NewServerClass()
|
||||
singleton.CronShared = singleton.NewCronClass()
|
||||
singleton.UserLock.Lock()
|
||||
singleton.UserInfoMap = map[uint64]model.UserInfo{100: {Role: model.RoleMember}}
|
||||
singleton.UserLock.Unlock()
|
||||
@@ -67,6 +69,8 @@ func setupCronPATWhitelistFixture(t *testing.T) (cronID7, cronID8 uint64) {
|
||||
singleton.CronShared.Update(cr8)
|
||||
|
||||
t.Cleanup(func() {
|
||||
singleton.CronShared.Close()
|
||||
_ = sqlDB.Close()
|
||||
singleton.DB = originalDB
|
||||
singleton.Cache = originalCache
|
||||
singleton.Loc = originalLoc
|
||||
|
||||
@@ -50,27 +50,28 @@ func setupCoverPATFixture(t *testing.T) {
|
||||
originalCron := singleton.CronShared
|
||||
originalServer := singleton.ServerShared
|
||||
originalUserInfo := singleton.UserInfoMap
|
||||
originalNotification := singleton.NotificationShared
|
||||
originalSentinel := singleton.ServiceSentinelShared
|
||||
|
||||
db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
|
||||
require.NoError(t, err)
|
||||
sqlDB, err := db.DB()
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, db.AutoMigrate(&model.Cron{}, &model.Server{}, &model.User{}, &model.Service{}, &model.NotificationGroup{}, &model.ServiceHistory{}))
|
||||
|
||||
singleton.DB = db
|
||||
singleton.Loc = time.UTC
|
||||
singleton.Cache = cache.New(time.Minute, time.Minute)
|
||||
singleton.Localizer = i18n.NewLocalizer("en_US", "nezha", "translations", i18n.Translations)
|
||||
singleton.NotificationShared = singleton.NewEmptyNotificationClassForTest()
|
||||
sc := singleton.NewEmptyServerClassForTest()
|
||||
singleton.ServerShared = sc
|
||||
singleton.CronShared = singleton.NewCronClass()
|
||||
|
||||
originalSentinel := singleton.ServiceSentinelShared
|
||||
sentinel, err := singleton.NewServiceSentinel(make(chan *model.Service, 4))
|
||||
require.NoError(t, err)
|
||||
singleton.ServiceSentinelShared = sentinel
|
||||
t.Cleanup(func() {
|
||||
sentinel.Close()
|
||||
singleton.ServiceSentinelShared = originalSentinel
|
||||
})
|
||||
|
||||
sc := singleton.NewEmptyServerClassForTest()
|
||||
for _, id := range []uint64{1, 2} {
|
||||
s := &model.Server{}
|
||||
s.ID = id
|
||||
@@ -84,12 +85,18 @@ func setupCoverPATFixture(t *testing.T) {
|
||||
singleton.UserLock.Unlock()
|
||||
|
||||
t.Cleanup(func() {
|
||||
// Background components must be joined before restoring process globals.
|
||||
sentinel.Close()
|
||||
singleton.CronShared.Close()
|
||||
_ = sqlDB.Close()
|
||||
singleton.DB = originalDB
|
||||
singleton.Cache = originalCache
|
||||
singleton.Loc = originalLoc
|
||||
singleton.Localizer = originalLocalizer
|
||||
singleton.CronShared = originalCron
|
||||
singleton.ServerShared = originalServer
|
||||
singleton.NotificationShared = originalNotification
|
||||
singleton.ServiceSentinelShared = originalSentinel
|
||||
singleton.UserLock.Lock()
|
||||
singleton.UserInfoMap = originalUserInfo
|
||||
singleton.UserLock.Unlock()
|
||||
|
||||
@@ -109,7 +109,7 @@ func TestForceUpdateServerOnlineForeignIDIndistinguishableFromUnknown(t *testing
|
||||
defer reset()
|
||||
|
||||
const bobID = uint64(200)
|
||||
foreignResp := decodeForceUpdate(t, runForceUpdate(t, bobID, []uint64{1})) // alice's online
|
||||
foreignResp := decodeForceUpdate(t, runForceUpdate(t, bobID, []uint64{1})) // alice's online
|
||||
unknownResp := decodeForceUpdate(t, runForceUpdate(t, bobID, []uint64{9999})) // does not exist
|
||||
|
||||
assert.Equal(t, foreignResp.Success, unknownResp.Success,
|
||||
|
||||
@@ -64,10 +64,10 @@ func issueJWTSession(c *gin.Context, user *model.User, jwtTimeoutHours int) (map
|
||||
|
||||
func initParams() *jwt.GinJWTMiddleware {
|
||||
return &jwt.GinJWTMiddleware{
|
||||
Realm: singleton.Conf.SiteName,
|
||||
Key: []byte(singleton.Conf.JWTSecretKey),
|
||||
CookieName: "nz-jwt",
|
||||
SendCookie: true,
|
||||
Realm: singleton.Conf.SiteName,
|
||||
Key: []byte(singleton.Conf.JWTSecretKey),
|
||||
CookieName: "nz-jwt",
|
||||
SendCookie: true,
|
||||
// Pin the signing algorithm so a future library default change (or an
|
||||
// `alg: none` confusion attempt) cannot weaken token validation.
|
||||
SigningAlgorithm: "HS256",
|
||||
|
||||
@@ -30,6 +30,8 @@ func setupJWTSessionTest(t *testing.T) (cleanup func()) {
|
||||
originalConf := singleton.Conf
|
||||
db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
|
||||
require.NoError(t, err)
|
||||
sqlDB, err := db.DB()
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, db.AutoMigrate(&model.User{}, &model.JWTSession{}, &model.WAF{}))
|
||||
singleton.DB = db
|
||||
singleton.Conf = &singleton.ConfigClass{Config: &model.Config{JWTTimeout: 1}}
|
||||
@@ -42,6 +44,7 @@ func setupJWTSessionTest(t *testing.T) (cleanup func()) {
|
||||
}).Error)
|
||||
|
||||
return func() {
|
||||
_ = sqlDB.Close()
|
||||
singleton.DB = originalDB
|
||||
singleton.Conf = originalConf
|
||||
}
|
||||
|
||||
@@ -135,6 +135,8 @@ func newValidationContext(t *testing.T, userID uint64, role model.Role) *gin.Con
|
||||
|
||||
db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
|
||||
assert.NoError(t, err)
|
||||
sqlDB, err := db.DB()
|
||||
assert.NoError(t, err)
|
||||
assert.NoError(t, db.AutoMigrate(
|
||||
&model.Cron{},
|
||||
&model.Server{},
|
||||
@@ -169,8 +171,8 @@ func newValidationContext(t *testing.T, userID uint64, role model.Role) *gin.Con
|
||||
singleton.DB = db
|
||||
singleton.Loc = time.Local
|
||||
singleton.Localizer = i18n.NewLocalizer("en_US", "nezha", "translations", i18n.Translations)
|
||||
singleton.CronShared = singleton.NewCronClass()
|
||||
singleton.ServerShared = singleton.NewServerClass()
|
||||
singleton.CronShared = singleton.NewCronClass()
|
||||
singleton.UserLock.Lock()
|
||||
singleton.UserInfoMap = map[uint64]model.UserInfo{
|
||||
1: {Role: model.RoleAdmin},
|
||||
@@ -178,6 +180,8 @@ func newValidationContext(t *testing.T, userID uint64, role model.Role) *gin.Con
|
||||
}
|
||||
singleton.UserLock.Unlock()
|
||||
t.Cleanup(func() {
|
||||
singleton.CronShared.Close()
|
||||
_ = sqlDB.Close()
|
||||
singleton.DB = originalDB
|
||||
singleton.Loc = originalLoc
|
||||
singleton.Localizer = originalLocalizer
|
||||
|
||||
@@ -307,6 +307,20 @@ func handleToolsCall(c *gin.Context, req *jsonRPCRequest, tok *model.APIToken) {
|
||||
}
|
||||
|
||||
finish := func(outcome, errCode, errMsg string, result any) {
|
||||
if outcome == model.MCPOutcomeOK {
|
||||
textPayload, err := marshalMCPToolResult(result)
|
||||
if err != nil {
|
||||
outcome = model.MCPOutcomeAgentError
|
||||
errCode = model.MCPOutcomeAgentError
|
||||
errMsg = "failed to encode tool result: " + err.Error()
|
||||
result = nil
|
||||
} else {
|
||||
writeJSONRPCResult(c, req.ID, mcpToolCallResult{
|
||||
Content: []mcpContent{{Type: "text", Text: textPayload}},
|
||||
StructuredContent: result,
|
||||
})
|
||||
}
|
||||
}
|
||||
audit.Outcome = outcome
|
||||
audit.ErrorCode = errCode
|
||||
audit.ErrorMsg = truncateString(errMsg, 512)
|
||||
@@ -315,24 +329,15 @@ func handleToolsCall(c *gin.Context, req *jsonRPCRequest, tok *model.APIToken) {
|
||||
mcpAuditWrite(audit, p.Arguments)
|
||||
|
||||
if outcome == model.MCPOutcomeOK {
|
||||
textPayload := "{}"
|
||||
if result != nil {
|
||||
if b, err := json.Marshal(result); err == nil {
|
||||
textPayload = string(b)
|
||||
}
|
||||
}
|
||||
writeJSONRPCResult(c, req.ID, mcpToolCallResult{
|
||||
Content: []mcpContent{{Type: "text", Text: textPayload}},
|
||||
StructuredContent: result,
|
||||
})
|
||||
return
|
||||
}
|
||||
// 错误结果不带 structuredContent:严格客户端会拿它去校验工具声明的
|
||||
// outputSchema(要求 exit_code/stdout/... 等成功字段),缺字段就整条
|
||||
// 响应报 -32602,把真正的 isError 文本掩盖掉。错误信息走 content[].text。
|
||||
// Semantic tool failures may retain a typed structured result (notably
|
||||
// server.exec) so clients can distinguish a non-zero command outcome from
|
||||
// transport, authorization, or deadline failures.
|
||||
writeJSONRPCResult(c, req.ID, mcpToolCallResult{
|
||||
Content: []mcpContent{{Type: "text", Text: errMsg}},
|
||||
IsError: true,
|
||||
Content: []mcpContent{{Type: "text", Text: errMsg}},
|
||||
StructuredContent: result,
|
||||
IsError: true,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -356,6 +361,11 @@ func handleToolsCall(c *gin.Context, req *jsonRPCRequest, tok *model.APIToken) {
|
||||
result, err := tool.Handler(c, p.Arguments)
|
||||
if err != nil {
|
||||
code, msg := classifyToolError(err)
|
||||
var structuredErr interface{ StructuredResult() any }
|
||||
if errors.As(err, &structuredErr) {
|
||||
finish(code, code, msg, structuredErr.StructuredResult())
|
||||
return
|
||||
}
|
||||
finish(code, code, msg, nil)
|
||||
return
|
||||
}
|
||||
|
||||
@@ -18,10 +18,14 @@ const MCPMinAgentVersion = "v2.1.0"
|
||||
// requireAgentSupportsMCP 在 tool handler 调 CallAgent 之前快速失败不支持的 agent。
|
||||
// 仅作为 UX 优化:真正的安全/正确性由 agent 端 task switch 的 default 分支保障。
|
||||
func requireAgentSupportsMCP(server *model.Server) error {
|
||||
if MCPMinAgentVersion == "" || server == nil || server.Host == nil {
|
||||
if MCPMinAgentVersion == "" || server == nil {
|
||||
return nil
|
||||
}
|
||||
if compareSemver(server.Host.Version, MCPMinAgentVersion) < 0 {
|
||||
runtime := server.RuntimeSnapshot()
|
||||
if runtime.Host == nil {
|
||||
return nil
|
||||
}
|
||||
if compareSemver(runtime.Host.Version, MCPMinAgentVersion) < 0 {
|
||||
return errMCPUnsupported
|
||||
}
|
||||
return nil
|
||||
|
||||
@@ -10,10 +10,7 @@ import (
|
||||
"github.com/nezhahq/nezha/service/singleton"
|
||||
)
|
||||
|
||||
// A tool error must not ship structuredContent: strict clients validate it
|
||||
// against the tool's outputSchema (which requires exec/fs result fields) and
|
||||
// reject the whole response with -32602, masking the real isError text.
|
||||
func TestServerExec_ErrorResult_OmitsStructuredContent(t *testing.T) {
|
||||
func TestServerExec_ErrorResult_PreservesStructuredContent(t *testing.T) {
|
||||
cleanup, uid := setupMCPTest(t)
|
||||
defer cleanup()
|
||||
|
||||
@@ -40,8 +37,12 @@ func TestServerExec_ErrorResult_OmitsStructuredContent(t *testing.T) {
|
||||
require.True(t, tcr.IsError)
|
||||
require.Contains(t, tcr.Content[0].Text, "agent disabled command execution",
|
||||
"error text must carry the real cause")
|
||||
require.Nil(t, tcr.StructuredContent,
|
||||
"error responses must omit structuredContent so strict clients don't validate it against outputSchema and mask the real error")
|
||||
var result model.ExecResult
|
||||
structuredJSON, err := json.Marshal(tcr.StructuredContent)
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, json.Unmarshal(structuredJSON, &result))
|
||||
require.Equal(t, -1, result.ExitCode)
|
||||
require.Equal(t, "agent disabled command execution", result.Error)
|
||||
}
|
||||
|
||||
func TestScopeDenied_OmitsStructuredContent(t *testing.T) {
|
||||
@@ -62,6 +63,5 @@ func TestScopeDenied_OmitsStructuredContent(t *testing.T) {
|
||||
_, tcr := decodeRPC(w)
|
||||
require.NotNil(t, tcr)
|
||||
require.True(t, tcr.IsError)
|
||||
require.Nil(t, tcr.StructuredContent,
|
||||
"scope-denied error must also omit structuredContent")
|
||||
require.Nil(t, tcr.StructuredContent)
|
||||
}
|
||||
|
||||
@@ -17,6 +17,7 @@ type MCPRateLimiter struct {
|
||||
secLimit int
|
||||
minLimit int
|
||||
lastPrune time.Time
|
||||
clock func() time.Time
|
||||
}
|
||||
|
||||
type tokenWindow struct {
|
||||
@@ -34,10 +35,15 @@ type tokenWindow struct {
|
||||
const mcpRateLimiterPruneInterval = time.Minute
|
||||
|
||||
func newMCPRateLimiter(secLimit, minLimit int) *MCPRateLimiter {
|
||||
return newMCPRateLimiterWithClock(secLimit, minLimit, time.Now)
|
||||
}
|
||||
|
||||
func newMCPRateLimiterWithClock(secLimit, minLimit int, clock func() time.Time) *MCPRateLimiter {
|
||||
return &MCPRateLimiter{
|
||||
perToken: make(map[uint64]*tokenWindow),
|
||||
secLimit: secLimit,
|
||||
minLimit: minLimit,
|
||||
clock: clock,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -58,9 +64,11 @@ func (r *MCPRateLimiter) Allow(tokenID uint64) bool {
|
||||
if tokenID == 0 {
|
||||
return true
|
||||
}
|
||||
now := time.Now()
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
// Sampling while holding the state lock linearizes the clock read with the
|
||||
// bucket update, so a delayed sample cannot commit after a newer rollover.
|
||||
now := r.clock()
|
||||
if now.Sub(r.lastPrune) >= mcpRateLimiterPruneInterval {
|
||||
r.pruneStaleLocked(now)
|
||||
r.lastPrune = now
|
||||
|
||||
@@ -0,0 +1,51 @@
|
||||
//go:build agentcompat
|
||||
|
||||
package controller
|
||||
|
||||
import (
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
type agentcompatMCPRateLimitProbeRequest struct {
|
||||
}
|
||||
|
||||
type agentcompatMCPRateLimitProbeResponse struct {
|
||||
SecondAllowedCount int `json:"second_allowed_count"`
|
||||
SecondRejectedAtCount int `json:"second_rejected_at_count"`
|
||||
MinuteAllowedCount int `json:"minute_allowed_count"`
|
||||
MinuteRejectedAtCount int `json:"minute_rejected_at_count"`
|
||||
}
|
||||
|
||||
func agentcompatMCPRateLimitProbeRoute(context *gin.Context) (agentcompatMCPRateLimitProbeResponse, error) {
|
||||
var request agentcompatMCPRateLimitProbeRequest
|
||||
if err := decodeAgentcompatJSON(context, &request); err != nil {
|
||||
return agentcompatMCPRateLimitProbeResponse{}, err
|
||||
}
|
||||
return runAgentcompatMCPRateLimitProbe(request)
|
||||
}
|
||||
|
||||
func runAgentcompatMCPRateLimitProbe(request agentcompatMCPRateLimitProbeRequest) (agentcompatMCPRateLimitProbeResponse, error) {
|
||||
response := agentcompatMCPRateLimitProbeResponse{}
|
||||
secondNow := time.Unix(1_700_000_000, 0)
|
||||
secondLimiter := newMCPRateLimiterWithClock(10, 120, func() time.Time { return secondNow })
|
||||
for requestNumber := 1; requestNumber <= 11; requestNumber++ {
|
||||
if secondLimiter.Allow(1) {
|
||||
response.SecondAllowedCount++
|
||||
} else if response.SecondRejectedAtCount == 0 {
|
||||
response.SecondRejectedAtCount = requestNumber
|
||||
}
|
||||
}
|
||||
minuteNow := time.Unix(1_700_000_000, 0)
|
||||
minuteLimiter := newMCPRateLimiterWithClock(10_000, 120, func() time.Time { return minuteNow })
|
||||
for requestNumber := 1; requestNumber <= 121; requestNumber++ {
|
||||
if minuteLimiter.Allow(1) {
|
||||
response.MinuteAllowedCount++
|
||||
} else if response.MinuteRejectedAtCount == 0 {
|
||||
response.MinuteRejectedAtCount = requestNumber
|
||||
}
|
||||
minuteNow = minuteNow.Add(100 * time.Millisecond)
|
||||
}
|
||||
return response, nil
|
||||
}
|
||||
@@ -0,0 +1,102 @@
|
||||
//go:build agentcompat
|
||||
|
||||
package controller
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"net/http/httptest"
|
||||
"sync"
|
||||
"testing"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
func TestAgentcompatMCPRateLimitProbe_returnsTypedBoundaryCountsWithoutSharedState(t *testing.T) {
|
||||
// Given
|
||||
request := agentcompatMCPRateLimitProbeRequest{}
|
||||
originalLimiter := mcpRateLimiterShared
|
||||
|
||||
// When
|
||||
response, err := runAgentcompatMCPRateLimitProbe(request)
|
||||
|
||||
// Then
|
||||
if err != nil {
|
||||
t.Fatalf("probe returned error: %v", err)
|
||||
}
|
||||
if response.SecondAllowedCount != 10 || response.SecondRejectedAtCount != 11 || response.MinuteAllowedCount != 120 || response.MinuteRejectedAtCount != 121 {
|
||||
t.Fatalf("probe result = %+v, want second=10/11 minute=120/121", response)
|
||||
}
|
||||
t.Logf("typed probe result: second=%d/%d minute=%d/%d", response.SecondAllowedCount, response.SecondRejectedAtCount, response.MinuteAllowedCount, response.MinuteRejectedAtCount)
|
||||
if mcpRateLimiterShared != originalLimiter {
|
||||
t.Fatal("probe mutated the shared production limiter")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAgentcompatMCPRateLimitProbe_isRepeatableAndConcurrentSafe(t *testing.T) {
|
||||
// Given
|
||||
request := agentcompatMCPRateLimitProbeRequest{}
|
||||
results := make(chan agentcompatMCPRateLimitProbeResponse, 8)
|
||||
errors := make(chan error, 8)
|
||||
|
||||
// When
|
||||
var waitGroup sync.WaitGroup
|
||||
for probeNumber := 0; probeNumber < 8; probeNumber++ {
|
||||
waitGroup.Add(1)
|
||||
go func() {
|
||||
defer waitGroup.Done()
|
||||
response, err := runAgentcompatMCPRateLimitProbe(request)
|
||||
if err != nil {
|
||||
errors <- err
|
||||
return
|
||||
}
|
||||
results <- response
|
||||
}()
|
||||
}
|
||||
waitGroup.Wait()
|
||||
close(results)
|
||||
close(errors)
|
||||
|
||||
// Then
|
||||
for err := range errors {
|
||||
t.Fatalf("concurrent probe returned error: %v", err)
|
||||
}
|
||||
for response := range results {
|
||||
if response.SecondAllowedCount != 10 || response.SecondRejectedAtCount != 11 || response.MinuteAllowedCount != 120 || response.MinuteRejectedAtCount != 121 {
|
||||
t.Fatalf("concurrent probe result = %+v, want second=10/11 minute=120/121", response)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestAgentcompatMCPRateLimitProbe_rejectsCallerControlledParameters(t *testing.T) {
|
||||
// Given
|
||||
gin.SetMode(gin.TestMode)
|
||||
context, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
context.Request = httptest.NewRequest("POST", "/", bytes.NewBufferString(`{"token_id":7}`))
|
||||
context.Request.Header.Set("Content-Type", "application/json")
|
||||
|
||||
// When
|
||||
var request agentcompatMCPRateLimitProbeRequest
|
||||
err := decodeAgentcompatJSON(context, &request)
|
||||
|
||||
// Then
|
||||
if err == nil {
|
||||
t.Fatal("caller-controlled rate probe parameters must be rejected")
|
||||
}
|
||||
context.Request = httptest.NewRequest("POST", "/", bytes.NewBufferString(`{}`))
|
||||
if err := decodeAgentcompatJSON(context, &request); err != nil {
|
||||
t.Fatalf("canonical empty probe request must be accepted: %v", err)
|
||||
}
|
||||
if _, err := json.Marshal(request); err != nil {
|
||||
t.Fatalf("canonical request must remain JSON encodable: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAgentcompatMCPRateLimitProbeRejectsTrailingJSON(t *testing.T) {
|
||||
context, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
context.Request = httptest.NewRequest("POST", "/", bytes.NewBufferString(`{}{}`))
|
||||
var request agentcompatMCPRateLimitProbeRequest
|
||||
if err := decodeAgentcompatJSON(context, &request); err == nil {
|
||||
t.Fatal("trailing JSON values must be rejected")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,197 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
type mcpRateLimitFakeClock struct {
|
||||
now time.Time
|
||||
}
|
||||
|
||||
func (clock *mcpRateLimitFakeClock) Now() time.Time {
|
||||
return clock.now
|
||||
}
|
||||
|
||||
func (clock *mcpRateLimitFakeClock) Advance(duration time.Duration) {
|
||||
clock.now = clock.now.Add(duration)
|
||||
}
|
||||
|
||||
func TestMCPRateLimiter_baselinePreservesAnonymousBypassAndBucketReset(t *testing.T) {
|
||||
// Given
|
||||
limiter := newMCPRateLimiter(3, 3)
|
||||
|
||||
// When
|
||||
anonymousAllowed := limiter.Allow(0)
|
||||
firstAllowed := limiter.Allow(7)
|
||||
secondAllowed := limiter.Allow(7)
|
||||
thirdAllowed := limiter.Allow(7)
|
||||
fourthRejected := limiter.Allow(7)
|
||||
// Then
|
||||
if !anonymousAllowed || !firstAllowed || !secondAllowed || !thirdAllowed || fourthRejected {
|
||||
t.Fatalf("baseline limiter semantics changed: anonymous=%t first=%t second=%t third=%t rejected=%t", anonymousAllowed, firstAllowed, secondAllowed, thirdAllowed, fourthRejected)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMCPRateLimiter_allowsTenAndRejectsElevenWithinOneSecond(t *testing.T) {
|
||||
// Given
|
||||
clock := &mcpRateLimitFakeClock{now: time.Unix(1_700_000_000, 0)}
|
||||
limiter := newMCPRateLimiterWithClock(10, 120, clock.Now)
|
||||
|
||||
// When
|
||||
allowed := 0
|
||||
for requestNumber := 1; requestNumber <= 11; requestNumber++ {
|
||||
if limiter.Allow(7) {
|
||||
allowed++
|
||||
}
|
||||
}
|
||||
|
||||
// Then
|
||||
if allowed != 10 {
|
||||
t.Fatalf("allowed count = %d, want 10", allowed)
|
||||
}
|
||||
if limiter.Allow(7) {
|
||||
t.Fatal("twelfth request must remain rejected in the same second")
|
||||
}
|
||||
}
|
||||
|
||||
func TestMCPRateLimiter_allowsOneAfterSecondBucketRollover(t *testing.T) {
|
||||
// Given
|
||||
clock := &mcpRateLimitFakeClock{now: time.Unix(1_700_000_000, 0)}
|
||||
limiter := newMCPRateLimiterWithClock(10, 120, clock.Now)
|
||||
for requestNumber := 1; requestNumber <= 10; requestNumber++ {
|
||||
if !limiter.Allow(7) {
|
||||
t.Fatalf("request %d must be allowed before rollover", requestNumber)
|
||||
}
|
||||
}
|
||||
|
||||
// When
|
||||
clock.Advance(time.Second)
|
||||
|
||||
// Then
|
||||
if !limiter.Allow(7) {
|
||||
t.Fatal("request after one-second bucket rollover must be allowed")
|
||||
}
|
||||
}
|
||||
|
||||
func TestMCPRateLimiter_allows120AndRejects121WithinOneMinute(t *testing.T) {
|
||||
// Given
|
||||
clock := &mcpRateLimitFakeClock{now: time.Unix(1_700_000_000, 0)}
|
||||
limiter := newMCPRateLimiterWithClock(10_000, 120, clock.Now)
|
||||
|
||||
// When
|
||||
allowed := 0
|
||||
for requestNumber := 1; requestNumber <= 121; requestNumber++ {
|
||||
if limiter.Allow(7) {
|
||||
allowed++
|
||||
}
|
||||
clock.Advance(100 * time.Millisecond)
|
||||
}
|
||||
|
||||
// Then
|
||||
if allowed != 120 {
|
||||
t.Fatalf("allowed count = %d, want 120", allowed)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMCPRateLimiter_independentTokensKeepIndependentBudgets(t *testing.T) {
|
||||
// Given
|
||||
clock := &mcpRateLimitFakeClock{now: time.Unix(1_700_000_000, 0)}
|
||||
limiter := newMCPRateLimiterWithClock(1, 120, clock.Now)
|
||||
|
||||
// When
|
||||
firstTokenAllowed := limiter.Allow(7)
|
||||
firstTokenRejected := limiter.Allow(7)
|
||||
secondTokenAllowed := limiter.Allow(8)
|
||||
|
||||
// Then
|
||||
if !firstTokenAllowed || firstTokenRejected || !secondTokenAllowed {
|
||||
t.Fatalf("token budgets are not independent: first allowed=%t rejected=%t second allowed=%t", firstTokenAllowed, firstTokenRejected, secondTokenAllowed)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMCPRateLimiter_clockControlledCallsDoNotMutateSharedLimiter(t *testing.T) {
|
||||
// Given
|
||||
originalLimiter := mcpRateLimiterShared
|
||||
clock := &mcpRateLimitFakeClock{now: time.Unix(1_700_000_000, 0)}
|
||||
limiter := newMCPRateLimiterWithClock(1, 1, clock.Now)
|
||||
|
||||
// When
|
||||
if !limiter.Allow(7) {
|
||||
t.Fatal("isolated limiter request must be allowed")
|
||||
}
|
||||
|
||||
// Then
|
||||
if mcpRateLimiterShared != originalLimiter {
|
||||
t.Fatal("isolated limiter changed shared production limiter ownership")
|
||||
}
|
||||
}
|
||||
|
||||
func TestMCPRateLimiter_concurrentCallsAtSecondBoundaryAllowExactlyLimit(t *testing.T) {
|
||||
// Given
|
||||
clock := &mcpRateLimitFakeClock{now: time.Unix(1_700_000_000, 0)}
|
||||
limiter := newMCPRateLimiterWithClock(10, 120, clock.Now)
|
||||
var allowed atomic.Int32
|
||||
var waitGroup sync.WaitGroup
|
||||
|
||||
// When
|
||||
for requestNumber := 0; requestNumber < 20; requestNumber++ {
|
||||
waitGroup.Add(1)
|
||||
go func() {
|
||||
defer waitGroup.Done()
|
||||
if limiter.Allow(7) {
|
||||
allowed.Add(1)
|
||||
}
|
||||
}()
|
||||
}
|
||||
waitGroup.Wait()
|
||||
|
||||
// Then
|
||||
if got := allowed.Load(); got != 10 {
|
||||
t.Fatalf("concurrent allowed count = %d, want 10", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMCPRateLimiter_minuteBucketRolloverRestoresBudget(t *testing.T) {
|
||||
// Given
|
||||
clock := &mcpRateLimitFakeClock{now: time.Unix(1_700_000_000, 0)}
|
||||
limiter := newMCPRateLimiterWithClock(10_000, 1, clock.Now)
|
||||
if !limiter.Allow(7) {
|
||||
t.Fatal("first request must be allowed")
|
||||
}
|
||||
if limiter.Allow(7) {
|
||||
t.Fatal("second request must be rejected before minute rollover")
|
||||
}
|
||||
|
||||
// When
|
||||
clock.Advance(time.Minute)
|
||||
|
||||
// Then
|
||||
if !limiter.Allow(7) {
|
||||
t.Fatal("request after minute rollover must be allowed")
|
||||
}
|
||||
}
|
||||
|
||||
func TestMCPRateLimiter_samplesClockWhileHoldingStateLock(t *testing.T) {
|
||||
// Given
|
||||
now := time.Unix(1_700_000_000, 0)
|
||||
var limiter *MCPRateLimiter
|
||||
clock := func() time.Time {
|
||||
if limiter.mu.TryLock() {
|
||||
limiter.mu.Unlock()
|
||||
t.Fatal("clock callback acquired limiter state lock; Allow sampled before locking")
|
||||
}
|
||||
return now
|
||||
}
|
||||
limiter = newMCPRateLimiterWithClock(1, 1, clock)
|
||||
|
||||
// When
|
||||
allowed := limiter.Allow(7)
|
||||
|
||||
// Then
|
||||
if !allowed {
|
||||
t.Fatal("initial request must be allowed")
|
||||
}
|
||||
}
|
||||
@@ -7,22 +7,39 @@ import (
|
||||
"testing"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/nezhahq/nezha/model"
|
||||
"github.com/nezhahq/nezha/service/singleton"
|
||||
)
|
||||
|
||||
func mcpEndpointTestCtx(t *testing.T, tok *model.APIToken, body any) (*gin.Context, *httptest.ResponseRecorder) {
|
||||
func mcpEndpointTestCtx(t *testing.T, tok *model.APIToken, body any) (*gin.Context, *httptest.ResponseRecorder, error) {
|
||||
t.Helper()
|
||||
gin.SetMode(gin.TestMode)
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
raw, _ := json.Marshal(body)
|
||||
c, engine := gin.CreateTestContext(w)
|
||||
require.NotNil(t, engine)
|
||||
raw, err := json.Marshal(body)
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
c.Request = httptest.NewRequest("POST", "/mcp", bytes.NewReader(raw))
|
||||
c.Request.Header.Set("Content-Type", "application/json")
|
||||
c.Set(apiTokenCtxKey, tok)
|
||||
c.Set(model.CtxKeyAPIToken, tok)
|
||||
return c, w
|
||||
return c, w, nil
|
||||
}
|
||||
|
||||
func TestMCPEndpointTestCtx_surfacesJSONMarshalError(t *testing.T) {
|
||||
// Given
|
||||
tok := &model.APIToken{ID: 4242, UserID: 1}
|
||||
unsupportedBody := map[string]any{"unsupported": func() {}}
|
||||
|
||||
// When
|
||||
_, _, err := mcpEndpointTestCtx(t, tok, unsupportedBody)
|
||||
|
||||
// Then
|
||||
require.Error(t, err)
|
||||
}
|
||||
|
||||
// An unknown tool name in tools/call must still consume the per-token rate
|
||||
@@ -53,7 +70,8 @@ func TestMCPUnknownToolCountsAgainstRateLimit(t *testing.T) {
|
||||
|
||||
var lastStatus int
|
||||
for i := 0; i < 5; i++ {
|
||||
c, w := mcpEndpointTestCtx(t, tok, body)
|
||||
c, w, err := mcpEndpointTestCtx(t, tok, body)
|
||||
require.NoError(t, err)
|
||||
mcpEndpoint(c)
|
||||
lastStatus = w.Code
|
||||
}
|
||||
|
||||
@@ -0,0 +1,45 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
type linearizationClock struct {
|
||||
mu sync.Mutex
|
||||
times []time.Time
|
||||
latest time.Time
|
||||
}
|
||||
|
||||
func newLinearizationClock(oldTime, newTime time.Time) *linearizationClock {
|
||||
return &linearizationClock{
|
||||
times: []time.Time{oldTime, oldTime, newTime},
|
||||
latest: newTime,
|
||||
}
|
||||
}
|
||||
|
||||
func (clock *linearizationClock) Now() time.Time {
|
||||
clock.mu.Lock()
|
||||
now := clock.latest
|
||||
if len(clock.times) > 0 {
|
||||
now = clock.times[0]
|
||||
clock.times = clock.times[1:]
|
||||
}
|
||||
clock.mu.Unlock()
|
||||
return now
|
||||
}
|
||||
|
||||
func TestLinearizationClockReturnsLatestTimeAfterSequenceIsConsumed(t *testing.T) {
|
||||
oldTime := time.Unix(1_700_000_000, 0)
|
||||
newTime := oldTime.Add(time.Second)
|
||||
clock := newLinearizationClock(oldTime, newTime)
|
||||
|
||||
for range 3 {
|
||||
clock.Now()
|
||||
}
|
||||
|
||||
if got := clock.Now(); !got.Equal(newTime) {
|
||||
t.Fatalf("clock returned %s after sequence was consumed, want %s", got, newTime)
|
||||
}
|
||||
}
|
||||
@@ -6,6 +6,7 @@ import (
|
||||
"testing"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/nezhahq/nezha/model"
|
||||
"github.com/nezhahq/nezha/service/singleton"
|
||||
@@ -15,7 +16,8 @@ func mcpEndpointRawCtx(t *testing.T, tok *model.APIToken, raw []byte) (*gin.Cont
|
||||
t.Helper()
|
||||
gin.SetMode(gin.TestMode)
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
c, engine := gin.CreateTestContext(w)
|
||||
require.NotNil(t, engine)
|
||||
c.Request = httptest.NewRequest("POST", "/mcp", bytes.NewReader(raw))
|
||||
c.Request.Header.Set("Content-Type", "application/json")
|
||||
c.Set(apiTokenCtxKey, tok)
|
||||
@@ -47,8 +49,9 @@ func TestMCPMalformedToolsCallParamsCountsAgainstRateLimit(t *testing.T) {
|
||||
raw := []byte(`{"jsonrpc":"2.0","id":1,"method":"tools/call","params":"not-an-object"}`)
|
||||
|
||||
for i := 0; i < 5; i++ {
|
||||
c, _ := mcpEndpointRawCtx(t, tok, raw)
|
||||
c, w := mcpEndpointRawCtx(t, tok, raw)
|
||||
mcpEndpoint(c)
|
||||
require.NotEmpty(t, w.Body.Bytes())
|
||||
}
|
||||
|
||||
if mcpRateLimiterShared.Allow(tok.ID) {
|
||||
@@ -64,8 +67,9 @@ func TestMCPMalformedEnvelopeCountsAgainstRateLimit(t *testing.T) {
|
||||
raw := []byte(`{not valid json`)
|
||||
|
||||
for i := 0; i < 5; i++ {
|
||||
c, _ := mcpEndpointRawCtx(t, tok, raw)
|
||||
c, w := mcpEndpointRawCtx(t, tok, raw)
|
||||
mcpEndpoint(c)
|
||||
require.NotEmpty(t, w.Body.Bytes())
|
||||
}
|
||||
|
||||
if mcpRateLimiterShared.Allow(tok.ID) {
|
||||
|
||||
@@ -0,0 +1,14 @@
|
||||
package controller
|
||||
|
||||
import "encoding/json"
|
||||
|
||||
func marshalMCPToolResult(result any) (string, error) {
|
||||
if result == nil {
|
||||
return "{}", nil
|
||||
}
|
||||
b, err := json.Marshal(result)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
return string(b), nil
|
||||
}
|
||||
@@ -0,0 +1,47 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/nezhahq/nezha/model"
|
||||
)
|
||||
|
||||
var registerUnmarshalableMCPTool sync.Once
|
||||
|
||||
func TestMCPToolCall_unmarshalableSuccessResultReturnsExplicitToolError(t *testing.T) {
|
||||
// Given
|
||||
cleanup, uid := setupMCPTest(t)
|
||||
defer cleanup()
|
||||
registerUnmarshalableMCPTool.Do(func() {
|
||||
registerMCPTool(&mcpTool{
|
||||
Name: "test.unmarshalable-success",
|
||||
RequiredScope: "",
|
||||
Handler: func(*gin.Context, json.RawMessage) (any, error) {
|
||||
return map[string]any{"unsupported": func() {}}, nil
|
||||
},
|
||||
})
|
||||
})
|
||||
tok, plainToken := mkToken(t, uid, []string{model.ScopeServerRead}, nil)
|
||||
require.NotEmpty(t, plainToken)
|
||||
c, w := mcpCallCtx(t, tok, uid, jsonRPCRequest{
|
||||
JSONRPC: "2.0", ID: json.RawMessage("1"), Method: "tools/call",
|
||||
Params: jsonObj(t, toolCallParams{Name: "test.unmarshalable-success", Arguments: json.RawMessage("{}")}),
|
||||
})
|
||||
|
||||
// When
|
||||
mcpEndpoint(c)
|
||||
|
||||
// Then
|
||||
_, result := decodeRPC(w)
|
||||
require.NotNil(t, result)
|
||||
require.True(t, result.IsError)
|
||||
require.Nil(t, result.StructuredContent)
|
||||
require.Len(t, result.Content, 1)
|
||||
require.True(t, strings.Contains(result.Content[0].Text, "encode"), result.Content[0].Text)
|
||||
}
|
||||
@@ -35,6 +35,8 @@ func setupMCPTest(t *testing.T) (func(), uint64) {
|
||||
patConnectionRegistryShared = newPATConnectionRegistry()
|
||||
db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
|
||||
require.NoError(t, err)
|
||||
sqlDB, err := db.DB()
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, db.AutoMigrate(&model.User{}, &model.APIToken{}, &model.MCPAuditLog{}, &model.Server{}, &model.WAF{}))
|
||||
singleton.DB = db
|
||||
singleton.Conf = &singleton.ConfigClass{Config: &model.Config{JWTTimeout: 1}}
|
||||
@@ -52,6 +54,7 @@ func setupMCPTest(t *testing.T) (func(), uint64) {
|
||||
singleton.ServerShared = sc
|
||||
|
||||
cleanup := func() {
|
||||
_ = sqlDB.Close()
|
||||
singleton.DB = originalDB
|
||||
singleton.ServerShared = originalServer
|
||||
singleton.Conf = originalConf
|
||||
|
||||
@@ -2,7 +2,7 @@ package controller
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
@@ -115,11 +115,22 @@ func handleServerExec(c *gin.Context, raw json.RawMessage) (any, error) {
|
||||
// handlers do, so MCP isError=true and audit outcome=agent_error. Non-zero
|
||||
// ExitCode alone is a normal command outcome, not a tool error.
|
||||
if res.Error != "" {
|
||||
return nil, errors.New(res.Error)
|
||||
return nil, &execToolError{mcpError: mcpError{Code: model.MCPOutcomeAgentError, Msg: res.Error}, result: res}
|
||||
}
|
||||
return res, nil
|
||||
}
|
||||
|
||||
type execToolError struct {
|
||||
mcpError
|
||||
result model.ExecResult
|
||||
}
|
||||
|
||||
func (err *execToolError) StructuredResult() any { return err.result }
|
||||
|
||||
func (err *execToolError) Error() string {
|
||||
return fmt.Sprintf("%s: %s", err.Code, err.Msg)
|
||||
}
|
||||
|
||||
// callAgentTimeout 给 dashboard 侧 CallAgent 计算等待上限。
|
||||
// 在用户请求的 timeout 基础上加 5s buffer,让 agent 端的 hard timeout 先触发,
|
||||
// 这样 dashboard 收到的总是结构化结果(包含 timed_out=true),
|
||||
|
||||
@@ -78,6 +78,12 @@ func TestServerExec_AgentReportedErrorBecomesToolError(t *testing.T) {
|
||||
"agent ExecResult.Error must propagate as MCP tool error; got %+v", tcr)
|
||||
require.Contains(t, tcr.Content[0].Text, "agent disabled command execution",
|
||||
"tool error text must surface the agent-reported error message")
|
||||
var result model.ExecResult
|
||||
structuredJSON, err := json.Marshal(tcr.StructuredContent)
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, json.Unmarshal(structuredJSON, &result))
|
||||
require.Equal(t, -1, result.ExitCode)
|
||||
require.Equal(t, "agent disabled command execution", result.Error)
|
||||
|
||||
require.Eventually(t, func() bool {
|
||||
var got model.MCPAuditLog
|
||||
|
||||
@@ -136,7 +136,8 @@ func handleServerList(c *gin.Context, raw json.RawMessage) (any, error) {
|
||||
if !tok.CanAccessServer(s.ID) {
|
||||
continue
|
||||
}
|
||||
online := !s.LastActive.IsZero() && now.Sub(s.LastActive) < onlineWindow
|
||||
runtime := s.RuntimeSnapshot()
|
||||
online := !runtime.LastActive.IsZero() && now.Sub(runtime.LastActive) < onlineWindow
|
||||
if args.OnlineOnly && !online {
|
||||
continue
|
||||
}
|
||||
@@ -145,11 +146,11 @@ func handleServerList(c *gin.Context, raw json.RawMessage) (any, error) {
|
||||
Name: s.Name,
|
||||
UUID: s.UUID,
|
||||
Online: online,
|
||||
LastActive: s.LastActive,
|
||||
LastActive: runtime.LastActive,
|
||||
}
|
||||
if s.Host != nil {
|
||||
item.Platform = s.Host.Platform
|
||||
item.Arch = s.Host.Arch
|
||||
if runtime.Host != nil {
|
||||
item.Platform = runtime.Host.Platform
|
||||
item.Arch = runtime.Host.Arch
|
||||
}
|
||||
if s.GeoIP != nil {
|
||||
item.IPv4 = s.GeoIP.IP.IPv4Addr
|
||||
@@ -187,15 +188,16 @@ func handleServerGet(c *gin.Context, raw json.RawMessage) (any, error) {
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
runtime := s.RuntimeSnapshot()
|
||||
return map[string]any{
|
||||
"id": s.ID,
|
||||
"name": s.Name,
|
||||
"uuid": s.UUID,
|
||||
"note": s.Note,
|
||||
"public_note": s.PublicNote,
|
||||
"host": s.Host,
|
||||
"state": s.State,
|
||||
"host": runtime.Host,
|
||||
"state": runtime.State,
|
||||
"geoip": s.GeoIP,
|
||||
"last_active": s.LastActive,
|
||||
"last_active": runtime.LastActive,
|
||||
}, nil
|
||||
}
|
||||
|
||||
@@ -19,12 +19,8 @@ import (
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/hashicorp/go-uuid"
|
||||
|
||||
"github.com/nezhahq/nezha/model"
|
||||
"github.com/nezhahq/nezha/pkg/utils"
|
||||
pb "github.com/nezhahq/nezha/proto"
|
||||
"github.com/nezhahq/nezha/service/rpc"
|
||||
"github.com/nezhahq/nezha/service/singleton"
|
||||
)
|
||||
|
||||
@@ -931,96 +927,6 @@ func xferSizeFromU64(raw uint64) (int64, error) {
|
||||
return int64(raw), nil
|
||||
}
|
||||
|
||||
// openFsTransferStream 走 IOStream 通道与目标 agent 建立一条专用大文件流。
|
||||
// 返回的 io.ReadWriteCloser 既能 Read(接收 agent→dashboard 字节)又能 Write
|
||||
// (发送 dashboard→agent 字节);调用方通过 readXferFixedHeader 解析控制帧。
|
||||
//
|
||||
// 内部步骤:
|
||||
// 1. 分配 streamId,CreateStream(streamId, 0, serverID) 在 NezhaHandler 注册
|
||||
// 一个 ioStreamContext,targetServerID 用于 agent 侧 stream 归属校验。
|
||||
// 2. 通过 server 当前的 RequestTask 流发 TaskTypeFsTransfer,把 streamId +
|
||||
// req JSON 下发给 agent。
|
||||
// 3. 等 agent 通过 IOStream() RPC 完成 magic 引导帧并 AgentConnected。
|
||||
// 4. 返回 agent 端流和 cleanup(CloseStream)。
|
||||
//
|
||||
// 任何步骤失败都会 CloseStream 释放资源;调用方只需在 defer cleanup() 即可。
|
||||
func openFsTransferStream(ctx context.Context, serverID uint64, req *model.FsTransferRequest) (io.ReadWriteCloser, func(), error) {
|
||||
if singleton.Conf == nil || !singleton.Conf.MCPEnabled() {
|
||||
return nil, func() {}, errors.New("MCP is disabled by the dashboard administrator")
|
||||
}
|
||||
server, _ := singleton.ServerShared.Get(serverID)
|
||||
if server == nil {
|
||||
return nil, func() {}, errors.New("server offline")
|
||||
}
|
||||
if server.GetTaskStream() == nil {
|
||||
return nil, func() {}, errors.New("server offline")
|
||||
}
|
||||
|
||||
streamId, err := uuid.GenerateUUID()
|
||||
if err != nil {
|
||||
return nil, func() {}, err
|
||||
}
|
||||
req.StreamID = streamId
|
||||
|
||||
if err := rpc.NezhaHandlerSingleton.CreateStreamWithPurpose(streamId, 0, serverID, rpc.PurposeMCPTransfer); err != nil {
|
||||
return nil, func() {}, err
|
||||
}
|
||||
cleanup := func() { _ = rpc.NezhaHandlerSingleton.CloseStream(streamId) }
|
||||
|
||||
body, err := json.Marshal(req)
|
||||
if err != nil {
|
||||
cleanup()
|
||||
return nil, func() {}, err
|
||||
}
|
||||
// 关闭 entry-check 与 SendTask 之间的 TOCTOU:stream 已注册后再复查一次
|
||||
// kill switch / ctx 取消,确保 disable sweep 要么扫到这条已注册 stream、
|
||||
// 要么这里读到 disabled,绝不会在禁用/吊销后仍把 transfer 任务发给 agent。
|
||||
if singleton.Conf == nil || !singleton.Conf.MCPEnabled() {
|
||||
cleanup()
|
||||
return nil, func() {}, errors.New("MCP is disabled by the dashboard administrator")
|
||||
}
|
||||
if err := ctx.Err(); err != nil {
|
||||
cleanup()
|
||||
return nil, func() {}, err
|
||||
}
|
||||
if err := server.SendTask(&pb.Task{
|
||||
Type: model.TaskTypeFsTransfer,
|
||||
Data: string(body),
|
||||
}); err != nil {
|
||||
cleanup()
|
||||
if errors.Is(err, model.ErrTaskStreamOffline) {
|
||||
return nil, func() {}, errors.New("server offline")
|
||||
}
|
||||
return nil, func() {}, err
|
||||
}
|
||||
|
||||
agentStream, ok := rpc.NezhaHandlerSingleton.WaitForAgent(ctx, streamId, 30*time.Second)
|
||||
if !ok {
|
||||
cleanup()
|
||||
return nil, func() {}, errors.New("agent did not attach within 30s")
|
||||
}
|
||||
|
||||
// After attach, the relay blocks in IOStreamWrapper.Read, which only
|
||||
// honours the gRPC stream context — not this per-transfer ctx. Without
|
||||
// the watcher below, a PAT revocation (deleteAPIToken cancels ctx) or a
|
||||
// client disconnect would leave a stalled/compromised agent pinning this
|
||||
// goroutine + IOStream until process restart. Closing the stream on
|
||||
// ctx.Done() unblocks the agent-side handler (iw.Wait) so Read returns.
|
||||
watcherDone := make(chan struct{})
|
||||
go func() {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
_ = rpc.NezhaHandlerSingleton.CloseStream(streamId)
|
||||
case <-watcherDone:
|
||||
}
|
||||
}()
|
||||
wrappedCleanup := func() {
|
||||
close(watcherDone)
|
||||
cleanup()
|
||||
}
|
||||
return agentStream, wrappedCleanup, nil
|
||||
}
|
||||
|
||||
// revalidateTransferEntry 在消费一次性 URL 时重新检查 mint 阶段的全部前置。
|
||||
// 这是 mint→consume 之间发生权限变化(PAT 吊销、scope/whitelist 收紧、
|
||||
// server 转手)时的兜底闸门:HMAC 签发与 sync.Map 一次性消费机制本身只能
|
||||
|
||||
@@ -3,6 +3,9 @@ package controller
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
@@ -23,16 +26,36 @@ func newFakeAgentIO() *grpcx.IOStreamWrapper {
|
||||
// wrapper is closed, exactly the post-attach state where nothing watches the
|
||||
// per-transfer context.
|
||||
type fakeAgentStream struct {
|
||||
closed chan struct{}
|
||||
closed chan struct{}
|
||||
closeOnce sync.Once
|
||||
closeSeen chan struct{}
|
||||
recvDone chan struct{}
|
||||
closeCall atomic.Int32
|
||||
}
|
||||
|
||||
func (f *fakeAgentStream) Recv() (*pb.IOStreamData, error) {
|
||||
<-f.closed
|
||||
close(f.recvDone)
|
||||
return nil, context.Canceled
|
||||
}
|
||||
func (f *fakeAgentStream) Send(*pb.IOStreamData) error { return nil }
|
||||
func (f *fakeAgentStream) Context() context.Context { return context.Background() }
|
||||
|
||||
func (f *fakeAgentStream) closeEndpoint() {
|
||||
f.closeOnce.Do(func() { close(f.closeSeen) })
|
||||
}
|
||||
|
||||
func (f *fakeAgentStream) Close() error {
|
||||
f.closeCall.Add(1)
|
||||
f.closeEndpoint()
|
||||
select {
|
||||
case <-f.closed:
|
||||
default:
|
||||
close(f.closed)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// transferRevokableContext only cancels a context; the post-attach relay
|
||||
// (readXferFixedHeader / relayDownloadFrames / io.CopyN) and IOStreamWrapper.Read
|
||||
// do not watch it. A revoked PAT (or a disconnected HTTP client) must still
|
||||
@@ -59,18 +82,17 @@ func TestOpenFsTransferStream_CancelClosesAttachedStream(t *testing.T) {
|
||||
t.Cleanup(func() { singleton.ServerShared = originalShared })
|
||||
|
||||
streamIDCh := make(chan string, 1)
|
||||
agentStreamCh := make(chan *fakeAgentStream, 1)
|
||||
attachReady := make(chan struct{})
|
||||
go func() {
|
||||
task := <-stream.sent
|
||||
var req model.FsTransferRequest
|
||||
_ = json.Unmarshal([]byte(task.GetData()), &req)
|
||||
require.NoError(t, json.Unmarshal([]byte(task.GetData()), &req))
|
||||
streamIDCh <- req.StreamID
|
||||
deadline := time.Now().Add(2 * time.Second)
|
||||
for time.Now().Before(deadline) {
|
||||
if err := rpc.NezhaHandlerSingleton.AgentConnected(req.StreamID, newFakeAgentIO()); err == nil {
|
||||
return
|
||||
}
|
||||
time.Sleep(5 * time.Millisecond)
|
||||
}
|
||||
fakeAgent := &fakeAgentStream{closed: make(chan struct{}), closeSeen: make(chan struct{}), recvDone: make(chan struct{})}
|
||||
agentStreamCh <- fakeAgent
|
||||
require.NoError(t, rpc.NezhaHandlerSingleton.AgentConnected(req.StreamID, grpcx.NewIOStreamWrapper(fakeAgent)))
|
||||
close(attachReady)
|
||||
}()
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
@@ -83,14 +105,65 @@ func TestOpenFsTransferStream_CancelClosesAttachedStream(t *testing.T) {
|
||||
defer cleanup()
|
||||
|
||||
streamID := <-streamIDCh
|
||||
<-attachReady
|
||||
_, getErr := rpc.NezhaHandlerSingleton.GetStream(streamID)
|
||||
require.NoError(t, getErr, "stream must be live before cancel")
|
||||
agentStream := <-agentStreamCh
|
||||
readDone := make(chan error, 1)
|
||||
go func() {
|
||||
_, readErr := streamIO.Read(make([]byte, 16))
|
||||
readDone <- readErr
|
||||
}()
|
||||
|
||||
cancel()
|
||||
|
||||
require.Eventually(t, func() bool {
|
||||
_, e := rpc.NezhaHandlerSingleton.GetStream(streamID)
|
||||
return e != nil
|
||||
}, 2*time.Second, 10*time.Millisecond,
|
||||
"cancelling the transfer context must tear down the attached IOStream")
|
||||
select {
|
||||
case <-readDone:
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("cancelling the transfer context must unblock the actual streamIO.Read")
|
||||
}
|
||||
select {
|
||||
case <-agentStream.closeSeen:
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("cancelling the transfer context must close the fake agent endpoint")
|
||||
}
|
||||
select {
|
||||
case <-agentStream.closed:
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("cancelling the transfer context must close the handler endpoint")
|
||||
}
|
||||
select {
|
||||
case <-agentStream.recvDone:
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("cancelling the transfer context must let the fake handler exit")
|
||||
}
|
||||
require.Equal(t, int32(1), agentStream.closeCall.Load(), "attached endpoint must be closed exactly once")
|
||||
require.Equal(t, 0, rpc.NezhaHandlerSingleton.StreamCount())
|
||||
for index := 0; index < 40; index++ {
|
||||
require.NoError(t, rpc.NezhaHandlerSingleton.CreateStream(fmt.Sprintf("cancel-reuse-%d", index), 0, 7))
|
||||
}
|
||||
require.ErrorIs(t, rpc.NezhaHandlerSingleton.CreateStream("cancel-reuse-over", 0, 7), rpc.ErrTooManyStreamsForServer)
|
||||
for index := 0; index < 40; index++ {
|
||||
require.NoError(t, rpc.NezhaHandlerSingleton.CloseStream(fmt.Sprintf("cancel-reuse-%d", index)))
|
||||
}
|
||||
require.Equal(t, 0, rpc.NezhaHandlerSingleton.StreamCount())
|
||||
|
||||
cleanupDone := make(chan struct{})
|
||||
go func() {
|
||||
var wg sync.WaitGroup
|
||||
for range 32 {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
cleanup()
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
close(cleanupDone)
|
||||
}()
|
||||
select {
|
||||
case <-cleanupDone:
|
||||
case <-time.After(2 * time.Second):
|
||||
t.Fatal("concurrent cleanup calls must complete")
|
||||
}
|
||||
}
|
||||
|
||||
@@ -586,5 +586,3 @@ func (b *bigReader) Read(p []byte) (int, error) {
|
||||
b.remaining -= int64(n)
|
||||
return n, nil
|
||||
}
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,84 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"io"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/hashicorp/go-uuid"
|
||||
|
||||
"github.com/nezhahq/nezha/model"
|
||||
pb "github.com/nezhahq/nezha/proto"
|
||||
"github.com/nezhahq/nezha/service/rpc"
|
||||
"github.com/nezhahq/nezha/service/singleton"
|
||||
)
|
||||
|
||||
// openFsTransferStream owns the task-to-agent IOStream lifecycle. The returned
|
||||
// cleanup is safe for concurrent callers and shares ownership with cancellation.
|
||||
func openFsTransferStream(ctx context.Context, serverID uint64, req *model.FsTransferRequest) (io.ReadWriteCloser, func(), error) {
|
||||
if singleton.Conf == nil || !singleton.Conf.MCPEnabled() {
|
||||
return nil, func() {}, errors.New("MCP is disabled by the dashboard administrator")
|
||||
}
|
||||
server, _ := singleton.ServerShared.Get(serverID)
|
||||
if server == nil || server.GetTaskStream() == nil {
|
||||
return nil, func() {}, errors.New("server offline")
|
||||
}
|
||||
handler := rpc.NezhaHandlerSingleton
|
||||
|
||||
streamID, err := uuid.GenerateUUID()
|
||||
if err != nil {
|
||||
return nil, func() {}, err
|
||||
}
|
||||
req.StreamID = streamID
|
||||
if err := handler.CreateStreamWithPurpose(streamID, 0, serverID, rpc.PurposeMCPTransfer); err != nil {
|
||||
return nil, func() {}, err
|
||||
}
|
||||
var cleanupOnce sync.Once
|
||||
cleanup := func() { cleanupOnce.Do(func() { _ = handler.CloseStream(streamID) }) }
|
||||
|
||||
body, err := json.Marshal(req)
|
||||
if err != nil {
|
||||
cleanup()
|
||||
return nil, func() {}, err
|
||||
}
|
||||
// The stream is owned by cleanup until the caller receives it; every failure path releases it.
|
||||
if singleton.Conf == nil || !singleton.Conf.MCPEnabled() {
|
||||
cleanup()
|
||||
return nil, func() {}, errors.New("MCP is disabled by the dashboard administrator")
|
||||
}
|
||||
if err := ctx.Err(); err != nil {
|
||||
cleanup()
|
||||
return nil, func() {}, err
|
||||
}
|
||||
if err := server.SendTask(&pb.Task{Type: model.TaskTypeFsTransfer, Data: string(body)}); err != nil {
|
||||
cleanup()
|
||||
if errors.Is(err, model.ErrTaskStreamOffline) {
|
||||
return nil, func() {}, errors.New("server offline")
|
||||
}
|
||||
return nil, func() {}, err
|
||||
}
|
||||
|
||||
agentStream, ok := handler.WaitForAgent(ctx, streamID, 30*time.Second)
|
||||
if !ok {
|
||||
cleanup()
|
||||
return nil, func() {}, errors.New("agent did not attach within 30s")
|
||||
}
|
||||
|
||||
watcherDone := make(chan struct{})
|
||||
var watcherOnce sync.Once
|
||||
go func() {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
cleanup()
|
||||
case <-watcherDone:
|
||||
}
|
||||
}()
|
||||
wrappedCleanup := func() {
|
||||
watcherOnce.Do(func() { close(watcherDone) })
|
||||
cleanup()
|
||||
}
|
||||
return agentStream, wrappedCleanup, nil
|
||||
}
|
||||
@@ -9,7 +9,6 @@ import (
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/goccy/go-json"
|
||||
"github.com/jinzhu/copier"
|
||||
"gorm.io/gorm"
|
||||
|
||||
"github.com/nezhahq/nezha/model"
|
||||
@@ -31,10 +30,13 @@ import (
|
||||
// @Router /server [get]
|
||||
func listServer(c *gin.Context) ([]*model.Server, error) {
|
||||
slist := singleton.ServerShared.GetSortedList()
|
||||
|
||||
var ssl []*model.Server
|
||||
if err := copier.Copy(&ssl, &slist); err != nil {
|
||||
return nil, err
|
||||
ssl := make([]*model.Server, 0, len(slist))
|
||||
for _, server := range slist {
|
||||
if server == nil {
|
||||
continue
|
||||
}
|
||||
runtime := server.RuntimeSnapshot()
|
||||
ssl = append(ssl, server.RuntimeCopy(runtime))
|
||||
}
|
||||
return ssl, nil
|
||||
}
|
||||
|
||||
@@ -32,6 +32,8 @@ func setupServerGroupVisibilityFixture(t *testing.T) {
|
||||
|
||||
db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
|
||||
require.NoError(t, err)
|
||||
sqlDB, err := db.DB()
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, db.AutoMigrate(&model.Server{}, &model.ServerGroup{}, &model.ServerGroupServer{}, &model.User{}))
|
||||
|
||||
singleton.DB = db
|
||||
@@ -55,6 +57,7 @@ func setupServerGroupVisibilityFixture(t *testing.T) {
|
||||
singleton.ServerShared = singleton.NewServerClass()
|
||||
|
||||
t.Cleanup(func() {
|
||||
_ = sqlDB.Close()
|
||||
singleton.DB = originalDB
|
||||
singleton.Cache = originalCache
|
||||
singleton.Loc = originalLoc
|
||||
|
||||
@@ -39,20 +39,14 @@ func setupServiceDispatchPATFixture(t *testing.T) {
|
||||
|
||||
db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
|
||||
require.NoError(t, err)
|
||||
sqlDB, err := db.DB()
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, db.AutoMigrate(&model.Service{}, &model.Server{}, &model.User{}, &model.ServiceHistory{}))
|
||||
|
||||
singleton.DB = db
|
||||
singleton.Loc = time.UTC
|
||||
singleton.Cache = cache.New(time.Minute, time.Minute)
|
||||
singleton.Localizer = i18n.NewLocalizer("en_US", "nezha", "translations", i18n.Translations)
|
||||
// ServiceSentinel 在构造时会调 CronShared.AddFunc 注册每日/每周维护任务,
|
||||
// 必须先于 NewServiceSentinel 装配。
|
||||
singleton.CronShared = singleton.NewCronClass()
|
||||
|
||||
sentinel, err := singleton.NewServiceSentinel(make(chan *model.Service, 4))
|
||||
require.NoError(t, err)
|
||||
singleton.ServiceSentinelShared = sentinel
|
||||
|
||||
sc := singleton.NewEmptyServerClassForTest()
|
||||
for _, id := range []uint64{1, 2} {
|
||||
s := &model.Server{}
|
||||
@@ -61,6 +55,13 @@ func setupServiceDispatchPATFixture(t *testing.T) {
|
||||
sc.InsertForTest(s)
|
||||
}
|
||||
singleton.ServerShared = sc
|
||||
// ServiceSentinel 在构造时会调 CronShared.AddFunc 注册每日/每周维护任务,
|
||||
// 必须先于 NewServiceSentinel 装配。
|
||||
singleton.CronShared = singleton.NewCronClass()
|
||||
|
||||
sentinel, err := singleton.NewServiceSentinel(make(chan *model.Service, 4))
|
||||
require.NoError(t, err)
|
||||
singleton.ServiceSentinelShared = sentinel
|
||||
|
||||
singleton.UserLock.Lock()
|
||||
singleton.UserInfoMap = map[uint64]model.UserInfo{100: {Role: model.RoleMember}}
|
||||
@@ -68,6 +69,8 @@ func setupServiceDispatchPATFixture(t *testing.T) {
|
||||
|
||||
t.Cleanup(func() {
|
||||
sentinel.Close()
|
||||
singleton.CronShared.Close()
|
||||
_ = sqlDB.Close()
|
||||
singleton.ServiceSentinelShared = originalSentinel
|
||||
singleton.CronShared = originalCron
|
||||
singleton.DB = originalDB
|
||||
|
||||
@@ -32,6 +32,8 @@ func setupTenancyTest(t *testing.T) func() {
|
||||
}
|
||||
db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
|
||||
require.NoError(t, err)
|
||||
sqlDB, err := db.DB()
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, db.AutoMigrate(
|
||||
&model.User{},
|
||||
&model.Cron{},
|
||||
@@ -47,6 +49,7 @@ func setupTenancyTest(t *testing.T) func() {
|
||||
singleton.DDNSShared = singleton.NewEmptyDDNSClassForTest()
|
||||
singleton.NotificationShared = singleton.NewEmptyNotificationClassForTest()
|
||||
return func() {
|
||||
_ = sqlDB.Close()
|
||||
singleton.DB = originalDB
|
||||
singleton.Localizer = originalLocalizer
|
||||
singleton.ServerShared = originalServer
|
||||
@@ -118,16 +121,16 @@ func TestTenancy_CreateDDNS_InjectedUserIDIgnored(t *testing.T) {
|
||||
defer setupTenancyTest(t)()
|
||||
|
||||
body := map[string]any{
|
||||
"name": "evil-ddns",
|
||||
"provider": "webhook",
|
||||
"access_id": "x",
|
||||
"access_secret": "y",
|
||||
"webhook_url": "http://127.0.0.1/",
|
||||
"webhook_method": "GET",
|
||||
"name": "evil-ddns",
|
||||
"provider": "webhook",
|
||||
"access_id": "x",
|
||||
"access_secret": "y",
|
||||
"webhook_url": "http://127.0.0.1/",
|
||||
"webhook_method": "GET",
|
||||
"webhook_request_type": "json",
|
||||
"webhook_request_body": "",
|
||||
"webhook_headers": "",
|
||||
"user_id": 999, // attacker
|
||||
"webhook_headers": "",
|
||||
"user_id": 999, // attacker
|
||||
}
|
||||
c := ctxAsMemberWithBody(10, body)
|
||||
_, err := createDDNS(c)
|
||||
@@ -155,12 +158,12 @@ func TestTenancy_UpdateDDNS_ForeignOwnerRejected(t *testing.T) {
|
||||
require.NoError(t, singleton.DB.Create(&foreign).Error)
|
||||
|
||||
c := ctxAsMemberWithBody(10, map[string]any{
|
||||
"name": "hijacked",
|
||||
"provider": "webhook",
|
||||
"access_id": "x",
|
||||
"access_secret": "y",
|
||||
"webhook_url": "http://attacker/",
|
||||
"webhook_method": "GET",
|
||||
"name": "hijacked",
|
||||
"provider": "webhook",
|
||||
"access_id": "x",
|
||||
"access_secret": "y",
|
||||
"webhook_url": "http://attacker/",
|
||||
"webhook_method": "GET",
|
||||
"webhook_request_type": "json",
|
||||
})
|
||||
c.Params = gin.Params{{Key: "id", Value: itoa(foreign.ID)}}
|
||||
@@ -237,7 +240,7 @@ func TestTenancy_UpdateNotificationGroup_ForeignOwnerRejected(t *testing.T) {
|
||||
require.NoError(t, singleton.DB.Create(&foreign).Error)
|
||||
|
||||
c := ctxAsMemberWithBody(10, map[string]any{
|
||||
"name": "hijacked",
|
||||
"name": "hijacked",
|
||||
"notifications": []uint64{},
|
||||
})
|
||||
c.Params = gin.Params{{Key: "id", Value: itoa(foreign.ID)}}
|
||||
|
||||
@@ -0,0 +1,79 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"context"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/gorilla/websocket"
|
||||
)
|
||||
|
||||
type websocketPingWriter interface {
|
||||
WriteMessage(messageType int, data []byte) error
|
||||
}
|
||||
|
||||
type websocketPingConnection interface {
|
||||
websocketPingWriter
|
||||
Close() error
|
||||
}
|
||||
|
||||
type websocketPingTransport struct {
|
||||
websocketPingWriter
|
||||
closeOnce sync.Once
|
||||
closeErr error
|
||||
close func() error
|
||||
}
|
||||
|
||||
func newWebsocketPingTransport(writer websocketPingWriter, close func() error) *websocketPingTransport {
|
||||
return &websocketPingTransport{websocketPingWriter: writer, close: close}
|
||||
}
|
||||
|
||||
func (transport *websocketPingTransport) Close() error {
|
||||
transport.closeOnce.Do(func() { transport.closeErr = transport.close() })
|
||||
return transport.closeErr
|
||||
}
|
||||
|
||||
func websocketPingLoop(ctx context.Context, ticks <-chan time.Time, writer websocketPingWriter) error {
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return nil
|
||||
case _, ok := <-ticks:
|
||||
if !ok {
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return nil
|
||||
default:
|
||||
}
|
||||
if err := writer.WriteMessage(websocket.PingMessage, []byte{}); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func startWebsocketPing(ctx context.Context, ticks <-chan time.Time, connection websocketPingConnection) func() {
|
||||
workerContext, cancel := context.WithCancel(ctx)
|
||||
workerDone := make(chan struct{})
|
||||
go func() {
|
||||
defer close(workerDone)
|
||||
_ = websocketPingLoop(workerContext, ticks, connection)
|
||||
}()
|
||||
return func() {
|
||||
_ = connection.Close()
|
||||
cancel()
|
||||
<-workerDone
|
||||
}
|
||||
}
|
||||
|
||||
func startWebsocketPingTicker(ctx context.Context, interval time.Duration, connection websocketPingConnection) func() {
|
||||
ticker := time.NewTicker(interval)
|
||||
stop := startWebsocketPing(ctx, ticker.C, connection)
|
||||
return func() {
|
||||
ticker.Stop()
|
||||
stop()
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,152 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
type pingWriterFake struct {
|
||||
mu sync.Mutex
|
||||
writeCalls int
|
||||
writeErr error
|
||||
writeStarted chan struct{}
|
||||
continueWrite chan struct{}
|
||||
}
|
||||
|
||||
func (writer *pingWriterFake) Close() error { return nil }
|
||||
|
||||
type permanentlyBlockedPingWriter struct {
|
||||
writeStarted chan struct{}
|
||||
transportClosed chan struct{}
|
||||
}
|
||||
|
||||
func (writer *permanentlyBlockedPingWriter) WriteMessage(int, []byte) error {
|
||||
close(writer.writeStarted)
|
||||
<-writer.transportClosed
|
||||
return errors.New("transport closed")
|
||||
}
|
||||
|
||||
func (writer *permanentlyBlockedPingWriter) Close() error {
|
||||
close(writer.transportClosed)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (writer *pingWriterFake) WriteMessage(int, []byte) error {
|
||||
writer.mu.Lock()
|
||||
writer.writeCalls++
|
||||
writer.mu.Unlock()
|
||||
if writer.writeStarted != nil {
|
||||
close(writer.writeStarted)
|
||||
<-writer.continueWrite
|
||||
}
|
||||
return writer.writeErr
|
||||
}
|
||||
|
||||
func (writer *pingWriterFake) calls() int {
|
||||
writer.mu.Lock()
|
||||
defer writer.mu.Unlock()
|
||||
return writer.writeCalls
|
||||
}
|
||||
|
||||
func TestWebsocketPingLoop_stopsAndJoinsWithoutWritingAfterStop(t *testing.T) {
|
||||
// Given
|
||||
ticks := make(chan time.Time, 1)
|
||||
writer := &pingWriterFake{writeStarted: make(chan struct{}), continueWrite: make(chan struct{})}
|
||||
stop := startWebsocketPing(context.Background(), ticks, writer)
|
||||
|
||||
// When
|
||||
ticks <- time.Time{}
|
||||
<-writer.writeStarted
|
||||
close(writer.continueWrite)
|
||||
stop()
|
||||
ticks <- time.Time{}
|
||||
|
||||
// Then
|
||||
require.Equal(t, 1, writer.calls())
|
||||
}
|
||||
|
||||
func TestWebsocketPingLoop_exitsOnWriteError(t *testing.T) {
|
||||
// Given
|
||||
ticks := make(chan time.Time, 1)
|
||||
writer := &pingWriterFake{writeErr: errors.New("closed")}
|
||||
done := make(chan error, 1)
|
||||
go func() { done <- websocketPingLoop(context.Background(), ticks, writer) }()
|
||||
|
||||
// When
|
||||
ticks <- time.Time{}
|
||||
|
||||
// Then
|
||||
require.Error(t, <-done)
|
||||
ticks <- time.Time{}
|
||||
require.Equal(t, 1, writer.calls())
|
||||
}
|
||||
|
||||
func TestWebsocketPingLoop_cleanupOverlapsTickAndJoinsWriter(t *testing.T) {
|
||||
// Given
|
||||
ticks := make(chan time.Time, 1)
|
||||
writer := &pingWriterFake{
|
||||
writeStarted: make(chan struct{}),
|
||||
continueWrite: make(chan struct{}),
|
||||
}
|
||||
stop := startWebsocketPing(context.Background(), ticks, writer)
|
||||
ticks <- time.Time{}
|
||||
<-writer.writeStarted
|
||||
|
||||
// When
|
||||
stopped := make(chan struct{})
|
||||
go func() {
|
||||
stop()
|
||||
close(stopped)
|
||||
}()
|
||||
select {
|
||||
case <-stopped:
|
||||
require.Fail(t, "ping worker stop returned before the in-flight write joined")
|
||||
default:
|
||||
}
|
||||
close(writer.continueWrite)
|
||||
<-stopped
|
||||
ticks <- time.Time{}
|
||||
|
||||
// Then
|
||||
require.Equal(t, 1, writer.calls())
|
||||
}
|
||||
|
||||
func TestWebsocketPingStop_unblocksPermanentlyBlockedWriteBeforeJoin(t *testing.T) {
|
||||
// Given
|
||||
ticks := make(chan time.Time, 1)
|
||||
writer := &permanentlyBlockedPingWriter{
|
||||
writeStarted: make(chan struct{}),
|
||||
transportClosed: make(chan struct{}),
|
||||
}
|
||||
stop := startWebsocketPing(context.Background(), ticks, writer)
|
||||
ticks <- time.Time{}
|
||||
<-writer.writeStarted
|
||||
|
||||
// When
|
||||
stopped := make(chan struct{})
|
||||
go func() {
|
||||
stop()
|
||||
close(stopped)
|
||||
}()
|
||||
|
||||
// Then
|
||||
deadline := time.NewTimer(time.Second)
|
||||
defer deadline.Stop()
|
||||
select {
|
||||
case <-writer.transportClosed:
|
||||
case <-deadline.C:
|
||||
require.Fail(t, "stop did not close the blocked ping transport")
|
||||
}
|
||||
select {
|
||||
case <-stopped:
|
||||
case <-deadline.C:
|
||||
require.Fail(t, "stop did not join after closing the blocked ping transport")
|
||||
}
|
||||
}
|
||||
|
||||
var _ websocketPingWriter = (*pingWriterFake)(nil)
|
||||
@@ -0,0 +1,77 @@
|
||||
package rpc
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"log"
|
||||
|
||||
"github.com/nezhahq/nezha/model"
|
||||
"github.com/nezhahq/nezha/proto"
|
||||
"github.com/nezhahq/nezha/service/singleton"
|
||||
)
|
||||
|
||||
func DispatchTask(serviceSentinelDispatchBus <-chan *model.Service) {
|
||||
for task := range serviceSentinelDispatchBus {
|
||||
if task == nil {
|
||||
continue
|
||||
}
|
||||
|
||||
switch task.Cover {
|
||||
case model.ServiceCoverIgnoreAll:
|
||||
for id, enabled := range task.SkipServers {
|
||||
if !enabled {
|
||||
continue
|
||||
}
|
||||
|
||||
server, _ := singleton.ServerShared.Get(id)
|
||||
if server == nil {
|
||||
continue
|
||||
}
|
||||
if !canSendTaskToServer(task, server) {
|
||||
continue
|
||||
}
|
||||
if err := server.SendTask(task.PB()); err != nil && !errors.Is(err, model.ErrTaskStreamOffline) {
|
||||
log.Printf("NEZHA>> DispatchTask send error (server=%d): %v", id, err)
|
||||
}
|
||||
}
|
||||
case model.ServiceCoverAll:
|
||||
for id, server := range singleton.ServerShared.GetList() {
|
||||
if server == nil || task.SkipServers[id] {
|
||||
continue
|
||||
}
|
||||
if !canSendTaskToServer(task, server) {
|
||||
continue
|
||||
}
|
||||
if err := server.SendTask(task.PB()); err != nil && !errors.Is(err, model.ErrTaskStreamOffline) {
|
||||
log.Printf("NEZHA>> DispatchTask send error (server=%d): %v", id, err)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func DispatchKeepalive() {
|
||||
singleton.CronShared.AddFunc("@every 20s", func() {
|
||||
list := singleton.ServerShared.GetSortedList()
|
||||
for _, s := range list {
|
||||
if s == nil {
|
||||
continue
|
||||
}
|
||||
if err := s.SendTask(&proto.Task{Type: model.TaskTypeKeepalive}); err != nil && !errors.Is(err, model.ErrTaskStreamOffline) {
|
||||
log.Printf("NEZHA>> Keepalive send error (server=%d): %v", s.ID, err)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func canSendTaskToServer(task *model.Service, server *model.Server) bool {
|
||||
var role model.Role
|
||||
singleton.UserLock.RLock()
|
||||
if u, ok := singleton.UserInfoMap[task.UserID]; !ok {
|
||||
role = model.RoleMember
|
||||
} else {
|
||||
role = u.Role
|
||||
}
|
||||
singleton.UserLock.RUnlock()
|
||||
|
||||
return task.UserID == server.GetUserID() || role.IsAdmin()
|
||||
}
|
||||
+10
-10
@@ -25,18 +25,18 @@ func (MCPAuditLog) TableName() string {
|
||||
}
|
||||
|
||||
const (
|
||||
MCPOutcomeOK = "ok"
|
||||
MCPOutcomeScopeDenied = "scope_denied"
|
||||
MCPOutcomePermDenied = "permission_denied"
|
||||
MCPOutcomeServerOffline = "server_offline"
|
||||
MCPOutcomeAgentTimeout = "agent_timeout"
|
||||
MCPOutcomeAgentError = "agent_error"
|
||||
MCPOutcomeOK = "ok"
|
||||
MCPOutcomeScopeDenied = "scope_denied"
|
||||
MCPOutcomePermDenied = "permission_denied"
|
||||
MCPOutcomeServerOffline = "server_offline"
|
||||
MCPOutcomeAgentTimeout = "agent_timeout"
|
||||
MCPOutcomeAgentError = "agent_error"
|
||||
// MCPOutcomeMCPDisabled 区分 “管理员按下 kill switch 把 MCP 关了” 与
|
||||
// “agent 真出故障” 两种语义:前者属 forbidden 类、不应该触发 agent
|
||||
// 故障告警;详见 service/rpc/mcp_rpc.go 里 ErrMCPDisabled 的注释。
|
||||
MCPOutcomeMCPDisabled = "mcp_disabled"
|
||||
MCPOutcomeInvalidArgs = "invalid_args"
|
||||
MCPOutcomeRateLimited = "rate_limited"
|
||||
MCPOutcomeMCPDisabled = "mcp_disabled"
|
||||
MCPOutcomeInvalidArgs = "invalid_args"
|
||||
MCPOutcomeRateLimited = "rate_limited"
|
||||
MCPOutcomeUnsupportedAgent = "unsupported_agent"
|
||||
MCPOutcomeInternalError = "internal_error"
|
||||
MCPOutcomeInternalError = "internal_error"
|
||||
)
|
||||
|
||||
@@ -16,9 +16,9 @@ type SettingForm struct {
|
||||
AgentRealIPHeader string `json:"agent_real_ip_header,omitempty" validate:"optional"` // Agent真实IP
|
||||
UserTemplate string `json:"user_template,omitempty" validate:"optional"`
|
||||
|
||||
AgentTLS bool `json:"tls,omitempty" validate:"optional"`
|
||||
EnableIPChangeNotification bool `json:"enable_ip_change_notification,omitempty" validate:"optional"`
|
||||
EnablePlainIPInNotification bool `json:"enable_plain_ip_in_notification,omitempty" validate:"optional"`
|
||||
AgentTLS bool `json:"tls,omitempty" validate:"optional"`
|
||||
EnableIPChangeNotification bool `json:"enable_ip_change_notification,omitempty" validate:"optional"`
|
||||
EnablePlainIPInNotification bool `json:"enable_plain_ip_in_notification,omitempty" validate:"optional"`
|
||||
EnableMCP *bool `json:"enable_mcp,omitempty" validate:"optional"`
|
||||
}
|
||||
|
||||
|
||||
@@ -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)
|
||||
|
||||
+50
-17
@@ -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 {
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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) {
|
||||
|
||||
Reference in New Issue
Block a user