feat: add in flight connection middleware

docs/content/middlewares/tcp/inflightconn.md

@@ -0,0 +1,63 @@
+# InFlightConn
+
+Limiting the Number of Simultaneous connections.
+{: .subtitle }
+
+To proactively prevent services from being overwhelmed with high load, the number of allowed simultaneous connections by IP can be limited.
+
+## Configuration Examples
+
+```yaml tab="Docker"
+labels:
+  - "traefik.tcp.middlewares.test-inflightconn.inflightconn.amount=10"
+```
+
+```yaml tab="Kubernetes"
+apiVersion: traefik.containo.us/v1alpha1
+kind: Middleware
+metadata:
+  name: test-inflightconn
+spec:
+  inFlightConn:
+    amount: 10
+```
+
+```yaml tab="Consul Catalog"
+# Limiting to 10 simultaneous connections
+- "traefik.tcp.middlewares.test-inflightconn.inflightconn.amount=10"
+```
+
+```json tab="Marathon"
+"labels": {
+  "traefik.tcp.middlewares.test-inflightconn.inflightconn.amount": "10"
+}
+```
+
+```yaml tab="Rancher"
+# Limiting to 10 simultaneous connections.
+labels:
+  - "traefik.tcp.middlewares.test-inflightconn.inflightconn.amount=10"
+```
+
+```yaml tab="File (YAML)"
+# Limiting to 10 simultaneous connections.
+tcp:
+  middlewares:
+    test-inflightconn:
+      inFlightConn:
+        amount: 10
+```
+
+```toml tab="File (TOML)"
+# Limiting to 10 simultaneous connections
+[tcp.middlewares]
+  [tcp.middlewares.test-inflightconn.inFlightConn]
+    amount = 10
+```
+
+## Configuration Options
+
+### `amount`
+
+The `amount` option defines the maximum amount of allowed simultaneous connections.
+The middleware closes the connection if there are already `amount` connections opened.

docs/content/middlewares/tcp/overview.md

@@ -131,4 +131,5 @@ tcp:
 
 | Middleware                                | Purpose                                           | Area                        |
 |-------------------------------------------|---------------------------------------------------|-----------------------------|
-| [IPWhiteList](ipwhitelist.md)             | Limit the allowed client IPs                      | Security, Request lifecycle |
+| [InFlightConn](inflightconn.md)           | Limits the number of simultaneous connections.    | Security, Request lifecycle |
+| [IPWhiteList](ipwhitelist.md)             | Limit the allowed client IPs.                     | Security, Request lifecycle |

docs/content/reference/dynamic-configuration/traefik.containo.us_middlewaretcps.yaml

@@ -36,6 +36,13 @@ spec:
           spec:
             description: MiddlewareTCPSpec holds the MiddlewareTCP configuration.
             properties:
+              inFlightConn:
+                description: TCPInFlightConn holds the TCP in flight connection configuration.
+                properties:
+                  amount:
+                    format: int64
+                    type: integer
+                type: object
               ipWhiteList:
                 description: TCPIPWhiteList holds the TCP ip white list configuration.
                 properties:

docs/mkdocs.yml

@@ -132,6 +132,7 @@ nav:
         - 'StripPrefixRegex': 'middlewares/http/stripprefixregex.md'
     - 'TCP':
         - 'Overview': 'middlewares/tcp/overview.md'
+        - 'InFlightConn': 'middlewares/tcp/inflightconn.md'
         - 'IpWhitelist': 'middlewares/tcp/ipwhitelist.md'
   - 'Plugins & Traefik Pilot': 'plugins/index.md'
   - 'Operations':

integration/fixtures/k8s/01-traefik-crd.yml

@@ -1050,6 +1050,13 @@ spec:
           spec:
             description: MiddlewareTCPSpec holds the MiddlewareTCP configuration.
             properties:
+              inFlightConn:
+                description: TCPInFlightConn holds the TCP in flight connection configuration.
+                properties:
+                  amount:
+                    format: int64
+                    type: integer
+                type: object
               ipWhiteList:
                 description: TCPIPWhiteList holds the TCP ip white list configuration.
                 properties:

pkg/config/dynamic/tcp_middlewares.go

