feat(ICRC-Rosetta): add ICRC-107 fee collector block handling

rs/ledger_suite/icrc1/test_utils/src/lib.rs

@@ -146,11 +146,21 @@ fn operation_strategy<Tokens: TokensType>(
                 fee,
             });
 
+        let fee_collector_strategy = (
+            prop::option::of(principal_strategy()),
+            prop::option::of(account_strategy()),
+        )
+            .prop_map(move |(caller, fee_collector)| Operation::FeeCollector {
+                fee_collector,
+                caller,
+            });
+
         prop_oneof![
             mint_strategy,
             burn_strategy,
             transfer_strategy,
             approve_strategy,
+            fee_collector_strategy,
         ]
     })
 }

rs/rosetta-api/icrc1/src/common/storage/schema.rs

@@ -91,6 +91,19 @@ pub fn create_tables(connection: &Connection) -> Result<()> {
         [],
     )?;
 
+    // The rosetta_metadata table is meant to store values like the index
+    // of the highest processed block. This is different than the metadata
+    // table above which stores the ICRC1 token metadata.
+    connection.execute(
+        r#"
+        CREATE TABLE IF NOT EXISTS rosetta_metadata (
+            key TEXT PRIMARY KEY,
+            value BLOB NOT NULL
+        );
+        "#,
+        [],
+    )?;
+
     create_indexes(connection)
 }
 

rs/rosetta-api/icrc1/src/common/storage/storage_client.rs

