-
-
Notifications
You must be signed in to change notification settings - Fork 450
Expand file tree
/
Copy pathhttp_reader.py
More file actions
166 lines (140 loc) · 5.43 KB
/
Copy pathhttp_reader.py
File metadata and controls
166 lines (140 loc) · 5.43 KB
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
"""
HTTP/1.1 response reader for keep-alive connections.
Reads exactly one HTTP response from an asyncio StreamReader, handling
chunked transfer-encoding, Content-Length framing, and streaming bodies.
Auto-decompresses the response body according to the Content-Encoding
header (gzip, deflate, brotli, zstd).
Usage::
status, headers, body = await read_http_response(reader, max_bytes=50_000_000)
"""
from __future__ import annotations
import asyncio
import re
from core import codec
__all__ = ["read_http_response"]
async def read_http_response(
reader: asyncio.StreamReader,
*,
max_bytes: int,
) -> tuple[int, dict[str, str], bytes]:
"""Read one HTTP/1.1 response. Keep-alive safe (no read-until-EOF).
Args:
reader: An ``asyncio.StreamReader`` positioned at the start of
an HTTP response.
max_bytes: Hard cap on the decompressed body size. Raises
``RuntimeError`` if exceeded.
Returns:
A ``(status_code, headers, body)`` triple. ``status_code`` is 0
and the other fields are empty/empty if the response is malformed.
"""
# ── Read until header boundary ────────────────────────────────
raw = b""
while b"\r\n\r\n" not in raw:
if len(raw) > 65536: # 64 KB header size limit
return 0, {}, b""
# 30s per-read: exit-node chain (Apps Script → VPS → target) needs
# time to fetch + process large responses before sending headers.
# The outer asyncio.wait_for in _relay_single caps total time.
chunk = await asyncio.wait_for(reader.read(8192), timeout=30)
if not chunk:
break
raw += chunk
if b"\r\n\r\n" not in raw:
return 0, {}, b""
header_section, body = raw.split(b"\r\n\r\n", 1)
lines = header_section.split(b"\r\n")
status_line = lines[0].decode(errors="replace")
m = re.search(r"\d{3}", status_line)
status = int(m.group()) if m else 0
headers: dict[str, str] = {}
for line in lines[1:]:
if b":" in line:
k, v = line.decode(errors="replace").split(":", 1)
headers[k.strip().lower()] = v.strip()
# ── Body framing ──────────────────────────────────────────────
content_length = headers.get("content-length")
transfer_encoding = headers.get("transfer-encoding", "")
if "chunked" in transfer_encoding:
body = await _read_chunked(reader, body, max_bytes=max_bytes)
elif content_length:
total = int(content_length)
if total > max_bytes:
raise RuntimeError(
"Relay response exceeds configured size cap "
f"({total} > {max_bytes} bytes)"
)
remaining = total - len(body)
while remaining > 0:
chunk = await asyncio.wait_for(
reader.read(min(remaining, 65536)), timeout=20
)
if not chunk:
break
body += chunk
if len(body) > max_bytes:
raise RuntimeError(
"Relay response exceeded configured size cap while reading body"
)
remaining -= len(chunk)
else:
# No framing — short timeout read (keep-alive safe)
while True:
try:
chunk = await asyncio.wait_for(reader.read(65536), timeout=2)
if not chunk:
break
body += chunk
if len(body) > max_bytes:
raise RuntimeError(
"Relay response exceeded configured size cap while streaming"
)
except asyncio.TimeoutError:
break
# ── Auto-decompress ───────────────────────────────────────────
enc = headers.get("content-encoding", "")
if enc:
body = codec.decode(body, enc)
if len(body) > max_bytes:
raise RuntimeError(
"Decoded relay response exceeded configured size cap"
)
return status, headers, body
async def _read_chunked(
reader: asyncio.StreamReader,
buf: bytes = b"",
*,
max_bytes: int,
) -> bytes:
"""Incrementally read a chunked-transfer-encoded body."""
result = b""
while True:
while b"\r\n" not in buf:
data = await asyncio.wait_for(reader.read(8192), timeout=20)
if not data:
return result
buf += data
end = buf.find(b"\r\n")
size_str = buf[:end].decode(errors="replace").strip()
buf = buf[end + 2:]
if not size_str:
continue
try:
size = int(size_str, 16)
except ValueError:
break
if size == 0:
break
if size > max_bytes or len(result) + size > max_bytes:
raise RuntimeError(
"Chunked relay response exceeded configured size cap "
f"({max_bytes} bytes)"
)
while len(buf) < size + 2:
data = await asyncio.wait_for(reader.read(65536), timeout=20)
if not data:
result += buf[:size]
return result
buf += data
result += buf[:size]
buf = buf[size + 2:]
return result