acid/firmware/acid-firmware/src/ui/mod.rs

// #![cfg_attr(not(feature = "simulator"), no_main)]

use core::{
    cell::RefCell,
    ffi::CStr,
    iter::Chain,
    ops::{Deref, DerefMut, Range},
};

use alloc::{
    borrow::Cow,
    boxed::Box,
    ffi::CString,
    format,
    rc::Rc,
    string::{String, ToString},
    vec,
    vec::Vec,
};
use chrono::{DateTime, NaiveDateTime};
use ekv::{Database, MountError, flash::PageID};
use embassy_embedded_hal::{adapter::BlockingAsync, flash::partition::Partition};
use embassy_sync::blocking_mutex::raw::{CriticalSectionRawMutex, RawMutex};
use embassy_time::{Duration, Instant};
use embedded_storage_async::nor_flash::{NorFlash, ReadNorFlash};
use esp_hal::rng::Trng;
use esp_storage::FlashStorage;
use itertools::Itertools;
use log::{info, warn};
use password_hash::Key;
use rmk::futures::TryFutureExt;
use serde::{Deserialize, Serialize};
use serde_bytes::Bytes;
use slint::{ModelRc, SharedString, StandardListViewItem, VecModel};
use spectre_api_sys::{
    SpectreAlgorithm, SpectreCounter, SpectreKeyID, SpectreKeyPurpose, SpectreResultType,
    SpectreUserKey,
};

#[cfg(feature = "limit-fps")]
use crate::FRAME_DURATION_MIN;
use crate::{
    PSRAM_ALLOCATOR, SIGNAL_LCD_SUBMIT, SIGNAL_UI_RENDER,
    db::{
        AcidDatabase, DbKey, DbPathSpectreUserSite, DbPathSpectreUserSites, DbPathSpectreUsers,
        PartitionAcid, ReadTransactionExt,
    },
    ffi::alloc::__spre_free,
    ui::{
        backend::SlintBackend,
        messages::{
            CallbackMessage, CallbackMessageLogin, CallbackMessageUserEdit, CallbackMessageUsers,
        },
        storage::SpectreUsersConfig,
    },
    util::DurationExt,
};

pub mod backend;
pub mod messages;
pub mod storage;
pub mod window_adapter;

slint::include_modules!();

fn spectre_derive_user_key(username: &CStr, password: &CStr) -> SpectreUserKey {
    let user_key_start = Instant::now();

    unsafe {
        let user_key = &*spectre_api_sys::spectre_user_key(
            username.as_ptr(),
            password.as_ptr(),
            SpectreAlgorithm::Current,
        );
        let user_key_duration = Instant::now().duration_since(user_key_start);
        warn!(
            "User key derived in {} seconds:\n{user_key:02x?}",
            user_key_duration.display_as_secs()
        );
        let user_key_stack = user_key.clone();

        // TODO: Erase memory before freeing
        __spre_free(user_key as *const _ as *mut _);

        user_key_stack
    }
}

