Fix binary responses, header case-insensitivity, Content-Type leak; add timeout

tests/integration/test_httpbin.py

@@ -3,6 +3,7 @@
 from collections.abc import Iterator
 
 import pytest
+from httpbin.core import resource as httpbin_resource
 
 from tls_client import Session
 from tls_client.exceptions import TLSClientException
@@ -78,13 +79,11 @@ def test_utf8_body_is_decoded(self, httpbin, session):
 
         assert "∮ E⋅da = Q" in response.text
 
-    @pytest.mark.xfail(reason="the library hands the body over as text, so bytes that are not UTF-8 get mangled")
     def test_binary_body_arrives_unchanged(self, httpbin, session):
-        payload = b"\x89PNG\r\n\x1a\n"
+        response = session.get(f"{httpbin.url}/image/png")
 
-        response = session.get(f"{httpbin.url}/base64/{base64.urlsafe_b64encode(payload).decode()}")
-
-        assert response.content == payload
+        # httpbin serves this file from its package data.
+        assert response.content == httpbin_resource("images/pig_icon.png")
 
 
 class TestRedirects:
@@ -130,11 +129,12 @@ def test_session_timeout_raises(self, httpbin, session):
             session.get(f"{httpbin.url}/delay/2")
 
     @pytest.mark.xfail(reason="the native session keeps the timeout of its first request")
-    def test_request_timeout_applies_after_the_first_request(self, httpbin, session):
+    @pytest.mark.parametrize("kwarg", ["timeout", "timeout_seconds"])
+    def test_request_timeout_applies_after_the_first_request(self, httpbin, session, kwarg):
         session.get(f"{httpbin.url}/get")
 
         with pytest.raises(TLSClientException, match="Timeout exceeded"):
-            session.get(f"{httpbin.url}/delay/2", timeout_seconds=1)
+            session.get(f"{httpbin.url}/delay/2", **{kwarg: 1})
 
 
 class TestCertificateVerification:

tests/integration/test_local_server.py

