fix(chain-spec): adjust hrmp channels json structure (#244)

When testing `zombienet-sdk` with cross-chain calls with Asset Hub, we
found that the hrmp channels passed to the chainspec caused the
following issue:

`Error: "Invalid JSON blob: invalid type: map, expected a tuple of size
4 at line 1 column 2794"`

This is because the values of `HrmpChannelConfig` serialise to a map
instead of a tuple, as expected by
https://github.com/paritytech/polkadot-sdk/blob/fc906d5d0fb7796ef54ba670101cf37b0aad6794/polkadot/runtime/parachains/src/hrmp.rs#L492.
Tests were updated accordingly to facilitate this fix.

Additional commits were also added to enable some of Asset Hub on
Polkadot functionality observed in the classic Zombienet repo, so that
we are able to launch a local Polkadot network with Asset Hub should
anyone require. Example configuration files can be found at the bottom
of https://github.com/r0gue-io/pop-cli/pull/278/files for anyone
interested, although that PR is still a work in progress and dependent
on this one being accepted.

PS - I added tests where I could, but additional tests might require
some minor refactoring to make it easier to test only the logic that is
changed by this PR, which I am not sure what the appetite is? An example
would be refactoring the file name generation in `keystore.rs` into its
own function to more easily test the outputs without the `scoped_fs`.

crates/orchestrator/src/generators/chain_spec.rs

@@ -881,7 +881,7 @@ fn add_balances(
     }
 }
 
