Merge ExternalKeyForSessionToSignBatch into InfoForBatch

processor/scanner/src/report/db.rs

@@ -10,12 +10,14 @@ use serai_db::{Get, DbTxn, create_db};
 use serai_primitives::Balance;
 use serai_validator_sets_primitives::Session;
 
+use primitives::EncodableG;
 use crate::{ScannerFeed, KeyFor, AddressFor};
 
 #[derive(BorshSerialize, BorshDeserialize)]
-pub(crate) struct BatchInfo {
+pub(crate) struct BatchInfo<K: BorshSerialize> {
   pub(crate) block_number: u64,
-  pub(crate) publisher: Session,
+  pub(crate) session_to_sign_batch: Session,
+  pub(crate) external_key_for_session_to_sign_batch: K,
   pub(crate) in_instructions_hash: [u8; 32],
 }
 
@@ -27,11 +29,7 @@ create_db!(
     NextBatchId: () -> u32,
 
     // The information needed to verify a batch
-    InfoForBatch: (batch: u32) -> BatchInfo,
-
-    // The external key for the session which should sign a batch
-    // TODO: Merge this with InfoForBatch
-    ExternalKeyForSessionToSignBatch: (batch: u32) -> Vec<u8>,
+    InfoForBatch: <G: GroupEncoding>(batch: u32) -> BatchInfo<EncodableG<G>>,
 
     // The return addresses for the InInstructions within a Batch
     SerializedReturnAddresses: (batch: u32) -> Vec<u8>,
@@ -65,37 +63,27 @@ impl<S: ScannerFeed> ReportDb<S> {
     txn: &mut impl DbTxn,
     id: u32,
     block_number: u64,
-    publisher: Session,
+    session_to_sign_batch: Session,
+    external_key_for_session_to_sign_batch: KeyFor<S>,
     in_instructions_hash: [u8; 32],
   ) {
-    InfoForBatch::set(txn, id, &BatchInfo { block_number, publisher, in_instructions_hash });
-  }
-
-  pub(crate) fn take_info_for_batch(txn: &mut impl DbTxn, id: u32) -> Option<BatchInfo> {
-    InfoForBatch::take(txn, id)
-  }
-
-  pub(crate) fn save_external_key_for_session_to_sign_batch(
-    txn: &mut impl DbTxn,
-    id: u32,
-    external_key_for_session_to_sign_batch: &KeyFor<S>,
-  ) {
-    ExternalKeyForSessionToSignBatch::set(
+    InfoForBatch::set(
       txn,
       id,
-      &external_key_for_session_to_sign_batch.to_bytes().as_ref().to_vec(),
+      &BatchInfo {
+        block_number,
+        session_to_sign_batch,
+        external_key_for_session_to_sign_batch: EncodableG(external_key_for_session_to_sign_batch),
+        in_instructions_hash,
+      },
     );
   }
 
-  pub(crate) fn take_external_key_for_session_to_sign_batch(
+  pub(crate) fn take_info_for_batch(
     txn: &mut impl DbTxn,
     id: u32,
-  ) -> Option<KeyFor<S>> {
-    ExternalKeyForSessionToSignBatch::get(txn, id).map(|key_vec| {
-      let mut key = <KeyFor<S> as GroupEncoding>::Repr::default();
-      key.as_mut().copy_from_slice(&key_vec);
-      KeyFor::<S>::from_bytes(&key).unwrap()
-    })
+  ) -> Option<BatchInfo<EncodableG<KeyFor<S>>>> {
+    InfoForBatch::take(txn, id)
   }
 
   pub(crate) fn save_return_information(

processor/scanner/src/report/mod.rs

@@ -7,7 +7,7 @@ use serai_db::{DbTxn, Db};
 
 use serai_in_instructions_primitives::{MAX_BATCH_SIZE, Batch};
 
-use primitives::task::ContinuallyRan;
+use primitives::{EncodableG, task::ContinuallyRan};
 use crate::{
   db::{Returnable, ScannerGlobalDb, InInstructionData, ScanToReportDb, Batches, BatchesToSign},
   scan::next_to_scan_for_outputs_block,
@@ -21,17 +21,10 @@ use db::ReportDb;
 pub(crate) fn take_info_for_batch<S: ScannerFeed>(
   txn: &mut impl DbTxn,
   id: u32,
-) -> Option<BatchInfo> {
+) -> Option<BatchInfo<EncodableG<KeyFor<S>>>> {
   ReportDb::<S>::take_info_for_batch(txn, id)
 }
 
-pub(crate) fn take_external_key_for_session_to_sign_batch<S: ScannerFeed>(
-  txn: &mut impl DbTxn,
-  id: u32,
-) -> Option<KeyFor<S>> {
-  ReportDb::<S>::take_external_key_for_session_to_sign_batch(txn, id)
-}
-
 pub(crate) fn take_return_information<S: ScannerFeed>(
   txn: &mut impl DbTxn,
   id: u32,
@@ -151,13 +144,9 @@ impl<D: Db, S: ScannerFeed> ContinuallyRan for ReportTask<D, S> {
               batch.id,
               block_number,
               session_to_sign_batch,
+              external_key_for_session_to_sign_batch,
               Blake2b::<U32>::digest(batch.instructions.encode()).into(),
             );
-            ReportDb::<S>::save_external_key_for_session_to_sign_batch(
-              &mut txn,
-              batch.id,
-              &external_key_for_session_to_sign_batch,
-            );
             ReportDb::<S>::save_return_information(&mut txn, batch.id, return_information);
           }
 

processor/scanner/src/substrate/mod.rs

@@ -81,7 +81,8 @@ impl<D: Db, S: ScannerFeed> ContinuallyRan for SubstrateTask<D, S> {
             // Check if we have the information for this batch
             let Some(report::BatchInfo {
               block_number,
-              publisher: expected_publisher,
+              session_to_sign_batch,
+              external_key_for_session_to_sign_batch,
               in_instructions_hash: expected_in_instructions_hash,
             }) = report::take_info_for_batch::<S>(&mut txn, batch_id)
             else {
@@ -90,7 +91,7 @@ impl<D: Db, S: ScannerFeed> ContinuallyRan for SubstrateTask<D, S> {
               return Ok(made_progress);
             };
             assert_eq!(
-              publisher, expected_publisher,
+              publisher, session_to_sign_batch,
               "batch acknowledged on-chain was acknowledged by an unexpected publisher"
             );
             assert_eq!(
@@ -98,16 +99,11 @@ impl<D: Db, S: ScannerFeed> ContinuallyRan for SubstrateTask<D, S> {
               "batch acknowledged on-chain was distinct"
             );
 
-            {
-              let external_key_for_session_to_sign_batch =
-                report::take_external_key_for_session_to_sign_batch::<S>(&mut txn, batch_id)
-                  .unwrap();
-              AcknowledgedBatches::send(
-                &mut txn,
-                &external_key_for_session_to_sign_batch,
-                batch_id,
-              );
-            }
+            AcknowledgedBatches::send(
+              &mut txn,
+              &external_key_for_session_to_sign_batch.0,
+              batch_id,
+            );
 
             // Mark we made progress and handle this
             made_progress = true;