fix: apply affiliate rebates on topups

Move rebate processing into model transactions, keep frozen rebates out of transferable quota until release, and trigger top-up rebates idempotently.

model/payment_method_guard_test.go

@@ -265,6 +265,63 @@ func TestRechargeWxpay_RejectsMismatchedPaymentProvider(t *testing.T) {
 	assert.Equal(t, 0, getUserQuotaForPaymentGuardTest(t, 183))
 }
 
+func TestRechargeWaffo_ProcessesAffiliateRebateIdempotently(t *testing.T) {
+	truncateTables(t)
+
+	common.OptionMapRWMutex.RLock()
+	originalOptionMap := common.OptionMap
+	common.OptionMapRWMutex.RUnlock()
+	common.OptionMapRWMutex.Lock()
+	common.OptionMap = map[string]string{
+		AffRebateEnabledKey:     "true",
+		AffRebateRatePercentKey: "10",
+		AffRebateFrozenDaysKey:  "0",
+	}
+	common.OptionMapRWMutex.Unlock()
+	t.Cleanup(func() {
+		common.OptionMapRWMutex.Lock()
+		common.OptionMap = originalOptionMap
+		common.OptionMapRWMutex.Unlock()
+	})
+
+	inviter := &User{Id: 184, Username: "rebate_inviter", AffCode: "rebate_inviter_code", Status: common.UserStatusEnabled}
+	invitee := &User{Id: 185, Username: "rebate_invitee", AffCode: "rebate_invitee_code", Status: common.UserStatusEnabled, InviterId: inviter.Id}
+	require.NoError(t, DB.Create(inviter).Error)
+	require.NoError(t, DB.Create(invitee).Error)
+	topUp := &TopUp{
+		UserId:          invitee.Id,
+		Amount:          2,
+		Money:           2,
+		TradeNo:         "waffo-rebate-guard",
+		PaymentMethod:   PaymentMethodWaffo,
+		PaymentProvider: PaymentProviderWaffo,
+		Status:          common.TopUpStatusPending,
+		CreateTime:      time.Now().Unix(),
+	}
+	require.NoError(t, topUp.Insert())
+
+	require.NoError(t, RechargeWaffo("waffo-rebate-guard", "127.0.0.1"))
+	require.NoError(t, RechargeWaffo("waffo-rebate-guard", "127.0.0.1"))
+
+	baseQuota := int(decimal.NewFromInt(2).Mul(decimal.NewFromFloat(common.QuotaPerUnit)).IntPart())
+	var record RebateRecord
+	require.NoError(t, DB.First(&record).Error)
+	assert.Equal(t, topUp.Id, record.OrderId)
+	assert.Equal(t, RebateOrderTypeTopUp, record.OrderType)
+	assert.Equal(t, baseQuota, record.OrderAmount)
+	assert.Equal(t, baseQuota/10, record.RebateAmount)
+	assert.Equal(t, "released", record.Status)
+
+	var reloadedInviter User
+	require.NoError(t, DB.First(&reloadedInviter, inviter.Id).Error)
+	assert.Equal(t, baseQuota/10, reloadedInviter.AffQuota)
+	assert.Equal(t, baseQuota/10, reloadedInviter.AffHistoryQuota)
+
+	var count int64
+	require.NoError(t, DB.Model(&RebateRecord{}).Count(&count).Error)
+	assert.Equal(t, int64(1), count)
+}
+
 func TestCompleteSubscriptionOrder_RejectsMismatchedPaymentProvider(t *testing.T) {
 	truncateTables(t)
 

model/rebate_process.go

@@ -0,0 +1,178 @@
+package model
+
+import (
+	"errors"
+	"strconv"
+	"time"
+
+	"github.com/QuantumNous/new-api/common"
+	"gorm.io/gorm"
+	"gorm.io/gorm/clause"
+)
+
+const RebateOrderTypeTopUp = "topup"
+
+func IsRebateEnabled() bool {
+	enabled, err := strconv.ParseBool(readRebateOption(AffRebateEnabledKey))
+	return err == nil && enabled
+}
+
+func GetGlobalRebateRatePercent() int {
+	return parseNonNegativeRebateInt(readRebateOption(AffRebateRatePercentKey))
+}
+
+func GetGlobalRebateFrozenDays() int {
+	return parseNonNegativeRebateInt(readRebateOption(AffRebateFrozenDaysKey))
+}
+
+func GetEffectiveRebateRate(inviter *User) int {
+	if inviter != nil && inviter.AffRebateRatePercent > 0 {
+		return inviter.AffRebateRatePercent
+	}
+	return GetGlobalRebateRatePercent()
+}
+
+func GetEffectiveFrozenDays(inviter *User) int {
+	if inviter != nil && inviter.AffRebateFrozenDays > 0 {
+		return inviter.AffRebateFrozenDays
+	}
+	return GetGlobalRebateFrozenDays()
+}
+
+func ProcessRebateAfterRecharge(invitee *User, orderId int, orderType string, orderAmount int) error {
+	return DB.Transaction(func(tx *gorm.DB) error {
+		return processRebateAfterRechargeTx(tx, invitee, orderId, orderType, orderAmount)
+	})
+}
+
+func processTopUpRebateTx(tx *gorm.DB, topUp *TopUp, quotaToAdd int) error {
+	if topUp == nil || quotaToAdd <= 0 {
+		return nil
+	}
+	var invitee User
+	if err := tx.First(&invitee, "id = ?", topUp.UserId).Error; err != nil {
+		return err
+	}
+	return processRebateAfterRechargeTx(tx, &invitee, topUp.Id, RebateOrderTypeTopUp, quotaToAdd)
+}
+
+func processRebateAfterRechargeTx(tx *gorm.DB, invitee *User, orderId int, orderType string, orderAmount int) error {
+	if invitee == nil || invitee.InviterId <= 0 || orderAmount <= 0 {
+		return nil
+	}
+	if !IsRebateEnabled() {
+		return nil
+	}
+	if orderType == "" {
+		return errors.New("order type is empty")
+	}
+	if orderId <= 0 {
+		return errors.New("order id must be positive")
+	}
+
+	var inviter User
+	if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).
+		First(&inviter, "id = ?", invitee.InviterId).Error; err != nil {
+		return err
+	}
+
+	var existing RebateRecord
+	err := tx.Where(
+		"order_type = ? AND order_id = ? AND inviter_id = ? AND invitee_id = ?",
+		orderType, orderId, inviter.Id, invitee.Id,
+	).First(&existing).Error
+	if err == nil {
+		return nil
+	}
+	if !errors.Is(err, gorm.ErrRecordNotFound) {
+		return err
+	}
+
+	rate := GetEffectiveRebateRate(&inviter)
+	if rate <= 0 {
+		return nil
+	}
+	rebateAmount := orderAmount * rate / 100
+	if rebateAmount <= 0 {
+		return nil
+	}
+
+	now := time.Now()
+	frozenDays := GetEffectiveFrozenDays(&inviter)
+	record := RebateRecord{
+		InviterId:    inviter.Id,
+		InviteeId:    invitee.Id,
+		OrderId:      orderId,
+		OrderType:    orderType,
+		OrderAmount:  orderAmount,
+		RebateAmount: rebateAmount,
+		RatePercent:  rate,
+		Status:       "released",
+		ReleasedAt:   &now,
+	}
+	if frozenDays > 0 {
+		frozenUntil := now.Add(time.Duration(frozenDays) * 24 * time.Hour)
+		record.Status = "frozen"
+		record.FrozenUntil = &frozenUntil
+		record.ReleasedAt = nil
+	}
+
+	if err := tx.Create(&record).Error; err != nil {
+		return err
+	}
+
+	updates := map[string]interface{}{
+		"aff_history": gorm.Expr("aff_history + ?", rebateAmount),
+	}
+	if frozenDays <= 0 {
+		updates["aff_quota"] = gorm.Expr("aff_quota + ?", rebateAmount)
+	}
+	return tx.Model(&User{}).Where("id = ?", inviter.Id).Updates(updates).Error
+}
+
+func ReleaseExpiredRebates() error {
+	return DB.Transaction(func(tx *gorm.DB) error {
+		var records []RebateRecord
+		if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).
+			Where("status = ? AND frozen_until IS NOT NULL AND frozen_until <= ?", "frozen", time.Now()).
+			Order("id ASC").
+			Find(&records).Error; err != nil {
+			return err
+		}
+
+		releasedAt := time.Now()
+		for _, record := range records {
+			result := tx.Model(&RebateRecord{}).
+				Where("id = ? AND status = ?", record.Id, "frozen").
+				Updates(map[string]interface{}{
+					"status":      "released",
+					"released_at": releasedAt,
+				})
+			if result.Error != nil {
+				return result.Error
+			}
+			if result.RowsAffected == 0 {
+				continue
+			}
+			if err := tx.Model(&User{}).Where("id = ?", record.InviterId).
+				Update("aff_quota", gorm.Expr("aff_quota + ?", record.RebateAmount)).Error; err != nil {
+				return err
+			}
+		}
+		return nil
+	})
+}
+
+func readRebateOption(key string) string {
+	common.OptionMapRWMutex.RLock()
+	defer common.OptionMapRWMutex.RUnlock()
+	return common.OptionMap[key]
+}
+
+func parseNonNegativeRebateInt(value string) int {
+	parsed, err := strconv.Atoi(value)
+	if err != nil || parsed < 0 {
+		return 0
+	}
+	return parsed
+}

