From ed5f2673cebc30094c219f4d7b0ceb90d981392f Mon Sep 17 00:00:00 2001
From: Jaco Greeff <jacogr@gmail.com>
Date: Thu, 8 Apr 2021 15:07:39 +0300
Subject: [PATCH] Detect Keys, inject SessionKeys (#3380)

---
 packages/api/src/base/util.ts                 | 105 ++++++++++++++----
 packages/types-known/src/spec/rococo.ts       |   2 +-
 packages/types/src/augment/registry.ts        |  23 +++-
 .../src/interfaces/session/definitions.ts     |  15 ++-
 .../types/src/interfaces/session/types.ts     |  23 +++-
 5 files changed, 141 insertions(+), 27 deletions(-)

diff --git a/packages/api/src/base/util.ts b/packages/api/src/base/util.ts
index 3008c17cb1..010439cd13 100644
--- a/packages/api/src/base/util.ts
+++ b/packages/api/src/base/util.ts
@@ -2,15 +2,18 @@
 // SPDX-License-Identifier: Apache-2.0
 
 import type { ApiInterfaceRx } from '@polkadot/api/types';
-import type { bool } from '@polkadot/types';
+import type { bool, Raw } from '@polkadot/types';
 import type { Releases } from '@polkadot/types/interfaces';
 import type { InterfaceTypes } from '@polkadot/types/types';
 
+import { assert, compactFromU8a } from '@polkadot/util';
 import { combineLatest, Observable, of } from '@polkadot/x-rxjs';
 import { map, take } from '@polkadot/x-rxjs/operators';
 
 // the order and types needs to map with the all array setup below
-type Extracted = [bool | null, bool | null, bool | null, Releases | null];
+type ExtractedQ = [bool | null, bool | null, bool | null, Releases | null];
+
+type ExtractedR = [Raw | null];
 
 type DetectedKeys = keyof Pick<InterfaceTypes, 'AccountInfo' | 'ValidatorPrefs'>;
 
@@ -18,11 +21,21 @@ type DetectedValues = keyof InterfaceTypes;
 
 interface DetectedTypes extends Record<DetectedKeys, DetectedValues> {
   AccountInfo: 'AccountInfoWithRefCount' | 'AccountInfoWithDualRefCount' | 'AccountInfoWithTripleRefCount',
+  Keys: 'SessionKeys1' | 'SessionKeys2' | 'SessionKeys3' | 'SessionKeys4' | 'SessionKeys5' | 'SessionKeys6' | 'SessionKeys7' | 'SessionKeys8' | 'SessionKeys9' | 'SessionKeys10';
   RefCount: 'u8' | 'u32',
   ValidatorPrefs: 'ValidatorPrefsWithBlocked' | 'ValidatorPrefsWithCommission';
 }
 
-function mapCapabilities ([systemRefcount32, systemRefcountDual, systemRefcountTriple, stakingVersion]: Extracted): Partial<DetectedTypes> {
+interface AvailableMap<T> {
+  filtered: T[];
+  included: boolean[];
+}
+
+interface Constants {
+  accountIdLength: number;
+}
+
+function mapCapabilities ({ accountIdLength }: Constants, [systemRefcount32, systemRefcountDual, systemRefcountTriple, stakingVersion]: ExtractedQ, [keys]: ExtractedR): Partial<DetectedTypes> {
   const types: Partial<DetectedTypes> = {};
 
   // AccountInfo
@@ -47,36 +60,90 @@ function mapCapabilities ([systemRefcount32, systemRefcountDual, systemRefcountT
     }
   }
 
+  // Keys
+  if (keys) {
+    try {
+      const [offset, numItems] = compactFromU8a(keys);
+      const tupleLength = (keys.length - offset) / numItems.toNumber();
+      const numIds = tupleLength / accountIdLength;
+      const numIdsRound = Math.floor(numIds);
+
+      assert(numIds >= 2 && numIds <= 11, `Detected ${numIds} in Keys, should be >= 2 and <= 11`);
+
+      if (numIdsRound !== numIds) {
+        // Beefy?
+        if ((((numIdsRound - 1) * accountIdLength) + 33) === tupleLength) {
+          // eslint-disable-next-line @typescript-eslint/ban-ts-comment
+          // @ts-ignore
+          types.Keys = `SessionKeys${numIdsRound - 1}B`;
+        } else {
+          assert(false, `Expected integer number of keys, found ${numIds.toFixed(2)}`);
+        }
+      } else {
+        // eslint-disable-next-line @typescript-eslint/ban-ts-comment
+        // @ts-ignore
+        types.Keys = `SessionKeys${numIds - 1}`;
+      }
+    } catch (error) {
+      console.error(error);
+    }
+  }
+
   return types;
 }
 
+function filterEntries <T> (entries: T[]): AvailableMap<T> {
+  const included = entries.map((c) => !!c);
+  const filtered = entries.filter((_, index) => included[index]);
+
+  return { filtered, included };
+}
+
+function extractResults <R, T = unknown> (results: unknown[], map: AvailableMap<T>): R {
+  let offset = -1;
+
+  return map.included.map((isIncluded) =>
+    isIncluded
+      ? results[++offset]
+      : null
+  ) as unknown as R;
+}
+
 /**
  * @description Query the chain for the specific capabilities
  */
 export function detectedCapabilities (api: ApiInterfaceRx, blockHash?: Uint8Array | string | undefined): Observable<Partial<DetectedTypes>> {
-  const all = [
+  const queries = filterEntries([
     api.query.system?.upgradedToU32RefCount,
     api.query.system?.upgradedToDualRefCount,
     api.query.system?.upgradedToTripleRefCount,
     api.query.staking?.storageVersion
-  ];
-  const included = all.map((c) => !!c);
-  const filtered = all.filter((_, index) => included[index]);
+  ]);
+  const keyed = filterEntries([
+    api.query.session?.queuedKeys
+  ]);
 
-  return (
-    filtered.length
+  return combineLatest([
+    queries.filtered.length
       ? blockHash
-        ? combineLatest(filtered.map((c) => c.at(blockHash)))
-        : api.queryMulti(filtered)
+        ? combineLatest(queries.filtered.map((c) => c.at(blockHash)))
+        : api.queryMulti(queries.filtered)
+      : of([]),
+    keyed.filtered.length
+      ? blockHash
+        ? combineLatest(keyed.filtered.map((k) => api.rpc.state.getStorage.raw(k.key(), blockHash)))
+        : combineLatest(keyed.filtered.map((k) => api.rpc.state.getStorage.raw(k.key())))
       : of([])
-  ).pipe(
-    map((results): Partial<DetectedTypes> => {
-      let offset = -1;
-
-      return mapCapabilities(
-        included.map((isIncluded) => isIncluded ? results[++offset] : null) as Extracted
-      );
-    }),
+  ]).pipe(
+    map(([qResults, rResults]): Partial<DetectedTypes> =>
+      mapCapabilities(
+        {
+          accountIdLength: api.registry.createType('AccountId').encodedLength
+        },
+        extractResults<ExtractedQ>(qResults, queries),
+        extractResults<ExtractedR>(rResults, keyed)
+      )
+    ),
     take(1)
   );
 }
diff --git a/packages/types-known/src/spec/rococo.ts b/packages/types-known/src/spec/rococo.ts
index 27471631af..2c3c301659 100644
--- a/packages/types-known/src/spec/rococo.ts
+++ b/packages/types-known/src/spec/rococo.ts
@@ -10,7 +10,7 @@ import type { OverrideVersionedType } from '@polkadot/types/types';
 
 const sharedTypes = {
   FullIdentification: '()', // No staking, only session (as per config)
-  Keys: 'SessionKeys7'
+  Keys: 'SessionKeys7B'
 };
 
 const versioned: OverrideVersionedType[] = [
diff --git a/packages/types/src/augment/registry.ts b/packages/types/src/augment/registry.ts
index 44267a1440..76d6761320 100644
--- a/packages/types/src/augment/registry.ts
+++ b/packages/types/src/augment/registry.ts
@@ -44,7 +44,7 @@ import type { RpcMethods } from '@polkadot/types/interfaces/rpc';
 import type { AccountId, AccountIdOf, AccountIndex, Address, AssetId, Balance, BalanceOf, Block, BlockNumber, Call, CallHash, CallHashOf, ChangesTrieConfiguration, CodecHash, Consensus, ConsensusEngineId, Digest, DigestItem, EncodedJustification, ExtrinsicsWeight, Fixed128, Fixed64, FixedI128, FixedI64, FixedU128, FixedU64, H1024, H128, H160, H2048, H256, H512, H64, Hash, Header, I32F32, Index, IndicesLookupSource, Justification, Justifications, KeyTypeId, KeyValue, LockIdentifier, LookupSource, LookupTarget, ModuleId, Moment, MultiAddress, MultiSigner, OpaqueCall, Origin, OriginCaller, PalletVersion, PalletsOrigin, Pays, PerU16, Perbill, Percent, Permill, Perquintill, Phantom, PhantomData, PreRuntime, Releases, RuntimeDbWeight, Seal, SealV0, SignedBlock, SignedBlockWithJustification, SignedBlockWithJustifications, Slot, StorageData, StorageProof, TransactionPriority, U32F32, ValidatorId, ValidatorIdOf, Weight, WeightMultiplier } from '@polkadot/types/interfaces/runtime';
 import type { SiField, SiLookupTypeId, SiPath, SiType, SiTypeDef, SiTypeDefArray, SiTypeDefComposite, SiTypeDefPrimitive, SiTypeDefSequence, SiTypeDefTuple, SiTypeDefVariant, SiVariant } from '@polkadot/types/interfaces/scaleInfo';
 import type { Period, Priority, SchedulePeriod, SchedulePriority, Scheduled, ScheduledTo254, TaskAddress } from '@polkadot/types/interfaces/scheduler';
-import type { FullIdentification, IdentificationTuple, Keys, MembershipProof, SessionIndex, SessionKeys1, SessionKeys2, SessionKeys3, SessionKeys4, SessionKeys5, SessionKeys6, SessionKeys7, SessionKeys8, SessionKeys9, ValidatorCount } from '@polkadot/types/interfaces/session';
+import type { BeefyKey, FullIdentification, IdentificationTuple, Keys, MembershipProof, SessionIndex, SessionKeys1, SessionKeys10, SessionKeys10B, SessionKeys2, SessionKeys3, SessionKeys4, SessionKeys5, SessionKeys6, SessionKeys6B, SessionKeys7, SessionKeys7B, SessionKeys8, SessionKeys8B, SessionKeys9, SessionKeys9B, ValidatorCount } from '@polkadot/types/interfaces/session';
 import type { Bid, BidKind, SocietyJudgement, SocietyVote, StrikeCount, VouchingStatus } from '@polkadot/types/interfaces/society';
 import type { ActiveEraInfo, CompactAssignments, CompactAssignmentsTo257, CompactScore, CompactScoreCompact, ElectionCompute, ElectionPhase, ElectionResult, ElectionScore, ElectionSize, ElectionStatus, EraIndex, EraPoints, EraRewardPoints, EraRewards, Exposure, ExtendedBalance, Forcing, IndividualExposure, KeyType, MomentOf, Nominations, NominatorIndex, NominatorIndexCompact, OffchainAccuracy, OffchainAccuracyCompact, PhragmenScore, Points, RawSolution, ReadySolution, RewardDestination, RewardPoint, RoundSnapshot, SeatHolder, SlashJournalEntry, SlashingSpans, SlashingSpansTo204, SolutionOrSnapshotSize, SolutionSupport, SolutionSupports, SpanIndex, SpanRecord, StakingLedger, StakingLedgerTo223, StakingLedgerTo240, UnappliedSlash, UnappliedSlashOther, UnlockChunk, ValidatorIndex, ValidatorIndexCompact, ValidatorPrefs, ValidatorPrefsTo145, ValidatorPrefsTo196, ValidatorPrefsWithBlocked, ValidatorPrefsWithCommission, VoteWeight, Voter } from '@polkadot/types/interfaces/staking';
 import type { ApiId, KeyValueOption, ReadProof, RuntimeVersion, RuntimeVersionApi, StorageChangeSet } from '@polkadot/types/interfaces/state';
@@ -213,6 +213,7 @@ declare module '@polkadot/types/types/registry' {
     'Option<BalanceOf>': Option<BalanceOf>;
     'Option<Balances>': Option<Balances>;
     'Option<BalanceStatus>': Option<BalanceStatus>;
+    'Option<BeefyKey>': Option<BeefyKey>;
     'Option<BeefyNextAuthoritySet>': Option<BeefyNextAuthoritySet>;
     'Option<Bid>': Option<Bid>;
     'Option<Bidder>': Option<Bidder>;
@@ -733,14 +734,20 @@ declare module '@polkadot/types/types/registry' {
     'Option<SessionInfo>': Option<SessionInfo>;
     'Option<SessionInfoValidatorGroup>': Option<SessionInfoValidatorGroup>;
     'Option<SessionKeys1>': Option<SessionKeys1>;
+    'Option<SessionKeys10>': Option<SessionKeys10>;
+    'Option<SessionKeys10B>': Option<SessionKeys10B>;
     'Option<SessionKeys2>': Option<SessionKeys2>;
     'Option<SessionKeys3>': Option<SessionKeys3>;
     'Option<SessionKeys4>': Option<SessionKeys4>;
     'Option<SessionKeys5>': Option<SessionKeys5>;
     'Option<SessionKeys6>': Option<SessionKeys6>;
+    'Option<SessionKeys6B>': Option<SessionKeys6B>;
     'Option<SessionKeys7>': Option<SessionKeys7>;
+    'Option<SessionKeys7B>': Option<SessionKeys7B>;
     'Option<SessionKeys8>': Option<SessionKeys8>;
+    'Option<SessionKeys8B>': Option<SessionKeys8B>;
     'Option<SessionKeys9>': Option<SessionKeys9>;
+    'Option<SessionKeys9B>': Option<SessionKeys9B>;
     'Option<SetId>': Option<SetId>;
     'Option<SetIndex>': Option<SetIndex>;
     'Option<SiField>': Option<SiField>;
@@ -978,6 +985,7 @@ declare module '@polkadot/types/types/registry' {
     'Vec<BalanceOf>': Vec<BalanceOf>;
     'Vec<Balances>': Vec<Balances>;
     'Vec<BalanceStatus>': Vec<BalanceStatus>;
+    'Vec<BeefyKey>': Vec<BeefyKey>;
     'Vec<BeefyNextAuthoritySet>': Vec<BeefyNextAuthoritySet>;
     'Vec<Bid>': Vec<Bid>;
     'Vec<Bidder>': Vec<Bidder>;
@@ -1498,14 +1506,20 @@ declare module '@polkadot/types/types/registry' {
     'Vec<SessionInfo>': Vec<SessionInfo>;
     'Vec<SessionInfoValidatorGroup>': Vec<SessionInfoValidatorGroup>;
     'Vec<SessionKeys1>': Vec<SessionKeys1>;
+    'Vec<SessionKeys10>': Vec<SessionKeys10>;
+    'Vec<SessionKeys10B>': Vec<SessionKeys10B>;
     'Vec<SessionKeys2>': Vec<SessionKeys2>;
     'Vec<SessionKeys3>': Vec<SessionKeys3>;
     'Vec<SessionKeys4>': Vec<SessionKeys4>;
     'Vec<SessionKeys5>': Vec<SessionKeys5>;
     'Vec<SessionKeys6>': Vec<SessionKeys6>;
+    'Vec<SessionKeys6B>': Vec<SessionKeys6B>;
     'Vec<SessionKeys7>': Vec<SessionKeys7>;
+    'Vec<SessionKeys7B>': Vec<SessionKeys7B>;
     'Vec<SessionKeys8>': Vec<SessionKeys8>;
+    'Vec<SessionKeys8B>': Vec<SessionKeys8B>;
     'Vec<SessionKeys9>': Vec<SessionKeys9>;
+    'Vec<SessionKeys9B>': Vec<SessionKeys9B>;
     'Vec<SetId>': Vec<SetId>;
     'Vec<SetIndex>': Vec<SetIndex>;
     'Vec<SiField>': Vec<SiField>;
@@ -1743,6 +1757,7 @@ declare module '@polkadot/types/types/registry' {
     BalanceOf: BalanceOf;
     Balances: Balances;
     BalanceStatus: BalanceStatus;
+    BeefyKey: BeefyKey;
     BeefyNextAuthoritySet: BeefyNextAuthoritySet;
     Bid: Bid;
     Bidder: Bidder;
@@ -2263,14 +2278,20 @@ declare module '@polkadot/types/types/registry' {
     SessionInfo: SessionInfo;
     SessionInfoValidatorGroup: SessionInfoValidatorGroup;
     SessionKeys1: SessionKeys1;
+    SessionKeys10: SessionKeys10;
+    SessionKeys10B: SessionKeys10B;
     SessionKeys2: SessionKeys2;
     SessionKeys3: SessionKeys3;
     SessionKeys4: SessionKeys4;
     SessionKeys5: SessionKeys5;
     SessionKeys6: SessionKeys6;
+    SessionKeys6B: SessionKeys6B;
     SessionKeys7: SessionKeys7;
+    SessionKeys7B: SessionKeys7B;
     SessionKeys8: SessionKeys8;
+    SessionKeys8B: SessionKeys8B;
     SessionKeys9: SessionKeys9;
+    SessionKeys9B: SessionKeys9B;
     SetId: SetId;
     SetIndex: SetIndex;
     SiField: SiField;
diff --git a/packages/types/src/interfaces/session/definitions.ts b/packages/types/src/interfaces/session/definitions.ts
index 4ceee57cf7..03cdd9f83f 100644
--- a/packages/types/src/interfaces/session/definitions.ts
+++ b/packages/types/src/interfaces/session/definitions.ts
@@ -12,22 +12,27 @@ import type { Definitions } from '../../types';
 //     pub struct SessionKeys {
 // Here we revert to tuples to keep the interfaces "opaque", as per the use
 const keyTypes = {
+  // key for beefy
+  BeefyKey: '[u8; 33]',
+
   // default to Substrate master defaults, 4 keys (polkadot master, 5 keys)
   Keys: 'SessionKeys4',
 
-  // shortcuts for 1-9 key tuples
   SessionKeys1: '(AccountId)',
   SessionKeys2: '(AccountId, AccountId)',
-  // older substrate master
   SessionKeys3: '(AccountId, AccountId, AccountId)',
-  // CC2, Substrate master
   SessionKeys4: '(AccountId, AccountId, AccountId, AccountId)',
-  // CC3
   SessionKeys5: '(AccountId, AccountId, AccountId, AccountId, AccountId)',
   SessionKeys6: '(AccountId, AccountId, AccountId, AccountId, AccountId, AccountId)',
+  SessionKeys6B: '(AccountId, AccountId, AccountId, AccountId, AccountId, BeefyKey)',
   SessionKeys7: '(AccountId, AccountId, AccountId, AccountId, AccountId, AccountId, AccountId)',
+  SessionKeys7B: '(AccountId, AccountId, AccountId, AccountId, AccountId, AccountId, BeefyKey)',
   SessionKeys8: '(AccountId, AccountId, AccountId, AccountId, AccountId, AccountId, AccountId, AccountId)',
-  SessionKeys9: '(AccountId, AccountId, AccountId, AccountId, AccountId, AccountId, AccountId, AccountId, AccountId)'
+  SessionKeys8B: '(AccountId, AccountId, AccountId, AccountId, AccountId, AccountId, AccountId, BeefyKey)',
+  SessionKeys9: '(AccountId, AccountId, AccountId, AccountId, AccountId, AccountId, AccountId, AccountId, AccountId)',
+  SessionKeys9B: '(AccountId, AccountId, AccountId, AccountId, AccountId, AccountId, AccountId, AccountId, BeefyKey)',
+  SessionKeys10: '(AccountId, AccountId, AccountId, AccountId, AccountId, AccountId, AccountId, AccountId, AccountId, AccountId)',
+  SessionKeys10B: '(AccountId, AccountId, AccountId, AccountId, AccountId, AccountId, AccountId, AccountId, AccountId, BeefyKey)'
 };
 
 export default {
diff --git a/packages/types/src/interfaces/session/types.ts b/packages/types/src/interfaces/session/types.ts
index ca5c2a6052..de12db0c63 100644
--- a/packages/types/src/interfaces/session/types.ts
+++ b/packages/types/src/interfaces/session/types.ts
@@ -1,11 +1,14 @@
 // Auto-generated via `yarn polkadot-types-from-defs`, do not edit
 /* eslint-disable */
 
-import type { Bytes, Struct, Vec, u32 } from '@polkadot/types';
+import type { Bytes, Struct, U8aFixed, Vec, u32 } from '@polkadot/types';
 import type { ITuple } from '@polkadot/types/types';
 import type { AccountId, ValidatorId } from '@polkadot/types/interfaces/runtime';
 import type { Exposure } from '@polkadot/types/interfaces/staking';
 
+/** @name BeefyKey */
+export interface BeefyKey extends U8aFixed {}
+
 /** @name FullIdentification */
 export interface FullIdentification extends Exposure {}
 
@@ -28,6 +31,12 @@ export interface SessionIndex extends u32 {}
 /** @name SessionKeys1 */
 export interface SessionKeys1 extends AccountId {}
 
+/** @name SessionKeys10 */
+export interface SessionKeys10 extends ITuple<[AccountId, AccountId, AccountId, AccountId, AccountId, AccountId, AccountId, AccountId, AccountId, AccountId]> {}
+
+/** @name SessionKeys10B */
+export interface SessionKeys10B extends ITuple<[AccountId, AccountId, AccountId, AccountId, AccountId, AccountId, AccountId, AccountId, AccountId, BeefyKey]> {}
+
 /** @name SessionKeys2 */
 export interface SessionKeys2 extends ITuple<[AccountId, AccountId]> {}
 
@@ -43,15 +52,27 @@ export interface SessionKeys5 extends ITuple<[AccountId, AccountId, AccountId, A
 /** @name SessionKeys6 */
 export interface SessionKeys6 extends ITuple<[AccountId, AccountId, AccountId, AccountId, AccountId, AccountId]> {}
 
+/** @name SessionKeys6B */
+export interface SessionKeys6B extends ITuple<[AccountId, AccountId, AccountId, AccountId, AccountId, BeefyKey]> {}
+
 /** @name SessionKeys7 */
 export interface SessionKeys7 extends ITuple<[AccountId, AccountId, AccountId, AccountId, AccountId, AccountId, AccountId]> {}
 
+/** @name SessionKeys7B */
+export interface SessionKeys7B extends ITuple<[AccountId, AccountId, AccountId, AccountId, AccountId, AccountId, BeefyKey]> {}
+
 /** @name SessionKeys8 */
 export interface SessionKeys8 extends ITuple<[AccountId, AccountId, AccountId, AccountId, AccountId, AccountId, AccountId, AccountId]> {}
 
+/** @name SessionKeys8B */
+export interface SessionKeys8B extends ITuple<[AccountId, AccountId, AccountId, AccountId, AccountId, AccountId, AccountId, BeefyKey]> {}
+
 /** @name SessionKeys9 */
 export interface SessionKeys9 extends ITuple<[AccountId, AccountId, AccountId, AccountId, AccountId, AccountId, AccountId, AccountId, AccountId]> {}
 
+/** @name SessionKeys9B */
+export interface SessionKeys9B extends ITuple<[AccountId, AccountId, AccountId, AccountId, AccountId, AccountId, AccountId, AccountId, BeefyKey]> {}
+
 /** @name ValidatorCount */
 export interface ValidatorCount extends u32 {}