Merge pull request #210 Efficient search for networks including ip fr…

…om the request, from egorgasay/efficient_ipnet_search

cmd/static/default-config.toml

@@ -101,9 +101,9 @@ TargetMap = []
 # ["IP:{{SOURCE_IP}}", "Proxy:lets-proxy", "Protocol:{{HTTP_PROTO}}" ]
 Headers = [ "X-Forwarded-Proto:{{HTTP_PROTO}}", "X-Forwarded-For:{{SOURCE_IP}}" ]
 
-
 # A map with an IP key/mask and a value with an array of strings separated by a colon Header:Value
 # to add to a request with matching ip address for backend.
+# You can use General.IncludeConfigs for load rules from dedicated rules config file.
 # Example:
 # [Proxy.HeadersByIP]
 #"192.168.0.0/24" = ["IPLocal:Test"]

go.mod

@@ -23,6 +23,7 @@ require (
 )
 
 require (
+	github.com/egorgasay/cidranger v1.0.1
 	github.com/hashicorp/golang-lru/v2 v2.0.1
 	github.com/jonboulle/clockwork v0.4.0
 	github.com/letsencrypt/pebble/v2 v2.4.0
@@ -43,6 +44,7 @@ require (
 	github.com/prometheus/common v0.34.0 // indirect
 	github.com/prometheus/procfs v0.7.3 // indirect
 	github.com/valyala/fastrand v1.1.0 // indirect
+	github.com/yl2chen/cidranger v1.0.2 // indirect
 	go.uber.org/atomic v1.9.0 // indirect
 	go.uber.org/multierr v1.8.0 // indirect
 	golang.org/x/mod v0.8.0 // indirect

go.sum

@@ -60,6 +60,8 @@ github.com/cncf/udpa/go v0.0.0-20191209042840-269d4d468f6f/go.mod h1:M8M6+tZqaGX
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/egorgasay/cidranger v1.0.1 h1:hSU9Yihw8Tx1QqVgjKY4udAJfjC/hQquhZC+9OUDnec=
+github.com/egorgasay/cidranger v1.0.1/go.mod h1:UHMmoDd2MvfRfOM+9lhqeakZW/pKy7APt49TztmcFjo=
 github.com/envoyproxy/go-control-plane v0.9.0/go.mod h1:YTl/9mNaCwkRvm6d1a2C3ymFceY/DCBVvsKhRF0iEA4=
 github.com/envoyproxy/go-control-plane v0.9.1-0.20191026205805-5f8ba28d4473/go.mod h1:YTl/9mNaCwkRvm6d1a2C3ymFceY/DCBVvsKhRF0iEA4=
 github.com/envoyproxy/go-control-plane v0.9.4/go.mod h1:6rpuAdCZL397s3pYoYcLgu1mIlRU8Am5FuJP05cCM98=
@@ -233,6 +235,8 @@ github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/
 github.com/stretchr/testify v1.8.2 h1:+h33VjcLVPDHtOdpUCuF+7gSuG3yGIftsP1YvFihtJ8=
 github.com/valyala/fastrand v1.1.0 h1:f+5HkLW4rsgzdNoleUOB69hyT9IlD2ZQh9GyDMfb5G8=
 github.com/valyala/fastrand v1.1.0/go.mod h1:HWqCzkrkg6QXT8V2EXWvXCoow7vLwOFN002oeRzjapQ=
+github.com/yl2chen/cidranger v1.0.2 h1:lbOWZVCG1tCRX4u24kuM1Tb4nHqWkDxwLdoS+SevawU=
+github.com/yl2chen/cidranger v1.0.2/go.mod h1:9U1yz7WPYDwf0vpNWFaeRh0bjwz5RVgRy/9UEQfHl0g=
 github.com/yuin/goldmark v1.1.25/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=
 github.com/yuin/goldmark v1.1.27/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=
 github.com/yuin/goldmark v1.1.32/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=

internal/proxy/config_test.go

@@ -1,7 +1,11 @@
 package proxy
 
 import (
+	"fmt"
+	"github.com/egorgasay/cidranger"
 	"net"
+	"strconv"
+	"strings"
 	"testing"
 
 	"github.com/rekby/lets-proxy2/internal/th"
@@ -295,18 +299,15 @@ func TestConfig_Apply(t *testing.T) {
 func TestConfig_getHeadersByIPDirector(t *testing.T) {
 	ctx, flush := th.TestContext(t)
 	defer flush()
-	td := testdeep.NewT(t)
 
 	tests := []struct {
 		name    string
 		c       Config
-		want    DirectorSetHeadersByIP
 		wantErr bool
 	}{
 		{
 			name: "empty",
 			c:    Config{},
-			want: nil,
 		},
 		{
 			name: "oneNetwork",
@@ -319,16 +320,6 @@ func TestConfig_getHeadersByIPDirector(t *testing.T) {
 					},
 				},
 			},
-			want: DirectorSetHeadersByIP{
-				NetHeaders{
-					IPNet: net.IPNet{IP: net.ParseIP("192.168.1.0"), Mask: net.CIDRMask(24, 32)},
-					Headers: HTTPHeaders{
-						{Name: "User-Agent", Value: "PostmanRuntime/7.29.2"},
-						{Name: "Accept", Value: "*/*"},
-						{Name: "Accept-Encoding", Value: "gzip, deflate, br"},
-					},
-				},
-			},
 		},
 		{
 			name: "configError1",
@@ -341,19 +332,18 @@ func TestConfig_getHeadersByIPDirector(t *testing.T) {
 					},
 				},
 			},
-			want:    DirectorSetHeadersByIP{},
 			wantErr: true,
 		},
 		{
 			name: "5Networks",
 			c: Config{
 				HeadersByIP: map[string][]string{
-					"10.0.0.0/8": {
+					"11.0.0.0/8": {
 						"User-Agent:PostmanRuntime/7.29.2",
 						"Accept:*/*",
 						"Accept-Encoding:gzip, deflate, br",
 					},
-					"10.0.0.0/24": {
+					"10.55.0.0/24": {
 						"Connection:Keep-Alive",
 						"Upgrade-Insecure-Requests:1",
 						"Cache-Control:no-cache",
@@ -375,47 +365,6 @@ func TestConfig_getHeadersByIPDirector(t *testing.T) {
 					},
 				},
 			},
-			want: DirectorSetHeadersByIP{
-				NetHeaders{
-					IPNet: net.IPNet{IP: net.ParseIP("10.0.0.0"), Mask: net.CIDRMask(8, 32)},
-					Headers: HTTPHeaders{
-						{Name: "User-Agent", Value: "PostmanRuntime/7.29.2"},
-						{Name: "Accept", Value: "*/*"},
-						{Name: "Accept-Encoding", Value: "gzip, deflate, br"},
-					},
-				},
-				NetHeaders{
-					IPNet: net.IPNet{IP: net.ParseIP("10.0.0.0"), Mask: net.CIDRMask(24, 32)},
-					Headers: HTTPHeaders{
-						{Name: "Connection", Value: "Keep-Alive"},
-						{Name: "Upgrade-Insecure-Requests", Value: "1"},
-						{Name: "Cache-Control", Value: "no-cache"},
-					},
-				},
-				NetHeaders{
-					IPNet: net.IPNet{IP: net.ParseIP("10.0.1.0"), Mask: net.CIDRMask(24, 32)},
-					Headers: HTTPHeaders{
-						{Name: "Origin", Value: "https://example.com"},
-						{Name: "Content-Type", Value: "application/json"},
-						{Name: "Content-Length", Value: "123"},
-					},
-				},
-				NetHeaders{
-					IPNet: net.IPNet{IP: net.ParseIP("10.2.0.0"), Mask: net.CIDRMask(24, 32)},
-					Headers: HTTPHeaders{
-						{Name: "Accept-Encoding", Value: "gzip, deflate, br"},
-						{Name: "Accept-Language", Value: "en-US,en;q=0.9"},
-					},
-				},
-				NetHeaders{
-					IPNet: net.IPNet{IP: net.ParseIP("fe80:0000:0000:0000::"), Mask: net.CIDRMask(64, 128)},
-					Headers: HTTPHeaders{
-						{Name: "Accept", Value: "*/*"},
-						{Name: "Accept-Encoding", Value: "gzip, deflate, br"},
-						{Name: "Accept-Language", Value: "en-US,en;q=0.9"},
-					},
-				},
-			},
 		},
 	}
 	for _, tt := range tests {
@@ -427,39 +376,89 @@ func TestConfig_getHeadersByIPDirector(t *testing.T) {
 
 			if tt.wantErr {
 				return
+			} else if got == nil {
+				return
 			}
 
-			getMapDir := func(d DirectorSetHeadersByIP) map[string]map[string]string {
-				var mp = make(map[string]map[string]string)
-				for _, netHeaders := range d {
-					if netHeaders.Headers == nil {
-						continue
-					}
-					if _, ok := mp[netHeaders.IPNet.String()]; !ok {
-						mp[netHeaders.IPNet.String()] = make(map[string]string)
-					}
-					for _, header := range netHeaders.Headers {
-						mp[netHeaders.IPNet.String()][header.Name] = header.Value
-					}
+			ranger := got.(cidranger.Ranger[HTTPHeaders])
+			for network, headers := range tt.c.HeadersByIP {
+				_, ipnet, err := net.ParseCIDR(network)
+				if err != nil {
+					t.Fatalf("ParseCIDR error %v", err)
 				}
 
-				t.Logf("getMapDir() got = %v", mp)
-				return mp
-			}
+				ip, err := netToIP(ipnet)
+				if err != nil {
+					t.Fatalf("netToIP error %v", err)
+				}
 
-			if got == nil && tt.want == nil {
-				return
-			}
+				if ok, err := ranger.Contains(ip); err != nil {
+					t.Fatalf("contains error %v", err)
+				} else if !ok {
+					t.Fatalf("network %s not found", network)
+				}
+
+				gotHeaders := make([]string, 0, len(headers))
+
+				err = ranger.IterByIncomingNetworks(ip, func(network net.IPNet, h HTTPHeaders) error {
+					if headers == nil {
+						return nil
+					}
+
+					for _, header := range h {
+						gotHeaders = append(gotHeaders, fmt.Sprintf("%s:%s", header.Name, header.Value))
+					}
+					return nil
+				})
+				if err != nil {
+					t.Fatalf("IterByIncomingNetworks error %v", err)
+				}
 
-			if gotDir, ok := got.(DirectorSetHeadersByIP); !ok {
-				t.Fatalf("can't lead to the type %v", got)
-			} else {
-				gotMap := getMapDir(gotDir)
-				wantMap := getMapDir(tt.want)
-				if !td.CmpDeeply(gotMap, wantMap) {
-					t.Fatalf("getHeadersByIPDirector() got = %v, want %v", gotMap, wantMap)
+				if !isTheSameArray(headers, gotHeaders) {
+					t.Fatalf("want \n%v \ngot \n%v", headers, gotHeaders)
 				}
 			}
 		})
 	}
 }
+
+func isTheSameArray[T comparable](a, b []T) bool {
+	if len(a) != len(b) {
+		return false
+	}
+
+	tmp := make(map[T]struct{})
+	for _, el := range a {
+		tmp[el] = struct{}{}
+	}
+	for _, el := range b {
+		if _, ok := tmp[el]; !ok {
+			return false
+		}
+	}
+	return true
+}
+
+func netToIP(ipnet *net.IPNet) (net.IP, error) {
+	ip := ipnet.IP.String()
+
+	sep := ":"
+	if ipnet.IP.To4() != nil {
+		sep = "."
+	}
+
+	split := strings.Split(ip, sep)
+
+	if sep == ":" && split[len(split)-1] == "" {
+		return net.ParseIP(strings.Join(split, sep) + "1"), nil
+	}
+	num, err := strconv.Atoi(split[len(split)-1])
+	if err != nil {
+		return nil, err
+	}
+
+	num++
+	split[3] = strconv.Itoa(num)
+
+	return net.ParseIP(strings.Join(split, sep)), nil
+}