openrewrite-remote 0.2.0__py3-none-any.whl

rewrite/remote/receiver.py

@@ -0,0 +1,386 @@
+from collections import OrderedDict
+from enum import Enum
+from pathlib import Path
+from typing import Protocol, TypeVar, Optional, Type, Dict, Callable, List, cast, Iterable, Any, Tuple, \
+    get_args, TYPE_CHECKING
+from uuid import UUID
+
+import cbor2
+from cbor2 import CBORDecoder
+from rewrite import Markers, Marker, ParseErrorVisitor
+
+from rewrite import Tree, TreeVisitor, Cursor, FileAttributes
+from rewrite.remote.event import DiffEvent, EventType
+from . import remote_utils, type_utils
+
+if TYPE_CHECKING:
+    from .remoting import RemotingContext
+
+A = TypeVar('T')
+T = TypeVar('T', bound=Tree)
+V = TypeVar('V')
+I = TypeVar('I')
+P = TypeVar('P')
+
+
+class Receiver(Protocol):
+    def fork(self, context: 'ReceiverContext') -> 'ReceiverContext':
+        ...
+
+    def receive(self, before: Optional[T], ctx: 'ReceiverContext') -> object:
+        ...
+
+
+class OmniReceiver(Receiver):
+    def fork(self, ctx: 'ReceiverContext') -> 'ReceiverContext':
+        raise NotImplementedError("Cannot fork OmniReceiver")
+
+    def receive(self, before: Optional['Tree'], ctx: 'ReceiverContext') -> 'Tree':
+        visitor = self.Visitor()
+        return visitor.visit(before, ctx)
+
+    class Visitor(TreeVisitor[Tree, 'ReceiverContext']):
+        def visit(self, tree: Optional[Tree], ctx: 'ReceiverContext', parent: Optional[Cursor] = None) -> Optional[
+            Tree]:
+            self.cursor = Cursor(self.cursor, tree)
+            tree = ctx.polymorphic_receive_tree(tree)
+            self.cursor = self.cursor.parent
+            return tree
+
+
+class TreeReceiver(Protocol):
+    def receive_node(self) -> DiffEvent:
+        ...
+
+    def receive_value(self, expected_type: Type) -> DiffEvent:
+        ...
+
+
+class ReceiverFactory(Protocol):
+    def create(self, type_name: str, ctx: 'ReceiverContext') -> Tree:
+        ...
+
+
+class DetailsReceiver(Protocol[T]):
+    def receive_details(self, before: Optional[T], type: Optional[Type[T]], ctx: 'ReceiverContext') -> T:
+        pass
+
+
+class ReceiverContext:
+    Registry: Dict[Type, Callable[[], Receiver]] = OrderedDict()
+
+    def __init__(self, receiver: TreeReceiver, visitor: Optional[TreeVisitor] = None,
+                 factory: ReceiverFactory = None):
+        self.receiver = receiver
+        self.visitor = visitor
+        self.factory = factory
+
+    def fork(self, visitor: TreeVisitor, factory: ReceiverFactory) -> 'ReceiverContext':
+        return ReceiverContext(self.receiver, visitor, factory)
+
+    def receive_any_tree(self, before) -> T:
+        return OmniReceiver().receive(before, self)
+
+    def receive_tree(self, before: Optional[Tree], tree_type: Optional[str], ctx: 'ReceiverContext') -> Tree:
+        if before:
+            return before.accept(self.visitor, ctx)
+        else:
+            return self.factory.create(tree_type, ctx)
+
+    def polymorphic_receive_tree(self, before: Optional[Tree]) -> Optional[Tree]:
+        diff_event = self.receiver.receive_node()
+        if diff_event.event_type in (EventType.Add, EventType.Update):
+            tree_receiver = self.new_receiver(diff_event.concrete_type or type(before).__name__)
+            forked = tree_receiver.fork(self)
+            return forked.receive_tree(
+                None if diff_event.event_type == EventType.Add else before,
+                diff_event.concrete_type, forked
+            )
+        elif diff_event.event_type == EventType.Delete:
+            return None
+        else:
+            return before
+
+    def new_receiver(self, type_name: str) -> Receiver:
+        type_ = type_utils.get_type(type_name)
+        for entry_type, factory in self.Registry.items():
+            if issubclass(type_, entry_type):
+                return factory()
+        raise ValueError(f"Unsupported receiver type: {type_name}")
+
+    def receive_node(self, before: Optional[A],
+                     details: Callable[[Optional[A], Optional[str], 'ReceiverContext'], A]) -> Optional[A]:
+        evt = self.receiver.receive_node()
+        if evt.event_type == EventType.Delete:
+            return None
+        elif evt.event_type == EventType.Add:
+            return details(None, evt.concrete_type, self)
+        elif evt.event_type == EventType.Update:
+            return details(before, evt.concrete_type, self)
+        return before
+
+    def receive_markers(self, before: Optional[Markers], type: Optional[str], ctx: 'ReceiverContext') -> Markers:
+        id_ = self.receive_value(getattr(before, 'id', None), UUID)
+        after_markers = self.receive_values(getattr(before, 'markers', None), Marker)
+        if before:
+            return before.with_id(id_).with_markers(after_markers)
+        else:
+            return Markers(id_, after_markers)
+
+    def receive_nodes(self, before: Optional[List[A]], details: Callable[[Optional[A], Optional[str], 'ReceiverContext'], A]) -> Optional[List[A]]:
+        return remote_utils.receive_nodes(before, details, self)
+
+    def receive_values(self, before: Optional[List[V]], type: Type) -> Optional[List[V]]:
+        return remote_utils.receive_values(before, type, self)
+
+    def receive_value(self, before: Optional[V], type: Type) -> Optional[V]:
+        return self.receive_value0(before, type)
+
+    def receive_value0(self, before: Optional[V], type: Type) -> Optional[V]:
+        evt = self.receiver.receive_value(type)
+        if evt.event_type in (EventType.Update, EventType.Add):
+            return evt.msg
+        elif evt.event_type == EventType.Delete:
+            return None
+        return before
+
+    @staticmethod
+    def register(type_: Type, receiver_factory: Callable[[], Receiver]):
+        ReceiverContext.Registry[type_] = receiver_factory
+
+
+class ValueDeserializer(Protocol):
+    def deserialize(self, type_: Optional[Type], reader: CBORDecoder, context: 'DeserializationContext') -> Optional[Any]:
+        ...
+
+
+class DefaultValueDeserializer(ValueDeserializer):
+    def deserialize(self, expected_type: Optional[Type], reader: CBORDecoder, context: 'DeserializationContext') -> Any:
+        cbor_map = reader.decode_map()
+        error_message = "No deserializer found for: " + ", ".join(f"{k}: {v}" for k, v in cbor_map.items())
+        raise NotImplementedError(error_message)
+
+
+class DeserializationContext:
+    DefaultDeserializer = DefaultValueDeserializer()
+
+    def __init__(self, remoting_context: 'RemotingContext', value_deserializers: Optional[List[Tuple[Type, ValueDeserializer]]] = None):
+        self.remoting_context = remoting_context
+        self.value_deserializers = value_deserializers or []
+
+    def deserialize(self, expected_type: Type, decoder: CBORDecoder) -> Any:
+        value = decoder.decode()
+        if value is None:
+            return None
+
+        if expected_type == UUID:
+            return UUID(bytes=value)
+
+        if expected_type == str:
+            return value
+
+        if expected_type == bool:
+            return value
+
+        if expected_type == list:
+            return value
+
+        if expected_type == int:
+            return value
+
+        if expected_type == Path:
+            return Path(value)
+
+        if expected_type == float:
+            return value
+
+        if issubclass(expected_type, Enum):
+            return expected_type(value)
+
+        state = decoder.peek_state()
+
+        if state in {cbor2.CborReaderState.BOOLEAN, cbor2.CborReaderState.UNSIGNED_INT,
+                     cbor2.CborReaderState.NEGATIVE_INT}:
+            result = decoder.read_int()
+            obj = self.remoting_context.get_object(result)
+            return obj if obj is not None else result
+
+        if state in {cbor2.CborReaderState.HALF_FLOAT, cbor2.CborReaderState.FLOAT, cbor2.CborReaderState.DOUBLE}:
+            return decoder.read_double()
+
+        if state == cbor2.CborReaderState.TEXT_STRING:
+            str_value = decoder.read_text_string()
+            if decoder.peek_state() == cbor2.CborReaderState.END_ARRAY:
+                return str_value
+
+            concrete_type = str_value if decoder.peek_state() != cbor2.CborReaderState.END_ARRAY else expected_type.__name__
+
+            if concrete_type == "org.openrewrite.FileAttributes":
+                map = decoder.read_cbor_map()
+                return FileAttributes(None, None, None, False, False, False, 0)
+
+            if concrete_type == "java.lang.String":
+                return decoder.decode()
+            if concrete_type == "java.lang.Boolean":
+                return decoder.decode()
+            if concrete_type == "java.lang.Integer":
+                return decoder.decode()
+            if concrete_type == "java.lang.Character":
+                return decoder.decode()[0]
+            if concrete_type == "java.lang.Long":
+                return decoder.decode()
+            if concrete_type == "java.lang.Double":
+                return decoder.decode()
+            if concrete_type == "java.lang.Float":
+                return decoder.decode()
+            if concrete_type == "java.math.BigInteger":
+                return decoder.decode()
+            if concrete_type == "java.math.BigDecimal":
+                return decoder.decode()
+
+            raise NotImplementedError(f"No deserialization implemented for: {concrete_type}")
+
+        if state == cbor2.CborReaderState.ARRAY:
+            decoder.read_array_start()
+            concrete_type = decoder.read_text_string()
+            actual_type = type_utils.get_type(concrete_type)
+            for type_, deserializer in self.value_deserializers:
+                if issubclass(actual_type, type_):
+                    return deserializer.deserialize(actual_type, decoder, self)
+
+        if state == cbor2.CborReaderState.MAP:
+            if issubclass(expected_type, Marker):
+                if decoder.peek_state() == cbor2.CborReaderState.UNSIGNED_INT:
+                    obj_id = decoder.read_int()
+                    return self.remoting_context.get_object(obj_id)
+
+                marker_map = ValueDeserializer.read_cbor_map(decoder)
+
+                if "@c" not in marker_map:
+                    raise NotImplementedError("Expected @c key")
+
+                concrete_type = marker_map["@c"]
+
+                if concrete_type in {"org.openrewrite.marker.SearchResult", "Rewrite.Core.Marker.SearchResult"}:
+                    desc = marker_map.get("description", None)
+                    marker = SearchResult(UUID(bytes=marker_map["id"]), desc)
+                else:
+                    marker = UnknownJavaMarker(UUID(bytes=marker_map["id"]), marker_map)
+
+                if "@ref" in marker_map:
+                    self.remoting_context.add(marker_map["@ref"], marker)
+
+                return marker
+
+            decoder.read_map_start()
+            if decoder.read_text_string() != "@c":
+                raise NotImplementedError("Expected @c key")
+            concrete_type = decoder.read_text_string()
+            actual_type = type_utils.get_type(concrete_type)
+            for type_, deserializer in self.value_deserializers:
+                if issubclass(actual_type, type_):
+                    return deserializer.deserialize(actual_type, decoder, self)
+        raise NotImplementedError(f"No deserialization implemented for: {expected_type}")
+
+
+class JsonReceiver(TreeReceiver):
+    DEBUG = False
+
+    def __init__(self, stream, context: DeserializationContext):
+        super().__init__()
+        self._stream = stream
+        self._decoder = cbor2.CBORDecoder(self._stream)
+        self._context = context
+        self._count = 0
+
+    def receive_node(self):
+        array = self._decoder.decode()
+        if isinstance(array, list):
+            event_type = EventType(cast(int, array[0]))
+            msg = None
+            concrete_type = None
+
+            if event_type in {EventType.Add, EventType.Update}:
+                if event_type == EventType.Add and isinstance(array[1], str):
+                    concrete_type = array[1]
+
+            elif event_type not in {EventType.Delete, EventType.NoChange, EventType.StartList, EventType.EndList}:
+                raise NotImplementedError(event_type)
+
+            if self.DEBUG:
+                print(f"[{self._count}] {DiffEvent(event_type, concrete_type, msg)}")
+            self._count += 1
+            return DiffEvent(event_type, concrete_type, msg)
+
+        else:
+            raise NotImplementedError(f"Unexpected state: {type(array)}")
+
+    def receive_value(self, expected_type: Type):
+        length = remote_utils.decode_array_start(self._decoder)
+        event_type = EventType(self._decoder.decode())
+        msg = None
+        concrete_type = None
+
+        if event_type in {EventType.Add, EventType.Update}:
+            if (bool(get_args(expected_type)) and
+                    issubclass(expected_type, (list, Iterable))):
+                # special case for list events
+                msg = self._decoder.decode()
+            else:
+                msg = self._context.deserialize(expected_type, self._decoder)
+
+        elif event_type not in {EventType.Delete, EventType.NoChange, EventType.StartList, EventType.EndList}:
+            raise NotImplementedError(event_type)
+
+        if length is None:
+            array_end = self._decoder.decode()
+            # assert array_end == break_marker
+
+        if self.DEBUG:
+            print(f"[{self._count}] {DiffEvent(event_type, concrete_type, msg)}")
+        self._count += 1
+        return DiffEvent(event_type, concrete_type, msg)
+
+
+class ParseErrorReceiver(Receiver):
+    def fork(self, ctx):
+        return ctx.fork(self.Visitor(), self.Factory())
+
+    def receive(self, before, ctx):
+        forked = self.fork(ctx)
+        return forked.visitor.visit(before, forked)
+
+    class Visitor(ParseErrorVisitor):
+        def visit(self, tree, ctx, parent: Optional[Cursor] = None):
+            self.cursor = Cursor(self.cursor, tree)
+            tree = ctx.receive_node(tree, ctx.receive_tree)
+            self.cursor = self.cursor.parent
+            return tree
+
+        def visit_parse_error(self, parse_error, ctx):
+            parse_error = parse_error.with_id(ctx.receive_value(parse_error.id))
+            parse_error = parse_error.with_markers(ctx.receive_node(parse_error.markers, ctx.receive_markers))
+            parse_error = parse_error.with_source_path(ctx.receive_value(parse_error.source_path))
+            parse_error = parse_error.with_file_attributes(ctx.receive_value(parse_error.file_attributes))
+            parse_error = parse_error.with_charset_name(ctx.receive_value(parse_error.charset_name))
+            parse_error = parse_error.with_charset_bom_marked(ctx.receive_value(parse_error.charset_bom_marked))
+            parse_error = parse_error.with_checksum(ctx.receive_value(parse_error.checksum))
+            parse_error = parse_error.with_text(ctx.receive_value(parse_error.text))
+            # parse_error = parse_error.with_erroneous(ctx.receive_tree(parse_error.erroneous))
+            return parse_error
+
+    class Factory(ReceiverFactory):
+        def create(self, type_, ctx):
+            if type_ in ["rewrite.parser.ParseError", "org.openrewrite.tree.ParseError"]:
+                return ParseError(
+                    ctx.receive_value(None),
+                    ctx.receive_node(None, ctx.receive_markers),
+                    ctx.receive_value(None),
+                    ctx.receive_value(None),
+                    ctx.receive_value(None),
+                    ctx.receive_value(None),
+                    ctx.receive_value(None),
+                    ctx.receive_value(None),
+                    None  # ctx.receive_tree(None)
+                )
+            raise NotImplementedError("No factory method for type: " + type_)

