From 293628455fe45456a5685d9329ff0604f45c115d Mon Sep 17 00:00:00 2001
From: realbigsean <seananderson33@GMAIL.com>
Date: Tue, 18 Jun 2024 11:10:32 -0400
Subject: [PATCH] Indexed att on disk (#35)

* indexed att on disk

* fix lints

* Update slasher/src/migrate.rs

Co-authored-by: ethDreamer <37123614+ethDreamer@users.noreply.github.com>

---------

Co-authored-by: Lion - dapplion <35266934+dapplion@users.noreply.github.com>
Co-authored-by: ethDreamer <37123614+ethDreamer@users.noreply.github.com>
---
 consensus/types/src/indexed_attestation.rs | 32 ++++++++++++++++++++++
 consensus/types/src/lib.rs                 |  3 +-
 slasher/src/database.rs                    | 17 ++++++------
 slasher/src/migrate.rs                     |  4 +++
 4 files changed, 46 insertions(+), 10 deletions(-)

diff --git a/consensus/types/src/indexed_attestation.rs b/consensus/types/src/indexed_attestation.rs
index 00cac09abf3..429d65898a9 100644
--- a/consensus/types/src/indexed_attestation.rs
+++ b/consensus/types/src/indexed_attestation.rs
@@ -239,6 +239,38 @@ mod quoted_variable_list_u64 {
     }
 }
 
+#[derive(Debug, Clone, Encode, Decode, PartialEq)]
+#[ssz(enum_behaviour = "union")]
+pub enum IndexedAttestationOnDisk<E: EthSpec> {
+    Base(IndexedAttestationBase<E>),
+    Electra(IndexedAttestationElectra<E>),
+}
+
+#[derive(Debug, Clone, Encode, PartialEq)]
+#[ssz(enum_behaviour = "union")]
+pub enum IndexedAttestationRefOnDisk<'a, E: EthSpec> {
+    Base(&'a IndexedAttestationBase<E>),
+    Electra(&'a IndexedAttestationElectra<E>),
+}
+
+impl<'a, E: EthSpec> From<&'a IndexedAttestation<E>> for IndexedAttestationRefOnDisk<'a, E> {
+    fn from(attestation: &'a IndexedAttestation<E>) -> Self {
+        match attestation {
+            IndexedAttestation::Base(attestation) => Self::Base(attestation),
+            IndexedAttestation::Electra(attestation) => Self::Electra(attestation),
+        }
+    }
+}
+
+impl<E: EthSpec> From<IndexedAttestationOnDisk<E>> for IndexedAttestation<E> {
+    fn from(attestation: IndexedAttestationOnDisk<E>) -> Self {
+        match attestation {
+            IndexedAttestationOnDisk::Base(attestation) => Self::Base(attestation),
+            IndexedAttestationOnDisk::Electra(attestation) => Self::Electra(attestation),
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
diff --git a/consensus/types/src/lib.rs b/consensus/types/src/lib.rs
index 716c8314788..213e09699d5 100644
--- a/consensus/types/src/lib.rs
+++ b/consensus/types/src/lib.rs
@@ -178,7 +178,8 @@ pub use crate::fork_versioned_response::{ForkVersionDeserialize, ForkVersionedRe
 pub use crate::graffiti::{Graffiti, GRAFFITI_BYTES_LEN};
 pub use crate::historical_batch::HistoricalBatch;
 pub use crate::indexed_attestation::{
-    IndexedAttestation, IndexedAttestationBase, IndexedAttestationElectra, IndexedAttestationRef,
+    IndexedAttestation, IndexedAttestationBase, IndexedAttestationElectra,
+    IndexedAttestationOnDisk, IndexedAttestationRef, IndexedAttestationRefOnDisk,
 };
 pub use crate::light_client_bootstrap::{
     LightClientBootstrap, LightClientBootstrapAltair, LightClientBootstrapCapella,
diff --git a/slasher/src/database.rs b/slasher/src/database.rs
index 16a67cd684f..16089706a00 100644
--- a/slasher/src/database.rs
+++ b/slasher/src/database.rs
@@ -18,12 +18,12 @@ use std::marker::PhantomData;
 use std::sync::Arc;
 use tree_hash::TreeHash;
 use types::{
-    Epoch, EthSpec, Hash256, IndexedAttestation, IndexedAttestationBase, ProposerSlashing,
-    SignedBeaconBlockHeader, Slot,
+    Epoch, EthSpec, Hash256, IndexedAttestation, IndexedAttestationOnDisk,
+    IndexedAttestationRefOnDisk, ProposerSlashing, SignedBeaconBlockHeader, Slot,
 };
 
 /// Current database schema version, to check compatibility of on-disk DB with software.
-pub const CURRENT_SCHEMA_VERSION: u64 = 3;
+pub const CURRENT_SCHEMA_VERSION: u64 = 4;
 
 /// Metadata about the slashing database itself.
 const METADATA_DB: &str = "metadata";
@@ -458,7 +458,9 @@ impl<E: EthSpec> SlasherDB<E> {
         };
 
         let attestation_key = IndexedAttestationId::new(indexed_att_id);
-        let data = indexed_attestation.as_ssz_bytes();
+        let indexed_attestation_on_disk: IndexedAttestationRefOnDisk<E> =
+            indexed_attestation.into();
+        let data = indexed_attestation_on_disk.as_ssz_bytes();
 
         cursor.put(attestation_key.as_ref(), &data)?;
         drop(cursor);
@@ -482,11 +484,8 @@ impl<E: EthSpec> SlasherDB<E> {
             .ok_or(Error::MissingIndexedAttestation {
                 id: indexed_attestation_id.as_u64(),
             })?;
-        // TODO(electra): make slasher fork-aware and return correct variant of attestation. Both
-        // Base and Electra variant can the same SSZ encoding, however the smaller list maximum of
-        // Base can error with an Electra attestation with heavy participation.
-        let indexed_attestation: IndexedAttestationBase<E> = ssz_decode(bytes)?;
-        Ok(IndexedAttestation::Base(indexed_attestation))
+        let indexed_attestation: IndexedAttestationOnDisk<E> = ssz_decode(bytes)?;
+        Ok(indexed_attestation.into())
     }
 
     fn get_attestation_data_root(
diff --git a/slasher/src/migrate.rs b/slasher/src/migrate.rs
index 674ab9c132d..ce298caaf2c 100644
--- a/slasher/src/migrate.rs
+++ b/slasher/src/migrate.rs
@@ -17,6 +17,10 @@ impl<E: EthSpec> SlasherDB<E> {
                     software_schema_version: CURRENT_SCHEMA_VERSION,
                 }),
                 (x, y) if x == y => Ok(self),
+                (3, 4) => {
+                    // TODO(electra): db migration due to `IndexedAttestationOnDisk`
+                    Ok(self)
+                }
                 (_, _) => Err(Error::IncompatibleSchemaVersion {
                     database_schema_version: schema_version,
                     software_schema_version: CURRENT_SCHEMA_VERSION,