Middleware/metadata v2 (#1050)
* add middleware metadata * add transport header * add transport carrier Co-authored-by: longXboy <longxboy@gmail.com>pull/1054/head
Tony Chen
4 years ago
committed by
GitHub
parent
7f2e3becbe
commit
8d8cd8c8a8
@ -0,0 +1,46 @@ |
|||||||
|
package metadata |
||||||
|
|
||||||
|
import ( |
||||||
|
"context" |
||||||
|
"fmt" |
||||||
|
) |
||||||
|
|
||||||
|
type serverMetadataKey struct{} |
||||||
|
|
||||||
|
// NewServerContext creates a new context with client md attached.
|
||||||
|
func NewServerContext(ctx context.Context, md Metadata) context.Context { |
||||||
|
return context.WithValue(ctx, serverMetadataKey{}, md) |
||||||
|
} |
||||||
|
|
||||||
|
// FromServerContext returns the server metadata in ctx if it exists.
|
||||||
|
func FromServerContext(ctx context.Context) (Metadata, bool) { |
||||||
|
md, ok := ctx.Value(serverMetadataKey{}).(Metadata) |
||||||
|
return md, ok |
||||||
|
} |
||||||
|
|
||||||
|
type clientMetadataKey struct{} |
||||||
|
|
||||||
|
// NewClientContext creates a new context with client md attached.
|
||||||
|
func NewClientContext(ctx context.Context, md Metadata) context.Context { |
||||||
|
return context.WithValue(ctx, clientMetadataKey{}, md) |
||||||
|
} |
||||||
|
|
||||||
|
// FromClientContext returns the client metadata in ctx if it exists.
|
||||||
|
func FromClientContext(ctx context.Context) (Metadata, bool) { |
||||||
|
md, ok := ctx.Value(clientMetadataKey{}).(Metadata) |
||||||
|
return md, ok |
||||||
|
} |
||||||
|
|
||||||
|
// AppendToClientContext returns a new context with the provided kv merged
|
||||||
|
// with any existing metadata in the context.
|
||||||
|
func AppendToClientContext(ctx context.Context, kv ...string) context.Context { |
||||||
|
if len(kv)%2 == 1 { |
||||||
|
panic(fmt.Sprintf("metadata: AppendToOutgoingContext got an odd number of input pairs for metadata: %d", len(kv))) |
||||||
|
} |
||||||
|
md, _ := FromClientContext(ctx) |
||||||
|
md = md.Clone() |
||||||
|
for i := 0; i < len(kv); i += 2 { |
||||||
|
md.Set(kv[i], kv[i+1]) |
||||||
|
} |
||||||
|
return NewClientContext(ctx, md) |
||||||
|
} |
||||||
@ -0,0 +1,95 @@ |
|||||||
|
package metadata |
||||||
|
|
||||||
|
import ( |
||||||
|
"context" |
||||||
|
"strings" |
||||||
|
|
||||||
|
"github.com/go-kratos/kratos/v2/metadata" |
||||||
|
"github.com/go-kratos/kratos/v2/middleware" |
||||||
|
"github.com/go-kratos/kratos/v2/transport" |
||||||
|
) |
||||||
|
|
||||||
|
// Option is metadata option.
|
||||||
|
type Option func(*options) |
||||||
|
|
||||||
|
type options struct { |
||||||
|
prefix []string |
||||||
|
md metadata.Metadata |
||||||
|
} |
||||||
|
|
||||||
|
// WithConstants is option with constant metadata key value.
|
||||||
|
func WithConstants(md metadata.Metadata) Option { |
||||||
|
return func(o *options) { |
||||||
|
o.md = md |
||||||
|
} |
||||||
|
} |
||||||
|
|
||||||
|
// WithPropagatedPrefix is option with global propagated key prefix.
|
||||||
|
func WithPropagatedPrefix(prefix ...string) Option { |
||||||
|
return func(o *options) { |
||||||
|
o.prefix = prefix |
||||||
|
} |
||||||
|
} |
||||||
|
|
||||||
|
// Server is middleware client-side metadata.
|
||||||
|
func Server(opts ...Option) middleware.Middleware { |
||||||
|
options := options{ |
||||||
|
prefix: []string{"x-md-global-", "x-md-local-"}, |
||||||
|
} |
||||||
|
for _, o := range opts { |
||||||
|
o(&options) |
||||||
|
} |
||||||
|
return func(handler middleware.Handler) middleware.Handler { |
||||||
|
return func(ctx context.Context, req interface{}) (reply interface{}, err error) { |
||||||
|
if tr, ok := transport.FromServerContext(ctx); ok { |
||||||
|
md := metadata.Metadata{} |
||||||
|
for _, k := range tr.Header().Keys() { |
||||||
|
key := strings.ToLower(k) |
||||||
|
for _, prefix := range options.prefix { |
||||||
|
if strings.HasPrefix(key, prefix) { |
||||||
|
md.Set(k, tr.Header().Get(k)) |
||||||
|
break |
||||||
|
} |
||||||
|
} |
||||||
|
} |
||||||
|
ctx = metadata.NewServerContext(ctx, md) |
||||||
|
} |
||||||
|
return handler(ctx, req) |
||||||
|
} |
||||||
|
} |
||||||
|
} |
||||||
|
|
||||||
|
// Client is middleware client-side metadata.
|
||||||
|
func Client(opts ...Option) middleware.Middleware { |
||||||
|
options := options{ |
||||||
|
prefix: []string{"x-md-global-"}, |
||||||
|
} |
||||||
|
for _, o := range opts { |
||||||
|
o(&options) |
||||||
|
} |
||||||
|
return func(handler middleware.Handler) middleware.Handler { |
||||||
|
return func(ctx context.Context, req interface{}) (reply interface{}, err error) { |
||||||
|
if tr, ok := transport.FromClientContext(ctx); ok { |
||||||
|
for k, v := range options.md { |
||||||
|
tr.Header().Set(k, v) |
||||||
|
} |
||||||
|
if md, ok := metadata.FromClientContext(ctx); ok { |
||||||
|
for k, v := range md { |
||||||
|
tr.Header().Set(k, v) |
||||||
|
} |
||||||
|
} |
||||||
|
if md, ok := metadata.FromServerContext(ctx); ok { |
||||||
|
for k, v := range md { |
||||||
|
for _, prefix := range options.prefix { |
||||||
|
if strings.HasPrefix(k, prefix) { |
||||||
|
tr.Header().Set(k, v) |
||||||
|
break |
||||||
|
} |
||||||
|
} |
||||||
|
} |
||||||
|
} |
||||||
|
} |
||||||
|
return handler(ctx, req) |
||||||
|
} |
||||||
|
} |
||||||
|
} |
||||||
Loading…
Reference in new issue