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
+55 -21
View File
@@ -8,7 +8,6 @@ import (
"log"
"net/http"
"os"
"path"
"regexp"
"slices"
"strings"
@@ -118,6 +117,11 @@ func routers(r *gin.Engine, frontendDist fs.FS) {
auth.POST("/batch-move/server", commonHandler(batchMoveServer))
auth.POST("/force-update/server", commonHandler(forceUpdateServer))
auth.GET("/transfer", listHandler(listServerTransfer))
auth.POST("/transfer/:id/cancel", commonHandler(cancelServerTransfer))
auth.POST("/transfer/:id/retry", commonHandler(retryServerTransfer))
auth.GET("/ws/transfer", commonHandler(transferStream))
auth.GET("/notification", listHandler(listNotification))
auth.POST("/notification", commonHandler(createNotification))
auth.PATCH("/notification/:id", commonHandler(updateNotification))
@@ -320,27 +324,52 @@ func getUid(c *gin.Context) uint64 {
}
func fallbackToFrontend(frontendDist fs.FS) func(*gin.Context) {
checkLocalFileOrFs := func(c *gin.Context, fs fs.FS, path string, customStatusCode int) bool {
if _, err := os.Stat(path); err == nil {
http.ServeFile(utils.NewGinCustomWriter(c, customStatusCode), c.Request, path)
return true
}
f, err := fs.Open(path)
if err != nil {
return false
}
defer f.Close()
fileStat, err := f.Stat()
serveFile := func(c *gin.Context, name string, file fs.File, customStatusCode int) bool {
defer file.Close()
fileStat, err := file.Stat()
if err != nil {
return false
}
if fileStat.IsDir() {
return false
}
http.ServeContent(utils.NewGinCustomWriter(c, customStatusCode), c.Request, path, fileStat.ModTime(), f.(io.ReadSeeker))
readSeeker, ok := file.(io.ReadSeeker)
if !ok {
return false
}
http.ServeContent(utils.NewGinCustomWriter(c, customStatusCode), c.Request, name, fileStat.ModTime(), readSeeker)
return true
}
checkLocalFileOrFs := func(c *gin.Context, frontendFS fs.FS, templateRoot, filePath string, customStatusCode int) bool {
if filePath != "" {
localRoot, err := os.OpenRoot(templateRoot)
if err == nil {
defer localRoot.Close()
// URL paths must stay inside the selected template root; never join them against the process cwd.
if file, err := localRoot.Open(filePath); err == nil && serveFile(c, filePath, file, customStatusCode) {
return true
}
}
}
if !fs.ValidPath(filePath) {
return false
}
templateFS, err := fs.Sub(frontendFS, templateRoot)
if err != nil {
return false
}
file, err := templateFS.Open(filePath)
if err != nil {
return false
}
if serveFile(c, filePath, file, customStatusCode) {
return true
}
return false
}
frontendPageUrlRegistry := []*regexp.Regexp{
// official user frontend
regexp.MustCompile(`^/$`),
@@ -361,6 +390,11 @@ func fallbackToFrontend(frontendDist fs.FS) func(*gin.Context) {
regexp.MustCompile(`^/dashboard/settings/user$`),
regexp.MustCompile(`^/dashboard/settings/online-user$`),
regexp.MustCompile(`^/dashboard/settings/waf$`),
// 注意:这里的白名单决定哪些 URL 走 index.html fallback;漏一条就会把
// 直接刷新该页面变成 404(HTTP 状态码层面,body 仍是 index.html,所以
// 浏览器内 SPA 看起来正常,但 monitoring / 链接预览会以为站点挂了)。
// 新增前端路由时必须在 admin-frontend/src/main.tsx 与这里同步加。
regexp.MustCompile(`^/dashboard/transfer$`),
}
getFallbackStatusCode := func(path string) int {
@@ -385,22 +419,22 @@ func fallbackToFrontend(frontendDist fs.FS) func(*gin.Context) {
}
fallbackStatusCode := getFallbackStatusCode(c.Request.URL.Path)
if strings.HasPrefix(c.Request.URL.Path, "/dashboard") {
stripPath := strings.TrimPrefix(c.Request.URL.Path, "/dashboard")
localFilePath := path.Join(singleton.Conf.AdminTemplate, stripPath)
if checkLocalFileOrFs(c, frontendDist, localFilePath, http.StatusOK) {
// Only /dashboard/ belongs to the admin frontend; /dashboard.. must not be trimmed into ../.
if strings.HasPrefix(c.Request.URL.Path, "/dashboard/") {
stripPath := strings.TrimPrefix(c.Request.URL.Path, "/dashboard/")
if checkLocalFileOrFs(c, frontendDist, singleton.Conf.AdminTemplate, stripPath, http.StatusOK) {
return
}
if !checkLocalFileOrFs(c, frontendDist, singleton.Conf.AdminTemplate+"/index.html", fallbackStatusCode) {
if !checkLocalFileOrFs(c, frontendDist, singleton.Conf.AdminTemplate, "index.html", fallbackStatusCode) {
c.JSON(http.StatusNotFound, newErrorResponse(errors.New("404 Not Found")))
}
return
}
localFilePath := path.Join(singleton.Conf.UserTemplate, c.Request.URL.Path)
if checkLocalFileOrFs(c, frontendDist, localFilePath, http.StatusOK) {
stripPath := strings.TrimPrefix(c.Request.URL.Path, "/")
if checkLocalFileOrFs(c, frontendDist, singleton.Conf.UserTemplate, stripPath, http.StatusOK) {
return
}
if !checkLocalFileOrFs(c, frontendDist, singleton.Conf.UserTemplate+"/index.html", fallbackStatusCode) {
if !checkLocalFileOrFs(c, frontendDist, singleton.Conf.UserTemplate, "index.html", fallbackStatusCode) {
c.JSON(http.StatusNotFound, newErrorResponse(errors.New("404 Not Found")))
}
}
+6 -2
View File
@@ -33,7 +33,11 @@ func createFM(c *gin.Context) (*model.CreateFMResponse, error) {
}
server, _ := singleton.ServerShared.Get(id)
if server == nil || server.TaskStream == nil {
if server == nil {
return nil, singleton.Localizer.ErrorT("server not found or not connected")
}
stream := server.GetTaskStream()
if stream == nil {
return nil, singleton.Localizer.ErrorT("server not found or not connected")
}
@@ -51,7 +55,7 @@ func createFM(c *gin.Context) (*model.CreateFMResponse, error) {
fmData, _ := json.Marshal(&model.TaskFM{
StreamID: streamId,
})
if err := server.TaskStream.Send(&proto.Task{
if err := stream.Send(&proto.Task{
Type: model.TaskTypeFM,
Data: string(fmData),
}); err != nil {
@@ -56,7 +56,7 @@ func setupServerOwnershipFixture(t *testing.T) (stream *fakeTaskStream, reset fu
alice, _ := singleton.ServerShared.Get(1)
stream = &fakeTaskStream{}
alice.TaskStream = stream
alice.SetTaskStream(stream)
return stream, func() {
singleton.DB = originalDB
@@ -0,0 +1,133 @@
package controller
import (
"io/fs"
"net/http"
"net/http/httptest"
"os"
"path/filepath"
"strings"
"testing"
"github.com/gin-gonic/gin"
"github.com/nezhahq/nezha/model"
"github.com/nezhahq/nezha/service/singleton"
)
func newFrontendFallbackTestRouter(t *testing.T) *gin.Engine {
t.Helper()
gin.SetMode(gin.TestMode)
originalConf := singleton.Conf
singleton.Conf = &singleton.ConfigClass{Config: &model.Config{
ConfigDashboard: model.ConfigDashboard{
AdminTemplate: "admin-dist",
UserTemplate: "user-dist",
},
}}
t.Cleanup(func() { singleton.Conf = originalConf })
writeFrontendFallbackTestFile(t, "admin-dist/index.html", "<html>admin index</html>")
writeFrontendFallbackTestFile(t, "admin-dist/assets/app.js", "console.log('admin asset')")
writeFrontendFallbackTestFile(t, "user-dist/index.html", "<html>user index</html>")
writeFrontendFallbackTestFile(t, "data/config.yaml", "jwt_secret_key: traversal-secret")
r := gin.New()
r.NoRoute(fallbackToFrontend(testFrontendDist{}))
return r
}
func writeFrontendFallbackTestFile(t *testing.T, name, content string) {
t.Helper()
if err := os.MkdirAll(filepath.Dir(name), 0o755); err != nil {
t.Fatalf("create fixture directory: %v", err)
}
if err := os.WriteFile(name, []byte(content), 0o644); err != nil {
t.Fatalf("write fixture file: %v", err)
}
}
type testFrontendDist struct{}
func (testFrontendDist) Open(string) (fs.File, error) {
return nil, fs.ErrNotExist
}
func performFrontendFallbackRequest(t *testing.T, router *gin.Engine, target string) *httptest.ResponseRecorder {
t.Helper()
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, target, nil)
router.ServeHTTP(w, req)
return w
}
func TestFallbackToFrontendBlocksDashboardTraversal(t *testing.T) {
t.Chdir(t.TempDir())
router := newFrontendFallbackTestRouter(t)
tests := []string{
"/dashboard../data/config.yaml",
"/dashboard%2e%2e/data/config.yaml",
"/dashboard%2e%2e%2fdata%2fconfig.yaml",
"/dashboard/../data/config.yaml",
"/dashboard/%2e%2e/data/config.yaml",
"/dashboard../assets/app.js",
}
for _, target := range tests {
t.Run(target, func(t *testing.T) {
w := performFrontendFallbackRequest(t, router, target)
body := w.Body.String()
if strings.Contains(body, "traversal-secret") || strings.Contains(body, "jwt_secret_key") || strings.Contains(body, "admin asset") {
t.Fatalf("%s leaked protected content with status %d: %q", target, w.Code, body)
}
})
}
}
func TestFallbackToFrontendBlocksUserTraversal(t *testing.T) {
t.Chdir(t.TempDir())
router := newFrontendFallbackTestRouter(t)
tests := []string{
"/../data/config.yaml",
"/%2e%2e/data/config.yaml",
"/%2e%2e%2fdata%2fconfig.yaml",
"/../admin-dist/assets/app.js",
"/%2e%2e/admin-dist/assets/app.js",
}
for _, target := range tests {
t.Run(target, func(t *testing.T) {
w := performFrontendFallbackRequest(t, router, target)
body := w.Body.String()
if strings.Contains(body, "traversal-secret") || strings.Contains(body, "jwt_secret_key") || strings.Contains(body, "admin asset") {
t.Fatalf("%s leaked protected content with status %d: %q", target, w.Code, body)
}
})
}
}
func TestFallbackToFrontendPreservesDashboardRoutes(t *testing.T) {
t.Chdir(t.TempDir())
router := newFrontendFallbackTestRouter(t)
w := performFrontendFallbackRequest(t, router, "/dashboard")
if w.Code != http.StatusMovedPermanently {
t.Fatalf("/dashboard status = %d, want %d", w.Code, http.StatusMovedPermanently)
}
if location := w.Header().Get("Location"); location != "/dashboard/" {
t.Fatalf("/dashboard Location = %q, want /dashboard/", location)
}
w = performFrontendFallbackRequest(t, router, "/dashboard/")
if w.Code != http.StatusOK || !strings.Contains(w.Body.String(), "admin index") {
t.Fatalf("/dashboard/ status = %d body = %q, want admin index", w.Code, w.Body.String())
}
w = performFrontendFallbackRequest(t, router, "/dashboard/assets/app.js")
if w.Code != http.StatusOK || !strings.Contains(w.Body.String(), "admin asset") {
t.Fatalf("/dashboard/assets/app.js status = %d body = %q, want admin asset", w.Code, w.Body.String())
}
}
@@ -478,6 +478,21 @@ func TestBatchMoveServerAllowsAdminCrossUser(t *testing.T) {
assert.NoError(t, err)
}
func TestBatchMoveServerMasksForeignServerIDsForMembers(t *testing.T) {
ctx := newMemberValidationContext(t)
assert.NoError(t, singleton.DB.Create(&model.Server{Common: model.Common{ID: 1, UserID: 1}, Name: "foreign", UUID: "foreign-server"}).Error)
singleton.ServerShared = singleton.NewServerClass()
ctx.Request = httptest.NewRequest(http.MethodPost, "/batch-move/server", strings.NewReader(`{"ids":[1,9999],"to_user":200}`))
ctx.Request.Header.Set("Content-Type", "application/json")
got, err := batchMoveServer(ctx)
assert.NoError(t, err)
assert.Len(t, got, 2)
assert.Equal(t, model.BatchMoveServerResultServerNotFound, got[0].Status)
assert.Equal(t, model.BatchMoveServerResultServerNotFound, got[1].Status)
}
func TestNATRejectsUnknownServerID(t *testing.T) {
ctx := newMemberValidationContext(t)
ctx.Request = httptest.NewRequest(http.MethodPost, "/nat", strings.NewReader(`{"name":"x","domain":"x.example","host":"127.0.0.1:80","server_id":9999}`))
+105 -33
View File
@@ -1,6 +1,7 @@
package controller
import (
"errors"
"slices"
"strconv"
"sync"
@@ -153,6 +154,14 @@ func batchDeleteServer(c *gin.Context) (any, error) {
singleton.DB.Unscoped().Delete(&model.Transfer{}, "server_id in (?)", servers)
singleton.AlertsLock.Unlock()
// Cancel any in-flight transfers BEFORE the in-memory ServerShared
// entry is dropped: the order shortens the window in which a
// concurrent Retry/Register could install a fresh pending entry for
// the same serverID and have it wiped by the cleanup. The
// transferID-guarded delete inside OnServersDeleted is the
// authoritative protection against that race; the ordering here is
// belt and braces.
singleton.ServerTransferShared.OnServersDeleted(servers)
singleton.ServerShared.Delete(servers)
return nil, nil
}
@@ -187,8 +196,8 @@ func forceUpdateServer(c *gin.Context) (*model.ServerTaskResponse, error) {
forceUpdateResp.Offline = append(forceUpdateResp.Offline, sid)
continue
}
if server.TaskStream != nil {
if err := server.TaskStream.Send(&pb.Task{
if stream := server.GetTaskStream(); stream != nil {
if err := stream.Send(&pb.Task{
Type: model.TaskTypeUpgrade,
}); err != nil {
forceUpdateResp.Failure = append(forceUpdateResp.Failure, sid)
@@ -220,7 +229,11 @@ func getServerConfig(c *gin.Context) (string, error) {
}
s, ok := singleton.ServerShared.Get(id)
if !ok || s.TaskStream == nil {
if !ok {
return "", nil
}
stream := s.GetTaskStream()
if stream == nil {
return "", nil
}
@@ -228,7 +241,7 @@ func getServerConfig(c *gin.Context) (string, error) {
return "", singleton.Localizer.ErrorT("permission denied")
}
if err := s.TaskStream.Send(&pb.Task{
if err := stream.Send(&pb.Task{
Type: model.TaskTypeReportConfig,
}); err != nil {
return "", err
@@ -276,7 +289,7 @@ func setServerConfig(c *gin.Context) (*model.ServerTaskResponse, error) {
if !s.HasPermission(c) {
return nil, singleton.Localizer.ErrorT("permission denied")
}
if s.TaskStream == nil {
if s.GetTaskStream() == nil {
resp.Offline = append(resp.Offline, s.ID)
continue
}
@@ -300,7 +313,14 @@ func setServerConfig(c *gin.Context) (*model.ServerTaskResponse, error) {
Type: model.TaskTypeApplyConfig,
Data: configForm.Config,
}
if err := s.TaskStream.Send(task); err != nil {
stream := s.GetTaskStream()
if stream == nil {
respMu.Lock()
resp.Offline = append(resp.Offline, s.ID)
respMu.Unlock()
continue
}
if err := stream.Send(task); err != nil {
respMu.Lock()
resp.Failure = append(resp.Failure, s.ID)
respMu.Unlock()
@@ -321,23 +341,29 @@ func setServerConfig(c *gin.Context) (*model.ServerTaskResponse, error) {
// @Summary Batch move servers to other user
// @Security BearerAuth
// @Schemes
// @Description Batch move servers to other user
// @Description Initiates one ServerTransfer per requested server and returns a
// @Description per-server result. The old behaviour flipped Server.UserID in
// @Description a single SQL UPDATE without telling the agent, so the agent
// @Description kept presenting its old AgentSecret — which now belonged to a
// @Description different user — and authorizeAgentForUUID dropped it. The
// @Description current flow writes a Pending ServerTransfer row and flips
// @Description Server.UserID to the target owner immediately; that row keeps
// @Description the old owner's AgentSecret acceptable for this UUID until the
// @Description agent reconnects under the new secret (MarkVerified clears the
// @Description pending window) or the transfer Cancel/Fail/Timeout-out and
// @Description reverts Server.UserID to the source owner.
// @Tags auth required
// @Accept json
// @Param request body model.BatchMoveServerForm true "BatchMoveServerForm"
// @Produce json
// @Success 200 {object} model.CommonResponse[any]
// @Success 200 {object} model.CommonResponse[[]model.BatchMoveServerResult]
// @Router /batch-move/server [post]
func batchMoveServer(c *gin.Context) (any, error) {
func batchMoveServer(c *gin.Context) ([]model.BatchMoveServerResult, error) {
var moveForm model.BatchMoveServerForm
if err := c.ShouldBindJSON(&moveForm); err != nil {
return nil, err
}
if !singleton.ServerShared.CheckPermission(c, slices.Values(moveForm.Ids)) {
return nil, singleton.Localizer.ErrorT("permission denied")
}
if moveForm.ToUser == 0 {
return nil, singleton.Localizer.ErrorT("user id is required")
}
@@ -347,35 +373,81 @@ func batchMoveServer(c *gin.Context) (any, error) {
}
singleton.UserLock.RLock()
defer singleton.UserLock.RUnlock()
if _, ok := singleton.UserInfoMap[moveForm.ToUser]; !ok {
_, toUserExists := singleton.UserInfoMap[moveForm.ToUser]
singleton.UserLock.RUnlock()
if !toUserExists {
return nil, singleton.Localizer.ErrorT("user id %d does not exist", moveForm.ToUser)
}
err := singleton.DB.Transaction(func(tx *gorm.DB) error {
if err := tx.Model(&model.Server{}).Where("id in (?)", moveForm.Ids).Update("user_id", moveForm.ToUser).Error; err != nil {
return err
}
return nil
})
results := make([]model.BatchMoveServerResult, 0, len(moveForm.Ids))
uid := getUid(c)
isAdmin := callerIsAdmin(c)
if err != nil {
return nil, newGormError("%v", err)
}
for _, sid := range moveForm.Ids {
res := model.BatchMoveServerResult{ServerID: sid}
idsMap := make(map[uint64]bool)
for _, id := range moveForm.Ids {
idsMap[id] = true
}
for _, s := range singleton.ServerShared.Range {
if s == nil || !idsMap[s.ID] {
srv, ok := singleton.ServerShared.Get(sid)
if !ok || srv == nil {
res.Status = model.BatchMoveServerResultServerNotFound
results = append(results, res)
continue
}
s.UserID = moveForm.ToUser
// Per-server permission: admin or current owner. We do NOT use the
// bulk CheckPermission because we want a partial-success response
// rather than rejecting the whole batch on the first unauthorized id.
//
// 必须走 GetUserID() 而不是裸读 srv.UserID — ServerTransfer.Register
// 和 revertTransition 会通过 atomic.StoreUint64 改写当前 Server.UserID
// 以反映新所有者。batchMoveServer 与 transfer 流程并发时(典型场景:两
// 个 operator 几乎同时发起 move),裸读会与 SetUserID 形成 data race
// 且可能在 transfer 切换瞬间读到过期值并据此做权限/同所有者/fromUser
// 判断。
currentOwner := srv.GetUserID()
if !isAdmin && currentOwner != uid {
// Match the unknown-id response for members. A distinct
// permission_denied result lets callers enumerate foreign server IDs.
res.Status = model.BatchMoveServerResultServerNotFound
results = append(results, res)
continue
}
return nil, nil
if currentOwner == moveForm.ToUser {
res.Status = model.BatchMoveServerResultSameOwner
results = append(results, res)
continue
}
// One active ServerTransfer per server. InitiateExclusive serializes
// the HasPending guard, the DB transaction, and the in-memory
// Register under a per-server claim so two concurrent operators
// can't both observe "no pending", both insert, and silently end up
// with two Pending rows for the same server.
fromUser := currentOwner
created, err := singleton.ServerTransferShared.InitiateExclusive(sid, fromUser, moveForm.ToUser, uid)
if err != nil {
switch {
case errors.Is(err, singleton.ErrServerAlreadyTransferring):
res.Status = model.BatchMoveServerResultAlreadyTransferring
case errors.Is(err, singleton.ErrAgentTooOldForTransfer):
res.Status = model.BatchMoveServerResultAgentTooOld
res.Error = err.Error()
default:
res.Status = model.BatchMoveServerResultServerNotFound
res.Error = err.Error()
}
results = append(results, res)
continue
}
singleton.ServerTransferShared.PushIfOnline(created)
res.Status = model.BatchMoveServerResultPending
res.TransferID = created.ID
results = append(results, res)
}
return results, nil
}
var serverMetricMap = map[string]tsdb.MetricType{
@@ -141,6 +141,39 @@ func TestJWTInitParamsPinsAlgorithmAndSameSite(t *testing.T) {
}
}
// MEDIUM security: an IOStream session created by the old owner (terminal,
// file-manager, NAT) must be torn down when the server's ownership rotates
// — Register on Initiate, revertTransition on Cancel/Fail/Timeout, and
// OnServersDeleted on delete. Otherwise the old owner keeps an open
// websocket attached to a server they no longer own, which is effectively
// post-transfer RCE / file-read.
func TestServerTransferTransitionRevokesActiveIOStreams(t *testing.T) {
gin.SetMode(gin.TestMode)
ensureLocalizerForStreamTests(t)
rpc.NezhaHandlerSingleton = rpc.NewNezhaHandler()
originalHook := singleton.ServerTransferStreamRevocationHook
singleton.ServerTransferStreamRevocationHook = rpc.NezhaHandlerSingleton.RevokeStreamsForServer
defer func() {
singleton.ServerTransferStreamRevocationHook = originalHook
}()
rpc.NezhaHandlerSingleton.CreateStream("term-server-1", 100, 1)
rpc.NezhaHandlerSingleton.CreateStream("fm-server-1", 100, 1)
rpc.NezhaHandlerSingleton.CreateStream("term-server-2", 100, 2)
singleton.ServerTransferRevokeStreamsForServer(1)
if _, exists := rpc.NezhaHandlerSingleton.StreamOwnership("term-server-1"); exists {
t.Fatal("terminal stream for transferred server 1 must be revoked on ownership rotation")
}
if _, exists := rpc.NezhaHandlerSingleton.StreamOwnership("fm-server-1"); exists {
t.Fatal("file-manager stream for transferred server 1 must be revoked on ownership rotation")
}
if _, exists := rpc.NezhaHandlerSingleton.StreamOwnership("term-server-2"); !exists {
t.Fatal("unrelated server's stream must NOT be revoked")
}
}
// nz-o2s carries the OAuth2 state binding that authenticates the callback.
// The frontend never reads it, so HttpOnly is safe to enable and shuts the
// door on XSS attempting to steal the state.
+6 -2
View File
@@ -31,7 +31,11 @@ func createTerminal(c *gin.Context) (*model.CreateTerminalResponse, error) {
}
server, _ := singleton.ServerShared.Get(createTerminalReq.ServerID)
if server == nil || server.TaskStream == nil {
if server == nil {
return nil, singleton.Localizer.ErrorT("server not found or not connected")
}
stream := server.GetTaskStream()
if stream == nil {
return nil, singleton.Localizer.ErrorT("server not found or not connected")
}
@@ -49,7 +53,7 @@ func createTerminal(c *gin.Context) (*model.CreateTerminalResponse, error) {
terminalData, _ := json.Marshal(&model.TerminalTask{
StreamID: streamId,
})
if err := server.TaskStream.Send(&proto.Task{
if err := stream.Send(&proto.Task{
Type: model.TaskTypeTerminalGRPC,
Data: string(terminalData),
}); err != nil {
+217
View File
@@ -0,0 +1,217 @@
package controller
import (
"strconv"
"time"
"github.com/gin-gonic/gin"
"github.com/goccy/go-json"
"github.com/gorilla/websocket"
"github.com/nezhahq/nezha/model"
"github.com/nezhahq/nezha/service/singleton"
)
// List server transfers
// @Summary List server transfers
// @Security BearerAuth
// @Schemes
// @Description Returns transfers visible to the caller. Admin sees all; a
// @Description member sees rows where they are FromUserID, ToUserID, or
// @Description InitiatorID. The same predicate is enforced both at the SQL
// @Description level (this handler) and by the listHandler post-filter
// @Description (ServerTransfer.HasPermission) — defence in depth.
// @Tags auth required
// @Produce json
// @Success 200 {object} model.CommonResponse[[]model.ServerTransfer]
// @Router /transfer [get]
func listServerTransfer(c *gin.Context) ([]*model.ServerTransfer, error) {
q := singleton.DB.Order("id DESC")
// ServerTransfer is the only listX endpoint that hits the DB — the others
// all serve in-memory caches — and it is an append-only audit log. Without
// this SQL-side filter, every member's page load scans the entire historical
// population only to have the listHandler post-filter throw most of it
// away. As the table grows (a single transfer per server-move adds a row
// forever) this degrades from cheap to dashboard-blocking. Mirror the
// HasPermission predicate at the WHERE clause for non-admins. The post-filter
// still runs unconditionally as a defence-in-depth guard.
if !callerIsAdmin(c) {
uid := getUid(c)
q = q.Where("from_user_id = ? OR to_user_id = ? OR initiator_id = ?", uid, uid, uid)
}
var transfers []*model.ServerTransfer
if err := q.Find(&transfers).Error; err != nil {
return nil, newGormError("%v", err)
}
return transfers, nil
}
// Cancel server transfer
// @Summary Cancel server transfer
// @Security BearerAuth
// @Schemes
// @Description Cancels a Pending transfer and reverts Server.UserID back to
// @Description FromUserID. Only admin or the original FromUserID may cancel
// @Description (the new owner cannot — that would be a denial primitive
// @Description against a server they don't own yet). No-op if the transfer is
// @Description already terminal.
// @Tags auth required
// @Param id path uint true "Transfer ID"
// @Produce json
// @Success 200 {object} model.CommonResponse[model.ServerTransfer]
// @Router /transfer/{id}/cancel [post]
func cancelServerTransfer(c *gin.Context) (*model.ServerTransfer, error) {
tid, err := strconv.ParseUint(c.Param("id"), 10, 64)
if err != nil {
return nil, err
}
// Avoid leaking transfer-row existence via response shape. Admin can
// look up any row; a member can only look up rows where they are the
// FromUserID. Both "row does not exist" and "row exists but caller is
// not FromUserID" must surface identically as permission denied.
q := singleton.DB
if !callerIsAdmin(c) {
q = q.Where("from_user_id = ?", getUid(c))
}
var t model.ServerTransfer
if err := q.First(&t, tid).Error; err != nil {
return nil, singleton.Localizer.ErrorT("permission denied")
}
updated, err := singleton.ServerTransferShared.Cancel(tid)
if err != nil {
return nil, err
}
if updated == nil {
// Already terminal — return current state so the UI can refresh.
return &t, nil
}
return updated, nil
}
// Retry server transfer
// @Summary Retry server transfer
// @Security BearerAuth
// @Schemes
// @Description Creates a fresh Pending transfer with the same From/To as a
// @Description terminal (Failed/Timeout/Cancelled) transfer. The previous row
// @Description is left intact for audit; this returns the new row. Admin-only:
// @Description non-admin transfer semantics are enforced by batchMoveServer's
// @Description "ToUser == self" rule, so allowing any historical From/To/
// @Description Initiator to retry would reintroduce the give-away path. Non-
// @Description admins receive permission denied before the transfer row is
// @Description read so the response cannot enumerate transfer ids.
// @Tags auth required
// @Param id path uint true "Transfer ID"
// @Produce json
// @Success 200 {object} model.CommonResponse[model.ServerTransfer]
// @Router /transfer/{id}/retry [post]
func retryServerTransfer(c *gin.Context) (*model.ServerTransfer, error) {
// Retry is admin-only. Non-admin transfer semantics are enforced by
// batchMoveServer's "ToUser == self" rule, which means a member can only
// receive a server, never give one away. Allowing a member to retry a
// historical row would reintroduce the give-away path: any prev.ToUserID
// on file becomes a one-click bypass of that policy. Members who want
// the server moved elsewhere ask an admin or use batch-move to pull it
// onto themselves.
//
// Refuse non-admins before reading the row so the response cannot be
// used to enumerate which transfer ids exist.
if !callerIsAdmin(c) {
return nil, singleton.Localizer.ErrorT("permission denied")
}
tid, err := strconv.ParseUint(c.Param("id"), 10, 64)
if err != nil {
return nil, err
}
var prev model.ServerTransfer
if err := singleton.DB.First(&prev, tid).Error; err != nil {
return nil, newGormError("%v", err)
}
return singleton.ServerTransferShared.Retry(&prev, getUid(c))
}
// transferStreamWriteTimeout caps a single WriteMessage so a stuck or
// half-open client cannot block the broker fan-out forever — once exceeded
// the connection is considered dead and dropped. Matches the cadence of the
// keepalive ping below.
const transferStreamWriteTimeout = 10 * time.Second
// transferStreamPingInterval is how often we send a ping to keep the
// connection alive through aggressive proxies. Independent of the event
// stream, so silent transfers still keep the socket warm.
const transferStreamPingInterval = 30 * time.Second
// Websocket server transfer stream
// @Summary Websocket server transfer stream
// @Security BearerAuth
// @Schemes
// @Description Pushes ServerTransfer state transitions (Pending → Verified /
// @Description Failed / Timeout / Cancelled) to the dashboard so the UI can
// @Description react without polling. Each frame is a single JSON-encoded
// @Description ServerTransfer. Subscribers see only transfers visible to
// @Description them (ServerTransfer.HasPermission).
// @tags common
// @Produce json
// @Success 200 {object} model.ServerTransfer
// @Router /ws/transfer [get]
func transferStream(c *gin.Context) (any, error) {
conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
if err != nil {
return nil, newWsError("%v", err)
}
defer conn.Close()
subID, ch := singleton.ServerTransferShared.Subscribe()
defer singleton.ServerTransferShared.Unsubscribe(subID)
// Pings keep the socket warm even when the broker is quiet. Without this
// a long idle period followed by a transfer event would race against
// upstream proxy idle-timeouts that may have already closed the conn.
ping := time.NewTicker(transferStreamPingInterval)
defer ping.Stop()
// Reader goroutine: needed only to surface client disconnects through
// SetReadDeadline / ReadMessage. We never expect inbound payloads.
closed := make(chan struct{})
go func() {
defer close(closed)
for {
if _, _, err := conn.NextReader(); err != nil {
return
}
}
}()
for {
select {
case <-closed:
return nil, newWsError("")
case <-ping.C:
if err := conn.WriteControl(websocket.PingMessage, nil, time.Now().Add(transferStreamWriteTimeout)); err != nil {
return nil, newWsError("%v", err)
}
case t, ok := <-ch:
if !ok {
return nil, newWsError("")
}
if !t.HasPermission(c) {
continue
}
payload, err := json.Marshal(t)
if err != nil {
continue
}
if err := conn.SetWriteDeadline(time.Now().Add(transferStreamWriteTimeout)); err != nil {
return nil, newWsError("%v", err)
}
if err := conn.WriteMessage(websocket.TextMessage, payload); err != nil {
return nil, newWsError("%v", err)
}
}
}
}
@@ -0,0 +1,231 @@
package controller
import (
"encoding/json"
"net/http"
"net/http/httptest"
"strings"
"sync"
"testing"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"gorm.io/driver/sqlite"
"gorm.io/gorm"
"github.com/nezhahq/nezha/model"
"github.com/nezhahq/nezha/pkg/i18n"
"github.com/nezhahq/nezha/service/singleton"
)
// listServerTransfer originally did `SELECT * FROM server_transfers ORDER BY
// id DESC` and relied on the listHandler post-filter (HasPermission) to drop
// rows the caller can't see. Functionally correct, but ServerTransfer is the
// only listX endpoint that hits the DB (the others serve in-memory caches),
// and it's an append-only audit table — every page load by a member who has
// participated in two transfers triggers a full table scan over the entire
// historical population. That cost is silent until the table is big and the
// dashboard slows for everyone at once. The fix pushes the same predicate
// HasPermission encodes down into the WHERE clause for non-admin callers.
//
// This test pins down the behavioural contract: regardless of the optimisation,
// a member must see only their own rows. It is intentionally written against
// the same response shape as the production handler so a regression in either
// the SQL filter OR the post-filter would fail it.
func TestListServerTransferReturnsOnlyCallerVisibleRowsForMember(t *testing.T) {
cleanup := setupListServerTransferFixture(t)
defer cleanup()
// alice=100, bob=200, charlie=300. Seed five rows covering every position
// a member could occupy plus one row the member must NEVER see.
aliceFrom := seedTerminalTransfer(t, 1, 100, 200, 100)
aliceTo := seedTerminalTransfer(t, 2, 300, 100, 300)
aliceInitiated := seedTerminalTransfer(t, 3, 200, 300, 100)
bobAlone := seedTerminalTransfer(t, 4, 200, 300, 200)
charlieAlone := seedTerminalTransfer(t, 5, 300, 200, 300)
ids := callListServerTransfer(t, 100, model.RoleMember)
assert.ElementsMatch(t,
[]uint64{aliceFrom, aliceTo, aliceInitiated},
ids,
"member must see exactly the rows where they are From/To/Initiator",
)
for _, forbidden := range []uint64{bobAlone, charlieAlone} {
assert.NotContains(t, ids, forbidden, "member must never see a row they do not participate in")
}
}
// Admin sees every row. This pins down that the SQL-level filter is gated on
// role and is not accidentally applied to admins (which would be a regression
// in the other direction).
func TestListServerTransferReturnsEveryRowForAdmin(t *testing.T) {
cleanup := setupListServerTransferFixture(t)
defer cleanup()
a := seedTerminalTransfer(t, 1, 100, 200, 100)
b := seedTerminalTransfer(t, 2, 200, 300, 200)
c := seedTerminalTransfer(t, 3, 300, 100, 300)
ids := callListServerTransfer(t, 999, model.RoleAdmin)
assert.ElementsMatch(t, []uint64{a, b, c}, ids, "admin sees every row")
}
// Pin down the optimisation directly: the SELECT that hits the audit table
// for a non-admin caller MUST include a per-user WHERE clause. The behavioural
// tests above would pass even if the SQL stayed `SELECT *` (the post-filter
// hides forbidden rows), so they cannot regress-detect the perf fix going
// away. This test captures the executed SQL and asserts the filter is pushed
// down to the database.
//
// Why pinning the optimisation matters: ServerTransfer is the only listX
// endpoint that hits the DB (the rest serve in-memory caches) and it grows
// unbounded as an audit log. Without SQL-side filtering, every member's page
// load scans the entire historical table. A future refactor that drops the
// per-caller WHERE clause would not be caught by any behavioural assertion —
// hence this explicit guard.
func TestListServerTransferPushesPerCallerFilterIntoSQLForMember(t *testing.T) {
cleanup, captured := setupListServerTransferFixtureWithSQLCapture(t)
defer cleanup()
seedTerminalTransfer(t, 1, 100, 200, 100)
seedTerminalTransfer(t, 2, 200, 300, 200)
_ = callListServerTransfer(t, 100, model.RoleMember)
stmt := findSelectAgainstServerTransfers(captured.Snapshot())
require.NotEmpty(t, stmt, "expected a SELECT against server_transfers to be issued")
low := strings.ToLower(stmt)
require.Contains(t, low, "where", "non-admin list must apply a per-caller WHERE filter at the SQL level")
for _, col := range []string{"from_user_id", "to_user_id", "initiator_id"} {
require.Contains(t, low, col, "WHERE clause must filter on %s", col)
}
}
// The admin path must NOT push the per-user filter — admins see everything.
// Without this assertion a refactor that always applies the filter would
// silently hide cross-tenant rows from admins (an availability regression of
// the admin observability surface).
func TestListServerTransferOmitsPerCallerFilterForAdmin(t *testing.T) {
cleanup, captured := setupListServerTransferFixtureWithSQLCapture(t)
defer cleanup()
seedTerminalTransfer(t, 1, 100, 200, 100)
_ = callListServerTransfer(t, 999, model.RoleAdmin)
stmt := findSelectAgainstServerTransfers(captured.Snapshot())
require.NotEmpty(t, stmt, "expected a SELECT against server_transfers to be issued")
low := strings.ToLower(stmt)
for _, col := range []string{"from_user_id", "to_user_id", "initiator_id"} {
require.NotContains(t, low, col, "admin list must NOT filter by %s — admins observe all transfers", col)
}
}
func setupListServerTransferFixture(t *testing.T) func() {
t.Helper()
if singleton.Localizer == nil {
singleton.Localizer = i18n.NewLocalizer("en_US", "nezha", "translations", i18n.Translations)
}
originalDB := singleton.DB
db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
assert.NoError(t, err)
assert.NoError(t, db.AutoMigrate(&model.Server{}, &model.ServerTransfer{}))
singleton.DB = db
return func() { singleton.DB = originalDB }
}
// sqlCapture records every SQL statement gorm executes against the test DB.
// Used by the optimisation guards above to assert that listServerTransfer
// actually pushes its per-caller filter into the WHERE clause.
type sqlCapture struct {
mu sync.Mutex
stmts []string
}
func (s *sqlCapture) record(stmt string) {
s.mu.Lock()
defer s.mu.Unlock()
s.stmts = append(s.stmts, stmt)
}
func (s *sqlCapture) Snapshot() []string {
s.mu.Lock()
defer s.mu.Unlock()
out := make([]string, len(s.stmts))
copy(out, s.stmts)
return out
}
func setupListServerTransferFixtureWithSQLCapture(t *testing.T) (func(), *sqlCapture) {
t.Helper()
cleanup := setupListServerTransferFixture(t)
cap := &sqlCapture{}
err := singleton.DB.Callback().Query().After("gorm:query").Register("test:capture_sql", func(tx *gorm.DB) {
cap.record(tx.Statement.SQL.String())
})
require.NoError(t, err)
return func() {
_ = singleton.DB.Callback().Query().Remove("test:capture_sql")
cleanup()
}, cap
}
// findSelectAgainstServerTransfers returns the first captured SELECT whose
// FROM clause is `server_transfers`. We don't care about ordering callbacks
// or AutoMigrate scaffolding queries — only the handler's own SELECT.
func findSelectAgainstServerTransfers(stmts []string) string {
for _, s := range stmts {
low := strings.ToLower(s)
if strings.HasPrefix(strings.TrimSpace(low), "select") && strings.Contains(low, "server_transfers") {
return s
}
}
return ""
}
func seedTerminalTransfer(t *testing.T, serverID, fromUserID, toUserID, initiatorID uint64) uint64 {
t.Helper()
tr := &model.ServerTransfer{
ServerID: serverID,
FromUserID: fromUserID,
ToUserID: toUserID,
InitiatorID: initiatorID,
Status: model.ServerTransferStatusVerified,
}
assert.NoError(t, singleton.DB.Create(tr).Error)
return tr.ID
}
func callListServerTransfer(t *testing.T, callerID uint64, role model.Role) []uint64 {
t.Helper()
r := gin.New()
r.Use(func(c *gin.Context) {
setAuthUser(c, callerID, role)
c.Next()
})
r.GET("/transfer", listHandler(listServerTransfer))
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/transfer", nil)
r.ServeHTTP(w, req)
assert.Equal(t, http.StatusOK, w.Code)
var resp struct {
Success bool `json:"success"`
Data []*model.ServerTransfer `json:"data"`
Error string `json:"error"`
}
assert.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp))
assert.True(t, resp.Success, "list call must succeed: %s", resp.Error)
ids := make([]uint64, 0, len(resp.Data))
for _, tr := range resp.Data {
ids = append(ids, tr.ID)
}
return ids
}
@@ -0,0 +1,219 @@
package controller
import (
"bytes"
"encoding/json"
"net/http"
"net/http/httptest"
"strconv"
"testing"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/assert"
"gorm.io/driver/sqlite"
"gorm.io/gorm"
"github.com/nezhahq/nezha/model"
"github.com/nezhahq/nezha/pkg/i18n"
"github.com/nezhahq/nezha/service/singleton"
)
// retryServerTransfer previously gated on prev.HasPermission(c), which honours
// the historical transfer row (FromUserID, ToUserID, InitiatorID). That lets
// any of those original parties re-initiate a transfer of the server long
// after ownership has moved on. Concretely: a stale "alice -> bob" Failed row
// stays visible to alice forever — even after she's transferred the server
// off to charlie — and the original endpoint would happily move it from
// charlie to bob without ever consulting the current owner.
//
// Authorization for an action that mutates the live server must use the
// live server, not a historical audit row.
func TestRetryServerTransferRejectsCallerWhoNoLongerOwnsServer(t *testing.T) {
cleanup := setupRetryServerTransferFixture(t)
defer cleanup()
// Seed: server originally owned by user 100 (alice). Failed transfer to
// 200 (bob) is recorded but server ownership has since moved to 300
// (charlie) — e.g. alice transferred elsewhere afterwards. Alice is no
// longer the owner, so retrying the stale row would be an unauthorized
// grab.
seedServer(t, 1, 300)
staleID := seedFailedTransfer(t, 1, 100 /*from*/, 200 /*to*/, 100 /*initiator*/)
resp, status := callRetryServerTransfer(t, staleID, 100, model.RoleMember)
assert.Equal(t, http.StatusOK, status)
assert.False(t, resp.Success, "alice no longer owns server 1; retry must be rejected")
assert.Contains(t, resp.Error, "permission denied")
var s model.Server
assert.NoError(t, singleton.DB.First(&s, 1).Error)
assert.Equal(t, uint64(300), s.UserID, "rejected retry must not flip ownership")
var count int64
assert.NoError(t, singleton.DB.Model(&model.ServerTransfer{}).Where("status = ?", model.ServerTransferStatusPending).Count(&count).Error)
assert.Equal(t, int64(0), count, "rejected retry must not create a Pending row")
}
// The historical ToUserID must also not be able to grab the server back via
// the stale row. Same root cause; this is the explicit assertion that the
// fix covers the To side, not just the From side.
func TestRetryServerTransferRejectsHistoricalTargetWhoNeverOwnedServer(t *testing.T) {
cleanup := setupRetryServerTransferFixture(t)
defer cleanup()
seedServer(t, 1, 300)
staleID := seedFailedTransfer(t, 1, 100, 200, 100)
resp, status := callRetryServerTransfer(t, staleID, 200, model.RoleMember)
assert.Equal(t, http.StatusOK, status)
assert.False(t, resp.Success, "bob was the failed transfer's target; he never owned server 1")
assert.Contains(t, resp.Error, "permission denied")
}
// Members never retry: batchMoveServer's "ToUser == self" policy means a
// member can only RECEIVE a server, not give one away. Retry of a failed
// "alice -> bob" by alice (member, current owner) is exactly the give-away
// case batch-move would refuse. Retry is admin-only.
func TestRetryServerTransferRejectsCurrentOwnerWhoIsMember(t *testing.T) {
cleanup := setupRetryServerTransferFixture(t)
defer cleanup()
seedServer(t, 1, 100)
failedID := seedFailedTransfer(t, 1, 100, 200, 100)
resp, status := callRetryServerTransfer(t, failedID, 100, model.RoleMember)
assert.Equal(t, http.StatusOK, status)
assert.False(t, resp.Success, "member retry is forbidden — the give-away semantics bypass batchMoveServer's ToUser==self policy")
assert.Contains(t, resp.Error, "permission denied")
}
// batchMoveServer enforces "non-admin caller may only move a server TO
// themselves" (controller/server.go: ToUser != getUid(c) returns permission
// denied). retryServerTransfer historically only checked the live owner
// and not the transfer's ToUserID, which let the current owner re-push the
// server to ANY historical ToUserID — bypassing the batch-move policy.
//
// Concretely: alice (member) currently owns server 1; she finds a Failed
// transfer whose ToUserID is bob and retries it. The server lands on bob
// even though batch-move would have refused "alice -> bob" from her.
func TestRetryServerTransferRejectsNonAdminPushingToForeignToUserID(t *testing.T) {
cleanup := setupRetryServerTransferFixture(t)
defer cleanup()
seedServer(t, 1, 100)
staleID := seedFailedTransfer(t, 1, 100 /*from*/, 200 /*to*/, 100 /*initiator*/)
resp, status := callRetryServerTransfer(t, staleID, 100, model.RoleMember)
assert.Equal(t, http.StatusOK, status)
assert.False(t, resp.Success, "non-admin owner cannot push their server to a historical foreign ToUserID — that would bypass batchMoveServer's ToUser==self policy")
assert.Contains(t, resp.Error, "permission denied")
var s model.Server
assert.NoError(t, singleton.DB.First(&s, 1).Error)
assert.Equal(t, uint64(100), s.UserID, "rejected retry must not flip ownership")
var count int64
assert.NoError(t, singleton.DB.Model(&model.ServerTransfer{}).Where("status = ?", model.ServerTransferStatusPending).Count(&count).Error)
assert.Equal(t, int64(0), count, "rejected retry must not create a Pending row")
}
// Admins must always be able to retry — they are the last-resort recovery
// path when an operator-cancelled transfer needs to be re-pushed.
func TestRetryServerTransferAllowsAdmin(t *testing.T) {
cleanup := setupRetryServerTransferFixture(t)
defer cleanup()
seedServer(t, 1, 300)
failedID := seedFailedTransfer(t, 1, 100, 200, 100)
resp, status := callRetryServerTransfer(t, failedID, 999, model.RoleAdmin)
assert.Equal(t, http.StatusOK, status)
assert.True(t, resp.Success, "admin must be able to retry any transfer: error=%s", resp.Error)
}
func setupRetryServerTransferFixture(t *testing.T) func() {
t.Helper()
if singleton.Localizer == nil {
singleton.Localizer = i18n.NewLocalizer("en_US", "nezha", "translations", i18n.Translations)
}
originalDB := singleton.DB
originalShared := singleton.ServerShared
originalTransferShared := singleton.ServerTransferShared
originalUserMap := singleton.UserInfoMap
db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
assert.NoError(t, err)
assert.NoError(t, db.AutoMigrate(&model.Server{}, &model.ServerTransfer{}))
singleton.DB = db
singleton.ServerShared = singleton.NewServerClass()
singleton.UserInfoMap = map[uint64]model.UserInfo{
100: {Role: model.RoleMember, AgentSecret: "alice-secret"},
200: {Role: model.RoleMember, AgentSecret: "bob-secret"},
300: {Role: model.RoleMember, AgentSecret: "charlie-secret"},
}
singleton.ServerTransferShared = singleton.NewServerTransferClass()
return func() {
if singleton.ServerTransferShared != nil {
singleton.ServerTransferShared.Stop()
}
singleton.DB = originalDB
singleton.ServerShared = originalShared
singleton.ServerTransferShared = originalTransferShared
singleton.UserInfoMap = originalUserMap
}
}
func seedServer(t *testing.T, id, ownerID uint64) {
t.Helper()
s := &model.Server{
Common: model.Common{ID: id, UserID: ownerID},
UUID: "uuid-" + strconv.FormatUint(id, 10),
Name: "seeded",
}
assert.NoError(t, singleton.DB.Create(s).Error)
model.InitServer(s)
singleton.ServerShared.Update(s, s.UUID)
}
func seedFailedTransfer(t *testing.T, serverID, fromUserID, toUserID, initiatorID uint64) uint64 {
t.Helper()
tr := &model.ServerTransfer{
ServerID: serverID,
FromUserID: fromUserID,
ToUserID: toUserID,
InitiatorID: initiatorID,
Status: model.ServerTransferStatusFailed,
LastError: "seeded",
}
assert.NoError(t, singleton.DB.Create(tr).Error)
return tr.ID
}
func callRetryServerTransfer(t *testing.T, transferID, callerID uint64, role model.Role) (commonResponseShape, int) {
t.Helper()
r := gin.New()
r.Use(func(c *gin.Context) {
setAuthUser(c, callerID, role)
c.Next()
})
r.POST("/transfer/:id/retry", commonHandler(retryServerTransfer))
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodPost, "/transfer/"+strconv.FormatUint(transferID, 10)+"/retry", bytes.NewReader(nil))
r.ServeHTTP(w, req)
var resp commonResponseShape
assert.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp))
return resp, w.Code
}
type commonResponseShape struct {
Success bool `json:"success"`
Error string `json:"error"`
}
+1 -1
View File
@@ -195,7 +195,7 @@ func getServerStat(withPublicNote bool, viewerUserID uint64, viewerIsAdmin bool)
func filterServersForViewer(servers []*model.Server, viewerUserID uint64, viewerIsAdmin bool, withPublicNote bool) []model.StreamServer {
out := make([]model.StreamServer, 0, len(servers))
for _, server := range servers {
isOwnerOrAdmin := viewerIsAdmin || (viewerUserID != 0 && server.UserID == viewerUserID)
isOwnerOrAdmin := viewerIsAdmin || (viewerUserID != 0 && server.GetUserID() == viewerUserID)
if server.HideForGuest && !isOwnerOrAdmin {
continue
}
+31 -9
View File
@@ -24,6 +24,10 @@ import (
func ServeRPC() *grpc.Server {
server := grpc.NewServer(grpc.ChainUnaryInterceptor(getRealIp, waf))
rpcService.NezhaHandlerSingleton = rpcService.NewNezhaHandler()
// Install the IOStream revocation hook so ServerTransferShared can tear
// down terminal/FM/NAT sessions held by the previous owner on every
// ownership rotation (Register/revertTransition/OnServersDeleted).
singleton.ServerTransferStreamRevocationHook = rpcService.NezhaHandlerSingleton.RevokeStreamsForServer
proto.RegisterNezhaServiceServer(server, rpcService.NezhaHandlerSingleton)
return server
}
@@ -89,22 +93,30 @@ func DispatchTask(serviceSentinelDispatchBus <-chan *model.Service) {
}
server, _ := singleton.ServerShared.Get(id)
if server == nil || server.TaskStream == nil {
if server == nil {
continue
}
stream := server.GetTaskStream()
if stream == nil {
continue
}
if canSendTaskToServer(task, server) {
server.TaskStream.Send(task.PB())
stream.Send(task.PB())
}
}
case model.ServiceCoverAll:
for id, server := range singleton.ServerShared.Range {
if server == nil || server.TaskStream == nil || task.SkipServers[id] {
if server == nil || task.SkipServers[id] {
continue
}
stream := server.GetTaskStream()
if stream == nil {
continue
}
if canSendTaskToServer(task, server) {
server.TaskStream.Send(task.PB())
stream.Send(task.PB())
}
}
}
@@ -115,17 +127,27 @@ func DispatchKeepalive() {
singleton.CronShared.AddFunc("@every 20s", func() {
list := singleton.ServerShared.GetSortedList()
for _, s := range list {
if s == nil || s.TaskStream == nil {
if s == nil {
continue
}
s.TaskStream.Send(&proto.Task{Type: model.TaskTypeKeepalive})
stream := s.GetTaskStream()
if stream == nil {
continue
}
stream.Send(&proto.Task{Type: model.TaskTypeKeepalive})
}
})
}
func ServeNAT(w http.ResponseWriter, r *http.Request, natConfig *model.NAT) {
server, _ := singleton.ServerShared.Get(natConfig.ServerID)
if server == nil || server.TaskStream == nil {
if server == nil {
w.WriteHeader(http.StatusServiceUnavailable)
w.Write([]byte("server not found or not connected"))
return
}
stream := server.GetTaskStream()
if stream == nil {
w.WriteHeader(http.StatusServiceUnavailable)
w.Write([]byte("server not found or not connected"))
return
@@ -157,7 +179,7 @@ func ServeNAT(w http.ResponseWriter, r *http.Request, natConfig *model.NAT) {
return
}
if err := server.TaskStream.Send(&proto.Task{
if err := stream.Send(&proto.Task{
Type: model.TaskTypeNAT,
Data: string(taskData),
}); err != nil {
@@ -192,5 +214,5 @@ func canSendTaskToServer(task *model.Service, server *model.Server) bool {
}
singleton.UserLock.RUnlock()
return task.UserID == server.UserID || role.IsAdmin()
return task.UserID == server.GetUserID() || role.IsAdmin()
}
+3 -3
View File
@@ -32,9 +32,9 @@ require (
github.com/swaggo/gin-swagger v1.6.1
github.com/swaggo/swag v1.16.6
github.com/tidwall/gjson v1.19.0
golang.org/x/crypto v0.51.0
golang.org/x/crypto v0.52.0
golang.org/x/exp v0.0.0-20260508232706-74f9aab9d74a
golang.org/x/net v0.54.0
golang.org/x/net v0.55.0
golang.org/x/oauth2 v0.36.0
golang.org/x/sync v0.20.0
google.golang.org/grpc v1.81.1
@@ -109,7 +109,7 @@ require (
go.yaml.in/yaml/v3 v3.0.4 // indirect
golang.org/x/arch v0.27.0 // indirect
golang.org/x/mod v0.36.0 // indirect
golang.org/x/sys v0.44.0 // indirect
golang.org/x/sys v0.45.0 // indirect
golang.org/x/text v0.37.0 // indirect
golang.org/x/time v0.15.0 // indirect
golang.org/x/tools v0.45.0 // indirect
+6 -6
View File
@@ -246,8 +246,8 @@ golang.org/x/arch v0.27.0 h1:0WNVcR8u9yFz8j5FvdHpgwNp3FS5U4guYdzHwEiGjoU=
golang.org/x/arch v0.27.0/go.mod h1:0X+GdSIP+kL5wPmpK7sdkEVTt2XoYP0cSjQSbZBwOi8=
golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w=
golang.org/x/crypto v0.0.0-20210921155107-089bfa567519/go.mod h1:GvvjBRRGRdwPK5ydBHafDWAxML/pGHZbMvKqRZ5+Abc=
golang.org/x/crypto v0.51.0 h1:IBPXwPfKxY7cWQZ38ZCIRPI50YLeevDLlLnyC5wRGTI=
golang.org/x/crypto v0.51.0/go.mod h1:8AdwkbraGNABw2kOX6YFPs3WM22XqI4EXEd8g+x7Oc8=
golang.org/x/crypto v0.52.0 h1:RMs7fP2rXdep0CftQlK8Uf+kibLm7qkCcradZWYz988=
golang.org/x/crypto v0.52.0/go.mod h1:1QgfPxDqh0T2M/elOJtp9RvuR95kVjir0e6/BvEmGbc=
golang.org/x/exp v0.0.0-20260508232706-74f9aab9d74a h1:+3jdDGGB8NGb1Zktc737jlt3/A5f6UlwSzmvqUuufxw=
golang.org/x/exp v0.0.0-20260508232706-74f9aab9d74a/go.mod h1:d2fgXJLVs4dYDHUk5lwMIfzRzSrWCfGZb0ZqeLa/Vcw=
golang.org/x/mod v0.6.0-dev.0.20220419223038-86c51ed26bb4/go.mod h1:jJ57K6gSWd91VN4djpZkiMVwK6gcyfeH4XE8wZrZaV4=
@@ -257,8 +257,8 @@ golang.org/x/net v0.0.0-20190620200207-3b0461eec859/go.mod h1:z5CRVTTTmAJ677TzLL
golang.org/x/net v0.0.0-20210226172049-e18ecbb05110/go.mod h1:m0MpNAwzfU5UDzcl9v0D8zg8gWTRqZa9RBIspLL5mdg=
golang.org/x/net v0.0.0-20220722155237-a158d28d115b/go.mod h1:XRhObCWvk6IyKnWLug+ECip1KBveYUHfp+8e9klMJ9c=
golang.org/x/net v0.7.0/go.mod h1:2Tu9+aMcznHK/AK1HMvgo6xiTLG5rD5rZLDS+rp2Bjs=
golang.org/x/net v0.54.0 h1:2zJIZAxAHV/OHCDTCOHAYehQzLfSXuf/5SoL/Dv6w/w=
golang.org/x/net v0.54.0/go.mod h1:Sj4oj8jK6XmHpBZU/zWHw3BV3abl4Kvi+Ut7cQcY+cQ=
golang.org/x/net v0.55.0 h1:bcvxaJn3e1U6InsFWt1JUq1aSjnRxLzT2rtD2KfkDF8=
golang.org/x/net v0.55.0/go.mod h1:L5U2KuzuOe1lY7Z+aWVIKK6qEeJXnXV9yzGA+WCHJww=
golang.org/x/oauth2 v0.36.0 h1:peZ/1z27fi9hUOFCAZaHyrpWG5lwe0RJEEEeH0ThlIs=
golang.org/x/oauth2 v0.36.0/go.mod h1:YDBUJMTkDnJS+A4BP4eZBjCqtokkg1hODuPjwiGPO7Q=
golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
@@ -271,8 +271,8 @@ golang.org/x/sys v0.0.0-20210615035016-665e8c7367d1/go.mod h1:oPkhp1MJrh7nUepCBc
golang.org/x/sys v0.0.0-20220520151302-bc2c85ada10a/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
golang.org/x/sys v0.0.0-20220722155257-8c9f86f7a55f/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
golang.org/x/sys v0.5.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
golang.org/x/sys v0.44.0 h1:ildZl3J4uzeKP07r2F++Op7E9B29JRUy+a27EibtBTQ=
golang.org/x/sys v0.44.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw=
golang.org/x/sys v0.45.0 h1:dO4czNzziLiiXplLQgBCEpCvXQ3dnkn0SdaZSYdQ+FY=
golang.org/x/sys v0.45.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw=
golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo=
golang.org/x/term v0.0.0-20210927222741-03fcf44c2211/go.mod h1:jbD1KX2456YbFQfuXm/mYQcufACuNUgVhRMnK/tPxf8=
golang.org/x/term v0.5.0/go.mod h1:jMB1sMXY+tzblOD4FWmEbocvup2/aLOaQEp7JmGp78k=
+23 -2
View File
@@ -6,6 +6,7 @@ import (
"slices"
"strconv"
"strings"
"sync/atomic"
"time"
"github.com/gin-gonic/gin"
@@ -37,8 +38,21 @@ func (c *Common) GetID() uint64 {
return c.ID
}
// GetUserID 原子读取所属用户 ID。Server.UserID 会在 ServerTransfer 的
// Register/revertTransition 流程里被实时改写以反映新所有者,同时 auth
// 热路径在每次 agent RPC 都会读它。任何并发读必须走 atomic,否则与 SetUserID
// 一起会被 go race detector 识别为 data race(见
// TestServerUserIDConcurrentAccessIsRaceFree)。
func (c *Common) GetUserID() uint64 {
return c.UserID
return atomic.LoadUint64(&c.UserID)
}
// SetUserID 原子改写所属用户 ID。仅在「server 已经在 in-memory cache 里」
// 的写入路径(ServerTransfer.Register / revertTransition)需要用 atomic
// 保证可见性;普通 GORM AfterFind / Create 因为没有并发读所以可以直接赋
// 值。配合 GetUserID 形成 atomic-only 的并发访问协议。
func (c *Common) SetUserID(uid uint64) {
atomic.StoreUint64(&c.UserID, uid)
}
func (c *Common) HasPermission(ctx *gin.Context) bool {
@@ -52,7 +66,14 @@ func (c *Common) HasPermission(ctx *gin.Context) bool {
return true
}
return user.ID == c.UserID
// 必须走 GetUserID 而不是裸读 c.UserID — Server.UserID 在
// ServerTransfer.Register / revertTransition 里会被 atomic.StoreUint64
// 改写,dashboard 各 controller 在 listHandler post-filter 这条热路径上
// 高频对同一 *Server 调 HasPermission。裸读会与 SetUserID 形成 data
// raceTestCommonHasPermissionConcurrentWithSetUserIDIsRaceFree 在
// -race 下钉死该不变量),并且在 transfer 切换瞬间可能给出错误的权限
// 判断。
return user.ID == c.GetUserID()
}
type CommonInterface interface {
+65 -1
View File
@@ -3,6 +3,7 @@ package model
import (
"os"
"path/filepath"
"strconv"
"strings"
"github.com/go-viper/mapstructure/v2"
@@ -17,7 +18,11 @@ import (
const (
ConfigUsePeerIP = "NZ::Use-Peer-IP"
ConfigCoverAll = iota
JWTSecretKeyRotationBaselineVersion = "v2.0.13"
)
const (
ConfigCoverAll = iota + 1
ConfigCoverIgnoreAll
)
@@ -61,6 +66,7 @@ type Config struct {
JWTTimeout int `koanf:"jwt_timeout" json:"jwt_timeout,omitempty"` // JWT token过期时间(小时)
JWTSecretKey string `koanf:"jwt_secret_key" json:"jwt_secret_key,omitempty"`
JWTSecretKeyLastRotatedVersion string `koanf:"jwt_secret_key_last_rotated_version" json:"jwt_secret_key_last_rotated_version,omitempty"`
ListenPort uint16 `koanf:"listen_port" json:"listen_port,omitempty"`
ListenHost string `koanf:"listen_host" json:"listen_host,omitempty"`
@@ -193,6 +199,30 @@ func (c *Config) Save() error {
return c.save()
}
func (c *Config) RotateJWTSecretKeyIfNeeded(currentVersion string) (bool, error) {
currentVersion = strings.TrimSpace(currentVersion)
if compareVersion(currentVersion, JWTSecretKeyRotationBaselineVersion) < 0 {
return false, nil
}
initialMarker := c.JWTSecretKeyLastRotatedVersion
shouldRotate := c.JWTSecretKeyLastRotatedVersion == "" || compareVersion(c.JWTSecretKeyLastRotatedVersion, JWTSecretKeyRotationBaselineVersion) < 0
if shouldRotate {
secret, err := utils.GenerateRandomString(1024)
if err != nil {
return false, err
}
c.JWTSecretKey = secret
}
c.JWTSecretKeyLastRotatedVersion = currentVersion
if !shouldRotate && c.JWTSecretKeyLastRotatedVersion == initialMarker {
return false, nil
}
return shouldRotate, c.Save()
}
func (c *Config) save() error {
data, err := yaml.Marshal(c)
if err != nil {
@@ -211,6 +241,40 @@ func (c *Config) write(data []byte) error {
return os.WriteFile(c.filePath, data, 0600)
}
func compareVersion(left, right string) int {
leftParts, leftOK := parseVersion(left)
rightParts, rightOK := parseVersion(right)
if !leftOK || !rightOK {
return -1
}
for i := range leftParts {
if leftParts[i] < rightParts[i] {
return -1
}
if leftParts[i] > rightParts[i] {
return 1
}
}
return 0
}
func parseVersion(version string) ([3]int, bool) {
version = strings.TrimPrefix(strings.TrimSpace(version), "v")
parts := strings.Split(version, ".")
if len(parts) != 3 {
return [3]int{}, false
}
var parsed [3]int
for i, part := range parts {
value, err := strconv.Atoi(part)
if err != nil {
return [3]int{}, false
}
parsed[i] = value
}
return parsed, true
}
func koanfConf(c any) koanf.UnmarshalConf {
return koanf.UnmarshalConf{
DecoderConfig: &mapstructure.DecoderConfig{
+105
View File
@@ -151,6 +151,111 @@ func TestReadConfig(t *testing.T) {
})
}
func TestRotateJWTSecretKeyIfNeeded(t *testing.T) {
tests := []struct {
name string
initialMarker string
currentVersion string
wantRotated bool
wantStoredVersion string
wantSecretChanged bool
wantSavedConfigKey bool
}{
{
name: "empty marker rotates leaked secret",
currentVersion: "v2.0.13",
wantRotated: true,
wantStoredVersion: "v2.0.13",
wantSecretChanged: true,
wantSavedConfigKey: true,
},
{
name: "old marker rotates leaked secret",
initialMarker: "v2.0.12",
currentVersion: "v2.0.14",
wantRotated: true,
wantStoredVersion: "v2.0.14",
wantSecretChanged: true,
wantSavedConfigKey: true,
},
{
name: "threshold marker keeps secret and advances marker",
initialMarker: "v2.0.13",
currentVersion: "v2.0.14",
wantStoredVersion: "v2.0.14",
wantSavedConfigKey: true,
},
{
name: "current marker keeps secret",
initialMarker: "v2.0.14",
currentVersion: "v2.0.14",
wantStoredVersion: "v2.0.14",
},
{
name: "debug version skips rotation and marker update",
currentVersion: "debug",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
file := newTempConfig(t, "")
t.Cleanup(func() { os.Remove(file) })
c := &Config{
JWTSecretKey: "leaked-secret",
JWTSecretKeyLastRotatedVersion: tt.initialMarker,
filePath: file,
}
rotated, err := c.RotateJWTSecretKeyIfNeeded(tt.currentVersion)
if err != nil {
t.Fatalf("rotate jwt secret key failed: %v", err)
}
if rotated != tt.wantRotated {
t.Fatalf("rotated = %v, want %v", rotated, tt.wantRotated)
}
if c.JWTSecretKeyLastRotatedVersion != tt.wantStoredVersion {
t.Fatalf("jwt secret key marker = %q, want %q", c.JWTSecretKeyLastRotatedVersion, tt.wantStoredVersion)
}
secretChanged := c.JWTSecretKey != "leaked-secret"
if secretChanged != tt.wantSecretChanged {
t.Fatalf("secret changed = %v, want %v", secretChanged, tt.wantSecretChanged)
}
saved, err := os.ReadFile(file)
if err != nil {
t.Fatalf("read saved config: %v", err)
}
hasMarker := strings.Contains(string(saved), "jwt_secret_key_last_rotated_version")
if hasMarker != tt.wantSavedConfigKey {
t.Fatalf("saved marker present = %v, want %v, config = %s", hasMarker, tt.wantSavedConfigKey, saved)
}
})
}
}
// Mirrors the upstream single-block declaration so iota lines up exactly:
// ConfigUsePeerIP occupies iota=0 (as a typed string), ConfigCoverAll=1,
// ConfigCoverIgnoreAll=2. Pins persisted `cover` semantics.
const (
originalConfigUsePeerIP = "NZ::Use-Peer-IP"
originalConfigCoverAll = iota
originalConfigCoverIgnoreAll
)
func TestConfigCoverConstantValues(t *testing.T) {
if ConfigUsePeerIP != originalConfigUsePeerIP {
t.Fatalf("ConfigUsePeerIP = %q, want %q", ConfigUsePeerIP, originalConfigUsePeerIP)
}
if ConfigCoverAll != originalConfigCoverAll {
t.Fatalf("ConfigCoverAll = %d, want original value %d", ConfigCoverAll, originalConfigCoverAll)
}
if ConfigCoverIgnoreAll != originalConfigCoverIgnoreAll {
t.Fatalf("ConfigCoverIgnoreAll = %d, want original value %d", ConfigCoverIgnoreAll, originalConfigCoverIgnoreAll)
}
}
func newTempConfig(t *testing.T, cfg string) string {
t.Helper()
-1
View File
@@ -80,7 +80,6 @@ func execCase(t *testing.T, item testSt) {
CountryCode: "",
},
LastActive: time.Time{},
TaskStream: nil,
PrevTransferInSnapshot: 0,
PrevTransferOutSnapshot: 0,
}
+104 -2
View File
@@ -3,6 +3,7 @@ package model
import (
"log"
"slices"
"sync/atomic"
"time"
"github.com/goccy/go-json"
@@ -32,13 +33,68 @@ type Server struct {
GeoIP *GeoIP `gorm:"-" json:"geoip,omitempty"`
LastActive time.Time `gorm:"-" json:"last_active,omitempty"`
TaskStream pb.NezhaService_RequestTaskServer `gorm:"-" json:"-"`
// taskStream MUST be accessed only via SetTaskStream / GetTaskStream. Direct
// field access from outside this file races with the gRPC RequestTask
// handler that reassigns the stream on every reconnect — a torn read of the
// two-word interface header would panic on a subsequent .Send call. The
// atomic.Pointer + holder struct lets us swap the stream lock-free while
// every reader observes a single, consistent value.
taskStream atomic.Pointer[taskStreamHolder]
ConfigCache chan any `gorm:"-" json:"-"`
PrevTransferInSnapshot uint64 `gorm:"-" json:"-"` // 上次数据点时的入站使用量
PrevTransferOutSnapshot uint64 `gorm:"-" json:"-"` // 上次数据点时的出站使用量
}
// taskStreamHolder wraps the interface so atomic.Pointer (which requires a
// concrete pointed-to type) can publish it atomically. The previous bare
// field `TaskStream pb.NezhaService_RequestTaskServer` was a plain interface
// value: two words on the heap (type ptr + data ptr). Concurrent assignment
// produced torn reads detectable by `go test -race` and crashable in production.
type taskStreamHolder struct {
s pb.NezhaService_RequestTaskServer
}
// SetTaskStream publishes the agent's RequestTask stream so other goroutines
// can deliver tasks to the agent. Pass nil to detach (e.g. on disconnect).
func (s *Server) SetTaskStream(stream pb.NezhaService_RequestTaskServer) {
if stream == nil {
s.taskStream.Store(nil)
return
}
s.taskStream.Store(&taskStreamHolder{s: stream})
}
// ClearTaskStreamIfCurrent detaches stream only if it is still the published
// RequestTask stream. Disconnect cleanup uses this guard so an old stream
// returning after a reconnect cannot erase the newer live stream.
func (s *Server) ClearTaskStreamIfCurrent(stream pb.NezhaService_RequestTaskServer) bool {
if stream == nil {
return false
}
for {
h := s.taskStream.Load()
if h == nil || h.s != stream {
return false
}
if s.taskStream.CompareAndSwap(h, nil) {
return true
}
}
}
// GetTaskStream returns the currently-published stream, or nil if the agent
// is offline. Callers MUST capture the return into a local variable before
// using it — re-reading via GetTaskStream() across a Send call reopens the
// race we're trying to close.
func (s *Server) GetTaskStream() pb.NezhaService_RequestTaskServer {
h := s.taskStream.Load()
if h == nil {
return nil
}
return h.s
}
func InitServer(s *Server) {
s.Host = &Host{}
s.State = &HostState{}
@@ -51,7 +107,9 @@ func (s *Server) CopyFromRunningServer(old *Server) {
s.State = old.State
s.GeoIP = old.GeoIP
s.LastActive = old.LastActive
s.TaskStream = old.TaskStream
// taskStream is an atomic.Pointer; copy the published value rather than
// the field itself (atomic.Pointer is not safe to copy by value).
s.SetTaskStream(old.GetTaskStream())
s.ConfigCache = old.ConfigCache
s.PrevTransferInSnapshot = old.PrevTransferInSnapshot
s.PrevTransferOutSnapshot = old.PrevTransferOutSnapshot
@@ -73,6 +131,50 @@ func (s *Server) AfterFind(tx *gorm.DB) error {
return nil
}
// ServerOwnerInfo carries the user-facing identity for Server.UserID. It is
// returned by the lookup function installed by the singleton layer; model
// must not import singleton (cycle), so the dependency flows through a
// package-level function variable instead.
type ServerOwnerInfo struct {
ID uint64 `json:"id"`
Username string `json:"username,omitempty"`
}
// ServerOwnerLookup is installed by singleton at startup to resolve a
// Server.UserID into a display-ready owner record. Returns ok=false when
// the uid does not map to a known user; the caller renders that as an
// "unknown user" placeholder so deleted-user rows stay debuggable. Left nil
// in tests / headless contexts so the JSON simply omits the owner field.
var ServerOwnerLookup func(uid uint64) (ServerOwnerInfo, bool)
type serverJSON Server
type serverWithOwner struct {
*serverJSON
Owner *ServerOwnerInfo `json:"owner,omitempty"`
}
// MarshalJSON projects Server.UserID into a structured owner field on the
// wire. Server.UserID itself stays `json:"-"` (set on Common) so callers
// that do not need owner info pay nothing and members do not accidentally
// receive raw uid integers. The lookup function is consulted only when
// installed; if absent we still emit a minimal {id} record so clients can
// at least distinguish ownership, except for uid=0 which is the legacy
// global-secret pseudo-owner and is best surfaced as such by the caller's
// translation table on the frontend.
func (s *Server) MarshalJSON() ([]byte, error) {
owner := &ServerOwnerInfo{ID: s.GetUserID()}
if ServerOwnerLookup != nil {
if info, ok := ServerOwnerLookup(owner.ID); ok {
owner.Username = info.Username
}
}
return json.Marshal(serverWithOwner{
serverJSON: (*serverJSON)(s),
Owner: owner,
})
}
func (s *Server) SplitList(x []*Server) ([]*Server, []*Server) {
pri := func(s *Server) bool {
return s.DisplayIndex == 0
+129
View File
@@ -0,0 +1,129 @@
package model
import (
"encoding/json"
"testing"
)
// Server.MarshalJSON projects Server.UserID into a public owner field while
// keeping the raw UserID json-hidden. The lookup function is package-level
// and shared across tests; each subtest installs its own stub and restores
// the original to avoid leaking state.
func TestServerMarshalJSONOwnerProjection(t *testing.T) {
original := ServerOwnerLookup
t.Cleanup(func() { ServerOwnerLookup = original })
tests := []struct {
name string
uid uint64
lookup func(uid uint64) (ServerOwnerInfo, bool)
wantID uint64
wantHasName bool
wantName string
}{
{
// uid=0 is the legacy global agent secret pseudo-owner. The
// lookup deliberately returns ok=false so the frontend can
// render it as "Global Agent" instead of a real username.
name: "uid_zero_has_no_username",
uid: 0,
lookup: func(uint64) (ServerOwnerInfo, bool) {
return ServerOwnerInfo{}, false
},
wantID: 0,
wantHasName: false,
},
{
// Known user → username flows through to the wire so the
// admin frontend can show it without a separate /user fetch
// (which members cannot call anyway).
name: "known_user_has_username",
uid: 42,
lookup: func(uid uint64) (ServerOwnerInfo, bool) {
return ServerOwnerInfo{ID: uid, Username: "alice"}, true
},
wantID: 42,
wantHasName: true,
wantName: "alice",
},
{
// Deleted user → lookup returns ok=false. The wire still
// carries owner.id so the frontend can render an "Unknown
// user (#id)" placeholder; otherwise the row would silently
// appear ownerless and ops would lose the audit trail.
name: "deleted_user_keeps_id_without_username",
uid: 999,
lookup: func(uint64) (ServerOwnerInfo, bool) {
return ServerOwnerInfo{}, false
},
wantID: 999,
wantHasName: false,
},
}
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
ServerOwnerLookup = tc.lookup
s := &Server{Common: Common{ID: 7, UserID: tc.uid}, Name: "srv"}
raw, err := json.Marshal(s)
if err != nil {
t.Fatalf("marshal: %v", err)
}
var got struct {
Owner *ServerOwnerInfo `json:"owner"`
// Owner must never appear as the raw uid via Common.UserID;
// the Common.UserID json tag is "-" and a regression that
// flips it to "user_id" would expose internal owner ids
// to the wire bypassing the lookup-controlled rendering.
UserID *uint64 `json:"user_id,omitempty"`
}
if err := json.Unmarshal(raw, &got); err != nil {
t.Fatalf("unmarshal: %v", err)
}
if got.UserID != nil {
t.Fatalf("Common.UserID must not appear on the wire as user_id, got %d", *got.UserID)
}
if got.Owner == nil {
t.Fatalf("owner field must always be present, raw=%s", raw)
}
if got.Owner.ID != tc.wantID {
t.Fatalf("owner.id=%d, want %d", got.Owner.ID, tc.wantID)
}
if tc.wantHasName {
if got.Owner.Username != tc.wantName {
t.Fatalf("owner.username=%q, want %q", got.Owner.Username, tc.wantName)
}
} else if got.Owner.Username != "" {
t.Fatalf("owner.username must be omitted for uid=%d, got %q", tc.uid, got.Owner.Username)
}
})
}
}
// When no lookup is installed (tests / headless tools), MarshalJSON must
// still emit a minimal owner record so consumers do not crash on missing
// fields. Without this guard a future refactor could silently drop the
// owner key entirely whenever the hook is nil.
func TestServerMarshalJSONEmitsOwnerWithoutLookup(t *testing.T) {
original := ServerOwnerLookup
t.Cleanup(func() { ServerOwnerLookup = original })
ServerOwnerLookup = nil
raw, err := json.Marshal(&Server{Common: Common{ID: 1, UserID: 17}, Name: "srv"})
if err != nil {
t.Fatalf("marshal: %v", err)
}
var got struct {
Owner *ServerOwnerInfo `json:"owner"`
}
if err := json.Unmarshal(raw, &got); err != nil {
t.Fatalf("unmarshal: %v", err)
}
if got.Owner == nil || got.Owner.ID != 17 || got.Owner.Username != "" {
t.Fatalf("expected bare owner record {id:17}, got %+v", got.Owner)
}
}
+86
View File
@@ -0,0 +1,86 @@
package model
import (
"net/http/httptest"
"sync"
"testing"
"github.com/gin-gonic/gin"
)
// Server.UserID 在 server-transfer rotation 流程里会被 ServerTransfer 的
// Register/revertTransition 改写以反映新所有者,同时 authorizeAgentForUUID
// 在每次 agent RPC 里读取它。原实现两处都是裸字段访问,race detector 会
// 报告 data race;这是 review 评分 75 的真实问题。
//
// 修复后所有并发读写都走 SetUserID/GetUserID 的 atomic 包装,本测试在
// `go test -race` 下应该完全跑干净。
func TestServerUserIDConcurrentAccessIsRaceFree(t *testing.T) {
s := &Server{}
const (
writers = 4
readers = 8
rounds = 500
)
var wg sync.WaitGroup
wg.Add(writers + readers)
for i := 0; i < writers; i++ {
uid := uint64(i + 1)
go func() {
defer wg.Done()
for j := 0; j < rounds; j++ {
s.SetUserID(uid)
}
}()
}
for i := 0; i < readers; i++ {
go func() {
defer wg.Done()
for j := 0; j < rounds; j++ {
_ = s.GetUserID()
}
}()
}
wg.Wait()
}
// Common.HasPermission 是 server-transfer 旋转下与 SetUserID 并发的主要读者
// 之一:dashboard 各 controller 的 listHandler post-filter 在 transfer 窗口
// 内不断对同一 *Server 调用 HasPermission,而 Register/revertTransition 同
// 时通过 SetUserID 改写所属用户。原实现的 `user.ID == c.UserID` 是裸读,会
// 与 atomic.StoreUint64 形成 data racego test -race 必爆)。修复后改成走
// GetUserID() 走 atomic 协议。这个测试就是用来在 -race 下钉死该不变量的。
func TestCommonHasPermissionConcurrentWithSetUserIDIsRaceFree(t *testing.T) {
s := &Server{Common: Common{ID: 1}}
const (
writers = 4
readers = 8
rounds = 500
)
var wg sync.WaitGroup
wg.Add(writers + readers)
for i := 0; i < writers; i++ {
uid := uint64(i + 1)
go func() {
defer wg.Done()
for j := 0; j < rounds; j++ {
s.SetUserID(uid)
}
}()
}
for i := 0; i < readers; i++ {
go func() {
defer wg.Done()
ctx, _ := gin.CreateTestContext(httptest.NewRecorder())
ctx.Set(CtxKeyAuthorizedUser, &User{Common: Common{ID: 2}, Role: RoleMember})
for j := 0; j < rounds; j++ {
_ = s.HasPermission(ctx)
}
}()
}
wg.Wait()
}
+91
View File
@@ -0,0 +1,91 @@
package model
import (
"context"
"sync"
"testing"
pb "github.com/nezhahq/nezha/proto"
)
// raceProbeStream is the smallest fake of pb.NezhaService_RequestTaskServer
// the race probe needs. We only call Send on it from the test; the embedded
// interface satisfies the rest of the contract with nil-panicking methods we
// never invoke.
type raceProbeStream struct {
pb.NezhaService_RequestTaskServer
}
func (raceProbeStream) Send(*pb.Task) error { return nil }
func (raceProbeStream) Context() context.Context { return context.Background() }
// model.Server.TaskStream is read from many goroutines (singleton cron pushes,
// transfer ApplyConfig pushes, terminal/fm proxies, dashboard rpc keepalives,
// per-server batch pushes) and written from exactly one (the gRPC RequestTask
// goroutine on every fresh agent connection). The bare-field access pattern
// `if s.TaskStream != nil { s.TaskStream.Send(...) }` is a data race on the
// interface header (two-word value) and can torn-read into a panic on a
// reconnect. This test pins down "concurrent set + send must be race-free"
// using the Go race detector — without the fix, `go test -race` reports a
// data race on TaskStream; with the fix the field is encapsulated behind
// atomic methods and the test runs clean. Without `-race` both versions are
// indistinguishable, so this test is only meaningful under the race flag —
// run it from CI as `go test -race ./model/`.
func TestServerTaskStreamConcurrentAccessIsRaceFree(t *testing.T) {
s := &Server{}
InitServer(s)
const (
writers = 4
readers = 8
rounds = 200
)
var wg sync.WaitGroup
wg.Add(writers + readers)
for i := 0; i < writers; i++ {
go func() {
defer wg.Done()
for j := 0; j < rounds; j++ {
s.SetTaskStream(raceProbeStream{})
s.SetTaskStream(nil)
}
}()
}
for i := 0; i < readers; i++ {
go func() {
defer wg.Done()
for j := 0; j < rounds; j++ {
if stream := s.GetTaskStream(); stream != nil {
_ = stream.Send(nil)
}
}
}()
}
wg.Wait()
}
func TestServerClearTaskStreamIfCurrentClearsOnlyMatchingStream(t *testing.T) {
s := &Server{}
InitServer(s)
first := &raceProbeStream{}
second := &raceProbeStream{}
s.SetTaskStream(first)
if !s.ClearTaskStreamIfCurrent(first) {
t.Fatal("matching current stream must be cleared")
}
if got := s.GetTaskStream(); got != nil {
t.Fatalf("expected cleared task stream, got %T", got)
}
s.SetTaskStream(first)
s.SetTaskStream(second)
if s.ClearTaskStreamIfCurrent(first) {
t.Fatal("stale stream cleanup must not clear a newer stream")
}
if got := s.GetTaskStream(); got != second {
t.Fatalf("expected newer stream to remain published, got %T", got)
}
}
+128
View File
@@ -0,0 +1,128 @@
package model
import (
"time"
"github.com/gin-gonic/gin"
)
// ServerTransferStatus represents the lifecycle state of a server ownership
// transfer. A transfer's life starts at Pending (server.user_id has been
// flipped to the new owner; agent still authenticates with the old owner's
// AgentSecret) and ends in exactly one of the terminal states.
type ServerTransferStatus uint8
const (
// ServerTransferStatusPending means the dashboard has flipped Server.UserID
// to the new owner and queued an ApplyConfig task to swap the agent's
// client_secret. Auth still accepts the old owner's AgentSecret for this
// UUID until verification arrives or the transfer times out.
ServerTransferStatusPending ServerTransferStatus = iota
// ServerTransferStatusVerified means the agent successfully reconnected
// using the new owner's AgentSecret. Auth no longer tolerates the old
// owner's secret on this UUID.
ServerTransferStatusVerified
// ServerTransferStatusFailed means the agent explicitly reported the
// ApplyConfig task as unsuccessful (e.g. DisableCommandExecute). The
// dashboard has rolled Server.UserID back to FromUserID.
ServerTransferStatusFailed
// ServerTransferStatusTimeout means the verification window expired
// without the agent reconnecting under the new secret. The dashboard has
// rolled Server.UserID back to FromUserID.
ServerTransferStatusTimeout
// ServerTransferStatusCancelled means an administrator cancelled the
// transfer before any verification event was observed. The dashboard has
// rolled Server.UserID back to FromUserID.
ServerTransferStatusCancelled
)
// IsTerminal reports whether the status represents a settled transfer. Only
// terminal transfers are eligible for retry and they will never be in the
// pending index.
func (s ServerTransferStatus) IsTerminal() bool {
return s != ServerTransferStatusPending
}
// ServerTransfer records a single attempt to transfer ownership of one server
// to another user. It is the source of truth for the auth-tolerance window
// during a transfer — service/rpc.authorizeAgentForUUID consults the pending
// index built from this table to decide whether to accept the old owner's
// AgentSecret on the affected UUID.
//
// Naming note: the existing model.Transfer records hourly traffic snapshots
// and is unrelated. This entity is named ServerTransfer to disambiguate.
type ServerTransfer struct {
Common
ServerID uint64 `json:"server_id" gorm:"index"`
FromUserID uint64 `json:"from_user_id"`
ToUserID uint64 `json:"to_user_id"`
InitiatorID uint64 `json:"initiator_id"`
Status ServerTransferStatus `json:"status" gorm:"index"`
LastError string `json:"last_error,omitempty"`
AckedAt *time.Time `json:"acked_at,omitempty"`
// HandshakeSecret is a per-transfer random credential that PushIfOnline
// delivers in place of the destination user's global AgentSecret. The
// agent treats it as a temporary handshake token: it rotates to this
// secret on the 10s reload, reconnects, and the dashboard's auth path
// recognises it as proof of transfer delivery (MarkVerified). It is
// scoped to this single transfer and to this single UUID — leaking it
// to the previous owner who hijacks the stream still does NOT expose
// the destination user's other agents. Never returned to API clients.
HandshakeSecret string `json:"-" gorm:"type:char(32)"`
// RevertHandshakeSecret is the same idea for the rollback path: when
// the dashboard pushes a revert ApplyConfig over a stream now held by
// the destination user, we must not embed the source user's global
// AgentSecret. Instead the agent rotates back through this token, which
// is recognised by the auth path during the revert window only.
RevertHandshakeSecret string `json:"-" gorm:"type:char(32)"`
}
// HasPermission overrides Common.HasPermission so a transfer is visible to
// admins, the source user, the destination user, and the initiator. Listing
// uses this to filter what the caller can see; mutating endpoints (cancel,
// retry) layer additional checks on top.
func (t *ServerTransfer) HasPermission(ctx *gin.Context) bool {
auth, ok := ctx.Get(CtxKeyAuthorizedUser)
if !ok {
return false
}
user := *auth.(*User)
if user.Role == RoleAdmin {
return true
}
return user.ID == t.FromUserID || user.ID == t.ToUserID || user.ID == t.InitiatorID
}
// BatchMoveServerResultStatus is the per-server outcome returned by the
// batch-move endpoint. It maps to TransferStatus for transfers that were
// successfully created, plus extra synchronous-failure modes (permission,
// duplicate active transfer, missing server) that never produce a row.
type BatchMoveServerResultStatus string
const (
// BatchMoveServerResultPending: ServerTransfer row created, agent push
// in progress. Callers should watch the WS for terminal status.
BatchMoveServerResultPending BatchMoveServerResultStatus = "pending"
// BatchMoveServerResultPermissionDenied: caller cannot move this server.
BatchMoveServerResultPermissionDenied BatchMoveServerResultStatus = "permission_denied"
// BatchMoveServerResultAlreadyTransferring: server already has an in-flight
// ServerTransfer row, cancel or wait first.
BatchMoveServerResultAlreadyTransferring BatchMoveServerResultStatus = "already_transferring"
// BatchMoveServerResultServerNotFound: server id does not exist.
BatchMoveServerResultServerNotFound BatchMoveServerResultStatus = "server_not_found"
// BatchMoveServerResultSameOwner: target user already owns this server.
BatchMoveServerResultSameOwner BatchMoveServerResultStatus = "same_owner"
// BatchMoveServerResultAgentTooOld: agent build does not understand
// TaskTypeServerTransferApply, so the rotation would never complete and
// dashboard refuses to start it. Operator must upgrade the agent.
BatchMoveServerResultAgentTooOld BatchMoveServerResultStatus = "agent_too_old"
)
// BatchMoveServerResult is one entry in the batchMoveServer response, one
// per requested server id, in the same order.
type BatchMoveServerResult struct {
ServerID uint64 `json:"server_id"`
Status BatchMoveServerResultStatus `json:"status"`
TransferID uint64 `json:"transfer_id,omitempty"`
Error string `json:"error,omitempty"`
}
+6 -1
View File
@@ -26,6 +26,10 @@ const (
TaskTypeFM
TaskTypeReportConfig
TaskTypeApplyConfig
// TaskTypeServerTransferApply: per-transfer credential rotation.
// Pre-transfer agents do not recognise this type — dashboard MUST gate
// transfers on agent capability before pushing.
TaskTypeServerTransferApply
)
type TerminalTask struct {
@@ -133,7 +137,8 @@ func IsServiceSentinelNeeded(t uint64) bool {
switch t {
case TaskTypeCommand, TaskTypeTerminalGRPC, TaskTypeUpgrade,
TaskTypeKeepalive, TaskTypeNAT, TaskTypeFM,
TaskTypeReportConfig, TaskTypeApplyConfig:
TaskTypeReportConfig, TaskTypeApplyConfig,
TaskTypeServerTransferApply:
return false
default:
return true
+1
View File
@@ -32,6 +32,7 @@ type User struct {
type UserInfo struct {
Role Role
Username string
AgentSecret string
}
+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)
}
}
+169 -11
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
}
// 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())
}
}
}
}