Uncommitted work in progress

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

Hide file tree

Showing 2 changed files with 65 additions and 9 deletions.
diff --git a/packages/nns/src/governance.canister.ts b/packages/nns/src/governance.canister.ts
@@ -90,6 +90,10 @@ import type {
   ProposalId,
   ProposalInfo,
 } from "./types/governance_converters";
+import {
+  getRandomMemo,
+  memoToNeuronAccountIdentifier,
+} from "./utils/neurons.utils";
 
 export class GovernanceCanister {
   private constructor(
@@ -208,7 +212,7 @@ export class GovernanceCanister {
    * Returns the latest reward event.
    *
    * If `certified` is true, the request is fetched as an update call, otherwise
-   * it's fetched using a query call.
+   * its fetched using a query call.
    *
    */
   public getLastestRewardEvent = async (
@@ -267,17 +271,16 @@ export class GovernanceCanister {
       throw new InsufficientAmountError(stake);
     }
 
-    const nonceBytes = new Uint8Array(randomBytes(8));
-    const nonce = uint8ArrayToBigInt(nonceBytes);
-    const toSubAccount = this.buildNeuronStakeSubAccount(nonceBytes, principal);
-    const accountIdentifier = AccountIdentifier.fromPrincipal({
-      principal: this.canisterId,
-      subAccount: toSubAccount,
+    const memo = getRandomMemo();
+    const accountIdentifier = memoToNeuronAccountIdentifier({
+      controller: principal,
+      memo,
+      governanceCanisterId: this.canisterId,
     });
 
     // Send amount to the ledger.
     await ledgerCanister.transfer({
-      memo: nonce,
+      memo,
       amount: stake,
       fromSubAccount,
       to: accountIdentifier,
@@ -289,7 +292,7 @@ export class GovernanceCanister {
     const neuronId: NeuronId | undefined =
       await this.claimOrRefreshNeuronFromAccount({
         controller: principal,
-        memo: nonce,
+        memo,
       });
 
     // Typescript was complaining with `neuronId || new NeuronNotFound()`:

diff --git a/packages/nns/src/utils/neurons.utils.ts b/packages/nns/src/utils/neurons.utils.ts
@@ -1,3 +1,13 @@
+import { AccountIdentifier, SubAccount } from "@dfinity/ledger-icp";
+import type { Principal } from "@dfinity/principal";
+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,
@@ -89,3 +99,46 @@ export const votedNeurons = ({
     (neuron: NeuronInfo) =>
       getNeuronVoteForProposal({ proposal, neuron }) !== Vote.Unspecified,
   );
+
+export const getRandomMemo = (): bigint => {
+  const bytes = randomBytes(8);
+  return uint8ArrayToBigInt(bytes);
+};
+
+export const memoToNeuronSubaccount = ({
+  controller,
+  memo,
+}: {
+  controller: Principal;
+  memo: bigint;
+}): SubAccount => {
+  const padding = asciiStringToByteArray("neuron-stake");
+  const shaObj = sha256.create();
+  shaObj.update(
+    arrayOfNumberToUint8Array([
+      0x0c, // Should be padding.length
+      ...padding,
+      ...controller.toUint8Array(),
+      ...bigIntToUint8Array(memo),
+    ]),
+  );
+
+  const ret = SubAccount.fromBytes(shaObj.digest());
+  return ret as SubAccount;
+};
+
+export const memoToNeuronAccountIdentifier = ({
+  controller,
+  memo,
+  governanceCanisterId,
+}: {
+  controller: Principal;
+  memo: bigint;
+  governanceCanisterId: Principal;
+}): AccountIdentifier => {
+  const subAccount = memoToNeuronSubaccount({ controller, memo });
+  return AccountIdentifier.fromPrincipal({
+    principal: governanceCanisterId,
+    subAccount,
+  });
+};