Connect-Go 项目中流式 RPC 的上下文传递问题解析

在 Connect-Go 项目中，开发者经常会遇到在流式 RPC 调用中无法正确传递上下文值的问题。本文将通过一个典型场景，深入分析问题原因并提供解决方案。

问题现象

开发者在使用 Connect-Go 实现认证中间件时，发现通过上下文传递的用户信息在流式 RPC 处理函数中无法获取。具体表现为：

1. 中间件通过 HTTP 头部获取 token 并验证
2. 验证成功后，将用户 ID 存入请求上下文
3. 在普通的 Unary RPC 中能够正常获取上下文中的用户信息
4. 但在流式 RPC（如 rpc AskGPT(AskGPTRequest) returns(stream AskGPTResponse)）中却无法获取

原因分析

问题的根本原因在于 Connect-Go 中间件的实现机制。Connect-Go 对 Unary 和 Streaming 两种 RPC 模式采用了不同的处理路径：

1. Unary RPC：使用 UnaryInterceptorFunc 处理，开发者只需实现简单的函数式拦截器
2. Streaming RPC：需要实现完整的 Interceptor 接口，特别是 WrapStreamingHandler 方法

当开发者只实现了 UnaryInterceptorFunc 而没有处理 Streaming 情况时，流式调用的请求会绕过中间件的上下文处理逻辑，导致上下文值丢失。

解决方案

要解决这个问题，需要实现完整的拦截器接口：

type authInterceptor struct{}

func (i *authInterceptor) WrapUnary(next connect.UnaryFunc) connect.UnaryFunc {
    return func(ctx context.Context, req connect.AnyRequest) (connect.AnyResponse, error) {
        // 原有的 Unary 处理逻辑
        tokenString := req.Header().Get("token")
        // ...验证逻辑
        ctx = ctxuser.SetInContext(ctx, ctxuser.RequestInfo{
            UserID: userID,
        })
        return next(ctx, req)
    }
}

func (i *authInterceptor) WrapStreamingHandler(next connect.StreamingHandlerFunc) connect.StreamingHandlerFunc {
    return func(ctx context.Context, conn connect.StreamingHandlerConn) error {
        // 流式调用的处理逻辑
        tokenString := conn.RequestHeader().Get("token")
        // ...同样的验证逻辑
        ctx = ctxuser.SetInContext(ctx, ctxuser.RequestInfo{
            UserID: userID,
        })
        return next(ctx, conn)
    }
}

func (i *authInterceptor) WrapStreamingClient(next connect.StreamingClientFunc) connect.StreamingClientFunc {
    // 客户端流式调用的处理（如不需要可留空）
    return next
}

最佳实践建议

1. 统一处理逻辑：将认证逻辑提取为独立函数，供 Unary 和 Streaming 共用
2. 上下文工具类优化：可以增强上下文工具类，添加日志和错误处理
3. 全面测试：同时测试 Unary 和 Streaming 两种调用方式
4. 考虑使用专业库：对于复杂的认证场景，可以考虑使用专门的认证库

总结

在 Connect-Go 项目中处理流式 RPC 时，开发者必须注意实现完整的拦截器接口，而不仅仅是 Unary 拦截器。这种设计虽然增加了些许复杂性，但提供了更大的灵活性和一致性。理解 Connect-Go 的这种设计模式，有助于开发者构建更健壮的 gRPC 服务。