feat: note the canonical bitcoin blocks in the database

signer/migrations/0021__add_is_canonical_to_bitcoin_blocks.sql

@@ -0,0 +1,9 @@
+-- Add is_canonical column to bitcoin_blocks table to track which blocks
+-- are on the canonical Bitcoin blockchain.
+
+ALTER TABLE sbtc_signer.bitcoin_blocks
+ADD COLUMN is_canonical BOOLEAN;
+
+-- Create an index on is_canonical for efficient queries filtering
+-- canonical vs non-canonical blocks.
+CREATE INDEX ix_bitcoin_blocks_is_canonical ON sbtc_signer.bitcoin_blocks(is_canonical);

signer/src/block_observer.rs

@@ -170,7 +170,7 @@ where
                     )
                     .increment(1);
 
-                    if let Err(error) = self.process_bitcoin_blocks_until(block_hash).await {
+                    if let Err(error) = self.process_bitcoin_chain_tip(block_hash).await {
                         tracing::warn!(%error, %block_hash, "could not process bitcoin blocks");
                     }
 
@@ -294,7 +294,8 @@ impl<C: Context, B> BlockObserver<C, B> {
     }
 
     /// Find the parent blocks from the given block that are also missing
-    /// from our database.
+    /// from our database. The results are returned from blocks with the
+    /// least height to the greatest height.
     ///
     /// # Notes
     ///
@@ -347,6 +348,27 @@ impl<C: Context, B> BlockObserver<C, B> {
         Ok(headers.into())
     }
 
+    /// Process the bitcoin chain tip by fetching all unknown block headers
+    /// from the given chain tip back until the nakamoto start height. Then
+    /// mark all blocks reachable from the chain tip as canonical in the
+    /// database.
+    ///
+    /// # Notes
+    ///
+    /// This function must only be called with the bitcoin chain tip since
+    /// it updates all blocks that are reachable from the given block hash
+    /// as canonical and may update blocks not reachable as non-canonical.
+    #[tracing::instrument(skip_all, fields(%chain_tip))]
+    async fn process_bitcoin_chain_tip(&self, chain_tip: BlockHash) -> Result<(), Error> {
+        self.process_bitcoin_blocks_until(chain_tip).await?;
+
+        let db = self.context.get_storage_mut();
+
+        tracing::info!("updating canonical bitcoin blockchain to chain tip");
+        let chain_tip: model::BitcoinBlockHash = chain_tip.into();
+        db.set_canonical_bitcoin_blockchain(&chain_tip).await
+    }
+
     /// Process bitcoin blocks until we get caught up to the given
     /// `block_hash`.
     ///

signer/src/storage/memory/store.rs

@@ -42,6 +42,9 @@ pub struct Store {
     /// Bitcoin blocks
     pub bitcoin_blocks: HashMap<model::BitcoinBlockHash, model::BitcoinBlock>,
 
+    /// Bitcoin blocks that are on the canonical bitcoin blockchain.
+    pub canonical_bitcoin_blocks: HashMap<model::BitcoinBlockHash, model::BitcoinBlock>,
+
     /// Stacks blocks
     pub stacks_blocks: HashMap<model::StacksBlockHash, model::StacksBlock>,
 

signer/src/storage/memory/write.rs

@@ -385,6 +385,41 @@ impl DbWrite for SharedStore {
 
         Ok(())
     }
+
+    async fn set_canonical_bitcoin_blockchain(
+        &self,
+        chain_tip: &model::BitcoinBlockHash,
+    ) -> Result<(), Error> {
+        let mut store = self.lock().await;
+        store.version += 1;
+
+        store.canonical_bitcoin_blocks.clear();
+        // Then, recursively mark all blocks reachable from the chain tip as canonical
+        if let Some(block) = store.bitcoin_blocks.get(chain_tip).cloned() {
+            store.canonical_bitcoin_blocks.insert(*chain_tip, block);
+        }
+
+        let mut current_block_hash = *chain_tip;
+        while let Some(block) = store
+            .canonical_bitcoin_blocks
+            .get(&current_block_hash)
+            .cloned()
+        {
+            match store.bitcoin_blocks.get(&block.parent_hash).cloned() {
+                Some(parent) => {
+                    store
+                        .canonical_bitcoin_blocks
+                        .insert(block.parent_hash, parent);
+                    current_block_hash = block.parent_hash;
+                }
+                None => {
+                    break;
+                }
+            }
+        }
+
+        Ok(())
+    }
 }
 
 impl DbWrite for InMemoryTransaction {
@@ -539,4 +574,11 @@ impl DbWrite for InMemoryTransaction {
             .update_peer_connection(pub_key, peer_id, address)
             .await
     }
+
+    async fn set_canonical_bitcoin_blockchain(
+        &self,
+        chain_tip: &model::BitcoinBlockHash,
+    ) -> Result<(), Error> {
+        self.store.set_canonical_bitcoin_blockchain(chain_tip).await
+    }
 }

signer/src/storage/mod.rs

@@ -624,4 +624,20 @@ pub trait DbWrite {
         peer_id: &PeerId,
         address: Multiaddr,
     ) -> impl Future<Output = Result<(), Error>> + Send;
+
+    /// Ensure that each and only blocks along the chain identified by the
+    /// given chain tip have their is_canonical set to TRUE.
+    ///
+    /// # Notes
+    ///
+    /// This function should mark all blocks that are reachable from the
+    /// given chain tip as canonical (is_canonical = TRUE), and no other
+    /// block should have is_canonical set to TRUE. That means that the
+    /// DbWrite::set_canonical_bitcoin_blockchain() function need not set
+    /// the is_canonical column to FALSE if it is a non-canonical block --
+    /// it can just leave the value of that column as NULL.
+    fn set_canonical_bitcoin_blockchain(
+        &self,
+        chain_tip: &model::BitcoinBlockHash,
+    ) -> impl Future<Output = Result<(), Error>> + Send;
 }

