proxy: fix proxy-protocol under originAddr (pingcap#219)

xhebox · Mar 13, 2023 · c3d8867 · c3d8867
1 parent f5a7080
commit c3d8867
Show file tree

Hide file tree

Showing 4 changed files with 27 additions and 15 deletions.
diff --git a/pkg/proxy/backend/backend_conn_mgr.go b/pkg/proxy/backend/backend_conn_mgr.go
@@ -228,7 +228,7 @@ func (mgr *BackendConnManager) getBackendIO(cctx ConnContext, auth *Authenticato
 			// NOTE: should use DNS name as much as possible
 			// Usually certs are signed with domain instead of IP addrs
 			// And `RemoteAddr()` will return IP addr
-			mgr.backendIO = pnet.NewPacketIO(cn, pnet.WithRemoteAddr(addr))
+			mgr.backendIO = pnet.NewPacketIO(cn, pnet.WithRemoteAddr(addr, cn.RemoteAddr()))
 			return mgr.backendIO, nil
 		},
 		backoff.WithContext(backoff.NewConstantBackOff(200*time.Millisecond), bctx),
@@ -429,7 +429,7 @@ func (mgr *BackendConnManager) tryRedirect(ctx context.Context) {
 		mgr.handshakeHandler.OnHandshake(mgr, rs.to, rs.err)
 		return
 	}
-	newBackendIO := pnet.NewPacketIO(cn, pnet.WithRemoteAddr(rs.to))
+	newBackendIO := pnet.NewPacketIO(cn, pnet.WithRemoteAddr(rs.to, cn.RemoteAddr()))
 
 	if rs.err = mgr.authenticator.handshakeSecondTime(mgr.logger, mgr.clientIO, newBackendIO, mgr.backendTLS, sessionToken); rs.err == nil {
 		rs.err = mgr.initSessionStates(newBackendIO, sessionStates)

diff --git a/pkg/proxy/net/packetio_options.go b/pkg/proxy/net/packetio_options.go
@@ -33,22 +33,26 @@ func WithWrapError(err error) func(pi *PacketIO) {
 }
 
 // WithRemoteAddr
-var _ net.Addr = &oriRemoteAddr{}
+var _ net.Addr = &originAddr{}
 
-type oriRemoteAddr struct {
+type originAddr struct {
+	net.Addr
 	addr string
 }
 
-func (o *oriRemoteAddr) Network() string {
-	return "tcp"
-}
-
-func (o *oriRemoteAddr) String() string {
+func (o *originAddr) String() string {
 	return o.addr
 }
 
-func WithRemoteAddr(readdr string) func(pi *PacketIO) {
+func WithRemoteAddr(readdr string, addr net.Addr) func(pi *PacketIO) {
 	return func(pi *PacketIO) {
-		pi.remoteAddr = &oriRemoteAddr{addr: readdr}
+		pi.remoteAddr = &originAddr{Addr: addr, addr: readdr}
+	}
+}
+
+func unwrapOriginAddr(addr net.Addr) net.Addr {
+	if oaddr, ok := addr.(*originAddr); ok {
+		return oaddr.Addr
 	}
+	return addr
 }
diff --git a/pkg/proxy/net/proxy.go b/pkg/proxy/net/proxy.go
@@ -93,14 +93,18 @@ func (p *Proxy) ToBytes() ([]byte, error) {
 
 	addressFamily := ProxyAFUnspec
 	network := ProxyNetworkUnspec
-	switch sadd := p.SrcAddress.(type) {
+
+	srcAddr := unwrapOriginAddr(p.SrcAddress)
+	dstAddr := unwrapOriginAddr(p.DstAddress)
+
+	switch sadd := srcAddr.(type) {
 	case *net.TCPAddr:
 		addressFamily = ProxyAFINet
 		if len(sadd.IP) == net.IPv6len {
 			addressFamily = ProxyAFINet6
 		}
 		network = ProxyNetworkStream
-		dadd, ok := p.DstAddress.(*net.TCPAddr)
+		dadd, ok := dstAddr.(*net.TCPAddr)
 		if !ok {
 			return nil, ErrAddressFamilyMismatch
 		}
@@ -114,7 +118,7 @@ func (p *Proxy) ToBytes() ([]byte, error) {
 			addressFamily = ProxyAFINet6
 		}
 		network = ProxyNetworkDgram
-		dadd, ok := p.DstAddress.(*net.UDPAddr)
+		dadd, ok := dstAddr.(*net.UDPAddr)
 		if !ok {
 			return nil, ErrAddressFamilyMismatch
 		}
@@ -130,7 +134,7 @@ func (p *Proxy) ToBytes() ([]byte, error) {
 		case "unixdgram":
 			network = ProxyNetworkDgram
 		}
-		dadd, ok := p.DstAddress.(*net.UnixAddr)
+		dadd, ok := dstAddr.(*net.UnixAddr)
 		if !ok {
 			return nil, ErrAddressFamilyMismatch
 		}

diff --git a/pkg/proxy/net/proxy_test.go b/pkg/proxy/net/proxy_test.go
@@ -84,4 +84,8 @@ func TestProxyToBytes(t *testing.T) {
 	hdr.DstAddress = &net.UDPAddr{}
 	_, err = hdr.ToBytes()
 	require.ErrorIs(t, err, ErrAddressFamilyMismatch)
+
+	hdr.DstAddress = &originAddr{Addr: &net.TCPAddr{IP: make(net.IP, net.IPv6len), Port: 0}}
+	_, err = hdr.ToBytes()
+	require.NoError(t, err)
 }