model/task_cas_test.go

@@ -45,6 +45,7 @@ func TestMain(m *testing.M) {
 		&UserSubscription{},
 		&PromoCode{},
 		&PromoCodeUsage{},
+		&RebateRecord{},
 	); err != nil {
 		panic("failed to migrate: " + err.Error())
 	}
@@ -66,6 +67,7 @@ func truncateTables(t *testing.T) {
 		DB.Exec("DELETE FROM user_subscriptions")
 		DB.Exec("DELETE FROM promo_code_usages")
 		DB.Exec("DELETE FROM promo_codes")
+		DB.Exec("DELETE FROM rebate_records")
 	})
 }
 

model/topup.go

@@ -286,8 +286,11 @@ func RechargeEpay(tradeNo string, actualPaymentMethod string, callerIp string) (
 		}
 
 		bonus, bonusErr := applyPromoCodeBonusTx(tx, topUp)
-		promoBonus = bonus
+		if bonusErr != nil {
 			return bonusErr
+		}
+		promoBonus = bonus
+		return processTopUpRebateTx(tx, topUp, quotaToAdd)
 	})
 
 	if err != nil {
@@ -345,8 +348,11 @@ func Recharge(referenceId string, customerId string, callerIp string) (err error
 		}
 
 		bonus, bonusErr := applyPromoCodeBonusTx(tx, topUp)
-		promoBonus = bonus
+		if bonusErr != nil {
 			return bonusErr
+		}
+		promoBonus = bonus
+		return processTopUpRebateTx(tx, topUp, int(quota))
 	})
 
 	if err != nil {
@@ -584,6 +590,9 @@ func ManualCompleteTopUp(tradeNo string, callerIp string) error {
 			return bonusErr
 		}
 		promoBonus = bonus
+		if err := processTopUpRebateTx(tx, topUp, quotaToAdd); err != nil {
+			return err
+		}
 
 		userId = topUp.UserId
 		payMoney = topUp.Money
@@ -670,8 +679,11 @@ func RechargeCreem(referenceId string, customerEmail string, customerName string
 		}
 
 		bonus, bonusErr := applyPromoCodeBonusTx(tx, topUp)
-		promoBonus = bonus
+		if bonusErr != nil {
 			return bonusErr
+		}
+		promoBonus = bonus
+		return processTopUpRebateTx(tx, topUp, int(quota))
 	})
 
 	if err != nil {
@@ -737,8 +749,11 @@ func RechargeWaffo(tradeNo string, callerIp string) (err error) {
 		}
 
 		bonus, bonusErr := applyPromoCodeBonusTx(tx, topUp)
-		promoBonus = bonus
+		if bonusErr != nil {
 			return bonusErr
+		}
+		promoBonus = bonus
+		return processTopUpRebateTx(tx, topUp, quotaToAdd)
 	})
 
 	if err != nil {
@@ -804,8 +819,11 @@ func RechargeAlipay(tradeNo string, callerIp string) (err error) {
 		}
 
 		bonus, bonusErr := applyPromoCodeBonusTx(tx, topUp)
-		promoBonus = bonus
+		if bonusErr != nil {
 			return bonusErr
+		}
+		promoBonus = bonus
+		return processTopUpRebateTx(tx, topUp, quotaToAdd)
 	})
 
 	if err != nil {
@@ -871,8 +889,11 @@ func RechargeWxpay(tradeNo string, callerIp string) (err error) {
 		}
 
 		bonus, bonusErr := applyPromoCodeBonusTx(tx, topUp)
-		promoBonus = bonus
+		if bonusErr != nil {
 			return bonusErr
+		}
+		promoBonus = bonus
+		return processTopUpRebateTx(tx, topUp, quotaToAdd)
 	})
 
 	if err != nil {
@@ -936,8 +957,11 @@ func RechargeWaffoPancake(tradeNo string) (err error) {
 		}
 
 		bonus, bonusErr := applyPromoCodeBonusTx(tx, topUp)
-		promoBonus = bonus
+		if bonusErr != nil {
 			return bonusErr
+		}
+		promoBonus = bonus
+		return processTopUpRebateTx(tx, topUp, quotaToAdd)
 	})
 
 	if err != nil {

service/rebate.go

@@ -1,151 +1,31 @@
 package service
 
-import (
+import "github.com/QuantumNous/new-api/model"
-	"errors"
-	"strconv"
-	"time"
-
-	"github.com/QuantumNous/new-api/common"
-	"github.com/QuantumNous/new-api/model"
-	"gorm.io/gorm"
-	"gorm.io/gorm/clause"
-)
 
 func IsRebateEnabled() bool {
-	enabled, err := strconv.ParseBool(readRebateOption(model.AffRebateEnabledKey))
+	return model.IsRebateEnabled()
-	return err == nil && enabled
 }
 
 func GetGlobalRebateRatePercent() int {
-	return parseNonNegativeRebateInt(readRebateOption(model.AffRebateRatePercentKey))
+	return model.GetGlobalRebateRatePercent()
 }
 
 func GetGlobalRebateFrozenDays() int {
-	return parseNonNegativeRebateInt(readRebateOption(model.AffRebateFrozenDaysKey))
+	return model.GetGlobalRebateFrozenDays()
 }
 
 func GetEffectiveRebateRate(inviter *model.User) int {
-	if inviter != nil && inviter.AffRebateRatePercent > 0 {
+	return model.GetEffectiveRebateRate(inviter)
-		return inviter.AffRebateRatePercent
-	}
-	return GetGlobalRebateRatePercent()
 }
 
 func GetEffectiveFrozenDays(inviter *model.User) int {
-	if inviter != nil && inviter.AffRebateFrozenDays > 0 {
+	return model.GetEffectiveFrozenDays(inviter)
-		return inviter.AffRebateFrozenDays
-	}
-	return GetGlobalRebateFrozenDays()
 }
 
 func ProcessRebateAfterRecharge(invitee *model.User, orderId int, orderType string, orderAmount int) error {
-	if invitee == nil || invitee.InviterId <= 0 || orderAmount <= 0 {
+	return model.ProcessRebateAfterRecharge(invitee, orderId, orderType, orderAmount)
-		return nil
-	}
-	if !IsRebateEnabled() {
-		return nil
-	}
-	if orderType == "" {
-		return errors.New("order type is empty")
-	}
-	if orderId <= 0 {
-		return errors.New("order id must be positive")
-	}
-
-	return model.DB.Transaction(func(tx *gorm.DB) error {
-		var inviter model.User
-		if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).
-			First(&inviter, "id = ?", invitee.InviterId).Error; err != nil {
-			return err
-		}
-
-		var existing model.RebateRecord
-		err := tx.Where(
-			"order_type = ? AND order_id = ? AND inviter_id = ? AND invitee_id = ?",
-			orderType, orderId, inviter.Id, invitee.Id,
-		).First(&existing).Error
-		if err == nil {
-			return nil
-		}
-		if !errors.Is(err, gorm.ErrRecordNotFound) {
-			return err
-		}
-
-		rate := GetEffectiveRebateRate(&inviter)
-		if rate <= 0 {
-			return nil
-		}
-		rebateAmount := orderAmount * rate / 100
-		if rebateAmount <= 0 {
-			return nil
-		}
-
-		now := time.Now()
-		record := model.RebateRecord{
-			InviterId:    inviter.Id,
-			InviteeId:    invitee.Id,
-			OrderId:      orderId,
-			OrderType:    orderType,
-			OrderAmount:  orderAmount,
-			RebateAmount: rebateAmount,
-			RatePercent:  rate,
-			Status:       "released",
-			ReleasedAt:   &now,
-		}
-		if frozenDays := GetEffectiveFrozenDays(&inviter); frozenDays > 0 {
-			frozenUntil := now.Add(time.Duration(frozenDays) * 24 * time.Hour)
-			record.Status = "frozen"
-			record.FrozenUntil = &frozenUntil
-			record.ReleasedAt = nil
-		}
-
-		if err := tx.Create(&record).Error; err != nil {
-			return err
-		}
-
-		return tx.Model(&model.User{}).Where("id = ?", inviter.Id).Updates(map[string]interface{}{
-			"aff_quota":   gorm.Expr("aff_quota + ?", rebateAmount),
-			"aff_history": gorm.Expr("aff_history + ?", rebateAmount),
-		}).Error
-	})
 }
 
 func ReleaseExpiredRebates() error {
-	return model.DB.Transaction(func(tx *gorm.DB) error {
+	return model.ReleaseExpiredRebates()
-		var records []model.RebateRecord
-		if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).
-			Where("status = ? AND frozen_until IS NOT NULL AND frozen_until <= ?", "frozen", time.Now()).
-			Order("id ASC").
-			Find(&records).Error; err != nil {
-			return err
-		}
-
-		releasedAt := time.Now()
-		for _, record := range records {
-			result := tx.Model(&model.RebateRecord{}).
-				Where("id = ? AND status = ?", record.Id, "frozen").
-				Updates(map[string]interface{}{
-					"status":      "released",
-					"released_at": releasedAt,
-				})
-			if result.Error != nil {
-				return result.Error
-			}
-		}
-		return nil
-	})
-}
-
-func readRebateOption(key string) string {
-	common.OptionMapRWMutex.RLock()
-	defer common.OptionMapRWMutex.RUnlock()
-	return common.OptionMap[key]
-}
-
-func parseNonNegativeRebateInt(value string) int {
-	parsed, err := strconv.Atoi(value)
-	if err != nil || parsed < 0 {
-		return 0
-	}
-	return parsed
 }

