diff --git a/crates/engines/src/ghidra/client.rs b/crates/engines/src/ghidra/client.rs index e44a57d3..63547c8d 100644 --- a/crates/engines/src/ghidra/client.rs +++ b/crates/engines/src/ghidra/client.rs @@ -4,17 +4,25 @@ use super::protocol::{GhidraError, Request, Response}; use serde_json::Value; use std::io::{BufRead, BufReader, Write}; -use std::net::{SocketAddr, TcpStream}; +use std::net::{Shutdown, SocketAddr, TcpStream}; use std::time::{Duration, Instant}; /// How long one socket read blocks before we re-check the deadline and the cancel flag. const READ_TICK: Duration = Duration::from_millis(200); +/// Largest reply line we read. The bridge answers with lists that the target decides the size +/// of (every function in the program), so a reply is refused rather than read without end. +pub const MAX_REPLY_BYTES: usize = 128 * 1024 * 1024; + pub struct Client { writer: TcpStream, reader: BufReader, token: String, next_id: u64, + /// The reply line read so far. It's kept between read ticks, so a long reply that arrives + /// over more than one tick isn't cut. + pending: Vec, + max_reply: usize, } impl std::fmt::Debug for Client { @@ -35,7 +43,7 @@ impl Client { Ok(stream) => { stream.set_read_timeout(Some(READ_TICK)).map_err(|e| GhidraError::Connection(e.to_string()))?; let reader = BufReader::new(stream.try_clone().map_err(|e| GhidraError::Connection(e.to_string()))?); - return Ok(Client { writer: stream, reader, token: token.to_string(), next_id: 1 }); + return Ok(Client { writer: stream, reader, token: token.to_string(), next_id: 1, pending: Vec::new(), max_reply: MAX_REPLY_BYTES }); } Err(e) => { last = e.to_string(); @@ -62,15 +70,13 @@ impl Client { if Instant::now() >= deadline { return Err(GhidraError::Timeout(format!("no reply to {method} within the deadline"))); } - let mut buf = String::new(); - match self.reader.read_line(&mut buf) { - Ok(0) => return Err(GhidraError::Connection("the bridge closed the connection".into())), - Ok(_) => { - let trimmed = buf.trim(); - if trimmed.is_empty() { + match self.read_line_part() { + Ok(None) => continue, + Ok(Some(line)) => { + if line.iter().all(u8::is_ascii_whitespace) { continue; } - let resp: Response = serde_json::from_str(trimmed).map_err(|e| GhidraError::Protocol(format!("bad response: {e}")))?; + let resp: Response = serde_json::from_slice(&line).map_err(|e| GhidraError::Protocol(format!("bad response: {e}")))?; if resp.id != id { // Out-of-order frame; keep reading for ours. continue; @@ -81,16 +87,48 @@ impl Client { let err = resp.error.unwrap_or(super::protocol::WireError { code: "bridge_error".into(), message: "unknown error".into() }); return Err(GhidraError::Bridge { code: err.code, message: err.message }); } - Err(e) if e.kind() == std::io::ErrorKind::WouldBlock || e.kind() == std::io::ErrorKind::TimedOut => { - // Read tick elapsed; loop re-checks deadline and cancel. - continue; - } - Err(e) if e.kind() == std::io::ErrorKind::Interrupted => continue, - Err(e) => return Err(GhidraError::Connection(e.to_string())), + Err(e) => return Err(e), } } } + /// Read what's available of the next reply line. `Some` is a whole line; `None` means the + /// read tick passed first (the deadline and cancel flag are checked again before the next). + fn read_line_part(&mut self) -> Result>, GhidraError> { + let room = self.max_reply.saturating_sub(self.pending.len()); + if room == 0 { + // The rest of this reply is still on the socket, so later replies can't be told + // apart from it: close the connection rather than read on. + let _ = self.writer.shutdown(Shutdown::Both); + self.pending = Vec::new(); + return Err(GhidraError::Protocol(format!("a reply from the bridge was larger than {} MiB", self.max_reply / (1024 * 1024)))); + } + let buf = match self.reader.fill_buf() { + Ok(buf) => buf, + Err(e) if matches!(e.kind(), std::io::ErrorKind::WouldBlock | std::io::ErrorKind::TimedOut | std::io::ErrorKind::Interrupted) => { + return Ok(None); + } + Err(e) => return Err(GhidraError::Connection(e.to_string())), + }; + if buf.is_empty() { + return Err(GhidraError::Connection("the bridge closed the connection".into())); + } + let (take, complete) = match buf.iter().position(|&b| b == b'\n') { + Some(i) if i < room => (i + 1, true), + _ => (buf.len().min(room), false), + }; + self.pending.extend_from_slice(buf.get(..take).unwrap_or_default()); + self.reader.consume(take); + Ok(complete.then(|| std::mem::take(&mut self.pending))) + } + + /// Use a smaller reply limit (tests). + #[cfg(test)] + pub(crate) fn with_max_reply(mut self, bytes: usize) -> Client { + self.max_reply = bytes; + self + } + /// The handshake: returns the bridge's session info. pub fn ping(&mut self, cancel: &dyn Fn() -> bool, deadline: Instant) -> Result { self.request("ping", Value::Null, cancel, deadline) diff --git a/crates/engines/src/ghidra/tests.rs b/crates/engines/src/ghidra/tests.rs index 4c172739..6c6fa4d3 100644 --- a/crates/engines/src/ghidra/tests.rs +++ b/crates/engines/src/ghidra/tests.rs @@ -16,6 +16,10 @@ enum Mode { Serve, /// Accept the connection and then stay silent forever. Silent, + /// Answer the first request with a reply line far longer than a test client accepts. + Huge, + /// Answer the first request with a reply sent in two halves, a few read ticks apart. + Split, } struct Fake { @@ -50,6 +54,33 @@ impl Fake { } } Mode::Serve => serve(stream, expected_token, &eof2), + Mode::Huge | Mode::Split => { + let mut writer = stream.try_clone().expect("clone"); + let mut reader = BufReader::new(stream); + let mut line = String::new(); + if reader.read_line(&mut line).unwrap_or(0) == 0 { + return; + } + let id = serde_json::from_str::(line.trim()).map(|r| r.id).unwrap_or(1); + let names = "f".repeat(64 * 1024); + let reply = format!("{{\"id\":{id},\"ok\":true,\"result\":{{\"functions\":[{{\"name\":\"{names}\",\"address\":\"0x1\"}}]}}}}\n"); + if matches!(mode, Mode::Split) { + let (a, b) = reply.as_bytes().split_at(reply.len() / 2); + let _ = writer.write_all(a); + let _ = writer.flush(); + std::thread::sleep(Duration::from_millis(600)); + let _ = writer.write_all(b); + } else { + let _ = writer.write_all(reply.as_bytes()); + } + let _ = writer.flush(); + // Wait for the client to go away. + line.clear(); + while reader.read_line(&mut line).map(|n| n > 0).unwrap_or(false) { + line.clear(); + } + eof2.store(true, Ordering::SeqCst); + } } } }); @@ -214,6 +245,35 @@ fn cancel_stops_a_request_quickly() { fake.join(); } +#[test] +fn an_oversized_reply_is_refused_and_the_connection_closed() { + let fake = Fake::start(Mode::Huge, "tok"); + let deadline = Instant::now() + Duration::from_secs(5); + let mut client = Client::connect(fake.addr, "tok", deadline).expect("connect").with_max_reply(4 * 1024); + match client.request("native_functions_list", serde_json::json!({}), &no_cancel(), deadline).unwrap_err() { + GhidraError::Protocol(m) => assert!(m.contains("larger than"), "{m}"), + other => panic!("expected a protocol error, got {other:?}"), + } + // The rest of that reply is still on the wire, so the client must not read on. + match client.request("ping", serde_json::json!({}), &no_cancel(), Instant::now() + Duration::from_secs(2)).unwrap_err() { + GhidraError::Connection(_) => {} + other => panic!("expected a connection error on the closed socket, got {other:?}"), + } + drop(client); + fake.join(); +} + +#[test] +fn a_reply_arriving_over_several_read_ticks_is_read_whole() { + let fake = Fake::start(Mode::Split, "tok"); + let deadline = Instant::now() + Duration::from_secs(10); + let mut client = Client::connect(fake.addr, "tok", deadline).expect("connect"); + let list = client.request("native_functions_list", serde_json::json!({}), &no_cancel(), deadline).expect("the whole reply"); + assert_eq!(list["functions"][0]["name"].as_str().map(str::len), Some(64 * 1024)); + drop(client); + fake.join(); +} + #[test] fn dropping_the_client_closes_the_connection() { let fake = Fake::start(Mode::Silent, "tok");