riftenlabs-indexer/src/chain.rs

// Copyright (C) 2022 Roman Zeyde
// Copyright (C) 2023 Bitcoin Unlimited
// Copyright (C) 2024 Riften Labs AS
//
// This software is licensed under the GNU Affero General Public License (AGPL), version 3.0 or later.
// A copy of the license can be found in the LICENSE file or at https://www.gnu.org/licenses/agpl-3.0.html

use anyhow::{bail, Context, Result};
use bitcoin_hashes::{hex::ToHex, Hash};
use bitcoincash::{consensus::deserialize, BlockHash, BlockHeader};
use electrum_client::{Batch, Client, ElectrumApi, Param};
use log::{info, trace};
use rayon::iter::{IndexedParallelIterator, IntoParallelIterator, ParallelIterator};

use serde_json::Value;
use std::{collections::HashMap, sync::RwLock};

use crate::{
    db::{config::config_set, delete_entries_for_block, DBPool},
    KEY_LAST_INDEXED,
};

/// A new header found, to be added to the chain at specific height
pub(crate) struct NewHeader {
    pub header: BlockHeader,
    pub hash: BlockHash,
    pub height: u64,
}

impl NewHeader {
    pub(crate) fn from_header_and_height(header: BlockHeader, height: u64) -> Self {
        let hash = header.block_hash();
        Self {
            header,
            hash,
            height,
        }
    }
}

pub trait BlockUndoer {
    fn undo_block(&self, blockheader: &BlockHeader) -> Result<()>;
}

pub struct StoreBlockUndoer {
    db: DBPool,
}

impl StoreBlockUndoer {
    pub fn new(db: DBPool) -> Result<Self> {
        Ok(Self { db })
    }
}

impl BlockUndoer for StoreBlockUndoer {
    fn undo_block(&self, blockheader: &BlockHeader) -> Result<()> {
        let conn = self.db.get().unwrap();
        if delete_entries_for_block(&conn, &blockheader.block_hash())? {
            // This block was indexed. Need to update indexing flag.
            config_set(
                &conn,
                KEY_LAST_INDEXED,
                &blockheader.prev_blockhash.to_hex(),
            );
        }
        Ok(())
    }
}

pub struct DummyBlockUndoer {
    allow_undo_calls: bool,
}

impl DummyBlockUndoer {
    pub fn new(allow_undo_calls: bool) -> Self {
        Self { allow_undo_calls }
    }
}
impl Default for DummyBlockUndoer {
    fn default() -> Self {
        Self::new(false)
    }
}

impl BlockUndoer for DummyBlockUndoer {
    fn undo_block(&self, _blockheader: &BlockHeader) -> Result<()> {
        if !self.allow_undo_calls {
            panic!("undo_block called on DummyBlockUndoer")
        } else {
            Ok(())
        }
    }
}

/// Current blockchain headers' list
pub struct Chain {
    headers: RwLock<Vec<(BlockHash, BlockHeader)>>,
    heights: RwLock<HashMap<BlockHash, u64>>,
}

impl Chain {
    // create an empty chain
    pub fn new(genesis: BlockHeader) -> Self {
        let genesis_hash = genesis.block_hash();
        Self {
            headers: RwLock::new(vec![(genesis_hash, genesis)]),
            heights: RwLock::new(std::iter::once((genesis_hash, 0)).collect()), // genesis header @ zero height
        }
    }

    // Constructor for unit tests.
    #[cfg(test)]
    pub fn new_regtest() -> Self {
        Self::new(
            bitcoincash::blockdata::constants::genesis_block(bitcoincash::Network::Regtest).header,
        )
    }

    #[cfg(test)]
    pub(crate) fn drop_last_headers(&mut self, undoer: impl BlockUndoer, n: u64) -> Result<()> {
        if n == 0 {
            return Ok(());
        }
        let new_height = self.height().saturating_sub(n);
        info!("Dropping {n} headers, undoing to block {new_height}");
        self.update(undoer, vec![], Some(new_height))
    }

    /// Load the chain from a collecion of headers, up to the given tip
    pub(crate) fn load(&mut self, headers: Vec<(BlockHeader, u64)>) -> Result<()> {
        let tip = headers.iter().max_by(|a, b| a.1.cmp(&b.1));

        #[allow(clippy::clone_on_copy)]
        let tip = if let Some(t) = tip {
            t.clone()
        } else {
            return Ok(());
        };

        // headers might contain blocks that have reorged out of the chain, so we need to build chain in reverse from chain.

        let mut header_map: HashMap<BlockHash, NewHeader> = headers
            .into_par_iter()
            .map(|(header, height)| {
                (
                    header.block_hash(),
                    NewHeader::from_header_and_height(header, height),
                )
            })
            .collect();

        let mut header_chain = vec![header_map.remove(&tip.0.block_hash()).unwrap()];

        loop {
            let prev = &header_chain.last().unwrap().header.prev_blockhash;
            if prev == &BlockHash::all_zeros() {
                // genesis
                break;
            }
            header_chain.push(header_map.remove(prev).context(format!(
                "missing header {} at height {}",
                prev.to_hex(),
                header_chain.last().unwrap().height - 1
            ))?);
        }

        header_chain.reverse();

        let dummy_undoer = DummyBlockUndoer::new(true);
        self.update(dummy_undoer, header_chain, None)
    }

