fix frontend

This commit is contained in:
王性驊 2026-07-15 15:23:59 +00:00
parent 377b7f06ea
commit fb8100825f
1202 changed files with 9706 additions and 108999 deletions

View File

@ -43,6 +43,7 @@ func main() {
}, },
"studio_outbox": { "studio_outbox": {
{Keys: bson.D{{Key: "status", Value: 1}, {Key: "updated_at", Value: 1}}, Options: options.Index().SetName("worker_outbox_status_updated")}, {Keys: bson.D{{Key: "status", Value: 1}, {Key: "updated_at", Value: 1}}, Options: options.Index().SetName("worker_outbox_status_updated")},
{Keys: bson.D{{Key: "steps.status", Value: 1}, {Key: "steps.scheduled_at", Value: 1}, {Key: "steps.lease_expires_at", Value: 1}}, Options: options.Index().SetName("claim_due_outbox_steps")},
{Keys: bson.D{{Key: "owner_uid", Value: 1}, {Key: "updated_at", Value: -1}}, Options: options.Index().SetName("owner_outbox_updated")}, {Keys: bson.D{{Key: "owner_uid", Value: 1}, {Key: "updated_at", Value: -1}}, Options: options.Index().SetName("owner_outbox_updated")},
}, },
} }

View File

@ -3,6 +3,7 @@ package main
import ( import (
"context" "context"
"encoding/json" "encoding/json"
"errors"
"flag" "flag"
"fmt" "fmt"
"os" "os"
@ -13,6 +14,7 @@ import (
"apps/backend/internal/config" "apps/backend/internal/config"
redislock "apps/backend/internal/lib/redislock" redislock "apps/backend/internal/lib/redislock"
"apps/backend/internal/lib/securelog"
"apps/backend/internal/module/ai" "apps/backend/internal/module/ai"
appnotifRepo "apps/backend/internal/module/appnotif/repository" appnotifRepo "apps/backend/internal/module/appnotif/repository"
appnotifUC "apps/backend/internal/module/appnotif/usecase" appnotifUC "apps/backend/internal/module/appnotif/usecase"
@ -40,6 +42,7 @@ import (
"github.com/zeromicro/go-zero/core/conf" "github.com/zeromicro/go-zero/core/conf"
"github.com/zeromicro/go-zero/core/logx" "github.com/zeromicro/go-zero/core/logx"
"github.com/zeromicro/go-zero/core/stores/redis"
) )
// Workerdemo + threads_token_renew + persona_analyze_* + outbox publish // Workerdemo + threads_token_renew + persona_analyze_* + outbox publish
@ -47,6 +50,8 @@ import (
// cd apps/backend && go build -o bin/worker ./cmd/worker // cd apps/backend && go build -o bin/worker ./cmd/worker
// ./bin/worker -f etc/gateway.yaml // ./bin/worker -f etc/gateway.yaml
func main() { func main() {
securelog.SilenceMongo()
configFile := flag.String("f", "etc/gateway.yaml", "config file") configFile := flag.String("f", "etc/gateway.yaml", "config file")
flag.Parse() flag.Parse()
@ -56,6 +61,13 @@ func main() {
if err := c.ValidateProductionDependencies(); err != nil { if err := c.ValidateProductionDependencies(); err != nil {
logx.Must(err) logx.Must(err)
} }
redisClient := redis.MustNewRedis(c.CacheRedis[0].RedisConf)
pingCtx, cancelPing := context.WithTimeout(context.Background(), 3*time.Second)
redisReady := redisClient.PingCtx(pingCtx)
cancelPing()
if !redisReady {
logx.Must(fmt.Errorf("redis startup ping failed"))
}
workerID := c.Worker.ID workerID := c.Worker.ID
if workerID == "" || workerID == "worker-1" || workerID == "worker-local-1" || workerID == "demo-worker-1" { if workerID == "" || workerID == "worker-1" || workerID == "worker-local-1" || workerID == "demo-worker-1" {
@ -86,7 +98,11 @@ func main() {
threadsSvc.Renew = jobs threadsSvc.Renew = jobs
// usage + AI keys與 gateway 對齊persona LLM 需真 key // usage + AI keys與 gateway 對齊persona LLM 需真 key
members := memberRepo.NewMoncStore(c.Mongo.URI, c.Mongo.Database, c.CacheRedis) settingsCodec, err := c.NewMemberSettingsCodec()
if err != nil {
logx.Must(err)
}
members := memberRepo.NewMoncStore(c.Mongo.URI, c.Mongo.Database, c.CacheRedis, c.CacheRedis[0].RedisConf, c.Redis.Namespace, settingsCodec)
keyRes := &usageUC.SettingsResolver{ keyRes := &usageUC.SettingsResolver{
Members: members, Members: members,
PlatformAI: c.Platform.AIKey, PlatformAI: c.Platform.AIKey,
@ -95,7 +111,9 @@ func main() {
PlatformExa: c.Platform.ExaKey, PlatformExa: c.Platform.ExaKey,
} }
usageSvc := usageUC.New(usageRepo.NewMonStore(c.Mongo.URI, c.Mongo.Database), keyRes) usageSvc := usageUC.New(usageRepo.NewMonStore(c.Mongo.URI, c.Mongo.Database), keyRes)
aiRegistry := ai.NewRegistry() usageSvc.Gate = usageUC.NewRedisPlatformGate(c.CacheRedis[0].RedisConf, 120, 60, c.Redis.Namespace)
modelsCache := ai.NewRedisModelsCache(c.CacheRedis[0].RedisConf, c.Redis.Namespace)
aiRegistry := ai.NewRegistry(modelsCache, c.Redis.AIModelCacheFingerprintSecret)
var pub studioPublish.Transport var pub studioPublish.Transport
if strings.TrimSpace(c.Platform.ThreadsAppId) != "" && strings.TrimSpace(c.Platform.ThreadsAppSecret) != "" { if strings.TrimSpace(c.Platform.ThreadsAppId) != "" && strings.TrimSpace(c.Platform.ThreadsAppSecret) != "" {
@ -105,6 +123,7 @@ func main() {
pub = studioPublish.NewUnavailable("set THREADS_APP_ID and THREADS_APP_SECRET") pub = studioPublish.NewUnavailable("set THREADS_APP_ID and THREADS_APP_SECRET")
} }
studio := studioUC.New(studioRepo.NewMonStore(c.Mongo.URI, c.Mongo.Database), pub) studio := studioUC.New(studioRepo.NewMonStore(c.Mongo.URI, c.Mongo.Database), pub)
studio.OutboxWorkerID = workerID
studio.Accounts = threadsSvc studio.Accounts = threadsSvc
studio.Usage = usageSvc studio.Usage = usageSvc
studio.AI = &ai.FakeClient{} studio.AI = &ai.FakeClient{}
@ -133,7 +152,9 @@ func main() {
defer tick.Stop() defer tick.Stop()
ctx := context.Background() ctx := context.Background()
outboxLock := redislock.New(c.CacheRedis[0].RedisConf, "haixun:worker:outbox", workerID, 60*time.Second) const outboxLockTTL = 3 * time.Minute
outboxLockKey := strings.Trim(c.Redis.Namespace, ":") + ":worker:outbox"
outboxLock := redislock.NewWithClient(redisClient, outboxLockKey, workerID, outboxLockTTL)
for { for {
select { select {
case <-sig: case <-sig:
@ -194,22 +215,8 @@ func main() {
_, _ = jobs.FailJob(ctx, j.ID, "unknown template: "+j.TemplateType) _, _ = jobs.FailJob(ctx, j.ID, "unknown template: "+j.TemplateType)
} }
} }
// 2) One worker owns an outbox tick. The Redis value is workerID and // 2) One worker owns an outbox tick and renews its lease while claims run.
// release is compare-and-delete, so a departing worker cannot unlock a processOutbox(ctx, studio, outboxLock, workerID, outboxLockTTL)
// successor's lease.
if locked, lerr := outboxLock.Acquire(ctx); lerr != nil {
logx.Errorf("worker %s outbox lock: %v", workerID, lerr)
} else if locked {
n, err := studio.ProcessDueSteps(ctx, 0)
if rerr := outboxLock.Release(ctx); rerr != nil {
logx.Errorf("worker %s outbox unlock: %v", workerID, rerr)
}
if err != nil {
logx.Errorf("worker %s outbox tick: %v", workerID, err)
} else if n > 0 {
logx.Infof("worker %s outbox published %d step(s)", workerID, n)
}
}
// 3) 終態任務超過 2 天自動清除 // 3) 終態任務超過 2 天自動清除
if purged, err := jobs.PurgeExpiredTerminal(ctx); err != nil { if purged, err := jobs.PurgeExpiredTerminal(ctx); err != nil {
logx.Errorf("worker %s purge jobs: %v", workerID, err) logx.Errorf("worker %s purge jobs: %v", workerID, err)
@ -220,6 +227,66 @@ func main() {
} }
} }
func processOutbox(ctx context.Context, studio *studioUC.Service, lock *redislock.Lock, workerID string, ttl time.Duration) {
locked, err := lock.Acquire(ctx)
if err != nil {
logx.Errorf("worker %s outbox lock acquire: %v", workerID, err)
return
}
if !locked {
return
}
runCtx, cancel := context.WithCancel(ctx)
renewFailure := make(chan error, 1)
renewStopped := make(chan struct{})
go func() {
defer close(renewStopped)
ticker := time.NewTicker(ttl / 3)
defer ticker.Stop()
for {
select {
case <-runCtx.Done():
return
case <-ticker.C:
renewed, renewErr := lock.Renew(runCtx)
if renewErr != nil {
renewFailure <- fmt.Errorf("renew: %w", renewErr)
cancel()
return
}
if !renewed {
renewFailure <- errors.New("lease ownership lost")
cancel()
return
}
}
}
}()
defer func() {
cancel()
<-renewStopped
releaseCtx, releaseCancel := context.WithTimeout(context.Background(), 5*time.Second)
defer releaseCancel()
if err := lock.Release(releaseCtx); err != nil {
logx.Errorf("worker %s outbox lock release: %v", workerID, err)
}
}()
n, processErr := studio.ProcessDueSteps(runCtx, 0)
select {
case renewErr := <-renewFailure:
logx.Errorf("worker %s outbox lock renewal failed; claims canceled: %v", workerID, renewErr)
return
default:
}
if processErr != nil {
logx.Errorf("worker %s outbox tick: %v", workerID, processErr)
} else if n > 0 {
logx.Infof("worker %s outbox published %d step(s)", workerID, n)
}
}
func runScoutScan(ctx context.Context, jobs *jobUC.Service, scout *scoutUC.Service, j *jobDomain.Job) error { func runScoutScan(ctx context.Context, jobs *jobUC.Service, scout *scoutUC.Service, j *jobDomain.Job) error {
var brief scoutDomain.RunBrief var brief scoutDomain.RunBrief
if err := json.Unmarshal([]byte(j.Payload), &brief); err != nil { if err := json.Unmarshal([]byte(j.Payload), &brief); err != nil {
@ -349,7 +416,10 @@ type workerAIKeys struct {
func (k *workerAIKeys) ResolveAI(ctx context.Context, ownerUID int64) (provider, model, apiKey string, err error) { func (k *workerAIKeys) ResolveAI(ctx context.Context, ownerUID int64) (provider, model, apiKey string, err error) {
st, err := k.Members.GetSettings(ctx, ownerUID) st, err := k.Members.GetSettings(ctx, ownerUID)
if err != nil || st == nil { if err != nil {
return "", "", "", err
}
if st == nil {
st = memberDomain.DefaultSettings(ownerUID) st = memberDomain.DefaultSettings(ownerUID)
} }
// 完全依會員 AI 設定provider / model / key不寫死模型 // 完全依會員 AI 設定provider / model / key不寫死模型
@ -360,6 +430,9 @@ func (k *workerAIKeys) ResolveAI(ctx context.Context, ownerUID int64) (provider,
} }
if k.Resolver != nil { if k.Resolver != nil {
_, key, rerr := k.Resolver.ResolveKey(ctx, ownerUID, usageDomain.MeterAICopy) _, key, rerr := k.Resolver.ResolveKey(ctx, ownerUID, usageDomain.MeterAICopy)
if rerr != nil && !errors.Is(rerr, usageDomain.ErrNoKey) {
return provider, model, "", rerr
}
if rerr == nil && strings.TrimSpace(key) != "" { if rerr == nil && strings.TrimSpace(key) != "" {
return provider, model, key, nil return provider, model, key, nil
} }
@ -379,7 +452,10 @@ func (d *workerDevMode) DevModeEnabled(ctx context.Context, uid int64) (bool, er
return false, nil return false, nil
} }
st, err := d.Members.GetSettings(ctx, uid) st, err := d.Members.GetSettings(ctx, uid)
if err != nil || st == nil { if err != nil {
return false, err
}
if st == nil {
return false, nil return false, nil
} }
return st.DevModeEnabled, nil return st.DevModeEnabled, nil

View File

@ -19,6 +19,13 @@ CacheRedis:
Type: node Type: node
Pass: "" Pass: ""
# Shared key prefix for this environment and a dedicated HMAC secret for
# credential-safe AI model-cache fingerprints. Prefer env REDIS_NAMESPACE and
# AI_MODEL_CACHE_FINGERPRINT_SECRET; do not reuse JWT, Scout, or encryption keys.
Redis:
Namespace: "haixun:dev:v1"
AIModelCacheFingerprintSecret: "REPLACE_WITH_A_DEDICATED_LONG_RANDOM_SECRET"
# Mongo — mon/monc 使用完整 URI含帳密時寫在 URI可用環境變數 MONGO_URI 覆寫) # Mongo — mon/monc 使用完整 URI含帳密時寫在 URI可用環境變數 MONGO_URI 覆寫)
# 本機有 auth 的例子: # 本機有 auth 的例子:
# export MONGO_URI='mongodb://USER:PASS@127.0.0.1:27017/?authSource=admin' # export MONGO_URI='mongodb://USER:PASS@127.0.0.1:27017/?authSource=admin'
@ -34,6 +41,13 @@ Auth:
RefreshSecret: "REPLACE_WITH_A_DIFFERENT_LONG_RANDOM_REFRESH_SECRET" RefreshSecret: "REPLACE_WITH_A_DIFFERENT_LONG_RANDOM_REFRESH_SECRET"
RefreshExpire: 604800 RefreshExpire: 604800
# AES-256-GCM for member BYOK at rest. Generate Key with: openssl rand -base64 32
# Prefer env MEMBER_SETTINGS_ENCRYPTION_KEY_ID / MEMBER_SETTINGS_ENCRYPTION_KEY.
# Keep this key ID/key unchanged until persisted values have been migrated to a replacement.
MemberSettingsEncryption:
CurrentKeyID: "v1"
Key: "REPLACE_WITH_BASE64_32_BYTE_KEY"
# 平台預設 key所有會員可用會員在設定頁填的 BYOK 優先) # 平台預設 key所有會員可用會員在設定頁填的 BYOK 優先)
# 也可用環境變數覆寫PLATFORM_XAI_KEY / PLATFORM_OPENCODE_KEY / PLATFORM_EXA_KEY / PLATFORM_AI_KEY # 也可用環境變數覆寫PLATFORM_XAI_KEY / PLATFORM_OPENCODE_KEY / PLATFORM_EXA_KEY / PLATFORM_AI_KEY
Platform: Platform:
@ -74,6 +88,18 @@ OAuth:
LineClientSecret: "" LineClientSecret: ""
LineRedirectURI: "" LineRedirectURI: ""
# Prefer STRIPE_* environment variables. Disabled/missing configuration never
# blocks startup; Stripe-dependent billing operations explicitly return 503.
Stripe:
Enabled: false
SecretKey: ""
WebhookSecret: ""
StarterPriceID: ""
ProPriceID: ""
SuccessURL: ""
CancelURL: ""
PortalReturnURL: ""
# 驗證碼/重設密碼寄信 # 驗證碼/重設密碼寄信
# #
# 本機MailpitMailgun 相容 SMTP 介面): # 本機MailpitMailgun 相容 SMTP 介面):

View File

@ -1,9 +1,11 @@
package main package main
import ( import (
"context"
"flag" "flag"
"fmt" "fmt"
"os" "os"
"time"
"apps/backend/internal/config" "apps/backend/internal/config"
"apps/backend/internal/handler" "apps/backend/internal/handler"
@ -41,6 +43,12 @@ func main() {
server.Use(middleware.NewCorsMiddleware(c.CORSAllowedOrigins).Handle) server.Use(middleware.NewCorsMiddleware(c.CORSAllowedOrigins).Handle)
ctx := svc.NewServiceContext(c) ctx := svc.NewServiceContext(c)
pingCtx, cancelPing := context.WithTimeout(context.Background(), 3*time.Second)
redisReady := ctx.Redis != nil && ctx.Redis.PingCtx(pingCtx)
cancelPing()
if !redisReady {
logx.Must(fmt.Errorf("redis startup ping failed"))
}
handler.RegisterHandlers(server, ctx) handler.RegisterHandlers(server, ctx)
// 靈感 SSE streamgoctl 不生成 streaming handler // 靈感 SSE streamgoctl 不生成 streaming handler
handler.RegisterInspireStream(server, ctx) handler.RegisterInspireStream(server, ctx)

View File

@ -0,0 +1,69 @@
syntax = "v1"
type (
BillingCheckoutCreateReq {
PlanId string `json:"plan_id"`
RequestId string `json:"request_id"`
}
BillingCheckoutData {
Id string `json:"id"`
SessionId string `json:"session_id"`
Url string `json:"url"`
ExpiresAt int64 `json:"expires_at"`
Status string `json:"status,optional"`
PaymentStatus string `json:"payment_status,optional"`
FulfillmentStatus string `json:"fulfillment_status,optional"`
PlanId string `json:"plan_id,optional"`
}
BillingCheckoutGetReq {
Id string `path:"id"`
}
BillingSubscriptionData {
PlanId string `json:"plan_id"`
PaidPlanId string `json:"paid_plan_id,optional"`
Status string `json:"status"`
CustomerId string `json:"customer_id,optional"`
SubscriptionId string `json:"subscription_id,optional"`
CurrentPeriodStart int64 `json:"current_period_start,optional"`
CurrentPeriodEnd int64 `json:"current_period_end,optional"`
CancelAtPeriodEnd bool `json:"cancel_at_period_end"`
AdminOverride bool `json:"admin_override"`
}
BillingPortalData {
Url string `json:"url"`
}
)
@server (
jwt: Auth
group: billing
prefix: /api/v1/billing
middleware: AuthJWT
)
service gateway {
@handler CreateBillingCheckoutSession
post /checkout-sessions (BillingCheckoutCreateReq) returns (BillingCheckoutData)
@handler GetBillingCheckoutSession
get /checkout-sessions/:id (BillingCheckoutGetReq) returns (BillingCheckoutData)
@handler GetBillingSubscription
get /subscription returns (BillingSubscriptionData)
@handler CreateBillingPortalSession
post /portal-sessions returns (BillingPortalData)
}
@server (
group: billing
prefix: /api/v1/billing
middleware: StripeRawBody
)
service gateway {
@handler HandleStripeWebhook
post /stripe/webhook
}

View File

@ -17,6 +17,7 @@ import "threads.api"
import "jobs.api" import "jobs.api"
import "notifications.api" import "notifications.api"
import "usage.api" import "usage.api"
import "billing.api"
import "proxy.api" import "proxy.api"
import "extension.api" import "extension.api"
import "studio.api" import "studio.api"

View File

@ -101,6 +101,54 @@ type (
UsageTenantSummaryReq { UsageTenantSummaryReq {
MonthKey string `form:"month_key,optional"` MonthKey string `form:"month_key,optional"`
} }
UsageTenantAnalyticsReq {
Granularity string `form:"granularity"` // day|month|year
From string `form:"from"` // inclusive YYYY-MM-DD UTC
To string `form:"to"` // inclusive YYYY-MM-DD UTC
}
UsageTenantAnalyticsBucket {
Key string `json:"key"`
Label string `json:"label"`
PurchasedCredits int `json:"purchased_credits"`
ConsumedCredits int `json:"consumed_credits"`
AiCalls int `json:"ai_calls"`
SearchCalls int `json:"search_calls"`
ByokCallCount int `json:"byok_call_count"`
}
UsageTenantAnalyticsMember {
Uid int64 `json:"uid"`
Email string `json:"email"`
DisplayName string `json:"display_name"`
Unlimited bool `json:"unlimited"`
PlanId string `json:"plan_id"`
PurchasedCredits int `json:"purchased_credits"`
TotalCredits int `json:"total_credits"`
AiCalls int `json:"ai_calls"`
SearchCalls int `json:"search_calls"`
ByokCallCount int `json:"byok_call_count"`
RemainingCredits int `json:"remaining_credits"`
Pct int `json:"pct"`
}
UsageTenantAnalyticsData {
Granularity string `json:"granularity"`
From string `json:"from"`
To string `json:"to"`
RangeLabel string `json:"range_label"`
PurchasedCredits int `json:"purchased_credits"`
ConsumedCredits int `json:"consumed_credits"`
RemainingCredits int `json:"remaining_credits"`
Pct int `json:"pct"`
AiCalls int `json:"ai_calls"`
SearchCalls int `json:"search_calls"`
ByMeter map[string]UsageMeterCount `json:"by_meter"`
Byok UsageByokBlock `json:"byok"`
Series []UsageTenantAnalyticsBucket `json:"series"`
Members []UsageTenantAnalyticsMember `json:"members"`
}
) )
@server ( @server (
@ -138,4 +186,7 @@ service gateway {
@handler GetTenantUsageSummary @handler GetTenantUsageSummary
get /tenant-summary (UsageTenantSummaryReq) returns (UsageTenantSummaryData) get /tenant-summary (UsageTenantSummaryReq) returns (UsageTenantSummaryData)
@handler GetTenantUsageAnalytics
get /tenant-analytics (UsageTenantAnalyticsReq) returns (UsageTenantAnalyticsData)
} }

View File

@ -49,16 +49,16 @@ make migrate-down
| 類型 | 放哪 | 範例 | | 類型 | 放哪 | 範例 |
|------|------|------| |------|------|------|
| Index / collection 約束 | `*.up.json` createIndexes | members.uid unique | | Index / collection 約束 | `*.up.json` createIndexes | members.uid unique |
| 種子資料 | 另開 version如 000003_seed_* | admin 帳 | | 種子資料 | 不放 migration | 使用 `cmd/seeder` 與 runtime secret |
| 業務邏輯、密碼 runtime 產生 | **不要**寫進 migration | 用 API / 測試 helper | | 業務邏輯、密碼 runtime 產生 | **不要**寫進 migration | 用 API / 測試 helper |
種子 admin 預設:`admin@haixun.local` / `admin123`hash 寫死在 up 檔;正式環境請改或不要跑 seed migration Migration 不建立任何預設帳號。首次管理員必須透過 `make seeder`,並由環境提供強密碼
## 與 gateway ## 與 gateway
`make migrate-up`,再 `go run . -f etc/gateway.yaml`(需 Mongo + Redis `make migrate-up`,再 `go run . -f etc/gateway.yaml`(需 Mongo + Redis
## 不要 ## 不要
- 手寫 Go seederinit - 在 migration 寫死帳號或密碼
- 手寫 migration runner - 手寫 migration runner
- 把 stand-alone 電商 `.txt` 當 migrate 來源(格式不同;要用 JSON runCommand - 把 stand-alone 電商 `.txt` 當 migrate 來源(格式不同;要用 JSON runCommand

View File

@ -1,11 +1 @@
[ []
{
"delete": "members",
"deletes": [
{
"q": { "uid": 1, "email": "admin@30cm.net" },
"limit": 1
}
]
}
]

View File

@ -1,28 +1,4 @@
[ [
{
"insert": "members",
"documents": [
{
"uid": 1000000,
"tenant_id": "default",
"email": "admin@30cm.net",
"display_name": "Harbor Admin",
"roles": ["admin", "member"],
"status": "active",
"email_verified": true,
"email_verified_at": 1700000000000000000,
"invite_code": "SEEDADMIN",
"password_hash": "$2a$10$iGGFNC5rM8MFhU8Wj6a7TOcQCsTeD03bkEODvQBh2ZxvCH.PcLl3O",
"timezone": "Asia/Taipei",
"notify_email": false,
"bio": "",
"created_at": 1700000000000000000,
"updated_at": 1700000000000000000,
"seed_note": "password is admin123 — change in production; uid=1000000 (8 digits from 1M)"
}
],
"ordered": false
},
{ {
"update": "counters", "update": "counters",
"updates": [ "updates": [

View File

@ -1,11 +1 @@
[ []
{
"delete": "identities",
"deletes": [
{
"q": { "uid": 1, "login_id": "admin@haixun.local", "platform": "platform" },
"limit": 1
}
]
}
]

View File

@ -1,17 +1 @@
[ []
{
"insert": "identities",
"documents": [
{
"uid": 1000000,
"login_id": "admin@haixun.local",
"platform": "platform",
"account_type": "email",
"password_hash": "$2a$10$iGGFNC5rM8MFhU8Wj6a7TOcQCsTeD03bkEODvQBh2ZxvCH.PcLl3O",
"created_at": 1700000000000000000,
"updated_at": 1700000000000000000
}
],
"ordered": false
}
]

View File

@ -0,0 +1,5 @@
[
{ "dropIndexes": "usage_monthly_counters", "index": "usage_counter_uid_month" },
{ "dropIndexes": "usage_events", "index": "usage_events_member_month_mode" },
{ "dropIndexes": "usage_events", "index": "usage_events_month" }
]

View File

@ -0,0 +1,25 @@
[
{
"createIndexes": "usage_monthly_counters",
"indexes": [
{
"key": { "uid": 1, "month_key": 1 },
"name": "usage_counter_uid_month",
"unique": true
}
]
},
{
"createIndexes": "usage_events",
"indexes": [
{
"key": { "uid": 1, "month_key": 1, "key_mode": 1, "created_at": -1 },
"name": "usage_events_member_month_mode"
},
{
"key": { "month_key": 1, "created_at": -1 },
"name": "usage_events_month"
}
]
}
]

View File

@ -0,0 +1,3 @@
[
{ "dropIndexes": "studio_outbox", "index": "claim_due_outbox_steps" }
]

View File

@ -0,0 +1,11 @@
[
{
"createIndexes": "studio_outbox",
"indexes": [
{
"key": { "steps.status": 1, "steps.scheduled_at": 1, "steps.lease_expires_at": 1 },
"name": "claim_due_outbox_steps"
}
]
}
]

View File

@ -0,0 +1,11 @@
[
{ "dropIndexes": "usage_prefs", "index": "unique_usage_prefs_uid" },
{ "dropIndexes": "usage_prefs", "index": "unique_stripe_customer" },
{ "dropIndexes": "usage_prefs", "index": "unique_stripe_subscription" },
{ "dropIndexes": "billing_checkout_attempts", "index": "unique_checkout_request" },
{ "dropIndexes": "billing_checkout_attempts", "index": "unique_checkout_session" },
{ "dropIndexes": "billing_webhook_events", "index": "stripe_webhook_event_id" },
{ "dropIndexes": "usage_purchases", "index": "unique_purchase_stripe_event" },
{ "dropIndexes": "usage_purchases", "index": "purchase_stripe_session" },
{ "dropIndexes": "usage_purchases", "index": "purchase_stripe_subscription" }
]

View File

@ -0,0 +1,31 @@
[
{
"createIndexes": "usage_prefs",
"indexes": [
{ "key": { "uid": 1 }, "name": "unique_usage_prefs_uid", "unique": true },
{ "key": { "stripe_customer_id": 1 }, "name": "unique_stripe_customer", "unique": true, "sparse": true },
{ "key": { "stripe_subscription_id": 1 }, "name": "unique_stripe_subscription", "unique": true, "sparse": true }
]
},
{
"createIndexes": "billing_checkout_attempts",
"indexes": [
{ "key": { "uid": 1, "request_id": 1 }, "name": "unique_checkout_request", "unique": true },
{ "key": { "session_id": 1 }, "name": "unique_checkout_session", "unique": true, "sparse": true }
]
},
{
"createIndexes": "billing_webhook_events",
"indexes": [
{ "key": { "stripe_event_id": 1 }, "name": "stripe_webhook_event_id", "unique": true }
]
},
{
"createIndexes": "usage_purchases",
"indexes": [
{ "key": { "stripe_event_id": 1 }, "name": "unique_purchase_stripe_event", "unique": true, "sparse": true },
{ "key": { "stripe_session_id": 1 }, "name": "purchase_stripe_session", "sparse": true },
{ "key": { "stripe_subscription_id": 1 }, "name": "purchase_stripe_subscription", "sparse": true }
]
}
]

View File

@ -3,8 +3,14 @@ module apps/backend
go 1.25.0 go 1.25.0
require ( require (
github.com/aws/aws-sdk-go-v2 v1.42.1
github.com/aws/aws-sdk-go-v2/credentials v1.19.28
github.com/aws/aws-sdk-go-v2/service/s3 v1.105.0
github.com/golang-jwt/jwt/v5 v5.2.1 github.com/golang-jwt/jwt/v5 v5.2.1
github.com/google/uuid v1.6.0
github.com/matcornic/hermes/v2 v2.1.0
github.com/stretchr/testify v1.10.0 github.com/stretchr/testify v1.10.0
github.com/stripe/stripe-go/v82 v82.5.1
github.com/zeromicro/go-zero v1.7.6 github.com/zeromicro/go-zero v1.7.6
go.mongodb.org/mongo-driver v1.17.1 go.mongodb.org/mongo-driver v1.17.1
golang.org/x/crypto v0.54.0 golang.org/x/crypto v0.54.0
@ -16,9 +22,7 @@ require (
github.com/PuerkitoBio/goquery v1.5.0 // indirect github.com/PuerkitoBio/goquery v1.5.0 // indirect
github.com/andybalholm/cascadia v1.0.0 // indirect github.com/andybalholm/cascadia v1.0.0 // indirect
github.com/aokoli/goutils v1.0.1 // indirect github.com/aokoli/goutils v1.0.1 // indirect
github.com/aws/aws-sdk-go-v2 v1.42.1 // indirect
github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.7.14 // indirect github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.7.14 // indirect
github.com/aws/aws-sdk-go-v2/credentials v1.19.28 // indirect
github.com/aws/aws-sdk-go-v2/internal/configsources v1.4.30 // indirect github.com/aws/aws-sdk-go-v2/internal/configsources v1.4.30 // indirect
github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.7.30 // indirect github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.7.30 // indirect
github.com/aws/aws-sdk-go-v2/internal/v4a v1.4.31 // indirect github.com/aws/aws-sdk-go-v2/internal/v4a v1.4.31 // indirect
@ -26,7 +30,6 @@ require (
github.com/aws/aws-sdk-go-v2/service/internal/checksum v1.9.23 // indirect github.com/aws/aws-sdk-go-v2/service/internal/checksum v1.9.23 // indirect
github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.13.30 // indirect github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.13.30 // indirect
github.com/aws/aws-sdk-go-v2/service/internal/s3shared v1.19.31 // indirect github.com/aws/aws-sdk-go-v2/service/internal/s3shared v1.19.31 // indirect
github.com/aws/aws-sdk-go-v2/service/s3 v1.105.0 // indirect
github.com/aws/smithy-go v1.27.3 // indirect github.com/aws/smithy-go v1.27.3 // indirect
github.com/beorn7/perks v1.0.1 // indirect github.com/beorn7/perks v1.0.1 // indirect
github.com/cenkalti/backoff/v4 v4.3.0 // indirect github.com/cenkalti/backoff/v4 v4.3.0 // indirect
@ -38,14 +41,12 @@ require (
github.com/go-logr/stdr v1.2.2 // indirect github.com/go-logr/stdr v1.2.2 // indirect
github.com/golang-jwt/jwt/v4 v4.5.1 // indirect github.com/golang-jwt/jwt/v4 v4.5.1 // indirect
github.com/golang/snappy v0.0.4 // indirect github.com/golang/snappy v0.0.4 // indirect
github.com/google/uuid v1.6.0 // indirect
github.com/gorilla/css v1.0.0 // indirect github.com/gorilla/css v1.0.0 // indirect
github.com/grpc-ecosystem/grpc-gateway/v2 v2.20.0 // indirect github.com/grpc-ecosystem/grpc-gateway/v2 v2.20.0 // indirect
github.com/huandu/xstrings v1.2.0 // indirect github.com/huandu/xstrings v1.2.0 // indirect
github.com/imdario/mergo v0.3.6 // indirect github.com/imdario/mergo v0.3.6 // indirect
github.com/jaytaylor/html2text v0.0.0-20180606194806-57d518f124b0 // indirect github.com/jaytaylor/html2text v0.0.0-20180606194806-57d518f124b0 // indirect
github.com/klauspost/compress v1.17.9 // indirect github.com/klauspost/compress v1.17.9 // indirect
github.com/matcornic/hermes/v2 v2.1.0 // indirect
github.com/mattn/go-colorable v0.1.13 // indirect github.com/mattn/go-colorable v0.1.13 // indirect
github.com/mattn/go-isatty v0.0.20 // indirect github.com/mattn/go-isatty v0.0.20 // indirect
github.com/mattn/go-runewidth v0.0.15 // indirect github.com/mattn/go-runewidth v0.0.15 // indirect

View File

@ -153,6 +153,8 @@ github.com/stretchr/testify v1.8.4/go.mod h1:sz/lmYIOXD/1dqDmKjjqLyZ2RngseejIcXl
github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY= github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
github.com/stretchr/testify v1.10.0 h1:Xv5erBjTwe/5IxqUQTdXv5kgmIvbHo3QQyRwhJsOfJA= github.com/stretchr/testify v1.10.0 h1:Xv5erBjTwe/5IxqUQTdXv5kgmIvbHo3QQyRwhJsOfJA=
github.com/stretchr/testify v1.10.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY= github.com/stretchr/testify v1.10.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
github.com/stripe/stripe-go/v82 v82.5.1 h1:05q6ZDKoe8PLMpQV072obF74HCgP4XJeJYoNuRSX2+8=
github.com/stripe/stripe-go/v82 v82.5.1/go.mod h1:majCQX6AfObAvJiHraPi/5udwHi4ojRvJnnxckvHrX8=
github.com/vanng822/css v0.0.0-20190504095207-a21e860bcd04 h1:L0rPdfzq43+NV8rfIx2kA4iSSLRj2jN5ijYHoeXRwvQ= github.com/vanng822/css v0.0.0-20190504095207-a21e860bcd04 h1:L0rPdfzq43+NV8rfIx2kA4iSSLRj2jN5ijYHoeXRwvQ=
github.com/vanng822/css v0.0.0-20190504095207-a21e860bcd04/go.mod h1:tcnB1voG49QhCrwq1W0w5hhGasvOg+VQp9i9H1rCM1w= github.com/vanng822/css v0.0.0-20190504095207-a21e860bcd04/go.mod h1:tcnB1voG49QhCrwq1W0w5hhGasvOg+VQp9i9H1rCM1w=
github.com/vanng822/go-premailer v0.0.0-20191214114701-be27abe028fe h1:9YnI5plmy+ad6BM+JCLJb2ZV7/TNiE5l7SNKfumYKgc= github.com/vanng822/go-premailer v0.0.0-20191214114701-be27abe028fe h1:9YnI5plmy+ad6BM+JCLJb2ZV7/TNiE5l7SNKfumYKgc=

View File

@ -2,9 +2,12 @@ package config
import ( import (
"fmt" "fmt"
"net/url"
"os" "os"
"strings" "strings"
"apps/backend/internal/lib/secretbox"
"github.com/zeromicro/go-zero/core/stores/cache" "github.com/zeromicro/go-zero/core/stores/cache"
"github.com/zeromicro/go-zero/rest" "github.com/zeromicro/go-zero/rest"
) )
@ -15,6 +18,10 @@ type Config struct {
// CacheRedis is go-zero cache.CacheConf (feeds monc Redis 快取). // CacheRedis is go-zero cache.CacheConf (feeds monc Redis 快取).
CacheRedis cache.CacheConf CacheRedis cache.CacheConf
Redis struct {
Namespace string `json:",optional,default=haixun:dev:v1"`
AIModelCacheFingerprintSecret string `json:",optional"`
}
Mongo struct { Mongo struct {
URI string URI string
@ -27,6 +34,12 @@ type Config struct {
RefreshSecret string RefreshSecret string
RefreshExpire int64 RefreshExpire int64
} }
// MemberSettingsEncryption is dedicated to member BYOK secrets. Never reuse
// Auth or Scout secrets here.
MemberSettingsEncryption struct {
CurrentKeyID string `json:",optional"`
Key string `json:",optional"`
}
// CORSAllowedOrigins is an explicit browser-origin allowlist. // CORSAllowedOrigins is an explicit browser-origin allowlist.
CORSAllowedOrigins []string `json:",optional"` CORSAllowedOrigins []string `json:",optional"`
Platform struct { Platform struct {
@ -65,6 +78,16 @@ type Config struct {
LineClientSecret string `json:",optional"` LineClientSecret string `json:",optional"`
LineRedirectURI string `json:",optional"` LineRedirectURI string `json:",optional"`
} }
Stripe struct {
Enabled bool `json:",optional"`
SecretKey string `json:",optional"`
WebhookSecret string `json:",optional"`
StarterPriceID string `json:",optional"`
ProPriceID string `json:",optional"`
SuccessURL string `json:",optional"`
CancelURL string `json:",optional"`
PortalReturnURL string `json:",optional"`
}
// Mail — 驗證碼重設密碼寄信SMTP空 Host 則只打 log + 可回 mock_code // Mail — 驗證碼重設密碼寄信SMTP空 Host 則只打 log + 可回 mock_code
Mail struct { Mail struct {
Sender string `json:",optional"` // e.g. "Harbor Desk <noreply@example.com>" Sender string `json:",optional"` // e.g. "Harbor Desk <noreply@example.com>"
@ -86,9 +109,9 @@ type Config struct {
PublicAPIBase string `json:",optional"` PublicAPIBase string `json:",optional"`
// Brand — 郵件 chromehermesLogoURL 空則用 PublicWebBase + /brand-mark.jpg // Brand — 郵件 chromehermesLogoURL 空則用 PublicWebBase + /brand-mark.jpg
Brand struct { Brand struct {
Name string `json:",optional,default=Harbor Desk"` Name string `json:",optional"`
LogoURL string `json:",optional"` // absolute URL LogoURL string `json:",optional"` // absolute URL
Copyright string `json:",optional,default=© Harbor Desk"` Copyright string `json:",optional"`
} `json:",optional"` } `json:",optional"`
// ObjectStorage — 頭像等二進位;**僅 MinIO / S3**(不落本機磁碟) // ObjectStorage — 頭像等二進位;**僅 MinIO / S3**(不落本機磁碟)
ObjectStorage struct { ObjectStorage struct {
@ -104,6 +127,15 @@ type Config struct {
// ApplyEnv overlays secrets from environment. // ApplyEnv overlays secrets from environment.
func (c *Config) ApplyEnv() { func (c *Config) ApplyEnv() {
if strings.TrimSpace(c.Redis.Namespace) == "" {
c.Redis.Namespace = "haixun:dev:v1"
}
if c.Brand.Name == "" {
c.Brand.Name = "Harbor Desk"
}
if c.Brand.Copyright == "" {
c.Brand.Copyright = "© Harbor Desk"
}
if v := os.Getenv("MONGO_URI"); v != "" { if v := os.Getenv("MONGO_URI"); v != "" {
c.Mongo.URI = v c.Mongo.URI = v
} }
@ -124,12 +156,48 @@ func (c *Config) ApplyEnv() {
c.CacheRedis[0].RedisConf.Host = v c.CacheRedis[0].RedisConf.Host = v
} }
} }
if v := firstEnv("REDIS_NAMESPACE", "HAIXUN_REDIS_NAMESPACE"); v != "" {
c.Redis.Namespace = strings.Trim(v, ":")
}
if v := firstEnv("AI_MODEL_CACHE_FINGERPRINT_SECRET", "HAIXUN_AI_MODEL_CACHE_FINGERPRINT_SECRET"); v != "" {
c.Redis.AIModelCacheFingerprintSecret = v
}
if v := os.Getenv("AUTH_ACCESS_SECRET"); v != "" { if v := os.Getenv("AUTH_ACCESS_SECRET"); v != "" {
c.Auth.AccessSecret = v c.Auth.AccessSecret = v
} }
if v := os.Getenv("AUTH_REFRESH_SECRET"); v != "" { if v := os.Getenv("AUTH_REFRESH_SECRET"); v != "" {
c.Auth.RefreshSecret = v c.Auth.RefreshSecret = v
} }
if v := os.Getenv("STRIPE_ENABLED"); v != "" {
c.Stripe.Enabled = v == "1" || strings.EqualFold(v, "true")
}
if v := os.Getenv("STRIPE_SECRET_KEY"); v != "" {
c.Stripe.SecretKey = v
}
if v := os.Getenv("STRIPE_WEBHOOK_SECRET"); v != "" {
c.Stripe.WebhookSecret = v
}
if v := os.Getenv("STRIPE_STARTER_PRICE_ID"); v != "" {
c.Stripe.StarterPriceID = v
}
if v := os.Getenv("STRIPE_PRO_PRICE_ID"); v != "" {
c.Stripe.ProPriceID = v
}
if v := os.Getenv("STRIPE_SUCCESS_URL"); v != "" {
c.Stripe.SuccessURL = v
}
if v := os.Getenv("STRIPE_CANCEL_URL"); v != "" {
c.Stripe.CancelURL = v
}
if v := os.Getenv("STRIPE_PORTAL_RETURN_URL"); v != "" {
c.Stripe.PortalReturnURL = v
}
if v := os.Getenv("MEMBER_SETTINGS_ENCRYPTION_KEY_ID"); v != "" {
c.MemberSettingsEncryption.CurrentKeyID = v
}
if v := os.Getenv("MEMBER_SETTINGS_ENCRYPTION_KEY"); v != "" {
c.MemberSettingsEncryption.Key = v
}
if v := os.Getenv("PLATFORM_AI_KEY"); v != "" { if v := os.Getenv("PLATFORM_AI_KEY"); v != "" {
c.Platform.AIKey = v c.Platform.AIKey = v
} }
@ -152,18 +220,32 @@ func (c *Config) ApplyEnv() {
if v := firstEnv("THREADS_APP_SECRET", "HAIXUN_THREADS_APP_SECRET"); v != "" { if v := firstEnv("THREADS_APP_SECRET", "HAIXUN_THREADS_APP_SECRET"); v != "" {
c.Platform.ThreadsAppSecret = v c.Platform.ThreadsAppSecret = v
} }
if v := os.Getenv("SCOUT_SESSION_SECRET"); v != "" {
c.Scout.SessionSecret = v
}
if v := os.Getenv("SCOUT_CRAWLER_ENDPOINT"); v != "" {
c.Scout.CrawlerEndpoint = v
}
if v := os.Getenv("SCOUT_CRAWLER_TOKEN"); v != "" {
c.Scout.CrawlerToken = v
}
if v := os.Getenv("SEED_ADMIN_PASSWORD"); v != "" {
c.Seed.AdminPassword = v
}
if v := os.Getenv("WORKER_ID"); v != "" { if v := os.Getenv("WORKER_ID"); v != "" {
c.Worker.ID = v c.Worker.ID = v
} }
if v := os.Getenv("PORT"); v != "" { if v := os.Getenv("PORT"); v != "" {
var p int var p int
valid := true
for _, ch := range v { for _, ch := range v {
if ch < '0' || ch > '9' { if ch < '0' || ch > '9' {
return valid = false
break
} }
p = p*10 + int(ch-'0') p = p*10 + int(ch-'0')
} }
if p > 0 { if valid && p > 0 {
c.Port = p c.Port = p
} }
} }
@ -246,6 +328,28 @@ func (c Config) ValidateProductionDependencies() error {
if c.Auth.AccessSecret == c.Auth.RefreshSecret { if c.Auth.AccessSecret == c.Auth.RefreshSecret {
return fmt.Errorf("AUTH_ACCESS_SECRET and AUTH_REFRESH_SECRET must differ") return fmt.Errorf("AUTH_ACCESS_SECRET and AUTH_REFRESH_SECRET must differ")
} }
fingerprintSecret := strings.TrimSpace(c.Redis.AIModelCacheFingerprintSecret)
if invalidSecret(fingerprintSecret) {
return fmt.Errorf("AI_MODEL_CACHE_FINGERPRINT_SECRET must be at least 32 bytes")
}
if fingerprintSecret == strings.TrimSpace(c.Auth.AccessSecret) ||
fingerprintSecret == strings.TrimSpace(c.Auth.RefreshSecret) ||
fingerprintSecret == strings.TrimSpace(c.Scout.SessionSecret) ||
fingerprintSecret == strings.TrimSpace(c.MemberSettingsEncryption.Key) {
return fmt.Errorf("AI model cache fingerprint secret must be dedicated")
}
if namespace := strings.TrimSpace(c.Redis.Namespace); namespace == "" || strings.ContainsAny(namespace, "{} \t\r\n") {
return fmt.Errorf("Redis.Namespace must be non-empty and must not contain whitespace or braces")
}
if _, err := c.NewMemberSettingsCodec(); err != nil {
return fmt.Errorf("member settings encryption: %w", err)
}
encryptionKey := strings.TrimSpace(c.MemberSettingsEncryption.Key)
if encryptionKey == strings.TrimSpace(c.Auth.AccessSecret) ||
encryptionKey == strings.TrimSpace(c.Auth.RefreshSecret) ||
encryptionKey == strings.TrimSpace(c.Scout.SessionSecret) {
return fmt.Errorf("member settings encryption key must not reuse Auth or Scout secrets")
}
if len(c.CacheRedis) == 0 || strings.TrimSpace(c.CacheRedis[0].Host) == "" { if len(c.CacheRedis) == 0 || strings.TrimSpace(c.CacheRedis[0].Host) == "" {
return fmt.Errorf("CacheRedis is required") return fmt.Errorf("CacheRedis is required")
} }
@ -255,14 +359,56 @@ func (c Config) ValidateProductionDependencies() error {
if len(c.CORSAllowedOrigins) == 0 { if len(c.CORSAllowedOrigins) == 0 {
return fmt.Errorf("CORSAllowedOrigins is required") return fmt.Errorf("CORSAllowedOrigins is required")
} }
if err := validateStripe(c.Stripe.Enabled, c.Stripe.SecretKey, c.Stripe.WebhookSecret, c.Stripe.StarterPriceID, c.Stripe.ProPriceID, c.Stripe.SuccessURL, c.Stripe.CancelURL, c.Stripe.PortalReturnURL); err != nil {
return err
}
return nil return nil
} }
// NewMemberSettingsCodec constructs the dedicated BYOK-at-rest codec.
func (c Config) NewMemberSettingsCodec() (*secretbox.Codec, error) {
return secretbox.New(c.MemberSettingsEncryption.CurrentKeyID, c.MemberSettingsEncryption.Key)
}
func invalidSecret(value string) bool { func invalidSecret(value string) bool {
value = strings.TrimSpace(value) value = strings.TrimSpace(value)
return len(value) < 32 || strings.Contains(value, "REPLACE_WITH") || strings.Contains(value, "CHANGE_ME") return len(value) < 32 || strings.Contains(value, "REPLACE_WITH") || strings.Contains(value, "CHANGE_ME")
} }
func validateStripe(enabled bool, secretKey, webhookSecret, starterPriceID, proPriceID, successURL, cancelURL, portalReturnURL string) error {
if !enabled {
return nil
}
if !strings.HasPrefix(strings.TrimSpace(secretKey), "sk_") {
return fmt.Errorf("STRIPE_SECRET_KEY is required when Stripe is enabled")
}
if !strings.HasPrefix(strings.TrimSpace(webhookSecret), "whsec_") {
return fmt.Errorf("STRIPE_WEBHOOK_SECRET is required when Stripe is enabled")
}
starterPriceID = strings.TrimSpace(starterPriceID)
proPriceID = strings.TrimSpace(proPriceID)
if !strings.HasPrefix(starterPriceID, "price_") || !strings.HasPrefix(proPriceID, "price_") {
return fmt.Errorf("STRIPE_STARTER_PRICE_ID and STRIPE_PRO_PRICE_ID are required when Stripe is enabled")
}
if starterPriceID == proPriceID {
return fmt.Errorf("Stripe Starter and Pro Price IDs must differ")
}
if !strings.Contains(successURL, "{CHECKOUT_ID}") {
return fmt.Errorf("STRIPE_SUCCESS_URL must contain {CHECKOUT_ID}")
}
for name, value := range map[string]string{
"STRIPE_SUCCESS_URL": successURL,
"STRIPE_CANCEL_URL": cancelURL,
"STRIPE_PORTAL_RETURN_URL": portalReturnURL,
} {
parsed, err := url.ParseRequestURI(strings.TrimSpace(value))
if err != nil || parsed.Scheme != "https" || parsed.Host == "" || parsed.User != nil {
return fmt.Errorf("%s must be an absolute HTTPS URL", name)
}
}
return nil
}
func firstEnv(keys ...string) string { func firstEnv(keys ...string) string {
for _, k := range keys { for _, k := range keys {
if v := os.Getenv(k); v != "" { if v := os.Getenv(k); v != "" {

View File

@ -0,0 +1,109 @@
package config
import (
"strings"
"testing"
)
func TestApplyEnvRedisCacheSettings(t *testing.T) {
t.Setenv("REDIS_NAMESPACE", "tenant:prod:v2")
t.Setenv("HAIXUN_REDIS_NAMESPACE", "")
t.Setenv("AI_MODEL_CACHE_FINGERPRINT_SECRET", "dedicated-model-cache-secret-at-least-32-bytes")
t.Setenv("HAIXUN_AI_MODEL_CACHE_FINGERPRINT_SECRET", "")
var c Config
c.ApplyEnv()
if c.Redis.Namespace != "tenant:prod:v2" {
t.Fatalf("namespace=%q", c.Redis.Namespace)
}
if c.Redis.AIModelCacheFingerprintSecret != "dedicated-model-cache-secret-at-least-32-bytes" {
t.Fatalf("fingerprint secret environment override was not applied")
}
}
func TestApplyEnvDefaultsRedisNamespace(t *testing.T) {
t.Setenv("REDIS_NAMESPACE", "")
t.Setenv("HAIXUN_REDIS_NAMESPACE", "")
var c Config
c.ApplyEnv()
if c.Redis.Namespace != "haixun:dev:v1" {
t.Fatalf("namespace=%q", c.Redis.Namespace)
}
}
func TestApplyEnvStripeSettings(t *testing.T) {
t.Setenv("STRIPE_ENABLED", "true")
t.Setenv("STRIPE_SECRET_KEY", "sk_test")
t.Setenv("STRIPE_WEBHOOK_SECRET", "whsec_test")
t.Setenv("STRIPE_STARTER_PRICE_ID", "price_starter")
t.Setenv("STRIPE_PRO_PRICE_ID", "price_pro")
t.Setenv("STRIPE_SUCCESS_URL", "https://app.test/success")
t.Setenv("STRIPE_CANCEL_URL", "https://app.test/cancel")
t.Setenv("STRIPE_PORTAL_RETURN_URL", "https://app.test/billing")
var c Config
c.ApplyEnv()
if !c.Stripe.Enabled || c.Stripe.SecretKey != "sk_test" || c.Stripe.WebhookSecret != "whsec_test" || c.Stripe.StarterPriceID != "price_starter" || c.Stripe.ProPriceID != "price_pro" || c.Stripe.SuccessURL != "https://app.test/success" || c.Stripe.CancelURL != "https://app.test/cancel" || c.Stripe.PortalReturnURL != "https://app.test/billing" {
t.Fatalf("Stripe environment was not fully applied: %+v", c.Stripe)
}
}
func TestApplyEnvScoutAndSeedSecrets(t *testing.T) {
t.Setenv("SCOUT_SESSION_SECRET", "dedicated-scout-secret")
t.Setenv("SCOUT_CRAWLER_ENDPOINT", "http://127.0.0.1:3010")
t.Setenv("SCOUT_CRAWLER_TOKEN", "crawler-token")
t.Setenv("SEED_ADMIN_PASSWORD", "seed-password")
var c Config
c.ApplyEnv()
if c.Scout.SessionSecret != "dedicated-scout-secret" || c.Scout.CrawlerEndpoint != "http://127.0.0.1:3010" || c.Scout.CrawlerToken != "crawler-token" || c.Seed.AdminPassword != "seed-password" {
t.Fatalf("Scout or seed environment was not fully applied")
}
}
func TestApplyEnvInvalidPortDoesNotSkipLaterOverrides(t *testing.T) {
t.Setenv("PORT", "invalid")
t.Setenv("PUBLIC_WEB_BASE", "https://app.test")
var c Config
c.ApplyEnv()
if c.PublicWebBase != "https://app.test" {
t.Fatalf("invalid PORT skipped later environment overrides")
}
}
func TestValidateProductionDependenciesRequiresDedicatedCacheSecret(t *testing.T) {
var c Config
c.Auth.AccessSecret = strings.Repeat("a", 40)
c.Auth.RefreshSecret = strings.Repeat("b", 40)
err := c.ValidateProductionDependencies()
if err == nil || !strings.Contains(err.Error(), "AI_MODEL_CACHE_FINGERPRINT_SECRET") {
t.Fatalf("unexpected validation error %v", err)
}
c.Redis.AIModelCacheFingerprintSecret = c.Auth.AccessSecret
err = c.ValidateProductionDependencies()
if err == nil || !strings.Contains(err.Error(), "must be dedicated") {
t.Fatalf("unexpected reused-secret validation error %v", err)
}
}
func TestValidateStripeEnabledConfiguration(t *testing.T) {
valid := func(successURL string) error {
return validateStripe(true, "sk_test_key", "whsec_test_secret", "price_starter", "price_pro", successURL, "https://app.test/app/usage/checkout?result=cancel", "https://app.test/app/usage/plans")
}
if err := valid("https://app.test/app/usage/checkout?result=success&checkout_id={CHECKOUT_ID}"); err != nil {
t.Fatalf("valid Stripe configuration rejected: %v", err)
}
if err := valid("https://app.test/app/usage/checkout?result=success"); err == nil || !strings.Contains(err.Error(), "{CHECKOUT_ID}") {
t.Fatalf("missing checkout token validation error = %v", err)
}
if err := validateStripe(true, "sk_test_key", "whsec_test_secret", "price_same", "price_same", "https://app.test/success?checkout_id={CHECKOUT_ID}", "https://app.test/cancel", "https://app.test/plans"); err == nil || !strings.Contains(err.Error(), "must differ") {
t.Fatalf("duplicate Price ID validation error = %v", err)
}
if err := validateStripe(true, "sk_test_key", "whsec_test_secret", "price_starter", "price_pro", "http://app.test/success?checkout_id={CHECKOUT_ID}", "https://app.test/cancel", "https://app.test/plans"); err == nil || !strings.Contains(err.Error(), "HTTPS") {
t.Fatalf("insecure URL validation error = %v", err)
}
}

View File

@ -0,0 +1,28 @@
// Code generated by goctl. DO NOT EDIT.
// goctl <no value>
package billing
import (
"net/http"
"apps/backend/internal/logic/billing"
"apps/backend/internal/response"
"apps/backend/internal/svc"
"apps/backend/internal/types"
"github.com/zeromicro/go-zero/rest/httpx"
)
func CreateBillingCheckoutSessionHandler(svcCtx *svc.ServiceContext) http.HandlerFunc {
return func(w http.ResponseWriter, r *http.Request) {
var req types.BillingCheckoutCreateReq
if err := httpx.Parse(r, &req); err != nil {
response.Write(r.Context(), w, nil, response.WrapRequestError(err))
return
}
l := billing.NewCreateBillingCheckoutSessionLogic(r.Context(), svcCtx)
data, err := l.CreateBillingCheckoutSession(&req)
response.Write(r.Context(), w, data, err)
}
}

View File

@ -0,0 +1,20 @@
// Code generated by goctl. DO NOT EDIT.
// goctl <no value>
package billing
import (
"net/http"
"apps/backend/internal/logic/billing"
"apps/backend/internal/response"
"apps/backend/internal/svc"
)
func CreateBillingPortalSessionHandler(svcCtx *svc.ServiceContext) http.HandlerFunc {
return func(w http.ResponseWriter, r *http.Request) {
l := billing.NewCreateBillingPortalSessionLogic(r.Context(), svcCtx)
data, err := l.CreateBillingPortalSession()
response.Write(r.Context(), w, data, err)
}
}

View File

@ -0,0 +1,28 @@
// Code generated by goctl. DO NOT EDIT.
// goctl <no value>
package billing
import (
"net/http"
"apps/backend/internal/logic/billing"
"apps/backend/internal/response"
"apps/backend/internal/svc"
"apps/backend/internal/types"
"github.com/zeromicro/go-zero/rest/httpx"
)
func GetBillingCheckoutSessionHandler(svcCtx *svc.ServiceContext) http.HandlerFunc {
return func(w http.ResponseWriter, r *http.Request) {
var req types.BillingCheckoutGetReq
if err := httpx.Parse(r, &req); err != nil {
response.Write(r.Context(), w, nil, response.WrapRequestError(err))
return
}
l := billing.NewGetBillingCheckoutSessionLogic(r.Context(), svcCtx)
data, err := l.GetBillingCheckoutSession(&req)
response.Write(r.Context(), w, data, err)
}
}

View File

@ -0,0 +1,20 @@
// Code generated by goctl. DO NOT EDIT.
// goctl <no value>
package billing
import (
"net/http"
"apps/backend/internal/logic/billing"
"apps/backend/internal/response"
"apps/backend/internal/svc"
)
func GetBillingSubscriptionHandler(svcCtx *svc.ServiceContext) http.HandlerFunc {
return func(w http.ResponseWriter, r *http.Request) {
l := billing.NewGetBillingSubscriptionLogic(r.Context(), svcCtx)
data, err := l.GetBillingSubscription()
response.Write(r.Context(), w, data, err)
}
}

View File

@ -0,0 +1,20 @@
// Code generated by goctl. DO NOT EDIT.
// goctl <no value>
package billing
import (
"net/http"
"apps/backend/internal/logic/billing"
"apps/backend/internal/response"
"apps/backend/internal/svc"
)
func HandleStripeWebhookHandler(svcCtx *svc.ServiceContext) http.HandlerFunc {
return func(w http.ResponseWriter, r *http.Request) {
l := billing.NewHandleStripeWebhookLogic(r.Context(), svcCtx)
err := l.HandleStripeWebhook()
response.Write(r.Context(), w, nil, err)
}
}

View File

@ -0,0 +1,60 @@
package ping
import (
"context"
"encoding/json"
"net/http"
"net/http/httptest"
"testing"
"apps/backend/internal/response"
"apps/backend/internal/svc"
)
type fakeRedisHealth bool
func (f fakeRedisHealth) PingCtx(context.Context) bool { return bool(f) }
func TestRedisUnavailableReturnsServiceUnavailable(t *testing.T) {
serviceContext := &svc.ServiceContext{Redis: fakeRedisHealth(false)}
tests := []struct {
name string
handler http.HandlerFunc
}{
{name: "ping", handler: PingHandler(serviceContext)},
{name: "health", handler: HealthHandler(serviceContext)},
}
for _, test := range tests {
t.Run(test.name, func(t *testing.T) {
recorder := httptest.NewRecorder()
test.handler.ServeHTTP(recorder, httptest.NewRequest(http.MethodGet, "/api/v1/"+test.name, nil))
if recorder.Code != http.StatusServiceUnavailable {
t.Fatalf("status = %d, want %d", recorder.Code, http.StatusServiceUnavailable)
}
var envelope response.Envelope
if err := json.Unmarshal(recorder.Body.Bytes(), &envelope); err != nil {
t.Fatal(err)
}
if envelope.Code != 503001 || envelope.Message != "redis unavailable" {
t.Fatalf("response = %#v", envelope)
}
})
}
}
func TestRedisAvailableKeepsPingContract(t *testing.T) {
serviceContext := &svc.ServiceContext{Redis: fakeRedisHealth(true)}
recorder := httptest.NewRecorder()
PingHandler(serviceContext).ServeHTTP(recorder, httptest.NewRequest(http.MethodGet, "/api/v1/ping", nil))
if recorder.Code != http.StatusOK {
t.Fatalf("status = %d, want %d", recorder.Code, http.StatusOK)
}
var envelope response.Envelope
if err := json.Unmarshal(recorder.Body.Bytes(), &envelope); err != nil {
t.Fatal(err)
}
if envelope.Code != 102000 {
t.Fatalf("code = %d, want 102000", envelope.Code)
}
}

View File

@ -8,6 +8,7 @@ import (
admin "apps/backend/internal/handler/admin" admin "apps/backend/internal/handler/admin"
auth "apps/backend/internal/handler/auth" auth "apps/backend/internal/handler/auth"
billing "apps/backend/internal/handler/billing"
compose "apps/backend/internal/handler/compose" compose "apps/backend/internal/handler/compose"
extension "apps/backend/internal/handler/extension" extension "apps/backend/internal/handler/extension"
inspire "apps/backend/internal/handler/inspire" inspire "apps/backend/internal/handler/inspire"
@ -187,6 +188,50 @@ func RegisterHandlers(server *rest.Server, serverCtx *svc.ServiceContext) {
rest.WithPrefix("/api/v1/auth"), rest.WithPrefix("/api/v1/auth"),
) )
server.AddRoutes(
rest.WithMiddlewares(
[]rest.Middleware{serverCtx.AuthJWT},
[]rest.Route{
{
Method: http.MethodPost,
Path: "/checkout-sessions",
Handler: billing.CreateBillingCheckoutSessionHandler(serverCtx),
},
{
Method: http.MethodGet,
Path: "/checkout-sessions/:id",
Handler: billing.GetBillingCheckoutSessionHandler(serverCtx),
},
{
Method: http.MethodPost,
Path: "/portal-sessions",
Handler: billing.CreateBillingPortalSessionHandler(serverCtx),
},
{
Method: http.MethodGet,
Path: "/subscription",
Handler: billing.GetBillingSubscriptionHandler(serverCtx),
},
}...,
),
rest.WithJwt(serverCtx.Config.Auth.AccessSecret),
rest.WithPrefix("/api/v1/billing"),
)
server.AddRoutes(
rest.WithMiddlewares(
[]rest.Middleware{serverCtx.StripeRawBody},
[]rest.Route{
{
Method: http.MethodPost,
Path: "/stripe/webhook",
Handler: billing.HandleStripeWebhookHandler(serverCtx),
},
}...,
),
rest.WithPrefix("/api/v1/billing"),
)
server.AddRoutes( server.AddRoutes(
rest.WithMiddlewares( rest.WithMiddlewares(
[]rest.Middleware{serverCtx.AuthJWT}, []rest.Middleware{serverCtx.AuthJWT},
@ -999,6 +1044,11 @@ func RegisterHandlers(server *rest.Server, serverCtx *svc.ServiceContext) {
Path: "/prefs", Path: "/prefs",
Handler: usage.SetUsagePrefsHandler(serverCtx), Handler: usage.SetUsagePrefsHandler(serverCtx),
}, },
{
Method: http.MethodGet,
Path: "/tenant-analytics",
Handler: usage.GetTenantUsageAnalyticsHandler(serverCtx),
},
{ {
Method: http.MethodGet, Method: http.MethodGet,
Path: "/tenant-summary", Path: "/tenant-summary",

View File

@ -0,0 +1,28 @@
// Code generated by goctl. DO NOT EDIT.
// goctl <no value>
package usage
import (
"net/http"
"apps/backend/internal/logic/usage"
"apps/backend/internal/response"
"apps/backend/internal/svc"
"apps/backend/internal/types"
"github.com/zeromicro/go-zero/rest/httpx"
)
func GetTenantUsageAnalyticsHandler(svcCtx *svc.ServiceContext) http.HandlerFunc {
return func(w http.ResponseWriter, r *http.Request) {
var req types.UsageTenantAnalyticsReq
if err := httpx.Parse(r, &req); err != nil {
response.Write(r.Context(), w, nil, response.WrapRequestError(err))
return
}
l := usage.NewGetTenantUsageAnalyticsLogic(r.Context(), svcCtx)
data, err := l.GetTenantUsageAnalytics(&req)
response.Write(r.Context(), w, data, err)
}
}

View File

@ -2,46 +2,155 @@ package redislock
import ( import (
"context" "context"
"crypto/rand"
"encoding/hex"
"fmt" "fmt"
"strings"
"sync"
"time" "time"
"github.com/zeromicro/go-zero/core/stores/redis" "github.com/zeromicro/go-zero/core/stores/redis"
) )
// Lock is a small Redis lease. Value must uniquely identify the worker that const (
// owns it; Release uses a Lua compare-and-delete to avoid deleting a renewed renewScript = `
// lease owned by another worker. if redis.call('GET', KEYS[1]) == ARGV[1] then
type Lock struct { return redis.call('PEXPIRE', KEYS[1], ARGV[2])
client *redis.Redis end
key string return 0
owner string `
ttl time.Duration releaseScript = `
}
func New(conf redis.RedisConf, key, owner string, ttl time.Duration) *Lock {
return &Lock{client: redis.MustNewRedis(conf), key: key, owner: owner, ttl: ttl}
}
func (l *Lock) Acquire(ctx context.Context) (bool, error) {
if l == nil || l.client == nil || l.key == "" || l.owner == "" {
return false, fmt.Errorf("invalid redis lock")
}
seconds := int(l.ttl.Seconds())
if seconds < 1 {
seconds = 1
}
return l.client.SetnxExCtx(ctx, l.key, l.owner, seconds)
}
func (l *Lock) Release(ctx context.Context) error {
if l == nil || l.client == nil {
return fmt.Errorf("invalid redis lock")
}
_, err := l.client.EvalCtx(ctx, `
if redis.call('GET', KEYS[1]) == ARGV[1] then if redis.call('GET', KEYS[1]) == ARGV[1] then
return redis.call('DEL', KEYS[1]) return redis.call('DEL', KEYS[1])
end end
return 0 return 0
`, []string{l.key}, l.owner) `
)
// Client is the subset of Redis used by Lock.
type Client interface {
SetnxExCtx(ctx context.Context, key, value string, seconds int) (bool, error)
EvalCtx(ctx context.Context, script string, keys []string, args ...any) (any, error)
}
// Lock is a small Redis lease. Each acquisition gets a unique token so an old
// owner cannot renew or release a successor's lease.
type Lock struct {
client Client
key string
prefix string
ttl time.Duration
mu sync.RWMutex
token string
}
func New(conf redis.RedisConf, key, ownerPrefix string, ttl time.Duration) *Lock {
return NewWithClient(redis.MustNewRedis(conf), key, ownerPrefix, ttl)
}
func NewWithClient(client Client, key, ownerPrefix string, ttl time.Duration) *Lock {
return &Lock{client: client, key: key, prefix: strings.TrimSpace(ownerPrefix), ttl: ttl}
}
func (l *Lock) Acquire(ctx context.Context) (bool, error) {
if err := l.validate(); err != nil {
return false, err
}
token, err := newToken(l.prefix)
if err != nil {
return false, err
}
acquired, err := l.client.SetnxExCtx(ctx, l.key, token, durationSeconds(l.ttl))
if err != nil || !acquired {
return acquired, err
}
l.mu.Lock()
l.token = token
l.mu.Unlock()
return true, nil
}
// Renew extends the lease only while this acquisition still owns it.
func (l *Lock) Renew(ctx context.Context) (bool, error) {
if err := l.validate(); err != nil {
return false, err
}
token := l.currentToken()
if token == "" {
return false, nil
}
ttlMilliseconds := l.ttl.Milliseconds()
if ttlMilliseconds < 1 {
ttlMilliseconds = 1
}
result, err := l.client.EvalCtx(ctx, renewScript, []string{l.key}, token, ttlMilliseconds)
if err != nil {
return false, err
}
return resultIsOne(result), nil
}
func (l *Lock) Release(ctx context.Context) error {
if err := l.validate(); err != nil {
return err
}
token := l.currentToken()
if token == "" {
return nil
}
result, err := l.client.EvalCtx(ctx, releaseScript, []string{l.key}, token)
if err == nil && resultIsOne(result) {
l.mu.Lock()
if l.token == token {
l.token = ""
}
l.mu.Unlock()
}
return err return err
} }
func (l *Lock) validate() error {
if l == nil || l.client == nil || l.key == "" || l.prefix == "" || l.ttl <= 0 {
return fmt.Errorf("invalid redis lock")
}
return nil
}
func (l *Lock) currentToken() string {
l.mu.RLock()
defer l.mu.RUnlock()
return l.token
}
func newToken(prefix string) (string, error) {
random := make([]byte, 16)
if _, err := rand.Read(random); err != nil {
return "", fmt.Errorf("generate redis lock token: %w", err)
}
return prefix + ":" + hex.EncodeToString(random), nil
}
func durationSeconds(ttl time.Duration) int {
seconds := int(ttl / time.Second)
if ttl%time.Second != 0 {
seconds++
}
if seconds < 1 {
seconds = 1
}
return seconds
}
func resultIsOne(result any) bool {
switch value := result.(type) {
case int64:
return value == 1
case int:
return value == 1
case uint64:
return value == 1
default:
return false
}
}

View File

@ -0,0 +1,105 @@
package redislock
import (
"context"
"fmt"
"strings"
"testing"
"time"
)
type fakeClient struct {
value string
ttlMS int64
}
func (f *fakeClient) SetnxExCtx(ctx context.Context, key, value string, seconds int) (bool, error) {
if err := ctx.Err(); err != nil {
return false, err
}
if f.value != "" {
return false, nil
}
f.value = value
f.ttlMS = int64(seconds) * 1000
return true, nil
}
func (f *fakeClient) EvalCtx(ctx context.Context, script string, keys []string, args ...any) (any, error) {
if err := ctx.Err(); err != nil {
return nil, err
}
owner := fmt.Sprint(args[0])
if f.value != owner {
return int64(0), nil
}
if script == renewScript {
f.ttlMS = args[1].(int64)
return int64(1), nil
}
if script == releaseScript {
f.value = ""
return int64(1), nil
}
return nil, fmt.Errorf("unexpected script")
}
func TestLockUsesUniqueTokenPerAcquisition(t *testing.T) {
client := &fakeClient{}
lock := NewWithClient(client, "locks:outbox", "worker-7", time.Minute)
acquired, err := lock.Acquire(context.Background())
if err != nil || !acquired {
t.Fatalf("Acquire() = %v, %v", acquired, err)
}
first := client.value
if !strings.HasPrefix(first, "worker-7:") {
t.Fatalf("token %q has no worker prefix", first)
}
if err := lock.Release(context.Background()); err != nil {
t.Fatal(err)
}
acquired, err = lock.Acquire(context.Background())
if err != nil || !acquired {
t.Fatalf("second Acquire() = %v, %v", acquired, err)
}
if client.value == first {
t.Fatal("successful acquisitions reused a token")
}
}
func TestRenewAndReleaseRequireCurrentOwnership(t *testing.T) {
client := &fakeClient{}
lock := NewWithClient(client, "locks:outbox", "worker-7", 3*time.Minute)
if acquired, err := lock.Acquire(context.Background()); err != nil || !acquired {
t.Fatalf("Acquire() = %v, %v", acquired, err)
}
client.ttlMS = 1
renewed, err := lock.Renew(context.Background())
if err != nil || !renewed || client.ttlMS != int64(3*time.Minute/time.Millisecond) {
t.Fatalf("Renew() = %v, %v, ttl=%d", renewed, err, client.ttlMS)
}
client.value = "worker-8:successor"
renewed, err = lock.Renew(context.Background())
if err != nil || renewed {
t.Fatalf("lost-owner Renew() = %v, %v", renewed, err)
}
if err := lock.Release(context.Background()); err != nil {
t.Fatal(err)
}
if client.value != "worker-8:successor" {
t.Fatalf("Release removed successor lease: %q", client.value)
}
}
func TestLockMethodsHonorCanceledContext(t *testing.T) {
client := &fakeClient{}
lock := NewWithClient(client, "locks:outbox", "worker-7", time.Minute)
ctx, cancel := context.WithCancel(context.Background())
cancel()
if _, err := lock.Acquire(ctx); err == nil {
t.Fatal("Acquire succeeded with canceled context")
}
}

View File

@ -0,0 +1,121 @@
// Package secretbox encrypts member-owned secrets for persistence.
package secretbox
import (
"crypto/aes"
"crypto/cipher"
"crypto/rand"
"encoding/base64"
"encoding/binary"
"errors"
"fmt"
"io"
"strings"
)
const prefix = "enc:v1:"
var (
ErrInvalidConfig = errors.New("invalid secretbox configuration")
ErrInvalidCiphertext = errors.New("invalid encrypted secret")
ErrUnsupportedVersion = errors.New("unsupported encrypted secret version")
)
// Codec is an AES-256-GCM codec for the current encryption key.
type Codec struct {
keyID string
aead cipher.AEAD
rand io.Reader
}
// New constructs a codec from a key ID and a base64-encoded 32-byte key.
func New(keyID, encodedKey string) (*Codec, error) {
keyID = strings.TrimSpace(keyID)
encodedKey = strings.TrimSpace(encodedKey)
if keyID == "" || strings.Contains(keyID, ":") || encodedKey == "" {
return nil, fmt.Errorf("%w: key ID and key are required", ErrInvalidConfig)
}
key, err := base64.StdEncoding.DecodeString(encodedKey)
if err != nil {
return nil, fmt.Errorf("%w: key must be base64: %v", ErrInvalidConfig, err)
}
if len(key) != 32 {
return nil, fmt.Errorf("%w: decoded key must be 32 bytes", ErrInvalidConfig)
}
block, err := aes.NewCipher(key)
if err != nil {
return nil, fmt.Errorf("%w: %v", ErrInvalidConfig, err)
}
aead, err := cipher.NewGCM(block)
if err != nil {
return nil, fmt.Errorf("%w: %v", ErrInvalidConfig, err)
}
return &Codec{keyID: keyID, aead: aead, rand: rand.Reader}, nil
}
// Encrypt encrypts plaintext and binds it to its member, field, and provider.
func (c *Codec) Encrypt(uid int64, field, provider, plaintext string) (string, error) {
if c == nil || c.aead == nil {
return "", ErrInvalidConfig
}
nonce := make([]byte, c.aead.NonceSize())
if _, err := io.ReadFull(c.rand, nonce); err != nil {
return "", fmt.Errorf("generate encryption nonce: %w", err)
}
sealed := c.aead.Seal(nil, nonce, []byte(plaintext), additionalData(uid, field, provider))
payload := append(nonce, sealed...)
return prefix + c.keyID + ":" + base64.RawURLEncoding.EncodeToString(payload), nil
}
// Decrypt decrypts encrypted values. Unprefixed values are legacy plaintext.
func (c *Codec) Decrypt(uid int64, field, provider, value string) (string, error) {
if !strings.HasPrefix(value, "enc:") {
return value, nil
}
if !strings.HasPrefix(value, prefix) {
return "", ErrUnsupportedVersion
}
if c == nil || c.aead == nil {
return "", ErrInvalidConfig
}
rest := strings.TrimPrefix(value, prefix)
keyID, payloadText, ok := strings.Cut(rest, ":")
if !ok || keyID == "" || payloadText == "" {
return "", ErrInvalidCiphertext
}
if keyID != c.keyID {
return "", fmt.Errorf("%w: unknown key ID %q", ErrInvalidCiphertext, keyID)
}
payload, err := base64.RawURLEncoding.DecodeString(payloadText)
if err != nil || len(payload) < c.aead.NonceSize()+c.aead.Overhead() {
return "", ErrInvalidCiphertext
}
nonce := payload[:c.aead.NonceSize()]
ciphertext := payload[c.aead.NonceSize():]
plaintext, err := c.aead.Open(nil, nonce, ciphertext, additionalData(uid, field, provider))
if err != nil {
return "", fmt.Errorf("%w: authentication failed", ErrInvalidCiphertext)
}
return string(plaintext), nil
}
// IsEncrypted reports whether a value uses the encrypted-value namespace.
func IsEncrypted(value string) bool {
return strings.HasPrefix(value, "enc:")
}
func additionalData(uid int64, field, provider string) []byte {
fieldBytes := []byte(field)
providerBytes := []byte(provider)
data := make([]byte, 8+4+len(fieldBytes)+4+len(providerBytes))
binary.BigEndian.PutUint64(data, uint64(uid))
offset := 8
binary.BigEndian.PutUint32(data[offset:], uint32(len(fieldBytes)))
offset += 4
copy(data[offset:], fieldBytes)
offset += len(fieldBytes)
binary.BigEndian.PutUint32(data[offset:], uint32(len(providerBytes)))
offset += 4
copy(data[offset:], providerBytes)
return data
}

View File

@ -0,0 +1,89 @@
package secretbox
import (
"encoding/base64"
"errors"
"strings"
"testing"
)
func testCodec(t *testing.T) *Codec {
t.Helper()
c, err := New("test-key", base64.StdEncoding.EncodeToString(make([]byte, 32)))
if err != nil {
t.Fatal(err)
}
return c
}
func TestCodecRoundTripUsesRandomNonce(t *testing.T) {
c := testCodec(t)
one, err := c.Encrypt(42, "api_key", "xai", "top-secret")
if err != nil {
t.Fatal(err)
}
two, err := c.Encrypt(42, "api_key", "xai", "top-secret")
if err != nil {
t.Fatal(err)
}
if one == two {
t.Fatal("ciphertexts must differ")
}
if !strings.HasPrefix(one, "enc:v1:test-key:") {
t.Fatalf("unexpected prefix: %q", one)
}
got, err := c.Decrypt(42, "api_key", "xai", one)
if err != nil || got != "top-secret" {
t.Fatalf("round trip = %q, %v", got, err)
}
}
func TestCodecAADAndTamperingFail(t *testing.T) {
c := testCodec(t)
value, err := c.Encrypt(42, "api_key", "xai", "top-secret")
if err != nil {
t.Fatal(err)
}
parts := strings.Split(value, ":")
payload, err := base64.RawURLEncoding.DecodeString(parts[3])
if err != nil {
t.Fatal(err)
}
payload[len(payload)-1] ^= 0x01
tampered := strings.Join([]string{parts[0], parts[1], parts[2], base64.RawURLEncoding.EncodeToString(payload)}, ":")
for _, tc := range []struct {
name string
uid int64
field string
provider string
value string
}{
{name: "uid", uid: 43, field: "api_key", provider: "xai", value: value},
{name: "field", uid: 42, field: "exa_api_key", provider: "xai", value: value},
{name: "provider", uid: 42, field: "api_key", provider: "other", value: value},
{name: "tamper", uid: 42, field: "api_key", provider: "xai", value: tampered},
} {
t.Run(tc.name, func(t *testing.T) {
if _, err := c.Decrypt(tc.uid, tc.field, tc.provider, tc.value); !errors.Is(err, ErrInvalidCiphertext) {
t.Fatalf("error = %v", err)
}
})
}
}
func TestCodecCompatibilityAndValidation(t *testing.T) {
c := testCodec(t)
got, err := c.Decrypt(1, "api_key", "xai", "legacy-plaintext")
if err != nil || got != "legacy-plaintext" {
t.Fatalf("legacy = %q, %v", got, err)
}
if _, err := c.Decrypt(1, "api_key", "xai", "enc:v2:key:data"); !errors.Is(err, ErrUnsupportedVersion) {
t.Fatalf("version error = %v", err)
}
if _, err := c.Decrypt(1, "api_key", "xai", "enc:v1:other:data"); !errors.Is(err, ErrInvalidCiphertext) {
t.Fatalf("key ID error = %v", err)
}
if _, err := New("key", base64.StdEncoding.EncodeToString(make([]byte, 31))); !errors.Is(err, ErrInvalidConfig) {
t.Fatalf("key size error = %v", err)
}
}

View File

@ -37,7 +37,10 @@ func (l *SetRolesLogic) SetRoles(req *types.AdminRolesReq) (resp *types.MemberPu
} }
} }
if wasAdmin && !willAdmin { if wasAdmin && !willAdmin {
n, _ := l.svcCtx.Members.CountActiveAdmins(l.ctx) n, countErr := l.svcCtx.Members.CountActiveAdmins(l.ctx)
if countErr != nil {
return nil, countErr
}
if n <= 1 { if n <= 1 {
return nil, domain.ErrLastAdmin return nil, domain.ErrLastAdmin
} }

View File

@ -4,7 +4,9 @@ import (
"context" "context"
"strconv" "strconv"
"apps/backend/internal/middleware"
"apps/backend/internal/module/member/domain" "apps/backend/internal/module/member/domain"
"apps/backend/internal/response"
"apps/backend/internal/svc" "apps/backend/internal/svc"
"apps/backend/internal/types" "apps/backend/internal/types"
@ -22,6 +24,7 @@ func NewSetStatusLogic(ctx context.Context, svcCtx *svc.ServiceContext) *SetStat
} }
func (l *SetStatusLogic) SetStatus(req *types.AdminStatusReq) (resp *types.MemberPublic, err error) { func (l *SetStatusLogic) SetStatus(req *types.AdminStatusReq) (resp *types.MemberPublic, err error) {
actor, _ := middleware.UIDFrom(l.ctx)
uid, err := strconv.ParseInt(req.Uid, 10, 64) uid, err := strconv.ParseInt(req.Uid, 10, 64)
if err != nil { if err != nil {
return nil, err return nil, err
@ -30,7 +33,25 @@ func (l *SetStatusLogic) SetStatus(req *types.AdminStatusReq) (resp *types.Membe
switch status { switch status {
case domain.StatusActive, domain.StatusSuspended, domain.StatusUnverified, domain.StatusUninitialized: case domain.StatusActive, domain.StatusSuspended, domain.StatusUnverified, domain.StatusUninitialized:
default: default:
return nil, domain.ErrNotFound // treat invalid as not found-ish; better weak password reuse? use ErrNotFound ok for now return nil, response.Biz(400, 400001, "invalid member status")
}
if status != domain.StatusActive {
if uid == actor {
return nil, domain.ErrSelfSuspend
}
member, findErr := l.svcCtx.Members.FindByUID(l.ctx, uid)
if findErr != nil {
return nil, findErr
}
if member.IsAdmin() {
n, countErr := l.svcCtx.Members.CountActiveAdmins(l.ctx)
if countErr != nil {
return nil, countErr
}
if n <= 1 {
return nil, domain.ErrLastAdmin
}
}
} }
m, err := l.svcCtx.Auth.UpdateStatus(l.ctx, uid, status) m, err := l.svcCtx.Auth.UpdateStatus(l.ctx, uid, status)
if err != nil { if err != nil {

View File

@ -36,7 +36,10 @@ func (l *SetSuspendedLogic) SetSuspended(req *types.AdminSuspendReq) (resp *type
return nil, err return nil, err
} }
if req.Suspended && m.IsAdmin() { if req.Suspended && m.IsAdmin() {
n, _ := l.svcCtx.Members.CountActiveAdmins(l.ctx) n, countErr := l.svcCtx.Members.CountActiveAdmins(l.ctx)
if countErr != nil {
return nil, countErr
}
if n <= 1 { if n <= 1 {
return nil, domain.ErrLastAdmin return nil, domain.ErrLastAdmin
} }

View File

@ -30,10 +30,7 @@ func (l *OAuthCallbackLogic) OAuthCallback(req *types.AuthOAuthCallbackReq) (res
var subject, email, name string var subject, email, name string
switch provider { switch provider {
case domain.PlatformGoogle: case domain.PlatformGoogle:
token := req.AccessToken token := strings.TrimSpace(req.IdToken)
if token == "" {
token = req.IdToken
}
if token == "" { if token == "" {
return nil, domain.ErrOAuthFailed return nil, domain.ErrOAuthFailed
} }

View File

@ -0,0 +1,42 @@
package billing
import (
"context"
"apps/backend/internal/middleware"
billingModule "apps/backend/internal/module/billing"
"apps/backend/internal/svc"
"apps/backend/internal/types"
"github.com/zeromicro/go-zero/core/logx"
)
type CreateBillingCheckoutSessionLogic struct {
logx.Logger
ctx context.Context
svcCtx *svc.ServiceContext
}
func NewCreateBillingCheckoutSessionLogic(ctx context.Context, svcCtx *svc.ServiceContext) *CreateBillingCheckoutSessionLogic {
return &CreateBillingCheckoutSessionLogic{
Logger: logx.WithContext(ctx),
ctx: ctx,
svcCtx: svcCtx,
}
}
func (l *CreateBillingCheckoutSessionLogic) CreateBillingCheckoutSession(req *types.BillingCheckoutCreateReq) (resp *types.BillingCheckoutData, err error) {
uid, ok := middleware.UIDFrom(l.ctx)
if !ok {
return nil, billingModule.ErrInvalidRequest
}
a, err := l.svcCtx.Billing.CreateCheckout(l.ctx, uid, req.PlanId, req.RequestId)
if err != nil {
return nil, err
}
return checkoutData(a), nil
}
func checkoutData(a *billingModule.CheckoutAttempt) *types.BillingCheckoutData {
return &types.BillingCheckoutData{Id: a.ID, SessionId: a.SessionID, Url: a.URL, ExpiresAt: a.ExpiresAt, Status: a.Status, PaymentStatus: a.PaymentStatus, FulfillmentStatus: a.FulfillmentStatus, PlanId: a.PlanID}
}

View File

@ -0,0 +1,38 @@
package billing
import (
"context"
"apps/backend/internal/middleware"
billingModule "apps/backend/internal/module/billing"
"apps/backend/internal/svc"
"apps/backend/internal/types"
"github.com/zeromicro/go-zero/core/logx"
)
type CreateBillingPortalSessionLogic struct {
logx.Logger
ctx context.Context
svcCtx *svc.ServiceContext
}
func NewCreateBillingPortalSessionLogic(ctx context.Context, svcCtx *svc.ServiceContext) *CreateBillingPortalSessionLogic {
return &CreateBillingPortalSessionLogic{
Logger: logx.WithContext(ctx),
ctx: ctx,
svcCtx: svcCtx,
}
}
func (l *CreateBillingPortalSessionLogic) CreateBillingPortalSession() (resp *types.BillingPortalData, err error) {
uid, ok := middleware.UIDFrom(l.ctx)
if !ok {
return nil, billingModule.ErrInvalidRequest
}
url, err := l.svcCtx.Billing.Portal(l.ctx, uid)
if err != nil {
return nil, err
}
return &types.BillingPortalData{Url: url}, nil
}

View File

@ -0,0 +1,38 @@
package billing
import (
"context"
"apps/backend/internal/middleware"
billingModule "apps/backend/internal/module/billing"
"apps/backend/internal/svc"
"apps/backend/internal/types"
"github.com/zeromicro/go-zero/core/logx"
)
type GetBillingCheckoutSessionLogic struct {
logx.Logger
ctx context.Context
svcCtx *svc.ServiceContext
}
func NewGetBillingCheckoutSessionLogic(ctx context.Context, svcCtx *svc.ServiceContext) *GetBillingCheckoutSessionLogic {
return &GetBillingCheckoutSessionLogic{
Logger: logx.WithContext(ctx),
ctx: ctx,
svcCtx: svcCtx,
}
}
func (l *GetBillingCheckoutSessionLogic) GetBillingCheckoutSession(req *types.BillingCheckoutGetReq) (resp *types.BillingCheckoutData, err error) {
uid, ok := middleware.UIDFrom(l.ctx)
if !ok {
return nil, billingModule.ErrInvalidRequest
}
a, err := l.svcCtx.Billing.GetCheckout(l.ctx, uid, req.Id)
if err != nil {
return nil, err
}
return checkoutData(a), nil
}

View File

@ -0,0 +1,38 @@
package billing
import (
"context"
"apps/backend/internal/middleware"
billingModule "apps/backend/internal/module/billing"
"apps/backend/internal/svc"
"apps/backend/internal/types"
"github.com/zeromicro/go-zero/core/logx"
)
type GetBillingSubscriptionLogic struct {
logx.Logger
ctx context.Context
svcCtx *svc.ServiceContext
}
func NewGetBillingSubscriptionLogic(ctx context.Context, svcCtx *svc.ServiceContext) *GetBillingSubscriptionLogic {
return &GetBillingSubscriptionLogic{
Logger: logx.WithContext(ctx),
ctx: ctx,
svcCtx: svcCtx,
}
}
func (l *GetBillingSubscriptionLogic) GetBillingSubscription() (resp *types.BillingSubscriptionData, err error) {
uid, ok := middleware.UIDFrom(l.ctx)
if !ok {
return nil, billingModule.ErrInvalidRequest
}
s, err := l.svcCtx.Billing.Subscription(l.ctx, uid)
if err != nil {
return nil, err
}
return &types.BillingSubscriptionData{PlanId: s.PlanID, PaidPlanId: s.PaidPlanID, Status: s.Status, CustomerId: s.CustomerID, SubscriptionId: s.SubscriptionID, CurrentPeriodStart: s.CurrentPeriodStart, CurrentPeriodEnd: s.CurrentPeriodEnd, CancelAtPeriodEnd: s.CancelAtPeriodEnd, AdminOverride: s.AdminOverride}, nil
}

View File

@ -0,0 +1,32 @@
package billing
import (
"context"
"apps/backend/internal/middleware"
billingModule "apps/backend/internal/module/billing"
"apps/backend/internal/svc"
"github.com/zeromicro/go-zero/core/logx"
)
type HandleStripeWebhookLogic struct {
logx.Logger
ctx context.Context
svcCtx *svc.ServiceContext
}
func NewHandleStripeWebhookLogic(ctx context.Context, svcCtx *svc.ServiceContext) *HandleStripeWebhookLogic {
return &HandleStripeWebhookLogic{
Logger: logx.WithContext(ctx),
ctx: ctx,
svcCtx: svcCtx,
}
}
func (l *HandleStripeWebhookLogic) HandleStripeWebhook() error {
payload, signature, ok := middleware.StripeWebhookPayload(l.ctx)
if !ok {
return billingModule.ErrInvalidSignature
}
return l.svcCtx.Billing.HandleWebhook(l.ctx, payload, signature)
}

View File

@ -20,5 +20,9 @@ func NewHealthLogic(ctx context.Context, svcCtx *svc.ServiceContext) *HealthLogi
} }
func (l *HealthLogic) Health() (resp *types.HealthData, err error) { func (l *HealthLogic) Health() (resp *types.HealthData, err error) {
if err := redisHealth(l.ctx, l.svcCtx); err != nil {
l.Errorf("health redis check failed: %v", err)
return nil, err
}
return &types.HealthData{Status: "ok", Service: "haixun-gateway"}, nil return &types.HealthData{Status: "ok", Service: "haixun-gateway"}, nil
} }

View File

@ -2,7 +2,10 @@ package ping
import ( import (
"context" "context"
"net/http"
"time"
"apps/backend/internal/response"
"apps/backend/internal/svc" "apps/backend/internal/svc"
"apps/backend/internal/types" "apps/backend/internal/types"
@ -20,5 +23,21 @@ func NewPingLogic(ctx context.Context, svcCtx *svc.ServiceContext) *PingLogic {
} }
func (l *PingLogic) Ping() (resp *types.HealthData, err error) { func (l *PingLogic) Ping() (resp *types.HealthData, err error) {
if err := redisHealth(l.ctx, l.svcCtx); err != nil {
l.Errorf("ping redis health failed: %v", err)
return nil, err
}
return &types.HealthData{Status: "ok", Service: "haixun-gateway"}, nil return &types.HealthData{Status: "ok", Service: "haixun-gateway"}, nil
} }
func redisHealth(ctx context.Context, svcCtx *svc.ServiceContext) error {
if svcCtx == nil || svcCtx.Redis == nil {
return response.Biz(http.StatusServiceUnavailable, 503001, "redis unavailable")
}
pingCtx, cancel := context.WithTimeout(ctx, time.Second)
defer cancel()
if !svcCtx.Redis.PingCtx(pingCtx) {
return response.Biz(http.StatusServiceUnavailable, 503001, "redis unavailable")
}
return nil
}

View File

@ -23,16 +23,19 @@ func NewGetPlacementLogic(ctx context.Context, svcCtx *svc.ServiceContext) *GetP
func (l *GetPlacementLogic) GetPlacement() (resp *types.PlacementSettingsData, err error) { func (l *GetPlacementLogic) GetPlacement() (resp *types.PlacementSettingsData, err error) {
uid, _ := middleware.UIDFrom(l.ctx) uid, _ := middleware.UIDFrom(l.ctx)
st, _ := l.svcCtx.Members.GetSettings(l.ctx, uid) st, err := l.svcCtx.Members.GetSettings(l.ctx, uid)
if err != nil {
return nil, err
}
if st == nil { if st == nil {
st = domain.DefaultSettings(uid) st = domain.DefaultSettings(uid)
} }
return &types.PlacementSettingsData{ return &types.PlacementSettingsData{
WebSearchProvider: "exa", WebSearchProvider: "exa",
ExpandStrategy: st.ExpandStrategy, ExpandStrategy: st.ExpandStrategy,
BraveApiKeyConfigured: false, BraveApiKeyConfigured: false,
ExaApiKeyConfigured: st.ExaAPIKeyConfigured, ExaApiKeyConfigured: st.ExaAPIKeyConfigured,
DevModeEnabled: st.DevModeEnabled, DevModeEnabled: st.DevModeEnabled,
PlatformKeyAvailable: l.svcCtx.Config.Platform.ExaKey != "", PlatformKeyAvailable: l.svcCtx.Config.Platform.ExaKey != "",
}, nil }, nil
} }

View File

@ -29,7 +29,10 @@ func (l *ListModelsLogic) ListModels(req *types.ListModelsReq) (*types.ListModel
provider = ai.ProviderXAI provider = ai.ProviderXAI
} }
uid, _ := middleware.UIDFrom(l.ctx) uid, _ := middleware.UIDFrom(l.ctx)
st, _ := l.svcCtx.Members.GetSettings(l.ctx, uid) st, err := l.svcCtx.Members.GetSettings(l.ctx, uid)
if err != nil {
return nil, err
}
if st == nil { if st == nil {
st = domain.DefaultSettings(uid) st = domain.DefaultSettings(uid)
} }

View File

@ -24,7 +24,10 @@ func NewSaveAiLogic(ctx context.Context, svcCtx *svc.ServiceContext) *SaveAiLogi
func (l *SaveAiLogic) SaveAi(req *types.SaveAiReq) (*types.AiSettingsData, error) { func (l *SaveAiLogic) SaveAi(req *types.SaveAiReq) (*types.AiSettingsData, error) {
uid, _ := middleware.UIDFrom(l.ctx) uid, _ := middleware.UIDFrom(l.ctx)
st, _ := l.svcCtx.Members.GetSettings(l.ctx, uid) st, err := l.svcCtx.Members.GetSettings(l.ctx, uid)
if err != nil {
return nil, err
}
if st == nil { if st == nil {
st = domain.DefaultSettings(uid) st = domain.DefaultSettings(uid)
} }

View File

@ -23,7 +23,10 @@ func NewSavePlacementLogic(ctx context.Context, svcCtx *svc.ServiceContext) *Sav
func (l *SavePlacementLogic) SavePlacement(req *types.SavePlacementReq) (resp *types.PlacementSettingsData, err error) { func (l *SavePlacementLogic) SavePlacement(req *types.SavePlacementReq) (resp *types.PlacementSettingsData, err error) {
uid, _ := middleware.UIDFrom(l.ctx) uid, _ := middleware.UIDFrom(l.ctx)
st, _ := l.svcCtx.Members.GetSettings(l.ctx, uid) st, err := l.svcCtx.Members.GetSettings(l.ctx, uid)
if err != nil {
return nil, err
}
if st == nil { if st == nil {
st = domain.DefaultSettings(uid) st = domain.DefaultSettings(uid)
} }
@ -41,6 +44,8 @@ func (l *SavePlacementLogic) SavePlacement(req *types.SavePlacementReq) (resp *t
st.ExaAPIKey = "" st.ExaAPIKey = ""
st.ExaAPIKeyConfigured = false st.ExaAPIKeyConfigured = false
} }
_ = l.svcCtx.Members.SaveSettings(l.ctx, uid, st) if err := l.svcCtx.Members.SaveSettings(l.ctx, uid, st); err != nil {
return nil, err
}
return NewGetPlacementLogic(l.ctx, l.svcCtx).GetPlacement() return NewGetPlacementLogic(l.ctx, l.svcCtx).GetPlacement()
} }

View File

@ -0,0 +1,90 @@
package usage
import (
"context"
"errors"
usageDomain "apps/backend/internal/module/usage/domain"
usageUC "apps/backend/internal/module/usage/usecase"
"apps/backend/internal/response"
"apps/backend/internal/svc"
"apps/backend/internal/types"
"github.com/zeromicro/go-zero/core/logx"
)
type GetTenantUsageAnalyticsLogic struct {
logx.Logger
ctx context.Context
svcCtx *svc.ServiceContext
}
func NewGetTenantUsageAnalyticsLogic(ctx context.Context, svcCtx *svc.ServiceContext) *GetTenantUsageAnalyticsLogic {
return &GetTenantUsageAnalyticsLogic{
Logger: logx.WithContext(ctx),
ctx: ctx,
svcCtx: svcCtx,
}
}
func (l *GetTenantUsageAnalyticsLogic) GetTenantUsageAnalytics(req *types.UsageTenantAnalyticsReq) (resp *types.UsageTenantAnalyticsData, err error) {
if l.svcCtx.Usage == nil || l.svcCtx.Members == nil {
return nil, response.Biz(503, 503001, "usage not configured")
}
query, err := usageUC.ParseAnalyticsQuery(req.Granularity, req.From, req.To)
if err != nil {
if errors.Is(err, usageDomain.ErrInvalidRange) {
return nil, response.Biz(400, 400001, err.Error())
}
return nil, err
}
memberList, err := l.svcCtx.Members.ListAllMembers(l.ctx)
if err != nil {
return nil, err
}
members := make([]usageDomain.AnalyticsMember, 0, len(memberList))
for _, member := range memberList {
members = append(members, usageDomain.AnalyticsMember{
UID: member.UID, Email: member.Email, DisplayName: member.DisplayName,
})
}
analytics, err := l.svcCtx.Usage.TenantAnalytics(l.ctx, query, members)
if err != nil {
return nil, err
}
byMeter := make(map[string]types.UsageMeterCount, len(analytics.ByMeter))
for meter, count := range analytics.ByMeter {
byMeter[meter] = types.UsageMeterCount{Count: count.Count, Credits: count.Credits}
}
byokMeter := make(map[string]types.UsageMeterCount, len(analytics.Byok.ByMeter))
for meter, count := range analytics.Byok.ByMeter {
byokMeter[meter] = types.UsageMeterCount{Count: count.Count, Credits: count.Credits}
}
series := make([]types.UsageTenantAnalyticsBucket, 0, len(analytics.Series))
for _, bucket := range analytics.Series {
series = append(series, types.UsageTenantAnalyticsBucket{
Key: bucket.Key, Label: bucket.Label, PurchasedCredits: bucket.PurchasedCredits,
ConsumedCredits: bucket.ConsumedCredits, AiCalls: bucket.AICalls,
SearchCalls: bucket.SearchCalls, ByokCallCount: bucket.ByokCallCount,
})
}
rows := make([]types.UsageTenantAnalyticsMember, 0, len(analytics.Members))
for _, member := range analytics.Members {
rows = append(rows, types.UsageTenantAnalyticsMember{
Uid: member.UID, Email: member.Email, DisplayName: member.DisplayName,
Unlimited: member.Unlimited, PlanId: member.PlanID,
PurchasedCredits: member.PurchasedCredits, TotalCredits: member.TotalCredits,
AiCalls: member.AICalls, SearchCalls: member.SearchCalls,
ByokCallCount: member.ByokCallCount, RemainingCredits: member.RemainingCredits, Pct: member.Pct,
})
}
return &types.UsageTenantAnalyticsData{
Granularity: analytics.Granularity, From: analytics.From, To: analytics.To,
RangeLabel: analytics.RangeLabel, PurchasedCredits: analytics.PurchasedCredits,
ConsumedCredits: analytics.ConsumedCredits, RemainingCredits: analytics.RemainingCredits,
Pct: analytics.Pct, AiCalls: analytics.AICalls, SearchCalls: analytics.SearchCalls,
ByMeter: byMeter, Byok: types.UsageByokBlock{CallCount: analytics.Byok.CallCount, ByMeter: byokMeter},
Series: series, Members: rows,
}, nil
}

View File

@ -62,12 +62,12 @@ func (m *mwRepo) DeleteIdentity(context.Context, string, string) error { return
func (m *mwRepo) UpdateIdentityPassword(context.Context, string, string, string) error { func (m *mwRepo) UpdateIdentityPassword(context.Context, string, string, string) error {
return nil return nil
} }
func (m *mwRepo) SetCode(string, string, string, time.Duration) {} func (m *mwRepo) SetCode(string, string, string, time.Duration) bool { return true }
func (m *mwRepo) CheckCode(string, string, string) bool { return false } func (m *mwRepo) CheckCode(string, string, string) bool { return false }
func (m *mwRepo) PeekCode(string, string, string) bool { return false } func (m *mwRepo) PeekCode(string, string, string) bool { return false }
func (m *mwRepo) SaveRefresh(string, int64, time.Time) {} func (m *mwRepo) SaveRefresh(string, int64, time.Time) {}
func (m *mwRepo) ConsumeRefresh(string) (int64, bool) { return 0, false } func (m *mwRepo) ConsumeRefresh(string) (int64, bool) { return 0, false }
func (m *mwRepo) RevokeRefresh(string) {} func (m *mwRepo) RevokeRefresh(string) {}
func (m *mwRepo) RevokeAllRefreshByUID(context.Context, int64) error { func (m *mwRepo) RevokeAllRefreshByUID(context.Context, int64) error {
return nil return nil
} }

View File

@ -0,0 +1,40 @@
package middleware
import (
"context"
"io"
"net/http"
"apps/backend/internal/response"
)
const stripeWebhookMaxBytes = 1 << 20
type stripeRawBodyKey struct{}
type stripeSignatureKey struct{}
type StripeRawBodyMiddleware struct {
}
func NewStripeRawBodyMiddleware() *StripeRawBodyMiddleware {
return &StripeRawBodyMiddleware{}
}
func (m *StripeRawBodyMiddleware) Handle(next http.HandlerFunc) http.HandlerFunc {
return func(w http.ResponseWriter, r *http.Request) {
body, err := io.ReadAll(http.MaxBytesReader(w, r.Body, stripeWebhookMaxBytes))
if err != nil {
response.Fail(w, http.StatusRequestEntityTooLarge, 413001, "webhook payload too large", nil)
return
}
ctx := context.WithValue(r.Context(), stripeRawBodyKey{}, body)
ctx = context.WithValue(ctx, stripeSignatureKey{}, r.Header.Get("Stripe-Signature"))
next(w, r.WithContext(ctx))
}
}
func StripeWebhookPayload(ctx context.Context) ([]byte, string, bool) {
payload, payloadOK := ctx.Value(stripeRawBodyKey{}).([]byte)
signature, signatureOK := ctx.Value(stripeSignatureKey{}).(string)
return payload, signature, payloadOK && signatureOK
}

View File

@ -0,0 +1,33 @@
package middleware
import (
"bytes"
"net/http"
"net/http/httptest"
"testing"
"github.com/stretchr/testify/require"
)
func TestStripeRawBodyPreservesPayloadAndSignature(t *testing.T) {
payload := []byte("{\n \"exact\": \"bytes\"\n}")
req := httptest.NewRequest(http.MethodPost, "/", bytes.NewReader(payload))
req.Header.Set("Stripe-Signature", "t=1,v1=abc")
rr := httptest.NewRecorder()
NewStripeRawBodyMiddleware().Handle(func(_ http.ResponseWriter, r *http.Request) {
got, sig, ok := StripeWebhookPayload(r.Context())
require.True(t, ok)
require.Equal(t, payload, got)
require.Equal(t, "t=1,v1=abc", sig)
})(rr, req)
require.Equal(t, http.StatusOK, rr.Code)
}
func TestStripeRawBodyRejectsOverOneMiB(t *testing.T) {
req := httptest.NewRequest(http.MethodPost, "/", bytes.NewReader(make([]byte, stripeWebhookMaxBytes+1)))
rr := httptest.NewRecorder()
called := false
NewStripeRawBodyMiddleware().Handle(func(http.ResponseWriter, *http.Request) { called = true })(rr, req)
require.False(t, called)
require.Equal(t, http.StatusRequestEntityTooLarge, rr.Code)
}

View File

@ -0,0 +1,83 @@
package ai
import (
"context"
"encoding/json"
"fmt"
"strings"
"time"
"github.com/zeromicro/go-zero/core/stores/redis"
)
const defaultRedisNamespace = "haixun:dev:v1"
// RedisModelsCache shares model-list entries and refresh locks across processes.
type RedisModelsCache struct {
client *redis.Redis
namespace string
}
func NewRedisModelsCache(conf redis.RedisConf, namespace string) *RedisModelsCache {
return &RedisModelsCache{client: redis.MustNewRedis(conf), namespace: normalizeRedisNamespace(namespace)}
}
func (c *RedisModelsCache) Get(ctx context.Context, key ModelsCacheKey) (ModelsCacheValue, bool, error) {
raw, err := c.client.GetCtx(ctx, c.valueKey(key))
if err != nil {
return ModelsCacheValue{}, false, err
}
if raw == "" {
return ModelsCacheValue{}, false, nil
}
var value ModelsCacheValue
if err := json.Unmarshal([]byte(raw), &value); err != nil {
return ModelsCacheValue{}, false, err
}
return value, true, nil
}
func (c *RedisModelsCache) Set(ctx context.Context, key ModelsCacheKey, value ModelsCacheValue, ttl time.Duration) error {
raw, err := json.Marshal(value)
if err != nil {
return err
}
return c.client.SetexCtx(ctx, c.valueKey(key), string(raw), durationSeconds(ttl))
}
func (c *RedisModelsCache) AcquireRefresh(ctx context.Context, key ModelsCacheKey, token string, ttl time.Duration) (bool, error) {
return c.client.SetnxExCtx(ctx, c.lockKey(key), token, durationSeconds(ttl))
}
func (c *RedisModelsCache) ReleaseRefresh(ctx context.Context, key ModelsCacheKey, token string) error {
_, err := c.client.EvalCtx(ctx, `
if redis.call('GET', KEYS[1]) == ARGV[1] then
return redis.call('DEL', KEYS[1])
end
return 0
`, []string{c.lockKey(key)}, token)
return err
}
func (c *RedisModelsCache) valueKey(key ModelsCacheKey) string {
return fmt.Sprintf("%s:ai:models:{%s:%s}:value", c.namespace, key.Provider, key.CredentialFingerprint)
}
func (c *RedisModelsCache) lockKey(key ModelsCacheKey) string {
return fmt.Sprintf("%s:ai:models:{%s:%s}:lock", c.namespace, key.Provider, key.CredentialFingerprint)
}
func normalizeRedisNamespace(namespace string) string {
if namespace = strings.Trim(strings.TrimSpace(namespace), ":"); namespace != "" {
return namespace
}
return defaultRedisNamespace
}
func durationSeconds(ttl time.Duration) int {
seconds := int(ttl.Seconds())
if seconds < 1 {
return 1
}
return seconds
}

View File

@ -0,0 +1,27 @@
package ai
import (
"strings"
"testing"
)
func TestRedisModelsCacheKeysAreNamespacedAndCoLocated(t *testing.T) {
cache := &RedisModelsCache{namespace: normalizeRedisNamespace("tenant:v2")}
key := ModelsCacheKey{Provider: ProviderXAI, CredentialFingerprint: strings.Repeat("a", 64)}
valueKey := cache.valueKey(key)
lockKey := cache.lockKey(key)
wantTag := "{" + ProviderXAI + ":" + strings.Repeat("a", 64) + "}"
if !strings.HasPrefix(valueKey, "tenant:v2:ai:models:") || !strings.Contains(valueKey, wantTag) {
t.Fatalf("unexpected value key %q", valueKey)
}
if !strings.Contains(lockKey, wantTag) {
t.Fatalf("lock key is not colocated with value key: %q", lockKey)
}
}
func TestRedisModelsCacheDefaultNamespace(t *testing.T) {
if got := normalizeRedisNamespace(""); got != "haixun:dev:v1" {
t.Fatalf("namespace=%q", got)
}
}

View File

@ -2,11 +2,12 @@ package ai
import ( import (
"context" "context"
"crypto/hmac"
"crypto/rand"
"crypto/sha256" "crypto/sha256"
"encoding/hex" "encoding/hex"
"fmt" "fmt"
"strings" "strings"
"sync"
"time" "time"
) )
@ -14,30 +15,61 @@ const (
ProviderXAI = "xai" ProviderXAI = "xai"
ProviderOpenCodeGo = "opencode-go" ProviderOpenCodeGo = "opencode-go"
// ModelsCacheTTL — list models cache window // ModelsCacheTTL is the fresh model-list window.
ModelsCacheTTL = 5 * time.Minute ModelsCacheTTL = 5 * time.Minute
// ModelsCacheStaleTTL is the maximum stale fallback window.
ModelsCacheStaleTTL = 15 * time.Minute
modelsRefreshLockTTL = 30 * time.Second
modelsRefreshWait = 25 * time.Millisecond
modelsRefreshChecks = 20
) )
// Registry holds provider clients + models cache. // ModelsCacheKey identifies one provider credential without containing the credential.
type ModelsCacheKey struct {
Provider string
CredentialFingerprint string
}
// ModelsCacheValue is the JSON payload persisted by a ModelsCache.
type ModelsCacheValue struct {
Models []string `json:"models"`
FetchedAt int64 `json:"fetchedAt"`
}
// ModelsCache stores model lists and coordinates refreshes across instances.
type ModelsCache interface {
Get(ctx context.Context, key ModelsCacheKey) (ModelsCacheValue, bool, error)
Set(ctx context.Context, key ModelsCacheKey, value ModelsCacheValue, ttl time.Duration) error
AcquireRefresh(ctx context.Context, key ModelsCacheKey, token string, ttl time.Duration) (bool, error)
ReleaseRefresh(ctx context.Context, key ModelsCacheKey, token string) error
}
// Registry holds the immutable provider client registry and a shared models cache.
type Registry struct { type Registry struct {
clients map[string]Client clients map[string]Client
mu sync.Mutex cache ModelsCache
cache map[string]modelsCacheEntry fingerprintSecret []byte
} }
type modelsCacheEntry struct { func NewRegistry(cache ModelsCache, fingerprintSecret string) *Registry {
Models []string return newRegistry(cache, fingerprintSecret, map[string]Client{
ExpiresAt time.Time ProviderXAI: NewOpenAICompatible(ProviderXAI, "https://api.x.ai/v1"),
ProviderOpenCodeGo: NewOpenAICompatible(ProviderOpenCodeGo, "https://opencode.ai/zen/go/v1"),
})
} }
func NewRegistry() *Registry { // NewRegistryWithClients injects provider clients for focused tests.
return &Registry{ func NewRegistryWithClients(cache ModelsCache, fingerprintSecret string, clients map[string]Client) *Registry {
clients: map[string]Client{ return newRegistry(cache, fingerprintSecret, clients)
ProviderXAI: NewOpenAICompatible(ProviderXAI, "https://api.x.ai/v1"), }
ProviderOpenCodeGo: NewOpenAICompatible(ProviderOpenCodeGo, "https://opencode.ai/zen/go/v1"),
}, func newRegistry(cache ModelsCache, fingerprintSecret string, clients map[string]Client) *Registry {
cache: map[string]modelsCacheEntry{}, immutableClients := make(map[string]Client, len(clients))
for provider, client := range clients {
immutableClients[NormalizeProvider(provider)] = client
} }
return &Registry{clients: immutableClients, cache: cache, fingerprintSecret: []byte(fingerprintSecret)}
} }
func NormalizeProvider(p string) string { func NormalizeProvider(p string) string {
@ -60,41 +92,105 @@ func (r *Registry) Client(provider string) (Client, error) {
return c, nil return c, nil
} }
// ListModelsCached fetches models via provider API; caches 5 minutes per provider+key. // ListModelsCached fetches models via provider API; caches per provider+credential.
// fromCache indicates whether result was served from cache. // fromCache indicates whether result was served from cache.
func (r *Registry) ListModelsCached(ctx context.Context, provider, apiKey string) (models []string, fromCache bool, err error) { func (r *Registry) ListModelsCached(ctx context.Context, provider, apiKey string) (models []string, fromCache bool, err error) {
id := NormalizeProvider(provider) id := NormalizeProvider(provider)
key := cacheKey(id, apiKey) if _, err := r.Client(id); err != nil {
return nil, false, err
r.mu.Lock() }
if ent, ok := r.cache[key]; ok && time.Now().Before(ent.ExpiresAt) { key := r.cacheKey(id, apiKey)
out := append([]string(nil), ent.Models...) now := time.Now().UTC()
r.mu.Unlock() stale, hasStale := r.cached(ctx, key, now)
return out, true, nil if hasStale && cacheAge(now, stale.FetchedAt) <= ModelsCacheTTL {
return cloneModels(stale.Models), true, nil
}
token, tokenErr := refreshToken()
locked := false
if tokenErr != nil {
err = tokenErr
} else if r.cache != nil {
locked, err = r.cache.AcquireRefresh(ctx, key, token, modelsRefreshLockTTL)
if err != nil {
locked = false
}
}
if locked {
defer func() {
releaseCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), time.Second)
defer cancel()
_ = r.cache.ReleaseRefresh(releaseCtx, key, token)
}()
} else if r.cache != nil && err == nil {
for range modelsRefreshChecks {
select {
case <-ctx.Done():
return nil, false, ctx.Err()
case <-time.After(modelsRefreshWait):
}
if current, ok := r.cached(ctx, key, time.Now().UTC()); ok && cacheAge(time.Now().UTC(), current.FetchedAt) <= ModelsCacheTTL {
return cloneModels(current.Models), true, nil
}
}
} }
r.mu.Unlock()
c, err := r.Client(id) c, err := r.Client(id)
if err != nil {
return nil, false, err
}
list, err := c.ListModels(ctx, apiKey) list, err := c.ListModels(ctx, apiKey)
if err != nil { if err != nil {
if hasStale {
return cloneModels(stale.Models), true, nil
}
return nil, false, err return nil, false, err
} }
r.mu.Lock() if r.cache != nil {
r.cache[key] = modelsCacheEntry{ _ = r.cache.Set(ctx, key, ModelsCacheValue{
Models: append([]string(nil), list...), Models: cloneModels(list),
ExpiresAt: time.Now().Add(ModelsCacheTTL), FetchedAt: time.Now().UTC().UnixNano(),
}, ModelsCacheStaleTTL)
} }
r.mu.Unlock()
return list, false, nil return list, false, nil
} }
func cacheKey(provider, apiKey string) string { func (r *Registry) cached(ctx context.Context, key ModelsCacheKey, now time.Time) (ModelsCacheValue, bool) {
sum := sha256.Sum256([]byte(provider + "\n" + apiKey)) if r.cache == nil {
return hex.EncodeToString(sum[:16]) return ModelsCacheValue{}, false
}
value, ok, err := r.cache.Get(ctx, key)
if err != nil || !ok || value.FetchedAt <= 0 {
return ModelsCacheValue{}, false
}
if cacheAge(now, value.FetchedAt) > ModelsCacheStaleTTL {
return ModelsCacheValue{}, false
}
return value, true
}
func (r *Registry) cacheKey(provider, apiKey string) ModelsCacheKey {
mac := hmac.New(sha256.New, r.fingerprintSecret)
_, _ = mac.Write([]byte(apiKey))
return ModelsCacheKey{Provider: provider, CredentialFingerprint: hex.EncodeToString(mac.Sum(nil))}
}
func refreshToken() (string, error) {
var token [16]byte
if _, err := rand.Read(token[:]); err != nil {
return "", err
}
return hex.EncodeToString(token[:]), nil
}
func cacheAge(now time.Time, fetchedAt int64) time.Duration {
age := now.Sub(time.Unix(0, fetchedAt).UTC())
if age < 0 {
return 0
}
return age
}
func cloneModels(models []string) []string {
return append([]string(nil), models...)
} }
// FallbackModels when no key / provider error (never empty for UI). // FallbackModels when no key / provider error (never empty for UI).

View File

@ -0,0 +1,236 @@
package ai
import (
"context"
"errors"
"fmt"
"strings"
"sync"
"sync/atomic"
"testing"
"time"
)
const testFingerprintSecret = "model-cache-test-fingerprint-secret-32-bytes"
type fakeModelsCache struct {
mu sync.Mutex
values map[ModelsCacheKey]ModelsCacheValue
locks map[ModelsCacheKey]string
seen []ModelsCacheKey
getErr error
lockErr error
}
func newFakeModelsCache() *fakeModelsCache {
return &fakeModelsCache{values: make(map[ModelsCacheKey]ModelsCacheValue), locks: make(map[ModelsCacheKey]string)}
}
func (c *fakeModelsCache) Get(_ context.Context, key ModelsCacheKey) (ModelsCacheValue, bool, error) {
c.mu.Lock()
defer c.mu.Unlock()
c.seen = append(c.seen, key)
if c.getErr != nil {
return ModelsCacheValue{}, false, c.getErr
}
value, ok := c.values[key]
value.Models = cloneModels(value.Models)
return value, ok, nil
}
func (c *fakeModelsCache) Set(_ context.Context, key ModelsCacheKey, value ModelsCacheValue, _ time.Duration) error {
c.mu.Lock()
defer c.mu.Unlock()
value.Models = cloneModels(value.Models)
c.values[key] = value
return nil
}
func (c *fakeModelsCache) AcquireRefresh(_ context.Context, key ModelsCacheKey, token string, _ time.Duration) (bool, error) {
c.mu.Lock()
defer c.mu.Unlock()
if c.lockErr != nil {
return false, c.lockErr
}
if _, exists := c.locks[key]; exists {
return false, nil
}
c.locks[key] = token
return true, nil
}
func (c *fakeModelsCache) ReleaseRefresh(_ context.Context, key ModelsCacheKey, token string) error {
c.mu.Lock()
defer c.mu.Unlock()
if c.locks[key] == token {
delete(c.locks, key)
}
return nil
}
type modelClient struct {
calls atomic.Int32
models []string
err error
started chan struct{}
release chan struct{}
}
func (c *modelClient) Complete(context.Context, string, string, string) (string, error) {
return "", nil
}
func (c *modelClient) CompleteStream(context.Context, string, string, string, func(string) error) (string, error) {
return "", nil
}
func (c *modelClient) ListModels(context.Context, string) ([]string, error) {
c.calls.Add(1)
if c.started != nil {
select {
case c.started <- struct{}{}:
default:
}
}
if c.release != nil {
<-c.release
}
return cloneModels(c.models), c.err
}
func TestRegistryListModelsCacheHit(t *testing.T) {
cache := newFakeModelsCache()
client := &modelClient{models: []string{"provider-model"}}
registry := NewRegistryWithClients(cache, testFingerprintSecret, map[string]Client{ProviderXAI: client})
cache.values[registry.cacheKey(ProviderXAI, "api-key")] = ModelsCacheValue{
Models: []string{"cached-model"}, FetchedAt: time.Now().UTC().UnixNano(),
}
models, fromCache, err := registry.ListModelsCached(context.Background(), ProviderXAI, "api-key")
if err != nil || !fromCache || fmt.Sprint(models) != "[cached-model]" {
t.Fatalf("unexpected result models=%v fromCache=%v err=%v", models, fromCache, err)
}
if client.calls.Load() != 0 {
t.Fatalf("provider called on cache hit")
}
}
func TestRegistrySeparatesProvidersAndCredentials(t *testing.T) {
cache := newFakeModelsCache()
xai := &modelClient{models: []string{"xai-model"}}
opencode := &modelClient{models: []string{"opencode-model"}}
registry := NewRegistryWithClients(cache, testFingerprintSecret, map[string]Client{
ProviderXAI: xai, ProviderOpenCodeGo: opencode,
})
for _, call := range []struct{ provider, key string }{
{ProviderXAI, "key-a"}, {ProviderXAI, "key-b"}, {ProviderOpenCodeGo, "key-a"},
} {
if _, _, err := registry.ListModelsCached(context.Background(), call.provider, call.key); err != nil {
t.Fatal(err)
}
}
cache.mu.Lock()
entries := len(cache.values)
cache.mu.Unlock()
if entries != 3 || xai.calls.Load() != 2 || opencode.calls.Load() != 1 {
t.Fatalf("separation failed entries=%d xai=%d opencode=%d", entries, xai.calls.Load(), opencode.calls.Load())
}
}
func TestRegistryReturnsStaleOnProviderFailure(t *testing.T) {
cache := newFakeModelsCache()
client := &modelClient{err: errors.New("provider unavailable")}
registry := NewRegistryWithClients(cache, testFingerprintSecret, map[string]Client{ProviderXAI: client})
cache.values[registry.cacheKey(ProviderXAI, "api-key")] = ModelsCacheValue{
Models: []string{"stale-model"}, FetchedAt: time.Now().UTC().Add(-10 * time.Minute).UnixNano(),
}
models, fromCache, err := registry.ListModelsCached(context.Background(), ProviderXAI, "api-key")
if err != nil || !fromCache || fmt.Sprint(models) != "[stale-model]" {
t.Fatalf("unexpected stale result models=%v fromCache=%v err=%v", models, fromCache, err)
}
}
func TestRegistryRedisFailureFallsOpenToProvider(t *testing.T) {
cache := newFakeModelsCache()
cache.getErr = errors.New("redis read failed")
cache.lockErr = errors.New("redis lock failed")
client := &modelClient{models: []string{"provider-model"}}
registry := NewRegistryWithClients(cache, testFingerprintSecret, map[string]Client{ProviderXAI: client})
models, fromCache, err := registry.ListModelsCached(context.Background(), ProviderXAI, "api-key")
if err != nil || fromCache || fmt.Sprint(models) != "[provider-model]" {
t.Fatalf("unexpected fail-open result models=%v fromCache=%v err=%v", models, fromCache, err)
}
if client.calls.Load() != 1 {
t.Fatalf("provider calls=%d, want 1", client.calls.Load())
}
}
func TestRegistryCoalescesConcurrentRefresh(t *testing.T) {
cache := newFakeModelsCache()
client := &modelClient{
models: []string{"fresh-model"}, started: make(chan struct{}, 1), release: make(chan struct{}),
}
registry := NewRegistryWithClients(cache, testFingerprintSecret, map[string]Client{ProviderXAI: client})
type result struct {
models []string
err error
}
results := make(chan result, 2)
call := func() {
models, _, err := registry.ListModelsCached(context.Background(), ProviderXAI, "api-key")
results <- result{models: models, err: err}
}
go call()
select {
case <-client.started:
case <-time.After(time.Second):
t.Fatal("provider did not start")
}
go call()
time.Sleep(50 * time.Millisecond)
close(client.release)
for range 2 {
select {
case got := <-results:
if got.err != nil || fmt.Sprint(got.models) != "[fresh-model]" {
t.Fatalf("unexpected result models=%v err=%v", got.models, got.err)
}
case <-time.After(2 * time.Second):
t.Fatal("concurrent refresh timed out")
}
}
if client.calls.Load() != 1 {
t.Fatalf("provider calls=%d, want 1", client.calls.Load())
}
}
func TestRegistryCacheKeyDoesNotExposeAPIKey(t *testing.T) {
cache := newFakeModelsCache()
client := &modelClient{models: []string{"model"}}
registry := NewRegistryWithClients(cache, testFingerprintSecret, map[string]Client{ProviderXAI: client})
apiKey := "sk-plain-secret-must-never-appear"
if _, _, err := registry.ListModelsCached(context.Background(), ProviderXAI, apiKey); err != nil {
t.Fatal(err)
}
cache.mu.Lock()
seen := fmt.Sprintf("%+v %+v", cache.seen, cache.values)
var fingerprint string
for key := range cache.values {
fingerprint = key.CredentialFingerprint
}
cache.mu.Unlock()
if strings.Contains(seen, apiKey) {
t.Fatalf("cache metadata exposed API key: %s", seen)
}
if len(fingerprint) != sha256HexLength || strings.Contains(fingerprint, apiKey) {
t.Fatalf("invalid credential fingerprint %q", fingerprint)
}
}
const sha256HexLength = 64

View File

@ -0,0 +1,148 @@
package billing
import (
"context"
"errors"
)
var (
ErrDisabled = errors.New("billing is not configured")
ErrInvalidPlan = errors.New("invalid billing plan")
ErrInvalidRequest = errors.New("invalid billing request")
ErrNotFound = errors.New("billing record not found")
ErrConflict = errors.New("billing request is already being created")
ErrInvalidSignature = errors.New("invalid stripe signature")
ErrCustomerRequired = errors.New("stripe customer is not available")
)
const (
PlanFree = "free"
PlanStarter = "starter"
PlanPro = "pro"
)
type CheckoutAttempt struct {
ID string `bson:"_id"`
UID int64 `bson:"uid"`
RequestID string `bson:"request_id"`
PlanID string `bson:"plan_id"`
PriceID string `bson:"price_id"`
SessionID string `bson:"session_id,omitempty"`
URL string `bson:"url,omitempty"`
ExpiresAt int64 `bson:"expires_at,omitempty"`
Status string `bson:"status"`
PaymentStatus string `bson:"payment_status,omitempty"`
FulfillmentStatus string `bson:"fulfillment_status"`
CustomerID string `bson:"customer_id,omitempty"`
SubscriptionID string `bson:"subscription_id,omitempty"`
Error string `bson:"error,omitempty"`
CreatedAt int64 `bson:"created_at"`
UpdatedAt int64 `bson:"updated_at"`
}
type SubscriptionState struct {
UID int64
PlanID string
PaidPlanID string
Status string
CustomerID string
SubscriptionID string
CurrentPeriodStart int64
CurrentPeriodEnd int64
CancelAtPeriodEnd bool
AdminOverride bool
}
type ProviderCheckout struct {
SessionID string
URL string
ExpiresAt int64
Status string
PaymentStatus string
PriceID string
CustomerID string
SubscriptionID string
}
type CheckoutParams struct {
UID int64
AttemptID string
PlanID string
PriceID string
CustomerID string
SuccessURL string
CancelURL string
IdempotencyKey string
}
type ProviderSubscription struct {
ID string
CustomerID string
Status string
PriceID string
CurrentPeriodStart int64
CurrentPeriodEnd int64
CancelAtPeriodEnd bool
}
type Event struct {
ID string
Type string
Created int64
ObjectID string
SubscriptionID string
CustomerID string
PriceID string
Status string
PeriodStart int64
PeriodEnd int64
CancelAtPeriodEnd bool
}
type EntitlementUpdate struct {
EventID string
EventCreatedAt int64
UID int64
PlanID string
PaidPlanID string
Status string
CustomerID string
SubscriptionID string
CurrentPeriodStart int64
CurrentPeriodEnd int64
CancelAtPeriodEnd bool
}
type PurchaseAudit struct {
ID string `bson:"_id"`
UID int64 `bson:"uid"`
PlanID string `bson:"plan_id"`
StripeEventID string `bson:"stripe_event_id"`
StripeSessionID string `bson:"stripe_session_id,omitempty"`
StripeCustomerID string `bson:"stripe_customer_id,omitempty"`
StripeSubscriptionID string `bson:"stripe_subscription_id,omitempty"`
Status string `bson:"status"`
CreatedAt int64 `bson:"created_at"`
}
type Provider interface {
CreateCheckout(context.Context, CheckoutParams) (*ProviderCheckout, error)
GetCheckout(context.Context, string) (*ProviderCheckout, error)
GetSubscription(context.Context, string) (*ProviderSubscription, error)
CreatePortal(context.Context, string, string) (string, error)
VerifyEvent([]byte, string, string) (*Event, error)
}
type Repository interface {
InsertCheckout(context.Context, *CheckoutAttempt) (bool, error)
GetCheckoutByRequest(context.Context, int64, string) (*CheckoutAttempt, error)
GetCheckout(context.Context, int64, string) (*CheckoutAttempt, error)
GetCheckoutBySession(context.Context, string) (*CheckoutAttempt, error)
UpdateCheckout(context.Context, *CheckoutAttempt) error
GetSubscription(context.Context, int64) (*SubscriptionState, error)
FindUIDByCustomer(context.Context, string) (int64, error)
ApplyEntitlement(context.Context, EntitlementUpdate) (bool, error)
WebhookProcessed(context.Context, string) (bool, error)
MarkWebhookProcessed(context.Context, string, string, int64) error
UpsertPurchase(context.Context, *PurchaseAudit) error
}

View File

@ -0,0 +1,144 @@
package billing
import (
"context"
"time"
libmongo "apps/backend/internal/lib/mongo"
usageDomain "apps/backend/internal/module/usage/domain"
"github.com/zeromicro/go-zero/core/stores/mon"
"go.mongodb.org/mongo-driver/bson"
"go.mongodb.org/mongo-driver/mongo"
"go.mongodb.org/mongo-driver/mongo/options"
)
type MongoRepository struct {
attempts *mon.Model
prefs *mon.Model
purchases *mon.Model
webhooks *mon.Model
}
func NewMongoRepository(uri, database string) *MongoRepository {
uri = libmongo.MustMongoURI(uri)
return &MongoRepository{
attempts: mon.MustNewModel(uri, database, "billing_checkout_attempts"),
prefs: mon.MustNewModel(uri, database, "usage_prefs"),
purchases: mon.MustNewModel(uri, database, "usage_purchases"),
webhooks: mon.MustNewModel(uri, database, "billing_webhook_events"),
}
}
func (r *MongoRepository) InsertCheckout(ctx context.Context, a *CheckoutAttempt) (bool, error) {
_, err := r.attempts.InsertOne(ctx, a)
if mongo.IsDuplicateKeyError(err) {
return false, nil
}
return err == nil, err
}
func (r *MongoRepository) GetCheckoutByRequest(ctx context.Context, uid int64, requestID string) (*CheckoutAttempt, error) {
return r.findAttempt(ctx, bson.M{"uid": uid, "request_id": requestID})
}
func (r *MongoRepository) GetCheckout(ctx context.Context, uid int64, id string) (*CheckoutAttempt, error) {
return r.findAttempt(ctx, bson.M{"_id": id, "uid": uid})
}
func (r *MongoRepository) GetCheckoutBySession(ctx context.Context, id string) (*CheckoutAttempt, error) {
return r.findAttempt(ctx, bson.M{"session_id": id})
}
func (r *MongoRepository) findAttempt(ctx context.Context, filter bson.M) (*CheckoutAttempt, error) {
var a CheckoutAttempt
if err := r.attempts.FindOne(ctx, &a, filter); err != nil {
if err == mon.ErrNotFound {
return nil, ErrNotFound
}
return nil, err
}
return &a, nil
}
func (r *MongoRepository) UpdateCheckout(ctx context.Context, a *CheckoutAttempt) error {
_, err := r.attempts.ReplaceOne(ctx, bson.M{"_id": a.ID}, a)
return err
}
func (r *MongoRepository) GetSubscription(ctx context.Context, uid int64) (*SubscriptionState, error) {
var p usageDomain.MemberPrefs
if err := r.prefs.FindOne(ctx, &p, bson.M{"uid": uid}); err != nil {
if err == mon.ErrNotFound {
return &SubscriptionState{UID: uid, PlanID: PlanFree, Status: "none"}, nil
}
return nil, err
}
return &SubscriptionState{UID: uid, PlanID: p.PlanID, PaidPlanID: p.PaidPlanID, Status: p.SubscriptionStatus, CustomerID: p.StripeCustomerID, SubscriptionID: p.StripeSubscriptionID, CurrentPeriodStart: p.CurrentPeriodStart, CurrentPeriodEnd: p.CurrentPeriodEnd, CancelAtPeriodEnd: p.CancelAtPeriodEnd, AdminOverride: p.AdminOverride}, nil
}
func (r *MongoRepository) FindUIDByCustomer(ctx context.Context, customerID string) (int64, error) {
var p usageDomain.MemberPrefs
if customerID == "" {
return 0, ErrNotFound
}
if err := r.prefs.FindOne(ctx, &p, bson.M{"stripe_customer_id": customerID}); err != nil {
if err == mon.ErrNotFound {
a, attemptErr := r.findAttempt(ctx, bson.M{"customer_id": customerID})
if attemptErr != nil {
return 0, ErrNotFound
}
return a.UID, nil
}
return 0, err
}
return p.UID, nil
}
func (r *MongoRepository) ApplyEntitlement(ctx context.Context, u EntitlementUpdate) (bool, error) {
filter := bson.M{"uid": u.UID, "$or": bson.A{
bson.M{"billing_event_created_at": bson.M{"$exists": false}},
bson.M{"billing_event_created_at": bson.M{"$lt": u.EventCreatedAt}},
bson.M{"billing_event_created_at": u.EventCreatedAt, "billing_event_id": bson.M{"$lt": u.EventID}},
}}
pipeline := mongo.Pipeline{bson.D{{Key: "$set", Value: bson.D{
{Key: "uid", Value: u.UID},
{Key: "paid_plan_id", Value: u.PaidPlanID},
{Key: "subscription_status", Value: u.Status},
{Key: "stripe_customer_id", Value: u.CustomerID},
{Key: "stripe_subscription_id", Value: u.SubscriptionID},
{Key: "current_period_start", Value: u.CurrentPeriodStart},
{Key: "current_period_end", Value: u.CurrentPeriodEnd},
{Key: "cancel_at_period_end", Value: u.CancelAtPeriodEnd},
{Key: "billing_event_id", Value: u.EventID},
{Key: "billing_event_created_at", Value: u.EventCreatedAt},
{Key: "updated_at", Value: time.Now().UTC().UnixNano()},
{Key: "plan_id", Value: bson.D{{Key: "$cond", Value: bson.A{
bson.D{{Key: "$eq", Value: bson.A{bson.D{{Key: "$ifNull", Value: bson.A{"$admin_override", false}}}, true}}},
bson.D{{Key: "$ifNull", Value: bson.A{"$plan_id", PlanFree}}}, u.PlanID,
}}}},
}}}}
res, err := r.prefs.UpdateOne(ctx, filter, pipeline, options.Update().SetUpsert(true))
if err != nil {
if mongo.IsDuplicateKeyError(err) {
return false, nil
}
return false, err
}
return res.MatchedCount > 0 || res.UpsertedCount > 0, nil
}
func (r *MongoRepository) WebhookProcessed(ctx context.Context, id string) (bool, error) {
var out bson.M
err := r.webhooks.FindOne(ctx, &out, bson.M{"_id": id})
if err == mon.ErrNotFound {
return false, nil
}
return err == nil, err
}
func (r *MongoRepository) MarkWebhookProcessed(ctx context.Context, id, eventType string, created int64) error {
_, err := r.webhooks.ReplaceOne(ctx, bson.M{"_id": id}, bson.M{"_id": id, "stripe_event_id": id, "type": eventType, "stripe_created_at": created * int64(time.Second), "processed_at": time.Now().UTC().UnixNano()}, options.Replace().SetUpsert(true))
return err
}
func (r *MongoRepository) UpsertPurchase(ctx context.Context, p *PurchaseAudit) error {
_, err := r.purchases.ReplaceOne(ctx, bson.M{"stripe_event_id": p.StripeEventID}, p, options.Replace().SetUpsert(true))
return err
}
var _ Repository = (*MongoRepository)(nil)

View File

@ -0,0 +1,315 @@
package billing
import (
"context"
"fmt"
"net/url"
"strconv"
"strings"
"time"
"github.com/google/uuid"
)
type Service struct {
Repo Repository
Provider Provider
Enabled bool
WebhookSecret string
StarterPriceID string
ProPriceID string
SuccessURL string
CancelURL string
PortalReturnURL string
Now func() time.Time
}
func (s *Service) available() error {
if s == nil || !s.Enabled || s.Repo == nil || s.Provider == nil || strings.TrimSpace(s.StarterPriceID) == "" || strings.TrimSpace(s.ProPriceID) == "" || strings.TrimSpace(s.SuccessURL) == "" || strings.TrimSpace(s.CancelURL) == "" || strings.TrimSpace(s.PortalReturnURL) == "" {
return ErrDisabled
}
return nil
}
func (s *Service) PriceForPlan(plan string) (string, error) {
switch strings.ToLower(strings.TrimSpace(plan)) {
case PlanStarter:
return s.StarterPriceID, nil
case PlanPro:
return s.ProPriceID, nil
default:
return "", ErrInvalidPlan
}
}
func (s *Service) PlanForPrice(price string) (string, error) {
switch price {
case s.StarterPriceID:
return PlanStarter, nil
case s.ProPriceID:
return PlanPro, nil
default:
return "", ErrInvalidPlan
}
}
func (s *Service) CreateCheckout(ctx context.Context, uid int64, planID, requestID string) (*CheckoutAttempt, error) {
if err := s.available(); err != nil {
return nil, err
}
requestID = strings.TrimSpace(requestID)
if uid <= 0 || requestID == "" || len(requestID) > 128 {
return nil, ErrInvalidRequest
}
priceID, err := s.PriceForPlan(planID)
if err != nil {
return nil, err
}
now := time.Now().UTC()
if s.Now != nil {
now = s.Now().UTC()
}
a := &CheckoutAttempt{ID: uuid.NewString(), UID: uid, RequestID: requestID, PlanID: strings.ToLower(strings.TrimSpace(planID)), PriceID: priceID, Status: "creating", FulfillmentStatus: "pending", CreatedAt: now.UnixNano(), UpdatedAt: now.UnixNano()}
inserted, err := s.Repo.InsertCheckout(ctx, a)
if err != nil {
return nil, err
}
if !inserted {
existing, err := s.Repo.GetCheckoutByRequest(ctx, uid, requestID)
if err != nil {
return nil, err
}
if existing.PlanID != a.PlanID {
return nil, ErrConflict
}
if existing.SessionID != "" {
return existing, nil
}
a = existing
}
idempotencyKey := fmt.Sprintf("billing-checkout-%d-%s", uid, requestID)
state, err := s.Repo.GetSubscription(ctx, uid)
if err != nil {
return nil, err
}
successURL := replaceURLTokens(s.SuccessURL, a.ID, a.PlanID)
cancelURL := replaceURLTokens(s.CancelURL, a.ID, a.PlanID)
co, err := s.Provider.CreateCheckout(ctx, CheckoutParams{UID: uid, AttemptID: a.ID, PlanID: a.PlanID, PriceID: priceID, CustomerID: state.CustomerID, SuccessURL: successURL, CancelURL: cancelURL, IdempotencyKey: idempotencyKey})
if err != nil {
a.Status, a.Error, a.UpdatedAt = "failed", err.Error(), time.Now().UTC().UnixNano()
_ = s.Repo.UpdateCheckout(ctx, a)
return nil, err
}
s.applyCheckout(a, co)
if err := s.Repo.UpdateCheckout(ctx, a); err != nil {
return nil, err
}
return a, nil
}
func (s *Service) GetCheckout(ctx context.Context, uid int64, id string) (*CheckoutAttempt, error) {
if err := s.available(); err != nil {
return nil, err
}
a, err := s.Repo.GetCheckout(ctx, uid, id)
if err != nil {
return nil, err
}
if a.SessionID == "" {
return a, nil
}
co, err := s.Provider.GetCheckout(ctx, a.SessionID)
if err != nil {
return nil, err
}
s.applyCheckout(a, co)
if co.PriceID != "" {
plan, mapErr := s.PlanForPrice(co.PriceID)
if mapErr != nil {
return nil, mapErr
}
a.PlanID = plan
}
_ = s.Repo.UpdateCheckout(ctx, a)
return a, nil
}
func (s *Service) applyCheckout(a *CheckoutAttempt, co *ProviderCheckout) {
a.SessionID, a.URL, a.ExpiresAt = co.SessionID, co.URL, co.ExpiresAt
a.Status, a.PaymentStatus = co.Status, co.PaymentStatus
a.CustomerID, a.SubscriptionID = co.CustomerID, co.SubscriptionID
if co.PriceID != "" {
a.PriceID = co.PriceID
}
a.UpdatedAt = time.Now().UTC().UnixNano()
}
func (s *Service) Subscription(ctx context.Context, uid int64) (*SubscriptionState, error) {
if s == nil || s.Repo == nil {
return nil, ErrDisabled
}
return s.Repo.GetSubscription(ctx, uid)
}
func (s *Service) Portal(ctx context.Context, uid int64) (string, error) {
if err := s.available(); err != nil {
return "", err
}
state, err := s.Repo.GetSubscription(ctx, uid)
if err != nil {
return "", err
}
if state.CustomerID == "" {
return "", ErrCustomerRequired
}
return s.Provider.CreatePortal(ctx, state.CustomerID, s.PortalReturnURL)
}
func (s *Service) HandleWebhook(ctx context.Context, payload []byte, signature string) error {
if err := s.available(); err != nil {
return err
}
if strings.TrimSpace(s.WebhookSecret) == "" {
return ErrDisabled
}
e, err := s.Provider.VerifyEvent(payload, signature, s.WebhookSecret)
if err != nil {
return err
}
done, err := s.Repo.WebhookProcessed(ctx, e.ID)
if err != nil || done {
return err
}
if err := s.processEvent(ctx, e); err != nil {
return err
}
return s.Repo.MarkWebhookProcessed(ctx, e.ID, e.Type, e.Created)
}
func (s *Service) processEvent(ctx context.Context, e *Event) error {
switch e.Type {
case "checkout.session.expired", "checkout.session.async_payment_failed":
a, err := s.Repo.GetCheckoutBySession(ctx, e.ObjectID)
if err != nil {
return err
}
a.Status = strings.TrimPrefix(e.Type, "checkout.session.")
a.FulfillmentStatus = "failed"
a.UpdatedAt = time.Now().UTC().UnixNano()
return s.Repo.UpdateCheckout(ctx, a)
case "checkout.session.completed", "checkout.session.async_payment_succeeded":
return s.fulfillCheckout(ctx, e)
case "customer.subscription.updated", "customer.subscription.deleted":
return s.applySubscriptionEvent(ctx, e)
case "invoice.paid":
return s.applyInvoiceEvent(ctx, e)
default:
return nil
}
}
func (s *Service) fulfillCheckout(ctx context.Context, e *Event) error {
a, err := s.Repo.GetCheckoutBySession(ctx, e.ObjectID)
if err != nil {
return err
}
co, err := s.Provider.GetCheckout(ctx, e.ObjectID)
if err != nil {
return err
}
plan, err := s.PlanForPrice(co.PriceID)
if err != nil {
return err
}
if co.PaymentStatus != "paid" && co.PaymentStatus != "no_payment_required" {
s.applyCheckout(a, co)
a.FulfillmentStatus = "pending"
return s.Repo.UpdateCheckout(ctx, a)
}
sub, err := s.Provider.GetSubscription(ctx, co.SubscriptionID)
if err != nil {
return err
}
if sub.PriceID != co.PriceID {
return ErrInvalidPlan
}
s.applyCheckout(a, co)
if err := s.apply(ctx, e, a.UID, plan, sub, a.SessionID); err != nil {
return err
}
a.PlanID, a.FulfillmentStatus = plan, "fulfilled"
return s.Repo.UpdateCheckout(ctx, a)
}
func (s *Service) applySubscriptionEvent(ctx context.Context, e *Event) error {
uid, err := s.Repo.FindUIDByCustomer(ctx, e.CustomerID)
if err != nil {
return err
}
plan := PlanFree
if e.Type != "customer.subscription.deleted" {
plan, err = s.PlanForPrice(e.PriceID)
if err != nil {
return err
}
} else if e.PriceID != "" {
if mapped, mapErr := s.PlanForPrice(e.PriceID); mapErr == nil {
plan = mapped
}
} else if state, stateErr := s.Repo.GetSubscription(ctx, uid); stateErr == nil && state.PaidPlanID != "" {
plan = state.PaidPlanID
}
sub := &ProviderSubscription{ID: e.ObjectID, CustomerID: e.CustomerID, Status: e.Status, PriceID: e.PriceID, CurrentPeriodStart: e.PeriodStart, CurrentPeriodEnd: e.PeriodEnd, CancelAtPeriodEnd: e.CancelAtPeriodEnd}
if e.Type == "customer.subscription.deleted" {
sub.Status = "canceled"
}
return s.apply(ctx, e, uid, plan, sub, "")
}
func (s *Service) applyInvoiceEvent(ctx context.Context, e *Event) error {
uid, err := s.Repo.FindUIDByCustomer(ctx, e.CustomerID)
if err != nil {
return err
}
subscriptionID := e.SubscriptionID
if subscriptionID == "" {
state, stateErr := s.Repo.GetSubscription(ctx, uid)
if stateErr != nil {
return stateErr
}
subscriptionID = state.SubscriptionID
}
sub, err := s.Provider.GetSubscription(ctx, subscriptionID)
if err != nil {
return err
}
plan, err := s.PlanForPrice(sub.PriceID)
if err != nil {
return err
}
if sub.CustomerID == "" {
sub.CustomerID = e.CustomerID
}
sub.Status = "active"
return s.apply(ctx, e, uid, plan, sub, "")
}
func (s *Service) apply(ctx context.Context, e *Event, uid int64, plan string, sub *ProviderSubscription, sessionID string) error {
activePlan := plan
if sub.Status != "active" && sub.Status != "trialing" {
activePlan = PlanFree
}
_, err := s.Repo.ApplyEntitlement(ctx, EntitlementUpdate{EventID: e.ID, EventCreatedAt: e.Created * int64(time.Second), UID: uid, PlanID: activePlan, PaidPlanID: plan, Status: sub.Status, CustomerID: sub.CustomerID, SubscriptionID: sub.ID, CurrentPeriodStart: sub.CurrentPeriodStart * int64(time.Second), CurrentPeriodEnd: sub.CurrentPeriodEnd * int64(time.Second), CancelAtPeriodEnd: sub.CancelAtPeriodEnd})
if err != nil {
return err
}
// Audit upsert must still run when entitlement was already applied.
return s.Repo.UpsertPurchase(ctx, &PurchaseAudit{ID: e.ID, UID: uid, PlanID: plan, StripeEventID: e.ID, StripeSessionID: sessionID, StripeCustomerID: sub.CustomerID, StripeSubscriptionID: sub.ID, Status: sub.Status, CreatedAt: e.Created * int64(time.Second)})
}
func ParseUID(v string) (int64, error) { return strconv.ParseInt(v, 10, 64) }
func replaceURLTokens(value, checkoutID, planID string) string {
value = strings.ReplaceAll(value, "{CHECKOUT_ID}", url.QueryEscape(checkoutID))
return strings.ReplaceAll(value, "{PLAN_ID}", url.QueryEscape(planID))
}

View File

@ -0,0 +1,311 @@
package billing
import (
"context"
"errors"
"sync"
"testing"
"github.com/stretchr/testify/require"
)
type fakeProvider struct {
mu sync.Mutex
creates int
lastCheckout CheckoutParams
event *Event
checkout ProviderCheckout
subscription ProviderSubscription
}
func (p *fakeProvider) CreateCheckout(_ context.Context, input CheckoutParams) (*ProviderCheckout, error) {
p.mu.Lock()
defer p.mu.Unlock()
p.creates++
p.lastCheckout = input
out := p.checkout
out.PriceID = input.PriceID
return &out, nil
}
func (p *fakeProvider) GetCheckout(context.Context, string) (*ProviderCheckout, error) {
out := p.checkout
return &out, nil
}
func (p *fakeProvider) GetSubscription(context.Context, string) (*ProviderSubscription, error) {
out := p.subscription
return &out, nil
}
func (p *fakeProvider) CreatePortal(context.Context, string, string) (string, error) {
return "https://portal.test", nil
}
func (p *fakeProvider) VerifyEvent([]byte, string, string) (*Event, error) {
if p.event == nil {
return nil, ErrInvalidSignature
}
out := *p.event
return &out, nil
}
type fakeRepo struct {
mu sync.Mutex
attempts map[string]*CheckoutAttempt
requests map[string]string
processed map[string]bool
state SubscriptionState
lastCreated int64
lastEvent string
purchases map[string]*PurchaseAudit
entitlementErr error
}
func newFakeRepo() *fakeRepo {
return &fakeRepo{attempts: map[string]*CheckoutAttempt{}, requests: map[string]string{}, processed: map[string]bool{}, purchases: map[string]*PurchaseAudit{}, state: SubscriptionState{PlanID: PlanFree}}
}
func (r *fakeRepo) InsertCheckout(_ context.Context, a *CheckoutAttempt) (bool, error) {
r.mu.Lock()
defer r.mu.Unlock()
key := formatUID(a.UID) + ":" + a.RequestID
if _, ok := r.requests[key]; ok {
return false, nil
}
cp := *a
r.attempts[a.ID] = &cp
r.requests[key] = a.ID
return true, nil
}
func (r *fakeRepo) GetCheckoutByRequest(_ context.Context, uid int64, request string) (*CheckoutAttempt, error) {
r.mu.Lock()
defer r.mu.Unlock()
return r.copyAttempt(r.requests[formatUID(uid)+":"+request])
}
func (r *fakeRepo) GetCheckout(_ context.Context, uid int64, id string) (*CheckoutAttempt, error) {
r.mu.Lock()
defer r.mu.Unlock()
a, err := r.copyAttempt(id)
if err != nil || a.UID != uid {
return nil, ErrNotFound
}
return a, nil
}
func (r *fakeRepo) GetCheckoutBySession(_ context.Context, session string) (*CheckoutAttempt, error) {
r.mu.Lock()
defer r.mu.Unlock()
for _, a := range r.attempts {
if a.SessionID == session {
cp := *a
return &cp, nil
}
}
return nil, ErrNotFound
}
func (r *fakeRepo) copyAttempt(id string) (*CheckoutAttempt, error) {
a := r.attempts[id]
if a == nil {
return nil, ErrNotFound
}
cp := *a
return &cp, nil
}
func (r *fakeRepo) UpdateCheckout(_ context.Context, a *CheckoutAttempt) error {
r.mu.Lock()
defer r.mu.Unlock()
cp := *a
r.attempts[a.ID] = &cp
return nil
}
func (r *fakeRepo) GetSubscription(context.Context, int64) (*SubscriptionState, error) {
r.mu.Lock()
defer r.mu.Unlock()
cp := r.state
return &cp, nil
}
func (r *fakeRepo) FindUIDByCustomer(_ context.Context, customer string) (int64, error) {
r.mu.Lock()
defer r.mu.Unlock()
if r.state.CustomerID == customer {
return r.state.UID, nil
}
for _, a := range r.attempts {
if a.CustomerID == customer {
return a.UID, nil
}
}
return 0, ErrNotFound
}
func (r *fakeRepo) ApplyEntitlement(_ context.Context, u EntitlementUpdate) (bool, error) {
r.mu.Lock()
defer r.mu.Unlock()
if r.entitlementErr != nil {
return false, r.entitlementErr
}
if u.EventCreatedAt < r.lastCreated || (u.EventCreatedAt == r.lastCreated && u.EventID <= r.lastEvent) {
return false, nil
}
r.lastCreated, r.lastEvent = u.EventCreatedAt, u.EventID
r.state.UID, r.state.Status, r.state.CustomerID, r.state.SubscriptionID = u.UID, u.Status, u.CustomerID, u.SubscriptionID
r.state.PaidPlanID = u.PaidPlanID
r.state.CurrentPeriodStart, r.state.CurrentPeriodEnd, r.state.CancelAtPeriodEnd = u.CurrentPeriodStart, u.CurrentPeriodEnd, u.CancelAtPeriodEnd
if !r.state.AdminOverride {
r.state.PlanID = u.PlanID
}
return true, nil
}
func (r *fakeRepo) WebhookProcessed(_ context.Context, id string) (bool, error) {
r.mu.Lock()
defer r.mu.Unlock()
return r.processed[id], nil
}
func (r *fakeRepo) MarkWebhookProcessed(_ context.Context, id, _ string, _ int64) error {
r.mu.Lock()
defer r.mu.Unlock()
r.processed[id] = true
return nil
}
func (r *fakeRepo) UpsertPurchase(_ context.Context, p *PurchaseAudit) error {
r.mu.Lock()
defer r.mu.Unlock()
cp := *p
r.purchases[p.StripeEventID] = &cp
return nil
}
func newService(repo Repository, provider Provider) *Service {
return &Service{Repo: repo, Provider: provider, Enabled: true, WebhookSecret: "whsec_test", StarterPriceID: "price_starter", ProPriceID: "price_pro", SuccessURL: "https://app/success", CancelURL: "https://app/cancel", PortalReturnURL: "https://app/billing"}
}
func TestPlanPriceMapping(t *testing.T) {
s := newService(newFakeRepo(), &fakeProvider{})
price, err := s.PriceForPlan(" STARTER ")
require.NoError(t, err)
require.Equal(t, "price_starter", price)
plan, err := s.PlanForPrice("price_pro")
require.NoError(t, err)
require.Equal(t, PlanPro, plan)
_, err = s.PriceForPlan(PlanFree)
require.ErrorIs(t, err, ErrInvalidPlan)
_, err = s.PlanForPrice("price_spoofed")
require.ErrorIs(t, err, ErrInvalidPlan)
}
func TestCreateCheckoutIsIdempotent(t *testing.T) {
r := newFakeRepo()
p := &fakeProvider{checkout: ProviderCheckout{SessionID: "cs_1", URL: "https://checkout", ExpiresAt: 123}}
s := newService(r, p)
a, err := s.CreateCheckout(context.Background(), 7, PlanStarter, "request-1")
require.NoError(t, err)
b, err := s.CreateCheckout(context.Background(), 7, PlanStarter, "request-1")
require.NoError(t, err)
require.Equal(t, a.ID, b.ID)
require.Equal(t, 1, p.creates)
_, err = s.CreateCheckout(context.Background(), 7, PlanPro, "request-1")
require.ErrorIs(t, err, ErrConflict)
}
func TestCreateCheckoutInjectsReturnURLTokens(t *testing.T) {
r := newFakeRepo()
p := &fakeProvider{checkout: ProviderCheckout{SessionID: "cs_1", URL: "https://checkout"}}
s := newService(r, p)
s.SuccessURL = "https://app.test/app/usage/checkout?result=success&checkout_id={CHECKOUT_ID}&plan={PLAN_ID}"
s.CancelURL = "https://app.test/app/usage/checkout?result=cancel&plan={PLAN_ID}"
a, err := s.CreateCheckout(context.Background(), 7, PlanStarter, "request with spaces")
require.NoError(t, err)
require.Contains(t, p.lastCheckout.SuccessURL, "checkout_id="+a.ID)
require.Contains(t, p.lastCheckout.SuccessURL, "plan=starter")
require.NotContains(t, p.lastCheckout.SuccessURL, "{")
require.Equal(t, "https://app.test/app/usage/checkout?result=cancel&plan=starter", p.lastCheckout.CancelURL)
}
func TestSubscriptionRemainsReadableWhenStripeDisabled(t *testing.T) {
r := newFakeRepo()
r.state = SubscriptionState{UID: 7, PlanID: PlanFree}
s := newService(r, &fakeProvider{})
s.Enabled = false
state, err := s.Subscription(context.Background(), 7)
require.NoError(t, err)
require.Equal(t, PlanFree, state.PlanID)
}
func TestWebhookDuplicateOutOfOrderDowngradeAndAdminOverride(t *testing.T) {
r := newFakeRepo()
r.state = SubscriptionState{UID: 7, PlanID: PlanFree, CustomerID: "cus_1"}
p := &fakeProvider{}
s := newService(r, p)
p.event = &Event{ID: "evt_new", Type: "customer.subscription.updated", Created: 20, ObjectID: "sub_1", CustomerID: "cus_1", PriceID: "price_pro", Status: "active"}
require.NoError(t, s.HandleWebhook(context.Background(), []byte("new"), "sig"))
require.Equal(t, PlanPro, r.state.PlanID)
require.NoError(t, s.HandleWebhook(context.Background(), []byte("duplicate"), "sig"))
require.Len(t, r.purchases, 1)
p.event = &Event{ID: "evt_old", Type: "customer.subscription.deleted", Created: 10, ObjectID: "sub_1", CustomerID: "cus_1", Status: "canceled"}
require.NoError(t, s.HandleWebhook(context.Background(), []byte("old"), "sig"))
require.Equal(t, PlanPro, r.state.PlanID)
p.event = &Event{ID: "evt_delete", Type: "customer.subscription.deleted", Created: 30, ObjectID: "sub_1", CustomerID: "cus_1", Status: "canceled"}
require.NoError(t, s.HandleWebhook(context.Background(), []byte("delete"), "sig"))
require.Equal(t, PlanFree, r.state.PlanID)
r.state.AdminOverride, r.state.PlanID = true, PlanStarter
p.event = &Event{ID: "evt_admin", Type: "customer.subscription.updated", Created: 40, ObjectID: "sub_1", CustomerID: "cus_1", PriceID: "price_pro", Status: "active"}
require.NoError(t, s.HandleWebhook(context.Background(), []byte("admin"), "sig"))
require.Equal(t, PlanStarter, r.state.PlanID)
require.Equal(t, "active", r.state.Status)
}
func TestPaidCheckoutActivationUsesRetrievedPrice(t *testing.T) {
r := newFakeRepo()
a := &CheckoutAttempt{ID: "attempt_1", UID: 9, RequestID: "r1", PlanID: PlanStarter, PriceID: "price_starter", SessionID: "cs_1", FulfillmentStatus: "pending"}
r.attempts[a.ID] = a
p := &fakeProvider{checkout: ProviderCheckout{SessionID: "cs_1", PaymentStatus: "paid", Status: "complete", PriceID: "price_pro", CustomerID: "cus_9", SubscriptionID: "sub_9"}, subscription: ProviderSubscription{ID: "sub_9", CustomerID: "cus_9", Status: "active", PriceID: "price_pro", CurrentPeriodStart: 10, CurrentPeriodEnd: 20}}
s := newService(r, p)
p.event = &Event{ID: "evt_checkout", Type: "checkout.session.completed", Created: 30, ObjectID: "cs_1"}
require.NoError(t, s.HandleWebhook(context.Background(), []byte("paid"), "sig"))
require.Equal(t, PlanPro, r.state.PlanID)
require.Equal(t, "fulfilled", r.attempts[a.ID].FulfillmentStatus)
}
func TestPaidCheckoutIsNotFulfilledWhenEntitlementWriteFails(t *testing.T) {
r := newFakeRepo()
r.entitlementErr = errors.New("entitlement unavailable")
a := &CheckoutAttempt{ID: "attempt_1", UID: 9, SessionID: "cs_1", FulfillmentStatus: "pending"}
r.attempts[a.ID] = a
p := &fakeProvider{checkout: ProviderCheckout{SessionID: "cs_1", PaymentStatus: "paid", Status: "complete", PriceID: "price_pro", CustomerID: "cus_9", SubscriptionID: "sub_9"}, subscription: ProviderSubscription{ID: "sub_9", CustomerID: "cus_9", Status: "active", PriceID: "price_pro"}}
s := newService(r, p)
p.event = &Event{ID: "evt_checkout", Type: "checkout.session.completed", Created: 30, ObjectID: "cs_1"}
err := s.HandleWebhook(context.Background(), []byte("paid"), "sig")
require.ErrorContains(t, err, "entitlement unavailable")
require.Equal(t, "pending", r.attempts[a.ID].FulfillmentStatus)
require.False(t, r.processed["evt_checkout"])
}
func TestUnpaidCompletedCheckoutWaitsForAsyncSuccess(t *testing.T) {
r := newFakeRepo()
a := &CheckoutAttempt{ID: "attempt_1", UID: 9, SessionID: "cs_1", FulfillmentStatus: "pending"}
r.attempts[a.ID] = a
p := &fakeProvider{checkout: ProviderCheckout{SessionID: "cs_1", PaymentStatus: "unpaid", Status: "complete", PriceID: "price_pro", CustomerID: "cus_9", SubscriptionID: "sub_9"}}
s := newService(r, p)
p.event = &Event{ID: "evt_unpaid", Type: "checkout.session.completed", Created: 30, ObjectID: "cs_1"}
require.NoError(t, s.HandleWebhook(context.Background(), []byte("unpaid"), "sig"))
require.Equal(t, PlanFree, r.state.PlanID)
require.Empty(t, r.purchases)
}
func TestCheckoutFailureAndInvoicePaid(t *testing.T) {
r := newFakeRepo()
a := &CheckoutAttempt{ID: "attempt_1", UID: 9, SessionID: "cs_1", FulfillmentStatus: "pending"}
r.attempts[a.ID] = a
r.state = SubscriptionState{UID: 9, PlanID: PlanStarter, CustomerID: "cus_9", SubscriptionID: "sub_9"}
p := &fakeProvider{subscription: ProviderSubscription{ID: "sub_9", CustomerID: "cus_9", Status: "past_due", PriceID: "price_pro", CurrentPeriodStart: 10, CurrentPeriodEnd: 20}}
s := newService(r, p)
p.event = &Event{ID: "evt_failed", Type: "checkout.session.async_payment_failed", Created: 10, ObjectID: "cs_1"}
require.NoError(t, s.HandleWebhook(context.Background(), []byte("failed"), "sig"))
require.Equal(t, "failed", r.attempts[a.ID].FulfillmentStatus)
// The provider's current subscription price is authoritative, not an invoice metadata/line hint.
p.event = &Event{ID: "evt_invoice", Type: "invoice.paid", Created: 20, ObjectID: "in_1", SubscriptionID: "sub_9", CustomerID: "cus_9", PriceID: "price_starter"}
require.NoError(t, s.HandleWebhook(context.Background(), []byte("invoice"), "sig"))
require.Equal(t, PlanPro, r.state.PlanID)
require.Equal(t, "active", r.state.Status)
}
func TestDisabled(t *testing.T) {
_, err := (&Service{}).CreateCheckout(context.Background(), 1, PlanStarter, "x")
require.True(t, errors.Is(err, ErrDisabled))
}

View File

@ -0,0 +1,199 @@
package billing
import (
"context"
"encoding/json"
"errors"
"strconv"
"time"
"github.com/stripe/stripe-go/v82"
"github.com/stripe/stripe-go/v82/client"
"github.com/stripe/stripe-go/v82/webhook"
)
type StripeProvider struct{ api *client.API }
func NewStripeProvider(secretKey string) *StripeProvider {
return &StripeProvider{api: client.New(secretKey, nil)}
}
func (p *StripeProvider) CreateCheckout(ctx context.Context, input CheckoutParams) (*ProviderCheckout, error) {
uidText := stripe.String(formatUID(input.UID))
params := &stripe.CheckoutSessionParams{
Params: stripe.Params{Context: ctx},
Mode: stripe.String(string(stripe.CheckoutSessionModeSubscription)),
SuccessURL: stripe.String(input.SuccessURL), CancelURL: stripe.String(input.CancelURL),
ClientReferenceID: uidText,
Metadata: map[string]string{"uid": formatUID(input.UID), "attempt_id": input.AttemptID, "plan_id": input.PlanID},
SubscriptionData: &stripe.CheckoutSessionSubscriptionDataParams{Metadata: map[string]string{"uid": formatUID(input.UID), "attempt_id": input.AttemptID}},
LineItems: []*stripe.CheckoutSessionLineItemParams{{Price: stripe.String(input.PriceID), Quantity: stripe.Int64(1)}},
}
if input.CustomerID != "" {
params.Customer = stripe.String(input.CustomerID)
}
params.SetIdempotencyKey(input.IdempotencyKey)
session, err := p.api.CheckoutSessions.New(params)
if err != nil {
return nil, err
}
return checkoutFromStripe(session), nil
}
func (p *StripeProvider) GetCheckout(ctx context.Context, id string) (*ProviderCheckout, error) {
params := &stripe.CheckoutSessionParams{Params: stripe.Params{Context: ctx}}
params.AddExpand("line_items.data.price")
session, err := p.api.CheckoutSessions.Get(id, params)
if err != nil {
return nil, err
}
return checkoutFromStripe(session), nil
}
func checkoutFromStripe(s *stripe.CheckoutSession) *ProviderCheckout {
co := &ProviderCheckout{SessionID: s.ID, URL: s.URL, ExpiresAt: s.ExpiresAt * int64(time.Second), Status: string(s.Status), PaymentStatus: string(s.PaymentStatus)}
if s.Customer != nil {
co.CustomerID = s.Customer.ID
}
if s.Subscription != nil {
co.SubscriptionID = s.Subscription.ID
}
if s.LineItems != nil {
for _, item := range s.LineItems.Data {
if item != nil && item.Price != nil {
co.PriceID = item.Price.ID
break
}
}
}
return co
}
func (p *StripeProvider) GetSubscription(ctx context.Context, id string) (*ProviderSubscription, error) {
params := &stripe.SubscriptionParams{Params: stripe.Params{Context: ctx}}
params.AddExpand("items.data.price")
s, err := p.api.Subscriptions.Get(id, params)
if err != nil {
return nil, err
}
out := &ProviderSubscription{ID: s.ID, Status: string(s.Status), CancelAtPeriodEnd: s.CancelAtPeriodEnd}
if s.Customer != nil {
out.CustomerID = s.Customer.ID
}
if s.Items != nil {
for _, item := range s.Items.Data {
if item == nil {
continue
}
if item.Price != nil && out.PriceID == "" {
out.PriceID = item.Price.ID
}
if item.CurrentPeriodStart > out.CurrentPeriodStart {
out.CurrentPeriodStart = item.CurrentPeriodStart
}
if item.CurrentPeriodEnd > out.CurrentPeriodEnd {
out.CurrentPeriodEnd = item.CurrentPeriodEnd
}
}
}
return out, nil
}
func (p *StripeProvider) CreatePortal(ctx context.Context, customerID, returnURL string) (string, error) {
s, err := p.api.BillingPortalSessions.New(&stripe.BillingPortalSessionParams{Params: stripe.Params{Context: ctx}, Customer: stripe.String(customerID), ReturnURL: stripe.String(returnURL)})
if err != nil {
return "", err
}
return s.URL, nil
}
func (p *StripeProvider) VerifyEvent(payload []byte, signature, secret string) (*Event, error) {
e, err := webhook.ConstructEventWithOptions(payload, signature, secret, webhook.ConstructEventOptions{Tolerance: 5 * time.Minute})
if err != nil {
return nil, errors.Join(ErrInvalidSignature, err)
}
out := &Event{ID: e.ID, Type: string(e.Type), Created: e.Created}
var obj stripeObject
if err := json.Unmarshal(e.Data.Raw, &obj); err != nil {
return nil, err
}
out.ObjectID = obj.ID
out.CustomerID = objectID(obj.Customer)
out.Status = obj.Status
out.PeriodStart, out.PeriodEnd = obj.PeriodStart, obj.PeriodEnd
out.CancelAtPeriodEnd = obj.CancelAtPeriodEnd
out.SubscriptionID = objectID(obj.Subscription)
if out.SubscriptionID == "" {
out.SubscriptionID = objectID(obj.Parent.SubscriptionDetails.Subscription)
}
if len(obj.Items.Data) > 0 {
out.PriceID = objectID(obj.Items.Data[0].Price)
out.PeriodStart, out.PeriodEnd = obj.Items.Data[0].CurrentPeriodStart, obj.Items.Data[0].CurrentPeriodEnd
}
if len(obj.Lines.Data) > 0 {
line := obj.Lines.Data[0]
out.PriceID = objectID(line.Price)
if out.PriceID == "" {
out.PriceID = objectID(line.Pricing.PriceDetails.Price)
}
if line.Period.Start != 0 {
out.PeriodStart, out.PeriodEnd = line.Period.Start, line.Period.End
}
}
return out, nil
}
type stripeObject struct {
ID string `json:"id"`
Customer json.RawMessage `json:"customer"`
Subscription json.RawMessage `json:"subscription"`
Status string `json:"status"`
PeriodStart int64 `json:"period_start"`
PeriodEnd int64 `json:"period_end"`
CancelAtPeriodEnd bool `json:"cancel_at_period_end"`
Parent struct {
SubscriptionDetails struct {
Subscription json.RawMessage `json:"subscription"`
} `json:"subscription_details"`
} `json:"parent"`
Items struct {
Data []struct {
Price json.RawMessage `json:"price"`
CurrentPeriodStart int64 `json:"current_period_start"`
CurrentPeriodEnd int64 `json:"current_period_end"`
} `json:"data"`
} `json:"items"`
Lines struct {
Data []struct {
Price json.RawMessage `json:"price"`
Period struct {
Start int64 `json:"start"`
End int64 `json:"end"`
} `json:"period"`
Pricing struct {
PriceDetails struct {
Price json.RawMessage `json:"price"`
} `json:"price_details"`
} `json:"pricing"`
} `json:"data"`
} `json:"lines"`
}
func objectID(raw json.RawMessage) string {
if len(raw) == 0 || string(raw) == "null" {
return ""
}
var id string
if json.Unmarshal(raw, &id) == nil {
return id
}
var obj struct {
ID string `json:"id"`
}
_ = json.Unmarshal(raw, &obj)
return obj.ID
}
func formatUID(uid int64) string { return strconv.FormatInt(uid, 10) }
var _ Provider = (*StripeProvider)(nil)

View File

@ -0,0 +1,20 @@
package billing
import (
"testing"
"time"
"github.com/stretchr/testify/require"
"github.com/stripe/stripe-go/v82/webhook"
)
func TestStripeOfficialSignatureVerifier(t *testing.T) {
payload := []byte(`{"id":"evt_1","object":"event","api_version":"2025-06-30.basil","created":123,"type":"customer.subscription.updated","data":{"object":{"id":"sub_1","customer":"cus_1","status":"active","items":{"data":[{"price":{"id":"price_pro"},"current_period_start":10,"current_period_end":20}]}}}}`)
signed := webhook.GenerateTestSignedPayload(&webhook.UnsignedPayload{Payload: payload, Secret: "whsec_test", Timestamp: time.Now()})
e, err := NewStripeProvider("sk_test_unused").VerifyEvent(signed.Payload, signed.Header, "whsec_test")
require.NoError(t, err)
require.Equal(t, "price_pro", e.PriceID)
require.Equal(t, "cus_1", e.CustomerID)
_, err = NewStripeProvider("sk_test_unused").VerifyEvent(signed.Payload, signed.Header, "wrong")
require.ErrorIs(t, err, ErrInvalidSignature)
}

View File

@ -45,6 +45,9 @@ type Job struct {
ProgressPercent int `bson:"progress_percent" json:"progress_percent"` ProgressPercent int `bson:"progress_percent" json:"progress_percent"`
Error string `bson:"error,omitempty" json:"error,omitempty"` Error string `bson:"error,omitempty" json:"error,omitempty"`
WorkerID string `bson:"worker_id,omitempty" json:"worker_id,omitempty"` WorkerID string `bson:"worker_id,omitempty" json:"worker_id,omitempty"`
LeaseOwner string `bson:"lease_owner,omitempty" json:"lease_owner,omitempty"`
LeaseExpiresAt int64 `bson:"lease_expires_at,omitempty" json:"lease_expires_at,omitempty"`
Attempt int `bson:"attempt,omitempty" json:"attempt,omitempty"`
// RefID e.g. threads account id for renew / persona id for analyze // RefID e.g. threads account id for renew / persona id for analyze
RefID string `bson:"ref_id,omitempty" json:"ref_id,omitempty"` RefID string `bson:"ref_id,omitempty" json:"ref_id,omitempty"`
// Payload JSON body for job worker (e.g. {"username":"x"}) // Payload JSON body for job worker (e.g. {"username":"x"})
@ -62,6 +65,13 @@ type Job struct {
// DefaultTokenRenewDelayNs — 連帳後約 30 天做第一次 renew長效 token ~60 天) // DefaultTokenRenewDelayNs — 連帳後約 30 天做第一次 renew長效 token ~60 天)
const DefaultTokenRenewDelayNs = int64(30 * 24 * time.Hour) const DefaultTokenRenewDelayNs = int64(30 * 24 * time.Hour)
// Job leases are renewed by the claiming service while work is running.
const (
DefaultLeaseDuration = time.Minute
DefaultLeaseDurationNs = int64(DefaultLeaseDuration)
DefaultHeartbeat = 20 * time.Second
)
// TerminalRetentionNs — 終態任務(已完成/失敗/取消)保留多久後自動清除 // TerminalRetentionNs — 終態任務(已完成/失敗/取消)保留多久後自動清除
const TerminalRetentionNs = int64(2 * 24 * time.Hour) const TerminalRetentionNs = int64(2 * 24 * time.Hour)

View File

@ -5,11 +5,14 @@ import "context"
type Repository interface { type Repository interface {
Insert(ctx context.Context, j *Job) error Insert(ctx context.Context, j *Job) error
Update(ctx context.Context, j *Job) error Update(ctx context.Context, j *Job) error
// UpdateOwned atomically requires the current running lease to belong to leaseOwner.
UpdateOwned(ctx context.Context, j *Job, leaseOwner string) error
RenewLease(ctx context.Context, id, leaseOwner string, leaseExpiresAt int64) error
FindByID(ctx context.Context, id string) (*Job, error) FindByID(ctx context.Context, id string) (*Job, error)
ListByOwner(ctx context.Context, ownerUID int64) ([]*Job, error) ListByOwner(ctx context.Context, ownerUID int64) ([]*Job, error)
// ListByOwnerTab 分頁分桶total 為該 tab 總筆數 // ListByOwnerTab 分頁分桶total 為該 tab 總筆數
ListByOwnerTab(ctx context.Context, ownerUID int64, tab string, page, pageSize int) (list []*Job, total int64, err error) ListByOwnerTab(ctx context.Context, ownerUID int64, tab string, page, pageSize int) (list []*Job, total int64, err error)
// ClaimNext picks oldest due pending/queued (run_after <= now) and sets running + workerID. // ClaimNext atomically claims a due queued job or a running job with an expired lease.
ClaimNext(ctx context.Context, workerID string) (*Job, error) ClaimNext(ctx context.Context, workerID string) (*Job, error)
// CancelPendingByRef cancels pending/queued jobs of template+ref (e.g. reschedule renew). // CancelPendingByRef cancels pending/queued jobs of template+ref (e.g. reschedule renew).
CancelPendingByRef(ctx context.Context, ownerUID int64, template, refID string) error CancelPendingByRef(ctx context.Context, ownerUID int64, template, refID string) error

View File

@ -25,6 +25,14 @@ func (s *MemoryStore) Insert(_ context.Context, j *domain.Job) error {
} }
func (s *MemoryStore) Update(_ context.Context, j *domain.Job) error { func (s *MemoryStore) Update(_ context.Context, j *domain.Job) error {
return s.update(j, "")
}
func (s *MemoryStore) UpdateOwned(_ context.Context, j *domain.Job, leaseOwner string) error {
return s.update(j, leaseOwner)
}
func (s *MemoryStore) update(j *domain.Job, leaseOwner string) error {
s.mu.Lock() s.mu.Lock()
defer s.mu.Unlock() defer s.mu.Unlock()
stored, ok := s.byID[j.ID] stored, ok := s.byID[j.ID]
@ -34,6 +42,9 @@ func (s *MemoryStore) Update(_ context.Context, j *domain.Job) error {
if stored.Version != j.Version { if stored.Version != j.Version {
return domain.ErrIllegalStatus return domain.ErrIllegalStatus
} }
if leaseOwner != "" && (stored.Status != domain.StatusRunning || stored.LeaseOwner != leaseOwner) {
return domain.ErrIllegalStatus
}
cp := *j cp := *j
cp.Version++ cp.Version++
s.byID[j.ID] = &cp s.byID[j.ID] = &cp
@ -41,6 +52,21 @@ func (s *MemoryStore) Update(_ context.Context, j *domain.Job) error {
return nil return nil
} }
func (s *MemoryStore) RenewLease(_ context.Context, id, leaseOwner string, leaseExpiresAt int64) error {
s.mu.Lock()
defer s.mu.Unlock()
j, ok := s.byID[id]
if !ok {
return domain.ErrNotFound
}
if j.Status != domain.StatusRunning || j.LeaseOwner != leaseOwner {
return domain.ErrIllegalStatus
}
j.LeaseExpiresAt = leaseExpiresAt
j.UpdatedAt = domain.NowNano()
return nil
}
func (s *MemoryStore) FindByID(_ context.Context, id string) (*domain.Job, error) { func (s *MemoryStore) FindByID(_ context.Context, id string) (*domain.Job, error) {
s.mu.Lock() s.mu.Lock()
defer s.mu.Unlock() defer s.mu.Unlock()
@ -115,10 +141,9 @@ func (s *MemoryStore) ClaimNext(_ context.Context, workerID string) (*domain.Job
now := domain.NowNano() now := domain.NowNano()
var best *domain.Job var best *domain.Job
for _, j := range s.byID { for _, j := range s.byID {
if j.Status != domain.StatusPending && j.Status != domain.StatusQueued { queuedDue := (j.Status == domain.StatusPending || j.Status == domain.StatusQueued) && (j.RunAfter <= 0 || j.RunAfter <= now)
continue staleRunning := j.Status == domain.StatusRunning && j.LeaseExpiresAt <= now
} if !queuedDue && !staleRunning {
if j.RunAfter > 0 && j.RunAfter > now {
continue continue
} }
if best == nil || j.CreatedAt < best.CreatedAt { if best == nil || j.CreatedAt < best.CreatedAt {
@ -129,13 +154,11 @@ func (s *MemoryStore) ClaimNext(_ context.Context, workerID string) (*domain.Job
if best == nil { if best == nil {
return nil, domain.ErrNotFound return nil, domain.ErrNotFound
} }
from := best.Status best.Status = domain.StatusRunning
to := domain.StatusRunning
if !domain.CanTransition(from, to) {
return nil, domain.ErrIllegalStatus
}
best.Status = to
best.WorkerID = workerID best.WorkerID = workerID
best.LeaseOwner = workerID
best.LeaseExpiresAt = now + domain.DefaultLeaseDurationNs
best.Attempt++
best.UpdatedAt = now best.UpdatedAt = now
best.Version++ best.Version++
if best.ProgressSummary == "" { if best.ProgressSummary == "" {

View File

@ -28,6 +28,15 @@ func (s *MonStore) Insert(ctx context.Context, j *domain.Job) error {
} }
func (s *MonStore) Update(ctx context.Context, j *domain.Job) error { func (s *MonStore) Update(ctx context.Context, j *domain.Job) error {
return s.update(ctx, j, nil)
}
func (s *MonStore) UpdateOwned(ctx context.Context, j *domain.Job, leaseOwner string) error {
filter := bson.M{"status": domain.StatusRunning, "lease_owner": leaseOwner}
return s.update(ctx, j, filter)
}
func (s *MonStore) update(ctx context.Context, j *domain.Job, guard bson.M) error {
if j == nil { if j == nil {
return domain.ErrNotFound return domain.ErrNotFound
} }
@ -35,7 +44,11 @@ func (s *MonStore) Update(ctx context.Context, j *domain.Job) error {
j.UpdatedAt = domain.NowNano() j.UpdatedAt = domain.NowNano()
candidate := *j candidate := *j
candidate.Version = expectedVersion + 1 candidate.Version = expectedVersion + 1
res, err := s.jobs.ReplaceOne(ctx, bson.M{"_id": j.ID, "version": expectedVersion}, &candidate) filter := bson.M{"_id": j.ID, "version": expectedVersion}
for key, value := range guard {
filter[key] = value
}
res, err := s.jobs.ReplaceOne(ctx, filter, &candidate)
if err != nil { if err != nil {
return err return err
} }
@ -46,6 +59,21 @@ func (s *MonStore) Update(ctx context.Context, j *domain.Job) error {
return nil return nil
} }
func (s *MonStore) RenewLease(ctx context.Context, id, leaseOwner string, leaseExpiresAt int64) error {
res, err := s.jobs.UpdateOne(ctx, bson.M{
"_id": id, "status": domain.StatusRunning, "lease_owner": leaseOwner,
}, bson.M{
"$set": bson.M{"lease_expires_at": leaseExpiresAt, "updated_at": domain.NowNano()},
})
if err != nil {
return err
}
if res.MatchedCount == 0 {
return domain.ErrIllegalStatus
}
return nil
}
func (s *MonStore) FindByID(ctx context.Context, id string) (*domain.Job, error) { func (s *MonStore) FindByID(ctx context.Context, id string) (*domain.Job, error) {
var j domain.Job var j domain.Job
err := s.jobs.FindOne(ctx, &j, bson.M{"_id": id}) err := s.jobs.FindOne(ctx, &j, bson.M{"_id": id})
@ -120,21 +148,34 @@ func (s *MonStore) ListByOwnerTab(ctx context.Context, ownerUID int64, tab strin
func (s *MonStore) ClaimNext(ctx context.Context, workerID string) (*domain.Job, error) { func (s *MonStore) ClaimNext(ctx context.Context, workerID string) (*domain.Job, error) {
now := domain.NowNano() now := domain.NowNano()
// due: run_after missing/0 OR run_after <= now
filter := bson.M{ filter := bson.M{
"status": bson.M{"$in": []string{domain.StatusPending, domain.StatusQueued}},
"$or": []bson.M{ "$or": []bson.M{
{"run_after": bson.M{"$exists": false}}, {
{"run_after": 0}, "status": bson.M{"$in": []string{domain.StatusPending, domain.StatusQueued}},
{"run_after": bson.M{"$lte": now}}, "$or": []bson.M{
{"run_after": bson.M{"$exists": false}},
{"run_after": 0},
{"run_after": bson.M{"$lte": now}},
},
},
{
"status": domain.StatusRunning,
"$or": []bson.M{
{"lease_expires_at": bson.M{"$exists": false}},
{"lease_expires_at": 0},
{"lease_expires_at": bson.M{"$lte": now}},
},
},
}, },
} }
update := bson.M{"$set": bson.M{ update := bson.M{"$set": bson.M{
"status": domain.StatusRunning, "status": domain.StatusRunning,
"worker_id": workerID, "worker_id": workerID,
"lease_owner": workerID,
"lease_expires_at": now + domain.DefaultLeaseDurationNs,
"updated_at": now, "updated_at": now,
"progress_summary": "已由 worker 領取(" + workerID + "", "progress_summary": "已由 worker 領取(" + workerID + "",
}, "$inc": bson.M{"version": 1}} }, "$inc": bson.M{"version": 1, "attempt": 1}}
opts := options.FindOneAndUpdate(). opts := options.FindOneAndUpdate().
SetSort(bson.D{{Key: "run_after", Value: 1}, {Key: "created_at", Value: 1}}). SetSort(bson.D{{Key: "run_after", Value: 1}, {Key: "created_at", Value: 1}}).
SetReturnDocument(options.After) SetReturnDocument(options.After)

View File

@ -5,6 +5,8 @@ import (
"encoding/json" "encoding/json"
"fmt" "fmt"
"strings" "strings"
"sync"
"time"
"apps/backend/internal/module/job/domain" "apps/backend/internal/module/job/domain"
@ -18,12 +20,84 @@ type Notifier interface {
} }
type Service struct { type Service struct {
Repo domain.Repository Repo domain.Repository
Notifier Notifier Notifier Notifier
LeaseDuration time.Duration
HeartbeatInterval time.Duration
mu sync.Mutex
workerID string
heartbeats map[string]context.CancelFunc
} }
func New(repo domain.Repository) *Service { func New(repo domain.Repository) *Service {
return &Service{Repo: repo} return &Service{
Repo: repo, LeaseDuration: domain.DefaultLeaseDuration,
HeartbeatInterval: domain.DefaultHeartbeat,
heartbeats: make(map[string]context.CancelFunc),
}
}
func (s *Service) leaseDuration() time.Duration {
if s.LeaseDuration > 0 {
return s.LeaseDuration
}
return domain.DefaultLeaseDuration
}
func (s *Service) leaseOwner() string {
s.mu.Lock()
defer s.mu.Unlock()
return s.workerID
}
func (s *Service) update(ctx context.Context, j *domain.Job) error {
if owner := s.leaseOwner(); owner != "" {
j.LeaseExpiresAt = domain.NowNano() + int64(s.leaseDuration())
return s.Repo.UpdateOwned(ctx, j, owner)
}
return s.Repo.Update(ctx, j)
}
func (s *Service) startHeartbeat(ctx context.Context, jobID, workerID string) {
heartbeatCtx, cancel := context.WithCancel(ctx)
s.mu.Lock()
if previous := s.heartbeats[jobID]; previous != nil {
previous()
}
s.heartbeats[jobID] = cancel
s.mu.Unlock()
go func() {
interval := s.HeartbeatInterval
if interval <= 0 {
interval = domain.DefaultHeartbeat
}
ticker := time.NewTicker(interval)
defer ticker.Stop()
for {
select {
case <-heartbeatCtx.Done():
return
case <-ticker.C:
expiresAt := domain.NowNano() + int64(s.leaseDuration())
if err := s.Repo.RenewLease(heartbeatCtx, jobID, workerID, expiresAt); err != nil {
s.stopHeartbeat(jobID)
return
}
}
}
}()
}
func (s *Service) stopHeartbeat(jobID string) {
s.mu.Lock()
cancel := s.heartbeats[jobID]
delete(s.heartbeats, jobID)
s.mu.Unlock()
if cancel != nil {
cancel()
}
} }
// StartDemo — JB-01 // StartDemo — JB-01
@ -112,9 +186,12 @@ func (s *Service) Transition(ctx context.Context, id, to, summary string, percen
} }
} }
} }
if err := s.Repo.Update(ctx, j); err != nil { if err := s.update(ctx, j); err != nil {
return nil, err return nil, err
} }
if domain.IsTerminal(to) {
s.stopHeartbeat(id)
}
if s.Notifier != nil && (to == domain.StatusSucceeded || to == domain.StatusFailed || to == domain.StatusCancelled) { if s.Notifier != nil && (to == domain.StatusSucceeded || to == domain.StatusFailed || to == domain.StatusCancelled) {
pct := j.ProgressPercent pct := j.ProgressPercent
if to == domain.StatusSucceeded { if to == domain.StatusSucceeded {
@ -126,10 +203,17 @@ func (s *Service) Transition(ctx context.Context, id, to, summary string, percen
} }
func (s *Service) ClaimNext(ctx context.Context, workerID string) (*domain.Job, error) { func (s *Service) ClaimNext(ctx context.Context, workerID string) (*domain.Job, error) {
workerID = strings.TrimSpace(workerID)
if workerID == "" {
return nil, domain.ErrForbidden
}
j, err := s.Repo.ClaimNext(ctx, workerID) j, err := s.Repo.ClaimNext(ctx, workerID)
if err != nil { if err != nil {
return nil, err return nil, err
} }
s.mu.Lock()
s.workerID = workerID
s.mu.Unlock()
// claim → running立刻寫鈴鐺列表進度所有模板共用 // claim → running立刻寫鈴鐺列表進度所有模板共用
if j.ProgressPercent < 5 { if j.ProgressPercent < 5 {
j.ProgressPercent = 5 j.ProgressPercent = 5
@ -138,7 +222,10 @@ func (s *Service) ClaimNext(ctx context.Context, workerID string) (*domain.Job,
j.ProgressSummary = "已由 worker 領取 · 執行中" j.ProgressSummary = "已由 worker 領取 · 執行中"
} }
j.UpdatedAt = domain.NowNano() j.UpdatedAt = domain.NowNano()
_ = s.Repo.Update(ctx, j) if err := s.Repo.UpdateOwned(ctx, j, workerID); err != nil {
return nil, err
}
s.startHeartbeat(ctx, j.ID, workerID)
s.notify(ctx, j) s.notify(ctx, j)
return j, nil return j, nil
} }
@ -159,7 +246,7 @@ func (s *Service) RunDemoToSuccess(ctx context.Context, jobID string) (*domain.J
j.ProgressSummary = "Demo 測試任務 · 執行中" j.ProgressSummary = "Demo 測試任務 · 執行中"
j.ProgressPercent = 5 j.ProgressPercent = 5
j.UpdatedAt = domain.NowNano() j.UpdatedAt = domain.NowNano()
_ = s.Repo.Update(ctx, j) _ = s.update(ctx, j)
s.notify(ctx, j) s.notify(ctx, j)
} }
if j.Status != domain.StatusRunning { if j.Status != domain.StatusRunning {
@ -351,9 +438,10 @@ func (s *Service) SucceedJobWithPayload(ctx context.Context, jobID, summary, pay
j.Status = domain.StatusSucceeded j.Status = domain.StatusSucceeded
j.CompletedAt = domain.NowNano() j.CompletedAt = domain.NowNano()
j.UpdatedAt = j.CompletedAt j.UpdatedAt = j.CompletedAt
if err := s.Repo.Update(ctx, j); err != nil { if err := s.update(ctx, j); err != nil {
return nil, err return nil, err
} }
s.stopHeartbeat(jobID)
s.notify(ctx, j) s.notify(ctx, j)
return j, nil return j, nil
} }
@ -406,7 +494,7 @@ func (s *Service) MarkRunningProgress(ctx context.Context, jobID string, pct int
if err := domain.ApplyProgress(j, pct, summary); err != nil { if err := domain.ApplyProgress(j, pct, summary); err != nil {
return nil, err return nil, err
} }
if err := s.Repo.Update(ctx, j); err != nil { if err := s.update(ctx, j); err != nil {
return nil, err return nil, err
} }
// 每次有進度就更新鈴鐺upsert 同一則) // 每次有進度就更新鈴鐺upsert 同一則)
@ -454,7 +542,7 @@ func (s *Service) FailJob(ctx context.Context, jobID, errMsg string) (*domain.Jo
if j.Status == domain.StatusPending || j.Status == domain.StatusQueued { if j.Status == domain.StatusPending || j.Status == domain.StatusQueued {
j.Status = domain.StatusRunning j.Status = domain.StatusRunning
j.UpdatedAt = domain.NowNano() j.UpdatedAt = domain.NowNano()
_ = s.Repo.Update(ctx, j) _ = s.update(ctx, j)
} }
if !domain.CanTransition(j.Status, domain.StatusFailed) { if !domain.CanTransition(j.Status, domain.StatusFailed) {
return nil, domain.ErrIllegalStatus return nil, domain.ErrIllegalStatus
@ -471,9 +559,10 @@ func (s *Service) FailJob(ctx context.Context, jobID, errMsg string) (*domain.Jo
} }
j.CompletedAt = domain.NowNano() j.CompletedAt = domain.NowNano()
j.UpdatedAt = j.CompletedAt j.UpdatedAt = j.CompletedAt
if err := s.Repo.Update(ctx, j); err != nil { if err := s.update(ctx, j); err != nil {
return nil, err return nil, err
} }
s.stopHeartbeat(jobID)
s.notify(ctx, j) s.notify(ctx, j)
return j, nil return j, nil
} }

View File

@ -3,6 +3,7 @@ package usecase_test
import ( import (
"context" "context"
"testing" "testing"
"time"
"apps/backend/internal/module/appnotif/repository" "apps/backend/internal/module/appnotif/repository"
appUC "apps/backend/internal/module/appnotif/usecase" appUC "apps/backend/internal/module/appnotif/usecase"
@ -274,3 +275,84 @@ func TestJB_ListTab_BucketsAndPaging(t *testing.T) {
require.NoError(t, err) require.NoError(t, err)
require.Len(t, page1, 1) require.Len(t, page1, 1)
} }
func TestJobLease_ReclaimsStaleRunningAndRejectsOldOwner(t *testing.T) {
ctx1, cancel1 := context.WithCancel(context.Background())
defer cancel1()
ctx2, cancel2 := context.WithCancel(context.Background())
defer cancel2()
repo := jobRepo.NewMemory()
worker1 := usecase.New(repo)
worker2 := usecase.New(repo)
queued, err := worker1.StartDemo(context.Background(), 1_000_300)
require.NoError(t, err)
first, err := worker1.ClaimNext(ctx1, "worker-1")
require.NoError(t, err)
require.Equal(t, queued.ID, first.ID)
require.Equal(t, "worker-1", first.LeaseOwner)
require.Equal(t, 1, first.Attempt)
require.Greater(t, first.LeaseExpiresAt, domain.NowNano())
_, err = worker2.ClaimNext(ctx2, "worker-2")
require.ErrorIs(t, err, domain.ErrNotFound)
stale, err := repo.FindByID(context.Background(), queued.ID)
require.NoError(t, err)
stale.LeaseExpiresAt = domain.NowNano() - 1
require.NoError(t, repo.Update(context.Background(), stale))
reclaimed, err := worker2.ClaimNext(ctx2, "worker-2")
require.NoError(t, err)
require.Equal(t, "worker-2", reclaimed.LeaseOwner)
require.Equal(t, 2, reclaimed.Attempt)
_, err = worker1.MarkRunningProgress(context.Background(), queued.ID, 30, "stale write")
require.ErrorIs(t, err, domain.ErrIllegalStatus)
_, err = worker1.SucceedJob(context.Background(), queued.ID, "stale completion")
require.ErrorIs(t, err, domain.ErrIllegalStatus)
_, err = worker1.FailJob(context.Background(), queued.ID, "stale failure")
require.ErrorIs(t, err, domain.ErrIllegalStatus)
updated, err := worker2.MarkRunningProgress(context.Background(), queued.ID, 30, "current write")
require.NoError(t, err)
require.Equal(t, "current write", updated.ProgressSummary)
completed, err := worker2.SucceedJob(context.Background(), queued.ID, "current completion")
require.NoError(t, err)
require.Equal(t, domain.StatusSucceeded, completed.Status)
}
func TestJobLease_HeartbeatRenewsWhileRunning(t *testing.T) {
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
repo := jobRepo.NewMemory()
svc := usecase.New(repo)
svc.HeartbeatInterval = 5 * time.Millisecond
svc.LeaseDuration = 2 * domain.DefaultLeaseDuration
_, err := svc.StartDemo(context.Background(), 1_000_301)
require.NoError(t, err)
claimed, err := svc.ClaimNext(ctx, "worker-heartbeat")
require.NoError(t, err)
initialExpiry := claimed.LeaseExpiresAt
require.Eventually(t, func() bool {
current, findErr := repo.FindByID(context.Background(), claimed.ID)
return findErr == nil && current.LeaseExpiresAt > initialExpiry
}, 200*time.Millisecond, 5*time.Millisecond)
}
func TestJobLease_ReclaimsLegacyRunningDocumentWithoutLease(t *testing.T) {
repo := jobRepo.NewMemory()
now := domain.NowNano()
legacy := &domain.Job{
ID: "legacy-running", OwnerUID: 1_000_302, TemplateType: domain.TemplateDemo,
Status: domain.StatusRunning, WorkerID: "dead-worker", CreatedAt: now, UpdatedAt: now,
}
require.NoError(t, repo.Insert(context.Background(), legacy))
reclaimed, err := repo.ClaimNext(context.Background(), "replacement-worker")
require.NoError(t, err)
require.Equal(t, "replacement-worker", reclaimed.LeaseOwner)
require.Equal(t, 1, reclaimed.Attempt)
require.Greater(t, reclaimed.LeaseExpiresAt, now)
}

View File

@ -24,12 +24,13 @@ var (
// ErrWeakPassword — API 400003Message = PasswordPolicyEN // ErrWeakPassword — API 400003Message = PasswordPolicyEN
ErrWeakPassword = errors.New(PasswordPolicyEN) ErrWeakPassword = errors.New(PasswordPolicyEN)
ErrInvalidCode = errors.New("invalid or expired code") ErrInvalidCode = errors.New("invalid or expired code")
ErrCodeRateLimited = errors.New("verification code requested too frequently")
ErrAlreadyBound = errors.New("already bound") ErrAlreadyBound = errors.New("already bound")
ErrInvalidPlatform = errors.New("invalid platform") ErrInvalidPlatform = errors.New("invalid platform")
ErrOAuthFailed = errors.New("oauth verification failed") ErrOAuthFailed = errors.New("oauth verification failed")
// 邀請網絡message = FE i18n key invite.err.* // 邀請網絡message = FE i18n key invite.err.*
ErrInviteAlreadyBound = errors.New("invite.err.alreadyBound") ErrInviteAlreadyBound = errors.New("invite.err.alreadyBound")
ErrInviteCodeRequired = errors.New("invite.err.codeRequired") ErrInviteCodeRequired = errors.New("invite.err.codeRequired")
ErrInviteCodeNotFound = errors.New("invite.err.codeNotFound") ErrInviteCodeNotFound = errors.New("invite.err.codeNotFound")
ErrInviteCodeSelf = errors.New("invite.err.codeSelf") ErrInviteCodeSelf = errors.New("invite.err.codeSelf")

View File

@ -29,7 +29,8 @@ type Repository interface {
UpdateIdentityPassword(ctx context.Context, loginID, platform, passwordHash string) error UpdateIdentityPassword(ctx context.Context, loginID, platform, passwordHash string) error
// Codes / refresh (Mongo durable) // Codes / refresh (Mongo durable)
SetCode(kind, key, code string, ttl time.Duration) // SetCode returns false when a recent code exists or persistence fails.
SetCode(kind, key, code string, ttl time.Duration) bool
// 驗證碼共用語意(信箱驗證 / 重設密碼相同): // 驗證碼共用語意(信箱驗證 / 重設密碼相同):
// CheckCode碼正確才刪成功消費錯誤不刪過期才刪 // CheckCode碼正確才刪成功消費錯誤不刪過期才刪
// PeekCode只檢查不消費業務前先 Peek成功後再 Check // PeekCode只檢查不消費業務前先 Peek成功後再 Check

View File

@ -2,20 +2,25 @@ package repository
import ( import (
"context" "context"
"encoding/json"
"errors" "errors"
"fmt" "fmt"
"strings" "strings"
"time" "time"
libmongo "apps/backend/internal/lib/mongo" libmongo "apps/backend/internal/lib/mongo"
"apps/backend/internal/lib/secretbox"
"apps/backend/internal/module/member/domain" "apps/backend/internal/module/member/domain"
"github.com/zeromicro/go-zero/core/logx"
"github.com/zeromicro/go-zero/core/stores/cache" "github.com/zeromicro/go-zero/core/stores/cache"
"github.com/zeromicro/go-zero/core/stores/mon" "github.com/zeromicro/go-zero/core/stores/mon"
"github.com/zeromicro/go-zero/core/stores/monc" "github.com/zeromicro/go-zero/core/stores/monc"
"github.com/zeromicro/go-zero/core/stores/redis"
"go.mongodb.org/mongo-driver/bson" "go.mongodb.org/mongo-driver/bson"
"go.mongodb.org/mongo-driver/mongo" "go.mongodb.org/mongo-driver/mongo"
"go.mongodb.org/mongo-driver/mongo/options" "go.mongodb.org/mongo-driver/mongo/options"
"golang.org/x/crypto/bcrypt"
) )
const ( const (
@ -25,33 +30,45 @@ const (
colCounters = "counters" colCounters = "counters"
colRefresh = "refresh_tokens" colRefresh = "refresh_tokens"
colAuthCodes = "auth_codes" colAuthCodes = "auth_codes"
settingsCacheVersion = 1
settingsCacheTTL = 30 * time.Minute
settingsNegativeCacheTTL = time.Minute
settingsCacheTimeout = 750 * time.Millisecond
settingsMongoReadTimeout = 5 * time.Second
) )
// MoncStore: members/identities/settings via monc; refresh/codes durable Mongo. // MoncStore: members/identities via monc; settings use explicit ciphertext Redis caching.
type MoncStore struct { type MoncStore struct {
uri string uri string
db string db string
member *monc.Model member *monc.Model
plain *mon.Model plain *mon.Model
identity *monc.Model identity *monc.Model
idPlain *mon.Model idPlain *mon.Model
settings *monc.Model settings *mon.Model
refresh *mon.Model settingsCache settingsCache
codes *mon.Model settingsCacheNamespace string
refresh *mon.Model
codes *mon.Model
secrets *secretbox.Codec
} }
func NewMoncStore(uri, database string, cacheConf cache.CacheConf) *MoncStore { func NewMoncStore(uri, database string, cacheConf cache.CacheConf, settingsRedisConf redis.RedisConf, redisNamespace string, secrets *secretbox.Codec) *MoncStore {
uri = libmongo.MustMongoURI(uri) uri = libmongo.MustMongoURI(uri)
return &MoncStore{ return &MoncStore{
uri: uri, uri: uri,
db: database, db: database,
member: monc.MustNewModel(uri, database, colMembers, cacheConf), member: monc.MustNewModel(uri, database, colMembers, cacheConf),
plain: mon.MustNewModel(uri, database, colMembers), plain: mon.MustNewModel(uri, database, colMembers),
identity: monc.MustNewModel(uri, database, colIdentities, cacheConf), identity: monc.MustNewModel(uri, database, colIdentities, cacheConf),
idPlain: mon.MustNewModel(uri, database, colIdentities), idPlain: mon.MustNewModel(uri, database, colIdentities),
settings: monc.MustNewModel(uri, database, colSettings, cacheConf), settings: mon.MustNewModel(uri, database, colSettings),
refresh: mon.MustNewModel(uri, database, colRefresh), settingsCache: &redisSettingsCache{client: redis.MustNewRedis(settingsRedisConf)},
codes: mon.MustNewModel(uri, database, colAuthCodes), settingsCacheNamespace: normalizeSettingsCacheNamespace(redisNamespace),
refresh: mon.MustNewModel(uri, database, colRefresh),
codes: mon.MustNewModel(uri, database, colAuthCodes),
secrets: secrets,
} }
} }
@ -62,7 +79,6 @@ func cacheKeyEmail(email string) string {
func cacheKeyPhone(phone string) string { func cacheKeyPhone(phone string) string {
return fmt.Sprintf("cache:member:phone:%s", strings.TrimSpace(phone)) return fmt.Sprintf("cache:member:phone:%s", strings.TrimSpace(phone))
} }
func cacheKeySettings(uid int64) string { return fmt.Sprintf("cache:member:settings:%d", uid) }
func cacheKeyIdentity(loginID, platform string) string { func cacheKeyIdentity(loginID, platform string) string {
return fmt.Sprintf("cache:identity:%s:%s", platform, strings.ToLower(loginID)) return fmt.Sprintf("cache:identity:%s:%s", platform, strings.ToLower(loginID))
} }
@ -347,19 +363,37 @@ type authCodeDoc struct {
ID string `bson:"_id"` ID string `bson:"_id"`
Kind string `bson:"kind"` Kind string `bson:"kind"`
Key string `bson:"key"` Key string `bson:"key"`
Code string `bson:"code"` Code string `bson:"code,omitempty"` // legacy plaintext reader
CodeHash string `bson:"code_hash,omitempty"`
Attempts int `bson:"attempts"`
CreatedAt time.Time `bson:"created_at"`
ExpiresAt time.Time `bson:"expires_at"` ExpiresAt time.Time `bson:"expires_at"`
} }
func (s *MoncStore) SetCode(kind, key, code string, ttl time.Duration) { func (s *MoncStore) SetCode(kind, key, code string, ttl time.Duration) bool {
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel() defer cancel()
id := kind + ":" + key hash, err := bcrypt.GenerateFromPassword([]byte(code), bcrypt.MinCost)
doc := authCodeDoc{ if err != nil {
ID: id, Kind: kind, Key: key, Code: code, return false
ExpiresAt: time.Now().UTC().Add(ttl),
} }
_, _ = s.codes.ReplaceOne(ctx, bson.M{"_id": id}, doc, options.Replace().SetUpsert(true)) now := time.Now().UTC()
id := kind + ":" + key
filter := bson.M{"_id": id, "$or": bson.A{
bson.M{"created_at": bson.M{"$lte": now.Add(-time.Minute)}},
bson.M{"created_at": bson.M{"$exists": false}},
bson.M{"expires_at": bson.M{"$lte": now}},
}}
update := bson.M{"$set": bson.M{
"kind": kind, "key": key, "code_hash": string(hash), "attempts": 0,
"created_at": now, "expires_at": now.Add(ttl),
}, "$unset": bson.M{"code": ""}}
var doc authCodeDoc
err = s.codes.FindOneAndUpdate(ctx, &doc, filter, update, options.FindOneAndUpdate().SetUpsert(true).SetReturnDocument(options.After))
if err != nil {
return false
}
return true
} }
// CheckCode 比對驗證碼: // CheckCode 比對驗證碼:
@ -367,41 +401,40 @@ func (s *MoncStore) SetCode(kind, key, code string, ttl time.Duration) {
// - 錯誤 → 不刪TTL 內可重試 // - 錯誤 → 不刪TTL 內可重試
// - 過期 → 刪除(清垃圾)並回 false // - 過期 → 刪除(清垃圾)並回 false
func (s *MoncStore) CheckCode(kind, key, code string) bool { func (s *MoncStore) CheckCode(kind, key, code string) bool {
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) return s.verifyCode(kind, key, code, true)
defer cancel()
id := kind + ":" + key
var doc authCodeDoc
err := s.codes.FindOne(ctx, &doc, bson.M{"_id": id})
if err != nil {
return false
}
if time.Now().UTC().After(doc.ExpiresAt) {
_, _ = s.codes.DeleteOne(ctx, bson.M{"_id": id})
return false
}
if doc.Code != code {
// 輸錯不銷毀N 分鐘內可繼續試,不必重寄
return false
}
_, _ = s.codes.DeleteOne(ctx, bson.M{"_id": id})
return true
} }
// PeekCode 只檢查、不消費;過期時順便刪除。 // PeekCode 只檢查、不消費;過期時順便刪除。
func (s *MoncStore) PeekCode(kind, key, code string) bool { func (s *MoncStore) PeekCode(kind, key, code string) bool {
return s.verifyCode(kind, key, code, false)
}
func (s *MoncStore) verifyCode(kind, key, code string, consume bool) bool {
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel() defer cancel()
id := kind + ":" + key id := kind + ":" + key
var doc authCodeDoc var doc authCodeDoc
err := s.codes.FindOne(ctx, &doc, bson.M{"_id": id}) now := time.Now().UTC()
err := s.codes.FindOneAndUpdate(ctx, &doc, bson.M{
"_id": id, "expires_at": bson.M{"$gt": now}, "attempts": bson.M{"$lt": 5},
}, bson.M{"$inc": bson.M{"attempts": 1}}, options.FindOneAndUpdate().SetReturnDocument(options.Before))
if err != nil { if err != nil {
return false return false
} }
if time.Now().UTC().After(doc.ExpiresAt) { matched := false
_, _ = s.codes.DeleteOne(ctx, bson.M{"_id": id}) if doc.CodeHash != "" {
matched = bcrypt.CompareHashAndPassword([]byte(doc.CodeHash), []byte(code)) == nil
} else {
matched = doc.Code == code
}
if !matched {
return false return false
} }
return doc.Code == code if consume {
_, err = s.codes.DeleteOne(ctx, bson.M{"_id": id})
return err == nil
}
return true
} }
// --- durable refresh_tokens --- // --- durable refresh_tokens ---
@ -446,26 +479,312 @@ func (s *MoncStore) RevokeAllRefreshByUID(ctx context.Context, uid int64) error
// --- settings --- // --- settings ---
type settingsCache interface {
Get(context.Context, string) (string, error)
Set(context.Context, string, string, time.Duration) error
Delete(context.Context, string) error
}
type redisSettingsCache struct {
client *redis.Redis
}
func (c *redisSettingsCache) Get(ctx context.Context, key string) (string, error) {
return c.client.GetCtx(ctx, key)
}
func (c *redisSettingsCache) Set(ctx context.Context, key, value string, ttl time.Duration) error {
return c.client.SetexCtx(ctx, key, value, durationSeconds(ttl))
}
func (c *redisSettingsCache) Delete(ctx context.Context, key string) error {
_, err := c.client.DelCtx(ctx, key)
return err
}
type settingsCacheDTO struct {
Version int `json:"version"`
Found bool `json:"found"`
UID int64 `json:"uid"`
Provider string `json:"provider,omitempty"`
Model string `json:"model,omitempty"`
APIKeyConfigured bool `json:"api_key_configured,omitempty"`
ResearchAPIKeyConfigured bool `json:"research_api_key_configured,omitempty"`
APIKey string `json:"api_key,omitempty"`
ProviderAPIKeys map[string]string `json:"provider_api_keys,omitempty"`
ExaAPIKey string `json:"exa_api_key,omitempty"`
WebSearchProvider string `json:"web_search_provider,omitempty"`
ExpandStrategy string `json:"expand_strategy,omitempty"`
ExaAPIKeyConfigured bool `json:"exa_api_key_configured,omitempty"`
DevModeEnabled bool `json:"dev_mode_enabled,omitempty"`
UpdatedAt int64 `json:"updated_at,omitempty"`
}
func settingsCacheKey(namespace string, uid int64) string {
return fmt.Sprintf("%s:member:settings:{%d}", normalizeSettingsCacheNamespace(namespace), uid)
}
func normalizeSettingsCacheNamespace(namespace string) string {
if namespace = strings.Trim(strings.TrimSpace(namespace), ":"); namespace != "" {
return namespace
}
return "haixun:dev:v1"
}
func newSettingsCacheDTO(st *domain.UserSettings) (settingsCacheDTO, error) {
if st == nil {
return settingsCacheDTO{}, errors.New("nil member settings cache value")
}
if err := validateCachedCiphertext(st); err != nil {
return settingsCacheDTO{}, err
}
return settingsCacheDTO{
Version: settingsCacheVersion, Found: true, UID: st.UID,
Provider: st.Provider, Model: st.Model,
APIKeyConfigured: st.APIKeyConfigured, ResearchAPIKeyConfigured: st.ResearchAPIKeyConfigured,
APIKey: st.APIKey, ProviderAPIKeys: cloneStringMap(st.ProviderAPIKeys), ExaAPIKey: st.ExaAPIKey,
WebSearchProvider: st.WebSearchProvider, ExpandStrategy: st.ExpandStrategy,
ExaAPIKeyConfigured: st.ExaAPIKeyConfigured, DevModeEnabled: st.DevModeEnabled, UpdatedAt: st.UpdatedAt,
}, nil
}
func (d settingsCacheDTO) settings() *domain.UserSettings {
return &domain.UserSettings{
UID: d.UID, Provider: d.Provider, Model: d.Model,
APIKeyConfigured: d.APIKeyConfigured, ResearchAPIKeyConfigured: d.ResearchAPIKeyConfigured,
APIKey: d.APIKey, ProviderAPIKeys: cloneStringMap(d.ProviderAPIKeys), ExaAPIKey: d.ExaAPIKey,
WebSearchProvider: d.WebSearchProvider, ExpandStrategy: d.ExpandStrategy,
ExaAPIKeyConfigured: d.ExaAPIKeyConfigured, DevModeEnabled: d.DevModeEnabled, UpdatedAt: d.UpdatedAt,
}
}
func validateCachedCiphertext(st *domain.UserSettings) error {
values := map[string]string{"api_key": st.APIKey, "exa_api_key": st.ExaAPIKey}
for provider, value := range st.ProviderAPIKeys {
values["provider_api_keys."+provider] = value
}
for field, value := range values {
if value != "" && !secretbox.IsEncrypted(value) {
return fmt.Errorf("member settings cache %s is not ciphertext", field)
}
}
return nil
}
func encodeSettingsCache(dto settingsCacheDTO) (string, error) {
raw, err := json.Marshal(dto)
return string(raw), err
}
func decodeSettingsCache(raw string, uid int64) (settingsCacheDTO, error) {
var dto settingsCacheDTO
if err := json.Unmarshal([]byte(raw), &dto); err != nil {
return settingsCacheDTO{}, err
}
if dto.Version != settingsCacheVersion || dto.UID != uid {
return settingsCacheDTO{}, errors.New("invalid member settings cache version or uid")
}
if dto.Found {
if err := validateCachedCiphertext(dto.settings()); err != nil {
return settingsCacheDTO{}, err
}
}
return dto, nil
}
func readSettingsCache(ctx context.Context, cache settingsCache, key string, uid int64, codec *secretbox.Codec) (*domain.UserSettings, bool, error) {
raw, err := cache.Get(ctx, key)
if err != nil || raw == "" {
return nil, false, err
}
dto, err := decodeSettingsCache(raw, uid)
if err != nil {
_ = cache.Delete(ctx, key)
return nil, false, err
}
if !dto.Found {
return domain.DefaultSettings(uid), true, nil
}
st, err := decryptSettings(codec, uid, dto.settings())
if err != nil {
_ = cache.Delete(ctx, key)
return nil, false, err
}
return st, true, nil
}
func (s *MoncStore) GetSettings(ctx context.Context, uid int64) (*domain.UserSettings, error) { func (s *MoncStore) GetSettings(ctx context.Context, uid int64) (*domain.UserSettings, error) {
var st domain.UserSettings key := settingsCacheKey(s.settingsCacheNamespace, uid)
err := s.settings.FindOne(ctx, cacheKeySettings(uid), &st, bson.M{"uid": uid}) cacheCtx, cacheCancel := context.WithTimeout(ctx, settingsCacheTimeout)
if errors.Is(err, monc.ErrNotFound) || errors.Is(err, mon.ErrNotFound) { st, hit, cacheErr := readSettingsCache(cacheCtx, s.settingsCache, key, uid, s.secrets)
cacheCancel()
if hit {
return st, nil
}
if cacheErr != nil {
logx.Errorf("member settings cache read failed uid=%d: %v", uid, cacheErr)
}
mongoCtx, mongoCancel := context.WithTimeout(ctx, settingsMongoReadTimeout)
defer mongoCancel()
var encrypted domain.UserSettings
err := s.settings.FindOne(mongoCtx, &encrypted, bson.M{"uid": uid})
if errors.Is(err, mon.ErrNotFound) || errors.Is(err, mongo.ErrNoDocuments) {
s.writeSettingsCache(key, settingsCacheDTO{Version: settingsCacheVersion, UID: uid}, settingsNegativeCacheTTL)
return domain.DefaultSettings(uid), nil return domain.DefaultSettings(uid), nil
} }
if err != nil { if err != nil {
return nil, err return nil, err
} }
return &st, nil st, err = decryptSettings(s.secrets, uid, &encrypted)
if err != nil {
return nil, err
}
cacheValue := &encrypted
if err := validateCachedCiphertext(cacheValue); err != nil {
cacheValue, err = encryptSettings(s.secrets, uid, st)
if err != nil {
logx.Errorf("member settings cache legacy refill skipped uid=%d: %v", uid, err)
return st, nil
}
}
dto, err := newSettingsCacheDTO(cacheValue)
if err == nil {
s.writeSettingsCache(key, dto, settingsCacheTTL)
}
return st, nil
} }
func (s *MoncStore) SaveSettings(ctx context.Context, uid int64, st *domain.UserSettings) error { func (s *MoncStore) SaveSettings(ctx context.Context, uid int64, st *domain.UserSettings) error {
st.UID = uid encrypted, err := encryptSettings(s.secrets, uid, st)
st.UpdatedAt = time.Now().UTC().UnixNano()
settingsPlain := mon.MustNewModel(s.uri, s.db, colSettings)
_, err := settingsPlain.ReplaceOne(ctx, bson.M{"uid": uid}, st, options.Replace().SetUpsert(true))
if err != nil { if err != nil {
return err return err
} }
_ = s.settings.DelCache(ctx, cacheKeySettings(uid)) encrypted.UID = uid
encrypted.UpdatedAt = time.Now().UTC().UnixNano()
_, err = s.settings.ReplaceOne(ctx, bson.M{"uid": uid}, encrypted, options.Replace().SetUpsert(true))
if err != nil {
return err
}
dto, cacheErr := newSettingsCacheDTO(encrypted)
if cacheErr == nil {
cacheErr = s.writeSettingsCache(settingsCacheKey(s.settingsCacheNamespace, uid), dto, settingsCacheTTL)
}
if cacheErr != nil {
key := settingsCacheKey(s.settingsCacheNamespace, uid)
cacheCtx, cancel := context.WithTimeout(context.Background(), settingsCacheTimeout)
deleteErr := s.settingsCache.Delete(cacheCtx, key)
cancel()
logx.Errorf("CRITICAL: member settings saved to Mongo but Redis ciphertext cache write failed uid=%d cache_err=%v delete_err=%v", uid, cacheErr, deleteErr)
}
return nil return nil
} }
func (s *MoncStore) writeSettingsCache(key string, dto settingsCacheDTO, ttl time.Duration) error {
raw, err := encodeSettingsCache(dto)
if err != nil {
return err
}
cacheCtx, cancel := context.WithTimeout(context.Background(), settingsCacheTimeout)
defer cancel()
if err := s.settingsCache.Set(cacheCtx, key, raw, ttl); err != nil {
logx.Errorf("member settings cache refill failed key=%s: %v", key, err)
return err
}
return nil
}
func durationSeconds(ttl time.Duration) int {
seconds := int(ttl.Seconds())
if seconds < 1 {
return 1
}
return seconds
}
func encryptSettings(codec *secretbox.Codec, uid int64, st *domain.UserSettings) (*domain.UserSettings, error) {
if st == nil {
return nil, domain.ErrNotFound
}
out := cloneSettings(st)
var err error
out.APIKey, err = encryptSetting(codec, uid, "api_key", out.Provider, out.APIKey)
if err != nil {
return nil, err
}
for provider, key := range out.ProviderAPIKeys {
out.ProviderAPIKeys[provider], err = encryptSetting(codec, uid, "provider_api_keys", provider, key)
if err != nil {
return nil, err
}
}
out.ExaAPIKey, err = encryptSetting(codec, uid, "exa_api_key", "exa", out.ExaAPIKey)
if err != nil {
return nil, err
}
return out, nil
}
func decryptSettings(codec *secretbox.Codec, uid int64, st *domain.UserSettings) (*domain.UserSettings, error) {
if st == nil {
return nil, domain.ErrNotFound
}
out := cloneSettings(st)
var err error
out.APIKey, err = decryptSetting(codec, uid, "api_key", out.Provider, out.APIKey)
if err != nil {
return nil, err
}
for provider, key := range out.ProviderAPIKeys {
out.ProviderAPIKeys[provider], err = decryptSetting(codec, uid, "provider_api_keys", provider, key)
if err != nil {
return nil, err
}
}
out.ExaAPIKey, err = decryptSetting(codec, uid, "exa_api_key", "exa", out.ExaAPIKey)
if err != nil {
return nil, err
}
return out, nil
}
func encryptSetting(codec *secretbox.Codec, uid int64, field, provider, value string) (string, error) {
if value == "" {
return "", nil
}
if codec == nil {
return "", fmt.Errorf("member settings encryption codec is required for %s", field)
}
return codec.Encrypt(uid, field, provider, value)
}
func decryptSetting(codec *secretbox.Codec, uid int64, field, provider, value string) (string, error) {
if value == "" || !secretbox.IsEncrypted(value) {
return value, nil
}
if codec == nil {
return "", fmt.Errorf("member settings encryption codec is required for %s", field)
}
plaintext, err := codec.Decrypt(uid, field, provider, value)
if err != nil {
return "", fmt.Errorf("decrypt member settings %s: %w", field, err)
}
return plaintext, nil
}
func cloneSettings(st *domain.UserSettings) *domain.UserSettings {
out := *st
out.ProviderAPIKeys = cloneStringMap(st.ProviderAPIKeys)
return &out
}
func cloneStringMap(values map[string]string) map[string]string {
if values == nil {
return nil
}
out := make(map[string]string, len(values))
for key, value := range values {
out[key] = value
}
return out
}

View File

@ -0,0 +1,102 @@
package repository
import (
"context"
"errors"
"strings"
"testing"
"time"
"apps/backend/internal/module/member/domain"
)
type fakeSettingsCache struct {
value string
getErr error
deleted bool
}
func (c *fakeSettingsCache) Get(context.Context, string) (string, error) {
return c.value, c.getErr
}
func (c *fakeSettingsCache) Set(context.Context, string, string, time.Duration) error {
return nil
}
func (c *fakeSettingsCache) Delete(context.Context, string) error {
c.deleted = true
return nil
}
func TestSettingsCacheDTORoundTripRetainsCiphertextWithoutPlaintext(t *testing.T) {
codec := repositoryTestCodec(t)
plaintext := &domain.UserSettings{
UID: 42, Provider: "xai", Model: "grok-3", APIKey: "legacy-plain-secret",
ProviderAPIKeys: map[string]string{"xai": "provider-plain-secret"}, ExaAPIKey: "exa-plain-secret",
WebSearchProvider: "exa", ExpandStrategy: "hybrid", UpdatedAt: 123,
}
encrypted, err := encryptSettings(codec, plaintext.UID, plaintext)
if err != nil {
t.Fatal(err)
}
dto, err := newSettingsCacheDTO(encrypted)
if err != nil {
t.Fatal(err)
}
raw, err := encodeSettingsCache(dto)
if err != nil {
t.Fatal(err)
}
for _, secret := range []string{plaintext.APIKey, plaintext.ProviderAPIKeys["xai"], plaintext.ExaAPIKey} {
if strings.Contains(raw, secret) {
t.Fatalf("serialized cache contains plaintext secret %q: %s", secret, raw)
}
}
decoded, err := decodeSettingsCache(raw, plaintext.UID)
if err != nil {
t.Fatal(err)
}
got := decoded.settings()
if got.APIKey != encrypted.APIKey || got.ProviderAPIKeys["xai"] != encrypted.ProviderAPIKeys["xai"] || got.ExaAPIKey != encrypted.ExaAPIKey {
t.Fatalf("ciphertext did not survive DTO round-trip: %#v", got)
}
}
func TestSettingsCacheRejectsPlaintextSecrets(t *testing.T) {
_, err := newSettingsCacheDTO(&domain.UserSettings{UID: 42, APIKey: "plaintext"})
if err == nil {
t.Fatal("plaintext cache DTO was accepted")
}
}
func TestReadSettingsCacheCorruptionDeletesAndFallsBack(t *testing.T) {
cache := &fakeSettingsCache{value: `{"version":1,"found":true,"uid":42,"api_key":"plaintext"}`}
got, hit, err := readSettingsCache(context.Background(), cache, "settings-key", 42, repositoryTestCodec(t))
if err == nil || hit || got != nil {
t.Fatalf("corrupt lookup = %#v, hit=%v, err=%v", got, hit, err)
}
if !cache.deleted {
t.Fatal("corrupt cache entry was not deleted")
}
}
func TestReadSettingsCacheOutageFallsBackWithoutDeleting(t *testing.T) {
cache := &fakeSettingsCache{getErr: errors.New("redis unavailable")}
got, hit, err := readSettingsCache(context.Background(), cache, "settings-key", 42, repositoryTestCodec(t))
if err == nil || hit || got != nil {
t.Fatalf("outage lookup = %#v, hit=%v, err=%v", got, hit, err)
}
if cache.deleted {
t.Fatal("Redis outage should not trigger cache deletion")
}
}
func TestSettingsCacheKeyUsesConfiguredNamespace(t *testing.T) {
if got, want := settingsCacheKey("tenant:prod:v2", 42), "tenant:prod:v2:member:settings:{42}"; got != want {
t.Fatalf("settings cache key=%q, want %q", got, want)
}
if got, want := settingsCacheKey("", 42), "haixun:dev:v1:member:settings:{42}"; got != want {
t.Fatalf("default settings cache key=%q, want %q", got, want)
}
}

View File

@ -0,0 +1,78 @@
package repository
import (
"encoding/base64"
"strings"
"testing"
"apps/backend/internal/lib/secretbox"
"apps/backend/internal/module/member/domain"
)
func repositoryTestCodec(t *testing.T) *secretbox.Codec {
t.Helper()
c, err := secretbox.New("repo-test", base64.StdEncoding.EncodeToString([]byte("0123456789abcdef0123456789abcdef")))
if err != nil {
t.Fatal(err)
}
return c
}
func TestSettingsTransformEncryptsAllSecretsWithoutMutatingCaller(t *testing.T) {
codec := repositoryTestCodec(t)
original := &domain.UserSettings{
UID: 7, Provider: "xai", APIKey: "legacy-key", ExaAPIKey: "exa-key",
ProviderAPIKeys: map[string]string{"xai": "xai-key", "opencode-go": "oc-key", "empty": ""},
}
encrypted, err := encryptSettings(codec, 7, original)
if err != nil {
t.Fatal(err)
}
if original.APIKey != "legacy-key" || original.ExaAPIKey != "exa-key" || original.ProviderAPIKeys["xai"] != "xai-key" {
t.Fatal("caller settings were mutated")
}
for name, value := range map[string]string{
"api": encrypted.APIKey, "exa": encrypted.ExaAPIKey,
"xai": encrypted.ProviderAPIKeys["xai"], "opencode": encrypted.ProviderAPIKeys["opencode-go"],
} {
if !strings.HasPrefix(value, "enc:v1:repo-test:") {
t.Fatalf("%s was not encrypted: %q", name, value)
}
}
if encrypted.ProviderAPIKeys["empty"] != "" {
t.Fatal("empty key should remain empty")
}
decrypted, err := decryptSettings(codec, 7, encrypted)
if err != nil {
t.Fatal(err)
}
if decrypted.APIKey != "legacy-key" || decrypted.ExaAPIKey != "exa-key" || decrypted.ProviderAPIKeys["opencode-go"] != "oc-key" {
t.Fatalf("unexpected plaintext: %#v", decrypted)
}
}
func TestSettingsTransformLegacyAndMissingCodecBehavior(t *testing.T) {
legacy := &domain.UserSettings{UID: 8, Provider: "xai", APIKey: "plain", ProviderAPIKeys: map[string]string{"xai": "plain-map"}}
got, err := decryptSettings(nil, 8, legacy)
if err != nil || got.APIKey != "plain" || got.ProviderAPIKeys["xai"] != "plain-map" {
t.Fatalf("legacy read = %#v, %v", got, err)
}
if _, err := encryptSettings(nil, 8, legacy); err == nil {
t.Fatal("nonempty secret save without codec must fail")
}
if _, err := encryptSettings(nil, 8, &domain.UserSettings{UID: 8, ProviderAPIKeys: map[string]string{}}); err != nil {
t.Fatalf("empty settings save should remain constructible: %v", err)
}
codec := repositoryTestCodec(t)
encrypted, err := encryptSettings(codec, 8, legacy)
if err != nil {
t.Fatal(err)
}
if _, err := decryptSettings(nil, 8, encrypted); err == nil {
t.Fatal("encrypted read without codec must fail")
}
if _, err := decryptSettings(codec, 9, encrypted); err == nil {
t.Fatal("ciphertext must not move between users")
}
}

View File

@ -21,10 +21,10 @@ type AccountService struct {
Issuer tokenDomain.Issuer Issuer tokenDomain.Issuer
BcryptCost int BcryptCost int
// OAuth optional config (empty = dev mock path) // OAuth optional config (empty = dev mock path)
GoogleClientID string GoogleClientID string
LineClientID string LineClientID string
LineClientSecret string LineClientSecret string
LineRedirectURI string LineRedirectURI string
} }
// AuthService kept as type alias for svc/logic compatibility. // AuthService kept as type alias for svc/logic compatibility.
@ -295,16 +295,18 @@ func (s *AccountService) UpdateUserToken(ctx context.Context, loginID, platform,
// sync primary password on member if platform password identity // sync primary password on member if platform password identity
ident, err := s.Repo.FindIdentity(ctx, loginID, platform) ident, err := s.Repo.FindIdentity(ctx, loginID, platform)
if err != nil { if err != nil {
return nil return err
} }
m, err := s.Repo.FindByUID(ctx, ident.UID) m, err := s.Repo.FindByUID(ctx, ident.UID)
if err != nil { if err != nil {
return nil return err
} }
m.PasswordHash = hash m.PasswordHash = hash
m.UpdatedAt = domain.NowNano() m.UpdatedAt = domain.NowNano()
_ = s.Repo.UpdateMember(ctx, m) if err := s.Repo.UpdateMember(ctx, m); err != nil {
return nil return err
}
return s.Repo.RevokeAllRefreshByUID(ctx, ident.UID)
} }
func (s *AccountService) UpdateUserInfo(ctx context.Context, uid int64, patch *domain.UpdateUserInfoPatch) (*domain.Member, error) { func (s *AccountService) UpdateUserInfo(ctx context.Context, uid int64, patch *domain.UpdateUserInfoPatch) (*domain.Member, error) {
@ -370,7 +372,8 @@ func (s *AccountService) UpdateUserInfo(ctx context.Context, uid int64, patch *d
m.PhoneVerified = false m.PhoneVerified = false
m.PhoneVerifiedAt = nil m.PhoneVerifiedAt = nil
} }
if patch.NewPassword != nil && *patch.NewPassword != "" { passwordChanged := patch.NewPassword != nil && *patch.NewPassword != ""
if passwordChanged {
if patch.CurrentPassword == nil || !CheckPasswordHash(*patch.CurrentPassword, m.PasswordHash) { if patch.CurrentPassword == nil || !CheckPasswordHash(*patch.CurrentPassword, m.PasswordHash) {
return nil, domain.ErrBadPassword return nil, domain.ErrBadPassword
} }
@ -388,6 +391,11 @@ func (s *AccountService) UpdateUserInfo(ctx context.Context, uid int64, patch *d
if err := s.Repo.UpdateMember(ctx, m); err != nil { if err := s.Repo.UpdateMember(ctx, m); err != nil {
return nil, err return nil, err
} }
if passwordChanged {
if err := s.Repo.RevokeAllRefreshByUID(ctx, uid); err != nil {
return nil, err
}
}
return m, nil return m, nil
} }
@ -546,7 +554,9 @@ func (s *AccountService) BindVerifyPhone(ctx context.Context, uid int64, phone s
func (s *AccountService) GenerateCode(ctx context.Context, loginID, codeType string) (string, error) { func (s *AccountService) GenerateCode(ctx context.Context, loginID, codeType string) (string, error) {
code := randomDigits(6) code := randomDigits(6)
s.Repo.SetCode(codeType, strings.ToLower(loginID), code, 10*time.Minute) if !s.Repo.SetCode(codeType, strings.ToLower(loginID), code, 10*time.Minute) {
return "", domain.ErrCodeRateLimited
}
return code, nil return code, nil
} }
@ -578,11 +588,14 @@ func (s *AccountService) VerifyPlatformAuthResult(ctx context.Context, loginID,
// ---------- OAuth providers ---------- // ---------- OAuth providers ----------
func (s *AccountService) VerifyGoogleIDToken(ctx context.Context, idToken string) (subject, email, name string, err error) { func (s *AccountService) VerifyGoogleIDToken(ctx context.Context, idToken string) (subject, email, name string, err error) {
info, e := googleUserInfo(ctx, idToken) if strings.TrimSpace(s.GoogleClientID) == "" {
return "", "", "", domain.ErrOAuthFailed
}
info, e := googleIDTokenInfo(ctx, idToken, s.GoogleClientID)
if e != nil { if e != nil {
return "", "", "", domain.ErrOAuthFailed return "", "", "", domain.ErrOAuthFailed
} }
return info.ID, info.Email, info.Name, nil return info.Subject, info.Email, info.Name, nil
} }
func (s *AccountService) LineCodeToProfile(ctx context.Context, code, redirectURI string) (subject, displayName, email string, err error) { func (s *AccountService) LineCodeToProfile(ctx context.Context, code, redirectURI string) (subject, displayName, email string, err error) {
@ -636,11 +649,15 @@ func (s *AccountService) RequestPasswordReset(ctx context.Context, email string)
code := randomDigits(6) code := randomDigits(6)
// try identity or member — unknown email must error (product: 無此信箱不寄) // try identity or member — unknown email must error (product: 無此信箱不寄)
if _, err := s.Repo.FindByEmail(ctx, email); err == nil { if _, err := s.Repo.FindByEmail(ctx, email); err == nil {
s.Repo.SetCode(domain.CodeTypeForgetPassword, email, code, 30*time.Minute) if !s.Repo.SetCode(domain.CodeTypeForgetPassword, email, code, 30*time.Minute) {
return "", domain.ErrCodeRateLimited
}
return code, nil return code, nil
} }
if _, err := s.Repo.FindIdentity(ctx, email, domain.PlatformPassword); err == nil { if _, err := s.Repo.FindIdentity(ctx, email, domain.PlatformPassword); err == nil {
s.Repo.SetCode(domain.CodeTypeForgetPassword, email, code, 30*time.Minute) if !s.Repo.SetCode(domain.CodeTypeForgetPassword, email, code, 30*time.Minute) {
return "", domain.ErrCodeRateLimited
}
return code, nil return code, nil
} }
return "", domain.ErrEmailNotRegistered return "", domain.ErrEmailNotRegistered
@ -661,6 +678,9 @@ func (s *AccountService) ResetPassword(ctx context.Context, email, code, newPass
} }
// 3) 業務先寫入;成功才消費碼 // 3) 業務先寫入;成功才消費碼
if err := s.UpdateUserToken(ctx, email, domain.PlatformPassword, newPassword); err != nil { if err := s.UpdateUserToken(ctx, email, domain.PlatformPassword, newPassword); err != nil {
if !errors.Is(err, domain.ErrIdentityNotFound) {
return err
}
// seed / legacy member without identity row // seed / legacy member without identity row
hash, herr := HashPassword(newPassword, s.BcryptCost) hash, herr := HashPassword(newPassword, s.BcryptCost)
if herr != nil { if herr != nil {
@ -680,6 +700,9 @@ func (s *AccountService) ResetPassword(ctx context.Context, email, code, newPass
AccountType: domain.AccountTypeEmail, PasswordHash: hash, AccountType: domain.AccountTypeEmail, PasswordHash: hash,
CreatedAt: domain.NowNano(), UpdatedAt: domain.NowNano(), CreatedAt: domain.NowNano(), UpdatedAt: domain.NowNano(),
}) })
if err := s.Repo.RevokeAllRefreshByUID(ctx, m.UID); err != nil {
return err
}
} }
// 4) 成功才刪碼(兩種 kind 都清) // 4) 成功才刪碼(兩種 kind 都清)
_ = s.Repo.CheckCode(domain.CodeTypeForgetPassword, email, code) _ = s.Repo.CheckCode(domain.CodeTypeForgetPassword, email, code)
@ -697,9 +720,13 @@ func (s *AccountService) SendEmailCode(ctx context.Context, uid int64) (string,
if key == "" { if key == "" {
key = fmt.Sprintf("%d", uid) key = fmt.Sprintf("%d", uid)
} }
s.Repo.SetCode(domain.CodeTypeEmail, strings.ToLower(key), code, 30*time.Minute) if !s.Repo.SetCode(domain.CodeTypeEmail, strings.ToLower(key), code, 30*time.Minute) {
return "", domain.ErrCodeRateLimited
}
// also store by uid for VerifyEmailCode // also store by uid for VerifyEmailCode
s.Repo.SetCode(domain.CodeTypeEmail, fmt.Sprintf("%d", uid), code, 30*time.Minute) if key != fmt.Sprintf("%d", uid) {
_ = s.Repo.SetCode(domain.CodeTypeEmail, fmt.Sprintf("%d", uid), code, 30*time.Minute)
}
return code, nil return code, nil
} }

View File

@ -244,10 +244,11 @@ func (f *fakeRepo) UpdateIdentityPassword(ctx context.Context, loginID, platform
return nil return nil
} }
func (f *fakeRepo) SetCode(kind, key, code string, ttl time.Duration) { func (f *fakeRepo) SetCode(kind, key, code string, ttl time.Duration) bool {
f.mu.Lock() f.mu.Lock()
defer f.mu.Unlock() defer f.mu.Unlock()
f.codes[kind+":"+key] = code f.codes[kind+":"+key] = code
return true
} }
func (f *fakeRepo) CheckCode(kind, key, code string) bool { func (f *fakeRepo) CheckCode(kind, key, code string) bool {
f.mu.Lock() f.mu.Lock()

View File

@ -7,23 +7,26 @@ import (
"io" "io"
"net/http" "net/http"
"net/url" "net/url"
"strconv"
"strings" "strings"
"time" "time"
) )
type googleUserInfoResp struct { type googleIDTokenInfoResp struct {
ID string `json:"id"` Audience string `json:"aud"`
Email string `json:"email"` Issuer string `json:"iss"`
Name string `json:"name"` Subject string `json:"sub"`
Email string `json:"email"`
Name string `json:"name"`
Expires string `json:"exp"`
} }
func googleUserInfo(ctx context.Context, accessOrIDToken string) (*googleUserInfoResp, error) { func googleIDTokenInfo(ctx context.Context, idToken, clientID string) (*googleIDTokenInfoResp, error) {
// Prefer userinfo with Bearer access token (stand-alone style) endpoint := "https://oauth2.googleapis.com/tokeninfo?id_token=" + url.QueryEscape(strings.TrimSpace(idToken))
req, err := http.NewRequestWithContext(ctx, http.MethodGet, "https://www.googleapis.com/oauth2/v2/userinfo", nil) req, err := http.NewRequestWithContext(ctx, http.MethodGet, endpoint, nil)
if err != nil { if err != nil {
return nil, err return nil, err
} }
req.Header.Set("Authorization", "Bearer "+accessOrIDToken)
client := &http.Client{Timeout: 10 * time.Second} client := &http.Client{Timeout: 10 * time.Second}
resp, err := client.Do(req) resp, err := client.Do(req)
if err != nil { if err != nil {
@ -32,18 +35,32 @@ func googleUserInfo(ctx context.Context, accessOrIDToken string) (*googleUserInf
defer resp.Body.Close() defer resp.Body.Close()
body, _ := io.ReadAll(resp.Body) body, _ := io.ReadAll(resp.Body)
if resp.StatusCode != http.StatusOK { if resp.StatusCode != http.StatusOK {
return nil, fmt.Errorf("google userinfo status %d: %s", resp.StatusCode, string(body)) return nil, fmt.Errorf("google tokeninfo status %d", resp.StatusCode)
} }
var info googleUserInfoResp var info googleIDTokenInfoResp
if err := json.Unmarshal(body, &info); err != nil { if err := json.Unmarshal(body, &info); err != nil {
return nil, err return nil, err
} }
if info.ID == "" { if err := validateGoogleIDTokenInfo(&info, clientID, time.Now()); err != nil {
return nil, fmt.Errorf("google userinfo missing id") return nil, err
} }
return &info, nil return &info, nil
} }
func validateGoogleIDTokenInfo(info *googleIDTokenInfoResp, clientID string, now time.Time) error {
if info == nil || info.Subject == "" || info.Audience != strings.TrimSpace(clientID) {
return fmt.Errorf("google ID token audience or subject is invalid")
}
if info.Issuer != "accounts.google.com" && info.Issuer != "https://accounts.google.com" {
return fmt.Errorf("google ID token issuer is invalid")
}
expires, err := strconv.ParseInt(info.Expires, 10, 64)
if err != nil || expires <= now.Unix() {
return fmt.Errorf("google ID token is expired")
}
return nil
}
type lineTokenResp struct { type lineTokenResp struct {
AccessToken string `json:"access_token"` AccessToken string `json:"access_token"`
} }

View File

@ -0,0 +1,31 @@
package usecase
import (
"testing"
"time"
"github.com/stretchr/testify/require"
)
func TestValidateGoogleIDTokenInfo(t *testing.T) {
now := time.Unix(1_700_000_000, 0)
valid := &googleIDTokenInfoResp{
Audience: "client-123",
Issuer: "https://accounts.google.com",
Subject: "google-user",
Expires: "1700000060",
}
require.NoError(t, validateGoogleIDTokenInfo(valid, "client-123", now))
wrongAudience := *valid
wrongAudience.Audience = "other-client"
require.Error(t, validateGoogleIDTokenInfo(&wrongAudience, "client-123", now))
wrongIssuer := *valid
wrongIssuer.Issuer = "https://attacker.example"
require.Error(t, validateGoogleIDTokenInfo(&wrongIssuer, "client-123", now))
expired := *valid
expired.Expires = "1699999999"
require.Error(t, validateGoogleIDTokenInfo(&expired, "client-123", now))
}

View File

@ -13,14 +13,15 @@ var (
ErrIllegalStatus = errors.New("illegal outbox status") ErrIllegalStatus = errors.New("illegal outbox status")
ErrNotDue = errors.New("step not due") ErrNotDue = errors.New("step not due")
ErrRootBlocked = errors.New("root not published; reply blocked") ErrRootBlocked = errors.New("root not published; reply blocked")
ErrLeaseLost = errors.New("outbox step lease lost")
ErrBadURL = errors.New("invalid external url") ErrBadURL = errors.New("invalid external url")
ErrFormulaRemove = errors.New("generateFromFormula removed") ErrFormulaRemove = errors.New("generateFromFormula removed")
) )
const ( const (
PersonaEmpty = "empty" PersonaEmpty = "empty"
PersonaAnalyzing = "analyzing" PersonaAnalyzing = "analyzing"
PersonaReady = "ready" PersonaReady = "ready"
PlayDraft = "draft" PlayDraft = "draft"
PlayReady = "ready" PlayReady = "ready"
@ -47,9 +48,9 @@ const (
StepCancelled = "cancelled" StepCancelled = "cancelled"
StepBlocked = "blocked" StepBlocked = "blocked"
MentionPending = "pending" MentionPending = "pending"
MentionReplied = "replied" MentionReplied = "replied"
MentionSkipped = "skipped" MentionSkipped = "skipped"
) )
func NowNano() int64 { return time.Now().UTC().UnixNano() } func NowNano() int64 { return time.Now().UTC().UnixNano() }
@ -70,14 +71,14 @@ type Persona struct {
} }
type PersonaStyle struct { type PersonaStyle struct {
Draft PersonaDraft `bson:"draft" json:"draft"` Draft PersonaDraft `bson:"draft" json:"draft"`
DraftText string `bson:"draft_text" json:"draft_text"` DraftText string `bson:"draft_text" json:"draft_text"`
Source string `bson:"source" json:"source"` // manual|benchmark|seed Source string `bson:"source" json:"source"` // manual|benchmark|seed
BenchmarkUsername string `bson:"benchmark_username,omitempty" json:"benchmark_username,omitempty"` BenchmarkUsername string `bson:"benchmark_username,omitempty" json:"benchmark_username,omitempty"`
SourceLabel string `bson:"source_label,omitempty" json:"source_label,omitempty"` SourceLabel string `bson:"source_label,omitempty" json:"source_label,omitempty"`
SampleCount int `bson:"sample_count" json:"sample_count"` SampleCount int `bson:"sample_count" json:"sample_count"`
AnalyzedAt int64 `bson:"analyzed_at,omitempty" json:"analyzed_at,omitempty"` AnalyzedAt int64 `bson:"analyzed_at,omitempty" json:"analyzed_at,omitempty"`
SamplePreviews []string `bson:"sample_previews,omitempty" json:"sample_previews,omitempty"` SamplePreviews []string `bson:"sample_previews,omitempty" json:"sample_previews,omitempty"`
// Dimensions — 8D 風格摘要(分析後必填,前端「分析」頁卡片) // Dimensions — 8D 風格摘要(分析後必填,前端「分析」頁卡片)
Dimensions map[string]StyleDimension `bson:"dimensions,omitempty" json:"dimensions,omitempty"` Dimensions map[string]StyleDimension `bson:"dimensions,omitempty" json:"dimensions,omitempty"`
} }
@ -128,42 +129,42 @@ type ExternalTarget struct {
} }
type PlayStep struct { type PlayStep struct {
ID string `bson:"id" json:"id"` ID string `bson:"id" json:"id"`
SortOrder int `bson:"sort_order" json:"sort_order"` SortOrder int `bson:"sort_order" json:"sort_order"`
Kind string `bson:"kind" json:"kind"` Kind string `bson:"kind" json:"kind"`
AccountID string `bson:"account_id" json:"account_id"` AccountID string `bson:"account_id" json:"account_id"`
Text string `bson:"text" json:"text"` Text string `bson:"text" json:"text"`
DelayFromPreviousSec int `bson:"delay_from_previous_sec" json:"delay_from_previous_sec"` DelayFromPreviousSec int `bson:"delay_from_previous_sec" json:"delay_from_previous_sec"`
PersonaID string `bson:"persona_id,omitempty" json:"persona_id,omitempty"` PersonaID string `bson:"persona_id,omitempty" json:"persona_id,omitempty"`
BrandID string `bson:"brand_id,omitempty" json:"brand_id,omitempty"` BrandID string `bson:"brand_id,omitempty" json:"brand_id,omitempty"`
ImageURLs []string `bson:"image_urls,omitempty" json:"image_urls,omitempty"` ImageURLs []string `bson:"image_urls,omitempty" json:"image_urls,omitempty"`
} }
type Play struct { type Play struct {
ID string `bson:"_id" json:"id"` ID string `bson:"_id" json:"id"`
OwnerUID int64 `bson:"owner_uid" json:"owner_uid"` OwnerUID int64 `bson:"owner_uid" json:"owner_uid"`
Title string `bson:"title" json:"title"` Title string `bson:"title" json:"title"`
Topic string `bson:"topic" json:"topic"` Topic string `bson:"topic" json:"topic"`
Status string `bson:"status" json:"status"` Status string `bson:"status" json:"status"`
LeadAccountID string `bson:"lead_account_id" json:"lead_account_id"` LeadAccountID string `bson:"lead_account_id" json:"lead_account_id"`
CastAccountIDs []string `bson:"cast_account_ids" json:"cast_account_ids"` CastAccountIDs []string `bson:"cast_account_ids" json:"cast_account_ids"`
TargetOwnPostID string `bson:"target_own_post_id,omitempty" json:"target_own_post_id,omitempty"` TargetOwnPostID string `bson:"target_own_post_id,omitempty" json:"target_own_post_id,omitempty"`
TargetExternal *ExternalTarget `bson:"target_external,omitempty" json:"target_external,omitempty"` TargetExternal *ExternalTarget `bson:"target_external,omitempty" json:"target_external,omitempty"`
Steps []PlayStep `bson:"steps" json:"steps"` Steps []PlayStep `bson:"steps" json:"steps"`
ScheduleStartAt int64 `bson:"schedule_start_at" json:"schedule_start_at"` ScheduleStartAt int64 `bson:"schedule_start_at" json:"schedule_start_at"`
IntervalMinutes int `bson:"interval_minutes,omitempty" json:"interval_minutes,omitempty"` IntervalMinutes int `bson:"interval_minutes,omitempty" json:"interval_minutes,omitempty"`
CreatedAt int64 `bson:"created_at" json:"created_at"` CreatedAt int64 `bson:"created_at" json:"created_at"`
UpdatedAt int64 `bson:"updated_at" json:"updated_at"` UpdatedAt int64 `bson:"updated_at" json:"updated_at"`
} }
// Outbox // Outbox
type OutboxStep struct { type OutboxStep struct {
ID string `bson:"id" json:"id"` ID string `bson:"id" json:"id"`
StepID string `bson:"step_id" json:"step_id"` StepID string `bson:"step_id" json:"step_id"`
SortOrder int `bson:"sort_order" json:"sort_order"` SortOrder int `bson:"sort_order" json:"sort_order"`
Kind string `bson:"kind" json:"kind"` Kind string `bson:"kind" json:"kind"`
AccountID string `bson:"account_id" json:"account_id"` AccountID string `bson:"account_id" json:"account_id"`
Text string `bson:"text" json:"text"` Text string `bson:"text" json:"text"`
// TopicTag — 主貼可用的 Threads 話題標籤 // TopicTag — 主貼可用的 Threads 話題標籤
TopicTag string `bson:"topic_tag,omitempty" json:"topic_tag,omitempty"` TopicTag string `bson:"topic_tag,omitempty" json:"topic_tag,omitempty"`
// ImageURLs — 公開 https 圖網址Meta image_url 可抓;勿存 data URL // ImageURLs — 公開 https 圖網址Meta image_url 可抓;勿存 data URL
@ -176,15 +177,19 @@ type OutboxStep struct {
MediaID string `bson:"media_id,omitempty" json:"media_id,omitempty"` MediaID string `bson:"media_id,omitempty" json:"media_id,omitempty"`
// ReplyTo is parent media id (external or previous root) // ReplyTo is parent media id (external or previous root)
ReplyTo string `bson:"reply_to,omitempty" json:"reply_to,omitempty"` ReplyTo string `bson:"reply_to,omitempty" json:"reply_to,omitempty"`
// Lease fields are internal worker fencing state. A unique owner is used for
// every claim so an expired publisher cannot finalize a reclaimed step.
LeaseOwner string `bson:"lease_owner,omitempty" json:"-"`
LeaseExpiresAt int64 `bson:"lease_expires_at,omitempty" json:"-"`
} }
type OutboxBundle struct { type OutboxBundle struct {
ID string `bson:"_id" json:"id"` ID string `bson:"_id" json:"id"`
OwnerUID int64 `bson:"owner_uid" json:"owner_uid"` OwnerUID int64 `bson:"owner_uid" json:"owner_uid"`
PlayID string `bson:"play_id" json:"play_id"` PlayID string `bson:"play_id" json:"play_id"`
Title string `bson:"title" json:"title"` Title string `bson:"title" json:"title"`
Status string `bson:"status" json:"status"` Status string `bson:"status" json:"status"`
Steps []OutboxStep `bson:"steps" json:"steps"` Steps []OutboxStep `bson:"steps" json:"steps"`
// Root reply target for under-post schemes // Root reply target for under-post schemes
ReplyToMediaID string `bson:"reply_to_media_id,omitempty" json:"reply_to_media_id,omitempty"` ReplyToMediaID string `bson:"reply_to_media_id,omitempty" json:"reply_to_media_id,omitempty"`
CreatedAt int64 `bson:"created_at" json:"created_at"` CreatedAt int64 `bson:"created_at" json:"created_at"`

View File

@ -23,8 +23,14 @@ type Repository interface {
GetOutbox(ctx context.Context, id string) (*OutboxBundle, error) GetOutbox(ctx context.Context, id string) (*OutboxBundle, error)
ListOutbox(ctx context.Context, ownerUID int64) ([]*OutboxBundle, error) ListOutbox(ctx context.Context, ownerUID int64) ([]*OutboxBundle, error)
DeleteOutbox(ctx context.Context, id string) error DeleteOutbox(ctx context.Context, id string) error
// ListDueSteps returns bundles that may have due scheduled steps (all owners for worker).
ListAllOutbox(ctx context.Context) ([]*OutboxBundle, error) ListAllOutbox(ctx context.Context) ([]*OutboxBundle, error)
// ClaimOutboxStep atomically moves one due scheduled (or expired publishing)
// step to publishing and assigns its fencing lease.
ClaimOutboxStep(ctx context.Context, bundleID, stepID, leaseOwner string, now, leaseExpiresAt int64) (*OutboxBundle, error)
RenewOutboxStepLease(ctx context.Context, bundleID, stepID, leaseOwner string, now, leaseExpiresAt int64) error
FinishOutboxStep(ctx context.Context, bundleID, stepID, leaseOwner string, step *OutboxStep, now int64) error
RecomputeOutboxStatus(ctx context.Context, bundleID string, now int64) error
RetryOutboxStep(ctx context.Context, bundleID, stepID string, now int64) (*OutboxBundle, error)
// OwnPosts // OwnPosts
SaveOwnPost(ctx context.Context, p *OwnPost) error SaveOwnPost(ctx context.Context, p *OwnPost) error

View File

@ -5,6 +5,7 @@ import (
"encoding/json" "encoding/json"
"fmt" "fmt"
"io" "io"
"net"
"net/http" "net/http"
"net/url" "net/url"
"strings" "strings"
@ -26,9 +27,18 @@ func NewMeta(graphBase string) *MetaTransport {
if graphBase == "" { if graphBase == "" {
graphBase = "https://graph.threads.net" graphBase = "https://graph.threads.net"
} }
transport := http.DefaultTransport.(*http.Transport).Clone()
transport.DialContext = dialPublicAddress
client := &http.Client{Timeout: 90 * time.Second, Transport: transport}
client.CheckRedirect = func(req *http.Request, via []*http.Request) error {
if len(via) >= 5 {
return fmt.Errorf("too many redirects")
}
return validatePublicHTTPURL(req.Context(), req.URL)
}
return &MetaTransport{ return &MetaTransport{
GraphBase: strings.TrimRight(graphBase, "/"), GraphBase: strings.TrimRight(graphBase, "/"),
HTTPClient: &http.Client{Timeout: 90 * time.Second}, HTTPClient: client,
} }
} }
@ -192,7 +202,14 @@ func publicImageURLs(raw []string) []string {
// preflightImageURL — 確認公開 URL 對匿名 GET 可讀Meta crawler 同樣匿名) // preflightImageURL — 確認公開 URL 對匿名 GET 可讀Meta crawler 同樣匿名)
func (t *MetaTransport) preflightImageURL(ctx context.Context, imageURL string) error { func (t *MetaTransport) preflightImageURL(ctx context.Context, imageURL string) error {
req, err := http.NewRequestWithContext(ctx, http.MethodHead, imageURL, nil) parsed, err := url.Parse(imageURL)
if err != nil {
return err
}
if err := validatePublicHTTPURL(ctx, parsed); err != nil {
return err
}
req, err := http.NewRequestWithContext(ctx, http.MethodHead, parsed.String(), nil)
if err != nil { if err != nil {
return err return err
} }
@ -226,6 +243,53 @@ func (t *MetaTransport) preflightImageURL(ctx context.Context, imageURL string)
return nil return nil
} }
func validatePublicHTTPURL(ctx context.Context, u *url.URL) error {
if u == nil || (u.Scheme != "http" && u.Scheme != "https") || u.Hostname() == "" {
return fmt.Errorf("invalid public HTTP URL")
}
if u.User != nil {
return fmt.Errorf("URL credentials are not allowed")
}
addrs, err := net.DefaultResolver.LookupIPAddr(ctx, u.Hostname())
if err != nil {
return fmt.Errorf("cannot resolve image host: %w", err)
}
if len(addrs) == 0 {
return fmt.Errorf("image host has no address")
}
for _, addr := range addrs {
if !isPublicIP(addr.IP) {
return fmt.Errorf("image URL resolves to a non-public address")
}
}
return nil
}
func dialPublicAddress(ctx context.Context, network, address string) (net.Conn, error) {
host, port, err := net.SplitHostPort(address)
if err != nil {
return nil, err
}
addrs, err := net.DefaultResolver.LookupIPAddr(ctx, host)
if err != nil {
return nil, err
}
for _, addr := range addrs {
if !isPublicIP(addr.IP) {
return nil, fmt.Errorf("refusing to connect to a non-public address")
}
}
if len(addrs) == 0 {
return nil, fmt.Errorf("host has no address")
}
dialer := &net.Dialer{Timeout: 15 * time.Second, KeepAlive: 30 * time.Second}
return dialer.DialContext(ctx, network, net.JoinHostPort(addrs[0].IP.String(), port))
}
func isPublicIP(ip net.IP) bool {
return ip.IsGlobalUnicast() && !ip.IsPrivate() && !ip.IsLoopback() && !ip.IsLinkLocalUnicast()
}
func (t *MetaTransport) createContainer(ctx context.Context, createURL string, form url.Values) (string, error) { func (t *MetaTransport) createContainer(ctx context.Context, createURL string, form url.Values) (string, error) {
cBody, cStatus, err := t.postForm(ctx, createURL, form) cBody, cStatus, err := t.postForm(ctx, createURL, form)
if err != nil { if err != nil {

View File

@ -0,0 +1,32 @@
package publish
import (
"context"
"net/url"
"testing"
"github.com/stretchr/testify/require"
)
func TestValidatePublicHTTPURLRejectsPrivateTargets(t *testing.T) {
t.Parallel()
for _, raw := range []string{
"http://127.0.0.1/image.jpg",
"http://10.0.0.1/image.jpg",
"http://169.254.169.254/latest/meta-data",
"http://[::1]/image.jpg",
} {
u, err := url.Parse(raw)
require.NoError(t, err)
require.Error(t, validatePublicHTTPURL(context.Background(), u), raw)
}
}
func TestValidatePublicHTTPURLRejectsCredentialsAndUnsupportedSchemes(t *testing.T) {
t.Parallel()
for _, raw := range []string{"file:///etc/passwd", "https://user:pass@example.com/a.jpg"} {
u, err := url.Parse(raw)
require.NoError(t, err)
require.Error(t, validatePublicHTTPURL(context.Background(), u), raw)
}
}

View File

@ -9,14 +9,14 @@ import (
) )
type MemoryStore struct { type MemoryStore struct {
mu sync.Mutex mu sync.Mutex
personas map[string]*domain.Persona personas map[string]*domain.Persona
active map[int64]string active map[int64]string
plays map[string]*domain.Play plays map[string]*domain.Play
outbox map[string]*domain.OutboxBundle outbox map[string]*domain.OutboxBundle
ownPosts map[string]*domain.OwnPost ownPosts map[string]*domain.OwnPost
mentions map[string]*domain.Mention mentions map[string]*domain.Mention
syncMeta map[int64]*domain.SyncMeta syncMeta map[int64]*domain.SyncMeta
} }
func NewMemory() *MemoryStore { func NewMemory() *MemoryStore {
@ -190,13 +190,156 @@ func (s *MemoryStore) ListOutbox(_ context.Context, ownerUID int64) ([]*domain.O
func (s *MemoryStore) DeleteOutbox(_ context.Context, id string) error { func (s *MemoryStore) DeleteOutbox(_ context.Context, id string) error {
s.mu.Lock() s.mu.Lock()
defer s.mu.Unlock() defer s.mu.Unlock()
if _, ok := s.outbox[id]; !ok { b, ok := s.outbox[id]
if !ok {
return domain.ErrNotFound return domain.ErrNotFound
} }
for _, step := range b.Steps {
if step.Status == domain.StepPublishing {
return domain.ErrIllegalStatus
}
}
delete(s.outbox, id) delete(s.outbox, id)
return nil return nil
} }
func (s *MemoryStore) ClaimOutboxStep(_ context.Context, bundleID, stepID, leaseOwner string, now, leaseExpiresAt int64) (*domain.OutboxBundle, error) {
s.mu.Lock()
defer s.mu.Unlock()
b, ok := s.outbox[bundleID]
if !ok {
return nil, domain.ErrNotFound
}
for i := range b.Steps {
step := &b.Steps[i]
if step.ID != stepID {
continue
}
due := step.Status == domain.StepScheduled && step.ScheduledAt <= now
stale := step.Status == domain.StepPublishing && step.LeaseExpiresAt <= now
if !due && !stale {
return nil, domain.ErrIllegalStatus
}
step.Status = domain.StepPublishing
step.Error = ""
step.LeaseOwner = leaseOwner
step.LeaseExpiresAt = leaseExpiresAt
b.Status = domain.OBActive
b.UpdatedAt = now
return copyOutbox(b), nil
}
return nil, domain.ErrNotFound
}
func (s *MemoryStore) RenewOutboxStepLease(_ context.Context, bundleID, stepID, leaseOwner string, now, leaseExpiresAt int64) error {
s.mu.Lock()
defer s.mu.Unlock()
b, ok := s.outbox[bundleID]
if !ok {
return domain.ErrNotFound
}
for i := range b.Steps {
step := &b.Steps[i]
if step.ID == stepID && step.Status == domain.StepPublishing && step.LeaseOwner == leaseOwner && step.LeaseExpiresAt > now {
step.LeaseExpiresAt = leaseExpiresAt
b.UpdatedAt = now
return nil
}
}
return domain.ErrLeaseLost
}
func (s *MemoryStore) FinishOutboxStep(_ context.Context, bundleID, stepID, leaseOwner string, result *domain.OutboxStep, now int64) error {
s.mu.Lock()
defer s.mu.Unlock()
b, ok := s.outbox[bundleID]
if !ok {
return domain.ErrNotFound
}
for i := range b.Steps {
step := &b.Steps[i]
if step.ID != stepID || step.Status != domain.StepPublishing || step.LeaseOwner != leaseOwner || step.LeaseExpiresAt <= now {
continue
}
*step = *result
step.LeaseOwner = ""
step.LeaseExpiresAt = 0
if step.Kind == domain.StepRoot {
for j := range b.Steps {
reply := &b.Steps[j]
if reply.Kind != domain.StepReply || reply.Status == domain.StepPublished || reply.Status == domain.StepPublishing {
continue
}
if step.Status == domain.StepFailed {
reply.Status, reply.Error = domain.StepBlocked, "waiting for root success"
} else if step.Status == domain.StepPublished && reply.Status == domain.StepBlocked {
reply.Status, reply.Error = domain.StepScheduled, ""
}
}
}
b.UpdatedAt = now
return nil
}
return domain.ErrLeaseLost
}
func (s *MemoryStore) RecomputeOutboxStatus(_ context.Context, bundleID string, now int64) error {
s.mu.Lock()
defer s.mu.Unlock()
b, ok := s.outbox[bundleID]
if !ok {
return domain.ErrNotFound
}
b.Status = outboxStatus(b.Steps)
b.UpdatedAt = now
return nil
}
func (s *MemoryStore) RetryOutboxStep(_ context.Context, bundleID, stepID string, now int64) (*domain.OutboxBundle, error) {
s.mu.Lock()
defer s.mu.Unlock()
b, ok := s.outbox[bundleID]
if !ok {
return nil, domain.ErrNotFound
}
for i := range b.Steps {
step := &b.Steps[i]
if step.ID != stepID {
continue
}
if step.Status != domain.StepFailed {
return nil, domain.ErrIllegalStatus
}
step.Status, step.Error, step.ScheduledAt = domain.StepScheduled, "", now
step.LeaseOwner, step.LeaseExpiresAt = "", 0
b.Status, b.UpdatedAt = domain.OBScheduling, now
return copyOutbox(b), nil
}
return nil, domain.ErrNotFound
}
func outboxStatus(steps []domain.OutboxStep) string {
if len(steps) == 0 {
return domain.OBCancelled
}
allPublished, anyFailed, anyPublishing := true, false, false
for _, step := range steps {
allPublished = allPublished && step.Status == domain.StepPublished
anyFailed = anyFailed || step.Status == domain.StepFailed
anyPublishing = anyPublishing || step.Status == domain.StepPublishing
}
if allPublished {
return domain.OBCompleted
}
if anyFailed {
return domain.OBPartial
}
if anyPublishing {
return domain.OBActive
}
return domain.OBScheduling
}
func (s *MemoryStore) ListAllOutbox(_ context.Context) ([]*domain.OutboxBundle, error) { func (s *MemoryStore) ListAllOutbox(_ context.Context) ([]*domain.OutboxBundle, error) {
s.mu.Lock() s.mu.Lock()
defer s.mu.Unlock() defer s.mu.Unlock()

View File

@ -8,6 +8,7 @@ import (
"github.com/zeromicro/go-zero/core/stores/mon" "github.com/zeromicro/go-zero/core/stores/mon"
"go.mongodb.org/mongo-driver/bson" "go.mongodb.org/mongo-driver/bson"
"go.mongodb.org/mongo-driver/mongo"
"go.mongodb.org/mongo-driver/mongo/options" "go.mongodb.org/mongo-driver/mongo/options"
) )
@ -147,16 +148,132 @@ func (s *MonStore) ListOutbox(ctx context.Context, ownerUID int64) ([]*domain.Ou
} }
func (s *MonStore) DeleteOutbox(ctx context.Context, id string) error { func (s *MonStore) DeleteOutbox(ctx context.Context, id string) error {
res, err := s.outbox.DeleteOne(ctx, bson.M{"_id": id}) res, err := s.outbox.DeleteOne(ctx, bson.M{
"_id": id, "steps": bson.M{"$not": bson.M{"$elemMatch": bson.M{"status": domain.StepPublishing}}},
})
if err != nil { if err != nil {
return err return err
} }
if res == 0 { if res == 0 {
if _, getErr := s.GetOutbox(ctx, id); getErr == nil {
return domain.ErrIllegalStatus
} else {
return getErr
}
}
return nil
}
func (s *MonStore) ClaimOutboxStep(ctx context.Context, bundleID, stepID, leaseOwner string, now, leaseExpiresAt int64) (*domain.OutboxBundle, error) {
claimable := bson.M{"id": stepID, "$or": bson.A{
bson.M{"status": domain.StepScheduled, "scheduled_at": bson.M{"$lte": now}},
bson.M{"status": domain.StepPublishing, "$or": bson.A{
bson.M{"lease_expires_at": bson.M{"$lte": now}}, bson.M{"lease_expires_at": bson.M{"$exists": false}},
}},
}}
filter := bson.M{"_id": bundleID, "steps": bson.M{"$elemMatch": claimable}}
update := bson.M{"$set": bson.M{
"steps.$.status": domain.StepPublishing, "steps.$.error": "",
"steps.$.lease_owner": leaseOwner, "steps.$.lease_expires_at": leaseExpiresAt,
"status": domain.OBActive, "updated_at": now,
}}
var b domain.OutboxBundle
err := s.outbox.FindOneAndUpdate(ctx, &b, filter, update, options.FindOneAndUpdate().SetReturnDocument(options.After))
if err == mon.ErrNotFound {
return nil, domain.ErrIllegalStatus
}
return &b, err
}
func (s *MonStore) RenewOutboxStepLease(ctx context.Context, bundleID, stepID, leaseOwner string, now, leaseExpiresAt int64) error {
filter := ownedPublishingStep(bundleID, stepID, leaseOwner, now)
res, err := s.outbox.UpdateOne(ctx, filter, bson.M{"$set": bson.M{
"steps.$.lease_expires_at": leaseExpiresAt, "updated_at": now,
}})
if err != nil {
return err
}
if res.MatchedCount == 0 {
return domain.ErrLeaseLost
}
return nil
}
func (s *MonStore) FinishOutboxStep(ctx context.Context, bundleID, stepID, leaseOwner string, step *domain.OutboxStep, now int64) error {
filter := ownedPublishingStep(bundleID, stepID, leaseOwner, now)
set := bson.M{
"steps.$[claimed].status": step.Status, "steps.$[claimed].error": step.Error,
"steps.$[claimed].published_at": step.PublishedAt, "steps.$[claimed].media_id": step.MediaID,
"steps.$[claimed].image_urls": step.ImageURLs, "updated_at": now,
}
unset := bson.M{"steps.$[claimed].lease_owner": "", "steps.$[claimed].lease_expires_at": ""}
filters := options.ArrayFilters{Filters: []interface{}{bson.M{
"claimed.id": stepID, "claimed.status": domain.StepPublishing, "claimed.lease_owner": leaseOwner,
}}}
if step.Kind == domain.StepRoot && (step.Status == domain.StepFailed || step.Status == domain.StepPublished) {
replyFilter := bson.M{"reply.kind": domain.StepReply, "reply.status": bson.M{"$nin": bson.A{domain.StepPublished, domain.StepPublishing}}}
if step.Status == domain.StepFailed {
set["steps.$[reply].status"], set["steps.$[reply].error"] = domain.StepBlocked, "waiting for root success"
} else {
replyFilter["reply.status"] = domain.StepBlocked
set["steps.$[reply].status"], set["steps.$[reply].error"] = domain.StepScheduled, ""
}
filters.Filters = append(filters.Filters, replyFilter)
}
res, err := s.outbox.UpdateOne(ctx, filter, bson.M{"$set": set, "$unset": unset}, options.Update().SetArrayFilters(filters))
if err != nil {
return err
}
if res.MatchedCount == 0 {
return domain.ErrLeaseLost
}
return nil
}
func ownedPublishingStep(bundleID, stepID, leaseOwner string, now int64) bson.M {
return bson.M{"_id": bundleID, "steps": bson.M{"$elemMatch": bson.M{
"id": stepID, "status": domain.StepPublishing, "lease_owner": leaseOwner,
"lease_expires_at": bson.M{"$gt": now},
}}}
}
func (s *MonStore) RecomputeOutboxStatus(ctx context.Context, bundleID string, now int64) error {
statuses := bson.M{"$map": bson.M{"input": "$steps", "as": "step", "in": "$$step.status"}}
pipeline := mongo.Pipeline{bson.D{{Key: "$set", Value: bson.M{
"status": bson.M{"$switch": bson.M{"branches": bson.A{
bson.M{"case": bson.M{"$eq": bson.A{statuses, bson.A{}}}, "then": domain.OBCancelled},
bson.M{"case": bson.M{"$allElementsTrue": bson.A{bson.M{"$map": bson.M{"input": "$steps", "as": "step", "in": bson.M{"$eq": bson.A{"$$step.status", domain.StepPublished}}}}}}, "then": domain.OBCompleted},
bson.M{"case": bson.M{"$in": bson.A{domain.StepFailed, statuses}}, "then": domain.OBPartial},
bson.M{"case": bson.M{"$in": bson.A{domain.StepPublishing, statuses}}, "then": domain.OBActive},
}, "default": domain.OBScheduling}}, "updated_at": now,
}}}}
res, err := s.outbox.UpdateOne(ctx, bson.M{"_id": bundleID}, pipeline)
if err != nil {
return err
}
if res.MatchedCount == 0 {
return domain.ErrNotFound return domain.ErrNotFound
} }
return nil return nil
} }
func (s *MonStore) RetryOutboxStep(ctx context.Context, bundleID, stepID string, now int64) (*domain.OutboxBundle, error) {
filter := bson.M{"_id": bundleID, "steps": bson.M{"$elemMatch": bson.M{"id": stepID, "status": domain.StepFailed}}}
update := bson.M{
"$set": bson.M{
"steps.$.status": domain.StepScheduled, "steps.$.error": "", "steps.$.scheduled_at": now,
"status": domain.OBScheduling, "updated_at": now,
},
"$unset": bson.M{"steps.$.lease_owner": "", "steps.$.lease_expires_at": ""},
}
var b domain.OutboxBundle
err := s.outbox.FindOneAndUpdate(ctx, &b, filter, update, options.FindOneAndUpdate().SetReturnDocument(options.After))
if err == mon.ErrNotFound {
return nil, domain.ErrIllegalStatus
}
return &b, err
}
func (s *MonStore) ListAllOutbox(ctx context.Context) ([]*domain.OutboxBundle, error) { func (s *MonStore) ListAllOutbox(ctx context.Context) ([]*domain.OutboxBundle, error) {
var list []*domain.OutboxBundle var list []*domain.OutboxBundle
err := s.outbox.Find(ctx, &list, bson.M{}, options.Find().SetLimit(500)) err := s.outbox.Find(ctx, &list, bson.M{}, options.Find().SetLimit(500))

View File

@ -0,0 +1,81 @@
package repository
import (
"context"
"sync"
"testing"
"time"
"apps/backend/internal/module/studio/domain"
"github.com/stretchr/testify/require"
)
func TestOutboxClaimIsAtomicPerStep(t *testing.T) {
repo := NewMemory()
now := domain.NowNano()
require.NoError(t, repo.SaveOutbox(context.Background(), &domain.OutboxBundle{
ID: "bundle", Steps: []domain.OutboxStep{{ID: "step", Status: domain.StepScheduled, ScheduledAt: now}},
}))
const workers = 16
start := make(chan struct{})
results := make(chan error, workers)
var wg sync.WaitGroup
for i := 0; i < workers; i++ {
wg.Add(1)
go func(owner string) {
defer wg.Done()
<-start
_, err := repo.ClaimOutboxStep(context.Background(), "bundle", "step", owner, now, now+int64(time.Minute))
results <- err
}(time.Duration(i).String())
}
close(start)
wg.Wait()
close(results)
claimed := 0
for err := range results {
if err == nil {
claimed++
} else {
require.ErrorIs(t, err, domain.ErrIllegalStatus)
}
}
require.Equal(t, 1, claimed)
}
func TestOutboxExpiredClaimFencesStaleOwner(t *testing.T) {
repo := NewMemory()
now := domain.NowNano()
require.NoError(t, repo.SaveOutbox(context.Background(), &domain.OutboxBundle{
ID: "bundle", Steps: []domain.OutboxStep{{ID: "step", Kind: domain.StepRoot, Status: domain.StepScheduled, ScheduledAt: now}},
}))
first, err := repo.ClaimOutboxStep(context.Background(), "bundle", "step", "worker-a", now, now+10)
require.NoError(t, err)
second, err := repo.ClaimOutboxStep(context.Background(), "bundle", "step", "worker-b", now+11, now+int64(time.Minute))
require.NoError(t, err)
staleResult := first.Steps[0]
staleResult.Status = domain.StepPublished
err = repo.FinishOutboxStep(context.Background(), "bundle", "step", "worker-a", &staleResult, now+12)
require.ErrorIs(t, err, domain.ErrLeaseLost)
result := second.Steps[0]
result.Status, result.MediaID = domain.StepPublished, "media-b"
require.NoError(t, repo.FinishOutboxStep(context.Background(), "bundle", "step", "worker-b", &result, now+12))
got, err := repo.GetOutbox(context.Background(), "bundle")
require.NoError(t, err)
require.Equal(t, "media-b", got.Steps[0].MediaID)
}
func TestDeleteOutboxRejectsPublishingStep(t *testing.T) {
repo := NewMemory()
require.NoError(t, repo.SaveOutbox(context.Background(), &domain.OutboxBundle{
ID: "bundle", Steps: []domain.OutboxStep{{ID: "step", Status: domain.StepPublishing}},
}))
require.ErrorIs(t, repo.DeleteOutbox(context.Background(), "bundle"), domain.ErrIllegalStatus)
_, err := repo.GetOutbox(context.Background(), "bundle")
require.NoError(t, err)
}

View File

@ -22,6 +22,21 @@ type memAccounts struct {
byID map[string]*threadsDomain.Account byID map[string]*threadsDomain.Account
} }
type blockingPublishTransport struct {
entered chan struct{}
release chan struct{}
}
func (t *blockingPublishTransport) Publish(ctx context.Context, _ domain.PublishRequest) (*domain.PublishResult, error) {
close(t.entered)
select {
case <-t.release:
return &domain.PublishResult{MediaID: "renewed-media"}, nil
case <-ctx.Done():
return nil, ctx.Err()
}
}
func (m *memAccounts) Get(_ context.Context, id string) (*threadsDomain.Account, error) { func (m *memAccounts) Get(_ context.Context, id string) (*threadsDomain.Account, error) {
a, ok := m.byID[id] a, ok := m.byID[id]
if !ok { if !ok {
@ -242,7 +257,7 @@ func TestPL_01_SaveGet(t *testing.T) {
addAccount(acc, uid, "acc1", "lead") addAccount(acc, uid, "acc1", "lead")
play, err := svc.SavePlay(context.Background(), uid, &domain.Play{ play, err := svc.SavePlay(context.Background(), uid, &domain.Play{
Title: "方案A", Topic: "主題", LeadAccountID: "acc1", Title: "方案A", Topic: "主題", LeadAccountID: "acc1",
Steps: []domain.PlayStep{{Kind: domain.StepRoot, AccountID: "acc1", Text: "主貼"}}, Steps: []domain.PlayStep{{Kind: domain.StepRoot, AccountID: "acc1", Text: "主貼"}},
ScheduleStartAt: domain.NowNano(), ScheduleStartAt: domain.NowNano(),
}) })
require.NoError(t, err) require.NoError(t, err)
@ -642,6 +657,35 @@ func TestOB_10_RemoveStopsWorker(t *testing.T) {
require.Equal(t, 0, tp.CallCount()) require.Equal(t, 0, tp.CallCount())
} }
func TestOB_PublishRenewsStepLease(t *testing.T) {
repo := repository.NewMemory()
transport := &blockingPublishTransport{entered: make(chan struct{}), release: make(chan struct{})}
svc := usecase.New(repo, transport)
svc.OutboxLeaseDuration = 90 * time.Millisecond
now := domain.NowNano()
require.NoError(t, repo.SaveOutbox(context.Background(), &domain.OutboxBundle{
ID: "lease-bundle", Steps: []domain.OutboxStep{{ID: "lease-step", Kind: domain.StepRoot, Status: domain.StepScheduled, ScheduledAt: now}},
}))
done := make(chan error, 1)
go func() {
_, err := svc.ProcessDueSteps(context.Background(), now)
done <- err
}()
<-transport.entered
claimed, err := repo.GetOutbox(context.Background(), "lease-bundle")
require.NoError(t, err)
initialExpiry := claimed.Steps[0].LeaseExpiresAt
require.Eventually(t, func() bool {
got, getErr := repo.GetOutbox(context.Background(), "lease-bundle")
return getErr == nil && got.Steps[0].LeaseExpiresAt > initialExpiry
}, time.Second, 10*time.Millisecond)
close(transport.release)
require.NoError(t, <-done)
got, err := repo.GetOutbox(context.Background(), "lease-bundle")
require.NoError(t, err)
require.Equal(t, domain.StepPublished, got.Steps[0].Status)
}
func TestOB_11_CrossUid(t *testing.T) { func TestOB_11_CrossUid(t *testing.T) {
svc, _, acc := newStudio() svc, _, acc := newStudio()
uid1, uid2 := int64(4_003_011), int64(4_003_012) uid1, uid2 := int64(4_003_011), int64(4_003_012)

View File

@ -4,6 +4,7 @@ import (
"context" "context"
"crypto/rand" "crypto/rand"
"encoding/json" "encoding/json"
"errors"
"fmt" "fmt"
"net/url" "net/url"
"strings" "strings"
@ -49,7 +50,7 @@ type ThreadsMediaSource interface {
// Fetched* 與 provider 解耦,避免 studio 依賴 provider 包。 // Fetched* 與 provider 解耦,避免 studio 依賴 provider 包。
type FetchedThread struct { type FetchedThread struct {
ID, Text, MediaType, MediaURL, ThumbnailURL, Permalink, Shortcode, TopicTag, Username string ID, Text, MediaType, MediaURL, ThumbnailURL, Permalink, Shortcode, TopicTag, Username string
PublishedAt int64 // unix ns PublishedAt int64 // unix ns
} }
type FetchedInsights struct { type FetchedInsights struct {
@ -67,8 +68,8 @@ type FetchedReply struct {
// FetchedMention — Graph /{user-id}/mentions // FetchedMention — Graph /{user-id}/mentions
type FetchedMention struct { type FetchedMention struct {
ID, Text, Username, Permalink, RootPostID, ParentID string ID, Text, Username, Permalink, RootPostID, ParentID string
IsReply, IsQuotePost bool IsReply, IsQuotePost bool
PublishedAt int64 PublishedAt int64
} }
// AIKeySource resolves provider/model/apiKey for a member (settings + platform). // AIKeySource resolves provider/model/apiKey for a member (settings + platform).
@ -107,10 +108,16 @@ type Service struct {
StoragePublicBase string StoragePublicBase string
// optional key resolver path via Usage.PrepareCall + Static/Settings // optional key resolver path via Usage.PrepareCall + Static/Settings
KeyModeDefault string // if Usage nil, use this for tests KeyModeDefault string // if Usage nil, use this for tests
OutboxWorkerID string
// OutboxLeaseDuration is configurable for focused lease-renewal tests.
OutboxLeaseDuration time.Duration
} }
func New(repo domain.Repository, transport publish.Transport) *Service { func New(repo domain.Repository, transport publish.Transport) *Service {
return &Service{Repo: repo, Transport: transport} return &Service{
Repo: repo, Transport: transport,
OutboxWorkerID: "studio-" + uuid.NewString(), OutboxLeaseDuration: 3 * time.Minute,
}
} }
// ---------- Personas (PE) ---------- // ---------- Personas (PE) ----------
@ -1110,11 +1117,15 @@ func (s *Service) RemoveOutbox(ctx context.Context, ownerUID int64, id string) e
if err != nil { if err != nil {
return err return err
} }
// 未發出/失敗取消:清掉暫存圖 // Delete has a repository-level no-publishing guard. Clean up only after it
// succeeds, otherwise a publisher may still need these URLs.
if err := s.Repo.DeleteOutbox(ctx, id); err != nil {
return err
}
for i := range b.Steps { for i := range b.Steps {
s.cleanupEphemeralImages(ctx, b.Steps[i].ImageURLs) s.cleanupEphemeralImages(ctx, b.Steps[i].ImageURLs)
} }
return s.Repo.DeleteOutbox(ctx, id) return nil
} }
func (s *Service) RetryStep(ctx context.Context, ownerUID int64, bundleID, stepID string) (*domain.OutboxBundle, error) { func (s *Service) RetryStep(ctx context.Context, ownerUID int64, bundleID, stepID string) (*domain.OutboxBundle, error) {
@ -1143,24 +1154,7 @@ func (s *Service) RetryStep(ctx context.Context, ownerUID int64, bundleID, stepI
return nil, domain.ErrRootBlocked return nil, domain.ErrRootBlocked
} }
} }
st.Status = domain.StepScheduled return s.Repo.RetryOutboxStep(ctx, bundleID, stepID, domain.NowNano())
st.Error = ""
st.ScheduledAt = domain.NowNano() // due now
// unblock later replies if root ok
if root := findRoot(b); root != nil && root.Status == domain.StepPublished {
for i := range b.Steps {
if b.Steps[i].Kind == domain.StepReply && b.Steps[i].Status == domain.StepBlocked {
b.Steps[i].Status = domain.StepScheduled
b.Steps[i].Error = ""
}
}
}
b.Status = recomputeBundle(b.Steps)
b.UpdatedAt = domain.NowNano()
if err := s.Repo.SaveOutbox(ctx, b); err != nil {
return nil, err
}
return b, nil
} }
// ProcessDueSteps is the worker tick — publishes via Transport only. // ProcessDueSteps is the worker tick — publishes via Transport only.
@ -1183,138 +1177,173 @@ func (s *Service) ProcessDueSteps(ctx context.Context, now int64) (published int
} }
func (s *Service) processBundle(ctx context.Context, b *domain.OutboxBundle, now int64) (int, error) { func (s *Service) processBundle(ctx context.Context, b *domain.OutboxBundle, now int64) (int, error) {
// skip removed (already not in list)
count := 0 count := 0
// if root failed, do not publish replies
root := findRoot(b)
rootFailed := root != nil && root.Status == domain.StepFailed
rootPublished := root != nil && root.Status == domain.StepPublished
// worker 重啟/中斷會留下 publishing逾時後收回可重試
const publishingStaleNs = int64(3 * time.Minute)
for i := range b.Steps { for i := range b.Steps {
st := &b.Steps[i] st := &b.Steps[i]
if st.Status == domain.StepPublishing && b.UpdatedAt > 0 && now-b.UpdatedAt > publishingStaleNs { claimable := st.Status == domain.StepScheduled && st.ScheduledAt <= now
logx.Infof("outbox reclaim stale publishing step id=%s bundle=%s", st.ID, b.ID) stale := st.Status == domain.StepPublishing && st.LeaseExpiresAt <= now
st.Status = domain.StepScheduled if !claimable && !stale {
st.Error = ""
}
}
// sort by sort_order
for i := range b.Steps {
st := &b.Steps[i]
if st.Status != domain.StepScheduled {
continue
}
if st.ScheduledAt > now {
// OB-01: not due
continue continue
} }
root := findRoot(b)
if st.Kind == domain.StepReply { if st.Kind == domain.StepReply {
if rootFailed { if root != nil && root.Status != domain.StepPublished {
st.Status = domain.StepBlocked
st.Error = "root failed; reply blocked"
continue
}
if root != nil && !rootPublished && root.Status != domain.StepPublished {
// wait for root
continue continue
} }
} }
// publish via transport
leaseDuration := s.OutboxLeaseDuration
if leaseDuration <= 0 {
leaseDuration = 3 * time.Minute
}
claimOwner := s.OutboxWorkerID + ":" + uuid.NewString()
claimNow := domain.NowNano()
claimed, claimErr := s.Repo.ClaimOutboxStep(ctx, b.ID, st.ID, claimOwner, now, claimNow+int64(leaseDuration))
if claimErr != nil {
if claimErr == domain.ErrIllegalStatus || claimErr == domain.ErrNotFound {
continue
}
return count, claimErr
}
if stale {
logx.Infof("outbox reclaimed expired publishing lease step=%s bundle=%s", st.ID, b.ID)
}
claimedStep := findOutboxStep(claimed, st.ID)
if claimedStep == nil {
return count, domain.ErrNotFound
}
result := *claimedStep
result.LeaseOwner, result.LeaseExpiresAt = "", 0
fail := func(message string) error {
result.Status, result.Error = domain.StepFailed, message
finishedAt := domain.NowNano()
if err := s.Repo.FinishOutboxStep(ctx, b.ID, result.ID, claimOwner, &result, finishedAt); err != nil {
return err
}
return s.Repo.RecomputeOutboxStatus(ctx, b.ID, finishedAt)
}
if s.Transport == nil { if s.Transport == nil {
st.Status = domain.StepFailed if err := fail("publish transport not configured"); err != nil {
st.Error = "publish transport not configured" return count, err
}
continue continue
} }
st.Status = domain.StepPublishing
b.Status = domain.OBActive
b.UpdatedAt = domain.NowNano()
_ = s.Repo.SaveOutbox(ctx, b)
token := "fake-token" token := "fake-token"
if s.Accounts != nil { if s.Accounts != nil {
acc, aerr := s.Accounts.Get(ctx, st.AccountID) acc, aerr := s.Accounts.Get(ctx, result.AccountID)
if aerr != nil || acc == nil || !acc.IsUsable { if aerr != nil || acc == nil || !acc.IsUsable {
st.Status = domain.StepFailed if err := fail("account not usable"); err != nil {
st.Error = "account not usable" return count, err
b.UpdatedAt = domain.NowNano() }
b.Status = recomputeBundle(b.Steps)
_ = s.Repo.SaveOutbox(ctx, b)
continue continue
} }
if t, terr := s.Accounts.AccessToken(ctx, acc); terr == nil && t != "" { if t, terr := s.Accounts.AccessToken(ctx, acc); terr != nil {
if err := fail(terr.Error()); err != nil {
return count, err
}
continue
} else if t != "" {
token = t token = t
} }
} }
replyTo := st.ReplyTo root = findRoot(claimed)
if replyTo == "" && st.Kind == domain.StepReply { replyTo := result.ReplyTo
if replyTo == "" && result.Kind == domain.StepReply {
if root != nil && root.MediaID != "" { if root != nil && root.MediaID != "" {
replyTo = root.MediaID replyTo = root.MediaID
} else if b.ReplyToMediaID != "" { } else if claimed.ReplyToMediaID != "" {
replyTo = b.ReplyToMediaID replyTo = claimed.ReplyToMediaID
} }
} }
res, perr := s.Transport.Publish(ctx, domain.PublishRequest{ request := domain.PublishRequest{
AccessToken: token, AccessToken: token,
AccountID: st.AccountID, AccountID: result.AccountID,
Text: st.Text, Text: result.Text,
ReplyTo: replyTo, ReplyTo: replyTo,
ImageURLs: st.ImageURLs, ImageURLs: result.ImageURLs,
// 僅主貼帶話題標籤 // 僅主貼帶話題標籤
TopicTag: func() string { TopicTag: func() string {
if st.Kind == domain.StepRoot { if result.Kind == domain.StepRoot {
return st.TopicTag return result.TopicTag
} }
return "" return ""
}(), }(),
}) }
res, perr := s.publishWithLease(ctx, b.ID, result.ID, claimOwner, leaseDuration, request)
if perr != nil { if perr != nil {
st.Status = domain.StepFailed if errors.Is(perr, domain.ErrLeaseLost) {
st.Error = perr.Error() return count, perr
// OB-04: block replies if root failed }
if st.Kind == domain.StepRoot { if err := fail(perr.Error()); err != nil {
for j := range b.Steps { return count, err
if b.Steps[j].Kind == domain.StepReply && b.Steps[j].Status != domain.StepPublished {
b.Steps[j].Status = domain.StepBlocked
b.Steps[j].Error = "waiting for root success"
}
}
} }
} else { } else {
st.Status = domain.StepPublished result.Status, result.Error = domain.StepPublished, ""
st.Error = "" result.PublishedAt = domain.NowNano()
st.PublishedAt = domain.NowNano()
if res != nil { if res != nil {
st.MediaID = res.MediaID result.MediaID = res.MediaID
} }
// Meta 已抓完圖:刪除 temp/* 暫存outbox 不再保留 image_urls images := append([]string(nil), result.ImageURLs...)
if len(st.ImageURLs) > 0 { result.ImageURLs = nil
s.cleanupEphemeralImages(ctx, st.ImageURLs) finishedAt := domain.NowNano()
st.ImageURLs = nil if err := s.Repo.FinishOutboxStep(ctx, b.ID, result.ID, claimOwner, &result, finishedAt); err != nil {
// The remote result may have succeeded. Never let a stale lease write
// final state; a reclaim can duplicate because Meta has no idempotency key.
return count, err
} }
if err := s.Repo.RecomputeOutboxStatus(ctx, b.ID, finishedAt); err != nil {
return count, err
}
s.cleanupEphemeralImages(ctx, images)
count++ count++
// after root success, unblock replies that were blocked only for root wait }
if st.Kind == domain.StepRoot { fresh, getErr := s.Repo.GetOutbox(ctx, b.ID)
for j := range b.Steps { if getErr != nil {
if b.Steps[j].Kind == domain.StepReply && b.Steps[j].Status == domain.StepBlocked { return count, getErr
b.Steps[j].Status = domain.StepScheduled }
b.Steps[j].Error = "" b = fresh
} }
return count, nil
}
func (s *Service) publishWithLease(ctx context.Context, bundleID, stepID, leaseOwner string, leaseDuration time.Duration, request domain.PublishRequest) (*domain.PublishResult, error) {
publishCtx, cancel := context.WithCancel(ctx)
defer cancel()
renewEvery := leaseDuration / 3
if renewEvery <= 0 {
renewEvery = time.Millisecond
}
renewErr := make(chan error, 1)
done := make(chan struct{})
stopped := make(chan struct{})
go func() {
defer close(stopped)
ticker := time.NewTicker(renewEvery)
defer ticker.Stop()
for {
select {
case <-done:
return
case <-ticker.C:
now := domain.NowNano()
if err := s.Repo.RenewOutboxStepLease(publishCtx, bundleID, stepID, leaseOwner, now, now+int64(leaseDuration)); err != nil {
renewErr <- err
cancel()
return
} }
} }
} }
b.Status = recomputeBundle(b.Steps) }()
b.UpdatedAt = domain.NowNano() res, err := s.Transport.Publish(publishCtx, request)
_ = s.Repo.SaveOutbox(ctx, b) close(done)
<-stopped
// re-read root flags after mutation select {
root = findRoot(b) case leaseErr := <-renewErr:
rootFailed = root != nil && root.Status == domain.StepFailed return res, fmt.Errorf("%w: renew failed: %v", domain.ErrLeaseLost, leaseErr)
rootPublished = root != nil && root.Status == domain.StepPublished default:
return res, err
} }
return count, nil
} }
// ---------- Compose (CP) ---------- // ---------- Compose (CP) ----------
@ -2655,7 +2684,7 @@ func playToOutbox(ownerUID int64, play *domain.Play) *domain.OutboxBundle {
ID: "obstep_" + uuid.NewString()[:10], StepID: st.ID, SortOrder: st.SortOrder, ID: "obstep_" + uuid.NewString()[:10], StepID: st.ID, SortOrder: st.SortOrder,
Kind: kind, AccountID: st.AccountID, Text: st.Text, Kind: kind, AccountID: st.AccountID, Text: st.Text,
ImageURLs: append([]string(nil), st.ImageURLs...), ImageURLs: append([]string(nil), st.ImageURLs...),
Status: domain.StepScheduled, ScheduledAt: cursor, ReplyTo: replyTo, Status: domain.StepScheduled, ScheduledAt: cursor, ReplyTo: replyTo,
}) })
} }
title := play.Title title := play.Title
@ -2719,6 +2748,15 @@ func findRoot(b *domain.OutboxBundle) *domain.OutboxStep {
return nil return nil
} }
func findOutboxStep(b *domain.OutboxBundle, stepID string) *domain.OutboxStep {
for i := range b.Steps {
if b.Steps[i].ID == stepID {
return &b.Steps[i]
}
}
return nil
}
func recomputeBundle(steps []domain.OutboxStep) string { func recomputeBundle(steps []domain.OutboxStep) string {
if len(steps) == 0 { if len(steps) == 0 {
return domain.OBCancelled return domain.OBCancelled

View File

@ -3,9 +3,15 @@ package domain
import "context" import "context"
type Repository interface { type Repository interface {
// ReservePlatform atomically guards both limits and increments one monthly counter document.
// A negative limit means unlimited. BYOK must never call this operation.
ReservePlatform(ctx context.Context, uid int64, monthKey, meter string, cost, totalLimit, meterLimit int) error
GetMonthlyCounter(ctx context.Context, uid int64, monthKey string) (*MonthlyCounter, error)
InsertEvent(ctx context.Context, e *Event) error InsertEvent(ctx context.Context, e *Event) error
ListEvents(ctx context.Context, uid int64, monthKey, keyMode string, limit int) ([]*Event, error) ListEvents(ctx context.Context, uid int64, monthKey, keyMode string, limit int) ([]*Event, error)
ListEventsAll(ctx context.Context, monthKey string, limit int) ([]*Event, error) ListEventsAll(ctx context.Context, monthKey string, limit int) ([]*Event, error)
ListEventsRange(ctx context.Context, fromNano, toExclusiveNano int64) ([]*Event, error)
GetPrefs(ctx context.Context, uid int64) (*MemberPrefs, error) GetPrefs(ctx context.Context, uid int64) (*MemberPrefs, error)
SavePrefs(ctx context.Context, p *MemberPrefs) error SavePrefs(ctx context.Context, p *MemberPrefs) error

View File

@ -30,6 +30,8 @@ var (
// 平台 key 無容量/限流:僅 BYOK 可繼續(不扣平台點) // 平台 key 無容量/限流:僅 BYOK 可繼續(不扣平台點)
ErrPlatformCapacity = errors.New("usage.err.platformCapacity") ErrPlatformCapacity = errors.New("usage.err.platformCapacity")
ErrInvalidPlan = errors.New("invalid plan") ErrInvalidPlan = errors.New("invalid plan")
ErrInvalidRange = errors.New("invalid usage analytics range")
ErrPurchaseDisabled = errors.New("self-service plan purchase is disabled; use billing checkout")
) )
type Event struct { type Event struct {
@ -44,11 +46,30 @@ type Event struct {
CreatedAt int64 `bson:"created_at" json:"created_at"` CreatedAt int64 `bson:"created_at" json:"created_at"`
} }
type MonthlyCounter struct {
ID string `bson:"_id" json:"id"`
UID int64 `bson:"uid" json:"uid"`
MonthKey string `bson:"month_key" json:"month_key"`
TotalCredits int `bson:"total_credits" json:"total_credits"`
ByMeter map[string]int `bson:"by_meter" json:"by_meter"`
UpdatedAt int64 `bson:"updated_at" json:"updated_at"`
}
type MemberPrefs struct { type MemberPrefs struct {
UID int64 `bson:"uid" json:"uid"` UID int64 `bson:"uid" json:"uid"`
PlanID string `bson:"plan_id" json:"plan_id"` PlanID string `bson:"plan_id" json:"plan_id"`
Unlimited bool `bson:"unlimited" json:"unlimited"` Unlimited bool `bson:"unlimited" json:"unlimited"`
UpdatedAt int64 `bson:"updated_at" json:"updated_at"` AdminOverride bool `bson:"admin_override" json:"admin_override"`
PaidPlanID string `bson:"paid_plan_id,omitempty" json:"paid_plan_id,omitempty"`
SubscriptionStatus string `bson:"subscription_status,omitempty" json:"subscription_status,omitempty"`
StripeCustomerID string `bson:"stripe_customer_id,omitempty" json:"stripe_customer_id,omitempty"`
StripeSubscriptionID string `bson:"stripe_subscription_id,omitempty" json:"stripe_subscription_id,omitempty"`
CurrentPeriodStart int64 `bson:"current_period_start,omitempty" json:"current_period_start,omitempty"`
CurrentPeriodEnd int64 `bson:"current_period_end,omitempty" json:"current_period_end,omitempty"`
CancelAtPeriodEnd bool `bson:"cancel_at_period_end" json:"cancel_at_period_end"`
BillingEventID string `bson:"billing_event_id,omitempty" json:"-"`
BillingEventCreatedAt int64 `bson:"billing_event_created_at,omitempty" json:"-"`
UpdatedAt int64 `bson:"updated_at" json:"updated_at"`
} }
type Purchase struct { type Purchase struct {
@ -152,16 +173,70 @@ type MonthSummary struct {
} }
type TenantRow struct { type TenantRow struct {
UID int64 `json:"uid"` UID int64 `json:"uid"`
PlanID string `json:"plan_id"` PlanID string `json:"plan_id"`
Unlimited bool `json:"unlimited"` Unlimited bool `json:"unlimited"`
PlatformCreditsUsed int `json:"platform_credits_used"` PlatformCreditsUsed int `json:"platform_credits_used"`
ByokCallCount int `json:"byok_call_count"` ByokCallCount int `json:"byok_call_count"`
} }
type TenantSummary struct { type TenantSummary struct {
MonthKey string `json:"month_key"` MonthKey string `json:"month_key"`
PlatformCreditsUsed int `json:"platform_credits_used"` PlatformCreditsUsed int `json:"platform_credits_used"`
ByokCallCount int `json:"byok_call_count"` ByokCallCount int `json:"byok_call_count"`
Members []TenantRow `json:"members"` Members []TenantRow `json:"members"`
}
type AnalyticsMember struct {
UID int64
Email string
DisplayName string
}
type AnalyticsQuery struct {
Granularity string
From time.Time
To time.Time
}
type AnalyticsBucket struct {
Key string
Label string
PurchasedCredits int
ConsumedCredits int
AICalls int
SearchCalls int
ByokCallCount int
}
type AnalyticsMemberRow struct {
UID int64
Email string
DisplayName string
Unlimited bool
PlanID string
PurchasedCredits int
TotalCredits int
AICalls int
SearchCalls int
ByokCallCount int
RemainingCredits int
Pct int
}
type TenantAnalytics struct {
Granularity string
From string
To string
RangeLabel string
PurchasedCredits int
ConsumedCredits int
RemainingCredits int
Pct int
AICalls int
SearchCalls int
ByMeter map[string]MeterCount
Byok ByokBlock
Series []AnalyticsBucket
Members []AnalyticsMemberRow
} }

View File

@ -2,6 +2,7 @@ package repository
import ( import (
"context" "context"
"fmt"
"sync" "sync"
"apps/backend/internal/module/usage/domain" "apps/backend/internal/module/usage/domain"
@ -10,16 +11,68 @@ import (
type MemoryStore struct { type MemoryStore struct {
mu sync.Mutex mu sync.Mutex
events []*domain.Event events []*domain.Event
counters map[string]*domain.MonthlyCounter
prefs map[int64]*domain.MemberPrefs prefs map[int64]*domain.MemberPrefs
purchases []*domain.Purchase purchases []*domain.Purchase
} }
func NewMemory() *MemoryStore { func NewMemory() *MemoryStore {
return &MemoryStore{ return &MemoryStore{
prefs: map[int64]*domain.MemberPrefs{}, counters: map[string]*domain.MonthlyCounter{},
prefs: map[int64]*domain.MemberPrefs{},
} }
} }
func memoryCounterID(uid int64, monthKey string) string {
return fmt.Sprintf("%d:%s", uid, monthKey)
}
func (s *MemoryStore) ReservePlatform(_ context.Context, uid int64, monthKey, meter string, cost, totalLimit, meterLimit int) error {
s.mu.Lock()
defer s.mu.Unlock()
id := memoryCounterID(uid, monthKey)
counter := s.counters[id]
if counter == nil {
counter = &domain.MonthlyCounter{ID: id, UID: uid, MonthKey: monthKey, ByMeter: map[string]int{}}
for _, event := range s.events {
if event.UID == uid && event.MonthKey == monthKey && event.KeyMode == domain.KeyModePlatform {
credits := event.Credits
if minimum := domain.MeterCost(event.Meter); credits < minimum {
credits = minimum
}
counter.TotalCredits += credits
counter.ByMeter[event.Meter] += credits
}
}
s.counters[id] = counter
}
if totalLimit >= 0 && counter.TotalCredits+cost > totalLimit {
return domain.ErrQuotaExceeded
}
if meterLimit >= 0 && counter.ByMeter[meter]+cost > meterLimit {
return domain.ErrMeterCapExceeded
}
counter.TotalCredits += cost
counter.ByMeter[meter] += cost
counter.UpdatedAt = domain.NowNano()
return nil
}
func (s *MemoryStore) GetMonthlyCounter(_ context.Context, uid int64, monthKey string) (*domain.MonthlyCounter, error) {
s.mu.Lock()
defer s.mu.Unlock()
counter := s.counters[memoryCounterID(uid, monthKey)]
if counter == nil {
return nil, domain.ErrNotFound
}
copyCounter := *counter
copyCounter.ByMeter = make(map[string]int, len(counter.ByMeter))
for meter, credits := range counter.ByMeter {
copyCounter.ByMeter[meter] = credits
}
return &copyCounter, nil
}
func (s *MemoryStore) InsertEvent(_ context.Context, e *domain.Event) error { func (s *MemoryStore) InsertEvent(_ context.Context, e *domain.Event) error {
s.mu.Lock() s.mu.Lock()
defer s.mu.Unlock() defer s.mu.Unlock()
@ -70,6 +123,20 @@ func (s *MemoryStore) ListEventsAll(_ context.Context, monthKey string, limit in
return out, nil return out, nil
} }
func (s *MemoryStore) ListEventsRange(_ context.Context, fromNano, toExclusiveNano int64) ([]*domain.Event, error) {
s.mu.Lock()
defer s.mu.Unlock()
out := make([]*domain.Event, 0)
for _, event := range s.events {
if event.CreatedAt < fromNano || event.CreatedAt >= toExclusiveNano {
continue
}
cp := *event
out = append(out, &cp)
}
return out, nil
}
func (s *MemoryStore) GetPrefs(_ context.Context, uid int64) (*domain.MemberPrefs, error) { func (s *MemoryStore) GetPrefs(_ context.Context, uid int64) (*domain.MemberPrefs, error) {
s.mu.Lock() s.mu.Lock()
defer s.mu.Unlock() defer s.mu.Unlock()

View File

@ -2,17 +2,20 @@ package repository
import ( import (
"context" "context"
"fmt"
libmongo "apps/backend/internal/lib/mongo" libmongo "apps/backend/internal/lib/mongo"
"apps/backend/internal/module/usage/domain" "apps/backend/internal/module/usage/domain"
"github.com/zeromicro/go-zero/core/stores/mon" "github.com/zeromicro/go-zero/core/stores/mon"
"go.mongodb.org/mongo-driver/bson" "go.mongodb.org/mongo-driver/bson"
"go.mongodb.org/mongo-driver/mongo"
"go.mongodb.org/mongo-driver/mongo/options" "go.mongodb.org/mongo-driver/mongo/options"
) )
type MonStore struct { type MonStore struct {
events *mon.Model events *mon.Model
counters *mon.Model
prefs *mon.Model prefs *mon.Model
purchases *mon.Model purchases *mon.Model
} }
@ -21,11 +24,116 @@ func NewMonStore(uri, database string) *MonStore {
uri = libmongo.MustMongoURI(uri) uri = libmongo.MustMongoURI(uri)
return &MonStore{ return &MonStore{
events: mon.MustNewModel(uri, database, "usage_events"), events: mon.MustNewModel(uri, database, "usage_events"),
counters: mon.MustNewModel(uri, database, "usage_monthly_counters"),
prefs: mon.MustNewModel(uri, database, "usage_prefs"), prefs: mon.MustNewModel(uri, database, "usage_prefs"),
purchases: mon.MustNewModel(uri, database, "usage_purchases"), purchases: mon.MustNewModel(uri, database, "usage_purchases"),
} }
} }
func monthlyCounterID(uid int64, monthKey string) string {
return fmt.Sprintf("%d:%s", uid, monthKey)
}
func meterCounterField(meter string) (string, error) {
switch meter {
case domain.MeterAICopy, domain.MeterAIResearch, domain.MeterWebSearch, domain.MeterAIImage:
return "by_meter." + meter, nil
default:
return "", fmt.Errorf("invalid usage meter %q", meter)
}
}
func (s *MonStore) ensureMonthlyCounter(ctx context.Context, uid int64, monthKey string) error {
id := monthlyCounterID(uid, monthKey)
var existing domain.MonthlyCounter
if err := s.counters.FindOne(ctx, &existing, bson.M{"_id": id}); err == nil {
return nil
} else if err != mon.ErrNotFound {
return err
}
events, err := s.ListEvents(ctx, uid, monthKey, domain.KeyModePlatform, 0)
if err != nil {
return err
}
counter := &domain.MonthlyCounter{
ID: id, UID: uid, MonthKey: monthKey, ByMeter: map[string]int{}, UpdatedAt: domain.NowNano(),
}
for _, event := range events {
credits := event.Credits
if minimum := domain.MeterCost(event.Meter); credits < minimum {
credits = minimum
}
counter.TotalCredits += credits
counter.ByMeter[event.Meter] += credits
}
if _, err := s.counters.InsertOne(ctx, counter); err != nil && !mongo.IsDuplicateKeyError(err) {
return err
}
return nil
}
func (s *MonStore) ReservePlatform(ctx context.Context, uid int64, monthKey, meter string, cost, totalLimit, meterLimit int) error {
field, err := meterCounterField(meter)
if err != nil {
return err
}
if totalLimit >= 0 && cost > totalLimit {
return domain.ErrQuotaExceeded
}
if meterLimit >= 0 && cost > meterLimit {
return domain.ErrMeterCapExceeded
}
if err := s.ensureMonthlyCounter(ctx, uid, monthKey); err != nil {
return err
}
filter := bson.M{"_id": monthlyCounterID(uid, monthKey)}
if totalLimit >= 0 {
filter["total_credits"] = bson.M{"$lte": totalLimit - cost}
}
if meterLimit >= 0 {
filter["$or"] = bson.A{
bson.M{field: bson.M{"$exists": false}},
bson.M{field: bson.M{"$lte": meterLimit - cost}},
}
}
res, err := s.counters.UpdateOne(ctx, filter, bson.M{
"$inc": bson.M{"total_credits": cost, field: cost},
"$set": bson.M{"updated_at": domain.NowNano()},
})
if err != nil {
return err
}
if res.MatchedCount != 0 {
return nil
}
counter, err := s.GetMonthlyCounter(ctx, uid, monthKey)
if err != nil {
return err
}
if totalLimit >= 0 && counter.TotalCredits+cost > totalLimit {
return domain.ErrQuotaExceeded
}
if meterLimit >= 0 && counter.ByMeter[meter]+cost > meterLimit {
return domain.ErrMeterCapExceeded
}
return fmt.Errorf("platform quota reservation did not match counter")
}
func (s *MonStore) GetMonthlyCounter(ctx context.Context, uid int64, monthKey string) (*domain.MonthlyCounter, error) {
var counter domain.MonthlyCounter
err := s.counters.FindOne(ctx, &counter, bson.M{"_id": monthlyCounterID(uid, monthKey)})
if err == mon.ErrNotFound {
return nil, domain.ErrNotFound
}
if err != nil {
return nil, err
}
if counter.ByMeter == nil {
counter.ByMeter = map[string]int{}
}
return &counter, nil
}
func (s *MonStore) InsertEvent(ctx context.Context, e *domain.Event) error { func (s *MonStore) InsertEvent(ctx context.Context, e *domain.Event) error {
_, err := s.events.InsertOne(ctx, e) _, err := s.events.InsertOne(ctx, e)
return err return err
@ -62,6 +170,14 @@ func (s *MonStore) ListEventsAll(ctx context.Context, monthKey string, limit int
return list, err return list, err
} }
func (s *MonStore) ListEventsRange(ctx context.Context, fromNano, toExclusiveNano int64) ([]*domain.Event, error) {
var list []*domain.Event
err := s.events.Find(ctx, &list, bson.M{
"created_at": bson.M{"$gte": fromNano, "$lt": toExclusiveNano},
}, options.Find().SetSort(bson.D{{Key: "created_at", Value: 1}}))
return list, err
}
func (s *MonStore) GetPrefs(ctx context.Context, uid int64) (*domain.MemberPrefs, error) { func (s *MonStore) GetPrefs(ctx context.Context, uid int64) (*domain.MemberPrefs, error) {
var p domain.MemberPrefs var p domain.MemberPrefs
err := s.prefs.FindOne(ctx, &p, bson.M{"uid": uid}) err := s.prefs.FindOne(ctx, &p, bson.M{"uid": uid})

View File

@ -0,0 +1,238 @@
package usecase
import (
"context"
"fmt"
"math"
"sort"
"time"
"apps/backend/internal/module/usage/domain"
)
const analyticsDateLayout = "2006-01-02"
func ParseAnalyticsQuery(granularity, fromRaw, toRaw string) (domain.AnalyticsQuery, error) {
if granularity != "day" && granularity != "month" && granularity != "year" {
return domain.AnalyticsQuery{}, fmt.Errorf("%w: granularity must be day, month, or year", domain.ErrInvalidRange)
}
from, err := time.Parse(analyticsDateLayout, fromRaw)
if err != nil {
return domain.AnalyticsQuery{}, fmt.Errorf("%w: from must be YYYY-MM-DD", domain.ErrInvalidRange)
}
to, err := time.Parse(analyticsDateLayout, toRaw)
if err != nil {
return domain.AnalyticsQuery{}, fmt.Errorf("%w: to must be YYYY-MM-DD", domain.ErrInvalidRange)
}
if from.After(to) {
return domain.AnalyticsQuery{}, fmt.Errorf("%w: from must not be after to", domain.ErrInvalidRange)
}
q := domain.AnalyticsQuery{Granularity: granularity, From: from.UTC(), To: to.UTC()}
if len(analyticsBuckets(q)) > map[string]int{"day": 93, "month": 36, "year": 8}[granularity] {
return domain.AnalyticsQuery{}, fmt.Errorf("%w: range exceeds %s limit", domain.ErrInvalidRange, granularity)
}
return q, nil
}
type analyticsBucketRange struct {
key string
label string
from time.Time
to time.Time
}
func analyticsBuckets(q domain.AnalyticsQuery) []analyticsBucketRange {
var buckets []analyticsBucketRange
switch q.Granularity {
case "day":
for cur := q.From; !cur.After(q.To); cur = cur.AddDate(0, 0, 1) {
buckets = append(buckets, analyticsBucketRange{key: cur.Format(analyticsDateLayout), label: cur.Format("01/02"), from: cur, to: cur})
}
case "month":
for cur := time.Date(q.From.Year(), q.From.Month(), 1, 0, 0, 0, 0, time.UTC); !cur.After(q.To); cur = cur.AddDate(0, 1, 0) {
end := cur.AddDate(0, 1, -1)
buckets = append(buckets, analyticsBucketRange{key: cur.Format("2006-01"), label: cur.Format("2006/01"), from: maxTime(cur, q.From), to: minTime(end, q.To)})
}
case "year":
for year := q.From.Year(); year <= q.To.Year(); year++ {
start := time.Date(year, 1, 1, 0, 0, 0, 0, time.UTC)
end := time.Date(year, 12, 31, 0, 0, 0, 0, time.UTC)
key := fmt.Sprintf("%04d", year)
buckets = append(buckets, analyticsBucketRange{key: key, label: key, from: maxTime(start, q.From), to: minTime(end, q.To)})
}
}
return buckets
}
func maxTime(a, b time.Time) time.Time {
if a.After(b) {
return a
}
return b
}
func minTime(a, b time.Time) time.Time {
if a.Before(b) {
return a
}
return b
}
func proratedCredits(monthly int, from, to time.Time) int {
credits := 0.0
for cur := from; !cur.After(to); cur = cur.AddDate(0, 0, 1) {
days := time.Date(cur.Year(), cur.Month()+1, 0, 0, 0, 0, 0, time.UTC).Day()
credits += float64(monthly) / float64(days)
}
return int(math.Round(credits))
}
func isAIMeter(meter string) bool {
return meter == domain.MeterAICopy || meter == domain.MeterAIResearch || meter == domain.MeterAIImage
}
func emptyMeterCounts() map[string]domain.MeterCount {
return map[string]domain.MeterCount{
domain.MeterAICopy: {}, domain.MeterAIResearch: {}, domain.MeterWebSearch: {}, domain.MeterAIImage: {},
}
}
func (s *Service) TenantAnalytics(ctx context.Context, q domain.AnalyticsQuery, tenantMembers []domain.AnalyticsMember) (*domain.TenantAnalytics, error) {
events, err := s.Repo.ListEventsRange(ctx, q.From.UnixNano(), q.To.AddDate(0, 0, 1).UnixNano())
if err != nil {
return nil, err
}
type memberAcc struct {
row domain.AnalyticsMemberRow
}
members := make(map[int64]*memberAcc, len(tenantMembers))
for _, member := range tenantMembers {
members[member.UID] = &memberAcc{row: domain.AnalyticsMemberRow{UID: member.UID, Email: member.Email, DisplayName: member.DisplayName}}
}
for _, event := range events {
if members[event.UID] == nil {
members[event.UID] = &memberAcc{row: domain.AnalyticsMemberRow{UID: event.UID}}
}
}
buckets := analyticsBuckets(q)
bucketByKey := make(map[string]int, len(buckets))
result := &domain.TenantAnalytics{
Granularity: q.Granularity, From: q.From.Format(analyticsDateLayout), To: q.To.Format(analyticsDateLayout),
RangeLabel: fmt.Sprintf("%s - %s", q.From.Format(analyticsDateLayout), q.To.Format(analyticsDateLayout)),
ByMeter: emptyMeterCounts(), Byok: domain.ByokBlock{ByMeter: emptyMeterCounts()},
}
for _, bucket := range buckets {
row := domain.AnalyticsBucket{Key: bucket.key, Label: bucket.label}
result.Series = append(result.Series, row)
bucketByKey[bucket.key] = len(result.Series) - 1
}
for _, acc := range members {
prefs, prefsErr := s.ensurePrefs(ctx, acc.row.UID)
if prefsErr != nil {
return nil, prefsErr
}
plan := domain.ResolvePlan(prefs.PlanID)
acc.row.PlanID, acc.row.Unlimited = plan.ID, prefs.Unlimited
acc.row.PurchasedCredits = proratedCredits(plan.MonthlyCredits, q.From, q.To)
result.PurchasedCredits += acc.row.PurchasedCredits
for i, bucket := range buckets {
result.Series[i].PurchasedCredits += proratedCredits(plan.MonthlyCredits, bucket.from, bucket.to)
}
}
for _, event := range events {
bucketIndex, hasBucket := bucketByKey[eventBucketKey(event.CreatedAt, q.Granularity)]
acc := members[event.UID]
if event.KeyMode == domain.KeyModeByok {
result.Byok.CallCount++
meter := result.Byok.ByMeter[event.Meter]
meter.Count++
result.Byok.ByMeter[event.Meter] = meter
acc.row.ByokCallCount++
if hasBucket {
result.Series[bucketIndex].ByokCallCount++
}
continue
}
if event.KeyMode != domain.KeyModePlatform {
continue
}
credits := event.Credits
if credits < 0 {
credits = 0
}
result.ConsumedCredits += credits
meter := result.ByMeter[event.Meter]
meter.Count++
meter.Credits += credits
result.ByMeter[event.Meter] = meter
acc.row.TotalCredits += credits
if isAIMeter(event.Meter) {
result.AICalls++
acc.row.AICalls++
if hasBucket {
result.Series[bucketIndex].AICalls++
}
} else if event.Meter == domain.MeterWebSearch {
result.SearchCalls++
acc.row.SearchCalls++
if hasBucket {
result.Series[bucketIndex].SearchCalls++
}
}
if hasBucket {
result.Series[bucketIndex].ConsumedCredits += credits
}
}
result.RemainingCredits = maxInt(0, result.PurchasedCredits-result.ConsumedCredits)
result.Pct = percentage(result.ConsumedCredits, result.PurchasedCredits)
for _, acc := range members {
acc.row.RemainingCredits = maxInt(0, acc.row.PurchasedCredits-acc.row.TotalCredits)
acc.row.Pct = percentage(acc.row.TotalCredits, acc.row.PurchasedCredits)
result.Members = append(result.Members, acc.row)
}
sort.Slice(result.Members, func(i, j int) bool {
if result.Members[i].TotalCredits == result.Members[j].TotalCredits {
return result.Members[i].UID < result.Members[j].UID
}
return result.Members[i].TotalCredits > result.Members[j].TotalCredits
})
return result, nil
}
func eventBucketKey(nano int64, granularity string) string {
t := time.Unix(0, nano).UTC()
switch granularity {
case "day":
return t.Format(analyticsDateLayout)
case "month":
return t.Format("2006-01")
default:
return t.Format("2006")
}
}
func percentage(used, total int) int {
if total <= 0 {
return 0
}
return minInt(999, int(math.Round(float64(used)*100/float64(total))))
}
func minInt(a, b int) int {
if a < b {
return a
}
return b
}
func maxInt(a, b int) int {
if a > b {
return a
}
return b
}

View File

@ -0,0 +1,73 @@
package usecase_test
import (
"context"
"testing"
"time"
"apps/backend/internal/module/usage/domain"
"apps/backend/internal/module/usage/repository"
"apps/backend/internal/module/usage/usecase"
"github.com/stretchr/testify/require"
)
func analyticsNano(t *testing.T, value string) int64 {
t.Helper()
parsed, err := time.Parse(time.RFC3339, value)
require.NoError(t, err)
return parsed.UnixNano()
}
func TestParseAnalyticsQueryValidationAndCaps(t *testing.T) {
_, err := usecase.ParseAnalyticsQuery("week", "2026-01-01", "2026-01-02")
require.ErrorIs(t, err, domain.ErrInvalidRange)
_, err = usecase.ParseAnalyticsQuery("day", "2026-01-03", "2026-01-02")
require.ErrorIs(t, err, domain.ErrInvalidRange)
_, err = usecase.ParseAnalyticsQuery("day", "2026-01-01", "2026-04-04")
require.ErrorIs(t, err, domain.ErrInvalidRange)
_, err = usecase.ParseAnalyticsQuery("month", "2023-01-01", "2025-12-31")
require.NoError(t, err)
}
func TestTenantAnalyticsUsesInclusiveUTCBoundsAndSeparateModes(t *testing.T) {
repo := repository.NewMemory()
for _, event := range []*domain.Event{
{ID: "before", UID: 1, Meter: domain.MeterAICopy, Credits: 9, KeyMode: domain.KeyModePlatform, CreatedAt: analyticsNano(t, "2026-01-31T23:59:59Z")},
{ID: "first", UID: 1, Meter: domain.MeterAICopy, Credits: 2, KeyMode: domain.KeyModePlatform, CreatedAt: analyticsNano(t, "2026-02-01T00:00:00Z")},
{ID: "byok", UID: 1, Meter: domain.MeterWebSearch, Credits: 99, KeyMode: domain.KeyModeByok, CreatedAt: analyticsNano(t, "2026-02-02T12:00:00Z")},
{ID: "last", UID: 1, Meter: domain.MeterWebSearch, Credits: 1, KeyMode: domain.KeyModePlatform, CreatedAt: analyticsNano(t, "2026-02-02T23:59:59.999999999Z")},
{ID: "after", UID: 1, Meter: domain.MeterAICopy, Credits: 9, KeyMode: domain.KeyModePlatform, CreatedAt: analyticsNano(t, "2026-02-03T00:00:00Z")},
} {
require.NoError(t, repo.InsertEvent(context.Background(), event))
}
query, err := usecase.ParseAnalyticsQuery("day", "2026-02-01", "2026-02-02")
require.NoError(t, err)
result, err := usecase.New(repo, nil).TenantAnalytics(context.Background(), query, []domain.AnalyticsMember{{UID: 1, Email: "one@example.test", DisplayName: "One"}})
require.NoError(t, err)
require.Equal(t, 3, result.ConsumedCredits)
require.Equal(t, 1, result.AICalls)
require.Equal(t, 1, result.SearchCalls)
require.Equal(t, 1, result.Byok.CallCount)
require.Equal(t, 0, result.Byok.ByMeter[domain.MeterWebSearch].Credits)
require.Equal(t, 1, result.Byok.ByMeter[domain.MeterWebSearch].Count)
require.Len(t, result.Series, 2)
require.Equal(t, 2, result.Series[0].ConsumedCredits)
require.Equal(t, 1, result.Series[1].ConsumedCredits)
require.Equal(t, 1, result.Series[1].ByokCallCount)
require.Equal(t, 3, result.Members[0].TotalCredits)
require.Equal(t, 1, result.Members[0].ByokCallCount)
}
func TestMemoryListEventsRangeUsesExclusiveUpperBound(t *testing.T) {
repo := repository.NewMemory()
for i, nano := range []int64{99, 100, 199, 200} {
require.NoError(t, repo.InsertEvent(context.Background(), &domain.Event{ID: string(rune('a' + i)), CreatedAt: nano}))
}
list, err := repo.ListEventsRange(context.Background(), 100, 200)
require.NoError(t, err)
require.Len(t, list, 2)
require.Equal(t, int64(100), list[0].CreatedAt)
require.Equal(t, int64(199), list[1].CreatedAt)
}

View File

@ -2,6 +2,7 @@ package usecase
import ( import (
"fmt" "fmt"
"strings"
"sync" "sync"
"time" "time"
@ -103,13 +104,20 @@ func (AlwaysAllowPlatform) MarkLimited(time.Time) {}
// than bypassing a shared safety limit. // than bypassing a shared safety limit.
type RedisPlatformGate struct { type RedisPlatformGate struct {
client *redis.Redis client *redis.Redis
namespace string
aiPerMin int aiPerMin int
searchPerMin int searchPerMin int
} }
func NewRedisPlatformGate(conf redis.RedisConf, aiPerMin, searchPerMin int) *RedisPlatformGate { func NewRedisPlatformGate(conf redis.RedisConf, aiPerMin, searchPerMin int, namespace ...string) *RedisPlatformGate {
prefix := "haixun:dev:v1"
if len(namespace) > 0 {
if configured := strings.Trim(strings.TrimSpace(namespace[0]), ":"); configured != "" {
prefix = configured
}
}
return &RedisPlatformGate{ return &RedisPlatformGate{
client: redis.MustNewRedis(conf), aiPerMin: aiPerMin, searchPerMin: searchPerMin, client: redis.MustNewRedis(conf), namespace: prefix, aiPerMin: aiPerMin, searchPerMin: searchPerMin,
} }
} }
@ -117,7 +125,7 @@ func (g *RedisPlatformGate) MarkLimited(until time.Time) {
if g == nil || g.client == nil { if g == nil || g.client == nil {
return return
} }
key := "haixun:usage:platform:limited" key := g.limitedKey()
if until.IsZero() || !until.After(time.Now()) { if until.IsZero() || !until.After(time.Now()) {
_, _ = g.client.Del(key) _, _ = g.client.Del(key)
return return
@ -138,7 +146,7 @@ func (g *RedisPlatformGate) AllowPlatform(meter string) bool {
limit, bucket = g.searchPerMin, "search" limit, bucket = g.searchPerMin, "search"
} }
if limit <= 0 { if limit <= 0 {
limited, err := g.client.Exists("haixun:usage:platform:limited") limited, err := g.client.Exists(g.limitedKey())
return err == nil && !limited return err == nil && !limited
} }
window := time.Now().UTC().Format("200601021504") window := time.Now().UTC().Format("200601021504")
@ -148,7 +156,7 @@ local n = redis.call('INCR', KEYS[2])
if n == 1 then redis.call('EXPIRE', KEYS[2], 120) end if n == 1 then redis.call('EXPIRE', KEYS[2], 120) end
if n > tonumber(ARGV[1]) then return 0 end if n > tonumber(ARGV[1]) then return 0 end
return 1 return 1
`, []string{"haixun:usage:platform:limited", fmt.Sprintf("haixun:usage:platform:%s:%s", bucket, window)}, limit) `, []string{g.limitedKey(), g.windowKey(bucket, window)}, limit)
if err != nil { if err != nil {
return false return false
} }
@ -161,3 +169,11 @@ return 1
return false return false
} }
} }
func (g *RedisPlatformGate) limitedKey() string {
return fmt.Sprintf("%s:usage:platform:{platform-gate}:limited", g.namespace)
}
func (g *RedisPlatformGate) windowKey(bucket, window string) string {
return fmt.Sprintf("%s:usage:platform:{platform-gate}:%s:%s", g.namespace, bucket, window)
}

View File

@ -0,0 +1,15 @@
package usecase
import "testing"
func TestRedisPlatformGateKeysShareClusterHashTag(t *testing.T) {
gate := &RedisPlatformGate{namespace: "haixun:test:v1"}
limited := gate.limitedKey()
window := gate.windowKey("ai", "202607141200")
if limited != "haixun:test:v1:usage:platform:{platform-gate}:limited" {
t.Fatalf("unexpected limited key %q", limited)
}
if window != "haixun:test:v1:usage:platform:{platform-gate}:ai:202607141200" {
t.Fatalf("unexpected window key %q", window)
}
}

View File

@ -37,7 +37,10 @@ func (r *SettingsResolver) platformAIKey(provider string) string {
func (r *SettingsResolver) Resolve(ctx context.Context, uid int64, meter string) (string, bool, error) { func (r *SettingsResolver) Resolve(ctx context.Context, uid int64, meter string) (string, bool, error) {
st, err := r.Members.GetSettings(ctx, uid) st, err := r.Members.GetSettings(ctx, uid)
if err != nil || st == nil { if err != nil {
return "", false, err
}
if st == nil {
st = memberDomain.DefaultSettings(uid) st = memberDomain.DefaultSettings(uid)
} }
provider := ai.NormalizeProvider(st.Provider) provider := ai.NormalizeProvider(st.Provider)
@ -72,7 +75,10 @@ func (r *SettingsResolver) ResolveKey(ctx context.Context, uid int64, meter stri
if !has { if !has {
return "", "", domain.ErrNoKey return "", "", domain.ErrNoKey
} }
st, _ := r.Members.GetSettings(ctx, uid) st, err := r.Members.GetSettings(ctx, uid)
if err != nil {
return "", "", err
}
if st == nil { if st == nil {
st = memberDomain.DefaultSettings(uid) st = memberDomain.DefaultSettings(uid)
} }

View File

@ -46,7 +46,8 @@ func New(repo domain.Repository, resolver KeyResolver) *Service {
return &Service{Repo: repo, Resolver: resolver} return &Service{Repo: repo, Resolver: resolver}
} }
// RecordCall after successful external call (or for meter tests). // RecordCall preserves the audit event after a call. Platform credits were already
// reserved by PrepareCall; this method must not charge them a second time.
func (s *Service) RecordCall(ctx context.Context, uid int64, meter, keyMode, label, source string) (*domain.Event, error) { func (s *Service) RecordCall(ctx context.Context, uid int64, meter, keyMode, label, source string) (*domain.Event, error) {
if keyMode != domain.KeyModePlatform && keyMode != domain.KeyModeByok { if keyMode != domain.KeyModePlatform && keyMode != domain.KeyModeByok {
return nil, fmt.Errorf("invalid key_mode") return nil, fmt.Errorf("invalid key_mode")
@ -54,9 +55,6 @@ func (s *Service) RecordCall(ctx context.Context, uid int64, meter, keyMode, lab
credits := 0 credits := 0
if keyMode == domain.KeyModePlatform { if keyMode == domain.KeyModePlatform {
credits = domain.MeterCost(meter) credits = domain.MeterCost(meter)
if err := s.checkPlatformQuota(ctx, uid, meter, credits); err != nil {
return nil, err
}
} }
now := domain.NowNano() now := domain.NowNano()
e := &domain.Event{ e := &domain.Event{
@ -87,8 +85,7 @@ func (s *Service) PrepareCall(ctx context.Context, uid int64, meter string) (key
if s.Gate != nil && !s.Gate.AllowPlatform(meter) { if s.Gate != nil && !s.Gate.AllowPlatform(meter) {
return "", domain.ErrPlatformCapacity return "", domain.ErrPlatformCapacity
} }
cost := domain.MeterCost(meter) if err := s.reservePlatform(ctx, uid, meter); err != nil {
if err := s.checkPlatformQuota(ctx, uid, meter, cost); err != nil {
return "", err return "", err
} }
} }
@ -104,55 +101,41 @@ func (s *Service) MarkPlatformLimited(until time.Time) {
s.Gate.MarkLimited(until) s.Gate.MarkLimited(until)
} }
// checkPlatformQuota — 總池 + 分項點數上限unlimited 略過。 func (s *Service) reservePlatform(ctx context.Context, uid int64, meter string) error {
func (s *Service) checkPlatformQuota(ctx context.Context, uid int64, meter string, cost int) error { prefs, err := s.ensurePrefs(ctx, uid)
if cost < 0 {
cost = 0
}
prefs, _ := s.ensurePrefs(ctx, uid)
if prefs.Unlimited {
return nil
}
sum, err := s.GetSummary(ctx, uid, domain.CurrentMonthKey())
if err != nil { if err != nil {
return err return err
} }
// 1) 整月 platform 總點
if sum.Platform.CreditsRemaining < cost {
return domain.ErrQuotaExceeded
}
// 2) 分項點數硬上限platform only加總 = MonthlyCredits
plan := domain.ResolvePlan(prefs.PlanID) plan := domain.ResolvePlan(prefs.PlanID)
totalLimit, meterLimit := plan.MonthlyCredits, -1
if capN, ok := plan.SoftCaps[meter]; ok && capN > 0 { if capN, ok := plan.SoftCaps[meter]; ok && capN > 0 {
usedCredits := 0 meterLimit = capN
usedCount := 0
if m, ok := sum.Platform.ByMeter[meter]; ok {
usedCredits = m.Credits
usedCount = m.Count
// 舊資料若 credits 漏記,用次數 × 單價估
if usedCredits < usedCount*domain.MeterCost(meter) {
usedCredits = usedCount * domain.MeterCost(meter)
}
}
if usedCredits+cost > capN {
return domain.ErrMeterCapExceeded
}
} }
return nil if prefs.Unlimited {
totalLimit, meterLimit = -1, -1
}
return s.Repo.ReservePlatform(ctx, uid, domain.CurrentMonthKey(), meter, domain.MeterCost(meter), totalLimit, meterLimit)
} }
func (s *Service) ensurePrefs(ctx context.Context, uid int64) (*domain.MemberPrefs, error) { func (s *Service) ensurePrefs(ctx context.Context, uid int64) (*domain.MemberPrefs, error) {
p, err := s.Repo.GetPrefs(ctx, uid) p, err := s.Repo.GetPrefs(ctx, uid)
if err != nil || p == nil { if err != nil {
return nil, err
}
if p == nil {
p = &domain.MemberPrefs{UID: uid, PlanID: domain.PlanFree, UpdatedAt: domain.NowNano()} p = &domain.MemberPrefs{UID: uid, PlanID: domain.PlanFree, UpdatedAt: domain.NowNano()}
_ = s.Repo.SavePrefs(ctx, p) if err := s.Repo.SavePrefs(ctx, p); err != nil {
return nil, err
}
} }
// 正規化 plan_id小寫有效方案 // 正規化 plan_id小寫有效方案
norm := domain.ResolvePlan(p.PlanID).ID norm := domain.ResolvePlan(p.PlanID).ID
if p.PlanID != norm { if p.PlanID != norm {
p.PlanID = norm p.PlanID = norm
p.UpdatedAt = domain.NowNano() p.UpdatedAt = domain.NowNano()
_ = s.Repo.SavePrefs(ctx, p) if err := s.Repo.SavePrefs(ctx, p); err != nil {
return nil, err
}
} }
return p, nil return p, nil
} }
@ -161,7 +144,10 @@ func (s *Service) GetSummary(ctx context.Context, uid int64, monthKey string) (*
if monthKey == "" { if monthKey == "" {
monthKey = domain.CurrentMonthKey() monthKey = domain.CurrentMonthKey()
} }
prefs, _ := s.ensurePrefs(ctx, uid) prefs, err := s.ensurePrefs(ctx, uid)
if err != nil {
return nil, err
}
plan := domain.ResolvePlan(prefs.PlanID) plan := domain.ResolvePlan(prefs.PlanID)
// 回寫正規化 plan_id避免 DB 殘留大小寫/空白導致永遠 fallback Free // 回寫正規化 plan_id避免 DB 殘留大小寫/空白導致永遠 fallback Free
if prefs.PlanID != plan.ID { if prefs.PlanID != plan.ID {
@ -191,6 +177,16 @@ func (s *Service) GetSummary(ctx context.Context, uid int64, monthKey string) (*
byokBy[e.Meter] = m byokBy[e.Meter] = m
} }
} }
if counter, counterErr := s.Repo.GetMonthlyCounter(ctx, uid, monthKey); counterErr == nil {
platUsed = counter.TotalCredits
for meter, credits := range counter.ByMeter {
m := platBy[meter]
m.Credits = credits
platBy[meter] = m
}
} else if counterErr != domain.ErrNotFound {
return nil, counterErr
}
remaining := plan.MonthlyCredits - platUsed remaining := plan.MonthlyCredits - platUsed
if remaining < 0 { if remaining < 0 {
remaining = 0 remaining = 0
@ -227,7 +223,10 @@ func (s *Service) SetPrefs(ctx context.Context, actorUID int64, targetUID int64,
// members cannot set unlimited // members cannot set unlimited
return nil, domain.ErrForbidden return nil, domain.ErrForbidden
} }
p, _ := s.ensurePrefs(ctx, targetUID) p, err := s.ensurePrefs(ctx, targetUID)
if err != nil {
return nil, err
}
if planID != "" { if planID != "" {
id := strings.ToLower(strings.TrimSpace(planID)) id := strings.ToLower(strings.TrimSpace(planID))
if _, ok := domain.Plans[id]; !ok { if _, ok := domain.Plans[id]; !ok {
@ -238,6 +237,7 @@ func (s *Service) SetPrefs(ctx context.Context, actorUID int64, targetUID int64,
return nil, domain.ErrForbidden return nil, domain.ErrForbidden
} }
p.PlanID = id p.PlanID = id
p.AdminOverride = true
} }
if unlimited != nil && isAdmin { if unlimited != nil && isAdmin {
p.Unlimited = *unlimited p.Unlimited = *unlimited
@ -250,24 +250,7 @@ func (s *Service) SetPrefs(ctx context.Context, actorUID int64, targetUID int64,
} }
func (s *Service) PurchasePlan(ctx context.Context, uid int64, planID, mockRef string) (*domain.Purchase, error) { func (s *Service) PurchasePlan(ctx context.Context, uid int64, planID, mockRef string) (*domain.Purchase, error) {
id := strings.ToLower(strings.TrimSpace(planID)) return nil, domain.ErrPurchaseDisabled
plan, ok := domain.Plans[id]
if !ok || plan.ID == "" {
return nil, domain.ErrInvalidPlan
}
p, _ := s.ensurePrefs(ctx, uid)
p.PlanID = plan.ID
p.UpdatedAt = domain.NowNano()
if err := s.Repo.SavePrefs(ctx, p); err != nil {
return nil, err
}
pur := &domain.Purchase{
ID: uuid.NewString(), UID: uid, PlanID: plan.ID, MockRef: mockRef, CreatedAt: domain.NowNano(),
}
if err := s.Repo.InsertPurchase(ctx, pur); err != nil {
return nil, err
}
return pur, nil
} }
func (s *Service) ListMyPurchases(ctx context.Context, uid int64, limit int) ([]*domain.Purchase, error) { func (s *Service) ListMyPurchases(ctx context.Context, uid int64, limit int) ([]*domain.Purchase, error) {
@ -307,7 +290,10 @@ func (s *Service) TenantSummary(ctx context.Context, monthKey string) (*domain.T
} }
var members []domain.TenantRow var members []domain.TenantRow
for uid, a := range byUID { for uid, a := range byUID {
prefs, _ := s.ensurePrefs(ctx, uid) prefs, err := s.ensurePrefs(ctx, uid)
if err != nil {
return nil, err
}
members = append(members, domain.TenantRow{ members = append(members, domain.TenantRow{
UID: uid, PlanID: prefs.PlanID, Unlimited: prefs.Unlimited, UID: uid, PlanID: prefs.PlanID, Unlimited: prefs.Unlimited,
PlatformCreditsUsed: a.plat, ByokCallCount: a.byok, PlatformCreditsUsed: a.plat, ByokCallCount: a.byok,

View File

@ -3,6 +3,7 @@ package usecase_test
import ( import (
"context" "context"
"fmt" "fmt"
"sync"
"testing" "testing"
"time" "time"
@ -145,7 +146,9 @@ func TestST_13_PlatformQuota(t *testing.T) {
svc := newUsage(1_000_001, domain.KeyModePlatform) svc := newUsage(1_000_001, domain.KeyModePlatform)
// free: ai_copy soft cap 60pt先燒滿分項 // free: ai_copy soft cap 60pt先燒滿分項
for i := 0; i < 60; i++ { for i := 0; i < 60; i++ {
_, err := svc.RecordCall(context.Background(), 1_000_001, domain.MeterAICopy, domain.KeyModePlatform, "x", "t") mode, err := svc.PrepareCall(context.Background(), 1_000_001, domain.MeterAICopy)
require.NoError(t, err)
_, err = svc.RecordCall(context.Background(), 1_000_001, domain.MeterAICopy, mode, "x", "t")
require.NoError(t, err) require.NoError(t, err)
} }
_, err := svc.PrepareCall(context.Background(), 1_000_001, domain.MeterAICopy) _, err := svc.PrepareCall(context.Background(), 1_000_001, domain.MeterAICopy)
@ -153,24 +156,78 @@ func TestST_13_PlatformQuota(t *testing.T) {
// 分項鎖死後,其他 meter 在總池仍有餘額時可繼續 // 分項鎖死後,其他 meter 在總池仍有餘額時可繼續
_, err = svc.PrepareCall(context.Background(), 1_000_001, domain.MeterWebSearch) _, err = svc.PrepareCall(context.Background(), 1_000_001, domain.MeterWebSearch)
require.NoError(t, err) require.NoError(t, err)
_, err = svc.RecordCall(context.Background(), 1_000_001, domain.MeterWebSearch, domain.KeyModePlatform, "s", "t")
require.NoError(t, err)
// 燒光剩餘總池search 30 + research 15 + image 15 = 60 // 燒光剩餘總池search 30 + research 15 + image 15 = 60
for i := 0; i < 30; i++ { for i := 1; i < 30; i++ {
_, err := svc.RecordCall(context.Background(), 1_000_001, domain.MeterWebSearch, domain.KeyModePlatform, "s", "t") mode, err := svc.PrepareCall(context.Background(), 1_000_001, domain.MeterWebSearch)
require.NoError(t, err)
_, err = svc.RecordCall(context.Background(), 1_000_001, domain.MeterWebSearch, mode, "s", "t")
require.NoError(t, err) require.NoError(t, err)
} }
for i := 0; i < 5; i++ { // 5*3=15 for i := 0; i < 5; i++ { // 5*3=15
_, err := svc.RecordCall(context.Background(), 1_000_001, domain.MeterAIResearch, domain.KeyModePlatform, "r", "t") mode, err := svc.PrepareCall(context.Background(), 1_000_001, domain.MeterAIResearch)
require.NoError(t, err)
_, err = svc.RecordCall(context.Background(), 1_000_001, domain.MeterAIResearch, mode, "r", "t")
require.NoError(t, err) require.NoError(t, err)
} }
// image 15pt = 3 張 // image 15pt = 3 張
for i := 0; i < 3; i++ { for i := 0; i < 3; i++ {
_, err := svc.RecordCall(context.Background(), 1_000_001, domain.MeterAIImage, domain.KeyModePlatform, "img", "t") mode, err := svc.PrepareCall(context.Background(), 1_000_001, domain.MeterAIImage)
require.NoError(t, err)
_, err = svc.RecordCall(context.Background(), 1_000_001, domain.MeterAIImage, mode, "img", "t")
require.NoError(t, err) require.NoError(t, err)
} }
_, err = svc.PrepareCall(context.Background(), 1_000_001, domain.MeterWebSearch) _, err = svc.PrepareCall(context.Background(), 1_000_001, domain.MeterWebSearch)
require.ErrorIs(t, err, domain.ErrQuotaExceeded) require.ErrorIs(t, err, domain.ErrQuotaExceeded)
} }
func TestPlatformReservation_ConcurrentOneCreditCannotBeOverspent(t *testing.T) {
repo := repository.NewMemory()
const attempts = 32
start := make(chan struct{})
results := make(chan error, attempts)
var wg sync.WaitGroup
for i := 0; i < attempts; i++ {
wg.Add(1)
go func() {
defer wg.Done()
<-start
results <- repo.ReservePlatform(context.Background(), 1_000_003, domain.CurrentMonthKey(), domain.MeterAICopy, 1, 1, 1)
}()
}
close(start)
wg.Wait()
close(results)
succeeded := 0
for err := range results {
if err == nil {
succeeded++
continue
}
require.ErrorIs(t, err, domain.ErrQuotaExceeded)
}
require.Equal(t, 1, succeeded)
counter, err := repo.GetMonthlyCounter(context.Background(), 1_000_003, domain.CurrentMonthKey())
require.NoError(t, err)
require.Equal(t, 1, counter.TotalCredits)
require.Equal(t, 1, counter.ByMeter[domain.MeterAICopy])
}
func TestPrepareCall_BYOKDoesNotCreatePlatformCounter(t *testing.T) {
repo := repository.NewMemory()
uid := int64(1_000_004)
svc := usecase.New(repo, &usecase.StaticResolver{Map: map[string]string{
fmt.Sprintf("%d:%s", uid, domain.MeterAICopy): domain.KeyModeByok,
}})
mode, err := svc.PrepareCall(context.Background(), uid, domain.MeterAICopy)
require.NoError(t, err)
require.Equal(t, domain.KeyModeByok, mode)
_, err = repo.GetMonthlyCounter(context.Background(), uid, domain.CurrentMonthKey())
require.ErrorIs(t, err, domain.ErrNotFound)
}
func TestUSG_MeterCap_Image(t *testing.T) { func TestUSG_MeterCap_Image(t *testing.T) {
svc := newUsage(1_000_001, domain.KeyModePlatform) svc := newUsage(1_000_001, domain.KeyModePlatform)
// free ai_image cap 15pt = 3 次;第 4 次因分項上限擋 // free ai_image cap 15pt = 3 次;第 4 次因分項上限擋
@ -210,32 +267,16 @@ func TestUSG_PlatformCapacity_RPM(t *testing.T) {
require.ErrorIs(t, err, domain.ErrPlatformCapacity) require.ErrorIs(t, err, domain.ErrPlatformCapacity)
} }
func TestUSG_11_PurchasePlan(t *testing.T) { func TestUSG_11_LegacyPurchaseCannotActivatePaidPlan(t *testing.T) {
svc := newUsage(1_000_001, domain.KeyModePlatform) svc := newUsage(1_000_001, domain.KeyModePlatform)
p, err := svc.PurchasePlan(context.Background(), 1_000_001, domain.PlanStarter, "mock-pay-1") p, err := svc.PurchasePlan(context.Background(), 1_000_001, domain.PlanStarter, "mock-pay-1")
require.NoError(t, err) require.Nil(t, p)
require.Equal(t, domain.PlanStarter, p.PlanID) require.ErrorIs(t, err, domain.ErrPurchaseDisabled)
prefs, _ := svc.GetPrefs(context.Background(), 1_000_001) prefs, _ := svc.GetPrefs(context.Background(), 1_000_001)
require.Equal(t, domain.PlanStarter, prefs.PlanID) require.Equal(t, domain.PlanFree, prefs.PlanID)
list, err := svc.ListMyPurchases(context.Background(), 1_000_001, 10) list, err := svc.ListMyPurchases(context.Background(), 1_000_001, 10)
require.NoError(t, err) require.NoError(t, err)
require.NotEmpty(t, list) require.Empty(t, list)
// 購買後 summary 配給必須是 Starter 600不是 Free
sum, err := svc.GetSummary(context.Background(), 1_000_001, "")
require.NoError(t, err)
require.Equal(t, domain.PlanStarter, sum.PlanID)
require.Equal(t, domain.Plans[domain.PlanStarter].MonthlyCredits, sum.Platform.CreditsTotal)
require.Equal(t, domain.Plans[domain.PlanStarter].MonthlyCredits, sum.TotalCredits)
require.Equal(t, domain.Plans[domain.PlanStarter].MonthlyCredits, sum.Platform.CreditsRemaining)
// 大小寫/空白 plan_id 也要生效
_, err = svc.PurchasePlan(context.Background(), 1_000_001, " Pro ", "mock-pay-2")
require.NoError(t, err)
sum, err = svc.GetSummary(context.Background(), 1_000_001, "")
require.NoError(t, err)
require.Equal(t, domain.PlanPro, sum.PlanID)
require.Equal(t, 2000, sum.Platform.CreditsTotal)
} }
func TestUSG_10_TenantAnalyticsSplit(t *testing.T) { func TestUSG_10_TenantAnalyticsSplit(t *testing.T) {

View File

@ -8,9 +8,10 @@ import (
"apps/backend/internal/domain" "apps/backend/internal/domain"
appnotifDomain "apps/backend/internal/module/appnotif/domain" appnotifDomain "apps/backend/internal/module/appnotif/domain"
billingDomain "apps/backend/internal/module/billing"
inspireDomain "apps/backend/internal/module/inspire/domain"
jobDomain "apps/backend/internal/module/job/domain" jobDomain "apps/backend/internal/module/job/domain"
memberDomain "apps/backend/internal/module/member/domain" memberDomain "apps/backend/internal/module/member/domain"
inspireDomain "apps/backend/internal/module/inspire/domain"
scoutDomain "apps/backend/internal/module/scout/domain" scoutDomain "apps/backend/internal/module/scout/domain"
studioDomain "apps/backend/internal/module/studio/domain" studioDomain "apps/backend/internal/module/studio/domain"
threadsDomain "apps/backend/internal/module/threads/domain" threadsDomain "apps/backend/internal/module/threads/domain"
@ -91,6 +92,18 @@ func mapError(err error) (int, Envelope) {
return be.http, Envelope{Code: be.code, Message: be.msg} return be.http, Envelope{Code: be.code, Message: be.msg}
} }
switch { switch {
case errors.Is(err, billingDomain.ErrDisabled):
return http.StatusServiceUnavailable, Envelope{Code: 503010, Message: err.Error()}
case errors.Is(err, billingDomain.ErrInvalidPlan), errors.Is(err, billingDomain.ErrInvalidRequest):
return http.StatusBadRequest, Envelope{Code: 400080, Message: err.Error()}
case errors.Is(err, billingDomain.ErrInvalidSignature):
return http.StatusBadRequest, Envelope{Code: 400081, Message: "invalid webhook signature"}
case errors.Is(err, billingDomain.ErrNotFound):
return http.StatusNotFound, Envelope{Code: 404080, Message: err.Error()}
case errors.Is(err, billingDomain.ErrConflict):
return http.StatusConflict, Envelope{Code: 409080, Message: err.Error()}
case errors.Is(err, billingDomain.ErrCustomerRequired):
return http.StatusConflict, Envelope{Code: 409081, Message: err.Error()}
case errors.Is(err, memberDomain.ErrBadPassword): case errors.Is(err, memberDomain.ErrBadPassword):
return http.StatusUnauthorized, Envelope{Code: 401010, Message: "invalid email or password"} return http.StatusUnauthorized, Envelope{Code: 401010, Message: "invalid email or password"}
case errors.Is(err, memberDomain.ErrSuspended): case errors.Is(err, memberDomain.ErrSuspended):
@ -103,6 +116,8 @@ func mapError(err error) (int, Envelope) {
case errors.Is(err, memberDomain.ErrInvalidCode): case errors.Is(err, memberDomain.ErrInvalidCode):
return http.StatusBadRequest, Envelope{Code: 400004, Message: "invalid or expired code"} return http.StatusBadRequest, Envelope{Code: 400004, Message: "invalid or expired code"}
case errors.Is(err, memberDomain.ErrCodeRateLimited):
return http.StatusTooManyRequests, Envelope{Code: 429001, Message: "verification code requested too frequently"}
case errors.Is(err, memberDomain.ErrEmailNotRegistered): case errors.Is(err, memberDomain.ErrEmailNotRegistered):
return http.StatusNotFound, Envelope{Code: 404002, Message: "email not registered"} return http.StatusNotFound, Envelope{Code: 404002, Message: "email not registered"}
case errors.Is(err, memberDomain.ErrNotFound): case errors.Is(err, memberDomain.ErrNotFound):
@ -152,6 +167,8 @@ func mapError(err error) (int, Envelope) {
return http.StatusForbidden, Envelope{Code: 403003, Message: "forbidden"} return http.StatusForbidden, Envelope{Code: 403003, Message: "forbidden"}
case errors.Is(err, usageDomain.ErrInvalidPlan): case errors.Is(err, usageDomain.ErrInvalidPlan):
return http.StatusBadRequest, Envelope{Code: 400041, Message: "invalid plan"} return http.StatusBadRequest, Envelope{Code: 400041, Message: "invalid plan"}
case errors.Is(err, usageDomain.ErrPurchaseDisabled):
return http.StatusGone, Envelope{Code: 410003, Message: err.Error()}
case errors.Is(err, studioDomain.ErrNotFound): case errors.Is(err, studioDomain.ErrNotFound):
return http.StatusNotFound, Envelope{Code: 404001, Message: "not found"} return http.StatusNotFound, Envelope{Code: 404001, Message: "not found"}
case errors.Is(err, studioDomain.ErrForbidden): case errors.Is(err, studioDomain.ErrForbidden):
@ -177,7 +194,7 @@ func mapError(err error) (int, Envelope) {
case errors.Is(err, scoutDomain.ErrHasProducts): case errors.Is(err, scoutDomain.ErrHasProducts):
return http.StatusConflict, Envelope{Code: 409030, Message: "brand has products; remove products first"} return http.StatusConflict, Envelope{Code: 409030, Message: "brand has products; remove products first"}
default: default:
return http.StatusInternalServerError, Envelope{Code: 500000, Message: err.Error()} return http.StatusInternalServerError, Envelope{Code: 500000, Message: "internal server error"}
} }
} }

View File

@ -3,6 +3,7 @@ package response
import ( import (
"context" "context"
"encoding/json" "encoding/json"
"errors"
"net/http" "net/http"
"net/http/httptest" "net/http/httptest"
"testing" "testing"
@ -48,3 +49,13 @@ func TestWrite_Error(t *testing.T) {
require.NoError(t, json.Unmarshal(rr.Body.Bytes(), &env)) require.NoError(t, json.Unmarshal(rr.Body.Bytes(), &env))
require.Equal(t, int64(400010), env.Code) require.Equal(t, int64(400010), env.Code)
} }
func TestWrite_InternalErrorDoesNotLeakDetails(t *testing.T) {
rr := httptest.NewRecorder()
Write(context.Background(), rr, nil, errors.New("mongodb password=secret"))
require.Equal(t, http.StatusInternalServerError, rr.Code)
var env Envelope
require.NoError(t, json.Unmarshal(rr.Body.Bytes(), &env))
require.Equal(t, "internal server error", env.Message)
require.NotContains(t, rr.Body.String(), "secret")
}

Some files were not shown because too many files have changed in this diff Show More