diff --git a/common/body_storage.go b/common/body_storage.go index 094dbda3..bd964b10 100644 --- a/common/body_storage.go +++ b/common/body_storage.go @@ -20,6 +20,15 @@ type BodyStorage interface { Size() int64 // IsDisk 是否是磁盘存储 IsDisk() bool + // NewReader returns an independent reader positioned at the start. + NewReader() (io.ReadCloser, error) +} + +// ReplayableBody exposes body metadata without transferring storage ownership. +type ReplayableBody interface { + io.Reader + Size() int64 + NewReader() (io.ReadCloser, error) } // ErrStorageClosed 存储已关闭错误 @@ -80,6 +89,15 @@ func (m *memoryStorage) Bytes() ([]byte, error) { return m.data, nil } +func (m *memoryStorage) NewReader() (io.ReadCloser, error) { + m.mu.Lock() + defer m.mu.Unlock() + if atomic.LoadInt32(&m.closed) == 1 { + return nil, ErrStorageClosed + } + return io.NopCloser(bytes.NewReader(m.data)), nil +} + func (m *memoryStorage) Size() int64 { return m.size } @@ -229,6 +247,19 @@ func (d *diskStorage) Bytes() ([]byte, error) { return data, nil } +func (d *diskStorage) NewReader() (io.ReadCloser, error) { + d.mu.Lock() + defer d.mu.Unlock() + if atomic.LoadInt32(&d.closed) == 1 { + return nil, ErrStorageClosed + } + file, err := os.Open(d.filePath) + if err != nil { + return nil, fmt.Errorf("failed to open body cache file for replay: %w", err) + } + return file, nil +} + func (d *diskStorage) Size() int64 { return d.size } @@ -302,9 +333,31 @@ func CreateBodyStorageFromReader(reader io.Reader, contentLength int64, maxBytes return storage, nil } -// ReaderOnly wraps an io.Reader to hide io.Closer, preventing http.NewRequest -// from type-asserting io.ReadCloser and closing the underlying BodyStorage. +type replayableBodyReader struct { + storage BodyStorage +} + +func (r replayableBodyReader) Read(p []byte) (int, error) { + return r.storage.Read(p) +} + +func (r replayableBodyReader) Size() int64 { + return r.storage.Size() +} + +func (r replayableBodyReader) NewReader() (io.ReadCloser, error) { + return r.storage.NewReader() +} + +func NewReplayableBodyReader(storage BodyStorage) ReplayableBody { + return replayableBodyReader{storage: storage} +} + +// ReaderOnly hides io.Closer while preserving replay metadata for BodyStorage. func ReaderOnly(r io.Reader) io.Reader { + if storage, ok := r.(BodyStorage); ok { + return NewReplayableBodyReader(storage) + } return struct{ io.Reader }{r} } diff --git a/common/body_storage_replay_test.go b/common/body_storage_replay_test.go new file mode 100644 index 00000000..e3cfba94 --- /dev/null +++ b/common/body_storage_replay_test.go @@ -0,0 +1,53 @@ +package common + +import ( + "io" + "testing" + + "github.com/stretchr/testify/require" +) + +func requireIndependentBodyReaders(t *testing.T, storage BodyStorage, payload []byte) { + t.Helper() + first, err := storage.NewReader() + require.NoError(t, err) + defer first.Close() + second, err := storage.NewReader() + require.NoError(t, err) + defer second.Close() + + prefix := make([]byte, 4) + _, err = io.ReadFull(first, prefix) + require.NoError(t, err) + secondBody, err := io.ReadAll(second) + require.NoError(t, err) + require.Equal(t, payload, secondBody) + firstRest, err := io.ReadAll(first) + require.NoError(t, err) + require.Equal(t, payload[4:], firstRest) +} + +func TestBodyStorageNewReaderUsesIndependentCursors(t *testing.T) { + payload := []byte("independent replay payload") + + t.Run("memory", func(t *testing.T) { + storage := newMemoryStorage(payload) + defer storage.Close() + requireIndependentBodyReaders(t, storage, payload) + }) + + t.Run("disk", func(t *testing.T) { + storage, err := newDiskStorage(payload, "") + require.NoError(t, err) + defer storage.Close() + requireIndependentBodyReaders(t, storage, payload) + }) +} + +func TestBodyStorageNewReaderRejectsClosedStorage(t *testing.T) { + storage := newMemoryStorage([]byte("closed")) + require.NoError(t, storage.Close()) + reader, err := storage.NewReader() + require.Nil(t, reader) + require.ErrorIs(t, err, ErrStorageClosed) +} diff --git a/common/email.go b/common/email.go index 9d65f3b1..50abee61 100644 --- a/common/email.go +++ b/common/email.go @@ -5,6 +5,7 @@ import ( "encoding/base64" "errors" "fmt" + "net" "net/smtp" "slices" "strings" @@ -12,6 +13,7 @@ import ( ) var ErrSMTPStartTLSUnsupported = errors.New("smtp server does not support STARTTLS") +var smtpOperationTimeout = 30 * time.Second func generateMessageID() (string, error) { split := strings.Split(SMTPFrom, "@") @@ -19,7 +21,11 @@ func generateMessageID() (string, error) { return "", fmt.Errorf("invalid SMTP account") } domain := strings.Split(SMTPFrom, "@")[1] - return fmt.Sprintf("<%d.%s@%s>", time.Now().UnixNano(), GetRandomString(12), domain), nil + randomPart, err := GenerateRandomCharsKey(12) + if err != nil { + return "", fmt.Errorf("generate message ID: %w", err) + } + return fmt.Sprintf("<%d.%s@%s>", time.Now().UnixNano(), randomPart, domain), nil } func shouldUseSMTPLoginAuth() bool { @@ -59,11 +65,17 @@ func smtpTLSConfig() *tls.Config { } func newSMTPClient(addr string) (*smtp.Client, error) { + dialer := &net.Dialer{Timeout: smtpOperationTimeout} + deadline := time.Now().Add(smtpOperationTimeout) if SMTPSSLEnabled || SMTPPort == 465 { - conn, err := tls.Dial("tcp", addr, smtpTLSConfig()) + conn, err := tls.DialWithDialer(dialer, "tcp", addr, smtpTLSConfig()) if err != nil { return nil, err } + if err := conn.SetDeadline(deadline); err != nil { + _ = conn.Close() + return nil, err + } client, err := smtp.NewClient(conn, SMTPServer) if err != nil { _ = conn.Close() @@ -72,8 +84,17 @@ func newSMTPClient(addr string) (*smtp.Client, error) { return client, nil } - client, err := smtp.Dial(addr) + conn, err := dialer.Dial("tcp", addr) + if err != nil { + return nil, err + } + if err := conn.SetDeadline(deadline); err != nil { + _ = conn.Close() + return nil, err + } + client, err := smtp.NewClient(conn, SMTPServer) if err != nil { + _ = conn.Close() return nil, err } diff --git a/common/email_test.go b/common/email_test.go index 13fd1d4c..75f89a0d 100644 --- a/common/email_test.go +++ b/common/email_test.go @@ -220,6 +220,53 @@ func TestSendEmailSkipsAuthWhenServerDoesNotAdvertiseAuth(t *testing.T) { } } +func TestNewSMTPClientUsesOperationDeadline(t *testing.T) { + restore := preserveSMTPGlobals() + defer restore() + + listener, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + defer listener.Close() + + SMTPServer = "localhost" + SMTPPort = mustPort(t, listener.Addr().String()) + SMTPSSLEnabled = false + SMTPStartTLSEnabled = false + + oldTimeout := smtpOperationTimeout + smtpOperationTimeout = 50 * time.Millisecond + defer func() { smtpOperationTimeout = oldTimeout }() + + done := make(chan error, 1) + go func() { + conn, acceptErr := listener.Accept() + if acceptErr != nil { + done <- acceptErr + return + } + defer conn.Close() + time.Sleep(200 * time.Millisecond) + done <- nil + }() + + started := time.Now() + client, err := newSMTPClient(listener.Addr().String()) + if client != nil { + _ = client.Close() + } + if err == nil { + t.Fatal("expected SMTP greeting timeout") + } + if time.Since(started) >= time.Second { + t.Fatalf("SMTP timeout took too long: %s", time.Since(started)) + } + if err := waitSMTPServer(done); err != nil { + t.Fatal(err) + } +} + func preserveSMTPGlobals() func() { old := struct { server string diff --git a/common/go-channel.go b/common/go-channel.go index f9168fc4..9863fb8d 100644 --- a/common/go-channel.go +++ b/common/go-channel.go @@ -34,12 +34,13 @@ func SafeSendString(ch chan string, value string) (closed bool) { return false } -// SafeSendStringTimeout send, return true, else return false -func SafeSendStringTimeout(ch chan string, value string, timeout int) (closed bool) { +// SafeSendStringTimeout returns true only when the value is sent before the timeout. +// A timeout or a closed channel returns false. +func SafeSendStringTimeout(ch chan string, value string, timeout int) (sent bool) { defer func() { - // Recover from panic if one occured. A panic would mean the channel was closed. + // Recover from panic if the channel was closed between selection and send. if recover() != nil { - closed = false + sent = false } }() diff --git a/common/go_channel_test.go b/common/go_channel_test.go new file mode 100644 index 00000000..a2a0864a --- /dev/null +++ b/common/go_channel_test.go @@ -0,0 +1,29 @@ +package common + +import ( + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +func TestSafeSendStringTimeoutReportsSendOutcome(t *testing.T) { + t.Run("sent", func(t *testing.T) { + ch := make(chan string, 1) + require.True(t, SafeSendStringTimeout(ch, "value", 1)) + require.Equal(t, "value", <-ch) + }) + + t.Run("timeout", func(t *testing.T) { + ch := make(chan string) + started := time.Now() + require.False(t, SafeSendStringTimeout(ch, "value", 0)) + require.Less(t, time.Since(started), time.Second) + }) + + t.Run("closed", func(t *testing.T) { + ch := make(chan string) + close(ch) + require.False(t, SafeSendStringTimeout(ch, "value", 1)) + }) +} diff --git a/common/init.go b/common/init.go index 2b28a566..8720cdc1 100644 --- a/common/init.go +++ b/common/init.go @@ -93,15 +93,10 @@ func InitEnv() { SessionCookieSecure = GetEnvOrDefaultBool("SESSION_COOKIE_SECURE", false) SMTPStartTLSEnabled = GetEnvOrDefaultBool("SMTP_STARTTLS_ENABLE", GetEnvOrDefaultBool("SMTP_STARTTLS_ENABLED", false)) SMTPInsecureSkipVerify = GetEnvOrDefaultBool("SMTP_INSECURE_SKIP_VERIFY", GetEnvOrDefaultBool("SMTP_TLS_INSECURE_SKIP_VERIFY", false)) - TrustedProxies = append([]string(nil), DefaultTrustedProxies...) - if trustedProxies := strings.TrimSpace(os.Getenv("TRUSTED_PROXIES")); trustedProxies != "" { - TrustedProxies = nil - for _, proxy := range strings.Split(trustedProxies, ",") { - proxy = strings.TrimSpace(proxy) - if proxy != "" { - TrustedProxies = append(TrustedProxies, proxy) - } - } + var trustedProxyErr error + TrustedProxies, trustedProxyErr = parseTrustedProxies(os.Getenv("TRUSTED_PROXIES")) + if trustedProxyErr != nil { + log.Fatal(trustedProxyErr) } IsMasterNode = os.Getenv("NODE_TYPE") != "slave" initNodeNameIdentity() diff --git a/common/pprof.go b/common/pprof.go index 74592653..14058166 100644 --- a/common/pprof.go +++ b/common/pprof.go @@ -1,6 +1,7 @@ package common import ( + "context" "fmt" "os" "runtime/pprof" @@ -9,37 +10,76 @@ import ( "github.com/shirou/gopsutil/cpu" ) +var cpuPercent = cpu.Percent + +const cpuMonitorInterval = 30 * time.Second + // Monitor 定时监控cpu使用率，超过阈值输出pprof文件 func Monitor() { + MonitorContext(context.Background()) +} + +func MonitorContext(ctx context.Context) { for { - percent, err := cpu.Percent(time.Second, false) - if err != nil { - panic(err) + if ctx.Err() != nil { + return } - if percent[0] > 80 { - fmt.Println("cpu usage too high") - // write pprof file - if _, err := os.Stat("./pprof"); os.IsNotExist(err) { - err := os.Mkdir("./pprof", os.ModePerm) - if err != nil { - SysLog("创建pprof文件夹失败 " + err.Error()) - continue + captureHighCPUProfileContext(ctx) + timer := time.NewTimer(cpuMonitorInterval) + select { + case <-ctx.Done(): + if !timer.Stop() { + select { + case <-timer.C: + default: } } - f, err := os.Create("./pprof/" + fmt.Sprintf("cpu-%s.pprof", time.Now().Format("20060102150405"))) - if err != nil { - SysLog("创建pprof文件失败 " + err.Error()) - continue - } - err = pprof.StartCPUProfile(f) - if err != nil { - SysLog("启动pprof失败 " + err.Error()) - continue + return + case <-timer.C: + } + } +} + +func captureHighCPUProfile() { + captureHighCPUProfileContext(context.Background()) +} + +func captureHighCPUProfileContext(ctx context.Context) { + percent, err := cpuPercent(time.Second, false) + if err != nil { + SysLog("获取CPU使用率失败 " + err.Error()) + return + } + if len(percent) == 0 || percent[0] <= 80 { + return + } + + fmt.Println("cpu usage too high") + if err := os.MkdirAll("./pprof", 0o755); err != nil { + SysLog("创建pprof文件夹失败 " + err.Error()) + return + } + f, err := os.Create("./pprof/" + fmt.Sprintf("cpu-%s.pprof", time.Now().Format("20060102150405"))) + if err != nil { + SysLog("创建pprof文件失败 " + err.Error()) + return + } + if err := pprof.StartCPUProfile(f); err != nil { + _ = f.Close() + SysLog("启动pprof失败 " + err.Error()) + return + } + timer := time.NewTimer(10 * time.Second) + select { + case <-ctx.Done(): + if !timer.Stop() { + select { + case <-timer.C: + default: } - time.Sleep(10 * time.Second) // profile for 30 seconds - pprof.StopCPUProfile() - f.Close() } - time.Sleep(30 * time.Second) + case <-timer.C: } + pprof.StopCPUProfile() + _ = f.Close() } diff --git a/common/pprof_test.go b/common/pprof_test.go new file mode 100644 index 00000000..310cfc79 --- /dev/null +++ b/common/pprof_test.go @@ -0,0 +1,57 @@ +package common + +import ( + "context" + "errors" + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +func TestCaptureHighCPUProfileHandlesSamplingFailureAndEmptyResult(t *testing.T) { + original := cpuPercent + t.Cleanup(func() { cpuPercent = original }) + + cpuPercent = func(time.Duration, bool) ([]float64, error) { + return nil, errors.New("sample failed") + } + require.NotPanics(t, captureHighCPUProfile) + + cpuPercent = func(time.Duration, bool) ([]float64, error) { + return nil, nil + } + require.NotPanics(t, captureHighCPUProfile) +} + +func TestMonitorStopsOnCancellation(t *testing.T) { + original := cpuPercent + t.Cleanup(func() { cpuPercent = original }) + sampled := make(chan struct{}, 1) + cpuPercent = func(time.Duration, bool) ([]float64, error) { + select { + case sampled <- struct{}{}: + default: + } + return nil, nil + } + + ctx, cancel := context.WithCancel(context.Background()) + done := make(chan struct{}) + go func() { + defer close(done) + MonitorContext(ctx) + }() + + select { + case <-sampled: + case <-time.After(time.Second): + t.Fatal("CPU monitor did not sample") + } + cancel() + select { + case <-done: + case <-time.After(time.Second): + t.Fatal("CPU monitor did not stop after cancellation") + } +} diff --git a/common/random_failure_test.go b/common/random_failure_test.go new file mode 100644 index 00000000..9a1f817f --- /dev/null +++ b/common/random_failure_test.go @@ -0,0 +1,46 @@ +package common + +import ( + crand "crypto/rand" + "errors" + "io" + "testing" + + "github.com/stretchr/testify/require" +) + +type failingRandomReader struct{} + +func (failingRandomReader) Read([]byte) (int, error) { + return 0, errors.New("random source unavailable") +} + +func TestRandomConvenienceHelpersDoNotPanicWhenEntropyFails(t *testing.T) { + originalReader := crand.Reader + originalSMTPFrom := SMTPFrom + crand.Reader = io.Reader(failingRandomReader{}) + SMTPFrom = "sender@example.com" + t.Cleanup(func() { + crand.Reader = originalReader + SMTPFrom = originalSMTPFrom + }) + + require.NotPanics(t, func() { + require.Equal(t, 0, GetRandomInt(10)) + require.Empty(t, GetRandomString(8)) + }) + _, err := GenerateRandomCharsKey(8) + require.ErrorContains(t, err, "random source unavailable") + _, err = GenerateRandomKey(8) + require.ErrorContains(t, err, "random source unavailable") + _, err = SecureRandomInt(10) + require.ErrorContains(t, err, "random source unavailable") + _, err = generateMessageID() + require.ErrorContains(t, err, "generate message ID") +} + +func TestGetRandomIntDoesNotPanicForInvalidBound(t *testing.T) { + require.NotPanics(t, func() { + require.Equal(t, 0, GetRandomInt(0)) + }) +} diff --git a/common/rate-limit.go b/common/rate-limit.go index 1ab14e96..9bdf9e68 100644 --- a/common/rate-limit.go +++ b/common/rate-limit.go @@ -1,6 +1,7 @@ package common import ( + "context" "sync" "time" ) @@ -9,35 +10,71 @@ type InMemoryRateLimiter struct { store map[string]*[]int64 mutex sync.Mutex expirationDuration time.Duration + initOnce sync.Once + stopOnce sync.Once + stop chan struct{} + done chan struct{} + stopped bool } func (l *InMemoryRateLimiter) Init(expirationDuration time.Duration) { - if l.store == nil { + l.initOnce.Do(func() { l.mutex.Lock() - if l.store == nil { - l.store = make(map[string]*[]int64) - l.expirationDuration = expirationDuration - if expirationDuration > 0 { - go l.clearExpiredItems() - } + defer l.mutex.Unlock() + if l.stopped { + return } - l.mutex.Unlock() - } + l.store = make(map[string]*[]int64) + l.expirationDuration = expirationDuration + if expirationDuration > 0 { + l.stop = make(chan struct{}) + l.done = make(chan struct{}) + go l.clearExpiredItems(expirationDuration, l.stop, l.done) + } + }) } -func (l *InMemoryRateLimiter) clearExpiredItems() { +func (l *InMemoryRateLimiter) clearExpiredItems(interval time.Duration, stop <-chan struct{}, done chan<- struct{}) { + defer close(done) + ticker := time.NewTicker(interval) + defer ticker.Stop() for { - time.Sleep(l.expirationDuration) - l.mutex.Lock() - now := time.Now().Unix() - for key := range l.store { - queue := l.store[key] - size := len(*queue) - if size == 0 || now-(*queue)[size-1] > int64(l.expirationDuration.Seconds()) { - delete(l.store, key) + select { + case <-stop: + return + case <-ticker.C: + l.mutex.Lock() + now := time.Now().Unix() + for key := range l.store { + queue := l.store[key] + size := len(*queue) + if size == 0 || now-(*queue)[size-1] > int64(interval.Seconds()) { + delete(l.store, key) + } } + l.mutex.Unlock() } - l.mutex.Unlock() + } +} + +func (l *InMemoryRateLimiter) Stop(ctx context.Context) error { + l.mutex.Lock() + l.stopped = true + stop := l.stop + done := l.done + l.mutex.Unlock() + if stop == nil || done == nil { + return nil + } + + l.stopOnce.Do(func() { + close(stop) + }) + select { + case <-done: + return nil + case <-ctx.Done(): + return ctx.Err() } } diff --git a/common/rate_limit_test.go b/common/rate_limit_test.go new file mode 100644 index 00000000..cf396cfe --- /dev/null +++ b/common/rate_limit_test.go @@ -0,0 +1,51 @@ +package common + +import ( + "context" + "fmt" + "sync" + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +func TestInMemoryRateLimiterConcurrentInitAndStop(t *testing.T) { + var limiter InMemoryRateLimiter + var wg sync.WaitGroup + errCh := make(chan error, 64) + + for i := 0; i < 64; i++ { + wg.Add(1) + go func(index int) { + defer wg.Done() + limiter.Init(10 * time.Millisecond) + if !limiter.Request(fmt.Sprintf("key-%d", index), 1, 1) { + errCh <- fmt.Errorf("first request for key %d was rejected", index) + } + }(i) + } + wg.Wait() + close(errCh) + for err := range errCh { + require.NoError(t, err) + } + + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + require.NoError(t, limiter.Stop(ctx)) +} + +func TestInMemoryRateLimiterStopBeforeInitPreventsLifecycleStart(t *testing.T) { + var limiter InMemoryRateLimiter + + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + require.NoError(t, limiter.Stop(ctx)) + + limiter.Init(10 * time.Millisecond) + + require.Nil(t, limiter.store) + require.Nil(t, limiter.stop) + require.Nil(t, limiter.done) +} diff --git a/common/str.go b/common/str.go index add09966..1daea64b 100644 --- a/common/str.go +++ b/common/str.go @@ -7,9 +7,8 @@ import ( "regexp" "strconv" "strings" + "unicode/utf8" "unsafe" - - "github.com/samber/lo" ) var ( @@ -21,6 +20,8 @@ var ( ) const LocalLogContentLimit = 2048 +const PersistedLogContentLimit = 4096 +const persistedLogContentTruncatedSuffix = "... [truncated]" // LocalLogPreview limits log-only content unless debug logging is enabled. func LocalLogPreview(content string) string { @@ -30,6 +31,43 @@ func LocalLogPreview(content string) string { return fmt.Sprintf("%s... [truncated, original_length=%d, limit=%d]", content[:LocalLogContentLimit], len(content), LocalLogContentLimit) } +func SanitizePersistedLogContent(content string) string { + if content == "" { + return "" + } + content = strings.ToValidUTF8(content, "") + var builder strings.Builder + builder.Grow(min(len(content), PersistedLogContentLimit)) + count := 0 + truncated := false + for _, r := range content { + switch r { + case '\r', '\n', '\t': + r = ' ' + default: + if r < 0x20 || r == 0x7f { + continue + } + } + if count >= PersistedLogContentLimit { + truncated = true + break + } + builder.WriteRune(r) + count++ + } + sanitized := strings.TrimSpace(builder.String()) + if truncated { + retainedRuneLimit := max(0, PersistedLogContentLimit-utf8.RuneCountInString(persistedLogContentTruncatedSuffix)) + sanitizedRunes := []rune(sanitized) + if len(sanitizedRunes) > retainedRuneLimit { + sanitized = strings.TrimSpace(string(sanitizedRunes[:retainedRuneLimit])) + } + sanitized += persistedLogContentTruncatedSuffix + } + return sanitized +} + func GetStringIfEmpty(str string, defaultValue string) string { if str == "" { return defaultValue @@ -41,7 +79,12 @@ func GetRandomString(length int) string { if length <= 0 { return "" } - return lo.RandomString(length, lo.AlphanumericCharset) + value, err := GenerateRandomCharsKey(length) + if err != nil { + SysError("failed to generate secure random string: " + err.Error()) + return "" + } + return value } func MapToJsonStr(m map[string]interface{}) string { diff --git a/common/str_test.go b/common/str_test.go new file mode 100644 index 00000000..cbce177d --- /dev/null +++ b/common/str_test.go @@ -0,0 +1,42 @@ +package common + +import ( + "strings" + "testing" + "unicode/utf8" + + "github.com/stretchr/testify/require" +) + +func TestSanitizePersistedLogContentNormalizesControlsAndInvalidUTF8(t *testing.T) { + input := " first\r\nsecond\tthird\x00" + string([]byte{0xff, 0xfe}) + " " + + got := SanitizePersistedLogContent(input) + + require.Equal(t, "first second third", got) + require.True(t, utf8.ValidString(got)) + require.NotContains(t, got, "\r") + require.NotContains(t, got, "\n") + require.NotContains(t, got, "\t") + require.NotContains(t, got, "\x00") +} + +func TestSanitizePersistedLogContentTruncatesByRune(t *testing.T) { + input := strings.Repeat("x", PersistedLogContentLimit+1) + + got := SanitizePersistedLogContent(input) + + suffixRuneCount := utf8.RuneCountInString(persistedLogContentTruncatedSuffix) + require.Equal(t, strings.Repeat("x", PersistedLogContentLimit-suffixRuneCount)+persistedLogContentTruncatedSuffix, got) + require.LessOrEqual(t, utf8.RuneCountInString(got), PersistedLogContentLimit) + require.True(t, strings.HasSuffix(got, persistedLogContentTruncatedSuffix)) +} + +func TestSanitizePersistedLogContentKeepsContentWithinLimit(t *testing.T) { + input := strings.Repeat("界", PersistedLogContentLimit) + + got := SanitizePersistedLogContent(input) + + require.Equal(t, input, got) + require.False(t, strings.HasSuffix(got, persistedLogContentTruncatedSuffix)) +} diff --git a/common/system_monitor.go b/common/system_monitor.go index 26710fae..a3981d74 100644 --- a/common/system_monitor.go +++ b/common/system_monitor.go @@ -1,6 +1,8 @@ package common import ( + "context" + "sync" "sync/atomic" "time" @@ -28,6 +30,11 @@ type SystemStatus struct { } var latestSystemStatus atomic.Value +var systemMonitorState struct { + sync.Mutex + cancel context.CancelFunc + done chan struct{} +} func init() { latestSystemStatus.Store(SystemStatus{}) @@ -35,18 +42,67 @@ func init() { // StartSystemMonitor 启动系统监控 func StartSystemMonitor() { + systemMonitorState.Lock() + defer systemMonitorState.Unlock() + if systemMonitorState.cancel != nil { + return + } + ctx, cancel := context.WithCancel(context.Background()) + done := make(chan struct{}) + systemMonitorState.cancel = cancel + systemMonitorState.done = done go func() { - for { - config := GetPerformanceMonitorConfig() - if !config.Enabled { - time.Sleep(30 * time.Second) - continue - } + defer close(done) + runSystemMonitor(ctx) + }() +} +func runSystemMonitor(ctx context.Context) { + for { + interval := 30 * time.Second + config := GetPerformanceMonitorConfig() + if config.Enabled { updateSystemStatus() - time.Sleep(5 * time.Second) + interval = 5 * time.Second } - }() + + timer := time.NewTimer(interval) + select { + case <-ctx.Done(): + if !timer.Stop() { + select { + case <-timer.C: + default: + } + } + return + case <-timer.C: + } + } +} + +func StopSystemMonitor(ctx context.Context) error { + systemMonitorState.Lock() + cancel := systemMonitorState.cancel + done := systemMonitorState.done + systemMonitorState.Unlock() + if cancel == nil || done == nil { + return nil + } + + cancel() + select { + case <-done: + systemMonitorState.Lock() + if systemMonitorState.done == done { + systemMonitorState.cancel = nil + systemMonitorState.done = nil + } + systemMonitorState.Unlock() + return nil + case <-ctx.Done(): + return ctx.Err() + } } func updateSystemStatus() { diff --git a/common/system_monitor_test.go b/common/system_monitor_test.go new file mode 100644 index 00000000..3b4f69ed --- /dev/null +++ b/common/system_monitor_test.go @@ -0,0 +1,18 @@ +package common + +import ( + "context" + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +func TestSystemMonitorCanStopAndRestart(t *testing.T) { + for i := 0; i < 2; i++ { + StartSystemMonitor() + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + require.NoError(t, StopSystemMonitor(ctx)) + cancel() + } +} diff --git a/common/trusted_proxy.go b/common/trusted_proxy.go new file mode 100644 index 00000000..65bc0113 --- /dev/null +++ b/common/trusted_proxy.go @@ -0,0 +1,147 @@ +package common + +import ( + "fmt" + "net" + "strings" +) + +func parseTrustedProxies(raw string) ([]string, error) { + values := DefaultTrustedProxies + if strings.TrimSpace(raw) != "" { + values = strings.Split(raw, ",") + } + + proxies := make([]string, 0, len(values)) + prefixes := make([]trustedProxyPrefix, 0, len(values)) + for _, value := range values { + proxy := strings.TrimSpace(value) + if proxy == "" { + continue + } + + prefix, err := parseTrustedProxyPrefix(proxy) + if err != nil { + return nil, err + } + if prefix.prefix == 0 { + return nil, unrestrictedTrustedProxyError(proxy, prefix.bits) + } + + prefixes = append(prefixes, prefix) + proxies = append(proxies, proxy) + } + + for _, bits := range []int{32, 128} { + if trustedProxyPrefixesCoverAll(prefixes, bits) { + return nil, fmt.Errorf("TRUSTED_PROXIES entries trust every %s address; configure a specific proxy address or network", trustedProxyAddressFamily(bits)) + } + } + + return proxies, nil +} + +type trustedProxyPrefix struct { + ip net.IP + bits int + prefix int +} + +func parseTrustedProxyPrefix(proxy string) (trustedProxyPrefix, error) { + if strings.Contains(proxy, "/") { + _, network, err := net.ParseCIDR(proxy) + if err != nil { + return trustedProxyPrefix{}, fmt.Errorf("invalid TRUSTED_PROXIES entry %q: %w", proxy, err) + } + + ip := network.IP + mask := network.Mask + if ipv4 := ip.To4(); ipv4 != nil { + if len(mask) == net.IPv6len { + mask = mask[net.IPv6len-net.IPv4len:] + } + ones, bits := mask.Size() + if bits != net.IPv4len*8 { + return trustedProxyPrefix{}, fmt.Errorf("invalid TRUSTED_PROXIES entry %q", proxy) + } + return trustedProxyPrefix{ip: ipv4, bits: bits, prefix: ones}, nil + } + + ones, bits := mask.Size() + if bits != net.IPv6len*8 { + return trustedProxyPrefix{}, fmt.Errorf("invalid TRUSTED_PROXIES entry %q", proxy) + } + return trustedProxyPrefix{ip: ip, bits: bits, prefix: ones}, nil + } + + ip := net.ParseIP(proxy) + if ip == nil { + return trustedProxyPrefix{}, fmt.Errorf("invalid TRUSTED_PROXIES entry %q", proxy) + } + if ipv4 := ip.To4(); ipv4 != nil { + return trustedProxyPrefix{ip: ipv4, bits: net.IPv4len * 8, prefix: net.IPv4len * 8}, nil + } + return trustedProxyPrefix{ip: ip, bits: net.IPv6len * 8, prefix: net.IPv6len * 8}, nil +} + +func unrestrictedTrustedProxyError(proxy string, bits int) error { + return fmt.Errorf("TRUSTED_PROXIES entry %q trusts every %s address; configure a specific proxy address or network", proxy, trustedProxyAddressFamily(bits)) +} + +func trustedProxyAddressFamily(bits int) string { + if bits == net.IPv4len*8 { + return "IPv4" + } + return "IPv6" +} + +type trustedProxyCoverageNode struct { + covered bool + children [2]*trustedProxyCoverageNode +} + +func trustedProxyPrefixesCoverAll(prefixes []trustedProxyPrefix, bits int) bool { + root := &trustedProxyCoverageNode{} + for _, prefix := range prefixes { + if prefix.bits != bits { + continue + } + addTrustedProxyPrefix(root, prefix) + } + return trustedProxyCoverageComplete(root) +} + +func addTrustedProxyPrefix(root *trustedProxyCoverageNode, prefix trustedProxyPrefix) { + if root.covered { + return + } + if prefix.prefix == 0 { + root.covered = true + return + } + + node := root + for bitIndex := 0; bitIndex < prefix.prefix; bitIndex++ { + byteIndex := bitIndex / 8 + bitOffset := 7 - (bitIndex % 8) + branch := (prefix.ip[byteIndex] >> bitOffset) & 1 + if node.children[branch] == nil { + node.children[branch] = &trustedProxyCoverageNode{} + } + node = node.children[branch] + if node.covered { + return + } + } + node.covered = true +} + +func trustedProxyCoverageComplete(node *trustedProxyCoverageNode) bool { + if node == nil { + return false + } + if node.covered { + return true + } + return trustedProxyCoverageComplete(node.children[0]) && trustedProxyCoverageComplete(node.children[1]) +} diff --git a/common/trusted_proxy_test.go b/common/trusted_proxy_test.go index 226e6189..ba95da7d 100644 --- a/common/trusted_proxy_test.go +++ b/common/trusted_proxy_test.go @@ -6,9 +6,43 @@ import ( ) func TestDefaultTrustedProxiesOnlyTrustLoopback(t *testing.T) { - for _, proxy := range []string{"10.0.0.0/8", "172.16.0.0/12", "192.168.0.0/16", "169.254.0.0/16", "fc00::/7", "fe80::/10"} { - if slices.Contains(DefaultTrustedProxies, proxy) { - t.Fatalf("default trusted proxies must not include broad private range %s", proxy) + expected := []string{"127.0.0.0/8", "::1/128"} + if !slices.Equal(DefaultTrustedProxies, expected) { + t.Fatalf("default trusted proxies must only include loopback networks: %v", DefaultTrustedProxies) + } +} + +func TestParseTrustedProxiesUsesLoopbackDefaults(t *testing.T) { + proxies, err := parseTrustedProxies("") + if err != nil { + t.Fatalf("parseTrustedProxies returned error: %v", err) + } + if !slices.Equal(proxies, DefaultTrustedProxies) { + t.Fatalf("unexpected default trusted proxies: %v", proxies) + } +} + +func TestParseTrustedProxiesRejectsUnrestrictedNetworks(t *testing.T) { + for _, value := range []string{ + "0.0.0.0/0", + "::/0", + "::ffff:0:0/96", + "0.0.0.0/1,128.0.0.0/1", + "::/1,8000::/1", + "not-a-network", + } { + if _, err := parseTrustedProxies(value); err == nil { + t.Fatalf("expected unrestricted proxy network %q to be rejected", value) } } } + +func TestParseTrustedProxiesAcceptsExplicitProxyNetworks(t *testing.T) { + proxies, err := parseTrustedProxies("127.0.0.1, 10.0.0.0/8, ::ffff:192.0.2.1/128") + if err != nil { + t.Fatalf("parseTrustedProxies rejected explicit proxy networks: %v", err) + } + if !slices.Equal(proxies, []string{"127.0.0.1", "10.0.0.0/8", "::ffff:192.0.2.1/128"}) { + t.Fatalf("unexpected explicit trusted proxies: %v", proxies) + } +} diff --git a/common/utils.go b/common/utils.go index 8d50903e..08878f34 100644 --- a/common/utils.go +++ b/common/utils.go @@ -8,7 +8,6 @@ import ( "io" "log" "math/big" - "math/rand" "net" "net/url" "os" @@ -241,7 +240,7 @@ func GenerateRandomCharsKey(length int) (string, error) { func GenerateRandomKey(length int) (string, error) { bytes := make([]byte, length*3/4) // 对于48位的输出，这里应该是36 - if _, err := crand.Read(bytes); err != nil { + if _, err := io.ReadFull(crand.Reader, bytes); err != nil { return "", err } return base64.StdEncoding.EncodeToString(bytes), nil @@ -253,8 +252,23 @@ func GenerateKey() (string, error) { } func GetRandomInt(max int) int { - //rand.Seed(time.Now().UnixNano()) - return rand.Intn(max) + n, err := SecureRandomInt(max) + if err != nil { + SysError("failed to generate secure random integer: " + err.Error()) + return 0 + } + return n +} + +func SecureRandomInt(max int) (int, error) { + if max <= 0 { + return 0, fmt.Errorf("random max must be greater than 0") + } + n, err := crand.Int(crand.Reader, big.NewInt(int64(max))) + if err != nil { + return 0, err + } + return int(n.Int64()), nil } func GetTimestamp() int64 { @@ -280,7 +294,11 @@ func MessageWithRequestId(message string, id string) string { func RandomSleep() { // Sleep for 0-3000 ms - time.Sleep(time.Duration(rand.Intn(3000)) * time.Millisecond) + delay, err := SecureRandomInt(3000) + if err != nil { + delay = 0 + } + time.Sleep(time.Duration(delay) * time.Millisecond) } func GetPointer[T any](v T) *T { diff --git a/controller/channel.go b/controller/channel.go index 345df4b8..386b27ac 100644 --- a/controller/channel.go +++ b/controller/channel.go @@ -920,29 +920,162 @@ func DeleteChannelBatch(c *gin.Context) { } type PatchChannel struct { - model.Channel - MultiKeyMode *string `json:"multi_key_mode"` - KeyMode *string `json:"key_mode"` // 多key模式下密钥覆盖或者追加 + Id int `json:"id"` + Type *int `json:"type"` + Key *string `json:"key"` + OpenAIOrganization *string `json:"openai_organization"` + TestModel *string `json:"test_model"` + Status *int `json:"status"` + Name *string `json:"name"` + Weight *uint `json:"weight"` + BaseURL *string `json:"base_url"` + Other *string `json:"other"` + Models *string `json:"models"` + Group *string `json:"group"` + ModelMapping *string `json:"model_mapping"` + StatusCodeMapping *string `json:"status_code_mapping"` + Priority *int64 `json:"priority"` + AutoBan *int `json:"auto_ban"` + Tag *string `json:"tag"` + Setting *string `json:"setting"` + ParamOverride *string `json:"param_override"` + HeaderOverride *string `json:"header_override"` + Remark *string `json:"remark"` + OtherSettings *string `json:"settings"` + MultiKeyMode *string `json:"multi_key_mode"` + KeyMode *string `json:"key_mode"` // 多key模式下密钥覆盖或者追加 +} + +func (patch *PatchChannel) applyTo(origin *model.Channel) (model.Channel, []string, error) { + channel := *origin + fields := make([]string, 0, 24) + if patch.Type != nil { + channel.Type = *patch.Type + fields = append(fields, "type") + } + if patch.Key != nil && strings.TrimSpace(*patch.Key) != "" { + channel.Key = *patch.Key + fields = append(fields, "key") + } + if patch.OpenAIOrganization != nil { + channel.OpenAIOrganization = patch.OpenAIOrganization + fields = append(fields, "openai_organization") + } + if patch.TestModel != nil { + channel.TestModel = patch.TestModel + fields = append(fields, "test_model") + } + if patch.Status != nil { + if !isEditableChannelStatus(*patch.Status) { + return channel, nil, fmt.Errorf("invalid channel status") + } + channel.Status = *patch.Status + fields = append(fields, "status") + } + if patch.Name != nil { + channel.Name = *patch.Name + fields = append(fields, "name") + } + if patch.Weight != nil { + channel.Weight = patch.Weight + fields = append(fields, "weight") + } + if patch.BaseURL != nil { + channel.BaseURL = patch.BaseURL + fields = append(fields, "base_url") + } + if patch.Other != nil { + channel.Other = *patch.Other + fields = append(fields, "other") + } + if patch.Models != nil { + channel.Models = *patch.Models + fields = append(fields, "models") + } + if patch.Group != nil { + channel.Group = *patch.Group + fields = append(fields, "group") + } + if patch.ModelMapping != nil { + channel.ModelMapping = patch.ModelMapping + fields = append(fields, "model_mapping") + } + if patch.StatusCodeMapping != nil { + channel.StatusCodeMapping = patch.StatusCodeMapping + fields = append(fields, "status_code_mapping") + } + if patch.Priority != nil { + channel.Priority = patch.Priority + fields = append(fields, "priority") + } + if patch.AutoBan != nil { + channel.AutoBan = patch.AutoBan + fields = append(fields, "auto_ban") + } + if patch.Tag != nil { + channel.Tag = patch.Tag + fields = append(fields, "tag") + } + if patch.Setting != nil { + channel.Setting = patch.Setting + fields = append(fields, "setting") + } + if patch.ParamOverride != nil { + channel.ParamOverride = patch.ParamOverride + fields = append(fields, "param_override") + } + if patch.HeaderOverride != nil { + channel.HeaderOverride = patch.HeaderOverride + fields = append(fields, "header_override") + } + if patch.Remark != nil { + channel.Remark = patch.Remark + fields = append(fields, "remark") + } + if patch.OtherSettings != nil { + channel.OtherSettings = *patch.OtherSettings + fields = append(fields, "settings") + } + return channel, fields, nil +} + +func isEditableChannelStatus(status int) bool { + switch status { + case common.ChannelStatusEnabled, common.ChannelStatusManuallyDisabled, common.ChannelStatusAutoDisabled: + return true + default: + return false + } } func UpdateChannel(c *gin.Context) { - channel := PatchChannel{} - err := c.ShouldBindJSON(&channel) + patch := PatchChannel{} + err := c.ShouldBindJSON(&patch) if err != nil { common.ApiError(c, err) return } - // 使用统一的校验函数 - if err := validateChannel(&channel.Channel, false); err != nil { + pollingLock := model.GetChannelPollingLock(patch.Id) + pollingLock.Lock() + pollingLocked := true + defer func() { + if pollingLocked { + pollingLock.Unlock() + } + }() + + // Preserve existing ChannelInfo to ensure multi-key channels keep correct state even if the client does not send ChannelInfo in the request. + originChannel, err := model.GetChannelById(patch.Id, true) + if err != nil { c.JSON(http.StatusOK, gin.H{ "success": false, "message": err.Error(), }) return } - // Preserve existing ChannelInfo to ensure multi-key channels keep correct state even if the client does not send ChannelInfo in the request. - originChannel, err := model.GetChannelById(channel.Id, true) + + channel, updateFields, err := patch.applyTo(originChannel) if err != nil { c.JSON(http.StatusOK, gin.H{ "success": false, @@ -955,13 +1088,23 @@ func UpdateChannel(c *gin.Context) { channel.ChannelInfo = originChannel.ChannelInfo // If the request explicitly specifies a new MultiKeyMode, apply it on top of the original info. - if channel.MultiKeyMode != nil && *channel.MultiKeyMode != "" { - channel.ChannelInfo.MultiKeyMode = constant.MultiKeyMode(*channel.MultiKeyMode) + if patch.MultiKeyMode != nil && *patch.MultiKeyMode != "" { + channel.ChannelInfo.MultiKeyMode = constant.MultiKeyMode(*patch.MultiKeyMode) + updateFields = append(updateFields, "channel_info") + } + + // 使用统一的校验函数 + if err := validateChannel(&channel, false); err != nil { + c.JSON(http.StatusOK, gin.H{ + "success": false, + "message": err.Error(), + }) + return } // 处理多key模式下的密钥追加/覆盖逻辑 - if channel.KeyMode != nil && channel.ChannelInfo.IsMultiKey { - switch *channel.KeyMode { + if patch.KeyMode != nil && channel.ChannelInfo.IsMultiKey && patch.Key != nil && strings.TrimSpace(*patch.Key) != "" { + switch *patch.KeyMode { case "append": // 追加模式：将新密钥添加到现有密钥列表 if originChannel.Key != "" { @@ -986,8 +1129,8 @@ func UpdateChannel(c *gin.Context) { // 处理 Vertex AI 的特殊情况 if channel.Type == constant.ChannelTypeVertexAi && channel.GetOtherSettings().VertexKeyType != dto.VertexKeyTypeAPIKey { // 尝试解析新密钥为JSON数组 - if strings.HasPrefix(strings.TrimSpace(channel.Key), "[") { - array, err := getVertexArrayKeys(channel.Key) + if strings.HasPrefix(strings.TrimSpace(*patch.Key), "[") { + array, err := getVertexArrayKeys(*patch.Key) if err != nil { c.JSON(http.StatusOK, gin.H{ "success": false, @@ -998,11 +1141,11 @@ func UpdateChannel(c *gin.Context) { newKeys = array } else { // 单个JSON密钥 - newKeys = []string{channel.Key} + newKeys = []string{*patch.Key} } } else { // 普通渠道的处理 - inputKeys := strings.Split(channel.Key, "\n") + inputKeys := strings.Split(*patch.Key, "\n") for _, key := range inputKeys { key = strings.TrimSpace(key) if key != "" { @@ -1034,12 +1177,16 @@ func UpdateChannel(c *gin.Context) { allKeys := append(existingKeys, dedupedNewKeys...) channel.Key = strings.Join(allKeys, "\n") + updateFields = append(updateFields, "key", "channel_info") } case "replace": // 覆盖模式：直接使用新密钥（默认行为，不需要特殊处理） + updateFields = append(updateFields, "channel_info") } } - err = channel.Update() + err = channel.UpdateFields(updateFields...) + pollingLock.Unlock() + pollingLocked = false if err != nil { common.ApiError(c, err) return @@ -1071,7 +1218,7 @@ func UpdateChannel(c *gin.Context) { "changed_fields": changedFields, }) channel.Key = "" - clearChannelInfo(&channel.Channel) + clearChannelInfo(&channel) c.JSON(http.StatusOK, gin.H{ "success": true, "message": "", @@ -1571,7 +1718,7 @@ func ManageMultiKeys(c *gin.Context) { channel.ChannelInfo.MultiKeyStatusList[keyIndex] = 2 // disabled - err = channel.Update() + err = channel.UpdateFields("channel_info") if err != nil { common.ApiError(c, err) return @@ -1614,7 +1761,7 @@ func ManageMultiKeys(c *gin.Context) { delete(channel.ChannelInfo.MultiKeyDisabledReason, keyIndex) } - err = channel.Update() + err = channel.UpdateFields("channel_info") if err != nil { common.ApiError(c, err) return @@ -1639,7 +1786,7 @@ func ManageMultiKeys(c *gin.Context) { channel.ChannelInfo.MultiKeyDisabledTime = make(map[int]int64) channel.ChannelInfo.MultiKeyDisabledReason = make(map[int]string) - err = channel.Update() + err = channel.UpdateFields("channel_info") if err != nil { common.ApiError(c, err) return @@ -1687,7 +1834,7 @@ func ManageMultiKeys(c *gin.Context) { return } - err = channel.Update() + err = channel.UpdateFields("channel_info") if err != nil { common.ApiError(c, err) return @@ -1768,7 +1915,7 @@ func ManageMultiKeys(c *gin.Context) { channel.ChannelInfo.MultiKeyDisabledTime = newDisabledTime channel.ChannelInfo.MultiKeyDisabledReason = newDisabledReason - err = channel.Update() + err = channel.UpdateFields("key", "channel_info") if err != nil { common.ApiError(c, err) return @@ -1837,7 +1984,7 @@ func ManageMultiKeys(c *gin.Context) { channel.ChannelInfo.MultiKeyDisabledTime = newDisabledTime channel.ChannelInfo.MultiKeyDisabledReason = newDisabledReason - err = channel.Update() + err = channel.UpdateFields("key", "channel_info") if err != nil { common.ApiError(c, err) return diff --git a/controller/channel_update_test.go b/controller/channel_update_test.go new file mode 100644 index 00000000..40776477 --- /dev/null +++ b/controller/channel_update_test.go @@ -0,0 +1,211 @@ +package controller + +import ( + "net/http" + "testing" + "time" + + "github.com/MAX-API-Next/MAX-API/common" + "github.com/MAX-API-Next/MAX-API/model" + "github.com/stretchr/testify/require" + "gorm.io/gorm" +) + +func setupChannelUpdateControllerTestDB(t *testing.T) (*gorm.DB, *model.Channel) { + t.Helper() + + db := openTokenControllerTestDB(t) + require.NoError(t, db.AutoMigrate(&model.Channel{}, &model.Ability{}, &model.Log{}, &model.User{})) + + priority := int64(10) + weight := uint(20) + baseURL := "https://origin.example.com" + tag := "stable" + channel := &model.Channel{ + Name: "origin-channel", + Type: 1, + Key: "origin-key", + Status: common.ChannelStatusManuallyDisabled, + Models: "gpt-origin", + Group: "default", + Priority: &priority, + Weight: &weight, + BaseURL: &baseURL, + Tag: &tag, + Balance: 12.5, + BalanceUpdatedTime: 111, + UsedQuota: 222, + CreatedTime: 333, + TestTime: 444, + ResponseTime: 555, + ChannelInfo: model.ChannelInfo{ + IsMultiKey: true, + MultiKeySize: 1, + }, + } + require.NoError(t, db.Create(channel).Error) + require.NoError(t, channel.AddAbilities(nil)) + return db, channel +} + +func TestUpdateChannelIgnoresMassAssignedRuntimeFields(t *testing.T) { + db, channel := setupChannelUpdateControllerTestDB(t) + + priority := int64(30) + weight := uint(40) + body := map[string]any{ + "id": channel.Id, + "name": "updated-channel", + "status": common.ChannelStatusEnabled, + "priority": priority, + "weight": weight, + "balance": 999999.99, + "balance_updated_time": 999, + "used_quota": 0, + "created_time": 999, + "test_time": 999, + "response_time": 999, + "other_info": `{"status_reason":"forged"}`, + "channel_info": map[string]any{ + "is_multi_key": false, + "multi_key_size": 99, + }, + } + + ctx, recorder := newAuthenticatedContext(t, http.MethodPut, "/api/channel/", body, 1) + UpdateChannel(ctx) + response := decodeAPIResponse(t, recorder) + require.True(t, response.Success, response.Message) + + var stored model.Channel + require.NoError(t, db.First(&stored, channel.Id).Error) + require.Equal(t, "updated-channel", stored.Name) + require.Equal(t, common.ChannelStatusEnabled, stored.Status) + require.EqualValues(t, priority, *stored.Priority) + require.EqualValues(t, weight, *stored.Weight) + require.Equal(t, 12.5, stored.Balance) + require.EqualValues(t, 111, stored.BalanceUpdatedTime) + require.EqualValues(t, 222, stored.UsedQuota) + require.EqualValues(t, 333, stored.CreatedTime) + require.EqualValues(t, 444, stored.TestTime) + require.EqualValues(t, 555, stored.ResponseTime) + require.Empty(t, stored.OtherInfo) + require.True(t, stored.ChannelInfo.IsMultiKey) + require.Equal(t, 1, stored.ChannelInfo.MultiKeySize) +} + +func TestUpdateChannelPartialStatusDoesNotClearConfiguration(t *testing.T) { + db, channel := setupChannelUpdateControllerTestDB(t) + + ctx, recorder := newAuthenticatedContext(t, http.MethodPut, "/api/channel/", map[string]any{ + "id": channel.Id, + "status": common.ChannelStatusEnabled, + }, 1) + UpdateChannel(ctx) + response := decodeAPIResponse(t, recorder) + require.True(t, response.Success, response.Message) + + var stored model.Channel + require.NoError(t, db.First(&stored, channel.Id).Error) + require.Equal(t, common.ChannelStatusEnabled, stored.Status) + require.Equal(t, "origin-channel", stored.Name) + require.Equal(t, "origin-key", stored.Key) + require.Equal(t, "gpt-origin", stored.Models) + require.Equal(t, "default", stored.Group) + require.NotNil(t, stored.BaseURL) + require.Equal(t, "https://origin.example.com", *stored.BaseURL) +} + +func TestUpdateChannelRejectsUnknownStatus(t *testing.T) { + db, channel := setupChannelUpdateControllerTestDB(t) + + ctx, recorder := newAuthenticatedContext(t, http.MethodPut, "/api/channel/", map[string]any{ + "id": channel.Id, + "status": common.ChannelStatusUnknown, + }, 1) + UpdateChannel(ctx) + response := decodeAPIResponse(t, recorder) + require.False(t, response.Success) + require.Contains(t, response.Message, "invalid channel status") + + var stored model.Channel + require.NoError(t, db.First(&stored, channel.Id).Error) + require.Equal(t, common.ChannelStatusManuallyDisabled, stored.Status) +} + +func TestUpdateChannelAppendsMultiKeysAndUpdatesSize(t *testing.T) { + db, channel := setupChannelUpdateControllerTestDB(t) + + ctx, recorder := newAuthenticatedContext(t, http.MethodPut, "/api/channel/", map[string]any{ + "id": channel.Id, + "key": "second-key\norigin-key", + "key_mode": "append", + }, 1) + UpdateChannel(ctx) + response := decodeAPIResponse(t, recorder) + require.True(t, response.Success, response.Message) + + var stored model.Channel + require.NoError(t, db.First(&stored, channel.Id).Error) + require.Equal(t, "origin-key\nsecond-key", stored.Key) + require.Equal(t, 2, stored.ChannelInfo.MultiKeySize) +} + +func TestUpdateChannelReplacesMultiKeysAndUpdatesSize(t *testing.T) { + db, channel := setupChannelUpdateControllerTestDB(t) + + ctx, recorder := newAuthenticatedContext(t, http.MethodPut, "/api/channel/", map[string]any{ + "id": channel.Id, + "key": "replacement-a\nreplacement-b", + "key_mode": "replace", + }, 1) + UpdateChannel(ctx) + response := decodeAPIResponse(t, recorder) + require.True(t, response.Success, response.Message) + + var stored model.Channel + require.NoError(t, db.First(&stored, channel.Id).Error) + require.Equal(t, "replacement-a\nreplacement-b", stored.Key) + require.Equal(t, 2, stored.ChannelInfo.MultiKeySize) +} + +func TestUpdateChannelWaitsForChannelPollingLock(t *testing.T) { + _, channel := setupChannelUpdateControllerTestDB(t) + + lock := model.GetChannelPollingLock(channel.Id) + lock.Lock() + locked := true + t.Cleanup(func() { + if locked { + lock.Unlock() + } + }) + + ctx, recorder := newAuthenticatedContext(t, http.MethodPut, "/api/channel/", map[string]any{ + "id": channel.Id, + "status": common.ChannelStatusEnabled, + }, 1) + done := make(chan struct{}) + go func() { + UpdateChannel(ctx) + close(done) + }() + + select { + case <-done: + lock.Unlock() + locked = false + t.Fatal("channel update completed while polling lock was held") + case <-time.After(150 * time.Millisecond): + } + + lock.Unlock() + locked = false + select { + case <-done: + response := decodeAPIResponse(t, recorder) + require.True(t, response.Success, response.Message) + case <-time.After(2 * time.Second): + t.Fatal("channel update did not finish after polling lock was released") + } +} diff --git a/controller/codex_oauth.go b/controller/codex_oauth.go index be2f96c6..ae83bbe4 100644 --- a/controller/codex_oauth.go +++ b/controller/codex_oauth.go @@ -164,8 +164,13 @@ func completeCodexOAuthWithChannelID(c *gin.Context, channelID int) { } session := sessions.Default(c) - expectedState, _ := session.Get(codexOAuthSessionKey(channelID, "state")).(string) - verifier, _ := session.Get(codexOAuthSessionKey(channelID, "verifier")).(string) + expectedState, stateOK := session.Get(codexOAuthSessionKey(channelID, "state")).(string) + verifier, verifierOK := session.Get(codexOAuthSessionKey(channelID, "verifier")).(string) + if !stateOK || !verifierOK { + common.SysError(fmt.Sprintf("invalid Codex OAuth session value types (channel_id=%d)", channelID)) + c.JSON(http.StatusOK, gin.H{"success": false, "message": "oauth flow not started or session expired"}) + return + } if strings.TrimSpace(expectedState) == "" || strings.TrimSpace(verifier) == "" { c.JSON(http.StatusOK, gin.H{"success": false, "message": "oauth flow not started or session expired"}) return diff --git a/controller/custom_oauth.go b/controller/custom_oauth.go index 6c71640c..e1a67538 100644 --- a/controller/custom_oauth.go +++ b/controller/custom_oauth.go @@ -12,6 +12,7 @@ import ( "github.com/MAX-API-Next/MAX-API/common" "github.com/MAX-API-Next/MAX-API/model" "github.com/MAX-API-Next/MAX-API/oauth" + "github.com/MAX-API-Next/MAX-API/service" "github.com/gin-gonic/gin" ) @@ -166,6 +167,11 @@ func FetchCustomOAuthDiscovery(c *gin.Context) { return } + if err := service.ValidateSSRFProtectedFetchURL(targetURL); err != nil { + common.ApiErrorMsg(c, "Discovery URL invalid: "+err.Error()) + return + } + ctx, cancel := context.WithTimeout(c.Request.Context(), 20*time.Second) defer cancel() @@ -176,8 +182,7 @@ func FetchCustomOAuthDiscovery(c *gin.Context) { } httpReq.Header.Set("Accept", "application/json") - client := &http.Client{Timeout: 20 * time.Second} - resp, err := client.Do(httpReq) + resp, err := service.GetSSRFProtectedHttpClient().Do(httpReq) if err != nil { common.ApiErrorMsg(c, "获取 Discovery 配置失败: "+err.Error()) return diff --git a/controller/custom_oauth_test.go b/controller/custom_oauth_test.go new file mode 100644 index 00000000..75e444c3 --- /dev/null +++ b/controller/custom_oauth_test.go @@ -0,0 +1,55 @@ +package controller + +import ( + "bytes" + "net/http" + "net/http/httptest" + "net/url" + "testing" + + "github.com/MAX-API-Next/MAX-API/common" + "github.com/MAX-API-Next/MAX-API/setting/system_setting" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +func TestFetchCustomOAuthDiscoveryRejectsPrivateURLWhenSSRFProtectionEnabled(t *testing.T) { + gin.SetMode(gin.TestMode) + serverHit := false + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + serverHit = true + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"issuer":"http://127.0.0.1","authorization_endpoint":"http://127.0.0.1/auth"}`)) + })) + defer server.Close() + + fetchSetting := system_setting.GetFetchSetting() + original := *fetchSetting + t.Cleanup(func() { + *fetchSetting = original + }) + fetchSetting.EnableSSRFProtection = true + fetchSetting.AllowPrivateIp = false + fetchSetting.DomainFilterMode = false + fetchSetting.IpFilterMode = false + fetchSetting.DomainList = nil + fetchSetting.IpList = nil + serverURL, err := url.Parse(server.URL) + require.NoError(t, err) + fetchSetting.AllowedPorts = []string{"80", "443", "8080", "8443", serverURL.Port()} + fetchSetting.ApplyIPFilterForDomain = true + + payload, err := common.Marshal(FetchCustomOAuthDiscoveryRequest{WellKnownURL: server.URL}) + require.NoError(t, err) + recorder := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(recorder) + ctx.Request = httptest.NewRequest(http.MethodPost, "/api/custom-oauth-provider/discovery", bytes.NewReader(payload)) + ctx.Request.Header.Set("Content-Type", "application/json") + + FetchCustomOAuthDiscovery(ctx) + + require.Equal(t, http.StatusOK, recorder.Code) + require.Contains(t, recorder.Body.String(), `"success":false`) + require.Contains(t, recorder.Body.String(), "private IP address not allowed") + require.False(t, serverHit) +} diff --git a/controller/midjourney.go b/controller/midjourney.go index 3a8e9fb3..3f71a713 100644 --- a/controller/midjourney.go +++ b/controller/midjourney.go @@ -3,7 +3,6 @@ package controller import ( "bytes" "context" - "encoding/json" "fmt" "io" "net/http" @@ -21,237 +20,130 @@ import ( ) func UpdateMidjourneyTaskBulk() { - //imageModel := "midjourney" - ctx := context.TODO() - for { - time.Sleep(time.Duration(15) * time.Second) + ctx := context.Background() + ticker := time.NewTicker(15 * time.Second) + defer ticker.Stop() + for range ticker.C { + func() { + defer func() { + if recovered := recover(); recovered != nil { + logger.LogError(ctx, fmt.Sprintf("Midjourney polling cycle panic: %v", recovered)) + } + }() + if err := updateMidjourneyTasksOnce(ctx); err != nil { + logger.LogError(ctx, "UpdateMidjourneyTask cycle error: "+err.Error()) + } + }() + } +} - tasks := model.GetAllUnFinishTasks() - if len(tasks) == 0 { +func updateMidjourneyTasksOnce(ctx context.Context) error { + tasks := model.GetAllUnFinishTasks() + if len(tasks) == 0 { + return nil + } + logger.LogInfo(ctx, fmt.Sprintf("检测到未完成的任务数有: %v", len(tasks))) + taskChannelM := make(map[int][]string) + taskM := make(map[int]map[string]*model.Midjourney) + for _, task := range tasks { + if task.MjId == "" { + if err := failMidjourneyTask(ctx, task, "上游任务 ID 为空，无法继续轮询"); err != nil { + logger.LogError(ctx, fmt.Sprintf("Fix null mj_id task error: %v", err)) + } continue } - - logger.LogInfo(ctx, fmt.Sprintf("检测到未完成的任务数有: %v", len(tasks))) - taskChannelM := make(map[int][]string) - taskM := make(map[string]*model.Midjourney) - nullTaskIds := make([]int, 0) - for _, task := range tasks { - if task.MjId == "" { - // 统计失败的未完成任务 - nullTaskIds = append(nullTaskIds, task.Id) - continue - } - taskM[task.MjId] = task - taskChannelM[task.ChannelId] = append(taskChannelM[task.ChannelId], task.MjId) + if taskM[task.ChannelId] == nil { + taskM[task.ChannelId] = make(map[string]*model.Midjourney) } - if len(nullTaskIds) > 0 { - err := model.MjBulkUpdateByTaskIds(nullTaskIds, map[string]any{ - "status": "FAILURE", - "progress": "100%", - }) - if err != nil { - logger.LogError(ctx, fmt.Sprintf("Fix null mj_id task error: %v", err)) - } else { - logger.LogInfo(ctx, fmt.Sprintf("Fix null mj_id task success: %v", nullTaskIds)) - } + taskM[task.ChannelId][task.MjId] = task + taskChannelM[task.ChannelId] = append(taskChannelM[task.ChannelId], task.MjId) + } + + for channelID, taskIDs := range taskChannelM { + logger.LogInfo(ctx, fmt.Sprintf("渠道 #%d 未完成的任务有: %d", channelID, len(taskIDs))) + channel, err := model.CacheGetChannel(channelID) + if err != nil { + logger.LogError(ctx, fmt.Sprintf("CacheGetChannel: %v", err)) + continue } - if len(taskChannelM) == 0 { + responseItems, err := fetchMidjourneyTaskUpdates(ctx, channel, taskIDs) + if err != nil { + logger.LogError(ctx, fmt.Sprintf("Get Midjourney tasks failed for channel #%d: %v", channelID, err)) continue } - - for channelId, taskIds := range taskChannelM { - logger.LogInfo(ctx, fmt.Sprintf("渠道 #%d 未完成的任务有: %d", channelId, len(taskIds))) - if len(taskIds) == 0 { - continue - } - midjourneyChannel, err := model.CacheGetChannel(channelId) - if err != nil { - logger.LogError(ctx, fmt.Sprintf("CacheGetChannel: %v", err)) - err := model.MjBulkUpdate(taskIds, map[string]any{ - "fail_reason": fmt.Sprintf("获取渠道信息失败，请联系管理员，渠道ID：%d", channelId), - "status": "FAILURE", - "progress": "100%", - }) - if err != nil { - logger.LogInfo(ctx, fmt.Sprintf("UpdateMidjourneyTask error: %v", err)) - } - continue - } - requestUrl := fmt.Sprintf("%s/mj/task/list-by-condition", *midjourneyChannel.BaseURL) - - body, _ := json.Marshal(map[string]any{ - "ids": taskIds, - }) - req, err := http.NewRequest("POST", requestUrl, bytes.NewBuffer(body)) - if err != nil { - logger.LogError(ctx, fmt.Sprintf("Get Task error: %v", err)) - continue - } - // 设置超时时间 - timeout := time.Second * 15 - ctx, cancel := context.WithTimeout(context.Background(), timeout) - // 使用带有超时的 context 创建新的请求 - req = req.WithContext(ctx) - req.Header.Set("Content-Type", "application/json") - req.Header.Set("mj-api-secret", midjourneyChannel.Key) - resp, err := service.GetHttpClient().Do(req) - if err != nil { - logger.LogError(ctx, fmt.Sprintf("Get Task Do req error: %v", err)) - continue - } - if resp.StatusCode != http.StatusOK { - logger.LogError(ctx, fmt.Sprintf("Get Task status code: %d", resp.StatusCode)) - continue - } - responseBody, err := io.ReadAll(resp.Body) - if err != nil { - logger.LogError(ctx, fmt.Sprintf("Get Mjp Task parse body error: %v", err)) - continue - } - var responseItems []dto.MidjourneyDto - err = json.Unmarshal(responseBody, &responseItems) - if err != nil { - logger.LogError(ctx, fmt.Sprintf("Get Mjp Task parse body error2: %v, body: %s", err, string(responseBody))) - continue - } - resp.Body.Close() - req.Body.Close() - cancel() - - for _, responseItem := range responseItems { - task := taskM[responseItem.MjId] - - useTime := (time.Now().UnixNano() / int64(time.Millisecond)) - task.SubmitTime - // 如果时间超过一小时，且进度不是100%，则认为任务失败 - if useTime > 3600000 && task.Progress != "100%" { - responseItem.FailReason = "上游任务超时（超过1小时）" - responseItem.Status = "FAILURE" - } - if !checkMjTaskNeedUpdate(task, responseItem) { - continue - } - preStatus := task.Status - task.Code = 1 - task.Progress = responseItem.Progress - task.PromptEn = responseItem.PromptEn - task.State = responseItem.State - task.SubmitTime = responseItem.SubmitTime - task.StartTime = responseItem.StartTime - task.FinishTime = responseItem.FinishTime - task.ImageUrl = responseItem.ImageUrl - task.Status = responseItem.Status - task.FailReason = responseItem.FailReason - if responseItem.Properties != nil { - propertiesStr, _ := json.Marshal(responseItem.Properties) - task.Properties = string(propertiesStr) - } - if responseItem.Buttons != nil { - buttonStr, _ := json.Marshal(responseItem.Buttons) - task.Buttons = string(buttonStr) - } - // 映射 VideoUrl - task.VideoUrl = responseItem.VideoUrl - - // 映射 VideoUrls - 将数组序列化为 JSON 字符串 - if responseItem.VideoUrls != nil && len(responseItem.VideoUrls) > 0 { - videoUrlsStr, err := json.Marshal(responseItem.VideoUrls) - if err != nil { - logger.LogError(ctx, fmt.Sprintf("序列化 VideoUrls 失败: %v", err)) - task.VideoUrls = "[]" // 失败时设置为空数组 - } else { - task.VideoUrls = string(videoUrlsStr) - } - } else { - task.VideoUrls = "" // 空值时清空字段 - } - - shouldReturnQuota := false - if (task.Progress != "100%" && responseItem.FailReason != "") || (task.Progress == "100%" && task.Status == "FAILURE") { - logger.LogInfo(ctx, task.MjId+" 构建失败，"+task.FailReason) - task.Progress = "100%" - if task.Quota != 0 { - shouldReturnQuota = true - } - } - won, err := task.UpdateWithStatus(preStatus) - if err != nil { - logger.LogError(ctx, "UpdateMidjourneyTask task error: "+err.Error()) - } else if won && shouldReturnQuota { - err = model.IncreaseUserQuota(task.UserId, task.Quota, false) - if err != nil { - logger.LogError(ctx, "fail to increase user quota: "+err.Error()) - } - model.RecordTaskBillingLog(model.RecordTaskBillingLogParams{ - UserId: task.UserId, - LogType: model.LogTypeRefund, - Content: "", - ChannelId: task.ChannelId, - ModelName: service.CovertMjpActionToModelName(task.Action), - Quota: task.Quota, - Other: map[string]interface{}{ - "task_id": task.MjId, - "reason": "构图失败", - }, - }) - } + for _, responseItem := range responseItems { + if err := applyMidjourneyTaskResponse(ctx, taskM[channelID], responseItem); err != nil { + logger.LogError(ctx, "UpdateMidjourneyTask task error: "+err.Error()) } } } + return nil } -func checkMjTaskNeedUpdate(oldTask *model.Midjourney, newTask dto.MidjourneyDto) bool { - if oldTask.Code != 1 { - return true - } - if oldTask.Progress != newTask.Progress { - return true - } - if oldTask.PromptEn != newTask.PromptEn { - return true - } - if oldTask.State != newTask.State { - return true - } - if oldTask.SubmitTime != newTask.SubmitTime { - return true - } - if oldTask.StartTime != newTask.StartTime { - return true - } - if oldTask.FinishTime != newTask.FinishTime { - return true - } - if oldTask.ImageUrl != newTask.ImageUrl { - return true - } - if oldTask.Status != newTask.Status { - return true - } - if oldTask.FailReason != newTask.FailReason { - return true - } - if oldTask.FinishTime != newTask.FinishTime { - return true - } - if oldTask.Progress != "100%" && newTask.FailReason != "" { - return true - } - // 检查 VideoUrl 是否需要更新 - if oldTask.VideoUrl != newTask.VideoUrl { - return true +func fetchMidjourneyTaskUpdates(parent context.Context, channel *model.Channel, taskIDs []string) ([]dto.MidjourneyDto, error) { + key, _, maxAPIError := channel.GetNextEnabledKey() + if maxAPIError != nil { + return nil, maxAPIError + } + body, err := common.Marshal(map[string]any{"ids": taskIDs}) + if err != nil { + return nil, err + } + requestURL := fmt.Sprintf("%s/mj/task/list-by-condition", channel.GetBaseURL()) + requestCtx, cancel := context.WithTimeout(parent, 15*time.Second) + defer cancel() + req, err := http.NewRequestWithContext(requestCtx, http.MethodPost, requestURL, bytes.NewReader(body)) + if err != nil { + return nil, err + } + defer req.Body.Close() + req.Header.Set("Content-Type", "application/json") + req.Header.Set("mj-api-secret", key) + resp, err := service.GetHttpClient().Do(req) + if err != nil { + return nil, err + } + defer resp.Body.Close() + if resp.StatusCode != http.StatusOK { + _, _ = io.Copy(io.Discard, resp.Body) + return nil, fmt.Errorf("unexpected status code: %d", resp.StatusCode) + } + responseBody, err := io.ReadAll(resp.Body) + if err != nil { + return nil, err + } + var responseItems []dto.MidjourneyDto + if err := common.Unmarshal(responseBody, &responseItems); err != nil { + return nil, fmt.Errorf("parse Midjourney task response: %w", err) + } + return responseItems, nil +} + +func applyMidjourneyTaskResponse(ctx context.Context, taskM map[string]*model.Midjourney, responseItem dto.MidjourneyDto) error { + task := taskM[responseItem.MjId] + if task == nil { + logger.LogWarn(ctx, fmt.Sprintf("忽略未知 Midjourney 任务 ID: %s", responseItem.MjId)) + return nil } - // 检查 VideoUrls 是否需要更新 - if newTask.VideoUrls != nil && len(newTask.VideoUrls) > 0 { - newVideoUrlsStr, _ := json.Marshal(newTask.VideoUrls) - if oldTask.VideoUrls != string(newVideoUrlsStr) { - return true - } - } else if oldTask.VideoUrls != "" { - // 如果新数据没有 VideoUrls 但旧数据有，需要更新（清空） - return true + useTime := time.Now().UnixMilli() - task.SubmitTime + if useTime > 3600000 && task.Progress != "100%" && responseItem.Progress != "100%" && responseItem.Status != "SUCCESS" { + responseItem.FailReason = "上游任务超时（超过1小时）" + responseItem.Status = "FAILURE" } + return updateMidjourneyTaskFromResponse(ctx, task, responseItem) +} + +func failMidjourneyTask(ctx context.Context, task *model.Midjourney, reason string) error { + return updateMidjourneyTaskFromResponse(ctx, task, dto.MidjourneyDto{ + MjId: task.MjId, Status: "FAILURE", Progress: "100%", FailReason: reason, + PromptEn: task.PromptEn, State: task.State, SubmitTime: task.SubmitTime, + StartTime: task.StartTime, FinishTime: time.Now().UnixMilli(), + ImageUrl: task.ImageUrl, VideoUrl: task.VideoUrl, + }) +} - return false +func updateMidjourneyTaskFromResponse(ctx context.Context, task *model.Midjourney, responseItem dto.MidjourneyDto) error { + return service.UpdateMidjourneyTaskFromResponse(ctx, task, responseItem) } func GetAllMidjourney(c *gin.Context) { diff --git a/controller/midjourney_test.go b/controller/midjourney_test.go new file mode 100644 index 00000000..d776dfbf --- /dev/null +++ b/controller/midjourney_test.go @@ -0,0 +1,292 @@ +package controller + +import ( + "context" + "fmt" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/MAX-API-Next/MAX-API/common" + "github.com/MAX-API-Next/MAX-API/constant" + "github.com/MAX-API-Next/MAX-API/dto" + "github.com/MAX-API-Next/MAX-API/model" + "github.com/MAX-API-Next/MAX-API/service" + "github.com/glebarez/sqlite" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "gorm.io/gorm" +) + +func TestApplyMidjourneyTaskResponseIgnoresUnknownUpstreamID(t *testing.T) { + require.NotPanics(t, func() { + err := applyMidjourneyTaskResponse(context.Background(), map[string]*model.Midjourney{}, dto.MidjourneyDto{ + MjId: "provider-task-not-in-local-batch", + }) + require.NoError(t, err) + }) +} + +func TestFetchMidjourneyTaskUpdatesUsesEnabledChannelKey(t *testing.T) { + service.InitHttpClient() + receivedKey := make(chan string, 1) + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + receivedKey <- r.Header.Get("mj-api-secret") + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`[]`)) + })) + t.Cleanup(upstream.Close) + + oldMemoryCache := common.MemoryCacheEnabled + common.MemoryCacheEnabled = false + t.Cleanup(func() { common.MemoryCacheEnabled = oldMemoryCache }) + baseURL := upstream.URL + channel := &model.Channel{ + Id: 701, Type: constant.ChannelTypeMidjourney, Key: "disabled-key\nenabled-key", + Status: common.ChannelStatusEnabled, BaseURL: &baseURL, + ChannelInfo: model.ChannelInfo{ + IsMultiKey: true, MultiKeyMode: constant.MultiKeyModeRandom, + MultiKeyStatusList: map[int]int{0: common.ChannelStatusManuallyDisabled}, + }, + } + + _, err := fetchMidjourneyTaskUpdates(context.Background(), channel, []string{"provider-task"}) + require.NoError(t, err) + select { + case key := <-receivedKey: + require.Equal(t, "enabled-key", key) + case <-time.After(time.Second): + t.Fatal("upstream request was not received") + } +} + +func TestUpdateMidjourneyTasksOnceScopesDuplicateProviderIDsByChannel(t *testing.T) { + service.InitHttpClient() + serverFor := func(progress, imageURL string) *httptest.Server { + return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`[{"id":"shared-provider-id","status":"IN_PROGRESS","progress":"` + progress + `","imageUrl":"` + imageURL + `"}]`)) + })) + } + upstreamA := serverFor("10%", "https://example.com/a.png") + upstreamB := serverFor("20%", "https://example.com/b.png") + t.Cleanup(upstreamA.Close) + t.Cleanup(upstreamB.Close) + + db, err := gorm.Open(sqlite.Open("file:midjourney_channel_scope?mode=memory&cache=shared"), &gorm.Config{}) + require.NoError(t, err) + require.NoError(t, db.AutoMigrate(&model.Channel{}, &model.Midjourney{}, &model.MidjourneyBillingClaim{})) + oldDB := model.DB + oldSQLite := common.UsingSQLite + oldMemoryCache := common.MemoryCacheEnabled + model.DB = db + common.UsingSQLite = true + common.MemoryCacheEnabled = false + t.Cleanup(func() { + model.DB = oldDB + common.UsingSQLite = oldSQLite + common.MemoryCacheEnabled = oldMemoryCache + if sqlDB, dbErr := db.DB(); dbErr == nil { + _ = sqlDB.Close() + } + }) + + baseURLA, baseURLB := upstreamA.URL, upstreamB.URL + require.NoError(t, db.Create(&model.Channel{ + Id: 711, Type: constant.ChannelTypeMidjourney, Key: "key-a", Status: common.ChannelStatusEnabled, BaseURL: &baseURLA, + }).Error) + require.NoError(t, db.Create(&model.Channel{ + Id: 712, Type: constant.ChannelTypeMidjourney, Key: "key-b", Status: common.ChannelStatusEnabled, BaseURL: &baseURLB, + }).Error) + submitTime := time.Now().UnixMilli() + taskA := model.Midjourney{UserId: 721, ChannelId: 711, MjId: "shared-provider-id", Progress: "0%", Status: "", SubmitTime: submitTime} + taskB := model.Midjourney{UserId: 722, ChannelId: 712, MjId: "shared-provider-id", Progress: "0%", Status: "", SubmitTime: submitTime} + require.NoError(t, db.Create(&taskA).Error) + require.NoError(t, db.Create(&taskB).Error) + + require.NoError(t, updateMidjourneyTasksOnce(context.Background())) + + var storedA, storedB model.Midjourney + require.NoError(t, db.First(&storedA, taskA.Id).Error) + require.NoError(t, db.First(&storedB, taskB.Id).Error) + assert.Equal(t, "10%", storedA.Progress) + assert.Equal(t, "https://example.com/a.png", storedA.ImageUrl) + assert.Equal(t, "20%", storedB.Progress) + assert.Equal(t, "https://example.com/b.png", storedB.ImageUrl) +} + +func setupMidjourneyBillingControllerTestDB(t *testing.T) *gorm.DB { + t.Helper() + dsn := fmt.Sprintf("file:%s?mode=memory&cache=shared", strings.ReplaceAll(t.Name(), "/", "_")) + db, err := gorm.Open(sqlite.Open(dsn), &gorm.Config{}) + require.NoError(t, err) + sqlDB, err := db.DB() + require.NoError(t, err) + + oldDB, oldLogDB := model.DB, model.LOG_DB + oldSQLite, oldRedis := common.UsingSQLite, common.RedisEnabled + model.DB, model.LOG_DB = db, db + common.UsingSQLite, common.RedisEnabled = true, false + t.Cleanup(func() { + model.DB, model.LOG_DB = oldDB, oldLogDB + common.UsingSQLite, common.RedisEnabled = oldSQLite, oldRedis + _ = sqlDB.Close() + }) + + require.NoError(t, db.AutoMigrate( + &model.User{}, &model.Token{}, &model.Task{}, &model.Midjourney{}, + &model.MidjourneyBillingClaim{}, &model.BillingSettlement{}, + &model.Log{}, &model.BillingLogReceipt{}, + )) + return db +} + +func TestMidjourneyFailureRefundsWalletAndTokenThroughDurableSettlement(t *testing.T) { + db := setupMidjourneyBillingControllerTestDB(t) + + user := model.User{Id: 801, Username: "midjourney-refund-user", AffCode: "mj-refund-user", Quota: 70, Status: common.UserStatusEnabled} + token := model.Token{ + Id: 802, UserId: user.Id, Key: "midjourney-refund-token", Name: "midjourney-refund-token", + Status: common.TokenStatusEnabled, RemainQuota: 70, UsedQuota: 30, + } + require.NoError(t, db.Create(&user).Error) + require.NoError(t, db.Create(&token).Error) + billingTask := &model.Task{ + TaskID: "task_midjourney_refund", Platform: constant.TaskPlatformMidjourney, + UserId: user.Id, ChannelId: 803, Quota: 30, Action: constant.MjActionImagine, + Status: model.TaskStatusNotStart, Progress: "0%", + PrivateData: model.TaskPrivateData{ + AwaitingUpstreamID: true, BillingSource: model.BillingSettlementSourceWallet, + BillingRequestId: "midjourney-refund-request", TokenId: token.Id, + }, + } + require.NoError(t, db.Create(billingTask).Error) + mj := &model.Midjourney{ + UserId: user.Id, ChannelId: 803, MjId: "provider-midjourney-refund", + Action: constant.MjActionImagine, Status: "", Progress: "0%", Quota: 30, + SubmitTime: 123000, ImageUrl: "https://example.com/original.png", + } + created, refundDuplicate, err := model.FinalizeMidjourneySubmission(mj, billingTask, &model.BillingSettlementInput{ + OperationKey: "request:midjourney-refund-request:finalize", Source: model.BillingSettlementSourceWallet, + UserID: user.Id, TokenID: token.Id, + }, nil) + require.NoError(t, err) + require.True(t, created) + require.False(t, refundDuplicate) + + err = updateMidjourneyTaskFromResponse(context.Background(), mj, dto.MidjourneyDto{ + MjId: mj.MjId, Status: "FAILURE", Progress: "100%", FailReason: "provider failed", + SubmitTime: mj.SubmitTime, FinishTime: 1000, + }) + require.NoError(t, err) + + var storedUser model.User + require.NoError(t, db.First(&storedUser, user.Id).Error) + assert.EqualValues(t, 100, storedUser.Quota) + var storedToken model.Token + require.NoError(t, db.First(&storedToken, token.Id).Error) + assert.EqualValues(t, 100, storedToken.RemainQuota) + assert.Zero(t, storedToken.UsedQuota) + var storedTask model.Task + require.NoError(t, db.First(&storedTask, billingTask.ID).Error) + assert.EqualValues(t, model.TaskStatusFailure, storedTask.Status) + assert.Zero(t, storedTask.Quota) + var storedMJ model.Midjourney + require.NoError(t, db.First(&storedMJ, mj.Id).Error) + assert.Equal(t, "FAILURE", storedMJ.Status) + assert.EqualValues(t, 123000, storedMJ.SubmitTime) + assert.Equal(t, "https://example.com/original.png", storedMJ.ImageUrl) +} + +func TestApplyMidjourneyTaskResponsePreservesLateCompletionWithoutRefund(t *testing.T) { + tests := []struct { + name string + responseStatus string + storedStatus string + storedProgress string + alreadyCompleted bool + }{ + {name: "explicit success clears stale failure", responseStatus: "SUCCESS", storedProgress: "50%"}, + {name: "completed progress without status is successful", storedProgress: "50%"}, + { + name: "repeated success clears the only stale field", responseStatus: "SUCCESS", + storedStatus: "SUCCESS", storedProgress: "100%", alreadyCompleted: true, + }, + } + + for i, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + db := setupMidjourneyBillingControllerTestDB(t) + suffix := fmt.Sprintf("%d", i) + finishTime := time.Now().UnixMilli() + imageURL := "https://example.com/late-success.png" + storedFinishTime := int64(0) + storedImageURL := "" + if tt.alreadyCompleted { + storedFinishTime = finishTime + storedImageURL = imageURL + } + user := model.User{Id: 811, Username: "midjourney-late-success-user", AffCode: "mj-late-success-user", Quota: 70, Status: common.UserStatusEnabled} + token := model.Token{ + Id: 812, UserId: user.Id, Key: "midjourney-late-success-token", Name: "midjourney-late-success-token", + Status: common.TokenStatusEnabled, RemainQuota: 70, UsedQuota: 30, + } + require.NoError(t, db.Create(&user).Error) + require.NoError(t, db.Create(&token).Error) + billingTask := &model.Task{ + TaskID: "task_midjourney_late_success_" + suffix, Platform: constant.TaskPlatformMidjourney, + UserId: user.Id, ChannelId: 813, Quota: 30, Action: constant.MjActionImagine, + Status: model.TaskStatusNotStart, Progress: "0%", + PrivateData: model.TaskPrivateData{ + AwaitingUpstreamID: true, BillingSource: model.BillingSettlementSourceWallet, + BillingRequestId: "midjourney-late-success-request-" + suffix, TokenId: token.Id, + }, + } + require.NoError(t, db.Create(billingTask).Error) + mj := &model.Midjourney{ + UserId: user.Id, ChannelId: billingTask.ChannelId, MjId: "provider-midjourney-late-success-" + suffix, + Action: constant.MjActionImagine, Code: 1, Status: tt.storedStatus, Progress: tt.storedProgress, + FailReason: "stale provider failure", Quota: 30, SubmitTime: time.Now().Add(-2 * time.Hour).UnixMilli(), + FinishTime: storedFinishTime, ImageUrl: storedImageURL, + } + created, refundDuplicate, err := model.FinalizeMidjourneySubmission(mj, billingTask, &model.BillingSettlementInput{ + OperationKey: "request:midjourney-late-success-request-" + suffix + ":finalize", Source: model.BillingSettlementSourceWallet, + UserID: user.Id, TokenID: token.Id, + }, nil) + require.NoError(t, err) + require.True(t, created) + require.False(t, refundDuplicate) + var settlementCountBefore int64 + require.NoError(t, db.Model(&model.BillingSettlement{}).Count(&settlementCountBefore).Error) + + err = applyMidjourneyTaskResponse(context.Background(), map[string]*model.Midjourney{mj.MjId: mj}, dto.MidjourneyDto{ + MjId: mj.MjId, Status: tt.responseStatus, Progress: "100%", FinishTime: finishTime, + ImageUrl: imageURL, + }) + require.NoError(t, err) + + var storedMJ model.Midjourney + require.NoError(t, db.First(&storedMJ, mj.Id).Error) + assert.Equal(t, "SUCCESS", storedMJ.Status) + assert.Empty(t, storedMJ.FailReason) + assert.Equal(t, imageURL, storedMJ.ImageUrl) + assert.Equal(t, finishTime, storedMJ.FinishTime) + var storedTask model.Task + require.NoError(t, db.First(&storedTask, billingTask.ID).Error) + assert.EqualValues(t, model.TaskStatusSuccess, storedTask.Status) + assert.Equal(t, 30, storedTask.Quota) + var storedUser model.User + require.NoError(t, db.First(&storedUser, user.Id).Error) + assert.EqualValues(t, 70, storedUser.Quota) + var storedToken model.Token + require.NoError(t, db.First(&storedToken, token.Id).Error) + assert.EqualValues(t, 70, storedToken.RemainQuota) + assert.EqualValues(t, 30, storedToken.UsedQuota) + var settlementCountAfter int64 + require.NoError(t, db.Model(&model.BillingSettlement{}).Count(&settlementCountAfter).Error) + assert.Equal(t, settlementCountBefore, settlementCountAfter) + }) + } +} diff --git a/controller/model_sync.go b/controller/model_sync.go index 1a8ac625..edd80ccf 100644 --- a/controller/model_sync.go +++ b/controller/model_sync.go @@ -6,7 +6,6 @@ import ( "errors" "fmt" "io" - "math/rand" "net" "net/http" "strings" @@ -26,6 +25,14 @@ const ( upstreamVendorsURL = "https://basellm.github.io/llm-metadata/api/newapi/vendors.json" ) +func modelSyncJitter(maxMillis int) time.Duration { + jitter, err := common.SecureRandomInt(maxMillis) + if err != nil { + return 0 + } + return time.Duration(jitter) * time.Millisecond +} + func normalizeLocale(locale string) (string, bool) { l := strings.ToLower(strings.TrimSpace(locale)) switch l { @@ -156,7 +163,7 @@ func fetchJSON[T any](ctx context.Context, url string, out *upstreamEnvelope[T]) lastErr = err // backoff with jitter sleep := baseDelay * time.Duration(1< 32 || + (request.Intent == model.AuthFlowIntentBind && request.Aff != "") { + common.ApiErrorI18n(c, i18n.MsgInvalidParams) + return + } + userID := 0 + if request.Intent == model.AuthFlowIntentBind { + userID = c.GetInt("id") + if userID <= 0 { + c.JSON(http.StatusUnauthorized, gin.H{"success": false, "message": i18n.T(c, i18n.MsgUnauthorized)}) + return + } + } + payload, err := common.Marshal(oauthFlowPayload{AffiliateCode: request.Aff}) if err != nil { common.ApiError(c, err) return } + expiresAt := time.Now().Add(oauthAuthFlowTTL) + state, _, err := model.CreateAuthFlow(model.AuthFlowCreate{ + Purpose: model.AuthFlowPurposeOAuth, + Provider: request.Provider, + Intent: request.Intent, + UserId: userID, + Payload: string(payload), + ExpiresAt: expiresAt, + }) + if err != nil { + common.ApiError(c, err) + return + } + session := sessions.Default(c) + session.Set(oauthStateSessionKey, state) + if err := session.Save(); err != nil { + if _, consumeErr := model.ConsumeAuthFlow(state, model.AuthFlowMatch{ + Purpose: model.AuthFlowPurposeOAuth, + Provider: request.Provider, + Intent: request.Intent, + UserId: userID, + }); consumeErr != nil { + common.SysError("failed to invalidate OAuth flow after session save failure: " + consumeErr.Error()) + } + common.ApiError(c, err) + return + } c.JSON(http.StatusOK, gin.H{ "success": true, "message": "", @@ -109,11 +186,20 @@ func HandleOAuth(c *gin.Context) { return } - session := sessions.Default(c) - // 1. Validate state (CSRF protection) state := c.Query("state") - if state == "" || session.Get("oauth_state") == nil || state != session.Get("oauth_state").(string) { + if !oauthSessionStateMatches(c, state) { + c.JSON(http.StatusForbidden, gin.H{ + "success": false, + "message": i18n.T(c, i18n.MsgOAuthStateInvalid), + }) + return + } + pendingFlow, err := model.GetAuthFlow(state, model.AuthFlowMatch{ + Purpose: model.AuthFlowPurposeOAuth, + Provider: providerName, + }) + if err != nil { c.JSON(http.StatusForbidden, gin.H{ "success": false, "message": i18n.T(c, i18n.MsgOAuthStateInvalid), @@ -121,10 +207,20 @@ func HandleOAuth(c *gin.Context) { return } - // 2. Check if user is already logged in (bind flow) - username := session.Get("username") - if username != nil { - handleOAuthBind(c, provider) + consumeMatch := model.AuthFlowMatch{ + Purpose: model.AuthFlowPurposeOAuth, + Provider: providerName, + Intent: pendingFlow.Intent, + } + if pendingFlow.Intent == model.AuthFlowIntentBind { + userID := c.GetInt("id") + if userID <= 0 || userID != pendingFlow.UserId { + c.JSON(http.StatusForbidden, gin.H{"success": false, "message": i18n.T(c, i18n.MsgOAuthStateInvalid)}) + return + } + consumeMatch.UserId = userID + } else if pendingFlow.Intent != model.AuthFlowIntentLogin { + common.ApiErrorI18n(c, i18n.MsgInvalidParams) return } @@ -137,13 +233,25 @@ func HandleOAuth(c *gin.Context) { // 4. Handle error from provider errorCode := c.Query("error") if errorCode != "" { + if _, err := model.ConsumeAuthFlow(state, consumeMatch); err != nil { + c.JSON(http.StatusForbidden, gin.H{"success": false, "message": i18n.T(c, i18n.MsgOAuthStateInvalid)}) + return + } + clearOAuthSessionState(c) errorDescription := c.Query("error_description") + if errorDescription == "" { + errorDescription = errorCode + } c.JSON(http.StatusOK, gin.H{ "success": false, "message": errorDescription, }) return } + if pendingFlow.Intent == model.AuthFlowIntentBind { + handleOAuthBind(c, provider, pendingFlow, state) + return + } // 5. Exchange code for token code := c.Query("code") @@ -159,9 +267,20 @@ func HandleOAuth(c *gin.Context) { handleOAuthError(c, err) return } + flow, err := model.ConsumeAuthFlow(state, consumeMatch) + if err != nil { + c.JSON(http.StatusForbidden, gin.H{"success": false, "message": i18n.T(c, i18n.MsgOAuthStateInvalid)}) + return + } + clearOAuthSessionState(c) // 7. Find or create user - user, err := findOrCreateOAuthUser(c, provider, oauthUser, session) + var payload oauthFlowPayload + if err := common.UnmarshalJsonStr(flow.Payload, &payload); err != nil { + common.ApiError(c, err) + return + } + user, err := findOrCreateOAuthUser(c, provider, oauthUser, payload.AffiliateCode) if err != nil { var lookupErr *oauthIdentityLookupError if errors.As(err, &lookupErr) { @@ -192,12 +311,7 @@ func HandleOAuth(c *gin.Context) { } // handleOAuthBind handles binding OAuth account to existing user -func handleOAuthBind(c *gin.Context, provider oauth.Provider) { - if !provider.IsEnabled() { - common.ApiErrorI18n(c, i18n.MsgOAuthNotEnabled, providerParams(provider.GetName())) - return - } - +func handleOAuthBind(c *gin.Context, provider oauth.Provider, pendingFlow *model.AuthFlow, flowToken string) { // Exchange code for token code := c.Query("code") token, err := provider.ExchangeToken(c.Request.Context(), code, c) @@ -234,14 +348,17 @@ func handleOAuthBind(c *gin.Context, provider oauth.Provider) { } } - // Get current user from session - session := sessions.Default(c) - id, ok := sessionUserID(session.Get("id")) - if !ok { - common.ApiErrorMsg(c, "用户未登录或登录状态已失效") + if _, err := model.ConsumeAuthFlow(flowToken, model.AuthFlowMatch{ + Purpose: model.AuthFlowPurposeOAuth, + Provider: pendingFlow.Provider, + Intent: model.AuthFlowIntentBind, + UserId: pendingFlow.UserId, + }); err != nil { + c.JSON(http.StatusForbidden, gin.H{"success": false, "message": i18n.T(c, i18n.MsgOAuthStateInvalid)}) return } - user := model.User{Id: id} + clearOAuthSessionState(c) + user := model.User{Id: pendingFlow.UserId} err = user.FillUserById() if err != nil { common.ApiError(c, err) @@ -277,7 +394,7 @@ func handleOAuthBind(c *gin.Context, provider oauth.Provider) { } // findOrCreateOAuthUser finds existing user or creates new user -func findOrCreateOAuthUser(c *gin.Context, provider oauth.Provider, oauthUser *oauth.OAuthUser, session sessions.Session) (*model.User, error) { +func findOrCreateOAuthUser(c *gin.Context, provider oauth.Provider, oauthUser *oauth.OAuthUser, affiliateCode string) (*model.User, error) { user := &model.User{} // Check if user already exists with new ID @@ -364,10 +481,9 @@ func findOrCreateOAuthUser(c *gin.Context, provider oauth.Provider, oauthUser *o user.Status = common.UserStatusEnabled // Handle affiliate code - affCode := session.Get("aff") inviterId := 0 - if affCode != nil { - inviterId, _ = model.GetUserIdByAffCode(affCode.(string)) + if affiliateCode != "" { + inviterId, _ = model.GetUserIdByAffCode(affiliateCode) } // Use transaction to ensure user creation and OAuth binding are atomic diff --git a/controller/oauth_test.go b/controller/oauth_test.go index 26fd1502..182a6af5 100644 --- a/controller/oauth_test.go +++ b/controller/oauth_test.go @@ -9,10 +9,15 @@ import ( "testing" "github.com/MAX-API-Next/MAX-API/common" + appi18n "github.com/MAX-API-Next/MAX-API/i18n" "github.com/MAX-API-Next/MAX-API/model" "github.com/MAX-API-Next/MAX-API/oauth" + "github.com/gin-contrib/sessions" + "github.com/gin-contrib/sessions/cookie" "github.com/gin-gonic/gin" + "github.com/glebarez/sqlite" "github.com/stretchr/testify/require" + "gorm.io/gorm" ) type deletedUserOAuthProvider struct{} @@ -22,6 +27,148 @@ func (*deletedUserOAuthProvider) IsEnabled() bool { return true } func (*deletedUserOAuthProvider) ExchangeToken(context.Context, string, *gin.Context) (*oauth.OAuthToken, error) { return nil, nil } + +func setupControllerAuthFlowDB(t *testing.T) { + t.Helper() + originalDB := model.DB + db, err := gorm.Open(sqlite.Open("file:controller_auth_flow?mode=memory&cache=shared"), &gorm.Config{}) + require.NoError(t, err) + require.NoError(t, db.AutoMigrate(&model.AuthFlow{})) + model.DB = db + t.Cleanup(func() { + model.DB = originalDB + if sqlDB, dbErr := db.DB(); dbErr == nil { + _ = sqlDB.Close() + } + }) +} + +func TestGenerateOAuthCodeCreatesProviderBoundOpaqueState(t *testing.T) { + setupControllerAuthFlowDB(t) + gin.SetMode(gin.TestMode) + router := gin.New() + router.Use(sessions.Sessions("session", cookie.NewStore([]byte("oauth-state-create-test")))) + router.POST("/api/oauth/state", GenerateOAuthCode) + + recorder := httptest.NewRecorder() + request := httptest.NewRequest(http.MethodPost, "/api/oauth/state", strings.NewReader(`{"provider":"github","intent":"login","aff":"invite"}`)) + request.Header.Set("Content-Type", "application/json") + router.ServeHTTP(recorder, request) + + var response struct { + Success bool `json:"success"` + Data string `json:"data"` + } + require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response)) + require.True(t, response.Success) + require.Len(t, response.Data, 43) + flow, err := model.GetAuthFlow(response.Data, model.AuthFlowMatch{ + Purpose: model.AuthFlowPurposeOAuth, Provider: "github", Intent: model.AuthFlowIntentLogin, + }) + require.NoError(t, err) + require.NotEqual(t, response.Data, flow.TokenHash) + require.Contains(t, flow.Payload, "invite") +} + +func TestHandleOAuthRejectsStateFromDifferentBrowserSession(t *testing.T) { + setupControllerAuthFlowDB(t) + require.NoError(t, appi18n.Init()) + originalGitHubOAuthEnabled := common.GitHubOAuthEnabled + common.GitHubOAuthEnabled = true + t.Cleanup(func() { common.GitHubOAuthEnabled = originalGitHubOAuthEnabled }) + gin.SetMode(gin.TestMode) + + router := gin.New() + router.Use(sessions.Sessions("session", cookie.NewStore([]byte("oauth-browser-binding-test")))) + router.POST("/api/oauth/state", GenerateOAuthCode) + router.GET("/api/oauth/:provider", HandleOAuth) + + stateRecorder := httptest.NewRecorder() + stateRequest := httptest.NewRequest(http.MethodPost, "/api/oauth/state", strings.NewReader(`{"provider":"github","intent":"login"}`)) + stateRequest.Header.Set("Content-Type", "application/json") + router.ServeHTTP(stateRecorder, stateRequest) + require.Equal(t, http.StatusOK, stateRecorder.Code) + + var stateResponse struct { + Success bool `json:"success"` + Data string `json:"data"` + } + require.NoError(t, common.Unmarshal(stateRecorder.Body.Bytes(), &stateResponse)) + require.True(t, stateResponse.Success) + require.NotEmpty(t, stateResponse.Data) + require.NotEmpty(t, stateRecorder.Result().Cookies()) + + callbackRecorder := httptest.NewRecorder() + callbackRequest := httptest.NewRequest(http.MethodGet, "/api/oauth/github?state="+stateResponse.Data+"&error=access_denied", nil) + // Deliberately omit browser A's signed session cookie. + router.ServeHTTP(callbackRecorder, callbackRequest) + + require.Equal(t, http.StatusForbidden, callbackRecorder.Code) + + callbackRecorder = httptest.NewRecorder() + callbackRequest = httptest.NewRequest(http.MethodGet, "/api/oauth/github?state="+stateResponse.Data+"&error=access_denied", nil) + for _, sessionCookie := range stateRecorder.Result().Cookies() { + callbackRequest.AddCookie(sessionCookie) + } + router.ServeHTTP(callbackRecorder, callbackRequest) + require.Equal(t, http.StatusOK, callbackRecorder.Code) + _, err := model.GetAuthFlow(stateResponse.Data, model.AuthFlowMatch{Purpose: model.AuthFlowPurposeOAuth}) + require.ErrorIs(t, err, model.ErrAuthFlowConsumed) +} + +func TestGenerateOAuthCodeRejectsAnonymousBind(t *testing.T) { + require.NoError(t, appi18n.Init()) + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/api/oauth/state", strings.NewReader(`{"provider":"github","intent":"bind"}`)) + c.Request.Header.Set("Accept-Language", "en") + + GenerateOAuthCode(c) + + require.Equal(t, http.StatusUnauthorized, recorder.Code) + var response struct { + Success bool `json:"success"` + Message string `json:"message"` + } + require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response)) + require.False(t, response.Success) + require.Equal(t, "Unauthorized", response.Message) +} + +func TestHandleOAuthRejectsConsumedStateBeforeProviderExchange(t *testing.T) { + setupControllerAuthFlowDB(t) + require.NoError(t, appi18n.Init()) + gin.SetMode(gin.TestMode) + router := gin.New() + router.Use(sessions.Sessions("session", cookie.NewStore([]byte("oauth-consumed-state-test")))) + router.POST("/api/oauth/state", GenerateOAuthCode) + router.GET("/api/oauth/:provider", HandleOAuth) + + stateRecorder := httptest.NewRecorder() + stateRequest := httptest.NewRequest(http.MethodPost, "/api/oauth/state", strings.NewReader(`{"provider":"github","intent":"login"}`)) + stateRequest.Header.Set("Content-Type", "application/json") + router.ServeHTTP(stateRecorder, stateRequest) + require.Equal(t, http.StatusOK, stateRecorder.Code) + var stateResponse struct { + Data string `json:"data"` + } + require.NoError(t, common.Unmarshal(stateRecorder.Body.Bytes(), &stateResponse)) + state := stateResponse.Data + require.NotEmpty(t, state) + + var err error + _, err = model.ConsumeAuthFlow(state, model.AuthFlowMatch{Purpose: model.AuthFlowPurposeOAuth, Provider: "github"}) + require.NoError(t, err) + + recorder := httptest.NewRecorder() + callbackRequest := httptest.NewRequest(http.MethodGet, "/api/oauth/github?state="+state+"&code=unused", nil) + for _, sessionCookie := range stateRecorder.Result().Cookies() { + callbackRequest.AddCookie(sessionCookie) + } + router.ServeHTTP(recorder, callbackRequest) + + require.Equal(t, http.StatusForbidden, recorder.Code) +} func (*deletedUserOAuthProvider) GetUserInfo(context.Context, *oauth.OAuthToken) (*oauth.OAuthUser, error) { return nil, nil } @@ -60,7 +207,7 @@ func TestOAuthProviderUserUpdateField(t *testing.T) { func TestFindOrCreateOAuthUserMapsDeletedUserDomainError(t *testing.T) { _, err := findOrCreateOAuthUser(nil, &deletedUserOAuthProvider{}, &oauth.OAuthUser{ ProviderUserID: "deleted-user-id", - }, nil) + }, "") var deletedErr *OAuthUserDeletedError if !errors.As(err, &deletedErr) { @@ -85,7 +232,7 @@ func TestFindOrCreateOAuthUserPropagatesUniquenessLookupError(t *testing.T) { provider := &failingLookupOAuthProvider{err: lookupErr} _, err := findOrCreateOAuthUser(nil, provider, &oauth.OAuthUser{ ProviderUserID: "unverified-user-id", - }, nil) + }, "") if err == nil || !errors.Is(err, lookupErr) || !strings.Contains(err.Error(), "identity lookup failed") { t.Fatalf("expected identity lookup domain error, got %v", err) @@ -109,7 +256,7 @@ func TestFindOrCreateOAuthUserWrapsProviderFillDatabaseError(t *testing.T) { _, err := findOrCreateOAuthUser(nil, &publicFillFailingOAuthProvider{}, &oauth.OAuthUser{ ProviderUserID: "existing-user-id", Extra: map[string]any{}, - }, nil) + }, "") var lookupErr *oauthIdentityLookupError if !errors.As(err, &lookupErr) { t.Fatalf("expected masked identity lookup error, got %v", err) @@ -136,7 +283,7 @@ func TestOAuthBindMasksUniquenessLookupDatabaseError(t *testing.T) { c, _ := gin.CreateTestContext(recorder) c.Request = httptest.NewRequest(http.MethodGet, "/oauth/test?code=test-code", nil) - handleOAuthBind(c, &publicLookupFailingOAuthProvider{}) + handleOAuthBind(c, &publicLookupFailingOAuthProvider{}, &model.AuthFlow{}, "") var response struct { Success bool `json:"success"` diff --git a/controller/subscription_payment_epay.go b/controller/subscription_payment_epay.go index 69e381c1..b738fdc2 100644 --- a/controller/subscription_payment_epay.go +++ b/controller/subscription_payment_epay.go @@ -9,6 +9,7 @@ import ( "github.com/Calcium-Ion/go-epay/epay" "github.com/MAX-API-Next/MAX-API/common" + "github.com/MAX-API-Next/MAX-API/i18n" "github.com/MAX-API-Next/MAX-API/model" "github.com/MAX-API-Next/MAX-API/service" "github.com/MAX-API-Next/MAX-API/setting/operation_setting" @@ -75,7 +76,13 @@ func SubscriptionRequestEpay(c *gin.Context) { return } - tradeNo := fmt.Sprintf("%s%d", common.GetRandomString(6), time.Now().Unix()) + randomPart, err := common.GenerateRandomCharsKey(6) + if err != nil { + common.SysError("failed to generate Epay subscription order number: " + err.Error()) + common.ApiErrorI18n(c, i18n.MsgGenerateFailed) + return + } + tradeNo := fmt.Sprintf("%s%d", randomPart, time.Now().Unix()) tradeNo = fmt.Sprintf("SUBUSR%dNO%s", userId, tradeNo) client := GetEpayClient() @@ -160,7 +167,8 @@ func SubscriptionEpayNotify(c *gin.Context) { LockOrder(verifyInfo.ServiceTradeNo) defer UnlockOrder(verifyInfo.ServiceTradeNo) - if err := model.CompleteSubscriptionOrder(verifyInfo.ServiceTradeNo, common.GetJsonString(verifyInfo), model.PaymentProviderEpay, verifyInfo.Type); err != nil { + if err := model.CompleteSubscriptionOrder(verifyInfo.ServiceTradeNo, common.GetJsonString(verifyInfo), model.PaymentProviderEpay, verifyInfo.Type, + model.PaymentValidationFromMajorString(verifyInfo.Money, "", "", false)); err != nil { _, _ = c.Writer.Write([]byte("fail")) return } @@ -209,7 +217,8 @@ func SubscriptionEpayReturn(c *gin.Context) { if verifyInfo.TradeStatus == epay.StatusTradeSuccess { LockOrder(verifyInfo.ServiceTradeNo) defer UnlockOrder(verifyInfo.ServiceTradeNo) - if err := model.CompleteSubscriptionOrder(verifyInfo.ServiceTradeNo, common.GetJsonString(verifyInfo), model.PaymentProviderEpay, verifyInfo.Type); err != nil { + if err := model.CompleteSubscriptionOrder(verifyInfo.ServiceTradeNo, common.GetJsonString(verifyInfo), model.PaymentProviderEpay, verifyInfo.Type, + model.PaymentValidationFromMajorString(verifyInfo.Money, "", "", false)); err != nil { c.Redirect(http.StatusFound, paymentReturnPath("/console/topup?pay=fail")) return } diff --git a/controller/subscription_payment_waffo_pancake.go b/controller/subscription_payment_waffo_pancake.go index 2471c5f4..bc0d18b5 100644 --- a/controller/subscription_payment_waffo_pancake.go +++ b/controller/subscription_payment_waffo_pancake.go @@ -109,8 +109,6 @@ func SubscriptionRequestWaffoPancakePay(c *gin.Context) { }) if err != nil { logger.LogError(c.Request.Context(), fmt.Sprintf("Waffo Pancake 订阅结账会话创建失败 user_id=%d plan_id=%d trade_no=%s error=%q", userId, plan.Id, tradeNo, err.Error())) - order.Status = common.TopUpStatusFailed - _ = order.Update() c.JSON(http.StatusOK, gin.H{"message": "error", "data": "拉起支付失败"}) return } diff --git a/controller/task_video.go b/controller/task_video.go index 19ce0639..ce2d0027 100644 --- a/controller/task_video.go +++ b/controller/task_video.go @@ -2,324 +2,14 @@ package controller import ( "context" - "fmt" - "io" - "time" - "github.com/MAX-API-Next/MAX-API/common" "github.com/MAX-API-Next/MAX-API/constant" - "github.com/MAX-API-Next/MAX-API/dto" - "github.com/MAX-API-Next/MAX-API/logger" "github.com/MAX-API-Next/MAX-API/model" - "github.com/MAX-API-Next/MAX-API/relay" - "github.com/MAX-API-Next/MAX-API/relay/channel" - "github.com/MAX-API-Next/MAX-API/relay/channel/task/taskcommon" - relaycommon "github.com/MAX-API-Next/MAX-API/relay/common" - "github.com/MAX-API-Next/MAX-API/setting/ratio_setting" + "github.com/MAX-API-Next/MAX-API/service" ) +// UpdateVideoTaskAll is kept for compatibility with older controller callers. +// The service implementation owns polling, CAS, and billing settlement. func UpdateVideoTaskAll(ctx context.Context, platform constant.TaskPlatform, taskChannelM map[int][]string, taskM map[string]*model.Task) error { - for channelId, taskIds := range taskChannelM { - if err := updateVideoTaskAll(ctx, platform, channelId, taskIds, taskM); err != nil { - logger.LogError(ctx, fmt.Sprintf("Channel #%d failed to update video async tasks: %s", channelId, err.Error())) - } - } - return nil -} - -func updateVideoTaskAll(ctx context.Context, platform constant.TaskPlatform, channelId int, taskIds []string, taskM map[string]*model.Task) error { - logger.LogInfo(ctx, fmt.Sprintf("Channel #%d pending video tasks: %d", channelId, len(taskIds))) - if len(taskIds) == 0 { - return nil - } - cacheGetChannel, err := model.CacheGetChannel(channelId) - if err != nil { - errUpdate := model.TaskBulkUpdate(taskIds, map[string]any{ - "fail_reason": fmt.Sprintf("Failed to get channel info, channel ID: %d", channelId), - "status": "FAILURE", - "progress": "100%", - }) - if errUpdate != nil { - common.SysLog(fmt.Sprintf("UpdateVideoTask error: %v", errUpdate)) - } - return fmt.Errorf("CacheGetChannel failed: %w", err) - } - adaptor := relay.GetTaskAdaptor(platform) - if adaptor == nil { - return fmt.Errorf("video adaptor not found") - } - info := &relaycommon.RelayInfo{} - info.ChannelMeta = &relaycommon.ChannelMeta{ - ChannelType: cacheGetChannel.Type, - ChannelId: cacheGetChannel.Id, - ChannelBaseUrl: cacheGetChannel.GetBaseURL(), - ChannelSetting: cacheGetChannel.GetSetting(), - ChannelOtherSettings: cacheGetChannel.GetOtherSettings(), - } - info.ApiKey = cacheGetChannel.Key - adaptor.Init(info) - for _, taskId := range taskIds { - if err := updateVideoSingleTask(ctx, adaptor, cacheGetChannel, taskId, taskM); err != nil { - logger.LogError(ctx, fmt.Sprintf("Failed to update video task %s: %s", taskId, err.Error())) - } - } - return nil -} - -func updateVideoSingleTask(ctx context.Context, adaptor channel.TaskAdaptor, channel *model.Channel, taskId string, taskM map[string]*model.Task) error { - baseURL := constant.ChannelBaseURLs[channel.Type] - if channel.GetBaseURL() != "" { - baseURL = channel.GetBaseURL() - } - proxy := channel.GetSetting().Proxy - - task := taskM[taskId] - if task == nil { - logger.LogError(ctx, fmt.Sprintf("Task %s not found in taskM", taskId)) - return fmt.Errorf("task %s not found", taskId) - } - key := channel.Key - - privateData := task.PrivateData - if privateData.Key != "" { - key = privateData.Key - } - resp, err := adaptor.FetchTask(baseURL, key, taskcommon.WithTaskProtocolConfig(map[string]any{ - "task_id": taskId, - "action": task.Action, - "model": task.Properties.UpstreamModelName, - }, channel.GetOtherSettings()), proxy) - if err != nil { - return fmt.Errorf("fetchTask failed for task %s: %w", taskId, err) - } - //if resp.StatusCode != http.StatusOK { - //return fmt.Errorf("get Video Task status code: %d", resp.StatusCode) - //} - defer resp.Body.Close() - responseBody, err := io.ReadAll(resp.Body) - if err != nil { - return fmt.Errorf("readAll failed for task %s: %w", taskId, err) - } - - logger.LogDebug(ctx, "UpdateVideoSingleTask response: %s", responseBody) - - taskResult := &relaycommon.TaskInfo{} - // try parse as MAX API response format - var responseItems dto.TaskResponse[model.Task] - if err = common.Unmarshal(responseBody, &responseItems); err == nil && responseItems.IsSuccess() { - logger.LogDebug(ctx, "UpdateVideoSingleTask parsed as compatible response format: %+v", responseItems) - t := responseItems.Data - taskResult.TaskID = t.TaskID - taskResult.Status = string(t.Status) - taskResult.Url = t.FailReason - taskResult.Progress = t.Progress - taskResult.Reason = t.FailReason - task.Data = t.Data - } else if configuredResult, ok, parseErr := taskcommon.ParseConfiguredTaskResult(responseBody, channel.GetOtherSettings()); parseErr != nil { - return fmt.Errorf("parse configured task result failed for task %s: %w", taskId, parseErr) - } else if ok { - taskResult = configuredResult - task.Data = redactVideoResponseBody(responseBody) - } else if taskResult, err = adaptor.ParseTaskResult(responseBody); err != nil { - return fmt.Errorf("parseTaskResult failed for task %s: %w", taskId, err) - } else { - task.Data = redactVideoResponseBody(responseBody) - } - - logger.LogDebug(ctx, "UpdateVideoSingleTask taskResult: %+v", taskResult) - - now := time.Now().Unix() - if taskResult.Status == "" { - //return fmt.Errorf("task %s status is empty", taskId) - taskResult = relaycommon.FailTaskInfo("upstream returned empty status") - } - - // 记录原本的状态，防止重复退款 - shouldRefund := false - quota := task.Quota - preStatus := task.Status - - task.Status = model.TaskStatus(taskResult.Status) - switch taskResult.Status { - case model.TaskStatusSubmitted: - task.Progress = "10%" - case model.TaskStatusQueued: - task.Progress = "20%" - case model.TaskStatusInProgress: - task.Progress = "30%" - if task.StartTime == 0 { - task.StartTime = now - } - case model.TaskStatusSuccess: - task.Progress = "100%" - if task.FinishTime == 0 { - task.FinishTime = now - } - if !(len(taskResult.Url) > 5 && taskResult.Url[:5] == "data:") { - task.FailReason = taskResult.Url - } - - // 如果返回了 total_tokens 并且配置了模型倍率(非固定价格),则重新计费 - if task.DeltaSettlementDisabledForChannel(channel.Type, channel.GetOtherSettings()) { - logger.LogInfo(ctx, fmt.Sprintf("视频任务 %s 已按渠道设置跳过完成态差额结算", task.TaskID)) - } else if taskResult.TotalTokens > 0 { - // 获取模型名称 - var taskData map[string]interface{} - if err := common.Unmarshal(task.Data, &taskData); err == nil { - if modelName, ok := taskData["model"].(string); ok && modelName != "" { - // 获取模型价格和倍率 - modelRatio, hasRatioSetting, _ := ratio_setting.GetModelRatio(modelName) - // 只有配置了倍率(非固定价格)时才按 token 重新计费 - if hasRatioSetting && modelRatio > 0 { - // 获取用户和组的倍率信息 - group := task.Group - if group == "" { - user, err := model.GetUserById(task.UserId, false) - if err == nil { - group = user.Group - } - } - if group != "" { - groupRatio := ratio_setting.GetGroupRatio(group) - userGroupRatio, hasUserGroupRatio := ratio_setting.GetGroupGroupRatio(group, group) - - var finalGroupRatio float64 - if hasUserGroupRatio { - finalGroupRatio = userGroupRatio - } else { - finalGroupRatio = groupRatio - } - - // 计算实际应扣费额度: totalTokens * modelRatio * groupRatio - actualQuota := common.QuotaFromFloat(float64(taskResult.TotalTokens) * modelRatio * finalGroupRatio) - - // 计算差额 - preConsumedQuota := task.Quota - quotaDelta := actualQuota - preConsumedQuota - - if quotaDelta > 0 { - // 需要补扣费 - logger.LogInfo(ctx, fmt.Sprintf("视频任务 %s 预扣费后补扣费：%s（实际消耗：%s，预扣费：%s，tokens：%d）", - task.TaskID, - logger.LogQuota(quotaDelta), - logger.LogQuota(actualQuota), - logger.LogQuota(preConsumedQuota), - taskResult.TotalTokens, - )) - if err := model.DecreaseUserQuota(task.UserId, quotaDelta, false); err != nil { - logger.LogError(ctx, fmt.Sprintf("补扣费失败: %s", err.Error())) - } else { - model.UpdateUserUsedQuotaAndRequestCount(task.UserId, quotaDelta) - model.UpdateChannelUsedQuota(task.ChannelId, quotaDelta) - task.Quota = actualQuota // 更新任务记录的实际扣费额度 - - // 记录消费日志 - logContent := fmt.Sprintf("视频任务成功补扣费，模型倍率 %.2f，分组倍率 %.2f，tokens %d，预扣费 %s，实际扣费 %s，补扣费 %s", - modelRatio, finalGroupRatio, taskResult.TotalTokens, - logger.LogQuota(preConsumedQuota), logger.LogQuota(actualQuota), logger.LogQuota(quotaDelta)) - model.RecordLog(task.UserId, model.LogTypeSystem, logContent) - } - } else if quotaDelta < 0 { - // 需要退还多扣的费用 - refundQuota := -quotaDelta - logger.LogInfo(ctx, fmt.Sprintf("视频任务 %s 预扣费后返还：%s（实际消耗：%s，预扣费：%s，tokens：%d）", - task.TaskID, - logger.LogQuota(refundQuota), - logger.LogQuota(actualQuota), - logger.LogQuota(preConsumedQuota), - taskResult.TotalTokens, - )) - if err := model.IncreaseUserQuota(task.UserId, refundQuota, false); err != nil { - logger.LogError(ctx, fmt.Sprintf("退还预扣费失败: %s", err.Error())) - } else { - task.Quota = actualQuota // 更新任务记录的实际扣费额度 - - // 记录退款日志 - logContent := fmt.Sprintf("视频任务成功退还多扣费用，模型倍率 %.2f，分组倍率 %.2f，tokens %d，预扣费 %s，实际扣费 %s，退还 %s", - modelRatio, finalGroupRatio, taskResult.TotalTokens, - logger.LogQuota(preConsumedQuota), logger.LogQuota(actualQuota), logger.LogQuota(refundQuota)) - model.RecordLog(task.UserId, model.LogTypeSystem, logContent) - } - } else { - // quotaDelta == 0, 预扣费刚好准确 - logger.LogInfo(ctx, fmt.Sprintf("视频任务 %s 预扣费准确（%s，tokens：%d）", - task.TaskID, logger.LogQuota(actualQuota), taskResult.TotalTokens)) - } - } - } - } - } - } - case model.TaskStatusFailure: - logger.LogJson(ctx, fmt.Sprintf("Task %s failed", taskId), task) - task.Status = model.TaskStatusFailure - task.Progress = "100%" - if task.FinishTime == 0 { - task.FinishTime = now - } - task.FailReason = taskResult.Reason - logger.LogInfo(ctx, fmt.Sprintf("Task %s failed: %s", task.TaskID, task.FailReason)) - taskResult.Progress = "100%" - if quota != 0 { - if preStatus != model.TaskStatusFailure { - shouldRefund = true - } else { - logger.LogWarn(ctx, fmt.Sprintf("Task %s already in failure status, skip refund", task.TaskID)) - } - } - default: - return fmt.Errorf("unknown task status %s for task %s", taskResult.Status, taskId) - } - if taskResult.Progress != "" { - task.Progress = taskResult.Progress - } - if err := task.Update(); err != nil { - common.SysLog("UpdateVideoTask task error: " + err.Error()) - shouldRefund = false - } - - if shouldRefund { - // 任务失败且之前状态不是失败才退还额度，防止重复退还 - if err := model.IncreaseUserQuota(task.UserId, quota, false); err != nil { - logger.LogWarn(ctx, "Failed to increase user quota: "+err.Error()) - } - logContent := fmt.Sprintf("Video async task failed %s, refund %s", task.TaskID, logger.LogQuota(quota)) - model.RecordLog(task.UserId, model.LogTypeSystem, logContent) - } - - return nil -} - -func redactVideoResponseBody(body []byte) []byte { - var m map[string]any - if err := common.Unmarshal(body, &m); err != nil { - return body - } - resp, _ := m["response"].(map[string]any) - if resp != nil { - delete(resp, "bytesBase64Encoded") - if v, ok := resp["video"].(string); ok { - resp["video"] = truncateBase64(v) - } - if vs, ok := resp["videos"].([]any); ok { - for i := range vs { - if vm, ok := vs[i].(map[string]any); ok { - delete(vm, "bytesBase64Encoded") - } - } - } - } - b, err := common.Marshal(m) - if err != nil { - return body - } - return b -} - -func truncateBase64(s string) string { - const maxKeep = 256 - if len(s) <= maxKeep { - return s - } - return s[:maxKeep] + "..." + return service.UpdateVideoTasks(ctx, platform, taskChannelM, taskM) } diff --git a/controller/topup.go b/controller/topup.go index e4250a29..a204191d 100644 --- a/controller/topup.go +++ b/controller/topup.go @@ -9,6 +9,7 @@ import ( "time" "github.com/MAX-API-Next/MAX-API/common" + "github.com/MAX-API-Next/MAX-API/i18n" "github.com/MAX-API-Next/MAX-API/logger" "github.com/MAX-API-Next/MAX-API/model" "github.com/MAX-API-Next/MAX-API/service" @@ -218,7 +219,13 @@ func RequestEpay(c *gin.Context) { callBackAddress := service.GetCallbackAddress() returnUrl, _ := url.Parse(paymentReturnPath("/console/log")) notifyUrl, _ := url.Parse(callBackAddress + "/api/user/epay/notify") - tradeNo := fmt.Sprintf("%s%d", common.GetRandomString(6), time.Now().Unix()) + randomPart, err := common.GenerateRandomCharsKey(6) + if err != nil { + logger.LogError(c.Request.Context(), "生成易支付订单号失败: "+err.Error()) + common.ApiErrorI18n(c, i18n.MsgGenerateFailed) + return + } + tradeNo := fmt.Sprintf("%s%d", randomPart, time.Now().Unix()) tradeNo = fmt.Sprintf("USR%dNO%s", id, tradeNo) client := GetEpayClient() if client == nil { @@ -351,16 +358,10 @@ func EpayNotify(c *gin.Context) { return } verifyInfo, err := client.Verify(params) - if err == nil && verifyInfo.VerifyStatus { - logger.LogInfo(c.Request.Context(), fmt.Sprintf("易支付 webhook 验签成功 trade_no=%s callback_type=%s trade_status=%s client_ip=%s verify_info=%q", verifyInfo.ServiceTradeNo, verifyInfo.Type, verifyInfo.TradeStatus, c.ClientIP(), common.GetJsonString(verifyInfo))) - _, err := c.Writer.Write([]byte("success")) - if err != nil { - logger.LogError(c.Request.Context(), fmt.Sprintf("易支付 webhook 响应写入失败 trade_no=%s client_ip=%s error=%q", verifyInfo.ServiceTradeNo, c.ClientIP(), err.Error())) - } - } else { - _, err := c.Writer.Write([]byte("fail")) - if err != nil { - logger.LogError(c.Request.Context(), fmt.Sprintf("易支付 webhook 响应写入失败 path=%q client_ip=%s error=%q", c.Request.RequestURI, c.ClientIP(), err.Error())) + if err != nil || !verifyInfo.VerifyStatus { + _, writeErr := c.Writer.Write([]byte("fail")) + if writeErr != nil { + logger.LogError(c.Request.Context(), fmt.Sprintf("易支付 webhook 响应写入失败 path=%q client_ip=%s error=%q", c.Request.RequestURI, c.ClientIP(), writeErr.Error())) } if err != nil { logger.LogWarn(c.Request.Context(), fmt.Sprintf("易支付 webhook 验签失败 path=%q client_ip=%s verify_error=%q", c.Request.RequestURI, c.ClientIP(), err.Error())) @@ -370,45 +371,23 @@ func EpayNotify(c *gin.Context) { return } - if verifyInfo.TradeStatus == epay.StatusTradeSuccess { - LockOrder(verifyInfo.ServiceTradeNo) - defer UnlockOrder(verifyInfo.ServiceTradeNo) - topUp := model.GetTopUpByTradeNo(verifyInfo.ServiceTradeNo) - if topUp == nil { - logger.LogWarn(c.Request.Context(), fmt.Sprintf("易支付 回调订单不存在 trade_no=%s callback_type=%s client_ip=%s verify_info=%q", verifyInfo.ServiceTradeNo, verifyInfo.Type, c.ClientIP(), common.GetJsonString(verifyInfo))) - return - } - if topUp.PaymentProvider != model.PaymentProviderEpay { - logger.LogWarn(c.Request.Context(), fmt.Sprintf("易支付 订单支付网关不匹配 trade_no=%s order_provider=%s callback_type=%s client_ip=%s", verifyInfo.ServiceTradeNo, topUp.PaymentProvider, verifyInfo.Type, c.ClientIP())) - return - } - if topUp.Status == common.TopUpStatusPending { - if topUp.PaymentMethod != verifyInfo.Type { - logger.LogInfo(c.Request.Context(), fmt.Sprintf("易支付 实际支付方式与订单不同 trade_no=%s order_payment_method=%s actual_type=%s client_ip=%s", verifyInfo.ServiceTradeNo, topUp.PaymentMethod, verifyInfo.Type, c.ClientIP())) - topUp.PaymentMethod = verifyInfo.Type - } - topUp.Status = common.TopUpStatusSuccess - err := topUp.Update() - if err != nil { - logger.LogError(c.Request.Context(), fmt.Sprintf("易支付 更新充值订单失败 trade_no=%s user_id=%d client_ip=%s error=%q topup=%q", topUp.TradeNo, topUp.UserId, c.ClientIP(), err.Error(), common.GetJsonString(topUp))) - return - } - //user, _ := model.GetUserById(topUp.UserId, false) - //user.Quota += topUp.Amount * 500000 - dAmount := decimal.NewFromInt(int64(topUp.Amount)) - dQuotaPerUnit := decimal.NewFromFloat(common.QuotaPerUnit) - quotaToAdd := int(dAmount.Mul(dQuotaPerUnit).IntPart()) - err = model.IncreaseUserQuota(topUp.UserId, quotaToAdd, true) - if err != nil { - logger.LogError(c.Request.Context(), fmt.Sprintf("易支付 更新用户额度失败 trade_no=%s user_id=%d client_ip=%s quota_to_add=%d error=%q topup=%q", topUp.TradeNo, topUp.UserId, c.ClientIP(), quotaToAdd, err.Error(), common.GetJsonString(topUp))) - return - } - logger.LogInfo(c.Request.Context(), fmt.Sprintf("易支付 充值成功 trade_no=%s user_id=%d client_ip=%s quota_to_add=%d money=%.2f topup=%q", topUp.TradeNo, topUp.UserId, c.ClientIP(), quotaToAdd, topUp.Money, common.GetJsonString(topUp))) - model.RecordTopupLog(topUp.UserId, fmt.Sprintf("使用在线充值成功，充值金额: %v，支付金额：%f", logger.LogQuota(quotaToAdd), topUp.Money), c.ClientIP(), topUp.PaymentMethod, "epay") - } - } else { - logger.LogInfo(c.Request.Context(), fmt.Sprintf("易支付 webhook 忽略事件 trade_no=%s callback_type=%s trade_status=%s client_ip=%s verify_info=%q", verifyInfo.ServiceTradeNo, verifyInfo.Type, verifyInfo.TradeStatus, c.ClientIP(), common.GetJsonString(verifyInfo))) + logger.LogInfo(c.Request.Context(), fmt.Sprintf("易支付 webhook 验签成功 trade_no=%s callback_type=%s trade_status=%s client_ip=%s verify_info=%q", verifyInfo.ServiceTradeNo, verifyInfo.Type, verifyInfo.TradeStatus, c.ClientIP(), common.GetJsonString(verifyInfo))) + if verifyInfo.TradeStatus != epay.StatusTradeSuccess { + logger.LogInfo(c.Request.Context(), fmt.Sprintf("易支付 webhook 忽略事件 trade_no=%s callback_type=%s trade_status=%s client_ip=%s", verifyInfo.ServiceTradeNo, verifyInfo.Type, verifyInfo.TradeStatus, c.ClientIP())) + _, _ = c.Writer.Write([]byte("success")) + return + } + + LockOrder(verifyInfo.ServiceTradeNo) + defer UnlockOrder(verifyInfo.ServiceTradeNo) + if err := model.RechargeEpay(verifyInfo.ServiceTradeNo, verifyInfo.Type, c.ClientIP(), + model.PaymentValidationFromMajorString(verifyInfo.Money, "", "", false)); err != nil { + logger.LogError(c.Request.Context(), fmt.Sprintf("易支付 充值处理失败 trade_no=%s callback_type=%s client_ip=%s error=%q", verifyInfo.ServiceTradeNo, verifyInfo.Type, c.ClientIP(), err.Error())) + _, _ = c.Writer.Write([]byte("fail")) + return } + + _, _ = c.Writer.Write([]byte("success")) } func RequestAmount(c *gin.Context) { diff --git a/controller/topup_creem.go b/controller/topup_creem.go index e2aa1c22..6a76b900 100644 --- a/controller/topup_creem.go +++ b/controller/topup_creem.go @@ -6,7 +6,6 @@ import ( "crypto/hmac" "crypto/sha256" "encoding/hex" - "encoding/json" "errors" "fmt" "github.com/MAX-API-Next/MAX-API/common" @@ -60,6 +59,39 @@ type CreemProduct struct { Quota int64 `json:"quota"` } +func expectedCreemCurrencyForProduct(productId string) string { + if productId == "" { + return "" + } + var products []CreemProduct + if err := common.UnmarshalJsonStr(setting.CreemProducts, &products); err != nil { + common.SysLog("failed to parse Creem products for payment validation: " + err.Error()) + return "" + } + for _, product := range products { + if product.ProductId == productId { + return product.Currency + } + } + return "" +} + +func creemTopUpPaymentValidation(event *CreemWebhookEvent) (model.PaymentValidation, error) { + if event == nil { + return model.PaymentValidation{}, errors.New("Creem webhook event is nil") + } + expectedCurrency := expectedCreemCurrencyForProduct(event.Object.Product.Id) + if expectedCurrency == "" { + return model.PaymentValidation{}, fmt.Errorf("Creem product currency is not configured: %s", event.Object.Product.Id) + } + return model.PaymentValidationFromMinorUnits( + int64(event.Object.Order.AmountPaid), + event.Object.Order.Currency, + expectedCurrency, + false, + ), nil +} + type CreemAdaptor struct { } @@ -76,7 +108,7 @@ func (*CreemAdaptor) RequestPay(c *gin.Context, req *CreemPayRequest) { // 解析产品列表 var products []CreemProduct - err := json.Unmarshal([]byte(setting.CreemProducts), &products) + err := common.UnmarshalJsonStr(setting.CreemProducts, &products) if err != nil { logger.LogError(c.Request.Context(), fmt.Sprintf("Creem 产品配置解析失败 user_id=%d error=%q", c.GetInt("id"), err.Error())) c.JSON(http.StatusOK, gin.H{"message": "error", "data": "产品配置错误"}) @@ -302,7 +334,8 @@ func handleCheckoutCompleted(c *gin.Context, event *CreemWebhookEvent) { // Try complete subscription order first LockOrder(referenceId) defer UnlockOrder(referenceId) - if err := model.CompleteSubscriptionOrder(referenceId, common.GetJsonString(event), model.PaymentProviderCreem, ""); err == nil { + paymentValidation := model.PaymentValidationFromMinorUnits(int64(event.Object.Order.AmountPaid), event.Object.Order.Currency, "", false) + if err := model.CompleteSubscriptionOrder(referenceId, common.GetJsonString(event), model.PaymentProviderCreem, "", paymentValidation); err == nil { logger.LogInfo(c.Request.Context(), fmt.Sprintf("Creem 订阅订单处理成功 trade_no=%s creem_order_id=%s", referenceId, event.Object.Order.Id)) c.Status(http.StatusOK) return @@ -347,7 +380,13 @@ func handleCheckoutCompleted(c *gin.Context, event *CreemWebhookEvent) { logger.LogWarn(c.Request.Context(), fmt.Sprintf("Creem 回调客户姓名为空 trade_no=%s creem_order_id=%s", referenceId, event.Object.Order.Id)) } - err := model.RechargeCreem(referenceId, customerEmail, customerName, c.ClientIP()) + topUpValidation, err := creemTopUpPaymentValidation(event) + if err != nil { + logger.LogError(c.Request.Context(), fmt.Sprintf("Creem 充值币种校验配置缺失 trade_no=%s creem_order_id=%s product_id=%s error=%q", referenceId, event.Object.Order.Id, event.Object.Product.Id, err.Error())) + c.AbortWithStatus(http.StatusInternalServerError) + return + } + err = model.RechargeCreem(referenceId, customerEmail, customerName, c.ClientIP(), topUpValidation) if err != nil { logger.LogError(c.Request.Context(), fmt.Sprintf("Creem 充值处理失败 trade_no=%s creem_order_id=%s client_ip=%s error=%q", referenceId, event.Object.Order.Id, c.ClientIP(), err.Error())) c.AbortWithStatus(http.StatusInternalServerError) @@ -402,7 +441,7 @@ func genCreemLink(ctx context.Context, referenceId string, product *CreemProduct } // 序列化请求数据 - jsonData, err := json.Marshal(requestData) + jsonData, err := common.Marshal(requestData) if err != nil { return "", fmt.Errorf("序列化请求数据失败: %v", err) } @@ -443,7 +482,7 @@ func genCreemLink(ctx context.Context, referenceId string, product *CreemProduct } // 解析响应 var checkoutResp CreemCheckoutResponse - err = json.Unmarshal(body, &checkoutResp) + err = common.Unmarshal(body, &checkoutResp) if err != nil { return "", fmt.Errorf("解析响应失败: %v", err) } diff --git a/controller/topup_creem_test.go b/controller/topup_creem_test.go new file mode 100644 index 00000000..396dd6e2 --- /dev/null +++ b/controller/topup_creem_test.go @@ -0,0 +1,41 @@ +package controller + +import ( + "testing" + + "github.com/MAX-API-Next/MAX-API/setting" + "github.com/stretchr/testify/require" +) + +func TestCreemTopUpPaymentValidationUsesConfiguredProductCurrency(t *testing.T) { + originalProducts := setting.CreemProducts + t.Cleanup(func() { setting.CreemProducts = originalProducts }) + setting.CreemProducts = `[{"productId":"prod-usd","currency":"USD"}]` + + event := &CreemWebhookEvent{} + event.Object.Product.Id = "prod-usd" + event.Object.Order.AmountPaid = 1234 + event.Object.Order.Currency = "usd" + + validation, err := creemTopUpPaymentValidation(event) + require.NoError(t, err) + require.EqualValues(t, 1234, validation.PaidAmountMinor) + require.Equal(t, "USD", validation.Currency) + require.Equal(t, "USD", validation.ExpectedCurrency) +} + +func TestCreemTopUpPaymentValidationRejectsUnknownProductCurrency(t *testing.T) { + originalProducts := setting.CreemProducts + t.Cleanup(func() { setting.CreemProducts = originalProducts }) + setting.CreemProducts = `[{"productId":"configured-product","currency":"USD"}]` + + event := &CreemWebhookEvent{} + event.Object.Product.Id = "unknown-product" + event.Object.Product.Currency = "EUR" + event.Object.Order.AmountPaid = 1234 + event.Object.Order.Currency = "EUR" + + _, err := creemTopUpPaymentValidation(event) + require.Error(t, err) + require.Contains(t, err.Error(), "currency is not configured") +} diff --git a/controller/topup_stripe.go b/controller/topup_stripe.go index 2c66a371..ef730e06 100644 --- a/controller/topup_stripe.go +++ b/controller/topup_stripe.go @@ -223,6 +223,30 @@ func sessionAsyncPaymentSucceeded(ctx context.Context, event stripe.Event, calle return fulfillOrder(ctx, event, referenceId, customerId, callerIp) } +func stripePaymentValidationFromEvent(event stripe.Event, allowDiscount bool) (model.PaymentValidation, error) { + amountTotal := strings.TrimSpace(event.GetObjectValue("amount_total")) + if amountTotal == "" { + return model.PaymentValidation{}, fmt.Errorf("stripe event missing amount_total") + } + amountMinor, err := strconv.ParseInt(amountTotal, 10, 64) + if err != nil { + return model.PaymentValidation{}, fmt.Errorf("parse stripe amount_total: %w", err) + } + return model.PaymentValidationFromMinorUnits(amountMinor, event.GetObjectValue("currency"), "", allowDiscount), nil +} + +func isStripeOrderAlreadyCompleted(referenceId string) bool { + if order := model.GetSubscriptionOrderByTradeNo(referenceId); order != nil && + order.PaymentProvider == model.PaymentProviderStripe && + order.Status == common.TopUpStatusSuccess { + return true + } + topUp, err := model.GetTopUpByTradeNoWithError(referenceId) + return err == nil && + topUp.PaymentProvider == model.PaymentProviderStripe && + topUp.Status == common.TopUpStatusSuccess +} + // sessionAsyncPaymentFailed marks orders as failed when delayed payment methods // ultimately fail (e.g. bank transfer not received, SEPA rejected). func sessionAsyncPaymentFailed(ctx context.Context, event stripe.Event, callerIp string) { @@ -271,13 +295,23 @@ func fulfillOrder(ctx context.Context, event stripe.Event, referenceId string, c LockOrder(referenceId) defer UnlockOrder(referenceId) + validation, err := stripePaymentValidationFromEvent(event, setting.StripePromotionCodesEnabled) + if err != nil { + if isStripeOrderAlreadyCompleted(referenceId) { + logger.LogInfo(ctx, fmt.Sprintf("Stripe 订单已完成，确认缺少金额字段的重放事件 trade_no=%s event_type=%s client_ip=%s", referenceId, string(event.Type), callerIp)) + return nil + } + return err + } payload := map[string]any{ "customer": customerId, "amount_total": event.GetObjectValue("amount_total"), "currency": strings.ToUpper(event.GetObjectValue("currency")), "event_type": string(event.Type), } - if err := model.CompleteSubscriptionOrder(referenceId, common.GetJsonString(payload), model.PaymentProviderStripe, ""); err == nil { + subscriptionValidation := validation + subscriptionValidation.AllowDiscount = false + if err := model.CompleteSubscriptionOrder(referenceId, common.GetJsonString(payload), model.PaymentProviderStripe, "", subscriptionValidation); err == nil { logger.LogInfo(ctx, fmt.Sprintf("Stripe 订阅订单处理成功 trade_no=%s event_type=%s client_ip=%s", referenceId, string(event.Type), callerIp)) return nil } else if !errors.Is(err, model.ErrSubscriptionOrderNotFound) { @@ -309,7 +343,9 @@ func fulfillOrder(ctx context.Context, event stripe.Event, referenceId string, c return nil } - if err := model.Recharge(referenceId, customerId, callerIp); err != nil { + topUpValidation := validation + topUpValidation.ExpectedCurrency = "USD" + if err := model.Recharge(referenceId, customerId, callerIp, topUpValidation); err != nil { latest, latestErr := model.GetTopUpByTradeNoWithError(referenceId) if latestErr == nil && latest.Status == common.TopUpStatusSuccess { logger.LogInfo(ctx, fmt.Sprintf("Stripe 充值订单已由并发回调完成 trade_no=%s client_ip=%s", referenceId, callerIp)) diff --git a/controller/topup_waffo.go b/controller/topup_waffo.go index 8e1b03a1..625fa865 100644 --- a/controller/topup_waffo.go +++ b/controller/topup_waffo.go @@ -6,6 +6,7 @@ import ( "io" "net/http" "strconv" + "strings" "time" "github.com/MAX-API-Next/MAX-API/common" @@ -59,6 +60,16 @@ func getWaffoCurrency() string { return "USD" } +func getWaffoPaidAmount(result *core.PaymentNotificationResult) string { + if result == nil { + return "" + } + if amount := strings.TrimSpace(result.FinalDealAmount); amount != "" { + return amount + } + return result.OrderAmount +} + // zeroDecimalCurrencies 零小数位币种，金额不能带小数点 var zeroDecimalCurrencies = map[string]bool{ "IDR": true, "JPY": true, "KRW": true, "VND": true, @@ -269,15 +280,11 @@ func RequestWaffoPay(c *gin.Context) { resp, err := sdk.Order().Create(c.Request.Context(), createParams, nil) if err != nil { logger.LogError(c.Request.Context(), fmt.Sprintf("Waffo 创建订单失败 user_id=%d trade_no=%s error=%q", id, merchantOrderId, err.Error())) - topUp.Status = common.TopUpStatusFailed - _ = topUp.Update() c.JSON(http.StatusOK, gin.H{"message": "error", "data": "拉起支付失败"}) return } if !resp.IsSuccess() { logger.LogWarn(c.Request.Context(), fmt.Sprintf("Waffo 创建订单业务失败 user_id=%d trade_no=%s code=%s message=%q response=%q", id, merchantOrderId, resp.Code, resp.Message, common.GetJsonString(resp))) - topUp.Status = common.TopUpStatusFailed - _ = topUp.Update() c.JSON(http.StatusOK, gin.H{"message": "error", "data": "拉起支付失败"}) return } @@ -394,7 +401,8 @@ func handleWaffoPayment(c *gin.Context, wh *core.WebhookHandler, result *core.Pa LockOrder(merchantOrderId) defer UnlockOrder(merchantOrderId) - if err := model.RechargeWaffo(merchantOrderId, c.ClientIP()); err != nil { + if err := model.RechargeWaffo(merchantOrderId, c.ClientIP(), + model.PaymentValidationFromMajorString(getWaffoPaidAmount(result), result.OrderCurrency, getWaffoCurrency(), false)); err != nil { logger.LogError(c.Request.Context(), fmt.Sprintf("Waffo 充值处理失败 trade_no=%s client_ip=%s error=%q", merchantOrderId, c.ClientIP(), err.Error())) sendWaffoWebhookResponse(c, wh, false, err.Error()) return diff --git a/controller/topup_waffo_pancake.go b/controller/topup_waffo_pancake.go index 82bd9f19..6b9e7fb8 100644 --- a/controller/topup_waffo_pancake.go +++ b/controller/topup_waffo_pancake.go @@ -411,8 +411,6 @@ func RequestWaffoPancakePay(c *gin.Context) { }) if err != nil { logger.LogError(c.Request.Context(), fmt.Sprintf("Waffo Pancake 创建结账会话失败 user_id=%d trade_no=%s error=%q", id, tradeNo, err.Error())) - topUp.Status = common.TopUpStatusFailed - _ = topUp.Update() c.JSON(http.StatusOK, gin.H{"message": "error", "data": "拉起支付失败"}) return } @@ -500,7 +498,8 @@ func WaffoPancakeWebhook(c *gin.Context) { } LockOrder(tradeNo) defer UnlockOrder(tradeNo) - if err := model.CompleteSubscriptionOrder(tradeNo, string(bodyBytes), model.PaymentProviderWaffoPancake, ""); err != nil { + if err := model.CompleteSubscriptionOrder(tradeNo, string(bodyBytes), model.PaymentProviderWaffoPancake, "", + model.PaymentValidationFromMajorString(event.Data.Amount, event.Data.Currency, "", false)); err != nil { logger.LogError(c.Request.Context(), fmt.Sprintf("Waffo Pancake 订阅完成失败 trade_no=%s event_id=%s order_id=%s client_ip=%s error=%q", tradeNo, event.ID, event.Data.OrderID, c.ClientIP(), err.Error())) c.String(http.StatusInternalServerError, "retry") return @@ -526,7 +525,8 @@ func WaffoPancakeWebhook(c *gin.Context) { LockOrder(tradeNo) defer UnlockOrder(tradeNo) - if err := model.RechargeWaffoPancake(tradeNo); err != nil { + if err := model.RechargeWaffoPancake(tradeNo, + model.PaymentValidationFromMajorString(event.Data.Amount, event.Data.Currency, "USD", false)); err != nil { logger.LogError(c.Request.Context(), fmt.Sprintf("Waffo Pancake 充值处理失败 trade_no=%s event_id=%s order_id=%s client_ip=%s error=%q", tradeNo, event.ID, event.Data.OrderID, c.ClientIP(), err.Error())) c.String(http.StatusInternalServerError, "retry") return diff --git a/controller/topup_waffo_pancake_test.go b/controller/topup_waffo_pancake_test.go index 5249ee78..097f043e 100644 --- a/controller/topup_waffo_pancake_test.go +++ b/controller/topup_waffo_pancake_test.go @@ -1,14 +1,71 @@ package controller import ( + "bytes" + "fmt" + "net/http" + "net/http/httptest" + "os" + "strings" "testing" "github.com/MAX-API-Next/MAX-API/common" + "github.com/MAX-API-Next/MAX-API/model" "github.com/MAX-API-Next/MAX-API/setting" "github.com/MAX-API-Next/MAX-API/setting/operation_setting" + "github.com/gin-gonic/gin" "github.com/stretchr/testify/require" + "gorm.io/gorm" ) +func setupWaffoPancakeControllerTestDB(t *testing.T) *gorm.DB { + t.Helper() + + gin.SetMode(gin.TestMode) + oldDB := model.DB + oldLOGDB := model.LOG_DB + oldSQLDSN := os.Getenv("SQL_DSN") + oldSQLitePath := common.SQLitePath + oldRedisEnabled := common.RedisEnabled + oldUsingSQLite := common.UsingSQLite + oldUsingMySQL := common.UsingMySQL + oldUsingPostgreSQL := common.UsingPostgreSQL + + common.RedisEnabled = false + common.UsingSQLite = true + common.UsingMySQL = false + common.UsingPostgreSQL = false + + dsn := fmt.Sprintf("file:%s?mode=memory&cache=shared", strings.ReplaceAll(t.Name(), "/", "_")) + common.SQLitePath = dsn + require.NoError(t, os.Setenv("SQL_DSN", "local")) + require.NoError(t, model.InitDB()) + require.NoError(t, model.InitLogDB()) + db := model.DB + require.NoError(t, db.AutoMigrate(&model.User{}, &model.TopUp{})) + + t.Cleanup(func() { + sqlDB, err := db.DB() + if err == nil { + _ = sqlDB.Close() + } + model.DB = oldDB + model.LOG_DB = oldLOGDB + common.SQLitePath = oldSQLitePath + if oldSQLDSN == "" { + _ = os.Unsetenv("SQL_DSN") + } else { + _ = os.Setenv("SQL_DSN", oldSQLDSN) + } + common.RedisEnabled = oldRedisEnabled + common.UsingSQLite = oldUsingSQLite + common.UsingMySQL = oldUsingMySQL + common.UsingPostgreSQL = oldUsingPostgreSQL + }) + + return db +} + func TestFormatWaffoPancakeAmount_UsesDisplayPriceString(t *testing.T) { testCases := []struct { name string @@ -89,3 +146,56 @@ func TestGetWaffoPancakePayMoney(t *testing.T) { }) } } + +func TestRequestWaffoPancakePayLeavesOrderPendingWhenSessionCreateFails(t *testing.T) { + db := setupWaffoPancakeControllerTestDB(t) + confirmPaymentComplianceForTest(t) + + originalMerchantID := setting.WaffoPancakeMerchantID + originalPrivateKey := setting.WaffoPancakePrivateKey + originalProductID := setting.WaffoPancakeProductID + originalUnitPrice := setting.WaffoPancakeUnitPrice + originalMinTopUp := setting.WaffoPancakeMinTopUp + originalQuotaDisplayType := operation_setting.GetGeneralSetting().QuotaDisplayType + t.Cleanup(func() { + setting.WaffoPancakeMerchantID = originalMerchantID + setting.WaffoPancakePrivateKey = originalPrivateKey + setting.WaffoPancakeProductID = originalProductID + setting.WaffoPancakeUnitPrice = originalUnitPrice + setting.WaffoPancakeMinTopUp = originalMinTopUp + operation_setting.GetGeneralSetting().QuotaDisplayType = originalQuotaDisplayType + }) + + setting.WaffoPancakeMerchantID = "MER_1234567890123456789012" + setting.WaffoPancakePrivateKey = "not-a-private-key" + setting.WaffoPancakeProductID = "PROD_1234567890123456789012" + setting.WaffoPancakeUnitPrice = 1 + setting.WaffoPancakeMinTopUp = 1 + operation_setting.GetGeneralSetting().QuotaDisplayType = operation_setting.QuotaDisplayTypeUSD + + user := &model.User{ + Id: 1201, + Username: "pancake-pending-user", + Group: "default", + Status: common.UserStatusEnabled, + } + require.NoError(t, db.Create(user).Error) + + payload, err := common.Marshal(WaffoPancakePayRequest{Amount: 10}) + require.NoError(t, err) + recorder := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(recorder) + ctx.Request = httptest.NewRequest(http.MethodPost, "/api/waffo-pancake/pay", bytes.NewReader(payload)) + ctx.Request.Header.Set("Content-Type", "application/json") + ctx.Set("id", user.Id) + + RequestWaffoPancakePay(ctx) + + require.Equal(t, http.StatusOK, recorder.Code) + require.Contains(t, recorder.Body.String(), `"data":"拉起支付失败"`) + + var topUps []model.TopUp + require.NoError(t, db.Find(&topUps).Error) + require.Len(t, topUps, 1) + require.Equal(t, common.TopUpStatusPending, topUps[0].Status) +} diff --git a/controller/user.go b/controller/user.go index 9c666962..7c40e762 100644 --- a/controller/user.go +++ b/controller/user.go @@ -388,7 +388,12 @@ func GenerateAccessToken(c *gin.Context) { return } // get rand int 28-32 - randI := common.GetRandomInt(4) + randI, err := common.SecureRandomInt(4) + if err != nil { + common.ApiErrorI18n(c, i18n.MsgGenerateFailed) + common.SysLog("failed to generate access token length: " + err.Error()) + return + } key, err := common.GenerateRandomKey(29 + randI) if err != nil { common.ApiErrorI18n(c, i18n.MsgGenerateFailed) @@ -451,7 +456,13 @@ func GetAffCode(c *gin.Context) { return } if user.AffCode == "" { - user.AffCode = common.GetRandomString(4) + affCode, err := common.GenerateRandomCharsKey(4) + if err != nil { + common.ApiErrorI18n(c, i18n.MsgGenerateFailed) + common.SysLog("failed to generate affiliate code: " + err.Error()) + return + } + user.AffCode = affCode if err := user.UpdateFields(false, model.UserUpdateFieldAffCode); err != nil { c.JSON(http.StatusOK, gin.H{ "success": false, @@ -468,6 +479,21 @@ func GetAffCode(c *gin.Context) { return } +func safeSelfUserSettingString(user *model.User) string { + if user == nil { + return "{}" + } + userSetting := user.GetSetting() + userSetting.WebhookSecret = "" + userSetting.GotifyToken = "" + settingBytes, err := common.Marshal(userSetting) + if err != nil { + common.SysLog("failed to marshal redacted user setting: " + err.Error()) + return "{}" + } + return string(settingBytes) +} + func GetSelf(c *gin.Context) { id := c.GetInt("id") userRole := c.GetInt("role") @@ -508,7 +534,7 @@ func GetSelf(c *gin.Context) { "aff_history_quota": user.AffHistoryQuota, "inviter_id": user.InviterId, "linux_do_id": user.LinuxDOId, - "setting": user.Setting, + "setting": safeSelfUserSettingString(user), "stripe_customer": user.StripeCustomer, "sidebar_modules": userSetting.SidebarModules, // 正确提取sidebar_modules字段 "permissions": permissions, // 新增权限字段 @@ -1130,13 +1156,10 @@ func ManageUser(c *gin.Context) { return } oldQuota := user.Quota - if err := model.DB.Model(&model.User{}).Where("id = ?", user.Id).Update("quota", req.Value).Error; err != nil { + if err := model.OverrideUserQuota(user.Id, req.Value); err != nil { common.ApiError(c, err) return } - if err := model.InvalidateUserCache(user.Id); err != nil { - common.SysLog(fmt.Sprintf("failed to invalidate user cache for user %d: %s", user.Id, err.Error())) - } recordManageAuditFor(c, user.Id, "user.quota_override", map[string]interface{}{ "from": logger.LogQuota(oldQuota), "to": logger.LogQuota(req.Value), @@ -1427,8 +1450,17 @@ func UpdateUserSetting(c *gin.Context) { return } if req.GotifyToken == "" { - common.ApiErrorI18n(c, i18n.MsgSettingGotifyTokenEmpty) - return + userId := c.GetInt("id") + user, err := model.GetUserById(userId, true) + if err != nil { + common.ApiError(c, err) + return + } + existingSettings := user.GetSetting() + if existingSettings.NotifyType != dto.NotifyTypeGotify || existingSettings.GotifyToken == "" { + common.ApiErrorI18n(c, i18n.MsgSettingGotifyTokenEmpty) + return + } } // 验证URL格式 if _, err := url.ParseRequestURI(req.GotifyUrl); err != nil { @@ -1474,6 +1506,8 @@ func UpdateUserSetting(c *gin.Context) { settings.WebhookUrl = req.WebhookUrl if req.WebhookSecret != "" { settings.WebhookSecret = req.WebhookSecret + } else if existingSettings.NotifyType == dto.NotifyTypeWebhook { + settings.WebhookSecret = existingSettings.WebhookSecret } } @@ -1490,7 +1524,11 @@ func UpdateUserSetting(c *gin.Context) { // 如果是Gotify类型，添加Gotify配置到设置中 if req.QuotaWarningType == dto.NotifyTypeGotify { settings.GotifyUrl = req.GotifyUrl - settings.GotifyToken = req.GotifyToken + if req.GotifyToken != "" { + settings.GotifyToken = req.GotifyToken + } else if existingSettings.NotifyType == dto.NotifyTypeGotify { + settings.GotifyToken = existingSettings.GotifyToken + } // Gotify优先级范围0-10，超出范围则使用默认值5 if req.GotifyPriority < 0 || req.GotifyPriority > 10 { settings.GotifyPriority = 5 diff --git a/controller/user_setting_test.go b/controller/user_setting_test.go index 55a8259f..2c8f0fb8 100644 --- a/controller/user_setting_test.go +++ b/controller/user_setting_test.go @@ -129,6 +129,325 @@ func TestUpdateUserSettingPreservesUnrelatedSettings(t *testing.T) { assert.Zero(t, setting.GotifyPriority) } +func TestGetSelfRedactsNotificationSecrets(t *testing.T) { + db := setupUserSettingControllerTestDB(t) + + initialSettings := dto.UserSetting{ + NotifyType: dto.NotifyTypeWebhook, + WebhookUrl: "https://example.com/webhook", + WebhookSecret: "stored-webhook-secret", + GotifyUrl: "https://gotify.example.com", + GotifyToken: "stored-gotify-token", + GotifyPriority: 7, + Language: "zh", + } + initialSettingBytes, err := common.Marshal(initialSettings) + require.NoError(t, err) + user := model.User{ + Id: 1101, + Username: "self-redaction-user", + Password: "password", + Role: common.RoleCommonUser, + Status: common.UserStatusEnabled, + Setting: string(initialSettingBytes), + } + require.NoError(t, db.Create(&user).Error) + + recorder := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(recorder) + ctx.Request = httptest.NewRequest(http.MethodGet, "/api/user/self", nil) + ctx.Set("id", user.Id) + + GetSelf(ctx) + + require.Equal(t, http.StatusOK, recorder.Code) + require.NotContains(t, recorder.Body.String(), "stored-webhook-secret") + require.NotContains(t, recorder.Body.String(), "stored-gotify-token") + + var response struct { + Data struct { + Setting string `json:"setting"` + } `json:"data"` + } + require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response)) + var returned dto.UserSetting + require.NoError(t, common.UnmarshalJsonStr(response.Data.Setting, &returned)) + assert.Equal(t, "https://example.com/webhook", returned.WebhookUrl) + assert.Empty(t, returned.WebhookSecret) + assert.Equal(t, "https://gotify.example.com", returned.GotifyUrl) + assert.Empty(t, returned.GotifyToken) + assert.Equal(t, 7, returned.GotifyPriority) +} + +func TestGetSelfIgnoresQueryUserID(t *testing.T) { + db := setupUserSettingControllerTestDB(t) + + authenticatedUser := model.User{ + Id: 1201, + Username: "self-authenticated", + Password: "password", + Role: common.RoleCommonUser, + Status: common.UserStatusEnabled, + AffCode: "self-authenticated", + } + otherUser := model.User{ + Id: 1202, + Username: "self-other", + Password: "password", + Role: common.RoleCommonUser, + Status: common.UserStatusEnabled, + AffCode: "self-other", + } + require.NoError(t, db.Create(&authenticatedUser).Error) + require.NoError(t, db.Create(&otherUser).Error) + + recorder := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(recorder) + ctx.Request = httptest.NewRequest(http.MethodGet, "/api/user/self?id=1202", nil) + ctx.Set("id", authenticatedUser.Id) + ctx.Set("role", authenticatedUser.Role) + + GetSelf(ctx) + + require.Equal(t, http.StatusOK, recorder.Code) + var response struct { + Data struct { + ID int `json:"id"` + Username string `json:"username"` + } `json:"data"` + } + require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response)) + assert.Equal(t, authenticatedUser.Id, response.Data.ID) + assert.Equal(t, authenticatedUser.Username, response.Data.Username) + assert.NotEqual(t, otherUser.Id, response.Data.ID) +} + +func TestUpdateUserSettingPreservesStoredSecretsWhenRedactedFieldsAreOmitted(t *testing.T) { + db := setupUserSettingControllerTestDB(t) + + initialSettings := dto.UserSetting{ + NotifyType: dto.NotifyTypeWebhook, + QuotaWarningThreshold: 0.5, + WebhookUrl: "https://example.com/webhook", + WebhookSecret: "stored-webhook-secret", + } + initialSettingBytes, err := common.Marshal(initialSettings) + require.NoError(t, err) + user := model.User{ + Id: 1102, + Username: "setting-secret-preserve-user", + Password: "password", + Role: common.RoleCommonUser, + Status: common.UserStatusEnabled, + Setting: string(initialSettingBytes), + } + require.NoError(t, db.Create(&user).Error) + + payload := map[string]any{ + "notify_type": dto.NotifyTypeWebhook, + "quota_warning_threshold": 0.8, + "webhook_url": "https://example.com/new-webhook", + "webhook_secret": "", + } + payloadBytes, err := common.Marshal(payload) + require.NoError(t, err) + recorder := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(recorder) + ctx.Request = httptest.NewRequest(http.MethodPut, "/api/user/setting", bytes.NewReader(payloadBytes)) + ctx.Request.Header.Set("Content-Type", "application/json") + ctx.Set("id", user.Id) + + UpdateUserSetting(ctx) + + require.Equal(t, http.StatusOK, recorder.Code) + var got model.User + require.NoError(t, db.First(&got, user.Id).Error) + setting := got.GetSetting() + assert.Equal(t, dto.NotifyTypeWebhook, setting.NotifyType) + assert.Equal(t, "https://example.com/new-webhook", setting.WebhookUrl) + assert.Equal(t, "stored-webhook-secret", setting.WebhookSecret) + assert.Equal(t, 0.8, setting.QuotaWarningThreshold) +} + +func TestUpdateUserSettingPreservesStoredGotifyTokenWhenOmitted(t *testing.T) { + db := setupUserSettingControllerTestDB(t) + + initialSettings := dto.UserSetting{ + NotifyType: dto.NotifyTypeGotify, + QuotaWarningThreshold: 0.5, + GotifyUrl: "https://gotify.example.com", + GotifyToken: "stored-gotify-token", + GotifyPriority: 3, + } + initialSettingBytes, err := common.Marshal(initialSettings) + require.NoError(t, err) + user := model.User{ + Id: 1103, + Username: "setting-gotify-preserve-user", + Password: "password", + Role: common.RoleCommonUser, + Status: common.UserStatusEnabled, + Setting: string(initialSettingBytes), + } + require.NoError(t, db.Create(&user).Error) + + payloadBytes, err := common.Marshal(map[string]any{ + "notify_type": dto.NotifyTypeGotify, + "quota_warning_threshold": 0.8, + "gotify_url": "https://gotify.example.com/new", + "gotify_token": "", + "gotify_priority": 5, + }) + require.NoError(t, err) + recorder := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(recorder) + ctx.Request = httptest.NewRequest(http.MethodPut, "/api/user/setting", bytes.NewReader(payloadBytes)) + ctx.Request.Header.Set("Content-Type", "application/json") + ctx.Set("id", user.Id) + + UpdateUserSetting(ctx) + + require.Equal(t, http.StatusOK, recorder.Code) + require.Contains(t, recorder.Body.String(), `"success":true`) + var got model.User + require.NoError(t, db.First(&got, user.Id).Error) + setting := got.GetSetting() + assert.Equal(t, dto.NotifyTypeGotify, setting.NotifyType) + assert.Equal(t, "https://gotify.example.com/new", setting.GotifyUrl) + assert.Equal(t, "stored-gotify-token", setting.GotifyToken) + assert.Equal(t, 5, setting.GotifyPriority) + assert.Equal(t, 0.8, setting.QuotaWarningThreshold) +} + +func TestUpdateUserSettingRejectsOmittedGotifyTokenFromDifferentStoredType(t *testing.T) { + db := setupUserSettingControllerTestDB(t) + + initialSettings := dto.UserSetting{ + NotifyType: dto.NotifyTypeWebhook, + QuotaWarningThreshold: 0.5, + WebhookUrl: "https://example.com/webhook", + WebhookSecret: "stored-webhook-secret", + GotifyToken: "stale-gotify-token", + } + initialSettingBytes, err := common.Marshal(initialSettings) + require.NoError(t, err) + user := model.User{ + Id: 1104, + Username: "setting-gotify-reject-user", + Password: "password", + Role: common.RoleCommonUser, + Status: common.UserStatusEnabled, + Setting: string(initialSettingBytes), + } + require.NoError(t, db.Create(&user).Error) + + payloadBytes, err := common.Marshal(map[string]any{ + "notify_type": dto.NotifyTypeGotify, + "quota_warning_threshold": 0.8, + "gotify_url": "https://gotify.example.com/new", + "gotify_token": "", + "gotify_priority": 5, + }) + require.NoError(t, err) + recorder := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(recorder) + ctx.Request = httptest.NewRequest(http.MethodPut, "/api/user/setting", bytes.NewReader(payloadBytes)) + ctx.Request.Header.Set("Content-Type", "application/json") + ctx.Set("id", user.Id) + + UpdateUserSetting(ctx) + + require.Equal(t, http.StatusOK, recorder.Code) + require.Contains(t, recorder.Body.String(), `"success":false`) + var got model.User + require.NoError(t, db.First(&got, user.Id).Error) + assert.Equal(t, initialSettings, got.GetSetting()) +} + +func TestUpdateUserSettingWebhookCredentialTransitions(t *testing.T) { + tests := []struct { + name string + initial dto.UserSetting + requestSecret string + wantWebhookSecret string + wantGotifyCleared bool + }{ + { + name: "preserves existing webhook secret when redacted value is empty", + initial: dto.UserSetting{ + NotifyType: dto.NotifyTypeWebhook, WebhookUrl: "https://example.com/old-webhook", WebhookSecret: "stored-webhook-secret", + }, + wantWebhookSecret: "stored-webhook-secret", + }, + { + name: "accepts explicit webhook secret", + initial: dto.UserSetting{ + NotifyType: dto.NotifyTypeWebhook, WebhookUrl: "https://example.com/old-webhook", WebhookSecret: "stored-webhook-secret", + }, + requestSecret: "replacement-webhook-secret", + wantWebhookSecret: "replacement-webhook-secret", + }, + { + name: "switching from gotify clears stale gotify credentials", + initial: dto.UserSetting{ + NotifyType: dto.NotifyTypeGotify, GotifyUrl: "https://gotify.example.com", GotifyToken: "stale-gotify-token", GotifyPriority: 7, + }, + requestSecret: "new-webhook-secret", + wantWebhookSecret: "new-webhook-secret", + wantGotifyCleared: true, + }, + { + name: "redacted webhook secret never reuses a gotify token", + initial: dto.UserSetting{ + NotifyType: dto.NotifyTypeGotify, GotifyUrl: "https://gotify.example.com", GotifyToken: "stale-gotify-token", GotifyPriority: 7, + }, + wantGotifyCleared: true, + }, + } + + for i, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + db := setupUserSettingControllerTestDB(t) + initial := tt.initial + initial.QuotaWarningThreshold = 0.5 + settingBytes, err := common.Marshal(initial) + require.NoError(t, err) + user := model.User{ + Id: 1200 + i, Username: fmt.Sprintf("webhook-transition-%d", i), Password: "password", + Role: common.RoleCommonUser, Status: common.UserStatusEnabled, Setting: string(settingBytes), + } + require.NoError(t, db.Create(&user).Error) + + payloadBytes, err := common.Marshal(map[string]any{ + "notify_type": dto.NotifyTypeWebhook, "quota_warning_threshold": 0.8, + "webhook_url": "https://example.com/new-webhook", "webhook_secret": tt.requestSecret, + }) + require.NoError(t, err) + recorder := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(recorder) + ctx.Request = httptest.NewRequest(http.MethodPut, "/api/user/setting", bytes.NewReader(payloadBytes)) + ctx.Request.Header.Set("Content-Type", "application/json") + ctx.Set("id", user.Id) + + UpdateUserSetting(ctx) + + require.Equal(t, http.StatusOK, recorder.Code) + require.Contains(t, recorder.Body.String(), `"success":true`) + var got model.User + require.NoError(t, db.First(&got, user.Id).Error) + setting := got.GetSetting() + assert.Equal(t, dto.NotifyTypeWebhook, setting.NotifyType) + assert.Equal(t, "https://example.com/new-webhook", setting.WebhookUrl) + assert.Equal(t, tt.wantWebhookSecret, setting.WebhookSecret) + if tt.wantGotifyCleared { + assert.Empty(t, setting.GotifyUrl) + assert.Empty(t, setting.GotifyToken) + assert.Zero(t, setting.GotifyPriority) + } + }) + } +} + func TestNormalizeNotificationEmailRejectsMalformedAddress(t *testing.T) { validEmail, ok := normalizeNotificationEmail("new@example.com") require.True(t, ok) diff --git a/controller/video_proxy.go b/controller/video_proxy.go index 48c2d6e6..4cf67708 100644 --- a/controller/video_proxy.go +++ b/controller/video_proxy.go @@ -7,6 +7,7 @@ import ( "io" "net/http" "net/url" + "strconv" "strings" "time" @@ -20,6 +21,15 @@ import ( "github.com/gin-gonic/gin" ) +var maxVideoDataURLDecodedBytes = 100 * 1024 * 1024 + +var allowedVideoDataURLMimeTypes = map[string]struct{}{ + "video/mp4": {}, + "video/webm": {}, + "video/quicktime": {}, + "video/x-matroska": {}, +} + // videoProxyError returns a standardized OpenAI-style error response. func videoProxyError(c *gin.Context, status int, errType, message string) { c.JSON(status, gin.H{ @@ -161,11 +171,7 @@ func VideoProxy(c *gin.Context) { return } - for key, values := range resp.Header { - for _, value := range values { - c.Writer.Header().Add(key, value) - } - } + copyVideoResponseHeaders(c, resp) c.Writer.Header().Set("Cache-Control", "public, max-age=86400") c.Writer.WriteHeader(resp.StatusCode) @@ -174,6 +180,16 @@ func VideoProxy(c *gin.Context) { } } +func copyVideoResponseHeaders(c *gin.Context, resp *http.Response) { + if resp == nil { + return + } + service.CopyUpstreamResponseHeaders(c, resp.Header) + if resp.ContentLength >= 0 { + c.Writer.Header().Set("Content-Length", strconv.FormatInt(resp.ContentLength, 10)) + } +} + func writeVideoDataURL(c *gin.Context, dataURL string) error { parts := strings.SplitN(dataURL, ",", 2) if len(parts) != 2 { @@ -182,14 +198,13 @@ func writeVideoDataURL(c *gin.Context, dataURL string) error { header := parts[0] payload := parts[1] - if !strings.HasPrefix(header, "data:") || !strings.Contains(header, ";base64") { - return fmt.Errorf("unsupported data url") + mimeType, err := parseVideoDataURLHeader(header) + if err != nil { + return err } - - mimeType := strings.TrimPrefix(header, "data:") - mimeType = strings.TrimSuffix(mimeType, ";base64") - if mimeType == "" { - mimeType = "video/mp4" + if len(payload) > base64.StdEncoding.EncodedLen(maxVideoDataURLDecodedBytes) || + base64.StdEncoding.DecodedLen(len(payload)) > maxVideoDataURLDecodedBytes { + return fmt.Errorf("video data url exceeds size limit") } videoBytes, err := base64.StdEncoding.DecodeString(payload) @@ -199,6 +214,9 @@ func writeVideoDataURL(c *gin.Context, dataURL string) error { return err } } + if len(videoBytes) > maxVideoDataURLDecodedBytes { + return fmt.Errorf("video data url exceeds size limit") + } c.Writer.Header().Set("Content-Type", mimeType) c.Writer.Header().Set("Cache-Control", "public, max-age=86400") @@ -206,3 +224,29 @@ func writeVideoDataURL(c *gin.Context, dataURL string) error { _, err = c.Writer.Write(videoBytes) return err } + +func parseVideoDataURLHeader(header string) (string, error) { + if !strings.HasPrefix(strings.ToLower(header), "data:") { + return "", fmt.Errorf("unsupported data url") + } + metadata := header[len("data:"):] + parts := strings.Split(metadata, ";") + mimeType := strings.ToLower(strings.TrimSpace(parts[0])) + if mimeType == "" { + mimeType = "video/mp4" + } + hasBase64 := false + for _, part := range parts[1:] { + if strings.EqualFold(strings.TrimSpace(part), "base64") { + hasBase64 = true + break + } + } + if !hasBase64 { + return "", fmt.Errorf("unsupported data url") + } + if _, ok := allowedVideoDataURLMimeTypes[mimeType]; !ok { + return "", fmt.Errorf("unsupported video data url mime type") + } + return mimeType, nil +} diff --git a/controller/video_proxy_test.go b/controller/video_proxy_test.go new file mode 100644 index 00000000..93ad3aae --- /dev/null +++ b/controller/video_proxy_test.go @@ -0,0 +1,77 @@ +package controller + +import ( + "encoding/base64" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +func TestWriteVideoDataURLAllowsVideoPayload(t *testing.T) { + gin.SetMode(gin.TestMode) + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + + err := writeVideoDataURL(c, "data:video/mp4;base64,"+base64.StdEncoding.EncodeToString([]byte("mp4"))) + + require.NoError(t, err) + require.Equal(t, http.StatusOK, recorder.Code) + require.Equal(t, "video/mp4", recorder.Header().Get("Content-Type")) + require.Equal(t, "mp4", recorder.Body.String()) +} + +func TestWriteVideoDataURLRejectsNonVideoPayload(t *testing.T) { + gin.SetMode(gin.TestMode) + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + + err := writeVideoDataURL(c, "data:text/html;base64,"+base64.StdEncoding.EncodeToString([]byte(""))) + + require.Error(t, err) + require.Contains(t, err.Error(), "unsupported video data url mime type") + require.Equal(t, http.StatusOK, recorder.Code) + require.Empty(t, recorder.Body.String()) +} + +func TestWriteVideoDataURLRejectsOversizedPayload(t *testing.T) { + gin.SetMode(gin.TestMode) + originalLimit := maxVideoDataURLDecodedBytes + maxVideoDataURLDecodedBytes = 4 + t.Cleanup(func() { + maxVideoDataURLDecodedBytes = originalLimit + }) + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + + payload := base64.StdEncoding.EncodeToString([]byte(strings.Repeat("x", 5))) + err := writeVideoDataURL(c, "data:video/mp4;base64,"+payload) + + require.Error(t, err) + require.Contains(t, err.Error(), "exceeds size limit") + require.Equal(t, http.StatusOK, recorder.Code) + require.Empty(t, recorder.Body.String()) +} + +func TestCopyVideoResponseHeadersPreservesKnownContentLength(t *testing.T) { + gin.SetMode(gin.TestMode) + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + resp := &http.Response{ + ContentLength: 1234, + Header: http.Header{ + "Content-Type": []string{"video/mp4"}, + "Content-Length": []string{"9999"}, + "Connection": []string{"keep-alive"}, + }, + } + + copyVideoResponseHeaders(c, resp) + + require.Equal(t, "video/mp4", recorder.Header().Get("Content-Type")) + require.Equal(t, "1234", recorder.Header().Get("Content-Length")) + require.Empty(t, recorder.Header().Get("Connection")) +} diff --git a/docs/openapi/api.json b/docs/openapi/api.json index 4527a57e..8505cd47 100644 --- a/docs/openapi/api.json +++ b/docs/openapi/api.json @@ -867,14 +867,44 @@ } }, "/api/oauth/state": { - "get": { - "summary": "生成OAuth State", + "post": { + "summary": "创建一次性 OAuth State", "deprecated": false, - "description": "🔓 无需鉴权", + "description": "登录无需鉴权；绑定流程需要有效用户会话", "tags": [ "OAuth" ], "parameters": [], + "requestBody": { + "required": true, + "content": { + "application/json": { + "schema": { + "type": "object", + "required": [ + "provider", + "intent" + ], + "properties": { + "provider": { + "type": "string" + }, + "intent": { + "type": "string", + "enum": [ + "login", + "bind" + ] + }, + "aff": { + "type": "string", + "maxLength": 32 + } + } + } + } + } + }, "responses": { "200": { "description": "成功", diff --git a/dto/openai_request.go b/dto/openai_request.go index 2e09795e..8220ba7d 100644 --- a/dto/openai_request.go +++ b/dto/openai_request.go @@ -490,14 +490,14 @@ func (m *Message) ParseToolCalls() []ToolCallRequest { return nil } var toolCalls []ToolCallRequest - if err := json.Unmarshal(m.ToolCalls, &toolCalls); err == nil { + if err := common.Unmarshal(m.ToolCalls, &toolCalls); err == nil { return toolCalls } return toolCalls } func (m *Message) SetToolCalls(toolCalls any) { - toolCallsJson, _ := json.Marshal(toolCalls) + toolCallsJson, _ := common.Marshal(toolCalls) m.ToolCalls = toolCallsJson } @@ -687,7 +687,7 @@ func (m *Message) ParseContent() []MediaContent { } var stringContent string - if err := json.Unmarshal(m.Content, &stringContent); err == nil { + if err := common.Unmarshal(m.Content, &stringContent); err == nil { m.parsedStringContent = &stringContent return stringContent } @@ -712,14 +712,14 @@ func (m *Message) SetNullContent() { } func (m *Message) SetStringContent(content string) { - jsonContent, _ := json.Marshal(content) + jsonContent, _ := common.Marshal(content) m.Content = jsonContent m.parsedStringContent = &content m.parsedContent = nil } func (m *Message) SetMediaContent(content []MediaContent) { - jsonContent, _ := json.Marshal(content) + jsonContent, _ := common.Marshal(content) m.Content = jsonContent m.parsedContent = nil m.parsedStringContent = nil @@ -730,7 +730,7 @@ func (m *Message) IsStringContent() bool { return true } var stringContent string - if err := json.Unmarshal(m.Content, &stringContent); err == nil { + if err := common.Unmarshal(m.Content, &stringContent); err == nil { m.parsedStringContent = &stringContent return true } @@ -746,7 +746,7 @@ func (m *Message) ParseContent() []MediaContent { // 先尝试解析为字符串 var stringContent string - if err := json.Unmarshal(m.Content, &stringContent); err == nil { + if err := common.Unmarshal(m.Content, &stringContent); err == nil { contentList = []MediaContent{{ Type: ContentTypeText, Text: stringContent, @@ -757,7 +757,7 @@ func (m *Message) ParseContent() []MediaContent { // 尝试解析为数组 var arrayContent []map[string]interface{} - if err := json.Unmarshal(m.Content, &arrayContent); err == nil { + if err := common.Unmarshal(m.Content, &arrayContent); err == nil { for _, contentItem := range arrayContent { contentType, ok := contentItem["type"].(string) if !ok { diff --git a/dto/values.go b/dto/values.go index 5d748264..c5ae374f 100644 --- a/dto/values.go +++ b/dto/values.go @@ -3,40 +3,42 @@ package dto import ( "encoding/json" "strconv" + + "github.com/MAX-API-Next/MAX-API/common" ) type StringValue string func (s *StringValue) UnmarshalJSON(data []byte) error { var str string - if err := json.Unmarshal(data, &str); err == nil { + if err := common.Unmarshal(data, &str); err == nil { *s = StringValue(str) return nil } var raw json.Number - if err := json.Unmarshal(data, &raw); err == nil { + if err := common.Unmarshal(data, &raw); err == nil { *s = StringValue(raw.String()) return nil } - return json.Unmarshal(data, &str) + return common.Unmarshal(data, &str) } func (s StringValue) MarshalJSON() ([]byte, error) { - return json.Marshal(string(s)) + return common.Marshal(string(s)) } type IntValue int func (i *IntValue) UnmarshalJSON(b []byte) error { var n int - if err := json.Unmarshal(b, &n); err == nil { + if err := common.Unmarshal(b, &n); err == nil { *i = IntValue(n) return nil } var s string - if err := json.Unmarshal(b, &s); err != nil { + if err := common.Unmarshal(b, &s); err != nil { return err } v, err := strconv.Atoi(s) @@ -48,19 +50,19 @@ func (i *IntValue) UnmarshalJSON(b []byte) error { } func (i IntValue) MarshalJSON() ([]byte, error) { - return json.Marshal(int(i)) + return common.Marshal(int(i)) } type BoolValue bool func (b *BoolValue) UnmarshalJSON(data []byte) error { var boolean bool - if err := json.Unmarshal(data, &boolean); err == nil { + if err := common.Unmarshal(data, &boolean); err == nil { *b = BoolValue(boolean) return nil } var str string - if err := json.Unmarshal(data, &str); err != nil { + if err := common.Unmarshal(data, &str); err != nil { return err } if str == "true" { @@ -68,10 +70,10 @@ func (b *BoolValue) UnmarshalJSON(data []byte) error { } else if str == "false" { *b = BoolValue(false) } else { - return json.Unmarshal(data, &boolean) + return common.Unmarshal(data, &boolean) } return nil } func (b BoolValue) MarshalJSON() ([]byte, error) { - return json.Marshal(bool(b)) + return common.Marshal(bool(b)) } diff --git a/main.go b/main.go index 53036daf..e8ee0e2a 100644 --- a/main.go +++ b/main.go @@ -70,6 +70,8 @@ func main() { common.FatalLog("failed to close database: " + err.Error()) } }() + backgroundCtx, stopBackground := context.WithCancel(context.Background()) + defer stopBackground() // Recover durable balance settlements after process restarts or transient DB failures. model.StartBillingSettlementTaskRunner() @@ -128,6 +130,7 @@ func main() { // Persistent system maintenance task runner service.StartSystemTaskRunner() + authFlowCleanupDone := service.StartAuthFlowCleanupWithContext(backgroundCtx) // Wire task polling adaptor factory (breaks service -> relay import cycle) service.GetTaskAdaptorFunc = func(platform constant.TaskPlatform) service.TaskPollingAdaptor { @@ -155,11 +158,17 @@ func main() { model.InitBatchUpdater() } + var cpuMonitorDone <-chan struct{} if os.Getenv("ENABLE_PPROF") == "true" { gopool.Go(func() { log.Println(http.ListenAndServe("0.0.0.0:8005", nil)) }) - go common.Monitor() + done := make(chan struct{}) + cpuMonitorDone = done + go func() { + defer close(done) + common.MonitorContext(backgroundCtx) + }() common.SysLog("pprof enabled") } @@ -193,13 +202,7 @@ func main() { common.FatalLog(err.Error()) } store := cookie.NewStore([]byte(common.SessionSecret)) - store.Options(sessions.Options{ - Path: "/", - MaxAge: 2592000, // 30 days - HttpOnly: true, - Secure: common.SessionCookieSecure, - SameSite: http.SameSiteStrictMode, - }) + store.Options(sessionCookieOptions()) server.Use(sessions.Sessions("session", store)) InjectUmamiAnalytics() @@ -241,17 +244,24 @@ func main() { shutdownTimeout := time.Duration(common.GetEnvOrDefault("SHUTDOWN_TIMEOUT_SECONDS", 120)) * time.Second ctx, cancel := context.WithTimeout(context.Background(), shutdownTimeout) defer cancel() + stopBackground() shutdownHTTPServer(ctx, srv) - stopBackgroundRunnerCtx, stopBackgroundRunnerCancel := context.WithTimeout(context.Background(), shutdownTimeout) - if err := model.StopBillingSettlementTaskRunner(stopBackgroundRunnerCtx); err != nil { + waitForBackgroundRunner(ctx, "auth flow cleanup", authFlowCleanupDone) + waitForBackgroundRunner(ctx, "CPU profile monitor", cpuMonitorDone) + if err := common.StopSystemMonitor(ctx); err != nil { + common.SysError(fmt.Sprintf("timed out stopping system monitor: %v", err)) + } + if err := middleware.StopInMemoryRateLimiter(ctx); err != nil { + common.SysError(fmt.Sprintf("timed out stopping in-memory rate limiter: %v", err)) + } + if err := model.StopBillingSettlementTaskRunner(ctx); err != nil { common.SysError(fmt.Sprintf("timed out stopping billing settlement runner: %v", err)) } if common.RedisEnabled { - if err := model.StopCacheInvalidationTaskRunner(stopBackgroundRunnerCtx); err != nil { + if err := model.StopCacheInvalidationTaskRunner(ctx); err != nil { common.SysError(fmt.Sprintf("timed out stopping cache invalidation runner: %v", err)) } } - stopBackgroundRunnerCancel() if common.DataExportEnabled { saveTimeout := time.Duration(common.GetEnvOrDefault("QUOTA_DATA_CACHE_SAVE_TIMEOUT_SECONDS", 30)) * time.Second if !runWithTimeout(saveTimeout, model.WaitPendingLogQuotaData) { @@ -271,6 +281,17 @@ func main() { common.SysLog("server exited") } +func sessionCookieOptions() sessions.Options { + return sessions.Options{ + Path: "/", + MaxAge: 2592000, // 30 days + HttpOnly: true, + Secure: common.SessionCookieSecure, + // OAuth providers return through a cross-site top-level GET redirect. + SameSite: http.SameSiteLaxMode, + } +} + func configureSessionCookieSecure(serverAddress string) error { value, explicitlyConfigured := os.LookupEnv("SESSION_COOKIE_SECURE") value = strings.TrimSpace(value) @@ -299,6 +320,22 @@ func shutdownHTTPServer(ctx context.Context, srv *http.Server) { } } +func waitForBackgroundRunner(ctx context.Context, name string, done <-chan struct{}) { + if done == nil { + return + } + select { + case <-done: + return + default: + } + select { + case <-done: + case <-ctx.Done(): + common.SysError(fmt.Sprintf("timed out stopping %s: %v", name, ctx.Err())) + } +} + func runWithTimeout(timeout time.Duration, fn func()) bool { done := make(chan struct{}) go func() { diff --git a/main_test.go b/main_test.go index 8a7753fc..35ee05d9 100644 --- a/main_test.go +++ b/main_test.go @@ -26,7 +26,7 @@ func TestRunWithTimeoutTimesOut(t *testing.T) { } func TestRunWithTimeoutRecoversPanic(t *testing.T) { - assert.True(t, runWithTimeout(time.Second, func() { + assert.True(t, runWithTimeout(5*time.Second, func() { panic("quota save failed") })) } @@ -127,3 +127,11 @@ func TestConfigureSessionCookieSecureRejectsInvalidValue(t *testing.T) { require.Error(t, configureSessionCookieSecure("https://example.com")) } + +func TestSessionCookieOptionsAllowOAuthTopLevelCallback(t *testing.T) { + options := sessionCookieOptions() + + assert.Equal(t, http.SameSiteLaxMode, options.SameSite) + assert.True(t, options.HttpOnly) + assert.Equal(t, "/", options.Path) +} diff --git a/middleware/audit.go b/middleware/audit.go index 010913b8..0448202b 100644 --- a/middleware/audit.go +++ b/middleware/audit.go @@ -7,7 +7,6 @@ import ( "github.com/MAX-API-Next/MAX-API/constant" "github.com/MAX-API-Next/MAX-API/model" - "github.com/bytedance/gopkg/util/gopool" "github.com/gin-gonic/gin" ) @@ -148,9 +147,7 @@ func finishAdminAudit(c *gin.Context, writer *auditResponseWriter) { auditInfo["params"] = routeParams } - gopool.Go(func() { - model.RecordOperationAuditLog(operatorID, content, ip, action, opParams, adminInfo, auditInfo) - }) + model.RecordOperationAuditLog(operatorID, content, ip, action, opParams, adminInfo, auditInfo) } func auditAuthMethod(c *gin.Context) string { diff --git a/middleware/auth.go b/middleware/auth.go index 1d7ca1a6..864dca40 100644 --- a/middleware/auth.go +++ b/middleware/auth.go @@ -192,7 +192,7 @@ func freshAuthHelper(c *gin.Context, minRole int) { return } } - if status == common.UserStatusDisabled { + if status != common.UserStatusEnabled { c.JSON(http.StatusOK, gin.H{ "success": false, "message": common.TranslateMessage(c, i18n.MsgAuthUserBanned), @@ -334,6 +334,14 @@ func TokenAuthReadOnly() func(c *gin.Context) { userCache, err := model.GetUserCache(token.UserId) if err != nil { + if errors.Is(err, gorm.ErrRecordNotFound) { + c.JSON(http.StatusForbidden, gin.H{ + "success": false, + "message": common.TranslateMessage(c, i18n.MsgAuthUserBanned), + }) + c.Abort() + return + } common.SysLog(fmt.Sprintf("TokenAuthReadOnly GetUserCache error for user %d: %v", token.UserId, err)) c.JSON(http.StatusInternalServerError, gin.H{ "success": false, @@ -451,6 +459,10 @@ func TokenAuth() func(c *gin.Context) { userCache, err := model.GetUserCache(token.UserId) if err != nil { + if errors.Is(err, gorm.ErrRecordNotFound) { + abortWithOpenAiMessage(c, http.StatusForbidden, common.TranslateMessage(c, i18n.MsgAuthUserBanned)) + return + } common.SysLog(fmt.Sprintf("TokenAuth GetUserCache error for user %d: %v", token.UserId, err)) abortWithOpenAiMessage(c, http.StatusInternalServerError, common.TranslateMessage(c, i18n.MsgDatabaseError)) diff --git a/middleware/auth_session_test.go b/middleware/auth_session_test.go index 35c56b5d..434b4f36 100644 --- a/middleware/auth_session_test.go +++ b/middleware/auth_session_test.go @@ -7,6 +7,7 @@ import ( "testing" "github.com/MAX-API-Next/MAX-API/common" + appi18n "github.com/MAX-API-Next/MAX-API/i18n" "github.com/MAX-API-Next/MAX-API/model" "github.com/gin-contrib/sessions" "github.com/gin-contrib/sessions/cookie" @@ -29,7 +30,7 @@ func TestMain(m *testing.M) { } model.DB = db model.LOG_DB = db - if err := db.AutoMigrate(&model.User{}); err != nil { + if err := db.AutoMigrate(&model.User{}, &model.Token{}); err != nil { panic("failed to migrate middleware test db: " + err.Error()) } @@ -39,6 +40,7 @@ func TestMain(m *testing.M) { func seedMiddlewareUser(t *testing.T, id int, role int, status int, group string) { t.Helper() + require.NoError(t, model.DB.Unscoped().Where("user_id = ?", id).Delete(&model.Token{}).Error) require.NoError(t, model.DB.Unscoped().Where("id = ?", id).Delete(&model.User{}).Error) user := &model.User{ Id: id, @@ -50,6 +52,26 @@ func seedMiddlewareUser(t *testing.T, id int, role int, status int, group string Group: group, } require.NoError(t, model.DB.Create(user).Error) + require.NoError(t, model.DB.Model(&model.User{}).Where("id = ?", id).UpdateColumn("status", status).Error) +} + +func seedMiddlewareToken(t *testing.T, userID int, key string) *model.Token { + t.Helper() + + token := &model.Token{ + UserId: userID, + Name: "test-token", + Key: key, + Status: common.TokenStatusEnabled, + CreatedTime: 1, + AccessedTime: 1, + ExpiredTime: -1, + RemainQuota: 100, + UnlimitedQuota: true, + Group: "default", + } + require.NoError(t, model.DB.Create(token).Error) + return token } func loginMiddlewareSession(t *testing.T, router *gin.Engine, user *model.User) []*http.Cookie { @@ -156,6 +178,39 @@ func TestUserAuthAllowsSessionWithStaleUserHeader(t *testing.T) { require.Equal(t, http.StatusNoContent, recorder.Code) } +func TestUserAuthRejectsNonEnabledStatus(t *testing.T) { + seedMiddlewareUser(t, 1, common.RoleCommonUser, 0, "default") + + router := gin.New() + router.Use(sessions.Sessions("session", cookie.NewStore([]byte("user-session-non-enabled-test")))) + cookies := loginMiddlewareSession(t, router, &model.User{ + Id: 1, + Username: "tester", + Role: common.RoleCommonUser, + Status: common.UserStatusEnabled, + Group: "default", + }) + router.GET("/self", UserAuth(), func(c *gin.Context) { + c.Status(http.StatusNoContent) + }) + + recorder := httptest.NewRecorder() + request := httptest.NewRequest(http.MethodGet, "/self", nil) + for _, cookie := range cookies { + request.AddCookie(cookie) + } + router.ServeHTTP(recorder, request) + + require.Equal(t, http.StatusOK, recorder.Code) + var response struct { + Success bool `json:"success"` + Message string `json:"message"` + } + require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response)) + require.False(t, response.Success) + require.Equal(t, appi18n.MsgAuthUserBanned, response.Message) +} + func TestTokenOrUserAuthRejectsDisabledStaleSession(t *testing.T) { seedMiddlewareUser(t, 1, common.RoleCommonUser, common.UserStatusEnabled, "default") @@ -183,3 +238,38 @@ func TestTokenOrUserAuthRejectsDisabledStaleSession(t *testing.T) { require.Equal(t, http.StatusUnauthorized, recorder.Code) } + +func TestTokenAuthRejectsDisabledTokenOwner(t *testing.T) { + seedMiddlewareUser(t, 1, common.RoleCommonUser, common.UserStatusDisabled, "default") + token := seedMiddlewareToken(t, 1, "disabledownertoken") + + router := gin.New() + router.GET("/v1/chat/completions", TokenAuth(), func(c *gin.Context) { + c.Status(http.StatusNoContent) + }) + + recorder := httptest.NewRecorder() + request := httptest.NewRequest(http.MethodGet, "/v1/chat/completions", nil) + request.Header.Set("Authorization", "Bearer sk-"+token.Key) + router.ServeHTTP(recorder, request) + + require.Equal(t, http.StatusForbidden, recorder.Code) +} + +func TestTokenAuthRejectsDeletedTokenOwner(t *testing.T) { + seedMiddlewareUser(t, 1, common.RoleCommonUser, common.UserStatusEnabled, "default") + token := seedMiddlewareToken(t, 1, "deletedownertoken") + require.NoError(t, model.DB.Delete(&model.User{}, 1).Error) + + router := gin.New() + router.GET("/v1/chat/completions", TokenAuth(), func(c *gin.Context) { + c.Status(http.StatusNoContent) + }) + + recorder := httptest.NewRecorder() + request := httptest.NewRequest(http.MethodGet, "/v1/chat/completions", nil) + request.Header.Set("Authorization", "Bearer sk-"+token.Key) + router.ServeHTTP(recorder, request) + + require.Equal(t, http.StatusForbidden, recorder.Code) +} diff --git a/middleware/email-verification-rate-limit.go b/middleware/email-verification-rate-limit.go index 0b68f217..4ae58505 100644 --- a/middleware/email-verification-rate-limit.go +++ b/middleware/email-verification-rate-limit.go @@ -55,6 +55,7 @@ func redisEmailVerificationRateLimiter(c *gin.Context) { } func memoryEmailVerificationRateLimiter(c *gin.Context) { + inMemoryRateLimiter.Init(common.RateLimitKeyExpirationDuration) key := EmailVerificationRateLimitMark + ":" + c.ClientIP() if !inMemoryRateLimiter.Request(key, EmailVerificationMaxRequests, EmailVerificationDuration) { @@ -74,7 +75,6 @@ func EmailVerificationRateLimit() gin.HandlerFunc { if common.RedisEnabled { redisEmailVerificationRateLimiter(c) } else { - inMemoryRateLimiter.Init(common.RateLimitKeyExpirationDuration) memoryEmailVerificationRateLimiter(c) } } diff --git a/middleware/model-rate-limit.go b/middleware/model-rate-limit.go index 0e97c903..2c512489 100644 --- a/middleware/model-rate-limit.go +++ b/middleware/model-rate-limit.go @@ -130,7 +130,7 @@ func redisRateLimitHandler(duration int64, totalMaxCount, successMaxCount int) g // 内存限流处理器 func memoryRateLimitHandler(duration int64, totalMaxCount, successMaxCount int) gin.HandlerFunc { - inMemoryRateLimiter.Init(time.Duration(setting.ModelRequestRateLimitDurationMinutes) * time.Minute) + inMemoryRateLimiter.Init(common.RateLimitKeyExpirationDuration) return func(c *gin.Context) { userId := strconv.Itoa(c.GetInt("id")) diff --git a/middleware/rate-limit.go b/middleware/rate-limit.go index 6250845b..23c26d5a 100644 --- a/middleware/rate-limit.go +++ b/middleware/rate-limit.go @@ -17,6 +17,10 @@ import ( var inMemoryRateLimiter common.InMemoryRateLimiter +func StopInMemoryRateLimiter(ctx context.Context) error { + return inMemoryRateLimiter.Stop(ctx) +} + var timeFormat = "2006-01-02T15:04:05.000Z" const rateLimitRedisTimeout = 500 * time.Millisecond diff --git a/middleware/turnstile-check.go b/middleware/turnstile-check.go index 8fbb490a..99930582 100644 --- a/middleware/turnstile-check.go +++ b/middleware/turnstile-check.go @@ -1,9 +1,12 @@ package middleware import ( + "context" + "fmt" "net/http" "net/url" "strings" + "time" "github.com/MAX-API-Next/MAX-API/common" "github.com/gin-contrib/sessions" @@ -16,6 +19,10 @@ type turnstileCheckResponse struct { const turnstileTokenHeader = "X-Turnstile-Token" +var turnstileVerificationTimeout = 10 * time.Second +var turnstileVerificationEndpoint = "https://challenges.cloudflare.com/turnstile/v0/siteverify" +var turnstileHTTPClient = &http.Client{Timeout: turnstileVerificationTimeout} + func getTurnstileToken(c *gin.Context) string { if token := strings.TrimSpace(c.GetHeader(turnstileTokenHeader)); token != "" { return token @@ -23,6 +30,36 @@ func getTurnstileToken(c *gin.Context) string { return strings.TrimSpace(c.Query("turnstile")) } +func verifyTurnstile(ctx context.Context, response string, remoteIP string) (bool, error) { + ctx, cancel := context.WithTimeout(ctx, turnstileVerificationTimeout) + defer cancel() + form := url.Values{ + "secret": {common.TurnstileSecretKey}, + "response": {response}, + "remoteip": {remoteIP}, + } + req, err := http.NewRequestWithContext(ctx, http.MethodPost, turnstileVerificationEndpoint, strings.NewReader(form.Encode())) + if err != nil { + return false, err + } + req.Header.Set("Content-Type", "application/x-www-form-urlencoded") + + rawRes, err := turnstileHTTPClient.Do(req) + if err != nil { + return false, err + } + defer rawRes.Body.Close() + if rawRes.StatusCode < http.StatusOK || rawRes.StatusCode >= http.StatusMultipleChoices { + return false, fmt.Errorf("turnstile verification returned HTTP %d", rawRes.StatusCode) + } + + var res turnstileCheckResponse + if err := common.DecodeJson(rawRes.Body, &res); err != nil { + return false, err + } + return res.Success, nil +} + func TurnstileCheck() gin.HandlerFunc { return func(c *gin.Context) { if common.TurnstileCheckEnabled { @@ -41,23 +78,7 @@ func TurnstileCheck() gin.HandlerFunc { c.Abort() return } - rawRes, err := http.PostForm("https://challenges.cloudflare.com/turnstile/v0/siteverify", url.Values{ - "secret": {common.TurnstileSecretKey}, - "response": {response}, - "remoteip": {c.ClientIP()}, - }) - if err != nil { - common.SysLog(err.Error()) - c.JSON(http.StatusOK, gin.H{ - "success": false, - "message": err.Error(), - }) - c.Abort() - return - } - defer rawRes.Body.Close() - var res turnstileCheckResponse - err = common.DecodeJson(rawRes.Body, &res) + verified, err := verifyTurnstile(c.Request.Context(), response, c.ClientIP()) if err != nil { common.SysLog(err.Error()) c.JSON(http.StatusOK, gin.H{ @@ -67,7 +88,7 @@ func TurnstileCheck() gin.HandlerFunc { c.Abort() return } - if !res.Success { + if !verified { c.JSON(http.StatusOK, gin.H{ "success": false, "message": "Turnstile 校验失败，请刷新重试！", diff --git a/middleware/turnstile_check_test.go b/middleware/turnstile_check_test.go index 34b2b811..4e584d1a 100644 --- a/middleware/turnstile_check_test.go +++ b/middleware/turnstile_check_test.go @@ -1,13 +1,23 @@ package middleware import ( + "context" + "net/http" "net/http/httptest" "testing" + "time" "github.com/gin-gonic/gin" "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" ) +type roundTripFunc func(*http.Request) (*http.Response, error) + +func (f roundTripFunc) RoundTrip(request *http.Request) (*http.Response, error) { + return f(request) +} + func TestGetTurnstileTokenPrefersHeaderWithQueryFallback(t *testing.T) { gin.SetMode(gin.TestMode) @@ -34,3 +44,24 @@ func TestGetTurnstileTokenPrefersHeaderWithQueryFallback(t *testing.T) { }) } } + +func TestVerifyTurnstileUsesBoundedClientTimeout(t *testing.T) { + oldClient := turnstileHTTPClient + oldTimeout := turnstileVerificationTimeout + turnstileVerificationTimeout = 50 * time.Millisecond + turnstileHTTPClient = &http.Client{ + Transport: roundTripFunc(func(request *http.Request) (*http.Response, error) { + <-request.Context().Done() + return nil, request.Context().Err() + }), + } + t.Cleanup(func() { + turnstileHTTPClient = oldClient + turnstileVerificationTimeout = oldTimeout + }) + + started := time.Now() + _, err := verifyTurnstile(context.Background(), "token", "127.0.0.1") + require.Error(t, err) + require.Less(t, time.Since(started), time.Second) +} diff --git a/model/ability.go b/model/ability.go index 18b20818..1157f7a6 100644 --- a/model/ability.go +++ b/model/ability.go @@ -94,7 +94,10 @@ func GetChannelExcluding(group string, model string, retry int, requestPath stri return nil, nil } - channelId, ok := selectChannelIdFromAbilities(abilities, retry, excludedChannelIDs) + channelId, ok, err := selectChannelIdFromAbilities(abilities, retry, excludedChannelIDs) + if err != nil { + return nil, err + } if !ok { return nil, nil } @@ -166,9 +169,9 @@ func abilityCol(name string) string { return "`abilities`.`" + name + "`" } -func selectChannelIdFromAbilities(abilities []Ability, retry int, excludedChannelIDs map[int]struct{}) (int, bool) { +func selectChannelIdFromAbilities(abilities []Ability, retry int, excludedChannelIDs map[int]struct{}) (int, bool, error) { if len(abilities) == 0 { - return 0, false + return 0, false, nil } priorities := make([]int64, 0, len(abilities)) seenPriorities := make(map[int64]struct{}, len(abilities)) @@ -187,7 +190,7 @@ func selectChannelIdFromAbilities(abilities []Ability, retry int, excludedChanne return priorities[i] > priorities[j] }) if len(priorities) == 0 { - return 0, false + return 0, false, nil } if retry >= len(priorities) { retry = len(priorities) - 1 @@ -201,21 +204,25 @@ func selectChannelIdFromAbilities(abilities []Ability, retry int, excludedChanne } } if len(targetAbilities) == 0 { - return 0, false + return 0, false, nil } weightSum := uint(0) for _, ability_ := range targetAbilities { weightSum += ability_.Weight + 10 } - weight := common.GetRandomInt(int(weightSum)) + weight, err := common.SecureRandomInt(int(weightSum)) + if err != nil { + return 0, false, fmt.Errorf("select channel by weight: %w", err) + } for _, ability_ := range targetAbilities { - weight -= int(ability_.Weight) + 10 - if weight <= 0 { - return ability_.ChannelId, true + candidateWeight := int(ability_.Weight) + 10 + if weight < candidateWeight { + return ability_.ChannelId, true, nil } + weight -= candidateWeight } - return 0, false + return 0, false, nil } func abilityPriority(ability Ability) int64 { diff --git a/model/auth_flow.go b/model/auth_flow.go new file mode 100644 index 00000000..6d458857 --- /dev/null +++ b/model/auth_flow.go @@ -0,0 +1,154 @@ +package model + +import ( + "crypto/rand" + "encoding/base64" + "errors" + "fmt" + "io" + "strings" + "time" + + "github.com/MAX-API-Next/MAX-API/common" + "gorm.io/gorm" +) + +const ( + AuthFlowPurposeOAuth = "oauth" + AuthFlowIntentLogin = "login" + AuthFlowIntentBind = "bind" + AuthFlowTokenBytes = 32 + AuthFlowDefaultCleanupRetention = 24 * time.Hour +) + +var ( + ErrAuthFlowInvalid = errors.New("auth flow is invalid") + ErrAuthFlowExpired = errors.New("auth flow has expired") + ErrAuthFlowConsumed = errors.New("auth flow has already been consumed") +) + +// AuthFlow stores one-time authentication state. The opaque token is never persisted. +type AuthFlow struct { + Id int64 `json:"id" gorm:"primaryKey"` + TokenHash string `json:"-" gorm:"type:char(64);not null;uniqueIndex"` + Purpose string `json:"purpose" gorm:"type:varchar(32);not null;index:idx_auth_flow_purpose_expiry"` + Provider string `json:"provider,omitempty" gorm:"type:varchar(64)"` + Intent string `json:"intent,omitempty" gorm:"type:varchar(16)"` + UserId int `json:"user_id,omitempty" gorm:"index"` + Payload string `json:"-" gorm:"type:text"` + CreatedAt time.Time `json:"created_at"` + ExpiresAt time.Time `json:"expires_at" gorm:"not null;index:idx_auth_flow_purpose_expiry"` + ConsumedAt *time.Time `json:"consumed_at,omitempty" gorm:"index"` +} + +func (AuthFlow) TableName() string { return "auth_flows" } + +type AuthFlowCreate struct { + Purpose string + Provider string + Intent string + UserId int + Payload string + ExpiresAt time.Time +} + +type AuthFlowMatch struct { + Purpose string + Provider string + Intent string + UserId int +} + +func authFlowTokenHash(token string) string { + return common.GenerateHMACWithKey([]byte("auth-flow-v1:"+common.SessionSecret), token) +} + +func applyAuthFlowMatch(query *gorm.DB, token string, match AuthFlowMatch) *gorm.DB { + query = query.Where("token_hash = ? AND purpose = ?", authFlowTokenHash(token), match.Purpose) + if match.Provider != "" { + query = query.Where("provider = ?", match.Provider) + } + if match.Intent != "" { + query = query.Where("intent = ?", match.Intent) + } + if match.UserId != 0 { + query = query.Where("user_id = ?", match.UserId) + } + return query +} + +func generateAuthFlowToken(randomReader io.Reader) (string, error) { + random := make([]byte, AuthFlowTokenBytes) + if _, err := io.ReadFull(randomReader, random); err != nil { + return "", fmt.Errorf("generate auth flow token: %w", err) + } + return base64.RawURLEncoding.EncodeToString(random), nil +} + +func CreateAuthFlow(input AuthFlowCreate) (string, *AuthFlow, error) { + if strings.TrimSpace(input.Purpose) == "" || input.ExpiresAt.IsZero() || !input.ExpiresAt.After(time.Now()) { + return "", nil, ErrAuthFlowInvalid + } + token, err := generateAuthFlowToken(rand.Reader) + if err != nil { + return "", nil, err + } + flow := &AuthFlow{ + TokenHash: authFlowTokenHash(token), + Purpose: input.Purpose, + Provider: input.Provider, + Intent: input.Intent, + UserId: input.UserId, + Payload: input.Payload, + ExpiresAt: input.ExpiresAt, + } + if err := DB.Create(flow).Error; err != nil { + return "", nil, err + } + return token, flow, nil +} + +func GetAuthFlow(token string, match AuthFlowMatch) (*AuthFlow, error) { + if token == "" || match.Purpose == "" { + return nil, ErrAuthFlowInvalid + } + var flow AuthFlow + if err := applyAuthFlowMatch(DB, token, match).First(&flow).Error; err != nil { + if errors.Is(err, gorm.ErrRecordNotFound) { + return nil, ErrAuthFlowInvalid + } + return nil, err + } + if flow.ConsumedAt != nil { + return nil, ErrAuthFlowConsumed + } + if !flow.ExpiresAt.After(time.Now()) { + return nil, ErrAuthFlowExpired + } + return &flow, nil +} + +func ConsumeAuthFlow(token string, match AuthFlowMatch) (*AuthFlow, error) { + flow, err := GetAuthFlow(token, match) + if err != nil { + return nil, err + } + now := time.Now() + result := DB.Model(&AuthFlow{}). + Where("id = ? AND consumed_at IS NULL AND expires_at > ?", flow.Id, now). + Update("consumed_at", now) + if result.Error != nil { + return nil, result.Error + } + if result.RowsAffected != 1 { + return nil, ErrAuthFlowConsumed + } + flow.ConsumedAt = &now + return flow, nil +} + +func DeleteExpiredAuthFlows(now time.Time) error { + cutoff := now.Add(-AuthFlowDefaultCleanupRetention) + return DB.Where("expires_at < ? OR (consumed_at IS NOT NULL AND consumed_at < ?)", cutoff, cutoff). + Delete(&AuthFlow{}).Error +} diff --git a/model/auth_flow_test.go b/model/auth_flow_test.go new file mode 100644 index 00000000..6658bda3 --- /dev/null +++ b/model/auth_flow_test.go @@ -0,0 +1,95 @@ +package model + +import ( + "errors" + "sync" + "testing" + "time" + + "github.com/glebarez/sqlite" + "github.com/stretchr/testify/require" + "gorm.io/gorm" +) + +type failingAuthFlowReader struct{} + +func (failingAuthFlowReader) Read([]byte) (int, error) { + return 0, errors.New("entropy unavailable") +} + +func setupAuthFlowTestDB(t *testing.T) { + t.Helper() + originalDB := DB + db, err := gorm.Open(sqlite.Open("file:auth_flow_test?mode=memory&cache=shared&_pragma=busy_timeout(5000)"), &gorm.Config{}) + require.NoError(t, err) + require.NoError(t, db.AutoMigrate(&AuthFlow{})) + DB = db + t.Cleanup(func() { + DB = originalDB + if sqlDB, dbErr := db.DB(); dbErr == nil { + _ = sqlDB.Close() + } + }) +} + +func TestAuthFlowIsProviderBoundAndConsumedOnce(t *testing.T) { + setupAuthFlowTestDB(t) + token, created, err := CreateAuthFlow(AuthFlowCreate{ + Purpose: AuthFlowPurposeOAuth, Provider: "github", Intent: AuthFlowIntentLogin, + Payload: `{"affiliate_code":"invite"}`, ExpiresAt: time.Now().Add(time.Minute), + }) + require.NoError(t, err) + require.Len(t, token, 43) + require.NotEqual(t, token, created.TokenHash) + + _, err = GetAuthFlow(token, AuthFlowMatch{Purpose: AuthFlowPurposeOAuth, Provider: "discord"}) + require.ErrorIs(t, err, ErrAuthFlowInvalid) + + const consumers = 2 + errorsByConsumer := make([]error, consumers) + var wg sync.WaitGroup + for i := 0; i < consumers; i++ { + wg.Add(1) + go func(index int) { + defer wg.Done() + _, errorsByConsumer[index] = ConsumeAuthFlow(token, AuthFlowMatch{ + Purpose: AuthFlowPurposeOAuth, Provider: "github", Intent: AuthFlowIntentLogin, + }) + }(i) + } + wg.Wait() + + successes := 0 + consumed := 0 + for _, consumeErr := range errorsByConsumer { + switch { + case consumeErr == nil: + successes++ + case errors.Is(consumeErr, ErrAuthFlowConsumed): + consumed++ + default: + t.Fatalf("unexpected consume error: %v", consumeErr) + } + } + require.Equal(t, 1, successes) + require.Equal(t, 1, consumed) +} + +func TestAuthFlowExpiryIsEnforced(t *testing.T) { + setupAuthFlowTestDB(t) + token, flow, err := CreateAuthFlow(AuthFlowCreate{ + Purpose: AuthFlowPurposeOAuth, Provider: "oidc", Intent: AuthFlowIntentLogin, + ExpiresAt: time.Now().Add(time.Minute), + }) + require.NoError(t, err) + require.NoError(t, DB.Model(&AuthFlow{}).Where("id = ?", flow.Id).Update("expires_at", time.Now().Add(-time.Second)).Error) + + _, err = ConsumeAuthFlow(token, AuthFlowMatch{Purpose: AuthFlowPurposeOAuth, Provider: "oidc"}) + require.ErrorIs(t, err, ErrAuthFlowExpired) +} + +func TestGenerateAuthFlowTokenReturnsEntropyFailure(t *testing.T) { + token, err := generateAuthFlowToken(failingAuthFlowReader{}) + require.Empty(t, token) + require.ErrorContains(t, err, "entropy unavailable") +} diff --git a/model/billing_settlement.go b/model/billing_settlement.go index 4063ce2b..90994a15 100644 --- a/model/billing_settlement.go +++ b/model/billing_settlement.go @@ -509,7 +509,9 @@ func billingSettlementEffectPayload(effect *BillingSettlementEffect) (string, er if effect == nil { return "", nil } - data, err := common.Marshal(effect) + safeEffect := *effect + safeEffect.Content = common.SanitizePersistedLogContent(safeEffect.Content) + data, err := common.Marshal(&safeEffect) if err != nil { return "", err } diff --git a/model/channel.go b/model/channel.go index b1f57056..fb2c834f 100644 --- a/model/channel.go +++ b/model/channel.go @@ -5,7 +5,7 @@ import ( "encoding/json" "errors" "fmt" - "math/rand" + "maps" "strings" "sync" @@ -258,7 +258,11 @@ func (channel *Channel) GetNextEnabledKey() (string, int, *types.MaxAPIError) { switch channel.ChannelInfo.MultiKeyMode { case constant.MultiKeyModeRandom: // Randomly pick one enabled key - selectedIdx := enabledIdx[rand.Intn(len(enabledIdx))] + randomIndex, err := common.SecureRandomInt(len(enabledIdx)) + if err != nil { + return "", 0, types.NewError(err, types.ErrorCodeChannelNoAvailableKey) + } + selectedIdx := enabledIdx[randomIndex] return keys[selectedIdx], selectedIdx, nil case constant.MultiKeyModePolling: // Use channel-specific lock to ensure thread-safe polling @@ -335,7 +339,7 @@ func (channel *Channel) GetOtherInfo() map[string]interface{} { } func (channel *Channel) SetOtherInfo(otherInfo map[string]interface{}) { - otherInfoBytes, err := json.Marshal(otherInfo) + otherInfoBytes, err := common.Marshal(otherInfo) if err != nil { common.SysLog(fmt.Sprintf("failed to marshal other info: channel_id=%d, tag=%s, name=%s, error=%v", channel.Id, channel.GetTag(), channel.Name, err)) return @@ -542,15 +546,87 @@ func (channel *Channel) Insert() error { return err } +var channelEditableUpdateFields = []string{ + "type", + "key", + "openai_organization", + "test_model", + "status", + "name", + "weight", + "base_url", + "other", + "models", + "group", + "model_mapping", + "status_code_mapping", + "priority", + "auto_ban", + "tag", + "setting", + "param_override", + "header_override", + "remark", + "channel_info", + "settings", +} + func (channel *Channel) Update() error { - // If this is a multi-key channel, recalculate MultiKeySize based on the current key list to avoid inconsistency after editing keys - if channel.ChannelInfo.IsMultiKey { + return channel.UpdateFields(channelEditableUpdateFields...) +} + +func cloneChannelValue[T any](value *T) *T { + if value == nil { + return nil + } + cloned := *value + return &cloned +} + +func cloneChannelForUpdate(channel *Channel) *Channel { + cloned := *channel + cloned.OpenAIOrganization = cloneChannelValue(channel.OpenAIOrganization) + cloned.TestModel = cloneChannelValue(channel.TestModel) + cloned.Weight = cloneChannelValue(channel.Weight) + cloned.BaseURL = cloneChannelValue(channel.BaseURL) + cloned.ModelMapping = cloneChannelValue(channel.ModelMapping) + cloned.StatusCodeMapping = cloneChannelValue(channel.StatusCodeMapping) + cloned.Priority = cloneChannelValue(channel.Priority) + cloned.AutoBan = cloneChannelValue(channel.AutoBan) + cloned.Tag = cloneChannelValue(channel.Tag) + cloned.Setting = cloneChannelValue(channel.Setting) + cloned.ParamOverride = cloneChannelValue(channel.ParamOverride) + cloned.HeaderOverride = cloneChannelValue(channel.HeaderOverride) + cloned.Remark = cloneChannelValue(channel.Remark) + cloned.ChannelInfo.MultiKeyStatusList = maps.Clone(channel.ChannelInfo.MultiKeyStatusList) + cloned.ChannelInfo.MultiKeyDisabledReason = maps.Clone(channel.ChannelInfo.MultiKeyDisabledReason) + cloned.ChannelInfo.MultiKeyDisabledTime = maps.Clone(channel.ChannelInfo.MultiKeyDisabledTime) + cloned.Keys = append([]string(nil), channel.Keys...) + return &cloned +} + +func (channel *Channel) UpdateFields(fields ...string) error { + if channel.Id == 0 { + return errors.New("channel ID is 0") + } + workingChannel := cloneChannelForUpdate(channel) + fieldSet := make(map[string]struct{}, len(fields)) + for _, field := range fields { + fieldSet[field] = struct{}{} + } + _, keySelected := fieldSet["key"] + _, channelInfoSelected := fieldSet["channel_info"] + multiKeyMetadataUpdated := false + + // Keep multi-key metadata consistent when either the key list or its metadata is persisted. + if workingChannel.ChannelInfo.IsMultiKey && (keySelected || channelInfoSelected) { + multiKeyMetadataUpdated = true var keyStr string - if channel.Key != "" { - keyStr = channel.Key + if keySelected || workingChannel.Key != "" { + keyStr = workingChannel.Key } else { // If key is not provided, read the existing key from the database - if existing, err := GetChannelById(channel.Id, true); err == nil { + if existing, err := GetChannelById(workingChannel.Id, true); err == nil { keyStr = existing.Key } } @@ -571,24 +647,145 @@ func (channel *Channel) Update() error { keys = strings.Split(strings.Trim(keyStr, "\n"), "\n") } } - channel.ChannelInfo.MultiKeySize = len(keys) + workingChannel.ChannelInfo.MultiKeySize = len(keys) // Clean up status data that exceeds the new key count to prevent index out of range - if channel.ChannelInfo.MultiKeyStatusList != nil { - for idx := range channel.ChannelInfo.MultiKeyStatusList { - if idx >= channel.ChannelInfo.MultiKeySize { - delete(channel.ChannelInfo.MultiKeyStatusList, idx) + if workingChannel.ChannelInfo.MultiKeyStatusList != nil { + for idx := range workingChannel.ChannelInfo.MultiKeyStatusList { + if idx >= workingChannel.ChannelInfo.MultiKeySize { + delete(workingChannel.ChannelInfo.MultiKeyStatusList, idx) + } + } + } + if workingChannel.ChannelInfo.MultiKeyDisabledTime != nil { + for idx := range workingChannel.ChannelInfo.MultiKeyDisabledTime { + if idx >= workingChannel.ChannelInfo.MultiKeySize { + delete(workingChannel.ChannelInfo.MultiKeyDisabledTime, idx) + } + } + } + if workingChannel.ChannelInfo.MultiKeyDisabledReason != nil { + for idx := range workingChannel.ChannelInfo.MultiKeyDisabledReason { + if idx >= workingChannel.ChannelInfo.MultiKeySize { + delete(workingChannel.ChannelInfo.MultiKeyDisabledReason, idx) } } } } - var err error - err = DB.Model(channel).Updates(channel).Error + updates := workingChannel.editableUpdateMap(fields) + if multiKeyMetadataUpdated && keySelected { + updates["channel_info"] = workingChannel.ChannelInfo + } + if len(updates) == 0 { + return nil + } + abilityFields := []string{"models", "group", "status", "priority", "weight", "tag"} + updateAbilities := false + for _, field := range abilityFields { + if _, ok := updates[field]; ok { + updateAbilities = true + break + } + } + + err := DB.Transaction(func(tx *gorm.DB) error { + if err := tx.Model(&Channel{}).Where("id = ?", workingChannel.Id).Updates(updates).Error; err != nil { + return err + } + if err := tx.Model(&Channel{}).First(workingChannel, "id = ?", workingChannel.Id).Error; err != nil { + return err + } + if keySelected { + workingChannel.Keys = nil + } + if !updateAbilities { + return nil + } + return workingChannel.UpdateAbilities(tx) + }) if err != nil { return err } - DB.Model(channel).First(channel, "id = ?", channel.Id) - err = channel.UpdateAbilities(nil) - return err + *channel = *workingChannel + return nil +} + +func (channel *Channel) editableUpdateMap(fields []string) map[string]any { + updates := make(map[string]any, len(fields)) + for _, field := range fields { + switch field { + case "type": + updates["type"] = channel.Type + case "key": + updates["key"] = channel.Key + case "openai_organization": + if channel.OpenAIOrganization != nil { + updates["openai_organization"] = channel.OpenAIOrganization + } + case "test_model": + if channel.TestModel != nil { + updates["test_model"] = channel.TestModel + } + case "status": + updates["status"] = channel.Status + case "name": + updates["name"] = channel.Name + case "weight": + if channel.Weight != nil { + updates["weight"] = channel.Weight + } + case "base_url": + if channel.BaseURL != nil { + updates["base_url"] = channel.BaseURL + } + case "other": + updates["other"] = channel.Other + case "models": + updates["models"] = channel.Models + case "group": + updates["group"] = channel.Group + case "model_mapping": + if channel.ModelMapping != nil { + updates["model_mapping"] = channel.ModelMapping + } + case "status_code_mapping": + if channel.StatusCodeMapping != nil { + updates["status_code_mapping"] = channel.StatusCodeMapping + } + case "priority": + if channel.Priority != nil { + updates["priority"] = channel.Priority + } + case "auto_ban": + if channel.AutoBan != nil { + updates["auto_ban"] = channel.AutoBan + } + case "tag": + if channel.Tag != nil { + updates["tag"] = channel.Tag + } + case "setting": + if channel.Setting != nil { + updates["setting"] = channel.Setting + } + case "param_override": + if channel.ParamOverride != nil { + updates["param_override"] = channel.ParamOverride + } + case "header_override": + if channel.HeaderOverride != nil { + updates["header_override"] = channel.HeaderOverride + } + case "remark": + if channel.Remark != nil { + updates["remark"] = channel.Remark + } + case "channel_info": + updates["channel_info"] = channel.ChannelInfo + case "settings": + updates["settings"] = channel.OtherSettings + } + } + return updates } func (channel *Channel) UpdateResponseTime(responseTime int64) { diff --git a/model/channel_cache.go b/model/channel_cache.go index 8896d1ad..832bced4 100644 --- a/model/channel_cache.go +++ b/model/channel_cache.go @@ -3,7 +3,6 @@ package model import ( "errors" "fmt" - "math/rand" "sort" "strings" "sync" @@ -223,7 +222,10 @@ func GetRandomSatisfiedChannelExcluding(group string, model string, retry int, r totalWeight := sumWeight * smoothingFactor // Generate a random value in the range [0, totalWeight) - randomWeight := rand.Intn(totalWeight) + randomWeight, err := common.SecureRandomInt(totalWeight) + if err != nil { + return nil, err + } // Find a channel based on its weight for _, channel := range targetChannels { diff --git a/model/channel_update_fields_test.go b/model/channel_update_fields_test.go new file mode 100644 index 00000000..36465cb6 --- /dev/null +++ b/model/channel_update_fields_test.go @@ -0,0 +1,168 @@ +package model + +import ( + "fmt" + "maps" + "strings" + "testing" + + "github.com/MAX-API-Next/MAX-API/common" + "github.com/glebarez/sqlite" + "github.com/stretchr/testify/require" + "gorm.io/gorm" +) + +func setupChannelUpdateFieldsTestDB(t *testing.T) (*gorm.DB, *Channel) { + t.Helper() + + originalDB := DB + originalLogDB := LOG_DB + originalUsingSQLite := common.UsingSQLite + originalUsingMySQL := common.UsingMySQL + originalUsingPostgreSQL := common.UsingPostgreSQL + originalRedisEnabled := common.RedisEnabled + + common.UsingSQLite = true + common.UsingMySQL = false + common.UsingPostgreSQL = false + common.RedisEnabled = false + + dsn := fmt.Sprintf("file:%s?mode=memory&cache=shared", strings.ReplaceAll(t.Name(), "/", "_")) + db, err := gorm.Open(sqlite.Open(dsn), &gorm.Config{}) + require.NoError(t, err) + DB = db + LOG_DB = db + require.NoError(t, db.AutoMigrate(&Channel{}, &Ability{})) + + sqlDB, err := db.DB() + require.NoError(t, err) + t.Cleanup(func() { + DB = originalDB + LOG_DB = originalLogDB + common.UsingSQLite = originalUsingSQLite + common.UsingMySQL = originalUsingMySQL + common.UsingPostgreSQL = originalUsingPostgreSQL + common.RedisEnabled = originalRedisEnabled + _ = sqlDB.Close() + }) + + priority := int64(10) + weight := uint(20) + tag := "stable" + channel := &Channel{ + Name: "channel-update-fields", + Key: "first-key\nsecond-key\nthird-key", + Status: common.ChannelStatusEnabled, + Models: "gpt-test", + Group: "default", + Priority: &priority, + Weight: &weight, + Tag: &tag, + ChannelInfo: ChannelInfo{ + IsMultiKey: true, + MultiKeySize: 3, + MultiKeyStatusList: map[int]int{2: 2}, + MultiKeyDisabledTime: map[int]int64{2: 123}, + MultiKeyDisabledReason: map[int]string{2: "failed"}, + }, + } + require.NoError(t, db.Create(channel).Error) + require.NoError(t, channel.AddAbilities(nil)) + return db, channel +} + +func installFailingAbilityDeleteTrigger(t *testing.T, db *gorm.DB) { + t.Helper() + require.NoError(t, db.Exec(`CREATE TRIGGER fail_ability_delete BEFORE DELETE ON abilities BEGIN SELECT RAISE(ABORT, 'ability delete blocked'); END`).Error) + t.Cleanup(func() { _ = db.Exec("DROP TRIGGER IF EXISTS fail_ability_delete").Error }) +} + +func TestUpdateFieldsSkipsAbilityRebuildForChannelInfoOnly(t *testing.T) { + db, channel := setupChannelUpdateFieldsTestDB(t) + installFailingAbilityDeleteTrigger(t, db) + + channel.ChannelInfo.MultiKeyPollingIndex = 1 + require.NoError(t, channel.UpdateFields("channel_info")) + + var stored Channel + require.NoError(t, db.First(&stored, channel.Id).Error) + require.Equal(t, 1, stored.ChannelInfo.MultiKeyPollingIndex) + var abilityCount int64 + require.NoError(t, db.Model(&Ability{}).Where("channel_id = ?", channel.Id).Count(&abilityCount).Error) + require.EqualValues(t, 1, abilityCount) +} + +func TestUpdateFieldsRollsBackChannelWhenAbilityRebuildFails(t *testing.T) { + db, channel := setupChannelUpdateFieldsTestDB(t) + installFailingAbilityDeleteTrigger(t, db) + + channel.Key = "first-key\nsecond-key" + channel.Status = common.ChannelStatusManuallyDisabled + priority := int64(99) + channel.Priority = &priority + expected := *channel + expectedPriority := *channel.Priority + expected.Priority = &expectedPriority + expected.ChannelInfo.MultiKeyStatusList = maps.Clone(channel.ChannelInfo.MultiKeyStatusList) + expected.ChannelInfo.MultiKeyDisabledTime = maps.Clone(channel.ChannelInfo.MultiKeyDisabledTime) + expected.ChannelInfo.MultiKeyDisabledReason = maps.Clone(channel.ChannelInfo.MultiKeyDisabledReason) + expected.Keys = append([]string(nil), channel.Keys...) + + require.Error(t, channel.UpdateFields("key", "status", "priority")) + require.Equal(t, expected, *channel) + + var stored Channel + require.NoError(t, db.First(&stored, channel.Id).Error) + require.Equal(t, common.ChannelStatusEnabled, stored.Status) + require.NotNil(t, stored.Priority) + require.EqualValues(t, 10, *stored.Priority) + require.Equal(t, "first-key\nsecond-key\nthird-key", stored.Key) + require.Equal(t, 3, stored.ChannelInfo.MultiKeySize) + require.Equal(t, map[int]int{2: 2}, stored.ChannelInfo.MultiKeyStatusList) + require.Equal(t, map[int]int64{2: 123}, stored.ChannelInfo.MultiKeyDisabledTime) + require.Equal(t, map[int]string{2: "failed"}, stored.ChannelInfo.MultiKeyDisabledReason) +} + +func TestUpdateFieldsPersistsMultiKeyInfoWhenOnlyKeyChanges(t *testing.T) { + db, channel := setupChannelUpdateFieldsTestDB(t) + + channel.Keys = channel.GetKeys() + require.Equal(t, []string{"first-key", "second-key", "third-key"}, channel.Keys) + channel.Key = "first-key\nsecond-key" + require.NoError(t, channel.UpdateFields("key")) + require.Equal(t, []string{"first-key", "second-key"}, channel.GetKeys()) + require.Equal(t, 2, channel.ChannelInfo.MultiKeySize) + require.NotContains(t, channel.ChannelInfo.MultiKeyStatusList, 2) + require.NotContains(t, channel.ChannelInfo.MultiKeyDisabledTime, 2) + require.NotContains(t, channel.ChannelInfo.MultiKeyDisabledReason, 2) + + var stored Channel + require.NoError(t, db.First(&stored, channel.Id).Error) + require.Equal(t, "first-key\nsecond-key", stored.Key) + require.Equal(t, 2, stored.ChannelInfo.MultiKeySize) + _, exists := stored.ChannelInfo.MultiKeyStatusList[2] + require.False(t, exists) + _, exists = stored.ChannelInfo.MultiKeyDisabledTime[2] + require.False(t, exists) + _, exists = stored.ChannelInfo.MultiKeyDisabledReason[2] + require.False(t, exists) +} + +func TestUpdateFieldsClearsFinalMultiKeyAndMetadata(t *testing.T) { + db, channel := setupChannelUpdateFieldsTestDB(t) + + channel.Key = "" + require.NoError(t, channel.UpdateFields("key")) + require.Zero(t, channel.ChannelInfo.MultiKeySize) + require.Empty(t, channel.ChannelInfo.MultiKeyStatusList) + require.Empty(t, channel.ChannelInfo.MultiKeyDisabledTime) + require.Empty(t, channel.ChannelInfo.MultiKeyDisabledReason) + + var stored Channel + require.NoError(t, db.First(&stored, channel.Id).Error) + require.Empty(t, stored.Key) + require.Zero(t, stored.ChannelInfo.MultiKeySize) + require.Empty(t, stored.ChannelInfo.MultiKeyStatusList) + require.Empty(t, stored.ChannelInfo.MultiKeyDisabledTime) + require.Empty(t, stored.ChannelInfo.MultiKeyDisabledReason) +} diff --git a/model/checkin.go b/model/checkin.go index f240bd38..1b1ec237 100644 --- a/model/checkin.go +++ b/model/checkin.go @@ -2,7 +2,6 @@ package model import ( "errors" - "math/rand" "time" "github.com/MAX-API-Next/MAX-API/common" @@ -70,7 +69,11 @@ func UserCheckin(userId int) (*Checkin, error) { // 计算随机额度奖励 quotaAwarded := setting.MinQuota if setting.MaxQuota > setting.MinQuota { - quotaAwarded = setting.MinQuota + rand.Intn(setting.MaxQuota-setting.MinQuota+1) + randomQuota, err := common.SecureRandomInt(setting.MaxQuota - setting.MinQuota + 1) + if err != nil { + return nil, err + } + quotaAwarded = setting.MinQuota + randomQuota } today := time.Now().Format("2006-01-02") diff --git a/model/log.go b/model/log.go index 5ecc4250..b27f6882 100644 --- a/model/log.go +++ b/model/log.go @@ -73,6 +73,7 @@ type BillingLogReceipt struct { func (log *Log) BeforeSave(*gorm.DB) error { log.syncRetryMarker() + log.Content = common.SanitizePersistedLogContent(log.Content) return nil } @@ -394,6 +395,7 @@ func RecordLog(userId int, logType int, content string) { if logType == LogTypeConsume && !common.LogConsumeEnabled { return } + content = common.SanitizePersistedLogContent(content) username, _ := GetUsernameById(userId, false) log := &Log{ UserId: userId, @@ -413,6 +415,7 @@ func RecordLogWithAdminInfo(userId int, logType int, content string, adminInfo m if logType == LogTypeConsume && !common.LogConsumeEnabled { return } + content = common.SanitizePersistedLogContent(content) username, _ := GetUsernameById(userId, false) log := &Log{ UserId: userId, @@ -443,6 +446,7 @@ func buildOpField(action string, params map[string]interface{}) map[string]inter } func RecordLoginLog(userId int, username string, content string, ip string, action string, params map[string]interface{}, extra map[string]interface{}) { + content = common.SanitizePersistedLogContent(content) other := map[string]interface{}{} for k, v := range extra { other[k] = v @@ -463,6 +467,7 @@ func RecordLoginLog(userId int, username string, content string, ip string, acti } func RecordOperationAuditLog(logUserId int, content string, ip string, action string, params map[string]interface{}, adminInfo map[string]interface{}, auditInfo map[string]interface{}) { + content = common.SanitizePersistedLogContent(content) username, _ := GetUsernameById(logUserId, false) other := map[string]interface{}{ "op": buildOpField(action, params), @@ -488,6 +493,7 @@ func RecordOperationAuditLog(logUserId int, content string, ip string, action st } func RecordTopupLog(userId int, content string, callerIp string, paymentMethod string, callbackPaymentMethod string) { + content = common.SanitizePersistedLogContent(content) username, _ := GetUsernameById(userId, false) adminInfo := map[string]interface{}{ "server_ip": common.GetIp(), @@ -517,6 +523,7 @@ func RecordTopupLog(userId int, content string, callerIp string, paymentMethod s func RecordErrorLog(c *gin.Context, userId int, channelId int, modelName string, tokenName string, content string, tokenId int, useTimeSeconds int, isStream bool, group string, other map[string]interface{}) { + content = common.SanitizePersistedLogContent(content) logger.LogInfo(c, fmt.Sprintf("record error log: userId=%d, channelId=%d, modelName=%s, tokenName=%s, content=%s", userId, channelId, modelName, tokenName, common.LocalLogPreview(content))) username := c.GetString("username") requestId := c.GetString(common.RequestIdKey) @@ -580,6 +587,7 @@ func RecordConsumeLog(c *gin.Context, userId int, params RecordConsumeLogParams) if !common.LogConsumeEnabled { return } + params.Content = common.SanitizePersistedLogContent(params.Content) logger.LogInfo(c, fmt.Sprintf("record consume log: userId=%d, params=%s", userId, common.GetJsonString(params))) username := c.GetString("username") requestId := c.GetString(common.RequestIdKey) @@ -672,6 +680,7 @@ func recordTaskBillingLog(operationKey string, params RecordTaskBillingLogParams if LOG_DB == nil { return errors.New("log database is not initialized") } + params.Content = common.SanitizePersistedLogContent(params.Content) username, _ := GetUsernameById(params.UserId, false) tokenName := "" if params.TokenId > 0 { @@ -1020,12 +1029,16 @@ func SumUsedToken(logType int, startTimestamp int64, endTimestamp int64, modelNa func CountOldLog(ctx context.Context, targetTimestamp int64) (int64, error) { var total int64 - if err := LOG_DB.WithContext(ctx).Model(&Log{}).Where("created_at < ?", targetTimestamp).Count(&total).Error; err != nil { + if err := oldLogCleanupScope(LOG_DB.WithContext(ctx).Model(&Log{}), targetTimestamp).Count(&total).Error; err != nil { return 0, err } return total, nil } +func oldLogCleanupScope(tx *gorm.DB, targetTimestamp int64) *gorm.DB { + return tx.Where("created_at < ? AND (type <> ? OR type IS NULL)", targetTimestamp, LogTypeManage) +} + func DeleteOldLogBatch(ctx context.Context, targetTimestamp int64, limit int) (int64, error) { if limit <= 0 { limit = 100 @@ -1035,9 +1048,7 @@ func DeleteOldLogBatch(ctx context.Context, targetTimestamp int64, limit int) (i } ids := make([]int, 0, limit) - err := LOG_DB.WithContext(ctx). - Model(&Log{}). - Where("created_at < ?", targetTimestamp). + err := oldLogCleanupScope(LOG_DB.WithContext(ctx).Model(&Log{}), targetTimestamp). Order("created_at ASC, id ASC"). Limit(limit). Pluck("id", &ids).Error @@ -1057,7 +1068,7 @@ func DeleteOldLogBatch(ctx context.Context, targetTimestamp int64, limit int) (i var rowsAffected int64 err = LOG_DB.WithContext(ctx).Transaction(func(tx *gorm.DB) error { - result := tx.Where("id IN ?", ids).Delete(&Log{}) + result := oldLogCleanupScope(tx.Where("id IN ?", ids), targetTimestamp).Delete(&Log{}) if nil != result.Error { return result.Error } @@ -1088,6 +1099,13 @@ func DeleteOldLog(ctx context.Context, targetTimestamp int64, limit int) (int64, } total += rowsAffected if rowsAffected < int64(limit) { + remaining, err := CountOldLog(ctx, targetTimestamp) + if err != nil { + return total, err + } + if remaining > 0 { + continue + } if _, err := DeleteOldBillingLogReceipts(ctx, targetTimestamp, limit); err != nil { return total, err } diff --git a/model/log_cleanup_test.go b/model/log_cleanup_test.go index 9f4a81fe..966ff296 100644 --- a/model/log_cleanup_test.go +++ b/model/log_cleanup_test.go @@ -41,6 +41,103 @@ func TestDeleteOldLogBatchDeletesOnlyLimitedRows(t *testing.T) { assert.EqualValues(t, 0, remainingOld) } +func TestDeleteOldLogBatchPreservesManagementAuditLogs(t *testing.T) { + truncateTables(t) + + logs := []Log{ + {UserId: 1, CreatedAt: 10, Type: LogTypeConsume}, + {UserId: 1, CreatedAt: 20, Type: LogTypeManage}, + {UserId: 1, CreatedAt: 1000, Type: LogTypeConsume}, + } + require.NoError(t, LOG_DB.Create(&logs).Error) + + ctx := context.Background() + remainingOld, err := CountOldLog(ctx, 100) + require.NoError(t, err) + assert.EqualValues(t, 1, remainingOld) + + rows, err := DeleteOldLogBatch(ctx, 100, 10) + require.NoError(t, err) + assert.EqualValues(t, 1, rows) + + remainingOld, err = CountOldLog(ctx, 100) + require.NoError(t, err) + assert.EqualValues(t, 0, remainingOld) + + var manageCount int64 + require.NoError(t, LOG_DB.Model(&Log{}).Where("type = ?", LogTypeManage).Count(&manageCount).Error) + assert.EqualValues(t, 1, manageCount) +} + +func TestDeleteOldLogBatchDeletesLegacyNullTypeLogs(t *testing.T) { + truncateTables(t) + + require.NoError(t, LOG_DB.Exec("INSERT INTO logs (user_id, created_at, type) VALUES (?, ?, NULL)", 1, 10).Error) + remainingOld, err := CountOldLog(context.Background(), 100) + require.NoError(t, err) + assert.EqualValues(t, 1, remainingOld) + + rows, err := DeleteOldLogBatch(context.Background(), 100, 10) + require.NoError(t, err) + assert.EqualValues(t, 1, rows) +} + +func TestDeleteOldLogBatchRechecksManagementTypeBeforeDelete(t *testing.T) { + truncateTables(t) + + logEntry := Log{UserId: 1, CreatedAt: 10, Type: LogTypeConsume} + require.NoError(t, LOG_DB.Create(&logEntry).Error) + callbackName := "test:promote_log_to_manage_before_cleanup_delete" + mutated := false + require.NoError(t, LOG_DB.Callback().Delete().Before("gorm:delete").Register(callbackName, func(tx *gorm.DB) { + if mutated || tx.Statement == nil || tx.Statement.Table != "logs" { + return + } + mutated = true + require.NoError(t, tx.Session(&gorm.Session{NewDB: true}).Model(&Log{}).Where("id = ?", logEntry.Id).Update("type", LogTypeManage).Error) + })) + t.Cleanup(func() { _ = LOG_DB.Callback().Delete().Remove(callbackName) }) + + rows, err := DeleteOldLogBatch(context.Background(), 100, 10) + require.NoError(t, err) + assert.Zero(t, rows) + + var stored Log + require.NoError(t, LOG_DB.First(&stored, logEntry.Id).Error) + assert.Equal(t, LogTypeManage, stored.Type) +} + +func TestDeleteOldLogContinuesAfterPromotedCandidateShortensBatch(t *testing.T) { + truncateTables(t) + + logs := []Log{ + {UserId: 1, CreatedAt: 10, Type: LogTypeConsume}, + {UserId: 1, CreatedAt: 20, Type: LogTypeConsume}, + } + require.NoError(t, LOG_DB.Create(&logs).Error) + callbackName := "test:promote_first_log_before_full_cleanup_delete" + mutated := false + require.NoError(t, LOG_DB.Callback().Delete().Before("gorm:delete").Register(callbackName, func(tx *gorm.DB) { + if mutated || tx.Statement == nil || tx.Statement.Table != "logs" { + return + } + mutated = true + require.NoError(t, tx.Session(&gorm.Session{NewDB: true}).Model(&Log{}).Where("id = ?", logs[0].Id).Update("type", LogTypeManage).Error) + })) + t.Cleanup(func() { _ = LOG_DB.Callback().Delete().Remove(callbackName) }) + + rows, err := DeleteOldLog(context.Background(), 100, 1) + require.NoError(t, err) + assert.EqualValues(t, 1, rows) + + var promoted Log + require.NoError(t, LOG_DB.First(&promoted, logs[0].Id).Error) + assert.Equal(t, LogTypeManage, promoted.Type) + var laterCount int64 + require.NoError(t, LOG_DB.Model(&Log{}).Where("id = ?", logs[1].Id).Count(&laterCount).Error) + assert.Zero(t, laterCount) +} + func TestDeleteOldLogBatchStopsWhenContextCancelledAfterSelect(t *testing.T) { truncateTables(t) diff --git a/model/log_test.go b/model/log_test.go index 6b4421df..bcddc89d 100644 --- a/model/log_test.go +++ b/model/log_test.go @@ -5,8 +5,10 @@ import ( "net/http" "net/http/httptest" "os" + "strings" "testing" "time" + "unicode/utf8" "github.com/MAX-API-Next/MAX-API/common" "github.com/gin-gonic/gin" @@ -103,6 +105,78 @@ func newLogTestContext() *gin.Context { return c } +func TestRecordConsumeLogSanitizesPersistedContent(t *testing.T) { + originalLogConsumeEnabled := common.LogConsumeEnabled + common.LogConsumeEnabled = true + t.Cleanup(func() { + common.LogConsumeEnabled = originalLogConsumeEnabled + }) + + require.NoError(t, LOG_DB.Where("1 = 1").Delete(&Log{}).Error) + t.Cleanup(func() { + require.NoError(t, LOG_DB.Where("1 = 1").Delete(&Log{}).Error) + }) + + content := "line1\r\nline2\tline3\x00" + strings.Repeat("x", common.PersistedLogContentLimit+1) + RecordConsumeLog(newLogTestContext(), 7, RecordConsumeLogParams{ + ModelName: "gpt-test", + Quota: 1, + Content: content, + }) + + var log Log + require.NoError(t, LOG_DB.Last(&log).Error) + require.True(t, utf8.ValidString(log.Content)) + require.NotContains(t, log.Content, "\r") + require.NotContains(t, log.Content, "\n") + require.NotContains(t, log.Content, "\t") + require.NotContains(t, log.Content, "\x00") + require.Contains(t, log.Content, "line1 line2 line3") + require.True(t, strings.HasSuffix(log.Content, "... [truncated]")) + require.LessOrEqual(t, utf8.RuneCountInString(log.Content), common.PersistedLogContentLimit) +} + +func TestBillingSettlementEffectPayloadSanitizesContent(t *testing.T) { + payload, err := billingSettlementEffectPayload(&BillingSettlementEffect{ + LogType: LogTypeConsume, + Content: "effect\r\ncontent\x00" + strings.Repeat("x", common.PersistedLogContentLimit+1), + }) + require.NoError(t, err) + + effect, err := decodeBillingSettlementEffect(payload) + require.NoError(t, err) + require.NotNil(t, effect) + require.NotContains(t, effect.Content, "\r") + require.NotContains(t, effect.Content, "\n") + require.NotContains(t, effect.Content, "\x00") + require.True(t, strings.HasSuffix(effect.Content, "... [truncated]")) +} + +func TestMarkTaskSubmitFailedSanitizesFailReason(t *testing.T) { + require.NoError(t, DB.Where("1 = 1").Delete(&Task{}).Error) + t.Cleanup(func() { + require.NoError(t, DB.Where("1 = 1").Delete(&Task{}).Error) + }) + + task := &Task{ + TaskID: "sanitize-fail-reason", + Status: TaskStatusSubmitted, + SubmitTime: time.Now().Unix(), + } + require.NoError(t, DB.Create(task).Error) + + reason := "upstream\r\nfailed\x00" + strings.Repeat("x", common.PersistedLogContentLimit+1) + require.NoError(t, MarkTaskSubmitFailed(task.ID, reason)) + + var reloaded Task + require.NoError(t, DB.First(&reloaded, task.ID).Error) + require.EqualValues(t, TaskStatusFailure, reloaded.Status) + require.NotContains(t, reloaded.FailReason, "\r") + require.NotContains(t, reloaded.FailReason, "\n") + require.NotContains(t, reloaded.FailReason, "\x00") + require.True(t, strings.HasSuffix(reloaded.FailReason, "... [truncated]")) +} + func TestLogQuotaFilterIndexes(t *testing.T) { db := newRetryBackfillTestDB(t, &Log{}) diff --git a/model/main.go b/model/main.go index 131bb2f1..34967397 100644 --- a/model/main.go +++ b/model/main.go @@ -338,9 +338,11 @@ func migrateDB() error { &SubscriptionPreConsumeRecord{}, &BillingSettlement{}, &BillingPreConsumeSelection{}, + &MidjourneyBillingClaim{}, &CacheInvalidationTask{}, &CustomOAuthProvider{}, &UserOAuthBinding{}, + &AuthFlow{}, &PerfMetric{}, &SystemTask{}, &SystemInstance{}, diff --git a/model/midjourney.go b/model/midjourney.go index 3abf6275..a6edaf5d 100644 --- a/model/midjourney.go +++ b/model/midjourney.go @@ -1,5 +1,16 @@ package model +import ( + "errors" + "fmt" + "time" + "unicode/utf8" + + "github.com/MAX-API-Next/MAX-API/common" + "gorm.io/gorm" + "gorm.io/gorm/clause" +) + type Midjourney struct { Id int `json:"id"` Code int `json:"code"` @@ -25,6 +36,28 @@ type Midjourney struct { Properties string `json:"properties"` } +func (midjourney *Midjourney) BeforeSave(*gorm.DB) error { + if midjourney != nil { + midjourney.FailReason = common.SanitizePersistedLogContent(midjourney.FailReason) + } + return nil +} + +// MidjourneyBillingClaim binds one provider task identity to the single local +// task that owns its pre-consumed quota. Historical Midjourney rows use a zero +// BillingTaskID tombstone so later duplicate submissions cannot be charged. +type MidjourneyBillingClaim struct { + ID int64 `gorm:"primaryKey"` + ChannelID int `gorm:"uniqueIndex:idx_midjourney_billing_claim,priority:1;not null"` + MjID string `gorm:"type:varchar(180);uniqueIndex:idx_midjourney_billing_claim,priority:2;not null"` + UserID int `gorm:"index;not null"` + BillingTaskID int64 `gorm:"index;not null;default:0"` + BillingRequestID string `gorm:"type:varchar(64);index;not null;default:''"` + CreatedAt int64 `gorm:"not null"` +} + +var ErrMidjourneyTaskAmbiguous = errors.New("midjourney task identity is ambiguous across channels") + // TaskQueryParams 用于包含所有搜索条件的结构体，可以根据需求添加更多字段 type TaskQueryParams struct { ChannelID string @@ -104,34 +137,43 @@ func GetAllUnFinishTasks() []*Midjourney { return tasks } -func GetByOnlyMJId(mjId string) *Midjourney { - var mj *Midjourney +func GetByMJIds(userId int, mjIds []string) []*Midjourney { + var mj []*Midjourney var err error - err = DB.Where("mj_id = ?", mjId).First(&mj).Error + err = DB.Where("user_id = ? and mj_id in (?)", userId, mjIds).Find(&mj).Error if err != nil { return nil } return mj } -func GetByMJId(userId int, mjId string) *Midjourney { - var mj *Midjourney - var err error - err = DB.Where("user_id = ? and mj_id = ?", userId, mjId).First(&mj).Error - if err != nil { - return nil +func GetUniqueMidjourneyByMJID(mjID string) (*Midjourney, error) { + if mjID == "" { + return nil, nil } - return mj + return findUniqueMidjourney(DB.Where("mj_id = ?", mjID)) } -func GetByMJIds(userId int, mjIds []string) []*Midjourney { - var mj []*Midjourney - var err error - err = DB.Where("user_id = ? and mj_id in (?)", userId, mjIds).Find(&mj).Error - if err != nil { - return nil +func GetUniqueMidjourneyByUserAndMJID(userID int, mjID string) (*Midjourney, error) { + if userID <= 0 || mjID == "" { + return nil, nil + } + return findUniqueMidjourney(DB.Where("user_id = ? AND mj_id = ?", userID, mjID)) +} + +func findUniqueMidjourney(query *gorm.DB) (*Midjourney, error) { + var tasks []Midjourney + if err := query.Order("id").Limit(2).Find(&tasks).Error; err != nil { + return nil, err + } + switch len(tasks) { + case 0: + return nil, nil + case 1: + return &tasks[0], nil + default: + return nil, ErrMidjourneyTaskAmbiguous } - return mj } func GetMjByuId(id int) *Midjourney { @@ -160,31 +202,271 @@ func (midjourney *Midjourney) Update() error { return err } -// UpdateWithStatus performs a conditional UPDATE guarded by fromStatus (CAS). -// Returns (true, nil) if this caller won the update, (false, nil) if -// another process already moved the task out of fromStatus. -// UpdateWithStatus performs a conditional UPDATE guarded by fromStatus (CAS). -// Uses Model().Select("*").Updates() to avoid GORM Save()'s INSERT fallback. -func (midjourney *Midjourney) UpdateWithStatus(fromStatus string) (bool, error) { - result := DB.Model(midjourney).Where("status = ?", fromStatus).Select("*").Updates(midjourney) +func midjourneyUpdateValues(midjourney *Midjourney) map[string]interface{} { + return map[string]interface{}{ + "code": midjourney.Code, "action": midjourney.Action, + "prompt": midjourney.Prompt, "prompt_en": midjourney.PromptEn, + "description": midjourney.Description, "state": midjourney.State, + "submit_time": midjourney.SubmitTime, "start_time": midjourney.StartTime, + "finish_time": midjourney.FinishTime, "image_url": midjourney.ImageUrl, + "video_url": midjourney.VideoUrl, "video_urls": midjourney.VideoUrls, + "status": midjourney.Status, "progress": midjourney.Progress, + "fail_reason": common.SanitizePersistedLogContent(midjourney.FailReason), "channel_id": midjourney.ChannelId, + "quota": midjourney.Quota, "buttons": midjourney.Buttons, + "properties": midjourney.Properties, + } +} + +func midjourneyTaskUpdateValues(task *Task, includeQuota bool) map[string]interface{} { + values := task.statusUpdateValues() + values["action"] = task.Action + values["channel_id"] = task.ChannelId + if includeQuota { + values["quota"] = task.Quota + } + return values +} + +// FinalizeMidjourneySubmission atomically claims the provider task ID, writes +// the legacy Midjourney row, advances the durable billing Task, and stages the +// matching settlement intent. It returns false when another local request (or +// a historical row) already owns the provider task; in that case the supplied +// duplicate refund intent is staged with the losing Task transition. +func FinalizeMidjourneySubmission(midjourney *Midjourney, billingTask *Task, settlementInput, duplicateRefundInput *BillingSettlementInput) (created bool, refundDuplicate bool, err error) { + if midjourney == nil || midjourney.UserId <= 0 || midjourney.MjId == "" { + return false, false, errors.New("midjourney user and provider task id are required") + } + if utf8.RuneCountInString(midjourney.MjId) > 180 { + return false, false, errors.New("midjourney provider task id exceeds 180 characters") + } + if billingTask == nil || billingTask.ID <= 0 { + return false, false, errors.New("persisted midjourney billing task is required") + } + if settlementInput != nil { + if err := validateBillingSettlementInput(*settlementInput); err != nil { + return false, false, err + } + } + if duplicateRefundInput != nil { + if err := validateBillingSettlementInput(*duplicateRefundInput); err != nil { + return false, false, err + } + } + + err = DB.Transaction(func(tx *gorm.DB) error { + var existing Midjourney + existingResult := tx.Where("channel_id = ? AND mj_id = ?", midjourney.ChannelId, midjourney.MjId). + Limit(1).Find(&existing) + if existingResult.Error != nil { + return existingResult.Error + } + + claimTaskID := billingTask.ID + claimRequestID := billingTask.PrivateData.BillingRequestId + if existingResult.RowsAffected > 0 { + claimTaskID = 0 + claimRequestID = "" + } + claim := MidjourneyBillingClaim{ + ChannelID: midjourney.ChannelId, MjID: midjourney.MjId, UserID: midjourney.UserId, + BillingTaskID: claimTaskID, BillingRequestID: claimRequestID, + CreatedAt: time.Now().Unix(), + } + if err := tx.Clauses(clause.OnConflict{ + Columns: []clause.Column{{Name: "channel_id"}, {Name: "mj_id"}}, + DoNothing: true, + }).Create(&claim).Error; err != nil { + return err + } + if err := midjourneyBillingClaimForUpdate(tx, midjourney.ChannelId, midjourney.MjId, &claim).Error; err != nil { + return err + } + + billingTask.PrivateData.UpstreamTaskID = midjourney.MjId + billingTask.PrivateData.AwaitingUpstreamID = false + if claim.BillingTaskID != billingTask.ID { + sameRequestReplay := claim.BillingRequestID != "" && claim.BillingRequestID == billingTask.PrivateData.BillingRequestId + refundDuplicate = billingTask.Quota != 0 && !sameRequestReplay + if refundDuplicate && duplicateRefundInput == nil { + return errors.New("duplicate midjourney submission requires a durable refund intent") + } + if refundDuplicate { + if _, _, err := ensureBillingSettlementRecordDB(tx, *duplicateRefundInput); err != nil { + return err + } + } + billingTask.Status = TaskStatusFailure + billingTask.Progress = "100%" + billingTask.FailReason = "provider task already claimed; duplicate pre-consume refund pending" + billingTask.Quota = 0 + result := tx.Model(&Task{}). + Where("id = ? AND status NOT IN ?", billingTask.ID, []string{TaskStatusFailure, TaskStatusSuccess}). + Updates(midjourneyTaskUpdateValues(billingTask, true)) + if result.Error != nil { + return result.Error + } + if result.RowsAffected != 1 { + return taskSubmitResultUpdateMissError(tx, billingTask.ID) + } + created = false + return nil + } + + if settlementInput != nil { + if _, _, err := ensureBillingSettlementRecordDB(tx, *settlementInput); err != nil { + return err + } + } + if existingResult.RowsAffected == 0 { + if err := tx.Create(midjourney).Error; err != nil { + return err + } + } else { + midjourney.Id = existing.Id + } + billingTask.Status = TaskStatusSubmitted + billingTask.Progress = midjourney.Progress + result := tx.Model(&Task{}). + Where("id = ? AND status NOT IN ?", billingTask.ID, []string{TaskStatusFailure, TaskStatusSuccess}). + Updates(midjourneyTaskUpdateValues(billingTask, true)) + if result.Error != nil { + return result.Error + } + if result.RowsAffected != 1 { + return taskSubmitResultUpdateMissError(tx, billingTask.ID) + } + created = true + refundDuplicate = false + return nil + }) + return created, refundDuplicate, err +} + +func midjourneyBillingClaimForUpdate(tx *gorm.DB, channelID int, mjID string, claim *MidjourneyBillingClaim) *gorm.DB { + return tx.Clauses(clause.Locking{Strength: "UPDATE"}). + Where("channel_id = ? AND mj_id = ?", channelID, mjID). + First(claim) +} + +func GetMidjourneyBillingTask(userID, channelID int, mjID string) (*Task, error) { + if userID <= 0 || channelID <= 0 || mjID == "" { + return nil, nil + } + var claim MidjourneyBillingClaim + result := DB.Where("channel_id = ? AND mj_id = ?", channelID, mjID).Limit(1).Find(&claim) if result.Error != nil { - return false, result.Error + return nil, result.Error } - return result.RowsAffected > 0, nil + if result.RowsAffected == 0 || claim.UserID != userID || claim.BillingTaskID <= 0 { + return nil, nil + } + var task Task + if err := DB.First(&task, claim.BillingTaskID).Error; err != nil { + return nil, err + } + return &task, nil +} + +// UpdateWithBillingTaskAndSettlement commits a Midjourney state transition, +// its shadow Task transition, and an optional terminal settlement intent in one +// transaction. Balance application happens after commit or through the durable +// settlement runner. +func (midjourney *Midjourney) UpdateWithBillingTaskAndSettlement(fromStatus string, billingTask *Task, input *BillingSettlementInput) (bool, error) { + if midjourney == nil || midjourney.Id <= 0 { + return false, errors.New("persisted midjourney task is required") + } + if input != nil { + if billingTask == nil || billingTask.ID <= 0 { + return false, errors.New("billing task is required for midjourney settlement") + } + if input.TaskID != billingTask.ID { + return false, fmt.Errorf("billing settlement task identity mismatch: task=%d input=%d", billingTask.ID, input.TaskID) + } + if err := validateBillingSettlementInput(*input); err != nil { + return false, err + } + } + + err := DB.Transaction(func(tx *gorm.DB) error { + if input != nil { + if _, _, err := ensureBillingSettlementRecordDB(tx, *input); err != nil { + return err + } + } + result := tx.Model(&Midjourney{}). + Where("id = ? AND status = ?", midjourney.Id, fromStatus). + Updates(midjourneyUpdateValues(midjourney)) + if result.Error != nil { + return result.Error + } + if result.RowsAffected == 0 { + var existing Midjourney + if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).Select("status").First(&existing, midjourney.Id).Error; err != nil { + return err + } + if existing.Status != fromStatus { + return errTaskStatusCASLost + } + } + if billingTask == nil { + return nil + } + result = tx.Model(&Task{}). + Where("id = ? AND (status NOT IN ? OR status = ?)", billingTask.ID, []string{TaskStatusFailure, TaskStatusSuccess}, billingTask.Status). + Updates(midjourneyTaskUpdateValues(billingTask, false)) + if result.Error != nil { + return result.Error + } + if result.RowsAffected != 1 { + if result.RowsAffected != 0 { + return fmt.Errorf("midjourney billing task did not advance: id=%d", billingTask.ID) + } + var existing Task + if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).Select("status").First(&existing, billingTask.ID).Error; err != nil { + return err + } + if existing.Status != billingTask.Status { + return fmt.Errorf("midjourney billing task did not advance: id=%d", billingTask.ID) + } + } + return nil + }) + if errors.Is(err, errTaskStatusCASLost) { + return false, nil + } + return err == nil, err +} + +// UpdateWithStatus performs a conditional UPDATE guarded by fromStatus (CAS). +// It delegates to UpdateWithBillingTaskAndSettlement, which uses an explicit +// update map to avoid GORM Save()'s INSERT fallback and zero-value skipping. +func (midjourney *Midjourney) UpdateWithStatus(fromStatus string) (bool, error) { + return midjourney.UpdateWithBillingTaskAndSettlement(fromStatus, nil, nil) } func MjBulkUpdate(mjIds []string, params map[string]any) error { + sanitizeFailReasonUpdateParam(params) return DB.Model(&Midjourney{}). Where("mj_id in (?)", mjIds). Updates(params).Error } func MjBulkUpdateByTaskIds(taskIDs []int, params map[string]any) error { + sanitizeFailReasonUpdateParam(params) return DB.Model(&Midjourney{}). Where("id in (?)", taskIDs). Updates(params).Error } +func sanitizeFailReasonUpdateParam(params map[string]any) { + if params == nil { + return + } + if reason, ok := params["fail_reason"].(string); ok { + params["fail_reason"] = common.SanitizePersistedLogContent(reason) + } +} + // CountAllTasks returns total midjourney tasks for admin query func CountAllTasks(queryParams TaskQueryParams) int64 { var total int64 diff --git a/model/midjourney_billing_test.go b/model/midjourney_billing_test.go new file mode 100644 index 00000000..2612c106 --- /dev/null +++ b/model/midjourney_billing_test.go @@ -0,0 +1,345 @@ +package model + +import ( + "database/sql" + "errors" + "fmt" + "strings" + "testing" + + "github.com/MAX-API-Next/MAX-API/constant" + "github.com/glebarez/sqlite" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + gormmysql "gorm.io/driver/mysql" + gormpostgres "gorm.io/driver/postgres" + "gorm.io/gorm" +) + +func newMidjourneyBillingTask(t *testing.T, suffix string, quota int) *Task { + t.Helper() + task := &Task{ + TaskID: "task_midjourney_" + suffix, + Platform: constant.TaskPlatformMidjourney, + UserId: 501, + ChannelId: 601, + Quota: quota, + Action: constant.MjActionImagine, + Status: TaskStatusNotStart, + Progress: "0%", + PrivateData: TaskPrivateData{ + AwaitingUpstreamID: true, + BillingSource: BillingSettlementSourceWallet, + BillingRequestId: "mj-request-" + suffix, + TokenId: 701, + }, + } + insertTask(t, task) + return task +} + +func TestFinalizeMidjourneySubmissionClaimsUpstreamTaskOnce(t *testing.T) { + truncateTables(t) + firstTask := newMidjourneyBillingTask(t, "first", 20) + secondTask := newMidjourneyBillingTask(t, "second", 20) + + first := &Midjourney{ + UserId: 501, ChannelId: 601, MjId: "provider-mj-1", Action: constant.MjActionImagine, + Status: "", Progress: "0%", Quota: 20, + } + created, refundDuplicate, err := FinalizeMidjourneySubmission(first, firstTask, + &BillingSettlementInput{ + OperationKey: "request:mj-request-first:finalize", Source: BillingSettlementSourceWallet, + UserID: 501, TokenID: 701, + }, + &BillingSettlementInput{ + OperationKey: "request:mj-request-first:refund", Source: BillingSettlementSourceWallet, + UserID: 501, TokenID: 701, FundingDelta: -20, TokenDelta: -20, + }, + ) + require.NoError(t, err) + require.True(t, created) + require.False(t, refundDuplicate) + + duplicate := &Midjourney{ + UserId: 501, ChannelId: 601, MjId: "provider-mj-1", Action: constant.MjActionImagine, + Status: "", Progress: "0%", Quota: 20, + } + created, refundDuplicate, err = FinalizeMidjourneySubmission(duplicate, secondTask, + &BillingSettlementInput{ + OperationKey: "request:mj-request-second:finalize", Source: BillingSettlementSourceWallet, + UserID: 501, TokenID: 701, + }, + &BillingSettlementInput{ + OperationKey: "request:mj-request-second:refund", Source: BillingSettlementSourceWallet, + UserID: 501, TokenID: 701, FundingDelta: -20, TokenDelta: -20, + }, + ) + require.NoError(t, err) + require.False(t, created) + require.True(t, refundDuplicate) + + var count int64 + require.NoError(t, DB.Model(&Midjourney{}).Where("user_id = ? AND mj_id = ?", 501, "provider-mj-1").Count(&count).Error) + assert.EqualValues(t, 1, count) + + var claim MidjourneyBillingClaim + require.NoError(t, DB.Where("channel_id = ? AND mj_id = ?", 601, "provider-mj-1").First(&claim).Error) + assert.Equal(t, firstTask.ID, claim.BillingTaskID) + + var reloadedSecond Task + require.NoError(t, DB.First(&reloadedSecond, secondTask.ID).Error) + assert.EqualValues(t, TaskStatusFailure, reloadedSecond.Status) + assert.Zero(t, reloadedSecond.Quota) + + var refund BillingSettlement + require.NoError(t, DB.Where("operation_key = ?", "request:mj-request-second:refund").First(&refund).Error) + assert.Equal(t, BillingSettlementStatusPending, refund.Status) + assert.EqualValues(t, -20, refund.FundingDelta) +} + +func TestMidjourneyUpdateTreatsUnchangedMatchedRowAsCASWin(t *testing.T) { + truncateTables(t) + midjourney := &Midjourney{ + UserId: 501, ChannelId: 601, MjId: "provider-mj-unchanged-cas", + Status: "IN_PROGRESS", Progress: "50%", Code: 1, + } + require.NoError(t, DB.Create(midjourney).Error) + + callbackName := "test:force-midjourney-zero-rows-affected" + require.NoError(t, DB.Callback().Update().After("gorm:update").Register(callbackName, func(tx *gorm.DB) { + if tx.Statement.Table == "midjourneys" { + tx.RowsAffected = 0 + } + })) + t.Cleanup(func() { _ = DB.Callback().Update().Remove(callbackName) }) + + won, err := midjourney.UpdateWithStatus("IN_PROGRESS") + require.NoError(t, err) + require.True(t, won) +} + +func TestMidjourneyUpdateTreatsAlreadyTerminalBillingTaskAsCASWin(t *testing.T) { + truncateTables(t) + midjourney := &Midjourney{ + UserId: 501, ChannelId: 601, MjId: "provider-mj-terminal-billing-cas", + Status: "IN_PROGRESS", Progress: "50%", Code: 1, + } + require.NoError(t, DB.Create(midjourney).Error) + billingTask := &Task{ + TaskID: "task_midjourney_terminal_billing_cas", Platform: constant.TaskPlatformMidjourney, + UserId: 501, ChannelId: 601, Status: TaskStatusSuccess, Progress: "100%", + } + insertTask(t, billingTask) + midjourney.Status = "SUCCESS" + midjourney.Progress = "100%" + + callbackName := "test:force-midjourney-billing-task-zero-rows-affected" + require.NoError(t, DB.Callback().Update().After("gorm:update").Register(callbackName, func(tx *gorm.DB) { + if tx.Statement.Table == "tasks" { + tx.RowsAffected = 0 + } + })) + t.Cleanup(func() { _ = DB.Callback().Update().Remove(callbackName) }) + + won, err := midjourney.UpdateWithBillingTaskAndSettlement("IN_PROGRESS", billingTask, nil) + require.NoError(t, err) + require.True(t, won) +} + +func TestMidjourneyBillingClaimLookupUsesCurrentReadAcrossDialects(t *testing.T) { + tests := []struct { + name string + open func(*testing.T) *gorm.DB + wantForUpdate bool + }{ + { + name: "sqlite", + open: func(t *testing.T) *gorm.DB { + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{DryRun: true}) + require.NoError(t, err) + return db + }, + }, + { + name: "mysql", + open: func(t *testing.T) *gorm.DB { + conn, err := sql.Open("mysql", "") + require.NoError(t, err) + t.Cleanup(func() { _ = conn.Close() }) + db, err := gorm.Open(gormmysql.New(gormmysql.Config{Conn: conn, SkipInitializeWithVersion: true}), &gorm.Config{ + DryRun: true, DisableAutomaticPing: true, + }) + require.NoError(t, err) + return db + }, + wantForUpdate: true, + }, + { + name: "postgres", + open: func(t *testing.T) *gorm.DB { + conn, err := sql.Open("pgx", "") + require.NoError(t, err) + t.Cleanup(func() { _ = conn.Close() }) + db, err := gorm.Open(gormpostgres.New(gormpostgres.Config{Conn: conn}), &gorm.Config{ + DryRun: true, DisableAutomaticPing: true, + }) + require.NoError(t, err) + return db + }, + wantForUpdate: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + db := tt.open(t) + statement := db.ToSQL(func(tx *gorm.DB) *gorm.DB { + var claim MidjourneyBillingClaim + return midjourneyBillingClaimForUpdate(tx, 601, "provider-mj-current-read", &claim) + }) + hasForUpdate := strings.Contains(strings.ToUpper(statement), "FOR UPDATE") + assert.Equal(t, tt.wantForUpdate, hasForUpdate, statement) + }) + } +} + +func TestFinalizeMidjourneySubmissionDoesNotRefundSameRequestReplay(t *testing.T) { + truncateTables(t) + firstTask := newMidjourneyBillingTask(t, "replay-first", 20) + firstTask.PrivateData.BillingRequestId = "mj-request-replay" + require.NoError(t, DB.Model(&Task{}).Where("id = ?", firstTask.ID).Update("private_data", firstTask.PrivateData).Error) + mj := &Midjourney{ + UserId: 501, ChannelId: 601, MjId: "provider-mj-replay", Action: constant.MjActionImagine, + Status: "", Progress: "0%", Quota: 20, + } + created, refundDuplicate, err := FinalizeMidjourneySubmission(mj, firstTask, + &BillingSettlementInput{ + OperationKey: "request:mj-request-replay:finalize", Source: BillingSettlementSourceWallet, + UserID: 501, TokenID: 701, + }, nil) + require.NoError(t, err) + require.True(t, created) + require.False(t, refundDuplicate) + + replayTask := newMidjourneyBillingTask(t, "replay-second", 20) + replayTask.PrivateData.BillingRequestId = "mj-request-replay" + require.NoError(t, DB.Model(&Task{}).Where("id = ?", replayTask.ID).Update("private_data", replayTask.PrivateData).Error) + created, refundDuplicate, err = FinalizeMidjourneySubmission(&Midjourney{ + UserId: 501, ChannelId: 601, MjId: mj.MjId, Action: constant.MjActionImagine, + Status: "", Progress: "0%", Quota: 20, + }, replayTask, + &BillingSettlementInput{ + OperationKey: "request:mj-request-replay:finalize", Source: BillingSettlementSourceWallet, + UserID: 501, TokenID: 701, + }, + &BillingSettlementInput{ + OperationKey: "request:mj-request-replay:finalize", Source: BillingSettlementSourceWallet, + UserID: 501, TokenID: 701, FundingDelta: -20, TokenDelta: -20, + }, + ) + require.NoError(t, err) + require.False(t, created) + require.False(t, refundDuplicate) + + var settlement BillingSettlement + require.NoError(t, DB.Where("operation_key = ?", "request:mj-request-replay:finalize").First(&settlement).Error) + assert.Zero(t, settlement.FundingDelta) +} + +func TestMidjourneyFailureCommitsRefundIntentWithTerminalState(t *testing.T) { + truncateTables(t) + billingTask := newMidjourneyBillingTask(t, "failure", 30) + mj := &Midjourney{ + UserId: 501, ChannelId: 601, MjId: "provider-mj-failure", Action: constant.MjActionImagine, + Status: "", Progress: "0%", Quota: 30, + } + created, refundDuplicate, err := FinalizeMidjourneySubmission(mj, billingTask, + &BillingSettlementInput{ + OperationKey: "request:mj-request-failure:finalize", Source: BillingSettlementSourceWallet, + UserID: 501, TokenID: 701, + }, + nil, + ) + require.NoError(t, err) + require.True(t, created) + require.False(t, refundDuplicate) + + loadedTask, err := GetMidjourneyBillingTask(mj.UserId, mj.ChannelId, mj.MjId) + require.NoError(t, err) + require.NotNil(t, loadedTask) + + fromStatus := mj.Status + mj.Status = "FAILURE" + mj.Progress = "100%" + mj.FailReason = "provider failed" + loadedTask.Status = TaskStatusFailure + loadedTask.Progress = "100%" + loadedTask.FailReason = mj.FailReason + operationKey := fmt.Sprintf("task:%d:refund", loadedTask.ID) + won, err := mj.UpdateWithBillingTaskAndSettlement(fromStatus, loadedTask, &BillingSettlementInput{ + OperationKey: operationKey, Source: BillingSettlementSourceWallet, + UserID: 501, TokenID: 701, FundingDelta: -30, TokenDelta: -30, + TaskID: loadedTask.ID, TaskQuota: 30, TaskQuotaTarget: 0, + }) + require.NoError(t, err) + require.True(t, won) + + var settlement BillingSettlement + require.NoError(t, DB.Where("operation_key = ?", operationKey).First(&settlement).Error) + assert.Equal(t, BillingSettlementStatusPending, settlement.Status) + + var storedMJ Midjourney + require.NoError(t, DB.First(&storedMJ, mj.Id).Error) + assert.Equal(t, "FAILURE", storedMJ.Status) + var storedTask Task + require.NoError(t, DB.First(&storedTask, loadedTask.ID).Error) + assert.EqualValues(t, TaskStatusFailure, storedTask.Status) +} + +func TestUniqueMidjourneyLookupRejectsCrossChannelAmbiguity(t *testing.T) { + truncateTables(t) + sharedID := "provider-mj-shared-across-channels" + require.NoError(t, DB.Create(&Midjourney{ + UserId: 501, ChannelId: 601, MjId: sharedID, Status: "SUCCESS", Progress: "100%", + }).Error) + require.NoError(t, DB.Create(&Midjourney{ + UserId: 501, ChannelId: 602, MjId: sharedID, Status: "SUCCESS", Progress: "100%", + }).Error) + + _, err := GetUniqueMidjourneyByUserAndMJID(501, sharedID) + require.ErrorIs(t, err, ErrMidjourneyTaskAmbiguous) + _, err = GetUniqueMidjourneyByMJID(sharedID) + require.True(t, errors.Is(err, ErrMidjourneyTaskAmbiguous)) +} + +func TestMidjourneyTerminalMetadataRefreshKeepsMatchingTaskTerminalState(t *testing.T) { + truncateTables(t) + billingTask := newMidjourneyBillingTask(t, "terminal-refresh", 20) + billingTask.Status = TaskStatusSuccess + billingTask.Progress = "100%" + billingTask.PrivateData.ResultURL = "https://example.com/old.png" + require.NoError(t, DB.Model(&Task{}).Where("id = ?", billingTask.ID).Updates(map[string]any{ + "status": billingTask.Status, "progress": billingTask.Progress, "private_data": billingTask.PrivateData, + }).Error) + mj := Midjourney{ + UserId: 501, ChannelId: 601, MjId: "provider-mj-terminal-refresh", + Status: "SUCCESS", Progress: "100%", ImageUrl: "https://example.com/old.png", Quota: 20, + } + require.NoError(t, DB.Create(&mj).Error) + + mj.ImageUrl = "https://example.com/final.png" + billingTask.PrivateData.ResultURL = mj.ImageUrl + billingTask.SetData(map[string]any{"imageUrl": mj.ImageUrl}) + won, err := mj.UpdateWithBillingTaskAndSettlement("SUCCESS", billingTask, nil) + require.NoError(t, err) + require.True(t, won) + + var storedTask Task + require.NoError(t, DB.First(&storedTask, billingTask.ID).Error) + require.EqualValues(t, TaskStatusSuccess, storedTask.Status) + require.Equal(t, mj.ImageUrl, storedTask.PrivateData.ResultURL) + var storedMJ Midjourney + require.NoError(t, DB.First(&storedMJ, mj.Id).Error) + require.Equal(t, mj.ImageUrl, storedMJ.ImageUrl) +} diff --git a/model/option.go b/model/option.go index 609f2541..9dc69724 100644 --- a/model/option.go +++ b/model/option.go @@ -354,7 +354,11 @@ func validateOptionUpdate(key string, value string) error { return validateJSONOption[map[string]float64](value) case "ModelRequestRateLimitGroup": return setting.CheckModelRequestRateLimitGroup(value) - case "ModelRatio", "ModelPrice", "CacheRatio", "CreateCacheRatio", "CompletionRatio", "ImageRatio", "AudioRatio", "AudioCompletionRatio": + case "ModelRatio": + return ratio_setting.ValidateModelRatioJSONString(value) + case "CompletionRatio": + return ratio_setting.ValidateCompletionRatioJSONString(value) + case "ModelPrice", "CacheRatio", "CreateCacheRatio", "ImageRatio", "AudioRatio", "AudioCompletionRatio": return ratio_setting.ValidatePricingMapJSONString(value) case "GroupRatio", "group_ratio_setting.group_ratio": return ratio_setting.CheckGroupRatio(value) diff --git a/model/option_test.go b/model/option_test.go index 8c9ce17b..9dbdc9e1 100644 --- a/model/option_test.go +++ b/model/option_test.go @@ -1,6 +1,8 @@ package model import ( + "fmt" + "strings" "testing" "github.com/MAX-API-Next/MAX-API/common" @@ -206,6 +208,27 @@ func TestValidateOptionUpdateRejectsUnsafePricingValues(t *testing.T) { } } +func TestValidateOptionUpdateRejectsOversizedPricingMaps(t *testing.T) { + const entryCount = 20001 + values := make(map[string]float64, entryCount) + for i := 0; i < entryCount; i++ { + values[fmt.Sprintf("model-%d", i)] = 1 + } + data, err := common.Marshal(values) + require.NoError(t, err) + + require.Error(t, validateOptionUpdate("ModelRatio", string(data))) +} + +func TestValidateOptionUpdateRejectsUnsafePricingKeys(t *testing.T) { + longModelName := strings.Repeat("a", 257) + + require.Error(t, validateOptionUpdate("CompletionRatio", fmt.Sprintf(`{"%s":1}`, longModelName))) + require.Error(t, validateOptionUpdate("CompletionRatio", `{" ":1}`)) + require.Error(t, validateOptionUpdate("ModelRatio", `{"gemini-2.5-flash-thinking-a":1,"gemini-2.5-flash-thinking-b":2}`)) + require.Error(t, validateOptionUpdate("CompletionRatio", `{"gemini-2.5-pro-thinking-a":1,"gemini-2.5-pro-thinking-b":2}`)) +} + func TestValidateOptionUpdateRejectsNullRWMapConfigs(t *testing.T) { for _, key := range []string{ "billing_setting.billing_mode", diff --git a/model/payment_method_guard_test.go b/model/payment_method_guard_test.go index 30107b82..3b10fc22 100644 --- a/model/payment_method_guard_test.go +++ b/model/payment_method_guard_test.go @@ -333,6 +333,105 @@ func TestRechargeCreemRejectsZeroQuotaBeforeCompletingOrder(t *testing.T) { assert.EqualValues(t, 0, getUserQuotaForPaymentGuardTest(t, 610)) } +func TestRechargeRejectsPaidAmountMismatchBeforeCompletingOrder(t *testing.T) { + truncateTables(t) + + insertUserForPaymentGuardTest(t, 613, 0) + topUp := &TopUp{ + UserId: 613, + Amount: 10, + Money: 9.99, + TradeNo: "stripe-amount-mismatch", + PaymentMethod: PaymentMethodStripe, + PaymentProvider: PaymentProviderStripe, + Status: common.TopUpStatusPending, + CreateTime: time.Now().Unix(), + } + require.NoError(t, topUp.Insert()) + + err := Recharge("stripe-amount-mismatch", "cus_test", "127.0.0.1", + PaymentValidationFromMinorUnits(998, "USD", "USD", false)) + require.ErrorIs(t, err, ErrPaymentAmountMismatch) + assert.Equal(t, common.TopUpStatusPending, getTopUpStatusForPaymentGuardTest(t, "stripe-amount-mismatch")) + assert.EqualValues(t, 0, getUserQuotaForPaymentGuardTest(t, 613)) +} + +func TestRechargeEpayCompletesOrderAndCreditsQuotaAtomically(t *testing.T) { + truncateTables(t) + + insertUserForPaymentGuardTest(t, 614, 0) + topUp := &TopUp{ + UserId: 614, + Amount: 3, + Money: 6.66, + TradeNo: "epay-atomic-success", + PaymentMethod: "alipay", + PaymentProvider: PaymentProviderEpay, + Status: common.TopUpStatusPending, + CreateTime: time.Now().Unix(), + } + require.NoError(t, topUp.Insert()) + + require.NoError(t, RechargeEpay("epay-atomic-success", "wxpay", "127.0.0.1", + PaymentValidationFromMajorString("6.66", "", "", false))) + + reloaded := GetTopUpByTradeNo("epay-atomic-success") + require.NotNil(t, reloaded) + assert.Equal(t, common.TopUpStatusSuccess, reloaded.Status) + assert.Equal(t, "wxpay", reloaded.PaymentMethod) + assert.EqualValues(t, int64(3*common.QuotaPerUnit), getUserQuotaForPaymentGuardTest(t, 614)) +} + +func TestRechargeEpayRejectsPaidAmountMismatchBeforeCompletingOrder(t *testing.T) { + truncateTables(t) + + insertUserForPaymentGuardTest(t, 615, 0) + topUp := &TopUp{ + UserId: 615, + Amount: 3, + Money: 6.66, + TradeNo: "epay-amount-mismatch", + PaymentMethod: "alipay", + PaymentProvider: PaymentProviderEpay, + Status: common.TopUpStatusPending, + CreateTime: time.Now().Unix(), + } + require.NoError(t, topUp.Insert()) + + err := RechargeEpay("epay-amount-mismatch", "wxpay", "127.0.0.1", + PaymentValidationFromMajorString("6.65", "", "", false)) + require.ErrorIs(t, err, ErrPaymentAmountMismatch) + assert.Equal(t, common.TopUpStatusPending, getTopUpStatusForPaymentGuardTest(t, "epay-amount-mismatch")) + reloaded := GetTopUpByTradeNo("epay-amount-mismatch") + require.NotNil(t, reloaded) + assert.Equal(t, "alipay", reloaded.PaymentMethod) + assert.EqualValues(t, 0, getUserQuotaForPaymentGuardTest(t, 615)) +} + +func TestCompleteSubscriptionOrderRejectsPaidCurrencyMismatch(t *testing.T) { + truncateTables(t) + + insertUserForPaymentGuardTest(t, 616, 0) + plan := insertSubscriptionPlanForPaymentGuardTest(t, 616) + insertSubscriptionOrderForPaymentGuardTest(t, "sub-currency-mismatch", 616, plan.Id, PaymentProviderStripe) + + err := CompleteSubscriptionOrder("sub-currency-mismatch", `{"provider":"stripe"}`, PaymentProviderStripe, "card", + PaymentValidationFromMinorUnits(999, "EUR", "USD", false)) + require.ErrorIs(t, err, ErrPaymentCurrencyMismatch) + insertSubscriptionOrderForPaymentGuardTest(t, "sub-plan-currency-mismatch", 616, plan.Id, PaymentProviderStripe) + err = CompleteSubscriptionOrder("sub-plan-currency-mismatch", `{"provider":"stripe"}`, PaymentProviderStripe, "card", + PaymentValidationFromMinorUnits(999, "EUR", "", false)) + require.ErrorIs(t, err, ErrPaymentCurrencyMismatch) + + order := GetSubscriptionOrderByTradeNo("sub-currency-mismatch") + require.NotNil(t, order) + assert.Equal(t, common.TopUpStatusPending, order.Status) + planCurrencyOrder := GetSubscriptionOrderByTradeNo("sub-plan-currency-mismatch") + require.NotNil(t, planCurrencyOrder) + assert.Equal(t, common.TopUpStatusPending, planCurrencyOrder.Status) + assert.Zero(t, countUserSubscriptionsForPaymentGuardTest(t, 616)) +} + func TestRechargeCreemSkipsDuplicateCustomerEmailBinding(t *testing.T) { truncateTables(t) diff --git a/model/payment_validation.go b/model/payment_validation.go new file mode 100644 index 00000000..33ba1b72 --- /dev/null +++ b/model/payment_validation.go @@ -0,0 +1,106 @@ +package model + +import ( + "errors" + "strings" + + "github.com/shopspring/decimal" +) + +var ( + ErrPaymentAmountMismatch = errors.New("payment amount mismatch") + ErrPaymentCurrencyMismatch = errors.New("payment currency mismatch") +) + +type PaymentValidation struct { + HasAmount bool + PaidAmountMinor int64 + Currency string + ExpectedCurrency string + AllowDiscount bool +} + +func PaymentValidationFromMinorUnits(amountMinor int64, currency string, expectedCurrency string, allowDiscount bool) PaymentValidation { + return PaymentValidation{ + HasAmount: true, + PaidAmountMinor: amountMinor, + Currency: normalizePaymentCurrency(currency), + ExpectedCurrency: normalizePaymentCurrency(expectedCurrency), + AllowDiscount: allowDiscount, + } +} + +func PaymentValidationFromMajorString(amount string, currency string, expectedCurrency string, allowDiscount bool) PaymentValidation { + paid, err := decimal.NewFromString(strings.TrimSpace(amount)) + if err != nil { + return PaymentValidation{ + HasAmount: true, + PaidAmountMinor: -1, + Currency: normalizePaymentCurrency(currency), + ExpectedCurrency: normalizePaymentCurrency(expectedCurrency), + AllowDiscount: allowDiscount, + } + } + return PaymentValidationFromMinorUnits(paid.Mul(decimal.NewFromInt(100)).Round(0).IntPart(), currency, expectedCurrency, allowDiscount) +} + +func normalizePaymentCurrency(currency string) string { + return strings.ToUpper(strings.TrimSpace(currency)) +} + +func validatePaymentAgainstExpectedAmount(expectedAmount float64, validation PaymentValidation) error { + if validation.ExpectedCurrency != "" && validation.Currency != "" && validation.Currency != validation.ExpectedCurrency { + return ErrPaymentCurrencyMismatch + } + if !validation.HasAmount { + return nil + } + if validation.PaidAmountMinor <= 0 { + return ErrPaymentAmountMismatch + } + expectedMinor := decimal.NewFromFloat(expectedAmount).Mul(decimal.NewFromInt(100)).Round(0).IntPart() + if expectedMinor <= 0 { + return ErrPaymentAmountMismatch + } + if validation.AllowDiscount { + if validation.PaidAmountMinor > expectedMinor { + return ErrPaymentAmountMismatch + } + return nil + } + if validation.PaidAmountMinor != expectedMinor { + return ErrPaymentAmountMismatch + } + return nil +} + +func validatePaymentAgainstTopUp(topUp *TopUp, validations []PaymentValidation) error { + if topUp == nil { + return ErrTopUpNotFound + } + for _, validation := range validations { + if err := validatePaymentAgainstExpectedAmount(topUp.Money, validation); err != nil { + return err + } + } + return nil +} + +func validatePaymentAgainstSubscriptionOrder(order *SubscriptionOrder, plan *SubscriptionPlan, validations []PaymentValidation) error { + if order == nil { + return ErrSubscriptionOrderNotFound + } + expectedCurrency := "" + if plan != nil { + expectedCurrency = normalizePaymentCurrency(plan.Currency) + } + for _, validation := range validations { + if validation.ExpectedCurrency == "" { + validation.ExpectedCurrency = expectedCurrency + } + if err := validatePaymentAgainstExpectedAmount(order.Money, validation); err != nil { + return err + } + } + return nil +} diff --git a/model/subscription.go b/model/subscription.go index 77dfd88f..04db9e92 100644 --- a/model/subscription.go +++ b/model/subscription.go @@ -542,6 +542,11 @@ func CreateUserSubscriptionFromPlanTx(tx *gorm.DB, userId int, plan *Subscriptio return nil, errors.New("invalid user id") } if plan.MaxPurchasePerUser > 0 { + var lockedUser User + if err := withRowLock(tx).Select("id").Where("id = ?", userId).First(&lockedUser).Error; err != nil { + return nil, err + } + var count int64 if err := tx.Model(&UserSubscription{}). Where("user_id = ? AND plan_id = ?", userId, plan.Id). @@ -604,7 +609,7 @@ func CreateUserSubscriptionFromPlanTx(tx *gorm.DB, userId int, plan *Subscriptio // Complete a subscription order (idempotent). Creates a UserSubscription snapshot from the plan. // expectedPaymentProvider guards against cross-gateway callback attacks (empty skips the check). // actualPaymentMethod updates the order's PaymentMethod to reflect the real payment type used (empty skips update). -func CompleteSubscriptionOrder(tradeNo string, providerPayload string, expectedPaymentProvider string, actualPaymentMethod string) error { +func CompleteSubscriptionOrder(tradeNo string, providerPayload string, expectedPaymentProvider string, actualPaymentMethod string, validations ...PaymentValidation) error { if tradeNo == "" { return errors.New("tradeNo is empty") } @@ -641,6 +646,9 @@ func CompleteSubscriptionOrder(tradeNo string, providerPayload string, expectedP if !plan.Enabled { // still allow completion for already purchased orders } + if err := validatePaymentAgainstSubscriptionOrder(&order, plan, validations); err != nil { + return err + } upgradeGroup = strings.TrimSpace(plan.UpgradeGroup) claimed, err := completePendingSubscriptionOrderTx(tx, &order, providerPayload, actualPaymentMethod) if err != nil { @@ -776,12 +784,16 @@ func PurchaseSubscriptionWithBalance(userId int, planId int) error { if userId <= 0 || planId <= 0 { return errors.New("invalid userId or planId") } + randomPart, err := common.GenerateRandomCharsKey(6) + if err != nil { + return fmt.Errorf("generate subscription order number: %w", err) + } var logPlanTitle string var logMoney float64 var chargedQuota int var upgradeGroup string - err := DB.Transaction(func(tx *gorm.DB) error { + err = DB.Transaction(func(tx *gorm.DB) error { plan, err := getSubscriptionPlanByIdTx(tx, planId) if err != nil { return err @@ -824,7 +836,7 @@ func PurchaseSubscriptionWithBalance(userId int, planId int) error { } now := common.GetTimestamp() - tradeNo := fmt.Sprintf("SUBBALUSR%dNO%s%d", userId, common.GetRandomString(6), time.Now().UnixNano()) + tradeNo := fmt.Sprintf("SUBBALUSR%dNO%s%d", userId, randomPart, time.Now().UnixNano()) order := &SubscriptionOrder{ UserId: userId, PlanId: plan.Id, diff --git a/model/task.go b/model/task.go index 1fe3624d..dd3c1302 100644 --- a/model/task.go +++ b/model/task.go @@ -72,6 +72,13 @@ type Task struct { includePrivateDataInUpdate bool `gorm:"-"` } +func (t *Task) BeforeSave(*gorm.DB) error { + if t != nil { + t.FailReason = common.SanitizePersistedLogContent(t.FailReason) + } + return nil +} + func (t *Task) SetData(data any) { b, _ := common.Marshal(data) t.Data = json.RawMessage(b) @@ -481,7 +488,7 @@ func (t *Task) statusUpdateValues() map[string]interface{} { "progress": t.Progress, "start_time": t.StartTime, "finish_time": t.FinishTime, - "fail_reason": t.FailReason, + "fail_reason": common.SanitizePersistedLogContent(t.FailReason), "updated_at": updatedAt, } if t.Data != nil || t.includeDataInUpdate { @@ -617,6 +624,7 @@ func MarkTaskSubmitFailed(taskID int64, reason string) error { if taskID <= 0 { return nil } + reason = common.SanitizePersistedLogContent(reason) return DB.Model(&Task{}). Where("id = ? AND status NOT IN ?", taskID, []string{TaskStatusFailure, TaskStatusSuccess}). Updates(map[string]interface{}{ @@ -629,6 +637,49 @@ func MarkTaskSubmitFailed(taskID int64, reason string) error { }).Error } +// MarkTaskSubmitFailedWithSettlement commits a definite submit rejection and +// its durable pre-consume refund intent atomically. The task is terminal with a +// zero displayed quota while the settlement runner applies or retries the +// already-recorded balance mutation. +func MarkTaskSubmitFailedWithSettlement(taskID int64, reason string, input *BillingSettlementInput) error { + if taskID <= 0 { + if input != nil { + return fmt.Errorf("cannot record settlement intent for unpersisted task: id=%d", taskID) + } + return nil + } + reason = common.SanitizePersistedLogContent(reason) + if input != nil { + if err := validateBillingSettlementInput(*input); err != nil { + return err + } + } + return DB.Transaction(func(tx *gorm.DB) error { + if input != nil { + if _, _, err := ensureBillingSettlementRecordDB(tx, *input); err != nil { + return err + } + } + result := tx.Model(&Task{}). + Where("id = ? AND status NOT IN ?", taskID, []string{TaskStatusFailure, TaskStatusSuccess}). + Updates(map[string]interface{}{ + "status": TaskStatusFailure, + "progress": "100%", + "finish_time": time.Now().Unix(), + "fail_reason": reason, + "quota": 0, + "updated_at": time.Now().Unix(), + }) + if result.Error != nil { + return result.Error + } + if result.RowsAffected != 1 { + return taskSubmitResultUpdateMissError(tx, taskID) + } + return nil + }) +} + // MarkTaskSubmitNeedsReview preserves the pre-consumed quota and any upstream // identity after the provider accepted a task but local durable finalization // failed. Refunding this state would create an unbilled upstream task. @@ -636,6 +687,7 @@ func MarkTaskSubmitNeedsReview(task *Task, reason string) error { if task == nil || task.ID <= 0 { return nil } + reason = common.SanitizePersistedLogContent(reason) task.PrivateData.AwaitingUpstreamID = false return DB.Model(&Task{}). Where("id = ? AND status <> ?", task.ID, TaskStatusSuccess). @@ -655,6 +707,7 @@ func MarkTaskSubmitAmbiguous(taskID int64, reason string) error { if taskID <= 0 { return nil } + reason = common.SanitizePersistedLogContent(reason) return DB.Model(&Task{}). Where("id = ? AND status NOT IN ?", taskID, []string{TaskStatusFailure, TaskStatusSuccess}). Updates(map[string]interface{}{ @@ -672,6 +725,7 @@ func TaskBulkUpdate(taskIds []string, params map[string]any) error { if len(taskIds) == 0 { return nil } + sanitizeFailReasonUpdateParam(params) return DB.Model(&Task{}). Where("task_id in (?)", taskIds). Updates(params).Error @@ -686,6 +740,7 @@ func TaskBulkUpdateByID(ids []int64, params map[string]any) error { if len(ids) == 0 { return nil } + sanitizeFailReasonUpdateParam(params) return DB.Model(&Task{}). Where("id in (?)", ids). Updates(params).Error diff --git a/model/task_cas_test.go b/model/task_cas_test.go index 1ee7dcc5..47057a5e 100644 --- a/model/task_cas_test.go +++ b/model/task_cas_test.go @@ -52,6 +52,8 @@ func TestMain(m *testing.M) { &SubscriptionPreConsumeRecord{}, &BillingSettlement{}, &BillingPreConsumeSelection{}, + &Midjourney{}, + &MidjourneyBillingClaim{}, &CacheInvalidationTask{}, &PerfMetric{}, &SystemTask{}, @@ -81,6 +83,8 @@ func truncateTables(t *testing.T) { DB.Exec("DELETE FROM subscription_pre_consume_records") DB.Exec("DELETE FROM billing_settlements") DB.Exec("DELETE FROM billing_pre_consume_selections") + DB.Exec("DELETE FROM midjourneys") + DB.Exec("DELETE FROM midjourney_billing_claims") DB.Exec("DELETE FROM cache_invalidation_tasks") DB.Exec("DELETE FROM perf_metrics") DB.Exec("DELETE FROM system_tasks") @@ -259,6 +263,25 @@ func TestUpdateWithStatusAndSettlementCASLossLeavesNoIntent(t *testing.T) { require.Zero(t, count) } +func TestMarkTaskSubmitFailedWithSettlementRejectsUnpersistedTask(t *testing.T) { + truncateTables(t) + operationKey := "request:unpersisted-midjourney:refund" + require.NoError(t, MarkTaskSubmitFailedWithSettlement(0, "submit rejected", nil)) + + err := MarkTaskSubmitFailedWithSettlement(0, "submit rejected", &BillingSettlementInput{ + OperationKey: operationKey, + Source: BillingSettlementSourceWallet, + UserID: 1, + FundingDelta: -10, + TokenDelta: -10, + }) + require.Error(t, err) + + var count int64 + require.NoError(t, DB.Model(&BillingSettlement{}).Where("operation_key = ?", operationKey).Count(&count).Error) + require.Zero(t, count) +} + func TestUpdateWithSettlementIntentCommitsTaskAndIntentTogether(t *testing.T) { truncateTables(t) task := &Task{ diff --git a/model/token.go b/model/token.go index f2876e8e..4f8603eb 100644 --- a/model/token.go +++ b/model/token.go @@ -9,6 +9,7 @@ import ( "github.com/MAX-API-Next/MAX-API/setting/operation_setting" "github.com/bytedance/gopkg/util/gopool" "gorm.io/gorm" + "gorm.io/gorm/clause" ) type Token struct { @@ -289,7 +290,11 @@ func GetTokenByKey(key string, fromDB bool) (token *Token, err error) { cacheVersion, err = common.RedisGetCacheVersion(getTokenCacheVersionKey(key)) cacheVersionValid = err == nil } - err = DB.Where(commonKeyCol+" = ?", key).First(&token).Error + var dbToken Token + err = DB.Where(clause.Eq{Column: clause.Column{Name: "key"}, Value: key}).First(&dbToken).Error + if err == nil { + token = &dbToken + } if err == nil && token != nil && cacheVersionValid { tokenSnapshot := *token gopool.Go(func() { @@ -498,27 +503,35 @@ func BatchDeleteTokens(ids []int, userId int) (int, error) { } tx := DB.Begin() + if tx.Error != nil { + return 0, tx.Error + } + defer tx.Rollback() var tokens []Token if err := tx.Where("user_id = ? AND id IN (?)", userId, ids).Find(&tokens).Error; err != nil { - tx.Rollback() return 0, err } + cacheTasks := make([]CacheInvalidationTask, 0, len(tokens)) + for _, token := range tokens { + if token.Key == "" { + continue + } + task, err := stageTokenCacheInvalidationTx(tx, token.Key, true) + if err != nil { + return 0, err + } + cacheTasks = append(cacheTasks, task) + } if err := tx.Where("user_id = ? AND id IN (?)", userId, ids).Delete(&Token{}).Error; err != nil { - tx.Rollback() return 0, err } if err := tx.Commit().Error; err != nil { return 0, err } - - if common.RedisEnabled { - for _, t := range tokens { - enqueueTokenCacheRetry(t.Key, true, nil) - } - } + dispatchStagedCacheInvalidations(cacheTasks) return len(tokens), nil } diff --git a/model/token_readonly_test.go b/model/token_readonly_test.go index 5572ed55..d3d79deb 100644 --- a/model/token_readonly_test.go +++ b/model/token_readonly_test.go @@ -5,6 +5,7 @@ import ( "github.com/MAX-API-Next/MAX-API/common" "github.com/stretchr/testify/require" + "gorm.io/gorm" ) func seedReadOnlyToken(t *testing.T, name string, rawKey string) *Token { @@ -49,3 +50,13 @@ func TestValidateUserTokenForReadOnlyStatusSemantics(t *testing.T) { _, err = ValidateUserTokenForReadOnly(expired.Key) require.ErrorIs(t, err, ErrTokenInvalid) } + +func TestGetTokenByKeyEmptyKeyDoesNotMatchAnotherToken(t *testing.T) { + truncateTables(t) + seedReadOnlyToken(t, "non-empty-token", "non-empty-key") + + token, err := GetTokenByKey("", true) + + require.ErrorIs(t, err, gorm.ErrRecordNotFound) + require.Nil(t, token) +} diff --git a/model/topup.go b/model/topup.go index 0486a609..9c9f4bb1 100644 --- a/model/topup.go +++ b/model/topup.go @@ -47,6 +47,13 @@ var ( ErrTopUpStatusInvalid = errors.New("topup status invalid") ) +func preservePaymentValidationError(err error) error { + if errors.Is(err, ErrPaymentAmountMismatch) || errors.Is(err, ErrPaymentCurrencyMismatch) || errors.Is(err, ErrPaymentMethodMismatch) { + return err + } + return nil +} + func updatePendingTopUpStatusTx(tx *gorm.DB, topUp *TopUp, targetStatus string) error { if tx == nil || topUp == nil || topUp.Id == 0 { return ErrTopUpNotFound @@ -152,7 +159,7 @@ func UpdatePendingTopUpStatus(tradeNo string, expectedPaymentProvider string, ta }) } -func Recharge(referenceId string, customerId string, callerIp string) (err error) { +func Recharge(referenceId string, customerId string, callerIp string, validations ...PaymentValidation) (err error) { if referenceId == "" { return errors.New("未提供支付单号") } @@ -179,6 +186,10 @@ func Recharge(referenceId string, customerId string, callerIp string) (err error return errors.New("充值订单状态错误") } + if err := validatePaymentAgainstTopUp(topUp, validations); err != nil { + return err + } + quota = topUp.Money * common.QuotaPerUnit if quota <= 0 { return errors.New("无效的充值额度") @@ -197,6 +208,9 @@ func Recharge(referenceId string, customerId string, callerIp string) (err error if err != nil { common.SysError("topup failed: " + err.Error()) + if preserved := preservePaymentValidationError(err); preserved != nil { + return preserved + } return errors.New("充值失败，请稍后重试") } @@ -213,6 +227,85 @@ func topUpQueryCutoff() int64 { return common.GetTimestamp() - topUpQueryWindowSeconds } +func RechargeEpay(tradeNo string, actualPaymentMethod string, callerIp string, validations ...PaymentValidation) (err error) { + if tradeNo == "" { + return errors.New("missing payment order") + } + + var quotaToAdd int64 + topUp := &TopUp{} + + refCol := "`trade_no`" + if common.UsingPostgreSQL { + refCol = `"trade_no"` + } + + err = DB.Transaction(func(tx *gorm.DB) error { + if err := withRowLock(tx).Where(refCol+" = ?", tradeNo).First(topUp).Error; err != nil { + return errors.New("top-up order not found") + } + + if topUp.PaymentProvider != PaymentProviderEpay { + return ErrPaymentMethodMismatch + } + + if topUp.Status == common.TopUpStatusSuccess { + return nil + } + + if topUp.Status != common.TopUpStatusPending { + return ErrTopUpStatusInvalid + } + + if err := validatePaymentAgainstTopUp(topUp, validations); err != nil { + return err + } + + quotaToAdd = decimal.NewFromInt(topUp.Amount).Mul(decimal.NewFromFloat(common.QuotaPerUnit)).IntPart() + if quotaToAdd <= 0 { + return errors.New("invalid top-up quota") + } + + if actualPaymentMethod != "" && topUp.PaymentMethod != actualPaymentMethod { + result := tx.Model(&TopUp{}). + Where("id = ? AND status = ?", topUp.Id, common.TopUpStatusPending). + Update("payment_method", actualPaymentMethod) + if result.Error != nil { + return result.Error + } + if result.RowsAffected != 1 { + return ErrTopUpStatusInvalid + } + topUp.PaymentMethod = actualPaymentMethod + } + + if err := markPendingTopUpSuccessTx(tx, topUp); err != nil { + return err + } + + result := tx.Model(&User{}).Where("id = ?", topUp.UserId).Update("quota", gorm.Expr("quota + ?", quotaToAdd)) + if err := ensureUserUpdateMatchedTx(tx, result, topUp.UserId, errors.New("top-up user not found")); err != nil { + return err + } + + return nil + }) + + if err != nil { + common.SysError("epay topup failed: " + err.Error()) + if preserved := preservePaymentValidationError(err); preserved != nil { + return preserved + } + return errors.New("top-up failed, please retry later") + } + + if quotaToAdd > 0 { + RecordTopupLog(topUp.UserId, fmt.Sprintf("使用在线充值成功，充值额度: %v，支付金额：%.2f", logger.FormatQuota(quotaToAdd), topUp.Money), callerIp, topUp.PaymentMethod, PaymentProviderEpay) + } + + return nil +} + func GetUserTopUps(userId int, pageInfo *common.PageInfo) (topups []*TopUp, total int64, err error) { // Start transaction tx := DB.Begin() @@ -433,7 +526,7 @@ func ManualCompleteTopUp(tradeNo string, callerIp string) error { RecordTopupLog(userId, fmt.Sprintf("管理员补单成功，充值金额: %v，支付金额：%f", logger.FormatQuota(quotaToAdd), payMoney), callerIp, paymentMethod, "admin") return nil } -func RechargeCreem(referenceId string, customerEmail string, customerName string, callerIp string) (err error) { +func RechargeCreem(referenceId string, customerEmail string, customerName string, callerIp string, validations ...PaymentValidation) (err error) { if referenceId == "" { return errors.New("未提供支付单号") } @@ -461,6 +554,10 @@ func RechargeCreem(referenceId string, customerEmail string, customerName string } // Creem 直接使用 Amount 作为充值额度（整数） + if err := validatePaymentAgainstTopUp(topUp, validations); err != nil { + return err + } + quota = topUp.Amount if quota <= 0 { return errors.New("无效的充值额度") @@ -489,6 +586,9 @@ func RechargeCreem(referenceId string, customerEmail string, customerName string if err != nil { common.SysError("creem topup failed: " + err.Error()) + if preserved := preservePaymentValidationError(err); preserved != nil { + return preserved + } return errors.New("充值失败，请稍后重试") } @@ -522,7 +622,7 @@ func addCreemCustomerEmailUpdateIfAvailable(tx *gorm.DB, userId int, customerEma return nil } -func RechargeWaffo(tradeNo string, callerIp string) (err error) { +func RechargeWaffo(tradeNo string, callerIp string, validations ...PaymentValidation) (err error) { if tradeNo == "" { return errors.New("未提供支付单号") } @@ -553,6 +653,10 @@ func RechargeWaffo(tradeNo string, callerIp string) (err error) { return errors.New("充值订单状态错误") } + if err := validatePaymentAgainstTopUp(topUp, validations); err != nil { + return err + } + dAmount := decimal.NewFromInt(topUp.Amount) dQuotaPerUnit := decimal.NewFromFloat(common.QuotaPerUnit) quotaToAdd = int(dAmount.Mul(dQuotaPerUnit).IntPart()) @@ -574,6 +678,9 @@ func RechargeWaffo(tradeNo string, callerIp string) (err error) { if err != nil { common.SysError("waffo topup failed: " + err.Error()) + if preserved := preservePaymentValidationError(err); preserved != nil { + return preserved + } return errors.New("充值失败，请稍后重试") } @@ -584,7 +691,7 @@ func RechargeWaffo(tradeNo string, callerIp string) (err error) { return nil } -func RechargeWaffoPancake(tradeNo string) (err error) { +func RechargeWaffoPancake(tradeNo string, validations ...PaymentValidation) (err error) { if tradeNo == "" { return errors.New("未提供支付单号") } @@ -615,6 +722,10 @@ func RechargeWaffoPancake(tradeNo string) (err error) { return errors.New("充值订单状态错误") } + if err := validatePaymentAgainstTopUp(topUp, validations); err != nil { + return err + } + quotaToAdd = int(decimal.NewFromInt(topUp.Amount).Mul(decimal.NewFromFloat(common.QuotaPerUnit)).IntPart()) if quotaToAdd <= 0 { return errors.New("无效的充值额度") @@ -634,6 +745,9 @@ func RechargeWaffoPancake(tradeNo string) (err error) { if err != nil { common.SysError("waffo pancake topup failed: " + err.Error()) + if preserved := preservePaymentValidationError(err); preserved != nil { + return preserved + } return errors.New("充值失败，请稍后重试") } diff --git a/model/twofa.go b/model/twofa.go index fc6c8bf0..f108cbd2 100644 --- a/model/twofa.go +++ b/model/twofa.go @@ -179,7 +179,7 @@ func ValidateBackupCode(userId int, code string) (bool, error) { // 查找未使用的备用码 var backupCodes []TwoFABackupCode - if err := DB.Where("user_id = ? AND is_used = false", userId).Find(&backupCodes).Error; err != nil { + if err := DB.Where("user_id = ? AND is_used = ?", userId, false).Find(&backupCodes).Error; err != nil { return false, err } @@ -205,7 +205,7 @@ func ValidateBackupCode(userId int, code string) (bool, error) { // GetUnusedBackupCodeCount 获取未使用的备用码数量 func GetUnusedBackupCodeCount(userId int) (int, error) { var count int64 - err := DB.Model(&TwoFABackupCode{}).Where("user_id = ? AND is_used = false", userId).Count(&count).Error + err := DB.Model(&TwoFABackupCode{}).Where("user_id = ? AND is_used = ?", userId, false).Count(&count).Error return int(count), err } @@ -304,7 +304,7 @@ func GetTwoFAStats() (map[string]interface{}, error) { } // 启用2FA的用户数 - if err := DB.Model(&TwoFA{}).Where("is_enabled = true").Count(&enabledUsers).Error; err != nil { + if err := DB.Model(&TwoFA{}).Where("is_enabled = ?", true).Count(&enabledUsers).Error; err != nil { return nil, err } diff --git a/model/user.go b/model/user.go index d00d2265..dfe167cc 100644 --- a/model/user.go +++ b/model/user.go @@ -702,10 +702,14 @@ func HardDeleteUserById(id int) error { } func inviteUser(inviterId int) (err error) { - result := DB.Model(&User{}).Where("id = ?", inviterId).Updates(map[string]interface{}{ + return inviteUserTx(DB, inviterId, int64(common.QuotaForInviter)) +} + +func inviteUserTx(tx *gorm.DB, inviterId int, quota int64) error { + result := tx.Model(&User{}).Where("id = ?", inviterId).Updates(map[string]interface{}{ "aff_count": gorm.Expr("aff_count + 1"), - "aff_quota": gorm.Expr("aff_quota + ?", common.QuotaForInviter), - "aff_history": gorm.Expr("aff_history + ?", common.QuotaForInviter), + "aff_quota": gorm.Expr("aff_quota + ?", quota), + "aff_history": gorm.Expr("aff_history + ?", quota), }) if result.Error != nil { return result.Error @@ -716,6 +720,44 @@ func inviteUser(inviterId int) (err error) { return nil } +type invitationRewardPlan struct { + inviteeQuota int64 + inviterQuota int64 +} + +func buildInvitationRewardPlan(inviterId int) invitationRewardPlan { + if inviterId == 0 || !operation_setting.IsPaymentComplianceConfirmed() { + return invitationRewardPlan{} + } + plan := invitationRewardPlan{} + if common.QuotaForInvitee > 0 { + plan.inviteeQuota = int64(common.QuotaForInvitee) + } + if common.QuotaForInviter > 0 { + plan.inviterQuota = int64(common.QuotaForInviter) + } + return plan +} + +func applyInviterRewardTx(tx *gorm.DB, inviterId int, plan invitationRewardPlan) (CacheInvalidationTask, error) { + if inviterId == 0 || plan.inviterQuota <= 0 { + return CacheInvalidationTask{}, nil + } + if err := inviteUserTx(tx, inviterId, plan.inviterQuota); err != nil { + return CacheInvalidationTask{}, err + } + return stageUserCacheInvalidationTx(tx, inviterId, false) +} + +func recordInvitationRewardLogs(inviteeId, inviterId int, plan invitationRewardPlan) { + if plan.inviteeQuota > 0 { + RecordLog(inviteeId, LogTypeSystem, fmt.Sprintf("使用邀请码赠送 %s", logger.LogQuota(int(plan.inviteeQuota)))) + } + if plan.inviterQuota > 0 { + RecordLog(inviterId, LogTypeSystem, fmt.Sprintf("邀请用户赠送 %s", logger.LogQuota(int(plan.inviterQuota)))) + } +} + func (user *User) TransferAffQuotaToQuota(quota int64) error { // 检查quota是否小于最小额度 if float64(quota) < common.QuotaPerUnit { @@ -822,12 +864,18 @@ func BindEmailToUser(user *User, email string) error { } func (user *User) Insert(inviterId int) error { + affCode, err := common.GenerateRandomCharsKey(4) + if err != nil { + return fmt.Errorf("generate affiliate code: %w", err) + } + rewardPlan := buildInvitationRewardPlan(inviterId) + var inviterCacheTask CacheInvalidationTask if err := WithNormalizedEmailWriteTx(user.Email, func(tx *gorm.DB) error { if err := user.prepareForInsert(tx); err != nil { return err } - user.Quota = int64(common.QuotaForNewUser) - user.AffCode = common.GetRandomString(4) + user.Quota = int64(common.QuotaForNewUser) + rewardPlan.inviteeQuota + user.AffCode = affCode // 初始化用户设置，包括默认的边栏配置 if user.Setting == "" { @@ -836,10 +884,16 @@ func (user *User) Insert(inviterId int) error { user.SetSetting(defaultSetting) } - return tx.Create(user).Error + if err := tx.Create(user).Error; err != nil { + return err + } + var err error + inviterCacheTask, err = applyInviterRewardTx(tx, inviterId, rewardPlan) + return err }); err != nil { return err } + dispatchStagedCacheInvalidation(inviterCacheTask) // 用户创建成功后，根据角色初始化边栏配置 // 需要重新获取用户以确保有正确的ID和Role @@ -851,25 +905,18 @@ func (user *User) Insert(inviterId int) error { currentSetting := createdUser.GetSetting() currentSetting.SidebarModules = defaultSidebarConfig createdUser.SetSetting(currentSetting) - createdUser.UpdateFields(false, UserUpdateFieldSetting) - common.SysLog(fmt.Sprintf("为新用户 %s (角色: %d) 初始化边栏配置", createdUser.Username, createdUser.Role)) + if err := createdUser.UpdateFields(false, UserUpdateFieldSetting); err != nil { + common.SysLog(fmt.Sprintf("初始化新用户 %s 边栏配置失败: %s", createdUser.Username, err.Error())) + } else { + common.SysLog(fmt.Sprintf("为新用户 %s (角色: %d) 初始化边栏配置", createdUser.Username, createdUser.Role)) + } } } if common.QuotaForNewUser > 0 { RecordLog(user.Id, LogTypeSystem, fmt.Sprintf("新用户注册赠送 %s", logger.LogQuota(common.QuotaForNewUser))) } - if inviterId != 0 && operation_setting.IsPaymentComplianceConfirmed() { - if common.QuotaForInvitee > 0 { - _ = IncreaseUserQuota(user.Id, common.QuotaForInvitee, true) - RecordLog(user.Id, LogTypeSystem, fmt.Sprintf("使用邀请码赠送 %s", logger.LogQuota(common.QuotaForInvitee))) - } - if common.QuotaForInviter > 0 { - //_ = IncreaseUserQuota(inviterId, common.QuotaForInviter) - RecordLog(inviterId, LogTypeSystem, fmt.Sprintf("邀请用户赠送 %s", logger.LogQuota(common.QuotaForInviter))) - _ = inviteUser(inviterId) - } - } + recordInvitationRewardLogs(user.Id, inviterId, rewardPlan) return nil } @@ -879,12 +926,17 @@ func (user *User) Insert(inviterId int) error { // This is used for OAuth registration where user creation and binding need to be atomic. // Post-creation tasks (sidebar config, logs, inviter rewards) are handled after the transaction commits. func (user *User) InsertWithTx(tx *gorm.DB, inviterId int) error { + affCode, err := common.GenerateRandomCharsKey(4) + if err != nil { + return fmt.Errorf("generate affiliate code: %w", err) + } + rewardPlan := buildInvitationRewardPlan(inviterId) return withNormalizedEmailLock(tx, user.Email, func(tx *gorm.DB) error { if err := user.prepareForInsert(tx); err != nil { return err } - user.Quota = int64(common.QuotaForNewUser) - user.AffCode = common.GetRandomString(4) + user.Quota = int64(common.QuotaForNewUser) + rewardPlan.inviteeQuota + user.AffCode = affCode // 初始化用户设置 if user.Setting == "" { @@ -892,13 +944,18 @@ func (user *User) InsertWithTx(tx *gorm.DB, inviterId int) error { user.SetSetting(defaultSetting) } - return tx.Create(user).Error + if err := tx.Create(user).Error; err != nil { + return err + } + _, err := applyInviterRewardTx(tx, inviterId, rewardPlan) + return err }) } // FinalizeOAuthUserCreation performs post-transaction tasks for OAuth user creation. // This should be called after the transaction commits successfully. func (user *User) FinalizeOAuthUserCreation(inviterId int) { + rewardPlan := buildInvitationRewardPlan(inviterId) // 用户创建成功后，根据角色初始化边栏配置 var createdUser User if err := DB.Where("id = ?", user.Id).First(&createdUser).Error; err == nil { @@ -907,24 +964,18 @@ func (user *User) FinalizeOAuthUserCreation(inviterId int) { currentSetting := createdUser.GetSetting() currentSetting.SidebarModules = defaultSidebarConfig createdUser.SetSetting(currentSetting) - createdUser.UpdateFields(false, UserUpdateFieldSetting) - common.SysLog(fmt.Sprintf("为新用户 %s (角色: %d) 初始化边栏配置", createdUser.Username, createdUser.Role)) + if err := createdUser.UpdateFields(false, UserUpdateFieldSetting); err != nil { + common.SysLog(fmt.Sprintf("初始化新用户 %s 边栏配置失败: %s", createdUser.Username, err.Error())) + } else { + common.SysLog(fmt.Sprintf("为新用户 %s (角色: %d) 初始化边栏配置", createdUser.Username, createdUser.Role)) + } } } if common.QuotaForNewUser > 0 { RecordLog(user.Id, LogTypeSystem, fmt.Sprintf("新用户注册赠送 %s", logger.LogQuota(common.QuotaForNewUser))) } - if inviterId != 0 && operation_setting.IsPaymentComplianceConfirmed() { - if common.QuotaForInvitee > 0 { - _ = IncreaseUserQuota(user.Id, common.QuotaForInvitee, true) - RecordLog(user.Id, LogTypeSystem, fmt.Sprintf("使用邀请码赠送 %s", logger.LogQuota(common.QuotaForInvitee))) - } - if common.QuotaForInviter > 0 { - RecordLog(inviterId, LogTypeSystem, fmt.Sprintf("邀请用户赠送 %s", logger.LogQuota(common.QuotaForInviter))) - _ = inviteUser(inviterId) - } - } + recordInvitationRewardLogs(user.Id, inviterId, rewardPlan) } func (user *User) Update(updatePassword bool) error { @@ -1635,6 +1686,30 @@ type quotaDeltaInteger interface { ~int | ~int64 } +func OverrideUserQuota(id int, quota int64) error { + if id <= 0 { + return ErrUserNotFound + } + if quota < 0 { + return errors.New("quota 不能为负数！") + } + var cacheTask CacheInvalidationTask + err := DB.Transaction(func(tx *gorm.DB) error { + result := tx.Model(&User{}).Where("id = ?", id).Update("quota", quota) + if err := ensureUserUpdateMatchedTx(tx, result, id, ErrUserNotFound); err != nil { + return err + } + var err error + cacheTask, err = stageUserCacheInvalidationTx(tx, id, false) + return err + }) + if err != nil { + return err + } + dispatchStagedCacheInvalidation(cacheTask) + return nil +} + // The final boolean is retained for caller compatibility. Accounting mutations // always write the database directly so conditional deductions remain atomic. func IncreaseUserQuota[T quotaDeltaInteger](id int, quota T, _ bool) (err error) { diff --git a/model/user_update_test.go b/model/user_update_test.go index 664080d8..862c6b74 100644 --- a/model/user_update_test.go +++ b/model/user_update_test.go @@ -16,6 +16,7 @@ import ( "github.com/MAX-API-Next/MAX-API/common" "github.com/MAX-API-Next/MAX-API/dto" + "github.com/MAX-API-Next/MAX-API/setting/operation_setting" miniredis "github.com/alicebob/miniredis/v2" "github.com/glebarez/sqlite" "github.com/go-redis/redis/v8" @@ -145,6 +146,75 @@ func TestInviteUserUpdatesOnlyAffiliateCounters(t *testing.T) { assert.Equal(t, common.UserStatusEnabled, got.Status) } +func TestUserInsertAppliesInvitationRewardsAtomically(t *testing.T) { + setupUserUpdateTestState(t) + oldNewUserQuota := common.QuotaForNewUser + oldInviteeQuota := common.QuotaForInvitee + oldInviterQuota := common.QuotaForInviter + oldPaymentSetting := *operation_setting.GetPaymentSetting() + common.QuotaForNewUser = 100 + common.QuotaForInvitee = 20 + common.QuotaForInviter = 30 + paymentSetting := operation_setting.GetPaymentSetting() + paymentSetting.ComplianceConfirmed = true + paymentSetting.ComplianceTermsVersion = operation_setting.CurrentComplianceTermsVersion + t.Cleanup(func() { + common.QuotaForNewUser = oldNewUserQuota + common.QuotaForInvitee = oldInviteeQuota + common.QuotaForInviter = oldInviterQuota + *operation_setting.GetPaymentSetting() = oldPaymentSetting + }) + + inviter := User{Id: 410, Username: "atomic-inviter", AffCode: "atomic-inviter", Status: common.UserStatusEnabled} + require.NoError(t, DB.Create(&inviter).Error) + invitee := User{Username: "atomic-invitee", AffCode: "", Status: common.UserStatusEnabled} + require.NoError(t, invitee.Insert(inviter.Id)) + + var storedInvitee User + require.NoError(t, DB.First(&storedInvitee, invitee.Id).Error) + assert.EqualValues(t, 120, storedInvitee.Quota) + var storedInviter User + require.NoError(t, DB.First(&storedInviter, inviter.Id).Error) + assert.Equal(t, 1, storedInviter.AffCount) + assert.EqualValues(t, 30, storedInviter.AffQuota) + assert.EqualValues(t, 30, storedInviter.AffHistoryQuota) +} + +func TestUserInsertRollsBackWhenInvitationRewardOutboxCannotPersist(t *testing.T) { + setupUserUpdateTestState(t) + oldInviteeQuota := common.QuotaForInvitee + oldInviterQuota := common.QuotaForInviter + oldPaymentSetting := *operation_setting.GetPaymentSetting() + common.QuotaForInvitee = 20 + common.QuotaForInviter = 30 + paymentSetting := operation_setting.GetPaymentSetting() + paymentSetting.ComplianceConfirmed = true + paymentSetting.ComplianceTermsVersion = operation_setting.CurrentComplianceTermsVersion + t.Cleanup(func() { + common.QuotaForInvitee = oldInviteeQuota + common.QuotaForInviter = oldInviterQuota + *operation_setting.GetPaymentSetting() = oldPaymentSetting + }) + + inviter := User{Id: 411, Username: "rollback-inviter", AffCode: "rollback-inviter", Status: common.UserStatusEnabled} + require.NoError(t, DB.Create(&inviter).Error) + useFailingUserUpdateRedis(t) + failCacheOutboxInserts(t) + + invitee := User{Username: "rollback-invitee", Status: common.UserStatusEnabled} + err := invitee.Insert(inviter.Id) + require.Error(t, err) + + var inviteeCount int64 + require.NoError(t, DB.Model(&User{}).Where("username = ?", invitee.Username).Count(&inviteeCount).Error) + assert.Zero(t, inviteeCount) + var storedInviter User + require.NoError(t, DB.First(&storedInviter, inviter.Id).Error) + assert.Zero(t, storedInviter.AffCount) + assert.Zero(t, storedInviter.AffQuota) + assert.Zero(t, storedInviter.AffHistoryQuota) +} + func TestFillUserMethodsReturnDatabaseErrors(t *testing.T) { db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) require.NoError(t, err) @@ -463,6 +533,21 @@ func TestUserQuotaMutationsRetryCacheInvalidation(t *testing.T) { } } +func TestOverrideUserQuotaRollsBackWhenCacheOutboxCannotPersist(t *testing.T) { + setupUserUpdateTestState(t) + user := User{Id: 399, Username: "quota-override-outbox-rollback", Quota: 100, Status: common.UserStatusEnabled} + require.NoError(t, DB.Create(&user).Error) + useFailingUserUpdateRedis(t) + failCacheOutboxInserts(t) + + err := OverrideUserQuota(user.Id, 25) + + require.Error(t, err) + var stored User + require.NoError(t, DB.First(&stored, user.Id).Error) + assert.EqualValues(t, 100, stored.Quota) +} + func TestUserAuthFieldMutationsRetryCacheInvalidation(t *testing.T) { setupUserUpdateTestState(t) user := User{Id: 82, Username: "auth-cache-retry-user", Quota: 10, Status: common.UserStatusEnabled, Role: common.RoleCommonUser} @@ -747,9 +832,26 @@ func TestBatchTokenDeleteRetriesCacheDeletion(t *testing.T) { deleted, err := BatchDeleteTokens([]int{token.Id}, token.UserId) require.NoError(t, err) require.Equal(t, 1, deleted) + processCacheInvalidationTasks() requireCacheKeyDeletedEventually(t, client, getTokenCacheKey(token.Key)) } +func TestBatchTokenDeleteRollsBackWhenCacheOutboxCannotPersist(t *testing.T) { + setupUserUpdateTestState(t) + token := Token{Id: 397, UserId: 398, Key: "batch-token-delete-outbox-rollback", RemainQuota: 10, Status: common.TokenStatusEnabled} + require.NoError(t, DB.Create(&token).Error) + useFailingUserUpdateRedis(t) + failCacheOutboxInserts(t) + + deleted, err := BatchDeleteTokens([]int{token.Id}, token.UserId) + + require.Error(t, err) + require.Zero(t, deleted) + var count int64 + require.NoError(t, DB.Model(&Token{}).Where("id = ?", token.Id).Count(&count).Error) + assert.EqualValues(t, 1, count) +} + func newBlockingCacheFillHook() *blockingCacheFillHook { return &blockingCacheFillHook{ started: make(chan struct{}), diff --git a/oauth/generic.go b/oauth/generic.go index 59d913fd..255a0c5b 100644 --- a/oauth/generic.go +++ b/oauth/generic.go @@ -92,8 +92,6 @@ func (p *GenericOAuthProvider) ExchangeToken(ctx context.Context, code string, c return nil, NewOAuthError(i18n.MsgOAuthInvalidCode, nil) } - logger.LogDebug(ctx, "[OAuth-Generic-%s] ExchangeToken: code=%s...", p.config.Slug, code[:min(len(code), 10)]) - redirectUri := fmt.Sprintf("%s/oauth/%s", system_setting.ServerAddress, p.config.Slug) values := url.Values{} values.Set("grant_type", "authorization_code") @@ -128,11 +126,11 @@ func (p *GenericOAuthProvider) ExchangeToken(ctx context.Context, code string, c req.Header.Set("Authorization", "Basic "+credentials) } - logger.LogDebug(ctx, "[OAuth-Generic-%s] ExchangeToken: token_endpoint=%s, redirect_uri=%s, auth_style=%d", - p.config.Slug, p.config.TokenEndpoint, redirectUri, authStyle) + logger.LogDebug(ctx, "[OAuth-Generic-%s] ExchangeToken: auth_style=%d", p.config.Slug, authStyle) - client := http.Client{ - Timeout: 20 * time.Second, + client, err := newOAuthHTTPClient(p.config.TokenEndpoint, 20*time.Second) + if err != nil { + return nil, NewOAuthErrorWithRaw(i18n.MsgOAuthConnectFailed, map[string]any{"Provider": p.config.Name}, err.Error()) } res, err := client.Do(req) if err != nil { @@ -150,7 +148,6 @@ func (p *GenericOAuthProvider) ExchangeToken(ctx context.Context, code string, c } bodyStr := string(body) - logger.LogDebug(ctx, "[OAuth-Generic-%s] ExchangeToken response body: %s", p.config.Slug, bodyStr[:min(len(bodyStr), 500)]) // Try to parse as JSON first var tokenResponse struct { @@ -200,7 +197,7 @@ func (p *GenericOAuthProvider) ExchangeToken(ctx context.Context, code string, c } func (p *GenericOAuthProvider) GetUserInfo(ctx context.Context, token *OAuthToken) (*OAuthUser, error) { - logger.LogDebug(ctx, "[OAuth-Generic-%s] GetUserInfo: fetching user info from %s", p.config.Slug, p.config.UserInfoEndpoint) + logger.LogDebug(ctx, "[OAuth-Generic-%s] GetUserInfo started", p.config.Slug) req, err := http.NewRequestWithContext(ctx, "GET", p.config.UserInfoEndpoint, nil) if err != nil { @@ -212,8 +209,9 @@ func (p *GenericOAuthProvider) GetUserInfo(ctx context.Context, token *OAuthToke req.Header.Set("Authorization", fmt.Sprintf("%s %s", tokenType, token.AccessToken)) req.Header.Set("Accept", "application/json") - client := http.Client{ - Timeout: 20 * time.Second, + client, err := newOAuthHTTPClient(p.config.UserInfoEndpoint, 20*time.Second) + if err != nil { + return nil, NewOAuthErrorWithRaw(i18n.MsgOAuthConnectFailed, map[string]any{"Provider": p.config.Name}, err.Error()) } res, err := client.Do(req) if err != nil { @@ -236,7 +234,6 @@ func (p *GenericOAuthProvider) GetUserInfo(ctx context.Context, token *OAuthToke } bodyStr := string(body) - logger.LogDebug(ctx, "[OAuth-Generic-%s] GetUserInfo response body: %s", p.config.Slug, bodyStr[:min(len(bodyStr), 500)]) // Extract fields using gjson (supports JSONPath-like syntax) userId := gjson.Get(bodyStr, p.config.UserIdField).String() @@ -260,8 +257,8 @@ func (p *GenericOAuthProvider) GetUserInfo(ctx context.Context, token *OAuthToke return nil, NewOAuthError(i18n.MsgOAuthUserInfoEmpty, map[string]any{"Provider": p.config.Name}) } - logger.LogDebug(ctx, "[OAuth-Generic-%s] GetUserInfo success: id=%s, username=%s, name=%s, email=%s", - p.config.Slug, userId, username, displayName, email) + logger.LogDebug(ctx, "[OAuth-Generic-%s] GetUserInfo success: username_present=%t name_present=%t email_present=%t", + p.config.Slug, username != "", displayName != "", email != "") policyRaw := strings.TrimSpace(p.config.AccessPolicy) if policyRaw != "" { diff --git a/oauth/http_client.go b/oauth/http_client.go new file mode 100644 index 00000000..13139b5a --- /dev/null +++ b/oauth/http_client.go @@ -0,0 +1,19 @@ +package oauth + +import ( + "fmt" + "net/http" + "time" + + "github.com/MAX-API-Next/MAX-API/service" +) + +func newOAuthHTTPClient(endpoint string, timeout time.Duration) (*http.Client, error) { + if err := service.ValidateSSRFProtectedFetchURL(endpoint); err != nil { + return nil, fmt.Errorf("OAuth endpoint rejected: %w", err) + } + base := service.GetSSRFProtectedHttpClient() + client := *base + client.Timeout = timeout + return &client, nil +} diff --git a/oauth/http_client_test.go b/oauth/http_client_test.go new file mode 100644 index 00000000..51ac5c03 --- /dev/null +++ b/oauth/http_client_test.go @@ -0,0 +1,119 @@ +package oauth + +import ( + "bytes" + "context" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/MAX-API-Next/MAX-API/common" + "github.com/MAX-API-Next/MAX-API/model" + "github.com/MAX-API-Next/MAX-API/setting/system_setting" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +func configureOAuthFetchTest(t *testing.T, enabled bool) { + t.Helper() + settings := system_setting.GetFetchSetting() + original := *settings + t.Cleanup(func() { *settings = original }) + settings.EnableSSRFProtection = enabled + settings.AllowPrivateIp = false + settings.DomainFilterMode = false + settings.IpFilterMode = false + settings.DomainList = nil + settings.IpList = nil + settings.AllowedPorts = []string{"80", "443"} + settings.ApplyIPFilterForDomain = true +} + +func TestNewOAuthHTTPClientRejectsPrivateEndpoint(t *testing.T) { + configureOAuthFetchTest(t, true) + + client, err := newOAuthHTTPClient("http://127.0.0.1/token", time.Second) + + require.Nil(t, client) + require.ErrorContains(t, err, "private IP address not allowed") +} + +func TestGenericOAuthDebugLogsDoNotContainResponseSecrets(t *testing.T) { + configureOAuthFetchTest(t, false) + originalDebug := common.DebugEnabled + common.DebugEnabled = true + t.Cleanup(func() { common.DebugEnabled = originalDebug }) + + const ( + accessToken = "access-secret-value" + refreshToken = "refresh-secret-value" + idToken = "id-secret-value" + email = "secret-email@example.com" + ) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + if r.URL.Path == "/token" { + _, _ = w.Write([]byte(`{"access_token":"` + accessToken + `","refresh_token":"` + refreshToken + `","id_token":"` + idToken + `","token_type":"Bearer"}`)) + return + } + _, _ = w.Write([]byte(`{"id":"provider-user","username":"tester","name":"Test User","email":"` + email + `"}`)) + })) + defer server.Close() + + var logs bytes.Buffer + common.LogWriterMu.Lock() + originalWriter := gin.DefaultErrorWriter + gin.DefaultErrorWriter = &logs + common.LogWriterMu.Unlock() + t.Cleanup(func() { + common.LogWriterMu.Lock() + gin.DefaultErrorWriter = originalWriter + common.LogWriterMu.Unlock() + }) + + provider := NewGenericOAuthProvider(&model.CustomOAuthProvider{ + Name: "Test OAuth", + Slug: "test-oauth", + Enabled: true, + TokenEndpoint: server.URL + "/token", + UserInfoEndpoint: server.URL + "/userinfo", + ClientId: "client-id", + ClientSecret: "client-secret", + UserIdField: "id", + UsernameField: "username", + DisplayNameField: "name", + EmailField: "email", + }) + token, err := provider.ExchangeToken(context.Background(), "authorization-code", nil) + require.NoError(t, err) + _, err = provider.GetUserInfo(context.Background(), token) + require.NoError(t, err) + + output := logs.String() + for _, secret := range []string{accessToken, refreshToken, idToken, email, "authorization-code"} { + require.False(t, strings.Contains(output, secret), "debug log leaked %q", secret) + } +} + +func TestOIDCGetUserInfoAllowsSubjectWithoutEmail(t *testing.T) { + configureOAuthFetchTest(t, false) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"sub":"subject-only-user","preferred_username":"tester"}`)) + })) + defer server.Close() + + settings := system_setting.GetOIDCSettings() + original := *settings + settings.UserInfoEndpoint = server.URL + t.Cleanup(func() { *settings = original }) + + user, err := (&OIDCProvider{}).GetUserInfo(context.Background(), &OAuthToken{AccessToken: "access-token"}) + + require.NoError(t, err) + require.Equal(t, "subject-only-user", user.ProviderUserID) + require.Equal(t, "tester", user.Username) + require.Empty(t, user.Email) +} diff --git a/oauth/oidc.go b/oauth/oidc.go index 3826e0aa..28a31904 100644 --- a/oauth/oidc.go +++ b/oauth/oidc.go @@ -2,13 +2,13 @@ package oauth import ( "context" - "encoding/json" "fmt" "net/http" "net/url" "strings" "time" + "github.com/MAX-API-Next/MAX-API/common" "github.com/MAX-API-Next/MAX-API/i18n" "github.com/MAX-API-Next/MAX-API/logger" "github.com/MAX-API-Next/MAX-API/model" @@ -53,8 +53,6 @@ func (p *OIDCProvider) ExchangeToken(ctx context.Context, code string, c *gin.Co return nil, NewOAuthError(i18n.MsgOAuthInvalidCode, nil) } - logger.LogDebug(ctx, "[OAuth-OIDC] ExchangeToken: code=%s...", code[:min(len(code), 10)]) - settings := system_setting.GetOIDCSettings() redirectUri := fmt.Sprintf("%s/oauth/oidc", system_setting.ServerAddress) values := url.Values{} @@ -64,8 +62,6 @@ func (p *OIDCProvider) ExchangeToken(ctx context.Context, code string, c *gin.Co values.Set("grant_type", "authorization_code") values.Set("redirect_uri", redirectUri) - logger.LogDebug(ctx, "[OAuth-OIDC] ExchangeToken: token_endpoint=%s, redirect_uri=%s", settings.TokenEndpoint, redirectUri) - req, err := http.NewRequestWithContext(ctx, "POST", settings.TokenEndpoint, strings.NewReader(values.Encode())) if err != nil { return nil, err @@ -73,8 +69,9 @@ func (p *OIDCProvider) ExchangeToken(ctx context.Context, code string, c *gin.Co req.Header.Set("Content-Type", "application/x-www-form-urlencoded") req.Header.Set("Accept", "application/json") - client := http.Client{ - Timeout: 5 * time.Second, + client, err := newOAuthHTTPClient(settings.TokenEndpoint, 5*time.Second) + if err != nil { + return nil, NewOAuthErrorWithRaw(i18n.MsgOAuthConnectFailed, map[string]any{"Provider": "OIDC"}, err.Error()) } res, err := client.Do(req) if err != nil { @@ -86,7 +83,7 @@ func (p *OIDCProvider) ExchangeToken(ctx context.Context, code string, c *gin.Co logger.LogDebug(ctx, "[OAuth-OIDC] ExchangeToken response status: %d", res.StatusCode) var oidcResponse oidcOAuthResponse - err = json.NewDecoder(res.Body).Decode(&oidcResponse) + err = common.DecodeJson(res.Body, &oidcResponse) if err != nil { logger.LogError(ctx, fmt.Sprintf("[OAuth-OIDC] ExchangeToken decode error: %s", err.Error())) return nil, err @@ -112,16 +109,15 @@ func (p *OIDCProvider) ExchangeToken(ctx context.Context, code string, c *gin.Co func (p *OIDCProvider) GetUserInfo(ctx context.Context, token *OAuthToken) (*OAuthUser, error) { settings := system_setting.GetOIDCSettings() - logger.LogDebug(ctx, "[OAuth-OIDC] GetUserInfo: userinfo_endpoint=%s", settings.UserInfoEndpoint) - req, err := http.NewRequestWithContext(ctx, "GET", settings.UserInfoEndpoint, nil) if err != nil { return nil, err } req.Header.Set("Authorization", "Bearer "+token.AccessToken) - client := http.Client{ - Timeout: 5 * time.Second, + client, err := newOAuthHTTPClient(settings.UserInfoEndpoint, 5*time.Second) + if err != nil { + return nil, NewOAuthErrorWithRaw(i18n.MsgOAuthConnectFailed, map[string]any{"Provider": "OIDC"}, err.Error()) } res, err := client.Do(req) if err != nil { @@ -138,18 +134,18 @@ func (p *OIDCProvider) GetUserInfo(ctx context.Context, token *OAuthToken) (*OAu } var oidcUser oidcUser - err = json.NewDecoder(res.Body).Decode(&oidcUser) + err = common.DecodeJson(res.Body, &oidcUser) if err != nil { logger.LogError(ctx, fmt.Sprintf("[OAuth-OIDC] GetUserInfo decode error: %s", err.Error())) return nil, err } - if oidcUser.OpenID == "" || oidcUser.Email == "" { - logger.LogError(ctx, fmt.Sprintf("[OAuth-OIDC] GetUserInfo failed: empty fields (sub=%s, email=%s)", oidcUser.OpenID, oidcUser.Email)) + if oidcUser.OpenID == "" { + logger.LogError(ctx, "[OAuth-OIDC] GetUserInfo failed: subject identifier is empty") return nil, NewOAuthError(i18n.MsgOAuthUserInfoEmpty, map[string]any{"Provider": "OIDC"}) } - logger.LogDebug(ctx, "[OAuth-OIDC] GetUserInfo success: sub=%s, username=%s, name=%s, email=%s", oidcUser.OpenID, oidcUser.PreferredUsername, oidcUser.Name, oidcUser.Email) + logger.LogDebug(ctx, "[OAuth-OIDC] GetUserInfo success: username_present=%t name_present=%t email_present=%t", oidcUser.PreferredUsername != "", oidcUser.Name != "", oidcUser.Email != "") return &OAuthUser{ ProviderUserID: oidcUser.OpenID, diff --git a/pkg/billingexpr/billingexpr_test.go b/pkg/billingexpr/billingexpr_test.go index 5ae29f6c..e0175b10 100644 --- a/pkg/billingexpr/billingexpr_test.go +++ b/pkg/billingexpr/billingexpr_test.go @@ -6,16 +6,17 @@ import ( "testing" "github.com/MAX-API-Next/MAX-API/pkg/billingexpr" + "github.com/stretchr/testify/require" ) // --------------------------------------------------------------------------- // Claude-style: fixed tiers, input > 200K changes both input & output price // --------------------------------------------------------------------------- -const claudeExpr = `p <= 200000 ? tier("standard", p * 1.5 + c * 7.5) : tier("long_context", p * 3.0 + c * 11.25)` +const claudeExpr = `len <= 200000 ? tier("standard", p * 1.5 + c * 7.5) : tier("long_context", p * 3.0 + c * 11.25)` func TestClaude_StandardTier(t *testing.T) { - cost, trace, err := billingexpr.RunExpr(claudeExpr, billingexpr.TokenParams{P: 100000, C: 5000}) + cost, trace, err := billingexpr.RunExpr(claudeExpr, billingexpr.TokenParams{P: 100000, C: 5000, Len: 100000}) if err != nil { t.Fatal(err) } @@ -29,7 +30,7 @@ func TestClaude_StandardTier(t *testing.T) { } func TestClaude_LongContextTier(t *testing.T) { - cost, trace, err := billingexpr.RunExpr(claudeExpr, billingexpr.TokenParams{P: 300000, C: 10000}) + cost, trace, err := billingexpr.RunExpr(claudeExpr, billingexpr.TokenParams{P: 300000, C: 10000, Len: 300000}) if err != nil { t.Fatal(err) } @@ -43,7 +44,7 @@ func TestClaude_LongContextTier(t *testing.T) { } func TestClaude_BoundaryExact(t *testing.T) { - cost, trace, err := billingexpr.RunExpr(claudeExpr, billingexpr.TokenParams{P: 200000, C: 1000}) + cost, trace, err := billingexpr.RunExpr(claudeExpr, billingexpr.TokenParams{P: 200000, C: 1000, Len: 200000}) if err != nil { t.Fatal(err) } @@ -323,7 +324,7 @@ func TestComputeTieredQuota_Basic(t *testing.T) { QuotaPerUnit: 500_000, } - result, err := billingexpr.ComputeTieredQuota(snap, billingexpr.TokenParams{P: 300000, C: 10000}) + result, err := billingexpr.ComputeTieredQuota(snap, billingexpr.TokenParams{P: 300000, C: 10000, Len: 300000}) if err != nil { t.Fatal(err) } @@ -354,7 +355,7 @@ func TestComputeTieredQuota_SameTier(t *testing.T) { QuotaPerUnit: 500_000, } - result, err := billingexpr.ComputeTieredQuota(snap, billingexpr.TokenParams{P: 80000, C: 2000}) + result, err := billingexpr.ComputeTieredQuota(snap, billingexpr.TokenParams{P: 80000, C: 2000, Len: 80000}) if err != nil { t.Fatal(err) } @@ -428,10 +429,10 @@ func TestExprHashString_Deterministic(t *testing.T) { // Cache variables: present // --------------------------------------------------------------------------- -const claudeWithCacheExpr = `p <= 200000 ? tier("standard", p * 1.5 + c * 7.5 + cr * 0.15 + cc * 1.875) : tier("long_context", p * 3.0 + c * 11.25 + cr * 0.3 + cc * 3.75)` +const claudeWithCacheExpr = `len <= 200000 ? tier("standard", p * 1.5 + c * 7.5 + cr * 0.15 + cc * 1.875) : tier("long_context", p * 3.0 + c * 11.25 + cr * 0.3 + cc * 3.75)` func TestCachePresent_StandardTier(t *testing.T) { - params := billingexpr.TokenParams{P: 100000, C: 5000, CR: 50000, CC: 10000} + params := billingexpr.TokenParams{P: 100000, C: 5000, Len: 160000, CR: 50000, CC: 10000} cost, trace, err := billingexpr.RunExpr(claudeWithCacheExpr, params) if err != nil { t.Fatal(err) @@ -446,7 +447,7 @@ func TestCachePresent_StandardTier(t *testing.T) { } func TestCachePresent_LongContextTier(t *testing.T) { - params := billingexpr.TokenParams{P: 300000, C: 10000, CR: 100000, CC: 20000} + params := billingexpr.TokenParams{P: 300000, C: 10000, Len: 420000, CR: 100000, CC: 20000} cost, trace, err := billingexpr.RunExpr(claudeWithCacheExpr, params) if err != nil { t.Fatal(err) @@ -465,7 +466,7 @@ func TestCachePresent_LongContextTier(t *testing.T) { // --------------------------------------------------------------------------- func TestCacheAbsent_ZeroCacheTokens(t *testing.T) { - params := billingexpr.TokenParams{P: 100000, C: 5000} + params := billingexpr.TokenParams{P: 100000, C: 5000, Len: 100000} cost, trace, err := billingexpr.RunExpr(claudeWithCacheExpr, params) if err != nil { t.Fatal(err) @@ -514,7 +515,7 @@ func TestMixedCacheFields_AllCacheZero(t *testing.T) { // --------------------------------------------------------------------------- func TestBackwardCompat_OldExprWithTokenParams(t *testing.T) { - params := billingexpr.TokenParams{P: 100000, C: 5000, CR: 99999, CC: 88888} + params := billingexpr.TokenParams{P: 100000, C: 5000, Len: 100000, CR: 99999, CC: 88888} cost, trace, err := billingexpr.RunExpr(claudeExpr, params) if err != nil { t.Fatal(err) @@ -546,7 +547,7 @@ func TestComputeTieredQuota_WithCache(t *testing.T) { QuotaPerUnit: 500_000, } - params := billingexpr.TokenParams{P: 100000, C: 5000, CR: 50000, CC: 10000} + params := billingexpr.TokenParams{P: 100000, C: 5000, Len: 160000, CR: 50000, CC: 10000} result, err := billingexpr.ComputeTieredQuota(snap, params) if err != nil { t.Fatal(err) @@ -578,7 +579,7 @@ func TestComputeTieredQuota_WithCacheCrossTier(t *testing.T) { QuotaPerUnit: 500_000, } - params := billingexpr.TokenParams{P: 300000, C: 10000, CR: 50000, CC: 10000} + params := billingexpr.TokenParams{P: 300000, C: 10000, Len: 360000, CR: 50000, CC: 10000} result, err := billingexpr.ComputeTieredQuota(snap, params) if err != nil { t.Fatal(err) @@ -644,12 +645,14 @@ func TestFuzz_SettlementConsistency(t *testing.T) { CR: math.Round(rng.Float64() * 100000), CC: math.Round(rng.Float64() * 30000), } + estParams.Len = estParams.P + estParams.CR + estParams.CC actParams := billingexpr.TokenParams{ P: math.Round(rng.Float64() * 500000), C: math.Round(rng.Float64() * 100000), CR: math.Round(rng.Float64() * 100000), CC: math.Round(rng.Float64() * 30000), } + actParams.Len = actParams.P + actParams.CR + actParams.CC groupRatio := 0.5 + rng.Float64()*2.0 estCost, estTrace, _ := billingexpr.RunExpr(claudeWithCacheExpr, estParams) @@ -682,6 +685,24 @@ func TestFuzz_SettlementConsistency(t *testing.T) { } } +func TestRunExprRejectsNonFiniteResults(t *testing.T) { + tests := []struct { + name string + expr string + }{ + {name: "infinity", expr: "p / c"}, + {name: "nan", expr: "(p - p) / c"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + _, _, err := billingexpr.RunExpr(tt.expr, billingexpr.TokenParams{P: 1, C: 0}) + require.Error(t, err) + require.Contains(t, err.Error(), "finite") + }) + } +} + // --------------------------------------------------------------------------- // Settlement-level tests for ComputeTieredQuota // --------------------------------------------------------------------------- diff --git a/pkg/billingexpr/run.go b/pkg/billingexpr/run.go index 04ff4456..6b53a7ea 100644 --- a/pkg/billingexpr/run.go +++ b/pkg/billingexpr/run.go @@ -112,6 +112,9 @@ func runProgram(prog *vm.Program, params TokenParams, request RequestInput) (flo if !ok { return 0, trace, fmt.Errorf("expr result is %T, want float64", out) } + if math.IsNaN(f) || math.IsInf(f, 0) { + return 0, trace, fmt.Errorf("expr result must be finite, got %v", f) + } return f, trace, nil } diff --git a/relay/channel/ali/dto.go b/relay/channel/ali/dto.go index aef32947..ad538077 100644 --- a/relay/channel/ali/dto.go +++ b/relay/channel/ali/dto.go @@ -26,11 +26,11 @@ type AliInput struct { } type AliParameters struct { - TopP float64 `json:"top_p,omitempty"` - TopK int `json:"top_k,omitempty"` - Seed uint64 `json:"seed,omitempty"` - EnableSearch bool `json:"enable_search,omitempty"` - IncrementalOutput bool `json:"incremental_output,omitempty"` + TopP *float64 `json:"top_p,omitempty"` + TopK *int `json:"top_k,omitempty"` + Seed *uint64 `json:"seed,omitempty"` + EnableSearch *bool `json:"enable_search,omitempty"` + IncrementalOutput *bool `json:"incremental_output,omitempty"` } type AliChatRequest struct { @@ -167,21 +167,21 @@ type AliImageRequest struct { Model string `json:"model"` Input any `json:"input"` Parameters AliImageParameters `json:"parameters,omitempty"` - ResponseFormat string `json:"response_format,omitempty"` + ResponseFormat *string `json:"response_format,omitempty"` } type AliImageParameters struct { - Size string `json:"size,omitempty"` - N int `json:"n,omitempty"` - Steps string `json:"steps,omitempty"` - Scale string `json:"scale,omitempty"` - Watermark *bool `json:"watermark,omitempty"` - PromptExtend *bool `json:"prompt_extend,omitempty"` - ThinkingMode *bool `json:"thinking_mode,omitempty"` - EnableSequential *bool `json:"enable_sequential,omitempty"` - BboxList any `json:"bbox_list,omitempty"` - ColorPalette any `json:"color_palette,omitempty"` - Seed *int `json:"seed,omitempty"` + Size *string `json:"size,omitempty"` + N *int `json:"n,omitempty"` + Steps *string `json:"steps,omitempty"` + Scale *string `json:"scale,omitempty"` + Watermark *bool `json:"watermark,omitempty"` + PromptExtend *bool `json:"prompt_extend,omitempty"` + ThinkingMode *bool `json:"thinking_mode,omitempty"` + EnableSequential *bool `json:"enable_sequential,omitempty"` + BboxList any `json:"bbox_list,omitempty"` + ColorPalette any `json:"color_palette,omitempty"` + Seed *int `json:"seed,omitempty"` } func (p *AliImageParameters) PromptExtendValue() bool { @@ -191,6 +191,13 @@ func (p *AliImageParameters) PromptExtendValue() bool { return false } +func (p *AliImageParameters) NValue() int { + if p == nil || p.N == nil { + return 0 + } + return *p.N +} + type AliImageInput struct { Prompt string `json:"prompt,omitempty"` NegativePrompt string `json:"negative_prompt,omitempty"` diff --git a/relay/channel/ali/image.go b/relay/channel/ali/image.go index f681c82b..ac59aebf 100644 --- a/relay/channel/ali/image.go +++ b/relay/channel/ali/image.go @@ -24,20 +24,18 @@ import ( func oaiImage2AliImageRequest(info *relaycommon.RelayInfo, request dto.ImageRequest, isSync bool) (*AliImageRequest, error) { var imageRequest AliImageRequest imageRequest.Model = request.Model - imageRequest.ResponseFormat = request.ResponseFormat + if request.ResponseFormat != "" { + imageRequest.ResponseFormat = &request.ResponseFormat + } + defaultN := int(lo.FromPtrOr(request.N, uint(1))) + hasCustomParameters := false if request.Extra != nil { if val, ok := request.Extra["parameters"]; ok { + hasCustomParameters = true err := common.Unmarshal(val, &imageRequest.Parameters) if err != nil { return nil, fmt.Errorf("invalid parameters field: %w", err) } - } else { - // 兼容没有parameters字段的情况，从openai标准字段中提取参数 - imageRequest.Parameters = AliImageParameters{ - Size: strings.Replace(request.Size, "x", "*", -1), - N: int(lo.FromPtrOr(request.N, uint(1))), - Watermark: request.Watermark, - } } if val, ok := request.Extra["input"]; ok { err := common.Unmarshal(val, &imageRequest.Input) @@ -46,6 +44,16 @@ func oaiImage2AliImageRequest(info *relaycommon.RelayInfo, request dto.ImageRequ } } } + if !hasCustomParameters { + imageRequest.Parameters.Watermark = request.Watermark + if request.Size != "" { + size := strings.Replace(request.Size, "x", "*", -1) + imageRequest.Parameters.Size = &size + } + } + if imageRequest.Parameters.N == nil { + imageRequest.Parameters.N = &defaultN + } if strings.Contains(request.Model, "z-image") { // z-image 开启prompt_extend后，按2倍计费 @@ -54,11 +62,12 @@ func oaiImage2AliImageRequest(info *relaycommon.RelayInfo, request dto.ImageRequ } } - if err := dto.ValidateImageN("parameters.n", imageRequest.Parameters.N); err != nil { + n := imageRequest.Parameters.NValue() + if err := dto.ValidateImageN("parameters.n", n); err != nil { return nil, err } - if imageRequest.Parameters.N != 0 { - info.PriceData.AddOtherRatio("n", float64(imageRequest.Parameters.N)) + if n != 0 { + info.PriceData.AddOtherRatio("n", float64(n)) } // 同步图片模型和异步图片模型请求格式不一样 @@ -158,7 +167,9 @@ func getImageBase64sFromForm(c *gin.Context, fieldName string) ([]string, error) func oaiFormEdit2AliImageEdit(c *gin.Context, info *relaycommon.RelayInfo, request dto.ImageRequest) (*AliImageRequest, error) { var imageRequest AliImageRequest imageRequest.Model = request.Model - imageRequest.ResponseFormat = request.ResponseFormat + if request.ResponseFormat != "" { + imageRequest.ResponseFormat = &request.ResponseFormat + } imageBase64s, err := getImageBase64sFromForm(c, "image") if err != nil { @@ -182,8 +193,9 @@ func oaiFormEdit2AliImageEdit(c *gin.Context, info *relaycommon.RelayInfo, reque }, }, } + n := int(lo.FromPtrOr(request.N, uint(1))) imageRequest.Parameters = AliImageParameters{ - N: int(lo.FromPtrOr(request.N, uint(1))), + N: &n, Watermark: request.Watermark, } return &imageRequest, nil diff --git a/relay/channel/ali/image_wan.go b/relay/channel/ali/image_wan.go index 8a256aba..0bb3e1d5 100644 --- a/relay/channel/ali/image_wan.go +++ b/relay/channel/ali/image_wan.go @@ -16,7 +16,9 @@ func oaiFormEdit2WanxImageEdit(c *gin.Context, info *relaycommon.RelayInfo, requ var err error var imageRequest AliImageRequest imageRequest.Model = request.Model - imageRequest.ResponseFormat = request.ResponseFormat + if request.ResponseFormat != "" { + imageRequest.ResponseFormat = &request.ResponseFormat + } wanInput := WanImageInput{ Prompt: request.Prompt, } @@ -31,10 +33,11 @@ func oaiFormEdit2WanxImageEdit(c *gin.Context, info *relaycommon.RelayInfo, requ // N: int(request.N), //} imageRequest.Input = wanInput + n := int(lo.FromPtrOr(request.N, uint(1))) imageRequest.Parameters = AliImageParameters{ - N: int(lo.FromPtrOr(request.N, uint(1))), + N: &n, } - info.PriceData.AddOtherRatio("n", float64(imageRequest.Parameters.N)) + info.PriceData.AddOtherRatio("n", float64(imageRequest.Parameters.NValue())) return &imageRequest, nil } diff --git a/relay/channel/ali/rerank.go b/relay/channel/ali/rerank.go index 8bc0b39a..7053c523 100644 --- a/relay/channel/ali/rerank.go +++ b/relay/channel/ali/rerank.go @@ -1,10 +1,10 @@ package ali import ( - "encoding/json" "io" "net/http" + "github.com/MAX-API-Next/MAX-API/common" "github.com/MAX-API-Next/MAX-API/dto" relaycommon "github.com/MAX-API-Next/MAX-API/relay/common" "github.com/MAX-API-Next/MAX-API/service" @@ -40,7 +40,7 @@ func RerankHandler(c *gin.Context, resp *http.Response, info *relaycommon.RelayI service.CloseResponseBodyGracefully(resp) var aliResponse AliRerankResponse - err = json.Unmarshal(responseBody, &aliResponse) + err = common.Unmarshal(responseBody, &aliResponse) if err != nil { return types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError), nil } @@ -64,7 +64,7 @@ func RerankHandler(c *gin.Context, resp *http.Response, info *relaycommon.RelayI Usage: usage, } - jsonResponse, err := json.Marshal(rerankResponse) + jsonResponse, err := common.Marshal(rerankResponse) if err != nil { return types.NewError(err, types.ErrorCodeBadResponseBody), nil } diff --git a/relay/channel/ali/text.go b/relay/channel/ali/text.go index 48649e34..5ce733bc 100644 --- a/relay/channel/ali/text.go +++ b/relay/channel/ali/text.go @@ -2,7 +2,6 @@ package ali import ( "github.com/MAX-API-Next/MAX-API/dto" - "github.com/samber/lo" ) // https://help.aliyun.com/document_detail/613695.html?spm=a2c4g.2399480.0.0.1adb778fAdzP9w#341800c0f8w0r @@ -10,11 +9,14 @@ import ( const EnableSearchModelSuffix = "-internet" func requestOpenAI2Ali(request dto.GeneralOpenAIRequest) *dto.GeneralOpenAIRequest { - topP := lo.FromPtrOr(request.TopP, 0) - if topP >= 1 { - request.TopP = lo.ToPtr(0.999) - } else if topP <= 0 { - request.TopP = lo.ToPtr(0.001) + if request.TopP != nil { + if *request.TopP >= 1 { + value := 0.99 + request.TopP = &value + } else if *request.TopP <= 0 { + value := 0.01 + request.TopP = &value + } } return &request } diff --git a/relay/channel/ali/text_test.go b/relay/channel/ali/text_test.go new file mode 100644 index 00000000..781b849e --- /dev/null +++ b/relay/channel/ali/text_test.go @@ -0,0 +1,132 @@ +package ali + +import ( + "encoding/json" + "testing" + + "github.com/MAX-API-Next/MAX-API/common" + "github.com/MAX-API-Next/MAX-API/dto" + relaycommon "github.com/MAX-API-Next/MAX-API/relay/common" + "github.com/samber/lo" + "github.com/stretchr/testify/assert" +) + +func TestRequestOpenAI2AliPreservesOptionalTopP(t *testing.T) { + tests := []struct { + name string + topP *float64 + want *float64 + }{ + {name: "omitted", topP: nil, want: nil}, + {name: "in range", topP: lo.ToPtr(0.5), want: lo.ToPtr(0.5)}, + {name: "zero", topP: lo.ToPtr(0.0), want: lo.ToPtr(0.01)}, + {name: "negative", topP: lo.ToPtr(-0.1), want: lo.ToPtr(0.01)}, + {name: "one", topP: lo.ToPtr(1.0), want: lo.ToPtr(0.99)}, + {name: "above one", topP: lo.ToPtr(1.2), want: lo.ToPtr(0.99)}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := requestOpenAI2Ali(dto.GeneralOpenAIRequest{Model: "qwen-plus", TopP: tt.topP}) + assert.Equal(t, tt.want, got.TopP) + }) + } +} + +func TestAliProviderRequestsPreserveExplicitZeroValues(t *testing.T) { + zeroFloat := 0.0 + zeroInt := 0 + zeroUint := uint64(0) + falseValue := false + empty := "" + + chatPayload, err := common.Marshal(AliParameters{ + TopP: &zeroFloat, + TopK: &zeroInt, + Seed: &zeroUint, + EnableSearch: &falseValue, + IncrementalOutput: &falseValue, + }) + assert.NoError(t, err) + assert.JSONEq(t, `{"top_p":0,"top_k":0,"seed":0,"enable_search":false,"incremental_output":false}`, string(chatPayload)) + + imagePayload, err := common.Marshal(AliImageRequest{ + Model: "wan", + Input: map[string]string{"prompt": "draw"}, + ResponseFormat: &empty, + Parameters: AliImageParameters{ + Size: &empty, + N: &zeroInt, + Steps: &empty, + Scale: &empty, + }, + }) + assert.NoError(t, err) + var imageBody map[string]json.RawMessage + assert.NoError(t, common.Unmarshal(imagePayload, &imageBody)) + assert.JSONEq(t, `""`, string(imageBody["response_format"])) + assert.JSONEq(t, `{"size":"","n":0,"steps":"","scale":""}`, string(imageBody["parameters"])) +} + +func TestOAIImageToAliOmitsAbsentSize(t *testing.T) { + var request dto.ImageRequest + assert.NoError(t, common.Unmarshal([]byte(`{"model":"wanx-v1","prompt":"draw"}`), &request)) + assert.NotNil(t, request.Extra) + + converted, err := oaiImage2AliImageRequest(&relaycommon.RelayInfo{}, request, false) + assert.NoError(t, err) + payload, err := common.Marshal(converted.Parameters) + assert.NoError(t, err) + + var parameters map[string]json.RawMessage + assert.NoError(t, common.Unmarshal(payload, ¶meters)) + assert.NotContains(t, parameters, "size") +} + +func TestOAIImageToAliUsesStandardParametersWhenExtraIsNil(t *testing.T) { + n := uint(3) + watermark := false + info := &relaycommon.RelayInfo{} + + converted, err := oaiImage2AliImageRequest(info, dto.ImageRequest{ + Model: "wanx-v1", + Prompt: "draw", + N: &n, + Size: "1024x1024", + Watermark: &watermark, + }, false) + + assert.NoError(t, err) + if assert.NotNil(t, converted) { + assert.Equal(t, 3, converted.Parameters.NValue()) + if assert.NotNil(t, converted.Parameters.Size) { + assert.Equal(t, "1024*1024", *converted.Parameters.Size) + } + if assert.NotNil(t, converted.Parameters.Watermark) { + assert.False(t, *converted.Parameters.Watermark) + } + } + assert.Equal(t, float64(3), info.PriceData.OtherRatios["n"]) +} + +func TestOAIImageToAliFallsBackToStandardNWithCustomParameters(t *testing.T) { + n := uint(3) + info := &relaycommon.RelayInfo{} + + converted, err := oaiImage2AliImageRequest(info, dto.ImageRequest{ + Model: "wanx-v1", + Prompt: "draw", + N: &n, + Extra: map[string]json.RawMessage{ + "parameters": json.RawMessage(`{"steps":"4"}`), + }, + }, false) + + assert.NoError(t, err) + if assert.NotNil(t, converted) { + assert.Equal(t, 3, converted.Parameters.NValue()) + assert.NotNil(t, converted.Parameters.Steps) + assert.Equal(t, "4", *converted.Parameters.Steps) + } + assert.Equal(t, float64(3), info.PriceData.OtherRatios["n"]) +} diff --git a/relay/channel/api_request.go b/relay/channel/api_request.go index 30f28a8f..2ce9b5b8 100644 --- a/relay/channel/api_request.go +++ b/relay/channel/api_request.go @@ -27,19 +27,20 @@ import ( "github.com/gorilla/websocket" ) -// applyUpstreamContentLength populates req.ContentLength when the upstream -// body is wrapped in a BodyStorage (see relay/common/outbound_body.go). -// -// net/http.NewRequest only auto-detects ContentLength for *bytes.Reader, -// *bytes.Buffer and *strings.Reader. When the body is a type-erased io.Reader -// (which is the case for ReaderOnly(BodyStorage)), the Content-Length header -// would otherwise be omitted, forcing chunked transfer encoding and breaking -// some upstreams that require an explicit Content-Length. -func applyUpstreamContentLength(req *http.Request, info *common.RelayInfo) { - if info == nil { +// applyUpstreamBodyMetadata restores metadata hidden when http.NewRequest +// wraps an arbitrary reader in req.Body. +func applyUpstreamBodyMetadata(req *http.Request, body io.Reader, info *common.RelayInfo) { + if replayable, ok := body.(common2.ReplayableBody); ok { + if _, rawStorage := body.(common2.BodyStorage); rawStorage { + req.Body = io.NopCloser(body) + } + req.ContentLength = replayable.Size() + if req.GetBody == nil { + req.GetBody = replayable.NewReader + } return } - if info.UpstreamRequestBodySize > 0 && req.ContentLength <= 0 { + if info != nil && info.UpstreamRequestBodySize > 0 && req.ContentLength <= 0 { req.ContentLength = info.UpstreamRequestBodySize } } @@ -322,7 +323,7 @@ func DoApiRequest(a Adaptor, c *gin.Context, info *common.RelayInfo, requestBody if err != nil { return nil, fmt.Errorf("new request failed: %w", err) } - applyUpstreamContentLength(req, info) + applyUpstreamBodyMetadata(req, requestBody, info) headers := req.Header err = a.SetupRequestHeader(c, &headers, info) if err != nil { @@ -352,7 +353,7 @@ func DoFormRequest(a Adaptor, c *gin.Context, info *common.RelayInfo, requestBod if err != nil { return nil, fmt.Errorf("new request failed: %w", err) } - applyUpstreamContentLength(req, info) + applyUpstreamBodyMetadata(req, requestBody, info) // set form data req.Header.Set("Content-Type", c.Request.Header.Get("Content-Type")) headers := req.Header @@ -559,44 +560,10 @@ func newTaskHTTPRequest(method string, fullRequestURL string, requestBody io.Rea if err != nil { return nil, err } - applyUpstreamContentLength(req, info) - if req.GetBody == nil { - attachSeekableGetBody(req, requestBody) - } + applyUpstreamBodyMetadata(req, requestBody, info) return req, nil } -func attachSeekableGetBody(req *http.Request, reader io.Reader) { - if req == nil { - return - } - seeker, ok := reader.(io.ReadSeeker) - if !ok { - return - } - start, err := seeker.Seek(0, io.SeekCurrent) - if err != nil { - return - } - end, err := seeker.Seek(0, io.SeekEnd) - if err != nil { - _, _ = seeker.Seek(start, io.SeekStart) - return - } - if _, err := seeker.Seek(start, io.SeekStart); err != nil { - return - } - if req.ContentLength <= 0 && end >= start { - req.ContentLength = end - start - } - req.GetBody = func() (io.ReadCloser, error) { - if _, err := seeker.Seek(start, io.SeekStart); err != nil { - return nil, err - } - return io.NopCloser(seeker), nil - } -} - func DoTaskApiRequest(a TaskAdaptor, c *gin.Context, info *common.RelayInfo, requestBody io.Reader) (*http.Response, error) { if info != nil && info.ChannelMeta == nil { info.InitChannelMeta(c) diff --git a/relay/channel/api_request_test.go b/relay/channel/api_request_test.go index a34a0608..06f4ec69 100644 --- a/relay/channel/api_request_test.go +++ b/relay/channel/api_request_test.go @@ -129,7 +129,7 @@ func TestNewTaskHTTPRequestDoesNotPreReadRequestBody(t *testing.T) { require.Equal(t, 1, body.reads) } -func TestNewTaskHTTPRequestUsesSeekableGetBodyWithoutBuffering(t *testing.T) { +func TestNewTaskHTTPRequestDoesNotInventReplayForUnknownSeeker(t *testing.T) { t.Parallel() body := &customReadSeeker{reader: strings.NewReader("abcdef")} @@ -139,14 +139,39 @@ func TestNewTaskHTTPRequestUsesSeekableGetBodyWithoutBuffering(t *testing.T) { req, err := newTaskHTTPRequest(http.MethodPost, "https://example.com/tasks", body, &relaycommon.RelayInfo{}) require.NoError(t, err) - require.Equal(t, int64(4), req.ContentLength) + require.Zero(t, req.ContentLength) + require.Nil(t, req.GetBody) +} + +func TestApplyUpstreamBodyMetadataUsesIndependentReplayReaders(t *testing.T) { + payload := []byte(`{"model":"test","input":"abcdef"}`) + storage, err := common2.CreateBodyStorage(payload) + require.NoError(t, err) + defer storage.Close() + body := common2.NewReplayableBodyReader(storage) + + req, err := http.NewRequest(http.MethodPost, "https://example.com", body) + require.NoError(t, err) + applyUpstreamBodyMetadata(req, body, nil) + require.Equal(t, int64(len(payload)), req.ContentLength) require.NotNil(t, req.GetBody) - replay, err := req.GetBody() + first, err := req.GetBody() + require.NoError(t, err) + defer first.Close() + second, err := req.GetBody() + require.NoError(t, err) + defer second.Close() + + prefix := make([]byte, 5) + _, err = io.ReadFull(first, prefix) + require.NoError(t, err) + secondBody, err := io.ReadAll(second) require.NoError(t, err) - replayed, err := io.ReadAll(replay) + require.Equal(t, payload, secondBody) + firstRest, err := io.ReadAll(first) require.NoError(t, err) - require.Equal(t, "cdef", string(replayed)) + require.Equal(t, payload[5:], firstRest) } func TestProcessHeaderOverride_ChannelTestSkipsPassthroughRules(t *testing.T) { diff --git a/relay/channel/baidu/dto.go b/relay/channel/baidu/dto.go index bf37efd8..67cfb94e 100644 --- a/relay/channel/baidu/dto.go +++ b/relay/channel/baidu/dto.go @@ -19,8 +19,8 @@ type BaiduChatRequest struct { PenaltyScore *float64 `json:"penalty_score,omitempty"` Stream *bool `json:"stream,omitempty"` System string `json:"system,omitempty"` - DisableSearch bool `json:"disable_search,omitempty"` - EnableCitation bool `json:"enable_citation,omitempty"` + DisableSearch *bool `json:"disable_search,omitempty"` + EnableCitation *bool `json:"enable_citation,omitempty"` MaxOutputTokens *int `json:"max_output_tokens,omitempty"` UserId json.RawMessage `json:"user_id,omitempty"` } diff --git a/relay/channel/baidu/relay-baidu.go b/relay/channel/baidu/relay-baidu.go index 1f27fac0..af6f81b3 100644 --- a/relay/channel/baidu/relay-baidu.go +++ b/relay/channel/baidu/relay-baidu.go @@ -1,7 +1,6 @@ package baidu import ( - "encoding/json" "errors" "fmt" "io" @@ -26,13 +25,11 @@ var baiduTokenStore sync.Map func requestOpenAI2Baidu(request dto.GeneralOpenAIRequest) *BaiduChatRequest { baiduRequest := BaiduChatRequest{ - Temperature: request.Temperature, - TopP: request.TopP, - PenaltyScore: request.FrequencyPenalty, - Stream: request.Stream, - DisableSearch: false, - EnableCitation: false, - UserId: request.User, + Temperature: request.Temperature, + TopP: request.TopP, + PenaltyScore: request.FrequencyPenalty, + Stream: request.Stream, + UserId: request.User, } if request.GetMaxTokens() != 0 { maxTokens := int(request.GetMaxTokens()) @@ -143,7 +140,7 @@ func baiduHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Respon return types.NewError(err, types.ErrorCodeBadResponseBody), nil } service.CloseResponseBodyGracefully(resp) - err = json.Unmarshal(responseBody, &baiduResponse) + err = common.Unmarshal(responseBody, &baiduResponse) if err != nil { return types.NewError(err, types.ErrorCodeBadResponseBody), nil } @@ -151,7 +148,7 @@ func baiduHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Respon return types.NewError(fmt.Errorf("%s", baiduResponse.ErrorMsg), types.ErrorCodeBadResponseBody), nil } fullTextResponse := responseBaidu2OpenAI(&baiduResponse) - jsonResponse, err := json.Marshal(fullTextResponse) + jsonResponse, err := common.Marshal(fullTextResponse) if err != nil { return types.NewError(err, types.ErrorCodeBadResponseBody), nil } @@ -171,7 +168,7 @@ func baiduEmbeddingHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *ht return types.NewError(err, types.ErrorCodeBadResponseBody), nil } service.CloseResponseBodyGracefully(resp) - err = json.Unmarshal(responseBody, &baiduResponse) + err = common.Unmarshal(responseBody, &baiduResponse) if err != nil { return types.NewError(err, types.ErrorCodeBadResponseBody), nil } @@ -179,7 +176,7 @@ func baiduEmbeddingHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *ht return types.NewError(fmt.Errorf("%s", baiduResponse.ErrorMsg), types.ErrorCodeBadResponseBody), nil } fullTextResponse := embeddingResponseBaidu2OpenAI(&baiduResponse) - jsonResponse, err := json.Marshal(fullTextResponse) + jsonResponse, err := common.Marshal(fullTextResponse) if err != nil { return types.NewError(err, types.ErrorCodeBadResponseBody), nil } diff --git a/relay/channel/baidu/relay_baidu_test.go b/relay/channel/baidu/relay_baidu_test.go index 4a6391bf..e3a6b851 100644 --- a/relay/channel/baidu/relay_baidu_test.go +++ b/relay/channel/baidu/relay_baidu_test.go @@ -40,8 +40,10 @@ func TestRequestOpenAI2BaiduPreservesExplicitZeroValues(t *testing.T) { FrequencyPenalty: &zero, Stream: &stream, }) + request.DisableSearch = &stream + request.EnableCitation = &stream payload, err := common.Marshal(request) require.NoError(t, err) - require.JSONEq(t, `{"messages":null,"top_p":0,"penalty_score":0,"stream":false}`, string(payload)) + require.JSONEq(t, `{"messages":null,"top_p":0,"penalty_score":0,"stream":false,"disable_search":false,"enable_citation":false}`, string(payload)) } diff --git a/relay/channel/coze/adaptor.go b/relay/channel/coze/adaptor.go index bc899ec6..c6e6f4b2 100644 --- a/relay/channel/coze/adaptor.go +++ b/relay/channel/coze/adaptor.go @@ -1,13 +1,13 @@ package coze import ( - "encoding/json" "errors" "fmt" "io" "net/http" "time" + common2 "github.com/MAX-API-Next/MAX-API/common" "github.com/MAX-API-Next/MAX-API/dto" "github.com/MAX-API-Next/MAX-API/relay/channel" "github.com/MAX-API-Next/MAX-API/relay/common" @@ -73,13 +73,19 @@ func (a *Adaptor) DoRequest(c *gin.Context, info *common.RelayInfo, requestBody if err != nil { return nil, err } + if resp == nil || resp.Body == nil { + return nil, errors.New("Coze create-message response body is nil") + } // 解析 resp - var cozeResponse CozeChatResponse - respBody, err := io.ReadAll(resp.Body) + respBody, readErr := io.ReadAll(resp.Body) + _ = resp.Body.Close() + if readErr != nil { + return nil, fmt.Errorf("read Coze create-message response: %w", readErr) + } + cozeResponse, err := decodeCozeCreateResponse(respBody) if err != nil { return nil, err } - err = json.Unmarshal(respBody, &cozeResponse) if cozeResponse.Code != 0 { return nil, errors.New(cozeResponse.Msg) } @@ -101,6 +107,14 @@ func (a *Adaptor) DoRequest(c *gin.Context, info *common.RelayInfo, requestBody return getChatDetail(a, c, info) } +func decodeCozeCreateResponse(body []byte) (CozeChatResponse, error) { + var response CozeChatResponse + if err := common2.Unmarshal(body, &response); err != nil { + return response, fmt.Errorf("decode Coze create-message response: %w", err) + } + return response, nil +} + // DoResponse implements channel.Adaptor. func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *common.RelayInfo) (usage any, err *types.MaxAPIError) { if info.IsStream { diff --git a/relay/channel/coze/dto.go b/relay/channel/coze/dto.go index da01bbb6..d8cba54a 100644 --- a/relay/channel/coze/dto.go +++ b/relay/channel/coze/dto.go @@ -19,9 +19,9 @@ type CozeChatRequest struct { BotId string `json:"bot_id"` UserId json.RawMessage `json:"user_id"` AdditionalMessages []CozeEnterMessage `json:"additional_messages,omitempty"` - Stream bool `json:"stream,omitempty"` + Stream *bool `json:"stream,omitempty"` CustomVariables json.RawMessage `json:"custom_variables,omitempty"` - AutoSaveHistory bool `json:"auto_save_history,omitempty"` + AutoSaveHistory *bool `json:"auto_save_history,omitempty"` MetaData json.RawMessage `json:"meta_data,omitempty"` ExtraParams json.RawMessage `json:"extra_params,omitempty"` ShortcutCommand json.RawMessage `json:"shortcut_command,omitempty"` diff --git a/relay/channel/coze/relay-coze.go b/relay/channel/coze/relay-coze.go index 3f8d20d6..739329a6 100644 --- a/relay/channel/coze/relay-coze.go +++ b/relay/channel/coze/relay-coze.go @@ -15,7 +15,6 @@ import ( "github.com/MAX-API-Next/MAX-API/relay/helper" "github.com/MAX-API-Next/MAX-API/service" "github.com/MAX-API-Next/MAX-API/types" - "github.com/samber/lo" "github.com/gin-gonic/gin" ) @@ -41,7 +40,7 @@ func convertCozeChatRequest(c *gin.Context, request dto.GeneralOpenAIRequest) *C BotId: c.GetString("bot_id"), UserId: user, AdditionalMessages: messages, - Stream: lo.FromPtrOr(request.Stream, false), + Stream: request.Stream, } return cozeRequest } @@ -56,7 +55,7 @@ func cozeChatHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Res var response dto.TextResponse var cozeResponse CozeChatDetailResponse response.Model = info.UpstreamModelName - err = json.Unmarshal(responseBody, &cozeResponse) + err = common.Unmarshal(responseBody, &cozeResponse) if err != nil { return nil, types.NewError(err, types.ErrorCodeBadResponseBody) } @@ -86,7 +85,7 @@ func cozeChatHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Res FinishReason: "stop", }, } - jsonResponse, err := json.Marshal(response) + jsonResponse, err := common.Marshal(response) if err != nil { return nil, types.NewError(err, types.ErrorCodeBadResponseBody) } @@ -154,7 +153,7 @@ func handleCozeEvent(c *gin.Context, event string, data string, responseText *st case "conversation.chat.completed": // 将 data 解析为 CozeChatResponseData var chatData CozeChatResponseData - err := json.Unmarshal([]byte(data), &chatData) + err := common.Unmarshal([]byte(data), &chatData) if err != nil { common.SysLog("error_unmarshalling_stream_response: " + err.Error()) return @@ -171,14 +170,14 @@ func handleCozeEvent(c *gin.Context, event string, data string, responseText *st case "conversation.message.delta": // 将 data 解析为 CozeChatV3MessageDetail var messageData CozeChatV3MessageDetail - err := json.Unmarshal([]byte(data), &messageData) + err := common.Unmarshal([]byte(data), &messageData) if err != nil { common.SysLog("error_unmarshalling_stream_response: " + err.Error()) return } var content string - err = json.Unmarshal(messageData.Content, &content) + err = common.Unmarshal(messageData.Content, &content) if err != nil { common.SysLog("error_unmarshalling_stream_response: " + err.Error()) return @@ -203,7 +202,7 @@ func handleCozeEvent(c *gin.Context, event string, data string, responseText *st case "error": var errorData CozeError - err := json.Unmarshal([]byte(data), &errorData) + err := common.Unmarshal([]byte(data), &errorData) if err != nil { common.SysLog("error_unmarshalling_stream_response: " + err.Error()) return @@ -242,7 +241,7 @@ func checkIfChatComplete(a *Adaptor, c *gin.Context, info *relaycommon.RelayInfo if err != nil { return fmt.Errorf("read response body failed: %w", err), false } - err = json.Unmarshal(responseBody, &cozeResponse) + err = common.Unmarshal(responseBody, &cozeResponse) if err != nil { return fmt.Errorf("unmarshal response body failed: %w", err), false } diff --git a/relay/channel/coze/relay_coze_test.go b/relay/channel/coze/relay_coze_test.go new file mode 100644 index 00000000..f8c8b095 --- /dev/null +++ b/relay/channel/coze/relay_coze_test.go @@ -0,0 +1,28 @@ +package coze + +import ( + "encoding/json" + "net/http/httptest" + "testing" + + "github.com/MAX-API-Next/MAX-API/common" + "github.com/MAX-API-Next/MAX-API/dto" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +func TestConvertCozeChatRequestPreservesExplicitFalseStream(t *testing.T) { + stream := false + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + request := convertCozeChatRequest(c, dto.GeneralOpenAIRequest{Stream: &stream, User: json.RawMessage(`"user"`)}) + request.AutoSaveHistory = &stream + + data, err := common.Marshal(request) + require.NoError(t, err) + require.JSONEq(t, `{"bot_id":"","user_id":"user","stream":false,"auto_save_history":false}`, string(data)) +} + +func TestDecodeCozeCreateResponseRejectsMalformedJSON(t *testing.T) { + _, err := decodeCozeCreateResponse([]byte(`{"code":`)) + require.ErrorContains(t, err, "decode Coze create-message response") +} diff --git a/relay/channel/gemini/relay-gemini.go b/relay/channel/gemini/relay-gemini.go index c47656bb..007776af 100644 --- a/relay/channel/gemini/relay-gemini.go +++ b/relay/channel/gemini/relay-gemini.go @@ -1779,6 +1779,13 @@ func GeminiImageHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http. if len(geminiResponse.Predictions) == 0 { return nil, types.NewOpenAIError(errors.New("no images generated"), types.ErrorCodeBadResponseBody, http.StatusInternalServerError) } + if len(geminiResponse.Predictions) > dto.MaxImageN { + return nil, types.NewOpenAIError( + fmt.Errorf("too many image predictions: got %d, maximum %d", len(geminiResponse.Predictions), dto.MaxImageN), + types.ErrorCodeBadResponseBody, + http.StatusBadGateway, + ) + } // convert to openai format response openAIResponse := dto.ImageResponse{ diff --git a/relay/channel/gemini/relay_gemini_usage_test.go b/relay/channel/gemini/relay_gemini_usage_test.go index 105f68ef..b9a72307 100644 --- a/relay/channel/gemini/relay_gemini_usage_test.go +++ b/relay/channel/gemini/relay_gemini_usage_test.go @@ -67,6 +67,28 @@ func TestGeminiChatHandlerCompletionTokensExcludeToolUsePromptTokens(t *testing. require.Equal(t, 1120, usage.CompletionTokenDetails.ReasoningTokens) } +func TestGeminiImageHandlerRejectsTooManyPredictions(t *testing.T) { + gin.SetMode(gin.TestMode) + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/generations", nil) + + payload := dto.GeminiImageResponse{ + Predictions: make([]dto.GeminiImagePrediction, dto.MaxImageN+1), + } + body, err := common.Marshal(payload) + require.NoError(t, err) + + resp := &http.Response{ + StatusCode: http.StatusOK, + Body: io.NopCloser(bytes.NewReader(body)), + } + + usage, maxAPIError := GeminiImageHandler(c, &relaycommon.RelayInfo{}, resp) + require.Nil(t, usage) + require.NotNil(t, maxAPIError) + require.Contains(t, maxAPIError.Error(), "too many image predictions") +} + func TestGeminiStreamHandlerCompletionTokensExcludeToolUsePromptTokens(t *testing.T) { gin.SetMode(gin.TestMode) c, _ := gin.CreateTestContext(httptest.NewRecorder()) diff --git a/relay/channel/minimax/tts.go b/relay/channel/minimax/tts.go index a19713f6..9f571289 100644 --- a/relay/channel/minimax/tts.go +++ b/relay/channel/minimax/tts.go @@ -190,15 +190,7 @@ func handleChatCompletionResponse(c *gin.Context, resp *http.Response, info *rel } defer resp.Body.Close() - // Set response headers - for key, values := range resp.Header { - if !service.ShouldCopyUpstreamHeader(c, key, values) { - continue - } - for _, value := range values { - c.Header(key, value) - } - } + service.CopyUpstreamResponseHeaders(c, resp.Header) c.Data(resp.StatusCode, "application/json", body) return nil, nil diff --git a/relay/channel/openai/audio.go b/relay/channel/openai/audio.go index f63977f8..06ddacde 100644 --- a/relay/channel/openai/audio.go +++ b/relay/channel/openai/audio.go @@ -46,12 +46,7 @@ func OpenaiTTSHandler(c *gin.Context, resp *http.Response, info *relaycommon.Rel usage := &dto.Usage{} usage.PromptTokens = info.GetEstimatePromptTokens() usage.TotalTokens = info.GetEstimatePromptTokens() - for k, v := range resp.Header { - if !service.ShouldCopyUpstreamHeader(c, k, v) { - continue - } - c.Writer.Header().Set(k, v[0]) - } + service.CopyUpstreamResponseHeaders(c, resp.Header) c.Writer.WriteHeader(resp.StatusCode) if info.IsStream { diff --git a/relay/channel/siliconflow/adaptor.go b/relay/channel/siliconflow/adaptor.go index b3679d10..9a3b1f32 100644 --- a/relay/channel/siliconflow/adaptor.go +++ b/relay/channel/siliconflow/adaptor.go @@ -15,7 +15,6 @@ import ( "github.com/MAX-API-Next/MAX-API/types" "github.com/gin-gonic/gin" - "github.com/samber/lo" ) type Adaptor struct { @@ -50,13 +49,11 @@ func (a *Adaptor) ConvertImageRequest(c *gin.Context, info *relaycommon.RelayInf sfRequest.Model = request.Model sfRequest.Prompt = request.Prompt // 优先使用image_size/batch_size，否则使用OpenAI标准的size/n - if sfRequest.ImageSize == "" { - sfRequest.ImageSize = request.Size + if sfRequest.ImageSize == nil && request.Size != "" { + sfRequest.ImageSize = &request.Size } - if sfRequest.BatchSize == 0 { - if request.N != nil { - sfRequest.BatchSize = lo.FromPtr(request.N) - } + if sfRequest.BatchSize == nil && request.N != nil { + sfRequest.BatchSize = request.N } return sfRequest, nil diff --git a/relay/channel/siliconflow/adaptor_test.go b/relay/channel/siliconflow/adaptor_test.go new file mode 100644 index 00000000..6f1e7779 --- /dev/null +++ b/relay/channel/siliconflow/adaptor_test.go @@ -0,0 +1,38 @@ +package siliconflow + +import ( + "encoding/json" + "net/http/httptest" + "testing" + + "github.com/MAX-API-Next/MAX-API/common" + "github.com/MAX-API-Next/MAX-API/dto" + relaycommon "github.com/MAX-API-Next/MAX-API/relay/common" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +func TestConvertImageRequestPreservesExplicitZeroProviderFields(t *testing.T) { + request := dto.ImageRequest{ + Model: "black-forest-labs/FLUX.1-schnell", + Prompt: "draw", + Extra: map[string]json.RawMessage{ + "batch_size": []byte("0"), + "seed": []byte("0"), + "num_inference_steps": []byte("0"), + "guidance_scale": []byte("0"), + "cfg": []byte("0"), + }, + } + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + converted, err := (&Adaptor{}).ConvertImageRequest(c, &relaycommon.RelayInfo{}, request) + require.NoError(t, err) + + data, err := common.Marshal(converted) + require.NoError(t, err) + var body map[string]json.RawMessage + require.NoError(t, common.Unmarshal(data, &body)) + for _, field := range []string{"batch_size", "seed", "num_inference_steps", "guidance_scale", "cfg"} { + require.JSONEq(t, "0", string(body[field]), field) + } +} diff --git a/relay/channel/siliconflow/dto.go b/relay/channel/siliconflow/dto.go index a217535a..1dffddfc 100644 --- a/relay/channel/siliconflow/dto.go +++ b/relay/channel/siliconflow/dto.go @@ -17,16 +17,16 @@ type SFRerankResponse struct { } type SFImageRequest struct { - Model string `json:"model"` - Prompt string `json:"prompt"` - NegativePrompt string `json:"negative_prompt,omitempty"` - ImageSize string `json:"image_size,omitempty"` - BatchSize uint `json:"batch_size,omitempty"` - Seed uint64 `json:"seed,omitempty"` - NumInferenceSteps uint `json:"num_inference_steps,omitempty"` - GuidanceScale float64 `json:"guidance_scale,omitempty"` - Cfg float64 `json:"cfg,omitempty"` - Image string `json:"image,omitempty"` - Image2 string `json:"image2,omitempty"` - Image3 string `json:"image3,omitempty"` + Model string `json:"model"` + Prompt string `json:"prompt"` + NegativePrompt *string `json:"negative_prompt,omitempty"` + ImageSize *string `json:"image_size,omitempty"` + BatchSize *uint `json:"batch_size,omitempty"` + Seed *uint64 `json:"seed,omitempty"` + NumInferenceSteps *uint `json:"num_inference_steps,omitempty"` + GuidanceScale *float64 `json:"guidance_scale,omitempty"` + Cfg *float64 `json:"cfg,omitempty"` + Image *string `json:"image,omitempty"` + Image2 *string `json:"image2,omitempty"` + Image3 *string `json:"image3,omitempty"` } diff --git a/relay/channel/siliconflow/relay-siliconflow.go b/relay/channel/siliconflow/relay-siliconflow.go index 7d16d806..a7749dca 100644 --- a/relay/channel/siliconflow/relay-siliconflow.go +++ b/relay/channel/siliconflow/relay-siliconflow.go @@ -1,10 +1,10 @@ package siliconflow import ( - "encoding/json" "io" "net/http" + "github.com/MAX-API-Next/MAX-API/common" "github.com/MAX-API-Next/MAX-API/dto" relaycommon "github.com/MAX-API-Next/MAX-API/relay/common" "github.com/MAX-API-Next/MAX-API/service" @@ -20,7 +20,7 @@ func siliconflowRerankHandler(c *gin.Context, info *relaycommon.RelayInfo, resp } service.CloseResponseBodyGracefully(resp) var siliconflowResp SFRerankResponse - err = json.Unmarshal(responseBody, &siliconflowResp) + err = common.Unmarshal(responseBody, &siliconflowResp) if err != nil { return nil, types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError) } @@ -34,7 +34,7 @@ func siliconflowRerankHandler(c *gin.Context, info *relaycommon.RelayInfo, resp Usage: *usage, } - jsonResponse, err := json.Marshal(rerankResp) + jsonResponse, err := common.Marshal(rerankResp) if err != nil { return nil, types.NewError(err, types.ErrorCodeBadResponseBody) } diff --git a/relay/channel/vertex/adaptor.go b/relay/channel/vertex/adaptor.go index 45538fc1..ccf2a846 100644 --- a/relay/channel/vertex/adaptor.go +++ b/relay/channel/vertex/adaptor.go @@ -1,7 +1,6 @@ package vertex import ( - "encoding/json" "errors" "fmt" "io" @@ -267,7 +266,7 @@ func (a *Adaptor) ConvertOpenAIRequest(c *gin.Context, info *relaycommon.RelayIn } if len(request.ExtraBody) > 0 { var extra map[string]any - if err := json.Unmarshal(request.ExtraBody, &extra); err == nil { + if err := common.Unmarshal(request.ExtraBody, &extra); err == nil { if n, ok := extra["n"].(float64); ok && n > 0 { imgReq.N = lo.ToPtr(uint(n)) } diff --git a/relay/channel/vertex/service_account.go b/relay/channel/vertex/service_account.go index 1571bead..56e3e31d 100644 --- a/relay/channel/vertex/service_account.go +++ b/relay/channel/vertex/service_account.go @@ -3,13 +3,13 @@ package vertex import ( "crypto/rsa" "crypto/x509" - "encoding/json" "encoding/pem" "errors" "net/http" "net/url" "strings" + "github.com/MAX-API-Next/MAX-API/common" relaycommon "github.com/MAX-API-Next/MAX-API/relay/common" "github.com/MAX-API-Next/MAX-API/service" @@ -46,7 +46,10 @@ func getAccessToken(a *Adaptor, info *relaycommon.RelayInfo) (string, error) { } val, err := Cache.Get(cacheKey) if err == nil { - return val.(string), nil + if token, ok := val.(string); ok && token != "" { + return token, nil + } + Cache.DeleteIf(func(key string) bool { return key == cacheKey }) } signedJWT, err := createSignedJWT(a.AccountCredentials.ClientEmail, a.AccountCredentials.PrivateKey) @@ -129,7 +132,7 @@ func exchangeJwtForAccessToken(signedJWT string, info *relaycommon.RelayInfo) (s defer resp.Body.Close() var result map[string]interface{} - if err := json.NewDecoder(resp.Body).Decode(&result); err != nil { + if err := common.DecodeJson(resp.Body, &result); err != nil { return "", err } @@ -172,7 +175,7 @@ func exchangeJwtForAccessTokenWithProxy(signedJWT string, proxy string) (string, defer resp.Body.Close() var result map[string]interface{} - if err := json.NewDecoder(resp.Body).Decode(&result); err != nil { + if err := common.DecodeJson(resp.Body, &result); err != nil { return "", err } diff --git a/relay/channel/vertex/service_account_test.go b/relay/channel/vertex/service_account_test.go new file mode 100644 index 00000000..37bd0cdd --- /dev/null +++ b/relay/channel/vertex/service_account_test.go @@ -0,0 +1,29 @@ +package vertex + +import ( + "testing" + + relaycommon "github.com/MAX-API-Next/MAX-API/relay/common" + "github.com/stretchr/testify/require" +) + +func TestGetAccessTokenRejectsUnexpectedCachedValueWithoutPanicking(t *testing.T) { + const channelID = 987654321 + cacheKey := "access-token-987654321" + Cache.DeleteIf(func(key string) bool { return key == cacheKey }) + require.False(t, Cache.SetDefault(cacheKey, 123)) + t.Cleanup(func() { Cache.DeleteIf(func(key string) bool { return key == cacheKey }) }) + + adaptor := &Adaptor{} + info := &relaycommon.RelayInfo{ChannelMeta: &relaycommon.ChannelMeta{ChannelId: channelID}} + var err error + require.NotPanics(t, func() { + _, err = getAccessToken(adaptor, info) + }) + require.ErrorContains(t, err, "failed to create signed JWT") +} + +func TestCreateSignedJWTRejectsMalformedPEM(t *testing.T) { + _, err := createSignedJWT("service@example.com", "not-a-private-key") + require.ErrorContains(t, err, "failed to parse PEM block") +} diff --git a/relay/channel/zhipu/adaptor.go b/relay/channel/zhipu/adaptor.go index e710310b..4131625f 100644 --- a/relay/channel/zhipu/adaptor.go +++ b/relay/channel/zhipu/adaptor.go @@ -50,7 +50,10 @@ func (a *Adaptor) GetRequestURL(info *relaycommon.RelayInfo) (string, error) { func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Header, info *relaycommon.RelayInfo) error { channel.SetupApiRequestHeader(info, c, req) - token := getZhipuToken(info.ApiKey) + token, err := getZhipuToken(info.ApiKey) + if err != nil { + return err + } req.Set("Authorization", token) return nil } diff --git a/relay/channel/zhipu/relay-zhipu.go b/relay/channel/zhipu/relay-zhipu.go index f160d676..8d71d926 100644 --- a/relay/channel/zhipu/relay-zhipu.go +++ b/relay/channel/zhipu/relay-zhipu.go @@ -2,7 +2,7 @@ package zhipu import ( "bufio" - "encoding/json" + "fmt" "io" "net/http" "strings" @@ -28,19 +28,21 @@ import ( var zhipuTokens sync.Map var expSeconds int64 = 24 * 3600 -func getZhipuToken(apikey string) string { +func getZhipuToken(apikey string) (string, error) { data, ok := zhipuTokens.Load(apikey) if ok { - tokenData := data.(zhipuTokenData) - if time.Now().Before(tokenData.ExpiryTime) { - return tokenData.Token + tokenData, valid := data.(zhipuTokenData) + if valid && time.Now().Before(tokenData.ExpiryTime) { + return tokenData.Token, nil + } + if !valid { + zhipuTokens.Delete(apikey) } } - split := strings.Split(apikey, ".") - if len(split) != 2 { - common.SysLog("invalid zhipu key: " + apikey) - return "" + split := strings.SplitN(apikey, ".", 2) + if len(split) != 2 || strings.TrimSpace(split[0]) == "" || strings.TrimSpace(split[1]) == "" { + return "", fmt.Errorf("invalid zhipu API key format") } id := split[0] @@ -64,7 +66,7 @@ func getZhipuToken(apikey string) string { tokenString, err := token.SignedString([]byte(secret)) if err != nil { - return "" + return "", fmt.Errorf("sign zhipu access token: %w", err) } zhipuTokens.Store(apikey, zhipuTokenData{ @@ -72,7 +74,7 @@ func getZhipuToken(apikey string) string { ExpiryTime: expiryTime, }) - return tokenString + return tokenString, nil } func requestOpenAI2Zhipu(request dto.GeneralOpenAIRequest) *ZhipuRequest { @@ -160,11 +162,24 @@ func zhipuStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http. dataChan := make(chan string) metaChan := make(chan string) stopChan := make(chan struct{}, 1) + producerErrChan := make(chan error, 1) done := make(chan struct{}) producerDone := make(chan struct{}) go func() { + var producerErr error defer close(producerDone) defer func() { + if recovered := recover(); recovered != nil { + producerErr = fmt.Errorf("panic while reading Zhipu stream: %v", recovered) + } + if producerErr != nil { + common.SysError(producerErr.Error()) + select { + case producerErrChan <- producerErr: + case <-done: + } + return + } select { case stopChan <- struct{}{}: case <-done: @@ -205,15 +220,16 @@ func zhipuStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http. return default: } - common.SysLog("error reading stream: " + err.Error()) + producerErr = fmt.Errorf("error reading Zhipu stream: %w", err) } }() helper.SetEventStreamHeaders(c) + var producerErr error c.Stream(func(w io.Writer) bool { select { case data := <-dataChan: response := streamResponseZhipu2OpenAI(data) - jsonResponse, err := json.Marshal(response) + jsonResponse, err := common.Marshal(response) if err != nil { common.SysLog("error marshalling stream response: " + err.Error()) return true @@ -222,13 +238,13 @@ func zhipuStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http. return true case data := <-metaChan: var zhipuResponse ZhipuStreamMetaResponse - err := json.Unmarshal([]byte(data), &zhipuResponse) + err := common.Unmarshal([]byte(data), &zhipuResponse) if err != nil { common.SysLog("error unmarshalling stream response: " + err.Error()) return true } response, zhipuUsage := streamMetaResponseZhipu2OpenAI(&zhipuResponse) - jsonResponse, err := json.Marshal(response) + jsonResponse, err := common.Marshal(response) if err != nil { common.SysLog("error marshalling stream response: " + err.Error()) return true @@ -239,6 +255,8 @@ func zhipuStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http. case <-stopChan: c.Render(-1, common.CustomEvent{Data: "data: [DONE]"}) return false + case producerErr = <-producerErrChan: + return false case <-c.Request.Context().Done(): return false } @@ -246,6 +264,9 @@ func zhipuStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http. close(done) service.CloseResponseBodyGracefully(resp) <-producerDone + if producerErr != nil { + return nil, types.NewOpenAIError(producerErr, types.ErrorCodeReadResponseBodyFailed, http.StatusInternalServerError) + } return usage, nil } @@ -256,7 +277,7 @@ func zhipuHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Respon return nil, types.NewOpenAIError(err, types.ErrorCodeReadResponseBodyFailed, http.StatusInternalServerError) } service.CloseResponseBodyGracefully(resp) - err = json.Unmarshal(responseBody, &zhipuResponse) + err = common.Unmarshal(responseBody, &zhipuResponse) if err != nil { return nil, types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError) } @@ -267,7 +288,7 @@ func zhipuHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Respon }, resp.StatusCode) } fullTextResponse := responseZhipu2OpenAI(&zhipuResponse) - jsonResponse, err := json.Marshal(fullTextResponse) + jsonResponse, err := common.Marshal(fullTextResponse) if err != nil { return nil, types.NewError(err, types.ErrorCodeBadResponseBody) } diff --git a/relay/channel/zhipu/relay_zhipu_test.go b/relay/channel/zhipu/relay_zhipu_test.go new file mode 100644 index 00000000..9247c3cc --- /dev/null +++ b/relay/channel/zhipu/relay_zhipu_test.go @@ -0,0 +1,26 @@ +package zhipu + +import ( + "testing" + + "github.com/stretchr/testify/require" +) + +func TestGetZhipuTokenIgnoresUnexpectedCachedValue(t *testing.T) { + const apiKey = "test-id.test-secret" + zhipuTokens.Store(apiKey, 123) + t.Cleanup(func() { zhipuTokens.Delete(apiKey) }) + + var token string + require.NotPanics(t, func() { + var err error + token, err = getZhipuToken(apiKey) + require.NoError(t, err) + }) + require.NotEmpty(t, token) +} + +func TestGetZhipuTokenRejectsMalformedAPIKey(t *testing.T) { + _, err := getZhipuToken("not-a-zhipu-key") + require.EqualError(t, err, "invalid zhipu API key format") +} diff --git a/relay/channel/zhipu/stream_cancel_test.go b/relay/channel/zhipu/stream_cancel_test.go index a264462b..946c583f 100644 --- a/relay/channel/zhipu/stream_cancel_test.go +++ b/relay/channel/zhipu/stream_cancel_test.go @@ -2,15 +2,18 @@ package zhipu import ( "context" + "errors" "io" "net/http" "net/http/httptest" + "strings" "testing" "time" "github.com/MAX-API-Next/MAX-API/common" "github.com/MAX-API-Next/MAX-API/dto" relaycommon "github.com/MAX-API-Next/MAX-API/relay/common" + "github.com/MAX-API-Next/MAX-API/types" "github.com/gin-gonic/gin" "github.com/stretchr/testify/require" ) @@ -24,6 +27,52 @@ func (r *closeNotifyRecorder) CloseNotify() <-chan bool { return r.closeNotify } +type panicReadCloser struct { + closed bool +} + +func (r *panicReadCloser) Read([]byte) (int, error) { + panic("upstream stream read panic") +} + +func (r *panicReadCloser) Close() error { + r.closed = true + return nil +} + +type errorReadCloser struct { + closed bool +} + +func (r *errorReadCloser) Read([]byte) (int, error) { + return 0, errors.New("upstream stream read failed") +} + +func (r *errorReadCloser) Close() error { + r.closed = true + return nil +} + +func assertZhipuStreamReadFailure(t *testing.T, body io.ReadCloser, isClosed func() bool, message string) { + t.Helper() + gin.SetMode(gin.TestMode) + recorder := &closeNotifyRecorder{ + ResponseRecorder: httptest.NewRecorder(), + closeNotify: make(chan bool), + } + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodGet, "/", nil) + + usage, maxErr := zhipuStreamHandler(c, &relaycommon.RelayInfo{}, &http.Response{Body: body}) + + require.Nil(t, usage) + require.NotNil(t, maxErr) + require.Equal(t, types.ErrorCodeReadResponseBodyFailed, maxErr.GetErrorCode()) + require.Contains(t, maxErr.Error(), message) + require.NotContains(t, recorder.Body.String(), "[DONE]") + require.True(t, isClosed()) +} + func TestZhipuStreamHandlerStopsAfterRequestCancellation(t *testing.T) { gin.SetMode(gin.TestMode) c, _ := gin.CreateTestContext(&closeNotifyRecorder{ @@ -62,6 +111,34 @@ func TestZhipuStreamHandlerStopsAfterRequestCancellation(t *testing.T) { } } +func TestZhipuStreamHandlerReturnsProducerPanic(t *testing.T) { + body := &panicReadCloser{} + assertZhipuStreamReadFailure(t, body, func() bool { return body.closed }, "upstream stream read panic") +} + +func TestZhipuStreamHandlerReturnsScannerError(t *testing.T) { + body := &errorReadCloser{} + assertZhipuStreamReadFailure(t, body, func() bool { return body.closed }, "upstream stream read failed") +} + +func TestZhipuStreamHandlerEmitsDoneOnNormalEOF(t *testing.T) { + gin.SetMode(gin.TestMode) + recorder := &closeNotifyRecorder{ + ResponseRecorder: httptest.NewRecorder(), + closeNotify: make(chan bool), + } + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodGet, "/", nil) + + usage, maxErr := zhipuStreamHandler(c, &relaycommon.RelayInfo{}, &http.Response{ + Body: io.NopCloser(strings.NewReader("")), + }) + + require.Nil(t, usage) + require.Nil(t, maxErr) + require.Contains(t, recorder.Body.String(), "[DONE]") +} + func TestRequestOpenAI2ZhipuPreservesExplicitTopPZero(t *testing.T) { zero := 0.0 request := requestOpenAI2Zhipu(dto.GeneralOpenAIRequest{TopP: &zero}) diff --git a/relay/claude_handler.go b/relay/claude_handler.go index 3dea41b4..6bec0e69 100644 --- a/relay/claude_handler.go +++ b/relay/claude_handler.go @@ -155,7 +155,15 @@ func ClaudeHelper(c *gin.Context, info *relaycommon.RelayInfo) (maxAPIError *typ if err != nil { return types.NewErrorWithStatusCode(err, types.ErrorCodeReadRequestBodyFailed, http.StatusBadRequest, types.ErrOptionWithSkipRetry()) } - requestBody = common.ReaderOnly(storage) + body, size, closer, err := relaycommon.PreparePassThroughJSONBody(storage, info.ChannelOtherSettings) + if err != nil { + return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry()) + } + if closer != nil { + defer closer.Close() + } + info.UpstreamRequestBodySize = size + requestBody = body } else { convertedRequest, err := adaptor.ConvertClaudeRequest(c, info, request) if err != nil { diff --git a/relay/common/outbound_body.go b/relay/common/outbound_body.go index 80280264..ed48b6bd 100644 --- a/relay/common/outbound_body.go +++ b/relay/common/outbound_body.go @@ -4,6 +4,7 @@ import ( "io" "github.com/MAX-API-Next/MAX-API/common" + "github.com/MAX-API-Next/MAX-API/dto" ) // NewOutboundJSONBody wraps the already-marshaled upstream request body into a @@ -18,14 +19,25 @@ import ( // The caller MUST invoke closer.Close() once the upstream call has finished // (typically via defer) to release the disk file / memory accounting. // -// The returned reader is wrapped with common.ReaderOnly to prevent the HTTP -// transport from prematurely closing the underlying BodyStorage. The returned -// size is meant to be propagated to http.Request.ContentLength because the -// type-erased io.Reader prevents net/http from auto-detecting it. +// The returned replayable reader hides storage ownership from net/http while +// exposing independent readers for transport-level retries. func NewOutboundJSONBody(data []byte) (body io.Reader, size int64, closer io.Closer, err error) { storage, err := common.CreateBodyStorage(data) if err != nil { return nil, 0, nil, err } - return common.ReaderOnly(storage), storage.Size(), storage, nil + return common.NewReplayableBodyReader(storage), storage.Size(), storage, nil +} + +// PreparePassThroughJSONBody preserves raw pass-through semantics while still +// enforcing channel permissions for fields that can change cost or privacy. +func PreparePassThroughJSONBody(storage common.BodyStorage, settings dto.ChannelOtherSettings) (body io.Reader, size int64, closer io.Closer, err error) { + filtered, err := newPassThroughFilteredBody(storage, settings) + if err != nil { + return nil, 0, nil, err + } + if filtered == nil { + return common.NewReplayableBodyReader(storage), storage.Size(), nil, nil + } + return filtered, filtered.Size(), filtered, nil } diff --git a/relay/common/outbound_body_test.go b/relay/common/outbound_body_test.go new file mode 100644 index 00000000..6a763914 --- /dev/null +++ b/relay/common/outbound_body_test.go @@ -0,0 +1,189 @@ +package common + +import ( + "bytes" + "errors" + "io" + "strings" + "testing" + + basecommon "github.com/MAX-API-Next/MAX-API/common" + "github.com/MAX-API-Next/MAX-API/dto" + "github.com/stretchr/testify/require" +) + +type noBytesBodyStorage struct { + payload []byte + reader *bytes.Reader + newReaderCalls int +} + +func newNoBytesBodyStorage(payload []byte) *noBytesBodyStorage { + return &noBytesBodyStorage{payload: payload, reader: bytes.NewReader(payload)} +} + +func (s *noBytesBodyStorage) Read(p []byte) (int, error) { return s.reader.Read(p) } +func (s *noBytesBodyStorage) Seek(offset int64, whence int) (int64, error) { + return s.reader.Seek(offset, whence) +} +func (s *noBytesBodyStorage) Close() error { return nil } +func (s *noBytesBodyStorage) Bytes() ([]byte, error) { + return nil, errors.New("full body read is forbidden") +} +func (s *noBytesBodyStorage) Size() int64 { return int64(len(s.payload)) } +func (s *noBytesBodyStorage) IsDisk() bool { return true } +func (s *noBytesBodyStorage) NewReader() (io.ReadCloser, error) { + s.newReaderCalls++ + return io.NopCloser(bytes.NewReader(s.payload)), nil +} + +func TestPreparePassThroughJSONBodyFiltersCostFields(t *testing.T) { + storage, err := basecommon.CreateBodyStorage([]byte(`{"model":"gpt-5","service_tier":"flex","speed":"fast"}`)) + require.NoError(t, err) + defer storage.Close() + + body, size, closer, err := PreparePassThroughJSONBody(storage, dto.ChannelOtherSettings{}) + require.NoError(t, err) + if closer != nil { + defer closer.Close() + } + data, err := io.ReadAll(body) + require.NoError(t, err) + require.JSONEq(t, `{"model":"gpt-5"}`, string(data)) + require.Equal(t, int64(len(data)), size) +} + +func TestPreparePassThroughJSONBodyReusesStorageWhenUnchanged(t *testing.T) { + payload := []byte(`{"model":"gpt-5","service_tier":"flex"}`) + storage, err := basecommon.CreateBodyStorage(payload) + require.NoError(t, err) + defer storage.Close() + + body, size, closer, err := PreparePassThroughJSONBody(storage, dto.ChannelOtherSettings{AllowServiceTier: true}) + require.NoError(t, err) + require.Nil(t, closer) + require.Equal(t, int64(len(payload)), size) + _, ok := body.(basecommon.ReplayableBody) + require.True(t, ok) +} + +func TestPreparePassThroughJSONBodyDoesNotMaterializeUnchangedStorage(t *testing.T) { + payload := []byte(`{"model":"gpt-5","input":"` + strings.Repeat("x", 1<<20) + `"}`) + storage := newNoBytesBodyStorage(payload) + + body, size, closer, err := PreparePassThroughJSONBody(storage, dto.ChannelOtherSettings{}) + require.NoError(t, err) + require.Nil(t, closer) + require.Equal(t, int64(len(payload)), size) + + data, err := io.ReadAll(body) + require.NoError(t, err) + require.Equal(t, payload, data) +} + +func TestPreparePassThroughJSONBodySkipsScanWhenAllControlledFieldsAllowed(t *testing.T) { + payload := []byte(`{"model":"gpt-5","service_tier":"flex"}`) + storage := newNoBytesBodyStorage(payload) + settings := dto.ChannelOtherSettings{ + AllowServiceTier: true, + AllowInferenceGeo: true, + AllowSpeed: true, + AllowSafetyIdentifier: true, + AllowIncludeObfuscation: true, + } + + body, size, closer, err := PreparePassThroughJSONBody(storage, settings) + require.NoError(t, err) + require.Nil(t, closer) + require.Equal(t, int64(len(payload)), size) + require.Zero(t, storage.newReaderCalls) + + data, err := io.ReadAll(body) + require.NoError(t, err) + require.Equal(t, payload, data) +} + +func TestPreparePassThroughJSONBodyFiltersWithoutMaterializingStorage(t *testing.T) { + payload := []byte(`{"service_tier":"flex","model":"gpt-5","input":"` + strings.Repeat("x", 1<<20) + `"}`) + storage := newNoBytesBodyStorage(payload) + + body, size, closer, err := PreparePassThroughJSONBody(storage, dto.ChannelOtherSettings{}) + require.NoError(t, err) + require.NotNil(t, closer) + defer closer.Close() + + data, err := io.ReadAll(body) + require.NoError(t, err) + require.NotContains(t, string(data), "service_tier") + require.JSONEq(t, `{"model":"gpt-5","input":"`+strings.Repeat("x", 1<<20)+`"}`, string(data)) + require.Equal(t, int64(len(data)), size) +} + +func TestPreparePassThroughJSONBodyPreservesValidStructureAcrossFieldPositions(t *testing.T) { + tests := []struct { + name string + payload string + expected string + }{ + {name: "first", payload: `{"service_tier":"flex","model":"gpt-5"}`, expected: `{"model":"gpt-5"}`}, + {name: "middle", payload: `{"model":"gpt-5","speed":"fast","input":"hi"}`, expected: `{"model":"gpt-5","input":"hi"}`}, + {name: "last", payload: `{"model":"gpt-5","speed":"fast"}`, expected: `{"model":"gpt-5"}`}, + {name: "all", payload: `{"service_tier":"flex","speed":"fast"}`, expected: `{}`}, + {name: "escaped key", payload: `{"service\u005ftier":"flex","model":"gpt-5"}`, expected: `{"model":"gpt-5"}`}, + {name: "nested root field is preserved", payload: `{"metadata":{"service_tier":"flex"},"model":"gpt-5"}`, expected: `{"metadata":{"service_tier":"flex"},"model":"gpt-5"}`}, + {name: "empty stream options removed", payload: `{"stream_options":{"include_obfuscation":false},"model":"gpt-5"}`, expected: `{"model":"gpt-5"}`}, + {name: "other stream options preserved", payload: `{"stream_options":{"include_usage":true,"include_obfuscation":false},"model":"gpt-5"}`, expected: `{"stream_options":{"include_usage":true},"model":"gpt-5"}`}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + storage := newNoBytesBodyStorage([]byte(tt.payload)) + body, size, closer, err := PreparePassThroughJSONBody(storage, dto.ChannelOtherSettings{}) + require.NoError(t, err) + if closer != nil { + defer closer.Close() + } + data, err := io.ReadAll(body) + require.NoError(t, err) + require.JSONEq(t, tt.expected, string(data)) + require.Equal(t, int64(len(data)), size) + }) + } +} + +func TestPreparePassThroughJSONBodyProvidesIndependentFilteredReplayReaders(t *testing.T) { + storage := newNoBytesBodyStorage([]byte(`{"model":"gpt-5","service_tier":"flex","input":"abcdef"}`)) + body, _, closer, err := PreparePassThroughJSONBody(storage, dto.ChannelOtherSettings{}) + require.NoError(t, err) + require.NotNil(t, closer) + + replayable, ok := body.(basecommon.ReplayableBody) + require.True(t, ok) + first, err := replayable.NewReader() + require.NoError(t, err) + defer first.Close() + second, err := replayable.NewReader() + require.NoError(t, err) + defer second.Close() + + prefix := make([]byte, 8) + _, err = io.ReadFull(first, prefix) + require.NoError(t, err) + secondData, err := io.ReadAll(second) + require.NoError(t, err) + firstRest, err := io.ReadAll(first) + require.NoError(t, err) + + expected := []byte(`{"model":"gpt-5","input":"abcdef"}`) + require.Equal(t, expected, secondData) + require.Equal(t, expected, append(prefix, firstRest...)) + require.NoError(t, closer.Close()) + _, err = body.Read(make([]byte, 1)) + require.ErrorIs(t, err, io.ErrClosedPipe) + third, err := replayable.NewReader() + require.NoError(t, err) + defer third.Close() + thirdData, err := io.ReadAll(third) + require.NoError(t, err) + require.Equal(t, expected, thirdData) +} diff --git a/relay/common/override_test.go b/relay/common/override_test.go index 1ae8fd01..c42121b8 100644 --- a/relay/common/override_test.go +++ b/relay/common/override_test.go @@ -1998,7 +1998,7 @@ func TestGetEffectiveHeaderOverrideUsesRuntimeOverrideAsFinalResult(t *testing.T } } -func TestRemoveDisabledFieldsSkipWhenChannelPassThroughEnabled(t *testing.T) { +func TestRemoveDisabledFieldsFiltersWhenChannelPassThroughEnabled(t *testing.T) { input := `{ "service_tier":"flex", "safety_identifier":"user-123", @@ -2011,10 +2011,10 @@ func TestRemoveDisabledFieldsSkipWhenChannelPassThroughEnabled(t *testing.T) { if err != nil { t.Fatalf("RemoveDisabledFields returned error: %v", err) } - assertJSONEqual(t, input, string(out)) + assertJSONEqual(t, `{"store":true}`, string(out)) } -func TestRemoveDisabledFieldsSkipWhenGlobalPassThroughEnabled(t *testing.T) { +func TestRemoveDisabledFieldsFiltersWhenGlobalPassThroughEnabled(t *testing.T) { original := model_setting.GetGlobalSettings().PassThroughRequestEnabled model_setting.GetGlobalSettings().PassThroughRequestEnabled = true t.Cleanup(func() { @@ -2032,7 +2032,21 @@ func TestRemoveDisabledFieldsSkipWhenGlobalPassThroughEnabled(t *testing.T) { if err != nil { t.Fatalf("RemoveDisabledFields returned error: %v", err) } - assertJSONEqual(t, input, string(out)) + assertJSONEqual(t, `{}`, string(out)) +} + +func TestRemoveDisabledFieldsPassThroughHonorsExplicitCostFieldOptIn(t *testing.T) { + input := `{"service_tier":"flex","speed":"fast","inference_geo":"eu"}` + settings := dto.ChannelOtherSettings{ + AllowServiceTier: true, + AllowSpeed: true, + } + + out, err := RemoveDisabledFields([]byte(input), settings, true) + if err != nil { + t.Fatalf("RemoveDisabledFields returned error: %v", err) + } + assertJSONEqual(t, `{"service_tier":"flex","speed":"fast"}`, string(out)) } func TestRemoveDisabledFieldsDefaultFiltering(t *testing.T) { diff --git a/relay/common/pass_through_body.go b/relay/common/pass_through_body.go new file mode 100644 index 00000000..941ebc46 --- /dev/null +++ b/relay/common/pass_through_body.go @@ -0,0 +1,567 @@ +package common + +import ( + "bufio" + "errors" + "fmt" + "io" + "sort" + "sync" + + basecommon "github.com/MAX-API-Next/MAX-API/common" + "github.com/MAX-API-Next/MAX-API/dto" +) + +const maxCapturedJSONKeyBytes = 256 + +type jsonFilterScope uint8 + +const ( + jsonFilterNone jsonFilterScope = iota + jsonFilterRoot + jsonFilterStreamOptions +) + +type jsonByteRange struct { + start int64 + end int64 +} + +type passThroughJSONScanner struct { + reader *bufio.Reader + offset int64 + peeked byte + hasPeeked bool +} + +func newPassThroughFilteredBody(storage basecommon.BodyStorage, settings dto.ChannelOtherSettings) (*filteredReplayableBody, error) { + if settings.AllowServiceTier && + settings.AllowInferenceGeo && + settings.AllowSpeed && + !settings.DisableStore && + settings.AllowSafetyIdentifier && + settings.AllowIncludeObfuscation { + return nil, nil + } + + reader, err := storage.NewReader() + if err != nil { + return nil, err + } + ranges, scannedSize, scanErr := scanPassThroughJSON(reader, settings) + closeErr := reader.Close() + if scanErr != nil { + return nil, scanErr + } + if closeErr != nil { + return nil, closeErr + } + if scannedSize != storage.Size() { + return nil, fmt.Errorf("pass-through body size changed while scanning: got %d, want %d", scannedSize, storage.Size()) + } + if len(ranges) == 0 { + return nil, nil + } + + ranges, filteredSize, err := normalizeJSONRanges(ranges, storage.Size()) + if err != nil { + return nil, err + } + return &filteredReplayableBody{storage: storage, ranges: ranges, size: filteredSize}, nil +} + +func scanPassThroughJSON(reader io.Reader, settings dto.ChannelOtherSettings) ([]jsonByteRange, int64, error) { + scanner := &passThroughJSONScanner{reader: bufio.NewReader(reader)} + if err := scanner.skipWhitespace(); err != nil && !errors.Is(err, io.EOF) { + return nil, scanner.offset, err + } + first, err := scanner.peekByte() + if err != nil { + return nil, scanner.offset, fmt.Errorf("scan pass-through JSON: %w", err) + } + if first != '{' { + if err := scanner.parseValue(); err != nil { + return nil, scanner.offset, err + } + if err := scanner.requireEOF(); err != nil { + return nil, scanner.offset, err + } + return nil, scanner.offset, nil + } + + ranges, _, _, err := scanner.parseObject(jsonFilterRoot, settings) + if err != nil { + return nil, scanner.offset, err + } + if err := scanner.requireEOF(); err != nil { + return nil, scanner.offset, err + } + return ranges, scanner.offset, nil +} + +func (s *passThroughJSONScanner) parseObject(scope jsonFilterScope, settings dto.ChannelOtherSettings) ([]jsonByteRange, int, bool, error) { + if err := s.expectByte('{'); err != nil { + return nil, 0, false, err + } + + var ranges []jsonByteRange + keptMembers := 0 + changed := false + lastKeptComma := int64(-1) + + for { + memberStart := s.offset + if err := s.skipWhitespace(); err != nil { + return nil, 0, false, err + } + next, err := s.peekByte() + if err != nil { + return nil, 0, false, fmt.Errorf("scan JSON object: %w", err) + } + if next == '}' { + _, _ = s.readByte() + return ranges, keptMembers, changed, nil + } + + key, err := s.parseString(scope != jsonFilterNone) + if err != nil { + return nil, 0, false, err + } + if err := s.skipWhitespace(); err != nil { + return nil, 0, false, err + } + if err := s.expectByte(':'); err != nil { + return nil, 0, false, err + } + if err := s.skipWhitespace(); err != nil { + return nil, 0, false, err + } + + removeMember := shouldRemoveJSONMember(scope, key, settings) + var childRanges []jsonByteRange + childChanged := false + if scope == jsonFilterRoot && key == "stream_options" && !settings.AllowIncludeObfuscation { + valueStart, peekErr := s.peekByte() + if peekErr != nil { + return nil, 0, false, peekErr + } + if valueStart == '{' { + var childKept int + childRanges, childKept, childChanged, err = s.parseObject(jsonFilterStreamOptions, settings) + if err != nil { + return nil, 0, false, err + } + if childChanged && childKept == 0 { + removeMember = true + childRanges = nil + } + } else if err := s.parseValue(); err != nil { + return nil, 0, false, err + } + } else if err := s.parseValue(); err != nil { + return nil, 0, false, err + } + + valueEnd := s.offset + if err := s.skipWhitespace(); err != nil { + return nil, 0, false, err + } + commaStart := s.offset + delimiter, err := s.readByte() + if err != nil { + return nil, 0, false, fmt.Errorf("scan JSON object delimiter: %w", err) + } + if delimiter != ',' && delimiter != '}' { + return nil, 0, false, fmt.Errorf("invalid JSON object delimiter %q at byte %d", delimiter, s.offset-1) + } + + if removeMember { + changed = true + if delimiter == ',' { + ranges = append(ranges, jsonByteRange{start: memberStart, end: s.offset}) + } else if lastKeptComma >= 0 { + ranges = append(ranges, jsonByteRange{start: lastKeptComma, end: valueEnd}) + } else { + ranges = append(ranges, jsonByteRange{start: memberStart, end: valueEnd}) + } + } else { + keptMembers++ + if childChanged { + changed = true + ranges = append(ranges, childRanges...) + } + if delimiter == ',' { + lastKeptComma = commaStart + } + } + + if delimiter == '}' { + return ranges, keptMembers, changed, nil + } + } +} + +func shouldRemoveJSONMember(scope jsonFilterScope, key string, settings dto.ChannelOtherSettings) bool { + switch scope { + case jsonFilterRoot: + switch key { + case "service_tier": + return !settings.AllowServiceTier + case "inference_geo": + return !settings.AllowInferenceGeo + case "speed": + return !settings.AllowSpeed + case "store": + return settings.DisableStore + case "safety_identifier": + return !settings.AllowSafetyIdentifier + } + case jsonFilterStreamOptions: + return key == "include_obfuscation" && !settings.AllowIncludeObfuscation + } + return false +} + +func (s *passThroughJSONScanner) parseValue() error { + if err := s.skipWhitespace(); err != nil { + return err + } + next, err := s.peekByte() + if err != nil { + return fmt.Errorf("scan JSON value: %w", err) + } + switch next { + case '{': + _, _, _, err = s.parseObject(jsonFilterNone, dto.ChannelOtherSettings{}) + return err + case '[': + return s.parseArray() + case '"': + _, err = s.parseString(false) + return err + case 't': + return s.consumeLiteral("true") + case 'f': + return s.consumeLiteral("false") + case 'n': + return s.consumeLiteral("null") + default: + if next == '-' || (next >= '0' && next <= '9') { + return s.consumeNumber() + } + return fmt.Errorf("invalid JSON value at byte %d", s.offset) + } +} + +func (s *passThroughJSONScanner) parseArray() error { + if err := s.expectByte('['); err != nil { + return err + } + if err := s.skipWhitespace(); err != nil { + return err + } + next, err := s.peekByte() + if err != nil { + return err + } + if next == ']' { + _, _ = s.readByte() + return nil + } + for { + if err := s.parseValue(); err != nil { + return err + } + if err := s.skipWhitespace(); err != nil { + return err + } + delimiter, err := s.readByte() + if err != nil { + return err + } + switch delimiter { + case ']': + return nil + case ',': + continue + default: + return fmt.Errorf("invalid JSON array delimiter %q at byte %d", delimiter, s.offset-1) + } + } +} + +func (s *passThroughJSONScanner) parseString(capture bool) (string, error) { + if err := s.expectByte('"'); err != nil { + return "", err + } + raw := make([]byte, 0, 32) + if capture { + raw = append(raw, '"') + } + overflow := false + escaped := false + for { + value, err := s.readByte() + if err != nil { + return "", fmt.Errorf("scan JSON string: %w", err) + } + if capture && !overflow { + if len(raw) >= maxCapturedJSONKeyBytes { + overflow = true + raw = nil + } else { + raw = append(raw, value) + } + } + if escaped { + escaped = false + continue + } + if value == '\\' { + escaped = true + continue + } + if value == '"' { + break + } + if value < 0x20 { + return "", fmt.Errorf("invalid control character in JSON string at byte %d", s.offset-1) + } + } + if !capture || overflow { + return "", nil + } + var decoded string + if err := basecommon.Unmarshal(raw, &decoded); err != nil { + return "", fmt.Errorf("decode JSON object key: %w", err) + } + return decoded, nil +} + +func (s *passThroughJSONScanner) consumeLiteral(literal string) error { + for i := 0; i < len(literal); i++ { + value, err := s.readByte() + if err != nil { + return err + } + if value != literal[i] { + return fmt.Errorf("invalid JSON literal at byte %d", s.offset-1) + } + } + return nil +} + +func (s *passThroughJSONScanner) consumeNumber() error { + consumed := false + for { + value, err := s.peekByte() + if errors.Is(err, io.EOF) { + if consumed { + return nil + } + return err + } + if err != nil { + return err + } + if (value >= '0' && value <= '9') || value == '-' || value == '+' || value == '.' || value == 'e' || value == 'E' { + _, _ = s.readByte() + consumed = true + continue + } + if isJSONWhitespace(value) || value == ',' || value == ']' || value == '}' { + return nil + } + return fmt.Errorf("invalid JSON number at byte %d", s.offset) + } +} + +func (s *passThroughJSONScanner) requireEOF() error { + if err := s.skipWhitespace(); err != nil && !errors.Is(err, io.EOF) { + return err + } + if _, err := s.peekByte(); !errors.Is(err, io.EOF) { + if err != nil { + return err + } + return fmt.Errorf("unexpected data after JSON value at byte %d", s.offset) + } + return nil +} + +func (s *passThroughJSONScanner) skipWhitespace() error { + for { + value, err := s.peekByte() + if err != nil { + return err + } + if !isJSONWhitespace(value) { + return nil + } + _, _ = s.readByte() + } +} + +func isJSONWhitespace(value byte) bool { + return value == ' ' || value == '\t' || value == '\r' || value == '\n' +} + +func (s *passThroughJSONScanner) expectByte(expected byte) error { + value, err := s.readByte() + if err != nil { + return err + } + if value != expected { + return fmt.Errorf("expected %q at byte %d, got %q", expected, s.offset-1, value) + } + return nil +} + +func (s *passThroughJSONScanner) peekByte() (byte, error) { + if s.hasPeeked { + return s.peeked, nil + } + value, err := s.reader.ReadByte() + if err != nil { + return 0, err + } + s.peeked = value + s.hasPeeked = true + return value, nil +} + +func (s *passThroughJSONScanner) readByte() (byte, error) { + if s.hasPeeked { + s.hasPeeked = false + s.offset++ + return s.peeked, nil + } + value, err := s.reader.ReadByte() + if err == nil { + s.offset++ + } + return value, err +} + +func normalizeJSONRanges(ranges []jsonByteRange, sourceSize int64) ([]jsonByteRange, int64, error) { + sort.Slice(ranges, func(i, j int) bool { + return ranges[i].start < ranges[j].start + }) + merged := make([]jsonByteRange, 0, len(ranges)) + for _, current := range ranges { + if current.start < 0 || current.end <= current.start || current.end > sourceSize { + return nil, 0, fmt.Errorf("invalid JSON filter range [%d,%d) for body size %d", current.start, current.end, sourceSize) + } + if len(merged) > 0 && current.start <= merged[len(merged)-1].end { + if current.end > merged[len(merged)-1].end { + merged[len(merged)-1].end = current.end + } + continue + } + merged = append(merged, current) + } + filteredSize := sourceSize + for _, current := range merged { + filteredSize -= current.end - current.start + } + return merged, filteredSize, nil +} + +type filteredReplayableBody struct { + storage basecommon.BodyStorage + ranges []jsonByteRange + size int64 + + mu sync.Mutex + current io.ReadCloser + closed bool +} + +func (b *filteredReplayableBody) Size() int64 { + return b.size +} + +func (b *filteredReplayableBody) Read(p []byte) (int, error) { + b.mu.Lock() + defer b.mu.Unlock() + if b.closed { + return 0, io.ErrClosedPipe + } + if b.current == nil { + reader, err := b.newReader() + if err != nil { + return 0, err + } + b.current = reader + } + return b.current.Read(p) +} + +func (b *filteredReplayableBody) Close() error { + b.mu.Lock() + defer b.mu.Unlock() + b.closed = true + if b.current == nil { + return nil + } + err := b.current.Close() + b.current = nil + return err +} + +func (b *filteredReplayableBody) NewReader() (io.ReadCloser, error) { + return b.newReader() +} + +func (b *filteredReplayableBody) newReader() (io.ReadCloser, error) { + source, err := b.storage.NewReader() + if err != nil { + return nil, err + } + return &rangeFilteringReader{source: source, ranges: b.ranges}, nil +} + +type rangeFilteringReader struct { + source io.ReadCloser + ranges []jsonByteRange + position int64 + index int +} + +func (r *rangeFilteringReader) Read(p []byte) (int, error) { + if len(p) == 0 { + return 0, nil + } + for { + for r.index < len(r.ranges) && r.position >= r.ranges[r.index].end { + r.index++ + } + if r.index < len(r.ranges) && r.position >= r.ranges[r.index].start { + toSkip := r.ranges[r.index].end - r.position + skipped, err := io.CopyN(io.Discard, r.source, toSkip) + r.position += skipped + if err != nil { + return 0, fmt.Errorf("skip filtered JSON field: %w", err) + } + continue + } + + limit := len(p) + if r.index < len(r.ranges) { + untilRange := r.ranges[r.index].start - r.position + if untilRange < int64(limit) { + limit = int(untilRange) + } + } + if limit == 0 { + continue + } + n, err := r.source.Read(p[:limit]) + r.position += int64(n) + return n, err + } +} + +func (r *rangeFilteringReader) Close() error { + return r.source.Close() +} diff --git a/relay/common/relay_info.go b/relay/common/relay_info.go index aa14a5c6..c4589af2 100644 --- a/relay/common/relay_info.go +++ b/relay/common/relay_info.go @@ -865,10 +865,7 @@ func FailTaskInfo(reason string) *TaskInfo { // store: 数据存储授权字段，涉及用户隐私（仅 OpenAI、Responses API 支持，默认允许透传，禁用后可能导致 Codex 无法使用） // safety_identifier: 安全标识符，用于向 OpenAI 报告违规用户（仅 OpenAI 支持，涉及用户隐私） // stream_options.include_obfuscation: 响应流混淆控制字段（仅 OpenAI Responses API 支持） -func RemoveDisabledFields(jsonData []byte, channelOtherSettings dto.ChannelOtherSettings, channelPassThroughEnabled bool) ([]byte, error) { - if model_setting.GetGlobalSettings().PassThroughRequestEnabled || channelPassThroughEnabled { - return jsonData, nil - } +func RemoveDisabledFields(jsonData []byte, channelOtherSettings dto.ChannelOtherSettings, _ bool) ([]byte, error) { if !hasRemovableDisabledField(jsonData, channelOtherSettings) { return jsonData, nil } diff --git a/relay/compatible_handler.go b/relay/compatible_handler.go index dcf43d1a..f0504ac2 100644 --- a/relay/compatible_handler.go +++ b/relay/compatible_handler.go @@ -104,7 +104,15 @@ func TextHelper(c *gin.Context, info *relaycommon.RelayInfo) (maxAPIError *types logger.LogDebug(c, "requestBody: %s", debugBytes) } } - requestBody = common.ReaderOnly(storage) + body, size, closer, err := relaycommon.PreparePassThroughJSONBody(storage, info.ChannelOtherSettings) + if err != nil { + return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry()) + } + if closer != nil { + defer closer.Close() + } + info.UpstreamRequestBodySize = size + requestBody = body } else { convertedRequest, err := adaptor.ConvertOpenAIRequest(c, info, request) if err != nil { diff --git a/relay/helper/price.go b/relay/helper/price.go index 73a35eb8..659fa3d7 100644 --- a/relay/helper/price.go +++ b/relay/helper/price.go @@ -19,7 +19,7 @@ import ( ) func modelPriceNotConfiguredError(modelName string, userId int) error { - if model.IsAdmin(userId) { + if userId > 0 && model.IsAdmin(userId) { return fmt.Errorf( "模型 %s 的价格未配置。请前往「系统设置 → 运营设置」开启自用模式，或在「系统设置 → 分组与模型定价设置」中为该模型配置价格；"+ "Model %s price not configured. Go to System Settings → Operation Settings to enable self-use mode, or configure the model price in System Settings → Group & Model Pricing.", @@ -33,6 +33,11 @@ func modelPriceNotConfiguredError(modelName string, userId int) error { ) } +func taskModelRequiresConfiguredPrice(modelName string) bool { + modelName = strings.ToLower(strings.TrimSpace(modelName)) + return modelName == "swap_face" || strings.HasPrefix(modelName, "mj_") || strings.HasPrefix(modelName, "suno_") +} + // https://docs.claude.com/en/docs/build-with-claude/prompt-caching#1-hour-cache-duration const claudeCacheCreation1hMultiplier = 6 / 3.75 @@ -174,12 +179,19 @@ func ModelPriceHelper(c *gin.Context, info *relaycommon.RelayInfo, promptTokens // ModelPriceHelperPerCall 按次/按量计费的 PriceHelper (MJ、Task) func ModelPriceHelperPerCall(c *gin.Context, info *relaycommon.RelayInfo) (types.PriceData, error) { + if info == nil { + return types.PriceData{}, fmt.Errorf("relay info is nil") + } groupRatioInfo := HandleGroupRatio(c, info) modelPrice, success := ratio_setting.GetModelPrice(info.OriginModelName, true) usePrice := success var modelRatio float64 - rateCardPriced := task_billing_setting.HasRateCard(info.OriginModelName, info.UpstreamModelName) + upstreamModelName := "" + if info.ChannelMeta != nil { + upstreamModelName = info.UpstreamModelName + } + rateCardPriced := task_billing_setting.HasRateCard(info.OriginModelName, upstreamModelName) if !success { defaultPrice, ok := ratio_setting.GetDefaultModelPriceMap()[info.OriginModelName] @@ -189,6 +201,8 @@ func ModelPriceHelperPerCall(c *gin.Context, info *relaycommon.RelayInfo) (types } else if rateCardPriced { modelPrice = 0 usePrice = true + } else if taskModelRequiresConfiguredPrice(info.OriginModelName) { + return types.PriceData{}, modelPriceNotConfiguredError(info.OriginModelName, info.UserId) } else { var ratioSuccess bool var matchName string diff --git a/relay/helper/price_test.go b/relay/helper/price_test.go index bcc8b1d3..b79ff335 100644 --- a/relay/helper/price_test.go +++ b/relay/helper/price_test.go @@ -6,10 +6,13 @@ import ( "testing" "github.com/MAX-API-Next/MAX-API/common" + "github.com/MAX-API-Next/MAX-API/dto" "github.com/MAX-API-Next/MAX-API/pkg/billingexpr" relaycommon "github.com/MAX-API-Next/MAX-API/relay/common" "github.com/MAX-API-Next/MAX-API/setting/billing_setting" "github.com/MAX-API-Next/MAX-API/setting/config" + "github.com/MAX-API-Next/MAX-API/setting/operation_setting" + "github.com/MAX-API-Next/MAX-API/setting/ratio_setting" "github.com/MAX-API-Next/MAX-API/types" "github.com/gin-gonic/gin" "github.com/stretchr/testify/require" @@ -60,3 +63,64 @@ func TestModelPriceHelperTieredUsesPreloadedRequestInput(t *testing.T) { require.Equal(t, billing_setting.BillingModeTieredExpr, info.TieredBillingSnapshot.BillingMode) require.Equal(t, common.QuotaPerUnit, info.TieredBillingSnapshot.QuotaPerUnit) } + +func TestModelPriceHelperPerCallRejectsUnpricedMJSunoTaskModels(t *testing.T) { + gin.SetMode(gin.TestMode) + originalSelfUseMode := operation_setting.SelfUseModeEnabled + operation_setting.SelfUseModeEnabled = true + t.Cleanup(func() { + operation_setting.SelfUseModeEnabled = originalSelfUseMode + }) + + for _, modelName := range []string{"mj_unmapped_action", "suno_unmapped_action"} { + t.Run(modelName, func(t *testing.T) { + ctx, _ := gin.CreateTestContext(httptest.NewRecorder()) + info := &relaycommon.RelayInfo{ + OriginModelName: modelName, + UserGroup: "default", + UsingGroup: "default", + ChannelMeta: &relaycommon.ChannelMeta{}, + UserSetting: dto.UserSetting{ + AcceptUnsetRatioModel: true, + }, + } + + _, err := ModelPriceHelperPerCall(ctx, info) + + require.Error(t, err) + require.Contains(t, err.Error(), "not been priced") + }) + } +} + +func TestModelPriceHelperPerCallUsesDefaultTaskPrice(t *testing.T) { + gin.SetMode(gin.TestMode) + originalQuotaPerUnit := common.QuotaPerUnit + originalGroupRatio := ratio_setting.GroupRatio2JSONString() + quotaSetting := operation_setting.GetQuotaSetting() + originalEnableFreeModelPreConsume := quotaSetting.EnableFreeModelPreConsume + t.Cleanup(func() { + common.QuotaPerUnit = originalQuotaPerUnit + quotaSetting.EnableFreeModelPreConsume = originalEnableFreeModelPreConsume + require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(originalGroupRatio)) + }) + common.QuotaPerUnit = 1000 + quotaSetting.EnableFreeModelPreConsume = true + require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"default":1}`)) + + ctx, _ := gin.CreateTestContext(httptest.NewRecorder()) + info := &relaycommon.RelayInfo{ + OriginModelName: "suno_music", + UserGroup: "default", + UsingGroup: "default", + ChannelMeta: &relaycommon.ChannelMeta{}, + } + + priceData, err := ModelPriceHelperPerCall(ctx, info) + + require.NoError(t, err) + require.True(t, priceData.UsePrice) + require.Equal(t, float64(1), priceData.GroupRatioInfo.GroupRatio) + require.Equal(t, 0.1, priceData.ModelPrice) + require.Equal(t, 100, priceData.Quota) +} diff --git a/relay/mjproxy_handler.go b/relay/mjproxy_handler.go index 0a5c2279..0054d263 100644 --- a/relay/mjproxy_handler.go +++ b/relay/mjproxy_handler.go @@ -2,7 +2,7 @@ package relay import ( "bytes" - "encoding/json" + "errors" "fmt" "io" "log" @@ -22,6 +22,7 @@ import ( "github.com/MAX-API-Next/MAX-API/service" "github.com/MAX-API-Next/MAX-API/setting" "github.com/MAX-API-Next/MAX-API/setting/system_setting" + "github.com/MAX-API-Next/MAX-API/types" "github.com/gin-gonic/gin" ) @@ -30,6 +31,61 @@ func writeMidjourneyStatusCode(c *gin.Context, upstreamStatusCode int) { c.Writer.WriteHeader(service.MapStatusCode(upstreamStatusCode, c.GetString("status_code_mapping"))) } +func midjourneyTaskLookupErrorDescription(err error) string { + if errors.Is(err, model.ErrMidjourneyTaskAmbiguous) { + return "midjourney_task_ambiguous" + } + return "midjourney_task_lookup_failed" +} + +func classifyMidjourneySubmission(statusCode int, response *dto.MidjourneyResponse) (accepted, ambiguous bool) { + if response == nil { + return false, statusCode >= http.StatusInternalServerError + } + acceptedCode := response.Code == 1 || response.Code == 21 || response.Code == 22 + hasTaskID := strings.TrimSpace(response.Result) != "" + if acceptedCode && hasTaskID { + return true, false + } + if acceptedCode || statusCode >= http.StatusInternalServerError { + return false, true + } + return false, false +} + +func applyMidjourneyOriginChannel(c *gin.Context, info *relaycommon.RelayInfo, channel *model.Channel) *types.MaxAPIError { + key, keyIndex, maxAPIError := channel.GetNextEnabledKey() + if maxAPIError != nil { + return maxAPIError + } + + common.SetContextKey(c, constant.ContextKeyChannelId, channel.Id) + common.SetContextKey(c, constant.ContextKeyChannelName, channel.Name) + common.SetContextKey(c, constant.ContextKeyChannelType, channel.Type) + common.SetContextKey(c, constant.ContextKeyChannelCreateTime, channel.CreatedTime) + common.SetContextKey(c, constant.ContextKeyChannelBaseUrl, channel.GetBaseURL()) + common.SetContextKey(c, constant.ContextKeyChannelKey, key) + common.SetContextKey(c, constant.ContextKeyChannelSetting, channel.GetSetting()) + common.SetContextKey(c, constant.ContextKeyChannelOtherSetting, channel.GetOtherSettings()) + common.SetContextKey(c, constant.ContextKeyChannelParamOverride, channel.GetParamOverride()) + common.SetContextKey(c, constant.ContextKeyChannelHeaderOverride, channel.GetHeaderOverride()) + common.SetContextKey(c, constant.ContextKeyChannelAutoBan, channel.GetAutoBan()) + common.SetContextKey(c, constant.ContextKeyChannelModelMapping, channel.GetModelMapping()) + common.SetContextKey(c, constant.ContextKeyChannelStatusCodeMapping, channel.GetStatusCodeMapping()) + common.SetContextKey(c, constant.ContextKeyChannelIsMultiKey, channel.ChannelInfo.IsMultiKey) + common.SetContextKey(c, constant.ContextKeyChannelMultiKeyIndex, keyIndex) + if channel.OpenAIOrganization != nil { + common.SetContextKey(c, constant.ContextKeyChannelOrganization, *channel.OpenAIOrganization) + } else { + common.SetContextKey(c, constant.ContextKeyChannelOrganization, "") + } + + if info != nil { + info.InitChannelMeta(c) + } + return nil +} + func RelayMidjourneyImage(c *gin.Context) { taskId := c.Param("id") userID, userErr := strconv.Atoi(c.Query("uid")) @@ -41,8 +97,15 @@ func RelayMidjourneyImage(c *gin.Context) { }) return } - midjourneyTask := model.GetByMJId(userID, taskId) - if midjourneyTask == nil { + midjourneyTask, lookupErr := model.GetUniqueMidjourneyByUserAndMJID(userID, taskId) + if lookupErr != nil && !errors.Is(lookupErr, model.ErrMidjourneyTaskAmbiguous) { + common.SysError("midjourney image task lookup failed: " + lookupErr.Error()) + c.JSON(http.StatusInternalServerError, gin.H{ + "error": "midjourney_task_lookup_failed", + }) + return + } + if lookupErr != nil || midjourneyTask == nil { c.JSON(http.StatusForbidden, gin.H{ "error": "midjourney_image_authorization_failed", }) @@ -113,7 +176,13 @@ func RelayMidjourneyNotify(c *gin.Context) *dto.MidjourneyResponse { Result: "", } } - midjourneyTask := model.GetByOnlyMJId(midjRequest.MjId) + midjourneyTask, lookupErr := model.GetUniqueMidjourneyByMJID(midjRequest.MjId) + if lookupErr != nil { + return &dto.MidjourneyResponse{ + Code: constant.MjRequestError, + Description: midjourneyTaskLookupErrorDescription(lookupErr), + } + } if midjourneyTask == nil { return &dto.MidjourneyResponse{ Code: 4, @@ -122,19 +191,7 @@ func RelayMidjourneyNotify(c *gin.Context) *dto.MidjourneyResponse { Result: "", } } - midjourneyTask.Progress = midjRequest.Progress - midjourneyTask.PromptEn = midjRequest.PromptEn - midjourneyTask.State = midjRequest.State - midjourneyTask.SubmitTime = midjRequest.SubmitTime - midjourneyTask.StartTime = midjRequest.StartTime - midjourneyTask.FinishTime = midjRequest.FinishTime - midjourneyTask.ImageUrl = midjRequest.ImageUrl - midjourneyTask.VideoUrl = midjRequest.VideoUrl - videoUrlsStr, _ := json.Marshal(midjRequest.VideoUrls) - midjourneyTask.VideoUrls = string(videoUrlsStr) - midjourneyTask.Status = midjRequest.Status - midjourneyTask.FailReason = midjRequest.FailReason - err = midjourneyTask.Update() + err = service.UpdateMidjourneyTaskFromResponse(c.Request.Context(), midjourneyTask, midjRequest) if err != nil { return &dto.MidjourneyResponse{ Code: 4, @@ -172,21 +229,21 @@ func coverMidjourneyTaskDto(c *gin.Context, originTask *model.Midjourney) (midjo midjourneyTask.Prompt = originTask.Prompt if originTask.Buttons != "" { var buttons []dto.ActionButton - err := json.Unmarshal([]byte(originTask.Buttons), &buttons) + err := common.Unmarshal([]byte(originTask.Buttons), &buttons) if err == nil { midjourneyTask.Buttons = buttons } } if originTask.VideoUrls != "" { var videoUrls []dto.ImgUrls - err := json.Unmarshal([]byte(originTask.VideoUrls), &videoUrls) + err := common.Unmarshal([]byte(originTask.VideoUrls), &videoUrls) if err == nil { midjourneyTask.VideoUrls = videoUrls } } if originTask.Properties != "" { var properties dto.Properties - err := json.Unmarshal([]byte(originTask.Properties), &properties) + err := common.Unmarshal([]byte(originTask.Properties), &properties) if err == nil { midjourneyTask.Properties = &properties } @@ -194,6 +251,158 @@ func coverMidjourneyTaskDto(c *gin.Context, originTask *model.Midjourney) (midjo return } +func prepareMidjourneyBillingTask(c *gin.Context, info *relaycommon.RelayInfo, action string, priceData types.PriceData, charge bool) (*model.Task, *dto.MidjourneyResponse) { + if info.TaskRelayInfo == nil { + info.TaskRelayInfo = &relaycommon.TaskRelayInfo{} + } + info.Action = action + info.OriginModelName = service.CovertMjpActionToModelName(action) + info.PriceData = priceData + if !charge || priceData.FreeModel { + info.PriceData.Quota = 0 + } + if info.PriceData.Quota > 0 { + info.ForcePreConsume = true + if apiErr := service.PreConsumeBilling(c, info.PriceData.Quota, info); apiErr != nil { + return nil, &dto.MidjourneyResponse{Code: constant.MjRequestError, Description: apiErr.Error()} + } + } + task, taskErr := ensureTaskPlaceholder(constant.TaskPlatformMidjourney, info) + if taskErr != nil { + if info.Billing != nil { + info.Billing.Refund(c) + } + return nil, &dto.MidjourneyResponse{Code: constant.MjRequestError, Description: taskErr.Message} + } + return task, nil +} + +func prepareMidjourneySettlementInput(info *relaycommon.RelayInfo, quota int) (*model.BillingSettlementInput, error) { + if info == nil || info.Billing == nil { + return nil, nil + } + preparer, ok := info.Billing.(service.SettlementPreparer) + if !ok { + return nil, fmt.Errorf("midjourney billing session cannot prepare settlement") + } + return preparer.PrepareSettlement(quota) +} + +func prepareMidjourneyRefundInput(info *relaycommon.RelayInfo) (*model.BillingSettlementInput, error) { + if info == nil || info.Billing == nil { + return nil, nil + } + preparer, ok := info.Billing.(service.RefundSettlementPreparer) + if !ok { + return nil, fmt.Errorf("midjourney billing session cannot prepare refund") + } + return preparer.PrepareRefundSettlement() +} + +func failMidjourneySubmission(c *gin.Context, info *relaycommon.RelayInfo, task *model.Task, reason string) { + if task == nil { + return + } + refundInput, err := prepareMidjourneyRefundInput(info) + if err != nil { + common.SysError("prepare midjourney submission refund error: " + err.Error()) + _ = model.MarkTaskSubmitAmbiguous(task.ID, reason+"（退款 intent 无法持久化，请人工核对）") + return + } + if err := model.MarkTaskSubmitFailedWithSettlement(task.ID, reason, refundInput); err != nil { + common.SysError("persist midjourney submission failure error: " + err.Error()) + if info != nil && info.Billing != nil { + info.Billing.Refund(c) + } + return + } + if info != nil && info.Billing != nil { + info.Billing.Refund(c) + } +} + +func markMidjourneySubmissionAmbiguous(info *relaycommon.RelayInfo, task *model.Task, reason string) { + if info != nil { + info.UpstreamTaskOutcomeUnknown = true + } + if task != nil { + if err := model.MarkTaskSubmitAmbiguous(task.ID, reason); err != nil { + common.SysError("mark midjourney submission ambiguous error: " + err.Error()) + } + } +} + +func markMidjourneySubmissionResponseAmbiguous(info *relaycommon.RelayInfo, task *model.Task, midjourneyTask *model.Midjourney, reason string) { + if info != nil { + info.UpstreamTaskOutcomeUnknown = true + } + if task == nil { + return + } + if midjourneyTask == nil || strings.TrimSpace(midjourneyTask.MjId) == "" { + if err := model.MarkTaskSubmitAmbiguous(task.ID, reason); err != nil { + common.SysError("mark midjourney submission ambiguous error: " + err.Error()) + } + return + } + task.PrivateData.UpstreamTaskID = midjourneyTask.MjId + task.PrivateData.AwaitingUpstreamID = false + task.SetData(midjourneyTask) + if err := model.MarkTaskSubmitNeedsReview(task, reason); err != nil { + common.SysError("mark midjourney submission for manual review error: " + err.Error()) + } +} + +func finalizeMidjourneySubmission(c *gin.Context, info *relaycommon.RelayInfo, task *model.Task, midjourneyTask *model.Midjourney) (bool, error) { + if task == nil || midjourneyTask == nil { + return false, errors.New("midjourney submission result is incomplete") + } + task.SetData(midjourneyTask) + settlementInput, err := prepareMidjourneySettlementInput(info, midjourneyTask.Quota) + if err != nil { + return false, err + } + settlementOperationKey := "" + if settlementInput != nil { + settlementInput.Effect = service.BuildTaskSubmissionSettlementEffect(c, info, midjourneyTask.Quota) + settlementOperationKey = settlementInput.OperationKey + } + refundInput, err := prepareMidjourneyRefundInput(info) + if err != nil { + return false, err + } + created, refundDuplicate, err := model.FinalizeMidjourneySubmission(midjourneyTask, task, settlementInput, refundInput) + if err != nil { + task.PrivateData.UpstreamTaskID = midjourneyTask.MjId + task.PrivateData.AwaitingUpstreamID = false + markErr := model.MarkTaskSubmitNeedsReview(task, "上游已接受 Midjourney 任务，但本地任务/结算持久化失败："+err.Error()) + if markErr != nil { + common.SysError("mark midjourney submission for manual review error: " + markErr.Error()) + } + if info != nil { + info.UpstreamTaskOutcomeUnknown = true + } + return false, err + } + if !created { + if refundDuplicate && info != nil && info.Billing != nil { + info.Billing.Refund(c) + } + return false, nil + } + if info != nil && info.Billing != nil { + if err := service.SettleBilling(c, info, midjourneyTask.Quota); err != nil { + common.SysError("midjourney billing settlement remains pending/manual: " + err.Error()) + } + if settlementOperationKey != "" { + if err := model.ProcessBillingSettlementEffect(settlementOperationKey); err != nil { + common.SysError("midjourney billing effect remains pending: " + err.Error()) + } + } + } + return true, nil +} + func RelaySwapFace(c *gin.Context, info *relaycommon.RelayInfo) *dto.MidjourneyResponse { var swapFaceRequest dto.SwapFaceRequest err := common.UnmarshalBodyReusable(c, &swapFaceRequest) @@ -206,8 +415,7 @@ func RelaySwapFace(c *gin.Context, info *relaycommon.RelayInfo) *dto.MidjourneyR if swapFaceRequest.SourceBase64 == "" || swapFaceRequest.TargetBase64 == "" { return service.MidjourneyErrorWrapper(constant.MjRequestError, "sour_base64_and_target_base64_is_required") } - modelName := service.CovertMjpActionToModelName(constant.MjActionSwapFace) - + info.OriginModelName = service.CovertMjpActionToModelName(constant.MjActionSwapFace) priceData, err := helper.ModelPriceHelperPerCall(c, info) if err != nil { return &dto.MidjourneyResponse{ @@ -216,51 +424,22 @@ func RelaySwapFace(c *gin.Context, info *relaycommon.RelayInfo) *dto.MidjourneyR } } - userQuota, err := model.GetUserQuota(info.UserId, false) - if err != nil { - return &dto.MidjourneyResponse{ - Code: 4, - Description: err.Error(), - } - } - - if userQuota-int64(priceData.Quota) < 0 { - return &dto.MidjourneyResponse{ - Code: 4, - Description: "quota_not_enough", - } + billingTask, billingErr := prepareMidjourneyBillingTask(c, info, constant.MjActionSwapFace, priceData, true) + if billingErr != nil { + return billingErr } requestURL := getMjRequestPath(c.Request.URL.String()) baseURL := c.GetString("base_url") fullRequestURL := fmt.Sprintf("%s%s", baseURL, requestURL) - mjResp, _, err := service.DoMidjourneyHttpRequest(c, time.Second*60, fullRequestURL) + mjResp, _, requestSent, err := service.DoMidjourneyHttpRequest(c, time.Second*60, fullRequestURL) if err != nil { + if requestSent { + markMidjourneySubmissionAmbiguous(info, billingTask, "Midjourney 换脸请求已发送但结果未知，请人工核对") + } else { + failMidjourneySubmission(c, info, billingTask, err.Error()) + } return &mjResp.Response } - defer func() { - if mjResp.StatusCode == 200 && mjResp.Response.Code == 1 { - err := service.PostConsumeQuotaOnce(info, "midjourney-success", priceData.Quota, 0, true) - if err != nil { - common.SysLog("error consuming token remain quota: " + err.Error()) - } - - tokenName := c.GetString("token_name") - logContent := fmt.Sprintf("模型固定价格 %.2f，分组倍率 %.2f，操作 %s", priceData.ModelPrice, priceData.GroupRatioInfo.GroupRatio, constant.MjActionSwapFace) - other := service.GenerateMjOtherInfo(info, priceData) - model.RecordConsumeLog(c, info.UserId, model.RecordConsumeLogParams{ - ChannelId: info.ChannelId, - ModelName: modelName, - TokenName: tokenName, - Quota: priceData.Quota, - Content: logContent, - TokenId: info.TokenId, - Group: info.UsingGroup, - Other: other, - }) - model.UpdateUserUsedQuotaAndRequestCount(info.UserId, priceData.Quota) - model.UpdateChannelUsedQuota(info.ChannelId, priceData.Quota) - } - }() midjResponse := &mjResp.Response midjourneyTask := &model.Midjourney{ UserId: info.UserId, @@ -279,14 +458,23 @@ func RelaySwapFace(c *gin.Context, info *relaycommon.RelayInfo) *dto.MidjourneyR Progress: "0%", FailReason: "", ChannelId: c.GetInt("channel_id"), - Quota: priceData.Quota, + Quota: info.PriceData.Quota, } - err = midjourneyTask.Insert() - if err != nil { - return service.MidjourneyErrorWrapper(constant.MjRequestError, "insert_midjourney_task_failed") + accepted, ambiguous := classifyMidjourneySubmission(mjResp.StatusCode, midjResponse) + if ambiguous { + reason := fmt.Sprintf("Midjourney 换脸响应结果不确定（status=%d code=%d），已保留预扣费待人工核对", mjResp.StatusCode, midjResponse.Code) + markMidjourneySubmissionResponseAmbiguous(info, billingTask, midjourneyTask, reason) + return service.MidjourneyErrorWrapper(constant.MjRequestError, "midjourney_submission_outcome_unknown") + } + if !accepted { + failMidjourneySubmission(c, info, billingTask, midjResponse.Description) + } else { + if _, err := finalizeMidjourneySubmission(c, info, billingTask, midjourneyTask); err != nil { + return service.MidjourneyErrorWrapper(constant.MjRequestError, "persist_midjourney_task_failed") + } } writeMidjourneyStatusCode(c, mjResp.StatusCode) - respBody, err := json.Marshal(midjResponse) + respBody, err := common.Marshal(midjResponse) if err != nil { return service.MidjourneyErrorWrapper(constant.MjRequestError, "unmarshal_response_body_failed") } @@ -300,7 +488,10 @@ func RelaySwapFace(c *gin.Context, info *relaycommon.RelayInfo) *dto.MidjourneyR func RelayMidjourneyTaskImageSeed(c *gin.Context) *dto.MidjourneyResponse { taskId := c.Param("id") userId := c.GetInt("id") - originTask := model.GetByMJId(userId, taskId) + originTask, lookupErr := model.GetUniqueMidjourneyByUserAndMJID(userId, taskId) + if lookupErr != nil { + return service.MidjourneyErrorWrapper(constant.MjRequestError, midjourneyTaskLookupErrorDescription(lookupErr)) + } if originTask == nil { return service.MidjourneyErrorWrapper(constant.MjRequestError, "task_no_found") } @@ -311,18 +502,19 @@ func RelayMidjourneyTaskImageSeed(c *gin.Context) *dto.MidjourneyResponse { if channel.Status != common.ChannelStatusEnabled { return service.MidjourneyErrorWrapper(constant.MjRequestError, "该任务所属渠道已被禁用") } - c.Set("channel_id", originTask.ChannelId) - c.Request.Header.Set("Authorization", fmt.Sprintf("Bearer %s", channel.Key)) + if maxAPIError := applyMidjourneyOriginChannel(c, nil, channel); maxAPIError != nil { + return service.MidjourneyErrorWrapper(constant.MjRequestError, maxAPIError.Error()) + } requestURL := getMjRequestPath(c.Request.URL.String()) fullRequestURL := fmt.Sprintf("%s%s", channel.GetBaseURL(), requestURL) - midjResponseWithStatus, _, err := service.DoMidjourneyHttpRequest(c, time.Second*30, fullRequestURL) + midjResponseWithStatus, _, _, err := service.DoMidjourneyHttpRequest(c, time.Second*30, fullRequestURL) if err != nil { return &midjResponseWithStatus.Response } midjResponse := &midjResponseWithStatus.Response writeMidjourneyStatusCode(c, midjResponseWithStatus.StatusCode) - respBody, err := json.Marshal(midjResponse) + respBody, err := common.Marshal(midjResponse) if err != nil { return service.MidjourneyErrorWrapper(constant.MjRequestError, "unmarshal_response_body_failed") } @@ -337,7 +529,13 @@ func RelayMidjourneyTask(c *gin.Context, relayMode int) *dto.MidjourneyResponse switch relayMode { case relayconstant.RelayModeMidjourneyTaskFetch: taskId := c.Param("id") - originTask := model.GetByMJId(userId, taskId) + originTask, lookupErr := model.GetUniqueMidjourneyByUserAndMJID(userId, taskId) + if lookupErr != nil { + return &dto.MidjourneyResponse{ + Code: constant.MjRequestError, + Description: midjourneyTaskLookupErrorDescription(lookupErr), + } + } if originTask == nil { return &dto.MidjourneyResponse{ Code: 4, @@ -345,7 +543,7 @@ func RelayMidjourneyTask(c *gin.Context, relayMode int) *dto.MidjourneyResponse } } midjourneyTask := coverMidjourneyTaskDto(c, originTask) - respBody, err = json.Marshal(midjourneyTask) + respBody, err = common.Marshal(midjourneyTask) if err != nil { return &dto.MidjourneyResponse{ Code: 4, @@ -374,7 +572,7 @@ func RelayMidjourneyTask(c *gin.Context, relayMode int) *dto.MidjourneyResponse if tasks == nil { tasks = make([]dto.MidjourneyDto, 0) } - respBody, err = json.Marshal(tasks) + respBody, err = common.Marshal(tasks) if err != nil { return &dto.MidjourneyResponse{ Code: 4, @@ -469,7 +667,10 @@ func RelayMidjourneySubmit(c *gin.Context, relayInfo *relaycommon.RelayInfo) *dt mjId = midjRequest.TaskId } - originTask := model.GetByMJId(relayInfo.UserId, mjId) + originTask, lookupErr := model.GetUniqueMidjourneyByUserAndMJID(relayInfo.UserId, mjId) + if lookupErr != nil { + return service.MidjourneyErrorWrapper(constant.MjRequestError, midjourneyTaskLookupErrorDescription(lookupErr)) + } if originTask == nil { return service.MidjourneyErrorWrapper(constant.MjRequestError, "task_not_found") } else { //原任务的Status=SUCCESS，则可以做放大UPSCALE、变换VARIATION等动作，此时必须使用原来的请求地址才能正确处理 @@ -485,9 +686,9 @@ func RelayMidjourneySubmit(c *gin.Context, relayInfo *relaycommon.RelayInfo) *dt if channel.Status != common.ChannelStatusEnabled { return service.MidjourneyErrorWrapper(constant.MjRequestError, "该任务所属渠道已被禁用") } - c.Set("base_url", channel.GetBaseURL()) - c.Set("channel_id", originTask.ChannelId) - c.Request.Header.Set("Authorization", fmt.Sprintf("Bearer %s", channel.Key)) + if maxAPIError := applyMidjourneyOriginChannel(c, relayInfo, channel); maxAPIError != nil { + return service.MidjourneyErrorWrapper(constant.MjRequestError, maxAPIError.Error()) + } logger.LogDebug(c, "Midjourney action uses origin channel: id=%s, base_url=%s", strconv.Itoa(originTask.ChannelId), channel.GetBaseURL()) } midjRequest.Prompt = originTask.Prompt @@ -513,8 +714,7 @@ func RelayMidjourneySubmit(c *gin.Context, relayInfo *relaycommon.RelayInfo) *dt fullRequestURL := fmt.Sprintf("%s%s", baseURL, requestURL) - modelName := service.CovertMjpActionToModelName(midjRequest.Action) - + relayInfo.OriginModelName = service.CovertMjpActionToModelName(midjRequest.Action) priceData, err := helper.ModelPriceHelperPerCall(c, relayInfo) if err != nil { return &dto.MidjourneyResponse{ @@ -523,51 +723,22 @@ func RelayMidjourneySubmit(c *gin.Context, relayInfo *relaycommon.RelayInfo) *dt } } - userQuota, err := model.GetUserQuota(relayInfo.UserId, false) - if err != nil { - return &dto.MidjourneyResponse{ - Code: 4, - Description: err.Error(), - } + billingTask, billingErr := prepareMidjourneyBillingTask(c, relayInfo, midjRequest.Action, priceData, consumeQuota) + if billingErr != nil { + return billingErr } - if consumeQuota && userQuota-int64(priceData.Quota) < 0 { - return &dto.MidjourneyResponse{ - Code: 4, - Description: "quota_not_enough", - } - } - - midjResponseWithStatus, responseBody, err := service.DoMidjourneyHttpRequest(c, time.Second*60, fullRequestURL) + midjResponseWithStatus, responseBody, requestSent, err := service.DoMidjourneyHttpRequest(c, time.Second*60, fullRequestURL) if err != nil { + if requestSent { + markMidjourneySubmissionAmbiguous(relayInfo, billingTask, "Midjourney 请求已发送但结果未知，请人工核对") + } else { + failMidjourneySubmission(c, relayInfo, billingTask, err.Error()) + } return &midjResponseWithStatus.Response } midjResponse := &midjResponseWithStatus.Response - defer func() { - if consumeQuota && midjResponseWithStatus.StatusCode == 200 { - err := service.PostConsumeQuotaOnce(relayInfo, "midjourney-success", priceData.Quota, 0, true) - if err != nil { - common.SysLog("error consuming token remain quota: " + err.Error()) - } - tokenName := c.GetString("token_name") - logContent := fmt.Sprintf("模型固定价格 %.2f，分组倍率 %.2f，操作 %s，ID %s", priceData.ModelPrice, priceData.GroupRatioInfo.GroupRatio, midjRequest.Action, midjResponse.Result) - other := service.GenerateMjOtherInfo(relayInfo, priceData) - model.RecordConsumeLog(c, relayInfo.UserId, model.RecordConsumeLogParams{ - ChannelId: relayInfo.ChannelId, - ModelName: modelName, - TokenName: tokenName, - Quota: priceData.Quota, - Content: logContent, - TokenId: relayInfo.TokenId, - Group: relayInfo.UsingGroup, - Other: other, - }) - model.UpdateUserUsedQuotaAndRequestCount(relayInfo.UserId, priceData.Quota) - model.UpdateChannelUsedQuota(relayInfo.ChannelId, priceData.Quota) - } - }() - // 文档：https://github.com/novicezk/midjourney-proxy/blob/main/docs/api.md //1-提交成功 // 21-任务已存在（处理中或者有结果了） {"code":21,"description":"任务已存在","result":"0741798445574458","properties":{"status":"SUCCESS","imageUrl":"https://xxxx"}} @@ -592,22 +763,20 @@ func RelayMidjourneySubmit(c *gin.Context, relayInfo *relaycommon.RelayInfo) *dt Progress: "0%", FailReason: "", ChannelId: c.GetInt("channel_id"), - Quota: priceData.Quota, + Quota: relayInfo.PriceData.Quota, } if midjResponse.Code == 3 { //无实例账号自动禁用渠道（No available account instance） channel, err := model.GetChannelById(midjourneyTask.ChannelId, true) if err != nil { common.SysLog("get_channel_null: " + err.Error()) - } - if channel.GetAutoBan() && common.AutomaticDisableChannelEnabled { + } else if channel.GetAutoBan() && common.AutomaticDisableChannelEnabled { model.UpdateChannelStatus(midjourneyTask.ChannelId, "", 2, "No available account instance") } } if midjResponse.Code != 1 && midjResponse.Code != 21 && midjResponse.Code != 22 { //非1-提交成功,21-任务已存在和22-排队中，则记录错误原因 midjourneyTask.FailReason = midjResponse.Description - consumeQuota = false } if midjResponse.Code == 21 { //21-任务已存在（处理中或者有结果了） @@ -637,19 +806,22 @@ func RelayMidjourneySubmit(c *gin.Context, relayInfo *relaycommon.RelayInfo) *dt midjourneyTask.Progress = "100%" midjourneyTask.Status = "SUCCESS" } - err = midjourneyTask.Insert() - if err != nil { - return &dto.MidjourneyResponse{ - Code: 4, - Description: "insert_midjourney_task_failed", - } - } - if midjResponse.Code == 22 { //22-排队中，说明任务已存在 //修改返回值 newBody := strings.Replace(string(responseBody), `"code":22`, `"code":1`, -1) responseBody = []byte(newBody) } + accepted, ambiguous := classifyMidjourneySubmission(midjResponseWithStatus.StatusCode, midjResponse) + if ambiguous { + reason := fmt.Sprintf("Midjourney 响应结果不确定（status=%d code=%d），已保留预扣费待人工核对", midjResponseWithStatus.StatusCode, midjResponse.Code) + markMidjourneySubmissionResponseAmbiguous(relayInfo, billingTask, midjourneyTask, reason) + return service.MidjourneyErrorWrapper(constant.MjRequestError, "midjourney_submission_outcome_unknown") + } + if !accepted { + failMidjourneySubmission(c, relayInfo, billingTask, midjResponse.Description) + } else if _, err := finalizeMidjourneySubmission(c, relayInfo, billingTask, midjourneyTask); err != nil { + return service.MidjourneyErrorWrapper(constant.MjRequestError, "persist_midjourney_task_failed") + } //resp.Body = io.NopCloser(bytes.NewBuffer(responseBody)) bodyReader := io.NopCloser(bytes.NewBuffer(responseBody)) diff --git a/relay/mjproxy_handler_test.go b/relay/mjproxy_handler_test.go index bb95259e..040e87a1 100644 --- a/relay/mjproxy_handler_test.go +++ b/relay/mjproxy_handler_test.go @@ -1,6 +1,8 @@ package relay import ( + "bytes" + "errors" "net/http" "net/http/httptest" "strconv" @@ -8,7 +10,10 @@ import ( "time" "github.com/MAX-API-Next/MAX-API/common" + "github.com/MAX-API-Next/MAX-API/constant" + "github.com/MAX-API-Next/MAX-API/dto" "github.com/MAX-API-Next/MAX-API/model" + relaycommon "github.com/MAX-API-Next/MAX-API/relay/common" "github.com/MAX-API-Next/MAX-API/service" "github.com/gin-gonic/gin" "github.com/glebarez/sqlite" @@ -64,4 +69,260 @@ func TestRelayMidjourneyImageRejectsUnsignedAndWrongUserURLs(t *testing.T) { router.ServeHTTP(wrongUser, httptest.NewRequest(http.MethodGet, wrongUserURL, nil)) require.Equal(t, http.StatusForbidden, wrongUser.Code) require.Contains(t, wrongUser.Body.String(), "midjourney_image_authorization_failed") + + require.NoError(t, db.Create(&model.Midjourney{ + Id: 2, UserId: 42, MjId: "mj-private-task", ImageUrl: "https://example.com/duplicate.png", + }).Error) + ambiguousURL := "/mj/image/mj-private-task?uid=42&expires=" + + strconv.FormatInt(expiresAt, 10) + "&signature=" + + service.SignMidjourneyImageURL("mj-private-task", 42, expiresAt) + ambiguous := httptest.NewRecorder() + router.ServeHTTP(ambiguous, httptest.NewRequest(http.MethodGet, ambiguousURL, nil)) + require.Equal(t, http.StatusForbidden, ambiguous.Code) + require.Contains(t, ambiguous.Body.String(), "midjourney_image_authorization_failed") +} + +func TestRelayMidjourneyImageReturnsServerErrorForLookupFailure(t *testing.T) { + gin.SetMode(gin.TestMode) + oldDB := model.DB + oldSecret := common.CryptoSecret + var db *gorm.DB + callbackRegistered := false + const callbackName = "test:midjourney-lookup-error" + t.Cleanup(func() { + model.DB = oldDB + common.CryptoSecret = oldSecret + if db == nil { + return + } + if callbackRegistered { + _ = db.Callback().Query().Remove(callbackName) + } + if sqlDB, dbErr := db.DB(); dbErr == nil { + _ = sqlDB.Close() + } + }) + common.CryptoSecret = "midjourney-handler-test-secret" + var err error + db, err = gorm.Open(sqlite.Open("file:mjproxy_lookup_error?mode=memory&cache=shared"), &gorm.Config{}) + require.NoError(t, err) + model.DB = db + require.NoError(t, db.AutoMigrate(&model.Midjourney{})) + callbackHit := false + require.NoError(t, db.Callback().Query().Before("gorm:query").Register(callbackName, func(tx *gorm.DB) { + if tx.Statement.Table == "midjourneys" { + callbackHit = true + tx.AddError(errors.New("database unavailable")) + } + })) + callbackRegistered = true + + expiresAt := time.Now().Add(service.MidjourneyImageURLTTL).Unix() + requestURL := "/mj/image/mj-lookup-error?uid=42&expires=" + + strconv.FormatInt(expiresAt, 10) + "&signature=" + + service.SignMidjourneyImageURL("mj-lookup-error", 42, expiresAt) + recorder := httptest.NewRecorder() + router := gin.New() + router.GET("/mj/image/:id", RelayMidjourneyImage) + router.ServeHTTP(recorder, httptest.NewRequest(http.MethodGet, requestURL, nil)) + + require.Equal(t, http.StatusInternalServerError, recorder.Code) + require.Contains(t, recorder.Body.String(), "midjourney_task_lookup_failed") + require.True(t, callbackHit, "lookup error callback was not invoked") +} + +func TestRelayMidjourneyTaskImageSeedUsesOriginChannelKey(t *testing.T) { + gin.SetMode(gin.TestMode) + receivedKey := make(chan string, 1) + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + receivedKey <- r.Header.Get("mj-api-secret") + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"code":1,"description":"ok","result":"seed"}`)) + })) + t.Cleanup(upstream.Close) + + oldDB := model.DB + oldSQLite := common.UsingSQLite + oldMemoryCache := common.MemoryCacheEnabled + db, err := gorm.Open(sqlite.Open("file:mjproxy_origin_channel?mode=memory&cache=shared"), &gorm.Config{}) + require.NoError(t, err) + model.DB = db + common.UsingSQLite = true + common.MemoryCacheEnabled = false + require.NoError(t, db.AutoMigrate(&model.Channel{}, &model.Midjourney{})) + baseURL := upstream.URL + require.NoError(t, db.Create(&model.Channel{ + Id: 91, Type: constant.ChannelTypeMidjourney, Key: "origin-channel-key", + Status: common.ChannelStatusEnabled, Name: "origin-midjourney", BaseURL: &baseURL, + }).Error) + require.NoError(t, db.Create(&model.Midjourney{ + Id: 92, UserId: 42, ChannelId: 91, MjId: "provider-task-for-seed", + }).Error) + service.InitHttpClient() + t.Cleanup(func() { + model.DB = oldDB + common.UsingSQLite = oldSQLite + common.MemoryCacheEnabled = oldMemoryCache + if sqlDB, dbErr := db.DB(); dbErr == nil { + _ = sqlDB.Close() + } + }) + + router := gin.New() + router.GET("/mj/task/:id/image-seed", func(c *gin.Context) { + c.Set("id", 42) + if mjErr := RelayMidjourneyTaskImageSeed(c); mjErr != nil { + c.JSON(http.StatusBadGateway, mjErr) + } + }) + recorder := httptest.NewRecorder() + router.ServeHTTP(recorder, httptest.NewRequest(http.MethodGet, "/mj/task/provider-task-for-seed/image-seed", nil)) + + require.Equal(t, http.StatusOK, recorder.Code) + select { + case key := <-receivedKey: + require.Equal(t, "origin-channel-key", key) + case <-time.After(time.Second): + t.Fatal("upstream request was not received") + } +} + +func TestApplyMidjourneyOriginChannelUpdatesBillingMetadata(t *testing.T) { + gin.SetMode(gin.TestMode) + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + common.SetContextKey(c, constant.ContextKeyChannelId, 11) + common.SetContextKey(c, constant.ContextKeyChannelType, constant.ChannelTypeMidjourneyPlus) + common.SetContextKey(c, constant.ContextKeyChannelBaseUrl, "https://selected.example") + common.SetContextKey(c, constant.ContextKeyChannelKey, "selected-key") + info := &relaycommon.RelayInfo{} + info.InitChannelMeta(c) + + baseURL := "https://origin.example" + channel := &model.Channel{ + Id: 21, Type: constant.ChannelTypeMidjourney, Key: "origin-key", + Status: common.ChannelStatusEnabled, Name: "origin", BaseURL: &baseURL, + } + require.Nil(t, applyMidjourneyOriginChannel(c, info, channel)) + + require.Equal(t, 21, common.GetContextKeyInt(c, constant.ContextKeyChannelId)) + require.Equal(t, constant.ChannelTypeMidjourney, common.GetContextKeyInt(c, constant.ContextKeyChannelType)) + require.Equal(t, "https://origin.example", common.GetContextKeyString(c, constant.ContextKeyChannelBaseUrl)) + require.Equal(t, "origin-key", common.GetContextKeyString(c, constant.ContextKeyChannelKey)) + require.NotNil(t, info.ChannelMeta) + require.Equal(t, 21, info.ChannelId) + require.Equal(t, constant.ChannelTypeMidjourney, info.ChannelType) + require.Equal(t, "https://origin.example", info.ChannelBaseUrl) + require.Equal(t, "origin-key", info.ApiKey) +} + +func TestClassifyMidjourneySubmissionPreservesUncertainOutcomes(t *testing.T) { + tests := []struct { + name string + status int + response dto.MidjourneyResponse + accepted bool + ambiguous bool + }{ + {name: "normal accepted", status: http.StatusOK, response: dto.MidjourneyResponse{Code: 1, Result: "task-id"}, accepted: true}, + {name: "accepted body wins over server status", status: http.StatusBadGateway, response: dto.MidjourneyResponse{Code: 1, Result: "task-id"}, accepted: true}, + {name: "accepted code without identity", status: http.StatusOK, response: dto.MidjourneyResponse{Code: 1}, ambiguous: true}, + {name: "server error without identity", status: http.StatusInternalServerError, response: dto.MidjourneyResponse{Code: 23}, ambiguous: true}, + {name: "explicit queue rejection", status: http.StatusOK, response: dto.MidjourneyResponse{Code: 23, Result: "queue-marker"}}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + accepted, ambiguous := classifyMidjourneySubmission(test.status, &test.response) + require.Equal(t, test.accepted, accepted) + require.Equal(t, test.ambiguous, ambiguous) + }) + } +} + +func TestRelayMidjourneyNotifyCompletesShadowTask(t *testing.T) { + gin.SetMode(gin.TestMode) + oldDB := model.DB + db, err := gorm.Open(sqlite.Open("file:mjproxy_notify_shadow_task?mode=memory&cache=shared"), &gorm.Config{}) + require.NoError(t, err) + model.DB = db + require.NoError(t, db.AutoMigrate( + &model.Midjourney{}, &model.MidjourneyBillingClaim{}, &model.Task{}, &model.BillingSettlement{}, + )) + t.Cleanup(func() { + model.DB = oldDB + if sqlDB, dbErr := db.DB(); dbErr == nil { + _ = sqlDB.Close() + } + }) + + shadowTask := model.Task{ + TaskID: "task_midjourney_notify_success", Platform: constant.TaskPlatformMidjourney, + UserId: 301, ChannelId: 401, Quota: 30, Status: model.TaskStatusSubmitted, + Progress: "0%", SubmitTime: time.Now().Add(-2 * time.Hour).Unix(), + PrivateData: model.TaskPrivateData{ + UpstreamTaskID: "provider-notify-success", BillingRequestId: "notify-success-request", + }, + } + require.NoError(t, db.Create(&shadowTask).Error) + midjourneyTask := model.Midjourney{ + UserId: 301, ChannelId: 401, MjId: "provider-notify-success", + Code: 1, Status: "SUCCESS", Progress: "100%", Quota: 30, + SubmitTime: shadowTask.SubmitTime * 1000, FinishTime: 1700000000000, + } + require.NoError(t, db.Create(&midjourneyTask).Error) + require.NoError(t, db.Create(&model.MidjourneyBillingClaim{ + ChannelID: 401, MjID: midjourneyTask.MjId, UserID: 301, + BillingTaskID: shadowTask.ID, BillingRequestID: "notify-success-request", + CreatedAt: time.Now().Unix(), + }).Error) + + requestBody := []byte(`{"id":"provider-notify-success","status":"SUCCESS","progress":"100%","finishTime":1700000000000}`) + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/mj/notify", bytes.NewReader(requestBody)) + c.Request.Header.Set("Content-Type", "application/json") + require.Nil(t, RelayMidjourneyNotify(c)) + + var storedTask model.Task + require.NoError(t, db.First(&storedTask, shadowTask.ID).Error) + require.EqualValues(t, model.TaskStatusSuccess, storedTask.Status) + require.Equal(t, "100%", storedTask.Progress) + require.Equal(t, 30, storedTask.Quota) + require.Empty(t, model.GetTimedOutUnfinishedTasks(time.Now().Unix(), 10)) +} + +func TestRelayMidjourneyNotifyRejectsAmbiguousProviderID(t *testing.T) { + gin.SetMode(gin.TestMode) + oldDB := model.DB + db, err := gorm.Open(sqlite.Open("file:mjproxy_notify_ambiguous?mode=memory&cache=shared"), &gorm.Config{}) + require.NoError(t, err) + model.DB = db + require.NoError(t, db.AutoMigrate(&model.Midjourney{})) + t.Cleanup(func() { + model.DB = oldDB + if sqlDB, dbErr := db.DB(); dbErr == nil { + _ = sqlDB.Close() + } + }) + + sharedID := "provider-notify-shared" + require.NoError(t, db.Create(&model.Midjourney{ + UserId: 501, ChannelId: 601, MjId: sharedID, Status: "", Progress: "0%", + }).Error) + require.NoError(t, db.Create(&model.Midjourney{ + UserId: 502, ChannelId: 602, MjId: sharedID, Status: "", Progress: "0%", + }).Error) + + requestBody := []byte(`{"id":"provider-notify-shared","status":"SUCCESS","progress":"100%"}`) + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + c.Request = httptest.NewRequest(http.MethodPost, "/mj/notify", bytes.NewReader(requestBody)) + c.Request.Header.Set("Content-Type", "application/json") + mjErr := RelayMidjourneyNotify(c) + require.NotNil(t, mjErr) + require.Equal(t, "midjourney_task_ambiguous", mjErr.Description) + + var tasks []model.Midjourney + require.NoError(t, db.Order("id").Find(&tasks).Error) + require.Len(t, tasks, 2) + require.Equal(t, "", tasks[0].Status) + require.Equal(t, "", tasks[1].Status) } diff --git a/relay/relay_task.go b/relay/relay_task.go index 5a7b6e43..a4450442 100644 --- a/relay/relay_task.go +++ b/relay/relay_task.go @@ -324,12 +324,7 @@ func RelayTaskSubmit(c *gin.Context, info *relaycommon.RelayInfo) (*TaskSubmitRe return nil, service.TaskErrorWrapper(err, "do_request_failed", http.StatusInternalServerError) } if resp != nil && resp.StatusCode != http.StatusOK { - if resp.Body != nil { - defer resp.Body.Close() - } - responseBody, _ := io.ReadAll(resp.Body) - taskErr := service.TaskErrorWrapper(fmt.Errorf("%s", string(responseBody)), "fail_to_fetch_task", resp.StatusCode) - return nil, mapUpstreamTaskError(c, taskErr) + return nil, mapUpstreamTaskError(c, service.TaskErrorFromUpstreamResponse(c, resp, "fail_to_fetch_task")) } if resp == nil { return nil, service.TaskErrorWrapperLocal(errors.New("task upstream response is nil"), "empty_upstream_response", http.StatusBadGateway) diff --git a/relay/responses_handler.go b/relay/responses_handler.go index 5ce1ee07..9a761c90 100644 --- a/relay/responses_handler.go +++ b/relay/responses_handler.go @@ -76,7 +76,15 @@ func ResponsesHelper(c *gin.Context, info *relaycommon.RelayInfo) (maxAPIError * if err != nil { return types.NewError(err, types.ErrorCodeReadRequestBodyFailed, types.ErrOptionWithSkipRetry()) } - requestBody = common.ReaderOnly(storage) + body, size, closer, err := relaycommon.PreparePassThroughJSONBody(storage, info.ChannelOtherSettings) + if err != nil { + return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry()) + } + if closer != nil { + defer closer.Close() + } + info.UpstreamRequestBodySize = size + requestBody = body } else { convertedRequest, err := adaptor.ConvertOpenAIResponsesRequest(c, info, *request) if err != nil { diff --git a/router/api-router.go b/router/api-router.go index 55caf4d4..90c465cc 100644 --- a/router/api-router.go +++ b/router/api-router.go @@ -43,7 +43,7 @@ func SetApiRouter(router *gin.Engine) { apiRouter.GET("/reset_password", middleware.CriticalRateLimit(), middleware.TurnstileCheck(), controller.SendPasswordResetEmail) apiRouter.POST("/user/reset", middleware.CriticalRateLimit(), anonymousRequestBodyLimit, controller.ResetPassword) // OAuth routes - specific routes must come before :provider wildcard - apiRouter.GET("/oauth/state", middleware.CriticalRateLimit(), controller.GenerateOAuthCode) + apiRouter.POST("/oauth/state", middleware.CriticalRateLimit(), middleware.DisableCache(), middleware.TryUserAuth(), anonymousRequestBodyLimit, controller.GenerateOAuthCode) apiRouter.POST("/oauth/email/bind", middleware.CriticalRateLimit(), anonymousRequestBodyLimit, controller.EmailBind) // Non-standard OAuth (WeChat, Telegram) - keep original routes apiRouter.GET("/oauth/wechat", middleware.CriticalRateLimit(), controller.WeChatAuth) @@ -51,7 +51,7 @@ func SetApiRouter(router *gin.Engine) { apiRouter.GET("/oauth/telegram/login", middleware.CriticalRateLimit(), controller.TelegramLogin) apiRouter.GET("/oauth/telegram/bind", middleware.CriticalRateLimit(), controller.TelegramBind) // Standard OAuth providers (GitHub, Discord, OIDC, LinuxDO) - unified route - apiRouter.GET("/oauth/:provider", middleware.CriticalRateLimit(), controller.HandleOAuth) + apiRouter.GET("/oauth/:provider", middleware.CriticalRateLimit(), middleware.DisableCache(), middleware.TryUserAuth(), controller.HandleOAuth) apiRouter.GET("/ratio_config", middleware.CriticalRateLimit(), controller.GetRatioConfig) apiRouter.POST("/stripe/webhook", anonymousRequestBodyLimit, controller.StripeWebhook) @@ -86,7 +86,7 @@ func SetApiRouter(router *gin.Engine) { selfRoute.GET("/models", controller.GetUserModels) selfRoute.PUT("/self", middleware.CriticalRateLimit(), controller.UpdateSelf) selfRoute.DELETE("/self", middleware.CriticalRateLimit(), middleware.DisableCache(), middleware.SecureVerificationRequired("account_delete"), controller.DeleteSelf) - selfRoute.POST("/token", middleware.CriticalRateLimit(), middleware.DisableCache(), middleware.SecureVerificationRequired("access_token"), controller.GenerateAccessToken) + selfRoute.POST("/token", middleware.CriticalRateLimit(), middleware.UserCriticalRateLimit(), middleware.DisableCache(), middleware.SecureVerificationRequired("access_token"), controller.GenerateAccessToken) selfRoute.GET("/passkey", controller.PasskeyStatus) selfRoute.POST("/passkey/register/begin", controller.PasskeyRegisterBegin) selfRoute.POST("/passkey/register/finish", controller.PasskeyRegisterFinish) @@ -106,7 +106,7 @@ func SetApiRouter(router *gin.Engine) { selfRoute.POST("/waffo/pay", middleware.CriticalRateLimit(), controller.RequestWaffoPay) selfRoute.POST("/waffo-pancake/amount", controller.RequestWaffoPancakeAmount) selfRoute.POST("/waffo-pancake/pay", middleware.CriticalRateLimit(), controller.RequestWaffoPancakePay) - selfRoute.POST("/aff_transfer", controller.TransferAffQuota) + selfRoute.POST("/aff_transfer", middleware.CriticalRateLimit(), middleware.UserCriticalRateLimit(), controller.TransferAffQuota) selfRoute.PUT("/setting", controller.UpdateUserSetting) // 2FA routes diff --git a/router/api_router_test.go b/router/api_router_test.go index 10c08c90..6a5b692c 100644 --- a/router/api_router_test.go +++ b/router/api_router_test.go @@ -49,6 +49,43 @@ func TestUserTokenRouteUsesPostOnly(t *testing.T) { } } +func TestOAuthStateRouteUsesPostOnly(t *testing.T) { + gin.SetMode(gin.TestMode) + engine := gin.New() + SetApiRouter(engine) + + var hasPost bool + var hasGet bool + for _, route := range engine.Routes() { + if route.Path != "/api/oauth/state" { + continue + } + hasPost = hasPost || route.Method == http.MethodPost + hasGet = hasGet || route.Method == http.MethodGet + } + require.True(t, hasPost) + require.False(t, hasGet) +} + +func TestOAuthStateOpenAPIUsesPostWithBoundRequest(t *testing.T) { + documentBytes, err := os.ReadFile("../docs/openapi/api.json") + require.NoError(t, err) + + var document struct { + Paths map[string]map[string]map[string]any `json:"paths"` + } + require.NoError(t, common.Unmarshal(documentBytes, &document)) + oauthPath := document.Paths["/api/oauth/state"] + require.Contains(t, oauthPath, "post") + require.NotContains(t, oauthPath, "get") + requestBody := oauthPath["post"]["requestBody"].(map[string]any) + content := requestBody["content"].(map[string]any) + schema := content["application/json"].(map[string]any)["schema"].(map[string]any) + properties := schema["properties"].(map[string]any) + require.Contains(t, properties, "provider") + require.Contains(t, properties, "intent") +} + func TestUserTokenOpenAPIUsesPostOnly(t *testing.T) { documentBytes, err := os.ReadFile("../docs/openapi/api.json") require.NoError(t, err) @@ -170,6 +207,93 @@ func TestUniversalVerifyRateLimitFollowsUserAcrossIPs(t *testing.T) { require.Equal(t, http.StatusTooManyRequests, performVerify(fmt.Sprintf("[2001:db8:%x::2]:1234", testRun)).Code) } +func TestCriticalAccountRoutesRateLimitSameUserAcrossIPs(t *testing.T) { + for _, path := range []string{"/api/user/token", "/api/user/aff_transfer"} { + path := path + t.Run(path, func(t *testing.T) { + gin.SetMode(gin.TestMode) + testRun := atomic.AddUint64(&universalVerifyRateLimitTestRun, 1) + + oldDB := model.DB + oldLogDB := model.LOG_DB + oldRedisEnabled := common.RedisEnabled + oldMemoryCacheEnabled := common.MemoryCacheEnabled + oldGlobalAPIRateLimitEnable := common.GlobalApiRateLimitEnable + oldCriticalRateLimitEnable := common.CriticalRateLimitEnable + oldCriticalRateLimitNum := common.CriticalRateLimitNum + oldCriticalRateLimitDuration := common.CriticalRateLimitDuration + + common.RedisEnabled = false + common.MemoryCacheEnabled = false + common.GlobalApiRateLimitEnable = false + common.CriticalRateLimitEnable = true + common.CriticalRateLimitNum = 1 + common.CriticalRateLimitDuration = 20 * 60 + + dsn := fmt.Sprintf("file:critical_account_route_%d?mode=memory&cache=shared", testRun) + db, err := gorm.Open(sqlite.Open(dsn), &gorm.Config{}) + require.NoError(t, err) + model.DB = db + model.LOG_DB = db + require.NoError(t, db.AutoMigrate(&model.User{})) + user := model.User{ + Id: 1200000 + int(testRun), Username: fmt.Sprintf("critical-user-%d", testRun), + Password: "hashed-password", Role: common.RoleCommonUser, + Status: common.UserStatusEnabled, Group: "default", + } + require.NoError(t, db.Create(&user).Error) + + t.Cleanup(func() { + if sqlDB, dbErr := db.DB(); dbErr == nil { + _ = sqlDB.Close() + } + model.DB = oldDB + model.LOG_DB = oldLogDB + common.RedisEnabled = oldRedisEnabled + common.MemoryCacheEnabled = oldMemoryCacheEnabled + common.GlobalApiRateLimitEnable = oldGlobalAPIRateLimitEnable + common.CriticalRateLimitEnable = oldCriticalRateLimitEnable + common.CriticalRateLimitNum = oldCriticalRateLimitNum + common.CriticalRateLimitDuration = oldCriticalRateLimitDuration + }) + + engine := gin.New() + engine.Use(sessions.Sessions("session", cookie.NewStore([]byte("critical-account-session")))) + engine.GET("/test/login", func(c *gin.Context) { + session := sessions.Default(c) + session.Set("id", user.Id) + session.Set("username", user.Username) + session.Set("role", user.Role) + session.Set("status", user.Status) + session.Set("group", user.Group) + require.NoError(t, session.Save()) + c.Status(http.StatusNoContent) + }) + SetApiRouter(engine) + + loginRecorder := httptest.NewRecorder() + engine.ServeHTTP(loginRecorder, httptest.NewRequest(http.MethodGet, "/test/login", nil)) + require.Equal(t, http.StatusNoContent, loginRecorder.Code) + + perform := func(remoteAddr string) *httptest.ResponseRecorder { + recorder := httptest.NewRecorder() + request := httptest.NewRequest(http.MethodPost, path, strings.NewReader(`{"quota":1}`)) + request.RemoteAddr = remoteAddr + request.Header.Set("Content-Type", "application/json") + for _, sessionCookie := range loginRecorder.Result().Cookies() { + request.AddCookie(sessionCookie) + } + engine.ServeHTTP(recorder, request) + return recorder + } + + first := perform(fmt.Sprintf("[2001:db8:%x::1]:1234", testRun)) + require.NotEqual(t, http.StatusTooManyRequests, first.Code) + require.Equal(t, http.StatusTooManyRequests, perform(fmt.Sprintf("[2001:db8:%x::2]:1234", testRun)).Code) + }) + } +} + func TestSecureVerificationOpenAPIIncludesScopeAndRequestBody(t *testing.T) { documentBytes, err := os.ReadFile("../docs/openapi/api.json") require.NoError(t, err) diff --git a/service/auth_flow_cleanup.go b/service/auth_flow_cleanup.go new file mode 100644 index 00000000..d5fc42ce --- /dev/null +++ b/service/auth_flow_cleanup.go @@ -0,0 +1,51 @@ +package service + +import ( + "context" + "time" + + "github.com/MAX-API-Next/MAX-API/common" + "github.com/MAX-API-Next/MAX-API/model" +) + +const authFlowCleanupInterval = time.Hour + +func StartAuthFlowCleanup() { + StartAuthFlowCleanupWithContext(context.Background()) +} + +func StartAuthFlowCleanupWithContext(ctx context.Context) <-chan struct{} { + done := make(chan struct{}) + if !common.IsMasterNode { + close(done) + return done + } + go func() { + defer close(done) + runAuthFlowCleanup(ctx, authFlowCleanupInterval, cleanupAuthFlows) + }() + return done +} + +func runAuthFlowCleanup(ctx context.Context, interval time.Duration, cleanup func()) { + if ctx.Err() != nil { + return + } + cleanup() + ticker := time.NewTicker(interval) + defer ticker.Stop() + for { + select { + case <-ctx.Done(): + return + case <-ticker.C: + cleanup() + } + } +} + +func cleanupAuthFlows() { + if err := model.DeleteExpiredAuthFlows(time.Now()); err != nil { + common.SysError("failed to delete expired OAuth flows: " + err.Error()) + } +} diff --git a/service/auth_flow_cleanup_test.go b/service/auth_flow_cleanup_test.go new file mode 100644 index 00000000..64443b88 --- /dev/null +++ b/service/auth_flow_cleanup_test.go @@ -0,0 +1,30 @@ +package service + +import ( + "context" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +func TestRunAuthFlowCleanupStopsOnCancellation(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + done := make(chan struct{}) + var calls atomic.Int32 + go func() { + defer close(done) + runAuthFlowCleanup(ctx, time.Millisecond, func() { + calls.Add(1) + }) + }() + + require.Eventually(t, func() bool { return calls.Load() >= 2 }, time.Second, time.Millisecond) + cancel() + select { + case <-done: + case <-time.After(time.Second): + t.Fatal("auth flow cleanup did not stop after cancellation") + } +} diff --git a/service/billing_session.go b/service/billing_session.go index c0869fbb..27d10e49 100644 --- a/service/billing_session.go +++ b/service/billing_session.go @@ -48,6 +48,12 @@ type SettlementPreparer interface { var _ SettlementPreparer = (*BillingSession)(nil) +type RefundSettlementPreparer interface { + PrepareRefundSettlement() (*model.BillingSettlementInput, error) +} + +var _ RefundSettlementPreparer = (*BillingSession)(nil) + type billingSettleIntent struct { input model.BillingSettlementInput } @@ -306,43 +312,61 @@ func (s *BillingSession) prepareRefundIntent() (*billingRefundIntent, bool) { if s.settled || s.refunded || s.refundInFlight || s.settleInFlight || !s.needsRefundLocked() { return nil, false } + input, err := s.refundSettlementInputLocked() + if err != nil { + common.SysLog(err.Error()) + return nil, false + } + + s.refundInFlight = true + return &billingRefundIntent{ + input: *input, + userID: s.relayInfo.UserId, + requestID: s.relayInfo.RequestId, + tokenConsumed: s.tokenConsumed, + fundingSource: s.funding.Source(), + }, true +} + +func (s *BillingSession) refundSettlementInputLocked() (*model.BillingSettlementInput, error) { refundFunding := s.refundFundingAmountLocked() refundToken := int64(s.tokenConsumed) if s.relayInfo.IsPlayground { refundToken = 0 } requestID := s.relayInfo.RequestId - funding := s.funding if requestID == "" || refundFunding <= 0 { - common.SysLog(fmt.Sprintf("billing refund requires manual review (userId=%d, requestId=%q, funding=%d)", s.relayInfo.UserId, requestID, refundFunding)) - return nil, false + return nil, fmt.Errorf("billing refund requires manual review (userId=%d, requestId=%q, funding=%d)", s.relayInfo.UserId, requestID, refundFunding) } - source, userID, subscriptionID, ok := durableFundingIdentity(funding) + source, userID, subscriptionID, ok := durableFundingIdentity(s.funding) if !ok { - common.SysLog(fmt.Sprintf("billing refund skipped for non-durable funding source (userId=%d, requestId=%s)", s.relayInfo.UserId, requestID)) - return nil, false + return nil, fmt.Errorf("billing refund skipped for non-durable funding source (userId=%d, requestId=%s)", s.relayInfo.UserId, requestID) } + return &model.BillingSettlementInput{ + OperationKey: "request:" + requestID + ":finalize", + Source: source, + UserID: userID, + SubscriptionID: subscriptionID, + TokenID: s.relayInfo.TokenId, + TokenKey: s.relayInfo.TokenKey, + FundingDelta: -refundFunding, + TokenDelta: -refundToken, + SubscriptionPreConsumeRequestID: subscriptionPreConsumeRequestID(s.funding), + FinalizeSubscriptionPreConsume: source == model.BillingSettlementSourceSubscription, + AllowMissingToken: true, + }, nil +} - s.refundInFlight = true - return &billingRefundIntent{ - input: model.BillingSettlementInput{ - OperationKey: "request:" + requestID + ":finalize", - Source: source, - UserID: userID, - SubscriptionID: subscriptionID, - TokenID: s.relayInfo.TokenId, - TokenKey: s.relayInfo.TokenKey, - FundingDelta: -refundFunding, - TokenDelta: -refundToken, - SubscriptionPreConsumeRequestID: subscriptionPreConsumeRequestID(funding), - FinalizeSubscriptionPreConsume: source == model.BillingSettlementSourceSubscription, - AllowMissingToken: true, - }, - userID: s.relayInfo.UserId, - requestID: requestID, - tokenConsumed: s.tokenConsumed, - fundingSource: funding.Source(), - }, true +// PrepareRefundSettlement returns the exact durable refund intent without +// applying it. Async submission paths use it to commit a failed task state and +// the refund operation atomically before attempting the balance mutation. +func (s *BillingSession) PrepareRefundSettlement() (*model.BillingSettlementInput, error) { + s.mu.Lock() + defer s.mu.Unlock() + if s.settled || s.refunded || s.refundInFlight || s.settleInFlight || !s.needsRefundLocked() { + return nil, nil + } + return s.refundSettlementInputLocked() } func (s *BillingSession) finishRefundIntent(applied bool) { diff --git a/service/channel_affinity.go b/service/channel_affinity.go index 6e45be1b..2af02fd8 100644 --- a/service/channel_affinity.go +++ b/service/channel_affinity.go @@ -37,9 +37,11 @@ var ( channelAffinityUsageCacheStatsOnce sync.Once channelAffinityUsageCacheStatsCache *cachex.HybridCache[ChannelAffinityUsageCacheCounters] - channelAffinityRegexCache sync.Map // map[string]*regexp.Regexp + channelAffinityRegexCache = newChannelAffinityRegexCache(channelAffinityRegexCacheCapacity) ) +const channelAffinityRegexCacheCapacity = 4096 + type channelAffinityMeta struct { CacheKey string TTLSeconds int @@ -253,22 +255,40 @@ func matchAnyRegexCached(patterns []string, s string) bool { if pattern == "" { continue } - re, ok := channelAffinityRegexCache.Load(pattern) - if !ok { - compiled, err := regexp.Compile(pattern) - if err != nil { - continue - } - re = compiled - channelAffinityRegexCache.Store(pattern, re) + re, err := loadChannelAffinityRegex(channelAffinityRegexCache, pattern) + if err != nil { + continue } - if re.(*regexp.Regexp).MatchString(s) { + if re.MatchString(s) { return true } } return false } +func newChannelAffinityRegexCache(capacity int) *hot.HotCache[string, *regexp.Regexp] { + if capacity <= 0 { + capacity = 1 + } + return hot.NewHotCache[string, *regexp.Regexp](hot.LRU, capacity).Build() +} + +func loadChannelAffinityRegex(cache *hot.HotCache[string, *regexp.Regexp], pattern string) (*regexp.Regexp, error) { + compiled, ok, err := cache.Get(pattern) + if err != nil { + return nil, err + } + if ok { + return compiled, nil + } + compiled, err = regexp.Compile(pattern) + if err != nil { + return nil, err + } + cache.Set(pattern, compiled) + return compiled, nil +} + func matchAnyIncludeFold(patterns []string, s string) bool { if len(patterns) == 0 || s == "" { return false diff --git a/service/channel_affinity_template_test.go b/service/channel_affinity_template_test.go index ac42bebb..d25629b8 100644 --- a/service/channel_affinity_template_test.go +++ b/service/channel_affinity_template_test.go @@ -332,3 +332,25 @@ func TestChannelAffinityHitCodexTemplatePassHeadersEffective(t *testing.T) { _, exists = info.RuntimeHeadersOverride["x-codex-turn-metadata"] require.False(t, exists) } + +func TestChannelAffinityRegexCacheEvictsLeastRecentlyUsedPattern(t *testing.T) { + cache := newChannelAffinityRegexCache(2) + + _, err := loadChannelAffinityRegex(cache, "^first$") + require.NoError(t, err) + _, err = loadChannelAffinityRegex(cache, "^second$") + require.NoError(t, err) + _, found, err := cache.Get("^first$") + require.NoError(t, err) + require.True(t, found) + _, err = loadChannelAffinityRegex(cache, "^third$") + require.NoError(t, err) + + _, found, err = cache.Get("^second$") + require.NoError(t, err) + require.False(t, found) + _, found, err = cache.Get("^first$") + require.NoError(t, err) + require.True(t, found) + require.Len(t, cache.Keys(), 2) +} diff --git a/service/error.go b/service/error.go index de9e3e8a..6100cbdd 100644 --- a/service/error.go +++ b/service/error.go @@ -135,7 +135,77 @@ func RelayErrorHandler(ctx context.Context, resp *http.Response, showBodyWhenFai return } +func TaskErrorFromUpstreamResponse(ctx context.Context, resp *http.Response, code string) *dto.TaskError { + if resp == nil { + return TaskErrorWrapperLocal(errors.New("task upstream response is nil"), code, http.StatusBadGateway) + } + responseBody, truncated, err := readUpstreamErrorBody(resp.Body) + CloseResponseBodyGracefully(resp) + if err != nil { + return TaskErrorWrapper(err, code, resp.StatusCode) + } + + message := safeUpstreamTaskErrorMessage(responseBody) + if message == "" { + logger.LogError(ctx, fmt.Sprintf("task upstream bad response status code %d, body_bytes=%d, truncated=%t", resp.StatusCode, len(responseBody), truncated)) + message = fmt.Sprintf("bad upstream response status code %d", resp.StatusCode) + } else if truncated { + message += fmt.Sprintf("...[truncated after %d bytes]", maxUpstreamErrorBodyBytes) + } + message = common.SanitizePersistedLogContent(common.MaskSensitiveInfo(message)) + if message == "" { + message = fmt.Sprintf("bad upstream response status code %d", resp.StatusCode) + } + return &dto.TaskError{ + Code: code, + Message: message, + StatusCode: resp.StatusCode, + Error: errors.New(message), + } +} + +func safeUpstreamTaskErrorMessage(responseBody []byte) string { + var errResponse dto.GeneralErrorResponse + if err := common.Unmarshal(responseBody, &errResponse); err != nil { + return "" + } + if openAIError := errResponse.TryToOpenAIError(); openAIError != nil { + return openAIError.Message + } + if errResponse.Message != "" { + return errResponse.Message + } + if errResponse.Msg != "" { + return errResponse.Msg + } + if errResponse.Err != "" { + return errResponse.Err + } + if errResponse.ErrorMsg != "" { + return errResponse.ErrorMsg + } + if errResponse.Detail != "" { + return errResponse.Detail + } + if errResponse.Header.Message != "" { + return errResponse.Header.Message + } + if errResponse.Response.Error.Message != "" { + return errResponse.Response.Error.Message + } + if len(errResponse.Error) > 0 && common.GetJsonType(errResponse.Error) == "string" { + var message string + if err := common.Unmarshal(errResponse.Error, &message); err == nil { + return message + } + } + return "" +} + func readUpstreamErrorBody(body io.Reader) ([]byte, bool, error) { + if body == nil { + return nil, false, nil + } responseBody, err := io.ReadAll(io.LimitReader(body, maxUpstreamErrorBodyBytes+1)) if err != nil { return nil, false, err diff --git a/service/error_test.go b/service/error_test.go index dc3467cd..93214d6f 100644 --- a/service/error_test.go +++ b/service/error_test.go @@ -197,6 +197,47 @@ func TestRelayErrorHandlerLimitsUpstreamErrorBodyRead(t *testing.T) { require.Contains(t, maxAPIError.Error(), "truncated after") } +func TestTaskErrorFromUpstreamResponseKeepsStructuredSafeMessage(t *testing.T) { + resp := &http.Response{ + StatusCode: http.StatusBadGateway, + Body: io.NopCloser(strings.NewReader(`{"error":{"message":"upstream failed\r\ntry later","type":"server_error","code":"server_error"}}`)), + } + + taskErr := TaskErrorFromUpstreamResponse(context.Background(), resp, "fail_to_fetch_task") + + require.NotNil(t, taskErr) + require.Equal(t, "fail_to_fetch_task", taskErr.Code) + require.Equal(t, http.StatusBadGateway, taskErr.StatusCode) + require.Equal(t, "upstream failed try later", taskErr.Message) + require.NotContains(t, taskErr.Message, "\r") + require.NotContains(t, taskErr.Message, "\n") +} + +func TestTaskErrorFromUpstreamResponseDoesNotExposeRawInvalidBody(t *testing.T) { + rawBody := `{"internal_prompt":"` + strings.Repeat("secret", 100) + `"}` + resp := &http.Response{ + StatusCode: http.StatusInternalServerError, + Body: io.NopCloser(strings.NewReader(rawBody)), + } + + taskErr := TaskErrorFromUpstreamResponse(context.Background(), resp, "fail_to_fetch_task") + + require.NotNil(t, taskErr) + require.Equal(t, "bad upstream response status code 500", taskErr.Message) + require.NotContains(t, taskErr.Message, "secret") +} + +func TestTaskErrorFromUpstreamResponseHandlesNilBody(t *testing.T) { + resp := &http.Response{ + StatusCode: http.StatusBadGateway, + } + + taskErr := TaskErrorFromUpstreamResponse(context.Background(), resp, "fail_to_fetch_task") + + require.NotNil(t, taskErr) + require.Equal(t, "bad upstream response status code 502", taskErr.Message) +} + type countingReadCloser struct { io.Reader read int diff --git a/service/http.go b/service/http.go index c3022b7e..116b95fb 100644 --- a/service/http.go +++ b/service/http.go @@ -13,6 +13,19 @@ import ( "github.com/gin-gonic/gin" ) +var blockedUpstreamResponseHeaders = map[string]struct{}{ + "connection": {}, + "content-length": {}, + "keep-alive": {}, + "proxy-authenticate": {}, + "proxy-authorization": {}, + "set-cookie": {}, + "te": {}, + "trailer": {}, + "transfer-encoding": {}, + "upgrade": {}, +} + func CloseResponseBodyGracefully(httpResponse *http.Response) { if httpResponse == nil || httpResponse.Body == nil { return @@ -23,21 +36,122 @@ func CloseResponseBodyGracefully(httpResponse *http.Response) { } } -// ShouldCopyUpstreamHeader checks whether a given upstream response header -// should be copied to the client response. It returns false for Content-Length -// (managed separately) and X-Oneapi-Request-Id (to preserve the local instance -// ID). When the upstream header is X-Oneapi-Request-Id, the value is captured -// into the Gin context for later logging. -func ShouldCopyUpstreamHeader(c *gin.Context, k string, v []string) bool { - if strings.EqualFold(k, "Content-Length") { +// ShouldCopyUpstreamHeader checks whether a given upstream response header can +// be copied to the client response. Hop-by-hop headers and Set-Cookie stay +// under MAX API control, and malformed names/values are dropped before they +// reach net/http's writer. +func ShouldCopyUpstreamHeader(_ *gin.Context, k string, v []string) bool { + return shouldCopyUpstreamHeader(k, v) +} + +func shouldCopyUpstreamHeader(k string, v []string) bool { + if len(v) == 0 { + return false + } + headerName := strings.TrimSpace(k) + if headerName == "" || headerName != k || !isSafeResponseHeaderName(headerName) || !isSafeResponseHeaderValue(v[0]) { return false } - if strings.EqualFold(k, common.RequestIdKey) { - if c != nil && len(v) > 0 { - c.Set(common.UpstreamRequestIdKey, v[0]) + if strings.EqualFold(headerName, common.RequestIdKey) { + return false + } + if _, blocked := blockedUpstreamResponseHeaders[strings.ToLower(headerName)]; blocked { + return false + } + return true +} + +func connectionScopedHeaderNames(header http.Header) map[string]struct{} { + var connectionScoped map[string]struct{} + for key, values := range header { + if !strings.EqualFold(key, "Connection") { + continue + } + for _, value := range values { + for _, name := range strings.Split(value, ",") { + name = strings.TrimSpace(name) + if name == "" || !isSafeResponseHeaderName(name) { + continue + } + if connectionScoped == nil { + connectionScoped = make(map[string]struct{}) + } + connectionScoped[strings.ToLower(name)] = struct{}{} + } + } + } + return connectionScoped +} + +func CopyUpstreamResponseHeaders(c *gin.Context, header http.Header) { + if c == nil || c.Writer == nil { + return + } + connectionScoped := connectionScopedHeaderNames(header) + for key, values := range header { + if connectionScoped != nil { + if _, blocked := connectionScoped[strings.ToLower(key)]; blocked { + continue + } + } + if len(values) > 0 && strings.EqualFold(key, common.RequestIdKey) && isSafeResponseHeaderName(key) && isSafeResponseHeaderValue(values[0]) { + c.Set(common.UpstreamRequestIdKey, values[0]) + } + if !shouldCopyUpstreamHeader(key, values) { + continue + } + copied := false + for _, value := range values { + if !isSafeResponseHeaderValue(value) { + continue + } + if copied { + c.Writer.Header().Add(key, value) + } else { + c.Writer.Header().Set(key, value) + copied = true + } + } + } +} + +func isSafeResponseHeaderName(name string) bool { + for i := 0; i < len(name); i++ { + if !isHTTPTokenChar(name[i]) { + return false } + } + return true +} + +func isHTTPTokenChar(c byte) bool { + if c >= 'a' && c <= 'z' { + return true + } + if c >= 'A' && c <= 'Z' { + return true + } + if c >= '0' && c <= '9' { + return true + } + switch c { + case '!', '#', '$', '%', '&', '\'', '*', '+', '-', '.', '^', '_', '`', '|', '~': + return true + default: return false } +} + +func isSafeResponseHeaderValue(value string) bool { + for i := 0; i < len(value); i++ { + switch value[i] { + case '\r', '\n', 0, 0x7f: + return false + } + if value[i] < 0x20 && value[i] != '\t' { + return false + } + } return true } @@ -53,12 +167,7 @@ func IOCopyBytesGracefully(c *gin.Context, src *http.Response, data []byte) { // So the httpClient will be confused by the response. // For example, Postman will report error, and we cannot check the response at all. if src != nil { - for k, v := range src.Header { - if !ShouldCopyUpstreamHeader(c, k, v) { - continue - } - c.Writer.Header().Set(k, v[0]) - } + CopyUpstreamResponseHeaders(c, src.Header) } // set Content-Length header manually BEFORE calling WriteHeader diff --git a/service/http_client_test.go b/service/http_client_test.go index 0bd24b27..59c7a026 100644 --- a/service/http_client_test.go +++ b/service/http_client_test.go @@ -2,7 +2,13 @@ package service import ( "net/http" + "net/http/httptest" + "strings" "testing" + + "github.com/MAX-API-Next/MAX-API/common" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" ) func TestNewSSRFProtectedHTTPClientDisablesEnvironmentProxy(t *testing.T) { @@ -25,3 +31,75 @@ func TestNewBaseTransportKeepsEnvironmentProxySupport(t *testing.T) { t.Fatalf("expected base transport to keep environment proxy support") } } + +func TestShouldCopyUpstreamHeaderRejectsUnsafeResponseHeaders(t *testing.T) { + gin.SetMode(gin.TestMode) + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + + require.False(t, ShouldCopyUpstreamHeader(c, "Content-Length", []string{"10"})) + require.False(t, ShouldCopyUpstreamHeader(c, "Connection", []string{"keep-alive"})) + require.False(t, ShouldCopyUpstreamHeader(c, "Set-Cookie", []string{"session=upstream"})) + require.False(t, ShouldCopyUpstreamHeader(c, "Bad Header", []string{"value"})) + require.False(t, ShouldCopyUpstreamHeader(c, "X-Bad", []string{"ok\r\nX-Injected: yes"})) + require.False(t, ShouldCopyUpstreamHeader(c, "X-Del", []string{"bad" + string(rune(0x7f))})) + require.False(t, ShouldCopyUpstreamHeader(c, "X-Empty", nil)) + require.False(t, ShouldCopyUpstreamHeader(c, common.RequestIdKey, []string{"upstream-request"})) + _, capturedRequestID := c.Get(common.UpstreamRequestIdKey) + require.False(t, capturedRequestID) + require.True(t, ShouldCopyUpstreamHeader(c, "X-Safe", []string{"safe"})) +} + +func TestCopyUpstreamResponseHeadersRejectsConnectionNominatedHeaders(t *testing.T) { + gin.SetMode(gin.TestMode) + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + + CopyUpstreamResponseHeaders(c, http.Header{ + "Connection": []string{"keep-alive, X-Transient", strings.ToLower(common.RequestIdKey)}, + "X-Transient": []string{"must-not-copy"}, + "X-Safe": []string{"safe"}, + common.RequestIdKey: []string{"must-not-capture"}, + }) + + require.Empty(t, recorder.Header().Get("Connection")) + require.Empty(t, recorder.Header().Get("X-Transient")) + require.Equal(t, "safe", recorder.Header().Get("X-Safe")) + _, capturedRequestID := c.Get(common.UpstreamRequestIdKey) + require.False(t, capturedRequestID) +} + +func TestIOCopyBytesGracefullyCopiesOnlySafeUpstreamHeaders(t *testing.T) { + gin.SetMode(gin.TestMode) + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + src := &http.Response{ + StatusCode: http.StatusAccepted, + Header: http.Header{ + "Connection": []string{"keep-alive"}, + "Set-Cookie": []string{"session=upstream"}, + "X-Bad": []string{"ok\r\nX-Injected: yes"}, + "X-Multi": []string{"one", "two\r\nX-Injected: yes", "three"}, + "X-Safe": []string{"safe"}, + common.RequestIdKey: []string{"upstream-request"}, + "Content-Type": []string{"application/json"}, + "Transfer-Encoding": []string{"chunked"}, + "Content-Length": []string{"999"}, + "X-Empty-Value-List": nil, + }, + } + + IOCopyBytesGracefully(c, src, []byte(`{"ok":true}`)) + + require.Equal(t, http.StatusAccepted, recorder.Code) + require.Equal(t, "safe", recorder.Header().Get("X-Safe")) + require.Equal(t, []string{"one", "three"}, recorder.Header().Values("X-Multi")) + require.Equal(t, "application/json", recorder.Header().Get("Content-Type")) + require.Equal(t, "11", recorder.Header().Get("Content-Length")) + require.Empty(t, recorder.Header().Get("Connection")) + require.Empty(t, recorder.Header().Get("Set-Cookie")) + require.Empty(t, recorder.Header().Get("X-Bad")) + require.Empty(t, recorder.Header().Get(common.RequestIdKey)) + require.Empty(t, recorder.Header().Get("Transfer-Encoding")) + require.Equal(t, "upstream-request", c.GetString(common.UpstreamRequestIdKey)) + require.JSONEq(t, `{"ok":true}`, recorder.Body.String()) +} diff --git a/service/midjourney.go b/service/midjourney.go index 9cb1fedd..79e5e902 100644 --- a/service/midjourney.go +++ b/service/midjourney.go @@ -2,11 +2,13 @@ package service import ( "context" - "encoding/json" + "errors" "io" "net/http" + "net/http/httptrace" "strconv" "strings" + "sync/atomic" "time" "github.com/MAX-API-Next/MAX-API/common" @@ -162,7 +164,7 @@ func ConvertSimpleChangeParams(content string) *dto.MidjourneyRequest { return changeParams } -func DoMidjourneyHttpRequest(c *gin.Context, timeout time.Duration, fullRequestURL string) (*dto.MidjourneyResponseWithStatusCode, []byte, error) { +func DoMidjourneyHttpRequest(c *gin.Context, timeout time.Duration, fullRequestURL string) (*dto.MidjourneyResponseWithStatusCode, []byte, bool, error) { var nullBytes []byte //var requestBody io.Reader //requestBody = c.Request.Body @@ -170,9 +172,9 @@ func DoMidjourneyHttpRequest(c *gin.Context, timeout time.Duration, fullRequestU var mapResult map[string]interface{} // if get request, no need to read request body if c.Request.Method != "GET" { - err := json.NewDecoder(c.Request.Body).Decode(&mapResult) + err := common.DecodeJson(c.Request.Body, &mapResult) if err != nil { - return MidjourneyErrorWithStatusCodeWrapper(constant.MjErrorUnknown, "read_request_body_failed", http.StatusInternalServerError), nullBytes, err + return MidjourneyErrorWithStatusCodeWrapper(constant.MjErrorUnknown, "read_request_body_failed", http.StatusInternalServerError), nullBytes, false, err } if !setting.MjAccountFilterEnabled { delete(mapResult, "accountFilter") @@ -192,14 +194,15 @@ func DoMidjourneyHttpRequest(c *gin.Context, timeout time.Duration, fullRequestU mapResult["prompt"] = prompt } } - reqBody, err := json.Marshal(mapResult) + reqBody, err := common.Marshal(mapResult) if err != nil { - return MidjourneyErrorWithStatusCodeWrapper(constant.MjErrorUnknown, "marshal_request_body_failed", http.StatusInternalServerError), nullBytes, err + return MidjourneyErrorWithStatusCodeWrapper(constant.MjErrorUnknown, "marshal_request_body_failed", http.StatusInternalServerError), nullBytes, false, err } req, err := http.NewRequest(c.Request.Method, fullRequestURL, strings.NewReader(string(reqBody))) if err != nil { - return MidjourneyErrorWithStatusCodeWrapper(constant.MjErrorUnknown, "create_request_failed", http.StatusInternalServerError), nullBytes, err + return MidjourneyErrorWithStatusCodeWrapper(constant.MjErrorUnknown, "create_request_failed", http.StatusInternalServerError), nullBytes, false, err } + defer req.Body.Close() ctx, cancel := context.WithTimeout(context.Background(), timeout) // 使用带有超时的 context 创建新的请求 req = req.WithContext(ctx) @@ -211,39 +214,40 @@ func DoMidjourneyHttpRequest(c *gin.Context, timeout time.Duration, fullRequestU req.Header.Set("mj-api-secret", auth) } defer cancel() + var requestSent atomic.Bool + req = req.WithContext(httptrace.WithClientTrace(req.Context(), &httptrace.ClientTrace{ + WroteRequest: func(httptrace.WroteRequestInfo) { + // Any write attempt makes the provider outcome ambiguous. A partial + // write must not be treated as a definite rejection and auto-refunded. + requestSent.Store(true) + }, + })) resp, err := GetHttpClient().Do(req) if err != nil { common.SysLog("do request failed: " + err.Error()) - return MidjourneyErrorWithStatusCodeWrapper(constant.MjErrorUnknown, "do_request_failed", http.StatusInternalServerError), nullBytes, err + return MidjourneyErrorWithStatusCodeWrapper(constant.MjErrorUnknown, "do_request_failed", http.StatusInternalServerError), nullBytes, requestSent.Load(), err } + requestSent.Store(true) + defer CloseResponseBodyGracefully(resp) statusCode := resp.StatusCode //if statusCode != 200 { // return MidjourneyErrorWithStatusCodeWrapper(constant.MjErrorUnknown, "bad_response_status_code", statusCode), nullBytes, nil //} - err = req.Body.Close() - if err != nil { - return MidjourneyErrorWithStatusCodeWrapper(constant.MjErrorUnknown, "close_request_body_failed", statusCode), nullBytes, err - } - err = c.Request.Body.Close() - if err != nil { - return MidjourneyErrorWithStatusCodeWrapper(constant.MjErrorUnknown, "close_request_body_failed", statusCode), nullBytes, err - } var midjResponse dto.MidjourneyResponse var midjourneyUploadsResponse dto.MidjourneyUploadResponse responseBody, err := io.ReadAll(resp.Body) if err != nil { - return MidjourneyErrorWithStatusCodeWrapper(constant.MjErrorUnknown, "read_response_body_failed", statusCode), nullBytes, err + return MidjourneyErrorWithStatusCodeWrapper(constant.MjErrorUnknown, "read_response_body_failed", statusCode), nullBytes, requestSent.Load(), err } - CloseResponseBodyGracefully(resp) logger.LogDebug(c, "midjourney response body: %s", responseBody) if len(responseBody) == 0 { - return MidjourneyErrorWithStatusCodeWrapper(constant.MjErrorUnknown, "empty_response_body", statusCode), responseBody, nil + return MidjourneyErrorWithStatusCodeWrapper(constant.MjErrorUnknown, "empty_response_body", statusCode), responseBody, requestSent.Load(), errors.New("midjourney upstream response body is empty") } else { - err = json.Unmarshal(responseBody, &midjResponse) + err = common.Unmarshal(responseBody, &midjResponse) if err != nil { - err2 := json.Unmarshal(responseBody, &midjourneyUploadsResponse) + err2 := common.Unmarshal(responseBody, &midjourneyUploadsResponse) if err2 != nil { - return MidjourneyErrorWithStatusCodeWrapper(constant.MjErrorUnknown, "unmarshal_response_body_failed", statusCode), responseBody, err + return MidjourneyErrorWithStatusCodeWrapper(constant.MjErrorUnknown, "unmarshal_response_body_failed", statusCode), responseBody, requestSent.Load(), err } } } @@ -253,5 +257,5 @@ func DoMidjourneyHttpRequest(c *gin.Context, timeout time.Duration, fullRequestU return &dto.MidjourneyResponseWithStatusCode{ StatusCode: statusCode, Response: midjResponse, - }, responseBody, nil + }, responseBody, requestSent.Load(), nil } diff --git a/service/midjourney_task.go b/service/midjourney_task.go new file mode 100644 index 00000000..e41ffe5b --- /dev/null +++ b/service/midjourney_task.go @@ -0,0 +1,216 @@ +package service + +import ( + "context" + "fmt" + + "github.com/MAX-API-Next/MAX-API/common" + "github.com/MAX-API-Next/MAX-API/dto" + "github.com/MAX-API-Next/MAX-API/logger" + "github.com/MAX-API-Next/MAX-API/model" +) + +// UpdateMidjourneyTaskFromResponse advances the legacy Midjourney row and its +// durable billing Task through one CAS-protected transaction. +func UpdateMidjourneyTaskFromResponse(ctx context.Context, task *model.Midjourney, responseItem dto.MidjourneyDto) error { + legacyNeedsUpdate := midjourneyTaskNeedsUpdate(task, responseItem) + preStatus := task.Status + task.Code = 1 + if responseItem.Progress != "" { + task.Progress = responseItem.Progress + } + if responseItem.PromptEn != "" { + task.PromptEn = responseItem.PromptEn + } + if responseItem.State != "" { + task.State = responseItem.State + } + if responseItem.SubmitTime > 0 { + task.SubmitTime = responseItem.SubmitTime + } + if responseItem.StartTime > 0 { + task.StartTime = responseItem.StartTime + } + if responseItem.FinishTime > 0 { + task.FinishTime = responseItem.FinishTime + } + if responseItem.ImageUrl != "" { + task.ImageUrl = responseItem.ImageUrl + } + if responseItem.VideoUrl != "" { + task.VideoUrl = responseItem.VideoUrl + } + if responseItem.Status != "" { + task.Status = responseItem.Status + } + if responseItem.FailReason != "" { + task.FailReason = common.SanitizePersistedLogContent(responseItem.FailReason) + } + if responseItem.Properties != nil { + properties, err := common.Marshal(responseItem.Properties) + if err != nil { + return err + } + task.Properties = string(properties) + } + if responseItem.Buttons != nil { + buttons, err := common.Marshal(responseItem.Buttons) + if err != nil { + return err + } + task.Buttons = string(buttons) + } + if responseItem.VideoUrls != nil { + videoURLs, err := common.Marshal(responseItem.VideoUrls) + if err != nil { + return err + } + task.VideoUrls = string(videoURLs) + } + + providerSucceeded := responseItem.Status == "SUCCESS" || + (responseItem.Status == "" && responseItem.Progress == "100%" && responseItem.FailReason == "") + if providerSucceeded { + legacyNeedsUpdate = legacyNeedsUpdate || preStatus != "SUCCESS" || task.FailReason != "" + task.Status = "SUCCESS" + task.FailReason = "" + } + isFailure := !providerSucceeded && (task.Status == "FAILURE" || task.FailReason != "") + if isFailure { + task.Status = "FAILURE" + task.Progress = "100%" + logger.LogInfo(ctx, task.MjId+" 构建失败，"+task.FailReason) + } + billingTask, err := model.GetMidjourneyBillingTask(task.UserId, task.ChannelId, task.MjId) + if err != nil { + return err + } + var settlement *model.BillingSettlementInput + if billingTask != nil { + billingFromStatus := billingTask.Status + billingSnapshotBefore := billingTask.Snapshot() + billingTask.Status = midjourneyBillingTaskStatus(task.Status, task.Progress) + billingTask.Progress = task.Progress + billingTask.StartTime = task.StartTime / 1000 + billingTask.FinishTime = task.FinishTime / 1000 + billingTask.FailReason = task.FailReason + if task.VideoUrl != "" { + billingTask.PrivateData.ResultURL = task.VideoUrl + } else if task.ImageUrl != "" { + billingTask.PrivateData.ResultURL = task.ImageUrl + } + billingTask.SetData(responseItem) + if isFailure { + settlement = BuildTaskRefundSettlementInput(billingTask, "Midjourney 任务失败："+task.FailReason) + } + if !legacyNeedsUpdate { + if billingSnapshotBefore.Equal(billingTask.Snapshot()) { + return nil + } + if (billingFromStatus == model.TaskStatusFailure || billingFromStatus == model.TaskStatusSuccess) && billingFromStatus != billingTask.Status { + return fmt.Errorf("midjourney billing task terminal state conflicts with provider state: task=%d status=%s provider=%s", billingTask.ID, billingFromStatus, task.Status) + } + var won bool + if settlement != nil { + won, err = billingTask.UpdateWithStatusAndSettlement(billingFromStatus, *settlement) + } else { + won, err = billingTask.UpdateWithStatus(billingFromStatus) + } + if err != nil || !won { + return err + } + if settlement != nil { + ApplyTaskBillingSettlement(ctx, billingTask, settlement) + } + return nil + } + } else if isFailure && task.Quota != 0 { + logger.LogError(ctx, fmt.Sprintf("Midjourney 历史任务缺少账务身份，未自动退款，请人工对账: user=%d task=%s quota=%d", task.UserId, task.MjId, task.Quota)) + } + if !legacyNeedsUpdate { + return nil + } + + won, err := task.UpdateWithBillingTaskAndSettlement(preStatus, billingTask, settlement) + if err != nil || !won { + return err + } + if settlement != nil { + ApplyTaskBillingSettlement(ctx, billingTask, settlement) + } + return nil +} + +func midjourneyBillingTaskStatus(status, progress string) model.TaskStatus { + switch status { + case "SUCCESS": + return model.TaskStatusSuccess + case "FAILURE": + return model.TaskStatusFailure + case "SUBMITTED", "QUEUED": + return model.TaskStatusQueued + default: + if progress == "0%" || progress == "" { + return model.TaskStatusSubmitted + } + return model.TaskStatusInProgress + } +} + +func midjourneyTaskNeedsUpdate(oldTask *model.Midjourney, newTask dto.MidjourneyDto) bool { + if oldTask.Code != 1 { + return true + } + if newTask.Progress != "" && oldTask.Progress != newTask.Progress { + return true + } + if newTask.PromptEn != "" && oldTask.PromptEn != newTask.PromptEn { + return true + } + if newTask.State != "" && oldTask.State != newTask.State { + return true + } + if newTask.SubmitTime > 0 && oldTask.SubmitTime != newTask.SubmitTime { + return true + } + if newTask.StartTime > 0 && oldTask.StartTime != newTask.StartTime { + return true + } + if newTask.FinishTime > 0 && oldTask.FinishTime != newTask.FinishTime { + return true + } + if newTask.ImageUrl != "" && oldTask.ImageUrl != newTask.ImageUrl { + return true + } + if newTask.Status != "" && oldTask.Status != newTask.Status { + return true + } + if newTask.FailReason != "" && oldTask.FailReason != newTask.FailReason { + return true + } + if oldTask.Progress != "100%" && newTask.FailReason != "" { + return true + } + if newTask.VideoUrl != "" && oldTask.VideoUrl != newTask.VideoUrl { + return true + } + if newTask.Properties != nil { + properties, _ := common.Marshal(newTask.Properties) + if oldTask.Properties != string(properties) { + return true + } + } + if newTask.Buttons != nil { + buttons, _ := common.Marshal(newTask.Buttons) + if oldTask.Buttons != string(buttons) { + return true + } + } + if newTask.VideoUrls != nil && len(newTask.VideoUrls) > 0 { + videoURLs, _ := common.Marshal(newTask.VideoUrls) + if oldTask.VideoUrls != string(videoURLs) { + return true + } + } + return false +} diff --git a/service/midjourney_test.go b/service/midjourney_test.go new file mode 100644 index 00000000..d7e3f706 --- /dev/null +++ b/service/midjourney_test.go @@ -0,0 +1,89 @@ +package service + +import ( + "bytes" + "errors" + "io" + "net/http" + "net/http/httptest" + "net/http/httptrace" + "testing" + "time" + + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +type midjourneyRoundTripFunc func(*http.Request) (*http.Response, error) + +func (f midjourneyRoundTripFunc) RoundTrip(req *http.Request) (*http.Response, error) { + return f(req) +} + +func TestDoMidjourneyHttpRequestTracksWhetherTransportWroteRequest(t *testing.T) { + gin.SetMode(gin.TestMode) + originalClient := httpClient + t.Cleanup(func() { httpClient = originalClient }) + + tests := []struct { + name string + wroteErr error + invokeTrace bool + wantSent bool + }{ + {name: "transport failure before write", wantSent: false}, + {name: "request fully written before response failure", invokeTrace: true, wantSent: true}, + {name: "partial write remains ambiguous", invokeTrace: true, wroteErr: errors.New("partial write"), wantSent: true}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + httpClient = &http.Client{Transport: midjourneyRoundTripFunc(func(req *http.Request) (*http.Response, error) { + if tt.invokeTrace { + trace := httptrace.ContextClientTrace(req.Context()) + require.NotNil(t, trace) + require.NotNil(t, trace.WroteRequest) + trace.WroteRequest(httptrace.WroteRequestInfo{Err: tt.wroteErr}) + } + return nil, errors.New("upstream transport failed") + })} + + recorder := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(recorder) + ctx.Request = httptest.NewRequest(http.MethodPost, "/mj/submit", bytes.NewBufferString(`{"prompt":"test"}`)) + ctx.Request.Header.Set("Content-Type", "application/json") + + _, _, requestSent, err := DoMidjourneyHttpRequest(ctx, time.Second, "https://midjourney.invalid/submit") + require.Error(t, err) + require.Equal(t, tt.wantSent, requestSent) + }) + } +} + +func TestDoMidjourneyHttpRequestTreatsReceivedResponseAsSent(t *testing.T) { + gin.SetMode(gin.TestMode) + originalClient := httpClient + httpClient = &http.Client{Transport: midjourneyRoundTripFunc(func(*http.Request) (*http.Response, error) { + return &http.Response{ + StatusCode: http.StatusOK, + Header: make(http.Header), + Body: io.NopCloser(midjourneyFailingReader{}), + }, nil + })} + t.Cleanup(func() { httpClient = originalClient }) + + recorder := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(recorder) + ctx.Request = httptest.NewRequest(http.MethodPost, "/mj/submit", bytes.NewBufferString(`{"prompt":"test"}`)) + ctx.Request.Header.Set("Content-Type", "application/json") + + _, _, requestSent, err := DoMidjourneyHttpRequest(ctx, time.Second, "https://midjourney.invalid/submit") + require.Error(t, err) + require.True(t, requestSent) +} + +type midjourneyFailingReader struct{} + +func (midjourneyFailingReader) Read([]byte) (int, error) { + return 0, errors.New("response body failed") +} diff --git a/service/system_task.go b/service/system_task.go index e1482a7b..7de3509d 100644 --- a/service/system_task.go +++ b/service/system_task.go @@ -4,6 +4,7 @@ import ( "context" "errors" "fmt" + "os" "sync" "time" @@ -97,7 +98,12 @@ func StartSystemTaskRunner() { return } - runnerID := fmt.Sprintf("%s-%s", common.NodeName, common.GetRandomString(8)) + randomPart, err := common.GenerateRandomCharsKey(8) + if err != nil { + logger.LogError(context.Background(), "failed to generate system task runner ID: "+err.Error()) + randomPart = fmt.Sprintf("%d-%d", os.Getpid(), time.Now().UnixNano()) + } + runnerID := fmt.Sprintf("%s-%s", common.NodeName, randomPart) gopool.Go(func() { logger.LogInfo(context.Background(), fmt.Sprintf("system task runner started: runner=%s idle_interval=%s", runnerID, systemTaskRunnerIdleInterval)) diff --git a/service/task_billing.go b/service/task_billing.go index 1f739fce..eb7b841f 100644 --- a/service/task_billing.go +++ b/service/task_billing.go @@ -145,10 +145,11 @@ func taskModelName(task *model.Task) string { return task.Properties.OriginModelName } -func buildTaskRefundSettlementInput(task *model.Task, reason string) *model.BillingSettlementInput { +func BuildTaskRefundSettlementInput(task *model.Task, reason string) *model.BillingSettlementInput { if task == nil || task.Quota == 0 || task.ID <= 0 { return nil } + reason = common.SanitizePersistedLogContent(reason) source := model.BillingSettlementSourceWallet if taskIsSubscription(task) { source = model.BillingSettlementSourceSubscription @@ -182,6 +183,7 @@ func buildTaskFinalSettlementInput(task *model.Task, actualQuota int, reason str if task == nil || task.ID <= 0 || actualQuota <= 0 || actualQuota == task.Quota { return nil } + reason = common.SanitizePersistedLogContent(reason) quotaDelta := actualQuota - task.Quota source := model.BillingSettlementSourceWallet if taskIsSubscription(task) { @@ -257,7 +259,11 @@ func RefundTaskQuota(ctx context.Context, task *model.Task, reason string) bool return false } - return applyTaskBillingSettlement(ctx, task, buildTaskRefundSettlementInput(task, reason)) + return applyTaskBillingSettlement(ctx, task, BuildTaskRefundSettlementInput(task, reason)) +} + +func ApplyTaskBillingSettlement(ctx context.Context, task *model.Task, input *model.BillingSettlementInput) bool { + return applyTaskBillingSettlement(ctx, task, input) } // RecalculateTaskQuota 通用的异步差额结算。 diff --git a/service/task_billing_test.go b/service/task_billing_test.go index b8c27528..7b5be2a9 100644 --- a/service/task_billing_test.go +++ b/service/task_billing_test.go @@ -563,6 +563,38 @@ func TestProcessSunoTaskResponsePersistsRefundIntentBeforeApplying(t *testing.T) assert.Equal(t, model.BillingSettlementStatusApplied, settlement.Status) } +func TestProcessSunoTaskResponseRefundsWhenFailureReasonSanitizesEmpty(t *testing.T) { + truncate(t) + ctx := context.Background() + + const userID, tokenID, channelID = 7, 7, 7 + const pendingRefund, userQuota, tokenRemain = 250, 1000, 2000 + seedUser(t, userID, userQuota) + seedToken(t, tokenID, userID, "sk-suno-empty-sanitized-reason", tokenRemain) + seedChannel(t, channelID) + task := makeTask(userID, channelID, pendingRefund, tokenID, BillingSourceWallet, 0) + persistTask(t, task) + + processSunoTaskResponse(ctx, task, dto.SunoDataResponse{ + TaskID: task.TaskID, + Status: string(model.TaskStatusInProgress), + FailReason: "\x00\t", + Data: task.Data, + }) + + assert.EqualValues(t, userQuota+pendingRefund, getUserQuota(t, userID)) + assert.Equal(t, tokenRemain+pendingRefund, getTokenRemainQuota(t, tokenID)) + var reloaded model.Task + require.NoError(t, model.DB.First(&reloaded, task.ID).Error) + assert.EqualValues(t, model.TaskStatusFailure, reloaded.Status) + assert.Equal(t, "100%", reloaded.Progress) + assert.Empty(t, reloaded.FailReason) + assert.Zero(t, reloaded.Quota) + var settlement model.BillingSettlement + require.NoError(t, model.DB.Where("operation_key = ?", fmt.Sprintf("task:%d:refund", task.ID)).First(&settlement).Error) + assert.Equal(t, model.BillingSettlementStatusApplied, settlement.Status) +} + // =========================================================================== // RecalculateTaskQuota tests // =========================================================================== diff --git a/service/task_polling.go b/service/task_polling.go index 8981e158..603157c8 100644 --- a/service/task_polling.go +++ b/service/task_polling.go @@ -73,7 +73,7 @@ func sweepTimedOutTasks(ctx context.Context) { var settlement *model.BillingSettlementInput if !isLegacy && !isUnconfirmedSubmit { - settlement = buildTaskRefundSettlementInput(task, reason) + settlement = BuildTaskRefundSettlementInput(task, reason) } var won bool var err error @@ -172,22 +172,7 @@ func updateSunoTasks(ctx context.Context, channelId int, taskIds []string, taskM ch, err := model.CacheGetChannel(channelId) if err != nil { common.SysLog(fmt.Sprintf("CacheGetChannel: %v", err)) - // Collect DB primary key IDs for bulk update (taskIds are upstream IDs, not task_id column values) - var failedIDs []int64 - for _, upstreamID := range taskIds { - if t, ok := taskM[upstreamID]; ok { - failedIDs = append(failedIDs, t.ID) - } - } - err = model.TaskBulkUpdateByID(failedIDs, map[string]any{ - "fail_reason": fmt.Sprintf("获取渠道信息失败，请联系管理员，渠道ID：%d", channelId), - "status": "FAILURE", - "progress": "100%", - }) - if err != nil { - common.SysLog(fmt.Sprintf("UpdateSunoTask error: %v", err)) - } - return err + return fmt.Errorf("CacheGetChannel failed: %w", err) } adaptor := GetTaskAdaptorFunc(constant.TaskPlatformSuno) if adaptor == nil { @@ -214,12 +199,17 @@ func updateSunoTasks(ctx context.Context, channelId int, taskIds []string, taskM var responseItems dto.TaskResponse[[]dto.SunoDataResponse] err = common.Unmarshal(responseBody, &responseItems) if err != nil { - logger.LogError(ctx, fmt.Sprintf("Get Suno Task parse body error2: %v, body: %s", err, string(responseBody))) + logger.LogError(ctx, fmt.Sprintf("Get Suno Task parse body error2: %v, body_bytes=%d", err, len(responseBody))) return err } if !responseItems.IsSuccess() { - common.SysLog(fmt.Sprintf("渠道 #%d 未完成的任务有: %d, 成功获取到任务数: %s", channelId, len(taskIds), string(responseBody))) - return err + common.SysLog(fmt.Sprintf("渠道 #%d 未完成的任务有: %d, response_bytes=%d", channelId, len(taskIds), len(responseBody))) + return fmt.Errorf( + "Suno task polling failed for channel #%d: code=%q, message=%q", + channelId, + responseItems.Code, + common.SanitizePersistedLogContent(common.MaskSensitiveInfo(responseItems.Message)), + ) } for _, responseItem := range responseItems.Data { @@ -240,7 +230,7 @@ func processSunoTaskResponse(ctx context.Context, task *model.Task, responseItem previousStatus := task.Status task.Status = lo.If(model.TaskStatus(responseItem.Status) != "", model.TaskStatus(responseItem.Status)).Else(task.Status) - task.FailReason = lo.If(responseItem.FailReason != "", responseItem.FailReason).Else(task.FailReason) + task.FailReason = lo.If(responseItem.FailReason != "", common.SanitizePersistedLogContent(responseItem.FailReason)).Else(task.FailReason) task.SubmitTime = lo.If(responseItem.SubmitTime != 0, responseItem.SubmitTime).Else(task.SubmitTime) task.StartTime = lo.If(responseItem.StartTime != 0, responseItem.StartTime).Else(task.StartTime) task.FinishTime = lo.If(responseItem.FinishTime != 0, responseItem.FinishTime).Else(task.FinishTime) @@ -257,7 +247,7 @@ func processSunoTaskResponse(ctx context.Context, task *model.Task, responseItem var settlement *model.BillingSettlementInput if isFailure { - settlement = buildTaskRefundSettlementInput(task, task.FailReason) + settlement = BuildTaskRefundSettlementInput(task, task.FailReason) } var won bool var err error @@ -289,7 +279,14 @@ func taskNeedsUpdate(oldTask *model.Task, newTask dto.SunoDataResponse) bool { if string(oldTask.Status) != newTask.Status { return true } - if oldTask.FailReason != newTask.FailReason { + if newTask.FailReason != "" && oldTask.Status != model.TaskStatusFailure { + return true + } + newFailReason := oldTask.FailReason + if newTask.FailReason != "" { + newFailReason = common.SanitizePersistedLogContent(newTask.FailReason) + } + if oldTask.FailReason != newFailReason { return true } @@ -330,21 +327,6 @@ func updateVideoTasks(ctx context.Context, platform constant.TaskPlatform, chann } cacheGetChannel, err := model.CacheGetChannel(channelId) if err != nil { - // Collect DB primary key IDs for bulk update (taskIds are upstream IDs, not task_id column values) - var failedIDs []int64 - for _, upstreamID := range taskIds { - if t, ok := taskM[upstreamID]; ok { - failedIDs = append(failedIDs, t.ID) - } - } - errUpdate := model.TaskBulkUpdateByID(failedIDs, map[string]any{ - "fail_reason": fmt.Sprintf("Failed to get channel info, channel ID: %d", channelId), - "status": "FAILURE", - "progress": "100%", - }) - if errUpdate != nil { - common.SysLog(fmt.Sprintf("UpdateVideoTask error: %v", errUpdate)) - } return fmt.Errorf("CacheGetChannel failed: %w", err) } adaptor := GetTaskAdaptorFunc(platform) @@ -403,7 +385,7 @@ func updateVideoSingleTask(ctx context.Context, adaptor TaskPollingAdaptor, ch * return fmt.Errorf("readAll failed for task %s: %w", taskId, err) } - logger.LogDebug(ctx, "updateVideoSingleTask response: %s", responseBody) + logger.LogDebug(ctx, "updateVideoSingleTask response bytes: %d", len(responseBody)) snap := task.Snapshot() @@ -448,7 +430,7 @@ func updateVideoSingleTask(ctx context.Context, adaptor TaskPollingAdaptor, ch * taskResult = relaycommon.FailTaskInfo("upstream returned error") } else { // unknown error format, log original response - logger.LogError(ctx, fmt.Sprintf("Task %s returned empty status with unrecognized error format, response: %s", taskId, string(responseBody))) + logger.LogError(ctx, fmt.Sprintf("Task %s returned empty status with unrecognized error format, body_bytes=%d", taskId, len(responseBody))) taskResult = relaycommon.FailTaskInfo("upstream returned unrecognized message") } } @@ -492,7 +474,7 @@ func updateVideoSingleTask(ctx context.Context, adaptor TaskPollingAdaptor, ch * if task.FinishTime == 0 { task.FinishTime = now } - task.FailReason = taskResult.Reason + task.FailReason = common.SanitizePersistedLogContent(taskResult.Reason) logger.LogInfo(ctx, fmt.Sprintf("Task %s failed: %s", task.TaskID, task.FailReason)) taskResult.Progress = taskcommon.ProgressComplete if quota != 0 { @@ -512,7 +494,7 @@ func updateVideoSingleTask(ctx context.Context, adaptor TaskPollingAdaptor, ch * if shouldSettle { settlement = prepareTaskCompletionSettlement(ctx, adaptor, task, taskResult, ch.Type, ch.GetOtherSettings()) } else if shouldRefund { - settlement = buildTaskRefundSettlementInput(task, task.FailReason) + settlement = BuildTaskRefundSettlementInput(task, task.FailReason) } var won bool var err error diff --git a/service/task_polling_test.go b/service/task_polling_test.go new file mode 100644 index 00000000..a718e234 --- /dev/null +++ b/service/task_polling_test.go @@ -0,0 +1,114 @@ +package service + +import ( + "bytes" + "context" + "io" + "net/http" + "strings" + "testing" + + "github.com/MAX-API-Next/MAX-API/common" + "github.com/MAX-API-Next/MAX-API/constant" + "github.com/MAX-API-Next/MAX-API/dto" + "github.com/MAX-API-Next/MAX-API/model" + relaycommon "github.com/MAX-API-Next/MAX-API/relay/common" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +type sunoPollingResponseAdaptor struct { + responseBody string +} + +func (a *sunoPollingResponseAdaptor) Init(*relaycommon.RelayInfo) {} + +func (a *sunoPollingResponseAdaptor) FetchTask(string, string, map[string]any, string) (*http.Response, error) { + return &http.Response{ + StatusCode: http.StatusOK, + Body: io.NopCloser(bytes.NewBufferString(a.responseBody)), + }, nil +} + +func (a *sunoPollingResponseAdaptor) ParseTaskResult([]byte) (*relaycommon.TaskInfo, error) { + return nil, nil +} + +func (a *sunoPollingResponseAdaptor) AdjustBillingOnComplete(*model.Task, *relaycommon.TaskInfo) int { + return 0 +} + +func TestUpdateVideoTasksLeavesChargedTasksPendingWhenChannelCacheFails(t *testing.T) { + truncate(t) + + task := makeTask(7101, 999999, 123, 0, BillingSourceWallet, 0) + task.TaskID = "video-cache-failure-upstream" + task.Status = model.TaskStatusInProgress + persistTask(t, task) + + err := updateVideoTasks(context.Background(), constant.TaskPlatform("video"), 999999, []string{task.TaskID}, map[string]*model.Task{ + task.TaskID: task, + }) + require.Error(t, err) + + var reloaded model.Task + require.NoError(t, model.DB.First(&reloaded, task.ID).Error) + assert.EqualValues(t, model.TaskStatusInProgress, reloaded.Status) + assert.Equal(t, 123, reloaded.Quota) + assert.Empty(t, reloaded.FailReason) +} + +func TestUpdateSunoTasksReturnsErrorForUnsuccessfulResponse(t *testing.T) { + truncate(t) + seedChannel(t, 7201) + baseURL := "https://suno.example.com" + require.NoError(t, model.DB.Model(&model.Channel{}).Where("id = ?", 7201).Update("base_url", baseURL).Error) + + originalGetTaskAdaptorFunc := GetTaskAdaptorFunc + GetTaskAdaptorFunc = func(platform constant.TaskPlatform) TaskPollingAdaptor { + require.Equal(t, constant.TaskPlatformSuno, platform) + return &sunoPollingResponseAdaptor{ + responseBody: `{"code":"failure","message":"poll https://api.suno.example.com/v1/tasks?api_key=url-secret api_key:header-secret","data":[]}`, + } + } + t.Cleanup(func() { + GetTaskAdaptorFunc = originalGetTaskAdaptorFunc + }) + + err := updateSunoTasks(context.Background(), 7201, []string{"suno-task"}, map[string]*model.Task{}) + + require.Error(t, err) + require.Contains(t, err.Error(), "https://***.com/") + require.Contains(t, err.Error(), "?api_key=***") + require.Contains(t, err.Error(), "api_key:***") + require.NotContains(t, err.Error(), "suno.example.com") + require.NotContains(t, err.Error(), "url-secret") + require.NotContains(t, err.Error(), "header-secret") +} + +func TestTaskNeedsUpdateComparesPersistedFailReasonNormalization(t *testing.T) { + t.Run("sanitized equivalent", func(t *testing.T) { + rawFailReason := "upstream\r\nfailed\x00" + strings.Repeat("界", common.PersistedLogContentLimit+1) + oldTask := &model.Task{ + Status: model.TaskStatusFailure, + Progress: "100%", + FailReason: common.SanitizePersistedLogContent(rawFailReason), + } + + require.False(t, taskNeedsUpdate(oldTask, dto.SunoDataResponse{ + Status: string(model.TaskStatusFailure), + FailReason: rawFailReason, + })) + }) + + t.Run("empty incoming reason keeps persisted value", func(t *testing.T) { + oldTask := &model.Task{ + Status: model.TaskStatusInProgress, + FailReason: "persisted failure reason", + } + + require.False(t, taskNeedsUpdate(oldTask, dto.SunoDataResponse{ + Status: string(model.TaskStatusInProgress), + })) + }) +} diff --git a/service/tiered_settle.go b/service/tiered_settle.go index 24a19dfb..260c5673 100644 --- a/service/tiered_settle.go +++ b/service/tiered_settle.go @@ -4,6 +4,7 @@ import ( "fmt" "net/http" + "github.com/MAX-API-Next/MAX-API/common" "github.com/MAX-API-Next/MAX-API/dto" "github.com/MAX-API-Next/MAX-API/pkg/billingexpr" relaycommon "github.com/MAX-API-Next/MAX-API/relay/common" @@ -188,9 +189,11 @@ func TryTieredSettle(relayInfo *relaycommon.RelayInfo, params billingexpr.TokenP tr, err := billingexpr.ComputeTieredQuotaWithRequest(snap, params, requestInput) if err != nil { quota = relayInfo.FinalPreConsumedQuota - if quota <= 0 { + if snap.EstimatedQuotaAfterGroup > quota { quota = snap.EstimatedQuotaAfterGroup } + common.SysError(fmt.Sprintf("tiered billing settlement expression failed; using frozen fallback quota (requestId=%s, model=%s, preConsumed=%d, estimated=%d, fallback=%d): %v", + relayInfo.RequestId, relayInfo.OriginModelName, relayInfo.FinalPreConsumedQuota, snap.EstimatedQuotaAfterGroup, quota, err)) return true, quota, nil } diff --git a/service/tiered_settle_test.go b/service/tiered_settle_test.go index 30ad0b14..8e3da08c 100644 --- a/service/tiered_settle_test.go +++ b/service/tiered_settle_test.go @@ -114,6 +114,24 @@ func TestTryTieredSettleFallsBackToFrozenPreConsumeOnExprError(t *testing.T) { } } +func TestTryTieredSettleErrorFallbackUsesHigherFrozenEstimate(t *testing.T) { + relayInfo := &relaycommon.RelayInfo{ + FinalPreConsumedQuota: 123, + TieredBillingSnapshot: &billingexpr.BillingSnapshot{ + BillingMode: "tiered_expr", + ExprString: `invalid +-+ expr`, + ExprHash: billingexpr.ExprHashString(`invalid +-+ expr`), + GroupRatio: 1.0, + EstimatedQuotaAfterGroup: 999, + }, + } + + ok, quota, result := TryTieredSettle(relayInfo, billingexpr.TokenParams{P: 100}) + require.True(t, ok) + require.Equal(t, 999, quota) + require.Nil(t, result) +} + // --------------------------------------------------------------------------- // Pre-consume vs Post-consume consistency // --------------------------------------------------------------------------- diff --git a/setting/ratio_setting/model_ratio.go b/setting/ratio_setting/model_ratio.go index 07011079..dd88da35 100644 --- a/setting/ratio_setting/model_ratio.go +++ b/setting/ratio_setting/model_ratio.go @@ -395,7 +395,10 @@ func GetModelPrice(name string, printErr bool) (float64, bool) { } func UpdateModelRatioByJSONString(jsonStr string) error { - return loadPricingMap(modelRatioMap, jsonStr) + return loadPricingMapWithOptions(modelRatioMap, jsonStr, pricingMapValidationOptions{ + name: "model ratio configuration", + normalizeKey: FormatMatchingModelName, + }) } // 处理带有思考预算的模型名称，方便统一定价 @@ -443,7 +446,10 @@ func CompletionRatio2JSONString() string { } func UpdateCompletionRatioByJSONString(jsonStr string) error { - return loadPricingMap(completionRatioMap, jsonStr) + return loadPricingMapWithOptions(completionRatioMap, jsonStr, pricingMapValidationOptions{ + name: "completion ratio configuration", + normalizeKey: FormatMatchingModelName, + }) } func GetCompletionRatio(name string) float64 { diff --git a/setting/ratio_setting/pricing_validation.go b/setting/ratio_setting/pricing_validation.go index 2cfc10c3..63d61e06 100644 --- a/setting/ratio_setting/pricing_validation.go +++ b/setting/ratio_setting/pricing_validation.go @@ -3,33 +3,118 @@ package ratio_setting import ( "fmt" "math" + "strings" + "unicode/utf8" "github.com/MAX-API-Next/MAX-API/common" "github.com/MAX-API-Next/MAX-API/types" ) +const ( + maxPricingMapJSONBytes = 2 * 1024 * 1024 + maxPricingMapEntries = 20000 + maxPricingMapKeyRunes = 256 +) + +type pricingMapValidationOptions struct { + name string + normalizeKey func(string) string +} + func ValidatePricingMapJSONString(jsonStr string) error { + return validatePricingMapJSONString(jsonStr, pricingMapValidationOptions{ + name: "pricing configuration", + }) +} + +func ValidateModelRatioJSONString(jsonStr string) error { + return validatePricingMapJSONString(jsonStr, pricingMapValidationOptions{ + name: "model ratio configuration", + normalizeKey: FormatMatchingModelName, + }) +} + +func ValidateCompletionRatioJSONString(jsonStr string) error { + return validatePricingMapJSONString(jsonStr, pricingMapValidationOptions{ + name: "completion ratio configuration", + normalizeKey: FormatMatchingModelName, + }) +} + +func validatePricingMapJSONString(jsonStr string, opts pricingMapValidationOptions) error { + _, err := normalizePricingMapJSONString(jsonStr, opts) + return err +} + +func normalizePricingMapJSONString(jsonStr string, opts pricingMapValidationOptions) (string, error) { + if opts.name == "" { + opts.name = "pricing configuration" + } + if len(jsonStr) > maxPricingMapJSONBytes { + return "", fmt.Errorf("%s is too large: %d bytes exceeds limit %d", opts.name, len(jsonStr), maxPricingMapJSONBytes) + } var values map[string]*float64 if err := common.UnmarshalJsonStr(jsonStr, &values); err != nil { - return err + return "", err } if values == nil { - return fmt.Errorf("pricing configuration must be a JSON object") + return "", fmt.Errorf("%s must be a JSON object", opts.name) + } + if len(values) > maxPricingMapEntries { + return "", fmt.Errorf("%s has too many entries: %d exceeds limit %d", opts.name, len(values), maxPricingMapEntries) + } + var normalizedValues map[string]float64 + var normalizedKeys map[string]string + if opts.normalizeKey != nil { + normalizedValues = make(map[string]float64, len(values)) + normalizedKeys = make(map[string]string, len(values)) } for modelName, value := range values { + trimmedName := strings.TrimSpace(modelName) + if trimmedName == "" { + return "", fmt.Errorf("%s contains an empty model name", opts.name) + } + if utf8.RuneCountInString(modelName) > maxPricingMapKeyRunes { + return "", fmt.Errorf("%s model name %q exceeds %d characters", opts.name, modelName, maxPricingMapKeyRunes) + } + normalizedName := trimmedName + if opts.normalizeKey != nil { + normalizedName = opts.normalizeKey(trimmedName) + if existing, ok := normalizedKeys[normalizedName]; ok && existing != modelName { + return "", fmt.Errorf("%s contains duplicate normalized model names %q and %q", opts.name, existing, modelName) + } + normalizedKeys[normalizedName] = modelName + } if value == nil { - return fmt.Errorf("pricing value for model %q must not be null", modelName) + return "", fmt.Errorf("%s value for model %q must not be null", opts.name, modelName) } if *value < 0 || math.IsNaN(*value) || math.IsInf(*value, 0) { - return fmt.Errorf("pricing value for model %q must be finite and non-negative", modelName) + return "", fmt.Errorf("%s value for model %q must be finite and non-negative", opts.name, modelName) + } + if normalizedValues != nil { + normalizedValues[normalizedName] = *value } } - return nil + if normalizedValues == nil { + return jsonStr, nil + } + normalizedJSON, err := common.Marshal(normalizedValues) + if err != nil { + return "", err + } + return string(normalizedJSON), nil } func loadPricingMap(m *types.RWMap[string, float64], jsonStr string) error { - if err := ValidatePricingMapJSONString(jsonStr); err != nil { + return loadPricingMapWithOptions(m, jsonStr, pricingMapValidationOptions{ + name: "pricing configuration", + }) +} + +func loadPricingMapWithOptions(m *types.RWMap[string, float64], jsonStr string, opts pricingMapValidationOptions) error { + normalizedJSON, err := normalizePricingMapJSONString(jsonStr, opts) + if err != nil { return err } - return types.LoadFromJsonStringWithCallback(m, jsonStr, InvalidateExposedDataCache) + return types.LoadFromJsonStringWithCallback(m, normalizedJSON, InvalidateExposedDataCache) } diff --git a/setting/ratio_setting/pricing_validation_test.go b/setting/ratio_setting/pricing_validation_test.go new file mode 100644 index 00000000..ba1925dc --- /dev/null +++ b/setting/ratio_setting/pricing_validation_test.go @@ -0,0 +1,55 @@ +package ratio_setting + +import ( + "strings" + "testing" + + "github.com/MAX-API-Next/MAX-API/types" + "github.com/stretchr/testify/require" +) + +func TestUpdateModelRatioRejectsNormalizedDuplicateWithoutReplacing(t *testing.T) { + original := ModelRatio2JSONString() + require.NoError(t, UpdateModelRatioByJSONString(`{"stable-model":1.5}`)) + t.Cleanup(func() { + require.NoError(t, UpdateModelRatioByJSONString(original)) + }) + + err := UpdateModelRatioByJSONString(`{"gemini-2.5-flash-thinking-a":1,"gemini-2.5-flash-thinking-b":2}`) + + require.Error(t, err) + ratio, ok, _ := GetModelRatio("stable-model") + require.True(t, ok) + require.Equal(t, 1.5, ratio) +} + +func TestUpdateModelRatioLoadsNormalizedAlias(t *testing.T) { + original := ModelRatio2JSONString() + t.Cleanup(func() { + require.NoError(t, UpdateModelRatioByJSONString(original)) + }) + + const alias = "gemini-2.5-flash-thinking-a" + require.NoError(t, UpdateModelRatioByJSONString(`{" `+alias+` ":1.75}`)) + + ratio, ok, normalizedName := GetModelRatio(alias) + require.True(t, ok) + require.Equal(t, "gemini-2.5-flash-thinking-*", normalizedName) + require.Equal(t, 1.75, ratio) +} + +func TestLoadPricingMapWithoutNormalizerPreservesRawKeys(t *testing.T) { + pricingMap := types.NewRWMap[string, float64]() + + require.NoError(t, loadPricingMap(pricingMap, `{" raw-model ":2.5}`)) + + value, ok := pricingMap.Get(" raw-model ") + require.True(t, ok) + require.Equal(t, 2.5, value) +} + +func TestValidatePricingMapRejectsOversizedJSONBeforeLoad(t *testing.T) { + raw := `{"` + strings.Repeat("a", maxPricingMapJSONBytes) + `":1}` + + require.Error(t, ValidatePricingMapJSONString(raw)) +} diff --git a/setting/task_billing_setting/rate_card.go b/setting/task_billing_setting/rate_card.go index 966c98f3..862a6c71 100644 --- a/setting/task_billing_setting/rate_card.go +++ b/setting/task_billing_setting/rate_card.go @@ -42,6 +42,9 @@ func init() { } func GetRateCardsCopy() map[string]RateCard { + if taskBillingSetting.RateCards == nil { + return map[string]RateCard{} + } return cloneRateCards(taskBillingSetting.RateCards.ReadAll()) } @@ -137,6 +140,9 @@ func validQuantity(value float64) bool { } func findRateCard(models ...string) (*RateCard, string) { + if taskBillingSetting.RateCards == nil { + return nil, "" + } rateCards := taskBillingSetting.RateCards.ReadAll() for _, model := range models { model = strings.TrimSpace(model) diff --git a/tools/jsonwrapcheck/allowlist.txt b/tools/jsonwrapcheck/allowlist.txt index b1473e14..e6285063 100644 --- a/tools/jsonwrapcheck/allowlist.txt +++ b/tools/jsonwrapcheck/allowlist.txt @@ -29,12 +29,6 @@ controller/github.go|getGitHubUserInfoByCode|NewDecoder|64ee82d1a886a8c84d5771d6 controller/github.go|getGitHubUserInfoByCode|NewDecoder|4d0c26e8fa58bdb0c5f537b0ace454c979f3c8f0a35ea2e5e74fda46f9ddc9a9 controller/linuxdo.go|getLinuxdoUserInfoByCode|NewDecoder|64ee82d1a886a8c84d5771d6a7733571f892d8b735ff5b5ffa0167dc5f1ecb39 controller/linuxdo.go|getLinuxdoUserInfoByCode|NewDecoder|4d0c26e8fa58bdb0c5f537b0ace454c979f3c8f0a35ea2e5e74fda46f9ddc9a9 -controller/midjourney.go|UpdateMidjourneyTaskBulk|Marshal|84c92152a5e604587b9e7a49d43015339d549d9b9608db43862859d33c2e9ec7 -controller/midjourney.go|UpdateMidjourneyTaskBulk|Unmarshal|1f43767e2ed9cc6c4d09233897cb129c60eac8e2f0ceb6da06822135279a2c76 -controller/midjourney.go|UpdateMidjourneyTaskBulk|Marshal|a94b7820053b86d12899b13132b7cb90fd495bfcb281601dbe5214437610ac56 -controller/midjourney.go|UpdateMidjourneyTaskBulk|Marshal|bcab14b3ab8142cf9c370ab4fa67ba523db802318c0238baaf97e88d98e24c59 -controller/midjourney.go|UpdateMidjourneyTaskBulk|Marshal|3bfaa1b054790e95bd8bd83e0ad1ac0241147d6178e0a08467e03f4f29db6a6c -controller/midjourney.go|checkMjTaskNeedUpdate|Marshal|20a71e5a83bf926ecfe69d5861a4dbbd2f976b5997798ef857aec390fa9a15e3 controller/model_meta.go|enrichModels|Marshal|a2da6b346ff93776c463c1afd5a6ad7a4e117499a61427d2b1bc5a014827e24a controller/model_meta.go|enrichModels|Marshal|a2da6b346ff93776c463c1afd5a6ad7a4e117499a61427d2b1bc5a014827e24a controller/model_sync.go|fetchJSON|Unmarshal|f56c52bcda72d3e86474db3392395083226b9fec13223b5f4489771613b9902b @@ -46,27 +40,10 @@ controller/oidc.go|getOidcUserInfoByCode|NewDecoder|4d0c26e8fa58bdb0c5f537b0ace4 controller/tiered_billing_test.go|performTieredBillingUpdate|Marshal|2af017bb4f8b2b336f7069e3a6f219c0fb04cddbcb36aea984cee19b0af05e40 controller/tiered_billing_test.go|TestUpdateTieredBillingConfigAtomicallyUpdatesModeAndExpr|Unmarshal|80c32abf8b55ed520ba61e08871b1a3bcf25aa90146d3aacb07ff427b8fcb308 controller/tiered_billing_test.go|TestUpdateTieredBillingConfigRejectsInvalidExpression|Unmarshal|80c32abf8b55ed520ba61e08871b1a3bcf25aa90146d3aacb07ff427b8fcb308 -controller/topup_creem.go|*CreemAdaptor.RequestPay|Unmarshal|77623a5221a2cabc4986ca29fe2cf0cbdffb83554fe5bd96abf38794bb21a7ad -controller/topup_creem.go|genCreemLink|Marshal|64d944a5ed933c0e1229638016ea7b948f3f7af016473b552da179ef724a2588 -controller/topup_creem.go|genCreemLink|Unmarshal|246f2230ef15513f964156c37938859a8f00a897f55e0c3cd0043470ebf40108 controller/uptime_kuma.go|getAndDecode|NewDecoder|4e62c918bd14f5107549183c31d628b8b7c5c0f6d37c3b60ee14c04233f62047 controller/wechat.go|getWeChatIdByCode|NewDecoder|032c2fadb2b3bd14e88241bb5e9c7e1103ba0b6141f0525f0a9674892e6c8d59 dto/claude.go|*ClaudeRequest.GetEfforts|Unmarshal|3525f5b3c4bd5cbad903f267e36a9250edb08eb82f5ea217b5c72de9f368a9a3 -dto/openai_request.go|*Message.ParseToolCalls|Unmarshal|2969087b249de385620df94374e7ded88cbe21b27c422ee1836cfcec6ef72d2a -dto/openai_request.go|*Message.SetToolCalls|Marshal|6fbd2e8a031c5a521cbdd73843d1599d5ce586be435295775ce81592e7035ed7 -dto/values.go|*StringValue.UnmarshalJSON|Unmarshal|696c84522bd5ca3e97cd2054e281a036fe1ccc61e61e07b613c6af97dad8f64b -dto/values.go|*StringValue.UnmarshalJSON|Unmarshal|b302d32cbbccd546a60c4b7ee3efd96ffd097e331aab7a511ea77f01fb3a9d54 -dto/values.go|*StringValue.UnmarshalJSON|Unmarshal|696c84522bd5ca3e97cd2054e281a036fe1ccc61e61e07b613c6af97dad8f64b -dto/values.go|StringValue.MarshalJSON|Marshal|bf55276c5dadb05a5567c2dc1c4275674fe68cba6e9dc1e5cf71c233b8b4c110 -dto/values.go|*IntValue.UnmarshalJSON|Unmarshal|74017e5e429f1a6e2e2e721b508fa0ac0646035d76fd670867b236a774fc3775 -dto/values.go|*IntValue.UnmarshalJSON|Unmarshal|a7b5a77e3c5c35fb27bbe1156e261507108920c432efc748ee5f208975d53213 -dto/values.go|IntValue.MarshalJSON|Marshal|46f8a02fb525b82abfd486bb64d7a5f98eba67bb8d7581c4a207ab204e1eb9ff -dto/values.go|*BoolValue.UnmarshalJSON|Unmarshal|20093b52db26f80edb520049d20008b7c1a1ec58019c1671c2b1ccb20ef7c7b6 -dto/values.go|*BoolValue.UnmarshalJSON|Unmarshal|696c84522bd5ca3e97cd2054e281a036fe1ccc61e61e07b613c6af97dad8f64b -dto/values.go|*BoolValue.UnmarshalJSON|Unmarshal|20093b52db26f80edb520049d20008b7c1a1ec58019c1671c2b1ccb20ef7c7b6 -dto/values.go|BoolValue.MarshalJSON|Marshal|bc53c95a352b7d8ffde9f9b4d44968a905dfc0a762fdb9659c059feef11ed2a4 middleware/jimeng_adapter.go|JimengRequestConvert|Marshal|418a4a72b6b9c61eb7c62b58bb0f4bb1e23e251bda4f5a8f3514d66ebfd8d962 -model/channel.go|*Channel.SetOtherInfo|Marshal|18d7bc9d32c6ec57f9b016303592cb5e327ffbe1abfc9b66cb7bd8033d1d36b0 model/passkey.go|*PasskeyCredential.TransportList|Unmarshal|8174884feacc1fdfef084adb2900b1a97565b6c6c5f072b2d390e1c1dbadaa7f model/passkey.go|*PasskeyCredential.SetTransports|Marshal|891aca816b90a7b0706fdd24581e477bee4335c8f916733130b54141eb6448c7 model/prefill_group.go|*JSONValue.Scan|Marshal|be383ea6cfaaf292483613ff481847bcf18cbddf173a0870a25e345b29d9f5eb @@ -79,15 +56,7 @@ oauth/github.go|*GitHubProvider.ExchangeToken|NewDecoder|64ee82d1a886a8c84d5771d oauth/github.go|*GitHubProvider.GetUserInfo|NewDecoder|64ee82d1a886a8c84d5771d6a7733571f892d8b735ff5b5ffa0167dc5f1ecb39 oauth/linuxdo.go|*LinuxDOProvider.ExchangeToken|NewDecoder|64ee82d1a886a8c84d5771d6a7733571f892d8b735ff5b5ffa0167dc5f1ecb39 oauth/linuxdo.go|*LinuxDOProvider.GetUserInfo|NewDecoder|64ee82d1a886a8c84d5771d6a7733571f892d8b735ff5b5ffa0167dc5f1ecb39 -oauth/oidc.go|*OIDCProvider.ExchangeToken|NewDecoder|64ee82d1a886a8c84d5771d6a7733571f892d8b735ff5b5ffa0167dc5f1ecb39 -oauth/oidc.go|*OIDCProvider.GetUserInfo|NewDecoder|64ee82d1a886a8c84d5771d6a7733571f892d8b735ff5b5ffa0167dc5f1ecb39 -relay/channel/ali/rerank.go|RerankHandler|Unmarshal|66eb02820ea009d9d095c26550faf57ec59dba488d801f0cf3a611d01b3cfcf7 -relay/channel/ali/rerank.go|RerankHandler|Marshal|ea6233db65ed7e55f6f9ae4ca1f4d9b9ca04f9e508ab5eef3dc68f139231720d relay/channel/aws/dto.go|formatRequest|Marshal|0811da817c96b432dc92949bfc702d029e13adb47ab075abd27b6b82eae30478 -relay/channel/baidu/relay-baidu.go|baiduHandler|Unmarshal|6c88b37cda6b69cb880852f57074d792626b82bd771acbab3a46754f80631684 -relay/channel/baidu/relay-baidu.go|baiduHandler|Marshal|92981973316a59184c1f9d4fdca6fe3a1045b5da6a30f89abb52b516f1849571 -relay/channel/baidu/relay-baidu.go|baiduEmbeddingHandler|Unmarshal|6c88b37cda6b69cb880852f57074d792626b82bd771acbab3a46754f80631684 -relay/channel/baidu/relay-baidu.go|baiduEmbeddingHandler|Marshal|92981973316a59184c1f9d4fdca6fe3a1045b5da6a30f89abb52b516f1849571 relay/channel/claude/relay-claude.go|ResponseClaude2OpenAI|Marshal|9f23cc0eecff96adc0410bf3360508a44815f9bc18567a76278a14aaa6e46e0a relay/channel/claude/relay-claude.go|HandleClaudeResponseData|Marshal|0e01a2a781af1a773195f919e8bc2c53345aed7dd99776d98cfe550576ea831a relay/channel/cloudflare/relay_cloudflare.go|cfStreamHandler|Unmarshal|5a9c4fba44410fd4d5ee918ce49cdd49a713d756c491be89ccc2d9f3de44cdc5 @@ -101,14 +70,6 @@ relay/channel/cohere/relay-cohere.go|cohereHandler|Unmarshal|60f8972547f5094bf2b relay/channel/cohere/relay-cohere.go|cohereHandler|Marshal|e835663a563231e050ffc86f98edc0d4ea411076890252d835d77715185f6317 relay/channel/cohere/relay-cohere.go|cohereRerankHandler|Unmarshal|60f8972547f5094bf2b8ef4f8aee5478f66cbaec849d110427da197c8f20a177 relay/channel/cohere/relay-cohere.go|cohereRerankHandler|Marshal|7175c817a0dede230ae6b179388d06512cf9cc2c1d598aafa1b52cf3a52c7ad7 -relay/channel/coze/adaptor.go|*Adaptor.DoRequest|Unmarshal|4a81d986450fdd36e06fe6b2f0daa87bb14ba3753e1b988d9f1de193ad0f19ed -relay/channel/coze/relay-coze.go|cozeChatHandler|Unmarshal|039c274556c751a9861fa6dc063c6e342175865072161dd9fa1f2191e9c77119 -relay/channel/coze/relay-coze.go|cozeChatHandler|Marshal|168eb97ce25654c31be502868a95131c40bf62f506fa17dca8eec1cd5263b090 -relay/channel/coze/relay-coze.go|handleCozeEvent|Unmarshal|46986195906ac6b619793a91119aa284a08c45efcc069da1734e4e7973bf1a27 -relay/channel/coze/relay-coze.go|handleCozeEvent|Unmarshal|ac93cdd9401220cc3b63839da7fca6a90f99e4c4756d815bf8b6921356ebbd75 -relay/channel/coze/relay-coze.go|handleCozeEvent|Unmarshal|3f2ecb150024389acdffdac06618c7ba95fc5595d6b22d11b5b40c89757304aa -relay/channel/coze/relay-coze.go|handleCozeEvent|Unmarshal|5ce6ed25c133490ae1a626176b5f9981d5fa7fd487cbe8309cc649c69c2594b4 -relay/channel/coze/relay-coze.go|checkIfChatComplete|Unmarshal|039c274556c751a9861fa6dc063c6e342175865072161dd9fa1f2191e9c77119 relay/channel/dify/relay-dify.go|uploadDifyFile|NewDecoder|4e62c918bd14f5107549183c31d628b8b7c5c0f6d37c3b60ee14c04233f62047 relay/channel/dify/relay-dify.go|requestOpenAI2Dify|Unmarshal|301e4e0d469fcbcb1f5fcfaf6db9a98aac7c3efbc95148406f22c7e1bb7f79eb relay/channel/dify/relay-dify.go|difyStreamHandler|Unmarshal|c7b78d67e7541bc2dcbaca687f311c4a80640b86227ff7660fa1593679d9ccda @@ -137,22 +98,12 @@ relay/channel/palm/relay-palm.go|palmStreamHandler|Unmarshal|9e6a91a6205bc2652fd relay/channel/palm/relay-palm.go|palmStreamHandler|Marshal|92981973316a59184c1f9d4fdca6fe3a1045b5da6a30f89abb52b516f1849571 relay/channel/palm/relay-palm.go|palmHandler|Unmarshal|9e6a91a6205bc2652fded36139425ad52ae7c49ef94ccd023817fca36089ced0 relay/channel/replicate/adaptor.go|*Adaptor.ConvertImageRequest|Unmarshal|8657e632acc4f987e7c3c0aed04a081ab0b5762843f4e8a91305c1640357a714 -relay/channel/siliconflow/relay-siliconflow.go|siliconflowRerankHandler|Unmarshal|09aba8238a97edabda1d5481e553802d2eeae9a5adf762c00e6c452f0d583106 -relay/channel/siliconflow/relay-siliconflow.go|siliconflowRerankHandler|Marshal|7175c817a0dede230ae6b179388d06512cf9cc2c1d598aafa1b52cf3a52c7ad7 relay/channel/tencent/relay-tencent.go|tencentHandler|Unmarshal|c5c85399232311a43e2b04dcf6ab971f636d174fd3fd7a9c10168b9f3f4a9259 relay/channel/tencent/relay-tencent.go|getTencentSign|Marshal|556c76f2141b112391c2965ec7de57f4ac5225bcabc9df6af104fb7f1bac4a05 -relay/channel/vertex/adaptor.go|*Adaptor.ConvertOpenAIRequest|Unmarshal|adb7d12895b84c0ef11647409a61da1141cbef0e6b761e335146ce0045cbed02 -relay/channel/vertex/service_account.go|exchangeJwtForAccessToken|NewDecoder|4e62c918bd14f5107549183c31d628b8b7c5c0f6d37c3b60ee14c04233f62047 -relay/channel/vertex/service_account.go|exchangeJwtForAccessTokenWithProxy|NewDecoder|4e62c918bd14f5107549183c31d628b8b7c5c0f6d37c3b60ee14c04233f62047 relay/channel/volcengine/adaptor.go|*Adaptor.ConvertAudioRequest|Unmarshal|21d62956ab82cf6356638ba2e98ab21e8970ef3fec36d51a2362ea78b972e685 relay/channel/volcengine/adaptor.go|*Adaptor.ConvertAudioRequest|Marshal|6df4bc98beb165214f3179900a7dc69f4c1aaa89816f2699cf59ee2bd83ef3d2 relay/channel/volcengine/tts.go|handleTTSResponse|Unmarshal|1c19ee75effdb9440195e572a78088232532eacc7938c6bf68a28030b8265b5f relay/channel/volcengine/tts.go|handleTTSWebSocketResponse|Marshal|6df4bc98beb165214f3179900a7dc69f4c1aaa89816f2699cf59ee2bd83ef3d2 -relay/channel/zhipu/relay-zhipu.go|zhipuStreamHandler|Marshal|168eb97ce25654c31be502868a95131c40bf62f506fa17dca8eec1cd5263b090 -relay/channel/zhipu/relay-zhipu.go|zhipuStreamHandler|Unmarshal|d6c2892f24d52e3e062e3eb924a5a5e290199f1238c21310625defc58a4d1cdf -relay/channel/zhipu/relay-zhipu.go|zhipuStreamHandler|Marshal|168eb97ce25654c31be502868a95131c40bf62f506fa17dca8eec1cd5263b090 -relay/channel/zhipu/relay-zhipu.go|zhipuHandler|Unmarshal|ecfb8a328f2c468cfb19c2ffc8ca3cc558b08d1dc290b0fcf5d5c4c503de0a75 -relay/channel/zhipu/relay-zhipu.go|zhipuHandler|Marshal|92981973316a59184c1f9d4fdca6fe3a1045b5da6a30f89abb52b516f1849571 relay/common/override_test.go|TestApplyParamOverrideDelete|Unmarshal|692246112d039d019f6ed42cbf2362d484af719d1b9abec92686330698a56345 relay/common/override_test.go|TestApplyParamOverrideSetWildcardPath|Unmarshal|692246112d039d019f6ed42cbf2362d484af719d1b9abec92686330698a56345 relay/common/override_test.go|TestApplyParamOverrideTrimSpaceWildcardPath|Unmarshal|692246112d039d019f6ed42cbf2362d484af719d1b9abec92686330698a56345 @@ -161,19 +112,7 @@ relay/common/override_test.go|TestApplyParamOverrideTrimSpaceMultiWildcardPath|U relay/common/override_test.go|assertJSONEqual|Unmarshal|1b6e8a2882cd33079a79d74aeb366ef371ac02cbffa6d90f82ab6027c4e989bd relay/common/override_test.go|assertJSONEqual|Unmarshal|444da0b2c667f90adf285fc754ed502890312ce4837a5ec13c0ffbfb5aaa77bb relay/helper/model_mapped.go|ModelMappedHelper|Unmarshal|397c955a41ac41e35fdf6469fe18566e192cf9a0762e2594fba342148e25e7b7 -relay/mjproxy_handler.go|RelayMidjourneyNotify|Marshal|9a70f351fbab5e1ed0c3573dec3f415740d755011757d03854eeceb0e1fe4ec8 -relay/mjproxy_handler.go|coverMidjourneyTaskDto|Unmarshal|24c991800ec0e31d5fecaa9b39cac52f390648fd382d3fa569328aaa9b8d5ce9 -relay/mjproxy_handler.go|coverMidjourneyTaskDto|Unmarshal|e9a7e8d8a2e85745ba88c5778c78e9e3780131045371f89e3e4271ad9757fcf6 -relay/mjproxy_handler.go|coverMidjourneyTaskDto|Unmarshal|a5f4f149ec0cf55ed8663e19972391150e8e718c9fa66d38ffc0eb73d31172c4 -relay/mjproxy_handler.go|RelaySwapFace|Marshal|43f556717f85f280992d3a2d0008565d6ddbdba5994fdce3edcd08c6a91dc320 -relay/mjproxy_handler.go|RelayMidjourneyTaskImageSeed|Marshal|43f556717f85f280992d3a2d0008565d6ddbdba5994fdce3edcd08c6a91dc320 -relay/mjproxy_handler.go|RelayMidjourneyTask|Marshal|a06ef6dffc8c25d4f764f6e368c1b7519640fbe49986e4072b764c3a82d19dc0 -relay/mjproxy_handler.go|RelayMidjourneyTask|Marshal|8b291e884ff2bd4bdff251dcb055060790f313706be033cae508d75859f79cac service/codex_oauth.go|decodeJWTClaims|Unmarshal|7ca24118dbcd41c0feb7b0767369b64fadf9b9a4b1bcb1e62d1ec99b54f487d0 -service/midjourney.go|DoMidjourneyHttpRequest|NewDecoder|dc2d92e0dd4c66591ccceefe086dca58c0af182c9e2c61f92c80cda6833d2f9a -service/midjourney.go|DoMidjourneyHttpRequest|Marshal|af957403b76886439ffd61be2eec9cef5f0fd1f3c0b8201b7a5e9565234b44ce -service/midjourney.go|DoMidjourneyHttpRequest|Unmarshal|a15ce743bd5d73944713700514fe95dfe90f903ea83e6508c1d05f916dae3d35 -service/midjourney.go|DoMidjourneyHttpRequest|Unmarshal|b2bd81a114cb2619945fabb0351b230f06b5787f7ac39f2c3009db9ad53a84d8 service/passkey/session.go|SaveSessionData|Marshal|5b9b772fb88d846779cb56c7c4ca8852b09c2b90d1e5ae47e5e29386bd45a06e service/passkey/session.go|PopSessionData|Unmarshal|337d2d0489c6f4c2daf3edbb46ea8101d8b9b64a46fce59768dbc2aaacf048f6 service/passkey/session.go|PopSessionData|Unmarshal|f2cd8d698a8b3baee9426fd333d1a8ef6eae1c2f3c53939906d9181818b00c76 diff --git a/web/default/src/components/ai-elements/code-block.test.ts b/web/default/src/components/ai-elements/code-block.test.ts new file mode 100644 index 00000000..0486d098 --- /dev/null +++ b/web/default/src/components/ai-elements/code-block.test.ts @@ -0,0 +1,69 @@ +/* +Copyright (C) 2023-2026 MAX-API-Next + +This program is free software: you can redistribute it and/or modify +it under the terms of the GNU Affero General Public License as +published by the Free Software Foundation, either version 3 of the +License, or (at your option) any later version. + +This program is distributed in the hope that it will be useful, +but WITHOUT ANY WARRANTY; without even the implied warranty of +MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +GNU Affero General Public License for more details. + +You should have received a copy of the GNU Affero General Public License +along with this program. If not, see . + +For commercial licensing, please contact https://github.com/MAX-API-Next/MAX-API/issues +*/ +import { JSDOM } from 'jsdom' +import assert from 'node:assert/strict' +import { after, describe, test } from 'node:test' +import { renderHighlightedCode, sanitizeHighlightedCode } from './code-block' + +const dom = new JSDOM('', { url: 'http://localhost/' }) +const previousWindowDescriptor = Object.getOwnPropertyDescriptor( + globalThis, + 'window' +) + +Object.defineProperty(globalThis, 'window', { + configurable: true, + value: dom.window, +}) + +after(() => { + if (previousWindowDescriptor) { + Object.defineProperty(globalThis, 'window', previousWindowDescriptor) + return + } + + delete (globalThis as Partial).window +}) + +describe('renderHighlightedCode', () => { + test('sanitizes pre-rendered HTML independently of Shiki escaping', () => { + const html = sanitizeHighlightedCode( + '
safe
' + ) + const fragment = JSDOM.fragment(html) + + const image = fragment.querySelector('img') + + assert.notEqual(image, null) + assert.equal(image?.hasAttribute('onerror'), false) + assert.equal(fragment.querySelector('script'), null) + assert.equal(fragment.textContent?.includes('safe'), true) + }) + + test('sanitizes highlighted code before it reaches innerHTML', async () => { + const payload = + '' + const html = await renderHighlightedCode(payload, 'html') + const fragment = JSDOM.fragment(html) + + assert.equal(fragment.querySelector('img'), null) + assert.equal(fragment.querySelector('script'), null) + assert.match(fragment.textContent ?? '', /onerror/) + }) +}) diff --git a/web/default/src/components/ai-elements/code-block.tsx b/web/default/src/components/ai-elements/code-block.tsx index f4cf3805..0269532b 100644 --- a/web/default/src/components/ai-elements/code-block.tsx +++ b/web/default/src/components/ai-elements/code-block.tsx @@ -33,6 +33,7 @@ import { codeToHtml, type ShikiTransformer, } from 'shiki/bundle/web' +import { sanitizeHtmlContent } from '@/lib/html-sanitizer' import { cn } from '@/lib/utils' import { Button } from '@/components/ui/button' @@ -75,7 +76,7 @@ export async function highlightCode( code: string, language: BundledLanguage, showLineNumbers = false -) { +): Promise { const transformers: ShikiTransformer[] = showLineNumbers ? [lineNumberTransformer] : [] @@ -90,6 +91,20 @@ export async function highlightCode( }) } +export async function renderHighlightedCode( + code: string, + language: BundledLanguage, + showLineNumbers = false +): Promise { + return sanitizeHighlightedCode( + await highlightCode(code, language, showLineNumbers) + ) +} + +export function sanitizeHighlightedCode(html: string): string { + return sanitizeHtmlContent(html) +} + export const CodeBlock = ({ code, language, @@ -102,7 +117,7 @@ export const CodeBlock = ({ useEffect(() => { let cancelled = false - highlightCode(code, language, showLineNumbers).then((next) => { + renderHighlightedCode(code, language, showLineNumbers).then((next) => { if (!cancelled) { setHtml(next) } diff --git a/web/default/src/components/ui/chart.test.tsx b/web/default/src/components/ui/chart.test.tsx new file mode 100644 index 00000000..ccc7fb36 --- /dev/null +++ b/web/default/src/components/ui/chart.test.tsx @@ -0,0 +1,59 @@ +/* +Copyright (C) 2023-2026 MAX-API-Next + +This program is free software: you can redistribute it and/or modify +it under the terms of the GNU Affero General Public License as +published by the Free Software Foundation, either version 3 of the +License, or (at your option) any later version. + +This program is distributed in the hope that it will be useful, +but WITHOUT ANY WARRANTY; without even the implied warranty of +MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +GNU Affero General Public License for more details. + +You should have received a copy of the GNU Affero General Public License +along with this program. If not, see . + +For commercial licensing, please contact https://github.com/MAX-API-Next/MAX-API/issues +*/ +import React from 'react' +import { JSDOM } from 'jsdom' +import assert from 'node:assert/strict' +import { describe, test } from 'node:test' +import { renderToStaticMarkup } from 'react-dom/server' +import { ChartStyle } from './chart' + +describe('ChartStyle', () => { + test('does not allow chart config CSS to escape the style element', () => { + const payload = + 'red;} } const ChartTooltip = RechartsPrimitive.Tooltip diff --git a/web/default/src/features/auth/api.ts b/web/default/src/features/auth/api.ts index 2055707d..1ff95d70 100644 --- a/web/default/src/features/auth/api.ts +++ b/web/default/src/features/auth/api.ts @@ -87,10 +87,14 @@ export async function githubOAuthStart(clientId: string, state: string) { } // Get OAuth state for CSRF protection -export async function getOAuthState(): Promise { +export async function getOAuthState(provider: string): Promise { const aff = typeof window !== 'undefined' ? (localStorage.getItem('aff') ?? '') : '' - const res = await api.get('/api/oauth/state', { params: { aff } }) + const res = await api.post('/api/oauth/state', { + provider, + intent: 'login', + aff, + }) if (res.data?.success) return res.data.data return '' } diff --git a/web/default/src/features/auth/hooks/use-oauth-login.ts b/web/default/src/features/auth/hooks/use-oauth-login.ts index 5e69454b..77d3b44d 100644 --- a/web/default/src/features/auth/hooks/use-oauth-login.ts +++ b/web/default/src/features/auth/hooks/use-oauth-login.ts @@ -93,7 +93,7 @@ export function useOAuthLogin(status: SystemStatus | null) { try { await resetSession() - const state = await getOAuthState() + const state = await getOAuthState('github') if (!state) { toast.error(t('Failed to initialize OAuth')) if (githubTimeoutRef.current) { @@ -124,7 +124,7 @@ export function useOAuthLogin(status: SystemStatus | null) { setIsLoading(true) try { await resetSession() - const state = await getOAuthState() + const state = await getOAuthState('discord') if (!state) { toast.error(t('Failed to initialize OAuth')) return @@ -145,7 +145,7 @@ export function useOAuthLogin(status: SystemStatus | null) { setIsLoading(true) try { await resetSession() - const state = await getOAuthState() + const state = await getOAuthState('oidc') if (!state) { toast.error(t('Failed to initialize OAuth')) return @@ -170,7 +170,7 @@ export function useOAuthLogin(status: SystemStatus | null) { setIsLoading(true) try { await resetSession() - const state = await getOAuthState() + const state = await getOAuthState('linuxdo') if (!state) { toast.error(t('Failed to initialize OAuth')) return @@ -195,7 +195,7 @@ export function useOAuthLogin(status: SystemStatus | null) { setIsLoading(true) try { await resetSession() - const state = await getOAuthState() + const state = await getOAuthState(provider.slug) if (!state) { toast.error(t('Failed to initialize OAuth')) return diff --git a/web/default/src/lib/oauth.ts b/web/default/src/lib/oauth.ts index 3c2de139..a43c1e4e 100644 --- a/web/default/src/lib/oauth.ts +++ b/web/default/src/lib/oauth.ts @@ -44,7 +44,10 @@ function markOAuthBindingState(provider: string, state: string): void { } } -export function isOAuthBindingState(provider?: string, state?: string): boolean { +export function isOAuthBindingState( + provider?: string, + state?: string +): boolean { if (!provider || !state) return false try { const key = getOAuthBindingStateKey(provider, state) @@ -145,14 +148,12 @@ export function buildLinuxDOOAuthUrl(clientId: string, state: string): string { * Get OAuth state token * Includes affiliate code from localStorage if available */ -export async function getOAuthState(): Promise { +export async function getOAuthState(provider: string): Promise { try { - let path = '/api/oauth/state' - const affCode = localStorage.getItem('aff') - if (affCode && affCode.length > 0) { - path += `?aff=${affCode}` - } - const res = await api.get(path) + const res = await api.post('/api/oauth/state', { + provider: normalizeOAuthProvider(provider), + intent: 'bind', + }) if (res.data.success) { return res.data.data } @@ -168,7 +169,7 @@ export async function getOAuthState(): Promise { * Handle GitHub OAuth binding/login */ export async function handleGitHubOAuth(clientId: string): Promise { - const state = await getOAuthState() + const state = await getOAuthState('github') if (!state) return const url = buildGitHubOAuthUrl(clientId, state) @@ -179,7 +180,7 @@ export async function handleGitHubOAuth(clientId: string): Promise { * Handle Discord OAuth binding/login */ export async function handleDiscordOAuth(clientId: string): Promise { - const state = await getOAuthState() + const state = await getOAuthState('discord') if (!state) return const url = buildDiscordOAuthUrl(clientId, state) @@ -193,7 +194,7 @@ export async function handleOIDCOAuth( authUrl: string, clientId: string ): Promise { - const state = await getOAuthState() + const state = await getOAuthState('oidc') if (!state) return const url = buildOIDCOAuthUrl(authUrl, clientId, state) @@ -204,7 +205,7 @@ export async function handleOIDCOAuth( * Handle LinuxDO OAuth binding/login */ export async function handleLinuxDOOAuth(clientId: string): Promise { - const state = await getOAuthState() + const state = await getOAuthState('linuxdo') if (!state) return const url = buildLinuxDOOAuthUrl(clientId, state)