chain: record NoteSources in CompactBlocks

chain/src/lib.rs

@@ -12,5 +12,5 @@ pub mod sync;
 pub use epoch::Epoch;
 pub use known_assets::KnownAssets;
 pub use note_source::NoteSource;
-pub use sync::CompactBlock;
+pub use sync::{AnnotatedNotePayload, CompactBlock};
 pub use view::View;

chain/src/note_source.rs

@@ -54,6 +54,13 @@ impl TryFrom<[u8; 32]> for NoteSource {
     }
 }
 
+impl TryFrom<&[u8]> for NoteSource {
+    type Error = anyhow::Error;
+    fn try_from(value: &[u8]) -> Result<Self, Self::Error> {
+        <[u8; 32]>::try_from(value)?.try_into()
+    }
+}
+
 impl Protobuf<pb::NoteSource> for NoteSource {}
 
 impl TryFrom<pb::NoteSource> for NoteSource {

chain/src/quarantined.rs

@@ -1,9 +1,11 @@
 use serde::{Deserialize, Serialize};
 use std::collections::BTreeMap;
 
-use penumbra_crypto::{IdentityKey, NotePayload, Nullifier};
+use penumbra_crypto::{IdentityKey, Nullifier};
 use penumbra_proto::{chain as pb, Protobuf};
 
+use crate::sync::AnnotatedNotePayload;
+
 /// All the things scheduled for unquarantine, grouped by unbonding epoch.
 #[derive(Debug, Clone, Default, Serialize, Deserialize)]
 #[serde(into = "pb::Quarantined", try_from = "pb::Quarantined")]
@@ -17,7 +19,7 @@ impl Quarantined {
         &mut self,
         epoch: u64,
         identity_key: IdentityKey,
-        note_payload: NotePayload,
+        note_payload: AnnotatedNotePayload,
     ) {
         self.quarantined
             .entry(epoch)
@@ -143,7 +145,7 @@ pub struct Scheduled {
 
 impl Scheduled {
     /// Schedule a note for unquarantine, tied to a given validator.
-    pub fn schedule_note(&mut self, identity_key: IdentityKey, note_payload: NotePayload) {
+    pub fn schedule_note(&mut self, identity_key: IdentityKey, note_payload: AnnotatedNotePayload) {
         self.scheduled
             .entry(identity_key)
             .or_default()
@@ -272,13 +274,13 @@ impl From<Scheduled> for pb::quarantined::Scheduled {
     try_from = "pb::quarantined::Unbonding"
 )]
 pub struct Unbonding {
-    pub note_payloads: Vec<NotePayload>,
+    pub note_payloads: Vec<AnnotatedNotePayload>,
     pub nullifiers: Vec<Nullifier>,
 }
 
 impl Unbonding {
     /// Add a new note.
-    pub fn schedule_note(&mut self, note_payload: NotePayload) {
+    pub fn schedule_note(&mut self, note_payload: AnnotatedNotePayload) {
         self.note_payloads.push(note_payload)
     }
 
@@ -293,8 +295,8 @@ impl Unbonding {
     }
 }
 
-impl Extend<NotePayload> for Unbonding {
-    fn extend<T: IntoIterator<Item = NotePayload>>(&mut self, iter: T) {
+impl Extend<AnnotatedNotePayload> for Unbonding {
+    fn extend<T: IntoIterator<Item = AnnotatedNotePayload>>(&mut self, iter: T) {
         self.note_payloads.extend(iter);
     }
 }
@@ -305,8 +307,8 @@ impl Extend<Nullifier> for Unbonding {
     }
 }
 
-impl FromIterator<NotePayload> for Unbonding {
-    fn from_iter<T: IntoIterator<Item = NotePayload>>(iter: T) -> Self {
+impl FromIterator<AnnotatedNotePayload> for Unbonding {
+    fn from_iter<T: IntoIterator<Item = AnnotatedNotePayload>>(iter: T) -> Self {
         let mut unbonding = Unbonding::default();
         unbonding.extend(iter);
         unbonding
@@ -331,7 +333,7 @@ impl TryFrom<pb::quarantined::Unbonding> for Unbonding {
             note_payloads: value
                 .note_payloads
                 .into_iter()
-                .map(NotePayload::try_from)
+                .map(AnnotatedNotePayload::try_from)
                 .collect::<Result<Vec<_>, _>>()?,
             nullifiers: value
                 .nullifiers

chain/src/sync.rs

@@ -6,16 +6,27 @@ use penumbra_proto::{chain as pb, Protobuf};
 use penumbra_tct::builder::{block, epoch};
 use serde::{Deserialize, Serialize};
 
-use crate::quarantined::Quarantined;
+use crate::{quarantined::Quarantined, NoteSource};
+
+/// A note payload annotated with the source of the note.
+#[derive(Clone, Debug, Serialize, Deserialize)]
+#[serde(
+    try_from = "pb::AnnotatedNotePayload",
+    into = "pb::AnnotatedNotePayload"
+)]
+pub struct AnnotatedNotePayload {
+    pub payload: NotePayload,
+    pub source: NoteSource,
+}
 
 /// A compressed delta update with the minimal data from a block required to
 /// synchronize private client state.
 #[derive(Clone, Debug, Serialize, Deserialize)]
 #[serde(try_from = "pb::CompactBlock", into = "pb::CompactBlock")]
 pub struct CompactBlock {
     pub height: u64,
-    // Note payloads describing new notes.
-    pub note_payloads: Vec<NotePayload>,
+    // Annotated note payloads describing new notes.
+    pub note_payloads: Vec<AnnotatedNotePayload>,
     // Nullifiers identifying spent notes.
     pub nullifiers: Vec<Nullifier>,
     // The block root of this block.
@@ -55,6 +66,34 @@ impl CompactBlock {
     }
 }
 
+impl Protobuf<pb::AnnotatedNotePayload> for AnnotatedNotePayload {}
+
+impl From<AnnotatedNotePayload> for pb::AnnotatedNotePayload {
+    fn from(v: AnnotatedNotePayload) -> Self {
+        pb::AnnotatedNotePayload {
+            payload: Some(v.payload.into()),
+            source: Some(v.source.into()),
+        }
+    }
+}
+
+impl TryFrom<pb::AnnotatedNotePayload> for AnnotatedNotePayload {
+    type Error = anyhow::Error;
+
+    fn try_from(value: pb::AnnotatedNotePayload) -> Result<Self, Self::Error> {
+        Ok(AnnotatedNotePayload {
+            payload: value
+                .payload
+                .ok_or_else(|| anyhow::anyhow!("missing note payload"))?
+                .try_into()?,
+            source: value
+                .source
+                .ok_or_else(|| anyhow::anyhow!("missing note source"))?
+                .try_into()?,
+        })
+    }
+}
+
 impl Protobuf<pb::CompactBlock> for CompactBlock {}
 
 impl From<CompactBlock> for pb::CompactBlock {
@@ -89,8 +128,8 @@ impl TryFrom<pb::CompactBlock> for CompactBlock {
             note_payloads: value
                 .note_payloads
                 .into_iter()
-                .map(NotePayload::try_from)
-                .collect::<Result<Vec<NotePayload>>>()?,
+                .map(AnnotatedNotePayload::try_from)
+                .collect::<Result<Vec<AnnotatedNotePayload>>>()?,
             nullifiers: value
                 .nullifiers
                 .into_iter()

component/src/shielded_pool/component.rs

@@ -11,7 +11,7 @@ use decaf377::{Fq, Fr};
 use penumbra_chain::{
     genesis,
     quarantined::{self, Slashed},
-    sync::CompactBlock,
+    sync::{AnnotatedNotePayload, CompactBlock},
     Epoch, KnownAssets, NoteSource, View as _,
 };
 use penumbra_crypto::{
@@ -198,8 +198,9 @@ impl Component for ShieldedPool {
                 ctx.record(event::quarantine_spend(quarantined_spent_nullifier));
             }
         } else {
-            for compact_output in tx.note_payloads() {
-                self.add_note(compact_output, source).await;
+            for payload in tx.note_payloads() {
+                self.add_note(AnnotatedNotePayload { payload, source })
+                    .await;
             }
             for spent_nullifier in tx.spent_nullifiers() {
                 self.spend_nullifier(spent_nullifier, source).await;
@@ -333,56 +334,60 @@ impl ShieldedPool {
         self.state
             .update_token_supply(&value.asset_id, value.amount as i64)
             .await?;
-        self.add_note(
-            NotePayload {
+        self.add_note(AnnotatedNotePayload {
+            payload: NotePayload {
                 note_commitment,
                 ephemeral_key,
                 encrypted_note,
             },
             source,
-        )
+        })
         .await;
 
         Ok(())
     }
 
-    #[instrument(skip(self, source, note_payload), fields(note_commitment = ?note_payload.note_commitment))]
-    async fn add_note(&mut self, note_payload: NotePayload, source: NoteSource) {
+    #[instrument(skip(self, source, payload), fields(note_commitment = ?payload.note_commitment))]
+    async fn add_note(&mut self, AnnotatedNotePayload { payload, source }: AnnotatedNotePayload) {
         tracing::debug!("adding note");
 
         // 1. Insert it into the NCT
         self.note_commitment_tree
-            .insert(tct::Witness::Forget, note_payload.note_commitment)
+            .insert(tct::Witness::Forget, payload.note_commitment)
             .expect("inserting into the note commitment tree never fails");
 
         // 2. Record its source in the JMT
         self.state
-            .set_note_source(note_payload.note_commitment, source)
+            .set_note_source(payload.note_commitment, source)
             .await;
 
         // 3. Finally, record it in the pending compact block.
-        self.compact_block.note_payloads.push(note_payload);
+        self.compact_block
+            .note_payloads
+            .push(AnnotatedNotePayload { payload, source });
     }
 
-    #[instrument(skip(self, source, note_payload), fields(note_commitment = ?note_payload.note_commitment))]
+    #[instrument(skip(self, source, payload), fields(note_commitment = ?payload.note_commitment))]
     async fn schedule_note(
         &mut self,
         epoch: u64,
         identity_key: IdentityKey,
-        note_payload: NotePayload,
+        payload: NotePayload,
         source: NoteSource,
     ) {
         tracing::debug!("scheduling note");
 
         // 1. Record its source in the JMT
         self.state
-            .set_note_source(note_payload.note_commitment, source)
+            .set_note_source(payload.note_commitment, source)
             .await;
 
         // 2. Schedule it in the compact block
-        self.compact_block
-            .quarantined
-            .schedule_note(epoch, identity_key, note_payload);
+        self.compact_block.quarantined.schedule_note(
+            epoch,
+            identity_key,
+            AnnotatedNotePayload { payload, source },
+        );
     }
 
     #[instrument(skip(self, source))]
@@ -611,14 +616,8 @@ impl ShieldedPool {
                 // For all the note payloads scheduled for unquarantine now, remove them from
                 // quarantine and add them to the proper notes for this block
                 for note_payload in per_validator.note_payloads {
-                    let note_source = self
-                        .state
-                        .note_source(note_payload.note_commitment)
-                        .await
-                        .expect("can try to unquarantine note")
-                        .expect("note payload to unquarantine has source");
                     tracing::debug!(?note_payload, "unquarantining note");
-                    self.add_note(note_payload, note_source).await;
+                    self.add_note(note_payload).await;
                 }
                 // For all the nullifiers scheduled for unquarantine now, remove them from
                 // quarantine and add them to the proper nullifiers for this block
@@ -922,7 +921,8 @@ pub trait View: StateExt {
                     self.unquarantine_nullifier(nullifier).await?;
                 }
                 for note_payload in unbonding.note_payloads.iter() {
-                    self.roll_back_note(note_payload.note_commitment).await?;
+                    self.roll_back_note(note_payload.payload.note_commitment)
+                        .await?;
                 }
             }
             // We're removed all the scheduled notes and nullifiers for this epoch and identity key:

proto/build.rs

@@ -140,6 +140,7 @@ static TYPE_ATTRIBUTES: &[(&str, &str)] = &[
     (".penumbra.crypto.SpendAuthSignature", SERDE_TRANSPARENT),
     (".penumbra.chain.ChainParams", SERIALIZE),
     (".penumbra.chain.CompactBlock", SERIALIZE),
+    (".penumbra.chain.AnnotatedNotePayload", SERIALIZE),
     (".penumbra.chain.KnownAssets", SERIALIZE),
     (".penumbra.chain.KnownAssets", SERDE_TRANSPARENT),
     (".penumbra.chain.NoteSource", SERIALIZE),

proto/proto/chain.proto

@@ -49,7 +49,7 @@ message AssetInfo {
 message CompactBlock {
   uint64 height = 1;
   // NotePayloads describing new notes.
-  repeated crypto.NotePayload note_payloads = 2;
+  repeated AnnotatedNotePayload note_payloads = 2;
   // Nullifiers identifying spent notes.
   repeated crypto.Nullifier nullifiers = 3;
   // The block root of this block.
@@ -59,7 +59,13 @@ message CompactBlock {
   // Newly quarantined things in this block.
   Quarantined quarantined = 6;
   // Validators slashed in this block.
-  repeated crypto.IdentityKey slashed = 7;
+  repeated crypto.IdentityKey slashed = 16;
+}
+
+// A note payload, annotated with the note source.
+message AnnotatedNotePayload {
+  crypto.NotePayload payload = 1;
+  NoteSource source = 2;
 }
 
 message KnownAssets {
@@ -90,7 +96,7 @@ message GenesisAppState {
 
 message Quarantined {
   message Unbonding {
-    repeated crypto.NotePayload note_payloads = 1;
+    repeated AnnotatedNotePayload note_payloads = 1;
     repeated crypto.Nullifier nullifiers = 2;
   }
 
@@ -113,4 +119,4 @@ message Quarantined {
 
 message Slashed {
   repeated crypto.IdentityKey validators = 1;
-}
+}

proto/proto/view.proto

@@ -104,6 +104,8 @@ message NoteRecord {
     optional uint64 height_spent = 6;
     // The note position.
     uint64 position = 7;
+    // The source of the note (a tx hash or otherwise)
+    chain.NoteSource source = 8;
 }
 
 // A query for notes known by the view service.
@@ -151,6 +153,8 @@ message QuarantinedNoteRecord {
     uint64 unbonding_epoch = 5;
     // The validator identity the quarantining is bound to.
     crypto.IdentityKey identity_key = 6;
+    // The source of the note (a tx hash or otherwise)
+    chain.NoteSource source = 7;
 }
 
 // A query for quarantined notes known by the view service.