feat: middleware seletor add ctx match (#1652)

* feat: middleware seletor add headers match

* fix

* add some annotation

* add ctx seletor

* fix match function

* update match

middleware/selector/selector.go

@@ -11,7 +11,7 @@ import (
 
 type (
 	transporter func(ctx context.Context) (transport.Transporter, bool)
-	MatchFunc   func(operation string) bool
+	MatchFunc   func(ctx context.Context, operation string) bool
 )
 
 var (
@@ -83,7 +83,13 @@ func (b *Builder) Build() middleware.Middleware {
 }
 
 // matchs is match operation compliance Builder
-func (b *Builder) matchs(operation string) bool {
+func (b *Builder) matchs(ctx context.Context, transporter transporter) bool {
+	info, ok := transporter(ctx)
+	if !ok {
+		return false
+	}
+
+	operation := info.Operation()
 	for _, prefix := range b.prefix {
 		if prefixMatch(prefix, operation) {
 			return true
@@ -101,21 +107,19 @@ func (b *Builder) matchs(operation string) bool {
 	}
 
 	if b.match != nil {
-		return b.match(operation)
+		if b.match(ctx, operation) {
+			return true
+		}
 	}
+
 	return false
 }
 
 // selector middleware
-func selector(transporter transporter, match MatchFunc, ms ...middleware.Middleware) middleware.Middleware {
+func selector(transporter transporter, match func(context.Context, transporter) bool, ms ...middleware.Middleware) middleware.Middleware {
 	return func(handler middleware.Handler) middleware.Handler {
 		return func(ctx context.Context, req interface{}) (reply interface{}, err error) {
-			info, ok := transporter(ctx)
-			if !ok {
-				return handler(ctx, req)
-			}
-
-			if !match(info.Operation()) {
+			if !match(ctx, transporter) {
 				return handler(ctx, req)
 			}
 			return middleware.Chain(ms...)(handler)(ctx, req)

middleware/selector/selector_test.go

@@ -17,6 +17,7 @@ type Transport struct {
 	kind      transport.Kind
 	endpoint  string
 	operation string
+	headers   *mockHeader
 }
 
 func (tr *Transport) Kind() transport.Kind {
@@ -32,13 +33,33 @@ func (tr *Transport) Operation() string {
 }
 
 func (tr *Transport) RequestHeader() transport.Header {
-	return nil
+	return tr.headers
 }
 
 func (tr *Transport) ReplyHeader() transport.Header {
 	return nil
 }
 
+type mockHeader struct {
+	m map[string]string
+}
+
+func (m *mockHeader) Get(key string) string {
+	return m.m[key]
+}
+
+func (m *mockHeader) Set(key, value string) {
+	m.m[key] = value
+}
+
+func (m *mockHeader) Keys() []string {
+	keys := make([]string, 0, len(m.m))
+	for k := range m.m {
+		keys = append(keys, k)
+	}
+	return keys
+}
+
 func TestMatch(t *testing.T) {
 	tests := []struct {
 		name string
@@ -139,7 +160,7 @@ func TestFunc(t *testing.T) {
 				t.Log(req)
 				return "reply", nil
 			}
-			next = Server(testMiddleware).Match(func(operation string) bool {
+			next = Server(testMiddleware).Match(func(ctx context.Context, operation string) bool {
 				if strings.HasPrefix(operation, "/go-kratos.dev") || strings.HasSuffix(operation, "world") {
 					return true
 				}
@@ -152,6 +173,66 @@ func TestFunc(t *testing.T) {
 	}
 }
 
+func TestHeaderFunc(t *testing.T) {
+	tests := []struct {
+		name string
+		ctx  context.Context
+	}{
+		{
+			name: "/hello.Update/world",
+			ctx: transport.NewServerContext(context.Background(), &Transport{
+				operation: "/hello.Update/world",
+				headers:   &mockHeader{map[string]string{"X-Test": "test"}},
+			}),
+		},
+		{
+			name: "/hi.Create/world",
+			ctx: transport.NewServerContext(context.Background(), &Transport{
+				operation: "/hi.Create/world",
+				headers:   &mockHeader{map[string]string{"X-Test": "test2", "go-kratos": "kratos"}},
+			}),
+		},
+		{
+			name: "/test.Name/1234",
+			ctx: transport.NewServerContext(context.Background(), &Transport{
+				operation: "/test.Name/1234",
+				headers:   &mockHeader{map[string]string{"X-Test": "test3"}},
+			}),
+		},
+		{
+			name: "/go-kratos.dev/kratos",
+			ctx: transport.NewServerContext(context.Background(), &Transport{
+				operation: "/go-kratos.dev/kratos",
+				headers:   &mockHeader{map[string]string{"X-Test": "test"}},
+			}),
+		},
+	}
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			next := func(ctx context.Context, req interface{}) (interface{}, error) {
+				t.Log(req)
+				return "reply", nil
+			}
+			next = Server(testMiddleware).Match(func(ctx context.Context, operation string) bool {
+				tr, ok := transport.FromServerContext(ctx)
+				if !ok {
+					return false
+				}
+				if tr.RequestHeader().Get("X-Test") == "test" {
+					return true
+				}
+				if tr.RequestHeader().Get("go-kratos") == "kratos" {
+					return true
+				}
+				return false
+			}).Build()(next)
+			reply, err := next(test.ctx, test.name)
+			assert.Equal(t, reply, "reply")
+			assert.Nil(t, err)
+		})
+	}
+}
+
 func testMiddleware(handler middleware.Handler) middleware.Handler {
 	return func(ctx context.Context, req interface{}) (reply interface{}, err error) {
 		fmt.Println("before")