use std::cell::RefCell;
use std::collections::HashMap;
use std::path::{Path, PathBuf};
use crate::runner::check_reported_path;
pub trait SnippetSource {
fn snippet(&self, path: &str, start: u32, end: u32) -> Option<String>;
}
#[derive(Debug, Clone, Copy, Default)]
pub struct NoSnippets;
impl SnippetSource for NoSnippets {
fn snippet(&self, _path: &str, _start: u32, _end: u32) -> Option<String> {
None
}
}
#[derive(Debug)]
pub struct WorktreeSnippets {
root: PathBuf,
cache: RefCell<HashMap<String, Option<Vec<u8>>>>,
}
impl WorktreeSnippets {
pub const MAX_FILE_BYTES: u64 = 8 << 20;
#[must_use]
pub fn new(root: impl Into<PathBuf>) -> Self {
Self {
root: root.into(),
cache: RefCell::new(HashMap::new()),
}
}
fn bytes(&self, path: &str) -> Option<Vec<u8>> {
if let Some(hit) = self.cache.borrow().get(path) {
return hit.clone();
}
let read = self.read(path);
self.cache
.borrow_mut()
.insert(path.to_owned(), read.clone());
read
}
fn read(&self, path: &str) -> Option<Vec<u8>> {
check_reported_path(path).ok()?;
let full = self.root.join(Path::new(path));
let meta = std::fs::metadata(&full).ok()?;
if !meta.is_file() || meta.len() > Self::MAX_FILE_BYTES {
return None;
}
std::fs::read(&full).ok()
}
}
impl SnippetSource for WorktreeSnippets {
fn snippet(&self, path: &str, start: u32, end: u32) -> Option<String> {
if end < start {
return None;
}
let bytes = self.bytes(path)?;
let (from, to) = (start as usize, end as usize);
if to > bytes.len() {
return None;
}
String::from_utf8(bytes[from..to].to_vec()).ok()
}
}
#[cfg(test)]
mod tests {
use super::{NoSnippets, SnippetSource, WorktreeSnippets};
struct Scratch(std::path::PathBuf);
impl Scratch {
fn new(name: &str) -> Self {
let dir = std::env::temp_dir().join(format!("rto-exec-snippet-{name}"));
std::fs::remove_dir_all(&dir).ok();
std::fs::create_dir_all(dir.join("src")).expect("create");
std::fs::write(dir.join("src/app.py"), b"import os\nos.system(cmd)\n").expect("write");
Self(dir)
}
}
impl Drop for Scratch {
fn drop(&mut self) {
std::fs::remove_dir_all(&self.0).ok();
}
}
#[test]
fn reads_the_bytes_a_finding_points_at() {
let scratch = Scratch::new("reads");
let snippets = WorktreeSnippets::new(&scratch.0);
assert_eq!(
snippets.snippet("src/app.py", 10, 24).as_deref(),
Some("os.system(cmd)")
);
assert_eq!(
snippets.snippet("src/app.py", 10, 24).as_deref(),
Some("os.system(cmd)")
);
}
#[test]
fn a_file_it_does_not_have_is_simply_unavailable() {
let scratch = Scratch::new("missing");
let snippets = WorktreeSnippets::new(&scratch.0);
assert!(snippets.snippet("src/nope.py", 0, 4).is_none());
}
#[test]
fn refuses_to_read_outside_the_worktree() {
let scratch = Scratch::new("escape");
let snippets = WorktreeSnippets::new(&scratch.0);
for hostile in ["../../../etc/passwd", "/etc/passwd", ""] {
assert!(snippets.snippet(hostile, 0, 4).is_none(), "{hostile:?}");
}
}
#[test]
fn offsets_past_the_end_yield_nothing_rather_than_a_truncated_slice() {
let scratch = Scratch::new("bounds");
let snippets = WorktreeSnippets::new(&scratch.0);
assert!(snippets.snippet("src/app.py", 0, 9_999).is_none());
assert!(snippets.snippet("src/app.py", 20, 5).is_none());
}
#[test]
fn the_empty_source_answers_nothing() {
assert!(NoSnippets.snippet("src/app.py", 0, 4).is_none());
}
}