You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
116 lines
2.5 KiB
116 lines
2.5 KiB
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
|
|
}
|
|
|
|
func (o *options) hasPrefix(key string) bool {
|
|
k := strings.ToLower(key)
|
|
for _, prefix := range o.prefix {
|
|
if strings.HasPrefix(k, prefix) {
|
|
return true
|
|
}
|
|
}
|
|
return false
|
|
}
|
|
|
|
// 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) {
|
|
tr, ok := transport.FromServerContext(ctx)
|
|
if !ok {
|
|
return handler(ctx, req)
|
|
}
|
|
|
|
md := options.md.Clone()
|
|
header := tr.RequestHeader()
|
|
for _, k := range header.Keys() {
|
|
if options.hasPrefix(k) {
|
|
for _, v := range header.Values(k) {
|
|
md.Add(k, v)
|
|
}
|
|
}
|
|
}
|
|
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) {
|
|
tr, ok := transport.FromClientContext(ctx)
|
|
if !ok {
|
|
return handler(ctx, req)
|
|
}
|
|
|
|
header := tr.RequestHeader()
|
|
// x-md-local-
|
|
for k, vList := range options.md {
|
|
for _, v := range vList {
|
|
header.Add(k, v)
|
|
}
|
|
}
|
|
if md, ok := metadata.FromClientContext(ctx); ok {
|
|
for k, vList := range md {
|
|
for _, v := range vList {
|
|
header.Add(k, v)
|
|
}
|
|
}
|
|
}
|
|
// x-md-global-
|
|
if md, ok := metadata.FromServerContext(ctx); ok {
|
|
for k, vList := range md {
|
|
if options.hasPrefix(k) {
|
|
for _, v := range vList {
|
|
header.Add(k, v)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
return handler(ctx, req)
|
|
}
|
|
}
|
|
}
|
|
|