refactor(nns): Some refactoring on how neuron subaccount addresses ar…

…e calculated (#1317) # Why The logic and documentation around NNS neuron subaccount addresses are confusing. # What * Update the link to dfinity_wallet to ic-js. * Refactor the calculation logic for different neuron subaccounts to using the same function. * Move the pointer to the ic-js link to a new test for the staking account.
dfinity · Sep 6, 2024 · 36c1976 · 36c1976
1 parent cd153a6
commit 36c1976
Show file tree

Hide file tree

Showing 3 changed files with 55 additions and 38 deletions.
diff --git a/rs/nervous_system/common/src/ledger.rs b/rs/nervous_system/common/src/ledger.rs
@@ -209,36 +209,38 @@ impl IcpLedger for IcpLedgerCanister {
 /// Computes the bytes of the subaccount to which neuron staking transfers are made. This
 /// function must be kept in sync with the Nervous System UI equivalent.
 pub fn compute_neuron_staking_subaccount_bytes(controller: PrincipalId, nonce: u64) -> [u8; 32] {
-    // The equivalent function in the NNS UI is
-    // https://github.com/dfinity/dfinity_wallet/blob/351e07d3e6d007b090117161a94ce8ec9d5a6b49/js-agent/src/canisters/createNeuron.ts#L63
-    const DOMAIN: &[u8] = b"neuron-stake";
-    const DOMAIN_LENGTH: [u8; 1] = [0x0c];
-
-    let mut hasher = Sha256::new();
-    hasher.write(&DOMAIN_LENGTH);
-    hasher.write(DOMAIN);
-    hasher.write(controller.as_slice());
-    hasher.write(&nonce.to_be_bytes());
-    hasher.finish()
+    compute_neuron_domain_subaccount_bytes(controller, b"neuron-stake", nonce)
 }
 
 /// Computes the subaccount to which neuron staking transfers are made. This
 /// function must be kept in sync with the Nervous System UI equivalent.
 pub fn compute_neuron_staking_subaccount(controller: PrincipalId, nonce: u64) -> IcpSubaccount {
-    // The equivalent function in the NNS UI is
-    // https://github.com/dfinity/dfinity_wallet/blob/351e07d3e6d007b090117161a94ce8ec9d5a6b49/js-agent/src/canisters/createNeuron.ts#L63
     IcpSubaccount(compute_neuron_staking_subaccount_bytes(controller, nonce))
 }
 
 /// Computes the subaccount to which locked token distributions are initialized to.
 pub fn compute_distribution_subaccount_bytes(principal_id: PrincipalId, nonce: u64) -> [u8; 32] {
-    const DOMAIN: &[u8] = b"token-distribution";
-    const DOMAIN_LENGTH: [u8; 1] = [0x12];
+    compute_neuron_domain_subaccount_bytes(principal_id, b"token-distribution", nonce)
+}
+
+// Computes the subaccount to which neuron disburse transfers are made.
+pub fn compute_neuron_disburse_subaccount_bytes(controller: PrincipalId, nonce: u64) -> [u8; 32] {
+    // The "domain" for neuron disburse was unfortunately chosen to be "neuron-split". It might be
+    // possible to change to a more meaningful name, but there is no strong reason to do so, and
+    // there is some risk that this behavior is depended on.
+    compute_neuron_domain_subaccount_bytes(controller, b"neuron-split", nonce)
+}
 
+fn compute_neuron_domain_subaccount_bytes(
+    controller: PrincipalId,
+    domain: &[u8],
+    nonce: u64,
+) -> [u8; 32] {
+    let domain_length: [u8; 1] = [domain.len() as u8];
     let mut hasher = Sha256::new();
-    hasher.write(&DOMAIN_LENGTH);
-    hasher.write(DOMAIN);
-    hasher.write(principal_id.as_slice());
+    hasher.write(&domain_length);
+    hasher.write(domain);
+    hasher.write(controller.as_slice());
     hasher.write(&nonce.to_be_bytes());
     hasher.finish()
 }

diff --git a/rs/nervous_system/common/src/tests.rs b/rs/nervous_system/common/src/tests.rs
@@ -1,5 +1,9 @@
 use super::*;
 
+use crate::ledger::compute_neuron_staking_subaccount_bytes;
+use ic_base_types::PrincipalId;
+use ic_crypto_sha2::Sha256;
+
 #[test]
 fn test_wide_range_of_u64_values() {
     assert!(WIDE_RANGE_OF_U64_VALUES.contains(&0));
@@ -22,3 +26,23 @@ fn test_e8s_to_tokens() {
         );
     }
 }
+
+#[test]
+fn test_compute_neuron_staking_subaccount_bytes() {
+    let principal_id = PrincipalId::new_user_test_id(1);
+    let nonce = 42u64;
+
+    // The equivalent implementation in the ic-js is at
+    // https://github.com/dfinity/ic-js/blob/0dd5c1954d94dad6911b73707c454f978624f607/packages/nns/src/governance.canister.ts#L952-L967.
+    let mut hasher = Sha256::new();
+    hasher.write(&[0x0c]);
+    hasher.write(b"neuron-stake");
+    hasher.write(principal_id.as_slice());
+    hasher.write(&nonce.to_be_bytes());
+    let hash = hasher.finish();
+
+    assert_eq!(
+        compute_neuron_staking_subaccount_bytes(principal_id, nonce),
+        hash
+    );
+}
diff --git a/rs/nns/governance/src/governance.rs b/rs/nns/governance/src/governance.rs
@@ -73,7 +73,6 @@ use dfn_core::api::spawn;
 use dfn_core::println;
 use dfn_protobuf::ToProto;
 use ic_base_types::{CanisterId, PrincipalId};
-use ic_crypto_sha2::Sha256;
 use ic_nervous_system_common::{
     cmc::CMC, ledger, ledger::IcpLedger, NervousSystemError, ONE_DAY_SECONDS, ONE_MONTH_SECONDS,
     ONE_YEAR_SECONDS,
@@ -3284,16 +3283,12 @@ impl Governance {
 
         // Validate that if a child neuron controller was provided, it is a valid
         // principal.
-        let child_controller = &disburse_to_neuron
-            .new_controller
-            .as_ref()
-            .ok_or_else(|| {
-                GovernanceError::new_with_message(
-                    ErrorType::InvalidCommand,
-                    "Must specify a new controller for disburse to neuron.",
-                )
-            })?
-            .clone();
+        let child_controller = disburse_to_neuron.new_controller.ok_or_else(|| {
+            GovernanceError::new_with_message(
+                ErrorType::InvalidCommand,
+                "Must specify a new controller for disburse to neuron.",
+            )
+        })?;
 
         let child_nid = self.neuron_store.new_neuron_id(&mut *self.env);
         let from_subaccount = parent_neuron.subaccount();
@@ -3302,14 +3297,10 @@ impl Governance {
         // the owner on the ledger. There is no need to length-prefix the
         // principal since the nonce is constant length, and so there is no risk
         // of ambiguity.
-        let to_subaccount = Subaccount({
-            let mut state = Sha256::new();
-            state.write(&[0x0c]);
-            state.write(b"neuron-split");
-            state.write(child_controller.as_slice());
-            state.write(&disburse_to_neuron.nonce.to_be_bytes());
-            state.finish()
-        });
+        let to_subaccount = Subaccount(ledger::compute_neuron_disburse_subaccount_bytes(
+            child_controller,
+            disburse_to_neuron.nonce,
+        ));
 
         // Make sure there isn't already a neuron with the same sub-account.
         if self.neuron_store.has_neuron_with_subaccount(to_subaccount) {
@@ -3344,7 +3335,7 @@ impl Governance {
         let child_neuron = NeuronBuilder::new(
             child_nid,
             to_subaccount,
-            *child_controller,
+            child_controller,
             DissolveStateAndAge::NotDissolving {
                 dissolve_delay_seconds,
                 aging_since_timestamp_seconds: created_timestamp_seconds,