Preflight invalid program upgrades

crates/leo/src/cli/commands/upgrade.rs

@@ -227,6 +227,13 @@ fn handle_upgrade<N: Network, A: Aleo<Network = N>>(
         command.env_override.network_retries,
     )?;
 
+    // Fail closed for proven-invalid upgrades by using the stricter active endpoint rules; fetch uncertainty remains a warning below.
+    let validation_consensus_version =
+        get_endpoint_consensus_version(&endpoint, network, command.env_override.network_retries)
+            .map_or(consensus_version, |network_version| consensus_version.max(network_version));
+    let remote_programs =
+        validate_upgrade_tasks(&endpoint, network, &local, &skipped, validation_consensus_version, command)?;
+
     // Build the config for JSON output.
     let config = Some(Config {
         address: address.to_string(),
@@ -244,7 +251,7 @@ fn handle_upgrade<N: Network, A: Aleo<Network = N>>(
         &local,
         &skipped,
         &remote,
-        &check_tasks_for_warnings(&endpoint, network, &local, consensus_version, command),
+        &check_tasks_for_warnings(&endpoint, network, &local, &remote_programs, consensus_version, command),
         consensus_version,
         &command.into(),
     );
@@ -467,6 +474,79 @@ fn handle_upgrade<N: Network, A: Aleo<Network = N>>(
     Ok(build_deploy_output(config, &transactions, &all_stats, &all_broadcasts))
 }
 
+fn validate_upgrade_tasks<N: Network>(
+    endpoint: &str,
+    network: NetworkName,
+    tasks: &[Task<N>],
+    skipped: &HashSet<ProgramID<N>>,
+    consensus_version: ConsensusVersion,
+    command: &LeoUpgrade,
+) -> Result<Vec<(ProgramID<N>, Program<N>)>> {
+    let mut remote_programs = Vec::with_capacity(tasks.len());
+
+    for Task { id, program, is_local, .. } in tasks {
+        // A proven-invalid upgrade is rejected before transaction construction in every output mode.
+        if !is_local || skipped.contains(id) {
+            continue;
+        }
+
+        let Ok(remote_program) =
+            fetch_program_from_network(&id.to_string(), endpoint, network, command.env_override.network_retries)
+        else {
+            // Fetch uncertainty remains in `check_tasks_for_warnings` so transient network state does not block planning.
+            continue;
+        };
+        let Ok(remote_program) = Program::<N>::from_str(&remote_program) else {
+            continue;
+        };
+        reject_invalid_upgrade(id, &remote_program, program, consensus_version)?;
+        remote_programs.push((*id, remote_program));
+    }
+
+    Ok(remote_programs)
+}
+
+fn get_endpoint_consensus_version(
+    endpoint: &str,
+    network: NetworkName,
+    network_retries: u32,
+) -> Option<ConsensusVersion> {
+    let response = leo_package::fetch_from_network(&format!("{endpoint}/{network}/consensus_version"), network_retries)
+        .ok()?
+        .parse::<u8>()
+        .ok()?;
+
+    number_to_consensus_version(response as usize).ok()
+}
+
+fn reject_invalid_upgrade<N: Network>(
+    id: &ProgramID<N>,
+    remote_program: &Program<N>,
+    new_program: &Program<N>,
+    consensus_version: ConsensusVersion,
+) -> Result<()> {
+    if !remote_program.contains_constructor() {
+        return Ok(());
+    }
+
+    Stack::check_upgrade_is_valid(remote_program, new_program)
+        .and_then(|_| {
+            if consensus_version >= ConsensusVersion::V10 {
+                snarkvm::synthesizer::vm::check_output_register_indices_unchanged(remote_program, new_program)
+            } else {
+                Ok(())
+            }
+        })
+        .map_err(|e| {
+            crate::errors::custom(format!("program '{id}' is not a valid upgrade: {e}"))
+                .with_help(
+                    "Try preserving the original interface and output registers, adding a new function for the changed \
+                    interface, or deploying a new program.",
+                )
+                .into()
+        })
+}
+
 /// Check the tasks to warn the user about any potential issues.
 /// The following properties are checked:
 /// - If the transaction is to be broadcast:
@@ -477,6 +557,7 @@ fn check_tasks_for_warnings<N: Network>(
     endpoint: &str,
     network: NetworkName,
     tasks: &[Task<N>],
+    remote_programs: &[(ProgramID<N>, Program<N>)],
     consensus_version: ConsensusVersion,
     command: &LeoUpgrade,
 ) -> Vec<String> {
@@ -487,7 +568,9 @@ fn check_tasks_for_warnings<N: Network>(
         }
 
         // Check if the program exists on the network.
-        if let Ok(remote_program) =
+        if let Some((_, remote_program)) = remote_programs.iter().find(|(remote_id, _)| remote_id == id) {
+            push_remote_upgrade_warnings(id, remote_program, program, consensus_version, &mut warnings);
+        } else if let Ok(remote_program) =
             fetch_program_from_network(&id.to_string(), endpoint, network, command.env_override.network_retries)
         {
             // Parse the program.
@@ -498,18 +581,7 @@ fn check_tasks_for_warnings<N: Network>(
                     continue;
                 }
             };
-            // Check if the program is a valid upgrade.
-            if remote_program.contains_constructor() {
-                if let Err(e) = Stack::check_upgrade_is_valid(&remote_program, program) {
-                    warnings.push(format!(
-                        "The program '{id}' is not a valid upgrade. The upgrade will likely fail. Error: {e}",
-                    ));
-                }
-            } else if consensus_version >= ConsensusVersion::V8 {
-                warnings.push(format!("The program '{id}' can only ever be upgraded once and its contents cannot be changed. Otherwise, the upgrade will likely fail."));
-            } else {
-                warnings.push(format!("The program '{id}' does not have a constructor and is not eligible for a one-time upgrade (>= `ConsensusVersion::V8`). The upgrade will likely fail."));
-            }
+            push_remote_upgrade_warnings(id, &remote_program, program, consensus_version, &mut warnings);
         } else {
             warnings.push(format!("The program '{id}' does not exist on the network. The upgrade will likely fail.",));
         }
@@ -550,6 +622,25 @@ fn check_tasks_for_warnings<N: Network>(
     warnings
 }
 
+fn push_remote_upgrade_warnings<N: Network>(
+    id: &ProgramID<N>,
+    remote_program: &Program<N>,
+    program: &Program<N>,
+    consensus_version: ConsensusVersion,
+    warnings: &mut Vec<String>,
+) {
+    // Check if the program is a valid upgrade.
+    if remote_program.contains_constructor() {
+        if let Err(e) = reject_invalid_upgrade(id, remote_program, program, consensus_version) {
+            warnings.push(e.to_string());
+        }
+    } else if consensus_version >= ConsensusVersion::V8 {
+        warnings.push(format!("The program '{id}' can only ever be upgraded once and its contents cannot be changed. Otherwise, the upgrade will likely fail."));
+    } else {
+        warnings.push(format!("The program '{id}' does not have a constructor and is not eligible for a one-time upgrade (>= `ConsensusVersion::V8`). The upgrade will likely fail."));
+    }
+}
+
 // Convert the `LeoUpgrade` into a `LeoDeploy` command.
 impl From<&LeoUpgrade> for LeoDeploy {
     fn from(upgrade: &LeoUpgrade) -> Self {
@@ -565,3 +656,153 @@ impl From<&LeoUpgrade> for LeoDeploy {
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use snarkvm::prelude::TestnetV0;
+
+    fn parse_program(source: &str) -> Program<TestnetV0> {
+        Program::<TestnetV0>::from_str(source).unwrap()
+    }
+
+    #[test]
+    fn rejects_upgrade_that_changes_record_output_register() {
+        let original = parse_program(
+            r"program upgrade_check.aleo;
+
+record Token:
+    owner as address.private;
+    amount as u64.private;
+
+function mint:
+    input r0 as address.private;
+    input r1 as u64.private;
+    cast r0 r1 into r2 as Token.record;
+    output r2 as Token.record;
+
+constructor:
+    assert.eq edition 0u16;
+",
+        );
+        let changed_output_register = parse_program(
+            r"program upgrade_check.aleo;
+
+record Token:
+    owner as address.private;
+    amount as u64.private;
+
+function mint:
+    input r0 as address.private;
+    input r1 as u64.private;
+    add r1 0u64 into r2;
+    cast r0 r2 into r3 as Token.record;
+    output r3 as Token.record;
+
+constructor:
+    assert.eq edition 0u16;
+",
+        );
+        let id = original.id();
+
+        let error = reject_invalid_upgrade(id, &original, &changed_output_register, ConsensusVersion::V10)
+            .unwrap_err()
+            .to_string();
+
+        assert!(error.contains("not a valid upgrade"));
+        assert!(error.contains("output register"));
+    }
+
+    #[test]
+    fn accepts_upgrade_that_preserves_record_output_register() {
+        let original = parse_program(
+            r"program upgrade_check.aleo;
+
+record Token:
+    owner as address.private;
+    amount as u64.private;
+
+function mint:
+    input r0 as address.private;
+    input r1 as u64.private;
+    cast r0 r1 into r2 as Token.record;
+    output r2 as Token.record;
+
+constructor:
+    assert.eq edition 0u16;
+",
+        );
+        let same_interface = parse_program(
+            r"program upgrade_check.aleo;
+
+record Token:
+    owner as address.private;
+    amount as u64.private;
+
+function mint:
+    input r0 as address.private;
+    input r1 as u64.private;
+    add r1 1u64 into r3;
+    cast r0 r1 into r2 as Token.record;
+    output r2 as Token.record;
+
+constructor:
+    assert.eq edition 0u16;
+",
+        );
+
+        assert!(reject_invalid_upgrade(original.id(), &original, &same_interface, ConsensusVersion::V10).is_ok());
+    }
+
+    #[test]
+    fn warns_for_upgrade_that_changes_record_output_register() {
+        let original = parse_program(
+            r"program upgrade_check.aleo;
+
+record Token:
+    owner as address.private;
+    amount as u64.private;
+
+function mint:
+    input r0 as address.private;
+    input r1 as u64.private;
+    cast r0 r1 into r2 as Token.record;
+    output r2 as Token.record;
+
+constructor:
+    assert.eq edition 0u16;
+",
+        );
+        let changed_output_register = parse_program(
+            r"program upgrade_check.aleo;
+
+record Token:
+    owner as address.private;
+    amount as u64.private;
+
+function mint:
+    input r0 as address.private;
+    input r1 as u64.private;
+    add r1 0u64 into r2;
+    cast r0 r2 into r3 as Token.record;
+    output r3 as Token.record;
+
+constructor:
+    assert.eq edition 0u16;
+",
+        );
+        let mut warnings = Vec::new();
+
+        push_remote_upgrade_warnings(
+            original.id(),
+            &original,
+            &changed_output_register,
+            ConsensusVersion::V10,
+            &mut warnings,
+        );
+
+        assert_eq!(warnings.len(), 1);
+        assert!(warnings[0].contains("not a valid upgrade"));
+        assert!(warnings[0].contains("Try preserving the original interface and output registers"));
+    }
+}

documentation/leo_by_example/battleship.md

@@ -664,7 +664,7 @@ If a ship is valid vertically or horizontally, then we know the ship is valid. W
 
 ### Bit Counting
 
-See the "c_bitcount" closure to follow along with the code. 50 years ago, MIT AI Laboratory published HAKMEM, which was a series of tricks and hacks to speed up processing for bitwise operations. <https://w3.pppl.gov/~hammett/work/2009/AIM-239-ocr.pdf> We turned to HAKMEM 169 for bitcounting inspiration, although we've tweaked our implementation to be (hopefully) easier to understand. Before diving into details, let's build some intuition.
+See the "c_bitcount" closure to follow along with the code. 50 years ago, MIT AI Laboratory published HAKMEM, which was a series of tricks and hacks to speed up processing for bitwise operations. <https://www.jjj.de/hakmem/hakmem.html> We turned to HAKMEM 169 for bitcounting inspiration, although we've tweaked our implementation to be (hopefully) easier to understand. Before diving into details, let's build some intuition.
 
 Let a,b,c,d be either 0 or 1. Given a polynomial 8a + 4b + 2c + d, how do we find the summation of a + b + c + d? If we subtract subsets of this polynomial, we'll be left with the summation.