    /// Get the block hash at specified height (if exists)
    #[cfg(test)]
    pub fn get_block_hash(&self, height: u64) -> Option<BlockHash> {
        self.headers
            .read()
            .unwrap()
            .get(height as usize)
            .map(|(hash, _header)| *hash)
    }

    /// Get the block header at specified height (if exists)
    #[cfg(test)]
    pub(crate) fn get_block_header(&self, height: u64) -> Option<BlockHeader> {
        self.headers
            .read()
            .unwrap()
            .get(height as usize)
            .map(|(_hash, header)| header.clone())
    }

    /// Get the block height given the specified hash (if exists)
    pub(crate) fn get_block_height(&self, blockhash: &BlockHash) -> Option<u64> {
        self.heights.read().unwrap().get(blockhash).copied()
    }

    pub(crate) fn get_mtp(&self, height: u64) -> Option<u64> {
        let mut times = Vec::with_capacity(11);

        {
            let headers = self.headers.read().unwrap();
            for h in height.saturating_sub(10)..=height {
                let (_, header) = headers.get(h as usize)?;
                times.push(header.time as u64);
            }
        }

        if times.is_empty() {
            debug_assert!(false, "should always get mtp if block exists");
            return None;
        }

        times.sort_unstable();
        let mtp = times[times.len() / 2];

        Some(mtp)
    }

    /// Update the chain with a list of new headers (possibly a reorg)
    /// Tip height parameter is needed when chain shrinks and there are
    /// no new header (happens when invalidateblock is called).
    pub(crate) fn update(
        &self,
        undoer: impl BlockUndoer,
        new_headers: Vec<NewHeader>,
        tip_height: Option<u64>,
    ) -> Result<()> {
        let mut headers = self.headers.write().unwrap();
        let mut heights = self.heights.write().unwrap();

        let rewind = match new_headers.first() {
            Some(first) => first.height,
            None => tip_height.expect("empty new_headers and no tip height") + 1,
        };

        for (hash, header) in headers.drain(rewind as usize..) {
            let height = heights.remove(&hash).expect("heights map missing entry");
            info!("Chain: Undoing block {} (height {})", hash, height);
            undoer.undo_block(&header)?;
        }

        if let Some(first_height) = new_headers.first().map(|h| h.height) {
            for (h, height) in new_headers.into_iter().zip(first_height..) {
                assert_eq!(h.height, height);
                assert_eq!(h.hash, h.header.block_hash());
                let existed = heights.insert(h.hash, h.height);
                assert!(existed.is_none());
                headers.push((h.hash, h.header));
            }
            info!(
                "chain updated: tip={}, height={}",
                headers.last().unwrap().0,
                headers.len() - 1
            );
        }
        Ok(())
    }

    /// Best block hash
    pub fn tip_hash(&self) -> BlockHash {
        self.headers.read().unwrap().last().expect("empty chain").0
    }

    /// Number of blocks (excluding genesis block)
    pub(crate) fn height(&self) -> u64 {
        (self.headers.read().unwrap().len() - 1) as u64
    }

    /// If we have block
    pub fn contains(&self, blockhash: &BlockHash) -> bool {
        self.heights.read().unwrap().contains_key(blockhash)
    }
}

pub fn download_all_headers(electrum: &Client) -> Result<Vec<NewHeader>> {
    let genesis = 0;
    let tip: Value =
        serde_json::from_str(&electrum.raw_call("blockchain.headers.tip", [])?.to_string())?;

    let tip = tip
        .get("height")
        .context("no height")?
        .as_i64()
        .context("no int")?;

    let mut headers: Vec<NewHeader> = Vec::new();

    for batch_start in (genesis..=tip).step_by(1000) {
        let batch_end = std::cmp::min(batch_start + 999, tip); // Ensure we don't exceed the tip
        let mut batch = Batch::default();

        info!("Fetching headers {} -> {}", batch_start, batch_end);

        for height in batch_start..=batch_end {
            batch.raw(
                "blockchain.block.header".to_owned(),
                vec![Param::U32(height as u32)],
            );
        }

        let headers_response = electrum.batch_call(&batch)?;

        let processed_headers: Result<Vec<NewHeader>> = headers_response
            .into_par_iter()
            .enumerate()
            .map(|(index, header)| {
                let header_hex = header.as_str().context("expected header str")?;
                let header: BlockHeader = deserialize(&hex::decode(header_hex)?)?;
                Ok(NewHeader::from_header_and_height(
                    header,
                    (batch_start as usize + index) as u64,
                ))
            })
            .collect();

        match processed_headers {
            Ok(mut processed) => {
                assert!(processed.first().unwrap().height == batch_start as u64);
                assert!(processed.last().unwrap().height == batch_end as u64);
                headers.append(&mut processed)
            }
            Err(e) => {
                bail!("Fetching headers failed {}", e);
            }
        }
    }

    Ok(headers)
}

