nezha_domains/model/server.go

package model

import (
	"errors"
	"log"
	"slices"
	"sync"
	"sync/atomic"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/goccy/go-json"
	"gorm.io/gorm"

	pb "github.com/nezhahq/nezha/proto"
)

type Server struct {
	Common

	Name                   string `json:"name"`
	UUID                   string `json:"uuid,omitempty" gorm:"unique"`
	Note                   string `json:"note,omitempty"`           // 管理员可见备注
	PublicNote             string `json:"public_note,omitempty"`    // 公开备注
	DisplayIndex           int    `json:"display_index"`            // 展示排序，越大越靠前
	HideForGuest           bool   `json:"hide_for_guest,omitempty"` // 对游客隐藏
	EnableDDNS             bool   `json:"enable_ddns,omitempty"`    // 启用DDNS
	DDNSProfilesRaw        string `gorm:"default:'[]';column:ddns_profiles_raw" json:"-"`
	OverrideDDNSDomainsRaw string `gorm:"default:'{}';column:override_ddns_domains_raw" json:"-"`

	DDNSProfiles        []uint64            `gorm:"-" json:"ddns_profiles,omitempty" validate:"optional"` // DDNS配置
	OverrideDDNSDomains map[uint64][]string `gorm:"-" json:"override_ddns_domains,omitempty" validate:"optional"`

	Host       *Host      `gorm:"-" json:"host,omitempty"`
	State      *HostState `gorm:"-" json:"state,omitempty"`
	GeoIP      *GeoIP     `gorm:"-" json:"geoip,omitempty"`
	LastActive time.Time  `gorm:"-" json:"last_active,omitempty"`

	// 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. The holder also carries
	// the send mutex so CopyFromRunningServer can share it across the old/new
	// *Server objects that briefly co-exist during edit/transfer rotations —
	// otherwise two *Server pointers would hold the same gRPC stream behind
	// two independent mutexes, defeating the "one SendMsg goroutine per stream"
	// invariant grpc-go requires.
	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.
//
// sendMu lives on the holder (not on *Server) so it is bound to the stream
// itself: CopyFromRunningServer shares the same holder pointer with the new
// *Server, and SendTask locks via the holder, guaranteeing serialized SendMsg
// even when old/new *Server objects briefly co-exist during edit/transfer.
type taskStreamHolder struct {
	s      pb.NezhaService_RequestTaskServer
	sendMu sync.Mutex
}

// 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})
}

// adoptTaskStreamHolder publishes an existing holder verbatim. Used by
// CopyFromRunningServer so the new *Server shares the send mutex (and the
// underlying stream identity) with the old *Server.
func (s *Server) adoptTaskStreamHolder(h *taskStreamHolder) {
	s.taskStream.Store(h)
}

// 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
}

// SendTask dispatches a task on the agent's RequestTask stream under the
// holder's sendMu so concurrent dispatchers (cron, server-transfer
// ApplyConfig, MCP CallAgent, MCP fs.transfer, force-update, report-config)
// cannot violate grpc-go's "one SendMsg goroutine per stream" rule. Returns
// ErrTaskStreamOffline if the agent has not published a stream yet; callers
// that need to distinguish offline from send failure should branch on that.
//
// The mutex is keyed by holder (= by stream) rather than by *Server so that
// edit/transfer rotations replacing *Server in the singleton map still share
// a single lock across the old and new objects pointing at the same stream.
func (s *Server) SendTask(task *pb.Task) error {
	h := s.taskStream.Load()
	if h == nil {
		return ErrTaskStreamOffline
	}
	h.sendMu.Lock()
	defer h.sendMu.Unlock()
	return h.s.Send(task)
}

// ErrTaskStreamOffline is returned by SendTask when the agent has no
// published RequestTask stream. Defined here (rather than in service/rpc)
// so model-layer callers can branch on it without an import cycle.
var ErrTaskStreamOffline = errors.New("agent task stream offline")

func InitServer(s *Server) {
	s.Host = &Host{}
	s.State = &HostState{}
	s.GeoIP = &GeoIP{}
	s.ConfigCache = make(chan any, 1)
}

func (s *Server) CopyFromRunningServer(old *Server) {
	s.Host = old.Host
	s.State = old.State
	s.GeoIP = old.GeoIP
	s.LastActive = old.LastActive
	// Adopt the holder pointer verbatim so the new *Server shares the send
	// mutex AND the stream identity with the old *Server; constructing a fresh
	// holder via SetTaskStream(GetTaskStream()) would give the new object its
	// own mutex, letting two *Server pointers race SendMsg on the same stream
	// during the edit/transfer rotation window.
	s.adoptTaskStreamHolder(old.taskStream.Load())
	s.ConfigCache = old.ConfigCache
	s.PrevTransferInSnapshot = old.PrevTransferInSnapshot
	s.PrevTransferOutSnapshot = old.PrevTransferOutSnapshot
}

