Usishchev Yury
105 lines
2.5 KiB
Go
105 lines
2.5 KiB
Go
package quota
|
|
|
|
import (
|
|
"strings"
|
|
"time"
|
|
|
|
"github.com/alecthomas/units"
|
|
"github.com/go-gost/core/logger"
|
|
"github.com/go-gost/x/config"
|
|
xquota "github.com/go-gost/x/limiter/quota"
|
|
)
|
|
|
|
// ParseQuotaLimiter is best-effort: malformed fields are logged and treated as unset.
|
|
func ParseQuotaLimiter(cfg *config.QuotaConfig) *xquota.Limiter {
|
|
if cfg == nil {
|
|
return nil
|
|
}
|
|
|
|
log := logger.Default().WithFields(map[string]any{
|
|
"kind": "quota",
|
|
"quota": cfg.Name,
|
|
})
|
|
|
|
var limit uint64
|
|
if s := strings.TrimSpace(cfg.Limit); s != "" && s != "0" {
|
|
if v, err := units.ParseBase2Bytes(s); err != nil {
|
|
log.Warnf("quota: parse limit %q: %v", s, err)
|
|
} else if v > 0 {
|
|
limit = uint64(v)
|
|
}
|
|
}
|
|
|
|
var dir xquota.Direction
|
|
switch strings.ToLower(strings.TrimSpace(cfg.Direction)) {
|
|
case "in":
|
|
dir = xquota.DirectionIn
|
|
case "out":
|
|
dir = xquota.DirectionOut
|
|
default:
|
|
dir = xquota.DirectionTotal
|
|
}
|
|
|
|
var flush time.Duration
|
|
if s := strings.TrimSpace(cfg.Flush); s != "" {
|
|
if d, err := time.ParseDuration(s); err != nil {
|
|
log.Warnf("quota: parse flush %q: %v", s, err)
|
|
} else {
|
|
flush = d
|
|
}
|
|
}
|
|
|
|
return xquota.NewLimiter(cfg.Name, xquota.Options{
|
|
Limit: limit,
|
|
StartsAt: parseQuotaTime(cfg.StartsAt, log, "startsAt"),
|
|
ExpiresAt: parseQuotaTime(cfg.ExpiresAt, log, "expiresAt"),
|
|
Direction: dir,
|
|
Flush: flush,
|
|
Store: parseQuotaStore(cfg.Store, log),
|
|
Logger: log,
|
|
})
|
|
}
|
|
|
|
func parseQuotaTime(s string, log logger.Logger, field string) time.Time {
|
|
s = strings.TrimSpace(s)
|
|
if s == "" {
|
|
return time.Time{}
|
|
}
|
|
t, err := time.Parse(time.RFC3339, s)
|
|
if err != nil {
|
|
log.Warnf("quota: parse %s %q: %v", field, s, err)
|
|
return time.Time{}
|
|
}
|
|
return t
|
|
}
|
|
|
|
func parseQuotaStore(sc *config.QuotaStoreConfig, log logger.Logger) xquota.Store {
|
|
const defaultFile = "gost-quota.json"
|
|
|
|
if sc == nil || sc.Type == "" || strings.EqualFold(sc.Type, "file") {
|
|
path := defaultFile
|
|
if sc != nil && strings.TrimSpace(sc.File) != "" {
|
|
path = sc.File
|
|
}
|
|
return xquota.NewFileStore(path)
|
|
}
|
|
|
|
if strings.EqualFold(sc.Type, "redis") {
|
|
log.Warnf("quota: redis store not implemented; counter will not be persisted")
|
|
var rc config.QuotaRedisConfig
|
|
if sc.Redis != nil {
|
|
rc = *sc.Redis
|
|
}
|
|
return xquota.NewRedisStore(xquota.RedisConfig{
|
|
Addr: rc.Addr,
|
|
Username: rc.Username,
|
|
Password: rc.Password,
|
|
DB: rc.DB,
|
|
Key: rc.Key,
|
|
})
|
|
}
|
|
|
|
log.Warnf("quota: unknown store type %q; using file store %q", sc.Type, defaultFile)
|
|
return xquota.NewFileStore(defaultFile)
|
|
}
|