SimpleRemoter/server/go/licensing/quota.go

package licensing

import (
	"encoding/json"
	"os"
	"sync"
	"time"
)

// Tier names. The free tier is hard-coded into the License Server (no JWT
// needed); trial / paid are encoded in the issued JWT's "tier" claim along
// with "max_devices".
const (
	TierFree  = "free"
	TierTrial = "trial"
	TierPaid  = "paid"
)

// FreeMaxDevices is the per-deployment device cap for unauthenticated
// (no JWT) callers. Trial defaults to 20 if the JWT omits max_devices;
// paid customers MUST have an explicit max_devices in their JWT.
const (
	FreeMaxDevices  = 2
	TrialMaxDevices = 20
)

// persistedQuota is the on-disk snapshot format. V=1 is the current schema.
type persistedQuota struct {
	V         int                 `json:"v"`         // schema version
	Customers map[string][]string `json:"customers"` // sub → []clientID
}

// quotaTracker maintains the active-device set per customer. Customers are
// identified by the JWT "sub" claim. The set is keyed by clientID (uint64
// from the device, stringified) — same device coming back through the
// same License Server is one slot, not two.
//
// Eviction: any clientID not seen in /sign or /license/heartbeat within
// the eviction window is silently dropped from the active set. Default
// window is 5 minutes (twice the heartbeat interval).
//
// Persistence: when statePath is set, the sub→clientID map is written
// atomically to disk on every structural change (device added or evicted).
// Load() restores the state on startup with fresh timestamps so a License
// Server restart does not open a quota-bypass window.
//
// Empty customer entries are reaped at the end of each mutation.
type quotaTracker struct {
	evictAfter time.Duration
	statePath  string // "" = no persistence
	log        Logger // nil = silent

	mu       sync.Mutex
	customer map[string]*customerState // sub claim → state
}

type customerState struct {
	devices map[string]time.Time // clientID → last activity
}

func newQuotaTracker(evictAfter time.Duration) *quotaTracker {
	return &quotaTracker{
		evictAfter: evictAfter,
		customer:   make(map[string]*customerState),
	}
}

// Load reads the persisted state from statePath and restores each clientID
// with timestamp time.Now() so restored devices survive the initial eviction
// window long enough to heartbeat or re-sign. A missing or corrupt file is
// silently ignored so the server starts cleanly on first run.
func (q *quotaTracker) Load() error {
	if q.statePath == "" {
		return nil
	}
	data, err := os.ReadFile(q.statePath)
	if err != nil {
		if os.IsNotExist(err) {
			return nil
		}
		return err
	}
	var p persistedQuota
	if err := json.Unmarshal(data, &p); err != nil {
		if q.log != nil {
			q.log.Warn("quota: corrupt state file %s (starting empty): %v", q.statePath, err)
		}
		return nil
	}
	q.mu.Lock()
	defer q.mu.Unlock()
	now := time.Now()
	restored := 0
	for sub, ids := range p.Customers {
		if len(ids) == 0 {
			continue
		}
		st := &customerState{devices: make(map[string]time.Time, len(ids))}
		for _, cid := range ids {
			st.devices[cid] = now
			restored++
		}
		q.customer[sub] = st
	}
	if q.log != nil && restored > 0 {
		q.log.Info("quota: restored %d device slot(s) from %s", restored, q.statePath)
	}
	return nil
}

// snapshotLocked returns a sub→[]clientID map of the current state.
// Caller must hold q.mu.
func (q *quotaTracker) snapshotLocked() map[string][]string {
	out := make(map[string][]string, len(q.customer))
	for sub, st := range q.customer {
		if len(st.devices) == 0 {
			continue
		}
		ids := make([]string, 0, len(st.devices))
		for cid := range st.devices {
			ids = append(ids, cid)
		}
		out[sub] = ids
	}
	return out
}

// save writes snap atomically (temp file + rename). No-op when statePath is
// empty or snap is nil.
func (q *quotaTracker) save(snap map[string][]string) {
	if q.statePath == "" || snap == nil {
		return
	}
	data, err := json.Marshal(persistedQuota{V: 1, Customers: snap})
	if err != nil {
		if q.log != nil {
			q.log.Warn("quota: marshal state: %v", err)
		}
		return
	}
	tmp := q.statePath + ".tmp"
	if err := os.WriteFile(tmp, data, 0600); err != nil {
		if q.log != nil {
			q.log.Warn("quota: write state to %s: %v", tmp, err)
		}
		return
	}
	if err := os.Rename(tmp, q.statePath); err != nil {
		if q.log != nil {
			q.log.Warn("quota: rename %s → %s: %v", tmp, q.statePath, err)
		}
	}
}

