Fix panic in "compact-note-decryption/invalid" check for ZSA and Vani…

…lla branches

- Added `extract_asset` method to `OrchardDomainCommon` trait to get the asset base from the note plaintext.
  - For `OrchardVanilla`, the method extracts native assets.
  - For `OrchardZSA`, the method extracts ZSA assets.
- Transformed `parse_note_version` to `validate_note_version` function that checks if the version matches the expected version for the domain and returns a boolean.
- Updated `parse_note_plaintext_without_memo` to validate the note version using `validate_note_version` before processing the plaintext.

benches/note_decryption.rs

@@ -28,6 +28,8 @@ fn bench_note_decryption<FL: OrchardFlavorBench>(c: &mut Criterion) {
     let recipient = valid_ivk.address_at(0u32);
     let valid_ivk = PreparedIncomingViewingKey::new(&valid_ivk);
 
+    // FIXME: should we add to the following comment that now for ZSA flavor the early
+    // rejection also happens when the asset is invalid?
     // Compact actions don't have the full AEAD ciphertext, so ZIP 307 trial-decryption
     // relies on an invalid ivk resulting in random noise for which the note commitment
     // is invalid. However, in practice we still get early rejection:

src/note_encryption/domain.rs

@@ -17,7 +17,7 @@ use crate::{
         DiversifiedTransmissionKey, Diversifier, EphemeralPublicKey, EphemeralSecretKey,
         OutgoingViewingKey, PreparedEphemeralPublicKey, PreparedIncomingViewingKey, SharedSecret,
     },
-    note::{AssetBase, ExtractedNoteCommitment, Note, RandomSeed, Rho},
+    note::{ExtractedNoteCommitment, Note, RandomSeed, Rho},
     value::{NoteValue, ValueCommitment},
 };
 
@@ -45,12 +45,6 @@ const ZSA_ASSET_SIZE: usize = 32;
 /// The size of a ZSA compact note.
 pub(super) const COMPACT_NOTE_SIZE_ZSA: usize = COMPACT_NOTE_SIZE_VANILLA + ZSA_ASSET_SIZE;
 
-/// The version byte for Vanilla.
-pub(super) const NOTE_VERSION_BYTE_V2: u8 = 0x02;
-
-/// The version byte for ZSA.
-pub(super) const NOTE_VERSION_BYTE_V3: u8 = 0x03;
-
 pub(super) type Memo = [u8; MEMO_SIZE];
 
 /// Defined in [Zcash Protocol Spec § 5.4.2: Pseudo Random Functions][concreteprfs].
@@ -78,27 +72,32 @@ pub(super) fn prf_ock_orchard(
     )
 }
 
