Repository: go-kratos/kratos
Pull request #1244 (2021-07-25 06:29:19)

middleware/selector
#1241

zwhyb:  需要支持client不？
shenqidebaozi: 要的
shenqidebaozi: 路由声明是否要支持数组类型？
zwhyb: 添加了client的middleware 用断言的方式支持数组匹配了 如果继续用字符串，知识一个类型断言的影响 性能影响应该很小
tonybase: if o, ok := err.(interface {
	Operation() string
}); ok {
	op = o.Operation()
}

好像可以直接通过匿名方法获取？
那个interface{}不太好，应该强类型，要是不用链式使用第一个版本时也可以。
tonybase: ```
http.Middleware(
    selector.Client(logging.Server(), authjwt.Server())
        .Prefix("/article.", "/login")
        .Regex("/article.{List*}"),
        .Full("/article.ListCategory"),
)

http.Middleware(
    selector.Server(logging.Server(), authjwt.Server())
        .Prefix("/article.", "/login")
        .Regex("/article.{List*}"),
        .Full("/article.ListCategory"),
)
```

Patch:
diff --git a/config/env/env.go b/config/env/env.go
index f8da4110..c7d018b6 100644
--- a/config/env/env.go
+++ b/config/env/env.go
@@ -31,7 +31,7 @@ func (e *env) load(envStrings []string) []*config.KeyValue {
 
 		if len(e.prefixs) > 0 {
 			p, ok := matchPrefix(e.prefixs, envstr)
-			if !ok {
+			if !ok || len(p) == len(k) {
 				continue
 			}
 			// trim prefix
@@ -41,10 +41,12 @@ func (e *env) load(envStrings []string) []*config.KeyValue {
 			}
 		}
 
-		kv = append(kv, &config.KeyValue{
-			Key:   k,
-			Value: []byte(v),
-		})
+		if len(k) != 0 {
+			kv = append(kv, &config.KeyValue{
+				Key:   k,
+				Value: []byte(v),
+			})
+		}
 	}
 	return kv
 }
diff --git a/middleware/selector/selector.go b/middleware/selector/selector.go
new file mode 100644
index 00000000..7c216392
--- /dev/null
+++ b/middleware/selector/selector.go
@@ -0,0 +1,118 @@
+package selector
+
+import (
+	"context"
+	"github.com/go-kratos/kratos/v2/middleware"
+	"github.com/go-kratos/kratos/v2/transport"
+	"regexp"
+	"strings"
+)
+
+type (
+	transporter func(ctx context.Context) (transport.Transporter, bool)
+	match       func(operation string) bool
+)
+
+var (
+	serverTransporter transporter = func(ctx context.Context) (transport.Transporter, bool) {
+		return transport.FromServerContext(ctx)
+	}
+	clientTransporter transporter = func(ctx context.Context) (transport.Transporter, bool) {
+		return transport.FromClientContext(ctx)
+	}
+)
+
+type Builder struct {
+	client bool
+
+	prefix []string
+	regex  []string
+	path   []string
+
+	ms []middleware.Middleware
+}
+
+func Server(ms ...middleware.Middleware) *Builder {
+	return &Builder{ms: ms}
+}
+
+func Client(ms ...middleware.Middleware) *Builder {
+	return &Builder{client: true, ms: ms}
+}
+
+func (b *Builder) Prefix(prefix ...string) *Builder {
+	b.prefix = prefix
+	return b
+}
+
+func (b *Builder) Regex(regex ...string) *Builder {
+	b.regex = regex
+	return b
+}
+func (b *Builder) Path(path ...string) *Builder {
+	b.path = path
+	return b
+}
+
+func (b *Builder) Build() middleware.Middleware {
+	var transporter func(ctx context.Context) (transport.Transporter, bool)
+	if b.client {
+		transporter = clientTransporter
+	} else {
+		transporter = serverTransporter
+	}
+
+	return selector(transporter, b.match, b.ms...)
+}
+
+func (b *Builder) match(operation string) bool {
+	for _, prefix := range b.prefix {
+		if prefixMatch(prefix, operation) {
+			return true
+		}
+	}
+	for _, regex := range b.regex {
+		if regexMatch(regex, operation) {
+			return true
+		}
+	}
+	for _, path := range b.path {
+		if pathMatch(path, operation) {
+			return true
+		}
+	}
+	return false
+}
+
+func selector(transporter transporter, match match, 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()) {
+				return handler(ctx, req)
+			}
+			return middleware.Chain(ms...)(handler)(ctx, req)
+		}
+	}
+}
+
+func pathMatch(path string, operation string) bool {
+	return path == operation
+}
+
+func prefixMatch(prefix string, operation string) bool {
+	return strings.HasPrefix(operation, prefix)
+}
+
+func regexMatch(regex string, operation string) bool {
+	r, err := regexp.Compile(regex)
+	if err != nil {
+		return false
+	}
+	return r.FindString(operation) == operation
+}


