go - 如何在流拦截器中访问请求？

Question

我有一个包含以下代码的一元拦截器:

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)
}

如果我知道流肯定只会从服务器流向客户端，我如何才能将其转换为流拦截器？另外，有什么办法让它只在流开始的那一刻拦截？

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

    return handler(srv, wrappedStream)
}

Answer 1

对于“服务器流式 RPC”，客户端发送一条消息(请求)，服务器响应多条消息。在流拦截器中，需要调用一次ServerStream.RecvMsg(...) 来获取客户端的请求。然后，您需要将“包装流”传递给 handler，后者稍后会在第一次调用 RecvMsg 时返回此消息。对于使用 protobuf 消息的 gRPC，您可以执行如下操作。此 serverStreamWrapper 将允许您在拦截器中调用 peekRequest(...) 以获取请求，然后将流传递给处理程序。

对于这个具体的例子，您也可以在包装的 RecvMsg(...) 函数本身中实现逻辑。 gRPC 调用此函数以读取从客户端到服务器的请求。这可能比尝试存储请求更简单。

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)
}