Skip to main content

gossan_graph/store/
json.rs

1//! JSON graph backend — stores nodes and edges as a single JSON document.
2//!
3//! For large graphs (>10K nodes) the backend automatically flushes to a
4//! streaming JSONL file instead of a monolithic array.
5
6use std::io::{BufRead, BufReader, Write};
7use std::path::{Path, PathBuf};
8
9use crate::store::GraphBackend;
10use crate::{schema::EdgeType, Edge, Node};
11
12/// In-memory + JSON file backend.
13pub struct JsonBackend {
14    path: PathBuf,
15    nodes: Vec<Node>,
16    edges: Vec<Edge>,
17}
18
19/// Threshold above which we prefer JSONL streaming for writes.
20const STREAMING_THRESHOLD: usize = 10_000;
21
22impl JsonBackend {
23    /// Open or create a JSON graph file.
24    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        // Write a tiny manifest so consumers know where the data is.
62        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        // The file is one of three shapes:
81        //   1. a multi-line pretty-printed `JsonDocOwned` (the small-graph
82        //      flush path — first line will just be "{")
83        //   2. a single-line manifest with `"format": "jsonl"` (the
84        //      streaming flush path; nodes/edges in sibling .jsonl files)
85        //   3. mixed JSONL — each line is a Node or Edge
86        // Empty files are valid (a fresh handle on a NamedTempFile); we
87        // shortcut on zero length to avoid serde failing with "EOF while
88        // parsing".
89        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            // Try the whole file as one JSON object first — covers
97            // cases (1) and (2). Fall back to per-line manifest parse.
98            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                // The manifest *names* the sibling jsonl files. A
106                // malicious manifest could try to point those names at
107                // /etc/passwd, ~/.ssh/id_rsa, etc., and surface their
108                // contents via serde parse-error messages. Constrain
109                // both to the manifest's parent directory and reject
110                // any path that escapes via .. or absolute prefix.
111                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            // Fall through to mixed-JSONL handling below.
142        }
143
144        // Mixed JSONL: each non-empty line is a Node or Edge,
145        // discriminated by presence of `source_id`.
146        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/// Error type for JSON backend operations.
193#[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        // Re-open and verify
282        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    /// Adversarial: a malicious manifest that names an absolute or
295    /// `..`-escaped path for `nodes_file` / `edges_file` MUST be
296    /// silently ignored — the loader falls back to the safe sibling
297    /// default — so we never read /etc/passwd or surface its content
298    /// through serde parse-error messages.
299    #[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        // Drop a benign sibling so the safe default load resolves to
305        // an empty corpus rather than an "no such file" error.
306        std::fs::write(manifest.with_extension("nodes.jsonl"), "").unwrap();
307        std::fs::write(manifest.with_extension("edges.jsonl"), "").unwrap();
308
309        // Sentinel target the attacker would love to leak.
310        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        // Write just over the threshold
340        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        // Manifest should exist
351        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}