feat: Add WithForwardResponseRewriter to allow easier/more useful response control

[nkcmr]

Context/Background

I am working with grpc-gateway to mimick an older REST API while
implementing the core as a gRPC server. This REST API, by convention,
emits a response envelope.

After extensively researching grpc-gateway's code and options, I think
that there is not a good enough way to address my use-case; for a few
more nuanced reasons.

The sum of my requirements are as follows:

• 1. Allow a particular field of a response message to be used as the main
     response. ✅ This is handled with response_body annotation.
• 2. Be able to run interface checks on the response to extract useful
     information, like next_page_token 0 and surface it in the final
     response envelope. (npt, ok := resp.(interface { GetNextPageToken() string })).
• 3. Take the true result and place it into a response envelope along with
     other parts of the response by convention and let that be encoded and
     sent as the response instead.

Implementing a response envelope with Marshaler

My first attempt at getting my gRPC server's responses in an envelope
led me to implement my own Marshaler, I have seen this approach
discussed in #4483.

This does satisfy requirements 1 and 3 just fine, since the HTTP
annotations helpfully allow the code to only receive the true result,
and the Marshal interface has enough capabilities to take that and wrap
it in a response envelope.

However, requirements 1 and 2 are not both satisfiable with the
current grpc-gateway code because of how the XXX_ResponseBody() is
called before passing to the Marshal(v) function. This strips out
the other fields that I would normally be able to detect and place in
the response envelope.

I even tried creating my own protobuf message extension that would let
me define another way of defining the "true result" field. But the
options for implementing that are either a ton of protoreflect at
runtime to detect and extract that, or I am writing another protobuf
generator plugin (which I have done before 1), but both of those
options seem quite complex.

Other non-starter options

Just to get ahead of the discussion, WithForwardResponseOption clearly
was not meant for this use-case. At best, it seems to only be a way to
take information that might be in the response and add it as a header.

In practice

This change fulfills my requirements by allowing logic to be inserted
right before the Marshal is called:

gatewayMux := runtime.NewServeMux(
  runtime.WithForwardResponseRewriter(func(ctx context.Context, response proto.Message) (interface{}, error) {
    if s, ok := response.(*statuspb.Status); ok {
      return rewriteStatusToErrorEnvelope(ctx, s)
    }
    return rewriteResultToEnvelope(ctx, response)
  }),
)

In this PR

This PR introduces a new ServeMuxOption called
WithForwardResponseRewriter that allows for a user-provided function
to be supplied that can take a response proto.Message and return any
during unary response forwarding, stream response forwarding, and error
response forwarding.

The implementation ensures that checks for XXX_ResponseBody() are
always done, even on the results returned from user-provided forward
response rewriters.

Also added GetResponseBody() to allow any users that use the
WithForwardResponseRewriter option to still rely on the runtime to
know how to extract the configured response_body.

The code generation was also updated to make the XXX_ResponseBody()
response wrappers embed the concrete type instead of just
proto.Message. This allows any code in response rewriter functions to
be able to have access to the original type, so that interface checks
against it should pass as if it was the original message.

Testing

Added some basic unit tests to ensure Unary/Stream and error handlers
invoke ForwardResponseRewriter correctly.

Looking for feedback on

• The function signature of ForwardResponseRewriter is very light at
  the moment. Input about what else might be needed here to serve
  potential future uses of this are appreciated.

[johanbrandhorst]

Just having a brief look at this before diving deeper, but you'll need to regenerate the files in the repo using the new template.go logic. See CONTRIBUTING.md in the root of the repo for the steps.

[johanbrandhorst]

Thanks for your thoughtful and detailed description. I'm hoping to exhaust all options before going for new functionality, as I'm sure you understand. I had some crazy ideas to throw out there:

1. With your custom marshaler, can't you just not use the response_body option? You could make the marshaler switch on the type and pick the field you want to be in the envelope. Even the custom extension doesn't sound all that bad (protoreflect is good, actually).
2. It seems like you want to add a static type wrapper around all messages? Have you considered using a gRPC client-side interceptor on the gateway side? I don't even know if it's possible to change the type returned from an interceptor, but thought it'd be worth suggesting.

[nkcmr]

Hi @johanbrandhorst!

Thanks for the feedback. I have run the code-generators and pushed those changes up (sorry for missing that step and just breezing by the contributing docs).

---

Let me address your other feedback:

Have you considered using a gRPC client-side interceptor on the gateway side?

So, I might be misunderstanding what exactly you're recommending here, but I think you're recommending grpc.UnaryClientInterceptor be configred on the Gateway → gRPC-Server end. However, I am confused by how exactly I would be able to go from Response A (the original response of the method) to Response B (My response envelope) while using this:

grpc.NewClient(
	"localhost:8334",
	grpc.WithUnaryInterceptor(func(ctx context.Context, method string, req, reply any, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
		err := invoker(ctx, method, req, reply, cc, opts...)
		if err != nil {
			// uh..., how do I return my error response envelope?
		}
		// uh..., how do return my normal response envelope?
	}),
)

The interceptor does not allow for a different response to be returned. Again, I could be misunderstanding or missing something obvious, apologies if so.

Moreover though, my current approach to using grpc-gateway makes use of the registration functions that let you pass in the server implementation directly without using a gRPC connection; thus there is no way to add an interceptor there.

---

With your custom marshaler, can't you just not use the response_body option? You could make the marshaler switch on the type and pick the field you want to be in the envelope.

This would technically work, but I think it runs contrary to my goals with this change. I'll elaborate: I am actually trying to find a way to make this pattern repeatable across a number of teams inside my company. A lot of teams implement microservices that are upstream of a Load Balancer, and all of those services implement the response envelope convention.

