Implement dial address mapping

Author: everesio

README.md

@@ -83,6 +83,7 @@ See:
           --debug-enable                                   Enable Debug endpoint
           --debug-listen-address string                    Debug listen address (default "0.0.0.0:6060")
           --default-listener-ip string                     Default listener IP (default "127.0.0.1")
+          --dial-address-mapping stringArray               Mapping of target broker address to new one (host:port,host:port). The mapping is performed during connection establishment
           --dynamic-listeners-disable                      Disable dynamic listeners.
           --external-server-mapping stringArray            Mapping of Kafka server address to external address (host:port,host:port). A listener for the external address is not started
           --forbidden-api-keys intSlice                    Forbidden Kafka request types. The restriction should prevent some Kafka operations e.g. 20 - DeleteTopics
@@ -378,7 +379,7 @@ spec:
           secretName: tls-client-key-file
 ```
 
-### Connect to Kafka running in Kubernetes example
+### Connect to Kafka running in Kubernetes example (kafka proxy runs in cluster)
 
 ```yaml
 
@@ -484,9 +485,29 @@ spec:
 kubectl port-forward kafka-proxy-0 32400:32400 32401:32401 32402:32402
 ```
 
-Use localhost:32400, localhost:32401 and localhost:32402 as boostrap servers
+Use localhost:32400, localhost:32401 and localhost:32402 as bootstrap servers
 
 
+### Connect to Kafka running in Kubernetes example (kafka proxy runs locally)
+
+kafka.properties of one node Kafka
+
+```
+broker.id=0
+advertised.listeners=PLAINTEXT://kafka-0.kafka-headless.kafka:9092
+...
+```
+
+```bash
+kubectl port-forward -n kafka kafka-0 9092:9092
+```
+
+```bash
+kafka-proxy server --bootstrap-server-mapping "127.0.0.1:9092,0.0.0.0:19092" --dial-address-mapping "kafka-0.kafka-headless.kafka:9092,0.0.0.0:9092"
+```
+
+Use localhost:19092 as bootstrap servers
+
 ### Embedded third-party source code 
 
 * [Cloud SQL Proxy](https://github.com/GoogleCloudPlatform/cloudsql-proxy)

cmd/kafka-proxy/server.go

@@ -40,6 +40,7 @@ var (
 
 	bootstrapServersMapping = make([]string, 0)
 	externalServersMapping  = make([]string, 0)
+	dialAddressMapping      = make([]string, 0)
 )
 
 var Server = &cobra.Command{
@@ -57,6 +58,9 @@ var Server = &cobra.Command{
 		if err := c.InitExternalServers(getOrEnvStringSlice(externalServersMapping, "EXTERNAL_SERVER_MAPPING")); err != nil {
 			return err
 		}
+		if err := c.InitDialAddressMappings(getOrEnvStringSlice(dialAddressMapping, "DIAL_ADDRESS_MAPPING")); err != nil {
+			return err
+		}
 		if err := c.Validate(); err != nil {
 			return err
 		}
@@ -81,6 +85,7 @@ func initFlags() {
 	Server.Flags().StringVar(&c.Proxy.DefaultListenerIP, "default-listener-ip", "127.0.0.1", "Default listener IP")
 	Server.Flags().StringArrayVar(&bootstrapServersMapping, "bootstrap-server-mapping", []string{}, "Mapping of Kafka bootstrap server address to local address (host:port,host:port(,advhost:advport))")
 	Server.Flags().StringArrayVar(&externalServersMapping, "external-server-mapping", []string{}, "Mapping of Kafka server address to external address (host:port,host:port). A listener for the external address is not started")
+	Server.Flags().StringArrayVar(&dialAddressMapping, "dial-address-mapping", []string{}, "Mapping of target broker address to new one (host:port,host:port). The mapping is performed during connection establishment")
 	Server.Flags().BoolVar(&c.Proxy.DisableDynamicListeners, "dynamic-listeners-disable", false, "Disable dynamic listeners.")
 
 	Server.Flags().IntVar(&c.Proxy.RequestBufferSize, "proxy-request-buffer-size", 4096, "Request buffer size pro tcp connection")
@@ -173,7 +178,6 @@ func initFlags() {
 	Server.Flags().StringVar(&c.Log.Level, "log-level", "info", "Log level debug, info, warning, error, fatal or panic")
 	Server.Flags().StringVar(&c.Log.LevelFieldName, "log-level-fieldname", "@level", "Log level fieldname for json format")
 
-
 	// Connect through Socks5 or HTTP CONNECT to Kafka
 	Server.Flags().StringVar(&c.ForwardProxy.Url, "forward-proxy", "", "URL of the forward proxy. Supported schemas are socks5 and http")
 

cmd/kafka-proxy/server_test.go

@@ -11,8 +11,9 @@ func setupBootstrapServersMappingTest() {
 	Server.ResetFlags()
 	c = new(config.Config)
 	initFlags()
-	os.Setenv("BOOTSTRAP_SERVER_MAPPING", "")
-	os.Setenv("EXTERNAL_SERVER_MAPPING", "")
+	_ = os.Setenv("BOOTSTRAP_SERVER_MAPPING", "")
+	_ = os.Setenv("EXTERNAL_SERVER_MAPPING", "")
+	_ = os.Setenv("DIAL_ADDRESS_MAPPING", "")
 }
 
 func TestBootstrapServersMappingFromFlags(t *testing.T) {
@@ -24,7 +25,7 @@ func TestBootstrapServersMappingFromFlags(t *testing.T) {
 		"--bootstrap-server-mapping", "kafka-2.example.com:9092,0.0.0.0:32403,kafka-2.grepplabs.com:9092",
 	}
 
-	Server.ParseFlags(args)
+	_ = Server.ParseFlags(args)
 	err := Server.PreRunE(nil, args)
 	a := assert.New(t)
 	a.Nil(err)
@@ -44,13 +45,34 @@ func TestBootstrapServersMappingFromFlags(t *testing.T) {
 
 }
 
+func TestDialMappingFromFlags(t *testing.T) {
+	setupBootstrapServersMappingTest()
+
+	args := []string{"cobra.test",
+		"--bootstrap-server-mapping", "192.168.99.100:32401,0.0.0.0:32401",
+		"--dial-address-mapping", "service-kafka-0.service-kafka-headless.service:9092,0.0.0.0:19092",
+		"--dial-address-mapping", "192.168.99.100:32402,0.0.0.0:32402",
+	}
+
+	_ = Server.ParseFlags(args)
+	err := Server.PreRunE(nil, args)
+	a := assert.New(t)
+	a.Nil(err)
+	a.Len(c.Proxy.DialAddressMappings, 2)
+
+	a.Equal(c.Proxy.DialAddressMappings[0].SourceAddress, "service-kafka-0.service-kafka-headless.service:9092")
+	a.Equal(c.Proxy.DialAddressMappings[0].DestinationAddress, "0.0.0.0:19092")
+
+	a.Equal(c.Proxy.DialAddressMappings[1].SourceAddress, "192.168.99.100:32402")
+	a.Equal(c.Proxy.DialAddressMappings[1].DestinationAddress, "0.0.0.0:32402")
+}
 func TestBootstrapServersMappingFromEnv(t *testing.T) {
 	setupBootstrapServersMappingTest()
 
-	os.Setenv("BOOTSTRAP_SERVER_MAPPING", "192.168.99.100:32404,0.0.0.0:32404 kafka-5.example.com:9092,0.0.0.0:32405,kafka-5.grepplabs.com:9092")
+	_ = os.Setenv("BOOTSTRAP_SERVER_MAPPING", "192.168.99.100:32404,0.0.0.0:32404 kafka-5.example.com:9092,0.0.0.0:32405,kafka-5.grepplabs.com:9092")
 
 	var args []string
-	Server.ParseFlags(args)
+	_ = Server.ParseFlags(args)
 	err := Server.PreRunE(nil, args)
 	a := assert.New(t)
 	a.Nil(err)
@@ -70,7 +92,7 @@ func TestEmptyBootstrapServersMapping(t *testing.T) {
 	setupBootstrapServersMappingTest()
 
 	var args []string
-	Server.ParseFlags(args)
+	_ = Server.ParseFlags(args)
 	err := Server.PreRunE(nil, args)
 	a := assert.New(t)
 	a.Error(err, "list of bootstrap-server-mapping must not be empty")
@@ -79,10 +101,10 @@ func TestEmptyBootstrapServersMapping(t *testing.T) {
 func TestBootstrapServersMappingFromEnvWithWhiteSpaces(t *testing.T) {
 	setupBootstrapServersMappingTest()
 
-	os.Setenv("BOOTSTRAP_SERVER_MAPPING", "   192.168.99.100:32404,0.0.0.0:32404   kafka-5.example.com:9092,0.0.0.0:32405,kafka-5.grepplabs.com:9092    ")
+	_ = os.Setenv("BOOTSTRAP_SERVER_MAPPING", "   192.168.99.100:32404,0.0.0.0:32404   kafka-5.example.com:9092,0.0.0.0:32405,kafka-5.grepplabs.com:9092    ")
 
 	var args []string
-	Server.ParseFlags(args)
+	_ = Server.ParseFlags(args)
 	err := Server.PreRunE(nil, args)
 	a := assert.New(t)
 	a.Nil(err)
@@ -101,11 +123,11 @@ func TestBootstrapServersMappingFromEnvWithWhiteSpaces(t *testing.T) {
 func TestExternalServersMappingFromEnv(t *testing.T) {
 	setupBootstrapServersMappingTest()
 
-	os.Setenv("BOOTSTRAP_SERVER_MAPPING", "	192.168.99.100:32401,0.0.0.0:32401")
-	os.Setenv("EXTERNAL_SERVER_MAPPING", "	192.168.99.100:32404,0.0.0.0:32404	kafka-5.example.com:9092,0.0.0.0:32405,kafka-5.grepplabs.com:9092")
+	_ = os.Setenv("BOOTSTRAP_SERVER_MAPPING", "	192.168.99.100:32401,0.0.0.0:32401")
+	_ = os.Setenv("EXTERNAL_SERVER_MAPPING", "	192.168.99.100:32404,0.0.0.0:32404	kafka-5.example.com:9092,0.0.0.0:32405,kafka-5.grepplabs.com:9092")
 
 	var args []string
-	Server.ParseFlags(args)
+	_ = Server.ParseFlags(args)
 	err := Server.PreRunE(nil, args)
 	a := assert.New(t)
 	a.Nil(err)

config/config.go

@@ -24,6 +24,10 @@ type ListenerConfig struct {
 	ListenerAddress   string
 	AdvertisedAddress string
 }
+type DialAddressMapping struct {
+	SourceAddress      string
+	DestinationAddress string
+}
 
 type Config struct {
 	Http struct {
@@ -46,6 +50,7 @@ type Config struct {
 		DefaultListenerIP       string
 		BootstrapServers        []ListenerConfig
 		ExternalServers         []ListenerConfig
+		DialAddressMappings     []DialAddressMapping
 		DisableDynamicListeners bool
 		RequestBufferSize       int
 		ResponseBufferSize      int
@@ -146,11 +151,17 @@ func (c *Config) InitBootstrapServers(bootstrapServersMapping []string) (err err
 	c.Proxy.BootstrapServers, err = getListenerConfigs(bootstrapServersMapping)
 	return err
 }
+
 func (c *Config) InitExternalServers(externalServersMapping []string) (err error) {
 	c.Proxy.ExternalServers, err = getListenerConfigs(externalServersMapping)
 	return err
 }
 
+func (c *Config) InitDialAddressMappings(dialMappings []string) (err error) {
+	c.Proxy.DialAddressMappings, err = getDialAddressMappings(dialMappings)
+	return err
+}
+
 func (c *Config) InitSASLCredentials() (err error) {
 	if c.Kafka.SASL.JaasConfigFile != "" {
 		credentials, err := NewJaasCredentialFromFile(c.Kafka.SASL.JaasConfigFile)
@@ -162,6 +173,30 @@ func (c *Config) InitSASLCredentials() (err error) {
 	}
 	return nil
 }
+func getDialAddressMappings(dialMapping []string) ([]DialAddressMapping, error) {
+	dialMappings := make([]DialAddressMapping, 0)
+	if dialMapping != nil {
+		for _, v := range dialMapping {
+			pair := strings.Split(v, ",")
+			if len(pair) != 2 {
+				return nil, errors.New("dial-mapping must be in form 'srchost:srcport,dsthost:dstport'")
+			}
+			srcHost, srcPort, err := util.SplitHostPort(pair[0])
+			if err != nil {
+				return nil, err
+			}
+			dstHost, dstPort, err := util.SplitHostPort(pair[1])
+			if err != nil {
+				return nil, err
+			}
+			dialMapping := DialAddressMapping{
+				SourceAddress:      net.JoinHostPort(srcHost, fmt.Sprint(srcPort)),
+				DestinationAddress: net.JoinHostPort(dstHost, fmt.Sprint(dstPort))}
+			dialMappings = append(dialMappings, dialMapping)
+		}
+	}
+	return dialMappings, nil
+}
 
 func getListenerConfigs(serversMapping []string) ([]ListenerConfig, error) {
 	listenerConfigs := make([]ListenerConfig, 0)

proxy/client.go

@@ -2,6 +2,7 @@ package proxy
 
 import (
 	"crypto/tls"
+	"fmt"
 	"github.com/grepplabs/kafka-proxy/config"
 	"github.com/grepplabs/kafka-proxy/pkg/apis"
 	"github.com/pkg/errors"
@@ -36,6 +37,8 @@ type Client struct {
 
 	saslAuthByProxy SASLAuthByProxy
 	authClient      *AuthClient
+
+	dialAddressMapping map[string]config.DialAddressMapping
 }
 
 func NewClient(conns *ConnSet, c *config.Config, netAddressMappingFunc config.NetAddressMappingFunc, localPasswordAuthenticator apis.PasswordAuthenticator, localTokenAuthenticator apis.TokenInfo, saslTokenProvider apis.TokenProvider, gatewayTokenProvider apis.TokenProvider, gatewayTokenInfo apis.TokenInfo) (*Client, error) {
@@ -105,6 +108,10 @@ func NewClient(conns *ConnSet, c *config.Config, netAddressMappingFunc config.Ne
 			return nil, errors.Errorf("SASL Mechanism not valid '%s'", c.Kafka.SASL.Method)
 		}
 	}
+	dialAddressMapping, err := getAddressToDialAddressMapping(c)
+	if err != nil {
+		return nil, err
+	}
 
 	return &Client{conns: conns, config: c, dialer: dialer, tcpConnOptions: tcpConnOptions, stopRun: make(chan struct{}, 1),
 		saslAuthByProxy: saslAuthByProxy,
@@ -136,7 +143,25 @@ func NewClient(conns *ConnSet, c *config.Config, netAddressMappingFunc config.Ne
 				tokenInfo: gatewayTokenInfo,
 			},
 			ForbiddenApiKeys: forbiddenApiKeys,
-		}}, nil
+		},
+		dialAddressMapping: dialAddressMapping,
+	}, nil
+}
+
+func getAddressToDialAddressMapping(cfg *config.Config) (map[string]config.DialAddressMapping, error) {
+	addressToDialAddressMapping := make(map[string]config.DialAddressMapping)
+
+	for _, v := range cfg.Proxy.DialAddressMappings {
+		if lc, ok := addressToDialAddressMapping[v.SourceAddress]; ok {
+			if lc.SourceAddress != v.SourceAddress || lc.DestinationAddress != v.DestinationAddress {
+				return nil, fmt.Errorf("dial address mapping %s configured twice: %v and %v", v.SourceAddress, v, lc)
+			}
+			continue
+		}
+		logrus.Infof("Dial address mapping src %s dst %s", v.SourceAddress, v.DestinationAddress)
+		addressToDialAddressMapping[v.SourceAddress] = v
+	}
+	return addressToDialAddressMapping, nil
 }
 
 func newDialer(c *config.Config, tlsConfig *tls.Config) (Dialer, error) {
@@ -219,9 +244,15 @@ func (c *Client) Close() {
 func (c *Client) handleConn(conn Conn) {
 	proxyConnectionsTotal.WithLabelValues(conn.BrokerAddress).Inc()
 
-	server, err := c.DialAndAuth(conn.BrokerAddress)
+	dialAddress := conn.BrokerAddress
+	if addressMapping, ok := c.dialAddressMapping[dialAddress]; ok {
+		dialAddress = addressMapping.DestinationAddress
+		logrus.Infof("Dial address changed from %s to %s", conn.BrokerAddress, dialAddress)
+	}
+
+	server, err := c.DialAndAuth(dialAddress)
 	if err != nil {
-		logrus.Infof("couldn't connect to %s: %v", conn.BrokerAddress, err)
+		logrus.Infof("couldn't connect to %s(%s): %v", dialAddress, conn.BrokerAddress, err)
 		_ = conn.LocalConnection.Close()
 		return
 	}

proxy/proxy.go

@@ -122,9 +122,11 @@ func (p *Listeners) GetNetAddressMapping(brokerHost string, brokerPort int32) (l
 	p.lock.RUnlock()
 
 	if ok {
+		logrus.Debugf("Address mappings broker=%s, listener=%s, advertised=%s", listenerConfig.BrokerAddress, listenerConfig.ListenerAddress, listenerConfig.AdvertisedAddress)
 		return util.SplitHostPort(listenerConfig.AdvertisedAddress)
 	}
 	if !p.disableDynamicListeners {
+		logrus.Infof("Starting dynamic listener for broker %s", brokerAddress)
 		return p.ListenDynamicInstance(brokerAddress)
 	}
 	return "", 0, fmt.Errorf("net address mapping for %s:%d was not found", brokerHost, brokerPort)

proxy/sasl_scram.go

@@ -98,7 +98,7 @@ func (b *SASLSCRAMAuth) sendAndReceiveSASLAuth(conn DeadlineReaderWriter) error
 
 		msg, err = scramConversation.Step(string(challenge))
 		if err != nil {
-			logrus.Debugf("SASL authentication failed", err)
+			logrus.Debugf("SASL authentication failed %s", err)
 			//Logger.Println("SASL authentication failed", err)
 			return err
 		}