Skip to content

Commit

Permalink
chore: move fs cheatcodes to own file
Browse files Browse the repository at this point in the history
  • Loading branch information
@DaniPopes
DaniPopes committed Apr 27, 2023
1 parent 4d83342 commit f9eb363
Show file tree
Hide file tree
Showing 3 changed files with 304 additions and 292 deletions.
298 changes: 7 additions & 291 deletions evm/src/executor/inspector/cheatcodes/ext.rs
View file
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
use crate::{
abi::{DirEntry, FsMetadata, HEVMCalls},
abi::HEVMCalls,
error,
executor::inspector::{cheatcodes::util, Cheatcodes},
};
Expand All @@ -15,17 +15,8 @@ use hex::FromHex;
use jsonpath_lib;
use serde::Deserialize;
use serde_json::Value;
use std::{
collections::BTreeMap,
env,
io::{BufRead, BufReader, Write},
path::Path,
process::Command,
str::FromStr,
time::UNIX_EPOCH,
};
use std::{collections::BTreeMap, env, path::Path, process::Command, str::FromStr};
use tracing::{error, trace};
use walkdir::WalkDir;

/// Invokes a `Command` with the given args and returns the abi encoded response
///
Expand Down Expand Up @@ -184,260 +175,6 @@ fn get_env(
}
}

fn project_root(state: &Cheatcodes) -> Result<Bytes, Bytes> {
let root = state.config.root.display().to_string();

Ok(abi::encode(&[Token::String(root)]).into())
}

fn read_file(state: &Cheatcodes, path: impl AsRef<Path>) -> Result<Bytes, Bytes> {
let path =
state.config.ensure_path_allowed(path, FsAccessKind::Read).map_err(error::encode_error)?;

let data = fs::read_to_string(path).map_err(error::encode_error)?;

Ok(abi::encode(&[Token::String(data)]).into())
}

fn read_file_binary(state: &Cheatcodes, path: impl AsRef<Path>) -> Result<Bytes, Bytes> {
let path =
state.config.ensure_path_allowed(path, FsAccessKind::Read).map_err(error::encode_error)?;

let data = fs::read(path).map_err(error::encode_error)?;

Ok(abi::encode(&[Token::Bytes(data)]).into())
}

fn read_line(state: &mut Cheatcodes, path: impl AsRef<Path>) -> Result<Bytes, Bytes> {
let path =
state.config.ensure_path_allowed(path, FsAccessKind::Read).map_err(error::encode_error)?;

// Get reader for previously opened file to continue reading OR initialize new reader
let reader = state
.context
.opened_read_files
.entry(path.clone())
.or_insert(BufReader::new(fs::open(path).map_err(error::encode_error)?));

let mut line: String = String::new();
reader.read_line(&mut line).map_err(error::encode_error)?;

// Remove trailing newline character, preserving others for cases where it may be important
if line.ends_with('\n') {
line.pop();
if line.ends_with('\r') {
line.pop();
}
}

Ok(abi::encode(&[Token::String(line)]).into())
}

/// Writes `content` to `path`.
///
/// This function will create a file if it does not exist, and will entirely replace its contents if
/// it does.
///
/// Caution: writing files is only allowed if the targeted path is allowed, (inside `<root>/` by
/// default)
fn write_file(
state: &Cheatcodes,
path: impl AsRef<Path>,
content: impl AsRef<[u8]>,
) -> Result<Bytes, Bytes> {
let path =
state.config.ensure_path_allowed(path, FsAccessKind::Write).map_err(error::encode_error)?;
// write access to foundry.toml is not allowed
state.config.ensure_not_foundry_toml(&path).map_err(error::encode_error)?;

if state.fs_commit {
fs::write(path, content.as_ref()).map_err(error::encode_error)?;
}

Ok(Bytes::new())
}

/// Writes a single line to the file.
///
/// This will create a file if it does not exist, and append the `line` if it does.
fn write_line(state: &Cheatcodes, path: impl AsRef<Path>, line: &str) -> Result<Bytes, Bytes> {
let path =
state.config.ensure_path_allowed(path, FsAccessKind::Write).map_err(error::encode_error)?;
state.config.ensure_not_foundry_toml(&path).map_err(error::encode_error)?;

if state.fs_commit {
let mut file = std::fs::OpenOptions::new()
.append(true)
.create(true)
.open(path)
.map_err(error::encode_error)?;

writeln!(file, "{line}").map_err(error::encode_error)?;
}

Ok(Bytes::new())
}

