Merge pull request #28081 from ProvableHQ/fix/too-many-reqs

[Fix] Network retrieval refactor

compiler/compiler/tests/compile.rs

@@ -140,8 +140,9 @@ fn run_test(test: Test, handler: &Handler, buf: &BufferEmitter) -> Result<Value,
             handler.extend_if_error(process.add_program(&aleo_program).map_err(LeoError::Anyhow))?;
 
             // Add the bytecode to the import stubs.
-            let stub =
-                handler.extend_if_error(disassemble_from_str::<CurrentNetwork>(&bytecode).map_err(|err| err.into()))?;
+            let stub = handler.extend_if_error(
+                disassemble_from_str::<CurrentNetwork>(&program_name, &bytecode).map_err(|err| err.into()),
+            )?;
             import_stubs.insert(Symbol::intern(&program_name), stub);
 
             // Hash the ast files.

compiler/compiler/tests/execute.rs

@@ -184,8 +184,9 @@ fn run_test(test: Test, handler: &Handler, buf: &BufferEmitter) -> Result<Value,
             }
 
             // Add the bytecode to the import stubs.
-            let stub =
-                handler.extend_if_error(disassemble_from_str::<CurrentNetwork>(&bytecode).map_err(|err| err.into()))?;
+            let stub = handler.extend_if_error(
+                disassemble_from_str::<CurrentNetwork>(&program_name, &bytecode).map_err(|err| err.into()),
+            )?;
             import_stubs.insert(Symbol::intern(&program_name), stub);
 
             // Hash the ast files.

errors/src/errors/utils/util_errors.rs

