feat(core-transactions): register wallet attributes before accessing …

…them (#2867)
ArkEcosystem · Aug 7, 2019 · 3c39180 · 3c39180
1 parent 79dcbe0
commit 3c39180
Show file tree

Hide file tree

Showing 18 changed files with 110 additions and 6 deletions.
diff --git a/__tests__/unit/core-state/wallets/wallet-manager-htlc.test.ts b/__tests__/unit/core-state/wallets/wallet-manager-htlc.test.ts
@@ -3,6 +3,7 @@ import { database } from "../mocks/database";
 import { state } from "../mocks/state";
 
 import { State } from "@arkecosystem/core-interfaces";
+import { Handlers } from "@arkecosystem/core-transactions";
 import { Crypto, Identities, Transactions, Utils } from "@arkecosystem/crypto";
 import { Wallet, WalletManager } from "../../../../packages/core-state/src/wallets";
 import { TransactionFactory } from "../../../helpers/transaction-factory";
@@ -15,6 +16,14 @@ let walletManager: State.IWalletManager;
 
 const makeTimestamp = (secondsRelativeToNow = 0) => Math.floor((Date.now() + secondsRelativeToNow * 1000) / 1000);
 
+beforeAll(() => {
+    jest.spyOn(Handlers.Registry, "isKnownWalletAttribute").mockReturnValue(true);
+});
+
+afterAll(() => {
+    jest.restoreAllMocks();
+});
+
 describe("Wallet Manager", () => {
     describe("HTLC claim", () => {
         const lockPassphrase = "craft imitate step mixture patch forest volcano business charge around girl confirm";

diff --git a/__tests__/unit/core-state/wallets/wallet-manager.test.ts b/__tests__/unit/core-state/wallets/wallet-manager.test.ts
@@ -2,6 +2,7 @@
 import "../../core-database/mocks/core-container";
 
 import { State } from "@arkecosystem/core-interfaces";
+import { Handlers } from "@arkecosystem/core-transactions";
 import { InsufficientBalanceError } from "@arkecosystem/core-transactions/src/errors";
 import { Blocks, Constants, Identities, Interfaces, Transactions, Utils } from "@arkecosystem/crypto";
 import { Address } from "@arkecosystem/crypto/src/identities";
@@ -409,10 +410,17 @@ describe("Wallet Manager", () => {
             });
 
             const wallet = new Wallet(walletData1.address);
-            wallet.setAttribute("custom.attribute", "something");
+            expect(() => wallet.setAttribute("custom.attribute", "something")).toThrow();
+
+            const spy = jest.spyOn(Handlers.Registry, "isKnownWalletAttribute").mockReturnValue(true);
+
+            expect(() => wallet.setAttribute("custom.attribute", "something")).not.toThrow();
+
             walletManager.reindex(wallet);
 
             expect(walletManager.findById("something")).toBe(wallet);
+
+            spy.mockRestore();
         });
 
         it("should unregister an index", () => {

diff --git a/__tests__/unit/core-transactions/handler-registry.test.ts b/__tests__/unit/core-transactions/handler-registry.test.ts
@@ -64,7 +64,11 @@ class TestTransaction extends Transactions.Transaction {
 
 // tslint:disable-next-line:max-classes-per-file
 class TestTransactionHandler extends TransactionHandler {
-    public dependencies(): TransactionHandlerConstructor[] {
+    public dependencies(): ReadonlyArray<TransactionHandlerConstructor> {
+        return [];
+    }
+
+    public walletAttributes(): ReadonlyArray<string> {
         return [];
     }
 

diff --git a/packages/core-state/src/wallets/wallet.ts b/packages/core-state/src/wallets/wallet.ts
@@ -1,6 +1,7 @@
 import { State } from "@arkecosystem/core-interfaces";
-import { Errors } from "@arkecosystem/core-transactions";
+import { Errors, Handlers } from "@arkecosystem/core-transactions";
 import { Crypto, Enums, Identities, Interfaces, Transactions, Utils } from "@arkecosystem/crypto";
+import assert from "assert";
 import dottie from "dottie";
 
 export class Wallet implements State.IWallet {
@@ -20,18 +21,22 @@ export class Wallet implements State.IWallet {
     }
 
     public hasAttribute(key: string): boolean {
+        this.assertKnownAttribute(key);
         return dottie.exists(this.attributes, key);
     }
 
     public getAttribute<T>(key: string, defaultValue?: T): T {
+        this.assertKnownAttribute(key);
         return dottie.get(this.attributes, key, defaultValue);
     }
 
     public setAttribute<T = any>(key: string, value: T): void {
+        this.assertKnownAttribute(key);
         dottie.set(this.attributes, key, value);
     }
 
     public forgetAttribute(key: string): void {
+        this.assertKnownAttribute(key);
         this.setAttribute(key, undefined);
     }
 
@@ -232,4 +237,8 @@ export class Wallet implements State.IWallet {
     public toString(): string {
         return `${this.address} (${Utils.formatSatoshi(this.balance)})`;
     }
+
+    private assertKnownAttribute(key: string): void {
+        assert(Handlers.Registry.isKnownWalletAttribute(key), `Tried to access unknown attribute: ${key}`);
+    }
 }
diff --git a/packages/core-transactions/src/handlers/delegate-registration.ts b/packages/core-transactions/src/handlers/delegate-registration.ts
@@ -20,6 +20,19 @@ export class DelegateRegistrationTransactionHandler extends TransactionHandler {
         return [];
     }
 
+    public walletAttributes(): ReadonlyArray<string> {
+        return [
+            "delegate",
+            "delegate.lastBlock",
+            "delegate.rank",
+            "delegate.round",
+            "delegate.username",
+            "delegate.voteBalance",
+            "delegate.forgedTotal",
+            "delegate.approval",
+        ];
+    }
+
     public async bootstrap(connection: Database.IConnection, walletManager: State.IWalletManager): Promise<void> {
         const transactions = await connection.transactionsRepository.getAssetsByType(this.getConstructor().type);
         const forgedBlocks = await connection.blocksRepository.getDelegatesForgedBlocks();

diff --git a/packages/core-transactions/src/handlers/delegate-resignation.ts b/packages/core-transactions/src/handlers/delegate-resignation.ts
@@ -14,6 +14,10 @@ export class DelegateResignationTransactionHandler extends TransactionHandler {
         return [DelegateRegistrationTransactionHandler];
     }
 
+    public walletAttributes(): ReadonlyArray<string> {
+        return ["delegate.resigned"];
+    }
+
     public async bootstrap(connection: Database.IConnection, walletManager: State.IWalletManager): Promise<void> {
         const transactions = await connection.transactionsRepository.getAssetsByType(this.getConstructor().type);
 

diff --git a/packages/core-transactions/src/handlers/handler-registry.ts b/packages/core-transactions/src/handlers/handler-registry.ts
@@ -1,5 +1,7 @@
 import { Enums, Errors, Transactions } from "@arkecosystem/crypto";
 
+import assert from "assert";
+import { InvalidTransactionTypeError } from "../errors";
 import { DelegateRegistrationTransactionHandler } from "./delegate-registration";
 import { DelegateResignationTransactionHandler } from "./delegate-resignation";
 import { HtlcClaimTransactionHandler } from "./htlc-claim";
@@ -9,18 +11,18 @@ import { IpfsTransactionHandler } from "./ipfs";
 import { MultiPaymentTransactionHandler } from "./multi-payment";
 import { MultiSignatureTransactionHandler } from "./multi-signature";
 import { SecondSignatureTransactionHandler } from "./second-signature";
+import { TransactionHandler, TransactionHandlerConstructor } from "./transaction";
 import { TransferTransactionHandler } from "./transfer";
 import { VoteTransactionHandler } from "./vote";
 
-import { InvalidTransactionTypeError } from "../errors";
-import { TransactionHandler, TransactionHandlerConstructor } from "./transaction";
-
 export class TransactionHandlerRegistry {
     private readonly registeredTransactionHandlers: Map<
         Transactions.InternalTransactionType,
         TransactionHandler
     > = new Map();
 
+    private readonly knownWalletAttributes: Map<string, boolean> = new Map();
+
     constructor() {
         this.registerTransactionHandler(TransferTransactionHandler);
         this.registerTransactionHandler(SecondSignatureTransactionHandler);
@@ -80,6 +82,12 @@ export class TransactionHandlerRegistry {
             Transactions.TransactionRegistry.registerTransactionType(transactionConstructor);
         }
 
+        const walletAttributes: ReadonlyArray<string> = service.walletAttributes();
+        for (const attribute of walletAttributes) {
+            assert(!this.knownWalletAttributes.has(attribute), `Wallet attribute is already known: ${attribute}`);
+            this.knownWalletAttributes.set(attribute, true);
+        }
+
         this.registeredTransactionHandlers.set(internalType, service);
     }
 
@@ -100,9 +108,18 @@ export class TransactionHandlerRegistry {
             throw new InvalidTransactionTypeError(internalType.toString());
         }
 
+        const walletAttributes: ReadonlyArray<string> = service.walletAttributes();
+        for (const attribute of walletAttributes) {
+            this.knownWalletAttributes.delete(attribute);
+        }
+
         Transactions.TransactionRegistry.deregisterTransactionType(transactionConstructor);
         this.registeredTransactionHandlers.delete(internalType);
     }
+
+    public isKnownWalletAttribute(attribute: string): boolean {
+        return this.knownWalletAttributes.has(attribute);
+    }
 }
 
 export const transactionHandlerRegistry = new TransactionHandlerRegistry();
diff --git a/packages/core-transactions/src/handlers/htlc-claim.ts b/packages/core-transactions/src/handlers/htlc-claim.ts
@@ -28,6 +28,10 @@ export class HtlcClaimTransactionHandler extends TransactionHandler {
         return [HtlcLockTransactionHandler];
     }
 
+    public walletAttributes(): ReadonlyArray<string> {
+        return [];
+    }
+
     public async bootstrap(connection: Database.IConnection, walletManager: State.IWalletManager): Promise<void> {
         const transactions = await connection.transactionsRepository.getAssetsByType(this.getConstructor().type);
         for (const transaction of transactions) {

diff --git a/packages/core-transactions/src/handlers/htlc-lock.ts b/packages/core-transactions/src/handlers/htlc-lock.ts
@@ -11,6 +11,10 @@ export class HtlcLockTransactionHandler extends TransactionHandler {
         return [];
     }
 
+    public walletAttributes(): ReadonlyArray<string> {
+        return ["htlc.locks", "htlc.lockedBalance"];
+    }
+
     public async bootstrap(connection: Database.IConnection, walletManager: State.IWalletManager): Promise<void> {
         const lockTransactions = await connection.transactionsRepository.getAssetsByType(this.getConstructor().type);
         for (const transaction of lockTransactions) {

diff --git a/packages/core-transactions/src/handlers/htlc-refund.ts b/packages/core-transactions/src/handlers/htlc-refund.ts
@@ -27,6 +27,10 @@ export class HtlcRefundTransactionHandler extends TransactionHandler {
         return [HtlcLockTransactionHandler];
     }
 
+    public walletAttributes(): ReadonlyArray<string> {
+        return [];
+    }
+
     public async bootstrap(connection: Database.IConnection, walletManager: State.IWalletManager): Promise<void> {
         const transactions = await connection.transactionsRepository.getAssetsByType(this.getConstructor().type);
         for (const transaction of transactions) {

diff --git a/packages/core-transactions/src/handlers/ipfs.ts b/packages/core-transactions/src/handlers/ipfs.ts
@@ -11,6 +11,10 @@ export class IpfsTransactionHandler extends TransactionHandler {
         return [];
     }
 
+    public walletAttributes(): ReadonlyArray<string> {
+        return ["ipfs", "ipfs.hashes"];
+    }
+
     public async bootstrap(connection: Database.IConnection, walletManager: State.IWalletManager): Promise<void> {
         const transactions = await connection.transactionsRepository.getAssetsByType(this.getConstructor().type);
 

diff --git a/packages/core-transactions/src/handlers/multi-payment.ts b/packages/core-transactions/src/handlers/multi-payment.ts
@@ -12,6 +12,10 @@ export class MultiPaymentTransactionHandler extends TransactionHandler {
         return [];
     }
 
+    public walletAttributes(): ReadonlyArray<string> {
+        return [];
+    }
+
     public async bootstrap(connection: Database.IConnection, walletManager: State.IWalletManager): Promise<void> {
         const transactions = await connection.transactionsRepository.getAssetsByType(this.getConstructor().type);
 

diff --git a/packages/core-transactions/src/handlers/multi-signature.ts b/packages/core-transactions/src/handlers/multi-signature.ts
@@ -17,6 +17,10 @@ export class MultiSignatureTransactionHandler extends TransactionHandler {
         return [];
     }
 
+    public walletAttributes(): ReadonlyArray<string> {
+        return ["multiSignature"];
+    }
+
     public async bootstrap(connection: Database.IConnection, walletManager: State.IWalletManager): Promise<void> {
         const transactions = await connection.transactionsRepository.getAssetsByType(this.getConstructor().type);
 

diff --git a/packages/core-transactions/src/handlers/second-signature.ts b/packages/core-transactions/src/handlers/second-signature.ts
@@ -12,6 +12,10 @@ export class SecondSignatureTransactionHandler extends TransactionHandler {
         return [];
     }
 
+    public walletAttributes(): ReadonlyArray<string> {
+        return ["secondPublicKey"];
+    }
+
     public async bootstrap(connection: Database.IConnection, walletManager: State.IWalletManager): Promise<void> {
         const transactions = await connection.transactionsRepository.getAssetsByType(this.getConstructor().type);
 

diff --git a/packages/core-transactions/src/handlers/transaction.ts b/packages/core-transactions/src/handlers/transaction.ts
@@ -22,6 +22,8 @@ export abstract class TransactionHandler implements ITransactionHandler {
 
     public abstract dependencies(): ReadonlyArray<TransactionHandlerConstructor>;
 
+    public abstract walletAttributes(): ReadonlyArray<string>;
+
     /**
      * Wallet logic
      */

diff --git a/packages/core-transactions/src/handlers/transfer.ts b/packages/core-transactions/src/handlers/transfer.ts
@@ -12,6 +12,10 @@ export class TransferTransactionHandler extends TransactionHandler {
         return [];
     }
 
+    public walletAttributes(): ReadonlyArray<string> {
+        return [];
+    }
+
     public async bootstrap(connection: Database.IConnection, walletManager: State.IWalletManager): Promise<void> {
         const transactions = await connection.transactionsRepository.getReceivedTransactions();
 

diff --git a/packages/core-transactions/src/handlers/vote.ts b/packages/core-transactions/src/handlers/vote.ts
@@ -19,6 +19,10 @@ export class VoteTransactionHandler extends TransactionHandler {
         return [DelegateRegistrationTransactionHandler];
     }
 
+    public walletAttributes(): ReadonlyArray<string> {
+        return ["vote"];
+    }
+
     public async bootstrap(connection: Database.IConnection, walletManager: State.IWalletManager): Promise<void> {
         const transactions = await connection.transactionsRepository.getAssetsByType(this.getConstructor().type);
 

diff --git a/packages/core-transactions/src/interfaces.ts b/packages/core-transactions/src/interfaces.ts
@@ -7,6 +7,8 @@ export interface ITransactionHandler {
 
     dependencies(): ReadonlyArray<TransactionHandlerConstructor>;
 
+    walletAttributes(): ReadonlyArray<string>;
+
     bootstrap(connection: Database.IConnection, walletManager: State.IWalletManager): Promise<void>;
 
     isActivated(): Promise<boolean>;