fn close_file(state: &mut Cheatcodes, path: impl AsRef<Path>) -> Result<Bytes, Bytes> {
let path =
state.config.ensure_path_allowed(path, FsAccessKind::Read).map_err(error::encode_error)?;

state.context.opened_read_files.remove(&path);

Ok(Bytes::new())
}

/// Removes a file from the filesystem.
///
/// Only files inside `<root>/` can be removed, `foundry.toml` excluded.
///
/// This will return an error if the path points to a directory, or the file does not exist
fn remove_file(state: &mut Cheatcodes, path: impl AsRef<Path>) -> Result<Bytes, Bytes> {
let path =
state.config.ensure_path_allowed(path, FsAccessKind::Write).map_err(error::encode_error)?;
state.config.ensure_not_foundry_toml(&path).map_err(error::encode_error)?;

// also remove from the set if opened previously
state.context.opened_read_files.remove(&path);

if state.fs_commit {
fs::remove_file(&path).map_err(error::encode_error)?;
}

Ok(Bytes::new())
}

/// Creates a new, empty directory at the provided path.
///
/// If `recursive` is true, it will also create all the parent directories if they don't exist.
///
/// # Errors
///
/// This function will return an error in the following situations, but is not limited to just these
/// cases:
///
/// - User lacks permissions to modify `path`.
/// - A parent of the given path doesn't exist and `recursive` is false.
/// - `path` already exists and `recursive` is false.
fn create_dir(state: &Cheatcodes, path: impl AsRef<Path>, recursive: bool) -> Result<Bytes, Bytes> {
let path =
state.config.ensure_path_allowed(path, FsAccessKind::Write).map_err(error::encode_error)?;
if recursive { fs::create_dir_all(path) } else { fs::create_dir(path) }
.map(|()| Bytes::new())
.map_err(error::encode_error)
}

/// Removes a directory at the provided path.
///
/// This will also remove all the directory's contents recursively if `recursive` is true.
///
/// # Errors
///
/// This function will return an error in the following situations, but is not limited to just these
/// cases:
///
/// - `path` doesn't exist.
/// - `path` isn't a directory.
/// - User lacks permissions to modify `path`.
/// - The directory is not empty and `recursive` is false.
fn remove_dir(state: &Cheatcodes, path: impl AsRef<Path>, recursive: bool) -> Result<Bytes, Bytes> {
let path =
state.config.ensure_path_allowed(path, FsAccessKind::Write).map_err(error::encode_error)?;
if recursive { fs::remove_dir_all(path) } else { fs::remove_dir(path) }
.map(|()| Bytes::new())
.map_err(error::encode_error)
}

/// Reads the directory at the given path recursively, up to `max_depth`.
///
/// Follows symbolic links if `follow_links` is true.
fn read_dir(
state: &Cheatcodes,
path: impl AsRef<Path>,
max_depth: u64,
follow_links: bool,
) -> Result<Bytes, Bytes> {
let root =
state.config.ensure_path_allowed(path, FsAccessKind::Read).map_err(error::encode_error)?;
let paths: Vec<Token> = WalkDir::new(root)
.min_depth(1)
.max_depth(max_depth.try_into().map_err(error::encode_error)?)
.follow_links(follow_links)
.into_iter()
.map(|entry| match entry {
Ok(entry) => {
let entry = DirEntry {
error: String::new(),
path: entry.path().display().to_string(),
depth: entry.depth() as u64,
is_dir: entry.file_type().is_dir(),
is_symlink: entry.path_is_symlink(),
};
eprintln!("OK {:?}", entry);
Token::Bytes(entry.encode())
}
Err(e) => {
let entry = DirEntry {
error: e.to_string(),
path: e.path().map(|p| p.display().to_string()).unwrap_or_default(),
depth: e.depth() as u64,
is_dir: false,
is_symlink: false,
};
eprintln!("ERR {:?}", entry);
Token::Bytes(entry.encode())
}
})
.collect();
eprintln!("GOT {} ENTRIES", paths.len());
Ok(abi::encode(&[Token::Array(paths)]).into())
}

