proxy: expose proxy protocol API (#25)

pingcap · Aug 3, 2022 · d86b629 · d86b629
1 parent a6db638
commit d86b629
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 8 deletions.
diff --git a/pkg/proxy/net/packetio.go b/pkg/proxy/net/packetio.go
@@ -45,6 +45,7 @@ import (
 	"github.com/pingcap/tidb/errno"
 	"github.com/pingcap/tidb/parser/mysql"
 	"github.com/pingcap/tidb/util/dbterror"
+	"go.uber.org/atomic"
 )
 
 var (
@@ -76,7 +77,7 @@ type PacketIO struct {
 	conn        net.Conn
 	buf         *bufio.ReadWriter
 	sequence    uint8
-	proxyInited bool
+	proxyInited *atomic.Bool
 	proxy       *Proxy
 }
 
@@ -91,14 +92,21 @@ func NewPacketIO(conn net.Conn) *PacketIO {
 			buf.Reader,
 		},
 		sequence: 0,
-		// TODO: enable proxy probe for clients only
-		// disable it by default now
-		proxyInited: true,
+		// TODO: disable it by default now
+		proxyInited: atomic.NewBool(true),
 		buf:         buf,
 	}
 	return p
 }
 
+// Proxy returned parsed proxy header from clients if any.
+func (p *PacketIO) Proxy() *Proxy {
+	if p.proxyInited.Load() {
+		return p.proxy
+	}
+	return nil
+}
+
 func (p *PacketIO) LocalAddr() net.Addr {
 	return p.conn.LocalAddr()
 }
@@ -120,7 +128,7 @@ func (p *PacketIO) ReadOnePacket() ([]byte, bool, error) {
 
 	// probe proxy V2
 	refill := false
-	if !p.proxyInited {
+	if !p.proxyInited.Load() {
 		if bytes.Compare(header[:], proxyV2Magic[:4]) == 0 {
 			proxyHeader, err := p.parseProxyV2()
 			if err != nil {
@@ -131,7 +139,7 @@ func (p *PacketIO) ReadOnePacket() ([]byte, bool, error) {
 				refill = true
 			}
 		}
-		p.proxyInited = true
+		p.proxyInited.Store(true)
 	}
 
 	// refill mysql headers

diff --git a/pkg/proxy/net/proxy.go b/pkg/proxy/net/proxy.go
@@ -267,9 +267,11 @@ func (p *PacketIO) parseProxyV2() (*Proxy, error) {
 	return m, nil
 }
 
-func (p *PacketIO) writeProxyV2(m *Proxy) error {
+// WriteProxyV2 should only be called at the beginning of connection, before any write operations.
+func (p *PacketIO) WriteProxyV2(m *Proxy) error {
 	if _, err := io.Copy(p.buf, bytes.NewReader(m.ToBytes())); err != nil {
 		return errors.WithStack(errors.Wrap(ErrWriteConn, err))
 	}
+	// according to the spec, we better flush to avoid server hanging
 	return p.Flush()
 }
diff --git a/pkg/proxy/net/proxy_test.go b/pkg/proxy/net/proxy_test.go
@@ -28,7 +28,7 @@ func TestProxy(t *testing.T) {
 
 	testPipeConn(t,
 		func(t *testing.T, cli *PacketIO) {
-			require.NoError(t, cli.writeProxyV2(&Proxy{
+			require.NoError(t, cli.WriteProxyV2(&Proxy{
 				Version:    ProxyVersion2,
 				Command:    ProxyCommandLocal,
 				SrcAddress: tcpaddr,