如何让Go语言gRPC服务器流拦截器验证失败时关闭流？

问题背景 #

我在Go中实现gRPC服务器流拦截器，需求是客户端未携带指定Header时，直接拒绝RPC并关闭流。参考grpc-go官方示例，只要在拦截器里返回错误就能实现这个逻辑。我写的拦截器和示例几乎一致，只是修改了Header检查逻辑——缺失指定Header时返回错误，但实际运行时，拦截器执行到return语句后，客户端和服务器都没记录错误状态，客户端调用RPC也收不到任何错误。之前我在服务器处理函数里直接做Header检查，缺失时客户端能正常收到流错误，现在想通过拦截器复现这个效果。

原拦截器代码 #

func ClientMetadataInterceptor(srv any, ss grpc.ServerStream, _ *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
    md, ok := metadata.FromIncomingContext(ss.Context())
    if !ok {
        return status.Errorf(codes.InvalidArgument, "no incoming metadata in rpc context")
    }

    var headerValue string
    if len(md["required-header"]) > 0 {
        headerValue = md["required-header"][0]
    } else {
        return status.Error(codes.InvalidArgument, "missing required header")
    }

    md.Append("new-header-key", headerValue)
    ctx := metadata.NewIncomingContext(ss.Context(), md)

    // Call the handler to complete the normal execution of the RPC
    err := handler(srv, &wrappedStream{ss, ctx})

    return err
}

问题分析与解决方案 #

核心问题 #

1. 自定义Stream包装类缺失正确实现：你用到了wrappedStream但未提供代码，这个类必须完整实现grpc.ServerStream接口，尤其是Context()方法要返回修改后的上下文，否则会导致上下文传递异常。
2. 拦截器注册可能有误：必须确保启动gRPC服务器时，通过grpc.StreamInterceptor()选项注册该拦截器，而非仅注册Unary拦截器。
3. Header检查逻辑可优化：使用md.Get()方法可以更简洁地获取Header值，避免手动判断切片长度。

修正后的完整代码 #

1. 补全wrappedStream的正确实现

import "context"
import "google.golang.org/grpc"

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

// 必须重写Context方法，返回自定义上下文
func (w *wrappedStream) Context() context.Context {
    return w.ctx
}

2. 优化后的拦截器代码

import (
    "google.golang.org/grpc/codes"
    "google.golang.org/grpc/metadata"
    "google.golang.org/grpc/status"
)

func ClientMetadataInterceptor(srv any, ss grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
    md, ok := metadata.FromIncomingContext(ss.Context())
    if !ok {
        return status.Errorf(codes.InvalidArgument, "rpc上下文无传入元数据")
    }

    // 使用md.Get简化Header检查
    headerValues := md.Get("required-header")
    if len(headerValues) == 0 {
        // 直接返回错误，gRPC框架会自动关闭流并通知客户端
        return status.Error(codes.InvalidArgument, "缺失必填Header: required-header")
    }

    headerValue := headerValues[0]
    md.Append("new-header-key", headerValue)
    ctx := metadata.NewIncomingContext(ss.Context(), md)

    // 调用后续handler执行RPC
    return handler(srv, &wrappedStream{ServerStream: ss, ctx: ctx})
}

3. 正确注册拦截器

import "google.golang.org/grpc"

func main() {
    // 注册流拦截器
    server := grpc.NewServer(
        grpc.StreamInterceptor(ClientMetadataInterceptor),
        // 若需支持Unary接口，可添加Unary拦截器：grpc.UnaryInterceptor(...)
    )

    // 注册服务并启动...
}

验证要点 #

• 客户端调用时确保未携带required-header（gRPC会自动将客户端发送的Header键转为小写，所以客户端用Required-Header也会被识别为required-header）。
• 客户端调用流RPC时，在调用Recv()或CloseSend()方法时会收到对应的错误。
• 可开启gRPC调试日志查看服务器端错误输出：

  export GRPC_GO_LOG_VERBOSITY_LEVEL=99 GRPC_GO_LOG_SEVERITY_LEVEL=info
  

内容的提问来源于stack exchange，提问作者Chris Jansson