From 5758b1d4c02ec07e246c7bf9baa43a508bfa2ff2 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Gianguido=20Sor=C3=A0?= <me@gsora.xyz>
Date: Fri, 12 Aug 2022 16:54:26 +0200
Subject: [PATCH] feat(vault): handle usb restore for >7601 bytes of payload

---
 apps/vault/src/main.rs | 44 ++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 42 insertions(+), 2 deletions(-)

diff --git a/apps/vault/src/main.rs b/apps/vault/src/main.rs
index 9e26479db..d5b56e011 100644
--- a/apps/vault/src/main.rs
+++ b/apps/vault/src/main.rs
@@ -30,6 +30,7 @@ use framework::ListItem;
 use actions::{ActionOp, start_actions_thread};
 
 use crate::prereqs::ntp_updater;
+use crate::vendor_commands::VendorSession;
 
 // CTAP2 testing notes:
 // run our branch and use this to forward the prompts on to the device:
@@ -184,6 +185,7 @@ fn main() -> ! {
     start_actions_thread(conn, actions_sid, mode.clone(), item_list.clone(), action_active.clone());
     let actions_conn = xous::connect(actions_sid).unwrap();
 
+
     // spawn the FIDO2 USB handler
     let _ = thread::spawn({
         move || {
@@ -191,6 +193,8 @@ fn main() -> ! {
             let tt = ticktimer_server::Ticktimer::new().unwrap();
             let boot_time = ClockValue::new(tt.elapsed_ms() as i64, 1000);
 
+            let mut vendor_session = VendorSession::default();
+
             let mut rng = ctap_crypto::rng256::XousRng256::new(&xns);
             // this call will block until the PDDB is mounted.
             let usb = usb_device_xous::UsbHid::new();
@@ -206,7 +210,43 @@ fn main() -> ! {
                             let reply = match ctap_hid.process_hid_packet(&msg.packet, now, &mut ctap_state) {
                                 ctap::hid::send::CTAPHIDResponse::StandardCommand(iter) => iter,
                                 ctap::hid::send::CTAPHIDResponse::VendorCommand(cmd, cid, payload) => {
-                                    vendor_commands::handle_vendor_command(cmd, cid, payload)
+                                    match vendor_commands::handle_vendor_data(cmd, cid, payload, &mut vendor_session) {
+                                        Ok(return_payload) => {
+                                            // if None, this means we've finished parsing all that
+                                            // was needed, and we handle/respond with real data
+
+                                            match return_payload {
+                                                Some(data) => data,
+                                                None => {
+                                                    log::debug!("starting processing of vendor data...");
+                                                    let resp = vendor_commands::handle_vendor_command(&mut vendor_session);
+                                                    log::debug!("finished processing of vendor data!");
+                                                    
+                                                    match vendor_session.is_backup() {
+                                                        true => {
+                                                            if vendor_session.has_backup_data() {
+                                                                resp
+                                                            } else {
+                                                                vendor_session = VendorSession::default();
+                                                                resp
+                                                            }
+                                                        },
+                                                        false => {
+                                                            vendor_session = VendorSession::default();
+                                                            resp
+                                                        }
+                                                    }
+                                                }
+                                            }
+                                        }
+                                        Err(session_error) => {
+                                            // reset the session
+                                            vendor_session = VendorSession::default();
+
+                                            session_error.ctaphid_error(cid)
+                                        }
+                                    }
+                                    //vendor_commands::handle_vendor_command(cmd, cid, payload)
                                 }
                             };
                             // This block handles sending packets.
@@ -488,4 +528,4 @@ pub(crate) fn basis_change() {
     xous::send_message(SELF_CONN.load(Ordering::SeqCst),
         Message::new_scalar(VaultOp::BasisChange.to_usize().unwrap(), 0, 0, 0, 0)
     ).unwrap();
-}
\ No newline at end of file
+}