mirror of
https://github.com/Buriburizaem0n/nezha_domains.git
synced 2026-09-19 17:50:12 +00:00
naibaandcloudcode
e8dabf5bc6
Introduce Personal Access Tokens (nzp_*) as a stateless auth path alongside
JWT, gated per-endpoint by a scope middleware (nezha:{resource}:{verb}) with
fail-closed empty-scope defaults and a server-id whitelist. Self-management
endpoints (profile, api-tokens, oauth2 bind, refresh-token) explicitly reject
PATs to block privilege-escalation chains. A revoke registry tears down active
long-lived connections (terminal, fm, ws, transfer, mcp) the moment a PAT is
deleted, with a tombstone closing the revoke->register race.
Add an MCP endpoint that proxies tool calls (exec, fs read/write/delete,
transfer) to agents over gRPC, guarded by origin/DNS-rebinding checks, a
per-token rate limiter, audit logging, and a kill switch. Serialize all
sends through the IOStream wrapper to honour grpc-go's concurrency contract.
Add CSRF double-submit protection on unsafe cookie-authenticated methods,
exempting authenticated PAT requests by context identity (not a forgeable
Authorization header). Apply visibility/whitelist filtering consistently
across list, get-by-id, and mutate paths to enforce tenant isolation.
Migrate legacy mcp:* scopes: rewrite read/exec to nezha:* equivalents and
drop dangerous write/delete/wildcard grants.
Co-authored-by: cloudcode <cloudcode@users.noreply.github.com>
304 lines
10 KiB
Go
304 lines
10 KiB
Go
package controller
|
|
|
|
import (
|
|
"bytes"
|
|
"encoding/json"
|
|
"net/http/httptest"
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/gin-gonic/gin"
|
|
"github.com/stretchr/testify/require"
|
|
"gorm.io/driver/sqlite"
|
|
"gorm.io/gorm"
|
|
|
|
"github.com/nezhahq/nezha/model"
|
|
"github.com/nezhahq/nezha/pkg/i18n"
|
|
"github.com/nezhahq/nezha/service/singleton"
|
|
)
|
|
|
|
func setupMCPTest(t *testing.T) (func(), uint64) {
|
|
t.Helper()
|
|
originalDB := singleton.DB
|
|
originalServer := singleton.ServerShared
|
|
originalConf := singleton.Conf
|
|
originalAuditSync := mcpAuditSync
|
|
originalLimiter := mcpRateLimiterShared
|
|
originalLocalizer := singleton.Localizer
|
|
originalPATRegistry := patConnectionRegistryShared
|
|
singleton.Localizer = i18n.NewLocalizer("en_US", "nezha", "translations", i18n.Translations)
|
|
mcpAuditSync = true
|
|
mcpRateLimiterShared = newMCPRateLimiter(1000, 10000)
|
|
// Fresh per test: the DB resets token IDs to 1 each run, so a stale
|
|
// revoke tombstone from a prior test would otherwise cancel a reused id.
|
|
patConnectionRegistryShared = newPATConnectionRegistry()
|
|
db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
|
|
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}}
|
|
singleton.Conf.SetMCPEnabled(true)
|
|
|
|
user := model.User{Common: model.Common{ID: 100}, Username: "alice", Role: model.RoleMember}
|
|
require.NoError(t, db.Create(&user).Error)
|
|
|
|
sc := singleton.NewEmptyServerClassForTest()
|
|
srv := &model.Server{}
|
|
srv.ID = 7
|
|
srv.Name = "alpha"
|
|
srv.SetUserID(100)
|
|
sc.InsertForTest(srv)
|
|
singleton.ServerShared = sc
|
|
|
|
cleanup := func() {
|
|
singleton.DB = originalDB
|
|
singleton.ServerShared = originalServer
|
|
singleton.Conf = originalConf
|
|
singleton.Localizer = originalLocalizer
|
|
mcpAuditSync = originalAuditSync
|
|
mcpRateLimiterShared = originalLimiter
|
|
patConnectionRegistryShared = originalPATRegistry
|
|
}
|
|
return cleanup, user.ID
|
|
}
|
|
|
|
func mkToken(t *testing.T, uid uint64, scopes []string, serverIDs []uint64) (*model.APIToken, string) {
|
|
t.Helper()
|
|
plain := "nzp_" + strings.Repeat("a", 32) + "_" + ctoa(uid)
|
|
tok := model.APIToken{UserID: uid, Name: "t", TokenHash: model.HashAPIToken(plain)}
|
|
tok.SetScopes(scopes)
|
|
if len(serverIDs) > 0 {
|
|
tok.SetServerIDs(serverIDs)
|
|
}
|
|
require.NoError(t, singleton.DB.Create(&tok).Error)
|
|
return &tok, plain
|
|
}
|
|
|
|
func mcpCallCtx(t *testing.T, tok *model.APIToken, uid uint64, body any) (*gin.Context, *httptest.ResponseRecorder) {
|
|
t.Helper()
|
|
gin.SetMode(gin.TestMode)
|
|
w := httptest.NewRecorder()
|
|
c, _ := gin.CreateTestContext(w)
|
|
b, _ := json.Marshal(body)
|
|
c.Request = httptest.NewRequest("POST", "/mcp", bytes.NewReader(b))
|
|
c.Request.Header.Set("Content-Type", "application/json")
|
|
c.Set(model.CtxKeyAuthorizedUser, &model.User{Common: model.Common{ID: uid}, Role: model.RoleMember})
|
|
c.Set(apiTokenCtxKey, tok)
|
|
c.Set(model.CtxKeyAPIToken, tok)
|
|
return c, w
|
|
}
|
|
|
|
func decodeRPC(w *httptest.ResponseRecorder) (jsonRPCResponse, *mcpToolCallResult) {
|
|
var env jsonRPCResponse
|
|
_ = json.Unmarshal(w.Body.Bytes(), &env)
|
|
if env.Result == nil {
|
|
return env, nil
|
|
}
|
|
rb, _ := json.Marshal(env.Result)
|
|
var tcr mcpToolCallResult
|
|
_ = json.Unmarshal(rb, &tcr)
|
|
return env, &tcr
|
|
}
|
|
|
|
func TestMCP_RejectsMissingToken(t *testing.T) {
|
|
cleanup, _ := setupMCPTest(t)
|
|
defer cleanup()
|
|
|
|
gin.SetMode(gin.TestMode)
|
|
w := httptest.NewRecorder()
|
|
c, _ := gin.CreateTestContext(w)
|
|
body, _ := json.Marshal(jsonRPCRequest{JSONRPC: "2.0", ID: json.RawMessage("1"), Method: "initialize"})
|
|
c.Request = httptest.NewRequest("POST", "/mcp", bytes.NewReader(body))
|
|
c.Request.Header.Set("Content-Type", "application/json")
|
|
mcpEndpoint(c)
|
|
var env jsonRPCResponse
|
|
require.NoError(t, json.Unmarshal(w.Body.Bytes(), &env))
|
|
require.NotNil(t, env.Error)
|
|
require.Equal(t, rpcErrUnauthorized, env.Error.Code)
|
|
}
|
|
|
|
func TestMCP_Initialize_ReturnsServerInfo(t *testing.T) {
|
|
cleanup, uid := setupMCPTest(t)
|
|
defer cleanup()
|
|
tok, _ := mkToken(t, uid, []string{model.ScopeServerRead}, nil)
|
|
c, w := mcpCallCtx(t, tok, uid, jsonRPCRequest{JSONRPC: "2.0", ID: json.RawMessage("1"), Method: "initialize"})
|
|
mcpEndpoint(c)
|
|
env, _ := decodeRPC(w)
|
|
require.Nil(t, env.Error)
|
|
rb, _ := json.Marshal(env.Result)
|
|
require.Contains(t, string(rb), "nezha-mcp")
|
|
require.Contains(t, string(rb), "protocolVersion")
|
|
}
|
|
|
|
func TestMCP_ToolsList_IncludesRegisteredTools(t *testing.T) {
|
|
cleanup, uid := setupMCPTest(t)
|
|
defer cleanup()
|
|
tok, _ := mkToken(t, uid, []string{model.ScopeServerRead}, nil)
|
|
c, w := mcpCallCtx(t, tok, uid, jsonRPCRequest{JSONRPC: "2.0", ID: json.RawMessage("1"), Method: "tools/list"})
|
|
mcpEndpoint(c)
|
|
env, _ := decodeRPC(w)
|
|
require.Nil(t, env.Error)
|
|
rb, _ := json.Marshal(env.Result)
|
|
for _, name := range []string{"meta.whoami", "server.list", "server.exec", "fs.list", "fs.read", "fs.write", "fs.delete", "fs.download_url", "fs.upload_url"} {
|
|
require.Contains(t, string(rb), name)
|
|
}
|
|
}
|
|
|
|
func TestMCP_Whoami_HappyPath(t *testing.T) {
|
|
cleanup, uid := setupMCPTest(t)
|
|
defer cleanup()
|
|
tok, _ := mkToken(t, uid, []string{model.ScopeServerRead, model.ScopeServerRead}, []uint64{7, 8})
|
|
c, w := mcpCallCtx(t, tok, uid, jsonRPCRequest{
|
|
JSONRPC: "2.0", ID: json.RawMessage("1"), Method: "tools/call",
|
|
Params: jsonObj(t, toolCallParams{Name: "meta.whoami", Arguments: json.RawMessage("{}")}),
|
|
})
|
|
mcpEndpoint(c)
|
|
_, tcr := decodeRPC(w)
|
|
require.NotNil(t, tcr)
|
|
require.False(t, tcr.IsError, "got error content: %v", tcr.Content)
|
|
scb, _ := json.Marshal(tcr.StructuredContent)
|
|
require.Contains(t, string(scb), "user_id")
|
|
require.Contains(t, string(scb), "scopes")
|
|
}
|
|
|
|
func TestMCP_ScopeDenied(t *testing.T) {
|
|
cleanup, uid := setupMCPTest(t)
|
|
defer cleanup()
|
|
tok, _ := mkToken(t, uid, []string{model.ScopeServerRead}, nil)
|
|
c, w := mcpCallCtx(t, tok, uid, jsonRPCRequest{
|
|
JSONRPC: "2.0", ID: json.RawMessage("1"), Method: "tools/call",
|
|
Params: jsonObj(t, toolCallParams{Name: "server.exec", Arguments: jsonRaw(map[string]any{"server_id": 7, "cmd": "echo"})}),
|
|
})
|
|
mcpEndpoint(c)
|
|
_, tcr := decodeRPC(w)
|
|
require.NotNil(t, tcr)
|
|
require.True(t, tcr.IsError)
|
|
require.Contains(t, tcr.Content[0].Text, "missing required scope")
|
|
}
|
|
|
|
func TestMCP_PermissionDenied_WhenWrongUserOwnsServer(t *testing.T) {
|
|
cleanup, _ := setupMCPTest(t)
|
|
defer cleanup()
|
|
require.NoError(t, singleton.DB.Create(&model.User{Common: model.Common{ID: 200}, Username: "bob", Role: model.RoleMember}).Error)
|
|
tok, _ := mkToken(t, 200, []string{model.ScopeServerRead}, nil)
|
|
c, w := mcpCallCtx(t, tok, 200, jsonRPCRequest{
|
|
JSONRPC: "2.0", ID: json.RawMessage("1"), Method: "tools/call",
|
|
Params: jsonObj(t, toolCallParams{Name: "fs.list", Arguments: jsonRaw(map[string]any{"server_id": 7, "path": "/tmp"})}),
|
|
})
|
|
mcpEndpoint(c)
|
|
_, tcr := decodeRPC(w)
|
|
require.NotNil(t, tcr)
|
|
require.True(t, tcr.IsError)
|
|
require.Contains(t, tcr.Content[0].Text, "permission denied")
|
|
}
|
|
|
|
func TestMCP_ServerWhitelist_DenyOutsideList(t *testing.T) {
|
|
cleanup, uid := setupMCPTest(t)
|
|
defer cleanup()
|
|
tok, _ := mkToken(t, uid, []string{model.ScopeServerRead}, []uint64{99})
|
|
c, w := mcpCallCtx(t, tok, uid, jsonRPCRequest{
|
|
JSONRPC: "2.0", ID: json.RawMessage("1"), Method: "tools/call",
|
|
Params: jsonObj(t, toolCallParams{Name: "fs.list", Arguments: jsonRaw(map[string]any{"server_id": 7, "path": "/tmp"})}),
|
|
})
|
|
mcpEndpoint(c)
|
|
_, tcr := decodeRPC(w)
|
|
require.NotNil(t, tcr)
|
|
require.True(t, tcr.IsError)
|
|
}
|
|
|
|
func TestMCP_UnknownTool(t *testing.T) {
|
|
cleanup, uid := setupMCPTest(t)
|
|
defer cleanup()
|
|
tok, _ := mkToken(t, uid, []string{model.ScopeServerRead}, nil)
|
|
c, w := mcpCallCtx(t, tok, uid, jsonRPCRequest{
|
|
JSONRPC: "2.0", ID: json.RawMessage("1"), Method: "tools/call",
|
|
Params: jsonObj(t, toolCallParams{Name: "does.not.exist", Arguments: json.RawMessage("{}")}),
|
|
})
|
|
mcpEndpoint(c)
|
|
env, _ := decodeRPC(w)
|
|
require.NotNil(t, env.Error)
|
|
require.Equal(t, rpcErrMethodNotFound, env.Error.Code)
|
|
}
|
|
|
|
func TestMCP_InvalidJSONEnvelope(t *testing.T) {
|
|
cleanup, uid := setupMCPTest(t)
|
|
defer cleanup()
|
|
tok, _ := mkToken(t, uid, []string{model.ScopeServerRead}, nil)
|
|
|
|
gin.SetMode(gin.TestMode)
|
|
w := httptest.NewRecorder()
|
|
c, _ := gin.CreateTestContext(w)
|
|
c.Request = httptest.NewRequest("POST", "/mcp", bytes.NewReader([]byte("garbage")))
|
|
c.Request.Header.Set("Content-Type", "application/json")
|
|
c.Set(model.CtxKeyAuthorizedUser, &model.User{Common: model.Common{ID: uid}, Role: model.RoleMember})
|
|
c.Set(apiTokenCtxKey, tok)
|
|
mcpEndpoint(c)
|
|
var env jsonRPCResponse
|
|
require.NoError(t, json.Unmarshal(w.Body.Bytes(), &env))
|
|
require.NotNil(t, env.Error)
|
|
require.Equal(t, rpcErrParse, env.Error.Code)
|
|
}
|
|
|
|
func TestMCP_AuditRowIsWritten(t *testing.T) {
|
|
cleanup, uid := setupMCPTest(t)
|
|
defer cleanup()
|
|
tok, _ := mkToken(t, uid, []string{model.ScopeServerRead}, nil)
|
|
c, _ := mcpCallCtx(t, tok, uid, jsonRPCRequest{
|
|
JSONRPC: "2.0", ID: json.RawMessage("1"), Method: "tools/call",
|
|
Params: jsonObj(t, toolCallParams{Name: "meta.whoami", Arguments: json.RawMessage("{}")}),
|
|
})
|
|
mcpEndpoint(c)
|
|
|
|
require.Eventually(t, func() bool {
|
|
var cnt int64
|
|
_ = singleton.DB.Model(&model.MCPAuditLog{}).Where("token_id = ?", tok.ID).Count(&cnt).Error
|
|
return cnt == 1
|
|
}, 2*time.Second, 20*time.Millisecond, "audit row never appeared")
|
|
}
|
|
|
|
func TestMCP_RateLimit(t *testing.T) {
|
|
cleanup, uid := setupMCPTest(t)
|
|
defer cleanup()
|
|
original := mcpRateLimiterShared
|
|
mcpRateLimiterShared = newMCPRateLimiter(2, 100)
|
|
defer func() { mcpRateLimiterShared = original }()
|
|
|
|
tok, _ := mkToken(t, uid, []string{model.ScopeServerRead}, nil)
|
|
|
|
for i := 0; i < 2; i++ {
|
|
c, w := mcpCallCtx(t, tok, uid, jsonRPCRequest{
|
|
JSONRPC: "2.0", ID: json.RawMessage("1"), Method: "tools/call",
|
|
Params: jsonObj(t, toolCallParams{Name: "meta.whoami", Arguments: json.RawMessage("{}")}),
|
|
})
|
|
mcpEndpoint(c)
|
|
_, tcr := decodeRPC(w)
|
|
require.False(t, tcr.IsError)
|
|
}
|
|
c, w := mcpCallCtx(t, tok, uid, jsonRPCRequest{
|
|
JSONRPC: "2.0", ID: json.RawMessage("1"), Method: "tools/call",
|
|
Params: jsonObj(t, toolCallParams{Name: "meta.whoami", Arguments: json.RawMessage("{}")}),
|
|
})
|
|
mcpEndpoint(c)
|
|
_, tcr := decodeRPC(w)
|
|
require.True(t, tcr.IsError)
|
|
require.Contains(t, tcr.Content[0].Text, "rate limit")
|
|
}
|
|
|
|
func jsonObj(t *testing.T, v any) json.RawMessage {
|
|
t.Helper()
|
|
b, err := json.Marshal(v)
|
|
require.NoError(t, err)
|
|
return b
|
|
}
|
|
|
|
func jsonRaw(v map[string]any) json.RawMessage {
|
|
b, _ := json.Marshal(v)
|
|
return b
|
|
}
|
|
|
|
func ctoa(v uint64) string {
|
|
b, _ := json.Marshal(v)
|
|
return string(b)
|
|
}
|