Initial open-source release (MIT): image2api AI gateway
Full Go backend + Vue 3 frontend, OpenAI-compatible API, multi-provider account pools, billing/admin, Docker one-command deploy with auto HTTPS. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
Claude Opus 4.8
@@ -0,0 +1,6 @@
|
||||
# Keep secrets, local data and build artifacts out of the Docker build context.
|
||||
.env
|
||||
.env.*
|
||||
bin/
|
||||
data/
|
||||
*.log
|
||||
@@ -0,0 +1,20 @@
|
||||
# Vivid AI backend config template. Copy to .env and fill in real values.
|
||||
# Real environment variables override .env.
|
||||
|
||||
APP_ENV=development
|
||||
APP_TITLE=Vivid AI
|
||||
HTTP_ADDR=:6666
|
||||
|
||||
# PostgreSQL (database must already exist; tables auto-migrate on boot)
|
||||
POSTGRES_DSN=host=127.0.0.1 user=postgres password=YOUR_PASSWORD dbname=vivid_ai port=5432 sslmode=disable TimeZone=Asia/Shanghai
|
||||
|
||||
# Redis
|
||||
REDIS_ADDR=127.0.0.1:6379
|
||||
REDIS_PASSWORD=
|
||||
REDIS_DB=0
|
||||
|
||||
# Frontend dev origins allowed for CORS
|
||||
CORS_ORIGINS=http://localhost:5173,http://127.0.0.1:5173
|
||||
|
||||
# Optional: override generated media root (defaults to ../../ai-gateway/data/generated)
|
||||
# GENERATED_ROOT=
|
||||
@@ -0,0 +1,9 @@
|
||||
.gocache/
|
||||
.gomodcache/
|
||||
backend.exe
|
||||
*.exe
|
||||
*.out
|
||||
*.test
|
||||
.env
|
||||
backend.out.log
|
||||
backend.err.log
|
||||
@@ -0,0 +1,24 @@
|
||||
# syntax=docker/dockerfile:1
|
||||
|
||||
# --- Stage 1: build the Go binary from source ---
|
||||
FROM golang:1.26-alpine AS build
|
||||
WORKDIR /src
|
||||
RUN apk add --no-cache git
|
||||
# Cache deps first for faster rebuilds.
|
||||
COPY go.mod go.sum ./
|
||||
RUN go mod download
|
||||
COPY . .
|
||||
RUN CGO_ENABLED=0 GOOS=linux go build -trimpath -ldflags="-s -w" -o /out/api ./cmd/api
|
||||
|
||||
# --- Stage 2: minimal runtime image ---
|
||||
FROM alpine:3.20
|
||||
# ca-certificates: outbound HTTPS to the AI providers. tzdata: POSTGRES_DSN sets
|
||||
# TimeZone=Asia/Shanghai. wget: container healthcheck.
|
||||
RUN apk add --no-cache ca-certificates tzdata wget
|
||||
WORKDIR /app
|
||||
COPY --from=build /out/api /app/api
|
||||
# Local fallback for generated media / reference uploads (RustFS/S3 is primary).
|
||||
RUN mkdir -p /app/data/generated && chmod +x /app/api
|
||||
ENV HTTP_ADDR=0.0.0.0:6666
|
||||
EXPOSE 6666
|
||||
ENTRYPOINT ["/app/api"]
|
||||
@@ -0,0 +1,55 @@
|
||||
# Vivid AI Backend
|
||||
|
||||
Go backend for `vivid-ai`, using:
|
||||
|
||||
- Gin
|
||||
- GORM
|
||||
- PostgreSQL
|
||||
- Redis
|
||||
|
||||
## Current scope
|
||||
|
||||
This is an in-progress rewrite. The current skeleton already includes:
|
||||
|
||||
- app bootstrap
|
||||
- PostgreSQL and Redis initialization
|
||||
- GORM auto-migrations
|
||||
- session storage in Redis
|
||||
- image access control for `/images/:user/:name`
|
||||
- public site endpoint: `/admin/api/site`
|
||||
- public showcase endpoint: `/admin/api/showcase`
|
||||
- session-based auth endpoint: `/admin/api/auth/me`
|
||||
|
||||
## Environment
|
||||
|
||||
Set these before running:
|
||||
|
||||
```powershell
|
||||
$env:POSTGRES_DSN="host=127.0.0.1 user=postgres password=postgres dbname=vivid_ai port=5432 sslmode=disable TimeZone=Asia/Shanghai"
|
||||
$env:REDIS_ADDR="127.0.0.1:6379"
|
||||
$env:HTTP_ADDR=":6061"
|
||||
```
|
||||
|
||||
Optional:
|
||||
|
||||
```powershell
|
||||
$env:APP_ENV="development"
|
||||
$env:APP_TITLE="Vivid AI"
|
||||
$env:SESSION_COOKIE_NAME="vivid_session"
|
||||
$env:CORS_ORIGINS="http://localhost:5173,http://127.0.0.1:5173"
|
||||
```
|
||||
|
||||
## Run
|
||||
|
||||
```powershell
|
||||
go run ./cmd/api
|
||||
```
|
||||
|
||||
## Notes
|
||||
|
||||
- Generated media defaults to `../../ai-gateway/data/generated` relative to the backend working directory.
|
||||
- Private images require either:
|
||||
- session cookie
|
||||
- bearer session token
|
||||
- bearer API key
|
||||
- Showcase images are public.
|
||||
@@ -0,0 +1,50 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"log"
|
||||
"net/http"
|
||||
"os"
|
||||
"os/signal"
|
||||
"syscall"
|
||||
"time"
|
||||
|
||||
"backend/internal/bootstrap"
|
||||
)
|
||||
|
||||
func main() {
|
||||
ctx := context.Background()
|
||||
|
||||
app, err := bootstrap.NewApp(ctx)
|
||||
if err != nil {
|
||||
log.Fatalf("bootstrap app: %v", err)
|
||||
}
|
||||
|
||||
srv := &http.Server{
|
||||
Addr: app.Config.HTTPAddr,
|
||||
Handler: app.Engine,
|
||||
ReadHeaderTimeout: 5 * time.Second,
|
||||
}
|
||||
|
||||
go func() {
|
||||
log.Printf("backend listening on %s", app.Config.HTTPAddr)
|
||||
if err := srv.ListenAndServe(); err != nil && err != http.ErrServerClosed {
|
||||
log.Fatalf("listen and serve: %v", err)
|
||||
}
|
||||
}()
|
||||
|
||||
stop := make(chan os.Signal, 1)
|
||||
signal.Notify(stop, syscall.SIGINT, syscall.SIGTERM)
|
||||
<-stop
|
||||
|
||||
shutdownCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||||
defer cancel()
|
||||
|
||||
if err := srv.Shutdown(shutdownCtx); err != nil {
|
||||
log.Printf("shutdown http server: %v", err)
|
||||
}
|
||||
|
||||
if err := app.Close(); err != nil {
|
||||
log.Printf("close app: %v", err)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,45 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"os"
|
||||
|
||||
"gorm.io/driver/postgres"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
type row struct {
|
||||
ID string
|
||||
Pool string
|
||||
Status string
|
||||
AccountEmail string
|
||||
ImageLimited bool
|
||||
VideoLimited bool
|
||||
}
|
||||
|
||||
func main() {
|
||||
dsn := os.Getenv("POSTGRES_DSN")
|
||||
if dsn == "" {
|
||||
fmt.Println("POSTGRES_DSN env is required, e.g. host=127.0.0.1 user=postgres password=... dbname=vivid_ai port=5432 sslmode=disable TimeZone=Asia/Shanghai")
|
||||
os.Exit(1)
|
||||
}
|
||||
db, err := gorm.Open(postgres.Open(dsn), &gorm.Config{})
|
||||
if err != nil {
|
||||
fmt.Println("open err:", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
|
||||
var pick row
|
||||
db.Raw(`SELECT id, pool, status, account_email, image_limited, video_limited FROM token_accounts WHERE pool='adobe' ORDER BY id LIMIT 1`).Scan(&pick)
|
||||
fmt.Printf("picked: id=%s email=%s status=%s image_limited=%v video_limited=%v\n", pick.ID, pick.AccountEmail, pick.Status, pick.ImageLimited, pick.VideoLimited)
|
||||
|
||||
if err := db.Exec(`UPDATE token_accounts SET video_limited=true, updated_at=now() WHERE id=?`, pick.ID).Error; err != nil {
|
||||
fmt.Println("update err:", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
fmt.Println("-> set video_limited=true")
|
||||
|
||||
var after row
|
||||
db.Raw(`SELECT id, pool, status, account_email, image_limited, video_limited FROM token_accounts WHERE id=?`, pick.ID).Scan(&after)
|
||||
fmt.Printf("after: id=%s status=%s image_limited=%v video_limited=%v\n", after.ID, after.Status, after.ImageLimited, after.VideoLimited)
|
||||
}
|
||||
@@ -0,0 +1,71 @@
|
||||
module backend
|
||||
|
||||
go 1.26.0
|
||||
|
||||
replace github.com/quic-go/qpack => github.com/quic-go/qpack v0.5.1
|
||||
|
||||
require (
|
||||
github.com/bogdanfinn/fhttp v0.6.8
|
||||
github.com/bogdanfinn/tls-client v1.11.2
|
||||
github.com/gin-contrib/cors v1.7.6
|
||||
github.com/gin-gonic/gin v1.11.0
|
||||
github.com/google/uuid v1.6.0
|
||||
github.com/matoous/go-nanoid/v2 v2.1.0
|
||||
github.com/redis/go-redis/v9 v9.16.0
|
||||
golang.org/x/crypto v0.46.0
|
||||
gorm.io/datatypes v1.2.7
|
||||
gorm.io/driver/postgres v1.6.0
|
||||
gorm.io/gorm v1.31.0
|
||||
)
|
||||
|
||||
require (
|
||||
filippo.io/edwards25519 v1.1.0 // indirect
|
||||
github.com/andybalholm/brotli v1.2.0 // indirect
|
||||
github.com/bdandy/go-errors v1.2.2 // indirect
|
||||
github.com/bdandy/go-socks4 v1.2.3 // indirect
|
||||
github.com/bogdanfinn/quic-go-utls v1.0.4-utls // indirect
|
||||
github.com/bogdanfinn/utls v1.7.7-barnius // indirect
|
||||
github.com/bogdanfinn/websocket v1.5.5-barnius // indirect
|
||||
github.com/bytedance/sonic v1.14.0 // indirect
|
||||
github.com/bytedance/sonic/loader v0.3.0 // indirect
|
||||
github.com/cespare/xxhash/v2 v2.3.0 // indirect
|
||||
github.com/cloudwego/base64x v0.1.6 // indirect
|
||||
github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f // indirect
|
||||
github.com/gabriel-vasile/mimetype v1.4.9 // indirect
|
||||
github.com/gin-contrib/sse v1.1.0 // indirect
|
||||
github.com/go-playground/locales v0.14.1 // indirect
|
||||
github.com/go-playground/universal-translator v0.18.1 // indirect
|
||||
github.com/go-playground/validator/v10 v10.27.0 // indirect
|
||||
github.com/go-sql-driver/mysql v1.8.1 // indirect
|
||||
github.com/goccy/go-json v0.10.5 // indirect
|
||||
github.com/goccy/go-yaml v1.18.0 // indirect
|
||||
github.com/jackc/pgpassfile v1.0.0 // indirect
|
||||
github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 // indirect
|
||||
github.com/jackc/pgx/v5 v5.6.0 // indirect
|
||||
github.com/jackc/puddle/v2 v2.2.2 // indirect
|
||||
github.com/jinzhu/inflection v1.0.0 // indirect
|
||||
github.com/jinzhu/now v1.1.5 // indirect
|
||||
github.com/json-iterator/go v1.1.12 // indirect
|
||||
github.com/klauspost/compress v1.18.2 // indirect
|
||||
github.com/klauspost/cpuid/v2 v2.3.0 // indirect
|
||||
github.com/leodido/go-urn v1.4.0 // indirect
|
||||
github.com/mattn/go-isatty v0.0.20 // indirect
|
||||
github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect
|
||||
github.com/modern-go/reflect2 v1.0.2 // indirect
|
||||
github.com/pelletier/go-toml/v2 v2.2.4 // indirect
|
||||
github.com/quic-go/qpack v0.5.1 // indirect
|
||||
github.com/quic-go/quic-go v0.54.0 // indirect
|
||||
github.com/tam7t/hpkp v0.0.0-20160821193359-2b70b4024ed5 // indirect
|
||||
github.com/twitchyliquid64/golang-asm v0.15.1 // indirect
|
||||
github.com/ugorji/go/codec v1.3.0 // indirect
|
||||
go.uber.org/mock v0.5.2 // indirect
|
||||
golang.org/x/arch v0.20.0 // indirect
|
||||
golang.org/x/mod v0.30.0 // indirect
|
||||
golang.org/x/net v0.48.0 // indirect
|
||||
golang.org/x/sync v0.19.0 // indirect
|
||||
golang.org/x/sys v0.39.0 // indirect
|
||||
golang.org/x/text v0.32.0 // indirect
|
||||
golang.org/x/tools v0.39.0 // indirect
|
||||
google.golang.org/protobuf v1.36.9 // indirect
|
||||
gorm.io/driver/mysql v1.5.6 // indirect
|
||||
)
|
||||
+178
@@ -0,0 +1,178 @@
|
||||
filippo.io/edwards25519 v1.1.0 h1:FNf4tywRC1HmFuKW5xopWpigGjJKiJSV0Cqo0cJWDaA=
|
||||
filippo.io/edwards25519 v1.1.0/go.mod h1:BxyFTGdWcka3PhytdK4V28tE5sGfRvvvRV7EaN4VDT4=
|
||||
github.com/andybalholm/brotli v1.2.0 h1:ukwgCxwYrmACq68yiUqwIWnGY0cTPox/M94sVwToPjQ=
|
||||
github.com/andybalholm/brotli v1.2.0/go.mod h1:rzTDkvFWvIrjDXZHkuS16NPggd91W3kUSvPlQ1pLaKY=
|
||||
github.com/bdandy/go-errors v1.2.2 h1:WdFv/oukjTJCLa79UfkGmwX7ZxONAihKu4V0mLIs11Q=
|
||||
github.com/bdandy/go-errors v1.2.2/go.mod h1:NkYHl4Fey9oRRdbB1CoC6e84tuqQHiqrOcZpqFEkBxM=
|
||||
github.com/bdandy/go-socks4 v1.2.3 h1:Q6Y2heY1GRjCtHbmlKfnwrKVU/k81LS8mRGLRlmDlic=
|
||||
github.com/bdandy/go-socks4 v1.2.3/go.mod h1:98kiVFgpdogR8aIGLWLvjDVZ8XcKPsSI/ypGrO+bqHI=
|
||||
github.com/bogdanfinn/fhttp v0.6.8 h1:LiQyHOY3i0QoxxNB7nq27/nGNNbtPj0fuBPozhR7Ws4=
|
||||
github.com/bogdanfinn/fhttp v0.6.8/go.mod h1:A+EKDzMx2hb4IUbMx4TlkoHnaJEiLl8r/1Ss1Y+5e5M=
|
||||
github.com/bogdanfinn/quic-go-utls v1.0.4-utls h1:zPjusVVNeJFA2ORMAP0rjnrZrBkV4Dnia4e6ToOfUDA=
|
||||
github.com/bogdanfinn/quic-go-utls v1.0.4-utls/go.mod h1:UONJOaHGWho08kZtkkgH7GjktEPjMemGxjTcNpVPZVA=
|
||||
github.com/bogdanfinn/quic-go-utls v1.0.9-utls h1:tV6eDEiRbRCcepALSzxR94JUVD3N3ACIiRLgyc2Ep8s=
|
||||
github.com/bogdanfinn/quic-go-utls v1.0.9-utls/go.mod h1:aHph9B9H9yPOt5xnhWKSOum27DJAqpiHzwX+gjvaXcg=
|
||||
github.com/bogdanfinn/tls-client v1.11.2 h1:o6qX0L1cEi+4MaBqujxqOeK254VZM20t3QR+A34/V6I=
|
||||
github.com/bogdanfinn/tls-client v1.11.2/go.mod h1:qQIsVGe35NdxYEozNh9JuDZ+aOaOEq2tKAsu2iYEGZg=
|
||||
github.com/bogdanfinn/tls-client v1.15.1 h1:KiFAlED55DJ8Fcocn+/1nX6PrDFcttIHAf/GDkV6KN8=
|
||||
github.com/bogdanfinn/tls-client v1.15.1/go.mod h1:LsU6mXVn8MOFDwTkyRfI7V1BZM1p0wf2ZfZsICW/1fM=
|
||||
github.com/bogdanfinn/utls v1.7.7-barnius h1:OuJ497cc7F3yKNVHRsYPQdGggmk5x6+V5ZlrCR7fOLU=
|
||||
github.com/bogdanfinn/utls v1.7.7-barnius/go.mod h1:aAK1VZQlpKZClF1WEQeq6kyclbkPq4hz6xTbB5xSlmg=
|
||||
github.com/bogdanfinn/websocket v1.5.5-barnius h1:bY+qnxpai1qe7Jmjx+Sds/cmOSpuuLoR8x61rWltjOI=
|
||||
github.com/bogdanfinn/websocket v1.5.5-barnius/go.mod h1:gvvEw6pTKHb7yOiFvIfAFTStQWyrm25BMVCTj5wRSsI=
|
||||
github.com/bsm/ginkgo/v2 v2.12.0 h1:Ny8MWAHyOepLGlLKYmXG4IEkioBysk6GpaRTLC8zwWs=
|
||||
github.com/bsm/ginkgo/v2 v2.12.0/go.mod h1:SwYbGRRDovPVboqFv0tPTcG1sN61LM1Z4ARdbAV9g4c=
|
||||
github.com/bsm/gomega v1.27.10 h1:yeMWxP2pV2fG3FgAODIY8EiRE3dy0aeFYt4l7wh6yKA=
|
||||
github.com/bsm/gomega v1.27.10/go.mod h1:JyEr/xRbxbtgWNi8tIEVPUYZ5Dzef52k01W3YH0H+O0=
|
||||
github.com/bytedance/sonic v1.14.0 h1:/OfKt8HFw0kh2rj8N0F6C/qPGRESq0BbaNZgcNXXzQQ=
|
||||
github.com/bytedance/sonic v1.14.0/go.mod h1:WoEbx8WTcFJfzCe0hbmyTGrfjt8PzNEBdxlNUO24NhA=
|
||||
github.com/bytedance/sonic/loader v0.3.0 h1:dskwH8edlzNMctoruo8FPTJDF3vLtDT0sXZwvZJyqeA=
|
||||
github.com/bytedance/sonic/loader v0.3.0/go.mod h1:N8A3vUdtUebEY2/VQC0MyhYeKUFosQU6FxH2JmUe6VI=
|
||||
github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs=
|
||||
github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs=
|
||||
github.com/cloudwego/base64x v0.1.6 h1:t11wG9AECkCDk5fMSoxmufanudBtJ+/HemLstXDLI2M=
|
||||
github.com/cloudwego/base64x v0.1.6/go.mod h1:OFcloc187FXDaYHvrNIjxSe8ncn0OOM8gEHfghB2IPU=
|
||||
github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
|
||||
github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
|
||||
github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
|
||||
github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f h1:lO4WD4F/rVNCu3HqELle0jiPLLBs70cWOduZpkS1E78=
|
||||
github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f/go.mod h1:cuUVRXasLTGF7a8hSLbxyZXjz+1KgoB3wDUb6vlszIc=
|
||||
github.com/gabriel-vasile/mimetype v1.4.9 h1:5k+WDwEsD9eTLL8Tz3L0VnmVh9QxGjRmjBvAG7U/oYY=
|
||||
github.com/gabriel-vasile/mimetype v1.4.9/go.mod h1:WnSQhFKJuBlRyLiKohA/2DtIlPFAbguNaG7QCHcyGok=
|
||||
github.com/gin-contrib/cors v1.7.6 h1:3gQ8GMzs1Ylpf70y8bMw4fVpycXIeX1ZemuSQIsnQQY=
|
||||
github.com/gin-contrib/cors v1.7.6/go.mod h1:Ulcl+xN4jel9t1Ry8vqph23a60FwH9xVLd+3ykmTjOk=
|
||||
github.com/gin-contrib/sse v1.1.0 h1:n0w2GMuUpWDVp7qSpvze6fAu9iRxJY4Hmj6AmBOU05w=
|
||||
github.com/gin-contrib/sse v1.1.0/go.mod h1:hxRZ5gVpWMT7Z0B0gSNYqqsSCNIJMjzvm6fqCz9vjwM=
|
||||
github.com/gin-gonic/gin v1.11.0 h1:OW/6PLjyusp2PPXtyxKHU0RbX6I/l28FTdDlae5ueWk=
|
||||
github.com/gin-gonic/gin v1.11.0/go.mod h1:+iq/FyxlGzII0KHiBGjuNn4UNENUlKbGlNmc+W50Dls=
|
||||
github.com/go-playground/assert/v2 v2.2.0 h1:JvknZsQTYeFEAhQwI4qEt9cyV5ONwRHC+lYKSsYSR8s=
|
||||
github.com/go-playground/assert/v2 v2.2.0/go.mod h1:VDjEfimB/XKnb+ZQfWdccd7VUvScMdVu0Titje2rxJ4=
|
||||
github.com/go-playground/locales v0.14.1 h1:EWaQ/wswjilfKLTECiXz7Rh+3BjFhfDFKv/oXslEjJA=
|
||||
github.com/go-playground/locales v0.14.1/go.mod h1:hxrqLVvrK65+Rwrd5Fc6F2O76J/NuW9t0sjnWqG1slY=
|
||||
github.com/go-playground/universal-translator v0.18.1 h1:Bcnm0ZwsGyWbCzImXv+pAJnYK9S473LQFuzCbDbfSFY=
|
||||
github.com/go-playground/universal-translator v0.18.1/go.mod h1:xekY+UJKNuX9WP91TpwSH2VMlDf28Uj24BCp08ZFTUY=
|
||||
github.com/go-playground/validator/v10 v10.27.0 h1:w8+XrWVMhGkxOaaowyKH35gFydVHOvC0/uWoy2Fzwn4=
|
||||
github.com/go-playground/validator/v10 v10.27.0/go.mod h1:I5QpIEbmr8On7W0TktmJAumgzX4CA1XNl4ZmDuVHKKo=
|
||||
github.com/go-sql-driver/mysql v1.7.0/go.mod h1:OXbVy3sEdcQ2Doequ6Z5BW6fXNQTmx+9S1MCJN5yJMI=
|
||||
github.com/go-sql-driver/mysql v1.8.1 h1:LedoTUt/eveggdHS9qUFC1EFSa8bU2+1pZjSRpvNJ1Y=
|
||||
github.com/go-sql-driver/mysql v1.8.1/go.mod h1:wEBSXgmK//2ZFJyE+qWnIsVGmvmEKlqwuVSjsCm7DZg=
|
||||
github.com/goccy/go-json v0.10.5 h1:Fq85nIqj+gXn/S5ahsiTlK3TmC85qgirsdTP/+DeaC4=
|
||||
github.com/goccy/go-json v0.10.5/go.mod h1:oq7eo15ShAhp70Anwd5lgX2pLfOS3QCiwU/PULtXL6M=
|
||||
github.com/goccy/go-yaml v1.18.0 h1:8W7wMFS12Pcas7KU+VVkaiCng+kG8QiFeFwzFb+rwuw=
|
||||
github.com/goccy/go-yaml v1.18.0/go.mod h1:XBurs7gK8ATbW4ZPGKgcbrY1Br56PdM69F7LkFRi1kA=
|
||||
github.com/golang-sql/civil v0.0.0-20220223132316-b832511892a9 h1:au07oEsX2xN0ktxqI+Sida1w446QrXBRJ0nee3SNZlA=
|
||||
github.com/golang-sql/civil v0.0.0-20220223132316-b832511892a9/go.mod h1:8vg3r2VgvsThLBIFL93Qb5yWzgyZWhEmBwUJWevAkK0=
|
||||
github.com/golang-sql/sqlexp v0.1.0 h1:ZCD6MBpcuOVfGVqsEmY5/4FtYiKz6tSyUv9LPEDei6A=
|
||||
github.com/golang-sql/sqlexp v0.1.0/go.mod h1:J4ad9Vo8ZCWQ2GMrC4UCQy1JpCbwU9m3EOqtpKwwwHI=
|
||||
github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8=
|
||||
github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU=
|
||||
github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg=
|
||||
github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
|
||||
github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
|
||||
github.com/jackc/pgpassfile v1.0.0 h1:/6Hmqy13Ss2zCq62VdNG8tM1wchn8zjSGOBJ6icpsIM=
|
||||
github.com/jackc/pgpassfile v1.0.0/go.mod h1:CEx0iS5ambNFdcRtxPj5JhEz+xB6uRky5eyVu/W2HEg=
|
||||
github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 h1:iCEnooe7UlwOQYpKFhBabPMi4aNAfoODPEFNiAnClxo=
|
||||
github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761/go.mod h1:5TJZWKEWniPve33vlWYSoGYefn3gLQRzjfDlhSJ9ZKM=
|
||||
github.com/jackc/pgx/v5 v5.6.0 h1:SWJzexBzPL5jb0GEsrPMLIsi/3jOo7RHlzTjcAeDrPY=
|
||||
github.com/jackc/pgx/v5 v5.6.0/go.mod h1:DNZ/vlrUnhWCoFGxHAG8U2ljioxukquj7utPDgtQdTw=
|
||||
github.com/jackc/puddle/v2 v2.2.2 h1:PR8nw+E/1w0GLuRFSmiioY6UooMp6KJv0/61nB7icHo=
|
||||
github.com/jackc/puddle/v2 v2.2.2/go.mod h1:vriiEXHvEE654aYKXXjOvZM39qJ0q+azkZFrfEOc3H4=
|
||||
github.com/jinzhu/inflection v1.0.0 h1:K317FqzuhWc8YvSVlFMCCUb36O/S9MCKRDI7QkRKD/E=
|
||||
github.com/jinzhu/inflection v1.0.0/go.mod h1:h+uFLlag+Qp1Va5pdKtLDYj+kHp5pxUVkryuEj+Srlc=
|
||||
github.com/jinzhu/now v1.1.5 h1:/o9tlHleP7gOFmsnYNz3RGnqzefHA47wQpKrrdTIwXQ=
|
||||
github.com/jinzhu/now v1.1.5/go.mod h1:d3SSVoowX0Lcu0IBviAWJpolVfI5UJVZZ7cO71lE/z8=
|
||||
github.com/json-iterator/go v1.1.12 h1:PV8peI4a0ysnczrg+LtxykD8LfKY9ML6u2jnxaEnrnM=
|
||||
github.com/json-iterator/go v1.1.12/go.mod h1:e30LSqwooZae/UwlEbR2852Gd8hjQvJoHmT4TnhNGBo=
|
||||
github.com/klauspost/compress v1.18.2 h1:iiPHWW0YrcFgpBYhsA6D1+fqHssJscY/Tm/y2Uqnapk=
|
||||
github.com/klauspost/compress v1.18.2/go.mod h1:R0h/fSBs8DE4ENlcrlib3PsXS61voFxhIs2DeRhCvJ4=
|
||||
github.com/klauspost/cpuid/v2 v2.3.0 h1:S4CRMLnYUhGeDFDqkGriYKdfoFlDnMtqTiI/sFzhA9Y=
|
||||
github.com/klauspost/cpuid/v2 v2.3.0/go.mod h1:hqwkgyIinND0mEev00jJYCxPNVRVXFQeu1XKlok6oO0=
|
||||
github.com/leodido/go-urn v1.4.0 h1:WT9HwE9SGECu3lg4d/dIA+jxlljEa1/ffXKmRjqdmIQ=
|
||||
github.com/leodido/go-urn v1.4.0/go.mod h1:bvxc+MVxLKB4z00jd1z+Dvzr47oO32F/QSNjSBOlFxI=
|
||||
github.com/matoous/go-nanoid/v2 v2.1.0 h1:P64+dmq21hhWdtvZfEAofnvJULaRR1Yib0+PnU669bE=
|
||||
github.com/matoous/go-nanoid/v2 v2.1.0/go.mod h1:KlbGNQ+FhrUNIHUxZdL63t7tl4LaPkZNpUULS8H4uVM=
|
||||
github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY=
|
||||
github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y=
|
||||
github.com/mattn/go-sqlite3 v1.14.22 h1:2gZY6PC6kBnID23Tichd1K+Z0oS6nE/XwU+Vz/5o4kU=
|
||||
github.com/mattn/go-sqlite3 v1.14.22/go.mod h1:Uh1q+B4BYcTPb+yiD3kU8Ct7aC0hY9fxUwlHK0RXw+Y=
|
||||
github.com/microsoft/go-mssqldb v1.7.2 h1:CHkFJiObW7ItKTJfHo1QX7QBBD1iV+mn1eOyRP3b/PA=
|
||||
github.com/microsoft/go-mssqldb v1.7.2/go.mod h1:kOvZKUdrhhFQmxLZqbwUV0rHkNkZpthMITIb2Ko1IoA=
|
||||
github.com/modern-go/concurrent v0.0.0-20180228061459-e0a39a4cb421/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q=
|
||||
github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd h1:TRLaZ9cD/w8PVh93nsPXa1VrQ6jlwL5oN8l14QlcNfg=
|
||||
github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q=
|
||||
github.com/modern-go/reflect2 v1.0.2 h1:xBagoLtFs94CBntxluKeaWgTMpvLxC4ur3nMaC9Gz0M=
|
||||
github.com/modern-go/reflect2 v1.0.2/go.mod h1:yWuevngMOJpCy52FWWMvUC8ws7m/LJsjYzDa0/r8luk=
|
||||
github.com/pelletier/go-toml/v2 v2.2.4 h1:mye9XuhQ6gvn5h28+VilKrrPoQVanw5PMw/TB0t5Ec4=
|
||||
github.com/pelletier/go-toml/v2 v2.2.4/go.mod h1:2gIqNv+qfxSVS7cM2xJQKtLSTLUE9V8t9Stt+h56mCY=
|
||||
github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
|
||||
github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
|
||||
github.com/quic-go/qpack v0.5.1 h1:giqksBPnT/HDtZ6VhtFKgoLOWmlyo9Ei6u9PqzIMbhI=
|
||||
github.com/quic-go/qpack v0.5.1/go.mod h1:+PC4XFrEskIVkcLzpEkbLqq1uCoxPhQuvK5rH1ZgaEg=
|
||||
github.com/quic-go/qpack v0.6.0 h1:g7W+BMYynC1LbYLSqRt8PBg5Tgwxn214ZZR34VIOjz8=
|
||||
github.com/quic-go/qpack v0.6.0/go.mod h1:lUpLKChi8njB4ty2bFLX2x4gzDqXwUpaO1DP9qMDZII=
|
||||
github.com/quic-go/quic-go v0.54.0 h1:6s1YB9QotYI6Ospeiguknbp2Znb/jZYjZLRXn9kMQBg=
|
||||
github.com/quic-go/quic-go v0.54.0/go.mod h1:e68ZEaCdyviluZmy44P6Iey98v/Wfz6HCjQEm+l8zTY=
|
||||
github.com/redis/go-redis/v9 v9.16.0 h1:OotgqgLSRCmzfqChbQyG1PHC3tLNR89DG4jdOERSEP4=
|
||||
github.com/redis/go-redis/v9 v9.16.0/go.mod h1:u410H11HMLoB+TP67dz8rL9s6QW2j76l0//kSOd3370=
|
||||
github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
|
||||
github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw=
|
||||
github.com/stretchr/objx v0.5.0/go.mod h1:Yh+to48EsGEfYuaHDzXPcE3xhTkx73EhmCGUpEOglKo=
|
||||
github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI=
|
||||
github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
|
||||
github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
|
||||
github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU=
|
||||
github.com/stretchr/testify v1.8.1/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4=
|
||||
github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U=
|
||||
github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U=
|
||||
github.com/tam7t/hpkp v0.0.0-20160821193359-2b70b4024ed5 h1:YqAladjX7xpA6BM04leXMWAEjS0mTZ5kUU9KRBriQJc=
|
||||
github.com/tam7t/hpkp v0.0.0-20160821193359-2b70b4024ed5/go.mod h1:2JjD2zLQYH5HO74y5+aE3remJQvl6q4Sn6aWA2wD1Ng=
|
||||
github.com/twitchyliquid64/golang-asm v0.15.1 h1:SU5vSMR7hnwNxj24w34ZyCi/FmDZTkS4MhqMhdFk5YI=
|
||||
github.com/twitchyliquid64/golang-asm v0.15.1/go.mod h1:a1lVb/DtPvCB8fslRZhAngC2+aY1QWCk3Cedj/Gdt08=
|
||||
github.com/ugorji/go/codec v1.3.0 h1:Qd2W2sQawAfG8XSvzwhBeoGq71zXOC/Q1E9y/wUcsUA=
|
||||
github.com/ugorji/go/codec v1.3.0/go.mod h1:pRBVtBSKl77K30Bv8R2P+cLSGaTtex6fsA2Wjqmfxj4=
|
||||
github.com/xyproto/randomstring v1.0.5 h1:YtlWPoRdgMu3NZtP45drfy1GKoojuR7hmRcnhZqKjWU=
|
||||
github.com/xyproto/randomstring v1.0.5/go.mod h1:rgmS5DeNXLivK7YprL0pY+lTuhNQW3iGxZ18UQApw/E=
|
||||
go.uber.org/mock v0.5.2 h1:LbtPTcP8A5k9WPXj54PPPbjcI4Y6lhyOZXn+VS7wNko=
|
||||
go.uber.org/mock v0.5.2/go.mod h1:wLlUxC2vVTPTaE3UD51E0BGOAElKrILxhVSDYQLld5o=
|
||||
golang.org/x/arch v0.20.0 h1:dx1zTU0MAE98U+TQ8BLl7XsJbgze2WnNKF/8tGp/Q6c=
|
||||
golang.org/x/arch v0.20.0/go.mod h1:bdwinDaKcfZUGpH09BB7ZmOfhalA8lQdzl62l8gGWsk=
|
||||
golang.org/x/crypto v0.46.0 h1:cKRW/pmt1pKAfetfu+RCEvjvZkA9RimPbh7bhFjGVBU=
|
||||
golang.org/x/crypto v0.46.0/go.mod h1:Evb/oLKmMraqjZ2iQTwDwvCtJkczlDuTmdJXoZVzqU0=
|
||||
golang.org/x/mod v0.30.0 h1:fDEXFVZ/fmCKProc/yAXXUijritrDzahmwwefnjoPFk=
|
||||
golang.org/x/mod v0.30.0/go.mod h1:lAsf5O2EvJeSFMiBxXDki7sCgAxEUcZHXoXMKT4GJKc=
|
||||
golang.org/x/net v0.0.0-20211104170005-ce137452f963/go.mod h1:9nx3DQGgdP8bBQD5qxJ1jj9UTztislL4KSBs9R2vV5Y=
|
||||
golang.org/x/net v0.48.0 h1:zyQRTTrjc33Lhh0fBgT/H3oZq9WuvRR5gPC70xpDiQU=
|
||||
golang.org/x/net v0.48.0/go.mod h1:+ndRgGjkh8FGtu1w1FGbEC31if4VrNVMuKTgcAAnQRY=
|
||||
golang.org/x/sync v0.19.0 h1:vV+1eWNmZ5geRlYjzm2adRgW2/mcpevXNg50YZtPCE4=
|
||||
golang.org/x/sync v0.19.0/go.mod h1:9KTHXmSnoGruLpwFjVSX0lNNA75CykiMECbovNTZqGI=
|
||||
golang.org/x/sys v0.0.0-20201119102817-f84b799fce68/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
|
||||
golang.org/x/sys v0.0.0-20210423082822-04245dca01da/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
|
||||
golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
|
||||
golang.org/x/sys v0.39.0 h1:CvCKL8MeisomCi6qNZ+wbb0DN9E5AATixKsvNtMoMFk=
|
||||
golang.org/x/sys v0.39.0/go.mod h1:OgkHotnGiDImocRcuBABYBEXf8A9a87e/uXjp9XT3ks=
|
||||
golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo=
|
||||
golang.org/x/text v0.3.6/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
|
||||
golang.org/x/text v0.32.0 h1:ZD01bjUt1FQ9WJ0ClOL5vxgxOI/sVCNgX1YtKwcY0mU=
|
||||
golang.org/x/text v0.32.0/go.mod h1:o/rUWzghvpD5TXrTIBuJU77MTaN0ljMWE47kxGJQ7jY=
|
||||
golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
|
||||
golang.org/x/tools v0.39.0 h1:ik4ho21kwuQln40uelmciQPp9SipgNDdrafrYA4TmQQ=
|
||||
golang.org/x/tools v0.39.0/go.mod h1:JnefbkDPyD8UU2kI5fuf8ZX4/yUeh9W877ZeBONxUqQ=
|
||||
google.golang.org/protobuf v1.36.9 h1:w2gp2mA27hUeUzj9Ex9FBjsBm40zfaDtEWow293U7Iw=
|
||||
google.golang.org/protobuf v1.36.9/go.mod h1:fuxRtAxBytpl4zzqUh6/eyUujkJdNiuEkXntxiD/uRU=
|
||||
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
|
||||
gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
|
||||
gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
|
||||
gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
|
||||
gorm.io/datatypes v1.2.7 h1:ww9GAhF1aGXZY3EB3cJPJ7//JiuQo7DlQA7NNlVaTdk=
|
||||
gorm.io/datatypes v1.2.7/go.mod h1:M2iO+6S3hhi4nAyYe444Pcb0dcIiOMJ7QHaUXxyiNZY=
|
||||
gorm.io/driver/mysql v1.5.6 h1:Ld4mkIickM+EliaQZQx3uOJDJHtrd70MxAUqWqlx3Y8=
|
||||
gorm.io/driver/mysql v1.5.6/go.mod h1:sEtPWMiqiN1N1cMXoXmBbd8C6/l+TESwriotuRRpkDM=
|
||||
gorm.io/driver/postgres v1.6.0 h1:2dxzU8xJ+ivvqTRph34QX+WrRaJlmfyPqXmoGVjMBa4=
|
||||
gorm.io/driver/postgres v1.6.0/go.mod h1:vUw0mrGgrTK+uPHEhAdV4sfFELrByKVGnaVRkXDhtWo=
|
||||
gorm.io/driver/sqlite v1.6.0 h1:WHRRrIiulaPiPFmDcod6prc4l2VGVWHz80KspNsxSfQ=
|
||||
gorm.io/driver/sqlite v1.6.0/go.mod h1:AO9V1qIQddBESngQUKWL9yoH93HIeA1X6V633rBwyT8=
|
||||
gorm.io/driver/sqlserver v1.6.0 h1:VZOBQVsVhkHU/NzNhRJKoANt5pZGQAS1Bwc6m6dgfnc=
|
||||
gorm.io/driver/sqlserver v1.6.0/go.mod h1:WQzt4IJo/WHKnckU9jXBLMJIVNMVeTu25dnOzehntWw=
|
||||
gorm.io/gorm v1.25.7/go.mod h1:hbnx/Oo0ChWMn1BIhpy1oYozzpM15i4YPuHDmfYtwg8=
|
||||
gorm.io/gorm v1.31.0 h1:0VlycGreVhK7RF/Bwt51Fk8v0xLiiiFdbGDPIZQ7mJY=
|
||||
gorm.io/gorm v1.31.0/go.mod h1:XyQVbO2k6YkOis7C2437jSit3SsDK72s7n7rsSHd+Gs=
|
||||
@@ -0,0 +1,174 @@
|
||||
package bootstrap
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"os"
|
||||
"time"
|
||||
|
||||
"backend/internal/config"
|
||||
"backend/internal/http/handler"
|
||||
"backend/internal/http/router"
|
||||
"backend/internal/model"
|
||||
"backend/internal/provider/adobe"
|
||||
"backend/internal/provider/chatgpt"
|
||||
"backend/internal/provider/imagine"
|
||||
"backend/internal/provider/krea"
|
||||
"backend/internal/provider/leonardo"
|
||||
"backend/internal/provider/runway"
|
||||
"backend/internal/repo"
|
||||
"backend/internal/service"
|
||||
"backend/internal/storage"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/redis/go-redis/v9"
|
||||
"gorm.io/driver/postgres"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
type App struct {
|
||||
Config *config.Config
|
||||
DB *gorm.DB
|
||||
Redis *redis.Client
|
||||
Engine *gin.Engine
|
||||
maintenanceCancel context.CancelFunc
|
||||
}
|
||||
|
||||
func NewApp(ctx context.Context) (*App, error) {
|
||||
cfg, err := config.Load()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// Ensure the media root (generated outputs + uploaded reference images)
|
||||
// exists from the first request — don't rely on lazy per-file MkdirAll.
|
||||
if err := os.MkdirAll(cfg.GeneratedRoot, 0o755); err != nil {
|
||||
return nil, fmt.Errorf("create generated root %s: %w", cfg.GeneratedRoot, err)
|
||||
}
|
||||
|
||||
// TranslateError: 把驱动层错误(如 Postgres 23505 唯一冲突)翻译成 gorm.ErrDuplicatedKey,
|
||||
// 否则各 import-*（krea/adobe/leonardo/runway）里的 errors.Is(err, gorm.ErrDuplicatedKey)
|
||||
// 兜底命不中,重复导入会直接抛原始错误 → 400,而不是按预期 Update 已有行。
|
||||
db, err := gorm.Open(postgres.Open(cfg.PostgresDSN), &gorm.Config{TranslateError: true})
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("open postgres: %w", err)
|
||||
}
|
||||
|
||||
sqlDB, err := db.DB()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("sql db: %w", err)
|
||||
}
|
||||
sqlDB.SetMaxIdleConns(5)
|
||||
sqlDB.SetMaxOpenConns(20)
|
||||
sqlDB.SetConnMaxLifetime(30 * time.Minute)
|
||||
|
||||
if err := db.WithContext(ctx).AutoMigrate(model.AutoMigrateModels()...); err != nil {
|
||||
return nil, fmt.Errorf("auto migrate: %w", err)
|
||||
}
|
||||
// Hard backstop for "one marketing code per user per batch": a partial unique
|
||||
// index. Concurrent double-redeems that slip past the in-tx count check still
|
||||
// fail here. AutoMigrate can't express partial indexes, so do it raw.
|
||||
if err := db.WithContext(ctx).Exec(`CREATE UNIQUE INDEX IF NOT EXISTS uniq_cdk_marketing_batch_user ` +
|
||||
`ON cdk_codes (batch_id, redeemed_by) WHERE type = 'marketing' AND redeemed_by IS NOT NULL`).Error; err != nil {
|
||||
return nil, fmt.Errorf("cdk marketing index: %w", err)
|
||||
}
|
||||
if err := seedDefaults(ctx, db); err != nil {
|
||||
return nil, fmt.Errorf("seed defaults: %w", err)
|
||||
}
|
||||
|
||||
rdb := redis.NewClient(&redis.Options{
|
||||
Addr: cfg.RedisAddr,
|
||||
Password: cfg.RedisPassword,
|
||||
DB: cfg.RedisDB,
|
||||
})
|
||||
if err := rdb.Ping(ctx).Err(); err != nil {
|
||||
return nil, fmt.Errorf("ping redis: %w", err)
|
||||
}
|
||||
|
||||
userRepo := repo.NewUserRepository(db)
|
||||
showcaseRepo := repo.NewShowcaseRepository(db)
|
||||
siteRepo := repo.NewSiteSettingRepository(db, rdb)
|
||||
modelRepo := repo.NewModelRepository(db)
|
||||
eventRepo := repo.NewEventRepository(db)
|
||||
cdkRepo := repo.NewCDKRepository(db)
|
||||
apiKeyRepo := repo.NewAPIKeyRepository(db)
|
||||
tokenRepo := repo.NewTokenRepository(db)
|
||||
refreshRepo := repo.NewRefreshProfileRepository(db)
|
||||
sessionSvc := service.NewSessionService(rdb, cfg.SessionTTL, cfg.SessionSlideAfter)
|
||||
emailCodeSvc := service.NewEmailCodeService(rdb)
|
||||
smtpSvc := service.NewSMTPService()
|
||||
rateLimitSvc := service.NewRateLimitService(rdb)
|
||||
rustfsClient := storage.New(cfg.RustFSEndpoint, cfg.RustFSBucket, cfg.RustFSAccessKey, cfg.RustFSSecretKey)
|
||||
authSvc := service.NewAuthService(userRepo, siteRepo, sessionSvc, emailCodeSvc, smtpSvc)
|
||||
appSettingsSvc := service.NewAppSettingsService(siteRepo, eventRepo, smtpSvc, rustfsClient)
|
||||
imageAccessSvc := service.NewImageAccessService(cfg.GeneratedRoot, showcaseRepo, authSvc)
|
||||
adobeClient := adobe.NewClient("clio-playground-web", "")
|
||||
chatGPTClient := chatgpt.NewClient("")
|
||||
runwayClient := runway.NewClient("")
|
||||
leonardoClient := leonardo.NewClient("")
|
||||
kreaClient := krea.NewClient("")
|
||||
imagineClient := imagine.NewClient("")
|
||||
v1Svc := service.NewV1Service(cfg, modelRepo, userRepo, eventRepo, tokenRepo, siteRepo, adobeClient, chatGPTClient, runwayClient, leonardoClient, kreaClient, imagineClient, rustfsClient)
|
||||
siteSvc := service.NewSiteService(siteRepo, cfg.AppTitle)
|
||||
showcaseSvc := service.NewShowcaseService(showcaseRepo)
|
||||
adminReadSvc := service.NewAdminReadService(cfg, userRepo, modelRepo, eventRepo, siteRepo, tokenRepo, cdkRepo, rustfsClient)
|
||||
adminWriteSvc := service.NewAdminWriteService(userRepo, showcaseRepo, modelRepo, eventRepo, apiKeyRepo)
|
||||
cdkSvc := service.NewCDKService(cdkRepo, userRepo)
|
||||
apiKeySvc := service.NewAPIKeyService(apiKeyRepo)
|
||||
tokenSvc := service.NewTokenService(tokenRepo, refreshRepo, eventRepo, siteRepo, adobeClient, chatGPTClient, runwayClient, leonardoClient, kreaClient, imagineClient)
|
||||
refreshSvc := service.NewRefreshProfileService(refreshRepo, tokenRepo, adobeClient)
|
||||
// Enable refresh-then-retry on a mid-request Adobe 401 (re-mint access token
|
||||
// from the cookie). Wired post-construction to avoid a ctor init cycle.
|
||||
v1Svc.SetRefresh(refreshSvc)
|
||||
userGenSvc := service.NewUserGenerationService(v1Svc, eventRepo, userRepo, modelRepo)
|
||||
|
||||
engine := router.New(cfg, authSvc, router.Handlers{
|
||||
Health: handler.NewHealthHandler(),
|
||||
Images: handler.NewImageHandler(cfg, imageAccessSvc, rustfsClient),
|
||||
V1: handler.NewV1Handler(v1Svc),
|
||||
Site: handler.NewSiteHandler(siteSvc),
|
||||
Showcase: handler.NewShowcaseHandler(showcaseSvc),
|
||||
Auth: handler.NewAuthHandler(cfg, authSvc, rateLimitSvc),
|
||||
SiteSettings: handler.NewSiteSettingsHandler(siteSvc),
|
||||
AppSettings: handler.NewAppSettingsHandler(appSettingsSvc),
|
||||
AdminRead: handler.NewAdminReadHandler(adminReadSvc),
|
||||
AdminWrite: handler.NewAdminWriteHandler(adminWriteSvc),
|
||||
CDK: handler.NewCDKHandler(cdkSvc),
|
||||
UserTools: handler.NewUserToolsHandler(apiKeySvc, cdkSvc),
|
||||
UserGen: handler.NewUserGenerationHandler(userGenSvc, adminReadSvc),
|
||||
ProviderAdmin: handler.NewProviderAdminHandler(tokenSvc, refreshSvc),
|
||||
})
|
||||
|
||||
// Background self-healing sweep (quota recovery, cookie refresh, stale-pending
|
||||
// cleanup, log retention) — the Go equivalent of the Python daemon thread.
|
||||
maintenanceSvc := service.NewMaintenanceService(tokenRepo, tokenSvc, eventRepo, userRepo, refreshSvc, siteRepo, rustfsClient, v1Svc.Inflight(), showcaseRepo)
|
||||
loopCtx, loopCancel := context.WithCancel(context.Background())
|
||||
go maintenanceSvc.Run(loopCtx)
|
||||
|
||||
return &App{
|
||||
Config: cfg,
|
||||
DB: db,
|
||||
Redis: rdb,
|
||||
Engine: engine,
|
||||
maintenanceCancel: loopCancel,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (a *App) Close() error {
|
||||
if a.maintenanceCancel != nil {
|
||||
a.maintenanceCancel()
|
||||
}
|
||||
if a.Redis != nil {
|
||||
if err := a.Redis.Close(); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
if a.DB != nil {
|
||||
sqlDB, err := a.DB.DB()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return sqlDB.Close()
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,51 @@
|
||||
package bootstrap
|
||||
|
||||
import (
|
||||
"context"
|
||||
|
||||
"backend/internal/model"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
func seedDefaults(ctx context.Context, db *gorm.DB) error {
|
||||
defaults := []model.SiteSetting{
|
||||
{Key: "site.title", Value: "Vivid"},
|
||||
{Key: "contact.qq", Value: "1114639355"},
|
||||
{Key: "contact.qq_link", Value: "https://qm.qq.com/q/ItgCcNA7ac"},
|
||||
{Key: "contact.qq_group", Value: "1106849765"},
|
||||
{Key: "contact.qq_group_link", Value: "https://qm.qq.com/q/976LeMFoHu"},
|
||||
{Key: "contact.email", Value: "vividairun@gmail.com"},
|
||||
{Key: "contact.shop", Value: "https://pay.ldxp.cn/shop/chiyi"},
|
||||
{Key: "auth.open", Value: "true"},
|
||||
{Key: "auth.email_code", Value: "false"},
|
||||
{Key: "auth.allow_password_reset", Value: "false"},
|
||||
{Key: "auth.allowed_email_domains", Value: ""},
|
||||
{Key: "auth.code_ttl_seconds", Value: "600"},
|
||||
{Key: "smtp.host", Value: ""},
|
||||
{Key: "smtp.port", Value: "587"},
|
||||
{Key: "smtp.username", Value: ""},
|
||||
{Key: "smtp.password", Value: ""},
|
||||
{Key: "smtp.from_addr", Value: ""},
|
||||
{Key: "smtp.use_tls", Value: "true"},
|
||||
{Key: "proxy.url", Value: ""},
|
||||
{Key: "credits.checkin_enabled", Value: "true"},
|
||||
{Key: "credits.checkin_reward", Value: "3"},
|
||||
{Key: "credits.invite_enabled", Value: "true"},
|
||||
{Key: "credits.invite_reward", Value: "3"},
|
||||
{Key: "logs.retention_days", Value: "30"},
|
||||
{Key: "media.retention_days", Value: "30"},
|
||||
}
|
||||
for _, item := range defaults {
|
||||
var count int64
|
||||
if err := db.WithContext(ctx).Model(&model.SiteSetting{}).Where("key = ?", item.Key).Count(&count).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if count > 0 {
|
||||
continue
|
||||
}
|
||||
if err := db.WithContext(ctx).Create(&item).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,174 @@
|
||||
package config
|
||||
|
||||
import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
type Config struct {
|
||||
AppEnv string
|
||||
HTTPAddr string
|
||||
AppTitle string
|
||||
PostgresDSN string
|
||||
RedisAddr string
|
||||
RedisPassword string
|
||||
RedisDB int
|
||||
SessionCookieName string
|
||||
CookieSecure bool
|
||||
SessionTTL time.Duration
|
||||
SessionSlideAfter time.Duration
|
||||
CORSOrigins []string
|
||||
GeneratedRoot string
|
||||
RustFSEndpoint string
|
||||
RustFSBucket string
|
||||
RustFSAccessKey string
|
||||
RustFSSecretKey string
|
||||
}
|
||||
|
||||
func Load() (*Config, error) {
|
||||
loadDotEnv()
|
||||
|
||||
wd, err := os.Getwd()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
cfg := &Config{
|
||||
AppEnv: envString("APP_ENV", "development"),
|
||||
HTTPAddr: envString("HTTP_ADDR", ":6061"),
|
||||
AppTitle: envString("APP_TITLE", "Vivid AI"),
|
||||
PostgresDSN: envString("POSTGRES_DSN", "host=127.0.0.1 user=postgres password=postgres dbname=vivid_ai port=5432 sslmode=disable TimeZone=Asia/Shanghai"),
|
||||
RedisAddr: envString("REDIS_ADDR", "127.0.0.1:6379"),
|
||||
RedisPassword: envString("REDIS_PASSWORD", ""),
|
||||
RedisDB: envInt("REDIS_DB", 0),
|
||||
SessionCookieName: envString("SESSION_COOKIE_NAME", "vivid_session"),
|
||||
CookieSecure: envBool("COOKIE_SECURE", false),
|
||||
SessionTTL: time.Duration(envInt("SESSION_TTL_HOURS", 24)) * time.Hour,
|
||||
SessionSlideAfter: time.Duration(envInt("SESSION_SLIDE_AFTER_HOURS", 22)) * time.Hour,
|
||||
CORSOrigins: envList("CORS_ORIGINS", []string{"http://localhost:5173", "http://127.0.0.1:5173"}),
|
||||
GeneratedRoot: filepath.Clean(envString(
|
||||
"GENERATED_ROOT",
|
||||
// vivid-ai's own data dir (backend/data/generated) — NOT the Python
|
||||
// original's tree. Generated outputs and user-uploaded reference
|
||||
// images both live here and are served (cookie-authed) via /images.
|
||||
filepath.Join(wd, "data", "generated"),
|
||||
)),
|
||||
RustFSEndpoint: envString("RUSTFS_ENDPOINT", ""),
|
||||
RustFSBucket: envString("RUSTFS_BUCKET", ""),
|
||||
RustFSAccessKey: envString("RUSTFS_ACCESS_KEY", ""),
|
||||
RustFSSecretKey: envString("RUSTFS_SECRET_KEY", ""),
|
||||
}
|
||||
|
||||
return cfg, nil
|
||||
}
|
||||
|
||||
// loadDotEnv loads a .env file (KEY=VALUE per line) into the process environment
|
||||
// before config is read. Real environment variables always win — .env only fills
|
||||
// keys that aren't already set. Searches ENV_FILE, then walks up from the working
|
||||
// directory so it works whether the binary runs from backend/ or the repo root.
|
||||
func loadDotEnv() {
|
||||
for _, path := range dotEnvCandidates() {
|
||||
data, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
applyDotEnv(string(data))
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
func dotEnvCandidates() []string {
|
||||
var out []string
|
||||
if v := strings.TrimSpace(os.Getenv("ENV_FILE")); v != "" {
|
||||
out = append(out, v)
|
||||
}
|
||||
wd, err := os.Getwd()
|
||||
if err != nil {
|
||||
return out
|
||||
}
|
||||
dir := wd
|
||||
for i := 0; i < 4; i++ {
|
||||
out = append(out, filepath.Join(dir, ".env"))
|
||||
parent := filepath.Dir(dir)
|
||||
if parent == dir {
|
||||
break
|
||||
}
|
||||
dir = parent
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func applyDotEnv(content string) {
|
||||
for _, line := range strings.Split(content, "\n") {
|
||||
line = strings.TrimSpace(strings.TrimPrefix(strings.TrimSpace(line), "export "))
|
||||
if line == "" || strings.HasPrefix(line, "#") {
|
||||
continue
|
||||
}
|
||||
eq := strings.IndexByte(line, '=')
|
||||
if eq < 0 {
|
||||
continue
|
||||
}
|
||||
key := strings.TrimSpace(line[:eq])
|
||||
val := strings.TrimSpace(line[eq+1:])
|
||||
if len(val) >= 2 {
|
||||
if (val[0] == '"' && val[len(val)-1] == '"') || (val[0] == '\'' && val[len(val)-1] == '\'') {
|
||||
val = val[1 : len(val)-1]
|
||||
}
|
||||
}
|
||||
if key == "" {
|
||||
continue
|
||||
}
|
||||
// Real env wins: only set keys that aren't already present.
|
||||
if _, ok := os.LookupEnv(key); !ok {
|
||||
_ = os.Setenv(key, val)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func envString(key, fallback string) string {
|
||||
if v, ok := os.LookupEnv(key); ok && strings.TrimSpace(v) != "" {
|
||||
return strings.TrimSpace(v)
|
||||
}
|
||||
return fallback
|
||||
}
|
||||
|
||||
func envInt(key string, fallback int) int {
|
||||
if v, ok := os.LookupEnv(key); ok && strings.TrimSpace(v) != "" {
|
||||
if n, err := strconv.Atoi(strings.TrimSpace(v)); err == nil {
|
||||
return n
|
||||
}
|
||||
}
|
||||
return fallback
|
||||
}
|
||||
|
||||
func envList(key string, fallback []string) []string {
|
||||
if v, ok := os.LookupEnv(key); ok && strings.TrimSpace(v) != "" {
|
||||
parts := strings.Split(v, ",")
|
||||
out := make([]string, 0, len(parts))
|
||||
for _, part := range parts {
|
||||
s := strings.TrimSpace(part)
|
||||
if s != "" {
|
||||
out = append(out, s)
|
||||
}
|
||||
}
|
||||
if len(out) > 0 {
|
||||
return out
|
||||
}
|
||||
}
|
||||
return fallback
|
||||
}
|
||||
|
||||
func envBool(key string, fallback bool) bool {
|
||||
if v, ok := os.LookupEnv(key); ok && strings.TrimSpace(v) != "" {
|
||||
switch strings.ToLower(strings.TrimSpace(v)) {
|
||||
case "1", "true", "yes", "on":
|
||||
return true
|
||||
case "0", "false", "no", "off":
|
||||
return false
|
||||
}
|
||||
}
|
||||
return fallback
|
||||
}
|
||||
@@ -0,0 +1,234 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"strconv"
|
||||
"time"
|
||||
|
||||
"backend/internal/model"
|
||||
"backend/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
type AdminReadHandler struct {
|
||||
admin *service.AdminReadService
|
||||
}
|
||||
|
||||
func NewAdminReadHandler(admin *service.AdminReadService) *AdminReadHandler {
|
||||
return &AdminReadHandler{admin: admin}
|
||||
}
|
||||
|
||||
func (h *AdminReadHandler) Users(c *gin.Context) {
|
||||
users, stats, err := h.admin.Users(c.Request.Context())
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"detail": "failed to load users"})
|
||||
return
|
||||
}
|
||||
|
||||
out := make([]gin.H, 0, len(users))
|
||||
generationCounts := map[string]int64{}
|
||||
if raw, ok := stats["generation_counts"].(map[string]int64); ok {
|
||||
generationCounts = raw
|
||||
}
|
||||
for _, user := range users {
|
||||
row := userPublic(user)
|
||||
row["generation_count"] = generationCounts[user.ID]
|
||||
out = append(out, row)
|
||||
}
|
||||
delete(stats, "generation_counts")
|
||||
c.JSON(http.StatusOK, gin.H{"data": out, "stats": stats})
|
||||
}
|
||||
|
||||
func (h *AdminReadHandler) Models(c *gin.Context) {
|
||||
items, err := h.admin.ModelsView(c.Request.Context())
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"detail": "failed to load models"})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"data": items})
|
||||
}
|
||||
|
||||
func (h *AdminReadHandler) Logs(c *gin.Context) {
|
||||
limit := parseInt(c.Query("limit"), 50)
|
||||
offset := parseInt(c.Query("offset"), 0)
|
||||
kind := c.Query("kind")
|
||||
status := c.Query("status")
|
||||
var since *time.Time
|
||||
if raw := c.Query("since"); raw != "" {
|
||||
if f, err := strconv.ParseFloat(raw, 64); err == nil {
|
||||
t := time.Unix(int64(f), 0)
|
||||
since = &t
|
||||
}
|
||||
}
|
||||
|
||||
items, total, stats, err := h.admin.Logs(c.Request.Context(), limit, offset, kind, status, since, "", "", c.Query("source"), false)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"detail": "failed to load logs"})
|
||||
return
|
||||
}
|
||||
// Resolve user_id -> display name once for the page (mirrors admin.py).
|
||||
nameByID, err := h.admin.UserNameMap(c.Request.Context())
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"detail": "failed to load logs"})
|
||||
return
|
||||
}
|
||||
out := make([]gin.H, 0, len(items))
|
||||
for _, item := range items {
|
||||
var userName any
|
||||
if item.UserID == "" {
|
||||
userName = "匿名"
|
||||
} else if name, ok := nameByID[item.UserID]; ok {
|
||||
userName = name
|
||||
} else {
|
||||
userName = item.UserID
|
||||
}
|
||||
out = append(out, gin.H{
|
||||
"id": item.ID,
|
||||
"ts": item.TS.Unix(),
|
||||
"kind": item.Kind,
|
||||
"status": item.Status,
|
||||
"model": item.Model,
|
||||
"provider": item.Provider,
|
||||
"prompt": item.Prompt,
|
||||
"ratio": item.Ratio,
|
||||
"resolution": item.Resolution,
|
||||
"duration": item.Duration,
|
||||
"refs": item.Refs,
|
||||
"source": item.Source,
|
||||
"user_id": emptyStringNil(item.UserID),
|
||||
"user_name": userName,
|
||||
"cost": item.Cost,
|
||||
"elapsed_ms": item.ElapsedMS,
|
||||
"file": emptyStringNil(item.File),
|
||||
"error": emptyStringNil(item.Error),
|
||||
"created_at": unixSec(item.CreatedAt),
|
||||
"updated_at": unixSec(item.UpdatedAt),
|
||||
})
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"data": out,
|
||||
"total": total,
|
||||
"limit": limit,
|
||||
"offset": offset,
|
||||
"stats": stats,
|
||||
})
|
||||
}
|
||||
|
||||
func (h *AdminReadHandler) Stats(c *gin.Context) {
|
||||
stats, err := h.admin.Stats(c.Request.Context())
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"detail": "failed to load stats"})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, stats)
|
||||
}
|
||||
|
||||
func (h *AdminReadHandler) Dashboard(c *gin.Context) {
|
||||
data, err := h.admin.Dashboard(c.Request.Context())
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"detail": "failed to load dashboard"})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, data)
|
||||
}
|
||||
|
||||
func (h *AdminReadHandler) Invites(c *gin.Context) {
|
||||
items, stats, err := h.admin.Invites(c.Request.Context())
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"detail": "failed to load invites"})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"data": items, "stats": stats})
|
||||
}
|
||||
|
||||
func (h *AdminReadHandler) Providers(c *gin.Context) {
|
||||
items, err := h.admin.Providers(c.Request.Context())
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"detail": "failed to load providers"})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"data": items})
|
||||
}
|
||||
|
||||
func (h *AdminReadHandler) Images(c *gin.Context) {
|
||||
limit := parseInt(c.Query("limit"), 30)
|
||||
offset := parseInt(c.Query("offset"), 0)
|
||||
kind := c.Query("kind")
|
||||
items, total, stats, err := h.admin.Images(c.Request.Context(), limit, offset, kind)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"detail": "failed to load images"})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"data": items,
|
||||
"total": total,
|
||||
"limit": limit,
|
||||
"offset": offset,
|
||||
"stats": stats,
|
||||
})
|
||||
}
|
||||
|
||||
func userPublic(user model.User) gin.H {
|
||||
keys := make([]gin.H, 0, len(user.APIKeys))
|
||||
for _, key := range user.APIKeys {
|
||||
keys = append(keys, gin.H{
|
||||
"id": key.ID,
|
||||
"name": key.Name,
|
||||
"key_preview": key.KeyPreview,
|
||||
"created_at": unixSec(key.CreatedAt),
|
||||
"last_used_at": unixSecPtr(key.LastUsedAt),
|
||||
})
|
||||
}
|
||||
return gin.H{
|
||||
"id": user.ID,
|
||||
"email": user.Email,
|
||||
"name": user.Name,
|
||||
"role": user.Role,
|
||||
"status": user.Status,
|
||||
"credits": user.Credits,
|
||||
"notes": user.Notes,
|
||||
"created_at": unixSec(user.CreatedAt),
|
||||
"last_login_at": unixSecPtr(user.LastLoginAt),
|
||||
"last_login_ip": user.LastLoginIP,
|
||||
"invite_code": user.InviteCode,
|
||||
"invited_by": user.InvitedBy,
|
||||
"checkin_last": user.CheckinLast,
|
||||
"checkin_streak": user.CheckinStreak,
|
||||
"api_keys": keys,
|
||||
"has_password": user.PasswordHash != "",
|
||||
}
|
||||
}
|
||||
|
||||
// unixSec / unixSecPtr render timestamps as unix SECONDS — the frontend's
|
||||
// fmtTs/fmtRelative expect seconds (matching the Python reference's time.time()),
|
||||
// not the RFC3339 string Go marshals a time.Time into (which parses to NaN → "—").
|
||||
func unixSec(t time.Time) any {
|
||||
if t.IsZero() {
|
||||
return nil
|
||||
}
|
||||
return t.Unix()
|
||||
}
|
||||
|
||||
func unixSecPtr(t *time.Time) any {
|
||||
if t == nil || t.IsZero() {
|
||||
return nil
|
||||
}
|
||||
return t.Unix()
|
||||
}
|
||||
|
||||
func parseInt(raw string, fallback int) int {
|
||||
if raw == "" {
|
||||
return fallback
|
||||
}
|
||||
if n, err := strconv.Atoi(raw); err == nil {
|
||||
return n
|
||||
}
|
||||
return fallback
|
||||
}
|
||||
|
||||
func emptyStringNil(v string) any {
|
||||
if v == "" {
|
||||
return nil
|
||||
}
|
||||
return v
|
||||
}
|
||||
@@ -0,0 +1,241 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"net/http"
|
||||
|
||||
"backend/internal/model"
|
||||
"backend/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
type AdminWriteHandler struct {
|
||||
admin *service.AdminWriteService
|
||||
}
|
||||
|
||||
func NewAdminWriteHandler(admin *service.AdminWriteService) *AdminWriteHandler {
|
||||
return &AdminWriteHandler{admin: admin}
|
||||
}
|
||||
|
||||
func (h *AdminWriteHandler) CreateUser(c *gin.Context) {
|
||||
var body map[string]any
|
||||
if err := c.ShouldBindJSON(&body); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": "invalid request body"})
|
||||
return
|
||||
}
|
||||
user, err := h.admin.CreateUser(c.Request.Context(), body)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"ok": true, "data": userPublic(*user)})
|
||||
}
|
||||
|
||||
func (h *AdminWriteHandler) UpdateUser(c *gin.Context) {
|
||||
var body map[string]any
|
||||
if err := c.ShouldBindJSON(&body); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": "invalid request body"})
|
||||
return
|
||||
}
|
||||
user, err := h.admin.UpdateUser(c.Request.Context(), c.Param("user_id"), body)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"ok": true, "data": userPublic(*user)})
|
||||
}
|
||||
|
||||
func (h *AdminWriteHandler) DeleteUser(c *gin.Context) {
|
||||
if err := h.admin.DeleteUser(c.Request.Context(), c.Param("user_id")); err != nil {
|
||||
if errors.Is(err, service.ErrNotFound) {
|
||||
c.JSON(http.StatusNotFound, gin.H{"detail": "user not found"})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"detail": "failed to delete user"})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"ok": true})
|
||||
}
|
||||
|
||||
// DeleteUsersBulk removes multiple users in one call (multi-select).
|
||||
func (h *AdminWriteHandler) DeleteUsersBulk(c *gin.Context) {
|
||||
var body struct {
|
||||
IDs []string `json:"ids"`
|
||||
}
|
||||
if err := c.ShouldBindJSON(&body); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": "invalid request body"})
|
||||
return
|
||||
}
|
||||
if len(body.IDs) == 0 {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": "未选择任何用户"})
|
||||
return
|
||||
}
|
||||
n, err := h.admin.DeleteUsers(c.Request.Context(), body.IDs)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"detail": "failed to delete users"})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"ok": true, "deleted": n})
|
||||
}
|
||||
|
||||
func (h *AdminWriteHandler) AdjustUserCredits(c *gin.Context) {
|
||||
var body struct {
|
||||
Delta float64 `json:"delta"`
|
||||
Set *float64 `json:"set"`
|
||||
}
|
||||
if err := c.ShouldBindJSON(&body); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": "invalid request body"})
|
||||
return
|
||||
}
|
||||
var (
|
||||
user *model.User
|
||||
err error
|
||||
)
|
||||
if body.Set != nil {
|
||||
// Absolute set takes precedence over delta (matches Python admin.py).
|
||||
user, err = h.admin.SetUserCredits(c.Request.Context(), c.Param("user_id"), *body.Set)
|
||||
} else {
|
||||
user, err = h.admin.AdjustUserCredits(c.Request.Context(), c.Param("user_id"), body.Delta)
|
||||
}
|
||||
if err != nil {
|
||||
if errors.Is(err, service.ErrNotFound) {
|
||||
c.JSON(http.StatusNotFound, gin.H{"detail": "user not found"})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"ok": true, "data": userPublic(*user)})
|
||||
}
|
||||
|
||||
func (h *AdminWriteHandler) CreateUserAPIKey(c *gin.Context) {
|
||||
var body struct {
|
||||
Name string `json:"name"`
|
||||
}
|
||||
if err := c.ShouldBindJSON(&body); err != nil && err.Error() != "EOF" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": "invalid request body"})
|
||||
return
|
||||
}
|
||||
key, plain, err := h.admin.CreateUserAPIKey(c.Request.Context(), c.Param("user_id"), body.Name)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"ok": true,
|
||||
"key": plain,
|
||||
"data": gin.H{
|
||||
"id": key.ID,
|
||||
"name": key.Name,
|
||||
"key_preview": key.KeyPreview,
|
||||
"created_at": key.CreatedAt,
|
||||
"last_used_at": key.LastUsedAt,
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
func (h *AdminWriteHandler) DeleteUserAPIKey(c *gin.Context) {
|
||||
if err := h.admin.DeleteUserAPIKey(c.Request.Context(), c.Param("user_id"), c.Param("key_id")); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"ok": true})
|
||||
}
|
||||
|
||||
func (h *AdminWriteHandler) CreateShowcase(c *gin.Context) {
|
||||
var body map[string]any
|
||||
if err := c.ShouldBindJSON(&body); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": "invalid request body"})
|
||||
return
|
||||
}
|
||||
item, err := h.admin.CreateShowcase(c.Request.Context(), body)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"ok": true, "data": item})
|
||||
}
|
||||
|
||||
func (h *AdminWriteHandler) UpdateShowcase(c *gin.Context) {
|
||||
var body map[string]any
|
||||
if err := c.ShouldBindJSON(&body); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": "invalid request body"})
|
||||
return
|
||||
}
|
||||
item, err := h.admin.UpdateShowcase(c.Request.Context(), c.Param("entry_id"), body)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"ok": true, "data": item})
|
||||
}
|
||||
|
||||
func (h *AdminWriteHandler) DeleteShowcase(c *gin.Context) {
|
||||
if err := h.admin.DeleteShowcase(c.Request.Context(), c.Param("entry_id")); err != nil {
|
||||
if errors.Is(err, service.ErrNotFound) {
|
||||
c.JSON(http.StatusNotFound, gin.H{"detail": "showcase not found"})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"detail": "failed to delete showcase"})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"ok": true})
|
||||
}
|
||||
|
||||
func (h *AdminWriteHandler) CreateModel(c *gin.Context) {
|
||||
var body map[string]any
|
||||
if err := c.ShouldBindJSON(&body); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": "invalid request body"})
|
||||
return
|
||||
}
|
||||
item, err := h.admin.CreateModel(c.Request.Context(), body)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"ok": true, "data": item})
|
||||
}
|
||||
|
||||
func (h *AdminWriteHandler) UpdateModel(c *gin.Context) {
|
||||
var body map[string]any
|
||||
if err := c.ShouldBindJSON(&body); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": "invalid request body"})
|
||||
return
|
||||
}
|
||||
item, err := h.admin.UpdateModel(c.Request.Context(), c.Param("model_id"), body)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"ok": true, "data": item})
|
||||
}
|
||||
|
||||
func (h *AdminWriteHandler) DeleteModel(c *gin.Context) {
|
||||
if err := h.admin.DeleteModel(c.Request.Context(), c.Param("model_id")); err != nil {
|
||||
if errors.Is(err, service.ErrNotFound) {
|
||||
c.JSON(http.StatusNotFound, gin.H{"detail": "model not found"})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"detail": "failed to delete model"})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"ok": true})
|
||||
}
|
||||
|
||||
func (h *AdminWriteHandler) ClearLogs(c *gin.Context) {
|
||||
removed, err := h.admin.ClearLogs(c.Request.Context())
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"detail": "failed to clear logs"})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"ok": true, "removed": removed})
|
||||
}
|
||||
|
||||
func (h *AdminWriteHandler) ClearPendingLogs(c *gin.Context) {
|
||||
removed, err := h.admin.ClearPendingLogs(c.Request.Context())
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"detail": "failed to clear pending logs"})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"ok": true, "removed": removed})
|
||||
}
|
||||
@@ -0,0 +1,191 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
|
||||
"backend/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
type AppSettingsHandler struct {
|
||||
settings *service.AppSettingsService
|
||||
}
|
||||
|
||||
func NewAppSettingsHandler(settings *service.AppSettingsService) *AppSettingsHandler {
|
||||
return &AppSettingsHandler{settings: settings}
|
||||
}
|
||||
|
||||
func (h *AppSettingsHandler) RegistrationGet(c *gin.Context) {
|
||||
data, err := h.settings.Registration(c.Request.Context())
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"detail": "failed to load registration settings"})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, data)
|
||||
}
|
||||
|
||||
func (h *AppSettingsHandler) RegistrationPut(c *gin.Context) {
|
||||
var body service.RegistrationSettings
|
||||
if err := c.ShouldBindJSON(&body); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": "invalid request body"})
|
||||
return
|
||||
}
|
||||
data, err := h.settings.SaveRegistration(c.Request.Context(), body)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"ok": true, "data": data})
|
||||
}
|
||||
|
||||
func (h *AppSettingsHandler) SMTPGet(c *gin.Context) {
|
||||
data, err := h.settings.SMTP(c.Request.Context())
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"detail": "failed to load smtp settings"})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, data)
|
||||
}
|
||||
|
||||
func (h *AppSettingsHandler) SMTPPut(c *gin.Context) {
|
||||
var body service.SMTPSettings
|
||||
if err := c.ShouldBindJSON(&body); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": "invalid request body"})
|
||||
return
|
||||
}
|
||||
data, err := h.settings.SaveSMTP(c.Request.Context(), body)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"ok": true, "data": data})
|
||||
}
|
||||
|
||||
func (h *AppSettingsHandler) SMTPTest(c *gin.Context) {
|
||||
var body struct {
|
||||
Email string `json:"email"`
|
||||
}
|
||||
if err := c.ShouldBindJSON(&body); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": "invalid request body"})
|
||||
return
|
||||
}
|
||||
if err := h.settings.TestSMTP(c.Request.Context(), body.Email); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"ok": true, "detail": "测试邮件已发送"})
|
||||
}
|
||||
|
||||
func (h *AppSettingsHandler) ProxyGet(c *gin.Context) {
|
||||
data, err := h.settings.Proxy(c.Request.Context())
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"detail": "failed to load proxy settings"})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, data)
|
||||
}
|
||||
|
||||
func (h *AppSettingsHandler) ProxyPut(c *gin.Context) {
|
||||
var body struct {
|
||||
Proxy string `json:"proxy"`
|
||||
}
|
||||
if err := c.ShouldBindJSON(&body); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": "invalid request body"})
|
||||
return
|
||||
}
|
||||
data, err := h.settings.SaveProxy(c.Request.Context(), body.Proxy)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"ok": true, "data": data})
|
||||
}
|
||||
|
||||
func (h *AppSettingsHandler) ProxyTest(c *gin.Context) {
|
||||
var body struct {
|
||||
Proxy string `json:"proxy"`
|
||||
}
|
||||
if err := c.ShouldBindJSON(&body); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": "invalid request body"})
|
||||
return
|
||||
}
|
||||
data, err := h.settings.TestProxy(c.Request.Context(), body.Proxy)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"ok": true, "data": data})
|
||||
}
|
||||
|
||||
func (h *AppSettingsHandler) CreditsGet(c *gin.Context) {
|
||||
data, err := h.settings.Credits(c.Request.Context())
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"detail": "failed to load credit settings"})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, data)
|
||||
}
|
||||
|
||||
func (h *AppSettingsHandler) CreditsPut(c *gin.Context) {
|
||||
var body service.CreditSettings
|
||||
if err := c.ShouldBindJSON(&body); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": "invalid request body"})
|
||||
return
|
||||
}
|
||||
data, err := h.settings.SaveCredits(c.Request.Context(), body)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"ok": true, "data": data})
|
||||
}
|
||||
|
||||
func (h *AppSettingsHandler) LogsGet(c *gin.Context) {
|
||||
data, err := h.settings.Logs(c.Request.Context())
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"detail": "failed to load log settings"})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, data)
|
||||
}
|
||||
|
||||
func (h *AppSettingsHandler) LogsPut(c *gin.Context) {
|
||||
var body struct {
|
||||
RetentionDays int `json:"retention_days"`
|
||||
}
|
||||
if err := c.ShouldBindJSON(&body); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": "invalid request body"})
|
||||
return
|
||||
}
|
||||
data, err := h.settings.SaveLogs(c.Request.Context(), body.RetentionDays)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"ok": true, "data": data})
|
||||
}
|
||||
|
||||
func (h *AppSettingsHandler) MediaGet(c *gin.Context) {
|
||||
data, err := h.settings.Media(c.Request.Context())
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"detail": "failed to load media settings"})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, data)
|
||||
}
|
||||
|
||||
func (h *AppSettingsHandler) MediaPut(c *gin.Context) {
|
||||
var body struct {
|
||||
RetentionDays int `json:"retention_days"`
|
||||
}
|
||||
if err := c.ShouldBindJSON(&body); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": "invalid request body"})
|
||||
return
|
||||
}
|
||||
data, err := h.settings.SaveMedia(c.Request.Context(), body.RetentionDays)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"ok": true, "data": data.Settings, "removed": data.Removed, "freed_bytes": data.FreedBytes})
|
||||
}
|
||||
@@ -0,0 +1,350 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"net/http"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"backend/internal/config"
|
||||
"backend/internal/model"
|
||||
"backend/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
type AuthHandler struct {
|
||||
cfg *config.Config
|
||||
auth *service.AuthService
|
||||
limiter *service.RateLimitService
|
||||
}
|
||||
|
||||
func NewAuthHandler(cfg *config.Config, auth *service.AuthService, limiter *service.RateLimitService) *AuthHandler {
|
||||
return &AuthHandler{
|
||||
cfg: cfg,
|
||||
auth: auth,
|
||||
limiter: limiter,
|
||||
}
|
||||
}
|
||||
|
||||
func (h *AuthHandler) Config(c *gin.Context) {
|
||||
data, err := h.auth.AuthConfig(c.Request.Context())
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"detail": "failed to load auth config"})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, data)
|
||||
}
|
||||
|
||||
func (h *AuthHandler) SendCode(c *gin.Context) {
|
||||
var body struct {
|
||||
Email string `json:"email"`
|
||||
Purpose string `json:"purpose"`
|
||||
}
|
||||
if err := c.ShouldBindJSON(&body); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": "invalid request body"})
|
||||
return
|
||||
}
|
||||
ip := clientIP(c)
|
||||
if err := h.enforceRateLimit(c, "auth:send-code:ip:"+ip, 5, time.Hour); err != nil {
|
||||
return
|
||||
}
|
||||
if email, err := service.ValidateEmail(body.Email); err == nil {
|
||||
if err := h.enforceRateLimit(c, "auth:send-code:email:"+email, 3, 10*time.Minute); err != nil {
|
||||
return
|
||||
}
|
||||
}
|
||||
if err := h.auth.SendCode(c.Request.Context(), body.Email, body.Purpose); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"ok": true})
|
||||
}
|
||||
|
||||
func (h *AuthHandler) Register(c *gin.Context) {
|
||||
var body struct {
|
||||
Email string `json:"email"`
|
||||
Username string `json:"username"`
|
||||
Name string `json:"name"`
|
||||
Password string `json:"password"`
|
||||
InviteCode string `json:"invite_code"`
|
||||
EmailCode string `json:"email_code"`
|
||||
Code string `json:"code"`
|
||||
}
|
||||
if err := c.ShouldBindJSON(&body); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": "invalid request body"})
|
||||
return
|
||||
}
|
||||
username := strings.TrimSpace(body.Username)
|
||||
if username == "" {
|
||||
username = strings.TrimSpace(body.Name)
|
||||
}
|
||||
if err := h.enforceRateLimit(c, "auth:register:ip:"+clientIP(c), 10, time.Hour); err != nil {
|
||||
return
|
||||
}
|
||||
emailCode := strings.TrimSpace(body.EmailCode)
|
||||
if emailCode == "" {
|
||||
emailCode = strings.TrimSpace(body.Code)
|
||||
}
|
||||
user, token, session, err := h.auth.Register(
|
||||
c.Request.Context(),
|
||||
body.Email,
|
||||
username,
|
||||
body.Password,
|
||||
body.InviteCode,
|
||||
emailCode,
|
||||
clientIP(c),
|
||||
)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": err.Error()})
|
||||
return
|
||||
}
|
||||
h.writeSession(c, token, session, user)
|
||||
}
|
||||
|
||||
func (h *AuthHandler) Login(c *gin.Context) {
|
||||
var body struct {
|
||||
Identifier string `json:"identifier"`
|
||||
Email string `json:"email"`
|
||||
Username string `json:"username"`
|
||||
Password string `json:"password"`
|
||||
}
|
||||
if err := c.ShouldBindJSON(&body); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": "invalid request body"})
|
||||
return
|
||||
}
|
||||
|
||||
identifier := strings.TrimSpace(body.Identifier)
|
||||
if identifier == "" {
|
||||
if strings.TrimSpace(body.Email) != "" {
|
||||
identifier = strings.TrimSpace(body.Email)
|
||||
} else {
|
||||
identifier = strings.TrimSpace(body.Username)
|
||||
}
|
||||
}
|
||||
if identifier == "" || body.Password == "" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": "账号或密码不能为空"})
|
||||
return
|
||||
}
|
||||
ip := clientIP(c)
|
||||
if err := h.enforceRateLimit(c, "auth:login:ip:"+ip, 20, 15*time.Minute); err != nil {
|
||||
return
|
||||
}
|
||||
if normalized, err := service.ValidateLoginIdentifier(identifier); err == nil {
|
||||
if err := h.enforceRateLimit(c, "auth:login:target:"+ip+":"+strings.ToLower(normalized), 8, 15*time.Minute); err != nil {
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
user, token, session, err := h.auth.Login(c.Request.Context(), identifier, body.Password, ip)
|
||||
if err != nil {
|
||||
if writeLoginLocked(c, err) {
|
||||
return
|
||||
}
|
||||
if err == service.ErrAuthFailed {
|
||||
c.JSON(http.StatusUnauthorized, gin.H{"detail": "账号或密码错误"})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
h.writeSession(c, token, session, user)
|
||||
}
|
||||
|
||||
func (h *AuthHandler) ResetPassword(c *gin.Context) {
|
||||
var body struct {
|
||||
Email string `json:"email"`
|
||||
Password string `json:"password"`
|
||||
EmailCode string `json:"email_code"`
|
||||
Code string `json:"code"`
|
||||
}
|
||||
if err := c.ShouldBindJSON(&body); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": "invalid request body"})
|
||||
return
|
||||
}
|
||||
ip := clientIP(c)
|
||||
if err := h.enforceRateLimit(c, "auth:reset:ip:"+ip, 5, time.Hour); err != nil {
|
||||
return
|
||||
}
|
||||
if email, err := service.ValidateEmail(body.Email); err == nil {
|
||||
if err := h.enforceRateLimit(c, "auth:reset:email:"+email, 5, time.Hour); err != nil {
|
||||
return
|
||||
}
|
||||
}
|
||||
emailCode := strings.TrimSpace(body.EmailCode)
|
||||
if emailCode == "" {
|
||||
emailCode = strings.TrimSpace(body.Code)
|
||||
}
|
||||
if err := h.auth.ResetPassword(c.Request.Context(), body.Email, body.Password, emailCode, ip); err != nil {
|
||||
if writeLoginLocked(c, err) {
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"ok": true})
|
||||
}
|
||||
|
||||
func (h *AuthHandler) ChangePassword(c *gin.Context) {
|
||||
user := currentUser(c)
|
||||
if user == nil {
|
||||
c.JSON(http.StatusUnauthorized, gin.H{"detail": "未登录或会话已过期"})
|
||||
return
|
||||
}
|
||||
if err := h.enforceRateLimit(c, "auth:change-password:user:"+user.ID, 10, 30*time.Minute); err != nil {
|
||||
return
|
||||
}
|
||||
var body struct {
|
||||
CurrentPassword string `json:"current_password"`
|
||||
Current string `json:"current"`
|
||||
NewPassword string `json:"new_password"`
|
||||
Password string `json:"password"`
|
||||
}
|
||||
if err := c.ShouldBindJSON(&body); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": "invalid request body"})
|
||||
return
|
||||
}
|
||||
current := strings.TrimSpace(body.CurrentPassword)
|
||||
if current == "" {
|
||||
current = strings.TrimSpace(body.Current)
|
||||
}
|
||||
next := strings.TrimSpace(body.NewPassword)
|
||||
if next == "" {
|
||||
next = body.Password
|
||||
}
|
||||
if err := h.auth.ChangePassword(c.Request.Context(), user.ID, current, next); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"ok": true})
|
||||
}
|
||||
|
||||
func (h *AuthHandler) Checkin(c *gin.Context) {
|
||||
user := currentUser(c)
|
||||
if user == nil {
|
||||
c.JSON(http.StatusUnauthorized, gin.H{"detail": "未登录或会话已过期"})
|
||||
return
|
||||
}
|
||||
result, err := h.auth.Checkin(c.Request.Context(), user.ID)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"ok": true,
|
||||
"already": result.Already,
|
||||
"awarded": result.Awarded,
|
||||
"streak": result.Streak,
|
||||
"credits": result.Credits,
|
||||
})
|
||||
}
|
||||
|
||||
func (h *AuthHandler) Invites(c *gin.Context) {
|
||||
user := currentUser(c)
|
||||
if user == nil {
|
||||
c.JSON(http.StatusUnauthorized, gin.H{"detail": "未登录或会话已过期"})
|
||||
return
|
||||
}
|
||||
items, err := h.auth.InviteList(c.Request.Context(), user.ID)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"detail": "failed to load invites"})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"data": items, "reward": h.auth.InviteReward(c.Request.Context())})
|
||||
}
|
||||
|
||||
func (h *AuthHandler) Logout(c *gin.Context) {
|
||||
token := service.ParseBearer(c.GetHeader("Authorization"))
|
||||
if token == "" {
|
||||
token = readCookie(c, h.cfg.SessionCookieName)
|
||||
}
|
||||
_ = h.auth.Logout(c.Request.Context(), token)
|
||||
c.SetCookie(h.cfg.SessionCookieName, "", -1, "/", "", h.cfg.CookieSecure, true)
|
||||
c.JSON(http.StatusOK, gin.H{"ok": true})
|
||||
}
|
||||
|
||||
func (h *AuthHandler) Me(c *gin.Context) {
|
||||
userValue, ok := c.Get("current_user")
|
||||
if !ok {
|
||||
c.JSON(http.StatusUnauthorized, gin.H{"detail": "未登录或会话已过期"})
|
||||
return
|
||||
}
|
||||
sessionValue, ok := c.Get("current_session")
|
||||
if !ok {
|
||||
c.JSON(http.StatusUnauthorized, gin.H{"detail": "未登录或会话已过期"})
|
||||
return
|
||||
}
|
||||
|
||||
user, _ := userValue.(*model.User)
|
||||
session, _ := sessionValue.(*service.SessionPayload)
|
||||
if user == nil || session == nil {
|
||||
c.JSON(http.StatusUnauthorized, gin.H{"detail": "账号或密码错误"})
|
||||
return
|
||||
}
|
||||
publicUser, err := h.auth.PublicUser(c.Request.Context(), user)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"detail": "failed to load user profile"})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"ok": true,
|
||||
"expires_at": session.ExpiresAt,
|
||||
"user": publicUser,
|
||||
})
|
||||
}
|
||||
|
||||
// writeLoginLocked maps a LoginGuard lockout error to HTTP 429 with a
|
||||
// Retry-After header (mirrors Python api/auth.py:226-237). Returns true when it
|
||||
// handled the error so the caller stops processing.
|
||||
func writeLoginLocked(c *gin.Context, err error) bool {
|
||||
var locked *service.LoginLockedError
|
||||
if errors.As(err, &locked) {
|
||||
c.Header("Retry-After", strconv.Itoa(locked.RetryAfter))
|
||||
c.JSON(http.StatusTooManyRequests, gin.H{"detail": locked.Error()})
|
||||
return true
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func clientIP(c *gin.Context) string {
|
||||
if fwd := strings.TrimSpace(c.GetHeader("X-Forwarded-For")); fwd != "" {
|
||||
parts := strings.Split(fwd, ",")
|
||||
return strings.TrimSpace(parts[0])
|
||||
}
|
||||
if real := strings.TrimSpace(c.GetHeader("X-Real-Ip")); real != "" {
|
||||
return real
|
||||
}
|
||||
return c.ClientIP()
|
||||
}
|
||||
|
||||
func (h *AuthHandler) enforceRateLimit(c *gin.Context, bucket string, limit int64, window time.Duration) error {
|
||||
if h.limiter == nil {
|
||||
return nil
|
||||
}
|
||||
if err := h.limiter.Enforce(c.Request.Context(), bucket, limit, window); err != nil {
|
||||
if errors.Is(err, service.ErrRateLimited) {
|
||||
c.JSON(http.StatusTooManyRequests, gin.H{"detail": err.Error()})
|
||||
return err
|
||||
}
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"detail": "rate limiter unavailable"})
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (h *AuthHandler) writeSession(c *gin.Context, token string, session *service.SessionPayload, user *model.User) {
|
||||
c.SetSameSite(http.SameSiteLaxMode)
|
||||
c.SetCookie(h.cfg.SessionCookieName, token, int(h.cfg.SessionTTL.Seconds()), "/", "", h.cfg.CookieSecure, true)
|
||||
publicUser, err := h.auth.PublicUser(c.Request.Context(), user)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"detail": "failed to load user profile"})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"ok": true,
|
||||
"token": token,
|
||||
"expires_at": session.ExpiresAt,
|
||||
"user": publicUser,
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,104 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"net/http"
|
||||
|
||||
"backend/internal/model"
|
||||
"backend/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
type CDKHandler struct {
|
||||
cdks *service.CDKService
|
||||
}
|
||||
|
||||
func NewCDKHandler(cdks *service.CDKService) *CDKHandler {
|
||||
return &CDKHandler{cdks: cdks}
|
||||
}
|
||||
|
||||
func (h *CDKHandler) List(c *gin.Context) {
|
||||
items, stats, names, err := h.cdks.List(c.Request.Context())
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"detail": "failed to load cdks"})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"data": cdkPublic(items, names), "stats": stats})
|
||||
}
|
||||
|
||||
func (h *CDKHandler) Create(c *gin.Context) {
|
||||
var body struct {
|
||||
Amount int `json:"amount"`
|
||||
Count int `json:"count"`
|
||||
Note string `json:"note"`
|
||||
Type string `json:"type"`
|
||||
}
|
||||
if err := c.ShouldBindJSON(&body); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": "invalid request body"})
|
||||
return
|
||||
}
|
||||
items, err := h.cdks.Generate(c.Request.Context(), body.Amount, body.Count, body.Note, body.Type)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"ok": true, "created": cdkPublic(items, nil)})
|
||||
}
|
||||
|
||||
func (h *CDKHandler) Delete(c *gin.Context) {
|
||||
if err := h.cdks.Delete(c.Request.Context(), c.Param("code")); err != nil {
|
||||
if errors.Is(err, service.ErrNotFound) {
|
||||
c.JSON(http.StatusNotFound, gin.H{"detail": "cdk not found"})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"detail": "failed to delete cdk"})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"ok": true})
|
||||
}
|
||||
|
||||
// DeleteBulk removes multiple CDK codes in one call (multi-select).
|
||||
func (h *CDKHandler) DeleteBulk(c *gin.Context) {
|
||||
var body struct {
|
||||
Codes []string `json:"codes"`
|
||||
}
|
||||
if err := c.ShouldBindJSON(&body); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": "invalid request body"})
|
||||
return
|
||||
}
|
||||
if len(body.Codes) == 0 {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": "未选择任何兑换码"})
|
||||
return
|
||||
}
|
||||
n, err := h.cdks.DeleteBulk(c.Request.Context(), body.Codes)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"detail": "failed to delete cdks"})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"ok": true, "deleted": n})
|
||||
}
|
||||
|
||||
func cdkPublic(items []model.CDKCode, nameByID map[string]string) []gin.H {
|
||||
out := make([]gin.H, 0, len(items))
|
||||
for _, item := range items {
|
||||
var redeemedByName any
|
||||
if item.RedeemedBy != nil && *item.RedeemedBy != "" {
|
||||
if name, ok := nameByID[*item.RedeemedBy]; ok {
|
||||
redeemedByName = name
|
||||
}
|
||||
}
|
||||
out = append(out, gin.H{
|
||||
"code": item.Code,
|
||||
"amount": item.Amount,
|
||||
"status": item.Status,
|
||||
"type": item.Type,
|
||||
"batch_id": item.BatchID,
|
||||
"note": item.Note,
|
||||
"redeemed_by": item.RedeemedBy,
|
||||
"redeemed_by_name": redeemedByName,
|
||||
"redeemed_at": unixSecPtr(item.RedeemedAt),
|
||||
"created_at": unixSec(item.CreatedAt),
|
||||
})
|
||||
}
|
||||
return out
|
||||
}
|
||||
@@ -0,0 +1,13 @@
|
||||
package handler
|
||||
|
||||
import "github.com/gin-gonic/gin"
|
||||
|
||||
type HealthHandler struct{}
|
||||
|
||||
func NewHealthHandler() *HealthHandler {
|
||||
return &HealthHandler{}
|
||||
}
|
||||
|
||||
func (h *HealthHandler) Handle(c *gin.Context) {
|
||||
c.JSON(200, gin.H{"ok": true})
|
||||
}
|
||||
@@ -0,0 +1,88 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"io"
|
||||
"net/http"
|
||||
|
||||
"backend/internal/config"
|
||||
"backend/internal/service"
|
||||
"backend/internal/storage"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
type ImageHandler struct {
|
||||
cfg *config.Config
|
||||
imageAccess *service.ImageAccessService
|
||||
store *storage.Client
|
||||
}
|
||||
|
||||
func NewImageHandler(cfg *config.Config, imageAccess *service.ImageAccessService, store *storage.Client) *ImageHandler {
|
||||
return &ImageHandler{
|
||||
cfg: cfg,
|
||||
imageAccess: imageAccess,
|
||||
store: store,
|
||||
}
|
||||
}
|
||||
|
||||
// Serve gates access (public showcase images, or a logged-in cookie — a regular
|
||||
// user only their own images, an admin anyone's) and then PROXIES the object
|
||||
// from RustFS. Nothing is read from local disk; the RustFS endpoint is never
|
||||
// exposed to the client.
|
||||
func (h *ImageHandler) Serve(c *gin.Context) {
|
||||
user := c.Param("user")
|
||||
name := c.Param("name")
|
||||
|
||||
rel, err := h.imageAccess.Resolve(user, name)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": "invalid path"})
|
||||
return
|
||||
}
|
||||
|
||||
public, err := h.imageAccess.IsPublic(c.Request.Context(), rel)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"detail": "failed to authorize image"})
|
||||
return
|
||||
}
|
||||
if !public {
|
||||
authorized, err := h.imageAccess.IsAuthorized(
|
||||
c.Request.Context(),
|
||||
readCookie(c, h.cfg.SessionCookieName),
|
||||
user,
|
||||
)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"detail": "failed to authorize image"})
|
||||
return
|
||||
}
|
||||
if !authorized {
|
||||
c.JSON(http.StatusUnauthorized, gin.H{"detail": "需要登录后访问"})
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
// Forward Range so the browser can seek within videos.
|
||||
resp, err := h.store.Get(c.Request.Context(), rel, c.GetHeader("Range"))
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadGateway, gin.H{"detail": "failed to fetch object"})
|
||||
return
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode == http.StatusNotFound {
|
||||
c.JSON(http.StatusNotFound, gin.H{"detail": "not found"})
|
||||
return
|
||||
}
|
||||
for _, hdr := range []string{"Content-Type", "Content-Length", "Accept-Ranges", "Content-Range", "Last-Modified", "ETag", "Cache-Control"} {
|
||||
if v := resp.Header.Get(hdr); v != "" {
|
||||
c.Header(hdr, v)
|
||||
}
|
||||
}
|
||||
c.Status(resp.StatusCode)
|
||||
_, _ = io.Copy(c.Writer, resp.Body)
|
||||
}
|
||||
|
||||
func readCookie(c *gin.Context, name string) string {
|
||||
v, err := c.Cookie(name)
|
||||
if err != nil {
|
||||
return ""
|
||||
}
|
||||
return v
|
||||
}
|
||||
@@ -0,0 +1,338 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"net/http"
|
||||
|
||||
"backend/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
type ProviderAdminHandler struct {
|
||||
tokens *service.TokenService
|
||||
refresh *service.RefreshProfileService
|
||||
}
|
||||
|
||||
func NewProviderAdminHandler(tokens *service.TokenService, refresh *service.RefreshProfileService) *ProviderAdminHandler {
|
||||
return &ProviderAdminHandler{
|
||||
tokens: tokens,
|
||||
refresh: refresh,
|
||||
}
|
||||
}
|
||||
|
||||
func (h *ProviderAdminHandler) TokensList(c *gin.Context) {
|
||||
data, err := h.tokens.List(c.Request.Context())
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"detail": "failed to load tokens"})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"data": data})
|
||||
}
|
||||
|
||||
func (h *ProviderAdminHandler) TokensCreate(c *gin.Context) {
|
||||
var body struct {
|
||||
Pool string `json:"pool"`
|
||||
Value string `json:"value"`
|
||||
ID string `json:"id"`
|
||||
}
|
||||
if err := c.ShouldBindJSON(&body); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": "invalid request body"})
|
||||
return
|
||||
}
|
||||
item, err := h.tokens.Add(c.Request.Context(), body.Pool, body.Value, body.ID)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"ok": true, "id": item.ID})
|
||||
}
|
||||
|
||||
func (h *ProviderAdminHandler) ImportChatGPTToken(c *gin.Context) {
|
||||
var body struct {
|
||||
AccessToken string `json:"access_token"`
|
||||
Value string `json:"value"`
|
||||
Name string `json:"name"`
|
||||
ID string `json:"id"`
|
||||
}
|
||||
if err := c.ShouldBindJSON(&body); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": "invalid request body"})
|
||||
return
|
||||
}
|
||||
token := body.AccessToken
|
||||
if token == "" {
|
||||
token = body.Value
|
||||
}
|
||||
name := body.Name
|
||||
if name == "" {
|
||||
name = body.ID
|
||||
}
|
||||
item, err := h.tokens.ImportChatGPTToken(c.Request.Context(), token, name)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"ok": true, "id": item.ID, "status": item.Status, "pending": item.Status == "pending"})
|
||||
}
|
||||
|
||||
func (h *ProviderAdminHandler) ImportRunwayToken(c *gin.Context) {
|
||||
var body struct {
|
||||
AccessToken string `json:"access_token"`
|
||||
Value string `json:"value"`
|
||||
Name string `json:"name"`
|
||||
ID string `json:"id"`
|
||||
}
|
||||
if err := c.ShouldBindJSON(&body); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": "invalid request body"})
|
||||
return
|
||||
}
|
||||
token := body.AccessToken
|
||||
if token == "" {
|
||||
token = body.Value
|
||||
}
|
||||
name := body.Name
|
||||
if name == "" {
|
||||
name = body.ID
|
||||
}
|
||||
item, err := h.tokens.ImportRunwayToken(c.Request.Context(), token, name)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"ok": true, "id": item.ID, "status": item.Status, "pending": item.Status == "pending"})
|
||||
}
|
||||
|
||||
func (h *ProviderAdminHandler) ImportKreaCookie(c *gin.Context) {
|
||||
var body struct {
|
||||
Cookie string `json:"cookie"`
|
||||
Value string `json:"value"`
|
||||
Name string `json:"name"`
|
||||
ID string `json:"id"`
|
||||
}
|
||||
if err := c.ShouldBindJSON(&body); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": "invalid request body"})
|
||||
return
|
||||
}
|
||||
cookie := body.Cookie
|
||||
if cookie == "" {
|
||||
cookie = body.Value
|
||||
}
|
||||
name := body.Name
|
||||
if name == "" {
|
||||
name = body.ID
|
||||
}
|
||||
item, err := h.tokens.ImportKreaCookie(c.Request.Context(), cookie, name)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"ok": true, "id": item.ID, "status": item.Status, "pending": item.Status == "pending"})
|
||||
}
|
||||
|
||||
func (h *ProviderAdminHandler) ImportImagineToken(c *gin.Context) {
|
||||
var body struct {
|
||||
Cookie string `json:"cookie"`
|
||||
Value string `json:"value"`
|
||||
Name string `json:"name"`
|
||||
ID string `json:"id"`
|
||||
}
|
||||
if err := c.ShouldBindJSON(&body); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": "invalid request body"})
|
||||
return
|
||||
}
|
||||
cred := body.Cookie
|
||||
if cred == "" {
|
||||
cred = body.Value
|
||||
}
|
||||
name := body.Name
|
||||
if name == "" {
|
||||
name = body.ID
|
||||
}
|
||||
item, err := h.tokens.ImportImagineToken(c.Request.Context(), cred, name)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"ok": true, "id": item.ID, "status": item.Status, "pending": item.Status == "pending"})
|
||||
}
|
||||
|
||||
func (h *ProviderAdminHandler) ImportLeonardoCookie(c *gin.Context) {
|
||||
var body struct {
|
||||
Cookie string `json:"cookie"`
|
||||
Value string `json:"value"`
|
||||
Name string `json:"name"`
|
||||
ID string `json:"id"`
|
||||
}
|
||||
if err := c.ShouldBindJSON(&body); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": "invalid request body"})
|
||||
return
|
||||
}
|
||||
cookie := body.Cookie
|
||||
if cookie == "" {
|
||||
cookie = body.Value
|
||||
}
|
||||
name := body.Name
|
||||
if name == "" {
|
||||
name = body.ID
|
||||
}
|
||||
item, err := h.tokens.ImportLeonardoCookie(c.Request.Context(), cookie, name)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"ok": true, "id": item.ID, "status": item.Status, "pending": item.Status == "pending"})
|
||||
}
|
||||
|
||||
func (h *ProviderAdminHandler) ImportAdobeCookie(c *gin.Context) {
|
||||
var body struct {
|
||||
Cookie string `json:"cookie"`
|
||||
Value string `json:"value"`
|
||||
Name string `json:"name"`
|
||||
ID string `json:"id"`
|
||||
}
|
||||
if err := c.ShouldBindJSON(&body); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": "invalid request body"})
|
||||
return
|
||||
}
|
||||
cookie := body.Cookie
|
||||
if cookie == "" {
|
||||
cookie = body.Value
|
||||
}
|
||||
name := body.Name
|
||||
if name == "" {
|
||||
name = body.ID
|
||||
}
|
||||
item, profile, err := h.tokens.ImportAdobeCookie(c.Request.Context(), cookie, name)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"ok": true,
|
||||
"profile_id": profile.ID,
|
||||
"id": item.ID,
|
||||
"status": item.Status,
|
||||
"pending": item.Status == "pending",
|
||||
})
|
||||
}
|
||||
|
||||
func (h *ProviderAdminHandler) TokenUpdate(c *gin.Context) {
|
||||
var body map[string]any
|
||||
if err := c.ShouldBindJSON(&body); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": "invalid request body"})
|
||||
return
|
||||
}
|
||||
item, err := h.tokens.Update(c.Request.Context(), c.Param("pool"), c.Param("id"), body)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"ok": true, "data": item})
|
||||
}
|
||||
|
||||
func (h *ProviderAdminHandler) TokenDelete(c *gin.Context) {
|
||||
if err := h.tokens.Delete(c.Request.Context(), c.Param("pool"), c.Param("id")); err != nil {
|
||||
if errors.Is(err, service.ErrNotFound) {
|
||||
c.JSON(http.StatusNotFound, gin.H{"detail": "token not found"})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"ok": true})
|
||||
}
|
||||
|
||||
// TokenDeleteBulk removes multiple accounts in one call (account multi-select).
|
||||
func (h *ProviderAdminHandler) TokenDeleteBulk(c *gin.Context) {
|
||||
var body struct {
|
||||
IDs []string `json:"ids"`
|
||||
}
|
||||
if err := c.ShouldBindJSON(&body); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": "invalid request body"})
|
||||
return
|
||||
}
|
||||
if len(body.IDs) == 0 {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": "未选择任何账号"})
|
||||
return
|
||||
}
|
||||
n, err := h.tokens.DeleteBulk(c.Request.Context(), body.IDs)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"ok": true, "deleted": n})
|
||||
}
|
||||
|
||||
func (h *ProviderAdminHandler) AccountsList(c *gin.Context) {
|
||||
data, err := h.tokens.Accounts(c.Request.Context())
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"detail": "failed to load accounts"})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"data": data})
|
||||
}
|
||||
|
||||
func (h *ProviderAdminHandler) AccountQuota(c *gin.Context) {
|
||||
data, err := h.tokens.Quota(c.Request.Context(), c.Param("pool"), c.Param("id"))
|
||||
if err != nil {
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
c.JSON(http.StatusNotFound, gin.H{"detail": "account not found"})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, data)
|
||||
}
|
||||
|
||||
func (h *ProviderAdminHandler) AccountEmail(c *gin.Context) {
|
||||
data, err := h.tokens.Email(c.Request.Context(), c.Param("pool"), c.Param("id"))
|
||||
if err != nil {
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
c.JSON(http.StatusNotFound, gin.H{"detail": "account not found"})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, data)
|
||||
}
|
||||
|
||||
func (h *ProviderAdminHandler) RefreshProfiles(c *gin.Context) {
|
||||
items, err := h.refresh.List(c.Request.Context())
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"detail": "failed to load refresh profiles"})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"data": items})
|
||||
}
|
||||
|
||||
func (h *ProviderAdminHandler) RefreshNow(c *gin.Context) {
|
||||
if err := h.refresh.RefreshNow(c.Request.Context(), c.Param("profile_id")); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"ok": true})
|
||||
}
|
||||
|
||||
func (h *ProviderAdminHandler) RefreshUpdate(c *gin.Context) {
|
||||
var body map[string]any
|
||||
if err := c.ShouldBindJSON(&body); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": "invalid request body"})
|
||||
return
|
||||
}
|
||||
item, err := h.refresh.Update(c.Request.Context(), c.Param("profile_id"), body)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"ok": true, "data": item})
|
||||
}
|
||||
|
||||
func (h *ProviderAdminHandler) RefreshDelete(c *gin.Context) {
|
||||
if err := h.refresh.Delete(c.Request.Context(), c.Param("profile_id")); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"ok": true})
|
||||
}
|
||||
@@ -0,0 +1,45 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
|
||||
"backend/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
type ShowcaseHandler struct {
|
||||
showcase *service.ShowcaseService
|
||||
}
|
||||
|
||||
func NewShowcaseHandler(showcase *service.ShowcaseService) *ShowcaseHandler {
|
||||
return &ShowcaseHandler{showcase: showcase}
|
||||
}
|
||||
|
||||
func (h *ShowcaseHandler) List(c *gin.Context) {
|
||||
grouped, err := h.showcase.Grouped(c.Request.Context())
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"detail": "failed to load showcase"})
|
||||
return
|
||||
}
|
||||
out := gin.H{}
|
||||
for kind, items := range grouped {
|
||||
rows := make([]gin.H, 0, len(items))
|
||||
for _, item := range items {
|
||||
rows = append(rows, gin.H{
|
||||
"id": item.ID,
|
||||
"kind": item.Kind,
|
||||
"title": item.Title,
|
||||
"subtitle": item.Subtitle,
|
||||
"prompt": item.Prompt,
|
||||
"gradient": item.Gradient,
|
||||
"span": item.Span,
|
||||
"image": item.Image,
|
||||
"weight": item.Weight,
|
||||
"created_at": item.CreatedAt,
|
||||
"updated_at": item.UpdatedAt,
|
||||
})
|
||||
}
|
||||
out[kind] = rows
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"data": out})
|
||||
}
|
||||
@@ -0,0 +1,25 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
|
||||
"backend/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
type SiteHandler struct {
|
||||
site *service.SiteService
|
||||
}
|
||||
|
||||
func NewSiteHandler(site *service.SiteService) *SiteHandler {
|
||||
return &SiteHandler{site: site}
|
||||
}
|
||||
|
||||
func (h *SiteHandler) Public(c *gin.Context) {
|
||||
title, err := h.site.Title(c.Request.Context())
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"detail": "failed to load site"})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"title": title, "contact": h.site.Contact(c.Request.Context())})
|
||||
}
|
||||
@@ -0,0 +1,52 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"strings"
|
||||
|
||||
"backend/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
type SiteSettingsHandler struct {
|
||||
site *service.SiteService
|
||||
}
|
||||
|
||||
func NewSiteSettingsHandler(site *service.SiteService) *SiteSettingsHandler {
|
||||
return &SiteSettingsHandler{site: site}
|
||||
}
|
||||
|
||||
func (h *SiteSettingsHandler) Get(c *gin.Context) {
|
||||
title, err := h.site.Title(c.Request.Context())
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"detail": "failed to load site settings"})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"title": title, "contact": h.site.Contact(c.Request.Context())})
|
||||
}
|
||||
|
||||
func (h *SiteSettingsHandler) Put(c *gin.Context) {
|
||||
var body struct {
|
||||
Title string `json:"title"`
|
||||
Contact service.Contact `json:"contact"`
|
||||
}
|
||||
if err := c.ShouldBindJSON(&body); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": "invalid request body"})
|
||||
return
|
||||
}
|
||||
title := strings.TrimSpace(body.Title)
|
||||
if title == "" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": "网页主标题不能为空"})
|
||||
return
|
||||
}
|
||||
updated, err := h.site.SetTitle(c.Request.Context(), title)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"detail": "failed to save site settings"})
|
||||
return
|
||||
}
|
||||
if err := h.site.SetContact(c.Request.Context(), body.Contact); err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"detail": "failed to save contact info"})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"ok": true, "data": gin.H{"title": updated, "contact": h.site.Contact(c.Request.Context())}})
|
||||
}
|
||||
@@ -0,0 +1,603 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"net/http"
|
||||
|
||||
"backend/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
type UserGenerationHandler struct {
|
||||
userGen *service.UserGenerationService
|
||||
admin *service.AdminReadService
|
||||
}
|
||||
|
||||
func NewUserGenerationHandler(userGen *service.UserGenerationService, admin *service.AdminReadService) *UserGenerationHandler {
|
||||
return &UserGenerationHandler{
|
||||
userGen: userGen,
|
||||
admin: admin,
|
||||
}
|
||||
}
|
||||
|
||||
// MyImages returns the current user's own recently generated images (scoped to
|
||||
// their owner directory) — used by the showcase "选择已生成" picker so an admin
|
||||
// only sees their own images, not everyone's.
|
||||
func (h *UserGenerationHandler) MyImages(c *gin.Context) {
|
||||
user := currentUser(c)
|
||||
if user == nil {
|
||||
c.JSON(http.StatusUnauthorized, gin.H{"detail": "未登录或会话已过期"})
|
||||
return
|
||||
}
|
||||
items, err := h.admin.RecentImagesOwned(c.Request.Context(), service.OwnerDir(user), 60)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"detail": "failed to load images"})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"data": items})
|
||||
}
|
||||
|
||||
func (h *UserGenerationHandler) Generate(c *gin.Context) {
|
||||
user := currentUser(c)
|
||||
if user == nil {
|
||||
c.JSON(http.StatusUnauthorized, gin.H{"detail": "未登录或会话已过期"})
|
||||
return
|
||||
}
|
||||
|
||||
var body struct {
|
||||
Model string `json:"model"`
|
||||
Prompt string `json:"prompt"`
|
||||
Ratio string `json:"ratio"`
|
||||
Resolution string `json:"resolution"`
|
||||
Duration string `json:"duration"`
|
||||
ReferenceImages []string `json:"reference_images"`
|
||||
}
|
||||
if err := c.ShouldBindJSON(&body); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": "invalid request body"})
|
||||
return
|
||||
}
|
||||
|
||||
resp, err := h.userGen.Generate(c.Request.Context(), user, service.UserGenerateRequest{
|
||||
Model: body.Model,
|
||||
Prompt: body.Prompt,
|
||||
Ratio: body.Ratio,
|
||||
Resolution: body.Resolution,
|
||||
Duration: body.Duration,
|
||||
ReferenceImages: body.ReferenceImages,
|
||||
})
|
||||
if err != nil {
|
||||
switch {
|
||||
case errors.Is(err, service.ErrUnknownModel):
|
||||
c.JSON(http.StatusNotFound, gin.H{"detail": err.Error()})
|
||||
case errors.Is(err, service.ErrUnsupportedParams):
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": err.Error()})
|
||||
case errors.Is(err, service.ErrInsufficientFunds):
|
||||
c.JSON(http.StatusPaymentRequired, gin.H{"detail": "积分不足"})
|
||||
case errors.Is(err, service.ErrNoProviderAccount):
|
||||
c.JSON(http.StatusServiceUnavailable, gin.H{"detail": err.Error()})
|
||||
case errors.Is(err, service.ErrProviderAuth), errors.Is(err, service.ErrProviderTemporary):
|
||||
c.JSON(http.StatusServiceUnavailable, gin.H{"detail": err.Error()})
|
||||
case errors.Is(err, service.ErrProviderQuota):
|
||||
c.JSON(http.StatusTooManyRequests, gin.H{"detail": err.Error()})
|
||||
case errors.Is(err, service.ErrConcurrencyFull):
|
||||
c.JSON(http.StatusTooManyRequests, gin.H{"detail": err.Error()})
|
||||
case errors.Is(err, service.ErrProviderExecution):
|
||||
c.JSON(http.StatusBadGateway, gin.H{"detail": err.Error()})
|
||||
default:
|
||||
if err.Error() == "已有正在生成的任务,请稍候" {
|
||||
c.JSON(http.StatusConflict, gin.H{"detail": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": err.Error()})
|
||||
}
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, resp)
|
||||
}
|
||||
|
||||
func (h *UserGenerationHandler) Test(c *gin.Context) {
|
||||
user := currentUser(c)
|
||||
if user == nil {
|
||||
c.JSON(http.StatusUnauthorized, gin.H{"detail": "未登录或会话已过期"})
|
||||
return
|
||||
}
|
||||
if user.Role != "admin" {
|
||||
c.JSON(http.StatusForbidden, gin.H{"detail": "需要管理员权限"})
|
||||
return
|
||||
}
|
||||
|
||||
var body struct {
|
||||
Model string `json:"model"`
|
||||
Prompt string `json:"prompt"`
|
||||
Ratio string `json:"ratio"`
|
||||
Resolution string `json:"resolution"`
|
||||
Duration string `json:"duration"`
|
||||
ReferenceImages []string `json:"reference_images"`
|
||||
}
|
||||
if err := c.ShouldBindJSON(&body); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": "invalid request body"})
|
||||
return
|
||||
}
|
||||
|
||||
resp, err := h.userGen.AdminTest(c.Request.Context(), user, service.UserGenerateRequest{
|
||||
Model: body.Model,
|
||||
Prompt: body.Prompt,
|
||||
Ratio: body.Ratio,
|
||||
Resolution: body.Resolution,
|
||||
Duration: body.Duration,
|
||||
ReferenceImages: body.ReferenceImages,
|
||||
})
|
||||
if err != nil {
|
||||
switch {
|
||||
case errors.Is(err, service.ErrUnknownModel):
|
||||
c.JSON(http.StatusNotFound, gin.H{"detail": err.Error()})
|
||||
case errors.Is(err, service.ErrUnsupportedParams):
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": err.Error()})
|
||||
case errors.Is(err, service.ErrProviderQuota):
|
||||
c.JSON(http.StatusTooManyRequests, gin.H{"detail": err.Error()})
|
||||
case errors.Is(err, service.ErrConcurrencyFull):
|
||||
c.JSON(http.StatusTooManyRequests, gin.H{"detail": err.Error()})
|
||||
case errors.Is(err, service.ErrNoProviderAccount):
|
||||
c.JSON(http.StatusServiceUnavailable, gin.H{"detail": err.Error()})
|
||||
case errors.Is(err, service.ErrProviderAuth), errors.Is(err, service.ErrProviderTemporary):
|
||||
c.JSON(http.StatusServiceUnavailable, gin.H{"detail": err.Error()})
|
||||
case errors.Is(err, service.ErrProviderExecution):
|
||||
c.JSON(http.StatusBadGateway, gin.H{"detail": err.Error()})
|
||||
default:
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": err.Error()})
|
||||
}
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, resp)
|
||||
}
|
||||
|
||||
func (h *UserGenerationHandler) MyJobs(c *gin.Context) {
|
||||
user := currentUser(c)
|
||||
if user == nil {
|
||||
c.JSON(http.StatusOK, gin.H{"pending": nil, "latest": nil})
|
||||
return
|
||||
}
|
||||
data, err := h.userGen.MyJobs(c.Request.Context(), user, c.Query("source"))
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"detail": "failed to load jobs"})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, data)
|
||||
}
|
||||
|
||||
func (h *UserGenerationHandler) Logs(c *gin.Context) {
|
||||
user := currentUser(c)
|
||||
if user == nil {
|
||||
c.JSON(http.StatusUnauthorized, gin.H{"detail": "未登录或会话已过期"})
|
||||
return
|
||||
}
|
||||
|
||||
limit := parseInt(c.Query("limit"), 50)
|
||||
offset := parseInt(c.Query("offset"), 0)
|
||||
kind := c.Query("kind")
|
||||
status := c.Query("status")
|
||||
// Secure-by-default: always scope to the caller's OWN records. This endpoint
|
||||
// serves the front-end 日志 / 创作记录 pages, so an admin viewing their personal
|
||||
// records must NOT see other users' work. Only an admin who explicitly opts
|
||||
// into the full view (?scope=all — the admin 日志 page) sees everyone's logs.
|
||||
// API-key ("v1") usage IS included for the caller's own records so the user
|
||||
// can audit their key's calls on /mylogs; the image-only 创作记录 gallery still
|
||||
// hides them client-side (they have no stored file).
|
||||
userID := user.ID
|
||||
excludeSource := ""
|
||||
if user.Role == "admin" && c.Query("scope") == "all" {
|
||||
userID = ""
|
||||
}
|
||||
// 来源筛选: "v1" = API key, "user" = 前台画图, "admin" = 测试模型. 始终生效 ——
|
||||
// 普通用户已被 userID 限定为本人记录,按来源服务端筛选 + 分页(/mylogs 翻全部历史)。
|
||||
source := c.Query("source")
|
||||
// 创作记录 gallery passes has_file=1 so server-side pagination counts only
|
||||
// rows with real media (success + stored file), not failed/pending events.
|
||||
hasFile := c.Query("has_file") == "1" || c.Query("has_file") == "true"
|
||||
|
||||
items, total, stats, err := h.admin.Logs(c.Request.Context(), limit, offset, kind, status, nil, userID, excludeSource, source, hasFile)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"detail": "failed to load logs"})
|
||||
return
|
||||
}
|
||||
// Resolve user_id -> display name (mirrors admin.py / AdminReadHandler.Logs).
|
||||
// Without this the log table showed every row as "匿名".
|
||||
nameByID, err := h.admin.UserNameMap(c.Request.Context())
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"detail": "failed to load logs"})
|
||||
return
|
||||
}
|
||||
|
||||
out := make([]gin.H, 0, len(items))
|
||||
for _, item := range items {
|
||||
var userName any
|
||||
if item.UserID == "" {
|
||||
userName = "匿名"
|
||||
} else if name, ok := nameByID[item.UserID]; ok {
|
||||
userName = name
|
||||
} else {
|
||||
userName = item.UserID
|
||||
}
|
||||
out = append(out, gin.H{
|
||||
"id": item.ID,
|
||||
"ts": item.TS.Unix(),
|
||||
"kind": item.Kind,
|
||||
"status": item.Status,
|
||||
"model": item.Model,
|
||||
"provider": item.Provider,
|
||||
"prompt": item.Prompt,
|
||||
"ratio": item.Ratio,
|
||||
"resolution": item.Resolution,
|
||||
"duration": item.Duration,
|
||||
"refs": item.Refs,
|
||||
"source": emptyStringNil(item.Source),
|
||||
"user_id": emptyStringNil(item.UserID),
|
||||
"user_name": userName,
|
||||
"cost": item.Cost,
|
||||
"elapsed_ms": item.ElapsedMS,
|
||||
"file": emptyStringNil(item.File),
|
||||
"error": emptyStringNil(item.Error),
|
||||
"created_at": unixSec(item.CreatedAt),
|
||||
"updated_at": unixSec(item.UpdatedAt),
|
||||
})
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"data": out,
|
||||
"total": total,
|
||||
"limit": limit,
|
||||
"offset": offset,
|
||||
"stats": stats,
|
||||
})
|
||||
}
|
||||
|
||||
func (h *UserGenerationHandler) VideoPresets(c *gin.Context) {
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"data": []gin.H{
|
||||
{
|
||||
"key": "gemini-veo31",
|
||||
"label": "Veo31",
|
||||
"type": "video",
|
||||
"provider": "adobe",
|
||||
"durations": []string{"4s", "6s", "8s"},
|
||||
"ratios": []string{"16x9", "9x16"},
|
||||
"resolutions": []string{"720p", "1080p"},
|
||||
"max_reference_images": 2,
|
||||
"reference_mode": "frame",
|
||||
},
|
||||
{
|
||||
"key": "firefly-ray",
|
||||
"label": "Luma Ray",
|
||||
"type": "video",
|
||||
"provider": "adobe",
|
||||
"durations": []string{"5s", "10s"},
|
||||
"ratios": []string{"21:9", "16:9", "4:3", "1:1", "3:4", "9:16", "9:21"},
|
||||
"resolutions": []string{"720p"},
|
||||
"max_reference_images": 2,
|
||||
"reference_mode": "frame",
|
||||
},
|
||||
{
|
||||
"key": "firefly-video",
|
||||
"label": "Firefly Video",
|
||||
"type": "video",
|
||||
"provider": "adobe",
|
||||
"durations": []string{"5s"},
|
||||
"ratios": []string{"16:9", "1:1", "9:16"},
|
||||
"resolutions": []string{"540p", "720p", "1080p"},
|
||||
"max_reference_images": 2,
|
||||
"reference_mode": "frame",
|
||||
},
|
||||
{
|
||||
"key": "runway-gen4-turbo",
|
||||
"label": "Runway Gen-4 Turbo",
|
||||
"type": "video",
|
||||
"provider": "runway",
|
||||
"durations": []string{"5s", "10s"},
|
||||
"ratios": []string{"16:9", "9:16", "1:1", "4:3", "3:4", "21:9"},
|
||||
"resolutions": []string{"2K"},
|
||||
"max_reference_images": 1,
|
||||
"reference_mode": "frame",
|
||||
// Runway is strictly image-to-video — a first-frame image is required
|
||||
// (no text2video), so the UI must block submit without one.
|
||||
"requires_reference": true,
|
||||
},
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
func (h *UserGenerationHandler) Catalog(c *gin.Context) {
|
||||
items, err := h.catalogEntries(c)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"detail": "failed to load catalog"})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"data": items,
|
||||
})
|
||||
}
|
||||
|
||||
func (h *UserGenerationHandler) Models(c *gin.Context) {
|
||||
items, err := h.publicModels()
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"detail": "failed to load models"})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"data": items})
|
||||
}
|
||||
|
||||
func (h *UserGenerationHandler) catalogEntries(c *gin.Context) ([]gin.H, error) {
|
||||
items := []gin.H{
|
||||
{
|
||||
"id": "gpt-image-2",
|
||||
"provider": "chatgpt",
|
||||
"type": "image",
|
||||
// ChatGPT web backend only reliably produces 1K and honors a limited
|
||||
// ratio set; size params are advisory prompt hints. Mirrors the Python
|
||||
// reference (providers/chatgpt/provider.py) — do not offer 2K/4K.
|
||||
"ratios": []string{"1:1", "16:9", "9:16", "4:3", "3:4"},
|
||||
"resolutions": []string{"1K"},
|
||||
"image_to_image": true,
|
||||
"max_reference_images": 3,
|
||||
"description": "ChatGPT image generation",
|
||||
},
|
||||
{
|
||||
"id": "firefly-gpt-image-2",
|
||||
"provider": "adobe",
|
||||
"type": "image",
|
||||
"ratios": []string{"1:1", "5:4", "9:16", "21:9", "16:9", "4:3", "3:2", "4:5", "3:4", "2:3"},
|
||||
"resolutions": []string{"1K", "2K", "4K"},
|
||||
"image_to_image": true,
|
||||
"max_reference_images": 6,
|
||||
"description": "Adobe Firefly GPT Image",
|
||||
},
|
||||
{
|
||||
"id": "firefly-image-5",
|
||||
"provider": "adobe",
|
||||
"type": "image",
|
||||
"ratios": []string{"1:1", "16:9", "9:16", "4:3", "3:4"},
|
||||
"resolutions": []string{"1K", "2K"},
|
||||
"image_to_image": true,
|
||||
"description": "Adobe Firefly Image 5",
|
||||
},
|
||||
{
|
||||
"id": "flux-kontext-max",
|
||||
"provider": "adobe",
|
||||
"type": "image",
|
||||
"ratios": []string{"1:1", "16:9", "9:16", "4:3", "3:4"},
|
||||
"resolutions": []string{"1K"},
|
||||
"image_to_image": true,
|
||||
"max_reference_images": 4,
|
||||
"description": "Adobe Flux Kontext Max",
|
||||
},
|
||||
{
|
||||
"id": "nano-banana-2",
|
||||
"provider": "adobe",
|
||||
"type": "image",
|
||||
"ratios": []string{"1:1", "16:9", "9:16", "4:3", "3:4"},
|
||||
"resolutions": []string{"1K"},
|
||||
"image_to_image": true,
|
||||
"description": "Adobe Gemini Flash Nano Banana",
|
||||
},
|
||||
{
|
||||
"id": "gemini-veo31",
|
||||
"provider": "adobe",
|
||||
"type": "video",
|
||||
"ratios": []string{"16x9", "9x16"},
|
||||
"resolutions": []string{"720p", "1080p"},
|
||||
"durations": []string{"4s", "6s", "8s"},
|
||||
"max_reference_images": 2,
|
||||
"reference_mode": "frame",
|
||||
"description": "Veo31 video",
|
||||
},
|
||||
{
|
||||
"id": "firefly-ray",
|
||||
"provider": "adobe",
|
||||
"type": "video",
|
||||
"ratios": []string{"21:9", "16:9", "4:3", "1:1", "3:4", "9:16", "9:21"},
|
||||
"resolutions": []string{"720p"},
|
||||
"durations": []string{"5s", "10s"},
|
||||
"max_reference_images": 2,
|
||||
"reference_mode": "frame",
|
||||
"description": "Luma Ray video",
|
||||
},
|
||||
{
|
||||
"id": "firefly-video",
|
||||
"provider": "adobe",
|
||||
"type": "video",
|
||||
"ratios": []string{"16:9", "1:1", "9:16"},
|
||||
"resolutions": []string{"540p", "720p", "1080p"},
|
||||
"durations": []string{"5s"},
|
||||
"max_reference_images": 2,
|
||||
"reference_mode": "frame",
|
||||
"description": "Adobe Firefly Video",
|
||||
},
|
||||
{
|
||||
"id": "runway-gen4-turbo",
|
||||
"provider": "runway",
|
||||
"type": "video",
|
||||
"ratios": []string{"16:9", "9:16", "1:1", "4:3", "3:4", "21:9"},
|
||||
"resolutions": []string{"2K"},
|
||||
"durations": []string{"5s", "10s"},
|
||||
"max_reference_images": 1,
|
||||
"reference_mode": "frame",
|
||||
"description": "Runway Gen-4 Turbo video (图生视频)",
|
||||
},
|
||||
{
|
||||
"id": "seedream-4.5",
|
||||
"provider": "leonardo",
|
||||
"type": "image",
|
||||
"ratios": []string{"2:3", "1:1", "16:9", "4:3", "4:5", "9:16", "2:1"},
|
||||
"resolutions": []string{"2K", "4K"},
|
||||
"image_to_image": true,
|
||||
"max_reference_images": 6,
|
||||
"description": "Leonardo Seedream 4.5 (生图 / 图生图)",
|
||||
},
|
||||
{
|
||||
"id": "flux-klein-2",
|
||||
"provider": "krea",
|
||||
"type": "image",
|
||||
"ratios": []string{"1:1", "4:3", "3:4", "16:9", "9:16"},
|
||||
"resolutions": []string{"1K", "2K"},
|
||||
"image_to_image": true,
|
||||
"max_reference_images": 4,
|
||||
"description": "Krea Flux Klein (生图 / 图生图)",
|
||||
},
|
||||
{
|
||||
"id": "imagine-1.5",
|
||||
"provider": "imagine",
|
||||
"type": "image",
|
||||
"ratios": []string{"1:3", "9:16", "2:3", "3:4", "1:1", "4:3", "3:2", "16:9", "3:1"},
|
||||
"resolutions": []string{"2K"},
|
||||
"max_reference_images": 0,
|
||||
"description": "Imagine 1.5 (文生图)",
|
||||
},
|
||||
{
|
||||
"id": "imagine-1.5pro",
|
||||
"provider": "imagine",
|
||||
"type": "image",
|
||||
"ratios": []string{"1:3", "9:16", "2:3", "3:4", "1:1", "4:3", "3:2", "16:9", "3:1"},
|
||||
"resolutions": []string{"4K"},
|
||||
"max_reference_images": 0,
|
||||
"description": "Imagine 1.5 Pro (文生图)",
|
||||
},
|
||||
}
|
||||
existing := map[string]bool{}
|
||||
if h.admin != nil {
|
||||
models, err := h.admin.Models(c.Request.Context())
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
for _, item := range models {
|
||||
existing[item.ID] = true
|
||||
}
|
||||
}
|
||||
for i := range items {
|
||||
items[i]["added"] = existing[items[i]["id"].(string)]
|
||||
}
|
||||
return items, nil
|
||||
}
|
||||
|
||||
func (h *UserGenerationHandler) publicModels() ([]gin.H, error) {
|
||||
items := []gin.H{
|
||||
{
|
||||
"id": "gpt-image-2",
|
||||
"provider": "chatgpt",
|
||||
"kind": "image",
|
||||
// See catalogEntries — ChatGPT only reliably does 1K and a limited
|
||||
// ratio set; matches the Python reference. Keep both lists in sync.
|
||||
"ratios": []string{"1:1", "16:9", "9:16", "4:3", "3:4"},
|
||||
"resolutions": []string{"1K"},
|
||||
"description": "ChatGPT image generation",
|
||||
"stub": false,
|
||||
},
|
||||
{
|
||||
"id": "firefly-gpt-image-2",
|
||||
"provider": "adobe",
|
||||
"kind": "image",
|
||||
"ratios": []string{"1:1", "5:4", "9:16", "21:9", "16:9", "4:3", "3:2", "4:5", "3:4", "2:3"},
|
||||
"resolutions": []string{"1K", "2K", "4K"},
|
||||
"description": "Adobe Firefly GPT Image",
|
||||
"stub": false,
|
||||
},
|
||||
{
|
||||
"id": "firefly-image-5",
|
||||
"provider": "adobe",
|
||||
"kind": "image",
|
||||
"ratios": []string{"1:1", "16:9", "9:16", "4:3", "3:4"},
|
||||
"resolutions": []string{"1K", "2K"},
|
||||
"description": "Adobe Firefly Image 5",
|
||||
"stub": false,
|
||||
},
|
||||
{
|
||||
"id": "flux-kontext-max",
|
||||
"provider": "adobe",
|
||||
"kind": "image",
|
||||
"ratios": []string{"1:1", "16:9", "9:16", "4:3", "3:4"},
|
||||
"resolutions": []string{"1K"},
|
||||
"description": "Adobe Flux Kontext Max",
|
||||
"stub": false,
|
||||
},
|
||||
{
|
||||
"id": "nano-banana-2",
|
||||
"provider": "adobe",
|
||||
"kind": "image",
|
||||
"ratios": []string{"1:1", "16:9", "9:16", "4:3", "3:4"},
|
||||
"resolutions": []string{"1K"},
|
||||
"description": "Adobe Gemini Flash Nano Banana",
|
||||
"stub": false,
|
||||
},
|
||||
{
|
||||
"id": "gemini-veo31",
|
||||
"provider": "adobe",
|
||||
"kind": "video",
|
||||
"ratios": []string{"16x9", "9x16"},
|
||||
"resolutions": []string{"720p", "1080p"},
|
||||
"description": "Veo31 video",
|
||||
"stub": false,
|
||||
},
|
||||
{
|
||||
"id": "firefly-ray",
|
||||
"provider": "adobe",
|
||||
"kind": "video",
|
||||
"ratios": []string{"21:9", "16:9", "4:3", "1:1", "3:4", "9:16", "9:21"},
|
||||
"resolutions": []string{"720p"},
|
||||
"description": "Luma Ray video",
|
||||
"stub": false,
|
||||
},
|
||||
{
|
||||
"id": "firefly-video",
|
||||
"provider": "adobe",
|
||||
"kind": "video",
|
||||
"ratios": []string{"16:9", "1:1", "9:16"},
|
||||
"resolutions": []string{"540p", "720p", "1080p"},
|
||||
"description": "Adobe Firefly Video",
|
||||
"stub": false,
|
||||
},
|
||||
{
|
||||
"id": "runway-gen4-turbo",
|
||||
"provider": "runway",
|
||||
"kind": "video",
|
||||
"ratios": []string{"16:9", "9:16", "1:1", "4:3", "3:4", "21:9"},
|
||||
"resolutions": []string{"2K"},
|
||||
"description": "Runway Gen-4 Turbo video",
|
||||
"stub": false,
|
||||
},
|
||||
{
|
||||
"id": "seedream-4.5",
|
||||
"provider": "leonardo",
|
||||
"kind": "image",
|
||||
"ratios": []string{"2:3", "1:1", "16:9", "4:3", "4:5", "9:16", "2:1"},
|
||||
"resolutions": []string{"2K", "4K"},
|
||||
"description": "Leonardo Seedream 4.5",
|
||||
"stub": false,
|
||||
},
|
||||
{
|
||||
"id": "flux-klein-2",
|
||||
"provider": "krea",
|
||||
"kind": "image",
|
||||
"ratios": []string{"1:1", "4:3", "3:4", "16:9", "9:16"},
|
||||
"resolutions": []string{"1K", "2K"},
|
||||
"description": "Krea Flux Klein",
|
||||
"stub": false,
|
||||
},
|
||||
{
|
||||
"id": "imagine-1.5",
|
||||
"provider": "imagine",
|
||||
"kind": "image",
|
||||
"ratios": []string{"1:3", "9:16", "2:3", "3:4", "1:1", "4:3", "3:2", "16:9", "3:1"},
|
||||
"resolutions": []string{"2K"},
|
||||
"description": "Imagine 1.5",
|
||||
"stub": false,
|
||||
},
|
||||
{
|
||||
"id": "imagine-1.5pro",
|
||||
"provider": "imagine",
|
||||
"kind": "image",
|
||||
"ratios": []string{"1:3", "9:16", "2:3", "3:4", "1:1", "4:3", "3:2", "16:9", "3:1"},
|
||||
"resolutions": []string{"4K"},
|
||||
"description": "Imagine 1.5 Pro",
|
||||
"stub": false,
|
||||
},
|
||||
}
|
||||
return items, nil
|
||||
}
|
||||
@@ -0,0 +1,92 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
|
||||
"backend/internal/model"
|
||||
"backend/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
type UserToolsHandler struct {
|
||||
keys *service.APIKeyService
|
||||
cdks *service.CDKService
|
||||
}
|
||||
|
||||
func NewUserToolsHandler(keys *service.APIKeyService, cdks *service.CDKService) *UserToolsHandler {
|
||||
return &UserToolsHandler{
|
||||
keys: keys,
|
||||
cdks: cdks,
|
||||
}
|
||||
}
|
||||
|
||||
func (h *UserToolsHandler) APIKeyGet(c *gin.Context) {
|
||||
user := currentUser(c)
|
||||
if user == nil {
|
||||
c.JSON(http.StatusUnauthorized, gin.H{"detail": "未登录或会话已过期"})
|
||||
return
|
||||
}
|
||||
data, err := h.keys.Current(c.Request.Context(), user.ID)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"detail": "failed to load api key"})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, data)
|
||||
}
|
||||
|
||||
func (h *UserToolsHandler) APIKeyMint(c *gin.Context) {
|
||||
user := currentUser(c)
|
||||
if user == nil {
|
||||
c.JSON(http.StatusUnauthorized, gin.H{"detail": "未登录或会话已过期"})
|
||||
return
|
||||
}
|
||||
data, err := h.keys.Mint(c.Request.Context(), user.ID)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"detail": "failed to mint api key"})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, data)
|
||||
}
|
||||
|
||||
func (h *UserToolsHandler) APIKeyDelete(c *gin.Context) {
|
||||
user := currentUser(c)
|
||||
if user == nil {
|
||||
c.JSON(http.StatusUnauthorized, gin.H{"detail": "未登录或会话已过期"})
|
||||
return
|
||||
}
|
||||
if err := h.keys.Revoke(c.Request.Context(), user.ID); err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"detail": "failed to revoke api key"})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"ok": true})
|
||||
}
|
||||
|
||||
func (h *UserToolsHandler) RedeemCDK(c *gin.Context) {
|
||||
user := currentUser(c)
|
||||
if user == nil {
|
||||
c.JSON(http.StatusUnauthorized, gin.H{"detail": "未登录或会话已过期"})
|
||||
return
|
||||
}
|
||||
var body struct {
|
||||
Code string `json:"code"`
|
||||
}
|
||||
if err := c.ShouldBindJSON(&body); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": "invalid request body"})
|
||||
return
|
||||
}
|
||||
data, err := h.cdks.Redeem(c.Request.Context(), user.ID, body.Code)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"ok": true, "amount": data["amount"], "credits": data["credits"]})
|
||||
}
|
||||
|
||||
func currentUser(c *gin.Context) *model.User {
|
||||
value, ok := c.Get("current_user")
|
||||
if !ok {
|
||||
return nil
|
||||
}
|
||||
user, _ := value.(*model.User)
|
||||
return user
|
||||
}
|
||||
@@ -0,0 +1,373 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
"backend/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
type V1Handler struct {
|
||||
v1 *service.V1Service
|
||||
}
|
||||
|
||||
func NewV1Handler(v1 *service.V1Service) *V1Handler {
|
||||
return &V1Handler{v1: v1}
|
||||
}
|
||||
|
||||
func (h *V1Handler) Models(c *gin.Context) {
|
||||
principal, err := h.v1.Authenticate(c.Request.Context(), c.GetHeader("Authorization"))
|
||||
if err != nil {
|
||||
h.writeAuthError(c, err)
|
||||
return
|
||||
}
|
||||
_ = principal
|
||||
|
||||
items, err := h.v1.ListModels(c.Request.Context())
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"detail": "failed to load models"})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"object": "list",
|
||||
"data": items,
|
||||
})
|
||||
}
|
||||
|
||||
// ImageGenerations — OpenAI POST /v1/images/generations (text-to-image only).
|
||||
// Accepts exactly OpenAI's fields; size→aspect ratio and quality→resolution tier
|
||||
// are mapped server-side. Returns {created, data:[{b64_json}]}.
|
||||
func (h *V1Handler) ImageGenerations(c *gin.Context) {
|
||||
principal, err := h.v1.Authenticate(c.Request.Context(), c.GetHeader("Authorization"))
|
||||
if err != nil {
|
||||
h.writeAuthError(c, err)
|
||||
return
|
||||
}
|
||||
|
||||
var body struct {
|
||||
Model string `json:"model"`
|
||||
Prompt string `json:"prompt"`
|
||||
N int `json:"n"`
|
||||
Size string `json:"size"`
|
||||
Quality string `json:"quality"`
|
||||
ResponseFormat string `json:"response_format"`
|
||||
Background string `json:"background"`
|
||||
OutputFormat string `json:"output_format"`
|
||||
User string `json:"user"`
|
||||
}
|
||||
if err := c.ShouldBindJSON(&body); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": "invalid request body"})
|
||||
return
|
||||
}
|
||||
|
||||
resp, err := h.v1.PrepareImageRequest(c.Request.Context(), principal, service.V1ImageRequest{
|
||||
Model: body.Model,
|
||||
Prompt: body.Prompt,
|
||||
N: body.N,
|
||||
Size: body.Size,
|
||||
Quality: body.Quality,
|
||||
BaseURL: requestBaseURL(c),
|
||||
})
|
||||
if err != nil {
|
||||
h.writeV1Error(c, err, resp)
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, openaiImageResponse(resp))
|
||||
}
|
||||
|
||||
// ImageEdits — OpenAI POST /v1/images/edits (image-to-image). multipart/form-data
|
||||
// only: image / image[] file uploads (+ optional mask), prompt, model, n, size,
|
||||
// quality. Files become reference images. Returns {created, data:[{b64_json}]}.
|
||||
func (h *V1Handler) ImageEdits(c *gin.Context) {
|
||||
principal, err := h.v1.Authenticate(c.Request.Context(), c.GetHeader("Authorization"))
|
||||
if err != nil {
|
||||
h.writeAuthError(c, err)
|
||||
return
|
||||
}
|
||||
if !strings.HasPrefix(c.GetHeader("Content-Type"), "multipart/form-data") {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": "images/edits requires multipart/form-data"})
|
||||
return
|
||||
}
|
||||
if err := c.Request.ParseMultipartForm(64 << 20); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": "invalid multipart form"})
|
||||
return
|
||||
}
|
||||
refs := readMultipartImages(c, "image", "image[]")
|
||||
if len(refs) == 0 {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": "images/edits requires at least one image file"})
|
||||
return
|
||||
}
|
||||
n, _ := strconv.Atoi(strings.TrimSpace(c.PostForm("n")))
|
||||
resp, err := h.v1.PrepareImageRequest(c.Request.Context(), principal, service.V1ImageRequest{
|
||||
Model: c.PostForm("model"),
|
||||
Prompt: c.PostForm("prompt"),
|
||||
N: n,
|
||||
Size: c.PostForm("size"),
|
||||
Quality: c.PostForm("quality"),
|
||||
ReferenceImages: refs,
|
||||
BaseURL: requestBaseURL(c),
|
||||
})
|
||||
if err != nil {
|
||||
h.writeV1Error(c, err, resp)
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, openaiImageResponse(resp))
|
||||
}
|
||||
|
||||
// CreateVideo — OpenAI POST /v1/videos. Creates an async job and returns the
|
||||
// video object immediately ({id, status:"queued"}). Accepts JSON {model, prompt,
|
||||
// seconds, size} or multipart (with an input_reference file). size→ratio+
|
||||
// resolution, seconds→duration.
|
||||
func (h *V1Handler) CreateVideo(c *gin.Context) {
|
||||
principal, err := h.v1.Authenticate(c.Request.Context(), c.GetHeader("Authorization"))
|
||||
if err != nil {
|
||||
h.writeAuthError(c, err)
|
||||
return
|
||||
}
|
||||
var modelID, prompt, seconds, size string
|
||||
var refs []string
|
||||
if strings.HasPrefix(c.GetHeader("Content-Type"), "multipart/form-data") {
|
||||
if err := c.Request.ParseMultipartForm(64 << 20); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": "invalid multipart form"})
|
||||
return
|
||||
}
|
||||
modelID = c.PostForm("model")
|
||||
prompt = c.PostForm("prompt")
|
||||
seconds = c.PostForm("seconds")
|
||||
size = c.PostForm("size")
|
||||
refs = readMultipartImages(c, "input_reference", "input_reference[]")
|
||||
} else {
|
||||
var body struct {
|
||||
Model string `json:"model"`
|
||||
Prompt string `json:"prompt"`
|
||||
Seconds json.RawMessage `json:"seconds"`
|
||||
Size string `json:"size"`
|
||||
}
|
||||
if err := c.ShouldBindJSON(&body); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": "invalid request body"})
|
||||
return
|
||||
}
|
||||
modelID, prompt, size = body.Model, body.Prompt, body.Size
|
||||
seconds = rawToString(body.Seconds)
|
||||
}
|
||||
duration := strings.TrimSpace(seconds)
|
||||
if duration != "" && !strings.HasSuffix(duration, "s") {
|
||||
duration += "s"
|
||||
}
|
||||
aspect, resolution := videoSizeToInternal(size)
|
||||
resp, err := h.v1.StartVideoJob(c.Request.Context(), principal, service.V1VideoRequest{
|
||||
Model: modelID,
|
||||
Prompt: prompt,
|
||||
Duration: duration,
|
||||
AspectRatio: aspect,
|
||||
Resolution: resolution,
|
||||
ReferenceImages: refs,
|
||||
BaseURL: requestBaseURL(c),
|
||||
})
|
||||
if err != nil {
|
||||
h.writeV1Error(c, err, nil)
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, resp)
|
||||
}
|
||||
|
||||
// GetVideo — OpenAI GET /v1/videos/{id}. Returns the job status object.
|
||||
func (h *V1Handler) GetVideo(c *gin.Context) {
|
||||
principal, err := h.v1.Authenticate(c.Request.Context(), c.GetHeader("Authorization"))
|
||||
if err != nil {
|
||||
h.writeAuthError(c, err)
|
||||
return
|
||||
}
|
||||
resp, err := h.v1.VideoJob(c.Request.Context(), principal, c.Param("id"))
|
||||
if err != nil {
|
||||
h.writeV1Error(c, err, nil)
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, resp)
|
||||
}
|
||||
|
||||
// GetVideoContent — OpenAI GET /v1/videos/{id}/content. Streams the rendered mp4
|
||||
// by proxying the stored upstream URL (downloaded on demand, never persisted).
|
||||
func (h *V1Handler) GetVideoContent(c *gin.Context) {
|
||||
principal, err := h.v1.Authenticate(c.Request.Context(), c.GetHeader("Authorization"))
|
||||
if err != nil {
|
||||
h.writeAuthError(c, err)
|
||||
return
|
||||
}
|
||||
body, contentType, err := h.v1.OpenVideoContent(c.Request.Context(), principal, c.Param("id"))
|
||||
if err != nil {
|
||||
h.writeV1Error(c, err, nil)
|
||||
return
|
||||
}
|
||||
defer body.Close()
|
||||
c.Header("Content-Type", contentType)
|
||||
c.Status(http.StatusOK)
|
||||
_, _ = io.Copy(c.Writer, body)
|
||||
}
|
||||
|
||||
// readMultipartImages reads the given file fields and returns each as base64.
|
||||
func readMultipartImages(c *gin.Context, keys ...string) []string {
|
||||
var out []string
|
||||
form := c.Request.MultipartForm
|
||||
if form == nil {
|
||||
return out
|
||||
}
|
||||
for _, key := range keys {
|
||||
for _, fh := range form.File[key] {
|
||||
f, e := fh.Open()
|
||||
if e != nil {
|
||||
continue
|
||||
}
|
||||
b, _ := io.ReadAll(io.LimitReader(f, 8<<20+1))
|
||||
f.Close()
|
||||
if len(b) > 0 {
|
||||
out = append(out, base64.StdEncoding.EncodeToString(b))
|
||||
}
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// rawToString accepts OpenAI's `seconds` whether sent as a JSON string or number.
|
||||
func rawToString(raw json.RawMessage) string {
|
||||
if len(raw) == 0 {
|
||||
return ""
|
||||
}
|
||||
var s string
|
||||
if json.Unmarshal(raw, &s) == nil {
|
||||
return s
|
||||
}
|
||||
var n json.Number
|
||||
if json.Unmarshal(raw, &n) == nil {
|
||||
return n.String()
|
||||
}
|
||||
return strings.Trim(string(raw), `"`)
|
||||
}
|
||||
|
||||
// videoSizeToInternal maps OpenAI's "WxH" size to our aspect ratio + resolution
|
||||
// tier (height ≥1080 → 1080p, else 720p).
|
||||
func videoSizeToInternal(size string) (ratio, resolution string) {
|
||||
var w, h int
|
||||
if s := strings.TrimSpace(strings.ToLower(size)); s != "" {
|
||||
_, _ = fmt.Sscanf(s, "%dx%d", &w, &h)
|
||||
}
|
||||
if w == 0 || h == 0 {
|
||||
return "16:9", "720p"
|
||||
}
|
||||
long := w
|
||||
if h > long {
|
||||
long = h
|
||||
}
|
||||
resolution = "720p"
|
||||
if long >= 1080 {
|
||||
resolution = "1080p"
|
||||
}
|
||||
return guessRatioWH(w, h), resolution
|
||||
}
|
||||
|
||||
func guessRatioWH(w, h int) string {
|
||||
if w == h {
|
||||
return "1:1"
|
||||
}
|
||||
r := float64(w) / float64(h)
|
||||
cands := []struct {
|
||||
name string
|
||||
v float64
|
||||
}{{"16:9", 16.0 / 9}, {"9:16", 9.0 / 16}, {"4:3", 4.0 / 3}, {"3:4", 3.0 / 4}, {"1:1", 1}}
|
||||
best, bestD := "16:9", 1e9
|
||||
for _, cd := range cands {
|
||||
d := r - cd.v
|
||||
if d < 0 {
|
||||
d = -d
|
||||
}
|
||||
if d < bestD {
|
||||
best, bestD = cd.name, d
|
||||
}
|
||||
}
|
||||
return best
|
||||
}
|
||||
|
||||
// openaiImageResponse strips our rich internal map down to OpenAI's image shape.
|
||||
func openaiImageResponse(m map[string]any) gin.H {
|
||||
out := gin.H{"created": m["created"]}
|
||||
if d, ok := m["data"]; ok && d != nil {
|
||||
out["data"] = d
|
||||
} else {
|
||||
out["data"] = []any{}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func (h *V1Handler) writeAuthError(c *gin.Context, err error) {
|
||||
switch {
|
||||
case errors.Is(err, service.ErrMissingAPIKey):
|
||||
c.JSON(http.StatusUnauthorized, gin.H{"detail": err.Error()})
|
||||
case errors.Is(err, service.ErrInvalidAPIKey):
|
||||
c.JSON(http.StatusUnauthorized, gin.H{"detail": err.Error()})
|
||||
default:
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"detail": "failed to validate api key"})
|
||||
}
|
||||
}
|
||||
|
||||
func (h *V1Handler) writeV1Error(c *gin.Context, err error, payload map[string]any) {
|
||||
switch {
|
||||
case errors.Is(err, service.ErrUnknownModel):
|
||||
c.JSON(http.StatusNotFound, gin.H{"detail": err.Error()})
|
||||
case errors.Is(err, service.ErrUnsupportedParams):
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": err.Error()})
|
||||
case errors.Is(err, service.ErrInsufficientFunds):
|
||||
c.JSON(http.StatusPaymentRequired, gin.H{"detail": err.Error()})
|
||||
case errors.Is(err, service.ErrReferenceTooLarge):
|
||||
c.JSON(http.StatusRequestEntityTooLarge, gin.H{"detail": err.Error()})
|
||||
case errors.Is(err, service.ErrNoProviderAccount):
|
||||
c.JSON(http.StatusServiceUnavailable, gin.H{"detail": err.Error()})
|
||||
case errors.Is(err, service.ErrProviderAuth):
|
||||
c.JSON(http.StatusServiceUnavailable, gin.H{"detail": err.Error()})
|
||||
case errors.Is(err, service.ErrProviderQuota):
|
||||
// Match the Python contract: provider quota exhaustion maps to 401
|
||||
// (QuotaExhaustedError is handled alongside AuthError in routes.py).
|
||||
c.JSON(http.StatusUnauthorized, gin.H{"detail": err.Error()})
|
||||
case errors.Is(err, service.ErrProviderTemporary):
|
||||
c.JSON(http.StatusServiceUnavailable, gin.H{"detail": err.Error()})
|
||||
case errors.Is(err, service.ErrConcurrencyFull):
|
||||
c.JSON(http.StatusTooManyRequests, gin.H{"detail": err.Error()})
|
||||
case errors.Is(err, service.ErrVideoJobNotFound):
|
||||
c.JSON(http.StatusNotFound, gin.H{"detail": err.Error()})
|
||||
case errors.Is(err, service.ErrVideoNotReady):
|
||||
c.JSON(http.StatusConflict, gin.H{"detail": err.Error()})
|
||||
case errors.Is(err, service.ErrProviderUnsupported):
|
||||
c.JSON(http.StatusNotImplemented, gin.H{"detail": err.Error()})
|
||||
case errors.Is(err, service.ErrProviderExecution):
|
||||
c.JSON(http.StatusBadGateway, gin.H{"detail": err.Error()})
|
||||
case errors.Is(err, service.ErrGenerationPending):
|
||||
c.JSON(http.StatusNotImplemented, payload)
|
||||
default:
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": err.Error()})
|
||||
}
|
||||
}
|
||||
|
||||
// requestBaseURL derives the scheme+host of the inbound request so the service
|
||||
// layer can build absolute, directly-downloadable output URLs. Honors
|
||||
// X-Forwarded-Proto (reverse-proxy / TLS termination) before falling back to
|
||||
// the connection's TLS state. Returns "" when the host is unknown, which makes
|
||||
// the service fall back to a relative path.
|
||||
func requestBaseURL(c *gin.Context) string {
|
||||
host := c.Request.Host
|
||||
if host == "" {
|
||||
return ""
|
||||
}
|
||||
scheme := "http"
|
||||
if proto := strings.TrimSpace(c.GetHeader("X-Forwarded-Proto")); proto != "" {
|
||||
scheme = strings.ToLower(strings.Split(proto, ",")[0])
|
||||
} else if c.Request.TLS != nil {
|
||||
scheme = "https"
|
||||
}
|
||||
return scheme + "://" + host
|
||||
}
|
||||
@@ -0,0 +1,70 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
|
||||
"backend/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
const currentUserKey = "current_user"
|
||||
const currentSessionKey = "current_session"
|
||||
|
||||
func RequireSession(auth *service.AuthService) gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
user, session, err := auth.CurrentUserFromRequest(
|
||||
c.Request.Context(),
|
||||
c.GetHeader("Authorization"),
|
||||
readCookie(c, "vivid_session"),
|
||||
)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"detail": "failed to validate session"})
|
||||
c.Abort()
|
||||
return
|
||||
}
|
||||
if user == nil || session == nil {
|
||||
c.JSON(http.StatusUnauthorized, gin.H{"detail": "未登录或会话已过期"})
|
||||
c.Abort()
|
||||
return
|
||||
}
|
||||
c.Set(currentUserKey, user)
|
||||
c.Set(currentSessionKey, session)
|
||||
c.Next()
|
||||
}
|
||||
}
|
||||
|
||||
func RequireAdminSession(auth *service.AuthService) gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
user, session, err := auth.CurrentUserFromRequest(
|
||||
c.Request.Context(),
|
||||
c.GetHeader("Authorization"),
|
||||
readCookie(c, "vivid_session"),
|
||||
)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"detail": "failed to validate session"})
|
||||
c.Abort()
|
||||
return
|
||||
}
|
||||
if user == nil || session == nil {
|
||||
c.JSON(http.StatusUnauthorized, gin.H{"detail": "未登录或会话已过期"})
|
||||
c.Abort()
|
||||
return
|
||||
}
|
||||
if user.Role != "admin" {
|
||||
c.JSON(http.StatusForbidden, gin.H{"detail": "需要管理员权限"})
|
||||
c.Abort()
|
||||
return
|
||||
}
|
||||
c.Set(currentUserKey, user)
|
||||
c.Set(currentSessionKey, session)
|
||||
c.Next()
|
||||
}
|
||||
}
|
||||
|
||||
func readCookie(c *gin.Context, name string) string {
|
||||
v, err := c.Cookie(name)
|
||||
if err != nil {
|
||||
return ""
|
||||
}
|
||||
return v
|
||||
}
|
||||
@@ -0,0 +1,20 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
const requestIDKey = "request_id"
|
||||
|
||||
func RequestID() gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
reqID := c.GetHeader("X-Request-Id")
|
||||
if reqID == "" {
|
||||
reqID = uuid.NewString()
|
||||
}
|
||||
c.Set(requestIDKey, reqID)
|
||||
c.Writer.Header().Set("X-Request-Id", reqID)
|
||||
c.Next()
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,166 @@
|
||||
package router
|
||||
|
||||
import (
|
||||
"backend/internal/config"
|
||||
"backend/internal/http/handler"
|
||||
"backend/internal/http/middleware"
|
||||
"backend/internal/service"
|
||||
|
||||
"github.com/gin-contrib/cors"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
type Handlers struct {
|
||||
Health *handler.HealthHandler
|
||||
Images *handler.ImageHandler
|
||||
V1 *handler.V1Handler
|
||||
Site *handler.SiteHandler
|
||||
Showcase *handler.ShowcaseHandler
|
||||
Auth *handler.AuthHandler
|
||||
SiteSettings *handler.SiteSettingsHandler
|
||||
AppSettings *handler.AppSettingsHandler
|
||||
AdminRead *handler.AdminReadHandler
|
||||
AdminWrite *handler.AdminWriteHandler
|
||||
CDK *handler.CDKHandler
|
||||
UserTools *handler.UserToolsHandler
|
||||
UserGen *handler.UserGenerationHandler
|
||||
ProviderAdmin *handler.ProviderAdminHandler
|
||||
}
|
||||
|
||||
func New(cfg *config.Config, auth *service.AuthService, handlers Handlers) *gin.Engine {
|
||||
if cfg.AppEnv != "development" {
|
||||
gin.SetMode(gin.ReleaseMode)
|
||||
}
|
||||
|
||||
engine := gin.New()
|
||||
engine.Use(gin.Recovery())
|
||||
engine.Use(middleware.RequestID())
|
||||
engine.Use(cors.New(cors.Config{
|
||||
AllowOrigins: cfg.CORSOrigins,
|
||||
AllowMethods: []string{"GET", "POST", "PUT", "PATCH", "DELETE", "OPTIONS"},
|
||||
AllowHeaders: []string{"Authorization", "Content-Type", "X-Request-Id"},
|
||||
AllowCredentials: true,
|
||||
}))
|
||||
|
||||
engine.GET("/health", handlers.Health.Handle)
|
||||
engine.GET("/images/:user/:name", handlers.Images.Serve)
|
||||
engine.GET("/v1/models", handlers.V1.Models)
|
||||
engine.POST("/v1/images/generations", handlers.V1.ImageGenerations)
|
||||
engine.POST("/v1/images/edits", handlers.V1.ImageEdits)
|
||||
// OpenAI Sora-style async video: create job → poll → stream content.
|
||||
engine.POST("/v1/videos", handlers.V1.CreateVideo)
|
||||
engine.GET("/v1/videos/:id", handlers.V1.GetVideo)
|
||||
engine.GET("/v1/videos/:id/content", handlers.V1.GetVideoContent)
|
||||
|
||||
publicAdmin := engine.Group("/admin/api")
|
||||
{
|
||||
publicAdmin.GET("/site", handlers.Site.Public)
|
||||
publicAdmin.GET("/showcase", handlers.Showcase.List)
|
||||
publicAdmin.GET("/managed-models", handlers.AdminRead.Models)
|
||||
publicAdmin.GET("/stats", handlers.AdminRead.Stats)
|
||||
publicAdmin.GET("/video-presets", handlers.UserGen.VideoPresets)
|
||||
publicAdmin.GET("/catalog", handlers.UserGen.Catalog)
|
||||
publicAdmin.GET("/models", handlers.UserGen.Models)
|
||||
}
|
||||
|
||||
authGroup := engine.Group("/admin/api/auth")
|
||||
{
|
||||
authGroup.GET("/config", handlers.Auth.Config)
|
||||
authGroup.POST("/send-code", handlers.Auth.SendCode)
|
||||
authGroup.POST("/register", handlers.Auth.Register)
|
||||
authGroup.POST("/login", handlers.Auth.Login)
|
||||
authGroup.POST("/logout", handlers.Auth.Logout)
|
||||
authGroup.POST("/reset-password", handlers.Auth.ResetPassword)
|
||||
}
|
||||
|
||||
userAuthed := engine.Group("/admin/api")
|
||||
userAuthed.Use(middleware.RequireSession(auth))
|
||||
{
|
||||
userAuthed.GET("/logs", handlers.UserGen.Logs)
|
||||
userAuthed.POST("/generate", handlers.UserGen.Generate)
|
||||
userAuthed.POST("/test", handlers.UserGen.Test)
|
||||
userAuthed.GET("/jobs/mine", handlers.UserGen.MyJobs)
|
||||
userAuthed.GET("/my-images", handlers.UserGen.MyImages)
|
||||
}
|
||||
|
||||
authed := engine.Group("/admin/api")
|
||||
authed.Use(middleware.RequireAdminSession(auth))
|
||||
{
|
||||
authed.GET("/dashboard", handlers.AdminRead.Dashboard)
|
||||
authed.GET("/users", handlers.AdminRead.Users)
|
||||
authed.GET("/invites", handlers.AdminRead.Invites)
|
||||
authed.POST("/users", handlers.AdminWrite.CreateUser)
|
||||
authed.POST("/users/delete-bulk", handlers.AdminWrite.DeleteUsersBulk)
|
||||
authed.PATCH("/users/:user_id", handlers.AdminWrite.UpdateUser)
|
||||
authed.DELETE("/users/:user_id", handlers.AdminWrite.DeleteUser)
|
||||
authed.POST("/users/:user_id/credits", handlers.AdminWrite.AdjustUserCredits)
|
||||
authed.POST("/users/:user_id/api-keys", handlers.AdminWrite.CreateUserAPIKey)
|
||||
authed.DELETE("/users/:user_id/api-keys/:key_id", handlers.AdminWrite.DeleteUserAPIKey)
|
||||
authed.GET("/cdks", handlers.CDK.List)
|
||||
authed.POST("/cdks", handlers.CDK.Create)
|
||||
authed.POST("/cdks/delete-bulk", handlers.CDK.DeleteBulk)
|
||||
authed.DELETE("/cdks/:code", handlers.CDK.Delete)
|
||||
authed.GET("/tokens", handlers.ProviderAdmin.TokensList)
|
||||
authed.POST("/tokens", handlers.ProviderAdmin.TokensCreate)
|
||||
authed.POST("/tokens/import-chatgpt-token", handlers.ProviderAdmin.ImportChatGPTToken)
|
||||
authed.POST("/tokens/import-adobe-cookie", handlers.ProviderAdmin.ImportAdobeCookie)
|
||||
authed.POST("/tokens/import-runway-token", handlers.ProviderAdmin.ImportRunwayToken)
|
||||
authed.POST("/tokens/import-leonardo-cookie", handlers.ProviderAdmin.ImportLeonardoCookie)
|
||||
authed.POST("/tokens/import-krea-cookie", handlers.ProviderAdmin.ImportKreaCookie)
|
||||
authed.POST("/tokens/import-imagine-token", handlers.ProviderAdmin.ImportImagineToken)
|
||||
authed.POST("/tokens/delete-bulk", handlers.ProviderAdmin.TokenDeleteBulk)
|
||||
authed.PATCH("/tokens/:pool/:id", handlers.ProviderAdmin.TokenUpdate)
|
||||
authed.DELETE("/tokens/:pool/:id", handlers.ProviderAdmin.TokenDelete)
|
||||
authed.GET("/accounts", handlers.ProviderAdmin.AccountsList)
|
||||
authed.GET("/accounts/:pool/:id/quota", handlers.ProviderAdmin.AccountQuota)
|
||||
authed.GET("/accounts/:pool/:id/email", handlers.ProviderAdmin.AccountEmail)
|
||||
authed.GET("/providers", handlers.AdminRead.Providers)
|
||||
authed.GET("/images", handlers.AdminRead.Images)
|
||||
authed.GET("/refresh/profiles", handlers.ProviderAdmin.RefreshProfiles)
|
||||
authed.POST("/refresh/profiles/:profile_id/refresh-now", handlers.ProviderAdmin.RefreshNow)
|
||||
authed.PATCH("/refresh/profiles/:profile_id", handlers.ProviderAdmin.RefreshUpdate)
|
||||
authed.DELETE("/refresh/profiles/:profile_id", handlers.ProviderAdmin.RefreshDelete)
|
||||
authed.POST("/managed-models", handlers.AdminWrite.CreateModel)
|
||||
authed.PATCH("/managed-models/:model_id", handlers.AdminWrite.UpdateModel)
|
||||
authed.DELETE("/managed-models/:model_id", handlers.AdminWrite.DeleteModel)
|
||||
authed.DELETE("/logs", handlers.AdminWrite.ClearLogs)
|
||||
authed.DELETE("/logs/pending", handlers.AdminWrite.ClearPendingLogs)
|
||||
authed.POST("/showcase", handlers.AdminWrite.CreateShowcase)
|
||||
authed.PATCH("/showcase/:entry_id", handlers.AdminWrite.UpdateShowcase)
|
||||
authed.DELETE("/showcase/:entry_id", handlers.AdminWrite.DeleteShowcase)
|
||||
|
||||
settings := authed.Group("/settings")
|
||||
{
|
||||
settings.GET("/site", handlers.SiteSettings.Get)
|
||||
settings.PUT("/site", handlers.SiteSettings.Put)
|
||||
settings.GET("/registration", handlers.AppSettings.RegistrationGet)
|
||||
settings.PUT("/registration", handlers.AppSettings.RegistrationPut)
|
||||
settings.GET("/smtp", handlers.AppSettings.SMTPGet)
|
||||
settings.PUT("/smtp", handlers.AppSettings.SMTPPut)
|
||||
settings.POST("/smtp/test", handlers.AppSettings.SMTPTest)
|
||||
settings.GET("/proxy", handlers.AppSettings.ProxyGet)
|
||||
settings.PUT("/proxy", handlers.AppSettings.ProxyPut)
|
||||
settings.POST("/proxy/test", handlers.AppSettings.ProxyTest)
|
||||
settings.GET("/credits", handlers.AppSettings.CreditsGet)
|
||||
settings.PUT("/credits", handlers.AppSettings.CreditsPut)
|
||||
settings.GET("/logs", handlers.AppSettings.LogsGet)
|
||||
settings.PUT("/logs", handlers.AppSettings.LogsPut)
|
||||
settings.GET("/media", handlers.AppSettings.MediaGet)
|
||||
settings.PUT("/media", handlers.AppSettings.MediaPut)
|
||||
}
|
||||
}
|
||||
|
||||
authGroup.Use(middleware.RequireSession(auth))
|
||||
{
|
||||
authGroup.GET("/me", handlers.Auth.Me)
|
||||
authGroup.GET("/invites", handlers.Auth.Invites)
|
||||
authGroup.POST("/checkin", handlers.Auth.Checkin)
|
||||
authGroup.POST("/change-password", handlers.Auth.ChangePassword)
|
||||
authGroup.GET("/api-key", handlers.UserTools.APIKeyGet)
|
||||
authGroup.POST("/api-key", handlers.UserTools.APIKeyMint)
|
||||
authGroup.DELETE("/api-key", handlers.UserTools.APIKeyDelete)
|
||||
authGroup.POST("/redeem-cdk", handlers.UserTools.RedeemCDK)
|
||||
}
|
||||
|
||||
return engine
|
||||
}
|
||||
@@ -0,0 +1,190 @@
|
||||
package model
|
||||
|
||||
import (
|
||||
"time"
|
||||
|
||||
"gorm.io/datatypes"
|
||||
)
|
||||
|
||||
type User struct {
|
||||
ID string `gorm:"primaryKey;size:32"`
|
||||
Email string `gorm:"size:255;uniqueIndex;not null"`
|
||||
Name string `gorm:"size:255"`
|
||||
PasswordHash string `gorm:"size:255"`
|
||||
Role string `gorm:"size:32;index;not null"`
|
||||
Status string `gorm:"size:32;index;not null"`
|
||||
Credits float64 `gorm:"not null;default:0"`
|
||||
Notes string `gorm:"type:text"`
|
||||
InviteCode string `gorm:"size:32;uniqueIndex"`
|
||||
InvitedBy *string `gorm:"size:32;index"`
|
||||
InviteRewardDone bool `gorm:"not null;default:false"`
|
||||
InviteRewardAt *time.Time
|
||||
CheckinLast string `gorm:"size:32"`
|
||||
CheckinStreak int `gorm:"not null;default:0"`
|
||||
LastLoginAt *time.Time
|
||||
LastLoginIP string `gorm:"size:128"`
|
||||
CreatedAt time.Time
|
||||
UpdatedAt time.Time
|
||||
APIKeys []APIKey `gorm:"foreignKey:UserID"`
|
||||
}
|
||||
|
||||
type APIKey struct {
|
||||
ID string `gorm:"primaryKey;size:32"`
|
||||
UserID string `gorm:"size:32;index;not null"`
|
||||
Name string `gorm:"size:100;not null"`
|
||||
KeyPreview string `gorm:"size:32;not null"`
|
||||
KeyHash string `gorm:"size:255;uniqueIndex;not null"`
|
||||
CreatedAt time.Time
|
||||
LastUsedAt *time.Time
|
||||
}
|
||||
|
||||
type ShowcaseItem struct {
|
||||
ID string `gorm:"primaryKey;size:32"`
|
||||
Kind string `gorm:"size:32;index;not null"`
|
||||
Title string `gorm:"size:255"`
|
||||
Subtitle string `gorm:"size:255"`
|
||||
Prompt string `gorm:"type:text"`
|
||||
Gradient string `gorm:"type:text"`
|
||||
Span string `gorm:"size:100"`
|
||||
Image string `gorm:"size:500;index"`
|
||||
Weight int `gorm:"not null;default:0"`
|
||||
CreatedAt time.Time
|
||||
UpdatedAt time.Time
|
||||
}
|
||||
|
||||
type EventLog struct {
|
||||
ID string `gorm:"primaryKey;size:32"`
|
||||
TS time.Time `gorm:"index;not null"`
|
||||
Kind string `gorm:"size:32;index;not null"`
|
||||
Status string `gorm:"size:32;index;not null"`
|
||||
Model string `gorm:"size:255;index"`
|
||||
Provider string `gorm:"size:100;index"`
|
||||
Prompt string `gorm:"type:text"`
|
||||
Ratio string `gorm:"size:32"`
|
||||
Resolution string `gorm:"size:32"`
|
||||
Duration string `gorm:"size:32"`
|
||||
Refs int `gorm:"not null;default:0"`
|
||||
RefFiles datatypes.JSON `gorm:"type:jsonb"` // relative paths of saved reference images, for回显 on reload
|
||||
Source string `gorm:"size:32;index"`
|
||||
// AccountID is the provider token/account chosen to fulfil this generation,
|
||||
// stamped when the upstream call begins. Drives the accounts view's live
|
||||
// in-flight count (pending events per account) and lets an abandoned-event
|
||||
// purge attribute the failure back to the account it was using.
|
||||
AccountID string `gorm:"size:64;index"`
|
||||
UserID string `gorm:"size:32;index"`
|
||||
Cost float64 `gorm:"not null;default:0"`
|
||||
// Refunded marks that this event's up-front charge has already been credited
|
||||
// back, so the normal failure path and the abandoned-purge sweep can never
|
||||
// double-refund the same generation.
|
||||
Refunded bool `gorm:"not null;default:false"`
|
||||
ElapsedMS int `gorm:"not null;default:0"`
|
||||
File string `gorm:"size:500;index"`
|
||||
Error string `gorm:"type:text"`
|
||||
CreatedAt time.Time
|
||||
UpdatedAt time.Time
|
||||
}
|
||||
|
||||
type ModelConfig struct {
|
||||
ID string `gorm:"primaryKey;size:255"`
|
||||
Type string `gorm:"size:32;index;not null"`
|
||||
Name string `gorm:"size:255;not null"`
|
||||
Provider string `gorm:"size:100;index;not null"`
|
||||
Enabled bool `gorm:"not null;default:true"`
|
||||
Ratios datatypes.JSON `gorm:"type:jsonb"`
|
||||
Prices datatypes.JSONMap `gorm:"type:jsonb"`
|
||||
Resolutions datatypes.JSON `gorm:"type:jsonb"`
|
||||
ImageToImage bool `gorm:"not null;default:false"`
|
||||
DurationPrices datatypes.JSONMap `gorm:"type:jsonb"`
|
||||
// Agent (代理) pricing — optional overlay over Prices/DurationPrices. A tier
|
||||
// left unset here means agent users pay the normal price for that tier; the
|
||||
// set of *supported* tiers is always driven by Prices, not these.
|
||||
PricesAgent datatypes.JSONMap `gorm:"type:jsonb;column:prices_agent"`
|
||||
DurationPricesAgent datatypes.JSONMap `gorm:"type:jsonb;column:duration_prices_agent"`
|
||||
Durations datatypes.JSON `gorm:"type:jsonb"`
|
||||
MaxReferenceImages int `gorm:"not null;default:0"`
|
||||
ReferenceMode string `gorm:"size:32;not null;default:'none'"`
|
||||
// Weight controls display order in the model dropdown / admin list: higher
|
||||
// weight floats to the top (matches ShowcaseItem.Weight semantics). Ties fall
|
||||
// back to created_at desc. Default 0.
|
||||
Weight int `gorm:"not null;default:0;index"`
|
||||
CreatedAt time.Time
|
||||
UpdatedAt time.Time
|
||||
}
|
||||
|
||||
type CDKCode struct {
|
||||
Code string `gorm:"primaryKey;size:32"`
|
||||
Amount int `gorm:"not null"`
|
||||
Status string `gorm:"size:32;index;not null"`
|
||||
Type string `gorm:"size:16;not null;default:normal;index"` // normal | marketing
|
||||
BatchID string `gorm:"size:32;index"` // groups one generate call
|
||||
Note string `gorm:"type:text"`
|
||||
RedeemedBy *string `gorm:"size:32;index"`
|
||||
RedeemedAt *time.Time
|
||||
CreatedAt time.Time
|
||||
UpdatedAt time.Time
|
||||
}
|
||||
|
||||
type TokenAccount struct {
|
||||
ID string `gorm:"primaryKey;size:64"`
|
||||
Pool string `gorm:"size:64;index;not null"`
|
||||
Value string `gorm:"type:text"`
|
||||
Status string `gorm:"size:32;index;not null"`
|
||||
Fails int `gorm:"not null;default:0"`
|
||||
FailTotal int `gorm:"not null;default:0"`
|
||||
SuccessTotal int `gorm:"not null;default:0"`
|
||||
Dead bool `gorm:"not null;default:false"`
|
||||
Meta datatypes.JSONMap `gorm:"type:jsonb"`
|
||||
AddedAt *time.Time
|
||||
LastUsedAt *time.Time
|
||||
CachedQuotaResetAfter string `gorm:"size:128"`
|
||||
QuotaRecoverAt *time.Time
|
||||
// Adobe quota is tracked separately for image vs video. An account only
|
||||
// enters the shared "quota" waiting status when BOTH are limited; a single
|
||||
// limit leaves the account usable for the other kind. Recovery time is shared
|
||||
// (QuotaRecoverAt / CachedQuotaResetAfter) since Adobe resets both at once.
|
||||
ImageLimited bool `gorm:"not null;default:false"`
|
||||
VideoLimited bool `gorm:"not null;default:false"`
|
||||
AccountEmail string `gorm:"size:255"`
|
||||
AccountDisplayName string `gorm:"size:255"`
|
||||
CreatedAt time.Time
|
||||
UpdatedAt time.Time
|
||||
}
|
||||
|
||||
type RefreshProfile struct {
|
||||
ID string `gorm:"primaryKey;size:64"`
|
||||
Name string `gorm:"size:255;not null"`
|
||||
Pool string `gorm:"size:64;index;not null"`
|
||||
Kind string `gorm:"size:64;index;not null"`
|
||||
Cookie string `gorm:"type:text"`
|
||||
Enabled bool `gorm:"not null;default:true"`
|
||||
IntervalSeconds int `gorm:"not null;default:54000"`
|
||||
ImportedAt *time.Time
|
||||
LastAttemptAt *time.Time
|
||||
LastSuccessAt *time.Time
|
||||
LastError string `gorm:"type:text"`
|
||||
NextRetryAt *time.Time
|
||||
ConsecutiveFailures int `gorm:"not null;default:0"`
|
||||
CreatedAt time.Time
|
||||
UpdatedAt time.Time
|
||||
}
|
||||
|
||||
type SiteSetting struct {
|
||||
Key string `gorm:"primaryKey;size:100"`
|
||||
Value string `gorm:"type:text"`
|
||||
CreatedAt time.Time
|
||||
UpdatedAt time.Time
|
||||
}
|
||||
|
||||
func AutoMigrateModels() []any {
|
||||
return []any{
|
||||
&User{},
|
||||
&APIKey{},
|
||||
&ShowcaseItem{},
|
||||
&EventLog{},
|
||||
&ModelConfig{},
|
||||
&CDKCode{},
|
||||
&TokenAccount{},
|
||||
&RefreshProfile{},
|
||||
&SiteSetting{},
|
||||
}
|
||||
}
|
||||
Binary file not shown.
@@ -0,0 +1,39 @@
|
||||
package adobe
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"net/http"
|
||||
"strings"
|
||||
)
|
||||
|
||||
const (
|
||||
refreshURL = "https://adobeid-na1.services.adobe.com/ims/check/v6/token?jslVersion=v2-v0.48.0-1-g1e322cb"
|
||||
clientID = "clio-playground-web"
|
||||
scopeValue = "AdobeID,firefly_api,openid,pps.read,pps.write,additional_info.projectedProductContext,additional_info.ownerOrg,uds_read,uds_write,ab.manage,read_organizations,additional_info.roles,account_cluster.read,creative_production,profile"
|
||||
)
|
||||
|
||||
var ErrAdobeCookieEmpty = errors.New("cookie is empty")
|
||||
|
||||
type CookieExchangeResult struct {
|
||||
AccessToken string
|
||||
ExpiresIn int
|
||||
Raw map[string]any
|
||||
}
|
||||
|
||||
func ExchangeCookieToAccessToken(ctx context.Context, client *http.Client, cookie string) (*CookieExchangeResult, error) {
|
||||
_ = client
|
||||
tlsClient, err := NewClient(clientID, "").newTLSClient()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return exchangeCookieWithTLSClient(ctx, tlsClient, cookie)
|
||||
}
|
||||
|
||||
func normalizeCookie(v string) string {
|
||||
v = strings.TrimSpace(v)
|
||||
if strings.HasPrefix(strings.ToLower(v), "cookie:") {
|
||||
v = strings.TrimSpace(v[len("cookie:"):])
|
||||
}
|
||||
return v
|
||||
}
|
||||
@@ -0,0 +1,875 @@
|
||||
package adobe
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/url"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
http "github.com/bogdanfinn/fhttp"
|
||||
tlsclient "github.com/bogdanfinn/tls-client"
|
||||
"github.com/bogdanfinn/tls-client/profiles"
|
||||
)
|
||||
|
||||
const (
|
||||
submitURL = "https://firefly-3p.ff.adobe.io/v2/3p-images/generate-async"
|
||||
image5SubmitURL = "https://image-v5.ff.adobe.io/v1/images/generate-async"
|
||||
videoSubmitURL = "https://firefly-3p.ff.adobe.io/v2/3p-videos/generate-async"
|
||||
// Firefly-native video model (project id "firefly-video"): distinct host,
|
||||
// submit path and storage host from the 3p (veo/luma) video flow.
|
||||
fireflyVideoSubmitURL = "https://video-v1.ff.adobe.io/v2/videos/generate"
|
||||
fireflyVideoUploadURL = "https://video-v1.ff.adobe.io/v2/storage/image"
|
||||
uploadURL = "https://firefly-3p.ff.adobe.io/v2/storage/image"
|
||||
creditsURL = "https://firefly.adobe.io/v1/credits/balance"
|
||||
creditsAPIKey = "SunbreakWebUI1"
|
||||
)
|
||||
|
||||
var (
|
||||
ErrAuth = errors.New("adobe auth failed")
|
||||
ErrQuotaExhausted = errors.New("adobe quota exhausted")
|
||||
ErrTemporaryUpstream = errors.New("adobe upstream temporary error")
|
||||
)
|
||||
|
||||
var profileURLs = []string{
|
||||
"https://ims-na1.adobelogin.com/ims/profile/v1",
|
||||
"https://adobeid-na1.services.adobe.com/ims/profile/v1",
|
||||
}
|
||||
|
||||
type Client struct {
|
||||
apiKey string
|
||||
proxy string
|
||||
}
|
||||
|
||||
func NewClient(apiKey, proxy string) *Client {
|
||||
return &Client{
|
||||
apiKey: defaultString(apiKey, clientID),
|
||||
proxy: strings.TrimSpace(proxy),
|
||||
}
|
||||
}
|
||||
|
||||
func (c *Client) SetProxy(proxy string) {
|
||||
c.proxy = strings.TrimSpace(proxy)
|
||||
}
|
||||
|
||||
func (c *Client) ExchangeCookie(ctx context.Context, cookie string) (*CookieExchangeResult, error) {
|
||||
client, err := c.newTLSClient()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return exchangeCookieWithTLSClient(ctx, client, cookie)
|
||||
}
|
||||
|
||||
func (c *Client) UploadImage(ctx context.Context, token string, content []byte, contentType, engine string) (string, error) {
|
||||
client, err := c.newTLSClient()
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
endpoint := uploadURL
|
||||
if engine == "firefly-video" {
|
||||
endpoint = fireflyVideoUploadURL
|
||||
}
|
||||
req, err := http.NewRequest(http.MethodPost, endpoint, bytes.NewReader(content))
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
req = req.WithContext(ctx)
|
||||
req.Header = http.Header{
|
||||
"authorization": {"Bearer " + strings.TrimSpace(token)},
|
||||
"x-api-key": {c.apiKey},
|
||||
"content-type": {defaultString(contentType, "image/png")},
|
||||
"accept": {"*/*"},
|
||||
"user-agent": {defaultUserAgent},
|
||||
http.HeaderOrderKey: {
|
||||
"authorization",
|
||||
"x-api-key",
|
||||
"content-type",
|
||||
"accept",
|
||||
"user-agent",
|
||||
},
|
||||
}
|
||||
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("adobe upload request: %w", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
body, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if resp.StatusCode == 401 || resp.StatusCode == 403 {
|
||||
return "", fmt.Errorf("%w (upload %d %s: %s)", ErrAuth, resp.StatusCode, resp.Header.Get("x-access-error"), clip(body, 300))
|
||||
}
|
||||
if resp.StatusCode != 200 {
|
||||
return "", fmt.Errorf("adobe upload failed: %d %s", resp.StatusCode, clip(body, 300))
|
||||
}
|
||||
|
||||
var payload map[string]any
|
||||
if err := json.Unmarshal(body, &payload); err != nil {
|
||||
return "", err
|
||||
}
|
||||
if images, ok := payload["images"].([]any); ok && len(images) > 0 {
|
||||
if first, ok := images[0].(map[string]any); ok {
|
||||
if id := strings.TrimSpace(stringValue(first["id"])); id != "" {
|
||||
return id, nil
|
||||
}
|
||||
}
|
||||
}
|
||||
if id := strings.TrimSpace(stringValue(payload["id"])); id != "" {
|
||||
return id, nil
|
||||
}
|
||||
return "", errors.New("adobe upload missing blob id")
|
||||
}
|
||||
|
||||
func (c *Client) GenerateImage(ctx context.Context, token, modelID, prompt, aspectRatio, resolution string, blobIDs []string) ([]byte, map[string]any, error) {
|
||||
client, err := c.newTLSClient()
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
|
||||
var lastBody []byte
|
||||
var lastErr error
|
||||
// Firefly Image 5 uses a different endpoint + request schema (modelVersion
|
||||
// "image5", resolutionLevel, top-level aspectRatio label, no modelId/size).
|
||||
endpoint := submitURL
|
||||
var candidates []map[string]any
|
||||
if modelID == "firefly-image-5" {
|
||||
endpoint = image5SubmitURL
|
||||
candidates = []map[string]any{buildImage5Payload(prompt, aspectRatio, resolution, blobIDs)}
|
||||
} else {
|
||||
candidates = BuildImagePayloadCandidates(modelID, prompt, aspectRatio, resolution, blobIDs)
|
||||
}
|
||||
for _, payload := range candidates {
|
||||
respBody, pollURL, err := c.submitImage(ctx, client, token, prompt, endpoint, payload)
|
||||
if err == nil {
|
||||
meta, data, pollErr := c.pollImage(ctx, client, token, pollURL)
|
||||
if pollErr != nil {
|
||||
return nil, nil, pollErr
|
||||
}
|
||||
return data, meta, nil
|
||||
}
|
||||
lastBody = respBody
|
||||
lastErr = err
|
||||
if errors.Is(err, ErrAuth) || errors.Is(err, ErrQuotaExhausted) {
|
||||
return nil, nil, err
|
||||
}
|
||||
}
|
||||
// Preserve the temporary classification so the pool retries (overload / 5xx /
|
||||
// rate-limit) instead of failing the request outright.
|
||||
if errors.Is(lastErr, ErrTemporaryUpstream) {
|
||||
return nil, nil, fmt.Errorf("%w: adobe submit: %s", ErrTemporaryUpstream, clip(lastBody, 300))
|
||||
}
|
||||
return nil, nil, fmt.Errorf("adobe submit failed: %s", clip(lastBody, 300))
|
||||
}
|
||||
|
||||
// GenerateVideo renders the clip and (when downloadResult) downloads the MP4.
|
||||
// With downloadResult=false it returns nil bytes and the upstream presigned URL
|
||||
// in meta["video_url"] — used by the async /v1/videos job, which proxies that URL
|
||||
// on /content instead of persisting the file.
|
||||
func (c *Client) GenerateVideo(ctx context.Context, token, engine, prompt, aspectRatio string, durationSeconds int, resolution, referenceMode, upstreamModel string, blobIDs []string, downloadResult bool) ([]byte, map[string]any, error) {
|
||||
client, err := c.newTLSClient()
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
|
||||
payload := BuildVideoPayload(engine, prompt, aspectRatio, durationSeconds, resolution, referenceMode, upstreamModel, blobIDs)
|
||||
endpoint := videoSubmitURL
|
||||
if engine == "firefly-video" {
|
||||
endpoint = fireflyVideoSubmitURL
|
||||
}
|
||||
respBody, pollURL, err := c.submitVideo(ctx, client, token, endpoint, payload)
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
_ = respBody
|
||||
meta, data, pollErr := c.pollVideo(ctx, client, token, pollURL, downloadResult)
|
||||
if pollErr != nil {
|
||||
return nil, nil, pollErr
|
||||
}
|
||||
return data, meta, nil
|
||||
}
|
||||
|
||||
func (c *Client) FetchAccountProfile(ctx context.Context, token string) (map[string]any, error) {
|
||||
token = strings.TrimSpace(token)
|
||||
if token == "" {
|
||||
return map[string]any{}, nil
|
||||
}
|
||||
client, err := c.newTLSClient()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
for _, rawURL := range profileURLs {
|
||||
req, err := http.NewRequest(http.MethodGet, rawURL, nil)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
req = req.WithContext(ctx)
|
||||
req.Header = http.Header{
|
||||
"authorization": {"Bearer " + token},
|
||||
"accept": {"application/json"},
|
||||
"user-agent": {defaultUserAgent},
|
||||
http.HeaderOrderKey: {
|
||||
"authorization",
|
||||
"accept",
|
||||
"user-agent",
|
||||
},
|
||||
}
|
||||
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
body, readErr := io.ReadAll(resp.Body)
|
||||
resp.Body.Close()
|
||||
if readErr != nil || resp.StatusCode != 200 {
|
||||
continue
|
||||
}
|
||||
|
||||
var payload map[string]any
|
||||
if err := json.Unmarshal(body, &payload); err != nil {
|
||||
continue
|
||||
}
|
||||
email := strings.TrimSpace(stringValue(payload["email"]))
|
||||
displayName := strings.TrimSpace(stringValue(payload["displayName"]))
|
||||
if displayName == "" {
|
||||
displayName = strings.TrimSpace(stringValue(payload["name"]))
|
||||
}
|
||||
if displayName == "" {
|
||||
displayName = strings.TrimSpace(stringValue(payload["fullName"]))
|
||||
}
|
||||
userID := strings.TrimSpace(stringValue(payload["userId"]))
|
||||
if userID == "" {
|
||||
userID = strings.TrimSpace(stringValue(payload["authId"]))
|
||||
}
|
||||
if email != "" || displayName != "" || userID != "" {
|
||||
return map[string]any{
|
||||
"email": emptyStringNil(email),
|
||||
"display_name": emptyStringNil(displayName),
|
||||
"user_id": emptyStringNil(userID),
|
||||
}, nil
|
||||
}
|
||||
}
|
||||
|
||||
return map[string]any{}, nil
|
||||
}
|
||||
|
||||
func (c *Client) FetchCreditsBalance(ctx context.Context, token string) (map[string]any, error) {
|
||||
token = strings.TrimSpace(token)
|
||||
if token == "" {
|
||||
return map[string]any{
|
||||
"remaining": nil,
|
||||
"used": nil,
|
||||
"total": nil,
|
||||
"available_until": nil,
|
||||
"unknown": true,
|
||||
"error": "empty token",
|
||||
}, nil
|
||||
}
|
||||
|
||||
accountID := ExtractAccountID(token)
|
||||
if accountID == "" {
|
||||
return map[string]any{
|
||||
"remaining": nil,
|
||||
"used": nil,
|
||||
"total": nil,
|
||||
"available_until": nil,
|
||||
"unknown": true,
|
||||
"error": "no account id",
|
||||
}, nil
|
||||
}
|
||||
|
||||
client, err := c.newTLSClient()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
req, err := http.NewRequest(http.MethodGet, creditsURL, nil)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
req = req.WithContext(ctx)
|
||||
req.Header = http.Header{
|
||||
"authorization": {"Bearer " + token},
|
||||
"x-api-key": {creditsAPIKey},
|
||||
"x-account-id": {accountID},
|
||||
"accept": {"application/json"},
|
||||
"content-type": {"application/json"},
|
||||
"user-agent": {defaultUserAgent},
|
||||
http.HeaderOrderKey: {
|
||||
"authorization",
|
||||
"x-api-key",
|
||||
"x-account-id",
|
||||
"accept",
|
||||
"content-type",
|
||||
"user-agent",
|
||||
},
|
||||
}
|
||||
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
return map[string]any{
|
||||
"remaining": nil,
|
||||
"used": nil,
|
||||
"total": nil,
|
||||
"available_until": nil,
|
||||
"unknown": true,
|
||||
"error": "network: " + err.Error(),
|
||||
}, nil
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
body, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if resp.StatusCode == 401 {
|
||||
return nil, ErrAuth
|
||||
}
|
||||
if resp.StatusCode != 200 {
|
||||
return map[string]any{
|
||||
"remaining": nil,
|
||||
"used": nil,
|
||||
"total": nil,
|
||||
"available_until": nil,
|
||||
"unknown": true,
|
||||
"error": fmt.Sprintf("http %d: %s", resp.StatusCode, clip(body, 160)),
|
||||
}, nil
|
||||
}
|
||||
|
||||
var payload map[string]any
|
||||
if err := json.Unmarshal(body, &payload); err != nil {
|
||||
return map[string]any{
|
||||
"remaining": nil,
|
||||
"used": nil,
|
||||
"total": nil,
|
||||
"available_until": nil,
|
||||
"unknown": true,
|
||||
"error": "non-json",
|
||||
}, nil
|
||||
}
|
||||
|
||||
totalInfo, _ := payload["total"].(map[string]any)
|
||||
quota, _ := totalInfo["quota"].(map[string]any)
|
||||
return map[string]any{
|
||||
"remaining": intOrNil(quota["available"]),
|
||||
"used": intOrNil(quota["used"]),
|
||||
"total": intOrNil(quota["total"]),
|
||||
"available_until": emptyStringNil(strings.TrimSpace(stringValue(totalInfo["availableUntil"]))),
|
||||
"unknown": false,
|
||||
"error": nil,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (c *Client) submitImage(ctx context.Context, client tlsclient.HttpClient, token, prompt, endpoint string, payload map[string]any) ([]byte, string, error) {
|
||||
body, err := json.Marshal(payload)
|
||||
if err != nil {
|
||||
return nil, "", err
|
||||
}
|
||||
req, err := http.NewRequest(http.MethodPost, endpoint, bytes.NewReader(body))
|
||||
if err != nil {
|
||||
return nil, "", err
|
||||
}
|
||||
req = req.WithContext(ctx)
|
||||
req.Header = http.Header{
|
||||
"authorization": {"Bearer " + strings.TrimSpace(token)},
|
||||
"x-api-key": {c.apiKey},
|
||||
"content-type": {"application/json"},
|
||||
"accept": {"*/*"},
|
||||
"origin": {"https://firefly.adobe.com"},
|
||||
"referer": {"https://firefly.adobe.com/"},
|
||||
"accept-language": {"en-US,en;q=0.9"},
|
||||
"sec-ch-ua": {defaultSecCHUA},
|
||||
"sec-ch-ua-mobile": {"?0"},
|
||||
"sec-ch-ua-platform": {`"Windows"`},
|
||||
"sec-fetch-site": {"same-site"},
|
||||
"sec-fetch-mode": {"cors"},
|
||||
"sec-fetch-dest": {"empty"},
|
||||
"user-agent": {defaultUserAgent},
|
||||
"x-arp-session-id": {buildARPSessionID()},
|
||||
http.HeaderOrderKey: {
|
||||
"authorization",
|
||||
"x-api-key",
|
||||
"content-type",
|
||||
"accept",
|
||||
"origin",
|
||||
"referer",
|
||||
"accept-language",
|
||||
"sec-ch-ua",
|
||||
"sec-ch-ua-mobile",
|
||||
"sec-ch-ua-platform",
|
||||
"sec-fetch-site",
|
||||
"sec-fetch-mode",
|
||||
"sec-fetch-dest",
|
||||
"user-agent",
|
||||
"x-nonce",
|
||||
"x-arp-session-id",
|
||||
},
|
||||
}
|
||||
if nonce := buildSubmitNonce(token, prompt); nonce != "" {
|
||||
req.Header.Set("x-nonce", nonce)
|
||||
}
|
||||
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
return nil, "", fmt.Errorf("%w: %v", ErrTemporaryUpstream, err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
respBody, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return nil, "", err
|
||||
}
|
||||
if resp.StatusCode == 401 || resp.StatusCode == 403 {
|
||||
if strings.EqualFold(resp.Header.Get("x-access-error"), "taste_exhausted") {
|
||||
return respBody, "", ErrQuotaExhausted
|
||||
}
|
||||
return respBody, "", fmt.Errorf("%w (submit %d %s: %s)", ErrAuth, resp.StatusCode, resp.Header.Get("x-access-error"), clip(respBody, 300))
|
||||
}
|
||||
if resp.StatusCode == 429 || resp.StatusCode == 451 || resp.StatusCode >= 500 {
|
||||
return respBody, "", ErrTemporaryUpstream
|
||||
}
|
||||
// "system under load" / timeout_error = adobe rate-limit/overload (can come on a
|
||||
// non-5xx) — treat as temporary so the pool retries instead of failing.
|
||||
if b := string(respBody); strings.Contains(b, "system under load") || strings.Contains(b, "timeout_error") {
|
||||
return respBody, "", ErrTemporaryUpstream
|
||||
}
|
||||
if resp.StatusCode != 200 {
|
||||
return respBody, "", errors.New("submit rejected")
|
||||
}
|
||||
|
||||
var payloadResp map[string]any
|
||||
if err := json.Unmarshal(respBody, &payloadResp); err != nil {
|
||||
return respBody, "", err
|
||||
}
|
||||
if override := strings.TrimSpace(resp.Header.Get("x-override-status-link")); override != "" {
|
||||
return respBody, override, nil
|
||||
}
|
||||
if links, ok := payloadResp["links"].(map[string]any); ok {
|
||||
if result, ok := links["result"].(map[string]any); ok {
|
||||
if href := strings.TrimSpace(stringValue(result["href"])); href != "" {
|
||||
return respBody, href, nil
|
||||
}
|
||||
}
|
||||
if href := strings.TrimSpace(stringValue(links["result"])); href != "" {
|
||||
return respBody, href, nil
|
||||
}
|
||||
}
|
||||
return respBody, "", errors.New("submit ok but no poll url")
|
||||
}
|
||||
|
||||
func (c *Client) pollImage(ctx context.Context, client tlsclient.HttpClient, token, pollURL string) (map[string]any, []byte, error) {
|
||||
start := time.Now()
|
||||
for {
|
||||
if time.Since(start) > 3*time.Minute {
|
||||
return nil, nil, errors.New("adobe generation timed out")
|
||||
}
|
||||
|
||||
req, err := http.NewRequest(http.MethodGet, pollURL, nil)
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
req = req.WithContext(ctx)
|
||||
req.Header = http.Header{
|
||||
"authorization": {"Bearer " + strings.TrimSpace(token)},
|
||||
"accept": {"*/*"},
|
||||
"origin": {"https://firefly.adobe.com"},
|
||||
"referer": {"https://firefly.adobe.com/"},
|
||||
"user-agent": {defaultUserAgent},
|
||||
http.HeaderOrderKey: {
|
||||
"authorization",
|
||||
"accept",
|
||||
"origin",
|
||||
"referer",
|
||||
"user-agent",
|
||||
},
|
||||
}
|
||||
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
return nil, nil, fmt.Errorf("%w: %v", ErrTemporaryUpstream, err)
|
||||
}
|
||||
body, readErr := io.ReadAll(resp.Body)
|
||||
resp.Body.Close()
|
||||
if readErr != nil {
|
||||
return nil, nil, readErr
|
||||
}
|
||||
if resp.StatusCode == 429 || resp.StatusCode == 451 || resp.StatusCode >= 500 {
|
||||
return nil, nil, ErrTemporaryUpstream
|
||||
}
|
||||
if resp.StatusCode != 200 {
|
||||
return nil, nil, fmt.Errorf("adobe poll failed: %d %s", resp.StatusCode, clip(body, 300))
|
||||
}
|
||||
|
||||
var payload map[string]any
|
||||
if err := json.Unmarshal(body, &payload); err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
if outputs, ok := payload["outputs"].([]any); ok && len(outputs) > 0 {
|
||||
if first, ok := outputs[0].(map[string]any); ok {
|
||||
if image, ok := first["image"].(map[string]any); ok {
|
||||
if url := strings.TrimSpace(stringValue(image["presignedUrl"])); url != "" {
|
||||
data, err := c.download(ctx, client, url)
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
return payload, data, nil
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
status := strings.ToUpper(strings.TrimSpace(stringValue(payload["status"])))
|
||||
if status == "FAILED" || status == "CANCELLED" || status == "ERROR" {
|
||||
return nil, nil, fmt.Errorf("adobe job failed: %s", clip(body, 300))
|
||||
}
|
||||
time.Sleep(3 * time.Second)
|
||||
}
|
||||
}
|
||||
|
||||
func (c *Client) submitVideo(ctx context.Context, client tlsclient.HttpClient, token, endpoint string, payload map[string]any) ([]byte, string, error) {
|
||||
body, err := json.Marshal(payload)
|
||||
if err != nil {
|
||||
return nil, "", err
|
||||
}
|
||||
req, err := http.NewRequest(http.MethodPost, endpoint, bytes.NewReader(body))
|
||||
if err != nil {
|
||||
return nil, "", err
|
||||
}
|
||||
req = req.WithContext(ctx)
|
||||
req.Header = http.Header{
|
||||
"authorization": {"Bearer " + strings.TrimSpace(token)},
|
||||
"x-api-key": {c.apiKey},
|
||||
"content-type": {"application/json"},
|
||||
"accept": {"*/*"},
|
||||
"origin": {"https://firefly.adobe.com"},
|
||||
"referer": {"https://firefly.adobe.com/"},
|
||||
"accept-language": {"en-US,en;q=0.9"},
|
||||
"sec-ch-ua": {defaultSecCHUA},
|
||||
"sec-ch-ua-mobile": {"?0"},
|
||||
"sec-ch-ua-platform": {`"Windows"`},
|
||||
"sec-fetch-site": {"same-site"},
|
||||
"sec-fetch-mode": {"cors"},
|
||||
"sec-fetch-dest": {"empty"},
|
||||
"user-agent": {defaultUserAgent},
|
||||
"x-arp-session-id": {buildARPSessionID()},
|
||||
http.HeaderOrderKey: {
|
||||
"authorization",
|
||||
"x-api-key",
|
||||
"content-type",
|
||||
"accept",
|
||||
"origin",
|
||||
"referer",
|
||||
"accept-language",
|
||||
"sec-ch-ua",
|
||||
"sec-ch-ua-mobile",
|
||||
"sec-ch-ua-platform",
|
||||
"sec-fetch-site",
|
||||
"sec-fetch-mode",
|
||||
"sec-fetch-dest",
|
||||
"user-agent",
|
||||
"x-nonce",
|
||||
"x-arp-session-id",
|
||||
},
|
||||
}
|
||||
// The working video submit (HAR) carries x-nonce just like the image submit.
|
||||
if prompt, _ := payload["prompt"].(string); prompt != "" {
|
||||
if nonce := buildSubmitNonce(token, prompt); nonce != "" {
|
||||
req.Header.Set("x-nonce", nonce)
|
||||
}
|
||||
}
|
||||
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
return nil, "", fmt.Errorf("%w: %v", ErrTemporaryUpstream, err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
respBody, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return nil, "", err
|
||||
}
|
||||
if resp.StatusCode == 401 || resp.StatusCode == 403 {
|
||||
if strings.EqualFold(resp.Header.Get("x-access-error"), "taste_exhausted") {
|
||||
return respBody, "", ErrQuotaExhausted
|
||||
}
|
||||
// Surface Adobe's response body — "adobe auth failed" alone hides whether
|
||||
// it's a bad token, a missing scope, or a WAF/fingerprint block.
|
||||
return respBody, "", fmt.Errorf("%w (%d %s: %s)", ErrAuth, resp.StatusCode, resp.Header.Get("x-access-error"), clip(respBody, 300))
|
||||
}
|
||||
if resp.StatusCode == 429 || resp.StatusCode == 451 || resp.StatusCode >= 500 {
|
||||
return respBody, "", ErrTemporaryUpstream
|
||||
}
|
||||
if resp.StatusCode != 200 {
|
||||
return respBody, "", fmt.Errorf("video submit rejected: %d %s", resp.StatusCode, clip(respBody, 300))
|
||||
}
|
||||
|
||||
var payloadResp map[string]any
|
||||
if err := json.Unmarshal(respBody, &payloadResp); err != nil {
|
||||
return respBody, "", err
|
||||
}
|
||||
if override := strings.TrimSpace(resp.Header.Get("x-override-status-link")); override != "" {
|
||||
return respBody, normalizeVideoPollURL(override), nil
|
||||
}
|
||||
if links, ok := payloadResp["links"].(map[string]any); ok {
|
||||
if result, ok := links["result"].(map[string]any); ok {
|
||||
if href := strings.TrimSpace(stringValue(result["href"])); href != "" {
|
||||
return respBody, normalizeVideoPollURL(href), nil
|
||||
}
|
||||
}
|
||||
if href := strings.TrimSpace(stringValue(links["result"])); href != "" {
|
||||
return respBody, normalizeVideoPollURL(href), nil
|
||||
}
|
||||
}
|
||||
return respBody, "", errors.New("video submit ok but no poll url")
|
||||
}
|
||||
|
||||
func (c *Client) pollVideo(ctx context.Context, client tlsclient.HttpClient, token, pollURL string, downloadResult bool) (map[string]any, []byte, error) {
|
||||
start := time.Now()
|
||||
for {
|
||||
if time.Since(start) > 10*time.Minute {
|
||||
return nil, nil, errors.New("adobe video generation timed out")
|
||||
}
|
||||
|
||||
req, err := http.NewRequest(http.MethodGet, pollURL, nil)
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
req = req.WithContext(ctx)
|
||||
req.Header = http.Header{
|
||||
"authorization": {"Bearer " + strings.TrimSpace(token)},
|
||||
"accept": {"*/*"},
|
||||
"origin": {"https://firefly.adobe.com"},
|
||||
"referer": {"https://firefly.adobe.com/"},
|
||||
"user-agent": {defaultUserAgent},
|
||||
http.HeaderOrderKey: {
|
||||
"authorization",
|
||||
"accept",
|
||||
"origin",
|
||||
"referer",
|
||||
"user-agent",
|
||||
},
|
||||
}
|
||||
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
return nil, nil, fmt.Errorf("%w: %v", ErrTemporaryUpstream, err)
|
||||
}
|
||||
body, readErr := io.ReadAll(resp.Body)
|
||||
resp.Body.Close()
|
||||
if readErr != nil {
|
||||
return nil, nil, readErr
|
||||
}
|
||||
if resp.StatusCode == 401 || resp.StatusCode == 403 {
|
||||
return nil, nil, fmt.Errorf("%w (%d %s: %s)", ErrAuth, resp.StatusCode, resp.Header.Get("x-access-error"), clip(body, 300))
|
||||
}
|
||||
if resp.StatusCode == 429 || resp.StatusCode == 451 || resp.StatusCode >= 500 {
|
||||
return nil, nil, ErrTemporaryUpstream
|
||||
}
|
||||
if resp.StatusCode != 200 {
|
||||
return nil, nil, fmt.Errorf("adobe video poll failed: %d %s", resp.StatusCode, clip(body, 300))
|
||||
}
|
||||
|
||||
var payload map[string]any
|
||||
if err := json.Unmarshal(body, &payload); err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
if outputs, ok := payload["outputs"].([]any); ok && len(outputs) > 0 {
|
||||
if first, ok := outputs[0].(map[string]any); ok {
|
||||
if video, ok := first["video"].(map[string]any); ok {
|
||||
if raw := strings.TrimSpace(stringValue(video["presignedUrl"])); raw != "" {
|
||||
payload["video_url"] = raw
|
||||
if !downloadResult {
|
||||
return payload, nil, nil
|
||||
}
|
||||
data, err := c.download(ctx, client, raw)
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
return payload, data, nil
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
status := strings.ToUpper(strings.TrimSpace(stringValue(payload["status"])))
|
||||
if status == "FAILED" || status == "CANCELLED" || status == "ERROR" {
|
||||
return nil, nil, fmt.Errorf("adobe video job failed: %s", clip(body, 300))
|
||||
}
|
||||
time.Sleep(3 * time.Second)
|
||||
}
|
||||
}
|
||||
|
||||
func (c *Client) download(ctx context.Context, client tlsclient.HttpClient, url string) ([]byte, error) {
|
||||
req, err := http.NewRequest(http.MethodGet, url, nil)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
req = req.WithContext(ctx)
|
||||
req.Header = http.Header{
|
||||
"accept": {"*/*"},
|
||||
"user-agent": {defaultUserAgent},
|
||||
http.HeaderOrderKey: {
|
||||
"accept",
|
||||
"user-agent",
|
||||
},
|
||||
}
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode != 200 {
|
||||
body, _ := io.ReadAll(resp.Body)
|
||||
return nil, fmt.Errorf("adobe download failed: %d %s", resp.StatusCode, clip(body, 200))
|
||||
}
|
||||
return io.ReadAll(resp.Body)
|
||||
}
|
||||
|
||||
func (c *Client) newTLSClient() (tlsclient.HttpClient, error) {
|
||||
options := []tlsclient.HttpClientOption{
|
||||
tlsclient.WithTimeoutSeconds(60),
|
||||
tlsclient.WithClientProfile(profiles.Chrome_133),
|
||||
tlsclient.WithNotFollowRedirects(),
|
||||
tlsclient.WithRandomTLSExtensionOrder(),
|
||||
}
|
||||
if c.proxy != "" {
|
||||
options = append(options, tlsclient.WithProxyUrl(c.proxy))
|
||||
}
|
||||
return tlsclient.NewHttpClient(tlsclient.NewNoopLogger(), options...)
|
||||
}
|
||||
|
||||
func exchangeCookieWithTLSClient(ctx context.Context, client tlsclient.HttpClient, cookie string) (*CookieExchangeResult, error) {
|
||||
cookie = normalizeCookie(cookie)
|
||||
if cookie == "" {
|
||||
return nil, ErrAdobeCookieEmpty
|
||||
}
|
||||
|
||||
body := "client_id=" + clientID + "&guest_allowed=true&scope=" + strings.ReplaceAll(scopeValue, ",", "%2C")
|
||||
req, err := http.NewRequest(http.MethodPost, refreshURL, strings.NewReader(body))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
req = req.WithContext(ctx)
|
||||
req.Header = http.Header{
|
||||
"accept": {"*/*"},
|
||||
"accept-language": {"zh-CN,zh;q=0.9"},
|
||||
"content-type": {"application/x-www-form-urlencoded;charset=UTF-8"},
|
||||
"cookie": {cookie},
|
||||
"origin": {"https://firefly.adobe.com"},
|
||||
"referer": {"https://firefly.adobe.com/"},
|
||||
"user-agent": {defaultUserAgent},
|
||||
http.HeaderOrderKey: {
|
||||
"accept",
|
||||
"accept-language",
|
||||
"content-type",
|
||||
"cookie",
|
||||
"origin",
|
||||
"referer",
|
||||
"user-agent",
|
||||
},
|
||||
}
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("adobe cookie exchange network error: %w", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
respBody, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if resp.StatusCode != 200 {
|
||||
return nil, fmt.Errorf("adobe cookie exchange upstream %d: %s", resp.StatusCode, clip(respBody, 200))
|
||||
}
|
||||
var payload map[string]any
|
||||
if err := json.Unmarshal(respBody, &payload); err != nil {
|
||||
return nil, fmt.Errorf("adobe cookie exchange invalid json: %w", err)
|
||||
}
|
||||
token := strings.TrimSpace(stringValue(payload["access_token"]))
|
||||
if token == "" {
|
||||
return nil, errors.New("adobe cookie exchange missing access_token")
|
||||
}
|
||||
return &CookieExchangeResult{
|
||||
AccessToken: token,
|
||||
ExpiresIn: intValue(payload["expires_in"]),
|
||||
Raw: payload,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func buildSubmitNonce(token, prompt string) string {
|
||||
claims := decodeJWTPayload(token)
|
||||
userID := strings.TrimSpace(stringValue(claims["user_id"]))
|
||||
if userID == "" {
|
||||
userID = strings.TrimSpace(stringValue(claims["aa_id"]))
|
||||
}
|
||||
if userID == "" {
|
||||
userID = strings.TrimSpace(stringValue(claims["sub"]))
|
||||
}
|
||||
prompt = strings.TrimSpace(prompt)
|
||||
if userID == "" || prompt == "" {
|
||||
return ""
|
||||
}
|
||||
if len(prompt) > 256 {
|
||||
prompt = prompt[:256]
|
||||
}
|
||||
sum := sha256.Sum256([]byte(userID + "-" + prompt))
|
||||
return hex.EncodeToString(sum[:])
|
||||
}
|
||||
|
||||
func ExtractAccountID(token string) string {
|
||||
claims := decodeJWTPayload(token)
|
||||
userID := strings.TrimSpace(stringValue(claims["user_id"]))
|
||||
if userID == "" {
|
||||
userID = strings.TrimSpace(stringValue(claims["aa_id"]))
|
||||
}
|
||||
if userID == "" {
|
||||
userID = strings.TrimSpace(stringValue(claims["sub"]))
|
||||
}
|
||||
return userID
|
||||
}
|
||||
|
||||
func normalizeVideoPollURL(raw string) string {
|
||||
if strings.TrimSpace(raw) == "" {
|
||||
return raw
|
||||
}
|
||||
parsed, err := url.Parse(raw)
|
||||
if err != nil {
|
||||
return raw
|
||||
}
|
||||
host := parsed.Hostname()
|
||||
if !strings.HasPrefix(host, "firefly-epo") {
|
||||
return raw
|
||||
}
|
||||
parts := strings.Split(strings.Trim(parsed.Path, "/"), "/")
|
||||
if len(parts) == 0 {
|
||||
return raw
|
||||
}
|
||||
jobID := strings.TrimSpace(parts[len(parts)-1])
|
||||
hostSuffix := strings.TrimPrefix(host, "firefly-epo")
|
||||
hostSuffix = strings.SplitN(hostSuffix, ".", 2)[0]
|
||||
if len(hostSuffix) != 4 {
|
||||
return raw
|
||||
}
|
||||
for _, ch := range hostSuffix {
|
||||
if ch < '0' || ch > '9' {
|
||||
return raw
|
||||
}
|
||||
}
|
||||
return "https://bks-epo" + hostSuffix + ".adobe.io/v2/jobs/result/" + jobID + "?host=" + parsed.Host + "/"
|
||||
}
|
||||
|
||||
func clip(v []byte, n int) string {
|
||||
s := strings.TrimSpace(string(v))
|
||||
if len(s) <= n {
|
||||
return s
|
||||
}
|
||||
return s[:n]
|
||||
}
|
||||
@@ -0,0 +1,466 @@
|
||||
package adobe
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
type modelSpec struct {
|
||||
UpstreamModelID string
|
||||
UpstreamModelVersion string
|
||||
}
|
||||
|
||||
var lumaSize = map[string]map[string][2]int{
|
||||
"720p": {
|
||||
"21:9": {1280, 548}, "16:9": {1280, 720}, "4:3": {960, 720},
|
||||
"1:1": {720, 720}, "3:4": {720, 960}, "9:16": {720, 1280}, "9:21": {548, 1280},
|
||||
},
|
||||
"1080p": {
|
||||
"21:9": {1920, 822}, "16:9": {1920, 1080}, "4:3": {1440, 1080},
|
||||
"1:1": {1080, 1080}, "3:4": {1080, 1440}, "9:16": {1080, 1920}, "9:21": {822, 1920},
|
||||
},
|
||||
"4k": {
|
||||
"21:9": {3840, 1646}, "16:9": {3840, 2160}, "4:3": {2880, 2160},
|
||||
"1:1": {2160, 2160}, "3:4": {2160, 2880}, "9:16": {2160, 3840}, "9:21": {1646, 3840},
|
||||
},
|
||||
}
|
||||
|
||||
var gptImageSize = map[string]map[string][2]int{
|
||||
"1K": {"1:1": {1024, 1024}, "5:4": {1120, 896}, "9:16": {720, 1280}, "21:9": {1456, 624}, "16:9": {1280, 720}, "4:3": {1152, 864}, "3:2": {1248, 832}, "4:5": {896, 1120}, "3:4": {864, 1152}, "2:3": {832, 1248}},
|
||||
"2K": {"1:1": {2048, 2048}, "5:4": {2240, 1792}, "9:16": {1440, 2560}, "21:9": {3024, 1296}, "16:9": {2560, 1440}, "4:3": {2304, 1728}, "3:2": {2496, 1664}, "4:5": {1792, 2240}, "3:4": {1728, 2304}, "2:3": {1664, 2496}},
|
||||
"4K": {"1:1": {2880, 2880}, "5:4": {3200, 2560}, "9:16": {2160, 3840}, "21:9": {3696, 1584}, "16:9": {3840, 2160}, "4:3": {3264, 2448}, "3:2": {3504, 2336}, "4:5": {2560, 3200}, "3:4": {2448, 3264}, "2:3": {2336, 3504}},
|
||||
}
|
||||
|
||||
var fluxSize = map[string][2]int{
|
||||
"1:1": {1024, 1024},
|
||||
"16:9": {1408, 768},
|
||||
"9:16": {768, 1408},
|
||||
"4:3": {1280, 896},
|
||||
"3:4": {896, 1280},
|
||||
}
|
||||
|
||||
var defaultSize = map[string]map[string][2]int{
|
||||
"1K": {"1:1": {1024, 1024}, "1:8": {384, 3072}, "1:4": {512, 2048}, "16:9": {1360, 768}, "9:16": {768, 1360}, "4:1": {2048, 512}, "4:3": {1152, 864}, "3:4": {864, 1152}, "8:1": {3072, 384}},
|
||||
"2K": {"1:1": {2048, 2048}, "1:8": {768, 6144}, "1:4": {1024, 4096}, "16:9": {2752, 1536}, "9:16": {1536, 2752}, "4:1": {4096, 1024}, "4:3": {2048, 1536}, "3:4": {1536, 2048}, "8:1": {6144, 768}},
|
||||
"4K": {"1:1": {4096, 4096}, "1:8": {1536, 12288}, "1:4": {2048, 8192}, "16:9": {5504, 3072}, "9:16": {3072, 5504}, "4:1": {8192, 2048}, "4:3": {4096, 3072}, "3:4": {3072, 4096}, "8:1": {12288, 1536}},
|
||||
}
|
||||
|
||||
func ResolveModelSpec(modelID string) modelSpec {
|
||||
switch modelID {
|
||||
case "firefly-gpt-image", "firefly-gpt-image-2":
|
||||
return modelSpec{UpstreamModelID: "gpt-image", UpstreamModelVersion: "2"}
|
||||
case "flux-kontext-max":
|
||||
return modelSpec{UpstreamModelID: "flux", UpstreamModelVersion: "fluxKontextMax"}
|
||||
default:
|
||||
return modelSpec{UpstreamModelID: "gemini-flash", UpstreamModelVersion: "nano-banana-3"}
|
||||
}
|
||||
}
|
||||
|
||||
// buildImage5Payload builds the Adobe Firefly Image 5 request. It uses a distinct
|
||||
// schema from the firefly-3p models: NO modelId/size, a top-level aspectRatio
|
||||
// string label and a resolutionLevel (1K→1MP, 2K→4MP). Mirrors a captured
|
||||
// working image-v5.ff.adobe.io request.
|
||||
func buildImage5Payload(prompt, aspectRatio, resolution string, blobIDs []string) map[string]any {
|
||||
p := map[string]any{
|
||||
"n": 1,
|
||||
"seeds": []int{int(time.Now().Unix()) % 999999},
|
||||
"output": map[string]any{"storeInputs": true},
|
||||
"prompt": prompt,
|
||||
"referenceBlobs": []any{},
|
||||
"modelSpecificPayload": map[string]any{"locale": "en-US", "prompt_reasoner": "quality"},
|
||||
"modelVersion": "image5",
|
||||
"resolutionLevel": image5ResolutionLevel(resolution),
|
||||
"generationMetadata": map[string]any{"module": "text2image", "submodule": "ff-image-generate"},
|
||||
}
|
||||
if len(blobIDs) > 0 {
|
||||
// Instruct-edit: aspect ratio is derived from the reference image; sending
|
||||
// aspectRatio is rejected with a validation_error.
|
||||
p["referenceBlobs"] = blobRefs(blobIDs, "general")
|
||||
} else {
|
||||
p["aspectRatio"] = defaultString(aspectRatio, "1:1")
|
||||
}
|
||||
return p
|
||||
}
|
||||
|
||||
// image5ResolutionLevel maps the UI resolution tier to Image 5's megapixel level.
|
||||
func image5ResolutionLevel(resolution string) string {
|
||||
switch strings.ToUpper(strings.TrimSpace(resolution)) {
|
||||
case "1K":
|
||||
return "1MP"
|
||||
case "2K":
|
||||
return "4MP"
|
||||
default:
|
||||
return "4MP"
|
||||
}
|
||||
}
|
||||
|
||||
func BuildImagePayloadCandidates(modelID, prompt, aspectRatio, outputResolution string, blobIDs []string) []map[string]any {
|
||||
spec := ResolveModelSpec(modelID)
|
||||
ratio := defaultString(aspectRatio, "1:1")
|
||||
resolution := defaultString(outputResolution, "2K")
|
||||
|
||||
switch spec.UpstreamModelID {
|
||||
case "gpt-image":
|
||||
return buildGPTImagePayloads(spec, prompt, ratio, resolution, blobIDs)
|
||||
case "flux":
|
||||
return buildFluxPayloads(spec, prompt, ratio, blobIDs)
|
||||
default:
|
||||
return buildDefaultPayloads(spec, prompt, ratio, resolution, blobIDs)
|
||||
}
|
||||
}
|
||||
|
||||
func buildGPTImagePayloads(spec modelSpec, prompt, ratio, resolution string, blobIDs []string) []map[string]any {
|
||||
size := getSize(gptImageSize, resolution, ratio, "1:1")
|
||||
// Mirrors the captured working gpt-image request shape: modelSpecificPayload.size,
|
||||
// generationSettings.detailLevel 3, and NO top-level size / outputResolution
|
||||
// (sending those got 403). Keeps the chosen size via modelSpecificPayload.size
|
||||
// ("WxH") rather than "auto".
|
||||
base := map[string]any{
|
||||
"modelId": spec.UpstreamModelID,
|
||||
"modelVersion": spec.UpstreamModelVersion,
|
||||
"n": 1,
|
||||
"prompt": prompt,
|
||||
"seeds": []int{int(time.Now().Unix()) % 999999},
|
||||
"output": map[string]any{"storeInputs": true},
|
||||
"referenceBlobs": []any{},
|
||||
"generationMetadata": map[string]any{"module": "text2image", "submodule": "ff-image-generate"},
|
||||
"modelSpecificPayload": map[string]any{"size": sizeString(size)},
|
||||
"generationSettings": map[string]any{"detailLevel": 3},
|
||||
}
|
||||
if len(blobIDs) == 0 {
|
||||
return []map[string]any{base}
|
||||
}
|
||||
subject := cloneMap(base)
|
||||
subject["referenceBlobs"] = blobRefs(blobIDs, "subject")
|
||||
return []map[string]any{subject}
|
||||
}
|
||||
|
||||
func buildFluxPayloads(spec modelSpec, prompt, ratio string, blobIDs []string) []map[string]any {
|
||||
size := fluxSize[ratio]
|
||||
if size == [2]int{} {
|
||||
size = fluxSize["1:1"]
|
||||
}
|
||||
base := map[string]any{
|
||||
"modelId": spec.UpstreamModelID,
|
||||
"modelVersion": spec.UpstreamModelVersion,
|
||||
"n": 1,
|
||||
"prompt": prompt,
|
||||
"size": map[string]any{"width": size[0], "height": size[1]},
|
||||
"seeds": []int{int(time.Now().Unix()) % 999999},
|
||||
"output": map[string]any{"storeInputs": true},
|
||||
"referenceBlobs": []any{},
|
||||
"modelSpecificPayload": map[string]any{
|
||||
"prompt_upsampling": true,
|
||||
"safety_tolerance": 2,
|
||||
"aspect_ratio": ratio,
|
||||
},
|
||||
"generationMetadata": map[string]any{"module": "text2image", "submodule": "ff-image-generate"},
|
||||
}
|
||||
if len(blobIDs) == 0 {
|
||||
return []map[string]any{base}
|
||||
}
|
||||
edited := cloneMap(base)
|
||||
edited["generationMetadata"] = map[string]any{"module": "image2image", "submodule": "ff-image-generate"}
|
||||
edited["referenceBlobs"] = blobRefs(blobIDs, "general")
|
||||
return []map[string]any{edited}
|
||||
}
|
||||
|
||||
func buildDefaultPayloads(spec modelSpec, prompt, ratio, resolution string, blobIDs []string) []map[string]any {
|
||||
size := getSize(defaultSize, resolution, ratio, "16:9")
|
||||
// Shape mirrors a captured working firefly.adobe.com request exactly: top-level
|
||||
// size object, modelSpecificPayload only {parameters:{addWatermark:false}},
|
||||
// groundSearch:false, module "text2image" (even with a reference blob). NO
|
||||
// skipCai and NO modelSpecificPayload.aspectRatio — sending those got 403.
|
||||
base := map[string]any{
|
||||
"modelId": spec.UpstreamModelID,
|
||||
"modelVersion": spec.UpstreamModelVersion,
|
||||
"n": 1,
|
||||
"prompt": prompt,
|
||||
"size": map[string]any{"width": size[0], "height": size[1]},
|
||||
"seeds": []int{int(time.Now().Unix()) % 999999},
|
||||
"groundSearch": false,
|
||||
"output": map[string]any{"storeInputs": true},
|
||||
"generationMetadata": map[string]any{
|
||||
"module": "text2image",
|
||||
"submodule": "ff-image-generate",
|
||||
},
|
||||
"modelSpecificPayload": map[string]any{
|
||||
"parameters": map[string]any{"addWatermark": false},
|
||||
},
|
||||
}
|
||||
if len(blobIDs) == 0 {
|
||||
base["referenceBlobs"] = []any{}
|
||||
return []map[string]any{base}
|
||||
}
|
||||
edited := cloneMap(base)
|
||||
edited["referenceBlobs"] = blobRefs(blobIDs, "general")
|
||||
return []map[string]any{edited}
|
||||
}
|
||||
|
||||
func getSize(table map[string]map[string][2]int, resolution, ratio, fallbackRatio string) [2]int {
|
||||
level := defaultString(resolution, "2K")
|
||||
levelTable, ok := table[level]
|
||||
if !ok {
|
||||
levelTable = table["2K"]
|
||||
}
|
||||
size, ok := levelTable[ratio]
|
||||
if !ok {
|
||||
size = levelTable[fallbackRatio]
|
||||
}
|
||||
return size
|
||||
}
|
||||
|
||||
func sizeString(size [2]int) string {
|
||||
return itoa(size[0]) + "x" + itoa(size[1])
|
||||
}
|
||||
|
||||
func blobRefs(ids []string, usage string) []any {
|
||||
out := make([]any, 0, len(ids))
|
||||
for _, id := range ids {
|
||||
out = append(out, map[string]any{"id": id, "usage": usage})
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func referenceImagesByID(ids []string) []any {
|
||||
out := make([]any, 0, len(ids))
|
||||
for _, id := range ids {
|
||||
out = append(out, map[string]any{"id": id})
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func referenceImagesByLocal(ids []string) []any {
|
||||
out := make([]any, 0, len(ids))
|
||||
for _, id := range ids {
|
||||
out = append(out, map[string]any{"localBlobRef": id})
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func cloneMap(in map[string]any) map[string]any {
|
||||
out := make(map[string]any, len(in))
|
||||
for k, v := range in {
|
||||
out[k] = v
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func BuildVideoPayload(engine, prompt, aspectRatio string, durationSeconds int, resolution, referenceMode, upstreamModel string, blobIDs []string) map[string]any {
|
||||
seedVal := int(time.Now().Unix()) % 999999
|
||||
engine = defaultString(engine, "sora2")
|
||||
resolution = defaultString(resolution, "720p")
|
||||
aspectRatio = defaultString(aspectRatio, "16:9")
|
||||
if durationSeconds <= 0 {
|
||||
durationSeconds = 5
|
||||
}
|
||||
|
||||
switch engine {
|
||||
case "firefly-video":
|
||||
// Firefly-native video model — a distinct schema (mirrors a captured
|
||||
// working video-v1.ff.adobe.io request): sizes[] carries width/height +
|
||||
// numFrames (numFrames encodes duration, ~25.6fps so 5s = 128), and
|
||||
// reference frames go under image.conditions with placement.start
|
||||
// (0 = first frame / 首帧, 1 = last frame / 末帧). NO modelId / version /
|
||||
// engine / duration / referenceBlobs fields.
|
||||
w, h, frames := fireflyVideoSize(aspectRatio, resolution, durationSeconds)
|
||||
payload := map[string]any{
|
||||
"addOnTransparentBackground": false,
|
||||
"prompt": prompt,
|
||||
"seeds": []int{seedVal},
|
||||
"sizes": []any{map[string]any{"width": w, "height": h, "numFrames": frames}},
|
||||
"videoSettings": map[string]any{},
|
||||
"locale": "en-US",
|
||||
"generationMetadata": map[string]any{"module": "text2video", "submodule": "ff-video-generate"},
|
||||
"output": map[string]any{"storeInputs": true},
|
||||
}
|
||||
if len(blobIDs) > 0 {
|
||||
conds := make([]any, 0, 2)
|
||||
conds = append(conds, map[string]any{
|
||||
"source": map[string]any{"id": blobIDs[0]},
|
||||
"placement": map[string]any{"start": 0},
|
||||
})
|
||||
if len(blobIDs) > 1 {
|
||||
conds = append(conds, map[string]any{
|
||||
"source": map[string]any{"id": blobIDs[1]},
|
||||
"placement": map[string]any{"start": 1},
|
||||
})
|
||||
}
|
||||
payload["image"] = map[string]any{"conditions": conds}
|
||||
}
|
||||
return payload
|
||||
case "veo31-fast", "veo31-standard":
|
||||
modelVersion := "3.1-fast-generate"
|
||||
if engine == "veo31-standard" {
|
||||
modelVersion = "3.1-generate"
|
||||
}
|
||||
// Shape mirrors a captured working firefly.adobe.com video request: flat
|
||||
// top-level duration / negativePrompt / generateAudio, submodule set, and
|
||||
// NO `n` / NO modelSpecificPayload (sending those got 403).
|
||||
payload := map[string]any{
|
||||
"modelId": "veo",
|
||||
"modelVersion": modelVersion,
|
||||
"size": videoSize(aspectRatio, resolution),
|
||||
"seeds": []int{seedVal},
|
||||
"prompt": prompt,
|
||||
"negativePrompt": "",
|
||||
"duration": durationSeconds,
|
||||
"generateAudio": false,
|
||||
"generationMetadata": map[string]any{
|
||||
"module": "text2video",
|
||||
"submodule": "ff-video-generate",
|
||||
},
|
||||
"output": map[string]any{"storeInputs": true},
|
||||
"referenceBlobs": []any{},
|
||||
}
|
||||
if len(blobIDs) > 0 {
|
||||
payload["generationMetadata"] = map[string]any{"module": "image2video", "submodule": "ff-video-generate"}
|
||||
refs := make([]any, 0, min(len(blobIDs), 2))
|
||||
for idx, id := range blobIDs[:min(len(blobIDs), 2)] {
|
||||
refs = append(refs, map[string]any{"id": id, "usage": "general", "promptReference": idx + 1})
|
||||
}
|
||||
payload["referenceBlobs"] = refs
|
||||
}
|
||||
return payload
|
||||
case "luma":
|
||||
payload := map[string]any{
|
||||
"modelId": "luma",
|
||||
"modelVersion": "3.14-ray",
|
||||
"size": lumaVideoSize(aspectRatio, resolution),
|
||||
"mode": "flex_2",
|
||||
"prompt": prompt,
|
||||
"negativePrompt": "",
|
||||
"duration": durationSeconds,
|
||||
"generationMetadata": map[string]any{
|
||||
"module": "text2video",
|
||||
"submodule": "ff-video-generate",
|
||||
},
|
||||
"modelSpecificPayload": map[string]any{
|
||||
"resolution": strings.ToLower(resolution),
|
||||
"aspect_ratio": aspectRatio,
|
||||
},
|
||||
"output": map[string]any{"storeInputs": true},
|
||||
}
|
||||
if len(blobIDs) > 0 {
|
||||
payload["generationMetadata"] = map[string]any{
|
||||
"module": "image2video",
|
||||
"submodule": "ff-video-generate",
|
||||
}
|
||||
refs := make([]any, 0, min(len(blobIDs), 2))
|
||||
for idx, id := range blobIDs[:min(len(blobIDs), 2)] {
|
||||
refs = append(refs, map[string]any{"id": id, "usage": "frame", "order": idx + 1})
|
||||
}
|
||||
payload["referenceBlobs"] = refs
|
||||
}
|
||||
return payload
|
||||
default:
|
||||
upstream := defaultString(upstreamModel, "openai:firefly:colligo:sora2")
|
||||
payload := map[string]any{
|
||||
"n": 1,
|
||||
"seeds": []int{seedVal},
|
||||
"modelId": "sora",
|
||||
"modelVersion": "sora-2",
|
||||
"size": videoSize(aspectRatio, resolution),
|
||||
"duration": durationSeconds,
|
||||
"fps": 24,
|
||||
"prompt": buildVideoPromptJSON(prompt, durationSeconds),
|
||||
"generationMetadata": map[string]any{"module": "text2video"},
|
||||
"model": upstream,
|
||||
"generateAudio": true,
|
||||
"generateLoop": false,
|
||||
"transparentBackground": false,
|
||||
"seed": itoa(seedVal),
|
||||
"locale": "en-US",
|
||||
"camera": map[string]any{"angle": "none", "shotSize": "none", "motion": nil, "promptStyle": nil},
|
||||
"negativePrompt": "",
|
||||
"jobMode": "standard",
|
||||
"debugGenerationEndpoint": "",
|
||||
"referenceBlobs": []any{},
|
||||
"referenceFrames": []any{},
|
||||
"referenceVideo": nil,
|
||||
"cameraMotionReferenceVideo": nil,
|
||||
"characterReference": nil,
|
||||
"editReferenceVideo": nil,
|
||||
"output": map[string]any{"storeInputs": true},
|
||||
}
|
||||
if len(blobIDs) > 0 {
|
||||
firstID := blobIDs[0]
|
||||
payload["generationMetadata"] = map[string]any{"module": "image2video"}
|
||||
payload["referenceBlobs"] = []any{
|
||||
map[string]any{"id": firstID, "usage": "general", "promptReference": 1},
|
||||
}
|
||||
payload["referenceFrames"] = []any{map[string]any{"localBlobRef": firstID}, nil}
|
||||
}
|
||||
return payload
|
||||
}
|
||||
}
|
||||
|
||||
// fireflyVideoSizeTable maps the firefly-video resolution tier + aspect ratio to
|
||||
// pixel dimensions. Only 1080p 9:16 (1080x1920) is HAR-confirmed; the rest follow
|
||||
// the standard 540p/720p/1080p grid for each ratio.
|
||||
var fireflyVideoSizeTable = map[string]map[string][2]int{
|
||||
"540p": {"16:9": {960, 540}, "1:1": {540, 540}, "9:16": {540, 960}},
|
||||
"720p": {"16:9": {1280, 720}, "1:1": {720, 720}, "9:16": {720, 1280}},
|
||||
"1080p": {"16:9": {1920, 1080}, "1:1": {1080, 1080}, "9:16": {1080, 1920}},
|
||||
}
|
||||
|
||||
// fireflyVideoSize returns width, height and numFrames. numFrames encodes the
|
||||
// clip length (~25.6fps; 5s = 128 frames, HAR-confirmed).
|
||||
func fireflyVideoSize(aspectRatio, resolution string, durationSeconds int) (int, int, int) {
|
||||
table, ok := fireflyVideoSizeTable[strings.ToLower(defaultString(resolution, "1080p"))]
|
||||
if !ok {
|
||||
table = fireflyVideoSizeTable["1080p"]
|
||||
}
|
||||
wh, ok := table[defaultString(aspectRatio, "9:16")]
|
||||
if !ok {
|
||||
wh = table["9:16"]
|
||||
}
|
||||
frames := durationSeconds * 128 / 5
|
||||
if frames <= 0 {
|
||||
frames = 128
|
||||
}
|
||||
return wh[0], wh[1], frames
|
||||
}
|
||||
|
||||
func videoSize(aspectRatio, resolution string) map[string]any {
|
||||
if strings.EqualFold(resolution, "1080p") {
|
||||
if aspectRatio == "16:9" {
|
||||
return map[string]any{"width": 1920, "height": 1080}
|
||||
}
|
||||
return map[string]any{"width": 1080, "height": 1920}
|
||||
}
|
||||
if aspectRatio == "16:9" {
|
||||
return map[string]any{"width": 1280, "height": 720}
|
||||
}
|
||||
return map[string]any{"width": 720, "height": 1280}
|
||||
}
|
||||
|
||||
func lumaVideoSize(aspectRatio, resolution string) map[string]any {
|
||||
table, ok := lumaSize[strings.ToLower(defaultString(resolution, "720p"))]
|
||||
if !ok {
|
||||
table = lumaSize["720p"]
|
||||
}
|
||||
size, ok := table[defaultString(aspectRatio, "16:9")]
|
||||
if !ok {
|
||||
size = table["16:9"]
|
||||
}
|
||||
return map[string]any{"width": size[0], "height": size[1]}
|
||||
}
|
||||
|
||||
func buildVideoPromptJSON(prompt string, durationSeconds int) string {
|
||||
payload := map[string]any{
|
||||
"id": 1,
|
||||
"duration_sec": durationSeconds,
|
||||
"prompt_text": prompt,
|
||||
}
|
||||
b, _ := json.Marshal(payload)
|
||||
return string(b)
|
||||
}
|
||||
|
||||
func min(a, b int) int {
|
||||
if a < b {
|
||||
return a
|
||||
}
|
||||
return b
|
||||
}
|
||||
@@ -0,0 +1,144 @@
|
||||
package adobe
|
||||
|
||||
import (
|
||||
"encoding/base64"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"os"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
const (
|
||||
defaultUserAgent = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/145.0.0.0 Safari/537.36"
|
||||
defaultSecCHUA = `"Not:A-Brand";v="99", "Google Chrome";v="145", "Chromium";v="145"`
|
||||
)
|
||||
|
||||
func stringValue(v any) string {
|
||||
switch x := v.(type) {
|
||||
case string:
|
||||
return x
|
||||
case nil:
|
||||
return ""
|
||||
default:
|
||||
return strings.TrimSpace(strings.ReplaceAll(toJSONScalar(x), "\n", " "))
|
||||
}
|
||||
}
|
||||
|
||||
func toJSONScalar(v any) string {
|
||||
b, err := json.Marshal(v)
|
||||
if err != nil {
|
||||
return ""
|
||||
}
|
||||
return string(b)
|
||||
}
|
||||
|
||||
func intValue(v any) int {
|
||||
switch x := v.(type) {
|
||||
case int:
|
||||
return x
|
||||
case int64:
|
||||
return int(x)
|
||||
case float64:
|
||||
return int(x)
|
||||
case float32:
|
||||
return int(x)
|
||||
case json.Number:
|
||||
n, _ := x.Int64()
|
||||
return int(n)
|
||||
case string:
|
||||
n, _ := strconv.Atoi(strings.TrimSpace(x))
|
||||
return n
|
||||
default:
|
||||
return 0
|
||||
}
|
||||
}
|
||||
|
||||
func defaultString(v, fallback string) string {
|
||||
v = strings.TrimSpace(v)
|
||||
if v == "" {
|
||||
return fallback
|
||||
}
|
||||
return v
|
||||
}
|
||||
|
||||
func itoa(v int) string {
|
||||
return strconv.Itoa(v)
|
||||
}
|
||||
|
||||
func decodeJWTPayload(token string) map[string]any {
|
||||
parts := strings.Split(strings.TrimSpace(token), ".")
|
||||
if len(parts) < 2 {
|
||||
return map[string]any{}
|
||||
}
|
||||
raw, err := base64.RawURLEncoding.DecodeString(parts[1])
|
||||
if err != nil {
|
||||
return map[string]any{}
|
||||
}
|
||||
var out map[string]any
|
||||
if err := json.Unmarshal(raw, &out); err != nil {
|
||||
return map[string]any{}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func buildARPSessionID() string {
|
||||
raw := map[string]any{
|
||||
"sid": uuid.NewString(),
|
||||
"ftr": randomHex(16) + "_" + strconv.FormatInt(time.Now().UnixMilli(), 10) + "_" + strconv.Itoa(os.Getpid()) + "_dUAL43-mnts-ants-d4_31ck__tt",
|
||||
}
|
||||
b, _ := json.Marshal(raw)
|
||||
return base64.StdEncoding.EncodeToString(b)
|
||||
}
|
||||
|
||||
func randomHex(n int) string {
|
||||
if n <= 0 {
|
||||
return ""
|
||||
}
|
||||
buf := make([]byte, n)
|
||||
now := time.Now().UnixNano()
|
||||
for i := range buf {
|
||||
buf[i] = byte(now >> ((i % 8) * 8))
|
||||
}
|
||||
return hex.EncodeToString(buf)
|
||||
}
|
||||
|
||||
func intOrNil(v any) any {
|
||||
switch x := v.(type) {
|
||||
case nil:
|
||||
return nil
|
||||
case int:
|
||||
return x
|
||||
case int64:
|
||||
return int(x)
|
||||
case float64:
|
||||
return int(x)
|
||||
case float32:
|
||||
return int(x)
|
||||
case json.Number:
|
||||
n, err := x.Int64()
|
||||
if err != nil {
|
||||
return nil
|
||||
}
|
||||
return int(n)
|
||||
case string:
|
||||
n, err := strconv.Atoi(strings.TrimSpace(x))
|
||||
if err != nil {
|
||||
return nil
|
||||
}
|
||||
return n
|
||||
default:
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
func emptyStringNil(v string) any {
|
||||
v = strings.TrimSpace(v)
|
||||
if v == "" {
|
||||
return nil
|
||||
}
|
||||
return v
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,151 @@
|
||||
package chatgpt
|
||||
|
||||
import (
|
||||
"encoding/base64"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"math/rand"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"golang.org/x/crypto/sha3"
|
||||
)
|
||||
|
||||
var (
|
||||
cores = []int{8, 16, 24, 32}
|
||||
documentKeys = []string{"__reactContainer$fzelfjyxej8", "_reactListening5dehydibo78", "location"}
|
||||
screenResolutions = [][2]int{{1920, 1080}, {1440, 900}, {2560, 1440}, {3840, 2160}}
|
||||
navKeys = []string{
|
||||
"registerProtocolHandler−function registerProtocolHandler() { [native code] }",
|
||||
"storage−[object StorageManager]",
|
||||
"locks−[object LockManager]",
|
||||
"appCodeName−Mozilla",
|
||||
"permissions−[object Permissions]",
|
||||
"share−function share() { [native code] }",
|
||||
"webdriver−false",
|
||||
"vendor−Google Inc.",
|
||||
"mediaDevices−[object MediaDevices]",
|
||||
"cookieEnabled−true",
|
||||
"onLine−true",
|
||||
"mimeTypes−[object MimeTypeArray]",
|
||||
"credentials−[object CredentialsContainer]",
|
||||
"serviceWorker−[object ServiceWorkerContainer]",
|
||||
"keyboard−[object Keyboard]",
|
||||
"gpu−[object GPU]",
|
||||
"doNotTrack",
|
||||
"language−zh-CN",
|
||||
"geolocation−[object Geolocation]",
|
||||
"hardwareConcurrency−32",
|
||||
}
|
||||
winKeys = []string{
|
||||
"0", "window", "self", "document", "name", "location", "history",
|
||||
"navigation", "innerWidth", "innerHeight", "screen", "chrome",
|
||||
"navigator", "performance", "crypto", "indexedDB", "sessionStorage",
|
||||
"localStorage", "fetch", "matchMedia", "postMessage", "setTimeout",
|
||||
"caches", "__NEXT_DATA__",
|
||||
}
|
||||
)
|
||||
|
||||
func buildLegacyRequirementsToken(userAgent string, scriptSources []string, dataBuild string) string {
|
||||
cfg := buildPOWConfig(userAgent, scriptSources, dataBuild)
|
||||
body, _ := json.Marshal(cfg)
|
||||
return "gAAAAAC" + base64.StdEncoding.EncodeToString(body)
|
||||
}
|
||||
|
||||
func buildProofToken(seed, difficulty, userAgent string, scriptSources []string, dataBuild string) (string, error) {
|
||||
cfg := buildPOWConfig(userAgent, scriptSources, dataBuild)
|
||||
answer, solved := powGenerate(seed, difficulty, cfg, 500000)
|
||||
if !solved {
|
||||
return "", errors.New("failed to solve proof token")
|
||||
}
|
||||
return "gAAAAAB" + answer, nil
|
||||
}
|
||||
|
||||
func buildPOWConfig(userAgent string, scriptSources []string, dataBuild string) []any {
|
||||
// scriptSources/dataBuild are no longer part of the sentinel config array
|
||||
// (the current chatgpt.com client dropped them); kept in the signature for
|
||||
// call-site compatibility.
|
||||
_ = scriptSources
|
||||
_ = dataBuild
|
||||
rng := rand.New(rand.NewSource(time.Now().UnixNano()))
|
||||
screen := screenResolutions[rng.Intn(len(screenResolutions))]
|
||||
loc := time.FixedZone("GMT+0800", 8*3600)
|
||||
nowLocal := time.Now().In(loc).Format("Mon Jan 02 2006 15:04:05") + " GMT+0800 (中国标准时间)"
|
||||
perf := float64(time.Now().UnixNano()%1_000_000_000) / 1_000_000
|
||||
return []any{
|
||||
screen[0] + screen[1], // [0]
|
||||
nowLocal, // [1] local time, JS Date.toString() shape
|
||||
4395630592, // [2]
|
||||
1, // [3] overwritten by powGenerate counter
|
||||
userAgent, // [4]
|
||||
nil, // [5] (was script source; now null)
|
||||
defaultClientVersion, // [6] oai-client-version, must match header
|
||||
"zh-CN", // [7] matches oai-language
|
||||
"zh-CN,en,en-GB,en-US", // [8]
|
||||
rng.Float64(), // [9] overwritten by powGenerate counter
|
||||
navKeys[rng.Intn(len(navKeys))],
|
||||
documentKeys[rng.Intn(len(documentKeys))],
|
||||
winKeys[rng.Intn(len(winKeys))],
|
||||
perf, // [13]
|
||||
newUUID(), // [14]
|
||||
"", // [15]
|
||||
cores[rng.Intn(len(cores))], // [16]
|
||||
float64(timeMillis()) - perf, // [17]
|
||||
0, 0, 0, 0, 0, 0,
|
||||
0,
|
||||
}
|
||||
}
|
||||
|
||||
func powGenerate(seed, difficulty string, cfg []any, limit int) (string, bool) {
|
||||
target, err := hex.DecodeString(strings.TrimSpace(difficulty))
|
||||
if err != nil {
|
||||
return "", false
|
||||
}
|
||||
diffLen := len(strings.TrimSpace(difficulty)) / 2
|
||||
seedBytes := []byte(seed)
|
||||
head1, _ := json.Marshal(cfg[:3])
|
||||
head2, _ := json.Marshal(cfg[4:9])
|
||||
head3, _ := json.Marshal(cfg[10:])
|
||||
static1 := []byte(string(head1[:len(head1)-1]) + ",")
|
||||
static2 := []byte("," + string(head2[1:len(head2)-1]) + ",")
|
||||
static3 := []byte("," + string(head3[1:]))
|
||||
|
||||
for i := 0; i < limit; i++ {
|
||||
finalJSON := append([]byte{}, static1...)
|
||||
finalJSON = append(finalJSON, []byte(strconvItoa(i))...)
|
||||
finalJSON = append(finalJSON, static2...)
|
||||
finalJSON = append(finalJSON, []byte(strconvItoa(i>>1))...)
|
||||
finalJSON = append(finalJSON, static3...)
|
||||
encoded := base64.StdEncoding.EncodeToString(finalJSON)
|
||||
sum := sha3.Sum512(append(seedBytes, []byte(encoded)...))
|
||||
if bytesCompare(sum[:diffLen], target) <= 0 {
|
||||
return encoded, true
|
||||
}
|
||||
}
|
||||
fallback := "wQ8Lk5FbGpA2NcR9dShT6gYjU7VxZ4D" + base64.StdEncoding.EncodeToString([]byte(`"`+seed+`"`))
|
||||
return fallback, false
|
||||
}
|
||||
|
||||
func bytesCompare(a, b []byte) int {
|
||||
for i := 0; i < len(a) && i < len(b); i++ {
|
||||
if a[i] < b[i] {
|
||||
return -1
|
||||
}
|
||||
if a[i] > b[i] {
|
||||
return 1
|
||||
}
|
||||
}
|
||||
if len(a) < len(b) {
|
||||
return -1
|
||||
}
|
||||
if len(a) > len(b) {
|
||||
return 1
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
func strconvItoa(v int) string {
|
||||
return strconv.Itoa(v)
|
||||
}
|
||||
@@ -0,0 +1,182 @@
|
||||
package chatgpt
|
||||
|
||||
import (
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"math/rand"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
type orderedMap struct {
|
||||
keys []string
|
||||
values map[string]any
|
||||
}
|
||||
|
||||
func newOrderedMap() *orderedMap {
|
||||
return &orderedMap{values: map[string]any{}}
|
||||
}
|
||||
|
||||
func (m *orderedMap) add(key string, value any) {
|
||||
if _, ok := m.values[key]; !ok {
|
||||
m.keys = append(m.keys, key)
|
||||
}
|
||||
m.values[key] = value
|
||||
}
|
||||
|
||||
func solveTurnstileToken(dx, p string) string {
|
||||
decoded, err := base64.StdEncoding.DecodeString(dx)
|
||||
if err != nil {
|
||||
return ""
|
||||
}
|
||||
var tokenList [][]any
|
||||
if err := json.Unmarshal([]byte(xorString(string(decoded), p)), &tokenList); err != nil {
|
||||
return ""
|
||||
}
|
||||
|
||||
processMap := map[int]any{16: p}
|
||||
start := time.Now()
|
||||
result := ""
|
||||
rng := rand.New(rand.NewSource(time.Now().UnixNano()))
|
||||
|
||||
toStr := func(value any) string {
|
||||
if value == nil {
|
||||
return "undefined"
|
||||
}
|
||||
if s, ok := value.(string); ok {
|
||||
special := map[string]string{
|
||||
"window.Math": "[object Math]",
|
||||
"window.Reflect": "[object Reflect]",
|
||||
"window.performance": "[object Performance]",
|
||||
"window.localStorage": "[object Storage]",
|
||||
"window.Object": "function Object() { [native code] }",
|
||||
"window.Reflect.set": "function set() { [native code] }",
|
||||
"window.performance.now": "function () { [native code] }",
|
||||
"window.Object.create": "function create() { [native code] }",
|
||||
"window.Object.keys": "function keys() { [native code] }",
|
||||
"window.Math.random": "function random() { [native code] }",
|
||||
}
|
||||
if specialValue, ok := special[s]; ok {
|
||||
return specialValue
|
||||
}
|
||||
return s
|
||||
}
|
||||
if list, ok := value.([]string); ok {
|
||||
return strings.Join(list, ",")
|
||||
}
|
||||
return stringValue(value)
|
||||
}
|
||||
|
||||
for _, token := range tokenList {
|
||||
if len(token) == 0 {
|
||||
continue
|
||||
}
|
||||
op := intValue(token[0])
|
||||
switch op {
|
||||
case 2:
|
||||
if len(token) >= 3 {
|
||||
processMap[intValue(token[1])] = token[2]
|
||||
}
|
||||
case 3:
|
||||
if len(token) >= 2 {
|
||||
result = base64.StdEncoding.EncodeToString([]byte(toStr(processMap[intValue(token[1])])))
|
||||
}
|
||||
case 5:
|
||||
if len(token) >= 3 {
|
||||
e := intValue(token[1])
|
||||
t := intValue(token[2])
|
||||
cur := processMap[e]
|
||||
inc := processMap[t]
|
||||
if list, ok := cur.([]any); ok {
|
||||
processMap[e] = append(list, inc)
|
||||
} else if _, ok := cur.(string); ok {
|
||||
processMap[e] = toStr(cur) + toStr(inc)
|
||||
} else {
|
||||
processMap[e] = "NaN"
|
||||
}
|
||||
}
|
||||
case 6, 24:
|
||||
if len(token) >= 4 {
|
||||
e := intValue(token[1])
|
||||
t := toStr(processMap[intValue(token[2])])
|
||||
n := toStr(processMap[intValue(token[3])])
|
||||
v := t + "." + n
|
||||
if op == 6 && v == "window.document.location" {
|
||||
v = "https://chatgpt.com/"
|
||||
}
|
||||
processMap[e] = v
|
||||
}
|
||||
case 8:
|
||||
if len(token) >= 3 {
|
||||
processMap[intValue(token[1])] = processMap[intValue(token[2])]
|
||||
}
|
||||
case 14:
|
||||
if len(token) >= 3 {
|
||||
var parsed any
|
||||
if err := json.Unmarshal([]byte(toStr(processMap[intValue(token[2])])), &parsed); err == nil {
|
||||
processMap[intValue(token[1])] = parsed
|
||||
}
|
||||
}
|
||||
case 15:
|
||||
if len(token) >= 3 {
|
||||
b, _ := json.Marshal(processMap[intValue(token[2])])
|
||||
processMap[intValue(token[1])] = string(b)
|
||||
}
|
||||
case 17:
|
||||
if len(token) >= 3 {
|
||||
e := intValue(token[1])
|
||||
target := toStr(processMap[intValue(token[2])])
|
||||
switch target {
|
||||
case "window.performance.now":
|
||||
processMap[e] = float64(time.Since(start).Nanoseconds())/1e6 + rng.Float64()
|
||||
case "window.Object.create":
|
||||
processMap[e] = newOrderedMap()
|
||||
case "window.Object.keys":
|
||||
processMap[e] = []string{
|
||||
"STATSIG_LOCAL_STORAGE_INTERNAL_STORE_V4",
|
||||
"STATSIG_LOCAL_STORAGE_STABLE_ID",
|
||||
"client-correlated-secret",
|
||||
"oai/apps/capExpiresAt",
|
||||
"oai-did",
|
||||
"STATSIG_LOCAL_STORAGE_LOGGING_REQUEST",
|
||||
"UiState.isNavigationCollapsed.1",
|
||||
}
|
||||
case "window.Math.random":
|
||||
processMap[e] = rng.Float64()
|
||||
}
|
||||
}
|
||||
case 18:
|
||||
if len(token) >= 2 {
|
||||
raw, err := base64.StdEncoding.DecodeString(toStr(processMap[intValue(token[1])]))
|
||||
if err == nil {
|
||||
processMap[intValue(token[1])] = string(raw)
|
||||
}
|
||||
}
|
||||
case 19:
|
||||
if len(token) >= 2 {
|
||||
processMap[intValue(token[1])] = base64.StdEncoding.EncodeToString([]byte(toStr(processMap[intValue(token[1])])))
|
||||
}
|
||||
case 20:
|
||||
if len(token) >= 4 {
|
||||
if toStr(processMap[intValue(token[1])]) == toStr(processMap[intValue(token[2])]) {
|
||||
if intValue(token[3]) == 3 && len(token) >= 5 {
|
||||
result = base64.StdEncoding.EncodeToString([]byte(toStr(processMap[intValue(token[4])])))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
func xorString(text, key string) string {
|
||||
if key == "" {
|
||||
return text
|
||||
}
|
||||
out := make([]rune, 0, len(text))
|
||||
keyRunes := []rune(key)
|
||||
for i, ch := range text {
|
||||
out = append(out, ch^keyRunes[i%len(keyRunes)])
|
||||
}
|
||||
return string(out)
|
||||
}
|
||||
@@ -0,0 +1,125 @@
|
||||
package chatgpt
|
||||
|
||||
import (
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"regexp"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
const (
|
||||
baseURL = "https://chatgpt.com"
|
||||
defaultUserAgent = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/149.0.0.0 Safari/537.36 Edg/149.0.0.0"
|
||||
defaultClientVersion = "prod-ab8a6348980a3e1d771c463b9f4f3e4e584f2769"
|
||||
defaultClientBuildNumber = "7624276"
|
||||
defaultPOWScript = "https://chatgpt.com/backend-api/sentinel/sdk.js"
|
||||
)
|
||||
|
||||
var (
|
||||
fileServiceIDPattern = regexp.MustCompile(`file-service://([A-Za-z0-9_-]+)`)
|
||||
sedimentIDPattern = regexp.MustCompile(`sediment://([A-Za-z0-9_-]+)`)
|
||||
realImageIDPattern = regexp.MustCompile(`\bfile_00000000[a-f0-9]{24}\b`)
|
||||
conversationIDRE = regexp.MustCompile(`"conversation_id"\s*:\s*"([^"]+)"`)
|
||||
scriptSrcRE = regexp.MustCompile(`<script[^>]+src="([^"]+)"`)
|
||||
dataBuildPathRE = regexp.MustCompile(`c/[^/]*/_`)
|
||||
htmlDataBuildRE = regexp.MustCompile(`<html[^>]*data-build="([^"]*)"`)
|
||||
)
|
||||
|
||||
func stringValue(v any) string {
|
||||
switch x := v.(type) {
|
||||
case string:
|
||||
return x
|
||||
case nil:
|
||||
return ""
|
||||
default:
|
||||
return fmt.Sprint(v)
|
||||
}
|
||||
}
|
||||
|
||||
func intValue(v any) int {
|
||||
switch x := v.(type) {
|
||||
case int:
|
||||
return x
|
||||
case int64:
|
||||
return int(x)
|
||||
case float64:
|
||||
return int(x)
|
||||
case float32:
|
||||
return int(x)
|
||||
case json.Number:
|
||||
n, _ := x.Int64()
|
||||
return int(n)
|
||||
case string:
|
||||
n, _ := strconv.Atoi(strings.TrimSpace(x))
|
||||
return n
|
||||
default:
|
||||
return 0
|
||||
}
|
||||
}
|
||||
|
||||
func decodeJWTPayload(token string) map[string]any {
|
||||
parts := strings.Split(strings.TrimSpace(token), ".")
|
||||
if len(parts) < 2 {
|
||||
return map[string]any{}
|
||||
}
|
||||
raw, err := base64.RawURLEncoding.DecodeString(parts[1])
|
||||
if err != nil {
|
||||
return map[string]any{}
|
||||
}
|
||||
var out map[string]any
|
||||
if err := json.Unmarshal(raw, &out); err != nil {
|
||||
return map[string]any{}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func newUUID() string {
|
||||
return uuid.NewString()
|
||||
}
|
||||
|
||||
func clip(v []byte, n int) string {
|
||||
s := strings.TrimSpace(string(v))
|
||||
if len(s) <= n {
|
||||
return s
|
||||
}
|
||||
return s[:n]
|
||||
}
|
||||
|
||||
func parsePOWResources(html string) ([]string, string) {
|
||||
matches := scriptSrcRE.FindAllStringSubmatch(html, -1)
|
||||
sources := make([]string, 0, len(matches))
|
||||
dataBuild := ""
|
||||
for _, match := range matches {
|
||||
if len(match) < 2 {
|
||||
continue
|
||||
}
|
||||
src := strings.TrimSpace(match[1])
|
||||
if src == "" {
|
||||
continue
|
||||
}
|
||||
sources = append(sources, src)
|
||||
if dataBuild == "" {
|
||||
if path := dataBuildPathRE.FindString(src); path != "" {
|
||||
dataBuild = path
|
||||
}
|
||||
}
|
||||
}
|
||||
if dataBuild == "" {
|
||||
if match := htmlDataBuildRE.FindStringSubmatch(html); len(match) >= 2 {
|
||||
dataBuild = strings.TrimSpace(match[1])
|
||||
}
|
||||
}
|
||||
if len(sources) == 0 {
|
||||
sources = []string{defaultPOWScript}
|
||||
}
|
||||
return sources, dataBuild
|
||||
}
|
||||
|
||||
func timeMillis() int64 {
|
||||
return time.Now().UnixMilli()
|
||||
}
|
||||
@@ -0,0 +1,441 @@
|
||||
// Package imagine implements the Imagine.art (vyro.ai) provider client. The
|
||||
// durable credential is a JSON blob {"token","refreshToken"}: `token` is a ~6h
|
||||
// access JWT used as Authorization: Bearer for the API, and `refreshToken` is a
|
||||
// ~7d JWT that mints a fresh pair via /apis/v1/auth/other/refresh/web when the
|
||||
// access token expires. Both rotate on refresh, so the new pair MUST be saved.
|
||||
// tls-client gives a Chrome JA3/JA4 so vyro's edge doesn't flag the requests.
|
||||
package imagine
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/rand"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
http "github.com/bogdanfinn/fhttp"
|
||||
tlsclient "github.com/bogdanfinn/tls-client"
|
||||
"github.com/bogdanfinn/tls-client/profiles"
|
||||
)
|
||||
|
||||
const (
|
||||
apiBase = "https://imagine.vyro.ai"
|
||||
teamsBase = "https://teams-imagine.vyro.ai"
|
||||
authBase = "https://auth.vyro.ai"
|
||||
webOrigin = "https://www.imagine.art"
|
||||
userAgent = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36"
|
||||
)
|
||||
|
||||
var (
|
||||
ErrAuth = errors.New("imagine auth failed")
|
||||
ErrQuotaExhausted = errors.New("imagine quota exhausted")
|
||||
ErrTemporaryUpstream = errors.New("imagine upstream temporary error")
|
||||
)
|
||||
|
||||
// refreshLeadSeconds renews the access token this many seconds BEFORE it expires
|
||||
// (proactive, not lazy at expiry) — the maintenance sweep keeps tokens fresh so a
|
||||
// dormant account's rotating refreshToken never lapses.
|
||||
const refreshLeadSeconds = 600 // 10 minutes
|
||||
|
||||
type Client struct {
|
||||
proxy string
|
||||
// freshest credential per account (key: user id) + a per-account refresh lock,
|
||||
// so concurrent callers don't each spend the rotating refresh_token — the first
|
||||
// refreshes, the rest reuse the cached fresh credential.
|
||||
mu sync.Mutex
|
||||
creds map[string]string
|
||||
locks map[string]*sync.Mutex
|
||||
}
|
||||
|
||||
func NewClient(proxy string) *Client {
|
||||
return &Client{proxy: strings.TrimSpace(proxy), creds: map[string]string{}, locks: map[string]*sync.Mutex{}}
|
||||
}
|
||||
|
||||
func (c *Client) SetProxy(proxy string) {
|
||||
c.proxy = strings.TrimSpace(proxy)
|
||||
}
|
||||
|
||||
func (c *Client) userLock(userID string) *sync.Mutex {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
m, ok := c.locks[userID]
|
||||
if !ok {
|
||||
m = &sync.Mutex{}
|
||||
c.locks[userID] = m
|
||||
}
|
||||
return m
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Credential helpers
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
type credential struct {
|
||||
Token string `json:"token"`
|
||||
RefreshToken string `json:"refreshToken"`
|
||||
// Email is the real account email — supplied at import, used for display and
|
||||
// (pool,email) dedup. It is NOT in the JWT (which only carries userId), so it
|
||||
// must be carried across refreshes (the refresh response omits it).
|
||||
Email string `json:"email,omitempty"`
|
||||
// ParentID is a canvas node the account OWNS, used as the generation's
|
||||
// parent_id. Imagine rejects any parent the account doesn't own ("user does
|
||||
// not have access to parent asset") and silently orphans a parent-less
|
||||
// generation (charged but never produced) — so it's supplied at import and
|
||||
// carried across refreshes.
|
||||
ParentID string `json:"parentId,omitempty"`
|
||||
}
|
||||
|
||||
func parseCred(s string) (credential, bool) {
|
||||
var cr credential
|
||||
if json.Unmarshal([]byte(strings.TrimSpace(s)), &cr) != nil {
|
||||
return cr, false
|
||||
}
|
||||
if strings.TrimSpace(cr.Token) == "" || strings.TrimSpace(cr.RefreshToken) == "" {
|
||||
return cr, false
|
||||
}
|
||||
return cr, true
|
||||
}
|
||||
|
||||
func buildCred(token, refresh, email, parentID string) string {
|
||||
b, _ := json.Marshal(credential{
|
||||
Token: strings.TrimSpace(token),
|
||||
RefreshToken: strings.TrimSpace(refresh),
|
||||
Email: strings.TrimSpace(email),
|
||||
ParentID: strings.TrimSpace(parentID),
|
||||
})
|
||||
return string(b)
|
||||
}
|
||||
|
||||
// ParentIDFromCred returns the canvas parent node id supplied at import.
|
||||
func ParentIDFromCred(cred string) string {
|
||||
cr, ok := parseCred(cred)
|
||||
if !ok {
|
||||
return ""
|
||||
}
|
||||
return strings.TrimSpace(cr.ParentID)
|
||||
}
|
||||
|
||||
func looksLikeJWT(s string) bool {
|
||||
return len(strings.Split(strings.TrimSpace(s), ".")) == 3
|
||||
}
|
||||
|
||||
// IsImagineToken reports whether a pasted credential is an Imagine.art account:
|
||||
// a JSON object carrying a non-empty token + refreshToken that both look like
|
||||
// JWTs. Distinguishes it from adobe/leonardo/krea cookies.
|
||||
func IsImagineToken(value string) bool {
|
||||
cr, ok := parseCred(value)
|
||||
if !ok {
|
||||
return false
|
||||
}
|
||||
return looksLikeJWT(cr.Token) && looksLikeJWT(cr.RefreshToken)
|
||||
}
|
||||
|
||||
// jwtClaims base64url-decodes the JWT payload (segment 1) into a claims map.
|
||||
func jwtClaims(token string) map[string]any {
|
||||
parts := strings.Split(strings.TrimSpace(token), ".")
|
||||
if len(parts) < 2 {
|
||||
return nil
|
||||
}
|
||||
raw, err := base64.RawURLEncoding.DecodeString(parts[1])
|
||||
if err != nil {
|
||||
// tolerate padded variants
|
||||
if raw, err = base64.URLEncoding.DecodeString(parts[1]); err != nil {
|
||||
return nil
|
||||
}
|
||||
}
|
||||
var m map[string]any
|
||||
if json.Unmarshal(raw, &m) != nil {
|
||||
return nil
|
||||
}
|
||||
return m
|
||||
}
|
||||
|
||||
func userIDFromToken(token string) string {
|
||||
claims := jwtClaims(token)
|
||||
if claims == nil {
|
||||
return ""
|
||||
}
|
||||
if v := strings.TrimSpace(stringValue(claims["userId"])); v != "" {
|
||||
return v
|
||||
}
|
||||
return strings.TrimSpace(stringValue(claims["sub"]))
|
||||
}
|
||||
|
||||
func tokenExp(token string) int64 {
|
||||
claims := jwtClaims(token)
|
||||
if claims == nil {
|
||||
return 0
|
||||
}
|
||||
return toInt64(claims["exp"])
|
||||
}
|
||||
|
||||
// EmailFromCred returns the real account email supplied at import; if absent it
|
||||
// falls back to the JWT userId so (pool,email) dedup still has a stable key.
|
||||
func EmailFromCred(cred string) string {
|
||||
cr, ok := parseCred(cred)
|
||||
if !ok {
|
||||
return ""
|
||||
}
|
||||
if e := strings.TrimSpace(cr.Email); e != "" {
|
||||
return e
|
||||
}
|
||||
return userIDFromToken(cr.Token)
|
||||
}
|
||||
|
||||
// UserIDFromCred returns the JWT userId (== org_id used for credit/generation).
|
||||
func UserIDFromCred(cred string) string {
|
||||
cr, ok := parseCred(cred)
|
||||
if !ok {
|
||||
return ""
|
||||
}
|
||||
return userIDFromToken(cr.Token)
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Refresh
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
// RefreshIfNeeded returns a credential whose access token is still valid: if the
|
||||
// stored one is (near) expired it spends the refreshToken to mint a fresh pair
|
||||
// and rebuilds the credential. Returns (cred, changed, err); changed=true means
|
||||
// the caller must persist the new credential (both tokens rotate). ErrAuth means
|
||||
// the refreshToken is dead → the account is gone.
|
||||
func (c *Client) RefreshIfNeeded(ctx context.Context, cred string) (string, bool, error) {
|
||||
cr, ok := parseCred(cred)
|
||||
if !ok {
|
||||
return cred, false, nil // unparseable — let the downstream call surface the error
|
||||
}
|
||||
userID := userIDFromToken(cr.Token)
|
||||
now := time.Now().Unix()
|
||||
|
||||
lk := c.userLock(userID)
|
||||
lk.Lock()
|
||||
defer lk.Unlock()
|
||||
|
||||
// A concurrent caller may already have refreshed this account.
|
||||
if userID != "" {
|
||||
c.mu.Lock()
|
||||
cached := c.creds[userID]
|
||||
c.mu.Unlock()
|
||||
if cc, ok := parseCred(cached); ok && tokenExp(cc.Token)-refreshLeadSeconds > now {
|
||||
return cached, cached != cred, nil
|
||||
}
|
||||
}
|
||||
if tokenExp(cr.Token)-refreshLeadSeconds > now {
|
||||
return cred, false, nil // still valid
|
||||
}
|
||||
|
||||
respBody, status, err := c.refreshPost(ctx, cr.RefreshToken)
|
||||
if err != nil {
|
||||
return "", false, fmt.Errorf("%w: refresh: %s", ErrTemporaryUpstream, err.Error())
|
||||
}
|
||||
if status == 400 || status == 401 || status == 403 {
|
||||
return "", false, ErrAuth
|
||||
}
|
||||
if status != 200 {
|
||||
return "", false, fmt.Errorf("%w: refresh http %d: %s", ErrTemporaryUpstream, status, clip(respBody, 120))
|
||||
}
|
||||
var rb struct {
|
||||
Result struct {
|
||||
SessionToken string `json:"sessionToken"`
|
||||
RefreshToken string `json:"refreshToken"`
|
||||
} `json:"result"`
|
||||
}
|
||||
if json.Unmarshal(respBody, &rb) != nil || strings.TrimSpace(rb.Result.SessionToken) == "" {
|
||||
return "", false, ErrAuth
|
||||
}
|
||||
newRefresh := rb.Result.RefreshToken
|
||||
if strings.TrimSpace(newRefresh) == "" {
|
||||
newRefresh = cr.RefreshToken // some responses may omit it — keep the old one
|
||||
}
|
||||
newCred := buildCred(rb.Result.SessionToken, newRefresh, cr.Email, cr.ParentID)
|
||||
if userID != "" {
|
||||
c.mu.Lock()
|
||||
c.creds[userID] = newCred
|
||||
c.mu.Unlock()
|
||||
}
|
||||
return newCred, true, nil
|
||||
}
|
||||
|
||||
func (c *Client) refreshPost(ctx context.Context, refreshToken string) ([]byte, int, error) {
|
||||
client, err := c.newTLSClient()
|
||||
if err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
req, err := http.NewRequest(http.MethodPost, authBase+"/apis/v1/auth/other/refresh/web", nil)
|
||||
if err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
req = req.WithContext(ctx)
|
||||
req.Header = http.Header{
|
||||
"accept": {"application/json, text/plain, */*"},
|
||||
"authorization": {"Bearer " + refreshToken},
|
||||
"origin": {webOrigin},
|
||||
"referer": {webOrigin + "/"},
|
||||
"user-agent": {userAgent},
|
||||
http.HeaderOrderKey: {
|
||||
"accept", "authorization", "origin", "referer", "user-agent",
|
||||
},
|
||||
}
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
b, err := io.ReadAll(resp.Body)
|
||||
return b, resp.StatusCode, err
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Credits
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
// FetchCreditsBalance reads the account's credit balance via /v1/credit.
|
||||
// remaining is the `total` field. 401/403 → ErrAuth (token dead). Returns the
|
||||
// normalized map shared by all providers.
|
||||
func (c *Client) FetchCreditsBalance(ctx context.Context, cred string) (map[string]any, error) {
|
||||
cr, ok := parseCred(cred)
|
||||
if !ok {
|
||||
return unknownBalance("bad credential"), nil
|
||||
}
|
||||
userID := userIDFromToken(cr.Token)
|
||||
body, status, err := c.apiGet(ctx, cr.Token, apiBase+"/v1/credit?org_id="+userID)
|
||||
if err != nil {
|
||||
return unknownBalance("network: " + err.Error()), nil
|
||||
}
|
||||
if status == 401 || status == 403 {
|
||||
return nil, ErrAuth
|
||||
}
|
||||
if status != 200 {
|
||||
return unknownBalance(fmt.Sprintf("http %d: %s", status, clip(body, 160))), nil
|
||||
}
|
||||
var cb struct {
|
||||
Status string `json:"status"`
|
||||
Total int `json:"total"`
|
||||
}
|
||||
if err := json.Unmarshal(body, &cb); err != nil {
|
||||
return unknownBalance("non-json"), nil
|
||||
}
|
||||
return map[string]any{
|
||||
"remaining": cb.Total,
|
||||
"used": nil,
|
||||
"total": nil,
|
||||
"unknown": false,
|
||||
"error": nil,
|
||||
"email": emptyStringNil(EmailFromCred(cred)),
|
||||
}, nil
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// HTTP helpers
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
func (c *Client) apiGet(ctx context.Context, token, url string) ([]byte, int, error) {
|
||||
client, err := c.newTLSClient()
|
||||
if err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
req, err := http.NewRequest(http.MethodGet, url, nil)
|
||||
if err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
req = req.WithContext(ctx)
|
||||
req.Header = http.Header{
|
||||
"accept": {"application/json, text/plain, */*"},
|
||||
"authorization": {"Bearer " + token},
|
||||
"origin": {webOrigin},
|
||||
"referer": {webOrigin + "/"},
|
||||
"user-agent": {userAgent},
|
||||
http.HeaderOrderKey: {
|
||||
"accept", "authorization", "origin", "referer", "user-agent",
|
||||
},
|
||||
}
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
b, err := io.ReadAll(resp.Body)
|
||||
return b, resp.StatusCode, err
|
||||
}
|
||||
|
||||
func (c *Client) newTLSClient() (tlsclient.HttpClient, error) {
|
||||
options := []tlsclient.HttpClientOption{
|
||||
tlsclient.WithTimeoutSeconds(60),
|
||||
tlsclient.WithClientProfile(profiles.Chrome_120),
|
||||
}
|
||||
if c.proxy != "" {
|
||||
options = append(options, tlsclient.WithProxyUrl(c.proxy))
|
||||
}
|
||||
return tlsclient.NewHttpClient(tlsclient.NewNoopLogger(), options...)
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Small util
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
func uuid4() string {
|
||||
var b [16]byte
|
||||
_, _ = rand.Read(b[:])
|
||||
b[6] = (b[6] & 0x0f) | 0x40
|
||||
b[8] = (b[8] & 0x3f) | 0x80
|
||||
return fmt.Sprintf("%x-%x-%x-%x-%x", b[0:4], b[4:6], b[6:8], b[8:10], b[10:16])
|
||||
}
|
||||
|
||||
func unknownBalance(reason string) map[string]any {
|
||||
return map[string]any{
|
||||
"remaining": nil, "used": nil, "total": nil, "unknown": true, "error": reason,
|
||||
}
|
||||
}
|
||||
|
||||
func stringValue(v any) string {
|
||||
switch x := v.(type) {
|
||||
case string:
|
||||
return x
|
||||
case nil:
|
||||
return ""
|
||||
default:
|
||||
b, _ := json.Marshal(x)
|
||||
return strings.TrimSpace(string(b))
|
||||
}
|
||||
}
|
||||
|
||||
func toInt64(v any) int64 {
|
||||
switch x := v.(type) {
|
||||
case float64:
|
||||
return int64(x)
|
||||
case int64:
|
||||
return x
|
||||
case int:
|
||||
return int64(x)
|
||||
case json.Number:
|
||||
n, _ := x.Int64()
|
||||
return n
|
||||
case string:
|
||||
n, _ := strconv.ParseInt(strings.TrimSpace(x), 10, 64)
|
||||
return n
|
||||
default:
|
||||
return 0
|
||||
}
|
||||
}
|
||||
|
||||
func emptyStringNil(v string) any {
|
||||
if strings.TrimSpace(v) == "" {
|
||||
return nil
|
||||
}
|
||||
return v
|
||||
}
|
||||
|
||||
func clip(b []byte, n int) string {
|
||||
s := strings.TrimSpace(string(b))
|
||||
if len(s) > n {
|
||||
return s[:n]
|
||||
}
|
||||
return s
|
||||
}
|
||||
@@ -0,0 +1,242 @@
|
||||
package imagine
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"mime/multipart"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
http "github.com/bogdanfinn/fhttp"
|
||||
)
|
||||
|
||||
// GenerateImage runs the full Imagine.art pipeline: submit the txt2img job, poll
|
||||
// the org objects feed until the batch finishes, then download the produced
|
||||
// image. styleID picks the model (41001 = 1.5 / 2K, 41004 = 1.5pro / 4K).
|
||||
// HTTP 402 → ErrQuotaExhausted. These models are pure text2img (no refs).
|
||||
func (c *Client) GenerateImage(ctx context.Context, cred string, styleID int, resolution, aspectRatio, prompt string) ([]byte, map[string]any, error) {
|
||||
cr, ok := parseCred(cred)
|
||||
if !ok {
|
||||
return nil, nil, ErrAuth
|
||||
}
|
||||
userID := userIDFromToken(cr.Token)
|
||||
|
||||
metadata, _ := json.Marshal(map[string]any{
|
||||
"placeholderUuid": uuid4(),
|
||||
"promptWithoutManipulation": prompt,
|
||||
"modeId": 0,
|
||||
})
|
||||
// parent_id MUST be a canvas node this account owns — the server rejects a
|
||||
// foreign id ("user does not have access to parent asset") and silently
|
||||
// orphans a parent-less generation (charged but never produced). It's supplied
|
||||
// at import (credential.parentId).
|
||||
fields := map[string]string{
|
||||
"style_id": strconv.Itoa(styleID),
|
||||
"aspect_ratio": aspectRatio,
|
||||
"resolution": resolution,
|
||||
"variation": "txt2img",
|
||||
"prompt": prompt,
|
||||
"is_enhance": "0",
|
||||
"count": "1",
|
||||
"clientVersion": "1",
|
||||
"org_id": userID,
|
||||
"use_plugin": "false",
|
||||
"metadata": string(metadata),
|
||||
}
|
||||
if pid := strings.TrimSpace(cr.ParentID); pid != "" {
|
||||
fields["parent_id"] = pid
|
||||
}
|
||||
|
||||
var buf bytes.Buffer
|
||||
w := multipart.NewWriter(&buf)
|
||||
for k, v := range fields {
|
||||
_ = w.WriteField(k, v)
|
||||
}
|
||||
_ = w.Close()
|
||||
|
||||
body, status, err := c.apiPost(ctx, cr.Token, apiBase+"/v1/image/generations/upload", w.FormDataContentType(), buf.Bytes())
|
||||
if err != nil {
|
||||
return nil, nil, fmt.Errorf("%w: submit: %s", ErrTemporaryUpstream, err.Error())
|
||||
}
|
||||
if status == 401 || status == 403 {
|
||||
return nil, nil, ErrAuth
|
||||
}
|
||||
if status == 402 {
|
||||
return nil, nil, ErrQuotaExhausted
|
||||
}
|
||||
if status != 200 && status != 201 {
|
||||
return nil, nil, fmt.Errorf("%w: submit http %d: %s", ErrTemporaryUpstream, status, clip(body, 200))
|
||||
}
|
||||
var jobs []struct {
|
||||
BatchID string `json:"batchId"`
|
||||
ID string `json:"id"`
|
||||
Status string `json:"status"`
|
||||
}
|
||||
if err := json.Unmarshal(body, &jobs); err != nil || len(jobs) == 0 || jobs[0].BatchID == "" {
|
||||
return nil, nil, fmt.Errorf("%w: no batch id: %s", ErrTemporaryUpstream, clip(body, 200))
|
||||
}
|
||||
batchID := jobs[0].BatchID
|
||||
|
||||
imageURL, err := c.pollImage(ctx, cr.Token, userID, batchID)
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
data, err := c.download(ctx, imageURL)
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
return data, map[string]any{"batch_id": batchID, "image_url": imageURL, "org_id": userID}, nil
|
||||
}
|
||||
|
||||
// pollImage polls the org objects feed until the entry for our batch finishes,
|
||||
// then extracts its asset URL (image_url is a JSON-encoded array string).
|
||||
func (c *Client) pollImage(ctx context.Context, token, userID, batchID string) (string, error) {
|
||||
ticker := time.NewTicker(3 * time.Second)
|
||||
defer ticker.Stop()
|
||||
deadline := time.Now().Add(4 * time.Minute)
|
||||
|
||||
url := teamsBase + "/v1/org/" + userID + "/objects?batch=true&limit=50&service=image,chat-image"
|
||||
for {
|
||||
body, status, err := c.apiGet(ctx, token, url)
|
||||
if err == nil && status == 200 {
|
||||
var resp struct {
|
||||
Data []struct {
|
||||
BatchID string `json:"batch_id"`
|
||||
Status string `json:"status"`
|
||||
Code int `json:"code"`
|
||||
// Current shape: the produced asset lives at url.generation[0].
|
||||
URL struct {
|
||||
Generation []string `json:"generation"`
|
||||
} `json:"url"`
|
||||
ImageURL string `json:"image_url"` // legacy fallback
|
||||
} `json:"data"`
|
||||
}
|
||||
if json.Unmarshal(body, &resp) == nil {
|
||||
for _, o := range resp.Data {
|
||||
if o.BatchID != batchID {
|
||||
continue
|
||||
}
|
||||
st := strings.ToLower(strings.TrimSpace(o.Status))
|
||||
switch {
|
||||
case st == "finished" || o.Code == 2:
|
||||
if u := firstNonEmpty(o.URL.Generation); u != "" {
|
||||
return u, nil
|
||||
}
|
||||
if u := firstImageURL(o.ImageURL); u != "" {
|
||||
return u, nil
|
||||
}
|
||||
case st == "failed" || st == "error":
|
||||
return "", fmt.Errorf("%w: job %s", ErrTemporaryUpstream, st)
|
||||
}
|
||||
}
|
||||
}
|
||||
} else if status == 401 || status == 403 {
|
||||
return "", ErrAuth
|
||||
}
|
||||
if time.Now().After(deadline) {
|
||||
return "", fmt.Errorf("%w: generation timed out", ErrTemporaryUpstream)
|
||||
}
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return "", ctx.Err()
|
||||
case <-ticker.C:
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// firstNonEmpty returns the first non-blank string in a slice.
|
||||
func firstNonEmpty(ss []string) string {
|
||||
for _, s := range ss {
|
||||
if strings.TrimSpace(s) != "" {
|
||||
return strings.TrimSpace(s)
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
// firstImageURL parses the image_url field — a JSON-encoded array of URLs — and
|
||||
// returns the first one. Tolerates a bare string too.
|
||||
func firstImageURL(raw string) string {
|
||||
raw = strings.TrimSpace(raw)
|
||||
if raw == "" {
|
||||
return ""
|
||||
}
|
||||
var urls []string
|
||||
if json.Unmarshal([]byte(raw), &urls) == nil {
|
||||
for _, u := range urls {
|
||||
if strings.TrimSpace(u) != "" {
|
||||
return strings.TrimSpace(u)
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
if strings.HasPrefix(raw, "http") {
|
||||
return raw
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func (c *Client) download(ctx context.Context, url string) ([]byte, error) {
|
||||
client, err := c.newTLSClient()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
req, err := http.NewRequest(http.MethodGet, url, nil)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
req = req.WithContext(ctx)
|
||||
req.Header = http.Header{
|
||||
"accept": {"image/avif,image/webp,image/png,image/*,*/*;q=0.8"},
|
||||
"user-agent": {userAgent},
|
||||
"referer": {webOrigin + "/"},
|
||||
}
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
b, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if resp.StatusCode != 200 {
|
||||
return nil, fmt.Errorf("%w: image download http %d", ErrTemporaryUpstream, resp.StatusCode)
|
||||
}
|
||||
return b, nil
|
||||
}
|
||||
|
||||
// apiPost issues a POST with a raw body + content-type, carrying the bearer token.
|
||||
func (c *Client) apiPost(ctx context.Context, token, url, contentType string, body []byte) ([]byte, int, error) {
|
||||
client, err := c.newTLSClient()
|
||||
if err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
req, err := http.NewRequest(http.MethodPost, url, bytes.NewReader(body))
|
||||
if err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
req = req.WithContext(ctx)
|
||||
req.Header = http.Header{
|
||||
"accept": {"application/json, text/plain, */*"},
|
||||
"authorization": {"Bearer " + token},
|
||||
"content-type": {contentType},
|
||||
"origin": {webOrigin},
|
||||
"referer": {webOrigin + "/"},
|
||||
"user-agent": {userAgent},
|
||||
http.HeaderOrderKey: {
|
||||
"accept", "authorization", "content-type", "origin", "referer", "user-agent",
|
||||
},
|
||||
}
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
b, err := io.ReadAll(resp.Body)
|
||||
return b, resp.StatusCode, err
|
||||
}
|
||||
@@ -0,0 +1,488 @@
|
||||
// Package krea implements the Krea.ai (krea.ai) provider client. The durable
|
||||
// credential is the browser cookie (Supabase "sb-superb-auth-token"); Krea's own
|
||||
// Next.js backend reads it directly, so quota and generation just forward the
|
||||
// cookie — there's no separate token-exchange step. tls-client gives a Chrome
|
||||
// JA3/JA4 so Krea's Cloudflare edge doesn't flag the requests.
|
||||
package krea
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
http "github.com/bogdanfinn/fhttp"
|
||||
tlsclient "github.com/bogdanfinn/tls-client"
|
||||
"github.com/bogdanfinn/tls-client/profiles"
|
||||
)
|
||||
|
||||
// kreaAnonKey is Krea's public Supabase anon key (fixed, embedded in their
|
||||
// frontend) — required as the apikey/bearer when refreshing a session token.
|
||||
const kreaAnonKey = "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJyb2xlIjoiYW5vbiIsImlzcyI6InN1cGFiYXNlIiwiaWF0IjoxNzc1Mjc4ODU3LCJleHAiOjE5MzI5NTg4NTd9.NUiqEOd__QsCCMjo3D1zrCAda5dLV2F5p6Kf584sZKc"
|
||||
|
||||
const (
|
||||
apiBase = "https://www.krea.ai"
|
||||
userAgent = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36"
|
||||
)
|
||||
|
||||
var (
|
||||
ErrAuth = errors.New("krea auth failed")
|
||||
ErrQuotaExhausted = errors.New("krea quota exhausted")
|
||||
ErrTemporaryUpstream = errors.New("krea upstream temporary error")
|
||||
)
|
||||
|
||||
// refreshLeadSeconds renews the access token this many seconds BEFORE it expires
|
||||
// (not lazily at expiry). The maintenance sweep refreshes proactively, so a
|
||||
// dormant account's token never lapses — once the rotating refresh_token is gone
|
||||
// (expired/consumed) the account can't recover, so we keep it perpetually fresh.
|
||||
const refreshLeadSeconds = 600 // 10 minutes
|
||||
|
||||
type Client struct {
|
||||
proxy string
|
||||
// freshest cookie per account (key: user id) + a per-account refresh lock, so
|
||||
// concurrent callers don't each spend the single-use (rotating) refresh_token —
|
||||
// the first refreshes, the rest reuse the cached fresh cookie.
|
||||
mu sync.Mutex
|
||||
cookies map[string]string
|
||||
locks map[string]*sync.Mutex
|
||||
// actAt = last /app activation time per account (key: user id); actLocks gives
|
||||
// a per-account lock so concurrent generations wait for the first to finish the
|
||||
// (once-per-daily-reset) activation instead of each loading /app.
|
||||
actAt map[string]int64
|
||||
actLocks map[string]*sync.Mutex
|
||||
}
|
||||
|
||||
func NewClient(proxy string) *Client {
|
||||
return &Client{
|
||||
proxy: strings.TrimSpace(proxy),
|
||||
cookies: map[string]string{},
|
||||
locks: map[string]*sync.Mutex{},
|
||||
actAt: map[string]int64{},
|
||||
actLocks: map[string]*sync.Mutex{},
|
||||
}
|
||||
}
|
||||
|
||||
func (c *Client) userLock(userID string) *sync.Mutex {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
m, ok := c.locks[userID]
|
||||
if !ok {
|
||||
m = &sync.Mutex{}
|
||||
c.locks[userID] = m
|
||||
}
|
||||
return m
|
||||
}
|
||||
|
||||
// RefreshIfNeeded returns a cookie whose access_token is still valid: if the
|
||||
// stored one is (near) expired it spends the refresh_token to mint a new session
|
||||
// and rebuilds the cookie. Returns (cookie, changed, err); changed=true means the
|
||||
// caller must persist the new cookie (the refresh_token rotated). ErrAuth means
|
||||
// the refresh_token is dead → the account is gone.
|
||||
func (c *Client) RefreshIfNeeded(ctx context.Context, cookie string) (string, bool, error) {
|
||||
authVal := authCookieValue(cookie)
|
||||
sess, ok := decodeSession(authVal)
|
||||
if !ok {
|
||||
return cookie, false, nil // unparseable — let the downstream call surface the error
|
||||
}
|
||||
userID := nestedStr(sess, "user", "id")
|
||||
now := time.Now().Unix()
|
||||
|
||||
lk := c.userLock(userID)
|
||||
lk.Lock()
|
||||
defer lk.Unlock()
|
||||
|
||||
// A concurrent caller may already have refreshed this account.
|
||||
if userID != "" {
|
||||
c.mu.Lock()
|
||||
cached := c.cookies[userID]
|
||||
c.mu.Unlock()
|
||||
if cs, ok := decodeSession(authCookieValue(cached)); ok && toInt64(cs["expires_at"])-refreshLeadSeconds > now {
|
||||
return cached, cached != cookie, nil
|
||||
}
|
||||
}
|
||||
if toInt64(sess["expires_at"])-refreshLeadSeconds > now {
|
||||
return cookie, false, nil // still valid
|
||||
}
|
||||
refreshTok := strings.TrimSpace(stringValue(sess["refresh_token"]))
|
||||
if refreshTok == "" {
|
||||
return "", false, ErrAuth
|
||||
}
|
||||
|
||||
respBody, status, err := c.refreshPost(ctx, refreshTok)
|
||||
if err != nil {
|
||||
return "", false, fmt.Errorf("%w: refresh: %s", ErrTemporaryUpstream, err.Error())
|
||||
}
|
||||
if status == 400 || status == 401 || status == 403 {
|
||||
return "", false, ErrAuth
|
||||
}
|
||||
if status != 200 {
|
||||
return "", false, fmt.Errorf("%w: refresh http %d: %s", ErrTemporaryUpstream, status, clip(respBody, 120))
|
||||
}
|
||||
var ns map[string]any
|
||||
if json.Unmarshal(respBody, &ns) != nil || strings.TrimSpace(stringValue(ns["access_token"])) == "" {
|
||||
return "", false, ErrAuth
|
||||
}
|
||||
newCookie := replaceAuthCookie(cookie, "base64-"+base64.StdEncoding.EncodeToString(respBody))
|
||||
if userID != "" {
|
||||
c.mu.Lock()
|
||||
c.cookies[userID] = newCookie
|
||||
c.mu.Unlock()
|
||||
}
|
||||
return newCookie, true, nil
|
||||
}
|
||||
|
||||
func (c *Client) refreshPost(ctx context.Context, refreshToken string) ([]byte, int, error) {
|
||||
client, err := c.newTLSClient()
|
||||
if err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
body, _ := json.Marshal(map[string]string{"refresh_token": refreshToken})
|
||||
req, err := http.NewRequest(http.MethodPost, apiBase+"/auth/v1/token?grant_type=refresh_token", bytes.NewReader(body))
|
||||
if err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
req = req.WithContext(ctx)
|
||||
req.Header = http.Header{
|
||||
"accept": {"*/*"},
|
||||
"content-type": {"application/json;charset=UTF-8"},
|
||||
"apikey": {kreaAnonKey},
|
||||
"authorization": {"Bearer " + kreaAnonKey},
|
||||
"x-client-info": {"supabase-ssr/0.6.1 createBrowserClient"},
|
||||
"x-supabase-api-version": {"2024-01-01"},
|
||||
"origin": {apiBase},
|
||||
"referer": {apiBase + "/"},
|
||||
"user-agent": {userAgent},
|
||||
http.HeaderOrderKey: {
|
||||
"accept", "content-type", "apikey", "authorization", "x-client-info",
|
||||
"x-supabase-api-version", "origin", "referer", "user-agent",
|
||||
},
|
||||
}
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
b, err := io.ReadAll(resp.Body)
|
||||
return b, resp.StatusCode, err
|
||||
}
|
||||
|
||||
// chunkSize is supabase-ssr's per-cookie chunk limit; larger sessions (e.g.
|
||||
// Google-OAuth accounts) are split into sb-superb-auth-token.0/.1/...
|
||||
const chunkSize = 3600
|
||||
|
||||
func cookieVal(cookie, name string) string {
|
||||
for _, p := range strings.Split(cookie, ";") {
|
||||
p = strings.TrimSpace(p)
|
||||
if v, ok := strings.CutPrefix(p, name+"="); ok {
|
||||
return v
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
// authCookieValue returns the full auth value, transparently reassembling a
|
||||
// chunked cookie (sb-superb-auth-token.0 + .1 + ...) or returning the single one.
|
||||
func authCookieValue(cookie string) string {
|
||||
if v := cookieVal(cookie, "sb-superb-auth-token"); v != "" {
|
||||
return v
|
||||
}
|
||||
var b strings.Builder
|
||||
for i := 0; ; i++ {
|
||||
v := cookieVal(cookie, fmt.Sprintf("sb-superb-auth-token.%d", i))
|
||||
if v == "" {
|
||||
break
|
||||
}
|
||||
b.WriteString(v)
|
||||
}
|
||||
return b.String()
|
||||
}
|
||||
|
||||
// replaceAuthCookie drops every sb-superb-auth-token[.N] cookie and re-adds the
|
||||
// new value (chunked the same way supabase-ssr would if it's large), preserving
|
||||
// all other cookies (krea-workspace-id, etc.).
|
||||
func replaceAuthCookie(cookie, newValue string) string {
|
||||
var out []string
|
||||
for _, p := range strings.Split(cookie, ";") {
|
||||
t := strings.TrimSpace(p)
|
||||
if t == "" || strings.HasPrefix(t, "sb-superb-auth-token=") || strings.HasPrefix(t, "sb-superb-auth-token.") {
|
||||
continue
|
||||
}
|
||||
out = append(out, t)
|
||||
}
|
||||
if len(newValue) <= chunkSize {
|
||||
out = append(out, "sb-superb-auth-token="+newValue)
|
||||
} else {
|
||||
for i, off := 0, 0; off < len(newValue); i++ {
|
||||
end := off + chunkSize
|
||||
if end > len(newValue) {
|
||||
end = len(newValue)
|
||||
}
|
||||
out = append(out, fmt.Sprintf("sb-superb-auth-token.%d=%s", i, newValue[off:end]))
|
||||
off = end
|
||||
}
|
||||
}
|
||||
return strings.Join(out, "; ")
|
||||
}
|
||||
|
||||
// decodeSession base64-decodes the auth cookie value into the session JSON map.
|
||||
func decodeSession(authValue string) (map[string]any, bool) {
|
||||
v := strings.TrimPrefix(strings.TrimSpace(authValue), "base64-")
|
||||
if v == "" {
|
||||
return nil, false
|
||||
}
|
||||
raw, err := base64.StdEncoding.DecodeString(v)
|
||||
if err != nil {
|
||||
raw, err = base64.RawURLEncoding.DecodeString(v)
|
||||
if err != nil {
|
||||
return nil, false
|
||||
}
|
||||
}
|
||||
var m map[string]any
|
||||
if json.Unmarshal(raw, &m) != nil {
|
||||
return nil, false
|
||||
}
|
||||
return m, true
|
||||
}
|
||||
|
||||
func nestedStr(m map[string]any, k1, k2 string) string {
|
||||
if sub, ok := m[k1].(map[string]any); ok {
|
||||
return strings.TrimSpace(stringValue(sub[k2]))
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func toInt64(v any) int64 {
|
||||
switch x := v.(type) {
|
||||
case float64:
|
||||
return int64(x)
|
||||
case int64:
|
||||
return x
|
||||
case int:
|
||||
return int64(x)
|
||||
case json.Number:
|
||||
n, _ := x.Int64()
|
||||
return n
|
||||
case string:
|
||||
n, _ := strconv.ParseInt(strings.TrimSpace(x), 10, 64)
|
||||
return n
|
||||
default:
|
||||
return 0
|
||||
}
|
||||
}
|
||||
|
||||
func (c *Client) SetProxy(proxy string) {
|
||||
c.proxy = strings.TrimSpace(proxy)
|
||||
}
|
||||
|
||||
// IsKreaCookie reports whether a pasted credential is a Krea cookie: it carries
|
||||
// the Supabase auth cookie. Distinguishes it from adobe/leonardo cookies.
|
||||
func IsKreaCookie(value string) bool {
|
||||
return strings.Contains(value, "sb-superb-auth-token")
|
||||
}
|
||||
|
||||
// EmailFromCookie decodes the account email straight out of the cookie's embedded
|
||||
// Supabase session (no network), handling chunked cookies too.
|
||||
func EmailFromCookie(cookie string) string {
|
||||
sess, ok := decodeSession(authCookieValue(cookie))
|
||||
if !ok {
|
||||
return ""
|
||||
}
|
||||
return nestedStr(sess, "user", "email")
|
||||
}
|
||||
|
||||
// FetchCreditsBalance reads the account's free-credit balance via /api/billing-data.
|
||||
// remaining is the integer floor of balance.free (per spec: 17.94 → 17). 401 →
|
||||
// ErrAuth (cookie dead). Returns the normalized map shared by all providers.
|
||||
func (c *Client) FetchCreditsBalance(ctx context.Context, cookie string) (map[string]any, error) {
|
||||
// Detach from the request ctx so a page refresh can't cancel the probe
|
||||
// mid-flight (which left accounts stuck at "—").
|
||||
probeCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 30*time.Second)
|
||||
defer cancel()
|
||||
// NOTE: no /app here — the heavy SSR activation is done separately (on recovery
|
||||
// and at generation via Activate). This probe just reads the current balance.
|
||||
body, status, err := c.apiGet(probeCtx, cookie, "/api/billing-data")
|
||||
if err != nil {
|
||||
return unknownBalance("network: " + err.Error()), nil
|
||||
}
|
||||
if status == 401 || status == 403 {
|
||||
return nil, ErrAuth
|
||||
}
|
||||
if status != 200 {
|
||||
return unknownBalance(fmt.Sprintf("http %d: %s", status, clip(body, 160))), nil
|
||||
}
|
||||
// 余额在 balance.free(真实剩余,小数,随用量递减)。krea 的这个字段一直都在,
|
||||
// 只是排在很长的 entitlements 之后 —— 只读它,不用套餐配额兜底。
|
||||
var bd struct {
|
||||
Balance struct {
|
||||
Free float64 `json:"free"`
|
||||
Total float64 `json:"total"`
|
||||
} `json:"balance"`
|
||||
}
|
||||
if err := json.Unmarshal(body, &bd); err != nil {
|
||||
return unknownBalance("non-json"), nil
|
||||
}
|
||||
remaining := int(bd.Balance.Free) // floor
|
||||
return map[string]any{
|
||||
"remaining": remaining,
|
||||
"used": nil,
|
||||
"total": int(bd.Balance.Total),
|
||||
"unknown": false,
|
||||
"error": nil,
|
||||
"email": emptyStringNil(EmailFromCookie(cookie)),
|
||||
}, nil
|
||||
}
|
||||
|
||||
// Activate loads the authenticated SSR app page (/app), which is what makes krea
|
||||
// grant the account's DAILY free balance — a cold API-only call (billing-data /
|
||||
// generate) otherwise sees balance.free=0 and 402s. Done at most ONCE per account
|
||||
// per daily reset, under a per-account lock: the first caller loads /app while
|
||||
// concurrent callers wait, then everyone proceeds (no redundant /app). Called
|
||||
// before each generation and by the daily activation sweep. Best-effort.
|
||||
func (c *Client) Activate(ctx context.Context, cookie string) {
|
||||
key := accountKey(cookie)
|
||||
lastReset := (time.Now().Unix() / 86400) * 86400
|
||||
c.mu.Lock()
|
||||
doneToday := key != "" && c.actAt[key] >= lastReset
|
||||
c.mu.Unlock()
|
||||
if doneToday {
|
||||
return
|
||||
}
|
||||
lk := c.actLock(key)
|
||||
lk.Lock()
|
||||
defer lk.Unlock()
|
||||
// Re-check after acquiring the lock — another caller may have just activated.
|
||||
c.mu.Lock()
|
||||
doneToday = key != "" && c.actAt[key] >= lastReset
|
||||
c.mu.Unlock()
|
||||
if doneToday {
|
||||
return
|
||||
}
|
||||
actCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 60*time.Second)
|
||||
defer cancel()
|
||||
_, _, _ = c.apiGet(actCtx, cookie, "/app")
|
||||
c.mu.Lock()
|
||||
c.actAt[key] = time.Now().Unix()
|
||||
c.mu.Unlock()
|
||||
}
|
||||
|
||||
func (c *Client) actLock(key string) *sync.Mutex {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
m, ok := c.actLocks[key]
|
||||
if !ok {
|
||||
m = &sync.Mutex{}
|
||||
c.actLocks[key] = m
|
||||
}
|
||||
return m
|
||||
}
|
||||
|
||||
// accountKey is the stable per-account id (the Supabase user id from the session
|
||||
// cookie) used to key activation state — survives cookie rotation.
|
||||
func accountKey(cookie string) string {
|
||||
if sess, ok := decodeSession(authCookieValue(cookie)); ok {
|
||||
return nestedStr(sess, "user", "id")
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
// apiGet issues a GET to a krea.ai API path carrying the account cookie.
|
||||
func (c *Client) apiGet(ctx context.Context, cookie, path string) ([]byte, int, error) {
|
||||
client, err := c.newTLSClient()
|
||||
if err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
req, err := http.NewRequest(http.MethodGet, apiBase+path, nil)
|
||||
if err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
req = req.WithContext(ctx)
|
||||
req.Header = http.Header{
|
||||
"accept": {"*/*"},
|
||||
"accept-language": {"en-US,en;q=0.9"},
|
||||
"cookie": {cookie},
|
||||
"referer": {apiBase + "/"},
|
||||
"user-agent": {userAgent},
|
||||
"sec-fetch-dest": {"empty"},
|
||||
"sec-fetch-mode": {"cors"},
|
||||
"sec-fetch-site": {"same-origin"},
|
||||
http.HeaderOrderKey: {
|
||||
"accept", "accept-language", "cookie", "referer", "user-agent",
|
||||
"sec-fetch-dest", "sec-fetch-mode", "sec-fetch-site",
|
||||
},
|
||||
}
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
b, err := io.ReadAll(resp.Body)
|
||||
return b, resp.StatusCode, err
|
||||
}
|
||||
|
||||
func (c *Client) newTLSClient() (tlsclient.HttpClient, error) {
|
||||
options := []tlsclient.HttpClientOption{
|
||||
tlsclient.WithTimeoutSeconds(60),
|
||||
tlsclient.WithClientProfile(profiles.Chrome_120),
|
||||
}
|
||||
if c.proxy != "" {
|
||||
options = append(options, tlsclient.WithProxyUrl(c.proxy))
|
||||
}
|
||||
return tlsclient.NewHttpClient(tlsclient.NewNoopLogger(), options...)
|
||||
}
|
||||
|
||||
func unknownBalance(reason string) map[string]any {
|
||||
return map[string]any{
|
||||
"remaining": nil, "used": nil, "total": nil, "unknown": true, "error": reason,
|
||||
}
|
||||
}
|
||||
|
||||
func stringValue(v any) string {
|
||||
switch x := v.(type) {
|
||||
case string:
|
||||
return x
|
||||
case nil:
|
||||
return ""
|
||||
default:
|
||||
b, _ := json.Marshal(x)
|
||||
return strings.TrimSpace(string(b))
|
||||
}
|
||||
}
|
||||
|
||||
func intValue(v any) int {
|
||||
switch x := v.(type) {
|
||||
case int:
|
||||
return x
|
||||
case float64:
|
||||
return int(x)
|
||||
case json.Number:
|
||||
n, _ := x.Int64()
|
||||
return int(n)
|
||||
case string:
|
||||
n, _ := strconv.Atoi(strings.TrimSpace(x))
|
||||
return n
|
||||
default:
|
||||
return 0
|
||||
}
|
||||
}
|
||||
|
||||
func emptyStringNil(v string) any {
|
||||
if strings.TrimSpace(v) == "" {
|
||||
return nil
|
||||
}
|
||||
return v
|
||||
}
|
||||
|
||||
func clip(b []byte, n int) string {
|
||||
s := strings.TrimSpace(string(b))
|
||||
if len(s) > n {
|
||||
return s[:n]
|
||||
}
|
||||
return s
|
||||
}
|
||||
@@ -0,0 +1,308 @@
|
||||
package krea
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"mime/multipart"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
http "github.com/bogdanfinn/fhttp"
|
||||
)
|
||||
|
||||
const (
|
||||
genModel = "flux2-klein4b" // the single exposed Krea model
|
||||
genEndpoint = "/api/jobs/v2/new/fluxKlein4b"
|
||||
refStrength = 0.4
|
||||
)
|
||||
|
||||
// ensureProject returns a flux project id for the account: the first existing
|
||||
// project, or a freshly created one. Generation requires a project.
|
||||
func (c *Client) ensureProject(ctx context.Context, cookie string) (string, error) {
|
||||
body, status, err := c.apiGet(ctx, cookie, "/api/flux-projects")
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("%w: list projects: %s", ErrTemporaryUpstream, err.Error())
|
||||
}
|
||||
if status == 401 || status == 403 {
|
||||
return "", ErrAuth
|
||||
}
|
||||
if status == 200 {
|
||||
var projs []struct {
|
||||
ID string `json:"id"`
|
||||
}
|
||||
if json.Unmarshal(body, &projs) == nil {
|
||||
for _, p := range projs {
|
||||
if strings.TrimSpace(p.ID) != "" {
|
||||
return p.ID, nil
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
// No project on this account (e.g. brand-new). Try to create one; if that
|
||||
// fails, fall back to generating WITHOUT a project (Krea assigns a default).
|
||||
cb, cs, cerr := c.apiPostJSON(ctx, cookie, "/api/flux-projects", map[string]any{"title": "vivid"})
|
||||
if cerr == nil && (cs == 200 || cs == 201) {
|
||||
var pr struct {
|
||||
ID string `json:"id"`
|
||||
}
|
||||
if json.Unmarshal(cb, &pr) == nil && pr.ID != "" {
|
||||
return pr.ID, nil
|
||||
}
|
||||
}
|
||||
return "", nil // generate without an explicit project
|
||||
}
|
||||
|
||||
// uploadImage uploads a reference image (i2i) and returns its app-uploads URL.
|
||||
func (c *Client) uploadImage(ctx context.Context, cookie string, img []byte) (string, error) {
|
||||
var buf bytes.Buffer
|
||||
w := multipart.NewWriter(&buf)
|
||||
fw, err := w.CreateFormFile("file", "image.png")
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if _, err := fw.Write(img); err != nil {
|
||||
return "", err
|
||||
}
|
||||
_ = w.Close()
|
||||
body, status, err := c.apiPost(ctx, cookie, "/api/upload?", w.FormDataContentType(), buf.Bytes())
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("%w: upload: %s", ErrTemporaryUpstream, err.Error())
|
||||
}
|
||||
if status == 401 || status == 403 {
|
||||
return "", ErrAuth
|
||||
}
|
||||
if status != 200 {
|
||||
return "", fmt.Errorf("%w: upload http %d: %s", ErrTemporaryUpstream, status, clip(body, 160))
|
||||
}
|
||||
var ur struct {
|
||||
ImageURL string `json:"imageUrl"`
|
||||
}
|
||||
if json.Unmarshal(body, &ur) != nil || ur.ImageURL == "" {
|
||||
return "", fmt.Errorf("%w: no imageUrl", ErrTemporaryUpstream)
|
||||
}
|
||||
return ur.ImageURL, nil
|
||||
}
|
||||
|
||||
// GenerateImage runs the full Krea image pipeline: ensure a project, (for i2i)
|
||||
// upload reference images, submit the job, poll until done, then resolve and
|
||||
// download the produced image. 402 INSUFFICIENT_BALANCE → ErrQuotaExhausted.
|
||||
func (c *Client) GenerateImage(ctx context.Context, cookie, prompt string, width, height int, refImages [][]byte) ([]byte, map[string]any, error) {
|
||||
// Ensure the daily free balance is granted (load /app) before generating, so a
|
||||
// not-yet-activated account doesn't 402 INSUFFICIENT_BALANCE. Lock-guarded and
|
||||
// once-per-daily-reset — concurrent gens wait for the first activation, already
|
||||
// activated ones skip straight through.
|
||||
c.Activate(ctx, cookie)
|
||||
|
||||
projectID, err := c.ensureProject(ctx, cookie)
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
|
||||
var styleImages []map[string]any
|
||||
for _, img := range refImages {
|
||||
if len(img) == 0 {
|
||||
continue
|
||||
}
|
||||
url, upErr := c.uploadImage(ctx, cookie, img)
|
||||
if upErr != nil {
|
||||
return nil, nil, upErr
|
||||
}
|
||||
styleImages = append(styleImages, map[string]any{"url": url, "strength": refStrength, "source": "upload"})
|
||||
}
|
||||
|
||||
payload := map[string]any{
|
||||
"provider": genModel,
|
||||
"prompt": prompt,
|
||||
"width": width,
|
||||
"height": height,
|
||||
"strength": 1,
|
||||
"steps": 28,
|
||||
"guidance_scale_flux": 3.5,
|
||||
"presetStyles": []any{},
|
||||
"batchSize": 2,
|
||||
"guidance": 3.5,
|
||||
}
|
||||
if projectID != "" {
|
||||
payload["project"] = projectID
|
||||
}
|
||||
if len(styleImages) > 0 {
|
||||
payload["styleImages"] = styleImages
|
||||
}
|
||||
payloadJSON, _ := json.Marshal(payload)
|
||||
|
||||
// Submit (multipart with a single "payload" field).
|
||||
var buf bytes.Buffer
|
||||
w := multipart.NewWriter(&buf)
|
||||
_ = w.WriteField("payload", string(payloadJSON))
|
||||
_ = w.Close()
|
||||
body, status, err := c.apiPost(ctx, cookie, genEndpoint, w.FormDataContentType(), buf.Bytes())
|
||||
if err != nil {
|
||||
return nil, nil, fmt.Errorf("%w: submit: %s", ErrTemporaryUpstream, err.Error())
|
||||
}
|
||||
if status == 401 || status == 403 {
|
||||
return nil, nil, ErrAuth
|
||||
}
|
||||
if status == 402 || strings.Contains(string(body), "INSUFFICIENT_BALANCE") {
|
||||
return nil, nil, ErrQuotaExhausted
|
||||
}
|
||||
if status != 200 && status != 201 {
|
||||
return nil, nil, fmt.Errorf("%w: submit http %d: %s", ErrTemporaryUpstream, status, clip(body, 200))
|
||||
}
|
||||
var jobs []struct {
|
||||
JobID string `json:"job_id"`
|
||||
}
|
||||
if err := json.Unmarshal(body, &jobs); err != nil || len(jobs) == 0 || jobs[0].JobID == "" {
|
||||
return nil, nil, fmt.Errorf("%w: no job id: %s", ErrTemporaryUpstream, clip(body, 200))
|
||||
}
|
||||
// batchSize=2 returns two jobs; keep the SECOND image and discard the first
|
||||
// (per spec). Fall back to the first if only one came back.
|
||||
jobID := jobs[0].JobID
|
||||
if len(jobs) >= 2 && jobs[1].JobID != "" {
|
||||
jobID = jobs[1].JobID
|
||||
}
|
||||
|
||||
// Poll until terminal, then resolve the produced image.
|
||||
imageURL, err := c.pollImage(ctx, cookie, jobID)
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
data, err := c.download(ctx, imageURL)
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
return data, map[string]any{"job_id": jobID, "image_url": imageURL, "project": projectID}, nil
|
||||
}
|
||||
|
||||
// pollImage polls job-status until the job leaves the queue, then matches the
|
||||
// produced asset by generation_job_id and returns its image URL.
|
||||
func (c *Client) pollImage(ctx context.Context, cookie, jobID string) (string, error) {
|
||||
ticker := time.NewTicker(3 * time.Second)
|
||||
defer ticker.Stop()
|
||||
deadline := time.Now().Add(4 * time.Minute)
|
||||
|
||||
for {
|
||||
body, status, err := c.apiGet(ctx, cookie, "/api/job-status?id="+jobID)
|
||||
if err == nil && status == 200 {
|
||||
var js struct {
|
||||
Status string `json:"status"`
|
||||
}
|
||||
if json.Unmarshal(body, &js) == nil {
|
||||
switch strings.ToLower(js.Status) {
|
||||
case "complete", "completed", "succeeded", "success", "done", "finished":
|
||||
if url, e := c.assetForJob(ctx, cookie, jobID); e == nil && url != "" {
|
||||
return url, nil
|
||||
}
|
||||
case "failed", "error", "cancelled", "canceled":
|
||||
return "", fmt.Errorf("%w: job %s", ErrTemporaryUpstream, js.Status)
|
||||
}
|
||||
}
|
||||
} else if status == 401 || status == 403 {
|
||||
return "", ErrAuth
|
||||
}
|
||||
if time.Now().After(deadline) {
|
||||
return "", fmt.Errorf("%w: generation timed out", ErrTemporaryUpstream)
|
||||
}
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return "", ctx.Err()
|
||||
case <-ticker.C:
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// assetForJob finds the generated asset produced by a job and returns its URL.
|
||||
func (c *Client) assetForJob(ctx context.Context, cookie, jobID string) (string, error) {
|
||||
body, status, err := c.apiGet(ctx, cookie, "/api/assets?filter=generated&offset=0")
|
||||
if err != nil || status != 200 {
|
||||
return "", fmt.Errorf("assets http %d", status)
|
||||
}
|
||||
var assets []struct {
|
||||
ImageURL string `json:"image_url"`
|
||||
Metadata struct {
|
||||
GenerationJobID string `json:"generation_job_id"`
|
||||
} `json:"metadata"`
|
||||
}
|
||||
if json.Unmarshal(body, &assets) != nil {
|
||||
return "", fmt.Errorf("assets non-json")
|
||||
}
|
||||
for _, a := range assets {
|
||||
if a.Metadata.GenerationJobID == jobID && strings.TrimSpace(a.ImageURL) != "" {
|
||||
return a.ImageURL, nil
|
||||
}
|
||||
}
|
||||
return "", fmt.Errorf("asset not found yet")
|
||||
}
|
||||
|
||||
func (c *Client) download(ctx context.Context, url string) ([]byte, error) {
|
||||
client, err := c.newTLSClient()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
req, err := http.NewRequest(http.MethodGet, url, nil)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
req = req.WithContext(ctx)
|
||||
req.Header = http.Header{
|
||||
"accept": {"image/avif,image/webp,image/png,image/*,*/*;q=0.8"},
|
||||
"user-agent": {userAgent},
|
||||
"referer": {apiBase + "/"},
|
||||
}
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
b, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if resp.StatusCode != 200 {
|
||||
return nil, fmt.Errorf("%w: image download http %d", ErrTemporaryUpstream, resp.StatusCode)
|
||||
}
|
||||
return b, nil
|
||||
}
|
||||
|
||||
// apiPost issues a POST with a raw body + content-type, carrying the cookie.
|
||||
func (c *Client) apiPost(ctx context.Context, cookie, path, contentType string, body []byte) ([]byte, int, error) {
|
||||
client, err := c.newTLSClient()
|
||||
if err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
req, err := http.NewRequest(http.MethodPost, apiBase+path, bytes.NewReader(body))
|
||||
if err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
req = req.WithContext(ctx)
|
||||
req.Header = http.Header{
|
||||
"accept": {"*/*"},
|
||||
"accept-language": {"en-US,en;q=0.9"},
|
||||
"content-type": {contentType},
|
||||
"cookie": {cookie},
|
||||
"origin": {apiBase},
|
||||
"referer": {apiBase + "/"},
|
||||
"user-agent": {userAgent},
|
||||
"sec-fetch-dest": {"empty"},
|
||||
"sec-fetch-mode": {"cors"},
|
||||
"sec-fetch-site": {"same-origin"},
|
||||
http.HeaderOrderKey: {
|
||||
"accept", "accept-language", "content-type", "cookie", "origin",
|
||||
"referer", "user-agent", "sec-fetch-dest", "sec-fetch-mode", "sec-fetch-site",
|
||||
},
|
||||
}
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
b, err := io.ReadAll(resp.Body)
|
||||
return b, resp.StatusCode, err
|
||||
}
|
||||
|
||||
func (c *Client) apiPostJSON(ctx context.Context, cookie, path string, payload any) ([]byte, int, error) {
|
||||
b, _ := json.Marshal(payload)
|
||||
return c.apiPost(ctx, cookie, path, "application/json", b)
|
||||
}
|
||||
@@ -0,0 +1,396 @@
|
||||
// Package leonardo implements the Leonardo.ai (app.leonardo.ai) provider client.
|
||||
// Unlike chatgpt/runway (whose JWT IS the stored credential), Leonardo's durable
|
||||
// credential is the browser COOKIE (better-auth session): the bearer access token
|
||||
// it mints lives only ~1h. So every call here takes the cookie and derives a
|
||||
// fresh JWT on the fly via /api/auth/get-session — there is no long-lived token to
|
||||
// store or a separate refresh profile to maintain. tls-client gives a Chrome
|
||||
// JA3/JA4 fingerprint so the requests aren't flagged.
|
||||
package leonardo
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/url"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
http "github.com/bogdanfinn/fhttp"
|
||||
tlsclient "github.com/bogdanfinn/tls-client"
|
||||
"github.com/bogdanfinn/tls-client/profiles"
|
||||
)
|
||||
|
||||
const (
|
||||
appBase = "https://app.leonardo.ai"
|
||||
graphqlURL = "https://api.leonardo.ai/v1/graphql"
|
||||
schemaVersion = "1.187.0"
|
||||
userAgent = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/133.0.0.0 Safari/537.36"
|
||||
)
|
||||
|
||||
var (
|
||||
ErrAuth = errors.New("leonardo auth failed")
|
||||
ErrQuotaExhausted = errors.New("leonardo quota exhausted")
|
||||
ErrTemporaryUpstream = errors.New("leonardo upstream temporary error")
|
||||
)
|
||||
|
||||
type Client struct {
|
||||
proxy string
|
||||
// sessions caches the short-lived access token per cookie so we don't hit
|
||||
// /api/auth/get-session on every call — Leonardo rate-limits that endpoint
|
||||
// (429) hard, so re-using the ~1h JWT is essential.
|
||||
mu sync.Mutex
|
||||
sessions map[string]*Session
|
||||
}
|
||||
|
||||
func NewClient(proxy string) *Client {
|
||||
return &Client{proxy: strings.TrimSpace(proxy), sessions: map[string]*Session{}}
|
||||
}
|
||||
|
||||
func (c *Client) SetProxy(proxy string) {
|
||||
c.proxy = strings.TrimSpace(proxy)
|
||||
}
|
||||
|
||||
// IsLeonardoCookie reports whether a pasted credential is a Leonardo cookie: it
|
||||
// carries the better-auth session cookie name. This is what disambiguates it from
|
||||
// an Adobe cookie at import time.
|
||||
func IsLeonardoCookie(value string) bool {
|
||||
return strings.Contains(value, "__Secure-better-auth.session_token") ||
|
||||
strings.Contains(value, "better-auth.session_data")
|
||||
}
|
||||
|
||||
// Session is the result of /api/auth/get-session: the short-lived bearer plus the
|
||||
// ids the GraphQL API needs (cognitoSub for the quota query, userId for the feed
|
||||
// and the CDN image path) and the human-facing account fields.
|
||||
type Session struct {
|
||||
AccessToken string
|
||||
CognitoSub string
|
||||
UserID string
|
||||
Email string
|
||||
Name string
|
||||
ExpiresAt int64
|
||||
}
|
||||
|
||||
// GetSession exchanges the cookie for a fresh access token + account ids. A 401/403
|
||||
// (or a response with no access token) means the cookie/session is dead → ErrAuth.
|
||||
func (c *Client) GetSession(ctx context.Context, cookie string) (*Session, error) {
|
||||
cookie = strings.TrimSpace(cookie)
|
||||
if cookie == "" {
|
||||
return nil, ErrAuth
|
||||
}
|
||||
// Re-use a cached, still-valid access token (keep a 60s safety margin) instead
|
||||
// of hitting the heavily rate-limited get-session endpoint again.
|
||||
c.mu.Lock()
|
||||
if cs, ok := c.sessions[cookie]; ok && cs.ExpiresAt-60 > time.Now().Unix() {
|
||||
c.mu.Unlock()
|
||||
return cs, nil
|
||||
}
|
||||
c.mu.Unlock()
|
||||
|
||||
client, err := c.newTLSClient()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
req, err := http.NewRequest(http.MethodGet, appBase+"/api/auth/get-session", nil)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
req = req.WithContext(ctx)
|
||||
req.Header = http.Header{
|
||||
"accept": {"*/*"},
|
||||
"accept-language": {"en-US,en;q=0.9"},
|
||||
"cookie": {cookie},
|
||||
"origin": {appBase},
|
||||
"referer": {appBase + "/"},
|
||||
"user-agent": {userAgent},
|
||||
"sec-fetch-dest": {"empty"},
|
||||
"sec-fetch-mode": {"cors"},
|
||||
"sec-fetch-site": {"same-origin"},
|
||||
http.HeaderOrderKey: {
|
||||
"accept", "accept-language", "cookie", "origin", "referer",
|
||||
"user-agent", "sec-fetch-dest", "sec-fetch-mode", "sec-fetch-site",
|
||||
},
|
||||
}
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("%w: %s", ErrTemporaryUpstream, err.Error())
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
body, _ := io.ReadAll(resp.Body)
|
||||
if resp.StatusCode == 401 || resp.StatusCode == 403 {
|
||||
return nil, ErrAuth
|
||||
}
|
||||
if resp.StatusCode != 200 {
|
||||
return nil, fmt.Errorf("%w: get-session http %d: %s", ErrTemporaryUpstream, resp.StatusCode, clip(body, 160))
|
||||
}
|
||||
var raw struct {
|
||||
Session struct {
|
||||
AccessToken string `json:"accessToken"`
|
||||
CognitoSub string `json:"cognitoSub"`
|
||||
UserID string `json:"userId"`
|
||||
HasuraUserID string `json:"hasuraUserId"`
|
||||
TokenExpiry int64 `json:"accessTokenExpiry"`
|
||||
} `json:"session"`
|
||||
User struct {
|
||||
ID string `json:"id"`
|
||||
Email string `json:"email"`
|
||||
Name string `json:"name"`
|
||||
} `json:"user"`
|
||||
}
|
||||
if err := json.Unmarshal(body, &raw); err != nil {
|
||||
return nil, fmt.Errorf("%w: get-session non-json", ErrTemporaryUpstream)
|
||||
}
|
||||
if strings.TrimSpace(raw.Session.AccessToken) == "" {
|
||||
// No bearer despite 200 → the cookie no longer authenticates.
|
||||
return nil, ErrAuth
|
||||
}
|
||||
uid := raw.Session.UserID
|
||||
if uid == "" {
|
||||
uid = raw.Session.HasuraUserID
|
||||
}
|
||||
if uid == "" {
|
||||
uid = raw.User.ID
|
||||
}
|
||||
sess := &Session{
|
||||
AccessToken: raw.Session.AccessToken,
|
||||
CognitoSub: raw.Session.CognitoSub,
|
||||
UserID: uid,
|
||||
Email: strings.TrimSpace(raw.User.Email),
|
||||
Name: strings.TrimSpace(raw.User.Name),
|
||||
ExpiresAt: raw.Session.TokenExpiry,
|
||||
}
|
||||
if sess.ExpiresAt > time.Now().Unix() {
|
||||
c.mu.Lock()
|
||||
c.sessions[cookie] = sess
|
||||
c.mu.Unlock()
|
||||
}
|
||||
return sess, nil
|
||||
}
|
||||
|
||||
const qGetTokens = `query GetUserTokensFromSub($sub: String) {
|
||||
user_details(where: {cognitoId: {_eq: $sub}}) {
|
||||
id
|
||||
plan
|
||||
subscriptionTokens
|
||||
paidTokens
|
||||
rolloverTokens
|
||||
tokenRenewalDate
|
||||
__typename
|
||||
}
|
||||
}`
|
||||
|
||||
// FetchCreditsBalance derives a JWT from the cookie then reads the account's image
|
||||
// token balance. Returns a normalized map mirroring the other providers so the
|
||||
// TokenService quota plumbing is uniform. remaining = subscription+paid+rollover
|
||||
// (the spendable image tokens); available_until carries the daily renewal time so
|
||||
// the maintenance sweep can auto-recover a 限额 account.
|
||||
func (c *Client) FetchCreditsBalance(ctx context.Context, cookie string) (map[string]any, error) {
|
||||
sess, err := c.GetSession(ctx, cookie)
|
||||
if err != nil {
|
||||
if errors.Is(err, ErrAuth) {
|
||||
return nil, ErrAuth
|
||||
}
|
||||
return unknownBalance(err.Error()), nil
|
||||
}
|
||||
if sess.CognitoSub == "" {
|
||||
return unknownBalance("no cognitoSub"), nil
|
||||
}
|
||||
|
||||
payload, _ := json.Marshal(map[string]any{
|
||||
"operationName": "GetUserTokensFromSub",
|
||||
"variables": map[string]any{"sub": sess.CognitoSub},
|
||||
"query": qGetTokens,
|
||||
})
|
||||
body, status, err := c.graphql(ctx, sess.AccessToken, payload)
|
||||
if err != nil {
|
||||
return unknownBalance("network: " + err.Error()), nil
|
||||
}
|
||||
if status == 401 || status == 403 {
|
||||
return nil, ErrAuth
|
||||
}
|
||||
if status != 200 {
|
||||
return unknownBalance(fmt.Sprintf("http %d: %s", status, clip(body, 160))), nil
|
||||
}
|
||||
var result struct {
|
||||
Data struct {
|
||||
UserDetails []struct {
|
||||
Plan string `json:"plan"`
|
||||
SubscriptionTokens int `json:"subscriptionTokens"`
|
||||
PaidTokens int `json:"paidTokens"`
|
||||
RolloverTokens int `json:"rolloverTokens"`
|
||||
TokenRenewalDate string `json:"tokenRenewalDate"`
|
||||
} `json:"user_details"`
|
||||
} `json:"data"`
|
||||
}
|
||||
if err := json.Unmarshal(body, &result); err != nil {
|
||||
return unknownBalance("non-json"), nil
|
||||
}
|
||||
if len(result.Data.UserDetails) == 0 {
|
||||
return unknownBalance("no user_details"), nil
|
||||
}
|
||||
ud := result.Data.UserDetails[0]
|
||||
remaining := ud.SubscriptionTokens + ud.PaidTokens + ud.RolloverTokens
|
||||
return map[string]any{
|
||||
"remaining": remaining,
|
||||
"used": nil,
|
||||
"total": nil,
|
||||
"unknown": false,
|
||||
"error": nil,
|
||||
"plan": ud.Plan,
|
||||
"available_until": strings.TrimSpace(ud.TokenRenewalDate),
|
||||
"email": emptyStringNil(sess.Email),
|
||||
"display_name": emptyStringNil(sess.Name),
|
||||
"user_id": emptyStringNil(sess.UserID),
|
||||
}, nil
|
||||
}
|
||||
|
||||
// graphql POSTs a GraphQL body to the Leonardo API with the bearer + schema header,
|
||||
// returning the raw response body and status.
|
||||
func (c *Client) graphql(ctx context.Context, accessToken string, payload []byte) ([]byte, int, error) {
|
||||
client, err := c.newTLSClient()
|
||||
if err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
req, err := http.NewRequest(http.MethodPost, graphqlURL, bytes.NewReader(payload))
|
||||
if err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
req = req.WithContext(ctx)
|
||||
req.Header = http.Header{
|
||||
"content-type": {"application/json"},
|
||||
"accept": {"*/*"},
|
||||
"accept-language": {"en-US,en;q=0.9"},
|
||||
"origin": {appBase},
|
||||
"referer": {appBase + "/"},
|
||||
"user-agent": {userAgent},
|
||||
"authorization": {"Bearer " + accessToken},
|
||||
"x-leo-schema-version": {schemaVersion},
|
||||
"sec-fetch-dest": {"empty"},
|
||||
"sec-fetch-mode": {"cors"},
|
||||
"sec-fetch-site": {"same-site"},
|
||||
http.HeaderOrderKey: {
|
||||
"content-type", "accept", "accept-language", "origin", "referer",
|
||||
"user-agent", "authorization", "x-leo-schema-version",
|
||||
"sec-fetch-dest", "sec-fetch-mode", "sec-fetch-site",
|
||||
},
|
||||
}
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
body, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return nil, resp.StatusCode, err
|
||||
}
|
||||
return body, resp.StatusCode, nil
|
||||
}
|
||||
|
||||
func unknownBalance(reason string) map[string]any {
|
||||
return map[string]any{
|
||||
"remaining": nil,
|
||||
"used": nil,
|
||||
"total": nil,
|
||||
"unknown": true,
|
||||
"error": reason,
|
||||
}
|
||||
}
|
||||
|
||||
func (c *Client) newTLSClient() (tlsclient.HttpClient, error) {
|
||||
// Match the fingerprint proven to work against Leonardo's Cloudflare edge:
|
||||
// Chrome_120, fixed extension order. A randomized JA3 (Chrome_133 +
|
||||
// WithRandomTLSExtensionOrder) gets flagged and 429'd at get-session.
|
||||
options := []tlsclient.HttpClientOption{
|
||||
tlsclient.WithTimeoutSeconds(60),
|
||||
tlsclient.WithClientProfile(profiles.Chrome_120),
|
||||
}
|
||||
if c.proxy != "" {
|
||||
options = append(options, tlsclient.WithProxyUrl(c.proxy))
|
||||
}
|
||||
return tlsclient.NewHttpClient(tlsclient.NewNoopLogger(), options...)
|
||||
}
|
||||
|
||||
// downloadImage fetches a generated image (cdn.leonardo.ai) and returns the bytes.
|
||||
func (c *Client) downloadImage(ctx context.Context, imageURL string) ([]byte, error) {
|
||||
if _, err := url.Parse(imageURL); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
client, err := c.newTLSClient()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
req, err := http.NewRequest(http.MethodGet, imageURL, nil)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
req = req.WithContext(ctx)
|
||||
req.Header = http.Header{
|
||||
"accept": {"image/avif,image/webp,image/png,image/*,*/*;q=0.8"},
|
||||
"user-agent": {userAgent},
|
||||
"referer": {appBase + "/"},
|
||||
}
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
body, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if resp.StatusCode != 200 {
|
||||
return nil, fmt.Errorf("%w: image download http %d", ErrTemporaryUpstream, resp.StatusCode)
|
||||
}
|
||||
return body, nil
|
||||
}
|
||||
|
||||
func stringValue(v any) string {
|
||||
switch x := v.(type) {
|
||||
case string:
|
||||
return x
|
||||
case nil:
|
||||
return ""
|
||||
default:
|
||||
b, _ := json.Marshal(x)
|
||||
return strings.TrimSpace(string(b))
|
||||
}
|
||||
}
|
||||
|
||||
func intValue(v any) int {
|
||||
switch x := v.(type) {
|
||||
case int:
|
||||
return x
|
||||
case int64:
|
||||
return int(x)
|
||||
case float64:
|
||||
return int(x)
|
||||
case json.Number:
|
||||
n, _ := x.Int64()
|
||||
return int(n)
|
||||
case string:
|
||||
n, _ := strconv.Atoi(strings.TrimSpace(x))
|
||||
return n
|
||||
default:
|
||||
return 0
|
||||
}
|
||||
}
|
||||
|
||||
func emptyStringNil(v string) any {
|
||||
v = strings.TrimSpace(v)
|
||||
if v == "" {
|
||||
return nil
|
||||
}
|
||||
return v
|
||||
}
|
||||
|
||||
func clip(b []byte, n int) string {
|
||||
s := strings.TrimSpace(string(b))
|
||||
if len(s) > n {
|
||||
return s[:n]
|
||||
}
|
||||
return s
|
||||
}
|
||||
@@ -0,0 +1,326 @@
|
||||
package leonardo
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"mime/multipart"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
http "github.com/bogdanfinn/fhttp"
|
||||
)
|
||||
|
||||
// defaultStyleID is the "Dynamic" style applied when the caller doesn't specify
|
||||
// one — Leonardo's Generate mutation expects a style_ids entry.
|
||||
const defaultStyleID = "111dc692-d470-4eec-b791-3475abac4c46"
|
||||
|
||||
const mGenerate = `mutation Generate($request: CreateGenerationRequest!) {
|
||||
generate(request: $request) {
|
||||
apiCreditCost
|
||||
generationId
|
||||
__typename
|
||||
}
|
||||
}`
|
||||
|
||||
// qGenerationImages polls one generation's status AND its produced images in a
|
||||
// single round-trip (where: id _in [genId]).
|
||||
const qGenerationImages = `query GenerationImages($where: generations_bool_exp = {}) {
|
||||
generations(where: $where) {
|
||||
id
|
||||
status
|
||||
generated_images {
|
||||
id
|
||||
url
|
||||
__typename
|
||||
}
|
||||
__typename
|
||||
}
|
||||
}`
|
||||
|
||||
const mUploadImage = `mutation UploadImage($uploadImageInput: UploadImageInput!) {
|
||||
uploadImage(arg1: $uploadImageInput) {
|
||||
uploadId
|
||||
url
|
||||
fields
|
||||
__typename
|
||||
}
|
||||
}`
|
||||
|
||||
// uploadInitImage uploads a reference (init) image for image-to-image: it asks
|
||||
// Leonardo for a presigned S3 POST, uploads the bytes, and returns the upload id
|
||||
// to reference in the Generate request's image_reference guidance.
|
||||
func (c *Client) uploadInitImage(ctx context.Context, accessToken string, img []byte) (string, error) {
|
||||
payload, _ := json.Marshal(map[string]any{
|
||||
"operationName": "UploadImage",
|
||||
"query": mUploadImage,
|
||||
"variables": map[string]any{"uploadImageInput": map[string]any{"uploadType": "INIT", "extension": "png"}},
|
||||
})
|
||||
body, status, err := c.graphql(ctx, accessToken, payload)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("%w: upload-init: %s", ErrTemporaryUpstream, err.Error())
|
||||
}
|
||||
if status == 401 || status == 403 {
|
||||
return "", ErrAuth
|
||||
}
|
||||
if status != 200 {
|
||||
return "", fmt.Errorf("%w: upload-init http %d: %s", ErrTemporaryUpstream, status, clip(body, 160))
|
||||
}
|
||||
if e := graphqlError(body); e != nil {
|
||||
return "", e
|
||||
}
|
||||
var ur struct {
|
||||
Data struct {
|
||||
UploadImage struct {
|
||||
UploadID string `json:"uploadId"`
|
||||
URL string `json:"url"`
|
||||
Fields string `json:"fields"`
|
||||
} `json:"uploadImage"`
|
||||
} `json:"data"`
|
||||
}
|
||||
if err := json.Unmarshal(body, &ur); err != nil {
|
||||
return "", fmt.Errorf("%w: upload-init non-json", ErrTemporaryUpstream)
|
||||
}
|
||||
up := ur.Data.UploadImage
|
||||
if up.UploadID == "" || up.URL == "" {
|
||||
return "", fmt.Errorf("%w: no upload url", ErrTemporaryUpstream)
|
||||
}
|
||||
var fields map[string]string
|
||||
if err := json.Unmarshal([]byte(up.Fields), &fields); err != nil {
|
||||
return "", fmt.Errorf("%w: bad upload fields", ErrTemporaryUpstream)
|
||||
}
|
||||
|
||||
// Presigned S3 POST: all policy fields first, the file part LAST.
|
||||
var buf bytes.Buffer
|
||||
w := multipart.NewWriter(&buf)
|
||||
for k, v := range fields {
|
||||
_ = w.WriteField(k, v)
|
||||
}
|
||||
fw, err := w.CreateFormFile("file", "image.png")
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if _, err := fw.Write(img); err != nil {
|
||||
return "", err
|
||||
}
|
||||
_ = w.Close()
|
||||
|
||||
client, err := c.newTLSClient()
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
req, err := http.NewRequest(http.MethodPost, up.URL, &buf)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
req = req.WithContext(ctx)
|
||||
req.Header = http.Header{
|
||||
"content-type": {w.FormDataContentType()},
|
||||
"user-agent": {userAgent},
|
||||
"origin": {appBase},
|
||||
"referer": {appBase + "/"},
|
||||
}
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("%w: s3 upload: %s", ErrTemporaryUpstream, err.Error())
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode != 204 && resp.StatusCode != 200 && resp.StatusCode != 201 {
|
||||
return "", fmt.Errorf("%w: s3 upload http %d", ErrTemporaryUpstream, resp.StatusCode)
|
||||
}
|
||||
return up.UploadID, nil
|
||||
}
|
||||
|
||||
// GenerateImage runs the full Leonardo image pipeline against one account cookie:
|
||||
// mint a JWT, (for image-to-image) upload each reference image, submit the
|
||||
// Generate mutation, poll until COMPLETE, then download the first produced image.
|
||||
// Returns the image bytes, an info map, and a classified error.
|
||||
func (c *Client) GenerateImage(ctx context.Context, cookie, model, prompt string, width, height int, styleIDs []string, refImages [][]byte) ([]byte, map[string]any, error) {
|
||||
sess, err := c.GetSession(ctx, cookie)
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
if len(styleIDs) == 0 {
|
||||
styleIDs = []string{defaultStyleID}
|
||||
}
|
||||
if strings.TrimSpace(model) == "" {
|
||||
model = "seedream-4.5"
|
||||
}
|
||||
|
||||
// Image-to-image: upload each reference and collect its guidance entry.
|
||||
var imageRefs []map[string]any
|
||||
for _, img := range refImages {
|
||||
if len(img) == 0 {
|
||||
continue
|
||||
}
|
||||
uploadID, upErr := c.uploadInitImage(ctx, sess.AccessToken, img)
|
||||
if upErr != nil {
|
||||
return nil, nil, upErr
|
||||
}
|
||||
imageRefs = append(imageRefs, map[string]any{
|
||||
"image": map[string]any{"id": uploadID, "type": "UPLOADED"},
|
||||
"strength": "MID",
|
||||
})
|
||||
}
|
||||
|
||||
promptEnhance := "AUTO"
|
||||
parameters := map[string]any{
|
||||
"height": height,
|
||||
"width": width,
|
||||
"prompt_enhance": promptEnhance,
|
||||
"quantity": 1,
|
||||
"style_ids": styleIDs,
|
||||
"prompt": prompt,
|
||||
}
|
||||
if len(imageRefs) > 0 {
|
||||
// Preserve the reference when image-guided (matches the web app).
|
||||
parameters["prompt_enhance"] = "OFF"
|
||||
parameters["guidances"] = map[string]any{"image_reference": imageRefs}
|
||||
}
|
||||
|
||||
// 1. submit
|
||||
genReq := map[string]any{
|
||||
"operationName": "Generate",
|
||||
"query": mGenerate,
|
||||
"variables": map[string]any{
|
||||
"request": map[string]any{
|
||||
"model": model,
|
||||
"public": true,
|
||||
"parameters": parameters,
|
||||
},
|
||||
},
|
||||
}
|
||||
payload, _ := json.Marshal(genReq)
|
||||
body, status, err := c.graphql(ctx, sess.AccessToken, payload)
|
||||
if err != nil {
|
||||
return nil, nil, fmt.Errorf("%w: %s", ErrTemporaryUpstream, err.Error())
|
||||
}
|
||||
if status == 401 || status == 403 {
|
||||
return nil, nil, ErrAuth
|
||||
}
|
||||
if status != 200 {
|
||||
return nil, nil, fmt.Errorf("%w: generate http %d: %s", ErrTemporaryUpstream, status, clip(body, 200))
|
||||
}
|
||||
if e := graphqlError(body); e != nil {
|
||||
return nil, nil, e
|
||||
}
|
||||
var genResp struct {
|
||||
Data struct {
|
||||
Generate struct {
|
||||
GenerationID string `json:"generationId"`
|
||||
} `json:"generate"`
|
||||
} `json:"data"`
|
||||
}
|
||||
if err := json.Unmarshal(body, &genResp); err != nil {
|
||||
return nil, nil, fmt.Errorf("%w: generate non-json", ErrTemporaryUpstream)
|
||||
}
|
||||
genID := strings.TrimSpace(genResp.Data.Generate.GenerationID)
|
||||
if genID == "" {
|
||||
return nil, nil, fmt.Errorf("%w: no generationId: %s", ErrTemporaryUpstream, clip(body, 200))
|
||||
}
|
||||
|
||||
// 2. poll until COMPLETE, then read the image url.
|
||||
imageURL, err := c.pollImage(ctx, sess.AccessToken, genID)
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
|
||||
// 3. download bytes
|
||||
data, err := c.downloadImage(ctx, imageURL)
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
info := map[string]any{
|
||||
"generation_id": genID,
|
||||
"image_url": imageURL,
|
||||
"user_id": sess.UserID,
|
||||
}
|
||||
return data, info, nil
|
||||
}
|
||||
|
||||
// pollImage polls one generation until it reports COMPLETE (returning the first
|
||||
// image url) or FAILED (error). Honors ctx cancellation / deadline.
|
||||
func (c *Client) pollImage(ctx context.Context, accessToken, genID string) (string, error) {
|
||||
payload, _ := json.Marshal(map[string]any{
|
||||
"operationName": "GenerationImages",
|
||||
"query": qGenerationImages,
|
||||
"variables": map[string]any{
|
||||
"where": map[string]any{"id": map[string]any{"_in": []string{genID}}},
|
||||
},
|
||||
})
|
||||
|
||||
ticker := time.NewTicker(3 * time.Second)
|
||||
defer ticker.Stop()
|
||||
// Cap the wait independent of the parent deadline so a stuck job can't hang.
|
||||
deadline := time.Now().Add(5 * time.Minute)
|
||||
|
||||
for {
|
||||
body, status, err := c.graphql(ctx, accessToken, payload)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("%w: poll: %s", ErrTemporaryUpstream, err.Error())
|
||||
}
|
||||
if status == 401 || status == 403 {
|
||||
return "", ErrAuth
|
||||
}
|
||||
if status == 200 {
|
||||
var pr struct {
|
||||
Data struct {
|
||||
Generations []struct {
|
||||
Status string `json:"status"`
|
||||
GeneratedImages []struct {
|
||||
URL string `json:"url"`
|
||||
} `json:"generated_images"`
|
||||
} `json:"generations"`
|
||||
} `json:"data"`
|
||||
}
|
||||
if err := json.Unmarshal(body, &pr); err == nil && len(pr.Data.Generations) > 0 {
|
||||
g := pr.Data.Generations[0]
|
||||
switch strings.ToUpper(g.Status) {
|
||||
case "COMPLETE":
|
||||
for _, img := range g.GeneratedImages {
|
||||
if u := strings.TrimSpace(img.URL); u != "" {
|
||||
return u, nil
|
||||
}
|
||||
}
|
||||
return "", fmt.Errorf("%w: complete but no image url", ErrTemporaryUpstream)
|
||||
case "FAILED":
|
||||
return "", fmt.Errorf("%w: generation failed", ErrTemporaryUpstream)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if time.Now().After(deadline) {
|
||||
return "", fmt.Errorf("%w: generation timed out", ErrTemporaryUpstream)
|
||||
}
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return "", ctx.Err()
|
||||
case <-ticker.C:
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// graphqlError inspects a GraphQL response body for an "errors" array and maps the
|
||||
// first message to a classified sentinel (auth / quota / temporary). Returns nil
|
||||
// when there are no errors.
|
||||
func graphqlError(body []byte) error {
|
||||
var env struct {
|
||||
Errors []struct {
|
||||
Message string `json:"message"`
|
||||
} `json:"errors"`
|
||||
}
|
||||
if err := json.Unmarshal(body, &env); err != nil || len(env.Errors) == 0 {
|
||||
return nil
|
||||
}
|
||||
msg := strings.TrimSpace(env.Errors[0].Message)
|
||||
low := strings.ToLower(msg)
|
||||
switch {
|
||||
case strings.Contains(low, "unauthor") || strings.Contains(low, "jwt") || strings.Contains(low, "token is") || strings.Contains(low, "forbidden"):
|
||||
return ErrAuth
|
||||
case strings.Contains(low, "token") || strings.Contains(low, "credit") || strings.Contains(low, "quota") || strings.Contains(low, "insufficient") || strings.Contains(low, "not enough"):
|
||||
return ErrQuotaExhausted
|
||||
default:
|
||||
return fmt.Errorf("leonardo: %s", clip([]byte(msg), 200))
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,254 @@
|
||||
// Package runway implements the Runway (runwayml.com) provider client. For now
|
||||
// it only covers account management — JWT detection, workspace/team id
|
||||
// extraction and credit-balance probing — mirroring the curl_cffi reference in
|
||||
// query_credits.py with tls-client so the JA3/JA4 fingerprint matches Chrome.
|
||||
package runway
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
http "github.com/bogdanfinn/fhttp"
|
||||
tlsclient "github.com/bogdanfinn/tls-client"
|
||||
"github.com/bogdanfinn/tls-client/profiles"
|
||||
)
|
||||
|
||||
const (
|
||||
apiBase = "https://api.runwayml.com"
|
||||
origin = "https://app.runwayml.com"
|
||||
)
|
||||
|
||||
var (
|
||||
ErrAuth = errors.New("runway auth failed")
|
||||
ErrQuotaExhausted = errors.New("runway quota exhausted")
|
||||
ErrTemporaryUpstream = errors.New("runway upstream temporary error")
|
||||
)
|
||||
|
||||
type Client struct {
|
||||
proxy string
|
||||
}
|
||||
|
||||
func NewClient(proxy string) *Client {
|
||||
return &Client{proxy: strings.TrimSpace(proxy)}
|
||||
}
|
||||
|
||||
func (c *Client) SetProxy(proxy string) {
|
||||
c.proxy = strings.TrimSpace(proxy)
|
||||
}
|
||||
|
||||
// IsRunwayToken reports whether a JWT looks like a Runway access token: a
|
||||
// top-level numeric "id" plus an "sso" claim, and crucially NO OpenAI
|
||||
// (https://api.openai.com/*) claims — that's what disambiguates it from a
|
||||
// ChatGPT token, which is otherwise also an opaque three-part JWT.
|
||||
func IsRunwayToken(token string) bool {
|
||||
claims := decodeJWTPayload(token)
|
||||
if len(claims) == 0 {
|
||||
return false
|
||||
}
|
||||
for k := range claims {
|
||||
if strings.HasPrefix(k, "https://api.openai.com/") {
|
||||
return false
|
||||
}
|
||||
}
|
||||
_, hasSSO := claims["sso"]
|
||||
return hasSSO && claims["id"] != nil
|
||||
}
|
||||
|
||||
// TeamIDFromToken returns the Runway workspace/team id, which equals the JWT
|
||||
// "id" claim (query_credits.py / gen_video.py both derive teamId this way).
|
||||
func TeamIDFromToken(token string) string {
|
||||
claims := decodeJWTPayload(token)
|
||||
switch v := claims["id"].(type) {
|
||||
case float64:
|
||||
return strconv.FormatInt(int64(v), 10)
|
||||
case json.Number:
|
||||
return v.String()
|
||||
case string:
|
||||
return strings.TrimSpace(v)
|
||||
default:
|
||||
return ""
|
||||
}
|
||||
}
|
||||
|
||||
// ExtractAccountInfo decodes the free (no-network) JWT claims for the accounts
|
||||
// view: email, team id and expiry.
|
||||
func ExtractAccountInfo(token string) map[string]any {
|
||||
claims := decodeJWTPayload(token)
|
||||
return map[string]any{
|
||||
"email": emptyStringNil(strings.TrimSpace(stringValue(claims["email"]))),
|
||||
"team_id": emptyStringNil(TeamIDFromToken(token)),
|
||||
"expires_at": claims["exp"],
|
||||
}
|
||||
}
|
||||
|
||||
// FetchCreditsBalance probes the account's plan credits via /v1/profile/features
|
||||
// (query_credits.py). Returns a normalized map mirroring the Adobe client so the
|
||||
// TokenService quota plumbing can treat all providers uniformly. A 401/403 maps
|
||||
// to ErrAuth (token dead); any other failure is reported as unknown without
|
||||
// killing the account.
|
||||
func (c *Client) FetchCreditsBalance(ctx context.Context, token string) (map[string]any, error) {
|
||||
token = strings.TrimSpace(token)
|
||||
if token == "" {
|
||||
return unknownBalance("empty token"), nil
|
||||
}
|
||||
teamID := TeamIDFromToken(token)
|
||||
if teamID == "" {
|
||||
return unknownBalance("no team id"), nil
|
||||
}
|
||||
|
||||
client, err := c.newTLSClient()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
url := apiBase + "/v1/profile/features?asTeamId=" + teamID
|
||||
req, err := http.NewRequest(http.MethodGet, url, nil)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
req = req.WithContext(ctx)
|
||||
req.Header = http.Header{
|
||||
"accept": {"application/json"},
|
||||
"content-type": {"application/json"},
|
||||
"origin": {origin},
|
||||
"referer": {origin + "/"},
|
||||
"authorization": {"Bearer " + token},
|
||||
"x-runway-workspace": {teamID},
|
||||
http.HeaderOrderKey: {
|
||||
"accept",
|
||||
"content-type",
|
||||
"origin",
|
||||
"referer",
|
||||
"authorization",
|
||||
"x-runway-workspace",
|
||||
},
|
||||
}
|
||||
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
return unknownBalance("network: " + err.Error()), nil
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
body, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if resp.StatusCode == 401 || resp.StatusCode == 403 {
|
||||
return nil, ErrAuth
|
||||
}
|
||||
if resp.StatusCode != 200 {
|
||||
return unknownBalance(fmt.Sprintf("http %d: %s", resp.StatusCode, clip(body, 160))), nil
|
||||
}
|
||||
|
||||
var payload map[string]any
|
||||
if err := json.Unmarshal(body, &payload); err != nil {
|
||||
return unknownBalance("non-json"), nil
|
||||
}
|
||||
features, _ := payload["features"].(map[string]any)
|
||||
permitted, _ := features["permitted"].(map[string]any)
|
||||
used, _ := features["used"].(map[string]any)
|
||||
total := intValue(permitted["numPlanCredits"])
|
||||
spent := intValue(used["numPlanCredits"])
|
||||
remaining := total - spent
|
||||
if remaining < 0 {
|
||||
remaining = 0
|
||||
}
|
||||
return map[string]any{
|
||||
"remaining": remaining,
|
||||
"used": spent,
|
||||
"total": total,
|
||||
"unknown": false,
|
||||
"error": nil,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func unknownBalance(reason string) map[string]any {
|
||||
return map[string]any{
|
||||
"remaining": nil,
|
||||
"used": nil,
|
||||
"total": nil,
|
||||
"unknown": true,
|
||||
"error": reason,
|
||||
}
|
||||
}
|
||||
|
||||
func (c *Client) newTLSClient() (tlsclient.HttpClient, error) {
|
||||
options := []tlsclient.HttpClientOption{
|
||||
tlsclient.WithTimeoutSeconds(30),
|
||||
tlsclient.WithClientProfile(profiles.Chrome_133),
|
||||
tlsclient.WithRandomTLSExtensionOrder(),
|
||||
}
|
||||
if c.proxy != "" {
|
||||
options = append(options, tlsclient.WithProxyUrl(c.proxy))
|
||||
}
|
||||
return tlsclient.NewHttpClient(tlsclient.NewNoopLogger(), options...)
|
||||
}
|
||||
|
||||
func decodeJWTPayload(token string) map[string]any {
|
||||
parts := strings.Split(strings.TrimSpace(strings.TrimPrefix(token, "Bearer ")), ".")
|
||||
if len(parts) < 2 {
|
||||
return map[string]any{}
|
||||
}
|
||||
raw, err := base64.RawURLEncoding.DecodeString(parts[1])
|
||||
if err != nil {
|
||||
return map[string]any{}
|
||||
}
|
||||
var out map[string]any
|
||||
if err := json.Unmarshal(raw, &out); err != nil {
|
||||
return map[string]any{}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func stringValue(v any) string {
|
||||
switch x := v.(type) {
|
||||
case string:
|
||||
return x
|
||||
case nil:
|
||||
return ""
|
||||
default:
|
||||
b, _ := json.Marshal(x)
|
||||
return strings.TrimSpace(string(b))
|
||||
}
|
||||
}
|
||||
|
||||
func intValue(v any) int {
|
||||
switch x := v.(type) {
|
||||
case int:
|
||||
return x
|
||||
case int64:
|
||||
return int(x)
|
||||
case float64:
|
||||
return int(x)
|
||||
case json.Number:
|
||||
n, _ := x.Int64()
|
||||
return int(n)
|
||||
case string:
|
||||
n, _ := strconv.Atoi(strings.TrimSpace(x))
|
||||
return n
|
||||
default:
|
||||
return 0
|
||||
}
|
||||
}
|
||||
|
||||
func emptyStringNil(v string) any {
|
||||
v = strings.TrimSpace(v)
|
||||
if v == "" {
|
||||
return nil
|
||||
}
|
||||
return v
|
||||
}
|
||||
|
||||
func clip(b []byte, n int) string {
|
||||
s := strings.TrimSpace(string(b))
|
||||
if len(s) > n {
|
||||
return s[:n]
|
||||
}
|
||||
return s
|
||||
}
|
||||
@@ -0,0 +1,391 @@
|
||||
package runway
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"image"
|
||||
_ "image/gif"
|
||||
_ "image/jpeg"
|
||||
_ "image/png"
|
||||
"io"
|
||||
mrand "math/rand/v2"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
http "github.com/bogdanfinn/fhttp"
|
||||
tlsclient "github.com/bogdanfinn/tls-client"
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
// ratioDimensions maps an aspect ratio to the Gen-4 Turbo native output size.
|
||||
// These are the only dimensions gen4_turbo accepts; "2K" is a UI label over this
|
||||
// native tier (see runway-video-gen-spec). Unknown ratios fall back to 16:9.
|
||||
func ratioDimensions(aspectRatio string) (int, int) {
|
||||
switch strings.TrimSpace(strings.ReplaceAll(aspectRatio, "x", ":")) {
|
||||
case "16:9":
|
||||
return 1280, 720
|
||||
case "9:16":
|
||||
return 720, 1280
|
||||
case "1:1":
|
||||
return 960, 960
|
||||
case "4:3":
|
||||
return 1104, 832
|
||||
case "3:4":
|
||||
return 832, 1104
|
||||
case "21:9":
|
||||
return 1584, 672
|
||||
default:
|
||||
return 1280, 720
|
||||
}
|
||||
}
|
||||
|
||||
// GenerateVideo runs the full i2v pipeline (gen_video.py): upload the first-frame
|
||||
// image (preview + dataset), create a dataset, create a gen4_turbo task and poll
|
||||
// it to completion, then download the rendered MP4. teamID is the workspace id
|
||||
// (meta["team_id"]); if empty it's derived from the token. seconds must be 5 or
|
||||
// 10; aspectRatio picks the native output size.
|
||||
// GenerateVideo renders the clip and (when downloadResult) downloads the MP4.
|
||||
// With downloadResult=false it returns nil bytes and the upstream artifact URL in
|
||||
// meta["video_url"] — used by the async /v1/videos job, which proxies that URL on
|
||||
// /content instead of persisting the file.
|
||||
func (c *Client) GenerateVideo(ctx context.Context, token, teamID, prompt, aspectRatio string, seconds int, frame []byte, downloadResult bool) ([]byte, map[string]any, error) {
|
||||
token = strings.TrimSpace(strings.TrimPrefix(token, "Bearer "))
|
||||
if token == "" {
|
||||
return nil, nil, ErrAuth
|
||||
}
|
||||
if teamID == "" {
|
||||
teamID = TeamIDFromToken(token)
|
||||
}
|
||||
if teamID == "" {
|
||||
return nil, nil, errors.New("runway: no team id")
|
||||
}
|
||||
if len(frame) == 0 {
|
||||
return nil, nil, errors.New("runway: first-frame image required")
|
||||
}
|
||||
cfg, _, err := image.DecodeConfig(bytes.NewReader(frame))
|
||||
if err != nil {
|
||||
return nil, nil, errors.New("runway: failed to decode first-frame image")
|
||||
}
|
||||
|
||||
client, err := c.newTLSClient()
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
|
||||
filename := "frame_" + time.Now().UTC().Format("20060102_150405") + ".png"
|
||||
previewUploadID, _, err := c.uploadFile(ctx, client, token, teamID, filename, "DATASET_PREVIEW", frame)
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
datasetUploadID, _, err := c.uploadFile(ctx, client, token, teamID, filename, "DATASET", frame)
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
|
||||
assetID, imageURL, err := c.createDataset(ctx, client, token, teamID, filename, datasetUploadID, previewUploadID, cfg.Width, cfg.Height)
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
assetGroupID, _ := c.assetGroupID(ctx, client, token, teamID) // best-effort
|
||||
|
||||
taskID, err := c.createTask(ctx, client, token, teamID, prompt, imageURL, assetID, assetGroupID, aspectRatio, seconds)
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
|
||||
artifactURL, err := c.pollTask(ctx, client, token, teamID, taskID)
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
meta := map[string]any{
|
||||
"provider": "runway",
|
||||
"task_id": taskID,
|
||||
"team_id": teamID,
|
||||
"video_url": artifactURL,
|
||||
}
|
||||
if !downloadResult {
|
||||
return nil, meta, nil
|
||||
}
|
||||
data, err := c.download(ctx, client, artifactURL)
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
return data, meta, nil
|
||||
}
|
||||
|
||||
// uploadFile mirrors gen_video.upload_file: register the upload, PUT the bytes to
|
||||
// the returned S3 URL, then complete. Returns the upload id and final url.
|
||||
func (c *Client) uploadFile(ctx context.Context, client tlsclient.HttpClient, token, teamID, filename, uploadType string, data []byte) (string, string, error) {
|
||||
info, err := c.apiJSON(ctx, client, token, teamID, http.MethodPost, "/v1/uploads", map[string]any{
|
||||
"filename": filename,
|
||||
"numberOfParts": 1,
|
||||
"type": uploadType,
|
||||
})
|
||||
if err != nil {
|
||||
return "", "", err
|
||||
}
|
||||
uploadID := strings.TrimSpace(stringValue(info["id"]))
|
||||
urls, _ := info["uploadUrls"].([]any)
|
||||
if uploadID == "" || len(urls) == 0 {
|
||||
return "", "", fmt.Errorf("%w: upload register missing fields", ErrTemporaryUpstream)
|
||||
}
|
||||
putURL := strings.TrimSpace(stringValue(urls[0]))
|
||||
contentType := "application/octet-stream"
|
||||
if hdrs, ok := info["uploadHeaders"].(map[string]any); ok {
|
||||
if ct := strings.TrimSpace(stringValue(hdrs["Content-Type"])); ct != "" {
|
||||
contentType = ct
|
||||
}
|
||||
}
|
||||
|
||||
etag, err := c.putBytes(ctx, client, putURL, contentType, data)
|
||||
if err != nil {
|
||||
return "", "", err
|
||||
}
|
||||
|
||||
res, err := c.apiJSON(ctx, client, token, teamID, http.MethodPost, "/v1/uploads/"+uploadID+"/complete", map[string]any{
|
||||
"parts": []map[string]any{{"PartNumber": 1, "ETag": etag}},
|
||||
})
|
||||
if err != nil {
|
||||
return "", "", err
|
||||
}
|
||||
return uploadID, strings.TrimSpace(stringValue(res["url"])), nil
|
||||
}
|
||||
|
||||
func (c *Client) createDataset(ctx context.Context, client tlsclient.HttpClient, token, teamID, filename, datasetUploadID, previewUploadID string, w, h int) (string, string, error) {
|
||||
teamIDNum := jsonNumberOrString(teamID)
|
||||
res, err := c.apiJSON(ctx, client, token, teamID, http.MethodPost, "/v1/datasets", map[string]any{
|
||||
"fileCount": 1,
|
||||
"name": filename,
|
||||
"uploadId": datasetUploadID,
|
||||
"previewUploadIds": []string{previewUploadID},
|
||||
"metadata": map[string]any{"size": map[string]any{"width": w, "height": h}},
|
||||
"type": map[string]any{"name": "image", "type": "image", "isDirectory": false},
|
||||
"asTeamId": teamIDNum,
|
||||
"privateInTeam": true,
|
||||
})
|
||||
if err != nil {
|
||||
return "", "", err
|
||||
}
|
||||
ds, _ := res["dataset"].(map[string]any)
|
||||
id := strings.TrimSpace(stringValue(ds["id"]))
|
||||
url := strings.TrimSpace(stringValue(ds["url"]))
|
||||
if id == "" || url == "" {
|
||||
return "", "", fmt.Errorf("%w: dataset missing fields", ErrTemporaryUpstream)
|
||||
}
|
||||
return id, url, nil
|
||||
}
|
||||
|
||||
func (c *Client) assetGroupID(ctx context.Context, client tlsclient.HttpClient, token, teamID string) (string, error) {
|
||||
res, err := c.apiJSON(ctx, client, token, teamID, http.MethodGet,
|
||||
"/v1/asset_groups/by_name?name=Generations&asTeamId="+teamID+"&privateInTeam=true", nil)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
ag, _ := res["assetGroup"].(map[string]any)
|
||||
return strings.TrimSpace(stringValue(ag["id"])), nil
|
||||
}
|
||||
|
||||
func (c *Client) createTask(ctx context.Context, client tlsclient.HttpClient, token, teamID, prompt, imageURL, assetID, assetGroupID, aspectRatio string, seconds int) (string, error) {
|
||||
w, h := ratioDimensions(aspectRatio)
|
||||
opts := map[string]any{
|
||||
"route": "i2v",
|
||||
"name": "Gen-4 Turbo - " + prompt,
|
||||
"text_prompt": prompt,
|
||||
"seconds": seconds,
|
||||
"width": w,
|
||||
"height": h,
|
||||
"init_image": imageURL,
|
||||
"imageAssetId": assetID,
|
||||
"exploreMode": false,
|
||||
"creationSource": "tool-mode",
|
||||
"seed": mrand.IntN(999999999) + 1,
|
||||
"watermark": true,
|
||||
}
|
||||
if assetGroupID != "" {
|
||||
opts["assetGroupId"] = assetGroupID
|
||||
}
|
||||
res, err := c.apiJSON(ctx, client, token, teamID, http.MethodPost, "/v1/tasks", map[string]any{
|
||||
"taskType": "gen4_turbo",
|
||||
"options": opts,
|
||||
"asTeamId": jsonNumberOrString(teamID),
|
||||
"sessionId": uuid.NewString(),
|
||||
})
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
task, _ := res["task"].(map[string]any)
|
||||
id := strings.TrimSpace(stringValue(task["id"]))
|
||||
if id == "" {
|
||||
return "", fmt.Errorf("%w: task missing id", ErrTemporaryUpstream)
|
||||
}
|
||||
return id, nil
|
||||
}
|
||||
|
||||
func (c *Client) pollTask(ctx context.Context, client tlsclient.HttpClient, token, teamID, taskID string) (string, error) {
|
||||
for {
|
||||
if err := ctx.Err(); err != nil {
|
||||
return "", err
|
||||
}
|
||||
res, err := c.apiJSON(ctx, client, token, teamID, http.MethodGet, "/v1/tasks/"+taskID+"?asTeamId="+teamID, nil)
|
||||
if err != nil {
|
||||
// A transient blip shouldn't kill a render that may still succeed.
|
||||
if errors.Is(err, ErrTemporaryUpstream) {
|
||||
if sleepCtx(ctx, 5*time.Second) != nil {
|
||||
return "", ctx.Err()
|
||||
}
|
||||
continue
|
||||
}
|
||||
return "", err
|
||||
}
|
||||
task, _ := res["task"].(map[string]any)
|
||||
status := strings.ToUpper(strings.TrimSpace(stringValue(task["status"])))
|
||||
switch status {
|
||||
case "SUCCEEDED":
|
||||
arts, _ := task["artifacts"].([]any)
|
||||
for _, raw := range arts {
|
||||
art, _ := raw.(map[string]any)
|
||||
if url := strings.TrimSpace(stringValue(art["url"])); url != "" {
|
||||
return url, nil
|
||||
}
|
||||
}
|
||||
return "", errors.New("runway: task succeeded with no artifact url")
|
||||
case "FAILED", "CANCELED":
|
||||
reason := strings.TrimSpace(stringValue(task["error"]))
|
||||
if isCreditError(reason) {
|
||||
return "", fmt.Errorf("%w: %s", ErrQuotaExhausted, reason)
|
||||
}
|
||||
return "", fmt.Errorf("runway: task %s: %s", status, reason)
|
||||
}
|
||||
if sleepCtx(ctx, 5*time.Second) != nil {
|
||||
return "", ctx.Err()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// apiJSON performs an authed JSON request against the Runway API and returns the
|
||||
// parsed body, mapping status codes to the shared provider error sentinels.
|
||||
func (c *Client) apiJSON(ctx context.Context, client tlsclient.HttpClient, token, teamID, method, path string, body any) (map[string]any, error) {
|
||||
var reader io.Reader
|
||||
if body != nil {
|
||||
raw, _ := json.Marshal(body)
|
||||
reader = bytes.NewReader(raw)
|
||||
}
|
||||
req, err := http.NewRequest(method, apiBase+path, reader)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
req = req.WithContext(ctx)
|
||||
req.Header = http.Header{
|
||||
"accept": {"application/json"},
|
||||
"content-type": {"application/json"},
|
||||
"origin": {origin},
|
||||
"referer": {origin + "/"},
|
||||
"authorization": {"Bearer " + token},
|
||||
"x-runway-workspace": {teamID},
|
||||
http.HeaderOrderKey: {
|
||||
"accept", "content-type", "origin", "referer", "authorization", "x-runway-workspace",
|
||||
},
|
||||
}
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("%w: %v", ErrTemporaryUpstream, err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
raw, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
switch {
|
||||
case resp.StatusCode == 401 || resp.StatusCode == 403:
|
||||
return nil, fmt.Errorf("%w: %s %d %s", ErrAuth, path, resp.StatusCode, clip(raw, 200))
|
||||
case resp.StatusCode == 429:
|
||||
return nil, fmt.Errorf("%w: %s 429 %s", ErrQuotaExhausted, path, clip(raw, 200))
|
||||
case resp.StatusCode >= 500:
|
||||
return nil, fmt.Errorf("%w: %s %d %s", ErrTemporaryUpstream, path, resp.StatusCode, clip(raw, 200))
|
||||
case resp.StatusCode < 200 || resp.StatusCode >= 300:
|
||||
if isCreditError(string(raw)) {
|
||||
return nil, fmt.Errorf("%w: %s", ErrQuotaExhausted, clip(raw, 200))
|
||||
}
|
||||
return nil, fmt.Errorf("runway: %s %d %s", path, resp.StatusCode, clip(raw, 200))
|
||||
}
|
||||
var out map[string]any
|
||||
if len(raw) == 0 {
|
||||
return map[string]any{}, nil
|
||||
}
|
||||
if err := json.Unmarshal(raw, &out); err != nil {
|
||||
return nil, fmt.Errorf("%w: %s non-json: %s", ErrTemporaryUpstream, path, clip(raw, 120))
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// putBytes uploads raw bytes to a presigned S3 URL (no auth) and returns the
|
||||
// ETag, mirroring the plain requests.Session().put in gen_video.py.
|
||||
func (c *Client) putBytes(ctx context.Context, client tlsclient.HttpClient, url, contentType string, data []byte) (string, error) {
|
||||
req, err := http.NewRequest(http.MethodPut, url, bytes.NewReader(data))
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
req = req.WithContext(ctx)
|
||||
req.Header = http.Header{"content-type": {contentType}}
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("%w: %v", ErrTemporaryUpstream, err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
_, _ = io.Copy(io.Discard, resp.Body)
|
||||
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
||||
return "", fmt.Errorf("%w: s3 put %d", ErrTemporaryUpstream, resp.StatusCode)
|
||||
}
|
||||
return strings.Trim(resp.Header.Get("ETag"), `"`), nil
|
||||
}
|
||||
|
||||
func (c *Client) download(ctx context.Context, client tlsclient.HttpClient, url string) ([]byte, error) {
|
||||
req, err := http.NewRequest(http.MethodGet, url, nil)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
req = req.WithContext(ctx)
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("%w: %v", ErrTemporaryUpstream, err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
||||
return nil, fmt.Errorf("%w: download %d", ErrTemporaryUpstream, resp.StatusCode)
|
||||
}
|
||||
data, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if len(data) == 0 {
|
||||
return nil, errors.New("runway: empty artifact download")
|
||||
}
|
||||
return data, nil
|
||||
}
|
||||
|
||||
// jsonNumberOrString returns the team id as a JSON number when it's purely
|
||||
// numeric (Runway's asTeamId is an integer in the reference payloads), else the
|
||||
// raw string.
|
||||
func jsonNumberOrString(teamID string) any {
|
||||
return json.Number(strings.TrimSpace(teamID))
|
||||
}
|
||||
|
||||
func isCreditError(s string) bool {
|
||||
s = strings.ToLower(s)
|
||||
return strings.Contains(s, "credit") || strings.Contains(s, "insufficient") || strings.Contains(s, "quota")
|
||||
}
|
||||
|
||||
// sleepCtx sleeps for d or until ctx is done; returns ctx.Err() if cancelled.
|
||||
func sleepCtx(ctx context.Context, d time.Duration) error {
|
||||
t := time.NewTimer(d)
|
||||
defer t.Stop()
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
case <-t.C:
|
||||
return nil
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,53 @@
|
||||
package repo
|
||||
|
||||
import (
|
||||
"context"
|
||||
"time"
|
||||
|
||||
"backend/internal/model"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
type APIKeyRepository struct {
|
||||
db *gorm.DB
|
||||
}
|
||||
|
||||
func NewAPIKeyRepository(db *gorm.DB) *APIKeyRepository {
|
||||
return &APIKeyRepository{db: db}
|
||||
}
|
||||
|
||||
func (r *APIKeyRepository) ListByUserID(ctx context.Context, userID string) ([]model.APIKey, error) {
|
||||
var keys []model.APIKey
|
||||
if err := r.db.WithContext(ctx).Where("user_id = ?", userID).Order("created_at asc").Find(&keys).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return keys, nil
|
||||
}
|
||||
|
||||
func (r *APIKeyRepository) ReplaceForUser(ctx context.Context, userID string, key *model.APIKey) error {
|
||||
return r.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
|
||||
if err := tx.Where("user_id = ?", userID).Delete(&model.APIKey{}).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
return tx.Create(key).Error
|
||||
})
|
||||
}
|
||||
|
||||
func (r *APIKeyRepository) DeleteByUserID(ctx context.Context, userID string) error {
|
||||
return r.db.WithContext(ctx).Where("user_id = ?", userID).Delete(&model.APIKey{}).Error
|
||||
}
|
||||
|
||||
func (r *APIKeyRepository) DeleteByID(ctx context.Context, userID, keyID string) error {
|
||||
return r.db.WithContext(ctx).
|
||||
Where("user_id = ? AND id = ?", userID, keyID).
|
||||
Delete(&model.APIKey{}).Error
|
||||
}
|
||||
|
||||
func (r *APIKeyRepository) Create(ctx context.Context, key *model.APIKey) error {
|
||||
return r.db.WithContext(ctx).Create(key).Error
|
||||
}
|
||||
|
||||
func (r *APIKeyRepository) TouchUsage(ctx context.Context, keyHash string) error {
|
||||
now := time.Now()
|
||||
return r.db.WithContext(ctx).Model(&model.APIKey{}).Where("key_hash = ?", keyHash).Update("last_used_at", now).Error
|
||||
}
|
||||
@@ -0,0 +1,134 @@
|
||||
package repo
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"time"
|
||||
|
||||
"backend/internal/model"
|
||||
"gorm.io/gorm"
|
||||
"gorm.io/gorm/clause"
|
||||
)
|
||||
|
||||
// ErrCDKBatchLimit is returned when a user tries to redeem a second code from
|
||||
// the same marketing batch (one per user per batch).
|
||||
var ErrCDKBatchLimit = errors.New("cdk marketing batch already redeemed by this user")
|
||||
|
||||
type CDKRepository struct {
|
||||
db *gorm.DB
|
||||
}
|
||||
|
||||
func NewCDKRepository(db *gorm.DB) *CDKRepository {
|
||||
return &CDKRepository{db: db}
|
||||
}
|
||||
|
||||
func (r *CDKRepository) List(ctx context.Context) ([]model.CDKCode, error) {
|
||||
var items []model.CDKCode
|
||||
if err := r.db.WithContext(ctx).Order("created_at desc").Find(&items).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return items, nil
|
||||
}
|
||||
|
||||
func (r *CDKRepository) Stats(ctx context.Context) (map[string]any, error) {
|
||||
var total, active, redeemed int64
|
||||
if err := r.db.WithContext(ctx).Model(&model.CDKCode{}).Count(&total).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := r.db.WithContext(ctx).Model(&model.CDKCode{}).Where("status = ?", "active").Count(&active).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := r.db.WithContext(ctx).Model(&model.CDKCode{}).Where("status = ?", "redeemed").Count(&redeemed).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
type sumRow struct {
|
||||
Total *float64 `gorm:"column:total"`
|
||||
}
|
||||
var activeAmount, redeemedAmount sumRow
|
||||
if err := r.db.WithContext(ctx).
|
||||
Model(&model.CDKCode{}).
|
||||
Select("SUM(amount) AS total").
|
||||
Where("status = ?", "active").
|
||||
Scan(&activeAmount).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := r.db.WithContext(ctx).
|
||||
Model(&model.CDKCode{}).
|
||||
Select("SUM(amount) AS total").
|
||||
Where("status = ?", "redeemed").
|
||||
Scan(&redeemedAmount).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
activeAmt := 0.0
|
||||
if activeAmount.Total != nil {
|
||||
activeAmt = *activeAmount.Total
|
||||
}
|
||||
redeemedAmt := 0.0
|
||||
if redeemedAmount.Total != nil {
|
||||
redeemedAmt = *redeemedAmount.Total
|
||||
}
|
||||
|
||||
return map[string]any{
|
||||
"total": total,
|
||||
"active": active,
|
||||
"redeemed": redeemed,
|
||||
"active_amount": activeAmt,
|
||||
"redeemed_amount": redeemedAmt,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (r *CDKRepository) CreateBatch(ctx context.Context, items []model.CDKCode) error {
|
||||
return r.db.WithContext(ctx).Create(&items).Error
|
||||
}
|
||||
|
||||
func (r *CDKRepository) Delete(ctx context.Context, code string) (int64, error) {
|
||||
res := r.db.WithContext(ctx).Delete(&model.CDKCode{}, "code = ?", code)
|
||||
return res.RowsAffected, res.Error
|
||||
}
|
||||
|
||||
func (r *CDKRepository) DeleteByCodes(ctx context.Context, codes []string) (int64, error) {
|
||||
if len(codes) == 0 {
|
||||
return 0, nil
|
||||
}
|
||||
res := r.db.WithContext(ctx).Delete(&model.CDKCode{}, "code IN ?", codes)
|
||||
return res.RowsAffected, res.Error
|
||||
}
|
||||
|
||||
func (r *CDKRepository) Redeem(ctx context.Context, code, userID string) (*model.CDKCode, error) {
|
||||
var out *model.CDKCode
|
||||
err := r.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
|
||||
var item model.CDKCode
|
||||
if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).First(&item, "code = ?", code).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if item.Status == "redeemed" {
|
||||
return gorm.ErrDuplicatedKey
|
||||
}
|
||||
// Marketing codes: a user may redeem only ONE code per batch. The partial
|
||||
// unique index (batch_id, redeemed_by) is the hard backstop against
|
||||
// concurrent double-redeems; this check gives a friendly error first.
|
||||
if item.Type == "marketing" && item.BatchID != "" {
|
||||
var cnt int64
|
||||
if err := tx.Model(&model.CDKCode{}).
|
||||
Where("batch_id = ? AND type = 'marketing' AND redeemed_by = ?", item.BatchID, userID).
|
||||
Count(&cnt).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if cnt > 0 {
|
||||
return ErrCDKBatchLimit
|
||||
}
|
||||
}
|
||||
now := time.Now()
|
||||
item.Status = "redeemed"
|
||||
item.RedeemedBy = &userID
|
||||
item.RedeemedAt = &now
|
||||
if err := tx.Save(&item).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
out = &item
|
||||
return nil
|
||||
})
|
||||
return out, err
|
||||
}
|
||||
@@ -0,0 +1,650 @@
|
||||
package repo
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"backend/internal/model"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
type EventRepository struct {
|
||||
db *gorm.DB
|
||||
}
|
||||
|
||||
type EventListFilter struct {
|
||||
Limit int
|
||||
Offset int
|
||||
Kind string
|
||||
Status string
|
||||
Since *time.Time
|
||||
UserID string
|
||||
ExcludeSource string // when set, omit rows with this source (e.g. hide API-key "v1" usage from the customer logs page)
|
||||
Source string // when set, keep ONLY rows with this source (admin 来源 filter): "v1" (API key) / "user" (前台) / "admin" (测试模型)
|
||||
HasFile bool // when true, keep ONLY rows with a non-empty file (the 创作记录 gallery — paginates over real media)
|
||||
}
|
||||
|
||||
type EventStats struct {
|
||||
Total int64 `json:"total"`
|
||||
Success int64 `json:"success"`
|
||||
Failed int64 `json:"failed"`
|
||||
Pending int64 `json:"pending"`
|
||||
AvgElapsedMS *int `json:"avg_elapsed_ms"`
|
||||
AvgElapsedMS24 *int `json:"avg_elapsed_ms_24h"`
|
||||
}
|
||||
|
||||
func NewEventRepository(db *gorm.DB) *EventRepository {
|
||||
return &EventRepository{db: db}
|
||||
}
|
||||
|
||||
func (r *EventRepository) List(ctx context.Context, filter EventListFilter) ([]model.EventLog, int64, error) {
|
||||
q := r.db.WithContext(ctx).Model(&model.EventLog{})
|
||||
if filter.Kind != "" {
|
||||
q = q.Where("kind = ?", filter.Kind)
|
||||
}
|
||||
if filter.Status != "" {
|
||||
q = q.Where("status = ?", filter.Status)
|
||||
}
|
||||
if filter.Since != nil {
|
||||
q = q.Where("ts > ?", *filter.Since)
|
||||
}
|
||||
if filter.UserID != "" {
|
||||
q = q.Where("user_id = ?", filter.UserID)
|
||||
}
|
||||
if filter.ExcludeSource != "" {
|
||||
q = q.Where("(source IS NULL OR source <> ?)", filter.ExcludeSource)
|
||||
}
|
||||
if filter.Source != "" {
|
||||
q = q.Where("source = ?", filter.Source)
|
||||
}
|
||||
if filter.HasFile {
|
||||
q = q.Where("file <> ''")
|
||||
}
|
||||
|
||||
var total int64
|
||||
if err := q.Count(&total).Error; err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
|
||||
var items []model.EventLog
|
||||
if err := q.Order("ts desc").
|
||||
Limit(filter.Limit).
|
||||
Offset(filter.Offset).
|
||||
Find(&items).Error; err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
return items, total, nil
|
||||
}
|
||||
|
||||
func (r *EventRepository) Stats(ctx context.Context) (*EventStats, error) {
|
||||
stats := &EventStats{}
|
||||
if err := r.db.WithContext(ctx).Model(&model.EventLog{}).Count(&stats.Total).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := r.db.WithContext(ctx).Model(&model.EventLog{}).Where("status = ?", "success").Count(&stats.Success).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := r.db.WithContext(ctx).Model(&model.EventLog{}).Where("status = ?", "failed").Count(&stats.Failed).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := r.db.WithContext(ctx).Model(&model.EventLog{}).Where("status = ?", "pending").Count(&stats.Pending).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
type avgRow struct {
|
||||
Avg *float64 `gorm:"column:avg"`
|
||||
}
|
||||
var all avgRow
|
||||
if err := r.db.WithContext(ctx).
|
||||
Model(&model.EventLog{}).
|
||||
Select("AVG(elapsed_ms) AS avg").
|
||||
Where("status = ? AND elapsed_ms > 0", "success").
|
||||
Scan(&all).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if all.Avg != nil {
|
||||
v := int(*all.Avg + 0.5)
|
||||
stats.AvgElapsedMS = &v
|
||||
}
|
||||
|
||||
var recent avgRow
|
||||
cutoff := time.Now().Add(-24 * time.Hour)
|
||||
if err := r.db.WithContext(ctx).
|
||||
Model(&model.EventLog{}).
|
||||
Select("AVG(elapsed_ms) AS avg").
|
||||
Where("status = ? AND elapsed_ms > 0 AND ts >= ?", "success", cutoff).
|
||||
Scan(&recent).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if recent.Avg != nil {
|
||||
v := int(*recent.Avg + 0.5)
|
||||
stats.AvgElapsedMS24 = &v
|
||||
}
|
||||
|
||||
return stats, nil
|
||||
}
|
||||
|
||||
// StatsByUser returns total / success / failed / pending counts scoped to a
|
||||
// single user — for the customer-facing 生成日志 (/mylogs) KPI strip, so it
|
||||
// reflects the caller's own history, not the whole site.
|
||||
func (r *EventRepository) StatsByUser(ctx context.Context, userID string) (*EventStats, error) {
|
||||
stats := &EventStats{}
|
||||
q := func() *gorm.DB {
|
||||
return r.db.WithContext(ctx).Model(&model.EventLog{}).Where("user_id = ?", userID)
|
||||
}
|
||||
if err := q().Count(&stats.Total).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := q().Where("status = ?", "success").Count(&stats.Success).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := q().Where("status = ?", "failed").Count(&stats.Failed).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := q().Where("status = ?", "pending").Count(&stats.Pending).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return stats, nil
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Dashboard aggregates — server-side GROUP BY / FILTER so the admin overview
|
||||
// no longer derives 7-day / DAU / trend / top-N numbers client-side from the
|
||||
// last 200 logs (which silently undercounts once volume passes that window).
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
// DashboardWindow is a single time-window aggregate (e.g. last 24h / 7d).
|
||||
type DashboardWindow struct {
|
||||
Total int64 `json:"total"`
|
||||
Success int64 `json:"success"`
|
||||
Failed int64 `json:"failed"`
|
||||
Pending int64 `json:"pending"`
|
||||
Image int64 `json:"image"`
|
||||
Video int64 `json:"video"`
|
||||
API int64 `json:"api"` // source = 'v1' (OpenAI-compatible key)
|
||||
Web int64 `json:"web"` // everything else (web / playground)
|
||||
Spent float64 `json:"spent"`
|
||||
}
|
||||
|
||||
type ModelUsage struct {
|
||||
Model string `json:"model"`
|
||||
Count int64 `json:"count"`
|
||||
AvgMS *int `json:"avg_ms"`
|
||||
}
|
||||
|
||||
type FailureReason struct {
|
||||
Reason string `json:"reason"`
|
||||
Count int64 `json:"count"`
|
||||
}
|
||||
|
||||
type UserSpend struct {
|
||||
UserID string `json:"user_id"`
|
||||
Name string `json:"name"` // resolved by the service from user_id
|
||||
Count int64 `json:"count"`
|
||||
Spent float64 `json:"spent"`
|
||||
}
|
||||
|
||||
type HourBucket struct {
|
||||
Image int64 `json:"image"`
|
||||
Video int64 `json:"video"`
|
||||
}
|
||||
|
||||
// WindowStats rolls up counts + spend over a single window in one query.
|
||||
func (r *EventRepository) WindowStats(ctx context.Context, since time.Time) (*DashboardWindow, error) {
|
||||
type row struct {
|
||||
Total int64 `gorm:"column:total"`
|
||||
Success int64 `gorm:"column:success"`
|
||||
Failed int64 `gorm:"column:failed"`
|
||||
Pending int64 `gorm:"column:pending"`
|
||||
Image int64 `gorm:"column:image"`
|
||||
Video int64 `gorm:"column:video"`
|
||||
API int64 `gorm:"column:api"`
|
||||
Spent float64 `gorm:"column:spent"`
|
||||
}
|
||||
var out row
|
||||
if err := r.db.WithContext(ctx).
|
||||
Model(&model.EventLog{}).
|
||||
Select(`
|
||||
COUNT(*) AS total,
|
||||
COUNT(*) FILTER (WHERE status = 'success') AS success,
|
||||
COUNT(*) FILTER (WHERE status = 'failed') AS failed,
|
||||
COUNT(*) FILTER (WHERE status = 'pending') AS pending,
|
||||
COUNT(*) FILTER (WHERE kind = 'image') AS image,
|
||||
COUNT(*) FILTER (WHERE kind = 'video') AS video,
|
||||
COUNT(*) FILTER (WHERE source = 'v1') AS api,
|
||||
COALESCE(SUM(cost) FILTER (WHERE status = 'success'), 0) AS spent`).
|
||||
Where("ts >= ?", since).
|
||||
Scan(&out).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &DashboardWindow{
|
||||
Total: out.Total, Success: out.Success, Failed: out.Failed, Pending: out.Pending,
|
||||
Image: out.Image, Video: out.Video, API: out.API, Web: out.Total - out.API,
|
||||
Spent: out.Spent,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// CountBetween counts events in (start, end] — used for the prev-24h delta.
|
||||
func (r *EventRepository) CountBetween(ctx context.Context, start, end time.Time) (int64, error) {
|
||||
var n int64
|
||||
err := r.db.WithContext(ctx).Model(&model.EventLog{}).
|
||||
Where("ts > ? AND ts <= ?", start, end).Count(&n).Error
|
||||
return n, err
|
||||
}
|
||||
|
||||
// DistinctUsersSince counts distinct (non-empty) user_ids active since `since`.
|
||||
func (r *EventRepository) DistinctUsersSince(ctx context.Context, since time.Time) (int64, error) {
|
||||
var n int64
|
||||
err := r.db.WithContext(ctx).Model(&model.EventLog{}).
|
||||
Where("ts >= ? AND user_id <> ''", since).
|
||||
Distinct("user_id").Count(&n).Error
|
||||
return n, err
|
||||
}
|
||||
|
||||
// HourlyBuckets returns 24 oldest→newest buckets (image/video split) for the
|
||||
// last 24h trend chart.
|
||||
func (r *EventRepository) HourlyBuckets(ctx context.Context) ([24]HourBucket, error) {
|
||||
var out [24]HourBucket
|
||||
type hourRow struct {
|
||||
HoursAgo int `gorm:"column:hours_ago"`
|
||||
Image int64 `gorm:"column:image"`
|
||||
Video int64 `gorm:"column:video"`
|
||||
}
|
||||
var rows []hourRow
|
||||
if err := r.db.WithContext(ctx).
|
||||
Model(&model.EventLog{}).
|
||||
Select(`
|
||||
FLOOR(EXTRACT(EPOCH FROM (NOW() - ts)) / 3600)::int AS hours_ago,
|
||||
COUNT(*) FILTER (WHERE kind = 'video') AS video,
|
||||
COUNT(*) FILTER (WHERE kind <> 'video') AS image`).
|
||||
Where("ts >= NOW() - INTERVAL '24 hours'").
|
||||
Group("hours_ago").
|
||||
Scan(&rows).Error; err != nil {
|
||||
return out, err
|
||||
}
|
||||
for _, hr := range rows {
|
||||
if hr.HoursAgo < 0 || hr.HoursAgo >= 24 {
|
||||
continue
|
||||
}
|
||||
out[23-hr.HoursAgo] = HourBucket{Image: hr.Image, Video: hr.Video}
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// ModelUsageSince returns the top models by volume since `since`, with the
|
||||
// success-only average latency.
|
||||
func (r *EventRepository) ModelUsageSince(ctx context.Context, since time.Time, limit int) ([]ModelUsage, error) {
|
||||
if limit <= 0 {
|
||||
limit = 6
|
||||
}
|
||||
type row struct {
|
||||
Model string `gorm:"column:model"`
|
||||
Count int64 `gorm:"column:count"`
|
||||
Avg *float64 `gorm:"column:avg_ms"`
|
||||
}
|
||||
var rows []row
|
||||
if err := r.db.WithContext(ctx).
|
||||
Model(&model.EventLog{}).
|
||||
Select(`
|
||||
model,
|
||||
COUNT(*) AS count,
|
||||
AVG(elapsed_ms) FILTER (WHERE status = 'success' AND elapsed_ms > 0) AS avg_ms`).
|
||||
Where("ts >= ? AND model <> ''", since).
|
||||
Group("model").
|
||||
Order("count DESC").
|
||||
Limit(limit).
|
||||
Scan(&rows).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out := make([]ModelUsage, 0, len(rows))
|
||||
for _, item := range rows {
|
||||
var avg *int
|
||||
if item.Avg != nil {
|
||||
v := int(*item.Avg + 0.5)
|
||||
avg = &v
|
||||
}
|
||||
out = append(out, ModelUsage{Model: item.Model, Count: item.Count, AvgMS: avg})
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// TopFailures groups failed events by (truncated) error reason since `since`.
|
||||
func (r *EventRepository) TopFailures(ctx context.Context, since time.Time, limit int) ([]FailureReason, error) {
|
||||
if limit <= 0 {
|
||||
limit = 5
|
||||
}
|
||||
var out []FailureReason
|
||||
if err := r.db.WithContext(ctx).
|
||||
Model(&model.EventLog{}).
|
||||
Select(`
|
||||
LEFT(COALESCE(NULLIF(error, ''), '未知错误'), 60) AS reason,
|
||||
COUNT(*) AS count`).
|
||||
Where("ts >= ? AND status = 'failed'", since).
|
||||
Group("reason").
|
||||
Order("count DESC").
|
||||
Limit(limit).
|
||||
Scan(&out).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// TopUserSpend ranks users by credits spent on SUCCESSFUL generations since
|
||||
// `since`. Names are resolved by the caller (UserID -> display name).
|
||||
func (r *EventRepository) TopUserSpend(ctx context.Context, since time.Time, limit int) ([]UserSpend, error) {
|
||||
if limit <= 0 {
|
||||
limit = 6
|
||||
}
|
||||
var out []UserSpend
|
||||
if err := r.db.WithContext(ctx).
|
||||
Model(&model.EventLog{}).
|
||||
Select(`
|
||||
user_id,
|
||||
COUNT(*) AS count,
|
||||
COALESCE(SUM(cost), 0) AS spent`).
|
||||
Where("ts >= ? AND status = 'success'", since).
|
||||
Group("user_id").
|
||||
Order("spent DESC").
|
||||
Limit(limit).
|
||||
Scan(&out).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func (r *EventRepository) PurgeOlderThan(ctx context.Context, maxAge time.Duration) (int64, error) {
|
||||
if maxAge <= 0 {
|
||||
return 0, nil
|
||||
}
|
||||
cutoff := time.Now().Add(-maxAge)
|
||||
result := r.db.WithContext(ctx).Where("ts < ?", cutoff).Delete(&model.EventLog{})
|
||||
if result.Error != nil {
|
||||
return 0, result.Error
|
||||
}
|
||||
return result.RowsAffected, nil
|
||||
}
|
||||
|
||||
// ClearFiles blanks the `file` column on any event rows that point at one of the
|
||||
// given relative paths. Called after media retention deletes the files on disk so
|
||||
// the log views don't dangle a 404 image — an emptied `file` reads as "no preview"
|
||||
// ("—" in admin logs; hidden in the customer records page).
|
||||
func (r *EventRepository) ClearFiles(ctx context.Context, relPaths []string) (int64, error) {
|
||||
if len(relPaths) == 0 {
|
||||
return 0, nil
|
||||
}
|
||||
result := r.db.WithContext(ctx).
|
||||
Model(&model.EventLog{}).
|
||||
Where("file IN ?", relPaths).
|
||||
Updates(map[string]any{"file": "", "updated_at": time.Now()})
|
||||
if result.Error != nil {
|
||||
return 0, result.Error
|
||||
}
|
||||
return result.RowsAffected, nil
|
||||
}
|
||||
|
||||
// ClearRefFiles blanks the ref_files paths on one event (called after a
|
||||
// successful generation once the reference images are deleted from storage, so no
|
||||
// dangling reference_urls remain). The `refs` COUNT is kept for the log record.
|
||||
func (r *EventRepository) ClearRefFiles(ctx context.Context, eventID string) error {
|
||||
return r.db.WithContext(ctx).
|
||||
Model(&model.EventLog{}).
|
||||
Where("id = ?", eventID).
|
||||
Update("ref_files", nil).Error
|
||||
}
|
||||
|
||||
// StaleEvent identifies a purged pending event so the caller can refund the
|
||||
// credits debited up-front AND attribute the failure to the account the
|
||||
// (now-abandoned) generation was using.
|
||||
type StaleEvent struct {
|
||||
ID string `gorm:"column:id"`
|
||||
UserID string `gorm:"column:user_id"`
|
||||
AccountID string `gorm:"column:account_id"`
|
||||
Cost float64 `gorm:"column:cost"`
|
||||
}
|
||||
|
||||
// PurgeStale marks long-pending entries as failed/abandoned and RETURNS them so
|
||||
// the caller can refund their up-front charge. A stuck pending row otherwise
|
||||
// blocks the per-user generation gate (PendingByUser) forever AND silently eats
|
||||
// the user's credits (the charge happens at submit; the normal failure-refund
|
||||
// path never runs for a process-restart orphan). Mirrors Python purge_stale.
|
||||
func (r *EventRepository) PurgeStale(ctx context.Context, maxAge time.Duration) ([]StaleEvent, error) {
|
||||
if maxAge <= 0 {
|
||||
maxAge = 600 * time.Second
|
||||
}
|
||||
cutoff := time.Now().Add(-maxAge)
|
||||
var stale []StaleEvent
|
||||
err := r.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
|
||||
// Snapshot who/what to refund BEFORE flipping status, so a concurrent
|
||||
// sweep can't double-count (the UPDATE in the same tx removes them from
|
||||
// the pending set).
|
||||
if err := tx.Model(&model.EventLog{}).
|
||||
Where("status = ? AND ts < ?", "pending", cutoff).
|
||||
Select("id", "user_id", "account_id", "cost").
|
||||
Scan(&stale).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if len(stale) == 0 {
|
||||
return nil
|
||||
}
|
||||
return tx.Model(&model.EventLog{}).
|
||||
Where("status = ? AND ts < ?", "pending", cutoff).
|
||||
Updates(map[string]any{
|
||||
"status": "failed",
|
||||
"error": gorm.Expr("COALESCE(NULLIF(error, ''), ?)", "abandoned (process restarted or request interrupted)"),
|
||||
"updated_at": time.Now(),
|
||||
}).Error
|
||||
})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return stale, nil
|
||||
}
|
||||
|
||||
func (r *EventRepository) Create(ctx context.Context, item *model.EventLog) error {
|
||||
return r.db.WithContext(ctx).Create(item).Error
|
||||
}
|
||||
|
||||
// GetByID fetches a single event (nil, nil when not found). Used by the async
|
||||
// /v1/videos job to look up status / the stored upstream URL.
|
||||
func (r *EventRepository) GetByID(ctx context.Context, id string) (*model.EventLog, error) {
|
||||
var e model.EventLog
|
||||
if err := r.db.WithContext(ctx).First(&e, "id = ?", id).Error; err != nil {
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
return nil, nil
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
return &e, nil
|
||||
}
|
||||
|
||||
// MarkVideoReady completes an async video job: status=success, file=upstream URL
|
||||
// (proxied on /content — never persisted), elapsed.
|
||||
func (r *EventRepository) MarkVideoReady(ctx context.Context, eventID, fileURL string, elapsedMS int) error {
|
||||
return r.db.WithContext(ctx).
|
||||
Model(&model.EventLog{}).
|
||||
Where("id = ?", eventID).
|
||||
Updates(map[string]any{
|
||||
"status": "success",
|
||||
"file": fileURL,
|
||||
"error": "",
|
||||
"elapsed_ms": elapsedMS,
|
||||
"updated_at": time.Now(),
|
||||
}).Error
|
||||
}
|
||||
|
||||
func (r *EventRepository) UpdateStatus(ctx context.Context, eventID, status, errMsg string, elapsedMS int) error {
|
||||
patch := map[string]any{
|
||||
"status": status,
|
||||
"elapsed_ms": elapsedMS,
|
||||
"updated_at": time.Now(),
|
||||
}
|
||||
if strings.TrimSpace(errMsg) != "" {
|
||||
patch["error"] = strings.TrimSpace(errMsg)
|
||||
} else if status == "success" {
|
||||
// A late-completing generation (one the maintenance sweep had already
|
||||
// stamped "abandoned") must shed that stale error, or the row reads as
|
||||
// "成功 + abandoned" at once.
|
||||
patch["error"] = ""
|
||||
}
|
||||
return r.db.WithContext(ctx).
|
||||
Model(&model.EventLog{}).
|
||||
Where("id = ?", eventID).
|
||||
Updates(patch).Error
|
||||
}
|
||||
|
||||
// MarkRefunded atomically claims the right to refund this event exactly once:
|
||||
// it flips refunded false→true and returns true ONLY for the caller that won the
|
||||
// race. Both the normal failure path and the abandoned-purge sweep call this
|
||||
// before crediting, so a generation can never be refunded twice.
|
||||
func (r *EventRepository) MarkRefunded(ctx context.Context, eventID string) (bool, error) {
|
||||
res := r.db.WithContext(ctx).
|
||||
Model(&model.EventLog{}).
|
||||
Where("id = ? AND refunded = ?", eventID, false).
|
||||
Updates(map[string]any{"refunded": true, "updated_at": time.Now()})
|
||||
if res.Error != nil {
|
||||
return false, res.Error
|
||||
}
|
||||
return res.RowsAffected == 1, nil
|
||||
}
|
||||
|
||||
// SetAccount stamps which provider account is fulfilling an in-flight event.
|
||||
// Called when generation commits to a token, so the accounts view can count
|
||||
// pending events per account and an abandoned-event purge can attribute back.
|
||||
func (r *EventRepository) SetAccount(ctx context.Context, eventID, accountID string) error {
|
||||
return r.db.WithContext(ctx).
|
||||
Model(&model.EventLog{}).
|
||||
Where("id = ?", eventID).
|
||||
Update("account_id", accountID).Error
|
||||
}
|
||||
|
||||
// InFlightByAccount counts pending (in-flight) events grouped by account_id, for
|
||||
// the accounts view's live "in-flight" column.
|
||||
func (r *EventRepository) InFlightByAccount(ctx context.Context) (map[string]int64, error) {
|
||||
type row struct {
|
||||
AccountID string `gorm:"column:account_id"`
|
||||
Count int64 `gorm:"column:count"`
|
||||
}
|
||||
var rows []row
|
||||
if err := r.db.WithContext(ctx).
|
||||
Model(&model.EventLog{}).
|
||||
Select("account_id, COUNT(*) AS count").
|
||||
Where("status = ? AND account_id <> ''", "pending").
|
||||
Group("account_id").
|
||||
Scan(&rows).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out := make(map[string]int64, len(rows))
|
||||
for _, item := range rows {
|
||||
out[item.AccountID] = item.Count
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func (r *EventRepository) RecentByFile(ctx context.Context, limit int) ([]model.EventLog, error) {
|
||||
if limit <= 0 {
|
||||
limit = 1000
|
||||
}
|
||||
var items []model.EventLog
|
||||
if err := r.db.WithContext(ctx).
|
||||
Where("file <> ''").
|
||||
Order("ts desc").
|
||||
Limit(limit).
|
||||
Find(&items).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return items, nil
|
||||
}
|
||||
|
||||
func (r *EventRepository) ModelSuccessCounts(ctx context.Context) (map[string]int64, error) {
|
||||
type row struct {
|
||||
Model string `gorm:"column:model"`
|
||||
Count int64 `gorm:"column:count"`
|
||||
}
|
||||
var rows []row
|
||||
if err := r.db.WithContext(ctx).
|
||||
Model(&model.EventLog{}).
|
||||
Select("model, COUNT(*) AS count").
|
||||
Where("status = ? AND model <> ''", "success").
|
||||
Group("model").
|
||||
Scan(&rows).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out := make(map[string]int64, len(rows))
|
||||
for _, item := range rows {
|
||||
out[item.Model] = item.Count
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func (r *EventRepository) UserSuccessCounts(ctx context.Context) (map[string]int64, error) {
|
||||
type row struct {
|
||||
UserID string `gorm:"column:user_id"`
|
||||
Count int64 `gorm:"column:count"`
|
||||
}
|
||||
var rows []row
|
||||
if err := r.db.WithContext(ctx).
|
||||
Model(&model.EventLog{}).
|
||||
Select("user_id, COUNT(*) AS count").
|
||||
Where("status = ? AND user_id <> ''", "success").
|
||||
Group("user_id").
|
||||
Scan(&rows).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out := make(map[string]int64, len(rows))
|
||||
for _, item := range rows {
|
||||
out[item.UserID] = item.Count
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func (r *EventRepository) DeleteAll(ctx context.Context) (int64, error) {
|
||||
result := r.db.WithContext(ctx).Where("1 = 1").Delete(&model.EventLog{})
|
||||
if result.Error != nil {
|
||||
return 0, result.Error
|
||||
}
|
||||
return result.RowsAffected, nil
|
||||
}
|
||||
|
||||
func (r *EventRepository) DeletePending(ctx context.Context) (int64, error) {
|
||||
result := r.db.WithContext(ctx).Where("status = ?", "pending").Delete(&model.EventLog{})
|
||||
if result.Error != nil {
|
||||
return 0, result.Error
|
||||
}
|
||||
return result.RowsAffected, nil
|
||||
}
|
||||
|
||||
// LatestByUser / PendingByUser take onlySource: when non-empty they match ONLY
|
||||
// that source. The playground passes "user" so it echoes ONLY the user's own
|
||||
// web generations — never admin model-tests ("admin") or API-key calls ("v1").
|
||||
func (r *EventRepository) LatestByUser(ctx context.Context, userID, onlySource string) (*model.EventLog, error) {
|
||||
var item model.EventLog
|
||||
q := r.db.WithContext(ctx).Model(&model.EventLog{}).Where("user_id = ?", userID)
|
||||
if strings.TrimSpace(onlySource) != "" {
|
||||
q = q.Where("source = ?", strings.TrimSpace(onlySource))
|
||||
}
|
||||
if err := q.Order("ts desc").First(&item).Error; err != nil {
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
return nil, nil
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
return &item, nil
|
||||
}
|
||||
|
||||
func (r *EventRepository) PendingByUser(ctx context.Context, userID, onlySource string) (*model.EventLog, error) {
|
||||
var item model.EventLog
|
||||
q := r.db.WithContext(ctx).Model(&model.EventLog{}).
|
||||
Where("user_id = ? AND status = ?", userID, "pending")
|
||||
if strings.TrimSpace(onlySource) != "" {
|
||||
q = q.Where("source = ?", strings.TrimSpace(onlySource))
|
||||
}
|
||||
if err := q.Order("ts desc").First(&item).Error; err != nil {
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
return nil, nil
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
return &item, nil
|
||||
}
|
||||
@@ -0,0 +1,69 @@
|
||||
package repo
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"time"
|
||||
|
||||
"backend/internal/model"
|
||||
"gorm.io/datatypes"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
type ModelRepository struct {
|
||||
db *gorm.DB
|
||||
}
|
||||
|
||||
func NewModelRepository(db *gorm.DB) *ModelRepository {
|
||||
return &ModelRepository{db: db}
|
||||
}
|
||||
|
||||
func (r *ModelRepository) List(ctx context.Context) ([]model.ModelConfig, error) {
|
||||
var items []model.ModelConfig
|
||||
// Higher weight floats to the top of the dropdown / admin list; ties fall
|
||||
// back to newest-first so order stays stable for equal-weight models.
|
||||
if err := r.db.WithContext(ctx).Order("weight desc, created_at desc").Find(&items).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return items, nil
|
||||
}
|
||||
|
||||
func (r *ModelRepository) Get(ctx context.Context, modelID string) (*model.ModelConfig, error) {
|
||||
var item model.ModelConfig
|
||||
if err := r.db.WithContext(ctx).First(&item, "id = ?", modelID).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &item, nil
|
||||
}
|
||||
|
||||
func JSONStrings(v datatypes.JSON) []string {
|
||||
if len(v) == 0 {
|
||||
return []string{}
|
||||
}
|
||||
var out []string
|
||||
if err := json.Unmarshal([]byte(v), &out); err == nil {
|
||||
return out
|
||||
}
|
||||
return []string{}
|
||||
}
|
||||
|
||||
func (r *ModelRepository) Create(ctx context.Context, item *model.ModelConfig) error {
|
||||
return r.db.WithContext(ctx).Create(item).Error
|
||||
}
|
||||
|
||||
func (r *ModelRepository) Update(ctx context.Context, modelID string, patch map[string]any) (*model.ModelConfig, error) {
|
||||
patch["updated_at"] = time.Now()
|
||||
if err := r.db.WithContext(ctx).Model(&model.ModelConfig{}).Where("id = ?", modelID).Updates(patch).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
var item model.ModelConfig
|
||||
if err := r.db.WithContext(ctx).First(&item, "id = ?", modelID).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &item, nil
|
||||
}
|
||||
|
||||
func (r *ModelRepository) Delete(ctx context.Context, modelID string) (int64, error) {
|
||||
res := r.db.WithContext(ctx).Delete(&model.ModelConfig{}, "id = ?", modelID)
|
||||
return res.RowsAffected, res.Error
|
||||
}
|
||||
@@ -0,0 +1,73 @@
|
||||
package repo
|
||||
|
||||
import (
|
||||
"context"
|
||||
"time"
|
||||
|
||||
"backend/internal/model"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
type RefreshProfileRepository struct {
|
||||
db *gorm.DB
|
||||
}
|
||||
|
||||
func NewRefreshProfileRepository(db *gorm.DB) *RefreshProfileRepository {
|
||||
return &RefreshProfileRepository{db: db}
|
||||
}
|
||||
|
||||
func (r *RefreshProfileRepository) List(ctx context.Context) ([]model.RefreshProfile, error) {
|
||||
var items []model.RefreshProfile
|
||||
if err := r.db.WithContext(ctx).
|
||||
Order("created_at desc").
|
||||
Find(&items).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return items, nil
|
||||
}
|
||||
|
||||
func (r *RefreshProfileRepository) Get(ctx context.Context, id string) (*model.RefreshProfile, error) {
|
||||
var item model.RefreshProfile
|
||||
if err := r.db.WithContext(ctx).First(&item, "id = ?", id).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &item, nil
|
||||
}
|
||||
|
||||
func (r *RefreshProfileRepository) Create(ctx context.Context, item *model.RefreshProfile) error {
|
||||
return r.db.WithContext(ctx).Create(item).Error
|
||||
}
|
||||
|
||||
func (r *RefreshProfileRepository) Update(ctx context.Context, id string, patch map[string]any) (*model.RefreshProfile, error) {
|
||||
patch["updated_at"] = time.Now()
|
||||
if err := r.db.WithContext(ctx).
|
||||
Model(&model.RefreshProfile{}).
|
||||
Where("id = ?", id).
|
||||
Updates(patch).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return r.Get(ctx, id)
|
||||
}
|
||||
|
||||
func (r *RefreshProfileRepository) Delete(ctx context.Context, id string) error {
|
||||
return r.db.WithContext(ctx).Delete(&model.RefreshProfile{}, "id = ?", id).Error
|
||||
}
|
||||
|
||||
func (r *RefreshProfileRepository) DeleteByIDs(ctx context.Context, ids []string) error {
|
||||
if len(ids) == 0 {
|
||||
return nil
|
||||
}
|
||||
return r.db.WithContext(ctx).Delete(&model.RefreshProfile{}, "id IN ?", ids).Error
|
||||
}
|
||||
|
||||
// ListDue returns enabled profiles whose next_retry_at has passed (or is unset,
|
||||
// e.g. freshly imported). The background maintenance loop refreshes these.
|
||||
func (r *RefreshProfileRepository) ListDue(ctx context.Context, now time.Time) ([]model.RefreshProfile, error) {
|
||||
var items []model.RefreshProfile
|
||||
if err := r.db.WithContext(ctx).
|
||||
Where("enabled = ? AND (next_retry_at IS NULL OR next_retry_at <= ?)", true, now).
|
||||
Find(&items).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return items, nil
|
||||
}
|
||||
@@ -0,0 +1,99 @@
|
||||
package repo
|
||||
|
||||
import (
|
||||
"context"
|
||||
"sort"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"backend/internal/model"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
type ShowcaseRepository struct {
|
||||
db *gorm.DB
|
||||
}
|
||||
|
||||
func NewShowcaseRepository(db *gorm.DB) *ShowcaseRepository {
|
||||
return &ShowcaseRepository{db: db}
|
||||
}
|
||||
|
||||
func (r *ShowcaseRepository) IsPublicFile(ctx context.Context, rel string) (bool, error) {
|
||||
normalized := strings.TrimLeft(strings.TrimSpace(rel), "/")
|
||||
if normalized == "" {
|
||||
return false, nil
|
||||
}
|
||||
|
||||
var count int64
|
||||
if err := r.db.WithContext(ctx).
|
||||
Model(&model.ShowcaseItem{}).
|
||||
Where("image = ? OR image = ?", normalized, "/"+normalized).
|
||||
Count(&count).Error; err != nil {
|
||||
return false, err
|
||||
}
|
||||
return count > 0, nil
|
||||
}
|
||||
|
||||
// PublicFileSet returns the set of image keys referenced by any showcase item
|
||||
// (normalized, no leading slash). The media-prune sweep uses it to never delete
|
||||
// a file the homepage still shows, regardless of how old the file is.
|
||||
func (r *ShowcaseRepository) PublicFileSet(ctx context.Context) (map[string]struct{}, error) {
|
||||
var images []string
|
||||
if err := r.db.WithContext(ctx).
|
||||
Model(&model.ShowcaseItem{}).
|
||||
Where("image <> ''").
|
||||
Pluck("image", &images).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
set := make(map[string]struct{}, len(images))
|
||||
for _, img := range images {
|
||||
n := strings.TrimLeft(strings.TrimSpace(img), "/")
|
||||
if n != "" {
|
||||
set[n] = struct{}{}
|
||||
}
|
||||
}
|
||||
return set, nil
|
||||
}
|
||||
|
||||
func (r *ShowcaseRepository) Grouped(ctx context.Context) (map[string][]model.ShowcaseItem, error) {
|
||||
var items []model.ShowcaseItem
|
||||
if err := r.db.WithContext(ctx).Find(&items).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
grouped := map[string][]model.ShowcaseItem{
|
||||
"hero": {},
|
||||
"bento": {},
|
||||
"work": {},
|
||||
}
|
||||
for _, item := range items {
|
||||
grouped[item.Kind] = append(grouped[item.Kind], item)
|
||||
}
|
||||
for kind := range grouped {
|
||||
sort.Slice(grouped[kind], func(i, j int) bool {
|
||||
return grouped[kind][i].Weight > grouped[kind][j].Weight
|
||||
})
|
||||
}
|
||||
return grouped, nil
|
||||
}
|
||||
|
||||
func (r *ShowcaseRepository) Create(ctx context.Context, item *model.ShowcaseItem) error {
|
||||
return r.db.WithContext(ctx).Create(item).Error
|
||||
}
|
||||
|
||||
func (r *ShowcaseRepository) Update(ctx context.Context, entryID string, patch map[string]any) (*model.ShowcaseItem, error) {
|
||||
patch["updated_at"] = time.Now()
|
||||
if err := r.db.WithContext(ctx).Model(&model.ShowcaseItem{}).Where("id = ?", entryID).Updates(patch).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
var item model.ShowcaseItem
|
||||
if err := r.db.WithContext(ctx).First(&item, "id = ?", entryID).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &item, nil
|
||||
}
|
||||
|
||||
func (r *ShowcaseRepository) Delete(ctx context.Context, entryID string) (int64, error) {
|
||||
res := r.db.WithContext(ctx).Delete(&model.ShowcaseItem{}, "id = ?", entryID)
|
||||
return res.RowsAffected, res.Error
|
||||
}
|
||||
@@ -0,0 +1,96 @@
|
||||
package repo
|
||||
|
||||
import (
|
||||
"context"
|
||||
"time"
|
||||
|
||||
"backend/internal/model"
|
||||
"github.com/redis/go-redis/v9"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
const siteSettingCachePrefix = "setting:"
|
||||
|
||||
// siteSettingCacheTTL is a safety-net expiry; writes invalidate eagerly, so this
|
||||
// only bounds staleness if an invalidation is ever missed (e.g. Redis blip).
|
||||
const siteSettingCacheTTL = 5 * time.Minute
|
||||
|
||||
type SiteSettingRepository struct {
|
||||
db *gorm.DB
|
||||
cache *redis.Client
|
||||
}
|
||||
|
||||
// NewSiteSettingRepository wires the config KV store. cache may be nil, in which
|
||||
// case the repository transparently falls back to DB-only access.
|
||||
func NewSiteSettingRepository(db *gorm.DB, cache *redis.Client) *SiteSettingRepository {
|
||||
return &SiteSettingRepository{db: db, cache: cache}
|
||||
}
|
||||
|
||||
func (r *SiteSettingRepository) cacheKey(key string) string {
|
||||
return siteSettingCachePrefix + key
|
||||
}
|
||||
|
||||
func (r *SiteSettingRepository) GetValue(ctx context.Context, key string) (string, error) {
|
||||
if r.cache != nil {
|
||||
if v, err := r.cache.Get(ctx, r.cacheKey(key)).Result(); err == nil {
|
||||
return v, nil
|
||||
}
|
||||
// redis.Nil (miss) or any transient cache error -> fall through to DB.
|
||||
}
|
||||
|
||||
value := ""
|
||||
var setting model.SiteSetting
|
||||
if err := r.db.WithContext(ctx).First(&setting, "key = ?", key).Error; err != nil {
|
||||
if err != gorm.ErrRecordNotFound {
|
||||
return "", err
|
||||
}
|
||||
// Not found stays as "" — still cached below to absorb repeated misses.
|
||||
} else {
|
||||
value = setting.Value
|
||||
}
|
||||
|
||||
if r.cache != nil {
|
||||
_ = r.cache.Set(ctx, r.cacheKey(key), value, siteSettingCacheTTL).Err()
|
||||
}
|
||||
return value, nil
|
||||
}
|
||||
|
||||
func (r *SiteSettingRepository) UpsertValue(ctx context.Context, key, value string) error {
|
||||
if err := r.db.WithContext(ctx).Save(&model.SiteSetting{
|
||||
Key: key,
|
||||
Value: value,
|
||||
}).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
r.invalidate(ctx, key)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *SiteSettingRepository) UpsertValues(ctx context.Context, values map[string]string) error {
|
||||
if err := r.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
|
||||
for key, value := range values {
|
||||
if err := tx.Save(&model.SiteSetting{
|
||||
Key: key,
|
||||
Value: value,
|
||||
}).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}); err != nil {
|
||||
return err
|
||||
}
|
||||
for key := range values {
|
||||
r.invalidate(ctx, key)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// invalidate drops the cached entry so the next read repopulates from the DB.
|
||||
// Deleting (rather than overwriting) keeps writes simple and race-tolerant.
|
||||
func (r *SiteSettingRepository) invalidate(ctx context.Context, key string) {
|
||||
if r.cache == nil {
|
||||
return
|
||||
}
|
||||
_ = r.cache.Del(ctx, r.cacheKey(key)).Err()
|
||||
}
|
||||
@@ -0,0 +1,390 @@
|
||||
package repo
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"backend/internal/model"
|
||||
"gorm.io/datatypes"
|
||||
"gorm.io/gorm"
|
||||
"gorm.io/gorm/clause"
|
||||
)
|
||||
|
||||
type TokenRepository struct {
|
||||
db *gorm.DB
|
||||
}
|
||||
|
||||
func NewTokenRepository(db *gorm.DB) *TokenRepository {
|
||||
return &TokenRepository{db: db}
|
||||
}
|
||||
|
||||
func (r *TokenRepository) List(ctx context.Context) ([]model.TokenAccount, error) {
|
||||
var items []model.TokenAccount
|
||||
if err := r.db.WithContext(ctx).
|
||||
Order("pool asc, created_at desc").
|
||||
Find(&items).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return items, nil
|
||||
}
|
||||
|
||||
func (r *TokenRepository) ListByPool(ctx context.Context, pool string) ([]model.TokenAccount, error) {
|
||||
var items []model.TokenAccount
|
||||
if err := r.db.WithContext(ctx).
|
||||
Where("pool = ?", pool).
|
||||
Order("created_at desc").
|
||||
Find(&items).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return items, nil
|
||||
}
|
||||
|
||||
func (r *TokenRepository) Get(ctx context.Context, pool, id string) (*model.TokenAccount, error) {
|
||||
var item model.TokenAccount
|
||||
if err := r.db.WithContext(ctx).
|
||||
First(&item, "pool = ? AND id = ?", pool, id).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &item, nil
|
||||
}
|
||||
|
||||
// GetByPoolEmail finds an account in a pool by its account_email (the logical
|
||||
// identity for import dedup). Returns (nil, nil) when none / email is blank.
|
||||
func (r *TokenRepository) GetByPoolEmail(ctx context.Context, pool, email string) (*model.TokenAccount, error) {
|
||||
email = strings.TrimSpace(email)
|
||||
if email == "" {
|
||||
return nil, nil
|
||||
}
|
||||
var item model.TokenAccount
|
||||
err := r.db.WithContext(ctx).
|
||||
Where("pool = ? AND account_email = ?", pool, email).
|
||||
First(&item).Error
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
return nil, nil
|
||||
}
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &item, nil
|
||||
}
|
||||
|
||||
func (r *TokenRepository) Create(ctx context.Context, item *model.TokenAccount) error {
|
||||
return r.db.WithContext(ctx).Create(item).Error
|
||||
}
|
||||
|
||||
func (r *TokenRepository) Update(ctx context.Context, pool, id string, patch map[string]any) (*model.TokenAccount, error) {
|
||||
patch["updated_at"] = time.Now()
|
||||
if err := r.db.WithContext(ctx).
|
||||
Model(&model.TokenAccount{}).
|
||||
Where("pool = ? AND id = ?", pool, id).
|
||||
Updates(patch).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return r.Get(ctx, pool, id)
|
||||
}
|
||||
|
||||
// ReserveQuota atomically pre-deducts `amount` from an account's cached image
|
||||
// token balance under a row lock, so concurrent picks of the same near-empty
|
||||
// account can never over-commit it. Returns:
|
||||
// - allowed=true, deducted=true: balance was known and ≥ amount → decremented.
|
||||
// - allowed=true, deducted=false: balance unknown → allowed without a hold
|
||||
// (benefit of the doubt; a post-render reconcile writes the real value).
|
||||
// - allowed=false: balance known and < amount → caller should fail over.
|
||||
// RefundQuota releases a hold made with deducted=true when the render fails.
|
||||
func (r *TokenRepository) ReserveQuota(ctx context.Context, pool, id string, amount int) (allowed, deducted bool, err error) {
|
||||
err = r.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
|
||||
var item model.TokenAccount
|
||||
if e := tx.Clauses(clause.Locking{Strength: "UPDATE"}).
|
||||
First(&item, "pool = ? AND id = ?", pool, id).Error; e != nil {
|
||||
return e
|
||||
}
|
||||
rem, known := metaInt(item.Meta, "cached_quota_remaining")
|
||||
if !known {
|
||||
allowed, deducted = true, false
|
||||
return nil
|
||||
}
|
||||
if rem < amount {
|
||||
allowed, deducted = false, false
|
||||
return nil
|
||||
}
|
||||
meta := cloneMeta(item.Meta)
|
||||
meta["cached_quota_remaining"] = rem - amount
|
||||
if e := tx.Model(&model.TokenAccount{}).
|
||||
Where("pool = ? AND id = ?", pool, id).
|
||||
Updates(map[string]any{"meta": meta, "updated_at": time.Now()}).Error; e != nil {
|
||||
return e
|
||||
}
|
||||
allowed, deducted = true, true
|
||||
return nil
|
||||
})
|
||||
return allowed, deducted, err
|
||||
}
|
||||
|
||||
// RefundQuota atomically adds `amount` back to cached_quota_remaining (releasing a
|
||||
// hold from a reservation whose render then failed). No-op if the balance is
|
||||
// unknown. Row-locked like ReserveQuota.
|
||||
func (r *TokenRepository) RefundQuota(ctx context.Context, pool, id string, amount int) error {
|
||||
return r.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
|
||||
var item model.TokenAccount
|
||||
if e := tx.Clauses(clause.Locking{Strength: "UPDATE"}).
|
||||
First(&item, "pool = ? AND id = ?", pool, id).Error; e != nil {
|
||||
return e
|
||||
}
|
||||
rem, known := metaInt(item.Meta, "cached_quota_remaining")
|
||||
if !known {
|
||||
return nil
|
||||
}
|
||||
meta := cloneMeta(item.Meta)
|
||||
meta["cached_quota_remaining"] = rem + amount
|
||||
return tx.Model(&model.TokenAccount{}).
|
||||
Where("pool = ? AND id = ?", pool, id).
|
||||
Updates(map[string]any{"meta": meta, "updated_at": time.Now()}).Error
|
||||
})
|
||||
}
|
||||
|
||||
func cloneMeta(m datatypes.JSONMap) datatypes.JSONMap {
|
||||
out := datatypes.JSONMap{}
|
||||
for k, v := range m {
|
||||
out[k] = v
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func metaInt(m datatypes.JSONMap, key string) (int, bool) {
|
||||
if m == nil {
|
||||
return 0, false
|
||||
}
|
||||
v, ok := m[key]
|
||||
if !ok || v == nil {
|
||||
return 0, false
|
||||
}
|
||||
switch x := v.(type) {
|
||||
case int:
|
||||
return x, true
|
||||
case int64:
|
||||
return int(x), true
|
||||
case float64:
|
||||
return int(x), true
|
||||
case json.Number:
|
||||
n, e := x.Int64()
|
||||
if e != nil {
|
||||
return 0, false
|
||||
}
|
||||
return int(n), true
|
||||
case string:
|
||||
n, e := strconv.Atoi(strings.TrimSpace(x))
|
||||
if e != nil {
|
||||
return 0, false
|
||||
}
|
||||
return n, true
|
||||
default:
|
||||
return 0, false
|
||||
}
|
||||
}
|
||||
|
||||
// TouchLastUsed stamps last_used_at at the moment a token is SELECTED, so the
|
||||
// accounts view reflects an accurate "last used" time. Rotation order is driven
|
||||
// by the in-memory strict round-robin cursor in the service layer (see
|
||||
// V1Service.rotateRoundRobin), not by this timestamp.
|
||||
func (r *TokenRepository) TouchLastUsed(ctx context.Context, id string) error {
|
||||
return r.db.WithContext(ctx).
|
||||
Model(&model.TokenAccount{}).
|
||||
Where("id = ?", id).
|
||||
Update("last_used_at", time.Now()).Error
|
||||
}
|
||||
|
||||
// IncrementFail bumps an account's failure counters by one. Used to attribute
|
||||
// an abandoned (purged) generation's failure back to the account it was using,
|
||||
// since that generation never reached the normal markTokenFailure path.
|
||||
func (r *TokenRepository) IncrementFail(ctx context.Context, id string) error {
|
||||
return r.db.WithContext(ctx).
|
||||
Model(&model.TokenAccount{}).
|
||||
Where("id = ?", id).
|
||||
Updates(map[string]any{
|
||||
"fail_total": gorm.Expr("fail_total + 1"),
|
||||
"fails": gorm.Expr("fails + 1"),
|
||||
"updated_at": time.Now(),
|
||||
}).Error
|
||||
}
|
||||
|
||||
func (r *TokenRepository) Delete(ctx context.Context, pool, id string) (int64, error) {
|
||||
res := r.db.WithContext(ctx).
|
||||
Delete(&model.TokenAccount{}, "pool = ? AND id = ?", pool, id)
|
||||
return res.RowsAffected, res.Error
|
||||
}
|
||||
|
||||
// DeleteByIDs removes accounts by id across pools (ids are globally unique),
|
||||
// for bulk delete. Returns the number of rows removed.
|
||||
func (r *TokenRepository) DeleteByIDs(ctx context.Context, ids []string) (int64, error) {
|
||||
if len(ids) == 0 {
|
||||
return 0, nil
|
||||
}
|
||||
res := r.db.WithContext(ctx).Delete(&model.TokenAccount{}, "id IN ?", ids)
|
||||
return res.RowsAffected, res.Error
|
||||
}
|
||||
|
||||
// leonardoDailyTokens is the free-tier daily allowance restored at each reset.
|
||||
// A paid account's true balance is reconciled on its next successful render.
|
||||
const leonardoDailyTokens = 150
|
||||
|
||||
// RecoverQuota reactivates quota-exhausted tokens whose reset time has passed.
|
||||
// Reset source: cached_quota_reset_after (upstream marker) first, else the
|
||||
// quota_recover_at fallback stamped when the token was marked quota-exhausted.
|
||||
// Mirrors Python TokenPool.recover_quota; returns the count reactivated.
|
||||
// RecoverQuota reactivates quota-exhausted tokens whose reset time has passed and
|
||||
// returns the accounts it recovered, so the caller can re-sync their real balance
|
||||
// (the providers only sync quota when accessed).
|
||||
func (r *TokenRepository) RecoverQuota(ctx context.Context) ([]model.TokenAccount, error) {
|
||||
// Also pick up accounts that are only single-kind limited (image_limited /
|
||||
// video_limited) — those keep status "active" and would otherwise never have
|
||||
// their per-kind flag cleared. Adobe resets both kinds at once, so the shared
|
||||
// reset time gates recovery for all of them.
|
||||
var items []model.TokenAccount
|
||||
if err := r.db.WithContext(ctx).
|
||||
Where("status = ? OR image_limited = ? OR video_limited = ?", "quota", true, true).
|
||||
Find(&items).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
now := time.Now()
|
||||
var recovered []model.TokenAccount
|
||||
for i := range items {
|
||||
t := &items[i]
|
||||
// Runway's reset marker is the JWT expiry, not a quota-refresh time, and
|
||||
// there's no way to refresh a bare JWT — so a runway account is never
|
||||
// "recovered"; it's expired-to-dead by ExpireByReset instead.
|
||||
if t.Pool == "runway" {
|
||||
continue
|
||||
}
|
||||
reset := parseResetMarker(t.CachedQuotaResetAfter)
|
||||
if reset == nil {
|
||||
reset = t.QuotaRecoverAt
|
||||
}
|
||||
if reset == nil || now.Before(*reset) {
|
||||
continue
|
||||
}
|
||||
patch := map[string]any{
|
||||
"fails": 0,
|
||||
"quota_recover_at": nil,
|
||||
"image_limited": false,
|
||||
"video_limited": false,
|
||||
}
|
||||
// Only flip status back to active if it was sunk to "quota" (both kinds
|
||||
// limited); a single-kind limit left status untouched.
|
||||
if t.Status == "quota" {
|
||||
patch["status"] = "active"
|
||||
}
|
||||
// Leonardo's free tokens fully renew at each daily reset — restore the
|
||||
// balance and advance the reset marker to the next 08:00 Beijing (== next
|
||||
// UTC midnight), so the account is immediately usable instead of stuck at a
|
||||
// stale 0. A paid account's real balance is corrected on its next render.
|
||||
if t.Pool == "leonardo" || t.Pool == "krea" || t.Pool == "imagine" {
|
||||
meta := cloneMeta(t.Meta)
|
||||
if t.Pool == "leonardo" {
|
||||
meta["cached_quota_remaining"] = leonardoDailyTokens
|
||||
} else {
|
||||
// Krea/Imagine balances re-sync from upstream (billing-data / v1/credit)
|
||||
// on next probe — drop the stale value so the account isn't shown as
|
||||
// empty after reset.
|
||||
delete(meta, "cached_quota_remaining")
|
||||
}
|
||||
meta["cached_quota_at"] = int(now.Unix())
|
||||
patch["meta"] = meta
|
||||
patch["cached_quota_reset_after"] = time.Unix((now.Unix()/86400+1)*86400, 0).UTC().Format(time.RFC3339)
|
||||
}
|
||||
if _, err := r.Update(ctx, t.Pool, t.ID, patch); err != nil {
|
||||
return recovered, err
|
||||
}
|
||||
recovered = append(recovered, *t)
|
||||
}
|
||||
return recovered, nil
|
||||
}
|
||||
|
||||
// RollResetMarkers advances a stale (past) daily-reset marker to its next future
|
||||
// occurrence — same time-of-day, +N whole days — for ACTIVE accounts of the given
|
||||
// daily-reset pools, so the 恢复时间 column always shows the upcoming reset rather
|
||||
// than yesterday's. Only active accounts are rolled: a 限额 account must keep its
|
||||
// past marker so RecoverQuota can recover it (rolling it forward early would
|
||||
// prevent recovery). Returns the number advanced.
|
||||
func (r *TokenRepository) RollResetMarkers(ctx context.Context, pools []string) (int, error) {
|
||||
var items []model.TokenAccount
|
||||
if err := r.db.WithContext(ctx).
|
||||
Where("pool IN ? AND dead = ? AND status = ? AND image_limited = ? AND video_limited = ? AND cached_quota_reset_after <> ''",
|
||||
pools, false, "active", false, false).
|
||||
Find(&items).Error; err != nil {
|
||||
return 0, err
|
||||
}
|
||||
now := time.Now()
|
||||
n := 0
|
||||
for i := range items {
|
||||
t := &items[i]
|
||||
reset := parseResetMarker(t.CachedQuotaResetAfter)
|
||||
if reset == nil || !reset.Before(now) {
|
||||
continue // unparseable or already in the future
|
||||
}
|
||||
next := *reset
|
||||
for !next.After(now) {
|
||||
next = next.Add(24 * time.Hour)
|
||||
}
|
||||
if _, err := r.Update(ctx, t.Pool, t.ID, map[string]any{
|
||||
"cached_quota_reset_after": next.UTC().Format(time.RFC3339),
|
||||
}); err != nil {
|
||||
return n, err
|
||||
}
|
||||
n++
|
||||
}
|
||||
return n, nil
|
||||
}
|
||||
|
||||
// ExpireByReset marks accounts of a pool dead once their reset marker has passed.
|
||||
// For runway the marker IS the JWT expiry and there's no refresh, so an expired
|
||||
// token can only 401 — we proactively flip it to disabled+dead (the same end
|
||||
// state a 401 would produce) instead of leaving a doomed account "active".
|
||||
func (r *TokenRepository) ExpireByReset(ctx context.Context, pool string) (int, error) {
|
||||
var items []model.TokenAccount
|
||||
if err := r.db.WithContext(ctx).
|
||||
Where("pool = ? AND dead = ?", pool, false).
|
||||
Find(&items).Error; err != nil {
|
||||
return 0, err
|
||||
}
|
||||
now := time.Now()
|
||||
expired := 0
|
||||
for i := range items {
|
||||
t := &items[i]
|
||||
reset := parseResetMarker(t.CachedQuotaResetAfter)
|
||||
if reset == nil || now.Before(*reset) {
|
||||
continue
|
||||
}
|
||||
if _, err := r.Update(ctx, t.Pool, t.ID, map[string]any{
|
||||
"status": "disabled",
|
||||
"dead": true,
|
||||
}); err != nil {
|
||||
return expired, err
|
||||
}
|
||||
expired++
|
||||
}
|
||||
return expired, nil
|
||||
}
|
||||
|
||||
// parseResetMarker best-effort parses a quota reset marker into a time. Accepts
|
||||
// epoch seconds (numeric string) or ISO-8601 (e.g. Adobe's available_until
|
||||
// "2026-06-16T23:59:59.999Z"). Returns nil if unparseable.
|
||||
func parseResetMarker(v string) *time.Time {
|
||||
v = strings.TrimSpace(v)
|
||||
if v == "" {
|
||||
return nil
|
||||
}
|
||||
if f, err := strconv.ParseFloat(v, 64); err == nil && f > 946684800 {
|
||||
t := time.Unix(int64(f), 0)
|
||||
return &t
|
||||
}
|
||||
for _, layout := range []string{time.RFC3339, "2006-01-02T15:04:05.999Z07:00", "2006-01-02T15:04:05Z07:00"} {
|
||||
if t, err := time.Parse(layout, v); err == nil {
|
||||
return &t
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,627 @@
|
||||
package repo
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"backend/internal/model"
|
||||
"gorm.io/gorm"
|
||||
"gorm.io/gorm/clause"
|
||||
)
|
||||
|
||||
type UserRepository struct {
|
||||
db *gorm.DB
|
||||
}
|
||||
|
||||
var ErrAlreadyCheckedInToday = errors.New("already checked in today")
|
||||
|
||||
type InviteStats struct {
|
||||
InviteCount int64 `json:"invite_count"`
|
||||
InviteEarned int `json:"invite_earned"`
|
||||
}
|
||||
|
||||
type InviteRecord struct {
|
||||
Name string `json:"name,omitempty"`
|
||||
Inviter string `json:"inviter,omitempty"`
|
||||
Invitee string `json:"invitee,omitempty"`
|
||||
Reward int `json:"reward"`
|
||||
RegisteredAt time.Time `json:"registered_at"`
|
||||
CompletedAt *time.Time `json:"completed_at"`
|
||||
Status string `json:"status"`
|
||||
}
|
||||
|
||||
type InviteLogStats struct {
|
||||
Total int64 `json:"total"`
|
||||
Completed int64 `json:"completed"`
|
||||
Pending int64 `json:"pending"`
|
||||
RewardPaid int64 `json:"reward_paid"`
|
||||
}
|
||||
|
||||
type CheckinResult struct {
|
||||
Already bool `json:"already"`
|
||||
Awarded int `json:"awarded"`
|
||||
Streak int `json:"streak"`
|
||||
Credits float64 `json:"credits"`
|
||||
}
|
||||
|
||||
func NewUserRepository(db *gorm.DB) *UserRepository {
|
||||
return &UserRepository{db: db}
|
||||
}
|
||||
|
||||
func (r *UserRepository) GetByID(ctx context.Context, userID string) (*model.User, error) {
|
||||
var user model.User
|
||||
if err := r.db.WithContext(ctx).Preload("APIKeys").First(&user, "id = ?", userID).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &user, nil
|
||||
}
|
||||
|
||||
func (r *UserRepository) GetByIdentifier(ctx context.Context, identifier string) (*model.User, error) {
|
||||
ident := strings.TrimSpace(identifier)
|
||||
if ident == "" {
|
||||
return nil, gorm.ErrRecordNotFound
|
||||
}
|
||||
|
||||
var user model.User
|
||||
q := r.db.WithContext(ctx).Preload("APIKeys")
|
||||
if strings.Contains(ident, "@") {
|
||||
if err := q.First(&user, "email = ?", strings.ToLower(ident)).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &user, nil
|
||||
}
|
||||
|
||||
if err := q.First(&user, "LOWER(name) = ?", strings.ToLower(ident)).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &user, nil
|
||||
}
|
||||
|
||||
func (r *UserRepository) GetByInviteCode(ctx context.Context, code string) (*model.User, error) {
|
||||
var user model.User
|
||||
if err := r.db.WithContext(ctx).First(&user, "invite_code = ?", strings.ToUpper(strings.TrimSpace(code))).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &user, nil
|
||||
}
|
||||
|
||||
func (r *UserRepository) List(ctx context.Context) ([]model.User, error) {
|
||||
var users []model.User
|
||||
if err := r.db.WithContext(ctx).Preload("APIKeys").Order("created_at desc").Find(&users).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return users, nil
|
||||
}
|
||||
|
||||
func (r *UserRepository) ExistsEmail(ctx context.Context, email, excludeUserID string) (bool, error) {
|
||||
var count int64
|
||||
q := r.db.WithContext(ctx).Model(&model.User{}).Where("email = ?", strings.ToLower(strings.TrimSpace(email)))
|
||||
if strings.TrimSpace(excludeUserID) != "" {
|
||||
q = q.Where("id <> ?", strings.TrimSpace(excludeUserID))
|
||||
}
|
||||
if err := q.Count(&count).Error; err != nil {
|
||||
return false, err
|
||||
}
|
||||
return count > 0, nil
|
||||
}
|
||||
|
||||
func (r *UserRepository) ExistsName(ctx context.Context, name, excludeUserID string) (bool, error) {
|
||||
var count int64
|
||||
q := r.db.WithContext(ctx).Model(&model.User{}).Where("LOWER(name) = ?", strings.ToLower(strings.TrimSpace(name)))
|
||||
if strings.TrimSpace(excludeUserID) != "" {
|
||||
q = q.Where("id <> ?", strings.TrimSpace(excludeUserID))
|
||||
}
|
||||
if err := q.Count(&count).Error; err != nil {
|
||||
return false, err
|
||||
}
|
||||
return count > 0, nil
|
||||
}
|
||||
|
||||
func (r *UserRepository) GetByAPIKeyHash(ctx context.Context, keyHash string) (*model.User, error) {
|
||||
var apiKey model.APIKey
|
||||
if err := r.db.WithContext(ctx).First(&apiKey, "key_hash = ?", keyHash).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
var user model.User
|
||||
if err := r.db.WithContext(ctx).Preload("APIKeys").First(&user, "id = ?", apiKey.UserID).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &user, nil
|
||||
}
|
||||
|
||||
func (r *UserRepository) TouchLogin(ctx context.Context, userID, ip string) error {
|
||||
now := time.Now()
|
||||
return r.db.WithContext(ctx).
|
||||
Model(&model.User{}).
|
||||
Where("id = ?", userID).
|
||||
Updates(map[string]any{
|
||||
"last_login_at": now,
|
||||
"last_login_ip": ip,
|
||||
}).Error
|
||||
}
|
||||
|
||||
func (r *UserRepository) HasAdmin(ctx context.Context) (bool, error) {
|
||||
var count int64
|
||||
if err := r.db.WithContext(ctx).
|
||||
Model(&model.User{}).
|
||||
Where("role = ?", "admin").
|
||||
Count(&count).Error; err != nil {
|
||||
return false, err
|
||||
}
|
||||
return count > 0, nil
|
||||
}
|
||||
|
||||
func (r *UserRepository) Stats(ctx context.Context) (map[string]any, error) {
|
||||
var total, active, disabled, admins int64
|
||||
if err := r.db.WithContext(ctx).Model(&model.User{}).Count(&total).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := r.db.WithContext(ctx).Model(&model.User{}).Where("status = ?", "active").Count(&active).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := r.db.WithContext(ctx).Model(&model.User{}).Where("status = ?", "disabled").Count(&disabled).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := r.db.WithContext(ctx).Model(&model.User{}).Where("role = ?", "admin").Count(&admins).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
type sumRow struct {
|
||||
Total *float64 `gorm:"column:total"`
|
||||
}
|
||||
var credits sumRow
|
||||
if err := r.db.WithContext(ctx).
|
||||
Model(&model.User{}).
|
||||
Select("SUM(credits) AS total").
|
||||
Scan(&credits).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
now := time.Now()
|
||||
dayCut := now.Add(-24 * time.Hour)
|
||||
weekCut := now.Add(-7 * 24 * time.Hour)
|
||||
var new24h, new7d, active24h int64
|
||||
if err := r.db.WithContext(ctx).Model(&model.User{}).Where("created_at >= ?", dayCut).Count(&new24h).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := r.db.WithContext(ctx).Model(&model.User{}).Where("created_at >= ?", weekCut).Count(&new7d).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := r.db.WithContext(ctx).Model(&model.User{}).Where("last_login_at >= ?", dayCut).Count(&active24h).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
creditsTotal := 0.0
|
||||
if credits.Total != nil {
|
||||
creditsTotal = *credits.Total
|
||||
}
|
||||
|
||||
return map[string]any{
|
||||
"total": total,
|
||||
"active": active,
|
||||
"disabled": disabled,
|
||||
"admins": admins,
|
||||
"credits_total": creditsTotal,
|
||||
"new_24h": new24h,
|
||||
"new_7d": new7d,
|
||||
"active_24h": active24h,
|
||||
}, nil
|
||||
}
|
||||
|
||||
type CheckinStats struct {
|
||||
TodayCount int64 `json:"today_count"`
|
||||
MaxStreak int64 `json:"max_streak"`
|
||||
}
|
||||
|
||||
// CheckinStats counts users who checked in today and the longest active streak —
|
||||
// a single-query summary for the admin dashboard's 签到 card.
|
||||
func (r *UserRepository) CheckinStats(ctx context.Context) (*CheckinStats, error) {
|
||||
today := time.Now().Format("2006-01-02")
|
||||
type row struct {
|
||||
TodayCount int64 `gorm:"column:today_count"`
|
||||
MaxStreak int64 `gorm:"column:max_streak"`
|
||||
}
|
||||
var out row
|
||||
if err := r.db.WithContext(ctx).
|
||||
Model(&model.User{}).
|
||||
Select("COUNT(*) FILTER (WHERE checkin_last = ?) AS today_count, COALESCE(MAX(checkin_streak), 0) AS max_streak", today).
|
||||
Scan(&out).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &CheckinStats{TodayCount: out.TodayCount, MaxStreak: out.MaxStreak}, nil
|
||||
}
|
||||
|
||||
type InviteSummary struct {
|
||||
Total int64 `json:"total"`
|
||||
Completed int64 `json:"completed"`
|
||||
}
|
||||
|
||||
// InviteSummary is a lightweight count of invited users (and how many have had
|
||||
// their reward granted). Cheaper than AllInvites — no JOIN, no record list —
|
||||
// for the dashboard which polls frequently.
|
||||
func (r *UserRepository) InviteSummary(ctx context.Context) (*InviteSummary, error) {
|
||||
type row struct {
|
||||
Total int64 `gorm:"column:total"`
|
||||
Completed int64 `gorm:"column:completed"`
|
||||
}
|
||||
var out row
|
||||
if err := r.db.WithContext(ctx).
|
||||
Model(&model.User{}).
|
||||
Select("COUNT(*) AS total, COUNT(*) FILTER (WHERE invite_reward_done) AS completed").
|
||||
Where("invited_by IS NOT NULL AND invited_by <> ''").
|
||||
Scan(&out).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &InviteSummary{Total: out.Total, Completed: out.Completed}, nil
|
||||
}
|
||||
|
||||
func (r *UserRepository) Create(ctx context.Context, user *model.User) error {
|
||||
return r.db.WithContext(ctx).Create(user).Error
|
||||
}
|
||||
|
||||
func (r *UserRepository) Update(ctx context.Context, userID string, patch map[string]any) (*model.User, error) {
|
||||
patch["updated_at"] = time.Now()
|
||||
if err := r.db.WithContext(ctx).Model(&model.User{}).Where("id = ?", userID).Updates(patch).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return r.GetByID(ctx, userID)
|
||||
}
|
||||
|
||||
func (r *UserRepository) Delete(ctx context.Context, userID string) (int64, error) {
|
||||
res := r.db.WithContext(ctx).Delete(&model.User{}, "id = ?", userID)
|
||||
return res.RowsAffected, res.Error
|
||||
}
|
||||
|
||||
func (r *UserRepository) DeleteByIDs(ctx context.Context, ids []string) (int64, error) {
|
||||
if len(ids) == 0 {
|
||||
return 0, nil
|
||||
}
|
||||
res := r.db.WithContext(ctx).Delete(&model.User{}, "id IN ?", ids)
|
||||
return res.RowsAffected, res.Error
|
||||
}
|
||||
|
||||
func (r *UserRepository) SetPasswordByEmail(ctx context.Context, email, passwordHash string) (*model.User, error) {
|
||||
if err := r.db.WithContext(ctx).
|
||||
Model(&model.User{}).
|
||||
Where("email = ?", strings.ToLower(strings.TrimSpace(email))).
|
||||
Updates(map[string]any{
|
||||
"password_hash": passwordHash,
|
||||
"updated_at": time.Now(),
|
||||
}).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
var user model.User
|
||||
if err := r.db.WithContext(ctx).Preload("APIKeys").First(&user, "email = ?", strings.ToLower(strings.TrimSpace(email))).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &user, nil
|
||||
}
|
||||
|
||||
func (r *UserRepository) TouchAPIKeyUsage(ctx context.Context, keyHash string) error {
|
||||
now := time.Now()
|
||||
return r.db.WithContext(ctx).
|
||||
Model(&model.APIKey{}).
|
||||
Where("key_hash = ?", keyHash).
|
||||
Update("last_used_at", now).Error
|
||||
}
|
||||
|
||||
func (r *UserRepository) InviteStats(ctx context.Context, userID string, reward int) (*InviteStats, error) {
|
||||
var inviteCount int64
|
||||
if err := r.db.WithContext(ctx).
|
||||
Model(&model.User{}).
|
||||
Where("invited_by = ?", userID).
|
||||
Count(&inviteCount).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
var rewardedCount int64
|
||||
if err := r.db.WithContext(ctx).
|
||||
Model(&model.User{}).
|
||||
Where("invited_by = ? AND invite_reward_done = ?", userID, true).
|
||||
Count(&rewardedCount).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return &InviteStats{
|
||||
InviteCount: inviteCount,
|
||||
InviteEarned: int(rewardedCount) * reward,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (r *UserRepository) InviteList(ctx context.Context, userID string, reward int) ([]InviteRecord, error) {
|
||||
type row struct {
|
||||
Name string
|
||||
CreatedAt time.Time
|
||||
InviteRewardDone bool
|
||||
InviteRewardAt *time.Time
|
||||
}
|
||||
|
||||
var rows []row
|
||||
if err := r.db.WithContext(ctx).
|
||||
Model(&model.User{}).
|
||||
Select("name, created_at, invite_reward_done, invite_reward_at").
|
||||
Where("invited_by = ?", userID).
|
||||
Order("created_at desc").
|
||||
Find(&rows).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
out := make([]InviteRecord, 0, len(rows))
|
||||
for _, item := range rows {
|
||||
status := "pending"
|
||||
rewardValue := 0
|
||||
if item.InviteRewardDone {
|
||||
status = "completed"
|
||||
rewardValue = reward
|
||||
}
|
||||
name := strings.TrimSpace(item.Name)
|
||||
if name == "" {
|
||||
name = "—"
|
||||
}
|
||||
out = append(out, InviteRecord{
|
||||
Name: name,
|
||||
Reward: rewardValue,
|
||||
RegisteredAt: item.CreatedAt,
|
||||
CompletedAt: item.InviteRewardAt,
|
||||
Status: status,
|
||||
})
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func (r *UserRepository) AllInvites(ctx context.Context, reward int) ([]InviteRecord, *InviteLogStats, error) {
|
||||
type row struct {
|
||||
InviterName string `gorm:"column:inviter_name"`
|
||||
InviterEmail string `gorm:"column:inviter_email"`
|
||||
InviteeName string `gorm:"column:invitee_name"`
|
||||
InviteeEmail string `gorm:"column:invitee_email"`
|
||||
CreatedAt time.Time `gorm:"column:created_at"`
|
||||
InviteRewardDone bool `gorm:"column:invite_reward_done"`
|
||||
InviteRewardAt *time.Time `gorm:"column:invite_reward_at"`
|
||||
}
|
||||
|
||||
var rows []row
|
||||
if err := r.db.WithContext(ctx).
|
||||
Table("users AS invitee").
|
||||
Select(`
|
||||
inviter.name AS inviter_name,
|
||||
inviter.email AS inviter_email,
|
||||
invitee.name AS invitee_name,
|
||||
invitee.email AS invitee_email,
|
||||
invitee.created_at,
|
||||
invitee.invite_reward_done,
|
||||
invitee.invite_reward_at
|
||||
`).
|
||||
Joins("JOIN users AS inviter ON inviter.id = invitee.invited_by").
|
||||
Order("invitee.created_at desc").
|
||||
Scan(&rows).Error; err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
|
||||
out := make([]InviteRecord, 0, len(rows))
|
||||
stats := &InviteLogStats{}
|
||||
for _, item := range rows {
|
||||
stats.Total++
|
||||
status := "pending"
|
||||
rewardValue := 0
|
||||
if item.InviteRewardDone {
|
||||
status = "completed"
|
||||
rewardValue = reward
|
||||
stats.Completed++
|
||||
stats.RewardPaid += int64(reward)
|
||||
} else {
|
||||
stats.Pending++
|
||||
}
|
||||
|
||||
inviter := strings.TrimSpace(item.InviterName)
|
||||
if inviter == "" {
|
||||
inviter = strings.TrimSpace(item.InviterEmail)
|
||||
}
|
||||
invitee := strings.TrimSpace(item.InviteeName)
|
||||
if invitee == "" {
|
||||
invitee = strings.TrimSpace(item.InviteeEmail)
|
||||
}
|
||||
|
||||
out = append(out, InviteRecord{
|
||||
Inviter: inviter,
|
||||
Invitee: invitee,
|
||||
Reward: rewardValue,
|
||||
RegisteredAt: item.CreatedAt,
|
||||
CompletedAt: item.InviteRewardAt,
|
||||
Status: status,
|
||||
})
|
||||
}
|
||||
return out, stats, nil
|
||||
}
|
||||
|
||||
func (r *UserRepository) DailyCheckin(ctx context.Context, userID string, reward int) (*CheckinResult, error) {
|
||||
today := time.Now().Format("2006-01-02")
|
||||
yesterday := time.Now().Add(-24 * time.Hour).Format("2006-01-02")
|
||||
|
||||
var result *CheckinResult
|
||||
err := r.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
|
||||
var user model.User
|
||||
if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).First(&user, "id = ?", userID).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if user.CheckinLast == today {
|
||||
result = &CheckinResult{
|
||||
Already: true,
|
||||
Awarded: 0,
|
||||
Streak: user.CheckinStreak,
|
||||
Credits: user.Credits,
|
||||
}
|
||||
return ErrAlreadyCheckedInToday
|
||||
}
|
||||
|
||||
streak := 1
|
||||
if user.CheckinLast == yesterday {
|
||||
streak = user.CheckinStreak + 1
|
||||
}
|
||||
credits := user.Credits + float64(reward)
|
||||
|
||||
if err := tx.Model(&model.User{}).
|
||||
Where("id = ?", userID).
|
||||
Updates(map[string]any{
|
||||
"credits": credits,
|
||||
"checkin_last": today,
|
||||
"checkin_streak": streak,
|
||||
"updated_at": time.Now(),
|
||||
}).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
result = &CheckinResult{
|
||||
Already: false,
|
||||
Awarded: reward,
|
||||
Streak: streak,
|
||||
Credits: credits,
|
||||
}
|
||||
return nil
|
||||
})
|
||||
if err != nil {
|
||||
if errors.Is(err, ErrAlreadyCheckedInToday) {
|
||||
return result, nil
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func (r *UserRepository) AdjustCredits(ctx context.Context, userID string, delta float64) (*model.User, error) {
|
||||
err := r.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
|
||||
var user model.User
|
||||
if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).First(&user, "id = ?", userID).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
nextCredits := user.Credits + delta
|
||||
if nextCredits < 0 {
|
||||
nextCredits = 0
|
||||
}
|
||||
return tx.Model(&model.User{}).
|
||||
Where("id = ?", userID).
|
||||
Updates(map[string]any{
|
||||
"credits": nextCredits,
|
||||
"updated_at": time.Now(),
|
||||
}).Error
|
||||
})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return r.GetByID(ctx, userID)
|
||||
}
|
||||
|
||||
// SetCredits sets a user's credit balance to an absolute (non-negative) value.
|
||||
// The row is locked for the duration of the transaction so it stays consistent
|
||||
// with concurrent AdjustCredits/TryDebitCredits operations.
|
||||
func (r *UserRepository) SetCredits(ctx context.Context, userID string, value float64) (*model.User, error) {
|
||||
if value < 0 {
|
||||
value = 0
|
||||
}
|
||||
err := r.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
|
||||
var user model.User
|
||||
if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).First(&user, "id = ?", userID).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
return tx.Model(&model.User{}).
|
||||
Where("id = ?", userID).
|
||||
Updates(map[string]any{
|
||||
"credits": value,
|
||||
"updated_at": time.Now(),
|
||||
}).Error
|
||||
})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return r.GetByID(ctx, userID)
|
||||
}
|
||||
|
||||
func (r *UserRepository) TryDebitCredits(ctx context.Context, userID string, amount float64) (*model.User, bool, error) {
|
||||
if amount <= 0 {
|
||||
user, err := r.GetByID(ctx, userID)
|
||||
return user, user != nil, err
|
||||
}
|
||||
|
||||
var result *model.User
|
||||
debited := false
|
||||
err := r.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
|
||||
var user model.User
|
||||
if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).Preload("APIKeys").First(&user, "id = ?", userID).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if user.Credits < amount {
|
||||
result = &user
|
||||
return nil
|
||||
}
|
||||
nextCredits := user.Credits - amount
|
||||
if err := tx.Model(&model.User{}).
|
||||
Where("id = ?", userID).
|
||||
Updates(map[string]any{
|
||||
"credits": nextCredits,
|
||||
"updated_at": time.Now(),
|
||||
}).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
user.Credits = nextCredits
|
||||
user.UpdatedAt = time.Now()
|
||||
result = &user
|
||||
debited = true
|
||||
return nil
|
||||
})
|
||||
if err != nil {
|
||||
return nil, false, err
|
||||
}
|
||||
return result, debited, nil
|
||||
}
|
||||
|
||||
func (r *UserRepository) GrantInviteReward(ctx context.Context, inviteeUserID string, reward int) (bool, error) {
|
||||
if reward <= 0 {
|
||||
return false, nil
|
||||
}
|
||||
|
||||
granted := false
|
||||
err := r.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
|
||||
var invitee model.User
|
||||
if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).First(&invitee, "id = ?", inviteeUserID).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if invitee.InvitedBy == nil || *invitee.InvitedBy == "" || invitee.InviteRewardDone {
|
||||
return nil
|
||||
}
|
||||
|
||||
var inviter model.User
|
||||
if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).First(&inviter, "id = ?", *invitee.InvitedBy).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
now := time.Now()
|
||||
if err := tx.Model(&model.User{}).
|
||||
Where("id = ?", invitee.ID).
|
||||
Updates(map[string]any{
|
||||
"invite_reward_done": true,
|
||||
"invite_reward_at": now,
|
||||
"updated_at": now,
|
||||
}).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
if err := tx.Model(&model.User{}).
|
||||
Where("id = ?", inviter.ID).
|
||||
Updates(map[string]any{
|
||||
"credits": inviter.Credits + float64(reward),
|
||||
"updated_at": now,
|
||||
}).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
granted = true
|
||||
return nil
|
||||
})
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
return granted, nil
|
||||
}
|
||||
@@ -0,0 +1,568 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"sort"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"backend/internal/config"
|
||||
"backend/internal/model"
|
||||
"backend/internal/repo"
|
||||
"backend/internal/storage"
|
||||
)
|
||||
|
||||
type AdminReadService struct {
|
||||
cfg *config.Config
|
||||
users *repo.UserRepository
|
||||
models *repo.ModelRepository
|
||||
events *repo.EventRepository
|
||||
settings *repo.SiteSettingRepository
|
||||
tokens *repo.TokenRepository
|
||||
cdks *repo.CDKRepository
|
||||
store *storage.Client
|
||||
}
|
||||
|
||||
func NewAdminReadService(cfg *config.Config, users *repo.UserRepository, models *repo.ModelRepository, events *repo.EventRepository, settings *repo.SiteSettingRepository, tokens *repo.TokenRepository, cdks *repo.CDKRepository, store *storage.Client) *AdminReadService {
|
||||
return &AdminReadService{
|
||||
cfg: cfg,
|
||||
users: users,
|
||||
models: models,
|
||||
events: events,
|
||||
settings: settings,
|
||||
tokens: tokens,
|
||||
cdks: cdks,
|
||||
store: store,
|
||||
}
|
||||
}
|
||||
|
||||
func (s *AdminReadService) Users(ctx context.Context) ([]model.User, map[string]any, error) {
|
||||
users, err := s.users.List(ctx)
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
counts, err := s.events.UserSuccessCounts(ctx)
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
for i := range users {
|
||||
meta := users[i].Notes
|
||||
_ = meta
|
||||
}
|
||||
stats, err := s.users.Stats(ctx)
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
stats["generation_counts"] = counts
|
||||
return users, stats, nil
|
||||
}
|
||||
|
||||
func (s *AdminReadService) Models(ctx context.Context) ([]model.ModelConfig, error) {
|
||||
return s.models.List(ctx)
|
||||
}
|
||||
|
||||
func (s *AdminReadService) ModelsView(ctx context.Context) ([]map[string]any, error) {
|
||||
items, err := s.models.List(ctx)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
counts, err := s.events.ModelSuccessCounts(ctx)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out := make([]map[string]any, 0, len(items))
|
||||
for _, item := range items {
|
||||
out = append(out, map[string]any{
|
||||
"id": item.ID,
|
||||
"type": item.Type,
|
||||
"name": item.Name,
|
||||
"provider": item.Provider,
|
||||
"enabled": item.Enabled,
|
||||
"ratios": repo.JSONStrings(item.Ratios),
|
||||
"prices": map[string]any(item.Prices),
|
||||
"resolutions": repo.JSONStrings(item.Resolutions),
|
||||
"image_to_image": item.ImageToImage,
|
||||
"duration_prices": map[string]any(item.DurationPrices),
|
||||
"prices_agent": map[string]any(item.PricesAgent),
|
||||
"duration_prices_agent": map[string]any(item.DurationPricesAgent),
|
||||
"durations": repo.JSONStrings(item.Durations),
|
||||
"max_reference_images": item.MaxReferenceImages,
|
||||
"reference_mode": item.ReferenceMode,
|
||||
"weight": item.Weight,
|
||||
"generation_count": counts[item.ID],
|
||||
"created_at": item.CreatedAt,
|
||||
"updated_at": item.UpdatedAt,
|
||||
})
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func (s *AdminReadService) Logs(ctx context.Context, limit, offset int, kind, status string, since *time.Time, userID, excludeSource, source string, hasFile bool) ([]model.EventLog, int64, *repo.EventStats, error) {
|
||||
items, total, err := s.events.List(ctx, repo.EventListFilter{
|
||||
Limit: limit,
|
||||
Offset: offset,
|
||||
Kind: kind,
|
||||
Status: status,
|
||||
Since: since,
|
||||
UserID: userID,
|
||||
ExcludeSource: excludeSource,
|
||||
Source: source,
|
||||
HasFile: hasFile,
|
||||
})
|
||||
if err != nil {
|
||||
return nil, 0, nil, err
|
||||
}
|
||||
// 用户自己的日志(userID 非空)→ 按本人统计;管理员全站视图 → 全站统计。
|
||||
var stats *repo.EventStats
|
||||
if userID != "" {
|
||||
stats, err = s.events.StatsByUser(ctx, userID)
|
||||
} else {
|
||||
stats, err = s.events.Stats(ctx)
|
||||
}
|
||||
if err != nil {
|
||||
return nil, 0, nil, err
|
||||
}
|
||||
return items, total, stats, nil
|
||||
}
|
||||
|
||||
// UserNameMap builds an id -> display name lookup (name, else email, else id)
|
||||
// used to annotate admin log rows with user_name (mirrors admin.py:584-596).
|
||||
func (s *AdminReadService) UserNameMap(ctx context.Context) (map[string]string, error) {
|
||||
users, err := s.users.List(ctx)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out := make(map[string]string, len(users))
|
||||
for _, u := range users {
|
||||
name := strings.TrimSpace(u.Name)
|
||||
if name == "" {
|
||||
name = strings.TrimSpace(u.Email)
|
||||
}
|
||||
if name == "" {
|
||||
name = u.ID
|
||||
}
|
||||
out[u.ID] = name
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func (s *AdminReadService) Stats(ctx context.Context) (map[string]any, error) {
|
||||
stats, err := s.events.Stats(ctx)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
recentFiles, _ := s.RecentImages(ctx, 24)
|
||||
files, fileStats, _ := s.scanGeneratedFiles(ctx)
|
||||
var size int64
|
||||
if v, ok := fileStats["size_bytes"].(int64); ok {
|
||||
size = v
|
||||
}
|
||||
return map[string]any{
|
||||
"generated_count": len(files),
|
||||
"generated_size_bytes": size,
|
||||
"recent": recentFiles,
|
||||
"avg_elapsed_ms": stats.AvgElapsedMS,
|
||||
"avg_elapsed_ms_24h": stats.AvgElapsedMS24,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// Dashboard assembles the admin overview's analytics entirely server-side
|
||||
// (event windows, hourly trend, top models/failures/spenders) plus CDK / invite
|
||||
// / checkin summaries. This replaces the old client-side math over the last 200
|
||||
// logs, which silently undercounted week/DAU/trend once volume grew.
|
||||
func (s *AdminReadService) Dashboard(ctx context.Context) (map[string]any, error) {
|
||||
now := time.Now()
|
||||
dayCut := now.Add(-24 * time.Hour)
|
||||
weekCut := now.Add(-7 * 24 * time.Hour)
|
||||
|
||||
day, err := s.events.WindowStats(ctx, dayCut)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
week, err := s.events.WindowStats(ctx, weekCut)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
prevDay, err := s.events.CountBetween(ctx, now.Add(-48*time.Hour), dayCut)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
dau, err := s.events.DistinctUsersSince(ctx, dayCut)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
wau, err := s.events.DistinctUsersSince(ctx, weekCut)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
hourly, err := s.events.HourlyBuckets(ctx)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// Per-window top-N analytics so the frontend can toggle 24h / 7d without a
|
||||
// re-fetch (the lists are small — top 6 / top 5).
|
||||
nameByID, err := s.UserNameMap(ctx)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
analytics := func(since time.Time) (map[string]any, error) {
|
||||
models, err := s.events.ModelUsageSince(ctx, since, 6)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
failures, err := s.events.TopFailures(ctx, since, 5)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
users, err := s.events.TopUserSpend(ctx, since, 6)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
for i := range users {
|
||||
if users[i].UserID == "" {
|
||||
users[i].Name = "匿名"
|
||||
} else if name, ok := nameByID[users[i].UserID]; ok {
|
||||
users[i].Name = name
|
||||
} else {
|
||||
users[i].Name = users[i].UserID
|
||||
}
|
||||
}
|
||||
return map[string]any{"models": models, "failures": failures, "top_users": users}, nil
|
||||
}
|
||||
dayAnalytics, err := analytics(dayCut)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
weekAnalytics, err := analytics(weekCut)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
cdkStats, err := s.cdks.Stats(ctx)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
inviteReward := parseIntSetting(s.mustSetting(ctx, "credits.invite_reward"), 3)
|
||||
inviteSummary, err := s.users.InviteSummary(ctx)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
checkinReward := parseIntSetting(s.mustSetting(ctx, "credits.checkin_reward"), 3)
|
||||
checkin, err := s.users.CheckinStats(ctx)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return map[string]any{
|
||||
"day": day,
|
||||
"week": week,
|
||||
"prev_day_total": prevDay,
|
||||
"dau": dau,
|
||||
"wau": wau,
|
||||
"hourly": hourly,
|
||||
"analytics": map[string]any{
|
||||
"day": dayAnalytics,
|
||||
"week": weekAnalytics,
|
||||
},
|
||||
"cdk": cdkStats,
|
||||
"invites": map[string]any{
|
||||
"total": inviteSummary.Total,
|
||||
"completed": inviteSummary.Completed,
|
||||
"reward": inviteReward,
|
||||
"reward_paid": inviteSummary.Completed * int64(inviteReward),
|
||||
},
|
||||
"checkin": map[string]any{
|
||||
"today": checkin.TodayCount,
|
||||
"max_streak": checkin.MaxStreak,
|
||||
"reward": checkinReward,
|
||||
"awarded_today": checkin.TodayCount * int64(checkinReward),
|
||||
},
|
||||
}, nil
|
||||
}
|
||||
|
||||
// mustSetting reads a site setting value, returning "" on error so the caller's
|
||||
// parseIntSetting default kicks in (a missing reward setting shouldn't 500 the
|
||||
// whole dashboard).
|
||||
func (s *AdminReadService) mustSetting(ctx context.Context, key string) string {
|
||||
v, err := s.settings.GetValue(ctx, key)
|
||||
if err != nil {
|
||||
return ""
|
||||
}
|
||||
return v
|
||||
}
|
||||
|
||||
func (s *AdminReadService) Invites(ctx context.Context) ([]repo.InviteRecord, *repo.InviteLogStats, error) {
|
||||
rewardRaw, err := s.settings.GetValue(ctx, "credits.invite_reward")
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
reward := parseIntSetting(rewardRaw, 3)
|
||||
return s.users.AllInvites(ctx, reward)
|
||||
}
|
||||
|
||||
func (s *AdminReadService) Providers(ctx context.Context) ([]map[string]any, error) {
|
||||
models, err := s.models.List(ctx)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
tokens, err := s.tokens.List(ctx)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
modelCounts := map[string]int{}
|
||||
for _, item := range models {
|
||||
modelCounts[item.Provider]++
|
||||
}
|
||||
type aggregate struct {
|
||||
active int
|
||||
disabled int
|
||||
quota int
|
||||
}
|
||||
tokenCounts := map[string]*aggregate{}
|
||||
for _, item := range tokens {
|
||||
if _, ok := tokenCounts[item.Pool]; !ok {
|
||||
tokenCounts[item.Pool] = &aggregate{}
|
||||
}
|
||||
switch item.Status {
|
||||
case "active":
|
||||
tokenCounts[item.Pool].active++
|
||||
case "quota":
|
||||
tokenCounts[item.Pool].quota++
|
||||
default:
|
||||
tokenCounts[item.Pool].disabled++
|
||||
}
|
||||
}
|
||||
providers := []struct {
|
||||
Name string
|
||||
Pool string
|
||||
Type string
|
||||
}{
|
||||
{Name: "chatgpt", Pool: "chatgpt", Type: "openai"},
|
||||
{Name: "adobe", Pool: "adobe", Type: "adobe"},
|
||||
}
|
||||
out := make([]map[string]any, 0, len(providers))
|
||||
for _, item := range providers {
|
||||
count := tokenCounts[item.Pool]
|
||||
if count == nil {
|
||||
count = &aggregate{}
|
||||
}
|
||||
out = append(out, map[string]any{
|
||||
"name": item.Name,
|
||||
"token_pool": item.Pool,
|
||||
"type": item.Type,
|
||||
"model_count": modelCounts[item.Name],
|
||||
"tokens_total": count.active + count.disabled + count.quota,
|
||||
"tokens_active": count.active,
|
||||
"tokens_disabled": count.disabled,
|
||||
"tokens_quota": count.quota,
|
||||
})
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func (s *AdminReadService) Images(ctx context.Context, limit, offset int, kind string) ([]map[string]any, int, map[string]any, error) {
|
||||
if limit <= 0 {
|
||||
limit = 30
|
||||
}
|
||||
if limit > 200 {
|
||||
limit = 200
|
||||
}
|
||||
if offset < 0 {
|
||||
offset = 0
|
||||
}
|
||||
allFiles, stats, err := s.scanGeneratedFiles(ctx)
|
||||
if err != nil {
|
||||
return nil, 0, nil, err
|
||||
}
|
||||
filtered := make([]generatedFile, 0, len(allFiles))
|
||||
for _, item := range allFiles {
|
||||
if kind == "" || item.Kind == kind {
|
||||
filtered = append(filtered, item)
|
||||
}
|
||||
}
|
||||
sort.SliceStable(filtered, func(i, j int) bool {
|
||||
return filtered[i].MTime > filtered[j].MTime
|
||||
})
|
||||
total := len(filtered)
|
||||
if offset > total {
|
||||
offset = total
|
||||
}
|
||||
end := offset + limit
|
||||
if end > total {
|
||||
end = total
|
||||
}
|
||||
page := filtered[offset:end]
|
||||
index, err := s.eventIndexByFile(ctx)
|
||||
if err != nil {
|
||||
return nil, 0, nil, err
|
||||
}
|
||||
out := make([]map[string]any, 0, len(page))
|
||||
for _, item := range page {
|
||||
row := map[string]any{
|
||||
"name": item.Name,
|
||||
"size": item.Size,
|
||||
"mtime": item.MTime,
|
||||
"kind": item.Kind,
|
||||
"prompt": "",
|
||||
"model": "",
|
||||
"resolution": "",
|
||||
"ratio": "",
|
||||
"duration": "",
|
||||
}
|
||||
if event, ok := index[item.Name]; ok {
|
||||
row["prompt"] = event.Prompt
|
||||
row["model"] = event.Model
|
||||
row["resolution"] = event.Resolution
|
||||
row["ratio"] = event.Ratio
|
||||
row["duration"] = event.Duration
|
||||
}
|
||||
out = append(out, row)
|
||||
}
|
||||
return out, total, stats, nil
|
||||
}
|
||||
|
||||
func (s *AdminReadService) RecentImages(ctx context.Context, limit int) ([]map[string]any, error) {
|
||||
if limit <= 0 {
|
||||
limit = 24
|
||||
}
|
||||
allFiles, _, err := s.scanGeneratedFiles(ctx)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
sort.SliceStable(allFiles, func(i, j int) bool {
|
||||
return allFiles[i].MTime > allFiles[j].MTime
|
||||
})
|
||||
if len(allFiles) > limit {
|
||||
allFiles = allFiles[:limit]
|
||||
}
|
||||
out := make([]map[string]any, 0, len(allFiles))
|
||||
for _, item := range allFiles {
|
||||
out = append(out, map[string]any{
|
||||
"name": item.Name,
|
||||
"size": item.Size,
|
||||
"mtime": item.MTime,
|
||||
"kind": item.Kind,
|
||||
})
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// RecentImagesOwned lists the most-recent generated images under a single owner
|
||||
// directory (used by the showcase picker so an admin sees only their OWN images).
|
||||
func (s *AdminReadService) RecentImagesOwned(ctx context.Context, owner string, limit int) ([]map[string]any, error) {
|
||||
if limit <= 0 {
|
||||
limit = 24
|
||||
}
|
||||
owner = strings.TrimSpace(owner)
|
||||
if owner == "" || s.store == nil || !s.store.Configured() {
|
||||
return []map[string]any{}, nil
|
||||
}
|
||||
objs, err := s.store.List(ctx, owner+"/")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
files := make([]generatedFile, 0, len(objs))
|
||||
for _, o := range objs {
|
||||
if isReferenceFile(o.Key) {
|
||||
continue
|
||||
}
|
||||
kind := mediaKind(o.Key)
|
||||
if kind == "" {
|
||||
continue
|
||||
}
|
||||
files = append(files, generatedFile{Name: o.Key, Size: o.Size, MTime: o.LastModified.Unix(), Kind: kind})
|
||||
}
|
||||
sort.SliceStable(files, func(i, j int) bool { return files[i].MTime > files[j].MTime })
|
||||
if len(files) > limit {
|
||||
files = files[:limit]
|
||||
}
|
||||
out := make([]map[string]any, 0, len(files))
|
||||
for _, f := range files {
|
||||
out = append(out, map[string]any{"name": f.Name, "size": f.Size, "mtime": f.MTime, "kind": f.Kind})
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func (s *AdminReadService) eventIndexByFile(ctx context.Context) (map[string]model.EventLog, error) {
|
||||
items, err := s.events.RecentByFile(ctx, 10000)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out := make(map[string]model.EventLog, len(items))
|
||||
for _, item := range items {
|
||||
if item.File == "" {
|
||||
continue
|
||||
}
|
||||
if _, ok := out[item.File]; ok {
|
||||
continue
|
||||
}
|
||||
out[item.File] = item
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
type generatedFile struct {
|
||||
Name string
|
||||
Size int64
|
||||
MTime int64
|
||||
Kind string
|
||||
}
|
||||
|
||||
// mediaKind classifies an object key by extension (image / video / "" = skip).
|
||||
func mediaKind(name string) string {
|
||||
i := strings.LastIndex(name, ".")
|
||||
if i < 0 {
|
||||
return ""
|
||||
}
|
||||
switch strings.ToLower(name[i+1:]) {
|
||||
case "png", "jpg", "jpeg", "webp", "gif":
|
||||
return "image"
|
||||
case "mp4", "webm", "mov":
|
||||
return "video"
|
||||
default:
|
||||
return ""
|
||||
}
|
||||
}
|
||||
|
||||
// isReferenceFile reports whether a key is an uploaded reference image (named
|
||||
// "...-ref-..."), so the gallery / picker can skip them — only generated outputs
|
||||
// are listed.
|
||||
func isReferenceFile(name string) bool {
|
||||
return strings.Contains(name, "-ref-")
|
||||
}
|
||||
|
||||
// scanGeneratedFiles lists media objects from RustFS (replacing the old local
|
||||
// directory walk). Keys ARE the relative paths the rest of the app expects.
|
||||
func (s *AdminReadService) scanGeneratedFiles(ctx context.Context) ([]generatedFile, map[string]any, error) {
|
||||
stats := map[string]any{"total": 0, "image": 0, "video": 0, "size_bytes": int64(0)}
|
||||
if s.store == nil || !s.store.Configured() {
|
||||
return nil, stats, nil
|
||||
}
|
||||
objs, err := s.store.List(ctx, "")
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
out := make([]generatedFile, 0, len(objs))
|
||||
for _, o := range objs {
|
||||
if isReferenceFile(o.Key) {
|
||||
continue // reference uploads are not generated outputs — hide from gallery
|
||||
}
|
||||
kind := mediaKind(o.Key)
|
||||
if kind == "" {
|
||||
continue
|
||||
}
|
||||
stats[kind] = stats[kind].(int) + 1
|
||||
stats["total"] = stats["total"].(int) + 1
|
||||
stats["size_bytes"] = stats["size_bytes"].(int64) + o.Size
|
||||
out = append(out, generatedFile{
|
||||
Name: o.Key,
|
||||
Size: o.Size,
|
||||
MTime: o.LastModified.Unix(),
|
||||
Kind: kind,
|
||||
})
|
||||
}
|
||||
return out, stats, nil
|
||||
}
|
||||
@@ -0,0 +1,635 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"backend/internal/model"
|
||||
"backend/internal/repo"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"gorm.io/datatypes"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
// ErrNotFound is returned by delete/adjust service methods when the target
|
||||
// row does not exist, so handlers can translate it into a 404 (GORM's Delete
|
||||
// does not error on a zero-row delete).
|
||||
var ErrNotFound = errors.New("not found")
|
||||
|
||||
type AdminWriteService struct {
|
||||
users *repo.UserRepository
|
||||
showcase *repo.ShowcaseRepository
|
||||
models *repo.ModelRepository
|
||||
events *repo.EventRepository
|
||||
apiKeys *repo.APIKeyRepository
|
||||
}
|
||||
|
||||
func NewAdminWriteService(users *repo.UserRepository, showcase *repo.ShowcaseRepository, models *repo.ModelRepository, events *repo.EventRepository, apiKeys *repo.APIKeyRepository) *AdminWriteService {
|
||||
return &AdminWriteService{
|
||||
users: users,
|
||||
showcase: showcase,
|
||||
models: models,
|
||||
events: events,
|
||||
apiKeys: apiKeys,
|
||||
}
|
||||
}
|
||||
|
||||
func (s *AdminWriteService) CreateUser(ctx context.Context, body map[string]any) (*model.User, error) {
|
||||
email, err := ValidateEmail(stringValue(body["email"]))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
name := strings.TrimSpace(stringValue(body["name"]))
|
||||
if name != "" {
|
||||
name, err = ValidateUsername(name)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
password := stringValue(body["password"])
|
||||
role := normalizedRole(stringValue(body["role"]))
|
||||
// 管理员唯一:不能通过用户管理创建新的 admin(只能是 user / agent)。
|
||||
if role == "admin" {
|
||||
role = "user"
|
||||
}
|
||||
status := normalizedStatus(stringValue(body["status"]))
|
||||
credits := maxFloat(0, floatValue(body["credits"]))
|
||||
notes := strings.TrimSpace(stringValue(body["notes"]))
|
||||
|
||||
exists, err := s.users.ExistsEmail(ctx, email, "")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if exists {
|
||||
return nil, errors.New("邮箱已存在")
|
||||
}
|
||||
if name != "" {
|
||||
exists, err = s.users.ExistsName(ctx, name, "")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if exists {
|
||||
return nil, errors.New("用户名已存在")
|
||||
}
|
||||
}
|
||||
|
||||
passwordHash := ""
|
||||
if strings.TrimSpace(password) != "" {
|
||||
if err := ValidatePassword(password); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
h, err := HashPassword(password)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
passwordHash = h
|
||||
}
|
||||
|
||||
user := &model.User{
|
||||
ID: "u-" + uuid.NewString()[:10],
|
||||
Email: email,
|
||||
Name: name,
|
||||
PasswordHash: passwordHash,
|
||||
Role: role,
|
||||
Status: status,
|
||||
Credits: credits,
|
||||
Notes: notes,
|
||||
InviteCode: randomInviteCode(),
|
||||
CreatedAt: time.Now(),
|
||||
UpdatedAt: time.Now(),
|
||||
}
|
||||
if err := s.users.Create(ctx, user); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return s.users.GetByID(ctx, user.ID)
|
||||
}
|
||||
|
||||
func (s *AdminWriteService) UpdateUser(ctx context.Context, userID string, body map[string]any) (*model.User, error) {
|
||||
patch := map[string]any{}
|
||||
if _, ok := body["email"]; ok {
|
||||
email, err := ValidateEmail(stringValue(body["email"]))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
exists, err := s.users.ExistsEmail(ctx, email, userID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if exists {
|
||||
return nil, errors.New("邮箱已存在")
|
||||
}
|
||||
patch["email"] = email
|
||||
}
|
||||
if _, ok := body["name"]; ok {
|
||||
name := strings.TrimSpace(stringValue(body["name"]))
|
||||
if name != "" {
|
||||
var err error
|
||||
name, err = ValidateUsername(name)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
exists, err := s.users.ExistsName(ctx, name, userID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if exists {
|
||||
return nil, errors.New("用户名已存在")
|
||||
}
|
||||
}
|
||||
patch["name"] = name
|
||||
}
|
||||
if _, ok := body["role"]; ok {
|
||||
newRole := normalizedRole(stringValue(body["role"]))
|
||||
// 管理员唯一:不能把任何人提升为 admin;也绝不改动现有 admin 的角色
|
||||
// (防止把唯一管理员误降级导致后台失去管理员)。
|
||||
cur, _ := s.users.GetByID(ctx, userID)
|
||||
if newRole != "admin" && (cur == nil || cur.Role != "admin") {
|
||||
patch["role"] = newRole
|
||||
}
|
||||
}
|
||||
if _, ok := body["status"]; ok {
|
||||
patch["status"] = normalizedStatus(stringValue(body["status"]))
|
||||
}
|
||||
if _, ok := body["credits"]; ok {
|
||||
patch["credits"] = maxFloat(0, floatValue(body["credits"]))
|
||||
}
|
||||
if _, ok := body["notes"]; ok {
|
||||
patch["notes"] = strings.TrimSpace(stringValue(body["notes"]))
|
||||
}
|
||||
if _, ok := body["password"]; ok && strings.TrimSpace(stringValue(body["password"])) != "" {
|
||||
if err := ValidatePassword(stringValue(body["password"])); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
h, err := HashPassword(stringValue(body["password"]))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
patch["password_hash"] = h
|
||||
}
|
||||
return s.users.Update(ctx, userID, patch)
|
||||
}
|
||||
|
||||
func (s *AdminWriteService) DeleteUser(ctx context.Context, userID string) error {
|
||||
rows, err := s.users.Delete(ctx, userID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if rows == 0 {
|
||||
return ErrNotFound
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// DeleteUsers removes many users in one call (multi-select). Returns the count
|
||||
// removed.
|
||||
func (s *AdminWriteService) DeleteUsers(ctx context.Context, ids []string) (int, error) {
|
||||
seen := make(map[string]struct{}, len(ids))
|
||||
clean := make([]string, 0, len(ids))
|
||||
for _, id := range ids {
|
||||
id = strings.TrimSpace(id)
|
||||
if id == "" {
|
||||
continue
|
||||
}
|
||||
if _, ok := seen[id]; ok {
|
||||
continue
|
||||
}
|
||||
seen[id] = struct{}{}
|
||||
clean = append(clean, id)
|
||||
}
|
||||
if len(clean) == 0 {
|
||||
return 0, nil
|
||||
}
|
||||
rows, err := s.users.DeleteByIDs(ctx, clean)
|
||||
return int(rows), err
|
||||
}
|
||||
|
||||
func (s *AdminWriteService) AdjustUserCredits(ctx context.Context, userID string, delta float64) (*model.User, error) {
|
||||
user, err := s.users.AdjustCredits(ctx, userID, delta)
|
||||
if err != nil {
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
return nil, ErrNotFound
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
return user, nil
|
||||
}
|
||||
|
||||
// SetUserCredits sets a user's credit balance to an absolute value (non-negative).
|
||||
// Mirrors Python users_store.adjust_credits set_to mode; the update runs inside a
|
||||
// transaction with a row lock so concurrent adjustments stay consistent.
|
||||
func (s *AdminWriteService) SetUserCredits(ctx context.Context, userID string, value float64) (*model.User, error) {
|
||||
if value < 0 {
|
||||
value = 0
|
||||
}
|
||||
user, err := s.users.SetCredits(ctx, userID, value)
|
||||
if err != nil {
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
return nil, ErrNotFound
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
return user, nil
|
||||
}
|
||||
|
||||
func (s *AdminWriteService) CreateUserAPIKey(ctx context.Context, userID, name string) (*model.APIKey, string, error) {
|
||||
plain, err := generatePlainAPIKey()
|
||||
if err != nil {
|
||||
return nil, "", err
|
||||
}
|
||||
name = strings.TrimSpace(name)
|
||||
if name == "" {
|
||||
name = "admin"
|
||||
}
|
||||
key := &model.APIKey{
|
||||
ID: "k-" + time.Now().Format("150405") + randomSuffix(2),
|
||||
UserID: userID,
|
||||
Name: name,
|
||||
KeyPreview: previewAPIKey(plain),
|
||||
KeyHash: hashAPIKey(plain),
|
||||
CreatedAt: time.Now(),
|
||||
}
|
||||
if err := s.apiKeys.Create(ctx, key); err != nil {
|
||||
return nil, "", err
|
||||
}
|
||||
return key, plain, nil
|
||||
}
|
||||
|
||||
func (s *AdminWriteService) DeleteUserAPIKey(ctx context.Context, userID, keyID string) error {
|
||||
if strings.TrimSpace(keyID) == "" {
|
||||
return errors.New("key id required")
|
||||
}
|
||||
return s.apiKeys.DeleteByID(ctx, userID, keyID)
|
||||
}
|
||||
|
||||
func (s *AdminWriteService) CreateShowcase(ctx context.Context, body map[string]any) (*model.ShowcaseItem, error) {
|
||||
kind := normalizedShowcaseKind(stringValue(body["kind"]))
|
||||
if kind == "" {
|
||||
return nil, errors.New("kind must be hero, bento or work")
|
||||
}
|
||||
image := strings.TrimSpace(stringValue(body["image"]))
|
||||
if image == "" {
|
||||
return nil, errors.New("请选择底图")
|
||||
}
|
||||
title := strings.TrimSpace(stringValue(body["title"]))
|
||||
prompt := strings.TrimSpace(stringValue(body["prompt"]))
|
||||
if kind != "work" {
|
||||
if title == "" {
|
||||
return nil, errors.New("请填写标题")
|
||||
}
|
||||
if prompt == "" {
|
||||
return nil, errors.New("请填写提示词")
|
||||
}
|
||||
}
|
||||
|
||||
item := &model.ShowcaseItem{
|
||||
ID: "sc-" + uuid.NewString()[:10],
|
||||
Kind: kind,
|
||||
Title: title,
|
||||
Subtitle: strings.TrimSpace(stringValue(body["subtitle"])),
|
||||
Prompt: prompt,
|
||||
Gradient: strings.TrimSpace(stringValue(body["gradient"])),
|
||||
Span: strings.TrimSpace(stringValue(body["span"])),
|
||||
Image: image,
|
||||
Weight: intValue(body["weight"]),
|
||||
CreatedAt: time.Now(),
|
||||
UpdatedAt: time.Now(),
|
||||
}
|
||||
if err := s.showcase.Create(ctx, item); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return item, nil
|
||||
}
|
||||
|
||||
func (s *AdminWriteService) UpdateShowcase(ctx context.Context, entryID string, body map[string]any) (*model.ShowcaseItem, error) {
|
||||
patch := map[string]any{}
|
||||
if _, ok := body["kind"]; ok {
|
||||
kind := normalizedShowcaseKind(stringValue(body["kind"]))
|
||||
if kind == "" {
|
||||
return nil, errors.New("kind must be hero, bento or work")
|
||||
}
|
||||
patch["kind"] = kind
|
||||
}
|
||||
for _, field := range []string{"title", "subtitle", "prompt", "gradient", "span", "image"} {
|
||||
if _, ok := body[field]; ok {
|
||||
patch[field] = strings.TrimSpace(stringValue(body[field]))
|
||||
}
|
||||
}
|
||||
if _, ok := body["weight"]; ok {
|
||||
patch["weight"] = intValue(body["weight"])
|
||||
}
|
||||
return s.showcase.Update(ctx, entryID, patch)
|
||||
}
|
||||
|
||||
func (s *AdminWriteService) DeleteShowcase(ctx context.Context, entryID string) error {
|
||||
rows, err := s.showcase.Delete(ctx, entryID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if rows == 0 {
|
||||
return ErrNotFound
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *AdminWriteService) CreateModel(ctx context.Context, body map[string]any) (*model.ModelConfig, error) {
|
||||
modelID := strings.TrimSpace(stringValue(body["id"]))
|
||||
modelType := normalizedModelType(stringValue(body["type"]))
|
||||
provider := strings.TrimSpace(stringValue(body["provider"]))
|
||||
if modelID == "" {
|
||||
return nil, errors.New("id required")
|
||||
}
|
||||
if modelType == "" {
|
||||
return nil, errors.New("type must be image or video")
|
||||
}
|
||||
if provider == "" {
|
||||
return nil, errors.New("provider required")
|
||||
}
|
||||
|
||||
prices := jsonMap(body["prices"])
|
||||
// image: tiers derive from the price keys (form omits resolutions);
|
||||
// video: resolutions come straight from the form (720p/1080p…). Python parity.
|
||||
resolutions := jsonArray(body["resolutions"])
|
||||
if modelType != "video" {
|
||||
resolutions = resolutionsFromPrices(prices)
|
||||
}
|
||||
|
||||
item := &model.ModelConfig{
|
||||
ID: modelID,
|
||||
Type: modelType,
|
||||
Name: defaultString(strings.TrimSpace(stringValue(body["name"])), modelID),
|
||||
Provider: provider,
|
||||
Enabled: boolValueWithDefault(body["enabled"], true),
|
||||
Ratios: jsonArray(body["ratios"]),
|
||||
Prices: prices,
|
||||
Resolutions: resolutions,
|
||||
ImageToImage: boolValueWithDefault(body["image_to_image"], false),
|
||||
DurationPrices: jsonMap(body["duration_prices"]),
|
||||
PricesAgent: jsonMap(body["prices_agent"]),
|
||||
DurationPricesAgent: jsonMap(body["duration_prices_agent"]),
|
||||
Durations: jsonArray(body["durations"]),
|
||||
MaxReferenceImages: intValue(body["max_reference_images"]),
|
||||
ReferenceMode: defaultString(strings.TrimSpace(stringValue(body["reference_mode"])), "none"),
|
||||
Weight: intValue(body["weight"]),
|
||||
CreatedAt: time.Now(),
|
||||
UpdatedAt: time.Now(),
|
||||
}
|
||||
if err := s.models.Create(ctx, item); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return item, nil
|
||||
}
|
||||
|
||||
func (s *AdminWriteService) UpdateModel(ctx context.Context, modelID string, body map[string]any) (*model.ModelConfig, error) {
|
||||
patch := map[string]any{}
|
||||
if _, ok := body["type"]; ok {
|
||||
modelType := normalizedModelType(stringValue(body["type"]))
|
||||
if modelType == "" {
|
||||
return nil, errors.New("type must be image or video")
|
||||
}
|
||||
patch["type"] = modelType
|
||||
}
|
||||
if _, ok := body["name"]; ok {
|
||||
patch["name"] = strings.TrimSpace(stringValue(body["name"]))
|
||||
}
|
||||
if _, ok := body["provider"]; ok {
|
||||
provider := strings.TrimSpace(stringValue(body["provider"]))
|
||||
if provider == "" {
|
||||
return nil, errors.New("provider required")
|
||||
}
|
||||
patch["provider"] = provider
|
||||
}
|
||||
// Only touch `enabled` when the caller explicitly sends a non-null value;
|
||||
// mirrors Python models_store.update ("enabled" in fields and is not None).
|
||||
// Without this guard a PATCH that omits the field would default it to false
|
||||
// and silently disable the model.
|
||||
if raw, ok := body["enabled"]; ok && raw != nil {
|
||||
patch["enabled"] = boolValueWithDefault(raw, true)
|
||||
}
|
||||
if _, ok := body["ratios"]; ok {
|
||||
patch["ratios"] = jsonArray(body["ratios"])
|
||||
}
|
||||
if _, ok := body["prices"]; ok {
|
||||
prices := jsonMap(body["prices"])
|
||||
patch["prices"] = prices
|
||||
// Python parity (models_store.update): recompute resolutions from the new
|
||||
// price keys. An explicit `resolutions` field below (video) overrides this.
|
||||
patch["resolutions"] = resolutionsFromPrices(prices)
|
||||
}
|
||||
if _, ok := body["resolutions"]; ok {
|
||||
patch["resolutions"] = jsonArray(body["resolutions"])
|
||||
}
|
||||
if _, ok := body["image_to_image"]; ok {
|
||||
patch["image_to_image"] = boolValueWithDefault(body["image_to_image"], false)
|
||||
}
|
||||
if _, ok := body["duration_prices"]; ok {
|
||||
patch["duration_prices"] = jsonMap(body["duration_prices"])
|
||||
}
|
||||
if _, ok := body["prices_agent"]; ok {
|
||||
patch["prices_agent"] = jsonMap(body["prices_agent"])
|
||||
}
|
||||
if _, ok := body["duration_prices_agent"]; ok {
|
||||
patch["duration_prices_agent"] = jsonMap(body["duration_prices_agent"])
|
||||
}
|
||||
if _, ok := body["durations"]; ok {
|
||||
patch["durations"] = jsonArray(body["durations"])
|
||||
}
|
||||
if _, ok := body["max_reference_images"]; ok {
|
||||
patch["max_reference_images"] = intValue(body["max_reference_images"])
|
||||
}
|
||||
if _, ok := body["reference_mode"]; ok {
|
||||
patch["reference_mode"] = defaultString(strings.TrimSpace(stringValue(body["reference_mode"])), "none")
|
||||
}
|
||||
if _, ok := body["weight"]; ok {
|
||||
patch["weight"] = intValue(body["weight"])
|
||||
}
|
||||
return s.models.Update(ctx, modelID, patch)
|
||||
}
|
||||
|
||||
func (s *AdminWriteService) DeleteModel(ctx context.Context, modelID string) error {
|
||||
rows, err := s.models.Delete(ctx, modelID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if rows == 0 {
|
||||
return ErrNotFound
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *AdminWriteService) ClearLogs(ctx context.Context) (int64, error) {
|
||||
return s.events.DeleteAll(ctx)
|
||||
}
|
||||
|
||||
func (s *AdminWriteService) ClearPendingLogs(ctx context.Context) (int64, error) {
|
||||
return s.events.DeletePending(ctx)
|
||||
}
|
||||
|
||||
func HashPassword(password string) (string, error) {
|
||||
hash, err := GeneratePasswordHash(password)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
return "bcrypt$" + hash, nil
|
||||
}
|
||||
|
||||
func normalizedRole(role string) string {
|
||||
switch strings.TrimSpace(role) {
|
||||
case "admin":
|
||||
return "admin"
|
||||
case "agent":
|
||||
return "agent"
|
||||
default:
|
||||
return "user"
|
||||
}
|
||||
}
|
||||
|
||||
func normalizedStatus(status string) string {
|
||||
if strings.TrimSpace(status) == "disabled" {
|
||||
return "disabled"
|
||||
}
|
||||
return "active"
|
||||
}
|
||||
|
||||
func normalizedShowcaseKind(kind string) string {
|
||||
switch strings.TrimSpace(kind) {
|
||||
case "hero", "bento", "work":
|
||||
return strings.TrimSpace(kind)
|
||||
default:
|
||||
return ""
|
||||
}
|
||||
}
|
||||
|
||||
func normalizedModelType(v string) string {
|
||||
switch strings.TrimSpace(v) {
|
||||
case "image", "video":
|
||||
return strings.TrimSpace(v)
|
||||
default:
|
||||
return ""
|
||||
}
|
||||
}
|
||||
|
||||
func stringValue(v any) string {
|
||||
if v == nil {
|
||||
return ""
|
||||
}
|
||||
switch x := v.(type) {
|
||||
case string:
|
||||
return x
|
||||
default:
|
||||
return fmt.Sprint(v)
|
||||
}
|
||||
}
|
||||
|
||||
func floatValue(v any) float64 {
|
||||
switch x := v.(type) {
|
||||
case float64:
|
||||
return x
|
||||
case float32:
|
||||
return float64(x)
|
||||
case int:
|
||||
return float64(x)
|
||||
case int64:
|
||||
return float64(x)
|
||||
case json.Number:
|
||||
f, _ := x.Float64()
|
||||
return f
|
||||
case string:
|
||||
var f float64
|
||||
_, _ = fmt.Sscanf(strings.TrimSpace(x), "%f", &f)
|
||||
return f
|
||||
default:
|
||||
return 0
|
||||
}
|
||||
}
|
||||
|
||||
func intValue(v any) int {
|
||||
return int(floatValue(v))
|
||||
}
|
||||
|
||||
func boolValueWithDefault(v any, fallback bool) bool {
|
||||
if v == nil {
|
||||
return fallback
|
||||
}
|
||||
switch x := v.(type) {
|
||||
case bool:
|
||||
return x
|
||||
case string:
|
||||
switch strings.ToLower(strings.TrimSpace(x)) {
|
||||
case "1", "true", "yes", "on":
|
||||
return true
|
||||
case "0", "false", "no", "off":
|
||||
return false
|
||||
}
|
||||
}
|
||||
return fallback
|
||||
}
|
||||
|
||||
// resolutionsFromPrices mirrors Python models_store._resolutions_from_prices:
|
||||
// an image model's quality tiers ARE its price keys (the admin form never sends
|
||||
// `resolutions` for images), returned in canonical 1K/2K/4K order. gpt-image-2,
|
||||
// for example, only ever has a "1K" price, so it resolves to exactly ["1K"].
|
||||
func resolutionsFromPrices(prices datatypes.JSONMap) datatypes.JSON {
|
||||
out := []string{}
|
||||
for _, r := range []string{"1K", "2K", "4K"} {
|
||||
if _, ok := prices[r]; ok {
|
||||
out = append(out, r)
|
||||
}
|
||||
}
|
||||
return jsonArray(out)
|
||||
}
|
||||
|
||||
func jsonArray(v any) datatypes.JSON {
|
||||
if v == nil {
|
||||
return datatypes.JSON([]byte("[]"))
|
||||
}
|
||||
b, err := json.Marshal(v)
|
||||
if err != nil {
|
||||
return datatypes.JSON([]byte("[]"))
|
||||
}
|
||||
return datatypes.JSON(b)
|
||||
}
|
||||
|
||||
func jsonMap(v any) datatypes.JSONMap {
|
||||
if v == nil {
|
||||
return datatypes.JSONMap{}
|
||||
}
|
||||
switch m := v.(type) {
|
||||
case map[string]any:
|
||||
return datatypes.JSONMap(m)
|
||||
default:
|
||||
b, err := json.Marshal(v)
|
||||
if err != nil {
|
||||
return datatypes.JSONMap{}
|
||||
}
|
||||
var out map[string]any
|
||||
if err := json.Unmarshal(b, &out); err != nil {
|
||||
return datatypes.JSONMap{}
|
||||
}
|
||||
return datatypes.JSONMap(out)
|
||||
}
|
||||
}
|
||||
|
||||
func defaultString(v, fallback string) string {
|
||||
if strings.TrimSpace(v) == "" {
|
||||
return fallback
|
||||
}
|
||||
return v
|
||||
}
|
||||
|
||||
func maxFloat(min, v float64) float64 {
|
||||
if v < min {
|
||||
return min
|
||||
}
|
||||
return v
|
||||
}
|
||||
|
||||
func randomInviteCode() string {
|
||||
return randomUpper(8)
|
||||
}
|
||||
|
||||
var _ = gorm.ErrRecordNotFound
|
||||
@@ -0,0 +1,20 @@
|
||||
package service
|
||||
|
||||
import nanoid "github.com/matoous/go-nanoid/v2"
|
||||
|
||||
const UpperAlphabet = "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
|
||||
|
||||
func randomUpper(n int) string {
|
||||
v, err := nanoid.Generate(UpperAlphabet, n)
|
||||
if err != nil {
|
||||
if n <= 0 {
|
||||
return ""
|
||||
}
|
||||
out := make([]byte, n)
|
||||
for i := range out {
|
||||
out[i] = 'A'
|
||||
}
|
||||
return string(out)
|
||||
}
|
||||
return v
|
||||
}
|
||||
@@ -0,0 +1,130 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"errors"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"backend/internal/model"
|
||||
"backend/internal/repo"
|
||||
)
|
||||
|
||||
type APIKeyService struct {
|
||||
keys *repo.APIKeyRepository
|
||||
}
|
||||
|
||||
func NewAPIKeyService(keys *repo.APIKeyRepository) *APIKeyService {
|
||||
return &APIKeyService{keys: keys}
|
||||
}
|
||||
|
||||
func (s *APIKeyService) Current(ctx context.Context, userID string) (map[string]any, error) {
|
||||
keys, err := s.keys.ListByUserID(ctx, userID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if len(keys) == 0 {
|
||||
return map[string]any{"key": nil}, nil
|
||||
}
|
||||
key := keys[0]
|
||||
return map[string]any{
|
||||
"key": map[string]any{
|
||||
"id": key.ID,
|
||||
"name": key.Name,
|
||||
"key_preview": key.KeyPreview,
|
||||
"created_at": key.CreatedAt,
|
||||
"last_used_at": key.LastUsedAt,
|
||||
},
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (s *APIKeyService) Mint(ctx context.Context, userID string) (map[string]any, error) {
|
||||
plain, err := generatePlainAPIKey()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
key := &model.APIKey{
|
||||
ID: "k-" + time.Now().Format("150405") + randomSuffix(2),
|
||||
UserID: userID,
|
||||
Name: "default",
|
||||
KeyPreview: previewAPIKey(plain),
|
||||
KeyHash: hashAPIKey(plain),
|
||||
CreatedAt: time.Now(),
|
||||
}
|
||||
if err := s.keys.ReplaceForUser(ctx, userID, key); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return map[string]any{
|
||||
"ok": true,
|
||||
"key": plain,
|
||||
"preview": key.KeyPreview,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (s *APIKeyService) Revoke(ctx context.Context, userID string) error {
|
||||
return s.keys.DeleteByUserID(ctx, userID)
|
||||
}
|
||||
|
||||
func (s *APIKeyService) MintNamed(ctx context.Context, userID, name string, replace bool) (map[string]any, error) {
|
||||
name = strings.TrimSpace(name)
|
||||
if name == "" {
|
||||
name = "default"
|
||||
}
|
||||
plain, err := generatePlainAPIKey()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
key := &model.APIKey{
|
||||
ID: "k-" + time.Now().Format("150405") + randomSuffix(2),
|
||||
UserID: userID,
|
||||
Name: name,
|
||||
KeyPreview: previewAPIKey(plain),
|
||||
KeyHash: hashAPIKey(plain),
|
||||
CreatedAt: time.Now(),
|
||||
}
|
||||
if replace {
|
||||
if err := s.keys.ReplaceForUser(ctx, userID, key); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
} else {
|
||||
if err := s.keys.Create(ctx, key); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
return map[string]any{
|
||||
"ok": true,
|
||||
"key": plain,
|
||||
"preview": key.KeyPreview,
|
||||
"id": key.ID,
|
||||
"name": key.Name,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (s *APIKeyService) DeleteOne(ctx context.Context, userID, keyID string) error {
|
||||
if strings.TrimSpace(keyID) == "" {
|
||||
return errors.New("key id required")
|
||||
}
|
||||
return s.keys.DeleteByID(ctx, userID, keyID)
|
||||
}
|
||||
|
||||
func generatePlainAPIKey() (string, error) {
|
||||
return "sk-" + randomUpper(38), nil
|
||||
}
|
||||
|
||||
func previewAPIKey(plain string) string {
|
||||
if len(plain) <= 4 {
|
||||
return strings.Repeat("•", len(plain))
|
||||
}
|
||||
return "…" + plain[len(plain)-4:]
|
||||
}
|
||||
|
||||
func hashAPIKey(plain string) string {
|
||||
sum := sha256.Sum256([]byte(plain))
|
||||
return "sha256:" + hex.EncodeToString(sum[:])
|
||||
}
|
||||
|
||||
func randomSuffix(n int) string {
|
||||
return randomUpper(n)
|
||||
}
|
||||
@@ -0,0 +1,438 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"backend/internal/repo"
|
||||
"backend/internal/storage"
|
||||
)
|
||||
|
||||
type AppSettingsService struct {
|
||||
settings *repo.SiteSettingRepository
|
||||
events *repo.EventRepository
|
||||
smtp *SMTPService
|
||||
store *storage.Client
|
||||
}
|
||||
|
||||
type RegistrationSettings struct {
|
||||
Open bool `json:"open"`
|
||||
EmailCode bool `json:"email_code"`
|
||||
AllowPasswordReset bool `json:"allow_password_reset"`
|
||||
AllowedDomains []string `json:"allowed_email_domains"`
|
||||
CodeTTLSeconds int `json:"code_ttl_seconds"`
|
||||
}
|
||||
|
||||
type SMTPSettings struct {
|
||||
Host string `json:"host"`
|
||||
Port int `json:"port"`
|
||||
Username string `json:"username"`
|
||||
Password string `json:"password"`
|
||||
FromAddr string `json:"from_addr"`
|
||||
UseTLS bool `json:"use_tls"`
|
||||
}
|
||||
|
||||
type CreditSettings struct {
|
||||
CheckinEnabled bool `json:"checkin_enabled"`
|
||||
CheckinReward int `json:"checkin_reward"`
|
||||
InviteEnabled bool `json:"invite_enabled"`
|
||||
InviteReward int `json:"invite_reward"`
|
||||
}
|
||||
|
||||
type ProxySettings struct {
|
||||
Proxy string `json:"proxy"`
|
||||
}
|
||||
|
||||
type RetentionSettings struct {
|
||||
RetentionDays int `json:"retention_days"`
|
||||
}
|
||||
|
||||
type MediaRetentionResult struct {
|
||||
Settings *RetentionSettings
|
||||
Removed int `json:"removed"`
|
||||
FreedBytes int64 `json:"freed_bytes"`
|
||||
}
|
||||
|
||||
func NewAppSettingsService(settings *repo.SiteSettingRepository, events *repo.EventRepository, smtp *SMTPService, store *storage.Client) *AppSettingsService {
|
||||
return &AppSettingsService{
|
||||
settings: settings,
|
||||
events: events,
|
||||
smtp: smtp,
|
||||
store: store,
|
||||
}
|
||||
}
|
||||
|
||||
func (s *AppSettingsService) Registration(ctx context.Context) (*RegistrationSettings, error) {
|
||||
openRaw, err := s.settings.GetValue(ctx, "auth.open")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
emailCodeRaw, err := s.settings.GetValue(ctx, "auth.email_code")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
resetRaw, err := s.settings.GetValue(ctx, "auth.allow_password_reset")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
domainsRaw, err := s.settings.GetValue(ctx, "auth.allowed_email_domains")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
ttlRaw, err := s.settings.GetValue(ctx, "auth.code_ttl_seconds")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
ttl, _ := strconv.Atoi(strings.TrimSpace(ttlRaw))
|
||||
if ttl < 60 {
|
||||
ttl = 600
|
||||
}
|
||||
return &RegistrationSettings{
|
||||
Open: parseBoolSetting(openRaw, true),
|
||||
EmailCode: parseBoolSetting(emailCodeRaw, false),
|
||||
AllowPasswordReset: parseBoolSetting(resetRaw, false),
|
||||
AllowedDomains: parseCSVSetting(domainsRaw),
|
||||
CodeTTLSeconds: ttl,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (s *AppSettingsService) SaveRegistration(ctx context.Context, in RegistrationSettings) (*RegistrationSettings, error) {
|
||||
// Empty list is allowed and means "no domain restriction": EmailDomainAllowed
|
||||
// returns true for everyone when the whitelist is empty.
|
||||
domains := ValidateAllowedEmailDomains(in.AllowedDomains)
|
||||
if in.CodeTTLSeconds < 60 {
|
||||
in.CodeTTLSeconds = 600
|
||||
}
|
||||
if err := s.settings.UpsertValues(ctx, map[string]string{
|
||||
"auth.open": strconv.FormatBool(in.Open),
|
||||
"auth.email_code": strconv.FormatBool(in.EmailCode),
|
||||
"auth.allow_password_reset": strconv.FormatBool(in.AllowPasswordReset),
|
||||
"auth.allowed_email_domains": strings.Join(domains, ","),
|
||||
"auth.code_ttl_seconds": strconv.Itoa(in.CodeTTLSeconds),
|
||||
}); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return s.Registration(ctx)
|
||||
}
|
||||
|
||||
func (s *AppSettingsService) SMTP(ctx context.Context) (*SMTPSettings, error) {
|
||||
host, err := s.settings.GetValue(ctx, "smtp.host")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
portRaw, err := s.settings.GetValue(ctx, "smtp.port")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
username, err := s.settings.GetValue(ctx, "smtp.username")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
password, err := s.settings.GetValue(ctx, "smtp.password")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
fromAddr, err := s.settings.GetValue(ctx, "smtp.from_addr")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
useTLSRaw, err := s.settings.GetValue(ctx, "smtp.use_tls")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
port, _ := strconv.Atoi(strings.TrimSpace(portRaw))
|
||||
if port <= 0 {
|
||||
port = 587
|
||||
}
|
||||
return &SMTPSettings{
|
||||
Host: strings.TrimSpace(host),
|
||||
Port: port,
|
||||
Username: strings.TrimSpace(username),
|
||||
Password: maskedSecret(password),
|
||||
FromAddr: strings.TrimSpace(fromAddr),
|
||||
UseTLS: parseBoolSetting(useTLSRaw, true),
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (s *AppSettingsService) SaveSMTP(ctx context.Context, in SMTPSettings) (*SMTPSettings, error) {
|
||||
host := strings.TrimSpace(in.Host)
|
||||
username := strings.TrimSpace(in.Username)
|
||||
fromAddr := strings.TrimSpace(in.FromAddr)
|
||||
if host == "" || username == "" || fromAddr == "" {
|
||||
return nil, errors.New("请填写 主机 / 用户名 / 发件地址")
|
||||
}
|
||||
if _, err := ValidateEmail(fromAddr); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if in.Port <= 0 {
|
||||
return nil, errors.New("port 必须是正整数")
|
||||
}
|
||||
|
||||
updates := map[string]string{
|
||||
"smtp.host": host,
|
||||
"smtp.port": strconv.Itoa(in.Port),
|
||||
"smtp.username": username,
|
||||
"smtp.from_addr": fromAddr,
|
||||
"smtp.use_tls": strconv.FormatBool(in.UseTLS),
|
||||
}
|
||||
if strings.TrimSpace(in.Password) != "" && strings.TrimSpace(in.Password) != "***" {
|
||||
updates["smtp.password"] = in.Password
|
||||
}
|
||||
if err := s.settings.UpsertValues(ctx, updates); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return s.SMTP(ctx)
|
||||
}
|
||||
|
||||
func (s *AppSettingsService) TestSMTP(ctx context.Context, to string) error {
|
||||
to, err := ValidateEmail(to)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
cfg, err := s.loadSMTPConfig(ctx)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return s.smtp.SendCode(ctx, cfg, to, "123456", "register")
|
||||
}
|
||||
|
||||
func (s *AppSettingsService) Proxy(ctx context.Context) (*ProxySettings, error) {
|
||||
proxy, err := s.settings.GetValue(ctx, "proxy.url")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &ProxySettings{Proxy: strings.TrimSpace(proxy)}, nil
|
||||
}
|
||||
|
||||
func (s *AppSettingsService) SaveProxy(ctx context.Context, proxy string) (*ProxySettings, error) {
|
||||
proxy = strings.TrimSpace(proxy)
|
||||
if err := s.settings.UpsertValue(ctx, "proxy.url", proxy); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &ProxySettings{Proxy: proxy}, nil
|
||||
}
|
||||
|
||||
// TestProxy routes a probe request through the given proxy to an IP-echo service
|
||||
// and reports the egress IP + latency. Tests the value passed in (so the admin
|
||||
// can verify before saving). Mirrors how generation calls go out — same HTTP
|
||||
// CONNECT through the proxy — so a green result means upstream calls will route.
|
||||
func (s *AppSettingsService) TestProxy(ctx context.Context, proxy string) (map[string]any, error) {
|
||||
proxy = strings.TrimSpace(proxy)
|
||||
if proxy == "" {
|
||||
return nil, errors.New("代理地址为空,请先填写")
|
||||
}
|
||||
parsed, err := url.Parse(proxy)
|
||||
if err != nil || parsed.Host == "" {
|
||||
return nil, fmt.Errorf("代理地址格式不正确(应形如 http://user:pass@host:port)")
|
||||
}
|
||||
|
||||
transport := &http.Transport{Proxy: http.ProxyURL(parsed)}
|
||||
defer transport.CloseIdleConnections()
|
||||
client := &http.Client{Transport: transport, Timeout: 12 * time.Second}
|
||||
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, "https://api.ipify.org?format=json", nil)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
start := time.Now()
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("通过代理请求失败:%v", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
elapsed := int(time.Since(start).Milliseconds())
|
||||
body, _ := io.ReadAll(io.LimitReader(resp.Body, 4096))
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return nil, fmt.Errorf("代理已连接,但探测返回 HTTP %d", resp.StatusCode)
|
||||
}
|
||||
var echo struct {
|
||||
IP string `json:"ip"`
|
||||
}
|
||||
_ = json.Unmarshal(body, &echo)
|
||||
return map[string]any{
|
||||
"exit_ip": echo.IP,
|
||||
"elapsed_ms": elapsed,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (s *AppSettingsService) Credits(ctx context.Context) (*CreditSettings, error) {
|
||||
checkinEnabledRaw, err := s.settings.GetValue(ctx, "credits.checkin_enabled")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
checkinRewardRaw, err := s.settings.GetValue(ctx, "credits.checkin_reward")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
inviteEnabledRaw, err := s.settings.GetValue(ctx, "credits.invite_enabled")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
inviteRewardRaw, err := s.settings.GetValue(ctx, "credits.invite_reward")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &CreditSettings{
|
||||
CheckinEnabled: parseBoolSetting(checkinEnabledRaw, true),
|
||||
CheckinReward: parseIntSetting(checkinRewardRaw, 3),
|
||||
InviteEnabled: parseBoolSetting(inviteEnabledRaw, true),
|
||||
InviteReward: parseIntSetting(inviteRewardRaw, 3),
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (s *AppSettingsService) SaveCredits(ctx context.Context, in CreditSettings) (*CreditSettings, error) {
|
||||
if in.CheckinReward < 0 {
|
||||
in.CheckinReward = 0
|
||||
}
|
||||
if in.InviteReward < 0 {
|
||||
in.InviteReward = 0
|
||||
}
|
||||
if err := s.settings.UpsertValues(ctx, map[string]string{
|
||||
"credits.checkin_enabled": strconv.FormatBool(in.CheckinEnabled),
|
||||
"credits.checkin_reward": strconv.Itoa(in.CheckinReward),
|
||||
"credits.invite_enabled": strconv.FormatBool(in.InviteEnabled),
|
||||
"credits.invite_reward": strconv.Itoa(in.InviteReward),
|
||||
}); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return s.Credits(ctx)
|
||||
}
|
||||
|
||||
func (s *AppSettingsService) Logs(ctx context.Context) (*RetentionSettings, error) {
|
||||
return s.retention(ctx, "logs.retention_days")
|
||||
}
|
||||
|
||||
func (s *AppSettingsService) SaveLogs(ctx context.Context, days int) (*RetentionSettings, error) {
|
||||
days, err := normalizeRetentionDays(days)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := s.settings.UpsertValue(ctx, "logs.retention_days", strconv.Itoa(days)); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if s.events != nil {
|
||||
_, _ = s.events.PurgeOlderThan(ctx, time.Duration(days)*24*time.Hour)
|
||||
}
|
||||
return s.Logs(ctx)
|
||||
}
|
||||
|
||||
func (s *AppSettingsService) Media(ctx context.Context) (*RetentionSettings, error) {
|
||||
return s.retention(ctx, "media.retention_days")
|
||||
}
|
||||
|
||||
func (s *AppSettingsService) SaveMedia(ctx context.Context, days int) (*MediaRetentionResult, error) {
|
||||
days, err := normalizeRetentionDays(days)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := s.settings.UpsertValue(ctx, "media.retention_days", strconv.Itoa(days)); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
removed, freed := s.pruneGeneratedFiles(ctx, time.Duration(days)*24*time.Hour)
|
||||
settings, err := s.Media(ctx)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &MediaRetentionResult{
|
||||
Settings: settings,
|
||||
Removed: removed,
|
||||
FreedBytes: freed,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (s *AppSettingsService) loadSMTPConfig(ctx context.Context) (SMTPConfig, error) {
|
||||
current, err := s.SMTP(ctx)
|
||||
if err != nil {
|
||||
return SMTPConfig{}, err
|
||||
}
|
||||
password, err := s.settings.GetValue(ctx, "smtp.password")
|
||||
if err != nil {
|
||||
return SMTPConfig{}, err
|
||||
}
|
||||
return SMTPConfig{
|
||||
Host: current.Host,
|
||||
Port: current.Port,
|
||||
Username: current.Username,
|
||||
Password: password,
|
||||
FromAddr: current.FromAddr,
|
||||
UseTLS: current.UseTLS,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func maskedSecret(v string) string {
|
||||
if strings.TrimSpace(v) == "" {
|
||||
return ""
|
||||
}
|
||||
return "***"
|
||||
}
|
||||
|
||||
func parseIntSetting(v string, fallback int) int {
|
||||
n, err := strconv.Atoi(strings.TrimSpace(v))
|
||||
if err != nil {
|
||||
return fallback
|
||||
}
|
||||
return n
|
||||
}
|
||||
|
||||
func (s *AppSettingsService) retention(ctx context.Context, key string) (*RetentionSettings, error) {
|
||||
raw, err := s.settings.GetValue(ctx, key)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
days := parseIntSetting(raw, 30)
|
||||
if days < 1 {
|
||||
days = 30
|
||||
}
|
||||
return &RetentionSettings{RetentionDays: days}, nil
|
||||
}
|
||||
|
||||
func normalizeRetentionDays(days int) (int, error) {
|
||||
if days < 1 {
|
||||
return 0, errors.New("留存天数至少为 1 天")
|
||||
}
|
||||
if days > 365 {
|
||||
return 0, errors.New("留存天数最多 365 天")
|
||||
}
|
||||
return days, nil
|
||||
}
|
||||
|
||||
// pruneGeneratedFiles deletes RustFS objects older than maxAge and blanks the
|
||||
// matching event_log.file refs. Returns how many were removed and bytes freed.
|
||||
// (The maintenance loop does the same automatically every 60s; this gives the
|
||||
// admin an immediate result when they shorten the media retention window.)
|
||||
func (s *AppSettingsService) pruneGeneratedFiles(ctx context.Context, maxAge time.Duration) (int, int64) {
|
||||
if s.store == nil || !s.store.Configured() || maxAge <= 0 {
|
||||
return 0, 0
|
||||
}
|
||||
objs, err := s.store.List(ctx, "")
|
||||
if err != nil {
|
||||
return 0, 0
|
||||
}
|
||||
cutoff := time.Now().Add(-maxAge)
|
||||
removed := 0
|
||||
var freed int64
|
||||
var clearedKeys []string
|
||||
for _, o := range objs {
|
||||
if !o.LastModified.Before(cutoff) {
|
||||
continue
|
||||
}
|
||||
if err := s.store.Delete(ctx, o.Key); err == nil {
|
||||
removed++
|
||||
freed += o.Size
|
||||
clearedKeys = append(clearedKeys, o.Key)
|
||||
}
|
||||
}
|
||||
if len(clearedKeys) > 0 {
|
||||
_, _ = s.events.ClearFiles(ctx, clearedKeys)
|
||||
}
|
||||
return removed, freed
|
||||
}
|
||||
@@ -0,0 +1,584 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"errors"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"backend/internal/model"
|
||||
"backend/internal/repo"
|
||||
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
var ErrAuthFailed = errors.New("auth failed")
|
||||
|
||||
type AuthService struct {
|
||||
users *repo.UserRepository
|
||||
settings *repo.SiteSettingRepository
|
||||
sessions *SessionService
|
||||
codes *EmailCodeService
|
||||
smtp *SMTPService
|
||||
loginGuard *LoginGuard
|
||||
}
|
||||
|
||||
type AuthSettings struct {
|
||||
Open bool
|
||||
EmailCode bool
|
||||
AllowPasswordReset bool
|
||||
AllowedDomains []string
|
||||
}
|
||||
|
||||
func NewAuthService(
|
||||
users *repo.UserRepository,
|
||||
settings *repo.SiteSettingRepository,
|
||||
sessions *SessionService,
|
||||
codes *EmailCodeService,
|
||||
smtp *SMTPService,
|
||||
) *AuthService {
|
||||
return &AuthService{
|
||||
users: users,
|
||||
settings: settings,
|
||||
sessions: sessions,
|
||||
codes: codes,
|
||||
smtp: smtp,
|
||||
loginGuard: NewLoginGuard(codes.Redis()),
|
||||
}
|
||||
}
|
||||
|
||||
func (s *AuthService) IsAuthorizedForPrivateImage(ctx context.Context, sessionCookie, owner string) (bool, error) {
|
||||
// Private images are viewable ONLY via a logged-in session cookie (no Bearer
|
||||
// token / API key). A regular user may view only their OWN images; an admin
|
||||
// may view anyone's. `owner` is the /images/<owner>/... path segment.
|
||||
if sessionCookie == "" {
|
||||
return false, nil
|
||||
}
|
||||
payload, err := s.sessions.Validate(ctx, sessionCookie)
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
if payload == nil {
|
||||
return false, nil
|
||||
}
|
||||
user, err := s.users.GetByID(ctx, payload.UserID)
|
||||
if err != nil {
|
||||
if err == gorm.ErrRecordNotFound {
|
||||
return false, nil
|
||||
}
|
||||
return false, err
|
||||
}
|
||||
if user.Role == "admin" {
|
||||
return true, nil
|
||||
}
|
||||
return ownsImageDir(user, owner), nil
|
||||
}
|
||||
|
||||
// ownsImageDir reports whether `owner` (the /images/<owner>/... directory) is one
|
||||
// of the names this user's outputs are stored under. Mirrors the candidates
|
||||
// V1Service.userDir picks from: sanitized name → sanitized email-local → id.
|
||||
func ownsImageDir(user *model.User, owner string) bool {
|
||||
owner = strings.TrimSpace(owner)
|
||||
if owner == "" || user == nil {
|
||||
return false
|
||||
}
|
||||
if owner == user.ID {
|
||||
return true
|
||||
}
|
||||
if d := sanitizeOwnerName(user.Name); d != "" && d == owner {
|
||||
return true
|
||||
}
|
||||
if d := sanitizeOwnerName(strings.Split(user.Email, "@")[0]); d != "" && d == owner {
|
||||
return true
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func (s *AuthService) CurrentUserFromBearer(ctx context.Context, authHeader string) (*model.User, *SessionPayload, error) {
|
||||
token := ParseBearer(authHeader)
|
||||
return s.currentUserFromToken(ctx, token)
|
||||
}
|
||||
|
||||
func (s *AuthService) CurrentUserFromRequest(ctx context.Context, authHeader, cookieToken string) (*model.User, *SessionPayload, error) {
|
||||
if user, session, err := s.CurrentUserFromBearer(ctx, authHeader); err != nil || user != nil || session != nil {
|
||||
return user, session, err
|
||||
}
|
||||
return s.currentUserFromToken(ctx, cookieToken)
|
||||
}
|
||||
|
||||
func (s *AuthService) CurrentUserFromToken(ctx context.Context, token string) (*model.User, *SessionPayload, error) {
|
||||
return s.currentUserFromToken(ctx, token)
|
||||
}
|
||||
|
||||
func (s *AuthService) currentUserFromToken(ctx context.Context, token string) (*model.User, *SessionPayload, error) {
|
||||
if token == "" {
|
||||
return nil, nil, nil
|
||||
}
|
||||
|
||||
payload, err := s.sessions.Validate(ctx, token)
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
if payload == nil {
|
||||
return nil, nil, nil
|
||||
}
|
||||
|
||||
user, err := s.users.GetByID(ctx, payload.UserID)
|
||||
if err != nil {
|
||||
if err == gorm.ErrRecordNotFound {
|
||||
return nil, nil, nil
|
||||
}
|
||||
return nil, nil, err
|
||||
}
|
||||
if user.Status != "active" {
|
||||
return nil, nil, nil
|
||||
}
|
||||
return user, payload, nil
|
||||
}
|
||||
|
||||
func (s *AuthService) Login(ctx context.Context, identifier, password, ip string) (*model.User, string, *SessionPayload, error) {
|
||||
normalizedIdentifier, err := ValidateLoginIdentifier(identifier)
|
||||
if err != nil {
|
||||
return nil, "", nil, err
|
||||
}
|
||||
if strings.TrimSpace(password) == "" {
|
||||
return nil, "", nil, errors.New("密码不能为空")
|
||||
}
|
||||
|
||||
// Exponential-backoff lockout per (ip, account) + per-ip spray (Python
|
||||
// api/auth.py:226-237 via core/login_guard.py).
|
||||
if err := s.loginGuard.Check(ctx, ip, normalizedIdentifier); err != nil {
|
||||
return nil, "", nil, err
|
||||
}
|
||||
|
||||
user, err := s.users.GetByIdentifier(ctx, normalizedIdentifier)
|
||||
if err != nil {
|
||||
if err == gorm.ErrRecordNotFound {
|
||||
if rerr := s.loginGuard.RecordFailure(ctx, ip, normalizedIdentifier); rerr != nil {
|
||||
return nil, "", nil, rerr
|
||||
}
|
||||
return nil, "", nil, ErrAuthFailed
|
||||
}
|
||||
return nil, "", nil, err
|
||||
}
|
||||
if user.Status != "active" || !VerifyPassword(password, user.PasswordHash) {
|
||||
if rerr := s.loginGuard.RecordFailure(ctx, ip, normalizedIdentifier); rerr != nil {
|
||||
return nil, "", nil, rerr
|
||||
}
|
||||
return nil, "", nil, ErrAuthFailed
|
||||
}
|
||||
if err := s.loginGuard.RecordSuccess(ctx, ip, normalizedIdentifier); err != nil {
|
||||
return nil, "", nil, err
|
||||
}
|
||||
|
||||
if err := s.users.TouchLogin(ctx, user.ID, ip); err != nil {
|
||||
return nil, "", nil, err
|
||||
}
|
||||
token, payload, err := s.sessions.Create(ctx, user.ID)
|
||||
if err != nil {
|
||||
return nil, "", nil, err
|
||||
}
|
||||
return user, token, payload, nil
|
||||
}
|
||||
|
||||
func (s *AuthService) SendCode(ctx context.Context, email, purpose string) error {
|
||||
cfg, err := s.loadAuthSettings(ctx)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if !cfg.EmailCode {
|
||||
return errors.New("未开启邮箱验证码")
|
||||
}
|
||||
|
||||
normalizedEmail, err := ValidateEmail(email)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
purpose = strings.ToLower(strings.TrimSpace(purpose))
|
||||
switch purpose {
|
||||
case "register", "reset":
|
||||
default:
|
||||
return errors.New("验证码用途不正确")
|
||||
}
|
||||
|
||||
if purpose == "register" && !EmailDomainAllowed(normalizedEmail, cfg.AllowedDomains) {
|
||||
return errors.New("该邮箱后缀不允许注册")
|
||||
}
|
||||
|
||||
code, err := s.codes.Issue(ctx, normalizedEmail, purpose)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return s.smtp.SendCode(ctx, s.loadSMTPSettings(ctx), normalizedEmail, code, purpose)
|
||||
}
|
||||
|
||||
func (s *AuthService) Register(ctx context.Context, email, username, password, inviteCode, emailCode, ip string) (*model.User, string, *SessionPayload, error) {
|
||||
normalizedEmail, err := ValidateEmail(email)
|
||||
if err != nil {
|
||||
return nil, "", nil, err
|
||||
}
|
||||
normalizedUsername, err := ValidateUsername(username)
|
||||
if err != nil {
|
||||
return nil, "", nil, err
|
||||
}
|
||||
if err := ValidatePassword(password); err != nil {
|
||||
return nil, "", nil, err
|
||||
}
|
||||
|
||||
settings, err := s.loadAuthSettings(ctx)
|
||||
if err != nil {
|
||||
return nil, "", nil, err
|
||||
}
|
||||
hasAdmin, err := s.users.HasAdmin(ctx)
|
||||
if err != nil {
|
||||
return nil, "", nil, err
|
||||
}
|
||||
// The very first account ever bootstraps the admin and skips the open
|
||||
// toggle, the email-domain whitelist, and the email-code gate (Python
|
||||
// api/auth.py:195-204). All three are only enforced once an admin exists.
|
||||
if hasAdmin && !settings.Open {
|
||||
return nil, "", nil, errors.New("当前未开放注册")
|
||||
}
|
||||
if hasAdmin && !EmailDomainAllowed(normalizedEmail, settings.AllowedDomains) {
|
||||
return nil, "", nil, errors.New("该邮箱后缀不允许注册")
|
||||
}
|
||||
if hasAdmin && settings.EmailCode {
|
||||
ok, err := s.codes.Verify(ctx, normalizedEmail, "register", emailCode)
|
||||
if err != nil {
|
||||
return nil, "", nil, err
|
||||
}
|
||||
if !ok {
|
||||
return nil, "", nil, errors.New("邮箱验证码错误或已过期")
|
||||
}
|
||||
}
|
||||
|
||||
exists, err := s.users.ExistsEmail(ctx, normalizedEmail, "")
|
||||
if err != nil {
|
||||
return nil, "", nil, err
|
||||
}
|
||||
if exists {
|
||||
return nil, "", nil, errors.New("邮箱已存在")
|
||||
}
|
||||
exists, err = s.users.ExistsName(ctx, normalizedUsername, "")
|
||||
if err != nil {
|
||||
return nil, "", nil, err
|
||||
}
|
||||
if exists {
|
||||
return nil, "", nil, errors.New("用户名已存在")
|
||||
}
|
||||
|
||||
passwordHash, err := HashPassword(password)
|
||||
if err != nil {
|
||||
return nil, "", nil, err
|
||||
}
|
||||
role := "user"
|
||||
if !hasAdmin {
|
||||
role = "admin"
|
||||
}
|
||||
|
||||
var invitedBy *string
|
||||
if strings.TrimSpace(inviteCode) != "" {
|
||||
inviter, err := s.users.GetByInviteCode(ctx, inviteCode)
|
||||
if err == nil {
|
||||
invitedBy = &inviter.ID
|
||||
}
|
||||
}
|
||||
|
||||
now := time.Now()
|
||||
user := &model.User{
|
||||
ID: "u-" + randomUpper(10),
|
||||
Email: normalizedEmail,
|
||||
Name: normalizedUsername,
|
||||
PasswordHash: passwordHash,
|
||||
Role: role,
|
||||
Status: "active",
|
||||
InviteCode: randomInviteCode(),
|
||||
InvitedBy: invitedBy,
|
||||
CreatedAt: now,
|
||||
UpdatedAt: now,
|
||||
}
|
||||
if err := s.users.Create(ctx, user); err != nil {
|
||||
return nil, "", nil, err
|
||||
}
|
||||
if err := s.users.TouchLogin(ctx, user.ID, ip); err != nil {
|
||||
return nil, "", nil, err
|
||||
}
|
||||
token, payload, err := s.sessions.Create(ctx, user.ID)
|
||||
if err != nil {
|
||||
return nil, "", nil, err
|
||||
}
|
||||
created, err := s.users.GetByID(ctx, user.ID)
|
||||
if err != nil {
|
||||
return nil, "", nil, err
|
||||
}
|
||||
return created, token, payload, nil
|
||||
}
|
||||
|
||||
func (s *AuthService) ResetPassword(ctx context.Context, email, password, emailCode, ip string) error {
|
||||
settings, err := s.loadAuthSettings(ctx)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if !settings.EmailCode || !settings.AllowPasswordReset {
|
||||
return errors.New("未开放找回密码")
|
||||
}
|
||||
normalizedEmail, err := ValidateEmail(email)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if err := ValidatePassword(password); err != nil {
|
||||
return err
|
||||
}
|
||||
// Rate-limit reset attempts per IP+email so the 6-digit code can't be ground
|
||||
// down even with the single-use + wrong-guess cap (Python api/auth.py:257-268).
|
||||
guardID := "reset:" + normalizedEmail
|
||||
if err := s.loginGuard.Check(ctx, ip, guardID); err != nil {
|
||||
return err
|
||||
}
|
||||
ok, err := s.codes.Verify(ctx, normalizedEmail, "reset", emailCode)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if !ok {
|
||||
if rerr := s.loginGuard.RecordFailure(ctx, ip, guardID); rerr != nil {
|
||||
return rerr
|
||||
}
|
||||
return errors.New("邮箱验证码错误或已过期")
|
||||
}
|
||||
if err := s.loginGuard.RecordSuccess(ctx, ip, guardID); err != nil {
|
||||
return err
|
||||
}
|
||||
passwordHash, err := HashPassword(password)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
_, err = s.users.SetPasswordByEmail(ctx, normalizedEmail, passwordHash)
|
||||
return err
|
||||
}
|
||||
|
||||
func (s *AuthService) ChangePassword(ctx context.Context, userID, currentPassword, newPassword string) error {
|
||||
if strings.TrimSpace(currentPassword) == "" {
|
||||
return errors.New("当前密码不能为空")
|
||||
}
|
||||
if err := ValidatePassword(newPassword); err != nil {
|
||||
return err
|
||||
}
|
||||
user, err := s.users.GetByID(ctx, userID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if !VerifyPassword(currentPassword, user.PasswordHash) {
|
||||
return errors.New("当前密码错误")
|
||||
}
|
||||
passwordHash, err := HashPassword(newPassword)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
_, err = s.users.Update(ctx, userID, map[string]any{
|
||||
"password_hash": passwordHash,
|
||||
})
|
||||
return err
|
||||
}
|
||||
|
||||
func (s *AuthService) Logout(ctx context.Context, token string) error {
|
||||
return s.sessions.Destroy(ctx, token)
|
||||
}
|
||||
|
||||
func (s *AuthService) AuthConfig(ctx context.Context) (map[string]any, error) {
|
||||
hasAdmin, err := s.users.HasAdmin(ctx)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
settings, err := s.loadAuthSettings(ctx)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
credits, err := s.loadCreditSettings(ctx)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return map[string]any{
|
||||
"open": settings.Open,
|
||||
"email_code": settings.EmailCode,
|
||||
"allow_password_reset": settings.AllowPasswordReset,
|
||||
"allowed_email_domains": settings.AllowedDomains,
|
||||
"has_admin": hasAdmin,
|
||||
"checkin_enabled": credits.CheckinEnabled,
|
||||
"checkin_reward": credits.CheckinReward,
|
||||
"invite_enabled": credits.InviteEnabled,
|
||||
"invite_reward": credits.InviteReward,
|
||||
"server_time": time.Now().Unix(),
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (s *AuthService) PublicUser(ctx context.Context, user *model.User) (map[string]any, error) {
|
||||
if user == nil {
|
||||
return nil, nil
|
||||
}
|
||||
credits, err := s.loadCreditSettings(ctx)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
stats, err := s.users.InviteStats(ctx, user.ID, credits.InviteReward)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return map[string]any{
|
||||
"id": user.ID,
|
||||
"email": user.Email,
|
||||
"name": user.Name,
|
||||
"role": user.Role,
|
||||
"status": user.Status,
|
||||
"credits": user.Credits,
|
||||
"checkin_last": user.CheckinLast,
|
||||
"checkin_streak": user.CheckinStreak,
|
||||
"checkin_today": user.CheckinLast == time.Now().Format("2006-01-02"),
|
||||
"invite_code": user.InviteCode,
|
||||
"invite_count": stats.InviteCount,
|
||||
"invite_earned": stats.InviteEarned,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (s *AuthService) Checkin(ctx context.Context, userID string) (*repo.CheckinResult, error) {
|
||||
credits, err := s.loadCreditSettings(ctx)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if !credits.CheckinEnabled {
|
||||
return nil, errors.New("签到功能未开启")
|
||||
}
|
||||
return s.users.DailyCheckin(ctx, userID, credits.CheckinReward)
|
||||
}
|
||||
|
||||
func (s *AuthService) InviteList(ctx context.Context, userID string) ([]repo.InviteRecord, error) {
|
||||
credits, err := s.loadCreditSettings(ctx)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return s.users.InviteList(ctx, userID, credits.InviteReward)
|
||||
}
|
||||
|
||||
func ParseBearer(header string) string {
|
||||
if header == "" {
|
||||
return ""
|
||||
}
|
||||
lower := strings.ToLower(header)
|
||||
if !strings.HasPrefix(lower, "bearer ") {
|
||||
return ""
|
||||
}
|
||||
return strings.TrimSpace(header[7:])
|
||||
}
|
||||
|
||||
func HashAPIKey(plaintext string) string {
|
||||
sum := sha256.Sum256([]byte(plaintext))
|
||||
return "sha256:" + hex.EncodeToString(sum[:])
|
||||
}
|
||||
|
||||
func (s *AuthService) loadAuthSettings(ctx context.Context) (*AuthSettings, error) {
|
||||
openRaw, err := s.settings.GetValue(ctx, "auth.open")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
emailCodeRaw, err := s.settings.GetValue(ctx, "auth.email_code")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
resetRaw, err := s.settings.GetValue(ctx, "auth.allow_password_reset")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
domainsRaw, err := s.settings.GetValue(ctx, "auth.allowed_email_domains")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &AuthSettings{
|
||||
Open: parseBoolSetting(openRaw, true),
|
||||
EmailCode: parseBoolSetting(emailCodeRaw, false),
|
||||
AllowPasswordReset: parseBoolSetting(resetRaw, false),
|
||||
AllowedDomains: parseCSVSetting(domainsRaw),
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (s *AuthService) loadSMTPSettings(ctx context.Context) SMTPConfig {
|
||||
host, _ := s.settings.GetValue(ctx, "smtp.host")
|
||||
portRaw, _ := s.settings.GetValue(ctx, "smtp.port")
|
||||
username, _ := s.settings.GetValue(ctx, "smtp.username")
|
||||
password, _ := s.settings.GetValue(ctx, "smtp.password")
|
||||
fromAddr, _ := s.settings.GetValue(ctx, "smtp.from_addr")
|
||||
useTLSRaw, _ := s.settings.GetValue(ctx, "smtp.use_tls")
|
||||
|
||||
port, _ := strconv.Atoi(strings.TrimSpace(portRaw))
|
||||
if port <= 0 {
|
||||
port = 587
|
||||
}
|
||||
// Fall back to username when from_addr is unset (Python core/email_codes.py:92).
|
||||
from := strings.TrimSpace(fromAddr)
|
||||
if from == "" {
|
||||
from = strings.TrimSpace(username)
|
||||
}
|
||||
return SMTPConfig{
|
||||
Host: strings.TrimSpace(host),
|
||||
Port: port,
|
||||
Username: strings.TrimSpace(username),
|
||||
Password: password,
|
||||
FromAddr: from,
|
||||
// use_tls defaults to true to match Python (core/email_codes.py:93).
|
||||
UseTLS: parseBoolSetting(useTLSRaw, true),
|
||||
}
|
||||
}
|
||||
|
||||
func parseBoolSetting(v string, fallback bool) bool {
|
||||
switch strings.ToLower(strings.TrimSpace(v)) {
|
||||
case "1", "true", "yes", "on":
|
||||
return true
|
||||
case "0", "false", "no", "off":
|
||||
return false
|
||||
default:
|
||||
return fallback
|
||||
}
|
||||
}
|
||||
|
||||
func parseCSVSetting(v string) []string {
|
||||
if strings.TrimSpace(v) == "" {
|
||||
return []string{}
|
||||
}
|
||||
return ValidateAllowedEmailDomains(strings.Split(v, ","))
|
||||
}
|
||||
|
||||
// InviteReward returns the admin-configured 积分 awarded per completed invite
|
||||
// (falls back to 3). Exposed so the invite page shows the real number.
|
||||
func (s *AuthService) InviteReward(ctx context.Context) int {
|
||||
cs, err := s.loadCreditSettings(ctx)
|
||||
if err != nil {
|
||||
return 3
|
||||
}
|
||||
return cs.InviteReward
|
||||
}
|
||||
|
||||
func (s *AuthService) loadCreditSettings(ctx context.Context) (*CreditSettings, error) {
|
||||
checkinEnabledRaw, err := s.settings.GetValue(ctx, "credits.checkin_enabled")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
checkinRewardRaw, err := s.settings.GetValue(ctx, "credits.checkin_reward")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
inviteEnabledRaw, err := s.settings.GetValue(ctx, "credits.invite_enabled")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
inviteRewardRaw, err := s.settings.GetValue(ctx, "credits.invite_reward")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &CreditSettings{
|
||||
CheckinEnabled: parseBoolSetting(checkinEnabledRaw, true),
|
||||
CheckinReward: parseIntSetting(checkinRewardRaw, 3),
|
||||
InviteEnabled: parseBoolSetting(inviteEnabledRaw, true),
|
||||
InviteReward: parseIntSetting(inviteRewardRaw, 3),
|
||||
}, nil
|
||||
}
|
||||
@@ -0,0 +1,167 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"backend/internal/model"
|
||||
"backend/internal/repo"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
type CDKService struct {
|
||||
cdks *repo.CDKRepository
|
||||
users *repo.UserRepository
|
||||
}
|
||||
|
||||
func NewCDKService(cdks *repo.CDKRepository, users *repo.UserRepository) *CDKService {
|
||||
return &CDKService{
|
||||
cdks: cdks,
|
||||
users: users,
|
||||
}
|
||||
}
|
||||
|
||||
func (s *CDKService) List(ctx context.Context) ([]model.CDKCode, map[string]any, map[string]string, error) {
|
||||
items, err := s.cdks.List(ctx)
|
||||
if err != nil {
|
||||
return nil, nil, nil, err
|
||||
}
|
||||
stats, err := s.cdks.Stats(ctx)
|
||||
if err != nil {
|
||||
return nil, nil, nil, err
|
||||
}
|
||||
// Build an id -> display name map (name, else email, else id) so the handler
|
||||
// can annotate redeemed codes with redeemed_by_name (mirrors admin.py).
|
||||
users, err := s.users.List(ctx)
|
||||
if err != nil {
|
||||
return nil, nil, nil, err
|
||||
}
|
||||
nameByID := make(map[string]string, len(users))
|
||||
for _, u := range users {
|
||||
name := strings.TrimSpace(u.Name)
|
||||
if name == "" {
|
||||
name = strings.TrimSpace(u.Email)
|
||||
}
|
||||
if name == "" {
|
||||
name = u.ID
|
||||
}
|
||||
nameByID[u.ID] = name
|
||||
}
|
||||
return items, stats, nameByID, nil
|
||||
}
|
||||
|
||||
func normalizeCDKType(t string) string {
|
||||
if strings.EqualFold(strings.TrimSpace(t), "marketing") {
|
||||
return "marketing"
|
||||
}
|
||||
return "normal"
|
||||
}
|
||||
|
||||
func (s *CDKService) Generate(ctx context.Context, amount, count int, note, cdkType string) ([]model.CDKCode, error) {
|
||||
if amount <= 0 {
|
||||
return nil, errors.New("金额必须大于 0")
|
||||
}
|
||||
if count < 1 {
|
||||
count = 1
|
||||
}
|
||||
if count > 500 {
|
||||
count = 500
|
||||
}
|
||||
|
||||
cdkType = normalizeCDKType(cdkType)
|
||||
// One batch id per generate call — marketing codes are "one per user per
|
||||
// batch", so codes created together must share it.
|
||||
batchID := randomUpper(20)
|
||||
items := make([]model.CDKCode, 0, count)
|
||||
for i := 0; i < count; i++ {
|
||||
items = append(items, model.CDKCode{
|
||||
Code: randomCDK(),
|
||||
Amount: amount,
|
||||
Status: "active",
|
||||
Type: cdkType,
|
||||
BatchID: batchID,
|
||||
Note: strings.TrimSpace(note),
|
||||
CreatedAt: time.Now(),
|
||||
UpdatedAt: time.Now(),
|
||||
})
|
||||
}
|
||||
if err := s.cdks.CreateBatch(ctx, items); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return items, nil
|
||||
}
|
||||
|
||||
func (s *CDKService) Delete(ctx context.Context, code string) error {
|
||||
rows, err := s.cdks.Delete(ctx, strings.TrimSpace(strings.ToUpper(code)))
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if rows == 0 {
|
||||
return ErrNotFound
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// DeleteBulk removes many CDK codes in one call (multi-select).
|
||||
func (s *CDKService) DeleteBulk(ctx context.Context, codes []string) (int, error) {
|
||||
seen := make(map[string]struct{}, len(codes))
|
||||
clean := make([]string, 0, len(codes))
|
||||
for _, code := range codes {
|
||||
code = strings.TrimSpace(strings.ToUpper(code))
|
||||
if code == "" {
|
||||
continue
|
||||
}
|
||||
if _, ok := seen[code]; ok {
|
||||
continue
|
||||
}
|
||||
seen[code] = struct{}{}
|
||||
clean = append(clean, code)
|
||||
}
|
||||
if len(clean) == 0 {
|
||||
return 0, nil
|
||||
}
|
||||
rows, err := s.cdks.DeleteByCodes(ctx, clean)
|
||||
return int(rows), err
|
||||
}
|
||||
|
||||
func (s *CDKService) Redeem(ctx context.Context, userID, code string) (map[string]any, error) {
|
||||
code = strings.TrimSpace(strings.ToUpper(code))
|
||||
if code == "" {
|
||||
return nil, errors.New("请输入兑换码")
|
||||
}
|
||||
|
||||
item, err := s.cdks.Redeem(ctx, code, userID)
|
||||
if err != nil {
|
||||
if err == gorm.ErrRecordNotFound {
|
||||
return nil, errors.New("兑换码无效")
|
||||
}
|
||||
if errors.Is(err, repo.ErrCDKBatchLimit) {
|
||||
return nil, errors.New("该营销活动每人限兑一次,你已兑换过本批次的兑换码")
|
||||
}
|
||||
if err == gorm.ErrDuplicatedKey {
|
||||
return nil, errors.New("兑换码已被使用")
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// Atomic, row-locked credit grant — never read-modify-write the balance, or a
|
||||
// concurrent debit/redeem would clobber it (lost update).
|
||||
updated, err := s.users.AdjustCredits(ctx, userID, float64(item.Amount))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return map[string]any{
|
||||
"amount": item.Amount,
|
||||
"credits": updated.Credits,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func randomCDK() string {
|
||||
seg := func() string {
|
||||
return randomUpper(4)
|
||||
}
|
||||
return seg() + "-" + seg() + "-" + seg() + "-" + seg()
|
||||
}
|
||||
@@ -0,0 +1,126 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/rand"
|
||||
"fmt"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/redis/go-redis/v9"
|
||||
)
|
||||
|
||||
// maxCodeAttempts caps wrong guesses per issued code before it's burned. With a
|
||||
// single 6-digit code (1e6 space) and only this many tries per send — and sends
|
||||
// throttled by the cooldown — brute force is infeasible. Mirrors the Python
|
||||
// EmailCodeStore.MAX_ATTEMPTS.
|
||||
const maxCodeAttempts = 5
|
||||
|
||||
type EmailCodeService struct {
|
||||
redis *redis.Client
|
||||
codeTTL time.Duration
|
||||
resendCooldown time.Duration
|
||||
}
|
||||
|
||||
func NewEmailCodeService(redis *redis.Client) *EmailCodeService {
|
||||
return &EmailCodeService{
|
||||
redis: redis,
|
||||
codeTTL: 6 * time.Minute, // CODE_TTL_SECONDS=360
|
||||
resendCooldown: 120 * time.Second, // CODE_COOLDOWN_SECONDS=120
|
||||
}
|
||||
}
|
||||
|
||||
// Redis exposes the underlying client so collaborators (e.g. LoginGuard) can be
|
||||
// built without threading the client through every constructor.
|
||||
func (s *EmailCodeService) Redis() *redis.Client {
|
||||
return s.redis
|
||||
}
|
||||
|
||||
func (s *EmailCodeService) Issue(ctx context.Context, email, purpose string) (string, error) {
|
||||
email = strings.ToLower(strings.TrimSpace(email))
|
||||
purpose = strings.ToLower(strings.TrimSpace(purpose))
|
||||
|
||||
ok, err := s.redis.SetNX(ctx, s.cooldownKey(email, purpose), "1", s.resendCooldown).Result()
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if !ok {
|
||||
return "", fmt.Errorf("请稍后再试")
|
||||
}
|
||||
|
||||
code, err := randomDigits(6)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if err := s.redis.Set(ctx, s.codeKey(email, purpose), code, s.codeTTL).Err(); err != nil {
|
||||
return "", err
|
||||
}
|
||||
// Reset the wrong-guess counter for this fresh code (same TTL as the code).
|
||||
if err := s.redis.Set(ctx, s.attemptsKey(email, purpose), "0", s.codeTTL).Err(); err != nil {
|
||||
return "", err
|
||||
}
|
||||
return code, nil
|
||||
}
|
||||
|
||||
func (s *EmailCodeService) Verify(ctx context.Context, email, purpose, code string) (bool, error) {
|
||||
normalizedCode, err := ValidateEmailCode(code)
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
email = strings.ToLower(strings.TrimSpace(email))
|
||||
purpose = strings.ToLower(strings.TrimSpace(purpose))
|
||||
|
||||
stored, err := s.redis.Get(ctx, s.codeKey(email, purpose)).Result()
|
||||
if err != nil {
|
||||
if err == redis.Nil {
|
||||
return false, nil
|
||||
}
|
||||
return false, err
|
||||
}
|
||||
|
||||
// Count the attempt first; burn the code (and its counter) once the cap is
|
||||
// hit so the attacker must request a new one and wait out the send cooldown.
|
||||
attempts, err := s.redis.Incr(ctx, s.attemptsKey(email, purpose)).Result()
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
if attempts > maxCodeAttempts {
|
||||
if err := s.redis.Del(ctx, s.codeKey(email, purpose), s.attemptsKey(email, purpose)).Err(); err != nil {
|
||||
return false, err
|
||||
}
|
||||
return false, nil
|
||||
}
|
||||
|
||||
if stored != normalizedCode {
|
||||
return false, nil
|
||||
}
|
||||
// Correct code: one-time use, clear both the code and its attempt counter.
|
||||
if err := s.redis.Del(ctx, s.codeKey(email, purpose), s.attemptsKey(email, purpose)).Err(); err != nil {
|
||||
return false, err
|
||||
}
|
||||
return true, nil
|
||||
}
|
||||
|
||||
func (s *EmailCodeService) codeKey(email, purpose string) string {
|
||||
return "email_code:" + purpose + ":" + email
|
||||
}
|
||||
|
||||
func (s *EmailCodeService) attemptsKey(email, purpose string) string {
|
||||
return "email_code_attempts:" + purpose + ":" + email
|
||||
}
|
||||
|
||||
func (s *EmailCodeService) cooldownKey(email, purpose string) string {
|
||||
return "email_code_cooldown:" + purpose + ":" + email
|
||||
}
|
||||
|
||||
func randomDigits(n int) (string, error) {
|
||||
buf := make([]byte, n)
|
||||
src := make([]byte, n)
|
||||
if _, err := rand.Read(src); err != nil {
|
||||
return "", err
|
||||
}
|
||||
for i := range src {
|
||||
buf[i] = byte('0' + (src[i] % 10))
|
||||
}
|
||||
return string(buf), nil
|
||||
}
|
||||
@@ -0,0 +1,48 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"strings"
|
||||
|
||||
"backend/internal/repo"
|
||||
)
|
||||
|
||||
type ImageAccessService struct {
|
||||
generatedRoot string
|
||||
showcase *repo.ShowcaseRepository
|
||||
auth *AuthService
|
||||
}
|
||||
|
||||
func NewImageAccessService(generatedRoot string, showcase *repo.ShowcaseRepository, auth *AuthService) *ImageAccessService {
|
||||
return &ImageAccessService{
|
||||
generatedRoot: generatedRoot,
|
||||
showcase: showcase,
|
||||
auth: auth,
|
||||
}
|
||||
}
|
||||
|
||||
// Resolve validates the path params and returns the object key (user/name).
|
||||
// Existence isn't checked here — that's the storage GET's job (404 if missing).
|
||||
func (s *ImageAccessService) Resolve(user, name string) (string, error) {
|
||||
user = strings.TrimSpace(user)
|
||||
name = strings.TrimSpace(name)
|
||||
if user == "" || name == "" {
|
||||
return "", errors.New("missing path params")
|
||||
}
|
||||
// :user and :name are single path segments (gin won't match "/"); guard
|
||||
// against traversal tokens anyway.
|
||||
if strings.Contains(user, "..") || strings.Contains(name, "..") ||
|
||||
strings.ContainsAny(user, `/\`) || strings.ContainsAny(name, `/\`) {
|
||||
return "", errors.New("invalid image path")
|
||||
}
|
||||
return user + "/" + name, nil
|
||||
}
|
||||
|
||||
func (s *ImageAccessService) IsPublic(ctx context.Context, rel string) (bool, error) {
|
||||
return s.showcase.IsPublicFile(ctx, rel)
|
||||
}
|
||||
|
||||
func (s *ImageAccessService) IsAuthorized(ctx context.Context, sessionCookie, owner string) (bool, error) {
|
||||
return s.auth.IsAuthorizedForPrivateImage(ctx, sessionCookie, owner)
|
||||
}
|
||||
@@ -0,0 +1,169 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/redis/go-redis/v9"
|
||||
)
|
||||
|
||||
// ErrLoginLocked is returned when a login/reset attempt is currently locked out
|
||||
// by the LoginGuard. The wait window (in seconds) is carried by LoginLockedError.
|
||||
var ErrLoginLocked = errors.New("login locked")
|
||||
|
||||
// LoginLockedError signals that the caller must wait RetryAfter seconds before
|
||||
// retrying. Handlers map this to HTTP 429 with a Retry-After header.
|
||||
type LoginLockedError struct {
|
||||
RetryAfter int
|
||||
}
|
||||
|
||||
func (e *LoginLockedError) Error() string {
|
||||
return "尝试过于频繁，请 " + strconv.Itoa(e.RetryAfter) + " 秒后再试"
|
||||
}
|
||||
|
||||
func (e *LoginLockedError) Is(target error) bool {
|
||||
return target == ErrLoginLocked
|
||||
}
|
||||
|
||||
// LoginGuard implements a Redis-backed login throttle mirroring the Python
|
||||
// core.login_guard: two independent counters per attempt, exponential backoff
|
||||
// lockout after a small number of free failures, and decay after a quiet period.
|
||||
//
|
||||
// id:<ip>|<identifier> — targeted guessing of one account from one IP (5 free).
|
||||
// ip:<ip> — spraying many accounts from one IP (20 free).
|
||||
//
|
||||
// Either counter being locked rejects the attempt.
|
||||
type LoginGuard struct {
|
||||
redis *redis.Client
|
||||
|
||||
freeAttempts int // per (ip, account) before lockout kicks in
|
||||
ipFreeAttempts int // coarser per-ip spray threshold
|
||||
baseLock time.Duration // first lock duration
|
||||
maxLock time.Duration // lock cap
|
||||
decay time.Duration // forget a counter after this quiet period
|
||||
}
|
||||
|
||||
func NewLoginGuard(rdb *redis.Client) *LoginGuard {
|
||||
return &LoginGuard{
|
||||
redis: rdb,
|
||||
freeAttempts: 5,
|
||||
ipFreeAttempts: 20,
|
||||
baseLock: 15 * time.Second,
|
||||
maxLock: 900 * time.Second,
|
||||
decay: 1800 * time.Second,
|
||||
}
|
||||
}
|
||||
|
||||
func (g *LoginGuard) keys(ip, identifier string) (ipKey, idKey string) {
|
||||
ident := strings.ToLower(strings.TrimSpace(identifier))
|
||||
return "login_guard:ip:" + ip, "login_guard:id:" + ip + "|" + ident
|
||||
}
|
||||
|
||||
// remaining returns the seconds the given counter is still locked for (0 = free).
|
||||
// Counters are stored with TTL = decay so quiet entries expire on their own,
|
||||
// matching the Python decay semantics.
|
||||
func (g *LoginGuard) remaining(ctx context.Context, key string, now int64) (int, error) {
|
||||
lockedRaw, err := g.redis.HGet(ctx, key, "locked_until").Result()
|
||||
if err != nil {
|
||||
if err == redis.Nil {
|
||||
return 0, nil
|
||||
}
|
||||
return 0, err
|
||||
}
|
||||
lockedUntil, _ := strconv.ParseInt(strings.TrimSpace(lockedRaw), 10, 64)
|
||||
if lockedUntil <= now {
|
||||
return 0, nil
|
||||
}
|
||||
return int(lockedUntil - now), nil
|
||||
}
|
||||
|
||||
// RetryAfter reports how many seconds the caller must wait (0 = allowed).
|
||||
func (g *LoginGuard) RetryAfter(ctx context.Context, ip, identifier string) (int, error) {
|
||||
if g == nil || g.redis == nil {
|
||||
return 0, nil
|
||||
}
|
||||
now := time.Now().Unix()
|
||||
ipKey, idKey := g.keys(ip, identifier)
|
||||
ipWait, err := g.remaining(ctx, ipKey, now)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
idWait, err := g.remaining(ctx, idKey, now)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
if ipWait > idWait {
|
||||
return ipWait, nil
|
||||
}
|
||||
return idWait, nil
|
||||
}
|
||||
|
||||
// Check returns a *LoginLockedError when the attempt is currently locked out.
|
||||
func (g *LoginGuard) Check(ctx context.Context, ip, identifier string) error {
|
||||
wait, err := g.RetryAfter(ctx, ip, identifier)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if wait > 0 {
|
||||
return &LoginLockedError{RetryAfter: wait}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// RecordFailure increments both counters and, once a counter passes its free
|
||||
// allowance, arms an exponentially growing lockout window (capped at maxLock).
|
||||
func (g *LoginGuard) RecordFailure(ctx context.Context, ip, identifier string) error {
|
||||
if g == nil || g.redis == nil {
|
||||
return nil
|
||||
}
|
||||
now := time.Now().Unix()
|
||||
ipKey, idKey := g.keys(ip, identifier)
|
||||
for _, kf := range []struct {
|
||||
key string
|
||||
free int
|
||||
}{
|
||||
{ipKey, g.ipFreeAttempts},
|
||||
{idKey, g.freeAttempts},
|
||||
} {
|
||||
count, err := g.redis.HIncrBy(ctx, kf.key, "count", 1).Result()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if count >= int64(kf.free) {
|
||||
over := count - int64(kf.free)
|
||||
lock := g.baseLock
|
||||
for i := int64(0); i < over; i++ {
|
||||
lock *= 2
|
||||
if lock >= g.maxLock {
|
||||
lock = g.maxLock
|
||||
break
|
||||
}
|
||||
}
|
||||
if lock > g.maxLock {
|
||||
lock = g.maxLock
|
||||
}
|
||||
lockedUntil := now + int64(lock.Seconds())
|
||||
if err := g.redis.HSet(ctx, kf.key, "locked_until", lockedUntil).Err(); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
// Refresh decay TTL on every failure (quiet counters expire on their own).
|
||||
if err := g.redis.Expire(ctx, kf.key, g.decay).Err(); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// RecordSuccess clears the targeted (id) counter on a genuine login; the coarse
|
||||
// per-ip counter is left to decay so one valid account can't reset spray tracking.
|
||||
func (g *LoginGuard) RecordSuccess(ctx context.Context, ip, identifier string) error {
|
||||
if g == nil || g.redis == nil {
|
||||
return nil
|
||||
}
|
||||
_, idKey := g.keys(ip, identifier)
|
||||
return g.redis.Del(ctx, idKey).Err()
|
||||
}
|
||||
@@ -0,0 +1,282 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"log"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"backend/internal/model"
|
||||
"backend/internal/repo"
|
||||
"backend/internal/storage"
|
||||
)
|
||||
|
||||
// MaintenanceService runs the periodic self-healing sweep that the Python
|
||||
// original did via a 60s daemon thread plus read-time lazy cleanup. Without it
|
||||
// the Go token pool only ever loses capacity: tokens never re-activate after a
|
||||
// quota reset, cookies never auto-renew, stale pending events permanently block
|
||||
// a user's generation gate, and old media/logs accumulate unbounded.
|
||||
type MaintenanceService struct {
|
||||
tokens *repo.TokenRepository
|
||||
tokenSvc *TokenService
|
||||
events *repo.EventRepository
|
||||
users *repo.UserRepository
|
||||
refresh *RefreshProfileService
|
||||
settings *repo.SiteSettingRepository
|
||||
store *storage.Client
|
||||
inflight *InflightRegistry
|
||||
showcase *repo.ShowcaseRepository
|
||||
interval time.Duration
|
||||
stalePending time.Duration
|
||||
mediaPruneEvery time.Duration
|
||||
lastMediaPrune time.Time
|
||||
}
|
||||
|
||||
func NewMaintenanceService(tokens *repo.TokenRepository, tokenSvc *TokenService, events *repo.EventRepository, users *repo.UserRepository, refresh *RefreshProfileService, settings *repo.SiteSettingRepository, store *storage.Client, inflight *InflightRegistry, showcase *repo.ShowcaseRepository) *MaintenanceService {
|
||||
return &MaintenanceService{
|
||||
tokens: tokens,
|
||||
tokenSvc: tokenSvc,
|
||||
events: events,
|
||||
users: users,
|
||||
refresh: refresh,
|
||||
settings: settings,
|
||||
store: store,
|
||||
inflight: inflight,
|
||||
showcase: showcase,
|
||||
interval: 60 * time.Second,
|
||||
stalePending: 600 * time.Second,
|
||||
mediaPruneEvery: 60 * time.Second,
|
||||
}
|
||||
}
|
||||
|
||||
// Run drives the sweep every interval until ctx is cancelled. It runs one sweep
|
||||
// immediately on startup so a freshly restarted process heals stuck state right
|
||||
// away rather than after the first tick.
|
||||
func (m *MaintenanceService) Run(ctx context.Context) {
|
||||
ticker := time.NewTicker(m.interval)
|
||||
defer ticker.Stop()
|
||||
m.tick(ctx)
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case <-ticker.C:
|
||||
m.tick(ctx)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// syncRecoveredQuota re-probes each just-recovered account so its displayed
|
||||
// balance reflects the post-reset value (these providers only sync quota when
|
||||
// accessed). krea additionally needs /app (Activate) to actually grant the daily
|
||||
// free balance before billing-data reports it. Bounded concurrency avoids a
|
||||
// thundering herd at the daily reset.
|
||||
func (m *MaintenanceService) syncRecoveredQuota(accs []model.TokenAccount) {
|
||||
sem := make(chan struct{}, 4)
|
||||
var wg sync.WaitGroup
|
||||
for _, acc := range accs {
|
||||
switch acc.Pool {
|
||||
case "chatgpt", "leonardo", "krea", "imagine":
|
||||
default:
|
||||
continue
|
||||
}
|
||||
wg.Add(1)
|
||||
sem <- struct{}{}
|
||||
go func(a model.TokenAccount) {
|
||||
defer wg.Done()
|
||||
defer func() { <-sem }()
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 90*time.Second)
|
||||
defer cancel()
|
||||
if a.Pool == "krea" && m.tokenSvc.krea != nil {
|
||||
m.tokenSvc.krea.Activate(ctx, a.Value)
|
||||
}
|
||||
_, _ = m.tokenSvc.Quota(ctx, a.Pool, a.ID)
|
||||
}(acc)
|
||||
}
|
||||
wg.Wait()
|
||||
}
|
||||
|
||||
func (m *MaintenanceService) tick(ctx context.Context) {
|
||||
// 1. Re-activate quota-exhausted tokens whose reset time has passed, then
|
||||
// auto-sync their real balance — these providers only refresh quota when
|
||||
// accessed, so recovery alone would leave a stale 0/—. For krea the sync
|
||||
// must first load /app (Activate) to grant the daily free balance.
|
||||
if recovered, err := m.tokens.RecoverQuota(ctx); err != nil {
|
||||
log.Printf("maintenance: recover_quota: %v", err)
|
||||
} else if len(recovered) > 0 {
|
||||
log.Printf("maintenance: recovered %d quota token(s)", len(recovered))
|
||||
if m.tokenSvc != nil {
|
||||
go m.syncRecoveredQuota(recovered)
|
||||
}
|
||||
}
|
||||
|
||||
// 1a. Roll the 恢复时间 marker of ACTIVE daily-reset accounts forward to the next
|
||||
// future reset (same time-of-day, +1 day) so the column never shows a stale
|
||||
// past time. Limited accounts are intentionally skipped (RecoverQuota owns
|
||||
// their marker). adobe/leonardo/krea/imagine all renew daily.
|
||||
if _, err := m.tokens.RollResetMarkers(ctx, []string{"adobe", "leonardo", "krea", "imagine"}); err != nil {
|
||||
log.Printf("maintenance: roll_reset: %v", err)
|
||||
}
|
||||
|
||||
// 1b. Runway tokens have no refresh — once the JWT expiry (its reset marker)
|
||||
// passes, mark them dead directly instead of letting them 401 on next use.
|
||||
if n, err := m.tokens.ExpireByReset(ctx, "runway"); err != nil {
|
||||
log.Printf("maintenance: expire_runway: %v", err)
|
||||
} else if n > 0 {
|
||||
log.Printf("maintenance: expired %d runway token(s)", n)
|
||||
}
|
||||
|
||||
// 1c. Proactively renew krea/imagine sessions ~10min before expiry so a
|
||||
// dormant account's rotating refresh_token never lapses (a dead token
|
||||
// can't be recovered and, for krea, blocks the daily free-credit meter
|
||||
// from being re-created). Only near-expiry accounts hit the network.
|
||||
if m.tokenSvc != nil {
|
||||
m.tokenSvc.RefreshExpiringTokens(ctx)
|
||||
// 1d. Once-per-day krea /app activation for accounts not yet synced since the
|
||||
// daily reset — krea only grants the free balance after /app loads, so an
|
||||
// always-active account (never went 限额) would otherwise read 0 / 402
|
||||
// after each reset. Self-guarded + background; no-op once all are done.
|
||||
m.tokenSvc.ActivateKreaDue(ctx)
|
||||
}
|
||||
|
||||
// 2. Auto-renew Adobe cookies whose refresh interval has elapsed.
|
||||
if m.refresh != nil {
|
||||
if n, err := m.refresh.RefreshDue(ctx); err != nil {
|
||||
log.Printf("maintenance: refresh_due: %v", err)
|
||||
} else if n > 0 {
|
||||
log.Printf("maintenance: refreshed %d cookie profile(s)", n)
|
||||
}
|
||||
}
|
||||
|
||||
// 3. Fail long-pending events so they stop blocking the per-user gate, and
|
||||
// refund the credits debited up-front for each abandoned generation (the
|
||||
// normal failure-refund path never ran for a process-restart orphan).
|
||||
if purged, err := m.events.PurgeStale(ctx, m.stalePending); err != nil {
|
||||
log.Printf("maintenance: purge_stale: %v", err)
|
||||
} else if len(purged) > 0 {
|
||||
refunded := 0
|
||||
cancelled := 0
|
||||
for _, e := range purged {
|
||||
// Stop the generation goroutine if it's still running, so it doesn't
|
||||
// keep grinding for minutes and surface a late "success" on this
|
||||
// just-abandoned event.
|
||||
if m.inflight != nil && m.inflight.Cancel(e.ID) {
|
||||
cancelled++
|
||||
}
|
||||
// Attribute the abandoned failure back to the account it was using
|
||||
// (the normal markTokenFailure path never ran for an orphaned job).
|
||||
if e.AccountID != "" {
|
||||
if err := m.tokens.IncrementFail(ctx, e.AccountID); err != nil {
|
||||
log.Printf("maintenance: fail-count abandoned event %s (account %s): %v", e.ID, e.AccountID, err)
|
||||
}
|
||||
}
|
||||
if e.UserID == "" || e.Cost <= 0 {
|
||||
continue
|
||||
}
|
||||
// Exactly-once: only refund if we win the claim (the in-flight request
|
||||
// may have already refunded itself on its own failure path).
|
||||
claimed, err := m.events.MarkRefunded(ctx, e.ID)
|
||||
if err != nil {
|
||||
log.Printf("maintenance: claim refund %s: %v", e.ID, err)
|
||||
continue
|
||||
}
|
||||
if !claimed {
|
||||
continue
|
||||
}
|
||||
if _, err := m.users.AdjustCredits(ctx, e.UserID, e.Cost); err != nil {
|
||||
log.Printf("maintenance: refund abandoned event %s (user %s, %.0f): %v", e.ID, e.UserID, e.Cost, err)
|
||||
} else {
|
||||
refunded++
|
||||
}
|
||||
}
|
||||
log.Printf("maintenance: marked %d stale pending event(s) failed, refunded %d, cancelled %d in-flight", len(purged), refunded, cancelled)
|
||||
}
|
||||
|
||||
// 4. Enforce the admin-configured log retention window.
|
||||
m.pruneLogs(ctx)
|
||||
|
||||
// 5. Enforce the media retention window. Runs every 60s like the log prune;
|
||||
// mediaPruneEvery still gates it in case the interval is ever shortened.
|
||||
if time.Since(m.lastMediaPrune) >= m.mediaPruneEvery {
|
||||
m.pruneMedia(ctx)
|
||||
m.lastMediaPrune = time.Now()
|
||||
}
|
||||
}
|
||||
|
||||
func (m *MaintenanceService) pruneLogs(ctx context.Context) {
|
||||
days := m.retentionDays(ctx, "logs.retention_days")
|
||||
if days <= 0 {
|
||||
return
|
||||
}
|
||||
if _, err := m.events.PurgeOlderThan(ctx, time.Duration(days)*24*time.Hour); err != nil {
|
||||
log.Printf("maintenance: purge_older_than: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func (m *MaintenanceService) pruneMedia(ctx context.Context) {
|
||||
if m.store == nil || !m.store.Configured() {
|
||||
return
|
||||
}
|
||||
days := m.retentionDays(ctx, "media.retention_days")
|
||||
if days <= 0 {
|
||||
return
|
||||
}
|
||||
cutoff := time.Now().Add(-time.Duration(days) * 24 * time.Hour)
|
||||
objs, err := m.store.List(ctx, "")
|
||||
if err != nil {
|
||||
log.Printf("maintenance: list media: %v", err)
|
||||
return
|
||||
}
|
||||
// Files referenced by the homepage showcase are kept forever, no matter how
|
||||
// old — deleting them would break the public landing page.
|
||||
var pinned map[string]struct{}
|
||||
if m.showcase != nil {
|
||||
if pinned, err = m.showcase.PublicFileSet(ctx); err != nil {
|
||||
log.Printf("maintenance: showcase file set: %v", err)
|
||||
pinned = nil
|
||||
}
|
||||
}
|
||||
removed, skipped := 0, 0
|
||||
var clearedKeys []string
|
||||
for _, o := range objs {
|
||||
if !o.LastModified.Before(cutoff) {
|
||||
continue
|
||||
}
|
||||
if _, ok := pinned[strings.TrimLeft(o.Key, "/")]; ok {
|
||||
skipped++
|
||||
continue
|
||||
}
|
||||
if err := m.store.Delete(ctx, o.Key); err != nil {
|
||||
log.Printf("maintenance: delete %s: %v", o.Key, err)
|
||||
continue
|
||||
}
|
||||
removed++
|
||||
// event_log.file stores the same key — blank those rows so the log views
|
||||
// don't dangle a 404 preview.
|
||||
clearedKeys = append(clearedKeys, o.Key)
|
||||
}
|
||||
if removed > 0 || skipped > 0 {
|
||||
log.Printf("maintenance: pruned %d expired media object(s), kept %d showcase-pinned", removed, skipped)
|
||||
}
|
||||
if len(clearedKeys) > 0 {
|
||||
if n, err := m.events.ClearFiles(ctx, clearedKeys); err != nil {
|
||||
log.Printf("maintenance: clear_files: %v", err)
|
||||
} else if n > 0 {
|
||||
log.Printf("maintenance: cleared file ref on %d log row(s)", n)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (m *MaintenanceService) retentionDays(ctx context.Context, key string) int {
|
||||
raw, err := m.settings.GetValue(ctx, key)
|
||||
if err != nil {
|
||||
return 0
|
||||
}
|
||||
days, err := strconv.Atoi(strings.TrimSpace(raw))
|
||||
if err != nil || days <= 0 {
|
||||
return 0
|
||||
}
|
||||
return days
|
||||
}
|
||||
@@ -0,0 +1,46 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"crypto/sha256"
|
||||
"crypto/subtle"
|
||||
"encoding/base64"
|
||||
"encoding/hex"
|
||||
"strings"
|
||||
|
||||
"golang.org/x/crypto/bcrypt"
|
||||
)
|
||||
|
||||
func GeneratePasswordHash(password string) (string, error) {
|
||||
hash, err := bcrypt.GenerateFromPassword(bcryptPrehash(password), 12)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
return string(hash), nil
|
||||
}
|
||||
|
||||
func VerifyPassword(password, stored string) bool {
|
||||
stored = strings.TrimSpace(stored)
|
||||
if stored == "" {
|
||||
return false
|
||||
}
|
||||
|
||||
if strings.HasPrefix(stored, "bcrypt$") {
|
||||
hash := stored[len("bcrypt$"):]
|
||||
return bcrypt.CompareHashAndPassword([]byte(hash), bcryptPrehash(password)) == nil
|
||||
}
|
||||
|
||||
parts := strings.SplitN(stored, "$", 3)
|
||||
if len(parts) != 3 || parts[0] != "sha256" {
|
||||
return false
|
||||
}
|
||||
expected := sha256.Sum256([]byte(parts[1] + password))
|
||||
expectedHex := hex.EncodeToString(expected[:])
|
||||
return subtle.ConstantTimeCompare([]byte(expectedHex), []byte(parts[2])) == 1
|
||||
}
|
||||
|
||||
func bcryptPrehash(password string) []byte {
|
||||
sum := sha256.Sum256([]byte(password))
|
||||
dst := make([]byte, base64.StdEncoding.EncodedLen(len(sum)))
|
||||
base64.StdEncoding.Encode(dst, sum[:])
|
||||
return dst
|
||||
}
|
||||
@@ -0,0 +1,79 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/redis/go-redis/v9"
|
||||
)
|
||||
|
||||
var ErrRateLimited = errors.New("rate limited")
|
||||
|
||||
type RateLimitService struct {
|
||||
redis *redis.Client
|
||||
prefix string
|
||||
}
|
||||
|
||||
type RateLimitResult struct {
|
||||
Allowed bool
|
||||
Count int64
|
||||
Limit int64
|
||||
RetryAfter time.Duration
|
||||
}
|
||||
|
||||
func NewRateLimitService(redis *redis.Client) *RateLimitService {
|
||||
return &RateLimitService{
|
||||
redis: redis,
|
||||
prefix: "rl:",
|
||||
}
|
||||
}
|
||||
|
||||
func (s *RateLimitService) Allow(ctx context.Context, bucket string, limit int64, window time.Duration) (*RateLimitResult, error) {
|
||||
if limit <= 0 || window <= 0 {
|
||||
return &RateLimitResult{Allowed: true, Limit: limit}, nil
|
||||
}
|
||||
|
||||
key := s.prefix + strings.TrimSpace(bucket)
|
||||
count, err := s.redis.Incr(ctx, key).Result()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if count == 1 {
|
||||
if err := s.redis.Expire(ctx, key, window).Err(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
|
||||
ttl, err := s.redis.TTL(ctx, key).Result()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if ttl < 0 {
|
||||
ttl = window
|
||||
}
|
||||
|
||||
return &RateLimitResult{
|
||||
Allowed: count <= limit,
|
||||
Count: count,
|
||||
Limit: limit,
|
||||
RetryAfter: ttl,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (s *RateLimitService) Enforce(ctx context.Context, bucket string, limit int64, window time.Duration) error {
|
||||
result, err := s.Allow(ctx, bucket, limit, window)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if result.Allowed {
|
||||
return nil
|
||||
}
|
||||
retry := int(result.RetryAfter.Seconds())
|
||||
if retry < 1 {
|
||||
retry = 1
|
||||
}
|
||||
return fmt.Errorf("%w: 请稍后再试（%d 秒后）", ErrRateLimited, retry)
|
||||
}
|
||||
@@ -0,0 +1,185 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"backend/internal/model"
|
||||
"backend/internal/provider/adobe"
|
||||
"backend/internal/repo"
|
||||
"gorm.io/datatypes"
|
||||
)
|
||||
|
||||
type RefreshProfileService struct {
|
||||
profiles *repo.RefreshProfileRepository
|
||||
tokens *repo.TokenRepository
|
||||
adobe *adobe.Client
|
||||
}
|
||||
|
||||
func NewRefreshProfileService(profiles *repo.RefreshProfileRepository, tokens *repo.TokenRepository, adobeClient *adobe.Client) *RefreshProfileService {
|
||||
return &RefreshProfileService{
|
||||
profiles: profiles,
|
||||
tokens: tokens,
|
||||
adobe: adobeClient,
|
||||
}
|
||||
}
|
||||
|
||||
func (s *RefreshProfileService) List(ctx context.Context) ([]model.RefreshProfile, error) {
|
||||
return s.profiles.List(ctx)
|
||||
}
|
||||
|
||||
func (s *RefreshProfileService) RefreshNow(ctx context.Context, id string) error {
|
||||
if s.adobe == nil || s.tokens == nil {
|
||||
return errors.New("refresh client not configured")
|
||||
}
|
||||
profile, err := s.profiles.Get(ctx, id)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if profile.Pool != "adobe" || profile.Kind != "adobe_cookie" {
|
||||
return errors.New("unsupported refresh profile")
|
||||
}
|
||||
|
||||
now := time.Now()
|
||||
_, _ = s.profiles.Update(ctx, id, map[string]any{
|
||||
"last_attempt_at": now,
|
||||
})
|
||||
|
||||
result, err := s.adobe.ExchangeCookie(ctx, profile.Cookie)
|
||||
if err != nil {
|
||||
failures := profile.ConsecutiveFailures + 1
|
||||
// Exponential backoff: 60s per consecutive failure, capped at 1h.
|
||||
secs := 60 * failures
|
||||
if secs > 3600 {
|
||||
secs = 3600
|
||||
}
|
||||
msg := err.Error()
|
||||
if len(msg) > 300 {
|
||||
msg = msg[:300]
|
||||
}
|
||||
_, _ = s.profiles.Update(ctx, id, map[string]any{
|
||||
"last_error": msg,
|
||||
"consecutive_failures": failures,
|
||||
"next_retry_at": now.Add(time.Duration(secs) * time.Second),
|
||||
})
|
||||
// After repeated failures the cookie can no longer mint a token — it's
|
||||
// genuinely dead (expired/revoked). Lock the pool token (disabled+dead)
|
||||
// so the UI flags it red. A single failure may be a transient blip, so
|
||||
// only escalate after a few in a row (mirrors Python RefreshManager).
|
||||
if failures >= 3 {
|
||||
_, _ = s.tokens.Update(ctx, profile.Pool, id, map[string]any{
|
||||
"status": "disabled",
|
||||
"dead": true,
|
||||
})
|
||||
}
|
||||
return err
|
||||
}
|
||||
|
||||
tokenPatch := map[string]any{
|
||||
"value": result.AccessToken,
|
||||
"status": "active",
|
||||
"dead": false,
|
||||
"fails": 0,
|
||||
"updated_at": now,
|
||||
}
|
||||
email, exp := parseJWTEmailExpiry(result.AccessToken)
|
||||
if email != "" {
|
||||
tokenPatch["account_email"] = email
|
||||
}
|
||||
if exp != nil {
|
||||
tokenPatch["cached_quota_reset_after"] = exp.Format(time.RFC3339)
|
||||
}
|
||||
if profileData, profileErr := s.adobe.FetchAccountProfile(ctx, result.AccessToken); profileErr == nil {
|
||||
if email := strings.TrimSpace(stringValue(profileData["email"])); email != "" {
|
||||
tokenPatch["account_email"] = email
|
||||
}
|
||||
if displayName := strings.TrimSpace(stringValue(profileData["display_name"])); displayName != "" {
|
||||
tokenPatch["account_display_name"] = displayName
|
||||
}
|
||||
}
|
||||
if quotaData, quotaErr := s.adobe.FetchCreditsBalance(ctx, result.AccessToken); quotaErr == nil {
|
||||
meta := datatypes.JSONMap{
|
||||
"cached_quota_at": int(time.Now().Unix()),
|
||||
}
|
||||
if remaining, ok := quotaData["remaining"].(int); ok {
|
||||
meta["cached_quota_remaining"] = remaining
|
||||
}
|
||||
if used, ok := quotaData["used"].(int); ok {
|
||||
meta["cached_quota_used"] = used
|
||||
}
|
||||
if total, ok := quotaData["total"].(int); ok {
|
||||
meta["cached_quota_total"] = total
|
||||
}
|
||||
tokenPatch["meta"] = meta
|
||||
if resetAfter := strings.TrimSpace(stringValue(quotaData["available_until"])); resetAfter != "" {
|
||||
tokenPatch["cached_quota_reset_after"] = resetAfter
|
||||
}
|
||||
}
|
||||
if _, err := s.tokens.Update(ctx, "adobe", id, tokenPatch); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
interval := profile.IntervalSeconds
|
||||
if interval <= 0 {
|
||||
interval = 54000
|
||||
}
|
||||
_, err = s.profiles.Update(ctx, id, map[string]any{
|
||||
"last_success_at": now,
|
||||
"next_retry_at": now.Add(time.Duration(interval) * time.Second),
|
||||
"last_error": "",
|
||||
"consecutive_failures": 0,
|
||||
})
|
||||
return err
|
||||
}
|
||||
|
||||
// RefreshDue refreshes every enabled profile whose next_retry_at has passed.
|
||||
// Driven by the background maintenance loop so Adobe cookies auto-renew without
|
||||
// an admin clicking "refresh". Individual failures are recorded on the profile
|
||||
// (backoff + dead escalation) and don't abort the sweep.
|
||||
func (s *RefreshProfileService) RefreshDue(ctx context.Context) (int, error) {
|
||||
if s.adobe == nil || s.tokens == nil {
|
||||
return 0, nil
|
||||
}
|
||||
due, err := s.profiles.ListDue(ctx, time.Now())
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
refreshed := 0
|
||||
for _, p := range due {
|
||||
if p.Pool != "adobe" || p.Kind != "adobe_cookie" {
|
||||
continue
|
||||
}
|
||||
if err := s.RefreshNow(ctx, p.ID); err != nil {
|
||||
continue
|
||||
}
|
||||
refreshed++
|
||||
}
|
||||
return refreshed, nil
|
||||
}
|
||||
|
||||
func (s *RefreshProfileService) Update(ctx context.Context, id string, body map[string]any) (*model.RefreshProfile, error) {
|
||||
patch := map[string]any{}
|
||||
if raw, ok := body["enabled"]; ok {
|
||||
patch["enabled"] = boolValueWithDefault(raw, false)
|
||||
}
|
||||
if raw, ok := body["name"]; ok {
|
||||
patch["name"] = stringValue(raw)
|
||||
}
|
||||
if raw, ok := body["interval_seconds"]; ok {
|
||||
n := intValue(raw)
|
||||
if n <= 0 {
|
||||
return nil, errors.New("interval_seconds must be positive")
|
||||
}
|
||||
patch["interval_seconds"] = n
|
||||
}
|
||||
if len(patch) == 0 {
|
||||
return s.profiles.Get(ctx, id)
|
||||
}
|
||||
return s.profiles.Update(ctx, id, patch)
|
||||
}
|
||||
|
||||
func (s *RefreshProfileService) Delete(ctx context.Context, id string) error {
|
||||
return s.profiles.Delete(ctx, id)
|
||||
}
|
||||
@@ -0,0 +1,103 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"time"
|
||||
|
||||
"github.com/redis/go-redis/v9"
|
||||
)
|
||||
|
||||
type SessionPayload struct {
|
||||
UserID string `json:"user_id"`
|
||||
ExpiresAt int64 `json:"expires_at"`
|
||||
}
|
||||
|
||||
type SessionService struct {
|
||||
client *redis.Client
|
||||
prefix string
|
||||
ttl time.Duration
|
||||
slideAfter time.Duration
|
||||
slideTo time.Duration
|
||||
}
|
||||
|
||||
func NewSessionService(client *redis.Client, ttl, slideAfter time.Duration) *SessionService {
|
||||
return &SessionService{
|
||||
client: client,
|
||||
prefix: "session:",
|
||||
ttl: ttl,
|
||||
slideAfter: slideAfter,
|
||||
slideTo: ttl,
|
||||
}
|
||||
}
|
||||
|
||||
func (s *SessionService) Create(ctx context.Context, userID string) (string, *SessionPayload, error) {
|
||||
token := randomUpper(48)
|
||||
|
||||
payload := &SessionPayload{
|
||||
UserID: userID,
|
||||
ExpiresAt: time.Now().Add(s.ttl).Unix(),
|
||||
}
|
||||
|
||||
raw, err := json.Marshal(payload)
|
||||
if err != nil {
|
||||
return "", nil, err
|
||||
}
|
||||
|
||||
if err := s.client.Set(ctx, s.key(token), raw, s.ttl).Err(); err != nil {
|
||||
return "", nil, err
|
||||
}
|
||||
return token, payload, nil
|
||||
}
|
||||
|
||||
func (s *SessionService) Validate(ctx context.Context, token string) (*SessionPayload, error) {
|
||||
if token == "" {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
raw, err := s.client.Get(ctx, s.key(token)).Bytes()
|
||||
if err != nil {
|
||||
if errors.Is(err, redis.Nil) {
|
||||
return nil, nil
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
|
||||
var payload SessionPayload
|
||||
if err := json.Unmarshal(raw, &payload); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
ttl, err := s.client.TTL(ctx, s.key(token)).Result()
|
||||
if err == nil && ttl > 0 && ttl < s.slideAfter {
|
||||
// Slide the expiry, but only update the in-memory payload after Redis
|
||||
// has actually persisted it — otherwise a failed Set would leave the
|
||||
// returned ExpiresAt out of sync with what's stored.
|
||||
renewed := payload
|
||||
renewed.ExpiresAt = time.Now().Add(s.slideTo).Unix()
|
||||
if updated, marshalErr := json.Marshal(&renewed); marshalErr == nil {
|
||||
if setErr := s.client.Set(ctx, s.key(token), updated, s.slideTo).Err(); setErr == nil {
|
||||
payload.ExpiresAt = renewed.ExpiresAt
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if payload.ExpiresAt <= time.Now().Unix() {
|
||||
_ = s.Destroy(ctx, token)
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
return &payload, nil
|
||||
}
|
||||
|
||||
func (s *SessionService) Destroy(ctx context.Context, token string) error {
|
||||
if token == "" {
|
||||
return nil
|
||||
}
|
||||
return s.client.Del(ctx, s.key(token)).Err()
|
||||
}
|
||||
|
||||
func (s *SessionService) key(token string) string {
|
||||
return s.prefix + token
|
||||
}
|
||||
@@ -0,0 +1,20 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
|
||||
"backend/internal/model"
|
||||
"backend/internal/repo"
|
||||
)
|
||||
|
||||
type ShowcaseService struct {
|
||||
repo *repo.ShowcaseRepository
|
||||
}
|
||||
|
||||
func NewShowcaseService(repo *repo.ShowcaseRepository) *ShowcaseService {
|
||||
return &ShowcaseService{repo: repo}
|
||||
}
|
||||
|
||||
func (s *ShowcaseService) Grouped(ctx context.Context) (map[string][]model.ShowcaseItem, error) {
|
||||
return s.repo.Grouped(ctx)
|
||||
}
|
||||
@@ -0,0 +1,81 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
|
||||
"backend/internal/repo"
|
||||
)
|
||||
|
||||
type SiteService struct {
|
||||
settings *repo.SiteSettingRepository
|
||||
fallback string
|
||||
}
|
||||
|
||||
func NewSiteService(settings *repo.SiteSettingRepository, fallback string) *SiteService {
|
||||
return &SiteService{
|
||||
settings: settings,
|
||||
fallback: fallback,
|
||||
}
|
||||
}
|
||||
|
||||
func (s *SiteService) Title(ctx context.Context) (string, error) {
|
||||
v, err := s.settings.GetValue(ctx, "site.title")
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
v = strings.TrimSpace(v)
|
||||
if v == "" {
|
||||
return s.fallback, nil
|
||||
}
|
||||
return v, nil
|
||||
}
|
||||
|
||||
func (s *SiteService) SetTitle(ctx context.Context, title string) (string, error) {
|
||||
title = strings.TrimSpace(title)
|
||||
if title == "" {
|
||||
return "", nil
|
||||
}
|
||||
if err := s.settings.UpsertValue(ctx, "site.title", title); err != nil {
|
||||
return "", err
|
||||
}
|
||||
return title, nil
|
||||
}
|
||||
|
||||
// Contact is the admin-editable "联系我们" info shown in the public 关于 section.
|
||||
type Contact struct {
|
||||
QQ string `json:"qq"`
|
||||
QQLink string `json:"qq_link"`
|
||||
QQGroup string `json:"qq_group"`
|
||||
QQGroupLink string `json:"qq_group_link"`
|
||||
Email string `json:"email"`
|
||||
Shop string `json:"shop"`
|
||||
}
|
||||
|
||||
func (s *SiteService) Contact(ctx context.Context) Contact {
|
||||
get := func(k string) string { v, _ := s.settings.GetValue(ctx, k); return strings.TrimSpace(v) }
|
||||
return Contact{
|
||||
QQ: get("contact.qq"),
|
||||
QQLink: get("contact.qq_link"),
|
||||
QQGroup: get("contact.qq_group"),
|
||||
QQGroupLink: get("contact.qq_group_link"),
|
||||
Email: get("contact.email"),
|
||||
Shop: get("contact.shop"),
|
||||
}
|
||||
}
|
||||
|
||||
func (s *SiteService) SetContact(ctx context.Context, c Contact) error {
|
||||
for k, v := range map[string]string{
|
||||
"contact.qq": strings.TrimSpace(c.QQ),
|
||||
"contact.qq_link": strings.TrimSpace(c.QQLink),
|
||||
"contact.qq_group": strings.TrimSpace(c.QQGroup),
|
||||
"contact.qq_group_link": strings.TrimSpace(c.QQGroupLink),
|
||||
"contact.email": strings.TrimSpace(c.Email),
|
||||
"contact.shop": strings.TrimSpace(c.Shop),
|
||||
} {
|
||||
if err := s.settings.UpsertValue(ctx, k, v); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,117 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/tls"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net"
|
||||
"net/smtp"
|
||||
"strconv"
|
||||
"strings"
|
||||
)
|
||||
|
||||
type SMTPConfig struct {
|
||||
Host string
|
||||
Port int
|
||||
Username string
|
||||
Password string
|
||||
FromAddr string
|
||||
UseTLS bool
|
||||
}
|
||||
|
||||
type SMTPService struct{}
|
||||
|
||||
func NewSMTPService() *SMTPService {
|
||||
return &SMTPService{}
|
||||
}
|
||||
|
||||
func (s *SMTPService) SendCode(ctx context.Context, cfg SMTPConfig, to, code, purpose string) error {
|
||||
_ = ctx
|
||||
if strings.TrimSpace(cfg.Host) == "" || cfg.Port <= 0 || strings.TrimSpace(cfg.FromAddr) == "" {
|
||||
return errors.New("SMTP 未配置")
|
||||
}
|
||||
action := "注册"
|
||||
if purpose == "reset" {
|
||||
action = "找回密码"
|
||||
}
|
||||
subject := "Vivid AI 邮箱验证码"
|
||||
body := fmt.Sprintf("你正在进行%s，验证码为：%s\n\n验证码 6 分钟内有效。", action, code)
|
||||
msg := buildSMTPMessage(cfg.FromAddr, to, subject, body)
|
||||
addr := net.JoinHostPort(cfg.Host, strconv.Itoa(cfg.Port))
|
||||
|
||||
if cfg.UseTLS || cfg.Port == 465 {
|
||||
return sendMailTLS(addr, cfg, to, msg)
|
||||
}
|
||||
return sendMailSTARTTLS(addr, cfg, to, msg)
|
||||
}
|
||||
|
||||
func buildSMTPMessage(from, to, subject, body string) []byte {
|
||||
lines := []string{
|
||||
"From: " + from,
|
||||
"To: " + to,
|
||||
"Subject: " + subject,
|
||||
"MIME-Version: 1.0",
|
||||
"Content-Type: text/plain; charset=UTF-8",
|
||||
"",
|
||||
body,
|
||||
}
|
||||
return []byte(strings.Join(lines, "\r\n"))
|
||||
}
|
||||
|
||||
func sendMailTLS(addr string, cfg SMTPConfig, to string, msg []byte) error {
|
||||
conn, err := tls.Dial("tcp", addr, &tls.Config{ServerName: cfg.Host})
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer conn.Close()
|
||||
|
||||
client, err := smtp.NewClient(conn, cfg.Host)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer client.Close()
|
||||
return doSMTP(client, cfg, to, msg)
|
||||
}
|
||||
|
||||
func sendMailSTARTTLS(addr string, cfg SMTPConfig, to string, msg []byte) error {
|
||||
client, err := smtp.Dial(addr)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer client.Close()
|
||||
|
||||
if ok, _ := client.Extension("STARTTLS"); ok {
|
||||
if err := client.StartTLS(&tls.Config{ServerName: cfg.Host}); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return doSMTP(client, cfg, to, msg)
|
||||
}
|
||||
|
||||
func doSMTP(client *smtp.Client, cfg SMTPConfig, to string, msg []byte) error {
|
||||
if strings.TrimSpace(cfg.Username) != "" {
|
||||
auth := smtp.PlainAuth("", cfg.Username, cfg.Password, cfg.Host)
|
||||
if err := client.Auth(auth); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
if err := client.Mail(cfg.FromAddr); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := client.Rcpt(to); err != nil {
|
||||
return err
|
||||
}
|
||||
w, err := client.Data()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if _, err := w.Write(msg); err != nil {
|
||||
_ = w.Close()
|
||||
return err
|
||||
}
|
||||
if err := w.Close(); err != nil {
|
||||
return err
|
||||
}
|
||||
return client.Quit()
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,200 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"strings"
|
||||
|
||||
"backend/internal/model"
|
||||
"backend/internal/repo"
|
||||
)
|
||||
|
||||
type UserGenerationService struct {
|
||||
v1 *V1Service
|
||||
events *repo.EventRepository
|
||||
users *repo.UserRepository
|
||||
models *repo.ModelRepository
|
||||
}
|
||||
|
||||
func NewUserGenerationService(v1 *V1Service, events *repo.EventRepository, users *repo.UserRepository, models *repo.ModelRepository) *UserGenerationService {
|
||||
return &UserGenerationService{
|
||||
v1: v1,
|
||||
events: events,
|
||||
users: users,
|
||||
models: models,
|
||||
}
|
||||
}
|
||||
|
||||
type UserGenerateRequest struct {
|
||||
Model string
|
||||
Prompt string
|
||||
Ratio string
|
||||
Resolution string
|
||||
Duration string
|
||||
ReferenceImages []string
|
||||
}
|
||||
|
||||
func (s *UserGenerationService) Generate(ctx context.Context, user *model.User, in UserGenerateRequest) (map[string]any, error) {
|
||||
if user == nil || strings.TrimSpace(user.ID) == "" {
|
||||
return nil, errors.New("未登录或会话已过期")
|
||||
}
|
||||
pending, err := s.events.PendingByUser(ctx, user.ID, "user")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if pending != nil {
|
||||
return nil, errors.New("已有正在生成的任务,请稍候")
|
||||
}
|
||||
|
||||
modelItem, err := s.models.Get(ctx, strings.TrimSpace(in.Model))
|
||||
if err != nil {
|
||||
return nil, ErrUnknownModel
|
||||
}
|
||||
|
||||
principal := &APIPrincipal{
|
||||
User: user,
|
||||
TokenType: "session",
|
||||
}
|
||||
|
||||
switch modelItem.Type {
|
||||
case "video":
|
||||
resp, err := s.v1.prepareSessionVideo(ctx, principal, V1VideoRequest{
|
||||
Model: in.Model,
|
||||
Prompt: in.Prompt,
|
||||
Duration: in.Duration,
|
||||
AspectRatio: in.Ratio,
|
||||
Resolution: in.Resolution,
|
||||
ReferenceImages: in.ReferenceImages,
|
||||
})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return resp, nil
|
||||
default:
|
||||
resp, err := s.v1.prepareSessionImage(ctx, principal, V1ImageRequest{
|
||||
Model: in.Model,
|
||||
Prompt: in.Prompt,
|
||||
AspectRatio: in.Ratio,
|
||||
Resolution: in.Resolution,
|
||||
ReferenceImages: in.ReferenceImages,
|
||||
})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return resp, nil
|
||||
}
|
||||
}
|
||||
|
||||
func (s *UserGenerationService) AdminTest(ctx context.Context, user *model.User, in UserGenerateRequest) (map[string]any, error) {
|
||||
if user == nil || strings.TrimSpace(user.ID) == "" {
|
||||
return nil, errors.New("未登录或会话已过期")
|
||||
}
|
||||
modelItem, err := s.models.Get(ctx, strings.TrimSpace(in.Model))
|
||||
if err != nil {
|
||||
return nil, ErrUnknownModel
|
||||
}
|
||||
principal := &APIPrincipal{
|
||||
User: user,
|
||||
TokenType: "session",
|
||||
}
|
||||
switch modelItem.Type {
|
||||
case "video":
|
||||
return s.v1.prepareAdminTestVideo(ctx, principal, V1VideoRequest{
|
||||
Model: in.Model,
|
||||
Prompt: in.Prompt,
|
||||
Duration: in.Duration,
|
||||
AspectRatio: in.Ratio,
|
||||
Resolution: in.Resolution,
|
||||
ReferenceImages: in.ReferenceImages,
|
||||
})
|
||||
default:
|
||||
return s.v1.prepareAdminTestImage(ctx, principal, V1ImageRequest{
|
||||
Model: in.Model,
|
||||
Prompt: in.Prompt,
|
||||
AspectRatio: in.Ratio,
|
||||
Resolution: in.Resolution,
|
||||
ReferenceImages: in.ReferenceImages,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func (s *UserGenerationService) MyJobs(ctx context.Context, user *model.User, source string) (map[string]any, error) {
|
||||
if user == nil || strings.TrimSpace(user.ID) == "" {
|
||||
return map[string]any{"pending": nil, "latest": nil}, nil
|
||||
}
|
||||
// source scopes the lookup: "user" = 画图台(默认),"admin" = 后台测试模型。
|
||||
// Both are this caller's own events; the admin-test poll uses "admin" so a
|
||||
// gateway-timed-out (524) test can still recover its result.
|
||||
if source != "admin" {
|
||||
source = "user"
|
||||
}
|
||||
pending, err := s.events.PendingByUser(ctx, user.ID, source)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
latest, err := s.events.LatestByUser(ctx, user.ID, source)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return map[string]any{
|
||||
"pending": shapeJobEvent(pending),
|
||||
"latest": shapeJobEvent(latest),
|
||||
}, nil
|
||||
}
|
||||
|
||||
func shapeJobEvent(item *model.EventLog) map[string]any {
|
||||
if item == nil {
|
||||
return nil
|
||||
}
|
||||
status := item.Status
|
||||
url := ""
|
||||
if strings.TrimSpace(item.File) != "" {
|
||||
url = "/images/" + strings.ReplaceAll(strings.TrimSpace(item.File), "\\", "/")
|
||||
}
|
||||
return map[string]any{
|
||||
"id": item.ID,
|
||||
"kind": item.Kind,
|
||||
"model": item.Model,
|
||||
"prompt": item.Prompt,
|
||||
"ratio": item.Ratio,
|
||||
"resolution": item.Resolution,
|
||||
"duration": item.Duration,
|
||||
"status": status,
|
||||
"file": emptyOrNil(item.File),
|
||||
"url": emptyOrNil(url),
|
||||
"reference_urls": referenceURLs(item.RefFiles),
|
||||
"elapsed_ms": item.ElapsedMS,
|
||||
"error": emptyOrNil(item.Error),
|
||||
"charged": item.Cost,
|
||||
"cost": item.Cost,
|
||||
"ts": item.TS.Unix(),
|
||||
}
|
||||
}
|
||||
|
||||
// referenceURLs turns the stored relative reference paths into /images URLs so
|
||||
// the playground can re-display the uploaded reference image(s) after a reload.
|
||||
func referenceURLs(raw []byte) []string {
|
||||
if len(raw) == 0 {
|
||||
return []string{}
|
||||
}
|
||||
var paths []string
|
||||
if err := json.Unmarshal(raw, &paths); err != nil {
|
||||
return []string{}
|
||||
}
|
||||
out := make([]string, 0, len(paths))
|
||||
for _, p := range paths {
|
||||
p = strings.ReplaceAll(strings.TrimSpace(p), "\\", "/")
|
||||
if p != "" {
|
||||
out = append(out, "/images/"+p)
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func emptyOrNil(v string) any {
|
||||
if strings.TrimSpace(v) == "" {
|
||||
return nil
|
||||
}
|
||||
return v
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,161 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"net/mail"
|
||||
"regexp"
|
||||
"strings"
|
||||
"unicode"
|
||||
"unicode/utf8"
|
||||
)
|
||||
|
||||
const (
|
||||
MinUsernameLength = 6
|
||||
MaxUsernameLength = 24
|
||||
MinPasswordLength = 8
|
||||
MaxPasswordLength = 24
|
||||
)
|
||||
|
||||
var (
|
||||
usernamePattern = regexp.MustCompile(`^[A-Za-z0-9]{6,24}$`)
|
||||
emailCodePattern = regexp.MustCompile(`^\d{6}$`)
|
||||
)
|
||||
|
||||
func ValidateEmail(email string) (string, error) {
|
||||
normalized := strings.TrimSpace(strings.ToLower(email))
|
||||
if normalized == "" {
|
||||
return "", errors.New("邮箱不能为空")
|
||||
}
|
||||
if len(normalized) > 254 {
|
||||
return "", errors.New("邮箱长度不能超过 254 个字符")
|
||||
}
|
||||
addr, err := mail.ParseAddress(normalized)
|
||||
if err != nil || strings.TrimSpace(strings.ToLower(addr.Address)) != normalized {
|
||||
return "", errors.New("邮箱格式不正确")
|
||||
}
|
||||
local, domain, ok := strings.Cut(normalized, "@")
|
||||
if !ok || local == "" || domain == "" || strings.Contains(domain, "..") || !strings.Contains(domain, ".") {
|
||||
return "", errors.New("邮箱格式不正确")
|
||||
}
|
||||
return normalized, nil
|
||||
}
|
||||
|
||||
func ValidateUsername(username string) (string, error) {
|
||||
normalized := strings.TrimSpace(username)
|
||||
if normalized == "" {
|
||||
return "", errors.New("用户名不能为空")
|
||||
}
|
||||
length := utf8.RuneCountInString(normalized)
|
||||
if length < MinUsernameLength || length > MaxUsernameLength {
|
||||
return "", errors.New("用户名长度需为 6 到 24 个字符")
|
||||
}
|
||||
if !usernamePattern.MatchString(normalized) {
|
||||
return "", errors.New("用户名只能使用字母和数字")
|
||||
}
|
||||
return normalized, nil
|
||||
}
|
||||
|
||||
func ValidatePassword(password string) error {
|
||||
length := utf8.RuneCountInString(password)
|
||||
if length < MinPasswordLength || length > MaxPasswordLength {
|
||||
return errors.New("密码长度需为 8 到 24 个字符")
|
||||
}
|
||||
|
||||
var hasLetter bool
|
||||
var hasUpper bool
|
||||
var hasLower bool
|
||||
var hasDigit bool
|
||||
var hasSymbol bool
|
||||
for _, r := range password {
|
||||
if unicode.IsSpace(r) {
|
||||
return errors.New("密码不能包含空白字符")
|
||||
}
|
||||
if !isAllowedPasswordRune(r) {
|
||||
return errors.New("密码包含不允许的字符")
|
||||
}
|
||||
if unicode.IsLetter(r) {
|
||||
hasLetter = true
|
||||
if unicode.IsUpper(r) {
|
||||
hasUpper = true
|
||||
}
|
||||
if unicode.IsLower(r) {
|
||||
hasLower = true
|
||||
}
|
||||
}
|
||||
if unicode.IsDigit(r) {
|
||||
hasDigit = true
|
||||
}
|
||||
if !unicode.IsLetter(r) && !unicode.IsDigit(r) {
|
||||
hasSymbol = true
|
||||
}
|
||||
}
|
||||
if !hasLetter || !hasUpper || !hasLower || !hasDigit || !hasSymbol {
|
||||
return errors.New("密码必须同时包含大写字母、小写字母、数字和符号")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func isAllowedPasswordRune(r rune) bool {
|
||||
if unicode.IsLetter(r) || unicode.IsDigit(r) {
|
||||
return true
|
||||
}
|
||||
switch r {
|
||||
case '(', ')', '~', '!', '@', '#', '$', '%', '^', '&', '*', '-', '_', '+', '=', '|',
|
||||
'{', '}', '[', ']', ':', ';', '\'', '<', '>', ',', '.', '?', '/':
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
func ValidateEmailCode(code string) (string, error) {
|
||||
normalized := strings.TrimSpace(code)
|
||||
if !emailCodePattern.MatchString(normalized) {
|
||||
return "", errors.New("邮箱验证码必须是 6 位纯数字")
|
||||
}
|
||||
return normalized, nil
|
||||
}
|
||||
|
||||
func ValidateLoginIdentifier(identifier string) (string, error) {
|
||||
normalized := strings.TrimSpace(identifier)
|
||||
if normalized == "" {
|
||||
return "", errors.New("账号不能为空")
|
||||
}
|
||||
if strings.Contains(normalized, "@") {
|
||||
return ValidateEmail(normalized)
|
||||
}
|
||||
return ValidateUsername(normalized)
|
||||
}
|
||||
|
||||
func ValidateAllowedEmailDomains(domains []string) []string {
|
||||
out := make([]string, 0, len(domains))
|
||||
seen := map[string]struct{}{}
|
||||
for _, raw := range domains {
|
||||
normalized := strings.TrimSpace(strings.ToLower(strings.TrimPrefix(raw, "@")))
|
||||
if normalized == "" || strings.Contains(normalized, " ") {
|
||||
continue
|
||||
}
|
||||
if _, ok := seen[normalized]; ok {
|
||||
continue
|
||||
}
|
||||
seen[normalized] = struct{}{}
|
||||
out = append(out, normalized)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func EmailDomainAllowed(email string, domains []string) bool {
|
||||
if len(domains) == 0 {
|
||||
return true
|
||||
}
|
||||
_, domain, ok := strings.Cut(strings.ToLower(strings.TrimSpace(email)), "@")
|
||||
if !ok {
|
||||
return false
|
||||
}
|
||||
for _, allowed := range ValidateAllowedEmailDomains(domains) {
|
||||
if domain == allowed {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
@@ -0,0 +1,294 @@
|
||||
// Package storage is a minimal S3-compatible client for RustFS, implemented with
|
||||
// AWS Signature V4 over the standard library only (no external SDK — the build
|
||||
// environment can't reach the Go module proxy). It's intentionally small: Put /
|
||||
// Get / Delete / List cover everything the app needs (store generated media,
|
||||
// proxy it back through /images, list for the admin gallery, prune by age).
|
||||
//
|
||||
// The surface mirrors what a thin wrapper over aws-sdk-go-v2 would expose, so it
|
||||
// can be swapped for the official SDK later by reimplementing this one file.
|
||||
package storage
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"crypto/hmac"
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"encoding/xml"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
const (
|
||||
region = "us-east-1" // RustFS ignores the value but SigV4 requires a fixed one
|
||||
service = "s3"
|
||||
)
|
||||
|
||||
type Client struct {
|
||||
endpoint string // e.g. http://154.9.26.140:9000 (no trailing slash)
|
||||
host string // e.g. 154.9.26.140:9000
|
||||
bucket string
|
||||
ak, sk string
|
||||
http *http.Client
|
||||
}
|
||||
|
||||
// Object is one entry returned by List.
|
||||
type Object struct {
|
||||
Key string
|
||||
Size int64
|
||||
LastModified time.Time
|
||||
}
|
||||
|
||||
// New builds a client. endpoint must include the scheme (http:// or https://).
|
||||
func New(endpoint, bucket, accessKey, secretKey string) *Client {
|
||||
endpoint = strings.TrimRight(strings.TrimSpace(endpoint), "/")
|
||||
host := endpoint
|
||||
if i := strings.Index(host, "://"); i >= 0 {
|
||||
host = host[i+3:]
|
||||
}
|
||||
return &Client{
|
||||
endpoint: endpoint,
|
||||
host: host,
|
||||
bucket: strings.TrimSpace(bucket),
|
||||
ak: strings.TrimSpace(accessKey),
|
||||
sk: strings.TrimSpace(secretKey),
|
||||
http: &http.Client{Timeout: 60 * time.Second},
|
||||
}
|
||||
}
|
||||
|
||||
// Configured reports whether the client has the minimum config to be usable.
|
||||
func (c *Client) Configured() bool {
|
||||
return c != nil && c.endpoint != "" && c.bucket != "" && c.ak != "" && c.sk != ""
|
||||
}
|
||||
|
||||
// PublicURL is the direct object URL (used only for reference/debugging — the app
|
||||
// serves through the authenticated /images proxy, not this).
|
||||
func (c *Client) PublicURL(key string) string {
|
||||
return c.endpoint + "/" + c.bucket + "/" + strings.TrimPrefix(key, "/")
|
||||
}
|
||||
|
||||
// Put uploads body under key with the given content type.
|
||||
func (c *Client) Put(ctx context.Context, key string, body []byte, contentType string) error {
|
||||
resp, err := c.do(ctx, http.MethodPut, c.bucket+"/"+key, nil, body, contentType, nil)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode/100 != 2 {
|
||||
return c.statusErr("put", key, resp)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Get fetches key. The caller owns resp.Body (must Close it) and streams it. A
|
||||
// non-empty rangeHeader is forwarded verbatim (for video seeking). Returns the
|
||||
// raw *http.Response so headers/status can be passed through by the proxy.
|
||||
func (c *Client) Get(ctx context.Context, key, rangeHeader string) (*http.Response, error) {
|
||||
extra := map[string]string{}
|
||||
if strings.TrimSpace(rangeHeader) != "" {
|
||||
extra["Range"] = rangeHeader
|
||||
}
|
||||
return c.do(ctx, http.MethodGet, c.bucket+"/"+key, nil, nil, "", extra)
|
||||
}
|
||||
|
||||
// Delete removes key. A missing object is not an error.
|
||||
func (c *Client) Delete(ctx context.Context, key string) error {
|
||||
resp, err := c.do(ctx, http.MethodDelete, c.bucket+"/"+key, nil, nil, "", nil)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode/100 != 2 && resp.StatusCode != http.StatusNotFound {
|
||||
return c.statusErr("delete", key, resp)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// List returns every object whose key starts with prefix (paginated internally).
|
||||
func (c *Client) List(ctx context.Context, prefix string) ([]Object, error) {
|
||||
var out []Object
|
||||
token := ""
|
||||
for {
|
||||
q := map[string]string{"list-type": "2", "max-keys": "1000"}
|
||||
if prefix != "" {
|
||||
q["prefix"] = prefix
|
||||
}
|
||||
if token != "" {
|
||||
q["continuation-token"] = token
|
||||
}
|
||||
resp, err := c.do(ctx, http.MethodGet, c.bucket, q, nil, "", nil)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
data, _ := io.ReadAll(resp.Body)
|
||||
resp.Body.Close()
|
||||
if resp.StatusCode/100 != 2 {
|
||||
return nil, fmt.Errorf("rustfs list: status %d: %s", resp.StatusCode, truncate(data))
|
||||
}
|
||||
var parsed struct {
|
||||
Contents []struct {
|
||||
Key string `xml:"Key"`
|
||||
Size int64 `xml:"Size"`
|
||||
LastModified time.Time `xml:"LastModified"`
|
||||
} `xml:"Contents"`
|
||||
IsTruncated bool `xml:"IsTruncated"`
|
||||
NextContinuationToken string `xml:"NextContinuationToken"`
|
||||
}
|
||||
if err := xml.Unmarshal(data, &parsed); err != nil {
|
||||
return nil, fmt.Errorf("rustfs list: parse: %w", err)
|
||||
}
|
||||
for _, it := range parsed.Contents {
|
||||
out = append(out, Object{Key: it.Key, Size: it.Size, LastModified: it.LastModified})
|
||||
}
|
||||
if !parsed.IsTruncated || parsed.NextContinuationToken == "" {
|
||||
break
|
||||
}
|
||||
token = parsed.NextContinuationToken
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func (c *Client) statusErr(op, key string, resp *http.Response) error {
|
||||
data, _ := io.ReadAll(io.LimitReader(resp.Body, 512))
|
||||
return fmt.Errorf("rustfs %s %q: status %d: %s", op, key, resp.StatusCode, truncate(data))
|
||||
}
|
||||
|
||||
func truncate(b []byte) string {
|
||||
s := strings.TrimSpace(string(b))
|
||||
if len(s) > 300 {
|
||||
return s[:300]
|
||||
}
|
||||
return s
|
||||
}
|
||||
|
||||
// do builds, signs (SigV4) and sends a request. resourcePath is the path after
|
||||
// the host WITHOUT a leading slash (e.g. "bucket/dir/file.png" or "bucket").
|
||||
func (c *Client) do(ctx context.Context, method, resourcePath string, query map[string]string, body []byte, contentType string, extraHeaders map[string]string) (*http.Response, error) {
|
||||
now := time.Now().UTC()
|
||||
amzDate := now.Format("20060102T150405Z")
|
||||
dateStamp := now.Format("20060102")
|
||||
|
||||
canonicalURI := "/" + uriEncode(resourcePath, true)
|
||||
canonicalQuery := canonicalQueryString(query)
|
||||
|
||||
payloadHash := hexSHA256(body)
|
||||
// Signed headers: always host + x-amz-content-sha256 + x-amz-date, plus
|
||||
// content-type on PUT. Range etc. are sent unsigned.
|
||||
signed := map[string]string{
|
||||
"host": c.host,
|
||||
"x-amz-content-sha256": payloadHash,
|
||||
"x-amz-date": amzDate,
|
||||
}
|
||||
if strings.TrimSpace(contentType) != "" {
|
||||
signed["content-type"] = contentType
|
||||
}
|
||||
names := sortedKeys(signed)
|
||||
var canonHeaders strings.Builder
|
||||
for _, k := range names {
|
||||
canonHeaders.WriteString(k + ":" + signed[k] + "\n")
|
||||
}
|
||||
signedHeaders := strings.Join(names, ";")
|
||||
|
||||
canonicalRequest := strings.Join([]string{
|
||||
method, canonicalURI, canonicalQuery, canonHeaders.String(), signedHeaders, payloadHash,
|
||||
}, "\n")
|
||||
|
||||
scope := dateStamp + "/" + region + "/" + service + "/aws4_request"
|
||||
stringToSign := strings.Join([]string{
|
||||
"AWS4-HMAC-SHA256", amzDate, scope, hexSHA256([]byte(canonicalRequest)),
|
||||
}, "\n")
|
||||
signature := hex.EncodeToString(hmacSHA256(signingKey(c.sk, dateStamp), []byte(stringToSign)))
|
||||
auth := fmt.Sprintf("AWS4-HMAC-SHA256 Credential=%s/%s, SignedHeaders=%s, Signature=%s",
|
||||
c.ak, scope, signedHeaders, signature)
|
||||
|
||||
url := c.endpoint + canonicalURI
|
||||
if canonicalQuery != "" {
|
||||
url += "?" + canonicalQuery
|
||||
}
|
||||
var rdr io.Reader
|
||||
if body != nil {
|
||||
rdr = bytes.NewReader(body)
|
||||
}
|
||||
req, err := http.NewRequestWithContext(ctx, method, url, rdr)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
req.Host = c.host
|
||||
req.Header.Set("Authorization", auth)
|
||||
req.Header.Set("x-amz-date", amzDate)
|
||||
req.Header.Set("x-amz-content-sha256", payloadHash)
|
||||
if ct := strings.TrimSpace(contentType); ct != "" {
|
||||
req.Header.Set("Content-Type", ct)
|
||||
}
|
||||
for k, v := range extraHeaders {
|
||||
req.Header.Set(k, v)
|
||||
}
|
||||
return c.http.Do(req)
|
||||
}
|
||||
|
||||
// ---- SigV4 helpers ----
|
||||
|
||||
func hexSHA256(b []byte) string {
|
||||
h := sha256.Sum256(b)
|
||||
return hex.EncodeToString(h[:])
|
||||
}
|
||||
|
||||
func hmacSHA256(key, msg []byte) []byte {
|
||||
h := hmac.New(sha256.New, key)
|
||||
h.Write(msg)
|
||||
return h.Sum(nil)
|
||||
}
|
||||
|
||||
func signingKey(secret, dateStamp string) []byte {
|
||||
kDate := hmacSHA256([]byte("AWS4"+secret), []byte(dateStamp))
|
||||
kRegion := hmacSHA256(kDate, []byte(region))
|
||||
kService := hmacSHA256(kRegion, []byte(service))
|
||||
return hmacSHA256(kService, []byte("aws4_request"))
|
||||
}
|
||||
|
||||
func sortedKeys(m map[string]string) []string {
|
||||
out := make([]string, 0, len(m))
|
||||
for k := range m {
|
||||
out = append(out, k)
|
||||
}
|
||||
// simple insertion sort (small n)
|
||||
for i := 1; i < len(out); i++ {
|
||||
for j := i; j > 0 && out[j-1] > out[j]; j-- {
|
||||
out[j-1], out[j] = out[j], out[j-1]
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func canonicalQueryString(q map[string]string) string {
|
||||
if len(q) == 0 {
|
||||
return ""
|
||||
}
|
||||
keys := sortedKeys(q)
|
||||
parts := make([]string, 0, len(keys))
|
||||
for _, k := range keys {
|
||||
parts = append(parts, uriEncode(k, false)+"="+uriEncode(q[k], false))
|
||||
}
|
||||
return strings.Join(parts, "&")
|
||||
}
|
||||
|
||||
// uriEncode applies AWS's URI encoding rules. When keepSlash is true, '/' is left
|
||||
// as-is (for object key paths); otherwise it's percent-encoded (for query parts).
|
||||
func uriEncode(s string, keepSlash bool) string {
|
||||
var b strings.Builder
|
||||
for _, r := range []byte(s) {
|
||||
switch {
|
||||
case (r >= 'A' && r <= 'Z') || (r >= 'a' && r <= 'z') || (r >= '0' && r <= '9'),
|
||||
r == '-', r == '_', r == '.', r == '~':
|
||||
b.WriteByte(r)
|
||||
case r == '/' && keepSlash:
|
||||
b.WriteByte('/')
|
||||
default:
|
||||
b.WriteString(fmt.Sprintf("%%%02X", r))
|
||||
}
|
||||
}
|
||||
return b.String()
|
||||
}
|
||||
Reference in New Issue
Block a user