7434e5195e
Anonymous callers get a free cumulative quota (1000 calls per IP); once exhausted they get 403 quota_exhausted and must register. Public users can self-register (email+password) to obtain a higher-quota API key, view usage, and regenerate the key. Quota counters live in Redis; the public API stays read-only except for the registration writes. - migration 0011: app_user table + api_key.quota_total + 'registered' tier - ratelimit: IncrTotal/TotalUsed/CopyTotal lifetime counters - middleware: enforce cumulative quota + X-Quota-* headers - store: RegisterUser/Authenticate/RegenerateKey (bcrypt) - handlers: POST /api/v1/register, /account, /account/regenerate - admin: quota_total column + registered tier - public: 'API 密钥' account page + API docs quota section Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
150 lines
4.7 KiB
Go
150 lines
4.7 KiB
Go
package handler
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"os"
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/jackc/pgx/v5/pgxpool"
|
|
"github.com/redis/go-redis/v9"
|
|
|
|
"github.com/baicai2026-baicai/goods/api/internal/ratelimit"
|
|
"github.com/baicai2026-baicai/goods/api/internal/store"
|
|
)
|
|
|
|
func cleanupCounter(t *testing.T, subject string) {
|
|
t.Helper()
|
|
redisURL := os.Getenv("OPENGOODS_REDIS_URL")
|
|
if redisURL == "" {
|
|
redisURL = "redis://localhost:6379/0"
|
|
}
|
|
opt, err := redis.ParseURL(redisURL)
|
|
if err != nil {
|
|
return
|
|
}
|
|
rdb := redis.NewClient(opt)
|
|
defer rdb.Close()
|
|
rdb.Del(context.Background(), "usage:total:"+subject)
|
|
}
|
|
|
|
// newQuotaHandler builds a handler backed by the test DB and a live Redis
|
|
// limiter, with a small anonymous quota so exhaustion is cheap to exercise.
|
|
func newQuotaHandler(t *testing.T, anonQuota int) *Handler {
|
|
t.Helper()
|
|
dsn := os.Getenv("OPENGOODS_DATABASE_URL")
|
|
if dsn == "" {
|
|
dsn = "postgres://opengoods:opengoods@localhost:5432/opengoods?sslmode=disable"
|
|
}
|
|
ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
|
|
defer cancel()
|
|
pool, err := pgxpool.New(ctx, dsn)
|
|
if err != nil {
|
|
t.Skipf("no database: %v", err)
|
|
}
|
|
if err := pool.Ping(ctx); err != nil {
|
|
pool.Close()
|
|
t.Skipf("database not reachable: %v", err)
|
|
}
|
|
var hasUser bool
|
|
if err := pool.QueryRow(ctx, "SELECT to_regclass('public.app_user') IS NOT NULL").Scan(&hasUser); err != nil || !hasUser {
|
|
pool.Close()
|
|
t.Skip("migrations not applied")
|
|
}
|
|
redisURL := os.Getenv("OPENGOODS_REDIS_URL")
|
|
if redisURL == "" {
|
|
redisURL = "redis://localhost:6379/0"
|
|
}
|
|
limiter := ratelimit.New(redisURL)
|
|
pingCtx, pingCancel := context.WithTimeout(context.Background(), time.Second)
|
|
defer pingCancel()
|
|
if err := limiter.Ping(pingCtx); err != nil {
|
|
pool.Close()
|
|
t.Skipf("redis not reachable: %v", err)
|
|
}
|
|
t.Cleanup(pool.Close)
|
|
return New(store.New(pool), nil).
|
|
WithRateLimit(limiter, 1000).
|
|
WithQuotas(anonQuota, 300, 100000)
|
|
}
|
|
|
|
func cleanupAccount(t *testing.T, email string) {
|
|
t.Helper()
|
|
dsn := os.Getenv("OPENGOODS_DATABASE_URL")
|
|
if dsn == "" {
|
|
dsn = "postgres://opengoods:opengoods@localhost:5432/opengoods?sslmode=disable"
|
|
}
|
|
ctx := context.Background()
|
|
pool, err := pgxpool.New(ctx, dsn)
|
|
if err != nil {
|
|
return
|
|
}
|
|
defer pool.Close()
|
|
_, _ = pool.Exec(ctx, "DELETE FROM app_user WHERE lower(email)=lower($1)", email)
|
|
_, _ = pool.Exec(ctx, "DELETE FROM api_key WHERE owner_email=$1", email)
|
|
}
|
|
|
|
func TestAnonTotalQuotaExhausts(t *testing.T) {
|
|
h := newQuotaHandler(t, 3)
|
|
// Unique client IP so the lifetime counter starts fresh for this test; the
|
|
// counter never expires, so drop it afterwards to keep runs independent.
|
|
n := time.Now().UnixNano()
|
|
ip := fmt.Sprintf("203.%d.%d.%d", n/65536%256, n/256%256, n%256)
|
|
t.Cleanup(func() { cleanupCounter(t, "ip:"+ip) })
|
|
call := func() *httptest.ResponseRecorder {
|
|
req := httptest.NewRequest(http.MethodGet, "/api/"+APIVersion+"/stats", nil)
|
|
req.RemoteAddr = ip + ":12345"
|
|
rec := httptest.NewRecorder()
|
|
h.Router().ServeHTTP(rec, req)
|
|
return rec
|
|
}
|
|
for i := 1; i <= 3; i++ {
|
|
if rec := call(); rec.Code != http.StatusOK {
|
|
t.Fatalf("call %d should be allowed, got %d (%s)", i, rec.Code, rec.Body.String())
|
|
}
|
|
}
|
|
rec := call()
|
|
if rec.Code != http.StatusForbidden {
|
|
t.Fatalf("4th call should be 403 quota_exhausted, got %d (%s)", rec.Code, rec.Body.String())
|
|
}
|
|
if !strings.Contains(rec.Body.String(), "quota_exhausted") {
|
|
t.Fatalf("expected quota_exhausted error, got %s", rec.Body.String())
|
|
}
|
|
}
|
|
|
|
func TestRegisterIssuesHigherQuotaKey(t *testing.T) {
|
|
h := newQuotaHandler(t, 1000)
|
|
email := fmt.Sprintf("h-user-%d@example.com", time.Now().UnixNano())
|
|
t.Cleanup(func() { cleanupAccount(t, email) })
|
|
|
|
body := fmt.Sprintf(`{"email":%q,"password":"supersecret"}`, email)
|
|
req := httptest.NewRequest(http.MethodPost, "/api/"+APIVersion+"/register", strings.NewReader(body))
|
|
req.RemoteAddr = "198.51.100.7:9999"
|
|
rec := httptest.NewRecorder()
|
|
h.Router().ServeHTTP(rec, req)
|
|
if rec.Code != http.StatusCreated {
|
|
t.Fatalf("register status = %d (%s)", rec.Code, rec.Body.String())
|
|
}
|
|
var resp keyResponse
|
|
if err := json.NewDecoder(rec.Body).Decode(&resp); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if resp.APIKey == "" || resp.QuotaTotal != 100000 || resp.RateLimitPerMin != 300 {
|
|
t.Fatalf("unexpected register response: %+v", resp)
|
|
}
|
|
|
|
// A second registration with the same email conflicts.
|
|
req2 := httptest.NewRequest(http.MethodPost, "/api/"+APIVersion+"/register", strings.NewReader(body))
|
|
req2.RemoteAddr = "198.51.100.7:9999"
|
|
rec2 := httptest.NewRecorder()
|
|
h.Router().ServeHTTP(rec2, req2)
|
|
if rec2.Code != http.StatusConflict {
|
|
t.Fatalf("duplicate register status = %d (%s)", rec2.Code, rec2.Body.String())
|
|
}
|
|
}
|