From 0c0cea7f6f56d1362278a8c4baf8d8971e91f069 Mon Sep 17 00:00:00 2001
From: Evan Feenstra <evanfeenstra@gmail.com>
Date: Tue, 27 Jun 2023 15:23:25 -0700
Subject: [PATCH] loops in init_lss to wait until correct message type and
 format

---
 sphinx-key/src/core/lss.rs | 41 ++++++++++++++++++++++++++++----------
 1 file changed, 31 insertions(+), 10 deletions(-)

diff --git a/sphinx-key/src/core/lss.rs b/sphinx-key/src/core/lss.rs
index 01baa24..0e2d1f4 100644
--- a/sphinx-key/src/core/lss.rs
+++ b/sphinx-key/src/core/lss.rs
@@ -5,7 +5,7 @@ use embedded_svc::mqtt::client::MessageImpl;
 use embedded_svc::utils::mqtt::client::ConnState;
 use esp_idf_svc::mqtt::client::EspMqttClient;
 use esp_idf_sys::EspError;
-use lss_connector::{secp256k1::PublicKey, LssSigner, Msg as LssMsg};
+use lss_connector::{secp256k1::PublicKey, BrokerMutations, LssSigner, Msg as LssMsg};
 use sphinx_signer::sphinx_glyph::topics;
 use sphinx_signer::{self, RootHandler, RootHandlerBuilder};
 use std::sync::mpsc;
@@ -18,26 +18,47 @@ pub fn init_lss(
     handler_builder: RootHandlerBuilder,
     mqtt: &mut EspMqttClient<ConnState<MessageImpl, EspError>>,
 ) -> Result<(RootHandler, LssSigner)> {
-    let first_lss_msg = match rx.recv()? {
-        Event::LssMessage(b) => b,
-        _ => return Err(anyhow!("not a lss msg")),
+    let server_pubkey = loop {
+        let event = rx.recv()?;
+        match server_pubkey_from_event(event) {
+            Ok(spk) => break spk,
+            Err(e) => log::warn!("failed server_pubkey_from_event {:?}", e),
+        }
     };
-    let init = LssMsg::from_slice(&first_lss_msg)?.as_init()?;
-    let server_pubkey = PublicKey::from_slice(&init.server_pubkey)?;
 
     let (lss_signer, res1) = LssSigner::new(&handler_builder, &server_pubkey);
     let lss_res_topic = format!("{}/{}", client_id, topics::INIT_RES);
     mqtt.publish(&lss_res_topic, QOS, false, &res1)
         .expect("could not publish LSS response");
 
-    let second_lss_msg = match rx.recv()? {
-        Event::LssMessage(b) => b,
-        _ => return Err(anyhow!("not a lss msg")),
+    let created = loop {
+        let event = rx.recv()?;
+        match created_from_event(event) {
+            Ok(c) => break c,
+            Err(e) => log::warn!("failed created_from_event {:?}", e),
+        }
     };
-    let created = LssMsg::from_slice(&second_lss_msg)?.as_created()?;
+
     let (root_handler, res2) = lss_signer.build_with_lss(created, handler_builder)?;
     mqtt.publish(&lss_res_topic, QOS, false, &res2)
         .expect("could not publish LSS response 2");
 
     Ok((root_handler, lss_signer))
 }
+
+fn server_pubkey_from_event(event: Event) -> anyhow::Result<PublicKey> {
+    match event {
+        Event::LssMessage(b) => {
+            let init = LssMsg::from_slice(&b)?.as_init()?;
+            let server_pubkey = PublicKey::from_slice(&init.server_pubkey)?;
+            Ok(server_pubkey)
+        }
+        _ => Err(anyhow!("not an LSS msg")),
+    }
+}
+fn created_from_event(event: Event) -> anyhow::Result<BrokerMutations> {
+    match event {
+        Event::LssMessage(b) => Ok(LssMsg::from_slice(&b)?.as_created()?),
+        _ => Err(anyhow!("not an LSS msg")),
+    }
+}