feat(membership): enforce tenant subscription plans with blocked UI

Add configurable membership plans (free/plus/pro) synced to DB on boot, a subscription guard middleware that blocks tenant endpoints on expired or missing plans, and a MembershipBlockedView that surfaces the reason so the admin can contact the tenant owner. Quota and brand-library reads now honor the active plan's policy JSON and expired subscriptions.
2026-04-18 20:56:05 +08:00
parent 9ec9314aea
commit 63667ed2d1
28 changed files with 1738 additions and 80 deletions
@@ -1,6 +1,7 @@
 package config

 import (
+	"encoding/json"
 	"fmt"
 	"os"
 	"strings"
@@ -16,6 +17,7 @@ type Config struct {
 	RabbitMQ           RabbitMQConfig      `mapstructure:"rabbitmq"`
 	Scheduler          SchedulerConfig     `mapstructure:"scheduler"`
 	MonitoringWorkers  MonitoringConfig    `mapstructure:"monitoring_workers"`
+	Membership         MembershipConfig    `mapstructure:"membership"`
 	BrandLibrary       BrandLibraryConfig  `mapstructure:"brand_library"`
 	Redis              RedisConfig         `mapstructure:"redis"`
 	Qdrant             QdrantConfig        `mapstructure:"qdrant"`
@@ -110,6 +112,64 @@ type MonitoringConfig struct {
 	ProjectionRebuildConcurrency int `mapstructure:"projection_rebuild_concurrency"`
 }

+const (
+	ArticleQuotaCycleLifetime = "lifetime"
+	ArticleQuotaCycleMonthly  = "monthly"
+)
+
+type MembershipConfig struct {
+	DefaultPlanCode string                 `mapstructure:"default_plan_code"`
+	Plans           []MembershipPlanConfig `mapstructure:"plans"`
+}
+
+type MembershipPlanConfig struct {
+	Code                 string        `mapstructure:"code"`
+	Name                 string        `mapstructure:"name"`
+	ArticleGeneration    int           `mapstructure:"article_generation"`
+	ArticleQuotaCycle    string        `mapstructure:"article_quota_cycle"`
+	ImageStorageBytes    int64         `mapstructure:"image_storage_bytes"`
+	CompanyLimit         int           `mapstructure:"company_limit"`
+	SubscriptionDuration time.Duration `mapstructure:"subscription_duration"`
+	ContactAdminOnExpiry bool          `mapstructure:"contact_admin_on_expiry"`
+}
+
+type PlanQuotaPolicy struct {
+	ArticleGeneration        int    `json:"article_generation"`
+	ArticleQuotaCycle        string `json:"article_quota_cycle"`
+	ImageStorageBytes        int64  `json:"image_storage_bytes"`
+	BrandLimit               int    `json:"brand_limit"`
+	SubscriptionDurationSecs int64  `json:"subscription_duration_seconds,omitempty"`
+	ContactAdminOnExpiry     bool   `json:"contact_admin_on_expiry,omitempty"`
+}
+
+func (c MembershipConfig) FindPlan(code string) (MembershipPlanConfig, bool) {
+	normalizedCode := strings.ToLower(strings.TrimSpace(code))
+	for _, plan := range c.Plans {
+		if strings.EqualFold(strings.TrimSpace(plan.Code), normalizedCode) {
+			return plan, true
+		}
+	}
+	return MembershipPlanConfig{}, false
+}
+
+func (p MembershipPlanConfig) QuotaPolicy() PlanQuotaPolicy {
+	policy := PlanQuotaPolicy{
+		ArticleGeneration:    maxInt(p.ArticleGeneration, 0),
+		ArticleQuotaCycle:    normalizeArticleQuotaCycle(p.ArticleQuotaCycle, p.Code),
+		ImageStorageBytes:    maxInt64(p.ImageStorageBytes, 0),
+		BrandLimit:           maxInt(p.CompanyLimit, 0),
+		ContactAdminOnExpiry: p.ContactAdminOnExpiry,
+	}
+	if p.SubscriptionDuration > 0 {
+		policy.SubscriptionDurationSecs = int64(p.SubscriptionDuration / time.Second)
+	}
+	return policy
+}
+
+func (p MembershipPlanConfig) QuotaPolicyJSON() ([]byte, error) {
+	return json.Marshal(p.QuotaPolicy())
+}
+
 type BrandLibraryConfig struct {
 	FreeBrandLimit         int `mapstructure:"free_brand_limit"`
 	PaidBrandLimit         int `mapstructure:"paid_brand_limit"`
@@ -196,17 +256,15 @@ type GenerationConfig struct {

 func Load(configPath string) (*Config, error) {
 	v := viper.New()
-	v.SetConfigFile(configPath)
 	v.AutomaticEnv()
 	v.SetEnvKeyReplacer(strings.NewReplacer(".", "_"))

-	if err := v.ReadInConfig(); err != nil {
+	if err := readConfigWithFallback(v, configPath); err != nil {
 		return nil, fmt.Errorf("read config: %w", err)
 	}

 	// Allow local overrides
-	v.SetConfigFile(strings.Replace(configPath, ".yaml", ".local.yaml", 1))
-	_ = v.MergeInConfig()
+	_ = mergeLocalOverrideWithFallback(v, configPath)

 	var cfg Config
 	if err := v.Unmarshal(&cfg); err != nil {
@@ -217,11 +275,77 @@ func Load(configPath string) (*Config, error) {
 	normalizeRabbitMQConfig(&cfg.RabbitMQ)
 	normalizeSchedulerConfig(&cfg.Scheduler)
 	normalizeMonitoringConfig(&cfg.MonitoringWorkers)
+	normalizeMembershipConfig(&cfg.Membership)
 	normalizeBrandLibraryConfig(&cfg.BrandLibrary)

 	return &cfg, nil
 }

+func readConfigWithFallback(v *viper.Viper, configPath string) error {
+	var lastErr error
+	for _, candidate := range candidateConfigPaths(configPath, false) {
+		v.SetConfigFile(candidate)
+		if err := v.ReadInConfig(); err == nil {
+			return nil
+		} else {
+			lastErr = err
+		}
+	}
+	return lastErr
+}
+
+func mergeLocalOverrideWithFallback(v *viper.Viper, configPath string) error {
+	for _, candidate := range candidateConfigPaths(configPath, true) {
+		v.SetConfigFile(candidate)
+		if err := v.MergeInConfig(); err == nil {
+			return nil
+		}
+	}
+	return nil
+}
+
+func candidateConfigPaths(configPath string, local bool) []string {
+	trimmed := strings.TrimSpace(configPath)
+	if trimmed == "" {
+		return nil
+	}
+
+	paths := make([]string, 0, 2)
+	switch {
+	case strings.HasSuffix(trimmed, ".yaml"):
+		if local {
+			paths = append(paths, strings.TrimSuffix(trimmed, ".yaml")+".local.yaml")
+		} else {
+			paths = append(paths, trimmed)
+		}
+		alt := strings.TrimSuffix(trimmed, ".yaml")
+		if local {
+			paths = append(paths, alt+".local.yml")
+		} else {
+			paths = append(paths, alt+".yml")
+		}
+	case strings.HasSuffix(trimmed, ".yml"):
+		if local {
+			paths = append(paths, strings.TrimSuffix(trimmed, ".yml")+".local.yml")
+		} else {
+			paths = append(paths, trimmed)
+		}
+		alt := strings.TrimSuffix(trimmed, ".yml")
+		if local {
+			paths = append(paths, alt+".local.yaml")
+		} else {
+			paths = append(paths, alt+".yaml")
+		}
+	default:
+		if local {
+			paths = append(paths, trimmed+".local.yaml", trimmed+".local.yml")
+		} else {
+			paths = append(paths, trimmed, trimmed+".yaml", trimmed+".yml")
+		}
+	}
+	return paths
+}
+
 func applyEnvOverrides(cfg *Config) {
 	if cfg == nil {
 		return
@@ -457,6 +581,119 @@ func normalizeBrandLibraryConfig(cfg *BrandLibraryConfig) {
 	}
 }

+func normalizeMembershipConfig(cfg *MembershipConfig) {
+	if cfg == nil {
+		return
+	}
+
+	if strings.TrimSpace(cfg.DefaultPlanCode) == "" {
+		cfg.DefaultPlanCode = "free"
+	}
+
+	if len(cfg.Plans) == 0 {
+		cfg.Plans = defaultMembershipPlans()
+	}
+
+	normalized := make([]MembershipPlanConfig, 0, len(cfg.Plans))
+	seen := make(map[string]struct{}, len(cfg.Plans))
+	for _, plan := range cfg.Plans {
+		plan.Code = strings.ToLower(strings.TrimSpace(plan.Code))
+		if plan.Code == "" {
+			continue
+		}
+		if _, exists := seen[plan.Code]; exists {
+			continue
+		}
+		if strings.TrimSpace(plan.Name) == "" {
+			plan.Name = defaultPlanName(plan.Code)
+		}
+		plan.ArticleQuotaCycle = normalizeArticleQuotaCycle(plan.ArticleQuotaCycle, plan.Code)
+		if plan.ArticleGeneration < 0 {
+			plan.ArticleGeneration = 0
+		}
+		if plan.ImageStorageBytes < 0 {
+			plan.ImageStorageBytes = 0
+		}
+		if plan.CompanyLimit < 0 {
+			plan.CompanyLimit = 0
+		}
+		if plan.SubscriptionDuration < 0 {
+			plan.SubscriptionDuration = 0
+		}
+		normalized = append(normalized, plan)
+		seen[plan.Code] = struct{}{}
+	}
+
+	if len(normalized) == 0 {
+		normalized = defaultMembershipPlans()
+	}
+	cfg.Plans = normalized
+
+	if _, ok := cfg.FindPlan(cfg.DefaultPlanCode); !ok {
+		cfg.DefaultPlanCode = cfg.Plans[0].Code
+	}
+}
+
+func defaultMembershipPlans() []MembershipPlanConfig {
+	return []MembershipPlanConfig{
+		{
+			Code:                 "free",
+			Name:                 "Free",
+			ArticleGeneration:    3,
+			ArticleQuotaCycle:    ArticleQuotaCycleLifetime,
+			ImageStorageBytes:    10 * 1024 * 1024,
+			CompanyLimit:         1,
+			SubscriptionDuration: 72 * time.Hour,
+			ContactAdminOnExpiry: true,
+		},
+		{
+			Code:                 "plus",
+			Name:                 "Plus",
+			ArticleGeneration:    200,
+			ArticleQuotaCycle:    ArticleQuotaCycleMonthly,
+			ImageStorageBytes:    500 * 1024 * 1024,
+			CompanyLimit:         1,
+			SubscriptionDuration: 720 * time.Hour,
+		},
+		{
+			Code:                 "pro",
+			Name:                 "Pro",
+			ArticleGeneration:    400,
+			ArticleQuotaCycle:    ArticleQuotaCycleMonthly,
+			ImageStorageBytes:    1024 * 1024 * 1024,
+			CompanyLimit:         2,
+			SubscriptionDuration: 720 * time.Hour,
+		},
+	}
+}
+
+func normalizeArticleQuotaCycle(value, planCode string) string {
+	switch strings.ToLower(strings.TrimSpace(value)) {
+	case ArticleQuotaCycleLifetime:
+		return ArticleQuotaCycleLifetime
+	case ArticleQuotaCycleMonthly:
+		return ArticleQuotaCycleMonthly
+	default:
+		if strings.EqualFold(strings.TrimSpace(planCode), "free") {
+			return ArticleQuotaCycleLifetime
+		}
+		return ArticleQuotaCycleMonthly
+	}
+}
+
+func defaultPlanName(code string) string {
+	switch strings.ToLower(strings.TrimSpace(code)) {
+	case "free":
+		return "Free"
+	case "plus":
+		return "Plus"
+	case "pro":
+		return "Pro"
+	default:
+		return strings.ToUpper(strings.TrimSpace(code))
+	}
+}
+
 func lookupFirstNonEmptyEnv(keys ...string) (string, bool) {
 	for _, key := range keys {
 		if value, ok := lookupNonEmptyEnv(key); ok {
@@ -495,3 +732,17 @@ func lookupBoolEnv(key string) (bool, bool) {
 		return false, false
 	}
 }
+
+func maxInt(value, fallback int) int {
+	if value < fallback {
+		return fallback
+	}
+	return value
+}
+
+func maxInt64(value, fallback int64) int64 {
+	if value < fallback {
+		return fallback
+	}
+	return value
+}
@@ -4,6 +4,7 @@ import (
 	"os"
 	"path/filepath"
 	"testing"
+	"time"
 )

 func TestLoadPrefersEnvLLMAPIKey(t *testing.T) {
@@ -113,6 +114,55 @@ brand_library: {}
 	}
 }

+func TestLoadAppliesMembershipDefaults(t *testing.T) {
+	configPath := writeTestConfig(t, `
+membership: {}
+`)
+
+	cfg, err := Load(configPath)
+	if err != nil {
+		t.Fatalf("Load() error = %v", err)
+	}
+
+	if cfg.Membership.DefaultPlanCode != "free" {
+		t.Fatalf("expected default plan code free, got %q", cfg.Membership.DefaultPlanCode)
+	}
+	if len(cfg.Membership.Plans) != 3 {
+		t.Fatalf("expected 3 default plans, got %d", len(cfg.Membership.Plans))
+	}
+
+	free, ok := cfg.Membership.FindPlan("free")
+	if !ok {
+		t.Fatalf("expected free plan to exist")
+	}
+	if free.ArticleGeneration != 3 {
+		t.Fatalf("expected free article generation 3, got %d", free.ArticleGeneration)
+	}
+	if free.ImageStorageBytes != 10*1024*1024 {
+		t.Fatalf("expected free image storage 10MB, got %d", free.ImageStorageBytes)
+	}
+	if free.CompanyLimit != 1 {
+		t.Fatalf("expected free company limit 1, got %d", free.CompanyLimit)
+	}
+	if free.SubscriptionDuration != 72*time.Hour {
+		t.Fatalf("expected free subscription duration 72h, got %s", free.SubscriptionDuration)
+	}
+	if !free.ContactAdminOnExpiry {
+		t.Fatalf("expected free plan to require admin contact on expiry")
+	}
+
+	pro, ok := cfg.Membership.FindPlan("pro")
+	if !ok {
+		t.Fatalf("expected pro plan to exist")
+	}
+	if pro.ArticleGeneration != 400 {
+		t.Fatalf("expected pro article generation 400, got %d", pro.ArticleGeneration)
+	}
+	if pro.CompanyLimit != 2 {
+		t.Fatalf("expected pro company limit 2, got %d", pro.CompanyLimit)
+	}
+}
+
 func writeTestConfig(t *testing.T, body string) string {
 	t.Helper()