fix upgrade runtime

e2e/__snapshots__/upgrade.test.ts.snap

@@ -0,0 +1,31 @@
+// Vitest Snapshot v1
+
+exports[`upgrade > setCode works 1`] = `
+{
+  "consumers": 0,
+  "data": {
+    "feeFrozen": 0,
+    "free": 1000000000000,
+    "miscFrozen": 0,
+    "reserved": 0,
+  },
+  "nonce": 0,
+  "providers": 1,
+  "sufficients": 0,
+}
+`;
+
+exports[`upgrade > setCode works 2`] = `
+{
+  "consumers": 0,
+  "data": {
+    "feeFrozen": 0,
+    "free": 2000000000000,
+    "miscFrozen": 0,
+    "reserved": 0,
+  },
+  "nonce": 0,
+  "providers": 1,
+  "sufficients": 0,
+}
+`;

e2e/upgrade.test.ts

@@ -3,15 +3,15 @@ import { describe, expect, it } from 'vitest'
 import { readFileSync } from 'node:fs'
 import path from 'node:path'
 
-import { api, chain, dev, setupApi, testingPairs } from './helper'
+import { api, chain, dev, expectJson, setupApi, testingPairs } from './helper'
 
 setupApi({
   endpoint: 'wss://acala-rpc-1.aca-api.network',
   blockHash: '0x663c25dc86521f4b7f74dcbc26224bb0fac40e316e6b0bcf6a51de373f37afac',
 })
 
 describe('upgrade', () => {
-  const { alice } = testingPairs()
+  const { alice, bob } = testingPairs()
   it('setCode works', async () => {
     await dev.setStorages({
       Sudo: {
@@ -31,5 +31,13 @@ describe('upgrade', () => {
     await dev.newBlock()
     await dev.newBlock()
     expect(await chain.head.runtimeVersion).toContain({ specVersion: 2101 })
+
+    await api.tx.balances.transfer(bob.address, 1e12).signAndSend(alice)
+    await dev.newBlock()
+    await expectJson(api.query.system.account(bob.address)).toMatchSnapshot()
+
+    await api.tx.balances.transfer(bob.address, 1e12).signAndSend(alice)
+    await dev.newBlock()
+    await expectJson(api.query.system.account(bob.address)).toMatchSnapshot()
   })
 })

executor/src/lib.rs

@@ -78,8 +78,12 @@ pub async fn create_proof(
 
     let root_trie_hash = serde_wasm_bindgen::from_value::<HashHexString>(root_trie_hash)?;
     let proof = serde_wasm_bindgen::from_value::<HexString>(nodes)?;
-    let entries = serde_wasm_bindgen::from_value::<Vec<(HexString, HexString)>>(entries)?;
-    let entries = BTreeMap::from_iter(entries.into_iter().map(|(key, value)| (key.0, value.0)));
+    let entries = serde_wasm_bindgen::from_value::<Vec<(HexString, Option<HexString>)>>(entries)?;
+    let entries = BTreeMap::from_iter(
+        entries
+            .into_iter()
+            .map(|(key, value)| (key.0, value.map(|x| x.0))),
+    );
     let proof = proof::create_proof(root_trie_hash, proof.0, entries)?;
     let result = serde_wasm_bindgen::to_value(&proof)?;
 

executor/src/proof.rs

@@ -37,7 +37,7 @@ pub fn decode_proof(
 pub fn create_proof(
     hash: HashHexString,
     proof: Vec<u8>,
-    entries: BTreeMap<Vec<u8>, Vec<u8>>,
+    entries: BTreeMap<Vec<u8>, Option<Vec<u8>>>,
 ) -> Result<(HashHexString, Vec<HexString>), String> {
     let config = Config::<Vec<u8>> {
         trie_root_hash: &hash.0,
@@ -48,12 +48,18 @@ pub fn create_proof(
 
     let mut trie = trie_structure::TrieStructure::new();
 
+    let mut deletes: Vec<Vec<u8>> = vec![];
+
     for (key, value) in entries {
-        trie.node(bytes_to_nibbles(key.iter().cloned()))
-            .into_vacant()
-            .unwrap()
-            .insert_storage_value()
-            .insert(value.to_vec(), vec![]);
+        if let Some(value) = value {
+            trie.node(bytes_to_nibbles(key.iter().cloned()))
+                .into_vacant()
+                .unwrap()
+                .insert_storage_value()
+                .insert(value.to_vec(), vec![]);
+        } else {
+            deletes.push(key);
+        }
     }
 
     for (key, value) in decoded.iter_ordered() {
@@ -66,6 +72,16 @@ pub fn create_proof(
         }
     }
 
+    for key in deletes {
+        if let trie_structure::Entry::Occupied(occupied) =
+            trie.node(bytes_to_nibbles(key.iter().cloned()))
+        {
+            if occupied.has_storage_value() {
+                occupied.into_storage().unwrap().remove();
+            }
+        }
+    }
+
     for node_index in trie.clone().iter_unordered() {
         let key = trie
             .node_full_key_by_index(node_index)
@@ -139,18 +155,18 @@ fn create_proof_works() {
         "4a8902b29241020b24b4a1620d0154f756b81ffbcf739a9f06d3447df8123ebd"
     ));
 
-    let entries = BTreeMap::<Vec<u8>, Vec<u8>>::from([
+    let entries = BTreeMap::<Vec<u8>, Option<Vec<u8>>>::from([
         (
             hex!("06de3d8a54d27e44a9d5ce189618f22db4b49d95320d9021994c850f25b8e385").to_vec(),
-            hex!("000030000080000008000000000010000000100005000000050000000a0000000a000000000050000000100000e8764817000000040000000400000000000000000000000000000000000000000000000000000000000000000000000800000000200000040000000400000000001000b0040000000000000000000014000000040000000400000000000000010100000000060000006400000002000000c8000000020000001900000000000000020000000200000000c817a804000000").to_vec()
+            Some(hex!("000030000080000008000000000010000000100005000000050000000a0000000a000000000050000000100000e8764817000000040000000400000000000000000000000000000000000000000000000000000000000000000000000800000000200000040000000400000000001000b0040000000000000000000014000000040000000400000000000000010100000000060000006400000002000000c8000000020000001900000000000000020000000200000000c817a804000000").to_vec())
         ),
         (
             hex!("cd710b30bd2eab0352ddcc26417aa1949e94c040f5e73d9b7addd6cb603d15d363f5a4efb16ffa83d0070000").to_vec(),
-            hex!("01").to_vec()
+            Some(hex!("01").to_vec())
         )
     ]);
 
-    let (hash, nodes) = create_proof(root, get_proof(), entries).unwrap();
+    let (hash, nodes) = create_proof(root.clone(), get_proof(), entries).unwrap();
 
     let keys = vec![
 		HexString(
@@ -177,6 +193,36 @@ fn create_proof_works() {
             hex!("d205bfd64a59c64fe84480fda7dafd773cb029530c4efe8441bf1f4332bfa48a").to_vec()
         ))
     );
+
+    // delete entries
+    let entries = BTreeMap::<Vec<u8>, Option<Vec<u8>>>::from([
+        (
+            hex!("63f78c98723ddc9073523ef3beefda0c4d7fefc408aac59dbfe80a72ac8e3ce563f5a4efb16ffa83d0070000").to_vec(),
+            None
+        ),
+    ]);
+
+    let (hash, nodes) = create_proof(root, get_proof(), entries).unwrap();
+    let keys = vec![
+		HexString(
+			hex!("63f78c98723ddc9073523ef3beefda0c4d7fefc408aac59dbfe80a72ac8e3ce563f5a4efb16ffa83d0070000").to_vec(),
+		),
+		HexString(
+        	hex!("06de3d8a54d27e44a9d5ce189618f22db4b49d95320d9021994c850f25b8e385").to_vec(),
+    	),
+		HexString(
+			hex!("cd710b30bd2eab0352ddcc26417aa1949e94c040f5e73d9b7addd6cb603d15d363f5a4efb16ffa83d0070000").to_vec()
+		)
+	];
+    let decoded = decode_proof(
+        hash,
+        keys,
+        encode_proofs(nodes.iter().map(|x| x.0.clone()).collect::<Vec<_>>()),
+    )
+    .unwrap();
+    let (key, value) = decoded[0].to_owned();
+    assert_eq!(key, HexString(hex!("63f78c98723ddc9073523ef3beefda0c4d7fefc408aac59dbfe80a72ac8e3ce563f5a4efb16ffa83d0070000").to_vec()));
+    assert_eq!(value, None);
 }
 
 #[cfg(test)]

src/blockchain/head-state.ts

@@ -1,3 +1,4 @@
+import { stringToHex } from '@polkadot/util'
 import _ from 'lodash'
 
 import { Block } from './block'
@@ -40,15 +41,16 @@ export class HeadState {
     delete this.#storageListeners[id]
   }
 
-  subscrubeRuntimeVersion(cb: (block: Block) => void) {
-    // TODO: actually subscribe
-    void cb
-    return randomId()
+  async subscrubeRuntimeVersion(cb: (block: Block) => void) {
+    const id = randomId()
+    const codeKey = stringToHex(':code')
+    this.#storageListeners[id] = [[codeKey], cb]
+    this.#oldValues[codeKey] = await this.#head.get(codeKey)
+    return id
   }
 
   unsubscribeRuntimeVersion(id: string) {
-    // TODO: actually unsubscribe
-    void id
+    delete this.#storageListeners[id]
   }
 
   async setHead(head: Block) {

src/blockchain/inherents.ts

@@ -107,14 +107,18 @@ export class SetValidationData implements CreateInherents {
         .createType<GenericExtrinsic>('GenericExtrinsic', validationDataExtrinsic)
         .args[0].toJSON() as typeof MOCK_VALIDATION_DATA
 
-      const newEntries: [HexString, HexString][] = []
-      const upgrade = await parentBlock.get(compactHex(meta.query.parachainSystem.pendingValidationCode()))
-      if (upgrade) {
-        const paraIdStorage = await parentBlock.get(compactHex(meta.query.parachainInfo.parachainId()))
-        const paraId = meta.registry.createType('u32', hexToU8a(paraIdStorage as any))
-        const upgradeKey = upgradeGoAheadSignal(paraId)
+      const newEntries: [HexString, HexString | null][] = []
+      const pendingUpgrade = await parentBlock.get(compactHex(meta.query.parachainSystem.pendingValidationCode()))
+      const paraIdStorage = await parentBlock.get(compactHex(meta.query.parachainInfo.parachainId()))
+      const paraId = meta.registry.createType('u32', hexToU8a(paraIdStorage as any))
+      const upgradeKey = upgradeGoAheadSignal(paraId)
+      if (pendingUpgrade) {
+        // send goAhead signal
         const goAhead = meta.registry.createType('UpgradeGoAhead', 'GoAhead')
         newEntries.push([upgradeKey, goAhead.toHex()])
+      } else {
+        // make sure previous goAhead is removed
+        newEntries.push([upgradeKey, null])
       }
 
       const { trieRootHash, nodes } = await createProof(

src/blockchain/txpool.ts

@@ -178,7 +178,7 @@ export class TxPool {
     finalBlock.pushStorageLayer().setAll(diff)
 
     this.#chain.unregisterBlock(newBlock)
-    this.#chain.setHead(finalBlock)
+    await this.#chain.setHead(finalBlock)
 
     logger.info({ hash: finalBlock.hash, number: finalBlock.number, prevHash: newBlock.hash }, 'Block built')
   }

src/executor.ts

@@ -59,7 +59,11 @@ export const decodeProof = async (trieRootHash: HexString, keys: HexString[], no
   }, {} as Record<HexString, HexString | null>)
 }
 
-export const createProof = async (trieRootHash: HexString, nodes: HexString[], entries: [HexString, HexString][]) => {
+export const createProof = async (
+  trieRootHash: HexString,
+  nodes: HexString[],
+  entries: [HexString, HexString | null][]
+) => {
   const result = await create_proof(trieRootHash, nodesAddLength(nodes), entries)
   return { trieRootHash: result[0] as HexString, nodes: result[1] as HexString[] }
 }

src/rpc/substrate/state.ts

@@ -43,7 +43,7 @@ const handlers: Handlers = {
   },
   state_subscribeRuntimeVersion: async (context, _params, { subscribe }) => {
     let update = (_block: Block) => {}
-    const id = context.chain.headState.subscrubeRuntimeVersion((block) => update(block))
+    const id = await context.chain.headState.subscrubeRuntimeVersion((block) => update(block))
     const callback = subscribe('state_runtimeVersion', id)
     update = async (block) => callback(await block.runtimeVersion)
     context.chain.head.runtimeVersion.then(callback)