My "vision" is to broaden adoption of gRPC but also remain accessible in the REST/JSON world. grpc-gateway is quite perfect to automate this backwards compatibility. The only missing piece is some company-internal library that can configure the *runtime.ServeMux just right to make everything else work "out-of-the-box." And thus I am here with this proposed change. It would fit perfectly into an off-the-shelf library that would allow teams to quickly/effortlessly adopt gRPC.

Exec Summary: That approach does work, but it does not "scale" well when the aim is to get large groups of separate people to adopt it; to have a manual tedium of needing to hand-implement this sort of neutralizes a lot of the benefits.

I fully understand that my "starry-eyed dreams" of some internal company experiment are not the concern of this project 😅. That being said, I think the proposal made here is a net-benefit for the usefulness/value of this project. Having skimmed the past issues, I do get a sense that it is a missing piece (just my perspective).

Lastly, protoreflect isn't bad, but the scale of some internal services would disallow so much unnecessary runtime CPU cycles to be spent doing that. I would more than likely end up opting for a new protoc plugin to stave off those concerns; but I still believe that would be much more complexity than desired.

[johanbrandhorst]

Thanks for elaborating. I expected that this would come down to wanting to make this scalable, which I agree that manual marshaler logic probably isn't.

I want to try to narrow down exactly what the proposal here is though, because it's also not a "give me a way to do anything with the response", which is essentially what the custom marshaler does. This seems to me to be almost an interceptor/middleware/decorator in the middle of the gateway logic, and I don't know how applicable it is to the general user. It seems like it solve this very particular problem of introducing a message envelope, dynamically (rather than statically in your protos), but I'm not yet convinced that it's generally useful. Do you think you could find issues in our backlog that could be resolve by this functionality?

I want this feature to be both generally useful and orthogonal to other functionality of the gateway. I'm worried that this option will be confusing when we have WithForwardResponseOption. How would you see them interacting? When would a user want to use both?

[nkcmr]

In my mental model of grpc-gateway's configuration surface, there are a few elements that allow control over what the response looks like:

• WithForwardResponseOption: This lets the user control headers and status if need be (maybe add X-Request-ID or switch the response code to be 201 Created)
• WithMarshalerOption: This lets the user control the byte-for-byte representation of the response; its encoding.
  • Note: Given what I was trying to do and what previous issue discussions were recommending, the requirement of configuring my custom marshaler with a MIME type was a source of confusion/friction. I was not in a mindset where I was needing to think about the actual encoding, just the shape/type.
• WithStreamErrorHandler / WithProtoErrorHandler: This lets the user control nearly every aspect of an HTTP response when an error occurs.
  • Note: This contrasts with the aforementioned things, in that it does give the user full control over the response. Other options are simply small escape-hatches tucked into immutable implementations.

My proposal of WithForwardResponseRewriter fills the need/gap in between WithForwardResponseOption and WithMarshalerOption.

ForwardResponseOption allows control over the headers and status, but not the response body. Marshaler is intended to let the gateway understand different Accept headers and adapt the response's encoding to meet that requirement. The missing piece is allowing for the response's shape/type to be configured at a server level.

So, with that, I do believe it is "orthogonal" to other functionality and fills a missing piece.

---

The other extreme alternative to all this (which I actually do not think is a good idea) is to allow full control over the response. A bad direction (IMO) would be to allow configuration over what gets assigned/used in the generated code here:

// myservice.pb.gw.go
var (
	// Let this be configured to be something else?
	forward_MyService_ListStuff_0 = runtime.ForwardResponseMessage
)

I don't think this is a good approach though. Letting grpc-gateway encapsulate most of the implementation details of handling trailing headers and picking the right marshaler is actually good. Then simply have escape-hatches to control high-level things like Supplemental Headers, Status Code, Response Shape/Type, and Encoding/Decoding.

---

Some past issues that I think this helps:
#1610: This person was trying to do pretty much the exact same thing as me. This proposal would have suited their ask perfectly.
#4483: I do think this person would have been helped by this as well, their use case did happen to fit into the Marshaler though, still awkwardly though IMO because of the difference between "encoding concerns" and "type/shape concerns."
#517: I think this person could use this proposal to take a special error *status.Status to make the response empty (switch it to *httpbody.HttpBody) and use ForwardResponseOption to set the right status code.

---

Sorry for the meandering response, but to summarize: There is (IMO) a missing piece to response control. I think this is the right direction to go in (Just one more escape-hatch, not necessarily allowing full response control), and this approach does not overlap with anything currently available. This proposal also I believe just lets you control the shape/type of the response, not "give me a way to do anything with the response." 🙂

Side-note: I see "change requested" still, did I address that? Also the CI / lint (pull_request) job is failing on things I did not seem to touch, is there a way to fix that?

[johanbrandhorst]

The lint failure is not your fault. I will take time to review this when I have a chance. Thanks for your reply.

[johanbrandhorst]

Alright, I buy it. BUT, I'm going to make you work for it. This is probably one of the biggest feature additions in a year or more. Could you please add a section on this new functionality to our docs page on customizing the gateway: https://github.com/grpc-ecosystem/grpc-gateway/blob/main/docs/docs/mapping/customizing_your_gateway.md? Thanks!

[johanbrandhorst]

This looks great, just a few minor changes

[nkcmr]

Thank you so much, @johanbrandhorst! Especially for putting up with my short novel comments. 😆

[johanbrandhorst]

Thank you for your contribution!

[johanbrandhorst]

I just cut v2.22.0 with this new feature. Thanks again :).