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

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

use core::{cell::RefCell, ffi::CStr, ops::DerefMut};

use alloc::{boxed::Box, ffi::CString, format, rc::Rc, vec};
use embassy_time::Instant;
use hex::FromHexError;
use log::{error, info, warn};
use password_hash::Key;
use slint::{Model, ModelExt, ModelRc, SharedString, VecModel};
use spectre_api_sys::{SpectreAlgorithm, SpectreKeyID, SpectreUserKey};

#[cfg(feature = "limit-fps")]
use crate::FRAME_DURATION_MIN;
use crate::{
    PSRAM_ALLOCATOR, SIGNAL_LCD_SUBMIT, SIGNAL_UI_RENDER,
    db::{AcidDatabase, DbKey, DbPathSpectreUsers, PartitionAcid, ReadTransactionExt},
    ffi::{alloc::__spre_free, crypto::ACTIVE_ENCRYPTED_USER_KEY},
    ui::{
        backend::SlintBackend,
        messages::{
            CallbackMessage, CallbackMessageLogin, CallbackMessageUserEdit, CallbackMessageUsers,
        },
        storage::{SpectreUserConfig, 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,
    encrypted_key: Option<Key>,
) -> SpectreUserKey {
    if let Some(encrypted_key) = encrypted_key {
        critical_section::with(|cs| {
            ACTIVE_ENCRYPTED_USER_KEY.borrow(cs).set(encrypted_key);
        });
    }

    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;

        // 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;
    let window = state.borrow().window.clone_strong();

    State::run_event_loop(window).await;
}

struct State {
    window: AppWindow,
    db: Rc<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 users = {
            let users = Self::load_users(&db).await;
            warn!("Users: {users:#?}");
            users
        };
        let usernames = users.users.clone().map(|user| user.username);

        let state = Rc::new(RefCell::new(State {
            window: main.clone_strong(),
            users,
            db: Rc::new(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, true);
            }
        });

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

        main.set_login_usernames(ModelRc::new(usernames));

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

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

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

        main.on_user_edit_compute_key_id({
            let state = state.clone();
            move |key| {
                State::process_callback_message(
                    &state,
                    CallbackMessage::UserEdit(CallbackMessageUserEdit::ComputeKeyId { key }),
                );
            }
        });

        main.on_user_edit_confirm({
            let state = state.clone();
            move |encrypted_key| {
                State::process_callback_message(
                    &state,
                    CallbackMessage::UserEdit(CallbackMessageUserEdit::ConfirmRequest {
                        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
    }

    async fn load_users(db: &AcidDatabase) -> SpectreUsersConfig {
        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:?}"),
        }
    }

    async fn save_users(db: &AcidDatabase, users: &SpectreUsersConfig) {
        let mut write = db.write_transaction().await;
        let buffer = postcard::to_allocvec(&users).unwrap();
        write
            .write(&DbKey::new(DbPathSpectreUsers), &buffer)
            .await
            .unwrap();
        write.commit().await.unwrap();
    }

    fn process_callback_message(state_rc: &Rc<RefCell<State>>, message: CallbackMessage) {
        let view = state_rc.borrow().view;
        match view {
            AppState::Login => StateLogin::process_callback_message(state_rc, message),
            AppState::Users => StateUsers::process_callback_message(state_rc, message),
            AppState::UserEdit => StateUserEdit::process_callback_message(state_rc, message),
            AppState::UserSites => StateUserSites::process_callback_message(state_rc, message),
        }
    }

    fn reset_view(&mut self) {
        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(),
        }
    }

    fn set_view(&mut self, view: AppState, reset_source: bool, reset_target: bool) {
        if reset_source {
            self.reset_view();
        }

        self.view = view;
        self.window.set_app_state(view);

        if reset_target {
            self.reset_view();
        }
    }

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

trait AppViewTrait {
    fn process_callback_message(_state_rc: &Rc<RefCell<State>>, _message: CallbackMessage) {}
}

#[derive(Default)]
struct StateLogin {}

impl AppViewTrait for StateLogin {
    fn process_callback_message(state_rc: &Rc<RefCell<State>>, message: CallbackMessage) {
        let mut state = state_rc.borrow_mut();
        if let CallbackMessage::Login(CallbackMessageLogin::PwAccepted {
            user_index,
            username: _,
            password,
        }) = message
        {
            let Some(user) = state.users.users.row_data(user_index as usize) else {
                error!("Failed to find a user with index {user_index}.");
                return;
            };
            let username_c =
                CString::new(&*user.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, Some(user.encrypted_key));

            // 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 _);

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

#[derive(Default)]
struct StateUsers {}

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

#[derive(Default)]
struct StateUserEdit {
    username: SharedString,
    new: bool,
    password: Option<SharedString>,
    encrypted_key: Option<(Key, SpectreUserKey)>,
}

impl AppViewTrait for StateUserEdit {
    fn process_callback_message(state_rc: &Rc<RefCell<State>>, message: CallbackMessage) {
        let state = state_rc.clone();
        let mut state = state.borrow_mut();
        match message {
            CallbackMessage::Escape => {
                state.set_view(AppState::Users, true, false);
            }
            CallbackMessage::UserEdit(CallbackMessageUserEdit::ComputeIdenticon { 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.password = Some(password);
            }
            CallbackMessage::UserEdit(CallbackMessageUserEdit::ComputeKeyId {
                key: key_string,
            }) => {
                let Some(password) = state.state_user_edit.password.as_ref() else {
                    warn!("Attempted to compute a key ID when no password has been entered.");
                    return;
                };

                let mut key: Key = [0; _];
                if let Err(key_decode_error) = hex::decode_to_slice(&*key_string, &mut key) {
                    let message = match key_decode_error {
                        FromHexError::InvalidStringLength | FromHexError::OddLength => {
                            let required_size = key.len() * 2;
                            let provided_size = key_string.len();
                            let delta = provided_size as i32 - required_size as i32;

                            if delta < 0 {
                                slint::format!("Missing {} characters.", -delta)
                            } else if delta > 0 {
                                slint::format!("{} too many characters.", delta)
                            } else {
                                slint::format!("Invalid key length.")
                            }
                        }
                        FromHexError::InvalidHexCharacter { c, index } => {
                            slint::format!("Invalid character {c:?} at position {index}.")
                        }
                    };
                    state.window.set_user_edit_key_error(message);
                    return;
                }

                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, Some(key));

                state.window.set_user_edit_key_error(SharedString::new());
                state.window.set_user_edit_key_id(
                    CStr::from_bytes_with_nul(&user_key.keyID.hex)
                        .unwrap()
                        .to_str()
                        .unwrap()
                        .into(),
                );
                state.state_user_edit.encrypted_key = Some((key, user_key));
            }
            CallbackMessage::UserEdit(CallbackMessageUserEdit::ConfirmRequest {
                encrypted_key: _,
            }) => {
                let Some((
                    encrypted_key,
                    SpectreUserKey {
                        keyID: SpectreKeyID { bytes: key_id, .. },
                        ..
                    },
                )) = state.state_user_edit.encrypted_key.take()
                else {
                    warn!("Encrypted key is not set.");
                    return;
                };

                // If a user with that username already exists, overwrite it.
                let user = SpectreUserConfig {
                    username: state.state_user_edit.username.clone(),
                    encrypted_key,
                    key_id,
                };

                let mut existing_index = None;

                for index in 0..state.users.users.row_count() {
                    if let Some(current_user) = state.users.users.row_data(index) {
                        if current_user.username == user.username {
                            existing_index = Some(index);
                        }
                    }
                }

                if let Some(existing_index) = existing_index {
                    state.users.users.set_row_data(existing_index, user);
                } else {
                    state.users.users.push(user);
                }

                slint::spawn_local({
                    let state_rc = state_rc.clone();
                    let db = state.db.clone();
                    let users = state.users.clone();
                    async move {
                        State::save_users(&db, &users).await;
                        State::process_callback_message(
                            &state_rc,
                            CallbackMessage::UserEdit(CallbackMessageUserEdit::ConfirmProcessed),
                        );
                    }
                })
                .unwrap();
            }
            CallbackMessage::UserEdit(CallbackMessageUserEdit::ConfirmProcessed) => {
                state.state_user_edit = Default::default();
                state.set_view(AppState::Users, true, true);
            }
            _ => (),
        }
    }
}

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

impl AppViewTrait for StateUserSites {}