feat: add quota limiter and quota API (#107)

2026-06-20 08:44:42 +02:00
parent b7d807b464
commit b23553d37a
16 changed files with 1343 additions and 4 deletions
@@ -0,0 +1,104 @@
+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)
+}