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"
|
"log"
|
||||||
"net/http"
|
"net/http"
|
||||||
"os"
|
"os"
|
||||||
"path"
|
|
||||||
"regexp"
|
"regexp"
|
||||||
"slices"
|
"slices"
|
||||||
"strings"
|
"strings"
|
||||||
@@ -118,6 +117,11 @@ func routers(r *gin.Engine, frontendDist fs.FS) {
|
|||||||
auth.POST("/batch-move/server", commonHandler(batchMoveServer))
|
auth.POST("/batch-move/server", commonHandler(batchMoveServer))
|
||||||
auth.POST("/force-update/server", commonHandler(forceUpdateServer))
|
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.GET("/notification", listHandler(listNotification))
|
||||||
auth.POST("/notification", commonHandler(createNotification))
|
auth.POST("/notification", commonHandler(createNotification))
|
||||||
auth.PATCH("/notification/:id", commonHandler(updateNotification))
|
auth.PATCH("/notification/:id", commonHandler(updateNotification))
|
||||||
@@ -320,27 +324,52 @@ func getUid(c *gin.Context) uint64 {
|
|||||||
}
|
}
|
||||||
|
|
||||||
func fallbackToFrontend(frontendDist fs.FS) func(*gin.Context) {
|
func fallbackToFrontend(frontendDist fs.FS) func(*gin.Context) {
|
||||||
checkLocalFileOrFs := func(c *gin.Context, fs fs.FS, path string, customStatusCode int) bool {
|
serveFile := func(c *gin.Context, name string, file fs.File, customStatusCode int) bool {
|
||||||
if _, err := os.Stat(path); err == nil {
|
defer file.Close()
|
||||||
http.ServeFile(utils.NewGinCustomWriter(c, customStatusCode), c.Request, path)
|
fileStat, err := file.Stat()
|
||||||
return true
|
|
||||||
}
|
|
||||||
f, err := fs.Open(path)
|
|
||||||
if err != nil {
|
|
||||||
return false
|
|
||||||
}
|
|
||||||
defer f.Close()
|
|
||||||
fileStat, err := f.Stat()
|
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return false
|
return false
|
||||||
}
|
}
|
||||||
if fileStat.IsDir() {
|
if fileStat.IsDir() {
|
||||||
return false
|
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
|
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{
|
frontendPageUrlRegistry := []*regexp.Regexp{
|
||||||
// official user frontend
|
// official user frontend
|
||||||
regexp.MustCompile(`^/$`),
|
regexp.MustCompile(`^/$`),
|
||||||
@@ -361,6 +390,11 @@ func fallbackToFrontend(frontendDist fs.FS) func(*gin.Context) {
|
|||||||
regexp.MustCompile(`^/dashboard/settings/user$`),
|
regexp.MustCompile(`^/dashboard/settings/user$`),
|
||||||
regexp.MustCompile(`^/dashboard/settings/online-user$`),
|
regexp.MustCompile(`^/dashboard/settings/online-user$`),
|
||||||
regexp.MustCompile(`^/dashboard/settings/waf$`),
|
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 {
|
getFallbackStatusCode := func(path string) int {
|
||||||
@@ -385,22 +419,22 @@ func fallbackToFrontend(frontendDist fs.FS) func(*gin.Context) {
|
|||||||
}
|
}
|
||||||
|
|
||||||
fallbackStatusCode := getFallbackStatusCode(c.Request.URL.Path)
|
fallbackStatusCode := getFallbackStatusCode(c.Request.URL.Path)
|
||||||
if strings.HasPrefix(c.Request.URL.Path, "/dashboard") {
|
// Only /dashboard/ belongs to the admin frontend; /dashboard.. must not be trimmed into ../.
|
||||||
stripPath := strings.TrimPrefix(c.Request.URL.Path, "/dashboard")
|
if strings.HasPrefix(c.Request.URL.Path, "/dashboard/") {
|
||||||
localFilePath := path.Join(singleton.Conf.AdminTemplate, stripPath)
|
stripPath := strings.TrimPrefix(c.Request.URL.Path, "/dashboard/")
|
||||||
if checkLocalFileOrFs(c, frontendDist, localFilePath, http.StatusOK) {
|
if checkLocalFileOrFs(c, frontendDist, singleton.Conf.AdminTemplate, stripPath, http.StatusOK) {
|
||||||
return
|
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")))
|
c.JSON(http.StatusNotFound, newErrorResponse(errors.New("404 Not Found")))
|
||||||
}
|
}
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
localFilePath := path.Join(singleton.Conf.UserTemplate, c.Request.URL.Path)
|
stripPath := strings.TrimPrefix(c.Request.URL.Path, "/")
|
||||||
if checkLocalFileOrFs(c, frontendDist, localFilePath, http.StatusOK) {
|
if checkLocalFileOrFs(c, frontendDist, singleton.Conf.UserTemplate, stripPath, http.StatusOK) {
|
||||||
return
|
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")))
|
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)
|
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")
|
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{
|
fmData, _ := json.Marshal(&model.TaskFM{
|
||||||
StreamID: streamId,
|
StreamID: streamId,
|
||||||
})
|
})
|
||||||
if err := server.TaskStream.Send(&proto.Task{
|
if err := stream.Send(&proto.Task{
|
||||||
Type: model.TaskTypeFM,
|
Type: model.TaskTypeFM,
|
||||||
Data: string(fmData),
|
Data: string(fmData),
|
||||||
}); err != nil {
|
}); err != nil {
|
||||||
|
|||||||
@@ -56,7 +56,7 @@ func setupServerOwnershipFixture(t *testing.T) (stream *fakeTaskStream, reset fu
|
|||||||
|
|
||||||
alice, _ := singleton.ServerShared.Get(1)
|
alice, _ := singleton.ServerShared.Get(1)
|
||||||
stream = &fakeTaskStream{}
|
stream = &fakeTaskStream{}
|
||||||
alice.TaskStream = stream
|
alice.SetTaskStream(stream)
|
||||||
|
|
||||||
return stream, func() {
|
return stream, func() {
|
||||||
singleton.DB = originalDB
|
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)
|
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) {
|
func TestNATRejectsUnknownServerID(t *testing.T) {
|
||||||
ctx := newMemberValidationContext(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}`))
|
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
|
package controller
|
||||||
|
|
||||||
import (
|
import (
|
||||||
|
"errors"
|
||||||
"slices"
|
"slices"
|
||||||
"strconv"
|
"strconv"
|
||||||
"sync"
|
"sync"
|
||||||
@@ -153,6 +154,14 @@ func batchDeleteServer(c *gin.Context) (any, error) {
|
|||||||
singleton.DB.Unscoped().Delete(&model.Transfer{}, "server_id in (?)", servers)
|
singleton.DB.Unscoped().Delete(&model.Transfer{}, "server_id in (?)", servers)
|
||||||
singleton.AlertsLock.Unlock()
|
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)
|
singleton.ServerShared.Delete(servers)
|
||||||
return nil, nil
|
return nil, nil
|
||||||
}
|
}
|
||||||
@@ -187,8 +196,8 @@ func forceUpdateServer(c *gin.Context) (*model.ServerTaskResponse, error) {
|
|||||||
forceUpdateResp.Offline = append(forceUpdateResp.Offline, sid)
|
forceUpdateResp.Offline = append(forceUpdateResp.Offline, sid)
|
||||||
continue
|
continue
|
||||||
}
|
}
|
||||||
if server.TaskStream != nil {
|
if stream := server.GetTaskStream(); stream != nil {
|
||||||
if err := server.TaskStream.Send(&pb.Task{
|
if err := stream.Send(&pb.Task{
|
||||||
Type: model.TaskTypeUpgrade,
|
Type: model.TaskTypeUpgrade,
|
||||||
}); err != nil {
|
}); err != nil {
|
||||||
forceUpdateResp.Failure = append(forceUpdateResp.Failure, sid)
|
forceUpdateResp.Failure = append(forceUpdateResp.Failure, sid)
|
||||||
@@ -220,7 +229,11 @@ func getServerConfig(c *gin.Context) (string, error) {
|
|||||||
}
|
}
|
||||||
|
|
||||||
s, ok := singleton.ServerShared.Get(id)
|
s, ok := singleton.ServerShared.Get(id)
|
||||||
if !ok || s.TaskStream == nil {
|
if !ok {
|
||||||
|
return "", nil
|
||||||
|
}
|
||||||
|
stream := s.GetTaskStream()
|
||||||
|
if stream == nil {
|
||||||
return "", nil
|
return "", nil
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -228,7 +241,7 @@ func getServerConfig(c *gin.Context) (string, error) {
|
|||||||
return "", singleton.Localizer.ErrorT("permission denied")
|
return "", singleton.Localizer.ErrorT("permission denied")
|
||||||
}
|
}
|
||||||
|
|
||||||
if err := s.TaskStream.Send(&pb.Task{
|
if err := stream.Send(&pb.Task{
|
||||||
Type: model.TaskTypeReportConfig,
|
Type: model.TaskTypeReportConfig,
|
||||||
}); err != nil {
|
}); err != nil {
|
||||||
return "", err
|
return "", err
|
||||||
@@ -276,7 +289,7 @@ func setServerConfig(c *gin.Context) (*model.ServerTaskResponse, error) {
|
|||||||
if !s.HasPermission(c) {
|
if !s.HasPermission(c) {
|
||||||
return nil, singleton.Localizer.ErrorT("permission denied")
|
return nil, singleton.Localizer.ErrorT("permission denied")
|
||||||
}
|
}
|
||||||
if s.TaskStream == nil {
|
if s.GetTaskStream() == nil {
|
||||||
resp.Offline = append(resp.Offline, s.ID)
|
resp.Offline = append(resp.Offline, s.ID)
|
||||||
continue
|
continue
|
||||||
}
|
}
|
||||||
@@ -300,7 +313,14 @@ func setServerConfig(c *gin.Context) (*model.ServerTaskResponse, error) {
|
|||||||
Type: model.TaskTypeApplyConfig,
|
Type: model.TaskTypeApplyConfig,
|
||||||
Data: configForm.Config,
|
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()
|
respMu.Lock()
|
||||||
resp.Failure = append(resp.Failure, s.ID)
|
resp.Failure = append(resp.Failure, s.ID)
|
||||||
respMu.Unlock()
|
respMu.Unlock()
|
||||||
@@ -321,23 +341,29 @@ func setServerConfig(c *gin.Context) (*model.ServerTaskResponse, error) {
|
|||||||
// @Summary Batch move servers to other user
|
// @Summary Batch move servers to other user
|
||||||
// @Security BearerAuth
|
// @Security BearerAuth
|
||||||
// @Schemes
|
// @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
|
// @Tags auth required
|
||||||
// @Accept json
|
// @Accept json
|
||||||
// @Param request body model.BatchMoveServerForm true "BatchMoveServerForm"
|
// @Param request body model.BatchMoveServerForm true "BatchMoveServerForm"
|
||||||
// @Produce json
|
// @Produce json
|
||||||
// @Success 200 {object} model.CommonResponse[any]
|
// @Success 200 {object} model.CommonResponse[[]model.BatchMoveServerResult]
|
||||||
// @Router /batch-move/server [post]
|
// @Router /batch-move/server [post]
|
||||||
func batchMoveServer(c *gin.Context) (any, error) {
|
func batchMoveServer(c *gin.Context) ([]model.BatchMoveServerResult, error) {
|
||||||
var moveForm model.BatchMoveServerForm
|
var moveForm model.BatchMoveServerForm
|
||||||
if err := c.ShouldBindJSON(&moveForm); err != nil {
|
if err := c.ShouldBindJSON(&moveForm); err != nil {
|
||||||
return nil, err
|
return nil, err
|
||||||
}
|
}
|
||||||
|
|
||||||
if !singleton.ServerShared.CheckPermission(c, slices.Values(moveForm.Ids)) {
|
|
||||||
return nil, singleton.Localizer.ErrorT("permission denied")
|
|
||||||
}
|
|
||||||
|
|
||||||
if moveForm.ToUser == 0 {
|
if moveForm.ToUser == 0 {
|
||||||
return nil, singleton.Localizer.ErrorT("user id is required")
|
return nil, singleton.Localizer.ErrorT("user id is required")
|
||||||
}
|
}
|
||||||
@@ -347,35 +373,81 @@ func batchMoveServer(c *gin.Context) (any, error) {
|
|||||||
}
|
}
|
||||||
|
|
||||||
singleton.UserLock.RLock()
|
singleton.UserLock.RLock()
|
||||||
defer singleton.UserLock.RUnlock()
|
_, toUserExists := singleton.UserInfoMap[moveForm.ToUser]
|
||||||
if _, ok := singleton.UserInfoMap[moveForm.ToUser]; !ok {
|
singleton.UserLock.RUnlock()
|
||||||
|
if !toUserExists {
|
||||||
return nil, singleton.Localizer.ErrorT("user id %d does not exist", moveForm.ToUser)
|
return nil, singleton.Localizer.ErrorT("user id %d does not exist", moveForm.ToUser)
|
||||||
}
|
}
|
||||||
|
|
||||||
err := singleton.DB.Transaction(func(tx *gorm.DB) error {
|
results := make([]model.BatchMoveServerResult, 0, len(moveForm.Ids))
|
||||||
if err := tx.Model(&model.Server{}).Where("id in (?)", moveForm.Ids).Update("user_id", moveForm.ToUser).Error; err != nil {
|
uid := getUid(c)
|
||||||
return err
|
isAdmin := callerIsAdmin(c)
|
||||||
}
|
|
||||||
return nil
|
|
||||||
})
|
|
||||||
|
|
||||||
if err != nil {
|
for _, sid := range moveForm.Ids {
|
||||||
return nil, newGormError("%v", err)
|
res := model.BatchMoveServerResult{ServerID: sid}
|
||||||
}
|
|
||||||
|
|
||||||
idsMap := make(map[uint64]bool)
|
srv, ok := singleton.ServerShared.Get(sid)
|
||||||
for _, id := range moveForm.Ids {
|
if !ok || srv == nil {
|
||||||
idsMap[id] = true
|
res.Status = model.BatchMoveServerResultServerNotFound
|
||||||
}
|
results = append(results, res)
|
||||||
|
|
||||||
for _, s := range singleton.ServerShared.Range {
|
|
||||||
if s == nil || !idsMap[s.ID] {
|
|
||||||
continue
|
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
|
||||||
|
}
|
||||||
|
|
||||||
|
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 nil, nil
|
return results, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
var serverMetricMap = map[string]tsdb.MetricType{
|
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.
|
// 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
|
// The frontend never reads it, so HttpOnly is safe to enable and shuts the
|
||||||
// door on XSS attempting to steal the state.
|
// 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)
|
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")
|
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{
|
terminalData, _ := json.Marshal(&model.TerminalTask{
|
||||||
StreamID: streamId,
|
StreamID: streamId,
|
||||||
})
|
})
|
||||||
if err := server.TaskStream.Send(&proto.Task{
|
if err := stream.Send(&proto.Task{
|
||||||
Type: model.TaskTypeTerminalGRPC,
|
Type: model.TaskTypeTerminalGRPC,
|
||||||
Data: string(terminalData),
|
Data: string(terminalData),
|
||||||
}); err != nil {
|
}); 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 {
|
func filterServersForViewer(servers []*model.Server, viewerUserID uint64, viewerIsAdmin bool, withPublicNote bool) []model.StreamServer {
|
||||||
out := make([]model.StreamServer, 0, len(servers))
|
out := make([]model.StreamServer, 0, len(servers))
|
||||||
for _, server := range servers {
|
for _, server := range servers {
|
||||||
isOwnerOrAdmin := viewerIsAdmin || (viewerUserID != 0 && server.UserID == viewerUserID)
|
isOwnerOrAdmin := viewerIsAdmin || (viewerUserID != 0 && server.GetUserID() == viewerUserID)
|
||||||
if server.HideForGuest && !isOwnerOrAdmin {
|
if server.HideForGuest && !isOwnerOrAdmin {
|
||||||
continue
|
continue
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -24,6 +24,10 @@ import (
|
|||||||
func ServeRPC() *grpc.Server {
|
func ServeRPC() *grpc.Server {
|
||||||
server := grpc.NewServer(grpc.ChainUnaryInterceptor(getRealIp, waf))
|
server := grpc.NewServer(grpc.ChainUnaryInterceptor(getRealIp, waf))
|
||||||
rpcService.NezhaHandlerSingleton = rpcService.NewNezhaHandler()
|
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)
|
proto.RegisterNezhaServiceServer(server, rpcService.NezhaHandlerSingleton)
|
||||||
return server
|
return server
|
||||||
}
|
}
|
||||||
@@ -89,22 +93,30 @@ func DispatchTask(serviceSentinelDispatchBus <-chan *model.Service) {
|
|||||||
}
|
}
|
||||||
|
|
||||||
server, _ := singleton.ServerShared.Get(id)
|
server, _ := singleton.ServerShared.Get(id)
|
||||||
if server == nil || server.TaskStream == nil {
|
if server == nil {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
stream := server.GetTaskStream()
|
||||||
|
if stream == nil {
|
||||||
continue
|
continue
|
||||||
}
|
}
|
||||||
|
|
||||||
if canSendTaskToServer(task, server) {
|
if canSendTaskToServer(task, server) {
|
||||||
server.TaskStream.Send(task.PB())
|
stream.Send(task.PB())
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
case model.ServiceCoverAll:
|
case model.ServiceCoverAll:
|
||||||
for id, server := range singleton.ServerShared.Range {
|
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
|
continue
|
||||||
}
|
}
|
||||||
|
|
||||||
if canSendTaskToServer(task, server) {
|
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() {
|
singleton.CronShared.AddFunc("@every 20s", func() {
|
||||||
list := singleton.ServerShared.GetSortedList()
|
list := singleton.ServerShared.GetSortedList()
|
||||||
for _, s := range list {
|
for _, s := range list {
|
||||||
if s == nil || s.TaskStream == nil {
|
if s == nil {
|
||||||
continue
|
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) {
|
func ServeNAT(w http.ResponseWriter, r *http.Request, natConfig *model.NAT) {
|
||||||
server, _ := singleton.ServerShared.Get(natConfig.ServerID)
|
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.WriteHeader(http.StatusServiceUnavailable)
|
||||||
w.Write([]byte("server not found or not connected"))
|
w.Write([]byte("server not found or not connected"))
|
||||||
return
|
return
|
||||||
@@ -157,7 +179,7 @@ func ServeNAT(w http.ResponseWriter, r *http.Request, natConfig *model.NAT) {
|
|||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
|
||||||
if err := server.TaskStream.Send(&proto.Task{
|
if err := stream.Send(&proto.Task{
|
||||||
Type: model.TaskTypeNAT,
|
Type: model.TaskTypeNAT,
|
||||||
Data: string(taskData),
|
Data: string(taskData),
|
||||||
}); err != nil {
|
}); err != nil {
|
||||||
@@ -192,5 +214,5 @@ func canSendTaskToServer(task *model.Service, server *model.Server) bool {
|
|||||||
}
|
}
|
||||||
singleton.UserLock.RUnlock()
|
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/gin-swagger v1.6.1
|
||||||
github.com/swaggo/swag v1.16.6
|
github.com/swaggo/swag v1.16.6
|
||||||
github.com/tidwall/gjson v1.19.0
|
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/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/oauth2 v0.36.0
|
||||||
golang.org/x/sync v0.20.0
|
golang.org/x/sync v0.20.0
|
||||||
google.golang.org/grpc v1.81.1
|
google.golang.org/grpc v1.81.1
|
||||||
@@ -109,7 +109,7 @@ require (
|
|||||||
go.yaml.in/yaml/v3 v3.0.4 // indirect
|
go.yaml.in/yaml/v3 v3.0.4 // indirect
|
||||||
golang.org/x/arch v0.27.0 // indirect
|
golang.org/x/arch v0.27.0 // indirect
|
||||||
golang.org/x/mod v0.36.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/text v0.37.0 // indirect
|
||||||
golang.org/x/time v0.15.0 // indirect
|
golang.org/x/time v0.15.0 // indirect
|
||||||
golang.org/x/tools v0.45.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/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-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.0.0-20210921155107-089bfa567519/go.mod h1:GvvjBRRGRdwPK5ydBHafDWAxML/pGHZbMvKqRZ5+Abc=
|
||||||
golang.org/x/crypto v0.51.0 h1:IBPXwPfKxY7cWQZ38ZCIRPI50YLeevDLlLnyC5wRGTI=
|
golang.org/x/crypto v0.52.0 h1:RMs7fP2rXdep0CftQlK8Uf+kibLm7qkCcradZWYz988=
|
||||||
golang.org/x/crypto v0.51.0/go.mod h1:8AdwkbraGNABw2kOX6YFPs3WM22XqI4EXEd8g+x7Oc8=
|
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 h1:+3jdDGGB8NGb1Zktc737jlt3/A5f6UlwSzmvqUuufxw=
|
||||||
golang.org/x/exp v0.0.0-20260508232706-74f9aab9d74a/go.mod h1:d2fgXJLVs4dYDHUk5lwMIfzRzSrWCfGZb0ZqeLa/Vcw=
|
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=
|
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-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.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.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.55.0 h1:bcvxaJn3e1U6InsFWt1JUq1aSjnRxLzT2rtD2KfkDF8=
|
||||||
golang.org/x/net v0.54.0/go.mod h1:Sj4oj8jK6XmHpBZU/zWHw3BV3abl4Kvi+Ut7cQcY+cQ=
|
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 h1:peZ/1z27fi9hUOFCAZaHyrpWG5lwe0RJEEEeH0ThlIs=
|
||||||
golang.org/x/oauth2 v0.36.0/go.mod h1:YDBUJMTkDnJS+A4BP4eZBjCqtokkg1hODuPjwiGPO7Q=
|
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=
|
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-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.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.5.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
|
||||||
golang.org/x/sys v0.44.0 h1:ildZl3J4uzeKP07r2F++Op7E9B29JRUy+a27EibtBTQ=
|
golang.org/x/sys v0.45.0 h1:dO4czNzziLiiXplLQgBCEpCvXQ3dnkn0SdaZSYdQ+FY=
|
||||||
golang.org/x/sys v0.44.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw=
|
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-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.0.0-20210927222741-03fcf44c2211/go.mod h1:jbD1KX2456YbFQfuXm/mYQcufACuNUgVhRMnK/tPxf8=
|
||||||
golang.org/x/term v0.5.0/go.mod h1:jMB1sMXY+tzblOD4FWmEbocvup2/aLOaQEp7JmGp78k=
|
golang.org/x/term v0.5.0/go.mod h1:jMB1sMXY+tzblOD4FWmEbocvup2/aLOaQEp7JmGp78k=
|
||||||
|
|||||||
+23
-2
@@ -6,6 +6,7 @@ import (
|
|||||||
"slices"
|
"slices"
|
||||||
"strconv"
|
"strconv"
|
||||||
"strings"
|
"strings"
|
||||||
|
"sync/atomic"
|
||||||
"time"
|
"time"
|
||||||
|
|
||||||
"github.com/gin-gonic/gin"
|
"github.com/gin-gonic/gin"
|
||||||
@@ -37,8 +38,21 @@ func (c *Common) GetID() uint64 {
|
|||||||
return c.ID
|
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 {
|
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 {
|
func (c *Common) HasPermission(ctx *gin.Context) bool {
|
||||||
@@ -52,7 +66,14 @@ func (c *Common) HasPermission(ctx *gin.Context) bool {
|
|||||||
return true
|
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 {
|
type CommonInterface interface {
|
||||||
|
|||||||
+69
-5
@@ -3,6 +3,7 @@ package model
|
|||||||
import (
|
import (
|
||||||
"os"
|
"os"
|
||||||
"path/filepath"
|
"path/filepath"
|
||||||
|
"strconv"
|
||||||
"strings"
|
"strings"
|
||||||
|
|
||||||
"github.com/go-viper/mapstructure/v2"
|
"github.com/go-viper/mapstructure/v2"
|
||||||
@@ -16,8 +17,12 @@ import (
|
|||||||
)
|
)
|
||||||
|
|
||||||
const (
|
const (
|
||||||
ConfigUsePeerIP = "NZ::Use-Peer-IP"
|
ConfigUsePeerIP = "NZ::Use-Peer-IP"
|
||||||
ConfigCoverAll = iota
|
JWTSecretKeyRotationBaselineVersion = "v2.0.13"
|
||||||
|
)
|
||||||
|
|
||||||
|
const (
|
||||||
|
ConfigCoverAll = iota + 1
|
||||||
ConfigCoverIgnoreAll
|
ConfigCoverIgnoreAll
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -60,9 +65,10 @@ type Config struct {
|
|||||||
AgentSecretKey string `koanf:"agent_secret_key" json:"agent_secret_key,omitempty"`
|
AgentSecretKey string `koanf:"agent_secret_key" json:"agent_secret_key,omitempty"`
|
||||||
JWTTimeout int `koanf:"jwt_timeout" json:"jwt_timeout,omitempty"` // JWT token过期时间(小时)
|
JWTTimeout int `koanf:"jwt_timeout" json:"jwt_timeout,omitempty"` // JWT token过期时间(小时)
|
||||||
|
|
||||||
JWTSecretKey string `koanf:"jwt_secret_key" json:"jwt_secret_key,omitempty"`
|
JWTSecretKey string `koanf:"jwt_secret_key" json:"jwt_secret_key,omitempty"`
|
||||||
ListenPort uint16 `koanf:"listen_port" json:"listen_port,omitempty"`
|
JWTSecretKeyLastRotatedVersion string `koanf:"jwt_secret_key_last_rotated_version" json:"jwt_secret_key_last_rotated_version,omitempty"`
|
||||||
ListenHost string `koanf:"listen_host" json:"listen_host,omitempty"`
|
ListenPort uint16 `koanf:"listen_port" json:"listen_port,omitempty"`
|
||||||
|
ListenHost string `koanf:"listen_host" json:"listen_host,omitempty"`
|
||||||
|
|
||||||
// oauth2 配置
|
// oauth2 配置
|
||||||
Oauth2 map[string]*Oauth2Config `koanf:"oauth2" json:"oauth2,omitempty"`
|
Oauth2 map[string]*Oauth2Config `koanf:"oauth2" json:"oauth2,omitempty"`
|
||||||
@@ -193,6 +199,30 @@ func (c *Config) Save() error {
|
|||||||
return c.save()
|
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 {
|
func (c *Config) save() error {
|
||||||
data, err := yaml.Marshal(c)
|
data, err := yaml.Marshal(c)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
@@ -211,6 +241,40 @@ func (c *Config) write(data []byte) error {
|
|||||||
return os.WriteFile(c.filePath, data, 0600)
|
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 {
|
func koanfConf(c any) koanf.UnmarshalConf {
|
||||||
return koanf.UnmarshalConf{
|
return koanf.UnmarshalConf{
|
||||||
DecoderConfig: &mapstructure.DecoderConfig{
|
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 {
|
func newTempConfig(t *testing.T, cfg string) string {
|
||||||
t.Helper()
|
t.Helper()
|
||||||
|
|
||||||
|
|||||||
@@ -80,7 +80,6 @@ func execCase(t *testing.T, item testSt) {
|
|||||||
CountryCode: "",
|
CountryCode: "",
|
||||||
},
|
},
|
||||||
LastActive: time.Time{},
|
LastActive: time.Time{},
|
||||||
TaskStream: nil,
|
|
||||||
PrevTransferInSnapshot: 0,
|
PrevTransferInSnapshot: 0,
|
||||||
PrevTransferOutSnapshot: 0,
|
PrevTransferOutSnapshot: 0,
|
||||||
}
|
}
|
||||||
|
|||||||
+105
-3
@@ -3,6 +3,7 @@ package model
|
|||||||
import (
|
import (
|
||||||
"log"
|
"log"
|
||||||
"slices"
|
"slices"
|
||||||
|
"sync/atomic"
|
||||||
"time"
|
"time"
|
||||||
|
|
||||||
"github.com/goccy/go-json"
|
"github.com/goccy/go-json"
|
||||||
@@ -32,13 +33,68 @@ type Server struct {
|
|||||||
GeoIP *GeoIP `gorm:"-" json:"geoip,omitempty"`
|
GeoIP *GeoIP `gorm:"-" json:"geoip,omitempty"`
|
||||||
LastActive time.Time `gorm:"-" json:"last_active,omitempty"`
|
LastActive time.Time `gorm:"-" json:"last_active,omitempty"`
|
||||||
|
|
||||||
TaskStream pb.NezhaService_RequestTaskServer `gorm:"-" json:"-"`
|
// taskStream MUST be accessed only via SetTaskStream / GetTaskStream. Direct
|
||||||
ConfigCache chan any `gorm:"-" json:"-"`
|
// 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:"-"` // 上次数据点时的入站使用量
|
PrevTransferInSnapshot uint64 `gorm:"-" json:"-"` // 上次数据点时的入站使用量
|
||||||
PrevTransferOutSnapshot 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) {
|
func InitServer(s *Server) {
|
||||||
s.Host = &Host{}
|
s.Host = &Host{}
|
||||||
s.State = &HostState{}
|
s.State = &HostState{}
|
||||||
@@ -51,7 +107,9 @@ func (s *Server) CopyFromRunningServer(old *Server) {
|
|||||||
s.State = old.State
|
s.State = old.State
|
||||||
s.GeoIP = old.GeoIP
|
s.GeoIP = old.GeoIP
|
||||||
s.LastActive = old.LastActive
|
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.ConfigCache = old.ConfigCache
|
||||||
s.PrevTransferInSnapshot = old.PrevTransferInSnapshot
|
s.PrevTransferInSnapshot = old.PrevTransferInSnapshot
|
||||||
s.PrevTransferOutSnapshot = old.PrevTransferOutSnapshot
|
s.PrevTransferOutSnapshot = old.PrevTransferOutSnapshot
|
||||||
@@ -73,6 +131,50 @@ func (s *Server) AfterFind(tx *gorm.DB) error {
|
|||||||
return nil
|
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) {
|
func (s *Server) SplitList(x []*Server) ([]*Server, []*Server) {
|
||||||
pri := func(s *Server) bool {
|
pri := func(s *Server) bool {
|
||||||
return s.DisplayIndex == 0
|
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
|
TaskTypeFM
|
||||||
TaskTypeReportConfig
|
TaskTypeReportConfig
|
||||||
TaskTypeApplyConfig
|
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 {
|
type TerminalTask struct {
|
||||||
@@ -133,7 +137,8 @@ func IsServiceSentinelNeeded(t uint64) bool {
|
|||||||
switch t {
|
switch t {
|
||||||
case TaskTypeCommand, TaskTypeTerminalGRPC, TaskTypeUpgrade,
|
case TaskTypeCommand, TaskTypeTerminalGRPC, TaskTypeUpgrade,
|
||||||
TaskTypeKeepalive, TaskTypeNAT, TaskTypeFM,
|
TaskTypeKeepalive, TaskTypeNAT, TaskTypeFM,
|
||||||
TaskTypeReportConfig, TaskTypeApplyConfig:
|
TaskTypeReportConfig, TaskTypeApplyConfig,
|
||||||
|
TaskTypeServerTransferApply:
|
||||||
return false
|
return false
|
||||||
default:
|
default:
|
||||||
return true
|
return true
|
||||||
|
|||||||
@@ -32,6 +32,7 @@ type User struct {
|
|||||||
|
|
||||||
type UserInfo struct {
|
type UserInfo struct {
|
||||||
Role Role
|
Role Role
|
||||||
|
Username string
|
||||||
AgentSecret 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)
|
||||||
|
}
|
||||||
|
}
|
||||||
+170
-12
@@ -3,6 +3,7 @@ package rpc
|
|||||||
import (
|
import (
|
||||||
"context"
|
"context"
|
||||||
"fmt"
|
"fmt"
|
||||||
|
"log"
|
||||||
"strings"
|
"strings"
|
||||||
|
|
||||||
petname "github.com/dustinkirkland/golang-petname"
|
petname "github.com/dustinkirkland/golang-petname"
|
||||||
@@ -21,6 +22,18 @@ type authHandler struct {
|
|||||||
}
|
}
|
||||||
|
|
||||||
func (a *authHandler) Check(ctx context.Context) (uint64, error) {
|
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)
|
md, ok := metadata.FromIncomingContext(ctx)
|
||||||
if !ok {
|
if !ok {
|
||||||
return 0, status.Errorf(codes.Unauthenticated, "获取 metaData 失败")
|
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)
|
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()
|
singleton.UserLock.RLock()
|
||||||
userId, ok := singleton.AgentSecretToUserId[clientSecret]
|
userId, ok := singleton.AgentSecretToUserId[clientSecret]
|
||||||
if !ok {
|
if !ok {
|
||||||
@@ -48,15 +162,6 @@ func (a *authHandler) Check(ctx context.Context) (uint64, error) {
|
|||||||
|
|
||||||
model.UnblockIP(singleton.DB, ip, model.BlockIDgRPC)
|
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)
|
clientID, hasID, err := authorizeAgentForUUID(userId, clientUUID)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return 0, status.Error(codes.Unauthenticated, err.Error())
|
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
|
// 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
|
// agent secret owner", it pins down which user's secret has been reused
|
||||||
// against a server they don't own.
|
// 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) {
|
func authorizeAgentForUUID(userId uint64, clientUUID string) (clientID uint64, hasID bool, err error) {
|
||||||
cid, found := singleton.ServerShared.UUIDToID(clientUUID)
|
cid, found := singleton.ServerShared.UUIDToID(clientUUID)
|
||||||
if !found {
|
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.
|
// agent secrets, so keep it compatible by allowing any existing UUID.
|
||||||
return cid, true, nil
|
return cid, true, nil
|
||||||
}
|
}
|
||||||
if server.UserID != userId {
|
if server.GetUserID() == userId {
|
||||||
return 0, false, fmt.Errorf("client UUID does not belong to the agent secret owner")
|
// SECURITY: while a transfer is Pending, Server.UserID has already
|
||||||
|
// been flipped to ToUserID by Register, so userId==Server.UserID
|
||||||
|
// here also matches the destination user's user-global AgentSecret.
|
||||||
|
// PushIfOnline only delivers the per-transfer HandshakeSecret on
|
||||||
|
// the wire; the destination user's global AgentSecret is never
|
||||||
|
// pushed to the agent, so a reconnect under that secret is not
|
||||||
|
// proof of agent rotation. Admitting it would let the destination
|
||||||
|
// user — who can see Server.UUID — authenticate as the agent
|
||||||
|
// during the Pending window. Reject the user-global secret until
|
||||||
|
// the transfer settles; the HandshakeSecret path in check() is
|
||||||
|
// the only valid promotion route.
|
||||||
|
if singleton.ServerTransferShared != nil {
|
||||||
|
if _, ok := singleton.ServerTransferShared.LookupPending(cid); ok {
|
||||||
|
return 0, false, fmt.Errorf("destination user's global AgentSecret cannot authenticate during a pending transfer; agent must rotate to per-transfer HandshakeSecret")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return cid, true, nil
|
||||||
}
|
}
|
||||||
return cid, true, nil
|
// server.UserID != userId — normally an impersonation attempt. Allow it
|
||||||
|
// only when a ServerTransfer for this server is Pending AND the secret in
|
||||||
|
// hand is the previous owner's (FromUserID), OR when a recently terminated
|
||||||
|
// transfer left a revert-delivery for FromUserID and the agent is still
|
||||||
|
// presenting its pre-transfer global secret.
|
||||||
|
//
|
||||||
|
// SECURITY: we deliberately do NOT accept the destination user's global
|
||||||
|
// AgentSecret on the LookupRevertDelivery path. PushIfOnline only ever
|
||||||
|
// delivers per-transfer HandshakeSecret / RevertHandshakeSecret to the
|
||||||
|
// agent — the ToUserID global secret never travels over the wire — so a
|
||||||
|
// reconnect under that credential is not proof of agent rotation; it can
|
||||||
|
// only come from the destination user themselves, who can see Server.UUID
|
||||||
|
// once Register flips Server.UserID. Admitting it would let that user
|
||||||
|
// impersonate the agent during the rollback window, trigger
|
||||||
|
// pushRevertIfOnline to leak RevertHandshakeSecret, and then be promoted
|
||||||
|
// into verifiedHandshakes via MarkRevertDelivered. The legitimate recovery
|
||||||
|
// paths are: FromUserID global secret (handled below), forward
|
||||||
|
// HandshakeSecret and RevertHandshakeSecret (handled by the
|
||||||
|
// terminalSecretRecovery / verifiedHandshakes lookups in check()).
|
||||||
|
if singleton.ServerTransferShared != nil {
|
||||||
|
if t, ok := singleton.ServerTransferShared.LookupRevertDelivery(cid); ok && t.FromUserID == userId {
|
||||||
|
return cid, true, nil
|
||||||
|
}
|
||||||
|
if t, ok := singleton.ServerTransferShared.LookupPending(cid); ok && t.FromUserID == userId {
|
||||||
|
return cid, true, nil
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return 0, false, fmt.Errorf("client UUID does not belong to the agent secret owner")
|
||||||
}
|
}
|
||||||
|
|||||||
+506
-1
@@ -1,15 +1,93 @@
|
|||||||
package rpc
|
package rpc
|
||||||
|
|
||||||
import (
|
import (
|
||||||
|
"context"
|
||||||
|
"errors"
|
||||||
"testing"
|
"testing"
|
||||||
|
|
||||||
|
"google.golang.org/grpc/metadata"
|
||||||
"gorm.io/driver/sqlite"
|
"gorm.io/driver/sqlite"
|
||||||
"gorm.io/gorm"
|
"gorm.io/gorm"
|
||||||
|
|
||||||
"github.com/nezhahq/nezha/model"
|
"github.com/nezhahq/nezha/model"
|
||||||
|
"github.com/nezhahq/nezha/pkg/utils"
|
||||||
"github.com/nezhahq/nezha/service/singleton"
|
"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
|
// setupAuthAgentFixture seeds an in-memory DB and ServerShared with two
|
||||||
// servers belonging to different users so we can assert that a secret bound
|
// 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.
|
// to user A cannot resolve a server UUID owned by user B.
|
||||||
@@ -17,12 +95,13 @@ func setupAuthAgentFixture(t *testing.T) func() {
|
|||||||
t.Helper()
|
t.Helper()
|
||||||
originalDB := singleton.DB
|
originalDB := singleton.DB
|
||||||
originalServerShared := singleton.ServerShared
|
originalServerShared := singleton.ServerShared
|
||||||
|
originalServerTransferShared := singleton.ServerTransferShared
|
||||||
|
|
||||||
db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
|
db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
|
||||||
if err != nil {
|
if err != nil {
|
||||||
t.Fatalf("open db: %v", err)
|
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)
|
t.Fatalf("migrate: %v", err)
|
||||||
}
|
}
|
||||||
if err := db.Create(&model.Server{
|
if err := db.Create(&model.Server{
|
||||||
@@ -41,10 +120,15 @@ func setupAuthAgentFixture(t *testing.T) func() {
|
|||||||
}
|
}
|
||||||
singleton.DB = db
|
singleton.DB = db
|
||||||
singleton.ServerShared = singleton.NewServerClass()
|
singleton.ServerShared = singleton.NewServerClass()
|
||||||
|
singleton.ServerTransferShared = singleton.NewServerTransferClass()
|
||||||
|
|
||||||
return func() {
|
return func() {
|
||||||
|
if singleton.ServerTransferShared != nil {
|
||||||
|
singleton.ServerTransferShared.Stop()
|
||||||
|
}
|
||||||
singleton.DB = originalDB
|
singleton.DB = originalDB
|
||||||
singleton.ServerShared = originalServerShared
|
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)
|
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")
|
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 {
|
func (s *NezhaHandler) CloseStream(streamId string) error {
|
||||||
s.ioStreamMutex.Lock()
|
s.ioStreamMutex.Lock()
|
||||||
defer s.ioStreamMutex.Unlock()
|
defer s.ioStreamMutex.Unlock()
|
||||||
@@ -136,6 +165,8 @@ func (s *NezhaHandler) CloseStream(streamId string) error {
|
|||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
func (s *NezhaHandler) UserConnected(streamId string, userIo io.ReadWriteCloser) error {
|
func (s *NezhaHandler) UserConnected(streamId string, userIo io.ReadWriteCloser) error {
|
||||||
stream, err := s.GetStream(streamId)
|
stream, err := s.GetStream(streamId)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
|
|||||||
+32
-2
@@ -40,12 +40,21 @@ func NewNezhaHandler() *NezhaHandler {
|
|||||||
func (s *NezhaHandler) RequestTask(stream pb.NezhaService_RequestTaskServer) error {
|
func (s *NezhaHandler) RequestTask(stream pb.NezhaService_RequestTaskServer) error {
|
||||||
var clientID uint64
|
var clientID uint64
|
||||||
var err error
|
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
|
return err
|
||||||
}
|
}
|
||||||
|
|
||||||
server, _ := singleton.ServerShared.Get(clientID)
|
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
|
var result *pb.TaskResult
|
||||||
for {
|
for {
|
||||||
result, err = stream.Recv()
|
result, err = stream.Recv()
|
||||||
@@ -83,6 +92,27 @@ func (s *NezhaHandler) RequestTask(stream pb.NezhaService_RequestTaskServer) err
|
|||||||
}
|
}
|
||||||
server.ConfigCache <- result.Data
|
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:
|
default:
|
||||||
if model.IsServiceSentinelNeeded(result.GetType()) {
|
if model.IsServiceSentinelNeeded(result.GetType()) {
|
||||||
singleton.ServiceSentinelShared.Dispatch(singleton.ReportData{
|
singleton.ServiceSentinelShared.Dispatch(singleton.ReportData{
|
||||||
|
|||||||
@@ -18,6 +18,7 @@ import (
|
|||||||
type requestTaskSecurityStream struct {
|
type requestTaskSecurityStream struct {
|
||||||
ctx context.Context
|
ctx context.Context
|
||||||
results []*pb.TaskResult
|
results []*pb.TaskResult
|
||||||
|
onRecv func()
|
||||||
onSend func(*pb.Task)
|
onSend func(*pb.Task)
|
||||||
sendErr error
|
sendErr error
|
||||||
}
|
}
|
||||||
@@ -31,6 +32,9 @@ func (s *requestTaskSecurityStream) Send(task *pb.Task) error {
|
|||||||
|
|
||||||
func (s *requestTaskSecurityStream) Recv() (*pb.TaskResult, error) {
|
func (s *requestTaskSecurityStream) Recv() (*pb.TaskResult, error) {
|
||||||
if len(s.results) == 0 {
|
if len(s.results) == 0 {
|
||||||
|
if s.onRecv != nil {
|
||||||
|
s.onRecv()
|
||||||
|
}
|
||||||
return nil, context.Canceled
|
return nil, context.Canceled
|
||||||
}
|
}
|
||||||
result := s.results[0]
|
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) {
|
func setupRequestTaskSecurityFixture(t *testing.T, servers []*model.Server, crons []*model.Cron, users map[uint64]model.UserInfo, agentSecrets map[string]uint64) {
|
||||||
t.Helper()
|
t.Helper()
|
||||||
|
|
||||||
@@ -305,22 +355,26 @@ func connectRequestTaskSecurityTaskStreamWithSendHook(t *testing.T, serverID uin
|
|||||||
if !ok {
|
if !ok {
|
||||||
t.Fatalf("server %d not found", serverID)
|
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) {
|
func runRequestTaskSecurityResult(t *testing.T, secret string, uuid string, result *pb.TaskResult) {
|
||||||
t.Helper()
|
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(
|
ctx: metadata.NewIncomingContext(context.Background(), metadata.Pairs(
|
||||||
"client_secret", secret,
|
"client_secret", secret,
|
||||||
"client_uuid", uuid,
|
"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
|
role = u.Role
|
||||||
}
|
}
|
||||||
UserLock.RUnlock()
|
UserLock.RUnlock()
|
||||||
if alert.UserID != server.UserID && !role.IsAdmin() {
|
if alert.UserID != server.GetUserID() && !role.IsAdmin() {
|
||||||
continue
|
continue
|
||||||
}
|
}
|
||||||
alertsStore[alert.ID][server.ID] = append(alertsStore[alert.
|
alertsStore[alert.ID][server.ID] = append(alertsStore[alert.
|
||||||
|
|||||||
@@ -1,6 +1,7 @@
|
|||||||
package singleton
|
package singleton
|
||||||
|
|
||||||
import (
|
import (
|
||||||
|
"log"
|
||||||
"strconv"
|
"strconv"
|
||||||
"strings"
|
"strings"
|
||||||
|
|
||||||
@@ -26,6 +27,13 @@ func InitConfigFromPath(path string) error {
|
|||||||
if err != nil {
|
if err != nil {
|
||||||
return err
|
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.updateIgnoredIPNotificationID()
|
||||||
Conf.Oauth2Providers = utils.MapKeysToSlice(Conf.Oauth2)
|
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) {
|
if !cronCanSendToServer(cr, s) {
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
if s.TaskStream != nil {
|
stream := s.GetTaskStream()
|
||||||
|
if stream != nil {
|
||||||
cronShared := CronShared
|
cronShared := CronShared
|
||||||
if cronShared != nil {
|
if cronShared != nil {
|
||||||
cronShared.reserveAlertTriggerCronResult(cr.ID, s.ID)
|
cronShared.reserveAlertTriggerCronResult(cr.ID, s.ID)
|
||||||
}
|
}
|
||||||
if err := s.TaskStream.Send(&pb.Task{
|
if err := stream.Send(&pb.Task{
|
||||||
Id: cr.ID,
|
Id: cr.ID,
|
||||||
Data: cr.Command,
|
Data: cr.Command,
|
||||||
Type: model.TaskTypeCommand,
|
Type: model.TaskTypeCommand,
|
||||||
@@ -296,8 +297,8 @@ func CronTrigger(cr *model.Cron, triggerServer ...uint64) func() {
|
|||||||
if cr.Cover == model.CronCoverIgnoreAll && !crIgnoreMap[s.ID] {
|
if cr.Cover == model.CronCoverIgnoreAll && !crIgnoreMap[s.ID] {
|
||||||
continue
|
continue
|
||||||
}
|
}
|
||||||
if s.TaskStream != nil {
|
if stream := s.GetTaskStream(); stream != nil {
|
||||||
s.TaskStream.Send(&pb.Task{
|
stream.Send(&pb.Task{
|
||||||
Id: cr.ID,
|
Id: cr.ID,
|
||||||
Data: cr.Command,
|
Data: cr.Command,
|
||||||
Type: model.TaskTypeCommand,
|
Type: model.TaskTypeCommand,
|
||||||
@@ -313,7 +314,7 @@ func CronTrigger(cr *model.Cron, triggerServer ...uint64) func() {
|
|||||||
}
|
}
|
||||||
|
|
||||||
func cronCanSendToServer(cr *model.Cron, server *model.Server) bool {
|
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 {
|
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) SendMsg(any) error { return nil }
|
||||||
func (s *capturedTaskStream) RecvMsg(any) error { return context.Canceled }
|
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) {
|
func replaceServerSharedForSecurityTest(t *testing.T, servers ...*model.Server) {
|
||||||
t.Helper()
|
t.Helper()
|
||||||
|
|
||||||
@@ -75,8 +85,8 @@ func TestCronTriggerSkipsServersOwnedByOtherUsers(t *testing.T) {
|
|||||||
firstStream := newCapturedTaskStream()
|
firstStream := newCapturedTaskStream()
|
||||||
secondStream := newCapturedTaskStream()
|
secondStream := newCapturedTaskStream()
|
||||||
replaceServerSharedForSecurityTest(t,
|
replaceServerSharedForSecurityTest(t,
|
||||||
&model.Server{Common: model.Common{ID: 1, UserID: 100}, Name: "member-server", TaskStream: firstStream},
|
withTaskStream(&model.Server{Common: model.Common{ID: 1, UserID: 100}, Name: "member-server"}, firstStream),
|
||||||
&model.Server{Common: model.Common{ID: 2, UserID: 200}, Name: "admin-server", TaskStream: secondStream},
|
withTaskStream(&model.Server{Common: model.Common{ID: 2, UserID: 200}, Name: "admin-server"}, secondStream),
|
||||||
)
|
)
|
||||||
|
|
||||||
cronTask := &model.Cron{
|
cronTask := &model.Cron{
|
||||||
@@ -95,7 +105,7 @@ func TestCronTriggerSkipsServersOwnedByOtherUsers(t *testing.T) {
|
|||||||
func TestSendTriggerTasksSkipsCronOwnedByAnotherUser(t *testing.T) {
|
func TestSendTriggerTasksSkipsCronOwnedByAnotherUser(t *testing.T) {
|
||||||
attackerStream := newCapturedTaskStream()
|
attackerStream := newCapturedTaskStream()
|
||||||
replaceServerSharedForSecurityTest(t,
|
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{
|
replaceUserInfoMapForSecurityTest(t, map[uint64]model.UserInfo{
|
||||||
1: {Role: model.RoleAdmin},
|
1: {Role: model.RoleAdmin},
|
||||||
@@ -147,7 +157,7 @@ func assertNoTask(t *testing.T, stream *capturedTaskStream) {
|
|||||||
func TestCronTriggerSendsToMemberOwnedServer(t *testing.T) {
|
func TestCronTriggerSendsToMemberOwnedServer(t *testing.T) {
|
||||||
memberStream := newCapturedTaskStream()
|
memberStream := newCapturedTaskStream()
|
||||||
replaceServerSharedForSecurityTest(t,
|
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{
|
replaceUserInfoMapForSecurityTest(t, map[uint64]model.UserInfo{
|
||||||
100: {Role: model.RoleMember},
|
100: {Role: model.RoleMember},
|
||||||
@@ -168,8 +178,8 @@ func TestCronTriggerAdminCronFansOutAcrossOwners(t *testing.T) {
|
|||||||
first := newCapturedTaskStream()
|
first := newCapturedTaskStream()
|
||||||
second := newCapturedTaskStream()
|
second := newCapturedTaskStream()
|
||||||
replaceServerSharedForSecurityTest(t,
|
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),
|
||||||
&model.Server{Common: model.Common{ID: 2, UserID: 200}, Name: "admin-server", TaskStream: second},
|
withTaskStream(&model.Server{Common: model.Common{ID: 2, UserID: 200}, Name: "admin-server"}, second),
|
||||||
)
|
)
|
||||||
replaceUserInfoMapForSecurityTest(t, map[uint64]model.UserInfo{
|
replaceUserInfoMapForSecurityTest(t, map[uint64]model.UserInfo{
|
||||||
1: {Role: model.RoleAdmin},
|
1: {Role: model.RoleAdmin},
|
||||||
@@ -192,7 +202,7 @@ func TestCronTriggerAdminCronFansOutAcrossOwners(t *testing.T) {
|
|||||||
func TestCronTriggerLegacyZeroOwnerFansOut(t *testing.T) {
|
func TestCronTriggerLegacyZeroOwnerFansOut(t *testing.T) {
|
||||||
first := newCapturedTaskStream()
|
first := newCapturedTaskStream()
|
||||||
replaceServerSharedForSecurityTest(t,
|
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{
|
replaceUserInfoMapForSecurityTest(t, map[uint64]model.UserInfo{
|
||||||
100: {Role: model.RoleMember},
|
100: {Role: model.RoleMember},
|
||||||
@@ -212,7 +222,7 @@ func TestCronTriggerLegacyZeroOwnerFansOut(t *testing.T) {
|
|||||||
func TestCronTriggerSkipsServersWhenOwnerNotKnown(t *testing.T) {
|
func TestCronTriggerSkipsServersWhenOwnerNotKnown(t *testing.T) {
|
||||||
stream := newCapturedTaskStream()
|
stream := newCapturedTaskStream()
|
||||||
replaceServerSharedForSecurityTest(t,
|
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{
|
replaceUserInfoMapForSecurityTest(t, map[uint64]model.UserInfo{
|
||||||
100: {Role: model.RoleMember},
|
100: {Role: model.RoleMember},
|
||||||
@@ -232,7 +242,7 @@ func TestCronTriggerSkipsServersWhenOwnerNotKnown(t *testing.T) {
|
|||||||
func TestSendTriggerTasksAllowsSelfOwnedCron(t *testing.T) {
|
func TestSendTriggerTasksAllowsSelfOwnedCron(t *testing.T) {
|
||||||
stream := newCapturedTaskStream()
|
stream := newCapturedTaskStream()
|
||||||
replaceServerSharedForSecurityTest(t,
|
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{
|
replaceUserInfoMapForSecurityTest(t, map[uint64]model.UserInfo{
|
||||||
200: {Role: model.RoleMember},
|
200: {Role: model.RoleMember},
|
||||||
@@ -257,7 +267,7 @@ func TestSendTriggerTasksAllowsSelfOwnedCron(t *testing.T) {
|
|||||||
func TestSendTriggerTasksAllowsAdminCallerToTriggerAny(t *testing.T) {
|
func TestSendTriggerTasksAllowsAdminCallerToTriggerAny(t *testing.T) {
|
||||||
stream := newCapturedTaskStream()
|
stream := newCapturedTaskStream()
|
||||||
replaceServerSharedForSecurityTest(t,
|
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{
|
replaceUserInfoMapForSecurityTest(t, map[uint64]model.UserInfo{
|
||||||
1: {Role: model.RoleAdmin},
|
1: {Role: model.RoleAdmin},
|
||||||
@@ -283,7 +293,7 @@ func TestSendTriggerTasksAllowsAdminCallerToTriggerAny(t *testing.T) {
|
|||||||
func TestSendTriggerTasksIgnoresUnknownTaskIDs(t *testing.T) {
|
func TestSendTriggerTasksIgnoresUnknownTaskIDs(t *testing.T) {
|
||||||
stream := newCapturedTaskStream()
|
stream := newCapturedTaskStream()
|
||||||
replaceServerSharedForSecurityTest(t,
|
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{
|
replaceUserInfoMapForSecurityTest(t, map[uint64]model.UserInfo{
|
||||||
200: {Role: model.RoleMember},
|
200: {Role: model.RoleMember},
|
||||||
@@ -304,7 +314,7 @@ func TestSendTriggerTasksIgnoresUnknownTaskIDs(t *testing.T) {
|
|||||||
func TestSendTriggerTasksMixedCronIDsOnlyFiresAllowed(t *testing.T) {
|
func TestSendTriggerTasksMixedCronIDsOnlyFiresAllowed(t *testing.T) {
|
||||||
stream := newCapturedTaskStream()
|
stream := newCapturedTaskStream()
|
||||||
replaceServerSharedForSecurityTest(t,
|
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{
|
replaceUserInfoMapForSecurityTest(t, map[uint64]model.UserInfo{
|
||||||
1: {Role: model.RoleAdmin},
|
1: {Role: model.RoleAdmin},
|
||||||
@@ -540,7 +550,7 @@ func (s *failingTaskStream) Send(task *pb.Task) error {
|
|||||||
func TestCronTriggerRevokesAlertTriggerAuthorizationOnSendFailure(t *testing.T) {
|
func TestCronTriggerRevokesAlertTriggerAuthorizationOnSendFailure(t *testing.T) {
|
||||||
failing := newFailingTaskStream(context.Canceled)
|
failing := newFailingTaskStream(context.Canceled)
|
||||||
replaceServerSharedForSecurityTest(t,
|
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{
|
replaceUserInfoMapForSecurityTest(t, map[uint64]model.UserInfo{
|
||||||
100: {Role: model.RoleMember},
|
100: {Role: model.RoleMember},
|
||||||
@@ -828,6 +838,12 @@ func newServiceMonitorSecurityHarness(t *testing.T, servers ...*model.Server) *S
|
|||||||
t.Fatal(err)
|
t.Fatal(err)
|
||||||
}
|
}
|
||||||
ServiceSentinelShared = ss
|
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
|
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天数据缓存
|
// 30天数据缓存
|
||||||
monthlyStatusLock sync.Mutex
|
monthlyStatusLock sync.Mutex
|
||||||
monthlyStatus map[uint64]*serviceResponseItem
|
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 创建服务监控器
|
// NewServiceSentinel 创建服务监控器
|
||||||
@@ -113,7 +123,11 @@ func NewServiceSentinel(serviceSentinelDispatchBus chan<- *model.Service) (*Serv
|
|||||||
ss.loadTodayStats(today)
|
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)
|
_, err = CronShared.AddFunc("0 0 0 * * *", ss.refreshMonthlyServiceStatus)
|
||||||
@@ -489,10 +503,34 @@ func canReportServiceResult(service *model.Service, reporter *model.Server, task
|
|||||||
return false
|
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 服务监控的实际工作流程
|
// 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() {
|
func (ss *ServiceSentinel) worker() {
|
||||||
// 从服务状态汇报管道获取汇报的服务数据
|
// 从服务状态汇报管道获取汇报的服务数据
|
||||||
for r := range ss.serviceReportChannel {
|
for r := range ss.serviceReportChannel {
|
||||||
|
|||||||
@@ -34,6 +34,10 @@ var (
|
|||||||
NotificationShared *NotificationClass
|
NotificationShared *NotificationClass
|
||||||
NATShared *NATClass
|
NATShared *NATClass
|
||||||
CronShared *CronClass
|
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
|
//go:embed frontend-templates.yaml
|
||||||
@@ -59,6 +63,7 @@ func LoadSingleton(bus chan<- *model.Service) (err error) {
|
|||||||
NotificationShared = NewNotificationClass()
|
NotificationShared = NewNotificationClass()
|
||||||
ServerShared = NewServerClass()
|
ServerShared = NewServerClass()
|
||||||
CronShared = NewCronClass()
|
CronShared = NewCronClass()
|
||||||
|
ServerTransferShared = NewServerTransferClass()
|
||||||
// 最后初始化 ServiceSentinel
|
// 最后初始化 ServiceSentinel
|
||||||
ServiceSentinelShared, err = NewServiceSentinel(bus)
|
ServiceSentinelShared, err = NewServiceSentinel(bus)
|
||||||
return
|
return
|
||||||
@@ -89,7 +94,7 @@ func InitDBFromPath(path string) error {
|
|||||||
model.Notification{}, model.AlertRule{}, model.Service{}, model.NotificationGroupNotification{},
|
model.Notification{}, model.AlertRule{}, model.Service{}, model.NotificationGroupNotification{},
|
||||||
model.Cron{}, model.Transfer{}, model.ServerGroupServer{},
|
model.Cron{}, model.Transfer{}, model.ServerGroupServer{},
|
||||||
model.NAT{}, model.DDNSProfile{}, model.NotificationGroupNotification{},
|
model.NAT{}, model.DDNSProfile{}, model.NotificationGroupNotification{},
|
||||||
model.WAF{}, model.Oauth2Bind{})
|
model.WAF{}, model.Oauth2Bind{}, model.ServerTransfer{})
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return err
|
return err
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -40,10 +40,34 @@ func initUser() {
|
|||||||
|
|
||||||
UserInfoMap[u.ID] = model.UserInfo{
|
UserInfoMap[u.ID] = model.UserInfo{
|
||||||
Role: u.Role,
|
Role: u.Role,
|
||||||
|
Username: u.Username,
|
||||||
AgentSecret: u.AgentSecret,
|
AgentSecret: u.AgentSecret,
|
||||||
}
|
}
|
||||||
AgentSecretToUserId[u.AgentSecret] = u.ID
|
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) {
|
func OnUserUpdate(u *model.User) {
|
||||||
@@ -56,6 +80,7 @@ func OnUserUpdate(u *model.User) {
|
|||||||
|
|
||||||
UserInfoMap[u.ID] = model.UserInfo{
|
UserInfoMap[u.ID] = model.UserInfo{
|
||||||
Role: u.Role,
|
Role: u.Role,
|
||||||
|
Username: u.Username,
|
||||||
AgentSecret: u.AgentSecret,
|
AgentSecret: u.AgentSecret,
|
||||||
}
|
}
|
||||||
AgentSecretToUserId[u.AgentSecret] = u.ID
|
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")
|
return Localizer.ErrorT("user id not specified")
|
||||||
}
|
}
|
||||||
|
|
||||||
|
if ServerTransferShared != nil {
|
||||||
|
ServerTransferShared.OnUsersDeleted(id)
|
||||||
|
}
|
||||||
|
|
||||||
var (
|
var (
|
||||||
cron, server bool
|
cron, server bool
|
||||||
crons, servers []uint64
|
crons, servers []uint64
|
||||||
@@ -127,6 +156,11 @@ func OnUserDelete(id []uint64, errorFunc func(string, ...any) error) error {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
AlertsLock.Unlock()
|
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)
|
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