pangolin/server/internal/agentd/singbox.go

package agentd

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io/fs"
	"os"
	"path/filepath"
	"sort"
	"sync"
	"time"

	agentv1 "github.com/wangjia/pangolin/server/internal/pb/agentv1"
)

// Restarter applies a freshly-rendered sing-box config. sing-box has no hot
// reload, so the production implementation restarts the systemd unit (a ~1-2s
// data-plane blip). Tests inject a fake.
type Restarter interface {
	Restart(ctx context.Context) error
}

// Cred is the agent's in-memory view of one data-plane credential. It is the
// authoritative local state together with the rendered sing-box config.
type Cred struct {
	DpUUID    string
	Protocol  agentv1.Protocol
	Flow      string
	ExpiresAt int64 // unix seconds; 0 = no expiry
}

// persistedCred is the on-disk shape. Per the no-state invariant we persist ONLY
// dp_uuid + expires_at — protocol/flow are reconstructed (defaults) on load and
// then corrected by the next Register reconciliation.
type persistedCred struct {
	DpUUID    string `json:"dp_uuid"`
	ExpiresAt int64  `json:"expires_at"`
}

type persistedState struct {
	Creds []persistedCred `json:"creds"`
}

// SingBox owns the credential table, state.json persistence, sing-box config
// rendering and the debounced restart loop.
type SingBox struct {
	cfg       Config
	restarter Restarter
	clock     func() time.Time

	// OnExpire, if set, is invoked (outside the lock) when the TTL sweep removes
	// a credential, so the agent can surface the removal upstream.
	OnExpire func(dpUUID string)

	mu            sync.Mutex
	creds         map[string]*Cred
	reality       *agentv1.RealityInbound
	hy2           *agentv1.Hy2Inbound
	configVersion int64

	// debounce machinery
	dirty   chan struct{}
	done    chan struct{}
	closed  bool
	startWG sync.WaitGroup
}

// NewSingBox builds a manager. restarter may be nil (renders only, no restart).
func NewSingBox(cfg Config, restarter Restarter) *SingBox {
	if restarter == nil {
		restarter = noopRestarter{}
	}
	return &SingBox{
		cfg:       cfg.withDefaults(),
		restarter: restarter,
		clock:     time.Now,
		creds:     make(map[string]*Cred),
		dirty:     make(chan struct{}, 1),
		done:      make(chan struct{}),
	}
}

// LoadState restores the credential table from state.json after a restart. Only
// dp_uuid + expires_at survive; protocol defaults to BOTH and flow to the
// configured default until the next Register overwrites them. Expired entries are
// dropped on load.
func (s *SingBox) LoadState() error {
	data, err := os.ReadFile(s.cfg.StatePath())
	if errors.Is(err, fs.ErrNotExist) {
		return nil
	}
	if err != nil {
		return fmt.Errorf("agentd: read state: %w", err)
	}
	var ps persistedState
	if err := json.Unmarshal(data, &ps); err != nil {
		return fmt.Errorf("agentd: parse state: %w", err)
	}
	now := s.clock().Unix()
	s.mu.Lock()
	defer s.mu.Unlock()
	for _, pc := range ps.Creds {
		if pc.ExpiresAt != 0 && pc.ExpiresAt <= now {
			continue
		}
		s.creds[pc.DpUUID] = &Cred{
			DpUUID:    pc.DpUUID,
			Protocol:  agentv1.ProtocolBoth,
			Flow:      DefaultFlow,
			ExpiresAt: pc.ExpiresAt,
		}
	}
	return nil
}

// saveStateLocked writes state.json atomically. Caller holds s.mu.
func (s *SingBox) saveStateLocked() error {
	ps := persistedState{Creds: make([]persistedCred, 0, len(s.creds))}
	for _, c := range s.creds {
		ps.Creds = append(ps.Creds, persistedCred{DpUUID: c.DpUUID, ExpiresAt: c.ExpiresAt})
	}
	sort.Slice(ps.Creds, func(i, j int) bool { return ps.Creds[i].DpUUID < ps.Creds[j].DpUUID })
	data, err := json.MarshalIndent(ps, "", "  ")
	if err != nil {
		return err
	}
	return atomicWrite(s.cfg.StatePath(), data, 0o600)
}

