feat(nns): Support upgrading root canister through install_code (#396)

# Why

Previously, `install_code` proposal action is added without supporting
upgrading root. This PR supports it.

# Changes

We support creating a proposal with action `install_code` while
specifying NNS Root as the target canister, and the proposal execution
code will call the correct `upgrade_root` endpoint correctly (unlike
upgrading other canisters which calls root). Note that `install` or
`reinstall` modes aren't supported. It might not be obvious why
`reinstall` isn't supported - HardResetRootToVersion supports
"resetting" root, but the reset is implemented as an `uninstall_code`
followed by `install_code` (with `Install` mode). This behavior seems to
exceed what can be reasonably thought of as the `install_code` proposal
action. Therefore, `HardResetRootToVersion` will remain as an
NnsFunction and will later be switched to `ProtocolCanisterManagement`
topic.

rs/nns/governance/src/proposals/install_code.rs

@@ -6,7 +6,7 @@ use crate::{
     proposals::call_canister::CallCanister,
 };
 
-use candid::Encode;
+use candid::{CandidType, Deserialize, Encode};
 use ic_base_types::CanisterId;
 use ic_management_canister_types::CanisterInstallMode as RootCanisterInstallMode;
 use ic_nervous_system_root::change_canister::ChangeCanisterRequest;
@@ -16,6 +16,7 @@ use ic_nns_constants::{
     LEDGER_CANISTER_ID, LEDGER_INDEX_CANISTER_ID, LIFELINE_CANISTER_ID, REGISTRY_CANISTER_ID,
     ROOT_CANISTER_ID, SUBNET_RENTAL_CANISTER_ID,
 };
+use serde::Serialize;
 
 const PROTOCOL_CANISTER_IDS: [&CanisterId; 12] = [
     &REGISTRY_CANISTER_ID,
@@ -32,6 +33,15 @@ const PROTOCOL_CANISTER_IDS: [&CanisterId; 12] = [
     &EXCHANGE_RATE_CANISTER_ID,
 ];
 
+// When calling lifeline's upgrade_root method, this is the request. Keep this in sync with
+// `rs/nns/handlers/lifeline/impl/lifeline.mo`.
+#[derive(Clone, Debug, PartialEq, Eq, CandidType, Deserialize, Serialize)]
+struct UpgradeRootProposalPayload {
+    wasm_module: Vec<u8>,
+    module_arg: Vec<u8>,
+    stop_upgrade_start: bool,
+}
+
 impl InstallCode {
     pub fn validate(&self) -> Result<(), GovernanceError> {
         if !cfg!(feature = "test") {
@@ -43,11 +53,12 @@ impl InstallCode {
 
         let _ = self.valid_canister_id()?;
         let _ = self.valid_install_mode()?;
+        let _ = self.valid_wasm_module()?;
         let _ = self.valid_topic()?;
+        let _ = self.canister_and_function()?;
 
-        if self.wasm_module.is_none() {
-            return Err(Self::invalid_proposal_error("Wasm module is required"));
-        }
+        // In the future, we could potentially validate the wasm module to see if it's a valid gzip
+        // or a valid WASM.
 
         Ok(())
     }
@@ -58,14 +69,7 @@ impl InstallCode {
             .ok_or(Self::invalid_proposal_error("Canister ID is required"))?;
         let canister_id = CanisterId::try_from(canister_principal_id)
             .map_err(|_| Self::invalid_proposal_error("Invalid canister ID"))?;
-        if canister_id == ROOT_CANISTER_ID {
-            // TODO(NNS1-3190): support changing root canister
-            Err(Self::invalid_proposal_error(
-                "InstallCode for root canister is not supported yet",
-            ))
-        } else {
-            Ok(canister_id)
-        }
+        Ok(canister_id)
     }
 
     fn valid_install_mode(&self) -> Result<RootCanisterInstallMode, GovernanceError> {
@@ -85,14 +89,23 @@ impl InstallCode {
         }
     }
 
+    fn valid_wasm_module(&self) -> Result<&Vec<u8>, GovernanceError> {
+        // We do not want to copy the (potentially large) wasm module when validating, so we return
+        // a reference and let the caller clone it if needed.
+        self.wasm_module
+            .as_ref()
+            .ok_or(Self::invalid_proposal_error("Wasm module is required"))
+    }
+
     pub fn valid_topic(&self) -> Result<Topic, GovernanceError> {
         let canister_id = self.valid_canister_id()?;
         if PROTOCOL_CANISTER_IDS.contains(&&canister_id) {
             Ok(Topic::ProtocolCanisterManagement)
         } else {
-            Err(Self::invalid_proposal_error(
-                "Canister ID is not a protocol canister",
-            ))
+            Err(Self::invalid_proposal_error(&format!(
+                "Canister id {:?} is not a protocol canister",
+                canister_id
+            )))
         }
     }
 
@@ -102,28 +115,25 @@ impl InstallCode {
             format!("InstallCode proposal invalid because of {}", reason),
         )
     }
-}
 
-impl CallCanister for InstallCode {
-    fn canister_and_function(&self) -> Result<(CanisterId, &str), GovernanceError> {
-        let canister_id = self.valid_canister_id()?;
-        if canister_id == ROOT_CANISTER_ID {
-            // TODO(NNS1-3190): support changing root canister
-            return Err(Self::invalid_proposal_error(
-                "InstallCode for root canister is not supported yet",
-            ));
-        }
-        Ok((ROOT_CANISTER_ID, "change_nns_canister"))
+    fn payload_to_upgrade_root(&self) -> Result<Vec<u8>, GovernanceError> {
+        let stop_upgrade_start = !self.skip_stopping_before_installing.unwrap_or(false);
+        let wasm_module = self.valid_wasm_module()?.clone();
+        let module_arg = self.arg.clone().unwrap_or_default();
+
+        Encode!(&UpgradeRootProposalPayload {
+            stop_upgrade_start,
+            wasm_module,
+            module_arg,
+        })
+        .map_err(|e| Self::invalid_proposal_error(&format!("Failed to encode payload: {}", e)))
     }
 
-    fn payload(&self) -> Result<Vec<u8>, GovernanceError> {
+    fn payload_to_upgrade_non_root(&self) -> Result<Vec<u8>, GovernanceError> {
         let stop_before_installing = !self.skip_stopping_before_installing.unwrap_or(false);
         let mode = self.valid_install_mode()?;
         let canister_id = self.valid_canister_id()?;
-        let wasm_module = self
-            .wasm_module
-            .clone()
-            .ok_or(Self::invalid_proposal_error("Wasm module is required"))?;
+        let wasm_module = self.valid_wasm_module()?.clone();
         let arg = self.arg.clone().unwrap_or_default();
         let compute_allocation = None;
         let memory_allocation = None;
@@ -141,6 +151,48 @@ impl CallCanister for InstallCode {
     }
 }
 
+impl CallCanister for InstallCode {
+    fn canister_and_function(&self) -> Result<(CanisterId, &str), GovernanceError> {
+        let canister_id = self.valid_canister_id()?;
+        // Most canisters are upgraded indirectly via root. In such cases, we call root's
+        // change_nns_canister method. The exception is when root is to be upgraded. In that case,
+        // upgrades are instead done via lifeline's upgrade_root method.
+        if canister_id != ROOT_CANISTER_ID {
+            return Ok((ROOT_CANISTER_ID, "change_nns_canister"));
+        }
+
+        let install_mode = self.valid_install_mode()?;
+        match install_mode {
+            RootCanisterInstallMode::Install | RootCanisterInstallMode::Reinstall => {
+                // We can potentially support those modes in the future by extending what the
+                // lifeline canister can do. However there is no reason to do so currently: (1) the
+                // install mode is only useful when root does not have any code, which we don't
+                // expect to happen. (2) as the root canister does not have state, there is no
+                // reason to do reinstall instead of upgrade; for getting out of open call context
+                // problems, only uninstalling and reinstalling the root canister would help
+                // (uninstall cancels open calls), and that is achieved by
+                // HardResetNnsRootToVersion.
+                Err(Self::invalid_proposal_error(&format!(
+                    "InstallCode mode {:?} is not supported for root canister, consider using \
+                     HardResetNnsRootToVersion proposal instead",
+                    install_mode
+                )))
+            }
+            RootCanisterInstallMode::Upgrade => Ok((LIFELINE_CANISTER_ID, "upgrade_root")),
+        }
+    }
+
+    fn payload(&self) -> Result<Vec<u8>, GovernanceError> {
+        let canister_id = self.valid_canister_id()?;
+
+        if canister_id == ROOT_CANISTER_ID {
+            self.payload_to_upgrade_root()
+        } else {
+            self.payload_to_upgrade_non_root()
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -179,87 +231,92 @@ mod tests {
             skip_stopping_before_installing: None,
         };
 
-        assert_eq!(
+        let is_invalid_proposal_with_keywords = |install_code: InstallCode, keywords: Vec<&str>| {
+            let error = install_code.validate().unwrap_err();
+            assert_eq!(error.error_type, ErrorType::InvalidProposal as i32);
+            for keyword in keywords {
+                let error_message = error.error_message.to_lowercase();
+                assert!(
+                    error_message.contains(keyword),
+                    "{} not found in {:#?}",
+                    keyword,
+                    error_message
+                );
+            }
+        };
+
+        is_invalid_proposal_with_keywords(
             InstallCode {
                 canister_id: None,
                 ..valid_install_code.clone()
-            }
-            .validate(),
-            Err(GovernanceError::new_with_message(
-                ErrorType::InvalidProposal,
-                "InstallCode proposal invalid because of Canister ID is required".to_string(),
-            ))
+            },
+            vec!["canister id", "required"],
         );
-        assert_eq!(
+
+        is_invalid_proposal_with_keywords(
             InstallCode {
                 install_mode: None,
                 ..valid_install_code.clone()
-            }
-            .validate(),
-            Err(GovernanceError::new_with_message(
-                ErrorType::InvalidProposal,
-                "InstallCode proposal invalid because of Install mode is required".to_string(),
-            ))
+            },
+            vec!["install mode", "required"],
         );
-        assert_eq!(
+
+        is_invalid_proposal_with_keywords(
             InstallCode {
                 install_mode: Some(1000),
                 ..valid_install_code.clone()
-            }
-            .validate(),
-            Err(GovernanceError::new_with_message(
-                ErrorType::InvalidProposal,
-                "InstallCode proposal invalid because of Unspecified or unrecognized install mode"
-                    .to_string(),
-            ))
+            },
+            vec!["unspecified or unrecognized", "install mode"],
         );
-        assert_eq!(
+
+        is_invalid_proposal_with_keywords(
             InstallCode {
                 install_mode: Some(CanisterInstallMode::Unspecified as i32),
                 ..valid_install_code.clone()
-            }
-            .validate(),
-            Err(GovernanceError::new_with_message(
-                ErrorType::InvalidProposal,
-                "InstallCode proposal invalid because of Unspecified or unrecognized install mode"
-                    .to_string(),
-            ))
+            },
+            vec!["unspecified or unrecognized", "install mode"],
         );
-        assert_eq!(
+
+        is_invalid_proposal_with_keywords(
             InstallCode {
                 wasm_module: None,
                 ..valid_install_code.clone()
-            }
-            .validate(),
-            Err(GovernanceError::new_with_message(
-                ErrorType::InvalidProposal,
-                "InstallCode proposal invalid because of Wasm module is required".to_string(),
-            ))
+            },
+            vec!["wasm module", "required"],
         );
-        assert_eq!(
+
+        is_invalid_proposal_with_keywords(
             InstallCode {
                 canister_id: Some(ROOT_CANISTER_ID.get()),
+                install_mode: Some(CanisterInstallMode::Install as i32),
                 ..valid_install_code.clone()
-            }
-            .validate(),
-            Err(GovernanceError::new_with_message(
-                ErrorType::InvalidProposal,
-                "InstallCode proposal invalid because of InstallCode for root canister is not \
-                 supported yet"
-                    .to_string(),
-            ))
+            },
+            vec![
+                "installcode mode install",
+                "not supported for root canister",
+                "hardresetnnsroottoversion",
+            ],
         );
-        assert_eq!(
+
+        is_invalid_proposal_with_keywords(
+            InstallCode {
+                canister_id: Some(ROOT_CANISTER_ID.get()),
+                install_mode: Some(CanisterInstallMode::Reinstall as i32),
+                ..valid_install_code.clone()
+            },
+            vec![
+                "installcode mode reinstall",
+                "not supported for root canister",
+                "hardresetnnsroottoversion",
+            ],
+        );
+
+        is_invalid_proposal_with_keywords(
             InstallCode {
                 canister_id: Some(ic_nns_constants::SNS_WASM_CANISTER_ID.get()),
                 ..valid_install_code.clone()
-            }
-            .validate(),
-            Err(GovernanceError::new_with_message(
-                ErrorType::InvalidProposal,
-                "InstallCode proposal invalid because of Canister ID is not a protocol canister"
-                    .to_string(),
-            ))
+            },
+            vec!["canister id", "not a protocol canister"],
         );
     }
 
@@ -299,6 +356,38 @@ mod tests {
         );
     }
 
+    #[cfg(feature = "test")]
+    #[test]
+    fn test_upgrade_root_protocol_canister() {
+        let install_code = InstallCode {
+            canister_id: Some(ROOT_CANISTER_ID.get()),
+            wasm_module: Some(vec![1, 2, 3]),
+            install_mode: Some(CanisterInstallMode::Upgrade as i32),
+            arg: Some(vec![4, 5, 6]),
+            skip_stopping_before_installing: None,
+        };
+
+        assert_eq!(install_code.validate(), Ok(()));
+        assert_eq!(
+            install_code.valid_topic(),
+            Ok(Topic::ProtocolCanisterManagement)
+        );
+        assert_eq!(
+            install_code.canister_and_function(),
+            Ok((LIFELINE_CANISTER_ID, "upgrade_root"))
+        );
+        let decoded_payload =
+            Decode!(&install_code.payload().unwrap(), UpgradeRootProposalPayload).unwrap();
+        assert_eq!(
+            decoded_payload,
+            UpgradeRootProposalPayload {
+                stop_upgrade_start: true,
+                wasm_module: vec![1, 2, 3],
+                module_arg: vec![4, 5, 6],
+            }
+        );
+    }
+
     #[cfg(feature = "test")]
     #[test]
     fn test_reinstall_code_non_root_protocol_canister() {

rs/nns/integration_tests/src/canister_upgrade.rs

@@ -52,8 +52,8 @@ fn test_upgrade_canister(canister_id: CanisterId, canister_wasm: Wasm, use_propo
 fn upgrade_canisters_by_proposal() {
     test_upgrade_canister(GOVERNANCE_CANISTER_ID, build_governance_wasm(), true);
     test_upgrade_canister(GOVERNANCE_CANISTER_ID, build_governance_wasm(), false);
-    // TODO(NNS1-3190): Test upgrading root with proposal action when it's supported.
     test_upgrade_canister(ROOT_CANISTER_ID, build_root_wasm(), false);
+    test_upgrade_canister(ROOT_CANISTER_ID, build_root_wasm(), true);
     test_upgrade_canister(LIFELINE_CANISTER_ID, build_lifeline_wasm(), true);
     test_upgrade_canister(LIFELINE_CANISTER_ID, build_lifeline_wasm(), false);
 }