From 79cca8cc18e49ce8f767dd50035a0298a2dc68d0 Mon Sep 17 00:00:00 2001
From: Jacob Quinn <quinn.jacobd@gmail.com>
Date: Mon, 11 Jul 2022 07:37:10 -0600
Subject: [PATCH] Preserve response body when retrying/redirecting

This is another iteration in our redirect/retry response body handling.
Previously, when the response body is streamed, we ensured only the
"final" response body is the one written to the user-provided stream.
The only issue is that sometimes it can be helpful to have access to
the intermediate response bodies that ended up being retried/redirected.
This PR refactors the `readbody` machinery a bit to ensure we always read
the response body, and in the case of a user-provided response stream,
we read the response body as a string and store it in the request context.

One motivation I'm working towards with this is providing more nuanced control
over the retry logic, even allowing users to pass in their own `check` function
to determine if a request should be retried or not based on the failed response.
In that scenario, you'll want access to the failed response body, so this will
enabled that.
---
 src/clientlayers/StreamRequest.jl | 34 +++++++++++++++----------------
 test/client.jl                    |  7 +++++++
 2 files changed, 24 insertions(+), 17 deletions(-)

diff --git a/src/clientlayers/StreamRequest.jl b/src/clientlayers/StreamRequest.jl
index 0fe94e8a2..ff718bee1 100644
--- a/src/clientlayers/StreamRequest.jl
+++ b/src/clientlayers/StreamRequest.jl
@@ -103,13 +103,7 @@ writechunk(stream, body::IO) = writebodystream(stream, body)
 writechunk(stream, body::Union{Dict, NamedTuple}) = writebodystream(stream, body)
 writechunk(stream, body) = write(stream, body)
 
-function readbody(stream::Stream, res::Response, decompress)
-    # Bail early if we are not going to read anything.
-    # If the request/response pair are going to be redirected or retried,
-    # we want to avoid "contaminating" our response body stream.
-    willread = isbytes(res.body) || (!isredirect(stream) && !retryable(stream))
-    willread || return
-
+function readbody(stream::Stream, res::Response, decompress::Bool)
     if decompress && header(res, "Content-Encoding") == "gzip"
         # Plug in a buffer stream in between so that we can (i) read the http stream in
         # chunks instead of byte-by-byte and (ii) make sure to stop reading the http stream
@@ -127,18 +121,24 @@ function readbody(stream::Stream, res::Response, decompress)
                 close(gzstream)
             end
         end
-        if isbytes(res.body)
-            res.body = read(buf)
-        else
-            write(res.body, buf)
-        end
+        readbody!(stream, res, buf)
         wait(tsk)
     else
-        if isbytes(res.body)
-            res.body = read(stream)
-        else
-            write(res.body, stream)
-        end
+        readbody!(stream, res, stream)
+    end
+end
+
+function readbody!(stream::Stream, res::Response, buf_or_stream)
+    if isbytes(res.body)
+        # normal response body path: read as Vector{UInt8} and store
+        res.body = read(buf_or_stream)
+    elseif isredirect(stream) || retryable(stream)
+        # if response body is a stream, but we're redirecting or
+        # retrying, store this "temporary" body in the request context
+        res.request.context[:response_body] = read(buf_or_stream)
+    else
+        # normal streaming response body path: write response body out directly
+        write(res.body, buf_or_stream)
     end
 end
 
diff --git a/test/client.jl b/test/client.jl
index 2d59237c7..ee3b44704 100644
--- a/test/client.jl
+++ b/test/client.jl
@@ -537,6 +537,13 @@ end
         seekstart(req_body)
         resp = HTTP.get("http://localhost:8080/retry"; body=req_body, response_stream=res_body, retry=false, status_exception=false)
         @test String(take!(res_body)) == "500 unexpected error"
+        # when retrying, we can still get access to the most recent failed response body in the response's request context
+        shouldfail[] = true
+        seekstart(req_body)
+        resp = HTTP.get("http://localhost:8080/retry"; body=req_body, response_stream=res_body)
+        @test resp.status == 200
+        @test String(take!(res_body)) == "hey there sailor"
+        @test String(resp.request.context[:response_body]) == "500 unexpected error"
     finally
         if server !== nothing
             close(server)