/// Reads a symbolic link, returning the path that the link points to.
///
/// # Errors
///
/// This function will return an error in the following situations, but is not limited to just these
/// cases:
///
/// - `path` is not a symbolic link.
/// - `path` does not exist.
fn read_link(state: &Cheatcodes, path: impl AsRef<Path>) -> Result<Bytes, Bytes> {
let path =
state.config.ensure_path_allowed(path, FsAccessKind::Read).map_err(error::encode_error)?;

let target = fs::read_link(path).map_err(error::encode_error)?;

Ok(abi::encode(&[Token::String(target.display().to_string())]).into())
}

/// Gets the metadata of a file/directory
///
/// This will return an error if no file/directory is found, or if the target path isn't allowed
fn fs_metadata(state: &Cheatcodes, path: impl AsRef<Path>) -> Result<Bytes, Bytes> {
let path =
state.config.ensure_path_allowed(path, FsAccessKind::Read).map_err(error::encode_error)?;

let metadata = path.metadata().map_err(error::encode_error)?;

// These fields not available on all platforms; default to 0
let [modified, accessed, created] =
[metadata.modified(), metadata.accessed(), metadata.created()].map(|time| {
time.unwrap_or(UNIX_EPOCH).duration_since(UNIX_EPOCH).unwrap_or_default().as_secs()
});

let metadata = FsMetadata {
is_dir: metadata.is_dir(),
is_symlink: metadata.is_symlink(),
length: metadata.len().into(),
read_only: metadata.permissions().readonly(),
modified: modified.into(),
accessed: accessed.into(),
created: created.into(),
};
Ok(metadata.encode().into())
}

/// Converts a serde_json::Value to an abi::Token
/// The function is designed to run recursively, so that in case of an object
/// it will call itself to convert each of it's value and encode the whole as a
Expand Down Expand Up @@ -644,21 +381,17 @@ fn write_json(
json
})
.map_err(error::encode_error)?;
write_file(_state, path, json_string)?;
super::fs::write_file(_state, path, json_string)?;
Ok(Bytes::new())
}

pub fn apply(
state: &mut Cheatcodes,
ffi_enabled: bool,
call: &HEVMCalls,
) -> Option<Result<Bytes, Bytes>> {
pub fn apply(state: &mut Cheatcodes, call: &HEVMCalls) -> Option<Result<Bytes, Bytes>> {
Some(match call {
HEVMCalls::Ffi(inner) => {
if !ffi_enabled {
Err(error::encode_error("FFI disabled: run again with `--ffi` if you want to allow tests to call external scripts."))
} else {
if state.config.ffi {
ffi(state, &inner.0)
} else {
Err(error::encode_error("FFI disabled: run again with `--ffi` if you want to allow tests to call external scripts."))
}
}
HEVMCalls::GetCode(inner) => get_code(state, &inner.0),
Expand Down Expand Up @@ -743,23 +476,6 @@ pub fn apply(
Some(inner.2.iter().map(hex::encode).collect::<Vec<_>>().join(&inner.1)),
),

HEVMCalls::ProjectRoot(_) => project_root(state),
HEVMCalls::ReadFile(inner) => read_file(state, &inner.0),
HEVMCalls::ReadFileBinary(inner) => read_file_binary(state, &inner.0),
HEVMCalls::ReadLine(inner) => read_line(state, &inner.0),
HEVMCalls::WriteFile(inner) => write_file(state, &inner.0, &inner.1),
HEVMCalls::WriteFileBinary(inner) => write_file(state, &inner.0, &inner.1),
HEVMCalls::WriteLine(inner) => write_line(state, &inner.0, &inner.1),
HEVMCalls::CloseFile(inner) => close_file(state, &inner.0),
HEVMCalls::RemoveFile(inner) => remove_file(state, &inner.0),
HEVMCalls::FsMetadata(inner) => fs_metadata(state, &inner.0),
HEVMCalls::ReadLink(inner) => read_link(state, &inner.0),
HEVMCalls::CreateDir(inner) => create_dir(state, &inner.0, inner.1),
HEVMCalls::RemoveDir(inner) => remove_dir(state, &inner.0, inner.1),
HEVMCalls::ReadDir0(inner) => read_dir(state, &inner.0, 1, false),
HEVMCalls::ReadDir1(inner) => read_dir(state, &inner.0, inner.1, false),
HEVMCalls::ReadDir2(inner) => read_dir(state, &inner.0, inner.1, inner.2),

// If no key argument is passed, return the whole JSON object.
// "$" is the JSONPath key for the root of the object
HEVMCalls::ParseJson0(inner) => parse_json(state, &inner.0, "$", None),
Expand Down
Loading

0 comments on commit f9eb363

Please sign in to comment.