diff --git a/cachecontrol/adapter.py b/cachecontrol/adapter.py index 4f4c185..d23d74a 100644 --- a/cachecontrol/adapter.py +++ b/cachecontrol/adapter.py @@ -144,6 +144,10 @@ def _update_chunk_length( super_update_chunk_length(self) if self.chunk_left == 0: self._fp._close() # type: ignore[union-attr] + elif self.chunk_left is not None: + self._fp._set_chunk_bytes_remaining( # type: ignore[union-attr] + self.chunk_left + ) response._update_chunk_length = functools.partial( # type: ignore[method-assign] _update_chunk_length, weakref.ref(response) diff --git a/cachecontrol/filewrapper.py b/cachecontrol/filewrapper.py index 6569fb5..5185cbb 100644 --- a/cachecontrol/filewrapper.py +++ b/cachecontrol/filewrapper.py @@ -37,6 +37,7 @@ def __init__( self.__buf = NamedTemporaryFile("rb+", delete=True) self.__fp = fp self.__callback = callback + self.__chunk_bytes_remaining = 0 def __getattr__(self, name: str) -> Any: # The vagaries of garbage collection means that self.__fp is @@ -107,14 +108,16 @@ def read(self, amt: int | None = None) -> bytes: return data + def _set_chunk_bytes_remaining(self, chunk_bytes_remaining: int) -> None: + self.__chunk_bytes_remaining = chunk_bytes_remaining + def _safe_read(self, amt: int) -> bytes: data: bytes = self.__fp._safe_read(amt) # type: ignore[attr-defined] - if amt == 2 and data == b"\r\n": - # urllib executes this read to toss the CRLF at the end - # of the chunk. + if self.__chunk_bytes_remaining == 0 and amt == 2 and data == b"\r\n": return data self.__buf.write(data) + self.__chunk_bytes_remaining -= len(data) if self.__is_fp_closed(): self._close() diff --git a/tests/conftest.py b/tests/conftest.py index 00f4da9..de55a33 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -101,6 +101,14 @@ def stream(self, env, start_response): for i in range(10): yield pformat(i).encode("utf8") + def stream_with_crlf(self, env, start_response): + headers = [("Content-Type", "text/plain"), ("Cache-Control", "max-age=5000")] + start_response("200 OK", headers) + + yield b"AA" + yield b"\r\n" + yield b"BB" + def fixed_length(self, env, start_response): body = b"0123456789" headers = [ diff --git a/tests/test_chunked_response.py b/tests/test_chunked_response.py index 8cc4496..a79a713 100644 --- a/tests/test_chunked_response.py +++ b/tests/test_chunked_response.py @@ -49,6 +49,14 @@ def test_stream_is_cached(self, url, sess): assert resp_2.from_cache assert content_1 == content_2 + def test_stream_with_crlf_chunk_is_cached_without_corruption(self, url, sess): + resp_1 = sess.get(url + "stream_with_crlf") + resp_2 = sess.get(url + "stream_with_crlf") + + assert resp_1.content == b"AA\r\nBB" + assert resp_2.from_cache + assert resp_2.content == resp_1.content + def test_stream_is_not_cached_when_content_is_not_read(self, url, sess): sess.get(url + "stream", stream=True) resp = sess.get(url + "stream", stream=True)