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
pull/1688/head
haiyux 3 years ago committed by GitHub
parent e1e8184768
commit 2e045c3e42
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
  1. 24
      middleware/selector/selector.go
  2. 85
      middleware/selector/selector_test.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)

@ -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")

Loading…
Cancel
Save