Test patch:
diff --git a/config/env/env_test.go b/config/env/env_test.go
index 87f6def7..13861da8 100644
--- a/config/env/env_test.go
+++ b/config/env/env_test.go
@@ -49,6 +49,9 @@ func TestEnvWithPrefix(t *testing.T) {
 		prefix1 + "SERVICE_NAME": "kratos_app",
 		prefix2 + "ADDR":         "192.168.0.1",
 		prefix1 + "AGE":          "20",
+		// only prefix
+		prefix2:       "foo",
+		prefix2 + "_": "foo_",
 	}
 
 	for k, v := range envs {
diff --git a/middleware/selector/selector_test.go b/middleware/selector/selector_test.go
new file mode 100644
index 00000000..2f8ec532
--- /dev/null
+++ b/middleware/selector/selector_test.go
@@ -0,0 +1,118 @@
+package selector
+
+import (
+	"context"
+	"fmt"
+	"github.com/go-kratos/kratos/v2/middleware"
+	"github.com/go-kratos/kratos/v2/transport"
+	"testing"
+)
+
+var (
+	_ transport.Transporter = &Transport{}
+)
+
+type Transport struct {
+	kind      transport.Kind
+	endpoint  string
+	operation string
+}
+
+func (tr *Transport) Kind() transport.Kind {
+	return tr.kind
+}
+func (tr *Transport) Endpoint() string {
+	return tr.endpoint
+}
+func (tr *Transport) Operation() string {
+	return tr.operation
+}
+func (tr *Transport) RequestHeader() transport.Header {
+	return nil
+}
+func (tr *Transport) ReplyHeader() transport.Header {
+	return nil
+}
+
+func TestMatch(t *testing.T) {
+
+	tests := []struct {
+		name string
+		ctx  context.Context
+	}{
+		// TODO: Add test cases.
+		{
+			name: "/hello/world",
+			ctx:  transport.NewServerContext(context.Background(), &Transport{operation: "/hello/world"}),
+		},
+		{
+			name: "/hi/world",
+			ctx:  transport.NewServerContext(context.Background(), &Transport{operation: "/hi/world"}),
+		},
+		{
+			name: "/test/1234",
+			ctx:  transport.NewServerContext(context.Background(), &Transport{operation: "/test/1234"}),
+		},
+		{
+			name: "/example/kratos",
+			ctx:  transport.NewServerContext(context.Background(), &Transport{operation: "/example/kratos"}),
+		},
+	}
+	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).Prefix("/hello/").Regex(`/test/[0-9]+`).
+				Path("/example/kratos").Build()(next)
+			next(test.ctx, test.name)
+		})
+	}
+}
+
+func TestMatchClient(t *testing.T) {
+
+	tests := []struct {
+		name string
+		ctx  context.Context
+	}{
+		// TODO: Add test cases.
+		{
+			name: "/hello/world",
+			ctx:  transport.NewClientContext(context.Background(), &Transport{operation: "/hello/world"}),
+		},
+		{
+			name: "/hi/world",
+			ctx:  transport.NewClientContext(context.Background(), &Transport{operation: "/hi/world"}),
+		},
+		{
+			name: "/test/1234",
+			ctx:  transport.NewClientContext(context.Background(), &Transport{operation: "/test/1234"}),
+		},
+		{
+			name: "/example/kratos",
+			ctx:  transport.NewClientContext(context.Background(), &Transport{operation: "/example/kratos"}),
+		},
+	}
+	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 = Client(testMiddleware).Prefix("/hello/").Regex(`/test/[0-9]+`).
+				Path("/example/kratos").Build()(next)
+			next(test.ctx, test.name)
+		})
+	}
+}
+
+func testMiddleware(handler middleware.Handler) middleware.Handler {
+	return func(ctx context.Context, req interface{}) (reply interface{}, err error) {
+		fmt.Println("before")
+		reply, err = handler(ctx, req)
+		fmt.Println("after")
+		return
+	}
+}