// evictLocked drops stale entries from st.devices. Returns the number removed.
// Caller must hold q.mu.
func (q *quotaTracker) evictLocked(st *customerState) int {
	cutoff := time.Now().Add(-q.evictAfter)
	removed := 0
	for cid, last := range st.devices {
		if last.Before(cutoff) {
			delete(st.devices, cid)
			removed++
		}
	}
	return removed
}

// reapEmptyLocked deletes sub entries whose device sets are empty. This
// prevents long-running License Servers from leaking memory across
// churned trial customers. Caller must hold q.mu.
func (q *quotaTracker) reapEmptyLocked(sub string) {
	if st, ok := q.customer[sub]; ok && len(st.devices) == 0 {
		delete(q.customer, sub)
	}
}

// Reserve atomically checks whether <clientID> can be admitted under <sub>
// without exceeding <maxDevices>, and if so records it as active. Returns
// (active_count_after, accepted).
//
// "Accepted" means the slot is now reserved. "Rejected" leaves the device
// map untouched — a denied request must NOT permanently consume a slot
// just because the caller asked.
//
// If <clientID> is already in the set, Reserve refreshes its activity
// timestamp and returns accepted=true regardless of maxDevices (same device
// re-signing is never a quota violation — caps only apply to ADDING).
func (q *quotaTracker) Reserve(sub, clientID string, maxDevices int) (int, bool) {
	q.mu.Lock()

	st, ok := q.customer[sub]
	if !ok {
		st = &customerState{devices: make(map[string]time.Time)}
		q.customer[sub] = st
	}

	evicted := q.evictLocked(st)

	if _, already := st.devices[clientID]; already {
		st.devices[clientID] = time.Now()
		count := len(st.devices)
		var snap map[string][]string
		if evicted > 0 {
			snap = q.snapshotLocked()
		}
		q.mu.Unlock()
		q.save(snap)
		return count, true
	}

	if len(st.devices)+1 > maxDevices {
		count := len(st.devices)
		var snap map[string][]string
		if evicted > 0 {
			snap = q.snapshotLocked()
		}
		q.mu.Unlock()
		q.save(snap)
		return count, false
	}

	// New device admitted: always persist so a restart sees this slot.
	st.devices[clientID] = time.Now()
	count := len(st.devices)
	snap := q.snapshotLocked()
	q.mu.Unlock()
	q.save(snap)
	return count, true
}

// RefreshExisting bumps the last-activity timestamp for any clientID in
// <clientIDs> that is ALREADY in <sub>'s set. New IDs are silently ignored
// — this is the key anti-tamper property of /license/heartbeat: a
// malicious customer cannot inflate their quota by reporting fake IDs.
//
// Returns the count of IDs that were actually refreshed (i.e. that were
// known to us from a prior Reserve).
func (q *quotaTracker) RefreshExisting(sub string, clientIDs []string) int {
	q.mu.Lock()

	st, ok := q.customer[sub]
	if !ok {
		q.mu.Unlock()
		return 0
	}

	evicted := q.evictLocked(st)

	now := time.Now()
	refreshed := 0
	for _, cid := range clientIDs {
		if _, exists := st.devices[cid]; exists {
			st.devices[cid] = now
			refreshed++
		}
	}

	q.reapEmptyLocked(sub) // eviction may have emptied us

	var snap map[string][]string
	if evicted > 0 {
		snap = q.snapshotLocked()
	}
	q.mu.Unlock()
	q.save(snap)
	return refreshed
}

// Snapshot returns the current active clientIDs for <sub>. Used by
// /license/heartbeat to report the server-side view.
func (q *quotaTracker) Snapshot(sub string) []string {
	q.mu.Lock()

	st, ok := q.customer[sub]
	if !ok {
		q.mu.Unlock()
		return nil
	}

	evicted := q.evictLocked(st)
	out := make([]string, 0, len(st.devices))
	for cid := range st.devices {
		out = append(out, cid)
	}
	q.reapEmptyLocked(sub)

	var snap map[string][]string
	if evicted > 0 {
		snap = q.snapshotLocked()
	}
	q.mu.Unlock()
	q.save(snap)
	return out
}