fix: Allow replacing request and response in client interceptors #158
Conversation
Summary of ChangesHello @sangmeshcp, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request significantly enhances the flexibility of client-side interceptors by allowing them to not only observe but also actively modify or replace the data flowing through the client. This change empowers developers to implement more advanced use cases such as data transformation, encryption, or dynamic mocking directly within the interceptor chain, without altering the core client logic. Highlights
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
sangmeshcp
changed the title
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a powerful feature allowing client interceptors to replace request and response payloads. The implementation, however, has a critical flaw across all modified methods in a2aclient/client.go: it uses direct, unsafe type assertions on the payloads returned from interceptors. This can lead to panics if an interceptor returns a payload of an incorrect type or a nil payload where one is expected.
I've added two detailed comments with suggestions for GetTask (as an example of a simple RPC) and SendStreamingMessage (as an example of a streaming RPC) to demonstrate how to fix this using safe type assertions. This pattern of checking types should be applied consistently to all other methods modified in this PR to ensure robustness.
| query = iReq.(*a2a.TaskQueryParams) | ||
|
|
||
| resp, err := c.transport.GetTask(ctx, query) | ||
| if errOverride := c.interceptAfter(ctx, method, resp, err); errOverride != nil { | ||
| return nil, errOverride | ||
| iResp, finalErr := c.interceptAfter(ctx, method, resp, err) | ||
| if finalErr != nil { | ||
| return nil, finalErr | ||
| } | ||
|
|
||
| return resp, err | ||
| return iResp.(*a2a.Task), nil |
There was a problem hiding this comment.
The direct type assertions for both the request and response payloads in this function are unsafe and can cause a panic. An interceptor could return a payload of an incorrect type, or a nil response, which would crash the application.
You should use the value, ok form of type assertion to validate the types and handle nil values gracefully before using them. This same issue exists in CancelTask, SendMessage, and the various ...TaskPushConfig methods.
Here is a suggested implementation for GetTask that addresses this. Note that you will need to import the fmt package. A similar approach should be taken for the other non-streaming methods.
newQuery, ok := iReq.(*a2a.TaskQueryParams)
if !ok {
return nil, fmt.Errorf("interceptor returned unexpected request type: got %T, want *a2a.TaskQueryParams", iReq)
}
query = newQuery
resp, err := c.transport.GetTask(ctx, query)
iResp, finalErr := c.interceptAfter(ctx, method, resp, err)
if finalErr != nil {
return nil, finalErr
}
task, ok := iResp.(*a2a.Task)
if !ok {
return nil, fmt.Errorf("interceptor returned unexpected response type: got %T, want *a2a.Task", iResp)
}
return task, nil| message = iReq.(*a2a.MessageSendParams) | ||
|
|
||
| if card := c.card.Load(); card != nil && !card.Capabilities.Streaming { | ||
| resp, err := c.transport.SendMessage(ctx, message) | ||
| if errOverride := c.interceptAfter(ctx, method, resp, err); errOverride != nil { | ||
| yield(nil, errOverride) | ||
| iResp, finalErr := c.interceptAfter(ctx, method, resp, err) | ||
| if finalErr != nil { | ||
| yield(nil, finalErr) | ||
| return | ||
| } | ||
| yield(resp, err) | ||
| yield(iResp.(a2a.Event), nil) | ||
| return | ||
| } | ||
|
|
||
| for resp, err := range c.transport.SendStreamingMessage(ctx, message) { | ||
| if errOverride := c.interceptAfter(ctx, method, resp, err); errOverride != nil { | ||
| yield(nil, errOverride) | ||
| iResp, finalErr := c.interceptAfter(ctx, method, resp, err) | ||
| if finalErr != nil { | ||
| yield(nil, finalErr) | ||
| return | ||
| } | ||
| // When err is nil, finalErr is nil (unless interceptor added error) | ||
| // But if err was NOT nil, transport loop yields error. | ||
| // interceptAfter will see err, and return finalErr = err (or modified). | ||
| // So if finalErr != nil, we yield it and return. | ||
|
|
||
| if err != nil { | ||
| yield(nil, err) | ||
| return | ||
| } | ||
| // If transport sends (nil, err), interceptAfter sees (nil, err). | ||
| // If it returns (nil, finalErr), we yield (nil, finalErr). | ||
|
|
||
| if !yield(resp, nil) { | ||
| if !yield(iResp.(a2a.Event), nil) { | ||
| return | ||
| } | ||
| } |
There was a problem hiding this comment.
As with the non-streaming methods, the type assertions in this streaming function are unsafe. This applies to the request payload and the response payloads in both the non-streaming fallback and the streaming loop. An incorrect type from an interceptor will cause a panic.
Please add safe type checks for all assertions. Here is a suggested implementation for this function. A similar approach should be used for ResubscribeToTask. Note that you will need to import the fmt package.
newMessage, ok := iReq.(*a2a.MessageSendParams)
if !ok {
yield(nil, fmt.Errorf("interceptor returned unexpected request type: got %T, want *a2a.MessageSendParams", iReq))
return
}
message = newMessage
if card := c.card.Load(); card != nil && !card.Capabilities.Streaming {
resp, err := c.transport.SendMessage(ctx, message)
iResp, finalErr := c.interceptAfter(ctx, method, resp, err)
if finalErr != nil {
yield(nil, finalErr)
return
}
event, ok := iResp.(a2a.Event)
if !ok {
yield(nil, fmt.Errorf("interceptor returned unexpected response type: got %T, want a2a.Event", iResp))
return
}
yield(event, nil)
return
}
for resp, err := range c.transport.SendStreamingMessage(ctx, message) {
iResp, finalErr := c.interceptAfter(ctx, method, resp, err)
if finalErr != nil {
yield(nil, finalErr)
return
}
// When err is nil, finalErr is nil (unless interceptor added error)
// But if err was NOT nil, transport loop yields error.
// interceptAfter will see err, and return finalErr = err (or modified).
// So if finalErr != nil, we yield it and return.
// If transport sends (nil, err), interceptAfter sees (nil, err).
// If it returns (nil, finalErr), we yield (nil, finalErr).
event, ok := iResp.(a2a.Event)
if !ok {
yield(nil, fmt.Errorf("interceptor returned unexpected response type: got %T, want a2a.Event", iResp))
return
}
if !yield(event, nil) {
return
}
}
| if err != nil { | ||
| return nil, err | ||
| } | ||
| query = iReq.(*a2a.TaskQueryParams) |
There was a problem hiding this comment.
I'd rather make interceptBefore and interceptAfter into free generic functions and do type assertions there, for example:
Details
diff --git a/a2aclient/client.go b/a2aclient/client.go
index a796319..c59986b 100644
--- a/a2aclient/client.go
+++ b/a2aclient/client.go
@@ -16,6 +16,7 @@ package a2aclient
import (
"context"
+ "fmt"
"iter"
"sync/atomic"
@@ -67,11 +68,10 @@ func (c *Client) AddCallInterceptor(ci CallInterceptor) {
func (c *Client) GetTask(ctx context.Context, query *a2a.TaskQueryParams) (*a2a.Task, error) {
method := "GetTask"
- ctx, iReq, err := c.interceptBefore(ctx, method, query)
+ ctx, query, err := interceptBefore(ctx, c, method, query)
if err != nil {
return nil, err
}
- query = iReq.(*a2a.TaskQueryParams)
resp, err := c.transport.GetTask(ctx, query)
iResp, finalErr := c.interceptAfter(ctx, method, resp, err)
@@ -316,7 +316,7 @@ func (c *Client) withDefaultSendConfig(message *a2a.MessageSendParams, blocking
return &result
}
-func (c *Client) interceptBefore(ctx context.Context, method string, payload any) (context.Context, any, error) {
+func interceptBefore[T any](ctx context.Context, c *Client, method string, payload *T) (context.Context, *T, error) {
req := Request{
Method: method,
BaseURL: c.baseURL,
@@ -337,7 +337,12 @@ func (c *Client) interceptBefore(ctx context.Context, method string, payload any
ctx = localCtx
}
- return withCallMeta(ctx, req.Meta), req.Payload, nil
+ newPayload, ok := req.Payload.(*T)
+ if !ok {
+ return ctx, nil, fmt.Errorf("request type changed from %T to %T", payload, req.Payload)
+ }
+
+ return withCallMeta(ctx, req.Meta), newPayload, nil
}
func (c *Client) interceptAfter(ctx context.Context, method string, payload any, err error) (any, error) {
2 participants
No description provided.