diff --git a/constant/midjourney.go b/constant/midjourney.go index 5934be2f..8d4a2f85 100644 --- a/constant/midjourney.go +++ b/constant/midjourney.go @@ -5,6 +5,8 @@ const ( MjRequestError = 4 ) +const MjBillingSettlementPending = "billing_settlement_pending" + const ( MjActionImagine = "IMAGINE" MjActionDescribe = "DESCRIBE" diff --git a/controller/midjourney_test.go b/controller/midjourney_test.go index d776dfbf..865507d3 100644 --- a/controller/midjourney_test.go +++ b/controller/midjourney_test.go @@ -175,6 +175,16 @@ func TestMidjourneyFailureRefundsWalletAndTokenThroughDurableSettlement(t *testi require.NoError(t, err) require.True(t, created) require.False(t, refundDuplicate) + settlementInput := model.BillingSettlementInput{ + OperationKey: "request:midjourney-refund-request:finalize", Source: model.BillingSettlementSourceWallet, + UserID: user.Id, TokenID: token.Id, + } + _, alreadyApplied, err := model.ApplyBillingSettlementOnce(settlementInput) + require.NoError(t, err) + require.False(t, alreadyApplied) + _, alreadyApplied, err = model.ApplyBillingSettlementOnce(settlementInput) + require.NoError(t, err) + require.True(t, alreadyApplied) err = updateMidjourneyTaskFromResponse(context.Background(), mj, dto.MidjourneyDto{ MjId: mj.MjId, Status: "FAILURE", Progress: "100%", FailReason: "provider failed", @@ -258,8 +268,18 @@ func TestApplyMidjourneyTaskResponsePreservesLateCompletionWithoutRefund(t *test require.NoError(t, err) require.True(t, created) require.False(t, refundDuplicate) + settlementInput := model.BillingSettlementInput{ + OperationKey: "request:midjourney-late-success-request-" + suffix + ":finalize", Source: model.BillingSettlementSourceWallet, + UserID: user.Id, TokenID: token.Id, + } + _, alreadyApplied, err := model.ApplyBillingSettlementOnce(settlementInput) + require.NoError(t, err) + require.False(t, alreadyApplied) var settlementCountBefore int64 require.NoError(t, db.Model(&model.BillingSettlement{}).Count(&settlementCountBefore).Error) + _, alreadyApplied, err = model.ApplyBillingSettlementOnce(settlementInput) + require.NoError(t, err) + require.True(t, alreadyApplied) err = applyMidjourneyTaskResponse(context.Background(), map[string]*model.Midjourney{mj.MjId: mj}, dto.MidjourneyDto{ MjId: mj.MjId, Status: tt.responseStatus, Progress: "100%", FinishTime: finishTime, diff --git a/controller/misc.go b/controller/misc.go index d7b81929..e1e0b2d6 100644 --- a/controller/misc.go +++ b/controller/misc.go @@ -408,9 +408,10 @@ func ResetPassword(c *gin.Context) { return } c.JSON(http.StatusOK, gin.H{ - "success": true, - "message": "", - "data": password, + "success": true, + "message": common.TranslateMessage(c, i18n.MsgUserPasswordResetSuccessTokensRevoked), + "data": password, + "api_tokens_revoked": true, }) return } diff --git a/controller/oauth.go b/controller/oauth.go index 61d8725a..5a2ab1bf 100644 --- a/controller/oauth.go +++ b/controller/oauth.go @@ -54,6 +54,13 @@ func providerParams(name string) map[string]any { return map[string]any{"Provider": name} } +func supportsOAuthState(provider string, intent string) bool { + if provider == "telegram" { + return common.TelegramOAuthEnabled && intent == model.AuthFlowIntentLogin + } + return oauth.GetProvider(provider) != nil +} + func oauthProviderUserUpdateField(provider oauth.Provider) (model.UserUpdateField, bool) { switch provider.(type) { case *oauth.GitHubProvider: @@ -120,7 +127,7 @@ func GenerateOAuthCode(c *gin.Context) { request.Provider = strings.ToLower(strings.TrimSpace(request.Provider)) request.Intent = strings.TrimSpace(request.Intent) request.Aff = strings.TrimSpace(request.Aff) - if oauth.GetProvider(request.Provider) == nil || + if !supportsOAuthState(request.Provider, request.Intent) || (request.Intent != model.AuthFlowIntentLogin && request.Intent != model.AuthFlowIntentBind) || len(request.Aff) > 32 || (request.Intent == model.AuthFlowIntentBind && request.Aff != "") { diff --git a/controller/oauth_test.go b/controller/oauth_test.go index 182a6af5..af8420e3 100644 --- a/controller/oauth_test.go +++ b/controller/oauth_test.go @@ -3,6 +3,7 @@ package controller import ( "context" "errors" + "fmt" "net/http" "net/http/httptest" "strings" @@ -31,18 +32,77 @@ func (*deletedUserOAuthProvider) ExchangeToken(context.Context, string, *gin.Con 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{}) + originalLogDB := model.LOG_DB + 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) - require.NoError(t, db.AutoMigrate(&model.AuthFlow{})) + require.NoError(t, db.AutoMigrate(&model.AuthFlow{}, &model.User{}, &model.Log{})) model.DB = db + model.LOG_DB = db t.Cleanup(func() { model.DB = originalDB + model.LOG_DB = originalLogDB if sqlDB, dbErr := db.DB(); dbErr == nil { _ = sqlDB.Close() } }) } +func TestGenerateOAuthCodeCreatesTelegramLoginState(t *testing.T) { + setupControllerAuthFlowDB(t) + originalEnabled := common.TelegramOAuthEnabled + common.TelegramOAuthEnabled = true + t.Cleanup(func() { common.TelegramOAuthEnabled = originalEnabled }) + gin.SetMode(gin.TestMode) + router := gin.New() + router.Use(sessions.Sessions("session", cookie.NewStore([]byte("telegram-state-create-test")))) + router.POST("/api/oauth/state", GenerateOAuthCode) + + recorder := httptest.NewRecorder() + request := httptest.NewRequest(http.MethodPost, "/api/oauth/state", strings.NewReader(`{"provider":"telegram","intent":"login"}`)) + 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.NotEmpty(t, response.Data) + _, err := model.GetAuthFlow(response.Data, model.AuthFlowMatch{ + Purpose: model.AuthFlowPurposeOAuth, Provider: "telegram", Intent: model.AuthFlowIntentLogin, + }) + require.NoError(t, err) +} + +func TestGenerateOAuthCodeRejectsTelegramLoginStateWhenDisabled(t *testing.T) { + setupControllerAuthFlowDB(t) + require.NoError(t, appi18n.Init()) + originalEnabled := common.TelegramOAuthEnabled + common.TelegramOAuthEnabled = false + t.Cleanup(func() { common.TelegramOAuthEnabled = originalEnabled }) + + router := gin.New() + router.Use(sessions.Sessions("session", cookie.NewStore([]byte("telegram-state-disabled-test")))) + router.POST("/api/oauth/state", GenerateOAuthCode) + + recorder := httptest.NewRecorder() + request := httptest.NewRequest(http.MethodPost, "/api/oauth/state", strings.NewReader(`{"provider":"telegram","intent":"login"}`)) + request.Header.Set("Content-Type", "application/json") + router.ServeHTTP(recorder, request) + + require.Equal(t, http.StatusOK, recorder.Code) + var response struct { + Success bool `json:"success"` + } + require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response)) + require.False(t, response.Success) + var count int64 + require.NoError(t, model.DB.Model(&model.AuthFlow{}).Count(&count).Error) + require.Zero(t, count) +} + func TestGenerateOAuthCodeCreatesProviderBoundOpaqueState(t *testing.T) { setupControllerAuthFlowDB(t) gin.SetMode(gin.TestMode) @@ -135,6 +195,32 @@ func TestGenerateOAuthCodeRejectsAnonymousBind(t *testing.T) { require.Equal(t, "Unauthorized", response.Message) } +func TestGenerateOAuthCodeRejectsTelegramBindIntent(t *testing.T) { + setupControllerAuthFlowDB(t) + require.NoError(t, appi18n.Init()) + originalEnabled := common.TelegramOAuthEnabled + common.TelegramOAuthEnabled = true + t.Cleanup(func() { common.TelegramOAuthEnabled = originalEnabled }) + + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/api/oauth/state", strings.NewReader(`{"provider":"telegram","intent":"bind"}`)) + c.Request.Header.Set("Content-Type", "application/json") + c.Set("id", 123) + + GenerateOAuthCode(c) + + require.Equal(t, http.StatusOK, recorder.Code) + var response struct { + Success bool `json:"success"` + } + require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response)) + require.False(t, response.Success) + var count int64 + require.NoError(t, model.DB.Model(&model.AuthFlow{}).Count(&count).Error) + require.Zero(t, count) +} + func TestHandleOAuthRejectsConsumedStateBeforeProviderExchange(t *testing.T) { setupControllerAuthFlowDB(t) require.NoError(t, appi18n.Init()) diff --git a/controller/passkey.go b/controller/passkey.go index d142ab8f..e1455fe3 100644 --- a/controller/passkey.go +++ b/controller/passkey.go @@ -36,10 +36,6 @@ func PasskeyRegisterBegin(c *gin.Context) { return } - if !requirePasskeyRegistrationVerification(c, user.Id) { - return - } - credential, err := model.GetPasskeyByUserID(user.Id) if err != nil && !errors.Is(err, model.ErrPasskeyNotFound) { common.ApiError(c, err) @@ -100,10 +96,6 @@ func PasskeyRegisterFinish(c *gin.Context) { return } - if !requirePasskeyRegistrationVerification(c, user.Id) { - return - } - wa, err := passkeysvc.BuildWebAuthn(c.Request) if err != nil { common.ApiError(c, err) @@ -138,10 +130,12 @@ func PasskeyRegisterFinish(c *gin.Context) { return } - if err := model.UpsertPasskeyCredential(passkeyCredential); err != nil { + generation, err := model.ReplacePasskeyCredentialAndBumpSessionGeneration(passkeyCredential) + if err != nil { common.ApiError(c, err) return } + preserveCurrentSessionAfterCommittedSecurityChange(c, user.Id, generation, "registering a passkey") recordUserSecurityAudit(c, user.Id, "user.passkey_register", nil) c.JSON(http.StatusOK, gin.H{ @@ -164,10 +158,12 @@ func PasskeyDelete(c *gin.Context) { return } - if err := model.DeletePasskeyByUserID(user.Id); err != nil { + generation, err := model.DeletePasskeyAndBumpSessionGeneration(user.Id) + if err != nil { common.ApiError(c, err) return } + preserveCurrentSessionAfterCommittedSecurityChange(c, user.Id, generation, "deleting a passkey") recordUserSecurityAudit(c, user.Id, "user.passkey_delete", nil) c.JSON(http.StatusOK, gin.H{ @@ -331,7 +327,7 @@ func PasskeyLoginFinish(c *gin.Context) { } now := time.Now() updatedCredential.LastUsedAt = &now - if err := model.UpsertPasskeyCredential(updatedCredential); err != nil { + if err := model.UpdatePasskeyCredentialAfterAuthentication(updatedCredential); err != nil { common.ApiError(c, err) return } @@ -370,7 +366,7 @@ func AdminResetPasskey(c *gin.Context) { return } - if err := model.DeletePasskeyByUserID(user.Id); err != nil { + if _, err := model.DeletePasskeyAndBumpSessionGeneration(user.Id); err != nil { common.ApiError(c, err) return } @@ -479,16 +475,21 @@ func PasskeyVerifyFinish(c *gin.Context) { } waUser := passkeysvc.NewWebAuthnUser(user, credential) - _, err = wa.FinishLogin(waUser, *sessionData, c.Request) + validatedCredential, err := wa.FinishLogin(waUser, *sessionData, c.Request) if err != nil { common.ApiError(c, err) return } + if validatedCredential == nil { + common.ApiErrorMsg(c, "Passkey 凭证更新失败") + return + } // 更新凭证的最后使用时间 + credential.ApplyValidatedCredential(validatedCredential) now := time.Now() credential.LastUsedAt = &now - if err := model.UpsertPasskeyCredential(credential); err != nil { + if err := model.UpdatePasskeyCredentialAfterAuthentication(credential); err != nil { common.ApiError(c, err) return } @@ -531,18 +532,6 @@ func getSessionUser(c *gin.Context) (*model.User, error) { return user, nil } -func requirePasskeyRegistrationVerification(c *gin.Context, userID int) bool { - twoFA, err := model.GetTwoFAByUserId(userID) - if err != nil { - common.ApiError(c, err) - return false - } - if twoFA == nil || !twoFA.IsEnabled { - return true - } - return requireSecureVerificationMethod(c, secureVerificationMethod2FA) -} - func requirePasskeyDeleteVerification(c *gin.Context, userID int) bool { twoFA, err := model.GetTwoFAByUserId(userID) if err != nil { @@ -550,7 +539,7 @@ func requirePasskeyDeleteVerification(c *gin.Context, userID int) bool { return false } if twoFA != nil && twoFA.IsEnabled { - return requireSecureVerificationMethod(c, secureVerificationMethod2FA) + return requireSecureVerificationMethod(c, model.SecureVerificationMethod2FA) } _, err = model.GetPasskeyByUserID(userID) @@ -566,7 +555,7 @@ func requirePasskeyDeleteVerification(c *gin.Context, userID int) bool { return false } - return requireSecureVerificationMethod(c, secureVerificationMethodPasskey) + return requireSecureVerificationMethod(c, model.SecureVerificationMethodPasskey) } func requireSecureVerificationMethod(c *gin.Context, method string) bool { diff --git a/controller/relay.go b/controller/relay.go index 5bc2486f..d05d87d8 100644 --- a/controller/relay.go +++ b/controller/relay.go @@ -67,6 +67,80 @@ func geminiRelayHandler(c *gin.Context, info *relaycommon.RelayInfo) *types.MaxA return err } +func prepareAlphaSearchPreConsumedQuota(priceData types.PriceData, relayInfo *relaycommon.RelayInfo) (types.PriceData, error) { + if relayInfo == nil || relayInfo.RelayMode != relayconstant.RelayModeAlphaSearch { + return priceData, nil + } + + baseQuota := priceData.QuotaToPreConsume + totalQuota, err := service.AlphaSearchPreConsumeQuota( + baseQuota, + relayInfo, + priceData.GroupRatioInfo.GroupRatio, + ) + if err != nil { + return types.PriceData{}, err + } + if totalQuota > baseQuota { + priceData.FreeModel = false + } + priceData.QuotaToPreConsume, err = helper.ApplyPreConsumedQuotaFloor( + totalQuota, + !priceData.FreeModel && priceData.GroupRatioInfo.GroupRatio > 0, + ) + if err != nil { + return types.PriceData{}, err + } + return priceData, nil +} + +func prepareBillingForSelectedGroup( + c *gin.Context, + relayInfo *relaycommon.RelayInfo, + promptTokens int, + meta *types.TokenCountMeta, +) *types.MaxAPIError { + if relayInfo.TieredBillingSnapshot != nil { + return service.PrepareTieredBillingForSelectedGroup(c, relayInfo) + } + + priceData, err := helper.ModelPriceHelper(c, relayInfo, promptTokens, meta) + if err != nil { + return types.NewErrorWithStatusCode( + err, + types.ErrorCodeModelPriceError, + http.StatusBadRequest, + types.ErrOptionWithSkipRetry(), + ) + } + priceData, err = prepareAlphaSearchPreConsumedQuota(priceData, relayInfo) + if err != nil { + return types.NewErrorWithStatusCode( + err, + types.ErrorCodeModelPriceError, + http.StatusBadRequest, + types.ErrOptionWithSkipRetry(), + ) + } + relayInfo.PriceData = priceData + if priceData.FreeModel { + // Keep any existing Billing session so a successful free retry can + // settle the prior reservation to zero, while preventing this free + // attempt from inheriting the paid attempt's fallback charge. + relayInfo.FinalPreConsumedQuota = 0 + return nil + } + + if relayInfo.Billing == nil { + return service.PreConsumeBilling(c, priceData.QuotaToPreConsume, relayInfo) + } + if err := relayInfo.Billing.Reserve(priceData.QuotaToPreConsume); err != nil { + return types.NewError(err, types.ErrorCodeUpdateDataError, types.ErrOptionWithSkipRetry()) + } + relayInfo.FinalPreConsumedQuota = relayInfo.Billing.GetPreConsumedQuota() + return nil +} + func Relay(c *gin.Context, relayFormat types.RelayFormat) { requestId := c.GetString(common.RequestIdKey) @@ -157,19 +231,11 @@ func Relay(c *gin.Context, relayFormat types.RelayFormat) { maxAPIError = types.NewError(err, types.ErrorCodeModelPriceError, types.ErrOptionWithStatusCode(http.StatusBadRequest)) return } - basePreConsumedQuota := priceData.QuotaToPreConsume - priceData.QuotaToPreConsume, err = service.AlphaSearchPreConsumeQuota( - basePreConsumedQuota, - relayInfo, - priceData.GroupRatioInfo.GroupRatio, - ) + priceData, err = prepareAlphaSearchPreConsumedQuota(priceData, relayInfo) if err != nil { maxAPIError = types.NewError(err, types.ErrorCodeModelPriceError, types.ErrOptionWithStatusCode(http.StatusBadRequest)) return } - if priceData.QuotaToPreConsume > basePreConsumedQuota { - priceData.FreeModel = false - } relayInfo.PriceData = priceData // common.SetContextKey(c, constant.ContextKeyTokenCountMeta, meta) @@ -211,7 +277,7 @@ func Relay(c *gin.Context, relayFormat types.RelayFormat) { } addUsedChannel(c, channel.Id) - if billingErr := service.PrepareTieredBillingForSelectedGroup(c, relayInfo); billingErr != nil { + if billingErr := prepareBillingForSelectedGroup(c, relayInfo, tokens, meta); billingErr != nil { maxAPIError = billingErr break } @@ -451,10 +517,9 @@ func RelayMidjourney(c *gin.Context) { //err = relayMidjourneySubmit(c, relayMode) log.Println(mjErr) if mjErr != nil { - statusCode := http.StatusBadRequest + statusCode := midjourneyRelayErrorStatusCode(mjErr) if mjErr.Code == 30 { mjErr.Result = "当前分组负载已饱和,请稍后再试,或升级账户以提升服务质量。" - statusCode = http.StatusTooManyRequests } c.JSON(statusCode, gin.H{ "description": fmt.Sprintf("%s %s", mjErr.Description, mjErr.Result), @@ -466,6 +531,19 @@ func RelayMidjourney(c *gin.Context) { } } +func midjourneyRelayErrorStatusCode(mjErr *dto.MidjourneyResponse) int { + if mjErr == nil { + return http.StatusBadRequest + } + if mjErr.Code == 30 { + return http.StatusTooManyRequests + } + if mjErr.Description == constant.MjBillingSettlementPending { + return http.StatusConflict + } + return http.StatusBadRequest +} + func RelayNotImplemented(c *gin.Context) { err := types.OpenAIError{ Message: "API not implemented", @@ -524,11 +602,12 @@ func RelayTask(c *gin.Context) { var result *relay.TaskSubmitResult var taskErr *dto.TaskError var settlementEffectOperationKey string + var settlementIntentPersisted bool defer func() { if taskErr != nil { upstreamAmbiguous := relayInfo.UpstreamTaskResponseReceived || relayInfo.UpstreamTaskOutcomeUnknown upstreamPersisted := result != nil && result.Task != nil - if relayInfo.Billing != nil && !upstreamAmbiguous { + if relayInfo.Billing != nil && !upstreamAmbiguous && !settlementIntentPersisted && taskErr.Code != taskResponseWriteFailedCode { relayInfo.Billing.Refund(c) } if relayInfo.PersistedTaskID > 0 { @@ -537,9 +616,11 @@ func RelayTask(c *gin.Context) { reason = taskErr.Error.Error() } if upstreamPersisted { - reason += "(上游已接受任务,本地结算记录未完成;已保留预扣费,请人工核对)" - if markErr := model.MarkTaskSubmitNeedsReview(result.Task, reason); markErr != nil { - common.SysError("mark task submit for manual review error: " + markErr.Error()) + if shouldMarkTaskSubmitNeedsReview(taskErr, upstreamPersisted, settlementIntentPersisted) { + reason += "(上游已接受任务,本地结算记录未完成;已保留预扣费,请人工核对)" + if markErr := model.MarkTaskSubmitNeedsReview(result.Task, reason); markErr != nil { + common.SysError("mark task submit for manual review error: " + markErr.Error()) + } } } else if upstreamAmbiguous { if relayInfo.UpstreamTaskOutcomeUnknown { @@ -652,23 +733,15 @@ func RelayTask(c *gin.Context) { } if err = result.Task.UpdateWithSettlementIntent(settlementIntent); err != nil { taskErr = relay.TaskPersistenceError(err, "persist_task_result_failed", "failed to persist task result") + } else { + settlementIntentPersisted = settlementIntent != nil } } } if taskErr == nil { - if settleErr := service.SettleBilling(c, relayInfo, result.Quota); settleErr != nil { - common.SysError("settle task billing remains pending/manual: " + settleErr.Error()) - } - if settlementEffectOperationKey != "" { - if effectErr := model.ProcessBillingSettlementEffect(settlementEffectOperationKey); effectErr != nil { - common.SysError("task billing effect remains pending/manual: " + effectErr.Error()) - } - } else { - service.LogTaskConsumption(c, relayInfo) - } - if writeErr := result.WriteResponse(c); writeErr != nil { - common.SysError("write task response error: " + writeErr.Error()) - } + taskErr = finalizeTaskSubmission(c, relayInfo, result, settlementEffectOperationKey, func() error { + return result.WriteResponse(c) + }) } } @@ -677,7 +750,65 @@ func RelayTask(c *gin.Context) { } } -// respondTaskError 统一输出 Task 错误响应(含 429 限流提示改写) +const taskResponseWriteFailedCode = "write_task_response_failed" + +// shouldMarkTaskSubmitNeedsReview keeps an upstream-accepted task pollable +// when its durable finalize intent already exists. The settlement runner owns +// recovery in that state; converting the task to FAILURE would permanently +// remove it from normal polling. +func shouldMarkTaskSubmitNeedsReview(taskErr *dto.TaskError, upstreamPersisted, settlementIntentPersisted bool) bool { + if taskErr == nil { + return false + } + if taskErr.Code == taskResponseWriteFailedCode { + return false + } + return taskErr.Code != constant.MjBillingSettlementPending || !upstreamPersisted || !settlementIntentPersisted +} + +// finalizeTaskSubmission releases a task success response only after its +// billing settlement succeeds. The upstream task may already exist when this +// fails, so callers persist it for reconciliation and return a local no-retry +// error instead of pretending the request completed normally. +func finalizeTaskSubmission(c *gin.Context, relayInfo *relaycommon.RelayInfo, result *relay.TaskSubmitResult, settlementEffectOperationKey string, writeResponse func() error) *dto.TaskError { + if result == nil || result.Task == nil { + return service.TaskErrorWrapperLocal(errors.New("task submit result is incomplete"), "persist_task_failed", http.StatusInternalServerError) + } + if settleErr := service.SettleBilling(c, relayInfo, result.Quota); settleErr != nil { + common.SysError("settle task billing requires reconciliation: " + settleErr.Error()) + return &dto.TaskError{ + Code: constant.MjBillingSettlementPending, + Message: "任务已被上游接受,但计费结算未完成;请勿重复提交,可使用 task_id 查询任务或联系管理员处理", + Data: map[string]string{"task_id": result.Task.TaskID}, + StatusCode: http.StatusConflict, + LocalError: true, + Error: settleErr, + } + } + if settlementEffectOperationKey != "" { + if effectErr := model.ProcessBillingSettlementEffect(settlementEffectOperationKey); effectErr != nil { + common.SysError("task billing effect remains pending/manual: " + effectErr.Error()) + } + } else { + service.LogTaskConsumption(c, relayInfo) + } + if writeResponse != nil { + if writeErr := writeResponse(); writeErr != nil { + common.SysError("write task response error: " + writeErr.Error()) + return &dto.TaskError{ + Code: taskResponseWriteFailedCode, + Message: "任务已提交并完成结算,但响应写入失败;请勿重复提交,可使用 task_id 查询任务", + Data: map[string]string{"task_id": result.Task.TaskID}, + StatusCode: http.StatusConflict, + LocalError: true, + Error: writeErr, + } + } + } + return nil +} + +// respondTaskError writes the task error response and normalizes the 429 message. func respondTaskError(c *gin.Context, taskErr *dto.TaskError) { if taskErr.StatusCode == http.StatusTooManyRequests { taskErr.Message = "当前分组上游负载已饱和,请稍后再试" diff --git a/controller/secure_verification.go b/controller/secure_verification.go index ffbb2952..3daccec8 100644 --- a/controller/secure_verification.go +++ b/controller/secure_verification.go @@ -14,15 +14,10 @@ import ( const ( // SecureVerificationSessionKey means the user has fully passed secure verification. - SecureVerificationSessionKey = "secure_verified_at" - secureVerificationMethodSessionKey = "secure_verified_method" - secureVerificationUserSessionKey = "secure_verified_user_id" - secureVerificationScopeSessionKey = "secure_verified_scope" - secureVerificationMethod2FA = "2fa" - secureVerificationMethodPasskey = "passkey" - secureVerificationMethodPassword = "password" - secureVerificationScopeAccessToken = "access_token" - secureVerificationScopeAccountDelete = "account_delete" + SecureVerificationSessionKey = "secure_verified_at" + secureVerificationMethodSessionKey = "secure_verified_method" + secureVerificationUserSessionKey = "secure_verified_user_id" + secureVerificationScopeSessionKey = "secure_verified_scope" // PasskeyReadySessionKey means WebAuthn finished and /api/verify can finalize step-up verification. PasskeyReadySessionKey = "secure_passkey_ready_at" // SecureVerificationTimeout 验证有效期(秒) @@ -161,7 +156,7 @@ func UniversalVerify(c *gin.Context) { var verifyMethod string switch req.Method { - case "2fa": + case model.SecureVerificationMethod2FA: if !methods.has2FA { common.ApiError(c, fmt.Errorf("用户未启用2FA")) return @@ -173,7 +168,7 @@ func UniversalVerify(c *gin.Context) { verified = validateTwoFactorAuth(methods.twoFA, req.Code) verifyMethod = "2FA" - case "passkey": + case model.SecureVerificationMethodPasskey: if !methods.hasPasskey { common.ApiError(c, fmt.Errorf("用户未启用Passkey")) return @@ -190,7 +185,7 @@ func UniversalVerify(c *gin.Context) { } verifyMethod = "Passkey" - case "password": + case model.SecureVerificationMethodPassword: if !methods.hasPassword { common.ApiError(c, fmt.Errorf("密码验证不可用,请使用2FA或Passkey")) return @@ -208,7 +203,7 @@ func UniversalVerify(c *gin.Context) { } if !verified { - if req.Method == secureVerificationMethodPassword { + if req.Method == model.SecureVerificationMethodPassword { common.ApiError(c, fmt.Errorf("密码错误")) } else { common.ApiError(c, fmt.Errorf("验证失败,请检查验证码")) @@ -238,7 +233,9 @@ func UniversalVerify(c *gin.Context) { func isSupportedSecureVerificationScope(scope string) bool { switch scope { - case "", secureVerificationScopeAccessToken, secureVerificationScopeAccountDelete: + case "", model.SecureVerificationScopeAccessToken, model.SecureVerificationScopeAccountDelete, + model.SecureVerificationScopeCredentials, model.SecureVerificationScopeAPIToken, + model.SecureVerificationScopePasskeyRegister: return true default: return false @@ -246,7 +243,11 @@ func isSupportedSecureVerificationScope(scope string) bool { } func passwordVerificationAllowed(scope string) bool { - return scope == secureVerificationScopeAccessToken || scope == secureVerificationScopeAccountDelete + return scope == model.SecureVerificationScopeAccessToken || + scope == model.SecureVerificationScopeAccountDelete || + scope == model.SecureVerificationScopeCredentials || + scope == model.SecureVerificationScopeAPIToken || + scope == model.SecureVerificationScopePasskeyRegister } func setSecureVerificationSession(c *gin.Context, userId int, method string, scope string) (int64, error) { diff --git a/controller/secure_verification_test.go b/controller/secure_verification_test.go index c9446839..946cb6e6 100644 --- a/controller/secure_verification_test.go +++ b/controller/secure_verification_test.go @@ -2,15 +2,36 @@ package controller import ( "bytes" + "crypto/hmac" + "crypto/sha256" + "encoding/hex" + "errors" "net/http" "net/http/httptest" + "net/url" + "sort" + "strconv" + "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" + appi18n "github.com/MAX-API-Next/MAX-API/i18n" "github.com/MAX-API-Next/MAX-API/model" + "github.com/MAX-API-Next/MAX-API/relay" + relaycommon "github.com/MAX-API-Next/MAX-API/relay/common" + relayconstant "github.com/MAX-API-Next/MAX-API/relay/constant" + "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-contrib/sessions" "github.com/gin-contrib/sessions/cookie" "github.com/gin-gonic/gin" + "github.com/pquerna/otp/totp" + "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) @@ -114,10 +135,17 @@ func TestUniversalVerifyPasswordRequiresAccessTokenScope(t *testing.T) { func TestSecureVerificationScopeAllowlist(t *testing.T) { require.True(t, isSupportedSecureVerificationScope("")) - require.True(t, isSupportedSecureVerificationScope(secureVerificationScopeAccessToken)) - require.True(t, isSupportedSecureVerificationScope(secureVerificationScopeAccountDelete)) + require.True(t, isSupportedSecureVerificationScope(model.SecureVerificationScopeAccessToken)) + require.True(t, isSupportedSecureVerificationScope(model.SecureVerificationScopeAccountDelete)) + require.True(t, isSupportedSecureVerificationScope(model.SecureVerificationScopeCredentials)) + require.True(t, isSupportedSecureVerificationScope(model.SecureVerificationScopeAPIToken)) + require.True(t, isSupportedSecureVerificationScope(model.SecureVerificationScopePasskeyRegister)) + require.False(t, isSupportedSecureVerificationScope(model.SecureVerificationScopeOAuthReauthentication)) require.False(t, isSupportedSecureVerificationScope("admin_delete")) - require.True(t, passwordVerificationAllowed(secureVerificationScopeAccountDelete)) + require.True(t, passwordVerificationAllowed(model.SecureVerificationScopeAccountDelete)) + require.True(t, passwordVerificationAllowed(model.SecureVerificationScopeCredentials)) + require.True(t, passwordVerificationAllowed(model.SecureVerificationScopeAPIToken)) + require.True(t, passwordVerificationAllowed(model.SecureVerificationScopePasskeyRegister)) require.False(t, passwordVerificationAllowed("")) } @@ -182,12 +210,12 @@ func TestSetupLoginClearsPreviousSecureVerification(t *testing.T) { router := gin.New() router.Use(sessions.Sessions("session", cookie.NewStore([]byte("login-clears-verification-test")))) - router.GET("/login", func(c *gin.Context) { + router.GET("/api/user/login", func(c *gin.Context) { session := sessions.Default(c) session.Set(SecureVerificationSessionKey, int64(123)) - session.Set(secureVerificationMethodSessionKey, secureVerificationMethod2FA) + session.Set(secureVerificationMethodSessionKey, model.SecureVerificationMethod2FA) session.Set(secureVerificationUserSessionKey, 999) - session.Set(secureVerificationScopeSessionKey, secureVerificationScopeAccessToken) + session.Set(secureVerificationScopeSessionKey, model.SecureVerificationScopeAccessToken) session.Set(PasskeyReadySessionKey, int64(123)) setupLogin(&user, c) }) @@ -203,7 +231,7 @@ func TestSetupLoginClearsPreviousSecureVerification(t *testing.T) { }) loginRecorder := httptest.NewRecorder() - router.ServeHTTP(loginRecorder, httptest.NewRequest(http.MethodGet, "/login", nil)) + router.ServeHTTP(loginRecorder, httptest.NewRequest(http.MethodGet, "/api/user/login", nil)) require.Equal(t, http.StatusOK, loginRecorder.Code) inspectRecorder := httptest.NewRecorder() @@ -222,3 +250,989 @@ func TestSetupLoginClearsPreviousSecureVerification(t *testing.T) { "passkey_ready_at": null }`, inspectRecorder.Body.String()) } + +func TestOAuthLoginCreatesNarrowReauthenticationScope(t *testing.T) { + db := setupUserSettingControllerTestDB(t) + require.NoError(t, db.AutoMigrate(&model.Log{}, &model.TwoFA{}, &model.PasskeyCredential{})) + + user := model.User{ + Id: 1004, + Username: "oauth-reauth-user", + DisplayName: "OAuth Reauth User", + Role: common.RoleCommonUser, + Status: common.UserStatusEnabled, + Group: "default", + } + require.NoError(t, db.Create(&user).Error) + + router := gin.New() + router.Use(sessions.Sessions("session", cookie.NewStore([]byte("oauth-credential-verification-test")))) + router.GET("/api/oauth/:provider", func(c *gin.Context) { + setupLogin(&user, c) + }) + router.GET("/verification-session", func(c *gin.Context) { + session := sessions.Default(c) + c.JSON(http.StatusOK, gin.H{ + "verified_at": session.Get(SecureVerificationSessionKey), + "verified_method": session.Get(secureVerificationMethodSessionKey), + "verified_user_id": session.Get(secureVerificationUserSessionKey), + "verified_scope": session.Get(secureVerificationScopeSessionKey), + }) + }) + + loginRecorder := httptest.NewRecorder() + router.ServeHTTP(loginRecorder, httptest.NewRequest(http.MethodGet, "/api/oauth/github", nil)) + require.Equal(t, http.StatusOK, loginRecorder.Code) + + inspectRecorder := httptest.NewRecorder() + inspectRequest := httptest.NewRequest(http.MethodGet, "/verification-session", nil) + for _, sessionCookie := range loginRecorder.Result().Cookies() { + inspectRequest.AddCookie(sessionCookie) + } + router.ServeHTTP(inspectRecorder, inspectRequest) + + require.Equal(t, http.StatusOK, inspectRecorder.Code) + var verificationSession struct { + VerifiedAt int64 `json:"verified_at"` + VerifiedMethod string `json:"verified_method"` + VerifiedUserID int `json:"verified_user_id"` + VerifiedScope string `json:"verified_scope"` + } + require.NoError(t, common.Unmarshal(inspectRecorder.Body.Bytes(), &verificationSession)) + require.WithinDuration(t, time.Now(), time.Unix(verificationSession.VerifiedAt, 0), 2*time.Second) + require.Equal(t, model.SecureVerificationMethodOAuth, verificationSession.VerifiedMethod) + require.Equal(t, user.Id, verificationSession.VerifiedUserID) + require.Equal(t, model.SecureVerificationScopeOAuthReauthentication, verificationSession.VerifiedScope) +} + +func TestOAuthLoginDoesNotGrantReauthenticationWhenCredentialsExist(t *testing.T) { + db := setupUserSettingControllerTestDB(t) + require.NoError(t, db.AutoMigrate(&model.Log{}, &model.TwoFA{}, &model.PasskeyCredential{})) + + user := model.User{ + Id: 1006, + Username: "oauth-existing-credential-user", + Password: "existing-password-hash", + DisplayName: "OAuth Existing Credential User", + Role: common.RoleCommonUser, + Status: common.UserStatusEnabled, + Group: "default", + } + require.NoError(t, db.Create(&user).Error) + + router := gin.New() + router.Use(sessions.Sessions("session", cookie.NewStore([]byte("oauth-existing-credential-test")))) + router.GET("/api/oauth/:provider", func(c *gin.Context) { + setupLogin(&user, c) + }) + router.GET("/verification-session", func(c *gin.Context) { + session := sessions.Default(c) + c.JSON(http.StatusOK, gin.H{ + "verified_at": session.Get(SecureVerificationSessionKey), + "verified_method": session.Get(secureVerificationMethodSessionKey), + "verified_scope": session.Get(secureVerificationScopeSessionKey), + }) + }) + + loginRecorder := httptest.NewRecorder() + router.ServeHTTP(loginRecorder, httptest.NewRequest(http.MethodGet, "/api/oauth/github", nil)) + require.Equal(t, http.StatusOK, loginRecorder.Code) + + inspectRecorder := httptest.NewRecorder() + inspectRequest := httptest.NewRequest(http.MethodGet, "/verification-session", nil) + for _, sessionCookie := range loginRecorder.Result().Cookies() { + inspectRequest.AddCookie(sessionCookie) + } + router.ServeHTTP(inspectRecorder, inspectRequest) + + require.Equal(t, http.StatusOK, inspectRecorder.Code) + require.JSONEq(t, `{"verified_at":null,"verified_method":null,"verified_scope":null}`, inspectRecorder.Body.String()) +} + +func TestTelegramLoginDoesNotCreateSecureVerificationGrant(t *testing.T) { + db := setupUserSettingControllerTestDB(t) + require.NoError(t, db.AutoMigrate(&model.Log{})) + user := model.User{ + Id: 1005, Username: "telegram-no-step-up-user", Role: common.RoleCommonUser, + Status: common.UserStatusEnabled, Group: "default", + } + require.NoError(t, db.Create(&user).Error) + + router := gin.New() + router.Use(sessions.Sessions("session", cookie.NewStore([]byte("telegram-no-step-up-test")))) + router.GET("/api/oauth/telegram/login", func(c *gin.Context) { setupLogin(&user, c) }) + router.GET("/verification-session", func(c *gin.Context) { + session := sessions.Default(c) + c.JSON(http.StatusOK, gin.H{ + "verified_at": session.Get(SecureVerificationSessionKey), + "verified_method": session.Get(secureVerificationMethodSessionKey), + "verified_scope": session.Get(secureVerificationScopeSessionKey), + }) + }) + + loginRecorder := httptest.NewRecorder() + router.ServeHTTP(loginRecorder, httptest.NewRequest(http.MethodGet, "/api/oauth/telegram/login", nil)) + require.Equal(t, http.StatusOK, loginRecorder.Code) + inspectRecorder := httptest.NewRecorder() + inspectRequest := httptest.NewRequest(http.MethodGet, "/verification-session", nil) + for _, sessionCookie := range loginRecorder.Result().Cookies() { + inspectRequest.AddCookie(sessionCookie) + } + router.ServeHTTP(inspectRecorder, inspectRequest) + require.JSONEq(t, `{"verified_at":null,"verified_method":null,"verified_scope":null}`, inspectRecorder.Body.String()) +} + +func TestResetPasswordReportsRevokedApiTokens(t *testing.T) { + require.NoError(t, appi18n.Init()) + db := setupUserSettingControllerTestDB(t) + require.NoError(t, db.AutoMigrate(&model.Token{})) + user := model.User{ + Id: 9910, + Username: "password-recovery-notice-user", + Email: "password-recovery-notice@example.com", + Role: common.RoleCommonUser, + Status: common.UserStatusEnabled, + } + require.NoError(t, db.Create(&user).Error) + require.NoError(t, db.Create(&model.Token{ + Id: 99101, UserId: user.Id, Name: "recovery-token", Key: "recovery-token-key", Status: common.TokenStatusEnabled, + }).Error) + + code := "password-recovery-notice-code" + common.RegisterVerificationCodeWithKey(user.Email, code, common.PasswordResetPurpose) + t.Cleanup(func() { common.DeleteKey(user.Email, common.PasswordResetPurpose) }) + payload, err := common.Marshal(PasswordResetRequest{Email: user.Email, Token: code}) + require.NoError(t, err) + recorder := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(recorder) + ctx.Request = httptest.NewRequest(http.MethodPost, "/api/user/reset", bytes.NewReader(payload)) + ctx.Request.Header.Set("Content-Type", "application/json") + ctx.Request.Header.Set("Accept-Language", "zh-CN") + + ResetPassword(ctx) + + require.Equal(t, http.StatusOK, recorder.Code) + var response struct { + Success bool `json:"success"` + Message string `json:"message"` + Data string `json:"data"` + ApiTokensRevoked bool `json:"api_tokens_revoked"` + } + require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response)) + require.True(t, response.Success) + require.NotEmpty(t, response.Data) + require.True(t, response.ApiTokensRevoked) + require.Contains(t, response.Message, "所有现有 API 令牌已被撤销") + require.NotContains(t, response.Message, "如有需要") + + var token model.Token + require.NoError(t, db.First(&token, 99101).Error) + require.Equal(t, common.TokenStatusDisabled, token.Status) +} + +func TestPreserveCurrentSessionAfterSecurityChangeUpdatesGenerationAndClearsStepUp(t *testing.T) { + gin.SetMode(gin.TestMode) + engine := gin.New() + engine.Use(sessions.Sessions("session", cookie.NewStore([]byte("preserve-current-session-test")))) + engine.GET("/seed", func(c *gin.Context) { + session := sessions.Default(c) + session.Set("id", 9911) + session.Set("session_generation", int64(1)) + session.Set(SecureVerificationSessionKey, time.Now().Unix()) + session.Set(secureVerificationMethodSessionKey, "password") + session.Set(secureVerificationUserSessionKey, 9911) + session.Set(secureVerificationScopeSessionKey, model.SecureVerificationScopeCredentials) + session.Set(PasskeyReadySessionKey, time.Now().Unix()) + require.NoError(t, session.Save()) + c.Status(http.StatusNoContent) + }) + engine.POST("/preserve", func(c *gin.Context) { + require.NoError(t, preserveCurrentSessionAfterSecurityChange(c, 9911, 2)) + c.Status(http.StatusNoContent) + }) + engine.GET("/inspect", func(c *gin.Context) { + session := sessions.Default(c) + c.JSON(http.StatusOK, gin.H{ + "id": session.Get("id"), + "generation": session.Get("session_generation"), + "secure_verified_at": session.Get(SecureVerificationSessionKey), + "secure_method": session.Get(secureVerificationMethodSessionKey), + "secure_user": session.Get(secureVerificationUserSessionKey), + "secure_scope": session.Get(secureVerificationScopeSessionKey), + "passkey_ready_at": session.Get(PasskeyReadySessionKey), + }) + }) + + seedRecorder := httptest.NewRecorder() + engine.ServeHTTP(seedRecorder, httptest.NewRequest(http.MethodGet, "/seed", nil)) + require.Equal(t, http.StatusNoContent, seedRecorder.Code) + + preserveRequest := httptest.NewRequest(http.MethodPost, "/preserve", nil) + for _, sessionCookie := range seedRecorder.Result().Cookies() { + preserveRequest.AddCookie(sessionCookie) + } + preserveRecorder := httptest.NewRecorder() + engine.ServeHTTP(preserveRecorder, preserveRequest) + require.Equal(t, http.StatusNoContent, preserveRecorder.Code) + + inspectRequest := httptest.NewRequest(http.MethodGet, "/inspect", nil) + for _, sessionCookie := range preserveRecorder.Result().Cookies() { + inspectRequest.AddCookie(sessionCookie) + } + inspectRecorder := httptest.NewRecorder() + engine.ServeHTTP(inspectRecorder, inspectRequest) + require.Equal(t, http.StatusOK, inspectRecorder.Code) + require.JSONEq(t, `{"generation":2,"id":9911,"secure_verified_at":null,"secure_method":null,"secure_user":null,"secure_scope":null,"passkey_ready_at":null}`, inspectRecorder.Body.String()) +} + +func TestPasswordChangePreservesCommittedSessionGeneration(t *testing.T) { + db := setupUserSettingControllerTestDB(t) + hash, err := common.Password2Hash("old-password-123") + require.NoError(t, err) + user := model.User{ + Id: 9914, Username: "password-generation-session-user", Password: hash, + Role: common.RoleCommonUser, Status: common.UserStatusEnabled, SessionGeneration: 4, + } + require.NoError(t, db.Create(&user).Error) + + engine := gin.New() + engine.Use(sessions.Sessions("session", cookie.NewStore([]byte("password-generation-session-test")))) + engine.POST("/api/user/self", func(c *gin.Context) { + c.Set("id", user.Id) + UpdateSelf(c) + }) + engine.GET("/session", func(c *gin.Context) { + session := sessions.Default(c) + c.JSON(http.StatusOK, gin.H{ + "id": session.Get("id"), + "generation": session.Get("session_generation"), + }) + }) + + updateRequest := httptest.NewRequest( + http.MethodPost, + "/api/user/self", + bytes.NewBufferString(`{"original_password":"old-password-123","password":"new-password-123"}`), + ) + updateRequest.Header.Set("Content-Type", "application/json") + updateRecorder := httptest.NewRecorder() + engine.ServeHTTP(updateRecorder, updateRequest) + require.Equal(t, http.StatusOK, updateRecorder.Code) + + inspectRequest := httptest.NewRequest(http.MethodGet, "/session", nil) + for _, sessionCookie := range updateRecorder.Result().Cookies() { + inspectRequest.AddCookie(sessionCookie) + } + inspectRecorder := httptest.NewRecorder() + engine.ServeHTTP(inspectRecorder, inspectRequest) + require.Equal(t, http.StatusOK, inspectRecorder.Code) + require.JSONEq(t, `{"generation":5,"id":9914}`, inspectRecorder.Body.String()) + + require.NoError(t, db.First(&user, user.Id).Error) + require.EqualValues(t, 5, user.SessionGeneration) + require.True(t, common.ValidatePasswordAndHash("new-password-123", user.Password)) +} + +func TestRevokeOtherSessionsReturnsCommittedSuccessWhenSessionPreservationFails(t *testing.T) { + db := setupUserSettingControllerTestDB(t) + require.NoError(t, db.AutoMigrate(&model.Log{})) + user := model.User{Id: 9912, Username: "revoke-session-user", Role: common.RoleCommonUser, Status: common.UserStatusEnabled} + require.NoError(t, db.Create(&user).Error) + + recorder := performUserSecurityChangeWithFailingSession(t, user.Id, `{}`, RevokeOtherSessions) + + require.Equal(t, http.StatusOK, recorder.Code) + require.Contains(t, recorder.Body.String(), `"success":true`) + require.NoError(t, db.First(&user, user.Id).Error) + require.EqualValues(t, 1, user.SessionGeneration) +} + +func TestPasswordChangeReturnsCommittedSuccessWhenSessionPreservationFails(t *testing.T) { + db := setupUserSettingControllerTestDB(t) + hash, err := common.Password2Hash("old-password-123") + require.NoError(t, err) + user := model.User{ + Id: 9913, Username: "password-session-user", Password: hash, + Role: common.RoleCommonUser, Status: common.UserStatusEnabled, + } + require.NoError(t, db.Create(&user).Error) + + recorder := performUserSecurityChangeWithFailingSession( + t, + user.Id, + `{"original_password":"old-password-123","password":"new-password-123"}`, + UpdateSelf, + ) + + require.Equal(t, http.StatusOK, recorder.Code) + require.Contains(t, recorder.Body.String(), `"success":true`) + require.NoError(t, db.First(&user, user.Id).Error) + require.True(t, common.ValidatePasswordAndHash("new-password-123", user.Password)) + require.EqualValues(t, 1, user.SessionGeneration) +} + +func performUserSecurityChangeWithFailingSession( + t *testing.T, + userID int, + body string, + handler gin.HandlerFunc, +) *httptest.ResponseRecorder { + t.Helper() + router := gin.New() + router.Use(func(c *gin.Context) { + c.Set(sessions.DefaultKey, &failingSaveSession{values: map[interface{}]interface{}{}}) + c.Next() + }) + router.POST("/security-change", func(c *gin.Context) { + c.Set("id", userID) + handler(c) + }) + recorder := httptest.NewRecorder() + request := httptest.NewRequest(http.MethodPost, "/security-change", bytes.NewBufferString(body)) + request.Header.Set("Content-Type", "application/json") + router.ServeHTTP(recorder, request) + return recorder +} + +type failingTaskBillingSettler struct { + err error +} + +func (s *failingTaskBillingSettler) Settle(int) error { return s.err } +func (s *failingTaskBillingSettler) Refund(*gin.Context) {} +func (s *failingTaskBillingSettler) NeedsRefund() bool { return true } +func (s *failingTaskBillingSettler) GetPreConsumedQuota() int { return 10 } +func (s *failingTaskBillingSettler) Reserve(int) error { return nil } + +type successfulTaskBillingSettler struct { + settleCalls int + refundCalls int +} + +func (s *successfulTaskBillingSettler) Settle(int) error { + s.settleCalls++ + return nil +} +func (s *successfulTaskBillingSettler) Refund(*gin.Context) { s.refundCalls++ } +func (s *successfulTaskBillingSettler) NeedsRefund() bool { return true } +func (s *successfulTaskBillingSettler) GetPreConsumedQuota() int { return 10 } +func (s *successfulTaskBillingSettler) Reserve(int) error { return nil } + +type recordingSelectedGroupBillingSettler struct { + preConsumed int + reserves []int +} + +func (s *recordingSelectedGroupBillingSettler) Settle(int) error { return nil } +func (s *recordingSelectedGroupBillingSettler) Refund(*gin.Context) {} +func (s *recordingSelectedGroupBillingSettler) NeedsRefund() bool { return true } +func (s *recordingSelectedGroupBillingSettler) GetPreConsumedQuota() int { + return s.preConsumed +} +func (s *recordingSelectedGroupBillingSettler) Reserve(targetQuota int) error { + s.reserves = append(s.reserves, targetQuota) + if targetQuota > s.preConsumed { + s.preConsumed = targetQuota + } + return nil +} + +func TestFinalizeTaskSubmissionDoesNotWriteSuccessAfterSettlementFailure(t *testing.T) { + gin.SetMode(gin.TestMode) + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + settlementErr := errors.New("user quota is not enough") + info := &relaycommon.RelayInfo{ + UserId: 71, + TokenId: 72, + Billing: &failingTaskBillingSettler{err: settlementErr}, + } + result := &relay.TaskSubmitResult{ + Quota: 20, + Task: &model.Task{TaskID: "task-settlement-review"}, + } + writeCount := 0 + + taskErr := finalizeTaskSubmission(c, info, result, "", func() error { + writeCount++ + return nil + }) + + require.NotNil(t, taskErr) + assert.Equal(t, constant.MjBillingSettlementPending, taskErr.Code) + assert.Equal(t, http.StatusConflict, taskErr.StatusCode) + assert.True(t, taskErr.LocalError) + assert.Equal(t, map[string]string{"task_id": "task-settlement-review"}, taskErr.Data) + assert.Zero(t, writeCount) +} + +func TestPendingTaskSettlementDoesNotMarkAcceptedTaskAsManualFailure(t *testing.T) { + taskErr := &dto.TaskError{Code: constant.MjBillingSettlementPending} + + assert.False(t, shouldMarkTaskSubmitNeedsReview(taskErr, true, true)) + assert.True(t, shouldMarkTaskSubmitNeedsReview(taskErr, false, true)) + assert.True(t, shouldMarkTaskSubmitNeedsReview(taskErr, true, false)) + assert.True(t, shouldMarkTaskSubmitNeedsReview(&dto.TaskError{Code: "persist_task_failed"}, true, true)) +} + +func TestFinalizeTaskSubmissionReportsCommittedResponseWriteFailure(t *testing.T) { + setupUserSettingControllerTestDB(t) + gin.SetMode(gin.TestMode) + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + settler := &successfulTaskBillingSettler{} + info := &relaycommon.RelayInfo{ + UserId: 81, + TokenId: 82, + Billing: settler, + ChannelMeta: &relaycommon.ChannelMeta{}, + TaskRelayInfo: &relaycommon.TaskRelayInfo{Action: "submit"}, + } + result := &relay.TaskSubmitResult{ + Quota: 20, + Task: &model.Task{TaskID: "task-response-write-failed"}, + } + + oldLogConsumeEnabled := common.LogConsumeEnabled + common.LogConsumeEnabled = false + t.Cleanup(func() { common.LogConsumeEnabled = oldLogConsumeEnabled }) + + writeErr := errors.New("client connection closed") + taskErr := finalizeTaskSubmission(c, info, result, "", func() error { + return writeErr + }) + + require.NotNil(t, taskErr) + assert.Equal(t, "write_task_response_failed", taskErr.Code) + assert.Equal(t, http.StatusConflict, taskErr.StatusCode) + assert.True(t, taskErr.LocalError) + assert.ErrorIs(t, taskErr.Error, writeErr) + assert.Equal(t, map[string]string{"task_id": "task-response-write-failed"}, taskErr.Data) + assert.False(t, shouldRetryTaskRelay(c, 1, taskErr, 1)) + assert.False(t, shouldMarkTaskSubmitNeedsReview(taskErr, true, true)) + assert.Equal(t, 1, settler.settleCalls) + assert.Zero(t, settler.refundCalls) + + respondTaskError(c, taskErr) + assert.Equal(t, http.StatusConflict, recorder.Code) + assert.Contains(t, recorder.Body.String(), `"code":"write_task_response_failed"`) + assert.Contains(t, recorder.Body.String(), `"task_id":"task-response-write-failed"`) +} + +func TestMidjourneyRelayErrorStatusCodeReportsSettlementPendingAsConflict(t *testing.T) { + assert.Equal(t, http.StatusConflict, midjourneyRelayErrorStatusCode(&dto.MidjourneyResponse{ + Code: constant.MjRequestError, + Description: constant.MjBillingSettlementPending, + })) + assert.Equal(t, http.StatusTooManyRequests, midjourneyRelayErrorStatusCode(&dto.MidjourneyResponse{Code: 30})) + assert.Equal(t, http.StatusBadRequest, midjourneyRelayErrorStatusCode(&dto.MidjourneyResponse{Code: constant.MjRequestError})) +} + +func TestPrepareAlphaSearchPreConsumedQuotaAppliesFloorAfterSurcharge(t *testing.T) { + originalQuotaPerUnit := common.QuotaPerUnit + originalPreConsumedQuota := common.PreConsumedQuota + common.QuotaPerUnit = 1_000 + common.PreConsumedQuota = 500 + + toolPrices := config.GlobalConfig.Get("tool_price_setting").(*operation_setting.ToolPriceSetting) + originalPrices := make(map[string]float64, len(toolPrices.Prices)) + for key, value := range toolPrices.Prices { + originalPrices[key] = value + } + toolPrices.Prices = map[string]float64{"web_search_preview": 10} + operation_setting.RebuildToolPriceIndex() + t.Cleanup(func() { + common.QuotaPerUnit = originalQuotaPerUnit + common.PreConsumedQuota = originalPreConsumedQuota + toolPrices.Prices = originalPrices + operation_setting.RebuildToolPriceIndex() + }) + + info := &relaycommon.RelayInfo{ + RelayMode: relayconstant.RelayModeAlphaSearch, + OriginModelName: "alpha-floor-test", + } + + for _, test := range []struct { + name string + baseQuota int + freeModel bool + expected int + expectedFree bool + }{ + {name: "total remains below floor", baseQuota: 200, expected: 500}, + {name: "surcharge crosses floor", baseQuota: 495, expected: 505}, + {name: "tool surcharge makes free model billable", freeModel: true, expected: 500}, + } { + t.Run(test.name, func(t *testing.T) { + priceData, err := prepareAlphaSearchPreConsumedQuota(types.PriceData{ + FreeModel: test.freeModel, + QuotaToPreConsume: test.baseQuota, + GroupRatioInfo: types.GroupRatioInfo{GroupRatio: 1}, + }, info) + + require.NoError(t, err) + require.Equal(t, test.expected, priceData.QuotaToPreConsume) + require.Equal(t, test.expectedFree, priceData.FreeModel) + }) + } +} + +func TestPrepareBillingForSelectedGroupRepricesAlphaSearchCrossGroupRetry(t *testing.T) { + const modelName = "alpha-cross-group-retry" + originalQuotaPerUnit := common.QuotaPerUnit + originalPreConsumedQuota := common.PreConsumedQuota + originalModelRatio := ratio_setting.ModelRatio2JSONString() + originalGroupRatio := ratio_setting.GroupRatio2JSONString() + common.QuotaPerUnit = 1_000 + common.PreConsumedQuota = 2_500 + require.NoError(t, ratio_setting.UpdateModelRatioByJSONString(`{"alpha-cross-group-retry":1}`)) + require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"initial":1,"retry":3}`)) + + toolPrices := config.GlobalConfig.Get("tool_price_setting").(*operation_setting.ToolPriceSetting) + originalPrices := make(map[string]float64, len(toolPrices.Prices)) + for key, value := range toolPrices.Prices { + originalPrices[key] = value + } + toolPrices.Prices = map[string]float64{"web_search_preview": 10} + operation_setting.RebuildToolPriceIndex() + t.Cleanup(func() { + common.QuotaPerUnit = originalQuotaPerUnit + common.PreConsumedQuota = originalPreConsumedQuota + require.NoError(t, ratio_setting.UpdateModelRatioByJSONString(originalModelRatio)) + require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(originalGroupRatio)) + toolPrices.Prices = originalPrices + operation_setting.RebuildToolPriceIndex() + }) + + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + common.SetContextKey(c, constant.ContextKeyAutoGroup, "retry") + common.SetContextKey(c, constant.ContextKeyUsingGroup, "retry") + billing := &recordingSelectedGroupBillingSettler{preConsumed: 2_500} + info := &relaycommon.RelayInfo{ + RelayMode: relayconstant.RelayModeAlphaSearch, + OriginModelName: modelName, + UsingGroup: "initial", + Billing: billing, + FinalPreConsumedQuota: 2_500, + } + + require.Nil(t, prepareBillingForSelectedGroup(c, info, 1_000, &types.TokenCountMeta{})) + + require.Equal(t, "retry", info.UsingGroup) + require.Equal(t, 3.0, info.PriceData.GroupRatioInfo.GroupRatio) + require.Equal(t, 3_030, info.PriceData.QuotaToPreConsume) + require.Equal(t, []int{3_030}, billing.reserves) + require.Equal(t, 3_030, info.FinalPreConsumedQuota) +} + +func TestPrepareBillingForSelectedGroupClearsFallbackQuotaForFreeRetry(t *testing.T) { + const modelName = "free-cross-group-retry" + originalModelRatio := ratio_setting.ModelRatio2JSONString() + originalGroupRatio := ratio_setting.GroupRatio2JSONString() + quotaSetting := operation_setting.GetQuotaSetting() + originalFreeModelPreConsume := quotaSetting.EnableFreeModelPreConsume + quotaSetting.EnableFreeModelPreConsume = false + require.NoError(t, ratio_setting.UpdateModelRatioByJSONString(`{"free-cross-group-retry":1}`)) + require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"initial":1,"free":0}`)) + t.Cleanup(func() { + quotaSetting.EnableFreeModelPreConsume = originalFreeModelPreConsume + require.NoError(t, ratio_setting.UpdateModelRatioByJSONString(originalModelRatio)) + require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(originalGroupRatio)) + }) + + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + common.SetContextKey(c, constant.ContextKeyAutoGroup, "free") + common.SetContextKey(c, constant.ContextKeyUsingGroup, "free") + billing := &recordingSelectedGroupBillingSettler{preConsumed: 2_500} + info := &relaycommon.RelayInfo{ + OriginModelName: modelName, + UsingGroup: "initial", + Billing: billing, + FinalPreConsumedQuota: 2_500, + } + + require.Nil(t, prepareBillingForSelectedGroup(c, info, 1_000, &types.TokenCountMeta{})) + + require.Equal(t, "free", info.UsingGroup) + require.True(t, info.PriceData.FreeModel) + require.Zero(t, info.FinalPreConsumedQuota) + require.Same(t, billing, info.Billing) + require.Equal(t, 2_500, billing.GetPreConsumedQuota()) + require.Empty(t, billing.reserves) +} + +func signedTelegramParams(t *testing.T, token string, authDate time.Time) url.Values { + t.Helper() + params := url.Values{ + "id": {"123456"}, + "first_name": {"Security"}, + "auth_date": {strconv.FormatInt(authDate.Unix(), 10)}, + } + signTelegramParams(t, token, params) + return params +} + +func signTelegramParams(t *testing.T, token string, params url.Values) { + t.Helper() + params.Del("hash") + keys := make([]string, 0, len(params)) + for key := range params { + keys = append(keys, key) + } + sort.Strings(keys) + lines := make([]string, 0, len(keys)) + for _, key := range keys { + lines = append(lines, key+"="+params.Get(key)) + } + secret := sha256.Sum256([]byte(token)) + mac := hmac.New(sha256.New, secret[:]) + _, err := mac.Write([]byte(strings.Join(lines, "\n"))) + require.NoError(t, err) + params.Set("hash", hex.EncodeToString(mac.Sum(nil))) +} + +func TestTelegramAuthorizationRejectsStaleSignedPayload(t *testing.T) { + const token = "123456:telegram-test-token" + params := signedTelegramParams(t, token, time.Now().Add(-time.Hour)) + require.False(t, checkTelegramAuthorization(params, token)) +} + +func TestTelegramAuthorizationAcceptsFreshSignedPayload(t *testing.T) { + const token = "123456:telegram-test-token" + params := signedTelegramParams(t, token, time.Now().Add(-time.Minute)) + require.True(t, checkTelegramAuthorization(params, token)) +} + +func TestTelegramAuthorizationIgnoresLocalBindState(t *testing.T) { + const token = "123456:telegram-test-token" + now := time.Unix(1_800_000_000, 0) + params := signedTelegramParams(t, token, now.Add(-time.Minute)) + params.Set("state", "local-bind-state") + + require.True(t, checkTelegramAuthorizationAt(params, token, now)) +} + +func TestTelegramAuthorizationRejectsTamperedPayload(t *testing.T) { + const token = "123456:telegram-test-token" + now := time.Unix(1_800_000_000, 0) + params := signedTelegramParams(t, token, now.Add(-time.Minute)) + params.Set("id", "654321") + + require.False(t, checkTelegramAuthorizationAt(params, token, now)) +} + +func TestTelegramAuthorizationRejectsDuplicateParameterValues(t *testing.T) { + const token = "123456:telegram-test-token" + now := time.Unix(1_800_000_000, 0) + params := signedTelegramParams(t, token, now.Add(-time.Minute)) + params.Add("first_name", "Duplicate") + + require.False(t, checkTelegramAuthorizationAt(params, token, now)) +} + +func TestTelegramAuthorizationRejectsForeignBotToken(t *testing.T) { + const token = "123456:telegram-test-token" + now := time.Unix(1_800_000_000, 0) + params := signedTelegramParams(t, token, now.Add(-time.Minute)) + + require.False(t, checkTelegramAuthorizationAt( + params, + "987654:another-telegram-test-token", + now, + )) +} + +func TestTelegramAuthPayloadPreservesUnknownSignedFields(t *testing.T) { + const token = "123456:telegram-forward-compatible-test" + now := time.Unix(1_800_000_000, 0) + params := url.Values{ + "id": {"123456"}, + "first_name": {"Security"}, + "auth_date": {strconv.FormatInt(now.Add(-time.Minute).Unix(), 10)}, + "future_auth_field": {"future-value"}, + } + signTelegramParams(t, token, params) + + body, err := common.Marshal(map[string]any{ + "id": int64(123456), + "first_name": "Security", + "auth_date": now.Add(-time.Minute).Unix(), + "future_auth_field": "future-value", + "hash": params.Get("hash"), + "state": "local-bind-state", + }) + require.NoError(t, err) + + var payload telegramAuthPayload + require.NoError(t, common.DecodeJson(bytes.NewReader(body), &payload)) + values := payload.values() + require.Equal(t, "future-value", values.Get("future_auth_field")) + require.NotContains(t, values, "state") + require.Equal(t, "local-bind-state", payload.State) + require.True(t, checkTelegramAuthorizationAt(values, token, now)) +} + +func TestTelegramAuthorizationRejectsAuthDateBeyondClockSkew(t *testing.T) { + const token = "123456:telegram-test-token" + now := time.Unix(1_800_000_000, 0) + params := signedTelegramParams(t, token, now.Add(telegramClockSkew+time.Second)) + + require.False(t, checkTelegramAuthorizationAt(params, token, now)) +} + +func TestGenerateTelegramBindStateRejectsNonPositiveUserID(t *testing.T) { + setupControllerAuthFlowDB(t) + require.NoError(t, appi18n.Init()) + originalEnabled := common.TelegramOAuthEnabled + common.TelegramOAuthEnabled = true + t.Cleanup(func() { common.TelegramOAuthEnabled = originalEnabled }) + + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/api/oauth/telegram/bind/state", nil) + + GenerateTelegramBindState(c) + + require.Equal(t, http.StatusUnauthorized, recorder.Code) + var response struct { + Success bool `json:"success"` + } + require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response)) + require.False(t, response.Success) + var count int64 + require.NoError(t, model.DB.Model(&model.AuthFlow{}).Count(&count).Error) + require.Zero(t, count) +} + +func TestTelegramLoginRequiresBrowserBoundStateAndConsumesItOnce(t *testing.T) { + setupControllerAuthFlowDB(t) + require.NoError(t, appi18n.Init()) + gin.SetMode(gin.TestMode) + + const token = "123456:telegram-login-state-test" + originalEnabled := common.TelegramOAuthEnabled + originalToken := common.TelegramBotToken + common.TelegramOAuthEnabled = true + common.TelegramBotToken = token + t.Cleanup(func() { + common.TelegramOAuthEnabled = originalEnabled + common.TelegramBotToken = originalToken + }) + + user := model.User{ + Id: 99001, Username: "telegram-login-user", TelegramId: "123456", + Role: common.RoleCommonUser, Status: common.UserStatusEnabled, Group: "default", + } + require.NoError(t, model.DB.Create(&user).Error) + + router := gin.New() + router.Use(sessions.Sessions("session", cookie.NewStore([]byte("telegram-login-browser-state")))) + router.POST("/api/oauth/state", GenerateOAuthCode) + router.GET("/api/oauth/telegram/login", TelegramLogin) + + stateRecorder := httptest.NewRecorder() + stateRequest := httptest.NewRequest(http.MethodPost, "/api/oauth/state", strings.NewReader(`{"provider":"telegram","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) + + params := signedTelegramParams(t, token, time.Now()) + params.Set("state", stateResponse.Data) + callbackURL := "/api/oauth/telegram/login?" + params.Encode() + + crossBrowserRecorder := httptest.NewRecorder() + router.ServeHTTP(crossBrowserRecorder, httptest.NewRequest(http.MethodGet, callbackURL, nil)) + require.Equal(t, http.StatusForbidden, crossBrowserRecorder.Code) + + callbackRequest := httptest.NewRequest(http.MethodGet, callbackURL, nil) + for _, sessionCookie := range stateRecorder.Result().Cookies() { + callbackRequest.AddCookie(sessionCookie) + } + callbackRecorder := httptest.NewRecorder() + router.ServeHTTP(callbackRecorder, callbackRequest) + require.Equal(t, http.StatusOK, callbackRecorder.Code, callbackRecorder.Body.String()) + require.Contains(t, callbackRecorder.Body.String(), `"success":true`) + + replayRequest := httptest.NewRequest(http.MethodGet, callbackURL, nil) + for _, sessionCookie := range stateRecorder.Result().Cookies() { + replayRequest.AddCookie(sessionCookie) + } + replayRecorder := httptest.NewRecorder() + router.ServeHTTP(replayRecorder, replayRequest) + require.Equal(t, http.StatusForbidden, replayRecorder.Code) + _, err := model.GetAuthFlow(stateResponse.Data, telegramLoginFlowMatch) + require.ErrorIs(t, err, model.ErrAuthFlowConsumed) +} + +type failingSaveSession struct { + values map[interface{}]interface{} +} + +func (*failingSaveSession) ID() string { return "failing-save-session" } + +func (session *failingSaveSession) Get(key interface{}) interface{} { + return session.values[key] +} + +func (session *failingSaveSession) Set(key interface{}, value interface{}) { + session.values[key] = value +} + +func (session *failingSaveSession) Delete(key interface{}) { + delete(session.values, key) +} + +func (session *failingSaveSession) Clear() { + clear(session.values) +} + +func (*failingSaveSession) AddFlash(interface{}, ...string) {} + +func (*failingSaveSession) Flashes(...string) []interface{} { return nil } + +func (*failingSaveSession) Options(sessions.Options) {} + +func (*failingSaveSession) Save() error { + return errors.New("forced session save failure") +} + +func setupTwoFASessionFailureTest( + t *testing.T, + enabled bool, +) (*model.User, string) { + t.Helper() + + db := setupUserSettingControllerTestDB(t) + require.NoError(t, db.AutoMigrate( + &model.TwoFA{}, + &model.TwoFABackupCode{}, + &model.Log{}, + )) + + user := &model.User{ + Id: 71001, + Username: "twofa-session-save-user", + Role: common.RoleCommonUser, + Status: common.UserStatusEnabled, + Group: "default", + } + require.NoError(t, db.Create(user).Error) + + key, err := common.GenerateTOTPSecret(user.Username) + require.NoError(t, err) + require.NoError(t, db.Create(&model.TwoFA{ + UserId: user.Id, + Secret: key.Secret(), + IsEnabled: enabled, + }).Error) + if enabled { + require.NoError(t, model.CreateBackupCodes(user.Id, []string{"OLD1-CODE"})) + } + + code, err := totp.GenerateCode(key.Secret(), time.Now()) + require.NoError(t, err) + return user, code +} + +func performTwoFARequestWithFailingSessionSave( + t *testing.T, + userID int, + code string, + handler gin.HandlerFunc, +) *httptest.ResponseRecorder { + t.Helper() + + router := gin.New() + router.Use(func(c *gin.Context) { + c.Set(sessions.DefaultKey, &failingSaveSession{ + values: map[interface{}]interface{}{}, + }) + c.Next() + }) + router.POST("/twofa", func(c *gin.Context) { + c.Set("id", userID) + handler(c) + }) + + recorder := httptest.NewRecorder() + request := httptest.NewRequest( + http.MethodPost, + "/twofa", + bytes.NewReader([]byte(`{"code":"`+code+`"}`)), + ) + request.Header.Set("Content-Type", "application/json") + router.ServeHTTP(recorder, request) + return recorder +} + +func TestEnable2FAReturnsCommittedSuccessWhenSessionPreservationFails(t *testing.T) { + user, code := setupTwoFASessionFailureTest(t, false) + + recorder := performTwoFARequestWithFailingSessionSave(t, user.Id, code, Enable2FA) + + require.Equal(t, http.StatusOK, recorder.Code) + require.Contains(t, recorder.Body.String(), `"success":true`) + twoFA, err := model.GetTwoFAByUserId(user.Id) + require.NoError(t, err) + require.NotNil(t, twoFA) + require.True(t, twoFA.IsEnabled) +} + +func TestDisable2FAReturnsCommittedSuccessWhenSessionPreservationFails(t *testing.T) { + user, code := setupTwoFASessionFailureTest(t, true) + + recorder := performTwoFARequestWithFailingSessionSave(t, user.Id, code, Disable2FA) + + require.Equal(t, http.StatusOK, recorder.Code) + require.Contains(t, recorder.Body.String(), `"success":true`) + twoFA, err := model.GetTwoFAByUserId(user.Id) + require.NoError(t, err) + require.Nil(t, twoFA) +} + +func TestRegenerateBackupCodesReturnsCodesWhenSessionPreservationFails(t *testing.T) { + user, code := setupTwoFASessionFailureTest(t, true) + + recorder := performTwoFARequestWithFailingSessionSave( + t, + user.Id, + code, + RegenerateBackupCodes, + ) + + require.Equal(t, http.StatusOK, recorder.Code) + require.Contains(t, recorder.Body.String(), `"success":true`) + require.Contains(t, recorder.Body.String(), `"backup_codes"`) + var stored model.User + require.NoError(t, model.DB.First(&stored, user.Id).Error) + require.EqualValues(t, 1, stored.SessionGeneration) +} + +func TestGet2FAStatusAlwaysReturnsBackupCodeCount(t *testing.T) { + user, _ := setupTwoFASessionFailureTest(t, false) + router := gin.New() + router.GET("/twofa/status", func(c *gin.Context) { + c.Set("id", user.Id) + Get2FAStatus(c) + }) + + recorder := httptest.NewRecorder() + router.ServeHTTP(recorder, httptest.NewRequest(http.MethodGet, "/twofa/status", nil)) + + require.Equal(t, http.StatusOK, recorder.Code) + require.JSONEq(t, `{ + "success": true, + "message": "", + "data": { + "enabled": false, + "locked": false, + "backup_codes_remaining": 0 + } + }`, recorder.Body.String()) +} diff --git a/controller/telegram.go b/controller/telegram.go index 528d463c..ba24c0f5 100644 --- a/controller/telegram.go +++ b/controller/telegram.go @@ -4,126 +4,221 @@ import ( "crypto/hmac" "crypto/sha256" "encoding/hex" - "io" + "encoding/json" + "errors" + "fmt" "net/http" + "net/url" "sort" + "strconv" + "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/model" "github.com/gin-contrib/sessions" "github.com/gin-gonic/gin" ) -func TelegramBind(c *gin.Context) { +const ( + telegramAuthorizationTTL = 5 * time.Minute + telegramClockSkew = time.Minute + telegramBindFlowTTL = 5 * time.Minute +) + +var telegramLoginFlowMatch = model.AuthFlowMatch{ + Purpose: model.AuthFlowPurposeOAuth, + Provider: "telegram", + Intent: model.AuthFlowIntentLogin, +} + +type telegramAuthPayload struct { + ID int64 `json:"id"` + FirstName string `json:"first_name,omitempty"` + LastName string `json:"last_name,omitempty"` + Username string `json:"username,omitempty"` + PhotoURL string `json:"photo_url,omitempty"` + AuthDate int64 `json:"auth_date"` + Hash string `json:"hash"` + State string `json:"state"` + authorizationValues url.Values +} + +func (p *telegramAuthPayload) UnmarshalJSON(data []byte) error { + type plainTelegramAuthPayload telegramAuthPayload + var decoded plainTelegramAuthPayload + if err := common.Unmarshal(data, &decoded); err != nil { + return err + } + + var rawFields map[string]json.RawMessage + if err := common.Unmarshal(data, &rawFields); err != nil { + return err + } + authorizationValues := make(url.Values, len(rawFields)) + for key, rawValue := range rawFields { + if key == "state" { + continue + } + switch common.GetJsonType(rawValue) { + case "string", "number", "boolean": + authorizationValues.Set(key, common.JsonRawMessageToString(rawValue)) + default: + return fmt.Errorf("Telegram authorization field %q must be a scalar value", key) + } + } + + *p = telegramAuthPayload(decoded) + p.authorizationValues = authorizationValues + return nil +} + +func (p telegramAuthPayload) values() url.Values { + values := make(url.Values, len(p.authorizationValues)) + for key, entries := range p.authorizationValues { + values[key] = append([]string(nil), entries...) + } + return values +} + +func GenerateTelegramBindState(c *gin.Context) { if !common.TelegramOAuthEnabled { - c.JSON(200, gin.H{ - "message": "管理员未开启通过 Telegram 登录以及注册", - "success": false, - }) + common.ApiErrorMsg(c, "管理员未开启 Telegram 登录") return } - params := c.Request.URL.Query() - if !checkTelegramAuthorization(params, common.TelegramBotToken) { - c.JSON(200, gin.H{ - "message": "无效的请求", - "success": false, - }) + userID := c.GetInt("id") + if userID <= 0 { + c.JSON(http.StatusUnauthorized, gin.H{"success": false, "message": i18n.T(c, i18n.MsgUnauthorized)}) return } - telegramId := params["id"][0] - taken, err := model.IsTelegramIdAlreadyTaken(telegramId) - if handleOAuthIdentityLookupError(c, "Telegram", err) { + state, _, err := model.CreateAuthFlow(model.AuthFlowCreate{ + Purpose: model.AuthFlowPurposeOAuth, + Provider: "telegram", + Intent: model.AuthFlowIntentBind, + UserId: userID, + ExpiresAt: time.Now().Add(telegramBindFlowTTL), + }) + if err != nil { + common.ApiError(c, err) return } - if taken { - c.JSON(200, gin.H{ - "message": "该 Telegram 账户已被绑定", - "success": false, - }) + 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: "telegram", + Intent: model.AuthFlowIntentBind, + UserId: userID, + }); consumeErr != nil { + common.SysError("failed to invalidate Telegram bind flow after session save failure: " + consumeErr.Error()) + } + common.ApiError(c, err) return } + common.ApiSuccess(c, gin.H{"state": state}) +} - session := sessions.Default(c) - id, ok := sessionUserID(session.Get("id")) - if !ok { - common.ApiErrorMsg(c, "用户未登录或登录状态已失效") +func TelegramBind(c *gin.Context) { + if !common.TelegramOAuthEnabled { + common.ApiErrorMsg(c, "管理员未开启 Telegram 登录") return } - user := model.User{Id: id} - if err := user.FillUserById(); err != nil { - c.JSON(200, gin.H{ - "message": err.Error(), - "success": false, - }) + var payload telegramAuthPayload + if err := common.DecodeJson(c.Request.Body, &payload); err != nil { + common.ApiErrorI18n(c, i18n.MsgInvalidParams) return } - if user.Id == 0 { - c.JSON(http.StatusOK, gin.H{ - "success": false, - "message": "用户已注销", - }) + if !checkTelegramAuthorization(payload.values(), common.TelegramBotToken) { + c.JSON(http.StatusForbidden, gin.H{"success": false, "message": "无效或已过期的 Telegram 授权"}) return } - user.TelegramId = telegramId - if err := user.UpdateFields(false, model.UserUpdateFieldTelegramId); err != nil { - c.JSON(200, gin.H{ - "message": err.Error(), - "success": false, - }) + if !oauthSessionStateMatches(c, payload.State) { + c.JSON(http.StatusForbidden, gin.H{"success": false, "message": i18n.T(c, i18n.MsgOAuthStateInvalid)}) return } - - c.Redirect(302, common.ThemeAwarePath("/console/personal")) + userID := c.GetInt("id") + generation, err := model.BindTelegramIdentityWithAuthFlow(userID, strconv.FormatInt(payload.ID, 10), payload.State) + if err != nil { + switch { + case errors.Is(err, model.ErrTelegramIdAlreadyTaken): + common.ApiErrorMsg(c, "该 Telegram 账户已被绑定") + case errors.Is(err, model.ErrAuthFlowInvalid), errors.Is(err, model.ErrAuthFlowExpired), errors.Is(err, model.ErrAuthFlowConsumed): + c.JSON(http.StatusForbidden, gin.H{"success": false, "message": i18n.T(c, i18n.MsgOAuthStateInvalid)}) + default: + common.ApiError(c, err) + } + return + } + clearOAuthSessionState(c) + preserveCurrentSessionAfterCommittedSecurityChange(c, userID, generation, "binding Telegram") + recordUserSecurityAudit(c, userID, "user.telegram_bind", nil) + common.ApiSuccess(c, nil) } func TelegramLogin(c *gin.Context) { if !common.TelegramOAuthEnabled { - c.JSON(200, gin.H{ - "message": "管理员未开启通过 Telegram 登录以及注册", - "success": false, - }) + common.ApiErrorMsg(c, "管理员未开启 Telegram 登录") return } params := c.Request.URL.Query() if !checkTelegramAuthorization(params, common.TelegramBotToken) { - c.JSON(200, gin.H{ - "message": "无效的请求", - "success": false, - }) + c.JSON(http.StatusForbidden, gin.H{"success": false, "message": "无效或已过期的 Telegram 授权"}) return } - - telegramId := params["id"][0] - user := model.User{TelegramId: telegramId} + state := params.Get("state") + if !oauthSessionStateMatches(c, state) { + c.JSON(http.StatusForbidden, gin.H{"success": false, "message": i18n.T(c, i18n.MsgOAuthStateInvalid)}) + return + } + user := model.User{TelegramId: params.Get("id")} if err := user.FillUserByTelegramId(); handleOAuthUserLookupError(c, err) { return } + if _, err := model.ConsumeAuthFlow(state, telegramLoginFlowMatch); err != nil { + c.JSON(http.StatusForbidden, gin.H{"success": false, "message": i18n.T(c, i18n.MsgOAuthStateInvalid)}) + return + } + clearOAuthSessionState(c) setupLogin(&user, c) } -func checkTelegramAuthorization(params map[string][]string, token string) bool { - strs := []string{} - var hash = "" - for k, v := range params { - if k == "hash" { - hash = v[0] - continue +func checkTelegramAuthorization(params url.Values, token string) bool { + return checkTelegramAuthorizationAt(params, token, time.Now()) +} + +func checkTelegramAuthorizationAt(params url.Values, token string, now time.Time) bool { + if strings.TrimSpace(token) == "" || params.Get("id") == "" || params.Get("hash") == "" { + return false + } + authDate, err := strconv.ParseInt(params.Get("auth_date"), 10, 64) + if err != nil { + return false + } + authorizedAt := time.Unix(authDate, 0) + if authorizedAt.After(now.Add(telegramClockSkew)) || now.Sub(authorizedAt) > telegramAuthorizationTTL { + return false + } + keys := make([]string, 0, len(params)) + for key := range params { + if key != "hash" && key != "state" { + keys = append(keys, key) } - strs = append(strs, k+"="+v[0]) } - sort.Strings(strs) - var imploded = "" - for _, s := range strs { - if imploded != "" { - imploded += "\n" + sort.Strings(keys) + lines := make([]string, 0, len(keys)) + for _, key := range keys { + values := params[key] + if len(values) != 1 { + return false } - imploded += s - } - sha256hash := sha256.New() - io.WriteString(sha256hash, token) - hmachash := hmac.New(sha256.New, sha256hash.Sum(nil)) - io.WriteString(hmachash, imploded) - ss := hex.EncodeToString(hmachash.Sum(nil)) - return hash == ss + lines = append(lines, key+"="+values[0]) + } + secret := sha256.Sum256([]byte(token)) + mac := hmac.New(sha256.New, secret[:]) + _, _ = mac.Write([]byte(strings.Join(lines, "\n"))) + provided, err := hex.DecodeString(params.Get("hash")) + return err == nil && hmac.Equal(provided, mac.Sum(nil)) } diff --git a/controller/twofa.go b/controller/twofa.go index f8b3c0f6..cf02e13c 100644 --- a/controller/twofa.go +++ b/controller/twofa.go @@ -186,10 +186,12 @@ func Enable2FA(c *gin.Context) { } // 启用2FA - if err := twoFA.Enable(); err != nil { + generation, err := twoFA.EnableAndBumpSessionGeneration() + if err != nil { common.ApiError(c, err) return } + preserveCurrentSessionAfterCommittedSecurityChange(c, userId, generation, "enabling 2FA") // 记录操作日志 model.RecordLog(userId, model.LogTypeSystem, "成功启用两步验证") @@ -258,10 +260,12 @@ func Disable2FA(c *gin.Context) { } // 禁用2FA - if err := model.DisableTwoFA(userId); err != nil { + generation, err := model.DisableTwoFAAndBumpSessionGeneration(userId) + if err != nil { common.ApiError(c, err) return } + preserveCurrentSessionAfterCommittedSecurityChange(c, userId, generation, "disabling 2FA") // 记录操作日志 model.RecordLog(userId, model.LogTypeSystem, "禁用两步验证") @@ -283,8 +287,9 @@ func Get2FAStatus(c *gin.Context) { } status := map[string]interface{}{ - "enabled": false, - "locked": false, + "enabled": false, + "locked": false, + "backup_codes_remaining": 0, } if twoFA != nil { @@ -373,7 +378,8 @@ func RegenerateBackupCodes(c *gin.Context) { } // 保存新的备用码 - if err := model.CreateBackupCodes(userId, backupCodes); err != nil { + generation, err := model.ReplaceBackupCodesAndBumpSessionGeneration(userId, backupCodes) + if err != nil { c.JSON(http.StatusOK, gin.H{ "success": false, "message": "保存备用码失败", @@ -381,6 +387,7 @@ func RegenerateBackupCodes(c *gin.Context) { common.SysLog("保存备用码失败: " + err.Error()) return } + preserveCurrentSessionAfterCommittedSecurityChange(c, userId, generation, "regenerating 2FA backup codes") // 记录操作日志 model.RecordLog(userId, model.LogTypeSystem, "重新生成两步验证备用码") @@ -529,7 +536,7 @@ func AdminDisable2FA(c *gin.Context) { } // 禁用2FA - if err := model.DisableTwoFA(userId); err != nil { + if _, err := model.DisableTwoFAAndBumpSessionGeneration(userId); err != nil { if errors.Is(err, model.ErrTwoFANotEnabled) { c.JSON(http.StatusOK, gin.H{ "success": false, diff --git a/controller/user.go b/controller/user.go index 7c40e762..43b932db 100644 --- a/controller/user.go +++ b/controller/user.go @@ -9,6 +9,7 @@ import ( "strconv" "strings" "sync" + "time" "github.com/MAX-API-Next/MAX-API/common" "github.com/MAX-API-Next/MAX-API/dto" @@ -134,20 +135,38 @@ func recordLoginAudit(user *model.User, c *gin.Context) { }, extra) } +func isOAuthCredentialReauthentication(method string) bool { + return method == "wechat" || method == "oauth" || + strings.HasPrefix(method, "oauth:") +} + // setup session & cookies and then return user info func setupLogin(user *model.User, c *gin.Context) { model.UpdateUserLastLoginAt(user.Id) session := sessions.Default(c) + loginMethod := loginMethodFromContext(c) session.Delete(SecureVerificationSessionKey) session.Delete(secureVerificationMethodSessionKey) session.Delete(secureVerificationUserSessionKey) session.Delete(secureVerificationScopeSessionKey) session.Delete(PasskeyReadySessionKey) + if isOAuthCredentialReauthentication(loginMethod) { + allowed, grantErr := model.CanUseOAuthReauthentication(user.Id) + if grantErr != nil { + common.SysError(fmt.Sprintf("failed to evaluate OAuth reauthentication grant for user %d: %v", user.Id, grantErr)) + } else if allowed { + session.Set(SecureVerificationSessionKey, time.Now().Unix()) + session.Set(secureVerificationMethodSessionKey, model.SecureVerificationMethodOAuth) + session.Set(secureVerificationUserSessionKey, user.Id) + session.Set(secureVerificationScopeSessionKey, model.SecureVerificationScopeOAuthReauthentication) + } + } 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) + session.Set("session_generation", user.SessionGeneration) err := session.Save() if err != nil { common.ApiErrorI18n(c, i18n.MsgUserSessionSaveFailed) @@ -185,6 +204,18 @@ func Logout(c *gin.Context) { }) } +func RevokeOtherSessions(c *gin.Context) { + userID := c.GetInt("id") + generation, err := model.BumpUserSessionGeneration(userID) + if err != nil { + common.ApiError(c, err) + return + } + preserveCurrentSessionAfterCommittedSecurityChange(c, userID, generation, "revoking other sessions") + recordUserSecurityAudit(c, userID, "user.sessions_revoke_other", nil) + common.ApiSuccess(c, nil) +} + func Register(c *gin.Context) { if !common.RegisterEnabled { common.ApiErrorI18n(c, i18n.MsgUserRegisterDisabled) @@ -886,10 +917,14 @@ func UpdateSelf(c *gin.Context) { common.ApiError(c, err) return } - if err := cleanUser.UpdateFields(updatePassword, updateFields...); err != nil { + generation, err := cleanUser.UpdateFieldsWithSessionGeneration(updatePassword, updateFields...) + if err != nil { common.ApiError(c, err) return } + if updatePassword { + preserveCurrentSessionAfterCommittedSecurityChange(c, cleanUser.Id, generation, "changing password") + } c.JSON(http.StatusOK, gin.H{ "success": true, @@ -898,6 +933,50 @@ func UpdateSelf(c *gin.Context) { return } +func preserveCurrentSessionAfterSecurityChange(c *gin.Context, userID int, generation int64) error { + session := sessions.Default(c) + session.Set("id", userID) + session.Set("session_generation", generation) + session.Delete(SecureVerificationSessionKey) + session.Delete(secureVerificationMethodSessionKey) + session.Delete(secureVerificationUserSessionKey) + session.Delete(secureVerificationScopeSessionKey) + session.Delete(PasskeyReadySessionKey) + return session.Save() +} + +func preserveCurrentSessionAfterCommittedSecurityChange(c *gin.Context, userID int, generation int64, action string) { + if err := preserveCurrentSessionAfterSecurityChange(c, userID, generation); err != nil { + common.SysError(fmt.Sprintf( + "failed to preserve current session after %s for user %d: %v", + action, + userID, + err, + )) + expireCurrentSessionAfterSecurityChange(c, userID, action) + } +} + +func expireCurrentSessionAfterSecurityChange(c *gin.Context, userID int, action string) { + session := sessions.Default(c) + session.Clear() + session.Options(sessions.Options{ + Path: "/", + MaxAge: -1, + HttpOnly: true, + Secure: common.SessionCookieSecure, + SameSite: http.SameSiteLaxMode, + }) + if err := session.Save(); err != nil { + common.SysError(fmt.Sprintf( + "failed to expire current session after %s for user %d: %v", + action, + userID, + err, + )) + } +} + func checkUpdatePassword(originalPassword string, newPassword string, userId int) (updatePassword bool, err error) { if newPassword == "" { return diff --git a/docs/openapi/api.json b/docs/openapi/api.json index 8505cd47..47809e78 100644 --- a/docs/openapi/api.json +++ b/docs/openapi/api.json @@ -1019,15 +1019,124 @@ "get": { "summary": "Telegram登录", "deprecated": false, - "description": "🔓 无需鉴权(OAuth回调)", + "description": "🔓 无需鉴权(OAuth回调);需要先通过 /api/oauth/state 创建并携带当前浏览器会话绑定的一次性 state。客户端必须原样转发 Telegram 返回的全部字段,因为签名校验会覆盖除 hash 和 state 外的所有单值字段。", "tags": [ "OAuth" ], - "parameters": [], + "parameters": [ + { + "name": "id", + "in": "query", + "description": "Telegram 用户 ID", + "required": true, + "schema": { + "type": "integer", + "format": "int64" + } + }, + { + "name": "first_name", + "in": "query", + "description": "Telegram 返回的名字;存在时必须原样转发。", + "required": false, + "schema": { + "type": "string" + } + }, + { + "name": "last_name", + "in": "query", + "description": "Telegram 返回的姓氏;存在时必须原样转发。", + "required": false, + "schema": { + "type": "string" + } + }, + { + "name": "username", + "in": "query", + "description": "Telegram 用户名;存在时必须原样转发。", + "required": false, + "schema": { + "type": "string" + } + }, + { + "name": "photo_url", + "in": "query", + "description": "Telegram 头像 URL;存在时必须原样转发。", + "required": false, + "schema": { + "type": "string" + } + }, + { + "name": "auth_date", + "in": "query", + "description": "Telegram 授权的 Unix 时间戳;必须在有效期内。", + "required": true, + "schema": { + "type": "integer", + "format": "int64" + } + }, + { + "name": "hash", + "in": "query", + "description": "Telegram Login Widget 签名", + "required": true, + "schema": { + "type": "string" + } + }, + { + "name": "state", + "in": "query", + "description": "由 /api/oauth/state 为 Telegram 登录创建的浏览器会话绑定一次性 state", + "required": true, + "schema": { + "type": "string" + } + } + ], "responses": { "200": { "description": "成功", "headers": {} + }, + "403": { + "description": "Telegram 签名或 state 无效、过期或已使用", + "headers": {} + } + }, + "security": [ + { + "Combination343": [] + }, + { + "Combination1243": [] + } + ] + } + }, + "/api/oauth/telegram/bind/state": { + "post": { + "operationId": "createTelegramBindState", + "summary": "创建 Telegram 绑定 state", + "deprecated": false, + "description": "需要登录,并完成 credentials 作用域安全验证或刚刚通过已绑定 OAuth 账户重新认证(oauth_reauthentication)。返回短期有效、与用户绑定的一次性 state。", + "tags": [ + "OAuth" + ], + "parameters": [], + "responses": { + "200": { + "description": "state 创建成功", + "headers": {} + }, + "403": { + "description": "需要 credentials 或 oauth_reauthentication 作用域安全验证,或验证已过期", + "headers": {} } }, "security": [ @@ -1041,10 +1150,69 @@ } }, "/api/oauth/telegram/bind": { - "get": { - "summary": "绑定Telegram", + "post": { + "operationId": "bindTelegramAccount", + "requestBody": { + "required": true, + "content": { + "application/json": { + "schema": { + "type": "object", + "required": [ + "id", + "auth_date", + "hash", + "state" + ], + "properties": { + "id": { + "type": "integer", + "format": "int64" + }, + "first_name": { + "type": "string" + }, + "last_name": { + "type": "string" + }, + "username": { + "type": "string" + }, + "photo_url": { + "type": "string" + }, + "auth_date": { + "type": "integer", + "format": "int64", + "description": "Telegram 授权的 Unix 时间戳;必须在有效期内。" + }, + "hash": { + "type": "string" + }, + "state": { + "type": "string" + } + }, + "additionalProperties": { + "oneOf": [ + { + "type": "string" + }, + { + "type": "number" + }, + { + "type": "boolean" + } + ] + } + } + } + } + }, + "summary": "绑定 Telegram 账户", "deprecated": false, - "description": "🔓 无需鉴权", + "description": "需要登录、credentials 或 oauth_reauthentication 作用域安全验证、5 分钟内的 Telegram 签名,以及由 /api/oauth/telegram/bind/state 生成的一次性 state。", "tags": [ "OAuth" ], @@ -1053,6 +1221,10 @@ "200": { "description": "成功", "headers": {} + }, + "403": { + "description": "签名、state 或安全验证无效、过期或已使用", + "headers": {} } }, "security": [ @@ -1234,6 +1406,36 @@ ] } }, + "/api/user/sessions/revoke": { + "post": { + "operationId": "revokeOtherSessions", + "summary": "退出其他设备", + "deprecated": false, + "description": "需要登录并完成 credentials 作用域安全验证。递增会话代次,使其他设备上的 Cookie 会话失效,同时保留当前会话。", + "tags": [ + "用户管理" + ], + "parameters": [], + "responses": { + "200": { + "description": "其他会话已撤销", + "headers": {} + }, + "403": { + "description": "需要 credentials 作用域安全验证", + "headers": {} + } + }, + "security": [ + { + "Combination343": [] + }, + { + "Combination1243": [] + } + ] + } + }, "/api/user/passkey": { "get": { "summary": "获取Passkey状态", @@ -1261,7 +1463,7 @@ "delete": { "summary": "删除Passkey", "deprecated": false, - "description": "🔐 需要登录(User权限)", + "description": "🔐 需要登录(User权限)并完成 credentials 作用域安全验证", "tags": [ "用户管理" ], @@ -1270,6 +1472,10 @@ "200": { "description": "成功", "headers": {} + }, + "403": { + "description": "需要 credentials 作用域安全验证或验证已过期", + "headers": {} } }, "security": [ @@ -1286,7 +1492,7 @@ "post": { "summary": "开始注册Passkey", "deprecated": false, - "description": "🔐 需要登录(User权限)", + "description": "需要登录并完成 passkey_register 或 credentials 作用域安全验证;仅使用关联 OAuth 登录且没有密码、2FA 或 Passkey 的账户,也可使用刚完成 OAuth 登录授予的 oauth_reauthentication 作用域。已有 Passkey 时,本流程表示显式替换,不能使用 oauth_reauthentication。", "tags": [ "用户管理" ], @@ -1295,6 +1501,10 @@ "200": { "description": "成功", "headers": {} + }, + "403": { + "description": "需要 passkey_register、credentials,或仅 OAuth 账户的 oauth_reauthentication 作用域安全验证;验证已过期或账户已有其他凭据时也会拒绝", + "headers": {} } }, "security": [ @@ -1311,7 +1521,7 @@ "post": { "summary": "完成注册Passkey", "deprecated": false, - "description": "🔐 需要登录(User权限)", + "description": "需要登录并完成 passkey_register 或 credentials 作用域安全验证;仅使用关联 OAuth 登录且没有密码、2FA 或 Passkey 的账户,也可使用刚完成 OAuth 登录授予的 oauth_reauthentication 作用域。成功注册或替换后撤销其他会话。", "tags": [ "用户管理" ], @@ -1320,6 +1530,10 @@ "200": { "description": "成功", "headers": {} + }, + "403": { + "description": "需要 passkey_register、credentials,或仅 OAuth 账户的 oauth_reauthentication 作用域安全验证;验证已过期或账户已有其他凭据时也会拒绝", + "headers": {} } }, "security": [ @@ -2166,7 +2380,7 @@ "post": { "summary": "设置2FA", "deprecated": false, - "description": "🔐 需要登录(User权限)", + "description": "需要登录并完成 credentials 作用域安全验证;仅使用关联 OAuth 登录且尚无其他验证方式的账户,也可使用刚完成 OAuth 登录授予的 oauth_reauthentication 作用域。返回待确认的 TOTP 密钥与备用码。", "tags": [ "两步验证" ], @@ -2175,6 +2389,10 @@ "200": { "description": "成功", "headers": {} + }, + "403": { + "description": "需要 credentials 或 oauth_reauthentication 作用域安全验证,或验证已过期", + "headers": {} } }, "security": [ @@ -2191,7 +2409,7 @@ "post": { "summary": "启用2FA", "deprecated": false, - "description": "🔐 需要登录(User权限)", + "description": "需要登录并完成 credentials 作用域安全验证;仅使用关联 OAuth 登录且尚无其他验证方式的账户,也可使用刚完成 OAuth 登录授予的 oauth_reauthentication 作用域。成功启用后撤销其他会话。", "tags": [ "两步验证" ], @@ -2214,6 +2432,10 @@ "200": { "description": "成功", "headers": {} + }, + "403": { + "description": "需要 credentials 或 oauth_reauthentication 作用域安全验证,或验证已过期", + "headers": {} } }, "security": [ @@ -2230,7 +2452,7 @@ "post": { "summary": "禁用2FA", "deprecated": false, - "description": "🔐 需要登录(User权限)", + "description": "需要登录。成功禁用后会递增会话代次,并使其他设备上的 Cookie 会话失效。", "tags": [ "两步验证" ], @@ -2269,7 +2491,7 @@ "post": { "summary": "重新生成备用码", "deprecated": false, - "description": "🔐 需要登录(User权限)", + "description": "需要登录。成功重新生成后会递增会话代次,并使其他设备上的 Cookie 会话失效。", "tags": [ "两步验证" ], @@ -2350,14 +2572,17 @@ "password": { "type": "string", "format": "password", - "description": "当前账户密码,仅method=password且scope=access_token或account_delete时使用" + "description": "当前账户密码,仅 method=password 且 scope 为受支持的敏感操作作用域时使用" }, "scope": { "type": "string", - "description": "验证作用域;密码验证仅支持access_token或account_delete", + "description": "验证作用域;验证结果仅可用于明确允许该作用域的敏感操作。oauth_reauthentication 仅在完成 OAuth 登录后由登录流程授予,不属于本接口可提交值", "enum": [ "access_token", - "account_delete" + "account_delete", + "credentials", + "api_token", + "passkey_register" ] } } @@ -2394,12 +2619,15 @@ "name": "scope", "in": "query", "required": false, - "description": "验证作用域;传access_token或account_delete时,无2FA或Passkey的用户可使用密码验证", + "description": "验证作用域;无 2FA 或 Passkey 的用户可使用当前密码完成对应作用域的验证。oauth_reauthentication 仅在完成 OAuth 登录后由登录流程授予,不属于本接口可查询值", "schema": { "type": "string", "enum": [ "access_token", - "account_delete" + "account_delete", + "credentials", + "api_token", + "passkey_register" ] } } @@ -3466,7 +3694,7 @@ "post": { "summary": "创建令牌", "deprecated": false, - "description": "🔐 需要登录(User权限)", + "description": "需要登录并完成 api_token 作用域安全验证。创建接口不返回完整密钥;后续查看仍需调用受 api_token 作用域保护的密钥查看接口。", "tags": [ "令牌管理" ], @@ -3484,6 +3712,10 @@ "200": { "description": "成功", "headers": {} + }, + "403": { + "description": "需要 api_token 作用域安全验证或验证已过期", + "headers": {} } }, "security": [ @@ -3633,6 +3865,98 @@ ] } }, + "/api/token/{id}/key": { + "post": { + "operationId": "revealApiTokenKey", + "summary": "查看完整 API Token 密钥", + "deprecated": false, + "description": "需要登录并完成 api_token 作用域安全验证。仅返回当前用户拥有的指定 Token 密钥。", + "tags": [ + "令牌管理" + ], + "parameters": [ + { + "name": "id", + "in": "path", + "required": true, + "schema": { + "type": "integer" + } + } + ], + "responses": { + "200": { + "description": "成功", + "headers": {} + }, + "403": { + "description": "需要 api_token 作用域安全验证", + "headers": {} + } + }, + "security": [ + { + "Combination343": [] + }, + { + "Combination1243": [] + } + ] + } + }, + "/api/token/batch/keys": { + "post": { + "operationId": "revealApiTokenKeysBatch", + "summary": "批量导出完整 API Token 密钥", + "deprecated": false, + "description": "需要登录并完成 api_token 作用域安全验证。仅返回当前用户拥有的指定 Token 密钥。", + "tags": [ + "令牌管理" + ], + "parameters": [], + "requestBody": { + "required": true, + "content": { + "application/json": { + "schema": { + "type": "object", + "required": [ + "ids" + ], + "properties": { + "ids": { + "type": "array", + "minItems": 1, + "maxItems": 100, + "items": { + "type": "integer" + } + } + } + } + } + } + }, + "responses": { + "200": { + "description": "成功", + "headers": {} + }, + "403": { + "description": "需要 api_token 作用域安全验证", + "headers": {} + } + }, + "security": [ + { + "Combination343": [] + }, + { + "Combination1243": [] + } + ] + } + }, "/api/token/batch": { "post": { "summary": "批量删除令牌", diff --git a/i18n/keys.go b/i18n/keys.go index 5c1d8f27..3a745ba4 100644 --- a/i18n/keys.go +++ b/i18n/keys.go @@ -121,6 +121,8 @@ const ( MsgUserQuotaChangeZero = "user.quota_change_zero" ) +const MsgUserPasswordResetSuccessTokensRevoked = "user.password_reset_success_tokens_revoked" + // Quota related messages const ( MsgQuotaNegative = "quota.negative" @@ -130,6 +132,9 @@ const ( MsgQuotaThresholdGtZero = "quota.threshold_gt_zero" ) +// Billing related messages +const MsgBillingReconciliationPending = "billing.reconciliation_pending" + // Subscription related messages const ( MsgSubscriptionNotEnabled = "subscription.not_enabled" diff --git a/i18n/locales/en.yaml b/i18n/locales/en.yaml index 5380e595..cf87d6e9 100644 --- a/i18n/locales/en.yaml +++ b/i18n/locales/en.yaml @@ -79,6 +79,7 @@ user.verification_code_error: "Verification code is incorrect or has expired" user.email_already_taken: "Email address is already in use" user.password_unset: "This account has no password set. Please use password reset or contact an administrator to reset it." user.password_reset_link_invalid: "Password reset link is invalid or has expired" +user.password_reset_success_tokens_revoked: "Password reset successfully. For your security, all existing API tokens were revoked. Sign in and create new API tokens." user.input_invalid: "Invalid input {{.Error}}" user.no_permission_same_level: "No permission to access users of same or higher level" user.no_permission_higher_level: "No permission to update users of same or higher permission level" @@ -115,6 +116,9 @@ quota.insufficient: "Insufficient quota" quota.warning_invalid: "Invalid warning type" quota.threshold_gt_zero: "Warning threshold must be greater than 0" +# Billing messages +billing.reconciliation_pending: "Billing reconciliation is incomplete. Do not submit another request; contact an administrator for assistance." + # Subscription messages subscription.not_enabled: "Subscription plan is not enabled" subscription.title_empty: "Subscription plan title cannot be empty" diff --git a/i18n/locales/zh-CN.yaml b/i18n/locales/zh-CN.yaml index 7798a643..b3fd4cab 100644 --- a/i18n/locales/zh-CN.yaml +++ b/i18n/locales/zh-CN.yaml @@ -80,6 +80,7 @@ user.verification_code_error: "验证码错误或已过期" user.email_already_taken: "邮箱地址已被占用" user.password_unset: "当前账号未设置密码,请使用密码重置或联系管理员重置密码" user.password_reset_link_invalid: "重置链接非法或已过期" +user.password_reset_success_tokens_revoked: "密码重置成功。为保障账户安全,所有现有 API 令牌已被撤销。请登录并重新创建 API 令牌。" user.input_invalid: "输入不合法 {{.Error}}" user.no_permission_same_level: "无权获取同级或更高等级用户的信息" user.no_permission_higher_level: "无权更新同权限等级或更高权限等级的用户信息" @@ -116,6 +117,9 @@ quota.insufficient: "额度不足" quota.warning_invalid: "无效的预警类型" quota.threshold_gt_zero: "预警阈值必须大于0" +# Billing messages +billing.reconciliation_pending: "存在未完成的计费对账,请勿重复提交请求并联系管理员处理" + # Subscription messages subscription.not_enabled: "套餐未启用" subscription.title_empty: "套餐标题不能为空" diff --git a/i18n/locales/zh-TW.yaml b/i18n/locales/zh-TW.yaml index 3d3aed05..74858f7a 100644 --- a/i18n/locales/zh-TW.yaml +++ b/i18n/locales/zh-TW.yaml @@ -80,6 +80,7 @@ user.verification_code_error: "驗證碼錯誤或已過期" user.email_already_taken: "信箱位址已被占用" user.password_unset: "目前帳號未設定密碼,請使用密碼重置或聯繫管理員重置密碼" user.password_reset_link_invalid: "重置連結非法或已過期" +user.password_reset_success_tokens_revoked: "密碼重設成功。為保障帳戶安全,所有現有 API 令牌已被撤銷。請登入並重新建立 API 令牌。" user.input_invalid: "輸入不合法 {{.Error}}" user.no_permission_same_level: "無權獲取同級或更高等級使用者的資訊" user.no_permission_higher_level: "無權更新同權限等級或更高權限等級的使用者資訊" @@ -116,6 +117,9 @@ quota.insufficient: "額度不足" quota.warning_invalid: "無效的預警類型" quota.threshold_gt_zero: "預警閾值必須大於0" +# Billing messages +billing.reconciliation_pending: "存在未完成的計費對帳,請勿重複提交請求並聯絡管理員處理" + # Subscription messages subscription.not_enabled: "訂閱方案未啟用" subscription.title_empty: "訂閱方案標題不能為空" diff --git a/middleware/auth.go b/middleware/auth.go index f7282904..22b169f9 100644 --- a/middleware/auth.go +++ b/middleware/auth.go @@ -45,6 +45,35 @@ func sessionInt(value any) (int, bool) { } } +func sessionInt64(value any) (int64, bool) { + switch v := value.(type) { + case int: + return int64(v), true + case int64: + return v, true + case float64: + return int64(v), true + default: + return 0, false + } +} + +func sessionGenerationMatches(session sessions.Session, user *model.UserBase) bool { + value := session.Get("session_generation") + if value == nil { + return user.SessionGeneration == 0 + } + generation, ok := sessionInt64(value) + return ok && generation == user.SessionGeneration +} + +func clearInvalidSession(session sessions.Session) { + session.Clear() + if err := session.Save(); err != nil { + common.SysLog("failed to clear invalid session: " + err.Error()) + } +} + func refreshSessionFromUser(session sessions.Session, user *model.UserBase) { changed := false setIfChanged := func(key string, value any) { @@ -58,6 +87,7 @@ func refreshSessionFromUser(session sessions.Session, user *model.UserBase) { setIfChanged("role", user.Role) setIfChanged("status", user.Status) setIfChanged("group", user.Group) + setIfChanged("session_generation", user.SessionGeneration) if changed { if err := session.Save(); err != nil { common.SysLog(fmt.Sprintf("failed to refresh session for user %d: %v", user.Id, err)) @@ -73,6 +103,18 @@ func getSessionUser(session sessions.Session) (*model.UserBase, error) { return model.GetUserCache(id) } +func sessionUserForRequest(session sessions.Session) (*model.UserBase, error) { + user, err := getSessionUser(session) + if err != nil { + return nil, err + } + if !sessionGenerationMatches(session, user) { + clearInvalidSession(session) + return nil, nil + } + return user, nil +} + func writeSessionUserContext(c *gin.Context, user *model.UserBase) { c.Set("id", user.Id) c.Set("username", user.Username) @@ -111,6 +153,15 @@ func freshAuthHelper(c *gin.Context, minRole int) { c.Abort() return } + if !sessionGenerationMatches(session, userCache) { + clearInvalidSession(session) + c.JSON(http.StatusUnauthorized, gin.H{ + "success": false, + "message": common.TranslateMessage(c, i18n.MsgAuthNotLoggedIn), + }) + c.Abort() + return + } refreshSessionFromUser(session, userCache) username = userCache.Username role = userCache.Role @@ -239,7 +290,7 @@ func TryUserAuth() func(c *gin.Context) { return func(c *gin.Context) { session := sessions.Default(c) if session.Get("id") != nil { - if userCache, err := getSessionUser(session); err == nil && userCache.Status == common.UserStatusEnabled { + if userCache, err := sessionUserForRequest(session); err == nil && userCache != nil && userCache.Status == common.UserStatusEnabled { refreshSessionFromUser(session, userCache) writeSessionUserContext(c, userCache) } @@ -277,8 +328,8 @@ func TokenOrUserAuth() func(c *gin.Context) { // Try session auth first (dashboard users) session := sessions.Default(c) if session.Get("id") != nil { - userCache, err := getSessionUser(session) - if err == nil && userCache.Status == common.UserStatusEnabled { + userCache, err := sessionUserForRequest(session) + if err == nil && userCache != nil && userCache.Status == common.UserStatusEnabled { refreshSessionFromUser(session, userCache) writeSessionUserContext(c, userCache) c.Next() diff --git a/middleware/secure_verification.go b/middleware/secure_verification.go index 03ba00e2..dd31bde6 100644 --- a/middleware/secure_verification.go +++ b/middleware/secure_verification.go @@ -1,9 +1,12 @@ package middleware import ( + "fmt" "net/http" "time" + "github.com/MAX-API-Next/MAX-API/common" + "github.com/MAX-API-Next/MAX-API/model" "github.com/gin-contrib/sessions" "github.com/gin-gonic/gin" ) @@ -14,7 +17,6 @@ const ( secureVerificationMethodSessionKey = "secure_verified_method" secureVerificationUserSessionKey = "secure_verified_user_id" secureVerificationScopeSessionKey = "secure_verified_scope" - secureVerificationMethodPassword = "password" // SecureVerificationTimeout 验证有效期(秒) SecureVerificationTimeout = 300 // 5分钟 ) @@ -88,7 +90,7 @@ func SecureVerificationRequired(requiredScopes ...string) gin.HandlerFunc { } verifiedMethod, ok := session.Get(secureVerificationMethodSessionKey).(string) - if !ok || verifiedMethod == "" { + if !ok || !isSupportedSecureVerificationMethod(verifiedMethod) { clearSecureVerificationSession(session) c.JSON(http.StatusForbidden, gin.H{ "success": false, @@ -98,12 +100,25 @@ func SecureVerificationRequired(requiredScopes ...string) gin.HandlerFunc { c.Abort() return } - requiredScope := "" - if len(requiredScopes) > 0 { - requiredScope = requiredScopes[0] - } verifiedScope, _ := session.Get(secureVerificationScopeSessionKey).(string) - if requiredScope != "" && verifiedScope != requiredScope { + hasRequiredScope := false + scopeMatches := false + for _, allowedScope := range requiredScopes { + if allowedScope == "" { + continue + } + hasRequiredScope = true + if verifiedScope == allowedScope { + scopeMatches = true + break + } + } + if !hasRequiredScope { + scopeMatches = true + } + if !scopeMatches || + ((verifiedMethod == model.SecureVerificationMethodPassword || + verifiedMethod == model.SecureVerificationMethodOAuth) && !hasRequiredScope) { c.JSON(http.StatusForbidden, gin.H{ "success": false, "message": "需要对应操作的安全验证", @@ -112,23 +127,31 @@ func SecureVerificationRequired(requiredScopes ...string) gin.HandlerFunc { c.Abort() return } - if verifiedMethod == secureVerificationMethodPassword { - requiredScope := "" - if len(requiredScopes) > 0 { - requiredScope = requiredScopes[0] + if verifiedMethod == model.SecureVerificationMethodOAuth && + verifiedScope == model.SecureVerificationScopeOAuthReauthentication { + allowed, err := model.CanUseOAuthReauthentication(userId) + if err != nil { + common.SysError(fmt.Sprintf("failed to validate OAuth reauthentication grant for user %d: %v", userId, err)) + clearSecureVerificationSession(session) + c.JSON(http.StatusForbidden, gin.H{ + "success": false, + "message": "验证状态异常,请重新验证", + "code": "VERIFICATION_INVALID", + }) + c.Abort() + return } - verifiedScope, _ := session.Get(secureVerificationScopeSessionKey).(string) - if requiredScope == "" || verifiedScope != requiredScope { + if !allowed { + clearSecureVerificationSession(session) c.JSON(http.StatusForbidden, gin.H{ "success": false, - "message": "需要对应操作的安全验证", + "message": "请使用现有密码、2FA 或 Passkey 重新验证", "code": "VERIFICATION_REQUIRED", }) c.Abort() return } } - c.Next() } } @@ -185,7 +208,7 @@ func OptionalSecureVerification() gin.HandlerFunc { return } verifiedMethod, ok := session.Get(secureVerificationMethodSessionKey).(string) - if !ok || verifiedMethod == "" || verifiedMethod == secureVerificationMethodPassword { + if !ok || !isStrongSecureVerificationMethod(verifiedMethod) { c.Set("secure_verified", false) c.Next() return @@ -197,6 +220,18 @@ func OptionalSecureVerification() gin.HandlerFunc { } } +func isSupportedSecureVerificationMethod(method string) bool { + return method == model.SecureVerificationMethod2FA || + method == model.SecureVerificationMethodPasskey || + method == model.SecureVerificationMethodPassword || + method == model.SecureVerificationMethodOAuth +} + +func isStrongSecureVerificationMethod(method string) bool { + return method == model.SecureVerificationMethod2FA || + method == model.SecureVerificationMethodPasskey +} + // ClearSecureVerification 清除安全验证状态 // 用于用户登出或需要强制重新验证的场景 func ClearSecureVerification(c *gin.Context) { diff --git a/middleware/secure_verification_test.go b/middleware/secure_verification_test.go index 59414bcd..95ae5527 100644 --- a/middleware/secure_verification_test.go +++ b/middleware/secure_verification_test.go @@ -1,15 +1,21 @@ package middleware import ( + "fmt" "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/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" ) func TestSecureVerificationRejectsMarkerFromAnotherUser(t *testing.T) { @@ -56,8 +62,8 @@ func TestPasswordVerificationIsRestrictedToMatchingScope(t *testing.T) { session := sessions.Default(c) session.Set(SecureVerificationSessionKey, time.Now().Unix()) session.Set(secureVerificationUserSessionKey, 1001) - session.Set(secureVerificationMethodSessionKey, secureVerificationMethodPassword) - session.Set(secureVerificationScopeSessionKey, "access_token") + session.Set(secureVerificationMethodSessionKey, model.SecureVerificationMethodPassword) + session.Set(secureVerificationScopeSessionKey, model.SecureVerificationScopeAccessToken) require.NoError(t, session.Save()) c.Status(http.StatusNoContent) }) @@ -105,8 +111,8 @@ func TestNonPasswordVerificationIsRestrictedToMatchingScope(t *testing.T) { session := sessions.Default(c) session.Set(SecureVerificationSessionKey, time.Now().Unix()) session.Set(secureVerificationUserSessionKey, 1001) - session.Set(secureVerificationMethodSessionKey, "passkey") - session.Set(secureVerificationScopeSessionKey, "access_token") + session.Set(secureVerificationMethodSessionKey, model.SecureVerificationMethodPasskey) + session.Set(secureVerificationScopeSessionKey, model.SecureVerificationScopeAccessToken) require.NoError(t, session.Save()) c.Status(http.StatusNoContent) }) @@ -140,3 +146,332 @@ func TestNonPasswordVerificationIsRestrictedToMatchingScope(t *testing.T) { require.Equal(t, http.StatusForbidden, deleteRecorder.Code) require.Contains(t, deleteRecorder.Body.String(), `"code":"VERIFICATION_REQUIRED"`) } + +func TestSecureVerificationIgnoresEmptyScopeEntries(t *testing.T) { + router := gin.New() + router.Use(sessions.Sessions("session", cookie.NewStore([]byte("secure-verification-empty-scope-test")))) + router.GET("/seed", func(c *gin.Context) { + session := sessions.Default(c) + session.Set(SecureVerificationSessionKey, time.Now().Unix()) + session.Set(secureVerificationUserSessionKey, 1001) + session.Set(secureVerificationMethodSessionKey, model.SecureVerificationMethodPasskey) + session.Set(secureVerificationScopeSessionKey, model.SecureVerificationScopeAccessToken) + require.NoError(t, session.Save()) + c.Status(http.StatusNoContent) + }) + setUser := func(c *gin.Context) { + c.Set("id", 1001) + c.Next() + } + router.GET( + "/access-token", + setUser, + SecureVerificationRequired("", model.SecureVerificationScopeAccessToken), + func(c *gin.Context) { c.Status(http.StatusNoContent) }, + ) + router.GET( + "/credentials", + setUser, + SecureVerificationRequired("", model.SecureVerificationScopeCredentials), + func(c *gin.Context) { c.Status(http.StatusNoContent) }, + ) + + seedRecorder := httptest.NewRecorder() + router.ServeHTTP(seedRecorder, httptest.NewRequest(http.MethodGet, "/seed", nil)) + require.Equal(t, http.StatusNoContent, seedRecorder.Code) + perform := func(path string) *httptest.ResponseRecorder { + recorder := httptest.NewRecorder() + request := httptest.NewRequest(http.MethodGet, path, nil) + for _, sessionCookie := range seedRecorder.Result().Cookies() { + request.AddCookie(sessionCookie) + } + router.ServeHTTP(recorder, request) + return recorder + } + + require.Equal(t, http.StatusNoContent, perform("/access-token").Code) + credentialsRecorder := perform("/credentials") + require.Equal(t, http.StatusForbidden, credentialsRecorder.Code) + require.Contains(t, credentialsRecorder.Body.String(), `"code":"VERIFICATION_REQUIRED"`) +} + +func TestPasskeyDeletionRequiresCredentialVerificationScope(t *testing.T) { + router := gin.New() + router.Use(sessions.Sessions("session", cookie.NewStore([]byte("secure-verification-passkey-delete-scope-test")))) + router.GET("/seed", func(c *gin.Context) { + session := sessions.Default(c) + session.Set(SecureVerificationSessionKey, time.Now().Unix()) + session.Set(secureVerificationUserSessionKey, 1001) + session.Set(secureVerificationMethodSessionKey, model.SecureVerificationMethodPasskey) + session.Set(secureVerificationScopeSessionKey, model.SecureVerificationScopeAccessToken) + require.NoError(t, session.Save()) + c.Status(http.StatusNoContent) + }) + router.DELETE( + "/passkey", + func(c *gin.Context) { + c.Set("id", 1001) + c.Next() + }, + SecureVerificationRequired("credentials"), + func(c *gin.Context) { c.Status(http.StatusNoContent) }, + ) + + seedRecorder := httptest.NewRecorder() + router.ServeHTTP(seedRecorder, httptest.NewRequest(http.MethodGet, "/seed", nil)) + require.Equal(t, http.StatusNoContent, seedRecorder.Code) + + recorder := httptest.NewRecorder() + request := httptest.NewRequest(http.MethodDelete, "/passkey", nil) + for _, sessionCookie := range seedRecorder.Result().Cookies() { + request.AddCookie(sessionCookie) + } + router.ServeHTTP(recorder, request) + + require.Equal(t, http.StatusForbidden, recorder.Code) + require.Contains(t, recorder.Body.String(), `"code":"VERIFICATION_REQUIRED"`) +} + +func TestOAuthVerificationIsRestrictedToNarrowReauthenticationScope(t *testing.T) { + oldDB := model.DB + oldRedisEnabled := common.RedisEnabled + oldMemoryCacheEnabled := common.MemoryCacheEnabled + common.RedisEnabled = false + common.MemoryCacheEnabled = 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) + model.DB = db + require.NoError(t, db.AutoMigrate(&model.User{}, &model.TwoFA{}, &model.PasskeyCredential{})) + user := model.User{ + Id: 1001, Username: "oauth-scope-user", Role: common.RoleCommonUser, + Status: common.UserStatusEnabled, Group: "default", + } + require.NoError(t, db.Create(&user).Error) + t.Cleanup(func() { + model.DB = oldDB + common.RedisEnabled = oldRedisEnabled + common.MemoryCacheEnabled = oldMemoryCacheEnabled + if sqlDB, dbErr := db.DB(); dbErr == nil { + _ = sqlDB.Close() + } + }) + + router := gin.New() + router.Use(sessions.Sessions("session", cookie.NewStore([]byte("secure-verification-oauth-scope-test")))) + router.GET("/seed", func(c *gin.Context) { + session := sessions.Default(c) + session.Set(SecureVerificationSessionKey, time.Now().Unix()) + session.Set(secureVerificationUserSessionKey, 1001) + session.Set(secureVerificationMethodSessionKey, model.SecureVerificationMethodOAuth) + session.Set(secureVerificationScopeSessionKey, model.SecureVerificationScopeOAuthReauthentication) + require.NoError(t, session.Save()) + c.Status(http.StatusNoContent) + }) + setUser := func(c *gin.Context) { + c.Set("id", 1001) + c.Next() + } + router.POST("/telegram-bind", setUser, SecureVerificationRequired( + model.SecureVerificationScopeCredentials, + model.SecureVerificationScopeOAuthReauthentication, + ), func(c *gin.Context) { + c.Status(http.StatusNoContent) + }) + router.POST("/passkey-register", setUser, SecureVerificationRequired( + model.SecureVerificationScopePasskeyRegister, + model.SecureVerificationScopeCredentials, + model.SecureVerificationScopeOAuthReauthentication, + ), func(c *gin.Context) { + c.Status(http.StatusNoContent) + }) + router.POST("/two-fa-setup", setUser, SecureVerificationRequired( + model.SecureVerificationScopeCredentials, + model.SecureVerificationScopeOAuthReauthentication, + ), func(c *gin.Context) { + c.Status(http.StatusNoContent) + }) + router.POST("/credentials", setUser, SecureVerificationRequired(model.SecureVerificationScopeCredentials), func(c *gin.Context) { + c.Status(http.StatusNoContent) + }) + router.POST("/api-token", setUser, SecureVerificationRequired(model.SecureVerificationScopeAPIToken), func(c *gin.Context) { + c.Status(http.StatusNoContent) + }) + router.POST("/unscoped", setUser, SecureVerificationRequired(), func(c *gin.Context) { + c.Status(http.StatusNoContent) + }) + + seedRecorder := httptest.NewRecorder() + router.ServeHTTP(seedRecorder, httptest.NewRequest(http.MethodGet, "/seed", nil)) + require.Equal(t, http.StatusNoContent, seedRecorder.Code) + + perform := func(path string) *httptest.ResponseRecorder { + recorder := httptest.NewRecorder() + request := httptest.NewRequest(http.MethodPost, path, nil) + for _, sessionCookie := range seedRecorder.Result().Cookies() { + request.AddCookie(sessionCookie) + } + router.ServeHTTP(recorder, request) + return recorder + } + + require.Equal(t, http.StatusNoContent, perform("/telegram-bind").Code) + require.Equal(t, http.StatusNoContent, perform("/passkey-register").Code) + require.Equal(t, http.StatusNoContent, perform("/two-fa-setup").Code) + credentialsRecorder := perform("/credentials") + require.Equal(t, http.StatusForbidden, credentialsRecorder.Code) + require.Contains(t, credentialsRecorder.Body.String(), `"code":"VERIFICATION_REQUIRED"`) + apiTokenRecorder := perform("/api-token") + require.Equal(t, http.StatusForbidden, apiTokenRecorder.Code) + require.Contains(t, apiTokenRecorder.Body.String(), `"code":"VERIFICATION_REQUIRED"`) + unscopedRecorder := perform("/unscoped") + require.Equal(t, http.StatusForbidden, unscopedRecorder.Code) + require.Contains(t, unscopedRecorder.Body.String(), `"code":"VERIFICATION_REQUIRED"`) + + require.NoError(t, db.Model(&model.User{}).Where("id = ?", user.Id).UpdateColumn("password", "new-password-hash").Error) + revokedGrantRecorder := perform("/passkey-register") + require.Equal(t, http.StatusForbidden, revokedGrantRecorder.Code) + require.Contains(t, revokedGrantRecorder.Body.String(), `"code":"VERIFICATION_REQUIRED"`) +} + +func TestSecureVerificationRejectsLoginMethodMarker(t *testing.T) { + router := gin.New() + router.Use(sessions.Sessions("session", cookie.NewStore([]byte("secure-verification-login-marker-test")))) + router.GET("/seed", func(c *gin.Context) { + session := sessions.Default(c) + session.Set(SecureVerificationSessionKey, time.Now().Unix()) + session.Set(secureVerificationUserSessionKey, 1001) + session.Set(secureVerificationMethodSessionKey, "login:password") + session.Set(secureVerificationScopeSessionKey, model.SecureVerificationScopeCredentials) + require.NoError(t, session.Save()) + c.Status(http.StatusNoContent) + }) + router.POST( + "/credentials", + func(c *gin.Context) { + c.Set("id", 1001) + c.Next() + }, + SecureVerificationRequired("credentials"), + func(c *gin.Context) { + c.Status(http.StatusNoContent) + }, + ) + + seedRecorder := httptest.NewRecorder() + router.ServeHTTP(seedRecorder, httptest.NewRequest(http.MethodGet, "/seed", nil)) + require.Equal(t, http.StatusNoContent, seedRecorder.Code) + + recorder := httptest.NewRecorder() + request := httptest.NewRequest(http.MethodPost, "/credentials", nil) + for _, sessionCookie := range seedRecorder.Result().Cookies() { + request.AddCookie(sessionCookie) + } + router.ServeHTTP(recorder, request) + + require.Equal(t, http.StatusForbidden, recorder.Code) + require.Contains(t, recorder.Body.String(), `"code":"VERIFICATION_INVALID"`) +} + +func TestUserAuthRejectsStaleSessionGeneration(t *testing.T) { + gin.SetMode(gin.TestMode) + oldDB := model.DB + oldRedisEnabled := common.RedisEnabled + oldMemoryCacheEnabled := common.MemoryCacheEnabled + common.RedisEnabled = false + common.MemoryCacheEnabled = 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) + model.DB = db + require.NoError(t, db.AutoMigrate(&model.User{})) + if !db.Migrator().HasColumn(&model.User{}, "session_generation") { + require.NoError(t, db.Migrator().AddColumn(&model.User{}, "SessionGeneration")) + } + user := model.User{ + Id: 73001, + Username: "stale-session-user", + Password: "hashed-password", + Role: common.RoleCommonUser, + Status: common.UserStatusEnabled, + Group: "default", + } + require.NoError(t, db.Create(&user).Error) + + t.Cleanup(func() { + model.DB = oldDB + common.RedisEnabled = oldRedisEnabled + common.MemoryCacheEnabled = oldMemoryCacheEnabled + if sqlDB, dbErr := db.DB(); dbErr == nil { + _ = sqlDB.Close() + } + }) + + engine := gin.New() + engine.Use(sessions.Sessions("session", cookie.NewStore([]byte("session-generation-test")))) + engine.GET("/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) + session.Set("session_generation", int64(0)) + require.NoError(t, session.Save()) + c.Status(http.StatusNoContent) + }) + engine.GET("/protected", UserAuth(), func(c *gin.Context) { + c.Status(http.StatusNoContent) + }) + engine.GET("/try-user", TryUserAuth(), func(c *gin.Context) { + c.JSON(http.StatusOK, gin.H{"id": c.GetInt("id")}) + }) + engine.GET("/token-or-user", TokenOrUserAuth(), func(c *gin.Context) { + c.Status(http.StatusNoContent) + }) + + loginRecorder := httptest.NewRecorder() + engine.ServeHTTP(loginRecorder, httptest.NewRequest(http.MethodGet, "/login", nil)) + require.Equal(t, http.StatusNoContent, loginRecorder.Code) + + currentRequest := httptest.NewRequest(http.MethodGet, "/protected", nil) + for _, sessionCookie := range loginRecorder.Result().Cookies() { + currentRequest.AddCookie(sessionCookie) + } + currentRecorder := httptest.NewRecorder() + engine.ServeHTTP(currentRecorder, currentRequest) + require.Equal(t, http.StatusNoContent, currentRecorder.Code) + + require.NoError(t, db.Model(&model.User{}). + Where("id = ?", user.Id). + UpdateColumn("session_generation", 1).Error) + + protectedRequest := httptest.NewRequest(http.MethodGet, "/protected", nil) + for _, sessionCookie := range loginRecorder.Result().Cookies() { + protectedRequest.AddCookie(sessionCookie) + } + protectedRecorder := httptest.NewRecorder() + engine.ServeHTTP(protectedRecorder, protectedRequest) + + require.Equal(t, http.StatusUnauthorized, protectedRecorder.Code) + + tryUserRequest := httptest.NewRequest(http.MethodGet, "/try-user", nil) + for _, sessionCookie := range loginRecorder.Result().Cookies() { + tryUserRequest.AddCookie(sessionCookie) + } + tryUserRecorder := httptest.NewRecorder() + engine.ServeHTTP(tryUserRecorder, tryUserRequest) + + require.Equal(t, http.StatusOK, tryUserRecorder.Code) + require.JSONEq(t, `{"id":0}`, tryUserRecorder.Body.String()) + + tokenOrUserRequest := httptest.NewRequest(http.MethodGet, "/token-or-user", nil) + for _, sessionCookie := range loginRecorder.Result().Cookies() { + tokenOrUserRequest.AddCookie(sessionCookie) + } + tokenOrUserRecorder := httptest.NewRecorder() + engine.ServeHTTP(tokenOrUserRecorder, tokenOrUserRequest) + + require.Equal(t, http.StatusUnauthorized, tokenOrUserRecorder.Code) +} diff --git a/model/auth_flow.go b/model/auth_flow.go index 6d458857..a436b714 100644 --- a/model/auth_flow.go +++ b/model/auth_flow.go @@ -19,6 +19,18 @@ const ( AuthFlowIntentBind = "bind" AuthFlowTokenBytes = 32 AuthFlowDefaultCleanupRetention = 24 * time.Hour + + SecureVerificationMethod2FA = "2fa" + SecureVerificationMethodPasskey = "passkey" + SecureVerificationMethodPassword = "password" + SecureVerificationMethodOAuth = "oauth" + + SecureVerificationScopeAccessToken = "access_token" + SecureVerificationScopeAccountDelete = "account_delete" + SecureVerificationScopeCredentials = "credentials" + SecureVerificationScopeAPIToken = "api_token" + SecureVerificationScopeOAuthReauthentication = "oauth_reauthentication" + SecureVerificationScopePasskeyRegister = "passkey_register" ) var ( @@ -109,11 +121,15 @@ func CreateAuthFlow(input AuthFlowCreate) (string, *AuthFlow, error) { } func GetAuthFlow(token string, match AuthFlowMatch) (*AuthFlow, error) { + return getAuthFlowWithDB(DB, token, match) +} + +func getAuthFlowWithDB(db *gorm.DB, 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 err := applyAuthFlowMatch(db, token, match).First(&flow).Error; err != nil { if errors.Is(err, gorm.ErrRecordNotFound) { return nil, ErrAuthFlowInvalid } @@ -129,12 +145,16 @@ func GetAuthFlow(token string, match AuthFlowMatch) (*AuthFlow, error) { } func ConsumeAuthFlow(token string, match AuthFlowMatch) (*AuthFlow, error) { - flow, err := GetAuthFlow(token, match) + return consumeAuthFlowWithDB(DB, token, match) +} + +func consumeAuthFlowWithDB(db *gorm.DB, token string, match AuthFlowMatch) (*AuthFlow, error) { + flow, err := getAuthFlowWithDB(db, token, match) if err != nil { return nil, err } now := time.Now() - result := DB.Model(&AuthFlow{}). + 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 { diff --git a/model/billing_settlement.go b/model/billing_settlement.go index 90994a15..51907c07 100644 --- a/model/billing_settlement.go +++ b/model/billing_settlement.go @@ -24,10 +24,16 @@ const ( BillingSettlementEffectApplied = "applied" ) +const billingRequestOperationPrefix = "request:" +const billingRequestFinalizeSuffix = ":finalize" + +const billingSettlementBacklogSampleInterval = 30 * time.Second + var ( ErrBillingSettlementManualReview = errors.New("billing settlement requires manual review") ErrBillingSettlementTaskConflict = errors.New("billing settlement task quota conflict") ErrBillingSettlementOperationConflict = errors.New("billing settlement operation conflict") + ErrBillingSettlementRecordNotDurable = errors.New("billing settlement record was not durably created") ErrSubscriptionSettlementUnbound = errors.New("subscription settlement is not bound to its pre-consume request") ErrSubscriptionSettlementPeriodChanged = errors.New("subscription settlement crossed a quota reset period") ) @@ -61,17 +67,17 @@ type BillingSettlement struct { ID int64 `gorm:"primaryKey"` OperationKey string `gorm:"type:varchar(191);uniqueIndex;not null"` Source string `gorm:"type:varchar(32);not null"` - UserID int `gorm:"not null;default:0"` + UserID int `gorm:"not null;default:0;index:idx_billing_settlement_admission,priority:1"` SubscriptionID int `gorm:"not null;default:0"` TokenID int `gorm:"not null;default:0"` - FundingDelta int64 `gorm:"not null;default:0"` + FundingDelta int64 `gorm:"not null;default:0;index:idx_billing_settlement_admission,priority:3;index:idx_billing_settlement_status_funding,priority:2"` AppliedFundingDelta int64 `gorm:"not null;default:0"` TokenDelta int64 `gorm:"not null;default:0"` AppliedTokenDelta int64 `gorm:"not null;default:0"` TaskID int64 `gorm:"not null;default:0"` TaskQuota int64 `gorm:"not null;default:0"` TaskQuotaTarget int64 `gorm:"not null;default:0"` - Status string `gorm:"type:varchar(16);index;not null;default:'applied'"` + Status string `gorm:"type:varchar(16);index;index:idx_billing_settlement_admission,priority:2;index:idx_billing_settlement_status_funding,priority:1;not null;default:'applied'"` Attempts int `gorm:"not null;default:0"` LastError string `gorm:"type:text"` NextAttempt int64 `gorm:"index;not null;default:0"` @@ -106,16 +112,31 @@ type BillingPreConsumeSelection struct { } type BillingSettlementEffect struct { - LogType int `json:"log_type"` - Content string `json:"content"` - ChannelID int `json:"channel_id"` - ModelName string `json:"model_name"` - TokenID int `json:"token_id"` - Group string `json:"group"` - Other map[string]interface{} `json:"other"` - NodeName string `json:"node_name"` - UpdateUsage bool `json:"update_usage"` - Quota int64 `json:"quota,omitempty"` + LogType int `json:"log_type"` + Content string `json:"content"` + ChannelID int `json:"channel_id"` + ModelName string `json:"model_name"` + TokenID int `json:"token_id"` + TokenName string `json:"token_name,omitempty"` + Group string `json:"group"` + Other map[string]interface{} `json:"other"` + NodeName string `json:"node_name"` + UpdateUsage bool `json:"update_usage"` + Quota int64 `json:"quota,omitempty"` + QuotaIsActual bool `json:"quota_is_actual,omitempty"` + PromptTokens int `json:"prompt_tokens,omitempty"` + CompletionTokens int `json:"completion_tokens,omitempty"` + UseTimeSeconds int `json:"use_time_seconds,omitempty"` + IsStream bool `json:"is_stream,omitempty"` + RequestID string `json:"request_id,omitempty"` + UpstreamRequestID string `json:"upstream_request_id,omitempty"` + Subscription *BillingSettlementSubscriptionEffect `json:"subscription,omitempty"` +} + +type BillingSettlementSubscriptionEffect struct { + PreConsumed int64 `json:"pre_consumed"` + AmountTotal int64 `json:"amount_total"` + AmountUsedAfterPreConsume int64 `json:"amount_used_after_pre_consume"` } type BillingSettlementInput struct { @@ -144,6 +165,20 @@ type BillingSettlementInput struct { PreConsumeEffectiveQuota int64 } +// BillingSettlementBacklogStats is a read-only operational projection of +// positive final settlements that currently block new paid requests. +type BillingSettlementBacklogStats struct { + Count int64 `gorm:"column:record_count"` + OldestCreatedAt int64 `gorm:"column:oldest_created_at"` +} + +func unresolvedPositiveFinalizeSettlementScope(db *gorm.DB) *gorm.DB { + return db.Model(&BillingSettlement{}). + Where("funding_delta > 0"). + Where("status IN ?", []string{BillingSettlementStatusPending, BillingSettlementStatusManual}). + Where("operation_key LIKE ?", BillingRequestFinalizeOperationKey("%")) +} + var billingSettlementRunnerOnce sync.Once var billingSettlementRunnerState struct { sync.Mutex @@ -151,6 +186,20 @@ var billingSettlementRunnerState struct { done chan struct{} } +var billingSettlementBacklogObserverState struct { + sync.RWMutex + observer func(BillingSettlementBacklogStats, time.Time) +} + +// SetBillingSettlementBacklogObserver installs the operational observer used +// by the background settlement runner. The observer must remain read-only: it +// may alert operators, but it must not mutate or retry settlement records. +func SetBillingSettlementBacklogObserver(observer func(BillingSettlementBacklogStats, time.Time)) { + billingSettlementBacklogObserverState.Lock() + billingSettlementBacklogObserverState.observer = observer + billingSettlementBacklogObserverState.Unlock() +} + // StartBillingSettlementTaskRunner retries durable settlement intents after a // process restart or a transient database failure. All balance mutations are // still protected by ApplyBillingSettlementOnce's operation key. @@ -168,14 +217,19 @@ func StartBillingSettlementTaskRunner() { gopool.Go(func() { defer close(done) ProcessPendingBillingSettlementsOnce() - ticker := time.NewTicker(time.Second) - defer ticker.Stop() + observeBillingSettlementBacklog(time.Now()) + settlementTicker := time.NewTicker(time.Second) + backlogTicker := time.NewTicker(billingSettlementBacklogSampleInterval) + defer settlementTicker.Stop() + defer backlogTicker.Stop() for { select { case <-ctx.Done(): return - case <-ticker.C: + case <-settlementTicker.C: ProcessPendingBillingSettlementsOnce() + case observedAt := <-backlogTicker.C: + observeBillingSettlementBacklog(observedAt) } } }) @@ -203,6 +257,85 @@ func ProcessPendingBillingSettlementsOnce() { processPendingBillingSettlements() } +func observeBillingSettlementBacklog(observedAt time.Time) { + stats, err := GetUnresolvedPositiveFinalizeSettlementStats() + if err != nil { + common.SysLog(fmt.Sprintf("failed to query unresolved positive final billing settlements: %s", err.Error())) + return + } + billingSettlementBacklogObserverState.RLock() + observer := billingSettlementBacklogObserverState.observer + billingSettlementBacklogObserverState.RUnlock() + if observer != nil { + observer(stats, observedAt) + } +} + +// GetUnresolvedPositiveFinalizeSettlementStats returns the global count and +// oldest creation time for positive request-finalize settlements in pending or +// manual state. It is intentionally read-only and portable across all supported +// databases. +func GetUnresolvedPositiveFinalizeSettlementStats() (BillingSettlementBacklogStats, error) { + if DB == nil { + return BillingSettlementBacklogStats{}, errors.New("database is not initialized") + } + var stats BillingSettlementBacklogStats + err := unresolvedPositiveFinalizeSettlementScope(DB). + Select("COUNT(*) AS record_count, COALESCE(MIN(created_at), 0) AS oldest_created_at"). + Scan(&stats).Error + return stats, err +} + +// HasUnresolvedPositiveFinalizeSettlement reports whether a user has an +// upstream-accepted request whose positive final settlement has not completed. +// New paid requests must stop before they can consume more service against the +// residual balance. Reserve failures are deliberately excluded: no upstream +// request has been released for those operation kinds. +func HasUnresolvedPositiveFinalizeSettlement(userID int) (bool, error) { + if DB == nil { + return false, errors.New("database is not initialized") + } + if userID <= 0 { + return false, nil + } + + var record BillingSettlement + result := unresolvedPositiveFinalizeSettlementScope(DB). + Select("id"). + Where("user_id = ?", userID). + Limit(1). + Find(&record) + if result.Error != nil { + return false, result.Error + } + return result.RowsAffected > 0, nil +} + +func BillingRequestFinalizeOperationKey(requestID string) string { + return billingRequestOperationPrefix + requestID + billingRequestFinalizeSuffix +} + +// GetBillingSettlementStatus returns the durable funding state for one stable +// operation key. Effect replay is intentionally separate: once funding is +// applied, an effect-only pending state must not block async task polling. +func GetBillingSettlementStatus(operationKey string) (status string, found bool, err error) { + if DB == nil { + return "", false, errors.New("database is not initialized") + } + if operationKey == "" { + return "", false, nil + } + var record BillingSettlement + result := DB.Select("status").Where("operation_key = ?", operationKey).Limit(1).Find(&record) + if result.Error != nil { + return "", false, result.Error + } + if result.RowsAffected == 0 { + return "", false, nil + } + return record.Status, true, nil +} + func ApplyBillingSettlementOnce(input BillingSettlementInput) (appliedFundingDelta int64, alreadyApplied bool, err error) { if input.OperationKey == "" { return 0, false, errors.New("billing settlement operation key is required") @@ -276,6 +409,22 @@ func ApplyBillingSettlementOnce(input BillingSettlementInput) (appliedFundingDel resolvedTokenKey = token.Key } } + if input.TaskID > 0 { + var task Task + taskErr := withRowLock(tx). + Select("quota"). + Where("id = ?", input.TaskID). + First(&task).Error + if taskErr != nil { + if errors.Is(taskErr, gorm.ErrRecordNotFound) { + return permanentBillingSettlement(fmt.Errorf("%w: task_id=%d expected_quota=%d", ErrBillingSettlementTaskConflict, input.TaskID, input.TaskQuota)) + } + return taskErr + } + if int64(task.Quota) != input.TaskQuota { + return permanentBillingSettlement(fmt.Errorf("%w: task_id=%d expected_quota=%d", ErrBillingSettlementTaskConflict, input.TaskID, input.TaskQuota)) + } + } var applyErr error appliedFundingDelta, applyErr = applyFundingDeltaTx(tx, input) @@ -294,14 +443,16 @@ func ApplyBillingSettlementOnce(input BillingSettlementInput) (appliedFundingDel if input.TaskID > 0 { appliedTaskQuotaTarget := input.TaskQuota + appliedFundingDelta - result := tx.Model(&Task{}). - Where("id = ? AND quota = ?", input.TaskID, input.TaskQuota). - Update("quota", appliedTaskQuotaTarget) - if result.Error != nil { - return result.Error - } - if result.RowsAffected != 1 { - return permanentBillingSettlement(fmt.Errorf("%w: task_id=%d expected_quota=%d", ErrBillingSettlementTaskConflict, input.TaskID, input.TaskQuota)) + if appliedTaskQuotaTarget != input.TaskQuota { + result := tx.Model(&Task{}). + Where("id = ? AND quota = ?", input.TaskID, input.TaskQuota). + Update("quota", appliedTaskQuotaTarget) + if result.Error != nil { + return result.Error + } + if result.RowsAffected != 1 { + return permanentBillingSettlement(fmt.Errorf("%w: task_id=%d expected_quota=%d", ErrBillingSettlementTaskConflict, input.TaskID, input.TaskQuota)) + } } } @@ -463,11 +614,11 @@ func ensureBillingSettlementRecord(input BillingSettlementInput) (BillingSettlem func ensureBillingSettlementRecordDB(db *gorm.DB, input BillingSettlementInput) (BillingSettlement, bool, error) { if db == nil { - return BillingSettlement{}, false, errors.New("database is not initialized") + return BillingSettlement{}, false, fmt.Errorf("%w: database is not initialized", ErrBillingSettlementRecordNotDurable) } effectPayload, err := billingSettlementEffectPayload(input.Effect) if err != nil { - return BillingSettlement{}, false, err + return BillingSettlement{}, false, fmt.Errorf("%w: %v", ErrBillingSettlementRecordNotDurable, err) } now := time.Now().Unix() record := BillingSettlement{ @@ -491,7 +642,7 @@ func ensureBillingSettlementRecordDB(db *gorm.DB, input BillingSettlementInput) DoNothing: true, }).Create(&record) if result.Error != nil { - return BillingSettlement{}, false, result.Error + return BillingSettlement{}, false, fmt.Errorf("%w: %v", ErrBillingSettlementRecordNotDurable, result.Error) } if err := db.Where("operation_key = ?", input.OperationKey).First(&record).Error; err != nil { return BillingSettlement{}, false, err @@ -704,15 +855,20 @@ func ProcessBillingSettlementEffect(operationKey string) error { } appliedDelta := record.AppliedFundingDelta + useActualQuota := effect.QuotaIsActual || effect.Quota > 0 logQuota := appliedDelta - if effect.Quota > 0 { + if useActualQuota { logQuota = effect.Quota } - if logQuota != 0 { + if effect.UpdateUsage && logQuota < 0 { + return fmt.Errorf("billing settlement effect usage quota cannot be negative: operation=%s quota=%d", operationKey, logQuota) + } + if logQuota != 0 || effect.QuotaIsActual { other := effect.Other if other == nil { other = make(map[string]interface{}) } + applySubscriptionSettlementEffectMetadata(record, effect, other) if record.TaskID > 0 { other["actual_quota"] = record.TaskQuota + appliedDelta } @@ -723,7 +879,11 @@ func ProcessBillingSettlementEffect(operationKey string) error { if err := RecordTaskBillingLogOnce(operationKey, RecordTaskBillingLogParams{ UserId: record.UserID, LogType: effect.LogType, Content: effect.Content, ChannelId: effect.ChannelID, ModelName: effect.ModelName, Quota: int(quota), - TokenId: effect.TokenID, Group: effect.Group, Other: other, NodeName: effect.NodeName, + TokenId: effect.TokenID, TokenName: effect.TokenName, + Group: effect.Group, Other: other, NodeName: effect.NodeName, + PromptTokens: effect.PromptTokens, CompletionTokens: effect.CompletionTokens, + UseTimeSeconds: effect.UseTimeSeconds, IsStream: effect.IsStream, + RequestId: effect.RequestID, UpstreamRequestId: effect.UpstreamRequestID, }); err != nil { return err } @@ -740,10 +900,10 @@ func ProcessBillingSettlementEffect(operationKey string) error { return nil } usageQuota := appliedDelta - if effect.Quota > 0 { + if useActualQuota { usageQuota = effect.Quota } - if effect.UpdateUsage && usageQuota > 0 { + if effect.UpdateUsage && (usageQuota > 0 || effect.QuotaIsActual) { userResult := tx.Model(&User{}).Where("id = ?", record.UserID).Updates(map[string]interface{}{ "used_quota": gorm.Expr("used_quota + ?", usageQuota), "request_count": gorm.Expr("request_count + ?", 1), @@ -774,6 +934,60 @@ func ProcessBillingSettlementEffect(operationKey string) error { }) } +func BillingSettlementOwnsEffect(operationKey string) (bool, error) { + if DB == nil || operationKey == "" { + return false, errors.New("billing settlement effect operation key is required") + } + var record BillingSettlement + if err := DB.Select("effect_payload").Where("operation_key = ?", operationKey).First(&record).Error; err != nil { + return false, err + } + return record.EffectPayload != "", nil +} + +func applySubscriptionSettlementEffectMetadata( + record BillingSettlement, + effect *BillingSettlementEffect, + other map[string]interface{}, +) { + if record.Source != BillingSettlementSourceSubscription || effect == nil || + effect.Subscription == nil || other == nil { + return + } + + info := effect.Subscription + postDelta := record.AppliedFundingDelta + if postDelta != 0 { + other["subscription_post_delta"] = postDelta + } else { + delete(other, "subscription_post_delta") + } + + consumed := info.PreConsumed + postDelta + if consumed < 0 { + consumed = 0 + } + if consumed > 0 { + other["subscription_consumed"] = consumed + } else { + delete(other, "subscription_consumed") + } + + used := info.AmountUsedAfterPreConsume + postDelta + if used < 0 { + used = 0 + } + if info.AmountTotal > 0 { + remain := info.AmountTotal - used + if remain < 0 { + remain = 0 + } + other["subscription_total"] = info.AmountTotal + other["subscription_used"] = used + other["subscription_remain"] = remain + } +} + func validateBillingSettlementInput(input BillingSettlementInput) error { if input.Source != BillingSettlementSourceWallet && input.Source != BillingSettlementSourceSubscription { return permanentBillingSettlement(fmt.Errorf("unsupported billing settlement source: %s", input.Source)) diff --git a/model/billing_settlement_regression_test.go b/model/billing_settlement_regression_test.go index 51397d16..b4a89b20 100644 --- a/model/billing_settlement_regression_test.go +++ b/model/billing_settlement_regression_test.go @@ -1,16 +1,46 @@ package model import ( + "encoding/base64" "errors" + "sync" "testing" "time" "github.com/MAX-API-Next/MAX-API/common" + "github.com/glebarez/sqlite" + "github.com/go-webauthn/webauthn/protocol" + "github.com/go-webauthn/webauthn/webauthn" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "gorm.io/gorm" + "gorm.io/gorm/schema" ) +func TestTaskSchemaUsesPortableAutoIncrementPrimaryKeyAndTimeoutCursorIndex(t *testing.T) { + parsed, err := schema.Parse(&Task{}, &sync.Map{}, schema.NamingStrategy{}) + require.NoError(t, err) + + id := parsed.LookUpField("ID") + require.NotNil(t, id) + assert.True(t, id.PrimaryKey) + assert.True(t, id.AutoIncrement) + + index, ok := parsed.ParseIndexes()["idx_task_timeout_cursor"] + require.True(t, ok) + require.Len(t, index.Fields, 2) + assert.Equal(t, "submit_time", index.Fields[0].DBName) + assert.Equal(t, "id", index.Fields[1].DBName) +} + +func TestTaskSQLiteCreateGeneratesID(t *testing.T) { + truncateTables(t) + task := &Task{TaskID: "portable-auto-increment", Status: TaskStatusNotStart} + + require.NoError(t, DB.Create(task).Error) + assert.Positive(t, task.ID) +} + func TestConcurrentBillingPreConsumeSelectionAllowsOnlyOneFundingSource(t *testing.T) { setupUserUpdateTestState(t) const requestID = "regression:concurrent-funding-source" @@ -52,6 +82,65 @@ func TestConcurrentBillingPreConsumeSelectionAllowsOnlyOneFundingSource(t *testi assert.Contains(t, []string{BillingSettlementSourceWallet, BillingSettlementSourceSubscription}, source) } +func TestHasUnresolvedPositiveFinalizeSettlement(t *testing.T) { + setupUserUpdateTestState(t) + require.True(t, DB.Migrator().HasIndex(&BillingSettlement{}, "idx_billing_settlement_admission")) + require.True(t, DB.Migrator().HasIndex(&BillingSettlement{}, "idx_billing_settlement_status_funding")) + + const blockingUserID = 941 + const nonBlockingUserID = 942 + now := time.Now().Unix() + oldest := now - int64(time.Hour/time.Second) + records := []BillingSettlement{ + { + OperationKey: "request:pending-finalize:finalize", Source: BillingSettlementSourceWallet, + UserID: blockingUserID, FundingDelta: 10, TokenDelta: 10, + Status: BillingSettlementStatusPending, NextAttempt: now, CreatedAt: oldest, UpdatedAt: now, Revision: 1, + }, + { + OperationKey: "request:manual-finalize:finalize", Source: BillingSettlementSourceWallet, + UserID: blockingUserID, FundingDelta: 5, TokenDelta: 5, + Status: BillingSettlementStatusManual, CreatedAt: now, UpdatedAt: now, Revision: 1, + }, + { + OperationKey: "request:refund-finalize:finalize", Source: BillingSettlementSourceWallet, + UserID: nonBlockingUserID, FundingDelta: -5, TokenDelta: -5, + Status: BillingSettlementStatusManual, CreatedAt: now, UpdatedAt: now, Revision: 1, + }, + { + OperationKey: "request:failed-reserve:reserve:1", Source: BillingSettlementSourceWallet, + UserID: nonBlockingUserID, FundingDelta: 5, TokenDelta: 5, + Status: BillingSettlementStatusManual, CreatedAt: now, UpdatedAt: now, Revision: 1, + }, + } + require.NoError(t, DB.Create(&records).Error) + + blocked, err := HasUnresolvedPositiveFinalizeSettlement(blockingUserID) + + require.NoError(t, err) + assert.True(t, blocked) + + blocked, err = HasUnresolvedPositiveFinalizeSettlement(nonBlockingUserID) + require.NoError(t, err) + assert.False(t, blocked) + + stats, err := GetUnresolvedPositiveFinalizeSettlementStats() + require.NoError(t, err) + assert.EqualValues(t, 2, stats.Count) + assert.Equal(t, oldest, stats.OldestCreatedAt) + + require.NoError(t, DB.Where("user_id = ?", blockingUserID).Delete(&BillingSettlement{}).Error) + stats, err = GetUnresolvedPositiveFinalizeSettlementStats() + require.NoError(t, err) + assert.Zero(t, stats.Count) + assert.Zero(t, stats.OldestCreatedAt) +} + +func TestBillingRequestFinalizeOperationKey(t *testing.T) { + assert.Equal(t, "request:example-request:finalize", BillingRequestFinalizeOperationKey("example-request")) + assert.Equal(t, "request:%:finalize", BillingRequestFinalizeOperationKey("%")) +} + func TestSubscriptionSettlementUsesActualAppliedRefundForToken(t *testing.T) { setupUserUpdateTestState(t) @@ -117,6 +206,29 @@ func TestFailedSettlementLeavesDurablePendingOperation(t *testing.T) { require.EqualValues(t, 90, getRegressionTokenRemainQuota(t, token.Id)) } +func TestBillingSettlementEffectSerializationFailureIsNotDurable(t *testing.T) { + setupUserUpdateTestState(t) + input := BillingSettlementInput{ + OperationKey: "regression:effect-serialization-failure", + Source: BillingSettlementSourceWallet, + UserID: 905, + FundingDelta: 10, + Effect: &BillingSettlementEffect{ + Other: map[string]interface{}{"unsupported": func() {}}, + }, + } + + _, _, err := ApplyBillingSettlementOnce(input) + require.Error(t, err) + assert.ErrorIs(t, err, ErrBillingSettlementRecordNotDurable) + + var count int64 + require.NoError(t, DB.Model(&BillingSettlement{}). + Where("operation_key = ?", input.OperationKey). + Count(&count).Error) + assert.Zero(t, count) +} + func TestBillingSettlementPendingInvalidationBypassesStaleUserCache(t *testing.T) { setupUserUpdateTestState(t) user := User{Id: 977, Username: "settlement-pending-cache-user", Quota: 100, Status: common.UserStatusEnabled} @@ -382,6 +494,31 @@ func TestBillingSettlementTaskConflictRollsBackAllLedgerLegs(t *testing.T) { assert.Equal(t, BillingSettlementStatusManual, settlement.Status) } +func TestBillingSettlementZeroDeltaTaskConflictIsRejected(t *testing.T) { + setupUserUpdateTestState(t) + user := User{Id: 912, Username: "zero-delta-task-conflict-user", Quota: 100, Status: common.UserStatusEnabled} + task := Task{ID: 913, TaskID: "zero-delta-task-conflict", UserId: user.Id, Quota: 20, Status: TaskStatusSubmitted} + require.NoError(t, DB.Create(&user).Error) + require.NoError(t, DB.Create(&task).Error) + + _, _, err := ApplyBillingSettlementOnce(BillingSettlementInput{ + OperationKey: "regression:zero-delta-task-conflict", + Source: BillingSettlementSourceWallet, + UserID: user.Id, + TaskID: task.ID, + TaskQuota: 10, + TaskQuotaTarget: 10, + }) + require.Error(t, err) + require.ErrorIs(t, err, ErrBillingSettlementTaskConflict) + assert.EqualValues(t, 20, getRegressionTaskQuota(t, task.ID)) + assert.EqualValues(t, 100, getRegressionUserQuota(t, user.Id)) + + var settlement BillingSettlement + require.NoError(t, DB.Where("operation_key = ?", "regression:zero-delta-task-conflict").First(&settlement).Error) + assert.Equal(t, BillingSettlementStatusManual, settlement.Status) +} + func TestSubscriptionSettlementRefusesCrossResetPeriodReplay(t *testing.T) { setupUserUpdateTestState(t) user := User{Id: 912, Username: "subscription-period-user", Quota: 0, Status: common.UserStatusEnabled} @@ -619,3 +756,425 @@ func getRegressionTaskQuota(t *testing.T, taskID int64) int { require.NoError(t, DB.First(&task, taskID).Error) return task.Quota } + +func setupSecurityCredentialTestState(t *testing.T) { + t.Helper() + setupUserUpdateTestState(t) + require.NoError(t, DB.AutoMigrate( + &PasskeyCredential{}, + &TwoFA{}, + &TwoFABackupCode{}, + &AuthFlow{}, + )) + clearSecurityCredentialTestState(t) + t.Cleanup(func() { + clearSecurityCredentialTestState(t) + }) +} + +func clearSecurityCredentialTestState(t *testing.T) { + t.Helper() + for _, target := range []interface{}{ + &PasskeyCredential{}, + &TwoFABackupCode{}, + &TwoFA{}, + &AuthFlow{}, + } { + require.NoError(t, DB.Unscoped().Where("1 = 1").Delete(target).Error) + } +} + +func TestPasswordSecurityChangesBumpGenerationAndRecoveryRevokesTokens(t *testing.T) { + setupSecurityCredentialTestState(t) + + user := User{ + Id: 8801, + Username: "password-security-user", + Password: "OldPassword123", + Email: "password-security@example.com", + Role: common.RoleCommonUser, + Status: common.UserStatusEnabled, + SessionGeneration: 3, + } + user.SetAccessToken("management-access-token") + require.NoError(t, DB.Create(&user).Error) + + loaded, err := GetUserById(user.Id, true) + require.NoError(t, err) + loaded.Password = "ChangedPassword123" + require.NoError(t, loaded.Update(true)) + require.EqualValues(t, 4, loaded.SessionGeneration) + require.True(t, common.ValidatePasswordAndHash("ChangedPassword123", loaded.Password)) + + tokens := []Token{ + {Id: 8811, UserId: user.Id, Name: "enabled", Key: "password-reset-enabled", Status: common.TokenStatusEnabled}, + {Id: 8812, UserId: user.Id, Name: "disabled", Key: "password-reset-disabled", Status: common.TokenStatusDisabled}, + } + require.NoError(t, DB.Create(&tokens).Error) + + require.NoError(t, ResetUserPasswordByEmail(user.Email, "RecoveredPassword123")) + + var recovered User + require.NoError(t, DB.First(&recovered, user.Id).Error) + require.EqualValues(t, 5, recovered.SessionGeneration) + require.Empty(t, recovered.GetAccessToken()) + require.True(t, common.ValidatePasswordAndHash("RecoveredPassword123", recovered.Password)) + + var storedTokens []Token + require.NoError(t, DB.Where("user_id = ?", user.Id).Order("id asc").Find(&storedTokens).Error) + require.Len(t, storedTokens, 2) + for _, token := range storedTokens { + require.Equal(t, common.TokenStatusDisabled, token.Status) + } +} + +func TestPasskeyAndTwoFASecurityChangesBumpSessionGeneration(t *testing.T) { + setupSecurityCredentialTestState(t) + + user := User{ + Id: 8821, + Username: "credential-generation-user", + Role: common.RoleCommonUser, + Status: common.UserStatusEnabled, + } + require.NoError(t, DB.Create(&user).Error) + + firstCredential := &PasskeyCredential{ + UserID: user.Id, + CredentialID: "Y3JlZGVudGlhbC0x", + PublicKey: "cHVibGljLWtleS0x", + } + generation, err := ReplacePasskeyCredentialAndBumpSessionGeneration(firstCredential) + require.NoError(t, err) + require.EqualValues(t, 1, generation) + + replacement := &PasskeyCredential{ + UserID: user.Id, + CredentialID: "Y3JlZGVudGlhbC0y", + PublicKey: "cHVibGljLWtleS0y", + } + generation, err = ReplacePasskeyCredentialAndBumpSessionGeneration(replacement) + require.NoError(t, err) + require.EqualValues(t, 2, generation) + + var credentials []PasskeyCredential + require.NoError(t, DB.Unscoped().Where("user_id = ?", user.Id).Find(&credentials).Error) + require.Len(t, credentials, 1) + require.Equal(t, replacement.CredentialID, credentials[0].CredentialID) + + generation, err = DeletePasskeyAndBumpSessionGeneration(user.Id) + require.NoError(t, err) + require.EqualValues(t, 3, generation) + + twoFA := &TwoFA{UserId: user.Id, Secret: "TESTSECRET", IsEnabled: false} + require.NoError(t, DB.Create(twoFA).Error) + generation, err = twoFA.EnableAndBumpSessionGeneration() + require.NoError(t, err) + require.EqualValues(t, 4, generation) + + generation, err = ReplaceBackupCodesAndBumpSessionGeneration(user.Id, []string{"ABCD-EFGH"}) + require.NoError(t, err) + require.EqualValues(t, 5, generation) + var backupCode TwoFABackupCode + require.NoError(t, DB.Where("user_id = ?", user.Id).First(&backupCode).Error) + require.NotEqual(t, "ABCD-EFGH", backupCode.CodeHash) + require.True(t, common.ValidatePasswordAndHash("ABCD-EFGH", backupCode.CodeHash)) + valid, err := ValidateBackupCode(user.Id, "ABCD-EFGH") + require.NoError(t, err) + require.True(t, valid) + valid, err = ValidateBackupCode(user.Id, "ABCD-EFGH") + require.NoError(t, err) + require.False(t, valid) + + generation, err = DisableTwoFAAndBumpSessionGeneration(user.Id) + require.NoError(t, err) + require.EqualValues(t, 6, generation) + + var stored User + require.NoError(t, DB.First(&stored, user.Id).Error) + require.EqualValues(t, 6, stored.SessionGeneration) +} + +func TestBackupCodeMutationsRequireCredentialOwner(t *testing.T) { + setupSecurityCredentialTestState(t) + const missingUserID = 8899 + + require.EqualError(t, (&TwoFA{Id: 1}).Delete(), "2FA用户ID不能为空") + require.EqualError(t, (&TwoFA{Id: 1, UserId: -1}).Delete(), "2FA用户ID不能为空") + require.ErrorIs(t, CreateBackupCodes(missingUserID, []string{"ABCD-EFGH"}), gorm.ErrRecordNotFound) + require.ErrorIs(t, (&TwoFA{Id: 1, UserId: missingUserID}).Delete(), gorm.ErrRecordNotFound) +} + +func TestPasskeyUsageUpdateCannotRestoreReplacedCredential(t *testing.T) { + setupSecurityCredentialTestState(t) + + user := User{ + Id: 8822, + Username: "passkey-usage-race-user", + Role: common.RoleCommonUser, + Status: common.UserStatusEnabled, + } + require.NoError(t, DB.Create(&user).Error) + + staleCredential := &PasskeyCredential{ + UserID: user.Id, + CredentialID: "c3RhbGUtY3JlZGVudGlhbA==", + PublicKey: "c3RhbGUtcHVibGljLWtleQ==", + SignCount: 1, + } + _, err := ReplacePasskeyCredentialAndBumpSessionGeneration(staleCredential) + require.NoError(t, err) + staleUsageUpdate := *staleCredential + + replacement := &PasskeyCredential{ + UserID: user.Id, + CredentialID: "cmVwbGFjZW1lbnQtY3JlZGVudGlhbA==", + PublicKey: "cmVwbGFjZW1lbnQtcHVibGljLWtleQ==", + SignCount: 10, + } + _, err = ReplacePasskeyCredentialAndBumpSessionGeneration(replacement) + require.NoError(t, err) + + now := time.Now() + staleUsageUpdate.SignCount = 2 + staleUsageUpdate.LastUsedAt = &now + require.ErrorIs(t, UpdatePasskeyCredentialAfterAuthentication(&staleUsageUpdate), ErrPasskeyCredentialChanged) + + stored, err := GetPasskeyByUserID(user.Id) + require.NoError(t, err) + require.Equal(t, replacement.CredentialID, stored.CredentialID) + require.Equal(t, replacement.PublicKey, stored.PublicKey) + require.EqualValues(t, 10, stored.SignCount) + + later := now.Add(time.Second) + matchingUsageUpdate := *replacement + matchingUsageUpdate.ID = 0 + matchingUsageUpdate.SignCount = 11 + matchingUsageUpdate.LastUsedAt = &later + require.NoError(t, UpdatePasskeyCredentialAfterAuthentication(&matchingUsageUpdate)) + stored, err = GetPasskeyByUserID(user.Id) + require.NoError(t, err) + require.Equal(t, replacement.CredentialID, stored.CredentialID) + require.EqualValues(t, 11, stored.SignCount) + require.WithinDuration(t, later, *stored.LastUsedAt, time.Second) +} + +func TestPasskeyUsageUpdateMergesOrderedAuthenticationStateMonotonically(t *testing.T) { + setupSecurityCredentialTestState(t) + + user := User{ + Id: 8824, + Username: "passkey-monotonic-auth-user", + Role: common.RoleCommonUser, + Status: common.UserStatusEnabled, + } + require.NoError(t, DB.Create(&user).Error) + + credential := &PasskeyCredential{ + UserID: user.Id, + CredentialID: "bW9ub3RvbmljLWNlcnQ=", + PublicKey: "bW9ub3RvbmljLXB1YmxpYy1rZXk=", + SignCount: 9, + } + _, err := ReplacePasskeyCredentialAndBumpSessionGeneration(credential) + require.NoError(t, err) + + older := time.Now() + newer := older.Add(time.Second) + higherCounter := *credential + higherCounter.SignCount = 11 + higherCounter.CloneWarning = true + higherCounter.LastUsedAt = &newer + lowerCounter := *credential + lowerCounter.SignCount = 10 + lowerCounter.CloneWarning = false + lowerCounter.LastUsedAt = &older + + require.NoError(t, UpdatePasskeyCredentialAfterAuthentication(&higherCounter)) + require.NoError(t, UpdatePasskeyCredentialAfterAuthentication(&lowerCounter)) + + stored, err := GetPasskeyByUserID(user.Id) + require.NoError(t, err) + require.EqualValues(t, 11, stored.SignCount) + require.True(t, stored.CloneWarning) + require.WithinDuration(t, newer, *stored.LastUsedAt, time.Second) +} + +func TestPasskeyAuthenticationUpdateKeysMatchSchema(t *testing.T) { + setupSecurityCredentialTestState(t) + now := time.Now() + updates := passkeyAuthenticationUpdates(&PasskeyCredential{ + PublicKey: "schema-public-key", CloneWarning: true, LastUsedAt: &now, + }) + statement := &gorm.Statement{DB: DB} + require.NoError(t, statement.Parse(&PasskeyCredential{})) + for key := range updates { + require.Contains(t, statement.Schema.DBNames, key) + } +} + +func TestTimedOutTaskCursorRemainsGroupedWithStatusPredicate(t *testing.T) { + truncateTables(t) + require.True(t, DB.Migrator().HasIndex(&Task{}, "idx_task_timeout_cursor")) + const submitTime = int64(100) + tasks := []Task{ + {ID: 401, TaskID: "cursor-unfinished", Status: TaskStatusInProgress, SubmitTime: submitTime}, + {ID: 402, TaskID: "cursor-success", Status: TaskStatusSuccess, SubmitTime: submitTime}, + {ID: 403, TaskID: "cursor-failure", Status: TaskStatusFailure, SubmitTime: submitTime}, + } + require.NoError(t, DB.Create(&tasks).Error) + + got, err := GetTimedOutUnfinishedTasksAfter(200, submitTime, 400, 10) + require.NoError(t, err) + require.Len(t, got, 1) + require.Equal(t, "cursor-unfinished", got[0].TaskID) +} + +func TestTimedOutTaskCursorReturnsQueryError(t *testing.T) { + failingDB, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + require.NoError(t, err) + sqlDB, err := failingDB.DB() + require.NoError(t, err) + require.NoError(t, sqlDB.Close()) + + originalDB := DB + DB = failingDB + t.Cleanup(func() { DB = originalDB }) + + tasks, queryErr := GetTimedOutUnfinishedTasksAfter(200, 0, 0, 10) + require.Error(t, queryErr) + require.Nil(t, tasks) +} + +func TestPasskeyValidatedCredentialStateIsPersisted(t *testing.T) { + setupSecurityCredentialTestState(t) + + user := User{ + Id: 8823, + Username: "passkey-validated-state-user", + Role: common.RoleCommonUser, + Status: common.UserStatusEnabled, + } + require.NoError(t, DB.Create(&user).Error) + + credentialID := []byte("validated-state-credential") + storedCredential := &PasskeyCredential{ + UserID: user.Id, + CredentialID: base64.StdEncoding.EncodeToString(credentialID), + PublicKey: base64.StdEncoding.EncodeToString([]byte("old-public-key")), + SignCount: 1, + } + _, err := ReplacePasskeyCredentialAndBumpSessionGeneration(storedCredential) + require.NoError(t, err) + + validatedCredential := &webauthn.Credential{ + ID: credentialID, + PublicKey: []byte("updated-public-key"), + AttestationType: "none", + Transport: []protocol.AuthenticatorTransport{protocol.USB}, + Flags: webauthn.CredentialFlags{ + UserPresent: true, + UserVerified: true, + BackupEligible: true, + BackupState: true, + }, + Authenticator: webauthn.Authenticator{ + AAGUID: []byte("updated-aaguid"), + SignCount: 2, + CloneWarning: true, + Attachment: protocol.Platform, + }, + } + + storedCredential.ApplyValidatedCredential(validatedCredential) + now := time.Now() + storedCredential.LastUsedAt = &now + require.NoError(t, UpdatePasskeyCredentialAfterAuthentication(storedCredential)) + + updated, err := GetPasskeyByUserID(user.Id) + require.NoError(t, err) + require.Equal(t, base64.StdEncoding.EncodeToString(credentialID), updated.CredentialID) + require.Equal(t, base64.StdEncoding.EncodeToString(validatedCredential.PublicKey), updated.PublicKey) + require.EqualValues(t, 2, updated.SignCount) + require.True(t, updated.CloneWarning) + require.True(t, updated.UserPresent) + require.True(t, updated.UserVerified) + require.True(t, updated.BackupEligible) + require.True(t, updated.BackupState) + require.Equal(t, string(protocol.Platform), updated.Attachment) + require.Equal(t, []protocol.AuthenticatorTransport{protocol.USB}, updated.TransportList()) + require.WithinDuration(t, now, *updated.LastUsedAt, time.Second) +} + +func TestTelegramBindingStateIsUserBoundAndConsumedOnce(t *testing.T) { + setupSecurityCredentialTestState(t) + + owner := User{Id: 8831, Username: "telegram-owner", AffCode: "telegram-owner", Role: common.RoleCommonUser, Status: common.UserStatusEnabled} + other := User{Id: 8832, Username: "telegram-other", AffCode: "telegram-other", Role: common.RoleCommonUser, Status: common.UserStatusEnabled} + require.NoError(t, DB.Create(&owner).Error) + require.NoError(t, DB.Create(&other).Error) + + state, _, err := CreateAuthFlow(AuthFlowCreate{ + Purpose: AuthFlowPurposeOAuth, + Provider: "telegram", + Intent: AuthFlowIntentBind, + UserId: owner.Id, + ExpiresAt: time.Now().Add(time.Minute), + }) + require.NoError(t, err) + + _, err = BindTelegramIdentityWithAuthFlow(other.Id, "123456", state) + require.ErrorIs(t, err, ErrAuthFlowInvalid) + + generation, err := BindTelegramIdentityWithAuthFlow(owner.Id, "123456", state) + require.NoError(t, err) + require.EqualValues(t, 1, generation) + + _, err = BindTelegramIdentityWithAuthFlow(owner.Id, "123456", state) + require.ErrorIs(t, err, ErrAuthFlowConsumed) + + var stored User + require.NoError(t, DB.First(&stored, owner.Id).Error) + require.Equal(t, "123456", stored.TelegramId) + require.EqualValues(t, 1, stored.SessionGeneration) +} + +func TestOAuthReauthenticationBootstrapRequiresNoExistingCredentials(t *testing.T) { + setupSecurityCredentialTestState(t) + + users := []User{ + {Id: 8841, Username: "oauth-bootstrap-user", AffCode: "oauth-bootstrap", Role: common.RoleCommonUser, Status: common.UserStatusEnabled}, + {Id: 8842, Username: "oauth-password-user", Password: "existing-password-hash", AffCode: "oauth-password", Role: common.RoleCommonUser, Status: common.UserStatusEnabled}, + {Id: 8843, Username: "oauth-two-fa-user", AffCode: "oauth-two-fa", Role: common.RoleCommonUser, Status: common.UserStatusEnabled}, + {Id: 8844, Username: "oauth-passkey-user", AffCode: "oauth-passkey", Role: common.RoleCommonUser, Status: common.UserStatusEnabled}, + } + require.NoError(t, DB.Create(&users).Error) + require.NoError(t, DB.Create(&TwoFA{ + UserId: users[2].Id, + Secret: "oauth-reauth-test-secret", + IsEnabled: true, + }).Error) + require.NoError(t, DB.Create(&PasskeyCredential{ + UserID: users[3].Id, + CredentialID: "oauth-reauth-credential-id", + PublicKey: "oauth-reauth-public-key", + }).Error) + + tests := []struct { + name string + userID int + allowed bool + }{ + {name: "oauth only", userID: users[0].Id, allowed: true}, + {name: "password", userID: users[1].Id, allowed: false}, + {name: "enabled 2FA", userID: users[2].Id, allowed: false}, + {name: "Passkey", userID: users[3].Id, allowed: false}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + allowed, err := CanUseOAuthReauthentication(test.userID) + require.NoError(t, err) + require.Equal(t, test.allowed, allowed) + }) + } +} diff --git a/model/log.go b/model/log.go index b27f6882..22fa70ed 100644 --- a/model/log.go +++ b/model/log.go @@ -588,10 +588,20 @@ func RecordConsumeLog(c *gin.Context, userId int, params RecordConsumeLogParams) 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) - upstreamRequestId := c.GetString(common.UpstreamRequestIdKey) + logCtx := context.Background() + if c != nil { + logCtx = c + } + logger.LogInfo(logCtx, fmt.Sprintf("record consume log: userId=%d, params=%s", userId, common.GetJsonString(params))) + username, requestId, upstreamRequestId := "", "", "" + if c != nil { + username = c.GetString("username") + requestId = c.GetString(common.RequestIdKey) + upstreamRequestId = c.GetString(common.UpstreamRequestIdKey) + } + if username == "" { + username, _ = GetUsernameById(userId, false) + } otherStr := common.MapToJsonStr(params.Other) // 判断是否需要记录 IP needRecordIp := false @@ -617,7 +627,7 @@ func RecordConsumeLog(c *gin.Context, userId int, params RecordConsumeLogParams) IsStream: params.IsStream, Group: params.Group, Ip: func() string { - if needRecordIp { + if needRecordIp && c != nil { return c.ClientIP() } return "" @@ -628,7 +638,7 @@ func RecordConsumeLog(c *gin.Context, userId int, params RecordConsumeLogParams) } err := LOG_DB.Create(log).Error if err != nil { - logger.LogError(c, "failed to record log: "+err.Error()) + logger.LogError(logCtx, "failed to record log: "+err.Error()) } if common.DataExportEnabled { createdAt := common.GetTimestamp() @@ -648,16 +658,23 @@ func RecordConsumeLog(c *gin.Context, userId int, params RecordConsumeLogParams) } type RecordTaskBillingLogParams struct { - UserId int - LogType int - Content string - ChannelId int - ModelName string - Quota int - TokenId int - Group string - Other map[string]interface{} - NodeName string // 任务发起节点;为空时回退当前节点 + UserId int + LogType int + Content string + ChannelId int + ModelName string + Quota int + TokenId int + TokenName string + Group string + Other map[string]interface{} + NodeName string // 任务发起节点;为空时回退当前节点 + PromptTokens int + CompletionTokens int + UseTimeSeconds int + IsStream bool + RequestId string + UpstreamRequestId string } func RecordTaskBillingLog(params RecordTaskBillingLogParams) { @@ -682,26 +699,32 @@ func recordTaskBillingLog(operationKey string, params RecordTaskBillingLogParams } params.Content = common.SanitizePersistedLogContent(params.Content) username, _ := GetUsernameById(params.UserId, false) - tokenName := "" + tokenName := params.TokenName if params.TokenId > 0 { - if token, err := GetTokenById(params.TokenId); err == nil { + if token, err := GetTokenById(params.TokenId); err == nil && token.Name != "" { tokenName = token.Name } } createdAt := common.GetTimestamp() log := &Log{ - UserId: params.UserId, - Username: username, - CreatedAt: createdAt, - Type: params.LogType, - Content: params.Content, - TokenName: tokenName, - ModelName: params.ModelName, - Quota: params.Quota, - ChannelId: params.ChannelId, - TokenId: params.TokenId, - Group: params.Group, - Other: common.MapToJsonStr(params.Other), + UserId: params.UserId, + Username: username, + CreatedAt: createdAt, + Type: params.LogType, + Content: params.Content, + TokenName: tokenName, + ModelName: params.ModelName, + Quota: params.Quota, + PromptTokens: params.PromptTokens, + CompletionTokens: params.CompletionTokens, + UseTime: params.UseTimeSeconds, + IsStream: params.IsStream, + ChannelId: params.ChannelId, + TokenId: params.TokenId, + Group: params.Group, + RequestId: params.RequestId, + UpstreamRequestId: params.UpstreamRequestId, + Other: common.MapToJsonStr(params.Other), } inserted := true if operationKey == "" { @@ -754,6 +777,7 @@ func recordTaskBillingLog(operationKey string, params RecordTaskBillingLogParams ModelName: params.ModelName, Quota: params.Quota, CreatedAt: createdAt, + TokenUsed: params.PromptTokens + params.CompletionTokens, UseGroup: params.Group, TokenID: params.TokenId, ChannelID: params.ChannelId, diff --git a/model/log_test.go b/model/log_test.go index bcddc89d..cf8714c5 100644 --- a/model/log_test.go +++ b/model/log_test.go @@ -75,27 +75,18 @@ func markRetryLogBackfillCompletedForTest(t *testing.T) { func resetQuotaDataCacheForTest(t *testing.T) { t.Helper() - resetLogQuotaDataShutdownForTest(t) - CacheQuotaDataLock.Lock() - CacheQuotaData = make(map[string]*QuotaData) - CacheQuotaDataLock.Unlock() - t.Cleanup(func() { + drainAndReset := func() { + WaitPendingLogQuotaData() CacheQuotaDataLock.Lock() CacheQuotaData = make(map[string]*QuotaData) CacheQuotaDataLock.Unlock() - }) -} -func resetLogQuotaDataShutdownForTest(t *testing.T) { - t.Helper() - logQuotaDataShutdownMu.Lock() - logQuotaDataShutdownStarted = false - logQuotaDataShutdownMu.Unlock() - t.Cleanup(func() { logQuotaDataShutdownMu.Lock() logQuotaDataShutdownStarted = false logQuotaDataShutdownMu.Unlock() - }) + } + drainAndReset() + t.Cleanup(drainAndReset) } func newLogTestContext() *gin.Context { @@ -136,6 +127,44 @@ func TestRecordConsumeLogSanitizesPersistedContent(t *testing.T) { require.LessOrEqual(t, utf8.RuneCountInString(log.Content), common.PersistedLogContentLimit) } +func TestRecordConsumeLogWithNilContextPreservesAccountingFields(t *testing.T) { + originalLogConsumeEnabled := common.LogConsumeEnabled + originalDataExportEnabled := common.DataExportEnabled + common.LogConsumeEnabled = true + common.DataExportEnabled = false + t.Cleanup(func() { + common.LogConsumeEnabled = originalLogConsumeEnabled + common.DataExportEnabled = originalDataExportEnabled + }) + const userID = 72007 + require.NoError(t, LOG_DB.Where("1 = 1").Delete(&Log{}).Error) + require.NoError(t, DB.Where("id = ?", userID).Delete(&User{}).Error) + require.NoError(t, DB.Create(&User{ + Id: userID, + Username: "nil-context-log-user", + AffCode: "nil-context-log-user-aff", + Status: common.UserStatusEnabled, + }).Error) + t.Cleanup(func() { + require.NoError(t, LOG_DB.Where("1 = 1").Delete(&Log{}).Error) + require.NoError(t, DB.Where("id = ?", userID).Delete(&User{}).Error) + }) + + RecordConsumeLog(nil, userID, RecordConsumeLogParams{ + ModelName: "nil-context-model", + Quota: 42, + PromptTokens: 3, + CompletionTokens: 5, + }) + + var log Log + require.NoError(t, LOG_DB.Where("model_name = ?", "nil-context-model").First(&log).Error) + require.Equal(t, 42, log.Quota) + require.Equal(t, 3, log.PromptTokens) + require.Equal(t, 5, log.CompletionTokens) + require.Equal(t, "nil-context-log-user", log.Username) +} + func TestBillingSettlementEffectPayloadSanitizesContent(t *testing.T) { payload, err := billingSettlementEffectPayload(&BillingSettlementEffect{ LogType: LogTypeConsume, @@ -908,14 +937,16 @@ func TestRecordTaskBillingLogQueuesQuotaDataAsync(t *testing.T) { }) RecordTaskBillingLog(RecordTaskBillingLogParams{ - UserId: 7, - LogType: LogTypeConsume, - Content: "task billing", - ChannelId: 13, - ModelName: "task-test", - Quota: 42, - TokenId: 11, - Group: "default", + UserId: 7, + LogType: LogTypeConsume, + Content: "task billing", + ChannelId: 13, + ModelName: "task-test", + Quota: 42, + TokenId: 11, + Group: "default", + PromptTokens: 3, + CompletionTokens: 5, }) require.Len(t, queued, 1) @@ -923,6 +954,262 @@ func TestRecordTaskBillingLogQueuesQuotaDataAsync(t *testing.T) { queued[0]() require.Len(t, CacheQuotaData, 1) + for _, quotaData := range CacheQuotaData { + require.Equal(t, 8, quotaData.TokenUsed) + } +} + +func TestBillingSettlementEffectReplayExportsTokenUsage(t *testing.T) { + const userID = 7001 + require.NoError(t, DB.Where("1 = 1").Delete(&BillingSettlement{}).Error) + require.NoError(t, LOG_DB.Where("1 = 1").Delete(&BillingLogReceipt{}).Error) + require.NoError(t, LOG_DB.Where("1 = 1").Delete(&Log{}).Error) + t.Cleanup(func() { + require.NoError(t, DB.Where("1 = 1").Delete(&BillingSettlement{}).Error) + require.NoError(t, DB.Where("id = ?", userID).Delete(&User{}).Error) + require.NoError(t, LOG_DB.Where("1 = 1").Delete(&BillingLogReceipt{}).Error) + require.NoError(t, LOG_DB.Where("1 = 1").Delete(&Log{}).Error) + }) + resetQuotaDataCacheForTest(t) + + originalDataExportEnabled := common.DataExportEnabled + originalRunner := logQuotaDataAsyncRunner + common.DataExportEnabled = true + logQuotaDataAsyncRunner = func(fn func()) { fn() } + t.Cleanup(func() { + common.DataExportEnabled = originalDataExportEnabled + logQuotaDataAsyncRunner = originalRunner + }) + + require.NoError(t, DB.Create(&User{Id: userID, Username: "settlement-export-user", Status: common.UserStatusEnabled}).Error) + const operationKey = "request:settlement-export-token-usage" + _, _, err := ApplyBillingSettlementOnce(BillingSettlementInput{ + OperationKey: operationKey, + Source: BillingSettlementSourceWallet, + UserID: userID, + Effect: &BillingSettlementEffect{ + LogType: LogTypeConsume, + Content: "settlement effect", + ModelName: "settlement-model", + TokenID: 7002, + TokenName: "request-time-deleted-token", + UpdateUsage: true, + Quota: 42, + PromptTokens: 3, + CompletionTokens: 5, + RequestID: "settlement-export-request", + }, + }) + require.NoError(t, err) + var settlement BillingSettlement + require.NoError(t, DB.Where("operation_key = ?", operationKey).First(&settlement).Error) + require.Equal(t, BillingSettlementStatusApplied, settlement.Status) + require.Equal(t, BillingSettlementEffectPending, settlement.EffectStatus) + require.NoError(t, ProcessBillingSettlementEffect(operationKey)) + require.NoError(t, DB.First(&settlement, settlement.ID).Error) + require.Equal(t, BillingSettlementEffectApplied, settlement.EffectStatus) + + var exportedTokenUsed int + require.Len(t, CacheQuotaData, 1) + for _, quotaData := range CacheQuotaData { + exportedTokenUsed = quotaData.TokenUsed + } + require.Equal(t, 8, exportedTokenUsed) + var log Log + require.NoError(t, LOG_DB.Where("request_id = ?", "settlement-export-request").First(&log).Error) + require.Equal(t, 3, log.PromptTokens) + require.Equal(t, 5, log.CompletionTokens) + require.Equal(t, "request-time-deleted-token", log.TokenName) +} + +func TestApplySubscriptionSettlementEffectMetadataUsesAppliedDelta(t *testing.T) { + tests := []struct { + name string + delta int64 + wantPost bool + wantConsumed bool + consumed int64 + used int64 + remain int64 + }{ + {name: "positive settlement", delta: 5, wantPost: true, wantConsumed: true, consumed: 15, used: 45, remain: 55}, + {name: "zero settlement", wantConsumed: true, consumed: 10, used: 40, remain: 60}, + {name: "full reservation refund", delta: -10, wantPost: true, used: 30, remain: 70}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + other := map[string]interface{}{ + "subscription_post_delta": 99, + "subscription_consumed": 99, + } + applySubscriptionSettlementEffectMetadata(BillingSettlement{ + Source: BillingSettlementSourceSubscription, + AppliedFundingDelta: tt.delta, + }, &BillingSettlementEffect{ + Subscription: &BillingSettlementSubscriptionEffect{ + PreConsumed: 10, + AmountTotal: 100, + AmountUsedAfterPreConsume: 40, + }, + }, other) + + if tt.wantPost { + require.EqualValues(t, tt.delta, other["subscription_post_delta"]) + } else { + require.NotContains(t, other, "subscription_post_delta") + } + if tt.wantConsumed { + require.EqualValues(t, tt.consumed, other["subscription_consumed"]) + } else { + require.NotContains(t, other, "subscription_consumed") + } + require.EqualValues(t, 100, other["subscription_total"]) + require.EqualValues(t, tt.used, other["subscription_used"]) + require.EqualValues(t, tt.remain, other["subscription_remain"]) + }) + } +} + +func TestBillingSettlementEffectReplayPrefersCurrentTokenName(t *testing.T) { + const ( + userID = 7011 + tokenID = 7012 + ) + require.NoError(t, DB.Where("1 = 1").Delete(&BillingSettlement{}).Error) + require.NoError(t, DB.Where("1 = 1").Delete(&Token{}).Error) + require.NoError(t, LOG_DB.Where("1 = 1").Delete(&BillingLogReceipt{}).Error) + require.NoError(t, LOG_DB.Where("1 = 1").Delete(&Log{}).Error) + t.Cleanup(func() { + require.NoError(t, DB.Where("1 = 1").Delete(&BillingSettlement{}).Error) + require.NoError(t, DB.Where("1 = 1").Delete(&Token{}).Error) + require.NoError(t, DB.Where("id = ?", userID).Delete(&User{}).Error) + require.NoError(t, LOG_DB.Where("1 = 1").Delete(&BillingLogReceipt{}).Error) + require.NoError(t, LOG_DB.Where("1 = 1").Delete(&Log{}).Error) + }) + + require.NoError(t, DB.Create(&User{ + Id: userID, Username: "settlement-token-user", AffCode: "settlement-token-user-aff", + Status: common.UserStatusEnabled, + }).Error) + require.NoError(t, DB.Create(&Token{ + Id: tokenID, UserId: userID, Name: "current-token-name", Key: "settlement-token-key", + Status: common.TokenStatusEnabled, + }).Error) + const operationKey = "request:settlement-current-token-name" + _, _, err := ApplyBillingSettlementOnce(BillingSettlementInput{ + OperationKey: operationKey, + Source: BillingSettlementSourceWallet, + UserID: userID, + Effect: &BillingSettlementEffect{ + LogType: LogTypeConsume, + Content: "settlement token name precedence", + TokenID: tokenID, + TokenName: "request-time-token-name", + Quota: 1, + QuotaIsActual: true, + RequestID: "settlement-token-name-request", + }, + }) + require.NoError(t, err) + require.NoError(t, ProcessBillingSettlementEffect(operationKey)) + + var log Log + require.NoError(t, LOG_DB.Where("request_id = ?", "settlement-token-name-request").First(&log).Error) + require.Equal(t, "current-token-name", log.TokenName) +} + +func TestBillingSettlementEffectReplayPreservesRequestTokenNameWhenCurrentNameEmpty(t *testing.T) { + const ( + userID = 7013 + tokenID = 7014 + ) + require.NoError(t, DB.Where("1 = 1").Delete(&BillingSettlement{}).Error) + require.NoError(t, DB.Where("1 = 1").Delete(&Token{}).Error) + require.NoError(t, LOG_DB.Where("1 = 1").Delete(&BillingLogReceipt{}).Error) + require.NoError(t, LOG_DB.Where("1 = 1").Delete(&Log{}).Error) + t.Cleanup(func() { + require.NoError(t, DB.Where("1 = 1").Delete(&BillingSettlement{}).Error) + require.NoError(t, DB.Where("1 = 1").Delete(&Token{}).Error) + require.NoError(t, DB.Where("id = ?", userID).Delete(&User{}).Error) + require.NoError(t, LOG_DB.Where("1 = 1").Delete(&BillingLogReceipt{}).Error) + require.NoError(t, LOG_DB.Where("1 = 1").Delete(&Log{}).Error) + }) + + require.NoError(t, DB.Create(&User{ + Id: userID, Username: "empty-token-name-user", AffCode: "empty-token-name-user-aff", + Status: common.UserStatusEnabled, + }).Error) + require.NoError(t, DB.Create(&Token{ + Id: tokenID, UserId: userID, Name: "", Key: "empty-name-token-key", + Status: common.TokenStatusEnabled, + }).Error) + const operationKey = "request:settlement-empty-current-token-name" + _, _, err := ApplyBillingSettlementOnce(BillingSettlementInput{ + OperationKey: operationKey, + Source: BillingSettlementSourceWallet, + UserID: userID, + Effect: &BillingSettlementEffect{ + LogType: LogTypeConsume, + Content: "settlement empty token name fallback", + TokenID: tokenID, + TokenName: "request-time-token-name", + Quota: 1, + QuotaIsActual: true, + RequestID: "settlement-empty-token-name-request", + }, + }) + require.NoError(t, err) + require.NoError(t, ProcessBillingSettlementEffect(operationKey)) + + var log Log + require.NoError(t, LOG_DB.Where("request_id = ?", "settlement-empty-token-name-request").First(&log).Error) + require.Equal(t, "request-time-token-name", log.TokenName) +} + +func TestBillingSettlementEffectRejectsNegativeUsageProjection(t *testing.T) { + const userID = 7021 + require.NoError(t, DB.Where("1 = 1").Delete(&BillingSettlement{}).Error) + require.NoError(t, LOG_DB.Where("1 = 1").Delete(&BillingLogReceipt{}).Error) + require.NoError(t, LOG_DB.Where("1 = 1").Delete(&Log{}).Error) + t.Cleanup(func() { + require.NoError(t, DB.Where("1 = 1").Delete(&BillingSettlement{}).Error) + require.NoError(t, DB.Where("id = ?", userID).Delete(&User{}).Error) + require.NoError(t, LOG_DB.Where("1 = 1").Delete(&BillingLogReceipt{}).Error) + require.NoError(t, LOG_DB.Where("1 = 1").Delete(&Log{}).Error) + }) + + require.NoError(t, DB.Create(&User{ + Id: userID, Username: "negative-effect-user", AffCode: "negative-effect-user-aff", + Status: common.UserStatusEnabled, UsedQuota: 12, RequestCount: 3, + }).Error) + const operationKey = "request:negative-usage-effect" + _, _, err := ApplyBillingSettlementOnce(BillingSettlementInput{ + OperationKey: operationKey, + Source: BillingSettlementSourceWallet, + UserID: userID, + Effect: &BillingSettlementEffect{ + LogType: LogTypeConsume, + Content: "negative usage must not project", + UpdateUsage: true, + Quota: -5, + QuotaIsActual: true, + RequestID: "negative-usage-effect-request", + }, + }) + require.NoError(t, err) + require.ErrorContains(t, ProcessBillingSettlementEffect(operationKey), "usage quota cannot be negative") + + var user User + require.NoError(t, DB.First(&user, userID).Error) + require.EqualValues(t, 12, user.UsedQuota) + require.Equal(t, 3, user.RequestCount) + var settlement BillingSettlement + require.NoError(t, DB.Where("operation_key = ?", operationKey).First(&settlement).Error) + require.Equal(t, BillingSettlementEffectPending, settlement.EffectStatus) + var logCount int64 + require.NoError(t, LOG_DB.Where("request_id = ?", "negative-usage-effect-request").Model(&Log{}).Count(&logCount).Error) + require.Zero(t, logCount) } func TestWaitPendingLogQuotaDataDrainsEnqueuedWork(t *testing.T) { @@ -970,6 +1257,56 @@ func TestWaitPendingLogQuotaDataDrainsEnqueuedWork(t *testing.T) { require.Len(t, CacheQuotaData, 1) } +func TestResetQuotaDataCacheForTestWaitsForPendingWorkers(t *testing.T) { + resetQuotaDataCacheForTest(t) + + originalRunner := logQuotaDataAsyncRunner + started := make(chan struct{}) + release := make(chan struct{}) + logQuotaDataAsyncRunner = func(fn func()) { + go func() { + close(started) + <-release + fn() + }() + } + t.Cleanup(func() { + logQuotaDataAsyncRunner = originalRunner + }) + + enqueueLogQuotaData(QuotaDataLogParams{ + UserID: 8, + Username: "quota-reset", + ModelName: "gpt-test", + Quota: 42, + CreatedAt: time.Now().Unix(), + }) + <-started + + resetDone := make(chan struct{}) + go func() { + resetQuotaDataCacheForTest(t) + close(resetDone) + }() + + returnedBeforeRelease := false + select { + case <-resetDone: + returnedBeforeRelease = true + case <-time.After(10 * time.Millisecond): + } + + close(release) + select { + case <-resetDone: + case <-time.After(time.Second): + t.Fatal("timed out resetting quota data cache") + } + + require.False(t, returnedBeforeRelease, "reset returned before pending quota data worker completed") + require.Empty(t, CacheQuotaData) +} + func TestEnqueueLogQuotaDataAfterShutdownPersistsSynchronously(t *testing.T) { resetQuotaDataCacheForTest(t) require.NoError(t, DB.AutoMigrate(&QuotaData{}, &QuotaDataSnapshot{})) diff --git a/model/option.go b/model/option.go index 9dc69724..80f13b4c 100644 --- a/model/option.go +++ b/model/option.go @@ -247,12 +247,16 @@ func UpdateOption(key string, value string) error { Key: key, } // https://gorm.io/docs/update.html#Save-All-Fields - DB.FirstOrCreate(&option, Option{Key: key}) + if err := DB.FirstOrCreate(&option, Option{Key: key}).Error; err != nil { + return err + } option.Value = value // Save is a combination function. // If save value does not contain primary key, it will execute Create, // otherwise it will execute Update (with all fields). - DB.Save(&option) + if err := DB.Save(&option).Error; err != nil { + return err + } // Update OptionMap return updateOptionMap(key, value) } @@ -317,6 +321,12 @@ func requireRegisteredOptionKey(key string) error { func normalizeOptionUpdateValue(key string, value string) (string, error) { switch key { + case "PreConsumedQuota": + quota, err := parsePreConsumedQuota(value) + if err != nil { + return "", err + } + return strconv.Itoa(quota), nil case "GroupRatio", "group_ratio_setting.group_ratio": return ratio_setting.NormalizeGroupRatioJSONString(value) case "DataExportInterval": @@ -338,11 +348,22 @@ func parseDataExportInterval(value string) (int, error) { return interval, nil } +func parsePreConsumedQuota(value string) (int, error) { + quota, err := strconv.Atoi(strings.TrimSpace(value)) + if err != nil || quota < 0 { + return 0, fmt.Errorf("PreConsumedQuota must be a non-negative integer") + } + return quota, nil +} + func validateOptionUpdate(key string, value string) error { if err := validateRegisteredConfigOption(key, value); err != nil { return err } switch key { + case "PreConsumedQuota": + _, err := parsePreConsumedQuota(value) + return err case "Chats": return validateJSONOption[[]map[string]string](value) case "AutoGroups": @@ -403,6 +424,13 @@ func updateOptionMap(key string, value string) (err error) { if err := validateOptionUpdate(key, value); err != nil { return err } + preConsumedQuota := 0 + if key == "PreConsumedQuota" { + preConsumedQuota, err = parsePreConsumedQuota(value) + if err != nil { + return err + } + } common.OptionMapRWMutex.Lock() defer common.OptionMapRWMutex.Unlock() @@ -681,7 +709,7 @@ func updateOptionMap(key string, value string) (err error) { case "QuotaRemindThreshold": common.QuotaRemindThreshold, _ = strconv.Atoi(value) case "PreConsumedQuota": - common.PreConsumedQuota, _ = strconv.Atoi(value) + common.PreConsumedQuota = preConsumedQuota case "ModelRequestRateLimitCount": setting.ModelRequestRateLimitCount, _ = strconv.Atoi(value) case "ModelRequestRateLimitDurationMinutes": diff --git a/model/option_test.go b/model/option_test.go index 9dbdc9e1..e8a11944 100644 --- a/model/option_test.go +++ b/model/option_test.go @@ -208,6 +208,64 @@ func TestValidateOptionUpdateRejectsUnsafePricingValues(t *testing.T) { } } +func TestValidateOptionUpdateRejectsInvalidPreConsumedQuota(t *testing.T) { + for _, value := range []string{"-1", "1.5", "not-a-number"} { + t.Run(value, func(t *testing.T) { + require.Error(t, validateOptionUpdate("PreConsumedQuota", value)) + }) + } + + for _, value := range []string{"0", "500", " 1000 "} { + t.Run("accepts-"+strings.TrimSpace(value), func(t *testing.T) { + require.NoError(t, validateOptionUpdate("PreConsumedQuota", value)) + }) + } +} + +func TestNormalizePreConsumedQuota(t *testing.T) { + normalized, err := normalizeOptionUpdateValue("PreConsumedQuota", " 1000 ") + require.NoError(t, err) + require.Equal(t, "1000", normalized) + + for _, value := range []string{"-1", "1.5", "not-a-number"} { + _, err := normalizeOptionUpdateValue("PreConsumedQuota", value) + require.Error(t, err) + } +} + +func TestUpdateOptionMapRejectsInvalidPreConsumedQuotaWithoutMutation(t *testing.T) { + setupOptionMapTestState(t) + registerOptionKeysForTest("PreConsumedQuota") + originalPreConsumedQuota := common.PreConsumedQuota + common.PreConsumedQuota = 123 + common.OptionMapRWMutex.Lock() + common.OptionMap["PreConsumedQuota"] = "123" + common.OptionMapRWMutex.Unlock() + t.Cleanup(func() { common.PreConsumedQuota = originalPreConsumedQuota }) + + err := updateOptionMap("PreConsumedQuota", "not-a-number") + + require.Error(t, err) + require.Equal(t, 123, common.PreConsumedQuota) + require.Equal(t, "123", optionMapValueForTest("PreConsumedQuota")) +} + +func TestUpdateOptionPersistsPreConsumedQuota(t *testing.T) { + setupOptionMapTestState(t) + registerOptionKeysForTest("PreConsumedQuota") + deleteOptionsForTest(t, "PreConsumedQuota") + originalPreConsumedQuota := common.PreConsumedQuota + t.Cleanup(func() { + common.PreConsumedQuota = originalPreConsumedQuota + deleteOptionsForTest(t, "PreConsumedQuota") + }) + + require.NoError(t, UpdateOption("PreConsumedQuota", " 1000 ")) + require.Equal(t, "1000", optionValueForTest(t, "PreConsumedQuota")) + require.Equal(t, "1000", optionMapValueForTest("PreConsumedQuota")) + require.Equal(t, 1000, common.PreConsumedQuota) +} + func TestValidateOptionUpdateRejectsOversizedPricingMaps(t *testing.T) { const entryCount = 20001 values := make(map[string]float64, entryCount) diff --git a/model/passkey.go b/model/passkey.go index 9f190187..8cd45947 100644 --- a/model/passkey.go +++ b/model/passkey.go @@ -16,8 +16,9 @@ import ( ) var ( - ErrPasskeyNotFound = errors.New("passkey credential not found") - ErrFriendlyPasskeyNotFound = errors.New("Passkey 验证失败,请重试或联系管理员") + ErrPasskeyNotFound = errors.New("passkey credential not found") + ErrPasskeyCredentialChanged = errors.New("Passkey 凭证已变化,请重试") + ErrFriendlyPasskeyNotFound = errors.New("Passkey 验证失败,请重试或联系管理员") ) type PasskeyCredential struct { @@ -26,7 +27,7 @@ type PasskeyCredential struct { CredentialID string `json:"credential_id" gorm:"type:varchar(512);uniqueIndex;not null"` // base64 encoded PublicKey string `json:"public_key" gorm:"type:text;not null"` // base64 encoded AttestationType string `json:"attestation_type" gorm:"type:varchar(255)"` - AAGUID string `json:"aaguid" gorm:"type:varchar(512)"` // base64 encoded + AAGUID string `json:"aaguid" gorm:"column:aa_guid;type:varchar(512)"` // base64 encoded SignCount uint32 `json:"sign_count" gorm:"default:0"` CloneWarning bool `json:"clone_warning"` UserPresent bool `json:"user_present"` @@ -177,34 +178,106 @@ func GetPasskeyByCredentialID(credentialID []byte) (*PasskeyCredential, error) { return &credential, nil } -func UpsertPasskeyCredential(credential *PasskeyCredential) error { - if credential == nil { - common.SysLog("UpsertPasskeyCredential: nil credential provided") +func UpdatePasskeyCredentialAfterAuthentication(credential *PasskeyCredential) error { + if credential == nil || credential.UserID <= 0 || credential.CredentialID == "" { + common.SysLog("UpdatePasskeyCredentialAfterAuthentication: nil credential provided") + return fmt.Errorf("Passkey 保存失败,请重试") + } + updates := passkeyAuthenticationUpdates(credential) + result := DB.Model(&PasskeyCredential{}). + Where("user_id = ? AND credential_id = ?", credential.UserID, credential.CredentialID). + Updates(updates) + if result.Error != nil { + common.SysLog(fmt.Sprintf("UpdatePasskeyCredentialAfterAuthentication: failed to update credential for user %d: %v", credential.UserID, result.Error)) + return fmt.Errorf("Passkey 保存失败,请重试") + } + if result.RowsAffected == 1 { + return nil + } + var count int64 + if err := DB.Model(&PasskeyCredential{}). + Where("user_id = ? AND credential_id = ?", credential.UserID, credential.CredentialID). + Count(&count).Error; err != nil { return fmt.Errorf("Passkey 保存失败,请重试") } - return DB.Transaction(func(tx *gorm.DB) error { - // 使用Unscoped()进行硬删除,避免唯一索引冲突 + if count == 1 { + return nil + } + return ErrPasskeyCredentialChanged +} + +func passkeyAuthenticationUpdates(credential *PasskeyCredential) map[string]interface{} { + updates := map[string]interface{}{ + "public_key": credential.PublicKey, + "attestation_type": credential.AttestationType, + "aa_guid": credential.AAGUID, + "user_present": credential.UserPresent, + "user_verified": credential.UserVerified, + "backup_eligible": credential.BackupEligible, + "backup_state": credential.BackupState, + "transports": credential.Transports, + "attachment": credential.Attachment, + "sign_count": gorm.Expr("CASE WHEN sign_count < ? THEN ? ELSE sign_count END", credential.SignCount, credential.SignCount), + } + if credential.CloneWarning { + updates["clone_warning"] = true + } + if credential.LastUsedAt != nil { + updates["last_used_at"] = gorm.Expr( + "CASE WHEN last_used_at IS NULL OR last_used_at < ? THEN ? ELSE last_used_at END", + *credential.LastUsedAt, + *credential.LastUsedAt, + ) + } + return updates +} + +func ReplacePasskeyCredentialAndBumpSessionGeneration(credential *PasskeyCredential) (int64, error) { + if credential == nil || credential.UserID <= 0 { + return 0, fmt.Errorf("Passkey 保存失败,请重试") + } + var generation int64 + var cacheTask CacheInvalidationTask + err := DB.Transaction(func(tx *gorm.DB) error { if err := tx.Unscoped().Where("user_id = ?", credential.UserID).Delete(&PasskeyCredential{}).Error; err != nil { - common.SysLog(fmt.Sprintf("UpsertPasskeyCredential: failed to delete existing credential for user %d: %v", credential.UserID, err)) - return fmt.Errorf("Passkey 保存失败,请重试") + return err } if err := tx.Create(credential).Error; err != nil { - common.SysLog(fmt.Sprintf("UpsertPasskeyCredential: failed to create credential for user %d: %v", credential.UserID, err)) - return fmt.Errorf("Passkey 保存失败,请重试") + return err } - return nil + var err error + generation, cacheTask, err = bumpUserSessionGenerationTx(tx, credential.UserID) + return err }) + if err != nil { + common.SysLog(fmt.Sprintf("ReplacePasskeyCredentialAndBumpSessionGeneration: failed for user %d: %v", credential.UserID, err)) + return 0, fmt.Errorf("Passkey 保存失败,请重试") + } + dispatchStagedCacheInvalidation(cacheTask) + return generation, nil } -func DeletePasskeyByUserID(userID int) error { - if userID == 0 { - common.SysLog("DeletePasskeyByUserID: empty user ID") - return fmt.Errorf("删除失败,请重试") +func DeletePasskeyAndBumpSessionGeneration(userID int) (int64, error) { + if userID <= 0 { + return 0, fmt.Errorf("删除失败,请重试") } - // 使用Unscoped()进行硬删除,避免唯一索引冲突 - if err := DB.Unscoped().Where("user_id = ?", userID).Delete(&PasskeyCredential{}).Error; err != nil { - common.SysLog(fmt.Sprintf("DeletePasskeyByUserID: failed to delete passkey for user %d: %v", userID, err)) - return fmt.Errorf("删除失败,请重试") + var generation int64 + var cacheTask CacheInvalidationTask + err := DB.Transaction(func(tx *gorm.DB) error { + result := tx.Unscoped().Where("user_id = ?", userID).Delete(&PasskeyCredential{}) + if result.Error != nil { + return result.Error + } + if result.RowsAffected == 0 { + return ErrPasskeyNotFound + } + var err error + generation, cacheTask, err = bumpUserSessionGenerationTx(tx, userID) + return err + }) + if err != nil { + return 0, err } - return nil + dispatchStagedCacheInvalidation(cacheTask) + return generation, nil } diff --git a/model/task.go b/model/task.go index dd3c1302..949977b8 100644 --- a/model/task.go +++ b/model/task.go @@ -46,7 +46,7 @@ const ( ) type Task struct { - ID int64 `json:"id" gorm:"primary_key;AUTO_INCREMENT"` + ID int64 `json:"id" gorm:"primaryKey;index:idx_task_timeout_cursor,priority:2"` CreatedAt int64 `json:"created_at" gorm:"index"` UpdatedAt int64 `json:"updated_at"` TaskID string `json:"task_id" gorm:"type:varchar(191);index"` // 第三方id,不一定有/ song id\ Task id @@ -58,7 +58,7 @@ type Task struct { Action string `json:"action" gorm:"type:varchar(40);index"` // 任务类型, song, lyrics, description-mode Status TaskStatus `json:"status" gorm:"type:varchar(20);index"` // 任务状态 FailReason string `json:"fail_reason"` - SubmitTime int64 `json:"submit_time" gorm:"index"` + SubmitTime int64 `json:"submit_time" gorm:"index;index:idx_task_timeout_cursor,priority:1"` StartTime int64 `json:"start_time" gorm:"index"` FinishTime int64 `json:"finish_time" gorm:"index"` Progress string `json:"progress" gorm:"type:varchar(20);index"` @@ -349,16 +349,31 @@ func TaskGetAllTasks(startIdx int, num int, queryParams SyncTaskQueryParams) []* } func GetTimedOutUnfinishedTasks(cutoffUnix int64, limit int) []*Task { + tasks, _ := GetTimedOutUnfinishedTasksAfter(cutoffUnix, 0, 0, limit) + return tasks +} + +// GetTimedOutUnfinishedTasksAfter returns timed-out non-terminal tasks after +// the supplied (submit_time, id) cursor in stable ascending order. +func GetTimedOutUnfinishedTasksAfter(cutoffUnix int64, afterSubmitTime int64, afterID int64, limit int) ([]*Task, error) { var tasks []*Task - err := DB.Where("status NOT IN ?", []string{TaskStatusFailure, TaskStatusSuccess}). - Where("submit_time < ?", cutoffUnix). - Order("submit_time"). + query := DB.Where("status NOT IN ?", []string{TaskStatusFailure, TaskStatusSuccess}). + Where("submit_time < ?", cutoffUnix) + if afterID > 0 { + query = query.Where( + "(submit_time > ? OR (submit_time = ? AND id > ?))", + afterSubmitTime, + afterSubmitTime, + afterID, + ) + } + err := query.Order("submit_time ASC, id ASC"). Limit(limit). Find(&tasks).Error if err != nil { - return nil + return nil, err } - return tasks + return tasks, nil } func GetAllUnFinishSyncTasks(limit int) []*Task { diff --git a/model/token.go b/model/token.go index 4f8603eb..408c8ed6 100644 --- a/model/token.go +++ b/model/token.go @@ -536,6 +536,19 @@ func BatchDeleteTokens(ids []int, userId int) (int, error) { return len(tokens), nil } +func revokeUserTokensTx(tx *gorm.DB, userID int) ([]CacheInvalidationTask, error) { + cacheTasks, err := stageUserTokenCacheInvalidationsTx(tx, userID, false) + if err != nil { + return nil, err + } + if err := tx.Model(&Token{}). + Where("user_id = ?", userID). + Update("status", common.TokenStatusDisabled).Error; err != nil { + return nil, err + } + return cacheTasks, nil +} + func GetTokenKeysByIds(ids []int, userId int) ([]Token, error) { var tokens []Token err := DB.Select("id", commonKeyCol). diff --git a/model/twofa.go b/model/twofa.go index f108cbd2..b634cf49 100644 --- a/model/twofa.go +++ b/model/twofa.go @@ -98,9 +98,15 @@ func (t *TwoFA) Delete() error { if t.Id == 0 { return errors.New("2FA记录ID不能为空") } + if t.UserId <= 0 { + return errors.New("2FA用户ID不能为空") + } // 使用事务确保原子性 return DB.Transaction(func(tx *gorm.DB) error { + if err := lockTwoFACredentialUserTx(tx, t.UserId); err != nil { + return err + } // 同时删除相关的备用码记录(硬删除) if err := tx.Unscoped().Where("user_id = ?", t.UserId).Delete(&TwoFABackupCode{}).Error; err != nil { return err @@ -141,32 +147,50 @@ func (t *TwoFA) IsLocked() bool { // CreateBackupCodes 创建备用码 func CreateBackupCodes(userId int, codes []string) error { + hashes, err := hashBackupCodes(codes) + if err != nil { + return err + } return DB.Transaction(func(tx *gorm.DB) error { - // 先删除现有的备用码 - if err := tx.Where("user_id = ?", userId).Delete(&TwoFABackupCode{}).Error; err != nil { + if err := lockTwoFACredentialUserTx(tx, userId); err != nil { return err } + return replaceBackupCodeHashesTx(tx, userId, hashes) + }) +} - // 创建新的备用码记录 - for _, code := range codes { - hashedCode, err := common.HashBackupCode(code) - if err != nil { - return err - } - - backupCode := TwoFABackupCode{ - UserId: userId, - CodeHash: hashedCode, - IsUsed: false, - } +func hashBackupCodes(codes []string) ([]string, error) { + hashes := make([]string, 0, len(codes)) + for _, code := range codes { + hashedCode, err := common.HashBackupCode(code) + if err != nil { + return nil, err + } + hashes = append(hashes, hashedCode) + } + return hashes, nil +} - if err := tx.Create(&backupCode).Error; err != nil { - return err - } +func replaceBackupCodeHashesTx(tx *gorm.DB, userID int, hashes []string) error { + if err := tx.Unscoped().Where("user_id = ?", userID).Delete(&TwoFABackupCode{}).Error; err != nil { + return err + } + for _, hashedCode := range hashes { + backupCode := TwoFABackupCode{ + UserId: userID, + CodeHash: hashedCode, + IsUsed: false, } + if err := tx.Create(&backupCode).Error; err != nil { + return err + } + } + return nil +} - return nil - }) +func lockTwoFACredentialUserTx(tx *gorm.DB, userID int) error { + var user User + return withRowLock(tx).Select("id").Where("id = ?", userID).First(&user).Error } // ValidateBackupCode 验证并使用备用码 @@ -177,29 +201,39 @@ func ValidateBackupCode(userId int, code string) (bool, error) { normalizedCode := common.NormalizeBackupCode(code) - // 查找未使用的备用码 var backupCodes []TwoFABackupCode if err := DB.Where("user_id = ? AND is_used = ?", userId, false).Find(&backupCodes).Error; err != nil { return false, err } - - // 验证备用码 - for _, bc := range backupCodes { - if common.ValidatePasswordAndHash(normalizedCode, bc.CodeHash) { - // 标记为已使用 - now := time.Now() - bc.IsUsed = true - bc.UsedAt = &now - - if err := DB.Save(&bc).Error; err != nil { - return false, err - } - - return true, nil + matchedID := 0 + matchedHash := "" + for _, backupCode := range backupCodes { + if common.ValidatePasswordAndHash(normalizedCode, backupCode.CodeHash) { + matchedID = backupCode.Id + matchedHash = backupCode.CodeHash + break } } + if matchedID == 0 { + return false, nil + } - return false, nil + validated := false + err := DB.Transaction(func(tx *gorm.DB) error { + if err := lockTwoFACredentialUserTx(tx, userId); err != nil { + return err + } + now := time.Now() + result := tx.Model(&TwoFABackupCode{}). + Where("id = ? AND user_id = ? AND code_hash = ? AND is_used = ?", matchedID, userId, matchedHash, false). + Updates(map[string]interface{}{"is_used": true, "used_at": &now}) + if result.Error != nil { + return result.Error + } + validated = result.RowsAffected == 1 + return nil + }) + return validated, err } // GetUnusedBackupCodeCount 获取未使用的备用码数量 @@ -223,6 +257,37 @@ func DisableTwoFA(userId int) error { return twoFA.Delete() } +func DisableTwoFAAndBumpSessionGeneration(userID int) (int64, error) { + var generation int64 + var cacheTask CacheInvalidationTask + err := DB.Transaction(func(tx *gorm.DB) error { + if err := lockTwoFACredentialUserTx(tx, userID); err != nil { + return err + } + var twoFA TwoFA + if err := tx.Where("user_id = ? AND is_enabled = ?", userID, true).First(&twoFA).Error; err != nil { + if errors.Is(err, gorm.ErrRecordNotFound) { + return ErrTwoFANotEnabled + } + return err + } + if err := tx.Unscoped().Where("user_id = ?", userID).Delete(&TwoFABackupCode{}).Error; err != nil { + return err + } + if err := tx.Unscoped().Delete(&twoFA).Error; err != nil { + return err + } + var err error + generation, cacheTask, err = bumpUserSessionGenerationTx(tx, userID) + return err + }) + if err != nil { + return 0, err + } + dispatchStagedCacheInvalidation(cacheTask) + return generation, nil +} + // EnableTwoFA 启用2FA func (t *TwoFA) Enable() error { t.IsEnabled = true @@ -231,6 +296,77 @@ func (t *TwoFA) Enable() error { return t.Update() } +func (t *TwoFA) EnableAndBumpSessionGeneration() (int64, error) { + if t == nil || t.Id == 0 || t.UserId <= 0 { + return 0, errors.New("2FA记录无效") + } + var generation int64 + var cacheTask CacheInvalidationTask + err := DB.Transaction(func(tx *gorm.DB) error { + if err := lockTwoFACredentialUserTx(tx, t.UserId); err != nil { + return err + } + result := tx.Model(&TwoFA{}). + Where("id = ? AND user_id = ? AND is_enabled = ?", t.Id, t.UserId, false). + Updates(map[string]interface{}{ + "is_enabled": true, + "failed_attempts": 0, + "locked_until": nil, + }) + if result.Error != nil { + return result.Error + } + if result.RowsAffected != 1 { + return errors.New("2FA状态已变化,请重试") + } + var err error + generation, cacheTask, err = bumpUserSessionGenerationTx(tx, t.UserId) + return err + }) + if err != nil { + return 0, err + } + t.IsEnabled = true + t.FailedAttempts = 0 + t.LockedUntil = nil + dispatchStagedCacheInvalidation(cacheTask) + return generation, nil +} + +func ReplaceBackupCodesAndBumpSessionGeneration(userID int, codes []string) (int64, error) { + hashes, err := hashBackupCodes(codes) + if err != nil { + return 0, err + } + var generation int64 + var cacheTask CacheInvalidationTask + err = DB.Transaction(func(tx *gorm.DB) error { + if err := lockTwoFACredentialUserTx(tx, userID); err != nil { + return err + } + var enabledCount int64 + if err := tx.Model(&TwoFA{}). + Where("user_id = ? AND is_enabled = ?", userID, true). + Count(&enabledCount).Error; err != nil { + return err + } + if enabledCount != 1 { + return ErrTwoFANotEnabled + } + if err := replaceBackupCodeHashesTx(tx, userID, hashes); err != nil { + return err + } + var err error + generation, cacheTask, err = bumpUserSessionGenerationTx(tx, userID) + return err + }) + if err != nil { + return 0, err + } + dispatchStagedCacheInvalidation(cacheTask) + return generation, nil +} + // ValidateTOTPAndUpdateUsage 验证TOTP并更新使用记录 func (t *TwoFA) ValidateTOTPAndUpdateUsage(code string) (bool, error) { // 检查是否被锁定 diff --git a/model/user.go b/model/user.go index fe9cbe56..6b20fa78 100644 --- a/model/user.go +++ b/model/user.go @@ -27,6 +27,8 @@ const userOAuthIdentityLockName = "maxapi:user-oauth-identity" var userOAuthIdentityLockMu sync.Mutex +var ErrTelegramIdAlreadyTaken = errors.New("telegram identity is already bound") + type UserUpdateField string const ( @@ -57,38 +59,39 @@ const ( // User if you add sensitive fields, don't forget to clean them in setupLogin function. // Otherwise, the sensitive information will be saved on local storage in plain text! type User struct { - Id int `json:"id"` - Username string `json:"username" gorm:"unique;index" validate:"max=20"` - Password string `json:"password" gorm:"not null;" validate:"min=8,max=20"` - OriginalPassword string `json:"original_password" gorm:"-:all"` // this field is only for Password change verification, don't save it to database! - DisplayName string `json:"display_name" gorm:"index" validate:"max=20"` - Role int `json:"role" gorm:"type:int;default:1"` // admin, common - Status int `json:"status" gorm:"type:int;default:1"` // enabled, disabled - Email string `json:"email" gorm:"index" validate:"max=50"` - NormalizedEmail string `json:"-" gorm:"column:normalized_email;size:50;index"` - GitHubId string `json:"github_id" gorm:"column:github_id;size:512;index" validate:"max=512"` - DiscordId string `json:"discord_id" gorm:"column:discord_id;size:512;index" validate:"max=512"` - OidcId string `json:"oidc_id" gorm:"column:oidc_id;size:512;index" validate:"max=512"` - WeChatId string `json:"wechat_id" gorm:"column:wechat_id;size:512;index" validate:"max=512"` - TelegramId string `json:"telegram_id" gorm:"column:telegram_id;size:512;index" validate:"max=512"` - VerificationCode string `json:"verification_code" gorm:"-:all"` // this field is only for Email verification, don't save it to database! - AccessToken *string `json:"-" gorm:"type:char(32);column:access_token;uniqueIndex"` // this token is for system management - Quota int64 `json:"quota" gorm:"type:bigint;default:0"` - UsedQuota int64 `json:"used_quota" gorm:"type:bigint;default:0;column:used_quota"` // used quota - RequestCount int `json:"request_count" gorm:"type:int;default:0;"` // request number - Group string `json:"group" gorm:"type:varchar(64);default:'default'"` - AffCode string `json:"aff_code" gorm:"type:varchar(32);column:aff_code;uniqueIndex"` - AffCount int `json:"aff_count" gorm:"type:int;default:0;column:aff_count"` - AffQuota int64 `json:"aff_quota" gorm:"type:bigint;default:0;column:aff_quota"` // 邀请剩余额度 - AffHistoryQuota int64 `json:"aff_history_quota" gorm:"type:bigint;default:0;column:aff_history"` // 邀请历史额度 - InviterId int `json:"inviter_id" gorm:"type:int;column:inviter_id;index"` - DeletedAt gorm.DeletedAt `gorm:"index"` - LinuxDOId string `json:"linux_do_id" gorm:"column:linux_do_id;size:512;index" validate:"max=512"` - Setting string `json:"setting" gorm:"type:text;column:setting"` - Remark string `json:"remark,omitempty" gorm:"type:varchar(255)" validate:"max=255"` - StripeCustomer string `json:"stripe_customer" gorm:"type:varchar(64);column:stripe_customer;index"` - CreatedAt int64 `json:"created_at" gorm:"autoCreateTime;column:created_at"` - LastLoginAt int64 `json:"last_login_at" gorm:"default:0;column:last_login_at"` + Id int `json:"id"` + Username string `json:"username" gorm:"unique;index" validate:"max=20"` + Password string `json:"password" gorm:"not null;" validate:"min=8,max=20"` + OriginalPassword string `json:"original_password" gorm:"-:all"` // this field is only for Password change verification, don't save it to database! + DisplayName string `json:"display_name" gorm:"index" validate:"max=20"` + Role int `json:"role" gorm:"type:int;default:1"` // admin, common + Status int `json:"status" gorm:"type:int;default:1"` // enabled, disabled + Email string `json:"email" gorm:"index" validate:"max=50"` + NormalizedEmail string `json:"-" gorm:"column:normalized_email;size:50;index"` + GitHubId string `json:"github_id" gorm:"column:github_id;size:512;index" validate:"max=512"` + DiscordId string `json:"discord_id" gorm:"column:discord_id;size:512;index" validate:"max=512"` + OidcId string `json:"oidc_id" gorm:"column:oidc_id;size:512;index" validate:"max=512"` + WeChatId string `json:"wechat_id" gorm:"column:wechat_id;size:512;index" validate:"max=512"` + TelegramId string `json:"telegram_id" gorm:"column:telegram_id;size:512;index" validate:"max=512"` + VerificationCode string `json:"verification_code" gorm:"-:all"` // this field is only for Email verification, don't save it to database! + AccessToken *string `json:"-" gorm:"type:char(32);column:access_token;uniqueIndex"` // this token is for system management + Quota int64 `json:"quota" gorm:"type:bigint;default:0"` + UsedQuota int64 `json:"used_quota" gorm:"type:bigint;default:0;column:used_quota"` // used quota + RequestCount int `json:"request_count" gorm:"type:int;default:0;"` // request number + Group string `json:"group" gorm:"type:varchar(64);default:'default'"` + AffCode string `json:"aff_code" gorm:"type:varchar(32);column:aff_code;uniqueIndex"` + AffCount int `json:"aff_count" gorm:"type:int;default:0;column:aff_count"` + AffQuota int64 `json:"aff_quota" gorm:"type:bigint;default:0;column:aff_quota"` // 邀请剩余额度 + AffHistoryQuota int64 `json:"aff_history_quota" gorm:"type:bigint;default:0;column:aff_history"` // 邀请历史额度 + InviterId int `json:"inviter_id" gorm:"type:int;column:inviter_id;index"` + DeletedAt gorm.DeletedAt `gorm:"index"` + LinuxDOId string `json:"linux_do_id" gorm:"column:linux_do_id;size:512;index" validate:"max=512"` + Setting string `json:"setting" gorm:"type:text;column:setting"` + Remark string `json:"remark,omitempty" gorm:"type:varchar(255)" validate:"max=255"` + StripeCustomer string `json:"stripe_customer" gorm:"type:varchar(64);column:stripe_customer;index"` + CreatedAt int64 `json:"created_at" gorm:"autoCreateTime;column:created_at"` + LastLoginAt int64 `json:"last_login_at" gorm:"default:0;column:last_login_at"` + SessionGeneration int64 `json:"-" gorm:"default:0;not null;column:session_generation"` } func (user *User) BeforeSave(_ *gorm.DB) error { @@ -103,14 +106,15 @@ func (user *User) normalizeEmailFields() { func (user *User) ToBaseUser() *UserBase { cache := &UserBase{ - Id: user.Id, - Group: user.Group, - Quota: user.Quota, - Role: user.Role, - Status: user.Status, - Username: user.Username, - Setting: user.Setting, - Email: user.Email, + Id: user.Id, + Group: user.Group, + Quota: user.Quota, + Role: user.Role, + Status: user.Status, + Username: user.Username, + Setting: user.Setting, + Email: user.Email, + SessionGeneration: user.SessionGeneration, } return cache } @@ -676,6 +680,34 @@ func GetUserById(id int, selectAll bool) (*User, error) { return &user, err } +func CanUseOAuthReauthentication(userID int) (bool, error) { + user, err := GetUserById(userID, true) + if err != nil { + return false, err + } + if user.Password != "" { + return false, nil + } + + twoFA, err := GetTwoFAByUserId(userID) + if err != nil { + return false, err + } + if twoFA != nil && twoFA.IsEnabled { + return false, nil + } + + _, err = GetPasskeyByUserID(userID) + switch { + case err == nil: + return false, nil + case errors.Is(err, ErrPasskeyNotFound): + return true, nil + default: + return false, err + } +} + func GetUserIdByAffCode(affCode string) (int, error) { if affCode == "" { return 0, errors.New("affCode 为空!") @@ -997,6 +1029,15 @@ func (user *User) Update(updatePassword bool) error { } func (user *User) UpdateFields(updatePassword bool, fields ...UserUpdateField) error { + _, err := user.updateFields(updatePassword, fields...) + return err +} + +func (user *User) UpdateFieldsWithSessionGeneration(updatePassword bool, fields ...UserUpdateField) (int64, error) { + return user.updateFields(updatePassword, fields...) +} + +func (user *User) updateFields(updatePassword bool, fields ...UserUpdateField) (int64, error) { var cacheTask CacheInvalidationTask update := func(db *gorm.DB) error { return db.Transaction(func(tx *gorm.DB) error { @@ -1015,10 +1056,10 @@ func (user *User) UpdateFields(updatePassword bool, fields ...UserUpdateField) e err = withUserOAuthIdentityMutationLock(DB, update) } if err != nil { - return err + return 0, err } dispatchStagedCacheInvalidation(cacheTask) - return nil + return user.SessionGeneration, nil } func (user *User) UpdateWithTx(tx *gorm.DB, updatePassword bool) error { @@ -1091,6 +1132,7 @@ func buildUserUpdateValues(current User, newUser User, updatePassword bool, fiel } if updatePassword { updates["password"] = newUser.Password + updates["session_generation"] = gorm.Expr("session_generation + ?", 1) } copyUnspecifiedUserUpdateValues(updates, current) @@ -1541,6 +1583,47 @@ func IsTelegramIdAlreadyTaken(telegramId string) (bool, error) { return result.RowsAffected > 0, result.Error } +func BindTelegramIdentityWithAuthFlow(userID int, telegramID string, state string) (int64, error) { + if userID <= 0 || strings.TrimSpace(telegramID) == "" || strings.TrimSpace(state) == "" { + return 0, ErrAuthFlowInvalid + } + var generation int64 + var cacheTask CacheInvalidationTask + err := withUserOAuthIdentityMutationLock(DB, func(db *gorm.DB) error { + return db.Transaction(func(tx *gorm.DB) error { + if _, err := consumeAuthFlowWithDB(tx, state, AuthFlowMatch{ + Purpose: AuthFlowPurposeOAuth, + Provider: "telegram", + Intent: AuthFlowIntentBind, + UserId: userID, + }); err != nil { + return err + } + var count int64 + if err := tx.Unscoped().Model(&User{}). + Where("telegram_id = ? AND id <> ?", telegramID, userID). + Count(&count).Error; err != nil { + return err + } + if count > 0 { + return ErrTelegramIdAlreadyTaken + } + result := tx.Model(&User{}).Where("id = ?", userID).Update("telegram_id", telegramID) + if err := ensureUserUpdateMatchedTx(tx, result, userID, ErrUserNotFound); err != nil { + return err + } + var err error + generation, cacheTask, err = bumpUserSessionGenerationTx(tx, userID) + return err + }) + }) + if err != nil { + return 0, err + } + dispatchStagedCacheInvalidation(cacheTask) + return generation, nil +} + func ResetUserPasswordByEmail(email string, password string) error { if email == "" || password == "" { return errors.New("邮箱地址或密码为空!") @@ -1553,8 +1636,66 @@ func ResetUserPasswordByEmail(email string, password string) error { if err != nil { return err } - err = DB.Model(&User{}).Where("id = ?", user.Id).Update("password", hashedPassword).Error - return err + var userCacheTask CacheInvalidationTask + var tokenCacheTasks []CacheInvalidationTask + err = DB.Transaction(func(tx *gorm.DB) error { + result := tx.Model(&User{}).Where("id = ?", user.Id).Updates(map[string]interface{}{ + "password": hashedPassword, + "access_token": nil, + "session_generation": gorm.Expr("session_generation + ?", 1), + }) + if err := ensureUserUpdateMatchedTx(tx, result, user.Id, ErrUserNotFound); err != nil { + return err + } + var err error + tokenCacheTasks, err = revokeUserTokensTx(tx, user.Id) + if err != nil { + return err + } + userCacheTask, err = stageUserCacheInvalidationTx(tx, user.Id, false) + return err + }) + if err != nil { + return err + } + dispatchStagedCacheInvalidation(userCacheTask) + dispatchStagedCacheInvalidations(tokenCacheTasks) + return nil +} + +func bumpUserSessionGenerationTx(tx *gorm.DB, userID int) (int64, CacheInvalidationTask, error) { + if userID <= 0 { + return 0, CacheInvalidationTask{}, ErrUserNotFound + } + result := tx.Model(&User{}). + Where("id = ?", userID). + UpdateColumn("session_generation", gorm.Expr("session_generation + ?", 1)) + if err := ensureUserUpdateMatchedTx(tx, result, userID, ErrUserNotFound); err != nil { + return 0, CacheInvalidationTask{}, err + } + var generation int64 + if err := tx.Model(&User{}). + Where("id = ?", userID). + Pluck("session_generation", &generation).Error; err != nil { + return 0, CacheInvalidationTask{}, err + } + cacheTask, err := stageUserCacheInvalidationTx(tx, userID, false) + return generation, cacheTask, err +} + +func BumpUserSessionGeneration(userID int) (int64, error) { + var generation int64 + var cacheTask CacheInvalidationTask + err := DB.Transaction(func(tx *gorm.DB) error { + var err error + generation, cacheTask, err = bumpUserSessionGenerationTx(tx, userID) + return err + }) + if err != nil { + return 0, err + } + dispatchStagedCacheInvalidation(cacheTask) + return generation, nil } func IsAdmin(userId int) bool { diff --git a/model/user_cache.go b/model/user_cache.go index 9b9273d2..9ed29e0d 100644 --- a/model/user_cache.go +++ b/model/user_cache.go @@ -52,14 +52,15 @@ var userCacheRetries = struct { // UserBase struct remains the same as it represents the cached data structure type UserBase struct { - Id int `json:"id"` - Group string `json:"group"` - Email string `json:"email"` - Quota int64 `json:"quota"` - Role int `json:"role"` - Status int `json:"status"` - Username string `json:"username"` - Setting string `json:"setting"` + Id int `json:"id"` + Group string `json:"group"` + Email string `json:"email"` + Quota int64 `json:"quota"` + Role int `json:"role"` + Status int `json:"status"` + Username string `json:"username"` + Setting string `json:"setting"` + SessionGeneration int64 `json:"session_generation"` } func (user *UserBase) WriteContext(c *gin.Context) { @@ -383,14 +384,15 @@ func GetUserCache(userId int) (userCache *UserBase, err error) { // Create cache object from user data userCache = &UserBase{ - Id: user.Id, - Group: user.Group, - Quota: user.Quota, - Role: user.Role, - Status: user.Status, - Username: user.Username, - Setting: user.Setting, - Email: user.Email, + Id: user.Id, + Group: user.Group, + Quota: user.Quota, + Role: user.Role, + Status: user.Status, + Username: user.Username, + Setting: user.Setting, + Email: user.Email, + SessionGeneration: user.SessionGeneration, } if cacheVersionValid { userSnapshot := *user diff --git a/relay/helper/price.go b/relay/helper/price.go index 659fa3d7..7c20e28a 100644 --- a/relay/helper/price.go +++ b/relay/helper/price.go @@ -9,6 +9,7 @@ import ( "github.com/MAX-API-Next/MAX-API/model" "github.com/MAX-API-Next/MAX-API/pkg/billingexpr" relaycommon "github.com/MAX-API-Next/MAX-API/relay/common" + relayconstant "github.com/MAX-API-Next/MAX-API/relay/constant" "github.com/MAX-API-Next/MAX-API/setting/billing_setting" "github.com/MAX-API-Next/MAX-API/setting/operation_setting" "github.com/MAX-API-Next/MAX-API/setting/ratio_setting" @@ -38,6 +39,45 @@ func taskModelRequiresConfiguredPrice(modelName string) bool { return modelName == "swap_face" || strings.HasPrefix(modelName, "mj_") || strings.HasPrefix(modelName, "suno_") } +func configuredPreConsumedQuota() (int, error) { + if common.PreConsumedQuota < 0 { + return 0, fmt.Errorf("pre-consumed quota cannot be negative: %d", common.PreConsumedQuota) + } + return common.PreConsumedQuota, nil +} + +// ApplyPreConsumedQuotaFloor makes the administrator's pre-consume setting +// effective for every paid pricing mode. The floor is applied to the quota +// amount, after model/group pricing, and never to a free group. +func ApplyPreConsumedQuotaFloor(quota int, billable bool) (int, error) { + if quota < 0 { + return 0, fmt.Errorf("calculated pre-consume quota cannot be negative: %d", quota) + } + configured, err := configuredPreConsumedQuota() + if err != nil { + return 0, err + } + if billable && quota < configured { + return configured, nil + } + return quota, nil +} + +func addNonNegativeInts(values ...int) (int, error) { + maxInt := int(^uint(0) >> 1) + total := 0 + for _, value := range values { + if value < 0 { + return 0, fmt.Errorf("billing estimate cannot be negative: %d", value) + } + if total > maxInt-value { + return 0, fmt.Errorf("billing estimate overflows integer range") + } + total += value + } + return total, nil +} + // https://docs.claude.com/en/docs/build-with-claude/prompt-caching#1-hour-cache-duration const claudeCacheCreation1hMultiplier = 6 / 3.75 @@ -71,6 +111,12 @@ func HandleGroupRatio(ctx *gin.Context, relayInfo *relaycommon.RelayInfo) types. } func ModelPriceHelper(c *gin.Context, info *relaycommon.RelayInfo, promptTokens int, meta *types.TokenCountMeta) (types.PriceData, error) { + if _, err := configuredPreConsumedQuota(); err != nil { + return types.PriceData{}, err + } + if promptTokens < 0 { + return types.PriceData{}, fmt.Errorf("prompt token estimate cannot be negative: %d", promptTokens) + } modelPrice, usePrice := ratio_setting.GetModelPrice(info.OriginModelName, false) groupRatioInfo := HandleGroupRatio(c, info) @@ -91,10 +137,14 @@ func ModelPriceHelper(c *gin.Context, info *relaycommon.RelayInfo, promptTokens var audioRatio float64 var audioCompletionRatio float64 var freeModel bool + var err error if !usePrice { - preConsumedTokens := common.Max(promptTokens, common.PreConsumedQuota) - if meta.MaxTokens != 0 { - preConsumedTokens += meta.MaxTokens + preConsumedTokens := promptTokens + if meta != nil && meta.MaxTokens != 0 { + preConsumedTokens, err = addNonNegativeInts(preConsumedTokens, meta.MaxTokens) + if err != nil { + return types.PriceData{}, err + } } var success bool var matchName string @@ -124,7 +174,7 @@ func ModelPriceHelper(c *gin.Context, info *relaycommon.RelayInfo, promptTokens } preConsumedQuota = quota } else { - if meta.ImagePriceRatio != 0 { + if meta != nil && meta.ImagePriceRatio != 0 { modelPrice = modelPrice * meta.ImagePriceRatio } quota, err := common.QuotaFromFloatStrict(modelPrice * common.QuotaPerUnit * groupRatioInfo.GroupRatio) @@ -152,6 +202,14 @@ func ModelPriceHelper(c *gin.Context, info *relaycommon.RelayInfo, promptTokens } } } + // Alpha Search adds its deterministic tool surcharge in the controller. + // Defer the reservation floor until that complete estimate is available. + if !freeModel && groupRatioInfo.GroupRatio > 0 && info.RelayMode != relayconstant.RelayModeAlphaSearch { + preConsumedQuota, err = ApplyPreConsumedQuotaFloor(preConsumedQuota, true) + if err != nil { + return types.PriceData{}, err + } + } priceData := types.PriceData{ FreeModel: freeModel, @@ -183,6 +241,9 @@ func ModelPriceHelperPerCall(c *gin.Context, info *relaycommon.RelayInfo) (types return types.PriceData{}, fmt.Errorf("relay info is nil") } groupRatioInfo := HandleGroupRatio(c, info) + if _, err := configuredPreConsumedQuota(); err != nil { + return types.PriceData{}, err + } modelPrice, success := ratio_setting.GetModelPrice(info.OriginModelName, true) usePrice := success @@ -218,10 +279,10 @@ func ModelPriceHelperPerCall(c *gin.Context, info *relaycommon.RelayInfo) (types } var quota int + var err error freeModel := false if usePrice { - var err error quota, err = common.QuotaFromFloatStrict(modelPrice * common.QuotaPerUnit * groupRatioInfo.GroupRatio) if err != nil { return types.PriceData{}, err @@ -234,7 +295,6 @@ func ModelPriceHelperPerCall(c *gin.Context, info *relaycommon.RelayInfo) (types } } else { // 按量计费:以模型倍率的一半作为预扣额度 - var err error quota, err = common.QuotaFromFloatStrict(modelRatio / 2 * common.QuotaPerUnit * groupRatioInfo.GroupRatio) if err != nil { return types.PriceData{}, err @@ -248,13 +308,24 @@ func ModelPriceHelperPerCall(c *gin.Context, info *relaycommon.RelayInfo) (types } } + // Keep the model price in Quota for final settlement. The configured + // minimum belongs only to the reservation made before the request. + preConsumedQuota := quota + if !freeModel && groupRatioInfo.GroupRatio > 0 { + preConsumedQuota, err = ApplyPreConsumedQuotaFloor(quota, true) + if err != nil { + return types.PriceData{}, err + } + } + priceData := types.PriceData{ - FreeModel: freeModel, - ModelPrice: modelPrice, - ModelRatio: modelRatio, - UsePrice: usePrice, - Quota: quota, - GroupRatioInfo: groupRatioInfo, + FreeModel: freeModel, + ModelPrice: modelPrice, + ModelRatio: modelRatio, + UsePrice: usePrice, + Quota: quota, + QuotaToPreConsume: preConsumedQuota, + GroupRatioInfo: groupRatioInfo, } return priceData, nil } @@ -280,7 +351,7 @@ func modelPriceHelperTiered(c *gin.Context, info *relaycommon.RelayInfo, promptT } estimatedCompletionTokens := 0 - if meta.MaxTokens != 0 { + if meta != nil && meta.MaxTokens != 0 { estimatedCompletionTokens = meta.MaxTokens } @@ -304,6 +375,7 @@ func modelPriceHelperTiered(c *gin.Context, info *relaycommon.RelayInfo, promptT if err != nil { return types.PriceData{}, err } + estimatedQuotaAfterGroup := preConsumedQuota freeModel := false if !operation_setting.GetQuotaSetting().EnableFreeModelPreConsume { @@ -312,6 +384,14 @@ func modelPriceHelperTiered(c *gin.Context, info *relaycommon.RelayInfo, promptT freeModel = true } } + // Alpha Search adds its deterministic tool surcharge in the controller. + // Defer the reservation floor until that complete estimate is available. + if !freeModel && groupRatioInfo.GroupRatio > 0 && info.RelayMode != relayconstant.RelayModeAlphaSearch { + preConsumedQuota, err = ApplyPreConsumedQuotaFloor(preConsumedQuota, true) + if err != nil { + return types.PriceData{}, err + } + } exprHash := billingexpr.ExprHashString(exprStr) snapshot := &billingexpr.BillingSnapshot{ @@ -323,7 +403,7 @@ func modelPriceHelperTiered(c *gin.Context, info *relaycommon.RelayInfo, promptT EstimatedPromptTokens: promptTokens, EstimatedCompletionTokens: estimatedCompletionTokens, EstimatedQuotaBeforeGroup: quotaBeforeGroup, - EstimatedQuotaAfterGroup: preConsumedQuota, + EstimatedQuotaAfterGroup: estimatedQuotaAfterGroup, EstimatedTier: trace.MatchedTier, QuotaPerUnit: common.QuotaPerUnit, ExprVersion: billingexpr.ExprVersion(exprStr), diff --git a/relay/helper/price_test.go b/relay/helper/price_test.go index b79ff335..4c372a04 100644 --- a/relay/helper/price_test.go +++ b/relay/helper/price_test.go @@ -9,6 +9,7 @@ import ( "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" + relayconstant "github.com/MAX-API-Next/MAX-API/relay/constant" "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" @@ -20,6 +21,9 @@ import ( func TestModelPriceHelperTieredUsesPreloadedRequestInput(t *testing.T) { gin.SetMode(gin.TestMode) + originalPreConsumedQuota := common.PreConsumedQuota + common.PreConsumedQuota = 2_000 + t.Cleanup(func() { common.PreConsumedQuota = originalPreConsumedQuota }) saved := map[string]string{} require.NoError(t, config.GlobalConfig.SaveToDB(func(key, value string) error { @@ -45,6 +49,7 @@ func TestModelPriceHelperTieredUsesPreloadedRequestInput(t *testing.T) { ctx.Set("group", "default") info := &relaycommon.RelayInfo{ + RelayMode: relayconstant.RelayModeAlphaSearch, OriginModelName: "tiered-test-model", UserGroup: "default", UsingGroup: "default", @@ -59,11 +64,44 @@ func TestModelPriceHelperTieredUsesPreloadedRequestInput(t *testing.T) { require.NoError(t, err) require.Equal(t, 1500, priceData.QuotaToPreConsume) require.NotNil(t, info.TieredBillingSnapshot) + require.Equal(t, 1500, info.TieredBillingSnapshot.EstimatedQuotaAfterGroup) require.Equal(t, "stream", info.TieredBillingSnapshot.EstimatedTier) require.Equal(t, billing_setting.BillingModeTieredExpr, info.TieredBillingSnapshot.BillingMode) require.Equal(t, common.QuotaPerUnit, info.TieredBillingSnapshot.QuotaPerUnit) } +func TestModelPriceHelperDefersAlphaSearchFloorUntilSurcharge(t *testing.T) { + gin.SetMode(gin.TestMode) + originalQuotaPerUnit := common.QuotaPerUnit + originalPreConsumedQuota := common.PreConsumedQuota + originalModelPrice := ratio_setting.ModelPrice2JSONString() + originalGroupRatio := ratio_setting.GroupRatio2JSONString() + t.Cleanup(func() { + common.QuotaPerUnit = originalQuotaPerUnit + common.PreConsumedQuota = originalPreConsumedQuota + require.NoError(t, ratio_setting.UpdateModelPriceByJSONString(originalModelPrice)) + require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(originalGroupRatio)) + }) + + common.QuotaPerUnit = 1_000 + common.PreConsumedQuota = 500 + require.NoError(t, ratio_setting.UpdateModelPriceByJSONString(`{"alpha-floor-test":0.2}`)) + require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"default":1}`)) + + ctx, _ := gin.CreateTestContext(httptest.NewRecorder()) + info := &relaycommon.RelayInfo{ + RelayMode: relayconstant.RelayModeAlphaSearch, + OriginModelName: "alpha-floor-test", + UserGroup: "default", + UsingGroup: "default", + } + + priceData, err := ModelPriceHelper(ctx, info, 1, &types.TokenCountMeta{}) + + require.NoError(t, err) + require.Equal(t, 200, priceData.QuotaToPreConsume) +} + func TestModelPriceHelperPerCallRejectsUnpricedMJSunoTaskModels(t *testing.T) { gin.SetMode(gin.TestMode) originalSelfUseMode := operation_setting.SelfUseModeEnabled @@ -96,15 +134,18 @@ func TestModelPriceHelperPerCallRejectsUnpricedMJSunoTaskModels(t *testing.T) { func TestModelPriceHelperPerCallUsesDefaultTaskPrice(t *testing.T) { gin.SetMode(gin.TestMode) originalQuotaPerUnit := common.QuotaPerUnit + originalPreConsumedQuota := common.PreConsumedQuota originalGroupRatio := ratio_setting.GroupRatio2JSONString() quotaSetting := operation_setting.GetQuotaSetting() originalEnableFreeModelPreConsume := quotaSetting.EnableFreeModelPreConsume t.Cleanup(func() { common.QuotaPerUnit = originalQuotaPerUnit + common.PreConsumedQuota = originalPreConsumedQuota quotaSetting.EnableFreeModelPreConsume = originalEnableFreeModelPreConsume require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(originalGroupRatio)) }) common.QuotaPerUnit = 1000 + common.PreConsumedQuota = 0 quotaSetting.EnableFreeModelPreConsume = true require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"default":1}`)) @@ -124,3 +165,124 @@ func TestModelPriceHelperPerCallUsesDefaultTaskPrice(t *testing.T) { require.Equal(t, 0.1, priceData.ModelPrice) require.Equal(t, 100, priceData.Quota) } + +func TestModelPriceHelperPerCallSeparatesFinalAndPreConsumedQuota(t *testing.T) { + gin.SetMode(gin.TestMode) + originalQuotaPerUnit := common.QuotaPerUnit + originalPreConsumedQuota := common.PreConsumedQuota + originalGroupRatio := ratio_setting.GroupRatio2JSONString() + t.Cleanup(func() { + common.QuotaPerUnit = originalQuotaPerUnit + common.PreConsumedQuota = originalPreConsumedQuota + require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(originalGroupRatio)) + }) + + common.QuotaPerUnit = 1000 + common.PreConsumedQuota = 1000 + 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.False(t, priceData.FreeModel) + // The final per-call price remains the model price; the configured value + // is only the amount reserved before the upstream request. + require.Equal(t, 100, priceData.Quota) + require.Equal(t, 1000, priceData.QuotaToPreConsume) +} + +func TestModelPriceHelperAppliesPreConsumedQuotaToFixedPriceModels(t *testing.T) { + gin.SetMode(gin.TestMode) + originalQuotaPerUnit := common.QuotaPerUnit + originalPreConsumedQuota := common.PreConsumedQuota + originalModelPrice := ratio_setting.ModelPrice2JSONString() + originalGroupRatio := ratio_setting.GroupRatio2JSONString() + t.Cleanup(func() { + common.QuotaPerUnit = originalQuotaPerUnit + common.PreConsumedQuota = originalPreConsumedQuota + require.NoError(t, ratio_setting.UpdateModelPriceByJSONString(originalModelPrice)) + require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(originalGroupRatio)) + }) + + common.QuotaPerUnit = 1000 + common.PreConsumedQuota = 1000 + require.NoError(t, ratio_setting.UpdateModelPriceByJSONString(`{"fixed-preconsume-test":0.1}`)) + require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"default":1}`)) + + ctx, _ := gin.CreateTestContext(httptest.NewRecorder()) + info := &relaycommon.RelayInfo{ + OriginModelName: "fixed-preconsume-test", + UserGroup: "default", + UsingGroup: "default", + } + + priceData, err := ModelPriceHelper(ctx, info, 1, &types.TokenCountMeta{}) + + require.NoError(t, err) + require.True(t, priceData.UsePrice) + require.Equal(t, 1000, priceData.QuotaToPreConsume) +} + +func TestModelPriceHelperUsesConfiguredPreConsumedQuotaAsPromptEstimateFloor(t *testing.T) { + gin.SetMode(gin.TestMode) + originalQuotaPerUnit := common.QuotaPerUnit + originalPreConsumedQuota := common.PreConsumedQuota + originalModelRatio := ratio_setting.ModelRatio2JSONString() + originalGroupRatio := ratio_setting.GroupRatio2JSONString() + t.Cleanup(func() { + common.QuotaPerUnit = originalQuotaPerUnit + common.PreConsumedQuota = originalPreConsumedQuota + require.NoError(t, ratio_setting.UpdateModelRatioByJSONString(originalModelRatio)) + require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(originalGroupRatio)) + }) + + common.QuotaPerUnit = 1 + common.PreConsumedQuota = 100 + require.NoError(t, ratio_setting.UpdateModelRatioByJSONString(`{"ratio-preconsume-test":1}`)) + require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"default":1}`)) + + ctx, _ := gin.CreateTestContext(httptest.NewRecorder()) + + for _, test := range []struct { + name string + prompt int + maxTokens int + expectedPre int + }{ + {name: "below floor", prompt: 50, expectedPre: 100}, + {name: "above floor", prompt: 250, expectedPre: 250}, + {name: "includes completion estimate", prompt: 50, maxTokens: 100, expectedPre: 150}, + } { + t.Run(test.name, func(t *testing.T) { + info := &relaycommon.RelayInfo{ + OriginModelName: "ratio-preconsume-test", + UserGroup: "default", + UsingGroup: "default", + } + priceData, err := ModelPriceHelper(ctx, info, test.prompt, &types.TokenCountMeta{MaxTokens: test.maxTokens}) + + require.NoError(t, err) + require.False(t, priceData.UsePrice) + // PreConsumedQuota is the minimum admission reservation, not an + // extra token estimate that should be multiplied by the ratio. + require.Equal(t, test.expectedPre, priceData.QuotaToPreConsume) + }) + } +} + +func TestAddNonNegativeIntsRejectsInvalidBillingEstimates(t *testing.T) { + _, err := addNonNegativeInts(-1, 1) + require.Error(t, err) + + maxInt := int(^uint(0) >> 1) + _, err = addNonNegativeInts(maxInt, 1) + require.Error(t, err) +} diff --git a/relay/mjproxy_handler.go b/relay/mjproxy_handler.go index 0054d263..4929a0f0 100644 --- a/relay/mjproxy_handler.go +++ b/relay/mjproxy_handler.go @@ -258,16 +258,26 @@ func prepareMidjourneyBillingTask(c *gin.Context, info *relaycommon.RelayInfo, a info.Action = action info.OriginModelName = service.CovertMjpActionToModelName(action) info.PriceData = priceData + preConsumedQuota := priceData.Quota if !charge || priceData.FreeModel { info.PriceData.Quota = 0 + info.PriceData.QuotaToPreConsume = 0 + preConsumedQuota = 0 + } else if priceData.GroupRatioInfo.GroupRatio > 0 { + var err error + preConsumedQuota, err = helper.ApplyPreConsumedQuotaFloor(priceData.Quota, true) + if err != nil { + return nil, &dto.MidjourneyResponse{Code: constant.MjRequestError, Description: err.Error()} + } + info.PriceData.QuotaToPreConsume = preConsumedQuota } - if info.PriceData.Quota > 0 { + if preConsumedQuota > 0 { info.ForcePreConsume = true - if apiErr := service.PreConsumeBilling(c, info.PriceData.Quota, info); apiErr != nil { + if apiErr := service.PreConsumeBilling(c, preConsumedQuota, info); apiErr != nil { return nil, &dto.MidjourneyResponse{Code: constant.MjRequestError, Description: apiErr.Error()} } } - task, taskErr := ensureTaskPlaceholder(constant.TaskPlatformMidjourney, info) + task, taskErr := ensureTaskPlaceholder(constant.TaskPlatformMidjourney, info, preConsumedQuota) if taskErr != nil { if info.Billing != nil { info.Billing.Refund(c) @@ -353,6 +363,35 @@ func markMidjourneySubmissionResponseAmbiguous(info *relaycommon.RelayInfo, task } } +type midjourneyBillingSettlementPendingError struct { + taskID string + err error +} + +func (e *midjourneyBillingSettlementPendingError) Error() string { + return fmt.Sprintf("midjourney billing settlement requires reconciliation: %v", e.err) +} + +func (e *midjourneyBillingSettlementPendingError) Unwrap() error { + return e.err +} + +func midjourneyBillingSettlementPendingResponse(taskID string) *dto.MidjourneyResponse { + return &dto.MidjourneyResponse{ + Code: constant.MjRequestError, + Description: constant.MjBillingSettlementPending, + Result: taskID, + } +} + +func midjourneySubmissionErrorResponse(err error) *dto.MidjourneyResponse { + var settlementErr *midjourneyBillingSettlementPendingError + if errors.As(err, &settlementErr) { + return midjourneyBillingSettlementPendingResponse(settlementErr.taskID) + } + return service.MidjourneyErrorWrapper(constant.MjRequestError, "persist_midjourney_task_failed") +} + 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") @@ -392,7 +431,8 @@ func finalizeMidjourneySubmission(c *gin.Context, info *relaycommon.RelayInfo, t } 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()) + common.SysError("midjourney billing settlement requires reconciliation: " + err.Error()) + return false, &midjourneyBillingSettlementPendingError{taskID: midjourneyTask.MjId, err: err} } if settlementOperationKey != "" { if err := model.ProcessBillingSettlementEffect(settlementOperationKey); err != nil { @@ -470,7 +510,7 @@ func RelaySwapFace(c *gin.Context, info *relaycommon.RelayInfo) *dto.MidjourneyR 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") + return midjourneySubmissionErrorResponse(err) } } writeMidjourneyStatusCode(c, mjResp.StatusCode) @@ -820,7 +860,7 @@ func RelayMidjourneySubmit(c *gin.Context, relayInfo *relaycommon.RelayInfo) *dt 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") + return midjourneySubmissionErrorResponse(err) } //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 040e87a1..178c9b24 100644 --- a/relay/mjproxy_handler_test.go +++ b/relay/mjproxy_handler_test.go @@ -3,9 +3,11 @@ package relay import ( "bytes" "errors" + "fmt" "net/http" "net/http/httptest" "strconv" + "strings" "testing" "time" @@ -15,12 +17,39 @@ import ( "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/MAX-API-Next/MAX-API/types" "github.com/gin-gonic/gin" "github.com/glebarez/sqlite" + "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "gorm.io/gorm" ) +type failingMidjourneyBillingSettler struct { + settlementErr error + refundCalls int + operationKey string +} + +func (s *failingMidjourneyBillingSettler) Settle(int) error { return s.settlementErr } +func (s *failingMidjourneyBillingSettler) Refund(*gin.Context) { s.refundCalls++ } +func (s *failingMidjourneyBillingSettler) NeedsRefund() bool { return true } +func (s *failingMidjourneyBillingSettler) GetPreConsumedQuota() int { return 10 } +func (s *failingMidjourneyBillingSettler) Reserve(int) error { return nil } +func (s *failingMidjourneyBillingSettler) PrepareSettlement(int) (*model.BillingSettlementInput, error) { + return &model.BillingSettlementInput{ + OperationKey: s.operationKey, + Source: model.BillingSettlementSourceWallet, + UserID: 901, + TokenID: 902, + FundingDelta: 10, + TokenDelta: 10, + }, nil +} +func (s *failingMidjourneyBillingSettler) PrepareRefundSettlement() (*model.BillingSettlementInput, error) { + return nil, nil +} + func TestWriteMidjourneyStatusCodeUsesChannelMapping(t *testing.T) { gin.SetMode(gin.TestMode) recorder := httptest.NewRecorder() @@ -32,6 +61,45 @@ func TestWriteMidjourneyStatusCodeUsesChannelMapping(t *testing.T) { require.Equal(t, http.StatusServiceUnavailable, c.Writer.Status()) } +func TestPrepareMidjourneyBillingTaskDoesNotFloorZeroRatio(t *testing.T) { + gin.SetMode(gin.TestMode) + oldDB := model.DB + oldPreConsumedQuota := common.PreConsumedQuota + 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) + model.DB = db + common.PreConsumedQuota = 500 + require.NoError(t, db.AutoMigrate(&model.Task{})) + t.Cleanup(func() { + model.DB = oldDB + common.PreConsumedQuota = oldPreConsumedQuota + if sqlDB, dbErr := db.DB(); dbErr == nil { + _ = sqlDB.Close() + } + }) + + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + info := &relaycommon.RelayInfo{ + RequestId: "midjourney-zero-ratio", + UserId: 901, + TokenId: 902, + ChannelMeta: &relaycommon.ChannelMeta{ChannelId: 903}, + TaskRelayInfo: &relaycommon.TaskRelayInfo{}, + } + task, response := prepareMidjourneyBillingTask(c, info, constant.MjActionImagine, types.PriceData{ + Quota: 0, + GroupRatioInfo: types.GroupRatioInfo{GroupRatio: 0}, + }, true) + + require.Nil(t, response) + require.NotNil(t, task) + require.Zero(t, task.Quota) + require.Zero(t, info.PriceData.QuotaToPreConsume) + require.False(t, info.ForcePreConsume) + require.Nil(t, info.Billing) +} + func TestRelayMidjourneyImageRejectsUnsignedAndWrongUserURLs(t *testing.T) { gin.SetMode(gin.TestMode) oldDB := model.DB @@ -238,6 +306,66 @@ func TestClassifyMidjourneySubmissionPreservesUncertainOutcomes(t *testing.T) { } } +func TestFinalizeMidjourneySubmissionRejectsSuccessWhenSettlementFails(t *testing.T) { + gin.SetMode(gin.TestMode) + oldDB := model.DB + 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) + model.DB = db + require.NoError(t, db.AutoMigrate(&model.Task{}, &model.Midjourney{}, &model.MidjourneyBillingClaim{}, &model.BillingSettlement{})) + t.Cleanup(func() { + model.DB = oldDB + if sqlDB, dbErr := db.DB(); dbErr == nil { + _ = sqlDB.Close() + } + }) + + settler := &failingMidjourneyBillingSettler{ + settlementErr: errors.New("user quota is not enough"), + operationKey: "request:midjourney-finalize-failure:finalize", + } + task := &model.Task{ + TaskID: "task-midjourney-settlement-review", Platform: constant.TaskPlatformMidjourney, + UserId: 901, ChannelId: 903, Quota: 10, Status: model.TaskStatusNotStart, + PrivateData: model.TaskPrivateData{BillingRequestId: "midjourney-finalize-failure", TokenId: 902, AwaitingUpstreamID: true}, + } + require.NoError(t, db.Create(task).Error) + midjourneyTask := &model.Midjourney{ + UserId: 901, ChannelId: 903, MjId: "provider-midjourney-settlement-review", + Quota: 20, Progress: "0%", + } + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + c.Request = httptest.NewRequest(http.MethodPost, "/mj/submit", nil) + info := &relaycommon.RelayInfo{ + Billing: settler, + TaskRelayInfo: &relaycommon.TaskRelayInfo{Action: constant.MjActionImagine}, + ChannelMeta: &relaycommon.ChannelMeta{}, + } + + created, err := finalizeMidjourneySubmission(c, info, task, midjourneyTask) + + require.False(t, created) + var settlementErr *midjourneyBillingSettlementPendingError + require.ErrorAs(t, err, &settlementErr) + assert.Equal(t, midjourneyTask.MjId, settlementErr.taskID) + assert.Zero(t, settler.refundCalls) + assert.Equal(t, constant.MjBillingSettlementPending, midjourneySubmissionErrorResponse(err).Description) + assert.Equal(t, midjourneyTask.MjId, midjourneySubmissionErrorResponse(err).Result) + + var storedTask model.Task + require.NoError(t, db.First(&storedTask, task.ID).Error) + assert.Equal(t, model.TaskStatus(model.TaskStatusSubmitted), storedTask.Status) + assert.EqualValues(t, 10, storedTask.Quota) + assert.Equal(t, midjourneyTask.MjId, storedTask.PrivateData.UpstreamTaskID) + assert.Len(t, model.GetAllUnFinishSyncTasks(10), 1) + + var settlement model.BillingSettlement + require.NoError(t, db.Where("operation_key = ?", settler.operationKey).First(&settlement).Error) + assert.Equal(t, model.BillingSettlementStatusPending, settlement.Status) + assert.EqualValues(t, 10, settlement.FundingDelta) +} + func TestRelayMidjourneyNotifyCompletesShadowTask(t *testing.T) { gin.SetMode(gin.TestMode) oldDB := model.DB diff --git a/relay/relay_task.go b/relay/relay_task.go index a4450442..8a6502bb 100644 --- a/relay/relay_task.go +++ b/relay/relay_task.go @@ -66,10 +66,13 @@ func populateTaskBillingMetadata(task *model.Task, info *relaycommon.RelayInfo) task.Action = info.Action } -func ensureTaskPlaceholder(platform constant.TaskPlatform, info *relaycommon.RelayInfo) (*model.Task, *dto.TaskError) { +func ensureTaskPlaceholder(platform constant.TaskPlatform, info *relaycommon.RelayInfo, preConsumedQuota int) (*model.Task, *dto.TaskError) { if info == nil || info.TaskRelayInfo == nil { return nil, service.TaskErrorWrapperLocal(errors.New("task relay metadata is unavailable"), "persist_task_failed", http.StatusInternalServerError) } + if preConsumedQuota < 0 { + return nil, service.TaskErrorWrapperLocal(errors.New("pre-consumed quota cannot be negative"), "persist_task_failed", http.StatusInternalServerError) + } if info.PersistedTaskID > 0 { var task model.Task if err := model.DB.First(&task, info.PersistedTaskID).Error; err != nil { @@ -82,7 +85,7 @@ func ensureTaskPlaceholder(platform constant.TaskPlatform, info *relaycommon.Rel task.Properties = fresh.Properties task.PrivateData.Key = fresh.PrivateData.Key task.PrivateData.AwaitingUpstreamID = true - task.Quota = info.PriceData.Quota + task.Quota = preConsumedQuota populateTaskBillingMetadata(&task, info) if err := task.Update(); err != nil { return nil, TaskPersistenceError(err, "update_persisted_task_failed", "failed to update persisted task") @@ -91,7 +94,7 @@ func ensureTaskPlaceholder(platform constant.TaskPlatform, info *relaycommon.Rel } task := model.InitTask(platform, info) - task.Quota = info.PriceData.Quota + task.Quota = preConsumedQuota task.PrivateData.AwaitingUpstreamID = true populateTaskBillingMetadata(task, info) if err := task.Insert(); err != nil { @@ -106,6 +109,36 @@ func TaskPersistenceError(err error, code string, safeMessage string) *dto.TaskE return service.TaskErrorWrapperLocal(errors.New(safeMessage), code, http.StatusInternalServerError) } +func reserveTaskQuota(c *gin.Context, info *relaycommon.RelayInfo, targetQuota int) (int, *dto.TaskError) { + if info == nil { + return 0, service.TaskErrorWrapperLocal(errors.New("task relay metadata is unavailable"), "reserve_task_quota_failed", http.StatusInternalServerError) + } + if !info.PriceData.FreeModel { + if info.Billing == nil { + info.ForcePreConsume = true + if apiErr := service.PreConsumeBilling(c, targetQuota, info); apiErr != nil { + return 0, service.TaskErrorFromAPIError(apiErr) + } + } else if reserveErr := info.Billing.Reserve(targetQuota); reserveErr != nil { + return 0, service.TaskErrorWrapperLocal(reserveErr, "reserve_task_quota_failed", http.StatusForbidden) + } + } + if info.Billing != nil { + // A free retry still owns any reservation from an earlier paid attempt + // until the upstream accepts it. Keep that reservation on the task + // placeholder, then settle the successful free result to quota zero; + // failed attempts may continue to another paid channel. + return info.Billing.GetPreConsumedQuota(), nil + } + if info.PriceData.FreeModel { + // No Billing session means this attempt did not reserve any quota. Keep + // the placeholder at zero even if a future estimator accidentally + // supplies a non-zero target for a free route. + return 0, nil + } + return targetQuota, nil +} + type taskBillingEstimator interface { EstimateTaskBilling(c *gin.Context, info *relaycommon.RelayInfo) (*types.TaskBillingResult, error) } @@ -303,14 +336,25 @@ func RelayTaskSubmit(c *gin.Context, info *relaycommon.RelayInfo) (*TaskSubmitRe } } - // 7. 预扣费(仅首次 — 重试时 info.Billing 已存在,跳过) - if info.Billing == nil && !info.PriceData.FreeModel { - info.ForcePreConsume = true - if apiErr := service.PreConsumeBilling(c, info.PriceData.Quota, info); apiErr != nil { - return nil, service.TaskErrorFromAPIError(apiErr) + // Keep the raw estimate for final settlement. The configured value is an + // admission reservation only and must not become the task's final charge. + taskEstimateQuota := info.PriceData.Quota + preConsumedQuota := taskEstimateQuota + if !info.PriceData.FreeModel && info.PriceData.GroupRatioInfo.GroupRatio > 0 { + preConsumedQuota, err = helper.ApplyPreConsumedQuotaFloor(taskEstimateQuota, true) + if err != nil { + return nil, service.TaskErrorWrapper(err, "pre_consume_quota_error", http.StatusBadRequest) } } - task, taskErr := ensureTaskPlaceholder(platform, info) + info.PriceData.QuotaToPreConsume = preConsumedQuota + + // 7. 预扣费。重试可能因路由/分组变化得到更高的目标额度,必须通过 + // BillingSession.Reserve 幂等补足,不能继续沿用较小的首次预留。 + actualReservedQuota, reservationErr := reserveTaskQuota(c, info, preConsumedQuota) + if reservationErr != nil { + return nil, reservationErr + } + task, taskErr := ensureTaskPlaceholder(platform, info, actualReservedQuota) if taskErr != nil { return nil, taskErr } @@ -355,11 +399,11 @@ func RelayTaskSubmit(c *gin.Context, info *relaycommon.RelayInfo) (*TaskSubmitRe } // 11. 提交后计费调整:让适配器根据上游实际返回调整 OtherRatios - finalQuota := info.PriceData.Quota + finalQuota := taskEstimateQuota if !taskBillingOverride { if adjustedRatios := adaptor.AdjustBillingOnSubmit(info, taskData); len(adjustedRatios) > 0 { // 基于调整后的 ratios 重新计算 quota - finalQuota = recalcQuotaFromRatios(info, adjustedRatios) + finalQuota = recalcQuotaFromRatios(info, taskEstimateQuota, adjustedRatios) info.PriceData.OtherRatios = adjustedRatios info.PriceData.Quota = finalQuota } @@ -369,7 +413,11 @@ func RelayTaskSubmit(c *gin.Context, info *relaycommon.RelayInfo) (*TaskSubmitRe } task.PrivateData.UpstreamTaskID = upstreamTaskID task.PrivateData.AwaitingUpstreamID = false - task.Quota = finalQuota + // Keep the persisted task at the actual reservation until the durable + // request-finalize operation applies the funding delta and task quota in one + // transaction. Writing finalQuota here would make an unpaid target look + // settled and would break settlement recovery's task-quota CAS. + task.Quota = actualReservedQuota task.Data = taskData populateTaskBillingMetadata(task, info) @@ -482,10 +530,12 @@ func buildTaskSubmitRequestBody(c *gin.Context, info *relaycommon.RelayInfo, ada } // recalcQuotaFromRatios 根据 adjustedRatios 重新计算 quota。 -// 公式: baseQuota × ∏(ratio) — 其中 baseQuota 是不含 OtherRatios 的基础额度。 -func recalcQuotaFromRatios(info *relaycommon.RelayInfo, ratios map[string]float64) int { - // 从 PriceData 获取不含 OtherRatios 的基础价格 - baseQuota := float64(info.PriceData.Quota) +// estimatedQuota 是应用原始 OtherRatios 后、预扣下限前的估算额度。 +// 先还原不含 OtherRatios 的基础额度,再应用上游返回的新倍率。 +func recalcQuotaFromRatios(info *relaycommon.RelayInfo, estimatedQuota int, ratios map[string]float64) int { + // estimatedQuota includes the ratios used for the original reservation. + // Remove those ratios before applying the upstream-adjusted values. + baseQuota := float64(estimatedQuota) // 先除掉原有的 OtherRatios 恢复基础额度 for _, ra := range info.PriceData.OtherRatios { if ra != 1.0 && ra > 0 { diff --git a/relay/relay_task_test.go b/relay/relay_task_test.go index 958e672c..c3ccdad5 100644 --- a/relay/relay_task_test.go +++ b/relay/relay_task_test.go @@ -22,6 +22,79 @@ import ( "github.com/stretchr/testify/require" ) +type recordingTaskReservationBilling struct { + preConsumed int + reserveTo []int + settleTo []int +} + +func (b *recordingTaskReservationBilling) Settle(actualQuota int) error { + b.settleTo = append(b.settleTo, actualQuota) + b.preConsumed = actualQuota + return nil +} +func (b *recordingTaskReservationBilling) Refund(*gin.Context) {} +func (b *recordingTaskReservationBilling) NeedsRefund() bool { return true } +func (b *recordingTaskReservationBilling) GetPreConsumedQuota() int { return b.preConsumed } +func (b *recordingTaskReservationBilling) Reserve(targetQuota int) error { + b.reserveTo = append(b.reserveTo, targetQuota) + if targetQuota > b.preConsumed { + b.preConsumed = targetQuota + } + return nil +} + +func TestReserveTaskQuotaExtendsExistingReservationToActualTarget(t *testing.T) { + billing := &recordingTaskReservationBilling{preConsumed: 100} + info := &relaycommon.RelayInfo{ + Billing: billing, + PriceData: types.PriceData{ + FreeModel: false, + }, + } + + reserved, taskErr := reserveTaskQuota(nil, info, 250) + + require.Nil(t, taskErr) + assert.Equal(t, 250, reserved) + assert.Equal(t, []int{250}, billing.reserveTo) +} + +func TestReserveTaskQuotaPreservesPaidReservationUntilFreeRetrySettlement(t *testing.T) { + billing := &recordingTaskReservationBilling{preConsumed: 100} + info := &relaycommon.RelayInfo{ + Billing: billing, + PriceData: types.PriceData{ + FreeModel: true, + }, + } + + reserved, taskErr := reserveTaskQuota(nil, info, 0) + + require.Nil(t, taskErr) + assert.Equal(t, 100, reserved) + assert.Same(t, billing, info.Billing) + assert.Empty(t, billing.reserveTo) + + require.NoError(t, info.Billing.Settle(0)) + assert.Equal(t, []int{0}, billing.settleTo) + assert.Zero(t, info.Billing.GetPreConsumedQuota()) +} + +func TestReserveTaskQuotaReturnsZeroForUnreservedFreeTask(t *testing.T) { + info := &relaycommon.RelayInfo{ + PriceData: types.PriceData{ + FreeModel: true, + }, + } + + reserved, taskErr := reserveTaskQuota(nil, info, 250) + + require.Nil(t, taskErr) + assert.Zero(t, reserved) + assert.Nil(t, info.Billing) +} + func withRelayTaskRateCards(t *testing.T, cards map[string]task_billing_setting.RateCard) { t.Helper() original := task_billing_setting.GetRateCardsCopy() @@ -192,6 +265,31 @@ func TestEstimateTaskBillingFallsBackToGenericRateCard(t *testing.T) { assert.Equal(t, 3000, got.Quota) } +func TestRecalcQuotaFromRatiosUsesRawEstimateBeforePreConsumeFloor(t *testing.T) { + const estimatedQuota = 200 + info := &relaycommon.RelayInfo{ + PriceData: types.PriceData{ + Quota: estimatedQuota, + QuotaToPreConsume: 500, + OtherRatios: map[string]float64{"duration": 2}, + }, + } + + for _, tc := range []struct { + name string + ratio float64 + expected int + }{ + {name: "adjusted below estimate and floor", ratio: 1, expected: 100}, + {name: "adjusted above reservation", ratio: 6, expected: 600}, + } { + t.Run(tc.name, func(t *testing.T) { + got := recalcQuotaFromRatios(info, estimatedQuota, map[string]float64{"duration": tc.ratio}) + assert.Equal(t, tc.expected, got) + }) + } +} + func TestPrepareTaskSubmitRequestBodyParamOverrideReturnErrorIsLocal(t *testing.T) { gin.SetMode(gin.TestMode) c, _ := gin.CreateTestContext(httptest.NewRecorder()) diff --git a/router/api-router.go b/router/api-router.go index a2452325..99c6efd1 100644 --- a/router/api-router.go +++ b/router/api-router.go @@ -3,6 +3,7 @@ package router import ( "github.com/MAX-API-Next/MAX-API/controller" "github.com/MAX-API-Next/MAX-API/middleware" + "github.com/MAX-API-Next/MAX-API/model" // Import oauth package to register providers via init() _ "github.com/MAX-API-Next/MAX-API/oauth" @@ -53,7 +54,8 @@ func SetApiRouter(router *gin.Engine) { apiRouter.GET("/oauth/wechat", middleware.CriticalRateLimit(), controller.WeChatAuth) apiRouter.POST("/oauth/wechat/bind", middleware.CriticalRateLimit(), anonymousRequestBodyLimit, controller.WeChatBind) apiRouter.GET("/oauth/telegram/login", middleware.CriticalRateLimit(), controller.TelegramLogin) - apiRouter.GET("/oauth/telegram/bind", middleware.CriticalRateLimit(), controller.TelegramBind) + apiRouter.POST("/oauth/telegram/bind/state", middleware.UserAuth(), middleware.CriticalRateLimit(), middleware.UserCriticalRateLimit(), middleware.DisableCache(), middleware.SecureVerificationRequired(model.SecureVerificationScopeCredentials, model.SecureVerificationScopeOAuthReauthentication), anonymousRequestBodyLimit, controller.GenerateTelegramBindState) + apiRouter.POST("/oauth/telegram/bind", middleware.UserAuth(), middleware.CriticalRateLimit(), middleware.UserCriticalRateLimit(), middleware.DisableCache(), middleware.SecureVerificationRequired(model.SecureVerificationScopeCredentials, model.SecureVerificationScopeOAuthReauthentication), anonymousRequestBodyLimit, controller.TelegramBind) // Standard OAuth providers (GitHub, Discord, OIDC, LinuxDO) - unified route apiRouter.GET("/oauth/:provider", middleware.CriticalRateLimit(), middleware.DisableCache(), middleware.TryUserAuth(), controller.HandleOAuth) apiRouter.GET("/ratio_config", middleware.CriticalRateLimit(), controller.GetRatioConfig) @@ -89,14 +91,15 @@ func SetApiRouter(router *gin.Engine) { selfRoute.GET("/self", controller.GetSelf) 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.UserCriticalRateLimit(), middleware.DisableCache(), middleware.SecureVerificationRequired("access_token"), controller.GenerateAccessToken) + selfRoute.DELETE("/self", middleware.CriticalRateLimit(), middleware.DisableCache(), middleware.SecureVerificationRequired(model.SecureVerificationScopeAccountDelete), controller.DeleteSelf) + selfRoute.POST("/token", middleware.CriticalRateLimit(), middleware.UserCriticalRateLimit(), middleware.DisableCache(), middleware.SecureVerificationRequired(model.SecureVerificationScopeAccessToken), controller.GenerateAccessToken) + selfRoute.POST("/sessions/revoke", middleware.CriticalRateLimit(), middleware.UserCriticalRateLimit(), middleware.DisableCache(), middleware.SecureVerificationRequired(model.SecureVerificationScopeCredentials), controller.RevokeOtherSessions) selfRoute.GET("/passkey", controller.PasskeyStatus) - selfRoute.POST("/passkey/register/begin", controller.PasskeyRegisterBegin) - selfRoute.POST("/passkey/register/finish", controller.PasskeyRegisterFinish) + selfRoute.POST("/passkey/register/begin", middleware.CriticalRateLimit(), middleware.UserCriticalRateLimit(), middleware.DisableCache(), middleware.SecureVerificationRequired(model.SecureVerificationScopePasskeyRegister, model.SecureVerificationScopeCredentials, model.SecureVerificationScopeOAuthReauthentication), controller.PasskeyRegisterBegin) + selfRoute.POST("/passkey/register/finish", middleware.CriticalRateLimit(), middleware.UserCriticalRateLimit(), middleware.DisableCache(), middleware.SecureVerificationRequired(model.SecureVerificationScopePasskeyRegister, model.SecureVerificationScopeCredentials, model.SecureVerificationScopeOAuthReauthentication), controller.PasskeyRegisterFinish) selfRoute.POST("/passkey/verify/begin", controller.PasskeyVerifyBegin) selfRoute.POST("/passkey/verify/finish", controller.PasskeyVerifyFinish) - selfRoute.DELETE("/passkey", controller.PasskeyDelete) + selfRoute.DELETE("/passkey", middleware.CriticalRateLimit(), middleware.UserCriticalRateLimit(), middleware.DisableCache(), middleware.SecureVerificationRequired(model.SecureVerificationScopeCredentials), controller.PasskeyDelete) selfRoute.GET("/aff", controller.GetAffCode) selfRoute.GET("/topup/info", controller.GetTopUpInfo) selfRoute.GET("/topup/self", controller.GetUserTopUps) @@ -115,10 +118,10 @@ func SetApiRouter(router *gin.Engine) { // 2FA routes selfRoute.GET("/2fa/status", controller.Get2FAStatus) - selfRoute.POST("/2fa/setup", controller.Setup2FA) - selfRoute.POST("/2fa/enable", controller.Enable2FA) - selfRoute.POST("/2fa/disable", controller.Disable2FA) - selfRoute.POST("/2fa/backup_codes", controller.RegenerateBackupCodes) + selfRoute.POST("/2fa/setup", middleware.CriticalRateLimit(), middleware.UserCriticalRateLimit(), middleware.DisableCache(), middleware.SecureVerificationRequired(model.SecureVerificationScopeCredentials, model.SecureVerificationScopeOAuthReauthentication), controller.Setup2FA) + selfRoute.POST("/2fa/enable", middleware.CriticalRateLimit(), middleware.UserCriticalRateLimit(), middleware.DisableCache(), middleware.SecureVerificationRequired(model.SecureVerificationScopeCredentials, model.SecureVerificationScopeOAuthReauthentication), controller.Enable2FA) + selfRoute.POST("/2fa/disable", middleware.CriticalRateLimit(), middleware.UserCriticalRateLimit(), middleware.DisableCache(), controller.Disable2FA) + selfRoute.POST("/2fa/backup_codes", middleware.CriticalRateLimit(), middleware.UserCriticalRateLimit(), middleware.DisableCache(), controller.RegenerateBackupCodes) // Check-in routes selfRoute.GET("/checkin", controller.GetCheckinStatus) @@ -291,12 +294,12 @@ func SetApiRouter(router *gin.Engine) { tokenRoute.GET("/", controller.GetAllTokens) tokenRoute.GET("/search", middleware.SearchRateLimit(), controller.SearchTokens) tokenRoute.GET("/:id", controller.GetToken) - tokenRoute.POST("/:id/key", middleware.CriticalRateLimit(), middleware.DisableCache(), controller.GetTokenKey) - tokenRoute.POST("/", controller.AddToken) - tokenRoute.PUT("/", controller.UpdateToken) - tokenRoute.DELETE("/:id", controller.DeleteToken) - tokenRoute.POST("/batch", controller.DeleteTokenBatch) - tokenRoute.POST("/batch/keys", middleware.CriticalRateLimit(), middleware.DisableCache(), controller.GetTokenKeysBatch) + tokenRoute.POST("/:id/key", middleware.CriticalRateLimit(), middleware.UserCriticalRateLimit(), middleware.DisableCache(), middleware.SecureVerificationRequired(model.SecureVerificationScopeAPIToken), controller.GetTokenKey) + tokenRoute.POST("/", middleware.CriticalRateLimit(), middleware.UserCriticalRateLimit(), middleware.DisableCache(), middleware.SecureVerificationRequired(model.SecureVerificationScopeAPIToken), controller.AddToken) + tokenRoute.PUT("/", middleware.CriticalRateLimit(), middleware.UserCriticalRateLimit(), controller.UpdateToken) + tokenRoute.DELETE("/:id", middleware.CriticalRateLimit(), middleware.UserCriticalRateLimit(), controller.DeleteToken) + tokenRoute.POST("/batch", middleware.CriticalRateLimit(), middleware.UserCriticalRateLimit(), controller.DeleteTokenBatch) + tokenRoute.POST("/batch/keys", middleware.CriticalRateLimit(), middleware.UserCriticalRateLimit(), middleware.DisableCache(), middleware.SecureVerificationRequired(model.SecureVerificationScopeAPIToken), controller.GetTokenKeysBatch) } usageRoute := apiRouter.Group("/usage") diff --git a/router/api_router_test.go b/router/api_router_test.go index be2d0a64..ef38a19c 100644 --- a/router/api_router_test.go +++ b/router/api_router_test.go @@ -5,6 +5,7 @@ import ( "net/http" "net/http/httptest" "os" + "slices" "strings" "sync/atomic" "testing" @@ -19,7 +20,7 @@ import ( "gorm.io/gorm" ) -var universalVerifyRateLimitTestRun uint64 +var universalVerifyRateLimitTestRun int32 func TestUserTokenRouteUsesPostOnly(t *testing.T) { gin.SetMode(gin.TestMode) @@ -141,7 +142,7 @@ func TestHealthRoutesAreRegistered(t *testing.T) { func TestUniversalVerifyRateLimitFollowsUserAcrossIPs(t *testing.T) { gin.SetMode(gin.TestMode) - testRun := atomic.AddUint64(&universalVerifyRateLimitTestRun, 1) + testRun := atomic.AddInt32(&universalVerifyRateLimitTestRun, 1) oldDB := model.DB oldLogDB := model.LOG_DB @@ -232,11 +233,25 @@ func TestUniversalVerifyRateLimitFollowsUserAcrossIPs(t *testing.T) { } 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) { + tests := []struct { + name string + method string + path string + body string + }{ + {name: "management token", method: http.MethodPost, path: "/api/user/token", body: `{"quota":1}`}, + {name: "affiliate transfer", method: http.MethodPost, path: "/api/user/aff_transfer", body: `{"quota":1}`}, + {name: "disable 2fa", method: http.MethodPost, path: "/api/user/2fa/disable", body: `{"code":"000000"}`}, + {name: "regenerate 2fa backup codes", method: http.MethodPost, path: "/api/user/2fa/backup_codes", body: `{"code":"000000"}`}, + {name: "delete passkey", method: http.MethodDelete, path: "/api/user/passkey"}, + {name: "update api token", method: http.MethodPut, path: "/api/token/", body: `{"id":1,"name":"updated"}`}, + {name: "delete api token", method: http.MethodDelete, path: "/api/token/1"}, + {name: "delete api token batch", method: http.MethodPost, path: "/api/token/batch", body: `{"ids":[1]}`}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { gin.SetMode(gin.TestMode) - testRun := atomic.AddUint64(&universalVerifyRateLimitTestRun, 1) + testRun := atomic.AddInt32(&universalVerifyRateLimitTestRun, 1) oldDB := model.DB oldLogDB := model.LOG_DB @@ -259,7 +274,7 @@ func TestCriticalAccountRoutesRateLimitSameUserAcrossIPs(t *testing.T) { require.NoError(t, err) model.DB = db model.LOG_DB = db - require.NoError(t, db.AutoMigrate(&model.User{})) + require.NoError(t, db.AutoMigrate(&model.User{}, &model.Token{})) user := model.User{ Id: 1200000 + int(testRun), Username: fmt.Sprintf("critical-user-%d", testRun), Password: "hashed-password", Role: common.RoleCommonUser, @@ -301,7 +316,7 @@ func TestCriticalAccountRoutesRateLimitSameUserAcrossIPs(t *testing.T) { perform := func(remoteAddr string) *httptest.ResponseRecorder { recorder := httptest.NewRecorder() - request := httptest.NewRequest(http.MethodPost, path, strings.NewReader(`{"quota":1}`)) + request := httptest.NewRequest(tt.method, tt.path, strings.NewReader(tt.body)) request.RemoteAddr = remoteAddr request.Header.Set("Content-Type", "application/json") for _, sessionCookie := range loginRecorder.Result().Cookies() { @@ -338,8 +353,15 @@ func TestSecureVerificationOpenAPIIncludesScopeAndRequestBody(t *testing.T) { } schema, schemaOK := value["schema"].(map[string]any) values, valuesOK := schema["enum"].([]any) - return schemaOK && valuesOK && len(values) == 2 && - values[0] == "access_token" && values[1] == "account_delete" + if !schemaOK || !valuesOK { + return false + } + for _, requiredScope := range []any{"access_token", "account_delete", "credentials", "api_token", "passkey_register"} { + if !slices.Contains(values, requiredScope) { + return false + } + } + return true } return false }, "expected supported scope query parameter for GET /api/verify/methods") @@ -355,3 +377,188 @@ func TestSecureVerificationOpenAPIIncludesScopeAndRequestBody(t *testing.T) { require.Contains(t, properties, property) } } + +func TestSensitiveCredentialOpenAPIContracts(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)) + + telegramBind := document.Paths["/api/oauth/telegram/bind"] + require.Contains(t, telegramBind, "post") + require.NotContains(t, telegramBind, "get") + requestBody := telegramBind["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) + for _, property := range []string{"id", "auth_date", "hash", "state"} { + require.Contains(t, properties, property) + } + + telegramLogin := document.Paths["/api/oauth/telegram/login"]["get"] + parameters := telegramLogin["parameters"].([]any) + parametersByName := make(map[string]map[string]any, len(parameters)) + for _, rawParameter := range parameters { + parameter := rawParameter.(map[string]any) + name, ok := parameter["name"].(string) + require.True(t, ok) + parametersByName[name] = parameter + } + for _, name := range []string{"id", "auth_date", "hash", "state"} { + parameter, ok := parametersByName[name] + require.True(t, ok, "expected required Telegram login parameter %q", name) + require.Equal(t, "query", parameter["in"]) + require.Equal(t, true, parameter["required"]) + } + for _, name := range []string{"first_name", "last_name", "username", "photo_url"} { + parameter, ok := parametersByName[name] + require.True(t, ok, "expected optional Telegram login parameter %q", name) + require.Equal(t, "query", parameter["in"]) + require.Equal(t, false, parameter["required"]) + schema := parameter["schema"].(map[string]any) + require.Equal(t, "string", schema["type"]) + } + telegramLoginResponses := telegramLogin["responses"].(map[string]any) + require.Contains(t, telegramLoginResponses, "200") + require.Contains(t, telegramLoginResponses, "403") + + for _, path := range []string{ + "/api/user/passkey/register/begin", + "/api/user/passkey/register/finish", + } { + operation := document.Paths[path]["post"] + responses := operation["responses"].(map[string]any) + require.Contains(t, responses, "403", path) + } + for _, path := range []string{ + "/api/user/2fa/setup", + "/api/user/2fa/enable", + } { + operation := document.Paths[path]["post"] + responses := operation["responses"].(map[string]any) + require.Contains(t, responses, "403", path) + } + + for path, method := range map[string]string{ + "/api/oauth/telegram/bind/state": "post", + "/api/user/sessions/revoke": "post", + "/api/token/{id}/key": "post", + "/api/token/batch/keys": "post", + } { + require.Contains(t, document.Paths[path], method, path) + } +} + +func TestSensitiveCredentialRoutesRequireStepUp(t *testing.T) { + gin.SetMode(gin.TestMode) + testRun := atomic.AddInt32(&universalVerifyRateLimitTestRun, 1) + + oldDB := model.DB + oldLogDB := model.LOG_DB + oldRedisEnabled := common.RedisEnabled + oldMemoryCacheEnabled := common.MemoryCacheEnabled + oldGlobalAPIRateLimitEnable := common.GlobalApiRateLimitEnable + oldCriticalRateLimitEnable := common.CriticalRateLimitEnable + + common.RedisEnabled = false + common.MemoryCacheEnabled = false + common.GlobalApiRateLimitEnable = false + common.CriticalRateLimitEnable = false + + dsn := fmt.Sprintf("file:sensitive_routes_%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{}, + &model.Token{}, + &model.TwoFA{}, + &model.TwoFABackupCode{}, + &model.PasskeyCredential{}, + &model.Log{}, + )) + user := model.User{ + Id: 1300000 + int(testRun), + Username: fmt.Sprintf("sensitive-route-user-%d", testRun), + Password: "hashed-password", + Role: common.RoleCommonUser, + Status: common.UserStatusEnabled, + Group: "default", + } + require.NoError(t, db.Create(&user).Error) + token := model.Token{ + Id: 1400000 + int(testRun), + UserId: user.Id, + Key: fmt.Sprintf("sensitive-route-token-%d", testRun), + Name: "sensitive-route-token", + Status: common.TokenStatusEnabled, + ExpiredTime: -1, + } + require.NoError(t, db.Create(&token).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 + }) + + engine := gin.New() + engine.Use(sessions.Sessions("session", cookie.NewStore([]byte("sensitive-route-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) + + tests := []struct { + name string + method string + path string + body string + }{ + {name: "passkey registration begin", method: http.MethodPost, path: "/api/user/passkey/register/begin", body: `{}`}, + {name: "passkey registration finish", method: http.MethodPost, path: "/api/user/passkey/register/finish", body: `{}`}, + {name: "passkey deletion", method: http.MethodDelete, path: "/api/user/passkey", body: `{}`}, + {name: "2fa setup", method: http.MethodPost, path: "/api/user/2fa/setup", body: `{}`}, + {name: "2fa enable", method: http.MethodPost, path: "/api/user/2fa/enable", body: `{"code":"000000"}`}, + {name: "session revocation", method: http.MethodPost, path: "/api/user/sessions/revoke", body: `{}`}, + {name: "telegram bind state", method: http.MethodPost, path: "/api/oauth/telegram/bind/state", body: `{}`}, + {name: "telegram bind", method: http.MethodPost, path: "/api/oauth/telegram/bind", body: `{}`}, + {name: "api token creation", method: http.MethodPost, path: "/api/token/", body: `{"name":"blocked","expired_time":-1,"unlimited_quota":true}`}, + {name: "api token reveal", method: http.MethodPost, path: fmt.Sprintf("/api/token/%d/key", token.Id), body: `{}`}, + {name: "api token batch export", method: http.MethodPost, path: "/api/token/batch/keys", body: fmt.Sprintf(`{"ids":[%d]}`, token.Id)}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + request := httptest.NewRequest(tt.method, tt.path, strings.NewReader(tt.body)) + request.Header.Set("Content-Type", "application/json") + for _, sessionCookie := range loginRecorder.Result().Cookies() { + request.AddCookie(sessionCookie) + } + recorder := httptest.NewRecorder() + engine.ServeHTTP(recorder, request) + require.Equal(t, http.StatusForbidden, recorder.Code, recorder.Body.String()) + require.Contains(t, recorder.Body.String(), "VERIFICATION_REQUIRED") + }) + } +} diff --git a/service/alpha_search_billing_test.go b/service/alpha_search_billing_test.go index d199a813..05cfeca3 100644 --- a/service/alpha_search_billing_test.go +++ b/service/alpha_search_billing_test.go @@ -9,6 +9,8 @@ import ( "github.com/MAX-API-Next/MAX-API/pkg/billingexpr" relaycommon "github.com/MAX-API-Next/MAX-API/relay/common" relayconstant "github.com/MAX-API-Next/MAX-API/relay/constant" + "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/types" "github.com/gin-gonic/gin" @@ -63,8 +65,27 @@ func TestAlphaSearchPreConsumeQuotaRejectsOverflow(t *testing.T) { require.Error(t, err) } -func TestPrepareTieredAlphaSearchBillingKeepsSurchargeAfterGroupChange(t *testing.T) { +func TestPrepareTieredAlphaSearchBillingAppliesFloorAfterSurcharge(t *testing.T) { const expr = `tier("base", p)` + originalQuotaPerUnit := common.QuotaPerUnit + originalPreConsumedQuota := common.PreConsumedQuota + toolPrices := config.GlobalConfig.Get("tool_price_setting").(*operation_setting.ToolPriceSetting) + originalToolPrice, hadOriginalToolPrice := toolPrices.Prices[dto.BuildInToolWebSearchPreview] + common.QuotaPerUnit = 500_000 + common.PreConsumedQuota = 20_000 + toolPrices.Prices[dto.BuildInToolWebSearchPreview] = 10 + operation_setting.RebuildToolPriceIndex() + t.Cleanup(func() { + common.QuotaPerUnit = originalQuotaPerUnit + common.PreConsumedQuota = originalPreConsumedQuota + if hadOriginalToolPrice { + toolPrices.Prices[dto.BuildInToolWebSearchPreview] = originalToolPrice + } else { + delete(toolPrices.Prices, dto.BuildInToolWebSearchPreview) + } + operation_setting.RebuildToolPriceIndex() + }) + billing := &recordingBillingSettler{preConsumed: 6_000} info := &relaycommon.RelayInfo{ RelayMode: relayconstant.RelayModeAlphaSearch, @@ -85,12 +106,13 @@ func TestPrepareTieredAlphaSearchBillingKeepsSurchargeAfterGroupChange(t *testin }, } - expected, err := AlphaSearchPreConsumeQuota(2_000, info, 2) + totalBeforeFloor, err := AlphaSearchPreConsumeQuota(2_000, info, 2) require.NoError(t, err) + require.Equal(t, 12_000, totalBeforeFloor) require.Nil(t, PrepareTieredBillingForSelectedGroup(nil, info)) - require.Equal(t, []int{expected}, billing.reserves) - require.Equal(t, expected, info.FinalPreConsumedQuota) - require.Equal(t, expected, info.PriceData.QuotaToPreConsume) + require.Equal(t, []int{20_000}, billing.reserves) + require.Equal(t, 20_000, info.FinalPreConsumedQuota) + require.Equal(t, 20_000, info.PriceData.QuotaToPreConsume) require.Equal(t, 2_000, info.TieredBillingSnapshot.EstimatedQuotaAfterGroup) } diff --git a/service/billing.go b/service/billing.go index cd3ac711..c6a31642 100644 --- a/service/billing.go +++ b/service/billing.go @@ -1,10 +1,14 @@ package service import ( + "context" + "errors" "fmt" "net/http" + "github.com/MAX-API-Next/MAX-API/common" "github.com/MAX-API-Next/MAX-API/logger" + "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/types" "github.com/gin-gonic/gin" @@ -15,6 +19,62 @@ const ( BillingSourceSubscription = "subscription" ) +func billingLogContext(ctx *gin.Context) context.Context { + if ctx == nil { + return context.Background() + } + return ctx +} + +func billingEffectRequestIDs(ctx *gin.Context, relayInfo *relaycommon.RelayInfo) (string, string) { + requestID, upstreamRequestID := "", "" + if ctx != nil { + requestID = ctx.GetString(common.RequestIdKey) + upstreamRequestID = ctx.GetString(common.UpstreamRequestIdKey) + } + if requestID == "" && relayInfo != nil { + requestID = relayInfo.RequestId + } + return requestID, upstreamRequestID +} + +func newConsumeBillingSettlementEffect( + relayInfo *relaycommon.RelayInfo, + params model.RecordConsumeLogParams, + requestID string, + upstreamRequestID string, + updateUsage bool, +) *model.BillingSettlementEffect { + effect := &model.BillingSettlementEffect{ + LogType: model.LogTypeConsume, + Content: params.Content, + ChannelID: params.ChannelId, + ModelName: params.ModelName, + TokenID: params.TokenId, + TokenName: params.TokenName, + Group: params.Group, + Other: params.Other, + NodeName: common.NodeName, + UpdateUsage: updateUsage, + Quota: int64(params.Quota), + QuotaIsActual: true, + PromptTokens: params.PromptTokens, + CompletionTokens: params.CompletionTokens, + UseTimeSeconds: params.UseTimeSeconds, + IsStream: params.IsStream, + RequestID: requestID, + UpstreamRequestID: upstreamRequestID, + } + if relayInfo != nil && relayInfo.BillingSource == BillingSourceSubscription { + effect.Subscription = &model.BillingSettlementSubscriptionEffect{ + PreConsumed: relayInfo.SubscriptionPreConsumed, + AmountTotal: relayInfo.SubscriptionAmountTotal, + AmountUsedAfterPreConsume: relayInfo.SubscriptionAmountUsedAfterPreConsume, + } + } + return effect +} + func validatePreConsumedQuota(preConsumedQuota int, relayInfo *relaycommon.RelayInfo) *types.MaxAPIError { if relayInfo != nil && relayInfo.QuotaClamp != nil { return types.NewErrorWithStatusCode( @@ -56,30 +116,69 @@ func PreConsumeBilling(c *gin.Context, preConsumedQuota int, relayInfo *relaycom // SettleBilling 执行计费结算。如果 RelayInfo 上有 BillingSession 则通过 session 结算, // 否则回退到旧的 PostConsumeQuota 路径(兼容按次计费等场景)。 func SettleBilling(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, actualQuota int) error { + _, err := SettleBillingWithEffect(ctx, relayInfo, actualQuota, nil) + return err +} + +// SettleBillingWithEffect settles funding before projecting usage and logs. +// The returned flag is true when the durable settlement owns the projection; +// callers must not write a second usage/log record in that case, including when +// effect processing is still pending after the funding mutation succeeded. +func SettleBillingWithEffect(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, actualQuota int, effect *model.BillingSettlementEffect) (bool, error) { + if relayInfo == nil { + return false, fmt.Errorf("relayInfo is nil") + } + logCtx := billingLogContext(ctx) if relayInfo.Billing != nil { preConsumed := relayInfo.Billing.GetPreConsumedQuota() delta := actualQuota - preConsumed if delta > 0 { - logger.LogInfo(ctx, fmt.Sprintf("预扣费后补扣费:%s(实际消耗:%s,预扣费:%s)", + logger.LogInfo(logCtx, fmt.Sprintf("预扣费后补扣费:%s(实际消耗:%s,预扣费:%s)", logger.FormatQuota(delta), logger.FormatQuota(actualQuota), logger.FormatQuota(preConsumed), )) } else if delta < 0 { - logger.LogInfo(ctx, fmt.Sprintf("预扣费后返还扣费:%s(实际消耗:%s,预扣费:%s)", + logger.LogInfo(logCtx, fmt.Sprintf("预扣费后返还扣费:%s(实际消耗:%s,预扣费:%s)", logger.FormatQuota(-delta), logger.FormatQuota(actualQuota), logger.FormatQuota(preConsumed), )) } else { - logger.LogInfo(ctx, fmt.Sprintf("预扣费与实际消耗一致,无需调整:%s(按次计费)", + logger.LogInfo(logCtx, fmt.Sprintf("预扣费与实际消耗一致,无需调整:%s(按次计费)", logger.FormatQuota(actualQuota), )) } + if effect != nil { + if settler, ok := relayInfo.Billing.(interface { + SettleWithEffect(int, *model.BillingSettlementEffect) error + }); ok { + if err := settler.SettleWithEffect(actualQuota, effect); err != nil { + if errors.Is(err, errBillingSettlementEffectNotOwned) { + return false, nil + } + if errors.Is(err, ErrBillingSettlementEffectNotDurable) || + errors.Is(err, model.ErrBillingSettlementRecordNotDurable) || + errors.Is(err, ErrBillingFundingOutcomeUnknown) { + return false, err + } + return true, err + } + if actualQuota != 0 { + if relayInfo.BillingSource == BillingSourceSubscription { + checkAndSendSubscriptionQuotaNotify(relayInfo) + } else { + checkAndSendQuotaNotify(relayInfo, actualQuota-preConsumed, preConsumed) + } + } + return true, nil + } + } + if err := relayInfo.Billing.Settle(actualQuota); err != nil { - return err + return false, err } // 发送额度通知(订阅计费使用订阅剩余额度) @@ -90,13 +189,50 @@ func SettleBilling(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, actualQuo checkAndSendQuotaNotify(relayInfo, actualQuota-preConsumed, preConsumed) } } - return nil + return false, nil } // 回退:无 BillingSession 时使用旧路径 quotaDelta := actualQuota - relayInfo.FinalPreConsumedQuota + if effect != nil && relayInfo.RequestId != "" { + return postConsumeQuotaOnceWithEffect(relayInfo, "finalize", quotaDelta, relayInfo.FinalPreConsumedQuota, true, effect) + } if quotaDelta != 0 { - return PostConsumeQuotaOnce(relayInfo, "finalize", quotaDelta, relayInfo.FinalPreConsumedQuota, true) + return false, PostConsumeQuotaOnce(relayInfo, "finalize", quotaDelta, relayInfo.FinalPreConsumedQuota, true) } - return nil + return false, nil +} + +// settleAndRecordConsume keeps the successful usage projection behind the +// funding settlement. Durable BillingSession effects replay both the log and +// counters after a transient failure; legacy/custom sessions project locally +// only after SettleBilling succeeds. +func settleAndRecordConsume(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, shouldUpdateUsage bool, params model.RecordConsumeLogParams) { + logCtx := billingLogContext(ctx) + if relayInfo == nil { + logger.LogError(logCtx, "error settling billing: relayInfo is nil") + return + } + if params.Other == nil { + params.Other = make(map[string]interface{}) + } + appendBillingInfo(relayInfo, params.Other) + requestID, upstreamRequestID := billingEffectRequestIDs(ctx, relayInfo) + effect := newConsumeBillingSettlementEffect(relayInfo, params, requestID, upstreamRequestID, shouldUpdateUsage) + + effectHandled, err := SettleBillingWithEffect(ctx, relayInfo, params.Quota, effect) + if err != nil { + logger.LogError(logCtx, "error settling billing: "+err.Error()) + return + } + if effectHandled { + return + } + // Refresh post-settlement subscription fields for the local projection path. + appendBillingInfo(relayInfo, params.Other) + if shouldUpdateUsage { + model.UpdateUserUsedQuotaAndRequestCount(relayInfo.UserId, params.Quota) + model.UpdateChannelUsedQuota(relayInfo.ChannelId, params.Quota) + } + model.RecordConsumeLog(ctx, relayInfo.UserId, params) } diff --git a/service/billing_session.go b/service/billing_session.go index 27d10e49..0798c588 100644 --- a/service/billing_session.go +++ b/service/billing_session.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" relaycommon "github.com/MAX-API-Next/MAX-API/relay/common" @@ -39,9 +40,136 @@ type BillingSession struct { refunded bool // Refund 已调用 refundInFlight bool settleInFlight bool + fundingOutcomeUnknown bool mu sync.Mutex } +// ErrBillingFundingOutcomeUnknown means a non-durable funding mutation did +// not return a fully confirmed result. Such a mutation must not be retried +// blindly because the provider may already have committed it. +var ErrBillingFundingOutcomeUnknown = errors.New("billing funding settlement outcome is unknown") + +// ErrBillingSettlementEffectNotDurable means no idempotent settlement record +// can own the usage/log projection. The settlement remains failed; callers must +// not mistake it for a durable pending effect or project successful usage. +var ErrBillingSettlementEffectNotDurable = errors.New("billing settlement effect requires a durable funding source and request id") + +var errBillingSettlementEffectNotOwned = errors.New("billing settlement record does not own the requested effect") + +const billingSettlementBacklogNotificationInterval = 15 * time.Minute + +var billingSettlementBacklogAlertState struct { + sync.Mutex + active bool + lastCount int64 + oldestCreatedAt int64 + lastNotifiedAt time.Time +} + +var billingSettlementBacklogAlertNotificationSender = enqueueSmartOpsAlertNotification +var billingFundingOutcomeUnknownAlertNotificationSender = enqueueSmartOpsAlertNotification + +func init() { + model.SetBillingSettlementBacklogObserver(observeBillingSettlementBacklog) +} + +func observeBillingSettlementBacklog(stats model.BillingSettlementBacklogStats, observedAt time.Time) { + if observedAt.IsZero() { + observedAt = time.Now() + } + + billingSettlementBacklogAlertState.Lock() + wasActive := billingSettlementBacklogAlertState.active + previousCount := billingSettlementBacklogAlertState.lastCount + previousOldestCreatedAt := billingSettlementBacklogAlertState.oldestCreatedAt + if stats.Count <= 0 { + billingSettlementBacklogAlertState.active = false + billingSettlementBacklogAlertState.lastCount = 0 + billingSettlementBacklogAlertState.oldestCreatedAt = 0 + billingSettlementBacklogAlertState.lastNotifiedAt = time.Time{} + billingSettlementBacklogAlertState.Unlock() + if wasActive { + billingSettlementBacklogAlertNotificationSender(SmartOpsAlert{ + Key: "billing_settlement_backlog", + Status: smartOpsAlertStatusResolved, + Severity: smartOpsAlertSeverityWarning, + Component: "billing", + Node: smartOpsAlertNodeName(), + CurrentValue: 0, + Threshold: 0, + ObservedAt: observedAt, + Message: fmt.Sprintf("未完成的正向最终计费对账已清空,此前共 %d 条。", previousCount), + }) + } + return + } + + oldestAge := time.Duration(0) + if stats.OldestCreatedAt > 0 { + oldestAge = observedAt.Sub(time.Unix(stats.OldestCreatedAt, 0)) + if oldestAge < 0 { + oldestAge = 0 + } + } + shouldNotify := !wasActive || + stats.Count != previousCount || + stats.OldestCreatedAt != previousOldestCreatedAt || + billingSettlementBacklogAlertState.lastNotifiedAt.IsZero() || + observedAt.Sub(billingSettlementBacklogAlertState.lastNotifiedAt) >= billingSettlementBacklogNotificationInterval + billingSettlementBacklogAlertState.active = true + billingSettlementBacklogAlertState.lastCount = stats.Count + billingSettlementBacklogAlertState.oldestCreatedAt = stats.OldestCreatedAt + if shouldNotify { + billingSettlementBacklogAlertState.lastNotifiedAt = observedAt + } + billingSettlementBacklogAlertState.Unlock() + if !shouldNotify { + return + } + + billingSettlementBacklogAlertNotificationSender(SmartOpsAlert{ + Key: "billing_settlement_backlog", + Status: smartOpsAlertStatusFiring, + Severity: smartOpsAlertSeverityWarning, + Component: "billing", + Node: smartOpsAlertNodeName(), + CurrentValue: float64(stats.Count), + Threshold: oldestAge.Seconds(), + ObservedAt: observedAt, + Message: fmt.Sprintf( + "存在 %d 条未完成的正向最终计费对账,最早已等待 %s。相关用户的新付费请求会被阻止,请尽快处理 pending/manual 记录。", + stats.Count, + oldestAge.Round(time.Second), + ), + }) +} + +func (s *BillingSession) notifyFundingOutcomeUnknown(requested, applied int64) { + userID := 0 + tokenID := 0 + if s.relayInfo != nil { + userID = s.relayInfo.UserId + tokenID = s.relayInfo.TokenId + } + billingFundingOutcomeUnknownAlertNotificationSender(SmartOpsAlert{ + Key: "billing_funding_outcome_unknown", + Status: smartOpsAlertStatusFiring, + Severity: smartOpsAlertSeverityWarning, + Component: "billing", + Node: smartOpsAlertNodeName(), + CurrentValue: float64(requested), + Threshold: float64(applied), + ObservedAt: time.Now(), + Message: fmt.Sprintf( + "计费资金变更结果不明确(userId=%d, tokenId=%d, requested=%d, applied=%d)。系统已停止自动重试,请立即人工核对。", + userID, + tokenID, + requested, + applied, + ), + }) +} + type SettlementPreparer interface { PrepareSettlement(actualQuota int) (*model.BillingSettlementInput, error) } @@ -58,6 +186,11 @@ type billingSettleIntent struct { input model.BillingSettlementInput } +type fundingOutcomeUnknownNotification struct { + requested int64 + applied int64 +} + // Settle 根据实际消耗额度进行结算。 // 资金来源和令牌额度分两步提交。若令牌调整失败,会优先反向补偿已提交的 // 资金差额。已确认的部分补偿会记录剩余差额并在重试令牌前补齐;补偿 @@ -78,6 +211,11 @@ func (s *BillingSession) settleWithEffect(actualQuota int, effect *model.Billing return nil } else { lastErr = err + if errors.Is(err, ErrBillingFundingOutcomeUnknown) || + errors.Is(err, ErrBillingSettlementEffectNotDurable) || + errors.Is(err, errBillingSettlementEffectNotOwned) { + return err + } } if attempt < maxAttempts-1 { time.Sleep(time.Duration(attempt+1) * 20 * time.Millisecond) @@ -95,11 +233,20 @@ func (s *BillingSession) settleAttempt(actualQuota int, effect *model.BillingSet } func (s *BillingSession) beginSettleAttempt(actualQuota int, effect *model.BillingSettlementEffect) (*billingSettleIntent, bool, error) { + var notification *fundingOutcomeUnknownNotification s.mu.Lock() - defer s.mu.Unlock() + defer func() { + s.mu.Unlock() + if notification != nil { + s.notifyFundingOutcomeUnknown(notification.requested, notification.applied) + } + }() if s.settled || s.refunded { return nil, false, nil } + if s.fundingOutcomeUnknown { + return nil, false, ErrBillingFundingOutcomeUnknown + } if s.refundInFlight { return nil, false, errors.New("billing refund is already in progress") } @@ -107,18 +254,32 @@ func (s *BillingSession) beginSettleAttempt(actualQuota int, effect *model.Billi return nil, false, errors.New("billing settlement is already in progress") } delta := actualQuota - s.preConsumedQuota - if delta == 0 && effect == nil { - s.settled = true - return nil, false, nil - } if input, ok := s.prepareDurableSettlementLocked(actualQuota, effect); ok { + if delta == 0 && effect == nil { + // A persisted zero-delta record still owns settlement completion and + // must be confirmed; no record means there is nothing left to apply. + _, found, lookupErr := model.GetBillingSettlementStatus(input.OperationKey) + if lookupErr != nil { + return nil, false, lookupErr + } + if !found { + s.settled = true + return nil, false, nil + } + } s.settleInFlight = true return &billingSettleIntent{input: *input}, true, nil } + if delta == 0 && effect == nil { + s.settled = true + return nil, false, nil + } if effect != nil { - return nil, false, errors.New("billing settlement effects require a durable funding source and request id") + return nil, false, ErrBillingSettlementEffectNotDurable } - return nil, false, s.settleNonDurableLocked(delta) + settleErr, outcomeNotification := s.settleNonDurableLocked(delta) + notification = outcomeNotification + return nil, false, settleErr } func (s *BillingSession) applyDurableSettleIntent(intent *billingSettleIntent) error { @@ -130,7 +291,14 @@ func (s *BillingSession) applyDurableSettleIntent(intent *billingSettleIntent) e var effectErr error if intent.input.Effect != nil { - effectErr = model.ProcessBillingSettlementEffect(intent.input.OperationKey) + ownsEffect, ownershipErr := model.BillingSettlementOwnsEffect(intent.input.OperationKey) + if ownershipErr != nil { + effectErr = ownershipErr + } else if !ownsEffect { + effectErr = errBillingSettlementEffectNotOwned + } else { + effectErr = model.ProcessBillingSettlementEffect(intent.input.OperationKey) + } } s.finishDurableSettleIntent(true, applied) return effectErr @@ -151,12 +319,22 @@ func (s *BillingSession) finishDurableSettleIntent(applied bool, appliedDelta in s.settled = true } -func (s *BillingSession) settleNonDurableLocked(delta int) error { +func (s *BillingSession) settleNonDurableLocked(delta int) (error, *fundingOutcomeUnknownNotification) { // 1) 调整资金来源(仅在尚未提交时执行,防止重复调用) if !s.fundingSettled { applied, err := s.funding.Settle(delta) - if err != nil { - return err + if err != nil || applied != int64(delta) { + // FundingSource has no durable idempotency contract. Treat both an + // error and a short/over-applied result as outcome unknown: the + // mutation may have committed and must never be repeated blindly. + s.fundingSettled = true + s.appliedFundingDelta = applied + s.fundingOutcomeUnknown = true + notification := &fundingOutcomeUnknownNotification{requested: int64(delta), applied: applied} + if err != nil { + return fmt.Errorf("%w: %v", ErrBillingFundingOutcomeUnknown, err), notification + } + return fmt.Errorf("%w: requested=%d applied=%d", ErrBillingFundingOutcomeUnknown, delta, applied), notification } s.appliedFundingDelta = applied s.fundingSettled = true @@ -168,20 +346,22 @@ func (s *BillingSession) settleNonDurableLocked(delta int) error { needed := int64(delta) - s.appliedFundingDelta if needed != 0 { applied, err := s.funding.Settle(int(needed)) - if err != nil { - // The operation outcome is ambiguous. Do not issue another funding - // mutation automatically; later attempts may only retry the token. - s.compensationFailed = true - s.fundingReconcilePending = false - return err - } s.appliedFundingDelta += applied - if applied != needed { - return fmt.Errorf("funding reconciliation incomplete: needed=%d applied=%d", needed, applied) + if err != nil || applied != needed { + // The funding source has no durable idempotency contract. A + // failed or partial reconciliation may already have committed; + // do not repeat it or proceed with the token leg. + s.fundingOutcomeUnknown = true + s.fundingReconcilePending = false + notification := &fundingOutcomeUnknownNotification{requested: needed, applied: applied} + if err != nil { + return fmt.Errorf("%w: %v", ErrBillingFundingOutcomeUnknown, err), notification + } + return fmt.Errorf("%w: reconciliation requested=%d applied=%d", ErrBillingFundingOutcomeUnknown, needed, applied), notification } } if s.appliedFundingDelta != int64(delta) { - return fmt.Errorf("funding reconciliation incomplete: target=%d applied=%d", delta, s.appliedFundingDelta) + return fmt.Errorf("funding reconciliation incomplete: target=%d applied=%d", delta, s.appliedFundingDelta), nil } s.fundingReconcilePending = false } @@ -194,10 +374,11 @@ func (s *BillingSession) settleNonDurableLocked(delta int) error { tokenErr = model.IncreaseTokenQuota(s.relayInfo.TokenId, s.relayInfo.TokenKey, -delta) } if tokenErr != nil { + var notification *fundingOutcomeUnknownNotification common.SysLog(fmt.Sprintf("error adjusting token quota after funding settled (userId=%d, tokenId=%d, delta=%d): %s", s.relayInfo.UserId, s.relayInfo.TokenId, delta, tokenErr.Error())) if s.compensationFailed { - return tokenErr + return tokenErr, nil } if s.appliedFundingDelta == 0 { s.fundingSettled = false @@ -206,15 +387,20 @@ func (s *BillingSession) settleNonDurableLocked(delta int) error { compensated, compensationErr := s.funding.Settle(-int(committed)) if compensationErr != nil { // A funding error may be an ambiguous commit. Retain the target - // delta and never repeat this non-idempotent compensation. + // delta and never repeat this non-idempotent compensation or + // continue with the token leg. s.compensationFailed = true + s.fundingOutcomeUnknown = true + notification = &fundingOutcomeUnknownNotification{requested: -committed, applied: compensated} s.fundingReconcilePending = false common.SysLog(fmt.Sprintf("error compensating funding after token settlement failure (userId=%d, tokenId=%d, delta=%d, applied=%d, compensated=%d): %v", s.relayInfo.UserId, s.relayInfo.TokenId, delta, committed, compensated, compensationErr)) } else if compensated != -committed { s.appliedFundingDelta = committed + compensated s.compensationFailed = false - s.fundingReconcilePending = true + s.fundingOutcomeUnknown = true + notification = &fundingOutcomeUnknownNotification{requested: -committed, applied: compensated} + s.fundingReconcilePending = false common.SysLog(fmt.Sprintf("incomplete funding compensation after token settlement failure (userId=%d, tokenId=%d, delta=%d, applied=%d, compensated=%d)", s.relayInfo.UserId, s.relayInfo.TokenId, delta, committed, compensated)) } else { @@ -224,7 +410,7 @@ func (s *BillingSession) settleNonDurableLocked(delta int) error { s.fundingReconcilePending = false } } - return tokenErr + return tokenErr, notification } } // 3) 更新 relayInfo 上的订阅 PostDelta(用于日志) @@ -234,7 +420,7 @@ func (s *BillingSession) settleNonDurableLocked(delta int) error { s.compensationFailed = false s.fundingReconcilePending = false s.settled = true - return nil + return nil, nil } // PrepareSettlement returns the exact durable request-finalize intent without @@ -273,8 +459,8 @@ func (s *BillingSession) prepareDurableSettlementLocked(actualQuota int, effect if s.relayInfo.IsPlayground { tokenDelta = 0 } - return &model.BillingSettlementInput{ - OperationKey: "request:" + s.relayInfo.RequestId + ":finalize", + input := &model.BillingSettlementInput{ + OperationKey: model.BillingRequestFinalizeOperationKey(s.relayInfo.RequestId), Source: source, UserID: userID, SubscriptionID: subscriptionID, @@ -284,7 +470,13 @@ func (s *BillingSession) prepareDurableSettlementLocked(actualQuota int, effect TokenDelta: tokenDelta, SubscriptionPreConsumeRequestID: subscriptionPreConsumeRequestID(s.funding), Effect: effect, - }, true + } + if s.relayInfo.TaskRelayInfo != nil && s.relayInfo.TaskRelayInfo.PersistedTaskID > 0 { + input.TaskID = s.relayInfo.TaskRelayInfo.PersistedTaskID + input.TaskQuota = int64(s.preConsumedQuota) + input.TaskQuotaTarget = int64(actualQuota) + } + return input, true } func durableFundingIdentity(funding FundingSource) (source string, userID int, subscriptionID int, ok bool) { @@ -343,7 +535,7 @@ func (s *BillingSession) refundSettlementInputLocked() (*model.BillingSettlement 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", + OperationKey: model.BillingRequestFinalizeOperationKey(requestID), Source: source, UserID: userID, SubscriptionID: subscriptionID, @@ -453,6 +645,9 @@ func (s *BillingSession) Reserve(targetQuota int) error { s.mu.Lock() defer s.mu.Unlock() + if s.fundingOutcomeUnknown { + return ErrBillingFundingOutcomeUnknown + } if s.settled || s.refunded || s.trusted || targetQuota <= s.preConsumedQuota { return nil } @@ -738,6 +933,17 @@ func NewBillingSession(c *gin.Context, relayInfo *relaycommon.RelayInfo, preCons if relayInfo == nil { return nil, types.NewError(fmt.Errorf("relayInfo is nil"), types.ErrorCodeInvalidRequest, types.ErrOptionWithSkipRetry()) } + if blocked, err := model.HasUnresolvedPositiveFinalizeSettlement(relayInfo.UserId); err != nil { + return nil, types.NewError(err, types.ErrorCodeUpdateDataError, types.ErrOptionWithSkipRetry()) + } else if blocked { + return nil, types.NewErrorWithStatusCode( + errors.New(i18n.Translate(relayInfo.UserSetting.Language, i18n.MsgBillingReconciliationPending)), + types.ErrorCodeBillingReconciliationPending, + http.StatusForbidden, + types.ErrOptionWithSkipRetry(), + types.ErrOptionWithNoRecordErrorLog(), + ) + } pref := common.NormalizeBillingPreference(relayInfo.UserSetting.BillingPreference) diff --git a/service/billing_session_test.go b/service/billing_session_test.go index 0c2da86d..d552b659 100644 --- a/service/billing_session_test.go +++ b/service/billing_session_test.go @@ -10,6 +10,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/i18n" "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/types" @@ -32,12 +33,496 @@ func (f *recordingFundingSource) Settle(delta int) (int64, error) { return int64(delta), nil } +type uncertainFundingSource struct { + deltas []int + applied int64 + err error +} + +func TestBillingSettlementBacklogAlertReportsCountAgeAndRecovery(t *testing.T) { + originalSender := billingSettlementBacklogAlertNotificationSender + billingSettlementBacklogAlertState.Lock() + originalState := struct { + active bool + lastCount int64 + oldestCreatedAt int64 + lastNotifiedAt time.Time + }{ + active: billingSettlementBacklogAlertState.active, + lastCount: billingSettlementBacklogAlertState.lastCount, + oldestCreatedAt: billingSettlementBacklogAlertState.oldestCreatedAt, + lastNotifiedAt: billingSettlementBacklogAlertState.lastNotifiedAt, + } + billingSettlementBacklogAlertState.active = false + billingSettlementBacklogAlertState.lastCount = 0 + billingSettlementBacklogAlertState.oldestCreatedAt = 0 + billingSettlementBacklogAlertState.lastNotifiedAt = time.Time{} + billingSettlementBacklogAlertState.Unlock() + t.Cleanup(func() { + billingSettlementBacklogAlertNotificationSender = originalSender + billingSettlementBacklogAlertState.Lock() + billingSettlementBacklogAlertState.active = originalState.active + billingSettlementBacklogAlertState.lastCount = originalState.lastCount + billingSettlementBacklogAlertState.oldestCreatedAt = originalState.oldestCreatedAt + billingSettlementBacklogAlertState.lastNotifiedAt = originalState.lastNotifiedAt + billingSettlementBacklogAlertState.Unlock() + }) + + var alerts []SmartOpsAlert + billingSettlementBacklogAlertNotificationSender = func(alert SmartOpsAlert) { + alerts = append(alerts, alert) + } + now := time.Unix(10_000, 0) + stats := model.BillingSettlementBacklogStats{Count: 2, OldestCreatedAt: now.Add(-2 * time.Hour).Unix()} + + observeBillingSettlementBacklog(stats, now) + require.Len(t, alerts, 1) + assert.Equal(t, smartOpsAlertStatusFiring, alerts[0].Status) + assert.Equal(t, "billing", alerts[0].Component) + assert.Equal(t, float64(2), alerts[0].CurrentValue) + assert.Equal(t, (2 * time.Hour).Seconds(), alerts[0].Threshold) + assert.Contains(t, alerts[0].Message, "2 条") + assert.Contains(t, alerts[0].Message, "2h0m0s") + + observeBillingSettlementBacklog(stats, now.Add(time.Minute)) + require.Len(t, alerts, 1, "an unchanged backlog must be deduplicated inside the reminder interval") + + stats.Count = 3 + observeBillingSettlementBacklog(stats, now.Add(2*time.Minute)) + require.Len(t, alerts, 2) + assert.Equal(t, float64(3), alerts[1].CurrentValue) + + observeBillingSettlementBacklog(stats, now.Add(2*time.Minute+billingSettlementBacklogNotificationInterval)) + require.Len(t, alerts, 3, "a persistent backlog must produce an age reminder") + + observeBillingSettlementBacklog(model.BillingSettlementBacklogStats{}, now.Add(20*time.Minute)) + require.Len(t, alerts, 4) + assert.Equal(t, smartOpsAlertStatusResolved, alerts[3].Status) + assert.Contains(t, alerts[3].Message, "此前共 3 条") +} + +func (f *uncertainFundingSource) Source() string { return BillingSourceWallet } +func (f *uncertainFundingSource) PreConsume(int) error { return nil } +func (f *uncertainFundingSource) Refund() error { return nil } +func (f *uncertainFundingSource) Settle(delta int) (int64, error) { + f.deltas = append(f.deltas, delta) + return f.applied, f.err +} + +func TestBillingSessionFailsClosedAfterUncertainNonDurableFundingSettlement(t *testing.T) { + originalSender := billingFundingOutcomeUnknownAlertNotificationSender + var alerts []SmartOpsAlert + var activeSession *BillingSession + var senderObservedUnlocked bool + billingFundingOutcomeUnknownAlertNotificationSender = func(alert SmartOpsAlert) { + if activeSession != nil && activeSession.mu.TryLock() { + senderObservedUnlocked = true + activeSession.mu.Unlock() + } + alerts = append(alerts, alert) + } + t.Cleanup(func() { + billingFundingOutcomeUnknownAlertNotificationSender = originalSender + }) + + cases := []struct { + name string + applied int64 + err error + }{ + {name: "error after possible commit", applied: 5, err: errors.New("provider timeout")}, + {name: "partial result", applied: 3}, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + alerts = nil + senderObservedUnlocked = false + funding := &uncertainFundingSource{applied: tc.applied, err: tc.err} + session := &BillingSession{ + relayInfo: &relaycommon.RelayInfo{UserId: 101, TokenId: 102, TokenKey: "uncertain-funding-token"}, + funding: funding, preConsumedQuota: 10, + } + activeSession = session + + err := session.Settle(15) + require.ErrorIs(t, err, ErrBillingFundingOutcomeUnknown) + assert.Equal(t, []int{5}, funding.deltas) + assert.True(t, session.fundingOutcomeUnknown) + require.Len(t, alerts, 1) + assert.True(t, senderObservedUnlocked, "funding alerts must be enqueued after releasing the billing session mutex") + assert.Equal(t, "billing_funding_outcome_unknown", alerts[0].Key) + assert.Equal(t, smartOpsAlertStatusFiring, alerts[0].Status) + assert.Contains(t, alerts[0].Message, "userId=101") + assert.Contains(t, alerts[0].Message, "tokenId=102") + assert.Contains(t, alerts[0].Message, "requested=5") + assert.Contains(t, alerts[0].Message, fmt.Sprintf("applied=%d", tc.applied)) + + // A second attempt must not repeat the non-idempotent funding call. + require.ErrorIs(t, session.Settle(15), ErrBillingFundingOutcomeUnknown) + require.ErrorIs(t, session.Reserve(20), ErrBillingFundingOutcomeUnknown) + assert.Equal(t, []int{5}, funding.deltas) + require.Len(t, alerts, 1, "fail-closed retries must not emit duplicate outcome-unknown alerts") + }) + } +} + +func TestBillingSessionReturnsNonDurableEffectErrorWithoutRetryWrapping(t *testing.T) { + funding := &recordingFundingSource{} + session := &BillingSession{ + relayInfo: &relaycommon.RelayInfo{ + RequestId: "non-durable-effect-request", + UserId: 103, + TokenId: 104, + TokenKey: "non-durable-effect-token", + }, + funding: funding, + preConsumedQuota: 10, + } + + err := session.SettleWithEffect(15, &model.BillingSettlementEffect{}) + + require.ErrorIs(t, err, ErrBillingSettlementEffectNotDurable) + require.Equal(t, ErrBillingSettlementEffectNotDurable, err) + assert.Empty(t, funding.deltas) + assert.False(t, session.settled) +} + type recordingBillingSettler struct { preConsumed int reserves []int + settleCalls int +} + +type recordingEffectBillingSettler struct { + recordingBillingSettler + actualQuota int + effect *model.BillingSettlementEffect + err error +} + +func TestSettleAndRecordConsumeHandlesNilContextWithoutPanic(t *testing.T) { + originalLogConsumeEnabled := common.LogConsumeEnabled + common.LogConsumeEnabled = false + t.Cleanup(func() { common.LogConsumeEnabled = originalLogConsumeEnabled }) + + assert.NotPanics(t, func() { + settleAndRecordConsume(nil, nil, false, model.RecordConsumeLogParams{}) + }) + assert.NotPanics(t, func() { + settler := &recordingBillingSettler{} + settleAndRecordConsume(nil, &relaycommon.RelayInfo{Billing: settler}, false, model.RecordConsumeLogParams{}) + assert.Equal(t, 1, settler.settleCalls) + }) +} + +func TestSettleAndRecordConsumeCarriesZeroUsageLogInDurableEffect(t *testing.T) { + truncate(t) + ctx, _ := gin.CreateTestContext(nil) + ctx.Set(common.RequestIdKey, "zero-usage-request") + ctx.Set(common.UpstreamRequestIdKey, "zero-usage-upstream-request") + settler := &recordingEffectBillingSettler{} + info := &relaycommon.RelayInfo{ + RequestId: "relay-fallback-request", + UserId: 601, + Billing: settler, + } + params := model.RecordConsumeLogParams{ + ChannelId: 602, + PromptTokens: 3, + CompletionTokens: 5, + ModelName: "zero-usage-model", + Quota: 0, + Content: "upstream omitted billable usage", + TokenId: 603, + TokenName: "request-time-token-name", + UseTimeSeconds: 7, + IsStream: true, + Group: "default", + Other: map[string]interface{}{"reason": "missing_usage"}, + } + + settleAndRecordConsume(ctx, info, false, params) + + require.NotNil(t, settler.effect) + assert.Zero(t, settler.settleCalls) + assert.Zero(t, settler.actualQuota) + assert.False(t, settler.effect.UpdateUsage) + assert.True(t, settler.effect.QuotaIsActual) + assert.Equal(t, params.Content, settler.effect.Content) + assert.Equal(t, params.ChannelId, settler.effect.ChannelID) + assert.Equal(t, params.ModelName, settler.effect.ModelName) + assert.Equal(t, params.TokenId, settler.effect.TokenID) + assert.Equal(t, params.TokenName, settler.effect.TokenName) + assert.Equal(t, params.PromptTokens, settler.effect.PromptTokens) + assert.Equal(t, params.CompletionTokens, settler.effect.CompletionTokens) + assert.Equal(t, params.UseTimeSeconds, settler.effect.UseTimeSeconds) + assert.Equal(t, params.IsStream, settler.effect.IsStream) + assert.Equal(t, "zero-usage-request", settler.effect.RequestID) + assert.Equal(t, "zero-usage-upstream-request", settler.effect.UpstreamRequestID) +} + +func TestSettleAndRecordConsumeAddsBillingMetadataBeforeDurableEffect(t *testing.T) { + settler := &recordingEffectBillingSettler{} + info := &relaycommon.RelayInfo{ + Billing: settler, + BillingSource: BillingSourceSubscription, + SubscriptionId: 701, + SubscriptionPlanId: 702, + SubscriptionPreConsumed: 10, + SubscriptionAmountTotal: 100, + SubscriptionAmountUsedAfterPreConsume: 10, + UserSetting: dto.UserSetting{ + BillingPreference: "subscription_first", + }, + } + + settleAndRecordConsume(nil, info, false, model.RecordConsumeLogParams{}) + + require.NotNil(t, settler.effect) + require.NotNil(t, settler.effect.Other) + assert.Equal(t, BillingSourceSubscription, settler.effect.Other["billing_source"]) + assert.Equal(t, "subscription_first", settler.effect.Other["billing_preference"]) + assert.Equal(t, 702, settler.effect.Other["subscription_plan_id"]) +} + +func TestSettleBillingWithEffectDoesNotClaimPrePersistenceFailure(t *testing.T) { + settler := &recordingEffectBillingSettler{err: ErrBillingSettlementEffectNotDurable} + info := &relaycommon.RelayInfo{Billing: settler} + + handled, err := SettleBillingWithEffect(nil, info, 10, &model.BillingSettlementEffect{}) + + require.ErrorIs(t, err, ErrBillingSettlementEffectNotDurable) + assert.False(t, handled) +} + +func TestSettleBillingWithEffectDoesNotClaimModelRecordPersistenceFailure(t *testing.T) { + settler := &recordingEffectBillingSettler{err: model.ErrBillingSettlementRecordNotDurable} + info := &relaycommon.RelayInfo{Billing: settler} + + handled, err := SettleBillingWithEffect(nil, info, 10, &model.BillingSettlementEffect{}) + + require.ErrorIs(t, err, model.ErrBillingSettlementRecordNotDurable) + assert.False(t, handled) +} + +func TestSettleBillingWithEffectDoesNotClaimUnknownFundingOutcome(t *testing.T) { + settler := &recordingEffectBillingSettler{err: ErrBillingFundingOutcomeUnknown} + info := &relaycommon.RelayInfo{Billing: settler} + + handled, err := SettleBillingWithEffect(nil, info, 10, &model.BillingSettlementEffect{}) + + require.ErrorIs(t, err, ErrBillingFundingOutcomeUnknown) + assert.False(t, handled) +} + +func TestSettleBillingWithEffectKeepsDurablyOwnedFailureHandled(t *testing.T) { + settlementErr := errors.New("durable settlement remains pending") + settler := &recordingEffectBillingSettler{err: settlementErr} + info := &relaycommon.RelayInfo{Billing: settler} + + handled, err := SettleBillingWithEffect(nil, info, 10, &model.BillingSettlementEffect{}) + + require.ErrorIs(t, err, settlementErr) + assert.True(t, handled) } -func (s *recordingBillingSettler) Settle(int) error { return nil } +func TestSettleAndRecordConsumeDoesNotProjectNonDurableSettlementFailure(t *testing.T) { + truncate(t) + originalLogConsumeEnabled := common.LogConsumeEnabled + originalBatchUpdateEnabled := common.BatchUpdateEnabled + common.LogConsumeEnabled = true + common.BatchUpdateEnabled = false + t.Cleanup(func() { + common.LogConsumeEnabled = originalLogConsumeEnabled + common.BatchUpdateEnabled = originalBatchUpdateEnabled + }) + + cases := []struct { + name string + err error + }{ + {name: "effect not durable", err: ErrBillingSettlementEffectNotDurable}, + {name: "record not durable", err: model.ErrBillingSettlementRecordNotDurable}, + {name: "funding outcome unknown", err: ErrBillingFundingOutcomeUnknown}, + } + + for i, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + userID := 621 + i + channelID := 631 + i + require.NoError(t, model.DB.Create(&model.User{ + Id: userID, + Username: fmt.Sprintf("non-durable-settlement-user-%d", i), + AffCode: fmt.Sprintf("non-durable-%d", i), + Quota: 100, + Status: common.UserStatusEnabled, + }).Error) + seedChannel(t, channelID) + settler := &recordingEffectBillingSettler{err: tc.err} + params := model.RecordConsumeLogParams{ + ChannelId: channelID, + ModelName: "non-durable-settlement-model", + Quota: 10, + Content: tc.name, + } + + settleAndRecordConsume(nil, &relaycommon.RelayInfo{ + UserId: userID, + Billing: settler, + }, true, params) + + require.NotNil(t, settler.effect) + var user model.User + require.NoError(t, model.DB.Select("used_quota", "request_count").First(&user, userID).Error) + assert.Zero(t, user.UsedQuota) + assert.Zero(t, user.RequestCount) + var channel model.Channel + require.NoError(t, model.DB.Select("used_quota").First(&channel, channelID).Error) + assert.Zero(t, channel.UsedQuota) + var logCount int64 + require.NoError(t, model.LOG_DB.Model(&model.Log{}). + Where("user_id = ? AND type = ?", userID, model.LogTypeConsume). + Count(&logCount).Error) + assert.Zero(t, logCount) + }) + } +} + +func TestSettleAndRecordConsumePersistsLegacyEffectBeforeFailedSettlement(t *testing.T) { + truncate(t) + const ( + userID = 611 + tokenID = 612 + ) + seedUser(t, userID, 0) + seedToken(t, tokenID, userID, "legacy-effect-token", 0) + info := &relaycommon.RelayInfo{ + RequestId: "legacy-effect-request", + UserId: userID, + TokenId: tokenID, + TokenKey: "legacy-effect-token", + } + params := model.RecordConsumeLogParams{ + ChannelId: 613, + ModelName: "legacy-effect-model", + TokenId: tokenID, + Group: "default", + Quota: 10, + Content: "legacy settlement waits for funding", + } + + settleAndRecordConsume(nil, info, true, params) + + operationKey := "request:legacy-effect-request:legacy-post:finalize" + var settlement model.BillingSettlement + require.NoError(t, model.DB.Where("operation_key = ?", operationKey).First(&settlement).Error) + require.Equal(t, model.BillingSettlementStatusManual, settlement.Status) + require.NotEmpty(t, settlement.EffectPayload) + require.Empty(t, settlement.EffectStatus) + var logCount int64 + require.NoError(t, model.LOG_DB.Model(&model.Log{}).Where("user_id = ? AND type = ?", userID, model.LogTypeConsume).Count(&logCount).Error) + require.Zero(t, logCount) + + require.NoError(t, model.DB.Model(&model.User{}).Where("id = ?", userID).Update("quota", 100).Error) + require.NoError(t, model.DB.Model(&model.Token{}).Where("id = ?", tokenID).Updates(map[string]interface{}{ + "remain_quota": 100, + "used_quota": 0, + }).Error) + require.NoError(t, model.DB.Model(&model.BillingSettlement{}). + Where("id = ?", settlement.ID). + Updates(map[string]interface{}{ + "status": model.BillingSettlementStatusPending, + "last_error": "", + "next_attempt": time.Now().Unix(), + }).Error) + model.ProcessPendingBillingSettlementsOnce() + + settleAndRecordConsume(nil, info, true, params) + settleAndRecordConsume(nil, info, true, params) + + require.EqualValues(t, 90, getUserQuota(t, userID)) + require.Equal(t, 90, getTokenRemainQuota(t, tokenID)) + require.NoError(t, model.LOG_DB.Model(&model.Log{}).Where("user_id = ? AND type = ?", userID, model.LogTypeConsume).Count(&logCount).Error) + require.EqualValues(t, 1, logCount) + var storedUser model.User + require.NoError(t, model.DB.Select("used_quota", "request_count").First(&storedUser, userID).Error) + require.EqualValues(t, 10, storedUser.UsedQuota) + require.Equal(t, 1, storedUser.RequestCount) +} + +func TestSettleAndRecordConsumeProjectsLocallyWhenExistingSettlementHasNoEffect(t *testing.T) { + truncate(t) + originalLogConsumeEnabled := common.LogConsumeEnabled + common.LogConsumeEnabled = true + t.Cleanup(func() { common.LogConsumeEnabled = originalLogConsumeEnabled }) + + const userID = 614 + seedUser(t, userID, 100) + const requestID = "legacy-effectless-settlement" + operationKey := "request:" + requestID + ":legacy-post:finalize" + _, _, err := model.ApplyBillingSettlementOnce(model.BillingSettlementInput{ + OperationKey: operationKey, + Source: model.BillingSettlementSourceWallet, + UserID: userID, + SubscriptionPreConsumeRequestID: requestID, + }) + require.NoError(t, err) + ownsEffect, err := model.BillingSettlementOwnsEffect(operationKey) + require.NoError(t, err) + require.False(t, ownsEffect) + + settleAndRecordConsume(nil, &relaycommon.RelayInfo{ + RequestId: requestID, + UserId: userID, + }, false, model.RecordConsumeLogParams{ + ModelName: "legacy-effectless-model", + Content: "effectless settlement projection", + }) + + var logs []model.Log + require.NoError(t, model.LOG_DB.Where("user_id = ? AND type = ?", userID, model.LogTypeConsume).Find(&logs).Error) + require.Len(t, logs, 1) + require.Equal(t, "effectless settlement projection", logs[0].Content) +} + +func TestSettleAndRecordConsumeProjectsLocallyWhenBillingSessionSettlementHasNoEffect(t *testing.T) { + truncate(t) + originalLogConsumeEnabled := common.LogConsumeEnabled + common.LogConsumeEnabled = true + t.Cleanup(func() { common.LogConsumeEnabled = originalLogConsumeEnabled }) + + const userID = 615 + seedUser(t, userID, 100) + const requestID = "billing-session-effectless-settlement" + operationKey := model.BillingRequestFinalizeOperationKey(requestID) + _, _, err := model.ApplyBillingSettlementOnce(model.BillingSettlementInput{ + OperationKey: operationKey, + Source: model.BillingSettlementSourceWallet, + UserID: userID, + }) + require.NoError(t, err) + + info := &relaycommon.RelayInfo{ + RequestId: requestID, + UserId: userID, + } + info.Billing = &BillingSession{ + relayInfo: info, + funding: &WalletFunding{userId: userID}, + } + settleAndRecordConsume(nil, info, false, model.RecordConsumeLogParams{ + ModelName: "billing-session-effectless-model", + Content: "billing session effectless projection", + }) + + var logs []model.Log + require.NoError(t, model.LOG_DB.Where("user_id = ? AND type = ?", userID, model.LogTypeConsume).Find(&logs).Error) + require.Len(t, logs, 1) + require.Equal(t, "billing session effectless projection", logs[0].Content) +} + +func (s *recordingBillingSettler) Settle(int) error { s.settleCalls++; return nil } func (s *recordingBillingSettler) Refund(*gin.Context) {} func (s *recordingBillingSettler) NeedsRefund() bool { return false } func (s *recordingBillingSettler) GetPreConsumedQuota() int { return s.preConsumed } @@ -47,6 +532,12 @@ func (s *recordingBillingSettler) Reserve(target int) error { return nil } +func (s *recordingEffectBillingSettler) SettleWithEffect(actualQuota int, effect *model.BillingSettlementEffect) error { + s.actualQuota = actualQuota + s.effect = effect + return s.err +} + func TestBillingSessionPrepareSettlementPersistsZeroDeltaTaskEffect(t *testing.T) { session := &BillingSession{ relayInfo: &relaycommon.RelayInfo{ @@ -69,6 +560,57 @@ func TestBillingSessionPrepareSettlementPersistsZeroDeltaTaskEffect(t *testing.T assert.Zero(t, input.TokenDelta) } +func TestBillingSessionPrepareSettlementBindsTaskReservationAndTarget(t *testing.T) { + session := &BillingSession{ + relayInfo: &relaycommon.RelayInfo{ + RequestId: "task-reservation-binding", + UserId: 43, + TokenId: 44, + TokenKey: "task-reservation-token", + TaskRelayInfo: &relaycommon.TaskRelayInfo{PersistedTaskID: 45}, + }, + funding: &WalletFunding{userId: 43}, + preConsumedQuota: 100, + } + + input, err := session.PrepareSettlement(175) + + require.NoError(t, err) + require.NotNil(t, input) + assert.EqualValues(t, 45, input.TaskID) + assert.EqualValues(t, 100, input.TaskQuota) + assert.EqualValues(t, 175, input.TaskQuotaTarget) + assert.EqualValues(t, 75, input.FundingDelta) +} + +func TestBillingSessionSettleAppliesPersistedZeroDeltaTaskFinalize(t *testing.T) { + truncate(t) + const userID = 46 + task := makeTask(userID, 0, 10, 0, BillingSourceWallet, 0) + task.Status = model.TaskStatusSubmitted + persistTask(t, task) + session := &BillingSession{ + relayInfo: &relaycommon.RelayInfo{ + RequestId: "task-zero-delta-finalize", + UserId: userID, + TaskRelayInfo: &relaycommon.TaskRelayInfo{PersistedTaskID: task.ID}, + }, + funding: &WalletFunding{userId: userID}, + preConsumedQuota: 10, + } + input, err := session.PrepareSettlement(10) + require.NoError(t, err) + require.NotNil(t, input) + require.NoError(t, task.UpdateWithSettlementIntent(input)) + + require.NoError(t, session.Settle(10)) + + var settlement model.BillingSettlement + require.NoError(t, model.DB.Where("operation_key = ?", input.OperationKey).First(&settlement).Error) + assert.Equal(t, model.BillingSettlementStatusApplied, settlement.Status) + assert.True(t, session.settled) +} + func TestBillingSessionCompensatesFundingWhenTokenSettlementFails(t *testing.T) { truncate(t) funding := &recordingFundingSource{} @@ -170,6 +712,101 @@ func TestBillingSessionWalletPreConsumeUsesDurableRequestIdentity(t *testing.T) assert.Equal(t, model.BillingSettlementStatusApplied, settlement.Status) } +func TestNewBillingSessionRejectsUnresolvedPositiveFinalizeSettlement(t *testing.T) { + require.NoError(t, i18n.Init()) + truncate(t) + const userID, tokenID = 851, 852 + seedUser(t, userID, 100) + seedToken(t, tokenID, userID, "unresolved-finalize-token", 100) + now := time.Now().Unix() + require.NoError(t, model.DB.Create(&model.BillingSettlement{ + OperationKey: "request:unresolved-finalize:finalize", + Source: model.BillingSettlementSourceWallet, + UserID: userID, + TokenID: tokenID, + FundingDelta: 10, + TokenDelta: 10, + Status: model.BillingSettlementStatusManual, + CreatedAt: now, + UpdatedAt: now, + Revision: 1, + }).Error) + + ctx, _ := gin.CreateTestContext(nil) + session, apiErr := NewBillingSession(ctx, &relaycommon.RelayInfo{ + RequestId: "blocked-after-finalize", + UserId: userID, + TokenId: tokenID, + TokenKey: "unresolved-finalize-token", + OriginModelName: "blocked-finalize-model", + UserSetting: dto.UserSetting{ + BillingPreference: "wallet_only", + Language: i18n.LangEn, + }, + }, 10) + + require.Nil(t, session) + require.NotNil(t, apiErr) + assert.Equal( + t, + types.ErrorCodeBillingReconciliationPending, + apiErr.GetErrorCode(), + ) + assert.Equal( + t, + "Billing reconciliation is incomplete. Do not submit another request; contact an administrator for assistance.", + apiErr.Error(), + ) + assert.EqualValues(t, 100, getUserQuota(t, userID)) + assert.EqualValues(t, 100, getTokenRemainQuota(t, tokenID)) +} + +func TestInsufficientFinalizeBlocksFurtherBillingSessions(t *testing.T) { + require.NoError(t, i18n.Init()) + truncate(t) + const userID, tokenID = 853, 854 + seedUser(t, userID, 15) + seedToken(t, tokenID, userID, "insufficient-finalize-token", 100) + ctx, _ := gin.CreateTestContext(nil) + newInfo := func(requestID string) *relaycommon.RelayInfo { + return &relaycommon.RelayInfo{ + RequestId: requestID, + UserId: userID, + TokenId: tokenID, + TokenKey: "insufficient-finalize-token", + OriginModelName: "insufficient-finalize-model", + UserSetting: dto.UserSetting{ + BillingPreference: "wallet_only", + Language: i18n.LangZhCN, + }, + } + } + + session, apiErr := NewBillingSession(ctx, newInfo("insufficient-finalize-request"), 10) + require.Nil(t, apiErr) + require.Error(t, session.Settle(20)) + assert.EqualValues(t, 5, getUserQuota(t, userID)) + assert.EqualValues(t, 90, getTokenRemainQuota(t, tokenID)) + + var settlement model.BillingSettlement + require.NoError(t, model.DB.Where("operation_key = ?", "request:insufficient-finalize-request:finalize").First(&settlement).Error) + assert.Equal(t, model.BillingSettlementStatusManual, settlement.Status) + assert.EqualValues(t, 10, settlement.FundingDelta) + assert.Zero(t, settlement.AppliedFundingDelta) + + nextSession, nextErr := NewBillingSession(ctx, newInfo("blocked-after-insufficient-finalize"), 5) + require.Nil(t, nextSession) + require.NotNil(t, nextErr) + assert.Equal( + t, + types.ErrorCodeBillingReconciliationPending, + nextErr.GetErrorCode(), + ) + assert.Equal(t, "存在未完成的计费对账,请勿重复提交请求并联系管理员处理", nextErr.Error()) + assert.EqualValues(t, 5, getUserQuota(t, userID)) + assert.EqualValues(t, 90, getTokenRemainQuota(t, tokenID)) +} + func TestBillingSessionWalletFirstReplayCannotSwitchFromWalletToSubscription(t *testing.T) { truncate(t) const userID, tokenID, planID, subscriptionID = 805, 806, 807, 808 @@ -231,6 +868,55 @@ func TestBillingSessionWalletFirstReplayCannotSwitchFromSubscriptionToWallet(t * assert.EqualValues(t, 1, countSubscriptionPreConsumeRecords(t, "subscription-replay-request")) } +func TestSettleAndRecordConsumeLogsPostSettlementSubscriptionUsage(t *testing.T) { + truncate(t) + const userID, tokenID, channelID, planID, subscriptionID = 819, 820, 821, 822, 823 + const preConsumedQuota, actualQuota int64 = 10, 15 + seedUser(t, userID, 0) + seedToken(t, tokenID, userID, "subscription-log-token", 100) + seedChannel(t, channelID) + seedBillingSubscription(t, planID, subscriptionID, userID, 100) + + ctx, _ := gin.CreateTestContext(nil) + info := &relaycommon.RelayInfo{ + RequestId: "subscription-log-request", UserId: userID, + TokenId: tokenID, TokenKey: "subscription-log-token", + OriginModelName: "subscription-log-model", UsingGroup: "default", + ChannelMeta: &relaycommon.ChannelMeta{ChannelId: channelID}, + UserSetting: dto.UserSetting{BillingPreference: "subscription_first"}, + } + session, apiErr := NewBillingSession(ctx, info, int(preConsumedQuota)) + require.Nil(t, apiErr) + require.NotNil(t, session) + info.Billing = session + other := map[string]interface{}{} + appendBillingInfo(info, other) + + settleAndRecordConsume(ctx, info, true, model.RecordConsumeLogParams{ + ChannelId: channelID, + ModelName: "subscription-log-model", + TokenId: tokenID, + Group: "default", + Quota: int(actualQuota), + Content: "subscription settlement metadata", + Other: other, + }) + + require.Equal(t, actualQuota, getSubscriptionAmountUsed(t, subscriptionID)) + log := getLastLog(t) + var metadata struct { + PostDelta int64 `json:"subscription_post_delta"` + Consumed int64 `json:"subscription_consumed"` + Used int64 `json:"subscription_used"` + Remain int64 `json:"subscription_remain"` + } + require.NoError(t, common.UnmarshalJsonStr(log.Other, &metadata)) + assert.Equal(t, actualQuota-preConsumedQuota, metadata.PostDelta) + assert.Equal(t, actualQuota, metadata.Consumed) + assert.Equal(t, actualQuota, metadata.Used) + assert.Equal(t, int64(100)-actualQuota, metadata.Remain) +} + func TestBillingSessionTrustedWalletReplayFailsClosedInsteadOfSwitchingFunding(t *testing.T) { truncate(t) const userID, tokenID, planID, subscriptionID = 815, 816, 817, 818 @@ -382,6 +1068,8 @@ func TestViolationFeeSettlesOriginalSubscriptionPreConsume(t *testing.T) { t.Cleanup(func() { common.QuotaPerUnit = oldQuotaPerUnit }) ctx, _ := gin.CreateTestContext(nil) + ctx.Set(common.RequestIdKey, "violation-client-request") + ctx.Set(common.UpstreamRequestIdKey, "violation-upstream-request") info := &relaycommon.RelayInfo{ RequestId: "violation-subscription-request", UserId: userID, TokenId: tokenID, TokenKey: "violation-subscription-token", @@ -411,7 +1099,10 @@ func TestViolationFeeSettlesOriginalSubscriptionPreConsume(t *testing.T) { require.EqualValues(t, feeQuota, getSubscriptionAmountUsed(t, subscriptionID)) require.EqualValues(t, 100-feeQuota, getTokenRemainQuota(t, tokenID)) require.Equal(t, int64(1), countLogs(t)) - require.Equal(t, feeQuota, getLastLog(t).Quota) + log := getLastLog(t) + require.Equal(t, feeQuota, log.Quota) + require.Equal(t, "violation-client-request", log.RequestId) + require.Equal(t, "violation-upstream-request", log.UpstreamRequestId) var finalize model.BillingSettlement require.NoError(t, model.DB.Where("operation_key = ?", "request:violation-subscription-request:finalize").First(&finalize).Error) @@ -422,6 +1113,76 @@ func TestViolationFeeSettlesOriginalSubscriptionPreConsume(t *testing.T) { require.Zero(t, legacyCount) } +func TestViolationFeeProjectsLocallyWhenExistingSettlementHasNoEffect(t *testing.T) { + truncate(t) + originalLogConsumeEnabled := common.LogConsumeEnabled + common.LogConsumeEnabled = true + t.Cleanup(func() { common.LogConsumeEnabled = originalLogConsumeEnabled }) + + const userID, tokenID, channelID = 835, 836, 837 + const requestID = "violation-effectless-request" + seedUser(t, userID, 100) + seedToken(t, tokenID, userID, "violation-effectless-token", 100) + seedChannel(t, channelID) + + oldQuotaPerUnit := common.QuotaPerUnit + common.QuotaPerUnit = 100 + t.Cleanup(func() { common.QuotaPerUnit = oldQuotaPerUnit }) + feeQuota := calcViolationFeeQuota(0.05, 1) + operationKey := model.BillingRequestFinalizeOperationKey(requestID) + _, _, err := model.ApplyBillingSettlementOnce(model.BillingSettlementInput{ + OperationKey: operationKey, + Source: model.BillingSettlementSourceWallet, + UserID: userID, + TokenID: tokenID, + TokenKey: "violation-effectless-token", + FundingDelta: int64(feeQuota), + TokenDelta: int64(feeQuota), + }) + require.NoError(t, err) + require.EqualValues(t, 100-feeQuota, getUserQuota(t, userID)) + require.EqualValues(t, 100-feeQuota, getTokenRemainQuota(t, tokenID)) + + ctx, _ := gin.CreateTestContext(nil) + ctx.Set(common.RequestIdKey, "violation-effectless-client-request") + ctx.Set(common.UpstreamRequestIdKey, "violation-effectless-upstream-request") + info := &relaycommon.RelayInfo{ + RequestId: requestID, UserId: userID, + TokenId: tokenID, TokenKey: "violation-effectless-token", + ChannelMeta: &relaycommon.ChannelMeta{ChannelId: channelID}, + OriginModelName: "grok-effectless-test", UsingGroup: "default", + BillingSource: BillingSourceWallet, + StartTime: time.Now(), + } + info.PriceData.GroupRatioInfo.GroupRatio = 1 + info.Billing = &BillingSession{ + relayInfo: info, + funding: &WalletFunding{userId: userID}, + } + apiErr := types.NewErrorWithStatusCode( + errors.New(CSAMViolationMarker), types.ErrorCodeBadResponseStatusCode, + http.StatusBadRequest, + ) + + HandleFailedBilling(ctx, info, NormalizeViolationFeeError(apiErr)) + HandleFailedBilling(ctx, info, NormalizeViolationFeeError(apiErr)) + + require.EqualValues(t, 100-feeQuota, getUserQuota(t, userID)) + require.EqualValues(t, 100-feeQuota, getTokenRemainQuota(t, tokenID)) + require.Equal(t, int64(1), countLogs(t)) + log := getLastLog(t) + require.Equal(t, feeQuota, log.Quota) + require.Equal(t, "violation-effectless-client-request", log.RequestId) + require.Equal(t, "violation-effectless-upstream-request", log.UpstreamRequestId) + var storedUser model.User + require.NoError(t, model.DB.Select("used_quota", "request_count").First(&storedUser, userID).Error) + require.EqualValues(t, feeQuota, storedUser.UsedQuota) + require.Equal(t, 1, storedUser.RequestCount) + var channel model.Channel + require.NoError(t, model.DB.Select("used_quota").First(&channel, channelID).Error) + require.EqualValues(t, feeQuota, channel.UsedQuota) +} + type compensationFailingFundingSource struct { deltas []int } @@ -437,7 +1198,7 @@ func (f *compensationFailingFundingSource) Settle(delta int) (int64, error) { return int64(delta), nil } -func TestBillingSessionKeepsCommittedFundingRetryableWhenCompensationFails(t *testing.T) { +func TestBillingSessionFailsClosedWhenCompensationOutcomeIsUnknown(t *testing.T) { truncate(t) funding := &compensationFailingFundingSource{} session := &BillingSession{ @@ -449,14 +1210,15 @@ func TestBillingSessionKeepsCommittedFundingRetryableWhenCompensationFails(t *te assert.Equal(t, []int{5, -5}, funding.deltas) assert.True(t, session.fundingSettled) assert.True(t, session.compensationFailed) + assert.True(t, session.fundingOutcomeUnknown) assert.False(t, session.settled) require.NoError(t, model.DB.Create(&model.Token{ Id: 52, UserId: 51, Key: "retry-token", Status: common.TokenStatusEnabled, RemainQuota: 20, }).Error) - require.NoError(t, session.Settle(15)) + require.ErrorIs(t, session.Settle(15), ErrBillingFundingOutcomeUnknown) assert.Equal(t, []int{5, -5}, funding.deltas) - assert.True(t, session.settled) + assert.False(t, session.settled) } type ambiguousCompensationFundingSource struct { @@ -486,6 +1248,9 @@ func TestBillingSessionDoesNotReapplyFundingAfterAmbiguousCompensationError(t *t assert.Equal(t, []int{5, -5}, funding.deltas) assert.EqualValues(t, 5, session.appliedFundingDelta) assert.True(t, session.compensationFailed) + assert.True(t, session.fundingOutcomeUnknown) + require.ErrorIs(t, session.Settle(15), ErrBillingFundingOutcomeUnknown) + assert.Equal(t, []int{5, -5}, funding.deltas) } type partialCompensationFundingSource struct { @@ -503,7 +1268,7 @@ func (f *partialCompensationFundingSource) Settle(delta int) (int64, error) { return int64(delta), nil } -func TestBillingSessionReconcilesPartialFundingCompensationBeforeRetry(t *testing.T) { +func TestBillingSessionFailsClosedAfterPartialFundingCompensation(t *testing.T) { truncate(t) funding := &partialCompensationFundingSource{} session := &BillingSession{ @@ -511,19 +1276,20 @@ func TestBillingSessionReconcilesPartialFundingCompensationBeforeRetry(t *testin funding: funding, preConsumedQuota: 10, } - require.Error(t, session.Settle(15)) - require.Equal(t, []int{5, -5, 2, -5, 2, -5}, funding.deltas) + require.ErrorIs(t, session.Settle(15), ErrBillingFundingOutcomeUnknown) + require.Equal(t, []int{5, -5}, funding.deltas) require.False(t, session.compensationFailed) - require.True(t, session.fundingReconcilePending) + require.True(t, session.fundingOutcomeUnknown) + require.False(t, session.fundingReconcilePending) require.EqualValues(t, 3, session.appliedFundingDelta) require.NoError(t, model.DB.Create(&model.Token{ Id: 57, UserId: 56, Key: "partial-compensation-token", Status: common.TokenStatusEnabled, RemainQuota: 20, }).Error) - require.NoError(t, session.Settle(15)) - assert.True(t, session.settled) - assert.EqualValues(t, 5, session.appliedFundingDelta) - assert.Equal(t, int64(2), int64(funding.deltas[len(funding.deltas)-1])) + require.ErrorIs(t, session.Settle(15), ErrBillingFundingOutcomeUnknown) + assert.False(t, session.settled) + assert.EqualValues(t, 3, session.appliedFundingDelta) + assert.Equal(t, []int{5, -5}, funding.deltas) } type tokenRecoveringFundingSource struct { diff --git a/service/funding_source.go b/service/funding_source.go index 2d55fa20..0ae7ee98 100644 --- a/service/funding_source.go +++ b/service/funding_source.go @@ -49,9 +49,15 @@ func (w *WalletFunding) Settle(delta int) (int64, error) { return 0, nil } if delta > 0 { - return int64(delta), model.DecreaseUserQuota(w.userId, delta, false) + if err := model.DecreaseUserQuota(w.userId, delta, false); err != nil { + return 0, err + } + return int64(delta), nil + } + if err := model.IncreaseUserQuota(w.userId, -delta, false); err != nil { + return 0, err } - return int64(delta), model.IncreaseUserQuota(w.userId, -delta, false) + return int64(delta), nil } func (w *WalletFunding) Refund() error { diff --git a/service/midjourney_task.go b/service/midjourney_task.go index e41ffe5b..1934a714 100644 --- a/service/midjourney_task.go +++ b/service/midjourney_task.go @@ -14,77 +14,32 @@ import ( // 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) + billingTask, err := model.GetMidjourneyBillingTask(task.UserId, task.ChannelId, task.MjId) + if err != nil { + return err } - if responseItem.Properties != nil { - properties, err := common.Marshal(responseItem.Properties) - if err != nil { - return err - } - task.Properties = string(properties) + preStatus := task.Status + mergedTask := *task + providerTerminal, isFailure, err := mergeMidjourneyTaskResponse(&mergedTask, responseItem) + if err != nil { + return err } - if responseItem.Buttons != nil { - buttons, err := common.Marshal(responseItem.Buttons) - if err != nil { - return err + legacyNeedsUpdate = legacyNeedsUpdate || mergedTask.Status != task.Status || + mergedTask.Progress != task.Progress || mergedTask.FailReason != task.FailReason + if billingTask != nil { + pending, settlementErr := taskTerminalSettlementPending(billingTask, providerTerminal) + if settlementErr != nil { + return settlementErr } - task.Buttons = string(buttons) - } - if responseItem.VideoUrls != nil { - videoURLs, err := common.Marshal(responseItem.VideoUrls) - if err != nil { - return err + if pending { + return nil } - 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 != "") + *task = mergedTask 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 @@ -141,6 +96,76 @@ func UpdateMidjourneyTaskFromResponse(ctx context.Context, task *model.Midjourne return nil } +func mergeMidjourneyTaskResponse(task *model.Midjourney, responseItem dto.MidjourneyDto) (providerTerminal, isFailure bool, err error) { + 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 false, false, err + } + task.Properties = string(properties) + } + if responseItem.Buttons != nil { + buttons, err := common.Marshal(responseItem.Buttons) + if err != nil { + return false, false, err + } + task.Buttons = string(buttons) + } + if responseItem.VideoUrls != nil { + videoURLs, err := common.Marshal(responseItem.VideoUrls) + if err != nil { + return false, false, err + } + task.VideoUrls = string(videoURLs) + } + + providerSucceeded := responseItem.Status == "SUCCESS" || + (responseItem.Status == "" && responseItem.Progress == "100%" && responseItem.FailReason == "") + if providerSucceeded { + task.Status = "SUCCESS" + task.FailReason = "" + } + isFailure = !providerSucceeded && (task.Status == "FAILURE" || task.FailReason != "") + if isFailure { + task.Status = "FAILURE" + task.Progress = "100%" + } + providerTerminal = task.Status == "SUCCESS" || task.Status == "FAILURE" || + task.FailReason != "" || (task.Status == "" && task.Progress == "100%") + return providerTerminal, isFailure, nil +} + func midjourneyBillingTaskStatus(status, progress string) model.TaskStatus { switch status { case "SUCCESS": diff --git a/service/quota.go b/service/quota.go index 91728f1f..d4443755 100644 --- a/service/quota.go +++ b/service/quota.go @@ -238,15 +238,6 @@ func PostWssConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, mod logContent += fmt.Sprintf(", stream ended abnormally (%s), billed pre-consumed quota %d", relayInfo.StreamStatus.EndReason, fallbackQuota) logger.LogError(ctx, fmt.Sprintf("stream ended abnormally without billable usage, billing fallback quota, userId %d, channelId %d, tokenId %d, model %s, reason %s, fallback quota %d", relayInfo.UserId, relayInfo.ChannelId, relayInfo.TokenId, modelName, relayInfo.StreamStatus.EndReason, fallbackQuota)) } - if shouldUpdateUsageStats { - model.UpdateUserUsedQuotaAndRequestCount(relayInfo.UserId, quota) - model.UpdateChannelUsedQuota(relayInfo.ChannelId, quota) - } - - if err := SettleBilling(ctx, relayInfo, quota); err != nil { - logger.LogError(ctx, "error settling billing: "+err.Error()) - } - logModel := modelName if extraContent != "" { logContent += ", " + extraContent @@ -257,7 +248,7 @@ func PostWssConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, mod InjectTieredBillingInfo(other, relayInfo, tieredResult) } attachQuotaSaturation(ctx, relayInfo, other) - model.RecordConsumeLog(ctx, relayInfo.UserId, model.RecordConsumeLogParams{ + settleAndRecordConsume(ctx, relayInfo, shouldUpdateUsageStats, model.RecordConsumeLogParams{ ChannelId: relayInfo.ChannelId, PromptTokens: usage.InputTokens, CompletionTokens: usage.OutputTokens, @@ -369,15 +360,6 @@ func PostAudioConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, u logContent += fmt.Sprintf(", stream ended abnormally (%s), billed pre-consumed quota %d", relayInfo.StreamStatus.EndReason, fallbackQuota) logger.LogError(ctx, fmt.Sprintf("stream ended abnormally without billable usage, billing fallback quota, userId %d, channelId %d, tokenId %d, model %s, reason %s, fallback quota %d", relayInfo.UserId, relayInfo.ChannelId, relayInfo.TokenId, relayInfo.OriginModelName, relayInfo.StreamStatus.EndReason, fallbackQuota)) } - if shouldUpdateUsageStats { - model.UpdateUserUsedQuotaAndRequestCount(relayInfo.UserId, quota) - model.UpdateChannelUsedQuota(relayInfo.ChannelId, quota) - } - - if err := SettleBilling(ctx, relayInfo, quota); err != nil { - logger.LogError(ctx, "error settling billing: "+err.Error()) - } - logModel := relayInfo.OriginModelName if extraContent != "" { logContent += ", " + extraContent @@ -388,7 +370,7 @@ func PostAudioConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, u InjectTieredBillingInfo(other, relayInfo, tieredResult) } attachQuotaSaturation(ctx, relayInfo, other) - model.RecordConsumeLog(ctx, relayInfo.UserId, model.RecordConsumeLogParams{ + settleAndRecordConsume(ctx, relayInfo, shouldUpdateUsageStats, model.RecordConsumeLogParams{ ChannelId: relayInfo.ChannelId, PromptTokens: usage.PromptTokens, CompletionTokens: usage.CompletionTokens, @@ -436,14 +418,26 @@ func PostConsumeQuota(relayInfo *relaycommon.RelayInfo, quota int, preConsumedQu } func PostConsumeQuotaOnce(relayInfo *relaycommon.RelayInfo, operationKind string, quota int, preConsumedQuota int, sendEmail bool) (err error) { + _, err = postConsumeQuotaOnceWithEffect(relayInfo, operationKind, quota, preConsumedQuota, sendEmail, nil) + return err +} + +func postConsumeQuotaOnceWithEffect( + relayInfo *relaycommon.RelayInfo, + operationKind string, + quota int, + preConsumedQuota int, + sendEmail bool, + effect *model.BillingSettlementEffect, +) (effectHandled bool, err error) { if relayInfo == nil { - return errors.New("relayInfo is nil") + return false, errors.New("relayInfo is nil") } if operationKind == "" { - return errors.New("post-consume operation kind is required") + return false, errors.New("post-consume operation kind is required") } - if quota == 0 { - return nil + if quota == 0 && effect == nil { + return false, nil } if relayInfo.RequestId != "" { source := model.BillingSettlementSourceWallet @@ -466,9 +460,26 @@ func PostConsumeQuotaOnce(relayInfo *relaycommon.RelayInfo, operationKind string TokenDelta: tokenDelta, SubscriptionPreConsumeRequestID: relayInfo.RequestId, AllowMissingToken: quota < 0, + Effect: effect, }) if durableErr != nil { - return durableErr + if effect != nil && !errors.Is(durableErr, model.ErrBillingSettlementRecordNotDurable) { + return true, durableErr + } + return false, durableErr + } + effectHandled := false + if effect != nil { + ownsEffect, ownershipErr := model.BillingSettlementOwnsEffect(operationKey) + if ownershipErr != nil { + return true, ownershipErr + } + if ownsEffect { + effectHandled = true + if effectErr := model.ProcessBillingSettlementEffect(operationKey); effectErr != nil { + return true, effectErr + } + } } if sendEmail && (quota+preConsumedQuota) != 0 { if relayInfo.BillingSource == BillingSourceSubscription { @@ -477,10 +488,10 @@ func PostConsumeQuotaOnce(relayInfo *relaycommon.RelayInfo, operationKind string checkAndSendQuotaNotify(relayInfo, quota, preConsumedQuota) } } - return nil + return effectHandled, nil } - return errors.New("balance mutation requires a stable request id") + return false, errors.New("balance mutation requires a stable request id") } func checkAndSendQuotaNotify(relayInfo *relaycommon.RelayInfo, quota int, preConsumedQuota int) { diff --git a/service/task_billing.go b/service/task_billing.go index eb7b841f..06299ed6 100644 --- a/service/task_billing.go +++ b/service/task_billing.go @@ -4,6 +4,7 @@ import ( "context" "fmt" "strings" + "time" "github.com/MAX-API-Next/MAX-API/common" "github.com/MAX-API-Next/MAX-API/constant" @@ -66,18 +67,27 @@ func BuildTaskSubmissionSettlementEffect(c *gin.Context, info *relaycommon.Relay return nil } logContent, other := buildTaskConsumptionLogData(c, info) - return &model.BillingSettlementEffect{ - LogType: model.LogTypeConsume, - Content: logContent, - ChannelID: info.ChannelId, - ModelName: info.OriginModelName, - TokenID: info.TokenId, - Group: info.UsingGroup, - Other: other, - NodeName: common.NodeName, - UpdateUsage: true, - Quota: int64(quota), + useTimeSeconds := 0 + if !info.StartTime.IsZero() { + useTimeSeconds = int(time.Since(info.StartTime).Seconds()) } + tokenName := "" + if c != nil { + tokenName = c.GetString("token_name") + } + requestID, upstreamRequestID := billingEffectRequestIDs(c, info) + return newConsumeBillingSettlementEffect(info, model.RecordConsumeLogParams{ + ChannelId: info.ChannelId, + ModelName: info.OriginModelName, + Quota: quota, + Content: logContent, + TokenId: info.TokenId, + TokenName: tokenName, + UseTimeSeconds: useTimeSeconds, + IsStream: info.IsStream, + Group: info.UsingGroup, + Other: other, + }, requestID, upstreamRequestID, true) } // LogTaskConsumption records task usage for legacy/custom funding paths that @@ -145,6 +155,37 @@ func taskModelName(task *model.Task) string { return task.Properties.OriginModelName } +// taskFinalSettlementPending reports whether the submission-finalize funding +// settlement is pending or requires manual review. Callers use true to defer a +// provider-driven terminal transition until the funding state is applied. +func taskFinalSettlementPending(task *model.Task) (bool, error) { + if task == nil || task.PrivateData.BillingRequestId == "" { + return false, nil + } + status, found, err := model.GetBillingSettlementStatus(model.BillingRequestFinalizeOperationKey(task.PrivateData.BillingRequestId)) + if err != nil || !found { + return false, err + } + switch status { + case "", model.BillingSettlementStatusApplied: + return false, nil + case model.BillingSettlementStatusPending, model.BillingSettlementStatusManual: + return true, nil + default: + return true, fmt.Errorf("unknown task submission settlement status %q", status) + } +} + +// taskTerminalSettlementPending gates the submission-settlement lookup on a +// provider transition that would make the task terminal. Callers keep their +// own error reporting and skip/return behavior. +func taskTerminalSettlementPending(task *model.Task, providerTerminal bool) (bool, error) { + if !providerTerminal { + return false, nil + } + return taskFinalSettlementPending(task) +} + func BuildTaskRefundSettlementInput(task *model.Task, reason string) *model.BillingSettlementInput { if task == nil || task.Quota == 0 || task.ID <= 0 { return nil diff --git a/service/task_billing_test.go b/service/task_billing_test.go index 7b5be2a9..c05e76ce 100644 --- a/service/task_billing_test.go +++ b/service/task_billing_test.go @@ -15,6 +15,7 @@ import ( "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/gin-gonic/gin" "github.com/glebarez/sqlite" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" @@ -54,6 +55,8 @@ func TestMain(m *testing.M) { &model.BillingSettlement{}, &model.BillingPreConsumeSelection{}, &model.CacheInvalidationTask{}, + &model.Midjourney{}, + &model.MidjourneyBillingClaim{}, ); err != nil { panic("failed to migrate: " + err.Error()) } @@ -76,8 +79,38 @@ func TestShouldApplyTaskResultProgressSkipsNonTerminalCompleteProgress(t *testin })) } +func TestBuildTaskSubmissionSettlementEffectCarriesRequestMetadata(t *testing.T) { + ctx, _ := gin.CreateTestContext(nil) + ctx.Set(common.RequestIdKey, "task-client-request") + ctx.Set(common.UpstreamRequestIdKey, "task-upstream-request") + ctx.Set("token_name", "request-time-task-token") + info := &relaycommon.RelayInfo{ + RequestId: "task-relay-fallback-request", + ChannelMeta: &relaycommon.ChannelMeta{ChannelId: 91}, + TaskRelayInfo: &relaycommon.TaskRelayInfo{Action: constant.TaskActionGenerate}, + OriginModelName: "task-model", + TokenId: 92, + UsingGroup: "default", + StartTime: time.Now().Add(-3 * time.Second), + IsStream: true, + } + + effect := BuildTaskSubmissionSettlementEffect(ctx, info, 123) + + require.NotNil(t, effect) + assert.True(t, effect.UpdateUsage) + assert.True(t, effect.QuotaIsActual) + assert.EqualValues(t, 123, effect.Quota) + assert.Equal(t, "task-client-request", effect.RequestID) + assert.Equal(t, "task-upstream-request", effect.UpstreamRequestID) + assert.Equal(t, "request-time-task-token", effect.TokenName) + assert.GreaterOrEqual(t, effect.UseTimeSeconds, 3) + assert.True(t, effect.IsStream) +} + func TestSweepTimedOutUnconfirmedSubmitRequiresReviewWithoutRefund(t *testing.T) { truncate(t) + resetTimedOutTaskSweepCursorForTest(t) originalTimeout := constant.TaskTimeoutMinutes constant.TaskTimeoutMinutes = 1 t.Cleanup(func() { constant.TaskTimeoutMinutes = originalTimeout }) @@ -107,6 +140,185 @@ func TestSweepTimedOutUnconfirmedSubmitRequiresReviewWithoutRefund(t *testing.T) require.Zero(t, settlementCount) } +func TestSweepTimedOutTaskWaitsForSubmissionSettlement(t *testing.T) { + truncate(t) + resetTimedOutTaskSweepCursorForTest(t) + const userID = 711 + seedUser(t, userID, 80) + originalTimeout := constant.TaskTimeoutMinutes + constant.TaskTimeoutMinutes = 1 + t.Cleanup(func() { constant.TaskTimeoutMinutes = originalTimeout }) + now := time.Now().Unix() + task := &model.Task{ + TaskID: "task_timeout_waits_for_settlement", Status: model.TaskStatusSubmitted, + UserId: userID, Quota: 20, SubmitTime: time.Now().Add(-2 * time.Minute).Unix(), + PrivateData: model.TaskPrivateData{ + UpstreamTaskID: "provider-timeout-waits-for-settlement", + BillingRequestId: "timeout-waits-for-settlement", + BillingSource: model.BillingSettlementSourceWallet, + }, + } + require.NoError(t, model.DB.Create(task).Error) + submissionInput := model.BillingSettlementInput{ + OperationKey: model.BillingRequestFinalizeOperationKey(task.PrivateData.BillingRequestId), + Source: model.BillingSettlementSourceWallet, + UserID: userID, + FundingDelta: 10, + TaskID: task.ID, + TaskQuota: 20, + TaskQuotaTarget: 30, + } + require.NoError(t, model.DB.Create(&model.BillingSettlement{ + OperationKey: submissionInput.OperationKey, + Source: submissionInput.Source, UserID: submissionInput.UserID, + FundingDelta: submissionInput.FundingDelta, TaskID: submissionInput.TaskID, + TaskQuota: submissionInput.TaskQuota, TaskQuotaTarget: submissionInput.TaskQuotaTarget, + Status: model.BillingSettlementStatusPending, + CreatedAt: now, UpdatedAt: now, Revision: 1, + }).Error) + + sweepTimedOutTasks(context.Background()) + + var reloaded model.Task + require.NoError(t, model.DB.First(&reloaded, task.ID).Error) + assert.EqualValues(t, model.TaskStatusSubmitted, reloaded.Status) + assert.Equal(t, 20, reloaded.Quota) + assert.EqualValues(t, 80, getUserQuota(t, userID)) + + applied, alreadyApplied, err := model.ApplyBillingSettlementOnce(submissionInput) + require.NoError(t, err) + assert.False(t, alreadyApplied) + assert.EqualValues(t, 10, applied) + require.NoError(t, model.DB.First(&reloaded, task.ID).Error) + assert.EqualValues(t, model.TaskStatusSubmitted, reloaded.Status) + assert.Equal(t, 30, reloaded.Quota) + assert.EqualValues(t, 70, getUserQuota(t, userID)) + + sweepTimedOutTasks(context.Background()) + + require.NoError(t, model.DB.First(&reloaded, task.ID).Error) + assert.EqualValues(t, model.TaskStatusFailure, reloaded.Status) + assert.Zero(t, reloaded.Quota) + assert.EqualValues(t, 100, getUserQuota(t, userID)) + var refund model.BillingSettlement + require.NoError(t, model.DB.Where("operation_key = ?", fmt.Sprintf("task:%d:refund", task.ID)).First(&refund).Error) + assert.Equal(t, model.BillingSettlementStatusApplied, refund.Status) + assert.EqualValues(t, -30, refund.AppliedFundingDelta) +} + +func TestSweepTimedOutTasksScansPastPendingSettlementPage(t *testing.T) { + truncate(t) + resetTimedOutTaskSweepCursorForTest(t) + originalTimeout := constant.TaskTimeoutMinutes + constant.TaskTimeoutMinutes = 1 + t.Cleanup(func() { constant.TaskTimeoutMinutes = originalTimeout }) + + now := time.Now().Unix() + pendingSubmitTime := legacyTaskRefundCutoff - 2 + tasks := make([]model.Task, 0, 101) + settlements := make([]model.BillingSettlement, 0, 100) + for i := 0; i < 100; i++ { + requestID := fmt.Sprintf("timeout-pending-page-%d", i) + tasks = append(tasks, model.Task{ + TaskID: fmt.Sprintf("task_timeout_pending_page_%d", i), + Status: model.TaskStatusSubmitted, + SubmitTime: pendingSubmitTime, + PrivateData: model.TaskPrivateData{ + BillingRequestId: requestID, + }, + }) + settlements = append(settlements, model.BillingSettlement{ + OperationKey: model.BillingRequestFinalizeOperationKey(requestID), + Source: model.BillingSettlementSourceWallet, + UserID: 9000 + i, + FundingDelta: 1, + Status: model.BillingSettlementStatusPending, + CreatedAt: now, + UpdatedAt: now, + Revision: 1, + }) + } + tasks = append(tasks, model.Task{ + TaskID: "task_timeout_after_pending_page", + Status: model.TaskStatusSubmitted, + SubmitTime: pendingSubmitTime + 1, + }) + require.NoError(t, model.DB.CreateInBatches(&tasks, 25).Error) + require.NoError(t, model.DB.CreateInBatches(&settlements, 25).Error) + + sweepTimedOutTasks(context.Background()) + + var actionable model.Task + require.NoError(t, model.DB.Where("task_id = ?", "task_timeout_after_pending_page").First(&actionable).Error) + assert.EqualValues(t, model.TaskStatusFailure, actionable.Status) + var pending model.Task + require.NoError(t, model.DB.Where("task_id = ?", "task_timeout_pending_page_0").First(&pending).Error) + assert.EqualValues(t, model.TaskStatusSubmitted, pending.Status) +} + +func TestSweepTimedOutTasksBoundsScansAndContinuesFromCursor(t *testing.T) { + truncate(t) + resetTimedOutTaskSweepCursorForTest(t) + originalTimeout := constant.TaskTimeoutMinutes + originalScanBudget := timedOutTaskScanBudget + constant.TaskTimeoutMinutes = 1 + timedOutTaskScanBudget = 2 + t.Cleanup(func() { + constant.TaskTimeoutMinutes = originalTimeout + timedOutTaskScanBudget = originalScanBudget + }) + + now := time.Now().Unix() + submitTime := legacyTaskRefundCutoff - 2 + for i := 0; i < 2; i++ { + requestID := fmt.Sprintf("bounded-timeout-pending-%d", i) + task := model.Task{ + TaskID: fmt.Sprintf("bounded_timeout_pending_%d", i), + Status: model.TaskStatusSubmitted, + SubmitTime: submitTime, + PrivateData: model.TaskPrivateData{BillingRequestId: requestID}, + } + require.NoError(t, model.DB.Create(&task).Error) + require.NoError(t, model.DB.Create(&model.BillingSettlement{ + OperationKey: model.BillingRequestFinalizeOperationKey(requestID), + Source: model.BillingSettlementSourceWallet, + UserID: 9100 + i, + FundingDelta: 1, + Status: model.BillingSettlementStatusPending, + CreatedAt: now, + UpdatedAt: now, + Revision: 1, + }).Error) + } + actionable := model.Task{ + TaskID: "bounded_timeout_actionable", + Status: model.TaskStatusSubmitted, + SubmitTime: submitTime + 1, + } + require.NoError(t, model.DB.Create(&actionable).Error) + + sweepTimedOutTasks(context.Background()) + var reloaded model.Task + require.NoError(t, model.DB.First(&reloaded, actionable.ID).Error) + assert.EqualValues(t, model.TaskStatusSubmitted, reloaded.Status) + + sweepTimedOutTasks(context.Background()) + require.NoError(t, model.DB.First(&reloaded, actionable.ID).Error) + assert.EqualValues(t, model.TaskStatusFailure, reloaded.Status) +} + +func resetTimedOutTaskSweepCursorForTest(t *testing.T) { + t.Helper() + reset := func() { + timedOutTaskSweepCursor.Lock() + timedOutTaskSweepCursor.afterSubmitTime = 0 + timedOutTaskSweepCursor.afterID = 0 + timedOutTaskSweepCursor.Unlock() + } + reset() + t.Cleanup(reset) +} + // --------------------------------------------------------------------------- // Seed helpers // --------------------------------------------------------------------------- @@ -127,6 +339,8 @@ func truncate(t *testing.T) { model.DB.Exec("DELETE FROM billing_settlements") model.DB.Exec("DELETE FROM billing_pre_consume_selections") model.DB.Exec("DELETE FROM cache_invalidation_tasks") + model.DB.Exec("DELETE FROM midjourneys") + model.DB.Exec("DELETE FROM midjourney_billing_claims") }) } @@ -595,6 +809,282 @@ func TestProcessSunoTaskResponseRefundsWhenFailureReasonSanitizesEmpty(t *testin assert.Equal(t, model.BillingSettlementStatusApplied, settlement.Status) } +func TestTaskFinalSettlementPendingBlocksTerminalPollingUntilFundingApplied(t *testing.T) { + truncate(t) + const userID = 701 + now := time.Now().Unix() + task := makeTask(userID, 0, 100, 0, BillingSourceWallet, 0) + task.PrivateData.BillingRequestId = "task-final-settlement-pending" + persistTask(t, task) + + require.NoError(t, model.DB.Create(&model.BillingSettlement{ + OperationKey: model.BillingRequestFinalizeOperationKey(task.PrivateData.BillingRequestId), + Source: model.BillingSettlementSourceWallet, + UserID: userID, + FundingDelta: 50, + Status: model.BillingSettlementStatusPending, + CreatedAt: now, + UpdatedAt: now, + Revision: 1, + }).Error) + + pending, err := taskFinalSettlementPending(task) + require.NoError(t, err) + assert.True(t, pending) + + require.NoError(t, model.DB.Model(&model.BillingSettlement{}). + Where("operation_key = ?", model.BillingRequestFinalizeOperationKey(task.PrivateData.BillingRequestId)). + Update("status", model.BillingSettlementStatusApplied).Error) + pending, err = taskFinalSettlementPending(task) + require.NoError(t, err) + assert.False(t, pending) +} + +func TestTaskFinalSettlementEffectPendingDoesNotBlockPolling(t *testing.T) { + truncate(t) + const userID = 702 + now := time.Now().Unix() + task := makeTask(userID, 0, 100, 0, BillingSourceWallet, 0) + task.PrivateData.BillingRequestId = "task-final-effect-pending" + persistTask(t, task) + + require.NoError(t, model.DB.Create(&model.BillingSettlement{ + OperationKey: model.BillingRequestFinalizeOperationKey(task.PrivateData.BillingRequestId), + Source: model.BillingSettlementSourceWallet, + UserID: userID, + FundingDelta: 0, + Status: model.BillingSettlementStatusApplied, + EffectStatus: model.BillingSettlementEffectPending, + CreatedAt: now, + UpdatedAt: now, + Revision: 1, + }).Error) + + pending, err := taskFinalSettlementPending(task) + require.NoError(t, err) + assert.False(t, pending) +} + +func TestProcessSunoTaskResponseWaitsForSubmissionSettlement(t *testing.T) { + truncate(t) + const userID = 703 + now := time.Now().Unix() + task := makeTask(userID, 0, 100, 0, BillingSourceWallet, 0) + task.Platform = constant.TaskPlatformSuno + task.PrivateData.BillingRequestId = "suno-waits-for-submission-settlement" + persistTask(t, task) + require.NoError(t, model.DB.Create(&model.BillingSettlement{ + OperationKey: model.BillingRequestFinalizeOperationKey(task.PrivateData.BillingRequestId), + Source: model.BillingSettlementSourceWallet, UserID: userID, + FundingDelta: 50, Status: model.BillingSettlementStatusPending, + CreatedAt: now, UpdatedAt: now, Revision: 1, + }).Error) + response := dto.SunoDataResponse{ + TaskID: task.TaskID, Status: string(model.TaskStatusSuccess), + FinishTime: now, Data: json.RawMessage(`{"status":"SUCCESS"}`), + } + + processSunoTaskResponse(context.Background(), task, response) + var reloaded model.Task + require.NoError(t, model.DB.First(&reloaded, task.ID).Error) + assert.EqualValues(t, model.TaskStatusInProgress, reloaded.Status) + + require.NoError(t, model.DB.Model(&model.BillingSettlement{}). + Where("operation_key = ?", model.BillingRequestFinalizeOperationKey(task.PrivateData.BillingRequestId)). + Update("status", model.BillingSettlementStatusApplied).Error) + processSunoTaskResponse(context.Background(), &reloaded, response) + require.NoError(t, model.DB.First(&reloaded, task.ID).Error) + assert.EqualValues(t, model.TaskStatusSuccess, reloaded.Status) +} + +func TestProcessSunoStoredFailureWaitsForSubmissionSettlement(t *testing.T) { + for index, settlementStatus := range []string{ + model.BillingSettlementStatusPending, + model.BillingSettlementStatusManual, + } { + t.Run(settlementStatus, func(t *testing.T) { + truncate(t) + userID := 720 + index + now := time.Now().Unix() + seedUser(t, userID, 900) + + task := makeTask(userID, 0, 100, 0, BillingSourceWallet, 0) + task.Platform = constant.TaskPlatformSuno + task.Status = model.TaskStatusFailure + task.Progress = "50%" + task.FailReason = "stored provider failure" + task.PrivateData.BillingRequestId = "suno-stored-failure-" + settlementStatus + persistTask(t, task) + require.NoError(t, model.DB.Create(&model.BillingSettlement{ + OperationKey: model.BillingRequestFinalizeOperationKey(task.PrivateData.BillingRequestId), + Source: model.BillingSettlementSourceWallet, UserID: userID, + Status: settlementStatus, CreatedAt: now, UpdatedAt: now, Revision: 1, + }).Error) + + response := dto.SunoDataResponse{TaskID: task.TaskID, Data: task.Data} + processSunoTaskResponse(context.Background(), task, response) + + var blocked model.Task + require.NoError(t, model.DB.First(&blocked, task.ID).Error) + assert.EqualValues(t, model.TaskStatusFailure, blocked.Status) + assert.Equal(t, "50%", blocked.Progress) + assert.Equal(t, 100, blocked.Quota) + assert.EqualValues(t, 900, getUserQuota(t, userID)) + var refundCount int64 + require.NoError(t, model.DB.Model(&model.BillingSettlement{}). + Where("operation_key = ?", fmt.Sprintf("task:%d:refund", task.ID)). + Count(&refundCount).Error) + assert.Zero(t, refundCount) + + require.NoError(t, model.DB.Model(&model.BillingSettlement{}). + Where("operation_key = ?", model.BillingRequestFinalizeOperationKey(task.PrivateData.BillingRequestId)). + Update("status", model.BillingSettlementStatusApplied).Error) + processSunoTaskResponse(context.Background(), &blocked, response) + + require.NoError(t, model.DB.First(&blocked, task.ID).Error) + assert.EqualValues(t, model.TaskStatusFailure, blocked.Status) + assert.Equal(t, "100%", blocked.Progress) + assert.Zero(t, blocked.Quota) + assert.EqualValues(t, 1000, getUserQuota(t, userID)) + require.NoError(t, model.DB.Model(&model.BillingSettlement{}). + Where("operation_key = ?", fmt.Sprintf("task:%d:refund", task.ID)). + Count(&refundCount).Error) + assert.EqualValues(t, 1, refundCount) + }) + } +} + +func TestMidjourneyTerminalCallbackWaitsForSubmissionSettlement(t *testing.T) { + truncate(t) + const userID, channelID = 704, 705 + now := time.Now().Unix() + legacy := &model.Midjourney{ + UserId: userID, ChannelId: channelID, MjId: "provider-mj-waits-for-settlement", + Status: "IN_PROGRESS", Progress: "50%", Code: 1, + } + require.NoError(t, model.DB.Create(legacy).Error) + billingTask := makeTask(userID, channelID, 100, 0, BillingSourceWallet, 0) + billingTask.Platform = constant.TaskPlatformMidjourney + billingTask.PrivateData.BillingRequestId = "mj-waits-for-submission-settlement" + persistTask(t, billingTask) + require.NoError(t, model.DB.Create(&model.MidjourneyBillingClaim{ + ChannelID: channelID, MjID: legacy.MjId, UserID: userID, + BillingTaskID: billingTask.ID, BillingRequestID: billingTask.PrivateData.BillingRequestId, + CreatedAt: now, + }).Error) + require.NoError(t, model.DB.Create(&model.BillingSettlement{ + OperationKey: model.BillingRequestFinalizeOperationKey(billingTask.PrivateData.BillingRequestId), + Source: model.BillingSettlementSourceWallet, UserID: userID, + FundingDelta: 50, Status: model.BillingSettlementStatusPending, + CreatedAt: now, UpdatedAt: now, Revision: 1, + }).Error) + response := dto.MidjourneyDto{ + MjId: legacy.MjId, Status: "SUCCESS", Progress: "100%", FinishTime: now * 1000, + } + + require.NoError(t, UpdateMidjourneyTaskFromResponse(context.Background(), legacy, response)) + var reloadedLegacy model.Midjourney + require.NoError(t, model.DB.First(&reloadedLegacy, legacy.Id).Error) + assert.Equal(t, "IN_PROGRESS", reloadedLegacy.Status) + var reloadedBilling model.Task + require.NoError(t, model.DB.First(&reloadedBilling, billingTask.ID).Error) + assert.EqualValues(t, model.TaskStatusInProgress, reloadedBilling.Status) + + require.NoError(t, model.DB.Model(&model.BillingSettlement{}). + Where("operation_key = ?", model.BillingRequestFinalizeOperationKey(billingTask.PrivateData.BillingRequestId)). + Update("status", model.BillingSettlementStatusApplied).Error) + require.NoError(t, UpdateMidjourneyTaskFromResponse(context.Background(), &reloadedLegacy, response)) + require.NoError(t, model.DB.First(&reloadedLegacy, legacy.Id).Error) + assert.Equal(t, "SUCCESS", reloadedLegacy.Status) + require.NoError(t, model.DB.First(&reloadedBilling, billingTask.ID).Error) + assert.EqualValues(t, model.TaskStatusSuccess, reloadedBilling.Status) +} + +func TestMidjourneyMergedFailureWaitsForSubmissionSettlement(t *testing.T) { + for index, settlementStatus := range []string{ + model.BillingSettlementStatusPending, + model.BillingSettlementStatusManual, + } { + t.Run(settlementStatus, func(t *testing.T) { + truncate(t) + userID := 706 + index + channelID := 708 + index + now := time.Now().Unix() + seedUser(t, userID, 900) + + legacy := &model.Midjourney{ + UserId: userID, ChannelId: channelID, + MjId: "provider-mj-merged-failure-" + settlementStatus, + Status: "IN_PROGRESS", Progress: "50%", FailReason: "stored provider failure", + Code: 1, Quota: 100, + } + require.NoError(t, model.DB.Create(legacy).Error) + + billingTask := makeTask(userID, channelID, 100, 0, BillingSourceWallet, 0) + billingTask.Platform = constant.TaskPlatformMidjourney + billingTask.Status = model.TaskStatusInProgress + billingTask.Progress = "50%" + billingTask.PrivateData.BillingRequestId = "mj-merged-failure-" + settlementStatus + persistTask(t, billingTask) + require.NoError(t, model.DB.Create(&model.MidjourneyBillingClaim{ + ChannelID: channelID, MjID: legacy.MjId, UserID: userID, + BillingTaskID: billingTask.ID, BillingRequestID: billingTask.PrivateData.BillingRequestId, + CreatedAt: now, + }).Error) + require.NoError(t, model.DB.Create(&model.BillingSettlement{ + OperationKey: model.BillingRequestFinalizeOperationKey(billingTask.PrivateData.BillingRequestId), + Source: model.BillingSettlementSourceWallet, UserID: userID, + Status: settlementStatus, CreatedAt: now, UpdatedAt: now, Revision: 1, + }).Error) + + response := dto.MidjourneyDto{ + MjId: legacy.MjId, Status: "IN_PROGRESS", Progress: "50%", + } + require.NoError(t, UpdateMidjourneyTaskFromResponse(context.Background(), legacy, response)) + + var blockedLegacy model.Midjourney + require.NoError(t, model.DB.First(&blockedLegacy, legacy.Id).Error) + assert.Equal(t, "IN_PROGRESS", blockedLegacy.Status) + assert.Equal(t, "50%", blockedLegacy.Progress) + assert.Equal(t, "stored provider failure", blockedLegacy.FailReason) + var blockedBilling model.Task + require.NoError(t, model.DB.First(&blockedBilling, billingTask.ID).Error) + assert.EqualValues(t, model.TaskStatusInProgress, blockedBilling.Status) + assert.Equal(t, 100, blockedBilling.Quota) + assert.EqualValues(t, 900, getUserQuota(t, userID)) + var refundCount int64 + require.NoError(t, model.DB.Model(&model.BillingSettlement{}). + Where("operation_key = ?", fmt.Sprintf("task:%d:refund", billingTask.ID)). + Count(&refundCount).Error) + assert.Zero(t, refundCount) + + require.NoError(t, model.DB.Model(&model.BillingSettlement{}). + Where("operation_key = ?", model.BillingRequestFinalizeOperationKey(billingTask.PrivateData.BillingRequestId)). + Update("status", model.BillingSettlementStatusApplied).Error) + require.NoError(t, UpdateMidjourneyTaskFromResponse(context.Background(), &blockedLegacy, response)) + + require.NoError(t, model.DB.First(&blockedLegacy, legacy.Id).Error) + assert.Equal(t, "FAILURE", blockedLegacy.Status) + assert.Equal(t, "100%", blockedLegacy.Progress) + require.NoError(t, model.DB.First(&blockedBilling, billingTask.ID).Error) + assert.EqualValues(t, model.TaskStatusFailure, blockedBilling.Status) + assert.Zero(t, blockedBilling.Quota) + assert.EqualValues(t, 1000, getUserQuota(t, userID)) + require.NoError(t, model.DB.Model(&model.BillingSettlement{}). + Where("operation_key = ?", fmt.Sprintf("task:%d:refund", billingTask.ID)). + Count(&refundCount).Error) + assert.EqualValues(t, 1, refundCount) + + require.NoError(t, UpdateMidjourneyTaskFromResponse(context.Background(), &blockedLegacy, response)) + assert.EqualValues(t, 1000, getUserQuota(t, userID)) + require.NoError(t, model.DB.Model(&model.BillingSettlement{}). + Where("operation_key = ?", fmt.Sprintf("task:%d:refund", billingTask.ID)). + Count(&refundCount).Error) + assert.EqualValues(t, 1, refundCount) + }) + } +} + // =========================================================================== // RecalculateTaskQuota tests // =========================================================================== diff --git a/service/task_polling.go b/service/task_polling.go index 603157c8..7349a71b 100644 --- a/service/task_polling.go +++ b/service/task_polling.go @@ -8,6 +8,7 @@ import ( "net/http" "sort" "strings" + "sync" "time" "github.com/MAX-API-Next/MAX-API/common" @@ -36,65 +37,113 @@ type TaskPollingAdaptor interface { var GetTaskAdaptorFunc func(platform constant.TaskPlatform) TaskPollingAdaptor const legacyTaskRefundCutoff int64 = 1740182400 // 2025-02-22 00:00:00 UTC +const timedOutTaskTransitionBudget = 100 + +var timedOutTaskScanBudget = 2000 + +var timedOutTaskSweepCursor struct { + sync.Mutex + afterSubmitTime int64 + afterID int64 +} // sweepTimedOutTasks 在主轮询之前独立清理超时任务。 -// 每次最多处理 100 条,剩余的下个周期继续处理。 +// 每次最多转换 100 条、扫描 2000 条,剩余的从游标位置在下个周期继续处理。 // 使用 per-task CAS (UpdateWithStatus) 防止覆盖被正常轮询已推进的任务。 func sweepTimedOutTasks(ctx context.Context) { if constant.TaskTimeoutMinutes <= 0 { return } cutoff := time.Now().Unix() - int64(constant.TaskTimeoutMinutes)*60 - tasks := model.GetTimedOutUnfinishedTasks(cutoff, 100) - if len(tasks) == 0 { - return - } - reason := fmt.Sprintf("任务超时(%d分钟)", constant.TaskTimeoutMinutes) legacyReason := "任务超时(旧系统遗留任务,不进行退款,请联系管理员)" now := time.Now().Unix() timedOutCount := 0 - - for _, task := range tasks { - isLegacy := task.SubmitTime > 0 && task.SubmitTime < legacyTaskRefundCutoff - isUnconfirmedSubmit := task.PrivateData.AwaitingUpstreamID - - oldStatus := task.Status - task.Status = model.TaskStatusFailure - task.Progress = "100%" - task.FinishTime = now - if isLegacy { - task.FailReason = legacyReason - } else if isUnconfirmedSubmit { - task.FailReason = "任务提交超时(无法确认上游是否已创建任务,未自动退款,请人工核对)" - } else { - task.FailReason = reason + timedOutTaskSweepCursor.Lock() + defer timedOutTaskSweepCursor.Unlock() + + remaining := timedOutTaskTransitionBudget + scanBudget := timedOutTaskScanBudget + afterSubmitTime := timedOutTaskSweepCursor.afterSubmitTime + afterID := timedOutTaskSweepCursor.afterID + for remaining > 0 && scanBudget > 0 { + queryLimit := 100 + if scanBudget < queryLimit { + queryLimit = scanBudget } - - var settlement *model.BillingSettlementInput - if !isLegacy && !isUnconfirmedSubmit { - settlement = BuildTaskRefundSettlementInput(task, reason) - } - var won bool - var err error - if settlement != nil { - won, err = task.UpdateWithStatusAndSettlement(oldStatus, *settlement) - } else { - won, err = task.UpdateWithStatus(oldStatus) + tasks, queryErr := model.GetTimedOutUnfinishedTasksAfter(cutoff, afterSubmitTime, afterID, queryLimit) + if queryErr != nil { + logger.LogError(ctx, fmt.Sprintf("sweepTimedOutTasks query error after cursor (%d, %d): %v", afterSubmitTime, afterID, queryErr)) + break } - if err != nil { - logger.LogError(ctx, fmt.Sprintf("sweepTimedOutTasks CAS update error for task %s: %v", task.TaskID, err)) - continue + if len(tasks) == 0 { + afterSubmitTime = 0 + afterID = 0 + break } - if !won { - logger.LogInfo(ctx, fmt.Sprintf("sweepTimedOutTasks: task %s already transitioned, skip", task.TaskID)) - continue + for _, task := range tasks { + if remaining <= 0 || scanBudget <= 0 { + break + } + afterSubmitTime = task.SubmitTime + afterID = task.ID + scanBudget-- + settlementPending, settlementErr := taskTerminalSettlementPending(task, true) + if settlementErr != nil { + logger.LogError(ctx, fmt.Sprintf("sweepTimedOutTasks settlement lookup error for task %s: %v", task.TaskID, settlementErr)) + continue + } + if settlementPending { + continue + } + remaining-- + isLegacy := task.SubmitTime > 0 && task.SubmitTime < legacyTaskRefundCutoff + isUnconfirmedSubmit := task.PrivateData.AwaitingUpstreamID + + oldStatus := task.Status + task.Status = model.TaskStatusFailure + task.Progress = "100%" + task.FinishTime = now + if isLegacy { + task.FailReason = legacyReason + } else if isUnconfirmedSubmit { + task.FailReason = "任务提交超时(无法确认上游是否已创建任务,未自动退款,请人工核对)" + } else { + task.FailReason = reason + } + + var settlement *model.BillingSettlementInput + if !isLegacy && !isUnconfirmedSubmit { + settlement = BuildTaskRefundSettlementInput(task, reason) + } + var won bool + var err error + if settlement != nil { + won, err = task.UpdateWithStatusAndSettlement(oldStatus, *settlement) + } else { + won, err = task.UpdateWithStatus(oldStatus) + } + if err != nil { + logger.LogError(ctx, fmt.Sprintf("sweepTimedOutTasks CAS update error for task %s: %v", task.TaskID, err)) + continue + } + if !won { + logger.LogInfo(ctx, fmt.Sprintf("sweepTimedOutTasks: task %s already transitioned, skip", task.TaskID)) + continue + } + timedOutCount++ + if settlement != nil { + applyTaskBillingSettlement(ctx, task, settlement) + } } - timedOutCount++ - if settlement != nil { - applyTaskBillingSettlement(ctx, task, settlement) + if len(tasks) < queryLimit && remaining > 0 && scanBudget > 0 { + afterSubmitTime = 0 + afterID = 0 + break } } + timedOutTaskSweepCursor.afterSubmitTime = afterSubmitTime + timedOutTaskSweepCursor.afterID = afterID if timedOutCount > 0 { logger.LogInfo(ctx, fmt.Sprintf("sweepTimedOutTasks: timed out %d tasks", timedOutCount)) @@ -227,23 +276,37 @@ func processSunoTaskResponse(ctx context.Context, task *model.Task, responseItem if !taskNeedsUpdate(task, responseItem) { return } + mergedTask := *task + mergedTask.Status = lo.If(model.TaskStatus(responseItem.Status) != "", model.TaskStatus(responseItem.Status)).Else(mergedTask.Status) + mergedTask.FailReason = lo.If(responseItem.FailReason != "", common.SanitizePersistedLogContent(responseItem.FailReason)).Else(mergedTask.FailReason) + mergedTask.SubmitTime = lo.If(responseItem.SubmitTime != 0, responseItem.SubmitTime).Else(mergedTask.SubmitTime) + mergedTask.StartTime = lo.If(responseItem.StartTime != 0, responseItem.StartTime).Else(mergedTask.StartTime) + mergedTask.FinishTime = lo.If(responseItem.FinishTime != 0, responseItem.FinishTime).Else(mergedTask.FinishTime) + isFailure := responseItem.FailReason != "" || mergedTask.Status == model.TaskStatusFailure + if isFailure { + mergedTask.Status = model.TaskStatusFailure + mergedTask.Progress = "100%" + } + if responseItem.Status == model.TaskStatusSuccess { + mergedTask.Progress = "100%" + } + mergedTask.Data = responseItem.Data + providerTerminal := mergedTask.Status == model.TaskStatusSuccess || + mergedTask.Status == model.TaskStatusFailure || mergedTask.FailReason != "" + pending, settlementErr := taskTerminalSettlementPending(task, providerTerminal) + if settlementErr != nil { + logger.LogError(ctx, fmt.Sprintf("Suno task %s settlement lookup failed: %v", task.TaskID, settlementErr)) + return + } + if pending { + return + } previousStatus := task.Status - task.Status = lo.If(model.TaskStatus(responseItem.Status) != "", model.TaskStatus(responseItem.Status)).Else(task.Status) - 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) - isFailure := responseItem.FailReason != "" || task.Status == model.TaskStatusFailure + *task = mergedTask if isFailure { logger.LogInfo(ctx, task.TaskID+" 构建失败,"+task.FailReason) - task.Status = model.TaskStatusFailure - task.Progress = "100%" - } - if responseItem.Status == model.TaskStatusSuccess { - task.Progress = "100%" } - task.Data = responseItem.Data var settlement *model.BillingSettlementInput if isFailure { @@ -435,6 +498,14 @@ func updateVideoSingleTask(ctx context.Context, adaptor TaskPollingAdaptor, ch * } } } + providerTerminal := taskResult.Status == string(model.TaskStatusSuccess) || taskResult.Status == string(model.TaskStatusFailure) + pending, settlementErr := taskTerminalSettlementPending(task, providerTerminal) + if settlementErr != nil { + return fmt.Errorf("load task submission settlement for task %s: %w", task.TaskID, settlementErr) + } + if pending { + return nil + } shouldRefund := false shouldSettle := false diff --git a/service/text_quota.go b/service/text_quota.go index a7825e0c..7f2049fd 100644 --- a/service/text_quota.go +++ b/service/text_quota.go @@ -443,15 +443,6 @@ func PostTextConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, us extraContent = append(extraContent, fmt.Sprintf("stream ended abnormally (%s), billed pre-consumed quota %d", relayInfo.StreamStatus.EndReason, fallbackQuota)) logger.LogError(ctx, fmt.Sprintf("stream ended abnormally without billable usage, billing fallback quota, userId %d, channelId %d, tokenId %d, model %s, reason %s, fallback quota %d", relayInfo.UserId, relayInfo.ChannelId, relayInfo.TokenId, summary.ModelName, relayInfo.StreamStatus.EndReason, fallbackQuota)) } - if shouldUpdateUsageStats { - model.UpdateUserUsedQuotaAndRequestCount(relayInfo.UserId, summary.Quota) - model.UpdateChannelUsedQuota(relayInfo.ChannelId, summary.Quota) - } - - if err := SettleBilling(ctx, relayInfo, summary.Quota); err != nil { - logger.LogError(ctx, "error settling billing: "+err.Error()) - } - logModel := summary.ModelName if strings.HasPrefix(logModel, "gpt-4-gizmo") { logModel = "gpt-4-gizmo-*" @@ -539,7 +530,7 @@ func PostTextConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, us } attachQuotaSaturation(ctx, relayInfo, other) - model.RecordConsumeLog(ctx, relayInfo.UserId, model.RecordConsumeLogParams{ + settleAndRecordConsume(ctx, relayInfo, shouldUpdateUsageStats, model.RecordConsumeLogParams{ ChannelId: relayInfo.ChannelId, PromptTokens: summary.PromptTokens, CompletionTokens: summary.CompletionTokens, diff --git a/service/text_quota_test.go b/service/text_quota_test.go index c247e07b..1586d6c6 100644 --- a/service/text_quota_test.go +++ b/service/text_quota_test.go @@ -20,6 +20,7 @@ import ( "github.com/gin-gonic/gin" "github.com/shopspring/decimal" + "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) @@ -243,6 +244,197 @@ func TestPostTextConsumeQuotaUsesFallbackForNilUsageWithEstimate(t *testing.T) { require.Equal(t, 20, log.PromptTokens) } +func TestSettleAndRecordConsumeLeavesUsageProjectionPendingWhenFinalChargeExceedsWallet(t *testing.T) { + truncate(t) + const userID, tokenID, channelID = 121, 122, 123 + const preConsumedQuota, actualQuota = 100, 200 + + seedUser(t, userID, preConsumedQuota) + seedToken(t, tokenID, userID, "final-charge-insufficient-token", preConsumedQuota) + seedChannel(t, channelID) + + ctx, _ := gin.CreateTestContext(httptest.NewRecorder()) + ctx.Set("token_quota", int64(preConsumedQuota)) + info := &relaycommon.RelayInfo{ + RequestId: "final-charge-insufficient-request", + UserId: userID, + TokenId: tokenID, + TokenKey: "final-charge-insufficient-token", + OriginModelName: "final-charge-insufficient-model", + UsingGroup: "default", + ForcePreConsume: true, + ChannelMeta: &relaycommon.ChannelMeta{ChannelId: channelID}, + } + + session, apiErr := NewBillingSession(ctx, info, preConsumedQuota) + require.Nil(t, apiErr) + require.NotNil(t, session) + info.Billing = session + + settleAndRecordConsume(ctx, info, true, model.RecordConsumeLogParams{ + ChannelId: channelID, + ModelName: "final-charge-insufficient-model", + TokenId: tokenID, + Group: "default", + Quota: actualQuota, + Content: "final charge exceeds remaining wallet", + }) + + var user model.User + require.NoError(t, model.DB.Select("used_quota", "request_count").First(&user, userID).Error) + assert.Zero(t, user.UsedQuota) + assert.Zero(t, user.RequestCount) + var channel model.Channel + require.NoError(t, model.DB.Select("used_quota").First(&channel, channelID).Error) + assert.Zero(t, channel.UsedQuota) + assert.Zero(t, countLogs(t)) + + var settlement model.BillingSettlement + require.NoError(t, model.DB.Where("operation_key = ?", model.BillingRequestFinalizeOperationKey("final-charge-insufficient-request")).First(&settlement).Error) + assert.Equal(t, model.BillingSettlementStatusManual, settlement.Status) + assert.Contains(t, settlement.LastError, "user quota is not enough") +} + +func TestSettleAndRecordConsumeProjectsExplicitZeroFinalQuota(t *testing.T) { + truncate(t) + const userID, tokenID, channelID = 124, 125, 126 + const initialQuota, preConsumedQuota = 1_000, 100 + + seedUser(t, userID, initialQuota) + seedToken(t, tokenID, userID, "zero-final-quota-token", initialQuota) + seedChannel(t, channelID) + + ctx, _ := gin.CreateTestContext(httptest.NewRecorder()) + ctx.Set(common.RequestIdKey, "zero-final-quota-request") + ctx.Set("token_quota", int64(initialQuota)) + info := &relaycommon.RelayInfo{ + RequestId: "zero-final-quota-request", + UserId: userID, + TokenId: tokenID, + TokenKey: "zero-final-quota-token", + OriginModelName: "zero-final-quota-model", + UsingGroup: "default", + ForcePreConsume: true, + ChannelMeta: &relaycommon.ChannelMeta{ChannelId: channelID}, + } + + session, apiErr := NewBillingSession(ctx, info, preConsumedQuota) + require.Nil(t, apiErr) + require.NotNil(t, session) + info.Billing = session + + settleAndRecordConsume(ctx, info, true, model.RecordConsumeLogParams{ + ChannelId: channelID, + PromptTokens: 1, + ModelName: "zero-final-quota-model", + TokenId: tokenID, + Group: "default", + Quota: 0, + Content: "billable usage rounded to zero quota", + }) + + assert.EqualValues(t, initialQuota, getUserQuota(t, userID)) + assert.EqualValues(t, initialQuota, getTokenRemainQuota(t, tokenID)) + var user model.User + require.NoError(t, model.DB.Select("used_quota", "request_count").First(&user, userID).Error) + assert.Zero(t, user.UsedQuota) + assert.Equal(t, 1, user.RequestCount) + var channel model.Channel + require.NoError(t, model.DB.Select("used_quota").First(&channel, channelID).Error) + assert.Zero(t, channel.UsedQuota) + assert.Equal(t, int64(1), countLogs(t)) + log := getLastLog(t) + assert.Zero(t, log.Quota) + assert.Equal(t, 1, log.PromptTokens) + assert.Equal(t, "zero-final-quota-request", log.RequestId) +} + +func TestPostTextConsumeQuotaDoesNotProjectUsageWhenFinalizeSettlementFails(t *testing.T) { + truncate(t) + gin.SetMode(gin.TestMode) + w := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(w) + ctx.Set(common.RequestIdKey, "failed-finalize-text-request") + + const userID, tokenID, channelID = 103, 203, 303 + seedUser(t, userID, 15) + seedToken(t, tokenID, userID, "failed-finalize-text-token", 100) + seedChannel(t, channelID) + + relayInfo := &relaycommon.RelayInfo{ + RequestId: "failed-finalize-text-request", + UserId: userID, + TokenId: tokenID, + TokenKey: "failed-finalize-text-token", + OriginModelName: "test-model", + StartTime: time.Now(), + UserSetting: dto.UserSetting{BillingPreference: "wallet_only"}, + FinalRequestRelayFormat: types.RelayFormatOpenAI, + UsingGroup: "default", + ChannelMeta: &relaycommon.ChannelMeta{ChannelId: channelID}, + PriceData: types.PriceData{ + ModelRatio: 1, + CompletionRatio: 1, + GroupRatioInfo: types.GroupRatioInfo{GroupRatio: 1}, + }, + } + session, apiErr := NewBillingSession(ctx, relayInfo, 10) + require.Nil(t, apiErr) + relayInfo.Billing = session + + PostTextConsumeQuota(ctx, relayInfo, &dto.Usage{ + PromptTokens: 10, + CompletionTokens: 10, + TotalTokens: 20, + }, nil) + + assert.EqualValues(t, 5, getUserQuota(t, userID)) + assert.EqualValues(t, 90, getTokenRemainQuota(t, tokenID)) + var user model.User + require.NoError(t, model.DB.Select("used_quota", "request_count").First(&user, userID).Error) + assert.Zero(t, user.UsedQuota) + assert.Zero(t, user.RequestCount) + var channel model.Channel + require.NoError(t, model.DB.Select("used_quota").First(&channel, channelID).Error) + assert.Zero(t, channel.UsedQuota) + assert.Zero(t, countLogs(t)) + + var settlement model.BillingSettlement + require.NoError(t, model.DB.Where("operation_key = ?", model.BillingRequestFinalizeOperationKey("failed-finalize-text-request")).First(&settlement).Error) + assert.Equal(t, model.BillingSettlementStatusManual, settlement.Status) + assert.NotEmpty(t, settlement.EffectPayload) + assert.Empty(t, settlement.EffectStatus) + + // Manual review/reconciliation makes the missing funding available and + // requeues the durable settlement. Its effect must then project exactly once. + require.NoError(t, model.DB.Model(&model.User{}).Where("id = ?", userID).Update("quota", 15).Error) + require.NoError(t, model.DB.Model(&model.BillingSettlement{}). + Where("id = ?", settlement.ID). + Updates(map[string]interface{}{"status": model.BillingSettlementStatusPending, "next_attempt": time.Now().Unix()}).Error) + model.ProcessPendingBillingSettlementsOnce() + + assert.EqualValues(t, 5, getUserQuota(t, userID)) + var projectedUser model.User + require.NoError(t, model.DB.Select("used_quota", "request_count").First(&projectedUser, userID).Error) + assert.EqualValues(t, 20, projectedUser.UsedQuota) + assert.Equal(t, 1, projectedUser.RequestCount) + var projectedChannel model.Channel + require.NoError(t, model.DB.Select("used_quota").First(&projectedChannel, channelID).Error) + assert.EqualValues(t, 20, projectedChannel.UsedQuota) + assert.Equal(t, int64(1), countLogs(t)) + log := getLastLog(t) + assert.Equal(t, 20, log.Quota) + assert.Equal(t, 10, log.PromptTokens) + assert.Equal(t, 10, log.CompletionTokens) + assert.Equal(t, "failed-finalize-text-request", log.RequestId) + + model.ProcessPendingBillingSettlementsOnce() + var replayedUser model.User + require.NoError(t, model.DB.Select("used_quota").First(&replayedUser, userID).Error) + assert.EqualValues(t, 20, replayedUser.UsedQuota) + assert.Equal(t, int64(1), countLogs(t)) +} + func TestCalculateTextQuotaSummaryUsesSplitClaudeCacheCreationRatios(t *testing.T) { gin.SetMode(gin.TestMode) w := httptest.NewRecorder() diff --git a/service/tiered_settle.go b/service/tiered_settle.go index cf36286e..c3c788a1 100644 --- a/service/tiered_settle.go +++ b/service/tiered_settle.go @@ -8,7 +8,7 @@ import ( "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" - relayconstant "github.com/MAX-API-Next/MAX-API/relay/constant" + pricehelper "github.com/MAX-API-Next/MAX-API/relay/helper" "github.com/MAX-API-Next/MAX-API/types" "github.com/gin-gonic/gin" ) @@ -161,9 +161,16 @@ func PrepareTieredBillingForSelectedGroup(c *gin.Context, relayInfo *relaycommon types.ErrOptionWithSkipRetry(), ) } - if relayInfo.RelayMode == relayconstant.RelayModeAlphaSearch { - relayInfo.PriceData.QuotaToPreConsume = targetQuota + targetQuota, quotaErr = pricehelper.ApplyPreConsumedQuotaFloor(targetQuota, snap.GroupRatio > 0) + if quotaErr != nil { + return types.NewErrorWithStatusCode( + quotaErr, + types.ErrorCodeModelPriceError, + http.StatusBadRequest, + types.ErrOptionWithSkipRetry(), + ) } + relayInfo.PriceData.QuotaToPreConsume = targetQuota if snap.GroupRatio == 0 { return nil } diff --git a/service/tiered_settle_test.go b/service/tiered_settle_test.go index 8e3da08c..8084f443 100644 --- a/service/tiered_settle_test.go +++ b/service/tiered_settle_test.go @@ -7,6 +7,7 @@ import ( "sync" "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" @@ -425,6 +426,9 @@ func TestPrepareTieredBillingForSelectedGroupFreeToPaidInitializesBilling(t *tes seedToken(t, tokenID, userID, tokenKey, initialQuota) ctx, _ := gin.CreateTestContext(nil) + // This fixture covers generic tiered free-to-paid group repricing. Alpha + // Search adds its surcharge before the floor and is covered separately in + // alpha_search_billing_test.go, so it intentionally keeps the normal mode. relayInfo := &relaycommon.RelayInfo{ RequestId: "tiered-free-to-paid-request", UserId: userID, @@ -461,6 +465,55 @@ func TestPrepareTieredBillingForSelectedGroupFreeToPaidInitializesBilling(t *tes require.Equal(t, initialQuota-expectedQuota, getTokenRemainQuota(t, tokenID)) } +func TestPrepareTieredBillingForSelectedGroupFreeToPaidAppliesPreConsumedQuota(t *testing.T) { + truncate(t) + const ( + userID = 923 + tokenID = 924 + tokenKey = "tiered-free-to-paid-preconsume-token" + initialQuota = 200_000 + configured = 150_000 + ) + originalPreConsumedQuota := common.PreConsumedQuota + common.PreConsumedQuota = configured + t.Cleanup(func() { common.PreConsumedQuota = originalPreConsumedQuota }) + seedUser(t, userID, initialQuota) + seedToken(t, tokenID, userID, tokenKey, initialQuota) + + ctx, _ := gin.CreateTestContext(nil) + relayInfo := &relaycommon.RelayInfo{ + RequestId: "tiered-free-to-paid-preconsume-request", + UserId: userID, + TokenId: tokenID, + TokenKey: tokenKey, + OriginModelName: "tiered-free-to-paid-preconsume-model", + UsingGroup: "paid", + UserSetting: dto.UserSetting{ + BillingPreference: "wallet_only", + }, + PriceData: types.PriceData{ + FreeModel: true, + GroupRatioInfo: types.GroupRatioInfo{GroupRatio: 0.20}, + }, + TieredBillingSnapshot: &billingexpr.BillingSnapshot{ + BillingMode: "tiered_expr", + ExprString: `tier("base", p)`, + ExprHash: billingexpr.ExprHashString(`tier("base", p)`), + GroupRatio: 0, + EstimatedQuotaBeforeGroup: 500_000, + EstimatedQuotaAfterGroup: 0, + QuotaPerUnit: testQuotaPerUnit, + }, + } + + require.Nil(t, PrepareTieredBillingForSelectedGroup(ctx, relayInfo)) + require.Equal(t, configured, relayInfo.FinalPreConsumedQuota) + require.Equal(t, configured, relayInfo.Billing.GetPreConsumedQuota()) + require.Equal(t, 100_000, relayInfo.TieredBillingSnapshot.EstimatedQuotaAfterGroup) + require.EqualValues(t, initialQuota-configured, getUserQuota(t, userID)) + require.Equal(t, initialQuota-configured, getTokenRemainQuota(t, tokenID)) +} + func TestPrepareTieredBillingForSelectedGroupFreeToPaidRequiresContext(t *testing.T) { const expr = `tier("base", p)` relayInfo := makeRelayInfo(expr, 0, 1_000_000, 0) diff --git a/service/violation_fee.go b/service/violation_fee.go index caa232e1..e4aac18c 100644 --- a/service/violation_fee.go +++ b/service/violation_fee.go @@ -155,31 +155,8 @@ func chargeViolationFeeIfNeeded(ctx *gin.Context, relayInfo *relaycommon.RelayIn "upstream_error_code": fmt.Sprintf("%v", oai.Code), "violation_fee_marker": CSAMViolationMarker, } - effect := &model.BillingSettlementEffect{ - LogType: model.LogTypeConsume, Content: "Violation fee charged", - ChannelID: relayInfo.ChannelId, ModelName: relayInfo.OriginModelName, - TokenID: relayInfo.TokenId, Group: relayInfo.UsingGroup, Other: other, - UpdateUsage: true, Quota: int64(feeQuota), - } - - if settler, ok := relayInfo.Billing.(interface { - SettleWithEffect(int, *model.BillingSettlementEffect) error - }); ok { - if err := settler.SettleWithEffect(feeQuota, effect); err != nil { - logger.LogError(ctx, fmt.Sprintf("violation fee settlement remains pending/manual: %s", err.Error())) - return violationFeeChargeResult{applicable: true} - } - return violationFeeChargeResult{applicable: true, settled: true} - } - - if err := SettleBilling(ctx, relayInfo, feeQuota); err != nil { - logger.LogError(ctx, fmt.Sprintf("violation fee settlement remains pending/manual: %s", err.Error())) - return violationFeeChargeResult{applicable: true} - } - model.UpdateUserUsedQuotaAndRequestCount(relayInfo.UserId, feeQuota) - model.UpdateChannelUsedQuota(relayInfo.ChannelId, feeQuota) - - model.RecordConsumeLog(ctx, relayInfo.UserId, model.RecordConsumeLogParams{ + requestID, upstreamRequestID := billingEffectRequestIDs(ctx, relayInfo) + logParams := model.RecordConsumeLogParams{ ChannelId: relayInfo.ChannelId, ModelName: relayInfo.OriginModelName, TokenName: tokenName, @@ -190,7 +167,21 @@ func chargeViolationFeeIfNeeded(ctx *gin.Context, relayInfo *relaycommon.RelayIn IsStream: relayInfo.IsStream, Group: relayInfo.UsingGroup, Other: other, - }) + } + effect := newConsumeBillingSettlementEffect(relayInfo, logParams, requestID, upstreamRequestID, true) + + effectHandled, err := SettleBillingWithEffect(ctx, relayInfo, feeQuota, effect) + if err != nil { + logger.LogError(ctx, fmt.Sprintf("violation fee settlement remains pending/manual: %s", err.Error())) + return violationFeeChargeResult{applicable: true} + } + if effectHandled { + return violationFeeChargeResult{applicable: true, settled: true} + } + model.UpdateUserUsedQuotaAndRequestCount(relayInfo.UserId, feeQuota) + model.UpdateChannelUsedQuota(relayInfo.ChannelId, feeQuota) + + model.RecordConsumeLog(ctx, relayInfo.UserId, logParams) return violationFeeChargeResult{applicable: true, settled: true} } diff --git a/types/price_data.go b/types/price_data.go index fe5d63aa..065369ab 100644 --- a/types/price_data.go +++ b/types/price_data.go @@ -5,6 +5,10 @@ import ( "math" ) +// ErrorCodeBillingReconciliationPending identifies admission blocks caused by +// a previous positive final settlement that still requires reconciliation. +const ErrorCodeBillingReconciliationPending ErrorCode = "billing_reconciliation_pending" + type GroupRatioInfo struct { GroupRatio float64 GroupSpecialRatio float64 @@ -26,7 +30,7 @@ type PriceData struct { OtherRatios map[string]float64 UsePrice bool Quota int // 按次计费的最终额度(MJ / Task) - QuotaToPreConsume int // 按量计费的预消耗额度 + QuotaToPreConsume int // 预消耗额度(真实估算额度应用预扣下限后) GroupRatioInfo GroupRatioInfo } diff --git a/web/default/src/components/layout/components/authenticated-layout.tsx b/web/default/src/components/layout/components/authenticated-layout.tsx index 0188d0cb..fe828943 100644 --- a/web/default/src/components/layout/components/authenticated-layout.tsx +++ b/web/default/src/components/layout/components/authenticated-layout.tsx @@ -16,6 +16,7 @@ along with this program. If not, see . For commercial licensing, please contact https://github.com/MAX-API-Next/MAX-API/issues */ +import type { ReactElement, ReactNode } from 'react' import { getCookie } from '@/lib/cookies' import { cn } from '@/lib/utils' import { LayoutProvider } from '@/context/layout-provider' @@ -23,37 +24,42 @@ import { SearchProvider } from '@/context/search-provider' import { SidebarInset, SidebarProvider } from '@/components/ui/sidebar' import { AnimatedOutlet } from '@/components/page-transition' import { SkipToMain } from '@/components/skip-to-main' +import { SecureVerificationProvider } from '@/features/auth/secure-verification' import { AppHeader } from './app-header' import { AppSidebar } from './app-sidebar' type AuthenticatedLayoutProps = { - children?: React.ReactNode + children?: ReactNode } -export function AuthenticatedLayout(props: AuthenticatedLayoutProps) { +export function AuthenticatedLayout( + props: AuthenticatedLayoutProps +): ReactElement { const defaultOpen = getCookie('sidebar_state') !== 'false' return ( - - - - - -
- - - {props.children ?? } - -
-
-
-
+ + + + + + +
+ + + {props.children ?? } + +
+
+
+
+
) } diff --git a/web/default/src/components/layout/components/chat-presets-item.tsx b/web/default/src/components/layout/components/chat-presets-item.tsx index 531189f8..4f4b3f28 100644 --- a/web/default/src/components/layout/components/chat-presets-item.tsx +++ b/web/default/src/components/layout/components/chat-presets-item.tsx @@ -21,6 +21,8 @@ import { Link, useLocation } from '@tanstack/react-router' import { ExternalLink, Loader2, ChevronRight } from 'lucide-react' import { useTranslation } from 'react-i18next' import { toast } from 'sonner' +import { handleServerError } from '@/lib/handle-server-error' +import { wasSecureVerificationErrorReported } from '@/lib/secure-verification' import { Collapsible, CollapsibleContent, @@ -40,6 +42,7 @@ import { SidebarMenuSubItem, useSidebar, } from '@/components/ui/sidebar' +import { useApiTokenVerification } from '@/features/auth/secure-verification' import { fetchActiveChatKey } from '@/features/chat/hooks/use-active-chat-key' import { useChatPresets } from '@/features/chat/hooks/use-chat-presets' import { @@ -156,6 +159,7 @@ export function ChatPresetsItem({ item }: { item: NavChatPresets }) { const { t } = useTranslation() const { chatPresets, serverAddress } = useChatPresets() const { state, isMobile, setOpenMobile } = useSidebar() + const withApiTokenVerification = useApiTokenVerification() const href = useLocation({ select: (location) => location.href }) const [loadingPresetId, setLoadingPresetId] = useState(null) const loadingPresetIdRef = useRef(null) @@ -171,25 +175,50 @@ export function ChatPresetsItem({ item }: { item: NavChatPresets }) { const needsKey = chatLinkRequiresApiKey(preset.url) let activeKey: string | undefined + let chatWindow: Window | null = null if (needsKey && loadingPresetIdRef.current) { toast.info(t('Preparing your chat link, please try again in a moment.')) return } + if (typeof window !== 'undefined') { + try { + chatWindow = window.open('about:blank', '_blank') + if (chatWindow) { + chatWindow.opener = null + if (needsKey) { + try { + window.focus() + } catch { + // The placeholder remains usable even when focus is restricted. + } + } + } + } catch { + chatWindow = null + } + } + if (needsKey) { loadingPresetIdRef.current = preset.id setLoadingPresetId(preset.id) try { - activeKey = await fetchActiveChatKey() + activeKey = + (await withApiTokenVerification(fetchActiveChatKey)) || undefined + if (!activeKey) { + chatWindow?.close() + return + } } catch (error) { - const message = - error instanceof Error - ? error.message - : t( - 'Unable to prepare chat link. Please ensure you have an enabled API key.' - ) - toast.error(message) + chatWindow?.close() + if (!wasSecureVerificationErrorReported(error)) { + handleServerError(error, { + fallback: t( + 'Unable to prepare chat link. Please ensure you have an enabled API key.' + ), + }) + } return } finally { loadingPresetIdRef.current = null @@ -204,16 +233,29 @@ export function ChatPresetsItem({ item }: { item: NavChatPresets }) { }) if (!url) { + chatWindow?.close() toast.error(t('Invalid chat link. Please contact the administrator.')) return } - if (typeof window === 'undefined') return + if (typeof window === 'undefined') { + chatWindow?.close() + return + } - window.open(url, '_blank', 'noopener,noreferrer') + try { + if (chatWindow) { + chatWindow.location.replace(url) + } else { + window.location.assign(url) + } + } catch { + chatWindow?.close() + window.location.assign(url) + } setOpenMobile(false) }, - [serverAddress, setOpenMobile, t] + [serverAddress, setOpenMobile, t, withApiTokenVerification] ) const normalizedHref = normalizeHref(href) diff --git a/web/default/src/features/auth/passkey/api.ts b/web/default/src/features/auth/passkey/api.ts index 9abad630..a2a9cf2e 100644 --- a/web/default/src/features/auth/passkey/api.ts +++ b/web/default/src/features/auth/passkey/api.ts @@ -16,7 +16,7 @@ along with this program. If not, see . For commercial licensing, please contact https://github.com/MAX-API-Next/MAX-API/issues */ -import { api } from '@/lib/api' +import { api, sensitiveActionConfig } from '@/lib/api' import type { ApiResponse, PasskeyOptionsPayload, PasskeyStatus } from './types' export async function getPasskeyStatus(): Promise> { @@ -28,7 +28,9 @@ export async function beginPasskeyRegistration(): Promise< ApiResponse > { const res = await api.post>( - '/api/user/passkey/register/begin' + '/api/user/passkey/register/begin', + undefined, + sensitiveActionConfig ) return res.data } @@ -38,7 +40,8 @@ export async function finishPasskeyRegistration( ): Promise { const res = await api.post( '/api/user/passkey/register/finish', - payload + payload, + sensitiveActionConfig ) return res.data } diff --git a/web/default/src/features/auth/passkey/hooks/use-passkey-management.ts b/web/default/src/features/auth/passkey/hooks/use-passkey-management.ts index 3b233a90..2cb3f65c 100644 --- a/web/default/src/features/auth/passkey/hooks/use-passkey-management.ts +++ b/web/default/src/features/auth/passkey/hooks/use-passkey-management.ts @@ -19,12 +19,14 @@ For commercial licensing, please contact https://github.com/MAX-API-Next/MAX-API import { useCallback, useEffect, useMemo, useState } from 'react' import i18next from 'i18next' import { toast } from 'sonner' +import { handleServerError } from '@/lib/handle-server-error' import { buildRegistrationResult, createCredential, isPasskeySupported as detectPasskeySupport, prepareCredentialCreationOptions, } from '@/lib/passkey' +import { isVerificationRequiredError } from '@/lib/secure-verification' import { beginPasskeyRegistration, deletePasskey, @@ -59,10 +61,10 @@ export function usePasskeyManagement( setStatus(null) toast.error(res.message || i18next.t('Failed to load Passkey status')) } - } catch (error) { - // eslint-disable-next-line no-console - console.error('[Passkey] Failed to fetch status', error) - toast.error(i18next.t('Failed to load Passkey status')) + } catch (error: unknown) { + handleServerError(error, { + fallback: i18next.t('Failed to load Passkey status'), + }) setStatus(null) } finally { setLoading(false) @@ -70,7 +72,8 @@ export function usePasskeyManagement( }, [onStatusChange]) useEffect(() => { - fetchStatus() + const timer = window.setTimeout(() => void fetchStatus(), 0) + return () => window.clearTimeout(timer) }, [fetchStatus]) useEffect(() => { @@ -79,7 +82,7 @@ export function usePasskeyManagement( .catch(() => setSupported(false)) }, []) - const register = useCallback(async () => { + const register = useCallback(async (): Promise => { if (!supported) { toast.error(i18next.t('This device does not support Passkey')) return false @@ -130,24 +133,23 @@ export function usePasskeyManagement( await fetchStatus() return true } catch (error: unknown) { + if (isVerificationRequiredError(error)) { + throw error + } if (error instanceof DOMException && error.name === 'NotAllowedError') { toast.info(i18next.t('Passkey registration was cancelled')) return false } - // eslint-disable-next-line no-console - console.error('[Passkey] Registration error', error) - toast.error( - error instanceof Error - ? error.message - : i18next.t('Failed to register Passkey') - ) + handleServerError(error, { + fallback: i18next.t('Failed to register Passkey'), + }) return false } finally { setRegistering(false) } }, [supported, fetchStatus]) - const remove = useCallback(async () => { + const remove = useCallback(async (): Promise => { setRemoving(true) try { const res = await deletePasskey() @@ -160,9 +162,12 @@ export function usePasskeyManagement( await fetchStatus() return true } catch (error) { - // eslint-disable-next-line no-console - console.error('[Passkey] Removal error', error) - toast.error(i18next.t('Failed to remove Passkey')) + if (isVerificationRequiredError(error)) { + throw error + } + handleServerError(error, { + fallback: i18next.t('Failed to remove Passkey'), + }) return false } finally { setRemoving(false) diff --git a/web/default/src/features/auth/reset-password-confirm/index.tsx b/web/default/src/features/auth/reset-password-confirm/index.tsx index 38cc4f0f..1f09edba 100644 --- a/web/default/src/features/auth/reset-password-confirm/index.tsx +++ b/web/default/src/features/auth/reset-password-confirm/index.tsx @@ -16,7 +16,7 @@ along with this program. If not, see . For commercial licensing, please contact https://github.com/MAX-API-Next/MAX-API/issues */ -import { useState } from 'react' +import { useState, type ReactElement } from 'react' import { useNavigate } from '@tanstack/react-router' import { CheckIcon, CopyIcon } from 'lucide-react' import { useTranslation } from 'react-i18next' @@ -37,13 +37,21 @@ export type ResetPasswordSearchParams = { type ResetPasswordConfirmProps = ResetPasswordSearchParams -export function ResetPasswordConfirm({ - email, - token, -}: ResetPasswordConfirmProps) { +interface ResetPasswordResponse { + success: boolean + data: string + message?: string + api_tokens_revoked?: boolean +} + +export function ResetPasswordConfirm( + props: ResetPasswordConfirmProps +): ReactElement { const { t } = useTranslation() const navigate = useNavigate() const [newPassword, setNewPassword] = useState('') + const [apiTokensRevoked, setApiTokensRevoked] = useState(false) + const [autoCopied, setAutoCopied] = useState(false) const [loading, setLoading] = useState(false) const [copied, setCopied] = useState(false) const { @@ -52,10 +60,10 @@ export function ResetPasswordConfirm({ start: startCountdown, } = useCountdown({ initialSeconds: 30 }) - const isValidResetLink = Boolean(email && token) + const isValidResetLink = Boolean(props.email && props.token) - async function handleSubmit() { - if (!isValidResetLink || !email || !token) { + async function handleSubmit(): Promise { + if (!isValidResetLink || !props.email || !props.token) { toast.error(t('Invalid reset link, please request a new password reset')) return } @@ -63,14 +71,18 @@ export function ResetPasswordConfirm({ startCountdown() setLoading(true) try { - const res = await api.post('/api/user/reset', { email, token }, { - skipBusinessError: true, - } as Record) + const res = await api.post( + '/api/user/reset', + { email: props.email, token: props.token }, + { skipBusinessError: true } + ) if (res?.data?.success) { const password = res.data.data setNewPassword(password) + setApiTokensRevoked(res.data.api_tokens_revoked === true) const copySuccess = await copyToClipboard(password) + setAutoCopied(copySuccess) if (copySuccess) { toast.success( t('Password reset and copied to clipboard: {{password}}', { @@ -80,6 +92,8 @@ export function ResetPasswordConfirm({ } else { toast.success(t('Password reset: {{password}}', { password })) } + } else { + toast.error(res?.data?.message || t('Reset failed')) } } catch { // Errors handled by global interceptor @@ -88,7 +102,7 @@ export function ResetPasswordConfirm({ } } - async function handleCopy() { + async function handleCopy(): Promise { if (!newPassword) return const copySuccess = await copyToClipboard(newPassword) @@ -131,38 +145,50 @@ export function ResetPasswordConfirm({ {newPassword && ( -
- -
- - +
+
+ +
+ + +
+ {(autoCopied || copied) && ( +

+ {t('Password has been copied to clipboard')} +

+ )}
-

- {t('Password has been copied to clipboard')} -

+ {apiTokensRevoked && ( + + + {t('auth.resetPasswordConfirm.apiTokensRevoked')} + + + )}
)} diff --git a/web/default/src/features/auth/secure-verification/api.ts b/web/default/src/features/auth/secure-verification/api.ts index 4d5415e7..3c51e5ef 100644 --- a/web/default/src/features/auth/secure-verification/api.ts +++ b/web/default/src/features/auth/secure-verification/api.ts @@ -50,30 +50,28 @@ export async function checkVerificationMethods( const [methodsResponse, passkeySupported] = await Promise.all([ api.get('/api/verify/methods', { params: scope ? { scope } : undefined, + skipBusinessError: true, + skipErrorHandler: true, }), detectPasskeySupport(), ]) const methods = methodsResponse.data?.data + if (!methodsResponse.data?.success || !methods) { + throw new Error( + methodsResponse.data?.message || i18next.t('Verification unavailable') + ) + } return { - has2FA: - Boolean(methodsResponse.data?.success) && Boolean(methods?.has_2fa), - hasPasskey: - Boolean(methodsResponse.data?.success) && Boolean(methods?.has_passkey), - hasPassword: - Boolean(methodsResponse.data?.success) && - Boolean(methods?.has_password), + has2FA: Boolean(methods.has_2fa), + hasPasskey: Boolean(methods.has_passkey), + hasPassword: Boolean(methods.has_password), passkeySupported, } } catch (error) { // eslint-disable-next-line no-console console.error('[Secure Verification] Failed to check methods', error) - return { - has2FA: false, - hasPasskey: false, - hasPassword: false, - passkeySupported: false, - } + throw error } } diff --git a/web/default/src/features/auth/secure-verification/components/secure-verification-dialog.tsx b/web/default/src/features/auth/secure-verification/components/secure-verification-dialog.tsx index 034167f7..c9c38cf6 100644 --- a/web/default/src/features/auth/secure-verification/components/secure-verification-dialog.tsx +++ b/web/default/src/features/auth/secure-verification/components/secure-verification-dialog.tsx @@ -72,14 +72,14 @@ export function SecureVerificationDialog({ const title = state.title ?? (availableTabs.length - ? 'Additional verification required' - : 'Verification unavailable') + ? t('Additional verification required') + : t('Verification unavailable')) const description = state.description ?? (availableTabs.length - ? 'Confirm your identity before accessing this sensitive action.' - : 'Enable Two-factor Authentication or Passkey in your profile settings to continue.') + ? t('Confirm your identity before accessing this sensitive action.') + : t('Add a password or sign in again with a linked account to continue.')) const handleVerify = () => { if (!activeMethod) return @@ -117,7 +117,7 @@ export function SecureVerificationDialog({

{t( - 'Enable Two-factor Authentication or Passkey in your profile to unlock sensitive operations.' + 'Add a password or sign in again with a linked account to unlock sensitive operations.' )}

diff --git a/web/default/src/features/auth/secure-verification/components/secure-verification-provider.tsx b/web/default/src/features/auth/secure-verification/components/secure-verification-provider.tsx new file mode 100644 index 00000000..21275666 --- /dev/null +++ b/web/default/src/features/auth/secure-verification/components/secure-verification-provider.tsx @@ -0,0 +1,91 @@ +/* +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 { + createContext, + useCallback, + useContext, + useMemo, + type ReactElement, + type ReactNode, +} from 'react' +import { useSecureVerification } from '../hooks/use-secure-verification' +import type { StartVerificationOptions, VerificationMethod } from '../types' +import { SecureVerificationDialog } from './secure-verification-dialog' + +type SecureVerificationProviderProps = { + children: ReactNode +} + +interface SecureVerificationContextValue { + withVerification: ( + apiCall: () => Promise, + config?: StartVerificationOptions + ) => Promise +} + +export const SecureVerificationContext = + createContext(null) + +export function useSecureVerificationGate(): SecureVerificationContextValue { + const context = useContext(SecureVerificationContext) + if (!context) { + throw new Error( + 'useSecureVerificationGate must be used within SecureVerificationProvider' + ) + } + return context +} + +export function SecureVerificationProvider( + props: SecureVerificationProviderProps +): ReactElement { + const secureVerification = useSecureVerification() + + const handleVerification = useCallback( + async (method: VerificationMethod, code?: string): Promise => { + try { + await secureVerification.executeVerification(method, code) + } catch { + // The shared verification hook reports the error and keeps retryable + // verification failures inside the dialog. + } + }, + [secureVerification.executeVerification] + ) + const contextValue = useMemo( + () => ({ withVerification: secureVerification.withVerification }), + [secureVerification.withVerification] + ) + + return ( + + {props.children} + + + ) +} diff --git a/web/default/src/features/auth/secure-verification/hooks/use-api-token-verification.ts b/web/default/src/features/auth/secure-verification/hooks/use-api-token-verification.ts new file mode 100644 index 00000000..2ea25b09 --- /dev/null +++ b/web/default/src/features/auth/secure-verification/hooks/use-api-token-verification.ts @@ -0,0 +1,41 @@ +/* +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 { useCallback } from 'react' +import { useTranslation } from 'react-i18next' +import { useSecureVerificationGate } from '../components/secure-verification-provider' + +const defaultApiTokenVerificationDescription = + 'Confirm your identity before sending an API key to a chat client.' + +export function useApiTokenVerification( + description?: string +): (apiCall: () => Promise) => Promise { + const { t } = useTranslation() + const secureVerificationGate = useSecureVerificationGate() + + return useCallback( + (apiCall: () => Promise): Promise => + secureVerificationGate.withVerification(apiCall, { + scope: 'api_token', + title: t('Security verification'), + description: description ?? t(defaultApiTokenVerificationDescription), + }), + [description, t, secureVerificationGate.withVerification] + ) +} diff --git a/web/default/src/features/auth/secure-verification/hooks/use-secure-verification.ts b/web/default/src/features/auth/secure-verification/hooks/use-secure-verification.ts index eaf4b740..73c0688e 100644 --- a/web/default/src/features/auth/secure-verification/hooks/use-secure-verification.ts +++ b/web/default/src/features/auth/secure-verification/hooks/use-secure-verification.ts @@ -16,12 +16,14 @@ along with this program. If not, see . For commercial licensing, please contact https://github.com/MAX-API-Next/MAX-API/issues */ -import { useCallback, useMemo, useState } from 'react' +import { useCallback, useEffect, useMemo, useRef, useState } from 'react' import i18next from 'i18next' import { toast } from 'sonner' +import { handleServerError } from '@/lib/handle-server-error' import { extractVerificationInfo, isVerificationRequiredError, + markSecureVerificationErrorReported, } from '@/lib/secure-verification' import { checkVerificationMethods, verify } from '../api' import { selectVerificationMethod } from '../method-selection' @@ -36,8 +38,15 @@ import type { type ApiCall = (() => Promise) | null +interface PendingVerification { + attemptId: number + resolve: (value: unknown | null) => void + reject: (reason?: unknown) => void +} + interface InternalState extends SecureVerificationState { apiCall: ApiCall + attemptId: number | null } const defaultMethods: VerificationMethods = { @@ -54,6 +63,21 @@ const initialState: InternalState = { title: undefined, description: undefined, apiCall: null, + attemptId: null, +} + +function restrictVerificationMethods( + methods: VerificationMethods, + allowedMethods?: VerificationMethod[] +): VerificationMethods { + if (!allowedMethods?.length) return methods + const allowed = new Set(allowedMethods) + return { + has2FA: methods.has2FA && allowed.has('2fa'), + hasPasskey: methods.hasPasskey && allowed.has('passkey'), + hasPassword: methods.hasPassword && allowed.has('password'), + passkeySupported: methods.passkeySupported, + } } export function useSecureVerification( @@ -63,29 +87,99 @@ export function useSecureVerification( const [methods, setMethods] = useState(defaultMethods) const [state, setState] = useState(initialState) - const [open, setOpen] = useState(false) + const [open, setDialogOpen] = useState(false) + const pendingVerificationRef = useRef(null) + const verificationAttemptRef = useRef(0) + const mountedRef = useRef(true) + + const loadVerificationMethods = useCallback( + async (scope?: VerificationScope): Promise => + checkVerificationMethods(scope), + [] + ) const fetchVerificationMethods = useCallback( - async (scope?: VerificationScope) => { - const result = await checkVerificationMethods(scope) - setMethods(result) + async (scope?: VerificationScope): Promise => { + const result = await loadVerificationMethods(scope) + if (mountedRef.current) { + setMethods(result) + } return result }, + [loadVerificationMethods] + ) + + const settlePendingVerification = useCallback( + (value: unknown | null, error?: unknown, attemptId?: number): void => { + const pending = pendingVerificationRef.current + if (attemptId !== undefined && pending?.attemptId !== attemptId) return + pendingVerificationRef.current = null + if (!pending) return + if (error !== undefined) { + pending.reject(error) + } else { + pending.resolve(value) + } + }, + [] + ) + + const isVerificationAttemptActive = useCallback( + (attemptId: number): boolean => + mountedRef.current && verificationAttemptRef.current === attemptId, [] ) - const reset = useCallback(() => { + const beginVerificationAttempt = useCallback((): number | null => { + if (!mountedRef.current) return null + verificationAttemptRef.current += 1 + settlePendingVerification(null) setState(initialState) - setOpen(false) - }, []) + setDialogOpen(false) + return verificationAttemptRef.current + }, [settlePendingVerification]) - const startVerification = useCallback( + const reset = useCallback((): void => { + verificationAttemptRef.current += 1 + settlePendingVerification(null) + if (!mountedRef.current) return + setState(initialState) + setDialogOpen(false) + }, [settlePendingVerification]) + + useEffect(() => { + mountedRef.current = true + return () => { + mountedRef.current = false + verificationAttemptRef.current += 1 + settlePendingVerification(null) + } + }, [settlePendingVerification]) + + const startVerificationAttempt = useCallback( async ( apiCall: () => Promise, - config: StartVerificationOptions = {} - ) => { - const { preferredMethod, title, description, scope } = config - const availableMethods = await fetchVerificationMethods(scope) + config: StartVerificationOptions, + attemptId: number + ): Promise => { + let availableMethods: VerificationMethods + try { + availableMethods = restrictVerificationMethods( + await loadVerificationMethods(config.scope), + config.allowedMethods + ) + } catch (error) { + if (!isVerificationAttemptActive(attemptId)) return false + handleServerError(error, { + fallback: i18next.t('Verification unavailable'), + }) + markSecureVerificationErrorReported(error) + onError?.(error) + throw error + } + + if (!isVerificationAttemptActive(attemptId)) return false + setMethods(availableMethods) if ( !availableMethods.has2FA && @@ -94,12 +188,12 @@ export function useSecureVerification( ) { toast.error( i18next.t( - 'Please enable Two-factor Authentication or Passkey before proceeding' + 'No verification method is available. Add a password or sign in again with a linked account.' ) ) onError?.( new Error( - 'No verification methods available. Enable 2FA or Passkey to continue.' + 'No verification method is available. Add a password or sign in again with a linked account.' ) ) return false @@ -107,41 +201,76 @@ export function useSecureVerification( const defaultMethod = selectVerificationMethod( availableMethods, - preferredMethod + config.preferredMethod ) setState((prev) => ({ ...prev, apiCall, + attemptId, method: defaultMethod, - title, - description, - scope, + title: config.title, + description: config.description, + scope: config.scope, })) - setOpen(true) + setDialogOpen(true) return true }, - [fetchVerificationMethods, onError] + [isVerificationAttemptActive, loadVerificationMethods, onError] + ) + + const startVerification = useCallback( + async ( + apiCall: () => Promise, + config: StartVerificationOptions = {} + ): Promise => { + const attemptId = beginVerificationAttempt() + if (attemptId === null) return false + return startVerificationAttempt(apiCall, config, attemptId) + }, + [beginVerificationAttempt, startVerificationAttempt] ) const executeVerification = useCallback( - async (method?: VerificationMethod, code?: string) => { - if (!state.apiCall) { + async ( + method?: VerificationMethod, + code?: string + ): Promise => { + if (!state.apiCall || state.attemptId === null) { toast.error(i18next.t('Verification is not configured properly')) return } + const attemptId = state.attemptId + const apiCall = state.apiCall const actualMethod = method ?? state.method if (!actualMethod) { toast.error(i18next.t('Select a verification method first')) return } + if (!isVerificationAttemptActive(attemptId)) return setState((prev) => ({ ...prev, loading: true })) try { await verify(actualMethod, code ?? state.code, state.scope) - const result = await state.apiCall() + } catch (error) { + if (!isVerificationAttemptActive(attemptId)) return + setState((prev) => ({ ...prev, loading: false })) + handleServerError(error, { + fallback: i18next.t('Verification failed'), + }) + markSecureVerificationErrorReported(error) + onError?.(error) + throw error + } + + if (!isVerificationAttemptActive(attemptId)) return + + try { + const result = await apiCall() + if (!isVerificationAttemptActive(attemptId)) return + settlePendingVerification(result, undefined, attemptId) if (successMessage) { toast.success(successMessage) @@ -155,18 +284,33 @@ export function useSecureVerification( return result } catch (error) { - const message = - error instanceof Error - ? error.message - : i18next.t('Verification failed') - toast.error(message) + if (!isVerificationAttemptActive(attemptId)) return + settlePendingVerification(null, error, attemptId) + handleServerError(error, { + fallback: i18next.t('Verification failed'), + }) + markSecureVerificationErrorReported(error) onError?.(error) + reset() throw error } finally { - setState((prev) => ({ ...prev, loading: false })) + if (isVerificationAttemptActive(attemptId)) { + setState((prev) => + prev.attemptId === attemptId ? { ...prev, loading: false } : prev + ) + } } }, - [state, successMessage, onSuccess, onError, autoReset, reset] + [ + state, + isVerificationAttemptActive, + successMessage, + onSuccess, + onError, + autoReset, + reset, + settlePendingVerification, + ] ) const setCode = useCallback((code: string) => { @@ -177,28 +321,98 @@ export function useSecureVerification( setState((prev) => ({ ...prev, method, code: '' })) }, []) - const cancel = useCallback(() => { + const cancel = useCallback((): void => { reset() }, [reset]) + const setOpen = useCallback( + (nextOpen: boolean): void => { + if (!nextOpen) { + cancel() + return + } + setDialogOpen(true) + }, + [cancel] + ) + + /** + * Returns null when the user cancels. Callers must handle rejections from + * method discovery, the original action, and the post-verification retry. + */ const withVerification = useCallback( - async ( - apiCall: () => Promise, + ( + apiCall: () => Promise, config: StartVerificationOptions = {} - ) => { - try { - return await apiCall() - } catch (error) { - if (isVerificationRequiredError(error)) { + ): Promise => { + const attemptId = beginVerificationAttempt() + if (attemptId === null) return Promise.resolve(null) + const continuation = new Promise((resolve, reject) => { + pendingVerificationRef.current = { + attemptId, + resolve: (value) => resolve(value as T | null), + reject, + } + }) + // Keep the original promise and rejection contract for awaiting callers, + // while ensuring fire-and-forget UI effects do not produce an unhandled + // rejection after the shared flow has already reported the failure. + void continuation.catch(() => undefined) + + void (async (): Promise => { + if (config.forceVerification) { + try { + const started = await startVerificationAttempt( + apiCall, + config, + attemptId + ) + if (!started) { + settlePendingVerification(null, undefined, attemptId) + } + } catch (error) { + settlePendingVerification(null, error, attemptId) + } + return + } + try { + const result = await apiCall() + if (!isVerificationAttemptActive(attemptId)) return + settlePendingVerification(result, undefined, attemptId) + return + } catch (error) { + if (!isVerificationAttemptActive(attemptId)) return + if (!isVerificationRequiredError(error)) { + settlePendingVerification(null, error, attemptId) + return + } + const info = extractVerificationInfo(error) toast.info(info.message) - await startVerification(apiCall, config) - return null } - throw error - } + + try { + const started = await startVerificationAttempt( + apiCall, + config, + attemptId + ) + if (!started) { + settlePendingVerification(null, undefined, attemptId) + } + } catch (error) { + settlePendingVerification(null, error, attemptId) + } + })() + + return continuation }, - [startVerification] + [ + beginVerificationAttempt, + isVerificationAttemptActive, + settlePendingVerification, + startVerificationAttempt, + ] ) const canUseMethod = useCallback( diff --git a/web/default/src/features/auth/secure-verification/index.ts b/web/default/src/features/auth/secure-verification/index.ts index 1cb20155..4568b0e7 100644 --- a/web/default/src/features/auth/secure-verification/index.ts +++ b/web/default/src/features/auth/secure-verification/index.ts @@ -18,5 +18,7 @@ For commercial licensing, please contact https://github.com/MAX-API-Next/MAX-API */ export * from './api' export * from './types' +export * from './hooks/use-api-token-verification' export * from './hooks/use-secure-verification' export * from './components/secure-verification-dialog' +export * from './components/secure-verification-provider' diff --git a/web/default/src/features/auth/secure-verification/types.ts b/web/default/src/features/auth/secure-verification/types.ts index c3b7e99a..a3529794 100644 --- a/web/default/src/features/auth/secure-verification/types.ts +++ b/web/default/src/features/auth/secure-verification/types.ts @@ -17,7 +17,12 @@ along with this program. If not, see . For commercial licensing, please contact https://github.com/MAX-API-Next/MAX-API/issues */ export type VerificationMethod = '2fa' | 'passkey' | 'password' -export type VerificationScope = 'access_token' | 'account_delete' +export type VerificationScope = + | 'access_token' + | 'account_delete' + | 'credentials' + | 'api_token' + | 'passkey_register' export interface VerificationMethods { has2FA: boolean @@ -44,6 +49,8 @@ export interface UseSecureVerificationOptions { export interface StartVerificationOptions { preferredMethod?: VerificationMethod + allowedMethods?: VerificationMethod[] + forceVerification?: boolean title?: string description?: string scope?: VerificationScope diff --git a/web/default/src/features/chat/hooks/use-active-chat-key.ts b/web/default/src/features/chat/hooks/use-active-chat-key.ts index 6baef163..2e8d1cde 100644 --- a/web/default/src/features/chat/hooks/use-active-chat-key.ts +++ b/web/default/src/features/chat/hooks/use-active-chat-key.ts @@ -16,8 +16,11 @@ along with this program. If not, see . For commercial licensing, please contact https://github.com/MAX-API-Next/MAX-API/issues */ +import { useCallback } from 'react' import { useQuery } from '@tanstack/react-query' +import i18next from 'i18next' import { useAuthStore } from '@/stores/auth-store' +import { useApiTokenVerification } from '@/features/auth/secure-verification' import { fetchTokenKey, getApiKeys } from '@/features/keys/api' import { API_KEY_STATUS } from '@/features/keys/constants' @@ -46,12 +49,23 @@ export async function fetchActiveChatKey() { */ export function useActiveChatKey(enabled: boolean) { const userId = useAuthStore((state) => state.auth.user?.id) + const withApiTokenVerification = useApiTokenVerification() + const queryFn = useCallback(async (): Promise => { + const key = await withApiTokenVerification(fetchActiveChatKey) + if (!key) { + throw new Error(i18next.t('API key verification was cancelled')) + } + return key + }, [withApiTokenVerification]) return useQuery({ queryKey: ['chat-active-key', userId], - queryFn: fetchActiveChatKey, + queryFn, enabled: enabled && Boolean(userId), - staleTime: 5 * 60 * 1000, - gcTime: 10 * 60 * 1000, + staleTime: Number.POSITIVE_INFINITY, + gcTime: 0, + retry: false, + refetchOnWindowFocus: false, + refetchOnReconnect: false, }) } diff --git a/web/default/src/features/dashboard/components/overview/overview-dashboard.tsx b/web/default/src/features/dashboard/components/overview/overview-dashboard.tsx index 0df546f1..87b202ba 100644 --- a/web/default/src/features/dashboard/components/overview/overview-dashboard.tsx +++ b/web/default/src/features/dashboard/components/overview/overview-dashboard.tsx @@ -16,7 +16,7 @@ along with this program. If not, see . For commercial licensing, please contact https://github.com/MAX-API-Next/MAX-API/issues */ -import { useMemo, useState, type ReactNode } from 'react' +import { useMemo, useState, type JSX, type ReactNode } from 'react' import { useQuery } from '@tanstack/react-query' import { Link } from '@tanstack/react-router' import { @@ -45,8 +45,10 @@ import { toast } from 'sonner' import { useAuthStore } from '@/stores/auth-store' import { getUserModels } from '@/lib/api' import { formatNumber, formatQuota } from '@/lib/format' +import { handleServerError } from '@/lib/handle-server-error' import { MOTION_TRANSITION } from '@/lib/motion' import { ROLE } from '@/lib/roles' +import { wasSecureVerificationErrorReported } from '@/lib/secure-verification' import { cn } from '@/lib/utils' import { useCopyToClipboard } from '@/hooks/use-copy-to-clipboard' import { Button } from '@/components/ui/button' @@ -54,6 +56,7 @@ import { CardStaggerContainer, CardStaggerItem, } from '@/components/page-transition' +import { useApiTokenVerification } from '@/features/auth/secure-verification' import { fetchTokenKey, getApiKeys } from '@/features/keys/api' import type { ApiKey } from '@/features/keys/types' import { @@ -373,14 +376,19 @@ function StartStepItem(props: { ) } -function RequestPreview(props: { +type RequestPreviewProps = { example: RequestExample signals: HeroSignal[] isRetrying: boolean onRetry: () => void -}) { +} + +export function RequestPreview(props: RequestPreviewProps): JSX.Element { const { t } = useTranslation() const shouldReduceMotion = useReducedMotion() + const withApiTokenVerification = useApiTokenVerification( + t('Confirm your identity before copying a request with your API key.') + ) const [isCopying, setIsCopying] = useState(false) const { copyToClipboard } = useCopyToClipboard({ notify: false }) const previewCurl = buildCurlCommand({ @@ -389,12 +397,14 @@ function RequestPreview(props: { model: props.example.model, }) const previewLines = previewCurl.split('\n') - const handleCopyRequest = async () => { + const handleCopyRequest = async (): Promise => { if (!props.example.keyId || isCopying) return + const keyId = props.example.keyId setIsCopying(true) try { - const result = await fetchTokenKey(props.example.keyId) + const result = await withApiTokenVerification(() => fetchTokenKey(keyId)) + if (!result) return const key = result.success && result.data?.key ? result.data.key : '' if (!key) { toast.error(result.message || t('Failed to copy to clipboard')) @@ -412,6 +422,12 @@ function RequestPreview(props: { } else { toast.error(t('Failed to copy to clipboard')) } + } catch (error) { + if (!wasSecureVerificationErrorReported(error)) { + handleServerError(error, { + fallback: t('Failed to copy to clipboard'), + }) + } } finally { setIsCopying(false) } diff --git a/web/default/src/features/keys/api.ts b/web/default/src/features/keys/api.ts index ca5bceaa..f86e8a70 100644 --- a/web/default/src/features/keys/api.ts +++ b/web/default/src/features/keys/api.ts @@ -16,7 +16,7 @@ along with this program. If not, see . For commercial licensing, please contact https://github.com/MAX-API-Next/MAX-API/issues */ -import { api } from '@/lib/api' +import { api, sensitiveActionConfig } from '@/lib/api' import type { ApiKey, ApiResponse, @@ -63,7 +63,7 @@ export async function getApiKey(id: number): Promise> { export async function createApiKey( data: ApiKeyFormData ): Promise> { - const res = await api.post('/api/token/', data) + const res = await api.post('/api/token/', data, sensitiveActionConfig) return res.data } @@ -102,7 +102,11 @@ export async function updateApiKeyStatus( export async function fetchTokenKey( id: number ): Promise<{ success: boolean; message?: string; data?: { key: string } }> { - const res = await api.post(`/api/token/${id}/key`) + const res = await api.post( + `/api/token/${id}/key`, + undefined, + sensitiveActionConfig + ) return res.data } @@ -112,6 +116,10 @@ export async function fetchTokenKeysBatch(ids: number[]): Promise<{ message?: string data?: { keys: Record } }> { - const res = await api.post('/api/token/batch/keys', { ids }) + const res = await api.post( + '/api/token/batch/keys', + { ids }, + sensitiveActionConfig + ) return res.data } diff --git a/web/default/src/features/keys/components/api-keys-cells.tsx b/web/default/src/features/keys/components/api-keys-cells.tsx index cfd473f3..fc838faa 100644 --- a/web/default/src/features/keys/components/api-keys-cells.tsx +++ b/web/default/src/features/keys/components/api-keys-cells.tsx @@ -41,6 +41,7 @@ export function ApiKeyCell({ apiKey }: { apiKey: ApiKey }) { const { resolveRealKey, resolvedKeys, + clearResolvedKey, loadingKeys, copiedKeyId, markKeyCopied, @@ -53,27 +54,27 @@ export function ApiKeyCell({ apiKey }: { apiKey: ApiKey }) { const maskedKey = `sk-${apiKey.key}` const handlePopoverOpen = useCallback( - (open: boolean) => { + (open: boolean): void => { setPopoverOpen(open) if (open && !resolvedFullKey) { - resolveRealKey(apiKey.id) + void resolveRealKey(apiKey.id, { cache: true }) + } else if (!open) { + clearResolvedKey(apiKey.id) } }, - [resolvedFullKey, resolveRealKey, apiKey.id] + [clearResolvedKey, resolvedFullKey, resolveRealKey, apiKey.id] ) - const handleCopy = useCallback(async () => { - const realKey = resolvedFullKey - if (!realKey) { - void resolveRealKey(apiKey.id) - toast.info(t('API key is loading, please try again in a moment')) - return + const handleCopy = useCallback(async (): Promise => { + const realKey = await resolveRealKey(apiKey.id) + if (!realKey) return + const ok = await copyToClipboard(realKey) + if (ok) { + markKeyCopied(apiKey.id) + } else { + toast.error(t('Failed to copy to clipboard')) } - if (realKey) { - const ok = await copyToClipboard(realKey) - if (ok) markKeyCopied(apiKey.id) - } - }, [resolvedFullKey, resolveRealKey, apiKey.id, markKeyCopied, t]) + }, [resolveRealKey, apiKey.id, markKeyCopied, t]) return (
@@ -122,13 +123,8 @@ export function ApiKeyCell({ apiKey }: { apiKey: ApiKey }) { size='icon' className='size-7 shrink-0' onClick={handleCopy} - onFocus={() => { - if (!resolvedFullKey) void resolveRealKey(apiKey.id) - }} - onPointerEnter={() => { - if (!resolvedFullKey) void resolveRealKey(apiKey.id) - }} disabled={isLoading} + aria-label={t('Copy API key')} /> } > diff --git a/web/default/src/features/keys/components/api-keys-dialogs.tsx b/web/default/src/features/keys/components/api-keys-dialogs.tsx index 7e448e8b..76bcf78c 100644 --- a/web/default/src/features/keys/components/api-keys-dialogs.tsx +++ b/web/default/src/features/keys/components/api-keys-dialogs.tsx @@ -16,25 +16,34 @@ along with this program. If not, see . For commercial licensing, please contact https://github.com/MAX-API-Next/MAX-API/issues */ +import type { ReactElement } from 'react' import { ApiKeysDeleteDialog } from './api-keys-delete-dialog' import { ApiKeysMutateDrawer } from './api-keys-mutate-drawer' import { useApiKeys } from './api-keys-provider' import { CCSwitchDialog } from './dialogs/cc-switch-dialog' -export function ApiKeysDialogs() { - const { open, setOpen, currentRow, resolvedKey } = useApiKeys() +export function ApiKeysDialogs(): ReactElement { + const { open, setOpen, currentRow, resolvedKey, setResolvedKey } = + useApiKeys() return ( <> !isOpen && setOpen(null)} + onOpenChange={(isOpen: boolean): void => { + if (!isOpen) setOpen(null) + }} currentRow={open === 'update' ? currentRow || undefined : undefined} /> !isOpen && setOpen(null)} + onOpenChange={(isOpen: boolean): void => { + if (!isOpen) { + setResolvedKey('') + setOpen(null) + } + }} tokenKey={resolvedKey} /> diff --git a/web/default/src/features/keys/components/api-keys-mutate-drawer.tsx b/web/default/src/features/keys/components/api-keys-mutate-drawer.tsx index 7b13a990..95f09d50 100644 --- a/web/default/src/features/keys/components/api-keys-mutate-drawer.tsx +++ b/web/default/src/features/keys/components/api-keys-mutate-drawer.tsx @@ -16,7 +16,7 @@ along with this program. If not, see . For commercial licensing, please contact https://github.com/MAX-API-Next/MAX-API/issues */ -import { useEffect, useMemo, useRef, useState } from 'react' +import { useEffect, useMemo, useRef, useState, type ReactElement } from 'react' import { useForm, type SubmitErrorHandler } from 'react-hook-form' import { zodResolver } from '@hookform/resolvers/zod' import { useQuery } from '@tanstack/react-query' @@ -27,6 +27,7 @@ import { getUserModels, getUserGroups } from '@/lib/api' import { DEFAULT_AUTO_ROUTE_KEY } from '@/lib/auto-routes' import { getCurrencyDisplay, getCurrencyLabel } from '@/lib/currency' import { handleServerError } from '@/lib/handle-server-error' +import { wasSecureVerificationErrorReported } from '@/lib/secure-verification' import { cn } from '@/lib/utils' import { useStatus } from '@/hooks/use-status' import { Button } from '@/components/ui/button' @@ -92,14 +93,12 @@ type ApiKeyMutateDrawerProps = { currentRow?: ApiKey } -export function ApiKeysMutateDrawer({ - open, - onOpenChange, - currentRow, -}: ApiKeyMutateDrawerProps) { +export function ApiKeysMutateDrawer( + props: ApiKeyMutateDrawerProps +): ReactElement { const { t } = useTranslation() - const isUpdate = !!currentRow - const { triggerRefresh } = useApiKeys() + const isUpdate = !!props.currentRow + const { triggerRefresh, withApiTokenVerification } = useApiKeys() const { status } = useStatus() const [isSubmitting, setIsSubmitting] = useState(false) const [advancedOpen, setAdvancedOpen] = useState(false) @@ -108,7 +107,7 @@ export function ApiKeysMutateDrawer({ const [preservedManualGroups, setPreservedManualGroups] = useState( [] ) - const currentRowId = currentRow?.id + const currentRowId = props.currentRow?.id const defaultAutoRoute = typeof status?.default_auto_route === 'string' ? status.default_auto_route @@ -221,7 +220,7 @@ export function ApiKeysMutateDrawer({ // Load existing data when updating useEffect(() => { - if (!open || groupsLoading) return + if (!props.open || groupsLoading) return let cancelled = false const { @@ -293,10 +292,11 @@ export function ApiKeysMutateDrawer({ return () => { cancelled = true } - }, [open, isUpdate, currentRowId, form, groupsLoading, t]) + }, [props.open, isUpdate, currentRowId, form, groupsLoading, t]) - const onSubmit = async (data: ApiKeyFormValues) => { + const onSubmit = async (data: ApiKeyFormValues): Promise => { setIsSubmitting(true) + let created = 0 try { const dirtyFields = form.formState.dirtyFields const routingChanged = Boolean( @@ -315,14 +315,14 @@ export function ApiKeysMutateDrawer({ ), }) - if (isUpdate && currentRow) { + if (isUpdate && props.currentRow) { const result = await updateApiKey({ ...basePayload, - id: currentRow.id, + id: props.currentRow.id, }) if (result.success) { toast.success(t(SUCCESS_MESSAGES.API_KEY_UPDATED)) - onOpenChange(false) + props.onOpenChange(false) triggerRefresh() } else { toast.error(result.message || t(ERROR_MESSAGES.UPDATE_FAILED)) @@ -330,38 +330,43 @@ export function ApiKeysMutateDrawer({ } else { // Create mode - handle batch creation const count = data.tokenCount || 1 - let successCount = 0 - - for (let i = 0; i < count; i++) { - const result = await createApiKey({ - ...basePayload, - name: - i === 0 && data.name - ? data.name - : `${data.name || 'default'}-${Math.random().toString(36).slice(2, 8)}`, - }) - if (result.success) { - successCount++ - } else { - toast.error(result.message || t(ERROR_MESSAGES.CREATE_FAILED)) - break + await withApiTokenVerification(async (): Promise => { + for (let i = created; i < count; i++) { + const result = await createApiKey({ + ...basePayload, + name: + i === 0 && data.name + ? data.name + : `${data.name || 'default'}-${Math.random().toString(36).slice(2, 8)}`, + }) + if (result.success) { + created++ + } else { + toast.error(result.message || t(ERROR_MESSAGES.CREATE_FAILED)) + break + } } - } - if (successCount > 0) { - toast.success( - t('Successfully created {{count}} API Key(s)', { - count: successCount, - }) - ) - onOpenChange(false) - triggerRefresh() - } + return created + }) + } + } catch (error) { + if (!wasSecureVerificationErrorReported(error)) { + handleServerError(error, { + fallback: t(ERROR_MESSAGES.UNEXPECTED), + }) } - } catch (_error) { - toast.error(t(ERROR_MESSAGES.UNEXPECTED)) } finally { setIsSubmitting(false) + if (!isUpdate && created > 0) { + toast.success( + t('Successfully created {{count}} API Key(s)', { + count: created, + }) + ) + props.onOpenChange(false) + triggerRefresh() + } } } @@ -399,9 +404,9 @@ export function ApiKeysMutateDrawer({ return ( { - onOpenChange(v) + open={props.open} + onOpenChange={(v: boolean): void => { + props.onOpenChange(v) if (!v) { setEditingLegacyRouting(false) setPreservedSmartRoute(undefined) diff --git a/web/default/src/features/keys/components/api-keys-provider.tsx b/web/default/src/features/keys/components/api-keys-provider.tsx index 0207d5ee..36f8a025 100644 --- a/web/default/src/features/keys/components/api-keys-provider.tsx +++ b/web/default/src/features/keys/components/api-keys-provider.tsx @@ -19,7 +19,10 @@ For commercial licensing, please contact https://github.com/MAX-API-Next/MAX-API import React, { useState, useCallback, useRef, useEffect } from 'react' import { useTranslation } from 'react-i18next' import { toast } from 'sonner' +import { handleServerError } from '@/lib/handle-server-error' +import { wasSecureVerificationErrorReported } from '@/lib/secure-verification' import useDialogState from '@/hooks/use-dialog' +import { useApiTokenVerification } from '@/features/auth/secure-verification' import { fetchTokenKey, fetchTokenKeysBatch } from '../api' import { ERROR_MESSAGES } from '../constants' import { type ApiKey, type ApiKeysDialogType } from '../types' @@ -33,18 +36,32 @@ type ApiKeysContextType = { triggerRefresh: () => void resolvedKey: string setResolvedKey: React.Dispatch> - resolveRealKey: (id: number) => Promise + resolveRealKey: ( + id: number, + options?: { cache?: boolean } + ) => Promise resolveRealKeysBatch: (ids: number[]) => Promise> resolvedKeys: Record + clearResolvedKey: (id: number) => void loadingKeys: Record copiedKeyId: number | null markKeyCopied: (id: number) => void + withApiTokenVerification: (apiCall: () => Promise) => Promise } const ApiKeysContext = React.createContext(null) -export function ApiKeysProvider({ children }: { children: React.ReactNode }) { +type ApiKeysProviderProps = { + children: React.ReactNode +} + +export function ApiKeysProvider( + props: ApiKeysProviderProps +): React.JSX.Element { const { t } = useTranslation() + const withApiTokenVerification = useApiTokenVerification( + t('Confirm your identity before creating or revealing API keys.') + ) const [open, setOpen] = useDialogState(null) const [currentRow, setCurrentRow] = useState(null) const [refreshTrigger, setRefreshTrigger] = useState(0) @@ -52,7 +69,9 @@ export function ApiKeysProvider({ children }: { children: React.ReactNode }) { const [resolvedKeys, setResolvedKeys] = useState>({}) const [loadingKeys, setLoadingKeys] = useState>({}) - const pendingRequests = useRef>>({}) + const pendingRequests = useRef< + Record; cacheRequested: boolean }> + >({}) const [copiedKeyId, setCopiedKeyId] = useState(null) const copiedTimerRef = useRef>(undefined) @@ -72,23 +91,37 @@ export function ApiKeysProvider({ children }: { children: React.ReactNode }) { }, []) const resolveRealKey = useCallback( - async (id: number): Promise => { + async ( + id: number, + options: { cache?: boolean } = {} + ): Promise => { if (resolvedKeys[id]) return resolvedKeys[id] - if (id in pendingRequests.current) return pendingRequests.current[id] + const pending = pendingRequests.current[id] + if (pending) { + pending.cacheRequested ||= Boolean(options.cache) + return pending.promise + } const request = (async () => { setLoadingKeys((prev) => ({ ...prev, [id]: true })) try { - const res = await fetchTokenKey(id) + const res = await withApiTokenVerification(() => fetchTokenKey(id)) + if (!res) return null if (res.success && res.data?.key) { const fullKey = `sk-${res.data.key}` - setResolvedKeys((prev) => ({ ...prev, [id]: fullKey })) + if (pendingRequests.current[id]?.cacheRequested) { + setResolvedKeys((prev) => ({ ...prev, [id]: fullKey })) + } return fullKey } toast.error(res.message || t(ERROR_MESSAGES.UNEXPECTED)) return null - } catch { - toast.error(t(ERROR_MESSAGES.UNEXPECTED)) + } catch (error) { + if (!wasSecureVerificationErrorReported(error)) { + handleServerError(error, { + fallback: t(ERROR_MESSAGES.UNEXPECTED), + }) + } return null } finally { delete pendingRequests.current[id] @@ -100,47 +133,55 @@ export function ApiKeysProvider({ children }: { children: React.ReactNode }) { } })() - pendingRequests.current[id] = request + pendingRequests.current[id] = { + promise: request, + cacheRequested: Boolean(options.cache), + } return request }, - [resolvedKeys, t] + [resolvedKeys, t, withApiTokenVerification] ) + const clearResolvedKey = useCallback((id: number): void => { + const pending = pendingRequests.current[id] + if (pending) pending.cacheRequested = false + setResolvedKeys((prev) => { + if (!(id in prev)) return prev + const next = { ...prev } + delete next[id] + return next + }) + }, []) + const resolveRealKeysBatch = useCallback( async (ids: number[]): Promise> => { - const uncachedIds = ids.filter((id) => !resolvedKeys[id]) - if (uncachedIds.length === 0) { - const result: Record = {} - for (const id of ids) result[id] = resolvedKeys[id] - return result - } - - for (const id of uncachedIds) { + for (const id of ids) { setLoadingKeys((prev) => ({ ...prev, [id]: true })) } try { - const res = await fetchTokenKeysBatch(uncachedIds) + const res = await withApiTokenVerification(() => + fetchTokenKeysBatch(ids) + ) + if (!res) return {} if (res.success && res.data?.keys) { const newKeys: Record = {} for (const [idStr, key] of Object.entries(res.data.keys)) { newKeys[Number(idStr)] = `sk-${key}` } - setResolvedKeys((prev) => ({ ...prev, ...newKeys })) - - const result: Record = { ...newKeys } - for (const id of ids) { - if (resolvedKeys[id]) result[id] = resolvedKeys[id] - } - return result + return newKeys } toast.error(res.message || t(ERROR_MESSAGES.UNEXPECTED)) return {} - } catch { - toast.error(t(ERROR_MESSAGES.UNEXPECTED)) + } catch (error) { + if (!wasSecureVerificationErrorReported(error)) { + handleServerError(error, { + fallback: t(ERROR_MESSAGES.UNEXPECTED), + }) + } return {} } finally { - for (const id of uncachedIds) { + for (const id of ids) { setLoadingKeys((prev) => { const next = { ...prev } delete next[id] @@ -149,7 +190,7 @@ export function ApiKeysProvider({ children }: { children: React.ReactNode }) { } } }, - [resolvedKeys, t] + [t, withApiTokenVerification] ) return ( @@ -166,12 +207,14 @@ export function ApiKeysProvider({ children }: { children: React.ReactNode }) { resolveRealKey, resolveRealKeysBatch, resolvedKeys, + clearResolvedKey, loadingKeys, copiedKeyId, markKeyCopied, + withApiTokenVerification, }} > - {children} + {props.children} ) } diff --git a/web/default/src/features/keys/components/data-table-row-actions.tsx b/web/default/src/features/keys/components/data-table-row-actions.tsx index 4eca6280..e975ff0c 100644 --- a/web/default/src/features/keys/components/data-table-row-actions.tsx +++ b/web/default/src/features/keys/components/data-table-row-actions.tsx @@ -52,7 +52,11 @@ import { TooltipTrigger, } from '@/components/ui/tooltip' import { useChatPresets } from '@/features/chat/hooks/use-chat-presets' -import { resolveChatUrl, type ChatPreset } from '@/features/chat/lib/chat-links' +import { + chatLinkRequiresApiKey, + resolveChatUrl, + type ChatPreset, +} from '@/features/chat/lib/chat-links' import { sendToFluent } from '@/features/chat/lib/send-to-fluent' import { updateApiKeyStatus } from '../api' import { API_KEY_STATUS, ERROR_MESSAGES, SUCCESS_MESSAGES } from '../constants' @@ -87,39 +91,37 @@ export function DataTableRowActions({ triggerRefresh, setResolvedKey, resolveRealKey, - resolvedKeys, - loadingKeys, } = useApiKeys() const isEnabled = apiKey.status === API_KEY_STATUS.ENABLED const { chatPresets, serverAddress } = useChatPresets() const [isTogglingStatus, setIsTogglingStatus] = useState(false) - const resolvedRealKey = resolvedKeys[apiKey.id] - const isRealKeyLoading = Boolean(loadingKeys[apiKey.id]) - const hasChatPresets = chatPresets.length > 0 - const handleMenuOpenChange = useCallback( - (open: boolean) => { - if (open && !resolvedRealKey && !isRealKeyLoading) { - void resolveRealKey(apiKey.id) - } - }, - [apiKey.id, isRealKeyLoading, resolvedRealKey, resolveRealKey] - ) - - const getCachedRealKey = useCallback(() => { - if (resolvedRealKey) return resolvedRealKey - void resolveRealKey(apiKey.id) - toast.info(t('API key is loading, please try again in a moment')) - return null - }, [apiKey.id, resolvedRealKey, resolveRealKey, t]) - const handleOpenChatPreset = useCallback( async (preset: ChatPreset) => { - const realKey = await resolveRealKey(apiKey.id) - if (!realKey) return + let chatWindow: Window | null = null + if (preset.type !== 'fluent' && typeof window !== 'undefined') { + try { + chatWindow = window.open('about:blank', '_blank') + if (chatWindow) chatWindow.opener = null + } catch { + chatWindow = null + } + } + + const needsApiKey = chatLinkRequiresApiKey(preset.url) + let realKey: string | undefined + if (needsApiKey || preset.type === 'fluent') { + const resolvedKey = await resolveRealKey(apiKey.id) + if (!resolvedKey) { + chatWindow?.close() + return + } + realKey = resolvedKey + } if (preset.type === 'fluent') { + if (!realKey) return const success = sendToFluent(realKey, serverAddress) if (success) { toast.success(t('Sent the API key to FluentRead.')) @@ -135,11 +137,12 @@ export function DataTableRowActions({ const resolvedUrl = resolveChatUrl({ template: preset.url, - apiKey: realKey, + apiKey: needsApiKey ? realKey : undefined, serverAddress, }) if (!resolvedUrl) { + chatWindow?.close() toast.error(t('Invalid chat link. Please contact your administrator.')) return } @@ -147,8 +150,13 @@ export function DataTableRowActions({ if (typeof window === 'undefined') return try { - window.open(resolvedUrl, '_blank', 'noopener,noreferrer') + if (chatWindow) { + chatWindow.location.replace(resolvedUrl) + return + } + window.location.href = resolvedUrl } catch { + chatWindow?.close() window.location.href = resolvedUrl } }, @@ -214,7 +222,7 @@ export function DataTableRowActions({ - + ({ { - const realKey = getCachedRealKey() + const realKey = await resolveRealKey(apiKey.id) if (!realKey) return const ok = await copyToClipboard(realKey) - if (ok) toast.success(t('Copied')) + if (ok) { + toast.success(t('Copied')) + } else { + toast.error(t('Failed to copy to clipboard')) + } }} > {t('Copy Key')} @@ -242,14 +254,18 @@ export function DataTableRowActions({ { - const realKey = getCachedRealKey() + const realKey = await resolveRealKey(apiKey.id) if (!realKey) return const connStr = encodeChannelConnectionInfo( realKey, getServerAddress() ) const ok = await copyToClipboard(connStr) - if (ok) toast.success(t('Copied')) + if (ok) { + toast.success(t('Copied')) + } else { + toast.error(t('Failed to copy to clipboard')) + } }} > {t('Copy Connection Info')} diff --git a/web/default/src/features/profile/api.ts b/web/default/src/features/profile/api.ts index b4624736..316b8d74 100644 --- a/web/default/src/features/profile/api.ts +++ b/web/default/src/features/profile/api.ts @@ -16,7 +16,7 @@ along with this program. If not, see . For commercial licensing, please contact https://github.com/MAX-API-Next/MAX-API/issues */ -import { api, type ApiRequestConfig } from '@/lib/api' +import { api, sensitiveActionConfig } from '@/lib/api' import { selfCustomOAuthUnbindPath, type CustomOAuthBinding } from '@/lib/oauth' import { getTurnstileHeaders } from '@/features/auth/lib/turnstile-request' import type { @@ -29,11 +29,6 @@ import type { CheckinResponse, } from './types' -const sensitiveActionConfig: ApiRequestConfig = { - skipBusinessError: true, - skipErrorHandler: true, -} - // ============================================================================ // User Profile APIs // ============================================================================ @@ -141,6 +136,47 @@ export async function bindWeChat(code: string): Promise { return res.data } +export interface TelegramAuthorizationPayload { + id: number + first_name?: string + last_name?: string + username?: string + photo_url?: string + auth_date: number + hash: string +} + +export async function createTelegramBindState(): Promise< + ApiResponse<{ state: string }> +> { + const res = await api.post>( + '/api/oauth/telegram/bind/state', + undefined, + sensitiveActionConfig + ) + return res.data +} + +export async function bindTelegramAccount( + payload: TelegramAuthorizationPayload & { state: string } +): Promise { + const res = await api.post( + '/api/oauth/telegram/bind', + payload, + sensitiveActionConfig + ) + return res.data +} + +export async function revokeOtherSessions(): Promise { + const res = await api.post( + '/api/user/sessions/revoke', + undefined, + sensitiveActionConfig + ) + return res.data +} + // ============================================================================ // Custom OAuth Binding APIs // ============================================================================ diff --git a/web/default/src/features/profile/components/dialogs/telegram-bind-dialog.tsx b/web/default/src/features/profile/components/dialogs/telegram-bind-dialog.tsx index 20d6262b..ec2268a6 100644 --- a/web/default/src/features/profile/components/dialogs/telegram-bind-dialog.tsx +++ b/web/default/src/features/profile/components/dialogs/telegram-bind-dialog.tsx @@ -16,9 +16,12 @@ along with this program. If not, see . For commercial licensing, please contact https://github.com/MAX-API-Next/MAX-API/issues */ -import { Send } from 'lucide-react' +import type { ReactElement } from 'react' +import { RefreshCw, Send } from 'lucide-react' import { useTranslation } from 'react-i18next' +import { cn } from '@/lib/utils' import { Alert, AlertDescription } from '@/components/ui/alert' +import { Button } from '@/components/ui/button' import { Dialog, DialogContent, @@ -26,6 +29,8 @@ import { DialogHeader, DialogTitle, } from '@/components/ui/dialog' +import { Spinner } from '@/components/ui/spinner' +import { useTelegramLoginWidget } from '../../hooks/use-telegram-login-widget' // ============================================================================ // Telegram Bind Dialog Component @@ -38,14 +43,20 @@ interface TelegramBindDialogProps { onSuccess: () => void } -export function TelegramBindDialog({ - open, - onOpenChange, - botName, -}: TelegramBindDialogProps) { +export function TelegramBindDialog( + props: TelegramBindDialogProps +): ReactElement { const { t } = useTranslation() + const { widgetContainerRef, status, errorMessage, retry, handleOpenChange } = + useTelegramLoginWidget({ + open: props.open, + onOpenChange: props.onOpenChange, + botName: props.botName, + onSuccess: props.onSuccess, + }) + return ( - + {t('Bind Telegram Account')} @@ -54,25 +65,30 @@ export function TelegramBindDialog({ -
+
- + -
-
- +
+
+

{t('Bot:')}{' '} - @{botName} + + @{props.botName.replace(/^@/, '')} +

{t( @@ -81,13 +97,45 @@ export function TelegramBindDialog({

- {/* Telegram Login Widget will be injected here by react-telegram-login */} -
- {/* This would require the react-telegram-login library */} -
- {t('Telegram Login Widget')} + {(status === 'loading' || status === 'binding') && ( +
+ + {status === 'binding' + ? t('Binding Telegram account...') + : t('Preparing Telegram authorization...')}
-
+ )} + + {status === 'error' && ( +
+

+ {errorMessage} +

+ +
+ )} + +
+ + {status === 'idle' && ( + + )}

diff --git a/web/default/src/features/profile/components/dialogs/two-fa-setup-dialog.tsx b/web/default/src/features/profile/components/dialogs/two-fa-setup-dialog.tsx index 2e0e7cbd..db2576f4 100644 --- a/web/default/src/features/profile/components/dialogs/two-fa-setup-dialog.tsx +++ b/web/default/src/features/profile/components/dialogs/two-fa-setup-dialog.tsx @@ -16,12 +16,14 @@ along with this program. If not, see . For commercial licensing, please contact https://github.com/MAX-API-Next/MAX-API/issues */ -import { useState, useEffect, useCallback } from 'react' +import { useState, useEffect, useCallback, useRef, type ReactNode } from 'react' import { Loader2 } from 'lucide-react' import { QRCodeSVG } from 'qrcode.react' import { useTranslation } from 'react-i18next' import { toast } from 'sonner' import { setup2FA, enable2FA } from '@/lib/api' +import { handleServerError } from '@/lib/handle-server-error' +import { wasSecureVerificationErrorReported } from '@/lib/secure-verification' import { Alert, AlertDescription } from '@/components/ui/alert' import { Button } from '@/components/ui/button' import { @@ -35,6 +37,7 @@ import { import { Input } from '@/components/ui/input' import { Label } from '@/components/ui/label' import { CopyButton } from '@/components/copy-button' +import { useSecureVerificationGate } from '@/features/auth/secure-verification' import type { TwoFASetupData } from '../../types' // ============================================================================ @@ -47,6 +50,111 @@ interface TwoFASetupDialogProps { onSuccess: () => void } +interface TwoFASetupStepsProps { + setupData: TwoFASetupData + step: number + code: string + loading: boolean + onCodeChange: (code: string) => void +} + +function TwoFASetupSteps(props: TwoFASetupStepsProps): ReactNode { + const { t } = useTranslation() + + if (props.step === 0) { + return ( +

+

+ {t( + 'Scan this QR code with your authenticator app (Google Authenticator, Microsoft Authenticator, etc.)' + )} +

+
+ +
+
+
+
+

+ {t('Or enter this key manually:')} +

+ + {props.setupData.secret} + +
+ +
+
+
+ ) + } + + if (props.step === 1) { + return ( +
+ + + {t( + 'Save these backup codes in a safe place. Each code can only be used once.' + )} + + +
+
+ {props.setupData.backup_codes.map((backupCode, index) => ( +
+ {backupCode} +
+ ))} +
+
+ + {t('Copy All Codes')} + +
+ ) + } + + if (props.step === 2) { + return ( +
+
+ + props.onCodeChange(event.target.value)} + placeholder={t('Enter 6-digit code')} + maxLength={6} + disabled={props.loading} + /> +

+ {t('Enter the 6-digit code from your authenticator app')} +

+
+
+ ) + } + + return null +} + export function TwoFASetupDialog({ open, onOpenChange, @@ -55,36 +163,60 @@ export function TwoFASetupDialog({ const { t } = useTranslation() const [loading, setLoading] = useState(false) const [initializing, setInitializing] = useState(false) + const [setupAttempted, setSetupAttempted] = useState(false) const [step, setStep] = useState(0) const [setupData, setSetupData] = useState(null) const [code, setCode] = useState('') + const setupInFlightRef = useRef(false) + const { withVerification } = useSecureVerificationGate() const stepLabels = [ t('Scan QR Code'), t('Save Backup Codes'), t('Verify Setup'), ] - const handleSetup = useCallback(async () => { + const handleSetup = useCallback(async (): Promise => { + if (setupInFlightRef.current) return + setupInFlightRef.current = true + setSetupAttempted(false) + setInitializing(true) + let shouldMarkSetupAttempted = true try { - setInitializing(true) - const response = await setup2FA() - - if (response.success && response.data) { - setSetupData(response.data) - setStep(0) - } else { + const response = await withVerification(setup2FA, { + scope: 'credentials', + title: t('Security verification'), + description: t( + 'Confirm your identity before setting up Two-factor Authentication.' + ), + }) + if (!response) { + shouldMarkSetupAttempted = false + setSetupAttempted(false) + onOpenChange(false) + return + } + if (!response.success || !response.data) { toast.error(response.message || t('Failed to setup 2FA')) + shouldMarkSetupAttempted = false + setSetupAttempted(false) onOpenChange(false) + return } + setSetupData(response.data) + setStep(0) } catch (error) { - // eslint-disable-next-line no-console - console.error('Setup 2FA error:', error) - toast.error(t('Failed to setup 2FA')) + if (!wasSecureVerificationErrorReported(error)) { + handleServerError(error, { fallback: t('Failed to setup 2FA') }) + } + shouldMarkSetupAttempted = false + setSetupAttempted(false) onOpenChange(false) } finally { + setupInFlightRef.current = false setInitializing(false) + setSetupAttempted(shouldMarkSetupAttempted) } - }, [onOpenChange, t]) + }, [onOpenChange, t, withVerification]) const handleEnable = async () => { if (!code) { @@ -94,7 +226,14 @@ export function TwoFASetupDialog({ try { setLoading(true) - const response = await enable2FA(code) + const response = await withVerification(() => enable2FA(code), { + scope: 'credentials', + title: t('Security verification'), + description: t( + 'Confirm your identity before enabling Two-factor Authentication.' + ), + }) + if (!response) return if (response.success) { toast.success(t('Two-factor authentication enabled successfully!')) @@ -104,25 +243,26 @@ export function TwoFASetupDialog({ setStep(0) setCode('') setSetupData(null) + setSetupAttempted(false) } else { toast.error(response.message || t('Failed to enable 2FA')) } - } catch (_error) { - toast.error(t('Failed to enable 2FA')) + } catch (error) { + if (!wasSecureVerificationErrorReported(error)) { + handleServerError(error, { fallback: t('Failed to enable 2FA') }) + } } finally { setLoading(false) } } - const handleOpenChange = (open: boolean) => { + const handleOpenChange = (open: boolean): void => { if (!loading && !initializing) { - if (open && !setupData) { - handleSetup() - } if (!open) { setStep(0) setCode('') setSetupData(null) + setSetupAttempted(false) } onOpenChange(open) } @@ -130,156 +270,80 @@ export function TwoFASetupDialog({ // Initialize when dialog opens useEffect(() => { - if (open && !setupData && !initializing) { - handleSetup() - } - }, [open, setupData, initializing, handleSetup]) + if (!open || setupData) return + const timer = window.setTimeout(() => void handleSetup(), 0) + return () => window.clearTimeout(timer) + }, [open, setupData, handleSetup]) - return ( - - - - {t('Setup Two-Factor Authentication')} - - {t('Step')} {step + 1} {t('of 3:')} {stepLabels[step]} - - + const renderSetupState = (): ReactNode => { + if (initializing || !setupAttempted) { + return ( +
+
+
+ {t('Setting up 2FA...')} +
+
+ ) + } + return null + } -
- {initializing ? ( -
-
-
- {t('Setting up 2FA...')} -
-
- ) : !setupData ? ( -
-
- {t('Failed to load setup data')} -
-
- ) : ( - <> - {/* Step 0: QR Code */} - {step === 0 && ( -
-

- {t( - 'Scan this QR code with your authenticator app (Google Authenticator, Microsoft Authenticator, etc.)' - )} -

-
- -
-
-
-
-

- {t('Or enter this key manually:')} -

- - {setupData.secret} - -
- -
-
-
- )} + const setupState = renderSetupState() - {/* Step 1: Backup Codes */} - {step === 1 && ( -
- - - {t( - 'Save these backup codes in a safe place. Each code can only be used once.' - )} - - -
-
- {setupData.backup_codes.map((code, index) => ( -
- {code} -
- ))} -
-
- - {t('Copy All Codes')} - -
- )} + return ( + <> + + + + {t('Setup Two-Factor Authentication')} + + {t('Step')} {step + 1} {t('of 3:')} {stepLabels[step]} + + - {/* Step 2: Verify */} - {step === 2 && ( -
-
- - setCode(e.target.value)} - placeholder={t('Enter 6-digit code')} - maxLength={6} - disabled={loading} - /> -

- {t('Enter the 6-digit code from your authenticator app')} -

-
-
- )} - - )} -
+
+ {setupState ?? + (setupData && ( + + ))} +
- - {step > 0 && ( - - )} - {step < 2 ? ( - - ) : ( - - )} - - -
+ + {step > 0 && ( + + )} + {step < 2 ? ( + + ) : ( + + )} + + +
+ ) } diff --git a/web/default/src/features/profile/components/passkey-card.tsx b/web/default/src/features/profile/components/passkey-card.tsx index 30c9c47e..25cea18b 100644 --- a/web/default/src/features/profile/components/passkey-card.tsx +++ b/web/default/src/features/profile/components/passkey-card.tsx @@ -16,11 +16,12 @@ along with this program. If not, see . For commercial licensing, please contact https://github.com/MAX-API-Next/MAX-API/issues */ -import { useCallback, useMemo, useState } from 'react' +import { useCallback, useState, type ReactElement } from 'react' import { AlertTriangle, KeyRound, Loader2, ShieldAlert } from 'lucide-react' import { useTranslation } from 'react-i18next' import { toast } from 'sonner' import dayjs from '@/lib/dayjs' +import { handleServerError } from '@/lib/handle-server-error' import { AlertDialog, AlertDialogAction, @@ -44,8 +45,8 @@ import { Skeleton } from '@/components/ui/skeleton' import { StatusBadge } from '@/components/status-badge' import { usePasskeyManagement } from '@/features/auth/passkey' import { - SecureVerificationDialog, - useSecureVerification, + checkVerificationMethods, + useSecureVerificationGate, type VerificationMethod, type VerificationMethods, } from '@/features/auth/secure-verification' @@ -54,11 +55,9 @@ interface PasskeyCardProps { loading: boolean } -export function PasskeyCard({ loading: pageLoading }: PasskeyCardProps) { +export function PasskeyCard(props: PasskeyCardProps): ReactElement { const { t } = useTranslation() const [confirmOpen, setConfirmOpen] = useState(false) - const [restrictedMethod, setRestrictedMethod] = - useState(null) const { status, @@ -72,59 +71,37 @@ export function PasskeyCard({ loading: pageLoading }: PasskeyCardProps) { remove, } = usePasskeyManagement() - const { - open: verificationOpen, - setOpen: setVerificationOpen, - methods: verificationMethods, - state: verificationState, - startVerification, - executeVerification, - cancel: cancelVerification, - setCode, - switchMethod, - fetchVerificationMethods, - } = useSecureVerification({ - onSuccess: () => { - setRestrictedMethod(null) - }, - }) - - const dialogMethods = useMemo(() => { - if (!restrictedMethod) return verificationMethods - return { - ...verificationMethods, - has2FA: restrictedMethod === '2fa' && verificationMethods.has2FA, - hasPasskey: - restrictedMethod === 'passkey' && verificationMethods.hasPasskey, - } - }, [restrictedMethod, verificationMethods]) + const { withVerification } = useSecureVerificationGate() - const handleRegister = useCallback(async () => { + const handleRegister = useCallback(async (): Promise => { if (!supported) { toast.info(t('This device does not support Passkey')) return } - const methods = await fetchVerificationMethods() - if (!methods.has2FA) { - // Without 2FA enabled, register directly. The browser-level Passkey prompt - // is itself a strong proof of presence, so no extra verification is needed. - await register() - return + try { + await withVerification(register, { + scope: 'passkey_register', + title: t('Security verification'), + description: t('Confirm your identity before registering a Passkey.'), + }) + } catch { + // The action and verification hook own user-facing error messages. The + // event handler still consumes the rejection to avoid an unhandled + // promise when method discovery or the retried action fails. } + }, [register, supported, t, withVerification]) - setRestrictedMethod('2fa') - await startVerification(register, { - preferredMethod: '2fa', - title: t('Security verification'), - description: t( - 'Confirm your identity with Two-factor Authentication before registering a Passkey.' - ), - }) - }, [fetchVerificationMethods, register, startVerification, supported, t]) - - const handleRemove = useCallback(async () => { - const methods = await fetchVerificationMethods() + const handleRemove = useCallback(async (): Promise => { + let methods: VerificationMethods + try { + methods = await checkVerificationMethods('credentials') + } catch (error) { + handleServerError(error, { + fallback: t('Verification unavailable'), + }) + return + } const required: VerificationMethod | null = methods.has2FA ? '2fa' : methods.hasPasskey @@ -146,46 +123,23 @@ export function PasskeyCard({ loading: pageLoading }: PasskeyCardProps) { } setConfirmOpen(false) - setRestrictedMethod(required) - await startVerification(remove, { - preferredMethod: required, - title: t('Security verification'), - description: t( - 'Confirm your identity before removing this Passkey from your account.' - ), - }) - }, [fetchVerificationMethods, remove, startVerification, t]) - - const handleVerificationCancel = useCallback(() => { - setRestrictedMethod(null) - cancelVerification() - }, [cancelVerification]) - - const handleVerificationOpenChange = useCallback( - (next: boolean) => { - if (!next) { - setRestrictedMethod(null) - } - setVerificationOpen(next) - }, - [setVerificationOpen] - ) - - // Adapt the hook's `Promise` return into the dialog's - // `void | Promise` signature without losing error propagation - // semantics (errors are surfaced via toast inside the hook). - const handleDialogVerify = useCallback( - async (method: VerificationMethod, code?: string) => { - try { - await executeVerification(method, code) - } catch { - // Errors are already surfaced by useSecureVerification via toast. - } - }, - [executeVerification] - ) + try { + await withVerification(remove, { + scope: 'credentials', + preferredMethod: required, + allowedMethods: [required], + forceVerification: true, + title: t('Security verification'), + description: t( + 'Confirm your identity before removing this Passkey from your account.' + ), + }) + } catch { + // The shared verification provider owns user-facing error messages. + } + }, [remove, t, withVerification]) - if (pageLoading || loading) { + if (props.loading || loading) { return ( @@ -205,155 +159,151 @@ export function PasskeyCard({ loading: pageLoading }: PasskeyCardProps) { : t('Not used yet') const showUnsupportedNotice = !supported && !enabled + let backupLabel = t('No backup') + let backupVariant: 'success' | 'warning' | 'neutral' = 'neutral' + if (status?.backup_eligible) { + if (status.backup_state) { + backupLabel = t('Backed up') + backupVariant = 'success' + } else { + backupLabel = t('Not backed up') + backupVariant = 'warning' + } + } return ( - <> - - - - {t('Passkey Login')} - - - {t('Use Passkey to sign in without entering your password.')} - - + + + + {t('Passkey Login')} + + + {t('Use Passkey to sign in without entering your password.')} + + - -
-
-
-
- -
-
-
-

{t('Passkey Authentication')}

+ +
+
+
+
+
+
+
+

{t('Passkey Authentication')}

+ + {status?.backup_eligible !== undefined && ( - {status?.backup_eligible !== undefined && ( - - )} -
-

- {t('Last used:')} {formattedLastUsed} -

+ )}
+

+ {t('Last used:')} {formattedLastUsed} +

+
+ + {!enabled && ( + + )} +
- {!enabled && ( - - )} + {t('Remove Passkey')} + + + + {t('Remove Passkey?')} + + {t( + 'Removing Passkey will require you to sign in with your password next time. You can re-register anytime.' + )} + + + + + {t('Cancel')} + + { + event.preventDefault() + handleRemove() + }} + > + {t('Remove')} + + + +
+ )} - {enabled && ( -
- - - } - > - {removing ? ( - - ) : ( - - )} - {t('Remove Passkey')} - - - - - {t('Remove Passkey?')} - - - {t( - 'Removing Passkey will require you to sign in with your password next time. You can re-register anytime.' - )} - - - - - {t('Cancel')} - - { - event.preventDefault() - handleRemove() - }} - > - {t('Remove')} - - - - -
- )} - - {showUnsupportedNotice && ( -
- -
-

- {t('Passkey not supported on this device')} -

-

- {t( - 'Use a compatible browser or device with biometric authentication or a security key to register a Passkey.' - )} -

-
+ {showUnsupportedNotice && ( +
+
- - - - - +
+ )} +
+ + ) } diff --git a/web/default/src/features/profile/components/profile-security-card.tsx b/web/default/src/features/profile/components/profile-security-card.tsx index c5fe49f8..8ea11d82 100644 --- a/web/default/src/features/profile/components/profile-security-card.tsx +++ b/web/default/src/features/profile/components/profile-security-card.tsx @@ -16,12 +16,19 @@ along with this program. If not, see . For commercial licensing, please contact https://github.com/MAX-API-Next/MAX-API/issues */ -import { Shield, Key, Trash2 } from 'lucide-react' +import { useState } from 'react' +import { Shield, Key, Trash2, LogOut, Loader2 } from 'lucide-react' import { useTranslation } from 'react-i18next' +import { toast } from 'sonner' +import { handleServerError } from '@/lib/handle-server-error' +import { wasSecureVerificationErrorReported } from '@/lib/secure-verification' +import { cn } from '@/lib/utils' import { useDialogs } from '@/hooks/use-dialog' import { Card, CardContent, CardHeader } from '@/components/ui/card' import { Skeleton } from '@/components/ui/skeleton' import { TitledCard } from '@/components/ui/titled-card' +import { useSecureVerificationGate } from '@/features/auth/secure-verification' +import { revokeOtherSessions } from '../api' import type { UserProfile } from '../types' import { AccessTokenDialog } from './dialogs/access-token-dialog' import { ChangePasswordDialog } from './dialogs/change-password-dialog' @@ -44,6 +51,35 @@ export function ProfileSecurityCard({ }: ProfileSecurityCardProps) { const { t } = useTranslation() const dialogs = useDialogs() + const { withVerification } = useSecureVerificationGate() + const [revokingSessions, setRevokingSessions] = useState(false) + + const handleRevokeOtherSessions = async (): Promise => { + setRevokingSessions(true) + try { + const response = await withVerification(revokeOtherSessions, { + scope: 'credentials', + title: t('Security verification'), + description: t( + 'Confirm your identity before signing out your other sessions.' + ), + }) + if (!response) return + if (response.success) { + toast.success(t('Other sessions signed out successfully')) + } else { + toast.error(response.message || t('Failed to sign out other sessions')) + } + } catch (error) { + if (!wasSecureVerificationErrorReported(error)) { + handleServerError(error, { + fallback: t('Failed to sign out other sessions'), + }) + } + } finally { + setRevokingSessions(false) + } + } if (loading) { return ( @@ -70,6 +106,7 @@ export function ProfileSecurityCard({ description: t('Update your password to keep your account secure'), action: () => dialogs.open('password'), variant: 'default' as const, + loading: false, }, { icon: Key, @@ -77,6 +114,17 @@ export function ProfileSecurityCard({ description: t('Generate and manage your API access token'), action: () => dialogs.open('token'), variant: 'default' as const, + loading: false, + }, + { + icon: LogOut, + title: t('Sign Out Other Devices'), + description: t( + 'Keep this session and revoke all other signed-in devices' + ), + action: handleRevokeOtherSessions, + variant: 'default' as const, + loading: revokingSessions, }, { icon: Trash2, @@ -84,6 +132,7 @@ export function ProfileSecurityCard({ description: t('Permanently delete your account and all data'), action: () => dialogs.open('delete'), variant: 'destructive' as const, + loading: false, }, ] @@ -94,30 +143,37 @@ export function ProfileSecurityCard({ description={t('Manage your security settings and account access')} icon={} > -
+
{securityActions.map((item) => (