server: add proxyprotocol to http (#270)

conf/proxy.toml

@@ -34,6 +34,9 @@
 # user = ""
 # password = ""
 
+# same as [proxy.proxy-protocol], but for HTTP port
+# proxy-protocol = ""
+
 [log]
 
 # level = "info"

lib/config/proxy.go

@@ -84,6 +84,7 @@ type API struct {
 	User            string `yaml:"user,omitempty" toml:"user,omitempty" json:"user,omitempty"`
 	Password        string `yaml:"password,omitempty" toml:"password,omitempty" json:"password,omitempty"`
 	EnableBasicAuth bool   `yaml:"enable-basic-auth,omitempty" toml:"enable-basic-auth,omitempty" json:"enable-basic-auth,omitempty"`
+	ProxyProtocol   string `yaml:"proxy-protocol,omitempty" toml:"proxy-protocol,omitempty" json:"proxy-protocol,omitempty"`
 }
 
 type Advance struct {

pkg/proxy/proxyprotocol/listener.go

@@ -0,0 +1,81 @@
+// Copyright 2023 PingCAP, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package proxyprotocol
+
+import (
+	"bytes"
+	"io"
+	"net"
+)
+
+var _ net.Listener = (*Listener)(nil)
+var _ net.Conn = (*proxyConn)(nil)
+
+type Listener struct {
+	net.Listener
+}
+
+func NewListener(o net.Listener) *Listener {
+	return &Listener{o}
+}
+
+func (n *Listener) Accept() (net.Conn, error) {
+	conn, err := n.Listener.Accept()
+	return &proxyConn{Conn: conn, buf: new(bytes.Buffer)}, err
+}
+
+type proxyConn struct {
+	net.Conn
+	buf    *bytes.Buffer
+	proxy  *Proxy
+	inited bool
+}
+
+func (c *proxyConn) Read(b []byte) (n int, err error) {
+	if !c.inited {
+		_, err = c.buf.ReadFrom(io.LimitReader(c.Conn, int64(len(MagicV2)-c.buf.Len())))
+		if err != nil {
+			return
+		}
+		if bytes.HasPrefix(MagicV2, c.buf.Bytes()) {
+			if !bytes.Equal(MagicV2, c.buf.Bytes()) {
+				// prefix matches, maybe proxy header
+				// read again later
+				return 0, nil
+			}
+			// it is proxy protocol
+			c.buf.Reset()
+			c.proxy, _, err = ParseProxyV2(c.Conn)
+			if err != nil {
+				return 0, err
+			}
+		}
+		// prefixes mismatched, or we have parsed PP header
+		c.inited = true
+	}
+	if c.buf.Len() > 0 {
+		n = copy(b, c.buf.Bytes())
+		_ = c.buf.Next(n)
+		return n, nil
+	}
+	return c.Conn.Read(b)
+}
+
+func (c *proxyConn) RemoteAddr() net.Addr {
+	if c.proxy != nil && c.proxy.DstAddress != nil {
+		return c.proxy.DstAddress
+	}
+	return c.Conn.RemoteAddr()
+}

pkg/proxy/proxyprotocol/listener_test.go

@@ -0,0 +1,84 @@
+// Copyright 2023 PingCAP, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package proxyprotocol
+
+import (
+	"bytes"
+	"io"
+	"net"
+	"strings"
+	"testing"
+
+	"github.com/pingcap/TiProxy/pkg/testkit"
+	"github.com/stretchr/testify/require"
+)
+
+func TestProxyListener(t *testing.T) {
+	tcpaddr, err := net.ResolveTCPAddr("tcp", "192.168.1.1:34")
+	require.NoError(t, err)
+
+	testkit.TestTCPConnWithListener(t,
+		func(t *testing.T, network, addr string) net.Listener {
+			ln, err := net.Listen(network, addr)
+			require.NoError(t, err)
+			return NewListener(ln)
+		},
+		func(t *testing.T, c net.Conn) {
+			p := &Proxy{
+				Version:    ProxyVersion2,
+				Command:    ProxyCommandLocal,
+				SrcAddress: tcpaddr,
+				DstAddress: tcpaddr,
+				TLV: []ProxyTlv{
+					{
+						typ:     ProxyTlvALPN,
+						content: nil,
+					},
+					{
+						typ:     ProxyTlvUniqueID,
+						content: []byte("test"),
+					},
+				},
+			}
+			b, err := p.ToBytes()
+			require.NoError(t, err)
+			_, err = io.Copy(c, bytes.NewReader(b))
+			require.NoError(t, err)
+			_, err = io.Copy(c, strings.NewReader("test"))
+			require.NoError(t, err)
+		},
+		func(t *testing.T, c net.Conn) {
+			all, err := io.ReadAll(c)
+			require.NoError(t, err)
+			require.Equal(t, []byte("test"), all)
+			require.Equal(t, tcpaddr.String(), c.RemoteAddr().String())
+		}, 1)
+
+	testkit.TestTCPConnWithListener(t,
+		func(t *testing.T, network, addr string) net.Listener {
+			ln, err := net.Listen(network, addr)
+			require.NoError(t, err)
+			return NewListener(ln)
+		},
+		func(t *testing.T, c net.Conn) {
+			_, err = io.Copy(c, strings.NewReader("test"))
+			require.NoError(t, err)
+		},
+		func(t *testing.T, c net.Conn) {
+			all, err := io.ReadAll(c)
+			require.NoError(t, err)
+			require.Equal(t, []byte("test"), all)
+		}, 1)
+}

pkg/server/api/http.go

@@ -29,6 +29,7 @@ import (
 	mgrcfg "github.com/pingcap/TiProxy/pkg/manager/config"
 	mgrns "github.com/pingcap/TiProxy/pkg/manager/namespace"
 	"github.com/pingcap/TiProxy/pkg/proxy"
+	"github.com/pingcap/TiProxy/pkg/proxy/proxyprotocol"
 	"go.uber.org/atomic"
 	"go.uber.org/ratelimit"
 	"go.uber.org/zap"
@@ -80,6 +81,10 @@ func NewHTTPServer(cfg config.API, lg *zap.Logger,
 	if err != nil {
 		return nil, err
 	}
+	switch cfg.ProxyProtocol {
+	case "v2":
+		h.listener = proxyprotocol.NewListener(h.listener)
+	}
 
 	gin.SetMode(gin.ReleaseMode)
 	engine := gin.New()

pkg/testkit/main.go

@@ -43,8 +43,15 @@ func TestPipeConn(t *testing.T, a, b func(*testing.T, net.Conn), loop int) {
 }
 
 func TestTCPConn(t *testing.T, a, b func(*testing.T, net.Conn), loop int) {
-	listener, err := net.Listen("tcp", "127.0.0.0:0")
-	require.NoError(t, err)
+	TestTCPConnWithListener(t, func(t *testing.T, network, addr string) net.Listener {
+		ln, err := net.Listen(network, addr)
+		require.NoError(t, err)
+		return ln
+	}, a, b, loop)
+}
+
+func TestTCPConnWithListener(t *testing.T, listen func(*testing.T, string, string) net.Listener, a, b func(*testing.T, net.Conn), loop int) {
+	listener := listen(t, "tcp", "localhost:0")
 	defer func() {
 		require.NoError(t, listener.Close())
 	}()