From 3c0bfcbfd1d05faf26af5b1930f5bad546644ee5 Mon Sep 17 00:00:00 2001
From: hrxi <hrrrxi@gmail.com>
Date: Thu, 23 Jan 2025 17:48:42 +0100
Subject: [PATCH] Store decompressed BLS keys for genesis block

---
 genesis-builder/src/lib.rs | 37 +++++++++++++++++--------------------
 genesis/build.rs           |  1 +
 genesis/src/networks.rs    | 28 +++++++++++++++++++++++++++-
 3 files changed, 45 insertions(+), 21 deletions(-)

diff --git a/genesis-builder/src/lib.rs b/genesis-builder/src/lib.rs
index 7dff82774c..fd9e13e78d 100644
--- a/genesis-builder/src/lib.rs
+++ b/genesis-builder/src/lib.rs
@@ -1,11 +1,7 @@
 #[macro_use]
 extern crate log;
 
-use std::{
-    fs::{read_to_string, OpenOptions},
-    io::Error as IoError,
-    path::Path,
-};
+use std::{fs, io::Error as IoError, path::Path};
 
 use nimiq_account::{
     Account, Accounts, BasicAccount, HashedTimeLockedContract, StakingContract,
@@ -217,7 +213,7 @@ impl GenesisBuilder {
     ///
     /// See `genesis/src/genesis/unit-albatross.toml` for an example.
     pub fn from_config_file<P: AsRef<Path>>(path: P) -> Result<Self, GenesisBuilderError> {
-        Self::from_config(toml::from_str(&read_to_string(path)?)?)
+        Self::from_config(toml::from_str(&fs::read_to_string(path)?)?)
     }
 
     pub fn from_config(config: config::GenesisConfig) -> Result<Self, GenesisBuilderError> {
@@ -647,6 +643,8 @@ impl GenesisBuilder {
         db: MdbxDatabase,
         directory: P,
     ) -> Result<(Blake2bHash, bool), GenesisBuilderError> {
+        let directory = directory.as_ref();
+
         let GenesisInfo {
             block,
             hash,
@@ -658,25 +656,24 @@ impl GenesisBuilder {
         debug!("Accounts:");
         debug!(?accounts);
 
-        let block_path = directory.as_ref().join("block.dat");
+        let block_path = directory.join("block.dat");
         info!(path = %block_path.display(), "Writing block to");
-        let mut file = OpenOptions::new()
-            .create(true)
-            .truncate(true)
-            .write(true)
-            .open(&block_path)?;
-        block.serialize_to_writer(&mut file)?;
+        fs::write(block_path, block.serialize_to_vec())?;
+
+        {
+            let decompressed_path = directory.join("decompressed_keys.dat");
+            let mut decompressed = Vec::new();
+            for key in block.validators().expect("must be election").voting_keys() {
+                decompressed.extend_from_slice(&key.trusted_serialize());
+            }
+            fs::write(decompressed_path, decompressed)?;
+        }
 
         let have_accounts = accounts.is_some();
         if let Some(accounts) = accounts {
-            let accounts_path = directory.as_ref().join("accounts.dat");
+            let accounts_path = directory.join("accounts.dat");
             info!(path = %accounts_path.display(), "Writing accounts to");
-            let mut file = OpenOptions::new()
-                .create(true)
-                .truncate(true)
-                .write(true)
-                .open(&accounts_path)?;
-            accounts.serialize_to_writer(&mut file)?;
+            fs::write(accounts_path, accounts.serialize_to_vec())?;
         }
 
         Ok((hash, have_accounts))
diff --git a/genesis/build.rs b/genesis/build.rs
index 3f7320399f..3e6f169c57 100644
--- a/genesis/build.rs
+++ b/genesis/build.rs
@@ -23,6 +23,7 @@ fn write_genesis_rs(directory: &Path, name: &str, genesis_hash: &Blake2bHash, ha
     let genesis_rs = format!(
         r#"GenesisData {{
             block: include_bytes!(concat!(env!("OUT_DIR"), "/genesis/{name}/block.dat")),
+            decompressed_keys: include_bytes!(concat!(env!("OUT_DIR"), "/genesis/{name}/decompressed_keys.dat")),
             hash: Blake2bHash([{hash}]),
             accounts: {accounts_expr},
     }}"#,
diff --git a/genesis/src/networks.rs b/genesis/src/networks.rs
index 80a88bb047..60696de64b 100644
--- a/genesis/src/networks.rs
+++ b/genesis/src/networks.rs
@@ -1,8 +1,9 @@
-use std::env;
 #[cfg(feature = "genesis-override")]
 use std::path::Path;
+use std::{env, sync::OnceLock};
 
 use nimiq_block::Block;
+use nimiq_bls::{LazyPublicKey as BlsLazyPublicKey, PublicKey as BlsPublicKey};
 #[cfg(feature = "genesis-override")]
 use nimiq_database::mdbx::MdbxDatabase;
 #[cfg(feature = "genesis-override")]
@@ -17,6 +18,7 @@ use nimiq_serde::Serialize;
 #[derive(Clone, Debug)]
 struct GenesisData {
     block: &'static [u8],
+    decompressed_keys: &'static [u8],
     hash: Blake2bHash,
     accounts: Option<&'static [u8]>,
 }
@@ -81,16 +83,40 @@ fn read_genesis_config(config: &Path) -> Result<GenesisData, GenesisBuilderError
 
     Ok(GenesisData {
         block: Box::leak(block.into_boxed_slice()),
+        decompressed_keys: &[],
         hash,
         accounts: accounts.map(|accounts| Box::leak(accounts.into_boxed_slice()) as &'static _),
     })
 }
 
+static KEYS_DEV: OnceLock<Box<[BlsLazyPublicKey]>> = OnceLock::new();
+static KEYS_TEST: OnceLock<Box<[BlsLazyPublicKey]>> = OnceLock::new();
+static KEYS_UNIT: OnceLock<Box<[BlsLazyPublicKey]>> = OnceLock::new();
+static KEYS_MAIN: OnceLock<Box<[BlsLazyPublicKey]>> = OnceLock::new();
+
 fn network(network_id: NetworkId) -> Option<&'static NetworkInfo> {
     let result = network_impl(network_id);
     if let Some(info) = result {
         assert_eq!(network_id, info.network_id);
         assert_eq!(network_id, info.genesis_block().network());
+        let keys = match network_id {
+            NetworkId::DevAlbatross => &KEYS_DEV,
+            NetworkId::TestAlbatross => &KEYS_TEST,
+            NetworkId::UnitAlbatross => &KEYS_UNIT,
+            NetworkId::MainAlbatross => &KEYS_MAIN,
+            _ => unreachable!(),
+        };
+        keys.get_or_init(|| {
+            info.genesis
+                .decompressed_keys
+                .chunks(BlsPublicKey::TRUSTED_SERIALIZATION_SIZE)
+                .map(|chunk| {
+                    BlsLazyPublicKey::from(BlsPublicKey::trusted_deserialize(
+                        &chunk.try_into().unwrap(),
+                    ))
+                })
+                .collect()
+        });
     }
     result
 }