API proposal for custom proofs in substrate

proof.rs

mod sp_io {
    /// Determines in which exection mode we are.
    enum Mode {
        /// Standalone chain that does not have to deal with PoV at all.
        Standalone,
        /// Running on a parachain node (block building).
        Parachain,
        /// Running on a relay chain validator (block verification).
        Pvf,
    }

    /// Determines in which exection mode we are.
    ///
    /// # Note
    ///
    /// I think we want to keep this function private and don't provide it to the rest of substrate.
    /// Otherwise runtime developers could easily produce code which gets incompatible with
    /// parachains. Instead, we want to provide abstractions like `get_with_custom_proof`.
    fn mode() -> Mode {
        // This would probably call some host function or access some global variable.
        unimplemented!()
    }

    /// Reads the custom proof indexed by key from the PoV,
    fn get_proof_for_key(_key: &[u8]) -> Option<Vec<u8>> {
        // This would look that up in some extra map in `ParachainHeadData`.
        unimplemented!()
    }

    /// Removes the access to `key` from the main witness and replaces it by a custom proof.
    fn replace_access(_key: &[u8], _proof: Vec<u8>) {
        // This would propably call some host function to make that happen.
        unimplemented!()
    }

    /// Normal storage access (redirected to PoV in PvF case).
    fn get(_key: &[u8]) -> Option<Vec<u8>> {
        unimplemented!()
    }

    /// This is the API used by the runtime to access a storage which should not use the automatic
    /// recording of storage accesses.
    ///
    /// # Note
    ///
    /// We might want to create a higher level abstraction on top of this (in `frame_support`) which
    /// does not require creating a custom proof but provides a default proof system (trie).
    pub fn access_with_custom_proof<R, B, V>(
        key: &[u8],
        build_proof: B,
        validate_proof: V,
    ) -> Option<R>
    where
        // storage_value, out_proof -> result
        B: Fn(Vec<u8>, Option<&mut Vec<u8>>) -> R,
        // proof -> storage_value
        V: Fn(Vec<u8>) -> R,
    {
        match mode() {
            // In a standalone case no proof is collected hence `None` is passed.
            Mode::Standalone => get(key).map(|value| build_proof(value, None)),
            // We have access to storage to build a block. Get the value from storage
            // and create a custom proof from it. This proof will be written to the PoV
            // through `replace_access`. The original access is removed.
            Mode::Parachain => get(key).map(|value| {
                let mut proof = Vec::new();
                let result = build_proof(value, Some(&mut proof));
                replace_access(key, proof);
                result
            }),
            // We have no access to the original storage. Reconstruct the storage value
            // from the custom proof which is read from the PoV by `get_proof_for_key`.
            Mode::Pvf => get_proof_for_key(key).map(validate_proof),
        }
    }
}

mod pallet_contracts {
    pub enum DispatchError {
        CodeNotFound,
    }

    type ContractResult = Result<u32, DispatchError>;

    /// Load the merkle root of a contract code hash from a seperate storage map.
    ///
    /// It was stored when the contract was deployed.
    fn code_merkle_root(_code_hash: &[u8]) -> Vec<u8> {
        unimplemented!()
    }

    /// Execute the supplied wasm blob and create the proof while doing so.
    ///
    /// We need to execute the code to create the proof: While executing we gather data
    /// about which parts of the contract where actually executed. From that information
    /// we can create a witness/proof with only the parts that were executed.
    fn execute(_wasm_code: &[u8]) -> (Vec<u8>, u32) {
        unimplemented!()
    }

    /// Create a proof/witness that can be emitted into the PoV.
    fn build_proof(_merkle_root: &[u8], _coverage: &[u8], _wasm_code: &[u8]) -> Vec<u8> {
        unimplemented!()
    }

    /// Reconstruct a wasm module from the witness.
    ///
    /// # Note
    ///
    /// The validation can fail here when a collator sends an invalid proof as part of the
    /// PoV. The problem is that the logic to verify this proof sits in the runtime (here).
    /// How would we tell the client that the transaction is invalid?
    fn validate_proof(_merkle_root: &[u8], _proof: &[u8]) -> Option<Vec<u8>> {
        unimplemented!()
    }

    /// This is the contract call dispatchable. It loads the contract from storage and executes it.
    ///
    /// `code_hash` is the hash of the wasm blob and used to index the wasm blobs in storage.
    pub fn call_contract(code_hash: &[u8]) -> ContractResult {
        // We assume the wasm blob is stored directly at `code_hash`.
        let exit_code = super::sp_io::access_with_custom_proof(
            code_hash,
            |wasm_code, proof| {
                let (coverage, exit_code) = execute(&wasm_code);
                if let Some(proof) = proof {
                    let merkle_root = code_merkle_root(code_hash);
                    *proof = build_proof(&merkle_root, &coverage, &wasm_code);
                }
                exit_code
            },
            |proof| {
                let merkle_root = code_merkle_root(code_hash);
                let wasm_code =
                    validate_proof(&merkle_root, &proof).expect("See comment on `validate_proof`");
                let (_, exit_code) = execute(&wasm_code);
                exit_code
            },
        )
        .ok_or(DispatchError::CodeNotFound)?;

        Ok(exit_code)
    }
}

fn main() {
    let code = Vec::new();
    pallet_contracts::call_contract(&code).ok();
}