pub fn download_block_header(electrum: &Client, blockhash: &BlockHash) -> Result<NewHeader> {
    let resp = electrum.raw_call(
        "blockchain.block.header_verbose",
        [Param::String(blockhash.to_hex())],
    );
    let header: Value = serde_json::from_str(&resp?.to_string())?;

    let height = header
        .get("height")
        .context("no header height")?
        .as_i64()
        .context("height not int")?;
    let hex = header
        .get("hex")
        .context("no header hex")?
        .as_str()
        .context("header hex not str")?;

    let header: BlockHeader = deserialize(&hex::decode(hex)?)?;

    Ok(NewHeader {
        header,
        height: height as u64,
        hash: header.block_hash(),
    })
}

// Fetches all headers not in our chain
pub fn get_new_headers(
    electrum: &Client,
    chain: &Chain,
    new_tip: &BlockHash,
) -> Result<Vec<NewHeader>> {
    // Iterate back over headers until known blockash is found:
    if chain.height() == 0 {
        info!("Fetching all headers");
        return download_all_headers(electrum);
    }

    if chain.contains(new_tip) {
        return Ok(vec![]);
    }

    info!(
        "Downloading new block headers ({} already indexed)",
        chain.height(),
    );

    let mut new_headers: Vec<NewHeader> = vec![];
    let null_hash = BlockHash::all_zeros();

    let mut blockhash = *new_tip;

    while blockhash != null_hash {
        if new_headers.len() % 1000 == 0 {
            info!(
                "Downloading headers progress: {} fetched... ",
                new_headers.len()
            );
        }
        if chain.contains(&blockhash) {
            break;
        }

        let header = download_block_header(electrum, &blockhash)
            .context(format!("failed to get {} header", blockhash))?;
        blockhash = header.header.prev_blockhash;
        new_headers.push(header);
    }
    trace!("downloaded {} block headers", new_headers.len());
    new_headers.reverse(); // so the tip is the last vector entry
    Ok(new_headers)
}

#[cfg(test)]
mod tests {
    use super::*;
    use bitcoincash::consensus::deserialize;
    use bitcoincash::hashes::hex::{FromHex, ToHex};

    #[test]
    fn test_genesis() {
        let regtest = Chain::new_regtest();
        assert_eq!(regtest.height(), 0);
        assert_eq!(
            regtest.tip_hash().to_hex(),
            "0f9188f13cb7b2c71f2a335e3a4fc328bf5beb436012afca590b1a11466e2206"
        );
    }

