在 Go GRPC 服务器流式拦截器上修改元数据

Update个个

更简单解决方案是只覆盖ServerStream的Context()方法

type serverStream struct {
    grpc.ServerStream
    ctx context.Context
}

func (s *serverStream) Context() context.Context {
    return s.ctx
}

func UserIDInterceptor(srv interface{}, ss grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
    md, ok := metadata.FromIncomingContext(ss.Context())
    if ok {
        md.Append("X-User-Id", "real_user_id")
    }
    newCtx := metadata.NewIncomingContext(ss.Context(), md)

    return handler(srv, &serverStream{ss, newCtx})
}

Update个个

另一个简单的解决方案是将一个包装器定义为grpc.ServerStream，如下所示

type serverStreamWrapper struct {
    ss  grpc.ServerStream
    ctx context.Context
}

func (w serverStreamWrapper) Context() context.Context        { return w.ctx }
func (w serverStreamWrapper) RecvMsg(msg interface{}) error   { return w.ss.RecvMsg(msg) }
func (w serverStreamWrapper) SendMsg(msg interface{}) error   { return w.ss.SendMsg(msg) }
func (w serverStreamWrapper) SendHeader(md metadata.MD) error { return w.ss.SendHeader(md) }
func (w serverStreamWrapper) SetHeader(md metadata.MD) error  { return w.ss.SetHeader(md) }
func (w serverStreamWrapper) SetTrailer(md metadata.MD)       { w.ss.SetTrailer(md) }

func UserIDInterceptor(srv interface{}, ss grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
    md, ok := metadata.FromIncomingContext(ss.Context())
    if ok {
        md.Append("X-User-Id", "real_user_id")
    }
    newCtx := metadata.NewIncomingContext(ss.Context(), md)

    return handler(srv, serverStreamWrapper{ss, newCtx})
}

您可以使用NewIncomingContext来创建流中当前上下文的副本.

由于没有方法来设置grpc.ServerStream的context，为了将上下文设置回ServerStream,wrappedStream被定义为context.Context，并且SetContext方法被设置为context.Context

type wrappedStream struct {
    grpc.ServerStream
    ctx context.Context
}

func (w *wrappedStream) SetContext(ctx context.Context) {
    w.ctx = ctx
}

完整的示例代码

type wrappedStream struct {
    grpc.ServerStream
    ctx context.Context
}

func (w *wrappedStream) Context() context.Context {
    return w.ctx
}

func (w *wrappedStream) SetContext(ctx context.Context) {
    w.ctx = ctx
}

func (w *wrappedStream) RecvMsg(m interface{}) error {
    return w.ServerStream.RecvMsg(m)
}

func (w *wrappedStream) SendMsg(m interface{}) error {
    return w.ServerStream.SendMsg(m)
}

type StreamContextWrapper interface {
    grpc.ServerStream
    SetContext(context.Context)
}

func newStreamContextWrapper(ss grpc.ServerStream) StreamContextWrapper {
    ctx := ss.Context()
    return &wrappedStream{
        ss,
        ctx,
    }
}

func UserIDInterceptor(srv interface{}, ss grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
    md, ok := metadata.FromIncomingContext(ss.Context())
    if ok {
        md.Append("X-User-Id", "real_user_id")
    }
    newCtx := metadata.NewIncomingContext(ss.Context(), md)

    sw := newStreamContextWrapper(ss)
    sw.SetContext(newCtx)

    return handler(srv, sw)
}