proxy, util: meter public IP addresses into public traffic (#1006)

Author: djshow832

pkg/balance/router/group.go

@@ -95,9 +95,14 @@ func (g *Group) Match(clientInfo ClientInfo) bool {
 		if g.matchType == MatchClientCIDR {
 			addr = clientInfo.ClientAddr
 		}
-		contains, err := netutil.CIDRContainsIP(g.cidrList, addr)
+		ip, err := netutil.NetAddr2IP(addr)
 		if err != nil {
-			g.lg.Error("checking CIDR failed", zap.String("addr", addr.String()), zap.Error(err))
+			g.lg.Error("checking CIDR failed", zap.Stringer("addr", addr), zap.Error(err))
+			return false
+		}
+		contains, err := netutil.CIDRContainsIP(g.cidrList, ip)
+		if err != nil {
+			g.lg.Error("checking CIDR failed", zap.Stringer("addr", addr), zap.Error(err))
 		}
 		return contains
 	}

pkg/proxy/proxy.go

@@ -6,6 +6,7 @@ package proxy
 import (
 	"context"
 	"net"
+	"reflect"
 	"strings"
 	"sync"
 	"time"
@@ -166,11 +167,12 @@ func (s *SQLServer) onConn(ctx context.Context, conn net.Conn, addr string) {
 			zap.String("addr", addr))
 		clientConn := client.NewClientConnection(logger.Named("conn"), conn, s.certMgr.ServerSQLTLS(), s.certMgr.SQLTLS(),
 			s.hsHandler, s.cpt, connID, addr, &backend.BCConfig{
-				ProxyProtocol:      s.mu.proxyProtocol,
-				RequireBackendTLS:  s.mu.requireBackendTLS,
-				HealthyKeepAlive:   s.mu.healthyKeepAlive,
-				UnhealthyKeepAlive: s.mu.unhealthyKeepAlive,
-				ConnBufferSize:     s.mu.connBufferSize,
+				ProxyProtocol:       s.mu.proxyProtocol,
+				RequireBackendTLS:   s.mu.requireBackendTLS,
+				HealthyKeepAlive:    s.mu.healthyKeepAlive,
+				UnhealthyKeepAlive:  s.mu.unhealthyKeepAlive,
+				ConnBufferSize:      s.mu.connBufferSize,
+				FromPublicEndpoints: s.fromPublicEndpoint,
 			}, s.meter)
 		s.mu.clients[connID] = clientConn
 		logger.Debug("new connection", zap.Bool("proxy-protocol", s.mu.proxyProtocol), zap.Bool("require_backend_tls", s.mu.requireBackendTLS))
@@ -204,6 +206,31 @@ func (s *SQLServer) onConn(ctx context.Context, conn net.Conn, addr string) {
 	clientConn.Run(ctx)
 }
 
+func (s *SQLServer) fromPublicEndpoint(addr net.Addr) bool {
+	if addr == nil || reflect.ValueOf(addr).IsNil() {
+		return false
+	}
+	s.mu.RLock()
+	publicEndpoints := s.mu.publicEndpoints
+	s.mu.RUnlock()
+	ip, err := netutil.NetAddr2IP(addr)
+	if err != nil {
+		s.logger.Warn("failed to check public endpoint", zap.Any("addr", addr), zap.Error(err))
+		return false
+	}
+	contains, err := netutil.CIDRContainsIP(publicEndpoints, ip)
+	if err != nil {
+		s.logger.Warn("failed to check public endpoint", zap.Any("ip", ip), zap.Error(err))
+		return false
+	}
+	if contains {
+		return true
+	}
+	// The public NLB may enable preserveIP, and the incoming address is the client address, which may be a public address.
+	// Even if the private NLB enables preserveIP, the client address is still a private address.
+	return !netutil.IsPrivate(ip)
+}
+
 func (s *SQLServer) PreClose() {
 	// Step 1: HTTP status returns unhealthy so that NLB takes this instance offline and then new connections won't come.
 	s.mu.Lock()

pkg/proxy/proxy_test.go

@@ -25,6 +25,7 @@ import (
 	"github.com/pingcap/tiproxy/pkg/proxy/client"
 	pnet "github.com/pingcap/tiproxy/pkg/proxy/net"
 	"github.com/stretchr/testify/require"
+	"go.uber.org/zap"
 )
 
 func TestCreateConn(t *testing.T) {
@@ -285,6 +286,44 @@ func TestRecoverPanic(t *testing.T) {
 	certManager.Close()
 }
 
+func TestPublicEndpoint(t *testing.T) {
+	tests := []struct {
+		publicEndpoints []string
+		publicIps       []string
+		privateIps      []string
+	}{
+		{
+			publicIps:  []string{"137.84.2.178"},
+			privateIps: []string{"10.10.10.10"},
+		},
+		{
+			publicEndpoints: []string{"10.10.10.0/24"},
+			publicIps:       []string{"137.84.2.178", "10.10.10.10"},
+			privateIps:      []string{"10.10.20.10"},
+		},
+		{
+			publicEndpoints: []string{"10.10.10.0/24", "10.10.20.10"},
+			publicIps:       []string{"137.84.2.178", "10.10.10.10", "10.10.20.10"},
+			privateIps:      []string{"10.10.20.11"},
+		},
+	}
+
+	server, err := NewSQLServer(zap.NewNop(), &config.Config{}, nil, id.NewIDManager(), nil, nil, backend.NewDefaultHandshakeHandler(nil))
+	require.NoError(t, err)
+	for i, test := range tests {
+		cfg := &config.Config{}
+		cfg.Proxy.PublicEndpoints = test.publicEndpoints
+		server.reset(cfg)
+		for j, ip := range test.publicIps {
+			require.True(t, server.fromPublicEndpoint(&net.TCPAddr{IP: net.ParseIP(ip), Port: 1000}), "test %d %d", i, j)
+		}
+		for j, ip := range test.privateIps {
+			require.False(t, server.fromPublicEndpoint(&net.TCPAddr{IP: net.ParseIP(ip), Port: 1000}), "test %d %d", i, j)
+		}
+		require.False(t, server.fromPublicEndpoint(nil))
+	}
+}
+
 type mockHsHandler struct {
 	backend.DefaultHandshakeHandler
 	handshakeResp func(ctx backend.ConnContext, _ *pnet.HandshakeResp) error

pkg/util/netutil/netutil.go

@@ -28,18 +28,25 @@ func ParseCIDRList(strList []string) ([]*net.IPNet, error) {
 	return cidrList, parseErr
 }
 
-func CIDRContainsIP(cidrList []*net.IPNet, addr net.Addr) (bool, error) {
+func NetAddr2IP(addr net.Addr) (net.IP, error) {
 	if addr == nil || reflect.ValueOf(addr).IsNil() {
-		return false, errors.New("address is nil")
+		return nil, errors.New("address is nil")
 	}
 	value := addr.String()
 	ipStr, _, err := net.SplitHostPort(value)
 	if err != nil {
-		return false, errors.Wrapf(err, "failed to parse address '%s'", value)
+		return nil, errors.Wrapf(err, "failed to parse address '%s'", value)
 	}
 	ip := net.ParseIP(ipStr)
 	if ip == nil {
-		return false, errors.Errorf("failed to parse IP '%s'", value)
+		return nil, errors.Errorf("failed to parse IP '%s'", value)
+	}
+	return ip, nil
+}
+
+func CIDRContainsIP(cidrList []*net.IPNet, ip net.IP) (bool, error) {
+	if ip == nil || reflect.ValueOf(ip).IsNil() {
+		return false, errors.New("address is nil")
 	}
 	for _, cidr := range cidrList {
 		if cidr.Contains(ip) {
@@ -48,3 +55,15 @@ func CIDRContainsIP(cidrList []*net.IPNet, addr net.Addr) (bool, error) {
 	}
 	return false, nil
 }
+
+func IsPrivate(ip net.IP) bool {
+	if ip.IsLoopback() || ip.IsLinkLocalUnicast() || ip.IsLinkLocalMulticast() {
+		return true
+	}
+	if ipv4 := ip.To4(); ipv4 != nil {
+		if ipv4[0] == 100 && ipv4[1] >= 64 && ipv4[1] <= 127 {
+			return true
+		}
+	}
+	return ip.IsPrivate()
+}

pkg/util/netutil/netutil_test.go

@@ -89,7 +89,60 @@ func TestContainIP(t *testing.T) {
 	for _, test := range tests {
 		list, err := ParseCIDRList(test.cidrs)
 		require.NoError(t, err, "ip: %s, cidrs: %v", test.ip, test.cidrs)
-		contain, _ := CIDRContainsIP(list, &net.TCPAddr{IP: net.ParseIP(test.ip), Port: 1000})
+		contain, _ := CIDRContainsIP(list, net.ParseIP(test.ip))
 		require.Equal(t, test.success, contain, "ip: %s, cidrs: %v", test.ip, test.cidrs)
 	}
 }
+
+func TestIsPrivate(t *testing.T) {
+	tests := []struct {
+		ip      string
+		private bool
+	}{
+		{
+			ip:      "8.8.8.8",
+			private: false,
+		},
+		{
+			ip:      "192.168.1.1",
+			private: true,
+		},
+		{
+			ip:      "172.16.0.1",
+			private: true,
+		},
+		{
+			ip:      "10.0.0.1",
+			private: true,
+		},
+		{
+			ip:      "127.0.0.1",
+			private: true,
+		},
+		{
+			ip:      "::1",
+			private: true,
+		},
+		{
+			ip:      "169.254.1.1",
+			private: true,
+		},
+		{
+			ip:      "100.64.1.1",
+			private: true,
+		},
+		{
+			ip:      "2001:4860:4860::8888",
+			private: false,
+		},
+		{
+			ip:      "fc00::1",
+			private: true,
+		},
+	}
+
+	for _, test := range tests {
+		ip := net.ParseIP(test.ip)
+		require.Equal(t, test.private, IsPrivate(ip), "ip: %s", test.ip)
+	}
+}