mirror of
https://github.com/Buriburizaem0n/nezha_domains.git
synced 2026-09-19 09:40:12 +00:00
feat: server transfer rotation
This commit is contained in:
@@ -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")))
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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}`))
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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"`
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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()
|
||||
}
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
@@ -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
|
||||
// race(TestCommonHasPermissionConcurrentWithSetUserIDIsRaceFree 在
|
||||
// -race 下钉死该不变量),并且在 transfer 切换瞬间可能给出错误的权限
|
||||
// 判断。
|
||||
return user.ID == c.GetUserID()
|
||||
}
|
||||
|
||||
type CommonInterface interface {
|
||||
|
||||
+65
-1
@@ -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{
|
||||
|
||||
@@ -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()
|
||||
|
||||
|
||||
@@ -80,7 +80,6 @@ func execCase(t *testing.T, item testSt) {
|
||||
CountryCode: "",
|
||||
},
|
||||
LastActive: time.Time{},
|
||||
TaskStream: nil,
|
||||
PrevTransferInSnapshot: 0,
|
||||
PrevTransferOutSnapshot: 0,
|
||||
}
|
||||
|
||||
+104
-2
@@ -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
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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 race(go 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()
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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
@@ -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
|
||||
|
||||
@@ -32,6 +32,7 @@ type User struct {
|
||||
|
||||
type UserInfo struct {
|
||||
Role Role
|
||||
Username string
|
||||
AgentSecret string
|
||||
}
|
||||
|
||||
|
||||
@@ -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
@@ -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
@@ -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")
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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
@@ -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{
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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 {
|
||||
|
||||
@@ -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
@@ -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 {
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
|
||||
@@ -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())
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user