mirror of
https://github.com/Buriburizaem0n/nezha_domains.git
synced 2026-09-19 09:40:12 +00:00
naibaandcloudcode
ab25662ddd
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>
97 lines
3.1 KiB
Go
97 lines
3.1 KiB
Go
package controller
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/stretchr/testify/require"
|
|
|
|
"github.com/nezhahq/nezha/model"
|
|
"github.com/nezhahq/nezha/pkg/grpcx"
|
|
pb "github.com/nezhahq/nezha/proto"
|
|
"github.com/nezhahq/nezha/service/rpc"
|
|
"github.com/nezhahq/nezha/service/singleton"
|
|
)
|
|
|
|
func newFakeAgentIO() *grpcx.IOStreamWrapper {
|
|
return grpcx.NewIOStreamWrapper(&fakeAgentStream{closed: make(chan struct{})})
|
|
}
|
|
|
|
// fakeAgentStream mimics an attached-but-silent agent: Recv blocks until the
|
|
// wrapper is closed, exactly the post-attach state where nothing watches the
|
|
// per-transfer context.
|
|
type fakeAgentStream struct {
|
|
closed chan struct{}
|
|
}
|
|
|
|
func (f *fakeAgentStream) Recv() (*pb.IOStreamData, error) {
|
|
<-f.closed
|
|
return nil, context.Canceled
|
|
}
|
|
func (f *fakeAgentStream) Send(*pb.IOStreamData) error { return nil }
|
|
func (f *fakeAgentStream) Context() context.Context { return context.Background() }
|
|
|
|
// 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
|
|
// tear down the attached stream, else a stalled/compromised agent pins a
|
|
// dashboard goroutine + IOStream until restart. openFsTransferStream must wire
|
|
// ctx cancellation to CloseStream.
|
|
func TestOpenFsTransferStream_CancelClosesAttachedStream(t *testing.T) {
|
|
cleanupMCP, _ := setupMCPTest(t)
|
|
defer cleanupMCP()
|
|
singleton.Conf.SetMCPEnabled(true)
|
|
|
|
originalHandler := rpc.NezhaHandlerSingleton
|
|
rpc.NezhaHandlerSingleton = rpc.NewNezhaHandler()
|
|
t.Cleanup(func() { rpc.NezhaHandlerSingleton = originalHandler })
|
|
|
|
stream := newKillSwitchStream()
|
|
sc := singleton.NewEmptyServerClassForTest()
|
|
srv := &model.Server{}
|
|
srv.ID = 7
|
|
srv.SetTaskStream(stream)
|
|
sc.InsertForTest(srv)
|
|
originalShared := singleton.ServerShared
|
|
singleton.ServerShared = sc
|
|
t.Cleanup(func() { singleton.ServerShared = originalShared })
|
|
|
|
streamIDCh := make(chan string, 1)
|
|
go func() {
|
|
task := <-stream.sent
|
|
var req model.FsTransferRequest
|
|
_ = 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)
|
|
}
|
|
}()
|
|
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
streamIO, cleanup, err := openFsTransferStream(ctx, 7, &model.FsTransferRequest{
|
|
Op: model.MCPFsTransferOpDownload,
|
|
Path: "/srv/file",
|
|
})
|
|
require.NoError(t, err, "agent must attach so openFsTransferStream returns a live stream")
|
|
require.NotNil(t, streamIO)
|
|
defer cleanup()
|
|
|
|
streamID := <-streamIDCh
|
|
_, getErr := rpc.NezhaHandlerSingleton.GetStream(streamID)
|
|
require.NoError(t, getErr, "stream must be live before cancel")
|
|
|
|
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")
|
|
}
|