Noise sv2 hardening

sv2/codec-sv2/src/decoder.rs

@@ -39,8 +39,6 @@ use framing_sv2::{
 #[cfg(feature = "noise_sv2")]
 use framing_sv2::{ENCRYPTED_SV2_FRAME_HEADER_SIZE, SV2_FRAME_CHUNK_SIZE, SV2_FRAME_HEADER_SIZE};
 #[cfg(feature = "noise_sv2")]
-use noise_sv2::NoiseEngine;
-#[cfg(feature = "noise_sv2")]
 use noise_sv2::NOISE_FRAME_HEADER_SIZE;
 
 #[cfg(feature = "noise_sv2")]
@@ -49,7 +47,7 @@ use crate::error::Result;
 
 use crate::Error::MissingBytes;
 #[cfg(feature = "noise_sv2")]
-use crate::State;
+use crate::{State, TransportDecryptState};
 
 #[cfg(not(feature = "with_buffer_pool"))]
 use buffer_sv2::{Buffer as IsBuffer, BufferFromSystemMemory as Buffer};
@@ -155,31 +153,49 @@ impl<'a, T: Serialize + GetSize + Deserialize<'a>, B: IsBuffer + AeadBuffer> Wit
                     }
                 }
             }
-            State::Transport(noise_engine) => {
-                let hint = if IsBuffer::len(&self.sv2_buffer) < SV2_FRAME_HEADER_SIZE {
-                    let len = IsBuffer::len(&self.noise_buffer);
-                    let src = self.noise_buffer.get_data_by_ref(len);
-                    if src.len() < ENCRYPTED_SV2_FRAME_HEADER_SIZE {
-                        ENCRYPTED_SV2_FRAME_HEADER_SIZE - src.len()
-                    } else {
-                        0
-                    }
-                } else {
-                    let src = self.sv2_buffer.get_data_by_ref(SV2_FRAME_HEADER_SIZE);
-                    let header = Header::from_bytes(src)?;
-                    header.encrypted_len() - IsBuffer::len(&self.noise_buffer)
-                };
+            State::Transport(engine) => {
+                self.next_transport(|buf| engine.decrypt(buf).map_err(Into::into))
+            }
+        }
+    }
 
-                match hint {
-                    0 => {
-                        self.missing_noise_b = ENCRYPTED_SV2_FRAME_HEADER_SIZE;
-                        self.decode_noise_frame(noise_engine)
-                    }
-                    _ => {
-                        self.missing_noise_b = hint;
-                        Err(Error::MissingBytes(hint))
-                    }
-                }
+    /// Attempts to decode the next Noise frame with the decrypting half of a split [`State`].
+    #[inline]
+    pub fn next_transport_frame(
+        &mut self,
+        state: &mut TransportDecryptState,
+    ) -> Result<Frame<T, B::Slice>> {
+        self.next_transport(|buf| state.decrypt(buf))
+    }
+
+    // Decodes a transport-mode frame, decrypting through `decrypt`.
+    #[inline]
+    fn next_transport(
+        &mut self,
+        decrypt: impl FnMut(&mut B) -> Result<()>,
+    ) -> Result<Frame<T, B::Slice>> {
+        let hint = if IsBuffer::len(&self.sv2_buffer) < SV2_FRAME_HEADER_SIZE {
+            let len = IsBuffer::len(&self.noise_buffer);
+            let src = self.noise_buffer.get_data_by_ref(len);
+            if src.len() < ENCRYPTED_SV2_FRAME_HEADER_SIZE {
+                ENCRYPTED_SV2_FRAME_HEADER_SIZE - src.len()
+            } else {
+                0
+            }
+        } else {
+            let src = self.sv2_buffer.get_data_by_ref(SV2_FRAME_HEADER_SIZE);
+            let header = Header::from_bytes(src)?;
+            header.encrypted_len() - IsBuffer::len(&self.noise_buffer)
+        };
+
+        match hint {
+            0 => {
+                self.missing_noise_b = ENCRYPTED_SV2_FRAME_HEADER_SIZE;
+                self.decode_noise_frame(decrypt)
+            }
+            _ => {
+                self.missing_noise_b = hint;
+                Err(Error::MissingBytes(hint))
             }
         }
     }
@@ -255,7 +271,32 @@ impl<'a, T: Serialize + GetSize + Deserialize<'a>, B: IsBuffer + AeadBuffer> Wit
     // additional bytes required to fully decrypt the frame. Once all bytes are available, the
     // decryption process completes and the frame can be successfully decoded.
     #[inline]
