radare2-r2pipe/python/r2pipe/open_async.py

# -*- coding: utf-8 -*-
"""open_async.py
This script use code from r2pipe-async/open_p3.py script.

"""
import asyncio
import json
import os
import re
import sys


from collections.abc import Iterable
from concurrent.futures import ThreadPoolExecutor
from contextlib import ContextDecorator, suppress
from types import SimpleNamespace
from urllib.parse import quote, urlparse

from r2pipe.open_base import OpenBase, get_radare_path, jo2po, normalize_open_target


class open(OpenBase, ContextDecorator):
    # --------------------------------------------------------------------------
    # Context manager functions
    # --------------------------------------------------------------------------
    def __enter__(self):
        if self._running_loop() is not None:
            raise RuntimeError("Use 'async with' for r2pipe inside an event loop")
        return self

    def __exit__(self, *exc):
        if self._running_loop() is not None:
            raise RuntimeError("Use 'async with' for r2pipe inside an event loop")
        self.close()
        return False

    async def __aenter__(self):
        self._bind_loop()
        self._ensure_open()
        return self

    async def __aexit__(self, *exc):
        await self.aclose()
        return False

    async def _close_all_connections(self):
        """Close all connections in the connection pool"""
        async with self._connection_lock:
            for connections in self._connection_pool.values():
                for _, writer in connections:
                    writer.close()
                    try:
                        await writer.wait_closed()
                    except Exception:
                        pass
            self._connection_pool = {}

    async def _close_process(self):
        """Terminate the async subprocess and close its transport before the loop."""
        process = getattr(self, "process", None)
        if process is None:
            return

        stdin = getattr(process, "stdin", None)
        if stdin is not None:
            with suppress(Exception):
                if not stdin.is_closing():
                    stdin.close()

        if process.returncode is None:
            with suppress(ProcessLookupError):
                process.terminate()
            try:
                await asyncio.wait_for(process.wait(), timeout=1.0)
            except asyncio.TimeoutError:
                with suppress(ProcessLookupError):
                    process.kill()
                with suppress(Exception):
                    await process.wait()
        else:
            with suppress(Exception):
                await process.wait()

        transport = getattr(process, "_transport", None)
        if transport is not None:
            with suppress(Exception):
                transport.close()

        delattr(self, "process")
        self._pending_output = b""

    async def _shutdown(self):
        """Tear down all async resources while the loop is still alive."""
        try:
            if self._command_tasks:
                await asyncio.gather(*tuple(self._command_tasks), return_exceptions=True)
            async with self._cmd_lock:
                await self._close_all_connections()
                await self._close_process()
                await self._close_native()
        finally:
            self._closed = True

    def _free_native(self):
        core = self.__dict__.pop("native", None)
        if core is not None:
            core.free()

    async def _close_native(self):
        executor = self._native_executor
        if executor is None:
            return
        try:
            await self._loop.run_in_executor(executor, self._free_native)
        finally:
            # The native call has finished; join the now-idle dedicated worker.
            executor.shutdown(wait=True)
            self._native_executor = None

    def _ensure_open(self):
        if self._closed or self._closing or self._loop.is_closed():
            raise RuntimeError("r2pipe is closed")

    @staticmethod
    def _running_loop():
        try:
            if hasattr(asyncio, "get_running_loop"):
                return asyncio.get_running_loop()
            loop = asyncio.get_event_loop()
            return loop if loop.is_running() else None
        except RuntimeError:
            return None

    def _bind_loop(self):
        """Adopt the caller's loop before first use; transports cannot move loops."""
        running = self._running_loop()
        if running is not None and running is not self._loop:
            if self._owns_loop and not self._used and not self._closing:
                self._loop.close()
                self._loop = running
                self._owns_loop = False
                self._cmd_lock = self._new_lock()
                self._connection_lock = self._new_lock()
            else:
                raise RuntimeError("r2pipe belongs to a different event loop")
        elif running is None and not self._owns_loop and not self._closed:
            raise RuntimeError("Use r2pipe from its owning event loop")
        return running

    def _dispatch(self, coroutine, *args):
        running = self._bind_loop()
        self._ensure_open()
        if running is not None:
            return coroutine(*args)
        return self._loop.run_until_complete(coroutine(*args))

    async def aclose(self):
        """Finish in-flight commands and close resources without closing the loop."""
        if self._closed:
            return
        self._bind_loop()
        if self._shutdown_task is None:
            self._closing = True
            self._shutdown_task = self._loop.create_task(self._shutdown())
        canceled = False
        while True:
            try:
                await asyncio.shield(self._shutdown_task)
                break
            except asyncio.CancelledError:
                if self._shutdown_task.cancelled():
                    raise
                # Repeated cancellation must still let resource cleanup finish.
                canceled = True
        if canceled:
            raise asyncio.CancelledError

    def close(self):
        """Close synchronously, or return an awaitable inside an event loop."""
        if self._closed:
            return self.aclose() if self._running_loop() is not None else None
        running = self._bind_loop()
        if running is not None:
            return self.aclose()
        try:
            self._loop.run_until_complete(self.aclose())
        finally:
            if self._owns_loop:
                self._loop.close()

    def quit(self):
        """Use this client's event loop to finish the asynchronous shutdown."""
        return self.close()

    def _new_lock(self):
        if sys.version_info < (3, 10):
            return asyncio.Lock(loop=self._loop)
        return asyncio.Lock()

    def __init__(self, filename="", flags=None, radare2home=None):
        filename, file_args = normalize_open_target(filename)

        super().__init__(filename, flags)
        if flags is None:
            flags = []

        self.r2home = radare2home

        self._loop = self._running_loop()
        self._owns_loop = self._loop is None
        if self._owns_loop and os.name == "nt":
            self._loop = asyncio.ProactorEventLoop()
            asyncio.set_event_loop(self._loop)
        elif self._owns_loop:
            self._loop = asyncio.new_event_loop()
            if sys.version_info < (3, 12) and hasattr(asyncio, "get_child_watcher"):
                try:
                    watcher = asyncio.get_child_watcher()
                    watcher.attach_loop(self._loop)
                except RuntimeError:
                    pass

        # Add a lock for synchronizing command execution
        self._cmd_lock = self._new_lock()

        self._next_cmd_id = 0
        self._command_tasks = set()
        self._used = False
        self._closing = False
        self._shutdown_task = None
        self._native_executor = None

        # Add connection management for HTTP/TCP modes
        self._connection_pool = {}
        self._pool_size = 5  # Maximum number of concurrent connections
        self._connection_lock = self._new_lock()
        self._closed = False

        self.asyn = True

        if "_cmd" in self.__dict__:
            # OpenBase selected the embedded #!pipe or r2lang transport.
            self.asyn = False
            return

        if filename.startswith("http://"):
            self._cmd_coro = self._cmd_http
            self.uri = "/cmd"

            _tmp = urlparse(filename)
            self._host = _tmp.hostname
            self._port = _tmp.port

        elif filename.startswith("ccall://"):
            self._cmd_coro = self._cmd_native_async
            self.uri = filename[8:]

        elif filename.startswith("tcp://"):

            r = re.match(r"tcp://(\d+\.\d+.\d+.\d+):(\d+)/?", filename)
            if not r:
                raise ValueError("String doesn't match tcp format")

            self._cmd_coro = self._cmd_tcp
            self._host = r.group(1)
            self._port = r.group(2)

        elif filename:

            self._cmd_coro = self._cmd_process

            r_args = [f"-Rarg{i+1}={a}" for i, a in enumerate(file_args)]
            cmd = ["-q0", filename]
            cmd = cmd[:1] + flags + r_args + cmd[1:]
            self._process_start_cmd = cmd

        else:
            self.asyn = False

    async def _execute_cmd_with_lock(self, cmd, cmd_id, callback, started):
        """Execute command with proper synchronization"""
        async with self._cmd_lock:
            self._ensure_open()
            started[0] = True
            self._used = True
            result = await self._cmd_coro(cmd, cmd_id)
            if callback is not None:
                try:
                    callback(result)
                except Exception as error:
                    self._loop.call_exception_handler({
                        "message": "Exception in r2pipe command callback",
                        "exception": error,
                    })
            return result

    def _cmd(self, cmd, **kwargs):
        return self._dispatch(self._command, cmd, kwargs)

    async def _command(self, cmd, kwargs):
        self._bind_loop()
        self._ensure_open()
        callback = kwargs.get("callback")
        cmd_id = self._next_cmd_id
        self._next_cmd_id += 1
        started = [False]
        task = self._loop.create_task(
            self._execute_cmd_with_lock(cmd, cmd_id, callback, started)
        )
        self._command_tasks.add(task)
        task.add_done_callback(self._command_done)
        try:
            # A canceled caller must not interrupt reading a command's response.
            return await asyncio.shield(task)
        except asyncio.CancelledError:
            if not started[0]:
                task.cancel()
            raise

    def _command_done(self, task):
        self._command_tasks.discard(task)
        self._consume_task_error(task)

    @staticmethod
    def _consume_task_error(task):
        if not task.cancelled():
            # Retrieve errors even when a canceled caller no longer awaits the task.
            task.exception()

    async def _json_command(self, cmd, kwargs, object_hook=None, report_error=False):
        result = await self._command(cmd, kwargs)
        if report_error and result == "":
            result = "{}"
        try:
            return json.loads(result, object_hook=object_hook)
        except (ValueError, KeyError, TypeError) as error:
            if not report_error:
                raise
            sys.stderr.write(f"r2pipe.cmdj.Error: {error}\n")
            return None

    def cmdj(self, cmd, **kwargs):
        if not self.asyn:
            return super().cmdj(cmd, **kwargs)
        return self._dispatch(self._json_command, cmd, kwargs, None, True)

    async def _object_command(self, cmd, kwargs):
        result = await self._command(cmd, kwargs)
        try:
            return jo2po(result)
        except (ValueError, KeyError, TypeError) as error:
            sys.stderr.write(f"r2pipe.cmdj.Error: {error}\n")
            return None

    def cmdJ(self, cmd, **kwargs):
        if not self.asyn:
            return super().cmdJ(cmd, **kwargs)
        return self._dispatch(self._object_command, cmd, kwargs)

    def _r2pipe2(self, cmd, kwargs):
        request, transport_kwargs = self._r2pipe2_request(cmd, kwargs)
        cmd = "'" + json.dumps(request)
        return self._dispatch(
            self._json_command, cmd, transport_kwargs, lambda value: SimpleNamespace(**value)
        )

    def cmd2(self, cmd, **kwargs):
        if not self.asyn:
            return super().cmd2(cmd, **kwargs)
        return self._r2pipe2(cmd, kwargs)

    async def _cmd_native_async(self, cmd, _cmd_id):
        if self._native_executor is None:
            self._native_executor = ThreadPoolExecutor(max_workers=1)
        return await self._loop.run_in_executor(
            self._native_executor, super()._cmd_native, cmd
        )

    async def _cmd_process(self, cmd, _cmd_id):
        # Process initialization if needed
        if not hasattr(self, "process"):
            if self.r2home is not None:
                if not os.path.isdir(self.r2home):
                    raise Exception(
                        "`radare2home` passed to `open` is invalid, leave it None or put a valid path to r2 folder"
                    )
                r2path = os.path.join(self.r2home, "radare2")
                if os.name == "nt":
                    r2path += ".exe"
            else:
                r2path = get_radare_path()

            create = asyncio.create_subprocess_exec(
                r2path,
                *self._process_start_cmd,
                stdin=asyncio.subprocess.PIPE,
                stdout=asyncio.subprocess.PIPE,
            )

            self.process = await create  # Init the process

            # Track command execution state
            if not hasattr(self, "_pending_output"):
                self._pending_output = b""

            await self.process.stdout.read(1)  # Reads initial \x00

        # Clean command formatting
        cmd = cmd.strip().replace("\n", ";")

        # Write command to stdin
        await self.process.stdin.drain()  # Ensure any previous writes are complete
        self.process.stdin.write(bytes(cmd + "\n", "utf-8"))
        await self.process.stdin.drain()  # Ensure this write is complete

        # Read response with proper handling of pending data
        out = []

        # First check if we have pending output from previous commands
        if hasattr(self, "_pending_output") and self._pending_output:
            # Check if pending output contains a complete command response
            null_pos = self._pending_output.find(b"\x00")
            if null_pos != -1:
                # Extract the complete response and update pending
                response = self._pending_output[:null_pos]
                self._pending_output = self._pending_output[null_pos+1:]
                return response.decode("utf-8")

            # Add pending output to our current response
            out.append(self._pending_output)
            self._pending_output = b""

        # Read response byte by byte until null terminator
        buffer_size = 1024  # Read in larger chunks for efficiency
        while True:
            data = await self.process.stdout.read(buffer_size)

            if not data:
                raise ConnectionError("r2pipe connection closed before response terminator")

            null_pos = data.find(b"\x00")
            if null_pos != -1:
                # Found the terminator
                out.append(data[:null_pos])

                # Save any data after the null for future commands
                if null_pos < len(data) - 1:
                    self._pending_output = data[null_pos+1:]

                break
            else:
                out.append(data)

        # Join the output and convert to string
        return b"".join(out).decode("utf-8")

    async def _get_http_connection(self):
        """Get a connection from the pool or create a new one"""
        async with self._connection_lock:
            # Create a connection pool key
            pool_key = f"{self._host}:{self._port}"

            # Get or create the connection queue
            if pool_key not in self._connection_pool:
                self._connection_pool[pool_key] = []

            # Check if we have an available connection
            if self._connection_pool[pool_key]:
                # Reuse existing connection
                while self._connection_pool[pool_key]:
                    reader, writer = self._connection_pool[pool_key].pop(0)
                    if not reader.at_eof() and not writer.is_closing():
                        return reader, writer
                    writer.close()
                    with suppress(Exception):
                        await writer.wait_closed()

            # Create a new connection
            return await asyncio.open_connection(
                self._host, self._port
            )

    async def _return_http_connection(self, reader, writer):
        """Return a connection to the pool"""
        if writer.is_closing():
            return

        async with self._connection_lock:
            pool_key = f"{self._host}:{self._port}"

            # Only keep connections if within pool size limit
            if pool_key in self._connection_pool and len(self._connection_pool[pool_key]) < self._pool_size:
                self._connection_pool[pool_key].append((reader, writer))
            else:
                # Close the connection if pool is full
                writer.close()
                await writer.wait_closed()

    async def _cmd_http(self, cmd, cmd_id):
        """HTTP command execution with connection pooling and request ID tracking"""
        writer = None
        pooled = False
        try:
            quocmd = quote(cmd)

            # Get a connection from the pool
            reader, writer = await self._get_http_connection()

            # Create HTTP request with request ID tracking
            message = "\r\n".join(
                [
                    "GET /cmd/%s HTTP/1.1" % quocmd,
                    "Host: %s:%s" % (self._host, self._port),
                    "User-Agent: r2pipe/Python Client",
                    "Accept: */*",
                    "X-R2Pipe-Request-ID: %s" % cmd_id,
                    "Connection: keep-alive",  # Keep connection alive for pooling
                    "",
                    "",
                ]
            ).encode()

            # Send the request
            writer.write(message)
            await writer.drain()  # Ensure the request is fully sent

            # Read response with timeout to prevent hanging
            try:
                # Use a timeout to prevent hanging on incomplete responses
                body = []

                # Read response with proper chunking
                buffer_size = 4096  # Larger buffer size for efficiency

                # Header boundaries can span multiple transport reads.
                data = await asyncio.wait_for(
                    reader.readuntil(b'\r\n\r\n'),
                    timeout=10.0
                )
                header_lines = data[:-4].decode('utf-8', errors='ignore').split('\r\n')
                content_length = None
                connection = ""
                for line in header_lines[1:]:
                    if line.lower().startswith('content-length:'):
                        content_length = int(line.split(':', 1)[1].strip())
                        if content_length < 0:
                            raise ValueError("Invalid HTTP Content-Length")
                    elif line.lower().startswith('connection:'):
                        connection = line.split(':', 1)[1].strip().lower()
                keep_alive = content_length is not None and "close" not in connection and (
                    header_lines[0].startswith("HTTP/1.1 ") or "keep-alive" in connection
                )

                if content_length is not None:
                    remaining = content_length
                    while remaining > 0:
                        chunk = await asyncio.wait_for(
                            reader.read(min(buffer_size, remaining)),
                            timeout=5.0
                        )
                        if not chunk:
                            raise ConnectionError("Incomplete HTTP response body")
                        body.append(chunk)
                        remaining -= len(chunk)
                else:
                    # Responses without a length are delimited by connection close.
                    while True:
                        chunk = await asyncio.wait_for(
                            reader.read(buffer_size),
                            timeout=5.0
                        )
                        if not chunk:
                            break
                        body.append(chunk)

                # Reuse only responses whose framing and headers allow persistence.
                if keep_alive and not reader.at_eof():
                    await self._return_http_connection(reader, writer)
                    pooled = True

            except asyncio.TimeoutError as error:
                # If timeout occurs, close the connection
                writer.close()
                raise TimeoutError("Timeout while reading HTTP response") from error

            body = b"".join(body)
            return body.decode("utf-8", errors="ignore")

        finally:
            if writer is not None and not pooled:
                writer.close()
                with suppress(Exception):
                    await writer.wait_closed()

    async def _get_tcp_connection(self):
        """Get a TCP connection from the pool or create a new one"""
        async with self._connection_lock:
            # Create a connection pool key
            pool_key = f"tcp:{self._host}:{self._port}"

            # Get or create the connection queue
            if pool_key not in self._connection_pool:
                self._connection_pool[pool_key] = []

            # Check if we have an available connection
            if self._connection_pool[pool_key]:
                # Reuse existing connection
                reader, writer = self._connection_pool[pool_key].pop(0)
                return reader, writer

            # Create a new connection
            return await asyncio.open_connection(
                self._host, self._port
            )

    async def _return_tcp_connection(self, reader, writer):
        """Return a TCP connection to the pool"""
        if writer.is_closing():
            return

        async with self._connection_lock:
            pool_key = f"tcp:{self._host}:{self._port}"

            # Only keep connections if within pool size limit
            if pool_key in self._connection_pool and len(self._connection_pool[pool_key]) < self._pool_size:
                self._connection_pool[pool_key].append((reader, writer))
            else:
                # Close the connection if pool is full
                writer.close()
                await writer.wait_closed()

    async def _cmd_tcp(self, cmd, _cmd_id):
        """TCP command execution with connection pooling"""
        writer = None
        pooled = False
        try:
            # Get a connection from the pool
            reader, writer = await self._get_tcp_connection()

            writer.write(cmd.encode("utf-8"))

            # Make sure the command is sent
            await writer.drain()

            try:
                # Read with proper timeout handling
                res = []
                buffer_size = 4096  # Larger buffer for efficiency

                # Read data with timeout
                first_chunk = await asyncio.wait_for(
                    reader.read(buffer_size),
                    timeout=10.0
                )

                if not first_chunk:
                    # Connection closed, create new one next time
                    writer.close()
                    raise ConnectionError("TCP connection closed by server")

                res.append(first_chunk)

                # Continue reading until we get all data
                # TCP mode indicates end of data by closing the connection
                # or by having no more data available after a timeout
                try:
                    while not reader.at_eof():
                        chunk = await asyncio.wait_for(
                            reader.read(buffer_size),
                            timeout=1.0  # Short timeout to detect end of response
                        )
                        if not chunk:
                            break
                        res.append(chunk)
                except asyncio.TimeoutError:
                    # A short timeout here means we've likely read everything
                    pass

                # TCP connections are often single-use in r2pipe
                # But we'll attempt to keep them alive for reuse
                if not reader.at_eof():
                    await self._return_tcp_connection(reader, writer)
                    pooled = True
                else:
                    # Connection was closed by server, close our end too
                    writer.close()

            except asyncio.TimeoutError as error:
                # If we timeout on the initial read, something is wrong
                writer.close()
                raise TimeoutError("Timeout while reading TCP response") from error

            # Process the response
            response_data = b"".join(res)

            return response_data.decode("utf-8", errors="ignore")

        finally:
            if writer is not None and not pooled:
                writer.close()
                with suppress(Exception):
                    await writer.wait_closed()

    def wait(self, task):
        """Wait for awaitables and return completed futures in input order."""
        tasks = list(task) if isinstance(task, Iterable) else [task]
        return self._dispatch(self._wait, tasks)

    async def _wait(self, tasks):
        self._bind_loop()
        self._ensure_open()
        futures = []
        for task in tasks:
            future = asyncio.ensure_future(task, loop=self._loop)
            # A canceled wait() caller no longer receives these wrapper tasks.
            future.add_done_callback(self._consume_task_error)
            futures.append(future)
        if futures:
            await asyncio.wait(futures)
        return futures