Merge pull request #21 from trailofbits/library-search

Add library search functionality

cli/src/main.rs

@@ -19,6 +19,10 @@ struct Cli {
     #[clap(name = "INPUT")]
     input_files: Vec<PathBuf>,
 
+    /// Library file paths
+    #[clap(short = 'L', long = "library", name = "LIBRARIES")]
+    libraries: Vec<PathBuf>,
+
     /// Output level (INFO, WARNING, or ERROR)
     #[clap(short = 'l', long = "level", name = "LEVEL", default_value = config::DEFAULT_LEVEL)]
     output_level: MessageCategory,
@@ -69,7 +73,9 @@ fn main() -> ExitCode {
     }
 
     // Set up analysis runner.
-    let (mut runner, reports) = AnalysisRunner::new(options.curve).with_files(&options.input_files);
+    let (mut runner, reports) = AnalysisRunner::new(options.curve)
+        .with_libraries(&options.libraries)
+        .with_files(&options.input_files);
 
     // Set up writer and write reports to `stdout`.
     let allow_list = options.allow_list.clone();

parser/src/include_logic.rs

@@ -1,44 +1,73 @@
 use crate::errors::FileOsError;
 
+use log::debug;
+
 use super::errors::IncludeError;
 use program_structure::ast::Include;
 use program_structure::report::{Report, ReportCollection};
 use std::collections::HashSet;
+use std::ffi::OsString;
 use std::fs;
 use std::path::PathBuf;
 
 pub struct FileStack {
     current_location: Option<PathBuf>,
     black_paths: HashSet<PathBuf>,
     user_inputs: HashSet<PathBuf>,
+    libraries: Vec<Library>,
     stack: Vec<PathBuf>,
 }
 
+#[derive(Debug)]
+struct Library {
+    dir: bool,
+    path: PathBuf,
+}
+
 impl FileStack {
-    pub fn new(paths: &[PathBuf], reports: &mut ReportCollection) -> FileStack {
+    pub fn new(paths: &[PathBuf], libs: &[PathBuf], reports: &mut ReportCollection) -> FileStack {
         let mut result = FileStack {
             current_location: None,
             black_paths: HashSet::new(),
             user_inputs: HashSet::new(),
+            libraries: Vec::new(),
             stack: Vec::new(),
         };
+        result.add_libraries(libs, reports);
         result.add_files(paths, reports);
         result.user_inputs = result.stack.iter().cloned().collect::<HashSet<_>>();
 
         result
     }
 
+    fn add_libraries(&mut self, libs: &[PathBuf], reports: &mut ReportCollection) {
+        for path in libs {
+            if path.is_dir() {
+                self.libraries.push(Library { dir: true, path: path.clone() });
+            } else if let Some(extension) = path.extension() {
+                // Add Circom files to file stack.
+                if extension == "circom" {
+                    match fs::canonicalize(path) {
+                        Ok(path) => self.libraries.push(Library { dir: false, path: path.clone() }),
+                        Err(_) => {
+                            reports.push(
+                                FileOsError { path: path.display().to_string() }.into_report(),
+                            );
+                        }
+                    }
+                }
+            }
+        }
+    }
+
     fn add_files(&mut self, paths: &[PathBuf], reports: &mut ReportCollection) {
         for path in paths {
             if path.is_dir() {
                 // Handle directories on a best effort basis only.
-                let mut paths = Vec::new();
                 if let Ok(entries) = fs::read_dir(path) {
-                    for entry in entries.flatten() {
-                        paths.push(entry.path())
-                    }
+                    let paths: Vec<_> = entries.flatten().map(|x| x.path()).collect();
+                    self.add_files(&paths, reports);
                 }
-                self.add_files(&paths, reports);
             } else if let Some(extension) = path.extension() {
                 // Add Circom files to file stack.
                 if extension == "circom" {
@@ -58,22 +87,56 @@ impl FileStack {
     pub fn add_include(&mut self, include: &Include) -> Result<(), Box<Report>> {
         let mut location = self.current_location.clone().expect("parsing file");
         location.push(include.path.clone());
-        match fs::canonicalize(location) {
+        match fs::canonicalize(&location) {
             Ok(path) => {
                 if !self.black_paths.contains(&path) {
+                    debug!("adding local or absolute include `{}`", location.display());
                     self.stack.push(path);
                 }
                 Ok(())
             }
-            Err(_) => {
-                let error = IncludeError {
-                    path: include.path.clone(),
-                    file_id: include.meta.file_id,
-                    file_location: include.meta.file_location(),
-                };
-                Err(Box::new(error.into_report()))
+            Err(_) => self.include_library(include),
+        }
+    }
+
+    fn include_library(&mut self, include: &Include) -> Result<(), Box<Report>> {
+        // try and perform library resolution on the include
+        // at this point any absolute path has been handled by the push in add_include
+        let pathos = OsString::from(include.path.clone());
+        for lib in &self.libraries {
+            if lib.dir {
+                // only match relative paths that do not start with .
+                if include.path.find('.') == Some(0) {
+                    continue;
+                }
+
+                let libpath = lib.path.join(&include.path);
+                debug!("searching for `{}` in `{}`", include.path, lib.path.display());
+                if fs::canonicalize(&libpath).is_ok() {
+                    debug!("adding include `{}` from directory", libpath.display());
+                    self.stack.push(libpath);
+                    return Ok(());
+                }
+            } else {
+                // only match include paths with a single component i.e. lib.circom and not dir/lib.circom or
+                // ./lib.circom
+                if include.path.find(std::path::MAIN_SEPARATOR) == None {
+                    debug!("checking if `{}` matches `{}`", include.path, lib.path.display());
+                    if lib.path.file_name().expect("good library file") == pathos {
+                        debug!("adding include `{}` from file", lib.path.display());
+                        self.stack.push(lib.path.clone());
+                        return Ok(());
+                    }
+                }
             }
         }
+
+        let error = IncludeError {
+            path: include.path.clone(),
+            file_id: include.meta.file_id,
+            file_location: include.meta.file_location(),
+        };
+        Err(Box::new(error.into_report()))
     }
 
     pub fn take_next(&mut self) -> Option<PathBuf> {

parser/src/lib.rs

@@ -36,9 +36,13 @@ pub enum ParseResult {
     Library(Box<TemplateLibrary>, ReportCollection),
 }
 
-pub fn parse_files(file_paths: &[PathBuf], compiler_version: &Version) -> ParseResult {
+pub fn parse_files(
+    file_paths: &[PathBuf],
+    libraries: &[PathBuf],
+    compiler_version: &Version,
+) -> ParseResult {
     let mut reports = ReportCollection::new();
-    let mut file_stack = FileStack::new(file_paths, &mut reports);
+    let mut file_stack = FileStack::new(file_paths, libraries, &mut reports);
     let mut file_library = FileLibrary::new();
     let mut definitions = HashMap::new();
     let mut main_components = Vec::new();
@@ -48,6 +52,11 @@ pub fn parse_files(file_paths: &[PathBuf], compiler_version: &Version) -> ParseR
                 if let Some(main_component) = program.main_component {
                     main_components.push((file_id, main_component, program.custom_gates));
                 }
+                debug!(
+                    "adding {} definitions from `{}`",
+                    program.definitions.iter().map(|x| x.name()).collect::<Vec<_>>().join(", "),
+                    file_path.display(),
+                );
                 definitions.insert(file_id, program.definitions);
                 reports.append(&mut warnings);
             }

program_analysis/src/analysis_runner.rs

@@ -30,6 +30,7 @@ type ReportCache = HashMap<String, ReportCollection>;
 #[derive(Default)]
 pub struct AnalysisRunner {
     curve: Curve,
+    libraries: Vec<PathBuf>,
     /// The corresponding file library including file includes.
     file_library: FileLibrary,
     /// Template ASTs generated by the parser.
@@ -51,21 +52,27 @@ impl AnalysisRunner {
         AnalysisRunner { curve, ..Default::default() }
     }
 
+    pub fn with_libraries(mut self, libraries: &[PathBuf]) -> Self {
+        self.libraries.extend_from_slice(libraries);
+        self
+    }
+
     pub fn with_files(mut self, input_files: &[PathBuf]) -> (Self, ReportCollection) {
-        let reports = match parser::parse_files(input_files, &config::COMPILER_VERSION) {
-            ParseResult::Program(program, warnings) => {
-                self.template_asts = program.templates;
-                self.function_asts = program.functions;
-                self.file_library = program.file_library;
-                warnings
-            }
-            ParseResult::Library(library, warnings) => {
-                self.template_asts = library.templates;
-                self.function_asts = library.functions;
-                self.file_library = library.file_library;
-                warnings
-            }
-        };
+        let reports =
+            match parser::parse_files(input_files, &self.libraries, &config::COMPILER_VERSION) {
+                ParseResult::Program(program, warnings) => {
+                    self.template_asts = program.templates;
+                    self.function_asts = program.functions;
+                    self.file_library = program.file_library;
+                    warnings
+                }
+                ParseResult::Library(library, warnings) => {
+                    self.template_asts = library.templates;
+                    self.function_asts = library.functions;
+                    self.file_library = library.file_library;
+                    warnings
+                }
+            };
         (self, reports)
     }
 

program_structure/src/abstract_syntax_tree/ast.rs

@@ -161,6 +161,15 @@ pub fn build_function(
     Definition::Function { meta, name, args, arg_location, body }
 }
 
+impl Definition {
+    pub fn name(&self) -> String {
+        match self {
+            Self::Template { name, .. } => name.clone(),
+            Self::Function { name, .. } => name.clone(),
+        }
+    }
+}
+
 #[derive(Clone)]
 pub enum Statement {
     IfThenElse {