From 2d5b99b3ae06b05096326262b1f4c0f63c2ebf74 Mon Sep 17 00:00:00 2001 From: szeka9 Date: Tue, 28 Apr 2026 18:40:42 +0200 Subject: [PATCH 1/4] Encode content-type header once Move content-type header encoding to set_response_body() to improve readability. --- src/pyrobusta/protocol/http.py | 9 +++------ src/pyrobusta/protocol/http_multipart.py | 7 +++---- 2 files changed, 6 insertions(+), 10 deletions(-) diff --git a/src/pyrobusta/protocol/http.py b/src/pyrobusta/protocol/http.py index 93d55ab..dfa7bfb 100644 --- a/src/pyrobusta/protocol/http.py +++ b/src/pyrobusta/protocol/http.py @@ -412,7 +412,7 @@ def write_response_head(self, tx): def set_response_body( self, body: bytes | str | dict | tuple | list, - content_type: bytes = b"text/plain", + content_type: str = "text/plain", ): """ Serialize and wrap the response body with a BytesIO @@ -436,7 +436,7 @@ def set_response_body( self.set_response_header( b"content-length", str(len(body_encoded)).encode("ascii") ) - self.set_response_header(b"content-type", content_type) + self.set_response_header(b"content-type", content_type.encode("ascii")) if self.method != self.HEAD: self.resp_handler = BytesIO(body_encoded) @@ -702,7 +702,6 @@ def _app_endpoint_st(self, rx): dtype, data = callback( self, bytes(rx.peek(self.headers["content-length"])) ) - dtype = dtype.encode("ascii") else: if not callable(callback): # Handle as a file path @@ -711,10 +710,8 @@ def _app_endpoint_st(self, rx): ) return dtype, data = callback(self, b"") - dtype = dtype.encode("ascii") - self.set_response_header(b"content-type", dtype) - if dtype.startswith(b"multipart/"): + if dtype.startswith("multipart/"): self.state = lambda _rx: self._generate_multipart_response(_rx, data, dtype) return diff --git a/src/pyrobusta/protocol/http_multipart.py b/src/pyrobusta/protocol/http_multipart.py index fbc0862..c79c8b2 100644 --- a/src/pyrobusta/protocol/http_multipart.py +++ b/src/pyrobusta/protocol/http_multipart.py @@ -8,7 +8,7 @@ from pyrobusta.utils.helpers import add_method -def _generate_multipart_response(self, _, callback: callable, dtype: bytes): +def _generate_multipart_response(self, _, callback: callable, dtype: str): """ Generate multipart response depening on the exact content type. The callback function is called without arguments, and it must return bytes-like objects. @@ -19,7 +19,7 @@ def _generate_multipart_response(self, _, callback: callable, dtype: bytes): raise ValueError("Invalid response handler") self.terminate(200, True) boundary = self.MULTIPART_BOUNDARY - self.set_response_header(b"content-type", dtype + b"; boundary=" + boundary) + self.set_response_header(b"content-type", dtype.encode("ascii") + b"; boundary=" + boundary) if self.method != self.HEAD: self.resp_handler = self._multipart_wrapper_factory(callback, boundary) @@ -135,9 +135,8 @@ def _parse_complete_part_st(self, rx): raise http.InvalidContentLength() self.mp_is_last = True dtype, data = callback(self, (part_headers, part_body)) - self.set_response_header(b"content-type", dtype.encode("ascii")) self.terminate(200, True) - self.set_response_body(data) + self.set_response_body(data, dtype) def apply_patches(): From 2adaf8bacfe3779dbf0eec7821bcada24b7c0756 Mon Sep 17 00:00:00 2001 From: szeka9 Date: Tue, 28 Apr 2026 20:02:47 +0200 Subject: [PATCH 2/4] Handle combined field values As per RFC9110, allow headers with the same field name, parsed and combined into a comma-separated list of non-empty values. --- src/pyrobusta/protocol/http.py | 5 ++++- tests/unit/test_http.py | 13 +++++++++++++ 2 files changed, 17 insertions(+), 1 deletion(-) diff --git a/src/pyrobusta/protocol/http.py b/src/pyrobusta/protocol/http.py index dfa7bfb..45ba366 100644 --- a/src/pyrobusta/protocol/http.py +++ b/src/pyrobusta/protocol/http.py @@ -319,7 +319,10 @@ def _parse_headers(cls, raw_headers: memoryview) -> dict[str, str | int]: value = int(value.strip()) else: value = value.strip().decode("ascii") - headers[name] = value + if name not in headers and value: + headers[name] = value + elif value: + headers[name] += ", " + value # Combined field value return headers @staticmethod diff --git a/tests/unit/test_http.py b/tests/unit/test_http.py index 949bc2f..96be22d 100644 --- a/tests/unit/test_http.py +++ b/tests/unit/test_http.py @@ -170,6 +170,19 @@ def test_header_parsing_error(self): with self.assertRaises(self.http_module.InvalidHeaders): self.engine._parse_headers(case) + def test_header_parsing_combined(self): + for case in ( + ( + b"field-name: value1\r\nfield-name: value2", + {"field-name": "value1, value2"}, + ), + ( + b"field-name: \r\nfield-name: value1\r\nfield-name:\r\nfield-name: value2", + {"field-name": "value1, value2"}, + ), + ): + self.assertEqual(self.engine._parse_headers(case[0]), case[1]) + def test_routing_unsupported_method(self): self.engine.state = self.engine._route_request_st self.engine.url = b"/api/test" From 88c032a411f51bdc7ce602006363361072745259 Mon Sep 17 00:00:00 2001 From: szeka9 Date: Tue, 28 Apr 2026 20:08:31 +0200 Subject: [PATCH 3/4] Move file-specific content-type definitions to file server module Remove content type definitions defined in the HTTP parser class to reduce heap usage when the file server feature is turned off. --- src/pyrobusta/protocol/http.py | 24 ------------------- src/pyrobusta/protocol/http_file_server.py | 28 ++++++++++++++++++++-- 2 files changed, 26 insertions(+), 26 deletions(-) diff --git a/src/pyrobusta/protocol/http.py b/src/pyrobusta/protocol/http.py index 45ba366..abcd007 100644 --- a/src/pyrobusta/protocol/http.py +++ b/src/pyrobusta/protocol/http.py @@ -94,30 +94,6 @@ class HttpEngine: 505, b"505 Version Not Supported", ) - CONTENT_TYPES = ( - b"raw", - b"application/octet-stream", - b"html", - b"text/html", - b"css", - b"text/css", - b"js", - b"application/javascript", - b"json", - b"application/json", - b"ico", - b"image/x-icon", - b"jpeg", - b"image/jpeg", - b"jpg", - b"image/jpeg", - b"png", - b"image/png", - b"txt", - b"text/plain", - b"gif", - b"image/gif", - ) DELETE = b"DELETE" GET = b"GET" diff --git a/src/pyrobusta/protocol/http_file_server.py b/src/pyrobusta/protocol/http_file_server.py index f128d37..8a00572 100644 --- a/src/pyrobusta/protocol/http_file_server.py +++ b/src/pyrobusta/protocol/http_file_server.py @@ -9,6 +9,30 @@ from pyrobusta.protocol import http from pyrobusta.utils.helpers import normalize_path, add_method +CONTENT_TYPES = ( + b"raw", + b"application/octet-stream", + b"html", + b"text/html", + b"css", + b"text/css", + b"js", + b"application/javascript", + b"json", + b"application/json", + b"ico", + b"image/x-icon", + b"jpeg", + b"image/jpeg", + b"jpg", + b"image/jpeg", + b"png", + b"image/png", + b"txt", + b"text/plain", + b"gif", + b"image/gif", +) def _send_file_st(self, _, file_path: bytes): """ @@ -38,9 +62,9 @@ def _send_file_st(self, _, file_path: bytes): self.terminate(404, True) return try: - content_type = self._lookup(self.CONTENT_TYPES, extension) + content_type = self._lookup(CONTENT_TYPES, extension) except ValueError: - content_type = self._lookup(self.CONTENT_TYPES, b"raw") + content_type = self._lookup(CONTENT_TYPES, b"raw") try: self.set_response_header( b"content-length", str(stat(norm_path)[6]).encode("ascii") From 0a3e156cab53146f4330cb50cc99f2a9971d4532 Mon Sep 17 00:00:00 2001 From: szeka9 Date: Tue, 28 Apr 2026 20:14:07 +0200 Subject: [PATCH 4/4] Fix formatting issues --- src/pyrobusta/protocol/http.py | 2 +- src/pyrobusta/protocol/http_file_server.py | 1 + src/pyrobusta/protocol/http_multipart.py | 4 +++- 3 files changed, 5 insertions(+), 2 deletions(-) diff --git a/src/pyrobusta/protocol/http.py b/src/pyrobusta/protocol/http.py index abcd007..ccf4e2b 100644 --- a/src/pyrobusta/protocol/http.py +++ b/src/pyrobusta/protocol/http.py @@ -298,7 +298,7 @@ def _parse_headers(cls, raw_headers: memoryview) -> dict[str, str | int]: if name not in headers and value: headers[name] = value elif value: - headers[name] += ", " + value # Combined field value + headers[name] += ", " + value # Combined field value return headers @staticmethod diff --git a/src/pyrobusta/protocol/http_file_server.py b/src/pyrobusta/protocol/http_file_server.py index 8a00572..a8c278b 100644 --- a/src/pyrobusta/protocol/http_file_server.py +++ b/src/pyrobusta/protocol/http_file_server.py @@ -34,6 +34,7 @@ b"image/gif", ) + def _send_file_st(self, _, file_path: bytes): """ State for returning a file. By default, /www is prepended to the path. diff --git a/src/pyrobusta/protocol/http_multipart.py b/src/pyrobusta/protocol/http_multipart.py index c79c8b2..7e8cb5d 100644 --- a/src/pyrobusta/protocol/http_multipart.py +++ b/src/pyrobusta/protocol/http_multipart.py @@ -19,7 +19,9 @@ def _generate_multipart_response(self, _, callback: callable, dtype: str): raise ValueError("Invalid response handler") self.terminate(200, True) boundary = self.MULTIPART_BOUNDARY - self.set_response_header(b"content-type", dtype.encode("ascii") + b"; boundary=" + boundary) + self.set_response_header( + b"content-type", dtype.encode("ascii") + b"; boundary=" + boundary + ) if self.method != self.HEAD: self.resp_handler = self._multipart_wrapper_factory(callback, boundary)