#[embassy_executor::task]
pub async fn run_renderer_task(backend: SlintBackend, flash_part_acid: PartitionAcid) {
    let db = AcidDatabase::mount(flash_part_acid).await;

    // let value = SpectreUsersConfig {
    //     users: vec![SpectreUserConfig {
    //         username: "test".to_string(),
    //         encrypted_key: [0; _],
    //         key_id: [0; _],
    //     }],
    // };

    // {
    //     let mut write = db.write_transaction().await;
    //     write
    //         .write(
    //             &DbKey::new(DbPathSpectreUserSite {
    //                 user_sites: DbPathSpectreUserSites {
    //                     username: "test".into(),
    //                 },
    //                 site: "example.org".into(),
    //             }),
    //             &postcard::to_allocvec(&SpectreSiteConfig::default()).unwrap(),
    //         )
    //         .await
    //         .unwrap();
    //     write
    //         .write(
    //             &DbKey::new(DbPathSpectreUserSite {
    //                 user_sites: DbPathSpectreUserSites {
    //                     username: "test".into(),
    //                 },
    //                 site: "sub.example.org".into(),
    //             }),
    //             &postcard::to_allocvec(&SpectreSiteConfig::default()).unwrap(),
    //         )
    //         .await
    //         .unwrap();
    //     write
    //         .write(
    //             &DbKey::new(DbPathSpectreUsers),
    //             &postcard::to_allocvec(&value).unwrap(),
    //         )
    //         .await
    //         .unwrap();
    //     write.commit().await.unwrap();
    // }

    // let read_value = {
    //     let read = db.read_transaction().await;
    //     // TODO: https://github.com/embassy-rs/ekv/issues/20
    //     let mut buffer = vec![0; 256];
    //     let slice = read
    //         .read_to_vec(&DbKey::new(DbPathSpectreUsers), &mut buffer)
    //         .await
    //         .unwrap();
    //     let read_value = postcard::from_bytes::<SpectreUsersConfig>(&slice).unwrap();

    //     let key_sites = DbKey::new(DbPathSpectreUserSites {
    //         username: "test".into(),
    //     });
    //     let mut key_buffer = [0_u8; ekv::config::MAX_KEY_SIZE];
    //     let mut cursor = read
    //         .read_range(key_sites.range_of_children())
    //         .await
    //         .unwrap();

    //     while let Some((key_len, value_len)) =
    //         cursor.next(&mut key_buffer, &mut buffer).await.unwrap()
    //     {
    //         let key = DbKey::from_raw(key_buffer[..key_len].into());
    //         let mut key_segments = key.segments();
    //         let value = &buffer[..value_len];
    //         let site_config = postcard::from_bytes::<SpectreSiteConfig>(&value).unwrap();
    //         let site = SpectreSite {
    //             config: site_config,
    //             username: key_segments.nth(2).unwrap().into(),
    //             site_name: key_segments.nth(1).unwrap().into(),
    //         };

    //         info!("site = {:#?}", site);
    //     }

    //     read_value
    // };

    // info!("read_value = {:#?}", read_value);
    // assert_eq!(value, read_value, "values do not match");

    // TODO:
    // * Store a config as a versioned postcard-serialized struct
    // * Store accounts and sites as ranges in the DB

    i_slint_core::properties::ALLOCATOR
        .set(&PSRAM_ALLOCATOR)
        .ok()
        .unwrap();
    slint::platform::set_platform(Box::new(backend)).expect("backend already initialized");

    let main = AppWindow::new().unwrap();
    let state = State::new(db, main).await;

    state.borrow().run_event_loop().await;
}

struct State {
    window: AppWindow,
    db: AcidDatabase,
    users: SpectreUsersConfig,
    /// Currently active view.
    view: AppState,
    // Retained state for each view.
    state_login: StateLogin,
    state_users: StateUsers,
    state_user_edit: StateUserEdit,
    state_user_sites: Option<StateUserSites>,
}

