90 lines
2.7 KiB
Go
90 lines
2.7 KiB
Go
package wrappermetadata
|
|
|
|
import (
|
|
"context"
|
|
|
|
"go.unistack.org/micro/v4/client"
|
|
"go.unistack.org/micro/v4/metadata"
|
|
"go.unistack.org/micro/v4/options"
|
|
"go.unistack.org/micro/v4/server"
|
|
)
|
|
|
|
type wrapper struct {
|
|
client.Client
|
|
exclude []string
|
|
}
|
|
|
|
func NewClientWrapper(exclude ...string) client.Wrapper {
|
|
return func(c client.Client) client.Client {
|
|
handler := &wrapper{
|
|
Client: c,
|
|
exclude: exclude,
|
|
}
|
|
return handler
|
|
}
|
|
}
|
|
|
|
func NewClientCallWrapper(exclude ...string) client.CallWrapper {
|
|
return func(fn client.CallFunc) client.CallFunc {
|
|
return func(ctx context.Context, addr string, req client.Request, rsp interface{}, opts client.CallOptions) error {
|
|
if exclude == nil {
|
|
return fn(ctx, addr, req, rsp, opts)
|
|
}
|
|
if imd, iok := metadata.FromIncomingContext(ctx); iok && imd != nil {
|
|
omd, ook := metadata.FromOutgoingContext(ctx)
|
|
if !ook || omd == nil {
|
|
omd = metadata.New(len(imd))
|
|
}
|
|
ctx = metadata.NewOutgoingContext(ctx, metadata.Merge(omd, metadata.Copy(imd, exclude...), false))
|
|
}
|
|
return fn(ctx, addr, req, rsp, opts)
|
|
}
|
|
}
|
|
}
|
|
|
|
func (w *wrapper) Call(ctx context.Context, req client.Request, rsp interface{}, opts ...options.Option) error {
|
|
if w.exclude == nil {
|
|
return w.Client.Call(ctx, req, rsp, opts...)
|
|
}
|
|
if imd, iok := metadata.FromIncomingContext(ctx); iok && imd != nil {
|
|
omd, ook := metadata.FromOutgoingContext(ctx)
|
|
if !ook || omd == nil {
|
|
omd = metadata.New(len(imd))
|
|
}
|
|
ctx = metadata.NewOutgoingContext(ctx, metadata.Merge(omd, metadata.Copy(imd, w.exclude...), false))
|
|
}
|
|
return w.Client.Call(ctx, req, rsp, opts...)
|
|
}
|
|
|
|
func (w *wrapper) Stream(ctx context.Context, req client.Request, opts ...options.Option) (client.Stream, error) {
|
|
if w.exclude == nil {
|
|
return w.Client.Stream(ctx, req, opts...)
|
|
}
|
|
if imd, iok := metadata.FromIncomingContext(ctx); iok && imd != nil {
|
|
omd, ook := metadata.FromOutgoingContext(ctx)
|
|
if !ook || omd == nil {
|
|
omd = metadata.New(len(imd))
|
|
}
|
|
ctx = metadata.NewOutgoingContext(ctx, metadata.Merge(omd, metadata.Copy(imd, w.exclude...), false))
|
|
}
|
|
return w.Client.Stream(ctx, req, opts...)
|
|
}
|
|
|
|
func NewServerHandlerWrapper(exclude ...string) server.HandlerWrapper {
|
|
return func(fn server.HandlerFunc) server.HandlerFunc {
|
|
return func(ctx context.Context, req server.Request, rsp interface{}) error {
|
|
if exclude == nil {
|
|
return fn(ctx, req, rsp)
|
|
}
|
|
if imd, iok := metadata.FromIncomingContext(ctx); iok && imd != nil {
|
|
omd, ook := metadata.FromOutgoingContext(ctx)
|
|
if !ook || omd == nil {
|
|
omd = metadata.New(len(imd))
|
|
}
|
|
ctx = metadata.NewOutgoingContext(ctx, metadata.Merge(omd, metadata.Copy(imd, exclude...), false))
|
|
}
|
|
return fn(ctx, req, rsp)
|
|
}
|
|
}
|
|
}
|