How to access request in stream interceptor?

user_13178827 0
Viewed 765

I have a unary interceptor that contains the following code:

func (m Middlewares) LocationInterceptor(c context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
    ctx := c.(HarmonyContext)
    location, ok := req.(interface{ GetLocation() *corev1.Location })
    if !ok {
        panic("location middleware used on message without a location")
    }
    
    ctx.Location := location.GetLocation()

    return handler(c, req)
}

How would I be able to convert this to a stream interceptor, if I know that the stream will definitely only stream from server to client? In addition, is there any way to make it only intercept when the moment the stream begins?

func (m Middlewares) LocationInterceptorStream(srv interface{}, ss grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
    wrappedStream := ss.(HarmonyWrappedServerStream)

    return handler(srv, wrappedStream)
}
grpc go
1 Answers

For a "server streaming RPC", the client sends one message (the request), and the server responds with multiple messages. In the stream interceptor, you need to call ServerStream.RecvMsg(...) once to get the request from the client. You will then need to pass a "wrapped stream" into handler which will later return this message on the first call to RecvMsg. For gRPC using protobuf messages, you can do something like the following. This serverStreamWrapper will allow you to call peekRequest(...) in the interceptor to get the request, then pass the stream on to the handler.

For this specific example, you may also be able to implement the logic in the wrapped RecvMsg(...) function itself. This function gets called by gRPC to read the request from the client to the server. This may be simpler than trying to store the request.

type serverStreamWrapper struct {
    peekedRequest proto.Message
    wrappedStream grpc.ServerStream
}

func newServerStreamWrapper(stream grpc.ServerStream) *serverStreamWrapper {
    return &serverStreamWrapper{nil, stream}
}

func (s *serverStreamWrapper) peekRequest(msg interface{}) error {
    protoMsg := msg.(proto.Message)
    if protoMsg == nil {
        panic("BUG: msg must not be nil")
    }
    if s.peekedRequest != nil {
        panic("BUG: Must only called peekRequest once")
    }

    err := s.wrappedStream.RecvMsg(protoMsg)
    if err == nil {
        s.peekedRequest = protoMsg
    }
    return err
}

func (s *serverStreamWrapper) RecvMsg(msg interface{}) error {
    if s.peekedRequest != nil {
        protoMsg := msg.(proto.Message)
        proto.Reset(protoMsg)
        proto.Merge(protoMsg, s.peekedRequest)
        s.peekedRequest = nil
        return nil
    }
    return s.wrappedStream.RecvMsg(msg)
}
user_160997320
Related