impl State {
    async fn new(db: AcidDatabase, main: AppWindow) -> Rc<RefCell<Self>> {
        let state = Rc::new(RefCell::new(State {
            window: main.clone_strong(),
            users: {
                let read = db.read_transaction().await;
                let mut buffer = vec![0_u8; 128];
                match read
                    .read_to_vec(&DbKey::new(DbPathSpectreUsers), &mut buffer)
                    .await
                {
                    Ok(bytes) => postcard::from_bytes::<SpectreUsersConfig>(&bytes).unwrap(),
                    Err(ekv::ReadError::KeyNotFound) => Default::default(),
                    Err(error) => panic!("Failed to read the users config: {error:?}"),
                }
            },
            db,
            view: AppState::Login,
            state_login: Default::default(),
            state_users: Default::default(),
            state_user_edit: Default::default(),
            state_user_sites: Default::default(),
        }));

        main.on_enter_view({
            let state = state.clone();
            move |view| {
                state.borrow_mut().set_view(view, true);
            }
        });

        main.on_escape({
            let state = state.clone();
            move || {
                state
                    .borrow_mut()
                    .process_callback_message(CallbackMessage::Escape);
            }
        });

        main.on_login_pw_accepted({
            let state = state.clone();
            move |username, password| {
                state
                    .borrow_mut()
                    .process_callback_message(CallbackMessage::Login(
                        CallbackMessageLogin::PwAccepted { username, password },
                    ));
            }
        });

        main.on_users_edit_user({
            let state = state.clone();
            move |username, new| {
                state
                    .borrow_mut()
                    .process_callback_message(CallbackMessage::Users(
                        CallbackMessageUsers::EditUser { username, new },
                    ));
            }
        });

        main.on_user_edit_compute_identicon({
            let state = state.clone();
            move |encrypted_key, password| {
                state
                    .borrow_mut()
                    .process_callback_message(CallbackMessage::UserEdit(
                        CallbackMessageUserEdit::ComputeIdenticon {
                            encrypted_key,
                            password,
                        },
                    ));
            }
        });

        main.on_user_edit_confirm({
            let state = state.clone();
            move |encrypted_key| {
                state
                    .borrow_mut()
                    .process_callback_message(CallbackMessage::UserEdit(
                        CallbackMessageUserEdit::Confirm { encrypted_key },
                    ));
            }
        });

        // let sites = Rc::new(VecModel::default());
        // sites.push("First".into());
        // sites.push("Second".into());
        // main.set_sites(ModelRc::new(ModelRc::new(sites.clone()).map(
        //     |mut site: StandardListViewItem| {
        //         site.text += "10";
        //         site
        //     },
        // )));

        state
    }

    fn process_callback_message(&mut self, message: CallbackMessage) {
        match self.view {
            AppState::Login => StateLogin::process_callback_message(self, message),
            AppState::Users => StateUsers::process_callback_message(self, message),
            AppState::UserEdit => StateUserEdit::process_callback_message(self, message),
            AppState::UserSites => StateUserSites::process_callback_message(self, message),
        }
    }

    fn set_view(&mut self, view: AppState, reset: bool) {
        self.view = view;
        self.window.set_app_state(view);

        if reset {
            match self.view {
                AppState::Login => self.state_login = Default::default(),
                AppState::Users => self.state_users = Default::default(),
                AppState::UserEdit => self.state_user_edit = Default::default(),
                AppState::UserSites => self.state_user_sites = Default::default(),
            }
        }
    }

    /// Instead of having a `loop` in the non-async `SlintBackend::run_event_loop`, we achieve
    /// async by having only one iteration of the loop run, and `await`ing here.
    /// The following block is analogous to `main.run()`.
    async fn run_event_loop(&self) -> ! {
        self.window.show().unwrap();

        loop {
            slint::run_event_loop().unwrap();
            SIGNAL_LCD_SUBMIT.signal(());
            #[cfg(feature = "limit-fps")]
            embassy_time::Timer::after(FRAME_DURATION_MIN).await;
            SIGNAL_UI_RENDER.wait().await;
        }

        #[expect(unreachable_code)]
        self.window.hide().unwrap();
    }
}

trait AppViewTrait {
    fn process_callback_message(state: &mut State, message: CallbackMessage) {}
}

#[derive(Default)]
struct StateLogin {}

