03fc89da00
This commit introduces a new function, MaskEmail, to mask user email addresses in logs, preventing PII leakage. Additionally, the RelayInfo logging has been updated to utilize this new masking function, ensuring sensitive information is properly handled. The channel test logic has also been improved to dynamically determine the relay format based on the request path.
79 lines
3.2 KiB
Go
79 lines
3.2 KiB
Go
package service
|
|
|
|
import (
|
|
"errors"
|
|
"fmt"
|
|
"net/http"
|
|
"one-api/common"
|
|
"one-api/logger"
|
|
"one-api/model"
|
|
relaycommon "one-api/relay/common"
|
|
"one-api/types"
|
|
|
|
"github.com/bytedance/gopkg/util/gopool"
|
|
"github.com/gin-gonic/gin"
|
|
)
|
|
|
|
func ReturnPreConsumedQuota(c *gin.Context, relayInfo *relaycommon.RelayInfo, preConsumedQuota int) {
|
|
if preConsumedQuota != 0 {
|
|
gopool.Go(func() {
|
|
relayInfoCopy := *relayInfo
|
|
|
|
err := PostConsumeQuota(&relayInfoCopy, -preConsumedQuota, 0, false)
|
|
if err != nil {
|
|
common.SysLog("error return pre-consumed quota: " + err.Error())
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
// PreConsumeQuota checks if the user has enough quota to pre-consume.
|
|
// It returns the pre-consumed quota if successful, or an error if not.
|
|
func PreConsumeQuota(c *gin.Context, preConsumedQuota int, relayInfo *relaycommon.RelayInfo) (int, *types.NewAPIError) {
|
|
userQuota, err := model.GetUserQuota(relayInfo.UserId, false)
|
|
if err != nil {
|
|
return 0, types.NewError(err, types.ErrorCodeQueryDataError, types.ErrOptionWithSkipRetry())
|
|
}
|
|
if userQuota <= 0 {
|
|
return 0, types.NewErrorWithStatusCode(errors.New("user quota is not enough"), types.ErrorCodeInsufficientUserQuota, http.StatusForbidden, types.ErrOptionWithSkipRetry(), types.ErrOptionWithNoRecordErrorLog())
|
|
}
|
|
if userQuota-preConsumedQuota < 0 {
|
|
return 0, types.NewErrorWithStatusCode(fmt.Errorf("预扣费额度失败, 用户剩余额度: %s, 需要预扣费额度: %s", logger.FormatQuota(userQuota), logger.FormatQuota(preConsumedQuota)), types.ErrorCodeInsufficientUserQuota, http.StatusForbidden, types.ErrOptionWithSkipRetry(), types.ErrOptionWithNoRecordErrorLog())
|
|
}
|
|
|
|
trustQuota := common.GetTrustQuota()
|
|
|
|
relayInfo.UserQuota = userQuota
|
|
if userQuota > trustQuota {
|
|
// 用户额度充足，判断令牌额度是否充足
|
|
if !relayInfo.TokenUnlimited {
|
|
// 非无限令牌，判断令牌额度是否充足
|
|
tokenQuota := c.GetInt("token_quota")
|
|
if tokenQuota > trustQuota {
|
|
// 令牌额度充足，信任令牌
|
|
preConsumedQuota = 0
|
|
logger.LogInfo(c, fmt.Sprintf("user %d quota %s and token %d quota %d are enough, trusted and no need to pre-consume", relayInfo.UserId, logger.FormatQuota(userQuota), relayInfo.TokenId, tokenQuota))
|
|
}
|
|
} else {
|
|
// in this case, we do not pre-consume quota
|
|
// because the user has enough quota
|
|
preConsumedQuota = 0
|
|
logger.LogInfo(c, fmt.Sprintf("user %d with unlimited token has enough quota %s, trusted and no need to pre-consume", relayInfo.UserId, logger.FormatQuota(userQuota)))
|
|
}
|
|
}
|
|
|
|
if preConsumedQuota > 0 {
|
|
err := PreConsumeTokenQuota(relayInfo, preConsumedQuota)
|
|
if err != nil {
|
|
return 0, types.NewErrorWithStatusCode(err, types.ErrorCodePreConsumeTokenQuotaFailed, http.StatusForbidden, types.ErrOptionWithSkipRetry(), types.ErrOptionWithNoRecordErrorLog())
|
|
}
|
|
err = model.DecreaseUserQuota(relayInfo.UserId, preConsumedQuota)
|
|
if err != nil {
|
|
return 0, types.NewError(err, types.ErrorCodeUpdateDataError, types.ErrOptionWithSkipRetry())
|
|
}
|
|
logger.LogInfo(c, fmt.Sprintf("用户 %d 预扣费 %s, 预扣费后剩余额度: %s", relayInfo.UserId, logger.FormatQuota(preConsumedQuota), logger.FormatQuota(userQuota-preConsumedQuota)))
|
|
}
|
|
relayInfo.FinalPreConsumedQuota = preConsumedQuota
|
|
return preConsumedQuota, nil
|
|
}
|