service/rebate_test.go

@@ -144,7 +144,7 @@ func TestProcessRebateAfterRechargeUsesGlobalRate(t *testing.T) {
 	require.Nil(t, record.ReleasedAt)
 
 	inviter := loadRebateUser(t, db, 1)
-	require.Equal(t, 100, inviter.AffQuota)
+	require.Equal(t, 0, inviter.AffQuota)
 	require.Equal(t, 100, inviter.AffHistoryQuota)
 }
 
@@ -167,7 +167,7 @@ func TestProcessRebateAfterRechargeUsesPerUserOverrides(t *testing.T) {
 	require.WithinDuration(t, time.Now().Add(24*time.Hour), *record.FrozenUntil, time.Minute)
 
 	inviter := loadRebateUser(t, db, 1)
-	require.Equal(t, 250, inviter.AffQuota)
+	require.Equal(t, 0, inviter.AffQuota)
 	require.Equal(t, 250, inviter.AffHistoryQuota)
 }
 
@@ -193,7 +193,6 @@ func TestReleaseExpiredRebatesMarksOnlyDueFrozenRecords(t *testing.T) {
 	past := time.Now().Add(-time.Hour)
 	future := time.Now().Add(time.Hour)
 	require.NoError(t, db.Model(inviter).Updates(map[string]interface{}{
-		"aff_quota":   50,
 		"aff_history": 50,
 	}).Error)
 	require.NoError(t, db.Create(&model.RebateRecord{