mirror of
https://github.com/Buriburizaem0n/nezha_domains.git
synced 2026-09-19 09:40:12 +00:00
239 lines
12 KiB
Go
239 lines
12 KiB
Go
//go:build agentcompat
|
|
|
|
package controller
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"encoding/base64"
|
|
"encoding/json"
|
|
"io"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/gin-gonic/gin"
|
|
"github.com/stretchr/testify/require"
|
|
|
|
"github.com/nezhahq/nezha/model"
|
|
"github.com/nezhahq/nezha/service/rpc"
|
|
"github.com/nezhahq/nezha/service/singleton"
|
|
)
|
|
|
|
func TestAgentcompatCapabilityCancelAndUnregisterAreUniformAndNonMutating(t *testing.T) {
|
|
cleanup, userID := setupMCPTest(t)
|
|
t.Cleanup(cleanup)
|
|
handler := rpc.NewNezhaHandler()
|
|
originalHandler := rpc.NezhaHandlerSingleton
|
|
rpc.NezhaHandlerSingleton = handler
|
|
t.Cleanup(func() { rpc.NezhaHandlerSingleton = originalHandler })
|
|
ownerToken, ownerPAT := mkToken(t, userID, []string{model.ScopeServerRead}, nil)
|
|
_, foreignPAT := mkDistinctCapabilityToken(t, userID, "foreign")
|
|
server := newAgentcompatCapabilityServer(t)
|
|
capability := registerAgentcompatCapability(t, server.URL, ownerPAT, agentcompatCapabilityRegisterRequest{Purpose: agentcompatCapabilityPurposeTerminal, ServerID: 7})
|
|
bindAgentcompatTerminalCapability(t, handler, capability, ownerToken.ID, userID, 7, "private-stream")
|
|
start := handler.SnapshotIOStreamState()
|
|
unknown := base64.RawURLEncoding.EncodeToString([]byte(strings.Repeat("u", 32)))
|
|
|
|
cases := []struct {
|
|
name string
|
|
path string
|
|
pat string
|
|
body string
|
|
}{
|
|
{name: "malformed cancel", path: agentcompatCapabilityCancelPath, pat: ownerPAT, body: `{"capability":`},
|
|
{name: "oversize cancel", path: agentcompatCapabilityCancelPath, pat: ownerPAT, body: strings.Repeat("x", 513)},
|
|
{name: "duplicate cancel", path: agentcompatCapabilityCancelPath, pat: ownerPAT, body: `{"capability":"x","capability":"y","purpose":"terminal","server_id":7}`},
|
|
{name: "unknown cancel", path: agentcompatCapabilityCancelPath, pat: ownerPAT, body: capabilityAccessJSON(unknown, "terminal", 7, 0)},
|
|
{name: "foreign cancel", path: agentcompatCapabilityCancelPath, pat: foreignPAT, body: capabilityAccessJSON(capability, "terminal", 7, 0)},
|
|
{name: "purpose cancel", path: agentcompatCapabilityCancelPath, pat: ownerPAT, body: capabilityAccessJSON(capability, "file_manager", 7, 0)},
|
|
{name: "server cancel", path: agentcompatCapabilityCancelPath, pat: ownerPAT, body: capabilityAccessJSON(capability, "terminal", 999999, 0)},
|
|
{name: "invalid unregister", path: agentcompatCapabilityUnregisterPath, pat: ownerPAT, body: capabilityAccessJSON("invalid", "terminal", 7, 0)},
|
|
{name: "oversize unregister", path: agentcompatCapabilityUnregisterPath, pat: ownerPAT, body: strings.Repeat("x", 513)},
|
|
{name: "duplicate unregister", path: agentcompatCapabilityUnregisterPath, pat: ownerPAT, body: `{"capability":"x","purpose":"terminal","server_id":7,"server_id":8}`},
|
|
{name: "foreign unregister", path: agentcompatCapabilityUnregisterPath, pat: foreignPAT, body: capabilityAccessJSON(capability, "terminal", 7, 0)},
|
|
{name: "resource unregister", path: agentcompatCapabilityUnregisterPath, pat: ownerPAT, body: capabilityAccessJSON(capability, "terminal", 7, 41)},
|
|
}
|
|
var uniformBody string
|
|
for _, testCase := range cases {
|
|
t.Run(testCase.name, func(t *testing.T) {
|
|
status, body := postAgentcompatRaw(t, server.URL+testCase.path, testCase.pat, testCase.body)
|
|
require.Equal(t, http.StatusOK, status)
|
|
if uniformBody == "" {
|
|
uniformBody = body
|
|
}
|
|
require.Equal(t, uniformBody, body)
|
|
require.Equal(t, start, handler.SnapshotIOStreamState())
|
|
require.NotContains(t, body, capability)
|
|
require.NotContains(t, body, "private-stream")
|
|
})
|
|
}
|
|
|
|
waitContext, cancel := context.WithTimeout(context.Background(), time.Second)
|
|
defer cancel()
|
|
waited, err := postAgentcompatCapability[agentcompatCapabilityWaitRequest, agentcompatCapabilityWaitResponse](waitContext, server.URL+agentcompatCapabilityWaitPath, ownerPAT, agentcompatCapabilityWaitRequest{Capability: capability, Purpose: "terminal", ServerID: 7})
|
|
require.NoError(t, err)
|
|
require.Equal(t, "private-stream", waited.StreamID)
|
|
}
|
|
|
|
func TestAgentcompatCapabilityPermissionRevocationIsRecheckedWithoutMutation(t *testing.T) {
|
|
cleanup, userID := setupMCPTest(t)
|
|
t.Cleanup(cleanup)
|
|
handler := rpc.NewNezhaHandler()
|
|
originalHandler := rpc.NezhaHandlerSingleton
|
|
rpc.NezhaHandlerSingleton = handler
|
|
t.Cleanup(func() { rpc.NezhaHandlerSingleton = originalHandler })
|
|
token, plaintext := mkToken(t, userID, []string{model.ScopeServerRead}, []uint64{7})
|
|
server := newAgentcompatCapabilityServer(t)
|
|
capability := registerAgentcompatCapability(t, server.URL, plaintext, agentcompatCapabilityRegisterRequest{Purpose: "terminal", ServerID: 7})
|
|
bindAgentcompatTerminalCapability(t, handler, capability, token.ID, userID, 7, "revoked-private-stream")
|
|
start := handler.SnapshotIOStreamState()
|
|
token.SetServerIDs([]uint64{99})
|
|
require.NoError(t, singleton.DB.Model(token).Update("servers_csv", token.ServersCSV).Error)
|
|
|
|
cancelStatus, cancelBody := postAgentcompatEmpty(t, server.URL+agentcompatCapabilityCancelPath, plaintext, agentcompatCapabilityCancelRequest{Capability: capability, Purpose: "terminal", ServerID: 7})
|
|
unregisterStatus, unregisterBody := postAgentcompatEmpty(t, server.URL+agentcompatCapabilityUnregisterPath, plaintext, agentcompatCapabilityUnregisterRequest{Capability: capability, Purpose: "terminal", ServerID: 7})
|
|
require.Equal(t, http.StatusOK, cancelStatus)
|
|
require.Equal(t, http.StatusOK, unregisterStatus)
|
|
require.Equal(t, cancelBody, unregisterBody)
|
|
require.Equal(t, start, handler.SnapshotIOStreamState())
|
|
|
|
token.SetServerIDs(nil)
|
|
require.NoError(t, singleton.DB.Model(token).Update("servers_csv", "").Error)
|
|
waitContext, cancel := context.WithTimeout(context.Background(), time.Second)
|
|
defer cancel()
|
|
waited, err := postAgentcompatCapability[agentcompatCapabilityWaitRequest, agentcompatCapabilityWaitResponse](waitContext, server.URL+agentcompatCapabilityWaitPath, plaintext, agentcompatCapabilityWaitRequest{Capability: capability, Purpose: "terminal", ServerID: 7})
|
|
require.NoError(t, err)
|
|
require.Equal(t, "revoked-private-stream", waited.StreamID)
|
|
}
|
|
|
|
func TestAgentcompatCapabilityForeignWaitAndRevokedPATCannotObserveBinding(t *testing.T) {
|
|
cleanup, userID := setupMCPTest(t)
|
|
t.Cleanup(cleanup)
|
|
handler := rpc.NewNezhaHandler()
|
|
originalHandler := rpc.NezhaHandlerSingleton
|
|
rpc.NezhaHandlerSingleton = handler
|
|
t.Cleanup(func() { rpc.NezhaHandlerSingleton = originalHandler })
|
|
ownerToken, ownerPAT := mkToken(t, userID, []string{model.ScopeServerRead}, nil)
|
|
_, foreignPAT := mkDistinctCapabilityToken(t, userID, "foreign-wait")
|
|
server := newAgentcompatCapabilityServer(t)
|
|
capability := registerAgentcompatCapability(t, server.URL, ownerPAT, agentcompatCapabilityRegisterRequest{Purpose: "terminal", ServerID: 7})
|
|
bindAgentcompatTerminalCapability(t, handler, capability, ownerToken.ID, userID, 7, "foreign-private-stream")
|
|
start := handler.SnapshotIOStreamState()
|
|
|
|
foreignContext, cancelForeign := context.WithTimeout(context.Background(), time.Second)
|
|
defer cancelForeign()
|
|
_, foreignErr := postAgentcompatCapability[agentcompatCapabilityWaitRequest, agentcompatCapabilityWaitResponse](foreignContext, server.URL+agentcompatCapabilityWaitPath, foreignPAT, agentcompatCapabilityWaitRequest{Capability: capability, Purpose: "terminal", ServerID: 7})
|
|
require.Error(t, foreignErr)
|
|
require.NotContains(t, foreignErr.Error(), capability)
|
|
require.NotContains(t, foreignErr.Error(), "foreign-private-stream")
|
|
require.Equal(t, start, handler.SnapshotIOStreamState())
|
|
|
|
require.NoError(t, singleton.DB.Delete(&model.APIToken{}, ownerToken.ID).Error)
|
|
status, body := postAgentcompatEmpty(t, server.URL+agentcompatCapabilityWaitPath, ownerPAT, agentcompatCapabilityAccessRequest{Capability: capability, Purpose: "terminal", ServerID: 7})
|
|
require.Equal(t, http.StatusUnauthorized, status)
|
|
require.NotContains(t, body, capability)
|
|
require.NotContains(t, body, "foreign-private-stream")
|
|
require.Equal(t, start, handler.SnapshotIOStreamState())
|
|
}
|
|
|
|
func TestAgentcompatCapabilityRegisterRequiresCurrentServerWhitelist(t *testing.T) {
|
|
cleanup, userID := setupMCPTest(t)
|
|
t.Cleanup(cleanup)
|
|
originalHandler := rpc.NezhaHandlerSingleton
|
|
rpc.NezhaHandlerSingleton = rpc.NewNezhaHandler()
|
|
t.Cleanup(func() { rpc.NezhaHandlerSingleton = originalHandler })
|
|
_, plaintext := mkToken(t, userID, []string{model.ScopeServerRead}, []uint64{99})
|
|
server := newAgentcompatCapabilityServer(t)
|
|
|
|
status, body := postAgentcompatEmpty(t, server.URL+agentcompatCapabilityRegisterPath, plaintext, agentcompatCapabilityRegisterRequest{Purpose: "file_manager", ServerID: 7})
|
|
|
|
require.Equal(t, http.StatusOK, status)
|
|
require.Contains(t, body, errAgentcompatCapabilityUnavailable.Error())
|
|
require.NotContains(t, body, "99")
|
|
require.NotContains(t, body, plaintext)
|
|
}
|
|
|
|
func TestAgentcompatCapabilityRegisterRejectsInvalidIdentityWithoutDisclosure(t *testing.T) {
|
|
cleanup, userID := setupMCPTest(t)
|
|
t.Cleanup(cleanup)
|
|
originalHandler := rpc.NezhaHandlerSingleton
|
|
rpc.NezhaHandlerSingleton = rpc.NewNezhaHandler()
|
|
t.Cleanup(func() { rpc.NezhaHandlerSingleton = originalHandler })
|
|
_, plaintext := mkToken(t, userID, []string{model.ScopeServerRead}, nil)
|
|
server := newAgentcompatCapabilityServer(t)
|
|
privateServer := "777777777"
|
|
cases := []string{
|
|
`{`,
|
|
`{"purpose":"unknown","server_id":7}`,
|
|
`{"purpose":"terminal","server_id":0}`,
|
|
`{"purpose":"terminal","server_id":7,"resource_id":41}`,
|
|
`{"purpose":"nat","server_id":7}`,
|
|
`{"purpose":"terminal","server_id":` + privateServer + `}`,
|
|
}
|
|
for _, body := range cases {
|
|
status, responseBody := postAgentcompatRaw(t, server.URL+agentcompatCapabilityRegisterPath, plaintext, body)
|
|
require.Equal(t, http.StatusOK, status)
|
|
require.NotContains(t, responseBody, privateServer)
|
|
require.NotContains(t, responseBody, plaintext)
|
|
var envelope model.CommonResponse[agentcompatCapabilityRegisterResponse]
|
|
require.NoError(t, json.Unmarshal([]byte(responseBody), &envelope))
|
|
require.False(t, envelope.Success)
|
|
require.Empty(t, envelope.Data.Capability)
|
|
}
|
|
}
|
|
|
|
func newAgentcompatCapabilityServer(t *testing.T) *httptest.Server {
|
|
t.Helper()
|
|
gin.SetMode(gin.TestMode)
|
|
router := gin.New()
|
|
registerAgentcompatRoutes(router)
|
|
server := httptest.NewServer(router)
|
|
t.Cleanup(server.Close)
|
|
return server
|
|
}
|
|
|
|
func mkDistinctCapabilityToken(t *testing.T, userID uint64, suffix string) (*model.APIToken, string) {
|
|
t.Helper()
|
|
plaintext := "nzp_" + strings.Repeat("z", 32) + "_" + suffix
|
|
token := &model.APIToken{UserID: userID, Name: suffix, TokenHash: model.HashAPIToken(plaintext)}
|
|
token.SetScopes([]string{model.ScopeServerRead})
|
|
require.NoError(t, singleton.DB.Create(token).Error)
|
|
return token, plaintext
|
|
}
|
|
|
|
func bindAgentcompatTerminalCapability(t *testing.T, handler *rpc.NezhaHandler, rawCapability string, tokenID, userID, serverID uint64, streamID string) {
|
|
t.Helper()
|
|
capability, err := rpc.ParseAgentCompatIOStreamCapability(rawCapability)
|
|
require.NoError(t, err)
|
|
access := rpc.AgentCompatCapabilityAccess{
|
|
Capability: capability, Owner: rpc.AgentCompatCapabilityOwner{PATID: tokenID, UserID: userID},
|
|
Purpose: rpc.AgentCompatCapabilityTerminal, TargetServerID: serverID, ServerAccessAllowed: true,
|
|
}
|
|
require.NoError(t, handler.CreateStreamWithPurpose(streamID, userID, serverID, rpc.PurposeTerminal))
|
|
require.NoError(t, handler.BindAgentCompatIOStreamCapability(rpc.AgentCompatCapabilityBinding{AgentCompatCapabilityAccess: access, StreamID: streamID}))
|
|
}
|
|
|
|
func capabilityAccessJSON(capability, purpose string, serverID, resourceID uint64) string {
|
|
body, _ := json.Marshal(agentcompatCapabilityAccessRequest{Capability: capability, Purpose: purpose, ServerID: serverID, ResourceID: resourceID})
|
|
return string(body)
|
|
}
|
|
|
|
func postAgentcompatRaw(t *testing.T, path, token, body string) (int, string) {
|
|
t.Helper()
|
|
requestContext, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
|
defer cancel()
|
|
request, err := http.NewRequestWithContext(requestContext, http.MethodPost, path, bytes.NewBufferString(body))
|
|
require.NoError(t, err)
|
|
request.Header.Set("Authorization", "Bearer "+token)
|
|
request.Header.Set("Content-Type", "application/json")
|
|
response, err := http.DefaultClient.Do(request)
|
|
require.NoError(t, err)
|
|
defer response.Body.Close()
|
|
responseBody, err := io.ReadAll(response.Body)
|
|
require.NoError(t, err)
|
|
return response.StatusCode, string(responseBody)
|
|
}
|