@@ -0,0 +1,165 @@
+"""Real requests through the native library against a stdlib HTTP server, for bodies and headers httpbin can't
+produce byte for byte."""
+
+import base64
+import json
+import os
+import threading
+import time
+from collections.abc import Iterator
+from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
+
+import pytest
+
+from tls_client import Session
+from tls_client.exceptions import TLSClientException
+
+pytestmark = pytest.mark.integration
+
+# A 1x1 PNG, the kind of download that used to arrive mangled.
+PNG = base64.b64decode(
+    "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg=="
+)
+LARGE = os.urandom(8 * 1024 * 1024)
+# How long /slow takes to answer.
+SLOW_SECONDS = 0.6
+
+# Path -> (Content-Type, body)
+BODIES = {
+    "/bytes": ("application/octet-stream", bytes(range(256))),
+    "/png": ("image/png", PNG),
+    "/large": ("application/octet-stream", LARGE),
+    "/empty": ("text/plain", b""),
+    "/json": ("application/json", json.dumps({"name": "tls", "ids": [1, 2]}).encode()),
+    "/utf8": ("text/plain", "∮ E⋅da = Q".encode()),
+    "/latin-1": ("text/plain; charset=iso-8859-1", "café".encode("latin-1")),
+}
+
+
+class Handler(BaseHTTPRequestHandler):
+    protocol_version = "HTTP/1.1"
+
+    def do_GET(self) -> None:
+        if self.path == "/headers":
+            self.echo_headers()
+        elif self.path == "/slow":
+            time.sleep(SLOW_SECONDS)
+            self.reply("text/plain", b"done")
+        else:
+            content_type, body = BODIES[self.path]
+            self.reply(content_type, body, {"x-MiXeD-case": "yes"})
+
+    def do_POST(self) -> None:
+        self.rfile.read(int(self.headers.get("Content-Length", 0)))
+        self.echo_headers()
+
+    def echo_headers(self) -> None:
+        received = {name.lower(): value for name, value in self.headers.items()}
+        self.reply("application/json", json.dumps(received).encode())
+
+    def reply(self, content_type: str, body: bytes, extra_headers: dict[str, str] | None = None) -> None:
+        self.send_response(200)
+        self.send_header("Content-Type", content_type)
+        self.send_header("Content-Length", str(len(body)))
+        for name, value in (extra_headers or {}).items():
+            self.send_header(name, value)
+        try:
+            self.end_headers()
+            self.wfile.write(body)
+        except ConnectionError:
+            # The client gave up waiting, as the timeout tests make it do.
+            pass
+
+    def log_message(self, format: str, *args: object) -> None:
+        pass
+
+
+@pytest.fixture(scope="module")
+def server_url() -> Iterator[str]:
+    server = ThreadingHTTPServer(("127.0.0.1", 0), Handler)
+    server.daemon_threads = True
+    thread = threading.Thread(target=server.serve_forever, daemon=True)
+    thread.start()
+    yield f"http://127.0.0.1:{server.server_port}"
+    server.shutdown()
+    server.server_close()
+
+
+@pytest.fixture
+def session() -> Iterator[Session]:
+    with Session() as s:
+        yield s
+
+
+class TestBody:
+    @pytest.mark.parametrize("path", ["/bytes", "/png", "/large"])
+    def test_binary_body_arrives_byte_for_byte(self, server_url, session, path):
+        response = session.get(f"{server_url}{path}")
+
+        assert response.status_code == 200
+        assert response.content == BODIES[path][1]
+
+    def test_empty_body(self, server_url, session):
+        response = session.get(f"{server_url}/empty")
+
+        assert response.content == b""
+        assert response.text == ""
+
+    def test_json_body(self, server_url, session):
+        assert session.get(f"{server_url}/json").json() == {"name": "tls", "ids": [1, 2]}
+
+    def test_text_defaults_to_utf8(self, server_url, session):
+        assert session.get(f"{server_url}/utf8").text == "∮ E⋅da = Q"
+
+    def test_text_uses_the_charset_from_the_content_type_header(self, server_url, session):
+        response = session.get(f"{server_url}/latin-1")
+
+        assert response.content == "café".encode("latin-1")
+        assert response.encoding == "iso-8859-1"
+        assert response.text == "café"
+
+
+class TestHeaders:
+    @pytest.mark.parametrize("name", ["X-Mixed-Case", "x-mixed-case", "X-MIXED-CASE"])
+    def test_response_header_lookup_ignores_case(self, server_url, session, name):
+        response = session.get(f"{server_url}/bytes")
+
+        assert response.headers[name] == "yes"
+        assert response.headers["content-type"] == "application/octet-stream"
+
+    def test_inferred_content_type_does_not_leak_into_later_requests(self, server_url, session):
+        json_post = session.post(f"{server_url}/headers", json={"a": 1}).json()
+        form_post = session.post(f"{server_url}/headers", data={"a": "1"}).json()
+        get = session.get(f"{server_url}/headers").json()
+
+        assert json_post["content-type"] == "application/json"
+        assert form_post["content-type"] == "application/x-www-form-urlencoded"
+        assert "content-type" not in get
+        assert "Content-Type" not in session.headers
+
+    def test_content_type_from_the_session_or_the_request_is_kept(self, server_url, session):
+        session.headers["content-type"] = "application/vnd.api+json"
+
+        from_session = session.post(f"{server_url}/headers", json={"a": 1}).json()
+        from_request = session.post(
+            f"{server_url}/headers", json={"a": 1}, headers={"CONTENT-TYPE": "text/plain"}
+        ).json()
+
+        assert from_session["content-type"] == "application/vnd.api+json"
+        assert from_request["content-type"] == "text/plain"
+
+
+class TestTimeout:
+    @pytest.mark.parametrize("kwarg", ["timeout", "timeout_seconds"])
+    def test_fractional_timeout_is_applied(self, server_url, session, kwarg):
+        with pytest.raises(TLSClientException, match="Timeout exceeded"):
+            session.get(f"{server_url}/slow", **{kwarg: SLOW_SECONDS / 3})
+
+    def test_response_within_the_timeout_arrives(self, server_url, session):
+        assert session.get(f"{server_url}/slow", timeout=SLOW_SECONDS * 4).text == "done"
+
+    def test_fractional_session_timeout_is_applied(self, server_url, session):
+        session.timeout_seconds = SLOW_SECONDS / 3
+
+        with pytest.raises(TLSClientException, match="Timeout exceeded"):
+            session.get(f"{server_url}/slow")

