From 0f008973bb90190206b8a66effe5b5f2b789f992 Mon Sep 17 00:00:00 2001 From: Hung Pham Date: Sat, 12 Sep 2026 10:16:44 +0700 Subject: [PATCH 1/2] Split storage of binary graph into different database --- Cargo.lock | 2 + crates/codegraph-binary/src/extract.rs | 239 ++++- crates/codegraph-binary/src/model.rs | 8 + crates/codegraph-extract/Cargo.toml | 2 + crates/codegraph-extract/src/bingraph.rs | 954 +++++++++++++++++++ crates/codegraph-extract/src/config.rs | 75 ++ crates/codegraph-extract/src/lib.rs | 6 + crates/codegraph-extract/src/orchestrator.rs | 35 +- crates/codegraph-graph/src/lib.rs | 14 +- crates/codegraph-mcp/src/lib.rs | 12 + crates/codegraph-mcp/src/tools.rs | 161 ++++ 11 files changed, 1488 insertions(+), 20 deletions(-) create mode 100644 crates/codegraph-extract/src/bingraph.rs diff --git a/Cargo.lock b/Cargo.lock index e51ed268c..5b785ce28 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -847,7 +847,9 @@ dependencies = [ "ignore", "indicatif", "rayon", + "rusqlite", "serde", + "serde_json", "tempfile", "tokio", "toml", diff --git a/crates/codegraph-binary/src/extract.rs b/crates/codegraph-binary/src/extract.rs index db394ad39..4403b8810 100644 --- a/crates/codegraph-binary/src/extract.rs +++ b/crates/codegraph-binary/src/extract.rs @@ -10,6 +10,8 @@ use serde_json::Value; use std::collections::{HashMap, HashSet}; use std::path::Path; +use crate::model::{EntryPoint, ExportEntry}; + /// Trích xuất toàn bộ thông tin từ binary thành `ParseResult`. /// Gọi `aaa` một lần trong session, rồi query. pub fn extract_binary( @@ -70,13 +72,24 @@ fn do_extract( // 1. Functions (`aflj`) let functions = parse_aflj(session)?; - // Parse exports (`iEj`) for JNI address-based detection (catches stripped binaries). + // Parse exports (`iEj`) — JNI detection + index export như entrypoint cho link chéo. let exports = parse_iej(session)?; let jni_export_map: HashMap = exports .iter() .filter(|e| is_jni_name(e.name.as_deref().unwrap_or(""))) .filter_map(|e| e.vaddr.map(|v| (v, e.name.clone().unwrap_or_default()))) .collect(); + // Export theo vaddr — function trùng địa chỉ chỉ cần gắn annotation. + let export_by_addr: HashMap = exports + .iter() + .filter_map(|e| e.vaddr.map(|v| (v, e))) + .collect(); + // Entry points (`iej`) — điểm bắt đầu phân tích executable. + let entrypoints = parse_entrypoints(session)?; + let entry_by_addr: HashMap = entrypoints + .iter() + .filter_map(|e| e.vaddr.map(|v| (v, e))) + .collect(); let mut symbols: Vec = Vec::new(); let mut chains: HashMap> = HashMap::new(); let mut calls: Vec = Vec::new(); @@ -103,7 +116,7 @@ fn do_extract( fn_id_to_name.insert(id, name.clone()); // r2 6.x tự sinh symbol C++: class.X, method.Class.foo, namespace.X, enum.X let (kind, name) = classify_symbol(&raw_name, &name); - // JNI enrichment: name-based + address-based (via iEj export table). + // Enrichment theo địa chỉ: JNI (Java_/JNI_), export table, entry point. let mut annotations = Vec::new(); if is_jni_name(&name) || jni_export_map.contains_key(&addr) { annotations.push(Annotation { @@ -112,6 +125,20 @@ fn do_extract( line: 0, }); } + if let Some(export) = export_by_addr.get(&addr) { + annotations.push(export_annotation(export)); + } + if let Some(ep) = entry_by_addr.get(&addr) { + let mut args = HashMap::new(); + if let Some(n) = &ep.name { + args.insert("name".to_string(), n.clone()); + } + annotations.push(Annotation { + name: "entrypoint".to_string(), + args, + line: 0, + }); + } symbols.push(Symbol { id, name, @@ -130,35 +157,85 @@ fn do_extract( }); } + // 1b. Exports không trùng function nào (data export, stripped binary…) — + // tạo symbol riêng để bên ngoài link vào được theo tên export. + for export in &exports { + let Some(vaddr) = export.vaddr else { continue }; + if fn_by_addr.contains_key(&vaddr) { + continue; + } + let raw_name = export + .name + .clone() + .unwrap_or_else(|| format!("exp.{vaddr:x}")); + let name = demangle(&strip_r2_prefix(&raw_name)); + let (kind, name) = classify_symbol(&raw_name, &name); + let id = next_id; + next_id += 1; + fn_by_addr.insert(vaddr, id); + fn_id_to_name.insert(id, name.clone()); + let mut annotations = Vec::new(); + if is_jni_name(&name) || jni_export_map.contains_key(&vaddr) { + annotations.push(Annotation { + name: "jni".to_string(), + args: HashMap::new(), + line: 0, + }); + } + annotations.push(export_annotation(export)); + symbols.push(Symbol { + id, + name, + kind, + scope: codegraph_core::ScopeLevel::Global, + scope_id: 0, + type_ref: 0, + type_name: None, + file: path_str.to_string(), + line: vaddr.try_into().unwrap_or(0), + end_line: vaddr + .saturating_add(export.size.unwrap_or(0)) + .try_into() + .unwrap_or(u32::MAX), + signature: Some(format!( + "export ({})", + export.type_.as_deref().unwrap_or("?") + )), + doc: None, + annotations, + language: "binary".to_string(), + }); + } + // 2. Imports (`iij`) — tạo symbol; bỏ qua function entry "sym.imp." let imports = parse_iij(session)?; let mut import_name_to_id: HashMap = HashMap::new(); let mut plt_by_addr: HashMap = HashMap::new(); for imp in &imports { + // Giữ tên thuần làm name — bên ngoài link vào theo đúng tên hàm library. + // Tên library đưa vào type_name/annotation args thay vì đổi name. let clean = imp.import.as_deref().unwrap_or("?"); - let count = imports - .iter() - .filter(|i| i.import.as_deref() == Some(clean)) - .count(); - let name = if count > 1 { - format!("{clean} ({})", imp.lib.as_deref().unwrap_or("?")) - } else { - clean.to_string() - }; let id = next_id; next_id += 1; import_name_to_id.insert(clean.to_string(), id); if let Some(plt) = imp.plt { - plt_by_addr.insert(plt, name.clone()); + plt_by_addr.insert(plt, clean.to_string()); + } + let mut import_args = HashMap::new(); + if let Some(lib) = &imp.lib { + import_args.insert("lib".to_string(), lib.clone()); + } + if let Some(bind) = &imp.bind { + import_args.insert("bind".to_string(), bind.clone()); } symbols.push(Symbol { id, - name, + name: clean.to_string(), kind: SymbolKind::Function, scope: codegraph_core::ScopeLevel::Global, scope_id: 0, type_ref: 0, - type_name: None, + type_name: imp.lib.clone(), file: path_str.to_string(), line: imp.plt.unwrap_or(0).try_into().unwrap_or(0), end_line: 0, @@ -166,7 +243,7 @@ fn do_extract( doc: imp.lib.clone(), annotations: vec![Annotation { name: "import".to_string(), - args: HashMap::new(), + args: import_args, line: 0, }], language: "binary".to_string(), @@ -248,6 +325,27 @@ fn parse_izj(session: &mut dyn R2Client) -> Result, Error> { parse_array(session.cmdj("izj")?) } +/// Parse entry points từ `iej` — điểm bắt đầu phân tích executable. +fn parse_entrypoints(session: &mut dyn R2Client) -> Result, Error> { + parse_array(session.cmdj("iej")?) +} + +/// Annotation `"export"` kèm bind/type nếu có — dùng cho link chéo giữa binary. +fn export_annotation(export: &ExportEntry) -> Annotation { + let mut args = HashMap::new(); + if let Some(bind) = &export.bind { + args.insert("bind".to_string(), bind.clone()); + } + if let Some(t) = &export.type_ { + args.insert("type".to_string(), t.clone()); + } + Annotation { + name: "export".to_string(), + args, + line: 0, + } +} + fn build_signature(addr: u64, size: u64, entry: &FnEntry) -> String { let mut parts = vec![format!("0x{addr:x}")]; if size > 0 { @@ -706,4 +804,115 @@ mod tests { assert!(jni_export_map.contains_key(&addr)); // The function with this addr would get jni annotation even if r2 renamed it } + + struct FullMock { + responses: HashMap, + } + impl R2Client for FullMock { + fn cmd(&mut self, _cmd: &str) -> Result { + Ok(String::new()) + } + fn cmdj(&mut self, cmd: &str) -> Result { + Ok(self.responses.get(cmd).cloned().unwrap_or(json!([]))) + } + } + + #[test] + fn test_extract_entrypoint_export_annotations() { + let mock = FullMock { + responses: HashMap::from([ + ( + "aflj".to_string(), + json!([ + {"addr": 4196, "name": "method.Foo.bar", "size": 16} + ]), + ), + ( + "iEj".to_string(), + json!([ + {"name": "method.Foo.bar", "vaddr": 4196, "bind": "GLOBAL", "type": "FUNC"}, + {"name": "exported_data", "vaddr": 8192, "bind": "GLOBAL", "type": "OBJ"} + ]), + ), + ( + "iej".to_string(), + json!([{"vaddr": 4196, "name": "entry0"}]), + ), + ]), + }; + let result = do_extract( + &mut FullMock { + responses: mock.responses.clone(), + }, + Path::new("/tmp/fake.so"), + 0, + false, + AnalysisDepth::default(), + ) + .unwrap(); + + let func = result + .symbols + .iter() + .find(|s| s.name == "method.Foo.bar") + .expect("function symbol phải tồn tại"); + assert!( + func.annotations.iter().any(|a| a.name == "export"), + "function trùng vaddr export phải gắn annotation export" + ); + assert!( + func.annotations.iter().any(|a| a.name == "entrypoint"), + "function trùng vaddr entrypoint phải gắn annotation entrypoint" + ); + // Export không trùng function → symbol riêng. + let data_export = result + .symbols + .iter() + .find(|s| s.name == "exported_data") + .expect("export-only symbol phải được tạo"); + assert!(data_export.annotations.iter().any(|a| a.name == "export")); + } + + #[test] + fn test_extract_import_keeps_clean_name() { + let mock = FullMock { + responses: HashMap::from([ + ("aflj".to_string(), json!([])), + ( + "iij".to_string(), + json!([ + {"import": "memcpy", "plt": 100, "lib": "libc.so"}, + {"import": "memcpy", "plt": 200, "lib": "libb.so"} + ]), + ), + ]), + }; + let result = do_extract( + &mut FullMock { + responses: mock.responses.clone(), + }, + Path::new("/tmp/fake.so"), + 0, + false, + AnalysisDepth::default(), + ) + .unwrap(); + + let imports: Vec<_> = result + .symbols + .iter() + .filter(|s| s.annotations.iter().any(|a| a.name == "import")) + .collect(); + assert_eq!(imports.len(), 2, "2 import entries → 2 symbol"); + assert!( + imports.iter().all(|s| s.name == "memcpy"), + "import phải giữ tên thuần (không đổi thành 'memcpy (lib)')" + ); + assert!( + imports + .iter() + .any(|s| s.type_name.as_deref() == Some("libc.so")), + "tên library phải nằm trong type_name" + ); + } } diff --git a/crates/codegraph-binary/src/model.rs b/crates/codegraph-binary/src/model.rs index 644d563aa..f5e193ef9 100644 --- a/crates/codegraph-binary/src/model.rs +++ b/crates/codegraph-binary/src/model.rs @@ -107,6 +107,14 @@ pub struct ExportEntry { pub type_: Option, } +/// Entry point từ `iej` (entry addresses của executable). +#[derive(Debug, Deserialize)] +pub struct EntryPoint { + pub vaddr: Option, + pub paddr: Option, + pub name: Option, +} + /// String từ `izj` / `izzj`. #[derive(Debug, Deserialize)] pub struct StrEntry { diff --git a/crates/codegraph-extract/Cargo.toml b/crates/codegraph-extract/Cargo.toml index 97354b9c4..82338e9b3 100644 --- a/crates/codegraph-extract/Cargo.toml +++ b/crates/codegraph-extract/Cargo.toml @@ -32,6 +32,8 @@ tree-sitter-swift = { workspace = true, optional = true } # tree-sitter-kotlin = { workspace = true, optional = true } tree-sitter-lua = { workspace = true, optional = true } codegraph-binary = { path = "../codegraph-binary", optional = true } +rusqlite = { workspace = true } +serde_json = { workspace = true } ignore = { workspace = true } rayon = { workspace = true } camino = { workspace = true } diff --git a/crates/codegraph-extract/src/bingraph.rs b/crates/codegraph-extract/src/bingraph.rs new file mode 100644 index 000000000..cf9aeb849 --- /dev/null +++ b/crates/codegraph-extract/src/bingraph.rs @@ -0,0 +1,954 @@ +//! Binary graph runtime — dataset **riêng** cho symbol binary (pattern +//! `codegraph-docs`), chạy trên trait [`Storage`] của codegraph-graph nên hỗ trợ +//! mọi backend: sqlite (`.codegraph/binary.sqlite`), lmdb, redis (keyspace +//! `codegraph:binary`), in-memory. Không đụng bảng/keys của code index lẫn docs. +//! +//! Lazy: open chỉ mở storage (không load symbol nào vào RAM). Query đi qua: +//! - **Name trie** (`Search` trên cùng dataset — record index riêng bắt đầu +//! từ [`RECORD_START`]): substring/prefix/exact search theo tên, persist. +//! - **Secondary index** trên record-meta stream (`set_meta`/`get_meta`, keyed +//! bằng hash của tên key): `all` / `kind:{k}` / `flag:{f}` / `addr:{a}` / +//! `ep` / `path:{p}` → danh sách symbol id (JSON). Mỗi danh sách chỉ được +//! load lúc query, phân trang ở bước cuối. +//! - **Symbol JSON** qua `save_symbol`/`load_symbol`, chain qua +//! `set_chain`/`get_chain`, call records qua `set_call_records`. + +use crate::config::ExtractConfig; +use camino::Utf8Path; +use codegraph_core::{CallRecord, Error, Result, Symbol, SymbolKind}; +use codegraph_graph::{ + open_keyspace_storage, ParseResult, Search, SearchError, Storage, StorageError, +}; +use serde::{Deserialize, Serialize}; +use std::collections::HashMap; +use std::sync::Arc; +use tokio::sync::RwLock; + +/// Base id mặc định cho symbol binary graph — tránh dải docs (1e9/3e9) và +/// dải code index. Override bằng `[bingraph] bin_base`. +pub const DEFAULT_BIN_BASE: u64 = 2_000_000_000; + +/// Sharding của name trie (GraphIndex dùng 64 cho chain engine). +const BIN_SHARDING: usize = 64; + +/// Record index đầu tiên của name trie — các số nhỏ hơn là dải của secondary +/// index (hash key). Trie record tăng dần từ đây. +const RECORD_START: usize = 10_000; + +type SharedStorage = Arc>; + +/// Flag chính của symbol binary (annotation → index key). +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "lowercase")] +pub enum BinFlag { + Import, + Export, + Entrypoint, + Jni, +} + +impl BinFlag { + pub fn as_str(&self) -> &'static str { + match self { + BinFlag::Import => "import", + BinFlag::Export => "export", + BinFlag::Entrypoint => "entrypoint", + BinFlag::Jni => "jni", + } + } + + pub fn parse(s: &str) -> Option { + match s { + "import" => Some(BinFlag::Import), + "export" => Some(BinFlag::Export), + "entrypoint" => Some(BinFlag::Entrypoint), + "jni" => Some(BinFlag::Jni), + _ => None, + } + } +} + +/// Tất cả flag của một symbol (mỗi annotation khớp một flag — một symbol có +/// thể nằm trong nhiều index, vd jni + export). +fn symbol_flags(sym: &Symbol) -> Vec { + let mut v = Vec::new(); + let has = |n: &str| sym.annotations.iter().any(|a| a.name == n); + for (n, f) in [ + ("jni", BinFlag::Jni), + ("entrypoint", BinFlag::Entrypoint), + ("export", BinFlag::Export), + ("import", BinFlag::Import), + ] { + if has(n) { + v.push(f); + } + } + v +} + +/// Flag đại diện hiển thị (ưu tiên jni > entrypoint > export > import). +fn annotation_flag(sym: &Symbol) -> Option { + symbol_flags(sym).into_iter().next() +} + +/// Một row symbol trả về từ query — decode từ `Symbol` (lazy theo id). +#[derive(Debug, Clone, Serialize)] +pub struct BinSymbolRow { + pub id: u64, + pub name: String, + pub kind: String, + pub addr: u64, + pub end_addr: u64, + pub path: String, + pub flag: Option, + pub lib: Option, + pub signature: Option, +} + +impl From for BinSymbolRow { + fn from(s: Symbol) -> Self { + BinSymbolRow { + id: s.id, + flag: annotation_flag(&s).map(|f| f.as_str().to_string()), + lib: s.type_name.clone().or_else(|| s.doc.clone()), + kind: format!("{:?}", s.kind), + name: s.name, + addr: u64::from(s.line), + end_addr: u64::from(s.end_line), + path: s.file, + signature: s.signature, + } + } +} + +/// Mode search theo tên. +#[derive(Debug, Clone, Copy, Default, PartialEq)] +pub enum NameMatch { + Exact, + Prefix, + Suffix, + /// Chứa ở giữa — đi qua name trie (radix DFS + KMP), không scan. + #[default] + Contains, +} + +/// Sort order cho list. +#[derive(Debug, Clone, Copy, Default)] +pub enum ListOrder { + #[default] + Name, + Addr, + Id, +} + +/// Một trang kết quả list/search. +#[derive(Debug, Clone, Serialize)] +pub struct BinPage { + pub rows: Vec, + pub total: u64, + pub offset: u64, +} + +#[derive(Debug, Clone, Serialize)] +pub struct BinStats { + pub symbols: u64, + pub entrypoints: u64, + pub imports: u64, + pub exports: u64, + pub binaries: u64, +} + +// ── Secondary index trên record-meta stream ── + +/// FNV-1a 64 — hash key secondary index thành record id trên meta stream. +/// Trie record (bắt đầu từ [`RECORD_START`]) và hash key có thể trùng số trong +/// lý thuyết nhưng xác suất ~0 (FNV phân bố đều trên 2^64). +fn kv_record(key: &str) -> usize { + let mut h: u64 = 0xcbf2_9ce4_8422_2325; + for b in key.as_bytes() { + h ^= u64::from(*b); + h = h.wrapping_mul(0x0000_0100_0000_01b3); + } + h as usize +} + +fn ids_json(ids: &[u64]) -> Result> { + serde_json::to_vec(ids).map_err(|e| Error::Db(format!("encode ids: {e}"))) +} + +async fn meta_ids(storage: &SharedStorage, key: &str) -> Result> { + let s = storage.read().await; + let bytes = s.get_meta(kv_record(key)).await.map_err(db_err)?; + Ok(bytes + .and_then(|b| serde_json::from_slice::>(&b).ok()) + .unwrap_or_default()) +} + +/// Đọc ids theo record index thô (dùng cho meta của name-trie record). +async fn meta_ids_at(storage: &SharedStorage, record: usize) -> Result> { + let s = storage.read().await; + let bytes = s.get_meta(record).await.map_err(db_err)?; + Ok(bytes + .and_then(|b| serde_json::from_slice::>(&b).ok()) + .unwrap_or_default()) +} + +async fn meta_set_ids(storage: &SharedStorage, key: &str, ids: &[u64]) -> Result<()> { + let mut s = storage.write().await; + s.set_meta(kv_record(key), &ids_json(ids)?) + .await + .map_err(db_err) +} + +async fn meta_set_ids_at(storage: &SharedStorage, record: usize, ids: &[u64]) -> Result<()> { + let mut s = storage.write().await; + s.set_meta(record, &ids_json(ids)?).await.map_err(db_err) +} + +async fn meta_add_id(storage: &SharedStorage, key: &str, id: u64) -> Result<()> { + let mut ids = meta_ids(storage, key).await?; + if !ids.contains(&id) { + ids.push(id); + meta_set_ids(storage, key, &ids).await?; + } + Ok(()) +} + +async fn meta_remove_id(storage: &SharedStorage, key: &str, id: u64) -> Result<()> { + let mut ids = meta_ids(storage, key).await?; + let before = ids.len(); + ids.retain(|&x| x != id); + if ids.len() != before { + meta_set_ids(storage, key, &ids).await?; + } + Ok(()) +} + +// ── BinaryGraph ── + +/// Binary graph — dataset riêng cho symbol binary trên trait [`Storage`]. +pub struct BinaryGraph { + storage: SharedStorage, + /// Name trie (substring/prefix search) — record index riêng từ + /// [`RECORD_START`], meta của record = danh sách symbol id mang tên đó + /// (tên trùng nhiều symbol / nhiều binary). + names: Arc>>, + bin_base: u64, +} + +impl BinaryGraph { + /// Mở (hoặc tạo) binary graph. `dsn` dạng `sqlite://`, `lmdb://`, + /// `redis://`; `None` → in-memory. Open là O(1): chỉ mở storage + + /// Search (trie persist trong storage) — KHÔNG load symbol nào vào RAM. + pub async fn open(dsn: Option<&str>, bin_base: u64) -> Result { + let storage: SharedStorage = match dsn { + Some(dsn) => open_keyspace_storage(dsn, "codegraph:binary").await?, + None => Arc::new(RwLock::new(codegraph_graph::InMemoryStorage::default())), + }; + Ok(Self { + names: Arc::new(RwLock::new(Search::new(BIN_SHARDING, storage.clone()))), + storage, + bin_base, + }) + } + + /// Mở theo config `[bingraph]` — dùng chung cho CLI và MCP. + /// Backend không khai báo dsn → in-memory + warn. + pub async fn open_from_config(root: &Utf8Path) -> Result { + let cfg = ExtractConfig::load(root); + if !cfg.bingraph.is_enabled() { + return Err(Error::Db( + "[bingraph] bị tắt trong .codegraph/config.toml".to_string(), + )); + } + let dsn = cfg.bingraph_dsn(root); + if dsn.is_none() { + tracing::warn!( + "[bingraph] không có DSN hợp lệ — dùng in-memory \ + (override bằng [bingraph.storage] dsn)" + ); + } + Self::open(dsn.as_deref(), cfg.bin_base()).await + } + + /// Base id đang dùng cho symbol binary graph. + pub fn bin_base(&self) -> u64 { + self.bin_base + } + + // ------------------------------------------------------------------ + // Ingest + // ------------------------------------------------------------------ + + /// Ingest một `ParseResult` binary (language = "binary"): remap id sang dải + /// `bin_base`, lưu symbols + secondary index + name trie + chains + calls. + /// Idempotent per path — index/symbol của path cũ bị gỡ trước khi ghi. + pub async fn ingest(&self, parsed: &ParseResult, bin_base: u64) -> Result<()> { + // 1. Gỡ index của path cũ (re-index thay thế). + let path_key = format!("path:{}", parsed.path); + let old_ids = meta_ids(&self.storage, &path_key).await?; + for old in &old_ids { + if let Some(sym) = self.load_symbol(*old).await? { + meta_remove_id(&self.storage, "all", *old).await?; + meta_remove_id(&self.storage, &format!("kind:{:?}", sym.kind), *old).await?; + for f in symbol_flags(&sym) { + meta_remove_id(&self.storage, &format!("flag:{}", f.as_str()), *old).await?; + } + meta_remove_id(&self.storage, &format!("addr:{}", sym.line), *old).await?; + if symbol_flags(&sym).contains(&BinFlag::Entrypoint) { + meta_remove_id(&self.storage, "ep", *old).await?; + } + self.name_index_remove(&sym.name, *old).await?; + } + } + + // 2. Lưu symbol + secondary index mới. + let mut new_ids = Vec::with_capacity(parsed.symbols.len()); + for sym in &parsed.symbols { + let mut stored = sym.clone(); + stored.id = bin_base + sym.id; + new_ids.push(stored.id); + self.storage + .write() + .await + .save_symbol(&stored) + .await + .map_err(db_err)?; + meta_add_id(&self.storage, "all", stored.id).await?; + meta_add_id(&self.storage, &format!("kind:{:?}", sym.kind), stored.id).await?; + for f in symbol_flags(sym) { + meta_add_id(&self.storage, &format!("flag:{}", f.as_str()), stored.id).await?; + if f == BinFlag::Entrypoint { + meta_add_id(&self.storage, "ep", stored.id).await?; + } + } + meta_add_id(&self.storage, &format!("addr:{}", sym.line), stored.id).await?; + } + meta_add_id(&self.storage, "paths", kv_record(&path_key) as u64).await?; + meta_set_ids(&self.storage, &path_key, &new_ids).await?; + + // 3. Name trie: mỗi tên distinct một record; meta record = ids. + let mut names_map: HashMap<&str, Vec> = HashMap::new(); + for sym in &parsed.symbols { + names_map + .entry(sym.name.as_str()) + .or_default() + .push(bin_base + sym.id); + } + let mut next_record: usize = { + let ids = meta_ids(&self.storage, "next_record").await?; + ids.first().copied().unwrap_or(RECORD_START as u64) as usize + }; + for (name, ids) in &names_map { + let existing = self.name_record_lookup(name).await?; + match existing { + Some(record) => { + let mut current = meta_ids_at(&self.storage, record).await?; + for id in ids { + if !current.contains(id) { + current.push(*id); + } + } + meta_set_ids_at(&self.storage, record, ¤t).await?; + } + None => { + let metas: Vec> = vec![None; name.len()]; + self.names + .write() + .await + .insert_chain(next_record, name.as_bytes(), &metas) + .await + .map_err(|e| match e { + SearchError::Duplicated => { + Error::Db("name trie duplicated".to_string()) + } + other => Error::Db(format!("name trie insert: {other}")), + })?; + meta_set_ids_at(&self.storage, next_record, ids).await?; + next_record += 1; + } + } + } + meta_set_ids(&self.storage, "next_record", &[next_record as u64]).await?; + + // 4. Chains (u64 native) + call records (JSON). + for (local_id, chain) in &parsed.chains { + let global: Vec = chain.iter().map(|v| bin_base + v).collect(); + self.storage + .write() + .await + .set_chain((bin_base + local_id) as usize, &global) + .await + .map_err(db_err)?; + } + for call in &parsed.calls { + let mut recs = self + .storage + .read() + .await + .get_call_records(bin_base + call.caller_id) + .await + .map_err(db_err)? + .and_then(|b| serde_json::from_slice::>(&b).ok()) + .unwrap_or_default(); + let mut rec = call.clone(); + rec.caller_id = bin_base + call.caller_id; + recs.push(rec); + let blob = serde_json::to_vec(&recs).map_err(|e| Error::Db(e.to_string()))?; + self.storage + .write() + .await + .set_call_records(bin_base + call.caller_id, &blob) + .await + .map_err(db_err)?; + } + Ok(()) + } + + /// Tìm record của tên trong trie (exact match trên key). + async fn name_record_lookup(&self, name: &str) -> Result> { + let trie = self.names.read().await; + let hits = trie + .search_prefix(name.as_bytes()) + .await + .map_err(|e| Error::Db(format!("name trie lookup: {e}")))?; + Ok(hits + .into_iter() + .find(|(key, _)| key.as_slice() == name.as_bytes()) + .map(|(_, record)| record)) + } + + /// Gỡ một symbol id khỏi meta của name record. + async fn name_index_remove(&self, name: &str, id: u64) -> Result<()> { + if let Some(record) = self.name_record_lookup(name).await? { + let mut ids = meta_ids_at(&self.storage, record).await?; + ids.retain(|&x| x != id); + meta_set_ids_at(&self.storage, record, &ids).await?; + } + Ok(()) + } + + async fn load_symbol(&self, id: u64) -> Result> { + self.storage + .read() + .await + .load_symbol(id) + .await + .map_err(db_err) + } + + /// Load symbols theo danh sách id, lọc kind/flag/path, sort theo order, + /// phân trang. + async fn load_page( + &self, + ids: &[u64], + kind: Option<&SymbolKind>, + flag: Option<&BinFlag>, + path: Option<&str>, + order: ListOrder, + offset: u64, + limit: u64, + ) -> Result { + let mut rows: Vec = Vec::new(); + for id in ids { + if let Some(sym) = self.load_symbol(*id).await? { + if let Some(k) = kind { + if sym.kind != *k { + continue; + } + } + if let Some(f) = flag { + if !symbol_flags(&sym).contains(f) { + continue; + } + } + if let Some(p) = path { + if sym.file != p { + continue; + } + } + rows.push(sym.into()); + } + } + match order { + ListOrder::Name => rows.sort_by(|a, b| a.name.cmp(&b.name).then(a.id.cmp(&b.id))), + ListOrder::Addr => rows.sort_by(|a, b| a.addr.cmp(&b.addr).then(a.id.cmp(&b.id))), + ListOrder::Id => rows.sort_by_key(|r| r.id), + } + let total = rows.len() as u64; + let rows = rows + .into_iter() + .skip(offset as usize) + .take(limit as usize) + .collect(); + Ok(BinPage { + rows, + total, + offset, + }) + } + + // ------------------------------------------------------------------ + // Lazy queries + // ------------------------------------------------------------------ + + /// List symbol theo kind/flag/path — phân trang, không load hết vào RAM + /// trừ khi không có filter nào (ids từ secondary index). + pub async fn list( + &self, + kind: Option, + flag: Option, + path: Option<&str>, + order: ListOrder, + offset: u64, + limit: u64, + ) -> Result { + // Chọn index đơn hẹp nhất có sẵn, phần còn lại lọc khi load symbol. + let ids = if let Some(k) = &kind { + meta_ids(&self.storage, &format!("kind:{k:?}")).await? + } else if let Some(f) = &flag { + meta_ids(&self.storage, &format!("flag:{}", f.as_str())).await? + } else if let Some(p) = path { + meta_ids(&self.storage, &format!("path:{p}")).await? + } else { + meta_ids(&self.storage, "all").await? + }; + self.load_page( + &ids, + kind.as_ref(), + flag.as_ref(), + path, + order, + offset, + limit, + ) + .await + } + + /// Search theo tên + mode. Contains/suffix đi qua name trie (substring); + /// exact/prefix đi qua prefix lookup của trie — mọi mode đều không scan. + pub async fn search_name( + &self, + pattern: &str, + mode: NameMatch, + kind: Option, + flag: Option, + offset: u64, + limit: u64, + ) -> Result { + if pattern.is_empty() { + return Ok(BinPage { + rows: Vec::new(), + total: 0, + offset, + }); + } + let mut ids: Vec = Vec::new(); + match mode { + NameMatch::Contains | NameMatch::Suffix => { + let page = self + .names + .read() + .await + .search_resumable(pattern.as_bytes(), None, None, None) + .await + .map_err(|e| Error::Db(format!("name trie search: {e}")))?; + for record in page.record_ids { + ids.extend(meta_ids_at(&self.storage, record).await?); + } + } + NameMatch::Exact | NameMatch::Prefix => { + let hits = self + .names + .read() + .await + .search_prefix(pattern.as_bytes()) + .await + .map_err(|e| Error::Db(format!("name trie prefix: {e}")))?; + for (key, record) in hits { + if mode == NameMatch::Exact && key.as_slice() != pattern.as_bytes() { + continue; + } + ids.extend(meta_ids_at(&self.storage, record).await?); + } + } + } + self.load_page( + &ids, + kind.as_ref(), + flag.as_ref(), + None, + ListOrder::Name, + offset, + limit, + ) + .await + } + + /// Tra cứu theo địa chỉ (secondary index `addr:{a}`) — điểm bắt đầu + /// phân tích binary. + pub async fn by_addr(&self, addr: u64, limit: u64) -> Result> { + let ids = meta_ids(&self.storage, &format!("addr:{addr}")).await?; + let page = self + .load_page(&ids, None, None, None, ListOrder::Name, 0, limit) + .await?; + Ok(page.rows) + } + + /// Danh sách entry point (toàn bộ hoặc lọc theo binary path) — điểm bắt + /// đầu phân tích thay cho grep với code. + pub async fn entrypoints( + &self, + path: Option<&str>, + limit: u64, + ) -> Result> { + let ids = meta_ids(&self.storage, "ep").await?; + let mut eps: Vec<(u64, String, String)> = Vec::new(); + for id in ids { + if let Some(sym) = self.load_symbol(id).await? { + if let Some(p) = path { + if sym.file != p { + continue; + } + } + eps.push((u64::from(sym.line), sym.file, sym.name)); + } + } + eps.sort(); + Ok(eps + .into_iter() + .take(limit as usize) + .map(|(_, path, name)| (path, name)) + .collect()) + } + + /// Lấy symbol đầy đủ theo id — lazy hydrate. + pub async fn get_symbol(&self, id: u64) -> Result> { + self.load_symbol(id).await + } + + /// Chain (flow) của một symbol id — native u64 trên Storage. + pub async fn get_chain(&self, id: u64) -> Result>> { + self.storage + .read() + .await + .get_chain(id as usize) + .await + .map_err(db_err) + } + + /// Call records của một caller id. + pub async fn get_calls( + &self, + caller: u64, + ) -> Result, Option)>> { + let blob = self + .storage + .read() + .await + .get_call_records(caller) + .await + .map_err(db_err)?; + let recs: Vec = blob + .and_then(|b| serde_json::from_slice(&b).ok()) + .unwrap_or_default(); + Ok(recs + .into_iter() + .map(|c| (c.position as i64, Some(c.call_name), c.condition)) + .collect()) + } + + /// Thống kê — đếm từ secondary index (chỉ đọc danh sách id, không load + /// symbol). + pub async fn stats(&self) -> Result { + Ok(BinStats { + symbols: meta_ids(&self.storage, "all").await?.len() as u64, + entrypoints: meta_ids(&self.storage, "ep").await?.len() as u64, + imports: meta_ids(&self.storage, "flag:import").await?.len() as u64, + exports: meta_ids(&self.storage, "flag:export").await?.len() as u64, + binaries: meta_ids(&self.storage, "paths").await?.len() as u64, + }) + } +} + +fn db_err(e: StorageError) -> Error { + Error::Db(format!("binary graph: {e}")) +} + +#[cfg(test)] +mod tests { + use super::*; + use codegraph_core::{Annotation, EffectType}; + use std::collections::HashMap; + + fn sym( + id: u64, + name: &str, + kind: SymbolKind, + line: u32, + annotations: Vec, + ) -> Symbol { + Symbol { + id, + name: name.to_string(), + kind, + scope: codegraph_core::ScopeLevel::Global, + scope_id: 0, + type_ref: 0, + type_name: None, + file: "/tmp/fake.so".to_string(), + line, + end_line: line, + signature: None, + doc: None, + annotations, + language: "binary".to_string(), + } + } + + fn ann(name: &str) -> Annotation { + Annotation { + name: name.to_string(), + args: HashMap::new(), + line: 0, + } + } + + fn sample() -> ParseResult { + ParseResult { + path: "/tmp/fake.so".to_string(), + language: "binary".to_string(), + bytes: 0, + lines: 0, + symbols: vec![ + sym( + 1, + "entry0", + SymbolKind::Function, + 4096, + vec![ann("entrypoint")], + ), + sym(2, "foo", SymbolKind::Function, 4200, vec![ann("export")]), + sym(3, "memcpy", SymbolKind::Function, 100, vec![ann("import")]), + sym(4, "local_fn", SymbolKind::Function, 5000, Vec::new()), + sym(5, "str:6000", SymbolKind::Constant, 6000, Vec::new()), + ], + chains: HashMap::from([(1u64, vec![1u64, 3u64])]), + calls: vec![CallRecord { + caller_id: 1, + call_name: "memcpy".to_string(), + position: 1, + arg_exprs: Vec::new(), + line: 4100, + condition: None, + is_loop_body: false, + effect: EffectType::None, + effect_desc: None, + target_class: None, + target_method: None, + }], + } + } + + async fn mem_graph() -> BinaryGraph { + BinaryGraph::open(None, DEFAULT_BIN_BASE).await.unwrap() + } + + #[tokio::test] + async fn ingest_and_lazy_queries() { + let g = mem_graph().await; + g.ingest(&sample(), DEFAULT_BIN_BASE).await.unwrap(); + + // list theo flag — import/export/entrypoint. + let page = g + .list(None, Some(BinFlag::Export), None, ListOrder::Name, 0, 50) + .await + .unwrap(); + assert_eq!(page.rows.len(), 1); + assert_eq!(page.rows[0].name, "foo"); + + let eps = g + .list( + None, + Some(BinFlag::Entrypoint), + None, + ListOrder::Name, + 0, + 50, + ) + .await + .unwrap(); + assert_eq!(eps.rows.len(), 1); + assert_eq!(eps.rows[0].name, "entry0"); + + // list theo kind — Constant chỉ có str:6000. + let page = g + .list( + Some(SymbolKind::Constant), + None, + None, + ListOrder::Name, + 0, + 50, + ) + .await + .unwrap(); + assert_eq!(page.rows.len(), 1); + assert_eq!(page.rows[0].name, "str:6000"); + + // by_addr. + let rows = g.by_addr(4200, 10).await.unwrap(); + assert_eq!(rows.len(), 1); + assert_eq!(rows[0].name, "foo"); + + // entrypoints listing. + let eps = g.entrypoints(Some("/tmp/fake.so"), 10).await.unwrap(); + assert_eq!(eps.len(), 1); + assert_eq!(eps[0].1, "entry0"); + + // get_symbol lazy hydrate + chain (u64 native trên Storage). + let s = g.get_symbol(DEFAULT_BIN_BASE + 2).await.unwrap().unwrap(); + assert_eq!(s.name, "foo"); + let chain = g.get_chain(DEFAULT_BIN_BASE + 1).await.unwrap().unwrap(); + assert_eq!(chain, vec![DEFAULT_BIN_BASE + 1, DEFAULT_BIN_BASE + 3]); + let calls = g.get_calls(DEFAULT_BIN_BASE + 1).await.unwrap(); + assert_eq!(calls.len(), 1); + assert_eq!(calls[0].1.as_deref(), Some("memcpy")); + + // stats. + let stats = g.stats().await.unwrap(); + assert_eq!(stats.symbols, 5); + assert_eq!(stats.entrypoints, 1); + assert_eq!(stats.imports, 1); + assert_eq!(stats.exports, 1); + assert_eq!(stats.binaries, 1); + } + + #[tokio::test] + async fn reingest_replaces_path() { + let g = mem_graph().await; + g.ingest(&sample(), DEFAULT_BIN_BASE).await.unwrap(); + let mut updated = sample(); + updated.symbols = vec![sym(9, "only_one", SymbolKind::Function, 1, Vec::new())]; + g.ingest(&updated, DEFAULT_BIN_BASE).await.unwrap(); + let stats = g.stats().await.unwrap(); + assert_eq!(stats.symbols, 1, "re-ingest cùng path phải thay thế index"); + assert_eq!(stats.binaries, 1, "path cũ vẫn là 1 binary"); + } + + #[tokio::test] + async fn pagination() { + let g = mem_graph().await; + g.ingest(&sample(), DEFAULT_BIN_BASE).await.unwrap(); + let page = g + .list(None, None, None, ListOrder::Name, 0, 2) + .await + .unwrap(); + assert_eq!(page.total, 5); + assert_eq!(page.rows.len(), 2); + let page2 = g + .list(None, None, None, ListOrder::Name, 2, 2) + .await + .unwrap(); + assert_eq!(page2.rows.len(), 2); + assert_ne!(page.rows[0].id, page2.rows[0].id); + } + + #[tokio::test] + async fn contains_and_suffix_via_trie() { + let g = mem_graph().await; + g.ingest(&sample(), DEFAULT_BIN_BASE).await.unwrap(); + // contains "cpy" khớp "memcpy" qua trie (substring DFS). + let page = g + .search_name("cpy", NameMatch::Contains, None, None, 0, 50) + .await + .unwrap(); + assert_eq!(page.total, 1); + assert_eq!(page.rows[0].name, "memcpy"); + // suffix "oo" trả foo. + let page = g + .search_name("oo", NameMatch::Suffix, None, None, 0, 50) + .await + .unwrap(); + assert_eq!(page.rows.len(), 1); + assert_eq!(page.rows[0].name, "foo"); + // prefix qua trie. + let page = g + .search_name("mem", NameMatch::Prefix, None, None, 0, 50) + .await + .unwrap(); + assert_eq!(page.rows.len(), 1); + // contains không khớp gì → trang rỗng. + let page = g + .search_name("zzz", NameMatch::Contains, None, None, 0, 50) + .await + .unwrap(); + assert_eq!(page.total, 0); + } + + #[tokio::test] + async fn reingest_does_not_return_stale_names() { + let g = mem_graph().await; + g.ingest(&sample(), DEFAULT_BIN_BASE).await.unwrap(); + let mut updated = sample(); + updated.symbols = vec![sym(9, "only_one", SymbolKind::Function, 1, Vec::new())]; + g.ingest(&updated, DEFAULT_BIN_BASE).await.unwrap(); + // "memcpy" đã bị gỡ khỏi index của path cũ → contains không trả row. + let page = g + .search_name("cpy", NameMatch::Contains, None, None, 0, 50) + .await + .unwrap(); + assert_eq!(page.total, 0); + let page = g + .search_name("only", NameMatch::Contains, None, None, 0, 50) + .await + .unwrap(); + assert_eq!(page.total, 1); + } + + #[tokio::test] + async fn duplicate_name_across_binaries() { + let g = mem_graph().await; + g.ingest(&sample(), DEFAULT_BIN_BASE).await.unwrap(); + // Binary thứ 2 cũng có symbol "memcpy" — cùng tên, khác path. + let mut other = sample(); + other.path = "/tmp/other.so".to_string(); + let mut sym_other = sym(7, "memcpy", SymbolKind::Function, 200, vec![ann("import")]); + sym_other.file = other.path.clone(); + other.symbols = vec![sym_other]; + g.ingest(&other, DEFAULT_BIN_BASE).await.unwrap(); + // contains vẫn chỉ 1 record tên "memcpy" nhưng trả 2 symbol id. + let page = g + .search_name("memcpy", NameMatch::Exact, None, None, 0, 50) + .await + .unwrap(); + assert_eq!(page.total, 2, "2 binary cùng tên → 2 row"); + assert!(page.rows.iter().any(|r| r.path == "/tmp/other.so")); + } + + #[tokio::test] + async fn persists_on_sqlite_backend() { + // Backend sqlite qua trait Storage — persist qua các lần open, dataset + // riêng (binary.sqlite) không đụng db.sqlite/docs.sqlite. + let dir = tempfile::tempdir().unwrap(); + let dsn = format!( + "sqlite://{}", + dir.path().join("binary.sqlite").to_str().unwrap() + ); + let g = BinaryGraph::open(Some(&dsn), DEFAULT_BIN_BASE) + .await + .unwrap(); + g.ingest(&sample(), DEFAULT_BIN_BASE).await.unwrap(); + assert!(dir.path().join("binary.sqlite").exists()); + let g2 = BinaryGraph::open(Some(&dsn), DEFAULT_BIN_BASE) + .await + .unwrap(); + let stats = g2.stats().await.unwrap(); + assert_eq!(stats.symbols, 5, "persist qua các lần open"); + let page = g2 + .search_name("cpy", NameMatch::Contains, None, None, 0, 50) + .await + .unwrap(); + assert_eq!(page.total, 1, "name trie persist trên storage"); + } +} diff --git a/crates/codegraph-extract/src/config.rs b/crates/codegraph-extract/src/config.rs index 98380699f..7874e88e1 100644 --- a/crates/codegraph-extract/src/config.rs +++ b/crates/codegraph-extract/src/config.rs @@ -71,6 +71,10 @@ struct ConfigFile { /// Document graph — ingest tài liệu cấu trúc lúc `codegraph init`. #[serde(default)] docgraph: DocGraphSection, + /// Binary graph — dataset riêng cho symbol binary (`[bingraph]`). + #[cfg(feature = "binary")] + #[serde(default)] + bingraph: BinGraphSection, /// Phân tích binary (radare2) — feature `binary`. #[cfg(feature = "binary")] @@ -141,6 +145,30 @@ pub struct DocGraphStorageSection { pub dsn: Option, } +/// Section `[bingraph]` — cấu hình binary graph: dataset riêng (mặc định +/// `.codegraph/binary.sqlite`) cho symbol binary, tách khỏi code index và +/// docs để query search/list chạy lazy trên SQL index không phải rebuild RAM. +#[derive(Debug, Clone, Default, Deserialize)] +pub struct BinGraphSection { + /// Bật binary graph (mặc định bật). + #[serde(default)] + enabled: Option, + /// Override storage — hiện chỉ hỗ trợ sqlite; backend khác → in-memory + warn. + #[serde(default)] + storage: Option, + /// Base id cho symbol binary graph (mặc định 2e9 — không đụng dải docs + /// 1e9/3e9 và dải code index). + #[serde(default)] + bin_base: Option, +} + +impl BinGraphSection { + /// Binary graph có bật hay không (mặc định bật). + pub fn is_enabled(&self) -> bool { + self.enabled.unwrap_or(true) + } +} + impl DocGraphSection { /// Ingest docs có bật hay không: `enabled` override, mặc định = có `paths`. pub fn is_enabled(&self) -> bool { @@ -193,6 +221,9 @@ pub struct ExtractConfig { pub embedding: codegraph_graph::embeddings::EmbeddingConfig, /// Cấu hình document graph — đọc từ `[docgraph]`. pub docgraph: DocGraphSection, + /// Cấu hình binary graph — đọc từ `[bingraph]`. + #[cfg(feature = "binary")] + pub bingraph: BinGraphSection, /// Cấu hình phân tích binary (radare2). #[cfg(feature = "binary")] pub binary: BinaryConfig, @@ -254,6 +285,8 @@ impl ExtractConfig { }, docgraph: file.docgraph, #[cfg(feature = "binary")] + bingraph: file.bingraph, + #[cfg(feature = "binary")] binary: file.binary.unwrap_or_default(), } } @@ -385,6 +418,48 @@ impl ExtractConfig { } } + /// DSN dataset **riêng** cho binary graph (`[bingraph]`) — dataset chạy trên + /// trait `Storage` nên hỗ trợ mọi backend local/remote: + /// - `[bingraph.storage] dsn` override → dùng nguyên văn. + /// - Mặc định theo backend kind (override được bằng `[bingraph.storage] type`): + /// - sqlite → `sqlite:///.codegraph/binary.sqlite` + /// - lmdb → `lmdb:///.codegraph/binary.lmdb` + /// - redis → DSN của `[storage]` (keyspace `codegraph:binary`) + /// - memory / RDBMS → `None` (in-memory + warn ở caller) + #[cfg(feature = "binary")] + pub fn bingraph_dsn(&self, root: &Utf8Path) -> Option { + if let Some(dsn) = self + .bingraph + .storage + .as_ref() + .and_then(|s| s.dsn.as_deref()) + { + return Some(dsn.to_string()); + } + let kind = self + .bingraph + .storage + .as_ref() + .and_then(|s| s.type_.as_deref()) + .map(StorageKind::parse) + .unwrap_or(self.storage.kind); + match kind { + StorageKind::Sqlite => Some(format!( + "sqlite://{}", + project_dir(root).join("binary.sqlite") + )), + StorageKind::Lmdb => Some(format!("lmdb://{}", project_dir(root).join("binary.lmdb"))), + StorageKind::Redis => self.storage.dsn.clone(), + _ => None, + } + } + + /// Base id cho symbol binary graph (mặc định 2e9). + #[cfg(feature = "binary")] + pub fn bin_base(&self) -> u64 { + self.bingraph.bin_base.unwrap_or(2_000_000_000) + } + /// Config document graph + danh sách file khớp glob `[docgraph] paths` /// (path kèm format override). Trả `None` khi `[docgraph]` không bật / /// không khai báo `paths`. diff --git a/crates/codegraph-extract/src/lib.rs b/crates/codegraph-extract/src/lib.rs index 33095513a..cb7dff45e 100644 --- a/crates/codegraph-extract/src/lib.rs +++ b/crates/codegraph-extract/src/lib.rs @@ -13,6 +13,12 @@ mod orchestrator; mod project; mod walker; +/// Binary graph — dataset riêng cho symbol binary (feature `binary`). +#[cfg(feature = "binary")] +pub mod bingraph; + +#[cfg(feature = "binary")] +pub use bingraph::{BinFlag, BinPage, BinSymbolRow, BinaryGraph, ListOrder, NameMatch}; pub use config::{DocGraphSection, ExtractConfig, HeaderLanguage, DEFAULT_CONFIG_TOML}; pub use docgraph::open_doc_graph; pub use orchestrator::{ExtractStats, Orchestrator}; diff --git a/crates/codegraph-extract/src/orchestrator.rs b/crates/codegraph-extract/src/orchestrator.rs index 1c4015901..2915ad8a4 100644 --- a/crates/codegraph-extract/src/orchestrator.rs +++ b/crates/codegraph-extract/src/orchestrator.rs @@ -84,11 +84,38 @@ impl Orchestrator { let (mut parsed, mut skipped) = self.parse_files(&files, progress.clone(), config.effect_classifier.clone()); + // Binary đi dataset riêng (`[bingraph]` → binary.sqlite) — KHÔNG ingest + // vào code index nữa (tránh n_symbol binary làm phình trie/RAM của + // GraphIndex). `[bingraph]` tắt → giữ hành vi cũ (đẩy vào code index). + #[cfg(feature = "binary")] + let mut bin_stats: Vec = Vec::new(); #[cfg(feature = "binary")] { let (bin, bin_skipped) = codegraph_binary::collect_binaries(root, &config.binary); - parsed.extend(bin); skipped += bin_skipped; + if config.bingraph.is_enabled() && !bin.is_empty() { + let bin_base = config.bin_base(); + let dsn = config.bingraph_dsn(root); + if dsn.is_none() { + tracing::warn!("[bingraph] backend không phải sqlite — fallback in-memory"); + } + match crate::bingraph::BinaryGraph::open(dsn.as_deref(), bin_base).await { + Ok(bg) => { + for r in &bin { + if let Err(e) = bg.ingest(r, bin_base).await { + tracing::warn!("binary ingest {} thất bại: {e}", r.path); + } + } + bin_stats = bin; + } + Err(e) => { + tracing::warn!("mở binary graph thất bại: {e} — fallback code index"); + parsed.extend(bin); + } + } + } else { + parsed.extend(bin); + } } // Đưa ProgressBar vào ingest (register → edges → files → engines) — phase @@ -102,7 +129,11 @@ impl Orchestrator { if let Some(bar) = progress { bar.finish_with_message("Indexing complete"); } - Ok(stats_of(&parsed, skipped)) + #[allow(unused_mut)] + let mut all_for_stats = parsed; + #[cfg(feature = "binary")] + all_for_stats.extend(bin_stats); + Ok(stats_of(&all_for_stats, skipped)) } /// Parse song song một danh sách file — trả về parsed + số file bị skip. diff --git a/crates/codegraph-graph/src/lib.rs b/crates/codegraph-graph/src/lib.rs index ae2653ef5..d4d04d075 100644 --- a/crates/codegraph-graph/src/lib.rs +++ b/crates/codegraph-graph/src/lib.rs @@ -50,7 +50,7 @@ pub use crate::storage::mysql::MySqlStorage; pub use crate::storage::postgres::PostgresStorage; #[cfg(feature = "sqlite")] pub use crate::storage::sqlite::SqliteStorage; -pub use crate::storage::{InMemoryStorage, IndexCounts, Storage, Tx}; +pub use crate::storage::{InMemoryStorage, IndexCounts, Storage, StorageError, Tx}; // Sub-traits of `Storage` — callers that need only one facet (e.g. a chain-engine // read path) can name it directly instead of taking the full umbrella. #[cfg(feature = "bloom-search")] @@ -131,6 +131,13 @@ fn serr_search(e: crate::search::Error) -> Error { /// - `redis://...` → `RedisStorage` với keyspace prefix `codegraph:docs` /// (feature `redis`) — tách khỏi index `codegraph:idx:` pub async fn open_doc_storage(dsn: &str) -> Result>> { + open_keyspace_storage(dsn, "codegraph:docs").await +} + +/// Mở storage cho một dataset theo DSN + keyspace (backend redis dùng prefix +/// này để tách dữ liệu; sqlite/lmdb tách bằng file riêng nên bỏ qua keyspace). +/// Dùng chung cho docs (`codegraph:docs`) và binary graph (`codegraph:binary`). +pub async fn open_keyspace_storage(dsn: &str, keyspace: &str) -> Result>> { if let Some(path) = dsn.strip_prefix("sqlite://") { #[cfg(feature = "sqlite")] { @@ -164,18 +171,19 @@ pub async fn open_doc_storage(dsn: &str) -> Result>> { { let client = redis::Client::open(dsn).map_err(|e| Error::Db(format!("redis client: {e}")))?; - let storage = crate::storage::redis::RedisStorage::new(client, "codegraph:docs") + let storage = crate::storage::redis::RedisStorage::new(client, keyspace) .await .map_err(serr)?; return Ok(Arc::new(RwLock::new(storage))); } #[cfg(not(feature = "redis"))] { + let _ = keyspace; return Err(backend_unavailable("redis")); } } Err(Error::Db(format!( - "open_doc_storage: DSN scheme không hỗ trợ: {dsn}" + "open_keyspace_storage: DSN scheme không hỗ trợ: {dsn}" ))) } diff --git a/crates/codegraph-mcp/src/lib.rs b/crates/codegraph-mcp/src/lib.rs index cb8b74050..3cc57766b 100644 --- a/crates/codegraph-mcp/src/lib.rs +++ b/crates/codegraph-mcp/src/lib.rs @@ -314,6 +314,18 @@ impl CodegraphServer { }; } + // Binary tools — dataset riêng, lazy; mở per-call (open là O(1), + // search contains đi radix trie persist). + if name.starts_with("codegraph_binary_") { + return match tools::dispatch_binary(&root, name, args).await { + Ok(text) => Ok(ToolOutput::Text { + text, + source_bytes: 0, + }), + Err(e) => Ok(ToolOutput::Error(e.to_string())), + }; + } + let dispatch = match name { "codegraph_sandbox" => { codegraph_api::tools::dispatch_sandbox(&root, sgi.clone(), args.clone()).await diff --git a/crates/codegraph-mcp/src/tools.rs b/crates/codegraph-mcp/src/tools.rs index 2fcd41f4a..c2c5fcd8e 100644 --- a/crates/codegraph-mcp/src/tools.rs +++ b/crates/codegraph-mcp/src/tools.rs @@ -312,6 +312,45 @@ fn tool_defs() -> Vec { "Show document graph statistics (number of documents and nodes).", json!({ "type": "object", "properties": {} }), ), + // ── Binary tools (dataset riêng .codegraph/binary.sqlite — lazy SQL) ── + tool( + "codegraph_binary_list", + "List binary symbols from the separate binary graph (entrypoints/exports/imports/functions/strings). Fast SQL-indexed listing with pagination — the starting point for binary analysis (entrypoints replace grep as the anchor).", + json!({ "type": "object", "properties": { + "flag": { "type": "string", "enum": ["entrypoint", "export", "import", "jni"], "description": "Filter by flag. Omit to list all symbols." }, + "kind": { "type": "string", "description": "Filter by symbol kind (Function, Method, Class, Module, Enum, Constant)." }, + "path": { "type": "string", "description": "Filter by binary file path." }, + "order": { "type": "string", "enum": ["name", "addr", "id"], "default": "name" }, + "offset": { "type": "integer", "default": 0 }, + "limit": { "type": "integer", "default": 50, "description": "Max rows per page." } + } }), + ), + tool( + "codegraph_binary_search", + "Search binary symbols by name (exact/prefix/suffix/contains), optionally filtered by kind/flag. Backed by SQL indexes on the separate binary dataset — no in-memory rebuild.", + json!({ "type": "object", "properties": { + "pattern": { "type": "string", "description": "Name pattern to search." }, + "match": { "type": "string", "enum": ["exact", "prefix", "suffix", "contains"], "default": "contains" }, + "kind": { "type": "string", "description": "Optional kind filter (Function, Method, Class, Module, Enum, Constant)." }, + "flag": { "type": "string", "enum": ["entrypoint", "export", "import", "jni"], "description": "Optional flag filter." }, + "offset": { "type": "integer", "default": 0 }, + "limit": { "type": "integer", "default": 50 } + }, "required": ["pattern"] }), + ), + tool( + "codegraph_binary_addr", + "Look up binary symbols at an address (O(1) point query) and list known entrypoints of a binary. Use to anchor binary analysis at entry addresses.", + json!({ "type": "object", "properties": { + "addr": { "type": "integer", "description": "Virtual address to look up (omit to list entrypoints)." }, + "path": { "type": "string", "description": "Binary path for entrypoint listing." }, + "limit": { "type": "integer", "default": 20 } + } }), + ), + tool( + "codegraph_binary_stats", + "Show binary graph statistics (symbols, entrypoints, imports, exports, binaries).", + json!({ "type": "object", "properties": {} }), + ), ] } @@ -1091,3 +1130,125 @@ pub async fn dispatch_doc_stats(doc_graph: Arc>) -> R let stats = doc_graph.read().await.stats(); Ok(format!("documents: {}\nnodes: {}", stats.docs, stats.nodes)) } + +// ── Binary tool dispatch ── +// Dataset riêng `.codegraph/binary.sqlite` — query lazy trên SQL index, không +// đụng GraphIndex (code search). Sync rusqlite: open O(1) + query có LIMIT. + +fn parse_bin_kind(s: &str) -> Option { + match s.to_ascii_lowercase().as_str() { + "function" => Some(SymbolKind::Function), + "method" => Some(SymbolKind::Method), + "class" => Some(SymbolKind::Class), + "interface" => Some(SymbolKind::Interface), + "enum" => Some(SymbolKind::Enum), + "variable" => Some(SymbolKind::Variable), + "constant" => Some(SymbolKind::Constant), + "parameter" => Some(SymbolKind::Parameter), + "field" => Some(SymbolKind::Field), + "module" => Some(SymbolKind::Module), + "file" => Some(SymbolKind::File), + _ => None, + } +} + +pub async fn dispatch_binary(root: &Utf8Path, name: &str, args: Value) -> Result { + let graph = codegraph_extract::BinaryGraph::open_from_config(root) + .await + .map_err(|e| Error::Other(e.to_string()))?; + match name { + "codegraph_binary_list" => { + let kind = args + .get("kind") + .and_then(|v| v.as_str()) + .and_then(parse_bin_kind); + let flag = args + .get("flag") + .and_then(|v| v.as_str()) + .and_then(codegraph_extract::BinFlag::parse); + let path = args.get("path").and_then(|v| v.as_str()); + let order = match args.get("order").and_then(|v| v.as_str()) { + Some("addr") => codegraph_extract::ListOrder::Addr, + Some("id") => codegraph_extract::ListOrder::Id, + _ => codegraph_extract::ListOrder::Name, + }; + let offset = args.get("offset").and_then(|v| v.as_u64()).unwrap_or(0); + let limit = args + .get("limit") + .and_then(|v| v.as_u64()) + .unwrap_or(50) + .min(500); + let page = graph + .list(kind, flag, path, order, offset, limit) + .await + .map_err(|e| Error::Other(e.to_string()))?; + serde_json::to_string_pretty(&page).map_err(|e| Error::Other(e.to_string())) + } + "codegraph_binary_search" => { + let pattern = args + .get("pattern") + .and_then(|v| v.as_str()) + .ok_or_else(|| { + Error::Invalid("codegraph_binary_search requires `pattern`".into()) + })?; + let mode = match args.get("match").and_then(|v| v.as_str()) { + Some("exact") => codegraph_extract::NameMatch::Exact, + Some("prefix") => codegraph_extract::NameMatch::Prefix, + Some("suffix") => codegraph_extract::NameMatch::Suffix, + _ => codegraph_extract::NameMatch::Contains, + }; + let kind = args + .get("kind") + .and_then(|v| v.as_str()) + .and_then(parse_bin_kind); + let flag = args + .get("flag") + .and_then(|v| v.as_str()) + .and_then(codegraph_extract::BinFlag::parse); + let offset = args.get("offset").and_then(|v| v.as_u64()).unwrap_or(0); + let limit = args + .get("limit") + .and_then(|v| v.as_u64()) + .unwrap_or(50) + .min(500); + let page = graph + .search_name(pattern, mode, kind, flag, offset, limit) + .await + .map_err(|e| Error::Other(e.to_string()))?; + serde_json::to_string_pretty(&page).map_err(|e| Error::Other(e.to_string())) + } + "codegraph_binary_addr" => { + let limit = args + .get("limit") + .and_then(|v| v.as_u64()) + .unwrap_or(20) + .min(500); + let path = args.get("path").and_then(|v| v.as_str()); + if let Some(addr) = args.get("addr").and_then(|v| v.as_u64()) { + let rows = graph + .by_addr(addr, limit) + .await + .map_err(|e| Error::Other(e.to_string()))?; + serde_json::to_string_pretty(&rows).map_err(|e| Error::Other(e.to_string())) + } else { + let eps = graph + .entrypoints(path, limit) + .await + .map_err(|e| Error::Other(e.to_string()))?; + let list: Vec<_> = eps + .iter() + .map(|(p, n)| json!({ "path": p, "name": n })) + .collect(); + serde_json::to_string_pretty(&list).map_err(|e| Error::Other(e.to_string())) + } + } + "codegraph_binary_stats" => { + let stats = graph + .stats() + .await + .map_err(|e| Error::Other(e.to_string()))?; + serde_json::to_string_pretty(&stats).map_err(|e| Error::Other(e.to_string())) + } + _ => Err(Error::Invalid(format!("unknown binary tool: {name}"))), + } +} From a479f9e31ad3361540666358d8cf54c98a89a119 Mon Sep 17 00:00:00 2001 From: Hung Pham Date: Sat, 12 Sep 2026 10:39:09 +0700 Subject: [PATCH 2/2] Fix lint --- crates/codegraph-extract/src/bingraph.rs | 58 ++++++++++++------------ 1 file changed, 28 insertions(+), 30 deletions(-) diff --git a/crates/codegraph-extract/src/bingraph.rs b/crates/codegraph-extract/src/bingraph.rs index cf9aeb849..2218dd053 100644 --- a/crates/codegraph-extract/src/bingraph.rs +++ b/crates/codegraph-extract/src/bingraph.rs @@ -141,6 +141,14 @@ pub enum ListOrder { Id, } +/// Filter áp khi load symbol thành row — dùng chung cho `list`/`search_name`. +#[derive(Debug, Clone, Default)] +pub struct PageFilter { + pub kind: Option, + pub flag: Option, + pub path: Option, +} + /// Một trang kết quả list/search. #[derive(Debug, Clone, Serialize)] pub struct BinPage { @@ -437,14 +445,11 @@ impl BinaryGraph { .map_err(db_err) } - /// Load symbols theo danh sách id, lọc kind/flag/path, sort theo order, - /// phân trang. + /// Load symbols theo danh sách id, áp filter + sort + phân trang. async fn load_page( &self, ids: &[u64], - kind: Option<&SymbolKind>, - flag: Option<&BinFlag>, - path: Option<&str>, + filter: &PageFilter, order: ListOrder, offset: u64, limit: u64, @@ -452,18 +457,18 @@ impl BinaryGraph { let mut rows: Vec = Vec::new(); for id in ids { if let Some(sym) = self.load_symbol(*id).await? { - if let Some(k) = kind { + if let Some(k) = &filter.kind { if sym.kind != *k { continue; } } - if let Some(f) = flag { + if let Some(f) = &filter.flag { if !symbol_flags(&sym).contains(f) { continue; } } - if let Some(p) = path { - if sym.file != p { + if let Some(p) = &filter.path { + if sym.file != *p { continue; } } @@ -513,16 +518,12 @@ impl BinaryGraph { } else { meta_ids(&self.storage, "all").await? }; - self.load_page( - &ids, - kind.as_ref(), - flag.as_ref(), - path, - order, - offset, - limit, - ) - .await + let filter = PageFilter { + kind, + flag, + path: path.map(str::to_string), + }; + self.load_page(&ids, &filter, order, offset, limit).await } /// Search theo tên + mode. Contains/suffix đi qua name trie (substring); @@ -573,16 +574,13 @@ impl BinaryGraph { } } } - self.load_page( - &ids, - kind.as_ref(), - flag.as_ref(), - None, - ListOrder::Name, - offset, - limit, - ) - .await + let filter = PageFilter { + kind, + flag, + path: None, + }; + self.load_page(&ids, &filter, ListOrder::Name, offset, limit) + .await } /// Tra cứu theo địa chỉ (secondary index `addr:{a}`) — điểm bắt đầu @@ -590,7 +588,7 @@ impl BinaryGraph { pub async fn by_addr(&self, addr: u64, limit: u64) -> Result> { let ids = meta_ids(&self.storage, &format!("addr:{addr}")).await?; let page = self - .load_page(&ids, None, None, None, ListOrder::Name, 0, limit) + .load_page(&ids, &PageFilter::default(), ListOrder::Name, 0, limit) .await?; Ok(page.rows) }