Search code examples
goprotocol-buffersgrpcconnect-go

Modifying responsebody in connect-go interceptor

I am using Buf's connect-go library to implement a gRPC server.

Many of the gRPC calls are time-sensitive, so they include a field that the client uses to send its current timestamp. The server compares the client timestamp with the local timestamp and returns the difference between them. Here is the an example from the .proto definitions:

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

Because I have this implemented for several RPC methods, I wanted to see if I could use an interceptor to handle it so I don't need to make sure it is being handled in all of the individual RPC method implementations.

Because of how the protoc-gen-go compiler defines getters for the fields, checking if the request message contains the now_on_device field is easily done by defining an interface and using type assertion:

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

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

This makes most of the interceptor very easy to write:

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

The problem I have (as noted in the comment above) is how to modify the response.

I can't define an interface for the response the same way I did for the request because protoc-gen-go does not define setters. Then I thought I could just use a type switch, like this (where the TODO comment is above):

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
}

There are three problems with this approach:

  1. 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.)
  2. Using type assertion requires me to repeat almost the same code block over and over for each type.
  3. This is no longer simpler than implementing this in each RPC method individually.

Is there an easier way to use an interceptor to modify a field in the response? Or is there some other way I should be doing this?

Solution

  • Deepankar outlines one solution, though I do see the appeal of keeping all the response data in the schema-defined response structure. This would certainly be simpler if protoc-gen-go generated setters to go along with the getters!

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

    You don't need to do this. In your example, res.Any() returns a pointer to the protobuf message - you can modify it in place. Your type switch can look like this:

    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.

    Unfortunately, your best bet here is likely reflection. You can choose between standard Go reflection or protobuf reflection - either should work. With protobuf reflection, something like this should do the trick:

    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

    It's up to you whether you think this is better or worse than the repetitive type switch - it's quite a bit more complex, but it does keep things a bit more DRY.