1use std::io::{BufRead, BufReader, Write};
7use std::path::{Path, PathBuf};
8
9use crate::store::GraphBackend;
10use crate::{schema::EdgeType, Edge, Node};
11
12pub struct JsonBackend {
14 path: PathBuf,
15 nodes: Vec<Node>,
16 edges: Vec<Edge>,
17}
18
19const STREAMING_THRESHOLD: usize = 10_000;
21
22impl JsonBackend {
23 pub fn open<P: AsRef<Path>>(path: P) -> Self {
25 Self {
26 path: path.as_ref().to_path_buf(),
27 nodes: Vec::new(),
28 edges: Vec::new(),
29 }
30 }
31
32 fn flush(&self) -> Result<(), std::io::Error> {
33 if self.nodes.len() + self.edges.len() > STREAMING_THRESHOLD {
34 self.flush_jsonl()?;
35 } else {
36 let doc = JsonDoc {
37 schema: crate::schema::GraphSchema::current(),
38 nodes: &self.nodes,
39 edges: &self.edges,
40 };
41 let mut file = std::fs::File::create(&self.path)?;
42 serde_json::to_writer_pretty(&mut file, &doc)?;
43 file.write_all(b"\n")?;
44 }
45 Ok(())
46 }
47
48 fn flush_jsonl(&self) -> Result<(), std::io::Error> {
49 let nodes_path = self.path.with_extension("nodes.jsonl");
50 let edges_path = self.path.with_extension("edges.jsonl");
51 let mut nf = std::fs::File::create(&nodes_path)?;
52 for n in &self.nodes {
53 serde_json::to_writer(&mut nf, n)?;
54 nf.write_all(b"\n")?;
55 }
56 let mut ef = std::fs::File::create(&edges_path)?;
57 for e in &self.edges {
58 serde_json::to_writer(&mut ef, e)?;
59 ef.write_all(b"\n")?;
60 }
61 let manifest = serde_json::json!({
63 "format": "jsonl",
64 "schema": crate::schema::GraphSchema::current(),
65 "nodes_file": nodes_path,
66 "edges_file": edges_path,
67 "node_count": self.nodes.len(),
68 "edge_count": self.edges.len(),
69 });
70 let mut mf = std::fs::File::create(&self.path)?;
71 serde_json::to_writer_pretty(&mut mf, &manifest)?;
72 mf.write_all(b"\n")?;
73 Ok(())
74 }
75
76 fn load(&mut self) -> Result<(), JsonError> {
77 if !self.path.exists() {
78 return Ok(());
79 }
80 let raw = std::fs::read_to_string(&self.path)?;
90 if raw.trim().is_empty() {
91 return Ok(());
92 }
93
94 let trimmed_full = raw.trim_start();
95 if trimmed_full.starts_with('{') {
96 if let Ok(doc) = serde_json::from_str::<JsonDocOwned>(&raw) {
99 self.nodes = doc.nodes;
100 self.edges = doc.edges;
101 return Ok(());
102 }
103 let val: serde_json::Value = serde_json::from_str(&raw)?;
104 if val.get("format").and_then(|v| v.as_str()) == Some("jsonl") {
105 let parent = self
112 .path
113 .parent()
114 .map(Path::to_path_buf)
115 .unwrap_or_else(|| PathBuf::from("."));
116 let resolve_sibling = |raw: Option<&str>, default: PathBuf| -> PathBuf {
117 let Some(raw) = raw else { return default };
118 let candidate = PathBuf::from(raw);
119 let file_name = candidate.file_name();
120 let stays_in_parent = !candidate.is_absolute()
121 && !candidate
122 .components()
123 .any(|c| matches!(c, std::path::Component::ParentDir));
124 match (file_name, stays_in_parent) {
125 (Some(name), true) => parent.join(name),
126 _ => default,
127 }
128 };
129 let nodes_file = resolve_sibling(
130 val.get("nodes_file").and_then(|v| v.as_str()),
131 self.path.with_extension("nodes.jsonl"),
132 );
133 let edges_file = resolve_sibling(
134 val.get("edges_file").and_then(|v| v.as_str()),
135 self.path.with_extension("edges.jsonl"),
136 );
137 self.nodes = read_jsonl(&nodes_file)?;
138 self.edges = read_jsonl(&edges_file)?;
139 return Ok(());
140 }
141 }
143
144 for line in raw.lines() {
147 if line.trim().is_empty() {
148 continue;
149 }
150 let val: serde_json::Value = serde_json::from_str(line)?;
151 if val.get("source_id").is_some() {
152 self.edges.push(serde_json::from_value(val)?);
153 } else {
154 self.nodes.push(serde_json::from_value(val)?);
155 }
156 }
157 Ok(())
158 }
159}
160
161fn read_jsonl<T: serde::de::DeserializeOwned>(path: &Path) -> Result<Vec<T>, JsonError> {
162 let mut out = Vec::new();
163 if !path.exists() {
164 return Ok(out);
165 }
166 let file = std::fs::File::open(path)?;
167 for line in BufReader::new(file).lines() {
168 let line = line?;
169 if line.trim().is_empty() {
170 continue;
171 }
172 out.push(serde_json::from_str(&line)?);
173 }
174 Ok(out)
175}
176
177#[derive(Debug, serde::Serialize)]
178struct JsonDoc<'a> {
179 schema: crate::schema::GraphSchema,
180 nodes: &'a [Node],
181 edges: &'a [Edge],
182}
183
184#[derive(Debug, serde::Deserialize)]
185struct JsonDocOwned {
186 #[allow(dead_code)]
187 schema: crate::schema::GraphSchema,
188 nodes: Vec<Node>,
189 edges: Vec<Edge>,
190}
191
192#[derive(Debug, thiserror::Error)]
194pub enum JsonError {
195 #[error("IO error: {0}")]
196 Io(#[from] std::io::Error),
197 #[error("JSON error: {0}")]
198 Json(#[from] serde_json::Error),
199}
200
201impl GraphBackend for JsonBackend {
202 type Error = JsonError;
203
204 fn init(&mut self) -> Result<(), Self::Error> {
205 self.load()?;
206 Ok(())
207 }
208
209 fn write_nodes(&mut self, nodes: &[Node]) -> Result<(), Self::Error> {
210 self.nodes.extend(nodes.iter().cloned());
211 self.flush()?;
212 Ok(())
213 }
214
215 fn write_edges(&mut self, edges: &[Edge]) -> Result<(), Self::Error> {
216 self.edges.extend(edges.iter().cloned());
217 self.flush()?;
218 Ok(())
219 }
220
221 fn read_nodes(&self) -> Result<Vec<Node>, Self::Error> {
222 Ok(self.nodes.clone())
223 }
224
225 fn read_edges(&self) -> Result<Vec<Edge>, Self::Error> {
226 Ok(self.edges.clone())
227 }
228
229 fn find_nodes_by_type(&self, kind: crate::schema::NodeType) -> Result<Vec<Node>, Self::Error> {
230 Ok(self
231 .nodes
232 .iter()
233 .filter(|n| n.kind == kind)
234 .cloned()
235 .collect())
236 }
237
238 fn neighbors(
239 &self,
240 node_id: &str,
241 edge_type: Option<EdgeType>,
242 ) -> Result<Vec<Edge>, Self::Error> {
243 Ok(self
244 .edges
245 .iter()
246 .filter(|e| {
247 e.source_id == node_id && edge_type.as_ref().map_or(true, |et| e.kind == *et)
248 })
249 .cloned()
250 .collect())
251 }
252
253 fn clear(&mut self) -> Result<(), Self::Error> {
254 self.nodes.clear();
255 self.edges.clear();
256 let _ = std::fs::remove_file(&self.path);
257 let _ = std::fs::remove_file(self.path.with_extension("nodes.jsonl"));
258 let _ = std::fs::remove_file(self.path.with_extension("edges.jsonl"));
259 Ok(())
260 }
261}
262
263#[cfg(test)]
264mod tests {
265 use super::*;
266 use crate::schema::NodeType;
267 use tempfile::NamedTempFile;
268
269 #[test]
270 fn json_roundtrip() {
271 let file = NamedTempFile::new().unwrap();
272 let mut backend = JsonBackend::open(file.path());
273 backend.init().unwrap();
274
275 let node = Node::new("n1", NodeType::Domain, "example.com");
276 backend.write_nodes(&[node.clone()]).unwrap();
277
278 let edge = Edge::new("n1", "n2", EdgeType::ResolvesTo);
279 backend.write_edges(&[edge.clone()]).unwrap();
280
281 let mut backend2 = JsonBackend::open(file.path());
283 backend2.init().unwrap();
284
285 let nodes = backend2.read_nodes().unwrap();
286 assert_eq!(nodes.len(), 1);
287 assert_eq!(nodes[0].id, "n1");
288
289 let edges = backend2.read_edges().unwrap();
290 assert_eq!(edges.len(), 1);
291 assert_eq!(edges[0].source_id, "n1");
292 }
293
294 #[test]
300 fn json_load_rejects_path_traversal_in_manifest() {
301 let dir = tempfile::tempdir().unwrap();
302 let manifest = dir.path().join("scan.json");
303
304 std::fs::write(manifest.with_extension("nodes.jsonl"), "").unwrap();
307 std::fs::write(manifest.with_extension("edges.jsonl"), "").unwrap();
308
309 let sentinel = dir.path().join("secret.jsonl");
311 std::fs::write(&sentinel, r#"{"id":"leaked","kind":"Domain","label":"x"}"#).unwrap();
312
313 let manifest_body = serde_json::json!({
314 "format": "jsonl",
315 "schema": crate::schema::GraphSchema::current(),
316 "nodes_file": sentinel.to_string_lossy(),
317 "edges_file": "../../../etc/passwd",
318 "node_count": 0,
319 "edge_count": 0,
320 });
321 std::fs::write(&manifest, serde_json::to_string(&manifest_body).unwrap()).unwrap();
322
323 let mut backend = JsonBackend::open(&manifest);
324 backend.init().expect("safe fallback load should succeed");
325
326 let nodes = backend.read_nodes().unwrap();
327 assert!(
328 !nodes.iter().any(|n| n.id == "leaked"),
329 "manifest-named absolute path was followed — path-traversal guard regressed"
330 );
331 }
332
333 #[test]
334 fn json_streaming_threshold() {
335 let file = NamedTempFile::new().unwrap();
336 let mut backend = JsonBackend::open(file.path());
337 backend.init().unwrap();
338
339 let mut nodes = Vec::new();
341 for i in 0..STREAMING_THRESHOLD + 1 {
342 nodes.push(Node::new(
343 format!("n{i}"),
344 NodeType::Subdomain,
345 format!("sub{i}.example.com"),
346 ));
347 }
348 backend.write_nodes(&nodes).unwrap();
349
350 assert!(file.path().exists());
352 assert!(file.path().with_extension("nodes.jsonl").exists());
353
354 let mut backend2 = JsonBackend::open(file.path());
355 backend2.init().unwrap();
356 assert_eq!(
357 backend2.read_nodes().unwrap().len(),
358 STREAMING_THRESHOLD + 1
359 );
360 }
361}