diff --git a/limiter/conn/conn.go b/limiter/conn/conn.go index a61d0dc8..869df0f2 100644 --- a/limiter/conn/conn.go +++ b/limiter/conn/conn.go @@ -1,3 +1,7 @@ +// Package conn implements connection limiters that control the number of +// concurrent connections per key (global, per-IP, or per-CIDR). Limits can +// be loaded from static configuration, files, Redis, or HTTP endpoints with +// periodic hot-reload support. package conn import ( @@ -18,8 +22,12 @@ import ( ) const ( + // GlobalLimitKey is the special key for a global connection limit that applies + // to all connections regardless of their source IP. GlobalLimitKey = "$" - IPLimitKey = "$$" + // IPLimitKey is the special key for a per-IP connection limit that applies + // as a default when no specific IP or CIDR rule matches. + IPLimitKey = "$$" ) type options struct { @@ -33,36 +41,44 @@ type options struct { type Option func(opts *options) +// LimitsOption sets the static limit rules. +// Each rule is a string in the form "key limit", e.g. "192.168.1.0/24 100". func LimitsOption(limits ...string) Option { return func(opts *options) { opts.limits = limits } } +// ReloadPeriodOption sets the period for hot-reloading limit rules from +// external loaders (file, Redis, HTTP). A period <= 0 disables periodic reload. func ReloadPeriodOption(period time.Duration) Option { return func(opts *options) { opts.period = period } } +// FileLoaderOption sets the file loader for reading limit rules from a file. func FileLoaderOption(fileLoader loader.Loader) Option { return func(opts *options) { opts.fileLoader = fileLoader } } +// RedisLoaderOption sets the Redis loader for reading limit rules from Redis. func RedisLoaderOption(redisLoader loader.Loader) Option { return func(opts *options) { opts.redisLoader = redisLoader } } +// HTTPLoaderOption sets the HTTP loader for reading limit rules from an HTTP endpoint. func HTTPLoaderOption(httpLoader loader.Loader) Option { return func(opts *options) { opts.httpLoader = httpLoader } } +// LoggerOption sets the logger for the connection limiter. func LoggerOption(logger logger.Logger) Option { return func(opts *options) { opts.logger = logger @@ -79,6 +95,8 @@ type connLimiter struct { logger logger.Logger } +// NewConnLimiter creates a new ConnLimiter with the given options. +// It starts a background goroutine for periodic reload of limit rules. func NewConnLimiter(opts ...Option) limiter.ConnLimiter { var options options for _, opt := range opts { @@ -146,13 +164,14 @@ func (l *connLimiter) Limiter(key string) limiter.Limiter { } } - var lim limiter.Limiter - if len(lims) > 0 { - lim = newLimiterGroup(lims...) + if len(lims) == 0 { + return nil } + + lim := newLimiterGroup(lims...) l.limits[key] = lim - if lim != nil && l.logger != nil { + if l.logger != nil { l.logger.Debugf("conn limit for %s: %d", key, lim.Limit()) } @@ -233,7 +252,7 @@ func (l *connLimiter) reload(ctx context.Context) error { return nil } -func (l *connLimiter) load(ctx context.Context) (patterns []string, err error) { +func (l *connLimiter) load(ctx context.Context) (patterns []string, _ error) { if l.options.fileLoader != nil { if lister, ok := l.options.fileLoader.(loader.Lister); ok { list, er := lister.List(ctx) @@ -336,6 +355,9 @@ func (l *connLimiter) Close() error { if l.options.redisLoader != nil { l.options.redisLoader.Close() } + if l.options.httpLoader != nil { + l.options.httpLoader.Close() + } return nil } diff --git a/limiter/conn/generator.go b/limiter/conn/generator.go index 07dd4e1a..7755cf6f 100644 --- a/limiter/conn/generator.go +++ b/limiter/conn/generator.go @@ -4,6 +4,8 @@ import ( limiter "github.com/go-gost/core/limiter/conn" ) +// ConnLimitGenerator creates individual Limiter instances for use in +// per-IP or per-CIDR connection limiting. type ConnLimitGenerator interface { Limiter() limiter.Limiter } @@ -12,6 +14,8 @@ type connLimitGenerator struct { n int } +// NewConnLimitGenerator returns a ConnLimitGenerator that creates a new +// Limiter on each call to Limiter(), using the given limit n. func NewConnLimitGenerator(n int) ConnLimitGenerator { return &connLimitGenerator{ n: n, @@ -29,6 +33,8 @@ type connLimitSingleGenerator struct { limiter limiter.Limiter } +// NewConnLimitSingleGenerator returns a ConnLimitGenerator that always +// returns the same Limiter instance, using the given limit n. func NewConnLimitSingleGenerator(n int) ConnLimitGenerator { p := &connLimitSingleGenerator{} if n > 0 { @@ -38,5 +44,8 @@ func NewConnLimitSingleGenerator(n int) ConnLimitGenerator { } func (p *connLimitSingleGenerator) Limiter() limiter.Limiter { + if p == nil { + return nil + } return p.limiter } diff --git a/limiter/conn/limiter.go b/limiter/conn/limiter.go index 5d2f926b..9db9dfdf 100644 --- a/limiter/conn/limiter.go +++ b/limiter/conn/limiter.go @@ -12,6 +12,7 @@ type llimiter struct { current int64 } +// NewLimiter creates a Limiter that allows up to n concurrent connections. func NewLimiter(n int) limiter.Limiter { return &llimiter{limit: n} } @@ -21,10 +22,12 @@ func (l *llimiter) Limit() int { } func (l *llimiter) Allow(n int) bool { + if n < 0 { + atomic.AddInt64(&l.current, int64(n)) + return true + } if atomic.AddInt64(&l.current, int64(n)) > int64(l.limit) { - if n > 0 { - atomic.AddInt64(&l.current, -int64(n)) - } + atomic.AddInt64(&l.current, -int64(n)) return false } return true diff --git a/limiter/conn/wrapper/conn.go b/limiter/conn/wrapper/conn.go index 61b4718c..9cfe089d 100644 --- a/limiter/conn/wrapper/conn.go +++ b/limiter/conn/wrapper/conn.go @@ -1,3 +1,5 @@ +// Package wrapper provides net.Conn and net.Listener wrappers that enforce +// connection limits. package wrapper import ( @@ -21,6 +23,9 @@ type serverConn struct { limiter limiter.Limiter } +// WrapConn wraps a net.Conn with a connection limiter. On Close, the +// limiter's Allow(-1) is called to decrement the current count. If limiter +// is nil, the original connection is returned unchanged. func WrapConn(limiter limiter.Limiter, c net.Conn) net.Conn { if limiter == nil { return c @@ -63,5 +68,5 @@ func (c *serverConn) Context() context.Context { if innerCtx, ok := c.Conn.(ctx.Context); ok { return innerCtx.Context() } - return nil + return context.Background() } diff --git a/limiter/conn/wrapper/listener.go b/limiter/conn/wrapper/listener.go index d77d84a7..105df40b 100644 --- a/limiter/conn/wrapper/listener.go +++ b/limiter/conn/wrapper/listener.go @@ -1,6 +1,7 @@ package wrapper import ( + "errors" "net" limiter "github.com/go-gost/core/limiter/conn" @@ -11,6 +12,10 @@ type listener struct { limiter limiter.ConnLimiter } +// WrapListener wraps a net.Listener with a ConnLimiter. Each accepted +// connection is checked against the limiter using the remote IP as the key. +// If the limit is exceeded, the connection is closed with an error. If +// limiter is nil, the original listener is returned unchanged. func WrapListener(limiter limiter.ConnLimiter, ln net.Listener) net.Listener { if limiter == nil { return ln @@ -34,6 +39,7 @@ func (ln *listener) Accept() (net.Conn, error) { return WrapConn(lim, c), nil } c.Close() + return nil, errors.New("connection limit exceeded") } return c, nil