@@ -49,8 +49,8 @@ create_messages!(
 
     @formatted
     snarkvm_parsing_error {
-        args: (),
-        msg: "SnarkVM failure to parse `.aleo` program".to_string(),
+        args: (name: impl Display),
+        msg: format!("Failed to parse the source file for `{name}.aleo` into a valid Aleo program."),
         help: None,
     }
 

leo/cli/commands/query/mod.rs

@@ -15,6 +15,7 @@
 // along with the Leo library. If not, see <https://www.gnu.org/licenses/>.
 
 use super::*;
+use snarkvm::prelude::{MainnetV0, TestnetV0};
 
 mod block;
 use block::Block;
@@ -41,6 +42,7 @@ mod utils;
 use utils::*;
 
 use leo_errors::UtilError;
+use leo_retriever::{fetch_from_network, verify_valid_program, NetworkName};
 
 ///  Query live data from the Aleo network.
 #[derive(Parser, Debug)]
@@ -72,49 +74,61 @@ impl Command for Query {
     }
 
     fn apply(self, context: Context, _: Self::Input) -> Result<Self::Output> {
-        let recursive = context.recursive;
-        let output = match self.command {
-            QueryCommands::Block { command } => command.apply(context, ())?,
-            QueryCommands::Transaction { command } => command.apply(context, ())?,
-            QueryCommands::Program { command } => command.apply(context, ())?,
-            QueryCommands::Stateroot { command } => command.apply(context, ())?,
-            QueryCommands::Committee { command } => command.apply(context, ())?,
-            QueryCommands::Mempool { command } => {
-                if self.endpoint == "https://api.explorer.aleo.org/v1" {
-                    tracing::warn!(
-                        "⚠️  `leo query mempool` is only valid when using a custom endpoint. Specify one using `--endpoint`."
-                    );
-                }
-                command.apply(context, ())?
-            }
-            QueryCommands::Peers { command } => {
-                if self.endpoint == "https://api.explorer.aleo.org/v1" {
-                    tracing::warn!(
-                        "⚠️  `leo query peers` is only valid when using a custom endpoint. Specify one using `--endpoint`."
-                    );
-                }
-                command.apply(context, ())?
+        // Parse the network.
+        let network = NetworkName::try_from(self.network.as_str())?;
+        match network {
+            NetworkName::MainnetV0 => handle_query::<MainnetV0>(self, context),
+            NetworkName::TestnetV0 => handle_query::<TestnetV0>(self, context),
+        }
+    }
+}
+
+// A helper function to handle the `query` command.
+fn handle_query<N: Network>(query: Query, context: Context) -> Result<<Query as Command>::Output> {
+    let recursive = context.recursive;
+    let (program, output) = match query.command {
+        QueryCommands::Block { command } => (None, command.apply(context, ())?),
+        QueryCommands::Transaction { command } => (None, command.apply(context, ())?),
+        QueryCommands::Program { command } => {
+            // Check if querying for program source code.
+            let program =
+                if command.mappings || command.mapping_value.is_some() { None } else { Some(command.name.clone()) };
+            (program, command.apply(context, ())?)
+        }
+        QueryCommands::Stateroot { command } => (None, command.apply(context, ())?),
+        QueryCommands::Committee { command } => (None, command.apply(context, ())?),
+        QueryCommands::Mempool { command } => {
+            if query.endpoint == "https://api.explorer.aleo.org/v1" {
+                tracing::warn!(
+                    "⚠️  `leo query mempool` is only valid when using a custom endpoint. Specify one using `--endpoint`."
+                );
             }
-        };
-
-        // Make GET request to retrieve on-chain state.
-        let url = format!("{}/{}/{}", self.endpoint, self.network, output);
-        let response = ureq::get(&url.clone())
-            .set(&format!("X-Aleo-Leo-{}", env!("CARGO_PKG_VERSION")), "true")
-            .call()
-            .map_err(|err| UtilError::failed_to_retrieve_from_endpoint(err, Default::default()))?;
-        if response.status() == 200 {
-            // Unescape the newlines.
-            let result = response.into_string().unwrap().replace("\\n", "\n").replace('\"', "");
-            if !recursive {
-                tracing::info!("✅ Successfully retrieved data from '{url}'.\n");
-                println!("{}\n", result);
+            (None, command.apply(context, ())?)
+        }
+        QueryCommands::Peers { command } => {
+            if query.endpoint == "https://api.explorer.aleo.org/v1" {
+                tracing::warn!(
+                    "⚠️  `leo query peers` is only valid when using a custom endpoint. Specify one using `--endpoint`."
+                );
             }
-            Ok(result)
-        } else {
-            Err(UtilError::network_error(url, response.status(), Default::default()).into())
+            (None, command.apply(context, ())?)
         }
+    };
+
+    // Make GET request to retrieve on-chain state.
+    let url = format!("{}/{}/{output}", query.endpoint, query.network);
+    let result = fetch_from_network(&url)?;
+    if !recursive {
+        tracing::info!("✅ Successfully retrieved data from '{url}'.\n");
+        println!("{}\n", result);
     }
+
+    // Verify that the source file parses into a valid Aleo program.
+    if let Some(name) = program {
+        verify_valid_program::<N>(&name, &result)?;
+    }
+
+    Ok(result)
 }
 
 #[derive(Parser, Debug)]

utils/disassembler/src/lib.rs

@@ -86,10 +86,10 @@ pub fn disassemble<N: Network, Instruction: InstructionTrait<N>, Command: Comman
     }
 }
 
-pub fn disassemble_from_str<N: Network>(program: &str) -> Result<Stub, UtilError> {
+pub fn disassemble_from_str<N: Network>(name: &str, program: &str) -> Result<Stub, UtilError> {
     match Program::<N>::from_str(program) {
         Ok(p) => Ok(disassemble(p)),
-        Err(_) => Err(UtilError::snarkvm_parsing_error(Default::default())),
+        Err(_) => Err(UtilError::snarkvm_parsing_error(name, Default::default())),
     }
 }
 
@@ -124,7 +124,8 @@ mod tests {
         create_session_if_not_set_then(|_| {
             let program_from_file =
                 fs::read_to_string("../tmp/.aleo/registry/mainnet/zk_bitwise_stack_v0_0_2.aleo").unwrap();
-            let _program = disassemble_from_str::<CurrentNetwork>(&program_from_file).unwrap();
+            let _program =
+                disassemble_from_str::<CurrentNetwork>("zk_bitwise_stack_v0_0_2", &program_from_file).unwrap();
         });
     }
 }

utils/retriever/src/retriever/mod.rs

@@ -22,7 +22,7 @@ use leo_errors::UtilError;
 use leo_passes::{common::DiGraph, DiGraphError};
 use leo_span::Symbol;
 
-use snarkvm::prelude::Network;
+use snarkvm::prelude::{Network, Program};
 
 use indexmap::{IndexMap, IndexSet};
 use std::{
@@ -31,6 +31,7 @@ use std::{
     io::Read,
     marker::PhantomData,
     path::{Path, PathBuf},
+    str::FromStr,
 };
 
 // Retriever is responsible for retrieving external programs
@@ -309,7 +310,7 @@ impl<N: Network> Retriever<N> {
             })?;
 
             // Cache the disassembled stub
-            let stub: Stub = disassemble_from_str::<N>(&content)?;
+            let stub: Stub = disassemble_from_str::<N>(&name.to_string(), &content)?;
             if cur_context.add_stub(stub.clone()) {
                 Err(UtilError::duplicate_dependency_name_error(stub.stub_id.name.name, Default::default()))?;
             }
@@ -439,7 +440,7 @@ fn retrieve_from_network<N: Network>(
     // Check if the file is already cached in `~/.aleo/registry/{network}/{program}`
     let move_to_path = home_path.join(network.to_string());
     let path = move_to_path.join(name.clone());
-    let mut file_str: String;
+    let file_str: String;
     if !path.exists() {
         // Create directories along the way if they don't exist
         std::fs::create_dir_all(&move_to_path).map_err(|err| {
@@ -450,13 +451,10 @@ fn retrieve_from_network<N: Network>(
             )
         })?;
 
-        // Construct the endpoint for the network.
-        let endpoint = format!("{endpoint}/{network}");
-
         // Fetch from network
-        println!("Retrieving {name} from {endpoint}.");
-        file_str = fetch_from_network(&endpoint, name)?;
-        file_str = file_str.replace("\\n", "\n").replace('\"', "");
+        println!("Retrieving {name} from {endpoint} on {network}.");
+        file_str = fetch_from_network(&format!("{endpoint}/{network}/program/{}", &name))?;
+        verify_valid_program::<N>(name, &file_str)?;
         println!("Successfully retrieved {} from {:?}!", name, endpoint);
 
         // Write file to cache
@@ -498,7 +496,7 @@ fn retrieve_from_network<N: Network>(
     })?;
 
     // Disassemble into Stub
-    let stub: Stub = disassemble_from_str::<N>(&file_str)?;
+    let stub: Stub = disassemble_from_str::<N>(name, &file_str)?;
 
     // Create entry for leo.lock
     Ok((
@@ -518,14 +516,23 @@ fn retrieve_from_network<N: Network>(
     ))
 }
 
-fn fetch_from_network(endpoint: &String, program: &String) -> Result<String, UtilError> {
-    let url = format!("{}/program/{}", endpoint, program);
-    let response = ureq::get(&url.clone())
+// Fetch the given endpoint url and return the sanitized response.
+pub fn fetch_from_network(url: &str) -> Result<String, UtilError> {
+    let response = ureq::get(url)
+        .set(&format!("X-Aleo-Leo-{}", env!("CARGO_PKG_VERSION")), "true")
         .call()
         .map_err(|err| UtilError::failed_to_retrieve_from_endpoint(err, Default::default()))?;
     if response.status() == 200 {
-        Ok(response.into_string().unwrap())
+        Ok(response.into_string().unwrap().replace("\\n", "\n").replace('\"', ""))
     } else {
         Err(UtilError::network_error(url, response.status(), Default::default()))
     }
 }
+
+// Verify that a fetched program is valid aleo instructions.
+pub fn verify_valid_program<N: Network>(name: &str, program: &str) -> Result<(), UtilError> {
+    match Program::<N>::from_str(program) {
+        Ok(_) => Ok(()),
+        Err(_) => Err(UtilError::snarkvm_parsing_error(name, Default::default())),
+    }
+}