diff --git a/client.go b/client.go index 13c8bc0..b4e2189 100644 --- a/client.go +++ b/client.go @@ -1,7 +1,9 @@ package typhon import ( + "bytes" "context" + "io" "net/http" "time" @@ -67,7 +69,18 @@ func (f *ResponseFuture) Response() Response { func HttpService(rt http.RoundTripper) Service { return func(req Request) Response { ctx := req.unwrappedContext() - httpRsp, err := rt.RoundTrip(req.Request.WithContext(ctx)) + httpReq := req.Request.WithContext(ctx) + // Expose GetBody for buffered bodies so the HTTP/2 transport can rewind and transparently + // retry the request when a connection is torn down before the response arrives (eg. a + // graceful-shutdown GOAWAY). Without GetBody the transport can't replay the body and the + // request fails outright. Streaming bodies aren't a *bufCloser, so they stay non-replayable. + if buf, ok := httpReq.Body.(*bufCloser); ok { + body := buf.Bytes() + httpReq.GetBody = func() (io.ReadCloser, error) { + return io.NopCloser(bytes.NewReader(body)), nil + } + } + httpRsp, err := rt.RoundTrip(httpReq) // When the calling context is cancelled, close the response body // This protects callers that forget to call Close(), or those which proxy responses upstream // diff --git a/client_test.go b/client_test.go new file mode 100644 index 0000000..e067177 --- /dev/null +++ b/client_test.go @@ -0,0 +1,66 @@ +package typhon + +import ( + "bytes" + "context" + "io" + "net/http" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/monzo/typhon/prototest" +) + +type roundTripperFunc func(*http.Request) (*http.Response, error) + +func (f roundTripperFunc) RoundTrip(r *http.Request) (*http.Response, error) { + return f(r) +} + +// TestHttpServiceSetsGetBody verifies that HttpService populates Request.GetBody for buffered bodies, +// so the HTTP/2 transport can rewind and transparently retry the request after a mid-flight connection +// teardown (eg. a graceful-shutdown GOAWAY). +func TestHttpServiceSetsGetBody(t *testing.T) { + t.Parallel() + + req := NewRequest(context.Background(), "POST", "/", nil) + req.EncodeAsProtobuf(&prototest.Greeting{Message: "Hello world!"}) + want, err := req.BodyBytes(false) + require.NoError(t, err) + + var captured *http.Request + svc := HttpService(roundTripperFunc(func(r *http.Request) (*http.Response, error) { + captured = r + return &http.Response{StatusCode: http.StatusOK, Body: io.NopCloser(bytes.NewReader(nil))}, nil + })) + rsp := svc(req) + require.NoError(t, rsp.Error) + + require.NotNil(t, captured.GetBody, "GetBody should be set for a buffered body") + rc, err := captured.GetBody() + require.NoError(t, err) + got, err := io.ReadAll(rc) + require.NoError(t, err) + assert.Equal(t, want, got) +} + +// TestHttpServiceStreamingBodyHasNoGetBody verifies that streaming (non-buffered) bodies remain +// non-replayable - we must not claim a stream can be rewound. +func TestHttpServiceStreamingBodyHasNoGetBody(t *testing.T) { + t.Parallel() + + req := NewRequest(context.Background(), "POST", "/", nil) + req.Encode(io.NopCloser(bytes.NewReader([]byte("streamed")))) + + var captured *http.Request + svc := HttpService(roundTripperFunc(func(r *http.Request) (*http.Response, error) { + captured = r + return &http.Response{StatusCode: http.StatusOK, Body: io.NopCloser(bytes.NewReader(nil))}, nil + })) + rsp := svc(req) + require.NoError(t, rsp.Error) + + assert.Nil(t, captured.GetBody, "GetBody must stay nil for a streaming body") +}