WGPS: Add ReadyTransport and friends

wgps/Cargo.toml

@@ -6,6 +6,10 @@ edition = "2021"
 [dependencies]
 willow-data-model = { path = "../data-model", version = "0.1.0" }
 ufotofu = { version = "0.4.2", features = ["std"] }
+either = "1.10.0"
+
+[dev-dependencies]
+smol = "2.0.0"
 
 [lints]
 workspace = true

wgps/src/lib.rs

@@ -7,6 +7,9 @@ use willow_data_model::{
     ResumptionFailedError, Store, StoreEvent, SubspaceId,
 };
 
+mod ready_transport;
+pub use ready_transport::*;
+
 /// Options to specify how ranges should be partitioned.
 #[derive(Debug, Clone, Copy)]
 pub struct PartitionOpts {

wgps/src/ready_transport.rs

@@ -0,0 +1,188 @@
+use either::Either;
+use ufotofu::local_nb::BulkProducer;
+
+/** When things go wrong while trying to make a WGPS transport ready. */
+#[derive(Debug)]
+pub enum ReadyTransportError<E: core::fmt::Display> {
+    /** The transport returned an error of its own. */
+    Transport(E),
+    /** The received max payload power was invalid, i.e. greater than 64. */
+    MaxPayloadInvalid,
+    /** The transport stopped producing bytes before it could be deemed ready. */
+    FinishedTooSoon,
+}
+
+impl<E: core::fmt::Display> core::fmt::Display for ReadyTransportError<E> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            ReadyTransportError::Transport(e) => write!(f, "{}", e),
+            ReadyTransportError::MaxPayloadInvalid => write!(
+                f,
+                "The received max payload power was invalid, i.e. greater than 64."
+            ),
+            ReadyTransportError::FinishedTooSoon => write!(
+                f,
+                "The transport stopped producing bytes before it could be deemed ready."
+            ),
+        }
+    }
+}
+
+/** The result of intercepting the first few bytes of a WGPS transport. */
+#[derive(Debug)]
+#[allow(dead_code)] // TODO: Remove when this is used.
+pub(crate) struct ReadyTransport<
+    const CHALLENGE_HASH_LENGTH: usize,
+    E: core::fmt::Display,
+    P: BulkProducer<Item = u8, Final = (), Error = E>,
+> {
+    /** The maximum payload size which may be sent without being explicitly requested.*/
+    pub maximum_payload_size: usize,
+    /** The challenge hash of a nonce. */
+    pub received_commitment: [u8; CHALLENGE_HASH_LENGTH],
+    /** A 'ready' transport set to immediately produce encoded WGPS messages. */
+    pub transport: P,
+}
+
+impl<
+        const CHALLENGE_HASH_LENGTH: usize,
+        E: core::fmt::Display,
+        P: BulkProducer<Item = u8, Final = (), Error = E>,
+    > ReadyTransport<CHALLENGE_HASH_LENGTH, E, P>
+{
+    #[allow(dead_code)] // TODO: Remove when this is used.
+    pub(crate) fn transport(&self) -> &P {
+        &self.transport
+    }
+}
+
+/** Given a producer of bytes which is to immediately produce the bytes corresponding to the WGPS' [maximum payload size](https://willowprotocol.org/specs/sync/index.html#peer_max_payload_size) and [received commitment](https://willowprotocol.org/specs/sync/index.html#received_commitment), returns the computed maximum payload size, received commitment, and a 'ready' transport set to produce encoded WGPS messages.
+*/
+#[allow(dead_code)] // TODO: Remove when this is used.
+pub(crate) async fn ready_transport<
+    const CHALLENGE_HASH_LENGTH: usize,
+    E: core::fmt::Display,
+    P: BulkProducer<Item = u8, Final = (), Error = E>,
+>(
+    mut transport: P,
+) -> Result<ReadyTransport<CHALLENGE_HASH_LENGTH, E, P>, ReadyTransportError<E>> {
+    let maximum_payload_power = match transport.produce().await? {
+        Either::Left(byte) => byte,
+        Either::Right(_) => return Err(ReadyTransportError::FinishedTooSoon),
+    };
+
+    if maximum_payload_power > 64 {
+        return Err(ReadyTransportError::MaxPayloadInvalid);
+    }
+
+    let maximum_payload_size = 2_usize.pow(maximum_payload_power as u32);
+
+    let mut received_commitment = [0_u8; CHALLENGE_HASH_LENGTH];
+
+    if let Err(e) = transport
+        .bulk_overwrite_full_slice(&mut received_commitment)
+        .await
+    {
+        match e.reason {
+            Either::Left(_) => return Err(ReadyTransportError::FinishedTooSoon),
+            Either::Right(e) => return Err(ReadyTransportError::Transport(e)),
+        }
+    };
+
+    Ok(ReadyTransport {
+        maximum_payload_size,
+        received_commitment,
+        transport,
+    })
+}
+
+impl<E: core::fmt::Display> From<E> for ReadyTransportError<E> {
+    fn from(value: E) -> Self {
+        ReadyTransportError::Transport(value)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+
+    use super::*;
+
+    use ufotofu::local_nb::producer::FromSlice;
+
+    #[test]
+    fn empty_producer() {
+        let empty_transport = FromSlice::new(&[]);
+
+        smol::block_on(async {
+            let result = ready_transport::<4, _, _>(empty_transport).await;
+
+            assert!(matches!(result, Err(ReadyTransportError::FinishedTooSoon)))
+        });
+    }
+
+    #[test]
+    fn only_power_producer() {
+        let only_power_transport = FromSlice::new(&[0_u8]);
+
+        smol::block_on(async {
+            let result = ready_transport::<4, _, _>(only_power_transport).await;
+
+            assert!(matches!(result, Err(ReadyTransportError::FinishedTooSoon)))
+        });
+    }
+
+    #[test]
+    fn invalid_power_producer() {
+        let only_power_transport = FromSlice::new(&[65_u8]);
+
+        smol::block_on(async {
+            let result = ready_transport::<4, _, _>(only_power_transport).await;
+
+            assert!(matches!(
+                result,
+                Err(ReadyTransportError::MaxPayloadInvalid)
+            ))
+        });
+    }
+
+    #[test]
+    fn invalid_power_producer_correct_length() {
+        let only_power_transport = FromSlice::new(&[65_u8, 0, 0, 0, 0]);
+
+        smol::block_on(async {
+            let result = ready_transport::<4, _, _>(only_power_transport).await;
+
+            assert!(matches!(
+                result,
+                Err(ReadyTransportError::MaxPayloadInvalid)
+            ))
+        });
+    }
+
+    #[test]
+    fn commitment_too_short() {
+        let only_power_transport = FromSlice::new(&[0_u8, 0]);
+
+        smol::block_on(async {
+            let result = ready_transport::<4, _, _>(only_power_transport).await;
+
+            assert!(matches!(result, Err(ReadyTransportError::FinishedTooSoon)))
+        });
+    }
+
+    #[test]
+    fn success() {
+        let only_power_transport = FromSlice::new(&[1_u8, 1, 2, 3, 4, 5]);
+
+        smol::block_on(async {
+            let result = ready_transport::<4, _, _>(only_power_transport).await;
+
+            if let Ok(ready) = result {
+                assert!(ready.maximum_payload_size == 2);
+                assert!(ready.received_commitment == [1, 2, 3, 4]);
+            } else {
+                panic!()
+            }
+        });
+    }
+}