Template dynamicAdvertisedListener, allow {{.brokerId}} for a dynamic advertited hostname based on brokerId. Use fixed port if provided.

Author: everesio

Diff for: proxy/proxy.go

@@ -1,11 +1,13 @@
 package proxy
 
 import (
+	"bytes"
 	"crypto/tls"
 	"fmt"
 	"net"
 	"strconv"
 	"sync"
+	"text/template"
 
 	"github.com/grepplabs/kafka-proxy/config"
 	"github.com/grepplabs/kafka-proxy/pkg/libs/util"
@@ -190,17 +192,57 @@ func (p *Listeners) ListenDynamicInstance(brokerAddress string, brokerId int32)
 	port := l.Addr().(*net.TCPAddr).Port
 	address := net.JoinHostPort(p.defaultListenerIP, fmt.Sprint(port))
 
-	dynamicAdvertisedListener := p.dynamicAdvertisedListener
-	if dynamicAdvertisedListener == "" {
-		dynamicAdvertisedListener = p.defaultListenerIP
+	dynamicAdvertisedHost, dynamicAdvertisedPort, err := p.getDynamicAdvertisedAddress(cfg.BrokerID, port)
+	if err != nil {
+		return "", 0, err
 	}
-	cfg.AdvertisedAddress = net.JoinHostPort(dynamicAdvertisedListener, fmt.Sprint(port))
+	cfg.AdvertisedAddress = net.JoinHostPort(dynamicAdvertisedHost, fmt.Sprint(dynamicAdvertisedPort))
 	cfg.ListenerAddress = address
 
 	p.brokerToListenerConfig[brokerAddress] = cfg
 	logrus.Infof("Dynamic listener %s for broker %s brokerId %d advertised as %s", cfg.ListenerAddress, cfg.GetBrokerAddress(), cfg.BrokerID, cfg.AdvertisedAddress)
 
-	return dynamicAdvertisedListener, int32(port), nil
+	return dynamicAdvertisedHost, int32(dynamicAdvertisedPort), nil
+}
+
+func (p *Listeners) getDynamicAdvertisedAddress(brokerID int32, port int) (string, int, error) {
+	dynamicAdvertisedListener := p.dynamicAdvertisedListener
+	if dynamicAdvertisedListener == "" {
+		return p.defaultListenerIP, port, nil
+	}
+	dynamicAdvertisedListener, err := p.templateDynamicAdvertisedAddress(brokerID)
+	if err != nil {
+		return "", 0, err
+	}
+	var (
+		dynamicAdvertisedHost = dynamicAdvertisedListener
+		dynamicAdvertisedPort = port
+	)
+	advHost, advPortStr, err := net.SplitHostPort(dynamicAdvertisedListener)
+	if err == nil {
+		if advPort, err := strconv.Atoi(advPortStr); err == nil {
+			dynamicAdvertisedHost = advHost
+			dynamicAdvertisedPort = advPort
+		}
+	}
+	return dynamicAdvertisedHost, dynamicAdvertisedPort, nil
+}
+
+func (p *Listeners) templateDynamicAdvertisedAddress(brokerID int32) (string, error) {
+	tmpl, err := template.New("dynamicAdvertisedHost").Option("missingkey=error").Parse(p.dynamicAdvertisedListener)
+	if err != nil {
+		return "", fmt.Errorf("failed to parse host template '%s': %w", p.dynamicAdvertisedListener, err)
+	}
+	var buf bytes.Buffer
+	data := map[string]any{
+		"brokerId": brokerID,
+		"brokerID": brokerID,
+	}
+	err = tmpl.Execute(&buf, data)
+	if err != nil {
+		return "", fmt.Errorf("failed to execute host template '%s': %w", p.dynamicAdvertisedListener, err)
+	}
+	return buf.String(), nil
 }
 
 func (p *Listeners) ListenInstances(cfgs []config.ListenerConfig) (<-chan Conn, error) {

Diff for: proxy/proxy_test.go

@@ -283,3 +283,135 @@ func TestGetBrokerToListenerConfig(t *testing.T) {
 		assert.ObjectsAreEqual(tt.mapping, mapping)
 	}
 }
+
+func TestGetDynamicAdvertisedAddress(t *testing.T) {
+	tests := []struct {
+		name                      string
+		dynamicAdvertisedListener string
+		defaultListenerIP         string
+		brokerID                  int32
+		port                      int
+		expectedHost              string
+		expectedPort              int
+		expectError               bool
+	}{
+		{
+			name:                      "Default listener IP is 127.0.0.1",
+			dynamicAdvertisedListener: "",
+			defaultListenerIP:         "127.0.0.1",
+			brokerID:                  1,
+			port:                      9092,
+			expectedHost:              "127.0.0.1",
+			expectedPort:              9092,
+			expectError:               false,
+		},
+		{
+			name:                      "Default listener IP is 0.0.0.0",
+			dynamicAdvertisedListener: "",
+			defaultListenerIP:         "0.0.0.0",
+			brokerID:                  1,
+			port:                      9092,
+			expectedHost:              "0.0.0.0",
+			expectedPort:              9092,
+			expectError:               false,
+		},
+		{
+			name:                      "Default listener IP is localhost",
+			dynamicAdvertisedListener: "",
+			defaultListenerIP:         "localhost",
+			brokerID:                  1,
+			port:                      9092,
+			expectedHost:              "localhost",
+			expectedPort:              9092,
+			expectError:               false,
+		},
+		{
+			name:                      "Dynamic listener no template, host is IP",
+			dynamicAdvertisedListener: "0.0.0.0",
+			defaultListenerIP:         "127.0.0.1",
+			brokerID:                  2,
+			port:                      9093,
+			expectedHost:              "0.0.0.0",
+			expectedPort:              9093,
+			expectError:               false,
+		},
+		{
+			name:                      "Dynamic listener no template, host is IP and port is provided",
+			dynamicAdvertisedListener: "0.0.0.0:30000",
+			defaultListenerIP:         "127.0.0.1",
+			brokerID:                  2,
+			port:                      9093,
+			expectedHost:              "0.0.0.0",
+			expectedPort:              30000,
+			expectError:               false,
+		},
+		{
+			name:                      "Dynamic listener no template, host is dns name",
+			dynamicAdvertisedListener: "kafka-proxy.provisionedmskclust.zgjvgc.c2.kafka.eu-central-1.amazonaws.com",
+			defaultListenerIP:         "127.0.0.1",
+			brokerID:                  2,
+			port:                      9093,
+			expectedHost:              "kafka-proxy.provisionedmskclust.zgjvgc.c2.kafka.eu-central-1.amazonaws.com",
+			expectedPort:              9093,
+			expectError:               false,
+		},
+		{
+			name:                      "Dynamic listener no template, host is dns name and port is provided",
+			dynamicAdvertisedListener: "kafka-proxy.grepplabs.com:30000",
+			defaultListenerIP:         "127.0.0.1",
+			brokerID:                  2,
+			port:                      9093,
+			expectedHost:              "kafka-proxy.grepplabs.com",
+			expectedPort:              30000,
+			expectError:               false,
+		},
+		{
+			name:                      "Dynamic listener with template",
+			dynamicAdvertisedListener: "b-{{.brokerId}}.provisionedmskclust.zgjvgc.c2.kafka.eu-central-1.amazonaws.com",
+			defaultListenerIP:         "127.0.0.1",
+			brokerID:                  2,
+			port:                      9093,
+			expectedHost:              "b-2.provisionedmskclust.zgjvgc.c2.kafka.eu-central-1.amazonaws.com",
+			expectedPort:              9093,
+			expectError:               false,
+		},
+		{
+			name:                      "Dynamic listener with template and port is provided",
+			dynamicAdvertisedListener: "b-{{.brokerId}}.provisionedmskclust.zgjvgc.c2.kafka.eu-central-1.amazonaws.com:30000",
+			defaultListenerIP:         "127.0.0.1",
+			brokerID:                  2,
+			port:                      9093,
+			expectedHost:              "b-2.provisionedmskclust.zgjvgc.c2.kafka.eu-central-1.amazonaws.com",
+			expectedPort:              30000,
+			expectError:               false,
+		},
+		{
+			name:                      "Invalid dynamic listener template",
+			dynamicAdvertisedListener: "broker-{{.invalid}}",
+			defaultListenerIP:         "127.0.0.1",
+			brokerID:                  3,
+			port:                      9094,
+			expectedHost:              "",
+			expectedPort:              0,
+			expectError:               true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			listeners := &Listeners{
+				dynamicAdvertisedListener: tt.dynamicAdvertisedListener,
+				defaultListenerIP:         tt.defaultListenerIP,
+			}
+
+			host, port, err := listeners.getDynamicAdvertisedAddress(tt.brokerID, tt.port)
+			if tt.expectError {
+				assert.Error(t, err)
+			} else {
+				assert.NoError(t, err)
+				assert.Equal(t, tt.expectedHost, host)
+				assert.Equal(t, tt.expectedPort, port)
+			}
+		})
+	}
+}