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:
+28
-6
@@ -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
|
||||
}
|
||||
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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()
|
||||
}
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user