tests/unit/test_response.py

@@ -1,3 +1,6 @@
+import base64
+import os
+import tracemalloc
 from typing import Any
 
 import pytest
@@ -6,16 +9,28 @@
 from tls_client.response import build_response
 
 
+def data_uri(body: bytes, mime_type: str = "text/plain; charset=utf-8") -> str:
+    """A body the way the Go library returns it for isByteResponse requests; it sniffs the mime type itself."""
+    return f"data:{mime_type};base64,{base64.b64encode(body).decode()}"
+
+
 def library_reply(**fields: Any) -> dict[str, Any]:
     """A reply shaped like the one the Go library returns."""
-    return {"id": "response-1", "status": 200, "target": "https://example.com/", "body": "", "headers": {}, **fields}
+    return {
+        "id": "response-1",
+        "status": 200,
+        "target": "https://example.com/",
+        "body": data_uri(b""),
+        "headers": {},
+        **fields,
+    }
 
 
 def test_reply_fields_are_mapped_onto_the_response():
     cookies = cookiejar_from_dict({"sid": "abc"})
 
     response = build_response(
-        library_reply(status=404, target="https://example.com/final", body='{"error": "missing"}'),
+        library_reply(status=404, target="https://example.com/final", body=data_uri(b'{"error": "missing"}')),
         cookies,
     )
 
@@ -40,8 +55,138 @@ def test_null_headers_become_an_empty_mapping():
     assert dict(response.headers) == {}
 
 
-@pytest.mark.xfail(reason="build_response replaces the CaseInsensitiveDict with a plain dict")
-def test_header_lookup_ignores_case():
+@pytest.mark.parametrize("name", ["Content-Type", "content-type", "CONTENT-TYPE"])
+def test_header_lookup_ignores_case(name):
     response = build_response(library_reply(headers={"Content-Type": ["text/html"]}), cookiejar_from_dict({}))
 