signer/src/storage/postgres/write.rs

@@ -28,8 +28,9 @@ impl PgWrite {
               ( block_hash
               , block_height
               , parent_hash
+              , is_canonical
               )
-            VALUES ($1, $2, $3)
+            VALUES ($1, $2, $3, NULL)
             ON CONFLICT DO NOTHING",
         )
         .bind(block.block_hash)
@@ -1001,6 +1002,73 @@ impl PgWrite {
 
         Ok(())
     }
+
+    /// Update the is_canonical status for all blocks with height greater
+    /// than the current "canonical root height" (the first block on the chain
+    /// reachable from the chain tip that is already marked as canonical).
+    ///
+    /// # Notes
+    ///
+    /// If a forked block is added to the database and the next chain tip
+    /// builds off of a block where is_canonical is TRUE, and that block
+    /// has height greater than this forked block, then this function will
+    /// not set the is_canonical column for the forked block. This is okay,
+    /// since we just need to make sure that it is not set to TRUE.
+    async fn set_canonical_bitcoin_blockchain<'e, E>(
+        executor: &'e mut E,
+        chain_tip: &model::BitcoinBlockHash,
+    ) -> Result<(), Error>
+    where
+        E: 'static,
+        for<'c> &'c mut E: sqlx::PgExecutor<'c>,
+    {
+        // Walk backwards from chain tip to find the canonical root and
+        // build the canonical chain in a single recursive query, then
+        // update all blocks with height >= the minimum height found.
+        sqlx::query(
+            r#"
+            -- set_canonical_bitcoin_blockchain
+            WITH RECURSIVE canonical_chain AS (
+                -- Start from the chain tip
+                SELECT
+                    block_hash
+                  , block_height
+                  , parent_hash
+                  , is_canonical
+                FROM sbtc_signer.bitcoin_blocks
+                WHERE block_hash = $1
+
+                UNION ALL
+
+                -- Recursively get parent blocks, stopping when we find a canonical block
+                SELECT
+                    parent.block_hash
+                  , parent.block_height
+                  , parent.parent_hash
+                  , parent.is_canonical
+                FROM sbtc_signer.bitcoin_blocks AS parent
+                JOIN canonical_chain AS child
+                  ON parent.block_hash = child.parent_hash
+                WHERE child.is_canonical IS NOT TRUE
+            ),
+            min_height AS (
+                SELECT COALESCE(MIN(block_height), 0) AS height
+                FROM canonical_chain
+            )
+            UPDATE sbtc_signer.bitcoin_blocks
+            SET is_canonical = (block_hash IN (
+                SELECT block_hash FROM canonical_chain
+            ))
+            WHERE block_height >= (SELECT height FROM min_height)
+            "#,
+        )
+        .bind(chain_tip)
+        .execute(executor)
+        .await
+        .map_err(Error::SqlxQuery)?;
+
+        Ok(())
+    }
 }
 
 impl DbWrite for PgStore {
@@ -1161,6 +1229,14 @@ impl DbWrite for PgStore {
         )
         .await
     }
+
+    async fn set_canonical_bitcoin_blockchain(
+        &self,
+        chain_tip: &model::BitcoinBlockHash,
+    ) -> Result<(), Error> {
+        PgWrite::set_canonical_bitcoin_blockchain(self.get_connection().await?.as_mut(), chain_tip)
+            .await
+    }
 }
 
 impl DbWrite for PgTransaction<'_> {
@@ -1332,4 +1408,12 @@ impl DbWrite for PgTransaction<'_> {
         let mut tx = self.tx.lock().await;
         PgWrite::update_peer_connection(tx.as_mut(), pub_key, peer_id, address).await
     }
+
+    async fn set_canonical_bitcoin_blockchain(
+        &self,
+        chain_tip: &model::BitcoinBlockHash,
+    ) -> Result<(), Error> {
+        let mut tx = self.tx.lock().await;
+        PgWrite::set_canonical_bitcoin_blockchain(tx.as_mut(), chain_tip).await
+    }
 }

signer/src/testing/blocks.rs

