From 5d70622819e453c938dd0230998c44b42997cbbd Mon Sep 17 00:00:00 2001
From: Jaco Greeff <jacogr@gmail.com>
Date: Tue, 9 Mar 2021 16:38:33 +0100
Subject: [PATCH] Extend signer payload with any additional extension fields
 (#3269)

* Extend payload with additional fields

* Extras

* Adjust
---
 packages/types/src/extrinsic/SignerPayload.ts | 119 ++++++++++++++----
 1 file changed, 97 insertions(+), 22 deletions(-)

diff --git a/packages/types/src/extrinsic/SignerPayload.ts b/packages/types/src/extrinsic/SignerPayload.ts
index eb54687a9b..c725e2b139 100644
--- a/packages/types/src/extrinsic/SignerPayload.ts
+++ b/packages/types/src/extrinsic/SignerPayload.ts
@@ -2,7 +2,7 @@
 // SPDX-License-Identifier: Apache-2.0
 
 import type { Address, Balance, BlockNumber, Call, ExtrinsicEra, Hash, Index, RuntimeVersion } from '../interfaces';
-import type { Codec, Constructor, ISignerPayload, SignerPayloadJSON, SignerPayloadRaw } from '../types';
+import type { Codec, InterfaceTypes, ISignerPayload, Registry, SignerPayloadJSON, SignerPayloadRaw } from '../types';
 
 import { u8aToHex } from '@polkadot/util';
 
@@ -26,9 +26,7 @@ export interface SignerPayloadType extends Codec {
   version: u8;
 }
 
-// We explicitly cast the type here to get the actual TypeScript exports right
-// We can ignore the properties, added via Struct.with
-const _Payload = Struct.with({
+const knownTypes: Record<string, keyof InterfaceTypes> = {
   address: 'Address',
   blockHash: 'Hash',
   blockNumber: 'BlockNumber',
@@ -40,33 +38,106 @@ const _Payload = Struct.with({
   signedExtensions: 'Vec<Text>',
   tip: 'Compact<Balance>',
   version: 'u8'
-}) as unknown as Constructor<SignerPayloadType>;
+};
 
 /**
  * @name GenericSignerPayload
  * @description
  * A generic signer payload that can be used for serialization between API and signer
  */
-export class GenericSignerPayload extends _Payload implements ISignerPayload {
+export class GenericSignerPayload extends Struct implements ISignerPayload, SignerPayloadType {
+  private readonly _extraTypes: Record<string, keyof InterfaceTypes>;
+
+  constructor (registry: Registry, value?: string | { [x: string]: any; } | Map<unknown, unknown> | unknown[]) {
+    const extensionTypes = {
+      ...registry.getSignedExtensionTypes(),
+      ...registry.getSignedExtensionExtra()
+    };
+
+    super(registry, {
+      ...extensionTypes,
+      ...knownTypes
+    }, value);
+
+    // add all extras that are not in the base types
+    this._extraTypes = Object.entries(extensionTypes).reduce<Record<string, keyof InterfaceTypes>>((map, [key, type]) => {
+      if (!knownTypes[key]) {
+        map[key] = type;
+      }
+
+      return map;
+    }, {});
+  }
+
+  get address (): Address {
+    return this.get('address') as Address;
+  }
+
+  get blockHash (): Hash {
+    return this.get('blockHash') as Hash;
+  }
+
+  get blockNumber (): BlockNumber {
+    return this.get('blockNumber') as BlockNumber;
+  }
+
+  get era (): ExtrinsicEra {
+    return this.get('era') as ExtrinsicEra;
+  }
+
+  get genesisHash (): Hash {
+    return this.get('genesisHash') as Hash;
+  }
+
+  get method (): Call {
+    return this.get('method') as Call;
+  }
+
+  get nonce (): Compact<Index> {
+    return this.get('nonce') as Compact<Index>;
+  }
+
+  get runtimeVersion (): RuntimeVersion {
+    return this.get('runtimeVersion') as RuntimeVersion;
+  }
+
+  get signedExtensions (): Vec<Text> {
+    return this.get('signedExtensions') as Vec<Text>;
+  }
+
+  get tip (): Compact<Balance> {
+    return this.get('tip') as Compact<Balance>;
+  }
+
+  get version (): u8 {
+    return this.get('version') as u8;
+  }
+
   /**
    * @description Creates an representation of the structure as an ISignerPayload JSON
    */
   public toPayload (): SignerPayloadJSON {
-    const { address, blockHash, blockNumber, era, genesisHash, method, nonce, runtimeVersion: { specVersion, transactionVersion }, signedExtensions, tip, version } = this;
-
     return {
-      address: address.toString(),
-      blockHash: blockHash.toHex(),
-      blockNumber: blockNumber.toHex(),
-      era: era.toHex(),
-      genesisHash: genesisHash.toHex(),
-      method: method.toHex(),
-      nonce: nonce.toHex(),
-      signedExtensions: signedExtensions.map((e) => e.toString()),
-      specVersion: specVersion.toHex(),
-      tip: tip.toHex(),
-      transactionVersion: transactionVersion.toHex(),
-      version: version.toNumber()
+      // add any explicit overrides we may have
+      ...(Object.keys(this._extraTypes).reduce<Record<string, string>>((map, key) => {
+        map[key] = (this.get(key) as Codec).toHex();
+
+        return map;
+      }, {})),
+      // the known defaults as managed explicitly and has different
+      // formatting in cases, e.g. we mostly expose a hex format here
+      address: this.address.toString(),
+      blockHash: this.blockHash.toHex(),
+      blockNumber: this.blockNumber.toHex(),
+      era: this.era.toHex(),
+      genesisHash: this.genesisHash.toHex(),
+      method: this.method.toHex(),
+      nonce: this.nonce.toHex(),
+      signedExtensions: this.signedExtensions.map((e) => e.toString()),
+      specVersion: this.runtimeVersion.specVersion.toHex(),
+      tip: this.tip.toHex(),
+      transactionVersion: this.runtimeVersion.transactionVersion.toHex(),
+      version: this.version.toNumber()
     };
   }
 
@@ -75,8 +146,12 @@ export class GenericSignerPayload extends _Payload implements ISignerPayload {
    */
   public toRaw (): SignerPayloadRaw {
     const payload = this.toPayload();
-    // NOTE Explicitly pass the bare flag so the method is encoded un-prefixed (non-decodable, for signing only)
-    const data = u8aToHex(this.registry.createType('ExtrinsicPayload', payload, { version: payload.version }).toU8a({ method: true }));
+    const data = u8aToHex(
+      this.registry
+        .createType('ExtrinsicPayload', payload, { version: payload.version })
+        // NOTE Explicitly pass the bare flag so the method is encoded un-prefixed (non-decodable, for signing only)
+        .toU8a({ method: true })
+    );
 
     return {
       address: payload.address,