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