simplify state logic for reduced prf even more

Author: th4s

crates/components/hmac-sha256/src/prf/function.rs

@@ -234,6 +234,7 @@ mod tests {
             .unwrap();
         }
 
+        assert_eq!(prf_out_leader.len(), 2);
         assert_eq!(prf_out_leader.len(), prf_out_follower.len());
 
         let prf_result_leader: Vec<u8> = prf_out_leader

crates/components/hmac-sha256/src/prf/function/reduced.rs

@@ -17,13 +17,31 @@ pub(crate) struct PrfFunction {
     label: &'static [u8],
     // The start seed and the label, e.g. client_random + server_random + "master_secret".
     start_seed_label: Vec<u8>,
-    // The current HMAC message needed for a[i]
-    a_msg: Vec<u8>,
-    inner_partial: InnerPartial,
+    iterations: usize,
+    state: PrfState,
     a: Vec<PHash>,
     p: Vec<PHash>,
 }
 
+#[derive(Debug)]
+enum PrfState {
+    InnerPartial {
+        inner_partial: DecodeFutureTyped<BitVec, [u32; 8]>,
+    },
+    ComputeA {
+        iter: usize,
+        inner_partial: [u32; 8],
+        msg: Vec<u8>,
+    },
+    ComputeP {
+        iter: usize,
+        inner_partial: [u32; 8],
+        a_output: DecodeFutureTyped<BitVec, [u8; 32]>,
+    },
+    ComputeLastP,
+    Done,
+}
+
 impl PrfFunction {
     const MS_LABEL: &[u8] = b"master secret";
     const KEY_LABEL: &[u8] = b"key expansion";
@@ -63,51 +81,68 @@ impl PrfFunction {
     }
 
     pub(crate) fn wants_flush(&mut self) -> bool {
-        let last_p = self.p.last().expect("Prf should be allocated");
-
-        if let State::Done = last_p.state {
+        if let PrfState::Done = self.state {
             return false;
         }
         true
     }
 
     pub(crate) fn flush(&mut self, vm: &mut dyn Vm<Binary>) -> Result<(), PrfError> {
-        let inner_partial = self.inner_partial.try_recv()?;
-        let Some(inner_partial) = inner_partial else {
-            return Ok(());
-        };
-
-        for (a, p) in self.a.iter_mut().zip(self.p.iter_mut()) {
-            match &mut a.state {
-                State::Init { .. } => {
-                    a.assign_inner_local(vm, inner_partial, &self.a_msg)?;
-                    break;
-                }
-                State::Assigned { output } => {
-                    if let Some(output) = output.try_recv().map_err(PrfError::vm)? {
-                        let output = output.to_vec();
-                        a.state = State::Decoded {
-                            output: output.clone(),
-                        };
-                        self.a_msg = output;
-                    }
-                }
-                _ => (),
+        match &mut self.state {
+            PrfState::InnerPartial { inner_partial } => {
+                let Some(inner_partial) = inner_partial.try_recv().map_err(PrfError::vm)? else {
+                    return Ok(());
+                };
+
+                self.state = PrfState::ComputeA {
+                    iter: 0,
+                    inner_partial,
+                    msg: self.start_seed_label.clone(),
+                };
+                self.flush(vm)?;
             }
-
-            match &mut p.state {
-                State::Init { .. } => {
-                    if let State::Decoded { output } = &a.state {
-                        let mut p_msg = output.to_vec();
-                        p_msg.extend_from_slice(&self.start_seed_label);
-                        p.assign_inner_local(vm, inner_partial, &p_msg)?;
+            PrfState::ComputeA {
+                iter,
+                inner_partial,
+                msg,
+            } => {
+                let a = &self.a[*iter];
+                assign_inner_local(vm, a.inner_local, *inner_partial, msg)?;
+
+                let a_output = vm.decode(a.output).map_err(PrfError::vm)?;
+                self.state = PrfState::ComputeP {
+                    iter: *iter,
+                    inner_partial: *inner_partial,
+                    a_output,
+                };
+            }
+            PrfState::ComputeP {
+                iter,
+                inner_partial,
+                a_output,
+            } => {
+                let Some(output) = a_output.try_recv().map_err(PrfError::vm)? else {
+                    return Ok(());
+                };
+                let p = &self.p[*iter];
+
+                let mut msg = output.to_vec();
+                msg.extend_from_slice(&self.start_seed_label);
+
+                assign_inner_local(vm, p.inner_local, *inner_partial, &msg)?;
+
+                if *iter == self.iterations {
+                    self.state = PrfState::ComputeLastP;
+                } else {
+                    self.state = PrfState::ComputeA {
+                        iter: *iter + 1,
+                        inner_partial: *inner_partial,
+                        msg: output.to_vec(),
                     }
-                }
-                State::Assigned { .. } => {
-                    p.state = State::Done;
-                }
-                _ => (),
+                };
             }
+            PrfState::ComputeLastP => self.state = PrfState::Done,
+            _ => (),
         }
 
         Ok(())
@@ -117,8 +152,7 @@ impl PrfFunction {
         let mut start_seed_label = self.label.to_vec();
         start_seed_label.extend_from_slice(&seed);
 
-        self.start_seed_label = start_seed_label.clone();
-        self.a_msg = start_seed_label;
+        self.start_seed_label = start_seed_label;
     }
 
     pub(crate) fn output(&self) -> Vec<Array<U8, 32>> {
@@ -132,6 +166,10 @@ impl PrfFunction {
         inner_partial: Sha256,
         len: usize,
     ) -> Result<Self, PrfError> {
+        assert!(len > 0, "cannot compute 0 bytes for prf");
+
+        let iterations = len / 32 + ((len % 32) != 0) as usize;
+
         let (inner_partial, _) = inner_partial
             .state()
             .expect("state should be set for inner_partial");
@@ -140,100 +178,55 @@ impl PrfFunction {
         let mut prf = Self {
             label,
             start_seed_label: vec![],
-            a_msg: vec![],
-            inner_partial: InnerPartial::Decoding(inner_partial),
+            // used for indexing, so we need to subtract one here
+            iterations: iterations - 1,
+            state: PrfState::InnerPartial { inner_partial },
             a: vec![],
             p: vec![],
         };
 
-        assert!(len > 0, "cannot compute 0 bytes for prf");
-
-        let iterations = len / 32 + ((len % 32) != 0) as usize;
-
         for _ in 0..iterations {
-            let a = PHash::alloc(vm, outer_partial.clone())?;
-            prf.a.push(a);
-
-            let p = PHash::alloc(vm, outer_partial.clone())?;
-            prf.p.push(p);
+            // setup A[i]
+            let inner_local: Array<U8, 32> = vm.alloc().map_err(PrfError::vm)?;
+            let output = hmac_sha256(vm, outer_partial.clone(), inner_local)?;
+            let p_hash = PHash {
+                inner_local,
+                output,
+            };
+            prf.a.push(p_hash);
+
+            // setup P[i]
+            let inner_local: Array<U8, 32> = vm.alloc().map_err(PrfError::vm)?;
+            let output = hmac_sha256(vm, outer_partial.clone(), inner_local)?;
+            let p_hash = PHash {
+                inner_local,
+                output,
+            };
+            prf.p.push(p_hash);
         }
 
         Ok(prf)
     }
 }
 
-#[derive(Debug)]
-struct PHash {
-    output: Array<U8, 32>,
-    state: State,
-}
-
-impl PHash {
-    fn alloc(vm: &mut dyn Vm<Binary>, outer_partial: Sha256) -> Result<Self, PrfError> {
-        let inner_local: Array<U8, 32> = vm.alloc().map_err(PrfError::vm)?;
-        let output = hmac_sha256(vm, outer_partial, inner_local)?;
-
-        let p_hash = Self {
-            state: State::Init { inner_local },
-            output,
-        };
-
-        Ok(p_hash)
-    }
-
-    fn assign_inner_local(
-        &mut self,
-        vm: &mut dyn Vm<Binary>,
-        inner_partial: [u32; 8],
-        msg: &[u8],
-    ) -> Result<(), PrfError> {
-        if let State::Init { inner_local, .. } = self.state {
-            let inner_local_value = sha256(inner_partial, 64, msg);
-
-            vm.mark_public(inner_local).map_err(PrfError::vm)?;
-            vm.assign(inner_local, state_to_bytes(inner_local_value))
-                .map_err(PrfError::vm)?;
-            vm.commit(inner_local).map_err(PrfError::vm)?;
-
-            let output = vm.decode(self.output).map_err(PrfError::vm)?;
-            self.state = State::Assigned { output };
-        }
-
-        Ok(())
-    }
-}
+fn assign_inner_local(
+    vm: &mut dyn Vm<Binary>,
+    inner_local: Array<U8, 32>,
+    inner_partial: [u32; 8],
+    msg: &[u8],
+) -> Result<(), PrfError> {
+    let inner_local_value = sha256(inner_partial, 64, msg);
 
-#[derive(Debug)]
-enum State {
-    Init {
-        inner_local: Array<U8, 32>,
-    },
-    Assigned {
-        output: DecodeFutureTyped<BitVec, [u8; 32]>,
-    },
-    Decoded {
-        output: Vec<u8>,
-    },
-    Done,
-}
+    vm.mark_public(inner_local).map_err(PrfError::vm)?;
+    vm.assign(inner_local, state_to_bytes(inner_local_value))
+        .map_err(PrfError::vm)?;
+    vm.commit(inner_local).map_err(PrfError::vm)?;
 
-#[derive(Debug)]
-enum InnerPartial {
-    Decoding(DecodeFutureTyped<BitVec, [u32; 8]>),
-    Finished([u32; 8]),
+    Ok(())
 }
 
-impl InnerPartial {
-    pub(crate) fn try_recv(&mut self) -> Result<Option<[u32; 8]>, PrfError> {
-        match self {
-            InnerPartial::Decoding(value) => {
-                let value = value.try_recv().map_err(PrfError::vm)?;
-                if let Some(value) = value {
-                    *self = InnerPartial::Finished(value);
-                }
-                Ok(value)
-            }
-            InnerPartial::Finished(value) => Ok(Some(*value)),
-        }
-    }
+#[derive(Debug, Clone, Copy)]
+struct PHash {
+    inner_local: Array<U8, 32>,
+    output: Array<U8, 32>,
 }