feat(usage): usage_daily 聚合 + 广告解锁 + free 额度校验 (tsk_1taxhtV2k3RP)
server/internal/usage 模块实现（仅字节/分钟，无目的地，符合无日志口径）： - aggregator.go：实现 #5 的 ReportUsage 回调接口（UsageReporter）， dp_uuid→user_id 走内存+Redis 短缓存，按上报批次 id Redis 去重防重放， 按 At 的 UTC 日期 INSERT ... ON DUPLICATE KEY UPDATE 累加，并发安全。 - store.go：usage_daily 累加/区间查询/当日查询、MarkAdUnlocked（事务+FOR UPDATE 幂等）、EffectivePlan（活跃订阅取最高 tier，无则回落 free）。 - ads.go：AdVerifier 接口 + AdMob SSV 实现（拉取并缓存 Google ECDSA 公钥， ECDSA-SHA256 验签 + custom_data/ad_unit 匹配），Unity 留接口占位。 - service.go：UsageCurve（空日补零、仅当前用户）、TodaySummary（/v1/me）、 UnlockAd（验签→ad_token 一次性 nonce 去重→写 ad_unlocked_at，当日二次幂等）。 - quota.go：CheckFreeConnect 供 #5 connect 使用：ad_gate=true 校验当日 ad_unlocked_at + minutes_used<daily_minutes(free=10)，返回剩余分钟； pro/team 不受 ad_gate 限制（返回 Unlimited）；带双语 code 错误。 - handler.go：GET /v1/usage?days=7、POST /v1/ads/unlock(204)。 - apierr：新增 AD_NOT_UNLOCKED/QUOTA_EXHAUSTED/AD_VERIFY_FAILED/AD_TOKEN_REPLAY。 - openapi：/ads/unlock 补 409 Conflict（重放）。 测试：ads_test.go 单测覆盖 AdMob 验签全链路（合法/伪造/custom_data 不符/ ad_unit 不允许/畸形 token）+ Unity 未实现；usage_integration_test.go (testcontainers, build tag integration) 覆盖并发多节点累加、批次重放去重、 跨 UTC 日界分桶、未知 dp_uuid 丢弃、free 额度四态、ads 解锁伪造/重放/幂等、 曲线补零/隔离、/me 摘要。 Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
wangjia
@@ -323,6 +323,8 @@ paths:
|
||||
$ref: "#/components/responses/Unauthorized"
|
||||
"403":
|
||||
$ref: "#/components/responses/Forbidden"
|
||||
"409":
|
||||
$ref: "#/components/responses/Conflict"
|
||||
"429":
|
||||
$ref: "#/components/responses/TooManyRequests"
|
||||
"500":
|
||||
|
||||
@@ -59,6 +59,28 @@ var (
|
||||
MessageEn: "Invalid request parameters",
|
||||
}
|
||||
|
||||
// Usage / ads / quota errors.
|
||||
ErrAdNotUnlocked = &Error{
|
||||
Code: "AD_NOT_UNLOCKED",
|
||||
MessageZH: "请先观看激励视频解锁当日时长",
|
||||
MessageEn: "Please watch the rewarded ad to unlock today's minutes",
|
||||
}
|
||||
ErrQuotaExhausted = &Error{
|
||||
Code: "QUOTA_EXHAUSTED",
|
||||
MessageZH: "今日免费时长已用尽，请明天再来或升级套餐",
|
||||
MessageEn: "Today's free minutes are used up, try tomorrow or upgrade",
|
||||
}
|
||||
ErrAdVerifyFailed = &Error{
|
||||
Code: "AD_VERIFY_FAILED",
|
||||
MessageZH: "广告回执校验失败",
|
||||
MessageEn: "Ad receipt verification failed",
|
||||
}
|
||||
ErrAdReplay = &Error{
|
||||
Code: "AD_TOKEN_REPLAY",
|
||||
MessageZH: "该广告回执已被使用",
|
||||
MessageEn: "This ad receipt has already been used",
|
||||
}
|
||||
|
||||
// Webhook-specific errors.
|
||||
ErrWebhookSignature = &Error{
|
||||
Code: "WEBHOOK_INVALID_SIGNATURE",
|
||||
|
||||
@@ -0,0 +1,267 @@
|
||||
package usage
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/ecdsa"
|
||||
"crypto/sha256"
|
||||
"crypto/x509"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"encoding/pem"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
// AdVerifyRequest carries the inputs an AdVerifier needs to authenticate a
|
||||
// rewarded-ad receipt.
|
||||
type AdVerifyRequest struct {
|
||||
// UserID is the authenticated user requesting the unlock.
|
||||
UserID int64
|
||||
// DeviceID is the device UUID that played the ad.
|
||||
DeviceID string
|
||||
// AdToken is the provider's server-side verification receipt. For AdMob
|
||||
// SSV this is the full callback query string (everything after the '?').
|
||||
AdToken string
|
||||
// ExpectedCustomData, when non-empty, must match the receipt's custom_data
|
||||
// field (binds the receipt to this user).
|
||||
ExpectedCustomData string
|
||||
}
|
||||
|
||||
// AdVerifier authenticates a rewarded-ad receipt with the ad network.
|
||||
// Implementations must be safe for concurrent use.
|
||||
type AdVerifier interface {
|
||||
// Verify returns nil if the receipt is genuine and bound to the request.
|
||||
Verify(ctx context.Context, req AdVerifyRequest) error
|
||||
// Provider names the ad network (e.g. "admob", "unity").
|
||||
Provider() string
|
||||
}
|
||||
|
||||
// --------------------------------------------------------------------------
|
||||
// AdMob Server-Side Verification (SSV)
|
||||
// --------------------------------------------------------------------------
|
||||
|
||||
// DefaultAdMobKeyServerURL is Google's public ECDSA key endpoint for AdMob SSV.
|
||||
const DefaultAdMobKeyServerURL = "https://www.gstatic.com/admob/reward/verifier-keys.json"
|
||||
|
||||
// AdMobVerifier verifies AdMob SSV callbacks. It fetches and caches Google's
|
||||
// ECDSA public keys, then checks the ECDSA-SHA256 signature over the callback
|
||||
// content, that the receipt's custom_data matches the user, and (optionally)
|
||||
// that the ad_unit is one we recognise.
|
||||
type AdMobVerifier struct {
|
||||
keyServerURL string
|
||||
httpClient *http.Client
|
||||
allowedAdUnits map[string]struct{}
|
||||
keyTTL time.Duration
|
||||
|
||||
mu sync.RWMutex
|
||||
keys map[string]*ecdsa.PublicKey
|
||||
fetchedAt time.Time
|
||||
}
|
||||
|
||||
// NewAdMobVerifier creates an AdMobVerifier. keyServerURL empty → the default
|
||||
// Google endpoint. keyTTL <= 0 → 12h. allowedAdUnits empty → any ad unit is
|
||||
// accepted (only the signature and custom_data are enforced).
|
||||
func NewAdMobVerifier(keyServerURL string, httpClient *http.Client, keyTTL time.Duration, allowedAdUnits []string) *AdMobVerifier {
|
||||
if keyServerURL == "" {
|
||||
keyServerURL = DefaultAdMobKeyServerURL
|
||||
}
|
||||
if httpClient == nil {
|
||||
httpClient = &http.Client{Timeout: 5 * time.Second}
|
||||
}
|
||||
if keyTTL <= 0 {
|
||||
keyTTL = 12 * time.Hour
|
||||
}
|
||||
allowed := make(map[string]struct{}, len(allowedAdUnits))
|
||||
for _, u := range allowedAdUnits {
|
||||
allowed[u] = struct{}{}
|
||||
}
|
||||
return &AdMobVerifier{
|
||||
keyServerURL: keyServerURL,
|
||||
httpClient: httpClient,
|
||||
allowedAdUnits: allowed,
|
||||
keyTTL: keyTTL,
|
||||
}
|
||||
}
|
||||
|
||||
// Provider implements AdVerifier.
|
||||
func (v *AdMobVerifier) Provider() string { return "admob" }
|
||||
|
||||
// errAdMobVerify is returned for any receipt that fails authentication.
|
||||
var errAdMobVerify = errors.New("admob: receipt verification failed")
|
||||
|
||||
// Verify implements AdVerifier for AdMob SSV.
|
||||
//
|
||||
// AdMob signs the callback by computing ECDSA-SHA256 over the query string up
|
||||
// to (but excluding) "&signature="; the signature (web-safe base64, ASN.1 DER)
|
||||
// and key_id follow. See Google's SSV documentation.
|
||||
func (v *AdMobVerifier) Verify(ctx context.Context, req AdVerifyRequest) error {
|
||||
raw := strings.TrimPrefix(req.AdToken, "?")
|
||||
if raw == "" {
|
||||
return errAdMobVerify
|
||||
}
|
||||
|
||||
// Content to verify = everything before "&signature=".
|
||||
sigMarker := strings.Index(raw, "&signature=")
|
||||
if sigMarker < 0 {
|
||||
return errAdMobVerify
|
||||
}
|
||||
content := raw[:sigMarker]
|
||||
|
||||
params, err := url.ParseQuery(raw)
|
||||
if err != nil {
|
||||
return errAdMobVerify
|
||||
}
|
||||
sigB64 := params.Get("signature")
|
||||
keyID := params.Get("key_id")
|
||||
if sigB64 == "" || keyID == "" {
|
||||
return errAdMobVerify
|
||||
}
|
||||
|
||||
// Bind the receipt to the user / ad unit.
|
||||
if req.ExpectedCustomData != "" && params.Get("custom_data") != req.ExpectedCustomData {
|
||||
return errAdMobVerify
|
||||
}
|
||||
if len(v.allowedAdUnits) > 0 {
|
||||
if _, ok := v.allowedAdUnits[params.Get("ad_unit")]; !ok {
|
||||
return errAdMobVerify
|
||||
}
|
||||
}
|
||||
|
||||
sig, err := base64.RawURLEncoding.DecodeString(sigB64)
|
||||
if err != nil {
|
||||
// Some senders include padding; fall back to standard URL encoding.
|
||||
sig, err = base64.URLEncoding.DecodeString(sigB64)
|
||||
if err != nil {
|
||||
return errAdMobVerify
|
||||
}
|
||||
}
|
||||
|
||||
pub, err := v.publicKey(ctx, keyID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
digest := sha256.Sum256([]byte(content))
|
||||
if !ecdsa.VerifyASN1(pub, digest[:], sig) {
|
||||
return errAdMobVerify
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// publicKey returns the cached ECDSA public key for keyID, refreshing the key
|
||||
// set from the server when the cache is empty, stale, or missing the key.
|
||||
func (v *AdMobVerifier) publicKey(ctx context.Context, keyID string) (*ecdsa.PublicKey, error) {
|
||||
v.mu.RLock()
|
||||
pub, ok := v.keys[keyID]
|
||||
fresh := time.Since(v.fetchedAt) < v.keyTTL
|
||||
v.mu.RUnlock()
|
||||
if ok && fresh {
|
||||
return pub, nil
|
||||
}
|
||||
|
||||
if err := v.refreshKeys(ctx); err != nil {
|
||||
// Serve a stale cached key rather than fail outright on a transient
|
||||
// fetch error.
|
||||
if ok {
|
||||
return pub, nil
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
|
||||
v.mu.RLock()
|
||||
pub, ok = v.keys[keyID]
|
||||
v.mu.RUnlock()
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("admob: unknown key_id %q", keyID)
|
||||
}
|
||||
return pub, nil
|
||||
}
|
||||
|
||||
// admobKeySet mirrors the JSON returned by the AdMob verifier-keys endpoint.
|
||||
type admobKeySet struct {
|
||||
Keys []struct {
|
||||
KeyID json.Number `json:"keyId"`
|
||||
PEM string `json:"pem"`
|
||||
Base64 string `json:"base64"`
|
||||
} `json:"keys"`
|
||||
}
|
||||
|
||||
// refreshKeys fetches and parses the public key set from the key server.
|
||||
func (v *AdMobVerifier) refreshKeys(ctx context.Context) error {
|
||||
httpReq, err := http.NewRequestWithContext(ctx, http.MethodGet, v.keyServerURL, nil)
|
||||
if err != nil {
|
||||
return fmt.Errorf("admob refreshKeys: %w", err)
|
||||
}
|
||||
resp, err := v.httpClient.Do(httpReq)
|
||||
if err != nil {
|
||||
return fmt.Errorf("admob refreshKeys: %w", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return fmt.Errorf("admob refreshKeys: status %d", resp.StatusCode)
|
||||
}
|
||||
|
||||
var ks admobKeySet
|
||||
if err := json.NewDecoder(resp.Body).Decode(&ks); err != nil {
|
||||
return fmt.Errorf("admob refreshKeys decode: %w", err)
|
||||
}
|
||||
|
||||
parsed := make(map[string]*ecdsa.PublicKey, len(ks.Keys))
|
||||
for _, k := range ks.Keys {
|
||||
pub, err := parseECDSAPublicKey(k.PEM)
|
||||
if err != nil {
|
||||
continue // skip unparseable keys rather than failing the whole set
|
||||
}
|
||||
parsed[k.KeyID.String()] = pub
|
||||
}
|
||||
if len(parsed) == 0 {
|
||||
return errors.New("admob refreshKeys: no usable keys")
|
||||
}
|
||||
|
||||
v.mu.Lock()
|
||||
v.keys = parsed
|
||||
v.fetchedAt = time.Now()
|
||||
v.mu.Unlock()
|
||||
return nil
|
||||
}
|
||||
|
||||
// parseECDSAPublicKey parses a PEM-encoded PKIX ECDSA public key.
|
||||
func parseECDSAPublicKey(pemStr string) (*ecdsa.PublicKey, error) {
|
||||
block, _ := pem.Decode([]byte(pemStr))
|
||||
if block == nil {
|
||||
return nil, errors.New("admob: invalid PEM")
|
||||
}
|
||||
anyKey, err := x509.ParsePKIXPublicKey(block.Bytes)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("admob: parse pkix: %w", err)
|
||||
}
|
||||
pub, ok := anyKey.(*ecdsa.PublicKey)
|
||||
if !ok {
|
||||
return nil, errors.New("admob: not an ECDSA key")
|
||||
}
|
||||
return pub, nil
|
||||
}
|
||||
|
||||
// --------------------------------------------------------------------------
|
||||
// Unity Ads (interface placeholder)
|
||||
// --------------------------------------------------------------------------
|
||||
|
||||
// UnityVerifier is a placeholder for Unity Ads SSV; the interface is wired now
|
||||
// and the verification logic will be implemented when Unity is integrated.
|
||||
type UnityVerifier struct{}
|
||||
|
||||
// Provider implements AdVerifier.
|
||||
func (UnityVerifier) Provider() string { return "unity" }
|
||||
|
||||
// errUnityUnimplemented signals that Unity SSV is not yet available.
|
||||
var errUnityUnimplemented = errors.New("unity: SSV verification not implemented")
|
||||
|
||||
// Verify implements AdVerifier.
|
||||
func (UnityVerifier) Verify(context.Context, AdVerifyRequest) error {
|
||||
return errUnityUnimplemented
|
||||
}
|
||||
@@ -0,0 +1,138 @@
|
||||
package usage
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/ecdsa"
|
||||
"crypto/elliptic"
|
||||
"crypto/rand"
|
||||
"crypto/sha256"
|
||||
"crypto/x509"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"encoding/pem"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
// adMobTestKeyServer spins up an httptest server that serves a single ECDSA
|
||||
// public key in the AdMob verifier-keys.json format, and returns the server,
|
||||
// the private key, and the key id.
|
||||
func adMobTestKeyServer(t *testing.T) (*httptest.Server, *ecdsa.PrivateKey, string) {
|
||||
t.Helper()
|
||||
priv, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
|
||||
if err != nil {
|
||||
t.Fatalf("generate key: %v", err)
|
||||
}
|
||||
der, err := x509.MarshalPKIXPublicKey(&priv.PublicKey)
|
||||
if err != nil {
|
||||
t.Fatalf("marshal pub: %v", err)
|
||||
}
|
||||
pemBytes := pem.EncodeToMemory(&pem.Block{Type: "PUBLIC KEY", Bytes: der})
|
||||
|
||||
const keyID = "3335741209"
|
||||
body := map[string]any{
|
||||
"keys": []map[string]any{
|
||||
{"keyId": 3335741209, "pem": string(pemBytes), "base64": base64.StdEncoding.EncodeToString(der)},
|
||||
},
|
||||
}
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
_ = json.NewEncoder(w).Encode(body)
|
||||
}))
|
||||
t.Cleanup(srv.Close)
|
||||
return srv, priv, keyID
|
||||
}
|
||||
|
||||
// signAdMob builds a signed AdMob SSV callback query string.
|
||||
func signAdMob(t *testing.T, priv *ecdsa.PrivateKey, keyID, content string) string {
|
||||
t.Helper()
|
||||
digest := sha256.Sum256([]byte(content))
|
||||
sig, err := ecdsa.SignASN1(rand.Reader, priv, digest[:])
|
||||
if err != nil {
|
||||
t.Fatalf("sign: %v", err)
|
||||
}
|
||||
sigB64 := base64.RawURLEncoding.EncodeToString(sig)
|
||||
return fmt.Sprintf("%s&signature=%s&key_id=%s", content, sigB64, keyID)
|
||||
}
|
||||
|
||||
func TestAdMobVerifyValid(t *testing.T) {
|
||||
srv, priv, keyID := adMobTestKeyServer(t)
|
||||
v := NewAdMobVerifier(srv.URL, srv.Client(), time.Hour, []string{"ca-app-pub-123/456"})
|
||||
|
||||
content := "ad_network=admob&ad_unit=ca-app-pub-123/456&custom_data=user-uuid-42&reward_amount=1&reward_item=minutes×tamp=1700000000&transaction_id=abc123&user_id=device-1"
|
||||
token := signAdMob(t, priv, keyID, content)
|
||||
|
||||
err := v.Verify(context.Background(), AdVerifyRequest{
|
||||
UserID: 42,
|
||||
DeviceID: "device-1",
|
||||
AdToken: token,
|
||||
ExpectedCustomData: "user-uuid-42",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("expected valid receipt, got %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAdMobVerifyForgedSignature(t *testing.T) {
|
||||
srv, priv, keyID := adMobTestKeyServer(t)
|
||||
v := NewAdMobVerifier(srv.URL, srv.Client(), time.Hour, nil)
|
||||
|
||||
content := "ad_unit=ca-app-pub-123/456&custom_data=u1&transaction_id=abc"
|
||||
token := signAdMob(t, priv, keyID, content)
|
||||
// Tamper with the signed content after signing → signature no longer matches.
|
||||
tampered := "ad_unit=ca-app-pub-123/456&custom_data=u1&transaction_id=EVIL" + token[len(content):]
|
||||
|
||||
if err := v.Verify(context.Background(), AdVerifyRequest{AdToken: tampered}); err == nil {
|
||||
t.Fatal("expected forged signature to be rejected")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAdMobVerifyCustomDataMismatch(t *testing.T) {
|
||||
srv, priv, keyID := adMobTestKeyServer(t)
|
||||
v := NewAdMobVerifier(srv.URL, srv.Client(), time.Hour, nil)
|
||||
|
||||
content := "ad_unit=ca-app-pub-123/456&custom_data=u1&transaction_id=abc"
|
||||
token := signAdMob(t, priv, keyID, content)
|
||||
|
||||
if err := v.Verify(context.Background(), AdVerifyRequest{
|
||||
AdToken: token,
|
||||
ExpectedCustomData: "someone-else",
|
||||
}); err == nil {
|
||||
t.Fatal("expected custom_data mismatch to be rejected")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAdMobVerifyDisallowedAdUnit(t *testing.T) {
|
||||
srv, priv, keyID := adMobTestKeyServer(t)
|
||||
v := NewAdMobVerifier(srv.URL, srv.Client(), time.Hour, []string{"ca-app-pub-allowed/1"})
|
||||
|
||||
content := "ad_unit=ca-app-pub-OTHER/9&custom_data=u1&transaction_id=abc"
|
||||
token := signAdMob(t, priv, keyID, content)
|
||||
|
||||
if err := v.Verify(context.Background(), AdVerifyRequest{AdToken: token}); err == nil {
|
||||
t.Fatal("expected disallowed ad_unit to be rejected")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAdMobVerifyMalformedToken(t *testing.T) {
|
||||
srv, _, _ := adMobTestKeyServer(t)
|
||||
v := NewAdMobVerifier(srv.URL, srv.Client(), time.Hour, nil)
|
||||
|
||||
for _, tok := range []string{"", "no-signature-here", "?garbage"} {
|
||||
if err := v.Verify(context.Background(), AdVerifyRequest{AdToken: tok}); err == nil {
|
||||
t.Errorf("expected malformed token %q to be rejected", tok)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestUnityVerifierUnimplemented(t *testing.T) {
|
||||
var v AdVerifier = UnityVerifier{}
|
||||
if v.Provider() != "unity" {
|
||||
t.Errorf("provider = %q, want unity", v.Provider())
|
||||
}
|
||||
if err := v.Verify(context.Background(), AdVerifyRequest{}); err == nil {
|
||||
t.Error("expected Unity verifier to report unimplemented")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,184 @@
|
||||
package usage
|
||||
|
||||
import (
|
||||
"context"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/redis/go-redis/v9"
|
||||
)
|
||||
|
||||
// Report is one usage delta pushed by the data plane (#5's ReportUsage path).
|
||||
// It identifies the user only by data-plane UUID; the aggregator resolves it to
|
||||
// a user_id internally. No destination, domain, or DNS data is ever carried.
|
||||
type Report struct {
|
||||
// DPUUID is the data-plane credential UUID the node observed.
|
||||
DPUUID string
|
||||
// BytesUp / BytesDown are the incremental byte counts since the last report.
|
||||
BytesUp uint64
|
||||
BytesDown uint64
|
||||
// Minutes is the incremental session minutes since the last report.
|
||||
Minutes int
|
||||
// At is the event time used for UTC day bucketing. Zero = time.Now().
|
||||
At time.Time
|
||||
// BatchID is #5's upload batch identifier, used for replay-idempotent
|
||||
// de-duplication. Empty disables dedup for this report.
|
||||
BatchID string
|
||||
}
|
||||
|
||||
// UsageReporter is the callback interface #5 invokes when a node reports usage.
|
||||
// Aggregator satisfies it, so #5 can depend on the interface and be wired to a
|
||||
// real or mock implementation.
|
||||
type UsageReporter interface {
|
||||
ReportUsage(ctx context.Context, r Report) error
|
||||
}
|
||||
|
||||
// Aggregator implements UsageReporter: it resolves dp_uuid → user_id (with a
|
||||
// short in-memory + Redis cache), de-duplicates replayed batches, and folds the
|
||||
// delta into usage_daily by UTC date.
|
||||
type Aggregator struct {
|
||||
store *Store
|
||||
rdb *redis.Client
|
||||
cache *dpCache
|
||||
batchTTL time.Duration
|
||||
}
|
||||
|
||||
// NewAggregator builds an Aggregator. rdb may be nil (batch dedup and the
|
||||
// Redis tier of the dp_uuid cache are then disabled; the in-memory cache and
|
||||
// DB lookups still work). dpCacheTTL <= 0 defaults to 5 minutes; batchTTL <= 0
|
||||
// defaults to 10 minutes.
|
||||
func NewAggregator(store *Store, rdb *redis.Client, dpCacheTTL, batchTTL time.Duration) *Aggregator {
|
||||
if dpCacheTTL <= 0 {
|
||||
dpCacheTTL = 5 * time.Minute
|
||||
}
|
||||
if batchTTL <= 0 {
|
||||
batchTTL = 10 * time.Minute
|
||||
}
|
||||
return &Aggregator{
|
||||
store: store,
|
||||
rdb: rdb,
|
||||
cache: newDPCache(dpCacheTTL),
|
||||
batchTTL: batchTTL,
|
||||
}
|
||||
}
|
||||
|
||||
// redisKeyBatch returns the Redis key marking a processed upload batch.
|
||||
func redisKeyBatch(batchID string) string { return "usage:batch:" + batchID }
|
||||
|
||||
// redisKeyDP returns the Redis key caching a dp_uuid → user_id resolution.
|
||||
func redisKeyDP(dpUUID string) string { return "usage:dp:" + dpUUID }
|
||||
|
||||
// ReportUsage folds one report into usage_daily. It is safe for concurrent
|
||||
// use and idempotent across replays of the same BatchID.
|
||||
func (a *Aggregator) ReportUsage(ctx context.Context, r Report) error {
|
||||
if r.DPUUID == "" {
|
||||
return nil
|
||||
}
|
||||
if r.BytesUp == 0 && r.BytesDown == 0 && r.Minutes == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
// Replay guard: a batch already seen is a no-op.
|
||||
if dup, err := a.batchSeen(ctx, r.BatchID); err != nil {
|
||||
return err
|
||||
} else if dup {
|
||||
return nil
|
||||
}
|
||||
|
||||
userID, err := a.resolveUser(ctx, r.DPUUID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if userID == 0 {
|
||||
// Unknown/rotated dp_uuid: drop silently (no log of the mapping).
|
||||
return nil
|
||||
}
|
||||
|
||||
at := r.At
|
||||
if at.IsZero() {
|
||||
at = time.Now()
|
||||
}
|
||||
return a.store.AggregateUsage(ctx, userID, at.UTC(), r.BytesUp, r.BytesDown, r.Minutes)
|
||||
}
|
||||
|
||||
// batchSeen marks the batch as processed and reports whether it was already
|
||||
// seen. Returns false when batchID is empty or Redis is unavailable (dedup is
|
||||
// best-effort; the DB accumulation itself is additive, not idempotent, so #5
|
||||
// must supply a BatchID + Redis to get exactly-once semantics).
|
||||
func (a *Aggregator) batchSeen(ctx context.Context, batchID string) (bool, error) {
|
||||
if batchID == "" || a.rdb == nil {
|
||||
return false, nil
|
||||
}
|
||||
set, err := a.rdb.SetNX(ctx, redisKeyBatch(batchID), "1", a.batchTTL).Result()
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
// SetNX returns true when the key was newly set (i.e. not a duplicate).
|
||||
return !set, nil
|
||||
}
|
||||
|
||||
// resolveUser maps a dp_uuid to a user_id via the in-memory cache, then Redis,
|
||||
// then the database (populating both caches on a DB hit).
|
||||
func (a *Aggregator) resolveUser(ctx context.Context, dpUUID string) (int64, error) {
|
||||
if id, ok := a.cache.get(dpUUID); ok {
|
||||
return id, nil
|
||||
}
|
||||
|
||||
if a.rdb != nil {
|
||||
if id, err := a.rdb.Get(ctx, redisKeyDP(dpUUID)).Int64(); err == nil {
|
||||
a.cache.set(dpUUID, id)
|
||||
return id, nil
|
||||
} else if err != redis.Nil {
|
||||
return 0, err
|
||||
}
|
||||
}
|
||||
|
||||
id, err := a.store.LookupUserIDByDPUUID(ctx, dpUUID)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
if id == 0 {
|
||||
return 0, nil
|
||||
}
|
||||
|
||||
a.cache.set(dpUUID, id)
|
||||
if a.rdb != nil {
|
||||
_ = a.rdb.Set(ctx, redisKeyDP(dpUUID), id, a.cache.ttl).Err()
|
||||
}
|
||||
return id, nil
|
||||
}
|
||||
|
||||
// --------------------------------------------------------------------------
|
||||
// in-memory dp_uuid → user_id cache (short TTL)
|
||||
// --------------------------------------------------------------------------
|
||||
|
||||
type dpCacheEntry struct {
|
||||
userID int64
|
||||
expires time.Time
|
||||
}
|
||||
|
||||
type dpCache struct {
|
||||
mu sync.RWMutex
|
||||
m map[string]dpCacheEntry
|
||||
ttl time.Duration
|
||||
}
|
||||
|
||||
func newDPCache(ttl time.Duration) *dpCache {
|
||||
return &dpCache{m: make(map[string]dpCacheEntry), ttl: ttl}
|
||||
}
|
||||
|
||||
func (c *dpCache) get(dpUUID string) (int64, bool) {
|
||||
c.mu.RLock()
|
||||
e, ok := c.m[dpUUID]
|
||||
c.mu.RUnlock()
|
||||
if !ok || time.Now().After(e.expires) {
|
||||
return 0, false
|
||||
}
|
||||
return e.userID, true
|
||||
}
|
||||
|
||||
func (c *dpCache) set(dpUUID string, userID int64) {
|
||||
c.mu.Lock()
|
||||
c.m[dpUUID] = dpCacheEntry{userID: userID, expires: time.Now().Add(c.ttl)}
|
||||
c.mu.Unlock()
|
||||
}
|
||||
@@ -0,0 +1,134 @@
|
||||
package usage
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"strconv"
|
||||
|
||||
"github.com/wangjia/pangolin/server/internal/apierr"
|
||||
)
|
||||
|
||||
// ctxKey is the context key type for request-scoped values. Defined here to
|
||||
// avoid an import cycle; the JWT auth middleware sets the same key.
|
||||
type ctxKey string
|
||||
|
||||
// CtxKeyUserID is the context key carrying the authenticated int64 user ID.
|
||||
const CtxKeyUserID ctxKey = "user_id"
|
||||
|
||||
// userIDFromContext extracts the authenticated user ID set by the auth
|
||||
// middleware, or 0 if absent.
|
||||
func userIDFromContext(r *http.Request) int64 {
|
||||
id, _ := r.Context().Value(CtxKeyUserID).(int64)
|
||||
return id
|
||||
}
|
||||
|
||||
// UsageHandler serves GET /v1/usage?days=N.
|
||||
type UsageHandler struct {
|
||||
svc *Service
|
||||
}
|
||||
|
||||
// NewUsageHandler creates a UsageHandler.
|
||||
func NewUsageHandler(svc *Service) *UsageHandler { return &UsageHandler{svc: svc} }
|
||||
|
||||
type usageResponse struct {
|
||||
Points []UsagePoint `json:"points"`
|
||||
}
|
||||
|
||||
// ServeHTTP implements http.Handler.
|
||||
func (h *UsageHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Method != http.MethodGet {
|
||||
http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
|
||||
return
|
||||
}
|
||||
userID := userIDFromContext(r)
|
||||
if userID == 0 {
|
||||
apierr.WriteJSON(w, http.StatusUnauthorized, &apierr.Error{
|
||||
Code: "UNAUTHORIZED",
|
||||
MessageZH: "请先登录",
|
||||
MessageEn: "Authentication required",
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
days := 7
|
||||
if v := r.URL.Query().Get("days"); v != "" {
|
||||
n, err := strconv.Atoi(v)
|
||||
if err != nil || n < 1 || n > 90 {
|
||||
apierr.WriteJSON(w, http.StatusBadRequest, apierr.ErrBadRequest)
|
||||
return
|
||||
}
|
||||
days = n
|
||||
}
|
||||
|
||||
points, apiErr := h.svc.UsageCurve(r.Context(), userID, days)
|
||||
if apiErr != nil {
|
||||
apierr.WriteJSON(w, http.StatusInternalServerError, apiErr)
|
||||
return
|
||||
}
|
||||
|
||||
w.Header().Set("Content-Type", "application/json; charset=utf-8")
|
||||
w.WriteHeader(http.StatusOK)
|
||||
_ = json.NewEncoder(w).Encode(usageResponse{Points: points})
|
||||
}
|
||||
|
||||
// AdsUnlockHandler serves POST /v1/ads/unlock.
|
||||
type AdsUnlockHandler struct {
|
||||
svc *Service
|
||||
}
|
||||
|
||||
// NewAdsUnlockHandler creates an AdsUnlockHandler.
|
||||
func NewAdsUnlockHandler(svc *Service) *AdsUnlockHandler { return &AdsUnlockHandler{svc: svc} }
|
||||
|
||||
type adsUnlockRequest struct {
|
||||
DeviceID string `json:"device_id"`
|
||||
AdToken string `json:"ad_token"`
|
||||
}
|
||||
|
||||
// ServeHTTP implements http.Handler. On success it returns 204 No Content
|
||||
// (matching the OpenAPI contract), including the idempotent already-unlocked
|
||||
// case.
|
||||
func (h *AdsUnlockHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Method != http.MethodPost {
|
||||
http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
|
||||
return
|
||||
}
|
||||
userID := userIDFromContext(r)
|
||||
if userID == 0 {
|
||||
apierr.WriteJSON(w, http.StatusUnauthorized, &apierr.Error{
|
||||
Code: "UNAUTHORIZED",
|
||||
MessageZH: "请先登录",
|
||||
MessageEn: "Authentication required",
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
var req adsUnlockRequest
|
||||
if err := json.NewDecoder(http.MaxBytesReader(w, r.Body, 16*1024)).Decode(&req); err != nil {
|
||||
apierr.WriteJSON(w, http.StatusBadRequest, apierr.ErrBadRequest)
|
||||
return
|
||||
}
|
||||
|
||||
_, apiErr := h.svc.UnlockAd(r.Context(), userID, req.DeviceID, req.AdToken)
|
||||
if apiErr != nil {
|
||||
apierr.WriteJSON(w, adsErrorStatus(apiErr.Code), apiErr)
|
||||
return
|
||||
}
|
||||
|
||||
w.WriteHeader(http.StatusNoContent)
|
||||
}
|
||||
|
||||
// adsErrorStatus maps an ads-unlock error code to an HTTP status.
|
||||
func adsErrorStatus(code string) int {
|
||||
switch code {
|
||||
case "AD_VERIFY_FAILED":
|
||||
return http.StatusForbidden
|
||||
case "AD_TOKEN_REPLAY":
|
||||
return http.StatusConflict
|
||||
case "BAD_REQUEST":
|
||||
return http.StatusBadRequest
|
||||
case "INTERNAL_ERROR":
|
||||
return http.StatusInternalServerError
|
||||
default:
|
||||
return http.StatusBadRequest
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,66 @@
|
||||
package usage
|
||||
|
||||
import (
|
||||
"context"
|
||||
|
||||
"github.com/wangjia/pangolin/server/internal/apierr"
|
||||
)
|
||||
|
||||
// defaultFreeDailyMinutes is the口径 fallback when the free plan row has a NULL
|
||||
// daily_minutes (should not happen given the seed, but we stay safe).
|
||||
const defaultFreeDailyMinutes = 10
|
||||
|
||||
// CheckFreeConnect enforces the free-plan connect gate and returns the user's
|
||||
// remaining minutes for today (UTC). It is called by #5's connect endpoint to
|
||||
// derive the free credential's TTL.
|
||||
//
|
||||
// Rules:
|
||||
// - plan.ad_gate == false (pro/team): not minute-gated. Returns Unlimited
|
||||
// when daily_minutes is NULL, otherwise the remaining minutes for the day.
|
||||
// - plan.ad_gate == true (free): today's ad_unlocked_at must be set and
|
||||
// minutes_used must be below daily_minutes (free = 10). Returns the
|
||||
// remaining minutes; otherwise a bilingual semantic error
|
||||
// (AD_NOT_UNLOCKED / QUOTA_EXHAUSTED).
|
||||
func (svc *Service) CheckFreeConnect(ctx context.Context, userID int64) (remainingMinutes int, apiErr *apierr.Error) {
|
||||
plan, err := svc.store.EffectivePlan(ctx, userID)
|
||||
if err != nil {
|
||||
return 0, apierr.ErrInternal
|
||||
}
|
||||
|
||||
limit := defaultFreeDailyMinutes
|
||||
if plan.DailyMinutes.Valid {
|
||||
limit = int(plan.DailyMinutes.Int64)
|
||||
}
|
||||
|
||||
// Paid plans (no ad gate).
|
||||
if !plan.AdGate {
|
||||
if !plan.DailyMinutes.Valid {
|
||||
return Unlimited, nil
|
||||
}
|
||||
day, err := svc.store.GetDay(ctx, userID, utcToday())
|
||||
if err != nil {
|
||||
return 0, apierr.ErrInternal
|
||||
}
|
||||
remaining := limit
|
||||
if day != nil {
|
||||
remaining = limit - day.MinutesUsed
|
||||
}
|
||||
if remaining < 0 {
|
||||
remaining = 0
|
||||
}
|
||||
return remaining, nil
|
||||
}
|
||||
|
||||
// Free plan: require ad unlock + remaining minutes.
|
||||
day, err := svc.store.GetDay(ctx, userID, utcToday())
|
||||
if err != nil {
|
||||
return 0, apierr.ErrInternal
|
||||
}
|
||||
if day == nil || !day.AdUnlockedAt.Valid {
|
||||
return 0, apierr.ErrAdNotUnlocked
|
||||
}
|
||||
if day.MinutesUsed >= limit {
|
||||
return 0, apierr.ErrQuotaExhausted
|
||||
}
|
||||
return limit - day.MinutesUsed, nil
|
||||
}
|
||||
@@ -0,0 +1,176 @@
|
||||
package usage
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"time"
|
||||
|
||||
"github.com/redis/go-redis/v9"
|
||||
"github.com/wangjia/pangolin/server/internal/apierr"
|
||||
)
|
||||
|
||||
// nowFunc is overridable in tests to pin "today".
|
||||
var nowFunc = time.Now
|
||||
|
||||
// utcToday returns the current UTC calendar date (time truncated).
|
||||
func utcToday() time.Time {
|
||||
n := nowFunc().UTC()
|
||||
return time.Date(n.Year(), n.Month(), n.Day(), 0, 0, 0, 0, time.UTC)
|
||||
}
|
||||
|
||||
// Service serves usage queries, the /v1/me usage summary, the ads-unlock flow,
|
||||
// and the free-plan connect quota check.
|
||||
type Service struct {
|
||||
store *Store
|
||||
rdb *redis.Client
|
||||
verifier AdVerifier
|
||||
adNonceTTL time.Duration
|
||||
}
|
||||
|
||||
// NewService builds a Service. rdb may be nil (ad-token replay protection is
|
||||
// then disabled — not recommended in production). verifier may be nil if the
|
||||
// ads-unlock endpoint is not mounted. adNonceTTL <= 0 defaults to 1h.
|
||||
func NewService(store *Store, rdb *redis.Client, verifier AdVerifier, adNonceTTL time.Duration) *Service {
|
||||
if adNonceTTL <= 0 {
|
||||
adNonceTTL = time.Hour
|
||||
}
|
||||
return &Service{store: store, rdb: rdb, verifier: verifier, adNonceTTL: adNonceTTL}
|
||||
}
|
||||
|
||||
// UsagePoint is one day of the usage curve.
|
||||
type UsagePoint struct {
|
||||
Date string `json:"date"` // YYYY-MM-DD (UTC)
|
||||
BytesUp uint64 `json:"bytes_up"`
|
||||
BytesDown uint64 `json:"bytes_down"`
|
||||
MinutesUsed int `json:"minutes_used"`
|
||||
}
|
||||
|
||||
// UsageCurve returns the last `days` daily points for userID (UTC), with
|
||||
// missing days zero-filled, oldest first. days is clamped to [1, 90].
|
||||
func (svc *Service) UsageCurve(ctx context.Context, userID int64, days int) ([]UsagePoint, *apierr.Error) {
|
||||
if days < 1 {
|
||||
days = 7
|
||||
}
|
||||
if days > 90 {
|
||||
days = 90
|
||||
}
|
||||
|
||||
today := utcToday()
|
||||
from := today.AddDate(0, 0, -(days - 1))
|
||||
|
||||
rows, err := svc.store.GetUsageRange(ctx, userID, from, today)
|
||||
if err != nil {
|
||||
return nil, apierr.ErrInternal
|
||||
}
|
||||
byDate := make(map[string]DailyUsage, len(rows))
|
||||
for _, r := range rows {
|
||||
byDate[r.Date.UTC().Format(dateLayout)] = r
|
||||
}
|
||||
|
||||
points := make([]UsagePoint, 0, days)
|
||||
for i := 0; i < days; i++ {
|
||||
d := from.AddDate(0, 0, i).Format(dateLayout)
|
||||
if u, ok := byDate[d]; ok {
|
||||
points = append(points, UsagePoint{
|
||||
Date: d,
|
||||
BytesUp: u.BytesUp,
|
||||
BytesDown: u.BytesDown,
|
||||
MinutesUsed: u.MinutesUsed,
|
||||
})
|
||||
} else {
|
||||
points = append(points, UsagePoint{Date: d})
|
||||
}
|
||||
}
|
||||
return points, nil
|
||||
}
|
||||
|
||||
// TodaySummary is the /v1/me today_usage block.
|
||||
type TodaySummary struct {
|
||||
MinutesUsed int `json:"minutes_used"`
|
||||
MinutesRemaining *int `json:"minutes_remaining"` // nil = unlimited (pro/team)
|
||||
AdUnlocked bool `json:"ad_unlocked"`
|
||||
}
|
||||
|
||||
// TodaySummary returns the current user's usage summary for today (UTC),
|
||||
// reflecting plan limits and ad-unlock state.
|
||||
func (svc *Service) TodaySummary(ctx context.Context, userID int64) (*TodaySummary, *apierr.Error) {
|
||||
plan, err := svc.store.EffectivePlan(ctx, userID)
|
||||
if err != nil {
|
||||
return nil, apierr.ErrInternal
|
||||
}
|
||||
day, err := svc.store.GetDay(ctx, userID, utcToday())
|
||||
if err != nil {
|
||||
return nil, apierr.ErrInternal
|
||||
}
|
||||
|
||||
used := 0
|
||||
adUnlocked := false
|
||||
if day != nil {
|
||||
used = day.MinutesUsed
|
||||
adUnlocked = day.AdUnlockedAt.Valid
|
||||
}
|
||||
|
||||
out := &TodaySummary{MinutesUsed: used, AdUnlocked: adUnlocked}
|
||||
if plan.DailyMinutes.Valid {
|
||||
remaining := int(plan.DailyMinutes.Int64) - used
|
||||
if remaining < 0 {
|
||||
remaining = 0
|
||||
}
|
||||
out.MinutesRemaining = &remaining
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// UnlockAd verifies a rewarded-ad receipt and records the day's ad-unlock.
|
||||
//
|
||||
// Flow: validate inputs → verify the receipt with the ad network → consume the
|
||||
// ad_token as a one-time nonce (replay-protected) → set ad_unlocked_at. A
|
||||
// second unlock on a day already unlocked is idempotent (alreadyUnlocked=true).
|
||||
func (svc *Service) UnlockAd(ctx context.Context, userID int64, deviceID, adToken string) (alreadyUnlocked bool, apiErr *apierr.Error) {
|
||||
if deviceID == "" || adToken == "" {
|
||||
return false, apierr.ErrBadRequest
|
||||
}
|
||||
if svc.verifier == nil {
|
||||
return false, apierr.ErrInternal
|
||||
}
|
||||
|
||||
// 1. Authenticate the receipt with the ad network.
|
||||
if err := svc.verifier.Verify(ctx, AdVerifyRequest{
|
||||
UserID: userID,
|
||||
DeviceID: deviceID,
|
||||
AdToken: adToken,
|
||||
}); err != nil {
|
||||
return false, apierr.ErrAdVerifyFailed
|
||||
}
|
||||
|
||||
// 2. One-time nonce: a genuine receipt may only be redeemed once.
|
||||
if dup, err := svc.consumeAdNonce(ctx, adToken); err != nil {
|
||||
return false, apierr.ErrInternal
|
||||
} else if dup {
|
||||
return false, apierr.ErrAdReplay
|
||||
}
|
||||
|
||||
// 3. Record the unlock (idempotent per UTC day).
|
||||
already, err := svc.store.MarkAdUnlocked(ctx, userID, utcToday())
|
||||
if err != nil {
|
||||
return false, apierr.ErrInternal
|
||||
}
|
||||
return already, nil
|
||||
}
|
||||
|
||||
// consumeAdNonce atomically records the ad_token so it cannot be reused.
|
||||
// Returns dup=true if the token was already consumed. No-ops (dup=false) when
|
||||
// Redis is not configured.
|
||||
func (svc *Service) consumeAdNonce(ctx context.Context, adToken string) (bool, error) {
|
||||
if svc.rdb == nil {
|
||||
return false, nil
|
||||
}
|
||||
sum := sha256.Sum256([]byte(adToken))
|
||||
key := "ads:nonce:" + hex.EncodeToString(sum[:])
|
||||
set, err := svc.rdb.SetNX(ctx, key, "1", svc.adNonceTTL).Result()
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
return !set, nil
|
||||
}
|
||||
@@ -0,0 +1,208 @@
|
||||
package usage
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"fmt"
|
||||
"time"
|
||||
)
|
||||
|
||||
// Unlimited is the sentinel returned by CheckFreeConnect for plans with no
|
||||
// daily-minute cap (pro / team). Callers (e.g. the connect endpoint) treat it
|
||||
// as "no minute limit; use the default paid credential TTL".
|
||||
const Unlimited = -1
|
||||
|
||||
// dateLayout is the canonical UTC day format used as the usage_daily.date key
|
||||
// and in the API response.
|
||||
const dateLayout = "2006-01-02"
|
||||
|
||||
// Plan is the minimal subset of the plans row needed for quota decisions.
|
||||
type Plan struct {
|
||||
Code string // free | pro | team
|
||||
DailyMinutes sql.NullInt64 // free=10; NULL = unlimited
|
||||
AdGate bool // free=true: daily rewarded-ad unlock required
|
||||
}
|
||||
|
||||
// DailyUsage mirrors a usage_daily row. It carries only aggregate byte and
|
||||
// minute counts plus the ad-unlock timestamp — never destinations or DNS, per
|
||||
// the no-log policy.
|
||||
type DailyUsage struct {
|
||||
Date time.Time
|
||||
BytesUp uint64
|
||||
BytesDown uint64
|
||||
MinutesUsed int
|
||||
AdUnlockedAt sql.NullTime
|
||||
}
|
||||
|
||||
// Store wraps a *sql.DB and exposes the usage_daily / plans / users queries the
|
||||
// usage package needs.
|
||||
type Store struct {
|
||||
db *sql.DB
|
||||
}
|
||||
|
||||
// NewStore creates a Store backed by the given MySQL connection pool.
|
||||
func NewStore(db *sql.DB) *Store { return &Store{db: db} }
|
||||
|
||||
// AggregateUsage accumulates one usage delta into usage_daily for (userID, day)
|
||||
// using INSERT … ON DUPLICATE KEY UPDATE so concurrent reports from multiple
|
||||
// nodes add up correctly. day is truncated to its UTC calendar date.
|
||||
func (s *Store) AggregateUsage(ctx context.Context, userID int64, day time.Time, bytesUp, bytesDown uint64, minutes int) error {
|
||||
d := day.UTC().Format(dateLayout)
|
||||
_, err := s.db.ExecContext(ctx,
|
||||
`INSERT INTO usage_daily (user_id, date, bytes_up, bytes_down, minutes_used)
|
||||
VALUES (?, ?, ?, ?, ?)
|
||||
ON DUPLICATE KEY UPDATE
|
||||
bytes_up = bytes_up + VALUES(bytes_up),
|
||||
bytes_down = bytes_down + VALUES(bytes_down),
|
||||
minutes_used = minutes_used + VALUES(minutes_used)`,
|
||||
userID, d, bytesUp, bytesDown, minutes)
|
||||
if err != nil {
|
||||
return fmt.Errorf("store.AggregateUsage: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// LookupUserIDByDPUUID resolves a data-plane credential UUID to the owning
|
||||
// user_id. Returns (0, nil) when no user matches (e.g. a rotated/stale dp_uuid)
|
||||
// so the caller can decide to silently drop the report.
|
||||
func (s *Store) LookupUserIDByDPUUID(ctx context.Context, dpUUID string) (int64, error) {
|
||||
var id int64
|
||||
err := s.db.QueryRowContext(ctx,
|
||||
`SELECT id FROM users WHERE dp_uuid = ? LIMIT 1`, dpUUID).Scan(&id)
|
||||
if err == sql.ErrNoRows {
|
||||
return 0, nil
|
||||
}
|
||||
if err != nil {
|
||||
return 0, fmt.Errorf("store.LookupUserIDByDPUUID: %w", err)
|
||||
}
|
||||
return id, nil
|
||||
}
|
||||
|
||||
// GetUsageRange returns the usage_daily rows for userID with date in
|
||||
// [from, to] (inclusive, UTC dates), ordered ascending. Missing days are NOT
|
||||
// filled here; zero-filling is the handler's responsibility.
|
||||
func (s *Store) GetUsageRange(ctx context.Context, userID int64, from, to time.Time) ([]DailyUsage, error) {
|
||||
rows, err := s.db.QueryContext(ctx,
|
||||
`SELECT date, bytes_up, bytes_down, minutes_used, ad_unlocked_at
|
||||
FROM usage_daily
|
||||
WHERE user_id = ? AND date BETWEEN ? AND ?
|
||||
ORDER BY date ASC`,
|
||||
userID, from.UTC().Format(dateLayout), to.UTC().Format(dateLayout))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("store.GetUsageRange: %w", err)
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
var out []DailyUsage
|
||||
for rows.Next() {
|
||||
var u DailyUsage
|
||||
if err := rows.Scan(&u.Date, &u.BytesUp, &u.BytesDown, &u.MinutesUsed, &u.AdUnlockedAt); err != nil {
|
||||
return nil, fmt.Errorf("store.GetUsageRange scan: %w", err)
|
||||
}
|
||||
out = append(out, u)
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
|
||||
// GetDay returns the usage_daily row for (userID, day), or nil if none exists.
|
||||
func (s *Store) GetDay(ctx context.Context, userID int64, day time.Time) (*DailyUsage, error) {
|
||||
var u DailyUsage
|
||||
err := s.db.QueryRowContext(ctx,
|
||||
`SELECT date, bytes_up, bytes_down, minutes_used, ad_unlocked_at
|
||||
FROM usage_daily
|
||||
WHERE user_id = ? AND date = ?`,
|
||||
userID, day.UTC().Format(dateLayout)).
|
||||
Scan(&u.Date, &u.BytesUp, &u.BytesDown, &u.MinutesUsed, &u.AdUnlockedAt)
|
||||
if err == sql.ErrNoRows {
|
||||
return nil, nil
|
||||
}
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("store.GetDay: %w", err)
|
||||
}
|
||||
return &u, nil
|
||||
}
|
||||
|
||||
// MarkAdUnlocked sets usage_daily.ad_unlocked_at for (userID, day) to now if it
|
||||
// is not already set. It returns alreadyUnlocked=true when the day was already
|
||||
// unlocked (making a second call idempotent). Runs inside a transaction with
|
||||
// SELECT … FOR UPDATE to be safe under concurrent unlock attempts.
|
||||
func (s *Store) MarkAdUnlocked(ctx context.Context, userID int64, day time.Time) (alreadyUnlocked bool, err error) {
|
||||
d := day.UTC().Format(dateLayout)
|
||||
|
||||
tx, err := s.db.BeginTx(ctx, &sql.TxOptions{Isolation: sql.LevelReadCommitted})
|
||||
if err != nil {
|
||||
return false, fmt.Errorf("store.MarkAdUnlocked begin: %w", err)
|
||||
}
|
||||
committed := false
|
||||
defer func() {
|
||||
if !committed {
|
||||
_ = tx.Rollback()
|
||||
}
|
||||
}()
|
||||
|
||||
var unlockedAt sql.NullTime
|
||||
row := tx.QueryRowContext(ctx,
|
||||
`SELECT ad_unlocked_at FROM usage_daily WHERE user_id = ? AND date = ? FOR UPDATE`,
|
||||
userID, d)
|
||||
switch err := row.Scan(&unlockedAt); err {
|
||||
case nil:
|
||||
if unlockedAt.Valid {
|
||||
// Already unlocked today → idempotent success.
|
||||
if cErr := tx.Commit(); cErr != nil {
|
||||
return false, fmt.Errorf("store.MarkAdUnlocked commit: %w", cErr)
|
||||
}
|
||||
committed = true
|
||||
return true, nil
|
||||
}
|
||||
if _, uErr := tx.ExecContext(ctx,
|
||||
`UPDATE usage_daily SET ad_unlocked_at = UTC_TIMESTAMP(6)
|
||||
WHERE user_id = ? AND date = ? AND ad_unlocked_at IS NULL`,
|
||||
userID, d); uErr != nil {
|
||||
return false, fmt.Errorf("store.MarkAdUnlocked update: %w", uErr)
|
||||
}
|
||||
case sql.ErrNoRows:
|
||||
if _, iErr := tx.ExecContext(ctx,
|
||||
`INSERT INTO usage_daily (user_id, date, ad_unlocked_at)
|
||||
VALUES (?, ?, UTC_TIMESTAMP(6))`,
|
||||
userID, d); iErr != nil {
|
||||
return false, fmt.Errorf("store.MarkAdUnlocked insert: %w", iErr)
|
||||
}
|
||||
default:
|
||||
return false, fmt.Errorf("store.MarkAdUnlocked select: %w", err)
|
||||
}
|
||||
|
||||
if cErr := tx.Commit(); cErr != nil {
|
||||
return false, fmt.Errorf("store.MarkAdUnlocked commit: %w", cErr)
|
||||
}
|
||||
committed = true
|
||||
return false, nil
|
||||
}
|
||||
|
||||
// EffectivePlan returns the plan that currently governs userID: the highest
|
||||
// tier among the user's non-expired subscriptions, falling back to the free
|
||||
// plan when the user has no active subscription (trial-expired free state).
|
||||
func (s *Store) EffectivePlan(ctx context.Context, userID int64) (*Plan, error) {
|
||||
var p Plan
|
||||
err := s.db.QueryRowContext(ctx,
|
||||
`SELECT p.code, p.daily_minutes, p.ad_gate
|
||||
FROM subscriptions s
|
||||
JOIN plans p ON p.id = s.plan_id
|
||||
WHERE s.user_id = ? AND s.expires_at > UTC_TIMESTAMP(6)
|
||||
ORDER BY FIELD(p.code, 'team', 'pro', 'free'), s.expires_at DESC
|
||||
LIMIT 1`,
|
||||
userID).Scan(&p.Code, &p.DailyMinutes, &p.AdGate)
|
||||
if err == nil {
|
||||
return &p, nil
|
||||
}
|
||||
if err != sql.ErrNoRows {
|
||||
return nil, fmt.Errorf("store.EffectivePlan: %w", err)
|
||||
}
|
||||
|
||||
// No active subscription → free plan.
|
||||
if err := s.db.QueryRowContext(ctx,
|
||||
`SELECT code, daily_minutes, ad_gate FROM plans WHERE code = 'free'`).
|
||||
Scan(&p.Code, &p.DailyMinutes, &p.AdGate); err != nil {
|
||||
return nil, fmt.Errorf("store.EffectivePlan free fallback: %w", err)
|
||||
}
|
||||
return &p, nil
|
||||
}
|
||||
@@ -0,0 +1,516 @@
|
||||
//go:build integration
|
||||
|
||||
package usage_test
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"fmt"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
_ "github.com/go-sql-driver/mysql"
|
||||
"github.com/redis/go-redis/v9"
|
||||
"github.com/testcontainers/testcontainers-go"
|
||||
tcmysql "github.com/testcontainers/testcontainers-go/modules/mysql"
|
||||
tcredis "github.com/testcontainers/testcontainers-go/modules/redis"
|
||||
|
||||
"github.com/wangjia/pangolin/server/internal/usage"
|
||||
)
|
||||
|
||||
// --------------------------------------------------------------------------
|
||||
// Container setup
|
||||
// --------------------------------------------------------------------------
|
||||
|
||||
func setupMySQL(t *testing.T) *sql.DB {
|
||||
t.Helper()
|
||||
ctx := context.Background()
|
||||
|
||||
ctr, err := tcmysql.Run(ctx, "mysql:8.0",
|
||||
tcmysql.WithDatabase("pangolin_test"),
|
||||
tcmysql.WithUsername("root"),
|
||||
tcmysql.WithPassword("test"),
|
||||
)
|
||||
testcontainers.CleanupContainer(t, ctr)
|
||||
if err != nil {
|
||||
t.Fatalf("mysql container: %v", err)
|
||||
}
|
||||
dsn, err := ctr.ConnectionString(ctx, "parseTime=true", "loc=UTC", "time_zone='+00:00'")
|
||||
if err != nil {
|
||||
t.Fatalf("mysql dsn: %v", err)
|
||||
}
|
||||
db, err := sql.Open("mysql", dsn)
|
||||
if err != nil {
|
||||
t.Fatalf("open mysql: %v", err)
|
||||
}
|
||||
t.Cleanup(func() { db.Close() })
|
||||
if err := applySchema(db); err != nil {
|
||||
t.Fatalf("schema: %v", err)
|
||||
}
|
||||
return db
|
||||
}
|
||||
|
||||
func setupRedis(t *testing.T) *redis.Client {
|
||||
t.Helper()
|
||||
ctx := context.Background()
|
||||
ctr, err := tcredis.Run(ctx, "redis:7-alpine")
|
||||
testcontainers.CleanupContainer(t, ctr)
|
||||
if err != nil {
|
||||
t.Fatalf("redis container: %v", err)
|
||||
}
|
||||
addr, err := ctr.ConnectionString(ctx)
|
||||
if err != nil {
|
||||
t.Fatalf("redis addr: %v", err)
|
||||
}
|
||||
for _, p := range []string{"redis://", "rediss://"} {
|
||||
if len(addr) > len(p) && addr[:len(p)] == p {
|
||||
addr = addr[len(p):]
|
||||
}
|
||||
}
|
||||
rdb := redis.NewClient(&redis.Options{Addr: addr})
|
||||
t.Cleanup(func() { rdb.Close() })
|
||||
return rdb
|
||||
}
|
||||
|
||||
func applySchema(db *sql.DB) error {
|
||||
stmts := []string{
|
||||
`CREATE TABLE IF NOT EXISTS users (
|
||||
id BIGINT UNSIGNED AUTO_INCREMENT PRIMARY KEY,
|
||||
uuid CHAR(36) NOT NULL UNIQUE,
|
||||
email VARCHAR(255) NOT NULL UNIQUE,
|
||||
pw_hash VARCHAR(255) NOT NULL DEFAULT '',
|
||||
dp_uuid CHAR(36) NOT NULL,
|
||||
status ENUM('active','banned') NOT NULL DEFAULT 'active',
|
||||
created_at DATETIME(6) NOT NULL DEFAULT CURRENT_TIMESTAMP(6),
|
||||
INDEX idx_dp (dp_uuid)
|
||||
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4`,
|
||||
`CREATE TABLE IF NOT EXISTS plans (
|
||||
id BIGINT UNSIGNED AUTO_INCREMENT PRIMARY KEY,
|
||||
code ENUM('free','pro','team') NOT NULL UNIQUE,
|
||||
max_devices INT NOT NULL DEFAULT 1,
|
||||
daily_minutes INT NULL,
|
||||
ad_gate BOOLEAN NOT NULL DEFAULT FALSE
|
||||
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4`,
|
||||
`CREATE TABLE IF NOT EXISTS subscriptions (
|
||||
id BIGINT UNSIGNED AUTO_INCREMENT PRIMARY KEY,
|
||||
user_id BIGINT UNSIGNED NOT NULL,
|
||||
plan_id BIGINT UNSIGNED NOT NULL,
|
||||
expires_at DATETIME(6) NOT NULL,
|
||||
source ENUM('trial','code') NOT NULL,
|
||||
created_at DATETIME(6) NOT NULL DEFAULT CURRENT_TIMESTAMP(6),
|
||||
INDEX idx_user_exp (user_id, expires_at)
|
||||
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4`,
|
||||
`CREATE TABLE IF NOT EXISTS usage_daily (
|
||||
user_id BIGINT UNSIGNED NOT NULL,
|
||||
date DATE NOT NULL,
|
||||
bytes_up BIGINT UNSIGNED NOT NULL DEFAULT 0,
|
||||
bytes_down BIGINT UNSIGNED NOT NULL DEFAULT 0,
|
||||
minutes_used INT NOT NULL DEFAULT 0,
|
||||
ad_unlocked_at DATETIME(6) NULL,
|
||||
PRIMARY KEY (user_id, date)
|
||||
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4`,
|
||||
`INSERT IGNORE INTO plans (code, max_devices, daily_minutes, ad_gate)
|
||||
VALUES ('free',1,10,TRUE), ('pro',5,NULL,FALSE), ('team',10,NULL,FALSE)`,
|
||||
}
|
||||
for _, s := range stmts {
|
||||
if _, err := db.Exec(s); err != nil {
|
||||
return fmt.Errorf("exec: %w\nSQL: %s", err, s)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// createUser inserts a user with the given dp_uuid and returns its id.
|
||||
func createUser(t *testing.T, db *sql.DB, dpUUID string) int64 {
|
||||
t.Helper()
|
||||
uuid := fmt.Sprintf("u-%d", time.Now().UnixNano())
|
||||
res, err := db.Exec(
|
||||
`INSERT INTO users (uuid, email, pw_hash, dp_uuid) VALUES (?,?,?,?)`,
|
||||
uuid, uuid+"@example.com", "x", dpUUID)
|
||||
if err != nil {
|
||||
t.Fatalf("createUser: %v", err)
|
||||
}
|
||||
id, _ := res.LastInsertId()
|
||||
return id
|
||||
}
|
||||
|
||||
// giveSubscription grants userID an active subscription on the given plan.
|
||||
func giveSubscription(t *testing.T, db *sql.DB, userID int64, plan string) {
|
||||
t.Helper()
|
||||
var planID int64
|
||||
if err := db.QueryRow(`SELECT id FROM plans WHERE code=?`, plan).Scan(&planID); err != nil {
|
||||
t.Fatalf("plan lookup: %v", err)
|
||||
}
|
||||
if _, err := db.Exec(
|
||||
`INSERT INTO subscriptions (user_id, plan_id, expires_at, source)
|
||||
VALUES (?,?,?,'code')`,
|
||||
userID, planID, time.Now().UTC().AddDate(0, 0, 30)); err != nil {
|
||||
t.Fatalf("giveSubscription: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// fakeVerifier is an AdVerifier test double.
|
||||
type fakeVerifier struct{ ok bool }
|
||||
|
||||
func (f fakeVerifier) Provider() string { return "fake" }
|
||||
func (f fakeVerifier) Verify(context.Context, usage.AdVerifyRequest) error {
|
||||
if f.ok {
|
||||
return nil
|
||||
}
|
||||
return fmt.Errorf("fake verifier: rejected")
|
||||
}
|
||||
|
||||
// --------------------------------------------------------------------------
|
||||
// Aggregation
|
||||
// --------------------------------------------------------------------------
|
||||
|
||||
func TestIntAggregateConcurrentMultiNode(t *testing.T) {
|
||||
db := setupMySQL(t)
|
||||
rdb := setupRedis(t)
|
||||
store := usage.NewStore(db)
|
||||
agg := usage.NewAggregator(store, rdb, time.Minute, time.Minute)
|
||||
|
||||
uid := createUser(t, db, "dp-agg-1")
|
||||
|
||||
const goroutines = 25
|
||||
const each = 4
|
||||
var wg sync.WaitGroup
|
||||
wg.Add(goroutines)
|
||||
for g := 0; g < goroutines; g++ {
|
||||
go func(g int) {
|
||||
defer wg.Done()
|
||||
for i := 0; i < each; i++ {
|
||||
err := agg.ReportUsage(context.Background(), usage.Report{
|
||||
DPUUID: "dp-agg-1",
|
||||
BytesUp: 100,
|
||||
BytesDown: 200,
|
||||
Minutes: 1,
|
||||
BatchID: fmt.Sprintf("b-%d-%d", g, i), // unique → all counted
|
||||
})
|
||||
if err != nil {
|
||||
t.Errorf("report: %v", err)
|
||||
}
|
||||
}
|
||||
}(g)
|
||||
}
|
||||
wg.Wait()
|
||||
|
||||
var up, down uint64
|
||||
var minutes int
|
||||
db.QueryRow(`SELECT bytes_up, bytes_down, minutes_used FROM usage_daily WHERE user_id=?`, uid).
|
||||
Scan(&up, &down, &minutes)
|
||||
|
||||
wantUp := uint64(goroutines * each * 100)
|
||||
wantMin := goroutines * each
|
||||
if up != wantUp || down != wantUp*2 || minutes != wantMin {
|
||||
t.Errorf("got up=%d down=%d min=%d, want up=%d down=%d min=%d",
|
||||
up, down, minutes, wantUp, wantUp*2, wantMin)
|
||||
}
|
||||
}
|
||||
|
||||
func TestIntAggregateReplayDedup(t *testing.T) {
|
||||
db := setupMySQL(t)
|
||||
rdb := setupRedis(t)
|
||||
store := usage.NewStore(db)
|
||||
agg := usage.NewAggregator(store, rdb, time.Minute, time.Minute)
|
||||
|
||||
uid := createUser(t, db, "dp-replay")
|
||||
|
||||
r := usage.Report{DPUUID: "dp-replay", BytesUp: 500, Minutes: 5, BatchID: "same-batch"}
|
||||
for i := 0; i < 4; i++ {
|
||||
if err := agg.ReportUsage(context.Background(), r); err != nil {
|
||||
t.Fatalf("report %d: %v", i, err)
|
||||
}
|
||||
}
|
||||
|
||||
var up uint64
|
||||
var minutes int
|
||||
db.QueryRow(`SELECT bytes_up, minutes_used FROM usage_daily WHERE user_id=?`, uid).Scan(&up, &minutes)
|
||||
if up != 500 || minutes != 5 {
|
||||
t.Errorf("replay not deduped: up=%d minutes=%d, want 500/5", up, minutes)
|
||||
}
|
||||
}
|
||||
|
||||
func TestIntAggregateCrossUTCDayBucketing(t *testing.T) {
|
||||
db := setupMySQL(t)
|
||||
store := usage.NewStore(db)
|
||||
agg := usage.NewAggregator(store, nil, time.Minute, time.Minute)
|
||||
|
||||
uid := createUser(t, db, "dp-days")
|
||||
|
||||
day1 := time.Date(2026, 6, 12, 23, 30, 0, 0, time.UTC)
|
||||
day2 := time.Date(2026, 6, 13, 0, 15, 0, 0, time.UTC)
|
||||
|
||||
must := func(at time.Time, minutes int) {
|
||||
if err := agg.ReportUsage(context.Background(), usage.Report{
|
||||
DPUUID: "dp-days", BytesUp: 10, Minutes: minutes, At: at,
|
||||
}); err != nil {
|
||||
t.Fatalf("report: %v", err)
|
||||
}
|
||||
}
|
||||
must(day1, 3)
|
||||
must(day1.Add(10*time.Minute), 2) // still 06-12
|
||||
must(day2, 7) // 06-13
|
||||
|
||||
check := func(date string, wantMin int) {
|
||||
var m int
|
||||
err := db.QueryRow(`SELECT minutes_used FROM usage_daily WHERE user_id=? AND date=?`, uid, date).Scan(&m)
|
||||
if err != nil {
|
||||
t.Fatalf("query %s: %v", date, err)
|
||||
}
|
||||
if m != wantMin {
|
||||
t.Errorf("date %s minutes=%d, want %d", date, m, wantMin)
|
||||
}
|
||||
}
|
||||
check("2026-06-12", 5)
|
||||
check("2026-06-13", 7)
|
||||
}
|
||||
|
||||
func TestIntAggregateUnknownDPUUIDDropped(t *testing.T) {
|
||||
db := setupMySQL(t)
|
||||
store := usage.NewStore(db)
|
||||
agg := usage.NewAggregator(store, nil, time.Minute, time.Minute)
|
||||
|
||||
err := agg.ReportUsage(context.Background(), usage.Report{
|
||||
DPUUID: "does-not-exist", BytesUp: 1, Minutes: 1,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("unknown dp_uuid should be a silent no-op, got %v", err)
|
||||
}
|
||||
var n int
|
||||
db.QueryRow(`SELECT COUNT(*) FROM usage_daily`).Scan(&n)
|
||||
if n != 0 {
|
||||
t.Errorf("expected no rows for unknown dp_uuid, got %d", n)
|
||||
}
|
||||
}
|
||||
|
||||
// --------------------------------------------------------------------------
|
||||
// Quota (CheckFreeConnect)
|
||||
// --------------------------------------------------------------------------
|
||||
|
||||
func TestIntQuotaFreeNotUnlocked(t *testing.T) {
|
||||
db := setupMySQL(t)
|
||||
svc := usage.NewService(usage.NewStore(db), nil, fakeVerifier{ok: true}, time.Hour)
|
||||
|
||||
uid := createUser(t, db, "dp-q1") // no subscription → free plan
|
||||
|
||||
_, apiErr := svc.CheckFreeConnect(context.Background(), uid)
|
||||
if apiErr == nil || apiErr.Code != "AD_NOT_UNLOCKED" {
|
||||
t.Fatalf("expected AD_NOT_UNLOCKED, got %v", apiErr)
|
||||
}
|
||||
if apiErr.MessageZH == "" || apiErr.MessageEn == "" {
|
||||
t.Error("expected bilingual error messages")
|
||||
}
|
||||
}
|
||||
|
||||
func TestIntQuotaFreeUnlockedRemainingDecrements(t *testing.T) {
|
||||
db := setupMySQL(t)
|
||||
store := usage.NewStore(db)
|
||||
agg := usage.NewAggregator(store, nil, time.Minute, time.Minute)
|
||||
svc := usage.NewService(store, nil, fakeVerifier{ok: true}, time.Hour)
|
||||
|
||||
uid := createUser(t, db, "dp-q2")
|
||||
|
||||
if _, err := store.MarkAdUnlocked(context.Background(), uid, time.Now().UTC()); err != nil {
|
||||
t.Fatalf("unlock: %v", err)
|
||||
}
|
||||
|
||||
rem, apiErr := svc.CheckFreeConnect(context.Background(), uid)
|
||||
if apiErr != nil {
|
||||
t.Fatalf("unexpected err: %v", apiErr)
|
||||
}
|
||||
if rem != 10 {
|
||||
t.Errorf("remaining=%d, want 10", rem)
|
||||
}
|
||||
|
||||
// Consume 4 minutes.
|
||||
agg.ReportUsage(context.Background(), usage.Report{DPUUID: "dp-q2", Minutes: 4})
|
||||
rem, _ = svc.CheckFreeConnect(context.Background(), uid)
|
||||
if rem != 6 {
|
||||
t.Errorf("after 4 min: remaining=%d, want 6", rem)
|
||||
}
|
||||
}
|
||||
|
||||
func TestIntQuotaFreeExhausted(t *testing.T) {
|
||||
db := setupMySQL(t)
|
||||
store := usage.NewStore(db)
|
||||
svc := usage.NewService(store, nil, fakeVerifier{ok: true}, time.Hour)
|
||||
|
||||
uid := createUser(t, db, "dp-q3")
|
||||
store.MarkAdUnlocked(context.Background(), uid, time.Now().UTC())
|
||||
store.AggregateUsage(context.Background(), uid, time.Now().UTC(), 0, 0, 10)
|
||||
|
||||
_, apiErr := svc.CheckFreeConnect(context.Background(), uid)
|
||||
if apiErr == nil || apiErr.Code != "QUOTA_EXHAUSTED" {
|
||||
t.Fatalf("expected QUOTA_EXHAUSTED, got %v", apiErr)
|
||||
}
|
||||
}
|
||||
|
||||
func TestIntQuotaPaidUnlimited(t *testing.T) {
|
||||
db := setupMySQL(t)
|
||||
store := usage.NewStore(db)
|
||||
svc := usage.NewService(store, nil, fakeVerifier{ok: true}, time.Hour)
|
||||
|
||||
for _, plan := range []string{"pro", "team"} {
|
||||
uid := createUser(t, db, "dp-"+plan)
|
||||
giveSubscription(t, db, uid, plan)
|
||||
// Even with lots of minutes used, paid plans are not gated.
|
||||
store.AggregateUsage(context.Background(), uid, time.Now().UTC(), 0, 0, 9999)
|
||||
|
||||
rem, apiErr := svc.CheckFreeConnect(context.Background(), uid)
|
||||
if apiErr != nil {
|
||||
t.Fatalf("%s: unexpected err %v", plan, apiErr)
|
||||
}
|
||||
if rem != usage.Unlimited {
|
||||
t.Errorf("%s: remaining=%d, want Unlimited(%d)", plan, rem, usage.Unlimited)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// --------------------------------------------------------------------------
|
||||
// Ads unlock
|
||||
// --------------------------------------------------------------------------
|
||||
|
||||
func TestIntAdsUnlockForgedRejected(t *testing.T) {
|
||||
db := setupMySQL(t)
|
||||
rdb := setupRedis(t)
|
||||
svc := usage.NewService(usage.NewStore(db), rdb, fakeVerifier{ok: false}, time.Hour)
|
||||
|
||||
uid := createUser(t, db, "dp-ad1")
|
||||
_, apiErr := svc.UnlockAd(context.Background(), uid, "device-1", "forged-token")
|
||||
if apiErr == nil || apiErr.Code != "AD_VERIFY_FAILED" {
|
||||
t.Fatalf("expected AD_VERIFY_FAILED, got %v", apiErr)
|
||||
}
|
||||
}
|
||||
|
||||
func TestIntAdsUnlockReplayRejected(t *testing.T) {
|
||||
db := setupMySQL(t)
|
||||
rdb := setupRedis(t)
|
||||
svc := usage.NewService(usage.NewStore(db), rdb, fakeVerifier{ok: true}, time.Hour)
|
||||
|
||||
uid := createUser(t, db, "dp-ad2")
|
||||
already, apiErr := svc.UnlockAd(context.Background(), uid, "device-1", "token-xyz")
|
||||
if apiErr != nil || already {
|
||||
t.Fatalf("first unlock: already=%v err=%v", already, apiErr)
|
||||
}
|
||||
// Same token again → replay.
|
||||
_, apiErr = svc.UnlockAd(context.Background(), uid, "device-1", "token-xyz")
|
||||
if apiErr == nil || apiErr.Code != "AD_TOKEN_REPLAY" {
|
||||
t.Fatalf("expected AD_TOKEN_REPLAY, got %v", apiErr)
|
||||
}
|
||||
}
|
||||
|
||||
func TestIntAdsUnlockSecondTimeIdempotent(t *testing.T) {
|
||||
db := setupMySQL(t)
|
||||
rdb := setupRedis(t)
|
||||
svc := usage.NewService(usage.NewStore(db), rdb, fakeVerifier{ok: true}, time.Hour)
|
||||
|
||||
uid := createUser(t, db, "dp-ad3")
|
||||
if _, apiErr := svc.UnlockAd(context.Background(), uid, "d", "tok-1"); apiErr != nil {
|
||||
t.Fatalf("first unlock: %v", apiErr)
|
||||
}
|
||||
// Different (valid) token, same day → idempotent success.
|
||||
already, apiErr := svc.UnlockAd(context.Background(), uid, "d", "tok-2")
|
||||
if apiErr != nil {
|
||||
t.Fatalf("second unlock err: %v", apiErr)
|
||||
}
|
||||
if !already {
|
||||
t.Error("expected already-unlocked idempotent success")
|
||||
}
|
||||
|
||||
// Exactly one unlock timestamp.
|
||||
var n int
|
||||
db.QueryRow(`SELECT COUNT(*) FROM usage_daily WHERE user_id=? AND ad_unlocked_at IS NOT NULL`, uid).Scan(&n)
|
||||
if n != 1 {
|
||||
t.Errorf("expected 1 unlocked day, got %d", n)
|
||||
}
|
||||
}
|
||||
|
||||
// --------------------------------------------------------------------------
|
||||
// Usage curve & today summary
|
||||
// --------------------------------------------------------------------------
|
||||
|
||||
func TestIntUsageCurveZeroFill(t *testing.T) {
|
||||
db := setupMySQL(t)
|
||||
store := usage.NewStore(db)
|
||||
svc := usage.NewService(store, nil, nil, time.Hour)
|
||||
|
||||
uid := createUser(t, db, "dp-curve")
|
||||
today := time.Now().UTC()
|
||||
// Only two of the last 7 days have data.
|
||||
store.AggregateUsage(context.Background(), uid, today, 1000, 2000, 5)
|
||||
store.AggregateUsage(context.Background(), uid, today.AddDate(0, 0, -3), 50, 60, 2)
|
||||
|
||||
points, apiErr := svc.UsageCurve(context.Background(), uid, 7)
|
||||
if apiErr != nil {
|
||||
t.Fatalf("curve: %v", apiErr)
|
||||
}
|
||||
if len(points) != 7 {
|
||||
t.Fatalf("expected 7 points, got %d", len(points))
|
||||
}
|
||||
// Oldest first, dates contiguous.
|
||||
for i := 1; i < len(points); i++ {
|
||||
if points[i].Date <= points[i-1].Date {
|
||||
t.Errorf("points not ascending: %s then %s", points[i-1].Date, points[i].Date)
|
||||
}
|
||||
}
|
||||
// Last point = today with the data.
|
||||
last := points[len(points)-1]
|
||||
if last.BytesUp != 1000 || last.BytesDown != 2000 || last.MinutesUsed != 5 {
|
||||
t.Errorf("today point wrong: %+v", last)
|
||||
}
|
||||
// Days with no data are zero-filled.
|
||||
zeros := 0
|
||||
for _, p := range points {
|
||||
if p.BytesUp == 0 && p.BytesDown == 0 && p.MinutesUsed == 0 {
|
||||
zeros++
|
||||
}
|
||||
}
|
||||
if zeros != 5 {
|
||||
t.Errorf("expected 5 zero-filled days, got %d", zeros)
|
||||
}
|
||||
}
|
||||
|
||||
func TestIntUsageCurveOnlyCurrentUser(t *testing.T) {
|
||||
db := setupMySQL(t)
|
||||
store := usage.NewStore(db)
|
||||
svc := usage.NewService(store, nil, nil, time.Hour)
|
||||
|
||||
uA := createUser(t, db, "dp-A")
|
||||
uB := createUser(t, db, "dp-B")
|
||||
store.AggregateUsage(context.Background(), uA, time.Now().UTC(), 111, 0, 1)
|
||||
store.AggregateUsage(context.Background(), uB, time.Now().UTC(), 999, 0, 9)
|
||||
|
||||
points, _ := svc.UsageCurve(context.Background(), uA, 1)
|
||||
if len(points) != 1 || points[0].BytesUp != 111 {
|
||||
t.Errorf("user A curve leaked other user data: %+v", points)
|
||||
}
|
||||
}
|
||||
|
||||
func TestIntTodaySummary(t *testing.T) {
|
||||
db := setupMySQL(t)
|
||||
store := usage.NewStore(db)
|
||||
svc := usage.NewService(store, nil, nil, time.Hour)
|
||||
|
||||
// Free user: limited + ad flag.
|
||||
uFree := createUser(t, db, "dp-sf")
|
||||
store.MarkAdUnlocked(context.Background(), uFree, time.Now().UTC())
|
||||
store.AggregateUsage(context.Background(), uFree, time.Now().UTC(), 0, 0, 3)
|
||||
|
||||
sum, apiErr := svc.TodaySummary(context.Background(), uFree)
|
||||
if apiErr != nil {
|
||||
t.Fatalf("summary: %v", apiErr)
|
||||
}
|
||||
if sum.MinutesUsed != 3 || sum.MinutesRemaining == nil || *sum.MinutesRemaining != 7 || !sum.AdUnlocked {
|
||||
t.Errorf("free summary wrong: %+v (remaining=%v)", sum, sum.MinutesRemaining)
|
||||
}
|
||||
|
||||
// Pro user: unlimited (nil remaining).
|
||||
uPro := createUser(t, db, "dp-sp")
|
||||
giveSubscription(t, db, uPro, "pro")
|
||||
sum2, _ := svc.TodaySummary(context.Background(), uPro)
|
||||
if sum2.MinutesRemaining != nil {
|
||||
t.Errorf("pro remaining should be nil(unlimited), got %v", *sum2.MinutesRemaining)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user