-    fn decode_noise_frame(&mut self, noise_engine: &mut NoiseEngine) -> Result<Frame<T, B::Slice>> {
+    fn decode_noise_frame(
+        &mut self,
+        decrypt: impl FnMut(&mut B) -> Result<()>,
+    ) -> Result<Frame<T, B::Slice>> {
+        let result = self.try_decode_noise_frame(decrypt);
+
+        match &result {
+            // `MissingBytes` is the normal way out of the header round, so the buffer has to be
+            // left exactly as it is.
+            Err(Error::MissingBytes(_)) | Ok(_) => {}
+            Err(_) => {
+                // Not a no-op: the decrypt offset and the plaintext decrypted so far persist across
+                // calls, so without this the next frame is decrypted at the failing chunk's offset.
+                self.sv2_buffer.danger_set_start(0);
+                self.sv2_buffer.get_data_owned();
+            }
+        }
+
+        result
+    }
+
+    #[inline]
+    fn try_decode_noise_frame(
+        &mut self,
+        mut decrypt: impl FnMut(&mut B) -> Result<()>,
+    ) -> Result<Frame<T, B::Slice>> {
         match (
             IsBuffer::len(&self.noise_buffer),
             IsBuffer::len(&self.sv2_buffer),
@@ -268,7 +309,7 @@ impl<'a, T: Serialize + GetSize + Deserialize<'a>, B: IsBuffer + AeadBuffer> Wit
                     .get_writable(ENCRYPTED_SV2_FRAME_HEADER_SIZE);
                 decrypted_header.copy_from_slice(src.as_ref());
                 self.sv2_buffer.as_ref();
-                noise_engine.decrypt(&mut self.sv2_buffer)?;
+                decrypt(&mut self.sv2_buffer)?;
                 let header =
                     Header::from_bytes(self.sv2_buffer.get_data_by_ref(SV2_FRAME_HEADER_SIZE))?;
                 self.missing_noise_b = header.encrypted_len();
@@ -292,7 +333,7 @@ impl<'a, T: Serialize + GetSize + Deserialize<'a>, B: IsBuffer + AeadBuffer> Wit
                     let decrypted_payload = self.sv2_buffer.get_writable(end - start);
                     decrypted_payload.copy_from_slice(&encrypted_payload.as_ref()[start..end]);
                     self.sv2_buffer.danger_set_start(decrypted_len);
-                    noise_engine.decrypt(&mut self.sv2_buffer)?;
+                    decrypt(&mut self.sv2_buffer)?;
                     start = end;
                     end = (start + SV2_FRAME_CHUNK_SIZE).min(encrypted_payload_len);
                     decrypted_len += self.sv2_buffer.as_ref().len();
@@ -798,6 +839,47 @@ mod prop_tests {
         }
     }
 
+    #[cfg(feature = "noise_sv2")]
+    #[test]
+    fn noise_decoder_recovers_from_a_failed_decryption() {
+        let (mut sender_state, mut receiver_state) = make_transport_state_pair();
+        let frame = Frame::Sv2(
+            Sv2Frame::<TestMessage, Slice>::from_message(TestMessage { value: 7 }, 0, 0, false)
+                .unwrap(),
+        );
+        let mut encoder = NoiseEncoder::<TestMessage>::new();
+        let encrypted = encoder.encode(frame, &mut sender_state).unwrap();
+        let encrypted: &[u8] = encrypted.as_ref();
+
+        let mut decoder = StandardNoiseDecoder::<TestMessage>::new();
+
+        // Fail on the encrypted header. The closure never touches `receiver_state`, so its nonce
+        // stays where it was and the same bytes can be replayed below.
+        let hint = decoder
+            .next_transport(|_| Err(crate::Error::UnexpectedNoiseState))
+            .unwrap_err();
+        assert!(matches!(hint, crate::Error::MissingBytes(_)));
+        let writable = decoder.writable();
+        let len = writable.len();
+        writable.copy_from_slice(&encrypted[..len]);
+        let failed = decoder
+            .next_transport(|_| Err(crate::Error::UnexpectedNoiseState))
+            .unwrap_err();
+        assert!(matches!(failed, crate::Error::UnexpectedNoiseState));
+        assert_eq!(IsBuffer::len(&decoder.sv2_buffer), 0);
+        assert!(decoder.sv2_buffer.as_ref().is_empty());
+
+        // The same decoder must now decode the frame from the start.
+        let decoded = decode_noise_frame(&mut decoder, &mut receiver_state, encrypted);
+        match decoded {
+            Some(mut f) => assert_eq!(
+                binary_sv2::from_bytes::<TestMessage>(f.payload()).unwrap(),
+                TestMessage { value: 7 }
+            ),
+            None => panic!("failed to decode the frame after a failed decryption"),
+        }
+    }
+
     /// Verifies that a single `StandardNoiseDecoder` instance correctly
     /// decodes two consecutive noise-encrypted frames in sequence using
     /// the same shared transport state.