feat: add support for custom genesis state

Cargo.toml

@@ -3,6 +3,7 @@ members = [
   # Binaries
   "core/bin/block_reverter",
   "core/bin/contract-verifier",
+  "core/bin/custom_genesis_export",
   "core/bin/external_node",
   "core/bin/merkle_tree_consistency_checker",
   "core/bin/snapshots_creator",

core/bin/custom_genesis_export/.gitignore

@@ -0,0 +1 @@
+*.bin

core/bin/custom_genesis_export/Cargo.toml

@@ -0,0 +1,29 @@
+[package]
+name = "custom_genesis_export"
+version.workspace = true
+edition.workspace = true
+authors.workspace = true
+homepage.workspace = true
+repository.workspace = true
+license.workspace = true
+keywords.workspace = true
+categories.workspace = true
+
+[dependencies]
+clap = { workspace = true, features = ["derive"] }
+futures.workspace = true
+sqlx = { workspace = true, features = [
+    "runtime-tokio",
+    "tls-native-tls",
+    "macros",
+    "postgres",
+] }
+tokio = { workspace = true, features = ["full"] }
+
+zksync_types.workspace = true
+zksync_node_genesis.workspace = true
+zksync_contracts.workspace = true
+zksync_core_leftovers.workspace = true
+zksync_protobuf_config.workspace = true
+anyhow.workspace = true
+

core/bin/custom_genesis_export/README.md

@@ -0,0 +1,62 @@
+# Custom Genesis Export
+
+The `custom_genesis_export` tool allows exporting a state from a zkSync PostgreSQL database in a format that can be
+included as a custom genesis state for a new chain.
+
+This is particularly useful in data migration scenarios where a large existing state needs to be applied to a newly
+created chain.
+
+A typical workflow could be:
+
+- Run a chain locally, not connected to the real L1, and add all required data to it.
+- Export the data using the `custom_genesis_export` tool.
+- Create a new chain connected to the real ecosystem using the exported data.
+
+## How it works
+
+The tool exports all entries from `initial_writes`, `storage_logs`, and `factory_deps`, except those related to the
+system context. The data is then written to a binary file using the Rust standard library following a simple
+serialisation format.
+
+`custom_genesis_export` can be built using the following command:
+
+```shell
+cargo build --release -p custom_genesis_export
+```
+
+And then executed using the following command, where:
+
+- `database-url` is the URL of the PostgreSQL database.
+- `genesis-config-path` is the path to the `genesis.yaml` configuration file, used to set up a new chain (located in the
+  `file_based` directory).
+- `output-path` is the path to the generated binary output file.
+
+```shell
+custom_genesis_export --database-url=postgres://postgres:notsecurepassword@localhost:5432/zksync_server_localhost_validium --genesis-config-path=/Users/ischasny/Dev/zksync-era/etc/env/file_based/genesis.yaml --output-path=export.bin
+```
+
+> Please make sure that the database is not written into before running data export.
+
+After the export is completed, the tool will make the following updates to the `genesis.yaml` file in-place:
+
+- Update `genesis_root_hash`, `rollup_last_leaf_index`, and `genesis_commitment` to match the contents of the export
+  file.
+- Add a `custom_genesis_state_path` property pointing to the data export.
+
+The modified genesis file can be used to bootstrap an ecosystem or initialize new chains. The data export will be
+automatically recognized by the server during the execution of `zkstack ecosystem init ...` and
+`zkstack chain create ...` commands.
+
+### Running considerations
+
+- All chains within the same ecosystem must be bootstrapped from the same genesis state. This is enforced at the
+  protocol level. If two chains require different states, this can only be achieved by bringing the chain into the
+  ecosystem through governance voting.
+  - If a chain is added to the ecosystem via a vote, ensure no assets are minted on the old bridge, as this would create
+    discrepancies with the new one. One can consider setting gas prices to zero during when generating a state to
+    account for that.
+- To calculate genesis parameters, the tool must load all VM logs into RAM. This is due to implementation specifics. For
+  larger states, ensure the VM has sufficient RAM capacity.
+- After the import, block numbers for all VM logs will be reset to zero - if the imported data has been indexed based on
+  block number, such indexes will break.
+- External Nodes will have to be bootstrapped from data snapshot (i.e. genesis can't be generated locally).

core/bin/custom_genesis_export/src/main.rs

@@ -0,0 +1,239 @@
+use std::{
+    fs,
+    fs::File,
+    io::{BufWriter, Write},
+    path::PathBuf,
+};
+
+use clap::Parser;
+use futures::TryStreamExt;
+use sqlx::{prelude::*, Connection, PgConnection};
+use zksync_contracts::BaseSystemContractsHashes;
+use zksync_core_leftovers::temp_config_store::read_yaml_repr;
+use zksync_node_genesis::make_genesis_batch_params;
+use zksync_protobuf_config::encode_yaml_repr;
+use zksync_types::{AccountTreeId, StorageKey, StorageLog, H160, H256};
+
+#[derive(Debug, Parser)]
+#[command(name = "Custom genesis export tool", author = "Matter Labs")]
+struct Args {
+    /// PostgreSQL connection string for the database to export.
+    #[arg(short, long)]
+    database_url: Option<String>,
+
+    /// Output file path.
+    #[arg(short, long, default_value = "genesis_export.bin")]
+    output_path: PathBuf,
+
+    /// Path to the genesis.yaml
+    #[arg(short, long)]
+    genesis_config_path: PathBuf,
+}
+#[derive(FromRow)]
+struct InitialWriteRow {
+    hashed_key: [u8; 32],
+    index: i64,
+}
+#[derive(FromRow)]
+struct StorageLogRow {
+    address: [u8; 20],
+    key: [u8; 32],
+    value: [u8; 32],
+}
+#[derive(FromRow)]
+struct FactoryDepRow {
+    bytecode_hash: [u8; 32],
+    bytecode: Vec<u8>,
+}
+
+/// The `custom_genesis_export` tool allows exporting initial writes, storage logs, and factory dependencies
+/// from the ZKSync PostgreSQL database in a way that they can be used as a custom genesis state for a new chain.
+///
+/// Inputs:
+///     * `database_url` - URL to the PostgreSQL database.
+///     * `output` - Path to the output file.
+///     * `genesis_config_path` - Path to the `genesis.yaml` configuration file, which will be used to set up a new chain (located in the `file_based` directory).
+///
+/// Given the inputs above, `custom_genesis_export` will perform the following:
+///     * Read initial writes, storage logs, and factory dependencies; filter out those related to the system context,
+///       and save the remaining data to the output file.
+///     * Calculate the new `genesis_root_hash`, `rollup_last_leaf_index`, and `genesis_commitment`, then update these
+///       in-place in the provided `genesis.yaml`. Additionally, the tool will add a `custom_genesis_state_path` property
+///       pointing to the genesis export.
+///
+/// Note: To calculate the new genesis parameters, the current implementation requires loading all storage logs
+/// into RAM. This is necessary due to the specific sorting and filtering that need to be applied.
+/// For larger states, keep this in mind and ensure you have a machine with sufficient RAM.
+#[tokio::main]
+async fn main() -> anyhow::Result<()> {
+    let args = Args::parse();
+
+    let mut out = BufWriter::new(File::create(&args.output_path)?);
+
+    println!(
+        "Export file: {}",
+        args.output_path.canonicalize()?.display(),
+    );
+
+    println!("Connecting to source database...");
+    let mut conn_source =
+        PgConnection::connect(&args.database_url.or_else(|| std::env::var("DATABASE_URL").ok()).expect("Specify the database connection string in either a CLI argument or in the DATABASE_URL environment variable."))
+            .await?;
+    println!("Connected to source database.");
+
+    println!("Reading initial writes...");
+    let count_initial_writes: i64 = sqlx::query("select count(*) from initial_writes;")
+        .fetch_one(&mut conn_source)
+        .await?
+        .get(0);
+    let mut initial_writes =
+        sqlx::query_as::<_, InitialWriteRow>("select hashed_key, index from initial_writes;")
+            .fetch(&mut conn_source);
+
+    // write count of initial writes
+    out.write_all(&i64::to_le_bytes(count_initial_writes))?;
+    let mut actual_initial_writes_count = 0;
+    while let Some(r) = initial_writes.try_next().await? {
+        out.write_all(&r.hashed_key)?;
+        out.write_all(&r.index.to_le_bytes())?;
+        actual_initial_writes_count += 1;
+    }
+    if actual_initial_writes_count != count_initial_writes {
+        panic!("Database reported {count_initial_writes} initial writes; only received {actual_initial_writes_count} for export.");
+    }
+    drop(initial_writes);
+
+    println!("Exported {count_initial_writes} initial writes.");
+
+    println!("Reading storage logs...");
+
+    // skipping system context-related entries
+    let count_storage_logs: i64 = sqlx::query(
+        r#"
+        select count(distinct hashed_key) from storage_logs
+        where address <> '\x000000000000000000000000000000000000800b'::bytea or
+            key in (
+                '\x0000000000000000000000000000000000000000000000000000000000000000'::bytea,
+                '\x0000000000000000000000000000000000000000000000000000000000000003'::bytea,
+                '\x0000000000000000000000000000000000000000000000000000000000000004'::bytea,
+                '\x0000000000000000000000000000000000000000000000000000000000000005'::bytea
+            );"#,
+    )
+    .fetch_one(&mut conn_source)
+    .await?
+    .get(0);
+    out.write_all(&i64::to_le_bytes(count_storage_logs))?;
+
+    let mut storage_logs = sqlx::query_as::<_, StorageLogRow>(
+        r#"
+        select address, key, value
+        from storage_logs sl
+        where miniblock_number = (select max(miniblock_number) from storage_logs where hashed_key = sl.hashed_key)
+          and (
+            address <> '\x000000000000000000000000000000000000800b'::bytea or
+            key in (
+                    '\x0000000000000000000000000000000000000000000000000000000000000000'::bytea,
+                    '\x0000000000000000000000000000000000000000000000000000000000000003'::bytea,
+                    '\x0000000000000000000000000000000000000000000000000000000000000004'::bytea,
+                    '\x0000000000000000000000000000000000000000000000000000000000000005'::bytea
+                )
+            );"#,
+    )
+    .fetch(&mut conn_source);
+
+    let mut actual_storage_logs_count = 0;
+
+    // we need to keep this collection in memory to calculate hashes for genesis in the end
+    let mut storage_logs_for_genesis: Vec<StorageLog> =
+        Vec::with_capacity(count_storage_logs as usize);
+
+    while let Some(r) = storage_logs.try_next().await? {
+        out.write_all(&r.address)?;
+        out.write_all(&r.key)?;
+        out.write_all(&r.value)?;
+        actual_storage_logs_count += 1;
+        storage_logs_for_genesis.push(r.into());
+    }
+    if actual_storage_logs_count != count_storage_logs {
+        panic!("Retrieved {actual_storage_logs_count} storage logs from the database; expected {count_storage_logs}.");
+    }
+
+    println!("Exported {count_storage_logs} storage logs from source database.");
+
+    drop(storage_logs);
+
+    println!("Loading factory deps from source database...");
+    let count_factory_deps: i64 = sqlx::query("select count(*) from factory_deps;")
+        .fetch_one(&mut conn_source)
+        .await?
+        .get(0);
+    out.write_all(&i64::to_le_bytes(count_factory_deps))?;
+
+    let mut factory_deps =
+        sqlx::query_as::<_, FactoryDepRow>("select bytecode_hash, bytecode from factory_deps;")
+            .fetch(&mut conn_source);
+
+    let mut actual_factory_deps_count = 0;
+    while let Some(r) = factory_deps.try_next().await? {
+        out.write_all(&r.bytecode_hash)?;
+        out.write_all(&(r.bytecode.len() as u64).to_le_bytes())?;
+        out.write_all(&r.bytecode)?;
+        actual_factory_deps_count += 1;
+    }
+    if actual_factory_deps_count != count_factory_deps {
+        panic!("Retrieved {actual_factory_deps_count} factory deps from the database; expected {count_factory_deps}.");
+    }
+    drop(factory_deps);
+
+    println!("Exported {count_factory_deps} factory deps from source database.");
+
+    conn_source.close().await?;
+
+    println!("Calculating new genesis parameters");
+
+    let mut genesis_config = read_yaml_repr::<zksync_protobuf_config::proto::genesis::Genesis>(
+        &args.genesis_config_path,
+    )?;
+
+    let base_system_contract_hashes = BaseSystemContractsHashes {
+        bootloader: genesis_config
+            .bootloader_hash
+            .ok_or(anyhow::anyhow!("No bootloader_hash specified"))?,
+        default_aa: genesis_config
+            .default_aa_hash
+            .ok_or(anyhow::anyhow!("No default_aa_hash specified"))?,
+        evm_emulator: genesis_config.evm_emulator_hash,
+    };
+
+    let (genesis_batch_params, _) = make_genesis_batch_params(
+        storage_logs_for_genesis.as_slice(),
+        base_system_contract_hashes,
+        genesis_config
+            .protocol_version
+            .ok_or(anyhow::anyhow!("No bootloader_hash specified"))?
+            .minor,
+    );
+
+    genesis_config.genesis_root_hash = Some(genesis_batch_params.root_hash);
+    genesis_config.rollup_last_leaf_index = Some(genesis_batch_params.rollup_last_leaf_index);
+    genesis_config.genesis_commitment = Some(genesis_batch_params.commitment);
+    genesis_config.custom_genesis_state_path =
+        args.output_path.canonicalize()?.to_str().map(String::from);
+
+    let bytes =
+        encode_yaml_repr::<zksync_protobuf_config::proto::genesis::Genesis>(&genesis_config)?;
+    fs::write(&args.genesis_config_path, &bytes)?;
+
+    println!("Done.");
+
+    Ok(())
+}
+
+impl From<StorageLogRow> for StorageLog {
+    fn from(value: StorageLogRow) -> Self {
+        StorageLog::new_write_log(
+            StorageKey::new(AccountTreeId::new(H160(value.address)), H256(value.key)),
+            H256(value.value),
+        )
+    }
+}

core/lib/config/src/configs/genesis.rs

@@ -32,6 +32,7 @@ pub struct GenesisConfig {
     pub fee_account: Address,
     pub dummy_verifier: bool,
     pub l1_batch_commit_data_generator_mode: L1BatchCommitmentMode,
+    pub custom_genesis_state_path: Option<String>,
 }
 
 impl GenesisConfig {
@@ -60,6 +61,7 @@ impl GenesisConfig {
             l2_chain_id: L2ChainId::default(),
             dummy_verifier: false,
             l1_batch_commit_data_generator_mode: L1BatchCommitmentMode::Rollup,
+            custom_genesis_state_path: None,
         }
     }
 }

core/lib/config/src/testonly.rs

@@ -746,6 +746,7 @@ impl Distribution<configs::GenesisConfig> for EncodeDist {
                 0 => L1BatchCommitmentMode::Rollup,
                 _ => L1BatchCommitmentMode::Validium,
             },
+            custom_genesis_state_path: None,
         }
     }
 }