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:
+23
-2
@@ -6,6 +6,7 @@ import (
|
||||
"slices"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
@@ -37,8 +38,21 @@ func (c *Common) GetID() uint64 {
|
||||
return c.ID
|
||||
}
|
||||
|
||||
// GetUserID 原子读取所属用户 ID。Server.UserID 会在 ServerTransfer 的
|
||||
// Register/revertTransition 流程里被实时改写以反映新所有者,同时 auth
|
||||
// 热路径在每次 agent RPC 都会读它。任何并发读必须走 atomic,否则与 SetUserID
|
||||
// 一起会被 go race detector 识别为 data race(见
|
||||
// TestServerUserIDConcurrentAccessIsRaceFree)。
|
||||
func (c *Common) GetUserID() uint64 {
|
||||
return c.UserID
|
||||
return atomic.LoadUint64(&c.UserID)
|
||||
}
|
||||
|
||||
// SetUserID 原子改写所属用户 ID。仅在「server 已经在 in-memory cache 里」
|
||||
// 的写入路径(ServerTransfer.Register / revertTransition)需要用 atomic
|
||||
// 保证可见性;普通 GORM AfterFind / Create 因为没有并发读所以可以直接赋
|
||||
// 值。配合 GetUserID 形成 atomic-only 的并发访问协议。
|
||||
func (c *Common) SetUserID(uid uint64) {
|
||||
atomic.StoreUint64(&c.UserID, uid)
|
||||
}
|
||||
|
||||
func (c *Common) HasPermission(ctx *gin.Context) bool {
|
||||
@@ -52,7 +66,14 @@ func (c *Common) HasPermission(ctx *gin.Context) bool {
|
||||
return true
|
||||
}
|
||||
|
||||
return user.ID == c.UserID
|
||||
// 必须走 GetUserID 而不是裸读 c.UserID — Server.UserID 在
|
||||
// ServerTransfer.Register / revertTransition 里会被 atomic.StoreUint64
|
||||
// 改写,dashboard 各 controller 在 listHandler post-filter 这条热路径上
|
||||
// 高频对同一 *Server 调 HasPermission。裸读会与 SetUserID 形成 data
|
||||
// race(TestCommonHasPermissionConcurrentWithSetUserIDIsRaceFree 在
|
||||
// -race 下钉死该不变量),并且在 transfer 切换瞬间可能给出错误的权限
|
||||
// 判断。
|
||||
return user.ID == c.GetUserID()
|
||||
}
|
||||
|
||||
type CommonInterface interface {
|
||||
|
||||
+69
-5
@@ -3,6 +3,7 @@ package model
|
||||
import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
"github.com/go-viper/mapstructure/v2"
|
||||
@@ -16,8 +17,12 @@ import (
|
||||
)
|
||||
|
||||
const (
|
||||
ConfigUsePeerIP = "NZ::Use-Peer-IP"
|
||||
ConfigCoverAll = iota
|
||||
ConfigUsePeerIP = "NZ::Use-Peer-IP"
|
||||
JWTSecretKeyRotationBaselineVersion = "v2.0.13"
|
||||
)
|
||||
|
||||
const (
|
||||
ConfigCoverAll = iota + 1
|
||||
ConfigCoverIgnoreAll
|
||||
)
|
||||
|
||||
@@ -60,9 +65,10 @@ type Config struct {
|
||||
AgentSecretKey string `koanf:"agent_secret_key" json:"agent_secret_key,omitempty"`
|
||||
JWTTimeout int `koanf:"jwt_timeout" json:"jwt_timeout,omitempty"` // JWT token过期时间(小时)
|
||||
|
||||
JWTSecretKey string `koanf:"jwt_secret_key" json:"jwt_secret_key,omitempty"`
|
||||
ListenPort uint16 `koanf:"listen_port" json:"listen_port,omitempty"`
|
||||
ListenHost string `koanf:"listen_host" json:"listen_host,omitempty"`
|
||||
JWTSecretKey string `koanf:"jwt_secret_key" json:"jwt_secret_key,omitempty"`
|
||||
JWTSecretKeyLastRotatedVersion string `koanf:"jwt_secret_key_last_rotated_version" json:"jwt_secret_key_last_rotated_version,omitempty"`
|
||||
ListenPort uint16 `koanf:"listen_port" json:"listen_port,omitempty"`
|
||||
ListenHost string `koanf:"listen_host" json:"listen_host,omitempty"`
|
||||
|
||||
// oauth2 配置
|
||||
Oauth2 map[string]*Oauth2Config `koanf:"oauth2" json:"oauth2,omitempty"`
|
||||
@@ -193,6 +199,30 @@ func (c *Config) Save() error {
|
||||
return c.save()
|
||||
}
|
||||
|
||||
func (c *Config) RotateJWTSecretKeyIfNeeded(currentVersion string) (bool, error) {
|
||||
currentVersion = strings.TrimSpace(currentVersion)
|
||||
if compareVersion(currentVersion, JWTSecretKeyRotationBaselineVersion) < 0 {
|
||||
return false, nil
|
||||
}
|
||||
|
||||
initialMarker := c.JWTSecretKeyLastRotatedVersion
|
||||
shouldRotate := c.JWTSecretKeyLastRotatedVersion == "" || compareVersion(c.JWTSecretKeyLastRotatedVersion, JWTSecretKeyRotationBaselineVersion) < 0
|
||||
if shouldRotate {
|
||||
secret, err := utils.GenerateRandomString(1024)
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
c.JWTSecretKey = secret
|
||||
}
|
||||
|
||||
c.JWTSecretKeyLastRotatedVersion = currentVersion
|
||||
|
||||
if !shouldRotate && c.JWTSecretKeyLastRotatedVersion == initialMarker {
|
||||
return false, nil
|
||||
}
|
||||
return shouldRotate, c.Save()
|
||||
}
|
||||
|
||||
func (c *Config) save() error {
|
||||
data, err := yaml.Marshal(c)
|
||||
if err != nil {
|
||||
@@ -211,6 +241,40 @@ func (c *Config) write(data []byte) error {
|
||||
return os.WriteFile(c.filePath, data, 0600)
|
||||
}
|
||||
|
||||
func compareVersion(left, right string) int {
|
||||
leftParts, leftOK := parseVersion(left)
|
||||
rightParts, rightOK := parseVersion(right)
|
||||
if !leftOK || !rightOK {
|
||||
return -1
|
||||
}
|
||||
for i := range leftParts {
|
||||
if leftParts[i] < rightParts[i] {
|
||||
return -1
|
||||
}
|
||||
if leftParts[i] > rightParts[i] {
|
||||
return 1
|
||||
}
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
func parseVersion(version string) ([3]int, bool) {
|
||||
version = strings.TrimPrefix(strings.TrimSpace(version), "v")
|
||||
parts := strings.Split(version, ".")
|
||||
if len(parts) != 3 {
|
||||
return [3]int{}, false
|
||||
}
|
||||
var parsed [3]int
|
||||
for i, part := range parts {
|
||||
value, err := strconv.Atoi(part)
|
||||
if err != nil {
|
||||
return [3]int{}, false
|
||||
}
|
||||
parsed[i] = value
|
||||
}
|
||||
return parsed, true
|
||||
}
|
||||
|
||||
func koanfConf(c any) koanf.UnmarshalConf {
|
||||
return koanf.UnmarshalConf{
|
||||
DecoderConfig: &mapstructure.DecoderConfig{
|
||||
|
||||
@@ -151,6 +151,111 @@ func TestReadConfig(t *testing.T) {
|
||||
})
|
||||
}
|
||||
|
||||
func TestRotateJWTSecretKeyIfNeeded(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
initialMarker string
|
||||
currentVersion string
|
||||
wantRotated bool
|
||||
wantStoredVersion string
|
||||
wantSecretChanged bool
|
||||
wantSavedConfigKey bool
|
||||
}{
|
||||
{
|
||||
name: "empty marker rotates leaked secret",
|
||||
currentVersion: "v2.0.13",
|
||||
wantRotated: true,
|
||||
wantStoredVersion: "v2.0.13",
|
||||
wantSecretChanged: true,
|
||||
wantSavedConfigKey: true,
|
||||
},
|
||||
{
|
||||
name: "old marker rotates leaked secret",
|
||||
initialMarker: "v2.0.12",
|
||||
currentVersion: "v2.0.14",
|
||||
wantRotated: true,
|
||||
wantStoredVersion: "v2.0.14",
|
||||
wantSecretChanged: true,
|
||||
wantSavedConfigKey: true,
|
||||
},
|
||||
{
|
||||
name: "threshold marker keeps secret and advances marker",
|
||||
initialMarker: "v2.0.13",
|
||||
currentVersion: "v2.0.14",
|
||||
wantStoredVersion: "v2.0.14",
|
||||
wantSavedConfigKey: true,
|
||||
},
|
||||
{
|
||||
name: "current marker keeps secret",
|
||||
initialMarker: "v2.0.14",
|
||||
currentVersion: "v2.0.14",
|
||||
wantStoredVersion: "v2.0.14",
|
||||
},
|
||||
{
|
||||
name: "debug version skips rotation and marker update",
|
||||
currentVersion: "debug",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
file := newTempConfig(t, "")
|
||||
t.Cleanup(func() { os.Remove(file) })
|
||||
|
||||
c := &Config{
|
||||
JWTSecretKey: "leaked-secret",
|
||||
JWTSecretKeyLastRotatedVersion: tt.initialMarker,
|
||||
filePath: file,
|
||||
}
|
||||
|
||||
rotated, err := c.RotateJWTSecretKeyIfNeeded(tt.currentVersion)
|
||||
if err != nil {
|
||||
t.Fatalf("rotate jwt secret key failed: %v", err)
|
||||
}
|
||||
if rotated != tt.wantRotated {
|
||||
t.Fatalf("rotated = %v, want %v", rotated, tt.wantRotated)
|
||||
}
|
||||
if c.JWTSecretKeyLastRotatedVersion != tt.wantStoredVersion {
|
||||
t.Fatalf("jwt secret key marker = %q, want %q", c.JWTSecretKeyLastRotatedVersion, tt.wantStoredVersion)
|
||||
}
|
||||
secretChanged := c.JWTSecretKey != "leaked-secret"
|
||||
if secretChanged != tt.wantSecretChanged {
|
||||
t.Fatalf("secret changed = %v, want %v", secretChanged, tt.wantSecretChanged)
|
||||
}
|
||||
|
||||
saved, err := os.ReadFile(file)
|
||||
if err != nil {
|
||||
t.Fatalf("read saved config: %v", err)
|
||||
}
|
||||
hasMarker := strings.Contains(string(saved), "jwt_secret_key_last_rotated_version")
|
||||
if hasMarker != tt.wantSavedConfigKey {
|
||||
t.Fatalf("saved marker present = %v, want %v, config = %s", hasMarker, tt.wantSavedConfigKey, saved)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// Mirrors the upstream single-block declaration so iota lines up exactly:
|
||||
// ConfigUsePeerIP occupies iota=0 (as a typed string), ConfigCoverAll=1,
|
||||
// ConfigCoverIgnoreAll=2. Pins persisted `cover` semantics.
|
||||
const (
|
||||
originalConfigUsePeerIP = "NZ::Use-Peer-IP"
|
||||
originalConfigCoverAll = iota
|
||||
originalConfigCoverIgnoreAll
|
||||
)
|
||||
|
||||
func TestConfigCoverConstantValues(t *testing.T) {
|
||||
if ConfigUsePeerIP != originalConfigUsePeerIP {
|
||||
t.Fatalf("ConfigUsePeerIP = %q, want %q", ConfigUsePeerIP, originalConfigUsePeerIP)
|
||||
}
|
||||
if ConfigCoverAll != originalConfigCoverAll {
|
||||
t.Fatalf("ConfigCoverAll = %d, want original value %d", ConfigCoverAll, originalConfigCoverAll)
|
||||
}
|
||||
if ConfigCoverIgnoreAll != originalConfigCoverIgnoreAll {
|
||||
t.Fatalf("ConfigCoverIgnoreAll = %d, want original value %d", ConfigCoverIgnoreAll, originalConfigCoverIgnoreAll)
|
||||
}
|
||||
}
|
||||
|
||||
func newTempConfig(t *testing.T, cfg string) string {
|
||||
t.Helper()
|
||||
|
||||
|
||||
@@ -80,7 +80,6 @@ func execCase(t *testing.T, item testSt) {
|
||||
CountryCode: "",
|
||||
},
|
||||
LastActive: time.Time{},
|
||||
TaskStream: nil,
|
||||
PrevTransferInSnapshot: 0,
|
||||
PrevTransferOutSnapshot: 0,
|
||||
}
|
||||
|
||||
+105
-3
@@ -3,6 +3,7 @@ package model
|
||||
import (
|
||||
"log"
|
||||
"slices"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"github.com/goccy/go-json"
|
||||
@@ -32,13 +33,68 @@ type Server struct {
|
||||
GeoIP *GeoIP `gorm:"-" json:"geoip,omitempty"`
|
||||
LastActive time.Time `gorm:"-" json:"last_active,omitempty"`
|
||||
|
||||
TaskStream pb.NezhaService_RequestTaskServer `gorm:"-" json:"-"`
|
||||
ConfigCache chan any `gorm:"-" json:"-"`
|
||||
// taskStream MUST be accessed only via SetTaskStream / GetTaskStream. Direct
|
||||
// field access from outside this file races with the gRPC RequestTask
|
||||
// handler that reassigns the stream on every reconnect — a torn read of the
|
||||
// two-word interface header would panic on a subsequent .Send call. The
|
||||
// atomic.Pointer + holder struct lets us swap the stream lock-free while
|
||||
// every reader observes a single, consistent value.
|
||||
taskStream atomic.Pointer[taskStreamHolder]
|
||||
ConfigCache chan any `gorm:"-" json:"-"`
|
||||
|
||||
PrevTransferInSnapshot uint64 `gorm:"-" json:"-"` // 上次数据点时的入站使用量
|
||||
PrevTransferOutSnapshot uint64 `gorm:"-" json:"-"` // 上次数据点时的出站使用量
|
||||
}
|
||||
|
||||
// taskStreamHolder wraps the interface so atomic.Pointer (which requires a
|
||||
// concrete pointed-to type) can publish it atomically. The previous bare
|
||||
// field `TaskStream pb.NezhaService_RequestTaskServer` was a plain interface
|
||||
// value: two words on the heap (type ptr + data ptr). Concurrent assignment
|
||||
// produced torn reads detectable by `go test -race` and crashable in production.
|
||||
type taskStreamHolder struct {
|
||||
s pb.NezhaService_RequestTaskServer
|
||||
}
|
||||
|
||||
// SetTaskStream publishes the agent's RequestTask stream so other goroutines
|
||||
// can deliver tasks to the agent. Pass nil to detach (e.g. on disconnect).
|
||||
func (s *Server) SetTaskStream(stream pb.NezhaService_RequestTaskServer) {
|
||||
if stream == nil {
|
||||
s.taskStream.Store(nil)
|
||||
return
|
||||
}
|
||||
s.taskStream.Store(&taskStreamHolder{s: stream})
|
||||
}
|
||||
|
||||
// ClearTaskStreamIfCurrent detaches stream only if it is still the published
|
||||
// RequestTask stream. Disconnect cleanup uses this guard so an old stream
|
||||
// returning after a reconnect cannot erase the newer live stream.
|
||||
func (s *Server) ClearTaskStreamIfCurrent(stream pb.NezhaService_RequestTaskServer) bool {
|
||||
if stream == nil {
|
||||
return false
|
||||
}
|
||||
for {
|
||||
h := s.taskStream.Load()
|
||||
if h == nil || h.s != stream {
|
||||
return false
|
||||
}
|
||||
if s.taskStream.CompareAndSwap(h, nil) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// GetTaskStream returns the currently-published stream, or nil if the agent
|
||||
// is offline. Callers MUST capture the return into a local variable before
|
||||
// using it — re-reading via GetTaskStream() across a Send call reopens the
|
||||
// race we're trying to close.
|
||||
func (s *Server) GetTaskStream() pb.NezhaService_RequestTaskServer {
|
||||
h := s.taskStream.Load()
|
||||
if h == nil {
|
||||
return nil
|
||||
}
|
||||
return h.s
|
||||
}
|
||||
|
||||
func InitServer(s *Server) {
|
||||
s.Host = &Host{}
|
||||
s.State = &HostState{}
|
||||
@@ -51,7 +107,9 @@ func (s *Server) CopyFromRunningServer(old *Server) {
|
||||
s.State = old.State
|
||||
s.GeoIP = old.GeoIP
|
||||
s.LastActive = old.LastActive
|
||||
s.TaskStream = old.TaskStream
|
||||
// taskStream is an atomic.Pointer; copy the published value rather than
|
||||
// the field itself (atomic.Pointer is not safe to copy by value).
|
||||
s.SetTaskStream(old.GetTaskStream())
|
||||
s.ConfigCache = old.ConfigCache
|
||||
s.PrevTransferInSnapshot = old.PrevTransferInSnapshot
|
||||
s.PrevTransferOutSnapshot = old.PrevTransferOutSnapshot
|
||||
@@ -73,6 +131,50 @@ func (s *Server) AfterFind(tx *gorm.DB) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
// ServerOwnerInfo carries the user-facing identity for Server.UserID. It is
|
||||
// returned by the lookup function installed by the singleton layer; model
|
||||
// must not import singleton (cycle), so the dependency flows through a
|
||||
// package-level function variable instead.
|
||||
type ServerOwnerInfo struct {
|
||||
ID uint64 `json:"id"`
|
||||
Username string `json:"username,omitempty"`
|
||||
}
|
||||
|
||||
// ServerOwnerLookup is installed by singleton at startup to resolve a
|
||||
// Server.UserID into a display-ready owner record. Returns ok=false when
|
||||
// the uid does not map to a known user; the caller renders that as an
|
||||
// "unknown user" placeholder so deleted-user rows stay debuggable. Left nil
|
||||
// in tests / headless contexts so the JSON simply omits the owner field.
|
||||
var ServerOwnerLookup func(uid uint64) (ServerOwnerInfo, bool)
|
||||
|
||||
type serverJSON Server
|
||||
|
||||
type serverWithOwner struct {
|
||||
*serverJSON
|
||||
Owner *ServerOwnerInfo `json:"owner,omitempty"`
|
||||
}
|
||||
|
||||
// MarshalJSON projects Server.UserID into a structured owner field on the
|
||||
// wire. Server.UserID itself stays `json:"-"` (set on Common) so callers
|
||||
// that do not need owner info pay nothing and members do not accidentally
|
||||
// receive raw uid integers. The lookup function is consulted only when
|
||||
// installed; if absent we still emit a minimal {id} record so clients can
|
||||
// at least distinguish ownership, except for uid=0 which is the legacy
|
||||
// global-secret pseudo-owner and is best surfaced as such by the caller's
|
||||
// translation table on the frontend.
|
||||
func (s *Server) MarshalJSON() ([]byte, error) {
|
||||
owner := &ServerOwnerInfo{ID: s.GetUserID()}
|
||||
if ServerOwnerLookup != nil {
|
||||
if info, ok := ServerOwnerLookup(owner.ID); ok {
|
||||
owner.Username = info.Username
|
||||
}
|
||||
}
|
||||
return json.Marshal(serverWithOwner{
|
||||
serverJSON: (*serverJSON)(s),
|
||||
Owner: owner,
|
||||
})
|
||||
}
|
||||
|
||||
func (s *Server) SplitList(x []*Server) ([]*Server, []*Server) {
|
||||
pri := func(s *Server) bool {
|
||||
return s.DisplayIndex == 0
|
||||
|
||||
@@ -0,0 +1,129 @@
|
||||
package model
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// Server.MarshalJSON projects Server.UserID into a public owner field while
|
||||
// keeping the raw UserID json-hidden. The lookup function is package-level
|
||||
// and shared across tests; each subtest installs its own stub and restores
|
||||
// the original to avoid leaking state.
|
||||
func TestServerMarshalJSONOwnerProjection(t *testing.T) {
|
||||
original := ServerOwnerLookup
|
||||
t.Cleanup(func() { ServerOwnerLookup = original })
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
uid uint64
|
||||
lookup func(uid uint64) (ServerOwnerInfo, bool)
|
||||
wantID uint64
|
||||
wantHasName bool
|
||||
wantName string
|
||||
}{
|
||||
{
|
||||
// uid=0 is the legacy global agent secret pseudo-owner. The
|
||||
// lookup deliberately returns ok=false so the frontend can
|
||||
// render it as "Global Agent" instead of a real username.
|
||||
name: "uid_zero_has_no_username",
|
||||
uid: 0,
|
||||
lookup: func(uint64) (ServerOwnerInfo, bool) {
|
||||
return ServerOwnerInfo{}, false
|
||||
},
|
||||
wantID: 0,
|
||||
wantHasName: false,
|
||||
},
|
||||
{
|
||||
// Known user → username flows through to the wire so the
|
||||
// admin frontend can show it without a separate /user fetch
|
||||
// (which members cannot call anyway).
|
||||
name: "known_user_has_username",
|
||||
uid: 42,
|
||||
lookup: func(uid uint64) (ServerOwnerInfo, bool) {
|
||||
return ServerOwnerInfo{ID: uid, Username: "alice"}, true
|
||||
},
|
||||
wantID: 42,
|
||||
wantHasName: true,
|
||||
wantName: "alice",
|
||||
},
|
||||
{
|
||||
// Deleted user → lookup returns ok=false. The wire still
|
||||
// carries owner.id so the frontend can render an "Unknown
|
||||
// user (#id)" placeholder; otherwise the row would silently
|
||||
// appear ownerless and ops would lose the audit trail.
|
||||
name: "deleted_user_keeps_id_without_username",
|
||||
uid: 999,
|
||||
lookup: func(uint64) (ServerOwnerInfo, bool) {
|
||||
return ServerOwnerInfo{}, false
|
||||
},
|
||||
wantID: 999,
|
||||
wantHasName: false,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tc := range tests {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
ServerOwnerLookup = tc.lookup
|
||||
s := &Server{Common: Common{ID: 7, UserID: tc.uid}, Name: "srv"}
|
||||
|
||||
raw, err := json.Marshal(s)
|
||||
if err != nil {
|
||||
t.Fatalf("marshal: %v", err)
|
||||
}
|
||||
|
||||
var got struct {
|
||||
Owner *ServerOwnerInfo `json:"owner"`
|
||||
// Owner must never appear as the raw uid via Common.UserID;
|
||||
// the Common.UserID json tag is "-" and a regression that
|
||||
// flips it to "user_id" would expose internal owner ids
|
||||
// to the wire bypassing the lookup-controlled rendering.
|
||||
UserID *uint64 `json:"user_id,omitempty"`
|
||||
}
|
||||
if err := json.Unmarshal(raw, &got); err != nil {
|
||||
t.Fatalf("unmarshal: %v", err)
|
||||
}
|
||||
|
||||
if got.UserID != nil {
|
||||
t.Fatalf("Common.UserID must not appear on the wire as user_id, got %d", *got.UserID)
|
||||
}
|
||||
if got.Owner == nil {
|
||||
t.Fatalf("owner field must always be present, raw=%s", raw)
|
||||
}
|
||||
if got.Owner.ID != tc.wantID {
|
||||
t.Fatalf("owner.id=%d, want %d", got.Owner.ID, tc.wantID)
|
||||
}
|
||||
if tc.wantHasName {
|
||||
if got.Owner.Username != tc.wantName {
|
||||
t.Fatalf("owner.username=%q, want %q", got.Owner.Username, tc.wantName)
|
||||
}
|
||||
} else if got.Owner.Username != "" {
|
||||
t.Fatalf("owner.username must be omitted for uid=%d, got %q", tc.uid, got.Owner.Username)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// When no lookup is installed (tests / headless tools), MarshalJSON must
|
||||
// still emit a minimal owner record so consumers do not crash on missing
|
||||
// fields. Without this guard a future refactor could silently drop the
|
||||
// owner key entirely whenever the hook is nil.
|
||||
func TestServerMarshalJSONEmitsOwnerWithoutLookup(t *testing.T) {
|
||||
original := ServerOwnerLookup
|
||||
t.Cleanup(func() { ServerOwnerLookup = original })
|
||||
ServerOwnerLookup = nil
|
||||
|
||||
raw, err := json.Marshal(&Server{Common: Common{ID: 1, UserID: 17}, Name: "srv"})
|
||||
if err != nil {
|
||||
t.Fatalf("marshal: %v", err)
|
||||
}
|
||||
|
||||
var got struct {
|
||||
Owner *ServerOwnerInfo `json:"owner"`
|
||||
}
|
||||
if err := json.Unmarshal(raw, &got); err != nil {
|
||||
t.Fatalf("unmarshal: %v", err)
|
||||
}
|
||||
if got.Owner == nil || got.Owner.ID != 17 || got.Owner.Username != "" {
|
||||
t.Fatalf("expected bare owner record {id:17}, got %+v", got.Owner)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,86 @@
|
||||
package model
|
||||
|
||||
import (
|
||||
"net/http/httptest"
|
||||
"sync"
|
||||
"testing"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
// Server.UserID 在 server-transfer rotation 流程里会被 ServerTransfer 的
|
||||
// Register/revertTransition 改写以反映新所有者,同时 authorizeAgentForUUID
|
||||
// 在每次 agent RPC 里读取它。原实现两处都是裸字段访问,race detector 会
|
||||
// 报告 data race;这是 review 评分 75 的真实问题。
|
||||
//
|
||||
// 修复后所有并发读写都走 SetUserID/GetUserID 的 atomic 包装,本测试在
|
||||
// `go test -race` 下应该完全跑干净。
|
||||
func TestServerUserIDConcurrentAccessIsRaceFree(t *testing.T) {
|
||||
s := &Server{}
|
||||
|
||||
const (
|
||||
writers = 4
|
||||
readers = 8
|
||||
rounds = 500
|
||||
)
|
||||
var wg sync.WaitGroup
|
||||
wg.Add(writers + readers)
|
||||
|
||||
for i := 0; i < writers; i++ {
|
||||
uid := uint64(i + 1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
for j := 0; j < rounds; j++ {
|
||||
s.SetUserID(uid)
|
||||
}
|
||||
}()
|
||||
}
|
||||
for i := 0; i < readers; i++ {
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
for j := 0; j < rounds; j++ {
|
||||
_ = s.GetUserID()
|
||||
}
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
}
|
||||
|
||||
// Common.HasPermission 是 server-transfer 旋转下与 SetUserID 并发的主要读者
|
||||
// 之一:dashboard 各 controller 的 listHandler post-filter 在 transfer 窗口
|
||||
// 内不断对同一 *Server 调用 HasPermission,而 Register/revertTransition 同
|
||||
// 时通过 SetUserID 改写所属用户。原实现的 `user.ID == c.UserID` 是裸读,会
|
||||
// 与 atomic.StoreUint64 形成 data race(go test -race 必爆)。修复后改成走
|
||||
// GetUserID() 走 atomic 协议。这个测试就是用来在 -race 下钉死该不变量的。
|
||||
func TestCommonHasPermissionConcurrentWithSetUserIDIsRaceFree(t *testing.T) {
|
||||
s := &Server{Common: Common{ID: 1}}
|
||||
|
||||
const (
|
||||
writers = 4
|
||||
readers = 8
|
||||
rounds = 500
|
||||
)
|
||||
var wg sync.WaitGroup
|
||||
wg.Add(writers + readers)
|
||||
|
||||
for i := 0; i < writers; i++ {
|
||||
uid := uint64(i + 1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
for j := 0; j < rounds; j++ {
|
||||
s.SetUserID(uid)
|
||||
}
|
||||
}()
|
||||
}
|
||||
for i := 0; i < readers; i++ {
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
ctx, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
ctx.Set(CtxKeyAuthorizedUser, &User{Common: Common{ID: 2}, Role: RoleMember})
|
||||
for j := 0; j < rounds; j++ {
|
||||
_ = s.HasPermission(ctx)
|
||||
}
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
}
|
||||
@@ -0,0 +1,91 @@
|
||||
package model
|
||||
|
||||
import (
|
||||
"context"
|
||||
"sync"
|
||||
"testing"
|
||||
|
||||
pb "github.com/nezhahq/nezha/proto"
|
||||
)
|
||||
|
||||
// raceProbeStream is the smallest fake of pb.NezhaService_RequestTaskServer
|
||||
// the race probe needs. We only call Send on it from the test; the embedded
|
||||
// interface satisfies the rest of the contract with nil-panicking methods we
|
||||
// never invoke.
|
||||
type raceProbeStream struct {
|
||||
pb.NezhaService_RequestTaskServer
|
||||
}
|
||||
|
||||
func (raceProbeStream) Send(*pb.Task) error { return nil }
|
||||
func (raceProbeStream) Context() context.Context { return context.Background() }
|
||||
|
||||
// model.Server.TaskStream is read from many goroutines (singleton cron pushes,
|
||||
// transfer ApplyConfig pushes, terminal/fm proxies, dashboard rpc keepalives,
|
||||
// per-server batch pushes) and written from exactly one (the gRPC RequestTask
|
||||
// goroutine on every fresh agent connection). The bare-field access pattern
|
||||
// `if s.TaskStream != nil { s.TaskStream.Send(...) }` is a data race on the
|
||||
// interface header (two-word value) and can torn-read into a panic on a
|
||||
// reconnect. This test pins down "concurrent set + send must be race-free"
|
||||
// using the Go race detector — without the fix, `go test -race` reports a
|
||||
// data race on TaskStream; with the fix the field is encapsulated behind
|
||||
// atomic methods and the test runs clean. Without `-race` both versions are
|
||||
// indistinguishable, so this test is only meaningful under the race flag —
|
||||
// run it from CI as `go test -race ./model/`.
|
||||
func TestServerTaskStreamConcurrentAccessIsRaceFree(t *testing.T) {
|
||||
s := &Server{}
|
||||
InitServer(s)
|
||||
|
||||
const (
|
||||
writers = 4
|
||||
readers = 8
|
||||
rounds = 200
|
||||
)
|
||||
var wg sync.WaitGroup
|
||||
wg.Add(writers + readers)
|
||||
|
||||
for i := 0; i < writers; i++ {
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
for j := 0; j < rounds; j++ {
|
||||
s.SetTaskStream(raceProbeStream{})
|
||||
s.SetTaskStream(nil)
|
||||
}
|
||||
}()
|
||||
}
|
||||
for i := 0; i < readers; i++ {
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
for j := 0; j < rounds; j++ {
|
||||
if stream := s.GetTaskStream(); stream != nil {
|
||||
_ = stream.Send(nil)
|
||||
}
|
||||
}
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
}
|
||||
|
||||
func TestServerClearTaskStreamIfCurrentClearsOnlyMatchingStream(t *testing.T) {
|
||||
s := &Server{}
|
||||
InitServer(s)
|
||||
|
||||
first := &raceProbeStream{}
|
||||
second := &raceProbeStream{}
|
||||
|
||||
s.SetTaskStream(first)
|
||||
if !s.ClearTaskStreamIfCurrent(first) {
|
||||
t.Fatal("matching current stream must be cleared")
|
||||
}
|
||||
if got := s.GetTaskStream(); got != nil {
|
||||
t.Fatalf("expected cleared task stream, got %T", got)
|
||||
}
|
||||
|
||||
s.SetTaskStream(first)
|
||||
s.SetTaskStream(second)
|
||||
if s.ClearTaskStreamIfCurrent(first) {
|
||||
t.Fatal("stale stream cleanup must not clear a newer stream")
|
||||
}
|
||||
if got := s.GetTaskStream(); got != second {
|
||||
t.Fatalf("expected newer stream to remain published, got %T", got)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,128 @@
|
||||
package model
|
||||
|
||||
import (
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
// ServerTransferStatus represents the lifecycle state of a server ownership
|
||||
// transfer. A transfer's life starts at Pending (server.user_id has been
|
||||
// flipped to the new owner; agent still authenticates with the old owner's
|
||||
// AgentSecret) and ends in exactly one of the terminal states.
|
||||
type ServerTransferStatus uint8
|
||||
|
||||
const (
|
||||
// ServerTransferStatusPending means the dashboard has flipped Server.UserID
|
||||
// to the new owner and queued an ApplyConfig task to swap the agent's
|
||||
// client_secret. Auth still accepts the old owner's AgentSecret for this
|
||||
// UUID until verification arrives or the transfer times out.
|
||||
ServerTransferStatusPending ServerTransferStatus = iota
|
||||
// ServerTransferStatusVerified means the agent successfully reconnected
|
||||
// using the new owner's AgentSecret. Auth no longer tolerates the old
|
||||
// owner's secret on this UUID.
|
||||
ServerTransferStatusVerified
|
||||
// ServerTransferStatusFailed means the agent explicitly reported the
|
||||
// ApplyConfig task as unsuccessful (e.g. DisableCommandExecute). The
|
||||
// dashboard has rolled Server.UserID back to FromUserID.
|
||||
ServerTransferStatusFailed
|
||||
// ServerTransferStatusTimeout means the verification window expired
|
||||
// without the agent reconnecting under the new secret. The dashboard has
|
||||
// rolled Server.UserID back to FromUserID.
|
||||
ServerTransferStatusTimeout
|
||||
// ServerTransferStatusCancelled means an administrator cancelled the
|
||||
// transfer before any verification event was observed. The dashboard has
|
||||
// rolled Server.UserID back to FromUserID.
|
||||
ServerTransferStatusCancelled
|
||||
)
|
||||
|
||||
// IsTerminal reports whether the status represents a settled transfer. Only
|
||||
// terminal transfers are eligible for retry and they will never be in the
|
||||
// pending index.
|
||||
func (s ServerTransferStatus) IsTerminal() bool {
|
||||
return s != ServerTransferStatusPending
|
||||
}
|
||||
|
||||
// ServerTransfer records a single attempt to transfer ownership of one server
|
||||
// to another user. It is the source of truth for the auth-tolerance window
|
||||
// during a transfer — service/rpc.authorizeAgentForUUID consults the pending
|
||||
// index built from this table to decide whether to accept the old owner's
|
||||
// AgentSecret on the affected UUID.
|
||||
//
|
||||
// Naming note: the existing model.Transfer records hourly traffic snapshots
|
||||
// and is unrelated. This entity is named ServerTransfer to disambiguate.
|
||||
type ServerTransfer struct {
|
||||
Common
|
||||
ServerID uint64 `json:"server_id" gorm:"index"`
|
||||
FromUserID uint64 `json:"from_user_id"`
|
||||
ToUserID uint64 `json:"to_user_id"`
|
||||
InitiatorID uint64 `json:"initiator_id"`
|
||||
Status ServerTransferStatus `json:"status" gorm:"index"`
|
||||
LastError string `json:"last_error,omitempty"`
|
||||
AckedAt *time.Time `json:"acked_at,omitempty"`
|
||||
// HandshakeSecret is a per-transfer random credential that PushIfOnline
|
||||
// delivers in place of the destination user's global AgentSecret. The
|
||||
// agent treats it as a temporary handshake token: it rotates to this
|
||||
// secret on the 10s reload, reconnects, and the dashboard's auth path
|
||||
// recognises it as proof of transfer delivery (MarkVerified). It is
|
||||
// scoped to this single transfer and to this single UUID — leaking it
|
||||
// to the previous owner who hijacks the stream still does NOT expose
|
||||
// the destination user's other agents. Never returned to API clients.
|
||||
HandshakeSecret string `json:"-" gorm:"type:char(32)"`
|
||||
// RevertHandshakeSecret is the same idea for the rollback path: when
|
||||
// the dashboard pushes a revert ApplyConfig over a stream now held by
|
||||
// the destination user, we must not embed the source user's global
|
||||
// AgentSecret. Instead the agent rotates back through this token, which
|
||||
// is recognised by the auth path during the revert window only.
|
||||
RevertHandshakeSecret string `json:"-" gorm:"type:char(32)"`
|
||||
}
|
||||
|
||||
// HasPermission overrides Common.HasPermission so a transfer is visible to
|
||||
// admins, the source user, the destination user, and the initiator. Listing
|
||||
// uses this to filter what the caller can see; mutating endpoints (cancel,
|
||||
// retry) layer additional checks on top.
|
||||
func (t *ServerTransfer) HasPermission(ctx *gin.Context) bool {
|
||||
auth, ok := ctx.Get(CtxKeyAuthorizedUser)
|
||||
if !ok {
|
||||
return false
|
||||
}
|
||||
user := *auth.(*User)
|
||||
if user.Role == RoleAdmin {
|
||||
return true
|
||||
}
|
||||
return user.ID == t.FromUserID || user.ID == t.ToUserID || user.ID == t.InitiatorID
|
||||
}
|
||||
|
||||
// BatchMoveServerResultStatus is the per-server outcome returned by the
|
||||
// batch-move endpoint. It maps to TransferStatus for transfers that were
|
||||
// successfully created, plus extra synchronous-failure modes (permission,
|
||||
// duplicate active transfer, missing server) that never produce a row.
|
||||
type BatchMoveServerResultStatus string
|
||||
|
||||
const (
|
||||
// BatchMoveServerResultPending: ServerTransfer row created, agent push
|
||||
// in progress. Callers should watch the WS for terminal status.
|
||||
BatchMoveServerResultPending BatchMoveServerResultStatus = "pending"
|
||||
// BatchMoveServerResultPermissionDenied: caller cannot move this server.
|
||||
BatchMoveServerResultPermissionDenied BatchMoveServerResultStatus = "permission_denied"
|
||||
// BatchMoveServerResultAlreadyTransferring: server already has an in-flight
|
||||
// ServerTransfer row, cancel or wait first.
|
||||
BatchMoveServerResultAlreadyTransferring BatchMoveServerResultStatus = "already_transferring"
|
||||
// BatchMoveServerResultServerNotFound: server id does not exist.
|
||||
BatchMoveServerResultServerNotFound BatchMoveServerResultStatus = "server_not_found"
|
||||
// BatchMoveServerResultSameOwner: target user already owns this server.
|
||||
BatchMoveServerResultSameOwner BatchMoveServerResultStatus = "same_owner"
|
||||
// BatchMoveServerResultAgentTooOld: agent build does not understand
|
||||
// TaskTypeServerTransferApply, so the rotation would never complete and
|
||||
// dashboard refuses to start it. Operator must upgrade the agent.
|
||||
BatchMoveServerResultAgentTooOld BatchMoveServerResultStatus = "agent_too_old"
|
||||
)
|
||||
|
||||
// BatchMoveServerResult is one entry in the batchMoveServer response, one
|
||||
// per requested server id, in the same order.
|
||||
type BatchMoveServerResult struct {
|
||||
ServerID uint64 `json:"server_id"`
|
||||
Status BatchMoveServerResultStatus `json:"status"`
|
||||
TransferID uint64 `json:"transfer_id,omitempty"`
|
||||
Error string `json:"error,omitempty"`
|
||||
}
|
||||
+6
-1
@@ -26,6 +26,10 @@ const (
|
||||
TaskTypeFM
|
||||
TaskTypeReportConfig
|
||||
TaskTypeApplyConfig
|
||||
// TaskTypeServerTransferApply: per-transfer credential rotation.
|
||||
// Pre-transfer agents do not recognise this type — dashboard MUST gate
|
||||
// transfers on agent capability before pushing.
|
||||
TaskTypeServerTransferApply
|
||||
)
|
||||
|
||||
type TerminalTask struct {
|
||||
@@ -133,7 +137,8 @@ func IsServiceSentinelNeeded(t uint64) bool {
|
||||
switch t {
|
||||
case TaskTypeCommand, TaskTypeTerminalGRPC, TaskTypeUpgrade,
|
||||
TaskTypeKeepalive, TaskTypeNAT, TaskTypeFM,
|
||||
TaskTypeReportConfig, TaskTypeApplyConfig:
|
||||
TaskTypeReportConfig, TaskTypeApplyConfig,
|
||||
TaskTypeServerTransferApply:
|
||||
return false
|
||||
default:
|
||||
return true
|
||||
|
||||
@@ -32,6 +32,7 @@ type User struct {
|
||||
|
||||
type UserInfo struct {
|
||||
Role Role
|
||||
Username string
|
||||
AgentSecret string
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user