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?

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.

Do you have the possibility use Headers instead of body. If Clients can send NowOnDevice through request headers, then you can send back the response in response headers instead. Unix Timestamp might be the best approach.

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

			var deviceTimeOffset time.Duration
			// Check the header message `now-on-device` field, instead of body
			reqWithNow := req.Header().Get("now-on-device")

			if reqWithNow != "" {
				val, err := strconv.Atoi(reqWithNow)
				if err != nil {
					return nil, connect.NewError(connect.CodeInvalidArgument, errors.New("invalid timestamp"))
				}

				deviceTime := time.Unix(int64(val), 0)
				deviceTimeOffset = now.Sub(deviceTime)
				ctxa = context.WithValue(ctxa, DeviceTimeDiffKey{}, deviceTimeOffset)
			}

			res, err := next(ctxa, req)

			// Set to response header if value is set
			if deviceTimeOffset != 0 {
				res.Header().Set("device-time-offset", fmt.Sprintf("%d", deviceTimeOffset))
			}

			return res, err
		}
	}
}

Then you have the response:

curl -v \
    --header "Content-Type: application/json" --header "now-on-device: 1656442814" \
    --data '{"name": "Jane"}' \
    http://localhost:8080/greet.v1.GreetService/Greet
*   Trying 127.0.0.1:8080...
* Connected to localhost (127.0.0.1) port 8080 (#0)
> POST /greet.v1.GreetService/Greet HTTP/1.1
> Host: localhost:8080
> User-Agent: curl/7.79.1
> Accept: */*
> Content-Type: application/json
> now-on-device: 1656442814
> Content-Length: 16
>
* Mark bundle as not supporting multiuse
< HTTP/1.1 200 OK
< Accept-Encoding: gzip
< Content-Type: application/json
< Device-Time-Offset: 7259524766000
< Greet-Version: v1
< Date: Tue, 28 Jun 2022 21:01:13 GMT
< Content-Length: 27
<
* Connection #0 to host localhost left intact
{"greeting":"Hello, Jane!"}