Feature(Go): issue-token subcommand for minting customer JWTs

server/go/licensing/quota.go

@@ -1,6 +1,8 @@
 package licensing
 
 import (
+	"encoding/json"
+	"os"
 	"sync"
 	"time"
 )
@@ -22,20 +24,31 @@ const (
 	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. This stops
-// a never-heartbeating customer from holding slots forever. Default
-// window is twice the heartbeat interval the customer reports at (5 min).
+// the eviction window is silently dropped from the active set. Default
+// window is 5 minutes (twice the heartbeat interval).
 //
-// Empty customer entries are reaped at the end of each mutation so the
-// outer map doesn't accumulate sub claims of expired contracts.
+// 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
@@ -52,14 +65,105 @@ func newQuotaTracker(evictAfter time.Duration) *quotaTracker {
 	}
 }
 
-// evictLocked drops stale entries from st.devices. Caller must hold q.mu.
-func (q *quotaTracker) evictLocked(st *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
@@ -84,7 +188,6 @@ func (q *quotaTracker) reapEmptyLocked(sub string) {
 // 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 {
@@ -92,21 +195,38 @@ func (q *quotaTracker) Reserve(sub, clientID string, maxDevices int) (int, bool)
 		q.customer[sub] = st
 	}
 
-	q.evictLocked(st)
+	evicted := q.evictLocked(st)
 
 	if _, already := st.devices[clientID]; already {
 		st.devices[clientID] = time.Now()
-		return len(st.devices), true
+		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 {
-		// 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
+		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()
-	return len(st.devices), true
+	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
@@ -118,14 +238,14 @@ func (q *quotaTracker) Reserve(sub, clientID string, maxDevices int) (int, bool)
 // 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 {
+		q.mu.Unlock()
 		return 0
 	}
 
-	q.evictLocked(st)
+	evicted := q.evictLocked(st)
 
 	now := time.Now()
 	refreshed := 0
@@ -137,6 +257,13 @@ func (q *quotaTracker) RefreshExisting(sub string, clientIDs []string) int {
 	}
 
 	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
 }
 
@@ -144,16 +271,25 @@ func (q *quotaTracker) RefreshExisting(sub string, clientIDs []string) int {
 // /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 {
+		q.mu.Unlock()
 		return nil
 	}
-	q.evictLocked(st)
+
+	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
 }