@@ -337,9 +337,9 @@ impl StorageClient {
 
     /// Retrieves the highest block index in the account balance table.
     /// Returns None if the account balance table is empty.
-    pub fn get_highest_block_idx_in_account_balance_table(&self) -> Result<Option<u64>> {
+    pub fn get_highest_processed_block_idx(&self) -> Result<Option<u64>> {
         let open_connection = self.storage_connection.lock().unwrap();
-        storage_operations::get_highest_block_idx_in_account_balance_table(&open_connection)
+        storage_operations::get_highest_processed_block_idx(&open_connection)
     }
 
     // Retrieves the account balance at a certain block height
@@ -405,7 +405,8 @@ mod tests {
     use ic_icrc1::blocks::encoded_block_to_generic_block;
     use ic_icrc1::blocks::generic_block_to_encoded_block;
     use ic_icrc1_test_utils::{
-        arb_amount, blocks_strategy, metadata_strategy, valid_blockchain_with_gaps_strategy,
+        arb_amount, blocks_strategy, metadata_strategy, valid_blockchain_strategy,
+        valid_blockchain_with_gaps_strategy,
     };
     use ic_icrc1_tokens_u64::U64;
     use ic_icrc1_tokens_u256::U256;
@@ -432,7 +433,7 @@ mod tests {
 
     proptest! {
           #[test]
-          fn test_read_and_write_blocks_u64(blockchain in prop::collection::vec(blocks_strategy::<U64>(arb_amount()),0..5)){
+          fn test_read_and_write_blocks_u64(blockchain in valid_blockchain_strategy::<U64>(5)){
            let storage_client_memory = StorageClient::new_in_memory().unwrap();
            let mut rosetta_blocks = vec![];
            for (index,block) in blockchain.into_iter().enumerate(){
@@ -459,7 +460,7 @@ mod tests {
        }
 
        #[test]
-       fn test_read_and_write_blocks_u256(blockchain in prop::collection::vec(blocks_strategy::<U256>(arb_amount()),0..5)){
+       fn test_read_and_write_blocks_u256(blockchain in valid_blockchain_strategy::<U256>(5)){
         let storage_client_memory = StorageClient::new_in_memory().unwrap();
         let mut rosetta_blocks = vec![];
         for (index,block) in blockchain.into_iter().enumerate(){
@@ -508,10 +509,11 @@ mod tests {
 
               // Duplicate the last transaction generated
               let duplicate_tx_block = RosettaBlock::from_generic_block(last_block.get_generic_block(), last_block.index + 1).unwrap();
+              let count_before = storage_client_memory.get_transactions_by_hash(duplicate_tx_block.clone().get_transaction_hash()).unwrap().len();
               storage_client_memory.store_blocks([duplicate_tx_block.clone()].to_vec()).unwrap();
 
-              // The hash of the duplicated transaction should still be the same --> There should be two transactions with the same transaction hash.
-              assert_eq!(storage_client_memory.get_transactions_by_hash(duplicate_tx_block.clone().get_transaction_hash()).unwrap().len(),2);
+              // The hash of the duplicated transaction should still be the same --> There should be one more transaction with the same transaction hash.
+              assert_eq!(storage_client_memory.get_transactions_by_hash(duplicate_tx_block.clone().get_transaction_hash()).unwrap().len(), count_before + 1);
               }
            }
 

rs/rosetta-api/icrc1/src/common/storage/storage_operations.rs

@@ -14,6 +14,9 @@ use std::collections::{BTreeMap, HashMap};
 use std::str::FromStr;
 use tracing::{info, trace};
 
+pub const METADATA_FEE_COL: &str = "fee_collector_107";
+pub const METADATA_BLOCK_IDX: &str = "highest_processed_block_index";
+
 /// Gets the current value of a counter from the database.
 /// Returns None if the counter doesn't exist.
 pub fn get_counter_value(
@@ -106,6 +109,20 @@ pub fn initialize_counter_if_missing(
     Ok(())
 }
 
+pub fn get_107_fee_collector_or_legacy(
+    rosetta_block: &RosettaBlock,
+    connection: &Connection,
+    fee_collector_107: Option<Option<Account>>,
+) -> anyhow::Result<Option<Account>> {
+    // First check if we have a 107 fee collector
+    if let Some(fee_collector_107) = fee_collector_107 {
+        return Ok(fee_collector_107);
+    }
+
+    // There is no 107 fee collector, check legacy fee collector in the block
+    get_fee_collector_from_block(rosetta_block, connection)
+}
+
 // Helper function to resolve the fee collector account from a block
 pub fn get_fee_collector_from_block(
     rosetta_block: &RosettaBlock,
@@ -170,6 +187,27 @@ pub fn get_metadata(connection: &Connection) -> anyhow::Result<Vec<MetadataEntry
     Ok(result)
 }
 
+pub fn get_rosetta_metadata(connection: &Connection, key: &str) -> anyhow::Result<Option<Vec<u8>>> {
+    let mut stmt_metadata = connection.prepare_cached(&format!(
+        "SELECT value FROM rosetta_metadata WHERE key = '{key}'"
+    ))?;
+    let rows = stmt_metadata.query_map(params![], |row| row.get(0))?;
+    let mut result = vec![];
+    for row in rows {
+        let entry: Vec<u8> = row?;
+        result.push(entry);
+    }
+    if result.len() == 1 {
+        Ok(Some(result[0].clone()))
+    } else if result.is_empty() {
+        // Return None if no metadata entry found
+        Ok(None)
+    } else {
+        // If more than one metadata entry was found return an error
+        bail!(format!("Multiple metadata entries found for key: {key}"))
+    }
+}
+
 pub fn update_account_balances(
     connection: &mut Connection,
     flush_cache_and_shrink_memory: bool,
@@ -250,7 +288,7 @@ pub fn update_account_balances(
 
     // The next block to be updated is the highest block index in the account balance table + 1 if the table is not empty and 0 otherwise
     let next_block_to_be_updated =
-        get_highest_block_idx_in_account_balance_table(connection)?.map_or(0, |idx| idx + 1);
+        get_highest_processed_block_idx(connection)?.map_or(0, |idx| idx + 1);
     let highest_block_idx =
         get_block_with_highest_block_idx(connection)?.map_or(0, |block| block.index);
 
@@ -266,9 +304,26 @@ pub fn update_account_balances(
     // This also makes the inserting of the account balances batchable and therefore faster
     let mut account_balances_cache: HashMap<Account, BTreeMap<u64, Nat>> = HashMap::new();
 
+    let mut current_fee_collector_107 = match get_rosetta_metadata(connection, METADATA_FEE_COL)? {
+        Some(value) => {
+            let fc: Option<Account> = candid::decode_one(&value)?;
+            Some(fc)
+        }
+        None => None,
+    };
+    let collector_before = current_fee_collector_107;
+
     // As long as there are blocks to be fetched, keep on iterating over the blocks in the database with the given BATCH_SIZE interval
     while !rosetta_blocks.is_empty() {
+        let mut last_block_index = batch_start_idx;
         for rosetta_block in rosetta_blocks {
+            if rosetta_block.index < last_block_index {
+                bail!(format!(
+                    "Processing blocks not in order, previous processed block: {last_block_index}, current block {}",
+                    rosetta_block.index
+                ));
+            }
+            last_block_index = rosetta_block.index;
             match rosetta_block.get_transaction().operation {
                 crate::common::storage::types::IcrcOperation::Burn {
                     from,
@@ -290,9 +345,11 @@ pub fn update_account_balances(
                         connection,
                         &mut account_balances_cache,
                     )?;
-                    if let Some(collector) =
-                        get_fee_collector_from_block(&rosetta_block, connection)?
-                    {
+                    if let Some(collector) = get_107_fee_collector_or_legacy(
+                        &rosetta_block,
+                        connection,
+                        current_fee_collector_107,
+                    )? {
                         credit(
                             collector,
                             fee,
@@ -317,9 +374,11 @@ pub fn update_account_balances(
                         connection,
                         &mut account_balances_cache,
                     )?;
-                    if let Some(collector) =
-                        get_fee_collector_from_block(&rosetta_block, connection)?
-                    {
+                    if let Some(collector) = get_107_fee_collector_or_legacy(
+                        &rosetta_block,
+                        connection,
+                        current_fee_collector_107,
+                    )? {
                         credit(
                             collector,
                             fee,
@@ -342,11 +401,21 @@ pub fn update_account_balances(
                         .unwrap_or(Nat(BigUint::zero()));
                     debit(
                         from,
-                        fee,
+                        fee.clone(),
                         rosetta_block.index,
                         connection,
                         &mut account_balances_cache,
                     )?;
+
+                    if let Some(Some(collector)) = current_fee_collector_107 {
+                        credit(
+                            collector,
+                            fee,
+                            rosetta_block.index,
+                            connection,
+                            &mut account_balances_cache,
+                        )?;
+                    }
                 }
                 crate::common::storage::types::IcrcOperation::Transfer {
                     from,
@@ -378,9 +447,11 @@ pub fn update_account_balances(
                         &mut account_balances_cache,
                     )?;
 
-                    if let Some(collector) =
-                        get_fee_collector_from_block(&rosetta_block, connection)?
-                    {
+                    if let Some(collector) = get_107_fee_collector_or_legacy(
+                        &rosetta_block,
+                        connection,
+                        current_fee_collector_107,
+                    )? {
                         credit(
                             collector,
                             fee,
@@ -390,6 +461,12 @@ pub fn update_account_balances(
                         )?;
                     }
                 }
+                crate::common::storage::types::IcrcOperation::FeeCollector {
+                    fee_collector,
+                    caller: _,
+                } => {
+                    current_fee_collector_107 = Some(fee_collector);
+                }
             }
         }
 
@@ -407,6 +484,13 @@ pub fn update_account_balances(
                     })?;
             }
         }
+        if collector_before != current_fee_collector_107
+            && let Some(collector) = current_fee_collector_107
+        {
+            insert_tx.prepare_cached("INSERT INTO rosetta_metadata (key, value) VALUES (?1, ?2) ON CONFLICT (key) DO UPDATE SET value = excluded.value;")?.execute(params![METADATA_FEE_COL, candid::encode_one(collector)?])?;
+        }
+        let last_block_index_bytes = last_block_index.to_le_bytes();
+        insert_tx.prepare_cached("INSERT INTO rosetta_metadata (key, value) VALUES (?1, ?2) ON CONFLICT (key) DO UPDATE SET value = excluded.value;")?.execute(params![METADATA_BLOCK_IDX, last_block_index_bytes])?;
         insert_tx.commit()?;
 
         if flush_cache_and_shrink_memory {
@@ -416,7 +500,7 @@ pub fn update_account_balances(
         }
 
         // Fetch the next batch of blocks
-        batch_start_idx = get_highest_block_idx_in_account_balance_table(connection)?
+        batch_start_idx = get_highest_processed_block_idx(connection)?
             .context("No blocks in account balance table after inserting")?
             + 1;
         batch_end_idx = batch_start_idx + batch_size;
@@ -514,6 +598,22 @@ pub fn store_blocks(
                 fee,
                 expires_at,
             ),
+            crate::common::storage::types::IcrcOperation::FeeCollector {
+                fee_collector: _,
+                caller: _,
+            } => (
+                "107feecol",
+                None,
+                None,
+                None,
+                None,
+                None,
+                None,
+                Nat::from(0u64),
+                None,
+                None,
+                None,
+            ),
         };
 
         // SQLite doesn't support unsigned 64-bit integers. We need to convert the timestamps to signed
@@ -666,16 +766,10 @@ pub fn get_blocks_by_transaction_hash(
     read_blocks(&mut stmt, params![hash.as_slice().to_vec()])
 }
 
-pub fn get_highest_block_idx_in_account_balance_table(
-    connection: &Connection,
-) -> anyhow::Result<Option<u64>> {
-    match connection
-        .prepare_cached("SELECT block_idx FROM account_balances WHERE block_idx = (SELECT MAX(block_idx) FROM account_balances)")?
-        .query_map(params![], |row| row.get(0))?
-        .next()
-    {
+pub fn get_highest_processed_block_idx(connection: &Connection) -> anyhow::Result<Option<u64>> {
+    match get_rosetta_metadata(connection, METADATA_BLOCK_IDX)? {
+        Some(value) => Ok(Some(u64::from_le_bytes(value.as_slice().try_into()?))),
         None => Ok(None),
-        Some(res) => Ok(res?),
     }
 }
 
@@ -865,6 +959,10 @@ pub fn repair_fee_collector_balances(
 
     info!("Starting balance reconciliation...");
     connection.execute("DELETE FROM account_balances", params![])?;
+    connection.execute(
+        &format!("DELETE FROM rosetta_metadata WHERE key = '{METADATA_BLOCK_IDX}' OR key = '{METADATA_FEE_COL}'"),
+        params![],
+    )?;
 
     if block_count > 0 {
         info!("Reprocessing all blocks...");

rs/rosetta-api/icrc1/src/common/storage/storage_operations_test.rs

@@ -54,6 +54,7 @@ fn create_test_rosetta_block(
         effective_fee: None,
         fee_collector: None,
         fee_collector_block_index: None,
+        btype: None,
     };
 
     RosettaBlock {
@@ -108,6 +109,7 @@ fn create_test_approve_block(
         effective_fee: None,
         fee_collector: None,
         fee_collector_block_index: None,
+        btype: None,
     };
 
     RosettaBlock {
@@ -439,6 +441,20 @@ fn test_fee_collector_resolution_and_repair() -> anyhow::Result<()> {
     // Manually create broken balances (missing fee collector credits for block 2)
     connection.execute("DELETE FROM account_balances", params![])?;
 
+    // Insert metadata that needs to be cleared
+    connection.execute(
+        "INSERT INTO rosetta_metadata (key, value) VALUES (?1, ?2)",
+        params![METADATA_BLOCK_IDX, 100_000_000u64.to_le_bytes()],
+    )?;
+    let no_fee_col: Option<Account> = None;
+    connection.execute(
+        "INSERT INTO rosetta_metadata (key, value) VALUES (?1, ?2)",
+        params![
+            METADATA_FEE_COL,
+            candid::encode_one(no_fee_col).expect("failed to encode fee collector")
+        ],
+    )?;
+
     // Correct balances for mint and block 1
     connection.execute("INSERT INTO account_balances (block_idx, principal, subaccount, amount) VALUES (0, ?1, ?2, '1000000000')",
         params![from_account.owner.as_slice(), from_account.effective_subaccount().as_slice()])?;