-// FIXME: consider returning enum instead of u8
-/// Retrieves the version of the note plaintext.
-/// Returns `Some(u8)` if the version is recognized, otherwise `None`.
-pub(super) fn parse_note_version(plaintext: &[u8]) -> Option<u8> {
-    plaintext.first().and_then(|version| match *version {
-        NOTE_VERSION_BYTE_V2 | NOTE_VERSION_BYTE_V3 => Some(*version),
-        _ => None,
-    })
+/// Checks if the version of the note plaintext matches the expected version for the domain.
+pub(super) fn validate_note_version<D: OrchardDomainCommon>(
+    plaintext: &D::CompactNotePlaintextBytes,
+) -> bool {
+    plaintext
+        .as_ref()
+        .first()
+        .map_or(false, |version| *version == D::NOTE_VERSION_BYTE)
 }
 
 /// Parses the note plaintext (excluding the memo) and extracts the note and address if valid.
 /// Domain-specific requirements:
 /// - If the note version is 3, the `plaintext` must contain a valid encoding of a ZSA asset type.
-pub(super) fn parse_note_plaintext_without_memo<Bytes: AsRef<[u8]>, F>(
+pub(super) fn parse_note_plaintext_without_memo<D: OrchardDomainCommon, F>(
     rho: Rho,
-    plaintext: &Bytes,
+    plaintext: &D::CompactNotePlaintextBytes,
     get_validated_pk_d: F,
 ) -> Option<(Note, Address)>
 where
     F: FnOnce(&Diversifier) -> Option<DiversifiedTransmissionKey>,
 {
+    if !validate_note_version::<D>(plaintext) {
+        return None;
+    }
+
+    // FIXME: Is the following comment correct and clear?
     // The unwraps below are guaranteed to succeed by the assertion above
     let diversifier = Diversifier::from_bytes(
         plaintext.as_ref()[NOTE_DIVERSIFIER_OFFSET..NOTE_VALUE_OFFSET]
@@ -121,19 +120,9 @@ where
 
     let pk_d = get_validated_pk_d(&diversifier)?;
     let recipient = Address::from_parts(diversifier, pk_d);
-
-    let asset = match parse_note_version(plaintext.as_ref())? {
-        NOTE_VERSION_BYTE_V2 => AssetBase::native(),
-        NOTE_VERSION_BYTE_V3 => {
-            let bytes = plaintext.as_ref()[COMPACT_NOTE_SIZE_VANILLA..COMPACT_NOTE_SIZE_ZSA]
-                .try_into()
-                .unwrap();
-            AssetBase::from_bytes(bytes).unwrap()
-        }
-        _ => panic!("invalid note version"),
-    };
-
+    let asset = D::extract_asset(plaintext)?;
     let note = Option::from(Note::from_parts(recipient, value, asset, rho, rseed))?;
+
     Some((note, recipient))
 }
 
@@ -251,7 +240,7 @@ impl<D: OrchardDomainCommon> Domain for OrchardDomain<D> {
         ivk: &Self::IncomingViewingKey,
         plaintext: &D::CompactNotePlaintextBytes,
     ) -> Option<(Self::Note, Self::Recipient)> {
-        parse_note_plaintext_without_memo(self.rho, plaintext, |diversifier| {
+        parse_note_plaintext_without_memo::<D, _>(self.rho, plaintext, |diversifier| {
             Some(DiversifiedTransmissionKey::derive(ivk, diversifier))
         })
     }
@@ -261,7 +250,7 @@ impl<D: OrchardDomainCommon> Domain for OrchardDomain<D> {
         pk_d: &Self::DiversifiedTransmissionKey,
         plaintext: &D::CompactNotePlaintextBytes,
     ) -> Option<(Self::Note, Self::Recipient)> {
-        parse_note_plaintext_without_memo(self.rho, plaintext, |_| Some(*pk_d))
+        parse_note_plaintext_without_memo::<D, _>(self.rho, plaintext, |_| Some(*pk_d))
     }
 
     fn extract_memo(

src/note_encryption/orchard_domain.rs

@@ -5,7 +5,11 @@ use core::fmt;
 
 use zcash_note_encryption_zsa::{AEAD_TAG_SIZE, MEMO_SIZE};
 
-use crate::{action::Action, note::Rho, Note};
+use crate::{
+    action::Action,
+    note::{AssetBase, Rho},
+    Note,
+};
 
 use super::{compact_action::CompactAction, domain::Memo};
 
@@ -62,6 +66,9 @@ pub trait OrchardDomainCommon: fmt::Debug + Clone {
     /// The size of an encrypted note ciphertext, accounting for additional AEAD tag space.
     const ENC_CIPHERTEXT_SIZE: usize = Self::NOTE_PLAINTEXT_SIZE + AEAD_TAG_SIZE;
 
+    /// The note version byte.
+    const NOTE_VERSION_BYTE: u8;
+
     /// The raw bytes of a note plaintext.
     type NotePlaintextBytes: NoteBytes;
     /// The raw bytes of an encrypted note plaintext.
@@ -73,6 +80,9 @@ pub trait OrchardDomainCommon: fmt::Debug + Clone {
 
     /// Builds NotePlaintextBytes from Note and Memo.
     fn build_note_plaintext_bytes(note: &Note, memo: &Memo) -> Self::NotePlaintextBytes;
+
+    /// Extracts the asset from the note plaintext.
+    fn extract_asset(plaintext: &Self::CompactNotePlaintextBytes) -> Option<AssetBase>;
 }
 
 /// Orchard-specific note encryption logic.

src/note_encryption/orchard_domain_vanilla.rs

@@ -1,29 +1,36 @@
 //! This module implements the note encryption logic specific for the `OrchardVanilla` flavor.
 
-use crate::{note::Note, orchard_flavor::OrchardVanilla};
+use crate::{
+    note::{AssetBase, Note},
+    orchard_flavor::OrchardVanilla,
+};
 
 use super::{
-    domain::{
-        build_base_note_plaintext_bytes, Memo, COMPACT_NOTE_SIZE_VANILLA, NOTE_VERSION_BYTE_V2,
-    },
+    domain::{build_base_note_plaintext_bytes, Memo, COMPACT_NOTE_SIZE_VANILLA},
     orchard_domain::{NoteBytesData, OrchardDomainCommon},
 };
 
 impl OrchardDomainCommon for OrchardVanilla {
     const COMPACT_NOTE_SIZE: usize = COMPACT_NOTE_SIZE_VANILLA;
 
+    const NOTE_VERSION_BYTE: u8 = 0x02;
+
     type NotePlaintextBytes = NoteBytesData<{ Self::NOTE_PLAINTEXT_SIZE }>;
     type NoteCiphertextBytes = NoteBytesData<{ Self::ENC_CIPHERTEXT_SIZE }>;
     type CompactNotePlaintextBytes = NoteBytesData<{ Self::COMPACT_NOTE_SIZE }>;
     type CompactNoteCiphertextBytes = NoteBytesData<{ Self::COMPACT_NOTE_SIZE }>;
 
     fn build_note_plaintext_bytes(note: &Note, memo: &Memo) -> Self::NotePlaintextBytes {
-        let mut np = build_base_note_plaintext_bytes(NOTE_VERSION_BYTE_V2, note);
+        let mut np = build_base_note_plaintext_bytes(Self::NOTE_VERSION_BYTE, note);
 
         np[COMPACT_NOTE_SIZE_VANILLA..].copy_from_slice(memo);
 
         NoteBytesData(np)
     }
+
+    fn extract_asset(_plaintext: &Self::CompactNotePlaintextBytes) -> Option<AssetBase> {
+        Some(AssetBase::native())
+    }
 }
 
 #[cfg(test)]
@@ -55,10 +62,7 @@ mod tests {
 
     use super::super::{
         compact_action::CompactAction,
-        domain::{
-            parse_note_plaintext_without_memo, parse_note_version, prf_ock_orchard,
-            NOTE_VERSION_BYTE_V2,
-        },
+        domain::{parse_note_plaintext_without_memo, prf_ock_orchard, validate_note_version},
         orchard_domain::{NoteBytesData, OrchardDomain},
     };
 
@@ -81,10 +85,11 @@ mod tests {
 
             // Decode.
             let domain = OrchardDomainVanilla::for_rho(rho);
-            let version = parse_note_version(plaintext.as_ref()).unwrap();
             let (compact, parsed_memo) = domain.extract_memo(&plaintext);
 
-            let (parsed_note, parsed_recipient) = parse_note_plaintext_without_memo(rho, &compact,
+            assert!(validate_note_version::<OrchardVanilla>(&compact));
+
+            let (parsed_note, parsed_recipient) = parse_note_plaintext_without_memo::<OrchardVanilla, _>(rho, &compact,
                 |diversifier| {
                     assert_eq!(diversifier, &note.recipient().diversifier());
                     Some(*note.recipient().pk_d())
@@ -95,7 +100,6 @@ mod tests {
             assert_eq!(parsed_note, note);
             assert_eq!(parsed_recipient, note.recipient());
             assert_eq!(&parsed_memo, memo);
-            assert_eq!(version, NOTE_VERSION_BYTE_V2);
         }
     }
 

src/note_encryption/orchard_domain_zsa.rs

@@ -1,11 +1,13 @@
 //! This module implements the note encryption logic specific for the `OrchardZSA` flavor.
 
-use crate::{note::Note, orchard_flavor::OrchardZSA};
+use crate::{
+    note::{AssetBase, Note},
+    orchard_flavor::OrchardZSA,
+};
 
 use super::{
     domain::{
         build_base_note_plaintext_bytes, Memo, COMPACT_NOTE_SIZE_VANILLA, COMPACT_NOTE_SIZE_ZSA,
-        NOTE_VERSION_BYTE_V3,
     },
     orchard_domain::{NoteBytesData, OrchardDomainCommon},
 };
@@ -18,15 +20,25 @@ impl OrchardDomainCommon for OrchardZSA {
     type CompactNotePlaintextBytes = NoteBytesData<{ Self::COMPACT_NOTE_SIZE }>;
     type CompactNoteCiphertextBytes = NoteBytesData<{ Self::COMPACT_NOTE_SIZE }>;
 
+    const NOTE_VERSION_BYTE: u8 = 0x03;
+
     fn build_note_plaintext_bytes(note: &Note, memo: &Memo) -> Self::NotePlaintextBytes {
-        let mut np = build_base_note_plaintext_bytes(NOTE_VERSION_BYTE_V3, note);
+        let mut np = build_base_note_plaintext_bytes(Self::NOTE_VERSION_BYTE, note);
 
         np[COMPACT_NOTE_SIZE_VANILLA..COMPACT_NOTE_SIZE_ZSA]
             .copy_from_slice(&note.asset().to_bytes());
         np[COMPACT_NOTE_SIZE_ZSA..].copy_from_slice(memo);
 
         NoteBytesData(np)
     }
+
+    fn extract_asset(plaintext: &Self::CompactNotePlaintextBytes) -> Option<AssetBase> {
+        let bytes = plaintext.as_ref()[COMPACT_NOTE_SIZE_VANILLA..COMPACT_NOTE_SIZE_ZSA]
+            .try_into()
+            .unwrap();
+
+        AssetBase::from_bytes(bytes).into()
+    }
 }
 
 #[cfg(test)]
@@ -58,10 +70,7 @@ mod tests {
 
     use super::super::{
         compact_action::CompactAction,
-        domain::{
-            parse_note_plaintext_without_memo, parse_note_version, prf_ock_orchard,
-            NOTE_VERSION_BYTE_V3,
-        },
+        domain::{parse_note_plaintext_without_memo, prf_ock_orchard, validate_note_version},
         orchard_domain::{NoteBytesData, OrchardDomain},
     };
 
@@ -84,10 +93,11 @@ mod tests {
 
             // Decode.
             let domain = OrchardDomainZSA::for_rho(rho);
-            let version = parse_note_version(plaintext.as_ref()).unwrap();
             let (compact, parsed_memo) = domain.extract_memo(&plaintext);
 
-            let (parsed_note, parsed_recipient) = parse_note_plaintext_without_memo(rho, &compact,
+            assert!(validate_note_version::<OrchardZSA>(&compact));
+
+            let (parsed_note, parsed_recipient) = parse_note_plaintext_without_memo::<OrchardZSA, _>(rho, &compact,
                 |diversifier| {
                     assert_eq!(diversifier, &note.recipient().diversifier());
                     Some(*note.recipient().pk_d())
@@ -98,7 +108,6 @@ mod tests {
             assert_eq!(parsed_note, note);
             assert_eq!(parsed_recipient, note.recipient());
             assert_eq!(&parsed_memo, memo);
-            assert_eq!(version, NOTE_VERSION_BYTE_V3);
         }
     }