blob: 597b862445fec504535e1f8fde052eb3379757bd [file] [log] [blame]
khenaidoo59ce9dd2019-11-11 13:05:32 -05001// Copyright 2016 Michal Witkowski. All Rights Reserved.
2// See LICENSE for licensing terms.
3
4package grpc_middleware
5
6import (
7 "golang.org/x/net/context"
8 "google.golang.org/grpc"
9)
10
11// WrappedServerStream is a thin wrapper around grpc.ServerStream that allows modifying context.
12type WrappedServerStream struct {
13 grpc.ServerStream
14 // WrappedContext is the wrapper's own Context. You can assign it.
15 WrappedContext context.Context
16}
17
18// Context returns the wrapper's WrappedContext, overwriting the nested grpc.ServerStream.Context()
19func (w *WrappedServerStream) Context() context.Context {
20 return w.WrappedContext
21}
22
23// WrapServerStream returns a ServerStream that has the ability to overwrite context.
24func WrapServerStream(stream grpc.ServerStream) *WrappedServerStream {
25 if existing, ok := stream.(*WrappedServerStream); ok {
26 return existing
27 }
28 return &WrappedServerStream{ServerStream: stream, WrappedContext: stream.Context()}
29}