ghidra_mcp/
execute.rs

1//! The request engine (spec §4.3/§4.4/§4.7/§5): the ordered pipeline every tool call runs through.
2//! Returns either a worker `result` Value (plus a snapshot of the active canonicalizer for the caller
3//! to canonicalize addresses) or a structured `ErrorEnvelope`.
4
5use crate::state::{boot_reattach, expect_success, mk, ServerState, WorkerHolder, WorkerSlot};
6use ghidra_ipc::address::AddressCanonicalizer;
7use ghidra_ipc::error::{ErrorCode, ErrorEnvelope};
8use serde_json::Value;
9use std::sync::Arc;
10use tokio::time::timeout;
11
12/// Terminal "reproducible crash" error (spec §4.3 poison-pill).
13fn poison_terminal() -> ErrorEnvelope {
14    ErrorEnvelope::new(
15        ErrorCode::WorkerUnavailable,
16        "this request reproducibly crashed the worker and has been blocked; change the selector",
17    )
18    .with_details(serde_json::json!({ "reproducible_crash": true }))
19}
20
21fn server_busy() -> ErrorEnvelope {
22    ErrorEnvelope::new(ErrorCode::ServerBusy, "worker at capacity")
23}
24
25fn worker_warming() -> ErrorEnvelope {
26    ErrorEnvelope::new(
27        ErrorCode::WorkerWarming,
28        "the Ghidra worker is still starting",
29    )
30}
31
32/// Run one tool call. `selector` is the poison/short-circuit key component (e.g. the `function` arg,
33/// or the `program_path`); pass `""` for selector-less calls. Returns `(result_value, canon_snapshot)`.
34pub async fn call_worker(
35    state: &Arc<ServerState>,
36    tool: &str,
37    selector: &str,
38    method: &str,
39    params: Value,
40) -> Result<(Value, AddressCanonicalizer), ErrorEnvelope> {
41    // 1) PRE-LOCK poison check (§4.3 round-7 seat-1).
42    // `into_inner` on poisoning: a panic while holding this guard must not wedge every later
43    // tool call — serve the (possibly stale) set instead of panicking the server.
44    {
45        let mut p = state.poison.lock().unwrap_or_else(|e| e.into_inner());
46        if p.contains(tool, selector) {
47            return Err(poison_terminal());
48        }
49    }
50    // 2) Readiness wait (§4.7): NO permit, NO lock held during the wait (R2 seat-A).
51    wait_until_ready(state).await?;
52    // 3) Bounded permit — acquired BEFORE contending for the worker mutex (R3 seat-F).
53    let _permit = state.sem.try_acquire().map_err(|_| server_busy())?;
54    // 4) Run the RPC with transparent crash respawn (§4.3), respawn BOOT OUTSIDE the lock (R3 seat-G).
55    run_with_respawn(state, tool, selector, method, params).await
56}
57
58/// Wait until the worker slot is `Ready`, kicking a single-flight background boot if `Empty`. Waits up
59/// to `warming_deadline` (in ≤250 ms slices, woken early by `boot_done`); a boot longer than the budget
60/// returns `WORKER_WARMING`. Holds the slot lock only for the instantaneous state checks, never across
61/// the wait, and never while holding a permit (R2 seat-A).
62async fn wait_until_ready(state: &Arc<ServerState>) -> Result<(), ErrorEnvelope> {
63    // Ahead of the boot loop on purpose. With no usable configuration there is nothing to wait for,
64    // and falling through would kick a doomed single-flight boot on every pass until the deadline
65    // expired, finally answering WORKER_WARMING — a timeout that hides the actual fault. The
66    // configuration error is deterministic, so report it immediately.
67    let cfg = state.cfg()?;
68    let deadline = tokio::time::Instant::now() + cfg.warming_deadline;
69    loop {
70        {
71            let mut slot = state.worker.lock().await;
72            match slot_kind(&slot) {
73                SlotKind::Ready => return Ok(()),
74                SlotKind::Fatal => {
75                    if let WorkerSlot::Fatal(e) = &*slot {
76                        return Err(e.clone());
77                    }
78                    unreachable!("slot_kind said Fatal");
79                }
80                SlotKind::Empty => {
81                    *slot = WorkerSlot::Booting;
82                    drop(slot);
83                    Arc::clone(state).spawn_boot();
84                }
85                SlotKind::Booting => { /* fall through to wait, lock released below */ }
86            }
87        }
88        let remaining = deadline.saturating_duration_since(tokio::time::Instant::now());
89        if remaining.is_zero() {
90            return Err(worker_warming());
91        }
92        let slice = remaining.min(std::time::Duration::from_millis(250));
93        let _ = tokio::time::timeout(slice, state.boot_done.notified()).await;
94    }
95}
96
97/// Copyable classification of the worker slot, so a caller can inspect state and THEN take a `&mut`
98/// borrow without an overlapping-borrow error (R2 seat-D).
99#[derive(Clone, Copy)]
100enum SlotKind {
101    Empty,
102    Booting,
103    Ready,
104    Fatal,
105}
106
107fn slot_kind(s: &WorkerSlot) -> SlotKind {
108    match s {
109        WorkerSlot::Empty => SlotKind::Empty,
110        WorkerSlot::Booting => SlotKind::Booting,
111        WorkerSlot::Ready(_) => SlotKind::Ready,
112        WorkerSlot::Fatal(_) => SlotKind::Fatal,
113    }
114}
115
116/// Acquire the worker mutex and run the RPC. On a transport crash, respawn OUTSIDE the lock (R3 seat-G)
117/// and retry once; a second crash on the same request poisons the selector (§4.3). Attach state commits
118/// atomically under the lock (R2 seat-C).
119async fn run_with_respawn(
120    state: &Arc<ServerState>,
121    tool: &str,
122    selector: &str,
123    method: &str,
124    params: Value,
125) -> Result<(Value, AddressCanonicalizer), ErrorEnvelope> {
126    let last_good = {
127        let mut slot = state.worker.lock().await;
128        let holder = match &mut *slot {
129            WorkerSlot::Ready(h) => h,
130            WorkerSlot::Fatal(e) => return Err(e.clone()),
131            _ => return Err(worker_warming()),
132        };
133        if !holder.booted.conn.is_poisoned() {
134            match rpc_once(holder, method, params.clone(), state.cfg()?.rpc_deadline).await {
135                RpcOutcome::Ok(v) => {
136                    commit_attach_if_needed(holder, method, &v);
137                    return Ok((v, holder.canon.clone()));
138                }
139                RpcOutcome::Domain(e) => return Err(e),
140                RpcOutcome::Crash => { /* fall through to respawn */ }
141            }
142        }
143        match std::mem::replace(&mut *slot, WorkerSlot::Booting) {
144            WorkerSlot::Ready(mut h) => {
145                h.booted.lifecycle.on_crash();
146                let lg = h.active_program_path.clone();
147                tracing::warn!(last_good = ?lg, %method, "worker crashed; respawning OUTSIDE lock (others get WORKER_WARMING)");
148                drop(h);
149                lg
150            }
151            other => {
152                *slot = other;
153                return Err(worker_warming());
154            }
155        }
156    };
157
158    respawn_and_retry(state, tool, selector, method, params, last_good).await
159}
160
161/// Reboot a fresh worker (OUTSIDE the lock; slot is `Booting` so concurrent callers get WORKER_WARMING),
162/// re-attach last good, and retry ONCE before publishing. A second crash poisons the selector.
163async fn respawn_and_retry(
164    state: &Arc<ServerState>,
165    tool: &str,
166    selector: &str,
167    method: &str,
168    params: Value,
169    last_good: Option<String>,
170) -> Result<(Value, AddressCanonicalizer), ErrorEnvelope> {
171    let mut guard = BootingResetGuard::new(state);
172    let mut fresh = match state.boot_and_probe().await {
173        Ok(h) => h,
174        Err(e) => {
175            reset_booting_to_empty(state, &mut guard).await;
176            return Err(e);
177        }
178    };
179    boot_reattach(&mut fresh, last_good).await;
180    match rpc_once(&mut fresh, method, params, state.cfg()?.rpc_deadline).await {
181        RpcOutcome::Ok(v) => {
182            commit_attach_if_needed(&mut fresh, method, &v);
183            let canon = fresh.canon.clone();
184            publish_ready(state, fresh, &mut guard).await;
185            Ok((v, canon))
186        }
187        RpcOutcome::Domain(e) => {
188            publish_ready(state, fresh, &mut guard).await;
189            Err(e)
190        }
191        RpcOutcome::Crash => {
192            drop(fresh);
193            state
194                .poison
195                .lock()
196                .unwrap_or_else(|e| e.into_inner())
197                .insert(tool, selector);
198            reset_booting_to_empty(state, &mut guard).await;
199            Err(poison_terminal())
200        }
201    }
202}
203
204/// RAII guard resetting a `Booting` slot back to `Empty` if a respawn future is dropped before reaching
205/// a terminal state (R5 cancellation-safety). Drop cannot await, so it uses a sync `try_lock`; on
206/// contention it detaches a cleanup task. Idempotent: only resets if the slot is still `Booting`.
207struct BootingResetGuard<'a> {
208    state: &'a Arc<ServerState>,
209    armed: bool,
210}
211
212impl<'a> BootingResetGuard<'a> {
213    fn new(state: &'a Arc<ServerState>) -> Self {
214        Self { state, armed: true }
215    }
216    fn disarm(&mut self) {
217        self.armed = false;
218    }
219}
220
221impl Drop for BootingResetGuard<'_> {
222    fn drop(&mut self) {
223        if !self.armed {
224            return;
225        }
226        if let Ok(mut slot) = self.state.worker.try_lock() {
227            if matches!(&*slot, WorkerSlot::Booting) {
228                *slot = WorkerSlot::Empty;
229            }
230            self.state.boot_done.notify_waiters();
231        } else {
232            let s = Arc::clone(self.state);
233            tokio::spawn(async move {
234                let mut slot = s.worker.lock().await;
235                if matches!(&*slot, WorkerSlot::Booting) {
236                    *slot = WorkerSlot::Empty;
237                }
238                s.boot_done.notify_waiters();
239            });
240        }
241    }
242}
243
244/// Publish a freshly-booted worker into the slot (from `Booting`) and wake warming waiters. Disarms the
245/// cancellation guard INSIDE the lock, after the write (R6 seat-P).
246async fn publish_ready(
247    state: &Arc<ServerState>,
248    holder: WorkerHolder,
249    guard: &mut BootingResetGuard<'_>,
250) {
251    {
252        let mut slot = state.worker.lock().await;
253        *slot = WorkerSlot::Ready(holder);
254        guard.disarm();
255    }
256    state.boot_done.notify_waiters();
257}
258
259/// Reset a `Booting` slot back to `Empty` (retryable) after a failed respawn; wake warming waiters.
260/// Only resets if still `Booting`. Disarms the guard INSIDE the lock, after the write (R6 seat-P).
261async fn reset_booting_to_empty(state: &Arc<ServerState>, guard: &mut BootingResetGuard<'_>) {
262    {
263        let mut slot = state.worker.lock().await;
264        if matches!(&*slot, WorkerSlot::Booting) {
265            *slot = WorkerSlot::Empty;
266        }
267        guard.disarm();
268    }
269    state.boot_done.notify_waiters();
270}
271
272/// If the just-succeeded call was `attach_program`, commit the new active program + canonicalizer into
273/// the holder WHILE THE SLOT LOCK IS STILL HELD (R2 seat-C atomicity). A parse failure leaves the prior
274/// attachment intact.
275fn commit_attach_if_needed(holder: &mut WorkerHolder, method: &str, result: &Value) {
276    if method != "attach_program" {
277        return;
278    }
279    if let Ok(a) = serde_json::from_value::<ghidra_ipc::protocol::AttachResult>(result.clone()) {
280        holder.active_program_path = a.attached;
281        holder.canon = AddressCanonicalizer::new(a.default_space, a.pointer_size);
282    }
283}
284
285enum RpcOutcome {
286    Ok(Value),
287    Domain(ErrorEnvelope),
288    Crash,
289}
290
291/// One RPC with poison-on-drop (§4.4) + rpc_deadline (§4.7). A transport error OR a deadline overrun is
292/// a crash (respawn); a worker `Failure` frame is a domain error (worker alive).
293async fn rpc_once(
294    holder: &mut WorkerHolder,
295    method: &str,
296    params: Value,
297    deadline: std::time::Duration,
298) -> RpcOutcome {
299    let req = mk(method, params);
300    holder.booted.conn.poison();
301    let res = timeout(deadline, holder.booted.conn.request_raw(&req)).await;
302    match res {
303        Ok(Ok(resp)) => {
304            holder.booted.conn.unpoison();
305            match expect_success(resp) {
306                Ok(v) => RpcOutcome::Ok(v),
307                Err(e) => RpcOutcome::Domain(e),
308            }
309        }
310        Ok(Err(_io)) => RpcOutcome::Crash,
311        Err(_elapsed) => RpcOutcome::Crash,
312    }
313}
314
315/// Attach a program. The active-program + canon commit happens ATOMICALLY inside the engine under the
316/// worker lock (via `commit_attach_if_needed`) — NOT here (R2 seat-C). This wrapper only supplies the
317/// poison-key selector and runs the same crash/respawn engine (§4.3).
318pub async fn attach_program(
319    state: &Arc<ServerState>,
320    program_path: &str,
321) -> Result<Value, ErrorEnvelope> {
322    let params = serde_json::json!({ "program_path": program_path });
323    let (v, _canon) = call_worker(
324        state,
325        "attach_program",
326        program_path,
327        "attach_program",
328        params,
329    )
330    .await?;
331    Ok(v)
332}
333
334/// Snapshot the active canonicalizer (for canonicalizing an error-path candidate list).
335pub async fn current_canon(state: &Arc<ServerState>) -> Option<AddressCanonicalizer> {
336    let slot = state.worker.lock().await;
337    match &*slot {
338        WorkerSlot::Ready(h) => Some(h.canon.clone()),
339        _ => None,
340    }
341}
ghidra_mcp/execute.rs

ghidra_mcp/
execute.rs