// ─── mutations (all idempotent, all schedule a render) ───────────────────────────

// Upsert adds or replaces a credential.
func (s *SingBox) Upsert(c *Cred) {
	s.mu.Lock()
	cp := *c
	if cp.Flow == "" {
		cp.Flow = DefaultFlow
	}
	s.creds[cp.DpUUID] = &cp
	_ = s.saveStateLocked()
	s.mu.Unlock()
	s.markDirty()
}

// DpUUIDs returns a snapshot of the currently provisioned dp_uuids.
// Used by the usage source to attribute node-total traffic to active users.
func (s *SingBox) DpUUIDs() []string {
	s.mu.Lock()
	defer s.mu.Unlock()
	out := make([]string, 0, len(s.creds))
	for u := range s.creds {
		out = append(out, u)
	}
	return out
}

// Revoke removes a credential. No-op if absent (idempotent).
func (s *SingBox) Revoke(dpUUID string) {
	s.mu.Lock()
	_, existed := s.creds[dpUUID]
	delete(s.creds, dpUUID)
	if existed {
		_ = s.saveStateLocked()
	}
	s.mu.Unlock()
	if existed {
		s.markDirty()
	}
}

// Rotate installs new alongside old, keeping old alive until graceUntil. The TTL
// sweep removes old when its grace expires, giving "宽限期内新旧并存".
func (s *SingBox) Rotate(oldDpUUID string, newCred *Cred, graceUntil int64) {
	s.mu.Lock()
	if old, ok := s.creds[oldDpUUID]; ok && oldDpUUID != newCred.DpUUID {
		// Clamp old credential's expiry to the grace deadline.
		if graceUntil > 0 && (old.ExpiresAt == 0 || graceUntil < old.ExpiresAt) {
			old.ExpiresAt = graceUntil
		}
	}
	nc := *newCred
	if nc.Flow == "" {
		nc.Flow = DefaultFlow
	}
	s.creds[nc.DpUUID] = &nc
	_ = s.saveStateLocked()
	s.mu.Unlock()
	s.markDirty()
}

// ApplyConfig replaces the node-wide inbound parameters and (if the snapshot
// carries credentials) the full credential table. Used by APPLY_CONFIG commands
// and full resync.
func (s *SingBox) ApplyConfig(snap *agentv1.ConfigSnapshot, replaceCreds bool) {
	s.mu.Lock()
	if snap.Reality != nil {
		s.reality = snap.Reality
	}
	if snap.Hy2 != nil {
		s.hy2 = snap.Hy2
	}
	if snap.ConfigVersion != 0 {
		s.configVersion = snap.ConfigVersion
	}
	if replaceCreds {
		s.creds = make(map[string]*Cred, len(snap.Credentials))
		for _, c := range snap.Credentials {
			flow := c.Flow
			if flow == "" {
				flow = DefaultFlow
			}
			s.creds[c.DpUUID] = &Cred{
				DpUUID:    c.DpUUID,
				Protocol:  c.Protocol,
				Flow:      flow,
				ExpiresAt: c.ExpiresAtUnix,
			}
		}
	}
	_ = s.saveStateLocked()
	s.mu.Unlock()
	s.markDirty()
}

// SetConfigVersion records the authoritative config version (e.g. from Register).
func (s *SingBox) SetConfigVersion(v int64) {
	s.mu.Lock()
	s.configVersion = v
	s.mu.Unlock()
}

// ─── read accessors ──────────────────────────────────────────────────────────

func (s *SingBox) ConfigVersion() int64 {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.configVersion
}

