pangolin/server/internal/devices/context.go

package devices

import (
	"context"
	"time"

	"github.com/wangjia/pangolin/server/internal/codes"
)

// ctxKey is a private type for context keys to avoid collisions.
type ctxKey string

// CtxKeyUserID is the context key under which the authenticated user's
// internal int64 ID is stored by the JWT auth middleware (auth.RequireAuth).
//
// It is the single canonical key shared with the codes and auth modules
// (auth.UserIDFromContext reads codes.CtxKeyUserID). Aliasing it here — rather
// than declaring a distinct devices.ctxKey("user_id") — ensures the devices
// middleware/handlers resolve the same value the auth middleware injects.
const CtxKeyUserID = codes.CtxKeyUserID

// ctxKeyPlan is the context key under which the resolved subscription Plan is
// stored by SubscriptionMiddleware.
const ctxKeyPlan ctxKey = "plan"

// WithUserID returns a copy of ctx carrying the authenticated user ID.
// Primarily used in tests and by the auth middleware.
func WithUserID(ctx context.Context, userID int64) context.Context {
	return context.WithValue(ctx, CtxKeyUserID, userID)
}

// UserIDFromContext extracts the authenticated user ID from ctx.
// ok is false when no (valid) user ID is present.
func UserIDFromContext(ctx context.Context) (int64, bool) {
	v, ok := ctx.Value(CtxKeyUserID).(int64)
	if !ok || v == 0 {
		return 0, false
	}
	return v, true
}

// Plan is the resolved effective subscription for a user, injected into the
// request context by SubscriptionMiddleware and consumed by nodes (catalogue
// filtering, connect), usage (daily caps), and the /v1/me summary.
type Plan struct {
	// PlanCode is one of "free", "pro", "team".
	PlanCode string
	// ExpiresAt is the effective subscription expiry (UTC); nil for the free
	// fallback when the user has no active paid/trial subscription.
	ExpiresAt *time.Time
	// MaxDevices is the per-plan device cap (free 1 / pro 5 / team 10).
	MaxDevices int
	// DailyMinutes is the free-plan daily time cap in minutes; nil = unlimited.
	DailyMinutes *int
	// AdGate is true when the user must watch a rewarded ad to unlock daily use.
	AdGate bool
	// Source is the subscription source: "trial", "code", "admin", or "free"
	// for the no-subscription fallback.
	Source string
}

// WithPlan returns a copy of ctx carrying the resolved Plan.
func WithPlan(ctx context.Context, p Plan) context.Context {
	return context.WithValue(ctx, ctxKeyPlan, p)
}

// PlanFromCtx extracts the resolved Plan from ctx. ok is false when
// SubscriptionMiddleware has not run.
func PlanFromCtx(ctx context.Context) (Plan, bool) {
	p, ok := ctx.Value(ctxKeyPlan).(Plan)
	return p, ok
}

// planTier maps a plan code to a numeric tier (higher = better) for comparison.
func planTier(code string) int {
	switch code {
	case "team":
		return 3
	case "pro":
		return 2
	default: // "free" or unknown
		return 1
	}
}