@@ -4,11 +4,19 @@ package dynamic
 
 // TCPMiddleware holds the TCPMiddleware configuration.
 type TCPMiddleware struct {
+	InFlightConn *TCPInFlightConn `json:"InFlightConn,omitempty" toml:"InFlightConn,omitempty" yaml:"InFlightConn,omitempty" export:"true"`
 	IPWhiteList  *TCPIPWhiteList  `json:"ipWhiteList,omitempty" toml:"ipWhiteList,omitempty" yaml:"ipWhiteList,omitempty" export:"true"`
 }
 
 // +k8s:deepcopy-gen=true
 
+// TCPInFlightConn holds the TCP in flight connection configuration.
+type TCPInFlightConn struct {
+	Amount int64 `json:"amount,omitempty" toml:"amount,omitempty" yaml:"amount,omitempty" export:"true"`
+}
+
+// +k8s:deepcopy-gen=true
+
 // TCPIPWhiteList holds the TCP ip white list configuration.
 type TCPIPWhiteList struct {
 	SourceRange []string `json:"sourceRange,omitempty" toml:"sourceRange,omitempty" yaml:"sourceRange,omitempty"`

pkg/config/dynamic/zz_generated.deepcopy.go

@@ -1350,9 +1350,30 @@ func (in *TCPIPWhiteList) DeepCopy() *TCPIPWhiteList {
 	return out
 }
 
+// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
+func (in *TCPInFlightConn) DeepCopyInto(out *TCPInFlightConn) {
+	*out = *in
+	return
+}
+
+// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new TCPInFlightConn.
+func (in *TCPInFlightConn) DeepCopy() *TCPInFlightConn {
+	if in == nil {
+		return nil
+	}
+	out := new(TCPInFlightConn)
+	in.DeepCopyInto(out)
+	return out
+}
+
 // DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
 func (in *TCPMiddleware) DeepCopyInto(out *TCPMiddleware) {
 	*out = *in
+	if in.InFlightConn != nil {
+		in, out := &in.InFlightConn, &out.InFlightConn
+		*out = new(TCPInFlightConn)
+		**out = **in
+	}
 	if in.IPWhiteList != nil {
 		in, out := &in.IPWhiteList, &out.IPWhiteList
 		*out = new(TCPIPWhiteList)

pkg/config/label/label_test.go

@@ -174,6 +174,7 @@ func TestDecodeConfiguration(t *testing.T) {
 		"traefik.http.services.Service1.loadbalancer.sticky.cookie.name":               "fui",
 
 		"traefik.tcp.middlewares.Middleware0.ipwhitelist.sourcerange":      "foobar, fiibar",
+		"traefik.tcp.middlewares.Middleware2.inflightconn.amount":          "42",
 		"traefik.tcp.routers.Router0.rule":                                 "foobar",
 		"traefik.tcp.routers.Router0.entrypoints":                          "foobar, fiibar",
 		"traefik.tcp.routers.Router0.service":                              "foobar",
@@ -236,6 +237,11 @@ func TestDecodeConfiguration(t *testing.T) {
 						SourceRange: []string{"foobar", "fiibar"},
 					},
 				},
+				"Middleware2": {
+					InFlightConn: &dynamic.TCPInFlightConn{
+						Amount: 42,
+					},
+				},
 			},
 			Services: map[string]*dynamic.TCPService{
 				"Service0": {
@@ -719,6 +725,11 @@ func TestEncodeConfiguration(t *testing.T) {
 						SourceRange: []string{"foobar", "fiibar"},
 					},
 				},
+				"Middleware2": {
+					InFlightConn: &dynamic.TCPInFlightConn{
+						Amount: 42,
+					},
+				},
 			},
 			Services: map[string]*dynamic.TCPService{
 				"Service0": {
@@ -1320,6 +1331,7 @@ func TestEncodeConfiguration(t *testing.T) {
 		"traefik.HTTP.Services.Service0.LoadBalancer.HealthCheck.Headers.name0":        "foobar",
 
 		"traefik.TCP.Middlewares.Middleware0.IPWhiteList.SourceRange": "foobar, fiibar",
+		"traefik.TCP.Middlewares.Middleware2.InFlightConn.Amount":     "42",
 		"traefik.TCP.Routers.Router0.Rule":                            "foobar",
 		"traefik.TCP.Routers.Router0.EntryPoints":                     "foobar, fiibar",
 		"traefik.TCP.Routers.Router0.Service":                         "foobar",

pkg/middlewares/tcp/inflightconn/inflight_conn.go

@@ -0,0 +1,92 @@
+package tcpinflightconn
+
+import (
+	"context"
+	"fmt"
+	"net"
+	"sync"
+
+	"github.com/traefik/traefik/v2/pkg/config/dynamic"
+	"github.com/traefik/traefik/v2/pkg/log"
+	"github.com/traefik/traefik/v2/pkg/middlewares"
+	"github.com/traefik/traefik/v2/pkg/tcp"
+)
+
+const typeName = "InFlightConnTCP"
+
+type inFlightConn struct {
+	name           string
+	next           tcp.Handler
+	maxConnections int64
+
+	mu          sync.Mutex
+	connections map[string]int64 // current number of connections by remote IP.
+}
+
+// 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.TCPInFlightConn, name string) (tcp.Handler, error) {
+	logger := log.FromContext(middlewares.GetLoggerCtx(ctx, name, typeName))
+	logger.Debug("Creating middleware")
+
+	return &inFlightConn{
+		name:           name,
+		next:           next,
+		connections:    make(map[string]int64),
+		maxConnections: config.Amount,
+	}, nil
+}
+
+// ServeTCP serves the given TCP connection.
+func (i *inFlightConn) ServeTCP(conn tcp.WriteCloser) {
+	ctx := middlewares.GetLoggerCtx(context.Background(), i.name, typeName)
+	logger := log.FromContext(ctx)
+
+	ip, _, err := net.SplitHostPort(conn.RemoteAddr().String())
+	if err != nil {
+		logger.Errorf("Cannot parse IP from remote addr: %v", err)
+		conn.Close()
+		return
+	}
+
+	if err = i.increment(ip); err != nil {
+		logger.Errorf("Connection rejected: %v", err)
+		conn.Close()
+		return
+	}
+
+	defer i.decrement(ip)
+
+	i.next.ServeTCP(conn)
+}
+
+// increment increases the counter for the number of connections tracked for the
+// given IP.
+// It returns an error if the counter would go above the max allowed number of
+// connections.
+func (i *inFlightConn) increment(ip string) error {
+	i.mu.Lock()
+	defer i.mu.Unlock()
+
+	if i.connections[ip] >= i.maxConnections {
+		return fmt.Errorf("max number of connections reached for %s", ip)
+	}
+
+	i.connections[ip]++
+
+	return nil
+}
+
+// decrement decreases the counter for the number of connections tracked for the
+// given IP.
+// It ensures that the counter does not go below zero.
+func (i *inFlightConn) decrement(ip string) {
+	i.mu.Lock()
+	defer i.mu.Unlock()
+
+	if i.connections[ip] <= 0 {
+		return
+	}
+
+	i.connections[ip]--
+}

pkg/middlewares/tcp/inflightconn/inflight_conn_test.go

@@ -0,0 +1,95 @@
+package tcpinflightconn
+
+import (
+	"context"
+	"net"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+	"github.com/traefik/traefik/v2/pkg/config/dynamic"
+	"github.com/traefik/traefik/v2/pkg/tcp"
+)
+
+func TestInFlightConn_ServeTCP(t *testing.T) {
+	proceedCh := make(chan struct{})
+	waitCh := make(chan struct{})
+	finishCh := make(chan struct{})
+
+	next := tcp.HandlerFunc(func(conn tcp.WriteCloser) {
+		proceedCh <- struct{}{}
+
+		if fc, ok := conn.(fakeConn); !ok || !fc.wait {
+			return
+		}
+
+		<-waitCh
+		finishCh <- struct{}{}
+	})
+
+	middleware, err := New(context.Background(), next, dynamic.TCPInFlightConn{Amount: 1}, "foo")
+	require.NoError(t, err)
+
+	// The first connection should succeed and wait.
+	go middleware.ServeTCP(fakeConn{addr: "127.0.0.1:9000", wait: true})
+	requireMessage(t, proceedCh)
+
+	closeCh := make(chan struct{})
+
+	// The second connection from the same remote address should be closed as the maximum number of connections is exceeded.
+	go middleware.ServeTCP(fakeConn{addr: "127.0.0.1:9000", closeCh: closeCh})
+	requireMessage(t, closeCh)
+
+	// The connection from another remote address should succeed.
+	go middleware.ServeTCP(fakeConn{addr: "127.0.0.2:9000"})
+	requireMessage(t, proceedCh)
+
+	// Once the first connection is closed, next connection with the same remote address should succeed.
+	close(waitCh)
+	requireMessage(t, finishCh)
+
+	go middleware.ServeTCP(fakeConn{addr: "127.0.0.1:9000"})
+	requireMessage(t, proceedCh)
+}
+
+func requireMessage(t *testing.T, c chan struct{}) {
+	t.Helper()
+	select {
+	case <-c:
+	case <-time.After(time.Second):
+		t.Fatal("Timeout waiting for message")
+	}
+}
+
+type fakeConn struct {
+	net.Conn
+
+	addr    string
+	wait    bool
+	closeCh chan struct{}
+}
+
+func (c fakeConn) RemoteAddr() net.Addr {
+	return fakeAddr{addr: c.addr}
+}
+
+func (c fakeConn) Close() error {
+	close(c.closeCh)
+	return nil
+}
+
+func (c fakeConn) CloseWrite() error {
+	panic("implement me")
+}
+
+type fakeAddr struct {
+	addr string
+}
+
+func (a fakeAddr) Network() string {
+	return "tcp"
+}
+
+func (a fakeAddr) String() string {
+	return a.addr
+}

pkg/provider/kubernetes/crd/kubernetes.go

@@ -273,6 +273,7 @@ func (p *Provider) loadConfigurationFromCRD(ctx context.Context, client Client)
 		id := provider.Normalize(makeID(middlewareTCP.Namespace, middlewareTCP.Name))
 
 		conf.TCP.Middlewares[id] = &dynamic.TCPMiddleware{
+			InFlightConn: middlewareTCP.Spec.InFlightConn,
 			IPWhiteList:  middlewareTCP.Spec.IPWhiteList,
 		}
 	}

pkg/provider/kubernetes/crd/traefik/v1alpha1/middlewaretcp.go

@@ -20,6 +20,7 @@ type MiddlewareTCP struct {
 
 // MiddlewareTCPSpec holds the MiddlewareTCP configuration.
 type MiddlewareTCPSpec struct {
+	InFlightConn *dynamic.TCPInFlightConn `json:"inFlightConn,omitempty"`
 	IPWhiteList  *dynamic.TCPIPWhiteList  `json:"ipWhiteList,omitempty"`
 }
 

pkg/provider/kubernetes/crd/traefik/v1alpha1/zz_generated.deepcopy.go

@@ -820,6 +820,11 @@ func (in *MiddlewareTCPList) DeepCopyObject() runtime.Object {
 // DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
 func (in *MiddlewareTCPSpec) DeepCopyInto(out *MiddlewareTCPSpec) {
 	*out = *in
+	if in.InFlightConn != nil {
+		in, out := &in.InFlightConn, &out.InFlightConn
+		*out = new(dynamic.TCPInFlightConn)
+		**out = **in
+	}
 	if in.IPWhiteList != nil {
 		in, out := &in.IPWhiteList, &out.IPWhiteList
 		*out = new(dynamic.TCPIPWhiteList)

pkg/server/middleware/tcp/middlewares.go

@@ -6,6 +6,7 @@ import (
 	"strings"
 
 	"github.com/traefik/traefik/v2/pkg/config/runtime"
+	inflightconn "github.com/traefik/traefik/v2/pkg/middlewares/tcp/inflightconn"
 	ipwhitelist "github.com/traefik/traefik/v2/pkg/middlewares/tcp/ipwhitelist"
 	"github.com/traefik/traefik/v2/pkg/server/provider"
 	"github.com/traefik/traefik/v2/pkg/tcp"
@@ -86,6 +87,13 @@ func (b *Builder) buildConstructor(ctx context.Context, middlewareName string) (
 
 	var middleware tcp.Constructor
 
+	// InFlightConn
+	if config.InFlightConn != nil {
+		middleware = func(next tcp.Handler) (tcp.Handler, error) {
+			return inflightconn.New(ctx, next, *config.InFlightConn, middlewareName)
+		}
+	}
+
 	// IPWhiteList
 	if config.IPWhiteList != nil {
 		middleware = func(next tcp.Handler) (tcp.Handler, error) {