diff --git a/restler/engine/transport_layer/messaging.py b/restler/engine/transport_layer/messaging.py index 9fd4d532..106d6620 100644 --- a/restler/engine/transport_layer/messaging.py +++ b/restler/engine/transport_layer/messaging.py @@ -8,8 +8,7 @@ import time import threading from importlib import util -from utils.logging.trace_db import (DB as TraceDatabase, - SequenceTracker) +import utils.logging.trace_db as trace_db from utils.formatting import iso_timestamp from utils.logger import raw_network_logging as RAW_LOGGING from engine.errors import TransportLayerException @@ -124,7 +123,7 @@ def sendRecv(self, message, req_timeout_sec, reconnect=False): response = self._sock.recv() RAW_LOGGING(f'Received: {response.to_str!r}\n') if Settings().use_trace_database: - TraceDatabase().log_request_response( + trace_db.DB().log_request_response( response=response.to_str, timestamp=iso_timestamp() ) @@ -204,7 +203,7 @@ def _append_to_header(message, content): # Send the RESTler user agent only if a custom user agent is not specified message = _append_to_header(message, f"User-Agent: restler/{Settings().version}") if self.connection_settings.include_unique_sequence_id: - sequence_id = SequenceTracker().get_sequence_id() + sequence_id = trace_db.SequenceTracker().get_sequence_id() if sequence_id is not None: message = _append_to_header(message, f"x-restler-sequence-id: {sequence_id}") @@ -277,8 +276,24 @@ def decode_buf(buf): RAW_LOGGING(f'Failed to decode data due to {ex}. \ Trying again while ignoring offending bytes.') return buf.decode(UTF8, "ignore") + elif DELIM in data: + RAW_LOGGING(f'Failed to decode body data due to {ex}. \ + Trying again while replacing offending bytes.') + return buf.decode(UTF8, "replace") + elif b"\r\n\r\n" in buf: + delim_idx = buf.index(b"\r\n\r\n") + 4 + header_bytes = buf[:delim_idx] + body_bytes = buf[delim_idx:] + try: + decoded_header = header_bytes.decode(UTF8) + except Exception: + decoded_header = header_bytes.decode(UTF8, "ignore") + RAW_LOGGING(f'Failed to decode body data following headers due to {ex}. \ + Trying again while replacing offending bytes.') + decoded_body = body_bytes.decode(UTF8, "replace") + return decoded_header + decoded_body else: - raise + raise TransportLayerException(f"Failed to decode data: {ex}") # get the data of header (and maybe more) bytes_received = 0 diff --git a/restler/unit_tests/test_messaging.py b/restler/unit_tests/test_messaging.py new file mode 100644 index 00000000..08c7bc6e --- /dev/null +++ b/restler/unit_tests/test_messaging.py @@ -0,0 +1,92 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +import unittest +import unittest.mock +import engine.dependencies +from restler_settings import ConnectionSettings, Settings, RestlerSettings +from engine.transport_layer.messaging import HttpSock, UTF8 +from engine.errors import TransportLayerException +from engine.transport_layer.response import HttpResponse + +class MockSocket(object): + def __init__(self, buffers): + self.buffers = list(buffers) + + def settimeout(self, timeout): + pass + + def recv(self, bufsize=2**20): + if self.buffers: + return self.buffers.pop(0) + return b"" + + def close(self): + pass + + +class MessagingDecodeTest(unittest.TestCase): + def setUp(self): + RestlerSettings.TEST_DeleteInstance() + self.raw_logging_patcher = unittest.mock.patch("engine.transport_layer.messaging.RAW_LOGGING") + self.raw_logging_patcher.start() + self.settings = RestlerSettings({}) + self.conn_settings = ConnectionSettings("localhost", 8888) + + def tearDown(self): + self.raw_logging_patcher.stop() + RestlerSettings.TEST_DeleteInstance() + + def test_recv_response_binary_body_in_single_buffer(self): + """ Tests that an HTTP 200 response with binary body in the first recv buffer + (e.g., zip or octet-stream file download) decodes successfully without raising + UnicodeDecodeError (Issue #976). + """ + http_sock = HttpSock(self.conn_settings) + # Headers followed by binary byte sequence that is invalid UTF-8 + raw_response = b"HTTP/1.1 200 OK\r\nContent-Type: application/octet-stream\r\nContent-Length: 4\r\n\r\n\xfa\xfb\xfc\xfd" + http_sock._sock = MockSocket([raw_response]) + + received = http_sock._recvResponse(10, "GET") + response = HttpResponse(received) + self.assertEqual("200", response.status_code) + self.assertIsNotNone(response.body) + + def test_recv_response_binary_body_in_subsequent_buffer(self): + """ Tests that an HTTP 200 response where the binary body is received + in a second recv call decodes successfully without crashing. + """ + http_sock = HttpSock(self.conn_settings) + header_buf = b"HTTP/1.1 200 OK\r\nContent-Type: application/zip\r\nContent-Length: 4\r\n\r\n" + body_buf = b"\xfa\xfb\xfc\xfd" + http_sock._sock = MockSocket([header_buf, body_buf]) + + received = http_sock._recvResponse(10, "GET") + response = HttpResponse(received) + self.assertEqual("200", response.status_code) + self.assertIsNotNone(response.body) + + def test_recv_response_invalid_header_raises_transport_layer_exception(self): + """ Tests that invalid UTF-8 in headers when ignore_decoding_failures is False + raises TransportLayerException (handled by sendRecv) instead of unhandled UnicodeDecodeError. + """ + http_sock = HttpSock(self.conn_settings) + raw_header = b"HTTP/1.1 200 \xfa\xfb\r\n" + http_sock._sock = MockSocket([raw_header]) + + with self.assertRaises(TransportLayerException): + http_sock._recvResponse(10, "GET") + + def test_recv_response_ignore_decoding_failures_setting(self): + """ Tests that ignore_decoding_failures=True continues to ignore decoding errors. """ + RestlerSettings.TEST_DeleteInstance() + settings = RestlerSettings({"ignore_decoding_failures": True}) + conn_settings = ConnectionSettings("localhost", 8888) + http_sock = HttpSock(conn_settings) + + raw_header = b"HTTP/1.1 200 \xfa\xfb\r\nContent-Length: 0\r\n\r\n" + http_sock._sock = MockSocket([raw_header]) + + received = http_sock._recvResponse(10, "GET") + response = HttpResponse(received) + self.assertEqual("200", response.status_code)