| // Copyright 2016 Michal Witkowski. All Rights Reserved. |
| // See LICENSE for licensing terms. |
| // WrappedServerStream is a thin wrapper around grpc.ServerStream that allows modifying context. |
| type WrappedServerStream struct { |
| // WrappedContext is the wrapper's own Context. You can assign it. |
| WrappedContext context.Context |
| // Context returns the wrapper's WrappedContext, overwriting the nested grpc.ServerStream.Context() |
| func (w *WrappedServerStream) Context() context.Context { |
| // WrapServerStream returns a ServerStream that has the ability to overwrite context. |
| func WrapServerStream(stream grpc.ServerStream) *WrappedServerStream { |
| if existing, ok := stream.(*WrappedServerStream); ok { |
| return &WrappedServerStream{ServerStream: stream, WrappedContext: stream.Context()} |