feat: Add regex matching for request headers (#17)

lovesupergames · web-flow · commit 37d32a6c8f94 · 2025-07-11T01:06:55.000+04:00
* add regex MD

* refactor: Join header values for regex matching
diff --git a/pkg/discovery/discovery.go b/pkg/discovery/discovery.go
@@ -4,15 +4,14 @@ package discovery
 import (
 	"context"
 	"regexp"
-	"slices"
 	"strconv"
 	"strings"
 
+	"google.golang.org/grpc"
 	"google.golang.org/grpc/metadata"
 	"google.golang.org/grpc/status"
 	"google.golang.org/protobuf/proto"
 	"google.golang.org/protobuf/protoadapt"
-	"google.golang.org/grpc"
 )
 
 // Provider provides routing rules for the Service.
@@ -95,7 +94,8 @@ type RequestMatcher struct {
 	URI *regexp.Regexp
 
 	// IncomingMetadata contains the metadata of the incoming request.
-	IncomingMetadata metadata.MD
+	// The key is the metadata key, and the value is the regexp to match against the metadata value.
+	IncomingMetadata map[string]*regexp.Regexp
 
 	// Message contains the expected first RECV message of the request.
 	Message proto.Message
@@ -107,8 +107,14 @@ func (r RequestMatcher) Matches(uri string, md metadata.MD) bool {
 		return false
 	}
 
-	for k, v := range r.IncomingMetadata {
-		if !slices.Equal(v, md.Get(k)) {
+	for k, re := range r.IncomingMetadata {
+		vals := md.Get(k)
+		if len(vals) == 0 {
+			return false
+		}
+
+		// Join multiple values into a single string to allow matching against the whole.
+		if !re.MatchString(strings.Join(vals, ",")) {
 			return false
 		}
 	}
diff --git a/pkg/discovery/discovery_test.go b/pkg/discovery/discovery_test.go
@@ -14,10 +14,10 @@ func TestRule_String(t *testing.T) {
 	got := (&Rule{
 		Name: "name",
 		Match: RequestMatcher{
-			IncomingMetadata: metadata.New(map[string]string{
-				"key":  "value",
-				"key2": "value2",
-			}),
+			IncomingMetadata: map[string]*regexp.Regexp{
+				"key":  regexp.MustCompile("value"),
+				"key2": regexp.MustCompile("value2"),
+			},
 			Message: &errdetails.RequestInfo{
 				RequestId:   "request-id",
 				ServingData: "serving-data",
@@ -32,10 +32,30 @@ func TestRule_String(t *testing.T) {
 }
 
 func TestRequestMatcher_Matches(t *testing.T) {
-	rm := RequestMatcher{
-		URI:              regexp.MustCompile(`^/v1/example/.*$`),
-		IncomingMetadata: metadata.New(map[string]string{"key": "value"}),
-	}
-	assert.True(t, rm.Matches("/v1/example/123", metadata.New(map[string]string{"key": "value"})))
-	assert.False(t, rm.Matches("/v1/example/123", metadata.New(map[string]string{"key": "value2"})))
+	t.Run("single value", func(t *testing.T) {
+		rm := RequestMatcher{
+			URI:              regexp.MustCompile(`^/v1/example/.*$`),
+			IncomingMetadata: map[string]*regexp.Regexp{"key": regexp.MustCompile("^value.*$")},
+		}
+		assert.True(t, rm.Matches("/v1/example/123", metadata.New(map[string]string{"key": "value"})))
+		assert.True(t, rm.Matches("/v1/example/123", metadata.New(map[string]string{"key": "value123"})))
+		assert.False(t, rm.Matches("/v1/example/123", metadata.New(map[string]string{"key": "v"})))
+		assert.False(t, rm.Matches("/v1/example/123", metadata.New(map[string]string{"another": "value"})))
+		assert.False(t, rm.Matches("/v2/example/123", metadata.New(map[string]string{"key": "value"})))
+	})
+
+	t.Run("multiple values", func(t *testing.T) {
+		rm := RequestMatcher{
+			IncomingMetadata: map[string]*regexp.Regexp{"key": regexp.MustCompile("^v1,v2$")},
+		}
+
+		md := metadata.Pairs("key", "v1", "key", "v2")
+		assert.True(t, rm.Matches("any-uri", md))
+
+		md = metadata.Pairs("key", "v1")
+		assert.False(t, rm.Matches("any-uri", md))
+
+		md = metadata.Pairs("key", "v2", "key", "v1")
+		assert.False(t, rm.Matches("any-uri", md), "should not match on wrong order")
+	})
 }
diff --git a/pkg/discovery/fileprovider/file.go b/pkg/discovery/fileprovider/file.go
@@ -9,19 +9,20 @@ import (
 	"regexp"
 	"time"
 
+	"crypto/tls"
+	"sort"
+
 	"github.com/Semior001/groxy/pkg/discovery"
+	"github.com/Semior001/groxy/pkg/grpcx"
 	"github.com/Semior001/groxy/pkg/protodef"
 	"github.com/cappuccinotm/slogx"
+	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/credentials"
+	"google.golang.org/grpc/credentials/insecure"
 	"google.golang.org/grpc/metadata"
 	"google.golang.org/grpc/status"
 	"gopkg.in/yaml.v3"
-	"google.golang.org/grpc"
-	"google.golang.org/grpc/credentials/insecure"
-	"google.golang.org/grpc/credentials"
-	"crypto/tls"
-	"sort"
-	"github.com/Semior001/groxy/pkg/grpcx"
 )
 
 // File discovers the changes in routing rules from a file.
@@ -219,7 +220,16 @@ func parseRule(r Rule, upstreams []discovery.Upstream) (result discovery.Rule, e
 		return discovery.Rule{}, fmt.Errorf("compile URI regexp: %w", err)
 	}
 
-	result.Match.IncomingMetadata = metadata.New(r.Match.Header)
+	if len(r.Match.Header) > 0 {
+		result.Match.IncomingMetadata = make(map[string]*regexp.Regexp, len(r.Match.Header))
+		for k, v := range r.Match.Header {
+			re, err := regexp.Compile(v)
+			if err != nil {
+				return discovery.Rule{}, fmt.Errorf("compile header %q regexp: %w", k, err)
+			}
+			result.Match.IncomingMetadata[k] = re
+		}
+	}
 
 	if r.Match.Body != nil {
 		if result.Match.Message, err = protodef.BuildMessage(*r.Match.Body); err != nil {
diff --git a/pkg/discovery/fileprovider/file_test.go b/pkg/discovery/fileprovider/file_test.go
@@ -96,31 +96,31 @@ func TestFile_Rules(t *testing.T) {
 			Name: "com.github.Semior001.groxy.example.mock.ExampleService/Stub",
 			Match: discovery.RequestMatcher{
 				URI:              regexp.MustCompile("com.github.Semior001.groxy.example.mock.ExampleService/Stub"),
-				IncomingMetadata: map[string][]string{"test": {"true"}},
+				IncomingMetadata: map[string]*regexp.Regexp{"test": regexp.MustCompile("true")},
 			},
 			Mock: &discovery.Mock{},
 		},
 		{
 			Name: "com.github.Semior001.groxy.example.mock.ExampleService/Stub",
 			Match: discovery.RequestMatcher{
 				URI:              regexp.MustCompile("com.github.Semior001.groxy.example.mock.ExampleService/Stub"),
-				IncomingMetadata: metadata.New(nil),
+				IncomingMetadata: nil,
 			},
 			Mock: &discovery.Mock{},
 		},
 		{
 			Name: "com.github.Semior001.groxy.example.mock.ExampleService/Stub",
 			Match: discovery.RequestMatcher{
 				URI:              regexp.MustCompile("com.github.Semior001.groxy.example.mock.ExampleService/Stub"),
-				IncomingMetadata: metadata.New(nil),
+				IncomingMetadata: nil,
 			},
 			Mock: &discovery.Mock{},
 		},
 		{
 			Name: "com.github.Semior001.groxy.example.mock.ExampleService/Error",
 			Match: discovery.RequestMatcher{
 				URI:              regexp.MustCompile("com.github.Semior001.groxy.example.mock.ExampleService/Error"),
-				IncomingMetadata: metadata.New(nil),
+				IncomingMetadata: nil,
 			},
 			Mock: &discovery.Mock{
 				Status:  status.New(codes.InvalidArgument, "invalid request"),
@@ -132,7 +132,7 @@ func TestFile_Rules(t *testing.T) {
 			Name: "com.github.Semior001.groxy.example.mock.Upstream/Get",
 			Match: discovery.RequestMatcher{
 				URI:              regexp.MustCompile("com.github.Semior001.groxy.example.mock.Upstream/Get"),
-				IncomingMetadata: metadata.New(nil),
+				IncomingMetadata: nil,
 			},
 		},
 		{
diff --git a/pkg/discovery/service_test.go b/pkg/discovery/service_test.go
@@ -1,16 +1,17 @@
 package discovery
 
 import (
+	"context"
+	"errors"
 	"regexp"
 	"testing"
+	"time"
+
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"google.golang.org/genproto/googleapis/rpc/errdetails"
 	"google.golang.org/grpc/metadata"
 	"google.golang.org/protobuf/proto"
-	"context"
-	"time"
-	"errors"
 )
 
 func TestMatches_NeedsDeeperMatch(t *testing.T) {
@@ -34,15 +35,15 @@ func TestService_MatchMetadata(t *testing.T) {
 		rules: []*Rule{
 			{Name: "1", Match: RequestMatcher{
 				URI:              regexp.MustCompile("test"),
-				IncomingMetadata: metadata.New(map[string]string{"uri": "not-match"}),
+				IncomingMetadata: map[string]*regexp.Regexp{"uri": regexp.MustCompile("not-match")},
 			}},
 			{Name: "2", Match: RequestMatcher{
 				URI:              regexp.MustCompile("uri"),
-				IncomingMetadata: metadata.New(map[string]string{"uri": "test"}),
+				IncomingMetadata: map[string]*regexp.Regexp{"uri": regexp.MustCompile("test")},
 			}},
 			{Name: "3", Match: RequestMatcher{
 				URI:              regexp.MustCompile("uri"),
-				IncomingMetadata: metadata.New(map[string]string{"uri": "test"}),
+				IncomingMetadata: map[string]*regexp.Regexp{"uri": regexp.MustCompile("test")},
 			}},
 		},
 	}
@@ -95,7 +96,7 @@ func TestService_Run(t *testing.T) {
 			StateFunc: func(context.Context) (*State, error) {
 				return &State{Rules: []*Rule{
 					{Name: "1", Match: RequestMatcher{}},
-					{Name: "2", Match: RequestMatcher{IncomingMetadata: metadata.New(map[string]string{"uri": "test"})}},
+					{Name: "2", Match: RequestMatcher{IncomingMetadata: map[string]*regexp.Regexp{"uri": regexp.MustCompile("test")}}},
 				}}, nil
 			},
 		}
@@ -106,10 +107,10 @@ func TestService_Run(t *testing.T) {
 			},
 			StateFunc: func(context.Context) (*State, error) {
 				return &State{Rules: []*Rule{
-					{Name: "3", Match: RequestMatcher{IncomingMetadata: metadata.New(map[string]string{
-						"uri":  "test",
-						"uri2": "test2",
-					})}},
+					{Name: "3", Match: RequestMatcher{IncomingMetadata: map[string]*regexp.Regexp{
+						"uri":  regexp.MustCompile("test"),
+						"uri2": regexp.MustCompile("test2"),
+					}}},
 				}}, nil
 			},
 		}
@@ -129,11 +130,11 @@ func TestService_Run(t *testing.T) {
 		require.Error(t, err)
 		assert.Equal(t, context.DeadlineExceeded, err)
 		assert.Equal(t, []*Rule{
-			{Name: "3", Match: RequestMatcher{IncomingMetadata: metadata.New(map[string]string{
-				"uri":  "test",
-				"uri2": "test2",
-			})}},
-			{Name: "2", Match: RequestMatcher{IncomingMetadata: metadata.New(map[string]string{"uri": "test"})}},
+			{Name: "3", Match: RequestMatcher{IncomingMetadata: map[string]*regexp.Regexp{
+				"uri":  regexp.MustCompile("test"),
+				"uri2": regexp.MustCompile("test2"),
+			}}},
+			{Name: "2", Match: RequestMatcher{IncomingMetadata: map[string]*regexp.Regexp{"uri": regexp.MustCompile("test")}}},
 			{Name: "1", Match: RequestMatcher{}},
 		}, svc.rules)
 	})
@@ -151,7 +152,7 @@ func TestService_Run(t *testing.T) {
 			StateFunc: func(context.Context) (*State, error) {
 				return &State{Rules: []*Rule{
 					{Name: "1", Match: RequestMatcher{}},
-					{Name: "2", Match: RequestMatcher{IncomingMetadata: metadata.New(map[string]string{"uri": "test"})}},
+					{Name: "2", Match: RequestMatcher{IncomingMetadata: map[string]*regexp.Regexp{"uri": regexp.MustCompile("test")}}},
 				}}, nil
 			},
 		}
@@ -162,10 +163,10 @@ func TestService_Run(t *testing.T) {
 			},
 			StateFunc: func(context.Context) (*State, error) {
 				return &State{Rules: []*Rule{
-					{Name: "3", Match: RequestMatcher{IncomingMetadata: metadata.New(map[string]string{
-						"uri":  "test",
-						"uri2": "test2",
-					})}},
+					{Name: "3", Match: RequestMatcher{IncomingMetadata: map[string]*regexp.Regexp{
+						"uri":  regexp.MustCompile("test"),
+						"uri2": regexp.MustCompile("test2"),
+					}}},
 				}}, nil
 			},
 		}
