fix: sign Attestation instead of AttestationData

Author: MegaRedHand

crates/blockchain/src/key_manager.rs

@@ -1,8 +1,8 @@
 use std::collections::HashMap;
 
 use ethlambda_types::{
-    attestation::XmssSignature,
-    primitives::H256,
+    attestation::{Attestation, XmssSignature},
+    primitives::{H256, TreeHash},
     signature::{ValidatorSecretKey, ValidatorSignature},
 };
 
@@ -51,7 +51,7 @@ impl KeyManager {
         self.keys.keys().copied().collect()
     }
 
-    /// Signs an attestation message for the specified validator.
+    /// Signs a message for the specified validator.
     ///
     /// # Arguments
     ///
@@ -68,13 +68,13 @@ impl KeyManager {
     /// # Example
     ///
     /// ```ignore
-    /// let signature = key_manager.sign_attestation(
+    /// let signature = key_manager.sign_message(
     ///     validator_id,
     ///     epoch,
     ///     &message_hash
     /// )?;
     /// ```
-    pub fn sign_attestation(
+    fn sign_message(
         &mut self,
         validator_id: u64,
         epoch: u32,
@@ -94,8 +94,22 @@ impl KeyManager {
         let xmss_sig = XmssSignature::try_from(sig_bytes)
             .map_err(|e| KeyManagerError::SignatureConversionError(e.to_string()))?;
 
+        tracing::info!(
+            %validator_id, %epoch, %message, ?xmss_sig,
+            "Publishing attestation",
+        );
+
         Ok(xmss_sig)
     }
+
+    pub fn sign_attestation(
+        &mut self,
+        attestation: &Attestation,
+    ) -> Result<XmssSignature, KeyManagerError> {
+        let message_hash = attestation.tree_hash_root();
+        let epoch = attestation.data.slot as u32;
+        self.sign_message(attestation.validator_id, epoch, &message_hash)
+    }
 }
 
 #[cfg(test)]
@@ -110,12 +124,12 @@ mod tests {
     }
 
     #[test]
-    fn test_sign_attestation_validator_not_found() {
+    fn test_sign_message_validator_not_found() {
         let keys = HashMap::new();
         let mut key_manager = KeyManager::new(keys);
         let message = H256::default();
 
-        let result = key_manager.sign_attestation(123, 0, &message);
+        let result = key_manager.sign_message(123, 0, &message);
         assert!(matches!(
             result,
             Err(KeyManagerError::ValidatorKeyNotFound(123))

crates/blockchain/src/lib.rs

@@ -144,12 +144,6 @@ impl BlockChainServer {
         // Produce attestation data once for all validators
         let attestation_data = self.store.produce_attestation_data(slot);
 
-        // Hash the attestation data for signing
-        let message_hash = attestation_data.tree_hash_root();
-
-        // Epoch for signing
-        let epoch = slot as u32;
-
         // For each registered validator, produce and publish attestation
         for validator_id in self.key_manager.validator_ids() {
             // Skip if this validator is the slot proposer
@@ -158,14 +152,16 @@ impl BlockChainServer {
                 continue;
             }
 
+            // Hash the attestation data for signing
+            let attestation = Attestation {
+                data: attestation_data.clone(),
+                validator_id,
+            };
+
             // Sign the attestation
-            let Ok(signature) = self
-                .key_manager
-                .sign_attestation(validator_id, epoch, &message_hash)
-                .inspect_err(
-                    |err| error!(%slot, %validator_id, %err, "Failed to sign attestation"),
-                )
-            else {
+            let Ok(signature) = self.key_manager.sign_attestation(&attestation).inspect_err(
+                |err| error!(%slot, %validator_id, %err, "Failed to sign attestation"),
+            ) else {
                 continue;
             };
 
@@ -218,11 +214,9 @@ impl BlockChainServer {
         };
 
         // Sign the proposer's attestation
-        let message_hash = proposer_attestation.data.tree_hash_root();
-        let epoch = slot as u32;
         let Ok(proposer_signature) = self
             .key_manager
-            .sign_attestation(validator_id, epoch, &message_hash)
+            .sign_attestation(&proposer_attestation)
             .inspect_err(
                 |err| error!(%slot, %validator_id, %err, "Failed to sign proposer attestation"),
             )
@@ -233,9 +227,8 @@ impl BlockChainServer {
         // Assemble flat signature list: [attestation_sig_0, ..., attestation_sig_n, proposer_sig]
         let mut signatures = attestation_signatures;
         signatures.push(proposer_signature);
-        let block_signatures: BlockSignatures = signatures
-            .try_into()
-            .expect("signatures within limit");
+        let block_signatures: BlockSignatures =
+            signatures.try_into().expect("signatures within limit");
 
         // Assemble SignedBlockWithAttestation
         let signed_block = SignedBlockWithAttestation {