diff --git a/server/api/openapi.yaml b/server/api/openapi.yaml index 9ff7fa7..2e75e91 100644 --- a/server/api/openapi.yaml +++ b/server/api/openapi.yaml @@ -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": diff --git a/server/internal/apierr/apierr.go b/server/internal/apierr/apierr.go index 81b6d53..c8f1dfe 100644 --- a/server/internal/apierr/apierr.go +++ b/server/internal/apierr/apierr.go @@ -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", diff --git a/server/internal/usage/ads.go b/server/internal/usage/ads.go new file mode 100644 index 0000000..023b318 --- /dev/null +++ b/server/internal/usage/ads.go @@ -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 +} diff --git a/server/internal/usage/ads_test.go b/server/internal/usage/ads_test.go new file mode 100644 index 0000000..faf975b --- /dev/null +++ b/server/internal/usage/ads_test.go @@ -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") + } +} diff --git a/server/internal/usage/aggregator.go b/server/internal/usage/aggregator.go new file mode 100644 index 0000000..bc04c2e --- /dev/null +++ b/server/internal/usage/aggregator.go @@ -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() +} diff --git a/server/internal/usage/handler.go b/server/internal/usage/handler.go new file mode 100644 index 0000000..156e487 --- /dev/null +++ b/server/internal/usage/handler.go @@ -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 + } +} diff --git a/server/internal/usage/quota.go b/server/internal/usage/quota.go new file mode 100644 index 0000000..731f305 --- /dev/null +++ b/server/internal/usage/quota.go @@ -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 +} diff --git a/server/internal/usage/service.go b/server/internal/usage/service.go new file mode 100644 index 0000000..f814aaf --- /dev/null +++ b/server/internal/usage/service.go @@ -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 +} diff --git a/server/internal/usage/store.go b/server/internal/usage/store.go new file mode 100644 index 0000000..6483fe0 --- /dev/null +++ b/server/internal/usage/store.go @@ -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 +} diff --git a/server/internal/usage/usage_integration_test.go b/server/internal/usage/usage_integration_test.go new file mode 100644 index 0000000..b8355b7 --- /dev/null +++ b/server/internal/usage/usage_integration_test.go @@ -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) + } +}