Add export_annotations and save_project to keep the work

* Everything an agent renamed, commented or patched was lost with
  `close_file`.
* Record every state-changing command the tools run (the annotation
  commands and verified writes, at resolved addresses) in a journal kept
  until the next open, also past `close_file`. `export_annotations`
  returns it as a Rizin script to replay with `. <script>`.
* `save_project` saves the session as a Rizin project with `Ps`; it writes
  a file, so it is offered only with `--allow-write`, and it checks that
  the file was written.
* A test replays an exported script on a fresh analysis and reopens a
  saved project with `rizin -p`.
This commit is contained in:
Anton Kochkov authored and NOT XVilka committed 2026-10-30 18:35:59 +08:00
1 parent aa53d2dbec
commit bb18b601ba
5 files changed
+183 -4

No files matched your search

+8
View File
@@ -129,6 +129,14 @@ produces the full listing, which the server then slices.
understood so later disassembly and decompilation show it. Each returns
the state Rizin stored, and fails when Rizin only logs an error.
**Keeping the work**
- `export_annotations` — the renames, comments, flags, signatures and byte
writes made through the tools since the file was opened, as a Rizin script
(also after `close_file`); replay it with `. <script>` on the same binary.
- `save_project` — save the session as a Rizin project (`.rzdb`, reopened with
`rizin -p`); writes a file, so it needs `--allow-write`. Bytes written to a
file opened with `write: true` go to the file directly.
**Writing** (disabled by default — see *Security*)
- `write_hex`, `write_string`, `write_assembly` (instructions separated by `;`
or newlines, assembled for the destination address). Every write is read
+7
View File
@@ -248,6 +248,13 @@ pub struct SetSignatureParams {
pub signature: String,
}
/// Parameters for `save_project`.
#[derive(Debug, Deserialize, JsonSchema)]
pub struct SaveProjectParams {
/// Where to write the Rizin project (`.rzdb`), e.g. `/tmp/sample.rzdb`.
pub path: String,
}
/// Parameters for `write_hex` (write mode only).
#[derive(Debug, Deserialize, JsonSchema)]
pub struct WriteHexParams {
+60 -1
View File
@@ -75,6 +75,7 @@ pub const WRITE_TOOLS: &[&str] = &[
"write_string",
"write_assembly",
"flirt_create",
"save_project",
];
/// The tools that change the analysis (names, comments, flags, types),
@@ -255,7 +256,8 @@ impl RizinServer {
}
/// Run a command that changes state, failing also when Rizin logs an
/// error but reports success, as `f`, `wx` and `Fs` do.
/// error but reports success, as `f`, `wx` and `Fs` do. A successful
/// command is recorded for `export_annotations`.
async fn cmd_strict(&self, cmd: &str) -> Result<String, ErrorData> {
let out = self.manager.exec(cmd).await.mcp()?;
let out = crate::backend::check(cmd, out).mcp()?;
@@ -273,6 +275,7 @@ impl RizinServer {
}
.to_error_data());
}
self.manager.record(cmd.to_string()).await;
Ok(out.output)
}
@@ -390,6 +393,10 @@ impl RizinServer {
to_hex(bytes)
)));
}
let addr = self.resolve(loc).await?;
self.manager
.record(format!("wx {} @ {addr}", to_hex(bytes)))
.await;
Ok(())
}
@@ -1217,6 +1224,58 @@ impl RizinServer {
Ok(self.json_result(self.manager.cmdj(&format!("afsj @ {addr}")).await.mcp()?))
}
// --- Persistence -----------------------------------------------------
/// The recorded changes as a Rizin script.
#[tool(
description = "Export the renames, comments, flags, signatures and byte writes made \
through the tools since the file was opened, as a Rizin script (also after \
close_file). Replay it with `. <script>` in Rizin (or the command tool) on the same \
binary to restore that work."
)]
async fn export_annotations(&self) -> Result<CallToolResult, ErrorData> {
let (file, commands) = self.manager.journal().await;
let mut script = format!(
"# rizin-mcp changes for {}\n",
file.as_deref().unwrap_or("(no file)")
);
for cmd in &commands {
script.push_str(cmd);
script.push('\n');
}
Ok(self.json_result(json!({
"file": file,
"commands": commands.len(),
"script": script,
})))
}
/// Save a Rizin project (`Ps`).
#[tool(
description = "Save the session (analysis, names, comments, flags, types) as a Rizin \
project file (.rzdb), to reopen later with `rizin -p <file>`. Requires --allow-write \
(it writes a file to disk) (Ps)."
)]
async fn save_project(
&self,
Parameters(p): Parameters<SaveProjectParams>,
) -> Result<CallToolResult, ErrorData> {
if !self.manager.config().allow_write {
return Err(RizinError::NotPermitted(
"save_project writes a file to disk; start the server with --allow-write".into(),
)
.to_error_data());
}
let path = quote_arg(&p.path).mcp()?;
let out = self.manager.exec(&format!("Ps {path}")).await.mcp()?;
crate::backend::check(&format!("Ps {path}"), out).mcp()?;
let bytes = std::fs::metadata(&p.path).map(|m| m.len()).map_err(|e| {
RizinError::Io(format!("the project was not written to {}: {e}", p.path))
.to_error_data()
})?;
Ok(self.json_result(json!({ "path": p.path, "bytes": bytes })))
}
// --- Writing (guarded) ------------------------------------------------
/// Write hex bytes (`wx`).
+23 -3
View File
@@ -72,6 +72,11 @@ pub struct SessionManager {
struct Slot {
session: Option<Session>,
ended: Option<String>,
/// The file the journal belongs to.
journal_file: Option<String>,
/// The state-changing commands run through the tools since the file was
/// opened. Kept past a close, so the work can still be exported.
journal: Vec<String>,
}
impl Slot {
@@ -193,6 +198,8 @@ impl SessionManager {
*self.inner.lock().await = Slot {
session: Some(session),
ended: None,
journal_file: Some(path.to_string()),
journal: Vec::new(),
};
Ok(OpenInfo {
@@ -245,9 +252,22 @@ impl SessionManager {
/// Close the current session, if any. Returns whether one was open.
pub async fn close(&self) -> bool {
std::mem::take(&mut *self.inner.lock().await)
.session
.is_some()
let mut slot = self.inner.lock().await;
slot.ended = None;
slot.session.take().is_some()
}
/// Record a command that changed the analysis or the bytes, for
/// [`Self::journal`].
pub async fn record(&self, cmd: String) {
self.inner.lock().await.journal.push(cmd);
}
/// The file of the last opened session and the state-changing commands
/// recorded for it, in order.
pub async fn journal(&self) -> (Option<String>, Vec<String>) {
let slot = self.inner.lock().await;
(slot.journal_file.clone(), slot.journal.clone())
}
/// A snapshot of the current session state. A session that died (e.g.
+85
View File
@@ -781,3 +781,88 @@ async fn annotations_change_the_analysis() {
assert_eq!(sig["args"][1]["name"], "args", "{t:?}: {sig}");
}
}
#[tokio::test]
async fn work_can_be_exported_and_saved() {
let tmp = std::env::temp_dir();
let id = std::process::id();
for t in transports() {
let script = tmp.join(format!("rizin-mcp-{id}-{t:?}.rz"));
let project = tmp.join(format!("rizin-mcp-{id}-{t:?}.rzdb"));
let s = Session::start(t, |c| c.allow_write = true).await;
s.ok(
"open_file",
json!({ "path": target(), "analysis": "basic" }),
)
.await;
s.ok(
"rename_function",
json!({ "at": "main", "name": "entry_logic" }),
)
.await;
s.ok(
"set_comment",
json!({ "address": "entry_logic", "text": "parses; argv @ start" }),
)
.await;
s.ok(
"set_flag",
json!({ "name": "mcp.mark", "address": "entry0" }),
)
.await;
let saved = s
.ok("save_project", json!({ "path": project.to_string_lossy() }))
.await;
assert!(saved["bytes"].as_u64().unwrap() > 0, "{t:?}");
// The journal survives closing the file.
s.ok("close_file", json!({})).await;
let export = s.ok("export_annotations", json!({})).await;
assert_eq!(export["commands"], 4, "{t:?}: {export}"); // afn, CC-, CC, f+
std::fs::write(&script, export["script"].as_str().unwrap()).unwrap();
// Replaying the script on a fresh analysis restores the work.
s.ok(
"open_file",
json!({ "path": target(), "analysis": "basic" }),
)
.await;
assert_eq!(s.ok("export_annotations", json!({})).await["commands"], 0);
let r = s
.call(
"command",
json!({ "command": format!(". {}", script.to_string_lossy()) }),
)
.await
.expect("replay");
assert_ne!(r.is_error, Some(true), "{t:?}: {r:?}");
let info = s
.items("function_info", json!({ "at": "entry_logic" }))
.await;
assert_eq!(info[0]["name"], "entry_logic", "{t:?}");
let c = s
.call("command", json!({ "command": "CC. @ entry_logic" }))
.await
.expect("CC.");
assert_eq!(
c.content[0].as_text().unwrap().text.trim(),
"parses; argv @ start",
"{t:?}"
);
s.ok("seek", json!({ "address": "mcp.mark" })).await;
// The project reopens in rizin with the renamed function.
if pipe_available() {
let out = std::process::Command::new(rizin_binary())
.args(["-q", "-p"])
.arg(&project)
.args(["-c", "afl~entry_logic"])
.output()
.expect("rizin -p");
let stdout = String::from_utf8_lossy(&out.stdout);
assert!(stdout.contains("entry_logic"), "{t:?}: {stdout:?}");
}
let _ = std::fs::remove_file(&script);
let _ = std::fs::remove_file(&project);
}
}