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::{
    SIGNAL_LCD_SUBMIT, SIGNAL_UI_RENDER,
    db::{
        AcidDatabase, DbKey, DbPathSpectreUserSite, DbPathSpectreUserSites, DbPathSpectreUsers,
        PartitionAcid, ReadTransactionExt,
    },
    ffi::alloc::__spre_free,
    ui::backend::SlintBackend,
    util::DurationExt,
};

pub mod backend;
pub mod window_adapter;

slint::include_modules!();

#[derive(Deserialize, Serialize, Default, Clone, Debug, PartialEq, Eq)]
struct SpectreUsersConfig {
    users: Vec<SpectreUserConfig>,
}

#[derive(Deserialize, Serialize, Clone, Debug, PartialEq, Eq)]
struct SpectreUserConfig {
    username: String,
    #[serde(with = "serde_bytes")]
    encrypted_key: Key,
    #[serde(with = "serde_bytes")]
    key_id: [u8; 32],
}

trait ReprConvert: Copy {
    type Repr: Copy;

    fn into_repr(self) -> Self::Repr;
    fn from_repr(repr: Self::Repr) -> Self;
}

impl ReprConvert for SpectreAlgorithm {
    type Repr = u32;

    fn from_repr(repr: Self::Repr) -> Self {
        unsafe { core::mem::transmute::<Self::Repr, Self>(repr) }
    }

    fn into_repr(self) -> Self::Repr {
        self as Self::Repr
    }
}

impl ReprConvert for SpectreResultType {
    type Repr = u32;

    fn from_repr(repr: Self::Repr) -> Self {
        unsafe { core::mem::transmute::<Self::Repr, Self>(repr) }
    }

    fn into_repr(self) -> Self::Repr {
        self as Self::Repr
    }
}

mod with_repr {
    use serde::{Deserialize, Deserializer, Serializer};

    use super::ReprConvert;

    pub fn serialize<T, S>(value: &T, serializer: S) -> Result<S::Ok, S::Error>
    where
        T: ReprConvert<Repr = u32>,
        S: Serializer,
    {
        serializer.serialize_u32(value.into_repr())
    }

    pub fn deserialize<'de, T, D>(deserializer: D) -> Result<T, D::Error>
    where
        T: ReprConvert<Repr = u32>,
        D: Deserializer<'de>,
    {
        <u32 as Deserialize>::deserialize(deserializer).map(T::from_repr)
    }
}

#[derive(Deserialize, Serialize, Clone, Debug, PartialEq, Eq)]
struct SpectreSiteConfig {
    #[serde(with = "with_repr")]
    algorithm: SpectreAlgorithm,
    counter: SpectreCounter::Type,
    #[serde(rename = "type")]
    #[serde(with = "with_repr")]
    result_type: SpectreResultType,
    password: Option<String>,
    #[serde(with = "with_repr")]
    login_type: SpectreResultType,
    login_name: Option<String>,
    uses: u32,
    last_used: NaiveDateTime,
}

impl Default for SpectreSiteConfig {
    fn default() -> Self {
        Self {
            algorithm: SpectreAlgorithm::Current,
            counter: SpectreCounter::Default,
            result_type: SpectreResultType::SpectreResultDefaultResult,
            password: None,
            login_type: SpectreResultType::SpectreResultDefaultLogin,
            login_name: None,
            uses: 0,
            last_used: Default::default(),
        }
    }
}

#[derive(Clone, Debug, PartialEq, Eq)]
struct SpectreSite {
    username: String,
    site_name: String,
    config: SpectreSiteConfig,
}

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

    slint::platform::set_platform(Box::new(backend)).expect("backend already initialized");

    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 {
        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(),
                }
            }
        }
    }

    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()
                    };

                    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 {}

    enum CallbackMessage {
        /// The escape key was pressed.
        Escape,
        Login(CallbackMessageLogin),
        Users(CallbackMessageUsers),
        UserEdit(CallbackMessageUserEdit),
        UserSites(CallbackMessageUserSites),
    }

    enum CallbackMessageLogin {
        PwAccepted {
            username: SharedString,
            password: SharedString,
        },
    }

    enum CallbackMessageUsers {
        EditUser { username: SharedString, new: bool },
    }

    enum CallbackMessageUserEdit {
        ComputeIdenticon {
            encrypted_key: SharedString,
            password: SharedString,
        },
        Confirm {
            encrypted_key: SharedString,
        },
    }

    enum CallbackMessageUserSites {}

    let main = AppWindow::new().unwrap();

    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
    //     },
    // )));

    run_event_loop(main).await;
}

/// 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(main: AppWindow) {
    main.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)]
    main.hide().unwrap();
}