diff --git a/go.mod b/go.mod new file mode 100644 index 0000000..cdebae8 --- /dev/null +++ b/go.mod @@ -0,0 +1,5 @@ +module go.unistack.org/micro-wrapper-metadata/v4 + +go 1.20 + +require go.unistack.org/micro/v4 v4.0.17 diff --git a/go.sum b/go.sum new file mode 100644 index 0000000..90a23b8 --- /dev/null +++ b/go.sum @@ -0,0 +1,2 @@ +go.unistack.org/micro/v4 v4.0.17 h1:mF7uM+J4ILdG+1fcwzKYCwDlxhdbF/e1WnGzKKLnIXc= +go.unistack.org/micro/v4 v4.0.17/go.mod h1:ZDgU9931vm2l7X6RN/6UuwRIVp24GRdmQ7dKmegArk4= diff --git a/metadata.go b/metadata.go new file mode 100644 index 0000000..f1d5482 --- /dev/null +++ b/metadata.go @@ -0,0 +1,89 @@ +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) + } + } +}