Source file src/github.com/grpc-ecosystem/go-grpc-middleware/wrappers.go

Documentation: github.com/grpc-ecosystem/go-grpc-middleware

     1  // Copyright 2016 Michal Witkowski. All Rights Reserved.
     2  // See LICENSE for licensing terms.
     3  
     4  package grpc_middleware
     5  
     6  import (
     7  	"context"
     8  
     9  	"google.golang.org/grpc"
    10  )
    11  
    12  // WrappedServerStream is a thin wrapper around grpc.ServerStream that allows modifying context.
    13  type 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()
    20  func (w *WrappedServerStream) Context() context.Context {
    21  	return w.WrappedContext
    22  }
    23  
    24  // WrapServerStream returns a ServerStream that has the ability to overwrite context.
    25  func 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  }
    31  

View as plain text