/
selector.go
144 lines (122 loc) · 3.21 KB
/
selector.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
package selector
import (
"context"
"regexp"
"strings"
"github.com/go-kratos/kratos/v2/middleware"
"github.com/go-kratos/kratos/v2/transport"
)
type (
transporter func(ctx context.Context) (transport.Transporter, bool)
MatchFunc func(ctx context.Context, operation string) bool
)
var (
// serverTransporter is get server transport.Transporter from ctx
serverTransporter transporter = func(ctx context.Context) (transport.Transporter, bool) {
return transport.FromServerContext(ctx)
}
// clientTransporter is get client transport.Transporter from ctx
clientTransporter transporter = func(ctx context.Context) (transport.Transporter, bool) {
return transport.FromClientContext(ctx)
}
)
// Builder is a selector builder
type Builder struct {
client bool
prefix []string
regex []string
path []string
match MatchFunc
ms []middleware.Middleware
}
// Server selector middleware
func Server(ms ...middleware.Middleware) *Builder {
return &Builder{ms: ms}
}
// Client selector middleware
func Client(ms ...middleware.Middleware) *Builder {
return &Builder{client: true, ms: ms}
}
// Prefix is with Builder's prefix
func (b *Builder) Prefix(prefix ...string) *Builder {
b.prefix = prefix
return b
}
// Regex is with Builder's regex
func (b *Builder) Regex(regex ...string) *Builder {
b.regex = regex
return b
}
// Path is with Builder's path
func (b *Builder) Path(path ...string) *Builder {
b.path = path
return b
}
// Match is with Builder's match
func (b *Builder) Match(fn MatchFunc) *Builder {
b.match = fn
return b
}
// Build is Builder's Build, for example: Server().Path(m1,m2).Build()
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.matchs, b.ms...)
}
// matchs is match operation compliance Builder
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
}
}
for _, regex := range b.regex {
if regexMatch(regex, operation) {
return true
}
}
for _, path := range b.path {
if pathMatch(path, operation) {
return true
}
}
if b.match != nil {
if b.match(ctx, operation) {
return true
}
}
return false
}
// selector 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) {
if !match(ctx, transporter) {
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
}