Cap the size of a Ghidra bridge reply (#588)

The Ghidra client refuses a reply line over 128 MiB and closes the connection, so a hostile binary can't make one reply exhaust memory; a long reply split across read ticks is still read whole.

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<TcpStream>,
     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<u8>,
+    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<Option<Vec<u8>>, 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<Value, GhidraError> {
         self.request("ping", Value::Null, cancel, deadline)

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::<Request>(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");