func (s *SingBox) OnlinePeers() int {
	s.mu.Lock()
	defer s.mu.Unlock()
	return len(s.creds)
}

// Has reports whether a dp_uuid is currently provisioned.
func (s *SingBox) Has(dpUUID string) bool {
	s.mu.Lock()
	defer s.mu.Unlock()
	_, ok := s.creds[dpUUID]
	return ok
}

// Snapshot returns a copy of the current credential set (sorted by dp_uuid).
func (s *SingBox) Snapshot() []Cred {
	s.mu.Lock()
	defer s.mu.Unlock()
	out := make([]Cred, 0, len(s.creds))
	for _, c := range s.creds {
		out = append(out, *c)
	}
	sort.Slice(out, func(i, j int) bool { return out[i].DpUUID < out[j].DpUUID })
	return out
}

// ─── TTL sweep ───────────────────────────────────────────────────────────────

// sweepExpired removes every credential whose expires_at has passed. Returns the
// removed dp_uuids and whether anything changed.
func (s *SingBox) sweepExpired() []string {
	now := s.clock().Unix()
	s.mu.Lock()
	var removed []string
	for id, c := range s.creds {
		if c.ExpiresAt != 0 && c.ExpiresAt <= now {
			delete(s.creds, id)
			removed = append(removed, id)
		}
	}
	if len(removed) > 0 {
		_ = s.saveStateLocked()
	}
	s.mu.Unlock()
	if len(removed) > 0 {
		s.markDirty()
		if s.OnExpire != nil {
			for _, id := range removed {
				s.OnExpire(id)
			}
		}
	}
	sort.Strings(removed)
	return removed
}

// ─── rendering ───────────────────────────────────────────────────────────────

// RenderConfig builds the sing-box server config JSON from the current state.
// Exported (and pure w.r.t. the snapshot) so tests can assert its contents.
func (s *SingBox) RenderConfig() ([]byte, error) {
	s.mu.Lock()
	creds := make([]Cred, 0, len(s.creds))
	for _, c := range s.creds {
		creds = append(creds, *c)
	}
	reality := s.reality
	hy2 := s.hy2
	s.mu.Unlock()
	sort.Slice(creds, func(i, j int) bool { return creds[i].DpUUID < creds[j].DpUUID })
	return renderSingboxConfig(creds, reality, hy2, s.cfg.DeriveKey)
}

// writeAndRestart renders, writes the config file and restarts sing-box.
func (s *SingBox) writeAndRestart(ctx context.Context) error {
	data, err := s.RenderConfig()
	if err != nil {
		return err
	}
	if err := os.MkdirAll(filepath.Dir(s.cfg.SingboxConfigPath), 0o755); err != nil {
		return fmt.Errorf("agentd: mkdir singbox cfg: %w", err)
	}
	if err := atomicWrite(s.cfg.SingboxConfigPath, data, 0o644); err != nil {
		return err
	}
	return s.restarter.Restart(ctx)
}

// markDirty signals the debounce loop that a render is pending (non-blocking).
func (s *SingBox) markDirty() {
	select {
	case s.dirty <- struct{}{}:
	default:
	}
}

// Run drives the debounce loop: it coalesces bursts of changes within
// DebounceWindow into a single write+restart. It blocks until ctx is cancelled.
func (s *SingBox) Run(ctx context.Context) {
	timer := time.NewTimer(time.Hour)
	timer.Stop()
	pending := false
	for {
		select {
		case <-ctx.Done():
			timer.Stop()
			return
		case <-s.dirty:
			if !pending {
				pending = true
				timer.Reset(s.cfg.DebounceWindow)
			}
		case <-timer.C:
			pending = false
			if err := s.writeAndRestart(ctx); err != nil {
				logf("singbox render/restart failed: %v", err)
			}
		}
	}
}

// noopRestarter renders without restarting (dev / config-check only).
type noopRestarter struct{}

func (noopRestarter) Restart(context.Context) error { return nil }