feat: add sticky-session support in port-forwarding

Author: Revolyssup

README.md

@@ -81,15 +81,14 @@ First time the k8sdebug command is called this file will be created.
 
 
 ### 🔄 Smart Port Forwarding
-
+![arch](./archport.png)
 ```bash
 k8sdebug port-forward \
   -n my-namespace \
   --labels "app=backend" \
   --policy round-robin \
   --hostport 8080 \
   --containerport 80
-
 # Supported policies:
 # - round-robin (default)
 

archport.png

@@ -1,5 +1,9 @@
 package forwarder
 
+import (
+	"net"
+)
+
 type Forwarder interface {
-	NextPort() string
+	NextPort(net.Conn) string
 }

pkg/forwarder/forwarder.go

@@ -1,6 +1,10 @@
 package mock
 
-import "github.com/revolyssup/k8sdebug/pkg/forwarder"
+import (
+	"net"
+
+	"github.com/revolyssup/k8sdebug/pkg/forwarder"
+)
 
 type MockForwarder struct {
 	ports []string
@@ -15,7 +19,7 @@ func New(ports ...string) forwarder.Forwarder {
 }
 
 // Port returns the next predefined port or an error if exhausted.
-func (m *MockForwarder) NextPort() string {
+func (m *MockForwarder) NextPort(_ net.Conn) string {
 	port := m.ports[m.index]
 	m.index = (m.index + 1) % len(m.ports)
 	return port

pkg/portforward/mock/mock.go

@@ -10,10 +10,10 @@ import (
 func TestRoundRobin(t *testing.T) {
 	mock := mock.New("8080", "8081", "8082")
 
-	assert.Equal(t, "8080", mock.NextPort())
-	assert.Equal(t, "8081", mock.NextPort())
-	assert.Equal(t, "8082", mock.NextPort())
-	assert.Equal(t, "8080", mock.NextPort())
-	assert.Equal(t, "8081", mock.NextPort())
-	assert.Equal(t, "8082", mock.NextPort())
+	assert.Equal(t, "8080", mock.NextPort(nil))
+	assert.Equal(t, "8081", mock.NextPort(nil))
+	assert.Equal(t, "8082", mock.NextPort(nil))
+	assert.Equal(t, "8080", mock.NextPort(nil))
+	assert.Equal(t, "8081", mock.NextPort(nil))
+	assert.Equal(t, "8082", mock.NextPort(nil))
 }

pkg/portforward/mock/mock_test.go

@@ -17,6 +17,7 @@ import (
 	"github.com/revolyssup/k8sdebug/pkg/forwarder"
 	"github.com/revolyssup/k8sdebug/pkg/portforward/mock"
 	"github.com/revolyssup/k8sdebug/pkg/portforward/roundrobin"
+	"github.com/revolyssup/k8sdebug/pkg/portforward/sticky"
 	"github.com/spf13/cobra"
 	v1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
@@ -48,8 +49,8 @@ func forwardToPod(hostConn net.Conn, podCon net.Conn) {
 	io.Copy(podCon, hostConn)
 }
 
-func getPodConnection(fw forwarder.Forwarder) (net.Conn, error) {
-	port := fw.NextPort()
+func getPodConnection(fw forwarder.Forwarder, hostConn net.Conn) (net.Conn, error) {
+	port := fw.NextPort(hostConn)
 	if port == "" {
 		return nil, fmt.Errorf("no available port")
 	}
@@ -66,6 +67,8 @@ func getForwarder(policy string) forwarder.Forwarder {
 		return roundrobin.New(&connPool)
 	case "mock":
 		return mock.New()
+	case "sticky":
+		return sticky.New(&connPool)
 	}
 	return nil
 }
@@ -107,7 +110,7 @@ func listenAndAccept(ctx context.Context, listener net.Listener, fw forwarder.Fo
 				fmt.Printf("Accept error: %v", err)
 				continue
 			}
-			podConn, err := getPodConnection(fw)
+			podConn, err := getPodConnection(fw, hostConn)
 			if err != nil {
 				fmt.Printf("Pod connection error: %v", err)
 				continue
@@ -247,7 +250,11 @@ func NewCommand() *cobra.Command {
 
 	cmd.PersistentFlags().StringVarP(&namespace, "namespace", "n", "default", "Name of the pod")
 	cmd.PersistentFlags().StringVarP(&typ, "type", "t", "pod", "Name of the pod")
-	cmd.PersistentFlags().StringVar(&policy, "policy", "round-robin", "policy to use while sending requests")
+	cmd.PersistentFlags().StringVar(&policy, "policy", "round-robin", `policy to use while sending requests.
+Default policy is round-robin. 
+round-robin: In this mode, requests are balanced across all the pods much like Kubernetes service.
+sticky: In this mode, requests from a particular source IP will always be directed to a single pod.
+`)
 	cmd.Flags().StringVarP(&labels, "labels", "l", "", "list of key value pairs to use as labels while filtering pods.")
 	cmd.Flags().StringVar(&hostport, "hostport", "3000", "host port on which requests will be sent")
 	cmd.Flags().StringVar(&containerPort, "containerport", "80", "container port on which requests will be sent")

pkg/portforward/portforward.go

@@ -19,7 +19,7 @@ func New(connPool *[]string) forwarder.Forwarder {
 		connPool: connPool,
 	}
 }
-func (rr *RoundRobin) NextPort() string {
+func (rr *RoundRobin) NextPort(_ net.Conn) string {
 	rr.mx.Lock()
 	defer rr.mx.Unlock()
 

pkg/portforward/roundrobin/roundrobin.go

@@ -0,0 +1,40 @@
+package sticky
+
+import (
+	"net"
+	"sync"
+
+	"github.com/revolyssup/k8sdebug/pkg/forwarder"
+	"github.com/revolyssup/k8sdebug/pkg/portforward/roundrobin"
+)
+
+type StickySession struct {
+	ipTracker  map[string]string // Source IP to Allotted port mapping
+	mx         sync.Mutex
+	connPool   *[]string
+	roundRobin forwarder.Forwarder
+}
+
+func New(connPool *[]string) forwarder.Forwarder {
+	return &StickySession{
+		connPool:   connPool,
+		ipTracker:  make(map[string]string),
+		roundRobin: roundrobin.New(connPool),
+	}
+}
+func (ss *StickySession) NextPort(conn net.Conn) string {
+	ss.mx.Lock()
+	defer ss.mx.Unlock()
+
+	remoteAddr := conn.RemoteAddr().String()
+
+	// Split into IP and port
+	host, _, _ := net.SplitHostPort(remoteAddr)
+	if ss.ipTracker[host] != "" {
+		return ss.ipTracker[host]
+	}
+	//fallback to roundrobin
+	port := ss.roundRobin.NextPort(conn)
+	ss.ipTracker[host] = port
+	return port
+}