From 1965f1163b4fd081ff53064602b848b5829bb8ee Mon Sep 17 00:00:00 2001 From: CSCITech Date: Thu, 27 Aug 2026 01:22:16 +0800 Subject: [PATCH 01/70] fix bugs --- constant/midjourney.go | 2 + controller/relay.go | 73 +++++++++++++---- controller/relay_task_billing_test.go | 64 +++++++++++++++ model/billing_settlement.go | 34 +++++++- model/billing_settlement_regression_test.go | 40 ++++++++++ relay/mjproxy_handler.go | 39 ++++++++- relay/mjproxy_handler_test.go | 87 +++++++++++++++++++++ service/billing_session.go | 11 +++ service/billing_session_test.go | 73 +++++++++++++++++ 9 files changed, 400 insertions(+), 23 deletions(-) create mode 100644 controller/relay_task_billing_test.go 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/relay.go b/controller/relay.go index 5bc2486f..9c2e0e3b 100644 --- a/controller/relay.go +++ b/controller/relay.go @@ -451,10 +451,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 +465,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 +536,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 { relayInfo.Billing.Refund(c) } if relayInfo.PersistedTaskID > 0 { @@ -652,23 +665,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 +682,41 @@ func RelayTask(c *gin.Context) { } } -// respondTaskError 统一输出 Task 错误响应（含 429 限流提示改写） +// 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: "billing_settlement_pending", + 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 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/relay_task_billing_test.go b/controller/relay_task_billing_test.go new file mode 100644 index 00000000..ad4d238b --- /dev/null +++ b/controller/relay_task_billing_test.go @@ -0,0 +1,64 @@ +package controller + +import ( + "errors" + "net/http" + "net/http/httptest" + "testing" + + "github.com/MAX-API-Next/MAX-API/constant" + "github.com/MAX-API-Next/MAX-API/dto" + "github.com/MAX-API-Next/MAX-API/model" + "github.com/MAX-API-Next/MAX-API/relay" + relaycommon "github.com/MAX-API-Next/MAX-API/relay/common" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +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 } + +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, "billing_settlement_pending", 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 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})) +} diff --git a/model/billing_settlement.go b/model/billing_settlement.go index 90994a15..52e7aaa7 100644 --- a/model/billing_settlement.go +++ b/model/billing_settlement.go @@ -61,17 +61,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"` 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;not null;default:'applied'"` Attempts int `gorm:"not null;default:0"` LastError string `gorm:"type:text"` NextAttempt int64 `gorm:"index;not null;default:0"` @@ -203,6 +203,34 @@ func ProcessPendingBillingSettlementsOnce() { processPendingBillingSettlements() } +// 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 := DB.Model(&BillingSettlement{}). + Select("id"). + Where("user_id = ?", userID). + Where("funding_delta > 0"). + Where("status IN ?", []string{BillingSettlementStatusPending, BillingSettlementStatusManual}). + Where("operation_key LIKE ?", "request:%:finalize"). + Limit(1). + Find(&record) + if result.Error != nil { + return false, result.Error + } + return result.RowsAffected > 0, 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") diff --git a/model/billing_settlement_regression_test.go b/model/billing_settlement_regression_test.go index 51397d16..71943f10 100644 --- a/model/billing_settlement_regression_test.go +++ b/model/billing_settlement_regression_test.go @@ -52,6 +52,46 @@ func TestConcurrentBillingPreConsumeSelectionAllowsOnlyOneFundingSource(t *testi assert.Contains(t, []string{BillingSettlementSourceWallet, BillingSettlementSourceSubscription}, source) } +func TestHasUnresolvedPositiveFinalizeSettlement(t *testing.T) { + setupUserUpdateTestState(t) + + const blockingUserID = 941 + const nonBlockingUserID = 942 + now := time.Now().Unix() + records := []BillingSettlement{ + { + OperationKey: "request:pending-finalize:finalize", Source: BillingSettlementSourceWallet, + UserID: blockingUserID, FundingDelta: 10, TokenDelta: 10, + Status: BillingSettlementStatusPending, NextAttempt: now, CreatedAt: now, 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) +} + func TestSubscriptionSettlementUsesActualAppliedRefundForToken(t *testing.T) { setupUserUpdateTestState(t) diff --git a/relay/mjproxy_handler.go b/relay/mjproxy_handler.go index 0054d263..2b2b47f9 100644 --- a/relay/mjproxy_handler.go +++ b/relay/mjproxy_handler.go @@ -353,6 +353,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 +421,11 @@ 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()) + if markErr := model.MarkTaskSubmitNeedsReview(task, "Midjourney 任务已被上游接受，但计费结算未完成；请勿重复提交，使用任务 ID 查询或联系管理员处理"); markErr != nil { + common.SysError("mark midjourney billing settlement for manual review error: " + markErr.Error()) + } + return false, &midjourneyBillingSettlementPendingError{taskID: midjourneyTask.MjId, err: err} } if settlementOperationKey != "" { if err := model.ProcessBillingSettlementEffect(settlementOperationKey); err != nil { @@ -470,7 +503,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 +853,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..2a8520a0 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" @@ -17,10 +19,36 @@ import ( "github.com/MAX-API-Next/MAX-API/service" "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() @@ -238,6 +266,65 @@ 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.TaskStatusFailure), storedTask.Status) + assert.EqualValues(t, 10, storedTask.Quota) + assert.Equal(t, midjourneyTask.MjId, storedTask.PrivateData.UpstreamTaskID) + + 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/service/billing_session.go b/service/billing_session.go index 27d10e49..a6bcca1f 100644 --- a/service/billing_session.go +++ b/service/billing_session.go @@ -738,6 +738,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("存在未完成的计费对账，请勿重复提交请求并联系管理员处理"), + types.ErrorCodeInsufficientUserQuota, + 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..a1c5e652 100644 --- a/service/billing_session_test.go +++ b/service/billing_session_test.go @@ -170,6 +170,79 @@ func TestBillingSessionWalletPreConsumeUsesDurableRequestIdentity(t *testing.T) assert.Equal(t, model.BillingSettlementStatusApplied, settlement.Status) } +func TestNewBillingSessionRejectsUnresolvedPositiveFinalizeSettlement(t *testing.T) { + 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"}, + }, 10) + + require.Nil(t, session) + require.NotNil(t, apiErr) + assert.Equal(t, types.ErrorCodeInsufficientUserQuota, apiErr.GetErrorCode()) + assert.EqualValues(t, 100, getUserQuota(t, userID)) + assert.EqualValues(t, 100, getTokenRemainQuota(t, tokenID)) +} + +func TestInsufficientFinalizeBlocksFurtherBillingSessions(t *testing.T) { + 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"}, + } + } + + 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.ErrorCodeInsufficientUserQuota, nextErr.GetErrorCode()) + 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 From e03afa84511fc28f3a0f6a607e3af66ea023eaaa Mon Sep 17 00:00:00 2001 From: CSCITech Date: Thu, 27 Aug 2026 02:20:40 +0800 Subject: [PATCH 02/70] fix bugs --- model/billing_settlement.go | 29 ++++++++---- model/log.go | 56 +++++++++++++--------- service/billing.go | 94 +++++++++++++++++++++++++++++++++++-- service/quota.go | 22 +-------- service/text_quota.go | 11 +---- service/text_quota_test.go | 87 ++++++++++++++++++++++++++++++++++ 6 files changed, 233 insertions(+), 66 deletions(-) diff --git a/model/billing_settlement.go b/model/billing_settlement.go index 52e7aaa7..66b9fe4b 100644 --- a/model/billing_settlement.go +++ b/model/billing_settlement.go @@ -106,16 +106,22 @@ 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"` + 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"` + 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"` } type BillingSettlementInput struct { @@ -752,6 +758,9 @@ func ProcessBillingSettlementEffect(operationKey string) error { 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, + PromptTokens: effect.PromptTokens, CompletionTokens: effect.CompletionTokens, + UseTimeSeconds: effect.UseTimeSeconds, IsStream: effect.IsStream, + RequestId: effect.RequestID, UpstreamRequestId: effect.UpstreamRequestID, }); err != nil { return err } diff --git a/model/log.go b/model/log.go index b27f6882..c5693da8 100644 --- a/model/log.go +++ b/model/log.go @@ -648,16 +648,22 @@ 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 + Group string + Other map[string]interface{} + NodeName string // 任务发起节点；为空时回退当前节点 + PromptTokens int + CompletionTokens int + UseTimeSeconds int + IsStream bool + RequestId string + UpstreamRequestId string } func RecordTaskBillingLog(params RecordTaskBillingLogParams) { @@ -690,18 +696,24 @@ func recordTaskBillingLog(operationKey string, params RecordTaskBillingLogParams } 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 == "" { diff --git a/service/billing.go b/service/billing.go index cd3ac711..23bb91bd 100644 --- a/service/billing.go +++ b/service/billing.go @@ -4,7 +4,9 @@ import ( "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" @@ -56,6 +58,18 @@ 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") + } if relayInfo.Billing != nil { preConsumed := relayInfo.Billing.GetPreConsumedQuota() delta := actualQuota - preConsumed @@ -78,8 +92,26 @@ func SettleBilling(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, actualQuo )) } + if effect != nil { + if settler, ok := relayInfo.Billing.(interface { + SettleWithEffect(int, *model.BillingSettlementEffect) error + }); ok { + if err := settler.SettleWithEffect(actualQuota, effect); err != nil { + 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 +122,67 @@ 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 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) { + if relayInfo == nil { + logger.LogError(ctx, "error settling billing: relayInfo is nil") + return + } + requestID, upstreamRequestID := "", "" + if ctx != nil { + requestID = ctx.GetString(common.RequestIdKey) + upstreamRequestID = ctx.GetString(common.UpstreamRequestIdKey) + } + if requestID == "" { + requestID = relayInfo.RequestId + } + var effect *model.BillingSettlementEffect + if shouldUpdateUsage { + effect = &model.BillingSettlementEffect{ + LogType: model.LogTypeConsume, + Content: params.Content, + ChannelID: params.ChannelId, + ModelName: params.ModelName, + TokenID: params.TokenId, + Group: params.Group, + Other: params.Other, + NodeName: common.NodeName, + UpdateUsage: true, + Quota: int64(params.Quota), + PromptTokens: params.PromptTokens, + CompletionTokens: params.CompletionTokens, + UseTimeSeconds: params.UseTimeSeconds, + IsStream: params.IsStream, + RequestID: requestID, + UpstreamRequestID: upstreamRequestID, + } + } + + effectHandled, err := SettleBillingWithEffect(ctx, relayInfo, params.Quota, effect) + if err != nil { + logger.LogError(ctx, "error settling billing: "+err.Error()) + return + } + if effectHandled { + return + } + if shouldUpdateUsage { + model.UpdateUserUsedQuotaAndRequestCount(relayInfo.UserId, params.Quota) + model.UpdateChannelUsedQuota(relayInfo.ChannelId, params.Quota) + } + model.RecordConsumeLog(ctx, relayInfo.UserId, params) } diff --git a/service/quota.go b/service/quota.go index 91728f1f..4d7ffdd1 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, 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..20b5bb37 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,92 @@ func TestPostTextConsumeQuotaUsesFallbackForNilUsageWithEstimate(t *testing.T) { require.Equal(t, 20, log.PromptTokens) } +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 = ?", "request:failed-finalize-text-request:finalize").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() From 439e406923bea085bb701d41c38a5628379ac1fc Mon Sep 17 00:00:00 2001 From: CSCITech Date: Thu, 27 Aug 2026 09:39:54 +0800 Subject: [PATCH 03/70] fix bugs --- model/log.go | 20 ++++-- model/log_test.go | 108 +++++++++++++++++++++++++++++--- model/option.go | 34 +++++++++- model/option_test.go | 58 +++++++++++++++++ relay/helper/price.go | 83 +++++++++++++++++++++++- relay/helper/price_test.go | 106 +++++++++++++++++++++++++++++++ relay/relay_task.go | 7 +++ service/billing.go | 20 ++++-- service/billing_session_test.go | 18 +++++- service/tiered_settle.go | 5 ++ service/tiered_settle_test.go | 48 ++++++++++++++ 11 files changed, 481 insertions(+), 26 deletions(-) diff --git a/model/log.go b/model/log.go index c5693da8..6bd3f570 100644 --- a/model/log.go +++ b/model/log.go @@ -588,10 +588,17 @@ 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.Context(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) + } otherStr := common.MapToJsonStr(params.Other) // 判断是否需要记录 IP needRecordIp := false @@ -617,7 +624,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 +635,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() @@ -766,6 +773,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..5e04d6c0 100644 --- a/model/log_test.go +++ b/model/log_test.go @@ -136,6 +136,34 @@ 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 + }) + require.NoError(t, LOG_DB.Where("1 = 1").Delete(&Log{}).Error) + t.Cleanup(func() { + require.NoError(t, LOG_DB.Where("1 = 1").Delete(&Log{}).Error) + }) + + RecordConsumeLog(nil, 7, 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) +} + func TestBillingSettlementEffectPayloadSanitizesContent(t *testing.T) { payload, err := billingSettlementEffectPayload(&BillingSettlementEffect{ LogType: LogTypeConsume, @@ -908,14 +936,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 +953,68 @@ 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) { + 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("1 = 1").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 + }) + + const userID = 7001 + 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", + 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 + 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) } func TestWaitPendingLogQuotaDataDrainsEnqueuedWork(t *testing.T) { 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/relay/helper/price.go b/relay/helper/price.go index 659fa3d7..22d76708 100644 --- a/relay/helper/price.go +++ b/relay/helper/price.go @@ -38,6 +38,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 +110,13 @@ 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) { + configuredPreConsumed, err := configuredPreConsumedQuota() + if 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) @@ -92,9 +138,15 @@ func ModelPriceHelper(c *gin.Context, info *relaycommon.RelayInfo, promptTokens var audioCompletionRatio float64 var freeModel bool if !usePrice { - preConsumedTokens := common.Max(promptTokens, common.PreConsumedQuota) + preConsumedTokens, err := addNonNegativeInts(promptTokens, configuredPreConsumed) + if err != nil { + return types.PriceData{}, err + } if meta.MaxTokens != 0 { - preConsumedTokens += meta.MaxTokens + preConsumedTokens, err = addNonNegativeInts(preConsumedTokens, meta.MaxTokens) + if err != nil { + return types.PriceData{}, err + } } var success bool var matchName string @@ -131,7 +183,10 @@ func ModelPriceHelper(c *gin.Context, info *relaycommon.RelayInfo, promptTokens if err != nil { return types.PriceData{}, err } - preConsumedQuota = quota + preConsumedQuota, err = ApplyPreConsumedQuotaFloor(quota, groupRatioInfo.GroupRatio > 0) + if err != nil { + return types.PriceData{}, err + } } // check if free model pre-consume is disabled @@ -152,6 +207,12 @@ func ModelPriceHelper(c *gin.Context, info *relaycommon.RelayInfo, promptTokens } } } + if !freeModel && groupRatioInfo.GroupRatio > 0 { + preConsumedQuota, err = ApplyPreConsumedQuotaFloor(preConsumedQuota, true) + if err != nil { + return types.PriceData{}, err + } + } priceData := types.PriceData{ FreeModel: freeModel, @@ -183,6 +244,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 @@ -247,6 +311,13 @@ func ModelPriceHelperPerCall(c *gin.Context, info *relaycommon.RelayInfo) (types } } } + if !freeModel && groupRatioInfo.GroupRatio > 0 { + var err error + quota, err = ApplyPreConsumedQuotaFloor(quota, true) + if err != nil { + return types.PriceData{}, err + } + } priceData := types.PriceData{ FreeModel: freeModel, @@ -312,6 +383,12 @@ func modelPriceHelperTiered(c *gin.Context, info *relaycommon.RelayInfo, promptT freeModel = true } } + if !freeModel && groupRatioInfo.GroupRatio > 0 { + preConsumedQuota, err = ApplyPreConsumedQuotaFloor(preConsumedQuota, true) + if err != nil { + return types.PriceData{}, err + } + } exprHash := billingexpr.ExprHashString(exprStr) snapshot := &billingexpr.BillingSnapshot{ diff --git a/relay/helper/price_test.go b/relay/helper/price_test.go index b79ff335..6aa44844 100644 --- a/relay/helper/price_test.go +++ b/relay/helper/price_test.go @@ -96,15 +96,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 +127,106 @@ func TestModelPriceHelperPerCallUsesDefaultTaskPrice(t *testing.T) { require.Equal(t, 0.1, priceData.ModelPrice) require.Equal(t, 100, priceData.Quota) } + +func TestModelPriceHelperPerCallAppliesPreConsumedQuota(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) + require.Equal(t, 1000, priceData.Quota) +} + +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 TestModelPriceHelperAddsConfiguredPreConsumedQuotaToPromptEstimate(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()) + info := &relaycommon.RelayInfo{ + OriginModelName: "ratio-preconsume-test", + UserGroup: "default", + UsingGroup: "default", + } + + priceData, err := ModelPriceHelper(ctx, info, 250, &types.TokenCountMeta{}) + + require.NoError(t, err) + require.False(t, priceData.UsePrice) + require.Equal(t, 350, 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/relay_task.go b/relay/relay_task.go index a4450442..3f5a3012 100644 --- a/relay/relay_task.go +++ b/relay/relay_task.go @@ -304,6 +304,13 @@ func RelayTaskSubmit(c *gin.Context, info *relaycommon.RelayInfo) (*TaskSubmitRe } // 7. 预扣费（仅首次 — 重试时 info.Billing 已存在，跳过） + if !info.PriceData.FreeModel && info.PriceData.GroupRatioInfo.GroupRatio > 0 { + info.PriceData.Quota, err = helper.ApplyPreConsumedQuotaFloor(info.PriceData.Quota, true) + if err != nil { + return nil, service.TaskErrorWrapper(err, "pre_consume_quota_error", http.StatusBadRequest) + } + } + if info.Billing == nil && !info.PriceData.FreeModel { info.ForcePreConsume = true if apiErr := service.PreConsumeBilling(c, info.PriceData.Quota, info); apiErr != nil { diff --git a/service/billing.go b/service/billing.go index 23bb91bd..17cef857 100644 --- a/service/billing.go +++ b/service/billing.go @@ -1,6 +1,7 @@ package service import ( + "context" "fmt" "net/http" @@ -17,6 +18,13 @@ const ( BillingSourceSubscription = "subscription" ) +func billingLogContext(ctx *gin.Context) context.Context { + if ctx == nil { + return context.Background() + } + return ctx +} + func validatePreConsumedQuota(preConsumedQuota int, relayInfo *relaycommon.RelayInfo) *types.MaxAPIError { if relayInfo != nil && relayInfo.QuotaClamp != nil { return types.NewErrorWithStatusCode( @@ -70,24 +78,25 @@ func SettleBillingWithEffect(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, 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), )) } @@ -138,8 +147,9 @@ func SettleBillingWithEffect(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, // 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(ctx, "error settling billing: relayInfo is nil") + logger.LogError(logCtx, "error settling billing: relayInfo is nil") return } requestID, upstreamRequestID := "", "" @@ -174,7 +184,7 @@ func settleAndRecordConsume(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, effectHandled, err := SettleBillingWithEffect(ctx, relayInfo, params.Quota, effect) if err != nil { - logger.LogError(ctx, "error settling billing: "+err.Error()) + logger.LogError(logCtx, "error settling billing: "+err.Error()) return } if effectHandled { diff --git a/service/billing_session_test.go b/service/billing_session_test.go index a1c5e652..b865baee 100644 --- a/service/billing_session_test.go +++ b/service/billing_session_test.go @@ -35,9 +35,25 @@ func (f *recordingFundingSource) Settle(delta int) (int64, error) { type recordingBillingSettler struct { preConsumed int reserves []int + settleCalls int } -func (s *recordingBillingSettler) Settle(int) error { return nil } +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 (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 } diff --git a/service/tiered_settle.go b/service/tiered_settle.go index cf36286e..829d8925 100644 --- a/service/tiered_settle.go +++ b/service/tiered_settle.go @@ -9,6 +9,7 @@ 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" + pricehelper "github.com/MAX-API-Next/MAX-API/relay/helper" "github.com/MAX-API-Next/MAX-API/types" "github.com/gin-gonic/gin" ) @@ -128,6 +129,10 @@ func refreshTieredBillingGroup(relayInfo *relaycommon.RelayInfo) (*billingexpr.B if err != nil { return nil, err } + estimatedQuota, err = pricehelper.ApplyPreConsumedQuotaFloor(estimatedQuota, groupRatio > 0) + if err != nil { + return nil, err + } snap.GroupRatio = groupRatio snap.EstimatedQuotaAfterGroup = estimatedQuota return snap, nil diff --git a/service/tiered_settle_test.go b/service/tiered_settle_test.go index 8e3da08c..a9cfe3d0 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" @@ -461,6 +462,53 @@ 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, configured, relayInfo.TieredBillingSnapshot.EstimatedQuotaAfterGroup) +} + func TestPrepareTieredBillingForSelectedGroupFreeToPaidRequiresContext(t *testing.T) { const expr = `tier("base", p)` relayInfo := makeRelayInfo(expr, 0, 1_000_000, 0) From e37eebefe5ad6ca8746f58fd6ada845e7b7b7457 Mon Sep 17 00:00:00 2001 From: CSCITech Date: Thu, 27 Aug 2026 12:59:22 +0800 Subject: [PATCH 04/70] fix: harden quota reservation and settlement --- controller/relay.go | 37 +++++++--- controller/relay_preconsume_test.go | 64 ++++++++++++++++ model/billing_settlement.go | 10 ++- model/log.go | 2 +- relay/helper/price.go | 52 ++++++------- relay/helper/price_test.go | 72 ++++++++++++++++-- relay/mjproxy_handler.go | 15 +++- relay/relay_task.go | 36 +++++---- relay/relay_task_test.go | 25 +++++++ service/alpha_search_billing_test.go | 20 +++-- service/billing.go | 1 + service/billing_session.go | 54 ++++++++++---- service/billing_session_test.go | 72 +++++++++++++++--- service/funding_source.go | 10 ++- service/task_billing.go | 21 +++--- service/text_quota_test.go | 105 +++++++++++++++++++++++++++ service/tiered_settle.go | 19 ++++- service/violation_fee.go | 2 +- types/price_data.go | 2 +- 19 files changed, 510 insertions(+), 109 deletions(-) create mode 100644 controller/relay_preconsume_test.go diff --git a/controller/relay.go b/controller/relay.go index 9c2e0e3b..45445459 100644 --- a/controller/relay.go +++ b/controller/relay.go @@ -67,6 +67,33 @@ 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 Relay(c *gin.Context, relayFormat types.RelayFormat) { requestId := c.GetString(common.RequestIdKey) @@ -157,19 +184,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) diff --git a/controller/relay_preconsume_test.go b/controller/relay_preconsume_test.go new file mode 100644 index 00000000..84c2b6f3 --- /dev/null +++ b/controller/relay_preconsume_test.go @@ -0,0 +1,64 @@ +package controller + +import ( + "testing" + + "github.com/MAX-API-Next/MAX-API/common" + 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/stretchr/testify/require" +) + +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) + }) + } +} diff --git a/model/billing_settlement.go b/model/billing_settlement.go index 66b9fe4b..32c29313 100644 --- a/model/billing_settlement.go +++ b/model/billing_settlement.go @@ -116,6 +116,7 @@ type BillingSettlementEffect struct { 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"` @@ -738,11 +739,12 @@ 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 logQuota != 0 || effect.QuotaIsActual { other := effect.Other if other == nil { other = make(map[string]interface{}) @@ -777,10 +779,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), diff --git a/model/log.go b/model/log.go index 6bd3f570..bd3101aa 100644 --- a/model/log.go +++ b/model/log.go @@ -588,7 +588,7 @@ func RecordConsumeLog(c *gin.Context, userId int, params RecordConsumeLogParams) return } params.Content = common.SanitizePersistedLogContent(params.Content) - logCtx := context.Context(context.Background()) + logCtx := context.Background() if c != nil { logCtx = c } diff --git a/relay/helper/price.go b/relay/helper/price.go index 22d76708..9bb96b3c 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" @@ -110,8 +111,7 @@ 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) { - configuredPreConsumed, err := configuredPreConsumedQuota() - if err != nil { + if _, err := configuredPreConsumedQuota(); err != nil { return types.PriceData{}, err } if promptTokens < 0 { @@ -137,12 +137,10 @@ func ModelPriceHelper(c *gin.Context, info *relaycommon.RelayInfo, promptTokens var audioRatio float64 var audioCompletionRatio float64 var freeModel bool + var err error if !usePrice { - preConsumedTokens, err := addNonNegativeInts(promptTokens, configuredPreConsumed) - if err != nil { - return types.PriceData{}, err - } - if meta.MaxTokens != 0 { + preConsumedTokens := promptTokens + if meta != nil && meta.MaxTokens != 0 { preConsumedTokens, err = addNonNegativeInts(preConsumedTokens, meta.MaxTokens) if err != nil { return types.PriceData{}, err @@ -176,17 +174,14 @@ 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) if err != nil { return types.PriceData{}, err } - preConsumedQuota, err = ApplyPreConsumedQuotaFloor(quota, groupRatioInfo.GroupRatio > 0) - if err != nil { - return types.PriceData{}, err - } + preConsumedQuota = quota } // check if free model pre-consume is disabled @@ -207,7 +202,9 @@ func ModelPriceHelper(c *gin.Context, info *relaycommon.RelayInfo, promptTokens } } } - if !freeModel && groupRatioInfo.GroupRatio > 0 { + // 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 @@ -282,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 @@ -298,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 @@ -311,21 +307,25 @@ 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 { - var err error - quota, err = ApplyPreConsumedQuotaFloor(quota, true) + 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 } @@ -351,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 } @@ -383,7 +383,9 @@ func modelPriceHelperTiered(c *gin.Context, info *relaycommon.RelayInfo, promptT freeModel = true } } - if !freeModel && groupRatioInfo.GroupRatio > 0 { + // 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 diff --git a/relay/helper/price_test.go b/relay/helper/price_test.go index 6aa44844..62488a6c 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 @@ -128,7 +166,7 @@ func TestModelPriceHelperPerCallUsesDefaultTaskPrice(t *testing.T) { require.Equal(t, 100, priceData.Quota) } -func TestModelPriceHelperPerCallAppliesPreConsumedQuota(t *testing.T) { +func TestModelPriceHelperPerCallSeparatesFinalAndPreConsumedQuota(t *testing.T) { gin.SetMode(gin.TestMode) originalQuotaPerUnit := common.QuotaPerUnit originalPreConsumedQuota := common.PreConsumedQuota @@ -155,7 +193,10 @@ func TestModelPriceHelperPerCallAppliesPreConsumedQuota(t *testing.T) { require.NoError(t, err) require.False(t, priceData.FreeModel) - require.Equal(t, 1000, priceData.Quota) + // 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) { @@ -190,7 +231,7 @@ func TestModelPriceHelperAppliesPreConsumedQuotaToFixedPriceModels(t *testing.T) require.Equal(t, 1000, priceData.QuotaToPreConsume) } -func TestModelPriceHelperAddsConfiguredPreConsumedQuotaToPromptEstimate(t *testing.T) { +func TestModelPriceHelperUsesConfiguredPreConsumedQuotaAsPromptEstimateFloor(t *testing.T) { gin.SetMode(gin.TestMode) originalQuotaPerUnit := common.QuotaPerUnit originalPreConsumedQuota := common.PreConsumedQuota @@ -215,11 +256,26 @@ func TestModelPriceHelperAddsConfiguredPreConsumedQuotaToPromptEstimate(t *testi UsingGroup: "default", } - priceData, err := ModelPriceHelper(ctx, info, 250, &types.TokenCountMeta{}) - - require.NoError(t, err) - require.False(t, priceData.UsePrice) - require.Equal(t, 350, priceData.QuotaToPreConsume) + 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) { + 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) { diff --git a/relay/mjproxy_handler.go b/relay/mjproxy_handler.go index 2b2b47f9..db8e3269 100644 --- a/relay/mjproxy_handler.go +++ b/relay/mjproxy_handler.go @@ -258,16 +258,25 @@ func prepareMidjourneyBillingTask(c *gin.Context, info *relaycommon.RelayInfo, a info.Action = action info.OriginModelName = service.CovertMjpActionToModelName(action) info.PriceData = priceData + preConsumedQuota := 0 if !charge || priceData.FreeModel { info.PriceData.Quota = 0 + info.PriceData.QuotaToPreConsume = 0 + } else { + 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) diff --git a/relay/relay_task.go b/relay/relay_task.go index 3f5a3012..c89a1177 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 { @@ -303,21 +306,26 @@ func RelayTaskSubmit(c *gin.Context, info *relaycommon.RelayInfo) (*TaskSubmitRe } } - // 7. 预扣费（仅首次 — 重试时 info.Billing 已存在，跳过） + // 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 { - info.PriceData.Quota, err = helper.ApplyPreConsumedQuotaFloor(info.PriceData.Quota, true) + preConsumedQuota, err = helper.ApplyPreConsumedQuotaFloor(taskEstimateQuota, true) if err != nil { return nil, service.TaskErrorWrapper(err, "pre_consume_quota_error", http.StatusBadRequest) } } + info.PriceData.QuotaToPreConsume = preConsumedQuota + // 7. 预扣费（仅首次 — 重试时 info.Billing 已存在，跳过） if info.Billing == nil && !info.PriceData.FreeModel { 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, service.TaskErrorFromAPIError(apiErr) } } - task, taskErr := ensureTaskPlaceholder(platform, info) + task, taskErr := ensureTaskPlaceholder(platform, info, preConsumedQuota) if taskErr != nil { return nil, taskErr } @@ -362,11 +370,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 } @@ -489,10 +497,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..542068f9 100644 --- a/relay/relay_task_test.go +++ b/relay/relay_task_test.go @@ -192,6 +192,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/service/alpha_search_billing_test.go b/service/alpha_search_billing_test.go index d199a813..f2cbdffb 100644 --- a/service/alpha_search_billing_test.go +++ b/service/alpha_search_billing_test.go @@ -63,8 +63,17 @@ 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 + common.QuotaPerUnit = 500_000 + common.PreConsumedQuota = 20_000 + t.Cleanup(func() { + common.QuotaPerUnit = originalQuotaPerUnit + common.PreConsumedQuota = originalPreConsumedQuota + }) + billing := &recordingBillingSettler{preConsumed: 6_000} info := &relaycommon.RelayInfo{ RelayMode: relayconstant.RelayModeAlphaSearch, @@ -85,12 +94,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 17cef857..35b1bfd2 100644 --- a/service/billing.go +++ b/service/billing.go @@ -173,6 +173,7 @@ func settleAndRecordConsume(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, NodeName: common.NodeName, UpdateUsage: true, Quota: int64(params.Quota), + QuotaIsActual: true, PromptTokens: params.PromptTokens, CompletionTokens: params.CompletionTokens, UseTimeSeconds: params.UseTimeSeconds, diff --git a/service/billing_session.go b/service/billing_session.go index a6bcca1f..21104b0a 100644 --- a/service/billing_session.go +++ b/service/billing_session.go @@ -39,9 +39,15 @@ 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") + type SettlementPreparer interface { PrepareSettlement(actualQuota int) (*model.BillingSettlementInput, error) } @@ -78,6 +84,9 @@ func (s *BillingSession) settleWithEffect(actualQuota int, effect *model.Billing return nil } else { lastErr = err + if errors.Is(err, ErrBillingFundingOutcomeUnknown) { + return err + } } if attempt < maxAttempts-1 { time.Sleep(time.Duration(attempt+1) * 20 * time.Millisecond) @@ -100,6 +109,9 @@ func (s *BillingSession) beginSettleAttempt(actualQuota int, effect *model.Billi 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") } @@ -155,8 +167,17 @@ func (s *BillingSession) settleNonDurableLocked(delta int) error { // 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 + if err != nil { + return fmt.Errorf("%w: %v", ErrBillingFundingOutcomeUnknown, err) + } + return fmt.Errorf("%w: requested=%d applied=%d", ErrBillingFundingOutcomeUnknown, delta, applied) } s.appliedFundingDelta = applied s.fundingSettled = true @@ -168,16 +189,17 @@ 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 + if err != nil { + return fmt.Errorf("%w: %v", ErrBillingFundingOutcomeUnknown, err) + } + return fmt.Errorf("%w: reconciliation requested=%d applied=%d", ErrBillingFundingOutcomeUnknown, needed, applied) } } if s.appliedFundingDelta != int64(delta) { @@ -206,15 +228,18 @@ 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 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 + 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 { @@ -453,6 +478,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 } diff --git a/service/billing_session_test.go b/service/billing_session_test.go index b865baee..975f41ec 100644 --- a/service/billing_session_test.go +++ b/service/billing_session_test.go @@ -32,6 +32,51 @@ func (f *recordingFundingSource) Settle(delta int) (int64, error) { return int64(delta), nil } +type uncertainFundingSource struct { + deltas []int + applied int64 + err error +} + +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) { + 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) { + 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, + } + + err := session.Settle(15) + require.ErrorIs(t, err, ErrBillingFundingOutcomeUnknown) + assert.Equal(t, []int{5}, funding.deltas) + assert.True(t, session.fundingOutcomeUnknown) + + // 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) + }) + } +} + type recordingBillingSettler struct { preConsumed int reserves []int @@ -526,7 +571,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{ @@ -538,14 +583,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 { @@ -575,6 +621,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 { @@ -592,7 +641,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{ @@ -600,19 +649,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/task_billing.go b/service/task_billing.go index eb7b841f..cb2b48c6 100644 --- a/service/task_billing.go +++ b/service/task_billing.go @@ -67,16 +67,17 @@ func BuildTaskSubmissionSettlementEffect(c *gin.Context, info *relaycommon.Relay } 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), + 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), + QuotaIsActual: true, } } diff --git a/service/text_quota_test.go b/service/text_quota_test.go index 20b5bb37..0824de1e 100644 --- a/service/text_quota_test.go +++ b/service/text_quota_test.go @@ -244,6 +244,111 @@ 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 = ?", "request:final-charge-insufficient-request:finalize").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) diff --git a/service/tiered_settle.go b/service/tiered_settle.go index 829d8925..2ebe6ef0 100644 --- a/service/tiered_settle.go +++ b/service/tiered_settle.go @@ -129,9 +129,13 @@ func refreshTieredBillingGroup(relayInfo *relaycommon.RelayInfo) (*billingexpr.B if err != nil { return nil, err } - estimatedQuota, err = pricehelper.ApplyPreConsumedQuotaFloor(estimatedQuota, groupRatio > 0) - if err != nil { - return nil, err + // Alpha Search has a deterministic tool surcharge. Its reservation floor + // is applied after that surcharge is added in PrepareTieredBillingForSelectedGroup. + if relayInfo.RelayMode != relayconstant.RelayModeAlphaSearch { + estimatedQuota, err = pricehelper.ApplyPreConsumedQuotaFloor(estimatedQuota, groupRatio > 0) + if err != nil { + return nil, err + } } snap.GroupRatio = groupRatio snap.EstimatedQuotaAfterGroup = estimatedQuota @@ -167,6 +171,15 @@ func PrepareTieredBillingForSelectedGroup(c *gin.Context, relayInfo *relaycommon ) } if relayInfo.RelayMode == relayconstant.RelayModeAlphaSearch { + 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 { diff --git a/service/violation_fee.go b/service/violation_fee.go index caa232e1..4fcd18c2 100644 --- a/service/violation_fee.go +++ b/service/violation_fee.go @@ -159,7 +159,7 @@ func chargeViolationFeeIfNeeded(ctx *gin.Context, relayInfo *relaycommon.RelayIn 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), + UpdateUsage: true, Quota: int64(feeQuota), QuotaIsActual: true, } if settler, ok := relayInfo.Billing.(interface { diff --git a/types/price_data.go b/types/price_data.go index fe5d63aa..6cd2f92b 100644 --- a/types/price_data.go +++ b/types/price_data.go @@ -26,7 +26,7 @@ type PriceData struct { OtherRatios map[string]float64 UsePrice bool Quota int // 按次计费的最终额度（MJ / Task） - QuotaToPreConsume int // 按量计费的预消耗额度 + QuotaToPreConsume int // 预消耗额度（真实估算额度应用预扣下限后） GroupRatioInfo GroupRatioInfo } From f5caf7ff80599b0caef6d04caa45c4b8124d0619 Mon Sep 17 00:00:00 2001 From: CSCITech Date: Thu, 27 Aug 2026 17:32:32 +0800 Subject: [PATCH 05/70] fix: keep async tasks pollable during settlement recovery --- controller/midjourney_test.go | 10 ++ controller/relay.go | 19 ++- controller/relay_task_billing_test.go | 8 ++ model/billing_settlement.go | 23 +++- relay/mjproxy_handler.go | 3 - relay/mjproxy_handler_test.go | 3 +- relay/relay_task.go | 39 ++++-- relay/relay_task_test.go | 33 +++++ service/billing_session.go | 28 ++++- service/billing_session_test.go | 51 ++++++++ service/midjourney_task.go | 19 ++- service/task_billing.go | 21 ++++ service/task_billing_test.go | 169 ++++++++++++++++++++++++++ service/task_polling.go | 29 +++++ 14 files changed, 429 insertions(+), 26 deletions(-) diff --git a/controller/midjourney_test.go b/controller/midjourney_test.go index d776dfbf..e5147f23 100644 --- a/controller/midjourney_test.go +++ b/controller/midjourney_test.go @@ -175,6 +175,11 @@ func TestMidjourneyFailureRefundsWalletAndTokenThroughDurableSettlement(t *testi require.NoError(t, err) require.True(t, created) require.False(t, refundDuplicate) + _, _, err = model.ApplyBillingSettlementOnce(model.BillingSettlementInput{ + OperationKey: "request:midjourney-refund-request:finalize", Source: model.BillingSettlementSourceWallet, + UserID: user.Id, TokenID: token.Id, + }) + require.NoError(t, err) err = updateMidjourneyTaskFromResponse(context.Background(), mj, dto.MidjourneyDto{ MjId: mj.MjId, Status: "FAILURE", Progress: "100%", FailReason: "provider failed", @@ -258,6 +263,11 @@ func TestApplyMidjourneyTaskResponsePreservesLateCompletionWithoutRefund(t *test require.NoError(t, err) require.True(t, created) require.False(t, refundDuplicate) + _, _, err = model.ApplyBillingSettlementOnce(model.BillingSettlementInput{ + OperationKey: "request:midjourney-late-success-request-" + suffix + ":finalize", Source: model.BillingSettlementSourceWallet, + UserID: user.Id, TokenID: token.Id, + }) + require.NoError(t, err) var settlementCountBefore int64 require.NoError(t, db.Model(&model.BillingSettlement{}).Count(&settlementCountBefore).Error) diff --git a/controller/relay.go b/controller/relay.go index 45445459..4ea51bbf 100644 --- a/controller/relay.go +++ b/controller/relay.go @@ -569,9 +569,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 { @@ -701,6 +703,17 @@ func RelayTask(c *gin.Context) { } } +// 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 + } + return taskErr.Code != "billing_settlement_pending" || !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 diff --git a/controller/relay_task_billing_test.go b/controller/relay_task_billing_test.go index ad4d238b..418818c5 100644 --- a/controller/relay_task_billing_test.go +++ b/controller/relay_task_billing_test.go @@ -54,6 +54,14 @@ func TestFinalizeTaskSubmissionDoesNotWriteSuccessAfterSettlementFailure(t *test assert.Zero(t, writeCount) } +func TestPendingTaskSettlementDoesNotMarkAcceptedTaskAsManualFailure(t *testing.T) { + taskErr := &dto.TaskError{Code: "billing_settlement_pending"} + + assert.False(t, shouldMarkTaskSubmitNeedsReview(taskErr, true, true)) + assert.True(t, shouldMarkTaskSubmitNeedsReview(taskErr, false, true)) + assert.True(t, shouldMarkTaskSubmitNeedsReview(&dto.TaskError{Code: "persist_task_failed"}, true, true)) +} + func TestMidjourneyRelayErrorStatusCodeReportsSettlementPendingAsConflict(t *testing.T) { assert.Equal(t, http.StatusConflict, midjourneyRelayErrorStatusCode(&dto.MidjourneyResponse{ Code: constant.MjRequestError, diff --git a/model/billing_settlement.go b/model/billing_settlement.go index 32c29313..c40f0bc5 100644 --- a/model/billing_settlement.go +++ b/model/billing_settlement.go @@ -238,6 +238,27 @@ func HasUnresolvedPositiveFinalizeSettlement(userID int) (bool, error) { return result.RowsAffected > 0, nil } +// 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") @@ -327,7 +348,7 @@ func ApplyBillingSettlementOnce(input BillingSettlementInput) (appliedFundingDel appliedTokenDelta = effectiveTokenDelta } - if input.TaskID > 0 { + if input.TaskID > 0 && input.TaskQuotaTarget != input.TaskQuota { appliedTaskQuotaTarget := input.TaskQuota + appliedFundingDelta result := tx.Model(&Task{}). Where("id = ? AND quota = ?", input.TaskID, input.TaskQuota). diff --git a/relay/mjproxy_handler.go b/relay/mjproxy_handler.go index db8e3269..0587adab 100644 --- a/relay/mjproxy_handler.go +++ b/relay/mjproxy_handler.go @@ -431,9 +431,6 @@ 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 requires reconciliation: " + err.Error()) - if markErr := model.MarkTaskSubmitNeedsReview(task, "Midjourney 任务已被上游接受，但计费结算未完成；请勿重复提交，使用任务 ID 查询或联系管理员处理"); markErr != nil { - common.SysError("mark midjourney billing settlement for manual review error: " + markErr.Error()) - } return false, &midjourneyBillingSettlementPendingError{taskID: midjourneyTask.MjId, err: err} } if settlementOperationKey != "" { diff --git a/relay/mjproxy_handler_test.go b/relay/mjproxy_handler_test.go index 2a8520a0..3e15adc1 100644 --- a/relay/mjproxy_handler_test.go +++ b/relay/mjproxy_handler_test.go @@ -315,9 +315,10 @@ func TestFinalizeMidjourneySubmissionRejectsSuccessWhenSettlementFails(t *testin var storedTask model.Task require.NoError(t, db.First(&storedTask, task.ID).Error) - assert.Equal(t, model.TaskStatus(model.TaskStatusFailure), storedTask.Status) + 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) diff --git a/relay/relay_task.go b/relay/relay_task.go index c89a1177..43871681 100644 --- a/relay/relay_task.go +++ b/relay/relay_task.go @@ -109,6 +109,26 @@ 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 { + return info.Billing.GetPreConsumedQuota(), nil + } + return targetQuota, nil +} + type taskBillingEstimator interface { EstimateTaskBilling(c *gin.Context, info *relaycommon.RelayInfo) (*types.TaskBillingResult, error) } @@ -318,14 +338,13 @@ func RelayTaskSubmit(c *gin.Context, info *relaycommon.RelayInfo) (*TaskSubmitRe } info.PriceData.QuotaToPreConsume = preConsumedQuota - // 7. 预扣费（仅首次 — 重试时 info.Billing 已存在，跳过） - if info.Billing == nil && !info.PriceData.FreeModel { - info.ForcePreConsume = true - if apiErr := service.PreConsumeBilling(c, preConsumedQuota, info); apiErr != nil { - return nil, service.TaskErrorFromAPIError(apiErr) - } + // 7. 预扣费。重试可能因路由/分组变化得到更高的目标额度，必须通过 + // BillingSession.Reserve 幂等补足，不能继续沿用较小的首次预留。 + actualReservedQuota, reservationErr := reserveTaskQuota(c, info, preConsumedQuota) + if reservationErr != nil { + return nil, reservationErr } - task, taskErr := ensureTaskPlaceholder(platform, info, preConsumedQuota) + task, taskErr := ensureTaskPlaceholder(platform, info, actualReservedQuota) if taskErr != nil { return nil, taskErr } @@ -384,7 +403,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) diff --git a/relay/relay_task_test.go b/relay/relay_task_test.go index 542068f9..e50d9991 100644 --- a/relay/relay_task_test.go +++ b/relay/relay_task_test.go @@ -22,6 +22,39 @@ import ( "github.com/stretchr/testify/require" ) +type recordingTaskReservationBilling struct { + preConsumed int + reserveTo []int +} + +func (b *recordingTaskReservationBilling) Settle(int) error { 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 TestReserveTaskQuotaUsesActualReservationAndExtendsRetries(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 withRelayTaskRateCards(t *testing.T, cards map[string]task_billing_setting.RateCard) { t.Helper() original := task_billing_setting.GetRateCardsCopy() diff --git a/service/billing_session.go b/service/billing_session.go index 21104b0a..bc970f3f 100644 --- a/service/billing_session.go +++ b/service/billing_session.go @@ -119,14 +119,24 @@ 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 { + _, 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") } @@ -298,7 +308,7 @@ func (s *BillingSession) prepareDurableSettlementLocked(actualQuota int, effect if s.relayInfo.IsPlayground { tokenDelta = 0 } - return &model.BillingSettlementInput{ + input := &model.BillingSettlementInput{ OperationKey: "request:" + s.relayInfo.RequestId + ":finalize", Source: source, UserID: userID, @@ -309,7 +319,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) { diff --git a/service/billing_session_test.go b/service/billing_session_test.go index 975f41ec..da677f47 100644 --- a/service/billing_session_test.go +++ b/service/billing_session_test.go @@ -130,6 +130,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{} diff --git a/service/midjourney_task.go b/service/midjourney_task.go index e41ffe5b..95c2d8b3 100644 --- a/service/midjourney_task.go +++ b/service/midjourney_task.go @@ -14,6 +14,21 @@ 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) + billingTask, err := model.GetMidjourneyBillingTask(task.UserId, task.ChannelId, task.MjId) + if err != nil { + return err + } + providerTerminal := responseItem.Status == "SUCCESS" || responseItem.Status == "FAILURE" || + responseItem.FailReason != "" || (responseItem.Status == "" && responseItem.Progress == "100%") + if billingTask != nil && providerTerminal { + pending, settlementErr := taskFinalSettlementPending(billingTask) + if settlementErr != nil { + return settlementErr + } + if pending { + return nil + } + } preStatus := task.Status task.Code = 1 if responseItem.Progress != "" { @@ -81,10 +96,6 @@ func UpdateMidjourneyTaskFromResponse(ctx context.Context, task *model.Midjourne 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 diff --git a/service/task_billing.go b/service/task_billing.go index cb2b48c6..bb48c38d 100644 --- a/service/task_billing.go +++ b/service/task_billing.go @@ -146,6 +146,27 @@ func taskModelName(task *model.Task) string { return task.Properties.OriginModelName } +// taskFinalSettlementPending reports only the funding state of the submission +// finalize operation. A pending settlement effect must not block provider +// polling because the debit/refund itself is already durably applied. +func taskFinalSettlementPending(task *model.Task) (bool, error) { + if task == nil || task.PrivateData.BillingRequestId == "" { + return false, nil + } + status, found, err := model.GetBillingSettlementStatus("request:" + task.PrivateData.BillingRequestId + ":finalize") + 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) + } +} + 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..ba0f77a8 100644 --- a/service/task_billing_test.go +++ b/service/task_billing_test.go @@ -54,6 +54,8 @@ func TestMain(m *testing.M) { &model.BillingSettlement{}, &model.BillingPreConsumeSelection{}, &model.CacheInvalidationTask{}, + &model.Midjourney{}, + &model.MidjourneyBillingClaim{}, ); err != nil { panic("failed to migrate: " + err.Error()) } @@ -107,6 +109,37 @@ func TestSweepTimedOutUnconfirmedSubmitRequiresReviewWithoutRefund(t *testing.T) require.Zero(t, settlementCount) } +func TestSweepTimedOutTaskWaitsForSubmissionSettlement(t *testing.T) { + truncate(t) + 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, + 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) + require.NoError(t, model.DB.Create(&model.BillingSettlement{ + OperationKey: "request:timeout-waits-for-settlement:finalize", + Source: model.BillingSettlementSourceWallet, UserID: 711, + FundingDelta: 10, 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) +} + // --------------------------------------------------------------------------- // Seed helpers // --------------------------------------------------------------------------- @@ -127,6 +160,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 +630,140 @@ 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: "request:task-final-settlement-pending:finalize", + 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 = ?", "request:task-final-settlement-pending:finalize"). + 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: "request:task-final-effect-pending:finalize", + 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: "request:suno-waits-for-submission-settlement:finalize", + 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 = ?", "request:suno-waits-for-submission-settlement:finalize"). + 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 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: "request:mj-waits-for-submission-settlement:finalize", + 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 = ?", "request:mj-waits-for-submission-settlement:finalize"). + 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) +} + // =========================================================================== // RecalculateTaskQuota tests // =========================================================================== diff --git a/service/task_polling.go b/service/task_polling.go index 603157c8..ae19f421 100644 --- a/service/task_polling.go +++ b/service/task_polling.go @@ -56,6 +56,14 @@ func sweepTimedOutTasks(ctx context.Context) { timedOutCount := 0 for _, task := range tasks { + settlementPending, settlementErr := taskFinalSettlementPending(task) + if settlementErr != nil { + logger.LogError(ctx, fmt.Sprintf("sweepTimedOutTasks settlement lookup error for task %s: %v", task.TaskID, settlementErr)) + continue + } + if settlementPending { + continue + } isLegacy := task.SubmitTime > 0 && task.SubmitTime < legacyTaskRefundCutoff isUnconfirmedSubmit := task.PrivateData.AwaitingUpstreamID @@ -227,6 +235,18 @@ func processSunoTaskResponse(ctx context.Context, task *model.Task, responseItem if !taskNeedsUpdate(task, responseItem) { return } + providerTerminal := responseItem.Status == string(model.TaskStatusSuccess) || + responseItem.Status == string(model.TaskStatusFailure) || responseItem.FailReason != "" + if providerTerminal { + pending, err := taskFinalSettlementPending(task) + if err != nil { + logger.LogError(ctx, fmt.Sprintf("Suno task %s settlement lookup failed: %v", task.TaskID, err)) + return + } + if pending { + return + } + } previousStatus := task.Status task.Status = lo.If(model.TaskStatus(responseItem.Status) != "", model.TaskStatus(responseItem.Status)).Else(task.Status) @@ -435,6 +455,15 @@ func updateVideoSingleTask(ctx context.Context, adaptor TaskPollingAdaptor, ch * } } } + if taskResult.Status == string(model.TaskStatusSuccess) || taskResult.Status == string(model.TaskStatusFailure) { + pending, settlementErr := taskFinalSettlementPending(task) + 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 From e4e80dd6b7ab8d1ce6d3b7726d4f123c38aadf84 Mon Sep 17 00:00:00 2001 From: CSCITech Date: Thu, 27 Aug 2026 17:49:25 +0800 Subject: [PATCH 06/70] refactor: centralize task settlement pending guard --- controller/relay_task_billing_test.go | 1 + service/midjourney_task.go | 4 ++-- service/task_billing.go | 10 ++++++++ service/task_polling.go | 33 ++++++++++++--------------- 4 files changed, 28 insertions(+), 20 deletions(-) diff --git a/controller/relay_task_billing_test.go b/controller/relay_task_billing_test.go index 418818c5..53d4353d 100644 --- a/controller/relay_task_billing_test.go +++ b/controller/relay_task_billing_test.go @@ -59,6 +59,7 @@ func TestPendingTaskSettlementDoesNotMarkAcceptedTaskAsManualFailure(t *testing. 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)) } diff --git a/service/midjourney_task.go b/service/midjourney_task.go index 95c2d8b3..bf4ce9b2 100644 --- a/service/midjourney_task.go +++ b/service/midjourney_task.go @@ -20,8 +20,8 @@ func UpdateMidjourneyTaskFromResponse(ctx context.Context, task *model.Midjourne } providerTerminal := responseItem.Status == "SUCCESS" || responseItem.Status == "FAILURE" || responseItem.FailReason != "" || (responseItem.Status == "" && responseItem.Progress == "100%") - if billingTask != nil && providerTerminal { - pending, settlementErr := taskFinalSettlementPending(billingTask) + if billingTask != nil { + pending, settlementErr := taskTerminalSettlementPending(billingTask, providerTerminal) if settlementErr != nil { return settlementErr } diff --git a/service/task_billing.go b/service/task_billing.go index bb48c38d..e51d929e 100644 --- a/service/task_billing.go +++ b/service/task_billing.go @@ -167,6 +167,16 @@ func taskFinalSettlementPending(task *model.Task) (bool, error) { } } +// 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_polling.go b/service/task_polling.go index ae19f421..fca488c0 100644 --- a/service/task_polling.go +++ b/service/task_polling.go @@ -56,7 +56,7 @@ func sweepTimedOutTasks(ctx context.Context) { timedOutCount := 0 for _, task := range tasks { - settlementPending, settlementErr := taskFinalSettlementPending(task) + 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 @@ -237,15 +237,13 @@ func processSunoTaskResponse(ctx context.Context, task *model.Task, responseItem } providerTerminal := responseItem.Status == string(model.TaskStatusSuccess) || responseItem.Status == string(model.TaskStatusFailure) || responseItem.FailReason != "" - if providerTerminal { - pending, err := taskFinalSettlementPending(task) - if err != nil { - logger.LogError(ctx, fmt.Sprintf("Suno task %s settlement lookup failed: %v", task.TaskID, err)) - return - } - if pending { - return - } + 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 @@ -455,14 +453,13 @@ func updateVideoSingleTask(ctx context.Context, adaptor TaskPollingAdaptor, ch * } } } - if taskResult.Status == string(model.TaskStatusSuccess) || taskResult.Status == string(model.TaskStatusFailure) { - pending, settlementErr := taskFinalSettlementPending(task) - if settlementErr != nil { - return fmt.Errorf("load task submission settlement for task %s: %w", task.TaskID, settlementErr) - } - if pending { - return nil - } + 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 From abdfbfb1319115fc8e3a7f35ea4201f7303ce2fa Mon Sep 17 00:00:00 2001 From: CSCITech Date: Thu, 27 Aug 2026 23:59:36 +0800 Subject: [PATCH 07/70] fix bugs --- controller/passkey.go | 36 +- controller/secure_verification.go | 10 +- controller/secure_verification_test.go | 4 + controller/telegram.go | 206 ++++++---- controller/telegram_test.go | 51 +++ controller/twofa.go | 23 +- controller/user.go | 41 ++ controller/user_session_generation_test.go | 63 +++ docs/openapi/api.json | 230 ++++++++++- middleware/auth.go | 48 ++- middleware/auth_session_generation_test.go | 110 ++++++ model/auth_flow.go | 14 +- model/passkey.go | 93 ++++- model/security_credentials_test.go | 227 +++++++++++ model/token.go | 13 + model/twofa.go | 131 +++++-- model/user.go | 188 +++++++-- model/user_cache.go | 34 +- router/api-router.go | 18 +- router/api_router_test.go | 142 ++++++- .../components/authenticated-layout.tsx | 45 ++- .../layout/components/chat-presets-item.tsx | 14 +- web/default/src/features/auth/passkey/api.ts | 14 +- .../passkey/hooks/use-passkey-management.ts | 7 +- .../components/secure-verification-context.ts | 40 ++ .../components/secure-verification-dialog.tsx | 10 +- .../secure-verification-provider.tsx | 69 ++++ .../hooks/use-secure-verification.ts | 95 ++++- .../auth/secure-verification/index.ts | 2 + .../auth/secure-verification/types.ts | 6 +- .../chat/hooks/use-active-chat-key.ts | 25 +- .../overview/overview-dashboard.tsx | 12 +- web/default/src/features/keys/api.ts | 21 +- .../keys/components/api-keys-cells.tsx | 23 +- .../keys/components/api-keys-dialogs.tsx | 10 +- .../components/api-keys-mutate-drawer.tsx | 42 +- .../keys/components/api-keys-provider.tsx | 92 +++-- .../components/data-table-row-actions.tsx | 27 +- web/default/src/features/profile/api.ts | 41 ++ .../dialogs/telegram-bind-dialog.tsx | 212 ++++++++++- .../dialogs/two-fa-setup-dialog.tsx | 358 ++++++++++-------- .../profile/components/passkey-card.tsx | 20 +- .../components/profile-security-card.tsx | 72 +++- web/default/src/i18n/locales/en.json | 29 +- web/default/src/i18n/locales/fr.json | 29 +- web/default/src/i18n/locales/ja.json | 29 +- web/default/src/i18n/locales/ru.json | 29 +- web/default/src/i18n/locales/vi.json | 29 +- web/default/src/i18n/locales/zh.json | 29 +- web/default/src/lib/api.ts | 17 +- .../api-key-plaintext-lifecycle.test.tsx | 104 +++++ .../secure-verification-continuation.test.tsx | 152 ++++++++ 52 files changed, 2819 insertions(+), 567 deletions(-) create mode 100644 controller/telegram_test.go create mode 100644 controller/user_session_generation_test.go create mode 100644 middleware/auth_session_generation_test.go create mode 100644 model/security_credentials_test.go create mode 100644 web/default/src/features/auth/secure-verification/components/secure-verification-context.ts create mode 100644 web/default/src/features/auth/secure-verification/components/secure-verification-provider.tsx create mode 100644 web/default/tests/api-key-plaintext-lifecycle.test.tsx create mode 100644 web/default/tests/secure-verification-continuation.test.tsx diff --git a/controller/passkey.go b/controller/passkey.go index d142ab8f..279f7b8e 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,7 +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 + } + if err := preserveCurrentSessionAfterSecurityChange(c, user.Id, generation); err != nil { common.ApiError(c, err) return } @@ -164,7 +161,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 + } + if err := preserveCurrentSessionAfterSecurityChange(c, user.Id, generation); err != nil { common.ApiError(c, err) return } @@ -370,7 +372,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 } @@ -531,18 +533,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 { diff --git a/controller/secure_verification.go b/controller/secure_verification.go index ffbb2952..ac8b5e34 100644 --- a/controller/secure_verification.go +++ b/controller/secure_verification.go @@ -23,6 +23,8 @@ const ( secureVerificationMethodPassword = "password" secureVerificationScopeAccessToken = "access_token" secureVerificationScopeAccountDelete = "account_delete" + secureVerificationScopeCredentials = "credentials" + secureVerificationScopeAPIToken = "api_token" // PasskeyReadySessionKey means WebAuthn finished and /api/verify can finalize step-up verification. PasskeyReadySessionKey = "secure_passkey_ready_at" // SecureVerificationTimeout 验证有效期（秒） @@ -238,7 +240,8 @@ func UniversalVerify(c *gin.Context) { func isSupportedSecureVerificationScope(scope string) bool { switch scope { - case "", secureVerificationScopeAccessToken, secureVerificationScopeAccountDelete: + case "", secureVerificationScopeAccessToken, secureVerificationScopeAccountDelete, + secureVerificationScopeCredentials, secureVerificationScopeAPIToken: return true default: return false @@ -246,7 +249,10 @@ func isSupportedSecureVerificationScope(scope string) bool { } func passwordVerificationAllowed(scope string) bool { - return scope == secureVerificationScopeAccessToken || scope == secureVerificationScopeAccountDelete + return scope == secureVerificationScopeAccessToken || + scope == secureVerificationScopeAccountDelete || + scope == secureVerificationScopeCredentials || + scope == secureVerificationScopeAPIToken } 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..76d0d324 100644 --- a/controller/secure_verification_test.go +++ b/controller/secure_verification_test.go @@ -116,8 +116,12 @@ func TestSecureVerificationScopeAllowlist(t *testing.T) { require.True(t, isSupportedSecureVerificationScope("")) require.True(t, isSupportedSecureVerificationScope(secureVerificationScopeAccessToken)) require.True(t, isSupportedSecureVerificationScope(secureVerificationScopeAccountDelete)) + require.True(t, isSupportedSecureVerificationScope("credentials")) + require.True(t, isSupportedSecureVerificationScope("api_token")) require.False(t, isSupportedSecureVerificationScope("admin_delete")) require.True(t, passwordVerificationAllowed(secureVerificationScopeAccountDelete)) + require.True(t, passwordVerificationAllowed("credentials")) + require.True(t, passwordVerificationAllowed("api_token")) require.False(t, passwordVerificationAllowed("")) } diff --git a/controller/telegram.go b/controller/telegram.go index 528d463c..d7f3e7ee 100644 --- a/controller/telegram.go +++ b/controller/telegram.go @@ -4,126 +4,182 @@ import ( "crypto/hmac" "crypto/sha256" "encoding/hex" - "io" + "errors" "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 +) + +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"` +} + +func (p telegramAuthPayload) values() url.Values { + values := url.Values{ + "id": {strconv.FormatInt(p.ID, 10)}, + "auth_date": {strconv.FormatInt(p.AuthDate, 10)}, + "hash": {p.Hash}, + "state": {p.State}, + } + for key, value := range map[string]string{ + "first_name": p.FirstName, + "last_name": p.LastName, + "username": p.Username, + "photo_url": p.PhotoURL, + } { + if value != "" { + values.Set(key, value) + } + } + 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") + 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 + } + session := sessions.Default(c) + session.Set(oauthStateSessionKey, state) + if err := session.Save(); err != nil { + _, _ = model.ConsumeAuthFlow(state, model.AuthFlowMatch{ + Purpose: model.AuthFlowPurposeOAuth, + Provider: "telegram", + Intent: model.AuthFlowIntentBind, + UserId: userID, }) + common.ApiError(c, err) return } - telegramId := params["id"][0] - taken, err := model.IsTelegramIdAlreadyTaken(telegramId) - if handleOAuthIdentityLookupError(c, "Telegram", err) { + common.ApiSuccess(c, gin.H{"state": state}) +} + +func TelegramBind(c *gin.Context) { + if !common.TelegramOAuthEnabled { + common.ApiErrorMsg(c, "管理员未开启 Telegram 登录") return } - if taken { - c.JSON(200, gin.H{ - "message": "该 Telegram 账户已被绑定", - "success": false, - }) + var payload telegramAuthPayload + if err := common.DecodeJson(c.Request.Body, &payload); err != nil { + common.ApiErrorI18n(c, i18n.MsgInvalidParams) return } - - session := sessions.Default(c) - id, ok := sessionUserID(session.Get("id")) - if !ok { - common.ApiErrorMsg(c, "用户未登录或登录状态已失效") + if !checkTelegramAuthorization(payload.values(), common.TelegramBotToken) { + c.JSON(http.StatusForbidden, gin.H{"success": false, "message": "无效或已过期的 Telegram 授权"}) return } - user := model.User{Id: id} - if err := user.FillUserById(); 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 } - if user.Id == 0 { - c.JSON(http.StatusOK, gin.H{ - "success": false, - "message": "用户已注销", - }) + 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 } - user.TelegramId = telegramId - if err := user.UpdateFields(false, model.UserUpdateFieldTelegramId); err != nil { - c.JSON(200, gin.H{ - "message": err.Error(), - "success": false, - }) + clearOAuthSessionState(c) + if err := preserveCurrentSessionAfterSecurityChange(c, userID, generation); err != nil { + common.ApiError(c, err) return } - - c.Redirect(302, common.ThemeAwarePath("/console/personal")) + 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} + user := model.User{TelegramId: params.Get("id")} if err := user.FillUserByTelegramId(); handleOAuthUserLookupError(c, err) { return } 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/telegram_test.go b/controller/telegram_test.go new file mode 100644 index 00000000..64be79b0 --- /dev/null +++ b/controller/telegram_test.go @@ -0,0 +1,51 @@ +package controller + +import ( + "crypto/hmac" + "crypto/sha256" + "encoding/hex" + "net/url" + "sort" + "strconv" + "strings" + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +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)}, + } + 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))) + return params +} + +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)) +} diff --git a/controller/twofa.go b/controller/twofa.go index f8b3c0f6..683507e9 100644 --- a/controller/twofa.go +++ b/controller/twofa.go @@ -186,7 +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 + } + if err := preserveCurrentSessionAfterSecurityChange(c, userId, generation); err != nil { common.ApiError(c, err) return } @@ -258,7 +263,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 + } + if err := preserveCurrentSessionAfterSecurityChange(c, userId, generation); err != nil { common.ApiError(c, err) return } @@ -373,7 +383,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 +392,10 @@ func RegenerateBackupCodes(c *gin.Context) { common.SysLog("保存备用码失败: " + err.Error()) return } + if err := preserveCurrentSessionAfterSecurityChange(c, userId, generation); err != nil { + common.ApiError(c, err) + return + } // 记录操作日志 model.RecordLog(userId, model.LogTypeSystem, "重新生成两步验证备用码") @@ -529,7 +544,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..da10ed52 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" @@ -148,6 +149,13 @@ func setupLogin(user *model.User, c *gin.Context) { session.Set("role", user.Role) session.Set("status", user.Status) session.Set("group", user.Group) + session.Set("session_generation", user.SessionGeneration) + if method := loginMethodFromContext(c); method != "unknown" { + session.Set(SecureVerificationSessionKey, time.Now().Unix()) + session.Set(secureVerificationMethodSessionKey, "login:"+method) + session.Set(secureVerificationUserSessionKey, user.Id) + session.Set(secureVerificationScopeSessionKey, secureVerificationScopeCredentials) + } err := session.Save() if err != nil { common.ApiErrorI18n(c, i18n.MsgUserSessionSaveFailed) @@ -185,6 +193,21 @@ 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 + } + if err := preserveCurrentSessionAfterSecurityChange(c, userID, generation); err != nil { + common.ApiError(c, err) + return + } + 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) @@ -890,6 +913,12 @@ func UpdateSelf(c *gin.Context) { common.ApiError(c, err) return } + if updatePassword { + if err := preserveCurrentSessionAfterSecurityChange(c, cleanUser.Id, cleanUser.SessionGeneration); err != nil { + common.ApiError(c, err) + return + } + } c.JSON(http.StatusOK, gin.H{ "success": true, @@ -898,6 +927,18 @@ 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 checkUpdatePassword(originalPassword string, newPassword string, userId int) (updatePassword bool, err error) { if newPassword == "" { return diff --git a/controller/user_session_generation_test.go b/controller/user_session_generation_test.go new file mode 100644 index 00000000..6a765a68 --- /dev/null +++ b/controller/user_session_generation_test.go @@ -0,0 +1,63 @@ +package controller + +import ( + "net/http" + "net/http/httptest" + "testing" + "time" + + "github.com/gin-contrib/sessions" + "github.com/gin-contrib/sessions/cookie" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +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, secureVerificationScopeCredentials) + 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), + }) + }) + + 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}`, inspectRecorder.Body.String()) +} diff --git a/docs/openapi/api.json b/docs/openapi/api.json index 8505cd47..2a537109 100644 --- a/docs/openapi/api.json +++ b/docs/openapi/api.json @@ -1040,11 +1040,87 @@ ] } }, + "/api/oauth/telegram/bind/state": { + "post": { + "operationId": "createTelegramBindState", + "summary": "Create Telegram binding state", + "deprecated": false, + "description": "Requires an authenticated user and credentials-scoped security verification. Returns a short-lived, user-bound, one-time state value.", + "tags": [ + "OAuth" + ], + "parameters": [], + "responses": { + "200": { + "description": "State created", + "headers": {} + }, + "403": { + "description": "Credentials verification is required or expired", + "headers": {} + } + }, + "security": [ + { + "Combination343": [] + }, + { + "Combination1243": [] + } + ] + } + }, "/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 authorization Unix timestamp; must be fresh." + }, + "hash": { + "type": "string" + }, + "state": { + "type": "string" + } + } + } + } + } + }, + "summary": "绑定 Telegram 账户", "deprecated": false, - "description": "🔓 无需鉴权", + "description": "需要登录、credentials 作用域安全验证、5 分钟内的 Telegram 签名，以及由 /api/oauth/telegram/bind/state 生成的一次性 state。", "tags": [ "OAuth" ], @@ -1053,6 +1129,10 @@ "200": { "description": "成功", "headers": {} + }, + "403": { + "description": "签名、state 或安全验证无效、过期或已使用", + "headers": {} } }, "security": [ @@ -1234,6 +1314,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状态", @@ -1286,7 +1396,7 @@ "post": { "summary": "开始注册Passkey", "deprecated": false, - "description": "🔐 需要登录（User权限）", + "description": "需要登录并完成 credentials 作用域安全验证。已有 Passkey 时，本流程表示显式替换。", "tags": [ "用户管理" ], @@ -1311,7 +1421,7 @@ "post": { "summary": "完成注册Passkey", "deprecated": false, - "description": "🔐 需要登录（User权限）", + "description": "需要登录并完成 credentials 作用域安全验证。成功注册或替换后撤销其他会话。", "tags": [ "用户管理" ], @@ -2166,7 +2276,7 @@ "post": { "summary": "设置2FA", "deprecated": false, - "description": "🔐 需要登录（User权限）", + "description": "需要登录并完成 credentials 作用域安全验证。返回待确认的 TOTP 密钥与备用码。", "tags": [ "两步验证" ], @@ -2191,7 +2301,7 @@ "post": { "summary": "启用2FA", "deprecated": false, - "description": "🔐 需要登录（User权限）", + "description": "需要登录并完成 credentials 作用域安全验证。成功启用后撤销其他会话。", "tags": [ "两步验证" ], @@ -2350,14 +2460,16 @@ "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": "验证作用域；验证结果仅可用于完全匹配的敏感操作作用域", "enum": [ "access_token", - "account_delete" + "account_delete", + "credentials", + "api_token" ] } } @@ -2394,12 +2506,14 @@ "name": "scope", "in": "query", "required": false, - "description": "验证作用域；传access_token或account_delete时，无2FA或Passkey的用户可使用密码验证", + "description": "验证作用域；无 2FA 或 Passkey 的用户可使用当前密码完成对应作用域的验证", "schema": { "type": "string", "enum": [ "access_token", - "account_delete" + "account_delete", + "credentials", + "api_token" ] } } @@ -3466,7 +3580,7 @@ "post": { "summary": "创建令牌", "deprecated": false, - "description": "🔐 需要登录（User权限）", + "description": "需要登录并完成 api_token 作用域安全验证。创建接口不返回完整密钥；后续查看仍需调用受 api_token 作用域保护的密钥查看接口。", "tags": [ "令牌管理" ], @@ -3633,6 +3747,96 @@ ] } }, + "/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", + "items": { + "type": "integer" + } + } + } + } + } + } + }, + "responses": { + "200": { + "description": "成功", + "headers": {} + }, + "403": { + "description": "需要 api_token 作用域安全验证", + "headers": {} + } + }, + "security": [ + { + "Combination343": [] + }, + { + "Combination1243": [] + } + ] + } + }, "/api/token/batch": { "post": { "summary": "批量删除令牌", diff --git a/middleware/auth.go b/middleware/auth.go index f7282904..2cd547b8 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)) @@ -111,6 +141,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,9 +278,11 @@ 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 := getSessionUser(session); err == nil && userCache.Status == common.UserStatusEnabled && sessionGenerationMatches(session, userCache) { refreshSessionFromUser(session, userCache) writeSessionUserContext(c, userCache) + } else if err == nil && userCache != nil && !sessionGenerationMatches(session, userCache) { + clearInvalidSession(session) } } c.Next() @@ -278,12 +319,15 @@ func TokenOrUserAuth() func(c *gin.Context) { session := sessions.Default(c) if session.Get("id") != nil { userCache, err := getSessionUser(session) - if err == nil && userCache.Status == common.UserStatusEnabled { + if err == nil && userCache.Status == common.UserStatusEnabled && sessionGenerationMatches(session, userCache) { refreshSessionFromUser(session, userCache) writeSessionUserContext(c, userCache) c.Next() return } + if err == nil && userCache != nil && !sessionGenerationMatches(session, userCache) { + clearInvalidSession(session) + } if err != nil && !errors.Is(err, gorm.ErrRecordNotFound) { common.SysLog("TokenOrUserAuth session user refresh error: " + err.Error()) } diff --git a/middleware/auth_session_generation_test.go b/middleware/auth_session_generation_test.go new file mode 100644 index 00000000..7a18240f --- /dev/null +++ b/middleware/auth_session_generation_test.go @@ -0,0 +1,110 @@ +package middleware + +import ( + "fmt" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "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 TestUserAuthRejectsStaleSessionGeneration(t *testing.T) { + gin.SetMode(gin.TestMode) + oldDB := model.DB + oldRedisEnabled := common.RedisEnabled + common.RedisEnabled = false + + dsn := fmt.Sprintf("file:%s?mode=memory&cache=shared", strings.ReplaceAll(t.Name(), "/", "_")) + db, err := gorm.Open(sqlite.Open(dsn), &gorm.Config{}) + require.NoError(t, err) + model.DB = db + require.NoError(t, db.AutoMigrate(&model.User{})) + if !db.Migrator().HasColumn(&model.User{}, "session_generation") { + require.NoError(t, db.Exec("ALTER TABLE users ADD COLUMN session_generation INTEGER NOT NULL DEFAULT 0").Error) + } + 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 + 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", 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) + + 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..1fe32dde 100644 --- a/model/auth_flow.go +++ b/model/auth_flow.go @@ -109,11 +109,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 +133,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/passkey.go b/model/passkey.go index 9f190187..e38b6728 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 { @@ -178,22 +179,69 @@ func GetPasskeyByCredentialID(credentialID []byte) (*PasskeyCredential, error) { } func UpsertPasskeyCredential(credential *PasskeyCredential) error { - if credential == nil { + if credential == nil || credential.UserID <= 0 || credential.CredentialID == "" { common.SysLog("UpsertPasskeyCredential: nil credential provided") return fmt.Errorf("Passkey 保存失败，请重试") } - return DB.Transaction(func(tx *gorm.DB) error { - // 使用Unscoped()进行硬删除，避免唯一索引冲突 + result := DB.Model(&PasskeyCredential{}). + Where("user_id = ? AND credential_id = ?", credential.UserID, credential.CredentialID). + Select( + "PublicKey", + "AttestationType", + "AAGUID", + "SignCount", + "CloneWarning", + "UserPresent", + "UserVerified", + "BackupEligible", + "BackupState", + "Transports", + "Attachment", + "LastUsedAt", + ). + Updates(credential) + if result.Error != nil { + common.SysLog(fmt.Sprintf("UpsertPasskeyCredential: 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 保存失败，请重试") + } + if count == 1 { + return nil + } + return ErrPasskeyCredentialChanged +} + +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 { @@ -208,3 +256,28 @@ func DeletePasskeyByUserID(userID int) error { } return nil } + +func DeletePasskeyAndBumpSessionGeneration(userID int) (int64, error) { + if userID <= 0 { + return 0, 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 + } + dispatchStagedCacheInvalidation(cacheTask) + return generation, nil +} diff --git a/model/security_credentials_test.go b/model/security_credentials_test.go new file mode 100644 index 00000000..0f3c88ed --- /dev/null +++ b/model/security_credentials_test.go @@ -0,0 +1,227 @@ +package model + +import ( + "errors" + "testing" + "time" + + "github.com/MAX-API-Next/MAX-API/common" + "github.com/stretchr/testify/require" +) + +func setupSecurityCredentialTestState(t *testing.T) { + t.Helper() + setupUserUpdateTestState(t) + require.NoError(t, DB.AutoMigrate( + &PasskeyCredential{}, + &TwoFA{}, + &TwoFABackupCode{}, + &AuthFlow{}, + )) + for _, table := range []string{ + "passkey_credentials", + "two_fas", + "two_fa_backup_codes", + "auth_flows", + } { + require.NoError(t, DB.Exec("DELETE FROM "+table).Error) + } + t.Cleanup(func() { + for _, table := range []string{ + "passkey_credentials", + "two_fas", + "two_fa_backup_codes", + "auth_flows", + } { + DB.Exec("DELETE FROM " + table) + } + }) +} + +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) + + 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 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.Error(t, UpsertPasskeyCredential(&staleUsageUpdate)) + + 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, UpsertPasskeyCredential(&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 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.True(t, errors.Is(err, ErrAuthFlowConsumed) || errors.Is(err, ErrAuthFlowInvalid)) + + var stored User + require.NoError(t, DB.First(&stored, owner.Id).Error) + require.Equal(t, "123456", stored.TelegramId) + require.EqualValues(t, 1, stored.SessionGeneration) +} 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..7751eec3 100644 --- a/model/twofa.go +++ b/model/twofa.go @@ -142,31 +142,29 @@ func (t *TwoFA) IsLocked() bool { // CreateBackupCodes 创建备用码 func CreateBackupCodes(userId int, codes []string) error { return DB.Transaction(func(tx *gorm.DB) error { - // 先删除现有的备用码 - if err := tx.Where("user_id = ?", userId).Delete(&TwoFABackupCode{}).Error; err != nil { + return replaceBackupCodesTx(tx, userId, codes) + }) +} + +func replaceBackupCodesTx(tx *gorm.DB, userID int, codes []string) error { + if err := tx.Unscoped().Where("user_id = ?", userID).Delete(&TwoFABackupCode{}).Error; err != nil { + return err + } + for _, code := range codes { + hashedCode, err := common.HashBackupCode(code) + if err != nil { return err } - - // 创建新的备用码记录 - for _, code := range codes { - hashedCode, err := common.HashBackupCode(code) - if err != nil { - return err - } - - backupCode := TwoFABackupCode{ - UserId: userId, - CodeHash: hashedCode, - IsUsed: false, - } - - if err := tx.Create(&backupCode).Error; err != nil { - return err - } + backupCode := TwoFABackupCode{ + UserId: userID, + CodeHash: hashedCode, + IsUsed: false, } - - return nil - }) + if err := tx.Create(&backupCode).Error; err != nil { + return err + } + } + return nil } // ValidateBackupCode 验证并使用备用码 @@ -223,6 +221,34 @@ 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 { + 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 +257,67 @@ 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 { + 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) { + var generation int64 + var cacheTask CacheInvalidationTask + err := DB.Transaction(func(tx *gorm.DB) error { + 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 := replaceBackupCodesTx(tx, userID, codes); 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..8fd29366 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 } @@ -1091,6 +1095,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 +1546,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 +1599,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/router/api-router.go b/router/api-router.go index a2452325..b0d7fdb6 100644 --- a/router/api-router.go +++ b/router/api-router.go @@ -53,7 +53,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("credentials"), anonymousRequestBodyLimit, controller.GenerateTelegramBindState) + apiRouter.POST("/oauth/telegram/bind", middleware.UserAuth(), middleware.CriticalRateLimit(), middleware.UserCriticalRateLimit(), middleware.DisableCache(), middleware.SecureVerificationRequired("credentials"), 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) @@ -91,9 +92,10 @@ func SetApiRouter(router *gin.Engine) { 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.POST("/sessions/revoke", middleware.CriticalRateLimit(), middleware.UserCriticalRateLimit(), middleware.DisableCache(), middleware.SecureVerificationRequired("credentials"), 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("credentials"), controller.PasskeyRegisterBegin) + selfRoute.POST("/passkey/register/finish", middleware.CriticalRateLimit(), middleware.UserCriticalRateLimit(), middleware.DisableCache(), middleware.SecureVerificationRequired("credentials"), controller.PasskeyRegisterFinish) selfRoute.POST("/passkey/verify/begin", controller.PasskeyVerifyBegin) selfRoute.POST("/passkey/verify/finish", controller.PasskeyVerifyFinish) selfRoute.DELETE("/passkey", controller.PasskeyDelete) @@ -115,8 +117,8 @@ 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/setup", middleware.CriticalRateLimit(), middleware.UserCriticalRateLimit(), middleware.DisableCache(), middleware.SecureVerificationRequired("credentials"), controller.Setup2FA) + selfRoute.POST("/2fa/enable", middleware.CriticalRateLimit(), middleware.UserCriticalRateLimit(), middleware.DisableCache(), middleware.SecureVerificationRequired("credentials"), controller.Enable2FA) selfRoute.POST("/2fa/disable", controller.Disable2FA) selfRoute.POST("/2fa/backup_codes", controller.RegenerateBackupCodes) @@ -291,12 +293,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.POST("/:id/key", middleware.CriticalRateLimit(), middleware.UserCriticalRateLimit(), middleware.DisableCache(), middleware.SecureVerificationRequired("api_token"), controller.GetTokenKey) + tokenRoute.POST("/", middleware.CriticalRateLimit(), middleware.UserCriticalRateLimit(), middleware.DisableCache(), middleware.SecureVerificationRequired("api_token"), 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("/batch/keys", middleware.CriticalRateLimit(), middleware.UserCriticalRateLimit(), middleware.DisableCache(), middleware.SecureVerificationRequired("api_token"), controller.GetTokenKeysBatch) } usageRoute := apiRouter.Group("/usage") diff --git a/router/api_router_test.go b/router/api_router_test.go index be2d0a64..2f93ebd1 100644 --- a/router/api_router_test.go +++ b/router/api_router_test.go @@ -338,8 +338,9 @@ 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" + return schemaOK && valuesOK && len(values) == 4 && + values[0] == "access_token" && values[1] == "account_delete" && + values[2] == "credentials" && values[3] == "api_token" } return false }, "expected supported scope query parameter for GET /api/verify/methods") @@ -355,3 +356,140 @@ 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) + } + + 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.AddUint64(&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 + path string + body string + }{ + {name: "passkey registration", path: "/api/user/passkey/register/begin", body: `{}`}, + {name: "2fa setup", path: "/api/user/2fa/setup", body: `{}`}, + {name: "session revocation", path: "/api/user/sessions/revoke", body: `{}`}, + {name: "telegram bind state", path: "/api/oauth/telegram/bind/state", body: `{}`}, + {name: "telegram bind", path: "/api/oauth/telegram/bind", body: `{}`}, + {name: "api token creation", path: "/api/token/", body: `{"name":"blocked","expired_time":-1,"unlimited_quota":true}`}, + {name: "api token reveal", path: fmt.Sprintf("/api/token/%d/key", token.Id), body: `{}`}, + {name: "api token batch export", 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(http.MethodPost, 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/web/default/src/components/layout/components/authenticated-layout.tsx b/web/default/src/components/layout/components/authenticated-layout.tsx index 0188d0cb..4f5d39eb 100644 --- a/web/default/src/components/layout/components/authenticated-layout.tsx +++ b/web/default/src/components/layout/components/authenticated-layout.tsx @@ -23,6 +23,7 @@ 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' @@ -34,26 +35,28 @@ export function AuthenticatedLayout(props: AuthenticatedLayoutProps) { 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..84b7880f 100644 --- a/web/default/src/components/layout/components/chat-presets-item.tsx +++ b/web/default/src/components/layout/components/chat-presets-item.tsx @@ -40,6 +40,7 @@ import { SidebarMenuSubItem, useSidebar, } from '@/components/ui/sidebar' +import { useSecureVerificationGate } 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 +157,7 @@ export function ChatPresetsItem({ item }: { item: NavChatPresets }) { const { t } = useTranslation() const { chatPresets, serverAddress } = useChatPresets() const { state, isMobile, setOpenMobile } = useSidebar() + const { withVerification } = useSecureVerificationGate() const href = useLocation({ select: (location) => location.href }) const [loadingPresetId, setLoadingPresetId] = useState(null) const loadingPresetIdRef = useRef(null) @@ -181,7 +183,15 @@ export function ChatPresetsItem({ item }: { item: NavChatPresets }) { loadingPresetIdRef.current = preset.id setLoadingPresetId(preset.id) try { - activeKey = await fetchActiveChatKey() + activeKey = + (await withVerification(fetchActiveChatKey, { + scope: 'api_token', + title: t('Security verification'), + description: t( + 'Confirm your identity before sending an API key to a chat client.' + ), + })) || undefined + if (!activeKey) return } catch (error) { const message = error instanceof Error @@ -213,7 +223,7 @@ export function ChatPresetsItem({ item }: { item: NavChatPresets }) { window.open(url, '_blank', 'noopener,noreferrer') setOpenMobile(false) }, - [serverAddress, setOpenMobile, t] + [serverAddress, setOpenMobile, t, withVerification] ) 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..1a7713e7 100644 --- a/web/default/src/features/auth/passkey/api.ts +++ b/web/default/src/features/auth/passkey/api.ts @@ -16,9 +16,14 @@ 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, type ApiRequestConfig } from '@/lib/api' import type { ApiResponse, PasskeyOptionsPayload, PasskeyStatus } from './types' +const credentialSensitiveActionConfig: ApiRequestConfig = { + skipBusinessError: true, + skipErrorHandler: true, +} + export async function getPasskeyStatus(): Promise> { const res = await api.get>('/api/user/passkey') return res.data @@ -28,7 +33,9 @@ export async function beginPasskeyRegistration(): Promise< ApiResponse > { const res = await api.post>( - '/api/user/passkey/register/begin' + '/api/user/passkey/register/begin', + undefined, + credentialSensitiveActionConfig ) return res.data } @@ -38,7 +45,8 @@ export async function finishPasskeyRegistration( ): Promise { const res = await api.post( '/api/user/passkey/register/finish', - payload + payload, + credentialSensitiveActionConfig ) 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..277ed9b4 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 @@ -25,6 +25,7 @@ import { isPasskeySupported as detectPasskeySupport, prepareCredentialCreationOptions, } from '@/lib/passkey' +import { isVerificationRequiredError } from '@/lib/secure-verification' import { beginPasskeyRegistration, deletePasskey, @@ -70,7 +71,8 @@ export function usePasskeyManagement( }, [onStatusChange]) useEffect(() => { - fetchStatus() + const timer = window.setTimeout(() => void fetchStatus(), 0) + return () => window.clearTimeout(timer) }, [fetchStatus]) useEffect(() => { @@ -130,6 +132,9 @@ 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 diff --git a/web/default/src/features/auth/secure-verification/components/secure-verification-context.ts b/web/default/src/features/auth/secure-verification/components/secure-verification-context.ts new file mode 100644 index 00000000..5fc8f113 --- /dev/null +++ b/web/default/src/features/auth/secure-verification/components/secure-verification-context.ts @@ -0,0 +1,40 @@ +/* +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, useContext } from 'react' +import type { StartVerificationOptions } from '../types' + +interface SecureVerificationContextValue { + withVerification: ( + apiCall: () => Promise, + config?: StartVerificationOptions + ) => Promise +} + +export const SecureVerificationContext = + createContext(null) + +export function useSecureVerificationGate() { + const context = useContext(SecureVerificationContext) + if (!context) { + throw new Error( + 'useSecureVerificationGate must be used within SecureVerificationProvider' + ) + } + return context +} 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..455f150f --- /dev/null +++ b/web/default/src/features/auth/secure-verification/components/secure-verification-provider.tsx @@ -0,0 +1,69 @@ +/* +Copyright (C) 2023-2026 MAX-API-Next + +This program is free software: you can redistribute it and/or modify +it under the terms of the GNU Affero General Public License as +published by the Free Software Foundation, either version 3 of the +License, or (at your option) any later version. + +This program is distributed in the hope that it will be useful, +but WITHOUT ANY WARRANTY; without even the implied warranty of +MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +GNU Affero General Public License for more details. + +You should have received a copy of the GNU Affero General Public License +along with this program. If not, see . + +For commercial licensing, please contact https://github.com/MAX-API-Next/MAX-API/issues +*/ +import { useCallback } from 'react' +import { useSecureVerification } from '../hooks/use-secure-verification' +import type { VerificationMethod } from '../types' +import { SecureVerificationContext } from './secure-verification-context' +import { SecureVerificationDialog } from './secure-verification-dialog' + +export function SecureVerificationProvider({ + children, +}: { + children: React.ReactNode +}) { + const { + open, + setOpen, + methods, + state, + executeVerification, + cancel, + setCode, + switchMethod, + withVerification, + } = useSecureVerification() + + const handleVerification = useCallback( + async (method: VerificationMethod, code?: string) => { + try { + await executeVerification(method, code) + } catch { + // The shared verification hook reports the error and keeps retryable + // verification failures inside the dialog. + } + }, + [executeVerification] + ) + + return ( + + {children} + + + ) +} 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..d9ebf1ec 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,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 { useCallback, useMemo, useState } from 'react' +import { useCallback, useEffect, useMemo, useRef, useState } from 'react' import i18next from 'i18next' import { toast } from 'sonner' import { @@ -36,6 +36,11 @@ import type { type ApiCall = (() => Promise) | null +interface PendingVerification { + resolve: (value: unknown | null) => void + reject: (reason?: unknown) => void +} + interface InternalState extends SecureVerificationState { apiCall: ApiCall } @@ -63,7 +68,8 @@ 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 fetchVerificationMethods = useCallback( async (scope?: VerificationScope) => { @@ -76,9 +82,27 @@ export function useSecureVerification( const reset = useCallback(() => { setState(initialState) - setOpen(false) + setDialogOpen(false) }, []) + const settlePendingVerification = useCallback( + (value: unknown | null, error?: unknown) => { + const pending = pendingVerificationRef.current + pendingVerificationRef.current = null + if (!pending) return + if (error !== undefined) { + pending.reject(error) + } else { + pending.resolve(value) + } + }, + [] + ) + + useEffect(() => { + return () => settlePendingVerification(null) + }, [settlePendingVerification]) + const startVerification = useCallback( async ( apiCall: () => Promise, @@ -94,12 +118,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 @@ -118,7 +142,7 @@ export function useSecureVerification( description, scope, })) - setOpen(true) + setDialogOpen(true) return true }, [fetchVerificationMethods, onError] @@ -141,7 +165,20 @@ export function useSecureVerification( try { await verify(actualMethod, code ?? state.code, state.scope) + } catch (error) { + setState((prev) => ({ ...prev, loading: false })) + const message = + error instanceof Error + ? error.message + : i18next.t('Verification failed') + toast.error(message) + onError?.(error) + throw error + } + + try { const result = await state.apiCall() + settlePendingVerification(result) if (successMessage) { toast.success(successMessage) @@ -155,18 +192,28 @@ export function useSecureVerification( return result } catch (error) { + settlePendingVerification(null, error) const message = error instanceof Error ? error.message : i18next.t('Verification failed') toast.error(message) onError?.(error) + reset() throw error } finally { setState((prev) => ({ ...prev, loading: false })) } }, - [state, successMessage, onSuccess, onError, autoReset, reset] + [ + state, + successMessage, + onSuccess, + onError, + autoReset, + reset, + settlePendingVerification, + ] ) const setCode = useCallback((code: string) => { @@ -178,27 +225,47 @@ export function useSecureVerification( }, []) const cancel = useCallback(() => { + settlePendingVerification(null) reset() - }, [reset]) + }, [reset, settlePendingVerification]) + + const setOpen = useCallback( + (nextOpen: boolean) => { + if (!nextOpen) { + cancel() + return + } + setDialogOpen(true) + }, + [cancel] + ) const withVerification = useCallback( - async ( - apiCall: () => Promise, + async ( + apiCall: () => Promise, config: StartVerificationOptions = {} - ) => { + ): Promise => { try { return await apiCall() } catch (error) { if (isVerificationRequiredError(error)) { const info = extractVerificationInfo(error) toast.info(info.message) - await startVerification(apiCall, config) - return null + const started = await startVerification(apiCall, config) + if (!started) return null + + settlePendingVerification(null) + return await new Promise((resolve, reject) => { + pendingVerificationRef.current = { + resolve: (value) => resolve(value as T | null), + reject, + } + }) } throw error } }, - [startVerification] + [settlePendingVerification, startVerification] ) 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..00923b7b 100644 --- a/web/default/src/features/auth/secure-verification/index.ts +++ b/web/default/src/features/auth/secure-verification/index.ts @@ -20,3 +20,5 @@ export * from './api' export * from './types' export * from './hooks/use-secure-verification' export * from './components/secure-verification-dialog' +export * from './components/secure-verification-context' +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..37a9aa96 100644 --- a/web/default/src/features/auth/secure-verification/types.ts +++ b/web/default/src/features/auth/secure-verification/types.ts @@ -17,7 +17,11 @@ 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' export interface VerificationMethods { has2FA: boolean 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..fb37c8dd 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 { useSecureVerificationGate } from '@/features/auth/secure-verification' import { fetchTokenKey, getApiKeys } from '@/features/keys/api' import { API_KEY_STATUS } from '@/features/keys/constants' @@ -46,12 +49,28 @@ export async function fetchActiveChatKey() { */ export function useActiveChatKey(enabled: boolean) { const userId = useAuthStore((state) => state.auth.user?.id) + const { withVerification } = useSecureVerificationGate() + const queryFn = useCallback(async () => { + const key = await withVerification(fetchActiveChatKey, { + scope: 'api_token', + title: i18next.t('Security verification'), + description: i18next.t( + 'Confirm your identity before sending an API key to a chat client.' + ), + }) + if (!key) { + throw new Error(i18next.t('API key verification was cancelled')) + } + return key + }, [withVerification]) 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, + 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..67be6c49 100644 --- a/web/default/src/features/dashboard/components/overview/overview-dashboard.tsx +++ b/web/default/src/features/dashboard/components/overview/overview-dashboard.tsx @@ -54,6 +54,7 @@ import { CardStaggerContainer, CardStaggerItem, } from '@/components/page-transition' +import { useSecureVerificationGate } from '@/features/auth/secure-verification' import { fetchTokenKey, getApiKeys } from '@/features/keys/api' import type { ApiKey } from '@/features/keys/types' import { @@ -381,6 +382,7 @@ function RequestPreview(props: { }) { const { t } = useTranslation() const shouldReduceMotion = useReducedMotion() + const { withVerification } = useSecureVerificationGate() const [isCopying, setIsCopying] = useState(false) const { copyToClipboard } = useCopyToClipboard({ notify: false }) const previewCurl = buildCurlCommand({ @@ -391,10 +393,18 @@ function RequestPreview(props: { const previewLines = previewCurl.split('\n') const handleCopyRequest = async () => { if (!props.example.keyId || isCopying) return + const keyId = props.example.keyId setIsCopying(true) try { - const result = await fetchTokenKey(props.example.keyId) + const result = await withVerification(() => fetchTokenKey(keyId), { + scope: 'api_token', + title: t('Security verification'), + description: t( + 'Confirm your identity before copying a request with your API key.' + ), + }) + 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')) diff --git a/web/default/src/features/keys/api.ts b/web/default/src/features/keys/api.ts index ca5bceaa..c860b91b 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, type ApiRequestConfig } from '@/lib/api' import type { ApiKey, ApiResponse, @@ -30,6 +30,11 @@ import type { // API Key Management // ============================================================================ +const apiTokenSensitiveActionConfig: ApiRequestConfig = { + skipBusinessError: true, + skipErrorHandler: true, +} + // Get paginated API keys list export async function getApiKeys( params: GetApiKeysParams = {} @@ -63,7 +68,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, apiTokenSensitiveActionConfig) return res.data } @@ -102,7 +107,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, + apiTokenSensitiveActionConfig + ) return res.data } @@ -112,6 +121,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 }, + apiTokenSensitiveActionConfig + ) 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..d8578719 100644 --- a/web/default/src/features/keys/components/api-keys-cells.tsx +++ b/web/default/src/features/keys/components/api-keys-cells.tsx @@ -19,7 +19,6 @@ For commercial licensing, please contact https://github.com/MAX-API-Next/MAX-API import { useState, useCallback } from 'react' import { Check, Copy, Loader2 } from 'lucide-react' import { useTranslation } from 'react-i18next' -import { toast } from 'sonner' import { copyToClipboard } from '@/lib/copy-to-clipboard' import { Button } from '@/components/ui/button' import { @@ -41,6 +40,7 @@ export function ApiKeyCell({ apiKey }: { apiKey: ApiKey }) { const { resolveRealKey, resolvedKeys, + clearResolvedKey, loadingKeys, copiedKeyId, markKeyCopied, @@ -56,24 +56,21 @@ export function ApiKeyCell({ apiKey }: { apiKey: ApiKey }) { (open: boolean) => { 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 realKey = await resolveRealKey(apiKey.id) if (realKey) { const ok = await copyToClipboard(realKey) if (ok) markKeyCopied(apiKey.id) } - }, [resolvedFullKey, resolveRealKey, apiKey.id, markKeyCopied, t]) + }, [resolveRealKey, apiKey.id, markKeyCopied]) return (
@@ -122,12 +119,6 @@ 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} /> } 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..b2822517 100644 --- a/web/default/src/features/keys/components/api-keys-dialogs.tsx +++ b/web/default/src/features/keys/components/api-keys-dialogs.tsx @@ -22,7 +22,8 @@ import { useApiKeys } from './api-keys-provider' import { CCSwitchDialog } from './dialogs/cc-switch-dialog' export function ApiKeysDialogs() { - const { open, setOpen, currentRow, resolvedKey } = useApiKeys() + const { open, setOpen, currentRow, resolvedKey, setResolvedKey } = + useApiKeys() return ( <> @@ -34,7 +35,12 @@ export function ApiKeysDialogs() { !isOpen && setOpen(null)} + onOpenChange={(isOpen) => { + 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..de782194 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 @@ -99,7 +99,7 @@ export function ApiKeysMutateDrawer({ }: ApiKeyMutateDrawerProps) { const { t } = useTranslation() const isUpdate = !!currentRow - const { triggerRefresh } = useApiKeys() + const { triggerRefresh, withApiTokenVerification } = useApiKeys() const { status } = useStatus() const [isSubmitting, setIsSubmitting] = useState(false) const [advancedOpen, setAdvancedOpen] = useState(false) @@ -329,26 +329,30 @@ 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 + const successCount = await withApiTokenVerification(async () => { + const count = data.tokenCount || 1 + let created = 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) { + created++ + } else { + toast.error(result.message || t(ERROR_MESSAGES.CREATE_FAILED)) + break + } } - } - if (successCount > 0) { + return created + }) + + if (successCount !== null && successCount > 0) { toast.success( t('Successfully created {{count}} API Key(s)', { count: successCount, 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..681e93ff 100644 --- a/web/default/src/features/keys/components/api-keys-provider.tsx +++ b/web/default/src/features/keys/components/api-keys-provider.tsx @@ -20,6 +20,7 @@ import React, { useState, useCallback, useRef, useEffect } from 'react' import { useTranslation } from 'react-i18next' import { toast } from 'sonner' import useDialogState from '@/hooks/use-dialog' +import { useSecureVerificationGate } from '@/features/auth/secure-verification' import { fetchTokenKey, fetchTokenKeysBatch } from '../api' import { ERROR_MESSAGES } from '../constants' import { type ApiKey, type ApiKeysDialogType } from '../types' @@ -33,18 +34,24 @@ 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 }) { const { t } = useTranslation() + const { withVerification } = useSecureVerificationGate() const [open, setOpen] = useDialogState(null) const [currentRow, setCurrentRow] = useState(null) const [refreshTrigger, setRefreshTrigger] = useState(0) @@ -52,7 +59,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) @@ -71,18 +80,40 @@ export function ApiKeysProvider({ children }: { children: React.ReactNode }) { setRefreshTrigger((prev) => prev + 1) }, []) + const withApiTokenVerification = useCallback( + (apiCall: () => Promise) => + withVerification(apiCall, { + scope: 'api_token', + title: t('Security verification'), + description: t( + 'Confirm your identity before creating or revealing API keys.' + ), + }), + [t, withVerification] + ) + 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)) @@ -100,47 +131,48 @@ 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) => { + 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)) - return {} } finally { - for (const id of uncachedIds) { + for (const id of ids) { setLoadingKeys((prev) => { const next = { ...prev } delete next[id] @@ -149,7 +181,7 @@ export function ApiKeysProvider({ children }: { children: React.ReactNode }) { } } }, - [resolvedKeys, t] + [t, withApiTokenVerification] ) return ( @@ -166,9 +198,11 @@ export function ApiKeysProvider({ children }: { children: React.ReactNode }) { resolveRealKey, resolveRealKeysBatch, resolvedKeys, + clearResolvedKey, loadingKeys, copiedKeyId, markKeyCopied, + withApiTokenVerification, }} > {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..e8b12e43 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 @@ -87,33 +87,12 @@ 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) @@ -214,7 +193,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')) @@ -242,7 +221,7 @@ export function DataTableRowActions({ { - const realKey = getCachedRealKey() + const realKey = await resolveRealKey(apiKey.id) if (!realKey) return const connStr = encodeChannelConnectionInfo( realKey, diff --git a/web/default/src/features/profile/api.ts b/web/default/src/features/profile/api.ts index b4624736..0b95c0af 100644 --- a/web/default/src/features/profile/api.ts +++ b/web/default/src/features/profile/api.ts @@ -141,6 +141,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..fd5ed7e4 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,13 @@ 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 { useCallback, useEffect, useRef, useState } from 'react' +import { RefreshCw, Send } from 'lucide-react' import { useTranslation } from 'react-i18next' +import { toast } from 'sonner' +import { cn } from '@/lib/utils' import { Alert, AlertDescription } from '@/components/ui/alert' +import { Button } from '@/components/ui/button' import { Dialog, DialogContent, @@ -26,6 +30,13 @@ import { DialogHeader, DialogTitle, } from '@/components/ui/dialog' +import { Spinner } from '@/components/ui/spinner' +import { useSecureVerificationGate } from '@/features/auth/secure-verification' +import { + bindTelegramAccount, + createTelegramBindState, + type TelegramAuthorizationPayload, +} from '../../api' // ============================================================================ // Telegram Bind Dialog Component @@ -42,10 +53,153 @@ export function TelegramBindDialog({ open, onOpenChange, botName, + onSuccess, }: TelegramBindDialogProps) { const { t } = useTranslation() + const { withVerification } = useSecureVerificationGate() + const widgetContainerRef = useRef(null) + const [callbackName] = useState( + () => `telegramAuthCallback_${Math.random().toString(36).slice(2)}` + ) + const [bindState, setBindState] = useState('') + const [status, setStatus] = useState< + 'idle' | 'loading' | 'ready' | 'binding' | 'error' + >('idle') + const [errorMessage, setErrorMessage] = useState('') + const [retryVersion, setRetryVersion] = useState(0) + + const initializeBinding = useCallback(async () => { + setStatus('loading') + setErrorMessage('') + try { + const response = await withVerification(createTelegramBindState, { + scope: 'credentials', + title: t('Security verification'), + description: t( + 'Confirm your identity before binding a Telegram account.' + ), + }) + if (!response) { + setStatus('idle') + return + } + if (!response.success || !response.data?.state) { + throw new Error( + response.message || t('Failed to initialize Telegram binding') + ) + } + setBindState(response.data.state) + } catch (error) { + setErrorMessage( + error instanceof Error + ? error.message + : t('Failed to initialize Telegram binding') + ) + setStatus('error') + } + }, [t, withVerification]) + + const resetBindingState = useCallback(() => { + setBindState('') + setStatus('idle') + setErrorMessage('') + }, []) + + const handleOpenChange = useCallback( + (nextOpen: boolean) => { + if (!nextOpen) resetBindingState() + onOpenChange(nextOpen) + }, + [onOpenChange, resetBindingState] + ) + + useEffect(() => { + if (!open) return + const timer = window.setTimeout(() => void initializeBinding(), 0) + return () => window.clearTimeout(timer) + }, [initializeBinding, open, retryVersion]) + + useEffect(() => { + if (!open || !bindState || !widgetContainerRef.current) return + + const container = widgetContainerRef.current + const windowCallbacks = window as unknown as Record + container.replaceChildren() + + windowCallbacks[callbackName] = async ( + authorization: TelegramAuthorizationPayload + ) => { + setStatus('binding') + setErrorMessage('') + try { + const response = await withVerification( + () => bindTelegramAccount({ ...authorization, state: bindState }), + { + scope: 'credentials', + title: t('Security verification'), + description: t( + 'Confirm your identity before binding a Telegram account.' + ), + } + ) + if (!response) { + setStatus('ready') + return + } + if (!response.success) { + throw new Error( + response.message || t('Failed to bind Telegram account') + ) + } + toast.success(t('Telegram account bound successfully')) + onSuccess() + handleOpenChange(false) + } catch (error) { + setErrorMessage( + error instanceof Error + ? error.message + : t('Failed to bind Telegram account') + ) + setStatus('error') + } + } + + const script = document.createElement('script') + script.async = true + script.src = 'https://telegram.org/js/telegram-widget.js?22' + script.setAttribute('data-telegram-login', botName.replace(/^@/, '')) + script.setAttribute('data-size', 'large') + script.setAttribute('data-userpic', 'false') + script.setAttribute('data-onauth', `${callbackName}(user)`) + script.onload = () => setStatus('ready') + script.onerror = () => { + setErrorMessage(t('Failed to load Telegram Login Widget')) + setStatus('error') + } + container.appendChild(script) + + return () => { + delete windowCallbacks[callbackName] + container.replaceChildren() + } + }, [ + bindState, + botName, + callbackName, + handleOpenChange, + onSuccess, + open, + t, + withVerification, + ]) + + const retry = () => { + setBindState('') + setRetryVersion((version) => version + 1) + } + return ( - + {t('Bind Telegram Account')} @@ -54,19 +208,19 @@ export function TelegramBindDialog({ -
+
- + {t( - 'You will be redirected to Telegram to complete the binding process.' + 'Telegram will ask you to confirm the account before it is bound.' )} -
-
- +
+
+
@@ -81,13 +235,43 @@ 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..1cc65f1d 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,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, useEffect, useCallback } from 'react' +import { useState, useEffect, useCallback, useRef } from 'react' import { Loader2 } from 'lucide-react' import { QRCodeSVG } from 'qrcode.react' import { useTranslation } from 'react-i18next' @@ -35,6 +35,11 @@ import { import { Input } from '@/components/ui/input' import { Label } from '@/components/ui/label' import { CopyButton } from '@/components/copy-button' +import { + SecureVerificationDialog, + useSecureVerification, + type VerificationMethod, +} from '@/features/auth/secure-verification' import type { TwoFASetupData } from '../../types' // ============================================================================ @@ -58,6 +63,18 @@ export function TwoFASetupDialog({ const [step, setStep] = useState(0) const [setupData, setSetupData] = useState(null) const [code, setCode] = useState('') + const setupInFlightRef = useRef(false) + const { + open: verificationOpen, + setOpen: setVerificationOpen, + methods: verificationMethods, + state: verificationState, + executeVerification, + cancel: cancelVerification, + setCode: setVerificationCode, + switchMethod, + withVerification, + } = useSecureVerification() const stepLabels = [ t('Scan QR Code'), t('Save Backup Codes'), @@ -65,26 +82,47 @@ export function TwoFASetupDialog({ ] const handleSetup = useCallback(async () => { + if (setupInFlightRef.current) return + setupInFlightRef.current = true + setInitializing(true) try { - setInitializing(true) - const response = await setup2FA() - - if (response.success && response.data) { - setSetupData(response.data) - setStep(0) - } else { - toast.error(response.message || t('Failed to setup 2FA')) + const response = await withVerification(setup2FA, { + scope: 'credentials', + title: t('Security verification'), + description: t( + 'Confirm your identity before setting up Two-factor Authentication.' + ), + }) + if (!response) { onOpenChange(false) + return } + if (!response.success || !response.data) { + throw new Error(response.message || t('Failed to setup 2FA')) + } + 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')) onOpenChange(false) } finally { + setupInFlightRef.current = false setInitializing(false) } - }, [onOpenChange, t]) + }, [onOpenChange, t, withVerification]) + + const handleVerification = useCallback( + async (method: VerificationMethod, value?: string) => { + try { + await executeVerification(method, value) + } catch { + // The shared verification hook reports the error. + } + }, + [executeVerification] + ) const handleEnable = async () => { if (!code) { @@ -94,7 +132,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!')) @@ -116,9 +161,6 @@ export function TwoFASetupDialog({ const handleOpenChange = (open: boolean) => { if (!loading && !initializing) { - if (open && !setupData) { - handleSetup() - } if (!open) { setStep(0) setCode('') @@ -130,156 +172,170 @@ 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]} - - + <> + + + + {t('Setup Two-Factor Authentication')} + + {t('Step')} {step + 1} {t('of 3:')} {stepLabels[step]} + + -
- {initializing ? ( -
-
-
- {t('Setting up 2FA...')} +
+ {initializing ? ( +
+
+
+ {t('Setting up 2FA...')} +
-
- ) : !setupData ? ( -
-
- {t('Failed to load setup data')} + ) : !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} - + ) : ( + <> + {/* 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} + +
+
-
-
- )} + )} - {/* 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} -
- ))} + {/* 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')} +
- - {t('Copy All Codes')} - -
- )} + )} - {/* 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')} -

+ {/* 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' + )} +

+
-
- )} - - )} -
+ )} + + )} +
- - {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..372408db 100644 --- a/web/default/src/features/profile/components/passkey-card.tsx +++ b/web/default/src/features/profile/components/passkey-card.tsx @@ -83,6 +83,7 @@ export function PasskeyCard({ loading: pageLoading }: PasskeyCardProps) { setCode, switchMethod, fetchVerificationMethods, + withVerification, } = useSecureVerification({ onSuccess: () => { setRestrictedMethod(null) @@ -105,23 +106,12 @@ export function PasskeyCard({ loading: pageLoading }: PasskeyCardProps) { 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 - } - - setRestrictedMethod('2fa') - await startVerification(register, { - preferredMethod: '2fa', + await withVerification(register, { + scope: 'credentials', title: t('Security verification'), - description: t( - 'Confirm your identity with Two-factor Authentication before registering a Passkey.' - ), + description: t('Confirm your identity before registering a Passkey.'), }) - }, [fetchVerificationMethods, register, startVerification, supported, t]) + }, [register, supported, t, withVerification]) const handleRemove = useCallback(async () => { const methods = await fetchVerificationMethods() 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..c40106b8 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,17 @@ 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 { 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 +49,31 @@ export function ProfileSecurityCard({ }: ProfileSecurityCardProps) { const { t } = useTranslation() const dialogs = useDialogs() + const { withVerification } = useSecureVerificationGate() + const [revokingSessions, setRevokingSessions] = useState(false) + + const handleRevokeOtherSessions = async () => { + 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 { + // The shared verification flow reports transport and verification errors. + } finally { + setRevokingSessions(false) + } + } if (loading) { return ( @@ -78,12 +108,23 @@ export function ProfileSecurityCard({ action: () => dialogs.open('token'), variant: 'default' as const, }, + { + 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, title: t('Delete Account'), description: t('Permanently delete your account and all data'), action: () => dialogs.open('delete'), variant: 'destructive' as const, + loading: false, }, ] @@ -94,30 +135,37 @@ export function ProfileSecurityCard({ description={t('Manage your security settings and account access')} icon={} > -
+
{securityActions.map((item) => ( + ), +})) + +mock.module('../src/components/ui/dialog', () => ({ + Dialog: TestContainer, + DialogContent: TestContainer, + DialogDescription: TestContainer, + DialogHeader: TestContainer, + DialogTitle: TestContainer, +})) + +mock.module('../src/components/ui/spinner', () => ({ + Spinner: () => , +})) + +const { TelegramBindDialog } = await import( + '../src/features/profile/components/dialogs/telegram-bind-dialog' +) +const testEnv = createReactTestEnvironment() + +beforeAll(() => testEnv.setup()) +afterEach(() => cleanup()) +afterAll(() => testEnv.teardown()) + +describe('TelegramBindDialog widget lifecycle', () => { + test('keeps the active widget mounted across callback-only rerenders', async () => { + const firstOpenChange = mock(() => undefined) + const secondOpenChange = mock(() => undefined) + const firstSuccess = mock(() => undefined) + const secondSuccess = mock(() => undefined) + + const view = render( + + ) + + await waitFor(() => + assert.equal(createTelegramBindState.mock.calls.length, 1) + ) + const initialScript = await waitFor(() => { + const script = document.querySelector( + 'script[data-onauth]' + ) + assert.ok(script) + return script + }) + const callbackName = initialScript + .getAttribute('data-onauth') + ?.replace(/\(user\)$/, '') + assert.ok(callbackName) + + view.rerender( + + ) + + await act(async () => Promise.resolve()) + assert.equal( + document.querySelector('script[data-onauth]'), + initialScript + ) + + const callback = ( + window as unknown as Record< + string, + (authorization: { + id: number + auth_date: number + hash: string + }) => Promise + > + )[callbackName] + assert.equal(typeof callback, 'function') + + await act(async () => { + await callback({ id: 123, auth_date: 456, hash: 'telegram-hash' }) + }) + + assert.equal(firstSuccess.mock.calls.length, 0) + assert.equal(firstOpenChange.mock.calls.length, 0) + assert.equal(secondSuccess.mock.calls.length, 1) + assert.deepEqual(secondOpenChange.mock.calls[0], [false]) + }) +}) From 1016df745944484b9daf0170943e0c8b2ce2dc3e Mon Sep 17 00:00:00 2001 From: CSCITech Date: Fri, 28 Aug 2026 03:04:28 +0800 Subject: [PATCH 09/70] fix: address PR 38 review follow-ups --- model/security_credentials_test.go | 4 + model/twofa.go | 32 ++- service/billing.go | 71 ++--- service/billing_session_test.go | 63 ++++- service/task_billing.go | 30 ++- service/task_billing_test.go | 28 ++ service/violation_fee.go | 19 +- .../hooks/use-secure-verification.ts | 4 + .../chat/hooks/use-active-chat-key.ts | 1 + .../keys/components/api-keys-cells.tsx | 13 +- .../dialogs/two-fa-setup-dialog.tsx | 248 ++++++++---------- .../profile/components/passkey-card.tsx | 18 +- web/default/src/i18n/locales/ru.json | 2 +- web/default/src/i18n/locales/vi.json | 2 +- .../active-chat-key-verification.test.tsx | 71 +++++ .../api-key-plaintext-lifecycle.test.tsx | 67 ++++- ...sskey-registration-error-handling.test.tsx | 100 +++++++ .../two-fa-setup-shared-verification.test.tsx | 82 ++++++ 18 files changed, 649 insertions(+), 206 deletions(-) create mode 100644 web/default/tests/active-chat-key-verification.test.tsx create mode 100644 web/default/tests/passkey-registration-error-handling.test.tsx create mode 100644 web/default/tests/two-fa-setup-shared-verification.test.tsx diff --git a/model/security_credentials_test.go b/model/security_credentials_test.go index 4267b32c..e2beeea0 100644 --- a/model/security_credentials_test.go +++ b/model/security_credentials_test.go @@ -126,6 +126,10 @@ func TestPasskeyAndTwoFASecurityChangesBumpSessionGeneration(t *testing.T) { 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)) generation, err = DisableTwoFAAndBumpSessionGeneration(user.Id) require.NoError(t, err) diff --git a/model/twofa.go b/model/twofa.go index 7751eec3..59ece320 100644 --- a/model/twofa.go +++ b/model/twofa.go @@ -141,20 +141,32 @@ 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 { - return replaceBackupCodesTx(tx, userId, codes) + return replaceBackupCodeHashesTx(tx, userId, hashes) }) } -func replaceBackupCodesTx(tx *gorm.DB, userID int, codes []string) error { - if err := tx.Unscoped().Where("user_id = ?", userID).Delete(&TwoFABackupCode{}).Error; err != nil { - return err - } +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 err + return nil, err } + hashes = append(hashes, hashedCode) + } + return hashes, nil +} + +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, @@ -292,9 +304,13 @@ func (t *TwoFA) EnableAndBumpSessionGeneration() (int64, error) { } 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 { + err = DB.Transaction(func(tx *gorm.DB) error { var enabledCount int64 if err := tx.Model(&TwoFA{}). Where("user_id = ? AND is_enabled = ?", userID, true). @@ -304,7 +320,7 @@ func ReplaceBackupCodesAndBumpSessionGeneration(userID int, codes []string) (int if enabledCount != 1 { return ErrTwoFANotEnabled } - if err := replaceBackupCodesTx(tx, userID, codes); err != nil { + if err := replaceBackupCodeHashesTx(tx, userID, hashes); err != nil { return err } var err error diff --git a/service/billing.go b/service/billing.go index 35b1bfd2..bd862d80 100644 --- a/service/billing.go +++ b/service/billing.go @@ -25,6 +25,45 @@ func billingLogContext(ctx *gin.Context) context.Context { 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( + params model.RecordConsumeLogParams, + requestID string, + upstreamRequestID string, + updateUsage bool, +) *model.BillingSettlementEffect { + return &model.BillingSettlementEffect{ + LogType: model.LogTypeConsume, + Content: params.Content, + ChannelID: params.ChannelId, + ModelName: params.ModelName, + TokenID: params.TokenId, + 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, + } +} + func validatePreConsumedQuota(preConsumedQuota int, relayInfo *relaycommon.RelayInfo) *types.MaxAPIError { if relayInfo != nil && relayInfo.QuotaClamp != nil { return types.NewErrorWithStatusCode( @@ -152,36 +191,8 @@ func settleAndRecordConsume(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, logger.LogError(logCtx, "error settling billing: relayInfo is nil") return } - requestID, upstreamRequestID := "", "" - if ctx != nil { - requestID = ctx.GetString(common.RequestIdKey) - upstreamRequestID = ctx.GetString(common.UpstreamRequestIdKey) - } - if requestID == "" { - requestID = relayInfo.RequestId - } - var effect *model.BillingSettlementEffect - if shouldUpdateUsage { - effect = &model.BillingSettlementEffect{ - LogType: model.LogTypeConsume, - Content: params.Content, - ChannelID: params.ChannelId, - ModelName: params.ModelName, - TokenID: params.TokenId, - Group: params.Group, - Other: params.Other, - NodeName: common.NodeName, - UpdateUsage: true, - Quota: int64(params.Quota), - QuotaIsActual: true, - PromptTokens: params.PromptTokens, - CompletionTokens: params.CompletionTokens, - UseTimeSeconds: params.UseTimeSeconds, - IsStream: params.IsStream, - RequestID: requestID, - UpstreamRequestID: upstreamRequestID, - } - } + requestID, upstreamRequestID := billingEffectRequestIDs(ctx, relayInfo) + effect := newConsumeBillingSettlementEffect(params, requestID, upstreamRequestID, shouldUpdateUsage) effectHandled, err := SettleBillingWithEffect(ctx, relayInfo, params.Quota, effect) if err != nil { diff --git a/service/billing_session_test.go b/service/billing_session_test.go index da677f47..a6288f9e 100644 --- a/service/billing_session_test.go +++ b/service/billing_session_test.go @@ -83,6 +83,12 @@ type recordingBillingSettler struct { settleCalls int } +type recordingEffectBillingSettler struct { + recordingBillingSettler + actualQuota int + effect *model.BillingSettlementEffect +} + func TestSettleAndRecordConsumeHandlesNilContextWithoutPanic(t *testing.T) { originalLogConsumeEnabled := common.LogConsumeEnabled common.LogConsumeEnabled = false @@ -98,6 +104,50 @@ func TestSettleAndRecordConsumeHandlesNilContextWithoutPanic(t *testing.T) { }) } +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, + 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.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 (s *recordingBillingSettler) Settle(int) error { s.settleCalls++; return nil } func (s *recordingBillingSettler) Refund(*gin.Context) {} func (s *recordingBillingSettler) NeedsRefund() bool { return false } @@ -108,6 +158,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 nil +} + func TestBillingSessionPrepareSettlementPersistsZeroDeltaTaskEffect(t *testing.T) { session := &BillingSession{ relayInfo: &relaycommon.RelayInfo{ @@ -567,6 +623,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", @@ -596,7 +654,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) diff --git a/service/task_billing.go b/service/task_billing.go index e51d929e..b6859eb1 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,19 +67,22 @@ 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), - QuotaIsActual: true, - } + useTimeSeconds := 0 + if !info.StartTime.IsZero() { + useTimeSeconds = int(time.Since(info.StartTime).Seconds()) + } + requestID, upstreamRequestID := billingEffectRequestIDs(c, info) + return newConsumeBillingSettlementEffect(model.RecordConsumeLogParams{ + ChannelId: info.ChannelId, + ModelName: info.OriginModelName, + Quota: quota, + Content: logContent, + TokenId: info.TokenId, + UseTimeSeconds: useTimeSeconds, + IsStream: info.IsStream, + Group: info.UsingGroup, + Other: other, + }, requestID, upstreamRequestID, true) } // LogTaskConsumption records task usage for legacy/custom funding paths that diff --git a/service/task_billing_test.go b/service/task_billing_test.go index ba0f77a8..c24a659b 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" @@ -78,6 +79,33 @@ 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") + 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.GreaterOrEqual(t, effect.UseTimeSeconds, 3) + assert.True(t, effect.IsStream) +} + func TestSweepTimedOutUnconfirmedSubmitRequiresReviewWithoutRefund(t *testing.T) { truncate(t) originalTimeout := constant.TaskTimeoutMinutes diff --git a/service/violation_fee.go b/service/violation_fee.go index 4fcd18c2..045ea2ed 100644 --- a/service/violation_fee.go +++ b/service/violation_fee.go @@ -155,12 +155,19 @@ 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), QuotaIsActual: true, - } + requestID, upstreamRequestID := billingEffectRequestIDs(ctx, relayInfo) + effect := newConsumeBillingSettlementEffect(model.RecordConsumeLogParams{ + ChannelId: relayInfo.ChannelId, + ModelName: relayInfo.OriginModelName, + TokenName: tokenName, + Quota: feeQuota, + Content: "Violation fee charged", + TokenId: relayInfo.TokenId, + UseTimeSeconds: int(useTimeSeconds), + IsStream: relayInfo.IsStream, + Group: relayInfo.UsingGroup, + Other: other, + }, requestID, upstreamRequestID, true) if settler, ok := relayInfo.Billing.(interface { SettleWithEffect(int, *model.BillingSettlementEffect) error 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 d9ebf1ec..605e1e31 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 @@ -240,6 +240,10 @@ export function useSecureVerification( [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, 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 fb37c8dd..02bee671 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 @@ -70,6 +70,7 @@ export function useActiveChatKey(enabled: boolean) { enabled: enabled && Boolean(userId), staleTime: Number.POSITIVE_INFINITY, gcTime: 0, + retry: false, refetchOnWindowFocus: false, refetchOnReconnect: false, }) 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 d8578719..ce450af8 100644 --- a/web/default/src/features/keys/components/api-keys-cells.tsx +++ b/web/default/src/features/keys/components/api-keys-cells.tsx @@ -19,6 +19,7 @@ For commercial licensing, please contact https://github.com/MAX-API-Next/MAX-API import { useState, useCallback } from 'react' import { Check, Copy, Loader2 } from 'lucide-react' import { useTranslation } from 'react-i18next' +import { toast } from 'sonner' import { copyToClipboard } from '@/lib/copy-to-clipboard' import { Button } from '@/components/ui/button' import { @@ -66,11 +67,14 @@ export function ApiKeyCell({ apiKey }: { apiKey: ApiKey }) { const handleCopy = useCallback(async () => { const realKey = await resolveRealKey(apiKey.id) - if (realKey) { - const ok = await copyToClipboard(realKey) - if (ok) markKeyCopied(apiKey.id) + if (!realKey) return + const ok = await copyToClipboard(realKey) + if (ok) { + markKeyCopied(apiKey.id) + } else { + toast.error(t('Failed to copy to clipboard')) } - }, [resolveRealKey, apiKey.id, markKeyCopied]) + }, [resolveRealKey, apiKey.id, markKeyCopied, t]) return (
@@ -120,6 +124,7 @@ export function ApiKeyCell({ apiKey }: { apiKey: ApiKey }) { className='size-7 shrink-0' onClick={handleCopy} disabled={isLoading} + aria-label={t('Copy API key')} /> } > 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 1cc65f1d..a2a639e1 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 @@ -35,11 +35,7 @@ import { import { Input } from '@/components/ui/input' import { Label } from '@/components/ui/label' import { CopyButton } from '@/components/copy-button' -import { - SecureVerificationDialog, - useSecureVerification, - type VerificationMethod, -} from '@/features/auth/secure-verification' +import { useSecureVerificationGate } from '@/features/auth/secure-verification' import type { TwoFASetupData } from '../../types' // ============================================================================ @@ -64,17 +60,7 @@ export function TwoFASetupDialog({ const [setupData, setSetupData] = useState(null) const [code, setCode] = useState('') const setupInFlightRef = useRef(false) - const { - open: verificationOpen, - setOpen: setVerificationOpen, - methods: verificationMethods, - state: verificationState, - executeVerification, - cancel: cancelVerification, - setCode: setVerificationCode, - switchMethod, - withVerification, - } = useSecureVerification() + const { withVerification } = useSecureVerificationGate() const stepLabels = [ t('Scan QR Code'), t('Save Backup Codes'), @@ -113,17 +99,6 @@ export function TwoFASetupDialog({ } }, [onOpenChange, t, withVerification]) - const handleVerification = useCallback( - async (method: VerificationMethod, value?: string) => { - try { - await executeVerification(method, value) - } catch { - // The shared verification hook reports the error. - } - }, - [executeVerification] - ) - const handleEnable = async () => { if (!code) { toast.error(t('Please enter the verification code')) @@ -177,6 +152,31 @@ export function TwoFASetupDialog({ return () => window.clearTimeout(timer) }, [open, setupData, handleSetup]) + const renderSetupState = () => { + if (initializing) { + return ( +
+
+
+ {t('Setting up 2FA...')} +
+
+ ) + } + if (!setupData) { + return ( +
+
+ {t('Failed to load setup data')} +
+
+ ) + } + return null + } + + const setupState = renderSetupState() + return ( <> @@ -189,112 +189,100 @@ export function TwoFASetupDialog({
- {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} - + {setupState ?? + (setupData && ( + <> + {/* 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} + +
+
-
-
- )} + )} - {/* 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} -
- ))} + {/* 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')} +
- - {t('Copy All Codes')} - -
- )} + )} - {/* 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' - )} -

+ {/* 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' + )} +

+
-
- )} - - )} + )} + + ))}
@@ -326,16 +314,6 @@ export function TwoFASetupDialog({
- ) } diff --git a/web/default/src/features/profile/components/passkey-card.tsx b/web/default/src/features/profile/components/passkey-card.tsx index 372408db..0471f3ef 100644 --- a/web/default/src/features/profile/components/passkey-card.tsx +++ b/web/default/src/features/profile/components/passkey-card.tsx @@ -106,11 +106,19 @@ export function PasskeyCard({ loading: pageLoading }: PasskeyCardProps) { return } - await withVerification(register, { - scope: 'credentials', - title: t('Security verification'), - description: t('Confirm your identity before registering a Passkey.'), - }) + try { + await withVerification(register, { + scope: 'credentials', + title: t('Security verification'), + description: t('Confirm your identity before registering a Passkey.'), + }) + } catch (error) { + // 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. + // eslint-disable-next-line no-console + console.error('[Passkey] Verification continuation failed', error) + } }, [register, supported, t, withVerification]) const handleRemove = useCallback(async () => { diff --git a/web/default/src/i18n/locales/ru.json b/web/default/src/i18n/locales/ru.json index 6e360f31..22be7041 100644 --- a/web/default/src/i18n/locales/ru.json +++ b/web/default/src/i18n/locales/ru.json @@ -5478,7 +5478,7 @@ "Failed to initialize Telegram binding": "Не удалось инициализировать привязку Telegram", "Failed to load Telegram Login Widget": "Не удалось загрузить виджет входа Telegram", "Failed to sign out other sessions": "Не удалось завершить другие сеансы", - "Keep this session and revoke all other signed-in devices": "Сохранить текущий сеанс и отозвать все остальные входы", + "Keep this session and revoke all other signed-in devices": "Сохранить текущий сеанс и завершить сеансы на всех остальных устройствах", "No verification method is available. Add a password or sign in again with a linked account.": "Нет доступного способа подтверждения. Добавьте пароль или снова войдите через привязанную учетную запись.", "Other sessions signed out successfully": "Другие сеансы успешно завершены", "Preparing Telegram authorization...": "Подготовка авторизации Telegram...", diff --git a/web/default/src/i18n/locales/vi.json b/web/default/src/i18n/locales/vi.json index 545beacc..da355e82 100644 --- a/web/default/src/i18n/locales/vi.json +++ b/web/default/src/i18n/locales/vi.json @@ -5486,6 +5486,6 @@ "Start Telegram authorization": "Bắt đầu ủy quyền Telegram", "Telegram account bound successfully": "Đã liên kết tài khoản Telegram thành công", "Telegram will ask you to confirm the account before it is bound.": "Telegram sẽ yêu cầu bạn xác nhận tài khoản trước khi liên kết.", - "Verification unavailable": "Không thể xác minh" + "Verification unavailable": "Xác minh không khả dụng" } } diff --git a/web/default/tests/active-chat-key-verification.test.tsx b/web/default/tests/active-chat-key-verification.test.tsx new file mode 100644 index 00000000..7629da8e --- /dev/null +++ b/web/default/tests/active-chat-key-verification.test.tsx @@ -0,0 +1,71 @@ +/* +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 type { ReactNode } from 'react' +import { createReactTestEnvironment } from '@/test/react' +import { cleanup, renderHook, waitFor } from '@testing-library/react' +import { QueryClient, QueryClientProvider } from '@tanstack/react-query' +import { afterAll, afterEach, beforeAll, beforeEach, mock, test } from 'bun:test' +import assert from 'node:assert/strict' + +const withVerification = mock(async () => null) + +mock.module('../src/stores/auth-store', () => ({ + useAuthStore: (selector: (state: unknown) => unknown) => + selector({ auth: { user: { id: 77 } } }), +})) + +mock.module('../src/features/auth/secure-verification', () => ({ + useSecureVerificationGate: () => ({ withVerification }), +})) + +mock.module('../src/features/keys/api', () => ({ + getApiKeys: mock(async () => ({ success: true, data: { items: [] } })), + fetchTokenKey: mock(async () => ({ success: false })), +})) + +mock.module('i18next', () => ({ + default: { t: (key: string) => key }, +})) + +const { useActiveChatKey } = await import( + '../src/features/chat/hooks/use-active-chat-key' +) +const testEnv = createReactTestEnvironment() + +beforeAll(() => testEnv.setup()) +beforeEach(() => withVerification.mockClear()) +afterEach(() => cleanup()) +afterAll(() => testEnv.teardown()) + +test('does not retry a cancelled chat API-key verification', async () => { + const queryClient = new QueryClient({ + defaultOptions: { + queries: { retry: 2, retryDelay: 0 }, + }, + }) + const wrapper = ({ children }: { children: ReactNode }) => ( + {children} + ) + + const { result } = renderHook(() => useActiveChatKey(true), { wrapper }) + + await waitFor(() => assert.equal(result.current.isError, true)) + assert.equal(withVerification.mock.calls.length, 1) + queryClient.clear() +}) diff --git a/web/default/tests/api-key-plaintext-lifecycle.test.tsx b/web/default/tests/api-key-plaintext-lifecycle.test.tsx index 0978d2fb..cf871b00 100644 --- a/web/default/tests/api-key-plaintext-lifecycle.test.tsx +++ b/web/default/tests/api-key-plaintext-lifecycle.test.tsx @@ -18,7 +18,14 @@ For commercial licensing, please contact https://github.com/MAX-API-Next/MAX-API */ import type { ReactNode } from 'react' import { createReactTestEnvironment } from '@/test/react' -import { act, cleanup, renderHook, waitFor } from '@testing-library/react' +import { + act, + cleanup, + fireEvent, + render, + renderHook, + waitFor, +} from '@testing-library/react' import { afterAll, afterEach, @@ -43,6 +50,8 @@ const fetchTokenKey = mock( resolveTokenKeyRequest = resolve }) ) +const copyToClipboard = mock(async () => false) +const toastError = mock(() => undefined) mock.module('../src/features/auth/secure-verification', () => ({ useSecureVerificationGate: () => ({ @@ -54,6 +63,14 @@ mock.module('react-i18next', () => ({ useTranslation: () => ({ t: (key: string) => key }), })) +mock.module('sonner', () => ({ + toast: { error: toastError }, +})) + +mock.module('../src/lib/copy-to-clipboard', () => ({ + copyToClipboard, +})) + mock.module('../src/features/keys/api', () => ({ fetchTokenKey, fetchTokenKeysBatch: mock(async () => ({ success: true, data: { keys: {} } })), @@ -62,6 +79,9 @@ mock.module('../src/features/keys/api', () => ({ const { ApiKeysProvider, useApiKeys } = await import( '../src/features/keys/components/api-keys-provider' ) +const { ApiKeyCell } = await import( + '../src/features/keys/components/api-keys-cells' +) const testEnv = createReactTestEnvironment() function wrapper(props: { children: ReactNode }) { @@ -72,10 +92,15 @@ beforeAll(() => testEnv.setup()) beforeEach(() => { fetchTokenKey.mockClear() + copyToClipboard.mockClear() + toastError.mockClear() resolveTokenKeyRequest = undefined }) -afterEach(() => cleanup()) +afterEach(async () => { + cleanup() + await new Promise((resolve) => setTimeout(resolve, 50)) +}) afterAll(() => testEnv.teardown()) @@ -101,4 +126,42 @@ describe('API key plaintext lifecycle', () => { assert.equal(result.current.resolvedKeys[101], undefined) }) + + test('reports clipboard failures without marking the API key as copied', async () => { + fetchTokenKey.mockImplementationOnce(async () => ({ + success: true, + data: { key: 'plaintext-copy-key' }, + })) + + const view = render( + + + + ) + + fireEvent.click(view.getByRole('button', { name: 'Copy API key' })) + + await waitFor(() => assert.equal(copyToClipboard.mock.calls.length, 1)) + assert.equal(toastError.mock.calls.length, 1) + assert.equal(toastError.mock.calls[0]?.[0], 'Failed to copy to clipboard') + }) }) diff --git a/web/default/tests/passkey-registration-error-handling.test.tsx b/web/default/tests/passkey-registration-error-handling.test.tsx new file mode 100644 index 00000000..7174068d --- /dev/null +++ b/web/default/tests/passkey-registration-error-handling.test.tsx @@ -0,0 +1,100 @@ +/* +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 { createReactTestEnvironment } from '@/test/react' +import { cleanup, fireEvent, render, waitFor } from '@testing-library/react' +import { afterAll, afterEach, beforeAll, beforeEach, mock, test } from 'bun:test' +import assert from 'node:assert/strict' + +const withVerification = mock(async () => { + throw new Error('verification method lookup failed') +}) + +mock.module('react-i18next', () => ({ + useTranslation: () => ({ t: (key: string) => key }), +})) + +mock.module('sonner', () => ({ + toast: { + error: mock(() => undefined), + info: mock(() => undefined), + success: mock(() => undefined), + }, +})) + +mock.module('../src/features/auth/passkey', () => ({ + usePasskeyManagement: () => ({ + status: null, + loading: false, + registering: false, + removing: false, + supported: true, + enabled: false, + lastUsed: null, + register: mock(async () => true), + remove: mock(async () => true), + }), +})) + +mock.module('../src/features/auth/secure-verification', () => ({ + SecureVerificationDialog: () => null, + useSecureVerification: () => ({ + open: false, + setOpen: mock(() => undefined), + methods: { + has2FA: false, + hasPasskey: false, + hasPassword: true, + passkeySupported: true, + }, + state: { method: null, loading: false, code: '' }, + startVerification: mock(async () => true), + executeVerification: mock(async () => undefined), + cancel: mock(() => undefined), + setCode: mock(() => undefined), + switchMethod: mock(() => undefined), + fetchVerificationMethods: mock(async () => ({ + has2FA: false, + hasPasskey: false, + hasPassword: true, + passkeySupported: true, + })), + withVerification, + }), +})) + +const { PasskeyCard } = await import( + '../src/features/profile/components/passkey-card' +) +const testEnv = createReactTestEnvironment() + +beforeAll(() => testEnv.setup()) +beforeEach(() => { + withVerification.mockClear() +}) +afterEach(() => cleanup()) +afterAll(() => testEnv.teardown()) + +test('handles a rejected Passkey verification continuation', async () => { + const view = render() + + fireEvent.click(view.getByRole('button', { name: 'Enable Passkey' })) + + await waitFor(() => assert.equal(withVerification.mock.calls.length, 1)) + await new Promise((resolve) => setTimeout(resolve, 0)) +}) diff --git a/web/default/tests/two-fa-setup-shared-verification.test.tsx b/web/default/tests/two-fa-setup-shared-verification.test.tsx new file mode 100644 index 00000000..7900b1bd --- /dev/null +++ b/web/default/tests/two-fa-setup-shared-verification.test.tsx @@ -0,0 +1,82 @@ +/* +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 { createReactTestEnvironment } from '@/test/react' +import { cleanup, render, waitFor } from '@testing-library/react' +import { afterAll, afterEach, beforeAll, beforeEach, mock, test } from 'bun:test' +import assert from 'node:assert/strict' + +const setup2FA = mock(async () => ({ + success: true, + data: { + qr_code_data: 'otpauth://totp/MAX-API:test', + secret: 'TESTSECRET', + backup_codes: ['ABCD-EFGH'], + }, +})) +const withVerification = mock((apiCall: () => Promise) => apiCall()) + +mock.module('react-i18next', () => ({ + useTranslation: () => ({ t: (key: string) => key }), +})) + +mock.module('sonner', () => ({ + toast: { + error: mock(() => undefined), + success: mock(() => undefined), + }, +})) + +mock.module('qrcode.react', () => ({ + QRCodeSVG: () =>
, +})) + +mock.module('../src/lib/api', () => ({ + setup2FA, + enable2FA: mock(async () => ({ success: true })), +})) + +mock.module('../src/features/auth/secure-verification', () => ({ + useSecureVerificationGate: () => ({ withVerification }), +})) + +const { TwoFASetupDialog } = await import( + '../src/features/profile/components/dialogs/two-fa-setup-dialog' +) +const testEnv = createReactTestEnvironment() + +beforeAll(() => testEnv.setup()) +beforeEach(() => { + setup2FA.mockClear() + withVerification.mockClear() +}) +afterEach(() => cleanup()) +afterAll(() => testEnv.teardown()) + +test('uses the authenticated layout verification gate for 2FA setup', async () => { + render( + undefined)} + onSuccess={mock(() => undefined)} + /> + ) + + await waitFor(() => assert.equal(setup2FA.mock.calls.length, 1)) + assert.equal(withVerification.mock.calls.length, 1) +}) From 771bdfd2a300d5d24c18251be5b46f2e0700f9a5 Mon Sep 17 00:00:00 2001 From: CSCITech Date: Fri, 28 Aug 2026 08:23:15 +0800 Subject: [PATCH 10/70] fix: harden secure verification error handling --- .../hooks/use-secure-verification.ts | 15 +++++--- .../dialogs/two-fa-setup-dialog.tsx | 10 ++++-- ...sskey-registration-error-handling.test.tsx | 24 ++++++++++--- .../secure-verification-continuation.test.tsx | 36 +++++++++++++++++++ 4 files changed, 75 insertions(+), 10 deletions(-) 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 605e1e31..3acb9ea1 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 @@ -86,7 +86,7 @@ export function useSecureVerification( }, []) const settlePendingVerification = useCallback( - (value: unknown | null, error?: unknown) => { + (value: unknown | null, error?: unknown): void => { const pending = pendingVerificationRef.current pendingVerificationRef.current = null if (!pending) return @@ -109,7 +109,14 @@ export function useSecureVerification( config: StartVerificationOptions = {} ) => { const { preferredMethod, title, description, scope } = config - const availableMethods = await fetchVerificationMethods(scope) + let availableMethods: VerificationMethods + try { + availableMethods = await fetchVerificationMethods(scope) + } catch (error) { + toast.error(i18next.t('Verification unavailable')) + onError?.(error) + throw error + } if ( !availableMethods.has2FA && @@ -224,13 +231,13 @@ export function useSecureVerification( setState((prev) => ({ ...prev, method, code: '' })) }, []) - const cancel = useCallback(() => { + const cancel = useCallback((): void => { settlePendingVerification(null) reset() }, [reset, settlePendingVerification]) const setOpen = useCallback( - (nextOpen: boolean) => { + (nextOpen: boolean): void => { if (!nextOpen) { cancel() return 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 a2a639e1..d64a4a0e 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,7 +16,13 @@ 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, useRef } 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' @@ -152,7 +158,7 @@ export function TwoFASetupDialog({ return () => window.clearTimeout(timer) }, [open, setupData, handleSetup]) - const renderSetupState = () => { + const renderSetupState = (): ReactNode => { if (initializing) { return (
diff --git a/web/default/tests/passkey-registration-error-handling.test.tsx b/web/default/tests/passkey-registration-error-handling.test.tsx index 7174068d..f5a09196 100644 --- a/web/default/tests/passkey-registration-error-handling.test.tsx +++ b/web/default/tests/passkey-registration-error-handling.test.tsx @@ -18,12 +18,21 @@ For commercial licensing, please contact https://github.com/MAX-API-Next/MAX-API */ import { createReactTestEnvironment } from '@/test/react' import { cleanup, fireEvent, render, waitFor } from '@testing-library/react' -import { afterAll, afterEach, beforeAll, beforeEach, mock, test } from 'bun:test' +import { + afterAll, + afterEach, + beforeAll, + beforeEach, + mock, + spyOn, + test, +} from 'bun:test' import assert from 'node:assert/strict' const withVerification = mock(async () => { throw new Error('verification method lookup failed') }) +const consoleError = spyOn(console, 'error').mockImplementation(() => undefined) mock.module('react-i18next', () => ({ useTranslation: () => ({ t: (key: string) => key }), @@ -83,12 +92,18 @@ const { PasskeyCard } = await import( ) const testEnv = createReactTestEnvironment() -beforeAll(() => testEnv.setup()) +beforeAll(() => { + testEnv.setup() +}) beforeEach(() => { withVerification.mockClear() + consoleError.mockClear() }) afterEach(() => cleanup()) -afterAll(() => testEnv.teardown()) +afterAll(() => { + consoleError.mockRestore() + testEnv.teardown() +}) test('handles a rejected Passkey verification continuation', async () => { const view = render() @@ -96,5 +111,6 @@ test('handles a rejected Passkey verification continuation', async () => { fireEvent.click(view.getByRole('button', { name: 'Enable Passkey' })) await waitFor(() => assert.equal(withVerification.mock.calls.length, 1)) - await new Promise((resolve) => setTimeout(resolve, 0)) + await waitFor(() => assert.equal(consoleError.mock.calls.length, 1)) + assert.match(String(consoleError.mock.calls[0]?.[1]), /method lookup failed/) }) diff --git a/web/default/tests/secure-verification-continuation.test.tsx b/web/default/tests/secure-verification-continuation.test.tsx index c32d7579..7129b95f 100644 --- a/web/default/tests/secure-verification-continuation.test.tsx +++ b/web/default/tests/secure-verification-continuation.test.tsx @@ -36,6 +36,21 @@ const checkVerificationMethods = mock(async () => ({ passkeySupported: false, })) const verify = mock(async () => undefined) +const toastError = mock((_message: string) => undefined) +const toastInfo = mock((_message: string) => undefined) +const toastSuccess = mock((_message: string) => undefined) + +mock.module('i18next', () => ({ + default: { t: (key: string) => key }, +})) + +mock.module('sonner', () => ({ + toast: { + error: toastError, + info: toastInfo, + success: toastSuccess, + }, +})) mock.module('../src/features/auth/secure-verification/api', () => ({ checkVerificationMethods, @@ -62,6 +77,9 @@ beforeAll(() => testEnv.setup()) beforeEach(() => { checkVerificationMethods.mockClear() verify.mockClear() + toastError.mockClear() + toastInfo.mockClear() + toastSuccess.mockClear() }) afterEach(() => cleanup()) @@ -116,6 +134,24 @@ describe('useSecureVerification', () => { assert.equal(result.current.open, false) }) + test('reports verification method discovery failures before rejecting', async () => { + checkVerificationMethods.mockImplementationOnce(async () => { + throw new Error('method discovery unavailable') + }) + const apiCall = mock(async () => { + throw verificationRequiredError + }) + const { result } = renderHook(() => useSecureVerification()) + + await assert.rejects( + result.current.withVerification(apiCall, { scope: 'credentials' }), + /method discovery unavailable/ + ) + + assert.equal(toastError.mock.calls.length, 1) + assert.equal(toastError.mock.calls[0]?.[0], 'Verification unavailable') + }) + test('clears loading after a retryable verification failure', async () => { verify.mockImplementationOnce(async () => { throw new Error('Incorrect password') From ac4dfaa941db3ea23169a36218204786d27a64af Mon Sep 17 00:00:00 2001 From: CSCITech Date: Fri, 28 Aug 2026 11:43:58 +0800 Subject: [PATCH 11/70] fix: settle verification continuation on reset --- .../hooks/use-secure-verification.ts | 14 ++++++------- .../secure-verification-continuation.test.tsx | 20 +++++++++++++++++++ 2 files changed, 27 insertions(+), 7 deletions(-) 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 3acb9ea1..07b820a8 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 @@ -80,11 +80,6 @@ export function useSecureVerification( [] ) - const reset = useCallback(() => { - setState(initialState) - setDialogOpen(false) - }, []) - const settlePendingVerification = useCallback( (value: unknown | null, error?: unknown): void => { const pending = pendingVerificationRef.current @@ -99,6 +94,12 @@ export function useSecureVerification( [] ) + const reset = useCallback((): void => { + settlePendingVerification(null) + setState(initialState) + setDialogOpen(false) + }, [settlePendingVerification]) + useEffect(() => { return () => settlePendingVerification(null) }, [settlePendingVerification]) @@ -232,9 +233,8 @@ export function useSecureVerification( }, []) const cancel = useCallback((): void => { - settlePendingVerification(null) reset() - }, [reset, settlePendingVerification]) + }, [reset]) const setOpen = useCallback( (nextOpen: boolean): void => { diff --git a/web/default/tests/secure-verification-continuation.test.tsx b/web/default/tests/secure-verification-continuation.test.tsx index 7129b95f..2cf6db98 100644 --- a/web/default/tests/secure-verification-continuation.test.tsx +++ b/web/default/tests/secure-verification-continuation.test.tsx @@ -134,6 +134,26 @@ describe('useSecureVerification', () => { assert.equal(result.current.open, false) }) + test('resolves a pending continuation when reset is called', async () => { + const apiCall = mock(async () => { + throw verificationRequiredError + }) + const { result } = renderHook(() => useSecureVerification()) + + let continuation: Promise | undefined + act(() => { + continuation = result.current.withVerification(apiCall, { + scope: 'credentials', + }) + }) + + await waitFor(() => assert.equal(result.current.open, true)) + act(() => result.current.reset()) + + assert.equal(await continuation, null) + assert.equal(result.current.open, false) + }) + test('reports verification method discovery failures before rejecting', async () => { checkVerificationMethods.mockImplementationOnce(async () => { throw new Error('method discovery unavailable') From 0ed722b636750d7f54c825d4c65791d0d0dd7d8d Mon Sep 17 00:00:00 2001 From: CSCITech Date: Fri, 28 Aug 2026 17:21:59 +0800 Subject: [PATCH 12/70] fix: harden secure verification cancellation --- .../features/auth/secure-verification/api.ts | 22 +- .../hooks/use-secure-verification.ts | 166 ++++++++++--- .../profile/components/passkey-card.tsx | 28 ++- ...sskey-registration-error-handling.test.tsx | 100 +++++++- .../tests/secure-verification-api.test.tsx | 121 ++++++++++ .../secure-verification-continuation.test.tsx | 222 +++++++++++++++++- 6 files changed, 590 insertions(+), 69 deletions(-) create mode 100644 web/default/tests/secure-verification-api.test.tsx diff --git a/web/default/src/features/auth/secure-verification/api.ts b/web/default/src/features/auth/secure-verification/api.ts index 4d5415e7..15cc137f 100644 --- a/web/default/src/features/auth/secure-verification/api.ts +++ b/web/default/src/features/auth/secure-verification/api.ts @@ -54,26 +54,22 @@ export async function checkVerificationMethods( 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/hooks/use-secure-verification.ts b/web/default/src/features/auth/secure-verification/hooks/use-secure-verification.ts index 07b820a8..0d39a4e4 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 @@ -37,12 +37,14 @@ 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 = { @@ -59,6 +61,7 @@ const initialState: InternalState = { title: undefined, description: undefined, apiCall: null, + attemptId: null, } export function useSecureVerification( @@ -70,19 +73,30 @@ export function useSecureVerification( const [state, setState] = useState(initialState) 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): void => { + (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) { @@ -94,31 +108,58 @@ export function useSecureVerification( [] ) + const isVerificationAttemptActive = useCallback( + (attemptId: number): boolean => + mountedRef.current && verificationAttemptRef.current === attemptId, + [] + ) + + const beginVerificationAttempt = useCallback((): number | null => { + if (!mountedRef.current) return null + verificationAttemptRef.current += 1 + settlePendingVerification(null) + setState(initialState) + setDialogOpen(false) + return verificationAttemptRef.current + }, [settlePendingVerification]) + const reset = useCallback((): void => { + verificationAttemptRef.current += 1 settlePendingVerification(null) + if (!mountedRef.current) return setState(initialState) setDialogOpen(false) }, [settlePendingVerification]) useEffect(() => { - return () => settlePendingVerification(null) + mountedRef.current = true + return () => { + mountedRef.current = false + verificationAttemptRef.current += 1 + settlePendingVerification(null) + } }, [settlePendingVerification]) - const startVerification = useCallback( + const startVerificationAttempt = useCallback( async ( apiCall: () => Promise, - config: StartVerificationOptions = {} - ) => { + config: StartVerificationOptions, + attemptId: number + ): Promise => { const { preferredMethod, title, description, scope } = config let availableMethods: VerificationMethods try { - availableMethods = await fetchVerificationMethods(scope) + availableMethods = await loadVerificationMethods(scope) } catch (error) { + if (!isVerificationAttemptActive(attemptId)) return false toast.error(i18next.t('Verification unavailable')) onError?.(error) throw error } + if (!isVerificationAttemptActive(attemptId)) return false + setMethods(availableMethods) + if ( !availableMethods.has2FA && !availableMethods.hasPasskey && @@ -145,6 +186,7 @@ export function useSecureVerification( setState((prev) => ({ ...prev, apiCall, + attemptId, method: defaultMethod, title, description, @@ -153,27 +195,46 @@ export function useSecureVerification( 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) } catch (error) { + if (!isVerificationAttemptActive(attemptId)) return setState((prev) => ({ ...prev, loading: false })) const message = error instanceof Error @@ -184,9 +245,12 @@ export function useSecureVerification( throw error } + if (!isVerificationAttemptActive(attemptId)) return + try { - const result = await state.apiCall() - settlePendingVerification(result) + const result = await apiCall() + if (!isVerificationAttemptActive(attemptId)) return + settlePendingVerification(result, undefined, attemptId) if (successMessage) { toast.success(successMessage) @@ -200,7 +264,8 @@ export function useSecureVerification( return result } catch (error) { - settlePendingVerification(null, error) + if (!isVerificationAttemptActive(attemptId)) return + settlePendingVerification(null, error, attemptId) const message = error instanceof Error ? error.message @@ -210,11 +275,16 @@ export function useSecureVerification( reset() throw error } finally { - setState((prev) => ({ ...prev, loading: false })) + if (isVerificationAttemptActive(attemptId)) { + setState((prev) => + prev.attemptId === attemptId ? { ...prev, loading: false } : prev + ) + } } }, [ state, + isVerificationAttemptActive, successMessage, onSuccess, onError, @@ -256,27 +326,55 @@ export function useSecureVerification( apiCall: () => Promise, config: StartVerificationOptions = {} ): Promise => { - try { - return await apiCall() - } catch (error) { - if (isVerificationRequiredError(error)) { + const attemptId = beginVerificationAttempt() + if (attemptId === null) return null + const continuation = new Promise((resolve, reject) => { + pendingVerificationRef.current = { + attemptId, + resolve: (value) => resolve(value as T | null), + reject, + } + }) + + void (async (): Promise => { + 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) - const started = await startVerification(apiCall, config) - if (!started) return null - - settlePendingVerification(null) - return await new Promise((resolve, reject) => { - pendingVerificationRef.current = { - resolve: (value) => resolve(value as T | null), - reject, - } - }) } - throw error - } + + try { + const started = await startVerificationAttempt( + apiCall, + config, + attemptId + ) + if (!started) { + settlePendingVerification(null, undefined, attemptId) + } + } catch (error) { + settlePendingVerification(null, error, attemptId) + } + })() + + return await continuation }, - [settlePendingVerification, startVerification] + [ + beginVerificationAttempt, + isVerificationAttemptActive, + settlePendingVerification, + startVerificationAttempt, + ] ) const canUseMethod = useCallback( diff --git a/web/default/src/features/profile/components/passkey-card.tsx b/web/default/src/features/profile/components/passkey-card.tsx index 0471f3ef..e165d673 100644 --- a/web/default/src/features/profile/components/passkey-card.tsx +++ b/web/default/src/features/profile/components/passkey-card.tsx @@ -121,8 +121,14 @@ export function PasskeyCard({ loading: pageLoading }: PasskeyCardProps) { } }, [register, supported, t, withVerification]) - const handleRemove = useCallback(async () => { - const methods = await fetchVerificationMethods() + const handleRemove = useCallback(async (): Promise => { + let methods: VerificationMethods + try { + methods = await fetchVerificationMethods() + } catch { + toast.error(t('Verification unavailable')) + return + } const required: VerificationMethod | null = methods.has2FA ? '2fa' : methods.hasPasskey @@ -145,13 +151,17 @@ 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.' - ), - }) + try { + await startVerification(remove, { + preferredMethod: required, + title: t('Security verification'), + description: t( + 'Confirm your identity before removing this Passkey from your account.' + ), + }) + } catch { + // useSecureVerification already reports method discovery failures. + } }, [fetchVerificationMethods, remove, startVerification, t]) const handleVerificationCancel = useCallback(() => { diff --git a/web/default/tests/passkey-registration-error-handling.test.tsx b/web/default/tests/passkey-registration-error-handling.test.tsx index f5a09196..66ca6026 100644 --- a/web/default/tests/passkey-registration-error-handling.test.tsx +++ b/web/default/tests/passkey-registration-error-handling.test.tsx @@ -28,11 +28,36 @@ import { test, } from 'bun:test' import assert from 'node:assert/strict' +import type { MouseEventHandler, ReactNode } from 'react' const withVerification = mock(async () => { throw new Error('verification method lookup failed') }) +const fetchVerificationMethods = mock(async () => ({ + has2FA: false, + hasPasskey: false, + hasPassword: true, + passkeySupported: true, +})) +const startVerification = mock(async () => true) +const register = mock(async () => true) +const remove = mock(async () => true) +const toastError = mock((_message: string) => undefined) const consoleError = spyOn(console, 'error').mockImplementation(() => undefined) +let passkeyEnabled = false + +interface DialogPartProps { + children?: ReactNode +} + +interface DialogButtonProps extends DialogPartProps { + disabled?: boolean + onClick?: MouseEventHandler +} + +function DialogPart(props: DialogPartProps) { + return
{props.children}
+} mock.module('react-i18next', () => ({ useTranslation: () => ({ t: (key: string) => key }), @@ -40,23 +65,51 @@ mock.module('react-i18next', () => ({ mock.module('sonner', () => ({ toast: { - error: mock(() => undefined), + error: toastError, info: mock(() => undefined), success: mock(() => undefined), }, })) +mock.module('../src/components/ui/alert-dialog', () => ({ + AlertDialog: (props: DialogPartProps) => <>{props.children}, + AlertDialogAction: (props: DialogButtonProps) => ( + + ), + AlertDialogCancel: (props: DialogButtonProps) => ( + + ), + AlertDialogContent: DialogPart, + AlertDialogDescription: DialogPart, + AlertDialogFooter: DialogPart, + AlertDialogHeader: DialogPart, + AlertDialogTitle: DialogPart, + AlertDialogTrigger: (props: DialogPartProps) => ( + + ), +})) + mock.module('../src/features/auth/passkey', () => ({ + beginPasskeyVerification: mock(async () => ({})), + finishPasskeyVerification: mock(async () => ({})), usePasskeyManagement: () => ({ status: null, loading: false, registering: false, removing: false, supported: true, - enabled: false, + enabled: passkeyEnabled, lastUsed: null, - register: mock(async () => true), - remove: mock(async () => true), + register, + remove, }), })) @@ -72,17 +125,12 @@ mock.module('../src/features/auth/secure-verification', () => ({ passkeySupported: true, }, state: { method: null, loading: false, code: '' }, - startVerification: mock(async () => true), + startVerification, executeVerification: mock(async () => undefined), cancel: mock(() => undefined), setCode: mock(() => undefined), switchMethod: mock(() => undefined), - fetchVerificationMethods: mock(async () => ({ - has2FA: false, - hasPasskey: false, - hasPassword: true, - passkeySupported: true, - })), + fetchVerificationMethods, withVerification, }), })) @@ -96,7 +144,19 @@ beforeAll(() => { testEnv.setup() }) beforeEach(() => { + passkeyEnabled = false withVerification.mockClear() + fetchVerificationMethods.mockClear() + fetchVerificationMethods.mockImplementation(async () => ({ + has2FA: false, + hasPasskey: false, + hasPassword: true, + passkeySupported: true, + })) + startVerification.mockClear() + register.mockClear() + remove.mockClear() + toastError.mockClear() consoleError.mockClear() }) afterEach(() => cleanup()) @@ -114,3 +174,21 @@ test('handles a rejected Passkey verification continuation', async () => { await waitFor(() => assert.equal(consoleError.mock.calls.length, 1)) assert.match(String(consoleError.mock.calls[0]?.[1]), /method lookup failed/) }) + +test('stops Passkey removal when verification method discovery fails', async () => { + passkeyEnabled = true + fetchVerificationMethods.mockImplementationOnce(async () => { + throw new Error('verification methods unavailable') + }) + const view = render() + + fireEvent.click(view.getByRole('button', { name: 'Remove' })) + + await waitFor(() => + assert.equal(fetchVerificationMethods.mock.calls.length, 1) + ) + await waitFor(() => assert.equal(toastError.mock.calls.length, 1)) + assert.equal(toastError.mock.calls[0]?.[0], 'Verification unavailable') + assert.equal(startVerification.mock.calls.length, 0) + assert.equal(remove.mock.calls.length, 0) +}) diff --git a/web/default/tests/secure-verification-api.test.tsx b/web/default/tests/secure-verification-api.test.tsx new file mode 100644 index 00000000..449340a7 --- /dev/null +++ b/web/default/tests/secure-verification-api.test.tsx @@ -0,0 +1,121 @@ +/* +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 { afterAll, beforeEach, mock, spyOn, test } from 'bun:test' +import assert from 'node:assert/strict' + +const apiGet = mock(async () => ({ + data: { + success: true, + data: { + has_2fa: false, + has_passkey: false, + has_password: true, + }, + }, +})) +const consoleError = spyOn(console, 'error').mockImplementation(() => undefined) + +const passkeyManagementStub = { + status: null, + loading: false, + registering: false, + removing: false, + supported: true, + enabled: false, + lastUsed: null, + register: mock(async () => true), + remove: mock(async () => true), +} + +mock.module('../src/lib/api', () => ({ + api: { + get: apiGet, + }, +})) + +mock.module('../src/lib/passkey', () => ({ + buildAssertionResult: mock(() => null), + prepareCredentialRequestOptions: mock(() => ({})), + isPasskeySupported: mock(async () => false), +})) + +mock.module('../src/features/auth/passkey', () => ({ + beginPasskeyVerification: mock(async () => ({})), + finishPasskeyVerification: mock(async () => ({})), + usePasskeyManagement: () => passkeyManagementStub, +})) + +const { checkVerificationMethods } = await import( + '../src/features/auth/secure-verification/api' +) + +beforeEach(() => { + apiGet.mockClear() + consoleError.mockClear() +}) + +afterAll(() => consoleError.mockRestore()) + +test('propagates verification method request failures', async () => { + apiGet.mockImplementationOnce(async () => { + throw new Error('verification methods unavailable') + }) + + await assert.rejects( + checkVerificationMethods('credentials'), + /verification methods unavailable/ + ) + assert.equal(consoleError.mock.calls.length, 1) +}) + +test('rejects unsuccessful verification method responses', async () => { + apiGet.mockImplementationOnce(async () => ({ + data: { + success: false, + message: 'failed to load verification methods', + }, + })) + + await assert.rejects( + checkVerificationMethods('credentials'), + /failed to load verification methods/ + ) + assert.equal(consoleError.mock.calls.length, 1) +}) + +test('preserves a successful response with no available methods', async () => { + apiGet.mockImplementationOnce(async () => ({ + data: { + success: true, + data: { + has_2fa: false, + has_passkey: false, + has_password: false, + }, + }, + })) + + assert.deepEqual(await checkVerificationMethods('credentials'), { + has2FA: false, + hasPasskey: false, + hasPassword: false, + passkeySupported: false, + }) + assert.equal(consoleError.mock.calls.length, 0) +}) diff --git a/web/default/tests/secure-verification-continuation.test.tsx b/web/default/tests/secure-verification-continuation.test.tsx index 2cf6db98..7b39a1df 100644 --- a/web/default/tests/secure-verification-continuation.test.tsx +++ b/web/default/tests/secure-verification-continuation.test.tsx @@ -29,17 +29,40 @@ import { } from 'bun:test' import assert from 'node:assert/strict' -const checkVerificationMethods = mock(async () => ({ +const availableVerificationMethods = { has2FA: false, hasPasskey: false, hasPassword: true, passkeySupported: false, -})) +} +const checkVerificationMethods = mock( + async () => availableVerificationMethods +) const verify = mock(async () => undefined) const toastError = mock((_message: string) => undefined) const toastInfo = mock((_message: string) => undefined) const toastSuccess = mock((_message: string) => undefined) +interface Deferred { + promise: Promise + resolve: (value: T) => void + reject: (reason?: unknown) => void +} + +function createDeferred(): Deferred { + let resolve: ((value: T) => void) | undefined + let reject: ((reason?: unknown) => void) | undefined + const promise = new Promise((resolver, rejecter) => { + resolve = resolver + reject = rejecter + }) + return { + promise, + resolve: (value) => resolve?.(value), + reject: (reason) => reject?.(reason), + } +} + mock.module('i18next', () => ({ default: { t: (key: string) => key }, })) @@ -76,7 +99,11 @@ beforeAll(() => testEnv.setup()) beforeEach(() => { checkVerificationMethods.mockClear() + checkVerificationMethods.mockImplementation( + async () => availableVerificationMethods + ) verify.mockClear() + verify.mockImplementation(async () => undefined) toastError.mockClear() toastInfo.mockClear() toastSuccess.mockClear() @@ -87,6 +114,36 @@ afterEach(() => cleanup()) afterAll(() => testEnv.teardown()) describe('useSecureVerification', () => { + test('returns an initial protected action success without discovery', async () => { + const apiCall = mock(async () => ({ success: true })) + const { result } = renderHook(() => useSecureVerification()) + + assert.deepEqual( + await result.current.withVerification(apiCall, { + scope: 'credentials', + }), + { success: true } + ) + assert.equal(apiCall.mock.calls.length, 1) + assert.equal(checkVerificationMethods.mock.calls.length, 0) + assert.equal(result.current.open, false) + }) + + test('rejects an initial protected action failure without discovery', async () => { + const apiCall = mock(async () => { + throw new Error('protected action failed') + }) + const { result } = renderHook(() => useSecureVerification()) + + await assert.rejects( + result.current.withVerification(apiCall, { scope: 'credentials' }), + /protected action failed/ + ) + assert.equal(apiCall.mock.calls.length, 1) + assert.equal(checkVerificationMethods.mock.calls.length, 0) + assert.equal(result.current.open, false) + }) + test('continues the original API call after verification succeeds', async () => { let callCount = 0 const apiCall = mock(async () => { @@ -154,6 +211,167 @@ describe('useSecureVerification', () => { assert.equal(result.current.open, false) }) + test('settles reset while the initial protected action is pending', async () => { + const initialAction = createDeferred() + const apiCall = mock(() => initialAction.promise) + const { result } = renderHook(() => useSecureVerification()) + + const continuation = result.current.withVerification(apiCall, { + scope: 'credentials', + }) + await waitFor(() => assert.equal(apiCall.mock.calls.length, 1)) + + act(() => result.current.reset()) + assert.equal(await continuation, null) + + initialAction.reject(verificationRequiredError) + await Promise.resolve() + assert.equal(checkVerificationMethods.mock.calls.length, 0) + assert.equal(result.current.open, false) + }) + + test('does not retry the protected action after reset during verification', async () => { + const verification = createDeferred() + verify.mockImplementationOnce(() => verification.promise) + let callCount = 0 + const apiCall = mock(async () => { + callCount += 1 + if (callCount === 1) throw verificationRequiredError + return { success: true } + }) + const { result } = renderHook(() => useSecureVerification()) + + const continuation = result.current.withVerification(apiCall, { + scope: 'credentials', + }) + await waitFor(() => assert.equal(result.current.open, true)) + + let execution: Promise | undefined + act(() => { + execution = result.current.executeVerification('password', 'secret') + }) + await waitFor(() => assert.equal(verify.mock.calls.length, 1)) + + act(() => result.current.reset()) + assert.equal(await continuation, null) + + await act(async () => { + verification.resolve() + await execution + }) + + assert.equal(apiCall.mock.calls.length, 1) + assert.equal(toastSuccess.mock.calls.length, 0) + }) + + test('settles reset before verification method discovery completes', async () => { + const discovery = createDeferred() + checkVerificationMethods.mockImplementationOnce(() => discovery.promise) + const apiCall = mock(async () => { + throw verificationRequiredError + }) + const { result } = renderHook(() => useSecureVerification()) + + const continuation = result.current.withVerification(apiCall, { + scope: 'credentials', + }) + await waitFor(() => + assert.equal(checkVerificationMethods.mock.calls.length, 1) + ) + + act(() => result.current.reset()) + const outcome = await Promise.race([ + continuation.then((value) => ({ settled: true, value })), + new Promise<{ settled: false }>((resolve) => { + setTimeout(() => resolve({ settled: false }), 50) + }), + ]) + + await act(async () => { + discovery.resolve(availableVerificationMethods) + await Promise.resolve() + }) + + assert.deepEqual(outcome, { settled: true, value: null }) + assert.equal(result.current.open, false) + }) + + test('settles unmount before verification method discovery completes', async () => { + const discovery = createDeferred() + checkVerificationMethods.mockImplementationOnce(() => discovery.promise) + const apiCall = mock(async () => { + throw verificationRequiredError + }) + const { result, unmount } = renderHook(() => useSecureVerification()) + + const continuation = result.current.withVerification(apiCall, { + scope: 'credentials', + }) + await waitFor(() => + assert.equal(checkVerificationMethods.mock.calls.length, 1) + ) + + unmount() + const outcome = await Promise.race([ + continuation.then((value) => ({ settled: true, value })), + new Promise<{ settled: false }>((resolve) => { + setTimeout(() => resolve({ settled: false }), 50) + }), + ]) + discovery.resolve(availableVerificationMethods) + await Promise.resolve() + + assert.deepEqual(outcome, { settled: true, value: null }) + assert.equal(apiCall.mock.calls.length, 1) + }) + + test('does not let an obsolete discovery settle a newer continuation', async () => { + const firstDiscovery = createDeferred() + const secondDiscovery = createDeferred< + typeof availableVerificationMethods + >() + checkVerificationMethods + .mockImplementationOnce(() => firstDiscovery.promise) + .mockImplementationOnce(() => secondDiscovery.promise) + const firstCall = mock(async () => { + throw verificationRequiredError + }) + const secondCall = mock(async () => { + throw verificationRequiredError + }) + const { result } = renderHook(() => useSecureVerification()) + + const firstContinuation = result.current.withVerification(firstCall, { + scope: 'credentials', + }) + await waitFor(() => + assert.equal(checkVerificationMethods.mock.calls.length, 1) + ) + const secondContinuation = result.current.withVerification(secondCall, { + scope: 'api_token', + }) + await waitFor(() => + assert.equal(checkVerificationMethods.mock.calls.length, 2) + ) + + assert.equal(await firstContinuation, null) + await act(async () => { + firstDiscovery.resolve(availableVerificationMethods) + await Promise.resolve() + }) + assert.equal(result.current.open, false) + + await act(async () => { + secondDiscovery.resolve(availableVerificationMethods) + await Promise.resolve() + }) + await waitFor(() => assert.equal(result.current.open, true)) + assert.equal(result.current.state.scope, 'api_token') + + act(() => result.current.cancel()) + assert.equal(await secondContinuation, null) + }) + test('reports verification method discovery failures before rejecting', async () => { checkVerificationMethods.mockImplementationOnce(async () => { throw new Error('method discovery unavailable') From 5cb92633f5f83c26083f64a00b34dd26f7ec7a42 Mon Sep 17 00:00:00 2001 From: CSCITech Date: Fri, 28 Aug 2026 21:21:12 +0800 Subject: [PATCH 13/70] fix: unify secure verification error handling --- .../features/auth/secure-verification/api.ts | 2 + .../hooks/use-secure-verification.ts | 7 +- .../profile/components/passkey-card.tsx | 7 +- ...sskey-registration-error-handling.test.tsx | 19 ++++- .../tests/secure-verification-api.test.tsx | 23 +++++- .../secure-verification-continuation.test.tsx | 71 +++++++++++++++++-- 6 files changed, 117 insertions(+), 12 deletions(-) diff --git a/web/default/src/features/auth/secure-verification/api.ts b/web/default/src/features/auth/secure-verification/api.ts index 15cc137f..3c51e5ef 100644 --- a/web/default/src/features/auth/secure-verification/api.ts +++ b/web/default/src/features/auth/secure-verification/api.ts @@ -50,6 +50,8 @@ export async function checkVerificationMethods( const [methodsResponse, passkeySupported] = await Promise.all([ api.get('/api/verify/methods', { params: scope ? { scope } : undefined, + skipBusinessError: true, + skipErrorHandler: true, }), detectPasskeySupport(), ]) 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 0d39a4e4..eb0488e9 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 @@ -19,6 +19,7 @@ For commercial licensing, please contact https://github.com/MAX-API-Next/MAX-API 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, @@ -152,7 +153,9 @@ export function useSecureVerification( availableMethods = await loadVerificationMethods(scope) } catch (error) { if (!isVerificationAttemptActive(attemptId)) return false - toast.error(i18next.t('Verification unavailable')) + handleServerError(error, { + fallback: i18next.t('Verification unavailable'), + }) onError?.(error) throw error } @@ -270,7 +273,7 @@ export function useSecureVerification( error instanceof Error ? error.message : i18next.t('Verification failed') - toast.error(message) + handleServerError(error, { fallback: message }) onError?.(error) reset() throw error diff --git a/web/default/src/features/profile/components/passkey-card.tsx b/web/default/src/features/profile/components/passkey-card.tsx index e165d673..9555eeb2 100644 --- a/web/default/src/features/profile/components/passkey-card.tsx +++ b/web/default/src/features/profile/components/passkey-card.tsx @@ -21,6 +21,7 @@ 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, @@ -125,8 +126,10 @@ export function PasskeyCard({ loading: pageLoading }: PasskeyCardProps) { let methods: VerificationMethods try { methods = await fetchVerificationMethods() - } catch { - toast.error(t('Verification unavailable')) + } catch (error) { + handleServerError(error, { + fallback: t('Verification unavailable'), + }) return } const required: VerificationMethod | null = methods.has2FA diff --git a/web/default/tests/passkey-registration-error-handling.test.tsx b/web/default/tests/passkey-registration-error-handling.test.tsx index 66ca6026..2868d1d3 100644 --- a/web/default/tests/passkey-registration-error-handling.test.tsx +++ b/web/default/tests/passkey-registration-error-handling.test.tsx @@ -43,6 +43,9 @@ const startVerification = mock(async () => true) const register = mock(async () => true) const remove = mock(async () => true) const toastError = mock((_message: string) => undefined) +const handleServerError = mock( + (_error: unknown, _options?: { fallback?: string }) => undefined +) const consoleError = spyOn(console, 'error').mockImplementation(() => undefined) let passkeyEnabled = false @@ -71,6 +74,10 @@ mock.module('sonner', () => ({ }, })) +mock.module('../src/lib/handle-server-error', () => ({ + handleServerError, +})) + mock.module('../src/components/ui/alert-dialog', () => ({ AlertDialog: (props: DialogPartProps) => <>{props.children}, AlertDialogAction: (props: DialogButtonProps) => ( @@ -157,6 +164,7 @@ beforeEach(() => { register.mockClear() remove.mockClear() toastError.mockClear() + handleServerError.mockClear() consoleError.mockClear() }) afterEach(() => cleanup()) @@ -187,8 +195,15 @@ test('stops Passkey removal when verification method discovery fails', async () await waitFor(() => assert.equal(fetchVerificationMethods.mock.calls.length, 1) ) - await waitFor(() => assert.equal(toastError.mock.calls.length, 1)) - assert.equal(toastError.mock.calls[0]?.[0], 'Verification unavailable') + await waitFor(() => assert.equal(handleServerError.mock.calls.length, 1)) + assert.match( + String(handleServerError.mock.calls[0]?.[0]), + /verification methods unavailable/ + ) + assert.deepEqual(handleServerError.mock.calls[0]?.[1], { + fallback: 'Verification unavailable', + }) + assert.equal(toastError.mock.calls.length, 0) assert.equal(startVerification.mock.calls.length, 0) assert.equal(remove.mock.calls.length, 0) }) diff --git a/web/default/tests/secure-verification-api.test.tsx b/web/default/tests/secure-verification-api.test.tsx index 449340a7..a27c9308 100644 --- a/web/default/tests/secure-verification-api.test.tsx +++ b/web/default/tests/secure-verification-api.test.tsx @@ -19,7 +19,15 @@ For commercial licensing, please contact https://github.com/MAX-API-Next/MAX-API import { afterAll, beforeEach, mock, spyOn, test } from 'bun:test' import assert from 'node:assert/strict' -const apiGet = mock(async () => ({ +interface VerificationMethodsRequestConfig { + params?: { + scope?: string + } + skipBusinessError?: boolean + skipErrorHandler?: boolean +} + +const successfulMethodsResponse = { data: { success: true, data: { @@ -28,7 +36,12 @@ const apiGet = mock(async () => ({ has_password: true, }, }, -})) +} + +const apiGet = mock( + async (_url: string, _config?: VerificationMethodsRequestConfig) => + successfulMethodsResponse +) const consoleError = spyOn(console, 'error').mockImplementation(() => undefined) const passkeyManagementStub = { @@ -117,5 +130,11 @@ test('preserves a successful response with no available methods', async () => { hasPassword: false, passkeySupported: false, }) + assert.equal(apiGet.mock.calls[0]?.[0], '/api/verify/methods') + assert.deepEqual(apiGet.mock.calls[0]?.[1]?.params, { + scope: 'credentials', + }) + assert.equal(apiGet.mock.calls[0]?.[1]?.skipBusinessError, true) + assert.equal(apiGet.mock.calls[0]?.[1]?.skipErrorHandler, true) assert.equal(consoleError.mock.calls.length, 0) }) diff --git a/web/default/tests/secure-verification-continuation.test.tsx b/web/default/tests/secure-verification-continuation.test.tsx index 7b39a1df..c2840457 100644 --- a/web/default/tests/secure-verification-continuation.test.tsx +++ b/web/default/tests/secure-verification-continuation.test.tsx @@ -42,6 +42,9 @@ const verify = mock(async () => undefined) const toastError = mock((_message: string) => undefined) const toastInfo = mock((_message: string) => undefined) const toastSuccess = mock((_message: string) => undefined) +const handleServerError = mock( + (_error: unknown, _options?: { fallback?: string }) => undefined +) interface Deferred { promise: Promise @@ -63,6 +66,12 @@ function createDeferred(): Deferred { } } +async function flushMacrotask(): Promise { + await new Promise((resolve) => { + setTimeout(resolve, 0) + }) +} + mock.module('i18next', () => ({ default: { t: (key: string) => key }, })) @@ -80,6 +89,10 @@ mock.module('../src/features/auth/secure-verification/api', () => ({ verify, })) +mock.module('../src/lib/handle-server-error', () => ({ + handleServerError, +})) + const { useSecureVerification } = await import( '../src/features/auth/secure-verification/hooks/use-secure-verification' ) @@ -107,6 +120,7 @@ beforeEach(() => { toastError.mockClear() toastInfo.mockClear() toastSuccess.mockClear() + handleServerError.mockClear() }) afterEach(() => cleanup()) @@ -225,7 +239,7 @@ describe('useSecureVerification', () => { assert.equal(await continuation, null) initialAction.reject(verificationRequiredError) - await Promise.resolve() + await flushMacrotask() assert.equal(checkVerificationMethods.mock.calls.length, 0) assert.equal(result.current.open, false) }) @@ -319,7 +333,7 @@ describe('useSecureVerification', () => { }), ]) discovery.resolve(availableVerificationMethods) - await Promise.resolve() + await flushMacrotask() assert.deepEqual(outcome, { settled: true, value: null }) assert.equal(apiCall.mock.calls.length, 1) @@ -386,8 +400,57 @@ describe('useSecureVerification', () => { /method discovery unavailable/ ) - assert.equal(toastError.mock.calls.length, 1) - assert.equal(toastError.mock.calls[0]?.[0], 'Verification unavailable') + assert.equal(handleServerError.mock.calls.length, 1) + assert.match( + String(handleServerError.mock.calls[0]?.[0]), + /method discovery unavailable/ + ) + assert.deepEqual(handleServerError.mock.calls[0]?.[1], { + fallback: 'Verification unavailable', + }) + assert.equal(toastError.mock.calls.length, 0) + }) + + test('reports a protected action retry failure through the shared handler', async () => { + const retryError = new Error('protected action retry failed') + let callCount = 0 + const apiCall = mock(async () => { + callCount += 1 + if (callCount === 1) throw verificationRequiredError + throw retryError + }) + const { result } = renderHook(() => useSecureVerification()) + + const continuation = result.current.withVerification(apiCall, { + scope: 'api_token', + }) + const continuationOutcome = continuation.then( + (value) => ({ status: 'resolved' as const, value }), + (error: unknown) => ({ status: 'rejected' as const, error }) + ) + await waitFor(() => assert.equal(result.current.open, true)) + + let executionError: unknown + await act(async () => { + try { + await result.current.executeVerification('password', 'secret') + } catch (error) { + executionError = error + } + }) + assert.equal(executionError, retryError) + + const outcome = await continuationOutcome + assert.equal(outcome.status, 'rejected') + if (outcome.status === 'rejected') { + assert.equal(outcome.error, retryError) + } + assert.equal(handleServerError.mock.calls.length, 1) + assert.equal(handleServerError.mock.calls[0]?.[0], retryError) + assert.deepEqual(handleServerError.mock.calls[0]?.[1], { + fallback: 'protected action retry failed', + }) + assert.equal(toastError.mock.calls.length, 0) }) test('clears loading after a retryable verification failure', async () => { From c65ec927b8eae729f7b6edd6d8688788e8f2f8f2 Mon Sep 17 00:00:00 2001 From: CSCITech Date: Fri, 28 Aug 2026 22:44:41 +0800 Subject: [PATCH 14/70] fix bugs --- .../src/lib/handle-server-error.test.ts | 83 ++++++++++++++++++- web/default/src/lib/handle-server-error.ts | 69 ++++++++++----- 2 files changed, 132 insertions(+), 20 deletions(-) diff --git a/web/default/src/lib/handle-server-error.test.ts b/web/default/src/lib/handle-server-error.test.ts index ae46a305..1e83a78e 100644 --- a/web/default/src/lib/handle-server-error.test.ts +++ b/web/default/src/lib/handle-server-error.test.ts @@ -19,7 +19,10 @@ For commercial licensing, please contact https://github.com/MAX-API-Next/MAX-API import { AxiosError, type InternalAxiosRequestConfig } from 'axios' import assert from 'node:assert/strict' import { describe, test } from 'node:test' -import { getSafeErrorDebugInfo } from './handle-server-error' +import { + getSafeErrorDebugInfo, + resolveServerErrorMessage, +} from './handle-server-error' describe('getSafeErrorDebugInfo', () => { test('does not include request headers or config', () => { @@ -38,3 +41,81 @@ describe('getSafeErrorDebugInfo', () => { assert.doesNotMatch(JSON.stringify(debugInfo), /secret-value/) }) }) + +describe('resolveServerErrorMessage', () => { + test('uses an ordinary Error message before the fallback', () => { + assert.equal( + resolveServerErrorMessage( + new Error('account disabled'), + 'Verification unavailable', + 'Content not found.' + ), + 'account disabled' + ) + }) + + test('uses an Axios response message before the fallback', () => { + const error = new AxiosError('request failed', 'ERR_BAD_REQUEST') + error.response = { + status: 503, + statusText: 'Service Unavailable', + headers: {}, + config: {} as InternalAxiosRequestConfig, + data: { message: 'verification service unavailable' }, + } + + assert.equal( + resolveServerErrorMessage( + error, + 'Verification unavailable', + 'Content not found.' + ), + 'verification service unavailable' + ) + }) + + test('keeps an Axios response title ahead of other messages', () => { + const error = new AxiosError('request failed', 'ERR_BAD_REQUEST') + error.response = { + status: 400, + statusText: 'Bad Request', + headers: {}, + config: {} as InternalAxiosRequestConfig, + data: { + title: 'Invalid verification request', + message: 'verification failed', + }, + } + + assert.equal( + resolveServerErrorMessage( + error, + 'Verification unavailable', + 'Content not found.' + ), + 'Invalid verification request' + ) + }) + + test('ignores empty error messages and uses the fallback', () => { + assert.equal( + resolveServerErrorMessage( + new Error(' '), + 'Verification unavailable', + 'Content not found.' + ), + 'Verification unavailable' + ) + }) + + test('preserves the content-not-found message for status 204', () => { + assert.equal( + resolveServerErrorMessage( + { status: 204 }, + 'Verification unavailable', + 'Content not found.' + ), + 'Content not found.' + ) + }) +}) diff --git a/web/default/src/lib/handle-server-error.ts b/web/default/src/lib/handle-server-error.ts index 59a6cd34..b4f6f5de 100644 --- a/web/default/src/lib/handle-server-error.ts +++ b/web/default/src/lib/handle-server-error.ts @@ -31,6 +31,51 @@ interface HandleServerErrorOptions { fallback?: string } +function isNonEmptyString(value: unknown): value is string { + return typeof value === 'string' && value.trim().length > 0 +} + +function getAxiosResponseText( + error: AxiosError, + property: 'title' | 'message' +): string | undefined { + const responseData = error.response?.data + if (!responseData || typeof responseData !== 'object') return undefined + + const value = (responseData as Record)[property] + return isNonEmptyString(value) ? value : undefined +} + +export function resolveServerErrorMessage( + error: unknown, + fallback: string, + contentNotFound: string +): string { + const isContentNotFound = + error !== null && + typeof error === 'object' && + 'status' in error && + Number(error.status) === 204 + + if (error instanceof AxiosError) { + const responseTitle = getAxiosResponseText(error, 'title') + if (responseTitle) return responseTitle + + const responseMessage = getAxiosResponseText(error, 'message') + if (responseMessage) return responseMessage + } + + if ( + !isContentNotFound && + error instanceof Error && + isNonEmptyString(error.message) + ) { + return error.message + } + + return isContentNotFound ? contentNotFound : fallback +} + export function getSafeErrorDebugInfo(error: unknown): SafeErrorDebugInfo { if (error instanceof AxiosError) { return { @@ -55,25 +100,11 @@ export function handleServerError( console.error('[handleServerError]', getSafeErrorDebugInfo(error)) } - let errMsg = options.fallback || i18next.t('Something went wrong!') - - if ( - error && - typeof error === 'object' && - 'status' in error && - Number(error.status) === 204 - ) { - errMsg = i18next.t('Content not found.') - } - - if (error instanceof AxiosError) { - const responseTitle = ( - error.response?.data as { title?: unknown } | undefined - )?.title - if (typeof responseTitle === 'string' && responseTitle) { - errMsg = responseTitle - } - } + const errMsg = resolveServerErrorMessage( + error, + options.fallback || i18next.t('Something went wrong!'), + i18next.t('Content not found.') + ) toast.error(errMsg) } From 64d179a090c5d4c5053b55d19728bb381d5450f6 Mon Sep 17 00:00:00 2001 From: CSCITech Date: Fri, 28 Aug 2026 23:35:18 +0800 Subject: [PATCH 15/70] fix: address final secure verification review findings --- controller/telegram_test.go | 17 +++ router/api_router_test.go | 3 +- .../components/secure-verification-context.ts | 2 +- .../hooks/use-secure-verification.ts | 4 + .../overview/overview-dashboard.tsx | 10 +- .../dialogs/telegram-bind-dialog.tsx | 88 ++++++++----- .../components/profile-security-card.tsx | 2 +- web/default/src/i18n/locales/fr.json | 2 +- web/default/src/i18n/locales/ja.json | 2 +- web/default/src/lib/api.ts | 23 +++- web/default/src/lib/secure-verification.ts | 16 +++ .../api-key-plaintext-lifecycle.test.tsx | 8 +- ...rview-request-copy-error-handling.test.tsx | 123 ++++++++++++++++++ .../tests/secure-verification-api.test.tsx | 24 +++- .../secure-verification-continuation.test.tsx | 10 +- .../telegram-bind-dialog-lifecycle.test.tsx | 71 +++++++++- .../two-fa-setup-shared-verification.test.tsx | 47 ++++++- 17 files changed, 398 insertions(+), 54 deletions(-) create mode 100644 web/default/tests/overview-request-copy-error-handling.test.tsx diff --git a/controller/telegram_test.go b/controller/telegram_test.go index 64be79b0..1616858f 100644 --- a/controller/telegram_test.go +++ b/controller/telegram_test.go @@ -49,3 +49,20 @@ func TestTelegramAuthorizationAcceptsFreshSignedPayload(t *testing.T) { 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 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)) +} diff --git a/router/api_router_test.go b/router/api_router_test.go index 3b4fe82e..43260208 100644 --- a/router/api_router_test.go +++ b/router/api_router_test.go @@ -470,7 +470,8 @@ func TestSensitiveCredentialRoutesRequireStepUp(t *testing.T) { path string body string }{ - {name: "passkey registration", path: "/api/user/passkey/register/begin", body: `{}`}, + {name: "passkey registration begin", path: "/api/user/passkey/register/begin", body: `{}`}, + {name: "passkey registration finish", path: "/api/user/passkey/register/finish", body: `{}`}, {name: "2fa setup", path: "/api/user/2fa/setup", body: `{}`}, {name: "session revocation", path: "/api/user/sessions/revoke", body: `{}`}, {name: "telegram bind state", path: "/api/oauth/telegram/bind/state", body: `{}`}, diff --git a/web/default/src/features/auth/secure-verification/components/secure-verification-context.ts b/web/default/src/features/auth/secure-verification/components/secure-verification-context.ts index 5fc8f113..6ee9296b 100644 --- a/web/default/src/features/auth/secure-verification/components/secure-verification-context.ts +++ b/web/default/src/features/auth/secure-verification/components/secure-verification-context.ts @@ -29,7 +29,7 @@ interface SecureVerificationContextValue { export const SecureVerificationContext = createContext(null) -export function useSecureVerificationGate() { +export function useSecureVerificationGate(): SecureVerificationContextValue { const context = useContext(SecureVerificationContext) if (!context) { throw new Error( 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 eb0488e9..e464682d 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 @@ -23,6 +23,7 @@ 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' @@ -156,6 +157,7 @@ export function useSecureVerification( handleServerError(error, { fallback: i18next.t('Verification unavailable'), }) + markSecureVerificationErrorReported(error) onError?.(error) throw error } @@ -244,6 +246,7 @@ export function useSecureVerification( ? error.message : i18next.t('Verification failed') toast.error(message) + markSecureVerificationErrorReported(error) onError?.(error) throw error } @@ -274,6 +277,7 @@ export function useSecureVerification( ? error.message : i18next.t('Verification failed') handleServerError(error, { fallback: message }) + markSecureVerificationErrorReported(error) onError?.(error) reset() throw error 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 67be6c49..3e363151 100644 --- a/web/default/src/features/dashboard/components/overview/overview-dashboard.tsx +++ b/web/default/src/features/dashboard/components/overview/overview-dashboard.tsx @@ -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' @@ -374,7 +376,7 @@ function StartStepItem(props: { ) } -function RequestPreview(props: { +export function RequestPreview(props: { example: RequestExample signals: HeroSignal[] isRetrying: boolean @@ -422,6 +424,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/profile/components/dialogs/telegram-bind-dialog.tsx b/web/default/src/features/profile/components/dialogs/telegram-bind-dialog.tsx index 1c1acae0..182bfb95 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 @@ -68,24 +68,57 @@ export function TelegramBindDialog({ const [errorMessage, setErrorMessage] = useState('') const [retryVersion, setRetryVersion] = useState(0) - const initializeBinding = useCallback(async () => { + const resetBindingState = useCallback(() => { + setBindState('') + setStatus('idle') + setErrorMessage('') + }, []) + + const handleOpenChange = useCallback( + (nextOpen: boolean) => { + if (!nextOpen) resetBindingState() + onOpenChange(nextOpen) + }, + [onOpenChange, resetBindingState] + ) + const callbackHandlersRef = useRef({ + handleOpenChange, + onSuccess, + t, + withVerification, + }) + useEffect(() => { + callbackHandlersRef.current = { + handleOpenChange, + onSuccess, + t, + withVerification, + } + }, [handleOpenChange, onSuccess, t, withVerification]) + + const initializeBinding = useCallback(async (): Promise => { + const handlers = callbackHandlersRef.current setStatus('loading') setErrorMessage('') try { - const response = await withVerification(createTelegramBindState, { - scope: 'credentials', - title: t('Security verification'), - description: t( - 'Confirm your identity before binding a Telegram account.' - ), - }) + const response = await handlers.withVerification( + createTelegramBindState, + { + scope: 'credentials', + title: handlers.t('Security verification'), + description: handlers.t( + 'Confirm your identity before binding a Telegram account.' + ), + } + ) if (!response) { setStatus('idle') return } if (!response.success || !response.data?.state) { throw new Error( - response.message || t('Failed to initialize Telegram binding') + response.message || + handlers.t('Failed to initialize Telegram binding') ) } setBindState(response.data.state) @@ -93,28 +126,12 @@ export function TelegramBindDialog({ setErrorMessage( error instanceof Error ? error.message - : t('Failed to initialize Telegram binding') + : handlers.t('Failed to initialize Telegram binding') ) setStatus('error') } - }, [t, withVerification]) - - const resetBindingState = useCallback(() => { - setBindState('') - setStatus('idle') - setErrorMessage('') }, []) - const handleOpenChange = useCallback( - (nextOpen: boolean) => { - if (!nextOpen) resetBindingState() - onOpenChange(nextOpen) - }, - [onOpenChange, resetBindingState] - ) - const callbackHandlersRef = useRef({ handleOpenChange, onSuccess }) - callbackHandlersRef.current = { handleOpenChange, onSuccess } - useEffect(() => { if (!open) return const timer = window.setTimeout(() => void initializeBinding(), 0) @@ -131,15 +148,16 @@ export function TelegramBindDialog({ windowCallbacks[callbackName] = async ( authorization: TelegramAuthorizationPayload ) => { + const handlers = callbackHandlersRef.current setStatus('binding') setErrorMessage('') try { - const response = await withVerification( + const response = await handlers.withVerification( () => bindTelegramAccount({ ...authorization, state: bindState }), { scope: 'credentials', - title: t('Security verification'), - description: t( + title: handlers.t('Security verification'), + description: handlers.t( 'Confirm your identity before binding a Telegram account.' ), } @@ -150,17 +168,17 @@ export function TelegramBindDialog({ } if (!response.success) { throw new Error( - response.message || t('Failed to bind Telegram account') + response.message || handlers.t('Failed to bind Telegram account') ) } - toast.success(t('Telegram account bound successfully')) + toast.success(handlers.t('Telegram account bound successfully')) callbackHandlersRef.current.onSuccess() callbackHandlersRef.current.handleOpenChange(false) } catch (error) { setErrorMessage( error instanceof Error ? error.message - : t('Failed to bind Telegram account') + : handlers.t('Failed to bind Telegram account') ) setStatus('error') } @@ -175,7 +193,9 @@ export function TelegramBindDialog({ script.setAttribute('data-onauth', `${callbackName}(user)`) script.onload = () => setStatus('ready') script.onerror = () => { - setErrorMessage(t('Failed to load Telegram Login Widget')) + setErrorMessage( + callbackHandlersRef.current.t('Failed to load Telegram Login Widget') + ) setStatus('error') } container.appendChild(script) @@ -184,7 +204,7 @@ export function TelegramBindDialog({ delete windowCallbacks[callbackName] container.replaceChildren() } - }, [bindState, botName, callbackName, open, t, withVerification]) + }, [bindState, botName, callbackName, open]) const retry = () => { setBindState('') 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 c40106b8..ae45ef19 100644 --- a/web/default/src/features/profile/components/profile-security-card.tsx +++ b/web/default/src/features/profile/components/profile-security-card.tsx @@ -52,7 +52,7 @@ export function ProfileSecurityCard({ const { withVerification } = useSecureVerificationGate() const [revokingSessions, setRevokingSessions] = useState(false) - const handleRevokeOtherSessions = async () => { + const handleRevokeOtherSessions = async (): Promise => { setRevokingSessions(true) try { const response = await withVerification(revokeOtherSessions, { diff --git a/web/default/src/i18n/locales/fr.json b/web/default/src/i18n/locales/fr.json index ffdb986b..81177511 100644 --- a/web/default/src/i18n/locales/fr.json +++ b/web/default/src/i18n/locales/fr.json @@ -5465,7 +5465,7 @@ "Additional verification required": "Vérification supplémentaire requise", "API key verification was cancelled": "La vérification de la clé API a été annulée", "Binding Telegram account...": "Association du compte Telegram...", - "Confirm your identity before accessing this sensitive action.": "Confirmez votre identité avant d’accéder à cette action sensible.", + "Confirm your identity before accessing this sensitive action.": "Confirmez votre identité avant d’effectuer cette action sensible.", "Confirm your identity before binding a Telegram account.": "Confirmez votre identité avant d’associer un compte Telegram.", "Confirm your identity before copying a request with your API key.": "Confirmez votre identité avant de copier une requête contenant votre clé API.", "Confirm your identity before creating or revealing API keys.": "Confirmez votre identité avant de créer ou d’afficher des clés API.", diff --git a/web/default/src/i18n/locales/ja.json b/web/default/src/i18n/locales/ja.json index 4fe5ed51..35e334c2 100644 --- a/web/default/src/i18n/locales/ja.json +++ b/web/default/src/i18n/locales/ja.json @@ -5461,7 +5461,7 @@ "No active alerts.": "アクティブなアラートはありません。", "This process is not currently reporting a sustained CPU, memory, or disk alert.": "このプロセスは現在、継続的な CPU、メモリ、ディスクのアラートを報告していません。", "Add a password or sign in again with a linked account to continue.": "パスワードを追加するか、連携済みアカウントで再度ログインして続行してください。", - "Add a password or sign in again with a linked account to unlock sensitive operations.": "パスワードを追加するか、連携済みアカウントで再度ログインして機密操作を有効にしてください。", + "Add a password or sign in again with a linked account to unlock sensitive operations.": "パスワードを追加するか、連携済みアカウントで再度ログインして機密操作のロックを解除してください。", "Additional verification required": "追加の本人確認が必要です", "API key verification was cancelled": "APIキーの確認がキャンセルされました", "Binding Telegram account...": "Telegramアカウントを連携しています...", diff --git a/web/default/src/lib/api.ts b/web/default/src/lib/api.ts index 31be1f3e..cdcc362a 100644 --- a/web/default/src/lib/api.ts +++ b/web/default/src/lib/api.ts @@ -239,6 +239,21 @@ export const sensitiveActionConfig: ApiRequestConfig = { skipErrorHandler: true, } +export interface TwoFASetupResponse { + success: boolean + message?: string + data?: { + secret: string + qr_code_data: string + backup_codes: string[] + } +} + +export interface TwoFAActionResponse { + success: boolean + message?: string +} + // Get 2FA status export async function get2FAStatus() { const res = await api.get('/api/user/2fa/status') @@ -246,8 +261,8 @@ export async function get2FAStatus() { } // Setup 2FA -export async function setup2FA() { - const res = await api.post( +export async function setup2FA(): Promise { + const res = await api.post( '/api/user/2fa/setup', undefined, sensitiveActionConfig @@ -256,8 +271,8 @@ export async function setup2FA() { } // Enable 2FA with verification code -export async function enable2FA(code: string) { - const res = await api.post( +export async function enable2FA(code: string): Promise { + const res = await api.post( '/api/user/2fa/enable', { code }, sensitiveActionConfig diff --git a/web/default/src/lib/secure-verification.ts b/web/default/src/lib/secure-verification.ts index 71c2296f..e347ed68 100644 --- a/web/default/src/lib/secure-verification.ts +++ b/web/default/src/lib/secure-verification.ts @@ -18,6 +18,22 @@ For commercial licensing, please contact https://github.com/MAX-API-Next/MAX-API */ import type { AxiosError } from 'axios' +const reportedSecureVerificationErrors = new WeakSet() + +function isTrackableError(error: unknown): error is object { + return ( + (typeof error === 'object' && error !== null) || typeof error === 'function' + ) +} + +export function markSecureVerificationErrorReported(error: unknown): void { + if (isTrackableError(error)) reportedSecureVerificationErrors.add(error) +} + +export function wasSecureVerificationErrorReported(error: unknown): boolean { + return isTrackableError(error) && reportedSecureVerificationErrors.has(error) +} + export interface VerificationRequiredInfo { code?: string message: string diff --git a/web/default/tests/api-key-plaintext-lifecycle.test.tsx b/web/default/tests/api-key-plaintext-lifecycle.test.tsx index cf871b00..a6c7076d 100644 --- a/web/default/tests/api-key-plaintext-lifecycle.test.tsx +++ b/web/default/tests/api-key-plaintext-lifecycle.test.tsx @@ -160,8 +160,10 @@ describe('API key plaintext lifecycle', () => { fireEvent.click(view.getByRole('button', { name: 'Copy API key' })) - await waitFor(() => assert.equal(copyToClipboard.mock.calls.length, 1)) - assert.equal(toastError.mock.calls.length, 1) - assert.equal(toastError.mock.calls[0]?.[0], 'Failed to copy to clipboard') + await waitFor(() => { + assert.equal(copyToClipboard.mock.calls.length, 1) + assert.equal(toastError.mock.calls.length, 1) + assert.equal(toastError.mock.calls[0]?.[0], 'Failed to copy to clipboard') + }) }) }) diff --git a/web/default/tests/overview-request-copy-error-handling.test.tsx b/web/default/tests/overview-request-copy-error-handling.test.tsx new file mode 100644 index 00000000..13eb9d28 --- /dev/null +++ b/web/default/tests/overview-request-copy-error-handling.test.tsx @@ -0,0 +1,123 @@ +/* +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 { createReactTestEnvironment } from '@/test/react' +import { markSecureVerificationErrorReported } from '@/lib/secure-verification' +import { cleanup, fireEvent, render, waitFor } from '@testing-library/react' +import { + afterAll, + afterEach, + beforeAll, + beforeEach, + mock, + test, +} from 'bun:test' +import assert from 'node:assert/strict' + +let verificationError = new Error('request failed') +const withVerification = mock(async () => { + throw verificationError +}) +const handleServerError = mock( + (_error: unknown, _options?: { fallback?: string }) => undefined +) + +mock.module('react-i18next', () => ({ + useTranslation: () => ({ t: (key: string) => key }), +})) + +mock.module('sonner', () => ({ + toast: { + error: mock(() => undefined), + success: mock(() => undefined), + }, +})) + +mock.module('@/features/auth/secure-verification', () => ({ + useSecureVerificationGate: () => ({ withVerification }), +})) + +mock.module('@/lib/handle-server-error', () => ({ handleServerError })) + +const { RequestPreview } = await import( + '../src/features/dashboard/components/overview/overview-dashboard' +) +const testEnv = createReactTestEnvironment() + +beforeAll(() => testEnv.setup()) +beforeEach(() => { + verificationError = new Error('request failed') + withVerification.mockClear() + handleServerError.mockClear() +}) +afterEach(() => cleanup()) +afterAll(() => testEnv.teardown()) + +function renderRequestPreview(): ReturnType { + return render( + undefined)} + /> + ) +} + +async function waitForCopyToSettle( + view: ReturnType +): Promise { + await waitFor(() => { + const button = view.getByRole('button', { + name: 'Copy ready-to-run curl', + }) + assert.equal(button.hasAttribute('disabled'), false) + }) +} + +test('reports an unreported request-copy rejection and consumes it', async () => { + const view = renderRequestPreview() + + fireEvent.click(view.getByRole('button', { name: 'Copy ready-to-run curl' })) + + await waitFor(() => assert.equal(handleServerError.mock.calls.length, 1)) + await waitForCopyToSettle(view) + assert.equal(handleServerError.mock.calls[0]?.[0], verificationError) + assert.deepEqual(handleServerError.mock.calls[0]?.[1], { + fallback: 'Failed to copy to clipboard', + }) +}) + +test('does not duplicate a rejection already reported by verification', async () => { + markSecureVerificationErrorReported(verificationError) + const view = renderRequestPreview() + + fireEvent.click(view.getByRole('button', { name: 'Copy ready-to-run curl' })) + + await waitFor(() => assert.equal(withVerification.mock.calls.length, 1)) + await waitForCopyToSettle(view) + assert.equal(handleServerError.mock.calls.length, 0) +}) diff --git a/web/default/tests/secure-verification-api.test.tsx b/web/default/tests/secure-verification-api.test.tsx index a27c9308..51be641e 100644 --- a/web/default/tests/secure-verification-api.test.tsx +++ b/web/default/tests/secure-verification-api.test.tsx @@ -27,6 +27,24 @@ interface VerificationMethodsRequestConfig { skipErrorHandler?: boolean } +type VerificationMethodsResponse = + | { + data: { + success: true + data: { + has_2fa: boolean + has_passkey: boolean + has_password: boolean + } + } + } + | { + data: { + success: false + message: string + } + } + const successfulMethodsResponse = { data: { success: true, @@ -39,8 +57,10 @@ const successfulMethodsResponse = { } const apiGet = mock( - async (_url: string, _config?: VerificationMethodsRequestConfig) => - successfulMethodsResponse + async ( + _url: string, + _config?: VerificationMethodsRequestConfig + ): Promise => successfulMethodsResponse ) const consoleError = spyOn(console, 'error').mockImplementation(() => undefined) diff --git a/web/default/tests/secure-verification-continuation.test.tsx b/web/default/tests/secure-verification-continuation.test.tsx index c2840457..8dc31313 100644 --- a/web/default/tests/secure-verification-continuation.test.tsx +++ b/web/default/tests/secure-verification-continuation.test.tsx @@ -17,6 +17,7 @@ along with this program. If not, see . For commercial licensing, please contact https://github.com/MAX-API-Next/MAX-API/issues */ import { createReactTestEnvironment } from '@/test/react' +import { wasSecureVerificationErrorReported } from '@/lib/secure-verification' import { act, cleanup, renderHook, waitFor } from '@testing-library/react' import { afterAll, @@ -144,8 +145,9 @@ describe('useSecureVerification', () => { }) test('rejects an initial protected action failure without discovery', async () => { + const actionError = new Error('protected action failed') const apiCall = mock(async () => { - throw new Error('protected action failed') + throw actionError }) const { result } = renderHook(() => useSecureVerification()) @@ -156,6 +158,7 @@ describe('useSecureVerification', () => { assert.equal(apiCall.mock.calls.length, 1) assert.equal(checkVerificationMethods.mock.calls.length, 0) assert.equal(result.current.open, false) + assert.equal(wasSecureVerificationErrorReported(actionError), false) }) test('continues the original API call after verification succeeds', async () => { @@ -387,8 +390,9 @@ describe('useSecureVerification', () => { }) test('reports verification method discovery failures before rejecting', async () => { + const discoveryError = new Error('method discovery unavailable') checkVerificationMethods.mockImplementationOnce(async () => { - throw new Error('method discovery unavailable') + throw discoveryError }) const apiCall = mock(async () => { throw verificationRequiredError @@ -408,6 +412,7 @@ describe('useSecureVerification', () => { assert.deepEqual(handleServerError.mock.calls[0]?.[1], { fallback: 'Verification unavailable', }) + assert.equal(wasSecureVerificationErrorReported(discoveryError), true) assert.equal(toastError.mock.calls.length, 0) }) @@ -450,6 +455,7 @@ describe('useSecureVerification', () => { assert.deepEqual(handleServerError.mock.calls[0]?.[1], { fallback: 'protected action retry failed', }) + assert.equal(wasSecureVerificationErrorReported(retryError), true) assert.equal(toastError.mock.calls.length, 0) }) diff --git a/web/default/tests/telegram-bind-dialog-lifecycle.test.tsx b/web/default/tests/telegram-bind-dialog-lifecycle.test.tsx index 06ff65ff..f07bb8dd 100644 --- a/web/default/tests/telegram-bind-dialog-lifecycle.test.tsx +++ b/web/default/tests/telegram-bind-dialog-lifecycle.test.tsx @@ -19,7 +19,15 @@ For commercial licensing, please contact https://github.com/MAX-API-Next/MAX-API import type { ReactNode } from 'react' import { createReactTestEnvironment } from '@/test/react' import { act, cleanup, render, waitFor } from '@testing-library/react' -import { afterAll, afterEach, beforeAll, describe, mock, test } from 'bun:test' +import { + afterAll, + afterEach, + beforeAll, + beforeEach, + describe, + mock, + test, +} from 'bun:test' import assert from 'node:assert/strict' const createTelegramBindState = mock(async () => ({ @@ -30,9 +38,14 @@ const bindTelegramAccount = mock(async () => ({ success: true })) const toastSuccess = mock(() => undefined) const translate = (key: string) => key const withVerification = (apiCall: () => Promise) => apiCall() +let useUnstableVerificationGate = false mock.module('../src/features/auth/secure-verification', () => ({ - useSecureVerificationGate: () => ({ withVerification }), + useSecureVerificationGate: () => ({ + withVerification: useUnstableVerificationGate + ? (apiCall: () => Promise) => apiCall() + : withVerification, + }), })) mock.module('../src/features/profile/api', () => ({ @@ -87,6 +100,12 @@ const { TelegramBindDialog } = await import( const testEnv = createReactTestEnvironment() beforeAll(() => testEnv.setup()) +beforeEach(() => { + useUnstableVerificationGate = false + createTelegramBindState.mockClear() + bindTelegramAccount.mockClear() + toastSuccess.mockClear() +}) afterEach(() => cleanup()) afterAll(() => testEnv.teardown()) @@ -157,4 +176,52 @@ describe('TelegramBindDialog widget lifecycle', () => { assert.equal(secondSuccess.mock.calls.length, 1) assert.deepEqual(secondOpenChange.mock.calls[0], [false]) }) + + test('keeps the active widget mounted across gate identity changes', async () => { + useUnstableVerificationGate = true + const view = render( + undefined)} + botName='max_api_bot' + onSuccess={mock(() => undefined)} + /> + ) + + const initialScript = await waitFor(() => { + const script = document.querySelector( + 'script[data-onauth]' + ) + assert.ok(script) + return script + }) + const callbackName = initialScript + .getAttribute('data-onauth') + ?.replace(/\(user\)$/, '') + assert.ok(callbackName) + const initialCallback = (window as unknown as Record)[ + callbackName + ] + assert.equal(typeof initialCallback, 'function') + + view.rerender( + undefined)} + botName='max_api_bot' + onSuccess={mock(() => undefined)} + /> + ) + + await act(async () => Promise.resolve()) + assert.equal( + document.querySelector('script[data-onauth]'), + initialScript + ) + assert.equal( + (window as unknown as Record)[callbackName], + initialCallback + ) + assert.equal(createTelegramBindState.mock.calls.length, 1) + }) }) diff --git a/web/default/tests/two-fa-setup-shared-verification.test.tsx b/web/default/tests/two-fa-setup-shared-verification.test.tsx index 7900b1bd..c9db0d48 100644 --- a/web/default/tests/two-fa-setup-shared-verification.test.tsx +++ b/web/default/tests/two-fa-setup-shared-verification.test.tsx @@ -17,7 +17,7 @@ along with this program. If not, see . For commercial licensing, please contact https://github.com/MAX-API-Next/MAX-API/issues */ import { createReactTestEnvironment } from '@/test/react' -import { cleanup, render, waitFor } from '@testing-library/react' +import { act, cleanup, render, waitFor } from '@testing-library/react' import { afterAll, afterEach, beforeAll, beforeEach, mock, test } from 'bun:test' import assert from 'node:assert/strict' @@ -80,3 +80,48 @@ test('uses the authenticated layout verification gate for 2FA setup', async () = await waitFor(() => assert.equal(setup2FA.mock.calls.length, 1)) assert.equal(withVerification.mock.calls.length, 1) }) + +test('does not restart 2FA setup while initialization is pending', async () => { + let resolveSetup: + | ((value: Awaited>) => void) + | undefined + setup2FA.mockImplementationOnce( + () => + new Promise((resolve) => { + resolveSetup = resolve + }) + ) + + const view = render( + undefined)} + onSuccess={mock(() => undefined)} + /> + ) + await waitFor(() => assert.equal(setup2FA.mock.calls.length, 1)) + + view.rerender( + undefined)} + onSuccess={mock(() => undefined)} + /> + ) + await act(async () => { + await new Promise((resolve) => setTimeout(resolve, 0)) + }) + assert.equal(setup2FA.mock.calls.length, 1) + + await act(async () => { + resolveSetup?.({ + success: true, + data: { + qr_code_data: 'otpauth://totp/MAX-API:test', + secret: 'TESTSECRET', + backup_codes: ['ABCD-EFGH'], + }, + }) + await Promise.resolve() + }) +}) From c48a5f757b65a6f61b85c4067e15339f69b0f053 Mon Sep 17 00:00:00 2001 From: CSCITech Date: Sat, 29 Aug 2026 01:41:22 +0800 Subject: [PATCH 16/70] fix: resolve final PR 38 review findings --- controller/oauth.go | 9 +- controller/oauth_test.go | 30 ++- controller/passkey.go | 14 +- controller/telegram.go | 21 +- controller/telegram_test.go | 74 ++++++ controller/twofa.go | 25 +- controller/user.go | 42 +++- controller/user_session_generation_test.go | 71 +++++- docs/openapi/api.json | 14 +- model/passkey.go | 6 +- model/security_credentials_test.go | 4 +- relay/relay_task_test.go | 28 ++- router/api_router_test.go | 8 + service/billing.go | 3 + service/billing_session_test.go | 62 +++++ service/quota.go | 32 ++- .../chat/hooks/use-active-chat-key.ts | 2 +- .../components/data-table-row-actions.tsx | 23 +- .../dialogs/telegram-bind-dialog.tsx | 165 +------------ .../dialogs/two-fa-setup-dialog.tsx | 15 +- .../hooks/use-telegram-login-widget.ts | 217 ++++++++++++++++++ web/default/src/i18n/locales/ja.json | 2 +- .../src/lib/handle-server-error.test.ts | 13 ++ web/default/src/lib/handle-server-error.ts | 2 + .../api-key-plaintext-lifecycle.test.tsx | 7 + ...i-key-row-actions-chat-navigation.test.tsx | 209 +++++++++++++++++ ...rview-request-copy-error-handling.test.tsx | 9 +- ...sskey-registration-error-handling.test.tsx | 2 + .../secure-verification-continuation.test.tsx | 2 +- .../two-fa-setup-shared-verification.test.tsx | 42 +++- 30 files changed, 912 insertions(+), 241 deletions(-) create mode 100644 web/default/src/features/profile/hooks/use-telegram-login-widget.ts create mode 100644 web/default/tests/api-key-row-actions-chat-navigation.test.tsx diff --git a/controller/oauth.go b/controller/oauth.go index 61d8725a..a552750e 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 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..0bc9507a 100644 --- a/controller/oauth_test.go +++ b/controller/oauth_test.go @@ -31,18 +31,46 @@ func (*deletedUserOAuthProvider) ExchangeToken(context.Context, string, *gin.Con func setupControllerAuthFlowDB(t *testing.T) { t.Helper() originalDB := model.DB + originalLogDB := model.LOG_DB db, err := gorm.Open(sqlite.Open("file:controller_auth_flow?mode=memory&cache=shared"), &gorm.Config{}) require.NoError(t, err) - require.NoError(t, db.AutoMigrate(&model.AuthFlow{})) + 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) + 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 TestGenerateOAuthCodeCreatesProviderBoundOpaqueState(t *testing.T) { setupControllerAuthFlowDB(t) gin.SetMode(gin.TestMode) diff --git a/controller/passkey.go b/controller/passkey.go index 279f7b8e..5ad63789 100644 --- a/controller/passkey.go +++ b/controller/passkey.go @@ -135,10 +135,7 @@ func PasskeyRegisterFinish(c *gin.Context) { common.ApiError(c, err) return } - if err := preserveCurrentSessionAfterSecurityChange(c, user.Id, generation); 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{ @@ -166,10 +163,7 @@ func PasskeyDelete(c *gin.Context) { common.ApiError(c, err) return } - if err := preserveCurrentSessionAfterSecurityChange(c, user.Id, generation); 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{ @@ -333,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 } @@ -490,7 +484,7 @@ func PasskeyVerifyFinish(c *gin.Context) { // 更新凭证的最后使用时间 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 } diff --git a/controller/telegram.go b/controller/telegram.go index d7f3e7ee..71dffa14 100644 --- a/controller/telegram.go +++ b/controller/telegram.go @@ -26,6 +26,12 @@ const ( 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"` @@ -121,10 +127,7 @@ func TelegramBind(c *gin.Context) { return } clearOAuthSessionState(c) - if err := preserveCurrentSessionAfterSecurityChange(c, userID, generation); err != nil { - common.ApiError(c, err) - return - } + preserveCurrentSessionAfterCommittedSecurityChange(c, userID, generation, "binding Telegram") recordUserSecurityAudit(c, userID, "user.telegram_bind", nil) common.ApiSuccess(c, nil) } @@ -139,10 +142,20 @@ func TelegramLogin(c *gin.Context) { c.JSON(http.StatusForbidden, gin.H{"success": false, "message": "无效或已过期的 Telegram 授权"}) return } + 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) } diff --git a/controller/telegram_test.go b/controller/telegram_test.go index 1616858f..a318fa83 100644 --- a/controller/telegram_test.go +++ b/controller/telegram_test.go @@ -4,6 +4,8 @@ import ( "crypto/hmac" "crypto/sha256" "encoding/hex" + "net/http" + "net/http/httptest" "net/url" "sort" "strconv" @@ -11,6 +13,12 @@ import ( "testing" "time" + "github.com/MAX-API-Next/MAX-API/common" + appi18n "github.com/MAX-API-Next/MAX-API/i18n" + "github.com/MAX-API-Next/MAX-API/model" + "github.com/gin-contrib/sessions" + "github.com/gin-contrib/sessions/cookie" + "github.com/gin-gonic/gin" "github.com/stretchr/testify/require" ) @@ -66,3 +74,69 @@ func TestTelegramAuthorizationRejectsAuthDateBeyondClockSkew(t *testing.T) { require.False(t, checkTelegramAuthorizationAt(params, token, now)) } + +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) +} diff --git a/controller/twofa.go b/controller/twofa.go index ca0c955b..bdee455d 100644 --- a/controller/twofa.go +++ b/controller/twofa.go @@ -2,7 +2,6 @@ package controller import ( "errors" - "fmt" "net/http" "strconv" @@ -192,13 +191,7 @@ func Enable2FA(c *gin.Context) { common.ApiError(c, err) return } - if err := preserveCurrentSessionAfterSecurityChange(c, userId, generation); err != nil { - common.SysLog(fmt.Sprintf( - "failed to preserve current session after enabling 2FA for user %d: %v", - userId, - err, - )) - } + preserveCurrentSessionAfterCommittedSecurityChange(c, userId, generation, "enabling 2FA") // 记录操作日志 model.RecordLog(userId, model.LogTypeSystem, "成功启用两步验证") @@ -272,13 +265,7 @@ func Disable2FA(c *gin.Context) { common.ApiError(c, err) return } - if err := preserveCurrentSessionAfterSecurityChange(c, userId, generation); err != nil { - common.SysLog(fmt.Sprintf( - "failed to preserve current session after disabling 2FA for user %d: %v", - userId, - err, - )) - } + preserveCurrentSessionAfterCommittedSecurityChange(c, userId, generation, "disabling 2FA") // 记录操作日志 model.RecordLog(userId, model.LogTypeSystem, "禁用两步验证") @@ -399,13 +386,7 @@ func RegenerateBackupCodes(c *gin.Context) { common.SysLog("保存备用码失败: " + err.Error()) return } - if err := preserveCurrentSessionAfterSecurityChange(c, userId, generation); err != nil { - common.SysLog(fmt.Sprintf( - "failed to preserve current session after regenerating 2FA backup codes for user %d: %v", - userId, - err, - )) - } + preserveCurrentSessionAfterCommittedSecurityChange(c, userId, generation, "regenerating 2FA backup codes") // 记录操作日志 model.RecordLog(userId, model.LogTypeSystem, "重新生成两步验证备用码") diff --git a/controller/user.go b/controller/user.go index eb1313f6..662e4a6b 100644 --- a/controller/user.go +++ b/controller/user.go @@ -193,10 +193,7 @@ func RevokeOtherSessions(c *gin.Context) { common.ApiError(c, err) return } - if err := preserveCurrentSessionAfterSecurityChange(c, userID, generation); 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) } @@ -907,10 +904,7 @@ func UpdateSelf(c *gin.Context) { return } if updatePassword { - if err := preserveCurrentSessionAfterSecurityChange(c, cleanUser.Id, cleanUser.SessionGeneration); err != nil { - common.ApiError(c, err) - return - } + preserveCurrentSessionAfterCommittedSecurityChange(c, cleanUser.Id, cleanUser.SessionGeneration, "changing password") } c.JSON(http.StatusOK, gin.H{ @@ -932,6 +926,38 @@ func preserveCurrentSessionAfterSecurityChange(c *gin.Context, userID int, gener 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/controller/user_session_generation_test.go b/controller/user_session_generation_test.go index 6a765a68..bd3b8b17 100644 --- a/controller/user_session_generation_test.go +++ b/controller/user_session_generation_test.go @@ -1,11 +1,14 @@ package controller import ( + "bytes" "net/http" "net/http/httptest" "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" @@ -24,6 +27,7 @@ func TestPreserveCurrentSessionAfterSecurityChangeUpdatesGenerationAndClearsStep session.Set(secureVerificationMethodSessionKey, "password") session.Set(secureVerificationUserSessionKey, 9911) session.Set(secureVerificationScopeSessionKey, secureVerificationScopeCredentials) + session.Set(PasskeyReadySessionKey, time.Now().Unix()) require.NoError(t, session.Save()) c.Status(http.StatusNoContent) }) @@ -37,6 +41,10 @@ func TestPreserveCurrentSessionAfterSecurityChangeUpdatesGenerationAndClearsStep "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), }) }) @@ -59,5 +67,66 @@ func TestPreserveCurrentSessionAfterSecurityChangeUpdatesGenerationAndClearsStep 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}`, inspectRecorder.Body.String()) + 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 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 } diff --git a/docs/openapi/api.json b/docs/openapi/api.json index 2a537109..fee3c17a 100644 --- a/docs/openapi/api.json +++ b/docs/openapi/api.json @@ -1019,11 +1019,21 @@ "get": { "summary": "Telegram登录", "deprecated": false, - "description": "🔓 无需鉴权（OAuth回调）", + "description": "🔓 无需鉴权（OAuth回调）；需要先通过 /api/oauth/state 创建并携带当前浏览器会话绑定的一次性 state", "tags": [ "OAuth" ], - "parameters": [], + "parameters": [ + { + "name": "state", + "in": "query", + "description": "由 /api/oauth/state 为 Telegram 登录创建的浏览器会话绑定一次性 state", + "required": true, + "schema": { + "type": "string" + } + } + ], "responses": { "200": { "description": "成功", diff --git a/model/passkey.go b/model/passkey.go index e38b6728..bcd134a1 100644 --- a/model/passkey.go +++ b/model/passkey.go @@ -178,9 +178,9 @@ func GetPasskeyByCredentialID(credentialID []byte) (*PasskeyCredential, error) { return &credential, nil } -func UpsertPasskeyCredential(credential *PasskeyCredential) error { +func UpdatePasskeyCredentialAfterAuthentication(credential *PasskeyCredential) error { if credential == nil || credential.UserID <= 0 || credential.CredentialID == "" { - common.SysLog("UpsertPasskeyCredential: nil credential provided") + common.SysLog("UpdatePasskeyCredentialAfterAuthentication: nil credential provided") return fmt.Errorf("Passkey 保存失败，请重试") } result := DB.Model(&PasskeyCredential{}). @@ -201,7 +201,7 @@ func UpsertPasskeyCredential(credential *PasskeyCredential) error { ). Updates(credential) if result.Error != nil { - common.SysLog(fmt.Sprintf("UpsertPasskeyCredential: failed to update credential for user %d: %v", credential.UserID, result.Error)) + 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 { diff --git a/model/security_credentials_test.go b/model/security_credentials_test.go index e2beeea0..8ced71a5 100644 --- a/model/security_credentials_test.go +++ b/model/security_credentials_test.go @@ -173,7 +173,7 @@ func TestPasskeyUsageUpdateCannotRestoreReplacedCredential(t *testing.T) { now := time.Now() staleUsageUpdate.SignCount = 2 staleUsageUpdate.LastUsedAt = &now - require.Error(t, UpsertPasskeyCredential(&staleUsageUpdate)) + require.Error(t, UpdatePasskeyCredentialAfterAuthentication(&staleUsageUpdate)) stored, err := GetPasskeyByUserID(user.Id) require.NoError(t, err) @@ -186,7 +186,7 @@ func TestPasskeyUsageUpdateCannotRestoreReplacedCredential(t *testing.T) { matchingUsageUpdate.ID = 0 matchingUsageUpdate.SignCount = 11 matchingUsageUpdate.LastUsedAt = &later - require.NoError(t, UpsertPasskeyCredential(&matchingUsageUpdate)) + require.NoError(t, UpdatePasskeyCredentialAfterAuthentication(&matchingUsageUpdate)) stored, err = GetPasskeyByUserID(user.Id) require.NoError(t, err) require.Equal(t, replacement.CredentialID, stored.CredentialID) diff --git a/relay/relay_task_test.go b/relay/relay_task_test.go index e50d9991..c89827ec 100644 --- a/relay/relay_task_test.go +++ b/relay/relay_task_test.go @@ -25,9 +25,14 @@ import ( type recordingTaskReservationBilling struct { preConsumed int reserveTo []int + settleTo []int } -func (b *recordingTaskReservationBilling) Settle(int) error { return nil } +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 } @@ -55,6 +60,27 @@ func TestReserveTaskQuotaUsesActualReservationAndExtendsRetries(t *testing.T) { 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 withRelayTaskRateCards(t *testing.T, cards map[string]task_billing_setting.RateCard) { t.Helper() original := task_billing_setting.GetRateCardsCopy() diff --git a/router/api_router_test.go b/router/api_router_test.go index 43260208..9bd6f648 100644 --- a/router/api_router_test.go +++ b/router/api_router_test.go @@ -377,6 +377,14 @@ func TestSensitiveCredentialOpenAPIContracts(t *testing.T) { require.Contains(t, properties, property) } + telegramLogin := document.Paths["/api/oauth/telegram/login"]["get"] + parameters := telegramLogin["parameters"].([]any) + require.Len(t, parameters, 1) + stateParameter := parameters[0].(map[string]any) + require.Equal(t, "state", stateParameter["name"]) + require.Equal(t, "query", stateParameter["in"]) + require.Equal(t, true, stateParameter["required"]) + for path, method := range map[string]string{ "/api/oauth/telegram/bind/state": "post", "/api/user/sessions/revoke": "post", diff --git a/service/billing.go b/service/billing.go index bd862d80..b6f08fba 100644 --- a/service/billing.go +++ b/service/billing.go @@ -175,6 +175,9 @@ func SettleBillingWithEffect(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, // 回退：无 BillingSession 时使用旧路径 quotaDelta := actualQuota - relayInfo.FinalPreConsumedQuota + if effect != nil && relayInfo.RequestId != "" { + return postConsumeQuotaOnceWithEffect(relayInfo, "finalize", quotaDelta, relayInfo.FinalPreConsumedQuota, true, effect) + } if quotaDelta != 0 { return false, PostConsumeQuotaOnce(relayInfo, "finalize", quotaDelta, relayInfo.FinalPreConsumedQuota, true) } diff --git a/service/billing_session_test.go b/service/billing_session_test.go index a6288f9e..057642f1 100644 --- a/service/billing_session_test.go +++ b/service/billing_session_test.go @@ -148,6 +148,68 @@ func TestSettleAndRecordConsumeCarriesZeroUsageLogInDurableEffect(t *testing.T) assert.Equal(t, "zero-usage-upstream-request", settler.effect.UpstreamRequestID) } +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 (s *recordingBillingSettler) Settle(int) error { s.settleCalls++; return nil } func (s *recordingBillingSettler) Refund(*gin.Context) {} func (s *recordingBillingSettler) NeedsRefund() bool { return false } diff --git a/service/quota.go b/service/quota.go index 4d7ffdd1..eef324ea 100644 --- a/service/quota.go +++ b/service/quota.go @@ -418,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 @@ -448,9 +460,15 @@ func PostConsumeQuotaOnce(relayInfo *relaycommon.RelayInfo, operationKind string TokenDelta: tokenDelta, SubscriptionPreConsumeRequestID: relayInfo.RequestId, AllowMissingToken: quota < 0, + Effect: effect, }) if durableErr != nil { - return durableErr + return effect != nil, durableErr + } + if effect != nil { + if effectErr := model.ProcessBillingSettlementEffect(operationKey); effectErr != nil { + return true, effectErr + } } if sendEmail && (quota+preConsumedQuota) != 0 { if relayInfo.BillingSource == BillingSourceSubscription { @@ -459,10 +477,10 @@ func PostConsumeQuotaOnce(relayInfo *relaycommon.RelayInfo, operationKind string checkAndSendQuotaNotify(relayInfo, quota, preConsumedQuota) } } - return nil + return effect != nil, 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/web/default/src/features/chat/hooks/use-active-chat-key.ts b/web/default/src/features/chat/hooks/use-active-chat-key.ts index 02bee671..bcd9abaf 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 @@ -50,7 +50,7 @@ export async function fetchActiveChatKey() { export function useActiveChatKey(enabled: boolean) { const userId = useAuthStore((state) => state.auth.user?.id) const { withVerification } = useSecureVerificationGate() - const queryFn = useCallback(async () => { + const queryFn = useCallback(async (): Promise => { const key = await withVerification(fetchActiveChatKey, { scope: 'api_token', title: i18next.t('Security verification'), 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 e8b12e43..50871359 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 @@ -95,8 +95,21 @@ export function DataTableRowActions({ const handleOpenChatPreset = useCallback( async (preset: ChatPreset) => { + let chatWindow: Window | null = null + if (preset.type === 'web' && typeof window !== 'undefined') { + try { + chatWindow = window.open('about:blank', '_blank') + if (chatWindow) chatWindow.opener = null + } catch { + chatWindow = null + } + } + const realKey = await resolveRealKey(apiKey.id) - if (!realKey) return + if (!realKey) { + chatWindow?.close() + return + } if (preset.type === 'fluent') { const success = sendToFluent(realKey, serverAddress) @@ -119,6 +132,7 @@ export function DataTableRowActions({ }) if (!resolvedUrl) { + chatWindow?.close() toast.error(t('Invalid chat link. Please contact your administrator.')) return } @@ -126,8 +140,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 } }, 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 182bfb95..65ebf49f 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,10 +16,8 @@ along with this program. If not, see . For commercial licensing, please contact https://github.com/MAX-API-Next/MAX-API/issues */ -import { useCallback, useEffect, useRef, useState } from 'react' import { RefreshCw, Send } from 'lucide-react' import { useTranslation } from 'react-i18next' -import { toast } from 'sonner' import { cn } from '@/lib/utils' import { Alert, AlertDescription } from '@/components/ui/alert' import { Button } from '@/components/ui/button' @@ -31,12 +29,7 @@ import { DialogTitle, } from '@/components/ui/dialog' import { Spinner } from '@/components/ui/spinner' -import { useSecureVerificationGate } from '@/features/auth/secure-verification' -import { - bindTelegramAccount, - createTelegramBindState, - type TelegramAuthorizationPayload, -} from '../../api' +import { useTelegramLoginWidget } from '../../hooks/use-telegram-login-widget' // ============================================================================ // Telegram Bind Dialog Component @@ -56,160 +49,8 @@ export function TelegramBindDialog({ onSuccess, }: TelegramBindDialogProps) { const { t } = useTranslation() - const { withVerification } = useSecureVerificationGate() - const widgetContainerRef = useRef(null) - const [callbackName] = useState( - () => `telegramAuthCallback_${Math.random().toString(36).slice(2)}` - ) - const [bindState, setBindState] = useState('') - const [status, setStatus] = useState< - 'idle' | 'loading' | 'ready' | 'binding' | 'error' - >('idle') - const [errorMessage, setErrorMessage] = useState('') - const [retryVersion, setRetryVersion] = useState(0) - - const resetBindingState = useCallback(() => { - setBindState('') - setStatus('idle') - setErrorMessage('') - }, []) - - const handleOpenChange = useCallback( - (nextOpen: boolean) => { - if (!nextOpen) resetBindingState() - onOpenChange(nextOpen) - }, - [onOpenChange, resetBindingState] - ) - const callbackHandlersRef = useRef({ - handleOpenChange, - onSuccess, - t, - withVerification, - }) - useEffect(() => { - callbackHandlersRef.current = { - handleOpenChange, - onSuccess, - t, - withVerification, - } - }, [handleOpenChange, onSuccess, t, withVerification]) - - const initializeBinding = useCallback(async (): Promise => { - const handlers = callbackHandlersRef.current - setStatus('loading') - setErrorMessage('') - try { - const response = await handlers.withVerification( - createTelegramBindState, - { - scope: 'credentials', - title: handlers.t('Security verification'), - description: handlers.t( - 'Confirm your identity before binding a Telegram account.' - ), - } - ) - if (!response) { - setStatus('idle') - return - } - if (!response.success || !response.data?.state) { - throw new Error( - response.message || - handlers.t('Failed to initialize Telegram binding') - ) - } - setBindState(response.data.state) - } catch (error) { - setErrorMessage( - error instanceof Error - ? error.message - : handlers.t('Failed to initialize Telegram binding') - ) - setStatus('error') - } - }, []) - - useEffect(() => { - if (!open) return - const timer = window.setTimeout(() => void initializeBinding(), 0) - return () => window.clearTimeout(timer) - }, [initializeBinding, open, retryVersion]) - - useEffect(() => { - if (!open || !bindState || !widgetContainerRef.current) return - - const container = widgetContainerRef.current - const windowCallbacks = window as unknown as Record - container.replaceChildren() - - windowCallbacks[callbackName] = async ( - authorization: TelegramAuthorizationPayload - ) => { - const handlers = callbackHandlersRef.current - setStatus('binding') - setErrorMessage('') - try { - const response = await handlers.withVerification( - () => bindTelegramAccount({ ...authorization, state: bindState }), - { - scope: 'credentials', - title: handlers.t('Security verification'), - description: handlers.t( - 'Confirm your identity before binding a Telegram account.' - ), - } - ) - if (!response) { - setStatus('ready') - return - } - if (!response.success) { - throw new Error( - response.message || handlers.t('Failed to bind Telegram account') - ) - } - toast.success(handlers.t('Telegram account bound successfully')) - callbackHandlersRef.current.onSuccess() - callbackHandlersRef.current.handleOpenChange(false) - } catch (error) { - setErrorMessage( - error instanceof Error - ? error.message - : handlers.t('Failed to bind Telegram account') - ) - setStatus('error') - } - } - - const script = document.createElement('script') - script.async = true - script.src = 'https://telegram.org/js/telegram-widget.js?22' - script.setAttribute('data-telegram-login', botName.replace(/^@/, '')) - script.setAttribute('data-size', 'large') - script.setAttribute('data-userpic', 'false') - script.setAttribute('data-onauth', `${callbackName}(user)`) - script.onload = () => setStatus('ready') - script.onerror = () => { - setErrorMessage( - callbackHandlersRef.current.t('Failed to load Telegram Login Widget') - ) - setStatus('error') - } - container.appendChild(script) - - return () => { - delete windowCallbacks[callbackName] - container.replaceChildren() - } - }, [bindState, botName, callbackName, open]) - - const retry = () => { - setBindState('') - setRetryVersion((version) => version + 1) - } + const { widgetContainerRef, status, errorMessage, retry, handleOpenChange } = + useTelegramLoginWidget({ open, onOpenChange, botName, onSuccess }) return ( 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 d64a4a0e..f99ee0c7 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,13 +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, - useEffect, - useCallback, - useRef, - type ReactNode, -} 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' @@ -62,6 +56,7 @@ 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('') @@ -76,6 +71,7 @@ export function TwoFASetupDialog({ const handleSetup = useCallback(async () => { if (setupInFlightRef.current) return setupInFlightRef.current = true + setSetupAttempted(false) setInitializing(true) try { const response = await withVerification(setup2FA, { @@ -102,6 +98,7 @@ export function TwoFASetupDialog({ } finally { setupInFlightRef.current = false setInitializing(false) + setSetupAttempted(true) } }, [onOpenChange, t, withVerification]) @@ -130,6 +127,7 @@ export function TwoFASetupDialog({ setStep(0) setCode('') setSetupData(null) + setSetupAttempted(false) } else { toast.error(response.message || t('Failed to enable 2FA')) } @@ -146,6 +144,7 @@ export function TwoFASetupDialog({ setStep(0) setCode('') setSetupData(null) + setSetupAttempted(false) } onOpenChange(open) } @@ -159,7 +158,7 @@ export function TwoFASetupDialog({ }, [open, setupData, handleSetup]) const renderSetupState = (): ReactNode => { - if (initializing) { + if (initializing || !setupAttempted) { return (
diff --git a/web/default/src/features/profile/hooks/use-telegram-login-widget.ts b/web/default/src/features/profile/hooks/use-telegram-login-widget.ts new file mode 100644 index 00000000..23edf185 --- /dev/null +++ b/web/default/src/features/profile/hooks/use-telegram-login-widget.ts @@ -0,0 +1,217 @@ +/* +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, useEffect, useRef, useState, type RefObject } from 'react' +import { useTranslation } from 'react-i18next' +import { toast } from 'sonner' +import { useSecureVerificationGate } from '@/features/auth/secure-verification' +import { + bindTelegramAccount, + createTelegramBindState, + type TelegramAuthorizationPayload, +} from '../api' + +type TelegramWidgetStatus = 'idle' | 'loading' | 'ready' | 'binding' | 'error' + +interface UseTelegramLoginWidgetOptions { + open: boolean + onOpenChange: (open: boolean) => void + botName: string + onSuccess: () => void +} + +interface UseTelegramLoginWidgetResult { + widgetContainerRef: RefObject + status: TelegramWidgetStatus + errorMessage: string + retry: () => void + handleOpenChange: (open: boolean) => void +} + +export function useTelegramLoginWidget( + options: UseTelegramLoginWidgetOptions +): UseTelegramLoginWidgetResult { + const { t } = useTranslation() + const { withVerification } = useSecureVerificationGate() + const widgetContainerRef = useRef(null) + const [callbackName] = useState( + () => `telegramAuthCallback_${Math.random().toString(36).slice(2)}` + ) + const [bindState, setBindState] = useState('') + const [status, setStatus] = useState('idle') + const [errorMessage, setErrorMessage] = useState('') + const [retryVersion, setRetryVersion] = useState(0) + + const resetBindingState = useCallback((): void => { + setBindState('') + setStatus('idle') + setErrorMessage('') + }, []) + + const handleOpenChange = useCallback( + (nextOpen: boolean): void => { + if (!nextOpen) resetBindingState() + options.onOpenChange(nextOpen) + }, + [options.onOpenChange, resetBindingState] + ) + + const callbackHandlersRef = useRef({ + handleOpenChange, + onSuccess: options.onSuccess, + t, + withVerification, + }) + + useEffect(() => { + callbackHandlersRef.current = { + handleOpenChange, + onSuccess: options.onSuccess, + t, + withVerification, + } + }, [handleOpenChange, options.onSuccess, t, withVerification]) + + const initializeBinding = useCallback(async (): Promise => { + const handlers = callbackHandlersRef.current + setStatus('loading') + setErrorMessage('') + try { + const response = await handlers.withVerification( + createTelegramBindState, + { + scope: 'credentials', + title: handlers.t('Security verification'), + description: handlers.t( + 'Confirm your identity before binding a Telegram account.' + ), + } + ) + if (!response) { + setStatus('idle') + return + } + if (!response.success || !response.data?.state) { + throw new Error( + response.message || + handlers.t('Failed to initialize Telegram binding') + ) + } + setBindState(response.data.state) + } catch (error) { + setErrorMessage( + error instanceof Error + ? error.message + : handlers.t('Failed to initialize Telegram binding') + ) + setStatus('error') + } + }, []) + + useEffect(() => { + if (!options.open) return + const timer = window.setTimeout(() => void initializeBinding(), 0) + return () => window.clearTimeout(timer) + }, [initializeBinding, options.open, retryVersion]) + + useEffect(() => { + if (!options.open || !bindState || !widgetContainerRef.current) return + + const container = widgetContainerRef.current + const windowCallbacks = window as unknown as Record + container.replaceChildren() + + windowCallbacks[callbackName] = async ( + authorization: TelegramAuthorizationPayload + ): Promise => { + const handlers = callbackHandlersRef.current + setStatus('binding') + setErrorMessage('') + try { + const response = await handlers.withVerification( + () => bindTelegramAccount({ ...authorization, state: bindState }), + { + scope: 'credentials', + title: handlers.t('Security verification'), + description: handlers.t( + 'Confirm your identity before binding a Telegram account.' + ), + } + ) + if (!response) { + setStatus('ready') + return + } + if (!response.success) { + throw new Error( + response.message || handlers.t('Failed to bind Telegram account') + ) + } + toast.success(handlers.t('Telegram account bound successfully')) + callbackHandlersRef.current.onSuccess() + callbackHandlersRef.current.handleOpenChange(false) + } catch (error) { + setErrorMessage( + error instanceof Error + ? error.message + : handlers.t('Failed to bind Telegram account') + ) + setStatus('error') + } + } + + const script = document.createElement('script') + script.async = true + script.src = 'https://telegram.org/js/telegram-widget.js?22' + script.setAttribute( + 'data-telegram-login', + options.botName.replace(/^@/, '') + ) + script.setAttribute('data-size', 'large') + script.setAttribute('data-userpic', 'false') + script.setAttribute('data-onauth', `${callbackName}(user)`) + script.onload = () => setStatus('ready') + script.onerror = () => { + setErrorMessage( + callbackHandlersRef.current.t('Failed to load Telegram Login Widget') + ) + setStatus('error') + } + container.appendChild(script) + + return () => { + delete windowCallbacks[callbackName] + container.replaceChildren() + } + }, [bindState, callbackName, options.botName, options.open]) + + const retry = useCallback((): void => { + setBindState('') + setStatus('idle') + setErrorMessage('') + setRetryVersion((version) => version + 1) + }, []) + + return { + widgetContainerRef, + status, + errorMessage, + retry, + handleOpenChange, + } +} diff --git a/web/default/src/i18n/locales/ja.json b/web/default/src/i18n/locales/ja.json index 35e334c2..b4b6f256 100644 --- a/web/default/src/i18n/locales/ja.json +++ b/web/default/src/i18n/locales/ja.json @@ -5478,7 +5478,7 @@ "Failed to initialize Telegram binding": "Telegram連携の初期化に失敗しました", "Failed to load Telegram Login Widget": "Telegramログインウィジェットの読み込みに失敗しました", "Failed to sign out other sessions": "他のセッションからのログアウトに失敗しました", - "Keep this session and revoke all other signed-in devices": "現在のセッションを保持し、他のすべての端末のログインを取り消す", + "Keep this session and revoke all other signed-in devices": "現在のセッションを保持し、他のすべての端末からログアウトする", "No verification method is available. Add a password or sign in again with a linked account.": "利用可能な確認方法がありません。パスワードを追加するか、連携済みアカウントで再度ログインしてください。", "Other sessions signed out successfully": "他のセッションから正常にログアウトしました", "Preparing Telegram authorization...": "Telegram認証を準備しています...", diff --git a/web/default/src/lib/handle-server-error.test.ts b/web/default/src/lib/handle-server-error.test.ts index 1e83a78e..2801ffb4 100644 --- a/web/default/src/lib/handle-server-error.test.ts +++ b/web/default/src/lib/handle-server-error.test.ts @@ -74,6 +74,19 @@ describe('resolveServerErrorMessage', () => { ) }) + test('uses the localized fallback for Axios transport failures', () => { + const error = new AxiosError('Network Error', 'ERR_NETWORK') + + assert.equal( + resolveServerErrorMessage( + error, + 'Verification unavailable', + 'Content not found.' + ), + 'Verification unavailable' + ) + }) + test('keeps an Axios response title ahead of other messages', () => { const error = new AxiosError('request failed', 'ERR_BAD_REQUEST') error.response = { diff --git a/web/default/src/lib/handle-server-error.ts b/web/default/src/lib/handle-server-error.ts index b4f6f5de..6371b8b4 100644 --- a/web/default/src/lib/handle-server-error.ts +++ b/web/default/src/lib/handle-server-error.ts @@ -63,6 +63,8 @@ export function resolveServerErrorMessage( const responseMessage = getAxiosResponseText(error, 'message') if (responseMessage) return responseMessage + + return isContentNotFound ? contentNotFound : fallback } if ( diff --git a/web/default/tests/api-key-plaintext-lifecycle.test.tsx b/web/default/tests/api-key-plaintext-lifecycle.test.tsx index a6c7076d..71d93ef6 100644 --- a/web/default/tests/api-key-plaintext-lifecycle.test.tsx +++ b/web/default/tests/api-key-plaintext-lifecycle.test.tsx @@ -84,6 +84,11 @@ const { ApiKeyCell } = await import( ) const testEnv = createReactTestEnvironment() +function CopiedKeyProbe() { + const copiedKeyId = useApiKeys().copiedKeyId + return {copiedKeyId ?? 'none'} +} + function wrapper(props: { children: ReactNode }) { return {props.children} } @@ -155,6 +160,7 @@ describe('API key plaintext lifecycle', () => { allow_ips: '', }} /> + ) @@ -164,6 +170,7 @@ describe('API key plaintext lifecycle', () => { assert.equal(copyToClipboard.mock.calls.length, 1) assert.equal(toastError.mock.calls.length, 1) assert.equal(toastError.mock.calls[0]?.[0], 'Failed to copy to clipboard') + assert.equal(view.getByTestId('copied-key-id').textContent, 'none') }) }) }) diff --git a/web/default/tests/api-key-row-actions-chat-navigation.test.tsx b/web/default/tests/api-key-row-actions-chat-navigation.test.tsx new file mode 100644 index 00000000..80d12711 --- /dev/null +++ b/web/default/tests/api-key-row-actions-chat-navigation.test.tsx @@ -0,0 +1,209 @@ +/* +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 type { ReactElement, ReactNode } from 'react' +import { createReactTestEnvironment } from '@/test/react' +import { cleanup, fireEvent, render, waitFor } from '@testing-library/react' +import { afterAll, afterEach, beforeAll, beforeEach, mock, test } from 'bun:test' +import assert from 'node:assert/strict' + +let resolveKeyRequest: ((value: string | null) => void) | undefined +const resolveRealKey = mock( + () => + new Promise((resolve) => { + resolveKeyRequest = resolve + }) +) +const replaceLocation = mock((_url: string) => undefined) +const closePopup = mock(() => undefined) +const popupWindow = { + closed: false, + close: closePopup, + location: { replace: replaceLocation }, + opener: {} as Window, +} as unknown as Window +const openWindow = mock(() => popupWindow) + +function Container(props: { children?: ReactNode }): ReactElement { + return
{props.children}
+} + +mock.module('react-i18next', () => ({ + useTranslation: () => ({ t: (key: string) => key }), +})) + +mock.module('sonner', () => ({ + toast: { + error: mock(() => undefined), + info: mock(() => undefined), + success: mock(() => undefined), + }, +})) + +mock.module('../src/components/ui/button', () => ({ + Button: (props: { + children?: ReactNode + onClick?: () => void + type?: 'button' | 'submit' | 'reset' + }) => ( + + ), +})) + +mock.module('../src/components/ui/dropdown-menu', () => ({ + DropdownMenu: Container, + DropdownMenuContent: Container, + DropdownMenuItem: (props: { + children?: ReactNode + onClick?: () => void + }) => , + DropdownMenuSeparator: () => null, + DropdownMenuSub: Container, + DropdownMenuSubContent: Container, + DropdownMenuSubTrigger: Container, + DropdownMenuShortcut: Container, + DropdownMenuTrigger: (props: { + children?: ReactNode + render?: ReactElement + }) => ( + <> + {props.render} + {props.children} + + ), +})) + +mock.module('../src/components/ui/tooltip', () => ({ + Tooltip: Container, + TooltipContent: Container, + TooltipTrigger: (props: { + children?: ReactNode + render?: ReactElement + }) => ( + <> + {props.render} + {props.children} + + ), +})) + +mock.module('../src/features/chat/hooks/use-chat-presets', () => ({ + useChatPresets: () => ({ + serverAddress: 'https://api.example.test', + chatPresets: [ + { + id: 'browser-chat', + name: 'Browser chat', + type: 'web', + url: 'https://chat.example.test/?key={key}', + }, + ], + }), +})) + +mock.module('../src/features/chat/lib/chat-links', () => ({ + resolveChatUrl: () => 'https://chat.example.test/?key=resolved-key', +})) + +mock.module('../src/features/chat/lib/send-to-fluent', () => ({ + sendToFluent: () => false, +})) + +mock.module('../src/features/keys/api', () => ({ + updateApiKeyStatus: mock(async () => ({ success: true })), +})) + +mock.module('../src/features/keys/components/api-keys-provider', () => ({ + useApiKeys: () => ({ + setOpen: mock(() => undefined), + setCurrentRow: mock(() => undefined), + triggerRefresh: mock(() => undefined), + setResolvedKey: mock(() => undefined), + resolveRealKey, + }), +})) + +const { DataTableRowActions } = await import( + '../src/features/keys/components/data-table-row-actions' +) +const testEnv = createReactTestEnvironment() +let originalWindowOpen: typeof window.open + +beforeAll(async () => { + await testEnv.setup() + originalWindowOpen = window.open + window.open = openWindow as typeof window.open +}) + +beforeEach(() => { + resolveKeyRequest = undefined + resolveRealKey.mockClear() + openWindow.mockClear() + replaceLocation.mockClear() + closePopup.mockClear() +}) + +afterEach(() => cleanup()) + +afterAll(() => { + window.open = originalWindowOpen + testEnv.teardown() +}) + +test('opens a placeholder before resolving the API key', async () => { + const view = render( + + ) + + fireEvent.click(view.getByRole('button', { name: /Browser chat/ })) + + assert.equal(openWindow.mock.calls.length, 1) + assert.equal(replaceLocation.mock.calls.length, 0) + + resolveKeyRequest?.('resolved-key') + await waitFor(() => assert.equal(replaceLocation.mock.calls.length, 1)) + assert.deepEqual(replaceLocation.mock.calls[0], [ + 'https://chat.example.test/?key=resolved-key', + ]) + assert.equal(popupWindow.opener, null) +}) diff --git a/web/default/tests/overview-request-copy-error-handling.test.tsx b/web/default/tests/overview-request-copy-error-handling.test.tsx index 13eb9d28..70ae6f77 100644 --- a/web/default/tests/overview-request-copy-error-handling.test.tsx +++ b/web/default/tests/overview-request-copy-error-handling.test.tsx @@ -30,9 +30,11 @@ import { import assert from 'node:assert/strict' let verificationError = new Error('request failed') -const withVerification = mock(async () => { - throw verificationError -}) +const withVerification = mock( + async (_apiCall: () => Promise, _config?: { scope?: string }) => { + throw verificationError + } +) const handleServerError = mock( (_error: unknown, _options?: { fallback?: string }) => undefined ) @@ -105,6 +107,7 @@ test('reports an unreported request-copy rejection and consumes it', async () => await waitFor(() => assert.equal(handleServerError.mock.calls.length, 1)) await waitForCopyToSettle(view) + assert.equal(withVerification.mock.calls[0]?.[1]?.scope, 'api_token') assert.equal(handleServerError.mock.calls[0]?.[0], verificationError) assert.deepEqual(handleServerError.mock.calls[0]?.[1], { fallback: 'Failed to copy to clipboard', diff --git a/web/default/tests/passkey-registration-error-handling.test.tsx b/web/default/tests/passkey-registration-error-handling.test.tsx index 2868d1d3..4e14d2dc 100644 --- a/web/default/tests/passkey-registration-error-handling.test.tsx +++ b/web/default/tests/passkey-registration-error-handling.test.tsx @@ -181,6 +181,8 @@ test('handles a rejected Passkey verification continuation', async () => { await waitFor(() => assert.equal(withVerification.mock.calls.length, 1)) await waitFor(() => assert.equal(consoleError.mock.calls.length, 1)) assert.match(String(consoleError.mock.calls[0]?.[1]), /method lookup failed/) + assert.equal(toastError.mock.calls.length, 0) + assert.equal(handleServerError.mock.calls.length, 0) }) test('stops Passkey removal when verification method discovery fails', async () => { diff --git a/web/default/tests/secure-verification-continuation.test.tsx b/web/default/tests/secure-verification-continuation.test.tsx index 8dc31313..95333187 100644 --- a/web/default/tests/secure-verification-continuation.test.tsx +++ b/web/default/tests/secure-verification-continuation.test.tsx @@ -374,7 +374,7 @@ describe('useSecureVerification', () => { assert.equal(await firstContinuation, null) await act(async () => { firstDiscovery.resolve(availableVerificationMethods) - await Promise.resolve() + await flushMacrotask() }) assert.equal(result.current.open, false) diff --git a/web/default/tests/two-fa-setup-shared-verification.test.tsx b/web/default/tests/two-fa-setup-shared-verification.test.tsx index c9db0d48..bf177c39 100644 --- a/web/default/tests/two-fa-setup-shared-verification.test.tsx +++ b/web/default/tests/two-fa-setup-shared-verification.test.tsx @@ -20,6 +20,7 @@ import { createReactTestEnvironment } from '@/test/react' import { act, cleanup, render, waitFor } from '@testing-library/react' import { afterAll, afterEach, beforeAll, beforeEach, mock, test } from 'bun:test' import assert from 'node:assert/strict' +import type { ReactNode } from 'react' const setup2FA = mock(async () => ({ success: true, @@ -46,6 +47,44 @@ mock.module('qrcode.react', () => ({ QRCodeSVG: () =>
, })) +function TestContainer(props: { children?: ReactNode }) { + return
{props.children}
+} + +mock.module('../src/components/ui/alert', () => ({ + Alert: TestContainer, + AlertDescription: TestContainer, +})) + +mock.module('../src/components/ui/button', () => ({ + Button: (props: { children?: ReactNode; onClick?: () => void }) => ( + + ), +})) + +mock.module('../src/components/ui/dialog', () => ({ + Dialog: TestContainer, + DialogContent: TestContainer, + DialogDescription: TestContainer, + DialogFooter: TestContainer, + DialogHeader: TestContainer, + DialogTitle: TestContainer, +})) + +mock.module('../src/components/ui/input', () => ({ + Input: (props: { value?: string }) => , +})) + +mock.module('../src/components/ui/label', () => ({ + Label: TestContainer, +})) + +mock.module('../src/components/copy-button', () => ({ + CopyButton: TestContainer, +})) + mock.module('../src/lib/api', () => ({ setup2FA, enable2FA: mock(async () => ({ success: true })), @@ -69,7 +108,7 @@ afterEach(() => cleanup()) afterAll(() => testEnv.teardown()) test('uses the authenticated layout verification gate for 2FA setup', async () => { - render( + const view = render( undefined)} @@ -77,6 +116,7 @@ test('uses the authenticated layout verification gate for 2FA setup', async () = /> ) + assert.ok(view.getByText('Setting up 2FA...')) await waitFor(() => assert.equal(setup2FA.mock.calls.length, 1)) assert.equal(withVerification.mock.calls.length, 1) }) From c03ec6501135d5060c56c5f06c89934965949fad Mon Sep 17 00:00:00 2001 From: CSCITech Date: Sat, 29 Aug 2026 04:12:41 +0800 Subject: [PATCH 17/70] fix bugs --- controller/relay_preconsume_test.go | 64 ----------- controller/relay_task_billing_test.go | 55 ++++++++++ middleware/auth_session_generation_test.go | 118 --------------------- middleware/secure_verification_test.go | 106 ++++++++++++++++++ 4 files changed, 161 insertions(+), 182 deletions(-) delete mode 100644 controller/relay_preconsume_test.go delete mode 100644 middleware/auth_session_generation_test.go diff --git a/controller/relay_preconsume_test.go b/controller/relay_preconsume_test.go deleted file mode 100644 index 84c2b6f3..00000000 --- a/controller/relay_preconsume_test.go +++ /dev/null @@ -1,64 +0,0 @@ -package controller - -import ( - "testing" - - "github.com/MAX-API-Next/MAX-API/common" - 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/stretchr/testify/require" -) - -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) - }) - } -} diff --git a/controller/relay_task_billing_test.go b/controller/relay_task_billing_test.go index 53d4353d..3b7c817e 100644 --- a/controller/relay_task_billing_test.go +++ b/controller/relay_task_billing_test.go @@ -6,11 +6,16 @@ import ( "net/http/httptest" "testing" + "github.com/MAX-API-Next/MAX-API/common" "github.com/MAX-API-Next/MAX-API/constant" "github.com/MAX-API-Next/MAX-API/dto" "github.com/MAX-API-Next/MAX-API/model" "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/types" "github.com/gin-gonic/gin" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" @@ -71,3 +76,53 @@ func TestMidjourneyRelayErrorStatusCodeReportsSettlementPendingAsConflict(t *tes 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) + }) + } +} diff --git a/middleware/auth_session_generation_test.go b/middleware/auth_session_generation_test.go deleted file mode 100644 index d960f08a..00000000 --- a/middleware/auth_session_generation_test.go +++ /dev/null @@ -1,118 +0,0 @@ -package middleware - -import ( - "fmt" - "net/http" - "net/http/httptest" - "strings" - "testing" - - "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 TestUserAuthRejectsStaleSessionGeneration(t *testing.T) { - gin.SetMode(gin.TestMode) - oldDB := model.DB - oldRedisEnabled := common.RedisEnabled - common.RedisEnabled = false - - dsn := fmt.Sprintf("file:%s?mode=memory&cache=shared", strings.ReplaceAll(t.Name(), "/", "_")) - db, err := gorm.Open(sqlite.Open(dsn), &gorm.Config{}) - require.NoError(t, err) - model.DB = db - require.NoError(t, db.AutoMigrate(&model.User{})) - if !db.Migrator().HasColumn(&model.User{}, "session_generation") { - require.NoError(t, db.Exec("ALTER TABLE users ADD COLUMN session_generation INTEGER NOT NULL DEFAULT 0").Error) - } - 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 - 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", 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/middleware/secure_verification_test.go b/middleware/secure_verification_test.go index 699218df..e35e2e1b 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) { @@ -179,3 +185,103 @@ func TestSecureVerificationRejectsLoginMethodMarker(t *testing.T) { 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 + common.RedisEnabled = false + + dsn := fmt.Sprintf("file:%s?mode=memory&cache=shared", strings.ReplaceAll(t.Name(), "/", "_")) + db, err := gorm.Open(sqlite.Open(dsn), &gorm.Config{}) + require.NoError(t, err) + model.DB = db + require.NoError(t, db.AutoMigrate(&model.User{})) + if !db.Migrator().HasColumn(&model.User{}, "session_generation") { + require.NoError(t, db.Exec("ALTER TABLE users ADD COLUMN session_generation INTEGER NOT NULL DEFAULT 0").Error) + } + 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 + 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", 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) +} From f8d2947f98700c7e6c2cfc4122d597dc0ef2f547 Mon Sep 17 00:00:00 2001 From: CSCITech Date: Sat, 29 Aug 2026 04:49:25 +0800 Subject: [PATCH 18/70] fix remaining review findings --- model/billing_settlement.go | 4 +- model/log.go | 3 +- model/log_test.go | 51 +++++++++++++++++++ service/billing.go | 1 + service/billing_session_test.go | 2 + service/task_billing_test.go | 42 +++++++++++++-- .../keys/components/api-keys-provider.tsx | 3 ++ .../api-key-plaintext-lifecycle.test.tsx | 41 ++++++++++++++- 8 files changed, 140 insertions(+), 7 deletions(-) diff --git a/model/billing_settlement.go b/model/billing_settlement.go index c40f0bc5..8371352c 100644 --- a/model/billing_settlement.go +++ b/model/billing_settlement.go @@ -111,6 +111,7 @@ type BillingSettlementEffect struct { 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"` @@ -780,7 +781,8 @@ 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, diff --git a/model/log.go b/model/log.go index bd3101aa..8db4d923 100644 --- a/model/log.go +++ b/model/log.go @@ -662,6 +662,7 @@ type RecordTaskBillingLogParams struct { ModelName string Quota int TokenId int + TokenName string Group string Other map[string]interface{} NodeName string // 任务发起节点；为空时回退当前节点 @@ -695,7 +696,7 @@ 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 { tokenName = token.Name diff --git a/model/log_test.go b/model/log_test.go index 5e04d6c0..e6dd1154 100644 --- a/model/log_test.go +++ b/model/log_test.go @@ -990,6 +990,8 @@ func TestBillingSettlementEffectReplayExportsTokenUsage(t *testing.T) { LogType: LogTypeConsume, Content: "settlement effect", ModelName: "settlement-model", + TokenID: 7002, + TokenName: "request-time-deleted-token", UpdateUsage: true, Quota: 42, PromptTokens: 3, @@ -1015,6 +1017,55 @@ func TestBillingSettlementEffectReplayExportsTokenUsage(t *testing.T) { 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 TestBillingSettlementEffectReplayPrefersCurrentTokenName(t *testing.T) { + 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("1 = 1").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) + }) + + const ( + userID = 7011 + tokenID = 7012 + ) + 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 TestWaitPendingLogQuotaDataDrainsEnqueuedWork(t *testing.T) { diff --git a/service/billing.go b/service/billing.go index b6f08fba..e651ef8f 100644 --- a/service/billing.go +++ b/service/billing.go @@ -49,6 +49,7 @@ func newConsumeBillingSettlementEffect( ChannelID: params.ChannelId, ModelName: params.ModelName, TokenID: params.TokenId, + TokenName: params.TokenName, Group: params.Group, Other: params.Other, NodeName: common.NodeName, diff --git a/service/billing_session_test.go b/service/billing_session_test.go index 057642f1..f16a9fb0 100644 --- a/service/billing_session_test.go +++ b/service/billing_session_test.go @@ -123,6 +123,7 @@ func TestSettleAndRecordConsumeCarriesZeroUsageLogInDurableEffect(t *testing.T) Quota: 0, Content: "upstream omitted billable usage", TokenId: 603, + TokenName: "request-time-token-name", UseTimeSeconds: 7, IsStream: true, Group: "default", @@ -140,6 +141,7 @@ func TestSettleAndRecordConsumeCarriesZeroUsageLogInDurableEffect(t *testing.T) 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) diff --git a/service/task_billing_test.go b/service/task_billing_test.go index c24a659b..5295cd2a 100644 --- a/service/task_billing_test.go +++ b/service/task_billing_test.go @@ -139,13 +139,15 @@ func TestSweepTimedOutUnconfirmedSubmitRequiresReviewWithoutRefund(t *testing.T) func TestSweepTimedOutTaskWaitsForSubmissionSettlement(t *testing.T) { truncate(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, - Quota: 20, SubmitTime: time.Now().Add(-2 * time.Minute).Unix(), + 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", @@ -153,10 +155,21 @@ func TestSweepTimedOutTaskWaitsForSubmissionSettlement(t *testing.T) { }, } require.NoError(t, model.DB.Create(task).Error) + submissionInput := model.BillingSettlementInput{ + OperationKey: "request:timeout-waits-for-settlement:finalize", + Source: model.BillingSettlementSourceWallet, + UserID: userID, + FundingDelta: 10, + TaskID: task.ID, + TaskQuota: 20, + TaskQuotaTarget: 30, + } require.NoError(t, model.DB.Create(&model.BillingSettlement{ - OperationKey: "request:timeout-waits-for-settlement:finalize", - Source: model.BillingSettlementSourceWallet, UserID: 711, - FundingDelta: 10, Status: model.BillingSettlementStatusPending, + 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) @@ -166,6 +179,27 @@ func TestSweepTimedOutTaskWaitsForSubmissionSettlement(t *testing.T) { 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) } // --------------------------------------------------------------------------- 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 681e93ff..c5f18510 100644 --- a/web/default/src/features/keys/components/api-keys-provider.tsx +++ b/web/default/src/features/keys/components/api-keys-provider.tsx @@ -171,6 +171,9 @@ export function ApiKeysProvider({ children }: { children: React.ReactNode }) { } toast.error(res.message || t(ERROR_MESSAGES.UNEXPECTED)) return {} + } catch { + toast.error(t(ERROR_MESSAGES.UNEXPECTED)) + return {} } finally { for (const id of ids) { setLoadingKeys((prev) => { diff --git a/web/default/tests/api-key-plaintext-lifecycle.test.tsx b/web/default/tests/api-key-plaintext-lifecycle.test.tsx index 71d93ef6..0ba93e0e 100644 --- a/web/default/tests/api-key-plaintext-lifecycle.test.tsx +++ b/web/default/tests/api-key-plaintext-lifecycle.test.tsx @@ -43,6 +43,7 @@ let resolveTokenKeyRequest: data: { key: string } }) => void) | undefined +let rejectTokenKeysBatchRequest: ((reason?: unknown) => void) | undefined const fetchTokenKey = mock( () => @@ -50,6 +51,12 @@ const fetchTokenKey = mock( resolveTokenKeyRequest = resolve }) ) +const fetchTokenKeysBatch = mock( + async (): Promise<{ + success: boolean + data: { keys: Record } + }> => ({ success: true, data: { keys: {} } }) +) const copyToClipboard = mock(async () => false) const toastError = mock(() => undefined) @@ -73,7 +80,7 @@ mock.module('../src/lib/copy-to-clipboard', () => ({ mock.module('../src/features/keys/api', () => ({ fetchTokenKey, - fetchTokenKeysBatch: mock(async () => ({ success: true, data: { keys: {} } })), + fetchTokenKeysBatch, })) const { ApiKeysProvider, useApiKeys } = await import( @@ -97,9 +104,11 @@ beforeAll(() => testEnv.setup()) beforeEach(() => { fetchTokenKey.mockClear() + fetchTokenKeysBatch.mockClear() copyToClipboard.mockClear() toastError.mockClear() resolveTokenKeyRequest = undefined + rejectTokenKeysBatchRequest = undefined }) afterEach(async () => { @@ -110,6 +119,36 @@ afterEach(async () => { afterAll(() => testEnv.teardown()) describe('API key plaintext lifecycle', () => { + test('reports batch reveal failures and clears every loading state', async () => { + fetchTokenKeysBatch.mockImplementationOnce( + () => + new Promise((_, reject) => { + rejectTokenKeysBatchRequest = reject + }) + ) + const { result } = renderHook(() => useApiKeys(), { wrapper }) + + let request!: Promise> + act(() => { + request = result.current.resolveRealKeysBatch([201, 202]) + }) + + await waitFor(() => { + assert.equal(result.current.loadingKeys[201], true) + assert.equal(result.current.loadingKeys[202], true) + }) + + await act(async () => { + rejectTokenKeysBatchRequest?.(new Error('batch request failed')) + assert.deepEqual(await request, {}) + }) + + assert.equal(toastError.mock.calls.length, 1) + assert.equal(toastError.mock.calls[0]?.[0], 'An unexpected error occurred') + assert.equal(result.current.loadingKeys[201], undefined) + assert.equal(result.current.loadingKeys[202], undefined) + }) + test('does not cache a key that resolves after its reveal popover closes', async () => { const { result } = renderHook(() => useApiKeys(), { wrapper }) From 2c84712a82af001ea192e188a37190097649f418 Mon Sep 17 00:00:00 2001 From: CSCITech Date: Sat, 29 Aug 2026 05:36:52 +0800 Subject: [PATCH 19/70] fix remaining security and billing review findings --- model/billing_settlement.go | 12 ++- model/billing_settlement_regression_test.go | 5 ++ model/log_test.go | 45 +++++++++++ model/security_credentials_test.go | 8 +- model/twofa.go | 53 +++++++----- router/api_router_test.go | 1 + service/billing_session.go | 4 +- service/task_billing.go | 2 +- service/violation_fee.go | 18 +---- .../layout/components/chat-presets-item.tsx | 14 +--- .../hooks/use-secure-verification.ts | 20 +++++ .../chat/hooks/use-active-chat-key.ts | 14 +--- .../overview/overview-dashboard.tsx | 14 ++-- .../keys/components/api-keys-provider.tsx | 18 +---- .../components/data-table-row-actions.tsx | 12 ++- .../active-chat-key-verification.test.tsx | 8 +- .../api-key-plaintext-lifecycle.test.tsx | 7 +- ...i-key-row-actions-chat-navigation.test.tsx | 81 +++++++++++++------ ...rview-request-copy-error-handling.test.tsx | 22 +++-- 19 files changed, 236 insertions(+), 122 deletions(-) diff --git a/model/billing_settlement.go b/model/billing_settlement.go index 8371352c..f6851f98 100644 --- a/model/billing_settlement.go +++ b/model/billing_settlement.go @@ -24,6 +24,9 @@ const ( BillingSettlementEffectApplied = "applied" ) +const billingRequestOperationPrefix = "request:" +const billingRequestFinalizeSuffix = ":finalize" + var ( ErrBillingSettlementManualReview = errors.New("billing settlement requires manual review") ErrBillingSettlementTaskConflict = errors.New("billing settlement task quota conflict") @@ -230,7 +233,7 @@ func HasUnresolvedPositiveFinalizeSettlement(userID int) (bool, error) { Where("user_id = ?", userID). Where("funding_delta > 0"). Where("status IN ?", []string{BillingSettlementStatusPending, BillingSettlementStatusManual}). - Where("operation_key LIKE ?", "request:%:finalize"). + Where("operation_key LIKE ?", BillingRequestFinalizeOperationKey("%")). Limit(1). Find(&record) if result.Error != nil { @@ -239,6 +242,10 @@ func HasUnresolvedPositiveFinalizeSettlement(userID int) (bool, 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. @@ -766,6 +773,9 @@ func ProcessBillingSettlementEffect(operationKey string) error { if useActualQuota { logQuota = effect.Quota } + 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 { diff --git a/model/billing_settlement_regression_test.go b/model/billing_settlement_regression_test.go index 71943f10..bea6c222 100644 --- a/model/billing_settlement_regression_test.go +++ b/model/billing_settlement_regression_test.go @@ -92,6 +92,11 @@ func TestHasUnresolvedPositiveFinalizeSettlement(t *testing.T) { assert.False(t, blocked) } +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) diff --git a/model/log_test.go b/model/log_test.go index e6dd1154..052a707e 100644 --- a/model/log_test.go +++ b/model/log_test.go @@ -1068,6 +1068,51 @@ func TestBillingSettlementEffectReplayPrefersCurrentTokenName(t *testing.T) { require.Equal(t, "current-token-name", log.TokenName) } +func TestBillingSettlementEffectRejectsNegativeUsageProjection(t *testing.T) { + 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("1 = 1").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) + }) + + const userID = 7021 + 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) { resetQuotaDataCacheForTest(t) diff --git a/model/security_credentials_test.go b/model/security_credentials_test.go index 8ced71a5..9361766a 100644 --- a/model/security_credentials_test.go +++ b/model/security_credentials_test.go @@ -130,6 +130,12 @@ func TestPasskeyAndTwoFASecurityChangesBumpSessionGeneration(t *testing.T) { 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) @@ -173,7 +179,7 @@ func TestPasskeyUsageUpdateCannotRestoreReplacedCredential(t *testing.T) { now := time.Now() staleUsageUpdate.SignCount = 2 staleUsageUpdate.LastUsedAt = &now - require.Error(t, UpdatePasskeyCredentialAfterAuthentication(&staleUsageUpdate)) + require.ErrorIs(t, UpdatePasskeyCredentialAfterAuthentication(&staleUsageUpdate), ErrPasskeyCredentialChanged) stored, err := GetPasskeyByUserID(user.Id) require.NoError(t, err) diff --git a/model/twofa.go b/model/twofa.go index 59ece320..bd45e16c 100644 --- a/model/twofa.go +++ b/model/twofa.go @@ -179,6 +179,11 @@ func replaceBackupCodeHashesTx(tx *gorm.DB, userID int, hashes []string) error { 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 验证并使用备用码 func ValidateBackupCode(userId int, code string) (bool, error) { if !common.ValidateBackupCode(code) { @@ -187,29 +192,33 @@ 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 - } + validated := false + err := DB.Transaction(func(tx *gorm.DB) error { + if err := lockTwoFACredentialUserTx(tx, userId); err != nil { + return err + } - // 验证备用码 - for _, bc := range backupCodes { - if common.ValidatePasswordAndHash(normalizedCode, bc.CodeHash) { - // 标记为已使用 + var backupCodes []TwoFABackupCode + if err := tx.Where("user_id = ? AND is_used = ?", userId, false).Find(&backupCodes).Error; err != nil { + return err + } + for _, bc := range backupCodes { + if !common.ValidatePasswordAndHash(normalizedCode, bc.CodeHash) { + continue + } now := time.Now() - bc.IsUsed = true - bc.UsedAt = &now - - if err := DB.Save(&bc).Error; err != nil { - return false, err + result := tx.Model(&TwoFABackupCode{}). + Where("id = ? AND user_id = ? AND is_used = ?", bc.Id, userId, false). + Updates(map[string]interface{}{"is_used": true, "used_at": &now}) + if result.Error != nil { + return result.Error } - - return true, nil + validated = result.RowsAffected == 1 + return nil } - } - - return false, nil + return nil + }) + return validated, err } // GetUnusedBackupCodeCount 获取未使用的备用码数量 @@ -237,6 +246,9 @@ 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) { @@ -311,6 +323,9 @@ func ReplaceBackupCodesAndBumpSessionGeneration(userID int, codes []string) (int 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). diff --git a/router/api_router_test.go b/router/api_router_test.go index 9bd6f648..0673eaf5 100644 --- a/router/api_router_test.go +++ b/router/api_router_test.go @@ -481,6 +481,7 @@ func TestSensitiveCredentialRoutesRequireStepUp(t *testing.T) { {name: "passkey registration begin", path: "/api/user/passkey/register/begin", body: `{}`}, {name: "passkey registration finish", path: "/api/user/passkey/register/finish", body: `{}`}, {name: "2fa setup", path: "/api/user/2fa/setup", body: `{}`}, + {name: "2fa enable", path: "/api/user/2fa/enable", body: `{"code":"000000"}`}, {name: "session revocation", path: "/api/user/sessions/revoke", body: `{}`}, {name: "telegram bind state", path: "/api/oauth/telegram/bind/state", body: `{}`}, {name: "telegram bind", path: "/api/oauth/telegram/bind", body: `{}`}, diff --git a/service/billing_session.go b/service/billing_session.go index bc970f3f..44a4587b 100644 --- a/service/billing_session.go +++ b/service/billing_session.go @@ -309,7 +309,7 @@ func (s *BillingSession) prepareDurableSettlementLocked(actualQuota int, effect tokenDelta = 0 } input := &model.BillingSettlementInput{ - OperationKey: "request:" + s.relayInfo.RequestId + ":finalize", + OperationKey: model.BillingRequestFinalizeOperationKey(s.relayInfo.RequestId), Source: source, UserID: userID, SubscriptionID: subscriptionID, @@ -384,7 +384,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, diff --git a/service/task_billing.go b/service/task_billing.go index b6859eb1..93c279f4 100644 --- a/service/task_billing.go +++ b/service/task_billing.go @@ -157,7 +157,7 @@ func taskFinalSettlementPending(task *model.Task) (bool, error) { if task == nil || task.PrivateData.BillingRequestId == "" { return false, nil } - status, found, err := model.GetBillingSettlementStatus("request:" + task.PrivateData.BillingRequestId + ":finalize") + status, found, err := model.GetBillingSettlementStatus(model.BillingRequestFinalizeOperationKey(task.PrivateData.BillingRequestId)) if err != nil || !found { return false, err } diff --git a/service/violation_fee.go b/service/violation_fee.go index 045ea2ed..51395150 100644 --- a/service/violation_fee.go +++ b/service/violation_fee.go @@ -156,7 +156,7 @@ func chargeViolationFeeIfNeeded(ctx *gin.Context, relayInfo *relaycommon.RelayIn "violation_fee_marker": CSAMViolationMarker, } requestID, upstreamRequestID := billingEffectRequestIDs(ctx, relayInfo) - effect := newConsumeBillingSettlementEffect(model.RecordConsumeLogParams{ + logParams := model.RecordConsumeLogParams{ ChannelId: relayInfo.ChannelId, ModelName: relayInfo.OriginModelName, TokenName: tokenName, @@ -167,7 +167,8 @@ func chargeViolationFeeIfNeeded(ctx *gin.Context, relayInfo *relaycommon.RelayIn IsStream: relayInfo.IsStream, Group: relayInfo.UsingGroup, Other: other, - }, requestID, upstreamRequestID, true) + } + effect := newConsumeBillingSettlementEffect(logParams, requestID, upstreamRequestID, true) if settler, ok := relayInfo.Billing.(interface { SettleWithEffect(int, *model.BillingSettlementEffect) error @@ -186,18 +187,7 @@ func chargeViolationFeeIfNeeded(ctx *gin.Context, relayInfo *relaycommon.RelayIn model.UpdateUserUsedQuotaAndRequestCount(relayInfo.UserId, feeQuota) model.UpdateChannelUsedQuota(relayInfo.ChannelId, feeQuota) - model.RecordConsumeLog(ctx, relayInfo.UserId, model.RecordConsumeLogParams{ - ChannelId: relayInfo.ChannelId, - ModelName: relayInfo.OriginModelName, - TokenName: tokenName, - Quota: feeQuota, - Content: "Violation fee charged", - TokenId: relayInfo.TokenId, - UseTimeSeconds: int(useTimeSeconds), - IsStream: relayInfo.IsStream, - Group: relayInfo.UsingGroup, - Other: other, - }) + model.RecordConsumeLog(ctx, relayInfo.UserId, logParams) return violationFeeChargeResult{applicable: true, settled: true} } 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 84b7880f..315030e9 100644 --- a/web/default/src/components/layout/components/chat-presets-item.tsx +++ b/web/default/src/components/layout/components/chat-presets-item.tsx @@ -40,7 +40,7 @@ import { SidebarMenuSubItem, useSidebar, } from '@/components/ui/sidebar' -import { useSecureVerificationGate } from '@/features/auth/secure-verification' +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 { @@ -157,7 +157,7 @@ export function ChatPresetsItem({ item }: { item: NavChatPresets }) { const { t } = useTranslation() const { chatPresets, serverAddress } = useChatPresets() const { state, isMobile, setOpenMobile } = useSidebar() - const { withVerification } = useSecureVerificationGate() + const withApiTokenVerification = useApiTokenVerification() const href = useLocation({ select: (location) => location.href }) const [loadingPresetId, setLoadingPresetId] = useState(null) const loadingPresetIdRef = useRef(null) @@ -184,13 +184,7 @@ export function ChatPresetsItem({ item }: { item: NavChatPresets }) { setLoadingPresetId(preset.id) try { activeKey = - (await withVerification(fetchActiveChatKey, { - scope: 'api_token', - title: t('Security verification'), - description: t( - 'Confirm your identity before sending an API key to a chat client.' - ), - })) || undefined + (await withApiTokenVerification(fetchActiveChatKey)) || undefined if (!activeKey) return } catch (error) { const message = @@ -223,7 +217,7 @@ export function ChatPresetsItem({ item }: { item: NavChatPresets }) { window.open(url, '_blank', 'noopener,noreferrer') setOpenMobile(false) }, - [serverAddress, setOpenMobile, t, withVerification] + [serverAddress, setOpenMobile, t, withApiTokenVerification] ) const normalizedHref = normalizeHref(href) 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 e464682d..da71514e 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 @@ -18,6 +18,7 @@ For commercial licensing, please contact https://github.com/MAX-API-Next/MAX-API */ import { useCallback, useEffect, useMemo, useRef, useState } from 'react' import i18next from 'i18next' +import { useTranslation } from 'react-i18next' import { toast } from 'sonner' import { handleServerError } from '@/lib/handle-server-error' import { @@ -26,6 +27,7 @@ import { markSecureVerificationErrorReported, } from '@/lib/secure-verification' import { checkVerificationMethods, verify } from '../api' +import { useSecureVerificationGate } from '../components/secure-verification-context' import { selectVerificationMethod } from '../method-selection' import type { SecureVerificationState, @@ -66,6 +68,24 @@ const initialState: InternalState = { attemptId: null, } +const defaultApiTokenVerificationDescription = + 'Confirm your identity before sending an API key to a chat client.' + +export function useApiTokenVerification(description?: string) { + const { t } = useTranslation() + const { withVerification } = useSecureVerificationGate() + + return useCallback( + (apiCall: () => Promise): Promise => + withVerification(apiCall, { + scope: 'api_token', + title: t('Security verification'), + description: description ?? t(defaultApiTokenVerificationDescription), + }), + [description, t, withVerification] + ) +} + export function useSecureVerification( options: UseSecureVerificationOptions = {} ) { 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 bcd9abaf..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 @@ -20,7 +20,7 @@ import { useCallback } from 'react' import { useQuery } from '@tanstack/react-query' import i18next from 'i18next' import { useAuthStore } from '@/stores/auth-store' -import { useSecureVerificationGate } from '@/features/auth/secure-verification' +import { useApiTokenVerification } from '@/features/auth/secure-verification' import { fetchTokenKey, getApiKeys } from '@/features/keys/api' import { API_KEY_STATUS } from '@/features/keys/constants' @@ -49,20 +49,14 @@ export async function fetchActiveChatKey() { */ export function useActiveChatKey(enabled: boolean) { const userId = useAuthStore((state) => state.auth.user?.id) - const { withVerification } = useSecureVerificationGate() + const withApiTokenVerification = useApiTokenVerification() const queryFn = useCallback(async (): Promise => { - const key = await withVerification(fetchActiveChatKey, { - scope: 'api_token', - title: i18next.t('Security verification'), - description: i18next.t( - 'Confirm your identity before sending an API key to a chat client.' - ), - }) + const key = await withApiTokenVerification(fetchActiveChatKey) if (!key) { throw new Error(i18next.t('API key verification was cancelled')) } return key - }, [withVerification]) + }, [withApiTokenVerification]) return useQuery({ queryKey: ['chat-active-key', userId], 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 3e363151..4f7212d7 100644 --- a/web/default/src/features/dashboard/components/overview/overview-dashboard.tsx +++ b/web/default/src/features/dashboard/components/overview/overview-dashboard.tsx @@ -56,7 +56,7 @@ import { CardStaggerContainer, CardStaggerItem, } from '@/components/page-transition' -import { useSecureVerificationGate } from '@/features/auth/secure-verification' +import { useApiTokenVerification } from '@/features/auth/secure-verification' import { fetchTokenKey, getApiKeys } from '@/features/keys/api' import type { ApiKey } from '@/features/keys/types' import { @@ -384,7 +384,9 @@ export function RequestPreview(props: { }) { const { t } = useTranslation() const shouldReduceMotion = useReducedMotion() - const { withVerification } = useSecureVerificationGate() + 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({ @@ -399,13 +401,7 @@ export function RequestPreview(props: { setIsCopying(true) try { - const result = await withVerification(() => fetchTokenKey(keyId), { - scope: 'api_token', - title: t('Security verification'), - description: t( - 'Confirm your identity before copying a request with your API key.' - ), - }) + const result = await withApiTokenVerification(() => fetchTokenKey(keyId)) if (!result) return const key = result.success && result.data?.key ? result.data.key : '' if (!key) { 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 c5f18510..d293b0b0 100644 --- a/web/default/src/features/keys/components/api-keys-provider.tsx +++ b/web/default/src/features/keys/components/api-keys-provider.tsx @@ -20,7 +20,7 @@ import React, { useState, useCallback, useRef, useEffect } from 'react' import { useTranslation } from 'react-i18next' import { toast } from 'sonner' import useDialogState from '@/hooks/use-dialog' -import { useSecureVerificationGate } from '@/features/auth/secure-verification' +import { useApiTokenVerification } from '@/features/auth/secure-verification' import { fetchTokenKey, fetchTokenKeysBatch } from '../api' import { ERROR_MESSAGES } from '../constants' import { type ApiKey, type ApiKeysDialogType } from '../types' @@ -51,7 +51,9 @@ const ApiKeysContext = React.createContext(null) export function ApiKeysProvider({ children }: { children: React.ReactNode }) { const { t } = useTranslation() - const { withVerification } = useSecureVerificationGate() + 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) @@ -80,18 +82,6 @@ export function ApiKeysProvider({ children }: { children: React.ReactNode }) { setRefreshTrigger((prev) => prev + 1) }, []) - const withApiTokenVerification = useCallback( - (apiCall: () => Promise) => - withVerification(apiCall, { - scope: 'api_token', - title: t('Security verification'), - description: t( - 'Confirm your identity before creating or revealing API keys.' - ), - }), - [t, withVerification] - ) - const resolveRealKey = useCallback( async ( id: number, 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 50871359..d00dac1b 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 @@ -230,7 +230,11 @@ export function DataTableRowActions({ 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')} @@ -247,7 +251,11 @@ export function DataTableRowActions({ 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/tests/active-chat-key-verification.test.tsx b/web/default/tests/active-chat-key-verification.test.tsx index 7629da8e..3a97c45e 100644 --- a/web/default/tests/active-chat-key-verification.test.tsx +++ b/web/default/tests/active-chat-key-verification.test.tsx @@ -23,7 +23,7 @@ import { QueryClient, QueryClientProvider } from '@tanstack/react-query' import { afterAll, afterEach, beforeAll, beforeEach, mock, test } from 'bun:test' import assert from 'node:assert/strict' -const withVerification = mock(async () => null) +const withApiTokenVerification = mock(async () => null) mock.module('../src/stores/auth-store', () => ({ useAuthStore: (selector: (state: unknown) => unknown) => @@ -31,7 +31,7 @@ mock.module('../src/stores/auth-store', () => ({ })) mock.module('../src/features/auth/secure-verification', () => ({ - useSecureVerificationGate: () => ({ withVerification }), + useApiTokenVerification: () => withApiTokenVerification, })) mock.module('../src/features/keys/api', () => ({ @@ -49,7 +49,7 @@ const { useActiveChatKey } = await import( const testEnv = createReactTestEnvironment() beforeAll(() => testEnv.setup()) -beforeEach(() => withVerification.mockClear()) +beforeEach(() => withApiTokenVerification.mockClear()) afterEach(() => cleanup()) afterAll(() => testEnv.teardown()) @@ -66,6 +66,6 @@ test('does not retry a cancelled chat API-key verification', async () => { const { result } = renderHook(() => useActiveChatKey(true), { wrapper }) await waitFor(() => assert.equal(result.current.isError, true)) - assert.equal(withVerification.mock.calls.length, 1) + assert.equal(withApiTokenVerification.mock.calls.length, 1) queryClient.clear() }) diff --git a/web/default/tests/api-key-plaintext-lifecycle.test.tsx b/web/default/tests/api-key-plaintext-lifecycle.test.tsx index 0ba93e0e..5c0f6dba 100644 --- a/web/default/tests/api-key-plaintext-lifecycle.test.tsx +++ b/web/default/tests/api-key-plaintext-lifecycle.test.tsx @@ -61,9 +61,10 @@ const copyToClipboard = mock(async () => false) const toastError = mock(() => undefined) mock.module('../src/features/auth/secure-verification', () => ({ - useSecureVerificationGate: () => ({ - withVerification: (apiCall: () => Promise) => apiCall(), - }), + useApiTokenVerification: + () => + (apiCall: () => Promise) => + apiCall(), })) mock.module('react-i18next', () => ({ diff --git a/web/default/tests/api-key-row-actions-chat-navigation.test.tsx b/web/default/tests/api-key-row-actions-chat-navigation.test.tsx index 80d12711..79e91129 100644 --- a/web/default/tests/api-key-row-actions-chat-navigation.test.tsx +++ b/web/default/tests/api-key-row-actions-chat-navigation.test.tsx @@ -38,6 +38,8 @@ const popupWindow = { opener: {} as Window, } as unknown as Window const openWindow = mock(() => popupWindow) +const copyToClipboard = mock(async () => false) +const toastError = mock(() => undefined) function Container(props: { children?: ReactNode }): ReactElement { return
{props.children}
@@ -49,12 +51,14 @@ mock.module('react-i18next', () => ({ mock.module('sonner', () => ({ toast: { - error: mock(() => undefined), + error: toastError, info: mock(() => undefined), success: mock(() => undefined), }, })) +mock.module('../src/lib/copy-to-clipboard', () => ({ copyToClipboard })) + mock.module('../src/components/ui/button', () => ({ Button: (props: { children?: ReactNode @@ -146,29 +150,8 @@ const { DataTableRowActions } = await import( const testEnv = createReactTestEnvironment() let originalWindowOpen: typeof window.open -beforeAll(async () => { - await testEnv.setup() - originalWindowOpen = window.open - window.open = openWindow as typeof window.open -}) - -beforeEach(() => { - resolveKeyRequest = undefined - resolveRealKey.mockClear() - openWindow.mockClear() - replaceLocation.mockClear() - closePopup.mockClear() -}) - -afterEach(() => cleanup()) - -afterAll(() => { - window.open = originalWindowOpen - testEnv.teardown() -}) - -test('opens a placeholder before resolving the API key', async () => { - const view = render( +function renderRowActions(): ReturnType { + return render( { } /> ) +} + +beforeAll(async () => { + await testEnv.setup() + originalWindowOpen = window.open + window.open = openWindow as typeof window.open +}) + +beforeEach(() => { + resolveKeyRequest = undefined + resolveRealKey.mockClear() + openWindow.mockClear() + replaceLocation.mockClear() + closePopup.mockClear() + copyToClipboard.mockClear() + toastError.mockClear() + popupWindow.opener = {} as Window +}) + +afterEach(() => cleanup()) + +afterAll(() => { + window.open = originalWindowOpen + testEnv.teardown() +}) + +test('opens a placeholder before resolving the API key', async () => { + const view = renderRowActions() fireEvent.click(view.getByRole('button', { name: /Browser chat/ })) @@ -207,3 +218,25 @@ test('opens a placeholder before resolving the API key', async () => { ]) assert.equal(popupWindow.opener, null) }) + +test('reports a clipboard failure when copying an API key', async () => { + resolveRealKey.mockImplementationOnce(async () => 'resolved-key') + const view = renderRowActions() + + fireEvent.click(view.getByRole('button', { name: /Copy Key/ })) + + await waitFor(() => assert.equal(copyToClipboard.mock.calls.length, 1)) + assert.deepEqual(toastError.mock.calls[0], ['Failed to copy to clipboard']) +}) + +test('reports a clipboard failure when copying connection info', async () => { + resolveRealKey.mockImplementationOnce(async () => 'resolved-key') + const view = renderRowActions() + + fireEvent.click( + view.getByRole('button', { name: /Copy Connection Info/ }) + ) + + await waitFor(() => assert.equal(copyToClipboard.mock.calls.length, 1)) + assert.deepEqual(toastError.mock.calls[0], ['Failed to copy to clipboard']) +}) diff --git a/web/default/tests/overview-request-copy-error-handling.test.tsx b/web/default/tests/overview-request-copy-error-handling.test.tsx index 70ae6f77..f620de52 100644 --- a/web/default/tests/overview-request-copy-error-handling.test.tsx +++ b/web/default/tests/overview-request-copy-error-handling.test.tsx @@ -30,10 +30,11 @@ import { import assert from 'node:assert/strict' let verificationError = new Error('request failed') -const withVerification = mock( - async (_apiCall: () => Promise, _config?: { scope?: string }) => { - throw verificationError - } +const withApiTokenVerification = mock(async () => { + throw verificationError +}) +const useApiTokenVerification = mock( + (_description?: string) => withApiTokenVerification ) const handleServerError = mock( (_error: unknown, _options?: { fallback?: string }) => undefined @@ -51,7 +52,7 @@ mock.module('sonner', () => ({ })) mock.module('@/features/auth/secure-verification', () => ({ - useSecureVerificationGate: () => ({ withVerification }), + useApiTokenVerification, })) mock.module('@/lib/handle-server-error', () => ({ handleServerError })) @@ -64,7 +65,8 @@ const testEnv = createReactTestEnvironment() beforeAll(() => testEnv.setup()) beforeEach(() => { verificationError = new Error('request failed') - withVerification.mockClear() + withApiTokenVerification.mockClear() + useApiTokenVerification.mockClear() handleServerError.mockClear() }) afterEach(() => cleanup()) @@ -107,7 +109,9 @@ test('reports an unreported request-copy rejection and consumes it', async () => await waitFor(() => assert.equal(handleServerError.mock.calls.length, 1)) await waitForCopyToSettle(view) - assert.equal(withVerification.mock.calls[0]?.[1]?.scope, 'api_token') + assert.equal(useApiTokenVerification.mock.calls[0]?.[0], + 'Confirm your identity before copying a request with your API key.' + ) assert.equal(handleServerError.mock.calls[0]?.[0], verificationError) assert.deepEqual(handleServerError.mock.calls[0]?.[1], { fallback: 'Failed to copy to clipboard', @@ -120,7 +124,9 @@ test('does not duplicate a rejection already reported by verification', async () fireEvent.click(view.getByRole('button', { name: 'Copy ready-to-run curl' })) - await waitFor(() => assert.equal(withVerification.mock.calls.length, 1)) + await waitFor(() => + assert.equal(withApiTokenVerification.mock.calls.length, 1) + ) await waitForCopyToSettle(view) assert.equal(handleServerError.mock.calls.length, 0) }) From 5e56f0232d4742a0000b877557c3b3d1ff7efcfa Mon Sep 17 00:00:00 2001 From: CSCITech Date: Sat, 29 Aug 2026 06:02:22 +0800 Subject: [PATCH 20/70] polish secure verification review follow-ups --- docs/openapi/api.json | 8 ++++---- .../hooks/use-secure-verification.ts | 12 ++++++------ .../features/keys/components/api-keys-provider.tsx | 10 ++++++++-- .../components/dialogs/telegram-bind-dialog.tsx | 7 +++++-- .../profile/components/profile-security-card.tsx | 4 ++-- .../tests/secure-verification-continuation.test.tsx | 12 ++++++++++-- 6 files changed, 35 insertions(+), 18 deletions(-) diff --git a/docs/openapi/api.json b/docs/openapi/api.json index fee3c17a..510ff87c 100644 --- a/docs/openapi/api.json +++ b/docs/openapi/api.json @@ -1053,20 +1053,20 @@ "/api/oauth/telegram/bind/state": { "post": { "operationId": "createTelegramBindState", - "summary": "Create Telegram binding state", + "summary": "创建 Telegram 绑定 state", "deprecated": false, - "description": "Requires an authenticated user and credentials-scoped security verification. Returns a short-lived, user-bound, one-time state value.", + "description": "需要登录并完成 credentials 作用域安全验证。返回短期有效、与用户绑定的一次性 state。", "tags": [ "OAuth" ], "parameters": [], "responses": { "200": { - "description": "State created", + "description": "state 创建成功", "headers": {} }, "403": { - "description": "Credentials verification is required or expired", + "description": "需要 credentials 作用域安全验证或验证已过期", "headers": {} } }, 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 da71514e..ab49c736 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 @@ -71,7 +71,9 @@ const initialState: InternalState = { const defaultApiTokenVerificationDescription = 'Confirm your identity before sending an API key to a chat client.' -export function useApiTokenVerification(description?: string) { +export function useApiTokenVerification( + description?: string +): (apiCall: () => Promise) => Promise { const { t } = useTranslation() const { withVerification } = useSecureVerificationGate() @@ -261,11 +263,9 @@ export function useSecureVerification( } catch (error) { if (!isVerificationAttemptActive(attemptId)) return setState((prev) => ({ ...prev, loading: false })) - const message = - error instanceof Error - ? error.message - : i18next.t('Verification failed') - toast.error(message) + handleServerError(error, { + fallback: i18next.t('Verification failed'), + }) markSecureVerificationErrorReported(error) onError?.(error) throw error 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 d293b0b0..b30b9556 100644 --- a/web/default/src/features/keys/components/api-keys-provider.tsx +++ b/web/default/src/features/keys/components/api-keys-provider.tsx @@ -49,7 +49,13 @@ type ApiKeysContextType = { 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.') @@ -198,7 +204,7 @@ export function ApiKeysProvider({ children }: { children: React.ReactNode }) { withApiTokenVerification, }} > - {children} + {props.children} ) } 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 65ebf49f..3539f162 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 @@ -64,7 +64,7 @@ export function TelegramBindDialog({
- +