    #[test]
    fn test_updates() {
        let hex_headers = vec![
"0000002006226e46111a0b59caaf126043eb5bbf28c34f3a5e332a1fc7b2b73cf188910f1d14d3c7ff12d6adf494ebbcfba69baa915a066358b68a2b8c37126f74de396b1d61cc60ffff7f2000000000",
"00000020d700ae5d3c705702e0a5d9ababd22ded079f8a63b880b1866321d6bfcb028c3fc816efcf0e84ccafa1dda26be337f58d41b438170c357cda33a68af5550590bc1e61cc60ffff7f2004000000",
"00000020d13731bc59bc0989e06a5e7cab9843a4e17ad65c7ca47cd77f50dfd24f1f55793f7f342526aca9adb6ce8f33d8a07662c97d29d83b9e18117fb3eceecb2ab99b1e61cc60ffff7f2001000000",
"00000020a603def3e1255cadfb6df072946327c58b344f9bfb133e8e3e280d1c2d55b31c731a68f70219472864a7cb010cd53dc7e0f67e57f7d08b97e5e092b0c3942ad51f61cc60ffff7f2001000000",
"0000002041dd202b3b2edcdd3c8582117376347d48ff79ff97c95e5ac814820462012e785142dc360975b982ca43eecd14b4ba6f019041819d4fc5936255d7a2c45a96651f61cc60ffff7f2000000000",
"0000002072e297a2d6b633c44f3c9b1a340d06f3ce4e6bcd79ebd4c4ff1c249a77e1e37c59c7be1ca0964452e1735c0d2740f0d98a11445a6140c36b55770b5c0bcf801f1f61cc60ffff7f2000000000",
"000000200c9eb5889a8e924d1c4e8e79a716514579e41114ef37d72295df8869d6718e4ac5840f28de43ff25c7b9200aaf7873b20587c92827eaa61943484ca828bdd2e11f61cc60ffff7f2000000000",
"000000205873f322b333933e656b07881bb399dae61a6c0fa74188b5fb0e3dd71c9e2442f9e2f433f54466900407cf6a9f676913dd54aad977f7b05afcd6dcd81e98ee752061cc60ffff7f2004000000",
"00000020fd1120713506267f1dba2e1856ca1d4490077d261cde8d3e182677880df0d856bf94cfa5e189c85462813751ab4059643759ed319a81e0617113758f8adf67bc2061cc60ffff7f2000000000",
"000000200030d7f9c11ef35b89a0eefb9a5e449909339b5e7854d99804ea8d6a49bf900a0304d2e55fe0b6415949cff9bca0f88c0717884a5e5797509f89f856af93624a2061cc60ffff7f2002000000",
        ];
        let headers: Vec<(BlockHeader, u64)> = hex_headers
            .iter()
            .enumerate()
            .map(|(height, hex_header)| {
                (
                    deserialize(&Vec::from_hex(hex_header).unwrap()).unwrap(),
                    1 + height as u64,
                )
            })
            .collect();

        for chunk_size in 1..hex_headers.len() {
            let regtest = Chain::new_regtest();
            let mut tip = regtest.tip_hash();
            for chunk in headers.chunks(chunk_size) {
                let mut update = vec![];
                let mut last_height = u64::MAX;
                for (header, height) in chunk {
                    tip = header.block_hash();
                    last_height = *height;
                    update.push(NewHeader::from_header_and_height(*header, *height))
                }
                regtest
                    .update(DummyBlockUndoer::new(true), update, None)
                    .unwrap();
                assert_eq!(regtest.tip_hash(), tip);
                assert_eq!(regtest.height(), last_height);
            }
            assert_eq!(regtest.tip_hash(), headers.last().unwrap().0.block_hash());
            assert_eq!(regtest.height(), headers.len() as u64);
        }

        // test loading from a list
        let mut regtest = Chain::new_regtest();
        let chain_with_genesis = {
            let genesis = Chain::new_regtest().get_block_header(0).unwrap();
            let mut chain = headers.clone();
            chain.push((genesis, 0 as u64));
            chain
        };
        regtest.load(chain_with_genesis.clone()).unwrap();

        assert_eq!(regtest.height(), headers.len() as u64);

        // test getters
        for (header, height) in headers.iter() {
            assert_eq!(regtest.get_block_header(*height), Some(*header));
            assert_eq!(regtest.get_block_hash(*height), Some(header.block_hash()));
            assert_eq!(
                regtest.get_block_height(&header.block_hash()),
                Some(*height as u64)
            );
        }

        // test chain shortening
        for i in (0..=headers.len()).rev() {
            let hash = regtest.get_block_hash(i as u64).unwrap();
            assert_eq!(regtest.get_block_height(&hash), Some(i as u64));
            assert_eq!(regtest.height(), i as u64);
            assert_eq!(regtest.tip_hash(), hash);
            regtest
                .drop_last_headers(DummyBlockUndoer::new(true), 1)
                .unwrap();
        }
        assert_eq!(regtest.height(), 0);
        assert_eq!(
            regtest.tip_hash().to_hex(),
            "0f9188f13cb7b2c71f2a335e3a4fc328bf5beb436012afca590b1a11466e2206"
        );

        regtest
            .drop_last_headers(DummyBlockUndoer::new(true), 1)
            .unwrap();
        assert_eq!(regtest.height(), 0);
        assert_eq!(
            regtest.tip_hash().to_hex(),
            "0f9188f13cb7b2c71f2a335e3a4fc328bf5beb436012afca590b1a11466e2206"
        );

        // test reorg
        let mut regtest = Chain::new_regtest();
        regtest.load(chain_with_genesis.clone()).unwrap();
        let height = regtest.height();

        let new_header: BlockHeader = deserialize(&Vec::from_hex("000000200030d7f9c11ef35b89a0eefb9a5e449909339b5e7854d99804ea8d6a49bf900a0304d2e55fe0b6415949cff9bca0f88c0717884a5e5797509f89f856af93624a7a6bcc60ffff7f2000000000").unwrap()).unwrap();
        regtest
            .update(
                DummyBlockUndoer::new(true),
                vec![NewHeader::from_header_and_height(new_header, height)],
                None,
            )
            .unwrap();
        assert_eq!(regtest.height(), height);
        assert_eq!(
            regtest.tip_hash().to_hex(),
            "0e16637fe0700a7c52e9a6eaa58bd6ac7202652103be8f778680c66f51ad2e9b"
        );
    }
}