wip

Author: Jon-Becker

crates/cache/src/error.rs

@@ -1,7 +1,12 @@
+//! Cache errors
+
+/// Generic error type for heimdall cache operations
 #[derive(Debug, thiserror::Error)]
 pub enum Error {
+    /// Generic error
     #[error("Error: {0}")]
     Generic(String),
+    /// An IO error occurred
     #[error("IO error: {0}")]
     IOError(#[from] std::io::Error),
 }

crates/cache/src/lib.rs

@@ -1,3 +1,7 @@
+//! A simple cache system for heimdall-rs
+//! Stores objects in ~/.bifrost/cache as bincode serialized files
+//! Objects are stored with an expiry time, and are deleted if they are expired
+
 use clap::Parser;
 use serde::{de::DeserializeOwned, Deserialize, Serialize};
 #[allow(deprecated)]
@@ -17,6 +21,7 @@ pub(crate) mod util;
     override_usage = "heimdall cache <SUBCOMMAND>"
 )]
 pub struct CacheArgs {
+    /// Cache subcommand
     #[clap(subcommand)]
     pub sub: Subcommands,
 }
@@ -33,12 +38,15 @@ pub struct NoArguments {}
 )]
 #[allow(clippy::large_enum_variant)]
 pub enum Subcommands {
+    /// Clear the cache, removing all objects
     #[clap(name = "clean", about = "Removes all cached objects in ~/.bifrost/cache")]
     Clean(NoArguments),
 
+    /// List all cached objects
     #[clap(name = "ls", about = "Lists all cached objects in ~/.bifrost/cache")]
     Ls(NoArguments),
 
+    /// Print the size of the cache in ~/.bifrost/cache
     #[clap(name = "size", about = "Prints the size of the cache in ~/.bifrost/cache")]
     Size(NoArguments),
 }
