How to safely add values to grpc ServerStream in interceptor

Question

I have a logging interceptor for my grpc server and want to add a value to the metadata (I want to track the request throughout its lifetime):

func (m *middleware) loggingInterceptor(srv interface{},
    ss grpc.ServerStream,
    info *grpc.StreamServerInfo,
    handler grpc.StreamHandler) 

    md, ok := metadata.FromIncomingContext(ss.Context())
    if !ok {
        return errors.New("could not get metadata from incoming stream context")
    }

    // add the transaction id to the metadata so that business logic can track it
    md.Append("trans-id", "some-transaction-id")

    // call the handler func
    return handler(srv, ss)
}

but the docs for FromIncomingContext state that:

// FromIncomingContext returns the incoming metadata in ctx if it exists.  The
// returned MD should not be modified. Writing to it may cause races.
// Modification should be made to copies of the returned MD.

Ok, so I look at the copy function and copy the metadata:

mdCopy := md.Copy()
mdCopy.Append("trans-id", "some-transaction-id")

and think "how do I attach this metadata back to the ServerStream context?", and I check if there's some ss.SetContext(newCtx), but I don't see anything of the sort. Am I thinking about this from the wrong perspective, or am I missing something else?

You could find the sample of `wrappedStream` from here https://stackoverflow.com/a/74148611/3011380 — zangw, Nov 11 '22 at 11:32

score 6 · Answer 1 · answered Apr 08 '20 at 18:28

You would need to use NewIncomingContext to create a copy of the current context in the stream.

Then you would have to create a wrappedStream type which overrides Context method in ServerStream to return the modified context. You would need to pass this wrappedStream to the handler that you received in your interceptor.

You can see an example of this here (it overrides other methods here, but the idea is the same): https://github.com/grpc/grpc-go/blob/master/examples/features/interceptor/server/main.go#L106-L124

Hope this helps.

score 1 · Answer 2 · answered Mar 19 '22 at 22:50

Easwar is right.

You can either create your own ServerStream implementation, override Context() method with your own context or there's struct inside grpc package which is WrappedServerStream (github.com/grpc-ecosystem/go-grpc-middleware) which you can pass context and original server stream object and use it inside handler.

Example:

    // This methods get the current context, and creates a new one
    newContext, err := interceptor.authorize(ss.Context(), info.FullMethod)
    if err != nil {
        log.Printf("authorization failed: %v", err)
        return err
    }

    
    err = handler(srv, &grpc_middleware.WrappedServerStream{
        ServerStream:   ss,
        WrappedContext: newContext,
    })

How to safely add values to grpc ServerStream in interceptor

2 Answers2