docs(limiter/conn): add doc comments, fix nil guards and missing Close

Add package and exported symbol doc comments across conn, generator,
limiter, and wrapper packages. Fix nil receiver guard in
connLimitSingleGenerator.Limiter, add missing httpLoader.Close call,
return context.Background instead of nil in serverConn.Context, return
explicit error on connection limit exceeded, and simplify Allow logic.
This commit is contained in:
ginuerzh
2026-05-24 19:38:07 +08:00
parent 0e96f602fa
commit fa708f4b5f
5 changed files with 55 additions and 10 deletions
+27 -5
View File
@@ -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 package conn
import ( import (
@@ -18,7 +22,11 @@ import (
) )
const ( const (
// GlobalLimitKey is the special key for a global connection limit that applies
// to all connections regardless of their source IP.
GlobalLimitKey = "$" GlobalLimitKey = "$"
// 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 = "$$" IPLimitKey = "$$"
) )
@@ -33,36 +41,44 @@ type options struct {
type Option func(opts *options) 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 { func LimitsOption(limits ...string) Option {
return func(opts *options) { return func(opts *options) {
opts.limits = limits 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 { func ReloadPeriodOption(period time.Duration) Option {
return func(opts *options) { return func(opts *options) {
opts.period = period opts.period = period
} }
} }
// FileLoaderOption sets the file loader for reading limit rules from a file.
func FileLoaderOption(fileLoader loader.Loader) Option { func FileLoaderOption(fileLoader loader.Loader) Option {
return func(opts *options) { return func(opts *options) {
opts.fileLoader = fileLoader opts.fileLoader = fileLoader
} }
} }
// RedisLoaderOption sets the Redis loader for reading limit rules from Redis.
func RedisLoaderOption(redisLoader loader.Loader) Option { func RedisLoaderOption(redisLoader loader.Loader) Option {
return func(opts *options) { return func(opts *options) {
opts.redisLoader = redisLoader opts.redisLoader = redisLoader
} }
} }
// HTTPLoaderOption sets the HTTP loader for reading limit rules from an HTTP endpoint.
func HTTPLoaderOption(httpLoader loader.Loader) Option { func HTTPLoaderOption(httpLoader loader.Loader) Option {
return func(opts *options) { return func(opts *options) {
opts.httpLoader = httpLoader opts.httpLoader = httpLoader
} }
} }
// LoggerOption sets the logger for the connection limiter.
func LoggerOption(logger logger.Logger) Option { func LoggerOption(logger logger.Logger) Option {
return func(opts *options) { return func(opts *options) {
opts.logger = logger opts.logger = logger
@@ -79,6 +95,8 @@ type connLimiter struct {
logger logger.Logger 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 { func NewConnLimiter(opts ...Option) limiter.ConnLimiter {
var options options var options options
for _, opt := range opts { for _, opt := range opts {
@@ -146,13 +164,14 @@ func (l *connLimiter) Limiter(key string) limiter.Limiter {
} }
} }
var lim limiter.Limiter if len(lims) == 0 {
if len(lims) > 0 { return nil
lim = newLimiterGroup(lims...)
} }
lim := newLimiterGroup(lims...)
l.limits[key] = lim 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()) 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 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 l.options.fileLoader != nil {
if lister, ok := l.options.fileLoader.(loader.Lister); ok { if lister, ok := l.options.fileLoader.(loader.Lister); ok {
list, er := lister.List(ctx) list, er := lister.List(ctx)
@@ -336,6 +355,9 @@ func (l *connLimiter) Close() error {
if l.options.redisLoader != nil { if l.options.redisLoader != nil {
l.options.redisLoader.Close() l.options.redisLoader.Close()
} }
if l.options.httpLoader != nil {
l.options.httpLoader.Close()
}
return nil return nil
} }
+9
View File
@@ -4,6 +4,8 @@ import (
limiter "github.com/go-gost/core/limiter/conn" 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 { type ConnLimitGenerator interface {
Limiter() limiter.Limiter Limiter() limiter.Limiter
} }
@@ -12,6 +14,8 @@ type connLimitGenerator struct {
n int 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 { func NewConnLimitGenerator(n int) ConnLimitGenerator {
return &connLimitGenerator{ return &connLimitGenerator{
n: n, n: n,
@@ -29,6 +33,8 @@ type connLimitSingleGenerator struct {
limiter limiter.Limiter limiter limiter.Limiter
} }
// NewConnLimitSingleGenerator returns a ConnLimitGenerator that always
// returns the same Limiter instance, using the given limit n.
func NewConnLimitSingleGenerator(n int) ConnLimitGenerator { func NewConnLimitSingleGenerator(n int) ConnLimitGenerator {
p := &connLimitSingleGenerator{} p := &connLimitSingleGenerator{}
if n > 0 { if n > 0 {
@@ -38,5 +44,8 @@ func NewConnLimitSingleGenerator(n int) ConnLimitGenerator {
} }
func (p *connLimitSingleGenerator) Limiter() limiter.Limiter { func (p *connLimitSingleGenerator) Limiter() limiter.Limiter {
if p == nil {
return nil
}
return p.limiter return p.limiter
} }
+6 -3
View File
@@ -12,6 +12,7 @@ type llimiter struct {
current int64 current int64
} }
// NewLimiter creates a Limiter that allows up to n concurrent connections.
func NewLimiter(n int) limiter.Limiter { func NewLimiter(n int) limiter.Limiter {
return &llimiter{limit: n} return &llimiter{limit: n}
} }
@@ -21,10 +22,12 @@ func (l *llimiter) Limit() int {
} }
func (l *llimiter) Allow(n int) bool { func (l *llimiter) Allow(n int) bool {
if atomic.AddInt64(&l.current, int64(n)) > int64(l.limit) { if n < 0 {
if n > 0 { atomic.AddInt64(&l.current, int64(n))
atomic.AddInt64(&l.current, -int64(n)) return true
} }
if atomic.AddInt64(&l.current, int64(n)) > int64(l.limit) {
atomic.AddInt64(&l.current, -int64(n))
return false return false
} }
return true return true
+6 -1
View File
@@ -1,3 +1,5 @@
// Package wrapper provides net.Conn and net.Listener wrappers that enforce
// connection limits.
package wrapper package wrapper
import ( import (
@@ -21,6 +23,9 @@ type serverConn struct {
limiter limiter.Limiter 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 { func WrapConn(limiter limiter.Limiter, c net.Conn) net.Conn {
if limiter == nil { if limiter == nil {
return c return c
@@ -63,5 +68,5 @@ func (c *serverConn) Context() context.Context {
if innerCtx, ok := c.Conn.(ctx.Context); ok { if innerCtx, ok := c.Conn.(ctx.Context); ok {
return innerCtx.Context() return innerCtx.Context()
} }
return nil return context.Background()
} }
+6
View File
@@ -1,6 +1,7 @@
package wrapper package wrapper
import ( import (
"errors"
"net" "net"
limiter "github.com/go-gost/core/limiter/conn" limiter "github.com/go-gost/core/limiter/conn"
@@ -11,6 +12,10 @@ type listener struct {
limiter limiter.ConnLimiter 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 { func WrapListener(limiter limiter.ConnLimiter, ln net.Listener) net.Listener {
if limiter == nil { if limiter == nil {
return ln return ln
@@ -34,6 +39,7 @@ func (ln *listener) Accept() (net.Conn, error) {
return WrapConn(lim, c), nil return WrapConn(lim, c), nil
} }
c.Close() c.Close()
return nil, errors.New("connection limit exceeded")
} }
return c, nil return c, nil