pangolin/server/internal/devices/middleware.go

package devices

import (
	"net/http"
	"time"

	"github.com/redis/go-redis/v9"
	"github.com/wangjia/pangolin/server/internal/apierr"
)

// Middleware resolves the caller's effective subscription and injects it into
// the request context. It MUST run after the JWT auth middleware (module #2),
// which is responsible for setting CtxKeyUserID.
type Middleware struct {
	svc *Service

	// rdb and cacheTTL reserve a future 60s per-user plan cache. When cacheTTL
	// > 0 and rdb != nil the resolved plan could be cached under
	// "sub:plan:<userID>"; this is intentionally left disabled (cacheTTL == 0)
	// for the initial direct-DB implementation (subscription volume is low).
	// NOTE: a cache must not mask account bans — ban status would need a
	// separate, uncached check before serving any cached plan.
	rdb      *redis.Client
	cacheTTL time.Duration
}

// NewMiddleware creates the subscription-resolving middleware. Pass rdb=nil and
// cacheTTL=0 to use the direct-DB path (current default).
func NewMiddleware(svc *Service, rdb *redis.Client, cacheTTL time.Duration) *Middleware {
	return &Middleware{svc: svc, rdb: rdb, cacheTTL: cacheTTL}
}

// Handler is the net/http middleware. It writes a JSON apierr and stops the
// chain on failure (missing user / banned / internal error); otherwise it
// injects the resolved Plan and calls next.
func (m *Middleware) Handler(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		userID, ok := UserIDFromContext(r.Context())
		if !ok {
			apierr.WriteJSON(w, http.StatusUnauthorized, apierr.ErrUnauthorized)
			return
		}

		plan, apiErr := m.svc.ResolvePlan(r.Context(), userID)
		if apiErr != nil {
			apierr.WriteJSON(w, StatusForError(apiErr), apiErr)
			return
		}

		ctx := WithPlan(r.Context(), plan)
		next.ServeHTTP(w, r.WithContext(ctx))
	})
}

// --------------------------------------------------------------------------
// Helpers consumed by other modules (nodes catalogue filtering / connect, etc.)
// --------------------------------------------------------------------------

// CheckDeviceQuota returns a bilingual error when currentDevices already meets
// or exceeds the plan's device cap, otherwise nil.
func CheckDeviceQuota(p Plan, currentDevices int) *apierr.Error {
	if p.MaxDevices > 0 && currentDevices >= p.MaxDevices {
		return errDeviceLimit(p.MaxDevices)
	}
	return nil
}

// RequirePaidTier returns a 403 error when the plan is the free tier; nil for
// pro/team. Used to gate paid-only nodes and features.
func RequirePaidTier(p Plan) *apierr.Error {
	if planTier(p.PlanCode) < planTier("pro") {
		return &apierr.Error{
			Code:      "PAID_TIER_REQUIRED",
			MessageZH: "该功能仅限付费会员，请升级后使用",
			MessageEn: "This feature requires a paid plan. Please upgrade to continue.",
		}
	}
	return nil
}

// StatusForError maps an apierr.Error code to an HTTP status code.
func StatusForError(e *apierr.Error) int {
	if e == nil {
		return http.StatusOK
	}
	switch e.Code {
	case "UNAUTHORIZED":
		return http.StatusUnauthorized
	case "FORBIDDEN", "ACCOUNT_BANNED", "PAID_TIER_REQUIRED":
		return http.StatusForbidden
	case "NOT_FOUND":
		return http.StatusNotFound
	case "DEVICE_LIMIT_EXCEEDED":
		return http.StatusForbidden
	case "BAD_REQUEST":
		return http.StatusBadRequest
	case "INTERNAL_ERROR":
		return http.StatusInternalServerError
	default:
		return http.StatusBadRequest
	}
}