ghidra_worker_ctl/
conn.rs

1//! Framed JSON-RPC 2.0 request/response over the worker's TCP stream (spec §5). One request in
2//! flight at a time (the worker services calls single-threaded on the Ghidra thread). Uses the
3//! existing ghidra-ipc frame codec (4-byte BE length prefix) and RpcRequest/RpcResponse types.
4
5use ghidra_ipc::frame::{encode_frame, FrameDecoder, FrameError};
6use ghidra_ipc::rpc::{RpcRequest, RpcResponse};
7use std::io;
8use tokio::io::{AsyncReadExt, AsyncWriteExt};
9use tokio::net::TcpStream;
10
11pub struct WorkerConn {
12    stream: TcpStream,
13    decoder: FrameDecoder,
14    scratch: Vec<u8>,
15    /// Set when the stream is known-corrupt (an id mismatch, a framing/EOF error, or an externally
16    /// signalled cancellation). A poisoned conn must NOT be reused; the holder kill+respawns (§4.4).
17    poisoned: bool,
18}
19
20impl WorkerConn {
21    pub fn new(stream: TcpStream) -> Self {
22        Self {
23            stream,
24            decoder: FrameDecoder::new(),
25            scratch: vec![0u8; 64 * 1024],
26            poisoned: false,
27        }
28    }
29
30    /// Send a request and read the matching response frame. On any stream/framing error, or a
31    /// response whose `id` does not equal the request's, the connection is POISONED (never reused).
32    pub async fn request(&mut self, req: &RpcRequest) -> io::Result<RpcResponse> {
33        if self.poisoned {
34            return Err(io::Error::new(
35                io::ErrorKind::BrokenPipe,
36                "connection poisoned; must respawn",
37            ));
38        }
39        let resp = self.request_inner(req).await;
40        if resp.is_err() {
41            self.poisoned = true;
42        }
43        resp
44    }
45
46    async fn request_inner(&mut self, req: &RpcRequest) -> io::Result<RpcResponse> {
47        let body = serde_json::to_vec(req).map_err(io::Error::other)?;
48        let frame = encode_frame(&body).map_err(frame_err)?;
49        self.stream.write_all(&frame).await?;
50        let body = self.read_frame().await?;
51        let resp: RpcResponse = serde_json::from_slice(&body).map_err(io::Error::other)?;
52        let resp_id = match &resp {
53            RpcResponse::Success { id, .. } => id,
54            RpcResponse::Failure { id, .. } => id,
55        };
56        if resp_id != &req.id {
57            return Err(io::Error::new(
58                io::ErrorKind::InvalidData,
59                format!(
60                    "response id {resp_id:?} does not match request id {:?}; stream desynced",
61                    req.id
62                ),
63            ));
64        }
65        Ok(resp)
66    }
67
68    /// Like `request` but does NOT early-return on the poison flag and does NOT set it on error — the
69    /// caller drives poisoning explicitly (poison-before-await, unpoison-on-success) so a DROPPED
70    /// future (cancellation) leaves the conn poisoned (spec §4.4 poison-on-drop). Still validates the
71    /// response id and still errors on a broken/desynced stream.
72    pub async fn request_raw(&mut self, req: &RpcRequest) -> io::Result<RpcResponse> {
73        self.request_inner(req).await
74    }
75
76    /// Clear the poison flag (used only after a `request_raw` completes successfully).
77    pub fn unpoison(&mut self) {
78        self.poisoned = false;
79    }
80
81    /// True once the connection has been poisoned and must be torn down.
82    pub fn is_poisoned(&self) -> bool {
83        self.poisoned
84    }
85
86    /// Mark the connection poisoned (e.g. the in-flight request future was dropped/cancelled — §4.4
87    /// poison-on-drop: the caller marks it before the next call so the orphaned response frame is
88    /// never mis-read).
89    pub fn poison(&mut self) {
90        self.poisoned = true;
91    }
92
93    /// Read exactly one complete frame body, pulling from the socket until the decoder yields one.
94    pub async fn read_frame(&mut self) -> io::Result<Vec<u8>> {
95        loop {
96            if let Some(b) = self.decoder.next_frame().map_err(frame_err)? {
97                return Ok(b);
98            }
99            let n = self.stream.read(&mut self.scratch).await?;
100            if n == 0 {
101                return Err(io::Error::new(
102                    io::ErrorKind::UnexpectedEof,
103                    "worker closed the connection",
104                ));
105            }
106            self.decoder.push(&self.scratch[..n]);
107        }
108    }
109
110    /// Send a request without waiting for a response (used for `shutdown`).
111    pub async fn send(&mut self, req: &RpcRequest) -> io::Result<()> {
112        let body = serde_json::to_vec(req).map_err(io::Error::other)?;
113        let frame = encode_frame(&body).map_err(frame_err)?;
114        self.stream.write_all(&frame).await
115    }
116}
117
118fn frame_err(e: FrameError) -> io::Error {
119    io::Error::new(io::ErrorKind::InvalidData, e)
120}
121
122#[cfg(test)]
123mod tests {
124    use super::*;
125    use ghidra_ipc::frame::{encode_frame, FrameDecoder};
126    use ghidra_ipc::rpc::{JsonRpcVersion, RpcId, RpcRequest};
127    use tokio::io::{AsyncReadExt, AsyncWriteExt};
128    use tokio::net::TcpListener;
129
130    #[tokio::test]
131    async fn request_response_roundtrips_over_tcp() {
132        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
133        let addr = listener.local_addr().unwrap();
134
135        // Mock worker: read one request frame, echo back a Success whose result mirrors the method.
136        let server = tokio::spawn(async move {
137            let (mut sock, _) = listener.accept().await.unwrap();
138            let mut dec = FrameDecoder::new();
139            let mut tmp = [0u8; 4096];
140            let body = loop {
141                let n = sock.read(&mut tmp).await.unwrap();
142                dec.push(&tmp[..n]);
143                if let Some(b) = dec.next_frame().unwrap() {
144                    break b;
145                }
146            };
147            let req: RpcRequest = serde_json::from_slice(&body).unwrap();
148            let resp = serde_json::json!({
149                "jsonrpc": "2.0", "id": 1, "result": { "echoed": req.method }
150            });
151            let frame = encode_frame(&serde_json::to_vec(&resp).unwrap()).unwrap();
152            sock.write_all(&frame).await.unwrap();
153        });
154
155        let stream = tokio::net::TcpStream::connect(addr).await.unwrap();
156        let mut conn = WorkerConn::new(stream);
157        let req = RpcRequest {
158            jsonrpc: JsonRpcVersion,
159            id: RpcId::Number(1),
160            method: "decompile_function".to_string(),
161            params: None,
162        };
163        let resp = conn.request(&req).await.unwrap();
164        match resp {
165            ghidra_ipc::rpc::RpcResponse::Success { result, .. } => {
166                assert_eq!(result["echoed"], "decompile_function");
167            }
168            other => panic!("expected Success, got {other:?}"),
169        }
170        server.await.unwrap();
171    }
172
173    #[tokio::test]
174    async fn mismatched_response_id_poisons_and_errors() {
175        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
176        let addr = listener.local_addr().unwrap();
177        // Mock worker replies with id:999 to a request sent with id:1.
178        let server = tokio::spawn(async move {
179            let (mut sock, _) = listener.accept().await.unwrap();
180            let mut dec = FrameDecoder::new();
181            let mut tmp = [0u8; 4096];
182            loop {
183                let n = sock.read(&mut tmp).await.unwrap();
184                dec.push(&tmp[..n]);
185                if dec.next_frame().unwrap().is_some() {
186                    break;
187                }
188            }
189            let resp = serde_json::json!({ "jsonrpc":"2.0","id":999,"result":{} });
190            let frame = encode_frame(&serde_json::to_vec(&resp).unwrap()).unwrap();
191            sock.write_all(&frame).await.unwrap();
192            tokio::time::sleep(std::time::Duration::from_millis(200)).await;
193        });
194        let stream = tokio::net::TcpStream::connect(addr).await.unwrap();
195        let mut conn = WorkerConn::new(stream);
196        let req = RpcRequest {
197            jsonrpc: JsonRpcVersion,
198            id: RpcId::Number(1),
199            method: "x".into(),
200            params: None,
201        };
202        let err = conn.request(&req).await.unwrap_err();
203        assert_eq!(err.kind(), std::io::ErrorKind::InvalidData);
204        assert!(
205            conn.is_poisoned(),
206            "a mismatched id must poison the connection"
207        );
208        server.await.unwrap();
209    }
210
211    #[tokio::test]
212    async fn poisoned_connection_refuses_further_requests() {
213        // A connection can be poisoned directly (e.g. after a cancelled call); the next request must
214        // fail fast without touching the socket.
215        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
216        let addr = listener.local_addr().unwrap();
217        let _server = tokio::spawn(async move {
218            let _ = listener.accept().await;
219        });
220        let stream = tokio::net::TcpStream::connect(addr).await.unwrap();
221        let mut conn = WorkerConn::new(stream);
222        conn.poison();
223        let req = RpcRequest {
224            jsonrpc: JsonRpcVersion,
225            id: RpcId::Number(1),
226            method: "x".into(),
227            params: None,
228        };
229        let err = conn.request(&req).await.unwrap_err();
230        assert!(conn.is_poisoned());
231        assert_eq!(err.kind(), std::io::ErrorKind::BrokenPipe);
232    }
233
234    #[tokio::test]
235    async fn request_raw_success_can_be_unpoisoned() {
236        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
237        let addr = listener.local_addr().unwrap();
238        let server = tokio::spawn(async move {
239            let (mut sock, _) = listener.accept().await.unwrap();
240            let mut dec = FrameDecoder::new();
241            let mut tmp = [0u8; 4096];
242            loop {
243                let n = sock.read(&mut tmp).await.unwrap();
244                dec.push(&tmp[..n]);
245                if dec.next_frame().unwrap().is_some() {
246                    break;
247                }
248            }
249            let resp = serde_json::json!({ "jsonrpc":"2.0","id":1,"result":{} });
250            sock.write_all(&encode_frame(&serde_json::to_vec(&resp).unwrap()).unwrap())
251                .await
252                .unwrap();
253            tokio::time::sleep(std::time::Duration::from_millis(150)).await;
254        });
255        let stream = tokio::net::TcpStream::connect(addr).await.unwrap();
256        let mut conn = WorkerConn::new(stream);
257        conn.poison();
258        let req = RpcRequest {
259            jsonrpc: JsonRpcVersion,
260            id: RpcId::Number(1),
261            method: "x".into(),
262            params: None,
263        };
264        assert!(conn.request_raw(&req).await.is_ok());
265        conn.unpoison();
266        assert!(!conn.is_poisoned());
267        server.await.unwrap();
268    }
269}
ghidra_worker_ctl/conn.rs

ghidra_worker_ctl/
conn.rs
