feat: server transfer rotation

This commit is contained in:
naiba
2026-05-25 10:17:34 +00:00
parent 636f4a9716
commit e816e94fca
43 changed files with 7072 additions and 134 deletions
+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)
}
}
+170 -12
View File
@@ -3,6 +3,7 @@ package rpc
import (
"context"
"fmt"
"log"
"strings"
petname "github.com/dustinkirkland/golang-petname"
@@ -21,6 +22,18 @@ 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 失败")
@@ -37,6 +50,107 @@ func (a *authHandler) Check(ctx context.Context) (uint64, error) {
ip, _ := ctx.Value(model.CtxKeyRealIP{}).(string)
var clientUUID string
if value, ok := md["client_uuid"]; ok {
clientUUID = value[0]
}
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 {
@@ -48,15 +162,6 @@ 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]
}
if _, err := uuid.ParseUUID(clientUUID); err != nil {
return 0, status.Error(codes.Unauthenticated, "客户端 UUID 不合法")
}
clientID, hasID, err := authorizeAgentForUUID(userId, clientUUID)
if err != nil {
return 0, status.Error(codes.Unauthenticated, err.Error())
@@ -90,6 +195,16 @@ func (a *authHandler) Check(ctx context.Context) (uint64, error) {
// 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 {
@@ -106,8 +221,51 @@ func authorizeAgentForUUID(userId uint64, clientUUID string) (clientID uint64, h
// agent secrets, so keep it compatible by allowing any existing UUID.
return cid, true, nil
}
if server.UserID != userId {
return 0, false, fmt.Errorf("client UUID does not belong to the agent secret owner")
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
}
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")
}
+506 -1
View File
@@ -1,15 +1,93 @@
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)
}
// 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()
}
}
// 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.
@@ -17,12 +95,13 @@ 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{}); err != nil {
if err := db.AutoMigrate(&model.Server{}, &model.ServerTransfer{}); err != nil {
t.Fatalf("migrate: %v", err)
}
if err := db.Create(&model.Server{
@@ -41,10 +120,15 @@ func setupAuthAgentFixture(t *testing.T) func() {
}
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
}
}
@@ -99,3 +183,424 @@ func TestAuthorizeAgentForUUIDPermitsUnknownUUIDForRegistration(t *testing.T) {
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")
}
}
+31
View File
@@ -119,6 +119,35 @@ func (s *NezhaHandler) GetStream(streamId string) (*ioStreamContext, error) {
return nil, errors.New("stream not found")
}
// RevokeStreamsForServer tears down every IOStream whose targetServerID
// matches serverID. Called by the singleton package via the
// ServerTransferStreamRevocationHook on every transfer ownership
// transition — a stream the previous owner had open against this server
// must not survive into the new tenant, otherwise terminal/file-manager/NAT
// sessions become post-transfer hijack channels (effectively RCE).
//
// Underlying IO pipes are closed inline so the dashboard websocket loop
// sees EOF immediately rather than at the next idle-timeout.
func (s *NezhaHandler) RevokeStreamsForServer(serverID uint64) {
if serverID == 0 {
return
}
s.ioStreamMutex.Lock()
defer s.ioStreamMutex.Unlock()
for streamId, ctx := range s.ioStreams {
if ctx.targetServerID != serverID {
continue
}
if ctx.userIo != nil {
ctx.userIo.Close()
}
if ctx.agentIo != nil {
ctx.agentIo.Close()
}
delete(s.ioStreams, streamId)
}
}
func (s *NezhaHandler) CloseStream(streamId string) error {
s.ioStreamMutex.Lock()
defer s.ioStreamMutex.Unlock()
@@ -136,6 +165,8 @@ func (s *NezhaHandler) CloseStream(streamId string) error {
return nil
}
func (s *NezhaHandler) UserConnected(streamId string, userIo io.ReadWriteCloser) error {
stream, err := s.GetStream(streamId)
if err != nil {
+32 -2
View File
@@ -40,12 +40,21 @@ func NewNezhaHandler() *NezhaHandler {
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.SetTaskStream(stream)
defer server.ClearTaskStreamIfCurrent(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()
@@ -83,6 +92,27 @@ 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.IsServiceSentinelNeeded(result.GetType()) {
singleton.ServiceSentinelShared.Dispatch(singleton.ReportData{
+61 -7
View File
@@ -18,6 +18,7 @@ import (
type requestTaskSecurityStream struct {
ctx context.Context
results []*pb.TaskResult
onRecv func()
onSend func(*pb.Task)
sendErr error
}
@@ -31,6 +32,9 @@ func (s *requestTaskSecurityStream) Send(task *pb.Task) error {
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]
@@ -201,6 +205,52 @@ func TestRequestTaskSkipsAlertTriggerCronResultAfterSendFailure(t *testing.T) {
}
}
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()
@@ -305,22 +355,26 @@ func connectRequestTaskSecurityTaskStreamWithSendHook(t *testing.T, serverID uin
if !ok {
t.Fatalf("server %d not found", serverID)
}
server.TaskStream = &requestTaskSecurityStream{ctx: context.Background(), sendErr: sendErr, onSend: onSend}
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 := &requestTaskSecurityStream{
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,
)),
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)
}
}
+1 -1
View File
@@ -149,7 +149,7 @@ 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.
+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)
+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)
}
}
+6 -5
View File
@@ -264,12 +264,13 @@ func CronTrigger(cr *model.Cron, triggerServer ...uint64) func() {
if !cronCanSendToServer(cr, s) {
return
}
if s.TaskStream != nil {
stream := s.GetTaskStream()
if stream != nil {
cronShared := CronShared
if cronShared != nil {
cronShared.reserveAlertTriggerCronResult(cr.ID, s.ID)
}
if err := s.TaskStream.Send(&pb.Task{
if err := stream.Send(&pb.Task{
Id: cr.ID,
Data: cr.Command,
Type: model.TaskTypeCommand,
@@ -296,8 +297,8 @@ func CronTrigger(cr *model.Cron, triggerServer ...uint64) func() {
if cr.Cover == model.CronCoverIgnoreAll && !crIgnoreMap[s.ID] {
continue
}
if s.TaskStream != nil {
s.TaskStream.Send(&pb.Task{
if stream := s.GetTaskStream(); stream != nil {
stream.Send(&pb.Task{
Id: cr.ID,
Data: cr.Command,
Type: model.TaskTypeCommand,
@@ -313,7 +314,7 @@ func CronTrigger(cr *model.Cron, triggerServer ...uint64) func() {
}
func cronCanSendToServer(cr *model.Cron, server *model.Server) bool {
return cr.UserID == server.UserID || userIsAdmin(cr.UserID)
return cr.UserID == server.GetUserID() || userIsAdmin(cr.UserID)
}
func userIsAdmin(userID uint64) bool {
+29 -13
View File
@@ -39,6 +39,16 @@ func (s *capturedTaskStream) Context() context.Context { return context.Bac
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()
@@ -75,8 +85,8 @@ func TestCronTriggerSkipsServersOwnedByOtherUsers(t *testing.T) {
firstStream := newCapturedTaskStream()
secondStream := newCapturedTaskStream()
replaceServerSharedForSecurityTest(t,
&model.Server{Common: model.Common{ID: 1, UserID: 100}, Name: "member-server", TaskStream: firstStream},
&model.Server{Common: model.Common{ID: 2, UserID: 200}, Name: "admin-server", TaskStream: secondStream},
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{
@@ -95,7 +105,7 @@ func TestCronTriggerSkipsServersOwnedByOtherUsers(t *testing.T) {
func TestSendTriggerTasksSkipsCronOwnedByAnotherUser(t *testing.T) {
attackerStream := newCapturedTaskStream()
replaceServerSharedForSecurityTest(t,
&model.Server{Common: model.Common{ID: 7, UserID: 200}, Name: "attacker-server", TaskStream: attackerStream},
withTaskStream(&model.Server{Common: model.Common{ID: 7, UserID: 200}, Name: "attacker-server"}, attackerStream),
)
replaceUserInfoMapForSecurityTest(t, map[uint64]model.UserInfo{
1: {Role: model.RoleAdmin},
@@ -147,7 +157,7 @@ func assertNoTask(t *testing.T, stream *capturedTaskStream) {
func TestCronTriggerSendsToMemberOwnedServer(t *testing.T) {
memberStream := newCapturedTaskStream()
replaceServerSharedForSecurityTest(t,
&model.Server{Common: model.Common{ID: 1, UserID: 100}, Name: "member-server", TaskStream: memberStream},
withTaskStream(&model.Server{Common: model.Common{ID: 1, UserID: 100}, Name: "member-server"}, memberStream),
)
replaceUserInfoMapForSecurityTest(t, map[uint64]model.UserInfo{
100: {Role: model.RoleMember},
@@ -168,8 +178,8 @@ func TestCronTriggerAdminCronFansOutAcrossOwners(t *testing.T) {
first := newCapturedTaskStream()
second := newCapturedTaskStream()
replaceServerSharedForSecurityTest(t,
&model.Server{Common: model.Common{ID: 1, UserID: 100}, Name: "member-server", TaskStream: first},
&model.Server{Common: model.Common{ID: 2, UserID: 200}, Name: "admin-server", TaskStream: second},
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},
@@ -192,7 +202,7 @@ func TestCronTriggerAdminCronFansOutAcrossOwners(t *testing.T) {
func TestCronTriggerLegacyZeroOwnerFansOut(t *testing.T) {
first := newCapturedTaskStream()
replaceServerSharedForSecurityTest(t,
&model.Server{Common: model.Common{ID: 1, UserID: 100}, Name: "member-server", TaskStream: first},
withTaskStream(&model.Server{Common: model.Common{ID: 1, UserID: 100}, Name: "member-server"}, first),
)
replaceUserInfoMapForSecurityTest(t, map[uint64]model.UserInfo{
100: {Role: model.RoleMember},
@@ -212,7 +222,7 @@ func TestCronTriggerLegacyZeroOwnerFansOut(t *testing.T) {
func TestCronTriggerSkipsServersWhenOwnerNotKnown(t *testing.T) {
stream := newCapturedTaskStream()
replaceServerSharedForSecurityTest(t,
&model.Server{Common: model.Common{ID: 1, UserID: 100}, Name: "member-server", TaskStream: stream},
withTaskStream(&model.Server{Common: model.Common{ID: 1, UserID: 100}, Name: "member-server"}, stream),
)
replaceUserInfoMapForSecurityTest(t, map[uint64]model.UserInfo{
100: {Role: model.RoleMember},
@@ -232,7 +242,7 @@ func TestCronTriggerSkipsServersWhenOwnerNotKnown(t *testing.T) {
func TestSendTriggerTasksAllowsSelfOwnedCron(t *testing.T) {
stream := newCapturedTaskStream()
replaceServerSharedForSecurityTest(t,
&model.Server{Common: model.Common{ID: 7, UserID: 200}, Name: "member-server", TaskStream: stream},
withTaskStream(&model.Server{Common: model.Common{ID: 7, UserID: 200}, Name: "member-server"}, stream),
)
replaceUserInfoMapForSecurityTest(t, map[uint64]model.UserInfo{
200: {Role: model.RoleMember},
@@ -257,7 +267,7 @@ func TestSendTriggerTasksAllowsSelfOwnedCron(t *testing.T) {
func TestSendTriggerTasksAllowsAdminCallerToTriggerAny(t *testing.T) {
stream := newCapturedTaskStream()
replaceServerSharedForSecurityTest(t,
&model.Server{Common: model.Common{ID: 9, UserID: 100}, Name: "any-server", TaskStream: stream},
withTaskStream(&model.Server{Common: model.Common{ID: 9, UserID: 100}, Name: "any-server"}, stream),
)
replaceUserInfoMapForSecurityTest(t, map[uint64]model.UserInfo{
1: {Role: model.RoleAdmin},
@@ -283,7 +293,7 @@ func TestSendTriggerTasksAllowsAdminCallerToTriggerAny(t *testing.T) {
func TestSendTriggerTasksIgnoresUnknownTaskIDs(t *testing.T) {
stream := newCapturedTaskStream()
replaceServerSharedForSecurityTest(t,
&model.Server{Common: model.Common{ID: 7, UserID: 200}, Name: "member-server", TaskStream: stream},
withTaskStream(&model.Server{Common: model.Common{ID: 7, UserID: 200}, Name: "member-server"}, stream),
)
replaceUserInfoMapForSecurityTest(t, map[uint64]model.UserInfo{
200: {Role: model.RoleMember},
@@ -304,7 +314,7 @@ func TestSendTriggerTasksIgnoresUnknownTaskIDs(t *testing.T) {
func TestSendTriggerTasksMixedCronIDsOnlyFiresAllowed(t *testing.T) {
stream := newCapturedTaskStream()
replaceServerSharedForSecurityTest(t,
&model.Server{Common: model.Common{ID: 7, UserID: 200}, Name: "member-server", TaskStream: stream},
withTaskStream(&model.Server{Common: model.Common{ID: 7, UserID: 200}, Name: "member-server"}, stream),
)
replaceUserInfoMapForSecurityTest(t, map[uint64]model.UserInfo{
1: {Role: model.RoleAdmin},
@@ -540,7 +550,7 @@ func (s *failingTaskStream) Send(task *pb.Task) error {
func TestCronTriggerRevokesAlertTriggerAuthorizationOnSendFailure(t *testing.T) {
failing := newFailingTaskStream(context.Canceled)
replaceServerSharedForSecurityTest(t,
&model.Server{Common: model.Common{ID: 7, UserID: 100}, Name: "broken-server", TaskStream: failing},
withTaskStream(&model.Server{Common: model.Common{ID: 7, UserID: 100}, Name: "broken-server"}, failing),
)
replaceUserInfoMapForSecurityTest(t, map[uint64]model.UserInfo{
100: {Role: model.RoleMember},
@@ -828,6 +838,12 @@ func newServiceMonitorSecurityHarness(t *testing.T, servers ...*model.Server) *S
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
}
File diff suppressed because it is too large Load Diff
File diff suppressed because it is too large Load Diff
+40 -2
View File
@@ -85,6 +85,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 +123,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)
@@ -489,10 +503,34 @@ func canReportServiceResult(service *model.Service, reporter *model.Server, task
return false
}
return service.UserID == reporter.UserID || userIsAdmin(service.UserID)
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 {
+6 -1
View File
@@ -34,6 +34,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 +63,7 @@ func LoadSingleton(bus chan<- *model.Service) (err error) {
NotificationShared = NewNotificationClass()
ServerShared = NewServerClass()
CronShared = NewCronClass()
ServerTransferShared = NewServerTransferClass()
// 最后初始化 ServiceSentinel
ServiceSentinelShared, err = NewServiceSentinel(bus)
return
@@ -89,7 +94,7 @@ 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.WAF{}, model.Oauth2Bind{}, model.ServerTransfer{})
if err != nil {
return err
}
+34
View File
@@ -40,10 +40,34 @@ func initUser() {
UserInfoMap[u.ID] = model.UserInfo{
Role: u.Role,
Username: u.Username,
AgentSecret: u.AgentSecret,
}
AgentSecretToUserId[u.AgentSecret] = u.ID
}
model.ServerOwnerLookup = lookupServerOwner
}
// lookupServerOwner resolves Server.UserID into a display-ready owner
// record for model.Server.MarshalJSON. uid=0 is the legacy global agent
// secret (a pseudo-owner with no User row) and intentionally returns
// ok=false with no username; the frontend renders it as "Global". Other
// uids return ok=false when the user has been deleted, so the JSON still
// carries the bare id and the frontend can render an "Unknown (#<uid>)"
// placeholder. RLock is required because OnUserUpdate / OnUserDelete may
// mutate UserInfoMap concurrently with serialization.
func lookupServerOwner(uid uint64) (model.ServerOwnerInfo, bool) {
if uid == 0 {
return model.ServerOwnerInfo{}, false
}
UserLock.RLock()
info, ok := UserInfoMap[uid]
UserLock.RUnlock()
if !ok {
return model.ServerOwnerInfo{}, false
}
return model.ServerOwnerInfo{ID: uid, Username: info.Username}, true
}
func OnUserUpdate(u *model.User) {
@@ -56,6 +80,7 @@ func OnUserUpdate(u *model.User) {
UserInfoMap[u.ID] = model.UserInfo{
Role: u.Role,
Username: u.Username,
AgentSecret: u.AgentSecret,
}
AgentSecretToUserId[u.AgentSecret] = u.ID
@@ -69,6 +94,10 @@ func OnUserDelete(id []uint64, errorFunc func(string, ...any) error) error {
return Localizer.ErrorT("user id not specified")
}
if ServerTransferShared != nil {
ServerTransferShared.OnUsersDeleted(id)
}
var (
cron, server bool
crons, servers []uint64
@@ -127,6 +156,11 @@ func OnUserDelete(id []uint64, errorFunc func(string, ...any) error) error {
}
}
AlertsLock.Unlock()
// Cancel pending transfers before ServerShared drops the
// in-memory entry: same ordering rationale as batchDeleteServer.
if ServerTransferShared != nil {
ServerTransferShared.OnServersDeleted(servers)
}
ServerShared.Delete(servers)
}
+90
View File
@@ -0,0 +1,90 @@
package singleton
import (
"testing"
"github.com/stretchr/testify/require"
"github.com/nezhahq/nezha/model"
"github.com/nezhahq/nezha/pkg/i18n"
)
func setupOnUserDeleteFixture(t *testing.T) (*ServerTransferClass, func()) {
t.Helper()
c, transferCleanup := setupTransferFixture(t)
require.NoError(t, DB.AutoMigrate(&model.Cron{}, &model.Transfer{}, &model.ServerGroupServer{}))
originalCronShared := CronShared
CronShared = &CronClass{
class: class[uint64, *model.Cron]{list: map[uint64]*model.Cron{}},
}
originalLocalizer := Localizer
Localizer = i18n.NewLocalizer("zh_CN", domain, "translations", i18n.Translations)
cleanup := func() {
Localizer = originalLocalizer
CronShared = originalCronShared
transferCleanup()
}
return c, cleanup
}
func TestOnUserDeleteCancelsPendingTransfersAwayFromDeletedUser(t *testing.T) {
c, cleanup := setupOnUserDeleteFixture(t)
defer cleanup()
const fromUser = uint64(100)
const toUser = uint64(200)
const serverID = uint64(1)
seedServerForTransfer(t, serverID, fromUser)
require.NoError(t, DB.AutoMigrate(&model.User{}))
require.NoError(t, DB.Create(&model.User{
Common: model.Common{ID: fromUser},
Username: "alice",
AgentSecret: "alice-secret",
}).Error)
require.NoError(t, DB.Create(&model.User{
Common: model.Common{ID: toUser},
Username: "bob",
AgentSecret: "bob-secret",
}).Error)
UserLock.Lock()
UserInfoMap[fromUser] = model.UserInfo{Role: model.RoleMember, AgentSecret: "alice-secret"}
UserInfoMap[toUser] = model.UserInfo{Role: model.RoleMember, AgentSecret: "bob-secret"}
UserLock.Unlock()
tr := initiateAndRegister(t, c, serverID, fromUser, toUser, fromUser)
require.True(t, c.HasPending(serverID), "precondition: pending transfer published")
srv, ok := ServerShared.Get(serverID)
require.True(t, ok)
require.Equal(t, toUser, srv.GetUserID(), "precondition: pending transfer flipped owner to ToUserID")
require.NoError(t, OnUserDelete([]uint64{fromUser}, func(format string, args ...any) error {
return nil
}))
if c.HasPending(serverID) {
t.Fatal("OnUserDelete on the transfer FromUserID must terminate the pending transfer so a later Cancel/Fail/Timeout cannot revert ownership to the deleted user")
}
if srv, ok := ServerShared.Get(serverID); ok {
require.NotEqual(t, fromUser, srv.GetUserID(),
"server owner must not be reverted to the deleted FromUserID; got owner=%d", srv.GetUserID())
}
if _, err := c.Cancel(tr.ID); err == nil {
var refreshed model.ServerTransfer
if err := DB.First(&refreshed, tr.ID).Error; err == nil {
require.NotEqual(t, model.ServerTransferStatusPending, refreshed.Status,
"after OnUserDelete a subsequent Cancel must not leave the transfer Pending")
if srv, ok := ServerShared.Get(serverID); ok {
require.NotEqual(t, fromUser, srv.GetUserID(),
"a late Cancel against the terminated transfer must not revert server.UserID to the deleted FromUserID; got owner=%d", srv.GetUserID())
}
}
}
}