[poc] middlewares: tcp: add rate limit

docs/content/reference/dynamic-configuration/docker-labels.yml

@@ -188,9 +188,12 @@
 - "traefik.http.services.service02.loadbalancer.server.port=foobar"
 - "traefik.http.services.service02.loadbalancer.server.scheme=foobar"
 - "traefik.http.services.service02.loadbalancer.server.weight=42"
-- "traefik.tcp.middlewares.tcpmiddleware01.ipallowlist.sourcerange=foobar, foobar"
-- "traefik.tcp.middlewares.tcpmiddleware02.ipwhitelist.sourcerange=foobar, foobar"
-- "traefik.tcp.middlewares.tcpmiddleware03.inflightconn.amount=42"
+- "traefik.tcp.middlewares.tcpmiddleware01.connratelimit.average=42"
+- "traefik.tcp.middlewares.tcpmiddleware01.connratelimit.burst=42"
+- "traefik.tcp.middlewares.tcpmiddleware01.connratelimit.period=42s"
+- "traefik.tcp.middlewares.tcpmiddleware02.ipallowlist.sourcerange=foobar, foobar"
+- "traefik.tcp.middlewares.tcpmiddleware03.ipwhitelist.sourcerange=foobar, foobar"
+- "traefik.tcp.middlewares.tcpmiddleware04.inflightconn.amount=42"
 - "traefik.tcp.routers.tcprouter0.entrypoints=foobar, foobar"
 - "traefik.tcp.routers.tcprouter0.middlewares=foobar, foobar"
 - "traefik.tcp.routers.tcprouter0.priority=42"

docs/content/reference/dynamic-configuration/file.toml

@@ -426,13 +426,18 @@
           weight = 42
   [tcp.middlewares]
     [tcp.middlewares.TCPMiddleware01]
-      [tcp.middlewares.TCPMiddleware01.ipAllowList]
-        sourceRange = ["foobar", "foobar"]
+      [tcp.middlewares.TCPMiddleware01.connRateLimit]
+        average = 42
+        period = "42s"
+        burst = 42
     [tcp.middlewares.TCPMiddleware02]
-      [tcp.middlewares.TCPMiddleware02.ipWhiteList]
+      [tcp.middlewares.TCPMiddleware02.ipAllowList]
         sourceRange = ["foobar", "foobar"]
     [tcp.middlewares.TCPMiddleware03]
-      [tcp.middlewares.TCPMiddleware03.inFlightConn]
+      [tcp.middlewares.TCPMiddleware03.ipWhiteList]
+        sourceRange = ["foobar", "foobar"]
+    [tcp.middlewares.TCPMiddleware04]
+      [tcp.middlewares.TCPMiddleware04.inFlightConn]
         amount = 42
   [tcp.serversTransports]
     [tcp.serversTransports.TCPServersTransport0]

docs/content/reference/dynamic-configuration/file.yaml

@@ -482,16 +482,21 @@ tcp:
             weight: 42
   middlewares:
     TCPMiddleware01:
+      connRateLimit:
+        average: 42
+        period: 42s
+        burst: 42
+    TCPMiddleware02:
       ipAllowList:
         sourceRange:
           - foobar
           - foobar
-    TCPMiddleware02:
+    TCPMiddleware03:
       ipWhiteList:
         sourceRange:
           - foobar
           - foobar
-    TCPMiddleware03:
+    TCPMiddleware04:
       inFlightConn:
         amount: 42
   serversTransports:

docs/content/reference/dynamic-configuration/kv-ref.md

@@ -274,11 +274,14 @@ THIS FILE MUST NOT BE EDITED BY HAND
 | `traefik/http/services/Service04/weighted/sticky/cookie/name` | `foobar` |
 | `traefik/http/services/Service04/weighted/sticky/cookie/sameSite` | `foobar` |
 | `traefik/http/services/Service04/weighted/sticky/cookie/secure` | `true` |
-| `traefik/tcp/middlewares/TCPMiddleware01/ipAllowList/sourceRange/0` | `foobar` |
-| `traefik/tcp/middlewares/TCPMiddleware01/ipAllowList/sourceRange/1` | `foobar` |
-| `traefik/tcp/middlewares/TCPMiddleware02/ipWhiteList/sourceRange/0` | `foobar` |
-| `traefik/tcp/middlewares/TCPMiddleware02/ipWhiteList/sourceRange/1` | `foobar` |
-| `traefik/tcp/middlewares/TCPMiddleware03/inFlightConn/amount` | `42` |
+| `traefik/tcp/middlewares/TCPMiddleware01/connRateLimit/average` | `42` |
+| `traefik/tcp/middlewares/TCPMiddleware01/connRateLimit/burst` | `42` |
+| `traefik/tcp/middlewares/TCPMiddleware01/connRateLimit/period` | `42s` |
+| `traefik/tcp/middlewares/TCPMiddleware02/ipAllowList/sourceRange/0` | `foobar` |
+| `traefik/tcp/middlewares/TCPMiddleware02/ipAllowList/sourceRange/1` | `foobar` |
+| `traefik/tcp/middlewares/TCPMiddleware03/ipWhiteList/sourceRange/0` | `foobar` |
+| `traefik/tcp/middlewares/TCPMiddleware03/ipWhiteList/sourceRange/1` | `foobar` |
+| `traefik/tcp/middlewares/TCPMiddleware04/inFlightConn/amount` | `42` |
 | `traefik/tcp/routers/TCPRouter0/entryPoints/0` | `foobar` |
 | `traefik/tcp/routers/TCPRouter0/entryPoints/1` | `foobar` |
 | `traefik/tcp/routers/TCPRouter0/middlewares/0` | `foobar` |

