Support proxy protocol v2 in MySQL (#12424)

#11684 added support for proxy protocol v2 for SSH and Postgres but MySQL uses different code path and it was missing. This change fixes that.
It also adds tests for v2 protocol support for MySQL, Postgres, Mongo and Redis

Author: probakowski

Diff for: lib/multiplexer/proxyline.go

@@ -60,6 +60,39 @@ func (p *ProxyLine) String() string {
 	return fmt.Sprintf("PROXY %s %s %s %d %d\r\n", p.Protocol, p.Source.IP.String(), p.Destination.IP.String(), p.Source.Port, p.Destination.Port)
 }
 
+// Bytes returns on-the wire bytes representation of proxy line conforming to the proxy v2 protocol
+func (p *ProxyLine) Bytes() []byte {
+	b := &bytes.Buffer{}
+	header := proxyV2Header{VersionCommand: (Version2 << 4) | ProxyCommand}
+	copy(header.Signature[:], proxyV2Prefix)
+	var addr interface{}
+	switch p.Protocol {
+	case TCP4:
+		header.Protocol = ProtocolTCP4
+		addr4 := proxyV2Address4{
+			SourcePort:      uint16(p.Source.Port),
+			DestinationPort: uint16(p.Destination.Port),
+		}
+		copy(addr4.Source[:], p.Source.IP.To4())
+		copy(addr4.Destination[:], p.Destination.IP.To4())
+		addr = addr4
+	case TCP6:
+		header.Protocol = ProtocolTCP6
+		addr6 := proxyV2Address6{
+			SourcePort:      uint16(p.Source.Port),
+			DestinationPort: uint16(p.Destination.Port),
+		}
+		copy(addr6.Source[:], p.Source.IP.To16())
+		copy(addr6.Destination[:], p.Destination.IP.To16())
+		addr = addr6
+	}
+	header.Length = uint16(binary.Size(addr))
+	binary.Write(b, binary.BigEndian, header)
+	binary.Write(b, binary.BigEndian, addr)
+
+	return b.Bytes()
+}
+
 // ReadProxyLine reads proxy line protocol from the reader
 func ReadProxyLine(reader *bufio.Reader) (*ProxyLine, error) {
 	line, err := reader.ReadString('\n')

Diff for: lib/multiplexer/testproxy.go

@@ -33,11 +33,12 @@ type TestProxy struct {
 	target   string
 	closeCh  chan (struct{})
 	log      logrus.FieldLogger
+	v2       bool
 }
 
 // NewTestProxy creates a new test proxy that sends a proxy-line when
 // proxying connections to the provided target address.
-func NewTestProxy(target string) (*TestProxy, error) {
+func NewTestProxy(target string, v2 bool) (*TestProxy, error) {
 	listener, err := net.Listen("tcp", "localhost:0")
 	if err != nil {
 		return nil, trace.Wrap(err)
@@ -47,6 +48,7 @@ func NewTestProxy(target string) (*TestProxy, error) {
 		target:   target,
 		closeCh:  make(chan struct{}),
 		log:      logrus.WithField(trace.Component, "test:proxy"),
+		v2:       v2,
 	}, nil
 }
 
@@ -128,7 +130,11 @@ func (p *TestProxy) sendProxyLine(clientConn, serverConn net.Conn) error {
 		Destination: net.TCPAddr{IP: net.ParseIP(serverAddr.Host()), Port: serverAddr.Port(0)},
 	}
 	p.log.Debugf("Sending %v to %v.", proxyLine.String(), serverConn.RemoteAddr().String())
-	_, err = serverConn.Write([]byte(proxyLine.String()))
+	if p.v2 {
+		_, err = serverConn.Write(proxyLine.Bytes())
+	} else {
+		_, err = serverConn.Write([]byte(proxyLine.String()))
+	}
 	if err != nil {
 		return trace.Wrap(err)
 	}

Diff for: lib/multiplexer/wrappers.go

@@ -84,7 +84,16 @@ func (c *Conn) Detect() (Protocol, error) {
 
 // ReadProxyLine reads proxy-line from the connection.
 func (c *Conn) ReadProxyLine() (*ProxyLine, error) {
-	proxyLine, err := ReadProxyLine(c.reader)
+	var proxyLine *ProxyLine
+	protocol, err := c.Detect()
+	if err != nil {
+		return nil, trace.Wrap(err)
+	}
+	if protocol == ProtoProxyV2 {
+		proxyLine, err = ReadProxyLineV2(c.reader)
+	} else {
+		proxyLine, err = ReadProxyLine(c.reader)
+	}
 	if err != nil {
 		return nil, trace.Wrap(err)
 	}

Diff for: lib/srv/db/mysql/proxy.go

@@ -212,7 +212,7 @@ func (p *Proxy) maybeReadProxyLine(conn *multiplexer.Conn) error {
 	if err != nil {
 		return trace.Wrap(err)
 	}
-	if proto != multiplexer.ProtoProxy {
+	if proto != multiplexer.ProtoProxy && proto != multiplexer.ProtoProxyV2 {
 		return nil
 	}
 	proxyLine, err := conn.ReadProxyLine()

Diff for: lib/srv/db/proxy_test.go

@@ -18,6 +18,7 @@ package db
 
 import (
 	"context"
+	"fmt"
 	"testing"
 	"time"
 
@@ -40,17 +41,21 @@ func TestProxyProtocolPostgres(t *testing.T) {
 
 	testCtx.createUserAndRole(ctx, t, "alice", "admin", []string{"postgres"}, []string{"postgres"})
 
-	// Point our proxy to the Teleport's db listener on the multiplexer.
-	proxy, err := multiplexer.NewTestProxy(testCtx.mux.DB().Addr().String())
-	require.NoError(t, err)
-	t.Cleanup(func() { proxy.Close() })
-	go proxy.Serve()
-
-	// Connect to the proxy instead of directly to Postgres listener and make
-	// sure the connection succeeds.
-	psql, err := testCtx.postgresClientWithAddr(ctx, proxy.Address(), "alice", "postgres", "postgres", "postgres")
-	require.NoError(t, err)
-	require.NoError(t, psql.Close(ctx))
+	for _, v2 := range []bool{false, true} {
+		t.Run(fmt.Sprintf("v2=%v", v2), func(t *testing.T) {
+			// Point our proxy to the Teleport's db listener on the multiplexer.
+			proxy, err := multiplexer.NewTestProxy(testCtx.mux.DB().Addr().String(), v2)
+			require.NoError(t, err)
+			t.Cleanup(func() { proxy.Close() })
+			go proxy.Serve()
+
+			// Connect to the proxy instead of directly to Postgres listener and make
+			// sure the connection succeeds.
+			psql, err := testCtx.postgresClientWithAddr(ctx, proxy.Address(), "alice", "postgres", "postgres", "postgres")
+			require.NoError(t, err)
+			require.NoError(t, psql.Close(ctx))
+		})
+	}
 }
 
 // TestProxyProtocolMySQL ensures that clients can successfully connect to a
@@ -63,17 +68,21 @@ func TestProxyProtocolMySQL(t *testing.T) {
 
 	testCtx.createUserAndRole(ctx, t, "alice", "admin", []string{"root"}, []string{types.Wildcard})
 
-	// Point our proxy to the Teleport's MySQL listener.
-	proxy, err := multiplexer.NewTestProxy(testCtx.mysqlListener.Addr().String())
-	require.NoError(t, err)
-	t.Cleanup(func() { proxy.Close() })
-	go proxy.Serve()
-
-	// Connect to the proxy instead of directly to MySQL listener and make
-	// sure the connection succeeds.
-	mysql, err := testCtx.mysqlClientWithAddr(proxy.Address(), "alice", "mysql", "root")
-	require.NoError(t, err)
-	require.NoError(t, mysql.Close())
+	for _, v2 := range []bool{false, true} {
+		t.Run(fmt.Sprintf("v2=%v", v2), func(t *testing.T) {
+			// Point our proxy to the Teleport's MySQL listener.
+			proxy, err := multiplexer.NewTestProxy(testCtx.mysqlListener.Addr().String(), v2)
+			require.NoError(t, err)
+			t.Cleanup(func() { proxy.Close() })
+			go proxy.Serve()
+
+			// Connect to the proxy instead of directly to MySQL listener and make
+			// sure the connection succeeds.
+			mysql, err := testCtx.mysqlClientWithAddr(proxy.Address(), "alice", "mysql", "root")
+			require.NoError(t, err)
+			require.NoError(t, mysql.Close())
+		})
+	}
 }
 
 // TestProxyProtocolMongo ensures that clients can successfully connect to a
@@ -86,17 +95,21 @@ func TestProxyProtocolMongo(t *testing.T) {
 
 	testCtx.createUserAndRole(ctx, t, "alice", "admin", []string{"admin"}, []string{types.Wildcard})
 
-	// Point our proxy to the Teleport's TLS listener.
-	proxy, err := multiplexer.NewTestProxy(testCtx.webListener.Addr().String())
-	require.NoError(t, err)
-	t.Cleanup(func() { proxy.Close() })
-	go proxy.Serve()
-
-	// Connect to the proxy instead of directly to Teleport listener and make
-	// sure the connection succeeds.
-	mongo, err := testCtx.mongoClientWithAddr(ctx, proxy.Address(), "alice", "mongo", "admin")
-	require.NoError(t, err)
-	require.NoError(t, mongo.Disconnect(ctx))
+	for _, v2 := range []bool{false, true} {
+		t.Run(fmt.Sprintf("v2=%v", v2), func(t *testing.T) {
+			// Point our proxy to the Teleport's TLS listener.
+			proxy, err := multiplexer.NewTestProxy(testCtx.webListener.Addr().String(), v2)
+			require.NoError(t, err)
+			t.Cleanup(func() { proxy.Close() })
+			go proxy.Serve()
+
+			// Connect to the proxy instead of directly to Teleport listener and make
+			// sure the connection succeeds.
+			mongo, err := testCtx.mongoClientWithAddr(ctx, proxy.Address(), "alice", "mongo", "admin")
+			require.NoError(t, err)
+			require.NoError(t, mongo.Disconnect(ctx))
+		})
+	}
 }
 
 func TestProxyProtocolRedis(t *testing.T) {
@@ -106,23 +119,27 @@ func TestProxyProtocolRedis(t *testing.T) {
 
 	testCtx.createUserAndRole(ctx, t, "alice", "admin", []string{"admin"}, []string{types.Wildcard})
 
-	// Point our proxy to the Teleport's TLS listener.
-	proxy, err := multiplexer.NewTestProxy(testCtx.webListener.Addr().String())
-	require.NoError(t, err)
-	t.Cleanup(func() { proxy.Close() })
-	go proxy.Serve()
-
-	// Connect to the proxy instead of directly to Teleport listener and make
-	// sure the connection succeeds.
-	redisClient, err := testCtx.redisClientWithAddr(ctx, proxy.Address(), "alice", "redis", "admin")
-	require.NoError(t, err)
-
-	// Send ECHO to Redis server and check if we get it back.
-	resp := redisClient.Echo(ctx, "hello")
-	require.NoError(t, resp.Err())
-	require.Equal(t, "hello", resp.Val())
-
-	require.NoError(t, redisClient.Close())
+	for _, v2 := range []bool{false, true} {
+		t.Run(fmt.Sprintf("v2=%v", v2), func(t *testing.T) {
+			// Point our proxy to the Teleport's TLS listener.
+			proxy, err := multiplexer.NewTestProxy(testCtx.webListener.Addr().String(), v2)
+			require.NoError(t, err)
+			t.Cleanup(func() { proxy.Close() })
+			go proxy.Serve()
+
+			// Connect to the proxy instead of directly to Teleport listener and make
+			// sure the connection succeeds.
+			redisClient, err := testCtx.redisClientWithAddr(ctx, proxy.Address(), "alice", "redis", "admin")
+			require.NoError(t, err)
+
+			// Send ECHO to Redis server and check if we get it back.
+			resp := redisClient.Echo(ctx, "hello")
+			require.NoError(t, resp.Err())
+			require.Equal(t, "hello", resp.Val())
+
+			require.NoError(t, redisClient.Close())
+		})
+	}
 }
 
 // TestProxyClientDisconnectDueToIdleConnection ensures that idle clients will be disconnected.