Rework PRF^expand to use explicit domains

`prf_expand{_vec}` have been replaced by the `PrfExpand` enum, which
has `PrfExpand::{expand, with_ad, with_ad_slices}` methods for use
within each domain as necessary.

src/keys.rs

@@ -15,8 +15,8 @@ use crate::{
     address::Address,
     primitives::redpallas::{self, SpendAuth},
     spec::{
-        commit_ivk, diversify_hash, extract_p, ka_orchard, prf_expand, prf_expand_vec, prf_nf,
-        to_base, to_scalar, NonIdentityPallasPoint, NonZeroPallasBase, NonZeroPallasScalar,
+        commit_ivk, diversify_hash, extract_p, ka_orchard, prf_nf, to_base, to_scalar,
+        NonIdentityPallasPoint, NonZeroPallasBase, NonZeroPallasScalar, PrfExpand,
     },
 };
 
@@ -73,7 +73,7 @@ pub struct SpendAuthorizingKey(redpallas::SigningKey<SpendAuth>);
 impl SpendAuthorizingKey {
     /// Derives ask from sk. Internal use only, does not enforce all constraints.
     fn derive_inner(sk: &SpendingKey) -> pallas::Scalar {
-        to_scalar(prf_expand(&sk.0, &[0x06]))
+        to_scalar(PrfExpand::OrchardAsk.expand(&sk.0))
     }
 
     /// Creates a spend authorization signature over the given message.
@@ -143,7 +143,7 @@ pub(crate) struct NullifierDerivingKey(pallas::Base);
 
 impl From<&SpendingKey> for NullifierDerivingKey {
     fn from(sk: &SpendingKey) -> Self {
-        NullifierDerivingKey(to_base(prf_expand(&sk.0, &[0x07])))
+        NullifierDerivingKey(to_base(PrfExpand::OrchardNk.expand(&sk.0)))
     }
 }
 
@@ -163,7 +163,7 @@ struct CommitIvkRandomness(pallas::Scalar);
 
 impl From<&SpendingKey> for CommitIvkRandomness {
     fn from(sk: &SpendingKey) -> Self {
-        CommitIvkRandomness(to_scalar(prf_expand(&sk.0, &[0x08])))
+        CommitIvkRandomness(to_scalar(PrfExpand::OrchardRivk.expand(&sk.0)))
     }
 }
 
@@ -209,7 +209,7 @@ impl FullViewingKey {
     fn derive_dk_ovk(&self) -> (DiversifierKey, OutgoingViewingKey) {
         let k = self.rivk.0.to_bytes();
         let b = [(&self.ak.0).into(), self.nk.0.to_bytes()];
-        let r = prf_expand_vec(&k, &[&[0x82], &b[0][..], &b[1][..]]);
+        let r = PrfExpand::OrchardDkOvk.with_ad_slices(&k, &[&b[0][..], &b[1][..]]);
         (
             DiversifierKey(r[..32].try_into().unwrap()),
             OutgoingViewingKey(r[32..].try_into().unwrap()),

src/note.rs

@@ -6,7 +6,7 @@ use subtle::CtOption;
 
 use crate::{
     keys::{FullViewingKey, SpendingKey},
-    spec::{prf_expand_vec, to_base, to_scalar},
+    spec::{to_base, to_scalar, PrfExpand},
     value::NoteValue,
     Address,
 };
@@ -32,24 +32,23 @@ impl RandomSeed {
     ///
     /// [orchardsend]: https://zips.z.cash/protocol/nu5.pdf#orchardsend
     fn psi(&self, rho: &Nullifier) -> pallas::Base {
-        to_base(prf_expand_vec(&self.0, &[&[0x09], &rho.to_bytes()[..]]))
+        to_base(PrfExpand::Psi.with_ad(&self.0, &rho.to_bytes()[..]))
     }
 
     /// Defined in [Zcash Protocol Spec § 4.7.3: Sending Notes (Orchard)][orchardsend].
     ///
     /// [orchardsend]: https://zips.z.cash/protocol/nu5.pdf#orchardsend
     fn esk(&self, rho: &Nullifier) -> pallas::Scalar {
-        to_scalar(prf_expand_vec(&self.0, &[&[0x04], &rho.to_bytes()[..]]))
+        to_scalar(PrfExpand::Esk.with_ad(&self.0, &rho.to_bytes()[..]))
     }
 
     /// Defined in [Zcash Protocol Spec § 4.7.3: Sending Notes (Orchard)][orchardsend].
     ///
     /// [orchardsend]: https://zips.z.cash/protocol/nu5.pdf#orchardsend
     fn rcm(&self, rho: &Nullifier) -> commitment::NoteCommitTrapdoor {
-        commitment::NoteCommitTrapdoor(to_scalar(prf_expand_vec(
-            &self.0,
-            &[&[0x05], &rho.to_bytes()[..]],
-        )))
+        commitment::NoteCommitTrapdoor(to_scalar(
+            PrfExpand::Rcm.with_ad(&self.0, &rho.to_bytes()[..]),
+        ))
     }
 }
 

src/spec.rs

@@ -15,7 +15,7 @@ use crate::{
 };
 
 mod prf_expand;
-pub(crate) use prf_expand::{prf_expand, prf_expand_vec};
+pub(crate) use prf_expand::PrfExpand;
 
 /// A Pallas point that is guaranteed to not be the identity.
 #[derive(Clone, Copy, Debug)]

src/spec/prf_expand.rs

@@ -2,23 +2,70 @@ use blake2b_simd::Params;
 
 const PRF_EXPAND_PERSONALIZATION: &[u8; 16] = b"Zcash_ExpandSeed";
 
-/// $PRF^\mathsf{expand}(sk, t) := BLAKE2b-512("Zcash_ExpandSeed", sk || t)$
-///
-/// Defined in [Zcash Protocol Spec § 5.4.2: Pseudo Random Functions][concreteprfs].
-///
-/// [concreteprfs]: https://zips.z.cash/protocol/nu5.pdf#concreteprfs
-pub(crate) fn prf_expand(sk: &[u8], t: &[u8]) -> [u8; 64] {
-    prf_expand_vec(sk, &[t])
+/// The set of domains in which $PRF^\mathsf{expand}$ is defined.
+pub(crate) enum PrfExpand {
+    Esk,
+    Rcm,
+    OrchardAsk,
+    OrchardNk,
+    OrchardRivk,
+    Psi,
+    OrchardDkOvk,
 }
 
-pub(crate) fn prf_expand_vec(sk: &[u8], ts: &[&[u8]]) -> [u8; 64] {
-    let mut h = Params::new()
-        .hash_length(64)
-        .personal(PRF_EXPAND_PERSONALIZATION)
-        .to_state();
-    h.update(sk);
-    for t in ts {
-        h.update(t);
+impl PrfExpand {
+    fn domain_separator(&self) -> u8 {
+        match self {
+            Self::Esk => 0x04,
+            Self::Rcm => 0x05,
+            Self::OrchardAsk => 0x06,
+            Self::OrchardNk => 0x07,
+            Self::OrchardRivk => 0x08,
+            Self::Psi => 0x09,
+            Self::OrchardDkOvk => 0x82,
+        }
+    }
+
+    /// Expands the given secret key in this domain, with no additional data.
+    ///
+    /// $PRF^\mathsf{expand}(sk, dst) := BLAKE2b-512("Zcash_ExpandSeed", sk || dst)$
+    ///
+    /// Defined in [Zcash Protocol Spec § 5.4.2: Pseudo Random Functions][concreteprfs].
+    ///
+    /// [concreteprfs]: https://zips.z.cash/protocol/nu5.pdf#concreteprfs
+    pub(crate) fn expand(self, sk: &[u8]) -> [u8; 64] {
+        self.with_ad_slices(sk, &[])
+    }
+
+    /// Expands the given secret key in this domain, with the given additional data.
+    ///
+    /// $PRF^\mathsf{expand}(sk, dst, t) := BLAKE2b-512("Zcash_ExpandSeed", sk || dst || t)$
+    ///
+    /// Defined in [Zcash Protocol Spec § 5.4.2: Pseudo Random Functions][concreteprfs].
+    ///
+    /// [concreteprfs]: https://zips.z.cash/protocol/nu5.pdf#concreteprfs
+    pub(crate) fn with_ad(self, sk: &[u8], t: &[u8]) -> [u8; 64] {
+        self.with_ad_slices(sk, &[t])
+    }
+
+    /// Expands the given secret key in this domain, with additional data concatenated
+    /// from the given slices.
+    ///
+    /// $PRF^\mathsf{expand}(sk, dst, a, b, ...) := BLAKE2b-512("Zcash_ExpandSeed", sk || dst || a || b || ...)$
+    ///
+    /// Defined in [Zcash Protocol Spec § 5.4.2: Pseudo Random Functions][concreteprfs].
+    ///
+    /// [concreteprfs]: https://zips.z.cash/protocol/nu5.pdf#concreteprfs
+    pub(crate) fn with_ad_slices(self, sk: &[u8], ts: &[&[u8]]) -> [u8; 64] {
+        let mut h = Params::new()
+            .hash_length(64)
+            .personal(PRF_EXPAND_PERSONALIZATION)
+            .to_state();
+        h.update(sk);
+        h.update(&[self.domain_separator()]);
+        for t in ts {
+            h.update(t);
+        }
+        *h.finalize().as_array()
     }
-    *h.finalize().as_array()
 }