fix frontend
This commit is contained in:
parent
377b7f06ea
commit
fb8100825f
|
|
@ -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")},
|
||||||
},
|
},
|
||||||
}
|
}
|
||||||
|
|
|
||||||
|
|
@ -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"
|
||||||
)
|
)
|
||||||
|
|
||||||
// Worker:demo + threads_token_renew + persona_analyze_* + outbox publish
|
// Worker:demo + 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
|
||||||
|
|
|
||||||
|
|
@ -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: ""
|
||||||
|
|
||||||
# 驗證碼/重設密碼寄信
|
# 驗證碼/重設密碼寄信
|
||||||
#
|
#
|
||||||
# 本機(Mailpit,Mailgun 相容 SMTP 介面):
|
# 本機(Mailpit,Mailgun 相容 SMTP 介面):
|
||||||
|
|
|
||||||
|
|
@ -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 stream(goctl 不生成 streaming handler)
|
// 靈感 SSE stream(goctl 不生成 streaming handler)
|
||||||
handler.RegisterInspireStream(server, ctx)
|
handler.RegisterInspireStream(server, ctx)
|
||||||
|
|
|
||||||
|
|
@ -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
|
||||||
|
}
|
||||||
|
|
@ -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"
|
||||||
|
|
|
||||||
|
|
@ -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)
|
||||||
}
|
}
|
||||||
|
|
|
||||||
|
|
@ -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 seeder/init
|
- 在 migration 寫死帳號或密碼
|
||||||
- 手寫 migration runner
|
- 手寫 migration runner
|
||||||
- 把 stand-alone 電商 `.txt` 當 migrate 來源(格式不同;要用 JSON runCommand)
|
- 把 stand-alone 電商 `.txt` 當 migrate 來源(格式不同;要用 JSON runCommand)
|
||||||
|
|
|
||||||
|
|
@ -1,11 +1 @@
|
||||||
[
|
[]
|
||||||
{
|
|
||||||
"delete": "members",
|
|
||||||
"deletes": [
|
|
||||||
{
|
|
||||||
"q": { "uid": 1, "email": "admin@30cm.net" },
|
|
||||||
"limit": 1
|
|
||||||
}
|
|
||||||
]
|
|
||||||
}
|
|
||||||
]
|
|
||||||
|
|
|
||||||
|
|
@ -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": [
|
||||||
|
|
|
||||||
|
|
@ -1,11 +1 @@
|
||||||
[
|
[]
|
||||||
{
|
|
||||||
"delete": "identities",
|
|
||||||
"deletes": [
|
|
||||||
{
|
|
||||||
"q": { "uid": 1, "login_id": "admin@haixun.local", "platform": "platform" },
|
|
||||||
"limit": 1
|
|
||||||
}
|
|
||||||
]
|
|
||||||
}
|
|
||||||
]
|
|
||||||
|
|
|
||||||
|
|
@ -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
|
|
||||||
}
|
|
||||||
]
|
|
||||||
|
|
|
||||||
|
|
@ -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" }
|
||||||
|
]
|
||||||
|
|
@ -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"
|
||||||
|
}
|
||||||
|
]
|
||||||
|
}
|
||||||
|
]
|
||||||
|
|
@ -0,0 +1,3 @@
|
||||||
|
[
|
||||||
|
{ "dropIndexes": "studio_outbox", "index": "claim_due_outbox_steps" }
|
||||||
|
]
|
||||||
|
|
@ -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"
|
||||||
|
}
|
||||||
|
]
|
||||||
|
}
|
||||||
|
]
|
||||||
|
|
@ -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" }
|
||||||
|
]
|
||||||
|
|
@ -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 }
|
||||||
|
]
|
||||||
|
}
|
||||||
|
]
|
||||||
|
|
@ -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
|
||||||
|
|
|
||||||
|
|
@ -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=
|
||||||
|
|
|
||||||
|
|
@ -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 — 郵件 chrome(hermes);LogoURL 空則用 PublicWebBase + /brand-mark.jpg
|
// Brand — 郵件 chrome(hermes);LogoURL 空則用 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 != "" {
|
||||||
|
|
|
||||||
|
|
@ -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)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
@ -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)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
@ -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)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
@ -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)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
@ -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)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
@ -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)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
@ -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)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
@ -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",
|
||||||
|
|
|
||||||
|
|
@ -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)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
@ -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
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
|
||||||
|
|
@ -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")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
@ -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
|
||||||
|
}
|
||||||
|
|
@ -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)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
@ -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
|
||||||
}
|
}
|
||||||
|
|
|
||||||
|
|
@ -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 {
|
||||||
|
|
|
||||||
|
|
@ -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
|
||||||
}
|
}
|
||||||
|
|
|
||||||
|
|
@ -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
|
||||||
}
|
}
|
||||||
|
|
|
||||||
|
|
@ -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}
|
||||||
|
}
|
||||||
|
|
@ -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
|
||||||
|
}
|
||||||
|
|
@ -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
|
||||||
|
}
|
||||||
|
|
@ -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
|
||||||
|
}
|
||||||
|
|
@ -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)
|
||||||
|
}
|
||||||
|
|
@ -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
|
||||||
}
|
}
|
||||||
|
|
|
||||||
|
|
@ -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
|
||||||
|
}
|
||||||
|
|
|
||||||
|
|
@ -23,7 +23,10 @@ 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)
|
||||||
}
|
}
|
||||||
|
|
|
||||||
|
|
@ -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)
|
||||||
}
|
}
|
||||||
|
|
|
||||||
|
|
@ -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)
|
||||||
}
|
}
|
||||||
|
|
|
||||||
|
|
@ -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()
|
||||||
}
|
}
|
||||||
|
|
|
||||||
|
|
@ -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
|
||||||
|
}
|
||||||
|
|
@ -62,7 +62,7 @@ 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) {}
|
||||||
|
|
|
||||||
|
|
@ -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
|
||||||
|
}
|
||||||
|
|
@ -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)
|
||||||
|
}
|
||||||
|
|
@ -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
|
||||||
|
}
|
||||||
|
|
@ -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)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
@ -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
|
|
||||||
}
|
|
||||||
|
|
||||||
func NewRegistry() *Registry {
|
|
||||||
return &Registry{
|
|
||||||
clients: map[string]Client{
|
|
||||||
ProviderXAI: NewOpenAICompatible(ProviderXAI, "https://api.x.ai/v1"),
|
ProviderXAI: NewOpenAICompatible(ProviderXAI, "https://api.x.ai/v1"),
|
||||||
ProviderOpenCodeGo: NewOpenAICompatible(ProviderOpenCodeGo, "https://opencode.ai/zen/go/v1"),
|
ProviderOpenCodeGo: NewOpenAICompatible(ProviderOpenCodeGo, "https://opencode.ai/zen/go/v1"),
|
||||||
},
|
})
|
||||||
cache: map[string]modelsCacheEntry{},
|
}
|
||||||
|
|
||||||
|
// NewRegistryWithClients injects provider clients for focused tests.
|
||||||
|
func NewRegistryWithClients(cache ModelsCache, fingerprintSecret string, clients map[string]Client) *Registry {
|
||||||
|
return newRegistry(cache, fingerprintSecret, clients)
|
||||||
|
}
|
||||||
|
|
||||||
|
func newRegistry(cache ModelsCache, fingerprintSecret string, clients map[string]Client) *Registry {
|
||||||
|
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).
|
||||||
|
|
|
||||||
|
|
@ -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
|
||||||
|
|
@ -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
|
||||||
|
}
|
||||||
|
|
@ -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)
|
||||||
|
|
@ -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))
|
||||||
|
}
|
||||||
|
|
@ -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))
|
||||||
|
}
|
||||||
|
|
@ -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)
|
||||||
|
|
@ -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)
|
||||||
|
}
|
||||||
|
|
@ -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)
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -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
|
||||||
|
|
|
||||||
|
|
@ -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 == "" {
|
||||||
|
|
|
||||||
|
|
@ -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{
|
||||||
|
"$or": []bson.M{
|
||||||
|
{
|
||||||
"status": bson.M{"$in": []string{domain.StatusPending, domain.StatusQueued}},
|
"status": bson.M{"$in": []string{domain.StatusPending, domain.StatusQueued}},
|
||||||
"$or": []bson.M{
|
"$or": []bson.M{
|
||||||
{"run_after": bson.M{"$exists": false}},
|
{"run_after": bson.M{"$exists": false}},
|
||||||
{"run_after": 0},
|
{"run_after": 0},
|
||||||
{"run_after": bson.M{"$lte": now}},
|
{"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)
|
||||||
|
|
|
||||||
|
|
@ -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"
|
||||||
|
|
||||||
|
|
@ -20,10 +22,82 @@ 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
|
||||||
}
|
}
|
||||||
|
|
|
||||||
|
|
@ -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)
|
||||||
|
}
|
||||||
|
|
|
||||||
|
|
@ -24,6 +24,7 @@ var (
|
||||||
// ErrWeakPassword — API 400003;Message = PasswordPolicyEN
|
// ErrWeakPassword — API 400003;Message = 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")
|
||||||
|
|
|
||||||
|
|
@ -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)
|
||||||
|
|
|
||||||
|
|
@ -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,9 +30,15 @@ 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
|
||||||
|
|
@ -35,12 +46,15 @@ type MoncStore struct {
|
||||||
plain *mon.Model
|
plain *mon.Model
|
||||||
identity *monc.Model
|
identity *monc.Model
|
||||||
idPlain *mon.Model
|
idPlain *mon.Model
|
||||||
settings *monc.Model
|
settings *mon.Model
|
||||||
|
settingsCache settingsCache
|
||||||
|
settingsCacheNamespace string
|
||||||
refresh *mon.Model
|
refresh *mon.Model
|
||||||
codes *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,
|
||||||
|
|
@ -49,9 +63,12 @@ func NewMoncStore(uri, database string, cacheConf cache.CacheConf) *MoncStore {
|
||||||
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),
|
||||||
|
settingsCache: &redisSettingsCache{client: redis.MustNewRedis(settingsRedisConf)},
|
||||||
|
settingsCacheNamespace: normalizeSettingsCacheNamespace(redisNamespace),
|
||||||
refresh: mon.MustNewModel(uri, database, colRefresh),
|
refresh: mon.MustNewModel(uri, database, colRefresh),
|
||||||
codes: mon.MustNewModel(uri, database, colAuthCodes),
|
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
|
||||||
|
}
|
||||||
|
|
|
||||||
|
|
@ -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)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
@ -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")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
@ -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
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -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()
|
||||||
|
|
|
||||||
|
|
@ -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"`
|
||||||
|
Issuer string `json:"iss"`
|
||||||
|
Subject string `json:"sub"`
|
||||||
Email string `json:"email"`
|
Email string `json:"email"`
|
||||||
Name string `json:"name"`
|
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"`
|
||||||
}
|
}
|
||||||
|
|
|
||||||
|
|
@ -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))
|
||||||
|
}
|
||||||
|
|
@ -13,6 +13,7 @@ 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")
|
||||||
)
|
)
|
||||||
|
|
@ -176,6 +177,10 @@ 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 {
|
||||||
|
|
|
||||||
|
|
@ -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
|
||||||
|
|
|
||||||
|
|
@ -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 {
|
||||||
|
|
|
||||||
|
|
@ -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)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
@ -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()
|
||||||
|
|
|
||||||
|
|
@ -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))
|
||||||
|
|
|
||||||
|
|
@ -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)
|
||||||
|
}
|
||||||
|
|
@ -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 {
|
||||||
|
|
@ -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)
|
||||||
|
|
|
||||||
|
|
@ -4,6 +4,7 @@ import (
|
||||||
"context"
|
"context"
|
||||||
"crypto/rand"
|
"crypto/rand"
|
||||||
"encoding/json"
|
"encoding/json"
|
||||||
|
"errors"
|
||||||
"fmt"
|
"fmt"
|
||||||
"net/url"
|
"net/url"
|
||||||
"strings"
|
"strings"
|
||||||
|
|
@ -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,140 +1177,175 @@ 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
|
for i := range b.Steps {
|
||||||
|
st := &b.Steps[i]
|
||||||
|
claimable := st.Status == domain.StepScheduled && st.ScheduledAt <= now
|
||||||
|
stale := st.Status == domain.StepPublishing && st.LeaseExpiresAt <= now
|
||||||
|
if !claimable && !stale {
|
||||||
|
continue
|
||||||
|
}
|
||||||
root := findRoot(b)
|
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 {
|
|
||||||
st := &b.Steps[i]
|
|
||||||
if st.Status == domain.StepPublishing && b.UpdatedAt > 0 && now-b.UpdatedAt > publishingStaleNs {
|
|
||||||
logx.Infof("outbox reclaim stale publishing step id=%s bundle=%s", st.ID, b.ID)
|
|
||||||
st.Status = domain.StepScheduled
|
|
||||||
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
|
|
||||||
}
|
|
||||||
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 {
|
|
||||||
for j := range b.Steps {
|
|
||||||
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"
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
if err := fail(perr.Error()); err != nil {
|
||||||
|
return count, err
|
||||||
}
|
}
|
||||||
} 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 {
|
|
||||||
for j := range b.Steps {
|
|
||||||
if b.Steps[j].Kind == domain.StepReply && b.Steps[j].Status == domain.StepBlocked {
|
|
||||||
b.Steps[j].Status = domain.StepScheduled
|
|
||||||
b.Steps[j].Error = ""
|
|
||||||
}
|
}
|
||||||
|
fresh, getErr := s.Repo.GetOutbox(ctx, b.ID)
|
||||||
|
if getErr != nil {
|
||||||
|
return count, getErr
|
||||||
}
|
}
|
||||||
}
|
b = fresh
|
||||||
}
|
|
||||||
b.Status = recomputeBundle(b.Steps)
|
|
||||||
b.UpdatedAt = domain.NowNano()
|
|
||||||
_ = s.Repo.SaveOutbox(ctx, b)
|
|
||||||
|
|
||||||
// re-read root flags after mutation
|
|
||||||
root = findRoot(b)
|
|
||||||
rootFailed = root != nil && root.Status == domain.StepFailed
|
|
||||||
rootPublished = root != nil && root.Status == domain.StepPublished
|
|
||||||
}
|
}
|
||||||
return count, nil
|
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
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}()
|
||||||
|
res, err := s.Transport.Publish(publishCtx, request)
|
||||||
|
close(done)
|
||||||
|
<-stopped
|
||||||
|
select {
|
||||||
|
case leaseErr := <-renewErr:
|
||||||
|
return res, fmt.Errorf("%w: renew failed: %v", domain.ErrLeaseLost, leaseErr)
|
||||||
|
default:
|
||||||
|
return res, err
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
// ---------- Compose (CP) ----------
|
// ---------- Compose (CP) ----------
|
||||||
|
|
||||||
func (s *Service) Mimic(ctx context.Context, ownerUID int64, sourceText, personaID, structureNotes string) (string, error) {
|
func (s *Service) Mimic(ctx context.Context, ownerUID int64, sourceText, personaID, structureNotes string) (string, error) {
|
||||||
|
|
@ -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
|
||||||
|
|
|
||||||
|
|
@ -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
|
||||||
|
|
|
||||||
|
|
@ -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,10 +46,29 @@ 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"`
|
||||||
|
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"`
|
UpdatedAt int64 `bson:"updated_at" json:"updated_at"`
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
@ -165,3 +186,57 @@ type TenantSummary struct {
|
||||||
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
|
||||||
|
}
|
||||||
|
|
|
||||||
|
|
@ -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{
|
||||||
|
counters: map[string]*domain.MonthlyCounter{},
|
||||||
prefs: map[int64]*domain.MemberPrefs{},
|
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 ©Counter, 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()
|
||||||
|
|
|
||||||
|
|
@ -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})
|
||||||
|
|
|
||||||
|
|
@ -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
|
||||||
|
}
|
||||||
|
|
@ -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)
|
||||||
|
}
|
||||||
|
|
@ -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)
|
||||||
|
}
|
||||||
|
|
|
||||||
|
|
@ -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)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
@ -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)
|
||||||
}
|
}
|
||||||
|
|
|
||||||
|
|
@ -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 prefs.Unlimited {
|
||||||
|
totalLimit, meterLimit = -1, -1
|
||||||
}
|
}
|
||||||
if usedCredits+cost > capN {
|
return s.Repo.ReservePlatform(ctx, uid, domain.CurrentMonthKey(), meter, domain.MeterCost(meter), totalLimit, meterLimit)
|
||||||
return domain.ErrMeterCapExceeded
|
|
||||||
}
|
|
||||||
}
|
|
||||||
return nil
|
|
||||||
}
|
}
|
||||||
|
|
||||||
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,
|
||||||
|
|
|
||||||
|
|
@ -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) {
|
||||||
|
|
|
||||||
|
|
@ -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"}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -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
Loading…
Reference in New Issue