30 lines
945 B
Go
30 lines
945 B
Go
|
// Copyright 2016 Michal Witkowski. All Rights Reserved.
|
||
|
|
// See LICENSE for licensing terms.
|
||
|
|
|
||
|
|
package grpc_middleware
|
||
|
|
|
||
|
|
import (
|
||
|
|
"golang.org/x/net/context"
|
||
|
|
"google.golang.org/grpc"
|
||
|
|
)
|
||
|
|
|
||
|
|
// WrappedServerStream is a thin wrapper around grpc.ServerStream that allows modifying context.
|
||
|
|
type WrappedServerStream struct {
|
||
|
|
grpc.ServerStream
|
||
|
|
// 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 {
|
||
|
|
return w.WrappedContext
|
||
|
|
}
|
||
|
|
|
||
|
|
// WrapServerStream returns a ServerStream that has the ability to overwrite context.
|
||
|
|
func WrapServerStream(stream grpc.ServerStream) *WrappedServerStream {
|
||
|
|
if existing, ok := stream.(*WrappedServerStream); ok {
|
||
|
|
return existing
|
||
|
|
}
|
||
|
|
return &WrappedServerStream{ServerStream: stream, WrappedContext: stream.Context()}
|
||
|
|
}
|