From b23553d37a6954dcd7569f59c7c5ac58dbd01942 Mon Sep 17 00:00:00 2001 From: Usishchev Yury Date: Sat, 20 Jun 2026 08:44:42 +0200 Subject: [PATCH] feat: add quota limiter and quota API (#107) --- api/api.go | 7 + api/config.go | 23 ++ api/config_quota.go | 350 ++++++++++++++++++++++++++++++ config/config.go | 63 +++++- config/loader/loader.go | 12 + config/parsing/quota/parse.go | 104 +++++++++ config/parsing/service/parse.go | 5 + limiter/quota/limiter.go | 326 ++++++++++++++++++++++++++++ limiter/quota/limiter_test.go | 195 +++++++++++++++++ limiter/quota/store.go | 16 ++ limiter/quota/store_file.go | 64 ++++++ limiter/quota/store_redis.go | 31 +++ limiter/quota/wrapper/conn.go | 78 +++++++ limiter/quota/wrapper/listener.go | 58 +++++ registry/quota.go | 8 + registry/registry.go | 7 + 16 files changed, 1343 insertions(+), 4 deletions(-) create mode 100644 api/config_quota.go create mode 100644 config/parsing/quota/parse.go create mode 100644 limiter/quota/limiter.go create mode 100644 limiter/quota/limiter_test.go create mode 100644 limiter/quota/store.go create mode 100644 limiter/quota/store_file.go create mode 100644 limiter/quota/store_redis.go create mode 100644 limiter/quota/wrapper/conn.go create mode 100644 limiter/quota/wrapper/listener.go create mode 100644 registry/quota.go diff --git a/api/api.go b/api/api.go index 826f8649..b409e4f2 100644 --- a/api/api.go +++ b/api/api.go @@ -145,6 +145,13 @@ func Register(r *gin.Engine, opts *Options) { config.PUT("/limiters/:limiter", updateLimiter) config.DELETE("/limiters/:limiter", deleteLimiter) + config.GET("/quotas", getQuotaList) + config.GET("/quotas/:quota", getQuota) + config.POST("/quotas", createQuota) + config.PUT("/quotas/:quota", updateQuota) + config.DELETE("/quotas/:quota", deleteQuota) + config.POST("/quotas/:quota/reset", resetQuota) + config.GET("/climiters", getConnLimiterList) config.GET("/climiters/:limiter", getConnLimiter) config.POST("/climiters", createConnLimiter) diff --git a/api/config.go b/api/config.go index cd58b177..c0a191f3 100644 --- a/api/config.go +++ b/api/config.go @@ -35,6 +35,29 @@ func fillServiceStatus(svc *config.ServiceConfig) { OutputBytes: st.Get(stats.KindOutputBytes), } } + var stopped bool + for _, qname := range svc.Quotas { + lim := registry.QuotaLimiterRegistry().Get(qname) + if lim == nil { + continue + } + qs := lim.Snapshot() + if qs.Blocked { + stopped = true + } + svc.Status.Quotas = append(svc.Status.Quotas, config.ServiceQuota{ + Name: qname, + Used: qs.Used, + Limit: qs.Limit, + StartsAt: qs.StartsAtUnix, + ExpiresAt: qs.ExpiresAtUnix, + Active: qs.Active, + Expired: qs.Expired, + Blocked: qs.Blocked, + Direction: qs.Direction, + }) + } + svc.Status.StoppedByLimit = stopped for _, ev := range status.Events() { if !ev.Time.IsZero() { svc.Status.Events = append(svc.Status.Events, config.ServiceEvent{ diff --git a/api/config_quota.go b/api/config_quota.go new file mode 100644 index 00000000..2244333b --- /dev/null +++ b/api/config_quota.go @@ -0,0 +1,350 @@ +package api + +import ( + "fmt" + "net/http" + "strings" + + "github.com/gin-gonic/gin" + "github.com/go-gost/x/config" + parser "github.com/go-gost/x/config/parsing/quota" + "github.com/go-gost/x/limiter/quota" + "github.com/go-gost/x/registry" +) + +func fillQuotaStatus(q *config.QuotaConfig) { + lim := registry.QuotaLimiterRegistry().Get(q.Name) + if lim == nil { + return + } + s := lim.Snapshot() + q.Status = &config.QuotaStatus{ + Used: s.Used, + Limit: s.Limit, + StartsAt: s.StartsAtUnix, + ExpiresAt: s.ExpiresAtUnix, + Active: s.Active, + Expired: s.Expired, + Blocked: s.Blocked, + Direction: s.Direction, + } +} + +// swagger:parameters getQuotaListRequest +type getQuotaListRequest struct { +} + +// successful operation. +// swagger:response getQuotaListResponse +type getQuotaListResponse struct { + // in: body + Data quotaList +} + +type quotaList struct { + Count int `json:"count"` + List []*config.QuotaConfig `json:"list"` +} + +func getQuotaList(ctx *gin.Context) { + // swagger:route GET /config/quotas Quota getQuotaListRequest + // + // Get quota list. + // + // Security: + // basicAuth: [] + // + // Responses: + // 200: getQuotaListResponse + + var req getQuotaListRequest + ctx.ShouldBindQuery(&req) + + list := config.Global().Quotas + for _, q := range list { + if q == nil { + continue + } + fillQuotaStatus(q) + } + + ctx.JSON(http.StatusOK, Response{ + Data: quotaList{ + Count: len(list), + List: list, + }, + }) +} + +// swagger:parameters getQuotaRequest +type getQuotaRequest struct { + // in: path + // required: true + Quota string `uri:"quota" json:"quota"` +} + +// successful operation. +// swagger:response getQuotaResponse +type getQuotaResponse struct { + // in: body + Data *config.QuotaConfig +} + +func getQuota(ctx *gin.Context) { + // swagger:route GET /config/quotas/{quota} Quota getQuotaRequest + // + // Get quota. + // + // Security: + // basicAuth: [] + // + // Responses: + // 200: getQuotaResponse + + var req getQuotaRequest + ctx.ShouldBindUri(&req) + + var resp getQuotaResponse + + for _, q := range config.Global().Quotas { + if q == nil { + continue + } + if q.Name == req.Quota { + fillQuotaStatus(q) + resp.Data = q + } + } + + ctx.JSON(http.StatusOK, Response{ + Data: resp.Data, + }) +} + +// swagger:parameters createQuotaRequest +type createQuotaRequest struct { + // in: body + Data config.QuotaConfig `json:"data"` +} + +// successful operation. +// swagger:response createQuotaResponse +type createQuotaResponse struct { + Data Response +} + +func createQuota(ctx *gin.Context) { + // swagger:route POST /config/quotas Quota createQuotaRequest + // + // Create a new quota, the name of the quota must be unique in quota list. + // + // Security: + // basicAuth: [] + // + // Responses: + // 200: createQuotaResponse + + var req createQuotaRequest + ctx.ShouldBindJSON(&req.Data) + + name := strings.TrimSpace(req.Data.Name) + if name == "" { + writeError(ctx, NewError(http.StatusBadRequest, ErrCodeInvalid, "quota name is required")) + return + } + req.Data.Name = name + req.Data.Status = nil + + if registry.QuotaLimiterRegistry().IsRegistered(name) { + writeError(ctx, NewError(http.StatusBadRequest, ErrCodeDup, fmt.Sprintf("quota %s already exists", name))) + return + } + + v := parser.ParseQuotaLimiter(&req.Data) + + if err := registry.QuotaLimiterRegistry().Register(name, v); err != nil { + writeError(ctx, NewError(http.StatusBadRequest, ErrCodeDup, fmt.Sprintf("quota %s already exists", name))) + return + } + + config.OnUpdate(func(c *config.Config) error { + c.Quotas = append(c.Quotas, &req.Data) + return nil + }) + + ctx.JSON(http.StatusOK, Response{ + Msg: "OK", + }) +} + +// swagger:parameters updateQuotaRequest +type updateQuotaRequest struct { + // in: path + // required: true + Quota string `uri:"quota" json:"quota"` + // in: body + Data config.QuotaConfig `json:"data"` +} + +// successful operation. +// swagger:response updateQuotaResponse +type updateQuotaResponse struct { + Data Response +} + +func updateQuota(ctx *gin.Context) { + // swagger:route PUT /config/quotas/{quota} Quota updateQuotaRequest + // + // Update quota by name, the quota must already exist. The cumulative counter + // is preserved across the update as long as the window is unchanged. + // + // Security: + // basicAuth: [] + // + // Responses: + // 200: updateQuotaResponse + + var req updateQuotaRequest + ctx.ShouldBindUri(&req) + ctx.ShouldBindJSON(&req.Data) + + name := strings.TrimSpace(req.Quota) + + if !registry.QuotaLimiterRegistry().IsRegistered(name) { + writeError(ctx, NewError(http.StatusBadRequest, ErrCodeNotFound, fmt.Sprintf("quota %s not found", name))) + return + } + + req.Data.Name = name + req.Data.Status = nil + + v := parser.ParseQuotaLimiter(&req.Data) + + registry.QuotaLimiterRegistry().Unregister(name) + + if err := registry.QuotaLimiterRegistry().Register(name, v); err != nil { + writeError(ctx, NewError(http.StatusBadRequest, ErrCodeDup, fmt.Sprintf("quota %s already exists", name))) + return + } + + config.OnUpdate(func(c *config.Config) error { + for i := range c.Quotas { + if c.Quotas[i].Name == name { + c.Quotas[i] = &req.Data + break + } + } + return nil + }) + + ctx.JSON(http.StatusOK, Response{ + Msg: "OK", + }) +} + +// swagger:parameters deleteQuotaRequest +type deleteQuotaRequest struct { + // in: path + // required: true + Quota string `uri:"quota" json:"quota"` +} + +// successful operation. +// swagger:response deleteQuotaResponse +type deleteQuotaResponse struct { + Data Response +} + +func deleteQuota(ctx *gin.Context) { + // swagger:route DELETE /config/quotas/{quota} Quota deleteQuotaRequest + // + // Delete quota by name. Services referencing it stop being limited (fail-open). + // + // Security: + // basicAuth: [] + // + // Responses: + // 200: deleteQuotaResponse + + var req deleteQuotaRequest + ctx.ShouldBindUri(&req) + + name := strings.TrimSpace(req.Quota) + + if !registry.QuotaLimiterRegistry().IsRegistered(name) { + writeError(ctx, NewError(http.StatusBadRequest, ErrCodeNotFound, fmt.Sprintf("quota %s not found", name))) + return + } + registry.QuotaLimiterRegistry().Unregister(name) + + config.OnUpdate(func(c *config.Config) error { + quotas := c.Quotas + c.Quotas = nil + for _, q := range quotas { + if q.Name == name { + continue + } + c.Quotas = append(c.Quotas, q) + } + return nil + }) + + ctx.JSON(http.StatusOK, Response{ + Msg: "OK", + }) +} + +// swagger:parameters resetQuotaRequest +type resetQuotaRequest struct { + // in: path + // required: true + Quota string `uri:"quota" json:"quota"` + // in: body + Data resetQuotaData `json:"data"` +} + +type resetQuotaData struct { + Used *uint64 `json:"used,omitempty"` +} + +// successful operation. +// swagger:response resetQuotaResponse +type resetQuotaResponse struct { + Data Response +} + +func resetQuota(ctx *gin.Context) { + // swagger:route POST /config/quotas/{quota}/reset Quota resetQuotaRequest + // + // Overwrite the cumulative counter of a quota (defaults to 0). The new value + // is applied immediately and persisted; if the quota was blocking, services + // resume once the value is below the limit. + // + // Security: + // basicAuth: [] + // + // Responses: + // 200: resetQuotaResponse + + var req resetQuotaRequest + ctx.ShouldBindUri(&req) + ctx.ShouldBindJSON(&req.Data) + + name := strings.TrimSpace(req.Quota) + + lim := registry.QuotaLimiterRegistry().Get(name) + if lim == nil { + writeError(ctx, NewError(http.StatusBadRequest, ErrCodeNotFound, fmt.Sprintf("quota %s not found", name))) + return + } + + used := uint64(0) + if req.Data.Used != nil { + used = *req.Data.Used + } + lim.Update(quota.Update{Used: &used}) + + ctx.JSON(http.StatusOK, Response{ + Msg: "OK", + }) +} diff --git a/config/config.go b/config/config.go index 306dd40e..fee91d51 100644 --- a/config/config.go +++ b/config/config.go @@ -470,6 +470,7 @@ type ServiceConfig struct { Resolver string `yaml:",omitempty" json:"resolver,omitempty"` Hosts string `yaml:",omitempty" json:"hosts,omitempty"` Limiter string `yaml:",omitempty" json:"limiter,omitempty"` + Quotas []string `yaml:",omitempty" json:"quotas,omitempty"` CLimiter string `yaml:"climiter,omitempty" json:"climiter,omitempty"` RLimiter string `yaml:"rlimiter,omitempty" json:"rlimiter,omitempty"` Logger string `yaml:",omitempty" json:"logger,omitempty"` @@ -485,10 +486,12 @@ type ServiceConfig struct { } type ServiceStatus struct { - CreateTime int64 `yaml:"createTime" json:"createTime"` - State string `yaml:"state" json:"state"` - Events []ServiceEvent `yaml:",omitempty" json:"events,omitempty"` - Stats *ServiceStats `yaml:",omitempty" json:"stats,omitempty"` + CreateTime int64 `yaml:"createTime" json:"createTime"` + State string `yaml:"state" json:"state"` + Events []ServiceEvent `yaml:",omitempty" json:"events,omitempty"` + Stats *ServiceStats `yaml:",omitempty" json:"stats,omitempty"` + StoppedByLimit bool `yaml:"stoppedByLimit,omitempty" json:"stoppedByLimit,omitempty"` + Quotas []ServiceQuota `yaml:",omitempty" json:"quotas,omitempty"` } type ServiceEvent struct { @@ -504,6 +507,57 @@ type ServiceStats struct { OutputBytes uint64 `yaml:"outputBytes" json:"outputBytes"` } +// QuotaConfig is a named cumulative traffic-volume limiter; the same name +// referenced by several services (ServiceConfig.Quotas) shares one counter. +type QuotaConfig struct { + Name string `json:"name"` + Limit string `yaml:",omitempty" json:"limit,omitempty"` + StartsAt string `yaml:"startsAt,omitempty" json:"startsAt,omitempty"` + ExpiresAt string `yaml:"expiresAt,omitempty" json:"expiresAt,omitempty"` + Direction string `yaml:",omitempty" json:"direction,omitempty"` + Flush string `yaml:",omitempty" json:"flush,omitempty"` + Store *QuotaStoreConfig `yaml:",omitempty" json:"store,omitempty"` + // read-only + Status *QuotaStatus `yaml:",omitempty" json:"status,omitempty"` +} + +type QuotaStoreConfig struct { + Type string `json:"type"` // "file" (default) | "redis" (stub, not implemented) + File string `yaml:",omitempty" json:"file,omitempty"` + Redis *QuotaRedisConfig `yaml:",omitempty" json:"redis,omitempty"` +} + +type QuotaRedisConfig struct { + Addr string `json:"addr"` + DB int `yaml:",omitempty" json:"db,omitempty"` + Username string `yaml:",omitempty" json:"username,omitempty"` + Password string `yaml:",omitempty" json:"password,omitempty"` + Key string `yaml:",omitempty" json:"key,omitempty"` +} + +type QuotaStatus struct { + Used uint64 `yaml:"used" json:"used"` + Limit uint64 `yaml:"limit" json:"limit"` + StartsAt int64 `yaml:"startsAt,omitempty" json:"startsAt,omitempty"` + ExpiresAt int64 `yaml:"expiresAt,omitempty" json:"expiresAt,omitempty"` + Active bool `yaml:"active" json:"active"` + Expired bool `yaml:"expired,omitempty" json:"expired,omitempty"` + Blocked bool `yaml:"blocked" json:"blocked"` + Direction string `yaml:"direction,omitempty" json:"direction,omitempty"` +} + +type ServiceQuota struct { + Name string `yaml:"name" json:"name"` + Used uint64 `yaml:"used" json:"used"` + Limit uint64 `yaml:"limit" json:"limit"` + StartsAt int64 `yaml:"startsAt,omitempty" json:"startsAt,omitempty"` + ExpiresAt int64 `yaml:"expiresAt,omitempty" json:"expiresAt,omitempty"` + Active bool `yaml:"active" json:"active"` + Expired bool `yaml:"expired,omitempty" json:"expired,omitempty"` + Blocked bool `yaml:"blocked" json:"blocked"` + Direction string `yaml:"direction,omitempty" json:"direction,omitempty"` +} + type ChainConfig struct { Name string `json:"name"` Hops []*HopConfig `json:"hops"` @@ -573,6 +627,7 @@ type Config struct { SDs []*SDConfig `yaml:"sds,omitempty" json:"sds,omitempty"` Recorders []*RecorderConfig `yaml:",omitempty" json:"recorders,omitempty"` Limiters []*LimiterConfig `yaml:",omitempty" json:"limiters,omitempty"` + Quotas []*QuotaConfig `yaml:",omitempty" json:"quotas,omitempty"` CLimiters []*LimiterConfig `yaml:"climiters,omitempty" json:"climiters,omitempty"` RLimiters []*LimiterConfig `yaml:"rlimiters,omitempty" json:"rlimiters,omitempty"` Observers []*ObserverConfig `yaml:",omitempty" json:"observers,omitempty"` diff --git a/config/loader/loader.go b/config/loader/loader.go index 0d5af67a..2e9c5f58 100644 --- a/config/loader/loader.go +++ b/config/loader/loader.go @@ -35,11 +35,13 @@ import ( limiter_parser "github.com/go-gost/x/config/parsing/limiter" logger_parser "github.com/go-gost/x/config/parsing/logger" observer_parser "github.com/go-gost/x/config/parsing/observer" + quota_parser "github.com/go-gost/x/config/parsing/quota" recorder_parser "github.com/go-gost/x/config/parsing/recorder" resolver_parser "github.com/go-gost/x/config/parsing/resolver" router_parser "github.com/go-gost/x/config/parsing/router" sd_parser "github.com/go-gost/x/config/parsing/sd" service_parser "github.com/go-gost/x/config/parsing/service" + quota "github.com/go-gost/x/limiter/quota" "github.com/go-gost/x/registry" ) @@ -252,6 +254,16 @@ func register(cfg *config.Config) error { } } + { + var entries []named[*quota.Limiter] + for _, c := range cfg.Quotas { + entries = append(entries, named[*quota.Limiter]{c.Name, quota_parser.ParseQuotaLimiter(c)}) + } + if err := registerGroup(entries, registry.QuotaLimiterRegistry()); err != nil { + return err + } + } + { var entries []named[conn.ConnLimiter] for _, c := range cfg.CLimiters { diff --git a/config/parsing/quota/parse.go b/config/parsing/quota/parse.go new file mode 100644 index 00000000..d5d827b2 --- /dev/null +++ b/config/parsing/quota/parse.go @@ -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) +} diff --git a/config/parsing/service/parse.go b/config/parsing/service/parse.go index f6dd87ed..e64ff70e 100644 --- a/config/parsing/service/parse.go +++ b/config/parsing/service/parse.go @@ -29,6 +29,7 @@ import ( logger_parser "github.com/go-gost/x/config/parsing/logger" selector_parser "github.com/go-gost/x/config/parsing/selector" tls_util "github.com/go-gost/x/internal/util/tls" + quota_wrapper "github.com/go-gost/x/limiter/quota/wrapper" cache_limiter "github.com/go-gost/x/limiter/traffic/cache" "github.com/go-gost/x/metadata" mdutil "github.com/go-gost/x/metadata/util" @@ -246,6 +247,10 @@ func ParseService(cfg *config.ServiceConfig) (service.Service, error) { return nil, err } + for _, qname := range cfg.Quotas { + ln = quota_wrapper.WrapListener(ln, strings.TrimSpace(qname)) + } + handlerLogger := serviceLogger.WithFields(map[string]any{ "kind": "handler", }) diff --git a/limiter/quota/limiter.go b/limiter/quota/limiter.go new file mode 100644 index 00000000..24e1a936 --- /dev/null +++ b/limiter/quota/limiter.go @@ -0,0 +1,326 @@ +// Package quota implements a persisted, long-term traffic-volume limiter that +// can be shared across services by name. Unlike the rate limiters in +// limiter/traffic and limiter/rate, it accumulates total bytes within a window +// [startsAt, expiresAt) and stops the referencing service(s) once a byte limit +// is reached. Enforcement is fail-open outside the window. +package quota + +import ( + "errors" + "sync" + "sync/atomic" + "time" + + "github.com/go-gost/core/logger" + xlogger "github.com/go-gost/x/logger" +) + +var ErrQuotaExceeded = errors.New("quota: traffic limit reached") + +const defaultFlushInterval = 10 * time.Second + +type Direction int + +const ( + DirectionTotal Direction = iota + DirectionIn + DirectionOut +) + +func (d Direction) String() string { + switch d { + case DirectionIn: + return "in" + case DirectionOut: + return "out" + default: + return "total" + } +} + +// Options seeds a Limiter. Limit and the window are config-authoritative; a +// persisted counter is restored only for a matching window (see NewLimiter). +type Options struct { + Limit uint64 + StartsAt time.Time + ExpiresAt time.Time + Direction Direction + Flush time.Duration + Store Store + Logger logger.Logger +} + +// Update overwrites runtime state; a nil field is left unchanged. +type Update struct { + Used *uint64 + Limit *uint64 + StartsAt *time.Time + ExpiresAt *time.Time +} + +type Snapshot struct { + Used uint64 + Limit uint64 + StartsAtUnix int64 + ExpiresAtUnix int64 + Active bool + Expired bool + Blocked bool + Direction string +} + +type Limiter struct { + name string + direction Direction + flush time.Duration + store Store + log logger.Logger + + used atomic.Uint64 + limit atomic.Uint64 + startsAt atomic.Int64 // unixnano; 0 = unset + expiresAt atomic.Int64 // unixnano; 0 = unset + blocked atomic.Bool + dirty atomic.Bool + closed atomic.Bool + + mu sync.Mutex + waitCh chan struct{} // closed+replaced to broadcast a state change + + closeOnce sync.Once + done chan struct{} +} + +func NewLimiter(name string, opts Options) *Limiter { + l := &Limiter{ + name: name, + direction: opts.Direction, + flush: opts.Flush, + store: opts.Store, + log: opts.Logger, + waitCh: make(chan struct{}), + done: make(chan struct{}), + } + if l.flush <= 0 { + l.flush = defaultFlushInterval + } + if l.log == nil { + l.log = xlogger.Nop() + } + + l.limit.Store(opts.Limit) + sa := unixNanoOrZero(opts.StartsAt) + ea := unixNanoOrZero(opts.ExpiresAt) + l.startsAt.Store(sa) + l.expiresAt.Store(ea) + + // Restore the counter only within the same window: a changed window (a new + // period pushed via config) starts fresh. + if l.store != nil { + if rec, ok, err := l.store.Load(name); err != nil { + l.log.Warnf("quota: load %s: %v", name, err) + } else if ok && rec.StartsAt == sa && rec.ExpiresAt == ea { + l.used.Store(rec.Used) + } + } + + l.blocked.Store(l.enforcing(time.Now()) && l.used.Load() >= l.limit.Load()) + go l.run() + return l +} + +func (l *Limiter) active(now time.Time) bool { + n := now.UnixNano() + if sa := l.startsAt.Load(); sa != 0 && n < sa { + return false + } + if ea := l.expiresAt.Load(); ea != 0 && n >= ea { + return false + } + return true +} + +func (l *Limiter) enforcing(now time.Time) bool { + return l.limit.Load() > 0 && l.active(now) +} + +func (l *Limiter) Blocked() bool { + if l.closed.Load() || !l.blocked.Load() { + return false + } + return l.enforcing(time.Now()) +} + +func (l *Limiter) AddIn(n int) { l.add(int64(n), DirectionIn) } +func (l *Limiter) AddOut(n int) { l.add(int64(n), DirectionOut) } + +func (l *Limiter) add(n int64, dir Direction) { + if n <= 0 || l.closed.Load() { + return + } + switch l.direction { + case DirectionIn: + if dir != DirectionIn { + return + } + case DirectionOut: + if dir != DirectionOut { + return + } + } + if !l.active(time.Now()) { + return + } + used := l.used.Add(uint64(n)) + l.dirty.Store(true) + if lim := l.limit.Load(); lim > 0 && used >= lim { + l.block() + } +} + +// block is hot-path: no lock, no I/O, no wakeup (parking only happens on the +// next Accept). +func (l *Limiter) block() { + if l.blocked.CompareAndSwap(false, true) { + l.log.Warnf("quota: %s reached limit %d bytes, stopping", l.name, l.limit.Load()) + } +} + +// WaitChan must be fetched before checking Blocked to avoid a missed wakeup. +func (l *Limiter) WaitChan() <-chan struct{} { + l.mu.Lock() + defer l.mu.Unlock() + return l.waitCh +} + +func (l *Limiter) notify() { + l.mu.Lock() + close(l.waitCh) + l.waitCh = make(chan struct{}) + l.mu.Unlock() +} + +func (l *Limiter) Update(u Update) { + if u.Used != nil { + l.used.Store(*u.Used) + } + if u.Limit != nil { + l.limit.Store(*u.Limit) + } + if u.StartsAt != nil { + l.startsAt.Store(unixNanoOrZero(*u.StartsAt)) + } + if u.ExpiresAt != nil { + l.expiresAt.Store(unixNanoOrZero(*u.ExpiresAt)) + } + l.reevaluate() +} + +func (l *Limiter) reevaluate() { + now := time.Now() + shouldBlock := l.enforcing(now) && l.used.Load() >= l.limit.Load() + l.blocked.Store(shouldBlock) + if !shouldBlock { + l.notify() + } + l.flushNow() +} + +func (l *Limiter) Snapshot() Snapshot { + now := time.Now() + ea := l.expiresAt.Load() + return Snapshot{ + Used: l.used.Load(), + Limit: l.limit.Load(), + StartsAtUnix: nanoToSec(l.startsAt.Load()), + ExpiresAtUnix: nanoToSec(ea), + Active: l.active(now), + Expired: ea != 0 && now.UnixNano() >= ea, + Blocked: l.Blocked(), + Direction: l.direction.String(), + } +} + +func (l *Limiter) run() { + flushT := time.NewTicker(l.flush) + defer flushT.Stop() + boundT := time.NewTicker(time.Second) + defer boundT.Stop() + + wasActive := l.active(time.Now()) + + for { + select { + case <-l.done: + return + + case <-flushT.C: + if l.dirty.Swap(false) { + l.persist() + } + + case <-boundT.C: + now := time.Now() + act := l.active(now) + if act == wasActive { + continue + } + wasActive = act + if act { + l.blocked.Store(l.limit.Load() > 0 && l.used.Load() >= l.limit.Load()) + } else { + l.blocked.Store(false) + } + l.notify() + l.persist() + } + } +} + +func (l *Limiter) flushNow() { + l.dirty.Store(false) + l.persist() +} + +func (l *Limiter) persist() { + if l.store == nil { + return + } + rec := Record{ + Used: l.used.Load(), + Limit: l.limit.Load(), + StartsAt: l.startsAt.Load(), + ExpiresAt: l.expiresAt.Load(), + UpdatedAt: time.Now().UnixNano(), + } + if err := l.store.Save(l.name, rec); err != nil { + l.log.Warnf("quota: save %s: %v", l.name, err) + } +} + +// Close makes the limiter inert (Blocked false, counting a no-op) so deleting a +// shared quota releases the referencing services instead of tearing them down. +// Called by the registry on Unregister. +func (l *Limiter) Close() error { + l.closeOnce.Do(func() { + l.closed.Store(true) + close(l.done) + l.flushNow() + l.notify() + }) + return nil +} + +func unixNanoOrZero(t time.Time) int64 { + if t.IsZero() { + return 0 + } + return t.UnixNano() +} + +func nanoToSec(n int64) int64 { + if n == 0 { + return 0 + } + return n / int64(time.Second) +} diff --git a/limiter/quota/limiter_test.go b/limiter/quota/limiter_test.go new file mode 100644 index 00000000..8e159402 --- /dev/null +++ b/limiter/quota/limiter_test.go @@ -0,0 +1,195 @@ +package quota + +import ( + "encoding/json" + "os" + "path/filepath" + "testing" + "time" +) + +func newTestLimiter(t *testing.T, opts Options) *Limiter { + t.Helper() + l := NewLimiter("test", opts) + t.Cleanup(func() { l.Close() }) + return l +} + +func TestLimiterBlocksAtLimit(t *testing.T) { + l := newTestLimiter(t, Options{Limit: 100}) + + l.AddIn(60) + if l.Blocked() { + t.Fatalf("should not block at 60/100") + } + l.AddOut(40) + if !l.Blocked() { + t.Fatalf("should block at 100/100") + } + if s := l.Snapshot(); !s.Blocked || s.Used != 100 { + t.Fatalf("snapshot=%+v, want blocked used=100", s) + } +} + +func TestLimiterUnblockOnOverwrite(t *testing.T) { + l := newTestLimiter(t, Options{Limit: 100}) + l.AddIn(150) + if !l.Blocked() { + t.Fatalf("expected blocked") + } + + zero := uint64(0) + l.Update(Update{Used: &zero}) + + if l.Blocked() { + t.Fatalf("expected unblocked after reset") + } + if got := l.Snapshot().Used; got != 0 { + t.Fatalf("used=%d, want 0", got) + } +} + +func TestLimiterDirection(t *testing.T) { + l := newTestLimiter(t, Options{Limit: 100, Direction: DirectionIn}) + + l.AddOut(500) + if u := l.Snapshot().Used; u != 0 { + t.Fatalf("outbound counted under DirectionIn: used=%d", u) + } + l.AddIn(100) + if !l.Blocked() { + t.Fatalf("expected blocked after inbound reaches limit") + } +} + +func TestLimiterWindowNotStarted(t *testing.T) { + l := newTestLimiter(t, Options{Limit: 100, StartsAt: time.Now().Add(time.Hour)}) + + l.AddIn(500) + s := l.Snapshot() + if s.Used != 0 { + t.Fatalf("counted before startsAt: used=%d", s.Used) + } + if s.Active { + t.Fatalf("should be inactive before startsAt") + } + if l.Blocked() { + t.Fatalf("should not block before startsAt") + } +} + +func TestLimiterFailOpenAfterExpiry(t *testing.T) { + l := newTestLimiter(t, Options{Limit: 100, ExpiresAt: time.Now().Add(-time.Hour)}) + + l.AddIn(500) + s := l.Snapshot() + if s.Used != 0 { + t.Fatalf("counted after expiry: used=%d", s.Used) + } + if !s.Expired { + t.Fatalf("snapshot should report expired") + } + if l.Blocked() { + t.Fatalf("should be fail-open (unblocked) after expiry") + } +} + +func TestLimiterPersistsToDisk(t *testing.T) { + path := filepath.Join(t.TempDir(), "quota.json") + store := NewFileStore(path) + + l := NewLimiter("svc", Options{Limit: 1000, Store: store}) + l.AddIn(300) + l.AddOut(200) + l.Close() + + b, err := os.ReadFile(path) + if err != nil { + t.Fatalf("read store file: %v", err) + } + var data map[string]Record + if err := json.Unmarshal(b, &data); err != nil { + t.Fatalf("unmarshal: %v", err) + } + if rec := data["svc"]; rec.Used != 500 { + t.Fatalf("persisted used=%d, want 500", rec.Used) + } +} + +func TestLimiterRestoresFromDisk(t *testing.T) { + path := filepath.Join(t.TempDir(), "quota.json") + if err := os.WriteFile(path, []byte(`{"svc":{"used":777,"limit":1000}}`), 0o644); err != nil { + t.Fatal(err) + } + + store := NewFileStore(path) + l := NewLimiter("svc", Options{Limit: 1000, Store: store}) + t.Cleanup(func() { l.Close() }) + + if got := l.Snapshot().Used; got != 777 { + t.Fatalf("restored used=%d, want 777", got) + } +} + +func TestWaitChanNotifiesOnReset(t *testing.T) { + l := newTestLimiter(t, Options{Limit: 100}) + l.AddIn(200) + if !l.Blocked() { + t.Fatal("expected blocked") + } + + ch := l.WaitChan() + zero := uint64(0) + l.Update(Update{Used: &zero}) + + select { + case <-ch: + case <-time.After(time.Second): + t.Fatal("WaitChan not notified after reset") + } + if l.Blocked() { + t.Fatal("expected unblocked after reset") + } +} + +func TestClosedLimiterIsInert(t *testing.T) { + l := NewLimiter("c", Options{Limit: 100}) + l.AddIn(200) + if !l.Blocked() { + t.Fatal("expected blocked") + } + + ch := l.WaitChan() + l.Close() + + select { + case <-ch: + case <-time.After(time.Second): + t.Fatal("WaitChan not notified on close") + } + if l.Blocked() { + t.Fatal("closed limiter must be inert (not blocking)") + } + l.AddIn(1000) + if u := l.Snapshot().Used; u != 200 { + t.Fatalf("closed limiter counted: used=%d, want 200", u) + } +} + +func TestLimiterResetsOnWindowChange(t *testing.T) { + path := filepath.Join(t.TempDir(), "quota.json") + store := NewFileStore(path) + + t1 := time.Now().Add(time.Hour) + l1 := NewLimiter("svc", Options{Limit: 1000, ExpiresAt: t1, Store: store}) + l1.AddIn(500) + l1.Close() + + t2 := time.Now().Add(2 * time.Hour) + l2 := NewLimiter("svc", Options{Limit: 1000, ExpiresAt: t2, Store: store}) + t.Cleanup(func() { l2.Close() }) + + if u := l2.Snapshot().Used; u != 0 { + t.Fatalf("counter not reset on window change: used=%d, want 0", u) + } +} diff --git a/limiter/quota/store.go b/limiter/quota/store.go new file mode 100644 index 00000000..601de4fb --- /dev/null +++ b/limiter/quota/store.go @@ -0,0 +1,16 @@ +package quota + +type Record struct { + Used uint64 `json:"used"` + Limit uint64 `json:"limit"` + StartsAt int64 `json:"startsAt"` // unixnano; 0 = unset + ExpiresAt int64 `json:"expiresAt"` // unixnano; 0 = unset + UpdatedAt int64 `json:"updatedAt"` +} + +// Store persists per-name quota records. Implementations must be concurrency-safe. +type Store interface { + // Load returns ok=false when no record exists yet. + Load(name string) (rec Record, ok bool, err error) + Save(name string, rec Record) error +} diff --git a/limiter/quota/store_file.go b/limiter/quota/store_file.go new file mode 100644 index 00000000..84ae781d --- /dev/null +++ b/limiter/quota/store_file.go @@ -0,0 +1,64 @@ +package quota + +import ( + "encoding/json" + "os" + "sync" +) + +// fileStore persists all records into one JSON file. A single instance is shared +// per path so concurrent per-name saves merge instead of clobbering each other. +type fileStore struct { + path string + mu sync.Mutex + data map[string]Record +} + +var fileStores sync.Map // path -> *fileStore + +func NewFileStore(path string) Store { + if v, ok := fileStores.Load(path); ok { + return v.(*fileStore) + } + fs := &fileStore{path: path, data: make(map[string]Record)} + fs.load() + actual, _ := fileStores.LoadOrStore(path, fs) + return actual.(*fileStore) +} + +func (fs *fileStore) load() { + b, err := os.ReadFile(fs.path) + if err != nil { + return + } + data := make(map[string]Record) + if json.Unmarshal(b, &data) == nil { + fs.data = data + } +} + +func (fs *fileStore) Load(name string) (Record, bool, error) { + fs.mu.Lock() + defer fs.mu.Unlock() + rec, ok := fs.data[name] + return rec, ok, nil +} + +func (fs *fileStore) Save(name string, rec Record) error { + fs.mu.Lock() + defer fs.mu.Unlock() + fs.data[name] = rec + return fs.writeLocked() +} + +func (fs *fileStore) writeLocked() error { + b, err := json.MarshalIndent(fs.data, "", " ") + if err != nil { + return err + } + tmp := fs.path + ".tmp" + if err := os.WriteFile(tmp, b, 0o644); err != nil { + return err + } + return os.Rename(tmp, fs.path) +} diff --git a/limiter/quota/store_redis.go b/limiter/quota/store_redis.go new file mode 100644 index 00000000..e889754e --- /dev/null +++ b/limiter/quota/store_redis.go @@ -0,0 +1,31 @@ +package quota + +type RedisConfig struct { + Addr string + Username string + Password string + DB int + Key string +} + +// redisStore is a no-op placeholder: a redis-configured quota behaves as +// in-memory only until the backend is implemented. +// +// TODO: implement using github.com/go-redis/redis/v8 (already a dependency). +// Suggested schema: HSET for Save, HGET for Load. +// See internal/loader/redis.go and recorder/redis.go for the patterns. +type redisStore struct { + cfg RedisConfig +} + +func NewRedisStore(cfg RedisConfig) Store { + return &redisStore{cfg: cfg} +} + +func (s *redisStore) Load(name string) (Record, bool, error) { + return Record{}, false, nil +} + +func (s *redisStore) Save(name string, rec Record) error { + return nil +} diff --git a/limiter/quota/wrapper/conn.go b/limiter/quota/wrapper/conn.go new file mode 100644 index 00000000..0c7d4c76 --- /dev/null +++ b/limiter/quota/wrapper/conn.go @@ -0,0 +1,78 @@ +package wrapper + +import ( + "context" + "net" + "syscall" + + "github.com/go-gost/x/ctx" + xio "github.com/go-gost/x/internal/io" + "github.com/go-gost/x/limiter/quota" + "github.com/go-gost/x/registry" +) + +type quotaConn struct { + net.Conn + name string +} + +func WrapConn(c net.Conn, name string) net.Conn { + if c == nil || name == "" { + return c + } + return "aConn{Conn: c, name: name} +} + +func (c *quotaConn) Read(b []byte) (n int, err error) { + lim := registry.QuotaLimiterRegistry().Get(c.name) + if lim != nil && lim.Blocked() { + return 0, quota.ErrQuotaExceeded + } + n, err = c.Conn.Read(b) + if lim != nil && n > 0 { + lim.AddIn(n) + } + return +} + +func (c *quotaConn) Write(b []byte) (n int, err error) { + lim := registry.QuotaLimiterRegistry().Get(c.name) + if lim != nil && lim.Blocked() { + return 0, quota.ErrQuotaExceeded + } + n, err = c.Conn.Write(b) + if lim != nil && n > 0 { + lim.AddOut(n) + } + return +} + +// Forward optional capabilities so wrapping does not hide them from handlers. + +func (c *quotaConn) SyscallConn() (syscall.RawConn, error) { + if sc, ok := c.Conn.(syscall.Conn); ok { + return sc.SyscallConn() + } + return nil, xio.ErrUnsupported +} + +func (c *quotaConn) CloseRead() error { + if sc, ok := c.Conn.(xio.CloseRead); ok { + return sc.CloseRead() + } + return xio.ErrUnsupported +} + +func (c *quotaConn) CloseWrite() error { + if sc, ok := c.Conn.(xio.CloseWrite); ok { + return sc.CloseWrite() + } + return xio.ErrUnsupported +} + +func (c *quotaConn) Context() context.Context { + if cc, ok := c.Conn.(ctx.Context); ok { + return cc.Context() + } + return nil +} diff --git a/limiter/quota/wrapper/listener.go b/limiter/quota/wrapper/listener.go new file mode 100644 index 00000000..e5c997ef --- /dev/null +++ b/limiter/quota/wrapper/listener.go @@ -0,0 +1,58 @@ +// Package wrapper adapts a named quota onto a listener and its connections, +// resolving the quota.Limiter by name from the registry on every call so +// create/update/delete takes effect live and several services can share one. +package wrapper + +import ( + "net" + "sync" + + "github.com/go-gost/core/listener" + "github.com/go-gost/x/registry" +) + +type quotaListener struct { + listener.Listener + name string + done chan struct{} + once sync.Once +} + +func WrapListener(ln listener.Listener, name string) listener.Listener { + if ln == nil || name == "" { + return ln + } + return "aListener{Listener: ln, name: name, done: make(chan struct{})} +} + +func (ln *quotaListener) Accept() (net.Conn, error) { + for { + lim := registry.QuotaLimiterRegistry().Get(ln.name) + if lim == nil { + break + } + ch := lim.WaitChan() // before Blocked, to avoid a missed wakeup + if !lim.Blocked() { + break + } + // Park (not error) so the service's accept loop survives; ln.done wakes + // it on listener close. + select { + case <-ch: + case <-ln.done: + return nil, net.ErrClosed + } + } + + c, err := ln.Listener.Accept() + if err != nil { + return nil, err + } + return WrapConn(c, ln.name), nil +} + +func (ln *quotaListener) Close() error { + // Wake a parked Accept, but do not close the (possibly shared) limiter. + ln.once.Do(func() { close(ln.done) }) + return ln.Listener.Close() +} diff --git a/registry/quota.go b/registry/quota.go new file mode 100644 index 00000000..fedc9631 --- /dev/null +++ b/registry/quota.go @@ -0,0 +1,8 @@ +package registry + +import "github.com/go-gost/x/limiter/quota" + +// quotaLimiterRegistry holds shared quota limiters; Unregister closes them. +type quotaLimiterRegistry struct { + registry[*quota.Limiter] +} diff --git a/registry/registry.go b/registry/registry.go index 4ffd15e2..f63c54c6 100644 --- a/registry/registry.go +++ b/registry/registry.go @@ -23,6 +23,7 @@ import ( "github.com/go-gost/core/router" "github.com/go-gost/core/sd" "github.com/go-gost/core/service" + "github.com/go-gost/x/limiter/quota" ) var ( @@ -49,6 +50,7 @@ var ( trafficLimiterReg reg.Registry[traffic.TrafficLimiter] = new(trafficLimiterRegistry) connLimiterReg reg.Registry[conn.ConnLimiter] = new(connLimiterRegistry) rateLimiterReg reg.Registry[rate.RateLimiter] = new(rateLimiterRegistry) + quotaLimiterReg reg.Registry[*quota.Limiter] = new(quotaLimiterRegistry) ingressReg reg.Registry[ingress.Ingress] = new(ingressRegistry) routerReg reg.Registry[router.Router] = new(routerRegistry) @@ -200,6 +202,11 @@ func RateLimiterRegistry() reg.Registry[rate.RateLimiter] { return rateLimiterReg } +// QuotaLimiterRegistry returns the global registry of quota limiter instances. +func QuotaLimiterRegistry() reg.Registry[*quota.Limiter] { + return quotaLimiterReg +} + // IngressRegistry returns the global registry of ingress instances. func IngressRegistry() reg.Registry[ingress.Ingress] { return ingressReg