File: test_chunked_transfer.py

package info (click to toggle)
python-urllib3 1.26.12-1%2Bdeb12u1
  • links: PTS, VCS
  • area: main
  • in suites: bookworm
  • size: 2,652 kB
  • sloc: python: 17,628; makefile: 130; sh: 18
file content (263 lines) | stat: -rw-r--r-- 10,051 bytes parent folder | download
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
# -*- coding: utf-8 -*-

import pytest

from dummyserver.testcase import (
    ConnectionMarker,
    SocketDummyServerTestCase,
    consume_socket,
)
from urllib3 import HTTPConnectionPool
from urllib3.util import SKIP_HEADER
from urllib3.util.retry import Retry

# Retry failed tests
pytestmark = pytest.mark.flaky


class TestChunkedTransfer(SocketDummyServerTestCase):
    def start_chunked_handler(self):
        self.buffer = b""

        def socket_handler(listener):
            sock = listener.accept()[0]

            while not self.buffer.endswith(b"\r\n0\r\n\r\n"):
                self.buffer += sock.recv(65536)

            sock.send(
                b"HTTP/1.1 200 OK\r\n"
                b"Content-type: text/plain\r\n"
                b"Content-Length: 0\r\n"
                b"\r\n"
            )
            sock.close()

        self._start_server(socket_handler)

    def test_chunks(self):
        self.start_chunked_handler()
        chunks = ["foo", "bar", "", "bazzzzzzzzzzzzzzzzzzzzzz"]
        with HTTPConnectionPool(self.host, self.port, retries=False) as pool:
            pool.urlopen("GET", "/", body=chunks, headers=dict(DNT="1"), chunked=True)

            assert b"Transfer-Encoding" in self.buffer
            body = self.buffer.split(b"\r\n\r\n", 1)[1]
            lines = body.split(b"\r\n")
            # Empty chunks should have been skipped, as this could not be distinguished
            # from terminating the transmission
            for i, chunk in enumerate([c for c in chunks if c]):
                assert lines[i * 2] == hex(len(chunk))[2:].encode("utf-8")
                assert lines[i * 2 + 1] == chunk.encode("utf-8")

    def _test_body(self, data):
        self.start_chunked_handler()
        with HTTPConnectionPool(self.host, self.port, retries=False) as pool:
            pool.urlopen("GET", "/", data, chunked=True)
            header, body = self.buffer.split(b"\r\n\r\n", 1)

            assert b"Transfer-Encoding: chunked" in header.split(b"\r\n")
            if data:
                bdata = data if isinstance(data, bytes) else data.encode("utf-8")
                assert b"\r\n" + bdata + b"\r\n" in body
                assert body.endswith(b"\r\n0\r\n\r\n")

                len_str = body.split(b"\r\n", 1)[0]
                stated_len = int(len_str, 16)
                assert stated_len == len(bdata)
            else:
                assert body == b"0\r\n\r\n"

    def test_bytestring_body(self):
        self._test_body(b"thisshouldbeonechunk\r\nasdf")

    def test_unicode_body(self):
        self._test_body(u"thisshouldbeonechunk\r\näöüß")

    def test_empty_body(self):
        self._test_body(None)

    def test_empty_string_body(self):
        self._test_body("")

    def test_empty_iterable_body(self):
        self._test_body([])

    def _get_header_lines(self, prefix):
        header_block = self.buffer.split(b"\r\n\r\n", 1)[0].lower()
        header_lines = header_block.split(b"\r\n")[1:]
        return [x for x in header_lines if x.startswith(prefix)]

    def test_removes_duplicate_host_header(self):
        self.start_chunked_handler()
        chunks = ["foo", "bar", "", "bazzzzzzzzzzzzzzzzzzzzzz"]
        with HTTPConnectionPool(self.host, self.port, retries=False) as pool:
            pool.urlopen(
                "GET", "/", body=chunks, headers={"Host": "test.org"}, chunked=True
            )

            host_headers = self._get_header_lines(b"host")
            assert len(host_headers) == 1

    def test_provides_default_host_header(self):
        self.start_chunked_handler()
        chunks = ["foo", "bar", "", "bazzzzzzzzzzzzzzzzzzzzzz"]
        with HTTPConnectionPool(self.host, self.port, retries=False) as pool:
            pool.urlopen("GET", "/", body=chunks, chunked=True)

            host_headers = self._get_header_lines(b"host")
            assert len(host_headers) == 1

    def test_provides_default_user_agent_header(self):
        self.start_chunked_handler()
        chunks = ["foo", "bar", "", "bazzzzzzzzzzzzzzzzzzzzzz"]
        with HTTPConnectionPool(self.host, self.port, retries=False) as pool:
            pool.urlopen("GET", "/", body=chunks, chunked=True)

            ua_headers = self._get_header_lines(b"user-agent")
            assert len(ua_headers) == 1

    def test_preserve_user_agent_header(self):
        self.start_chunked_handler()
        chunks = ["foo", "bar", "", "bazzzzzzzzzzzzzzzzzzzzzz"]
        with HTTPConnectionPool(self.host, self.port, retries=False) as pool:
            pool.urlopen(
                "GET",
                "/",
                body=chunks,
                headers={"user-Agent": "test-agent"},
                chunked=True,
            )

            ua_headers = self._get_header_lines(b"user-agent")
            # Validate that there is only one User-Agent header.
            assert len(ua_headers) == 1
            # Validate that the existing User-Agent header is the one that was
            # provided.
            assert ua_headers[0] == b"user-agent: test-agent"

    def test_remove_user_agent_header(self):
        self.start_chunked_handler()
        chunks = ["foo", "bar", "", "bazzzzzzzzzzzzzzzzzzzzzz"]
        with HTTPConnectionPool(self.host, self.port, retries=False) as pool:
            pool.urlopen(
                "GET",
                "/",
                body=chunks,
                headers={"User-Agent": SKIP_HEADER},
                chunked=True,
            )

            ua_headers = self._get_header_lines(b"user-agent")
            assert len(ua_headers) == 0

    def test_provides_default_transfer_encoding_header(self):
        self.start_chunked_handler()
        chunks = ["foo", "bar", "", "bazzzzzzzzzzzzzzzzzzzzzz"]
        with HTTPConnectionPool(self.host, self.port, retries=False) as pool:
            pool.urlopen("GET", "/", body=chunks, chunked=True)

            te_headers = self._get_header_lines(b"transfer-encoding")
            assert len(te_headers) == 1

    def test_preserve_transfer_encoding_header(self):
        self.start_chunked_handler()
        chunks = ["foo", "bar", "", "bazzzzzzzzzzzzzzzzzzzzzz"]
        with HTTPConnectionPool(self.host, self.port, retries=False) as pool:
            pool.urlopen(
                "GET",
                "/",
                body=chunks,
                headers={"transfer-Encoding": "test-transfer-encoding"},
                chunked=True,
            )

            te_headers = self._get_header_lines(b"transfer-encoding")
            # Validate that there is only one Transfer-Encoding header.
            assert len(te_headers) == 1
            # Validate that the existing Transfer-Encoding header is the one that
            # was provided.
            assert te_headers[0] == b"transfer-encoding: test-transfer-encoding"

    def test_preserve_chunked_on_retry_after(self):
        self.chunked_requests = 0
        self.socks = []

        def socket_handler(listener):
            for _ in range(2):
                sock = listener.accept()[0]
                self.socks.append(sock)
                request = consume_socket(sock)
                if b"Transfer-Encoding: chunked" in request.split(b"\r\n"):
                    self.chunked_requests += 1

                sock.send(
                    b"HTTP/1.1 429 Too Many Requests\r\n"
                    b"Content-Type: text/plain\r\n"
                    b"Retry-After: 1\r\n"
                    b"Content-Length: 0\r\n"
                    b"Connection: close\r\n"
                    b"\r\n"
                )

        self._start_server(socket_handler)
        with HTTPConnectionPool(self.host, self.port) as pool:
            retries = Retry(total=1)
            pool.urlopen("GET", "/", chunked=True, retries=retries)
            for sock in self.socks:
                sock.close()
        assert self.chunked_requests == 2

    def test_preserve_chunked_on_redirect(self, monkeypatch):
        self.chunked_requests = 0

        def socket_handler(listener):
            for i in range(2):
                sock = listener.accept()[0]
                request = ConnectionMarker.consume_request(sock)
                if b"Transfer-Encoding: chunked" in request.split(b"\r\n"):
                    self.chunked_requests += 1

                if i == 0:
                    sock.sendall(
                        b"HTTP/1.1 301 Moved Permanently\r\n"
                        b"Location: /redirect\r\n\r\n"
                    )
                else:
                    sock.sendall(b"HTTP/1.1 200 OK\r\n\r\n")
                sock.close()

        self._start_server(socket_handler)
        with ConnectionMarker.mark(monkeypatch):
            with HTTPConnectionPool(self.host, self.port) as pool:
                retries = Retry(redirect=1)
                pool.urlopen(
                    "GET", "/", chunked=True, preload_content=False, retries=retries
                )
        assert self.chunked_requests == 2

    def test_preserve_chunked_on_broken_connection(self, monkeypatch):
        self.chunked_requests = 0

        def socket_handler(listener):
            for i in range(2):
                sock = listener.accept()[0]
                request = ConnectionMarker.consume_request(sock)
                if b"Transfer-Encoding: chunked" in request.split(b"\r\n"):
                    self.chunked_requests += 1

                if i == 0:
                    # Bad HTTP version will trigger a connection close
                    sock.sendall(b"HTTP/0.5 200 OK\r\n\r\n")
                else:
                    sock.sendall(b"HTTP/1.1 200 OK\r\n\r\n")
                sock.close()

        self._start_server(socket_handler)
        with ConnectionMarker.mark(monkeypatch):
            with HTTPConnectionPool(self.host, self.port) as pool:
                retries = Retry(read=1)
                pool.urlopen(
                    "GET", "/", chunked=True, preload_content=False, retries=retries
                )
            assert self.chunked_requests == 2