func (s *Server) AfterFind(tx *gorm.DB) error {
	if s.DDNSProfilesRaw != "" {
		if err := json.Unmarshal([]byte(s.DDNSProfilesRaw), &s.DDNSProfiles); err != nil {
			log.Println("NEZHA>> Server.AfterFind:", err)
			return nil
		}
	}
	if s.OverrideDDNSDomainsRaw != "" {
		if err := json.Unmarshal([]byte(s.OverrideDDNSDomainsRaw), &s.OverrideDDNSDomains); err != nil {
			log.Println("NEZHA>> Server.AfterFind:", err)
			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)

// OwnerServerIDsLookup is installed by singleton at startup to enumerate the
// IDs of every in-memory Server whose UserID == ownerUID. It exists so that
// Cron.HasPermission / Service.HasPermission can faithfully replay the
// dispatch-side "CoverAll deny-list must cover every PAT-whitelisted-out
// owner server" rule without depending on controller helpers (model must
// not import service/singleton — cycle).
//
// Left nil in tests / headless contexts; callers MUST treat a nil hook as
// "unknown owner topology" and fall back to a conservative decision (the
// existing model.Cron / model.Service code rejects non-trivial CoverAll
// configs for limited PATs when the hook is nil, matching the historical
// behaviour for empty deny-lists).
var OwnerServerIDsLookup func(ownerUID uint64) []uint64

// OwnerIsAdminLookup reports whether ownerUID is an admin user. When the
// owner is admin the runtime dispatch path (CronTrigger, DispatchTask) gates
// on userIsAdmin(cr.UserID) / userIsAdmin(svc.UserID) and fans out across
// EVERY in-memory server — not just the owner's. DenyListSafeForLimitedPAT
// must mirror that fan-out widening or a limited PAT can pass safety check
// with a deny-list that covers only the admin's own servers while the
// runtime still ships the task to foreign-owned servers.
//
// Left nil in tests / headless contexts; callers fall back to
// "owner-set only" which matches the pre-C1 behaviour.
var OwnerIsAdminLookup func(ownerUID uint64) bool

// AllServerIDsLookup returns every in-memory server ID, regardless of
// owner. It is the system-wide fan-out set the runtime uses for
// admin-owned CoverAll cron/service dispatch and is the only correct
// containment set for a server-limited PAT operating on an admin-owned
// resource. Left nil in tests / headless contexts.
var AllServerIDsLookup func() []uint64

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) HasPermission(ctx *gin.Context) bool {
	if !s.Common.HasPermission(ctx) {
		return false
	}
	v, ok := ctx.Get(CtxKeyAPIToken)
	if !ok {
		return true
	}
	tok, ok := v.(APITokenAccessor)
	if !ok || tok == nil {
		return true
	}
	return tok.CanAccessServer(s.GetID())
}

// APITokenWhitelistView is the optional shape an APITokenAccessor can
// implement so DenyListSafeForLimitedPAT can tell unscoped PATs (no
// whitelist → not limited) apart from server-limited ones. Accessors that
// do NOT expose ServerIDs() are treated as potentially limited; the safe
// dispatch path then requires denyList to cover every owner-visible server
// outside what the PAT can reach.
type APITokenWhitelistView interface {
	ServerIDs() []uint64
}

// DenyListSafeForLimitedPAT reports whether a CoverAll/SkipServers deny-list
// keeps a server-limited PAT inside its server_ids whitelist. The runtime
// dispatch path (CronTrigger, DispatchTask) iterates every owner-visible
// server minus denyList; for the PAT to stay contained, every owner server
// outside its whitelist must already appear in denyList. JWT requests and
// PATs with no whitelist are unaffected. Nil OwnerServerIDsLookup forces
// the conservative "reject" branch instead of silently allowing a config
// the runtime would dispatch outside the whitelist.
func DenyListSafeForLimitedPAT(tok APITokenAccessor, ownerUID uint64, denyServers []uint64) bool {
	if tok == nil {
		return true
	}
	if wl, ok := tok.(APITokenWhitelistView); ok && len(wl.ServerIDs()) == 0 {
		return true
	}
	fanout := ownerEffectiveFanoutServerIDs(ownerUID)
	if fanout == nil {
		return false
	}
	denySet := make(map[uint64]struct{}, len(denyServers))
	for _, id := range denyServers {
		denySet[id] = struct{}{}
	}
	for _, id := range fanout {
		if tok.CanAccessServer(id) {
			continue
		}
		if _, denied := denySet[id]; !denied {
			return false
		}
	}
	return true
}

// ownerEffectiveFanoutServerIDs returns the server set the runtime dispatch
// will actually fan out to for a resource owned by ownerUID. Admin owners
// short-circuit cronCanSendToServer / canSendServiceTask via userIsAdmin,
// so the safe containment set is the WHOLE system, not just the admin's
// own servers. Member owners stay bounded to their own server set.
//
// Returns nil to signal "topology unknown" — callers (DenyListSafeForLimitedPAT)
// fall back to fail-closed in that case, matching the historical conservative
// branch when OwnerServerIDsLookup was nil.
func ownerEffectiveFanoutServerIDs(ownerUID uint64) []uint64 {
	if OwnerIsAdminLookup != nil && OwnerIsAdminLookup(ownerUID) {
		if AllServerIDsLookup == nil {
			return nil
		}
		return AllServerIDsLookup()
	}
	if OwnerServerIDsLookup == nil {
		return nil
	}
	return OwnerServerIDsLookup(ownerUID)
}

func (s *Server) SplitList(x []*Server) ([]*Server, []*Server) {
	pri := func(s *Server) bool {
		return s.DisplayIndex == 0
	}

	i := slices.IndexFunc(x, pri)
	if i == -1 {
		return nil, x
	}

	return x[:i], x[i:]
}