-fn get_node_keys(node: &NodeSpec, use_stash: bool) -> GenesisNodeKey {
+fn get_node_keys(node: &NodeSpec, use_stash: bool, asset_hub_polkadot: bool) -> GenesisNodeKey {
     let sr_account = node.accounts.accounts.get("sr").unwrap();
     let sr_stash = node.accounts.accounts.get("sr_stash").unwrap();
     let ed_account = node.accounts.accounts.get("ed").unwrap();
@@ -898,6 +898,10 @@ fn get_node_keys(node: &NodeSpec, use_stash: bool) -> GenesisNodeKey {
         "nimbus",
         "vrf",
     ] {
+        if k == "aura" && asset_hub_polkadot {
+            keys.insert(k.to_string(), ed_account.address.clone());
+            continue;
+        }
         keys.insert(k.to_string(), sr_account.address.clone());
     }
 
@@ -917,10 +921,15 @@ fn add_authorities(
     nodes: &[&NodeSpec],
     use_stash: bool,
 ) {
+    let asset_hub_polkadot = chain_spec_json
+        .get("id")
+        .and_then(|v| v.as_str())
+        .map(|id| id.starts_with("asset-hub-polkadot"))
+        .unwrap_or_default();
     if let Some(val) = chain_spec_json.pointer_mut(runtime_config_ptr) {
         let keys: Vec<GenesisNodeKey> = nodes
             .iter()
-            .map(|node| get_node_keys(node, use_stash))
+            .map(|node| get_node_keys(node, use_stash, asset_hub_polkadot))
             .collect();
         val["session"]["keys"] = json!(keys);
     } else {
@@ -934,6 +943,17 @@ fn add_hrmp_channels(
 ) {
     if let Some(val) = chain_spec_json.pointer_mut(runtime_config_ptr) {
         if let Some(preopen_hrmp_channels) = val.pointer_mut("/hrmp/preopenHrmpChannels") {
+            let hrmp_channels = hrmp_channels
+                .iter()
+                .map(|c| {
+                    (
+                        c.sender(),
+                        c.recipient(),
+                        c.max_capacity(),
+                        c.max_message_size(),
+                    )
+                })
+                .collect::<Vec<_>>();
             *preopen_hrmp_channels = json!(hrmp_channels);
         } else {
             warn!("⚠️  'hrmp/preopenHrmpChannels' key not present in runtime config.");
@@ -1174,8 +1194,10 @@ mod tests {
             .unwrap();
 
         assert_eq!(new_hrmp_channels.len(), 2);
-        assert_eq!(new_hrmp_channels.first().unwrap()["sender"], 100);
-        assert_eq!(new_hrmp_channels.first().unwrap()["recipient"], 101);
+        assert_eq!(new_hrmp_channels.first().unwrap()[0], 100);
+        assert_eq!(new_hrmp_channels.first().unwrap()[1], 101);
+        assert_eq!(new_hrmp_channels.last().unwrap()[0], 101);
+        assert_eq!(new_hrmp_channels.last().unwrap()[1], 100);
     }
 
     #[test]
@@ -1203,4 +1225,71 @@ mod tests {
         // assert 'preopenHrmpChannels' is not created
         assert_eq!(new_hrmp_channels, None);
     }
+
+    #[test]
+    fn get_node_keys_works() {
+        let mut name = String::from("luca");
+        let seed = format!("//{}{name}", name.remove(0).to_uppercase());
+        let accounts = NodeAccounts {
+            accounts: generators::generate_node_keys(&seed).unwrap(),
+            seed,
+        };
+        let node = NodeSpec {
+            name,
+            accounts,
+            ..Default::default()
+        };
+
+        let sr = &node.accounts.accounts["sr"];
+        let keys = [
+            ("babe".into(), sr.address.clone()),
+            ("im_online".into(), sr.address.clone()),
+            ("parachain_validator".into(), sr.address.clone()),
+            ("authority_discovery".into(), sr.address.clone()),
+            ("para_validator".into(), sr.address.clone()),
+            ("para_assignment".into(), sr.address.clone()),
+            ("aura".into(), sr.address.clone()),
+            ("nimbus".into(), sr.address.clone()),
+            ("vrf".into(), sr.address.clone()),
+            (
+                "grandpa".into(),
+                node.accounts.accounts["ed"].address.clone(),
+            ),
+            ("beefy".into(), node.accounts.accounts["ec"].address.clone()),
+        ]
+        .into();
+
+        // Stash
+        let sr_stash = &node.accounts.accounts["sr_stash"];
+        let node_key = get_node_keys(&node, true, false);
+        assert_eq!(node_key.0, sr_stash.address);
+        assert_eq!(node_key.1, sr_stash.address);
+        assert_eq!(node_key.2, keys);
+        // Non-stash
+        let node_key = get_node_keys(&node, false, false);
+        assert_eq!(node_key.0, sr.address);
+        assert_eq!(node_key.1, sr.address);
+        assert_eq!(node_key.2, keys);
+    }
+
+    #[test]
+    fn get_node_keys_supports_asset_hub_polkadot() {
+        let mut name = String::from("luca");
+        let seed = format!("//{}{name}", name.remove(0).to_uppercase());
+        let accounts = NodeAccounts {
+            accounts: generators::generate_node_keys(&seed).unwrap(),
+            seed,
+        };
+        let node = NodeSpec {
+            name,
+            accounts,
+            ..Default::default()
+        };
+
+        let node_key = get_node_keys(&node, false, false);
+        assert_eq!(node_key.2["aura"], node.accounts.accounts["sr"].address);
+
+        let node_key = get_node_keys(&node, false, true);
+        assert_eq!(node_key.2["aura"], node.accounts.accounts["ed"].address);
+    }
 }

crates/orchestrator/src/generators/keystore.rs

@@ -20,6 +20,7 @@ pub async fn generate<'a, T>(
     acc: &NodeAccounts,
     node_files_path: impl AsRef<Path>,
     scoped_fs: &ScopedFilesystem<'a, T>,
+    asset_hub_polkadot: bool,
 ) -> Result<Vec<PathBuf>, GeneratorError>
 where
     T: FileSystem,
@@ -32,10 +33,15 @@ where
         // let filename = encode(k);
 
         let filename = match k {
-            // TODO: add logic for isAssetHubPolkadot
-            // "aura" => {
-            //     ""
-            // },
+            "aura" if asset_hub_polkadot => {
+                let pk = acc
+                    .accounts
+                    .get("ed")
+                    .expect(&format!("Key 'ed' should be set for node {THIS_IS_A_BUG}"))
+                    .public_key
+                    .as_str();
+                format!("{}{}", encode(k), pk)
+            },
             "babe" | "imon" | "audi" | "asgn" | "para" | "nmbs" | "rand" | "aura" => {
                 let pk = acc
                     .accounts

crates/orchestrator/src/spawner.rs

@@ -61,10 +61,18 @@ where
         } else {
             node.name.clone()
         };
-        let key_filenames =
-            generators::generate_node_keystore(&node.accounts, &node_files_path, ctx.scoped_fs)
-                .await
-                .unwrap();
+        let asset_hub_polkadot = ctx
+            .parachain_id
+            .map(|id| id.starts_with("asset-hub-polkadot"))
+            .unwrap_or_default();
+        let key_filenames = generators::generate_node_keystore(
+            &node.accounts,
+            &node_files_path,
+            ctx.scoped_fs,
+            asset_hub_polkadot,
+        )
+        .await
+        .unwrap();
 
         // Paths returned are relative to the base dir, we need to convert into
         // fullpaths to inject them in the nodes.