@@ -47,7 +55,9 @@ pub enum Subcommands {
 /// The expiry time is a unix timestamp
 #[derive(Debug, Clone, Deserialize, Serialize)]
 pub struct Cache<T> {
+    /// The value stored in the cache
     pub value: T,
+    /// The expiry time of the cache object
     pub expiry: u64,
 }
 
@@ -215,7 +225,8 @@ pub fn delete_cache(key: &str) -> Result<(), Error> {
 #[allow(deprecated)]
 pub fn read_cache<T>(key: &str) -> Result<Option<T>, Error>
 where
-    T: 'static + DeserializeOwned, {
+    T: 'static + DeserializeOwned,
+{
     let home = home_dir().ok_or(Error::Generic(
         "failed to get home directory. does your os support `std::env::home_dir()`?".to_string(),
     ))?;
@@ -238,8 +249,8 @@ where
         .map_err(|e| Error::Generic(format!("failed to deserialize cache object: {:?}", e)))?;
 
     // check if the cache has expired, if so, delete it and return None
-    if cache.expiry <
-        std::time::SystemTime::now()
+    if cache.expiry
+        < std::time::SystemTime::now()
             .duration_since(std::time::UNIX_EPOCH)
             .map_err(|e| Error::Generic(format!("failed to get current time: {:?}", e)))?
             .as_secs()
@@ -266,7 +277,8 @@ where
 #[allow(deprecated)]
 pub fn store_cache<T>(key: &str, value: T, expiry: Option<u64>) -> Result<(), Error>
 where
-    T: Serialize, {
+    T: Serialize,
+{
     let home = home_dir().ok_or(Error::Generic(
         "failed to get home directory. does your os support `std::env::home_dir()`?".to_string(),
     ))?;
@@ -278,8 +290,8 @@ where
         std::time::SystemTime::now()
             .duration_since(std::time::UNIX_EPOCH)
             .map_err(|e| Error::Generic(format!("failed to get current time: {:?}", e)))?
-            .as_secs() +
-            60 * 60 * 24 * 90,
+            .as_secs()
+            + 60 * 60 * 24 * 90,
     );
 
     let cache = Cache { value, expiry };
@@ -304,7 +316,8 @@ pub async fn with_cache<T, F, Fut>(key: &str, func: F) -> eyre::Result<T>
 where
     T: 'static + Serialize + DeserializeOwned + Send + Sync,
     F: FnOnce() -> Fut + Send,
-    Fut: std::future::Future<Output = Result<T, eyre::Report>> + Send, {
+    Fut: std::future::Future<Output = Result<T, eyre::Report>> + Send,
+{
     // Try to read from cache
     match read_cache::<T>(key) {
         Ok(Some(cached_value)) => {

crates/cache/src/util.rs

@@ -10,20 +10,20 @@ use std::{
 use crate::error::Error;
 
 /// Decode a hex string into a bytearray
-pub fn decode_hex(s: &str) -> Result<Vec<u8>, ParseIntError> {
+pub(crate) fn decode_hex(s: &str) -> Result<Vec<u8>, ParseIntError> {
     (0..s.len()).step_by(2).map(|i| u8::from_str_radix(&s[i..i + 2], 16)).collect()
 }
 
 /// Encode a bytearray into a hex string
-pub fn encode_hex(s: Vec<u8>) -> String {
+pub(crate) fn encode_hex(s: Vec<u8>) -> String {
     s.iter().fold(String::new(), |mut acc: String, b| {
         write!(acc, "{b:02x}", b = b).expect("unable to write");
         acc
     })
 }
 
-/// Prettify bytes into a human-readable format \
-pub fn prettify_bytes(bytes: u64) -> String {
+/// Prettify bytes into a human-readable format
+pub(crate) fn prettify_bytes(bytes: u64) -> String {
     if bytes < 1024 {
         format!("{bytes} B")
     } else if bytes < 1024 * 1024 {
@@ -39,7 +39,8 @@ pub fn prettify_bytes(bytes: u64) -> String {
 }
 
 /// Write contents to a file on the disc
-pub fn write_file(path_str: &str, contents: &str) -> Result<(), Error> {
+/// If the parent directory does not exist, it will be created
+pub(crate) fn write_file(path_str: &str, contents: &str) -> Result<(), Error> {
     let path = Path::new(path_str);
 
     if let Some(prefix) = path.parent() {
@@ -58,7 +59,8 @@ pub fn write_file(path_str: &str, contents: &str) -> Result<(), Error> {
 }
 
 /// Read contents from a file on the disc
-pub fn read_file(path: &str) -> Result<String, Error> {
+/// Returns the contents as a string
+pub(crate) fn read_file(path: &str) -> Result<String, Error> {
     let path = Path::new(path);
     let mut file = File::open(path)
         .map_err(|e| Error::IOError(std::io::Error::new(std::io::ErrorKind::Other, e)))?;
@@ -68,7 +70,8 @@ pub fn read_file(path: &str) -> Result<String, Error> {
 }
 
 /// Delete a file or directory on the disc
-pub fn delete_path(_path: &str) -> bool {
+/// Returns true if the operation was successful
+pub(crate) fn delete_path(_path: &str) -> bool {
     let path = match std::path::Path::new(_path).to_str() {
         Some(path) => path,
         None => return false,

crates/cfg/src/core/graph.rs

@@ -9,8 +9,7 @@ use petgraph::{matrix_graph::NodeIndex, Graph};
 
 /// convert a symbolic execution [`VMTrace`] into a [`Graph`] of blocks, illustrating the
 /// control-flow graph found by the symbolic execution engine.
-// TODO: should this be a trait for VMTrace to implement?
-pub fn build_cfg(
+pub(crate) fn build_cfg(
     vm_trace: &VMTrace,
     contract_cfg: &mut Graph<String, String>,
     parent_node: Option<NodeIndex<u32>>,
@@ -61,8 +60,8 @@ pub fn build_cfg(
                 .first()
                 .ok_or_eyre("failed to get first operation")?
                 .last_instruction
-                .opcode ==
-                JUMPDEST,
+                .opcode
+                == JUMPDEST,
         )?;
     }
 

crates/cfg/src/core/mod.rs

@@ -13,12 +13,15 @@ use super::CfgArgs;
 use crate::{core::graph::build_cfg, error::Error};
 use tracing::{debug, info};
 
+/// The result of the cfg command. Contains the generated control flow graph.
 #[derive(Debug, Clone)]
 pub struct CfgResult {
+    /// The generated control flow graph of the contract.
     pub graph: Graph<String, String>,
 }
 
 impl CfgResult {
+    /// Returns the control flow graph as a graphviz formatted string.
     pub fn as_dot(&self, color_edges: bool) -> String {
         let output = format!("{}", Dot::with_config(&self.graph, &[]));
 
@@ -44,6 +47,7 @@ impl CfgResult {
     }
 }
 
+/// Generates a control flow graph for the target contract.
 pub async fn cfg(args: CfgArgs) -> Result<CfgResult, Error> {
     // init
     let start_time = Instant::now();

crates/cfg/src/error.rs

@@ -1,9 +1,15 @@
+//! CFG Errors
+
+/// Generic error type for the CFG Module
 #[derive(Debug, thiserror::Error)]
 pub enum Error {
+    /// Error when trying to fetch information from the chain
     #[error("Fetch error: {0}")]
     FetchError(String),
+    /// Error when disassembling contract bytecode
     #[error("Disassembly error: {0}")]
     DisassemblyError(#[from] heimdall_disassembler::Error),
+    /// Generic error
     #[error("Internal error: {0}")]
     Eyre(#[from] eyre::Report),
 }

crates/cfg/src/interfaces/args.rs

@@ -4,6 +4,7 @@ use eyre::Result;
 use heimdall_common::ether::bytecode::get_bytecode_from_target;
 use heimdall_config::parse_url_arg;
 
+/// Arguments for the CFG subcommand
 #[derive(Debug, Clone, Parser, Builder)]
 #[clap(
     about = "Generate a visual control flow graph for EVM bytecode",
@@ -43,12 +44,14 @@ pub struct CfgArgs {
 }
 
 impl CfgArgs {
+    /// Get the bytecode for the target
     pub async fn get_bytecode(&self) -> Result<Vec<u8>> {
         get_bytecode_from_target(&self.target, &self.rpc_url).await
     }
 }
 
 impl CfgArgsBuilder {
+    /// Create a new instance of the [`CfgArgsBuilder`]
     pub fn new() -> Self {
         Self {
             target: Some(String::new()),

crates/cfg/src/lib.rs

@@ -1,3 +1,6 @@
+//! The CFG module is responsible for generating control-flow graphs from the given
+//! contract's source code via symbolic execution.
+
 mod error;
 
 mod core;

crates/vm/src/core/opcodes/wrapped.rs

@@ -26,7 +26,7 @@ impl WrappedOpcode {
 }
 
 impl std::fmt::Display for WrappedOpcode {
-    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         write!(
             f,
             "{}({})",
@@ -49,7 +49,7 @@ impl WrappedInput {
 }
 
 impl std::fmt::Display for WrappedInput {
-    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         match self {
             WrappedInput::Raw(u256) => write!(f, "{u256}"),
             WrappedInput::Opcode(opcode) => write!(f, "{opcode}"),