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
3 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