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 with constant metadata key value. func WithConstants(md metadata.Metadata) Option { return func(o *options) { o.md = md } } // WithPropagatedPrefix with propagated key prefix. func WithPropagatedPrefix(prefix ...string) Option { return func(o *options) { o.prefix = prefix } } // Server is middleware server-side metadata. func Server(opts ...Option) middleware.Middleware { options := options{ prefix: []string{"x-md-"}, // 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.New() for _, k := range tr.Header().Keys() { for _, prefix := range options.prefix { if strings.HasPrefix(strings.ToLower(k), 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 { // x-md-local- 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) } } // x-md-global- 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) } } }