@@ -1,6 +1,9 @@
 //! Various utilities for generating and manipulating chains of bitcoin and
 //! stacks blocks for testing purposes.
 
+use std::collections::BTreeMap;
+use std::collections::btree_map;
+
 use fake::Fake as _;
 use fake::Faker;
 
@@ -13,14 +16,15 @@ use crate::storage::model::StacksBlockHeight;
 
 /// Represents a naive, sequential chain of bitcoin blocks and provides basic
 /// functionality for manipulation. Does not handle forks/branches.
-pub struct BitcoinChain(Vec<BitcoinBlock>);
+#[derive(Debug, Clone)]
+pub struct BitcoinChain(BTreeMap<BitcoinBlockHeight, BitcoinBlock>);
 
 impl<'a> IntoIterator for &'a BitcoinChain {
     type Item = &'a BitcoinBlock;
-    type IntoIter = std::slice::Iter<'a, BitcoinBlock>;
+    type IntoIter = std::collections::btree_map::Values<'a, BitcoinBlockHeight, BitcoinBlock>;
 
     fn into_iter(self) -> Self::IntoIter {
-        self.0.iter()
+        self.0.values()
     }
 }
 
@@ -40,7 +44,7 @@ impl BitcoinChain {
     /// block will have a height one greater than the previous block and a
     /// parent hash equal to the hash of the previous block.
     pub fn new() -> Self {
-        Self(vec![BitcoinBlock::new_genesis()])
+        Self(BTreeMap::from([(0u64.into(), BitcoinBlock::new_genesis())]))
     }
 
     /// Generate a new chain of bitcoin blocks with a length equal to `length`.
@@ -63,32 +67,64 @@ impl BitcoinChain {
     pub fn generate_blocks(&mut self, length: usize) -> Vec<&BitcoinBlock> {
         for _ in 0..length {
             let new_block = self.chain_tip().new_child();
-            self.0.push(new_block);
+            self.0.insert(new_block.block_height, new_block);
         }
 
-        self.0[(self.0.len() - length)..].iter().collect()
+        let start: BitcoinBlockHeight = ((self.0.len() - length) as u64).into();
+        self.0.range(start..).map(|(_, block)| block).collect()
     }
 
     /// Gets the first block in the chain.
     pub fn first_block(&self) -> &BitcoinBlock {
-        self.0.first().unwrap()
+        self.0.first_key_value().unwrap().1
     }
 
     /// Gets the last block in the chain.
     pub fn chain_tip(&self) -> &BitcoinBlock {
-        self.0.last().unwrap()
+        self.0.last_key_value().unwrap().1
     }
 
     /// Gets the nth block in the chain, if it exists.
     pub fn nth_block_checked(&self, height: BitcoinBlockHeight) -> Option<&BitcoinBlock> {
-        self.0.get(*height as usize)
+        self.0.get(&height)
     }
 
     /// Gets the nth block in the chain, panicking if it does not exist.
     pub fn nth_block(&self, height: BitcoinBlockHeight) -> &BitcoinBlock {
         self.nth_block_checked(height)
             .expect("no nth bitcoin block (index out of range)")
     }
+
+    /// Create a new fork of the chain at the given height.
+    ///
+    /// The fork will contain the blocks up to, but not including, the
+    /// given height. It will then generate `num_blocks` new blocks, and
+    /// add them to the chain starting at the given height, returning the
+    /// resulting chain.
+    pub fn fork_at_height<H>(&self, height: H, num_blocks: usize) -> Self
+    where
+        H: Into<BitcoinBlockHeight>,
+    {
+        let fork: BTreeMap<BitcoinBlockHeight, BitcoinBlock> = self
+            .0
+            .range(..height.into())
+            .map(|(height, block)| (*height, block.clone()))
+            .collect();
+
+        let mut fork_chain = Self(fork);
+
+        fork_chain.generate_blocks(num_blocks);
+
+        fork_chain
+    }
+
+    /// Get a range of blocks from the chain.
+    pub fn range<R>(&self, range: R) -> btree_map::Range<'_, BitcoinBlockHeight, BitcoinBlock>
+    where
+        R: std::ops::RangeBounds<BitcoinBlockHeight>,
+    {
+        self.0.range(range)
+    }
 }
 
 impl BitcoinBlock {

signer/src/testing/storage/postgres.rs

@@ -36,4 +36,19 @@ impl PgStore {
         .await
         .map_err(Error::SqlxQuery)
     }
+
+    /// Returns the value of the is_canonical column for the given block
+    /// hash. If the given block hash is missing then an error is returned.
+    pub async fn is_block_canonical(
+        &self,
+        block_hash: &model::BitcoinBlockHash,
+    ) -> Result<Option<bool>, Error> {
+        sqlx::query_scalar(
+            "SELECT is_canonical FROM sbtc_signer.bitcoin_blocks WHERE block_hash = $1",
+        )
+        .bind(block_hash)
+        .fetch_one(self.pool())
+        .await
+        .map_err(Error::SqlxQuery)
+    }
 }