Merge pull request #128 from chrysn-pull-requests/typestatify

Typestatification

examples/coap/src/bin/coapclient.rs

@@ -26,7 +26,7 @@ fn main() {
     // Send Message 1 over CoAP and convert the response to byte
     let mut msg_1_buf = Vec::from([0xf5u8]); // EDHOC message_1 when transported over CoAP is prepended with CBOR true
     let c_i = generate_connection_identifier_cbor();
-    let message_1 = initiator.prepare_message_1(c_i).unwrap();
+    let (initiator, message_1) = initiator.prepare_message_1(c_i).unwrap();
     msg_1_buf.extend_from_slice(&message_1.content[..message_1.len]);
     println!("message_1 len = {}", msg_1_buf.len());
 
@@ -37,15 +37,15 @@ fn main() {
     println!("response_vec = {:02x?}", response.message.payload);
     println!("message_2 len = {}", response.message.payload.len());
 
-    let c_r = initiator.process_message_2(
+    let m2result = initiator.process_message_2(
         &response.message.payload[..]
             .try_into()
             .expect("wrong length"),
     );
 
-    if c_r.is_ok() {
-        let mut msg_3 = Vec::from([c_r.unwrap()]);
-        let (message_3, prk_out) = initiator.prepare_message_3().unwrap();
+    if let Ok((initiator, c_r)) = m2result {
+        let mut msg_3 = Vec::from([c_r]);
+        let (mut initiator, message_3, prk_out) = initiator.prepare_message_3().unwrap();
         msg_3.extend_from_slice(&message_3.content[..message_3.len]);
         println!("message_3 len = {}", msg_3.len());
 
@@ -76,6 +76,6 @@ fn main() {
         println!("OSCORE secret after key update: {:02x?}", _oscore_secret);
         println!("OSCORE salt after key update: {:02x?}", _oscore_salt);
     } else {
-        panic!("Message 2 processing error: {:#?}", c_r);
+        panic!("Message 2 processing error: {:#?}", m2result);
     }
 }

examples/coap/src/bin/coapserver-coaphandler.rs

@@ -15,16 +15,18 @@ const R: &[u8] = &hex!("72cc4761dbd4c78f758931aa589d348d1ef874a7e303ede2f140dcf3
 
 #[derive(Default, Debug)]
 struct EdhocHandler {
-    connections: Vec<(u8, EdhocResponder<'static>)>,
+    connections: Vec<(u8, EdhocResponderWaitM3<'static>)>,
 }
 
 impl EdhocHandler {
-    fn connection_by_c_r(&mut self, c_r: u8) -> Option<&mut EdhocResponder<'static>> {
-        self.connections
-            .iter_mut()
-            .filter(|(current_c_r, _)| current_c_r == &c_r)
-            .map(|(_, responder)| responder)
-            .next()
+    fn take_connection_by_c_r(&mut self, c_r: u8) -> Option<EdhocResponderWaitM3<'static>> {
+        let index = self
+            .connections
+            .iter()
+            .position(|(current_c_r, _)| current_c_r == &c_r)?;
+        let last = self.connections.len() - 1;
+        self.connections.swap(index, last);
+        Some(self.connections.pop().unwrap().1)
     }
 
     fn new_c_r(&self) -> u8 {
@@ -40,7 +42,12 @@ impl EdhocHandler {
 }
 
 enum EdhocResponse {
-    OkSend2 { c_r: u8 },
+    // We could also store the responder in the Vec (once we're done rendering the response, we'll
+    // take up a slot there anyway) if we make it an enum.
+    OkSend2 {
+        c_r: u8,
+        responder: EdhocResponderBuildM2<'static>,
+    },
     Message3Processed,
 }
 
@@ -54,33 +61,33 @@ impl coap_handler::Handler for EdhocHandler {
 
         if starts_with_true {
             let state = EdhocState::default();
-            let mut responder = EdhocResponder::new(state, &R, &CRED_R, Some(&CRED_I));
+            let responder = EdhocResponder::new(state, &R, &CRED_R, Some(&CRED_I));
 
-            let error = responder
+            let response = responder
                 .process_message_1(&request.payload()[1..].try_into().expect("wrong length"));
 
-            if error.is_ok() {
+            if let Ok(responder) = response {
                 let c_r = self.new_c_r();
-                // save edhoc connection
-                self.connections.push((c_r, responder));
-                EdhocResponse::OkSend2 { c_r }
+                EdhocResponse::OkSend2 { c_r, responder }
             } else {
                 panic!("How to respond to non-OK?")
             }
         } else {
             // potentially message 3
             let c_r_rcvd = request.payload()[0];
-            let mut responder = self.connection_by_c_r(c_r_rcvd).expect("No such C_R found");
+            let responder = self
+                .take_connection_by_c_r(c_r_rcvd)
+                .expect("No such C_R found");
 
             println!("Found state with connection identifier {:?}", c_r_rcvd);
-            let prk_out = responder
+            let result = responder
                 .process_message_3(&request.payload()[1..].try_into().expect("wrong length"));
 
-            if prk_out.is_err() {
-                println!("EDHOC processing error: {:?}", prk_out);
+            let Ok((mut responder, prk_out)) = result else {
+                println!("EDHOC processing error: {:?}", result);
                 // FIXME remove state from edhoc_connections
                 panic!("Handler can't just not respond");
-            }
+            };
 
             println!("EDHOC exchange successfully completed");
             println!("PRK_out: {:02x?}", prk_out);
@@ -115,9 +122,9 @@ impl coap_handler::Handler for EdhocHandler {
     ) {
         response.set_code(coap_numbers::code::CHANGED.try_into().ok().unwrap());
         match req {
-            EdhocResponse::OkSend2 { c_r } => {
-                let responder = self.connection_by_c_r(c_r).unwrap();
-                let message_2 = responder.prepare_message_2(c_r).unwrap();
+            EdhocResponse::OkSend2 { c_r, responder } => {
+                let (responder, message_2) = responder.prepare_message_2(c_r).unwrap();
+                self.connections.push((c_r, responder));
                 response.set_payload(&message_2.content[..message_2.len]);
             }
             EdhocResponse::Message3Processed => (), // "send empty ack back"?
@@ -128,7 +135,7 @@ impl coap_handler::Handler for EdhocHandler {
 fn build_handler() -> impl coap_handler::Handler {
     use coap_handler_implementations::{HandlerBuilder, ReportingHandlerBuilder};
 
-    let mut edhoc: EdhocHandler = Default::default();
+    let edhoc: EdhocHandler = Default::default();
 
     coap_handler_implementations::new_dispatcher()
         .at_with_attributes(&[".well-known", "edhoc"], &[], edhoc)

examples/coap/src/bin/coapserver.rs

@@ -32,17 +32,17 @@ fn main() {
             // This is an EDHOC message
             if request.message.payload[0] == 0xf5 {
                 let state = EdhocState::default();
-                let mut responder = EdhocResponder::new(state, &R, &CRED_R, Some(&CRED_I));
+                let responder = EdhocResponder::new(state, &R, &CRED_R, Some(&CRED_I));
 
-                let error = responder.process_message_1(
+                let result = responder.process_message_1(
                     &request.message.payload[1..]
                         .try_into()
                         .expect("wrong length"),
                 );
 
-                if error.is_ok() {
+                if let Ok(responder) = result {
                     let c_r = generate_connection_identifier_cbor();
-                    let message_2 = responder.prepare_message_2(c_r).unwrap();
+                    let (responder, message_2) = responder.prepare_message_2(c_r).unwrap();
                     response.message.payload = Vec::from(&message_2.content[..message_2.len]);
                     // save edhoc connection
                     edhoc_connections.push((c_r, responder));
@@ -51,24 +51,21 @@ fn main() {
                 // potentially message 3
                 println!("Received message 3");
                 let c_r_rcvd = request.message.payload[0];
-                let (index, mut responder, ec) = lookup_state(c_r_rcvd, edhoc_connections).unwrap();
-                edhoc_connections = ec;
+                // FIXME let's better not *panic here
+                let responder = take_state(c_r_rcvd, &mut edhoc_connections).unwrap();
 
                 println!("Found state with connection identifier {:?}", c_r_rcvd);
-                let prk_out = responder.process_message_3(
+                let result = responder.process_message_3(
                     &request.message.payload[1..]
                         .try_into()
                         .expect("wrong length"),
                 );
-
-                if prk_out.is_err() {
-                    println!("EDHOC processing error: {:?}", prk_out);
-                    // FIXME remove state from edhoc_connections
+                let Ok((mut responder, prk_out)) = result else {
+                    println!("EDHOC processing error: {:?}", response);
+                    // We don't get another chance, it's popped and can't be used any further
+                    // anyway legally
                     continue;
-                }
-
-                // update edhoc connection
-                edhoc_connections[index] = (c_r_rcvd, responder);
+                };
 
                 // send empty ack back
                 response.message.payload = b"".to_vec();
@@ -106,14 +103,16 @@ fn main() {
     }
 }
 
-fn lookup_state<'a>(
-    c_r_rcvd: u8,
-    edhoc_protocol_states: Vec<(u8, EdhocResponder<'a>)>,
-) -> Result<(usize, EdhocResponder<'a>, Vec<(u8, EdhocResponder)>), EDHOCError> {
+fn take_state<R>(c_r_rcvd: u8, edhoc_protocol_states: &mut Vec<(u8, R)>) -> Result<R, EDHOCError> {
     for (i, element) in edhoc_protocol_states.iter().enumerate() {
         let (c_r, responder) = element;
         if *c_r == c_r_rcvd {
-            return Ok((i, *responder, edhoc_protocol_states));
+            let max_index = edhoc_protocol_states.len() - 1;
+            edhoc_protocol_states.swap(i, max_index);
+            let Some((_c_r, responder)) = edhoc_protocol_states.pop() else {
+                unreachable!();
+            };
+            return Ok(responder);
         }
     }
     return Err(EDHOCError::WrongState);