package model import ( "strings" "sync" "github.com/QuantumNous/new-api/common" ) type EmailQuotaRule struct { Id int `json:"id" gorm:"primaryKey"` EmailSuffix string `json:"email_suffix" gorm:"size:128;not null;uniqueIndex"` Quota int64 `json:"quota" gorm:"not null"` Enabled bool `json:"enabled" gorm:"default:1"` Description string `json:"description" gorm:"size:256"` CreatedTime int64 `json:"created_time" gorm:"bigint"` UpdatedTime int64 `json:"updated_time" gorm:"bigint"` } var ( emailQuotaCache map[string]int64 emailQuotaCacheMu sync.RWMutex ) // LoadEmailQuotaCache loads all enabled rules into memory cache. func LoadEmailQuotaCache() { var rules []EmailQuotaRule DB.Where("enabled = ?", true).Find(&rules) cache := make(map[string]int64, len(rules)) for _, r := range rules { cache[strings.ToLower(r.EmailSuffix)] = r.Quota } emailQuotaCacheMu.Lock() emailQuotaCache = cache emailQuotaCacheMu.Unlock() } // MatchEmailQuotaRule checks if the email matches any enabled suffix rule. // Returns the matched quota, or -1 if no match. func MatchEmailQuotaRule(email string) int64 { if email == "" { return -1 } at := strings.LastIndex(email, "@") if at < 0 { return -1 } suffix := strings.ToLower(email[at:]) emailQuotaCacheMu.RLock() defer emailQuotaCacheMu.RUnlock() if quota, ok := emailQuotaCache[suffix]; ok { return quota } return -1 } func boolToInt(b bool) int { if b { return 1 } return 0 } func (r *EmailQuotaRule) Insert() error { r.CreatedTime = common.GetTimestamp() r.UpdatedTime = r.CreatedTime err := DB.Model(&EmailQuotaRule{}).Create(map[string]interface{}{ "email_suffix": r.EmailSuffix, "quota": r.Quota, "enabled": boolToInt(r.Enabled), "description": r.Description, "created_time": r.CreatedTime, "updated_time": r.UpdatedTime, }).Error if err != nil { return err } var last EmailQuotaRule if err := DB.Where("email_suffix = ?", r.EmailSuffix).First(&last).Error; err == nil { r.Id = last.Id } LoadEmailQuotaCache() return nil } func (r *EmailQuotaRule) Update() error { r.UpdatedTime = common.GetTimestamp() err := DB.Model(&EmailQuotaRule{}).Where("id = ?", r.Id).Updates(map[string]interface{}{ "email_suffix": r.EmailSuffix, "quota": r.Quota, "enabled": boolToInt(r.Enabled), "description": r.Description, "updated_time": r.UpdatedTime, }).Error if err != nil { return err } LoadEmailQuotaCache() return nil } func (r *EmailQuotaRule) Delete() error { err := DB.Delete(r).Error if err != nil { return err } LoadEmailQuotaCache() return nil } func GetAllEmailQuotaRules() ([]EmailQuotaRule, error) { var list []EmailQuotaRule err := DB.Order("id ASC").Find(&list).Error return list, err } func GetEmailQuotaRuleById(id int) (*EmailQuotaRule, error) { var r EmailQuotaRule err := DB.First(&r, id).Error if err != nil { return nil, err } return &r, nil } func GetEmailQuotaRuleBySuffix(suffix string) (*EmailQuotaRule, error) { var r EmailQuotaRule err := DB.Where("email_suffix = ?", suffix).First(&r).Error if err != nil { return nil, err } return &r, nil }