ghidra_worker_ctl/
boot.rs

1//! Boot orchestration (spec §4/§5). Host is the TCP server; the worker connects back and proves
2//! itself with a token+version init handshake. A racing local process without the token is rejected
3//! and the accept loop keeps going until the boot deadline (agy risk note A: a silent stall here is
4//! often an EDR/firewall intercept of the unsigned Java process, not a code bug — check that first).
5//!
6//! Project-lock stance (spec §4/§8.1, spike D2 — DELIBERATE, not an omission): boot does NOT parse,
7//! reclaim, or delete Ghidra's own project `.lock`. That file carries no PID/owner token, so it
8//! cannot be proven stale; on channel-locking filesystems the OS releases the lock on worker death,
9//! so a same-host respawn self-heals with zero cleanup. The host's own sidecar dual-instance lock
10//! and network-share reclaim are M1 scope (this thin slice boots `-readOnly`, i.e. no transactions).
11
12use crate::config::WorkerConfig;
13use crate::conn::WorkerConn;
14use crate::launch::WorkerProcess;
15use crate::state::WorkerLifecycle;
16use crate::token::ConnectionToken;
17use ghidra_ipc::frame::FrameDecoder;
18use ghidra_ipc::rpc::{InitAnnounce, RpcNotification, INIT_METHOD, PROTOCOL_VERSION};
19use std::io;
20use std::time::Duration;
21use thiserror::Error;
22use tokio::io::AsyncReadExt;
23use tokio::net::{TcpListener, TcpStream};
24
25#[derive(Debug, Error)]
26pub enum BootError {
27    #[error("no valid worker handshake before the boot deadline")]
28    Timeout,
29    #[error("worker process exited during boot (code {code:?}) — check the worker log on stderr")]
30    WorkerExited { code: Option<i32> },
31    #[error("worker protocol {observed} incompatible with host {expected}")]
32    Incompatible { expected: u32, observed: u32 },
33    #[error(transparent)]
34    Io(#[from] io::Error),
35}
36
37/// The host's handshake listener (bound before the worker is spawned so the port is known).
38pub struct HandshakeListener {
39    listener: TcpListener,
40}
41
42impl HandshakeListener {
43    pub async fn bind() -> io::Result<Self> {
44        Ok(Self {
45            listener: TcpListener::bind("127.0.0.1:0").await?,
46        })
47    }
48
49    pub fn port(&self) -> u16 {
50        self.listener.local_addr().map(|a| a.port()).unwrap_or(0)
51    }
52
53    /// Accept connections until one presents a valid init (correct token + matching protocol) or
54    /// `deadline` elapses. Discards the authenticated stream (used by the version/token unit tests).
55    pub async fn await_valid_worker(
56        &self,
57        expected: &ConnectionToken,
58        deadline: Duration,
59    ) -> Result<InitAnnounce, BootError> {
60        let (announce, _stream) = self.accept_valid(expected, deadline).await?;
61        Ok(announce)
62    }
63
64    /// As `await_valid_worker`, but also hands back the authenticated stream (for real boots).
65    pub async fn accept_valid(
66        &self,
67        expected: &ConnectionToken,
68        deadline: Duration,
69    ) -> Result<(InitAnnounce, TcpStream), BootError> {
70        // A single 5 s bound on each connection's handshake read; the whole accept+read is raced
71        // against `overall` so a peer that connects and sends NOTHING can neither wedge the loop
72        // nor bypass the overall deadline (agy panel R1 — the accept branch used to `.await` the
73        // read outside `select!`, freezing the deadline).
74        const PER_CONN_READ: Duration = Duration::from_secs(5);
75        let overall = tokio::time::sleep(deadline);
76        tokio::pin!(overall);
77        loop {
78            // One connection attempt: accept, then a bounded init-frame read. Built fresh each
79            // iteration and raced against the overall deadline below.
80            let one_connection = async {
81                let (mut sock, _peer) = self.listener.accept().await?;
82                let ann = tokio::time::timeout(PER_CONN_READ, read_init_frame(&mut sock))
83                    .await
84                    .map_err(|_| {
85                        io::Error::new(io::ErrorKind::TimedOut, "handshake read timed out")
86                    })??;
87                Ok::<(InitAnnounce, TcpStream), io::Error>((ann, sock))
88            };
89            tokio::select! {
90                _ = &mut overall => return Err(BootError::Timeout),
91                res = one_connection => {
92                    match res {
93                        Ok((ann, sock)) => {
94                            if !expected.verify(&ann.token) {
95                                continue; // wrong/absent token: not our worker, keep listening
96                            }
97                            if ann.protocol_version != PROTOCOL_VERSION {
98                                return Err(BootError::Incompatible {
99                                    expected: PROTOCOL_VERSION,
100                                    observed: ann.protocol_version,
101                                });
102                            }
103                            return Ok((ann, sock));
104                        }
105                        // Accept error, malformed frame 1, or read timeout: drop and keep listening
106                        // (the overall deadline still governs via the outer select).
107                        Err(_) => continue,
108                    }
109                }
110            }
111        }
112    }
113}
114
115/// Read frame 1 and parse it as the worker/init notification carrying an InitAnnounce.
116async fn read_init_frame(sock: &mut TcpStream) -> io::Result<InitAnnounce> {
117    let mut dec = FrameDecoder::new();
118    let mut tmp = [0u8; 8192];
119    let body = loop {
120        let n = sock.read(&mut tmp).await?;
121        if n == 0 {
122            return Err(io::Error::new(
123                io::ErrorKind::UnexpectedEof,
124                "closed before init",
125            ));
126        }
127        dec.push(&tmp[..n]);
128        if let Some(b) = dec.next_frame().map_err(io::Error::other)? {
129            break b;
130        }
131    };
132    let note: RpcNotification = serde_json::from_slice(&body).map_err(io::Error::other)?;
133    if note.method != INIT_METHOD {
134        return Err(io::Error::new(
135            io::ErrorKind::InvalidData,
136            format!(
137                "frame 1 method is {:?}, expected {INIT_METHOD}",
138                note.method
139            ),
140        ));
141    }
142    let params = note
143        .params
144        .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "init has no params"))?;
145    serde_json::from_value(params).map_err(io::Error::other)
146}
147
148/// A fully booted worker: an authenticated framed connection, the live process handle, and the
149/// lifecycle marked Ready.
150pub struct BootedWorker {
151    pub conn: WorkerConn,
152    pub process: WorkerProcess,
153    pub lifecycle: WorkerLifecycle,
154    pub announce: InitAnnounce,
155}
156
157/// RAII cleanup for the per-boot token directory: removes it (and the plaintext token inside) when
158/// dropped, so the secret is erased on EVERY exit from `boot_worker` — success, `?` error, or async
159/// cancellation (the future dropped mid-await). A manual cleanup on the return paths is skipped on
160/// cancellation and leaks the token on disk (merge-gate finding).
161struct TokenDir(std::path::PathBuf);
162
163impl Drop for TokenDir {
164    fn drop(&mut self) {
165        let _ = std::fs::remove_dir_all(&self.0);
166    }
167}
168
169/// Full boot: bind the listener, write the token file into a private per-boot dir, spawn the
170/// launcher, await a valid handshake, delete the token file, and mark the lifecycle Ready.
171pub async fn boot_worker(cfg: &WorkerConfig) -> Result<BootedWorker, BootError> {
172    let listener = HandshakeListener::bind().await?;
173    let port = listener.port();
174
175    let token = ConnectionToken::generate();
176    let boot_dir = std::env::temp_dir().join(format!("ghidra-mcp-boot-{}", uuid::Uuid::new_v4()));
177    let token_file = boot_dir.join("worker.token");
178    // Arm the token-dir cleanup BEFORE writing the secret, so it fires on every subsequent exit
179    // (incl. the `?`s below and any async cancellation). The dir must outlive the handshake — the
180    // worker reads the token during boot — so the guard drops only when boot_worker returns/unwinds,
181    // exactly when the token is no longer needed (review #2 + merge-gate cancellation finding).
182    let _token_dir = TokenDir(boot_dir.clone());
183
184    std::fs::create_dir_all(&boot_dir)?;
185    // Owner-only the boot dir BEFORE writing the secret: on Windows a fresh temp dir inherits
186    // ambient `Users` read grants, so without this any local user could read worker.token during
187    // the handshake window. Same helper as the server log dir (ghidra_ipc::owner_only).
188    ghidra_ipc::owner_only::restrict_to_owner(&boot_dir)?;
189    token.write_to_file(&token_file)?;
190    let argv = cfg.build_argv(port, &token_file);
191    let mut process =
192        WorkerProcess::spawn(&cfg.analyze_headless_path(), &argv, cfg.max_heap.as_deref())?;
193
194    // Race the handshake against the child dying early: a crashed / mis-launched worker (bad MAXMEM,
195    // EDR-blocked connect, script compile error) would otherwise hang the WHOLE boot_timeout before
196    // yielding Timeout (agy panel R2). Poll try_wait() every 250 ms alongside the accept loop and
197    // fail fast with WorkerExited if the direct child (cmd.exe → java.exe) has already exited.
198    let handshake = listener.accept_valid(&token, cfg.boot_timeout);
199    tokio::pin!(handshake);
200    let outcome = loop {
201        tokio::select! {
202            r = &mut handshake => break r,
203            _ = tokio::time::sleep(Duration::from_millis(250)) => {
204                if let Ok(Some(status)) = process.child.try_wait() {
205                    break Err(BootError::WorkerExited { code: status.code() });
206                }
207            }
208        }
209    };
210    // On error the `?` unwinds through `_token_dir` (removes boot_dir); on success the guard drops at
211    // end of scope, after the token has been consumed by the handshake.
212    let (announce, stream) = outcome?;
213
214    let mut lifecycle = WorkerLifecycle::new();
215    lifecycle
216        .on_ready()
217        .expect("fresh lifecycle must accept on_ready");
218
219    Ok(BootedWorker {
220        conn: WorkerConn::new(stream),
221        process,
222        lifecycle,
223        announce,
224    })
225}
226
227#[cfg(test)]
228mod tests {
229    use super::*;
230    use ghidra_ipc::frame::encode_frame;
231    use ghidra_ipc::rpc::{
232        InitAnnounce, JsonRpcVersion, RpcNotification, INIT_METHOD, PROTOCOL_VERSION,
233    };
234    use tokio::io::AsyncWriteExt;
235    use tokio::net::TcpStream;
236
237    // Regression guard (merge-gate): BootedWorker must stay Send so M1 can hold it across spawned
238    // tokio tasks / in an Arc<Mutex<_>>. Fails to compile if a !Send field creeps back in (e.g. the
239    // windows HANDLE in JobObject losing its `unsafe impl Send`).
240    #[test]
241    fn booted_worker_is_send() {
242        fn is_send<T: Send>() {}
243        is_send::<BootedWorker>();
244    }
245
246    async fn send_init(port: u16, token: &str, protocol: u32) {
247        let mut s = TcpStream::connect(("127.0.0.1", port)).await.unwrap();
248        let ann = InitAnnounce {
249            protocol_version: protocol,
250            ghidra_version: "12.1.2".to_string(),
251            worker_script_version: "0.0.0".to_string(),
252            token: token.to_string(),
253        };
254        let note = RpcNotification {
255            jsonrpc: JsonRpcVersion,
256            method: INIT_METHOD.to_string(),
257            params: Some(serde_json::to_value(&ann).unwrap()),
258        };
259        let frame = encode_frame(&serde_json::to_vec(&note).unwrap()).unwrap();
260        s.write_all(&frame).await.unwrap();
261        // keep the stream open a moment so the host reads the frame
262        tokio::time::sleep(std::time::Duration::from_millis(200)).await;
263    }
264
265    #[tokio::test]
266    async fn accepts_valid_token_and_version() {
267        let h = HandshakeListener::bind().await.unwrap();
268        let port = h.port();
269        let token = crate::token::ConnectionToken::generate();
270        let expect = token.as_str().to_string();
271        let joined = tokio::spawn(async move { send_init(port, &expect, PROTOCOL_VERSION).await });
272        let announce = h
273            .await_valid_worker(&token, std::time::Duration::from_secs(3))
274            .await
275            .expect("handshake ok");
276        assert_eq!(announce.protocol_version, PROTOCOL_VERSION);
277        joined.await.unwrap();
278    }
279
280    #[tokio::test]
281    async fn rejects_bad_token_then_times_out() {
282        let h = HandshakeListener::bind().await.unwrap();
283        let port = h.port();
284        let token = crate::token::ConnectionToken::generate();
285        let joined =
286            tokio::spawn(async move { send_init(port, "not-the-token", PROTOCOL_VERSION).await });
287        let res = h
288            .await_valid_worker(&token, std::time::Duration::from_millis(600))
289            .await;
290        assert!(matches!(res, Err(BootError::Timeout)));
291        joined.await.unwrap();
292    }
293
294    #[tokio::test]
295    async fn rejects_version_mismatch() {
296        let h = HandshakeListener::bind().await.unwrap();
297        let port = h.port();
298        let token = crate::token::ConnectionToken::generate();
299        let expect = token.as_str().to_string();
300        let joined =
301            tokio::spawn(async move { send_init(port, &expect, PROTOCOL_VERSION + 1).await });
302        let res = h
303            .await_valid_worker(&token, std::time::Duration::from_millis(600))
304            .await;
305        assert!(matches!(res, Err(BootError::Incompatible { .. })));
306        joined.await.unwrap();
307    }
308}
ghidra_worker_ctl/boot.rs

ghidra_worker_ctl/
boot.rs