pkg/config/dynamic/tcp_middlewares.go

@@ -1,10 +1,13 @@
 package dynamic
 
+import ptypes "github.com/traefik/paerser/types"
+
 // +k8s:deepcopy-gen=true
 
 // TCPMiddleware holds the TCPMiddleware configuration.
 type TCPMiddleware struct {
-	InFlightConn *TCPInFlightConn `json:"inFlightConn,omitempty" toml:"inFlightConn,omitempty" yaml:"inFlightConn,omitempty" export:"true"`
+	InFlightConn  *TCPInFlightConn  `json:"inFlightConn,omitempty" toml:"inFlightConn,omitempty" yaml:"inFlightConn,omitempty" export:"true"`
+	ConnRateLimit *TCPConnRateLimit `json:"connRateLimit,omitempty" toml:"connRateLimit,omitempty" yaml:"connRateLimit,omitempty" export:"true"`
 	// Deprecated: please use IPAllowList instead.
 	IPWhiteList *TCPIPWhiteList `json:"ipWhiteList,omitempty" toml:"ipWhiteList,omitempty" yaml:"ipWhiteList,omitempty" export:"true"`
 	IPAllowList *TCPIPAllowList `json:"ipAllowList,omitempty" toml:"ipAllowList,omitempty" yaml:"ipAllowList,omitempty" export:"true"`
@@ -24,6 +27,28 @@ type TCPInFlightConn struct {
 
 // +k8s:deepcopy-gen=true
 
+// TCPConnRateLimit holds the TCP ConnRateLimit middleware configuration.
+// This middleware prevents services from being overwhelmed with high load,
+// by limiting the number of allowed simultaneous connections for one IP.
+// More info: https://doc.traefik.io/traefik/v3.0/middlewares/tcp/inflightconn/
+type TCPConnRateLimit struct {
+	// Average is the maximum rate, by default in requests/s, allowed for the given source.
+	// It defaults to 0, which means no rate limiting.
+	// The rate is actually defined by dividing Average by Period. So for a rate below 1req/s,
+	// one needs to define a Period larger than a second.
+	Average int64 `json:"average,omitempty" toml:"average,omitempty" yaml:"average,omitempty" export:"true"`
+
+	// Period, in combination with Average, defines the actual maximum rate, such as:
+	// r = Average / Period. It defaults to a second.
+	Period ptypes.Duration `json:"period,omitempty" toml:"period,omitempty" yaml:"period,omitempty" export:"true"`
+
+	// Burst is the maximum number of requests allowed to arrive in the same arbitrarily small period of time.
+	// It defaults to 1.
+	Burst int64 `json:"burst,omitempty" toml:"burst,omitempty" yaml:"burst,omitempty" export:"true"`
+}
+
+// +k8s:deepcopy-gen=true
+
 // TCPIPWhiteList holds the TCP IPWhiteList middleware configuration.
 // Deprecated: please use IPAllowList instead.
 type TCPIPWhiteList struct {

pkg/middlewares/tcp/connratelimit/connratelimit.go

@@ -0,0 +1,151 @@
+package connratelimit
+
+import (
+	"context"
+	"fmt"
+	"net"
+	"sync"
+	"time"
+
+	"github.com/mailgun/ttlmap"
+	"github.com/traefik/traefik/v3/pkg/config/dynamic"
+	"github.com/traefik/traefik/v3/pkg/middlewares"
+	"github.com/traefik/traefik/v3/pkg/tcp"
+	"golang.org/x/time/rate"
+)
+
+const (
+	typeName   = "ConnRateLimitTCP"
+	maxSources = 65536
+)
+
+// Mainly derived from HTTP RateLimiter
+type connRateLimitTCP struct {
+	name string
+	next tcp.Handler
+
+	rate  rate.Limit // conns/s
+	burst int64
+
+	// maxDelay is the maximum duration we're willing to wait for a bucket reservation to become effective, in nanoseconds.
+	// For now it is somewhat arbitrarily set to 1/(2*rate).
+	maxDelay time.Duration
+
+	mu      sync.Mutex
+	ttl     int
+	buckets *ttlmap.TtlMap
+}
+
+// New creates a max connections middleware.
+// The connections are identified and grouped by remote IP.
+func New(ctx context.Context, next tcp.Handler, config dynamic.TCPConnRateLimit, name string) (tcp.Handler, error) {
+	logger := middlewares.GetLogger(ctx, name, typeName)
+	logger.Debug().Msg("Creating middleware")
+
+	buckets, err := ttlmap.NewConcurrent(maxSources)
+	if err != nil {
+		return nil, err
+	}
+
+	burst := config.Burst
+	if burst < 1 {
+		burst = 1
+	}
+
+	period := time.Duration(config.Period)
+	if period < 0 {
+		return nil, fmt.Errorf("negative value not valid for period: %v", period)
+	}
+	if period == 0 {
+		period = time.Second
+	}
+
+	// Initialized at rate.Inf to enforce no rate limiting when config.Average == 0
+	rtl := float64(rate.Inf)
+	// No need to set any particular value for maxDelay as the reservation's delay
+	// will be <= 0 in the Inf case (i.e. the average == 0 case).
+	var maxDelay time.Duration
+
+	if config.Average > 0 {
+		rtl = float64(config.Average*int64(time.Second)) / float64(period)
+		// maxDelay does not scale well for rates below 1,
+		// so we just cap it to the corresponding value, i.e. 0.5s, in order to keep the effective rate predictable.
+		// One alternative would be to switch to a no-reservation mode (Allow() method) whenever we are in such a low rate regime.
+		if rtl < 1 {
+			maxDelay = 500 * time.Millisecond
+		} else {
+			maxDelay = time.Second / (time.Duration(rtl) * 2)
+		}
+	}
+
+	// Make the ttl inversely proportional to how often a rate limiter is supposed to see any activity (when maxed out),
+	// for low rate limiters.
+	// Otherwise just make it a second for all the high rate limiters.
+	// Add an extra second in both cases for continuity between the two cases.
+	ttl := 1
+	if rtl >= 1 {
+		ttl++
+	} else if rtl > 0 {
+		ttl += int(1 / rtl)
+	}
+
+	return &connRateLimitTCP{
+		name: name,
+		next: next,
+
+		rate:     rate.Limit(rtl),
+		maxDelay: maxDelay,
+		burst:    burst,
+		ttl:      ttl,
+
+		mu:      sync.Mutex{},
+		buckets: buckets,
+	}, nil
+}
+
+// ServeTCP serves the given TCP connection.
+func (rl *connRateLimitTCP) ServeTCP(conn tcp.WriteCloser) {
+	logger := middlewares.GetLogger(context.Background(), rl.name, typeName)
+
+	ip, port, err := net.SplitHostPort(conn.RemoteAddr().String())
+	if err != nil {
+		logger.Error().Err(err).Msg("Cannot parse IP from remote addr")
+		conn.Close()
+		return
+	}
+
+	var bucket *rate.Limiter
+	if rlSource, exists := rl.buckets.Get(ip); exists {
+		bucket = rlSource.(*rate.Limiter)
+	} else {
+		bucket = rate.NewLimiter(rl.rate, int(rl.burst))
+	}
+
+	// We Set even in the case where the source already exists,
+	// because we want to update the expiryTime everytime we get the source,
+	// as the expiryTime is supposed to reflect the activity (or lack thereof) on that source.
+	if err := rl.buckets.Set(ip, bucket, rl.ttl); err != nil {
+		logger.Error().Err(err).Msg("Could not insert/update bucket")
+		conn.Close()
+		return
+	}
+
+	res := bucket.Reserve()
+	if !res.OK() {
+		logger.Debug().Msgf("Dropper bursty traffic from %s:%s", ip, port)
+		conn.Close()
+		return
+	}
+
+	delay := res.Delay()
+	if delay > rl.maxDelay {
+		res.Cancel()
+		logger.Debug().Msgf("Connection from %s:%s rejected", ip, port)
+		conn.Close()
+		return
+	}
+	logger.Debug().Msgf("Connection from %s:%s accepted", ip, port)
+
+	time.Sleep(delay)
+	rl.next.ServeTCP(conn)
+}

pkg/server/middleware/tcp/middlewares.go

@@ -7,6 +7,7 @@ import (
 
 	"github.com/rs/zerolog/log"
 	"github.com/traefik/traefik/v3/pkg/config/runtime"
+	"github.com/traefik/traefik/v3/pkg/middlewares/tcp/connratelimit"
 	"github.com/traefik/traefik/v3/pkg/middlewares/tcp/inflightconn"
 	"github.com/traefik/traefik/v3/pkg/middlewares/tcp/ipallowlist"
 	"github.com/traefik/traefik/v3/pkg/middlewares/tcp/ipwhitelist"
@@ -96,6 +97,13 @@ func (b *Builder) buildConstructor(ctx context.Context, middlewareName string) (
 		}
 	}
 
+	// ConnRateLimit
+	if config.ConnRateLimit != nil {
+		middleware = func(next tcp.Handler) (tcp.Handler, error) {
+			return connratelimit.New(ctx, next, *config.ConnRateLimit, middlewareName)
+		}
+	}
+
 	// IPWhiteList
 	if config.IPWhiteList != nil {
 		log.Warn().Msg("IPWhiteList is deprecated, please use IPAllowList instead.")