Merge pull request #208 Director SetHeadersByIP from egorgasay/master

README.md

@@ -53,8 +53,6 @@ Use --help key for details:
 ==================
 Сайт программы: https://github.com/rekby/lets-proxy2
 
-Сейчас это тестовая версия, программа в процессе разработчик и она не готова для реального использования.
-
 Возможности:
 * Авторизация доменов по протоколам http-01 and tls-alpn-01
 * Проксирование HTTPS (с автовыпуском сертификата) and HTTP

cmd/static/default-config.toml

@@ -101,6 +101,17 @@ 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.
+# Example:
+# [Proxy.HeadersByIP]
+#"192.168.0.0/24" = ["IPLocal:Test"]
+#"fe80:0000:0000:0000::/64" = [
+#   "IPLocal:Test2",
+#   "Local:False"
+# ]
+
 # Use https requests to backend instead of http
 HTTPSBackend = false
 

internal/proxy/config.go

@@ -17,11 +17,17 @@ import (
 
 const defaultHTTPPort = 80
 
+type IPHeaders struct {
+	Headers map[string]string
+	IP      string
+}
+
 //nolint:lll
 type Config struct {
 	DefaultTarget           string
 	TargetMap               []string
 	Headers                 []string
+	HeadersByIP             map[string][]string
 	KeepAliveTimeoutSeconds int
 	HTTPSBackend            bool
 	HTTPSBackendIgnoreCert  bool
@@ -57,6 +63,7 @@ func (c *Config) Apply(ctx context.Context, p *HTTPProxy) error {
 	appendDirector(c.getMapDirector)
 	appendDirector(c.getHeadersDirector)
 	appendDirector(c.getSchemaDirector)
+	appendDirector(c.getHeadersByIPDirector)
 	p.HTTPTransport = Transport{
 		IgnoreHTTPSCertificate: c.HTTPSBackendIgnoreCert,
 		RateLimiter:            rateLimiter,
@@ -157,6 +164,83 @@ func (c *Config) getSchemaDirector(ctx context.Context) (Director, error) {
 	return NewSetSchemeDirector(ProtocolHTTP), nil
 }
 
+// getHeadersByIPDirector transform array to DirectorSetHeadersByIP
+// can return nil,nil
+// example:
+//
+// [Proxy.HeadersByIP]
+// "192.168.1.0/24" = [
+//
+//		"User-Agent:PostmanRuntime/7.29.2",
+//		"Accept:*/*",
+//		"Accept-Encoding:gzip, deflate, br",
+//	]
+//	"192.168.132.0/30" = [
+//		"Accept-Encoding:gzip",
+//
+// ]
+//
+// out:
+//
+//	DirectorSetHeadersByIP {
+//		HTTPHeader{
+//			IPNet: 192.168.1.0/24,
+//			Headers: []HTTPHeader{
+//				HTTPHeader{
+//					Name: "User-Agent",
+//					Value: "PostmanRuntime/7.29.2",
+//				},
+//				HTTPHeader{
+//					Name: "Accept": "*/*",
+//					Value: "Accept-Encoding": "gzip, deflate, br",
+//				},
+//			},
+//		},
+//		HTTPHeader{
+//	    	IPNet:  192.168.132.0/30,
+//			Headers: []HTTPHeader{
+//				HTTPHeader{
+//					Name: "Accept-Encoding",
+//					Value: "gzip",
+//				},
+//			},
+//		},
+//	}
+func (c *Config) getHeadersByIPDirector(ctx context.Context) (Director, error) {
+	logger := zc.L(ctx)
+
+	if len(c.HeadersByIP) == 0 {
+		return nil, nil
+	}
+
+	m := make(map[string]HTTPHeaders)
+	for ipNet, headers := range c.HeadersByIP {
+		ipNet = strings.TrimSpace(ipNet)
+		for _, header := range headers {
+			lineParts := strings.SplitN(header, ":", 2)
+			if len(lineParts) < 2 {
+				logger.Error("Can't split header line to parts", zap.String("line", header))
+				return nil, errors.New("can't parse headers proxy config")
+			}
+
+			name := lineParts[0]
+			value := lineParts[1]
+
+			if m[ipNet] == nil {
+				m[ipNet] = make(HTTPHeaders, 0)
+			}
+
+			m[ipNet] = append(m[ipNet], HTTPHeader{
+				Name:  name,
+				Value: value,
+			})
+		}
+	}
+
+	logger.Info("Create headers by ip director", zap.Any("headers", m))
+	return NewDirectorSetHeadersByIP(m)
+}
+
 func parseTCPMapPair(line string) (from, to string, err error) {
 	line = strings.TrimSpace(line)
 	lineParts := strings.Split(line, "-")

internal/proxy/config_test.go

@@ -1,6 +1,7 @@
 package proxy
 
 import (
+	"net"
 	"testing"
 
 	"github.com/rekby/lets-proxy2/internal/th"
@@ -290,3 +291,175 @@ func TestConfig_Apply(t *testing.T) {
 	transport = p.HTTPTransport.(Transport)
 	transport.IgnoreHTTPSCertificate = true
 }
+
+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",
+			c: Config{
+				HeadersByIP: map[string][]string{
+					"192.168.1.0/24": {
+						"User-Agent:PostmanRuntime/7.29.2",
+						"Accept:*/*",
+						"Accept-Encoding:gzip, deflate, br",
+					},
+				},
+			},
+			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",
+			c: Config{
+				HeadersByIP: map[string][]string{
+					"192.168.1.0/24": {
+						"User-AgentPostmanRuntime/7.29.2",
+						"Accept:*/*",
+						"Accept-Encoding:gzip, deflate, br",
+					},
+				},
+			},
+			want:    DirectorSetHeadersByIP{},
+			wantErr: true,
+		},
+		{
+			name: "5Networks",
+			c: Config{
+				HeadersByIP: map[string][]string{
+					"10.0.0.0/8": {
+						"User-Agent:PostmanRuntime/7.29.2",
+						"Accept:*/*",
+						"Accept-Encoding:gzip, deflate, br",
+					},
+					"10.0.0.0/24": {
+						"Connection:Keep-Alive",
+						"Upgrade-Insecure-Requests:1",
+						"Cache-Control:no-cache",
+					},
+					"10.0.1.0/24": {
+						"Origin:https://example.com",
+						"Content-Type:application/json",
+						"Content-Length:123",
+					},
+
+					"10.2.0.0/24": {
+						"Accept-Encoding:gzip, deflate, br",
+						"Accept-Language:en-US,en;q=0.9",
+					},
+					"fe80:0000:0000:0000::/64": {
+						"Accept:*/*",
+						"Accept-Encoding:gzip, deflate, br",
+						"Accept-Language:en-US,en;q=0.9",
+					},
+				},
+			},
+			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 {
+		t.Run(tt.name, func(t *testing.T) {
+			got, err := tt.c.getHeadersByIPDirector(ctx)
+			if (err != nil) != tt.wantErr {
+				t.Fatalf("getHeadersByIPDirector() error = %v, wantErr %v", err, tt.wantErr)
+			}
+
+			if tt.wantErr {
+				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
+					}
+				}
+
+				t.Logf("getMapDir() got = %v", mp)
+				return mp
+			}
+
+			if got == nil && tt.want == nil {
+				return
+			}
+
+			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)
+				}
+			}
+		})
+	}
+}