go - 在 connect-go 拦截器中修改 responsebody

Question

我正在使用 Buf 的 connect-go用于实现 gRPC 服务器的库。

许多 gRPC 调用都是时间敏感的，因此它们包含一个字段，客户端使用该字段发送其当前时间戳。服务器将客户端时间戳与本地时间戳进行比较，并返回它们之间的差异。以下是 .proto 定义的示例:

service EventService {
    // Start performing a task
    rpc Start (StartRequest) returns (StartResponse);
}

message StartRequest {
    int64 location_id = 1;
    int64 task_id = 2;
    Location user_latlng = 3;
    google.protobuf.Timestamp now_on_device = 4;
}

message StartResponse {
    TaskPerformanceInfo info = 1;
    google.protobuf.Duration device_offset = 2;
}

因为我已经为几个 RPC 方法实现了这个，所以我想看看我是否可以使用 interceptor处理它，所以我不需要确保它在所有单独的 RPC 方法实现中都被处理。

由于 protoc-gen-go 编译器如何为字段定义 getter，检查请求消息是否包含 now_on_device 字段很容易通过定义接口(interface)和使用类型断言:

type hasNowOnDevice interface {
    GetNowOnDevice() *timestamppb.Timestamp
}

if reqWithNow, ok := req.Any().(hasNowOnDevice); ok {
   // ...
}

这使得大多数拦截器非常容易编写:

func MakeDeviceTimeInterceptor() func(connect.UnaryFunc) connect.UnaryFunc {
    return connect.UnaryInterceptorFunc(
        func(next connect.UnaryFunc) connect.UnaryFunc {
            return connect.UnaryFunc(func(ctx context.Context, req connect.AnyRequest) (connect.AnyResponse, error) {
                now := time.Now().UTC()
                ctxa := context.WithValue(ctx, CurrentTimestampKey{}, now)

                var deviceTimeOffset time.Duration
                // If the protobuf message has a `NowOnDevice` field, use it
                // to get the difference betweent the device time and server time.
                if reqWithNow, ok := req.Any().(hasNowOnDevice); ok {
                    deviceTime := reqWithNow.GetNowOnDevice().AsTime()
                    deviceTimeOffset = now.Sub(deviceTime)
                    ctxa = context.WithValue(ctxa, DeviceTimeDiffKey{}, deviceTimeOffset)
                }

                res, err := next(ctxa, req)

                // TODO: How do I modify the response here?

                return res, err
            })
        },
    )
}

我遇到的问题(如上面评论中所述)是如何修改响应。

我无法像为请求那样为响应定义接口(interface)，因为 protoc-gen-go 没有定义 setter。然后我想我可以只使用一个类型开关，像这样(TODO 注释在上面):

switch resMsg := res.Any().(type) {
case *livev1.StartResponse:
    resMsg.DeviceOffset = durationpb.New(deviceTimeOffset)
    return &connect.Response[livev1.StartResponse]{
        Msg: resMsg,
    }, err
case *livev1.StatusResponse:
    resMsg.DeviceOffset = durationpb.New(deviceTimeOffset)
    return &connect.Response[livev1.StatusResponse]{
        Msg: resMsg,
    }, err
}

这种方法存在三个问题:

1. 我找不到将旧响应中的 header /尾部复制到新响应中的方法。 (我认为此时它们还没有真正确定，但我不确定。)
2. 使用类型断言需要我为每种类型一遍又一遍地重复几乎相同的代码块。
3. 这不再比在每个 RPC 方法中单独实现它更简单。

有没有更简单的方法来使用拦截器修改响应中的字段？还是有其他方法可以做到这一点？

Answer 1

Deepankar 概述了一种解决方案，但我确实看到了将所有响应数据保存在模式定义的响应结构中的吸引力。如果 protoc-gen-go 生成 setter 与 getter 一起使用，这肯定会更简单!

I can't find a way to copy the headers/trailers from the old response into this new response. (I don't think they are actually set yet at this point, but I don't know that for sure.)

您不需要这样做。在您的示例中，res.Any() 返回指向 protobuf 消息的指针 - 您可以就地修改它。您的类型开关可能如下所示:

switch resMsg := res.Any().(type) {
case *livev1.StartResponse:
    resMsg.DeviceOffset = durationpb.New(deviceTimeOffset)
case *livev1.StatusResponse:
    resMsg.DeviceOffset = durationpb.New(deviceTimeOffset)
}
return res, err

Using type assertion requires me to repeat almost the same code block over and over for each type.

不幸的是，这里最好的选择可能是反射。您可以在标准 Go 反射或 protobuf 反射之间进行选择——两者都可以。使用 protobuf 反射，像这样的事情应该可以解决问题:

res, err := next(ctx, req)
if err != nil {
    return nil, err
}
msg, ok := res.Any().(proto.Message)
if !ok {
    return res, nil
}

// Keep your logic to calculate offset!
var deviceTimeOffset time.Duration

// You could make this a global.
durationName := (*durationpb.Duration)(nil).ProtoReflect().Descriptor().FullName()

refMsg := msg.ProtoReflect()
offsetFD := refMsg.Descriptor().Fields().ByName("DeviceOffset")
if offsetFD != nil &&
    offsetFD.Message() != nil &&
    offsetFD.Message().FullName() == durationName {
    refOffset := durationpb.New(deviceTimeOffset).ProtoReflect()
    refMsg.Set(
        offsetFD, 
        protoreflect.ValueOf(refOffset),
    )
}
return res, nil

这取决于您是否认为这比重复的类型切换更好或更坏——它要复杂得多，但它确实让事情变得更干。