-    assert response.headers["content-type"] == "text/html"
+    assert response.headers[name] == "text/html"
+    assert response.headers.get(name) == "text/html"
+    assert name in response.headers
+
+
+class TestBody:
+    def test_binary_body_is_decoded_byte_for_byte(self):
+        body = bytes(range(256))
+
+        response = build_response(
+            library_reply(body=data_uri(body, "application/octet-stream")), cookiejar_from_dict({})
+        )
+
+        assert response.content == body
+
+    def test_empty_body(self):
+        response = build_response(library_reply(), cookiejar_from_dict({}))
+
+        assert response.content == b""
+        assert response.text == ""
+
+    def test_text_uses_the_charset_from_the_content_type_header(self):
+        # The library labels the data URI with a sniffed type; only the header names the real charset.
+        response = build_response(
+            library_reply(
+                body=data_uri("café".encode("latin-1")),
+                headers={"content-type": ['text/plain; Charset="ISO-8859-1"']},
+            ),
+            cookiejar_from_dict({}),
+        )
+
+        assert response.encoding == "ISO-8859-1"
+        assert response.text == "café"
+
+    def test_text_defaults_to_utf8_and_replaces_undecodable_bytes(self):
+        response = build_response(
+            library_reply(body=data_uri("∮ ok ".encode() + b"\xff"), headers={"Content-Type": ["text/plain"]}),
+            cookiejar_from_dict({}),
+        )
+
+        assert response.encoding is None
+        assert response.text == "∮ ok �"
+
+    def test_unknown_charset_falls_back_to_utf8(self):
+        response = build_response(
+            library_reply(body=data_uri("é".encode()), headers={"Content-Type": ["text/plain; charset=made-up"]}),
+            cookiejar_from_dict({}),
+        )
+
+        assert response.text == "é"
+
+    def test_text_follows_an_encoding_set_by_the_caller(self):
+        response = build_response(library_reply(body=data_uri("café".encode("cp1252"))), cookiejar_from_dict({}))
+
+        response.encoding = "cp1252"
+
+        assert response.text == "café"
+
+    @pytest.mark.parametrize(
+        "html",
+        [
+            '<html><head><meta charset="gbk"><title>{}</title>',
+            "<html><head><META HTTP-EQUIV='Content-Type' CONTENT='text/html; charset=GBK'><title>{}</title>",
+            "<html><head><meta charset=gb2312><title>{}</title>",
+        ],
+    )
+    def test_text_uses_the_charset_declared_in_a_meta_tag(self, html):
+        response = build_response(
+            library_reply(body=data_uri(html.format("中文").encode("gbk")), headers={"Content-Type": ["text/html"]}),
+            cookiejar_from_dict({}),
+        )
+
+        assert response.encoding.lower() == "gbk"
+        assert "<title>中文</title>" in response.text
+
+    def test_text_uses_the_encoding_of_an_xml_declaration(self):
+        body = '<?xml version="1.0" encoding="Shift_JIS"?><a>日本</a>'.encode("shift_jis")
+
+        response = build_response(
+            library_reply(body=data_uri(body), headers={"Content-Type": ["application/xml"]}), cookiejar_from_dict({})
+        )
+
+        assert response.text.endswith("<a>日本</a>")
+
+    def test_header_charset_wins_over_a_meta_tag(self):
+        body = '<meta charset="gbk">café'.encode("latin-1")
+
+        response = build_response(
+            library_reply(body=data_uri(body), headers={"Content-Type": ["text/html; charset=ISO-8859-1"]}),
+            cookiejar_from_dict({}),
+        )
+
+        assert response.encoding == "ISO-8859-1"
+        assert response.text.endswith("café")
+
+    def test_meta_tags_are_ignored_outside_html_and_xml(self):
+        body = b'{"html": "<meta charset=gbk>"}'
+
+        response = build_response(
+            library_reply(body=data_uri(body), headers={"Content-Type": ["application/json"]}), cookiejar_from_dict({})
+        )
+
+        assert response.encoding is None
+
+    @pytest.mark.parametrize("encoding", ["utf-8-sig", "utf-16"])
+    def test_text_uses_and_drops_a_byte_order_mark(self, encoding):
+        response = build_response(library_reply(body=data_uri("héllo".encode(encoding))), cookiejar_from_dict({}))
+
+        assert response.text == "héllo"
+
+    def test_body_that_is_not_a_data_uri_is_taken_as_text(self):
+        response = build_response(library_reply(body='{"plain": true}'), cookiejar_from_dict({}))
+
+        assert response.json() == {"plain": True}
+
+    def test_large_body_is_not_held_in_several_copies_at_once(self):
+        body = os.urandom(8_000_000)
+        tracemalloc.start()
+        try:
+            reply = library_reply(body=data_uri(body))
+            tracemalloc.reset_peak()
+            response = build_response(reply, cookiejar_from_dict({}))
+            _, peak = tracemalloc.get_traced_memory()
+        finally:
+            tracemalloc.stop()
+
+        assert response.content == body
+        assert "body" not in reply
+        # The data URI and its payload are each 4/3 of the body. Freeing the data URI before decoding keeps the
+        # peak near 8/3 of the body; holding it, the payload and the decoded bytes together would reach 11/3.
+        assert peak < 3 * len(body)