Implement connection limiter

LICENSE.md

@@ -0,0 +1,22 @@
+The MIT License (MIT)
+
+Copyright (c) 2025 idk
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+

Makefile

@@ -0,0 +1,2 @@
+fmt:
+	find . -name '*.go' -exec gofumpt -s -w -extra {} \;

README.md

@@ -0,0 +1,67 @@
+# go-limit
+
+A simple, thread-safe connection limiter for Go's `net.Listener` that manages concurrent connections and rate limiting.
+
+## Install
+
+```bash
+go get github.com/go-i2p/go-limit
+```
+
+## Quick Start
+
+```go
+package main
+
+import (
+    "log"
+    "net"
+    "github.com/go-i2p/go-limit"
+)
+
+func main() {
+    base, _ := net.Listen("tcp", ":8080")
+    
+    listener := limitedlistener.NewLimitedListener(base,
+        limitedlistener.WithMaxConnections(1000),  // max concurrent
+        limitedlistener.WithRateLimit(100))        // per second
+    defer listener.Close()
+
+    for {
+        conn, err := listener.Accept()
+        if err != nil {
+            log.Printf("Accept error: %v", err)
+            continue
+        }
+        go handleConnection(conn)
+    }
+}
+```
+
+## Features
+
+- Limit concurrent connections
+- Rate limiting (connections per second)
+- Connection tracking
+- Real-time statistics
+
+## Configuration
+
+```go
+// Set maximum concurrent connections
+WithMaxConnections(1000)
+
+// Set rate limit (connections per second)
+WithRateLimit(100.0)
+
+// Get current stats
+stats := listener.GetStats()
+```
+
+## License
+
+MIT License
+
+## Support
+
+Issues and PRs welcome at [github.com/go-i2p/go-limit](https://github.com/go-i2p/go-limit)

go.mod

@@ -0,0 +1,5 @@
+module github.com/go-i2p/go-limit
+
+go 1.23.5
+
+require golang.org/x/time v0.9.0

go.sum

@@ -0,0 +1,2 @@
+golang.org/x/time v0.9.0 h1:EsRrnYcQiGH+5FfbgvV4AP7qEZstoyrHB0DzarOQ4ZY=
+golang.org/x/time v0.9.0/go.mod h1:3BpzKBy/shNhVucY/MWOyx10tF3SFh9QdLuxbVysPQM=

limitedListener.go

@@ -0,0 +1,29 @@
+package limitedlistener
+
+import (
+	"fmt"
+	"net"
+	"sync"
+
+	"golang.org/x/time/rate"
+)
+
+// LimitedListener wraps a net.Listener with connection limiting capabilities
+type LimitedListener struct {
+	net.Listener
+	maxConns    int
+	activeConns int64
+	limiter     *rate.Limiter
+	mu          sync.Mutex
+	done        chan struct{}
+	activeSet   map[net.Conn]struct{}
+}
+
+// Option defines the functional option pattern for configuration
+type Option func(*LimitedListener)
+
+// ErrMaxConnsReached is returned when connection limit is reached
+var ErrMaxConnsReached = fmt.Errorf("maximum connections reached")
+
+// ErrRateLimitExceeded is returned when rate limit is exceeded
+var ErrRateLimitExceeded = fmt.Errorf("rate limit exceeded")

listener.go

@@ -0,0 +1,58 @@
+package limitedlistener
+
+import "net"
+
+// NewLimitedListener creates a new connection-limited listener
+func NewLimitedListener(listener net.Listener, opts ...Option) *LimitedListener {
+	l := &LimitedListener{
+		Listener:  listener,
+		maxConns:  1000, // default limit
+		done:      make(chan struct{}),
+		activeSet: make(map[net.Conn]struct{}),
+	}
+
+	for _, opt := range opts {
+		opt(l)
+	}
+
+	return l
+}
+
+// Accept accepts a connection with limiting logic
+func (l *LimitedListener) Accept() (net.Conn, error) {
+	for {
+		// Check if rate limit is exceeded
+		if l.limiter != nil {
+			if !l.limiter.Allow() {
+				return nil, ErrRateLimitExceeded
+			}
+		}
+
+		// Check concurrent connection limit
+		l.mu.Lock()
+		if l.maxConns > 0 && l.activeConns >= int64(l.maxConns) {
+			l.mu.Unlock()
+			return nil, ErrMaxConnsReached
+		}
+		l.mu.Unlock()
+
+		// Accept the connection
+		conn, err := l.Listener.Accept()
+		if err != nil {
+			return nil, err
+		}
+
+		// Wrap the connection for tracking
+		tracked := &trackedConn{
+			Conn:     conn,
+			listener: l,
+		}
+
+		l.mu.Lock()
+		l.activeConns++
+		l.activeSet[tracked] = struct{}{}
+		l.mu.Unlock()
+
+		return tracked, nil
+	}
+}

options.go

@@ -0,0 +1,17 @@
+package limitedlistener
+
+import "golang.org/x/time/rate"
+
+// WithMaxConnections sets the maximum number of concurrent connections
+func WithMaxConnections(max int) Option {
+	return func(l *LimitedListener) {
+		l.maxConns = max
+	}
+}
+
+// WithRateLimit sets connections per second limit
+func WithRateLimit(perSecond float64) Option {
+	return func(l *LimitedListener) {
+		l.limiter = rate.NewLimiter(rate.Limit(perSecond), int(perSecond))
+	}
+}

stats.go

@@ -0,0 +1,39 @@
+package limitedlistener
+
+// Stats provides current listener statistics
+type Stats struct {
+	ActiveConnections int64
+	MaxConnections    int
+	RateLimit         float64
+}
+
+// GetStats returns current listener statistics
+func (l *LimitedListener) GetStats() Stats {
+	l.mu.Lock()
+	defer l.mu.Unlock()
+
+	var rateLimit float64
+	if l.limiter != nil {
+		rateLimit = float64(l.limiter.Limit())
+	}
+
+	return Stats{
+		ActiveConnections: l.activeConns,
+		MaxConnections:    l.maxConns,
+		RateLimit:         rateLimit,
+	}
+}
+
+// Close implements graceful shutdown
+func (l *LimitedListener) Close() error {
+	l.mu.Lock()
+	defer l.mu.Unlock()
+
+	// Close all active connections
+	for conn := range l.activeSet {
+		conn.Close()
+	}
+
+	close(l.done)
+	return l.Listener.Close()
+}

trackedConn.go

@@ -0,0 +1,24 @@
+package limitedlistener
+
+import (
+	"net"
+	"sync"
+)
+
+// trackedConn wraps a net.Conn to track its lifecycle
+type trackedConn struct {
+	net.Conn
+	listener *LimitedListener
+	once     sync.Once
+}
+
+// Close implements net.Conn Close with connection tracking
+func (c *trackedConn) Close() error {
+	c.once.Do(func() {
+		c.listener.mu.Lock()
+		delete(c.listener.activeSet, c)
+		c.listener.activeConns--
+		c.listener.mu.Unlock()
+	})
+	return c.Conn.Close()
+}