Expose memoToNeuronSubaccount and memoToNeuronAccountIdentifier

dfinity · Nov 1, 2024 · 06ad654 · 06ad654
1 parent 07d436b
commit 06ad654
Show file tree

Hide file tree

Showing 3 changed files with 52 additions and 53 deletions.
diff --git a/packages/nns/src/governance.canister.ts b/packages/nns/src/governance.canister.ts
@@ -1,23 +1,16 @@
 import type { ActorSubclass, Agent } from "@dfinity/agent";
 import { Actor } from "@dfinity/agent";
 import type { LedgerCanister } from "@dfinity/ledger-icp";
-import {
-  AccountIdentifier,
-  SubAccount,
-  checkAccountId,
-} from "@dfinity/ledger-icp";
+import { AccountIdentifier, checkAccountId } from "@dfinity/ledger-icp";
 import type { Principal } from "@dfinity/principal";
 import {
-  arrayOfNumberToUint8Array,
-  asciiStringToByteArray,
   assertPercentageNumber,
   createServices,
   fromNullable,
   isNullish,
   nonNullish,
   uint8ArrayToBigInt,
 } from "@dfinity/utils";
-import { sha256 } from "@noble/hashes/sha256";
 import randomBytes from "randombytes";
 import type {
   Command_1,
@@ -90,10 +83,7 @@ import type {
   ProposalId,
   ProposalInfo,
 } from "./types/governance_converters";
-import {
-  getRandomMemo,
-  memoToNeuronAccountIdentifier,
-} from "./utils/neurons.utils";
+import { memoToNeuronAccountIdentifier } from "./utils/neurons.utils";
 
 export class GovernanceCanister {
   private constructor(
@@ -212,7 +202,7 @@ export class GovernanceCanister {
    * Returns the latest reward event.
    *
    * If `certified` is true, the request is fetched as an update call, otherwise
-   * its fetched using a query call.
+   * it's fetched using a query call.
    *
    */
   public getLastestRewardEvent = async (
@@ -271,16 +261,17 @@ export class GovernanceCanister {
       throw new InsufficientAmountError(stake);
     }
 
-    const memo = getRandomMemo();
+    const nonceBytes = new Uint8Array(randomBytes(8));
+    const nonce = uint8ArrayToBigInt(nonceBytes);
     const accountIdentifier = memoToNeuronAccountIdentifier({
       controller: principal,
-      memo,
+      memo: nonce,
       governanceCanisterId: this.canisterId,
     });
 
     // Send amount to the ledger.
     await ledgerCanister.transfer({
-      memo,
+      memo: nonce,
       amount: stake,
       fromSubAccount,
       to: accountIdentifier,
@@ -292,7 +283,7 @@ export class GovernanceCanister {
     const neuronId: NeuronId | undefined =
       await this.claimOrRefreshNeuronFromAccount({
         controller: principal,
-        memo,
+        memo: nonce,
       });
 
     // Typescript was complaining with `neuronId || new NeuronNotFound()`:
@@ -873,32 +864,6 @@ export class GovernanceCanister {
     );
   };
 
-  private buildNeuronStakeSubAccount = (
-    nonce: Uint8Array,
-    principal: Principal,
-  ): SubAccount => {
-    return SubAccount.fromBytes(
-      this.getNeuronStakeSubAccountBytes(nonce, principal),
-    ) as SubAccount;
-  };
-
-  private getNeuronStakeSubAccountBytes = (
-    nonce: Uint8Array,
-    principal: Principal,
-  ): Uint8Array => {
-    const padding = asciiStringToByteArray("neuron-stake");
-    const shaObj = sha256.create();
-    shaObj.update(
-      arrayOfNumberToUint8Array([
-        0x0c,
-        ...padding,
-        ...principal.toUint8Array(),
-        ...nonce,
-      ]),
-    );
-    return shaObj.digest();
-  };
-
   private getGovernanceService(certified: boolean): GovernanceService {
     return certified ? this.certifiedService : this.service;
   }

diff --git a/packages/nns/src/utils/neurons.utils.spec.ts b/packages/nns/src/utils/neurons.utils.spec.ts
@@ -1,7 +1,11 @@
+import { Principal } from "@dfinity/principal";
+import { uint8ArrayToHexString } from "@dfinity/utils";
 import { Vote } from "../enums/governance.enums";
 import { NeuronInfo, ProposalInfo } from "../types/governance_converters";
 import {
   ineligibleNeurons,
+  memoToNeuronAccountIdentifier,
+  memoToNeuronSubaccount,
   votableNeurons,
   votedNeurons,
 } from "./neurons.utils";
@@ -202,4 +206,42 @@ describe("neurons-utils", () => {
       expect(voted).toEqual([eligibleNeuronsData[1], eligibleNeuronsData[2]]);
     });
   });
+
+  describe("memoToNeuronSubaccount", () => {
+    it("should calculate the neuron subaccount", () => {
+      // It's not that useful to do the same hashing here so we just use some
+      // hardcoded values for a "change detector" test.
+      const controller = Principal.fromText("6czrj-7isge-rrema");
+      const memo = 123123123n;
+
+      expect(
+        uint8ArrayToHexString(
+          memoToNeuronSubaccount({ controller, memo }).toUint8Array(),
+        ),
+      ).toEqual(
+        "21c6b6f1c9be307a956fd950687037203c6ceb3557ae5cc49b52010d4497b2d3",
+      );
+    });
+  });
+
+  describe("memoToNeuronAccountIdentifier", () => {
+    it("should calculate the neuron account identifier", () => {
+      // It's not that useful to do the same hashing here so we just use some
+      // hardcoded values for a "change detector" test.
+      const governanceCanisterId = Principal.fromText(
+        "rrkah-fqaaa-aaaaa-aaaaq-cai",
+      );
+      const controller = Principal.fromText("jl4mq-2sfmr-lekya");
+      const memo = 456456456n;
+      expect(
+        memoToNeuronAccountIdentifier({
+          controller,
+          memo,
+          governanceCanisterId,
+        }).toHex(),
+      ).toEqual(
+        "2a8523d6488286c937f25a855cc74e640da7115824eb6c78d89d503a37e6b3d8",
+      );
+    });
+  });
 });
diff --git a/packages/nns/src/utils/neurons.utils.ts b/packages/nns/src/utils/neurons.utils.ts
@@ -4,10 +4,8 @@ import {
   arrayOfNumberToUint8Array,
   asciiStringToByteArray,
   bigIntToUint8Array,
-  uint8ArrayToBigInt,
 } from "@dfinity/utils";
 import { sha256 } from "@noble/hashes/sha256";
-import randomBytes from "randombytes";
 import { Vote } from "../enums/governance.enums";
 import type {
   Ballot,
@@ -100,11 +98,6 @@ export const votedNeurons = ({
       getNeuronVoteForProposal({ proposal, neuron }) !== Vote.Unspecified,
   );
 
-export const getRandomMemo = (): bigint => {
-  const bytes = randomBytes(8);
-  return uint8ArrayToBigInt(bytes);
-};
-
 export const memoToNeuronSubaccount = ({
   controller,
   memo,
@@ -116,15 +109,14 @@ export const memoToNeuronSubaccount = ({
   const shaObj = sha256.create();
   shaObj.update(
     arrayOfNumberToUint8Array([
-      0x0c, // Should be padding.length
+      padding.length,
       ...padding,
       ...controller.toUint8Array(),
       ...bigIntToUint8Array(memo),
     ]),
   );
 
-  const ret = SubAccount.fromBytes(shaObj.digest());
-  return ret as SubAccount;
+  return SubAccount.fromBytes(shaObj.digest()) as SubAccount;
 };
 
 export const memoToNeuronAccountIdentifier = ({