impl AppViewTrait for StateLogin {
    fn process_callback_message(state: &mut State, message: CallbackMessage) {
        match message {
            CallbackMessage::Login(CallbackMessageLogin::PwAccepted { username, password }) => {
                let username_c =
                    CString::new(&*username).expect("Username cannot be converted to a C string.");
                let password_c =
                    CString::new(&*password).expect("Password cannot be converted to a C string.");
                let user_key = spectre_derive_user_key(&username_c, &password_c);

                // let site_key_start = Instant::now();
                // let site_key = &*spectre_api_sys::spectre_site_key(
                //     user_key as *const SpectreUserKey,
                //     c"example.org".as_ptr(),
                //     SpectreCounter::Initial,
                //     SpectreKeyPurpose::Authentication,
                //     c"".as_ptr(),
                // );
                // let site_key_duration = Instant::now().duration_since(site_key_start);
                // warn!(
                //     "Site key derived in {} seconds:\n{site_key:02x?}",
                //     site_key_duration.display_as_secs()
                // );
                // TODO: Erase memory before freeing
                // __spre_free(site_key as *const _ as *mut _);
                let Some(user) = state
                    .users
                    .users
                    .iter()
                    .find(|user| &*user.username == &*username)
                else {
                    return;
                };

                if user.key_id == user_key.keyID.bytes {
                    info!("Correct password entered for user {username:?}.");
                    state.state_user_sites = Some(StateUserSites { username, user_key });
                    state.set_view(AppState::UserSites, true);
                } else {
                    warn!("Incorrect password entered for user {username:?}.");
                    // TODO: Clear the input
                }
            }
            _ => (),
        }
    }
}

#[derive(Default)]
struct StateUsers {}

impl AppViewTrait for StateUsers {
    fn process_callback_message(state: &mut State, message: CallbackMessage) {
        match message {
            CallbackMessage::Escape => {
                state.set_view(AppState::Login, false);
            }
            CallbackMessage::Users(CallbackMessageUsers::EditUser { username, new }) => {
                state.state_user_edit = StateUserEdit {
                    username: username.clone(),
                    new,
                    hashed: None,
                };
                state.window.set_user_edit_username(username);
                state.set_view(AppState::UserEdit, false);
            }
            _ => (),
        }
    }
}

#[derive(Default)]
struct StateUserEdit {
    username: SharedString,
    new: bool,
    hashed: Option<ProposedPassword>,
}

impl AppViewTrait for StateUserEdit {
    fn process_callback_message(state: &mut State, message: CallbackMessage) {
        match message {
            CallbackMessage::Escape => {
                state.set_view(AppState::Users, false);
            }
            CallbackMessage::UserEdit(CallbackMessageUserEdit::ComputeIdenticon {
                encrypted_key,
                password,
            }) => {
                let username_c = CString::new(&*state.state_user_edit.username)
                    .expect("Username cannot be converted to a C string.");
                let password_c =
                    CString::new(&*password).expect("Password cannot be converted to a C string.");
                // let user_key = spectre_derive_user_key(&username_c, &password_c);
                let identicon: SharedString = unsafe {
                    let identicon = spectre_api_sys::spectre_identicon(
                        username_c.as_ptr(),
                        password_c.as_ptr(),
                    );
                    // TODO: identicon.color
                    format!(
                        "{}{}{}{}",
                        CStr::from_ptr(identicon.leftArm).to_str().unwrap(),
                        CStr::from_ptr(identicon.body).to_str().unwrap(),
                        CStr::from_ptr(identicon.rightArm).to_str().unwrap(),
                        CStr::from_ptr(identicon.accessory).to_str().unwrap()
                    )
                    .into()
                };

                warn!("Identicon: {identicon} ({identicon:?})");
                state.window.set_user_edit_identicon(identicon.clone());
                state.state_user_edit.hashed = Some(ProposedPassword {
                    encrypted_key,
                    identicon,
                })
            }
            CallbackMessage::UserEdit(CallbackMessageUserEdit::Confirm { encrypted_key }) => {
                // TODO
            }
            _ => (),
        }
    }
}

struct ProposedPassword {
    encrypted_key: SharedString,
    identicon: SharedString,
}

struct StateUserSites {
    username: SharedString,
    user_key: SpectreUserKey,
}

impl AppViewTrait for StateUserSites {}