rewrite/remote/remote_utils.py

@@ -0,0 +1,280 @@
+import socket
+import struct
+from enum import Enum, auto
+from io import BytesIO
+from typing import Callable, List, Optional, TypeVar, Dict, TYPE_CHECKING, Iterable, cast, Type
+
+from cbor2 import CBORDecoder, CBORDecodeValueError
+
+from .event import EventType
+
+if TYPE_CHECKING:
+    from .receiver import ReceiverContext, DetailsReceiver
+
+T = TypeVar('T')
+I = TypeVar('I')
+
+
+class Operation(Enum):
+    Add = auto()
+    Delete = auto()
+    NoChange = auto()
+    Update = auto()
+    Move = auto()
+
+
+def receive_nodes(before: Optional[List[T]],
+                  details: Callable[[Optional[T], Optional[str], 'ReceiverContext'], T],
+                  ctx: 'ReceiverContext') -> Optional[List[T]]:
+    list_event = ctx.receiver.receive_value(list)
+    if list_event.event_type == EventType.NoChange:
+        return before
+    elif list_event.event_type == EventType.Delete:
+        return None
+    elif list_event.event_type == EventType.Add:
+        after_size = list_event.msg
+        after = [None] * after_size  # Preallocate list
+        for i in range(after_size):
+            diff_event = ctx.receiver.receive_node()
+            if diff_event.event_type == EventType.Add:
+                after[i] = details.receive_details(None, diff_event.concrete_type, ctx)
+            elif diff_event.event_type == EventType.NoChange:
+                after[i] = None  # Or some default value
+            else:
+                raise NotImplementedError(f"Unexpected operation: {diff_event.event_type}")
+        return after
+    elif list_event.event_type == EventType.Update:
+        return _receive_updated_nodes(before, list_event.msg, details, ctx)
+    else:
+        raise NotImplementedError(f"Unexpected operation: {list_event.event_type}")
+
+
+def _receive_updated_nodes(before: List[T],
+                           after_size: int,
+                           details: 'DetailsReceiver[T]',
+                           ctx: 'ReceiverContext') -> List[T]:
+    modified = False
+    after_list = before
+    evt = ctx.receiver.receive_node()
+    if evt.event_type != EventType.StartList:
+        raise ValueError(f"Expected start list event: {evt.event_type}")
+
+    before_idx = 0
+    while True:
+        evt = ctx.receiver.receive_node()
+        if evt.event_type in (EventType.NoChange, EventType.EndList):
+            break
+
+        if evt.event_type in (EventType.Delete, EventType.Update, EventType.Add):
+            if not modified:
+                after_list = _copy_range(before, before_idx)
+                modified = True
+
+        if evt.event_type == EventType.NoChange:
+            if modified:
+                after_list.append(before[before_idx])
+            before_idx += 1
+        elif evt.event_type == EventType.Delete:
+            before_idx += 1
+        elif evt.event_type == EventType.Update:
+            after_list.append(details.receive_details(before[before_idx], evt.concrete_type, ctx))
+            before_idx += 1
+        elif evt.event_type == EventType.Add:
+            after_list.append(details.receive_details(None, evt.concrete_type, ctx))
+
+        if evt.event_type == EventType.EndList:
+            break
+
+    return after_list[:after_size] if len(after_list) > after_size else after_list
+
+
+def receive_values(before: Optional[List[T]], type: Type, ctx: 'ReceiverContext') -> Optional[List[T]]:
+    list_event = ctx.receiver.receive_value(list)
+    if list_event.event_type == EventType.NoChange:
+        return before
+    elif list_event.event_type == EventType.Delete:
+        return None
+    elif list_event.event_type == EventType.Add:
+        after_size = list_event.msg
+        after = [None] * after_size  # Preallocate list
+        for i in range(after_size):
+            diff_event = ctx.receiver.receive_value(type)
+            if diff_event.event_type == EventType.Add:
+                after[i] = diff_event.msg
+            elif diff_event.event_type == EventType.NoChange:
+                after[i] = None  # Or some default value
+            else:
+                raise NotImplementedError(f"Unexpected operation: {diff_event.event_type}")
+        return after
+    elif list_event.event_type == EventType.Update:
+        return _receive_updated_values(before, list_event.msg, type, ctx)
+    else:
+        raise NotImplementedError(f"Unexpected operation: {list_event.event_type}")
+
+
+def _receive_updated_values(before: List[T], after_size: int, type: Type, ctx: 'ReceiverContext') -> List[T]:
+    modified = False
+    after_list = before
+    evt = ctx.receiver.receive_node()
+    if evt.event_type != EventType.StartList:
+        raise ValueError(f"Expected start list event: {evt.event_type}")
+
+    before_idx = 0
+    while True:
+        evt = ctx.receiver.receive_value(type)
+        if evt.event_type in (EventType.NoChange, EventType.EndList):
+            break
+
+        if evt.event_type in (EventType.Delete, EventType.Update, EventType.Add):
+            if not modified:
+                after_list = _copy_range(before, before_idx)
+                modified = True
+
+        if evt.event_type == EventType.NoChange:
+            if modified:
+                after_list.append(before[before_idx])
+            before_idx += 1
+        elif evt.event_type == EventType.Delete:
+            before_idx += 1
+        elif evt.event_type in (EventType.Update, EventType.Add):
+            after_list.append(evt.msg)
+            if evt.event_type == EventType.Update:
+                before_idx += 1
+
+        if evt.event_type == EventType.EndList:
+            break
+
+    return after_list[:after_size] if len(after_list) > after_size else after_list
+
+
+def _copy_range(before: Iterable[T], j: int) -> List[T]:
+    if isinstance(before, list):
+        return before[:j]
+    elif hasattr(before, 'getrange'):  # If the object has a 'getrange' method (e.g., an immutable list)
+        return before.getrange(0, j)
+    else:
+        return list(before)[:j]
+
+
+def calculate_list_diff(
+        before: List[T],
+        after: List[T],
+        id_function: Callable[[T], I],
+        consumer: Callable[[Operation, int, int, Optional[T], Optional[T]], None]
+) -> None:
+    before_idx, after_idx = 0, 0
+    before_size, after_size = len(before), len(after)
+    after_map = None
+
+    while before_idx < before_size or after_idx < after_size:
+        # Check if we've reached the end of either of the lists
+        if before_idx >= before_size:
+            consumer(Operation.Add, -1, after_idx, None, after[after_idx])
+            after_idx += 1
+            continue
+        elif after_idx >= after_size:
+            consumer(Operation.Delete, before_idx, -1, before[before_idx], None)
+            before_idx += 1
+            continue
+
+        if before[before_idx] == after[after_idx]:
+            consumer(
+                Operation.NoChange,
+                before_idx,
+                after_idx,
+                before[before_idx],
+                after[after_idx]
+            )
+            before_idx += 1
+            after_idx += 1
+        else:
+            before_id = id_function(before[before_idx])
+            after_id = id_function(after[after_idx])
+
+            if before_id == after_id:
+                consumer(
+                    Operation.Update,
+                    before_idx,
+                    after_idx,
+                    before[before_idx],
+                    after[after_idx]
+                )
+                before_idx += 1
+                after_idx += 1
+            else:
+                if after_map is None:
+                    after_map = create_index_map(after, after_idx, id_function)
+
+                # If elements at current indices are not equal, figure out the operation
+                if before_id not in after_map:
+                    consumer(Operation.Delete, before_idx, -1, before[before_idx], None)
+                    before_idx += 1
+                else:
+                    consumer(Operation.Add, -1, after_idx, None, after[after_idx])
+                    after_idx += 1
+
+
+def create_index_map(
+        lst: List[T],
+        from_index: int,
+        id_function: Callable[[T], I]
+) -> Dict[I, int]:
+    result = {}
+    for i in range(from_index, len(lst)):
+        result[id_function(lst[i])] = i
+    return result
+
+
+def _decode_length(decoder: CBORDecoder, subtype: int, allow_indefinite: bool = False):
+    if subtype < 24:
+        return subtype
+    elif subtype == 24:
+        return decoder.read(1)[0]
+    elif subtype == 25:
+        return cast(int, struct.unpack(">H", decoder.read(2))[0])
+    elif subtype == 26:
+        return cast(int, struct.unpack(">L", decoder.read(4))[0])
+    elif subtype == 27:
+        return cast(int, struct.unpack(">Q", decoder.read(8))[0])
+    elif subtype == 31 and allow_indefinite:
+        return None
+    else:
+        raise CBORDecodeValueError(f"unknown unsigned integer subtype 0x{subtype:x}")
+
+def decode_array_start(decoder: CBORDecoder):
+    initial_byte = decoder.read(1)[0]
+    major_type = initial_byte >> 5
+    assert major_type == 4
+    subtype = initial_byte & 31
+    return _decode_length(decoder, subtype, allow_indefinite=True)
+
+
+COPY_BUFFER_SIZE = 4096
+COPY_BUFFER = bytearray(COPY_BUFFER_SIZE)
+COMMAND_END = bytes([0x81, 0x17])
+
+def read_to_command_end(sock: socket.socket) -> BytesIO:
+    memory_stream = BytesIO()
+
+    try:
+        while True:
+            bytes_read = sock.recv_into(COPY_BUFFER, len(COPY_BUFFER))
+            if bytes_read == 0:
+                break
+            memory_stream.write(COPY_BUFFER[:bytes_read])
+            if (bytes_read > 1 and
+                    COPY_BUFFER[bytes_read - 2] == COMMAND_END[0] and
+                    COPY_BUFFER[bytes_read - 1] == COMMAND_END[1]):
+                break
+            elif bytes_read == 1:
+                original_position = memory_stream.tell()
+                memory_stream.seek(-2, 1)  # Move back by 2 bytes
+                if (memory_stream.read(1)[0] == COMMAND_END[0] and
+                        memory_stream.read(1)[0] == COMMAND_END[1]):
+                    memory_stream.seek(original_position)
+                    break
+    except socket.error as e:
+        print(f"Socket error: {e}")
+
+    memory_stream.seek(0)
+    return memory_stream