mirror of
https://github.com/radareorg/radare2-r2pipe.git
synced 2026-11-09 12:41:28 -10:00
673 lines
24 KiB
Python
673 lines
24 KiB
Python
# -*- coding: utf-8 -*-
|
|
"""open_async.py
|
|
This script use code from r2pipe-async/open_p3.py script.
|
|
|
|
"""
|
|
import asyncio
|
|
import os
|
|
import re
|
|
import sys
|
|
|
|
|
|
# whole file doesn't have any profit of asyncio usage, TODO: refactor
|
|
|
|
from collections.abc import Iterable
|
|
from contextlib import ContextDecorator, suppress
|
|
from urllib.parse import quote, urlparse
|
|
|
|
from r2pipe.open_base import OpenBase, get_radare_path, normalize_open_target
|
|
|
|
|
|
class open(OpenBase, ContextDecorator):
|
|
# --------------------------------------------------------------------------
|
|
# Contenxt manager functions
|
|
# --------------------------------------------------------------------------
|
|
def __enter__(self):
|
|
return self
|
|
|
|
def __exit__(self, *exc):
|
|
self.close()
|
|
return False
|
|
|
|
async def _close_all_connections(self):
|
|
"""Close all connections in the connection pool"""
|
|
async with self._connection_lock:
|
|
for pool_key, connections in self._connection_pool.items():
|
|
for reader, 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."""
|
|
await self._close_all_connections()
|
|
await self._close_process()
|
|
|
|
def _ensure_open(self):
|
|
if getattr(self, "_closed", False) or self._loop.is_closed():
|
|
raise RuntimeError("r2pipe is closed")
|
|
|
|
def close(self):
|
|
"""Close the r2pipe instance and clean up resources"""
|
|
if getattr(self, "_closed", False):
|
|
return
|
|
|
|
if self._loop.is_closed():
|
|
self._closed = True
|
|
return
|
|
|
|
if self._loop.is_running():
|
|
raise RuntimeError("Cannot close r2pipe while its event loop is running")
|
|
|
|
self._closed = True
|
|
try:
|
|
self._loop.run_until_complete(self._shutdown())
|
|
finally:
|
|
self._connection_pool = {}
|
|
self._loop.close()
|
|
|
|
def quit(self):
|
|
"""Use this client's event loop to finish the asynchronous shutdown."""
|
|
self.close()
|
|
|
|
def __init__(self, filename="", flags=None, radare2home=None):
|
|
file_args = []
|
|
filename, file_args = normalize_open_target(filename)
|
|
|
|
super(open, self).__init__(filename, flags)
|
|
if flags is None:
|
|
flags = []
|
|
|
|
self.r2home = radare2home
|
|
|
|
if os.name == "nt":
|
|
self._loop = asyncio.ProactorEventLoop()
|
|
asyncio.set_event_loop(self._loop)
|
|
else:
|
|
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 = asyncio.Lock()
|
|
|
|
# Add a task queue for managing command sequence
|
|
self._cmd_queue = []
|
|
self._next_cmd_id = 0
|
|
|
|
# Add connection management for HTTP/TCP modes
|
|
self._connection_pool = {}
|
|
self._pool_size = 5 # Maximum number of concurrent connections
|
|
self._connection_lock = asyncio.Lock()
|
|
self._closed = False
|
|
|
|
self.asyn = True
|
|
|
|
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
|
|
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
|
|
|
|
def _callback_wrapper(self, future):
|
|
result_data = future.result()
|
|
|
|
# Check if result contains command ID for ordered execution
|
|
if isinstance(result_data, tuple) and len(result_data) == 3:
|
|
cmd_id, result, callback = result_data
|
|
|
|
# Ensure commands are processed in sequence
|
|
if self._cmd_queue and self._cmd_queue[0] == cmd_id:
|
|
self._cmd_queue.pop(0) # Remove this command from queue
|
|
|
|
# Process callback if provided
|
|
if callback:
|
|
callback(result)
|
|
|
|
# Process any waiting callbacks that can now be executed
|
|
self._process_pending_callbacks()
|
|
else:
|
|
# Fallback for backward compatibility
|
|
result, callback = result_data
|
|
if callback:
|
|
callback(result)
|
|
|
|
def _process_pending_callbacks(self):
|
|
"""Process any callbacks that are ready based on command ID ordering"""
|
|
# This would handle pending callbacks when their turn comes
|
|
# Implementation would depend on how we store pending results
|
|
pass
|
|
|
|
async def _execute_cmd_with_lock(self, cmd, future, callback, cmd_id):
|
|
"""Execute command with proper synchronization"""
|
|
async with self._cmd_lock:
|
|
result = await self._cmd_coro(cmd, future, callback)
|
|
return result, cmd_id
|
|
|
|
def _cmd(self, cmd, **kwargs):
|
|
self._ensure_open()
|
|
|
|
# Get callback, if available
|
|
callback = kwargs.get("callback")
|
|
future = self._loop.create_future()
|
|
|
|
# Assign a sequential ID to this command for ordered execution
|
|
cmd_id = self._next_cmd_id
|
|
self._next_cmd_id += 1
|
|
self._cmd_queue.append(cmd_id)
|
|
|
|
# Store the command ID in the future for tracking
|
|
future._cmd_id = cmd_id
|
|
|
|
# Add the callback to process the result
|
|
future.add_done_callback(self._callback_wrapper)
|
|
|
|
coro = self._execute_cmd_with_lock(cmd, future, callback, cmd_id)
|
|
try:
|
|
result, _ = self._loop.run_until_complete(coro)
|
|
return result
|
|
except Exception:
|
|
coro.close()
|
|
raise
|
|
|
|
async def _cmd_process(self, cmd, future, callback):
|
|
# 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:]
|
|
result = response.decode("utf-8")
|
|
|
|
# Set result and return
|
|
if hasattr(future, "_cmd_id"):
|
|
future.set_result((future._cmd_id, result, callback))
|
|
else:
|
|
future.set_result((result, callback))
|
|
return result
|
|
|
|
# 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:
|
|
# Connection closed or error
|
|
if out:
|
|
break
|
|
return None
|
|
|
|
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
|
|
result = b"".join(out).decode("utf-8")
|
|
|
|
# Set result with command ID if available
|
|
if hasattr(future, "_cmd_id"):
|
|
future.set_result((future._cmd_id, result, callback))
|
|
else:
|
|
future.set_result((result, callback))
|
|
|
|
return result
|
|
|
|
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
|
|
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_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, future, callback):
|
|
"""HTTP command execution with connection pooling and request ID tracking"""
|
|
try:
|
|
quocmd = quote(cmd)
|
|
|
|
# Get a connection from the pool
|
|
reader, writer = await self._get_http_connection()
|
|
|
|
# Add request ID as a custom header for tracking
|
|
cmd_id = None
|
|
if hasattr(future, '_cmd_id'):
|
|
cmd_id = future._cmd_id
|
|
|
|
# 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 if cmd_id is not None else "none"),
|
|
"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
|
|
res = []
|
|
|
|
# Read response with proper chunking
|
|
buffer_size = 4096 # Larger buffer size for efficiency
|
|
|
|
# Read first chunk to get headers
|
|
data = await asyncio.wait_for(
|
|
reader.read(buffer_size),
|
|
timeout=10.0
|
|
)
|
|
|
|
if not data:
|
|
# Connection closed, need to create a new one next time
|
|
writer.close()
|
|
raise ConnectionError("Connection closed by server")
|
|
|
|
res.append(data)
|
|
|
|
# Check for Content-Length header to determine how much to read
|
|
headers_end = data.find(b'\r\n\r\n')
|
|
if headers_end != -1:
|
|
headers = data[:headers_end].decode('utf-8', errors='ignore')
|
|
content_length = None
|
|
|
|
# Parse content length
|
|
for line in headers.split('\r\n'):
|
|
if line.lower().startswith('content-length:'):
|
|
try:
|
|
content_length = int(line.split(':', 1)[1].strip())
|
|
break
|
|
except ValueError:
|
|
pass
|
|
|
|
# Calculate how much more data we need to read
|
|
if content_length is not None:
|
|
bytes_read = len(data) - headers_end - 4 # subtract header size
|
|
remaining = content_length - bytes_read
|
|
|
|
# Read remaining content if needed
|
|
while remaining > 0 and not reader.at_eof():
|
|
chunk = await asyncio.wait_for(
|
|
reader.read(min(buffer_size, remaining)),
|
|
timeout=5.0
|
|
)
|
|
if not chunk:
|
|
break
|
|
res.append(chunk)
|
|
remaining -= len(chunk)
|
|
else:
|
|
# If no content length, read until connection closes
|
|
while not reader.at_eof():
|
|
chunk = await asyncio.wait_for(
|
|
reader.read(buffer_size),
|
|
timeout=5.0
|
|
)
|
|
if not chunk:
|
|
break
|
|
res.append(chunk)
|
|
|
|
# Return connection to the pool for reuse
|
|
await self._return_http_connection(reader, writer)
|
|
|
|
except asyncio.TimeoutError:
|
|
# If timeout occurs, close the connection
|
|
writer.close()
|
|
raise TimeoutError("Timeout while reading HTTP response")
|
|
|
|
# Process the response
|
|
res = b"".join(res)
|
|
|
|
# Remove HTTP headers
|
|
start = 0
|
|
header_end = res.find(b'\r\n\r\n')
|
|
if header_end != -1:
|
|
start = header_end + 4
|
|
else:
|
|
# Fallback to line-by-line parsing if we can't find the header end
|
|
for x in res.splitlines():
|
|
if not x:
|
|
start += 1
|
|
break
|
|
start += len(x) + 1 # +1 because we must be count '\n'
|
|
|
|
# Parse the result
|
|
result = res[start:].decode("utf-8", errors="ignore")
|
|
|
|
# Set result with command ID if available
|
|
if cmd_id is not None:
|
|
future.set_result((cmd_id, result, callback))
|
|
else:
|
|
future.set_result((result, callback))
|
|
return result
|
|
|
|
except Exception as e:
|
|
# Set error result with command ID if available
|
|
if hasattr(future, '_cmd_id'):
|
|
future.set_result((future._cmd_id, str(e), callback))
|
|
else:
|
|
future.set_result((str(e), callback))
|
|
|
|
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, future, callback):
|
|
"""TCP command execution with connection pooling and command tracking"""
|
|
try:
|
|
# Get a connection from the pool
|
|
reader, writer = await self._get_tcp_connection()
|
|
|
|
# Get command ID if available
|
|
cmd_id = None
|
|
if hasattr(future, '_cmd_id'):
|
|
cmd_id = future._cmd_id
|
|
|
|
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)
|
|
else:
|
|
# Connection was closed by server, close our end too
|
|
writer.close()
|
|
|
|
except asyncio.TimeoutError:
|
|
# If we timeout on the initial read, something is wrong
|
|
writer.close()
|
|
raise TimeoutError("Timeout while reading TCP response")
|
|
|
|
# Process the response
|
|
response_data = b"".join(res)
|
|
|
|
# Check if the response contains our tracking ID
|
|
result = response_data.decode("utf-8", errors="ignore")
|
|
|
|
# Set result with command ID if available
|
|
if cmd_id is not None:
|
|
future.set_result((cmd_id, result, callback))
|
|
else:
|
|
future.set_result((result, callback))
|
|
|
|
return result
|
|
|
|
except Exception as e:
|
|
# Set error result with command ID if available
|
|
if hasattr(future, '_cmd_id'):
|
|
future.set_result((future._cmd_id, str(e), callback))
|
|
else:
|
|
future.set_result((str(e), callback))
|
|
|
|
def wait(self, task):
|
|
"""Wait until task finish with proper ordering
|
|
|
|
This method ensures tasks are completed in the correct sequence
|
|
even when multiple commands are executed concurrently.
|
|
"""
|
|
_tasks = task
|
|
if not isinstance(task, Iterable):
|
|
_tasks = [task]
|
|
|
|
# Create a function to wait for all tasks in order
|
|
async def wait_ordered():
|
|
# Wait for all tasks to complete
|
|
done, _ = await asyncio.wait(_tasks)
|
|
|
|
# Sort results by command ID if available
|
|
results = []
|
|
for future in done:
|
|
try:
|
|
# Extract result and command ID if available
|
|
result = future.result()
|
|
if isinstance(result, tuple) and len(result) > 1:
|
|
# If this is a tuple with cmd_id
|
|
cmd_id, data = result[0], result[1]
|
|
results.append((cmd_id, data, future))
|
|
else:
|
|
# No command ID available
|
|
results.append((None, result, future))
|
|
except Exception as e:
|
|
# Handle any exceptions during task execution
|
|
pass
|
|
|
|
# Sort results by command ID if available
|
|
results.sort(key=lambda x: x[0] if x[0] is not None else float('inf'))
|
|
|
|
# Return sorted results
|
|
return [r[2] for r in results]
|
|
|
|
# Run the ordered wait function
|
|
if self._loop.is_running():
|
|
return asyncio.ensure_future(wait_ordered())
|
|
else:
|
|
return self._loop.run_until_complete(wait_ordered())
|