SimpleRemoter/server/go/licensing/quota.go

package licensing

import (
	"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
)

// 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. This stops
// a never-heartbeating customer from holding slots forever. Default
// window is twice the heartbeat interval the customer reports at (5 min).
//
// Empty customer entries are reaped at the end of each mutation so the
// outer map doesn't accumulate sub claims of expired contracts.
type quotaTracker struct {
	evictAfter time.Duration

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

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

// 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()
	defer q.mu.Unlock()

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

	q.evictLocked(st)

	if _, already := st.devices[clientID]; already {
		st.devices[clientID] = time.Now()
		return len(st.devices), true
	}

	if len(st.devices)+1 > maxDevices {
		// Don't reap on rejection — the customer might be at exactly cap
		// with valid devices, and an empty map would lose info.
		return len(st.devices), false
	}

	st.devices[clientID] = time.Now()
	return len(st.devices), 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()
	defer q.mu.Unlock()

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

	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
	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()
	defer q.mu.Unlock()
	st, ok := q.customer[sub]
	if !ok {
		return nil
	}
	q.evictLocked(st)
	out := make([]string, 0, len(st.devices))
	for cid := range st.devices {
		out = append(out, cid)
	}
	q.reapEmptyLocked(sub)
	return out
}