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

Merge branch 'upstream/master' into master and preserve domain extensions

Browse Source
This commit is contained in:
Bot
2026-08-31 02:46:42 +08:00
parent c778472571 1405ba76e5
commit 00f5777112
758 changed files with 89269 additions and 1366 deletions
Download Patch File Download Diff File Expand all files Collapse all files
service/rpc/apply_config_authz_test.go
+565
View File
@@ -0,0 +1,565 @@
package rpc
import (
"context"
"errors"
"strings"
"testing"
"time"
"google.golang.org/grpc/metadata"
"gorm.io/driver/sqlite"
"gorm.io/gorm"
"github.com/nezhahq/nezha/model"
pb "github.com/nezhahq/nezha/proto"
"github.com/nezhahq/nezha/service/singleton"
)
// A malicious or buggy agent owning server A must NOT be able to fail a
// ServerTransfer row belonging to server B by reporting a TaskResult whose
// Id is set to B's transfer ID. The agent-task-result authorization
// invariant (commit 02129f1) requires the dashboard to verify the result's
// addressed object actually belongs to the reporting agent before acting
// on it. Without the cross-check, any compromised agent could cancel/fail
// every in-flight transfer in the system.
func TestRequestTaskApplyConfigIgnoresForeignTransferFailure(t *testing.T) {
// Two distinct servers with different owners. attackerSrv reports the
// failure; victimSrv is the one a pending transfer points at.
attackerSrv := &model.Server{
Common: model.Common{ID: 7, UserID: 100},
UUID: "cccccccc-cccc-cccc-cccc-cccccccccccc",
Name: "attacker",
}
victimSrv := &model.Server{
Common: model.Common{ID: 8, UserID: 200},
UUID: "dddddddd-dddd-dddd-dddd-dddddddddddd",
Name: "victim",
}
users := map[uint64]model.UserInfo{
100: {Role: model.RoleMember},
200: {Role: model.RoleMember},
300: {Role: model.RoleMember, AgentSecret: "to-user-secret"},
}
secrets := map[string]uint64{
"attacker-secret": 100,
"to-user-secret": 300,
}
setupApplyConfigAuthzFixture(t, []*model.Server{attackerSrv, victimSrv}, users, secrets)
// Pending transfer for victimSrv (200 -> 300). attackerSrv is unrelated.
tr := initiateAndRegisterPendingTransfer(t, victimSrv.ID, 200, 300, 1)
// Attacker reports a failed ApplyConfig carrying the victim's transfer ID.
runApplyConfigAuthzResult(t, "attacker-secret", attackerSrv.UUID, &pb.TaskResult{
Id: tr.ID,
Type: model.TaskTypeServerTransferApply,
Successful: false,
Data: "spoofed failure",
})
var refreshed model.ServerTransfer
if err := singleton.DB.First(&refreshed, tr.ID).Error; err != nil {
t.Fatalf("re-read transfer: %v", err)
}
if refreshed.Status != model.ServerTransferStatusPending {
t.Fatalf("foreign-server ApplyConfig failure must leave transfer Pending, got status=%d last_error=%q",
refreshed.Status, refreshed.LastError)
}
var vs model.Server
if err := singleton.DB.First(&vs, victimSrv.ID).Error; err != nil {
t.Fatalf("re-read victim server: %v", err)
}
if vs.UserID != 300 {
t.Fatalf("victim server ownership must remain at ToUserID, got %d", vs.UserID)
}
}
// The legitimate path must still mark the transfer Failed: the reporter is
// the actual transfer subject. This guards against an over-tight ownership
// check that would also break the working flow.
func TestRequestTaskApplyConfigAcceptsOwnTransferFailure(t *testing.T) {
srv := &model.Server{
Common: model.Common{ID: 9, UserID: 200},
UUID: "eeeeeeee-eeee-eeee-eeee-eeeeeeeeeeee",
Name: "subject",
}
users := map[uint64]model.UserInfo{
200: {Role: model.RoleMember, AgentSecret: "from-user-secret"},
300: {Role: model.RoleMember, AgentSecret: "to-user-secret"},
}
secrets := map[string]uint64{
// During Pending the agent still authenticates with the previous
// owner's secret — that's exactly the auth-tolerance window the
// transfer feature exists for.
"from-user-secret": 200,
"to-user-secret": 300,
}
setupApplyConfigAuthzFixture(t, []*model.Server{srv}, users, secrets)
tr := initiateAndRegisterPendingTransfer(t, srv.ID, 200, 300, 1)
runApplyConfigAuthzResult(t, "from-user-secret", srv.UUID, &pb.TaskResult{
Id: tr.ID,
Type: model.TaskTypeServerTransferApply,
Successful: false,
Data: "DisableCommandExecute=true",
})
var refreshed model.ServerTransfer
if err := singleton.DB.First(&refreshed, tr.ID).Error; err != nil {
t.Fatalf("re-read transfer: %v", err)
}
if refreshed.Status != model.ServerTransferStatusFailed {
t.Fatalf("own-server ApplyConfig failure must mark transfer Failed, got status=%d", refreshed.Status)
}
}
func TestRequestTaskCancelledTransferAllowsForwardHandshakeReconnectForRevert(t *testing.T) {
srv := &model.Server{
Common: model.Common{ID: 12, UserID: 200},
UUID: "12121212-1212-1212-1212-121212121212",
Name: "cancelled-revert",
}
users := map[uint64]model.UserInfo{
200: {Role: model.RoleMember, AgentSecret: "cancel-from-secret"},
300: {Role: model.RoleMember, AgentSecret: "cancel-to-secret"},
}
secrets := map[string]uint64{
"cancel-from-secret": 200,
"cancel-to-secret": 300,
}
setupApplyConfigAuthzFixture(t, []*model.Server{srv}, users, secrets)
tr := initiateAndRegisterPendingTransfer(t, srv.ID, 200, 300, 1)
forward := tr.HandshakeSecret
if forward == "" {
t.Fatal("precondition: pending transfer must carry a forward HandshakeSecret")
}
if _, err := singleton.ServerTransferShared.Cancel(tr.ID); err != nil {
t.Fatalf("cancel transfer: %v", err)
}
sent := runApplyConfigAuthzReconnect(t, forward, srv.UUID)
if len(sent) != 1 {
t.Fatalf("expected one revert ApplyConfig task, got %d", len(sent))
}
if sent[0].Type != model.TaskTypeServerTransferApply {
t.Fatalf("expected ApplyConfig task, got type=%d", sent[0].Type)
}
var settled model.ServerTransfer
if err := singleton.DB.First(&settled, tr.ID).Error; err != nil {
t.Fatalf("reload transfer: %v", err)
}
if !strings.Contains(sent[0].Data, settled.RevertHandshakeSecret) {
t.Fatalf("cancelled transfer rollback must push the per-transfer RevertHandshakeSecret, got payload %q", sent[0].Data)
}
if strings.Contains(sent[0].Data, "cancel-from-secret") || strings.Contains(sent[0].Data, "cancel-to-secret") {
t.Fatalf("user-global AgentSecrets must never appear in transfer payloads, got %q", sent[0].Data)
}
}
func TestRequestTaskTimedOutTransferAllowsForwardHandshakeReconnectForRevert(t *testing.T) {
srv := &model.Server{
Common: model.Common{ID: 17, UserID: 200},
UUID: "17171717-1717-1717-1717-171717171717",
Name: "timeout-revert",
}
users := map[uint64]model.UserInfo{
200: {Role: model.RoleMember, AgentSecret: "timeout-from-secret"},
300: {Role: model.RoleMember, AgentSecret: "timeout-to-secret"},
}
secrets := map[string]uint64{
"timeout-from-secret": 200,
"timeout-to-secret": 300,
}
setupApplyConfigAuthzFixture(t, []*model.Server{srv}, users, secrets)
tr := initiateAndRegisterPendingTransfer(t, srv.ID, 200, 300, 1)
forward := tr.HandshakeSecret
if forward == "" {
t.Fatal("precondition: pending transfer must carry a forward HandshakeSecret")
}
staleUpdatedAt := time.Now().Add(-25 * time.Hour)
if err := singleton.DB.Model(&model.ServerTransfer{}).
Where("id = ?", tr.ID).
UpdateColumn("updated_at", staleUpdatedAt).Error; err != nil {
t.Fatalf("stale transfer update: %v", err)
}
if _, err := singleton.ServerTransferShared.MarkTimeout(tr.ID); err != nil {
t.Fatalf("timeout transfer: %v", err)
}
sent := runApplyConfigAuthzReconnect(t, forward, srv.UUID)
if len(sent) != 1 {
t.Fatalf("expected one timeout revert ApplyConfig task, got %d", len(sent))
}
var settled model.ServerTransfer
if err := singleton.DB.First(&settled, tr.ID).Error; err != nil {
t.Fatalf("reload transfer: %v", err)
}
if !strings.Contains(sent[0].Data, settled.RevertHandshakeSecret) {
t.Fatalf("timeout rollback must push the per-transfer RevertHandshakeSecret, got payload %q", sent[0].Data)
}
if strings.Contains(sent[0].Data, "timeout-from-secret") || strings.Contains(sent[0].Data, "timeout-to-secret") {
t.Fatalf("user-global AgentSecrets must never appear in transfer payloads, got %q", sent[0].Data)
}
}
func TestRequestTaskRejectsToUserGlobalSecretEvenWithLiveRevertDelivery(t *testing.T) {
srv := &model.Server{
Common: model.Common{ID: 16, UserID: 200},
UUID: "16161616-1616-1616-1616-161616161616",
Name: "to-user-global-rejected",
}
users := map[uint64]model.UserInfo{
200: {Role: model.RoleMember, AgentSecret: "rejected-from-secret"},
300: {Role: model.RoleMember, AgentSecret: "rejected-to-secret"},
}
secrets := map[string]uint64{
"rejected-from-secret": 200,
"rejected-to-secret": 300,
}
setupApplyConfigAuthzFixture(t, []*model.Server{srv}, users, secrets)
tr := initiateAndRegisterPendingTransfer(t, srv.ID, 200, 300, 1)
if _, err := singleton.ServerTransferShared.Cancel(tr.ID); err != nil {
t.Fatalf("cancel transfer: %v", err)
}
if _, ok := singleton.ServerTransferShared.LookupRevertDelivery(srv.ID); !ok {
t.Fatal("precondition: cancel must register a revert delivery")
}
sent := 0
stream := &requestTaskSecurityStream{
ctx: metadata.NewIncomingContext(context.Background(), metadata.Pairs(
"client_secret", "rejected-to-secret",
"client_uuid", srv.UUID,
)),
onSend: func(*pb.Task) {
sent++
},
}
if err := NewNezhaHandler().RequestTask(stream); err == nil || errors.Is(err, context.Canceled) {
t.Fatal("ToUserID global AgentSecret must never authenticate via revert recovery; PushIfOnline only delivers per-transfer secrets to the real agent")
}
if sent != 0 {
t.Fatalf("rejected ToUserID auth must not trigger any ApplyConfig push, got %d sends", sent)
}
if _, ok := singleton.ServerTransferShared.LookupRevertDelivery(srv.ID); !ok {
t.Fatal("rejected ToUserID auth must not consume the revert delivery — the real agent still needs it for the eventual per-transfer recovery")
}
}
// Whether or not a revert delivery is still in flight, the destination
// user's global AgentSecret must be rejected on every auth path —
// PushIfOnline never sends that secret to the agent so a reconnect under
// it cannot come from the real agent. This pins the post-fix invariant.
func TestReportSystemInfoRejectsCancelledTransferToUserSecret(t *testing.T) {
srv := &model.Server{
Common: model.Common{ID: 15, UserID: 200},
UUID: "15151515-1515-1515-1515-151515151515",
Name: "cancelled-report",
}
users := map[uint64]model.UserInfo{
200: {Role: model.RoleMember, AgentSecret: "report-from-secret"},
300: {Role: model.RoleMember, AgentSecret: "report-to-secret"},
}
secrets := map[string]uint64{
"report-from-secret": 200,
"report-to-secret": 300,
}
setupApplyConfigAuthzFixture(t, []*model.Server{srv}, users, secrets)
tr := initiateAndRegisterPendingTransfer(t, srv.ID, 200, 300, 1)
if _, err := singleton.ServerTransferShared.Cancel(tr.ID); err != nil {
t.Fatalf("cancel transfer: %v", err)
}
ctx := metadata.NewIncomingContext(context.Background(), metadata.Pairs(
"client_secret", "report-to-secret",
"client_uuid", srv.UUID,
))
if _, err := NewNezhaHandler().ReportSystemInfo(ctx, &pb.Host{}); err == nil {
t.Fatal("ReportSystemInfo must reject the destination user's global AgentSecret during revert recovery; PushIfOnline never delivers that credential to the real agent")
}
if _, ok := singleton.ServerTransferShared.LookupRevertDelivery(srv.ID); !ok {
t.Fatal("rejected non-RequestTask auth must not consume the revert delivery")
}
}
func setupApplyConfigAuthzFixture(t *testing.T, servers []*model.Server, users map[uint64]model.UserInfo, agentSecrets map[string]uint64) {
t.Helper()
originalDB := singleton.DB
originalConf := singleton.Conf
originalLoc := singleton.Loc
originalServerShared := singleton.ServerShared
originalUserInfoMap := singleton.UserInfoMap
originalAgentSecretToUserID := singleton.AgentSecretToUserId
originalServerTransferShared := singleton.ServerTransferShared
db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
if err != nil {
t.Fatal(err)
}
sqlDB, err := db.DB()
if err != nil {
t.Fatal(err)
}
sqlDB.SetMaxOpenConns(1)
singleton.DB = db
singleton.Conf = &singleton.ConfigClass{Config: &model.Config{}}
singleton.Loc = time.UTC
if err := singleton.DB.AutoMigrate(model.Server{}, model.ServerTransfer{}); err != nil {
t.Fatal(err)
}
for _, server := range servers {
if err := singleton.DB.Create(server).Error; err != nil {
t.Fatal(err)
}
}
singleton.UserLock.Lock()
singleton.UserInfoMap = users
singleton.AgentSecretToUserId = agentSecrets
singleton.UserLock.Unlock()
singleton.ServerShared = singleton.NewServerClass()
for _, server := range servers {
model.InitServer(server)
singleton.ServerShared.Update(server, server.UUID)
}
singleton.ServerTransferShared = singleton.NewServerTransferClass()
t.Cleanup(func() {
if singleton.ServerTransferShared != nil {
singleton.ServerTransferShared.Stop()
}
sqlDB.Close()
singleton.DB = originalDB
singleton.Conf = originalConf
singleton.Loc = originalLoc
singleton.ServerShared = originalServerShared
singleton.ServerTransferShared = originalServerTransferShared
singleton.UserLock.Lock()
singleton.UserInfoMap = originalUserInfoMap
singleton.AgentSecretToUserId = originalAgentSecretToUserID
singleton.UserLock.Unlock()
})
}
func initiateAndRegisterPendingTransfer(t *testing.T, serverID, fromUserID, toUserID, initiatorID uint64) *model.ServerTransfer {
t.Helper()
var created *model.ServerTransfer
err := singleton.DB.Transaction(func(tx *gorm.DB) error {
var err error
created, err = singleton.ServerTransferShared.Initiate(tx, serverID, fromUserID, toUserID, initiatorID)
return err
})
if err != nil {
t.Fatalf("initiate transfer: %v", err)
}
singleton.ServerTransferShared.Register(created)
return created
}
func runApplyConfigAuthzResult(t *testing.T, secret, uuid string, result *pb.TaskResult) {
t.Helper()
stream := &requestTaskSecurityStream{
ctx: metadata.NewIncomingContext(context.Background(), metadata.Pairs(
"client_secret", secret,
"client_uuid", uuid,
)),
results: []*pb.TaskResult{result},
}
err := NewNezhaHandler().RequestTask(stream)
if !errors.Is(err, context.Canceled) {
t.Fatalf("expected RequestTask to finish after test result, got %v", err)
}
}
func runApplyConfigAuthzReconnect(t *testing.T, secret, uuid string) []*pb.Task {
t.Helper()
var sent []*pb.Task
stream := &requestTaskSecurityStream{
ctx: metadata.NewIncomingContext(context.Background(), metadata.Pairs(
"client_secret", secret,
"client_uuid", uuid,
)),
onSend: func(task *pb.Task) {
sent = append(sent, task)
},
}
err := NewNezhaHandler().RequestTask(stream)
if !errors.Is(err, context.Canceled) {
t.Fatalf("expected RequestTask to finish after reconnect probe, got %v", err)
}
return sent
}
// Finding B regression: during the agent's 10s delayed ApplyConfig swap
// window, the agent still talks to the dashboard with the OLD (FromUserID)
// secret. After a cancel/fail/timeout, a registered revert delivery is the
// only signal that lets the eventually-arriving new-secret reconnect
// recover. The previous implementation cleared revertDeliveries from ANY
// successful old-secret authentication — including ReportSystemInfo2 from
// the periodic reportHost path — so a single old-secret RPC during the
// timer window could destroy the rollback record before the agent ever
// actually swapped secrets. Clearing the delivery is only safe when the
// auth call also gets a chance to consume it by pushing the rollback,
// which only the RequestTask handler does via OnAgentReconnect.
func TestReportSystemInfoDoesNotClearRevertDeliveryForOldSecret(t *testing.T) {
srv := &model.Server{
Common: model.Common{ID: 23, UserID: 200},
UUID: "23232323-2323-2323-2323-232323232323",
Name: "preserve-revert-delivery",
}
users := map[uint64]model.UserInfo{
200: {Role: model.RoleMember, AgentSecret: "from-secret-23"},
300: {Role: model.RoleMember, AgentSecret: "to-secret-23"},
}
secrets := map[string]uint64{
"from-secret-23": 200,
"to-secret-23": 300,
}
setupApplyConfigAuthzFixture(t, []*model.Server{srv}, users, secrets)
tr := initiateAndRegisterPendingTransfer(t, srv.ID, 200, 300, 1)
if _, err := singleton.ServerTransferShared.Cancel(tr.ID); err != nil {
t.Fatalf("cancel transfer: %v", err)
}
if _, ok := singleton.ServerTransferShared.LookupRevertDelivery(srv.ID); !ok {
t.Fatal("precondition: cancel must have registered a revert delivery")
}
// Simulate the agent's periodic reportHost calling ReportSystemInfo2
// with the still-current (FromUserID) secret during the 10s pending
// ApplyConfig window. Must succeed (server already reverted to
// FromUserID) but must NOT clear the revert delivery — the agent has
// not yet swapped secrets, and destroying the only recovery record
// now would lock the agent out once its timer fires.
ctx := metadata.NewIncomingContext(context.Background(), metadata.Pairs(
"client_secret", "from-secret-23",
"client_uuid", srv.UUID,
))
if _, err := NewNezhaHandler().ReportSystemInfo2(ctx, &pb.Host{}); err != nil {
t.Fatalf("ReportSystemInfo2 with old (FromUserID) secret must succeed after revert, got %v", err)
}
if _, ok := singleton.ServerTransferShared.LookupRevertDelivery(srv.ID); !ok {
t.Fatal("non-RequestTask auth with old secret must NOT clear revert delivery; it cannot push the rollback, so destroying the record locks out the eventually-switched agent")
}
}
// Regression: when cancel/fail/timeout happens while the agent is offline
// (its only TaskStream is gone), pushRevertIfOnline is a no-op and the
// revertDelivery is the only signal we have left. The agent will reconnect
// *with the original FromUserID secret* (its in-memory liveCredentials still
// points at the secret it had before the swap), and that very reconnect must
// be the one that delivers the rollback ApplyConfig — otherwise the agent's
// 10s reload timer eventually commits the new secret and the dashboard, which
// already restored ownership to FromUserID, rejects every subsequent connect.
//
// The previous implementation cleared the revertDelivery inside
// authorizeAgentForUUIDWithRevertRecovery *before* RequestTask reached
// OnAgentReconnect, so the rollback push that OnAgentReconnect relies on
// (LookupRevertDelivery → pushRevertIfOnline) found nothing and the agent
// got no rollback at all.
func TestRequestTaskCancelledTransferDeliversRollbackOnOldSecretReconnect(t *testing.T) {
srv := &model.Server{
Common: model.Common{ID: 24, UserID: 200},
UUID: "24242424-2424-2424-2424-242424242424",
Name: "old-secret-rollback",
}
users := map[uint64]model.UserInfo{
200: {Role: model.RoleMember, AgentSecret: "rollback-from-secret"},
300: {Role: model.RoleMember, AgentSecret: "rollback-to-secret"},
}
secrets := map[string]uint64{
"rollback-from-secret": 200,
"rollback-to-secret": 300,
}
setupApplyConfigAuthzFixture(t, []*model.Server{srv}, users, secrets)
tr := initiateAndRegisterPendingTransfer(t, srv.ID, 200, 300, 1)
// Cancel while the agent is offline — the in-memory TaskStream is nil
// (we never attached one), so pushRevertIfOnline silently no-ops.
if _, err := singleton.ServerTransferShared.Cancel(tr.ID); err != nil {
t.Fatalf("cancel transfer: %v", err)
}
if _, ok := singleton.ServerTransferShared.LookupRevertDelivery(srv.ID); !ok {
t.Fatal("precondition: cancel while offline must leave a revert delivery for the eventual reconnect")
}
// Agent now reconnects with its original FromUserID secret (it never
// received the new-secret ApplyConfig because it was offline). This
// RequestTask must deliver the rollback so the agent's reload timer
// supersedes onto the correct credential.
sent := runApplyConfigAuthzReconnect(t, "rollback-from-secret", srv.UUID)
if len(sent) != 1 {
t.Fatalf("expected one rollback ApplyConfig task on old-secret reconnect, got %d", len(sent))
}
var settled model.ServerTransfer
if err := singleton.DB.First(&settled, tr.ID).Error; err != nil {
t.Fatalf("reload transfer: %v", err)
}
if !strings.Contains(sent[0].Data, settled.RevertHandshakeSecret) {
t.Fatalf("old-secret reconnect rollback must carry the per-transfer RevertHandshakeSecret, got %q", sent[0].Data)
}
if strings.Contains(sent[0].Data, "rollback-from-secret") || strings.Contains(sent[0].Data, "rollback-to-secret") {
t.Fatalf("user-global AgentSecrets must never appear in transfer payloads, got %q", sent[0].Data)
}
}
// FORWARD-RECOVERY end-to-end: the exact production scenario the fix
// targets. PushIfOnline only ever delivers t.HandshakeSecret, the agent's
// 10s timer commits it to disk, the operator Cancels in that 10s window
// (revert push misses because the stream had no agent yet, or arrived
// before the forward apply finished). The agent reconnects with the
// forward HandshakeSecret it has on disk. RequestTask MUST accept that
// auth and then deliver one rollback ApplyConfig carrying the per-transfer
// RevertHandshakeSecret so the agent's next reload rotates onto the correct
// credential. Without this, the agent has no path back into the dashboard.
func TestRequestTaskForwardHandshakeSecretReconnectAfterCancelDeliversRollback(t *testing.T) {
srv := &model.Server{
Common: model.Common{ID: 31, UserID: 200},
UUID: "31313131-3131-3131-3131-313131313131",
Name: "forward-recovery",
}
users := map[uint64]model.UserInfo{
200: {Role: model.RoleMember, AgentSecret: "fr-from-secret"},
300: {Role: model.RoleMember, AgentSecret: "fr-to-secret"},
}
secrets := map[string]uint64{
"fr-from-secret": 200,
"fr-to-secret": 300,
}
setupApplyConfigAuthzFixture(t, []*model.Server{srv}, users, secrets)
tr := initiateAndRegisterPendingTransfer(t, srv.ID, 200, 300, 1)
forward := tr.HandshakeSecret
if forward == "" {
t.Fatal("precondition: pending transfer must carry a forward HandshakeSecret")
}
if _, err := singleton.ServerTransferShared.Cancel(tr.ID); err != nil {
t.Fatalf("dashboard Cancel must succeed: %v", err)
}
sent := runApplyConfigAuthzReconnect(t, forward, srv.UUID)
if len(sent) != 1 {
t.Fatalf("forward-secret reconnect after Cancel must deliver one rollback ApplyConfig task, got %d", len(sent))
}
var settled model.ServerTransfer
if err := singleton.DB.First(&settled, tr.ID).Error; err != nil {
t.Fatalf("reload transfer: %v", err)
}
if !strings.Contains(sent[0].Data, settled.RevertHandshakeSecret) {
t.Fatalf("rollback delivered after forward-secret recovery must carry the per-transfer RevertHandshakeSecret, got %q", sent[0].Data)
}
if strings.Contains(sent[0].Data, "fr-from-secret") || strings.Contains(sent[0].Data, "fr-to-secret") {
t.Fatalf("user-global AgentSecrets must never appear in transfer payloads, got %q", sent[0].Data)
}
}
service/rpc/auth.go
+215 -13
View File
@@ -2,6 +2,8 @@ package rpc
import (
"context"
"fmt"
"log"
"strings"
petname "github.com/dustinkirkland/golang-petname"
@@ -20,15 +22,24 @@ type authHandler struct {
}
func (a *authHandler) Check(ctx context.Context) (uint64, error) {
return a.check(ctx)
}
func (a *authHandler) CheckRequestTask(ctx context.Context) (uint64, error) {
return a.check(ctx)
}
// 所有 auth caller 走完全相同的 ServerTransfer dual-secret 容忍策略。
// revertDelivery 不在 auth 阶段消费 —— 真正派发 rollback ApplyConfig 的
// pushRevertIfOnline 才有资格清理它，否则 auth 提前清就会让 OnAgentReconnect
// 找不到 recovery 记录，agent 10s timer 一到就锁死在被拒绝的新 secret 上。
func (a *authHandler) check(ctx context.Context) (uint64, error) {
md, ok := metadata.FromIncomingContext(ctx)
if !ok {
return 0, status.Errorf(codes.Unauthenticated, "获取 metaData 失败")
}
var clientSecret string
if value, ok := md["client_secret"]; ok {
clientSecret = strings.TrimSpace(value[0])
}
clientSecret := firstMetadataValue(md, "client-secret", "client_secret")
if clientSecret == "" {
return 0, status.Error(codes.Unauthenticated, "客户端认证失败")
@@ -36,6 +47,104 @@ func (a *authHandler) Check(ctx context.Context) (uint64, error) {
ip, _ := ctx.Value(model.CtxKeyRealIP{}).(string)
clientUUID := firstMetadataValue(md, "client-uuid", "client_uuid")
if _, err := uuid.ParseUUID(clientUUID); err != nil {
// Keep this counter on the same trigger surface as the
// unknown-secret path below: an attacker who pairs a bad secret
// with a malformed/missing UUID otherwise bypasses
// WAFBlockReasonTypeAgentAuthFail entirely and gets unbounded
// retries (TestAuthBadSecret*InvalidUUIDStillIncrementsAgentAuthFailWAF).
model.BlockIP(singleton.DB, ip, model.WAFBlockReasonTypeAgentAuthFail, model.BlockIDgRPC)
return 0, status.Error(codes.Unauthenticated, "客户端 UUID 不合法")
}
// Per-transfer handshake secret path: ApplyConfig delivers a random
// per-transfer token instead of the destination user's global AgentSecret
// (see PushIfOnline). When the agent reconnects under that token the auth
// layer recognises it here, scoped to the matching server UUID, and
// promotes the transfer to Verified. The user-global secret lookup below
// continues to handle every non-transfer agent, plus the still-tolerated
// previous-owner secret during the Pending window. Checked before the
// global lookup so the handshake-secret token can never collide with
// some other user's accidental match.
if singleton.ServerTransferShared != nil {
if t, ok := singleton.ServerTransferShared.LookupByHandshakeSecret(clientSecret); ok {
cid, found := singleton.ServerShared.UUIDToID(clientUUID)
if !found || cid != t.ServerID {
return 0, status.Error(codes.Unauthenticated, "transfer handshake secret bound to a different server")
}
// Auth via per-transfer HandshakeSecret succeeds only when
// MarkVerified actually performs the Pending → Verified
// transition. A lost CAS (concurrent Cancel/Fail/Timeout)
// means the credential is stale; the verifiedHandshakes
// fallthrough below will still admit it if it had been
// promoted by a successful previous reconnect, otherwise it
// is rejected.
verified, _, err := singleton.ServerTransferShared.MarkVerified(t.ServerID, t.ID)
if err != nil {
log.Printf("NEZHA>> ServerTransfer MarkVerified(cid=%d) via handshake secret failed: %v", t.ServerID, err)
return 0, status.Error(codes.Unauthenticated, "transfer handshake verification failed")
}
if verified {
model.UnblockIP(singleton.DB, ip, model.BlockIDgRPC)
return t.ServerID, nil
}
}
// Bounded terminal-recovery window: a transfer was Cancel/Fail/
// Timeout-ed and the agent may still be presenting either of its
// per-transfer secrets. Single lookup + kind switch:
//
// forward — agent committed t.HandshakeSecret to disk before
// the dashboard observed MarkVerified. Admit so
// RequestTask → OnAgentReconnect can deliver the
// rollback ApplyConfig. DO NOT call MarkVerified
// (transfer is terminal) and DO NOT promote into
// verifiedHandshakes (the agent's stable post-rollback
// credential will be the revert secret, not this one).
//
// revert — agent has applied the rollback and presented
// t.RevertHandshakeSecret. Promote via
// MarkRevertDelivered so the credential survives
// past the recovery window (~24h sweep).
//
// SECURITY: terminalSecretRecovery is only populated by
// revertTransition. A stolen per-transfer secret on a transfer
// whose terminal status was forged in the DB never reaches this
// table — TestAuthHandshakeSecretRejectedAfterTransferTerminated
// pins that path closed.
if t, kind, ok := singleton.ServerTransferShared.LookupByTerminalSecretRecovery(clientSecret); ok {
cid, found := singleton.ServerShared.UUIDToID(clientUUID)
if !found || cid != t.ServerID {
return 0, status.Error(codes.Unauthenticated, "transfer terminal-recovery secret bound to a different server")
}
model.UnblockIP(singleton.DB, ip, model.BlockIDgRPC)
if kind == singleton.TerminalRecoveryRevert {
if err := singleton.ServerTransferShared.MarkRevertDelivered(t.ServerID, t.ID); err != nil {
log.Printf("NEZHA>> ServerTransfer MarkRevertDelivered(server=%d transfer=%d) failed: %v", t.ServerID, t.ID, err)
}
}
return t.ServerID, nil
}
// Post-MarkVerified path: the agent's persisted client_secret is
// the per-transfer HandshakeSecret (PushIfOnline never delivers a
// user-global secret), and no follow-up ApplyConfig swaps it back
// out. So every reconnect after the first one — stream drop, agent
// restart, etc. — must still match this credential, bound strictly
// to (serverID, UUID). The match is constrained to a single server
// because the handshake secret was generated per-transfer; it does
// not unlock any other agent. A new transfer for the same server
// invalidates the entry inside Register, closing this acceptance
// window before the next HandshakeSecret takes over.
if cid, ok := singleton.ServerTransferShared.LookupServerByVerifiedHandshakeSecret(clientSecret); ok {
if uuidCID, found := singleton.ServerShared.UUIDToID(clientUUID); found && uuidCID == cid {
model.UnblockIP(singleton.DB, ip, model.BlockIDgRPC)
return cid, nil
}
return 0, status.Error(codes.Unauthenticated, "transfer verified handshake secret bound to a different server")
}
}
singleton.UserLock.RLock()
userId, ok := singleton.AgentSecretToUserId[clientSecret]
if !ok {
@@ -47,16 +156,10 @@ func (a *authHandler) Check(ctx context.Context) (uint64, error) {
model.UnblockIP(singleton.DB, ip, model.BlockIDgRPC)
var clientUUID string
if value, ok := md["client_uuid"]; ok {
clientUUID = value[0]
clientID, hasID, err := authorizeAgentForUUID(userId, clientUUID)
if err != nil {
return 0, status.Error(codes.Unauthenticated, err.Error())
}
if _, err := uuid.ParseUUID(clientUUID); err != nil {
return 0, status.Error(codes.Unauthenticated, "客户端 UUID 不合法")
}
clientID, hasID := singleton.ServerShared.UUIDToID(clientUUID)
if !hasID {
s := model.Server{UUID: clientUUID, Name: petname.Generate(2, "-"), Common: model.Common{
UserID: userId,
@@ -73,3 +176,102 @@ func (a *authHandler) Check(ctx context.Context) (uint64, error) {
return clientID, nil
}
func firstMetadataValue(md metadata.MD, keys ...string) string {
for _, key := range keys {
if value, ok := md[key]; ok && len(value) > 0 {
return strings.TrimSpace(value[0])
}
}
return ""
}
// authorizeAgentForUUID resolves a client UUID to the dashboard's internal
// server ID, ensuring the resolved server is actually owned by the agent
// secret's owner. Previously Check returned the resolved server ID without
// verifying ownership, allowing an agent that knew another user's server
// UUID to impersonate it (poisoning monitoring state, triggering alerts).
// hasID=false means the UUID is unknown and the caller may register it as
// a new server for the secret owner.
//
// The error path also doubles as a leak-detection signal for operators: if
// an agent persistently fails with "client UUID does not belong to the
// agent secret owner", it pins down which user's secret has been reused
// against a server they don't own.
//
// Server transfer interaction: while a ServerTransfer is Pending for this
// server, the agent is still authenticating with the previous owner's
// AgentSecret (the new secret has not yet propagated). To keep that agent
// online during the rollover, accept userId==FromUserID for the duration of
// the pending window. The dual-secret tolerance is narrowly scoped to the
// affected server only — every other agent of either user is unaffected.
// Once the agent reconnects under the new owner's secret (userId==ToUserID
// matching server.UserID), MarkVerified promotes the transfer and closes
// the tolerance window.
func authorizeAgentForUUID(userId uint64, clientUUID string) (clientID uint64, hasID bool, err error) {
cid, found := singleton.ServerShared.UUIDToID(clientUUID)
if !found {
return 0, false, nil
}
server, _ := singleton.ServerShared.Get(cid)
if server == nil {
// Cache inconsistency: UUID maps to an ID, but no server record exists.
// Treat as unknown (registration path) rather than impersonation.
return 0, false, nil
}
if userId == 0 {
// The legacy global agent secret maps to user 0. It predates per-user
// agent secrets, so keep it compatible by allowing any existing UUID.
// Possession of this deployment-wide master credential is therefore not
// a tenant-scoped authorization claim. Removal must follow an inventory and
// credential-rotation migration or legacy Agents will be locked out.
return cid, true, nil
}
if server.GetUserID() == userId {
// SECURITY: while a transfer is Pending, Server.UserID has already
// been flipped to ToUserID by Register, so userId==Server.UserID
// here also matches the destination user's user-global AgentSecret.
// PushIfOnline only delivers the per-transfer HandshakeSecret on
// the wire; the destination user's global AgentSecret is never
// pushed to the agent, so a reconnect under that secret is not
// proof of agent rotation. Admitting it would let the destination
// user — who can see Server.UUID — authenticate as the agent
// during the Pending window. Reject the user-global secret until
// the transfer settles; the HandshakeSecret path in check() is
// the only valid promotion route.
if singleton.ServerTransferShared != nil {
if _, ok := singleton.ServerTransferShared.LookupPending(cid); ok {
return 0, false, fmt.Errorf("destination user's global AgentSecret cannot authenticate during a pending transfer; agent must rotate to per-transfer HandshakeSecret")
}
}
return cid, true, nil
}
// server.UserID != userId — normally an impersonation attempt. Allow it
// only when a ServerTransfer for this server is Pending AND the secret in
// hand is the previous owner's (FromUserID), OR when a recently terminated
// transfer left a revert-delivery for FromUserID and the agent is still
// presenting its pre-transfer global secret.
//
// SECURITY: we deliberately do NOT accept the destination user's global
// AgentSecret on the LookupRevertDelivery path. PushIfOnline only ever
// delivers per-transfer HandshakeSecret / RevertHandshakeSecret to the
// agent — the ToUserID global secret never travels over the wire — so a
// reconnect under that credential is not proof of agent rotation; it can
// only come from the destination user themselves, who can see Server.UUID
// once Register flips Server.UserID. Admitting it would let that user
// impersonate the agent during the rollback window, trigger
// pushRevertIfOnline to leak RevertHandshakeSecret, and then be promoted
// into verifiedHandshakes via MarkRevertDelivered. The legitimate recovery
// paths are: FromUserID global secret (handled below), forward
// HandshakeSecret and RevertHandshakeSecret (handled by the
// terminalSecretRecovery / verifiedHandshakes lookups in check()).
if singleton.ServerTransferShared != nil {
if t, ok := singleton.ServerTransferShared.LookupRevertDelivery(cid); ok && t.FromUserID == userId {
return cid, true, nil
}
if t, ok := singleton.ServerTransferShared.LookupPending(cid); ok && t.FromUserID == userId {
return cid, true, nil
}
}
return 0, false, fmt.Errorf("client UUID does not belong to the agent secret owner")
}
service/rpc/auth_test.go
+669
View File
@@ -0,0 +1,669 @@
package rpc
import (
"context"
"errors"
"testing"
"google.golang.org/grpc/metadata"
"gorm.io/driver/sqlite"
"gorm.io/gorm"
"github.com/nezhahq/nezha/model"
"github.com/nezhahq/nezha/pkg/utils"
"github.com/nezhahq/nezha/service/singleton"
)
// authCheckWithSecret drives (*authHandler).check end-to-end via the same
// gRPC metadata path the real RPC handler uses. Tests rely on it to assert
// what a real reconnect — secret + UUID supplied on the wire — would do.
func authCheckWithSecret(secret, uuid string) (uint64, error) {
ctx := metadata.NewIncomingContext(context.Background(), metadata.Pairs(
"client_secret", secret,
"client_uuid", uuid,
))
return (&authHandler{}).Check(ctx)
}
func authCheckWithHyphenatedSecret(secret, uuid string) (uint64, error) {
ctx := metadata.NewIncomingContext(context.Background(), metadata.Pairs(
"client-secret", secret,
"client-uuid", uuid,
))
return (&authHandler{}).Check(ctx)
}
// authCheckWithBothKeyStyles reproduces the real post-upgrade wire state: a
// new agent (PR #244) emits BOTH hyphenated and underscore metadata, so a new
// dashboard receives both at once. hyphenSecret/underscoreSecret may differ so
// a test can assert which key wins.
func authCheckWithBothKeyStyles(hyphenSecret, underscoreSecret, uuid string) (uint64, error) {
ctx := metadata.NewIncomingContext(context.Background(), metadata.Pairs(
"client-secret", hyphenSecret,
"client_secret", underscoreSecret,
"client-uuid", uuid,
"client_uuid", uuid,
))
return (&authHandler{}).Check(ctx)
}
// authHandshakeUUID is RFC4122-shaped so it survives the uuid.ParseUUID gate
// at the top of check(); setupAuthAgentFixture's "uuid-alice" / "uuid-bob"
// only work for callers that bypass check() and exercise the inner helpers.
const authHandshakeUUID = "11111111-1111-1111-1111-111111111111"
// setupAuthHandshakeFixture seeds a single server (id=11, owner=user 100,
// real UUID) plus the user-secret tables so the global-secret fall-through
// in check() has something to match. Mirrors setupAuthAgentFixture's reset
// discipline but additionally restores AgentSecretToUserId / UserInfoMap.
func setupAuthHandshakeFixture(t *testing.T) func() {
t.Helper()
originalDB := singleton.DB
originalServerShared := singleton.ServerShared
originalServerTransferShared := singleton.ServerTransferShared
originalUserInfoMap := singleton.UserInfoMap
originalAgentSecretToUserId := singleton.AgentSecretToUserId
db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
if err != nil {
t.Fatalf("open db: %v", err)
}
if err := db.AutoMigrate(&model.Server{}, &model.ServerTransfer{}, &model.WAF{}); err != nil {
t.Fatalf("migrate: %v", err)
}
if err := db.Create(&model.Server{
Common: model.Common{ID: 11, UserID: 100},
UUID: authHandshakeUUID,
Name: "handshake-srv",
}).Error; err != nil {
t.Fatalf("create handshake server: %v", err)
}
singleton.DB = db
singleton.ServerShared = singleton.NewServerClass()
srv := &model.Server{Common: model.Common{ID: 11, UserID: 100}, UUID: authHandshakeUUID, Name: "handshake-srv"}
model.InitServer(srv)
singleton.ServerShared.Update(srv, authHandshakeUUID)
singleton.ServerTransferShared = singleton.NewServerTransferClass()
singleton.UserLock.Lock()
singleton.UserInfoMap = map[uint64]model.UserInfo{
100: {Role: model.RoleMember, AgentSecret: "alice-global"},
200: {Role: model.RoleMember, AgentSecret: "bob-global"},
}
singleton.AgentSecretToUserId = map[string]uint64{
"alice-global": 100,
"bob-global": 200,
}
singleton.UserLock.Unlock()
return func() {
if singleton.ServerTransferShared != nil {
singleton.ServerTransferShared.Stop()
}
singleton.DB = originalDB
singleton.ServerShared = originalServerShared
singleton.ServerTransferShared = originalServerTransferShared
singleton.UserLock.Lock()
singleton.UserInfoMap = originalUserInfoMap
singleton.AgentSecretToUserId = originalAgentSecretToUserId
singleton.UserLock.Unlock()
}
}
func TestAuthCheckAcceptsHyphenatedMetadata(t *testing.T) {
defer setupAuthHandshakeFixture(t)()
cid, err := authCheckWithHyphenatedSecret("alice-global", authHandshakeUUID)
if err != nil {
t.Fatalf("hyphenated metadata must authenticate: %v", err)
}
if cid != 11 {
t.Fatalf("expected server ID 11, got %d", cid)
}
}
// The everyday post-upgrade case (new agent + new dashboard): both key styles
// arrive together carrying the same secret and must authenticate normally.
func TestAuthCheckAcceptsBothKeyStylesPresent(t *testing.T) {
defer setupAuthHandshakeFixture(t)()
cid, err := authCheckWithBothKeyStyles("alice-global", "alice-global", authHandshakeUUID)
if err != nil {
t.Fatalf("an agent emitting both key styles must authenticate: %v", err)
}
if cid != 11 {
t.Fatalf("expected server ID 11, got %d", cid)
}
}
// When both styles are present the hyphenated key wins (firstMetadataValue
// lists it first). This pins the precedence so a future reorder can't silently
// start trusting the underscore alias that Caddy strips.
func TestAuthCheckHyphenatedKeyTakesPrecedence(t *testing.T) {
defer setupAuthHandshakeFixture(t)()
cid, err := authCheckWithBothKeyStyles("alice-global", "garbage-underscore", authHandshakeUUID)
if err != nil {
t.Fatalf("hyphenated secret must be the one used, so auth must succeed: %v", err)
}
if cid != 11 {
t.Fatalf("expected server ID 11 from the hyphenated secret, got %d", cid)
}
}
// setupAuthAgentFixture seeds an in-memory DB and ServerShared with two
// servers belonging to different users so we can assert that a secret bound
// to user A cannot resolve a server UUID owned by user B.
func setupAuthAgentFixture(t *testing.T) func() {
t.Helper()
originalDB := singleton.DB
originalServerShared := singleton.ServerShared
originalServerTransferShared := singleton.ServerTransferShared
db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
if err != nil {
t.Fatalf("open db: %v", err)
}
if err := db.AutoMigrate(&model.Server{}, &model.ServerTransfer{}); err != nil {
t.Fatalf("migrate: %v", err)
}
if err := db.Create(&model.Server{
Common: model.Common{ID: 1, UserID: 100},
UUID: "uuid-alice",
Name: "alice-srv",
}).Error; err != nil {
t.Fatalf("create alice: %v", err)
}
if err := db.Create(&model.Server{
Common: model.Common{ID: 2, UserID: 200},
UUID: "uuid-bob",
Name: "bob-srv",
}).Error; err != nil {
t.Fatalf("create bob: %v", err)
}
singleton.DB = db
singleton.ServerShared = singleton.NewServerClass()
singleton.ServerTransferShared = singleton.NewServerTransferClass()
return func() {
if singleton.ServerTransferShared != nil {
singleton.ServerTransferShared.Stop()
}
singleton.DB = originalDB
singleton.ServerShared = originalServerShared
singleton.ServerTransferShared = originalServerTransferShared
}
}
func TestAuthorizeAgentForUUIDAcceptsOwnedServer(t *testing.T) {
defer setupAuthAgentFixture(t)()
cid, hasID, err := authorizeAgentForUUID(100, "uuid-alice")
if err != nil {
t.Fatalf("alice with her own server UUID must not error, got %v", err)
}
if !hasID || cid != 1 {
t.Fatalf("expected (cid=1, hasID=true), got (cid=%d, hasID=%v)", cid, hasID)
}
}
// Core regression: an agent presenting user A's secret but user B's server
// UUID must be rejected. Previously the code returned the resolved server ID
// without verifying the UserID matched the secret owner, allowing same-tenant
// (and worse — cross-tenant if UUID leaks) server impersonation.
func TestAuthorizeAgentForUUIDRejectsForeignServerUUID(t *testing.T) {
defer setupAuthAgentFixture(t)()
_, _, err := authorizeAgentForUUID(100, "uuid-bob") // alice's secret + bob's UUID
if err == nil {
t.Fatalf("UUID owned by another user must be rejected")
}
}
func TestAuthorizeAgentForUUIDAllowsGlobalDefaultSecret(t *testing.T) {
defer setupAuthAgentFixture(t)()
cid, hasID, err := authorizeAgentForUUID(0, "uuid-bob")
if err != nil {
t.Fatalf("global default secret must be allowed to use existing UUIDs, got %v", err)
}
if !hasID || cid != 2 {
t.Fatalf("expected (cid=2, hasID=true), got (cid=%d, hasID=%v)", cid, hasID)
}
}
// An unknown UUID must NOT be treated as an impersonation attempt — it is
// the normal first-time registration path and the caller (Check) creates a
// new server bound to the secret owner.
func TestAuthorizeAgentForUUIDPermitsUnknownUUIDForRegistration(t *testing.T) {
defer setupAuthAgentFixture(t)()
cid, hasID, err := authorizeAgentForUUID(100, "uuid-never-seen-before")
if err != nil {
t.Fatalf("unknown UUID must be permitted for new registration, got %v", err)
}
if hasID {
t.Fatalf("hasID must be false for unknown UUID, got cid=%d", cid)
}
}
// initiatePendingTransfer mirrors the controller flow used by the batch-move
// endpoint to drive ownership through ServerTransferShared. Tests use it to
// set up the auth-tolerance window with the Server row already flipped to
// ToUserID. Returns nothing; callers use ServerTransferShared.LookupPending
// to fetch the row if they need it.
func initiatePendingTransfer(t *testing.T, serverID, fromUserID, toUserID uint64) {
t.Helper()
var created *model.ServerTransfer
err := singleton.DB.Transaction(func(tx *gorm.DB) error {
var err error
created, err = singleton.ServerTransferShared.Initiate(tx, serverID, fromUserID, toUserID, fromUserID)
return err
})
if err != nil {
t.Fatalf("initiate pending transfer: %v", err)
}
singleton.ServerTransferShared.Register(created)
}
// The auth-tolerance window: while a Pending transfer exists for this server,
// the old owner's AgentSecret must still authenticate this UUID — the agent
// hasn't received the new secret yet via ApplyConfig. Without this, every
// in-flight transfer would knock the affected agent offline immediately.
func TestAuthorizeAgentForUUIDAcceptsFromUserDuringPendingTransfer(t *testing.T) {
defer setupAuthAgentFixture(t)()
// Alice initiates: server 1 moves from alice (100) to bob (200).
// Server.UserID is now 200; alice's agent still presents secret==100.
initiatePendingTransfer(t, 1, 100, 200)
cid, hasID, err := authorizeAgentForUUID(100, "uuid-alice")
if err != nil {
t.Fatalf("FromUserID secret must be accepted during pending window, got %v", err)
}
if !hasID || cid != 1 {
t.Fatalf("expected (cid=1, hasID=true), got (cid=%d, hasID=%v)", cid, hasID)
}
}
// Tolerance is narrowly scoped: an unrelated user's secret must NOT be
// accepted just because *some* transfer is in flight. Specifically, only
// secrets matching FromUserID or ToUserID get through.
func TestAuthorizeAgentForUUIDRejectsThirdPartyDuringPendingTransfer(t *testing.T) {
defer setupAuthAgentFixture(t)()
initiatePendingTransfer(t, 1, 100, 200)
// userId=999 has nothing to do with this transfer.
_, _, err := authorizeAgentForUUID(999, "uuid-alice")
if err == nil {
t.Fatalf("third-party secret must be rejected even while a transfer is pending")
}
}
// SECURITY: during a Pending transfer the destination user's user-global
// AgentSecret must NOT close the pending window. PushIfOnline only delivers
// the per-transfer HandshakeSecret on the wire, so a reconnect under the
// destination user's global AgentSecret is not proof of agent rotation —
// it could just be the destination user authenticating with their own
// secret + the now-visible Server.UUID. Reject it; only the per-transfer
// HandshakeSecret path may promote to Verified.
func TestAuthorizeAgentForUUIDRejectsToUserGlobalSecretDuringPendingTransfer(t *testing.T) {
defer setupAuthAgentFixture(t)()
initiatePendingTransfer(t, 1, 100, 200)
if _, _, err := authorizeAgentForUUID(200, "uuid-alice"); err == nil {
t.Fatal("destination user's global AgentSecret must NOT authenticate during pending transfer; only per-transfer HandshakeSecret may close the window")
}
if !singleton.ServerTransferShared.HasPending(1) {
t.Fatal("pending transfer must survive a destination-user global AgentSecret reconnect")
}
if _, _, err := authorizeAgentForUUID(100, "uuid-alice"); err != nil {
t.Fatalf("FromUser tolerance window must remain open while transfer is still Pending, got %v", err)
}
}
// During the revert recovery window the destination user's global
// AgentSecret must NOT be accepted by authorizeAgentForUUID. PushIfOnline
// only delivers per-transfer HandshakeSecret / RevertHandshakeSecret on
// the wire, so a reconnect under the ToUserID global secret cannot come
// from the real agent — it can only come from the destination user
// themselves, who can see Server.UUID and would otherwise impersonate the
// agent during rollback, trigger pushRevertIfOnline to leak
// RevertHandshakeSecret, and get promoted via MarkRevertDelivered.
// Legitimate recovery goes through FromUserID's global secret, the
// forward HandshakeSecret, or the RevertHandshakeSecret.
func TestAuthorizeAgentForUUIDRejectsToUserGlobalSecretDuringRevertRecovery(t *testing.T) {
defer setupAuthAgentFixture(t)()
initiatePendingTransfer(t, 1, 100, 200)
pending, ok := singleton.ServerTransferShared.LookupPending(1)
if !ok {
t.Fatal("expected pending transfer")
}
if _, err := singleton.ServerTransferShared.Cancel(pending.ID); err != nil {
t.Fatalf("cancel transfer: %v", err)
}
if _, _, err := authorizeAgentForUUID(200, "uuid-alice"); err == nil {
t.Fatal("destination user's global AgentSecret must NOT authenticate during revert recovery; only per-transfer HandshakeSecret / RevertHandshakeSecret may close the window")
}
if _, ok := singleton.ServerTransferShared.LookupRevertDelivery(1); !ok {
t.Fatal("rejected ToUserID auth must not consume the revert delivery — the real agent still needs it for the eventual per-transfer recovery")
}
}
// Regression for finding A: after MarkVerified deletes the pending entry,
// the agent's persisted ClientSecret is still the per-transfer
// HandshakeSecret (PushIfOnline only ever delivered that value). The very
// next reconnect — gRPC stream drop, agent restart, network blip — must
// keep authenticating, otherwise the agent silently locks itself out on
// the now-orphaned handshake token. There is no follow-up ApplyConfig
// path that swaps the agent over to the destination user's stable
// AgentSecret, so auth itself has to keep treating the post-Verified
// HandshakeSecret as a valid credential for that server.
func TestAuthHandshakeSecretStillAuthenticatesAfterMarkVerified(t *testing.T) {
defer setupAuthHandshakeFixture(t)()
initiatePendingTransfer(t, 11, 100, 200)
pending, ok := singleton.ServerTransferShared.LookupPending(11)
if !ok {
t.Fatal("expected pending transfer")
}
handshakeSecret := pending.HandshakeSecret
if handshakeSecret == "" {
t.Fatal("precondition: pending transfer must carry a HandshakeSecret")
}
cid, err := authCheckWithSecret(handshakeSecret, authHandshakeUUID)
if err != nil {
t.Fatalf("first reconnect with HandshakeSecret must promote the transfer, got %v", err)
}
if cid != 11 {
t.Fatalf("first reconnect must resolve to server 11, got %d", cid)
}
if singleton.ServerTransferShared.HasPending(11) {
t.Fatal("MarkVerified must have cleared the pending index after the handshake reconnect")
}
if _, err := authCheckWithSecret(handshakeSecret, authHandshakeUUID); err != nil {
t.Fatalf("second reconnect with the same HandshakeSecret must still authenticate (the agent has no other credential to present until a final hand-off completes); got %v", err)
}
}
// First successful auth with RevertHandshakeSecret proves the agent has
// applied the rollback (10s reload + applyPendingReload have committed
// the secret to disk). At that point the auth path must promote the
// secret into the long-term verifiedHandshakes map and consume the
// temporary revertDeliveries entry — otherwise the only acceptance path
// is LookupByRevertHandshakeSecret, which prunes after
// defaultRevertDeliveryRecoveryWindow and leaves the agent locked out
// ~24h later. See ServerTransferClass.MarkRevertDelivered.
func TestAuthRevertHandshakeSecretPromotesToVerifiedAndKeepsAuthenticating(t *testing.T) {
defer setupAuthHandshakeFixture(t)()
initiatePendingTransfer(t, 11, 100, 200)
pending, ok := singleton.ServerTransferShared.LookupPending(11)
if !ok {
t.Fatal("expected pending transfer")
}
if _, err := singleton.ServerTransferShared.Cancel(pending.ID); err != nil {
t.Fatalf("cancel transfer to register a revert delivery: %v", err)
}
revert, ok := singleton.ServerTransferShared.LookupRevertDelivery(11)
if !ok {
t.Fatal("precondition: cancel must have registered a revert delivery")
}
revertHandshake := revert.RevertHandshakeSecret
if revertHandshake == "" {
t.Fatal("precondition: revert delivery must carry a RevertHandshakeSecret")
}
if _, err := authCheckWithSecret(revertHandshake, authHandshakeUUID); err != nil {
t.Fatalf("first auth with RevertHandshakeSecret must succeed, got %v", err)
}
if _, ok := singleton.ServerTransferShared.LookupRevertDelivery(11); ok {
t.Fatal("first successful auth must consume the temporary revertDelivery — the credential is now promoted to the long-term map")
}
sid, ok := singleton.ServerTransferShared.LookupServerByVerifiedHandshakeSecret(revertHandshake)
if !ok || sid != 11 {
t.Fatalf("RevertHandshakeSecret must be promoted into verifiedHandshakes; lookup got (sid=%d, ok=%v)", sid, ok)
}
if _, err := authCheckWithSecret(revertHandshake, authHandshakeUUID); err != nil {
t.Fatalf("second auth via the promoted verifiedHandshakes path must still succeed, got %v", err)
}
}
// HIGH security regression: if a transfer has already been Cancelled/Failed/
// Timed out, its HandshakeSecret must NEVER authenticate. Today auth.check
// calls MarkVerified on the lookup result and treats RowsAffected==0 as
// success, so an attacker who learned the per-transfer HandshakeSecret
// (e.g. previous owner whose stream was hijacked during Pending) can
// authenticate inside the narrow race window where revertTransition has
// changed DB status but not yet deleted the in-memory pending entry, or
// after that window simply because the swallowed return is `return
// t.ServerID, nil`.
//
// Expected: when the transfer row is no longer Pending, auth must reject
// the HandshakeSecret entirely.
func TestAuthHandshakeSecretRejectedAfterTransferTerminated(t *testing.T) {
defer setupAuthHandshakeFixture(t)()
initiatePendingTransfer(t, 11, 100, 200)
pending, ok := singleton.ServerTransferShared.LookupPending(11)
if !ok {
t.Fatal("expected pending transfer")
}
handshakeSecret := pending.HandshakeSecret
// Settle the DB row to Cancelled WITHOUT touching the in-memory
// pending entry. This reproduces the race window in revertTransition
// between the DB CAS and the c.mu.Lock that deletes the pending
// entry; LookupByHandshakeSecret still hits.
if err := singleton.DB.Model(&model.ServerTransfer{}).
Where("id = ?", pending.ID).
Update("status", model.ServerTransferStatusCancelled).Error; err != nil {
t.Fatalf("simulate concurrent cancel: %v", err)
}
_, err := authCheckWithSecret(handshakeSecret, authHandshakeUUID)
if err == nil {
t.Fatal("HandshakeSecret on a terminated transfer must be rejected — auth swallowed MarkVerified RowsAffected==0 and returned success, enabling auth bypass with a stale per-transfer secret")
}
}
// HIGH security regression: the auth tolerance window for the old owner's
// global AgentSecret must close in lockstep with MarkVerified. Holding c.mu
// across the DB CAS, the c.pending delete and the verifiedHandshakes write
// inside MarkVerified makes those three steps a single observable event for
// any auth-path lookup taking c.mu.RLock; once MarkVerified returns
// verified=true, no later authorizeAgentForUUID can still see the pending
// entry that previously admitted FromUserID.
func TestAuthOldOwnerSecretRejectedOnceTransferIsVerifiedInDB(t *testing.T) {
defer setupAuthHandshakeFixture(t)()
initiatePendingTransfer(t, 11, 100, 200)
pending, ok := singleton.ServerTransferShared.LookupPending(11)
if !ok {
t.Fatal("expected pending transfer")
}
verified, _, err := singleton.ServerTransferShared.MarkVerified(11, pending.ID)
if err != nil {
t.Fatalf("MarkVerified must succeed for a fresh pending: %v", err)
}
if !verified {
t.Fatal("MarkVerified must report verified=true for a fresh pending")
}
if _, _, err := authorizeAgentForUUID(100, authHandshakeUUID); err == nil {
t.Fatal("old owner's global AgentSecret must be rejected once MarkVerified has returned — the auth tolerance window must not outlive the verified transition")
}
}
// FORWARD-RECOVERY (HIGH): symmetric to TestAuthHandshakeSecretRejectedAfter
// TransferTerminated. That test pokes the DB directly to simulate an
// attacker who learned the per-transfer forward HandshakeSecret outside
// of any dashboard-driven cancellation; auth must reject. This test
// exercises the OTHER scenario: a legitimate agent that already wrote the
// forward HandshakeSecret to disk via the 10s reload timer, and the
// dashboard cancels the transfer via the normal Cancel API (which goes
// through revertTransition). The agent's next reconnect presents the
// forward HandshakeSecret. Auth must authenticate it so RequestTask can
// run OnAgentReconnect and push the RevertHandshakeSecret rollback —
// otherwise the agent is permanently locked out and the operator has to
// SSH in and edit the config by hand.
//
// The distinguishing signal is whether revertTransition was the one that
// settled the row: it populates terminalForwardRecovery; a direct DB
// poke does not.
func TestAuthForwardHandshakeSecretAcceptedAfterDashboardCancel(t *testing.T) {
defer setupAuthHandshakeFixture(t)()
initiatePendingTransfer(t, 11, 100, 200)
pending, ok := singleton.ServerTransferShared.LookupPending(11)
if !ok {
t.Fatal("expected pending transfer")
}
forward := pending.HandshakeSecret
if _, err := singleton.ServerTransferShared.Cancel(pending.ID); err != nil {
t.Fatalf("dashboard Cancel must succeed: %v", err)
}
cid, err := authCheckWithSecret(forward, authHandshakeUUID)
if err != nil {
t.Fatalf("forward HandshakeSecret must authenticate after dashboard Cancel so RequestTask can deliver the rollback; got %v", err)
}
if cid != 11 {
t.Fatalf("forward HandshakeSecret must resolve to its bound server, got cid=%d", cid)
}
if _, ok := singleton.ServerTransferShared.LookupServerByVerifiedHandshakeSecret(forward); ok {
t.Fatal("forward HandshakeSecret on a terminated transfer must NOT be promoted into verifiedHandshakes — promotion would outlive the bounded recovery window and turn a cancelled credential into a permanent one")
}
}
// wafAgentAuthFailCount returns the recorded WAF count for the given IP +
// gRPC block identifier. Used by the bad-credential WAF tests to assert
// FirstOrCreate / UPDATE actually fired.
func wafAgentAuthFailCount(t *testing.T, ip string) uint64 {
t.Helper()
bin, err := utils.IPStringToBinary(ip)
if err != nil {
t.Fatalf("ip parse: %v", err)
}
var w model.WAF
res := singleton.DB.Where("ip = ? AND block_identifier = ?", bin, model.BlockIDgRPC).First(&w)
if res.Error != nil {
if errors.Is(res.Error, gorm.ErrRecordNotFound) {
return 0
}
t.Fatalf("query waf: %v", res.Error)
}
return w.Count
}
// authCheckFromIP feeds an attacker IP through the real Check entry point
// so the WAF BlockIP path observes a non-empty CtxKeyRealIP. authCheckWithSecret
// uses a bare context.Background which keeps the IP empty and short-circuits
// BlockIP(ip == ""), masking the very regression these tests want to pin.
func authCheckFromIP(secret, uuid, ip string) (uint64, error) {
ctx := metadata.NewIncomingContext(context.Background(), metadata.Pairs(
"client_secret", secret,
"client_uuid", uuid,
))
ctx = context.WithValue(ctx, model.CtxKeyRealIP{}, ip)
return (&authHandler{}).Check(ctx)
}
// REGRESSION: the new per-transfer handshake path moved the client_uuid
// validation in front of the global AgentSecretToUserId lookup. A bad
// secret paired with a malformed/missing UUID now short-circuits to
// "客户端 UUID 不合法" and skips the BlockIP(WAFBlockReasonTypeAgentAuthFail)
// counter the previous implementation incremented. That counter is the
// only thing throttling brute-force on agent secrets — losing it lets an
// attacker enumerate secrets indefinitely just by also corrupting the
// UUID metadata. Both the missing-secret and bad-secret cases must still
// count toward AgentAuthFail when the UUID is unusable.
func TestAuthBadSecretInvalidUUIDStillIncrementsAgentAuthFailWAF(t *testing.T) {
defer setupAuthHandshakeFixture(t)()
const attackerIP = "203.0.113.7"
if _, err := authCheckFromIP("definitely-not-a-real-secret", "not-a-uuid", attackerIP); err == nil {
t.Fatal("Check must reject bogus credentials")
}
if got := wafAgentAuthFailCount(t, attackerIP); got == 0 {
t.Fatalf("bad client_secret + invalid client_uuid must still count toward WAFBlockReasonTypeAgentAuthFail; got count=%d", got)
}
}
// Mirror of the above for the empty-UUID metadata path. uuid.ParseUUID("")
// also errors out, so the same auth-fail counting must apply — otherwise an
// attacker can just omit the metadata key entirely.
func TestAuthBadSecretEmptyUUIDStillIncrementsAgentAuthFailWAF(t *testing.T) {
defer setupAuthHandshakeFixture(t)()
const attackerIP = "203.0.113.8"
if _, err := authCheckFromIP("another-bad-secret", "", attackerIP); err == nil {
t.Fatal("Check must reject bogus credentials")
}
if got := wafAgentAuthFailCount(t, attackerIP); got == 0 {
t.Fatalf("bad client_secret + empty client_uuid must still count toward WAFBlockReasonTypeAgentAuthFail; got count=%d", got)
}
}
// FORWARD-RECOVERY: forward secret bound to server A must not authenticate
// when presented with server B's UUID. Defence against an attacker who
// learns one server's forward secret and tries to attach it to a different
// agent during the recovery window.
func TestAuthForwardHandshakeSecretRejectedForDifferentUUID(t *testing.T) {
defer setupAuthHandshakeFixture(t)()
initiatePendingTransfer(t, 11, 100, 200)
pending, ok := singleton.ServerTransferShared.LookupPending(11)
if !ok {
t.Fatal("expected pending transfer")
}
forward := pending.HandshakeSecret
if _, err := singleton.ServerTransferShared.Cancel(pending.ID); err != nil {
t.Fatalf("dashboard Cancel must succeed: %v", err)
}
const otherUUID = "22222222-2222-2222-2222-222222222222"
if _, err := authCheckWithSecret(forward, otherUUID); err == nil {
t.Fatal("forward HandshakeSecret must be rejected when paired with a different server UUID even during recovery — token is per-(server, transfer)")
}
}
// SECURITY (P1): PushIfOnline only ever delivers the per-transfer
// HandshakeSecret to the real agent; the destination user's global
// AgentSecret is never sent on the wire and is therefore not proof of
// agent rotation. Server.UUID is visible to the destination user once
// Register flips Server.UserID, so admitting (ToUser global secret, real
// UUID) and calling MarkVerified would let the destination user clear
// the auth tolerance window for FromUser's secret (locking the real
// agent out) and flip transfer state to Verified without the agent ever
// applying the new credential. Only LookupByHandshakeSecret may promote.
func TestAuthDestinationUserGlobalSecretDoesNotVerifyPendingTransfer(t *testing.T) {
defer setupAuthHandshakeFixture(t)()
initiatePendingTransfer(t, 11, 100, 200)
if !singleton.ServerTransferShared.HasPending(11) {
t.Fatal("precondition: pending transfer must be registered")
}
cid, err := authCheckWithSecret("bob-global", authHandshakeUUID)
if err == nil {
t.Fatalf("destination user's global AgentSecret must not close the transfer's pending window; got cid=%d", cid)
}
if !singleton.ServerTransferShared.HasPending(11) {
t.Fatal("pending transfer must survive a destination-user global AgentSecret reconnect; only the per-transfer HandshakeSecret may promote to Verified")
}
}
service/rpc/geoip_rpc.go
+57
View File
@@ -0,0 +1,57 @@
package rpc
import (
"context"
"fmt"
"log"
"net"
"github.com/nezhahq/nezha/model"
geoipx "github.com/nezhahq/nezha/pkg/geoip"
pb "github.com/nezhahq/nezha/proto"
"github.com/nezhahq/nezha/service/singleton"
)
func (s *NezhaHandler) ReportGeoIP(ctx context.Context, report *pb.GeoIP) (*pb.GeoIP, error) {
clientID, err := s.Auth.Check(ctx)
if err != nil {
return nil, err
}
geoIP := model.PB2GeoIP(report)
if geoIP.IP.IPv4Addr == "" && geoIP.IP.IPv6Addr == "" {
ip, _ := ctx.Value(model.CtxKeyRealIP{}).(string)
if ip == "" {
ip, _ = ctx.Value(model.CtxKeyConnectingIP{}).(string)
}
geoIP.IP.IPv4Addr = ip
}
joinedIP := geoIP.IP.Join()
server, ok := singleton.ServerShared.Get(clientID)
if !ok || server == nil {
return nil, fmt.Errorf("server not found")
}
if server.EnableDDNS && joinedIP != "" && (server.GeoIP == nil || server.GeoIP.IP != geoIP.IP) {
if err := singleton.ServerShared.UpdateDDNS(server, &model.IP{IPv4Addr: geoIP.IP.IPv4Addr, IPv6Addr: geoIP.IP.IPv6Addr}); err != nil {
log.Printf("NEZHA>> Failed to update DDNS for server %d: %v", server.ID, err)
}
}
if server.GeoIP != nil && singleton.Conf.EnableIPChangeNotification &&
((singleton.Conf.Cover == model.ConfigCoverAll && !singleton.Conf.IgnoredIPNotificationServerIDs[clientID]) ||
(singleton.Conf.Cover == model.ConfigCoverIgnoreAll && singleton.Conf.IgnoredIPNotificationServerIDs[clientID])) &&
server.GeoIP.IP.Join() != "" && joinedIP != "" && server.GeoIP.IP != geoIP.IP {
singleton.NotificationShared.SendNotification(singleton.Conf.IPChangeNotificationGroupID,
fmt.Sprintf("[%s] %s, %s => %s", singleton.Localizer.T("IP Changed"), server.Name,
singleton.IPDesensitize(server.GeoIP.IP.Join()), singleton.IPDesensitize(joinedIP)), "")
}
ip := geoIP.IP.IPv4Addr
if geoIP.IP.IPv6Addr != "" && (report.GetUse6() || ip == "") {
ip = geoIP.IP.IPv6Addr
}
location, err := geoipx.Lookup(net.ParseIP(ip))
if err != nil {
log.Printf("NEZHA>> geoip.Lookup: %v", err)
}
geoIP.CountryCode = location
server.GeoIP = &geoIP
return &pb.GeoIP{Ip: nil, CountryCode: location, DashboardBootTime: singleton.DashboardBootTime}, nil
}
service/rpc/io_stream.go
+73 -115
View File
@@ -4,97 +4,57 @@ import (
"errors"
"io"
"sync"
"sync/atomic"
"time"
"github.com/nezhahq/nezha/service/singleton"
)
type StreamPurpose uint8
const (
PurposeLegacy StreamPurpose = iota
PurposeMCPTransfer
PurposeTerminal
PurposeFileManager
PurposeNAT
)
type ioStreamContext struct {
creatorUserID uint64
targetServerID uint64
purpose StreamPurpose
userIo io.ReadWriteCloser
agentIo io.ReadWriteCloser
userIoConnectCh chan struct{}
agentIoConnectCh chan struct{}
userIoChOnce sync.Once
agentIoChOnce sync.Once
revokedCh chan struct{}
revokedOnce sync.Once
waitStartedCh chan struct{}
waitStartedOnce sync.Once
startCaptureCh chan struct{}
startCaptureOnce sync.Once
}
type bp struct {
buf []byte
}
var bufPool = sync.Pool{
New: func() any {
return &bp{
buf: make([]byte, 1024*1024),
}
},
}
func (s *NezhaHandler) CreateStream(streamId string) {
s.ioStreamMutex.Lock()
defer s.ioStreamMutex.Unlock()
s.ioStreams[streamId] = &ioStreamContext{
userIoConnectCh: make(chan struct{}),
agentIoConnectCh: make(chan struct{}),
func newIOStreamContext(creatorUserID, targetServerID uint64, purpose StreamPurpose) *ioStreamContext {
return &ioStreamContext{
creatorUserID: creatorUserID, targetServerID: targetServerID, purpose: purpose,
userIoConnectCh: make(chan struct{}), agentIoConnectCh: make(chan struct{}),
revokedCh: make(chan struct{}), waitStartedCh: make(chan struct{}), startCaptureCh: make(chan struct{}),
}
}
func (s *NezhaHandler) GetStream(streamId string) (*ioStreamContext, error) {
s.ioStreamMutex.RLock()
defer s.ioStreamMutex.RUnlock()
if ctx, ok := s.ioStreams[streamId]; ok {
return ctx, nil
}
return nil, errors.New("stream not found")
func (stream *ioStreamContext) revoke() {
stream.revokedOnce.Do(func() { close(stream.revokedCh) })
}
func (s *NezhaHandler) CloseStream(streamId string) error {
s.ioStreamMutex.Lock()
defer s.ioStreamMutex.Unlock()
type bp struct{ buf []byte }
if ctx, ok := s.ioStreams[streamId]; ok {
if ctx.userIo != nil {
ctx.userIo.Close()
}
if ctx.agentIo != nil {
ctx.agentIo.Close()
}
delete(s.ioStreams, streamId)
}
var bufPool = sync.Pool{New: func() any { return &bp{buf: make([]byte, 1024*1024)} }}
return nil
}
func (s *NezhaHandler) UserConnected(streamId string, userIo io.ReadWriteCloser) error {
stream, err := s.GetStream(streamId)
if err != nil {
return err
}
stream.userIo = userIo
stream.userIoChOnce.Do(func() {
close(stream.userIoConnectCh)
})
return nil
}
func (s *NezhaHandler) AgentConnected(streamId string, agentIo io.ReadWriteCloser) error {
stream, err := s.GetStream(streamId)
if err != nil {
return err
}
stream.agentIo = agentIo
stream.agentIoChOnce.Do(func() {
close(stream.agentIoConnectCh)
})
return nil
func isValidIOStreamMagic(data []byte) bool {
return len(data) >= 4 && data[0] == 0xff && data[1] == 0x05 && data[2] == 0xff && data[3] == 0x05
}
func (s *NezhaHandler) StartStream(streamId string, timeout time.Duration) error {
@@ -102,64 +62,62 @@ func (s *NezhaHandler) StartStream(streamId string, timeout time.Duration) error
if err != nil {
return err
}
return s.startStreamContext(streamId, stream, timeout)
}
func (s *NezhaHandler) startStreamContext(streamId string, stream *ioStreamContext, timeout time.Duration) error {
timeoutTimer := time.NewTimer(timeout)
LOOP:
defer timeoutTimer.Stop()
userConnected := stream.userIoConnectCh
agentConnected := stream.agentIoConnectCh
for {
s.ioStreamMutex.RLock()
if current, exists := s.ioStreams[streamId]; !exists || current != stream {
s.ioStreamMutex.RUnlock()
return errors.New("stream revoked")
}
userIo, agentIo := stream.userIo, stream.agentIo
s.ioStreamMutex.RUnlock()
stream.startCaptureOnce.Do(func() { close(stream.startCaptureCh) })
if userIo != nil {
userConnected = nil
}
if agentIo != nil {
agentConnected = nil
}
if userIo != nil && agentIo != nil {
break
}
select {
case <-stream.userIoConnectCh:
if stream.agentIo != nil {
timeoutTimer.Stop()
break LOOP
}
case <-stream.agentIoConnectCh:
if stream.userIo != nil {
timeoutTimer.Stop()
break LOOP
}
case <-time.After(timeout):
break LOOP
case <-userConnected:
userConnected = nil
case <-agentConnected:
agentConnected = nil
case <-stream.revokedCh:
return errors.New("stream revoked")
case <-timeoutTimer.C:
return singleton.Localizer.ErrorT("timeout: stream endpoints not established")
}
time.Sleep(time.Millisecond * 500)
}
if stream.userIo == nil && stream.agentIo == nil {
return singleton.Localizer.ErrorT("timeout: no connection established")
}
if stream.userIo == nil {
return singleton.Localizer.ErrorT("timeout: user connection not established")
}
if stream.agentIo == nil {
return singleton.Localizer.ErrorT("timeout: agent connection not established")
s.ioStreamMutex.RLock()
if current, exists := s.ioStreams[streamId]; !exists || current != stream {
s.ioStreamMutex.RUnlock()
return errors.New("stream revoked")
}
isDone := new(atomic.Bool)
endCh := make(chan struct{})
userIo, agentIo := stream.userIo, stream.agentIo
s.ioStreamMutex.RUnlock()
errCh := make(chan error, 2)
go func() {
bp := bufPool.Get().(*bp)
defer bufPool.Put(bp)
_, innerErr := io.CopyBuffer(stream.userIo, stream.agentIo, bp.buf)
if innerErr != nil {
err = innerErr
}
if isDone.CompareAndSwap(false, true) {
close(endCh)
}
_, copyErr := io.CopyBuffer(userIo, agentIo, bp.buf)
errCh <- copyErr
}()
go func() {
bp := bufPool.Get().(*bp)
defer bufPool.Put(bp)
_, innerErr := io.CopyBuffer(stream.agentIo, stream.userIo, bp.buf)
if innerErr != nil {
err = innerErr
}
if isDone.CompareAndSwap(false, true) {
close(endCh)
}
_, copyErr := io.CopyBuffer(agentIo, userIo, bp.buf)
errCh <- copyErr
}()
<-endCh
return err
return <-errCh
}
service/rpc/io_stream_capability_agentcompat_test.go
+229
View File
@@ -0,0 +1,229 @@
//go:build agentcompat
package rpc
import (
"context"
"encoding/base64"
"errors"
"io"
"sync"
"sync/atomic"
"testing"
"github.com/stretchr/testify/require"
)
func capabilityOwner(patID, userID uint64) AgentCompatCapabilityOwner {
return AgentCompatCapabilityOwner{PATID: patID, UserID: userID, IsAdmin: false}
}
func capabilityRegistration(owner AgentCompatCapabilityOwner, purpose AgentCompatCapabilityPurpose, serverID, resourceID uint64) AgentCompatCapabilityRegistration {
return AgentCompatCapabilityRegistration{
Owner: owner, Purpose: purpose, TargetServerID: serverID, ResourceID: resourceID, ServerAccessAllowed: true,
}
}
func capabilityAccess(capability AgentCompatIOStreamCapability, registration AgentCompatCapabilityRegistration) AgentCompatCapabilityAccess {
return AgentCompatCapabilityAccess{
Capability: capability, Owner: registration.Owner, Purpose: registration.Purpose,
TargetServerID: registration.TargetServerID, ResourceID: registration.ResourceID, ServerAccessAllowed: true,
}
}
func TestAgentCompatCapabilityMintUsesURLSafe256BitTokens(t *testing.T) {
handler := NewNezhaHandler()
registration := capabilityRegistration(capabilityOwner(1, 2), AgentCompatCapabilityTerminal, 3, 0)
capability, err := handler.RegisterAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), registration)
require.NoError(t, err)
raw, err := base64.RawURLEncoding.DecodeString(capability.String())
require.NoError(t, err)
require.Len(t, raw, 32)
parsed, err := ParseAgentCompatIOStreamCapability(capability.String())
require.NoError(t, err)
require.Equal(t, capability, parsed)
_, err = ParseAgentCompatIOStreamCapability("not-a-capability")
require.ErrorIs(t, err, ErrAgentCompatCapabilityHidden)
}
func TestAgentCompatCapabilityMintRetriesActiveAndUsedCollisions(t *testing.T) {
handler := NewNezhaHandler()
first := make([]byte, 32)
second := make([]byte, 32)
third := make([]byte, 32)
first[0], second[0], third[0] = 1, 2, 3
var calls atomic.Int32
handler.setAgentCompatCapabilityTokenSourceForTest(func(destination []byte) error {
switch calls.Add(1) {
case 1, 2, 4:
copy(destination, first)
return nil
case 3:
copy(destination, second)
return nil
default:
copy(destination, third)
return nil
}
})
registration := capabilityRegistration(capabilityOwner(1, 2), AgentCompatCapabilityTerminal, 3, 0)
firstCapability, err := handler.RegisterAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), registration)
require.NoError(t, err)
activeCollisionCapability, err := handler.RegisterAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), registration)
require.NoError(t, err)
require.NotEqual(t, firstCapability, activeCollisionCapability)
require.NoError(t, handler.UnregisterAgentCompatIOStreamCapability(capabilityAccess(firstCapability, registration)))
tombstoneCollisionCapability, err := handler.RegisterAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), registration)
require.NoError(t, err)
require.NotEqual(t, firstCapability, tombstoneCollisionCapability)
require.Equal(t, int32(5), calls.Load())
}
func TestAgentCompatCapabilityRegistrationRequiresServerAccessProof(t *testing.T) {
handler := NewNezhaHandler()
registration := capabilityRegistration(capabilityOwner(1, 2), AgentCompatCapabilityTerminal, 3, 0)
registration.ServerAccessAllowed = false
_, err := handler.RegisterAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), registration)
require.ErrorIs(t, err, ErrAgentCompatCapabilityHidden)
}
func TestAgentCompatCapabilityWaitRequiresExactOwnerAndRetainsBinding(t *testing.T) {
handler := NewNezhaHandler()
registration := capabilityRegistration(capabilityOwner(10, 20), AgentCompatCapabilityTerminal, 30, 0)
capability, err := handler.RegisterAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), registration)
require.NoError(t, err)
require.NoError(t, handler.CreateStreamWithPurpose("terminal-bound", 20, 30, PurposeTerminal))
require.NoError(t, handler.BindAgentCompatIOStreamCapability(AgentCompatCapabilityBinding{
AgentCompatCapabilityAccess: capabilityAccess(capability, registration), StreamID: "terminal-bound",
}))
streamID, err := handler.WaitAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), capabilityAccess(capability, registration))
require.NoError(t, err)
require.Equal(t, "terminal-bound", streamID)
foreign := capabilityAccess(capability, registration)
foreign.Owner.PATID++
_, err = handler.WaitAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), foreign)
require.ErrorIs(t, err, ErrAgentCompatCapabilityHidden)
}
type capabilityCloseEndpoint struct {
handler *NezhaHandler
streamID string
err error
closed atomic.Int32
}
func (endpoint *capabilityCloseEndpoint) Read([]byte) (int, error) { return 0, io.EOF }
func (endpoint *capabilityCloseEndpoint) Write(data []byte) (int, error) { return len(data), nil }
func (endpoint *capabilityCloseEndpoint) Close() error {
endpoint.closed.Add(1)
endpoint.handler.StreamOwnership(endpoint.streamID)
return endpoint.err
}
func TestAgentCompatCapabilityCancelClosesOutsideLockAndJoinsErrors(t *testing.T) {
handler, registration, capability := boundCapabilityFixture(t, AgentCompatCapabilityFileManager, "fm-close", 41)
firstErr := errors.New("user close")
secondErr := errors.New("agent close")
first := &capabilityCloseEndpoint{handler: handler, streamID: "fm-close", err: firstErr}
second := &capabilityCloseEndpoint{handler: handler, streamID: "fm-close", err: secondErr}
require.NoError(t, handler.UserConnected("fm-close", first))
require.NoError(t, handler.AgentConnected("fm-close", second))
err := handler.CancelAgentCompatIOStreamCapability(capabilityAccess(capability, registration))
require.ErrorIs(t, err, firstErr)
require.ErrorIs(t, err, secondErr)
require.Equal(t, int32(1), first.closed.Load())
require.Equal(t, int32(1), second.closed.Load())
}
func boundCapabilityFixture(t *testing.T, purpose AgentCompatCapabilityPurpose, streamID string, serverID uint64) (*NezhaHandler, AgentCompatCapabilityRegistration, AgentCompatIOStreamCapability) {
t.Helper()
handler := NewNezhaHandler()
registration := capabilityRegistration(capabilityOwner(11, 21), purpose, serverID, 0)
capability, err := handler.RegisterAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), registration)
require.NoError(t, err)
require.NoError(t, handler.CreateStreamWithPurpose(streamID, 21, serverID, purpose.streamPurpose()))
require.NoError(t, handler.BindAgentCompatIOStreamCapability(AgentCompatCapabilityBinding{
AgentCompatCapabilityAccess: capabilityAccess(capability, registration), StreamID: streamID,
}))
return handler, registration, capability
}
func TestAgentCompatCapabilityCancelRacingCloseChangesGenerationOnce(t *testing.T) {
handler, registration, capability := boundCapabilityFixture(t, AgentCompatCapabilityTerminal, "race-close", 51)
endpoint := &capabilityCloseEndpoint{handler: handler, streamID: "race-close"}
require.NoError(t, handler.AgentConnected("race-close", endpoint))
start := handler.SnapshotIOStreamState()
ready := make(chan struct{})
raceCtx := agentCompatCapabilityTestContext(t)
var waitGroup sync.WaitGroup
waitGroup.Add(2)
go func() {
defer waitGroup.Done()
select {
case <-ready:
_ = handler.CloseStream("race-close")
case <-raceCtx.Done():
}
}()
go func() {
defer waitGroup.Done()
select {
case <-ready:
_ = handler.CancelAgentCompatIOStreamCapability(capabilityAccess(capability, registration))
case <-raceCtx.Done():
}
}()
close(ready)
raceDone := make(chan struct{})
go func() {
waitGroup.Wait()
close(raceDone)
}()
awaitAgentCompatCapabilitySignal(t, raceDone, "cancel/close race did not complete")
require.NoError(t, raceCtx.Err())
state := handler.SnapshotIOStreamState()
require.Equal(t, start.Generation+1, state.Generation)
require.Equal(t, int32(1), endpoint.closed.Load())
}
func TestAgentCompatCapabilityWaitTimeoutKeepsRegistration(t *testing.T) {
handler := NewNezhaHandler()
registration := capabilityRegistration(capabilityOwner(1, 2), AgentCompatCapabilityTerminal, 3, 0)
capability, err := handler.RegisterAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), registration)
require.NoError(t, err)
ctx, cancel := context.WithCancel(context.Background())
cancel()
_, err = handler.WaitAgentCompatIOStreamCapability(ctx, capabilityAccess(capability, registration))
require.ErrorIs(t, err, context.Canceled)
require.NoError(t, handler.UnregisterAgentCompatIOStreamCapability(capabilityAccess(capability, registration)))
}
func TestAgentCompatCapabilityWaitWakesAfterUnregister(t *testing.T) {
handler := NewNezhaHandler()
registration := capabilityRegistration(capabilityOwner(1, 2), AgentCompatCapabilityTerminal, 3, 0)
capability, err := handler.RegisterAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), registration)
require.NoError(t, err)
result := make(chan error, 1)
started := make(chan struct{})
handler.setAgentCompatCapabilityWaitObserverForTest(func() { close(started) })
waitCtx := agentCompatCapabilityTestContext(t)
go func() {
_, waitErr := handler.WaitAgentCompatIOStreamCapability(waitCtx, capabilityAccess(capability, registration))
result <- waitErr
}()
awaitAgentCompatCapabilitySignal(t, started, "wait observer did not start")
require.NoError(t, handler.UnregisterAgentCompatIOStreamCapability(capabilityAccess(capability, registration)))
require.ErrorIs(t, receiveAgentCompatCapabilityError(t, result, "unregister did not wake waiter"), ErrAgentCompatCapabilityHidden)
}
service/rpc/io_stream_capability_bind_agentcompat.go
+74
View File
@@ -0,0 +1,74 @@
//go:build agentcompat
package rpc
import "context"
func (s *NezhaHandler) BindAgentCompatIOStreamCapability(binding AgentCompatCapabilityBinding) error {
s.ioStreamMutex.Lock()
defer s.ioStreamMutex.Unlock()
registration, allowed := s.agentCompatRegistrationLocked(binding.AgentCompatCapabilityAccess)
if !allowed || registration.phase != agentCompatCapabilityRegistered || registration.registration.Purpose == AgentCompatCapabilityNAT {
return ErrAgentCompatCapabilityHidden
}
stream, exists := s.ioStreams[binding.StreamID]
stored := registration.registration
if !exists || binding.StreamID == "" || stream.creatorUserID != stored.Owner.UserID ||
stream.targetServerID != stored.TargetServerID || stream.purpose != stored.Purpose.streamPurpose() {
return ErrAgentCompatCapabilityHidden
}
if registration.stream != nil {
if registration.stream == stream && registration.streamID == binding.StreamID {
return nil
}
return ErrAgentCompatCapabilityConflict
}
registration.streamID = binding.StreamID
registration.stream = stream
registration.publishLocked()
return nil
}
func (s *NezhaHandler) WaitAgentCompatIOStreamCapability(ctx context.Context, access AgentCompatCapabilityAccess) (string, error) {
for {
s.ioStreamMutex.RLock()
registration, allowed := s.agentCompatRegistrationLocked(access)
if !allowed {
s.ioStreamMutex.RUnlock()
return "", ErrAgentCompatCapabilityHidden
}
if registration.streamID != "" {
streamID := registration.streamID
stream := registration.stream
stored := registration.registration
current, live := s.ioStreams[streamID]
if stored.Purpose == AgentCompatCapabilityNAT && registration.phase == agentCompatCapabilityPublished && stream != nil {
s.ioStreamMutex.RUnlock()
return streamID, nil
}
// A reused StreamID must not turn a retained capability into authority over a replacement stream.
creatorMatches := stream != nil && stream.creatorUserID == stored.Owner.UserID
if stored.Purpose == AgentCompatCapabilityNAT {
creatorMatches = stream != nil && stream.creatorUserID == 0
}
valid := live && current == stream && creatorMatches &&
stream.targetServerID == stored.TargetServerID && stream.purpose == stored.Purpose.streamPurpose()
s.ioStreamMutex.RUnlock()
if !valid {
return "", ErrAgentCompatCapabilityHidden
}
return streamID, nil
}
notify := registration.notify
observer := s.agentCompatCapabilities.waitObserver
s.ioStreamMutex.RUnlock()
if observer != nil {
observer()
}
select {
case <-ctx.Done():
return "", ctx.Err()
case <-notify:
}
}
}
service/rpc/io_stream_capability_binding_agentcompat_test.go
+193
View File
@@ -0,0 +1,193 @@
//go:build agentcompat
package rpc
import (
"errors"
"testing"
"github.com/stretchr/testify/require"
)
func TestAgentCompatCapabilityCancelLostCreateResponseDeletesOnlyExactStream(t *testing.T) {
handler, registration, capability := boundCapabilityFixture(t, AgentCompatCapabilityTerminal, "lost-response", 61)
require.NoError(t, handler.CreateStreamWithPurpose("other-stream", 21, 61, PurposeTerminal))
start := handler.SnapshotIOStreamState()
err := handler.CancelAgentCompatIOStreamCapability(capabilityAccess(capability, registration))
require.NoError(t, err)
state := handler.SnapshotIOStreamState()
require.Equal(t, start.Generation+1, state.Generation)
require.Equal(t, 1, state.Count)
_, found := handler.StreamOwnership("other-stream")
require.True(t, found)
}
func TestAgentCompatCapabilityCancelOneOfConcurrentCapabilitiesKeepsOthers(t *testing.T) {
handler := NewNezhaHandler()
registration := capabilityRegistration(capabilityOwner(12, 22), AgentCompatCapabilityTerminal, 62, 0)
first, err := handler.RegisterAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), registration)
require.NoError(t, err)
second, err := handler.RegisterAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), registration)
require.NoError(t, err)
for streamID, capability := range map[string]AgentCompatIOStreamCapability{"first": first, "second": second} {
require.NoError(t, handler.CreateStreamWithPurpose(streamID, 22, 62, PurposeTerminal))
require.NoError(t, handler.BindAgentCompatIOStreamCapability(AgentCompatCapabilityBinding{
AgentCompatCapabilityAccess: capabilityAccess(capability, registration), StreamID: streamID,
}))
}
require.NoError(t, handler.CancelAgentCompatIOStreamCapability(capabilityAccess(first, registration)))
streamID, err := handler.WaitAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), capabilityAccess(second, registration))
require.NoError(t, err)
require.Equal(t, "second", streamID)
}
func TestAgentCompatCapabilityBindValidatesStoredIdentityAndStream(t *testing.T) {
tests := []struct {
name string
mutateAccess func(*AgentCompatCapabilityAccess)
streamOwner uint64
streamServer uint64
streamPurpose StreamPurpose
}{
{name: "foreign PAT", mutateAccess: func(access *AgentCompatCapabilityAccess) { access.Owner.PATID++ }, streamOwner: 23, streamServer: 63, streamPurpose: PurposeTerminal},
{name: "user mismatch", mutateAccess: func(access *AgentCompatCapabilityAccess) { access.Owner.UserID++ }, streamOwner: 23, streamServer: 63, streamPurpose: PurposeTerminal},
{name: "admin mismatch", mutateAccess: func(access *AgentCompatCapabilityAccess) { access.Owner.IsAdmin = true }, streamOwner: 23, streamServer: 63, streamPurpose: PurposeTerminal},
{name: "purpose mismatch", mutateAccess: func(access *AgentCompatCapabilityAccess) { access.Purpose = AgentCompatCapabilityFileManager }, streamOwner: 23, streamServer: 63, streamPurpose: PurposeTerminal},
{name: "target mismatch", mutateAccess: func(access *AgentCompatCapabilityAccess) { access.TargetServerID++ }, streamOwner: 23, streamServer: 63, streamPurpose: PurposeTerminal},
{name: "resource mismatch", mutateAccess: func(access *AgentCompatCapabilityAccess) { access.ResourceID++ }, streamOwner: 23, streamServer: 63, streamPurpose: PurposeTerminal},
{name: "access denied", mutateAccess: func(access *AgentCompatCapabilityAccess) { access.ServerAccessAllowed = false }, streamOwner: 23, streamServer: 63, streamPurpose: PurposeTerminal},
{name: "stream creator mismatch", mutateAccess: func(*AgentCompatCapabilityAccess) {}, streamOwner: 24, streamServer: 63, streamPurpose: PurposeTerminal},
{name: "stream server mismatch", mutateAccess: func(*AgentCompatCapabilityAccess) {}, streamOwner: 23, streamServer: 64, streamPurpose: PurposeTerminal},
{name: "stream purpose mismatch", mutateAccess: func(*AgentCompatCapabilityAccess) {}, streamOwner: 23, streamServer: 63, streamPurpose: PurposeFileManager},
}
for _, testCase := range tests {
t.Run(testCase.name, func(t *testing.T) {
handler := NewNezhaHandler()
registration := capabilityRegistration(capabilityOwner(13, 23), AgentCompatCapabilityTerminal, 63, 0)
capability, err := handler.RegisterAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), registration)
require.NoError(t, err)
require.NoError(t, handler.CreateStreamWithPurpose("candidate", testCase.streamOwner, testCase.streamServer, testCase.streamPurpose))
access := capabilityAccess(capability, registration)
testCase.mutateAccess(&access)
err = handler.BindAgentCompatIOStreamCapability(AgentCompatCapabilityBinding{AgentCompatCapabilityAccess: access, StreamID: "candidate"})
require.ErrorIs(t, err, ErrAgentCompatCapabilityHidden)
})
}
}
func TestAgentCompatCapabilityBindIsIdempotentButRejectsConflict(t *testing.T) {
handler, registration, capability := boundCapabilityFixture(t, AgentCompatCapabilityTerminal, "original", 64)
binding := AgentCompatCapabilityBinding{AgentCompatCapabilityAccess: capabilityAccess(capability, registration), StreamID: "original"}
require.NoError(t, handler.BindAgentCompatIOStreamCapability(binding))
require.NoError(t, handler.CreateStreamWithPurpose("conflict", 21, 64, PurposeTerminal))
binding.StreamID = "conflict"
err := handler.BindAgentCompatIOStreamCapability(binding)
require.ErrorIs(t, err, ErrAgentCompatCapabilityConflict)
streamID, err := handler.WaitAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), capabilityAccess(capability, registration))
require.NoError(t, err)
require.Equal(t, "original", streamID)
}
func TestAgentCompatCapabilityCancelMismatchOrReplacementDoesNotDetach(t *testing.T) {
t.Run("target mismatch", func(t *testing.T) {
handler, registration, capability := boundCapabilityFixture(t, AgentCompatCapabilityTerminal, "target-mismatch", 65)
start := handler.SnapshotIOStreamState()
access := capabilityAccess(capability, registration)
access.TargetServerID++
err := handler.CancelAgentCompatIOStreamCapability(access)
require.NoError(t, err)
require.Equal(t, start, handler.SnapshotIOStreamState())
})
t.Run("entry replacement", func(t *testing.T) {
handler, registration, capability := boundCapabilityFixture(t, AgentCompatCapabilityTerminal, "replaced", 66)
require.NoError(t, handler.CloseStream("replaced"))
require.NoError(t, handler.CreateStreamWithPurpose("replaced", 21, 66, PurposeTerminal))
start := handler.SnapshotIOStreamState()
err := handler.CancelAgentCompatIOStreamCapability(capabilityAccess(capability, registration))
require.NoError(t, err)
require.Equal(t, start, handler.SnapshotIOStreamState())
})
}
func TestAgentCompatCapabilityCancelIsIdentityHidingIdempotentForAbsentAndUnbound(t *testing.T) {
handler := NewNezhaHandler()
registration := capabilityRegistration(capabilityOwner(14, 24), AgentCompatCapabilityTerminal, 67, 0)
capability, err := handler.RegisterAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), registration)
require.NoError(t, err)
start := handler.SnapshotIOStreamState()
require.NoError(t, handler.CancelAgentCompatIOStreamCapability(capabilityAccess(capability, registration)))
require.NoError(t, handler.CancelAgentCompatIOStreamCapability(capabilityAccess(capability, registration)))
require.NoError(t, handler.CancelAgentCompatIOStreamCapability(AgentCompatCapabilityAccess{}))
require.Equal(t, start, handler.SnapshotIOStreamState())
}
func TestAgentCompatCapabilityCancelAfterNormalCloseIsIdempotent(t *testing.T) {
handler, registration, capability := boundCapabilityFixture(t, AgentCompatCapabilityTerminal, "normally-closed", 69)
require.NoError(t, handler.CloseStream("normally-closed"))
start := handler.SnapshotIOStreamState()
require.NoError(t, handler.CancelAgentCompatIOStreamCapability(capabilityAccess(capability, registration)))
require.NoError(t, handler.CancelAgentCompatIOStreamCapability(capabilityAccess(capability, registration)))
require.Equal(t, start, handler.SnapshotIOStreamState())
}
func TestAgentCompatCapabilityForeignCancelDoesNotMutate(t *testing.T) {
handler, registration, capability := boundCapabilityFixture(t, AgentCompatCapabilityTerminal, "foreign-cancel", 70)
start := handler.SnapshotIOStreamState()
access := capabilityAccess(capability, registration)
access.Owner.PATID++
err := handler.CancelAgentCompatIOStreamCapability(access)
require.NoError(t, err)
require.Equal(t, start, handler.SnapshotIOStreamState())
_, found := handler.StreamOwnership("foreign-cancel")
require.True(t, found)
}
func TestAgentCompatCapabilityUnregisterRejectsBoundLiveStream(t *testing.T) {
handler, registration, capability := boundCapabilityFixture(t, AgentCompatCapabilityTerminal, "bound-unregister", 68)
start := handler.SnapshotIOStreamState()
err := handler.UnregisterAgentCompatIOStreamCapability(capabilityAccess(capability, registration))
require.ErrorIs(t, err, ErrAgentCompatCapabilityBound)
require.Equal(t, start, handler.SnapshotIOStreamState())
}
func TestAgentCompatCapabilityUnregisterRequiresSamePATAndIsIdempotent(t *testing.T) {
handler := NewNezhaHandler()
registration := capabilityRegistration(capabilityOwner(15, 25), AgentCompatCapabilityTerminal, 71, 0)
capability, err := handler.RegisterAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), registration)
require.NoError(t, err)
foreign := capabilityAccess(capability, registration)
foreign.Owner.PATID++
require.NoError(t, handler.UnregisterAgentCompatIOStreamCapability(foreign))
require.NoError(t, handler.UnregisterAgentCompatIOStreamCapability(capabilityAccess(capability, registration)))
require.NoError(t, handler.UnregisterAgentCompatIOStreamCapability(capabilityAccess(capability, registration)))
}
func TestAgentCompatCapabilityTokenSourceErrorIsVisible(t *testing.T) {
handler := NewNezhaHandler()
sourceErr := errors.New("token source failed")
handler.setAgentCompatCapabilityTokenSourceForTest(func([]byte) error { return sourceErr })
_, err := handler.RegisterAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), capabilityRegistration(capabilityOwner(1, 2), AgentCompatCapabilityTerminal, 3, 0))
require.ErrorIs(t, err, sourceErr)
}
service/rpc/io_stream_capability_cancel_agentcompat.go
+99
View File
@@ -0,0 +1,99 @@
//go:build agentcompat
package rpc
import (
"errors"
"io"
)
func (s *NezhaHandler) CancelAgentCompatIOStreamCapability(access AgentCompatCapabilityAccess) error {
s.ioStreamMutex.Lock()
registration, exists := s.agentCompatCapabilities.active[access.Capability.value]
if !exists {
s.ioStreamMutex.Unlock()
return nil
}
if !agentCompatAccessMatches(access, registration) {
s.ioStreamMutex.Unlock()
// Foreign and absent capabilities intentionally share the same inert result to prevent enumeration.
return nil
}
if registration.stream == nil || registration.streamID == "" {
s.removeAgentCompatCapabilityLocked(access.Capability.value, registration)
s.ioStreamMutex.Unlock()
return nil
}
stream := registration.stream
stored := registration.registration
current, live := s.ioStreams[registration.streamID]
if !live {
s.removeAgentCompatCapabilityLocked(access.Capability.value, registration)
s.ioStreamMutex.Unlock()
return nil
}
creatorMatches := stream.creatorUserID == stored.Owner.UserID
if stored.Purpose == AgentCompatCapabilityNAT {
creatorMatches = stream.creatorUserID == 0
}
if !access.ServerAccessAllowed || current != stream || !creatorMatches ||
stream.targetServerID != stored.TargetServerID || stream.purpose != stored.Purpose.streamPurpose() {
s.ioStreamMutex.Unlock()
return nil
}
stream.revoke()
endpoints := make([]io.ReadWriteCloser, 0, 2)
if stream.userIo != nil {
endpoints = append(endpoints, stream.userIo)
}
if stream.agentIo != nil {
endpoints = append(endpoints, stream.agentIo)
}
delete(s.ioStreams, registration.streamID)
s.publishIOStreamStateChangeLocked()
s.removeAgentCompatCapabilityLocked(access.Capability.value, registration)
s.ioStreamMutex.Unlock()
closeErrors := make([]error, 0, len(endpoints))
for _, endpoint := range endpoints {
if err := endpoint.Close(); err != nil {
closeErrors = append(closeErrors, err)
}
}
return errors.Join(closeErrors...)
}
func (s *NezhaHandler) UnregisterAgentCompatIOStreamCapability(access AgentCompatCapabilityAccess) error {
s.ioStreamMutex.Lock()
defer s.ioStreamMutex.Unlock()
registration, exists := s.agentCompatCapabilities.active[access.Capability.value]
if !exists {
return nil
}
if registration.registration.Owner.PATID != access.Owner.PATID || !agentCompatAccessMatches(access, registration) {
return nil
}
if registration.stream != nil {
if current, live := s.ioStreams[registration.streamID]; live && current == registration.stream {
return ErrAgentCompatCapabilityBound
}
}
s.removeAgentCompatCapabilityLocked(access.Capability.value, registration)
return nil
}
func (s *NezhaHandler) removeAgentCompatCapabilityLocked(capability string, registration *agentCompatCapabilityRegistration) {
current, active := s.agentCompatCapabilities.active[capability]
if !active || current != registration {
return
}
delete(s.agentCompatCapabilities.active, capability)
patID := registration.registration.Owner.PATID
remaining := s.agentCompatCapabilities.activeByPAT[patID] - 1
if remaining == 0 {
delete(s.agentCompatCapabilities.activeByPAT, patID)
} else {
s.agentCompatCapabilities.activeByPAT[patID] = remaining
}
registration.publishLocked()
}
service/rpc/io_stream_capability_default.go
+45
View File
@@ -0,0 +1,45 @@
//go:build !agentcompat
package rpc
import "context"
func (*NezhaHandler) RegisterAgentCompatIOStreamCapability(context.Context, AgentCompatCapabilityRegistration) (AgentCompatIOStreamCapability, error) {
return AgentCompatIOStreamCapability{}, ErrAgentCompatCapabilityUnavailable
}
func (*NezhaHandler) BindAgentCompatIOStreamCapability(AgentCompatCapabilityBinding) error {
return ErrAgentCompatCapabilityUnavailable
}
func (*NezhaHandler) ConsumeAgentCompatNATCapability(AgentCompatCapabilityAccess) (AgentCompatNATPublishHandle, error) {
return AgentCompatNATPublishHandle{}, ErrAgentCompatCapabilityUnavailable
}
func (*NezhaHandler) ConsumeAgentCompatNATCapabilityForProfile(string, uint64, uint64) (AgentCompatCapabilityAccess, AgentCompatNATPublishHandle, error) {
return AgentCompatCapabilityAccess{}, AgentCompatNATPublishHandle{}, ErrAgentCompatCapabilityUnavailable
}
func (*NezhaHandler) PublishAgentCompatNATStream(AgentCompatNATPublishHandle, AgentCompatNATPublication) error {
return ErrAgentCompatCapabilityUnavailable
}
func (*NezhaHandler) WaitAgentCompatIOStreamCapability(context.Context, AgentCompatCapabilityAccess) (string, error) {
return "", ErrAgentCompatCapabilityUnavailable
}
func (*NezhaHandler) CancelAgentCompatIOStreamCapability(AgentCompatCapabilityAccess) error {
return nil
}
func (*NezhaHandler) UnregisterAgentCompatIOStreamCapability(AgentCompatCapabilityAccess) error {
return nil
}
func (*NezhaHandler) CreateAgentCompatNATStream(AgentCompatNATPublishHandle, string) (*AgentCompatNATStreamLease, error) {
return nil, ErrAgentCompatCapabilityUnavailable
}
func (*NezhaHandler) CloseAgentCompatNATStreamLease(*AgentCompatNATStreamLease) error {
return nil
}
service/rpc/io_stream_capability_default_test.go
+64
View File
@@ -0,0 +1,64 @@
//go:build !agentcompat
package rpc
import (
"context"
"errors"
"reflect"
"testing"
"github.com/stretchr/testify/require"
)
func TestAgentCompatCapabilityDefaultBuildHasNoRegistryState(t *testing.T) {
handler := NewNezhaHandler()
state := reflect.ValueOf(handler.agentCompatCapabilities)
require.Equal(t, 0, state.NumField())
}
func TestAgentCompatCapabilityDefaultBuildUsesStableUnavailableAndNoopContracts(t *testing.T) {
handler := NewNezhaHandler()
registration := AgentCompatCapabilityRegistration{}
capability, err := handler.RegisterAgentCompatIOStreamCapability(context.Background(), registration)
require.Empty(t, capability.String())
require.ErrorIs(t, err, ErrAgentCompatCapabilityUnavailable)
access := AgentCompatCapabilityAccess{}
require.ErrorIs(t, handler.BindAgentCompatIOStreamCapability(AgentCompatCapabilityBinding{}), ErrAgentCompatCapabilityUnavailable)
_, err = handler.ConsumeAgentCompatNATCapability(access)
require.ErrorIs(t, err, ErrAgentCompatCapabilityUnavailable)
require.ErrorIs(t, handler.PublishAgentCompatNATStream(AgentCompatNATPublishHandle{}, AgentCompatNATPublication{}), ErrAgentCompatCapabilityUnavailable)
_, err = handler.WaitAgentCompatIOStreamCapability(context.Background(), access)
require.True(t, errors.Is(err, ErrAgentCompatCapabilityUnavailable))
require.NoError(t, handler.CancelAgentCompatIOStreamCapability(access))
require.NoError(t, handler.UnregisterAgentCompatIOStreamCapability(access))
}
func TestAgentCompatNATCapabilityForProfileDefaultBuildIsUnavailableAndNoOp(t *testing.T) {
handler := NewNezhaHandler()
access, handle, err := handler.ConsumeAgentCompatNATCapabilityForProfile("not-a-capability", 1, 2)
require.Equal(t, AgentCompatCapabilityAccess{}, access)
require.Equal(t, AgentCompatNATPublishHandle{}, handle)
require.ErrorIs(t, err, ErrAgentCompatCapabilityUnavailable)
}
func TestAgentCompatNATAtomicStartDefaultBuildIsUnavailableAndStateless(t *testing.T) {
handler := NewNezhaHandler()
publicationOwned, err := handler.StartAgentCompatNATStream(AgentCompatNATPublishHandle{}, 0)
require.False(t, publicationOwned)
require.ErrorIs(t, err, ErrAgentCompatCapabilityUnavailable)
require.Equal(t, 0, handler.StreamCount())
}
func TestAgentCompatNATLeaseDefaultBuildIsUnavailableAndStateless(t *testing.T) {
handler := NewNezhaHandler()
lease, err := handler.CreateAgentCompatNATStream(AgentCompatNATPublishHandle{}, "known")
require.Nil(t, lease)
require.ErrorIs(t, err, ErrAgentCompatCapabilityUnavailable)
require.NoError(t, handler.CloseAgentCompatNATStreamLease(nil))
require.Equal(t, 0, handler.StreamCount())
}
service/rpc/io_stream_capability_nat_agentcompat.go
+84
View File
@@ -0,0 +1,84 @@
//go:build agentcompat
package rpc
func (s *NezhaHandler) ConsumeAgentCompatNATCapability(access AgentCompatCapabilityAccess) (AgentCompatNATPublishHandle, error) {
s.ioStreamMutex.Lock()
defer s.ioStreamMutex.Unlock()
registration, allowed := s.agentCompatRegistrationLocked(access)
if !allowed {
return AgentCompatNATPublishHandle{}, ErrAgentCompatCapabilityHidden
}
return s.consumeAgentCompatNATCapabilityLocked(registration, access.Capability.value)
}
func (s *NezhaHandler) ConsumeAgentCompatNATCapabilityForProfile(value string, targetServerID, resourceID uint64) (AgentCompatCapabilityAccess, AgentCompatNATPublishHandle, error) {
capability, err := ParseAgentCompatIOStreamCapability(value)
if err != nil {
return AgentCompatCapabilityAccess{}, AgentCompatNATPublishHandle{}, ErrAgentCompatCapabilityHidden
}
s.ioStreamMutex.Lock()
defer s.ioStreamMutex.Unlock()
registration, exists := s.agentCompatCapabilities.active[capability.value]
if !exists || registration.registration.Purpose != AgentCompatCapabilityNAT ||
registration.registration.TargetServerID != targetServerID || registration.registration.ResourceID != resourceID {
return AgentCompatCapabilityAccess{}, AgentCompatNATPublishHandle{}, ErrAgentCompatCapabilityHidden
}
handle, err := s.consumeAgentCompatNATCapabilityLocked(registration, capability.value)
if err != nil {
return AgentCompatCapabilityAccess{}, AgentCompatNATPublishHandle{}, err
}
stored := registration.registration
return AgentCompatCapabilityAccess{
Capability: capability, Owner: stored.Owner, Purpose: stored.Purpose,
TargetServerID: stored.TargetServerID, ResourceID: stored.ResourceID,
ServerAccessAllowed: stored.ServerAccessAllowed,
}, handle, nil
}
func (s *NezhaHandler) consumeAgentCompatNATCapabilityLocked(registration *agentCompatCapabilityRegistration, capability string) (AgentCompatNATPublishHandle, error) {
if registration == nil || registration.registration.Purpose != AgentCompatCapabilityNAT || registration.phase != agentCompatCapabilityRegistered {
return AgentCompatNATPublishHandle{}, ErrAgentCompatCapabilityHidden
}
registration.phase = agentCompatCapabilityConsumed
return AgentCompatNATPublishHandle{
registration: registration, generation: registration.generation,
capability: capability,
}, nil
}
func (s *NezhaHandler) PublishAgentCompatNATStream(handle AgentCompatNATPublishHandle, publication AgentCompatNATPublication) error {
s.ioStreamMutex.RLock()
publishObserver := s.agentCompatCapabilities.publishObserver
s.ioStreamMutex.RUnlock()
if publishObserver != nil {
publishObserver()
}
s.ioStreamMutex.Lock()
defer s.ioStreamMutex.Unlock()
registration := handle.registration
// Pointer identity plus generation makes a late publisher inert after unregister/cancel.
if registration == nil || registration.generation != handle.generation {
return nil
}
current, active := s.agentCompatCapabilities.active[handle.capability]
if !active || current != registration {
return nil
}
if registration.phase == agentCompatCapabilityPublished {
return nil
}
stored := registration.registration
stream := registration.stream
exists := publication.StreamID != "" && registration.streamID == publication.StreamID && stream != nil && s.ioStreams[publication.StreamID] == stream
if registration.phase != agentCompatCapabilityConsumed || publication.Purpose != stored.Purpose ||
publication.TargetServerID != stored.TargetServerID || publication.ResourceID != stored.ResourceID ||
!exists || stream.creatorUserID != 0 ||
stream.targetServerID != stored.TargetServerID || stream.purpose != PurposeNAT {
return ErrAgentCompatCapabilityHidden
}
registration.phase = agentCompatCapabilityPublished
registration.publishLocked()
return nil
}
service/rpc/io_stream_capability_nat_agentcompat_test.go
+169
View File
@@ -0,0 +1,169 @@
//go:build agentcompat
package rpc
import (
"testing"
"github.com/stretchr/testify/require"
)
func natCapabilityFixture(t *testing.T, patID, userID, serverID, profileID uint64) (*NezhaHandler, AgentCompatCapabilityRegistration, AgentCompatIOStreamCapability) {
t.Helper()
handler := NewNezhaHandler()
registration := capabilityRegistration(capabilityOwner(patID, userID), AgentCompatCapabilityNAT, serverID, profileID)
capability, err := handler.RegisterAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), registration)
require.NoError(t, err)
return handler, registration, capability
}
func TestAgentCompatNATCapabilityTransitionsAndRetainsFirstPublication(t *testing.T) {
handler, registration, capability := natCapabilityFixture(t, 21, 31, 71, 81)
handle, err := handler.ConsumeAgentCompatNATCapability(capabilityAccess(capability, registration))
require.NoError(t, err)
_, err = handler.ConsumeAgentCompatNATCapability(capabilityAccess(capability, registration))
require.ErrorIs(t, err, ErrAgentCompatCapabilityHidden)
_, err = handler.CreateAgentCompatNATStream(handle, "nat-first")
require.NoError(t, err)
require.NoError(t, handler.CreateStreamWithPurpose("nat-second", 0, 71, PurposeNAT))
publication := AgentCompatNATPublication{Purpose: AgentCompatCapabilityNAT, TargetServerID: 71, ResourceID: 81, StreamID: "nat-first"}
require.NoError(t, handler.PublishAgentCompatNATStream(handle, publication))
publication.StreamID = "nat-second"
require.NoError(t, handler.PublishAgentCompatNATStream(handle, publication))
streamID, err := handler.WaitAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), capabilityAccess(capability, registration))
require.NoError(t, err)
require.Equal(t, "nat-first", streamID)
}
func TestAgentCompatNATCapabilityPublicationBeforeWaitWorks(t *testing.T) {
handler, registration, capability := natCapabilityFixture(t, 22, 32, 72, 82)
handle, err := handler.ConsumeAgentCompatNATCapability(capabilityAccess(capability, registration))
require.NoError(t, err)
_, err = handler.CreateAgentCompatNATStream(handle, "nat-published")
require.NoError(t, err)
require.NoError(t, handler.PublishAgentCompatNATStream(handle, AgentCompatNATPublication{
Purpose: AgentCompatCapabilityNAT, TargetServerID: 72, ResourceID: 82, StreamID: "nat-published",
}))
streamID, err := handler.WaitAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), capabilityAccess(capability, registration))
require.NoError(t, err)
require.Equal(t, "nat-published", streamID)
}
func TestAgentCompatNATCapabilityValidatesConsumeAndPublishIdentity(t *testing.T) {
handler, registration, capability := natCapabilityFixture(t, 23, 33, 73, 83)
access := capabilityAccess(capability, registration)
access.ResourceID++
_, err := handler.ConsumeAgentCompatNATCapability(access)
require.ErrorIs(t, err, ErrAgentCompatCapabilityHidden)
handle, err := handler.ConsumeAgentCompatNATCapability(capabilityAccess(capability, registration))
require.NoError(t, err)
_, err = handler.CreateAgentCompatNATStream(handle, "nat-identity")
require.NoError(t, err)
err = handler.PublishAgentCompatNATStream(handle, AgentCompatNATPublication{
Purpose: AgentCompatCapabilityNAT, TargetServerID: 73, ResourceID: 84, StreamID: "nat-identity",
})
require.ErrorIs(t, err, ErrAgentCompatCapabilityHidden)
}
func TestAgentCompatNATCapabilityLatePublishAfterUnregisterIsIgnored(t *testing.T) {
handler, registration, capability := natCapabilityFixture(t, 24, 34, 74, 84)
handle, err := handler.ConsumeAgentCompatNATCapability(capabilityAccess(capability, registration))
require.NoError(t, err)
require.NoError(t, handler.UnregisterAgentCompatIOStreamCapability(capabilityAccess(capability, registration)))
require.NoError(t, handler.CreateStreamWithPurpose("nat-late", 0, 74, PurposeNAT))
err = handler.PublishAgentCompatNATStream(handle, AgentCompatNATPublication{
Purpose: AgentCompatCapabilityNAT, TargetServerID: 74, ResourceID: 84, StreamID: "nat-late",
})
require.NoError(t, err)
_, err = handler.WaitAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), capabilityAccess(capability, registration))
require.ErrorIs(t, err, ErrAgentCompatCapabilityHidden)
}
func TestAgentCompatNATCapabilityLatePublishAfterCancelIsIgnored(t *testing.T) {
handler, registration, capability := natCapabilityFixture(t, 27, 37, 77, 87)
handle, err := handler.ConsumeAgentCompatNATCapability(capabilityAccess(capability, registration))
require.NoError(t, err)
require.NoError(t, handler.CancelAgentCompatIOStreamCapability(capabilityAccess(capability, registration)))
require.NoError(t, handler.CreateStreamWithPurpose("nat-after-cancel", 0, 77, PurposeNAT))
err = handler.PublishAgentCompatNATStream(handle, AgentCompatNATPublication{
Purpose: AgentCompatCapabilityNAT, TargetServerID: 77, ResourceID: 87, StreamID: "nat-after-cancel",
})
require.NoError(t, err)
_, found := handler.StreamOwnership("nat-after-cancel")
require.True(t, found)
}
func TestAgentCompatNATCapabilityReusedTokenCannotBindAnotherStream(t *testing.T) {
handler, registration, capability := natCapabilityFixture(t, 28, 38, 78, 88)
handle, err := handler.ConsumeAgentCompatNATCapability(capabilityAccess(capability, registration))
require.NoError(t, err)
_, err = handler.CreateAgentCompatNATStream(handle, "nat-original")
require.NoError(t, err)
require.NoError(t, handler.PublishAgentCompatNATStream(handle, AgentCompatNATPublication{
Purpose: AgentCompatCapabilityNAT, TargetServerID: 78, ResourceID: 88, StreamID: "nat-original",
}))
require.NoError(t, handler.CancelAgentCompatIOStreamCapability(capabilityAccess(capability, registration)))
require.NoError(t, handler.CreateStreamWithPurpose("nat-reuse", 0, 78, PurposeNAT))
_, err = handler.ConsumeAgentCompatNATCapability(capabilityAccess(capability, registration))
require.ErrorIs(t, err, ErrAgentCompatCapabilityHidden)
require.NoError(t, handler.PublishAgentCompatNATStream(handle, AgentCompatNATPublication{
Purpose: AgentCompatCapabilityNAT, TargetServerID: 78, ResourceID: 88, StreamID: "nat-reuse",
}))
_, found := handler.StreamOwnership("nat-reuse")
require.True(t, found)
}
func TestAgentCompatNATCapabilityCancelDetachesPublishedStream(t *testing.T) {
handler, registration, capability := natCapabilityFixture(t, 25, 35, 75, 85)
handle, err := handler.ConsumeAgentCompatNATCapability(capabilityAccess(capability, registration))
require.NoError(t, err)
_, err = handler.CreateAgentCompatNATStream(handle, "nat-cancel")
require.NoError(t, err)
require.NoError(t, handler.PublishAgentCompatNATStream(handle, AgentCompatNATPublication{
Purpose: AgentCompatCapabilityNAT, TargetServerID: 75, ResourceID: 85, StreamID: "nat-cancel",
}))
start := handler.SnapshotIOStreamState()
require.NoError(t, handler.CancelAgentCompatIOStreamCapability(capabilityAccess(capability, registration)))
state := handler.SnapshotIOStreamState()
require.Equal(t, start.Generation+1, state.Generation)
require.Equal(t, 0, state.Count)
}
func TestAgentCompatNATCapabilitiesRemainSeparatedAcrossProfiles(t *testing.T) {
handler := NewNezhaHandler()
owner := capabilityOwner(26, 36)
firstRegistration := capabilityRegistration(owner, AgentCompatCapabilityNAT, 76, 86)
secondRegistration := capabilityRegistration(owner, AgentCompatCapabilityNAT, 76, 87)
first, err := handler.RegisterAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), firstRegistration)
require.NoError(t, err)
second, err := handler.RegisterAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), secondRegistration)
require.NoError(t, err)
firstHandle, err := handler.ConsumeAgentCompatNATCapability(capabilityAccess(first, firstRegistration))
require.NoError(t, err)
secondHandle, err := handler.ConsumeAgentCompatNATCapability(capabilityAccess(second, secondRegistration))
require.NoError(t, err)
_, err = handler.CreateAgentCompatNATStream(firstHandle, "nat-profile-first")
require.NoError(t, err)
_, err = handler.CreateAgentCompatNATStream(secondHandle, "nat-profile-second")
require.NoError(t, err)
require.NoError(t, handler.PublishAgentCompatNATStream(firstHandle, AgentCompatNATPublication{Purpose: AgentCompatCapabilityNAT, TargetServerID: 76, ResourceID: 86, StreamID: "nat-profile-first"}))
require.NoError(t, handler.PublishAgentCompatNATStream(secondHandle, AgentCompatNATPublication{Purpose: AgentCompatCapabilityNAT, TargetServerID: 76, ResourceID: 87, StreamID: "nat-profile-second"}))
firstStream, err := handler.WaitAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), capabilityAccess(first, firstRegistration))
require.NoError(t, err)
secondStream, err := handler.WaitAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), capabilityAccess(second, secondRegistration))
require.NoError(t, err)
require.Equal(t, "nat-profile-first", firstStream)
require.Equal(t, "nat-profile-second", secondStream)
}
service/rpc/io_stream_capability_nat_atomic_start_agentcompat_test.go
+180
View File
@@ -0,0 +1,180 @@
//go:build agentcompat
package rpc
import (
"bytes"
"io"
"sync"
"sync/atomic"
"testing"
"time"
"github.com/stretchr/testify/require"
)
type atomicNATEndpoint struct {
handler *NezhaHandler
data *bytes.Reader
written bytes.Buffer
mu sync.Mutex
closed atomic.Int32
readSeen atomic.Int32
writeSeen chan struct{}
}
func (endpoint *atomicNATEndpoint) Read(data []byte) (int, error) {
endpoint.readSeen.Add(1)
endpoint.handler.SnapshotIOStreamState()
return endpoint.data.Read(data)
}
func (endpoint *atomicNATEndpoint) Write(data []byte) (int, error) {
endpoint.mu.Lock()
defer endpoint.mu.Unlock()
endpoint.handler.SnapshotIOStreamState()
n, err := endpoint.written.Write(data)
if endpoint.writeSeen != nil {
select {
case <-endpoint.writeSeen:
default:
close(endpoint.writeSeen)
}
}
return n, err
}
func (endpoint *atomicNATEndpoint) Close() error {
endpoint.closed.Add(1)
endpoint.handler.SnapshotIOStreamState()
return nil
}
func newAtomicNATEndpoint(handler *NezhaHandler, payload string) *atomicNATEndpoint {
return &atomicNATEndpoint{handler: handler, data: bytes.NewReader([]byte(payload)), writeSeen: make(chan struct{})}
}
func publishAtomicNATStream(t *testing.T, streamID string) (*NezhaHandler, AgentCompatCapabilityAccess, AgentCompatNATPublishHandle, AgentCompatCapabilityRegistration) {
t.Helper()
handler, registration, capability := natCapabilityFixture(t, 301, 302, 303, 304)
access := capabilityAccess(capability, registration)
handle, err := handler.ConsumeAgentCompatNATCapability(access)
require.NoError(t, err)
_, err = handler.CreateAgentCompatNATStream(handle, streamID)
require.NoError(t, err)
require.NoError(t, handler.PublishAgentCompatNATStream(handle, AgentCompatNATPublication{
Purpose: AgentCompatCapabilityNAT, TargetServerID: 303, ResourceID: 304, StreamID: streamID,
}))
return handler, access, handle, registration
}
func TestAgentCompatNATAtomicStartWhenCanceledBeforeCaptureDoesNotTouchReplacement(t *testing.T) {
handler, access, handle, _ := publishAtomicNATStream(t, "atomic-replacement-before-capture")
require.NoError(t, handler.CancelAgentCompatIOStreamCapability(access))
require.NoError(t, handler.CreateStreamWithPurpose("atomic-replacement-before-capture", 0, 303, PurposeNAT))
replacement := newAtomicNATEndpoint(handler, "replacement")
require.NoError(t, handler.UserConnected("atomic-replacement-before-capture", replacement))
require.NoError(t, handler.AgentConnected("atomic-replacement-before-capture", replacement))
publicationOwned, err := handler.StartAgentCompatNATStream(handle, time.Millisecond)
require.True(t, publicationOwned)
require.ErrorIs(t, err, ErrAgentCompatCapabilityHidden)
require.Equal(t, int32(0), replacement.readSeen.Load())
require.Equal(t, int32(0), replacement.closed.Load())
_, found := handler.StreamOwnership("atomic-replacement-before-capture")
require.True(t, found)
t.Logf("replacement after cancel-before-capture: read=%d write=%d close=%d registered=%t", replacement.readSeen.Load(), replacement.written.Len(), replacement.closed.Load(), found)
}
func TestAgentCompatNATAtomicStartWhenCanceledAfterCaptureDoesNotCloseReplacement(t *testing.T) {
handler, access, handle, _ := publishAtomicNATStream(t, "atomic-replacement-after-capture")
old := newAtomicNATEndpoint(handler, "old")
require.NoError(t, handler.UserConnected("atomic-replacement-after-capture", old))
result := make(chan error, 1)
go func() {
_, err := handler.StartAgentCompatNATStream(handle, time.Second)
result <- err
}()
stream := mustGetStream(t, handler, "atomic-replacement-after-capture")
select {
case <-stream.startCaptureCh:
case <-time.After(time.Second):
t.Fatal("atomic start did not capture retained stream")
}
require.NoError(t, handler.CancelAgentCompatIOStreamCapability(access))
require.NoError(t, handler.CreateStreamWithPurpose("atomic-replacement-after-capture", 0, 303, PurposeNAT))
replacement := newAtomicNATEndpoint(handler, "replacement")
require.NoError(t, handler.UserConnected("atomic-replacement-after-capture", replacement))
require.NoError(t, handler.AgentConnected("atomic-replacement-after-capture", replacement))
require.EqualError(t, receiveAtomicNATError(t, result), "stream revoked")
require.Equal(t, int32(1), old.closed.Load())
require.Equal(t, int32(0), replacement.readSeen.Load())
require.Equal(t, int32(0), replacement.closed.Load())
_, found := handler.StreamOwnership("atomic-replacement-after-capture")
require.True(t, found)
t.Logf("replacement after cancel-after-capture: read=%d write=%d close=%d registered=%t", replacement.readSeen.Load(), replacement.written.Len(), replacement.closed.Load(), found)
}
func TestAgentCompatNATAtomicStartDetachesOnlyRetainedStreamAfterRelay(t *testing.T) {
handler, _, handle, registration := publishAtomicNATStream(t, "atomic-normal-completion")
user := newAtomicNATEndpoint(handler, "request-bytes")
agent := newAtomicNATEndpoint(handler, "")
require.NoError(t, handler.UserConnected("atomic-normal-completion", user))
require.NoError(t, handler.AgentConnected("atomic-normal-completion", agent))
result := make(chan error, 1)
var publicationOwned bool
go func() {
var err error
publicationOwned, err = handler.StartAgentCompatNATStream(handle, time.Second)
result <- err
}()
select {
case <-agent.writeSeen:
case <-time.After(time.Second):
t.Fatal("atomic relay did not transfer request bytes")
}
err := receiveAtomicNATError(t, result)
require.True(t, publicationOwned)
require.NoError(t, err)
require.Equal(t, int32(1), user.closed.Load())
require.Equal(t, int32(1), agent.closed.Load())
require.Equal(t, "request-bytes", agent.written.String())
streamID, waitErr := handler.WaitAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), capabilityAccessFromRegistration(handle, registration))
require.NoError(t, waitErr)
require.Equal(t, "atomic-normal-completion", streamID)
require.NoError(t, handler.CreateStreamWithPurpose("atomic-normal-completion", 0, 303, PurposeNAT))
replacement := newAtomicNATEndpoint(handler, "replacement")
require.NoError(t, handler.UserConnected("atomic-normal-completion", replacement))
require.NoError(t, handler.AgentConnected("atomic-normal-completion", replacement))
require.Equal(t, int32(0), replacement.readSeen.Load())
require.Equal(t, int32(0), replacement.closed.Load())
t.Logf("replacement after normal retained teardown: read=%d write=%d close=%d registered=true", replacement.readSeen.Load(), replacement.written.Len(), replacement.closed.Load())
}
func capabilityAccessFromRegistration(handle AgentCompatNATPublishHandle, registration AgentCompatCapabilityRegistration) AgentCompatCapabilityAccess {
return AgentCompatCapabilityAccess{Capability: AgentCompatIOStreamCapability{value: handle.capability}, Owner: registration.Owner, Purpose: registration.Purpose, TargetServerID: registration.TargetServerID, ResourceID: registration.ResourceID, ServerAccessAllowed: registration.ServerAccessAllowed}
}
func mustGetStream(t *testing.T, handler *NezhaHandler, streamID string) *ioStreamContext {
t.Helper()
stream, err := handler.GetStream(streamID)
require.NoError(t, err)
return stream
}
func receiveAtomicNATError(t *testing.T, result <-chan error) error {
t.Helper()
select {
case err := <-result:
return err
case <-time.After(time.Second):
t.Fatal("atomic start did not return")
return nil
}
}
var _ io.ReadWriteCloser = (*atomicNATEndpoint)(nil)
service/rpc/io_stream_capability_nat_barrier_agentcompat_test.go
+117
View File
@@ -0,0 +1,117 @@
//go:build agentcompat
package rpc
import (
"sync"
"testing"
"github.com/stretchr/testify/require"
)
func TestAgentCompatNATCapabilityUnregisterBarrierMakesQueuedPublishInert(t *testing.T) {
handler, registration, capability := natCapabilityFixture(t, 38, 48, 60, 70)
handle, err := handler.ConsumeAgentCompatNATCapability(capabilityAccess(capability, registration))
require.NoError(t, err)
_, err = handler.CreateAgentCompatNATStream(handle, "nat-barrier")
require.NoError(t, err)
require.NoError(t, handler.detachExactStream("nat-barrier", handle.registration.stream))
stateBeforeRace := handler.SnapshotIOStreamState()
publishEntered := make(chan struct{})
publishRelease := make(chan struct{})
publishObserverCtx := agentCompatCapabilityTestContext(t)
var observeOnce sync.Once
handler.setAgentCompatCapabilityPublishObserverForTest(func() {
observeOnce.Do(func() {
close(publishEntered)
select {
case <-publishRelease:
case <-publishObserverCtx.Done():
}
})
})
t.Cleanup(func() { handler.setAgentCompatCapabilityPublishObserverForTest(nil) })
publishResult := make(chan error, 1)
go func() {
publishResult <- handler.PublishAgentCompatNATStream(handle, AgentCompatNATPublication{
Purpose: AgentCompatCapabilityNAT, TargetServerID: 60, ResourceID: 70, StreamID: "nat-barrier",
})
}()
awaitAgentCompatCapabilitySignal(t, publishEntered, "publish did not enter production path before unregister")
require.NoError(t, handler.UnregisterAgentCompatIOStreamCapability(capabilityAccess(capability, registration)))
close(publishRelease)
require.NoError(t, receiveAgentCompatCapabilityError(t, publishResult, "queued publish did not return after release"))
_, err = handler.WaitAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), capabilityAccess(capability, registration))
require.ErrorIs(t, err, ErrAgentCompatCapabilityHidden)
require.Equal(t, stateBeforeRace, handler.SnapshotIOStreamState())
_, found := handler.StreamOwnership("nat-barrier")
require.False(t, found)
handler.ioStreamMutex.RLock()
_, active := handler.agentCompatCapabilities.active[capability.value]
retainedStreamID := handle.registration.streamID
retainedStream := handle.registration.stream
handler.ioStreamMutex.RUnlock()
require.False(t, active)
require.Equal(t, "nat-barrier", retainedStreamID)
require.NotNil(t, retainedStream)
}
func TestAgentCompatNATCapabilityPublishObserverCanReenterRegistry(t *testing.T) {
handler, registration, capability := natCapabilityFixture(t, 39, 49, 61, 71)
handle, err := handler.ConsumeAgentCompatNATCapability(capabilityAccess(capability, registration))
require.NoError(t, err)
_, err = handler.CreateAgentCompatNATStream(handle, "nat-observer-reentry")
require.NoError(t, err)
observerEntered := make(chan struct{})
var observeOnce sync.Once
handler.setAgentCompatCapabilityPublishObserverForTest(func() {
handler.SnapshotIOStreamState()
observeOnce.Do(func() { close(observerEntered) })
})
t.Cleanup(func() { handler.setAgentCompatCapabilityPublishObserverForTest(nil) })
publishResult := make(chan error, 1)
go func() {
publishResult <- handler.PublishAgentCompatNATStream(handle, AgentCompatNATPublication{
Purpose: AgentCompatCapabilityNAT, TargetServerID: 61, ResourceID: 71, StreamID: "nat-observer-reentry",
})
}()
awaitAgentCompatCapabilitySignal(t, observerEntered, "publish observer did not reenter registry")
require.NoError(t, receiveAgentCompatCapabilityError(t, publishResult, "publish observer reentry deadlocked"))
streamID, err := handler.WaitAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), capabilityAccess(capability, registration))
require.NoError(t, err)
require.Equal(t, "nat-observer-reentry", streamID)
}
func TestAgentCompatNATCapabilityPublishObserverIsHandlerScoped(t *testing.T) {
first, firstRegistration, firstCapability := natCapabilityFixture(t, 40, 50, 62, 72)
second, secondRegistration, secondCapability := natCapabilityFixture(t, 41, 51, 63, 73)
firstHandle, err := first.ConsumeAgentCompatNATCapability(capabilityAccess(firstCapability, firstRegistration))
require.NoError(t, err)
secondHandle, err := second.ConsumeAgentCompatNATCapability(capabilityAccess(secondCapability, secondRegistration))
require.NoError(t, err)
_, err = first.CreateAgentCompatNATStream(firstHandle, "nat-scoped-first")
require.NoError(t, err)
_, err = second.CreateAgentCompatNATStream(secondHandle, "nat-scoped-second")
require.NoError(t, err)
firstObserved := make(chan struct{})
secondObserved := make(chan struct{})
first.setAgentCompatCapabilityPublishObserverForTest(func() { close(firstObserved) })
second.setAgentCompatCapabilityPublishObserverForTest(func() { close(secondObserved) })
require.NoError(t, first.PublishAgentCompatNATStream(firstHandle, AgentCompatNATPublication{
Purpose: AgentCompatCapabilityNAT, TargetServerID: 62, ResourceID: 72, StreamID: "nat-scoped-first",
}))
awaitAgentCompatCapabilitySignal(t, firstObserved, "first handler observer did not run")
select {
case <-secondObserved:
t.Fatal("second handler observer ran for first handler publish")
default:
}
require.NoError(t, second.PublishAgentCompatNATStream(secondHandle, AgentCompatNATPublication{
Purpose: AgentCompatCapabilityNAT, TargetServerID: 63, ResourceID: 73, StreamID: "nat-scoped-second",
}))
awaitAgentCompatCapabilitySignal(t, secondObserved, "second handler observer did not run")
}
service/rpc/io_stream_capability_nat_handle_binding_agentcompat_test.go
+75
View File
@@ -0,0 +1,75 @@
//go:build agentcompat
package rpc
import (
"testing"
"github.com/stretchr/testify/require"
)
func TestAgentCompatNATHandleCreationBindsEachHandleToItsOwnStream(t *testing.T) {
handler := NewNezhaHandler()
firstRegistration := capabilityRegistration(capabilityOwner(501, 502), AgentCompatCapabilityNAT, 503, 504)
secondRegistration := capabilityRegistration(capabilityOwner(505, 506), AgentCompatCapabilityNAT, 503, 507)
firstCapability, err := handler.RegisterAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), firstRegistration)
require.NoError(t, err)
secondCapability, err := handler.RegisterAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), secondRegistration)
require.NoError(t, err)
firstHandle, err := handler.ConsumeAgentCompatNATCapability(capabilityAccess(firstCapability, firstRegistration))
require.NoError(t, err)
secondHandle, err := handler.ConsumeAgentCompatNATCapability(capabilityAccess(secondCapability, secondRegistration))
require.NoError(t, err)
_, err = handler.CreateAgentCompatNATStream(firstHandle, "handle-bound-first")
require.NoError(t, err)
_, err = handler.CreateAgentCompatNATStream(secondHandle, "handle-bound-second")
require.NoError(t, err)
beforeCrossPublish := snapshotAgentCompatNATCreationState(handler, firstRegistration.Owner.PATID, secondRegistration.Owner.PATID)
require.ErrorIs(t, handler.PublishAgentCompatNATStream(firstHandle, AgentCompatNATPublication{
Purpose: AgentCompatCapabilityNAT, TargetServerID: 503, ResourceID: 504, StreamID: "handle-bound-second",
}), ErrAgentCompatCapabilityHidden)
require.ErrorIs(t, handler.PublishAgentCompatNATStream(secondHandle, AgentCompatNATPublication{
Purpose: AgentCompatCapabilityNAT, TargetServerID: 503, ResourceID: 507, StreamID: "handle-bound-first",
}), ErrAgentCompatCapabilityHidden)
requireUnchangedAgentCompatNATCreationState(t, handler, beforeCrossPublish)
requireNATHandleBindingsIntact(t, handler, firstHandle, "handle-bound-first")
requireNATHandleBindingsIntact(t, handler, secondHandle, "handle-bound-second")
require.NoError(t, handler.PublishAgentCompatNATStream(firstHandle, AgentCompatNATPublication{
Purpose: AgentCompatCapabilityNAT, TargetServerID: 503, ResourceID: 504, StreamID: "handle-bound-first",
}))
require.NoError(t, handler.PublishAgentCompatNATStream(secondHandle, AgentCompatNATPublication{
Purpose: AgentCompatCapabilityNAT, TargetServerID: 503, ResourceID: 507, StreamID: "handle-bound-second",
}))
firstLease, err := handler.CreateAgentCompatNATStream(firstHandle, "handle-bound-again")
require.Nil(t, firstLease)
require.ErrorIs(t, err, ErrAgentCompatCapabilityHidden)
require.Equal(t, 2, handler.StreamCount())
}
func TestAgentCompatNATHandleCreationCancelReleasesOnlyItsBoundStream(t *testing.T) {
handler := NewNezhaHandler()
firstRegistration := capabilityRegistration(capabilityOwner(508, 509), AgentCompatCapabilityNAT, 510, 511)
secondRegistration := capabilityRegistration(capabilityOwner(512, 513), AgentCompatCapabilityNAT, 510, 514)
firstCapability, err := handler.RegisterAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), firstRegistration)
require.NoError(t, err)
secondCapability, err := handler.RegisterAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), secondRegistration)
require.NoError(t, err)
firstAccess := capabilityAccess(firstCapability, firstRegistration)
secondAccess := capabilityAccess(secondCapability, secondRegistration)
firstHandle, err := handler.ConsumeAgentCompatNATCapability(firstAccess)
require.NoError(t, err)
secondHandle, err := handler.ConsumeAgentCompatNATCapability(secondAccess)
require.NoError(t, err)
_, err = handler.CreateAgentCompatNATStream(firstHandle, "bound-first")
require.NoError(t, err)
_, err = handler.CreateAgentCompatNATStream(secondHandle, "bound-second")
require.NoError(t, err)
require.NoError(t, handler.CancelAgentCompatIOStreamCapability(firstAccess))
_, firstFound := handler.StreamOwnership("bound-first")
_, secondFound := handler.StreamOwnership("bound-second")
require.False(t, firstFound)
require.True(t, secondFound)
require.NoError(t, handler.CancelAgentCompatIOStreamCapability(secondAccess))
}
service/rpc/io_stream_capability_nat_lease_authority_agentcompat_test.go
+202
View File
@@ -0,0 +1,202 @@
//go:build agentcompat
package rpc
import (
"errors"
"strconv"
"testing"
"github.com/stretchr/testify/require"
)
type agentCompatNATCreationState struct {
streamState IOStreamState
active int
used int
patState map[uint64]agentCompatNATPATState
}
type agentCompatNATPATState struct {
active uint16
exists bool
}
func snapshotAgentCompatNATCreationState(handler *NezhaHandler, patIDs ...uint64) agentCompatNATCreationState {
active, used := agentCompatCapabilityRegistryCounts(handler)
patState := make(map[uint64]agentCompatNATPATState, len(patIDs))
for _, patID := range patIDs {
patActive, patExists := agentCompatCapabilityActiveForPAT(handler, patID)
patState[patID] = agentCompatNATPATState{active: patActive, exists: patExists}
}
return agentCompatNATCreationState{
streamState: handler.SnapshotIOStreamState(),
active: active,
used: used,
patState: patState,
}
}
func requireUnchangedAgentCompatNATCreationState(t *testing.T, handler *NezhaHandler, before agentCompatNATCreationState) {
t.Helper()
ids := make([]uint64, 0, len(before.patState))
for patID := range before.patState {
ids = append(ids, patID)
}
after := snapshotAgentCompatNATCreationState(handler, ids...)
require.Equal(t, before, after)
}
func requireHiddenCreateAgentCompatNATStream(t *testing.T, handler *NezhaHandler, handle AgentCompatNATPublishHandle, streamID string, before agentCompatNATCreationState) {
t.Helper()
lease, err := handler.CreateAgentCompatNATStream(handle, streamID)
require.Nil(t, lease)
require.True(t, errors.Is(err, ErrAgentCompatCapabilityHidden) || errors.Is(err, ErrAgentCompatCapabilityUnavailable))
requireUnchangedAgentCompatNATCreationState(t, handler, before)
}
func requireNATHandleBindingsIntact(t *testing.T, handler *NezhaHandler, handle AgentCompatNATPublishHandle, streamID string) {
t.Helper()
handler.ioStreamMutex.RLock()
defer handler.ioStreamMutex.RUnlock()
registration := handle.registration
require.NotNil(t, registration)
require.Equal(t, agentCompatCapabilityConsumed, registration.phase)
require.Equal(t, streamID, registration.streamID)
stream, exists := handler.ioStreams[streamID]
require.True(t, exists)
require.Same(t, stream, registration.stream)
}
func TestAgentCompatNATHandleCreationAuthorityIsBoundToExactRegistration(t *testing.T) {
handler := NewNezhaHandler()
firstRegistration := capabilityRegistration(capabilityOwner(601, 602), AgentCompatCapabilityNAT, 603, 604)
secondRegistration := capabilityRegistration(capabilityOwner(605, 606), AgentCompatCapabilityNAT, 603, 607)
firstCapability, err := handler.RegisterAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), firstRegistration)
if err != nil {
t.Fatal("first capability registration failed")
}
secondCapability, err := handler.RegisterAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), secondRegistration)
if err != nil {
t.Fatal("second capability registration failed")
}
firstHandle, err := handler.ConsumeAgentCompatNATCapability(capabilityAccess(firstCapability, firstRegistration))
if err != nil {
t.Fatal("first capability consume failed")
}
secondHandle, err := handler.ConsumeAgentCompatNATCapability(capabilityAccess(secondCapability, secondRegistration))
if err != nil {
t.Fatal("second capability consume failed")
}
firstLease, err := handler.CreateAgentCompatNATStream(firstHandle, "bound-first")
if err != nil || firstLease == nil {
t.Fatal("first capability did not create its stream")
}
secondLease, err := handler.CreateAgentCompatNATStream(secondHandle, "bound-second")
if err != nil || secondLease == nil {
t.Fatal("second capability did not create its stream")
}
beforeCrossPublish := snapshotAgentCompatNATCreationState(handler, firstRegistration.Owner.PATID, secondRegistration.Owner.PATID)
if err := handler.PublishAgentCompatNATStream(firstHandle, AgentCompatNATPublication{Purpose: AgentCompatCapabilityNAT, TargetServerID: 603, ResourceID: 604, StreamID: "bound-second"}); err != ErrAgentCompatCapabilityHidden {
t.Fatal("first capability published the second stream")
}
if err := handler.PublishAgentCompatNATStream(secondHandle, AgentCompatNATPublication{Purpose: AgentCompatCapabilityNAT, TargetServerID: 603, ResourceID: 607, StreamID: "bound-first"}); err != ErrAgentCompatCapabilityHidden {
t.Fatal("second capability published the first stream")
}
requireUnchangedAgentCompatNATCreationState(t, handler, beforeCrossPublish)
requireNATHandleBindingsIntact(t, handler, firstHandle, "bound-first")
requireNATHandleBindingsIntact(t, handler, secondHandle, "bound-second")
require.NoError(t, handler.PublishAgentCompatNATStream(firstHandle, AgentCompatNATPublication{Purpose: AgentCompatCapabilityNAT, TargetServerID: 603, ResourceID: 604, StreamID: "bound-first"}))
publishedBeforeRepeat := snapshotAgentCompatNATCreationState(handler, firstRegistration.Owner.PATID)
requireHiddenCreateAgentCompatNATStream(t, handler, firstHandle, "bound-after-publish", publishedBeforeRepeat)
if lease, err := handler.CreateAgentCompatNATStream(firstHandle, "bound-again"); lease != nil || err != ErrAgentCompatCapabilityHidden {
t.Fatal("repeated creation mutated first capability state")
}
if handler.StreamCount() != 2 {
t.Fatal("repeated creation changed stream accounting")
}
if err := handler.CancelAgentCompatIOStreamCapability(capabilityAccess(firstCapability, firstRegistration)); err != nil {
t.Fatal("first capability cancellation failed")
}
if _, found := handler.StreamOwnership("bound-first"); found {
t.Fatal("first stream remained after cancellation")
}
if _, found := handler.StreamOwnership("bound-second"); !found {
t.Fatal("second stream was affected by first cancellation")
}
if err := handler.CloseAgentCompatNATStreamLease(secondLease); err != nil {
t.Fatal("second exact lease close failed")
}
}
func TestAgentCompatNATHandleCreationRejectsInvalidAuthorityWithoutMutation(t *testing.T) {
handler, registration, capability := natCapabilityFixture(t, 701, 702, 703, 704)
access := capabilityAccess(capability, registration)
before := snapshotAgentCompatNATCreationState(handler, registration.Owner.PATID)
requireHiddenCreateAgentCompatNATStream(t, handler, AgentCompatNATPublishHandle{}, "invalid-zero", before)
foreignHandler, foreignRegistration, foreignCapability := natCapabilityFixture(t, 705, 706, 703, 707)
foreignHandle, err := foreignHandler.ConsumeAgentCompatNATCapability(capabilityAccess(foreignCapability, foreignRegistration))
require.NoError(t, err)
foreignBefore := snapshotAgentCompatNATCreationState(foreignHandler, foreignRegistration.Owner.PATID)
requireHiddenCreateAgentCompatNATStream(t, handler, foreignHandle, "invalid-foreign", before)
requireUnchangedAgentCompatNATCreationState(t, foreignHandler, foreignBefore)
registeredCapability := registerAgentCompatCapability(t, handler, capabilityRegistration(capabilityOwner(708, 709), AgentCompatCapabilityNAT, 703, 710))
registeredParsed, err := ParseAgentCompatIOStreamCapability(registeredCapability.String())
require.NoError(t, err)
registeredHandle := AgentCompatNATPublishHandle{capability: registeredParsed.value}
handler.ioStreamMutex.RLock()
registeredHandle.registration = handler.agentCompatCapabilities.active[registeredParsed.value]
registeredHandle.generation = registeredHandle.registration.generation
handler.ioStreamMutex.RUnlock()
registeredBefore := snapshotAgentCompatNATCreationState(handler, registration.Owner.PATID, 708)
requireHiddenCreateAgentCompatNATStream(t, handler, registeredHandle, "invalid-registered", registeredBefore)
staleHandle, err := handler.ConsumeAgentCompatNATCapability(access)
require.NoError(t, err)
require.NoError(t, handler.UnregisterAgentCompatIOStreamCapability(access))
staleBefore := snapshotAgentCompatNATCreationState(handler, registration.Owner.PATID)
requireHiddenCreateAgentCompatNATStream(t, handler, staleHandle, "invalid-unregistered", staleBefore)
cancelRegistration := capabilityRegistration(capabilityOwner(711, 712), AgentCompatCapabilityNAT, 703, 713)
cancelCapability := registerAgentCompatCapability(t, handler, cancelRegistration)
cancelAccess := capabilityAccess(cancelCapability, cancelRegistration)
cancelHandle, err := handler.ConsumeAgentCompatNATCapability(cancelAccess)
require.NoError(t, err)
require.NoError(t, handler.CancelAgentCompatIOStreamCapability(cancelAccess))
cancelledBefore := snapshotAgentCompatNATCreationState(handler, registration.Owner.PATID, cancelRegistration.Owner.PATID)
requireHiddenCreateAgentCompatNATStream(t, handler, cancelHandle, "invalid-cancelled", cancelledBefore)
wrongPurposeRegistration := capabilityRegistration(capabilityOwner(714, 715), AgentCompatCapabilityTerminal, 703, 0)
wrongPurposeCapability := registerAgentCompatCapability(t, handler, wrongPurposeRegistration)
wrongPurposeParsed, err := ParseAgentCompatIOStreamCapability(wrongPurposeCapability.String())
require.NoError(t, err)
handler.ioStreamMutex.RLock()
wrongPurposeRegistrationState := handler.agentCompatCapabilities.active[wrongPurposeParsed.value]
handler.ioStreamMutex.RUnlock()
wrongPurposeHandle := AgentCompatNATPublishHandle{registration: wrongPurposeRegistrationState, generation: wrongPurposeRegistrationState.generation, capability: wrongPurposeParsed.value}
wrongPurposeBefore := snapshotAgentCompatNATCreationState(handler, registration.Owner.PATID, wrongPurposeRegistration.Owner.PATID)
requireHiddenCreateAgentCompatNATStream(t, handler, wrongPurposeHandle, "invalid-purpose", wrongPurposeBefore)
}
func TestAgentCompatNATHandleCreationRepeatedCreatePreservesAccountingAndQuota(t *testing.T) {
handler, registration, capability := natCapabilityFixture(t, 721, 722, 723, 724)
handle, err := handler.ConsumeAgentCompatNATCapability(capabilityAccess(capability, registration))
require.NoError(t, err)
lease, err := handler.CreateAgentCompatNATStream(handle, "repeated-create-first")
require.NoError(t, err)
stateBeforeRepeat := snapshotAgentCompatNATCreationState(handler, registration.Owner.PATID)
requireHiddenCreateAgentCompatNATStream(t, handler, handle, "repeated-create-second", stateBeforeRepeat)
for index := 0; index < maxStreamsPerServer-1; index++ {
require.NoError(t, handler.CreateStreamWithPurpose("quota-boundary-"+strconv.Itoa(index), 0, registration.TargetServerID, PurposeNAT))
}
require.ErrorIs(t, handler.CreateStreamWithPurpose("quota-boundary-overflow", 0, registration.TargetServerID, PurposeNAT), ErrTooManyStreamsForServer)
require.NoError(t, handler.CloseAgentCompatNATStreamLease(lease))
for index := 0; index < maxStreamsPerServer-1; index++ {
require.NoError(t, handler.CloseStream("quota-boundary-"+strconv.Itoa(index)))
}
require.NoError(t, handler.CancelAgentCompatIOStreamCapability(capabilityAccess(capability, registration)))
}
service/rpc/io_stream_capability_nat_profile_agentcompat_test.go
+99
View File
@@ -0,0 +1,99 @@
//go:build agentcompat
package rpc
import (
"errors"
"strings"
"testing"
"github.com/stretchr/testify/require"
)
func TestAgentCompatNATCapabilityForProfileConsumesStoredRegistration(t *testing.T) {
handler, registration, capability := natCapabilityFixture(t, 61, 71, 81, 91)
access, handle, err := handler.ConsumeAgentCompatNATCapabilityForProfile(capability.String(), 81, 91)
require.NoError(t, err)
require.Equal(t, registration.Owner, access.Owner)
require.Equal(t, registration.Purpose, access.Purpose)
require.Equal(t, registration.TargetServerID, access.TargetServerID)
require.Equal(t, registration.ResourceID, access.ResourceID)
require.True(t, access.ServerAccessAllowed)
require.NotEmpty(t, handle.capability)
}
func TestAgentCompatNATCapabilityForProfileHidesMalformedUnknownAndForeignTuples(t *testing.T) {
handler, _, capability := natCapabilityFixture(t, 62, 72, 82, 92)
terminalRegistration := capabilityRegistration(capabilityOwner(66, 76), AgentCompatCapabilityTerminal, 82, 0)
terminalCapability := registerAgentCompatCapability(t, handler, terminalRegistration)
cases := []struct {
name string
value string
serverID uint64
resourceID uint64
}{
{name: "malformed", value: "not-a-capability", serverID: 82, resourceID: 92},
{name: "unknown", value: strings.Repeat("a", 43), serverID: 82, resourceID: 92},
{name: "wrong server", value: capability.String(), serverID: 83, resourceID: 92},
{name: "wrong profile", value: capability.String(), serverID: 82, resourceID: 93},
{name: "wrong purpose", value: terminalCapability.String(), serverID: 82, resourceID: 0},
}
for _, testCase := range cases {
t.Run(testCase.name, func(t *testing.T) {
_, _, err := handler.ConsumeAgentCompatNATCapabilityForProfile(testCase.value, testCase.serverID, testCase.resourceID)
require.True(t, errors.Is(err, ErrAgentCompatCapabilityHidden))
require.NotContains(t, err.Error(), testCase.value)
})
}
}
func TestAgentCompatNATCapabilityForProfileHidesRepeatedAndInactiveConsume(t *testing.T) {
handler, _, capability := natCapabilityFixture(t, 63, 73, 83, 93)
activeBefore, usedBefore := agentCompatCapabilityRegistryCounts(handler)
_, _, err := handler.ConsumeAgentCompatNATCapabilityForProfile(capability.String(), 83, 93)
require.NoError(t, err)
activeAfter, usedAfter := agentCompatCapabilityRegistryCounts(handler)
require.Equal(t, activeBefore, activeAfter)
require.Equal(t, usedBefore, usedAfter)
_, _, err = handler.ConsumeAgentCompatNATCapabilityForProfile(capability.String(), 83, 93)
require.True(t, errors.Is(err, ErrAgentCompatCapabilityHidden))
activeAfterRepeat, usedAfterRepeat := agentCompatCapabilityRegistryCounts(handler)
require.Equal(t, activeAfter, activeAfterRepeat)
require.Equal(t, usedAfter, usedAfterRepeat)
require.NoError(t, handler.CancelAgentCompatIOStreamCapability(AgentCompatCapabilityAccess{}))
}
func TestAgentCompatNATCapabilityForProfileHidesCancelledAndUnregistered(t *testing.T) {
tests := []struct {
name string
cleanup func(*NezhaHandler, AgentCompatCapabilityAccess) error
}{
{name: "cancelled", cleanup: (*NezhaHandler).CancelAgentCompatIOStreamCapability},
{name: "unregistered", cleanup: (*NezhaHandler).UnregisterAgentCompatIOStreamCapability},
}
for _, testCase := range tests {
t.Run(testCase.name, func(t *testing.T) {
handler, registration, capability := natCapabilityFixture(t, 64, 74, 84, 94)
access := capabilityAccess(capability, registration)
require.NoError(t, testCase.cleanup(handler, access))
_, _, err := handler.ConsumeAgentCompatNATCapabilityForProfile(capability.String(), 84, 94)
require.True(t, errors.Is(err, ErrAgentCompatCapabilityHidden))
})
}
}
func TestAgentCompatNATCapabilityForProfileDoesNotLeakSensitiveValues(t *testing.T) {
handler, registration, capability := natCapabilityFixture(t, 65, 75, 85, 95)
_, _, err := handler.ConsumeAgentCompatNATCapabilityForProfile(capability.String(), 86, 95)
require.Error(t, err)
message := err.Error()
for _, sensitive := range []string{capability.String(), "65", "75", "85", "95", "nat"} {
require.NotContains(t, message, sensitive)
}
require.Equal(t, AgentCompatCapabilityNAT, registration.Purpose)
}
service/rpc/io_stream_capability_publication_agentcompat.go
+66
View File
@@ -0,0 +1,66 @@
//go:build agentcompat
package rpc
import (
"errors"
"time"
)
func (s *NezhaHandler) CreateAgentCompatNATStream(handle AgentCompatNATPublishHandle, streamID string) (*AgentCompatNATStreamLease, error) {
s.ioStreamMutex.Lock()
defer s.ioStreamMutex.Unlock()
registration := handle.registration
if registration == nil || registration.generation != handle.generation {
return nil, ErrAgentCompatCapabilityHidden
}
currentRegistration, active := s.agentCompatCapabilities.active[handle.capability]
stored := registration.registration
if !active || currentRegistration != registration || registration.phase != agentCompatCapabilityConsumed ||
stored.Purpose != AgentCompatCapabilityNAT || registration.stream != nil || streamID == "" {
return nil, ErrAgentCompatCapabilityHidden
}
if err := s.createStreamLocked(streamID, 0, stored.TargetServerID, PurposeNAT); err != nil {
if err == ErrStreamAlreadyExists {
return nil, ErrAgentCompatCapabilityHidden
}
return nil, err
}
stream := s.ioStreams[streamID]
registration.streamID = streamID
registration.stream = stream
return &AgentCompatNATStreamLease{streamID: streamID, stream: stream}, nil
}
func (s *NezhaHandler) CloseAgentCompatNATStreamLease(lease *AgentCompatNATStreamLease) error {
if lease == nil {
return nil
}
return s.detachExactStream(lease.streamID, lease.stream)
}
func (s *NezhaHandler) StartAgentCompatNATStream(handle AgentCompatNATPublishHandle, timeout time.Duration) (bool, error) {
s.ioStreamMutex.RLock()
registration := handle.registration
publicationOwned := registration != nil && registration.generation == handle.generation &&
registration.phase == agentCompatCapabilityPublished && registration.streamID != "" && registration.stream != nil
if registration == nil || registration.generation != handle.generation {
s.ioStreamMutex.RUnlock()
return publicationOwned, ErrAgentCompatCapabilityHidden
}
current, active := s.agentCompatCapabilities.active[handle.capability]
stored := registration.registration
streamID := registration.streamID
stream := registration.stream
valid := active && current == registration && registration.phase == agentCompatCapabilityPublished &&
streamID != "" && stream != nil && s.ioStreams[streamID] == stream &&
stream.creatorUserID == 0 && stream.targetServerID == stored.TargetServerID &&
stream.purpose == PurposeNAT && stored.Purpose == AgentCompatCapabilityNAT
s.ioStreamMutex.RUnlock()
if !valid {
return publicationOwned, ErrAgentCompatCapabilityHidden
}
startErr := s.startStreamContext(streamID, stream, timeout)
closeErr := s.detachExactStream(streamID, stream)
return publicationOwned, errors.Join(startErr, closeErr)
}
service/rpc/io_stream_capability_publication_default.go
+9
View File
@@ -0,0 +1,9 @@
//go:build !agentcompat
package rpc
import "time"
func (*NezhaHandler) StartAgentCompatNATStream(AgentCompatNATPublishHandle, time.Duration) (bool, error) {
return false, ErrAgentCompatCapabilityUnavailable
}
service/rpc/io_stream_capability_quota_accounting_agentcompat_test.go
+230
View File
@@ -0,0 +1,230 @@
//go:build agentcompat
package rpc
import (
"errors"
"sync"
"testing"
"github.com/stretchr/testify/require"
)
func TestAgentCompatCapabilityConcurrentRegistrationEnforcesExactPerPATQuota(t *testing.T) {
handler := NewNezhaHandler()
setUniqueAgentCompatCapabilityTokens(handler)
registration := capabilityRegistration(capabilityOwner(104, 204), AgentCompatCapabilityTerminal, 305, 0)
start := make(chan struct{})
results := make(chan error, 64)
var waitGroup sync.WaitGroup
waitGroup.Add(64)
for range 64 {
go func() {
defer waitGroup.Done()
<-start
_, err := handler.RegisterAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), registration)
results <- err
}()
}
close(start)
waitGroup.Wait()
close(results)
succeeded, unavailable := 0, 0
for err := range results {
switch {
case err == nil:
succeeded++
case errors.Is(err, ErrAgentCompatCapabilityUnavailable):
unavailable++
default:
t.Fatalf("unexpected registration error: %v", err)
}
}
require.Equal(t, 16, succeeded)
require.Equal(t, 48, unavailable)
active, used := agentCompatCapabilityRegistryCounts(handler)
require.Equal(t, 16, active)
require.Equal(t, 16, used)
}
func TestAgentCompatCapabilityConcurrentRegistrationEnforcesExactGlobalQuota(t *testing.T) {
handler := NewNezhaHandler()
setUniqueAgentCompatCapabilityTokens(handler)
start := make(chan struct{})
results := make(chan error, 256)
var waitGroup sync.WaitGroup
waitGroup.Add(256)
for index := range 256 {
go func() {
defer waitGroup.Done()
<-start
registration := capabilityRegistration(capabilityOwner(uint64(index+1), uint64(index+1001)), AgentCompatCapabilityTerminal, 306, 0)
_, err := handler.RegisterAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), registration)
results <- err
}()
}
close(start)
waitGroup.Wait()
close(results)
succeeded, unavailable := 0, 0
for err := range results {
if err == nil {
succeeded++
continue
}
require.ErrorIs(t, err, ErrAgentCompatCapabilityUnavailable)
unavailable++
}
require.Equal(t, 128, succeeded)
require.Equal(t, 128, unavailable)
active, used := agentCompatCapabilityRegistryCounts(handler)
require.Equal(t, 128, active)
require.Equal(t, 128, used)
}
func TestAgentCompatCapabilityConcurrentRegistrationEnforcesExactProcessMintQuota(t *testing.T) {
handler := NewNezhaHandler()
setUniqueAgentCompatCapabilityTokens(handler)
handler.ioStreamMutex.Lock()
for index := range agentCompatCapabilityMaxProcessMints - 1 {
handler.agentCompatCapabilities.used[string(rune(index+1))] = struct{}{}
}
handler.ioStreamMutex.Unlock()
start := make(chan struct{})
results := make(chan error, 2)
var waitGroup sync.WaitGroup
waitGroup.Add(2)
for index := range 2 {
go func() {
defer waitGroup.Done()
<-start
registration := capabilityRegistration(capabilityOwner(uint64(index+201), uint64(index+301)), AgentCompatCapabilityTerminal, 312, 0)
_, err := handler.RegisterAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), registration)
results <- err
}()
}
close(start)
waitGroup.Wait()
close(results)
succeeded, unavailable := 0, 0
for err := range results {
if err == nil {
succeeded++
continue
}
require.ErrorIs(t, err, ErrAgentCompatCapabilityUnavailable)
unavailable++
}
require.Equal(t, 1, succeeded)
require.Equal(t, 1, unavailable)
active, used := agentCompatCapabilityRegistryCounts(handler)
require.Equal(t, 1, active)
require.Equal(t, agentCompatCapabilityMaxProcessMints, used)
}
func TestAgentCompatCapabilityRemovalRequiresExactActiveRegistration(t *testing.T) {
handler := NewNezhaHandler()
setUniqueAgentCompatCapabilityTokens(handler)
registration := capabilityRegistration(capabilityOwner(105, 205), AgentCompatCapabilityTerminal, 307, 0)
capability := registerAgentCompatCapability(t, handler, registration)
handler.ioStreamMutex.Lock()
activeRegistration := handler.agentCompatCapabilities.active[capability.value]
staleRegistration := &agentCompatCapabilityRegistration{registration: registration, notify: make(chan struct{})}
handler.removeAgentCompatCapabilityLocked(capability.value, staleRegistration)
handler.ioStreamMutex.Unlock()
for range 15 {
registerAgentCompatCapability(t, handler, registration)
}
_, err := handler.RegisterAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), registration)
require.ErrorIs(t, err, ErrAgentCompatCapabilityUnavailable)
active, used := agentCompatCapabilityRegistryCounts(handler)
require.Equal(t, 16, active)
require.Equal(t, 16, used)
handler.ioStreamMutex.Lock()
handler.removeAgentCompatCapabilityLocked(capability.value, activeRegistration)
handler.removeAgentCompatCapabilityLocked(capability.value, activeRegistration)
handler.ioStreamMutex.Unlock()
replacement := registerAgentCompatCapability(t, handler, registration)
require.NotEmpty(t, replacement.String())
active, used = agentCompatCapabilityRegistryCounts(handler)
require.Equal(t, 16, active)
require.Equal(t, 17, used)
}
func TestAgentCompatCapabilityForeignRemovalDoesNotReleasePerPATQuota(t *testing.T) {
handler := NewNezhaHandler()
setUniqueAgentCompatCapabilityTokens(handler)
registration := capabilityRegistration(capabilityOwner(106, 206), AgentCompatCapabilityTerminal, 308, 0)
capability := registerAgentCompatCapability(t, handler, registration)
for range 15 {
registerAgentCompatCapability(t, handler, registration)
}
foreign := capabilityAccess(capability, registration)
foreign.Owner.PATID++
require.NoError(t, handler.CancelAgentCompatIOStreamCapability(foreign))
require.NoError(t, handler.UnregisterAgentCompatIOStreamCapability(foreign))
require.NoError(t, handler.CancelAgentCompatIOStreamCapability(AgentCompatCapabilityAccess{}))
require.NoError(t, handler.UnregisterAgentCompatIOStreamCapability(AgentCompatCapabilityAccess{}))
_, err := handler.RegisterAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), registration)
require.ErrorIs(t, err, ErrAgentCompatCapabilityUnavailable)
active, used := agentCompatCapabilityRegistryCounts(handler)
require.Equal(t, 16, active)
require.Equal(t, 16, used)
}
func TestAgentCompatCapabilityCancelReleasesQuotaBeforeEndpointCloseFailure(t *testing.T) {
handler, registration, capability := boundCapabilityFixture(t, AgentCompatCapabilityTerminal, "quota-close-failure", 309)
setUniqueAgentCompatCapabilityTokens(handler)
closeErr := errors.New("endpoint close failed")
endpoint := &capabilityCloseEndpoint{handler: handler, streamID: "quota-close-failure", err: closeErr}
require.NoError(t, handler.AgentConnected("quota-close-failure", endpoint))
for range 15 {
registerAgentCompatCapability(t, handler, registration)
}
err := handler.CancelAgentCompatIOStreamCapability(capabilityAccess(capability, registration))
require.ErrorIs(t, err, closeErr)
replacement := registerAgentCompatCapability(t, handler, registration)
require.NotEmpty(t, replacement.String())
activeForPAT, exists := agentCompatCapabilityActiveForPAT(handler, registration.Owner.PATID)
require.True(t, exists)
require.Equal(t, uint16(16), activeForPAT)
}
func TestAgentCompatCapabilityBoundUnregisterConflictRetainsQuota(t *testing.T) {
handler, registration, capability := boundCapabilityFixture(t, AgentCompatCapabilityTerminal, "quota-bound-conflict", 310)
setUniqueAgentCompatCapabilityTokens(handler)
for range 15 {
registerAgentCompatCapability(t, handler, registration)
}
err := handler.UnregisterAgentCompatIOStreamCapability(capabilityAccess(capability, registration))
_, registerErr := handler.RegisterAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), registration)
require.ErrorIs(t, err, ErrAgentCompatCapabilityBound)
require.ErrorIs(t, registerErr, ErrAgentCompatCapabilityUnavailable)
activeForPAT, exists := agentCompatCapabilityActiveForPAT(handler, registration.Owner.PATID)
require.True(t, exists)
require.Equal(t, uint16(16), activeForPAT)
}
func TestAgentCompatCapabilityLastRemovalDeletesPerPATAccountingEntry(t *testing.T) {
handler := NewNezhaHandler()
setUniqueAgentCompatCapabilityTokens(handler)
registration := capabilityRegistration(capabilityOwner(107, 207), AgentCompatCapabilityTerminal, 311, 0)
capability := registerAgentCompatCapability(t, handler, registration)
require.NoError(t, handler.UnregisterAgentCompatIOStreamCapability(capabilityAccess(capability, registration)))
activeForPAT, exists := agentCompatCapabilityActiveForPAT(handler, registration.Owner.PATID)
require.False(t, exists)
require.Zero(t, activeForPAT)
}
service/rpc/io_stream_capability_quota_boundaries_agentcompat_test.go
+94
View File
@@ -0,0 +1,94 @@
//go:build agentcompat
package rpc
import (
"errors"
"testing"
"github.com/stretchr/testify/require"
)
func TestAgentCompatCapabilityRegistrationEnforcesPerPATActiveQuotaAndReusesReleasedSlot(t *testing.T) {
handler := NewNezhaHandler()
issued := setUniqueAgentCompatCapabilityTokens(handler)
registration := capabilityRegistration(capabilityOwner(101, 201), AgentCompatCapabilityTerminal, 301, 0)
capabilities := make([]AgentCompatIOStreamCapability, 0, 16)
for range 16 {
capabilities = append(capabilities, registerAgentCompatCapability(t, handler, registration))
}
_, err := handler.RegisterAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), registration)
require.ErrorIs(t, err, ErrAgentCompatCapabilityUnavailable)
require.Equal(t, uint64(16), issued.Load())
require.NoError(t, handler.UnregisterAgentCompatIOStreamCapability(capabilityAccess(capabilities[0], registration)))
replacement := registerAgentCompatCapability(t, handler, registration)
require.NotEmpty(t, replacement.String())
active, used := agentCompatCapabilityRegistryCounts(handler)
require.Equal(t, 16, active)
require.Equal(t, 17, used)
}
func TestAgentCompatCapabilityRegistrationEnforcesGlobalActiveQuotaAndReusesReleasedSlot(t *testing.T) {
handler := NewNezhaHandler()
issued := setUniqueAgentCompatCapabilityTokens(handler)
registrations := make([]AgentCompatCapabilityRegistration, 0, 128)
capabilities := make([]AgentCompatIOStreamCapability, 0, 128)
for index := range 128 {
registration := capabilityRegistration(capabilityOwner(uint64(index+1), uint64(index+1001)), AgentCompatCapabilityTerminal, 302, 0)
registrations = append(registrations, registration)
capabilities = append(capabilities, registerAgentCompatCapability(t, handler, registration))
}
overflow := capabilityRegistration(capabilityOwner(10000, 20000), AgentCompatCapabilityTerminal, 302, 0)
_, err := handler.RegisterAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), overflow)
require.ErrorIs(t, err, ErrAgentCompatCapabilityUnavailable)
require.Equal(t, uint64(128), issued.Load())
require.NoError(t, handler.UnregisterAgentCompatIOStreamCapability(capabilityAccess(capabilities[0], registrations[0])))
replacement := registerAgentCompatCapability(t, handler, overflow)
require.NotEmpty(t, replacement.String())
active, used := agentCompatCapabilityRegistryCounts(handler)
require.Equal(t, 128, active)
require.Equal(t, 129, used)
}
func TestAgentCompatCapabilityRegistrationEnforcesProcessLifetimeMintQuota(t *testing.T) {
handler := NewNezhaHandler()
issued := setUniqueAgentCompatCapabilityTokens(handler)
registration := capabilityRegistration(capabilityOwner(102, 202), AgentCompatCapabilityTerminal, 303, 0)
for range 4096 {
capability := registerAgentCompatCapability(t, handler, registration)
require.NoError(t, handler.UnregisterAgentCompatIOStreamCapability(capabilityAccess(capability, registration)))
}
_, err := handler.RegisterAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), registration)
require.ErrorIs(t, err, ErrAgentCompatCapabilityUnavailable)
require.Equal(t, uint64(4096), issued.Load())
active, used := agentCompatCapabilityRegistryCounts(handler)
require.Zero(t, active)
require.Equal(t, 4096, used)
}
func TestAgentCompatCapabilityCollisionRetriesDoNotConsumeQuota(t *testing.T) {
handler := NewNezhaHandler()
registration := capabilityRegistration(capabilityOwner(103, 203), AgentCompatCapabilityTerminal, 304, 0)
fixedToken := make([]byte, 32)
fixedToken[0] = 1
handler.setAgentCompatCapabilityTokenSourceForTest(func(destination []byte) error {
copy(destination, fixedToken)
return nil
})
first := registerAgentCompatCapability(t, handler, registration)
require.NoError(t, handler.UnregisterAgentCompatIOStreamCapability(capabilityAccess(first, registration)))
_, err := handler.RegisterAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), registration)
require.ErrorIs(t, err, ErrAgentCompatCapabilityTokenExhausted)
active, used := agentCompatCapabilityRegistryCounts(handler)
require.Zero(t, active)
require.Equal(t, 1, used)
require.False(t, errors.Is(err, ErrAgentCompatCapabilityUnavailable))
}
service/rpc/io_stream_capability_quota_test_helpers_agentcompat_test.go
+40
View File
@@ -0,0 +1,40 @@
//go:build agentcompat
package rpc
import (
"encoding/binary"
"sync/atomic"
"testing"
"github.com/stretchr/testify/require"
)
func setUniqueAgentCompatCapabilityTokens(handler *NezhaHandler) *atomic.Uint64 {
var issued atomic.Uint64
handler.setAgentCompatCapabilityTokenSourceForTest(func(destination []byte) error {
binary.LittleEndian.PutUint64(destination, issued.Add(1))
return nil
})
return &issued
}
func registerAgentCompatCapability(t *testing.T, handler *NezhaHandler, registration AgentCompatCapabilityRegistration) AgentCompatIOStreamCapability {
t.Helper()
capability, err := handler.RegisterAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), registration)
require.NoError(t, err)
return capability
}
func agentCompatCapabilityRegistryCounts(handler *NezhaHandler) (active, used int) {
handler.ioStreamMutex.RLock()
defer handler.ioStreamMutex.RUnlock()
return len(handler.agentCompatCapabilities.active), len(handler.agentCompatCapabilities.used)
}
func agentCompatCapabilityActiveForPAT(handler *NezhaHandler, patID uint64) (uint16, bool) {
handler.ioStreamMutex.RLock()
defer handler.ioStreamMutex.RUnlock()
active, exists := handler.agentCompatCapabilities.activeByPAT[patID]
return active, exists
}
service/rpc/io_stream_capability_register_agentcompat.go
+96
View File
@@ -0,0 +1,96 @@
//go:build agentcompat
package rpc
import (
"context"
"encoding/base64"
)
const agentCompatCapabilityTokenAttempts = 32
func validAgentCompatRegistration(registration AgentCompatCapabilityRegistration) bool {
if !registration.ServerAccessAllowed || registration.Owner.PATID == 0 || registration.Owner.UserID == 0 || registration.TargetServerID == 0 {
return false
}
switch registration.Purpose {
case AgentCompatCapabilityTerminal, AgentCompatCapabilityFileManager:
return registration.ResourceID == 0
case AgentCompatCapabilityNAT:
return registration.ResourceID != 0
default:
return false
}
}
func (s *NezhaHandler) RegisterAgentCompatIOStreamCapability(ctx context.Context, registration AgentCompatCapabilityRegistration) (AgentCompatIOStreamCapability, error) {
if !validAgentCompatRegistration(registration) {
return AgentCompatIOStreamCapability{}, ErrAgentCompatCapabilityHidden
}
if err := ctx.Err(); err != nil {
return AgentCompatIOStreamCapability{}, err
}
s.ioStreamMutex.RLock()
tokenSource := s.agentCompatCapabilities.tokenSource
quotaAvailable := s.agentCompatCapabilityQuotaAvailableLocked(registration.Owner.PATID)
s.ioStreamMutex.RUnlock()
if !quotaAvailable {
return AgentCompatIOStreamCapability{}, ErrAgentCompatCapabilityUnavailable
}
for range agentCompatCapabilityTokenAttempts {
if err := ctx.Err(); err != nil {
return AgentCompatIOStreamCapability{}, err
}
// Token generation may block or reenter the registry, so it must never run under ioStreamMutex.
raw := make([]byte, 32)
if err := tokenSource(raw); err != nil {
return AgentCompatIOStreamCapability{}, err
}
capability := AgentCompatIOStreamCapability{value: base64.RawURLEncoding.EncodeToString(raw)}
if err := ctx.Err(); err != nil {
return AgentCompatIOStreamCapability{}, err
}
s.ioStreamMutex.Lock()
// Recheck every quota under the insertion lock so concurrent mints cannot oversubscribe any bound.
if !s.agentCompatCapabilityQuotaAvailableLocked(registration.Owner.PATID) {
s.ioStreamMutex.Unlock()
return AgentCompatIOStreamCapability{}, ErrAgentCompatCapabilityUnavailable
}
if _, used := s.agentCompatCapabilities.used[capability.value]; used {
s.ioStreamMutex.Unlock()
continue
}
s.agentCompatCapabilities.used[capability.value] = struct{}{}
s.agentCompatCapabilities.nextIdentity++
s.agentCompatCapabilities.activeByPAT[registration.Owner.PATID]++
s.agentCompatCapabilities.active[capability.value] = &agentCompatCapabilityRegistration{
registration: registration, phase: agentCompatCapabilityRegistered,
generation: s.agentCompatCapabilities.nextIdentity, notify: make(chan struct{}),
}
s.ioStreamMutex.Unlock()
return capability, nil
}
return AgentCompatIOStreamCapability{}, ErrAgentCompatCapabilityTokenExhausted
}
func (s *NezhaHandler) agentCompatCapabilityQuotaAvailableLocked(patID uint64) bool {
return s.agentCompatCapabilities.activeByPAT[patID] < agentCompatCapabilityMaxActivePerPAT &&
len(s.agentCompatCapabilities.active) < agentCompatCapabilityMaxActiveGlobal &&
len(s.agentCompatCapabilities.used) < agentCompatCapabilityMaxProcessMints
}
func sameAgentCompatOwner(left, right AgentCompatCapabilityOwner) bool {
return left == right
}
func agentCompatAccessMatches(access AgentCompatCapabilityAccess, registration *agentCompatCapabilityRegistration) bool {
stored := registration.registration
return access.ServerAccessAllowed && sameAgentCompatOwner(access.Owner, stored.Owner) &&
access.Purpose == stored.Purpose && access.TargetServerID == stored.TargetServerID && access.ResourceID == stored.ResourceID
}
func (s *NezhaHandler) agentCompatRegistrationLocked(access AgentCompatCapabilityAccess) (*agentCompatCapabilityRegistration, bool) {
registration, exists := s.agentCompatCapabilities.active[access.Capability.value]
return registration, exists && agentCompatAccessMatches(access, registration)
}
service/rpc/io_stream_capability_security_agentcompat_test.go
+234
View File
@@ -0,0 +1,234 @@
//go:build agentcompat
package rpc
import (
"context"
"errors"
"sync"
"sync/atomic"
"testing"
"github.com/stretchr/testify/require"
)
func TestAgentCompatCapabilityRegistrationExhaustsPermanentCollisionWithoutBlockingRegistry(t *testing.T) {
handler := NewNezhaHandler()
registration := capabilityRegistration(capabilityOwner(31, 41), AgentCompatCapabilityTerminal, 51, 0)
fixedToken := make([]byte, 32)
fixedToken[0] = 1
handler.setAgentCompatCapabilityTokenSourceForTest(func(destination []byte) error {
copy(destination, fixedToken)
return nil
})
_, err := handler.RegisterAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), registration)
require.NoError(t, err)
entered := make(chan struct{})
release := make(chan struct{})
releaseCtx := agentCompatCapabilityTestContext(t)
var once sync.Once
handler.setAgentCompatCapabilityTokenSourceForTest(func(destination []byte) error {
once.Do(func() {
close(entered)
select {
case <-release:
case <-releaseCtx.Done():
}
})
copy(destination, fixedToken)
return nil
})
result := make(chan error, 1)
go func() {
_, registerErr := handler.RegisterAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), registration)
result <- registerErr
}()
awaitAgentCompatCapabilitySignal(t, entered, "token source did not enter")
registryRead := make(chan struct{})
go func() {
handler.SnapshotIOStreamState()
close(registryRead)
}()
awaitAgentCompatCapabilitySignal(t, registryRead, "token source blocked unrelated registry operation")
close(release)
require.ErrorIs(t, receiveAgentCompatCapabilityError(t, result, "permanent token collision did not terminate"), ErrAgentCompatCapabilityTokenExhausted)
}
func TestAgentCompatCapabilityRegistrationPreservesCanceledContext(t *testing.T) {
handler := NewNezhaHandler()
ctx, cancel := context.WithCancel(context.Background())
cancel()
_, err := handler.RegisterAgentCompatIOStreamCapability(ctx, capabilityRegistration(capabilityOwner(32, 42), AgentCompatCapabilityTerminal, 52, 0))
require.ErrorIs(t, err, context.Canceled)
}
func TestAgentCompatCapabilityTokenSourceCanReenterRegistry(t *testing.T) {
handler := NewNezhaHandler()
handler.setAgentCompatCapabilityTokenSourceForTest(func(destination []byte) error {
handler.SnapshotIOStreamState()
destination[0] = 1
return nil
})
result := make(chan error, 1)
go func() {
_, err := handler.RegisterAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), capabilityRegistration(capabilityOwner(33, 43), AgentCompatCapabilityTerminal, 53, 0))
result <- err
}()
require.NoError(t, receiveAgentCompatCapabilityError(t, result, "reentrant token source deadlocked"))
}
func TestAgentCompatCapabilityWaitObserverCanReenterRegistry(t *testing.T) {
handler := NewNezhaHandler()
registration := capabilityRegistration(capabilityOwner(34, 44), AgentCompatCapabilityTerminal, 54, 0)
capability, err := handler.RegisterAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), registration)
require.NoError(t, err)
access := capabilityAccess(capability, registration)
handler.setAgentCompatCapabilityWaitObserverForTest(func() {
require.NoError(t, handler.UnregisterAgentCompatIOStreamCapability(access))
})
result := make(chan error, 1)
waitCtx := agentCompatCapabilityTestContext(t)
go func() {
_, waitErr := handler.WaitAgentCompatIOStreamCapability(waitCtx, access)
result <- waitErr
}()
require.ErrorIs(t, receiveAgentCompatCapabilityError(t, result, "reentrant wait observer deadlocked"), ErrAgentCompatCapabilityHidden)
}
func TestAgentCompatCapabilityWaitRejectsSameIDReplacement(t *testing.T) {
handler, registration, capability := boundCapabilityFixture(t, AgentCompatCapabilityTerminal, "reused-stream-id", 55)
require.NoError(t, handler.CloseStream("reused-stream-id"))
require.NoError(t, handler.CreateStreamWithPurpose("reused-stream-id", 21, 55, PurposeTerminal))
_, err := handler.WaitAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), capabilityAccess(capability, registration))
require.ErrorIs(t, err, ErrAgentCompatCapabilityHidden)
}
func TestAgentCompatCapabilityCancelAndUnregisterDoNotEnumerateForeignIdentity(t *testing.T) {
operations := []struct {
name string
run func(*NezhaHandler, AgentCompatCapabilityAccess) error
}{
{name: "cancel", run: (*NezhaHandler).CancelAgentCompatIOStreamCapability},
{name: "unregister", run: (*NezhaHandler).UnregisterAgentCompatIOStreamCapability},
}
for _, operation := range operations {
t.Run(operation.name, func(t *testing.T) {
handler, registration, capability := boundCapabilityFixture(t, AgentCompatCapabilityTerminal, operation.name+"-foreign", 56)
foreign := capabilityAccess(capability, registration)
foreign.Owner.PATID++
before := handler.SnapshotIOStreamState()
foreignErr := operation.run(handler, foreign)
unknownErr := operation.run(handler, AgentCompatCapabilityAccess{})
require.NoError(t, foreignErr)
require.NoError(t, unknownErr)
require.Equal(t, before, handler.SnapshotIOStreamState())
_, found := handler.StreamOwnership(operation.name + "-foreign")
require.True(t, found)
})
}
}
func TestAgentCompatCapabilityAccessMismatchMatrixIsHiddenOrInert(t *testing.T) {
mutations := []struct {
name string
mutate func(*AgentCompatCapabilityAccess)
}{
{name: "PAT", mutate: func(access *AgentCompatCapabilityAccess) { access.Owner.PATID++ }},
{name: "user", mutate: func(access *AgentCompatCapabilityAccess) { access.Owner.UserID++ }},
{name: "admin", mutate: func(access *AgentCompatCapabilityAccess) { access.Owner.IsAdmin = !access.Owner.IsAdmin }},
{name: "purpose", mutate: func(access *AgentCompatCapabilityAccess) { access.Purpose = AgentCompatCapabilityFileManager }},
{name: "resource", mutate: func(access *AgentCompatCapabilityAccess) { access.ResourceID++ }},
{name: "server", mutate: func(access *AgentCompatCapabilityAccess) { access.TargetServerID++ }},
{name: "access proof", mutate: func(access *AgentCompatCapabilityAccess) { access.ServerAccessAllowed = false }},
}
for _, mutation := range mutations {
t.Run(mutation.name, func(t *testing.T) {
handler := NewNezhaHandler()
registration := capabilityRegistration(capabilityOwner(35, 45), AgentCompatCapabilityTerminal, 57, 0)
capability, err := handler.RegisterAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), registration)
require.NoError(t, err)
require.NoError(t, handler.CreateStreamWithPurpose("matrix", 45, 57, PurposeTerminal))
access := capabilityAccess(capability, registration)
mutation.mutate(&access)
_, waitErr := handler.WaitAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), access)
bindErr := handler.BindAgentCompatIOStreamCapability(AgentCompatCapabilityBinding{AgentCompatCapabilityAccess: access, StreamID: "matrix"})
cancelErr := handler.CancelAgentCompatIOStreamCapability(access)
unregisterErr := handler.UnregisterAgentCompatIOStreamCapability(access)
require.ErrorIs(t, waitErr, ErrAgentCompatCapabilityHidden)
require.ErrorIs(t, bindErr, ErrAgentCompatCapabilityHidden)
require.NoError(t, cancelErr)
require.NoError(t, unregisterErr)
require.NoError(t, handler.BindAgentCompatIOStreamCapability(AgentCompatCapabilityBinding{
AgentCompatCapabilityAccess: capabilityAccess(capability, registration), StreamID: "matrix",
}))
streamID, err := handler.WaitAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), capabilityAccess(capability, registration))
require.NoError(t, err)
require.Equal(t, "matrix", streamID)
})
}
}
func TestAgentCompatNATCapabilityConsumeMismatchMatrixIsHidden(t *testing.T) {
mutations := []struct {
name string
mutate func(*AgentCompatCapabilityAccess)
}{
{name: "PAT", mutate: func(access *AgentCompatCapabilityAccess) { access.Owner.PATID++ }},
{name: "user", mutate: func(access *AgentCompatCapabilityAccess) { access.Owner.UserID++ }},
{name: "admin", mutate: func(access *AgentCompatCapabilityAccess) { access.Owner.IsAdmin = !access.Owner.IsAdmin }},
{name: "purpose", mutate: func(access *AgentCompatCapabilityAccess) { access.Purpose = AgentCompatCapabilityTerminal }},
{name: "resource", mutate: func(access *AgentCompatCapabilityAccess) { access.ResourceID++ }},
{name: "server", mutate: func(access *AgentCompatCapabilityAccess) { access.TargetServerID++ }},
{name: "access proof", mutate: func(access *AgentCompatCapabilityAccess) { access.ServerAccessAllowed = false }},
}
for _, mutation := range mutations {
t.Run(mutation.name, func(t *testing.T) {
handler, registration, capability := natCapabilityFixture(t, 36, 46, 58, 68)
access := capabilityAccess(capability, registration)
mutation.mutate(&access)
_, err := handler.ConsumeAgentCompatNATCapability(access)
require.True(t, errors.Is(err, ErrAgentCompatCapabilityHidden))
})
}
}
func TestAgentCompatCapabilityCancelBeforeBindWakesWaiterAndPreventsBind(t *testing.T) {
handler := NewNezhaHandler()
registration := capabilityRegistration(capabilityOwner(37, 47), AgentCompatCapabilityTerminal, 59, 0)
capability, err := handler.RegisterAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), registration)
require.NoError(t, err)
access := capabilityAccess(capability, registration)
started := make(chan struct{})
var observed atomic.Bool
handler.setAgentCompatCapabilityWaitObserverForTest(func() {
if observed.CompareAndSwap(false, true) {
close(started)
}
})
result := make(chan error, 1)
waitCtx := agentCompatCapabilityTestContext(t)
go func() {
_, waitErr := handler.WaitAgentCompatIOStreamCapability(waitCtx, access)
result <- waitErr
}()
awaitAgentCompatCapabilitySignal(t, started, "wait observer did not start")
require.NoError(t, handler.CancelAgentCompatIOStreamCapability(access))
require.ErrorIs(t, receiveAgentCompatCapabilityError(t, result, "canceled waiter did not return"), ErrAgentCompatCapabilityHidden)
require.NoError(t, handler.CreateStreamWithPurpose("after-cancel", 47, 59, PurposeTerminal))
require.ErrorIs(t, handler.BindAgentCompatIOStreamCapability(AgentCompatCapabilityBinding{AgentCompatCapabilityAccess: access, StreamID: "after-cancel"}), ErrAgentCompatCapabilityHidden)
}
service/rpc/io_stream_capability_state_agentcompat.go
+78
View File
@@ -0,0 +1,78 @@
//go:build agentcompat
package rpc
import (
"crypto/rand"
)
type agentCompatCapabilityPhase uint8
const (
agentCompatCapabilityRegistered agentCompatCapabilityPhase = iota + 1
agentCompatCapabilityConsumed
agentCompatCapabilityPublished
)
const (
agentCompatCapabilityMaxActivePerPAT = 16
agentCompatCapabilityMaxActiveGlobal = 128
agentCompatCapabilityMaxProcessMints = 4096
)
type agentCompatCapabilityRegistration struct {
registration AgentCompatCapabilityRegistration
phase agentCompatCapabilityPhase
generation uint64
streamID string
stream *ioStreamContext
notify chan struct{}
}
type agentCompatCapabilityState struct {
active map[string]*agentCompatCapabilityRegistration
activeByPAT map[uint64]uint16
// Used tokens are process-lifetime tombstones; deletion never makes a capability reusable.
used map[string]struct{}
tokenSource func([]byte) error
nextIdentity uint64
waitObserver func()
publishObserver func()
}
func (s *NezhaHandler) initializeAgentCompatCapabilities() {
s.agentCompatCapabilities.active = make(map[string]*agentCompatCapabilityRegistration)
s.agentCompatCapabilities.activeByPAT = make(map[uint64]uint16)
s.agentCompatCapabilities.used = make(map[string]struct{})
s.agentCompatCapabilities.tokenSource = func(destination []byte) error {
_, err := rand.Read(destination)
return err
}
}
func (s *NezhaHandler) setAgentCompatCapabilityTokenSourceForTest(source func([]byte) error) {
s.ioStreamMutex.Lock()
defer s.ioStreamMutex.Unlock()
s.agentCompatCapabilities.tokenSource = source
}
func (s *NezhaHandler) setAgentCompatCapabilityWaitObserverForTest(observer func()) {
s.ioStreamMutex.Lock()
defer s.ioStreamMutex.Unlock()
s.agentCompatCapabilities.waitObserver = observer
}
func (s *NezhaHandler) setAgentCompatCapabilityPublishObserverForTest(observer func()) {
s.ioStreamMutex.Lock()
defer s.ioStreamMutex.Unlock()
s.agentCompatCapabilities.publishObserver = observer
}
func (s *NezhaHandler) SetAgentCompatCapabilityPublishObserverForTest(observer func()) {
s.setAgentCompatCapabilityPublishObserverForTest(observer)
}
func (registration *agentCompatCapabilityRegistration) publishLocked() {
close(registration.notify)
registration.notify = make(chan struct{})
}
service/rpc/io_stream_capability_state_default.go
+9
View File
@@ -0,0 +1,9 @@
//go:build !agentcompat
package rpc
type agentCompatCapabilityState struct{}
type agentCompatCapabilityRegistration struct{}
func (*NezhaHandler) initializeAgentCompatCapabilities() {}
service/rpc/io_stream_capability_test_helpers_agentcompat_test.go
+38
View File
@@ -0,0 +1,38 @@
//go:build agentcompat
package rpc
import (
"context"
"testing"
"time"
)
const agentCompatCapabilityTestTimeout = 5 * time.Second
func agentCompatCapabilityTestContext(t *testing.T) context.Context {
t.Helper()
ctx, cancel := context.WithTimeout(context.Background(), agentCompatCapabilityTestTimeout)
t.Cleanup(cancel)
return ctx
}
func awaitAgentCompatCapabilitySignal(t *testing.T, signal <-chan struct{}, failureMessage string) {
t.Helper()
select {
case <-signal:
case <-agentCompatCapabilityTestContext(t).Done():
t.Fatal(failureMessage)
}
}
func receiveAgentCompatCapabilityError(t *testing.T, result <-chan error, failureMessage string) error {
t.Helper()
select {
case err := <-result:
return err
case <-agentCompatCapabilityTestContext(t).Done():
t.Fatal(failureMessage)
return nil
}
}
service/rpc/io_stream_capability_types.go
+93
View File
@@ -0,0 +1,93 @@
package rpc
import (
"encoding/base64"
"errors"
)
var (
ErrAgentCompatCapabilityUnavailable = errors.New("agentcompat IOStream capability unavailable")
ErrAgentCompatCapabilityHidden = errors.New("agentcompat IOStream capability unavailable")
ErrAgentCompatCapabilityConflict = errors.New("agentcompat IOStream capability conflict")
ErrAgentCompatCapabilityBound = errors.New("agentcompat IOStream capability has a live bound stream")
ErrAgentCompatCapabilityTokenExhausted = errors.New("agentcompat IOStream capability token attempts exhausted")
)
type AgentCompatCapabilityPurpose uint8
const (
AgentCompatCapabilityTerminal AgentCompatCapabilityPurpose = iota + 1
AgentCompatCapabilityFileManager
AgentCompatCapabilityNAT
)
func (purpose AgentCompatCapabilityPurpose) streamPurpose() StreamPurpose {
switch purpose {
case AgentCompatCapabilityTerminal:
return PurposeTerminal
case AgentCompatCapabilityFileManager:
return PurposeFileManager
case AgentCompatCapabilityNAT:
return PurposeNAT
default:
return PurposeLegacy
}
}
type AgentCompatCapabilityOwner struct {
PATID uint64
UserID uint64
IsAdmin bool
}
type AgentCompatCapabilityRegistration struct {
Owner AgentCompatCapabilityOwner
Purpose AgentCompatCapabilityPurpose
TargetServerID uint64
ResourceID uint64
ServerAccessAllowed bool
}
type AgentCompatIOStreamCapability struct{ value string }
func (capability AgentCompatIOStreamCapability) String() string { return capability.value }
func ParseAgentCompatIOStreamCapability(value string) (AgentCompatIOStreamCapability, error) {
raw, err := base64.RawURLEncoding.DecodeString(value)
if err != nil || len(raw) != 32 {
return AgentCompatIOStreamCapability{}, ErrAgentCompatCapabilityHidden
}
return AgentCompatIOStreamCapability{value: value}, nil
}
type AgentCompatCapabilityAccess struct {
Capability AgentCompatIOStreamCapability
Owner AgentCompatCapabilityOwner
Purpose AgentCompatCapabilityPurpose
TargetServerID uint64
ResourceID uint64
ServerAccessAllowed bool
}
type AgentCompatCapabilityBinding struct {
AgentCompatCapabilityAccess
StreamID string
}
type AgentCompatNATPublishHandle struct {
registration *agentCompatCapabilityRegistration
generation uint64
capability string
}
type AgentCompatNATStreamLease struct {
streamID string
stream *ioStreamContext
}
type AgentCompatNATPublication struct {
Purpose AgentCompatCapabilityPurpose
TargetServerID uint64
ResourceID uint64
StreamID string
}
service/rpc/io_stream_leak_test.go
+67
View File
@@ -0,0 +1,67 @@
package rpc
import (
"runtime"
"testing"
"time"
)
// settleGoroutines lets transient goroutines wind down so the count reflects
// only durable leaks, not in-flight teardown.
func settleGoroutines() int {
var n int
for i := 0; i < 50; i++ {
runtime.GC()
time.Sleep(20 * time.Millisecond)
n = runtime.NumGoroutine()
}
return n
}
// TestStartStream_NoGoroutineLeakAfterClose verifies the bidirectional relay in
// StartStream does not strand a goroutine. StartStream launches two
// io.CopyBuffer goroutines (user<-agent and agent<-user) but returns after the
// first one finishes. The second goroutine stays blocked in CopyBuffer until
// its endpoints are closed. CloseStream closes both endpoints, which must
// unblock and drain that second goroutine. If it doesn't, every terminal / fm /
// NAT session leaks one goroutine for the lifetime of the dashboard.
func TestStartStream_NoGoroutineLeakAfterClose(t *testing.T) {
base := settleGoroutines()
const n = 20
for i := 0; i < n; i++ {
h := NewNezhaHandler()
const id = "leak-stream"
if err := h.CreateStream(id, 1, 1); err != nil {
t.Fatalf("CreateStream: %v", err)
}
userIo, agentIo := newPipeReadWriter(), newPipeReadWriter()
h.AgentConnected(id, agentIo)
h.UserConnected(id, userIo)
done := make(chan struct{})
go func() {
_ = h.StartStream(id, time.Second*5)
close(done)
}()
// Close one endpoint so the first CopyBuffer returns and StartStream
// unblocks, mirroring a peer disconnect.
time.Sleep(10 * time.Millisecond)
userIo.Close()
<-done
// The caller's defer CloseStream closes both endpoints, which must
// drain the still-blocked second copy goroutine.
_ = h.CloseStream(id)
agentIo.Close()
}
after := settleGoroutines()
if grew := after - base; grew > 2 {
t.Fatalf("goroutine leak in StartStream relay: ran %d streams, goroutines grew by %d (base=%d after=%d)",
n, grew, base, after)
}
}
service/rpc/io_stream_lifecycle.go
+128
View File
@@ -0,0 +1,128 @@
package rpc
import (
"context"
"errors"
"io"
"time"
)
type ioStreamDetach struct {
stream *ioStreamContext
endpoints []io.ReadWriteCloser
}
func detachStreamLocked(streamID string, retainedStream *ioStreamContext, streams map[string]*ioStreamContext) (ioStreamDetach, bool) {
current, live := streams[streamID]
if streamID == "" || !live || current != retainedStream {
return ioStreamDetach{}, false
}
retainedStream.revoke()
endpoints := make([]io.ReadWriteCloser, 0, 2)
if retainedStream.userIo != nil {
endpoints = append(endpoints, retainedStream.userIo)
}
if retainedStream.agentIo != nil {
endpoints = append(endpoints, retainedStream.agentIo)
}
delete(streams, streamID)
return ioStreamDetach{stream: retainedStream, endpoints: endpoints}, true
}
func (s *NezhaHandler) detachExactStream(streamID string, retainedStream *ioStreamContext) error {
s.ioStreamMutex.Lock()
detached, ok := detachStreamLocked(streamID, retainedStream, s.ioStreams)
if !ok {
s.ioStreamMutex.Unlock()
return nil
}
s.publishIOStreamStateChangeLocked()
s.ioStreamMutex.Unlock()
var closeErrors []error
for _, endpoint := range detached.endpoints {
if err := endpoint.Close(); err != nil {
closeErrors = append(closeErrors, err)
}
}
return errors.Join(closeErrors...)
}
func (s *NezhaHandler) detachStreams(shouldDetach func(*ioStreamContext) bool) (int, error) {
s.ioStreamMutex.Lock()
detached := make([]ioStreamDetach, 0)
for streamID, stream := range s.ioStreams {
if !shouldDetach(stream) {
continue
}
item, ok := detachStreamLocked(streamID, stream, s.ioStreams)
if ok {
detached = append(detached, item)
}
}
if len(detached) > 0 {
s.publishIOStreamStateChangeLocked()
}
s.ioStreamMutex.Unlock()
var closeErrors []error
for _, item := range detached {
// Registry publication must precede endpoint Close so Close implementations may reenter safely.
for _, endpoint := range item.endpoints {
if err := endpoint.Close(); err != nil {
closeErrors = append(closeErrors, err)
}
}
}
return len(detached), errors.Join(closeErrors...)
}
func (s *NezhaHandler) CloseStream(streamID string) error {
_, err := s.detachStreams(func(stream *ioStreamContext) bool {
return stream != nil && streamID != "" && stream == s.ioStreams[streamID]
})
return err
}
func (s *NezhaHandler) WaitForAgent(ctx context.Context, streamID string, timeout time.Duration) (io.ReadWriteCloser, bool) {
deadline := time.NewTimer(timeout)
defer deadline.Stop()
for {
s.ioStreamMutex.RLock()
stream, ok := s.ioStreams[streamID]
if ok && stream.agentIo != nil {
agentIo := stream.agentIo
s.ioStreamMutex.RUnlock()
return agentIo, true
}
if !ok {
s.ioStreamMutex.RUnlock()
return nil, false
}
revokedCh := stream.revokedCh
agentIoConnectCh := stream.agentIoConnectCh
stream.waitStartedOnce.Do(func() { close(stream.waitStartedCh) })
s.ioStreamMutex.RUnlock()
select {
case <-ctx.Done():
return nil, false
case <-deadline.C:
return nil, false
case <-revokedCh:
return nil, false
case <-agentIoConnectCh:
}
}
}
func (s *NezhaHandler) RevokeStreamsForServer(serverID uint64) {
if serverID == 0 {
return
}
_, _ = s.detachStreams(func(stream *ioStreamContext) bool { return stream.targetServerID == serverID })
}
func (s *NezhaHandler) RevokeStreamsForPurpose(purpose StreamPurpose) int {
revoked, _ := s.detachStreams(func(stream *ioStreamContext) bool { return stream.purpose == purpose })
return revoked
}
service/rpc/io_stream_lifecycle_test.go
+259
View File
@@ -0,0 +1,259 @@
package rpc
import (
"context"
"errors"
"io"
"testing"
"time"
)
type lifecycleRWC struct {
closed chan struct{}
}
type reenteringErrorRWC struct {
handler *NezhaHandler
streamID string
err error
}
func (stream *reenteringErrorRWC) Read([]byte) (int, error) { return 0, io.EOF }
func (stream *reenteringErrorRWC) Write(data []byte) (int, error) { return len(data), nil }
func (stream *reenteringErrorRWC) Close() error {
if _, ok := stream.handler.StreamOwnership(stream.streamID); ok {
return errors.Join(stream.err, errors.New("stream remained registered during endpoint close"))
}
return stream.err
}
func newLifecycleRWC() *lifecycleRWC {
return &lifecycleRWC{closed: make(chan struct{})}
}
func (stream *lifecycleRWC) Read([]byte) (int, error) { return 0, io.EOF }
func (stream *lifecycleRWC) Write(data []byte) (int, error) { return len(data), nil }
func (stream *lifecycleRWC) Close() error {
select {
case <-stream.closed:
default:
close(stream.closed)
}
return nil
}
func TestIOStreamValidCreateAttachCloseLifecycle(t *testing.T) {
handler := NewNezhaHandler()
user := newLifecycleRWC()
agent := newLifecycleRWC()
if err := handler.CreateStream("valid-lifecycle", 11, 22); err != nil {
t.Fatalf("Given a new stream, CreateStream failed: %v", err)
}
if err := handler.UserConnected("valid-lifecycle", user); err != nil {
t.Fatalf("Given a tracked stream, UserConnected failed: %v", err)
}
if err := handler.AgentConnected("valid-lifecycle", agent); err != nil {
t.Fatalf("Given a tracked stream, AgentConnected failed: %v", err)
}
if _, ok := handler.StreamOwnership("valid-lifecycle"); !ok {
t.Fatal("Then a valid attached stream must remain tracked")
}
if err := handler.CloseStream("valid-lifecycle"); err != nil {
t.Fatalf("When closing the valid stream, CloseStream failed: %v", err)
}
if _, ok := handler.StreamOwnership("valid-lifecycle"); ok {
t.Fatal("Then CloseStream must remove the tracked stream")
}
}
func TestCreateStreamKeepsExistingStreamWhenIDIsDuplicated(t *testing.T) {
handler := NewNezhaHandler()
original := newLifecycleRWC()
if err := handler.CreateStream("duplicate-id", 11, 22); err != nil {
t.Fatalf("Given a new stream ID, CreateStream failed: %v", err)
}
if err := handler.AgentConnected("duplicate-id", original); err != nil {
t.Fatalf("Given a live stream, AgentConnected failed: %v", err)
}
err := handler.CreateStream("duplicate-id", 33, 44)
if !errors.Is(err, ErrStreamAlreadyExists) {
t.Fatalf("When reusing a live ID, expected ErrStreamAlreadyExists, got %v", err)
}
owner, found := handler.StreamOwnership("duplicate-id")
if !found || owner != 11 {
t.Fatalf("Then the original stream ownership must remain, found=%v owner=%d", found, owner)
}
select {
case <-original.closed:
t.Fatal("Then duplicate creation must not close the original endpoint")
default:
}
}
func TestAgentConnectedRejectsDuplicateEndpointWithoutReplacingLiveRelay(t *testing.T) {
handler := NewNezhaHandler()
first := newLifecycleRWC()
second := newLifecycleRWC()
if err := handler.CreateStream("agent-once", 11, 22); err != nil {
t.Fatalf("Given a new stream, CreateStream failed: %v", err)
}
if err := handler.AgentConnected("agent-once", first); err != nil {
t.Fatalf("Given no agent endpoint, AgentConnected failed: %v", err)
}
if err := handler.AgentConnected("agent-once", second); !errors.Is(err, ErrAgentStreamAlreadyConnected) {
t.Fatalf("When attaching a second agent endpoint, expected ErrAgentStreamAlreadyConnected, got %v", err)
}
endpoints, err := handler.GetStream("agent-once")
if err != nil || endpoints.agentIo != first {
t.Fatalf("Then the first endpoint must remain attached, err=%v", err)
}
select {
case <-second.closed:
default:
t.Fatal("Then the rejected duplicate endpoint must be closed")
}
}
func TestCloseStreamWakesWaitForAgentAndAllowsSlotReuse(t *testing.T) {
handler := NewNezhaHandler()
if err := handler.CreateStream("wait-close", 11, 22); err != nil {
t.Fatalf("Given a pending stream, CreateStream failed: %v", err)
}
stream, err := handler.GetStream("wait-close")
if err != nil {
t.Fatalf("Given a created stream, GetStream failed: %v", err)
}
result := make(chan bool, 1)
go func() {
_, ok := handler.WaitForAgent(context.Background(), "wait-close", time.Minute)
result <- ok
}()
select {
case <-stream.waitStartedCh:
case <-time.After(time.Second):
t.Fatal("WaitForAgent did not enter its blocking select")
}
if err := handler.CloseStream("wait-close"); err != nil {
t.Fatalf("When closing a pending stream, CloseStream failed: %v", err)
}
select {
case ok := <-result:
if ok {
t.Fatal("Then WaitForAgent must report no attached agent")
}
case <-time.After(time.Second):
t.Fatal("Then CloseStream must wake WaitForAgent")
}
if err := handler.CreateStream("wait-close-reused", 11, 22); err != nil {
t.Fatalf("Then the released user/server slot must be reusable: %v", err)
}
}
func TestRevokeStreamsForPurposeWakesWaitForAgentAndIsRepeatable(t *testing.T) {
handler := NewNezhaHandler()
if err := handler.CreateStreamWithPurpose("revoke-wait", 0, 22, PurposeMCPTransfer); err != nil {
t.Fatalf("Given a pending MCP stream, CreateStream failed: %v", err)
}
stream, err := handler.GetStream("revoke-wait")
if err != nil {
t.Fatalf("Given a created stream, GetStream failed: %v", err)
}
result := make(chan bool, 1)
go func() {
_, ok := handler.WaitForAgent(context.Background(), "revoke-wait", time.Minute)
result <- ok
}()
select {
case <-stream.waitStartedCh:
case <-time.After(time.Second):
t.Fatal("WaitForAgent did not enter its blocking select")
}
if revoked := handler.RevokeStreamsForPurpose(PurposeMCPTransfer); revoked != 1 {
t.Fatalf("When revoking the purpose, expected one stream, got %d", revoked)
}
if revoked := handler.RevokeStreamsForPurpose(PurposeMCPTransfer); revoked != 0 {
t.Fatalf("When repeating revocation, expected zero streams, got %d", revoked)
}
select {
case ok := <-result:
if ok {
t.Fatal("Then WaitForAgent must report no attached agent")
}
case <-time.After(time.Second):
t.Fatal("Then revocation must wake WaitForAgent")
}
}
func TestCloseStreamDetachesBeforeReenteringEndpointCloseAndJoinsErrors(t *testing.T) {
handler := NewNezhaHandler()
if err := handler.CreateStream("close-errors", 1, 1); err != nil {
t.Fatal(err)
}
firstErr := errors.New("first close error")
secondErr := errors.New("second close error")
if err := handler.UserConnected("close-errors", &reenteringErrorRWC{handler: handler, streamID: "close-errors", err: firstErr}); err != nil {
t.Fatal(err)
}
if err := handler.AgentConnected("close-errors", &reenteringErrorRWC{handler: handler, streamID: "close-errors", err: secondErr}); err != nil {
t.Fatal(err)
}
err := handler.CloseStream("close-errors")
if !errors.Is(err, firstErr) || !errors.Is(err, secondErr) {
t.Fatalf("close errors were not joined: %v", err)
}
}
func TestStartStreamReturnsImmediatelyWhenRevoked(t *testing.T) {
handler := NewNezhaHandler()
if err := handler.CreateStream("start-revoked", 1, 1); err != nil {
t.Fatal(err)
}
result := make(chan error, 1)
go func() { result <- handler.StartStream("start-revoked", time.Minute) }()
if revoked := handler.RevokeStreamsForPurpose(PurposeLegacy); revoked != 1 {
t.Fatalf("revoked streams: %d", revoked)
}
select {
case err := <-result:
if err == nil {
t.Fatal("revoked StartStream must return an error")
}
case <-time.After(time.Second):
t.Fatal("StartStream did not wake on revoke")
}
}
func TestConcurrentCloseAndRevokePublishOneGeneration(t *testing.T) {
handler := NewNezhaHandler()
if err := handler.CreateStream("single-generation", 1, 1); err != nil {
t.Fatal(err)
}
start := handler.SnapshotIOStreamState()
closeDone := make(chan struct{})
revokeDone := make(chan struct{})
go func() {
_ = handler.CloseStream("single-generation")
close(closeDone)
}()
go func() {
handler.RevokeStreamsForPurpose(PurposeLegacy)
close(revokeDone)
}()
select {
case <-closeDone:
case <-time.After(time.Second):
t.Fatal("CloseStream did not complete")
}
select {
case <-revokeDone:
case <-time.After(time.Second):
t.Fatal("RevokeStreamsForPurpose did not complete")
}
state := handler.SnapshotIOStreamState()
if state.Count != 0 || state.Generation != start.Generation+1 {
t.Fatalf("single detach publication: start=%+v final=%+v", start, state)
}
}
service/rpc/io_stream_quota.go
+62
View File
@@ -0,0 +1,62 @@
package rpc
import "errors"
const (
maxStreamsPerUser = 20
maxStreamsPerServer = 40
)
var (
ErrTooManyStreamsForUser = errors.New("too many concurrent streams for this user")
ErrTooManyStreamsForServer = errors.New("too many concurrent streams for this server")
ErrStreamAlreadyExists = errors.New("stream already exists")
)
func (s *NezhaHandler) CreateStream(streamId string, creatorUserID uint64, targetServerID uint64) error {
return s.CreateStreamWithPurpose(streamId, creatorUserID, targetServerID, PurposeLegacy)
}
func (s *NezhaHandler) CreateStreamWithPurpose(streamId string, creatorUserID uint64, targetServerID uint64, purpose StreamPurpose) error {
s.ioStreamMutex.Lock()
defer s.ioStreamMutex.Unlock()
return s.createStreamLocked(streamId, creatorUserID, targetServerID, purpose)
}
func (s *NezhaHandler) createStreamLocked(streamId string, creatorUserID uint64, targetServerID uint64, purpose StreamPurpose) error {
if _, exists := s.ioStreams[streamId]; exists {
// Stream IDs identify live relay ownership; never overwrite one or orphan its endpoint.
return ErrStreamAlreadyExists
}
var perUser, perServer int
for _, ctx := range s.ioStreams {
if creatorUserID != 0 && ctx.creatorUserID == creatorUserID {
perUser++
}
if ctx.targetServerID == targetServerID {
perServer++
}
}
// creatorUserID==0 is a dashboard-internal stream (NAT, server transfer,
// MCP transfer); only end-user-initiated streams are capped per user, but
// every stream counts toward the per-server cap so one server cannot be
// flooded regardless of who opened the streams.
if creatorUserID != 0 && perUser >= maxStreamsPerUser {
return ErrTooManyStreamsForUser
}
if perServer >= maxStreamsPerServer {
return ErrTooManyStreamsForServer
}
s.ioStreams[streamId] = newIOStreamContext(creatorUserID, targetServerID, purpose)
s.publishIOStreamStateChangeLocked()
return nil
}
// StreamCount reports the registry size under the same lock used by lifecycle mutations.
func (s *NezhaHandler) StreamCount() int {
s.ioStreamMutex.RLock()
defer s.ioStreamMutex.RUnlock()
return len(s.ioStreams)
}
service/rpc/io_stream_quota_agentcompat.go
+144
View File
@@ -0,0 +1,144 @@
//go:build agentcompat
package rpc
import (
"context"
"errors"
"fmt"
"time"
)
type IOStreamQuotaProbeResult struct {
UserAccepted int
UserRejected int
ServerAccepted int
ServerRejected int
TrackedStreams int
WaitForAgentWokeOnClose bool
UserSlotReused bool
UserBoundaryError error
ServerBoundaryError error
Err error
}
func RunIOStreamQuotaProbe(ctx context.Context) IOStreamQuotaProbeResult {
if err := ctx.Err(); err != nil {
return IOStreamQuotaProbeResult{Err: err}
}
h := NewNezhaHandler()
result := IOStreamQuotaProbeResult{}
defer func() {
h.ioStreamMutex.RLock()
streamIDs := make([]string, 0, len(h.ioStreams))
for streamID := range h.ioStreams {
streamIDs = append(streamIDs, streamID)
}
h.ioStreamMutex.RUnlock()
for _, streamID := range streamIDs {
_ = h.CloseStream(streamID)
}
}()
for i := 0; i < maxStreamsPerUser; i++ {
if err := h.CreateStream(fmt.Sprintf("probe-user-%d", i), 101, uint64(i+1)); err != nil {
result.Err = fmt.Errorf("create user stream %d: %w", i, err)
return result
}
result.UserAccepted++
}
result.UserBoundaryError = h.CreateStream("probe-user-over", 101, 500)
if !errors.Is(result.UserBoundaryError, ErrTooManyStreamsForUser) {
result.Err = fmt.Errorf("user boundary returned %v", result.UserBoundaryError)
return result
}
result.UserRejected = 1
for i := 0; i < maxStreamsPerServer; i++ {
if err := h.CreateStream(fmt.Sprintf("probe-server-%d", i), uint64(i+1000), 700); err != nil {
result.Err = fmt.Errorf("create server stream %d: %w", i, err)
return result
}
result.ServerAccepted++
}
result.ServerBoundaryError = h.CreateStream("probe-server-over", 2000, 700)
if !errors.Is(result.ServerBoundaryError, ErrTooManyStreamsForServer) {
result.Err = fmt.Errorf("server boundary returned %v", result.ServerBoundaryError)
return result
}
result.ServerRejected = 1
if err := h.CloseStream("probe-user-0"); err != nil {
result.Err = fmt.Errorf("close stale user slot: %w", err)
return result
}
if err := h.CreateStream("probe-user-reused", 101, 501); err != nil {
result.Err = fmt.Errorf("reuse stale user slot: %w", err)
return result
}
result.UserSlotReused = true
if err := h.CreateStream("probe-wait", 0, 502); err != nil {
result.Err = fmt.Errorf("create cancellation probe stream: %w", err)
return result
}
waitStream, err := h.GetStream("probe-wait")
if err != nil {
result.Err = fmt.Errorf("get cancellation probe stream: %w", err)
return result
}
waitResult := make(chan bool, 1)
go func() {
_, ok := h.WaitForAgent(ctx, "probe-wait", 30*time.Second)
waitResult <- ok
}()
select {
case <-waitStream.waitStartedCh:
case <-ctx.Done():
result.Err = ctx.Err()
return result
}
if err := h.CloseStream("probe-wait"); err != nil {
result.Err = fmt.Errorf("close cancellation probe stream: %w", err)
return result
}
select {
case ok := <-waitResult:
result.WaitForAgentWokeOnClose = !ok
case <-ctx.Done():
result.Err = ctx.Err()
return result
}
if !result.WaitForAgentWokeOnClose {
result.Err = errors.New("WaitForAgent did not wake after stream close")
return result
}
for i := 0; i < maxStreamsPerUser; i++ {
if err := h.CloseStream(fmt.Sprintf("probe-user-%d", i)); err != nil {
result.Err = fmt.Errorf("close user stream %d: %w", i, err)
return result
}
if err := h.CloseStream(fmt.Sprintf("probe-user-%d", i)); err != nil {
result.Err = fmt.Errorf("repeat close user stream %d: %w", i, err)
return result
}
}
if err := h.CloseStream("probe-user-reused"); err != nil {
result.Err = fmt.Errorf("close reused user slot: %w", err)
return result
}
for i := 0; i < maxStreamsPerServer; i++ {
if err := h.CloseStream(fmt.Sprintf("probe-server-%d", i)); err != nil {
result.Err = fmt.Errorf("close server stream %d: %w", i, err)
return result
}
}
if err := ctx.Err(); err != nil {
result.Err = err
return result
}
h.ioStreamMutex.RLock()
result.TrackedStreams = len(h.ioStreams)
h.ioStreamMutex.RUnlock()
return result
}
service/rpc/io_stream_quota_agentcompat_test.go
+77
View File
@@ -0,0 +1,77 @@
//go:build agentcompat
package rpc
import (
"context"
"errors"
"fmt"
"sync"
"testing"
)
func TestAgentCompatIOStreamQuotaProbe(t *testing.T) {
result := RunIOStreamQuotaProbe(context.Background())
if result.Err != nil {
t.Fatalf("quota probe failed: %v", result.Err)
}
if result.UserAccepted != maxStreamsPerUser || result.UserRejected != 1 {
t.Fatalf("unexpected user boundary counts: accepted=%d rejected=%d", result.UserAccepted, result.UserRejected)
}
if result.ServerAccepted != maxStreamsPerServer || result.ServerRejected != 1 {
t.Fatalf("unexpected server boundary counts: accepted=%d rejected=%d", result.ServerAccepted, result.ServerRejected)
}
if result.TrackedStreams != 0 {
t.Fatalf("probe left tracked streams: %d", result.TrackedStreams)
}
if !result.WaitForAgentWokeOnClose {
t.Fatal("probe did not prove WaitForAgent wakes after real stream close")
}
if !result.UserSlotReused {
t.Fatal("probe did not prove a released user slot was reusable")
}
}
func TestAgentCompatIOStreamQuotaProbeUsesProductionSeam(t *testing.T) {
result := RunIOStreamQuotaProbe(context.Background())
if !errors.Is(result.UserBoundaryError, ErrTooManyStreamsForUser) {
t.Fatalf("user rejection must preserve production error, got %v", result.UserBoundaryError)
}
if !errors.Is(result.ServerBoundaryError, ErrTooManyStreamsForServer) {
t.Fatalf("server rejection must preserve production error, got %v", result.ServerBoundaryError)
}
}
func TestAgentCompatIOStreamQuotaProbeConcurrentBoundaryCalls(t *testing.T) {
h := NewNezhaHandler()
const userID, serverID = uint64(701), uint64(901)
var wg sync.WaitGroup
results := make(chan error, maxStreamsPerUser+1)
for i := 0; i < maxStreamsPerUser+1; i++ {
wg.Add(1)
go func(index int) {
defer wg.Done()
results <- h.CreateStream(fmt.Sprintf("concurrent-user-%d", index), userID, serverID+uint64(index))
}(i)
}
wg.Wait()
close(results)
accepted, rejected := 0, 0
for err := range results {
if err == nil {
accepted++
continue
}
if errors.Is(err, ErrTooManyStreamsForUser) {
rejected++
continue
}
t.Fatalf("unexpected concurrent boundary error: %v", err)
}
if accepted != maxStreamsPerUser || rejected != 1 {
t.Fatalf("unexpected concurrent boundary counts: accepted=%d rejected=%d", accepted, rejected)
}
for i := 0; i < maxStreamsPerUser+1; i++ {
_ = h.CloseStream(fmt.Sprintf("concurrent-user-%d", i))
}
}
service/rpc/io_stream_quota_test.go
+106
View File
@@ -0,0 +1,106 @@
package rpc
import (
"errors"
"fmt"
"testing"
)
func TestCreateStreamExactUserBoundary(t *testing.T) {
h := NewNezhaHandler()
for i := 0; i < maxStreamsPerUser; i++ {
if err := h.CreateStream(fmt.Sprintf("quota-user-%d", i), 1, uint64(i+1)); err != nil {
t.Fatalf("20th user stream must succeed: %v", err)
}
}
if err := h.CreateStream("quota-user-21", 1, 100); !errors.Is(err, ErrTooManyStreamsForUser) {
t.Fatalf("21st user stream must be rejected: %v", err)
}
}
func TestCreateStreamNormalUserEverydayUseSucceeds(t *testing.T) {
h := NewNezhaHandler()
if err := h.CreateStream("term", 7, 1); err != nil {
t.Fatal(err)
}
if err := h.CreateStream("fm", 7, 1); err != nil {
t.Fatal(err)
}
}
func TestCreateStreamNormalUsersAreIndependent(t *testing.T) {
h := NewNezhaHandler()
for userID := uint64(1); userID <= 5; userID++ {
for i := 0; i < maxStreamsPerUser; i++ {
if err := h.CreateStream(fmt.Sprintf("independent-%d-%d", userID, i), userID, 100+userID); err != nil {
t.Fatalf("user %d stream %d: %v", userID, i, err)
}
}
}
}
func TestCreateStreamExemptsInternalStreamsFromPerUserCap(t *testing.T) {
h := NewNezhaHandler()
for i := 0; i < maxStreamsPerUser*3; i++ {
if err := h.CreateStream(fmt.Sprintf("internal-user-%d", i), 0, uint64(i+1)); err != nil {
t.Fatal(err)
}
}
}
func TestCreateStreamInternalStreamsStillCountTowardPerServerCap(t *testing.T) {
h := NewNezhaHandler()
for i := 0; i < maxStreamsPerServer; i++ {
if err := h.CreateStream(fmt.Sprintf("internal-server-%d", i), 0, 9); err != nil {
t.Fatal(err)
}
}
if err := h.CreateStream("internal-server-over", 0, 9); !errors.Is(err, ErrTooManyStreamsForServer) {
t.Fatal(err)
}
}
func TestCreateStreamExactServerBoundary(t *testing.T) {
h := NewNezhaHandler()
for i := 0; i < maxStreamsPerServer; i++ {
if err := h.CreateStream(fmt.Sprintf("quota-server-%d", i), uint64(i+1), 2); err != nil {
t.Fatalf("40th server stream must succeed: %v", err)
}
}
if err := h.CreateStream("quota-server-41", 100, 2); !errors.Is(err, ErrTooManyStreamsForServer) {
t.Fatalf("41st server stream must be rejected: %v", err)
}
}
func TestCreateStreamReleasesUserAndServerSlots(t *testing.T) {
h := NewNezhaHandler()
for i := 0; i < maxStreamsPerUser; i++ {
if err := h.CreateStream(fmt.Sprintf("reuse-user-%d", i), 1, uint64(i+10)); err != nil {
t.Fatalf("user setup stream %d failed: %v", i, err)
}
}
if !errors.Is(h.CreateStream("reuse-user-over", 1, 100), ErrTooManyStreamsForUser) {
t.Fatal("user cap was not enforced")
}
if err := h.CloseStream("reuse-user-0"); err != nil {
t.Fatal(err)
}
if err := h.CreateStream("reuse-user-new", 1, 101); err != nil {
t.Fatalf("closed user slot must be reusable: %v", err)
}
for i := 0; i < maxStreamsPerServer; i++ {
if err := h.CreateStream(fmt.Sprintf("reuse-server-%d", i), uint64(i+2), 2); err != nil {
t.Fatalf("server setup stream %d failed: %v", i, err)
}
}
if !errors.Is(h.CreateStream("reuse-server-over", 100, 2), ErrTooManyStreamsForServer) {
t.Fatal("server cap was not enforced")
}
if err := h.CloseStream("reuse-server-0"); err != nil {
t.Fatal(err)
}
if err := h.CreateStream("reuse-server-new", 100, 2); err != nil {
t.Fatalf("closed server slot must be reusable: %v", err)
}
}
service/rpc/io_stream_race_test.go
+92
View File
@@ -0,0 +1,92 @@
package rpc
import (
"io"
"sync"
"testing"
"time"
)
// nopRWC is a minimal io.ReadWriteCloser used for race tests; Close is a
// no-op so the racer goroutines do not panic on shared state.
type nopRWC struct{}
func (nopRWC) Read(p []byte) (int, error) { return 0, io.EOF }
func (nopRWC) Write(p []byte) (int, error) { return len(p), nil }
func (nopRWC) Close() error { return nil }
// H10 regression: UserConnected/AgentConnected mutate stream.userIo /
// stream.agentIo without holding ioStreamMutex, while WaitForAgent /
// RevokeStreamsForPurpose / RevokeStreamsForServer read & close the same
// fields under the lock. The go race detector catches it deterministically
// under -race; without the fix this test fails.
func TestIOStream_AgentConnectedIsRaceFreeUnderLock(t *testing.T) {
h := NewNezhaHandler()
const streamId = "race-test"
h.CreateStream(streamId, 1, 1)
t.Cleanup(func() { _ = h.CloseStream(streamId) })
var wg sync.WaitGroup
wg.Add(3)
go func() {
defer wg.Done()
// repeatedly attach an agent
for i := 0; i < 200; i++ {
_ = h.AgentConnected(streamId, nopRWC{})
}
}()
go func() {
defer wg.Done()
// concurrently attach a user
for i := 0; i < 200; i++ {
_ = h.UserConnected(streamId, nopRWC{})
}
}()
go func() {
defer wg.Done()
// Revoker takes the write lock and reads the same userIo/agentIo
// fields the unsynchronised writers above are setting. Use a real
// targetServerID (1) so RevokeStreamsForServer actually inspects
// the entry's userIo/agentIo before deleting.
for i := 0; i < 200; i++ {
h.RevokeStreamsForServer(1)
h.CreateStream(streamId, 1, 1)
}
}()
wg.Wait()
}
// StartStream reads stream.userIo/agentIo while it waits for both endpoints.
// Those reads must be lock-protected against the concurrent writes done by
// UserConnected/AgentConnected; otherwise -race flags the data race on the
// interface fields.
func TestIOStream_StartStreamReadsAreRaceFree(t *testing.T) {
h := NewNezhaHandler()
const streamId = "startstream-race"
h.CreateStream(streamId, 2, 2)
t.Cleanup(func() { _ = h.CloseStream(streamId) })
var wg sync.WaitGroup
wg.Add(3)
go func() {
defer wg.Done()
_ = h.StartStream(streamId, 50*time.Millisecond)
}()
go func() {
defer wg.Done()
time.Sleep(5 * time.Millisecond)
_ = h.AgentConnected(streamId, nopRWC{})
}()
go func() {
defer wg.Done()
time.Sleep(5 * time.Millisecond)
_ = h.UserConnected(streamId, nopRWC{})
}()
wg.Wait()
}
service/rpc/io_stream_registry.go
+85
View File
@@ -0,0 +1,85 @@
package rpc
import (
"errors"
"io"
)
var ErrAgentStreamAlreadyConnected = errors.New("agent stream already connected")
func (s *NezhaHandler) IsStreamAuthorizedForAgent(streamId string, agentServerID uint64) bool {
s.ioStreamMutex.RLock()
defer s.ioStreamMutex.RUnlock()
ctx, ok := s.ioStreams[streamId]
return ok && ctx.targetServerID != 0 && ctx.targetServerID == agentServerID
}
func (s *NezhaHandler) IsStreamAuthorizedForUser(streamId string, userID uint64, isAdmin bool) bool {
creator, found := s.StreamOwnership(streamId)
return found && (isAdmin || creator == userID)
}
func (s *NezhaHandler) StreamOwnership(streamId string) (uint64, bool) {
s.ioStreamMutex.RLock()
defer s.ioStreamMutex.RUnlock()
ctx, ok := s.ioStreams[streamId]
if !ok {
return 0, false
}
return ctx.creatorUserID, true
}
func (s *NezhaHandler) StreamTarget(streamId string) (uint64, bool) {
s.ioStreamMutex.RLock()
defer s.ioStreamMutex.RUnlock()
ctx, ok := s.ioStreams[streamId]
if !ok {
return 0, false
}
return ctx.targetServerID, true
}
func (s *NezhaHandler) GetStream(streamId string) (*ioStreamContext, error) {
s.ioStreamMutex.RLock()
defer s.ioStreamMutex.RUnlock()
if ctx, ok := s.ioStreams[streamId]; ok {
return ctx, nil
}
return nil, errors.New("stream not found")
}
func (s *NezhaHandler) UserConnected(streamId string, userIo io.ReadWriteCloser) error {
s.ioStreamMutex.Lock()
stream, ok := s.ioStreams[streamId]
if !ok {
s.ioStreamMutex.Unlock()
return errors.New("stream not found")
}
stream.userIo = userIo
s.ioStreamMutex.Unlock()
stream.userIoChOnce.Do(func() { close(stream.userIoConnectCh) })
return nil
}
func (s *NezhaHandler) AgentConnected(streamId string, agentIo io.ReadWriteCloser) error {
s.ioStreamMutex.Lock()
stream, ok := s.ioStreams[streamId]
if !ok {
s.ioStreamMutex.Unlock()
return errors.Join(errors.New("stream not found"), agentIo.Close())
}
if stream.agentIo != nil {
s.ioStreamMutex.Unlock()
return errors.Join(ErrAgentStreamAlreadyConnected, agentIo.Close())
}
stream.agentIo = agentIo
s.ioStreamMutex.Unlock()
stream.agentIoChOnce.Do(func() { close(stream.agentIoConnectCh) })
return nil
}
func (s *NezhaHandler) streamEndpoints(stream *ioStreamContext) (io.ReadWriteCloser, io.ReadWriteCloser) {
s.ioStreamMutex.RLock()
defer s.ioStreamMutex.RUnlock()
return stream.userIo, stream.agentIo
}
service/rpc/io_stream_rpc.go
+58
View File
@@ -0,0 +1,58 @@
package rpc
import (
"fmt"
"log"
"time"
"github.com/nezhahq/nezha/pkg/grpcx"
pb "github.com/nezhahq/nezha/proto"
)
func (s *NezhaHandler) IOStream(stream pb.NezhaService_IOStreamServer) error {
clientID, err := s.Auth.Check(stream.Context())
if err != nil {
return err
}
id, err := stream.Recv()
if err != nil {
return err
}
if id == nil || !isValidIOStreamMagic(id.Data) {
return fmt.Errorf("invalid stream id")
}
streamID := string(id.Data[4:])
if !s.IsStreamAuthorizedForAgent(streamID, clientID) {
return fmt.Errorf("stream not authorized for agent")
}
if _, err := s.GetStream(streamID); err != nil {
return err
}
wrapper := grpcx.NewIOStreamWrapper(stream)
keepaliveDone := make(chan struct{})
go func() {
defer close(keepaliveDone)
ticker := time.NewTicker(30 * time.Second)
defer ticker.Stop()
for {
select {
case <-wrapper.Context().Done():
return
case <-wrapper.Done():
return
case <-ticker.C:
if err := wrapper.SendKeepalive(); err != nil {
log.Printf("NEZHA>> IOStream keepAlive error: %v\n", err)
return
}
}
}
}()
if err := s.AgentConnected(streamID, wrapper); err != nil {
_ = wrapper.Close()
return err
}
wrapper.Wait()
<-keepaliveDone
return nil
}
service/rpc/io_stream_state.go
+95
View File
@@ -0,0 +1,95 @@
package rpc
import (
"context"
"errors"
)
var ErrInvalidIOStreamStateExpectation = errors.New("invalid IOStream state expectation")
type IOStreamState struct {
Count int `json:"count"`
Generation uint64 `json:"generation"`
}
type IOStreamStateExpectation struct {
// A pointer distinguishes an omitted count from an explicit zero count.
ExpectedCount *int `json:"expected_count,omitempty"`
PresentStreamID string `json:"present_stream_id,omitempty"`
AbsentStreamID string `json:"absent_stream_id,omitempty"`
}
func ExpectedIOStreamCount(count int) *int {
return &count
}
func (s IOStreamStateExpectation) validate() error {
if s.ExpectedCount == nil && s.PresentStreamID == "" && s.AbsentStreamID == "" {
return ErrInvalidIOStreamStateExpectation
}
if s.ExpectedCount != nil && *s.ExpectedCount < 0 {
return ErrInvalidIOStreamStateExpectation
}
if s.PresentStreamID != "" && s.PresentStreamID == s.AbsentStreamID {
return ErrInvalidIOStreamStateExpectation
}
return nil
}
func (s *NezhaHandler) SnapshotIOStreamState() IOStreamState {
s.ioStreamMutex.RLock()
defer s.ioStreamMutex.RUnlock()
return s.snapshotIOStreamStateLocked()
}
func (s *NezhaHandler) snapshotIOStreamStateLocked() IOStreamState {
return IOStreamState{Count: len(s.ioStreams), Generation: s.ioStreamGeneration}
}
func (s *NezhaHandler) ioStreamStateExpectationSatisfiedLocked(expectation IOStreamStateExpectation) bool {
if expectation.ExpectedCount != nil && len(s.ioStreams) != *expectation.ExpectedCount {
return false
}
if expectation.PresentStreamID != "" {
if _, exists := s.ioStreams[expectation.PresentStreamID]; !exists {
return false
}
}
if expectation.AbsentStreamID != "" {
if _, exists := s.ioStreams[expectation.AbsentStreamID]; exists {
return false
}
}
return true
}
func (s *NezhaHandler) WaitForIOStreamState(ctx context.Context, expectation IOStreamStateExpectation) (IOStreamState, error) {
if err := expectation.validate(); err != nil {
return IOStreamState{}, err
}
for {
s.ioStreamMutex.RLock()
notify := s.ioStreamNotify
state := s.snapshotIOStreamStateLocked()
satisfied := s.ioStreamStateExpectationSatisfiedLocked(expectation)
observer := s.ioStreamWaitLockedHook
s.ioStreamMutex.RUnlock()
if observer != nil {
observer()
}
if satisfied {
return state, nil
}
select {
case <-ctx.Done():
return IOStreamState{}, ctx.Err()
case <-notify:
}
}
}
func (s *NezhaHandler) publishIOStreamStateChangeLocked() {
s.ioStreamGeneration++
close(s.ioStreamNotify)
s.ioStreamNotify = make(chan struct{})
}
service/rpc/io_stream_state_agentcompat.go
+11
View File
@@ -0,0 +1,11 @@
//go:build agentcompat
package rpc
// SetIOStreamStateWaitObserverForAgentcompat installs a deterministic harness
// seam for observing that a waiter captured its notification channel.
func (s *NezhaHandler) SetIOStreamStateWaitObserverForAgentcompat(observer func()) {
s.ioStreamMutex.Lock()
defer s.ioStreamMutex.Unlock()
s.ioStreamWaitLockedHook = observer
}
service/rpc/io_stream_state_expectation_test.go
+142
View File
@@ -0,0 +1,142 @@
package rpc
import (
"context"
"errors"
"strings"
"testing"
)
func TestWaitForIOStreamStateRejectsZeroValueExpectation(t *testing.T) {
handler := NewNezhaHandler()
if _, err := handler.WaitForIOStreamState(context.Background(), IOStreamStateExpectation{}); !errors.Is(err, ErrInvalidIOStreamStateExpectation) {
t.Fatalf("zero-value expectation error: %v", err)
}
}
func TestWaitForIOStreamStateAcceptsExplicitZeroCount(t *testing.T) {
handler := NewNezhaHandler()
state, err := handler.WaitForIOStreamState(context.Background(), IOStreamStateExpectation{ExpectedCount: ExpectedIOStreamCount(0)})
if err != nil {
t.Fatal(err)
}
if state != (IOStreamState{}) {
t.Fatalf("explicit zero state: %+v", state)
}
}
func TestWaitForIOStreamStateAcceptsPresentOnlyExpectation(t *testing.T) {
handler := NewNezhaHandler()
if err := handler.CreateStream("present-only", 1, 1); err != nil {
t.Fatal(err)
}
state, err := handler.WaitForIOStreamState(context.Background(), IOStreamStateExpectation{PresentStreamID: "present-only"})
if err != nil {
t.Fatal(err)
}
if state.Count != 1 || state.Generation != 1 {
t.Fatalf("present-only state: %+v", state)
}
}
func TestWaitForIOStreamStateRejectsSamePresentAndAbsentID(t *testing.T) {
handler := NewNezhaHandler()
_, err := handler.WaitForIOStreamState(context.Background(), IOStreamStateExpectation{PresentStreamID: "same", AbsentStreamID: "same"})
if !errors.Is(err, ErrInvalidIOStreamStateExpectation) {
t.Fatalf("same identity expectation error: %v", err)
}
}
func TestWaitForIOStreamStateRequiresAllSpecifiedConditions(t *testing.T) {
handler := NewNezhaHandler()
if err := handler.CreateStream("present", 1, 1); err != nil {
t.Fatal(err)
}
if err := handler.CreateStream("other", 1, 2); err != nil {
t.Fatal(err)
}
if err := handler.CreateStream("absent", 1, 3); err != nil {
t.Fatal(err)
}
ctx, cancel := context.WithCancel(context.Background())
cancel()
_, err := handler.WaitForIOStreamState(ctx, IOStreamStateExpectation{
ExpectedCount: ExpectedIOStreamCount(2),
PresentStreamID: "present",
AbsentStreamID: "absent",
})
if !errors.Is(err, context.Canceled) {
t.Fatalf("combined expectation cancellation: %v", err)
}
}
func TestWaitForIOStreamStateAbsenceOnlyIgnoresUnrelatedStreams(t *testing.T) {
handler := NewNezhaHandler()
if err := handler.CreateStream("unrelated", 1, 1); err != nil {
t.Fatal(err)
}
state, err := handler.WaitForIOStreamState(context.Background(), IOStreamStateExpectation{AbsentStreamID: "absent"})
if err != nil {
t.Fatal(err)
}
if state.Count != 1 || state.Generation != 1 {
t.Fatalf("absence-only state: %+v", state)
}
}
func TestWaitForIOStreamStateRejectsNegativeCountWithoutPrivateID(t *testing.T) {
handler := NewNezhaHandler()
_, err := handler.WaitForIOStreamState(context.Background(), IOStreamStateExpectation{ExpectedCount: ExpectedIOStreamCount(-1), AbsentStreamID: "private-stream-id"})
if !errors.Is(err, ErrInvalidIOStreamStateExpectation) {
t.Fatalf("negative expectation error: %v", err)
}
if err != nil && strings.Contains(err.Error(), "private-stream-id") {
t.Fatalf("private stream ID leaked: %v", err)
}
}
func TestWaitForIOStreamStateRequiresCombinedCountAndAbsence(t *testing.T) {
handler := NewNezhaHandler()
if err := handler.CreateStream("present", 1, 1); err != nil {
t.Fatal(err)
}
state, err := handler.WaitForIOStreamState(context.Background(), IOStreamStateExpectation{ExpectedCount: ExpectedIOStreamCount(1), AbsentStreamID: "absent"})
if err != nil {
t.Fatal(err)
}
if state.Count != 1 || state.Generation != 1 {
t.Fatalf("combined expectation state: %+v", state)
}
}
func TestWaitForIOStreamStateCancellationRemainsValidForUnsatisfiedExpectation(t *testing.T) {
handler := NewNezhaHandler()
ctx, cancel := context.WithCancel(context.Background())
cancel()
if _, err := handler.WaitForIOStreamState(ctx, IOStreamStateExpectation{ExpectedCount: ExpectedIOStreamCount(1)}); !errors.Is(err, context.Canceled) {
t.Fatalf("cancel error: %v", err)
}
}
func TestWaitForIOStreamStateAlreadySatisfiedReturnsSnapshot(t *testing.T) {
handler := NewNezhaHandler()
state, err := handler.WaitForIOStreamState(context.Background(), IOStreamStateExpectation{ExpectedCount: ExpectedIOStreamCount(0)})
if err != nil {
t.Fatal(err)
}
if state != (IOStreamState{}) {
t.Fatalf("already satisfied state: %+v", state)
}
}
func TestWaitForIOStreamStateRejectsInvalidAndHonorsCancellation(t *testing.T) {
handler := NewNezhaHandler()
if _, err := handler.WaitForIOStreamState(context.Background(), IOStreamStateExpectation{ExpectedCount: ExpectedIOStreamCount(-1)}); !errors.Is(err, ErrInvalidIOStreamStateExpectation) {
t.Fatalf("invalid count error: %v", err)
}
ctx, cancel := context.WithCancel(context.Background())
cancel()
if _, err := handler.WaitForIOStreamState(ctx, IOStreamStateExpectation{ExpectedCount: ExpectedIOStreamCount(1)}); !errors.Is(err, context.Canceled) {
t.Fatalf("cancel error: %v", err)
}
}
service/rpc/io_stream_state_observer_agentcompat_test.go
+111
View File
@@ -0,0 +1,111 @@
//go:build agentcompat
package rpc
import (
"context"
"sync"
"testing"
"time"
)
func TestWaitForIOStreamStateObserverCanReenterWriteLock(t *testing.T) {
handler := NewNezhaHandler()
observerCalled := make(chan struct{})
var observerOnce sync.Once
handler.SetIOStreamStateWaitObserverForAgentcompat(func() {
observerOnce.Do(func() {
if err := handler.CreateStreamWithPurpose("observer-reentrant", 1, 1, PurposeLegacy); err != nil {
t.Errorf("observer create: %v", err)
}
close(observerCalled)
})
})
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
defer cancel()
result := make(chan error, 1)
go func() {
_, err := handler.WaitForIOStreamState(ctx, IOStreamStateExpectation{PresentStreamID: "observer-reentrant"})
result <- err
}()
select {
case <-observerCalled:
case <-ctx.Done():
t.Fatalf("observer remained blocked by read lock: %v", ctx.Err())
}
select {
case err := <-result:
if err != nil {
t.Fatalf("wait after reentrant observer: %v", err)
}
case <-ctx.Done():
t.Fatalf("wait did not observe observer mutation: %v", ctx.Err())
}
}
func TestWaitForIOStreamStateObserverMutationWakesCapturedNotification(t *testing.T) {
handler := NewNezhaHandler()
observerCalled := make(chan struct{})
var observerOnce sync.Once
handler.SetIOStreamStateWaitObserverForAgentcompat(func() {
observerOnce.Do(func() {
if err := handler.CreateStreamWithPurpose("observer-wake", 1, 1, PurposeLegacy); err != nil {
t.Errorf("observer create: %v", err)
}
close(observerCalled)
})
})
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
defer cancel()
state, err := handler.WaitForIOStreamState(ctx, IOStreamStateExpectation{ExpectedCount: ExpectedIOStreamCount(1)})
if err != nil {
t.Fatalf("wait after observer mutation: %v", err)
}
select {
case <-observerCalled:
default:
t.Fatal("observer did not run")
}
if state.Count != 1 {
t.Fatalf("observer mutation state: %+v", state)
}
}
func TestWaitForIOStreamStateNoOpObserverKeepsMutationBetweenSnapshotAndSelect(t *testing.T) {
handler := NewNezhaHandler()
observerCalled := make(chan struct{})
var observerOnce sync.Once
handler.SetIOStreamStateWaitObserverForAgentcompat(func() {
observerOnce.Do(func() { close(observerCalled) })
})
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
defer cancel()
result := make(chan IOStreamState, 1)
resultErr := make(chan error, 1)
go func() {
state, err := handler.WaitForIOStreamState(ctx, IOStreamStateExpectation{ExpectedCount: ExpectedIOStreamCount(1)})
if err != nil {
resultErr <- err
return
}
result <- state
}()
select {
case <-observerCalled:
case <-ctx.Done():
t.Fatalf("observer did not run: %v", ctx.Err())
}
if err := handler.CreateStreamWithPurpose("observer-noop-create", 1, 1, PurposeLegacy); err != nil {
t.Fatalf("mutation between snapshot and select: %v", err)
}
select {
case state := <-result:
if state.Count != 1 {
t.Fatalf("mutation state: %+v", state)
}
case err := <-resultErr:
t.Fatalf("wait after no-op observer mutation: %v", err)
case <-ctx.Done():
t.Fatalf("wait missed mutation after no-op observer: %v", ctx.Err())
}
}
service/rpc/io_stream_state_test.go
+68
View File
@@ -0,0 +1,68 @@
package rpc
import (
"errors"
"testing"
)
func TestIOStreamStateSnapshotAndGeneration(t *testing.T) {
handler := NewNezhaHandler()
initial := handler.SnapshotIOStreamState()
if initial.Count != 0 || initial.Generation != 0 {
t.Fatalf("unexpected initial state: %+v", initial)
}
if err := handler.CreateStream("state-stream", 1, 1); err != nil {
t.Fatal(err)
}
created := handler.SnapshotIOStreamState()
if created.Count != 1 || created.Generation != 1 {
t.Fatalf("unexpected created state: %+v", created)
}
if err := handler.CreateStream("state-stream", 2, 2); !errors.Is(err, ErrStreamAlreadyExists) {
t.Fatalf("duplicate create error: %v", err)
}
if got := handler.SnapshotIOStreamState(); got != created {
t.Fatalf("duplicate create changed state: %+v", got)
}
if err := handler.CloseStream("unknown"); err != nil {
t.Fatal(err)
}
if got := handler.SnapshotIOStreamState(); got != created {
t.Fatalf("unknown close changed state: %+v", got)
}
}
func TestIOStreamStateRevocationPublishesOncePerBatch(t *testing.T) {
handler := NewNezhaHandler()
if err := handler.CreateStreamWithPurpose("purpose-a", 0, 1, PurposeMCPTransfer); err != nil {
t.Fatal(err)
}
if err := handler.CreateStreamWithPurpose("purpose-b", 0, 1, PurposeMCPTransfer); err != nil {
t.Fatal(err)
}
if err := handler.CreateStream("server-a", 0, 2); err != nil {
t.Fatal(err)
}
before := handler.SnapshotIOStreamState()
if revoked := handler.RevokeStreamsForPurpose(PurposeMCPTransfer); revoked != 2 {
t.Fatalf("revoked purpose streams: %d", revoked)
}
afterPurpose := handler.SnapshotIOStreamState()
if afterPurpose.Generation != before.Generation+1 || afterPurpose.Count != 1 {
t.Fatalf("purpose revocation state: before=%+v after=%+v", before, afterPurpose)
}
if revoked := handler.RevokeStreamsForPurpose(PurposeMCPTransfer); revoked != 0 {
t.Fatalf("repeat purpose revocation: %d", revoked)
}
if got := handler.SnapshotIOStreamState(); got != afterPurpose {
t.Fatalf("empty purpose revocation changed state: %+v", got)
}
handler.RevokeStreamsForServer(2)
if got := handler.SnapshotIOStreamState(); got.Generation != afterPurpose.Generation+1 || got.Count != 0 {
t.Fatalf("server revocation state: %+v", got)
}
handler.RevokeStreamsForServer(2)
if got := handler.SnapshotIOStreamState(); got.Generation != afterPurpose.Generation+1 {
t.Fatalf("empty server revocation changed generation: %+v", got)
}
}
service/rpc/io_stream_state_wait_test.go
+220
View File
@@ -0,0 +1,220 @@
package rpc
import (
"context"
"sync"
"testing"
"time"
)
func TestWaitForIOStreamStateWakesOnCloseAndAbsence(t *testing.T) {
handler := NewNezhaHandler()
if err := handler.CreateStream("wait-state", 1, 1); err != nil {
t.Fatal(err)
}
result := make(chan IOStreamState, 1)
go func() {
state, err := handler.WaitForIOStreamState(context.Background(), IOStreamStateExpectation{ExpectedCount: ExpectedIOStreamCount(0), AbsentStreamID: "wait-state"})
if err != nil {
t.Errorf("wait failed: %v", err)
return
}
result <- state
}()
if err := handler.CloseStream("wait-state"); err != nil {
t.Fatal(err)
}
state := <-result
if state.Count != 0 || state.Generation != 2 {
t.Fatalf("unexpected waited state: %+v", state)
}
}
func TestWaitForIOStreamStateCreateWakeUsesCapturedNotification(t *testing.T) {
handler := NewNezhaHandler()
waitReady := make(chan struct{})
handler.ioStreamWaitLockedHook = func() {
select {
case <-waitReady:
default:
close(waitReady)
}
}
result := make(chan IOStreamState, 1)
go func() {
state, err := handler.WaitForIOStreamState(context.Background(), IOStreamStateExpectation{ExpectedCount: ExpectedIOStreamCount(1)})
if err == nil {
result <- state
}
}()
select {
case <-waitReady:
case <-time.After(time.Second):
t.Fatal("waiter did not capture its notification channel")
}
if err := handler.CreateStream("create-wake", 1, 1); err != nil {
t.Fatal(err)
}
select {
case state := <-result:
if state.Count != 1 || state.Generation != 1 {
t.Fatalf("unexpected created state: %+v", state)
}
case <-time.After(time.Second):
t.Fatal("create did not wake waiter")
}
}
func TestWaitForIOStreamStateCloseWakeUsesCapturedNotification(t *testing.T) {
handler := NewNezhaHandler()
if err := handler.CreateStream("close-wake", 1, 1); err != nil {
t.Fatal(err)
}
waitReady := make(chan struct{})
handler.ioStreamWaitLockedHook = func() {
select {
case <-waitReady:
default:
close(waitReady)
}
}
result := make(chan IOStreamState, 1)
go func() {
state, err := handler.WaitForIOStreamState(context.Background(), IOStreamStateExpectation{ExpectedCount: ExpectedIOStreamCount(0), AbsentStreamID: "close-wake"})
if err == nil {
result <- state
}
}()
select {
case <-waitReady:
case <-time.After(time.Second):
t.Fatal("waiter did not capture its notification channel")
}
if err := handler.CloseStream("close-wake"); err != nil {
t.Fatal(err)
}
select {
case state := <-result:
if state.Count != 0 || state.Generation != 2 {
t.Fatalf("unexpected closed state: %+v", state)
}
case <-time.After(time.Second):
t.Fatal("close did not wake waiter")
}
}
func TestWaitForIOStreamStateDoesNotMissMutationBetweenSnapshotAndWait(t *testing.T) {
handler := NewNezhaHandler()
hookCalled := make(chan struct{})
mutationDone := make(chan error, 1)
var hookOnce sync.Once
handler.ioStreamWaitLockedHook = func() {
hookOnce.Do(func() {
close(hookCalled)
go func() {
mutationDone <- handler.CreateStream("lost-wakeup", 1, 1)
}()
})
}
result := make(chan IOStreamState, 1)
go func() {
state, err := handler.WaitForIOStreamState(context.Background(), IOStreamStateExpectation{ExpectedCount: ExpectedIOStreamCount(1)})
if err == nil {
result <- state
}
}()
select {
case <-hookCalled:
case <-time.After(time.Second):
t.Fatal("waiter did not reach deterministic mutation seam")
}
select {
case err := <-mutationDone:
if err != nil {
t.Fatal(err)
}
case <-time.After(time.Second):
t.Fatal("mutation did not complete")
}
select {
case state := <-result:
if state.Count != 1 || state.Generation != 1 {
t.Fatalf("unexpected mutation state: %+v", state)
}
case <-time.After(time.Second):
t.Fatal("waiter missed mutation published during wait setup")
}
}
func TestWaitForIOStreamStateConcurrentCreateCloseWaiters(t *testing.T) {
handler := NewNezhaHandler()
created := make(chan IOStreamState, 1)
closed := make(chan IOStreamState, 1)
go func() {
state, err := handler.WaitForIOStreamState(context.Background(), IOStreamStateExpectation{ExpectedCount: ExpectedIOStreamCount(1)})
if err == nil {
created <- state
}
}()
if err := handler.CreateStream("concurrent", 1, 1); err != nil {
t.Fatal(err)
}
select {
case state := <-created:
if state.Count != 1 {
t.Fatalf("created waiter state: %+v", state)
}
case <-time.After(time.Second):
t.Fatal("created waiter did not wake")
}
go func() {
state, err := handler.WaitForIOStreamState(context.Background(), IOStreamStateExpectation{ExpectedCount: ExpectedIOStreamCount(0), AbsentStreamID: "concurrent"})
if err == nil {
closed <- state
}
}()
if err := handler.CloseStream("concurrent"); err != nil {
t.Fatal(err)
}
select {
case state := <-closed:
if state.Count != 0 {
t.Fatalf("closed waiter state: %+v", state)
}
case <-time.After(time.Second):
t.Fatal("closed waiter did not wake")
}
}
func TestWaitForIOStreamStateDoesNotAcceptUnrelatedSameCountForPresentID(t *testing.T) {
handler := NewNezhaHandler()
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
result := make(chan error, 1)
go func() {
_, err := handler.WaitForIOStreamState(ctx, IOStreamStateExpectation{
ExpectedCount: ExpectedIOStreamCount(1),
PresentStreamID: "wanted",
})
result <- err
}()
if err := handler.CreateStream("unrelated", 1, 1); err != nil {
t.Fatal(err)
}
select {
case err := <-result:
if err == nil {
t.Fatal("same-count unrelated stream satisfied identity expectation")
}
default:
}
cancel()
select {
case err := <-result:
if err == nil {
t.Fatal("identity waiter unexpectedly succeeded")
}
case <-time.After(time.Second):
t.Fatal("identity waiter did not observe cancellation")
}
}
service/rpc/io_stream_test.go
+162 -1
View File
@@ -12,7 +12,7 @@ func TestIOStream(t *testing.T) {
const testStreamID = "ffffffff-ffff-ffff-ffff-ffffffffffff"
handler.CreateStream(testStreamID)
handler.CreateStream(testStreamID, 0, 0)
userIo, agentIo := newPipeReadWriter(), newPipeReadWriter()
defer func() {
userIo.Close()
@@ -105,3 +105,164 @@ func newPipeReadWriter() io.ReadWriteCloser {
io.WriteCloser
}{r, w}
}
func TestStreamOwnershipReturnsCreatorUserID(t *testing.T) {
h := NewNezhaHandler()
h.CreateStream("alice-stream", 100, 0)
creator, found := h.StreamOwnership("alice-stream")
if !found {
t.Fatalf("expected stream to be found after CreateStream")
}
if creator != 100 {
t.Fatalf("expected creator user ID 100, got %d", creator)
}
}
func TestStreamOwnershipReturnsNotFoundForUnknownID(t *testing.T) {
h := NewNezhaHandler()
if _, found := h.StreamOwnership("nonexistent"); found {
t.Fatalf("expected unknown stream id to report not-found")
}
}
func TestStreamOwnershipPreservesPerStreamCreator(t *testing.T) {
h := NewNezhaHandler()
h.CreateStream("alice-stream", 100, 0)
h.CreateStream("bob-stream", 200, 0)
aliceCreator, _ := h.StreamOwnership("alice-stream")
bobCreator, _ := h.StreamOwnership("bob-stream")
if aliceCreator != 100 || bobCreator != 200 {
t.Fatalf("expected per-stream creator IDs alice=100 bob=200, got alice=%d bob=%d",
aliceCreator, bobCreator)
}
}
func TestIsStreamAuthorizedForUserAllowsCreator(t *testing.T) {
h := NewNezhaHandler()
h.CreateStream("alice-stream", 100, 0)
if !h.IsStreamAuthorizedForUser("alice-stream", 100, false) {
t.Fatalf("creator must be authorized to attach to their own stream")
}
}
func TestIsStreamAuthorizedForUserDeniesForeignMember(t *testing.T) {
h := NewNezhaHandler()
h.CreateStream("alice-stream", 100, 0)
if h.IsStreamAuthorizedForUser("alice-stream", 200, false) {
t.Fatalf("foreign member must not be authorized — session hijack would be possible")
}
}
func TestIsStreamAuthorizedForUserAllowsAdmin(t *testing.T) {
h := NewNezhaHandler()
h.CreateStream("alice-stream", 100, 0)
if !h.IsStreamAuthorizedForUser("alice-stream", 999, true) {
t.Fatalf("admin must be authorized to attach regardless of creator")
}
}
func TestIsStreamAuthorizedForUserDeniesUnknownStream(t *testing.T) {
h := NewNezhaHandler()
if h.IsStreamAuthorizedForUser("nonexistent", 100, true) {
t.Fatalf("unknown stream id must not authorize even admin")
}
}
// IOStream init messages begin with the magic marker ff05ff05. The inline
// check previously used && between byte inequalities, which due to short-
// circuit evaluation accepted almost every non-magic payload (any payload
// whose byte0 == 0xff was silently let through). These tests pin down the
// correct semantics: all four bytes must match exactly.
func TestIsValidIOStreamMagicAcceptsExactMagic(t *testing.T) {
if !isValidIOStreamMagic([]byte{0xff, 0x05, 0xff, 0x05}) {
t.Fatal("exact ff05ff05 magic must be accepted")
}
if !isValidIOStreamMagic([]byte{0xff, 0x05, 0xff, 0x05, 'p', 'a', 'y', 'l', 'o', 'a', 'd'}) {
t.Fatal("ff05ff05 followed by payload must be accepted")
}
}
func TestIsValidIOStreamMagicRejectsShortData(t *testing.T) {
if isValidIOStreamMagic([]byte{}) {
t.Fatal("empty data must be rejected")
}
if isValidIOStreamMagic([]byte{0xff, 0x05, 0xff}) {
t.Fatal("3-byte payload must be rejected")
}
}
// Agent-side stream authorization is the dual of IsStreamAuthorizedForUser:
// only the server the dashboard selected when CreateStream was called may
// attach via IOStream(). Without it, any authenticated agent that learns an
// active streamId (task-stream observation, leaked logs) can race in and
// serve a terminal/fm/NAT session originally addressed to a different
// server — a session-hijack RCE intermediation primitive.
func TestIsStreamAuthorizedForAgentAllowsBoundServer(t *testing.T) {
h := NewNezhaHandler()
h.CreateStream("terminal-for-server-100", 1, 100)
if !h.IsStreamAuthorizedForAgent("terminal-for-server-100", 100) {
t.Fatalf("the bound target server must be authorized to attach")
}
}
func TestIsStreamAuthorizedForAgentDeniesForeignServer(t *testing.T) {
h := NewNezhaHandler()
h.CreateStream("terminal-for-server-100", 1, 100)
if h.IsStreamAuthorizedForAgent("terminal-for-server-100", 200) {
t.Fatalf("a foreign agent must not be able to attach — session hijack would be possible")
}
}
func TestIsStreamAuthorizedForAgentDeniesUnboundStream(t *testing.T) {
h := NewNezhaHandler()
// targetServerID == 0 means the stream was created without a bound agent
// — no agent should be allowed to attach.
h.CreateStream("unbound-stream", 1, 0)
if h.IsStreamAuthorizedForAgent("unbound-stream", 100) {
t.Fatalf("unbound stream must not authorize any agent")
}
if h.IsStreamAuthorizedForAgent("unbound-stream", 0) {
t.Fatalf("unbound stream must not authorize a zero clientID either")
}
}
func TestIsStreamAuthorizedForAgentDeniesUnknownStreamID(t *testing.T) {
h := NewNezhaHandler()
if h.IsStreamAuthorizedForAgent("nonexistent", 100) {
t.Fatalf("unknown stream id must not authorize any agent")
}
}
func TestIsValidIOStreamMagicRejectsPartialOrWrongMagic(t *testing.T) {
// Each case has at least one byte that does NOT match the magic. The
// previous && short-circuit bug let cases like {0xff, 0, 0, 0} pass
// because byte0 alone matched. Correct semantics: any single byte off
// → reject.
cases := [][]byte{
{0x00, 0x00, 0x00, 0x00},
{0xff, 0x00, 0x00, 0x00},
{0x00, 0x05, 0x00, 0x00},
{0x00, 0x00, 0xff, 0x00},
{0x00, 0x00, 0x00, 0x05},
{0xff, 0x05, 0xff, 0x00},
{0xff, 0x05, 0x00, 0x05},
{0xff, 0x00, 0xff, 0x05},
{0x00, 0x05, 0xff, 0x05},
{0xff, 0xff, 0xff, 0xff},
}
for _, c := range cases {
if isValidIOStreamMagic(c) {
t.Fatalf("non-magic payload %v must be rejected (regression: && short-circuit bug)", c)
}
}
}
service/rpc/mcp_cancel_double_close_test.go
+55
View File
@@ -0,0 +1,55 @@
package rpc
import (
"sync"
"sync/atomic"
"testing"
pb "github.com/nezhahq/nezha/proto"
)
// updateConfig has no serialization, so two concurrent admin PATCH /setting
// requests can both flip EnableMCP true->false and both invoke
// CancelAllMCPInflight concurrently. The sweep must close each entry's cancel
// channel at most once; a non-atomic check-then-close double-closes the same
// channel and panics, crashing the dashboard.
func TestCancelAllMCPInflight_ConcurrentSweepsDoNotDoubleClose(t *testing.T) {
mcpInflight.Range(func(key, _ any) bool {
mcpInflight.Delete(key)
return true
})
t.Cleanup(func() {
mcpInflight.Range(func(key, _ any) bool {
mcpInflight.Delete(key)
return true
})
})
const entries = 256
for i := 0; i < entries; i++ {
mcpInflight.Store(uint64(i+1), &mcpInflightEntry{
serverID: uint64(i + 1),
result: make(chan *pb.TaskResult, 1),
cancel: make(chan struct{}),
cancelled: new(atomic.Bool),
})
}
const sweepers = 8
var wg sync.WaitGroup
wg.Add(sweepers)
for i := 0; i < sweepers; i++ {
go func() {
defer wg.Done()
// A double-close inside CancelAllMCPInflight panics here and
// fails the test (panic in a goroutine aborts the test binary).
CancelAllMCPInflight()
}()
}
wg.Wait()
mcpInflight.Range(func(key, _ any) bool {
t.Fatalf("inflight entry %v survived the sweep", key)
return false
})
}
service/rpc/mcp_kill_switch_race_test.go
+35
View File
@@ -0,0 +1,35 @@
package rpc
import (
"context"
"testing"
"time"
"github.com/nezhahq/nezha/model"
)
// H9 regression: CallAgent must consult mcpKillSwitchObserved before any
// side-effects. Without this gate, mcpEndpoint's EnableMCP read and
// CancelAllMCPInflight race against a fresh CallAgent that registers
// AFTER the cancel sweep, surviving the disabled state.
func TestCallAgent_RefusesWhenKillSwitchObserved(t *testing.T) {
prevCheck := mcpKillSwitchObserver()
SetMCPKillSwitchObserver(func() bool { return true })
t.Cleanup(func() { SetMCPKillSwitchObserver(prevCheck) })
_, err := CallAgent(context.Background(), 1, model.TaskTypeExec, struct{}{}, 50*time.Millisecond)
if err != ErrMCPDisabled {
t.Fatalf("CallAgent must short-circuit to ErrMCPDisabled when kill switch is observed, got %v", err)
}
}
// The default hook must keep production behaviour disarmed so tests and
// unconfigured deployments do not short-circuit CallAgent.
func TestCallAgent_KillSwitchHookDefaultsToDisarmed(t *testing.T) {
if mcpKillSwitchObserver() == nil {
t.Fatal("mcpKillSwitchObserver must always return a non-nil probe so dashboard can wire it")
}
if mcpKillSwitchObserver()() {
t.Fatal("default hook must return false so unconfigured dashboards / tests don't accidentally short-circuit CallAgent")
}
}
service/rpc/mcp_kill_switch_registration_race_test.go
+81
View File
@@ -0,0 +1,81 @@
package rpc
import (
"context"
"errors"
"sync"
"testing"
"time"
"github.com/nezhahq/nezha/model"
)
// Registration-after-sweep race (review issue #1): CancelAllMCPInflight only
// cancels entries already present in the inflight map. A CallAgent that passes
// the upfront kill-switch check but has not yet Store()d its entry is invisible
// to the sweep, so without a post-registration re-check it goes on to SendTask
// 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.
//
// With the race present, CallAgent sends the task and blocks until timeout
// (ErrAgentTimeout) — the agent received a fresh task past the kill switch.
func TestCallAgent_KillSwitchBeatsRegistrationAfterSweep(t *testing.T) {
const target uint64 = 7401
stream := newFakeStream()
cleanup := installFakeServer(t, target, stream)
defer cleanup()
var killed bool
var mu sync.Mutex
prev := mcpKillSwitchObserver()
// Observer returns the operator-controlled flag. CallAgent reads it both
// before and (with the fix) after registering the inflight entry.
SetMCPKillSwitchObserver(func() bool {
mu.Lock()
defer mu.Unlock()
return killed
})
t.Cleanup(func() { SetMCPKillSwitchObserver(prev) })
// Fail loudly if the agent ever receives a task: that means a fresh call
// leaked past the kill switch.
leaked := make(chan *struct{}, 1)
go func() {
select {
case <-stream.sent:
leaked <- nil
case <-time.After(2 * time.Second):
}
}()
// 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.
hook := func() {
mu.Lock()
killed = true
mu.Unlock()
CancelAllMCPInflight()
}
testKillSwitchAfterUpfrontCheck.Store(&hook)
t.Cleanup(func() { testKillSwitchAfterUpfrontCheck.Store(nil) })
_, err := CallAgent(context.Background(), target, model.TaskTypeExec,
model.ExecRequest{Cmd: "x"}, 1*time.Second)
if !errors.Is(err, ErrMCPDisabled) {
t.Fatalf("CallAgent must return ErrMCPDisabled when the kill switch fires during registration; got %v", err)
}
select {
case <-leaked:
t.Fatal("a fresh MCP task leaked to the agent past the kill switch")
default:
}
}
service/rpc/mcp_receipt_agentcompat_test.go
+197
View File
@@ -0,0 +1,197 @@
//go:build agentcompat
package rpc
import (
"bufio"
"context"
"errors"
"fmt"
"net"
"strconv"
"strings"
"testing"
"time"
"github.com/nezhahq/nezha/model"
pb "github.com/nezhahq/nezha/proto"
"github.com/stretchr/testify/require"
)
func TestMCPReceiptGate_FormatsTaskAndResultWithGeneration(t *testing.T) {
// Given
serverConn, clientConn := net.Pipe()
defer clientConn.Close()
gate := installReceiptGateForTest(serverConn)
defer clearReceiptGateForTest()
reader := bufio.NewReader(clientConn)
// When
taskDone := make(chan struct{})
go func() {
notifyMCPTaskDispatched(7, 9, model.TaskTypeExec)
close(taskDone)
}()
taskLine := mustReadLine(t, reader)
<-taskDone
resultDone := make(chan struct{})
go func() {
notifyMCPTaskResultAccepted(7, 9, model.TaskTypeExec)
close(resultDone)
}()
resultLine := mustReadLine(t, reader)
<-resultDone
// Then
require.Equal(t, "task "+itoa(gate.generation)+" 7 9 "+itoa(model.TaskTypeExec)+"\n", taskLine)
require.Equal(t, "result "+itoa(gate.generation)+" 7 9 "+itoa(model.TaskTypeExec)+"\n", resultLine)
}
func TestCallAgent_EmitsOneTaskAndOneAcceptedResult(t *testing.T) {
// Given
serverConn, clientConn := net.Pipe()
defer clientConn.Close()
gate := installReceiptGateForTest(serverConn)
defer clearReceiptGateForTest()
stream := newFakeStream()
cleanup := installFakeServer(t, 801, stream)
defer cleanup()
reader := bufio.NewReader(clientConn)
lines := make(chan string, 2)
go func() {
line, err := reader.ReadString('\n')
if err != nil {
return
}
lines <- line
line, err = reader.ReadString('\n')
if err == nil {
lines <- line
}
}()
go func() {
sent := <-stream.sent
deliverMCPResult(&pb.TaskResult{Id: sent.GetId(), Type: sent.GetType(), Successful: true, Data: "{}"})
deliverMCPResult(&pb.TaskResult{Id: sent.GetId(), Type: sent.GetType(), Successful: true, Data: "{}"})
}()
// When
_, err := CallAgent(context.Background(), 801, model.TaskTypeExec, model.ExecRequest{Cmd: "x"}, time.Second)
// Then
require.NoError(t, err)
taskLine := <-lines
resultLine := <-lines
require.Equal(t, "task "+itoa(gate.generation)+" 801 "+itoa(parseReceiptTaskID(taskLine))+" "+itoa(model.TaskTypeExec)+"\n", taskLine)
require.Equal(t, "result "+itoa(gate.generation)+" 801 "+itoa(parseReceiptTaskID(resultLine))+" "+itoa(model.TaskTypeExec)+"\n", resultLine)
require.Equal(t, parseReceiptTaskID(taskLine), parseReceiptTaskID(resultLine))
clientConn.SetReadDeadline(time.Now().Add(20 * time.Millisecond))
_, readErr := reader.ReadString('\n')
require.Error(t, readErr)
}
func TestCallAgent_SendFailureEmitsNoTaskReceipt(t *testing.T) {
// Given
serverConn, clientConn := net.Pipe()
defer clientConn.Close()
installReceiptGateForTest(serverConn)
defer clearReceiptGateForTest()
stream := &fakeTaskStream{sent: make(chan *pb.Task, 1), err: errors.New("send failed")}
cleanup := installFakeServer(t, 802, stream)
defer cleanup()
// When
_, err := CallAgent(context.Background(), 802, model.TaskTypeExec, model.ExecRequest{Cmd: "x"}, time.Second)
// Then
require.EqualError(t, err, "send failed")
clientConn.SetReadDeadline(time.Now().Add(20 * time.Millisecond))
_, readErr := bufio.NewReader(clientConn).ReadString('\n')
require.Error(t, readErr)
}
func TestCallAgent_LateDuplicateAndCancelledResultsEmitNoAcceptedReceipt(t *testing.T) {
// Given
serverConn, clientConn := net.Pipe()
defer clientConn.Close()
installReceiptGateForTest(serverConn)
defer clearReceiptGateForTest()
stream := newFakeStream()
cleanup := installFakeServer(t, 803, stream)
defer cleanup()
reader := bufio.NewReader(clientConn)
taskLineCh := make(chan string, 1)
// When
taskIDCh := make(chan uint64, 1)
go func() {
sent := <-stream.sent
taskIDCh <- sent.GetId()
line, _ := reader.ReadString('\n')
taskLineCh <- line
}()
_, err := CallAgent(context.Background(), 803, model.TaskTypeFsRead, model.FsReadRequest{Path: "/x"}, 20*time.Millisecond)
require.ErrorIs(t, err, ErrAgentTimeout)
taskID := <-taskIDCh
deliverMCPResult(&pb.TaskResult{Id: taskID, Type: model.TaskTypeFsRead, Successful: true, Data: "{}"})
deliverMCPResult(&pb.TaskResult{Id: taskID, Type: model.TaskTypeFsRead, Successful: true, Data: "{}"})
// Then
require.Contains(t, <-taskLineCh, "task ")
clientConn.SetReadDeadline(time.Now().Add(20 * time.Millisecond))
_, readErr := reader.ReadString('\n')
require.Error(t, readErr)
}
func TestCallAgent_CancelledResultEmitsNoAcceptedReceipt(t *testing.T) {
// Given
serverConn, clientConn := net.Pipe()
defer clientConn.Close()
installReceiptGateForTest(serverConn)
defer clearReceiptGateForTest()
stream := newFakeStream()
cleanup := installFakeServer(t, 804, stream)
defer cleanup()
reader := bufio.NewReader(clientConn)
taskLineCh := make(chan string, 1)
go func() {
line, _ := reader.ReadString('\n')
taskLineCh <- line
}()
taskIDCh := make(chan uint64, 1)
go func() {
sent := <-stream.sent
taskIDCh <- sent.GetId()
}()
// When
errCh := make(chan error, 1)
go func() {
_, err := CallAgent(context.Background(), 804, model.TaskTypeExec, model.ExecRequest{Cmd: "x"}, time.Second)
errCh <- err
}()
taskID := <-taskIDCh
CancelAllMCPInflight()
deliverMCPResult(&pb.TaskResult{Id: taskID, Type: model.TaskTypeExec, Successful: true, Data: "{}"})
// Then
require.ErrorIs(t, <-errCh, ErrMCPDisabled)
require.Contains(t, <-taskLineCh, "task ")
clientConn.SetReadDeadline(time.Now().Add(20 * time.Millisecond))
_, readErr := reader.ReadString('\n')
require.Error(t, readErr)
}
func itoa(value uint64) string {
return strconv.FormatUint(value, 10)
}
func parseReceiptTaskID(line string) uint64 {
fields := strings.Fields(line)
value, err := strconv.ParseUint(fields[3], 10, 64)
if err != nil {
panic(fmt.Sprintf("invalid receipt line %q: %v", line, err))
}
return value
}
service/rpc/mcp_rpc.go
+384
View File
@@ -0,0 +1,384 @@
package rpc
import (
"context"
"encoding/json"
"errors"
"log"
"sync"
"sync/atomic"
"time"
"github.com/nezhahq/nezha/model"
pb "github.com/nezhahq/nezha/proto"
"github.com/nezhahq/nezha/service/singleton"
)
// MCP 的"调用-响应"模式复用了 RequestTask 双向流：
// - dashboard 发 Task（带新分配的 taskID + JSON params）
// - agent 执行后回 TaskResult（同 taskID + JSON result）
// - RequestTask 接收循环把这种 TaskType 识别后路由到 inflight 等待方
//
// 不污染 model.Server 字段：用本包内的全局 inflight 表按 taskID 关联，
// 跨 server 共享单一命名空间。
var (
mcpTaskIDCounter atomic.Uint64
mcpInflight sync.Map // key: uint64 (taskID), value: chan *pb.TaskResult
)
// ErrMCPDisabled 是 CallAgent 在 MCP kill switch 被触发时返回的哨兵错误。
// 与 ErrAgentTimeout / ErrAgentOffline 平级，便于 controller 把它映射到
// MCPOutcomeForbidden 之类的审计 code 而不是误报 agent 故障。
var ErrMCPDisabled = errors.New("MCP is disabled by the dashboard administrator")
// mcpKillSwitchObserved is a process-level hook the dashboard wires to
// singleton.Conf.EnableMCP. CallAgent consults it before any side-effects so
// the entry-check / cancel-sweep / registration race cannot leak a fresh
// call past EnableMCP=false. Defaults to "disarmed" so tests and headless
// builds are unaffected.
//
// Stored behind atomic.Pointer because SetMCPKillSwitchObserver (startup +
// tests) and CallAgent (any RPC goroutine) touch it concurrently; a plain
// func variable is a data race under -race.
var mcpKillSwitchObserved atomic.Pointer[func() bool]
// disarmedKillSwitch is the default probe: never trips the kill switch.
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.
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
// restores the default disarmed hook (used by tests to undo overrides).
func SetMCPKillSwitchObserver(fn func() bool) {
if fn == nil {
mcpKillSwitchObserved.Store(&disarmedKillSwitch)
return
}
mcpKillSwitchObserved.Store(&fn)
}
// mcpKillSwitchObserver returns the currently installed probe, never nil.
func mcpKillSwitchObserver() func() bool {
if p := mcpKillSwitchObserved.Load(); p != nil {
return *p
}
return disarmedKillSwitch
}
// allocateMCPTaskID 分配下一个 MCP 用的 task ID。
// 取 1<<32 起步以与可能存在的 cron/transfer 等已有 ID 空间错开（cron.id 由
// DB 自增，常量级，不会触及 1<<32）。
func allocateMCPTaskID() uint64 {
const base uint64 = 1 << 32
v := mcpTaskIDCounter.Add(1)
return base + v
}
// CallAgent 给 serverID 对应的 agent 发一条 MCP-RPC 风格的 Task，并阻塞等待 TaskResult 回包。
//
// taskType 必须是 model.IsMCPRPCResult 返回 true 的类型；params 会被 JSON 编码进 Task.Data。
// 超时由调用方控制；触发超时后从 inflight 表移除等待 slot（晚到的回包会被丢弃）。
//
// 错误语义：
// - server 未在线 / 未连接 task stream → ErrAgentOffline
// - 超时 → ctx.Err 或 ErrAgentTimeout
// - agent 回包 successful=false → 把 result.Data 当错误字符串返回
// - CancelAllMCPInflight 期间被中断 → ErrMCPDisabled
// - 任何 send 失败、序列化失败 → 原始 error
//
// 返回的 raw JSON 是 agent 端 TaskResult.Data 的原文。
func CallAgent(ctx context.Context, serverID uint64, taskType uint64, params any, timeout time.Duration) (json.RawMessage, error) {
if !model.IsMCPRPCResult(taskType) {
return nil, errors.New("CallAgent: task type is not registered as MCP RPC")
}
killSwitch := mcpKillSwitchObserver()
if killSwitch() {
return nil, ErrMCPDisabled
}
server, _ := singleton.ServerShared.Get(serverID)
if server == nil {
return nil, ErrAgentOffline
}
if server.GetTaskStream() == nil {
return nil, ErrAgentOffline
}
body, err := json.Marshal(params)
if err != nil {
return nil, err
}
taskID := allocateMCPTaskID()
resultCh := make(chan *pb.TaskResult, 1)
cancelCh := make(chan struct{})
entry := &mcpInflightEntry{
serverID: serverID,
result: resultCh,
cancel: cancelCh,
cancelled: new(atomic.Bool),
}
if hook := testKillSwitchAfterUpfrontCheck.Load(); hook != nil {
(*hook)()
}
mcpInflight.Store(taskID, entry)
defer mcpInflight.Delete(taskID)
// Close the registration-after-sweep window: a kill switch that fired
// between the upfront check and this Store is invisible to
// CancelAllMCPInflight (our entry was not in the map yet). Because the
// operator sets EnableMCP=false BEFORE running the sweep, re-reading the
// observer here after Store guarantees we either see it disabled, or the
// sweep saw our now-registered entry and flipped entry.cancelled.
if killSwitch() || entry.cancelled.Load() {
return nil, ErrMCPDisabled
}
if err := server.SendTask(&pb.Task{
Id: taskID,
Type: taskType,
Data: string(body),
}); err != nil {
if errors.Is(err, model.ErrTaskStreamOffline) {
return nil, ErrAgentOffline
}
return nil, err
}
notifyMCPTaskDispatched(serverID, taskID, taskType)
waitCtx := ctx
var cancel context.CancelFunc
if timeout > 0 {
waitCtx, cancel = context.WithTimeout(ctx, timeout)
defer cancel()
}
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.
// Re-check the cancel flag and prefer ErrMCPDisabled, matching the
// contract documented above ("CancelAllMCPInflight 期间被中断 →
// ErrMCPDisabled") and what TestUpdateConfig_DisablingMCPInvokesKillSwitch
// expects.
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")
}
if !res.GetSuccessful() {
if res.GetData() != "" {
return nil, errors.New(res.GetData())
}
return nil, errors.New("agent returned unsuccessful result")
}
return json.RawMessage(res.GetData()), nil
case <-cancelCh:
return nil, ErrMCPDisabled
case <-waitCtx.Done():
if errors.Is(waitCtx.Err(), context.DeadlineExceeded) {
return nil, ErrAgentTimeout
}
return nil, waitCtx.Err()
}
}
// mcpInflightEntry binds an in-flight MCP call to its target serverID and
// pairs the result channel with a per-call cancel channel so the kill switch
// can break out of CallAgent without leaving the result channel dangling for
// the next late agent reply. The serverID is the authoritative reporter
// identity check at delivery time — without it deliverMCPResult would route
// purely by attacker-controlled TaskResult.Id (same bug class as commit
// 02129f1 in the cron path).
//
// 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
// panic on the second close.
func (e *mcpInflightEntry) closeCancel() {
e.closeOnce.Do(func() { close(e.cancel) })
}
// CancelAllMCPInflight closes every in-flight CallAgent so they return
// ErrMCPDisabled immediately. Used by the EnableMCP=false transition: by
// itself the inflight table holds the dashboard goroutine hostage until
// the agent replies (or the per-call timeout fires, up to ~305s for
// server.exec). Returns the number of calls cancelled for audit.
//
// Implementation notes:
// - Set the cancelled flag BEFORE closing cancelCh so any goroutine that
// already woke on resultCh observes it on the post-select re-check.
// - Delete the entry from mcpInflight immediately. Late agent replies via
// deliverMCPResult* would otherwise still find it (their own cancelled
// check covers concurrent delete, but evicting eagerly keeps the table
// small under repeated kill switch / re-enable cycles).
func CancelAllMCPInflight() int {
cancelled := 0
mcpInflight.Range(func(key, value any) bool {
entry, ok := value.(*mcpInflightEntry)
if !ok {
return true
}
entry.cancelCall()
mcpInflight.Delete(key)
cancelled++
return true
})
return cancelled
}
// DeliverMCPResultForTest 暴露 deliverMCPResult 给跨包测试用：这是显式的
// "信任路径 / 不做 reporter 校验"入口，专给不关心来源的旧测试用。
// 安全敏感测试请用 DeliverMCPResultFromReporterForTest 并传入真实 reporterID。
func DeliverMCPResultForTest(res *pb.TaskResult) { deliverMCPResult(res) }
// DeliverMCPResultFromReporterForTest 暴露带 reporter 校验的投递入口给跨包
// 测试用，与生产 RequestTask 路径同语义：reporterID 必须等于 inflight 条目
// 登记的目标 serverID 才会投递。reporterID == 0 视为 "未知 reporter" 并被
// 拒绝；要绕过 reporter 校验请改用 DeliverMCPResultForTest。
func DeliverMCPResultFromReporterForTest(res *pb.TaskResult, reporterID uint64) {
deliverMCPResultFromReporter(res, reporterID)
}
// inflightServerIDForTest 返回某个 taskID 当前挂载的目标 serverID。用于安全
// 回归测试断言 inflight 条目确实把目标 server 绑进了路由表。
// 未找到时返回 (0, false)。
func inflightServerIDForTest(taskID uint64) (uint64, bool) {
v, ok := mcpInflight.Load(taskID)
if !ok {
return 0, false
}
entry, ok := v.(*mcpInflightEntry)
if !ok {
return 0, false
}
return entry.serverID, true
}
// deliverMCPResult 把 RequestTask 收到的 MCP-RPC TaskResult 路由到等待方。
// 找不到等待 slot（已超时被移除）则丢弃。
//
// 此变体不做 reporter 校验，仅用于不关心 reporter 的内部/测试路径。生产
// RequestTask 接收循环必须走 deliverMCPResultFromReporter，把 stream 上
// 已认证的 clientID 作为 reporter 传入。
func deliverMCPResult(res *pb.TaskResult) {
if res == nil {
return
}
v, ok := mcpInflight.Load(res.GetId())
if !ok {
return
}
entry, ok := v.(*mcpInflightEntry)
if !ok {
return
}
entry.mu.Lock()
defer entry.mu.Unlock()
if entry.cancelled.Load() {
return
}
select {
case entry.result <- res:
default:
}
}
// deliverMCPResultFromReporter 是生产路径的入口：要求 reporterID 与 inflight
// 条目登记的目标 serverID 一致才投递；否则丢弃并打日志。reporterID == 0
// 视为“未知 reporter”，安全起见也丢弃。
//
// 这条校验是必要的：mcpInflight 用全局递增 taskID 做键，跨 server 共享
// 单一命名空间；如果不在投递时核对上报 agent 是 CallAgent 的目标 server，
// 任何已认证的恶意/失陷 agent 都能用猜到的 taskID 抢答其他 server 的
// MCP 调用（resultCh 容量 1，先到者覆盖真正回包）——和 commit 02129f1
// 在 cron 路径修过的攻击面同类。
func deliverMCPResultFromReporter(res *pb.TaskResult, reporterID uint64) {
if res == nil {
return
}
v, ok := mcpInflight.Load(res.GetId())
if !ok {
return
}
entry, ok := v.(*mcpInflightEntry)
if !ok {
return
}
if reporterID == 0 || entry.serverID != reporterID {
log.Printf("NEZHA>> MCP result ignored: taskID=%d targetServerID=%d reporterID=%d",
res.GetId(), entry.serverID, reporterID)
return
}
entry.mu.Lock()
defer entry.mu.Unlock()
if entry.cancelled.Load() {
return
}
select {
case entry.result <- res:
default:
}
}
// 错误类型
var (
ErrAgentOffline = errors.New("agent offline or task stream not connected")
ErrAgentTimeout = errors.New("agent did not respond within timeout")
)
service/rpc/mcp_rpc_helper_doc_test.go
+61
View File
@@ -0,0 +1,61 @@
package rpc
import (
"strings"
"sync/atomic"
"testing"
pb "github.com/nezhahq/nezha/proto"
)
// 把"测试 helper 的注释与运行时语义"钉成测试，避免 helper 文档骗读者：
//
// 1. DeliverMCPResultForTest 是显式的"信任路径 / 不做 reporter 校验"入口。
// 2. DeliverMCPResultFromReporterForTest 是带 reporter 校验的入口；
// reporterID == 0 视为"未知 reporter"，必须被拒绝，不能像旧注释暗示的
// 那样当作"未知/不校验"放行。
//
// 这条契约决定了任何安全敏感的跨包测试调用方式：要绕过 reporter，
// 必须用 DeliverMCPResultForTest，而不是 reporterID=0 通过 reporter 入口。
func TestDeliverMCPResultFromReporterForTest_ZeroReporterIDIsRejected(t *testing.T) {
taskID := allocateMCPTaskID()
resultCh := make(chan *pb.TaskResult, 1)
cancelCh := make(chan struct{})
mcpInflight.Store(taskID, &mcpInflightEntry{
serverID: 7,
result: resultCh,
cancel: cancelCh,
cancelled: new(atomic.Bool),
})
t.Cleanup(func() { mcpInflight.Delete(taskID) })
DeliverMCPResultFromReporterForTest(&pb.TaskResult{Id: taskID, Data: "x", Successful: true}, 0)
select {
case <-resultCh:
t.Fatalf("reporterID==0 must be rejected by the reporter-checked helper; expected no delivery")
default:
}
}
// 同时把"测试 helper 自身的文档约束"钉到代码里：注释必须明确说出
// "reporterID == 0 视为未知 reporter 并被拒绝"，否则未来维护者很容易看着
// "不校验"的旧措辞写出绕过 reporter 的安全敏感测试。
func TestDeliverMCPResultFromReporterForTest_DocStatesZeroIsRejected(t *testing.T) {
src := mustReadFile(t, "mcp_rpc.go")
if !strings.Contains(src, "DeliverMCPResultFromReporterForTest") {
t.Fatalf("expected helper to live in mcp_rpc.go")
}
// 提取 helper 上方的注释块：从 helper 名字往上找到第一段连续的 // 行。
idx := strings.Index(src, "func DeliverMCPResultFromReporterForTest(")
if idx < 0 {
t.Fatalf("helper not found in source")
}
prefix := src[:idx]
if !strings.Contains(prefix, "reporterID == 0") {
t.Fatalf("doc must mention reporterID == 0 contract explicitly")
}
if strings.Contains(prefix, "不校验") {
t.Fatalf("doc still claims reporterID==0 is 不校验; this contradicts deliverMCPResultFromReporter which drops it")
}
}
service/rpc/mcp_rpc_kill_switch_race_test.go
+156
View File
@@ -0,0 +1,156 @@
package rpc
import (
"context"
"errors"
"sync/atomic"
"testing"
"time"
"github.com/nezhahq/nezha/model"
pb "github.com/nezhahq/nezha/proto"
)
// Kill switch must beat a late agent reply. Without the cancelled-flag
// re-check in CallAgent, the following sequence surfaces success after
// EnableMCP=false:
//
// t0 agent puts TaskResult into resultCh (capacity 1, non-blocking)
// t1 admin flips EnableMCP=false → CancelAllMCPInflight closes cancelCh
// t2 CallAgent's select sees BOTH cases ready; Go picks one at random;
// if it picks resultCh, the call returns the agent's payload even
// though the operator's kill switch fired.
//
// The fix is to mark the entry cancelled BEFORE closing cancelCh and have
// the resultCh branch re-check that flag. This test pins the contract by
// driving the worst-case ordering: result is delivered FIRST, then the
// kill switch fires, then CallAgent observes both. With the race in place
// this would flake (random select); with the fix it always returns
// ErrMCPDisabled.
func TestCallAgent_KillSwitchBeatsConcurrentLateResult(t *testing.T) {
const target uint64 = 7301
stream := newFakeStream()
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
deliverMCPResultFromReporter(&pb.TaskResult{
Id: sent.GetId(),
Type: model.TaskTypeExec,
Successful: true,
Data: `{"exit_code":0,"stdout":"should-not-surface"}`,
}, target)
close(delivered)
}()
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
// deliverMCPResultFromReporter; the flag guarantees the late delivery is
// silently dropped even if the caller has not returned yet.
func TestCancelAllMCPInflight_LaterResultIsSwallowed(t *testing.T) {
const target uint64 = 7302
stream := newFakeStream()
cleanup := installFakeServer(t, target, stream)
defer cleanup()
taskIDCh := make(chan uint64, 1)
go func() {
sent := <-stream.sent
taskIDCh <- sent.GetId()
}()
resultCh := make(chan error, 1)
go func() {
_, err := CallAgent(context.Background(), target, model.TaskTypeExec,
model.ExecRequest{Cmd: "x"}, 5*time.Second)
resultCh <- err
}()
taskID := <-taskIDCh
CancelAllMCPInflight()
if err := <-resultCh; !errors.Is(err, ErrMCPDisabled) {
t.Fatalf("CallAgent must return ErrMCPDisabled after kill switch; got %v", err)
}
deliverMCPResultFromReporter(&pb.TaskResult{
Id: taskID,
Type: model.TaskTypeExec,
Successful: true,
Data: `{"exit_code":0}`,
}, target)
}
service/rpc/mcp_rpc_spoof_test.go
+149
View File
@@ -0,0 +1,149 @@
package rpc
import (
"context"
"encoding/json"
"errors"
"testing"
"time"
"github.com/nezhahq/nezha/model"
pb "github.com/nezhahq/nezha/proto"
)
// These tests pin the security invariant that an MCP TaskResult delivered
// back through RequestTask must come from the SAME agent the CallAgent was
// targeted at. The receive loop in service/rpc/nezha.go has the authenticated
// clientID in scope; deliverMCPResult must consume it and reject mismatches.
//
// Why the invariant matters: mcpInflight is keyed by a globally increasing
// counter (allocateMCPTaskID) and the lookup table is shared across servers.
// Without binding the inflight entry to the target serverID and verifying it
// against the reporter clientID, any compromised agent A can race a forged
// TaskResult for server B's CallAgent (resultCh capacity is 1; first reply
// wins, real reply is dropped). The same class of attack motivated the cron
// path's CanReportCronResult and the transfer path's pending.ID == result.Id
// check in this very file's RequestTask switch.
// TestDeliverMCPResult_RejectsForeignReporter is the security regression: a
// reporter that is NOT the call target must not be able to deliver into
// another server's inflight slot, even with a correctly-guessed taskID.
func TestDeliverMCPResult_RejectsForeignReporter(t *testing.T) {
const (
targetServerID uint64 = 6101
foreignAgentID uint64 = 6102
)
stream := newFakeStream()
cleanup := installFakeServer(t, targetServerID, stream)
defer cleanup()
captured := make(chan uint64, 1)
go func() {
sent := <-stream.sent
// Foreign agent racing a forged TaskResult with the right taskID.
DeliverMCPResultFromReporterForTest(&pb.TaskResult{
Id: sent.GetId(),
Type: model.TaskTypeExec,
Successful: true,
Data: `{"exit_code":0,"stdout":"forged"}`,
}, foreignAgentID)
captured <- sent.GetId()
}()
_, err := CallAgent(context.Background(), targetServerID, model.TaskTypeExec,
model.ExecRequest{Cmd: "x"}, 200*time.Millisecond)
if !errors.Is(err, ErrAgentTimeout) {
t.Fatalf("forged result from foreign reporter must NOT deliver; want ErrAgentTimeout, got %v", err)
}
select {
case <-captured:
case <-time.After(time.Second):
t.Fatalf("test stream never observed the dispatched task")
}
}
// TestDeliverMCPResult_AcceptsMatchingReporter is the green companion: when
// the reporter clientID matches the inflight target, the result must still
// route correctly (we are not breaking the happy path).
func TestDeliverMCPResult_AcceptsMatchingReporter(t *testing.T) {
const targetServerID uint64 = 6103
stream := newFakeStream()
cleanup := installFakeServer(t, targetServerID, stream)
defer cleanup()
want := model.ExecResult{ExitCode: 0, Stdout: "ok"}
payload, _ := json.Marshal(want)
go func() {
sent := <-stream.sent
DeliverMCPResultFromReporterForTest(&pb.TaskResult{
Id: sent.GetId(),
Type: model.TaskTypeExec,
Successful: true,
Data: string(payload),
}, targetServerID)
}()
raw, err := CallAgent(context.Background(), targetServerID, model.TaskTypeExec,
model.ExecRequest{Cmd: "x"}, 2*time.Second)
if err != nil {
t.Fatalf("matching reporter must deliver, got %v", err)
}
var got model.ExecResult
if err := json.Unmarshal(raw, &got); err != nil {
t.Fatalf("bad result json: %v", err)
}
if got.Stdout != "ok" {
t.Fatalf("payload not propagated, got %+v", got)
}
}
// TestDeliverMCPResult_InflightEntryBoundToServerID locks in the structural
// requirement that the inflight table records the target serverID. Without
// this binding deliverMCPResult cannot perform the reporter check above.
// Probing via reflection avoids exporting mcpInflight just for tests.
func TestDeliverMCPResult_InflightEntryBoundToServerID(t *testing.T) {
const targetServerID uint64 = 6104
stream := newFakeStream()
cleanup := installFakeServer(t, targetServerID, stream)
defer cleanup()
gotEntry := make(chan struct {
taskID uint64
serverID uint64
found bool
}, 1)
go func() {
sent := <-stream.sent
taskID := sent.GetId()
serverID, ok := inflightServerIDForTest(taskID)
gotEntry <- struct {
taskID uint64
serverID uint64
found bool
}{taskID, serverID, ok}
// Unblock CallAgent so the inflight slot is cleaned up.
DeliverMCPResultFromReporterForTest(&pb.TaskResult{
Id: taskID,
Type: model.TaskTypeExec,
Successful: true,
Data: "{}",
}, targetServerID)
}()
_, err := CallAgent(context.Background(), targetServerID, model.TaskTypeExec,
model.ExecRequest{Cmd: "x"}, 2*time.Second)
if err != nil {
t.Fatalf("unexpected CallAgent error: %v", err)
}
probe := <-gotEntry
if !probe.found {
t.Fatalf("inflight entry for taskID=%d not found while CallAgent was blocking", probe.taskID)
}
if probe.serverID != targetServerID {
t.Fatalf("inflight entry must carry target serverID=%d, got %d", targetServerID, probe.serverID)
}
}
service/rpc/mcp_rpc_test.go
+172
View File
@@ -0,0 +1,172 @@
package rpc
import (
"context"
"encoding/json"
"errors"
"sync/atomic"
"testing"
"time"
"google.golang.org/grpc/metadata"
"github.com/nezhahq/nezha/model"
pb "github.com/nezhahq/nezha/proto"
"github.com/nezhahq/nezha/service/singleton"
)
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 {
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()
original := singleton.ServerShared
sc := singleton.NewEmptyServerClassForTest()
srv := &model.Server{}
srv.ID = id
srv.SetTaskStream(stream)
sc.InsertForTest(srv)
singleton.ServerShared = sc
return func() { singleton.ServerShared = original }
}
func TestCallAgent_RejectsNonMCPType(t *testing.T) {
_, err := CallAgent(context.Background(), 1, model.TaskTypeCommand, struct{}{}, time.Second)
if err == nil {
t.Fatalf("expected error for non-MCP type")
}
}
func TestCallAgent_OfflineWhenNoStream(t *testing.T) {
original := singleton.ServerShared
sc := singleton.NewEmptyServerClassForTest()
srv := &model.Server{}
srv.ID = 7
sc.InsertForTest(srv)
singleton.ServerShared = sc
t.Cleanup(func() { singleton.ServerShared = original })
_, err := CallAgent(context.Background(), 7, model.TaskTypeExec, struct{}{}, time.Second)
if !errors.Is(err, ErrAgentOffline) {
t.Fatalf("expected ErrAgentOffline, got %v", err)
}
}
func TestCallAgent_HappyPath_DelivlersResultByTaskID(t *testing.T) {
stream := newFakeStream()
cleanup := installFakeServer(t, 42, stream)
defer cleanup()
resultPayload, _ := json.Marshal(model.ExecResult{ExitCode: 0, Stdout: "hello"})
var captured atomic.Uint64
done := make(chan struct{})
go func() {
sent := <-stream.sent
captured.Store(sent.GetId())
deliverMCPResult(&pb.TaskResult{
Id: sent.GetId(),
Type: model.TaskTypeExec,
Data: string(resultPayload),
Successful: true,
})
close(done)
}()
raw, err := CallAgent(context.Background(), 42, model.TaskTypeExec, model.ExecRequest{Cmd: "x"}, 5*time.Second)
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
<-done
if captured.Load() == 0 {
t.Fatalf("task id never captured")
}
var got model.ExecResult
if err := json.Unmarshal(raw, &got); err != nil {
t.Fatalf("bad result json: %v", err)
}
if got.Stdout != "hello" {
t.Fatalf("payload not propagated, got %+v", got)
}
}
func TestCallAgent_Timeout(t *testing.T) {
stream := newFakeStream()
cleanup := installFakeServer(t, 43, stream)
defer cleanup()
go func() { <-stream.sent }()
_, err := CallAgent(context.Background(), 43, model.TaskTypeFsRead, model.FsReadRequest{Path: "/x"}, 50*time.Millisecond)
if !errors.Is(err, ErrAgentTimeout) {
t.Fatalf("expected ErrAgentTimeout, got %v", err)
}
}
func TestCallAgent_LateResultIsDropped(t *testing.T) {
stream := newFakeStream()
cleanup := installFakeServer(t, 44, stream)
defer cleanup()
var taskID uint64
got := make(chan struct{})
go func() {
sent := <-stream.sent
taskID = sent.GetId()
close(got)
}()
_, err := CallAgent(context.Background(), 44, model.TaskTypeFsDelete, model.FsDeleteRequest{Path: "/x"}, 50*time.Millisecond)
if !errors.Is(err, ErrAgentTimeout) {
t.Fatalf("expected timeout")
}
<-got
deliverMCPResult(&pb.TaskResult{Id: taskID, Type: model.TaskTypeFsDelete, Successful: true, Data: "{}"})
if _, ok := mcpInflight.Load(taskID); ok {
t.Fatalf("inflight entry must be cleaned up after timeout")
}
}
func TestCallAgent_UnsuccessfulIsError(t *testing.T) {
stream := newFakeStream()
cleanup := installFakeServer(t, 45, stream)
defer cleanup()
go func() {
sent := <-stream.sent
deliverMCPResult(&pb.TaskResult{
Id: sent.GetId(),
Type: sent.GetType(),
Successful: false,
Data: "agent says nope",
})
}()
_, err := CallAgent(context.Background(), 45, model.TaskTypeFsWrite, model.FsWriteRequest{Path: "/x", Content: "y"}, time.Second)
if err == nil || err.Error() != "agent says nope" {
t.Fatalf("expected agent error message, got %v", err)
}
}
service/rpc/nezha.go
+177 -186
View File
@@ -5,13 +5,10 @@ import (
"errors"
"fmt"
"log"
"net"
"sync"
"time"
"github.com/jinzhu/copier"
geoipx "github.com/nezhahq/nezha/pkg/geoip"
"github.com/nezhahq/nezha/pkg/grpcx"
"github.com/nezhahq/nezha/pkg/tsdb"
"github.com/nezhahq/nezha/model"
@@ -23,29 +20,101 @@ var _ pb.NezhaServiceServer = (*NezhaHandler)(nil)
var NezhaHandlerSingleton *NezhaHandler
// ErrRequestTaskStreamSuperseded is returned when a RequestTask result arrives
// after its stream is no longer the live stream for the authenticated server.
var ErrRequestTaskStreamSuperseded = errors.New("request task stream superseded")
type NezhaHandler struct {
Auth *authHandler
ioStreams map[string]*ioStreamContext
ioStreamMutex *sync.RWMutex
Auth *authHandler
ioStreams map[string]*ioStreamContext
ioStreamMutex *sync.RWMutex
ioStreamGeneration uint64
ioStreamNotify chan struct{}
ioStreamWaitLockedHook func()
// Capability authorization and exact stream deletion share ioStreamMutex to avoid TOCTOU.
agentCompatCapabilities agentCompatCapabilityState
}
type serverMetricsWriter func(*tsdb.ServerMetrics) error
var writeServerMetrics serverMetricsWriter = writeServerMetricsToTSDB
func writeServerMetricsToTSDB(metrics *tsdb.ServerMetrics) error {
if !singleton.TSDBEnabled() {
return nil
}
return singleton.TSDBShared.WriteServerMetrics(metrics)
}
func NewNezhaHandler() *NezhaHandler {
return &NezhaHandler{
Auth: &authHandler{},
ioStreamMutex: new(sync.RWMutex),
ioStreams: make(map[string]*ioStreamContext),
handler := &NezhaHandler{
Auth: &authHandler{},
ioStreamMutex: new(sync.RWMutex),
ioStreams: make(map[string]*ioStreamContext),
ioStreamNotify: make(chan struct{}),
}
handler.initializeAgentCompatCapabilities()
return handler
}
// attachRequestTaskStream resolves the server for clientID and publishes the
// task stream. It mirrors the !ok || server == nil guard the other RPC entry
// points use: the server can be deleted between CheckRequestTask and this
// lookup, in which case Get returns a nil *Server and SetTaskStream would
// panic.
func attachRequestTaskStream(clientID uint64, stream pb.NezhaService_RequestTaskServer) (*model.Server, bool) {
server, ok := singleton.ServerShared.Get(clientID)
if !ok || server == nil {
return nil, false
}
server.SetTaskStream(stream)
return server, true
}
// clearRequestTaskStream detaches the dropped stream from whichever *Server is
// currently published for clientID. Edit and transfer rotation publish a new
// *Server that adopts the same stream holder, so cleanup must target the live
// map entry; the captured server is only the fallback for a removed entry.
func clearRequestTaskStream(clientID uint64, captured *model.Server, stream pb.NezhaService_RequestTaskServer) {
if current, ok := singleton.ServerShared.Get(clientID); ok && current != nil {
current.ClearTaskStreamIfCurrent(stream)
return
}
captured.ClearTaskStreamIfCurrent(stream)
}
// currentRequestTaskServer authorizes a received result against the live
// ServerShared entry. Server pointer replacement is valid when it inherited
// the same task stream holder; only a missing entry or different stream makes
// a received result stale.
func currentRequestTaskServer(clientID uint64, stream pb.NezhaService_RequestTaskServer) (*model.Server, error) {
current, ok := singleton.ServerShared.Get(clientID)
if !ok || current == nil || current.GetTaskStream() != stream {
return nil, ErrRequestTaskStreamSuperseded
}
return current, nil
}
func (s *NezhaHandler) RequestTask(stream pb.NezhaService_RequestTaskServer) error {
var clientID uint64
var err error
if clientID, err = s.Auth.Check(stream.Context()); err != nil {
if clientID, err = s.Auth.CheckRequestTask(stream.Context()); err != nil {
return err
}
server, _ := singleton.ServerShared.Get(clientID)
server.TaskStream = stream
server, ok := attachRequestTaskStream(clientID, stream)
if !ok {
return nil
}
defer clearRequestTaskStream(clientID, server, stream)
// If a transfer is mid-flight for this server, the agent has just brought
// up a fresh bidi stream — this is the moment to (re)deliver the
// ApplyConfig task carrying the new owner's AgentSecret. Pushes from
// dashboard mutation time are best-effort; this hook is the reliable
// re-delivery point that closes the offline-during-transfer gap.
if singleton.ServerTransferShared != nil {
singleton.ServerTransferShared.OnAgentReconnect(clientID)
}
var result *pb.TaskResult
for {
result, err = stream.Recv()
@@ -53,11 +122,16 @@ func (s *NezhaHandler) RequestTask(stream pb.NezhaService_RequestTaskServer) err
log.Printf("NEZHA>> RequestTask error: %v, clientID: %d\n", err, clientID)
return err
}
server, err = currentRequestTaskServer(clientID, stream)
if err != nil {
return err
}
switch result.GetType() {
case model.TaskTypeCommand:
// 处理上报的计划任务
cr, _ := singleton.CronShared.Get(result.GetId())
if cr != nil {
// 任务结果 ID 来自 agent，必须确认该 cron 本应派发给当前 reporter。
if singleton.CanReportCronResult(cr, server) {
// 保存当前服务器状态信息
var curServer model.Server
copier.Copy(&curServer, server)
@@ -82,7 +156,32 @@ func (s *NezhaHandler) RequestTask(stream pb.NezhaService_RequestTaskServer) err
}
server.ConfigCache <- result.Data
}
case model.TaskTypeServerTransferApply:
// Authorization: TaskResult.Id is attacker-controlled. Without
// the pending.ID == result.Id check below, agent A could cancel
// server B's in-flight transfer by spoofing B's transfer ID —
// same class of bug as commit 02129f1 in the cron path.
// Successful=true here is best-effort only; the authoritative
// verification is the agent's reconnect under the new secret.
if singleton.ServerTransferShared == nil {
continue
}
pending, ok := singleton.ServerTransferShared.LookupPending(clientID)
if !ok || pending.ID != result.GetId() {
log.Printf("NEZHA>> ServerTransferApply result ignored: clientID=%d reported transferID=%d but no matching pending transfer", clientID, result.GetId())
continue
}
if result.GetSuccessful() {
continue
}
if _, err := singleton.ServerTransferShared.MarkFailed(result.GetId(), result.GetData()); err != nil {
log.Printf("NEZHA>> ServerTransfer MarkFailed(%d) failed: %v", result.GetId(), err)
}
default:
if model.IsMCPRPCResult(result.GetType()) {
deliverMCPResultFromReporter(result, clientID)
continue
}
if model.IsServiceSentinelNeeded(result.GetType()) {
singleton.ServiceSentinelShared.Dispatch(singleton.ReportData{
Data: result,
@@ -98,84 +197,91 @@ func (s *NezhaHandler) ReportSystemState(stream pb.NezhaService_ReportSystemStat
if err != nil {
return err
}
server, ok := singleton.ServerShared.Get(clientID)
if !ok || server == nil {
return errors.New("server not found")
}
lease := server.AttachStateStream(stream)
defer lease.Clear()
var state *pb.State
var stateCount uint64
for {
state, err = stream.Recv()
if err != nil {
log.Printf("NEZHA>> ReportSystemState error: %v, clientID: %d\n", err, clientID)
return err
}
stateCount++
innerState := model.PB2State(state)
server, ok := singleton.ServerShared.Get(clientID)
if !ok || server == nil {
return errors.New("server not found")
}
server.LastActive = time.Now()
server.State = &innerState
if singleton.TSDBEnabled() {
maxTemp := 0.0
for _, t := range innerState.Temperatures {
if t.Temperature > maxTemp {
maxTemp = t.Temperature
lastActive := time.Now()
accepted := lease.UpdateStateWithSideEffect(&innerState, lastActive, func() error {
{
maxTemp := 0.0
for _, t := range innerState.Temperatures {
if t.Temperature > maxTemp {
maxTemp = t.Temperature
}
}
maxGPU := 0.0
for _, g := range innerState.GPU {
if g > maxGPU {
maxGPU = g
}
}
if err := writeServerMetrics(&tsdb.ServerMetrics{
ServerID: clientID,
Timestamp: lastActive,
CPU: innerState.CPU,
MemUsed: innerState.MemUsed,
SwapUsed: innerState.SwapUsed,
DiskUsed: innerState.DiskUsed,
NetInSpeed: innerState.NetInSpeed,
NetOutSpeed: innerState.NetOutSpeed,
NetInTransfer: innerState.NetInTransfer,
NetOutTransfer: innerState.NetOutTransfer,
Load1: innerState.Load1,
Load5: innerState.Load5,
Load15: innerState.Load15,
TCPConnCount: innerState.TcpConnCount,
UDPConnCount: innerState.UdpConnCount,
ProcessCount: innerState.ProcessCount,
Temperature: maxTemp,
Uptime: innerState.Uptime,
GPU: maxGPU,
}); err != nil {
log.Printf("NEZHA>> Failed to write server metrics to TSDB: %v", err)
}
}
maxGPU := 0.0
for _, g := range innerState.GPU {
if g > maxGPU {
maxGPU = g
}
}
if err := singleton.TSDBShared.WriteServerMetrics(&tsdb.ServerMetrics{
ServerID: clientID,
Timestamp: time.Now(),
CPU: innerState.CPU,
MemUsed: innerState.MemUsed,
SwapUsed: innerState.SwapUsed,
DiskUsed: innerState.DiskUsed,
NetInSpeed: innerState.NetInSpeed,
NetOutSpeed: innerState.NetOutSpeed,
NetInTransfer: innerState.NetInTransfer,
NetOutTransfer: innerState.NetOutTransfer,
Load1: innerState.Load1,
Load5: innerState.Load5,
Load15: innerState.Load15,
TCPConnCount: innerState.TcpConnCount,
UDPConnCount: innerState.UdpConnCount,
ProcessCount: innerState.ProcessCount,
Temperature: maxTemp,
Uptime: innerState.Uptime,
GPU: maxGPU,
}); err != nil {
log.Printf("NEZHA>> Failed to write server metrics to TSDB: %v", err)
}
return nil
})
if !accepted {
return errors.New("state stream superseded")
}
// 应对 dashboard / agent 重启的情况，如果从未记录过，先打点，等到小时时间点时入库
if server.PrevTransferInSnapshot == 0 || server.PrevTransferOutSnapshot == 0 {
server.PrevTransferInSnapshot = state.NetInTransfer
server.PrevTransferOutSnapshot = state.NetOutTransfer
if err := notifyStateReceived(clientID, server.UUID, lease.Generation(), stateCount); err != nil {
return err
}
if err := notifyReceiptAccepted(clientID, server.UUID, lease.Generation(), stateCount); err != nil {
return err
}
if err = stream.Send(&pb.Receipt{Proced: true}); err != nil {
return err
}
}
}
func (s *NezhaHandler) onReportSystemInfo(c context.Context, r *pb.Host) error {
func (s *NezhaHandler) onReportSystemInfo(c context.Context, r *pb.Host) (model.HostReportResult, error) {
var clientID uint64
var err error
if clientID, err = s.Auth.Check(c); err != nil {
return err
return model.HostReportResult{}, err
}
host := model.PB2Host(r)
server, ok := singleton.ServerShared.Get(clientID)
if !ok || server == nil {
return errors.New("server not found")
return model.HostReportResult{}, errors.New("server not found")
}
/**
@@ -183,138 +289,23 @@ func (s *NezhaHandler) onReportSystemInfo(c context.Context, r *pb.Host) error {
* 当 agent 重启时，bootTime 变大，agent 端会先上报 host 信息，然后上报 state 信息
* 这时可以借助上报顺序的空档，立即记录停机前的数据并重置 Prev* 数据，并由接下来的 state 方法重新赋值
*/
if !server.LastActive.IsZero() && host.BootTime > server.Host.BootTime {
singleton.RecordTransferHourlyUsage(server)
server.PrevTransferInSnapshot = 0
server.PrevTransferOutSnapshot = 0
}
server.Host = &host
return nil
return server.RuntimeHandle().ApplyHostReport(&host, time.Now(), singleton.PersistTransfer)
}
func (s *NezhaHandler) ReportSystemInfo(c context.Context, r *pb.Host) (*pb.Receipt, error) {
if err := s.onReportSystemInfo(c, r); err != nil {
if _, err := s.onReportSystemInfo(c, r); err != nil {
return nil, err
}
return &pb.Receipt{Proced: true}, nil
}
func (s *NezhaHandler) ReportSystemInfo2(c context.Context, r *pb.Host) (*pb.Uint64Receipt, error) {
if err := s.onReportSystemInfo(c, r); err != nil {
result, err := s.onReportSystemInfo(c, r)
if err != nil {
return nil, err
}
if err := notifyInfo2(result.ServerID, result.UUID); err != nil {
return nil, err
}
return &pb.Uint64Receipt{Data: singleton.DashboardBootTime}, nil
}
func (s *NezhaHandler) IOStream(stream pb.NezhaService_IOStreamServer) error {
if _, err := s.Auth.Check(stream.Context()); err != nil {
return err
}
id, err := stream.Recv()
if err != nil {
return err
}
// ff05ff05 是 Nezha 的魔数，用于标识流 ID
if id == nil || len(id.Data) < 4 || (id.Data[0] != 0xff && id.Data[1] != 0x05 && id.Data[2] != 0xff && id.Data[3] == 0x05) {
return fmt.Errorf("invalid stream id")
}
go func() {
for {
if err := stream.Send(&pb.IOStreamData{Data: []byte{}}); err != nil {
log.Printf("NEZHA>> IOStream keepAlive error: %v\n", err)
return
}
time.Sleep(time.Second * 30)
}
}()
streamId := string(id.Data[4:])
if _, err := s.GetStream(streamId); err != nil {
return err
}
iw := grpcx.NewIOStreamWrapper(stream)
if err := s.AgentConnected(streamId, iw); err != nil {
return err
}
iw.Wait()
return nil
}
func (s *NezhaHandler) ReportGeoIP(c context.Context, r *pb.GeoIP) (*pb.GeoIP, error) {
var clientID uint64
var err error
if clientID, err = s.Auth.Check(c); err != nil {
return nil, err
}
geoip := model.PB2GeoIP(r)
use6 := r.GetUse6()
if geoip.IP.IPv4Addr == "" && geoip.IP.IPv6Addr == "" {
ip, _ := c.Value(model.CtxKeyRealIP{}).(string)
if ip == "" {
ip, _ = c.Value(model.CtxKeyConnectingIP{}).(string)
}
geoip.IP.IPv4Addr = ip
}
joinedIP := geoip.IP.Join()
server, ok := singleton.ServerShared.Get(clientID)
if !ok || server == nil {
return nil, fmt.Errorf("server not found")
}
// 检查并更新DDNS
if server.EnableDDNS && joinedIP != "" &&
(server.GeoIP == nil || server.GeoIP.IP != geoip.IP) {
ipv4 := geoip.IP.IPv4Addr
ipv6 := geoip.IP.IPv6Addr
if err := singleton.ServerShared.UpdateDDNS(server, &model.IP{IPv4Addr: ipv4, IPv6Addr: ipv6}); err != nil {
log.Printf("NEZHA>> Failed to update DDNS for server %d: %v", err, server.ID)
}
}
// 发送IP变动通知
if server.GeoIP != nil && singleton.Conf.EnableIPChangeNotification &&
((singleton.Conf.Cover == model.ConfigCoverAll && !singleton.Conf.IgnoredIPNotificationServerIDs[clientID]) ||
(singleton.Conf.Cover == model.ConfigCoverIgnoreAll && singleton.Conf.IgnoredIPNotificationServerIDs[clientID])) &&
server.GeoIP.IP.Join() != "" &&
joinedIP != "" &&
server.GeoIP.IP != geoip.IP {
singleton.NotificationShared.SendNotification(singleton.Conf.IPChangeNotificationGroupID,
fmt.Sprintf(
"[%s] %s, %s => %s",
singleton.Localizer.T("IP Changed"),
server.Name, singleton.IPDesensitize(server.GeoIP.IP.Join()),
singleton.IPDesensitize(joinedIP),
),
"")
}
// 根据内置数据库查询 IP 地理位置
var ip string
if geoip.IP.IPv6Addr != "" && (use6 || geoip.IP.IPv4Addr == "") {
ip = geoip.IP.IPv6Addr
} else {
ip = geoip.IP.IPv4Addr
}
netIP := net.ParseIP(ip)
location, err := geoipx.Lookup(netIP)
if err != nil {
log.Printf("NEZHA>> geoip.Lookup: %v", err)
}
geoip.CountryCode = location
// 将地区码写入到 Host
server.GeoIP = &geoip
return &pb.GeoIP{Ip: nil, CountryCode: location, DashboardBootTime: singleton.DashboardBootTime}, nil
}
service/rpc/receipt_gate_agentcompat.go
+251
View File
@@ -0,0 +1,251 @@
//go:build agentcompat
package rpc
import (
"bufio"
"context"
"errors"
"fmt"
"net"
"strings"
"sync"
"time"
)
const receiptGateCommandTimeout = 30 * time.Second
type receiptGate struct {
conn net.Conn
read *bufio.Reader
generation uint64
stateMu sync.Mutex
ioMu sync.Mutex
closeOnce sync.Once
context context.Context
cancel context.CancelFunc
hold bool
acceptedCount uint64
}
var activeReceiptGate *receiptGate
var activeReceiptGateMu sync.RWMutex
var receiptGateListener net.Listener
var receiptGateGeneration uint64
var receiptGateCancel context.CancelFunc
var receiptGateWaitGroup sync.WaitGroup
func newReceiptGate(conn net.Conn, generation uint64) *receiptGate {
ctx, cancel := context.WithCancel(context.Background())
return &receiptGate{conn: conn, read: bufio.NewReader(conn), generation: generation, context: ctx, cancel: cancel, hold: true}
}
func SetReceiptGateListener(listener net.Listener) {
if listener == nil {
return
}
activeReceiptGateMu.Lock()
previousListener := receiptGateListener
previousCancel := receiptGateCancel
receiptGateListener = listener
listenerContext, cancel := context.WithCancel(context.Background())
receiptGateCancel = cancel
activeReceiptGateMu.Unlock()
if previousCancel != nil {
previousCancel()
}
if previousListener != nil {
_ = previousListener.Close()
}
receiptGateWaitGroup.Add(1)
go acceptReceiptGateConnections(listenerContext, listener)
}
func acceptReceiptGateConnections(ctx context.Context, listener net.Listener) {
defer receiptGateWaitGroup.Done()
for {
connection, err := listener.Accept()
if err != nil {
select {
case <-ctx.Done():
return
default:
}
return
}
activeReceiptGateMu.Lock()
receiptGateGeneration++
generation := receiptGateGeneration
previous := activeReceiptGate
gate := newReceiptGate(connection, generation)
activeReceiptGate = gate
activeReceiptGateMu.Unlock()
if previous != nil {
previous.close()
}
if err := connection.SetWriteDeadline(time.Now().Add(receiptGateCommandTimeout)); err != nil {
resetReceiptGate(gate)
continue
}
if _, err := fmt.Fprintln(connection, "ready"); err != nil {
resetReceiptGate(gate)
continue
}
_ = connection.SetWriteDeadline(time.Time{})
}
}
func (gate *receiptGate) close() {
gate.closeOnce.Do(func() {
gate.cancel()
_ = gate.conn.Close()
})
}
func CloseReceiptGate() {
activeReceiptGateMu.Lock()
listener := receiptGateListener
cancel := receiptGateCancel
gate := activeReceiptGate
receiptGateListener = nil
receiptGateCancel = nil
activeReceiptGate = nil
activeReceiptGateMu.Unlock()
if cancel != nil {
cancel()
}
if listener != nil {
_ = listener.Close()
}
if gate != nil {
gate.close()
}
receiptGateWaitGroup.Wait()
}
func resetReceiptGate(gate *receiptGate) {
activeReceiptGateMu.Lock()
if activeReceiptGate == gate {
activeReceiptGate = nil
}
activeReceiptGateMu.Unlock()
gate.close()
}
func currentReceiptGate() *receiptGate {
activeReceiptGateMu.RLock()
defer activeReceiptGateMu.RUnlock()
return activeReceiptGate
}
func (gate *receiptGate) sendAccepted(serverID uint64, uuid string, generation, count uint64) error {
gate.stateMu.Lock()
gate.acceptedCount++
count = gate.acceptedCount
hold := gate.hold
gate.stateMu.Unlock()
gate.ioMu.Lock()
defer gate.ioMu.Unlock()
if err := gate.conn.SetDeadline(time.Now().Add(receiptGateCommandTimeout)); err != nil {
resetReceiptGate(gate)
return err
}
if _, err := fmt.Fprintf(gate.conn, "accepted %d %s %d %d %d\n", serverID, uuid, gate.generation, generation, count); err != nil {
resetReceiptGate(gate)
return err
}
if !hold {
_ = gate.conn.SetDeadline(time.Time{})
return nil
}
command, err := gate.read.ReadString('\n')
if err != nil {
resetReceiptGate(gate)
return err
}
if strings.TrimSpace(command) != "release" {
err := errors.New("receipt gate received unexpected command")
resetReceiptGate(gate)
return err
}
gate.stateMu.Lock()
gate.hold = false
gate.stateMu.Unlock()
if err := gate.conn.SetDeadline(time.Time{}); err != nil {
resetReceiptGate(gate)
return err
}
return nil
}
func notifyReceiptAccepted(serverID uint64, uuid string, generation, count uint64) error {
gate := currentReceiptGate()
if gate == nil {
return nil
}
return gate.sendAccepted(serverID, uuid, generation, count)
}
func notifyStateReceived(serverID uint64, uuid string, generation, count uint64) error {
gate := currentReceiptGate()
if gate == nil {
return nil
}
gate.ioMu.Lock()
defer gate.ioMu.Unlock()
if err := gate.conn.SetWriteDeadline(time.Now().Add(receiptGateCommandTimeout)); err != nil {
resetReceiptGate(gate)
return err
}
if _, err := fmt.Fprintf(gate.conn, "state %d %s %d %d\n", serverID, uuid, generation, count); err != nil {
resetReceiptGate(gate)
return err
}
return gate.conn.SetWriteDeadline(time.Time{})
}
func notifyInfo2(serverID uint64, uuid string) error {
gate := currentReceiptGate()
if gate == nil {
return nil
}
gate.ioMu.Lock()
defer gate.ioMu.Unlock()
if err := gate.conn.SetWriteDeadline(time.Now().Add(receiptGateCommandTimeout)); err != nil {
resetReceiptGate(gate)
return err
}
if _, err := fmt.Fprintf(gate.conn, "info2 %d %d %s\n", gate.generation, serverID, uuid); err != nil {
resetReceiptGate(gate)
return err
}
return gate.conn.SetWriteDeadline(time.Time{})
}
func notifyMCPTaskDispatched(serverID, taskID, taskType uint64) {
notifyMCPReceipt("task", serverID, taskID, taskType)
}
func notifyMCPTaskResultAccepted(serverID, taskID, taskType uint64) {
notifyMCPReceipt("result", serverID, taskID, taskType)
}
func notifyMCPReceipt(kind string, serverID, taskID, taskType uint64) {
gate := currentReceiptGate()
if gate == nil {
return
}
gate.ioMu.Lock()
defer gate.ioMu.Unlock()
if err := gate.conn.SetWriteDeadline(time.Now().Add(receiptGateCommandTimeout)); err != nil {
resetReceiptGate(gate)
return
}
if _, err := fmt.Fprintf(gate.conn, "%s %d %d %d %d\n", kind, gate.generation, serverID, taskID, taskType); err != nil {
resetReceiptGate(gate)
return
}
if err := gate.conn.SetWriteDeadline(time.Time{}); err != nil {
resetReceiptGate(gate)
}
}
service/rpc/receipt_gate_agentcompat_test.go
+224
View File
@@ -0,0 +1,224 @@
//go:build agentcompat
package rpc
import (
"bufio"
"fmt"
"net"
"strings"
"testing"
"time"
"github.com/stretchr/testify/require"
)
func installReceiptGateForTest(conn net.Conn) *receiptGate {
activeReceiptGateMu.Lock()
receiptGateGeneration++
generation := receiptGateGeneration
activeReceiptGateMu.Unlock()
gate := newReceiptGate(conn, generation)
activeReceiptGateMu.Lock()
activeReceiptGate = gate
activeReceiptGateMu.Unlock()
return gate
}
func clearReceiptGateForTest() {
activeReceiptGateMu.Lock()
gate := activeReceiptGate
activeReceiptGate = nil
activeReceiptGateMu.Unlock()
if gate != nil {
gate.close()
}
}
func TestReceiptGate_EOFResetsGate(t *testing.T) {
// Given
serverConn, clientConn := net.Pipe()
defer clientConn.Close()
installReceiptGateForTest(serverConn)
defer clearReceiptGateForTest()
gate := currentReceiptGate()
require.NotNil(t, gate)
go func() {
reader := bufio.NewReader(clientConn)
_, _ = reader.ReadString('\n')
_ = clientConn.Close()
}()
// When
err := notifyReceiptAccepted(7, "uuid", 1, 1)
// Then
require.Error(t, err)
activeReceiptGateMu.RLock()
active := activeReceiptGate
activeReceiptGateMu.RUnlock()
require.Nil(t, active)
}
func TestReceiptGate_ListenerAcceptsAndReplacesConnections(t *testing.T) {
listener, err := net.Listen("tcp", "127.0.0.1:0")
require.NoError(t, err)
defer CloseReceiptGate()
SetReceiptGateListener(listener)
oldClient, err := net.Dial("tcp", listener.Addr().String())
require.NoError(t, err)
oldReader := bufio.NewReader(oldClient)
require.Equal(t, "ready\n", mustReadLine(t, oldReader))
newClient, err := net.Dial("tcp", listener.Addr().String())
require.NoError(t, err)
defer newClient.Close()
newReader := bufio.NewReader(newClient)
require.Equal(t, "ready\n", mustReadLine(t, newReader))
_ = oldClient.SetReadDeadline(time.Now().Add(time.Second))
_, oldErr := oldReader.ReadString('\n')
require.Error(t, oldErr)
}
func TestReceiptGate_CloseInterruptsHeldReadAndQueuedWrite(t *testing.T) {
// Given
serverConn, clientConn := net.Pipe()
defer clientConn.Close()
installReceiptGateForTest(serverConn)
acceptedStarted := make(chan struct{})
acceptedDone := make(chan error, 1)
go func() {
close(acceptedStarted)
acceptedDone <- notifyReceiptAccepted(7, "uuid", 1, 1)
}()
<-acceptedStarted
reader := bufio.NewReader(clientConn)
require.Equal(t, "accepted 7 uuid "+fmt.Sprint(currentReceiptGate().generation)+" 1 1\n", mustReadLine(t, reader))
infoStarted := make(chan struct{})
infoDone := make(chan error, 1)
go func() {
close(infoStarted)
infoDone <- notifyInfo2(9, "held")
}()
<-infoStarted
// When
CloseReceiptGate()
// Then
select {
case <-acceptedDone:
case <-time.After(time.Second):
t.Fatal("held receipt read was not interrupted")
}
select {
case <-infoDone:
case <-time.After(time.Second):
t.Fatal("queued notification write was not released")
}
}
func mustReadLine(t *testing.T, reader *bufio.Reader) string {
t.Helper()
line, err := reader.ReadString('\n')
require.NoError(t, err)
return line
}
func TestReceiptGate_MalformedCommandResetsGate(t *testing.T) {
// Given
serverConn, clientConn := net.Pipe()
defer clientConn.Close()
installReceiptGateForTest(serverConn)
defer clearReceiptGateForTest()
go func() {
reader := bufio.NewReader(clientConn)
_, _ = reader.ReadString('\n')
_, _ = clientConn.Write([]byte("hold\n"))
}()
// When
err := notifyReceiptAccepted(7, "uuid", 1, 1)
// Then
require.EqualError(t, err, "receipt gate received unexpected command")
activeReceiptGateMu.RLock()
active := activeReceiptGate
activeReceiptGateMu.RUnlock()
require.Nil(t, active)
}
func TestReceiptGate_ReplacementClosesOldConnection(t *testing.T) {
t.Run("replacement closes old connection", func(t *testing.T) {
// Given
oldServer, oldClient := net.Pipe()
newServer, newClient := net.Pipe()
t.Cleanup(func() { require.NoError(t, oldClient.Close()) })
t.Cleanup(func() { require.NoError(t, newClient.Close()) })
oldGate := installReceiptGateForTest(oldServer)
t.Cleanup(oldGate.close)
t.Cleanup(clearReceiptGateForTest)
newGate := newReceiptGate(newServer, oldGate.generation+1)
activeReceiptGateMu.Lock()
activeReceiptGate = newGate
activeReceiptGateMu.Unlock()
oldDone := make(chan error, 1)
go func() { oldDone <- oldGate.sendAccepted(7, "uuid", 1, 1) }()
reader := bufio.NewReader(oldClient)
_, _ = reader.ReadString('\n')
// When
oldGate.close()
// Then
select {
case err := <-oldDone:
require.Error(t, err)
case <-time.After(time.Second):
t.Fatal("old receipt gate remained blocked after replacement")
}
})
require.Nil(t, currentReceiptGate())
}
func TestReceiptGate_Info2AndReceiptNotificationsSerialize(t *testing.T) {
// Given
serverConn, clientConn := net.Pipe()
defer clientConn.Close()
installReceiptGateForTest(serverConn)
defer clearReceiptGateForTest()
gate := currentReceiptGate()
require.NotNil(t, gate)
lines := make(chan string, 2)
go func() {
reader := bufio.NewReader(clientConn)
line, err := reader.ReadString('\n')
if err != nil {
return
}
lines <- strings.TrimSpace(line)
_, _ = clientConn.Write([]byte("release\n"))
line, err = reader.ReadString('\n')
if err == nil {
lines <- strings.TrimSpace(line)
}
}()
// When
acceptedDone := make(chan error, 1)
go func() { acceptedDone <- notifyReceiptAccepted(7, "uuid", 1, 1) }()
select {
case err := <-acceptedDone:
require.NoError(t, err)
case <-time.After(time.Second):
require.NoError(t, <-acceptedDone)
}
require.NoError(t, notifyInfo2(7, "uuid"))
// Then
require.Equal(t, "accepted 7 uuid "+fmt.Sprint(gate.generation)+" 1 1", <-lines)
require.Equal(t, "info2 "+fmt.Sprint(gate.generation)+" 7 uuid", <-lines)
}
service/rpc/receipt_gate_default.go
+19
View File
@@ -0,0 +1,19 @@
//go:build !agentcompat
package rpc
import "net"
func SetReceiptGateListener(net.Listener) {}
func CloseReceiptGate() {}
func notifyReceiptAccepted(uint64, string, uint64, uint64) error { return nil }
func notifyStateReceived(uint64, string, uint64, uint64) error { return nil }
func notifyInfo2(uint64, string) error { return nil }
func notifyMCPTaskDispatched(uint64, uint64, uint64) {}
func notifyMCPTaskResultAccepted(uint64, uint64, uint64) {}
service/rpc/request_task_missing_server_test.go
+25
View File
@@ -0,0 +1,25 @@
package rpc
import (
"testing"
"github.com/nezhahq/nezha/model"
"github.com/nezhahq/nezha/service/singleton"
)
func TestAttachRequestTaskStream_MissingServerDoesNotPanic(t *testing.T) {
reporter := requestTaskSecurityServer(7, 200, "eeeeeeee-eeee-eeee-eeee-eeeeeeeeeeee")
setupRequestTaskSecurityFixture(t, []*model.Server{reporter}, nil, map[uint64]model.UserInfo{
200: {Role: model.RoleMember},
}, map[string]uint64{"reporter-secret": 200})
singleton.ServerShared.Delete([]uint64{reporter.ID})
srv, ok := attachRequestTaskStream(reporter.ID, nil)
if ok {
t.Fatal("attach must report not-ok when the server was deleted between auth and lookup")
}
if srv != nil {
t.Fatalf("attach must return a nil server for a deleted id, got %#v", srv)
}
}
service/rpc/request_task_security_test.go
+402
View File
@@ -0,0 +1,402 @@
package rpc
import (
"context"
"errors"
"testing"
"time"
"google.golang.org/grpc/metadata"
"gorm.io/driver/sqlite"
"gorm.io/gorm"
"github.com/nezhahq/nezha/model"
"github.com/nezhahq/nezha/pkg/i18n"
pb "github.com/nezhahq/nezha/proto"
"github.com/nezhahq/nezha/service/singleton"
)
type requestTaskSecurityStream struct {
ctx context.Context
results []*pb.TaskResult
onRecv func()
onResult func()
onSend func(*pb.Task)
sendErr error
}
func (s *requestTaskSecurityStream) Send(task *pb.Task) error {
if s.onSend != nil {
s.onSend(task)
}
return s.sendErr
}
func (s *requestTaskSecurityStream) Recv() (*pb.TaskResult, error) {
if len(s.results) == 0 {
if s.onRecv != nil {
s.onRecv()
}
return nil, context.Canceled
}
result := s.results[0]
s.results = s.results[1:]
if s.onResult != nil {
onResult := s.onResult
s.onResult = nil
onResult()
}
return result, nil
}
func (s *requestTaskSecurityStream) SetHeader(metadata.MD) error { return nil }
func (s *requestTaskSecurityStream) SendHeader(metadata.MD) error { return nil }
func (s *requestTaskSecurityStream) SetTrailer(metadata.MD) {}
func (s *requestTaskSecurityStream) Context() context.Context { return s.ctx }
func (s *requestTaskSecurityStream) SendMsg(any) error { return nil }
func (s *requestTaskSecurityStream) RecvMsg(any) error { return context.Canceled }
func TestRequestTaskSkipsCronResultOwnedByAnotherUser(t *testing.T) {
reporter := requestTaskSecurityServer(7, 200, "11111111-1111-1111-1111-111111111111")
victimCron := requestTaskSecurityCron(42, 100, model.CronCoverAll, nil)
setupRequestTaskSecurityFixture(t, []*model.Server{reporter}, []*model.Cron{victimCron}, map[uint64]model.UserInfo{
100: {Role: model.RoleMember},
200: {Role: model.RoleMember},
}, map[string]uint64{"reporter-secret": 200})
runRequestTaskSecurityResult(t, "reporter-secret", reporter.UUID, cronTaskResult(victimCron.ID, true))
if cronLastResult(t, victimCron.ID) {
t.Fatal("foreign cron result must not update victim cron status")
}
}
func TestRequestTaskSkipsCronResultOutsideReporterCover(t *testing.T) {
reporter := requestTaskSecurityServer(7, 200, "22222222-2222-2222-2222-222222222222")
coveredServerID := uint64(8)
cronTask := requestTaskSecurityCron(42, 200, model.CronCoverIgnoreAll, []uint64{coveredServerID})
setupRequestTaskSecurityFixture(t, []*model.Server{reporter}, []*model.Cron{cronTask}, map[uint64]model.UserInfo{
200: {Role: model.RoleMember},
}, map[string]uint64{"reporter-secret": 200})
runRequestTaskSecurityResult(t, "reporter-secret", reporter.UUID, cronTaskResult(cronTask.ID, true))
if cronLastResult(t, cronTask.ID) {
t.Fatal("cron result from a server outside cron cover must not update cron status")
}
}
func TestRequestTaskSkipsCronCoverAllExcludedReporter(t *testing.T) {
reporter := requestTaskSecurityServer(7, 200, "88888888-8888-8888-8888-888888888888")
cronTask := requestTaskSecurityCron(42, 200, model.CronCoverAll, []uint64{reporter.ID})
setupRequestTaskSecurityFixture(t, []*model.Server{reporter}, []*model.Cron{cronTask}, map[uint64]model.UserInfo{
200: {Role: model.RoleMember},
}, map[string]uint64{"reporter-secret": 200})
runRequestTaskSecurityResult(t, "reporter-secret", reporter.UUID, cronTaskResult(cronTask.ID, true))
if cronLastResult(t, cronTask.ID) {
t.Fatal("cron result from a server excluded by CronCoverAll must not update cron status")
}
}
func TestRequestTaskAllowsCronCoverAllReporter(t *testing.T) {
reporter := requestTaskSecurityServer(7, 200, "99999999-9999-9999-9999-999999999999")
cronTask := requestTaskSecurityCron(42, 200, model.CronCoverAll, []uint64{8})
setupRequestTaskSecurityFixture(t, []*model.Server{reporter}, []*model.Cron{cronTask}, map[uint64]model.UserInfo{
200: {Role: model.RoleMember},
}, map[string]uint64{"reporter-secret": 200})
runRequestTaskSecurityResult(t, "reporter-secret", reporter.UUID, cronTaskResult(cronTask.ID, true))
if !cronLastResult(t, cronTask.ID) {
t.Fatal("CronCoverAll reporter not in the exclusion list must update cron status")
}
}
func TestRequestTaskAllowsCronResultForCoveredOwnerServer(t *testing.T) {
reporter := requestTaskSecurityServer(7, 200, "33333333-3333-3333-3333-333333333333")
cronTask := requestTaskSecurityCron(42, 200, model.CronCoverIgnoreAll, []uint64{reporter.ID})
setupRequestTaskSecurityFixture(t, []*model.Server{reporter}, []*model.Cron{cronTask}, map[uint64]model.UserInfo{
200: {Role: model.RoleMember},
}, map[string]uint64{"reporter-secret": 200})
runRequestTaskSecurityResult(t, "reporter-secret", reporter.UUID, cronTaskResult(cronTask.ID, true))
if !cronLastResult(t, cronTask.ID) {
t.Fatal("covered owner cron result must update cron status")
}
}
func TestRequestTaskAllowsCronResultForCoveredAdminOwnedCron(t *testing.T) {
reporter := requestTaskSecurityServer(7, 200, "44444444-4444-4444-4444-444444444444")
cronTask := requestTaskSecurityCron(42, 1, model.CronCoverIgnoreAll, []uint64{reporter.ID})
setupRequestTaskSecurityFixture(t, []*model.Server{reporter}, []*model.Cron{cronTask}, map[uint64]model.UserInfo{
1: {Role: model.RoleAdmin},
200: {Role: model.RoleMember},
}, map[string]uint64{"reporter-secret": 200})
runRequestTaskSecurityResult(t, "reporter-secret", reporter.UUID, cronTaskResult(cronTask.ID, true))
if !cronLastResult(t, cronTask.ID) {
t.Fatal("covered admin-owned cron result must update cron status")
}
}
func TestRequestTaskSkipsAlertTriggerCronResultFromUntriggeredReporter(t *testing.T) {
reporter := requestTaskSecurityServer(7, 200, "55555555-5555-5555-5555-555555555555")
triggerServer := requestTaskSecurityServer(8, 200, "66666666-6666-6666-6666-666666666666")
cronTask := requestTaskSecurityCron(42, 200, model.CronCoverAlertTrigger, nil)
setupRequestTaskSecurityFixture(t, []*model.Server{reporter, triggerServer}, []*model.Cron{cronTask}, map[uint64]model.UserInfo{
200: {Role: model.RoleMember},
}, map[string]uint64{"reporter-secret": 200, "trigger-secret": 200})
connectRequestTaskSecurityTaskStream(t, triggerServer.ID)
singleton.CronTrigger(cronTask, triggerServer.ID)()
runRequestTaskSecurityResult(t, "reporter-secret", reporter.UUID, cronTaskResult(cronTask.ID, true))
if cronLastResult(t, cronTask.ID) {
t.Fatal("alert-trigger cron result from a non-triggered server must not update cron status")
}
}
func TestRequestTaskAllowsAlertTriggerCronResultForTriggeredReporter(t *testing.T) {
reporter := requestTaskSecurityServer(7, 200, "77777777-7777-7777-7777-777777777777")
cronTask := requestTaskSecurityCron(42, 200, model.CronCoverAlertTrigger, nil)
setupRequestTaskSecurityFixture(t, []*model.Server{reporter}, []*model.Cron{cronTask}, map[uint64]model.UserInfo{
200: {Role: model.RoleMember},
}, map[string]uint64{"reporter-secret": 200})
connectRequestTaskSecurityTaskStream(t, reporter.ID)
singleton.CronTrigger(cronTask, reporter.ID)()
runRequestTaskSecurityResult(t, "reporter-secret", reporter.UUID, cronTaskResult(cronTask.ID, true))
if !cronLastResult(t, cronTask.ID) {
t.Fatal("alert-trigger cron result from the triggered server must update cron status")
}
}
func TestRequestTaskAllowsAlertTriggerCronResultReportedDuringSend(t *testing.T) {
reporter := requestTaskSecurityServer(7, 200, "aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa")
cronTask := requestTaskSecurityCron(42, 200, model.CronCoverAlertTrigger, nil)
setupRequestTaskSecurityFixture(t, []*model.Server{reporter}, []*model.Cron{cronTask}, map[uint64]model.UserInfo{
200: {Role: model.RoleMember},
}, map[string]uint64{"reporter-secret": 200})
connectRequestTaskSecurityTaskStreamWithSendHook(t, reporter.ID, nil, func(task *pb.Task) {
if task.GetId() != cronTask.ID {
t.Fatalf("expected alert-trigger task %d, got %d", cronTask.ID, task.GetId())
}
runRequestTaskSecurityResult(t, "reporter-secret", reporter.UUID, cronTaskResult(cronTask.ID, true))
})
singleton.CronTrigger(cronTask, reporter.ID)()
if !cronLastResult(t, cronTask.ID) {
t.Fatal("alert-trigger cron result reported during Send must update cron status")
}
}
func TestRequestTaskSkipsAlertTriggerCronResultAfterSendFailure(t *testing.T) {
reporter := requestTaskSecurityServer(7, 200, "bbbbbbbb-bbbb-bbbb-bbbb-bbbbbbbbbbbb")
cronTask := requestTaskSecurityCron(42, 200, model.CronCoverAlertTrigger, nil)
setupRequestTaskSecurityFixture(t, []*model.Server{reporter}, []*model.Cron{cronTask}, map[uint64]model.UserInfo{
200: {Role: model.RoleMember},
}, map[string]uint64{"reporter-secret": 200})
connectRequestTaskSecurityTaskStreamWithSendHook(t, reporter.ID, errors.New("send failed"), nil)
singleton.CronTrigger(cronTask, reporter.ID)()
runRequestTaskSecurityResult(t, "reporter-secret", reporter.UUID, cronTaskResult(cronTask.ID, true))
if cronLastResult(t, cronTask.ID) {
t.Fatal("alert-trigger cron result after failed dispatch must not update cron status")
}
}
func TestRequestTaskClearsTaskStreamOnRecvError(t *testing.T) {
reporter := requestTaskSecurityServer(7, 200, "cccccccc-cccc-cccc-cccc-cccccccccccc")
setupRequestTaskSecurityFixture(t, []*model.Server{reporter}, nil, map[uint64]model.UserInfo{
200: {Role: model.RoleMember},
}, map[string]uint64{"reporter-secret": 200})
stream := requestTaskSecurityAuthedStream("reporter-secret", reporter.UUID)
err := NewNezhaHandler().RequestTask(stream)
if !errors.Is(err, context.Canceled) {
t.Fatalf("expected RequestTask to finish after Recv error, got %v", err)
}
server, ok := singleton.ServerShared.Get(reporter.ID)
if !ok {
t.Fatalf("server %d not found", reporter.ID)
}
if got := server.GetTaskStream(); got != nil {
t.Fatalf("dead RequestTask stream must be cleared, got %T", got)
}
}
func TestRequestTaskKeepsNewerTaskStreamOnOldRecvError(t *testing.T) {
reporter := requestTaskSecurityServer(7, 200, "dddddddd-dddd-dddd-dddd-dddddddddddd")
setupRequestTaskSecurityFixture(t, []*model.Server{reporter}, nil, map[uint64]model.UserInfo{
200: {Role: model.RoleMember},
}, map[string]uint64{"reporter-secret": 200})
server, ok := singleton.ServerShared.Get(reporter.ID)
if !ok {
t.Fatalf("server %d not found", reporter.ID)
}
newer := &requestTaskSecurityStream{ctx: context.Background()}
old := requestTaskSecurityAuthedStream("reporter-secret", reporter.UUID)
old.onRecv = func() {
server.SetTaskStream(newer)
}
err := NewNezhaHandler().RequestTask(old)
if !errors.Is(err, context.Canceled) {
t.Fatalf("expected RequestTask to finish after Recv error, got %v", err)
}
if got := server.GetTaskStream(); got != newer {
t.Fatalf("old stream cleanup must keep newer stream, got %T", got)
}
}
func setupRequestTaskSecurityFixture(t *testing.T, servers []*model.Server, crons []*model.Cron, users map[uint64]model.UserInfo, agentSecrets map[string]uint64) {
t.Helper()
originalDB := singleton.DB
originalConf := singleton.Conf
originalLoc := singleton.Loc
originalLocalizer := singleton.Localizer
originalNotification := singleton.NotificationShared
originalServerShared := singleton.ServerShared
originalServiceSentinel := singleton.ServiceSentinelShared
originalCronShared := singleton.CronShared
originalUserInfoMap := singleton.UserInfoMap
originalAgentSecretToUserID := singleton.AgentSecretToUserId
db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
if err != nil {
t.Fatal(err)
}
sqlDB, err := db.DB()
if err != nil {
t.Fatal(err)
}
sqlDB.SetMaxOpenConns(1)
singleton.DB = db
singleton.Conf = &singleton.ConfigClass{Config: &model.Config{}}
singleton.Loc = time.UTC
singleton.Localizer = i18n.NewLocalizer("en_US", "nezha", "translations", i18n.Translations)
singleton.NotificationShared = singleton.NewEmptyNotificationClassForTest()
if err := singleton.DB.AutoMigrate(model.Server{}, model.Cron{}); err != nil {
t.Fatal(err)
}
for _, server := range servers {
if err := singleton.DB.Create(server).Error; err != nil {
t.Fatal(err)
}
}
for _, cronTask := range crons {
if err := singleton.DB.Create(cronTask).Error; err != nil {
t.Fatal(err)
}
}
singleton.UserLock.Lock()
singleton.UserInfoMap = users
singleton.AgentSecretToUserId = agentSecrets
singleton.UserLock.Unlock()
singleton.ServerShared = singleton.NewServerClass()
singleton.CronShared = singleton.NewCronClass()
t.Cleanup(func() {
singleton.CronShared.Close()
_ = sqlDB.Close()
singleton.DB = originalDB
singleton.Conf = originalConf
singleton.Loc = originalLoc
singleton.Localizer = originalLocalizer
singleton.NotificationShared = originalNotification
singleton.ServiceSentinelShared = originalServiceSentinel
singleton.ServerShared = originalServerShared
singleton.CronShared = originalCronShared
singleton.UserLock.Lock()
singleton.UserInfoMap = originalUserInfoMap
singleton.AgentSecretToUserId = originalAgentSecretToUserID
singleton.UserLock.Unlock()
})
}
func requestTaskSecurityServer(id, userID uint64, uuid string) *model.Server {
return &model.Server{
Common: model.Common{ID: id, UserID: userID},
UUID: uuid,
Name: "request-task-security-server",
}
}
func requestTaskSecurityCron(id, userID uint64, cover uint8, servers []uint64) *model.Cron {
return &model.Cron{
Common: model.Common{ID: id, UserID: userID},
Name: "request-task-security-cron",
Command: "id",
Scheduler: "@every 1h",
Cover: cover,
Servers: servers,
}
}
func cronTaskResult(cronID uint64, successful bool) *pb.TaskResult {
return &pb.TaskResult{
Id: cronID,
Type: model.TaskTypeCommand,
Delay: 1,
Data: "cron result",
Successful: successful,
}
}
func connectRequestTaskSecurityTaskStream(t *testing.T, serverID uint64) {
t.Helper()
connectRequestTaskSecurityTaskStreamWithSendHook(t, serverID, nil, nil)
}
func connectRequestTaskSecurityTaskStreamWithSendHook(t *testing.T, serverID uint64, sendErr error, onSend func(*pb.Task)) {
t.Helper()
server, ok := singleton.ServerShared.Get(serverID)
if !ok {
t.Fatalf("server %d not found", serverID)
}
server.SetTaskStream(&requestTaskSecurityStream{ctx: context.Background(), sendErr: sendErr, onSend: onSend})
}
func runRequestTaskSecurityResult(t *testing.T, secret string, uuid string, result *pb.TaskResult) {
t.Helper()
stream := requestTaskSecurityAuthedStream(secret, uuid)
stream.results = []*pb.TaskResult{result}
err := NewNezhaHandler().RequestTask(stream)
if !errors.Is(err, context.Canceled) {
t.Fatalf("expected RequestTask to finish after test result, got %v", err)
}
}
func requestTaskSecurityAuthedStream(secret string, uuid string) *requestTaskSecurityStream {
return &requestTaskSecurityStream{
ctx: metadata.NewIncomingContext(context.Background(), metadata.Pairs(
"client_secret", secret,
"client_uuid", uuid,
)),
}
}
func cronLastResult(t *testing.T, cronID uint64) bool {
t.Helper()
var cronTask model.Cron
if err := singleton.DB.First(&cronTask, cronID).Error; err != nil {
t.Fatal(err)
}
return cronTask.LastResult
}
service/rpc/request_task_stale_stream_test.go
+135
View File
@@ -0,0 +1,135 @@
package rpc
import (
"context"
"errors"
"testing"
"github.com/nezhahq/nezha/model"
pb "github.com/nezhahq/nezha/proto"
"github.com/nezhahq/nezha/service/singleton"
)
// When a server is edited mid-session, updateServer swaps a new *Server into
// ServerShared that adopts the live stream holder. The agent's RequestTask
// cleanup must detach the stream from whichever *Server is currently published,
// not the stale object captured when the stream attached — otherwise the new
// object keeps reporting the agent as online on a dead stream.
func TestRequestTaskCleanupDetachesStreamFromCurrentServerAfterEdit(t *testing.T) {
reporter := requestTaskSecurityServer(7, 200, "ffffffff-ffff-ffff-ffff-ffffffffffff")
setupRequestTaskSecurityFixture(t, []*model.Server{reporter}, nil, map[uint64]model.UserInfo{
200: {Role: model.RoleMember},
}, map[string]uint64{"reporter-secret": 200})
old, ok := singleton.ServerShared.Get(reporter.ID)
if !ok {
t.Fatalf("server %d not found", reporter.ID)
}
stream := requestTaskSecurityAuthedStream("reporter-secret", reporter.UUID)
stream.onRecv = func() {
edited := &model.Server{Common: model.Common{ID: old.ID, UserID: old.UserID}, UUID: old.UUID, Name: "edited"}
edited.CopyFromRunningServer(old)
singleton.ServerShared.Update(edited, "")
}
if err := NewNezhaHandler().RequestTask(stream); !errors.Is(err, context.Canceled) {
t.Fatalf("expected RequestTask to finish after Recv error, got %v", err)
}
current, ok := singleton.ServerShared.Get(reporter.ID)
if !ok {
t.Fatalf("server %d not found after edit", reporter.ID)
}
if got := current.GetTaskStream(); got != nil {
t.Fatalf("edited server must report offline after the agent stream dropped, got %T", got)
}
}
func TestRequestTaskRejectsResultWhenServerDeletedAfterRecv(t *testing.T) {
reporter := requestTaskSecurityServer(7, 200, "10101010-1010-1010-1010-101010101010")
cronTask := requestTaskSecurityCron(42, 200, model.CronCoverAll, nil)
setupRequestTaskSecurityFixture(t, []*model.Server{reporter}, []*model.Cron{cronTask}, map[uint64]model.UserInfo{
200: {Role: model.RoleMember},
}, map[string]uint64{"reporter-secret": 200})
stream := requestTaskSecurityAuthedStream("reporter-secret", reporter.UUID)
stream.results = []*pb.TaskResult{cronTaskResult(cronTask.ID, true)}
stream.onResult = func() {
singleton.ServerShared.Delete([]uint64{reporter.ID})
}
err := NewNezhaHandler().RequestTask(stream)
if !errors.Is(err, ErrRequestTaskStreamSuperseded) {
t.Fatalf("expected stale RequestTask stream error, got %v", err)
}
assertCronResultNotUpdated(t, cronTask.ID)
}
func TestRequestTaskRejectsResultWhenNewerStreamSupersedesOld(t *testing.T) {
reporter := requestTaskSecurityServer(7, 200, "20202020-2020-2020-2020-202020202020")
cronTask := requestTaskSecurityCron(42, 200, model.CronCoverAll, nil)
setupRequestTaskSecurityFixture(t, []*model.Server{reporter}, []*model.Cron{cronTask}, map[uint64]model.UserInfo{
200: {Role: model.RoleMember},
}, map[string]uint64{"reporter-secret": 200})
current, ok := singleton.ServerShared.Get(reporter.ID)
if !ok {
t.Fatalf("server %d not found", reporter.ID)
}
newer := &requestTaskSecurityStream{ctx: context.Background()}
stream := requestTaskSecurityAuthedStream("reporter-secret", reporter.UUID)
stream.results = []*pb.TaskResult{cronTaskResult(cronTask.ID, true)}
stream.onResult = func() {
current.SetTaskStream(newer)
}
err := NewNezhaHandler().RequestTask(stream)
if !errors.Is(err, ErrRequestTaskStreamSuperseded) {
t.Fatalf("expected superseded RequestTask stream error, got %v", err)
}
if got := current.GetTaskStream(); got != newer {
t.Fatalf("old stream cleanup must preserve newer stream, got %T", got)
}
assertCronResultNotUpdated(t, cronTask.ID)
}
func TestRequestTaskAcceptsResultAfterServerPointerReplacementWithSameStream(t *testing.T) {
reporter := requestTaskSecurityServer(7, 200, "30303030-3030-3030-3030-303030303030")
cronTask := requestTaskSecurityCron(42, 200, model.CronCoverAll, nil)
setupRequestTaskSecurityFixture(t, []*model.Server{reporter}, []*model.Cron{cronTask}, map[uint64]model.UserInfo{
200: {Role: model.RoleMember},
}, map[string]uint64{"reporter-secret": 200})
old, ok := singleton.ServerShared.Get(reporter.ID)
if !ok {
t.Fatalf("server %d not found", reporter.ID)
}
stream := requestTaskSecurityAuthedStream("reporter-secret", reporter.UUID)
stream.results = []*pb.TaskResult{cronTaskResult(cronTask.ID, true)}
stream.onResult = func() {
replacement := &model.Server{Common: model.Common{ID: old.ID, UserID: old.UserID}, UUID: old.UUID, Name: "replacement"}
replacement.CopyFromRunningServer(old)
singleton.ServerShared.Update(replacement, "")
}
err := NewNezhaHandler().RequestTask(stream)
if !errors.Is(err, context.Canceled) {
t.Fatalf("expected RequestTask to finish after accepted result, got %v", err)
}
if !cronLastResult(t, cronTask.ID) {
t.Fatal("result on a replacement server that inherited the stream must be accepted")
}
}
func assertCronResultNotUpdated(t *testing.T, cronID uint64) {
t.Helper()
var cronTask model.Cron
if err := singleton.DB.First(&cronTask, cronID).Error; err != nil {
t.Fatal(err)
}
if cronTask.LastResult || !cronTask.LastExecutedAt.IsZero() {
t.Fatalf("stale RequestTask result must not mutate cron, got last_result=%t last_executed_at=%s", cronTask.LastResult, cronTask.LastExecutedAt)
}
}
service/rpc/state_metrics_test.go
+42
View File
@@ -0,0 +1,42 @@
package rpc
import (
"testing"
"time"
"github.com/stretchr/testify/require"
"github.com/nezhahq/nezha/model"
"github.com/nezhahq/nezha/pkg/tsdb"
)
func TestStateMetricsWriterRunsOnlyForCurrentGeneration(t *testing.T) {
// Given
server := &model.Server{}
model.InitServer(server)
oldLease := server.AttachStateStream(stateGenerationStream{})
newLease := server.AttachStateStream(stateGenerationStream{})
oldCalls := 0
newCalls := 0
oldWriter := writeServerMetrics
writeServerMetrics = func(*tsdb.ServerMetrics) error {
newCalls++
return nil
}
t.Cleanup(func() { writeServerMetrics = oldWriter })
// When
oldAccepted := server.UpdateStateIfCurrentWithSideEffect(oldLease, &model.HostState{Uptime: 11}, time.Unix(100, 0), func() error {
oldCalls++
return writeServerMetrics(&tsdb.ServerMetrics{ServerID: 7, Timestamp: time.Unix(100, 0)})
})
newAccepted := server.UpdateStateIfCurrentWithSideEffect(newLease, &model.HostState{Uptime: 22}, time.Unix(200, 0), func() error {
return writeServerMetrics(&tsdb.ServerMetrics{ServerID: 7, Timestamp: time.Unix(200, 0)})
})
// Then
require.False(t, oldAccepted)
require.True(t, newAccepted)
require.Zero(t, oldCalls)
require.Equal(t, 1, newCalls)
}
service/rpc/state_stream_generation_test.go
+206
View File
@@ -0,0 +1,206 @@
package rpc
import (
"context"
"sync"
"testing"
"time"
"github.com/stretchr/testify/require"
"google.golang.org/grpc/metadata"
"github.com/nezhahq/nezha/model"
"github.com/nezhahq/nezha/pkg/tsdb"
pb "github.com/nezhahq/nezha/proto"
"github.com/nezhahq/nezha/service/singleton"
)
func TestReportSystemState_HandlerWaitsForMetricsBeforeReceipt(t *testing.T) {
// Given
reporter := requestTaskSecurityServer(9, 200, "ffffffff-ffff-ffff-ffff-ffffffffffff")
setupRequestTaskSecurityFixture(t, []*model.Server{reporter}, nil, map[uint64]model.UserInfo{200: {Role: model.RoleMember}}, map[string]uint64{"reporter-secret": 200})
stop := make(chan struct{})
stream := &stateGenerationHandlerStream{
ctx: metadata.NewIncomingContext(context.Background(), metadata.Pairs("client_secret", "reporter-secret", "client_uuid", reporter.UUID)),
states: make(chan *pb.State, 1), receipts: make(chan *pb.Receipt, 1), stop: stop,
}
stream.states <- &pb.State{Uptime: 44}
metricsStarted := make(chan *tsdb.ServerMetrics, 1)
metricsRelease := make(chan struct{})
oldWriter := writeServerMetrics
writeServerMetrics = func(metrics *tsdb.ServerMetrics) error {
metricsStarted <- metrics
<-metricsRelease
return nil
}
t.Cleanup(func() { writeServerMetrics = oldWriter })
// When
done := make(chan error, 1)
go func() { done <- NewNezhaHandler().ReportSystemState(stream) }()
metrics := <-metricsStarted
select {
case <-stream.receipts:
t.Fatal("receipt sent before metrics writer completed")
default:
}
close(metricsRelease)
// Then
require.Equal(t, reporter.ID, metrics.ServerID)
require.Equal(t, uint64(44), metrics.Uptime)
require.NotNil(t, <-stream.receipts)
current, ok := singleton.ServerShared.Get(reporter.ID)
require.True(t, ok)
require.Equal(t, current.RuntimeSnapshot().LastActive, metrics.Timestamp)
close(stop)
require.ErrorIs(t, <-done, context.Canceled)
}
type stateGenerationStream struct{}
func (stateGenerationStream) Send(*pb.Receipt) error { return nil }
func (stateGenerationStream) Recv() (*pb.State, error) { return nil, nil }
func (stateGenerationStream) SetHeader(metadata.MD) error { return nil }
func (stateGenerationStream) SendHeader(metadata.MD) error { return nil }
func (stateGenerationStream) SetTrailer(metadata.MD) {}
func (stateGenerationStream) Context() context.Context { return context.Background() }
func (stateGenerationStream) SendMsg(any) error { return nil }
func (stateGenerationStream) RecvMsg(any) error { return nil }
type stateGenerationHandlerStream struct {
ctx context.Context
states chan *pb.State
receipts chan *pb.Receipt
stop <-chan struct{}
}
func (s *stateGenerationHandlerStream) Send(receipt *pb.Receipt) error {
s.receipts <- receipt
return nil
}
func (s *stateGenerationHandlerStream) Recv() (*pb.State, error) {
select {
case state := <-s.states:
return state, nil
case <-s.stop:
return nil, context.Canceled
}
}
func (s *stateGenerationHandlerStream) SetHeader(metadata.MD) error { return nil }
func (s *stateGenerationHandlerStream) SendHeader(metadata.MD) error { return nil }
func (s *stateGenerationHandlerStream) SetTrailer(metadata.MD) {}
func (s *stateGenerationHandlerStream) Context() context.Context { return s.ctx }
func (s *stateGenerationHandlerStream) SendMsg(any) error { return nil }
func (s *stateGenerationHandlerStream) RecvMsg(any) error { return nil }
func TestReportSystemState_HandlerOldStreamCannotClearNewerState(t *testing.T) {
// Given
reporter := requestTaskSecurityServer(7, 200, "eeeeeeee-eeee-eeee-eeee-eeeeeeeeeeee")
setupRequestTaskSecurityFixture(t, []*model.Server{reporter}, nil, map[uint64]model.UserInfo{
200: {Role: model.RoleMember},
}, map[string]uint64{"reporter-secret": 200})
oldStop := make(chan struct{})
newStop := make(chan struct{})
oldStream := &stateGenerationHandlerStream{
ctx: metadata.NewIncomingContext(context.Background(), metadata.Pairs("client_secret", "reporter-secret", "client_uuid", reporter.UUID)),
states: make(chan *pb.State, 1), receipts: make(chan *pb.Receipt, 1), stop: oldStop,
}
newStream := &stateGenerationHandlerStream{
ctx: metadata.NewIncomingContext(context.Background(), metadata.Pairs("client_secret", "reporter-secret", "client_uuid", reporter.UUID)),
states: make(chan *pb.State, 1), receipts: make(chan *pb.Receipt, 1), stop: newStop,
}
oldStream.states <- &pb.State{Uptime: 11}
newStream.states <- &pb.State{Uptime: 22}
handler := NewNezhaHandler()
oldDone := make(chan error, 1)
newDone := make(chan error, 1)
go func() { oldDone <- handler.ReportSystemState(oldStream) }()
<-oldStream.receipts
// When
go func() { newDone <- handler.ReportSystemState(newStream) }()
<-newStream.receipts
close(newStop)
require.ErrorIs(t, <-newDone, context.Canceled)
close(oldStop)
require.ErrorIs(t, <-oldDone, context.Canceled)
// Then
server, ok := singleton.ServerShared.Get(reporter.ID)
require.True(t, ok)
require.Equal(t, uint64(22), server.State.Uptime)
require.True(t, server.LastActive.IsZero())
}
func TestReportSystemState_OldStreamCannotUpdateNewerGeneration(t *testing.T) {
// Given
server := &model.Server{}
model.InitServer(server)
oldStream := stateGenerationStream{}
newStream := stateGenerationStream{}
oldLease := server.AttachStateStream(oldStream)
updateGate := make(chan struct{})
updateDone := make(chan bool, 1)
oldState := &model.HostState{Uptime: 11}
newState := &model.HostState{Uptime: 22}
oldTime := time.Unix(100, 0)
newTime := time.Unix(200, 0)
var waitGroup sync.WaitGroup
waitGroup.Add(1)
go func() {
defer waitGroup.Done()
<-updateGate
updateDone <- server.UpdateStateIfCurrent(oldLease, oldState, oldTime)
}()
// When
newLease := server.AttachStateStream(newStream)
close(updateGate)
oldUpdateAccepted := <-updateDone
newUpdateAccepted := server.UpdateStateIfCurrent(newLease, newState, newTime)
waitGroup.Wait()
// Then
require.False(t, oldUpdateAccepted)
require.True(t, newUpdateAccepted)
require.Equal(t, newState, server.State)
require.Equal(t, newTime, server.LastActive)
}
func TestReportSystemState_OldCleanupCannotClearNewerGeneration(t *testing.T) {
// Given
server := &model.Server{}
model.InitServer(server)
oldLease := server.AttachStateStream(stateGenerationStream{})
newLease := server.AttachStateStream(stateGenerationStream{})
state := &model.HostState{Uptime: 22}
activeAt := time.Unix(200, 0)
require.True(t, server.UpdateStateIfCurrent(newLease, state, activeAt))
// When
oldCleanup := server.ClearStateStreamIfCurrent(oldLease)
// Then
require.False(t, oldCleanup)
require.Equal(t, state, server.State)
require.Equal(t, activeAt, server.LastActive)
}
func TestReportSystemState_CurrentCleanupClearsOnlineVisibility(t *testing.T) {
// Given
server := &model.Server{}
model.InitServer(server)
lease := server.AttachStateStream(stateGenerationStream{})
activeAt := time.Unix(300, 0)
require.True(t, server.UpdateStateIfCurrent(lease, &model.HostState{Uptime: 33}, activeAt))
// When
cleared := server.ClearStateStreamIfCurrent(lease)
// Then
require.True(t, cleared)
require.True(t, server.LastActive.IsZero())
}
service/rpc/testdata_helper_test.go
+20
View File
@@ -0,0 +1,20 @@
package rpc
import (
"os"
"path/filepath"
"testing"
)
func mustReadFile(t *testing.T, name string) string {
t.Helper()
wd, err := os.Getwd()
if err != nil {
t.Fatalf("getwd: %v", err)
}
b, err := os.ReadFile(filepath.Join(wd, name))
if err != nil {
t.Fatalf("read %s: %v", name, err)
}
return string(b)
}
service/rpc/wait_for_agent_revoke_test.go
+108
View File
@@ -0,0 +1,108 @@
package rpc
import (
"context"
"testing"
"time"
)
// WaitForAgent 必须在 stream 被 RevokeStreamsForPurpose 强制下线时立刻返回，
// 否则 EnableMCP=false 的 kill switch 不能真的“立即”切断那些还卡在
// “等待 agent attach” 阶段的 transfer 请求 —— 它们会一直等到 timeout
// （生产路径上是 30 秒）。
//
// 期望行为：调用 RevokeStreamsForPurpose 后，WaitForAgent 在远小于 timeout
// 的时间内返回 (nil, false)。
func TestWaitForAgent_RevokeWakesUpWaiter(t *testing.T) {
h := NewNezhaHandler()
const streamID = "kill-switch-wait"
if err := h.CreateStreamWithPurpose(streamID, 0, 7, PurposeMCPTransfer); err != nil {
t.Fatalf("create waiter stream: %v", err)
}
done := make(chan struct {
io any
ok bool
dur time.Duration
}, 1)
start := time.Now()
go func() {
// 给一个明显大于 revoke 触发延时的 timeout；如果 revoke 没唤醒，
// WaitForAgent 会一直等到这里，下面的 assertion 就会失败。
stream, ok := h.WaitForAgent(context.Background(), streamID, 5*time.Second)
done <- struct {
io any
ok bool
dur time.Duration
}{stream, ok, time.Since(start)}
}()
waiter, err := h.GetStream(streamID)
if err != nil {
t.Fatalf("get waiter context: %v", err)
}
select {
case <-waiter.waitStartedCh:
case <-time.After(time.Second):
t.Fatal("WaitForAgent did not enter its blocking select")
}
if revoked := h.RevokeStreamsForPurpose(PurposeMCPTransfer); revoked != 1 {
t.Fatalf("expected to revoke exactly 1 MCP stream, got %d", revoked)
}
select {
case res := <-done:
if res.ok {
t.Fatalf("WaitForAgent must return ok=false after revoke; got ok=true")
}
if res.dur > time.Second {
t.Fatalf("WaitForAgent did not wake up promptly after revoke (took %s); kill switch is not immediate", res.dur)
}
case <-time.After(2 * time.Second):
t.Fatalf("WaitForAgent never returned after revoke; kill switch did not wake the waiter")
}
}
func TestRevokeStreamsForServerWakesWaitForAgentAndPreservesNewGeneration(t *testing.T) {
h := NewNezhaHandler()
const streamID = "server-revoke-generation"
if err := h.CreateStream(streamID, 0, 7); err != nil {
t.Fatalf("create waiter stream: %v", err)
}
done := make(chan bool, 1)
go func() {
_, ok := h.WaitForAgent(context.Background(), streamID, time.Minute)
done <- ok
}()
waiter, err := h.GetStream(streamID)
if err != nil {
t.Fatalf("get waiter stream: %v", err)
}
select {
case <-waiter.waitStartedCh:
case <-time.After(time.Second):
t.Fatal("WaitForAgent did not reach its blocking select")
}
h.RevokeStreamsForServer(7)
select {
case ok := <-done:
if ok {
t.Fatal("WaitForAgent must return false after server revocation")
}
case <-time.After(time.Second):
t.Fatal("server revocation did not wake WaitForAgent")
}
h.RevokeStreamsForServer(7)
if err := h.CreateStream(streamID, 0, 8); err != nil {
t.Fatalf("new generation must reuse released ID: %v", err)
}
h.RevokeStreamsForServer(7)
if h.StreamCount() != 1 {
t.Fatalf("new generation must remain tracked, got %d streams", h.StreamCount())
}
if err := h.CloseStream(streamID); err != nil {
t.Fatalf("cleanup new generation: %v", err)
}
}
service/singleton/alertsentinel.go
+13 -8
View File
@@ -149,13 +149,13 @@ func checkStatus() {
role = u.Role
}
UserLock.RUnlock()
if alert.UserID != server.UserID && !role.IsAdmin() {
if alert.UserID != server.GetUserID() && !role.IsAdmin() {
continue
}
alertsStore[alert.ID][server.ID] = append(alertsStore[alert.
ID][server.ID], alert.Snapshot(AlertsCycleTransferStatsStore[alert.ID], server, DB))
// 发送通知，分为触发报警和恢复通知
max, passed := alert.Check(alertsStore[alert.ID][server.ID])
_, passed := alert.Check(alertsStore[alert.ID][server.ID])
// 保存当前服务器状态信息
curServer := model.Server{}
copier.Copy(&curServer, server)
@@ -167,7 +167,7 @@ func checkStatus() {
alertsPrevState[alert.ID][server.ID] = _RuleCheckFail
message := fmt.Sprintf("[%s] %s(%s) %s", Localizer.T("Incident"),
server.Name, IPDesensitize(server.GeoIP.IP.Join()), alert.Name)
go CronShared.SendTriggerTasks(alert.FailTriggerTasks, curServer.ID)
go CronShared.SendTriggerTasks(alert.FailTriggerTasks, curServer.ID, alert.UserID)
go NotificationShared.SendNotification(alert.NotificationGroupID, message, NotificationMuteLabel.ServerIncident(server.ID, alert.ID), &curServer)
// 清除恢复通知的静音缓存
NotificationShared.UnMuteNotification(alert.NotificationGroupID, NotificationMuteLabel.ServerIncidentResolved(server.ID, alert.ID))
@@ -177,17 +177,22 @@ func checkStatus() {
if alertsPrevState[alert.ID][server.ID] == _RuleCheckFail {
message := fmt.Sprintf("[%s] %s(%s) %s", Localizer.T("Resolved"),
server.Name, IPDesensitize(server.GeoIP.IP.Join()), alert.Name)
go CronShared.SendTriggerTasks(alert.RecoverTriggerTasks, curServer.ID)
go CronShared.SendTriggerTasks(alert.RecoverTriggerTasks, curServer.ID, alert.UserID)
go NotificationShared.SendNotification(alert.NotificationGroupID, message, NotificationMuteLabel.ServerIncidentResolved(server.ID, alert.ID), &curServer)
// 清除失败通知的静音缓存
NotificationShared.UnMuteNotification(alert.NotificationGroupID, NotificationMuteLabel.ServerIncident(server.ID, alert.ID))
}
alertsPrevState[alert.ID][server.ID] = _RuleCheckPass
}
// 清理旧数据
if max > 0 && max < len(alertsStore[alert.ID][server.ID]) {
index := len(alertsStore[alert.ID][server.ID]) - max
alertsStore[alert.ID][server.ID] = alertsStore[alert.ID][server.ID][index:]
// 清理旧数据：保留窗口由规则定义决定（各规则 Duration 的最大值），
// 而非 Check 的判定结果。window==0 表示没有任何有效规则需要回看历史
// （例如全部 Duration<=0），此时清空采样避免切片无限增长。
window := alert.RetentionWindow()
samples := alertsStore[alert.ID][server.ID]
if window <= 0 {
alertsStore[alert.ID][server.ID] = samples[:0]
} else if window < len(samples) {
alertsStore[alert.ID][server.ID] = samples[len(samples)-window:]
}
}
}
service/singleton/alertsentinel_test.go
+119
View File
@@ -0,0 +1,119 @@
package singleton
import (
"testing"
"github.com/nezhahq/nezha/model"
)
// notifyDecision replays the exact send-gate from checkStatus (lines 164-186)
// for a single (alert, server) pair: given the current Check verdict and the
// previous stored state, it reports whether an incident or recovery notification
// would be dispatched and what the next stored state becomes. Kept in lockstep
// with checkStatus so the end-to-end "does it actually notify" path is testable
// without the DB/global singletons checkStatus pulls in.
func notifyDecision(triggerMode uint8, passed bool, prev uint8) (incident, recover bool, next uint8) {
if !passed {
if triggerMode == model.ModeAlwaysTrigger || prev != _RuleCheckFail {
return true, false, _RuleCheckFail
}
return false, false, _RuleCheckFail
}
if prev == _RuleCheckFail {
return false, true, _RuleCheckPass
}
return false, false, _RuleCheckPass
}
// driveCheckStatus simulates the checkStatus tick loop end-to-end: each tick
// appends one sample, runs the real Check, applies the real RetentionWindow
// trim, then runs the send-gate. It returns how many incident notifications
// would have been dispatched across all ticks and the final sample window size.
func driveCheckStatus(rule *model.AlertRule, triggerMode uint8, ticks int, sample []bool) (incidents int, finalWindow int) {
incidents, finalWindow, _ = driveCheckStatusCap(rule, triggerMode, ticks, sample)
return incidents, finalWindow
}
// driveCheckStatusCap additionally reports the peak length and capacity the
// sample slice ever reached, so tests can assert memory stays bounded.
func driveCheckStatusCap(rule *model.AlertRule, triggerMode uint8, ticks int, sample []bool) (incidents, finalLen, peakCap int) {
var samples [][]bool
prev := uint8(_RuleCheckNoData)
for i := 0; i < ticks; i++ {
samples = append(samples, append([]bool(nil), sample...))
_, passed := rule.Check(samples)
w := rule.RetentionWindow()
if w <= 0 {
samples = samples[:0]
} else if w < len(samples) {
samples = samples[len(samples)-w:]
}
if cap(samples) > peakCap {
peakCap = cap(samples)
}
incident, _, next := notifyDecision(triggerMode, passed, prev)
if incident {
incidents++
}
prev = next
}
return incidents, len(samples), peakCap
}
// TestCheckStatus_GeneralRuleFiresIncident is the end-to-end guard for the
// regression: a Duration:10 rule on a server that fails every tick must,
// after the window fills, reach passed=false and actually dispatch an incident
// notification. Before the fix the window was wiped each tick, so passed never
// became false and zero notifications were sent.
func TestCheckStatus_GeneralRuleFiresIncident(t *testing.T) {
rule := &model.AlertRule{Rules: []*model.Rule{{Type: "cpu", Duration: 10}}}
t.Run("AlwaysTrigger fires repeatedly once window fills", func(t *testing.T) {
incidents, window := driveCheckStatus(rule, model.ModeAlwaysTrigger, 30, []bool{false})
if window < 10 {
t.Fatalf("window never filled: got %d want >= 10", window)
}
if incidents == 0 {
t.Fatalf("AlwaysTrigger rule never dispatched an incident notification")
}
})
t.Run("OnetimeTrigger fires exactly once", func(t *testing.T) {
incidents, _ := driveCheckStatus(rule, model.ModeOnetimeTrigger, 30, []bool{false})
if incidents != 1 {
t.Fatalf("OnetimeTrigger must dispatch exactly one incident, got %d", incidents)
}
})
}
// TestCheckStatus_HealthyServerStaysSilent guards the other direction: a server
// passing every tick must never dispatch an incident.
func TestCheckStatus_HealthyServerStaysSilent(t *testing.T) {
rule := &model.AlertRule{Rules: []*model.Rule{{Type: "cpu", Duration: 10}}}
incidents, _ := driveCheckStatus(rule, model.ModeAlwaysTrigger, 30, []bool{true})
if incidents != 0 {
t.Fatalf("a healthy server must never trigger an incident, got %d", incidents)
}
}
// TestCheckStatus_SampleMemoryBounded pins the no-memory-leak invariant: no
// matter how many ticks run, the per-(alert,server) sample slice length and
// capacity stay bounded by the rule's retention window, never growing with
// elapsed time. Runs far more ticks than the window to expose any unbounded
// growth.
func TestCheckStatus_SampleMemoryBounded(t *testing.T) {
const duration = 10
rule := &model.AlertRule{Rules: []*model.Rule{{Type: "cpu", Duration: duration}}}
_, finalLen, peakCap := driveCheckStatusCap(rule, model.ModeAlwaysTrigger, 100000, []bool{false})
if finalLen > duration {
t.Fatalf("sample length exceeded retention window after many ticks: got %d want <= %d", finalLen, duration)
}
// append grows capacity geometrically; with length pinned at window+1 the
// backing array stabilises at a small constant. A generous 4x window bound
// catches any reintroduced unbounded growth without being flaky.
if peakCap > duration*4 {
t.Fatalf("sample capacity grew unbounded: peak cap %d exceeds 4x window %d", peakCap, duration*4)
}
}
service/singleton/clean_monitor_history_test.go
+57
View File
@@ -0,0 +1,57 @@
package singleton
import (
"path/filepath"
"testing"
"github.com/stretchr/testify/require"
"gorm.io/gorm"
"github.com/nezhahq/nezha/model"
)
func setupCleanMonitorHistoryTestDB(t *testing.T) {
t.Helper()
previousDB := DB
var err error
DB, err = gorm.Open(openSQLiteDialector(filepath.Join(t.TempDir(), "dashboard.sqlite")), &gorm.Config{})
require.NoError(t, err)
sqlDB, err := DB.DB()
require.NoError(t, err)
t.Cleanup(func() {
DB = previousDB
if err := sqlDB.Close(); err != nil {
t.Errorf("close transfer cleanup test database: %v", err)
}
})
require.NoError(t, DB.AutoMigrate(&model.Server{}, &model.Transfer{}, &model.AlertRule{}))
require.NoError(t, DB.Exec("INSERT INTO servers (id, name, uuid) VALUES (1, 'server', 'clean-monitor-history-test')").Error)
}
func TestCleanMonitorHistoryWithoutRulesDeletesAllTransfers(t *testing.T) {
setupCleanMonitorHistoryTestDB(t)
require.NoError(t, DB.Create(&model.Transfer{ServerID: 1, In: 1}).Error)
CleanMonitorHistory()
var count int64
require.NoError(t, DB.Model(&model.Transfer{}).Count(&count).Error)
require.Zero(t, count)
}
func TestCleanMonitorHistoryPreservesTransfersWhenAlertRulesCannotBeLoaded(t *testing.T) {
setupCleanMonitorHistoryTestDB(t)
require.NoError(t, DB.Create(&model.Transfer{ServerID: 1, In: 1}).Error)
require.NoError(t, DB.Exec("INSERT INTO alert_rules (id, name, rules_raw, fail_trigger_tasks_raw, recover_trigger_tasks_raw) VALUES (1, 'broken', '{', '[]', '[]')").Error)
var alerts []model.AlertRule
require.Error(t, DB.Find(&alerts).Error, "precondition: malformed rules_raw must fail AlertRule.AfterFind")
CleanMonitorHistory()
var count int64
require.NoError(t, DB.Model(&model.Transfer{}).Count(&count).Error)
require.EqualValues(t, 1, count)
}
service/singleton/config.go
+8
View File
@@ -1,6 +1,7 @@
package singleton
import (
"log"
"strconv"
"strings"
@@ -26,6 +27,13 @@ func InitConfigFromPath(path string) error {
if err != nil {
return err
}
rotated, err := Conf.RotateJWTSecretKeyIfNeeded(Version)
if err != nil {
return err
}
if rotated {
log.Printf("NEZHA>> Rotated jwt_secret_key for dashboard version %s", Version)
}
Conf.updateIgnoredIPNotificationID()
Conf.Oauth2Providers = utils.MapKeysToSlice(Conf.Oauth2)
service/singleton/config_test.go
+54
View File
@@ -0,0 +1,54 @@
package singleton
import (
"os"
"strings"
"testing"
"github.com/nezhahq/nezha/model"
)
func TestInitConfigFromPathRotatesJWTSecretKey(t *testing.T) {
file, err := os.CreateTemp(t.TempDir(), "nezha-config-*.yaml")
if err != nil {
t.Fatalf("create temp config: %v", err)
}
if _, err := file.WriteString("jwt_secret_key: leaked-secret\nagent_secret_key: agent-secret\njwt_secret_key_last_rotated_version: v2.0.12\n"); err != nil {
t.Fatalf("write temp config: %v", err)
}
if err := file.Close(); err != nil {
t.Fatalf("close temp config: %v", err)
}
originalConf := Conf
originalVersion := Version
originalTemplates := FrontendTemplates
Version = "v2.0.13"
FrontendTemplates = nil
t.Cleanup(func() {
Conf = originalConf
Version = originalVersion
FrontendTemplates = originalTemplates
})
if err := InitConfigFromPath(file.Name()); err != nil {
t.Fatalf("init config: %v", err)
}
if Conf.JWTSecretKey == "leaked-secret" {
t.Fatal("jwt_secret_key was not rotated")
}
if Conf.JWTSecretKeyLastRotatedVersion != model.JWTSecretKeyRotationBaselineVersion {
t.Fatalf("jwt secret key marker = %q, want %q", Conf.JWTSecretKeyLastRotatedVersion, model.JWTSecretKeyRotationBaselineVersion)
}
saved, err := os.ReadFile(file.Name())
if err != nil {
t.Fatalf("read saved config: %v", err)
}
if strings.Contains(string(saved), "leaked-secret") {
t.Fatalf("saved config still contains leaked jwt_secret_key: %s", saved)
}
if !strings.Contains(string(saved), "jwt_secret_key_last_rotated_version: v2.0.13") {
t.Fatalf("saved config did not persist jwt secret key marker: %s", saved)
}
}
service/singleton/crontask.go
+175 -9
View File
@@ -5,6 +5,8 @@ import (
"fmt"
"slices"
"strings"
"sync"
"time"
"github.com/jinzhu/copier"
@@ -15,9 +17,25 @@ import (
pb "github.com/nezhahq/nezha/proto"
)
const alertTriggerCronResultAuthorizationTTL = 24 * time.Hour
type CronClass struct {
class[uint64, *model.Cron]
*cron.Cron
pendingAlertTriggerTasksMu sync.Mutex
pendingAlertTriggerTasks map[uint64]map[uint64][]time.Time
closeOnce sync.Once
}
// Close stops the scheduler and joins every job before a test restores globals.
// The embedded cron.Stop only exposes the completion context; callers must await it.
func (c *CronClass) Close() {
if c == nil || c.Cron == nil {
return
}
c.closeOnce.Do(func() {
<-c.Cron.Stop().Done()
})
}
func NewCronClass() *CronClass {
@@ -64,7 +82,8 @@ func NewCronClass() *CronClass {
list: list,
sortedList: sortedList,
},
Cron: cronx,
Cron: cronx,
pendingAlertTriggerTasks: make(map[uint64]map[uint64][]time.Time),
}
}
@@ -78,6 +97,7 @@ func (c *CronClass) Update(cr *model.Cron) {
delete(c.list, cr.ID)
c.list[cr.ID] = cr
c.listMu.Unlock()
c.deleteAlertTriggerCronResultAuthorizations([]uint64{cr.ID})
c.sortList()
}
@@ -92,6 +112,7 @@ func (c *CronClass) Delete(idList []uint64) {
delete(c.list, id)
}
c.listMu.Unlock()
c.deleteAlertTriggerCronResultAuthorizations(idList)
c.sortList()
}
@@ -110,11 +131,11 @@ func (c *CronClass) sortList() {
c.sortedList = sortedList
}
func (c *CronClass) SendTriggerTasks(taskIDs []uint64, triggerServer uint64) {
func (c *CronClass) SendTriggerTasks(taskIDs []uint64, triggerServer uint64, triggerOwner uint64) {
c.listMu.RLock()
var cronLists []*model.Cron
for _, taskID := range taskIDs {
if c, ok := c.list[taskID]; ok {
if c, ok := c.list[taskID]; ok && cronCanBeTriggeredByOwner(c, triggerOwner) {
cronLists = append(cronLists, c)
}
}
@@ -126,6 +147,117 @@ func (c *CronClass) SendTriggerTasks(taskIDs []uint64, triggerServer uint64) {
}
}
func cronCanBeTriggeredByOwner(cr *model.Cron, triggerOwner uint64) bool {
return cr.UserID == triggerOwner || userIsAdmin(triggerOwner)
}
func CanReportCronResult(cr *model.Cron, reporter *model.Server) bool {
if cr == nil || reporter == nil || !cronCanSendToServer(cr, reporter) {
return false
}
if cr.Cover == model.CronCoverAll {
return !slices.Contains(cr.Servers, reporter.ID)
}
if cr.Cover == model.CronCoverIgnoreAll {
return slices.Contains(cr.Servers, reporter.ID)
}
if cr.Cover == model.CronCoverAlertTrigger {
return CronShared != nil && CronShared.consumeAlertTriggerCronResult(cr.ID, reporter.ID)
}
return false
}
func (c *CronClass) reserveAlertTriggerCronResult(cronID uint64, serverID uint64) {
c.pendingAlertTriggerTasksMu.Lock()
defer c.pendingAlertTriggerTasksMu.Unlock()
now := time.Now()
c.pruneExpiredAlertTriggerCronResultsLocked(now)
if c.pendingAlertTriggerTasks == nil {
c.pendingAlertTriggerTasks = make(map[uint64]map[uint64][]time.Time)
}
if c.pendingAlertTriggerTasks[cronID] == nil {
c.pendingAlertTriggerTasks[cronID] = make(map[uint64][]time.Time)
}
c.pendingAlertTriggerTasks[cronID][serverID] = append(c.pendingAlertTriggerTasks[cronID][serverID], now.Add(alertTriggerCronResultAuthorizationTTL))
}
func (c *CronClass) revokeAlertTriggerCronResult(cronID uint64, serverID uint64) {
c.pendingAlertTriggerTasksMu.Lock()
defer c.pendingAlertTriggerTasksMu.Unlock()
serverTasks := c.pendingAlertTriggerTasks[cronID]
expiresAtList := serverTasks[serverID]
if len(expiresAtList) == 0 {
return
}
expiresAtList = expiresAtList[:len(expiresAtList)-1]
if len(expiresAtList) == 0 {
delete(serverTasks, serverID)
} else {
serverTasks[serverID] = expiresAtList
}
if len(serverTasks) == 0 {
delete(c.pendingAlertTriggerTasks, cronID)
}
}
func (c *CronClass) consumeAlertTriggerCronResult(cronID uint64, serverID uint64) bool {
c.pendingAlertTriggerTasksMu.Lock()
defer c.pendingAlertTriggerTasksMu.Unlock()
c.pruneExpiredAlertTriggerCronResultsLocked(time.Now())
return c.consumeAlertTriggerCronResultLocked(cronID, serverID)
}
func (c *CronClass) consumeAlertTriggerCronResultLocked(cronID uint64, serverID uint64) bool {
serverTasks := c.pendingAlertTriggerTasks[cronID]
expiresAtList := serverTasks[serverID]
if len(expiresAtList) == 0 {
return false
}
expiresAtList = expiresAtList[1:]
if len(expiresAtList) == 0 {
delete(serverTasks, serverID)
} else {
serverTasks[serverID] = expiresAtList
}
if len(serverTasks) == 0 {
delete(c.pendingAlertTriggerTasks, cronID)
}
return true
}
func (c *CronClass) pruneExpiredAlertTriggerCronResultsLocked(now time.Time) {
for cronID, serverTasks := range c.pendingAlertTriggerTasks {
for serverID, expiresAtList := range serverTasks {
validExpiresAtList := expiresAtList[:0]
for _, expiresAt := range expiresAtList {
if expiresAt.After(now) {
validExpiresAtList = append(validExpiresAtList, expiresAt)
}
}
if len(validExpiresAtList) == 0 {
delete(serverTasks, serverID)
} else {
serverTasks[serverID] = validExpiresAtList
}
}
if len(serverTasks) == 0 {
delete(c.pendingAlertTriggerTasks, cronID)
}
}
}
func (c *CronClass) deleteAlertTriggerCronResultAuthorizations(cronIDs []uint64) {
c.pendingAlertTriggerTasksMu.Lock()
defer c.pendingAlertTriggerTasksMu.Unlock()
for _, cronID := range cronIDs {
delete(c.pendingAlertTriggerTasks, cronID)
}
}
func ManualTrigger(cr *model.Cron) {
CronTrigger(cr)()
}
@@ -141,12 +273,21 @@ func CronTrigger(cr *model.Cron, triggerServer ...uint64) func() {
return
}
if s, ok := ServerShared.Get(triggerServer[0]); ok {
if s.TaskStream != nil {
s.TaskStream.Send(&pb.Task{
if !cronCanSendToServer(cr, s) {
return
}
if s.GetTaskStream() != nil {
cronShared := CronShared
if cronShared != nil {
cronShared.reserveAlertTriggerCronResult(cr.ID, s.ID)
}
if err := s.SendTask(&pb.Task{
Id: cr.ID,
Data: cr.Command,
Type: model.TaskTypeCommand,
})
}); err != nil && cronShared != nil {
cronShared.revokeAlertTriggerCronResult(cr.ID, s.ID)
}
} else {
// 保存当前服务器状态信息
curServer := model.Server{}
@@ -157,15 +298,24 @@ func CronTrigger(cr *model.Cron, triggerServer ...uint64) func() {
return
}
for _, s := range ServerShared.Range {
// 先在锁内快照 server 列表再逐个 SendTask：ServerShared.Range 会在整个
// 回调期间持 listMu.RLock，而 SendTask 走阻塞 gRPC，一个卡死的 agent
// 会让需要写锁的 server 编辑/删除被拖死。GetList 克隆后即释放锁。
for _, s := range ServerShared.GetList() {
if s == nil {
continue
}
if !cronCanSendToServer(cr, s) {
continue
}
if cr.Cover == model.CronCoverAll && crIgnoreMap[s.ID] {
continue
}
if cr.Cover == model.CronCoverIgnoreAll && !crIgnoreMap[s.ID] {
continue
}
if s.TaskStream != nil {
s.TaskStream.Send(&pb.Task{
if s.GetTaskStream() != nil {
_ = s.SendTask(&pb.Task{
Id: cr.ID,
Data: cr.Command,
Type: model.TaskTypeCommand,
@@ -179,3 +329,19 @@ func CronTrigger(cr *model.Cron, triggerServer ...uint64) func() {
}
}
}
func cronCanSendToServer(cr *model.Cron, server *model.Server) bool {
return cr.UserID == server.GetUserID() || userIsAdmin(cr.UserID)
}
func userIsAdmin(userID uint64) bool {
if userID == 0 {
return true
}
UserLock.RLock()
defer UserLock.RUnlock()
userInfo, ok := UserInfoMap[userID]
return ok && userInfo.Role.IsAdmin()
}
service/singleton/crontask_lifecycle_test.go
+97
View File
@@ -0,0 +1,97 @@
package singleton
import (
"context"
"testing"
"time"
"github.com/robfig/cron/v3"
"github.com/stretchr/testify/require"
)
const lifecycleTestTimeout = time.Second
func TestCronClassClose_waitsForRunningJobs(t *testing.T) {
// Given
started := make(chan struct{})
release := make(chan struct{})
events := make(chan string, 9)
cronClass := &CronClass{Cron: cron.New(cron.WithSeconds())}
_, err := cronClass.AddFunc("@every 1ns", func() {
defer func() { events <- "job" }()
close(started)
<-release
})
require.NoError(t, err)
cronClass.Start()
<-started
// When
closed := make(chan struct{})
for range 8 {
go func() {
cronClass.Close()
events <- "close"
closed <- struct{}{}
}()
}
// Then
close(release)
firstEvent := awaitCronLifecycleEvent(t, events, "cron lifecycle did not complete")
if firstEvent != "job" {
t.Fatalf("Close returned before the running cron job returned: first event=%q", firstEvent)
}
for range 8 {
awaitCronLifecycleSignal(t, closed, "concurrent Close call did not return")
}
for range 8 {
if event := awaitCronLifecycleEvent(t, events, "concurrent Close call did not complete"); event != "close" {
t.Fatalf("unexpected cron lifecycle event: %q", event)
}
}
}
func TestCronClassClose_isIdempotentAndNilSafe(t *testing.T) {
cronClass := &CronClass{Cron: cron.New(cron.WithSeconds())}
cronClass.Start()
closed := make(chan struct{})
for range 8 {
go func() {
cronClass.Close()
closed <- struct{}{}
}()
}
for range 8 {
awaitCronLifecycleSignal(t, closed, "concurrent Close call did not return")
}
cronClass.Close()
var nilCronClass *CronClass
nilCronClass.Close()
(&CronClass{}).Close()
}
func awaitCronLifecycleSignal(t *testing.T, signal <-chan struct{}, message string) {
t.Helper()
ctx, cancel := context.WithTimeout(context.Background(), lifecycleTestTimeout)
defer cancel()
select {
case <-signal:
case <-ctx.Done():
t.Fatal(message)
}
}
func awaitCronLifecycleEvent(t *testing.T, events <-chan string, message string) string {
t.Helper()
ctx, cancel := context.WithTimeout(context.Background(), lifecycleTestTimeout)
defer cancel()
select {
case event := <-events:
return event
case <-ctx.Done():
t.Fatal(message)
return ""
}
}
service/singleton/ddns.go
+20 -1
View File
@@ -3,6 +3,7 @@ package singleton
import (
"cmp"
"fmt"
"log"
"slices"
"github.com/libdns/cloudflare"
@@ -56,12 +57,30 @@ func (c *DDNSClass) Delete(idList []uint64) {
c.sortList()
}
func (c *DDNSClass) GetDDNSProvidersFromProfiles(profileId []uint64, ip *model.IP) ([]*ddns2.Provider, error) {
// profileOwnedByRealAdmin reports whether uid is a genuine admin user that
// may share its DDNS profiles globally. userIsAdmin(0) returns true as a
// "system resource" shortcut, but a profile with UserID==0 is a migration /
// default-value artifact, not an admin grant — sharing it with foreign server
// owners reopens GHSA-39g2-8x68-pmx8. A real admin always has a non-zero ID.
func profileOwnedByRealAdmin(uid uint64) bool {
return uid != 0 && userIsAdmin(uid)
}
// GHSA-39g2-8x68-pmx8: bind-time CheckPermission 对「不存在的 profile ID」放行，
// 攻击者可预绑定将来才会被受害者创建的自增 ID。worker 解析时必须按 ownerUID
// 重新校验归属，跳过非 server owner（且非管理员）所有的 profile。
func (c *DDNSClass) GetDDNSProvidersFromProfiles(profileId []uint64, ip *model.IP, ownerUID uint64) ([]*ddns2.Provider, error) {
profiles := make([]*model.DDNSProfile, 0, len(profileId))
c.listMu.RLock()
for _, id := range profileId {
if profile, ok := c.list[id]; ok {
if profile.UserID != ownerUID && !profileOwnedByRealAdmin(profile.UserID) {
// Fail-closed skip: an admin may bind a member-owned profile,
// but worker-time only runs same-owner or real-admin profiles.
log.Printf("NEZHA>> Skipping DDNS profile %d (owner %d) for server owner %d: not owned by server owner or a real admin", profile.ID, profile.UserID, ownerUID)
continue
}
profiles = append(profiles, profile)
} else {
c.listMu.RUnlock()
service/singleton/ddns_worker_authz_test.go
+121
View File
@@ -0,0 +1,121 @@
package singleton
import (
"testing"
"github.com/nezhahq/nezha/model"
)
// newDDNSClassForTest builds a DDNSClass backed by an in-memory profile map,
// mirroring the production cache layout without touching the database.
func newDDNSClassForTest(profiles ...*model.DDNSProfile) *DDNSClass {
list := make(map[uint64]*model.DDNSProfile, len(profiles))
for _, p := range profiles {
list[p.ID] = p
}
return &DDNSClass{
class: class[uint64, *model.DDNSProfile]{
list: list,
sortedList: profiles,
},
}
}
// GHSA-39g2-8x68-pmx8: a server owned by the attacker must not be able to
// drive a DDNS update through a DDNS profile owned by another (victim) user.
// The worker-time resolution must skip foreign-owned profiles.
func TestGetDDNSProvidersSkipsForeignOwnedProfile(t *testing.T) {
replaceUserInfoMapForSecurityTest(t, map[uint64]model.UserInfo{
100: {Role: model.RoleMember}, // attacker / server owner
200: {Role: model.RoleMember}, // victim / profile owner
})
victimProfile := &model.DDNSProfile{
Common: model.Common{ID: 1, UserID: 200},
Provider: model.ProviderDummy,
Name: "victim-profile",
AccessSecret: "victim-secret",
}
dc := newDDNSClassForTest(victimProfile)
providers, err := dc.GetDDNSProvidersFromProfiles([]uint64{1}, &model.IP{}, 100)
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if len(providers) != 0 {
t.Fatalf("expected foreign-owned profile to be skipped, got %d provider(s)", len(providers))
}
}
// A server owner using their own DDNS profile must still resolve normally.
func TestGetDDNSProvidersAllowsOwnedProfile(t *testing.T) {
replaceUserInfoMapForSecurityTest(t, map[uint64]model.UserInfo{
100: {Role: model.RoleMember},
})
ownProfile := &model.DDNSProfile{
Common: model.Common{ID: 5, UserID: 100},
Provider: model.ProviderDummy,
Name: "own-profile",
}
dc := newDDNSClassForTest(ownProfile)
providers, err := dc.GetDDNSProvidersFromProfiles([]uint64{5}, &model.IP{}, 100)
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if len(providers) != 1 {
t.Fatalf("expected owned profile to resolve, got %d provider(s)", len(providers))
}
}
// An admin-owned profile may be shared across servers (admin resources are
// global), so an admin profile resolves regardless of the server owner.
func TestGetDDNSProvidersAllowsAdminOwnedProfile(t *testing.T) {
replaceUserInfoMapForSecurityTest(t, map[uint64]model.UserInfo{
1: {Role: model.RoleAdmin}, // admin / profile owner
100: {Role: model.RoleMember},
})
adminProfile := &model.DDNSProfile{
Common: model.Common{ID: 9, UserID: 1},
Provider: model.ProviderDummy,
Name: "admin-profile",
}
dc := newDDNSClassForTest(adminProfile)
providers, err := dc.GetDDNSProvidersFromProfiles([]uint64{9}, &model.IP{}, 100)
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if len(providers) != 1 {
t.Fatalf("expected admin-owned profile to resolve, got %d provider(s)", len(providers))
}
}
// GHSA-39g2-8x68-pmx8 (UserID==0 variant): userIsAdmin(0) returns true as a
// "system resource" shortcut, but a DDNS profile with UserID==0 is not a real
// admin grant — it is a migration/default-value artifact. A foreign server
// owner must NOT be able to drive an update through such a profile, so the
// worker must skip a UserID==0 profile that the caller does not own.
func TestGetDDNSProvidersSkipsUnownedZeroUserProfile(t *testing.T) {
replaceUserInfoMapForSecurityTest(t, map[uint64]model.UserInfo{
100: {Role: model.RoleMember}, // attacker / server owner
})
orphanProfile := &model.DDNSProfile{
Common: model.Common{ID: 3, UserID: 0},
Provider: model.ProviderDummy,
Name: "orphan-profile",
AccessSecret: "orphan-secret",
}
dc := newDDNSClassForTest(orphanProfile)
providers, err := dc.GetDDNSProvidersFromProfiles([]uint64{3}, &model.IP{}, 100)
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if len(providers) != 0 {
t.Fatalf("expected UserID==0 foreign profile to be skipped, got %d provider(s)", len(providers))
}
}
service/singleton/frontend-templates.yaml
+18 -4
View File
@@ -2,20 +2,34 @@
name: "OfficialAdmin"
repository: "https://github.com/nezhahq/admin-frontend"
author: "nezhahq"
version: "v2.0.6"
version: "v2.3.4"
is_admin: true
is_official: true
- path: "user-dist"
name: "Official"
repository: "https://github.com/hamster1963/nezha-dash-v2"
author: "hamster1963"
version: "v2.0.3"
version: "v2.4.2"
is_official: true
- path: "nezha-pixel-dist"
name: "Nezha-Pixel"
repository: "https://github.com/karllao/nezha-pixel"
author: "karllao"
version: "v1.6.0"
# Third-party user themes consume the opaque Server.PublicNote field. Theme
# maintainers must validate URL schemes (including after decoding) before using
# values such as customData.orderLink in href or window.open; the Dashboard
# backend and admin frontend do not execute those fields.
- path: "nazhua-dist"
name: "Nazhua"
repository: "https://github.com/hi2shark/nazhua"
author: "hi2hi"
version: "v0.9.1"
author: "hi2shark"
version: "v1.2.0"
- path: "aobobo-dist"
name: "Aobobo"
repository: "https://github.com/hi2shark/aobobo"
author: "hi2shark"
version: "v1.5.1"
- path: "nezha-ascii-dist"
name: "Nezha-ASCII"
repository: "https://github.com/hamster1963/nezha-ascii"
service/singleton/jwt_session.go
+55
View File
@@ -0,0 +1,55 @@
package singleton
import (
"log"
"time"
"github.com/nezhahq/nezha/model"
)
const (
JWTSessionGCSchedule = "@every 10m"
JWTSessionRevokedRetention = 24 * time.Hour
JWTSessionExpiredGrace = 1 * time.Hour
)
func StartJWTSessionGC() error {
if _, err := CronShared.AddFunc(JWTSessionGCSchedule, RunJWTSessionGC); err != nil {
return err
}
RunJWTSessionGC()
return nil
}
func RunJWTSessionGC() {
if DB == nil {
return
}
now := time.Now()
if err := DB.
Where("expires_at < ?", now.Add(-JWTSessionExpiredGrace)).
Delete(&model.JWTSession{}).Error; err != nil {
log.Printf("NEZHA>> JWTSession GC delete expired failed: %v", err)
}
if err := DB.
Where("revoked_at IS NOT NULL AND revoked_at < ?", now.Add(-JWTSessionRevokedRetention)).
Delete(&model.JWTSession{}).Error; err != nil {
log.Printf("NEZHA>> JWTSession GC delete revoked failed: %v", err)
}
}
func RevokeJWTSession(keyID string) error {
now := time.Now()
return DB.Model(&model.JWTSession{}).
Where("key_id = ? AND revoked_at IS NULL", keyID).
Update("revoked_at", &now).Error
}
func RevokeJWTSessionsByUser(userID uint64) error {
now := time.Now()
return DB.Model(&model.JWTSession{}).
Where("user_id = ? AND revoked_at IS NULL", userID).
Update("revoked_at", &now).Error
}
service/singleton/nat.go
+70
View File
@@ -2,12 +2,81 @@ package singleton
import (
"cmp"
"net"
"net/netip"
"slices"
"strings"
"github.com/nezhahq/nezha/model"
"github.com/nezhahq/nezha/pkg/utils"
)
// GHSA-x6fg-52vr-hj4w: NAT 是 commonHandler，任意认证成员可创建。newHTTPandGRPCMux
// 在分发 dashboard/gRPC 之前先按 r.Host 命中 NAT，故成员若把 Domain 设成 dashboard
// 自身 host 即可抢占全局路由（disabled 触发 DoS，enabled 把请求隧道到攻击者 agent）。
// 把 dashboard 的 InstallHost、ListenHost 以及运维声明的 ReservedHosts 列为
// 保留 host：create/update 时拒绝，启动建表时丢弃，确保补丁前已植入的恶意记录
// 在升级后不再生效。每个 host 拆成 hostname 后比较（忽略端口与大小写），反代/
// 默认端口下的端口变体也拦得住。反代部署时进程看不到对外域名，运维把它配进
// Conf.ReservedHosts（逗号分隔）即可让此处覆盖到公网入口。
func IsReservedDashboardHost(domain string) bool {
if Conf == nil {
return false
}
target := splitDashboardHostname(domain)
if target == "" {
return false
}
hosts := []string{Conf.InstallHost, Conf.DashboardHost, Conf.ListenHost}
hosts = append(hosts, strings.Split(Conf.ReservedHosts, ",")...)
for _, host := range hosts {
if reserved := splitDashboardHostname(host); reserved != "" && reserved == target {
return true
}
}
return false
}
// splitDashboardHostname 归一化为小写 hostname，对 bracketed IPv6（[::1]、
// [::1]:8008）与裸 host:port 一视同仁，避免 candidate 与 reserved 解析形态不
// 一致导致漏拦。两类等价形态也必须收敛，否则 guard 放行而 r.Host 精确命中
// 仍能劫持路由：
// - DNS absolute name 的尾点（panel.example.com. 与 panel.example.com 指向同
// 一主机），去掉单个尾点；
// - IP literal 的压缩/展开写法（::1 与 0:0:0:0:0:0:0:1），用 netip 归一到
// 规范文本。
func splitDashboardHostname(host string) string {
host = strings.ToLower(strings.TrimSpace(host))
if host == "" {
return ""
}
if h, _, err := net.SplitHostPort(host); err == nil && h != "" {
host = h
} else {
host = strings.Trim(host, "[]")
}
host = strings.TrimSuffix(host, ".")
if addr, err := netip.ParseAddr(host); err == nil {
return addr.String()
}
return host
}
// filterReservedNATProfiles 丢弃 Domain 命中 dashboard 保留 host 的 NAT 记录，
// 让 NewNATClass 启动建表时不把补丁前植入的劫持记录加载进路由表。
func filterReservedNATProfiles(in []*model.NAT) []*model.NAT {
out := in[:0]
for _, profile := range in {
if profile == nil || IsReservedDashboardHost(profile.Domain) {
continue
}
out = append(out, profile)
}
return out
}
type NATClass struct {
class[string, *model.NAT]
@@ -18,6 +87,7 @@ func NewNATClass() *NATClass {
var sortedList []*model.NAT
DB.Find(&sortedList)
sortedList = filterReservedNATProfiles(sortedList)
list := make(map[string]*model.NAT, len(sortedList))
idToDomain := make(map[uint64]string, len(sortedList))
for _, profile := range sortedList {
service/singleton/nat_reserved_host_test.go
+139
View File
@@ -0,0 +1,139 @@
package singleton
import (
"testing"
"github.com/nezhahq/nezha/model"
)
func withReservedHostConf(t *testing.T, c *model.Config) {
t.Helper()
original := Conf
Conf = &ConfigClass{Config: c}
t.Cleanup(func() { Conf = original })
}
// GHSA-x6fg-52vr-hj4w: the reserved-host check is the single source of truth
// for both the create/update guard and the startup cache filter. It must
// reject any NAT domain whose hostname collides with the dashboard's own
// InstallHost / ListenHost, regardless of port or case.
func TestIsReservedDashboardHost(t *testing.T) {
withReservedHostConf(t, &model.Config{
ConfigDashboard: model.ConfigDashboard{InstallHost: "dashboard.example:8008"},
ListenHost: "10.0.0.5",
ListenPort: 8008,
})
cases := []struct {
name string
domain string
want bool
}{
{"exact install host", "dashboard.example:8008", true},
{"install host case-insensitive", "Dashboard.Example:8008", true},
{"install host without port", "dashboard.example", true},
{"install host arbitrary port", "dashboard.example:8443", true},
{"listen host and port", "10.0.0.5:8008", true},
{"listen host bare", "10.0.0.5", true},
{"unrelated domain", "tunnel.member.example", false},
{"empty domain", "", false},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
if got := IsReservedDashboardHost(tc.domain); got != tc.want {
t.Fatalf("IsReservedDashboardHost(%q) = %v, want %v", tc.domain, got, tc.want)
}
})
}
}
// GHSA-x6fg-52vr-hj4w (reverse-proxy coverage): InstallHost/ListenHost alone
// cannot cover a dashboard reached through a reverse proxy on a public domain
// that the dashboard process never sees. ReservedHosts lets the operator
// declare those extra hostnames (comma-separated) so members still cannot
// register a NAT domain that collides with the public entry point.
func TestIsReservedDashboardHostHonoursReservedHostsList(t *testing.T) {
withReservedHostConf(t, &model.Config{
ConfigDashboard: model.ConfigDashboard{
InstallHost: "internal.example:8008",
ReservedHosts: "panel.example.com, Admin.Example.COM:443 , ",
},
})
reserved := []string{
"panel.example.com",
"panel.example.com:8443",
"admin.example.com",
"ADMIN.EXAMPLE.COM:443",
"internal.example",
}
for _, d := range reserved {
if !IsReservedDashboardHost(d) {
t.Errorf("IsReservedDashboardHost(%q) = false, want true (declared reserved host)", d)
}
}
if IsReservedDashboardHost("tunnel.member.example") {
t.Error("unrelated member domain must not be reserved")
}
if IsReservedDashboardHost("") {
t.Error("empty domain must not be reserved")
}
}
// The startup cache must not load a NAT record whose domain is reserved, so a
// malicious record planted before the patch cannot keep hijacking dashboard
// routing after upgrade. filterReservedNATProfiles is the gate NewNATClass
// runs over the DB result set.
func TestFilterReservedNATProfilesDropsReserved(t *testing.T) {
withReservedHostConf(t, &model.Config{
ConfigDashboard: model.ConfigDashboard{InstallHost: "dashboard.example:8008"},
})
in := []*model.NAT{
{Common: model.Common{ID: 1}, Domain: "dashboard.example", Enabled: true},
{Common: model.Common{ID: 2}, Domain: "tunnel.member.example", Enabled: true},
{Common: model.Common{ID: 3}, Domain: "Dashboard.Example:9999", Enabled: false},
}
out := filterReservedNATProfiles(in)
if len(out) != 1 {
t.Fatalf("expected only the non-reserved profile to survive, got %d", len(out))
}
if out[0].Domain != "tunnel.member.example" {
t.Fatalf("surviving profile must be the member tunnel, got %q", out[0].Domain)
}
}
// GHSA-x6fg-52vr-hj4w (canonical-host coverage): the routing match is an exact
// lookup on r.Host, so a member who registers a NAT Domain that is a DNS/IP
// *equivalent* of the dashboard host — but a different literal string — still
// hijacks the matching r.Host. The guard must collapse the trailing DNS dot and
// the IPv6 compressed/expanded forms, or these variants slip past create/update.
func TestIsReservedDashboardHostCollapsesEquivalentForms(t *testing.T) {
withReservedHostConf(t, &model.Config{
ConfigDashboard: model.ConfigDashboard{
InstallHost: "panel.example.com",
ReservedHosts: "[::1]:8008",
},
})
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
}
for _, d := range reserved {
if !IsReservedDashboardHost(d) {
t.Errorf("IsReservedDashboardHost(%q) = false, want true (equivalent of reserved host)", d)
}
}
if IsReservedDashboardHost("tunnel.member.example.") {
t.Error("unrelated member domain with trailing dot must not be reserved")
}
}
service/singleton/security_regression_test.go
+940
View File
@@ -0,0 +1,940 @@
package singleton
import (
"context"
"net/http/httptest"
"slices"
"testing"
"time"
"github.com/gin-gonic/gin"
"github.com/patrickmn/go-cache"
"github.com/robfig/cron/v3"
"gorm.io/driver/sqlite"
"gorm.io/gorm"
"github.com/nezhahq/nezha/model"
pb "github.com/nezhahq/nezha/proto"
"google.golang.org/grpc/metadata"
)
type capturedTaskStream struct {
tasks chan *pb.Task
}
func newCapturedTaskStream() *capturedTaskStream {
return &capturedTaskStream{tasks: make(chan *pb.Task, 4)}
}
func (s *capturedTaskStream) Send(task *pb.Task) error {
s.tasks <- task
return nil
}
func (s *capturedTaskStream) Recv() (*pb.TaskResult, error) { return nil, context.Canceled }
func (s *capturedTaskStream) SetHeader(metadata.MD) error { return nil }
func (s *capturedTaskStream) SendHeader(metadata.MD) error { return nil }
func (s *capturedTaskStream) SetTrailer(metadata.MD) {}
func (s *capturedTaskStream) Context() context.Context { return context.Background() }
func (s *capturedTaskStream) SendMsg(any) error { return nil }
func (s *capturedTaskStream) RecvMsg(any) error { return context.Canceled }
// withTaskStream attaches a TaskStream to a freshly constructed Server using the
// new atomic accessor. The field itself is unexported (see Fix #12) precisely
// because direct struct-literal access invited torn interface reads on hot
// paths — tests use this helper rather than reaching in, mirroring production
// callsites.
func withTaskStream(s *model.Server, stream pb.NezhaService_RequestTaskServer) *model.Server {
s.SetTaskStream(stream)
return s
}
func replaceServerSharedForSecurityTest(t *testing.T, servers ...*model.Server) {
t.Helper()
original := ServerShared
serverClass := &ServerClass{
class: class[uint64, *model.Server]{
list: make(map[uint64]*model.Server),
},
uuidToID: make(map[string]uint64),
}
for _, server := range servers {
serverClass.list[server.ID] = server
}
ServerShared = serverClass
t.Cleanup(func() { ServerShared = original })
}
func replaceUserInfoMapForSecurityTest(t *testing.T, users map[uint64]model.UserInfo) {
t.Helper()
UserLock.Lock()
original := UserInfoMap
UserInfoMap = users
UserLock.Unlock()
t.Cleanup(func() {
UserLock.Lock()
UserInfoMap = original
UserLock.Unlock()
})
}
func TestCronTriggerSkipsServersOwnedByOtherUsers(t *testing.T) {
firstStream := newCapturedTaskStream()
secondStream := newCapturedTaskStream()
replaceServerSharedForSecurityTest(t,
withTaskStream(&model.Server{Common: model.Common{ID: 1, UserID: 100}, Name: "member-server"}, firstStream),
withTaskStream(&model.Server{Common: model.Common{ID: 2, UserID: 200}, Name: "admin-server"}, secondStream),
)
cronTask := &model.Cron{
Common: model.Common{ID: 99, UserID: 100},
Command: "id",
Cover: model.CronCoverAll,
Servers: []uint64{},
}
CronTrigger(cronTask)()
assertTaskCommand(t, firstStream, "id")
assertNoTask(t, secondStream)
}
func TestSendTriggerTasksSkipsCronOwnedByAnotherUser(t *testing.T) {
attackerStream := newCapturedTaskStream()
replaceServerSharedForSecurityTest(t,
withTaskStream(&model.Server{Common: model.Common{ID: 7, UserID: 200}, Name: "attacker-server"}, attackerStream),
)
replaceUserInfoMapForSecurityTest(t, map[uint64]model.UserInfo{
1: {Role: model.RoleAdmin},
200: {Role: model.RoleMember},
})
adminCron := &model.Cron{
Common: model.Common{ID: 42, UserID: 1},
Command: "admin-maintenance",
Cover: model.CronCoverAlertTrigger,
}
cronClass := &CronClass{
class: class[uint64, *model.Cron]{
list: map[uint64]*model.Cron{adminCron.ID: adminCron},
},
}
cronClass.SendTriggerTasks([]uint64{adminCron.ID}, 7, 200)
assertNoTask(t, attackerStream)
}
func assertTaskCommand(t *testing.T, stream *capturedTaskStream, expectedCommand string) {
t.Helper()
select {
case task := <-stream.tasks:
if task.GetType() != model.TaskTypeCommand {
t.Fatalf("expected command task type, got %v", task.GetType())
}
if task.GetData() != expectedCommand {
t.Fatalf("expected command %q, got %q", expectedCommand, task.GetData())
}
case <-time.After(time.Second):
t.Fatalf("expected command %q to be sent", expectedCommand)
}
}
func assertNoTask(t *testing.T, stream *capturedTaskStream) {
t.Helper()
select {
case task := <-stream.tasks:
t.Fatalf("expected no task to be sent, got command %q", task.GetData())
case <-time.After(50 * time.Millisecond):
}
}
func TestCronTriggerSendsToMemberOwnedServer(t *testing.T) {
memberStream := newCapturedTaskStream()
replaceServerSharedForSecurityTest(t,
withTaskStream(&model.Server{Common: model.Common{ID: 1, UserID: 100}, Name: "member-server"}, memberStream),
)
replaceUserInfoMapForSecurityTest(t, map[uint64]model.UserInfo{
100: {Role: model.RoleMember},
})
cronTask := &model.Cron{
Common: model.Common{ID: 99, UserID: 100},
Command: "id",
Cover: model.CronCoverAll,
}
CronTrigger(cronTask)()
assertTaskCommand(t, memberStream, "id")
}
func TestCronTriggerAdminCronFansOutAcrossOwners(t *testing.T) {
first := newCapturedTaskStream()
second := newCapturedTaskStream()
replaceServerSharedForSecurityTest(t,
withTaskStream(&model.Server{Common: model.Common{ID: 1, UserID: 100}, Name: "member-server"}, first),
withTaskStream(&model.Server{Common: model.Common{ID: 2, UserID: 200}, Name: "admin-server"}, second),
)
replaceUserInfoMapForSecurityTest(t, map[uint64]model.UserInfo{
1: {Role: model.RoleAdmin},
100: {Role: model.RoleMember},
200: {Role: model.RoleAdmin},
})
cronTask := &model.Cron{
Common: model.Common{ID: 99, UserID: 1},
Command: "maintenance",
Cover: model.CronCoverAll,
}
CronTrigger(cronTask)()
assertTaskCommand(t, first, "maintenance")
assertTaskCommand(t, second, "maintenance")
}
func TestCronTriggerLegacyZeroOwnerFansOut(t *testing.T) {
first := newCapturedTaskStream()
replaceServerSharedForSecurityTest(t,
withTaskStream(&model.Server{Common: model.Common{ID: 1, UserID: 100}, Name: "member-server"}, first),
)
replaceUserInfoMapForSecurityTest(t, map[uint64]model.UserInfo{
100: {Role: model.RoleMember},
})
cronTask := &model.Cron{
Common: model.Common{ID: 99, UserID: 0},
Command: "legacy",
Cover: model.CronCoverAll,
}
CronTrigger(cronTask)()
assertTaskCommand(t, first, "legacy")
}
func TestCronTriggerSkipsServersWhenOwnerNotKnown(t *testing.T) {
stream := newCapturedTaskStream()
replaceServerSharedForSecurityTest(t,
withTaskStream(&model.Server{Common: model.Common{ID: 1, UserID: 100}, Name: "member-server"}, stream),
)
replaceUserInfoMapForSecurityTest(t, map[uint64]model.UserInfo{
100: {Role: model.RoleMember},
})
cronTask := &model.Cron{
Common: model.Common{ID: 99, UserID: 999},
Command: "ghost",
Cover: model.CronCoverAll,
}
CronTrigger(cronTask)()
assertNoTask(t, stream)
}
func TestSendTriggerTasksAllowsSelfOwnedCron(t *testing.T) {
stream := newCapturedTaskStream()
replaceServerSharedForSecurityTest(t,
withTaskStream(&model.Server{Common: model.Common{ID: 7, UserID: 200}, Name: "member-server"}, stream),
)
replaceUserInfoMapForSecurityTest(t, map[uint64]model.UserInfo{
200: {Role: model.RoleMember},
})
memberCron := &model.Cron{
Common: model.Common{ID: 42, UserID: 200},
Command: "member-task",
Cover: model.CronCoverAlertTrigger,
}
cronClass := &CronClass{
class: class[uint64, *model.Cron]{
list: map[uint64]*model.Cron{memberCron.ID: memberCron},
},
}
cronClass.SendTriggerTasks([]uint64{memberCron.ID}, 7, 200)
assertTaskCommand(t, stream, "member-task")
}
func TestSendTriggerTasksAllowsAdminCallerToTriggerAny(t *testing.T) {
stream := newCapturedTaskStream()
replaceServerSharedForSecurityTest(t,
withTaskStream(&model.Server{Common: model.Common{ID: 9, UserID: 100}, Name: "any-server"}, stream),
)
replaceUserInfoMapForSecurityTest(t, map[uint64]model.UserInfo{
1: {Role: model.RoleAdmin},
100: {Role: model.RoleMember},
})
memberCron := &model.Cron{
Common: model.Common{ID: 42, UserID: 100},
Command: "member-task",
Cover: model.CronCoverAlertTrigger,
}
cronClass := &CronClass{
class: class[uint64, *model.Cron]{
list: map[uint64]*model.Cron{memberCron.ID: memberCron},
},
}
cronClass.SendTriggerTasks([]uint64{memberCron.ID}, 9, 1)
assertTaskCommand(t, stream, "member-task")
}
func TestSendTriggerTasksIgnoresUnknownTaskIDs(t *testing.T) {
stream := newCapturedTaskStream()
replaceServerSharedForSecurityTest(t,
withTaskStream(&model.Server{Common: model.Common{ID: 7, UserID: 200}, Name: "member-server"}, stream),
)
replaceUserInfoMapForSecurityTest(t, map[uint64]model.UserInfo{
200: {Role: model.RoleMember},
})
cronClass := &CronClass{
class: class[uint64, *model.Cron]{
list: map[uint64]*model.Cron{},
},
}
cronClass.SendTriggerTasks([]uint64{12345}, 7, 200)
cronClass.SendTriggerTasks(nil, 7, 200)
assertNoTask(t, stream)
}
func TestSendTriggerTasksMixedCronIDsOnlyFiresAllowed(t *testing.T) {
stream := newCapturedTaskStream()
replaceServerSharedForSecurityTest(t,
withTaskStream(&model.Server{Common: model.Common{ID: 7, UserID: 200}, Name: "member-server"}, stream),
)
replaceUserInfoMapForSecurityTest(t, map[uint64]model.UserInfo{
1: {Role: model.RoleAdmin},
200: {Role: model.RoleMember},
})
memberCron := &model.Cron{
Common: model.Common{ID: 7, UserID: 200},
Command: "member-task",
Cover: model.CronCoverAlertTrigger,
}
adminCron := &model.Cron{
Common: model.Common{ID: 8, UserID: 1},
Command: "admin-task",
Cover: model.CronCoverAlertTrigger,
}
cronClass := &CronClass{
class: class[uint64, *model.Cron]{
list: map[uint64]*model.Cron{memberCron.ID: memberCron, adminCron.ID: adminCron},
},
}
cronClass.SendTriggerTasks([]uint64{memberCron.ID, adminCron.ID}, 7, 200)
select {
case task := <-stream.tasks:
if task.GetData() != "member-task" {
t.Fatalf("expected member-task, got %q", task.GetData())
}
case <-time.After(time.Second):
t.Fatalf("expected member-task to be sent")
}
assertNoTask(t, stream)
}
func TestAlertTriggerCronResultAuthorizationConsumesOneDispatch(t *testing.T) {
cronClass := &CronClass{}
cronClass.reserveAlertTriggerCronResult(42, 7)
cronClass.reserveAlertTriggerCronResult(42, 7)
if !cronClass.consumeAlertTriggerCronResult(42, 7) {
t.Fatal("expected first alert-trigger authorization to be consumed")
}
if !cronClass.consumeAlertTriggerCronResult(42, 7) {
t.Fatal("expected second alert-trigger authorization to be consumed")
}
if cronClass.consumeAlertTriggerCronResult(42, 7) {
t.Fatal("expected alert-trigger authorization to be consumed only once per dispatch")
}
}
func TestAlertTriggerCronResultAuthorizationExpires(t *testing.T) {
cronClass := &CronClass{
pendingAlertTriggerTasks: map[uint64]map[uint64][]time.Time{
42: {7: {time.Now().Add(-time.Second)}},
},
}
if cronClass.consumeAlertTriggerCronResult(42, 7) {
t.Fatal("expired alert-trigger authorization must not be accepted")
}
if len(cronClass.pendingAlertTriggerTasks) != 0 {
t.Fatal("expired alert-trigger authorization must be pruned")
}
}
func TestAlertTriggerCronResultAuthorizationRevokeRemovesLatestDispatch(t *testing.T) {
existingAuthorizationExpiresAt := time.Now().Add(time.Hour)
cronClass := &CronClass{
pendingAlertTriggerTasks: map[uint64]map[uint64][]time.Time{
42: {7: {existingAuthorizationExpiresAt}},
},
}
cronClass.reserveAlertTriggerCronResult(42, 7)
cronClass.revokeAlertTriggerCronResult(42, 7)
authorizations := cronClass.pendingAlertTriggerTasks[42][7]
if len(authorizations) != 1 {
t.Fatalf("expected one previous alert-trigger authorization to remain, got %d", len(authorizations))
}
if !authorizations[0].Equal(existingAuthorizationExpiresAt) {
t.Fatal("send failure rollback must remove the newest reserved authorization")
}
}
func TestCronClassUpdatePrunesAlertTriggerCronResultAuthorization(t *testing.T) {
cronClass := &CronClass{
Cron: cron.New(cron.WithSeconds()),
class: class[uint64, *model.Cron]{
list: map[uint64]*model.Cron{42: {Common: model.Common{ID: 42}}},
},
pendingAlertTriggerTasks: map[uint64]map[uint64][]time.Time{
42: {7: {time.Now().Add(time.Hour)}},
},
}
cronClass.Update(&model.Cron{Common: model.Common{ID: 42}})
if len(cronClass.pendingAlertTriggerTasks) != 0 {
t.Fatal("cron update must prune old alert-trigger result authorizations")
}
}
func TestCronClassDeletePrunesAlertTriggerCronResultAuthorization(t *testing.T) {
cronClass := &CronClass{
Cron: cron.New(cron.WithSeconds()),
class: class[uint64, *model.Cron]{
list: map[uint64]*model.Cron{42: {Common: model.Common{ID: 42}}},
},
pendingAlertTriggerTasks: map[uint64]map[uint64][]time.Time{
42: {7: {time.Now().Add(time.Hour)}},
},
}
cronClass.Delete([]uint64{42})
if len(cronClass.pendingAlertTriggerTasks) != 0 {
t.Fatal("cron delete must prune alert-trigger result authorizations")
}
}
// CanReportCronResult is the cron-side dual of canReportServiceResult: it gates
// agent-reported TaskTypeCommand results to only the cron/server pairs the
// dashboard actually fanned the task out to. Without these inbound checks any
// authenticated agent could fabricate a TaskResult for an arbitrary cron ID and
// poison LastResult / fire success/failure notifications belonging to another
// tenant. The tests below pin each Cover branch end-to-end against the dispatch
// logic in CronTrigger so the two sides stay symmetric.
func TestCanReportCronResultRejectsNilCronOrReporter(t *testing.T) {
cr := &model.Cron{Common: model.Common{ID: 7, UserID: 100}, Cover: model.CronCoverAll}
reporter := &model.Server{Common: model.Common{ID: 1, UserID: 100}}
if CanReportCronResult(nil, reporter) {
t.Fatal("nil cron must be rejected — would dereference inside cover branches")
}
if CanReportCronResult(cr, nil) {
t.Fatal("nil reporter must be rejected")
}
}
func TestCanReportCronResultRejectsForeignReporter(t *testing.T) {
replaceUserInfoMapForSecurityTest(t, map[uint64]model.UserInfo{
100: {Role: model.RoleMember},
200: {Role: model.RoleMember},
})
cr := &model.Cron{
Common: model.Common{ID: 7, UserID: 100},
Cover: model.CronCoverAll,
}
foreign := &model.Server{Common: model.Common{ID: 1, UserID: 200}}
if CanReportCronResult(cr, foreign) {
t.Fatal("foreign-user reporter must be rejected: CronTrigger never dispatched to it")
}
}
func TestCanReportCronResultCronCoverAllRejectsReporterInDenyList(t *testing.T) {
cr := &model.Cron{
Common: model.Common{ID: 7, UserID: 100},
Cover: model.CronCoverAll,
Servers: []uint64{1},
}
reporter := &model.Server{Common: model.Common{ID: 1, UserID: 100}}
if CanReportCronResult(cr, reporter) {
t.Fatal("CronCoverAll treats Servers as deny-list; reporter in the list must be rejected")
}
}
func TestCanReportCronResultCronCoverAllAcceptsReporterNotInDenyList(t *testing.T) {
cr := &model.Cron{
Common: model.Common{ID: 7, UserID: 100},
Cover: model.CronCoverAll,
Servers: []uint64{99},
}
reporter := &model.Server{Common: model.Common{ID: 1, UserID: 100}}
if !CanReportCronResult(cr, reporter) {
t.Fatal("CronCoverAll with reporter NOT in Servers must accept — CronTrigger dispatches to it")
}
}
func TestCanReportCronResultCronCoverIgnoreAllAcceptsReporterInAllowList(t *testing.T) {
cr := &model.Cron{
Common: model.Common{ID: 7, UserID: 100},
Cover: model.CronCoverIgnoreAll,
Servers: []uint64{1},
}
reporter := &model.Server{Common: model.Common{ID: 1, UserID: 100}}
if !CanReportCronResult(cr, reporter) {
t.Fatal("CronCoverIgnoreAll treats Servers as allow-list; reporter in the list must be accepted")
}
}
func TestCanReportCronResultCronCoverIgnoreAllRejectsReporterOutsideAllowList(t *testing.T) {
cr := &model.Cron{
Common: model.Common{ID: 7, UserID: 100},
Cover: model.CronCoverIgnoreAll,
Servers: []uint64{99},
}
reporter := &model.Server{Common: model.Common{ID: 1, UserID: 100}}
if CanReportCronResult(cr, reporter) {
t.Fatal("CronCoverIgnoreAll with reporter NOT in Servers must reject — CronTrigger never dispatched to it")
}
}
// failingTaskStream simulates a TaskStream whose Send always errors. CronTrigger
// uses this signal to revoke a reserved alert-trigger authorization, so the
// agent can't later attach to the cron via CanReportCronResult based on a
// dispatch that never actually reached the wire.
type failingTaskStream struct {
capturedTaskStream
sendErr error
}
func newFailingTaskStream(err error) *failingTaskStream {
return &failingTaskStream{
capturedTaskStream: capturedTaskStream{tasks: make(chan *pb.Task, 4)},
sendErr: err,
}
}
func (s *failingTaskStream) Send(task *pb.Task) error {
s.tasks <- task
return s.sendErr
}
func TestCronTriggerRevokesAlertTriggerAuthorizationOnSendFailure(t *testing.T) {
failing := newFailingTaskStream(context.Canceled)
replaceServerSharedForSecurityTest(t,
withTaskStream(&model.Server{Common: model.Common{ID: 7, UserID: 100}, Name: "broken-server"}, failing),
)
replaceUserInfoMapForSecurityTest(t, map[uint64]model.UserInfo{
100: {Role: model.RoleMember},
})
originalCronShared := CronShared
t.Cleanup(func() { CronShared = originalCronShared })
CronShared = &CronClass{
class: class[uint64, *model.Cron]{list: map[uint64]*model.Cron{}},
pendingAlertTriggerTasks: map[uint64]map[uint64][]time.Time{},
}
cr := &model.Cron{
Common: model.Common{ID: 42, UserID: 100},
Cover: model.CronCoverAlertTrigger,
}
CronTrigger(cr, 7)()
// drain the dispatched task — Send error is what we care about, not the payload
select {
case <-failing.tasks:
case <-time.After(time.Second):
t.Fatal("expected CronTrigger to call Send before reacting to the error")
}
if CronShared.consumeAlertTriggerCronResult(42, 7) {
t.Fatal("Send failure must revoke the reserved alert-trigger authorization; otherwise a foreign agent could later report a result for a dispatch that never reached the wire")
}
if len(CronShared.pendingAlertTriggerTasks) != 0 {
t.Fatalf("expected pendingAlertTriggerTasks to be empty after revoke, got %d entries", len(CronShared.pendingAlertTriggerTasks))
}
}
func TestClassCheckPermission(t *testing.T) {
replaceUserInfoMapForSecurityTest(t, map[uint64]model.UserInfo{
1: {Role: model.RoleAdmin},
200: {Role: model.RoleMember},
})
sharedClass := &ServerClass{
class: class[uint64, *model.Server]{
list: map[uint64]*model.Server{
1: {Common: model.Common{ID: 1, UserID: 200}},
2: {Common: model.Common{ID: 2, UserID: 1}},
},
},
uuidToID: map[string]uint64{},
}
memberCtx, _ := gin.CreateTestContext(httptest.NewRecorder())
memberCtx.Set(model.CtxKeyAuthorizedUser, &model.User{
Common: model.Common{ID: 200},
Role: model.RoleMember,
})
adminCtx, _ := gin.CreateTestContext(httptest.NewRecorder())
adminCtx.Set(model.CtxKeyAuthorizedUser, &model.User{
Common: model.Common{ID: 1},
Role: model.RoleAdmin,
})
if !sharedClass.CheckPermission(memberCtx, slices.Values([]uint64{1})) {
t.Fatal("expected member to access own resource")
}
if sharedClass.CheckPermission(memberCtx, slices.Values([]uint64{2})) {
t.Fatal("expected member to be denied foreign resource")
}
if !sharedClass.CheckPermission(memberCtx, slices.Values([]uint64{})) {
t.Fatal("expected empty iterator to be allowed")
}
if !sharedClass.CheckPermission(memberCtx, slices.Values([]uint64{999})) {
t.Fatal("expected unknown id to be ignored (vacuous true)")
}
if !sharedClass.CheckPermission(adminCtx, slices.Values([]uint64{1, 2})) {
t.Fatal("expected admin to access any resource")
}
}
func TestServiceMonitorResultSkipsReporterOutsideServiceCover(t *testing.T) {
ss := newServiceMonitorSecurityHarness(t,
&model.Server{Common: model.Common{ID: 1, UserID: 100}, Name: "covered-server"},
&model.Server{Common: model.Common{ID: 2, UserID: 100}, Name: "uncovered-server"},
)
addServiceMonitorSecurityService(t, ss, &model.Service{
Common: model.Common{ID: 10, UserID: 100},
Name: "selected-only-service",
Type: model.TaskTypeTCPPing,
Target: "example.invalid:443",
Duration: 3600,
Cover: model.ServiceCoverIgnoreAll,
SkipServers: map[uint64]bool{1: true},
})
ss.Dispatch(serviceMonitorResult(2, 10, model.TaskTypeTCPPing, true))
ss.Dispatch(serviceMonitorResult(1, 10, model.TaskTypeTCPPing, true))
waitForServiceHistory(t, 10, 1)
assertNoServiceHistory(t, 10, 2)
}
func TestServiceMonitorResultSkipsCoveredReporterOwnedByAnotherUser(t *testing.T) {
ss := newServiceMonitorSecurityHarness(t,
&model.Server{Common: model.Common{ID: 1, UserID: 100}, Name: "owner-server"},
&model.Server{Common: model.Common{ID: 2, UserID: 200}, Name: "foreign-server"},
)
addServiceMonitorSecurityService(t, ss, &model.Service{
Common: model.Common{ID: 10, UserID: 100},
Name: "owner-only-service",
Type: model.TaskTypeTCPPing,
Target: "example.invalid:443",
Duration: 3600,
Cover: model.ServiceCoverIgnoreAll,
SkipServers: map[uint64]bool{1: true, 2: true},
})
ss.Dispatch(serviceMonitorResult(2, 10, model.TaskTypeTCPPing, true))
ss.Dispatch(serviceMonitorResult(1, 10, model.TaskTypeTCPPing, true))
waitForServiceHistory(t, 10, 1)
assertNoServiceHistory(t, 10, 2)
}
func TestServiceMonitorResultSkipsMismatchedTaskType(t *testing.T) {
ss := newServiceMonitorSecurityHarness(t,
&model.Server{Common: model.Common{ID: 1, UserID: 100}, Name: "owner-server"},
)
addServiceMonitorSecurityService(t, ss, &model.Service{
Common: model.Common{ID: 10, UserID: 100},
Name: "http-service",
Type: model.TaskTypeHTTPGet,
Target: "https://example.invalid",
Duration: 3600,
Cover: model.ServiceCoverIgnoreAll,
SkipServers: map[uint64]bool{1: true},
})
ss.Dispatch(serviceMonitorResult(1, 10, model.TaskTypeTCPPing, false))
ss.Dispatch(serviceMonitorResult(1, 10, model.TaskTypeHTTPGet, true))
waitForTodayStats(t, ss, 10, 1, 0)
}
func TestServiceMonitorResultSkipsUnknownReporter(t *testing.T) {
ss := newServiceMonitorSecurityHarness(t,
&model.Server{Common: model.Common{ID: 1, UserID: 100}, Name: "owner-server"},
)
addServiceMonitorSecurityService(t, ss, &model.Service{
Common: model.Common{ID: 10, UserID: 100},
Name: "known-reporter-service",
Type: model.TaskTypeTCPPing,
Target: "example.invalid:443",
Duration: 3600,
Cover: model.ServiceCoverIgnoreAll,
SkipServers: map[uint64]bool{1: true},
})
ss.Dispatch(serviceMonitorResult(999, 10, model.TaskTypeTCPPing, true))
ss.Dispatch(serviceMonitorResult(1, 10, model.TaskTypeTCPPing, true))
waitForServiceHistory(t, 10, 1)
assertNoServiceHistory(t, 10, 999)
}
func TestServiceMonitorResultAllowsCoveredReporterOwnedByServiceOwner(t *testing.T) {
ss := newServiceMonitorSecurityHarness(t,
&model.Server{Common: model.Common{ID: 1, UserID: 100}, Name: "owner-server"},
)
addServiceMonitorSecurityService(t, ss, &model.Service{
Common: model.Common{ID: 10, UserID: 100},
Name: "owner-service",
Type: model.TaskTypeTCPPing,
Target: "example.invalid:443",
Duration: 3600,
Cover: model.ServiceCoverIgnoreAll,
SkipServers: map[uint64]bool{1: true},
})
ss.Dispatch(serviceMonitorResult(1, 10, model.TaskTypeTCPPing, true))
waitForServiceHistory(t, 10, 1)
}
func TestServiceMonitorResultAllowsCoveredReporterForAdminOwnedService(t *testing.T) {
replaceUserInfoMapForSecurityTest(t, map[uint64]model.UserInfo{
1: {Role: model.RoleAdmin},
200: {Role: model.RoleMember},
})
ss := newServiceMonitorSecurityHarness(t,
&model.Server{Common: model.Common{ID: 2, UserID: 200}, Name: "member-server"},
)
addServiceMonitorSecurityService(t, ss, &model.Service{
Common: model.Common{ID: 10, UserID: 1},
Name: "admin-service",
Type: model.TaskTypeTCPPing,
Target: "example.invalid:443",
Duration: 3600,
Cover: model.ServiceCoverIgnoreAll,
SkipServers: map[uint64]bool{2: true},
})
ss.Dispatch(serviceMonitorResult(2, 10, model.TaskTypeTCPPing, true))
waitForServiceHistory(t, 10, 2)
}
func newServiceMonitorSecurityHarness(t *testing.T, servers ...*model.Server) *ServiceSentinel {
t.Helper()
originalDB := DB
originalConf := Conf
originalCache := Cache
originalCronShared := CronShared
originalServerShared := ServerShared
originalServiceSentinelShared := ServiceSentinelShared
originalNotificationShared := NotificationShared
originalTSDBShared := TSDBShared
originalLoc := Loc
var sqlDBClose func() error
t.Cleanup(func() {
DB = originalDB
Conf = originalConf
Cache = originalCache
CronShared = originalCronShared
ServerShared = originalServerShared
ServiceSentinelShared = originalServiceSentinelShared
NotificationShared = originalNotificationShared
TSDBShared = originalTSDBShared
Loc = originalLoc
if sqlDBClose != nil {
_ = sqlDBClose()
}
})
db, err := gorm.Open(sqlite.Open("file:"+t.Name()+"?mode=memory&cache=shared"), &gorm.Config{})
if err != nil {
t.Fatal(err)
}
sqlDB, err := db.DB()
if err != nil {
t.Fatal(err)
}
sqlDB.SetMaxOpenConns(1)
sqlDBClose = sqlDB.Close
DB = db
if err := DB.AutoMigrate(
model.Server{},
model.Service{},
model.ServiceHistory{},
model.Notification{},
model.NotificationGroup{},
model.NotificationGroupNotification{},
); err != nil {
t.Fatal(err)
}
Conf = &ConfigClass{Config: &model.Config{AvgPingCount: 1}}
Cache = cache.New(time.Minute, time.Minute)
CronShared = &CronClass{
Cron: cron.New(cron.WithSeconds()),
class: class[uint64, *model.Cron]{list: map[uint64]*model.Cron{}},
}
NotificationShared = &NotificationClass{
class: class[uint64, *model.Notification]{list: map[uint64]*model.Notification{}},
groupToIDList: map[uint64]map[uint64]*model.Notification{},
idToGroupList: map[uint64]map[uint64]struct{}{},
groupList: map[uint64]string{},
}
TSDBShared = nil
Loc = time.UTC
serverClass := &ServerClass{
class: class[uint64, *model.Server]{
list: make(map[uint64]*model.Server),
},
uuidToID: make(map[string]uint64),
}
for _, server := range servers {
serverClass.list[server.ID] = server
}
ServerShared = serverClass
bus := make(chan *model.Service, 1)
ss, err := NewServiceSentinel(bus)
if err != nil {
t.Fatal(err)
}
ServiceSentinelShared = ss
// LIFO Cleanup ordering: this Close() runs BEFORE the earlier t.Cleanup that
// restores Conf/Cache/CronShared/NotificationShared/TSDBShared, so the
// worker has fully exited before we swap those globals out. Skipping this
// step causes `go test -race` to flag the write-vs-read between the
// teardown and the still-running worker.
t.Cleanup(func() { ss.Close() })
return ss
}
func addServiceMonitorSecurityService(t *testing.T, ss *ServiceSentinel, service *model.Service) {
t.Helper()
if err := DB.Create(service).Error; err != nil {
t.Fatal(err)
}
if err := ss.Update(service); err != nil {
t.Fatal(err)
}
}
func serviceMonitorResult(reporter, serviceID uint64, taskType uint8, successful bool) ReportData {
return ReportData{
Reporter: reporter,
Data: &pb.TaskResult{
Id: serviceID,
Type: uint64(taskType),
Delay: 12,
Data: "service monitor result",
Successful: successful,
},
}
}
func waitForServiceHistory(t *testing.T, serviceID, serverID uint64) {
t.Helper()
deadline := time.After(time.Second)
for {
var count int64
if err := DB.Model(&model.ServiceHistory{}).
Where("service_id = ? AND server_id = ?", serviceID, serverID).
Count(&count).Error; err != nil {
t.Fatal(err)
}
if count > 0 {
return
}
select {
case <-deadline:
t.Fatalf("expected service history for service %d from server %d", serviceID, serverID)
default:
time.Sleep(10 * time.Millisecond)
}
}
}
func assertNoServiceHistory(t *testing.T, serviceID, serverID uint64) {
t.Helper()
var count int64
if err := DB.Model(&model.ServiceHistory{}).
Where("service_id = ? AND server_id = ?", serviceID, serverID).
Count(&count).Error; err != nil {
t.Fatal(err)
}
if count != 0 {
t.Fatalf("expected no service history for service %d from server %d, got %d", serviceID, serverID, count)
}
}
func waitForTodayStats(t *testing.T, ss *ServiceSentinel, serviceID uint64, wantUp, wantDown uint64) {
t.Helper()
deadline := time.After(time.Second)
for {
ss.serviceResponseDataStoreLock.RLock()
stats := ss.serviceStatusToday[serviceID]
var up, down uint64
if stats != nil {
up = stats.Up
down = stats.Down
}
ss.serviceResponseDataStoreLock.RUnlock()
if up == wantUp && down == wantDown {
return
}
if down > wantDown {
t.Fatalf("expected service %d down count %d, got %d", serviceID, wantDown, down)
}
select {
case <-deadline:
t.Fatalf("expected service %d stats up=%d down=%d", serviceID, wantUp, wantDown)
default:
time.Sleep(10 * time.Millisecond)
}
}
}
service/singleton/server.go
+78 -7
View File
@@ -6,6 +6,7 @@ import (
"log"
"slices"
"strings"
"sync"
"github.com/nezhahq/nezha/model"
"github.com/nezhahq/nezha/pkg/ddns"
@@ -15,6 +16,10 @@ import (
type ServerClass struct {
class[uint64, *model.Server]
// lifecycleMu serializes changes to the authoritative server entries with
// synchronous ServiceSentinel report processing.
lifecycleMu sync.RWMutex
uuidToID map[string]uint64
sortedListForGuest []*model.Server
@@ -30,18 +35,67 @@ func NewServerClass() *ServerClass {
var servers []model.Server
DB.Find(&servers)
for _, s := range servers {
innerS := s
model.InitServer(&innerS)
sc.list[innerS.ID] = &innerS
for i := range servers {
innerS := &servers[i]
model.InitServer(innerS)
sc.list[innerS.ID] = innerS
sc.uuidToID[innerS.UUID] = innerS.ID
}
sc.sortList()
model.OwnerServerIDsLookup = sc.ownerServerIDs
model.AllServerIDsLookup = sc.allServerIDs
model.OwnerIsAdminLookup = ownerIsAdmin
return sc
}
func (c *ServerClass) ownerServerIDs(ownerUID uint64) []uint64 {
var ids []uint64
c.Range(func(id uint64, s *model.Server) bool {
if s != nil && s.GetUserID() == ownerUID {
ids = append(ids, id)
}
return true
})
return ids
}
func (c *ServerClass) allServerIDs() []uint64 {
var ids []uint64
c.Range(func(id uint64, s *model.Server) bool {
if s != nil {
ids = append(ids, id)
}
return true
})
return ids
}
func ownerIsAdmin(ownerUID uint64) bool {
return userIsAdmin(ownerUID)
}
func (c *ServerClass) lockLifecycleRead() {
c.lifecycleMu.RLock()
}
func (c *ServerClass) unlockLifecycleRead() {
c.lifecycleMu.RUnlock()
}
func (c *ServerClass) lockLifecycleWrite() {
c.lifecycleMu.Lock()
}
func (c *ServerClass) unlockLifecycleWrite() {
c.lifecycleMu.Unlock()
}
func (c *ServerClass) Update(s *model.Server, uuid string) {
c.lockLifecycleWrite()
defer c.unlockLifecycleWrite()
c.listMu.Lock()
c.list[s.ID] = s
@@ -61,11 +115,17 @@ func (c *ServerClass) Update(s *model.Server, uuid string) {
}
func (c *ServerClass) Delete(idList []uint64) {
c.lockLifecycleWrite()
defer c.unlockLifecycleWrite()
c.listMu.Lock()
for _, id := range idList {
serverUUID := c.list[id].UUID
delete(c.uuidToID, serverUUID)
s, ok := c.list[id]
if !ok {
continue
}
delete(c.uuidToID, s.UUID)
delete(c.list, id)
}
@@ -74,6 +134,17 @@ func (c *ServerClass) Delete(idList []uint64) {
c.sortList()
}
// setUserID updates in-memory ownership under the server lifecycle lock so a
// transfer cannot change authorization during synchronous report processing.
func (c *ServerClass) setUserID(id, userID uint64) {
c.lockLifecycleWrite()
defer c.unlockLifecycleWrite()
if s, ok := c.Get(id); ok && s != nil {
s.SetUserID(userID)
}
}
func (c *ServerClass) GetSortedListForGuest() []*model.Server {
c.sortedListMu.RLock()
defer c.sortedListMu.RUnlock()
@@ -93,7 +164,7 @@ func (c *ServerClass) UpdateDDNS(server *model.Server, ip *model.IP) error {
confServers := strings.Split(Conf.DNSServers, ",")
ctx := context.WithValue(context.Background(), ddns.DNSServerKey{}, utils.IfOr(confServers[0] != "", confServers, utils.DNSServers))
providers, err := DDNSShared.GetDDNSProvidersFromProfiles(server.DDNSProfiles, utils.IfOr(ip != nil, ip, &server.GeoIP.IP))
providers, err := DDNSShared.GetDDNSProvidersFromProfiles(server.DDNSProfiles, utils.IfOr(ip != nil, ip, &server.GeoIP.IP), server.GetUserID())
if err != nil {
return err
}
service/singleton/server_delete_missing_test.go
+32
View File
@@ -0,0 +1,32 @@
package singleton
import (
"testing"
"github.com/nezhahq/nezha/model"
)
func TestServerClassDeleteMissingIDNoPanic(t *testing.T) {
c := &ServerClass{
class: class[uint64, *model.Server]{
list: map[uint64]*model.Server{
1: {Common: model.Common{ID: 1}, UUID: "uuid-1"},
},
},
uuidToID: map[string]uint64{"uuid-1": 1},
}
c.Delete([]uint64{999999})
if _, ok := c.list[1]; !ok {
t.Fatalf("existing server 1 must remain after deleting a non-existent id")
}
c.Delete([]uint64{1, 424242})
if _, ok := c.list[1]; ok {
t.Fatalf("server 1 should be removed")
}
if _, ok := c.uuidToID["uuid-1"]; ok {
t.Fatalf("uuid mapping for server 1 should be removed")
}
}
service/singleton/server_transfer.go
+1447
View File
File diff suppressed because it is too large Load Diff
service/singleton/server_transfer_test.go
+2194
View File
File diff suppressed because it is too large Load Diff
service/singleton/service_type_security_test.go
+37
View File
@@ -0,0 +1,37 @@
package singleton
import (
"testing"
"github.com/stretchr/testify/require"
"github.com/nezhahq/nezha/model"
)
func TestServiceSentinelUpdateRejectsNonProbeTaskTypes(t *testing.T) {
ss := &ServiceSentinel{}
require.Error(t, ss.Update(nil))
for _, taskType := range []uint8{0, model.TaskTypeCommand, model.TaskTypeApplyConfig, model.TaskTypeExec, 255} {
require.Error(t, ss.Update(&model.Service{Type: taskType}), "type %d must not be scheduled", taskType)
}
}
func TestServiceSentinelQuarantinesInvalidPersistedTypes(t *testing.T) {
ss := newServiceMonitorSecurityHarness(t)
insert := `INSERT INTO services
(id, user_id, name, type, target, duration, cover, skip_servers_raw, fail_trigger_tasks_raw, recover_trigger_tasks_raw)
VALUES (?, 100, ?, ?, 'example.invalid:443', 3600, ?, '{}', '[]', '[]')`
require.NoError(t, DB.Exec(insert, 91, "legacy-command", model.TaskTypeCommand, model.ServiceCoverIgnoreAll).Error)
require.NoError(t, DB.Exec(insert, 92, "legacy-apply-config", model.TaskTypeApplyConfig, model.ServiceCoverIgnoreAll).Error)
require.NoError(t, DB.Exec(insert, 93, "valid-probe", model.TaskTypeTCPPing, model.ServiceCoverIgnoreAll).Error)
require.NoError(t, ss.loadServiceHistory())
_, commandLoaded := ss.Get(91)
_, applyConfigLoaded := ss.Get(92)
valid, validLoaded := ss.Get(93)
require.False(t, commandLoaded)
require.False(t, applyConfigLoaded)
require.True(t, validLoaded)
require.Equal(t, uint8(model.TaskTypeTCPPing), valid.Type)
}
service/singleton/servicesentinel.go
+338 -209
View File
@@ -74,8 +74,11 @@ type ServiceSentinel struct {
serviceCurrentStatusData map[uint64]*serviceTaskStatus // 当前任务结果缓存
serviceResponseDataStore map[uint64]serviceResponseData // 当前数据
serviceResponsePing map[uint64]map[uint64]*pingStore // [service_id] -> ClientID -> delay
tlsCertCache map[uint64]string
serviceResponsePing map[uint64]map[uint64]*pingStore // guarded by serviceResponseDataStoreLock; [service_id] -> ClientID -> delay
tlsCertCache map[uint64]string // guarded by serviceResponseDataStoreLock
serviceReportValidatedHook func(uint64)
loadStatsResponseLockedHook func()
serviceReportBeforeTLSSideEffectsHook func(uint64)
servicesLock sync.RWMutex
serviceListLock sync.RWMutex
@@ -85,6 +88,16 @@ type ServiceSentinel struct {
// 30天数据缓存
monthlyStatusLock sync.Mutex
monthlyStatus map[uint64]*serviceResponseItem
// closeOnce + workerWG together let Close() wait for the worker goroutine
// to fully exit. Without this, a test that swaps ServiceSentinelShared back
// to its original value in t.Cleanup races against the still-running
// worker, which keeps reading globals like Conf/CronShared/NotificationShared.
// Production never calls Close() — the process exits while the worker is
// still running and that is fine — but tests must drain the worker before
// restoring globals.
closeOnce sync.Once
workerWG sync.WaitGroup
}
// NewServiceSentinel 创建服务监控器
@@ -113,7 +126,11 @@ func NewServiceSentinel(serviceSentinelDispatchBus chan<- *model.Service) (*Serv
ss.loadTodayStats(today)
// 启动服务监控器
go ss.worker()
ss.workerWG.Add(1)
go func() {
defer ss.workerWG.Done()
ss.worker()
}()
// 每日将游标往后推一天
_, err = CronShared.AddFunc("0 0 0 * * *", ss.refreshMonthlyServiceStatus)
@@ -192,7 +209,15 @@ func (ss *ServiceSentinel) loadServiceHistory() error {
return err
}
validServices := services[:0]
for _, service := range services {
if err := model.ValidateServiceMonitorType(uint64(service.Type)); err != nil {
// Existing databases may contain values written before Service.Type was
// constrained. Quarantine them in the database for operator review, but
// never register a cron job that could dispatch a privileged Agent task.
log.Printf("NEZHA>> quarantining service %d: %v", service.ID, err)
continue
}
task := service
// 通过cron定时将服务监控任务传递给任务调度管道
service.CronJobID, err = CronShared.AddFunc(task.CronSpec(), func() {
@@ -205,7 +230,9 @@ func (ss *ServiceSentinel) loadServiceHistory() error {
ss.serviceCurrentStatusData[service.ID] = new(serviceTaskStatus)
ss.serviceCurrentStatusData[service.ID].result = make([]*pb.TaskResult, 0, _CurrentStatusSize)
ss.serviceStatusToday[service.ID] = &_TodayStatsOfService{}
validServices = append(validServices, service)
}
services = validServices
ss.serviceList = services
sortServices(ss.serviceList)
@@ -322,6 +349,13 @@ func (ss *ServiceSentinel) loadTodayStats(today time.Time) {
}
func (ss *ServiceSentinel) Update(m *model.Service) error {
if m == nil {
return fmt.Errorf("service is nil")
}
if err := model.ValidateServiceMonitorType(uint64(m.Type)); err != nil {
return err
}
ss.serviceResponseDataStoreLock.Lock()
defer ss.serviceResponseDataStoreLock.Unlock()
ss.monthlyStatusLock.Lock()
@@ -372,11 +406,21 @@ func (ss *ServiceSentinel) Delete(ids []uint64) {
for _, id := range ids {
delete(ss.serviceCurrentStatusData, id)
delete(ss.serviceResponseDataStore, id)
delete(ss.serviceResponsePing, id)
delete(ss.tlsCertCache, id)
delete(ss.serviceStatusToday, id)
// 停掉定时任务
CronShared.Remove(ss.services[id].CronJobID)
// GHSA-jx78-55p5-rwv5 (Finding 2): guard against a caller supplying an id
// that does not exist in the in-memory registry. CheckPermission returns
// vacuously true for unknown ids, so the controller layer cannot prevent
// this. Without the guard, ss.services[id] is nil and the .CronJobID
// field access panics, aborting the Delete loop before the remaining valid
// ids are cleaned from memory — their service records were already deleted
// from the database, producing zombie services.
if svc := ss.services[id]; svc != nil {
CronShared.Remove(svc.CronJobID)
}
delete(ss.services, id)
delete(ss.monthlyStatus, id)
@@ -384,12 +428,15 @@ func (ss *ServiceSentinel) Delete(ids []uint64) {
}
func (ss *ServiceSentinel) LoadStats() map[uint64]*serviceResponseItem {
ss.servicesLock.RLock()
defer ss.servicesLock.RUnlock()
ss.serviceResponseDataStoreLock.RLock()
defer ss.serviceResponseDataStoreLock.RUnlock()
if ss.loadStatsResponseLockedHook != nil {
ss.loadStatsResponseLockedHook()
}
ss.monthlyStatusLock.Lock()
defer ss.monthlyStatusLock.Unlock()
ss.servicesLock.RLock()
defer ss.servicesLock.RUnlock()
// 刷新最新一天的数据
for k := range ss.services {
@@ -424,11 +471,6 @@ func (ss *ServiceSentinel) CopyStats() map[uint64]model.ServiceResponseItem {
sri := make(map[uint64]model.ServiceResponseItem)
for k, service := range stats {
if !service.service.EnableShowInService {
delete(stats, k)
continue
}
service.ServiceName = service.service.Name
sri[k] = service.ServiceResponseItem
}
@@ -472,223 +514,297 @@ func (ss *ServiceSentinel) CheckPermission(c *gin.Context, idList iter.Seq[uint6
return true
}
func canReportServiceResult(service *model.Service, reporter *model.Server, taskType uint64) bool {
if service == nil || reporter == nil || uint64(service.Type) != taskType {
return false
}
switch service.Cover {
case model.ServiceCoverAll:
if service.SkipServers[reporter.ID] {
return false
}
case model.ServiceCoverIgnoreAll:
if !service.SkipServers[reporter.ID] {
return false
}
default:
return false
}
return service.UserID == reporter.GetUserID() || userIsAdmin(service.UserID)
}
// Close shuts down the ServiceSentinel worker goroutine and waits for it to
// exit. It is idempotent and safe to call more than once.
//
// Why this exists: the worker reads multiple package-level globals during
// each report (Conf, CronShared via notifyCheck, NotificationShared via
// UnMuteNotification, ServerShared, TSDBShared). A test fixture that swaps
// those globals out in t.Cleanup MUST first call Close() — otherwise the
// cleanup write races the still-running worker's read and `go test -race`
// fires (see security_regression_test.go newServiceMonitorSecurityHarness).
// Production never calls Close because the process exits with the worker
// still running, which is fine.
func (ss *ServiceSentinel) Close() {
ss.closeOnce.Do(func() {
close(ss.serviceReportChannel)
ss.workerWG.Wait()
})
}
// worker 服务监控的实际工作流程
//
// IMPORTANT: this loop reads several package-level globals (Conf, CronShared,
// NotificationShared, ServerShared, TSDBShared). Any test that replaces those
// globals via t.Cleanup must first call ServiceSentinel.Close() so the worker
// drains and exits before the swap, otherwise the race detector trips. See
// the Close() comment above for the full rationale.
func (ss *ServiceSentinel) worker() {
// 从服务状态汇报管道获取汇报的服务数据
for r := range ss.serviceReportChannel {
css, _ := ss.Get(r.Data.GetId())
if css == nil || css.ID == 0 {
log.Printf("NEZHA>> Incorrect service monitor report %+v", r)
continue
}
css = nil
mh := r.Data
if mh.Type == model.TaskTypeTCPPing || mh.Type == model.TaskTypeICMPPing {
// TCP/ICMP Ping 使用平均值计算后再写入
serviceTcpMap, ok := ss.serviceResponsePing[mh.GetId()]
if !ok {
serviceTcpMap = make(map[uint64]*pingStore)
ss.serviceResponsePing[mh.GetId()] = serviceTcpMap
}
ts, ok := serviceTcpMap[r.Reporter]
if !ok {
ts = &pingStore{}
}
ts.count++
ts.ping = (ts.ping*float64(ts.count-1) + float64(mh.Delay)) / float64(ts.count)
if mh.Successful {
ts.successCount++
}
if ts.count == Conf.AvgPingCount {
if TSDBEnabled() {
if err := TSDBShared.WriteServiceMetrics(&tsdb.ServiceMetrics{
ServiceID: mh.GetId(),
ServerID: r.Reporter,
Timestamp: time.Now(),
Delay: ts.ping,
Successful: ts.successCount*2 >= ts.count,
}); err != nil {
log.Printf("NEZHA>> Failed to save service monitor metrics to TSDB: %v", err)
}
} else {
if err := DB.Create(&model.ServiceHistory{
ServiceID: mh.GetId(),
AvgDelay: ts.ping,
Data: mh.Data,
ServerID: r.Reporter,
}).Error; err != nil {
log.Printf("NEZHA>> Failed to save service monitor metrics: %v", err)
}
serverShared := ServerShared
func() {
defer func() {
if recovered := recover(); recovered != nil {
log.Printf("NEZHA>> Service monitor report processing panicked: %v", recovered)
}
ts.count = 0
ts.ping = 0
ts.successCount = 0
}
serviceTcpMap[r.Reporter] = ts
} else {
}()
ss.processReport(r, serverShared)
}()
}
}
func (ss *ServiceSentinel) processReport(r ReportData, serverShared *ServerClass) {
serverShared.lockLifecycleRead()
defer serverShared.unlockLifecycleRead()
cs, _ := ss.Get(r.Data.GetId())
reporter, _ := serverShared.Get(r.Reporter)
// 入站结果必须匹配出站任务派发边界，避免 agent 伪造其他服务 ID 写入监控状态。
if !canReportServiceResult(cs, reporter, r.Data.GetType()) {
log.Printf("NEZHA>> Incorrect service monitor report %+v", r)
return
}
if ss.serviceReportValidatedHook != nil {
ss.serviceReportValidatedHook(r.Data.GetId())
}
mh := r.Data
m := serverShared.GetList()
// Serialize Delete and Update before this accepted report causes any side effect.
ss.serviceResponseDataStoreLock.Lock()
defer ss.serviceResponseDataStoreLock.Unlock()
serviceStatusToday := ss.serviceStatusToday[mh.GetId()]
serviceCurrentStatusData := ss.serviceCurrentStatusData[mh.GetId()]
currentService, serviceExists := ss.Get(mh.GetId())
if serviceStatusToday == nil || serviceCurrentStatusData == nil || !serviceExists ||
!canReportServiceResult(currentService, reporter, mh.GetType()) {
return
}
cs = currentService
if mh.Type == model.TaskTypeTCPPing || mh.Type == model.TaskTypeICMPPing {
// TCP/ICMP Ping 使用平均值计算后再写入
serviceTcpMap, ok := ss.serviceResponsePing[mh.GetId()]
if !ok {
serviceTcpMap = make(map[uint64]*pingStore)
ss.serviceResponsePing[mh.GetId()] = serviceTcpMap
}
ts, ok := serviceTcpMap[r.Reporter]
if !ok {
ts = &pingStore{}
}
ts.count++
ts.ping = (ts.ping*float64(ts.count-1) + float64(mh.Delay)) / float64(ts.count)
if mh.Successful {
ts.successCount++
}
if ts.count == Conf.AvgPingCount {
if TSDBEnabled() {
if err := TSDBShared.WriteServiceMetrics(&tsdb.ServiceMetrics{
ServiceID: mh.GetId(),
ServerID: r.Reporter,
Timestamp: time.Now(),
Delay: float64(mh.Delay),
Successful: mh.Successful,
Delay: ts.ping,
Successful: ts.successCount*2 >= ts.count,
}); err != nil {
log.Printf("NEZHA>> Failed to save service monitor metrics to TSDB: %v", err)
}
}
}
ss.serviceResponseDataStoreLock.Lock()
// 写入当天状态
if mh.Successful {
ss.serviceStatusToday[mh.GetId()].Delay = (ss.serviceStatusToday[mh.
GetId()].Delay*float64(ss.serviceStatusToday[mh.GetId()].Up) +
float64(mh.Delay)) / float64(ss.serviceStatusToday[mh.GetId()].Up+1)
ss.serviceStatusToday[mh.GetId()].Up++
} else {
ss.serviceStatusToday[mh.GetId()].Down++
}
currentTime := time.Now()
if ss.serviceCurrentStatusData[mh.GetId()].t.IsZero() {
ss.serviceCurrentStatusData[mh.GetId()].t = currentTime
}
// 写入当前数据
if ss.serviceCurrentStatusData[mh.GetId()].t.Before(currentTime) {
ss.serviceCurrentStatusData[mh.GetId()].t = currentTime.Add(30 * time.Second)
ss.serviceCurrentStatusData[mh.GetId()].result = append(ss.serviceCurrentStatusData[mh.GetId()].result, mh)
}
// 更新当前状态
ss.serviceResponseDataStore[mh.GetId()] = serviceResponseData{}
// 永远是最新的 30 个数据的状态 [01:00, 02:00, 03:00] -> [04:00, 02:00, 03: 00]
for _, cs := range ss.serviceCurrentStatusData[mh.GetId()].result {
if cs.GetId() > 0 {
rd := ss.serviceResponseDataStore[mh.GetId()]
if cs.Successful {
rd.Up++
rd.Delay = (rd.Delay*float64(rd.Up-1) + float64(cs.Delay)) / float64(rd.Up)
} else {
rd.Down++
}
ss.serviceResponseDataStore[mh.GetId()] = rd
}
}
// 计算在线率，
var stateCode uint8
{
upPercent := uint64(0)
rd := ss.serviceResponseDataStore[mh.GetId()]
if rd.Down+rd.Up > 0 {
upPercent = rd.Up * 100 / (rd.Down + rd.Up)
}
stateCode = GetStatusCode(upPercent)
}
if len(ss.serviceCurrentStatusData[mh.GetId()].result) == _CurrentStatusSize {
ss.serviceCurrentStatusData[mh.GetId()].t = currentTime
if !TSDBEnabled() {
rd := ss.serviceResponseDataStore[mh.GetId()]
} else {
if err := DB.Create(&model.ServiceHistory{
ServiceID: mh.GetId(),
AvgDelay: rd.Delay,
AvgDelay: ts.ping,
Data: mh.Data,
Up: rd.Up,
Down: rd.Down,
ServerID: r.Reporter,
}).Error; err != nil {
log.Printf("NEZHA>> Failed to save service monitor metrics: %v", err)
}
}
ss.serviceCurrentStatusData[mh.GetId()].result = ss.serviceCurrentStatusData[mh.GetId()].result[:0]
ts.count = 0
ts.ping = 0
ts.successCount = 0
}
cs, _ := ss.Get(mh.GetId())
m := ServerShared.GetList()
// 延迟报警
if mh.Delay > 0 {
delayCheck(&r, m, cs, mh)
}
// 状态变更报警+触发任务执行
if stateCode == StatusDown || stateCode != ss.serviceCurrentStatusData[mh.GetId()].lastStatus {
lastStatus := ss.serviceCurrentStatusData[mh.GetId()].lastStatus
// 存储新的状态值
ss.serviceCurrentStatusData[mh.GetId()].lastStatus = stateCode
notifyCheck(&r, m, cs, mh, lastStatus, stateCode)
}
ss.serviceResponseDataStoreLock.Unlock()
// TLS 证书报警
var errMsg string
if strings.HasPrefix(mh.Data, "SSL证书错误：") {
// i/o timeout、connection timeout、EOF 错误
if !strings.HasSuffix(mh.Data, "timeout") &&
!strings.HasSuffix(mh.Data, "EOF") &&
!strings.HasSuffix(mh.Data, "timed out") {
errMsg = mh.Data
if cs.Notify {
muteLabel := NotificationMuteLabel.ServiceTLS(mh.GetId(), "network")
go NotificationShared.SendNotification(cs.NotificationGroupID, Localizer.Tf("[TLS] Fetch cert info failed, Reporter: %s, Error: %s", cs.Name, errMsg), muteLabel)
}
serviceTcpMap[r.Reporter] = ts
} else {
if TSDBEnabled() {
if err := TSDBShared.WriteServiceMetrics(&tsdb.ServiceMetrics{
ServiceID: mh.GetId(),
ServerID: r.Reporter,
Timestamp: time.Now(),
Delay: float64(mh.Delay),
Successful: mh.Successful,
}); err != nil {
log.Printf("NEZHA>> Failed to save service monitor metrics to TSDB: %v", err)
}
} else {
// 清除网络错误静音缓存
NotificationShared.UnMuteNotification(cs.NotificationGroupID, NotificationMuteLabel.ServiceTLS(mh.GetId(), "network"))
}
}
var newCert = strings.Split(mh.Data, "|")
if len(newCert) > 1 {
enableNotify := cs.Notify
// 写入当天状态
if mh.Successful {
serviceStatusToday.Delay = (serviceStatusToday.Delay*float64(serviceStatusToday.Up) +
float64(mh.Delay)) / float64(serviceStatusToday.Up+1)
serviceStatusToday.Up++
} else {
serviceStatusToday.Down++
}
// 首次获取证书信息时，缓存证书信息
if ss.tlsCertCache[mh.GetId()] == "" {
ss.tlsCertCache[mh.GetId()] = mh.Data
currentTime := time.Now()
if serviceCurrentStatusData.t.IsZero() {
serviceCurrentStatusData.t = currentTime
}
// 写入当前数据
if serviceCurrentStatusData.t.Before(currentTime) {
serviceCurrentStatusData.t = currentTime.Add(30 * time.Second)
serviceCurrentStatusData.result = append(serviceCurrentStatusData.result, mh)
}
// 更新当前状态
ss.serviceResponseDataStore[mh.GetId()] = serviceResponseData{}
// 永远是最新的 30 个数据的状态 [01:00, 02:00, 03:00] -> [04:00, 02:00, 03: 00]
for _, cs := range serviceCurrentStatusData.result {
if cs.GetId() > 0 {
rd := ss.serviceResponseDataStore[mh.GetId()]
if cs.Successful {
rd.Up++
rd.Delay = (rd.Delay*float64(rd.Up-1) + float64(cs.Delay)) / float64(rd.Up)
} else {
rd.Down++
}
ss.serviceResponseDataStore[mh.GetId()] = rd
}
}
// 计算在线率，
var stateCode uint8
{
upPercent := uint64(0)
rd := ss.serviceResponseDataStore[mh.GetId()]
if rd.Down+rd.Up > 0 {
upPercent = rd.Up * 100 / (rd.Down + rd.Up)
}
stateCode = GetStatusCode(upPercent)
}
if len(serviceCurrentStatusData.result) == _CurrentStatusSize {
serviceCurrentStatusData.t = currentTime
if !TSDBEnabled() {
rd := ss.serviceResponseDataStore[mh.GetId()]
if err := DB.Create(&model.ServiceHistory{
ServiceID: mh.GetId(),
AvgDelay: rd.Delay,
Data: mh.Data,
Up: rd.Up,
Down: rd.Down,
}).Error; err != nil {
log.Printf("NEZHA>> Failed to save service monitor metrics: %v", err)
}
}
serviceCurrentStatusData.result = serviceCurrentStatusData.result[:0]
}
// 延迟报警
if mh.Delay > 0 {
delayCheck(&r, m, cs, mh)
}
// 状态变更报警+触发任务执行
if stateCode == StatusDown || stateCode != serviceCurrentStatusData.lastStatus {
lastStatus := serviceCurrentStatusData.lastStatus
// 存储新的状态值
serviceCurrentStatusData.lastStatus = stateCode
notifyCheck(&r, m, cs, mh, lastStatus, stateCode)
}
// TLS 证书报警
if ss.serviceReportBeforeTLSSideEffectsHook != nil {
ss.serviceReportBeforeTLSSideEffectsHook(mh.GetId())
}
var errMsg string
if strings.HasPrefix(mh.Data, "SSL证书错误：") {
// i/o timeout、connection timeout、EOF 错误
if !strings.HasSuffix(mh.Data, "timeout") &&
!strings.HasSuffix(mh.Data, "EOF") &&
!strings.HasSuffix(mh.Data, "timed out") {
errMsg = mh.Data
if cs.Notify {
muteLabel := NotificationMuteLabel.ServiceTLS(mh.GetId(), "network")
go NotificationShared.SendNotification(cs.NotificationGroupID, Localizer.Tf("[TLS] Fetch cert info failed, Reporter: %s, Error: %s", cs.Name, errMsg), muteLabel)
}
}
} else {
// 清除网络错误静音缓存
NotificationShared.UnMuteNotification(cs.NotificationGroupID, NotificationMuteLabel.ServiceTLS(mh.GetId(), "network"))
var newCert = strings.Split(mh.Data, "|")
if len(newCert) > 1 {
enableNotify := cs.Notify
// 首次获取证书信息时，缓存证书信息
if ss.tlsCertCache[mh.GetId()] == "" {
ss.tlsCertCache[mh.GetId()] = mh.Data
}
oldCert := strings.Split(ss.tlsCertCache[mh.GetId()], "|")
isCertChanged := false
expiresOld, _ := time.Parse("2006-01-02 15:04:05 -0700 MST", oldCert[1])
expiresNew, _ := time.Parse("2006-01-02 15:04:05 -0700 MST", newCert[1])
// 证书变更时，更新缓存
if oldCert[0] != newCert[0] && !expiresNew.Equal(expiresOld) {
isCertChanged = true
ss.tlsCertCache[mh.GetId()] = mh.Data
}
notificationGroupID := cs.NotificationGroupID
serviceName := cs.Name
// 需要发送提醒
if enableNotify {
// 证书过期提醒
if expiresNew.Before(time.Now().AddDate(0, 0, 7)) {
expiresTimeStr := expiresNew.Format("2006-01-02 15:04:05")
errMsg = Localizer.Tf(
"The TLS certificate will expire within seven days. Expiration time: %s",
expiresTimeStr,
)
// 静音规则： 服务id+证书过期时间
// 用于避免多个监测点对相同证书同时报警
muteLabel := NotificationMuteLabel.ServiceTLS(mh.GetId(), fmt.Sprintf("expire_%s", expiresTimeStr))
go NotificationShared.SendNotification(notificationGroupID, fmt.Sprintf("[TLS] %s %s", serviceName, errMsg), muteLabel)
}
oldCert := strings.Split(ss.tlsCertCache[mh.GetId()], "|")
isCertChanged := false
expiresOld, _ := time.Parse("2006-01-02 15:04:05 -0700 MST", oldCert[1])
expiresNew, _ := time.Parse("2006-01-02 15:04:05 -0700 MST", newCert[1])
// 证书变更提醒
if isCertChanged {
errMsg = Localizer.Tf(
"TLS certificate changed, old: issuer %s, expires at %s; new: issuer %s, expires at %s",
oldCert[0], expiresOld.Format("2006-01-02 15:04:05"), newCert[0], expiresNew.Format("2006-01-02 15:04:05"))
// 证书变更时，更新缓存
if oldCert[0] != newCert[0] && !expiresNew.Equal(expiresOld) {
isCertChanged = true
ss.tlsCertCache[mh.GetId()] = mh.Data
}
notificationGroupID := cs.NotificationGroupID
serviceName := cs.Name
// 需要发送提醒
if enableNotify {
// 证书过期提醒
if expiresNew.Before(time.Now().AddDate(0, 0, 7)) {
expiresTimeStr := expiresNew.Format("2006-01-02 15:04:05")
errMsg = Localizer.Tf(
"The TLS certificate will expire within seven days. Expiration time: %s",
expiresTimeStr,
)
// 静音规则： 服务id+证书过期时间
// 用于避免多个监测点对相同证书同时报警
muteLabel := NotificationMuteLabel.ServiceTLS(mh.GetId(), fmt.Sprintf("expire_%s", expiresTimeStr))
go NotificationShared.SendNotification(notificationGroupID, fmt.Sprintf("[TLS] %s %s", serviceName, errMsg), muteLabel)
}
// 证书变更提醒
if isCertChanged {
errMsg = Localizer.Tf(
"TLS certificate changed, old: issuer %s, expires at %s; new: issuer %s, expires at %s",
oldCert[0], expiresOld.Format("2006-01-02 15:04:05"), newCert[0], expiresNew.Format("2006-01-02 15:04:05"))
// 证书变更后会自动更新缓存，所以不需要静音
go NotificationShared.SendNotification(notificationGroupID, fmt.Sprintf("[TLS] %s %s", serviceName, errMsg), "")
}
// 证书变更后会自动更新缓存，所以不需要静音
go NotificationShared.SendNotification(notificationGroupID, fmt.Sprintf("[TLS] %s %s", serviceName, errMsg), "")
}
}
}
@@ -700,17 +816,25 @@ func delayCheck(r *ReportData, m map[uint64]*model.Server, ss *model.Service, mh
return
}
// GHSA-jx78-55p5-rwv5 (incomplete fix of GHSA-qjpp-gffx-2wm9): the server
// map snapshot m is taken outside serviceResponseDataStoreLock and
// ServerShared has its own independent lock, so a concurrent batch-delete of
// the reporter's server can remove the entry between the pre-lock validation
// and this point. Guard against the nil pointer before using the server.
reporterServer := m[r.Reporter]
if reporterServer == nil {
return
}
notificationGroupID := ss.NotificationGroupID
minMuteLabel := NotificationMuteLabel.ServiceLatencyMin(mh.GetId())
maxMuteLabel := NotificationMuteLabel.ServiceLatencyMax(mh.GetId())
if mh.Delay > ss.MaxLatency {
// 延迟超过最大值
reporterServer := m[r.Reporter]
msg := Localizer.Tf("[Latency] %s %2f > %2f, Reporter: %s", ss.Name, mh.Delay, ss.MaxLatency, reporterServer.Name)
go NotificationShared.SendNotification(notificationGroupID, msg, minMuteLabel)
} else if mh.Delay < ss.MinLatency {
// 延迟低于最小值
reporterServer := m[r.Reporter]
msg := Localizer.Tf("[Latency] %s %2f < %2f, Reporter: %s", ss.Name, mh.Delay, ss.MinLatency, reporterServer.Name)
go NotificationShared.SendNotification(notificationGroupID, msg, maxMuteLabel)
} else {
@@ -722,10 +846,16 @@ func delayCheck(r *ReportData, m map[uint64]*model.Server, ss *model.Service, mh
func notifyCheck(r *ReportData, m map[uint64]*model.Server,
ss *model.Service, mh *pb.TaskResult, lastStatus, stateCode uint8) {
// GHSA-jx78-55p5-rwv5: guard against concurrent server deletion (same TOCTOU
// class as the 2026-07-21 fix, a few dozen lines lower in the same worker).
// ServerShared has its own lock; m is a snapshot taken outside
// serviceResponseDataStoreLock, so the server may have been removed between
// the pre-lock validation and here.
reporterServer := m[r.Reporter]
// 判断是否需要发送通知
isNeedSendNotification := ss.Notify && (lastStatus != 0 || stateCode == StatusDown)
if isNeedSendNotification {
reporterServer := m[r.Reporter]
if isNeedSendNotification && reporterServer != nil {
notificationGroupID := ss.NotificationGroupID
notificationMsg := Localizer.Tf("[%s] %s Reporter: %s, Error: %s", StatusCodeToString(stateCode), ss.Name, reporterServer.Name, mh.Data)
muteLabel := NotificationMuteLabel.ServiceStateChanged(mh.GetId())
@@ -740,14 +870,13 @@ func notifyCheck(r *ReportData, m map[uint64]*model.Server,
// 判断是否需要触发任务
isNeedTriggerTask := ss.EnableTriggerTask && lastStatus != 0
if isNeedTriggerTask {
reporterServer := m[r.Reporter]
if isNeedTriggerTask && reporterServer != nil {
if stateCode == StatusGood && lastStatus != stateCode {
// 当前状态正常 前序状态非正常时 触发恢复任务
go CronShared.SendTriggerTasks(ss.RecoverTriggerTasks, reporterServer.ID)
go CronShared.SendTriggerTasks(ss.RecoverTriggerTasks, reporterServer.ID, ss.UserID)
} else if lastStatus == StatusGood && lastStatus != stateCode {
// 前序状态正常 当前状态非正常时 触发失败任务
go CronShared.SendTriggerTasks(ss.FailTriggerTasks, reporterServer.ID)
go CronShared.SendTriggerTasks(ss.FailTriggerTasks, reporterServer.ID, ss.UserID)
}
}
}
service/singleton/servicesentinel_lifecycle_test.go
+646
View File
@@ -0,0 +1,646 @@
package singleton
import (
"context"
"fmt"
"os"
"os/exec"
"strings"
"sync"
"testing"
"time"
"github.com/nezhahq/nezha/model"
)
// Regression markers for Finding 1 and Finding 2 of GHSA-jx78-55p5-rwv5
// (incomplete fix of GHSA-qjpp-gffx-2wm9).
const (
concurrentServerDeleteSuccessMarker = "ghsa-jx78-55p5-rwv5-finding1-no-crash"
deleteUnknownIDSuccessMarker = "ghsa-jx78-55p5-rwv5-finding2-no-zombie"
)
const serviceSentinelLifecycleSuccessMarker = "service-sentinel-stale-report-lifecycle-success"
func TestServiceSentinelReporterDeleteWaitsForSynchronousReportProcessing(t *testing.T) {
ctx, cancel := context.WithTimeout(t.Context(), time.Second)
defer cancel()
ss := newServiceMonitorSecurityHarness(t,
&model.Server{Common: model.Common{ID: 1, UserID: 1}, Name: "reporter"},
)
service := &model.Service{
Common: model.Common{ID: 10, UserID: 1},
Name: "lifecycle-service",
Type: model.TaskTypeTCPPing,
Target: "lifecycle.example.invalid:443",
Duration: 3600,
Cover: model.ServiceCoverIgnoreAll,
SkipServers: map[uint64]bool{1: true},
}
addServiceMonitorSecurityService(t, ss, service)
reportValidated := make(chan struct{})
releaseReport := make(chan struct{})
var releaseOnce sync.Once
release := func() { releaseOnce.Do(func() { close(releaseReport) }) }
ss.serviceReportValidatedHook = func(serviceID uint64) {
if serviceID == service.ID {
close(reportValidated)
<-releaseReport
}
}
t.Cleanup(func() {
release()
ss.Close()
})
ss.Dispatch(serviceMonitorResult(1, service.ID, model.TaskTypeTCPPing, true))
select {
case <-reportValidated:
case <-ctx.Done():
t.Fatal(ctx.Err())
}
deleteDone := make(chan struct{})
go func() {
ServerShared.Delete([]uint64{1})
close(deleteDone)
}()
select {
case <-deleteDone:
t.Fatal("server deletion returned before the accepted report completed")
case <-time.After(25 * time.Millisecond):
}
release()
select {
case <-deleteDone:
case <-ctx.Done():
t.Fatal(ctx.Err())
}
var historyCount int64
if err := DB.Model(&model.ServiceHistory{}).
Where("service_id = ? AND server_id = ?", service.ID, 1).
Count(&historyCount).Error; err != nil {
t.Fatal(err)
}
if historyCount != 1 {
t.Fatalf("expected report side effects before deletion returned, got %d history rows", historyCount)
}
if _, ok := ServerShared.Get(1); ok {
t.Fatal("expected reporter to be deleted after the report completed")
}
}
func TestServiceSentinelWorkerRejectsReportAfterReporterDeletion(t *testing.T) {
ss := newServiceMonitorSecurityHarness(t,
&model.Server{Common: model.Common{ID: 1, UserID: 1}, Name: "reporter"},
)
service := &model.Service{
Common: model.Common{ID: 10, UserID: 1},
Name: "deleted-reporter-service",
Type: model.TaskTypeTCPPing,
Target: "deleted-reporter.example.invalid:443",
Duration: 3600,
Cover: model.ServiceCoverIgnoreAll,
SkipServers: map[uint64]bool{1: true},
}
addServiceMonitorSecurityService(t, ss, service)
ServerShared.Delete([]uint64{1})
ss.Dispatch(serviceMonitorResult(1, service.ID, model.TaskTypeTCPPing, true))
ss.Close()
var historyCount int64
if err := DB.Model(&model.ServiceHistory{}).
Where("service_id = ? AND server_id = ?", service.ID, 1).
Count(&historyCount).Error; err != nil {
t.Fatal(err)
}
if historyCount != 0 {
t.Fatalf("expected no history after reporter deletion, got %d rows", historyCount)
}
ss.serviceResponseDataStoreLock.RLock()
_, pingCached := ss.serviceResponsePing[service.ID]
_, responseCached := ss.serviceResponseDataStore[service.ID]
stats := ss.serviceStatusToday[service.ID]
ss.serviceResponseDataStoreLock.RUnlock()
if pingCached {
t.Fatal("expected no ping cache side effect after reporter deletion")
}
if responseCached {
t.Fatal("expected no response cache side effect after reporter deletion")
}
if stats == nil || stats.Up != 0 || stats.Down != 0 {
t.Fatalf("expected no stats side effect after reporter deletion, got %+v", stats)
}
}
func TestServiceSentinelWorkerRecoversPerReportPanic(t *testing.T) {
ctx, cancel := context.WithTimeout(t.Context(), time.Second)
defer cancel()
ss := newServiceMonitorSecurityHarness(t,
&model.Server{Common: model.Common{ID: 1, UserID: 1}, Name: "reporter"},
)
panicService := &model.Service{
Common: model.Common{ID: 10, UserID: 1},
Name: "panic-service",
Type: model.TaskTypeHTTPGet,
Target: "https://panic.example.invalid",
Duration: 3600,
Cover: model.ServiceCoverIgnoreAll,
SkipServers: map[uint64]bool{1: true},
}
validService := &model.Service{
Common: model.Common{ID: 20, UserID: 1},
Name: "valid-service",
Type: model.TaskTypeTCPPing,
Target: "valid.example.invalid:443",
Duration: 3600,
Cover: model.ServiceCoverIgnoreAll,
SkipServers: map[uint64]bool{1: true},
}
addServiceMonitorSecurityService(t, ss, panicService)
addServiceMonitorSecurityService(t, ss, validService)
ss.serviceReportBeforeTLSSideEffectsHook = func(serviceID uint64) {
if serviceID == panicService.ID {
panic("test service report panic")
}
}
ss.Dispatch(serviceMonitorResult(1, panicService.ID, model.TaskTypeHTTPGet, true))
ss.Dispatch(serviceMonitorResult(1, validService.ID, model.TaskTypeTCPPing, true))
waitForServiceHistory(t, validService.ID, 1)
ss.Close()
if !ss.serviceResponseDataStoreLock.TryLock() {
t.Fatal("panic leaked the service response lock")
}
ss.serviceResponseDataStoreLock.Unlock()
deleteDone := make(chan struct{})
go func() {
ServerShared.Delete([]uint64{1})
close(deleteDone)
}()
select {
case <-deleteDone:
case <-ctx.Done():
t.Fatal("panic leaked a lifecycle lock: " + ctx.Err().Error())
}
}
func TestServiceSentinelWorkerIgnoresStaleReportAfterDeletion(t *testing.T) {
if os.Getenv("NEZHA_SERVICE_SENTINEL_LIFECYCLE_CHILD") == "1" {
testServiceSentinelWorkerIgnoresStaleReportAfterDeletionChild(t)
return
}
// Given
ctx, cancel := context.WithTimeout(t.Context(), 5*time.Second)
defer cancel()
child := exec.CommandContext(ctx, os.Args[0], "-test.run=^TestServiceSentinelWorkerIgnoresStaleReportAfterDeletion$")
child.Env = append(os.Environ(), "NEZHA_SERVICE_SENTINEL_LIFECYCLE_CHILD=1")
// When
output, err := child.CombinedOutput()
// Then
if ctx.Err() != nil {
t.Fatalf("service sentinel lifecycle child timed out: %v\n%s", ctx.Err(), output)
}
if err != nil {
t.Fatalf("service sentinel lifecycle child failed: %v\n%s", err, output)
}
if !strings.Contains(string(output), serviceSentinelLifecycleSuccessMarker) {
t.Fatalf("service sentinel lifecycle child did not report success:\n%s", output)
}
}
func testServiceSentinelWorkerIgnoresStaleReportAfterDeletionChild(t *testing.T) {
// Given
ctx, cancel := context.WithTimeout(t.Context(), time.Second)
defer cancel()
ss := newServiceMonitorSecurityHarness(t,
&model.Server{Common: model.Common{ID: 1, UserID: 1}, Name: "reporter"},
)
for _, service := range []*model.Service{
{
Common: model.Common{ID: 10, UserID: 1},
Name: "stale-service",
Type: model.TaskTypeTCPPing,
Target: "stale.example.invalid:443",
Duration: 3600,
Cover: model.ServiceCoverIgnoreAll,
SkipServers: map[uint64]bool{1: true},
},
{
Common: model.Common{ID: 20, UserID: 1},
Name: "valid-service",
Type: model.TaskTypeTCPPing,
Target: "valid.example.invalid:443",
Duration: 3600,
Cover: model.ServiceCoverIgnoreAll,
SkipServers: map[uint64]bool{1: true},
},
} {
addServiceMonitorSecurityService(t, ss, service)
}
acceptedStaleReport := make(chan struct{})
releaseWorker := make(chan struct{})
var releaseOnce sync.Once
releaseWorkerHook := func() {
releaseOnce.Do(func() { close(releaseWorker) })
}
ss.serviceReportValidatedHook = func(serviceID uint64) {
if serviceID == 10 {
close(acceptedStaleReport)
<-releaseWorker
}
}
t.Cleanup(func() {
releaseWorkerHook()
ss.Close()
})
// When
ss.Dispatch(serviceMonitorResult(1, 10, model.TaskTypeTCPPing, true))
select {
case <-acceptedStaleReport:
case <-ctx.Done():
t.Fatal(ctx.Err())
}
ss.Delete([]uint64{10})
releaseWorkerHook()
ss.Dispatch(serviceMonitorResult(1, 20, model.TaskTypeTCPPing, true))
ss.Close()
// Then
var staleHistoryCount int64
if err := DB.Model(&model.ServiceHistory{}).
Where("service_id = ? AND server_id = ?", 10, 1).
Count(&staleHistoryCount).Error; err != nil {
t.Fatal(err)
}
if staleHistoryCount != 0 {
t.Fatalf("expected stale service to write zero per-reporter history rows, got %d", staleHistoryCount)
}
var validHistoryCount int64
if err := DB.Model(&model.ServiceHistory{}).
Where("service_id = ? AND server_id = ?", 20, 1).
Count(&validHistoryCount).Error; err != nil {
t.Fatal(err)
}
if validHistoryCount != 1 {
t.Fatalf("expected exactly one valid service history row, got %d", validHistoryCount)
}
ss.serviceResponseDataStoreLock.RLock()
_, stalePingCached := ss.serviceResponsePing[10]
validStats := ss.serviceStatusToday[20]
ss.serviceResponseDataStoreLock.RUnlock()
if stalePingCached {
t.Fatal("expected stale service ping cache to be deleted")
}
if validStats == nil || validStats.Up != 1 || validStats.Down != 0 {
t.Fatalf("expected valid service stats up=1 down=0, got %+v", validStats)
}
if _, err := fmt.Fprintln(os.Stdout, serviceSentinelLifecycleSuccessMarker); err != nil {
t.Fatal(err)
}
}
func TestServiceSentinelWorkerRevalidatesReportAfterUpdate(t *testing.T) {
// Given
ctx, cancel := context.WithTimeout(t.Context(), time.Second)
defer cancel()
ss := newServiceMonitorSecurityHarness(t,
&model.Server{Common: model.Common{ID: 1, UserID: 1}, Name: "reporter"},
)
service := &model.Service{
Common: model.Common{ID: 10, UserID: 1},
Name: "updatable-service",
Type: model.TaskTypeTCPPing,
Target: "updatable.example.invalid:443",
Duration: 3600,
Cover: model.ServiceCoverIgnoreAll,
SkipServers: map[uint64]bool{1: true},
}
addServiceMonitorSecurityService(t, ss, service)
ss.serviceResponseDataStoreLock.Lock()
ss.serviceStatusToday[service.ID] = &_TodayStatsOfService{Up: 7, Down: 3, Delay: 12.5}
ss.serviceResponseDataStoreLock.Unlock()
acceptedReport := make(chan struct{})
releaseWorker := make(chan struct{})
var releaseOnce sync.Once
releaseWorkerHook := func() {
releaseOnce.Do(func() { close(releaseWorker) })
}
ss.serviceReportValidatedHook = func(serviceID uint64) {
if serviceID == service.ID {
close(acceptedReport)
<-releaseWorker
}
}
t.Cleanup(func() {
releaseWorkerHook()
ss.Close()
})
// When
ss.Dispatch(serviceMonitorResult(1, service.ID, model.TaskTypeTCPPing, true))
select {
case <-acceptedReport:
case <-ctx.Done():
t.Fatal(ctx.Err())
}
updatedService := *service
updatedService.Name = "updated-service"
updatedService.SkipServers = map[uint64]bool{}
if err := ss.Update(&updatedService); err != nil {
t.Fatal(err)
}
releaseWorkerHook()
ss.Close()
// Then
var historyCount int64
if err := DB.Model(&model.ServiceHistory{}).
Where("service_id = ? AND server_id = ?", service.ID, 1).
Count(&historyCount).Error; err != nil {
t.Fatal(err)
}
if historyCount != 0 {
t.Fatalf("expected updated service to write zero per-reporter history rows, got %d", historyCount)
}
ss.serviceResponseDataStoreLock.RLock()
_, pingCached := ss.serviceResponsePing[service.ID]
stats := ss.serviceStatusToday[service.ID]
ss.serviceResponseDataStoreLock.RUnlock()
if pingCached {
t.Fatal("expected updated service report to leave no ping cache entry")
}
if stats == nil || stats.Up != 7 || stats.Down != 3 || stats.Delay != 12.5 {
t.Fatalf("expected existing service stats to remain unchanged, got %+v", stats)
}
currentService, ok := ss.Get(service.ID)
if !ok || currentService.Name != updatedService.Name || currentService.SkipServers[1] {
t.Fatalf("expected updated service configuration, got %+v", currentService)
}
}
func TestServiceSentinelLoadStatsFollowsLifecycleLockOrder(t *testing.T) {
// Given
ss := &ServiceSentinel{
serviceStatusToday: make(map[uint64]*_TodayStatsOfService),
serviceResponseDataStore: make(map[uint64]serviceResponseData),
services: make(map[uint64]*model.Service),
monthlyStatus: make(map[uint64]*serviceResponseItem),
}
ss.loadStatsResponseLockedHook = func() {
if ss.serviceResponseDataStoreLock.TryLock() {
ss.serviceResponseDataStoreLock.Unlock()
t.Fatal("LoadStats invoked the hook before acquiring the response read lock")
}
if !ss.monthlyStatusLock.TryLock() {
t.Fatal("LoadStats acquired monthlyStatusLock before the response lock hook")
}
ss.monthlyStatusLock.Unlock()
if !ss.servicesLock.TryLock() {
t.Fatal("LoadStats acquired servicesLock before the response lock hook")
}
ss.servicesLock.Unlock()
}
// When / Then
ss.LoadStats()
}
func TestServiceSentinelWorkerHoldsResponseLockDuringTLSSideEffects(t *testing.T) {
// Given
ctx, cancel := context.WithTimeout(t.Context(), time.Second)
defer cancel()
ss := newServiceMonitorSecurityHarness(t,
&model.Server{Common: model.Common{ID: 1, UserID: 1}, Name: "reporter"},
)
service := &model.Service{
Common: model.Common{ID: 10, UserID: 1},
Name: "tls-service",
Type: model.TaskTypeHTTPGet,
Target: "https://tls.example.invalid",
Duration: 3600,
Cover: model.ServiceCoverIgnoreAll,
SkipServers: map[uint64]bool{1: true},
}
addServiceMonitorSecurityService(t, ss, service)
tlsSideEffectsReady := make(chan struct{})
releaseTLSSideEffects := make(chan struct{})
var releaseOnce sync.Once
releaseTLSSideEffectsHook := func() {
releaseOnce.Do(func() { close(releaseTLSSideEffects) })
}
ss.serviceReportBeforeTLSSideEffectsHook = func(serviceID uint64) {
if serviceID == service.ID {
close(tlsSideEffectsReady)
<-releaseTLSSideEffects
}
}
t.Cleanup(func() {
releaseTLSSideEffectsHook()
ss.Close()
})
report := serviceMonitorResult(1, service.ID, model.TaskTypeHTTPGet, true)
report.Data.Data = "issuer|2030-01-02 15:04:05 +0000 UTC"
// When
ss.Dispatch(report)
select {
case <-tlsSideEffectsReady:
case <-ctx.Done():
t.Fatal(ctx.Err())
}
responseLockAcquired := ss.serviceResponseDataStoreLock.TryLock()
if responseLockAcquired {
ss.serviceResponseDataStoreLock.Unlock()
t.Fatal("worker released the response lock before TLS side effects")
}
releaseTLSSideEffectsHook()
ss.Close()
// Then
ss.serviceResponseDataStoreLock.RLock()
cachedCertificate := ss.tlsCertCache[service.ID]
ss.serviceResponseDataStoreLock.RUnlock()
if cachedCertificate != report.Data.Data {
t.Fatalf("expected TLS cache %q, got %q", report.Data.Data, cachedCertificate)
}
}
// TestServiceSentinelWorkerSurvivesConcurrentReporterServerDelete is a
// regression test for GHSA-jx78-55p5-rwv5 Finding 1 (incomplete fix of
// GHSA-qjpp-gffx-2wm9).
//
// The vulnerability: after the 2026-07-21 fix, the worker re-validates the
// service under serviceResponseDataStoreLock, but then takes a fresh snapshot
// m := ServerShared.GetList() with no guard. A concurrent batch-delete of the
// reporter's own server removes it between the pre-lock validation and the
// GetList call, so m[r.Reporter] is nil. delayCheck and notifyCheck then
// dereference m[r.Reporter].Name unconditionally — SIGSEGV.
//
// The subprocess-isolation pattern is used because the pre-fix code path
// panicked (nil pointer dereference in an unrecovered goroutine), which would
// crash the whole test binary rather than simply failing a single test.
func TestServiceSentinelWorkerSurvivesConcurrentReporterServerDelete(t *testing.T) {
if os.Getenv("NEZHA_SENTINEL_CONCURRENT_DELETE_CHILD") == "1" {
testServiceSentinelWorkerSurvivesConcurrentReporterServerDeleteChild(t)
return
}
ctx, cancel := context.WithTimeout(t.Context(), 5*time.Second)
defer cancel()
child := exec.CommandContext(ctx, os.Args[0],
"-test.run=^TestServiceSentinelWorkerSurvivesConcurrentReporterServerDelete$",
"-test.v",
)
child.Env = append(os.Environ(), "NEZHA_SENTINEL_CONCURRENT_DELETE_CHILD=1")
output, err := child.CombinedOutput()
if ctx.Err() != nil {
t.Fatalf("child process timed out: %v\n%s", ctx.Err(), output)
}
if err != nil {
t.Fatalf("child process crashed (likely nil deref in delayCheck/notifyCheck): %v\n%s", err, output)
}
if !strings.Contains(string(output), concurrentServerDeleteSuccessMarker) {
t.Fatalf("child did not print success marker:\n%s", output)
}
}
func testServiceSentinelWorkerSurvivesConcurrentReporterServerDeleteChild(t *testing.T) {
// Given: a reporter server and a service with latency-alerting enabled so
// that delayCheck (the vulnerable sink at line 785) is exercised on every
// dispatch. MaxLatency=1 ensures delay=12 always exceeds the threshold and
// the notification branch (not just the mute-clear branch) is taken.
ss := newServiceMonitorSecurityHarness(t,
&model.Server{Common: model.Common{ID: 1, UserID: 1}, Name: "reporter"},
)
service := &model.Service{
Common: model.Common{ID: 10, UserID: 1},
Name: "latency-service",
Type: model.TaskTypeTCPPing,
Target: "example.invalid:443",
Duration: 3600,
Cover: model.ServiceCoverIgnoreAll,
SkipServers: map[uint64]bool{1: true},
LatencyNotify: true,
MaxLatency: 1,
}
addServiceMonitorSecurityService(t, ss, service)
reportProcessing := make(chan struct{})
releaseWorker := make(chan struct{})
var releaseOnce sync.Once
releaseWorkerFn := func() { releaseOnce.Do(func() { close(releaseWorker) }) }
// serviceReportValidatedHook runs while the report holds the lifecycle read
// lock. Deletion must therefore run in another goroutine and wait until this
// hook releases; attempting Delete here would try to upgrade the RWMutex.
ss.serviceReportValidatedHook = func(serviceID uint64) {
if serviceID == service.ID {
close(reportProcessing)
<-releaseWorker
}
}
t.Cleanup(func() {
releaseWorkerFn()
ss.Close()
})
// When
ss.Dispatch(serviceMonitorResult(1, service.ID, model.TaskTypeTCPPing, true))
select {
case <-reportProcessing:
case <-t.Context().Done():
t.Fatal(t.Context().Err())
}
deleteDone := make(chan struct{})
go func() {
ServerShared.Delete([]uint64{1})
close(deleteDone)
}()
releaseWorkerFn()
select {
case <-deleteDone:
case <-t.Context().Done():
t.Fatal(t.Context().Err())
}
ss.Close()
// Then: no crash; the worker handled the nil reporter gracefully.
if _, err := fmt.Fprintln(os.Stdout, concurrentServerDeleteSuccessMarker); err != nil {
t.Fatal(err)
}
}
// TestServiceSentinelDeleteWithUnknownIDDoesNotLeaveZombies is a regression
// test for GHSA-jx78-55p5-rwv5 Finding 2 (low severity).
//
// The vulnerability: ServiceSentinel.Delete iterates the caller-supplied id
// slice and does CronShared.Remove(ss.services[id].CronJobID) without checking
// whether id is present in ss.services. CheckPermission returns vacuously
// true for unknown ids, so the controller layer cannot block this path.
// ss.services[unknownID] returns nil, and .CronJobID panics. Because the
// panic aborts the loop, every id ordered AFTER the bogus one is never removed
// from the in-memory registry even though its database row was already deleted,
// producing zombie services that keep dispatching cron probes.
func TestServiceSentinelDeleteWithUnknownIDDoesNotLeaveZombies(t *testing.T) {
// Given: one legitimate service (ID 10) registered in the sentinel.
ss := newServiceMonitorSecurityHarness(t,
&model.Server{Common: model.Common{ID: 1, UserID: 1}, Name: "reporter"},
)
service := &model.Service{
Common: model.Common{ID: 10, UserID: 1},
Name: "real-service",
Type: model.TaskTypeTCPPing,
Target: "example.invalid:443",
Duration: 3600,
Cover: model.ServiceCoverIgnoreAll,
SkipServers: map[uint64]bool{1: true},
}
addServiceMonitorSecurityService(t, ss, service)
// When: Delete is called with a bogus ID first, then the real service ID.
// Before the fix this panicked on ss.services[99999].CronJobID and left
// service 10 as a zombie.
ss.Delete([]uint64{99999, service.ID})
// Then: the real service must be fully removed from every in-memory map.
ss.serviceResponseDataStoreLock.RLock()
_, todayPresent := ss.serviceStatusToday[service.ID]
_, pingPresent := ss.serviceResponsePing[service.ID]
ss.serviceResponseDataStoreLock.RUnlock()
ss.servicesLock.RLock()
_, servicePresent := ss.services[service.ID]
ss.servicesLock.RUnlock()
ss.monthlyStatusLock.Lock()
_, monthlyPresent := ss.monthlyStatus[service.ID]
ss.monthlyStatusLock.Unlock()
if todayPresent {
t.Error("zombie: serviceStatusToday still contains the deleted service")
}
if pingPresent {
t.Error("zombie: serviceResponsePing still contains the deleted service")
}
if servicePresent {
t.Error("zombie: services map still contains the deleted service")
}
if monthlyPresent {
t.Error("zombie: monthlyStatus still contains the deleted service")
}
if _, err := fmt.Fprintln(os.Stdout, deleteUnknownIDSuccessMarker); err != nil {
t.Fatal(err)
}
}
service/singleton/singleton.go
+40 -8
View File
@@ -11,7 +11,6 @@ import (
"github.com/gin-gonic/gin"
"github.com/patrickmn/go-cache"
"gorm.io/driver/sqlite"
"gorm.io/gorm"
"sigs.k8s.io/yaml"
@@ -34,6 +33,10 @@ var (
NotificationShared *NotificationClass
NATShared *NATClass
CronShared *CronClass
// ServerTransferShared is initialized in LoadSingleton AFTER ServerShared
// (so the in-memory pending index can write back into ServerShared.UserID
// on transitions) and AFTER initUser (so PushIfOnline can read secrets
// from UserInfoMap).
)
//go:embed frontend-templates.yaml
@@ -59,6 +62,7 @@ func LoadSingleton(bus chan<- *model.Service) (err error) {
NotificationShared = NewNotificationClass()
ServerShared = NewServerClass()
CronShared = NewCronClass()
ServerTransferShared = NewServerTransferClass()
// 最后初始化 ServiceSentinel
ServiceSentinelShared, err = NewServiceSentinel(bus)
if err == nil {
@@ -79,7 +83,7 @@ func InitFrontendTemplates() error {
// InitDBFromPath 从给出的文件路径中加载数据库
func InitDBFromPath(path string) error {
var err error
DB, err = gorm.Open(sqlite.Open(path), &gorm.Config{
DB, err = gorm.Open(openSQLiteDialector(path), &gorm.Config{
CreateBatchSize: 200,
})
if err != nil {
@@ -92,11 +96,22 @@ func InitDBFromPath(path string) error {
model.Notification{}, model.AlertRule{}, model.Service{}, model.NotificationGroupNotification{},
model.Cron{}, model.Transfer{}, model.ServerGroupServer{},
model.NAT{}, model.DDNSProfile{}, model.NotificationGroupNotification{},
model.WAF{}, model.Oauth2Bind{}, model.Domain{})
model.WAF{}, model.Oauth2Bind{}, model.Domain{}, model.ServerTransfer{}, model.JWTSession{},
model.APIToken{}, model.MCPAuditLog{})
if err != nil {
return err
}
// 旧 mcp:* scope 与 nezha:* 并行了一段时间，HasScope 通过别名让 mcp:fs:write
// 静默扩到 REST nezha:server:write。统一命名后这里把残留旧 scope 一次性
// 归一化（或在仅剩危险旧 scope 时整张 PAT 删除），保证运行时不再依赖别名。
if rewritten, deleted, mErr := model.MigrateLegacyMCPScopes(DB); mErr != nil {
log.Printf("NEZHA>> MigrateLegacyMCPScopes failed: %v", mErr)
} else if rewritten > 0 || deleted > 0 {
log.Printf("NEZHA>> Migrated legacy mcp:* api token scopes: rewritten=%d deleted=%d", rewritten, deleted)
}
return nil
}
@@ -114,16 +129,15 @@ func RecordTransferHourlyUsage(servers ...*model.Server) {
}
for server := range slist {
_, _, deltaIn, deltaOut := server.TransferDeltaAndAdvance()
tx := model.Transfer{
ServerID: server.ID,
In: utils.SubUintChecked(server.State.NetInTransfer, server.PrevTransferInSnapshot),
Out: utils.SubUintChecked(server.State.NetOutTransfer, server.PrevTransferOutSnapshot),
In: deltaIn,
Out: deltaOut,
}
if tx.In == 0 && tx.Out == 0 {
continue
}
server.PrevTransferInSnapshot = server.State.NetInTransfer
server.PrevTransferOutSnapshot = server.State.NetOutTransfer
tx.CreatedAt = nowTrimSeconds
txs = append(txs, tx)
}
@@ -134,6 +148,13 @@ func RecordTransferHourlyUsage(servers ...*model.Server) {
log.Printf("NEZHA>> Saved traffic metrics to database. Affected %d row(s), Error: %v", len(txs), DB.Create(txs).Error)
}
func PersistTransfer(transfer model.Transfer) error {
if transfer.In == 0 && transfer.Out == 0 {
return nil
}
return DB.Create(&transfer).Error
}
// CleanMonitorHistory 清理流量记录（TSDB 有自己的保留策略）
func CleanMonitorHistory() {
// 清理已被删除的服务器的流量记录
@@ -143,7 +164,10 @@ func CleanMonitorHistory() {
specialServerKeep := make(map[uint64]time.Time)
var specialServerIDs []uint64
var alerts []model.AlertRule
DB.Find(&alerts)
if err := DB.Find(&alerts).Error; err != nil {
log.Printf("NEZHA>> Failed to load alert rules while cleaning transfer history: %v", err)
return
}
for _, alert := range alerts {
for _, rule := range alert.Rules {
// 是不是流量记录规则
@@ -171,6 +195,14 @@ func CleanMonitorHistory() {
for id, couldRemove := range specialServerKeep {
DB.Unscoped().Delete(&model.Transfer{}, "server_id = ? AND datetime(`created_at`) < datetime(?)", id, couldRemove)
}
if len(specialServerIDs) == 0 {
if allServerKeep.IsZero() {
DB.Unscoped().Session(&gorm.Session{AllowGlobalUpdate: true}).Delete(&model.Transfer{})
} else {
DB.Unscoped().Delete(&model.Transfer{}, "datetime(`created_at`) < datetime(?)", allServerKeep)
}
return
}
if allServerKeep.IsZero() {
DB.Unscoped().Delete(&model.Transfer{}, "server_id NOT IN (?)", specialServerIDs)
} else {
service/singleton/sqlite_attribution_boundary_agentcompat_linux_test.go
+83
View File
@@ -0,0 +1,83 @@
//go:build agentcompat && linux
package singleton
import (
"errors"
"path/filepath"
"strings"
"testing"
)
func TestSQLiteAttributionErrorsDoNotExposeDatabasePath(t *testing.T) {
// Given
databasePath := filepath.Join(t.TempDir(), "private-dashboard.sqlite")
unsupportedDSN := "file:" + databasePath + "?mode=memory"
missingJournal := filepath.Join(t.TempDir(), "missing-journal")
// When
memoryDatabase, dsnErr := openSQLiteAttributionTestDB(unsupportedDSN)
if dsnErr == nil {
dsnErr = memoryDatabase.Ping()
}
if memoryDatabase != nil {
t.Cleanup(func() {
if closeErr := memoryDatabase.Close(); closeErr != nil {
t.Error(closeErr)
}
})
}
_, journalErr := sqliteAttributionOpenJournalDescriptor(missingJournal)
// Then
if !errors.Is(dsnErr, ErrSQLiteAttributionUnsupportedDSN) {
t.Fatal("file URI error does not wrap the typed unsupported DSN error")
}
if !errors.Is(journalErr, ErrSQLiteAttributionJournalIdentity) {
t.Fatal("journal error does not wrap the typed journal identity error")
}
if strings.Contains(dsnErr.Error(), databasePath) || strings.Contains(dsnErr.Error(), unsupportedDSN) {
t.Fatal("unsupported DSN error exposes its database path")
}
if strings.Contains(journalErr.Error(), missingJournal) {
t.Fatal("journal identity error exposes its database path")
}
}
func TestSQLiteAttributionAcceptsOnDiskFileURIAndRejectsInMemoryDSN(t *testing.T) {
// Given
resetSQLiteAttributionForTest()
databasePath := filepath.Join(t.TempDir(), "dashboard.sqlite")
// When
fileDatabase, fileErr := openSQLiteAttributionTestDB("file:" + databasePath)
if fileDatabase != nil {
t.Cleanup(func() {
if closeErr := fileDatabase.Close(); closeErr != nil {
t.Error(closeErr)
}
})
}
if fileErr == nil {
fileErr = fileDatabase.Ping()
}
memoryDatabase, memoryErr := openSQLiteAttributionTestDB(":memory:")
if memoryDatabase != nil {
t.Cleanup(func() {
if closeErr := memoryDatabase.Close(); closeErr != nil {
t.Error(closeErr)
}
})
}
if memoryErr == nil {
memoryErr = memoryDatabase.Ping()
}
// Then
if fileErr != nil {
t.Fatal("file URI for an on-disk database was rejected")
}
if !errors.Is(memoryErr, ErrSQLiteAttributionUnsupportedDSN) {
t.Fatal("in-memory DSN does not return the typed unsupported DSN error")
}
}
service/singleton/sqlite_attribution_close_agentcompat_linux_test.go
+180
View File
@@ -0,0 +1,180 @@
//go:build agentcompat && linux
package singleton
import (
"context"
"database/sql"
"database/sql/driver"
"errors"
"testing"
"golang.org/x/sys/unix"
)
func TestSQLiteAttributionConnectionCloseFinalizesActiveWriteTransaction(t *testing.T) {
// Given
resetSQLiteAttributionForTest()
databasePath := sqliteAttributionTestDatabasePath(t)
rawConnection, err := sqliteAttributionDriver{}.Open(databasePath)
if err != nil {
t.Fatal(err)
}
connection := rawConnection.(*sqliteAttributionConnection)
if _, err := connection.connection.Exec("CREATE TABLE settings (id INTEGER PRIMARY KEY, value TEXT)", nil); err != nil {
t.Fatal(err)
}
enableSQLiteAttribution()
transaction, err := connection.Begin()
if err != nil {
t.Fatal(err)
}
statement, err := connection.Prepare("INSERT INTO settings (value) VALUES (?)")
if err != nil {
t.Fatal(err)
}
if _, err := statement.Exec([]driver.Value{"close-active"}); err != nil {
t.Fatal(err)
}
if err := statement.Close(); err != nil {
t.Fatal(err)
}
identity, descriptor, active := sqliteAttributionTransactionState(t, connection)
if !active {
t.Fatal("active transaction is missing before Close")
}
// When
firstCloseErr := connection.Close()
secondCloseErr := connection.Close()
commitErr := transaction.Commit()
rollbackErr := transaction.Rollback()
_, descriptorErr := unix.FcntlInt(uintptr(descriptor), unix.F_GETFD, 0)
standard, openErr := sql.Open("sqlite3", databasePath)
if openErr != nil {
t.Fatal(openErr)
}
defer standard.Close()
var count int
countErr := standard.QueryRow("SELECT COUNT(*) FROM settings").Scan(&count)
tracker := sqliteAttributionTracker.Load()
tracker.mu.Lock()
_, active = tracker.transactions[identity]
tracker.mu.Unlock()
// Then
if firstCloseErr != nil || secondCloseErr != nil {
t.Fatalf("Close errors = %v / %v", firstCloseErr, secondCloseErr)
}
if !errors.Is(commitErr, driver.ErrBadConn) {
t.Fatalf("Commit after Close error = %v, want driver.ErrBadConn", commitErr)
}
if !errors.Is(rollbackErr, driver.ErrBadConn) {
t.Fatalf("Rollback after Close error = %v, want driver.ErrBadConn", rollbackErr)
}
if !errors.Is(descriptorErr, unix.EBADF) {
t.Fatalf("journal descriptor after Close = %v, want EBADF", descriptorErr)
}
if countErr != nil || count != 0 {
t.Fatalf("closed active transaction persisted count=%d err=%v", count, countErr)
}
if active {
t.Fatal("Close left the tracker transaction active")
}
}
func TestSQLiteAttributionConnectionCloseWakesSelectedCommit(t *testing.T) {
// Given
resetSQLiteAttributionForTest()
databasePath := sqliteAttributionTestDatabasePath(t)
rawConnection, err := sqliteAttributionDriver{}.Open(databasePath)
if err != nil {
t.Fatal(err)
}
connection := rawConnection.(*sqliteAttributionConnection)
if _, err := connection.connection.Exec("CREATE TABLE settings (id INTEGER PRIMARY KEY, value TEXT)", nil); err != nil {
t.Fatal(err)
}
enableSQLiteAttribution()
transaction, err := connection.Begin()
if err != nil {
t.Fatal(err)
}
statement, err := connection.Prepare("INSERT INTO settings (value) VALUES (?)")
if err != nil {
t.Fatal(err)
}
if _, err := statement.Exec([]driver.Value{"close-wakes-commit"}); err != nil {
t.Fatal(err)
}
if err := statement.Close(); err != nil {
t.Fatal(err)
}
_, descriptor, active := sqliteAttributionTransactionState(t, connection)
if !active {
t.Fatal("active transaction is missing before selected Commit")
}
tracker := sqliteAttributionTracker.Load()
session, err := tracker.ArmNextSQLiteHold()
if err != nil {
t.Fatal(err)
}
finalizing := make(chan error, 1)
commit := make(chan error, 1)
go func() {
_, waitErr := tracker.WaitSQLiteHold(context.Background(), session, SQLiteHoldWaitFinalizing)
finalizing <- waitErr
}()
// When
go func() { commit <- transaction.Commit() }()
if finalizingErr := <-finalizing; finalizingErr != nil {
t.Fatalf("Commit finalization wait error = %v", finalizingErr)
}
closeErr := connection.Close()
commitErr := <-commit
_, descriptorErr := unix.FcntlInt(uintptr(descriptor), unix.F_GETFD, 0)
standard, err := sql.Open("sqlite3", databasePath)
if err != nil {
t.Fatal(err)
}
defer standard.Close()
var count int
if err := standard.QueryRow("SELECT COUNT(*) FROM settings").Scan(&count); err != nil {
t.Fatal(err)
}
// Then
if closeErr != nil {
t.Fatal(closeErr)
}
var holdErr *SQLiteHoldError
if !errors.As(commitErr, &holdErr) || !errors.Is(commitErr, ErrSQLiteHoldAborted) {
t.Fatalf("Commit error after Close = %v", commitErr)
}
if !errors.Is(descriptorErr, unix.EBADF) {
t.Fatalf("journal descriptor after Close = %v, want EBADF", descriptorErr)
}
if count != 0 {
t.Fatalf("Close while Commit waited persisted %d rows", count)
}
}
type sqliteAttributionBadConnRows struct{}
func (sqliteAttributionBadConnRows) Columns() []string { return []string{"value"} }
func (sqliteAttributionBadConnRows) Close() error { return nil }
func (sqliteAttributionBadConnRows) Next([]driver.Value) error { return driver.ErrBadConn }
func TestSQLiteAttributionReadonlyRowsReturnBadConnWithoutPanic(t *testing.T) {
// Given
rows := &sqliteAttributionRows{rows: sqliteAttributionBadConnRows{}}
// When
err := rows.Next(make([]driver.Value, 1))
// Then
if !errors.Is(err, driver.ErrBadConn) {
t.Fatalf("readonly rows Next error = %v, want driver.ErrBadConn", err)
}
}
service/singleton/sqlite_attribution_commit_linearization_agentcompat_linux_test.go
+102
View File
@@ -0,0 +1,102 @@
//go:build agentcompat && linux
package singleton
import (
"context"
"errors"
"testing"
"golang.org/x/sys/unix"
)
func TestSQLiteAttributionCommitReleaseLinearizesBeforeContextCancellation(t *testing.T) {
// Given
transactionContext, cancel := context.WithCancel(context.Background())
defer cancel()
connection, transaction, session, databasePath := sqliteAttributionHeldTransaction(t, "released-before-cancel", transactionContext)
identity, descriptor, active := sqliteAttributionTransactionState(t, connection)
if !active {
t.Fatal("released transaction is not active")
}
// When
commit := sqliteAttributionStartHeldCommit(t, transaction, session)
if err := sqliteAttributionTracker.Load().ReleaseSQLiteHold(session); err != nil {
t.Fatal(err)
}
cancel()
commitErr := <-commit
terminal, terminalErr := sqliteAttributionTracker.Load().WaitSQLiteHold(context.Background(), session, SQLiteHoldWaitFinalizing)
_, descriptorErr := unix.FcntlInt(uintptr(descriptor), unix.F_GETFD, 0)
// Then
if commitErr != nil {
t.Fatalf("released Commit error after context cancellation = %v", commitErr)
}
if terminalErr != nil || !terminal.Released || !terminal.Selected || !terminal.Finalizing {
t.Fatalf("released terminal=%+v err=%v", terminal, terminalErr)
}
if !errors.Is(descriptorErr, unix.EBADF) {
t.Fatalf("journal descriptor after released Commit = %v, want EBADF", descriptorErr)
}
if count := sqliteAttributionPersistedCount(t, databasePath); count != 1 {
t.Fatalf("released Commit persisted %d rows", count)
}
if _, _, active := sqliteAttributionTransactionState(t, connection); active {
t.Fatal("released Commit left the connection transaction active")
}
if sqliteAttributionTrackerTransactionActive(sqliteAttributionTracker.Load(), identity) {
t.Fatal("released Commit left the tracker transaction active")
}
}
func TestSQLiteAttributionCommitCancellationAbortsBeforeRelease(t *testing.T) {
// Given
transactionContext, cancel := context.WithCancel(context.Background())
defer cancel()
connection, transaction, session, databasePath := sqliteAttributionHeldTransaction(t, "cancelled-before-release", transactionContext)
identity, descriptor, active := sqliteAttributionTransactionState(t, connection)
if !active {
t.Fatal("cancelled transaction is not active")
}
// When
commit := sqliteAttributionStartHeldCommit(t, transaction, session)
cancel()
commitErr := <-commit
releaseErr := sqliteAttributionTracker.Load().ReleaseSQLiteHold(session)
_, terminalErr := sqliteAttributionTracker.Load().WaitSQLiteHold(context.Background(), session, SQLiteHoldWaitFinalizing)
_, descriptorErr := unix.FcntlInt(uintptr(descriptor), unix.F_GETFD, 0)
tracker := sqliteAttributionTracker.Load()
tracker.mu.Lock()
terminal := tracker.terminal
tracker.mu.Unlock()
// Then
var holdErr *SQLiteHoldError
if !errors.As(commitErr, &holdErr) || !errors.Is(commitErr, ErrSQLiteHoldAborted) || !errors.Is(commitErr, context.Canceled) {
t.Fatalf("cancelled Commit error = %v", commitErr)
}
if !errors.Is(releaseErr, ErrSQLiteHoldStaleSession) {
t.Fatalf("release after cancellation-owned abort = %v", releaseErr)
}
if !errors.Is(terminalErr, ErrSQLiteHoldAborted) {
t.Fatalf("cancelled terminal wait error = %v, want aborted", terminalErr)
}
if terminal == nil || terminal.released {
t.Fatalf("cancelled terminal state = %+v, want aborted and unreleased", terminal)
}
if !errors.Is(descriptorErr, unix.EBADF) {
t.Fatalf("journal descriptor after cancelled Commit = %v, want EBADF", descriptorErr)
}
if count := sqliteAttributionPersistedCount(t, databasePath); count != 0 {
t.Fatalf("cancelled Commit persisted %d rows", count)
}
if _, _, active := sqliteAttributionTransactionState(t, connection); active {
t.Fatal("cancelled Commit left the connection transaction active")
}
if sqliteAttributionTrackerTransactionActive(sqliteAttributionTracker.Load(), identity) {
t.Fatal("cancelled Commit left the tracker transaction active")
}
}
service/singleton/sqlite_attribution_completion_agentcompat_linux_test.go
+154
View File
@@ -0,0 +1,154 @@
//go:build agentcompat && linux
package singleton
import (
"database/sql/driver"
"errors"
"sync"
"sync/atomic"
"testing"
)
type sqliteAttributionBlockingTx struct {
commitStarted chan struct{}
allowCommit chan struct{}
commitCalls atomic.Int32
rollbackCalls atomic.Int32
lifecycleMu *sync.Mutex
lockFree atomic.Bool
}
func (transaction *sqliteAttributionBlockingTx) Commit() error {
transaction.commitCalls.Add(1)
// Probe before publishing entry so losing terminal calls cannot contend with this boundary check.
if transaction.lifecycleMu.TryLock() {
transaction.lockFree.Store(true)
transaction.lifecycleMu.Unlock()
}
close(transaction.commitStarted)
<-transaction.allowCommit
return nil
}
func (transaction *sqliteAttributionBlockingTx) Rollback() error {
transaction.rollbackCalls.Add(1)
return nil
}
func TestSQLiteAttributionTransactionCompletionRunsRawCommitExactlyOnce(t *testing.T) {
// Given
tracker := NewSQLiteHoldTracker()
identity := sqliteHoldTestTransaction(201)
if err := tracker.BeginSQLiteTransaction(identity); err != nil {
t.Fatal(err)
}
raw := &sqliteAttributionBlockingTx{commitStarted: make(chan struct{}), allowCommit: make(chan struct{})}
state := &sqliteAttributionTransaction{
transaction: identity,
raw: raw,
tracker: tracker,
journalFD: -1,
done: make(chan struct{}),
}
var rawCloseCalls atomic.Int32
connection := &sqliteAttributionConnection{transaction: state}
connection.closeRawConnection = func() error {
if !connection.lifecycleMu.TryLock() {
return errors.New("raw connection Close ran while lifecycle lock was held")
}
connection.lifecycleMu.Unlock()
rawCloseCalls.Add(1)
return nil
}
raw.lifecycleMu = &connection.lifecycleMu
owner := &sqliteAttributionTx{connection: connection, state: state}
secondCommit := &sqliteAttributionTx{connection: connection, state: state}
commitResult := make(chan error, 1)
secondCommitResult := make(chan error, 1)
rollbackResult := make(chan error, 1)
closeResult := make(chan error, 1)
secondCommitStarted := make(chan struct{})
rollbackStarted := make(chan struct{})
closeStarted := make(chan struct{})
// When
go func() { commitResult <- owner.Commit() }()
<-raw.commitStarted
go func() {
close(secondCommitStarted)
secondCommitResult <- secondCommit.Commit()
}()
go func() {
close(rollbackStarted)
rollbackResult <- owner.Rollback()
}()
go func() {
close(closeStarted)
closeResult <- connection.Close()
}()
<-secondCommitStarted
<-rollbackStarted
<-closeStarted
for _, result := range []<-chan error{secondCommitResult, rollbackResult} {
select {
case err := <-result:
t.Fatalf("completion loser returned before raw Commit completion: %v", err)
default:
}
}
select {
case <-state.done:
t.Fatal("completion published before raw Commit was allowed to finish")
default:
}
if calls := rawCloseCalls.Load(); calls != 0 {
t.Fatalf("raw connection Close calls while raw Commit was in flight = %d, want 0", calls)
}
if !raw.lockFree.Load() {
t.Fatal("raw Commit ran while lifecycle lock was held")
}
if !sqliteAttributionTrackerTransactionActive(tracker, identity) {
t.Fatal("tracker transaction became inactive while raw Commit was blocked")
}
close(raw.allowCommit)
commitErr := <-commitResult
secondCommitErr := <-secondCommitResult
rollbackErr := <-rollbackResult
closeErr := <-closeResult
// Then
if commitErr != nil {
t.Fatalf("raw Commit error = %v", commitErr)
}
for _, loserErr := range []error{secondCommitErr, rollbackErr} {
if !errors.Is(loserErr, driver.ErrBadConn) {
t.Fatalf("completion loser error = %v, want driver.ErrBadConn", loserErr)
}
}
if closeErr != nil && !errors.Is(closeErr, driver.ErrBadConn) {
t.Fatalf("connection Close error = %v, want nil or driver.ErrBadConn", closeErr)
}
select {
case <-state.done:
default:
t.Fatal("completion signal did not close after raw Commit finished")
}
if calls := raw.commitCalls.Load(); calls != 1 {
t.Fatalf("raw Commit calls = %d, want 1", calls)
}
if calls := rawCloseCalls.Load(); calls != 1 {
t.Fatalf("raw connection Close calls after raw Commit completed = %d, want 1", calls)
}
if calls := raw.rollbackCalls.Load(); calls != 0 {
t.Fatalf("raw Rollback calls while raw Commit owned completion = %d, want 0", calls)
}
if _, _, active := sqliteAttributionTransactionState(t, connection); active {
t.Fatal("completed transaction remained attached to the connection")
}
if sqliteAttributionTrackerTransactionActive(tracker, identity) {
t.Fatal("completed transaction remained active in the tracker")
}
}
var _ driver.Tx = (*sqliteAttributionBlockingTx)(nil)
service/singleton/sqlite_attribution_connection_close_agentcompat_linux.go
+26
View File
@@ -0,0 +1,26 @@
//go:build agentcompat && linux
package singleton
import (
"errors"
)
func (connection *sqliteAttributionConnection) Close() error {
connection.lifecycleMu.Lock()
state := connection.transaction
connection.lifecycleMu.Unlock()
if state == nil {
return connection.closeRaw()
}
// database/sql may close Conn before Tx reaches Rollback; converge the attribution state here.
rollbackErr := state.rollback(connection)
return errors.Join(rollbackErr, connection.closeRaw())
}
func (connection *sqliteAttributionConnection) closeRaw() error {
if connection.closeRawConnection != nil {
return connection.closeRawConnection()
}
return connection.connection.Close()
}
service/singleton/sqlite_attribution_connection_dispatch_agentcompat_linux.go
+52
View File
@@ -0,0 +1,52 @@
//go:build agentcompat && linux
package singleton
import (
"context"
"database/sql/driver"
"errors"
)
func (connection *sqliteAttributionConnection) Exec(query string, values []driver.Value) (driver.Result, error) {
statement, err := connection.Prepare(query)
if err != nil {
return nil, err
}
defer statement.Close()
return statement.Exec(values)
}
func (connection *sqliteAttributionConnection) ExecContext(ctx context.Context, query string, values []driver.NamedValue) (driver.Result, error) {
statement, err := connection.PrepareContext(ctx, query)
if err != nil {
return nil, err
}
defer statement.Close()
return statement.(driver.StmtExecContext).ExecContext(ctx, values)
}
func (connection *sqliteAttributionConnection) Query(query string, values []driver.Value) (driver.Rows, error) {
statement, err := connection.Prepare(query)
if err != nil {
return nil, err
}
wrapped := statement.(*sqliteAttributionStatement)
return wrapped.queryOwned(func() (driver.Rows, error) { return wrapped.statement.Query(values) }, statement)
}
func (connection *sqliteAttributionConnection) QueryContext(ctx context.Context, query string, values []driver.NamedValue) (driver.Rows, error) {
statement, err := connection.PrepareContext(ctx, query)
if err != nil {
return nil, err
}
return connection.queryContextStatement(ctx, statement.(*sqliteAttributionStatement), values, statement)
}
func (connection *sqliteAttributionConnection) queryContextStatement(ctx context.Context, statement *sqliteAttributionStatement, values []driver.NamedValue, owner driver.Stmt) (driver.Rows, error) {
contextStatement, ok := statement.statement.(driver.StmtQueryContext)
if !ok {
return nil, errors.Join(driver.ErrSkip, owner.Close())
}
return statement.queryOwned(func() (driver.Rows, error) { return contextStatement.QueryContext(ctx, values) }, owner)
}
service/singleton/sqlite_attribution_evidence_agentcompat_linux.go
+149
View File
@@ -0,0 +1,149 @@
//go:build agentcompat && linux
package singleton
import (
"errors"
"runtime"
"strings"
"golang.org/x/sys/unix"
)
func (connection *sqliteAttributionConnection) beforeWrite(classification sqliteAttributionClassification) error {
if !sqliteAttributionEnabled.Load() || classification.readonly {
return nil
}
if !classification.valid() {
return &SQLiteAttributionError{Cause: ErrSQLiteAttributionUnsupportedWrite}
}
connection.lifecycleMu.Lock()
state := connection.transaction
if state == nil {
connection.lifecycleMu.Unlock()
return &SQLiteAttributionError{Cause: ErrSQLiteAttributionUnboundWrite}
}
poison := state.poison
connection.lifecycleMu.Unlock()
if poison != nil {
return poison
}
connection.execution = &sqliteAttributionExecution{classification: classification, origin: sqliteAttributionOrigin()}
return nil
}
func (connection *sqliteAttributionConnection) discardExecution() { connection.execution = nil }
func (connection *sqliteAttributionConnection) publishExecution() error {
execution := connection.execution
connection.execution = nil
if execution == nil {
return nil
}
connection.lifecycleMu.Lock()
state := connection.transaction
if state == nil {
connection.lifecycleMu.Unlock()
return &SQLiteAttributionError{Cause: ErrSQLiteAttributionUnboundWrite}
}
if execution.hook.mismatch || !execution.hook.seen {
err := &SQLiteAttributionError{Cause: ErrSQLiteAttributionUnsupportedWrite}
state.poison = err
connection.lifecycleMu.Unlock()
return err
}
if state.journalFD < 0 {
descriptor, err := sqliteAttributionOpenJournalDescriptor(connection.journal)
if err != nil {
state.poison = err
connection.lifecycleMu.Unlock()
return err
}
state.journalFD = descriptor
}
journal, err := sqliteAttributionJournalIdentityFromDescriptor(state.journalFD)
if err != nil {
state.poison = err
connection.lifecycleMu.Unlock()
return err
}
transaction := state.transaction
tracker := state.tracker
connection.lifecycleMu.Unlock()
origin := execution.origin
origin.Operation = execution.classification.operation
origin.Table = execution.classification.table
err = tracker.RecordSQLiteWrite(transaction, SQLiteWriteObservation{
Origin: origin,
Update: SQLiteUpdateObservation{Operation: execution.classification.operation, Table: execution.classification.table, Journal: journal},
})
if err != nil {
return connection.poison(err)
}
return nil
}
func (connection *sqliteAttributionConnection) poison(err error) error {
connection.lifecycleMu.Lock()
state := connection.transaction
if state != nil {
state.poison = err
}
connection.lifecycleMu.Unlock()
return err
}
// The adapter owns the explicit main transaction and DELETE journal; path-only stat is racy, so retain this O_PATH descriptor through finalization.
func sqliteAttributionOpenJournalDescriptor(path string) (int, error) {
descriptor, err := unix.Open(path, unix.O_PATH|unix.O_NOFOLLOW|unix.O_CLOEXEC, 0)
if err != nil {
return -1, &SQLiteAttributionError{Cause: errors.Join(ErrSQLiteAttributionJournalIdentity, err)}
}
return descriptor, nil
}
func sqliteAttributionCloseJournalDescriptor(descriptor int) error {
if descriptor < 0 {
return nil
}
if err := unix.Close(descriptor); err != nil {
return &SQLiteAttributionError{Cause: errors.Join(ErrSQLiteAttributionJournalIdentity, err)}
}
return nil
}
func sqliteAttributionJournalIdentityFromDescriptor(descriptor int) (SQLiteJournalIdentity, error) {
var status unix.Statx_t
mask := uint32(unix.STATX_BASIC_STATS | unix.STATX_BTIME | unix.STATX_MNT_ID)
if err := unix.Statx(descriptor, "", unix.AT_EMPTY_PATH|unix.AT_STATX_SYNC_AS_STAT, int(mask), &status); err != nil {
return SQLiteJournalIdentity{}, &SQLiteAttributionError{Cause: errors.Join(ErrSQLiteAttributionJournalIdentity, err)}
}
required := uint32(unix.STATX_MNT_ID | unix.STATX_BTIME)
if status.Mask&required != required {
return SQLiteJournalIdentity{}, &SQLiteAttributionError{Cause: ErrSQLiteAttributionJournalIdentity}
}
return SQLiteJournalIdentity{MountID: status.Mnt_id, DeviceMajor: status.Dev_major, DeviceMinor: status.Dev_minor, Inode: status.Ino, BirthSeconds: status.Btime.Sec, BirthNanoseconds: status.Btime.Nsec}, nil
}
func sqliteAttributionOrigin() SQLiteExecutionOrigin {
programCounters := make([]uintptr, 16)
count := runtime.Callers(3, programCounters)
programCounters = programCounters[:count]
frames := runtime.CallersFrames(programCounters)
frame, more := frames.Next()
for more && !strings.Contains(frame.Function, "github.com/nezhahq/nezha/") {
frame, more = frames.Next()
}
return SQLiteExecutionOrigin{StackHash: sqliteAttributionStackHash(programCounters), FirstNezhaFrame: frame.Function}
}
func sqliteAttributionStackHash(programCounters []uintptr) uint64 {
const offsetBasis uint64 = 14695981039346656037
const prime uint64 = 1099511628211
hash := offsetBasis
for _, programCounter := range programCounters {
hash ^= uint64(programCounter)
hash *= prime
}
return hash
}
service/singleton/sqlite_attribution_hold_lifecycle_agentcompat_linux_test.go
+186
View File
@@ -0,0 +1,186 @@
//go:build agentcompat && linux
package singleton
import (
"context"
"errors"
"path/filepath"
"testing"
"time"
"github.com/nezhahq/nezha/model"
"gorm.io/gorm"
)
func TestSQLiteAttributionHoldFacadeEnablesOnArmAndDisablesOnAbort(t *testing.T) {
// Given
resetSQLiteAttributionForTest()
t.Cleanup(resetSQLiteAttributionForTest)
// When
receipt, armErr := ArmNextSQLiteHold()
enabledAfterArm := sqliteAttributionEnabled.Load()
_, abortErr := AbortSQLiteHold(receipt)
// Then
if armErr != nil || abortErr != nil {
t.Fatalf("arm=%v abort=%v", armErr, abortErr)
}
if !enabledAfterArm {
t.Fatal("successful hold arm did not enable SQLite attribution")
}
if sqliteAttributionEnabled.Load() {
t.Fatal("successful hold abort left SQLite attribution enabled")
}
}
func TestSQLiteAttributionHoldFacadeDisablesAfterRelease(t *testing.T) {
// Given
resetSQLiteAttributionForTest()
t.Cleanup(resetSQLiteAttributionForTest)
receipt, err := ArmNextSQLiteHold()
if err != nil {
t.Fatal(err)
}
transaction := sqliteHoldTestTransaction(204)
tracker := sqliteAttributionTracker.Load()
if err := tracker.BeginSQLiteTransaction(transaction); err != nil {
t.Fatal(err)
}
if err := tracker.RecordSQLiteUpdate(transaction, SQLiteUpdateObservation{Operation: SQLiteOperationUpdate, Table: "api_tokens", Journal: sqliteHoldTestJournal}); err != nil {
t.Fatal(err)
}
if _, err := tracker.BeginSQLiteFinalization(transaction); err != nil {
t.Fatal(err)
}
// When
result, releaseErr := ReleaseSQLiteHold(receipt)
// Then
if releaseErr != nil || result.State != SQLiteHoldControlStateReleased {
t.Fatalf("release=%+v err=%v", result, releaseErr)
}
if sqliteAttributionEnabled.Load() {
t.Fatal("successful hold release left SQLite attribution enabled")
}
}
func TestSQLiteAttributionHoldFacadeKeepsEnabledAfterCanceledWaitUntilAbort(t *testing.T) {
// Given
resetSQLiteAttributionForTest()
t.Cleanup(resetSQLiteAttributionForTest)
receipt, err := ArmNextSQLiteHold()
if err != nil {
t.Fatal(err)
}
ctx, cancel := context.WithCancel(context.Background())
cancel()
// When
_, waitErr := WaitSQLiteHoldSelected(ctx, receipt)
enabledAfterWait := sqliteAttributionEnabled.Load()
_, abortErr := AbortSQLiteHold(receipt)
// Then
if !errors.Is(waitErr, context.Canceled) || abortErr != nil {
t.Fatalf("wait=%v abort=%v", waitErr, abortErr)
}
if !enabledAfterWait {
t.Fatal("canceled wait disabled attribution before the active hold was aborted")
}
if sqliteAttributionEnabled.Load() {
t.Fatal("abort after canceled wait left SQLite attribution enabled")
}
}
func TestSQLiteAttributionHoldFacadeStaleReceiptCannotDisableNewHold(t *testing.T) {
// Given
resetSQLiteAttributionForTest()
t.Cleanup(resetSQLiteAttributionForTest)
first, err := ArmNextSQLiteHold()
if err != nil {
t.Fatal(err)
}
if _, err := AbortSQLiteHold(first); err != nil {
t.Fatal(err)
}
second, err := ArmNextSQLiteHold()
if err != nil {
t.Fatal(err)
}
// When
_, staleErr := AbortSQLiteHold(first)
// Then
if !errors.Is(staleErr, ErrSQLiteHoldStaleSession) {
t.Fatalf("stale abort error = %v", staleErr)
}
if !sqliteAttributionEnabled.Load() {
t.Fatal("stale receipt disabled attribution owned by the current hold")
}
if _, err := AbortSQLiteHold(second); err != nil {
t.Fatal(err)
}
}
func TestSQLiteAttributionHoldFacadeSelectsGORMAPITokenUsageUpdate(t *testing.T) {
// Given
resetSQLiteAttributionForTest()
t.Cleanup(resetSQLiteAttributionForTest)
database, err := gorm.Open(openSQLiteDialector(filepath.Join(t.TempDir(), "dashboard.sqlite")), &gorm.Config{})
if err != nil {
t.Fatal(err)
}
sqlDatabase, err := database.DB()
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() {
if closeErr := sqlDatabase.Close(); closeErr != nil {
t.Error(closeErr)
}
})
if err := database.AutoMigrate(&model.APIToken{}); err != nil {
t.Fatal(err)
}
token := model.APIToken{UserID: 1, Name: "usage-update", TokenHash: "hash"}
if err := database.Create(&token).Error; err != nil {
t.Fatal(err)
}
receipt, err := ArmNextSQLiteHold()
if err != nil {
t.Fatal(err)
}
ctx, cancel := context.WithCancel(t.Context())
defer cancel()
writerDone := make(chan error, 1)
usageTime := time.Unix(1_700_000_000, 0)
go func() {
writerDone <- database.Model(&model.APIToken{}).Where("id = ?", token.ID).Updates(map[string]any{
"last_used_at": usageTime,
"last_used_ip": "127.0.0.1",
}).Error
cancel()
}()
// When
selected, selectedErr := WaitSQLiteHoldSelected(ctx, receipt)
finalizing, finalizingErr := WaitSQLiteHoldFinalizing(ctx, selected)
_, releaseErr := ReleaseSQLiteHold(finalizing)
writerErr := <-writerDone
// Then
if selectedErr != nil || finalizingErr != nil || releaseErr != nil || writerErr != nil {
t.Fatalf("selected=%v finalizing=%v release=%v writer=%v", selectedErr, finalizingErr, releaseErr, writerErr)
}
var updated model.APIToken
if err := database.First(&updated, token.ID).Error; err != nil {
t.Fatal(err)
}
if updated.LastUsedAt == nil || !updated.LastUsedAt.Equal(usageTime) || updated.LastUsedIP != "127.0.0.1" {
t.Fatalf("persisted usage update = %+v", updated)
}
}
service/singleton/sqlite_attribution_journal_fd_agentcompat_linux_test.go
+42
View File
@@ -0,0 +1,42 @@
//go:build agentcompat && linux
package singleton
import (
"path/filepath"
"testing"
"golang.org/x/sys/unix"
)
func TestSQLiteAttributionDerivesJournalIdentityFromOpenedDescriptor(t *testing.T) {
// Given
journalPath := filepath.Join(t.TempDir(), "dashboard.sqlite-journal")
descriptor, err := unix.Open(journalPath, unix.O_CREAT|unix.O_WRONLY|unix.O_CLOEXEC, 0o600)
if err != nil {
t.Fatal(err)
}
if err := unix.Close(descriptor); err != nil {
t.Fatal(err)
}
descriptor, err = unix.Open(journalPath, unix.O_PATH|unix.O_NOFOLLOW|unix.O_CLOEXEC, 0)
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() {
if closeErr := unix.Close(descriptor); closeErr != nil {
t.Error(closeErr)
}
})
// When
identity, err := sqliteAttributionJournalIdentityFromDescriptor(descriptor)
// Then
if err != nil {
t.Fatal(err)
}
if identity.MountID == 0 || identity.Inode == 0 || identity.BirthSeconds == 0 {
t.Fatal("descriptor-derived journal identity is incomplete")
}
}
service/singleton/sqlite_attribution_lifecycle_agentcompat_linux_test.go
+211
View File
@@ -0,0 +1,211 @@
//go:build agentcompat && linux
package singleton
import (
"context"
"database/sql/driver"
"errors"
"testing"
"golang.org/x/sys/unix"
)
func TestSQLiteAttributionReturningCompletesExactlyOneRow(t *testing.T) {
// Given
database := openSQLiteAttributionTestDatabase(t)
enableSQLiteAttribution()
transaction, err := database.BeginTx(context.Background(), nil)
if err != nil {
t.Fatal(err)
}
// When
rows, err := transaction.QueryContext(context.Background(), "INSERT INTO settings (value) VALUES (?) RETURNING id", "one-row")
if err != nil {
t.Fatal(err)
}
if !rows.Next() {
t.Fatal("RETURNING did not yield its inserted row")
}
var identifier int64
if err := rows.Scan(&identifier); err != nil {
t.Fatal(err)
}
if rows.Next() {
t.Fatal("RETURNING yielded more than one row")
}
rowsErr := rows.Err()
closeErr := rows.Close()
commitErr := transaction.Commit()
count := sqliteAttributionSettingsCount(t, database)
// Then
if rowsErr != nil || closeErr != nil || commitErr != nil {
t.Fatal("completed RETURNING did not finish successfully")
}
if count != 1 {
t.Fatalf("completed RETURNING persisted %d rows", count)
}
}
func TestSQLiteAttributionEarlyReturningCloseRollsBackTransaction(t *testing.T) {
// Given
database := openSQLiteAttributionTestDatabase(t)
enableSQLiteAttribution()
transaction, err := database.BeginTx(context.Background(), nil)
if err != nil {
t.Fatal(err)
}
// When
rows, err := transaction.QueryContext(context.Background(), "INSERT INTO settings (value) VALUES (?) RETURNING id", "early-close")
if err != nil {
t.Fatal(err)
}
if !rows.Next() {
t.Fatal("RETURNING did not yield its inserted row")
}
if err := rows.Close(); err != nil {
t.Fatal(err)
}
commitErr := transaction.Commit()
count := sqliteAttributionSettingsCount(t, database)
// Then
if commitErr == nil {
t.Fatal("early RETURNING Close allowed Commit")
}
if count != 0 {
t.Fatalf("early RETURNING Close persisted %d rows", count)
}
}
func TestSQLiteAttributionZeroRowUpdateDoesNotPublishEvidence(t *testing.T) {
// Given
database := openSQLiteAttributionTestDatabase(t)
enableSQLiteAttribution()
transaction, err := database.BeginTx(context.Background(), nil)
if err != nil {
t.Fatal(err)
}
// When
if _, err := transaction.Exec("UPDATE settings SET value = ? WHERE id = ?", "zero", -1); err != nil {
t.Fatal(err)
}
commitErr := transaction.Commit()
evidence := sqliteAttributionTrackerWriteEvidence()
// Then
if commitErr != nil {
t.Fatal(commitErr)
}
if evidence.hasWrite {
t.Fatal("zero-row UPDATE published write evidence")
}
}
func TestSQLiteAttributionDirectReadQueryClosesOwnedStatement(t *testing.T) {
// Given
database := openSQLiteAttributionTestDatabase(t)
// When
rows, err := database.QueryContext(context.Background(), "SELECT id FROM settings")
if err != nil {
t.Fatal(err)
}
next := rows.Next()
rowsErr := rows.Err()
closeErr := rows.Close()
// Then
if next {
t.Fatal("empty SELECT returned a row")
}
if rowsErr != nil || closeErr != nil {
t.Fatal("direct readonly Query did not complete cleanly")
}
}
func TestSQLiteAttributionCommitClosesRetainedJournalDescriptor(t *testing.T) {
// Given
resetSQLiteAttributionForTest()
rawConnection, err := sqliteAttributionDriver{}.Open(sqliteAttributionTestDatabasePath(t))
if err != nil {
t.Fatal(err)
}
connection := rawConnection.(*sqliteAttributionConnection)
t.Cleanup(func() {
if closeErr := connection.Close(); closeErr != nil {
t.Error(closeErr)
}
})
create, err := connection.Prepare("CREATE TABLE settings (id INTEGER PRIMARY KEY, value TEXT)")
if err != nil {
t.Fatal(err)
}
if _, err := create.Exec(nil); err != nil {
t.Fatal(err)
}
if err := create.Close(); err != nil {
t.Fatal(err)
}
enableSQLiteAttribution()
transaction, err := connection.Begin()
if err != nil {
t.Fatal(err)
}
insert, err := connection.Prepare("INSERT INTO settings (value) VALUES (?)")
if err != nil {
t.Fatal(err)
}
if _, err := insert.Exec([]driver.Value{"retained-descriptor"}); err != nil {
t.Fatal(err)
}
if err := insert.Close(); err != nil {
t.Fatal(err)
}
_, descriptor, active := sqliteAttributionTransactionState(t, connection)
if !active {
t.Fatal("active transaction is missing before Commit")
}
// When
commitErr := transaction.Commit()
_, descriptorErr := unix.FcntlInt(uintptr(descriptor), unix.F_GETFD, 0)
// Then
if commitErr != nil {
t.Fatal(commitErr)
}
if !errors.Is(descriptorErr, unix.EBADF) {
t.Fatalf("retained journal descriptor remains open: %v", descriptorErr)
}
}
func TestSQLiteAttributionQueryRowReturningPoisonsTransaction(t *testing.T) {
// Given
database := openSQLiteAttributionTestDatabase(t)
enableSQLiteAttribution()
transaction, err := database.BeginTx(context.Background(), nil)
if err != nil {
t.Fatal(err)
}
// When
var identifier int64
scanErr := transaction.QueryRowContext(context.Background(), "INSERT INTO settings (value) VALUES (?) RETURNING id", "query-row").Scan(&identifier)
commitErr := transaction.Commit()
// Then
if scanErr != nil {
t.Fatal(scanErr)
}
if commitErr == nil {
t.Fatal("QueryRowContext RETURNING committed without reaching EOF")
}
if sqliteAttributionTrackerHasWrite() {
t.Fatal("QueryRowContext RETURNING published evidence without EOF")
}
}
service/singleton/sqlite_attribution_query_lifecycle_agentcompat_linux_test.go
+145
View File
@@ -0,0 +1,145 @@
//go:build agentcompat && linux
package singleton
import (
"context"
"database/sql/driver"
"errors"
"testing"
)
type sqliteAttributionLegacyQueryStmt struct {
closeCount int
}
func (statement *sqliteAttributionLegacyQueryStmt) Close() error { statement.closeCount++; return nil }
func (sqliteAttributionLegacyQueryStmt) NumInput() int { return -1 }
func (sqliteAttributionLegacyQueryStmt) Exec([]driver.Value) (driver.Result, error) {
return nil, driver.ErrSkip
}
func (sqliteAttributionLegacyQueryStmt) Query([]driver.Value) (driver.Rows, error) {
return nil, driver.ErrSkip
}
type sqliteAttributionCountingStmt struct {
driver.Stmt
closeCount int
}
func (statement *sqliteAttributionCountingStmt) Close() error {
statement.closeCount++
return statement.Stmt.Close()
}
type sqliteAttributionQueryProbeStmt struct {
closeCount int
queryCount int
}
func (statement *sqliteAttributionQueryProbeStmt) Close() error { statement.closeCount++; return nil }
func (sqliteAttributionQueryProbeStmt) NumInput() int { return -1 }
func (sqliteAttributionQueryProbeStmt) Exec([]driver.Value) (driver.Result, error) {
return nil, driver.ErrSkip
}
func (statement *sqliteAttributionQueryProbeStmt) Query([]driver.Value) (driver.Rows, error) {
statement.queryCount++
return nil, driver.ErrSkip
}
func TestSQLiteAttributionDirectQueryClosesOwnedStatementWhenPreStepRejects(t *testing.T) {
// Given
resetSQLiteAttributionForTest()
rawConnection, err := sqliteAttributionDriver{}.Open(sqliteAttributionTestDatabasePath(t))
if err != nil {
t.Fatal(err)
}
connection := rawConnection.(*sqliteAttributionConnection)
t.Cleanup(func() {
if closeErr := connection.Close(); closeErr != nil {
t.Error(closeErr)
}
})
if _, err := connection.connection.Exec("CREATE TABLE settings (id INTEGER PRIMARY KEY, value TEXT)", nil); err != nil {
t.Fatal(err)
}
prepared, err := connection.Prepare("INSERT INTO settings (value) VALUES (?) RETURNING id")
if err != nil {
t.Fatal(err)
}
statement := prepared.(*sqliteAttributionStatement)
owner := &sqliteAttributionCountingStmt{Stmt: prepared}
enableSQLiteAttribution()
// When
rows, queryErr := statement.queryOwned(func() (driver.Rows, error) {
return statement.statement.Query([]driver.Value{"rejected"})
}, owner)
// Then
if rows != nil {
t.Fatal("rejected direct Query returned rows")
}
if !errors.Is(queryErr, ErrSQLiteAttributionUnboundWrite) {
t.Fatalf("rejected direct Query error = %v", queryErr)
}
if owner.closeCount != 1 {
t.Fatalf("owned statement Close count = %d, want 1", owner.closeCount)
}
}
func TestSQLiteAttributionDirectQueryRejectsUnboundWriteThroughPublicConnection(t *testing.T) {
resetSQLiteAttributionForTest()
probe := &sqliteAttributionQueryProbeStmt{}
var connection *sqliteAttributionConnection
connection = &sqliteAttributionConnection{
prepareStatement: func(context.Context, string) (driver.Stmt, error) {
return &sqliteAttributionStatement{
connection: connection,
statement: probe,
classification: sqliteAttributionClassification{
hasRowDML: true,
operation: SQLiteOperationInsert,
table: "settings",
},
}, nil
},
}
enableSQLiteAttribution()
rows, queryErr := connection.Query("INSERT INTO settings (value) VALUES (?) RETURNING id", []driver.Value{"public-rejected"})
if rows != nil {
t.Fatal("public direct Query returned rows after pre-step rejection")
}
if !errors.Is(queryErr, ErrSQLiteAttributionUnboundWrite) {
t.Fatalf("public direct Query error = %v", queryErr)
}
if probe.closeCount != 1 {
t.Fatalf("public direct Query statement Close count = %d, want 1", probe.closeCount)
}
if probe.queryCount != 0 {
t.Fatalf("public direct Query invoked underlying Query %d times, want 0", probe.queryCount)
}
}
func TestSQLiteAttributionQueryContextClosesUnsupportedStatementBeforeErrSkip(t *testing.T) {
// Given
statement := &sqliteAttributionLegacyQueryStmt{}
connection := &sqliteAttributionConnection{}
wrapped := &sqliteAttributionStatement{connection: connection, statement: statement}
// When
rows, queryErr := connection.queryContextStatement(context.Background(), wrapped, nil, statement)
// Then
if rows != nil {
t.Fatal("unsupported QueryContext returned rows")
}
if !errors.Is(queryErr, driver.ErrSkip) {
t.Fatalf("unsupported QueryContext error = %v", queryErr)
}
if statement.closeCount != 1 {
t.Fatalf("unsupported QueryContext Close count = %d, want 1", statement.closeCount)
}
}
service/singleton/sqlite_attribution_release_terminal_ownership_agentcompat_linux_test.go
+227
View File
@@ -0,0 +1,227 @@
//go:build agentcompat && linux
package singleton
import (
"context"
"database/sql"
"database/sql/driver"
"errors"
"testing"
"golang.org/x/sys/unix"
)
func sqliteAttributionReleasedCommitFixture(t *testing.T, value string) (*sqliteAttributionConnection, *sqliteAttributionTx, SQLiteHoldSession, string, SQLiteTransaction, int) {
t.Helper()
resetSQLiteAttributionForTest()
databasePath := sqliteAttributionTestDatabasePath(t)
rawConnection, err := sqliteAttributionDriver{}.Open(databasePath)
if err != nil {
t.Fatal(err)
}
connection := rawConnection.(*sqliteAttributionConnection)
if _, err := connection.connection.Exec("CREATE TABLE settings (id INTEGER PRIMARY KEY, value TEXT)", nil); err != nil {
t.Fatal(err)
}
enableSQLiteAttribution()
transaction, err := connection.BeginTx(context.Background(), driver.TxOptions{})
if err != nil {
t.Fatal(err)
}
statement, err := connection.Prepare("INSERT INTO settings (value) VALUES (?)")
if err != nil {
t.Fatal(err)
}
if _, err := statement.Exec([]driver.Value{value}); err != nil {
t.Fatal(err)
}
if err := statement.Close(); err != nil {
t.Fatal(err)
}
session, err := sqliteAttributionTracker.Load().ArmNextSQLiteHold()
if err != nil {
t.Fatal(err)
}
wrapped := transaction.(*sqliteAttributionTx)
identity, descriptor, active := sqliteAttributionTransactionState(t, connection)
if !active {
t.Fatal("selected transaction is not active")
}
return connection, wrapped, session, databasePath, identity, descriptor
}
func sqliteAttributionReleasedTerminalSnapshot(t *testing.T, session SQLiteHoldSession) SQLiteHoldSnapshot {
t.Helper()
terminal, err := sqliteAttributionTracker.Load().WaitSQLiteHold(context.Background(), session, SQLiteHoldWaitFinalizing)
if err != nil {
t.Fatal(err)
}
return terminal
}
func TestSQLiteAttributionReleasedCommitOwnsTerminalAgainstConcurrentRollback(t *testing.T) {
// Given
connection, transaction, session, databasePath, identity, descriptor := sqliteAttributionReleasedCommitFixture(t, "released-rollback-owner")
t.Cleanup(func() {
if err := connection.Close(); err != nil && !errors.Is(err, driver.ErrBadConn) {
t.Error(err)
}
})
releasedBoundary := make(chan struct{})
allowCommitClaim := make(chan struct{})
loserWaiting := make(chan struct{})
transaction.state.releasedCommitBoundary = func() {
close(releasedBoundary)
<-allowCommitClaim
}
transaction.state.terminalLoserWaitBoundary = func() { close(loserWaiting) }
finalizing := make(chan error, 1)
commitResult := make(chan error, 1)
go func() {
_, err := sqliteAttributionTracker.Load().WaitSQLiteHold(context.Background(), session, SQLiteHoldWaitFinalizing)
finalizing <- err
}()
go func() { commitResult <- transaction.Commit() }()
if err := <-finalizing; err != nil {
t.Fatal(err)
}
// When
if err := sqliteAttributionTracker.Load().ReleaseSQLiteHold(session); err != nil {
t.Fatal(err)
}
<-releasedBoundary
rollbackResult := make(chan error, 1)
go func() { rollbackResult <- transaction.Rollback() }()
loserReturned := false
var rollbackErr error
select {
case <-loserWaiting:
case rollbackErr = <-rollbackResult:
loserReturned = true
}
close(allowCommitClaim)
commitErr := <-commitResult
if !loserReturned {
rollbackErr = <-rollbackResult
}
terminal := sqliteAttributionReleasedTerminalSnapshot(t, session)
_, descriptorErr := unix.FcntlInt(uintptr(descriptor), unix.F_GETFD, 0)
count := sqliteAttributionPersistedCount(t, databasePath)
// Then
if loserReturned {
t.Errorf("released Commit lost terminal ownership: Rollback returned %v before waiting", rollbackErr)
}
if commitErr != nil {
t.Errorf("released Commit error = %v, want nil", commitErr)
}
if !errors.Is(rollbackErr, driver.ErrBadConn) {
t.Errorf("Rollback after successful Release error = %v, want driver.ErrBadConn", rollbackErr)
}
if !terminal.Selected || !terminal.Finalizing || !terminal.Released {
t.Errorf("released terminal snapshot = %+v", terminal)
}
if !errors.Is(descriptorErr, unix.EBADF) {
t.Errorf("journal descriptor after released Commit = %v, want EBADF", descriptorErr)
}
if count != 1 {
t.Errorf("released Commit persisted %d rows, want 1", count)
}
if _, _, active := sqliteAttributionTransactionState(t, connection); active {
t.Error("released Commit left the connection transaction active")
}
if sqliteAttributionTrackerTransactionActive(sqliteAttributionTracker.Load(), identity) {
t.Error("released Commit left the tracker transaction active")
}
}
func TestSQLiteAttributionReleasedCommitOwnsTerminalAgainstConcurrentConnectionClose(t *testing.T) {
// Given
connection, transaction, session, databasePath, identity, descriptor := sqliteAttributionReleasedCommitFixture(t, "released-close-owner")
connectionClosed := false
t.Cleanup(func() {
if !connectionClosed {
if err := connection.Close(); err != nil && !errors.Is(err, driver.ErrBadConn) {
t.Error(err)
}
}
})
releasedBoundary := make(chan struct{})
allowCommitClaim := make(chan struct{})
loserWaiting := make(chan struct{})
transaction.state.releasedCommitBoundary = func() {
close(releasedBoundary)
<-allowCommitClaim
}
transaction.state.terminalLoserWaitBoundary = func() { close(loserWaiting) }
finalizing := make(chan error, 1)
commitResult := make(chan error, 1)
go func() {
_, err := sqliteAttributionTracker.Load().WaitSQLiteHold(context.Background(), session, SQLiteHoldWaitFinalizing)
finalizing <- err
}()
go func() { commitResult <- transaction.Commit() }()
if err := <-finalizing; err != nil {
t.Fatal(err)
}
// When
if err := sqliteAttributionTracker.Load().ReleaseSQLiteHold(session); err != nil {
t.Fatal(err)
}
<-releasedBoundary
closeResult := make(chan error, 1)
go func() { closeResult <- connection.Close() }()
loserReturned := false
var closeErr error
select {
case <-loserWaiting:
case closeErr = <-closeResult:
loserReturned = true
}
close(allowCommitClaim)
commitErr := <-commitResult
if !loserReturned {
closeErr = <-closeResult
}
connectionClosed = true
terminal := sqliteAttributionReleasedTerminalSnapshot(t, session)
_, descriptorErr := unix.FcntlInt(uintptr(descriptor), unix.F_GETFD, 0)
database, err := sql.Open("sqlite3", databasePath)
if err != nil {
t.Fatal(err)
}
defer database.Close()
var count int
if err := database.QueryRow("SELECT COUNT(*) FROM settings").Scan(&count); err != nil {
t.Fatal(err)
}
// Then
if loserReturned {
t.Errorf("released Commit lost terminal ownership: connection Close returned %v before waiting", closeErr)
}
if commitErr != nil {
t.Errorf("released Commit error = %v, want nil", commitErr)
}
if closeErr != nil && !errors.Is(closeErr, driver.ErrBadConn) {
t.Errorf("connection Close after successful Release error = %v, want nil or driver.ErrBadConn", closeErr)
}
if !terminal.Selected || !terminal.Finalizing || !terminal.Released {
t.Errorf("released terminal snapshot = %+v", terminal)
}
if !errors.Is(descriptorErr, unix.EBADF) {
t.Errorf("journal descriptor after released Commit = %v, want EBADF", descriptorErr)
}
if count != 1 {
t.Errorf("released Commit persisted %d rows, want 1", count)
}
if _, _, active := sqliteAttributionTransactionState(t, connection); active {
t.Error("released Commit left the connection transaction active")
}
if sqliteAttributionTrackerTransactionActive(sqliteAttributionTracker.Load(), identity) {
t.Error("released Commit left the tracker transaction active")
}
}

Some files were not shown because too many files have changed in this diff Show More