Feat/combine uis (#2633)

* wip

* restructure backend for new ui structure

* new patchdb bootstrap, single websocket api, local storage migration, more

* update db websocket

* init apis

* update patch-db

* setup progress

* feat: implement state service, alert and routing

Signed-off-by: waterplea <alexander@inkin.ru>

* update setup wizard for new types

* feat: add init page

Signed-off-by: waterplea <alexander@inkin.ru>

* chore: refactor message, patch-db source stream and connection service

Signed-off-by: waterplea <alexander@inkin.ru>

* fix method not found on state

* fix backend bugs

* fix compat assets

* address comments

* remove unneeded styling

* cleaner progress

* bugfixes

* fix init logs

* fix progress reporting

* fix navigation by getting state after init

* remove patch dependency from live api

* fix caching

* re-add patchDB to live api

* fix metrics values

* send close frame

* add bootId and fix polling

---------

Signed-off-by: waterplea <alexander@inkin.ru>
Co-authored-by: Aiden McClelland <me@drbonez.dev>
Co-authored-by: waterplea <alexander@inkin.ru>

core/startos/src/context/cli.rs

@@ -18,7 +18,7 @@ use tracing::instrument;
 
 use super::setup::CURRENT_SECRET;
 use crate::context::config::{local_config_path, ClientConfig};
-use crate::context::{DiagnosticContext, InstallContext, RpcContext, SetupContext};
+use crate::context::{DiagnosticContext, InitContext, InstallContext, RpcContext, SetupContext};
 use crate::middleware::auth::LOCAL_AUTH_COOKIE_PATH;
 use crate::prelude::*;
 use crate::rpc_continuations::Guid;
@@ -271,6 +271,11 @@ impl CallRemote<DiagnosticContext> for CliContext {
         call_remote_http(&self.client, self.rpc_url.clone(), method, params).await
     }
 }
+impl CallRemote<InitContext> for CliContext {
+    async fn call_remote(&self, method: &str, params: Value, _: Empty) -> Result<Value, RpcError> {
+        call_remote_http(&self.client, self.rpc_url.clone(), method, params).await
+    }
+}
 impl CallRemote<SetupContext> for CliContext {
     async fn call_remote(&self, method: &str, params: Value, _: Empty) -> Result<Value, RpcError> {
         call_remote_http(&self.client, self.rpc_url.clone(), method, params).await

core/startos/src/context/diagnostic.rs

@@ -14,7 +14,7 @@ use crate::Error;
 
 pub struct DiagnosticContextSeed {
     pub datadir: PathBuf,
-    pub shutdown: Sender<Option<Shutdown>>,
+    pub shutdown: Sender<Shutdown>,
     pub error: Arc<RpcError>,
     pub disk_guid: Option<Arc<String>>,
     pub rpc_continuations: RpcContinuations,

core/startos/src/context/init.rs

@@ -0,0 +1,47 @@
+use std::ops::Deref;
+use std::sync::Arc;
+
+use rpc_toolkit::Context;
+use tokio::sync::broadcast::Sender;
+use tracing::instrument;
+
+use crate::context::config::ServerConfig;
+use crate::progress::FullProgressTracker;
+use crate::rpc_continuations::RpcContinuations;
+use crate::Error;
+
+pub struct InitContextSeed {
+    pub config: ServerConfig,
+    pub progress: FullProgressTracker,
+    pub shutdown: Sender<()>,
+    pub rpc_continuations: RpcContinuations,
+}
+
+#[derive(Clone)]
+pub struct InitContext(Arc<InitContextSeed>);
+impl InitContext {
+    #[instrument(skip_all)]
+    pub async fn init(cfg: &ServerConfig) -> Result<Self, Error> {
+        let (shutdown, _) = tokio::sync::broadcast::channel(1);
+        Ok(Self(Arc::new(InitContextSeed {
+            config: cfg.clone(),
+            progress: FullProgressTracker::new(),
+            shutdown,
+            rpc_continuations: RpcContinuations::new(),
+        })))
+    }
+}
+
+impl AsRef<RpcContinuations> for InitContext {
+    fn as_ref(&self) -> &RpcContinuations {
+        &self.rpc_continuations
+    }
+}
+
+impl Context for InitContext {}
+impl Deref for InitContext {
+    type Target = InitContextSeed;
+    fn deref(&self) -> &Self::Target {
+        &*self.0
+    }
+}

core/startos/src/context/install.rs

@@ -6,11 +6,13 @@ use tokio::sync::broadcast::Sender;
 use tracing::instrument;
 
 use crate::net::utils::find_eth_iface;
+use crate::rpc_continuations::RpcContinuations;
 use crate::Error;
 
 pub struct InstallContextSeed {
     pub ethernet_interface: String,
     pub shutdown: Sender<()>,
+    pub rpc_continuations: RpcContinuations,
 }
 
 #[derive(Clone)]
@@ -22,10 +24,17 @@ impl InstallContext {
         Ok(Self(Arc::new(InstallContextSeed {
             ethernet_interface: find_eth_iface().await?,
             shutdown,
+            rpc_continuations: RpcContinuations::new(),
         })))
     }
 }
 
+impl AsRef<RpcContinuations> for InstallContext {
+    fn as_ref(&self) -> &RpcContinuations {
+        &self.rpc_continuations
+    }
+}
+
 impl Context for InstallContext {}
 impl Deref for InstallContext {
     type Target = InstallContextSeed;

core/startos/src/context/mod.rs

@@ -1,12 +1,14 @@
 pub mod cli;
 pub mod config;
 pub mod diagnostic;
+pub mod init;
 pub mod install;
 pub mod rpc;
 pub mod setup;
 
 pub use cli::CliContext;
 pub use diagnostic::DiagnosticContext;
+pub use init::InitContext;
 pub use install::InstallContext;
 pub use rpc::RpcContext;
 pub use setup::SetupContext;

core/startos/src/context/rpc.rs

@@ -6,11 +6,12 @@ use std::sync::atomic::{AtomicBool, Ordering};
 use std::sync::Arc;
 use std::time::Duration;
 
+use imbl_value::InternedString;
 use josekit::jwk::Jwk;
 use reqwest::{Client, Proxy};
 use rpc_toolkit::yajrc::RpcError;
 use rpc_toolkit::{CallRemote, Context, Empty};
-use tokio::sync::{broadcast, oneshot, Mutex, RwLock};
+use tokio::sync::{broadcast, Mutex, RwLock};
 use tokio::time::Instant;
 use tracing::instrument;
 
@@ -22,12 +23,12 @@ use crate::dependencies::compute_dependency_config_errs;
 use crate::disk::OsPartitionInfo;
 use crate::init::check_time_is_synchronized;
 use crate::lxc::{ContainerId, LxcContainer, LxcManager};
-use crate::middleware::auth::HashSessionToken;
-use crate::net::net_controller::NetController;
+use crate::net::net_controller::{NetController, PreInitNetController};
 use crate::net::utils::{find_eth_iface, find_wifi_iface};
 use crate::net::wifi::WpaCli;
 use crate::prelude::*;
-use crate::rpc_continuations::RpcContinuations;
+use crate::progress::{FullProgressTracker, PhaseProgressTrackerHandle};
+use crate::rpc_continuations::{OpenAuthedContinuations, RpcContinuations};
 use crate::service::ServiceMap;
 use crate::shutdown::Shutdown;
 use crate::system::get_mem_info;
@@ -49,7 +50,7 @@ pub struct RpcContextSeed {
     pub shutdown: broadcast::Sender<Option<Shutdown>>,
     pub tor_socks: SocketAddr,
     pub lxc_manager: Arc<LxcManager>,
-    pub open_authed_websockets: Mutex<BTreeMap<HashSessionToken, Vec<oneshot::Sender<()>>>>,
+    pub open_authed_continuations: OpenAuthedContinuations<InternedString>,
     pub rpc_continuations: RpcContinuations,
     pub wifi_manager: Option<Arc<RwLock<WpaCli>>>,
     pub current_secret: Arc<Jwk>,
@@ -68,45 +69,103 @@ pub struct Hardware {
     pub ram: u64,
 }
 
+pub struct InitRpcContextPhases {
+    load_db: PhaseProgressTrackerHandle,
+    init_net_ctrl: PhaseProgressTrackerHandle,
+    read_device_info: PhaseProgressTrackerHandle,
+    cleanup_init: CleanupInitPhases,
+}
+impl InitRpcContextPhases {
+    pub fn new(handle: &FullProgressTracker) -> Self {
+        Self {
+            load_db: handle.add_phase("Loading database".into(), Some(5)),
+            init_net_ctrl: handle.add_phase("Initializing network".into(), Some(1)),
+            read_device_info: handle.add_phase("Reading device information".into(), Some(1)),
+            cleanup_init: CleanupInitPhases::new(handle),
+        }
+    }
+}
+
+pub struct CleanupInitPhases {
+    init_services: PhaseProgressTrackerHandle,
+    check_dependencies: PhaseProgressTrackerHandle,
+}
+impl CleanupInitPhases {
+    pub fn new(handle: &FullProgressTracker) -> Self {
+        Self {
+            init_services: handle.add_phase("Initializing services".into(), Some(10)),
+            check_dependencies: handle.add_phase("Checking dependencies".into(), Some(1)),
+        }
+    }
+}
+
 #[derive(Clone)]
 pub struct RpcContext(Arc<RpcContextSeed>);
 impl RpcContext {
     #[instrument(skip_all)]
-    pub async fn init(config: &ServerConfig, disk_guid: Arc<String>) -> Result<Self, Error> {
-        tracing::info!("Loaded Config");
+    pub async fn init(
+        config: &ServerConfig,
+        disk_guid: Arc<String>,
+        net_ctrl: Option<PreInitNetController>,
+        InitRpcContextPhases {
+            mut load_db,
+            mut init_net_ctrl,
+            mut read_device_info,
+            cleanup_init,
+        }: InitRpcContextPhases,
+    ) -> Result<Self, Error> {
         let tor_proxy = config.tor_socks.unwrap_or(SocketAddr::V4(SocketAddrV4::new(
             Ipv4Addr::new(127, 0, 0, 1),
             9050,
         )));
         let (shutdown, _) = tokio::sync::broadcast::channel(1);
 
-        let db = TypedPatchDb::<Database>::load(config.db().await?).await?;
+        load_db.start();
+        let db = if let Some(net_ctrl) = &net_ctrl {
+            net_ctrl.db.clone()
+        } else {
+            TypedPatchDb::<Database>::load(config.db().await?).await?
+        };
         let peek = db.peek().await;
         let account = AccountInfo::load(&peek)?;
+        load_db.complete();
         tracing::info!("Opened PatchDB");
+
+        init_net_ctrl.start();
         let net_controller = Arc::new(
             NetController::init(
-                db.clone(),
-                config
-                    .tor_control
-                    .unwrap_or(SocketAddr::from(([127, 0, 0, 1], 9051))),
-                tor_proxy,
+                if let Some(net_ctrl) = net_ctrl {
+                    net_ctrl
+                } else {
+                    PreInitNetController::init(
+                        db.clone(),
+                        config
+                            .tor_control
+                            .unwrap_or(SocketAddr::from(([127, 0, 0, 1], 9051))),
+                        tor_proxy,
+                        &account.hostname,
+                        account.tor_key.clone(),
+                    )
+                    .await?
+                },
                 config
                     .dns_bind
                     .as_deref()
                     .unwrap_or(&[SocketAddr::from(([127, 0, 0, 1], 53))]),
-                &account.hostname,
-                account.tor_key.clone(),
             )
             .await?,
         );
+        init_net_ctrl.complete();
         tracing::info!("Initialized Net Controller");
+
         let services = ServiceMap::default();
         let metrics_cache = RwLock::<Option<crate::system::Metrics>>::new(None);
-        tracing::info!("Initialized Notification Manager");
         let tor_proxy_url = format!("socks5h://{tor_proxy}");
+
+        read_device_info.start();
         let devices = lshw().await?;
         let ram = get_mem_info().await?.total.0 as u64 * 1024 * 1024;
+        read_device_info.complete();
 
         if !db
             .peek()
@@ -163,7 +222,7 @@ impl RpcContext {
             shutdown,
             tor_socks: tor_proxy,
             lxc_manager: Arc::new(LxcManager::new()),
-            open_authed_websockets: Mutex::new(BTreeMap::new()),
+            open_authed_continuations: OpenAuthedContinuations::new(),
             rpc_continuations: RpcContinuations::new(),
             wifi_manager: wifi_interface
                 .clone()
@@ -196,7 +255,7 @@ impl RpcContext {
         });
 
         let res = Self(seed.clone());
-        res.cleanup_and_initialize().await?;
+        res.cleanup_and_initialize(cleanup_init).await?;
         tracing::info!("Cleaned up transient states");
         Ok(res)
     }
@@ -210,11 +269,18 @@ impl RpcContext {
         Ok(())
     }
 
-    #[instrument(skip(self))]
-    pub async fn cleanup_and_initialize(&self) -> Result<(), Error> {
-        self.services.init(&self).await?;
+    #[instrument(skip_all)]
+    pub async fn cleanup_and_initialize(
+        &self,
+        CleanupInitPhases {
+            init_services,
+            mut check_dependencies,
+        }: CleanupInitPhases,
+    ) -> Result<(), Error> {
+        self.services.init(&self, init_services).await?;
         tracing::info!("Initialized Package Managers");
 
+        check_dependencies.start();
         let mut updated_current_dependents = BTreeMap::new();
         let peek = self.db.peek().await;
         for (package_id, package) in peek.as_public().as_package_data().as_entries()?.into_iter() {
@@ -238,6 +304,7 @@ impl RpcContext {
                 Ok(())
             })
             .await?;
+        check_dependencies.complete();
 
         Ok(())
     }
@@ -274,6 +341,11 @@ impl AsRef<RpcContinuations> for RpcContext {
         &self.rpc_continuations
     }
 }
+impl AsRef<OpenAuthedContinuations<InternedString>> for RpcContext {
+    fn as_ref(&self) -> &OpenAuthedContinuations<InternedString> {
+        &self.open_authed_continuations
+    }
+}
 impl Context for RpcContext {}
 impl Deref for RpcContext {
     type Target = RpcContextSeed;

core/startos/src/context/setup.rs

@@ -1,23 +1,31 @@
 use std::ops::Deref;
 use std::path::PathBuf;
 use std::sync::Arc;
+use std::time::Duration;
 
+use futures::{Future, StreamExt};
+use helpers::NonDetachingJoinHandle;
 use josekit::jwk::Jwk;
 use patch_db::PatchDb;
-use rpc_toolkit::yajrc::RpcError;
 use rpc_toolkit::Context;
 use serde::{Deserialize, Serialize};
 use sqlx::postgres::PgConnectOptions;
 use sqlx::PgPool;
 use tokio::sync::broadcast::Sender;
-use tokio::sync::RwLock;
+use tokio::sync::OnceCell;
 use tracing::instrument;
+use ts_rs::TS;
 
+use crate::account::AccountInfo;
 use crate::context::config::ServerConfig;
+use crate::context::RpcContext;
 use crate::disk::OsPartitionInfo;
 use crate::init::init_postgres;
 use crate::prelude::*;
-use crate::setup::SetupStatus;
+use crate::progress::FullProgressTracker;
+use crate::rpc_continuations::{Guid, RpcContinuation, RpcContinuations};
+use crate::setup::SetupProgress;
+use crate::util::net::WebSocketExt;
 
 lazy_static::lazy_static! {
     pub static ref CURRENT_SECRET: Jwk = Jwk::generate_ec_key(josekit::jwk::alg::ec::EcCurve::P256).unwrap_or_else(|e| {
@@ -27,30 +35,35 @@ lazy_static::lazy_static! {
     });
 }
 
-#[derive(Clone, Serialize, Deserialize)]
+#[derive(Debug, Clone, Deserialize, Serialize, TS)]
 #[serde(rename_all = "camelCase")]
+#[ts(export)]
 pub struct SetupResult {
     pub tor_address: String,
     pub lan_address: String,
     pub root_ca: String,
 }
+impl TryFrom<&AccountInfo> for SetupResult {
+    type Error = Error;
+    fn try_from(value: &AccountInfo) -> Result<Self, Self::Error> {
+        Ok(Self {
+            tor_address: format!("https://{}", value.tor_key.public().get_onion_address()),
+            lan_address: value.hostname.lan_address(),
+            root_ca: String::from_utf8(value.root_ca_cert.to_pem()?)?,
+        })
+    }
+}
 
 pub struct SetupContextSeed {
     pub config: ServerConfig,
     pub os_partitions: OsPartitionInfo,
     pub disable_encryption: bool,
+    pub progress: FullProgressTracker,
+    pub task: OnceCell<NonDetachingJoinHandle<()>>,
+    pub result: OnceCell<Result<(SetupResult, RpcContext), Error>>,
     pub shutdown: Sender<()>,
     pub datadir: PathBuf,
-    pub selected_v2_drive: RwLock<Option<PathBuf>>,
-    pub cached_product_key: RwLock<Option<Arc<String>>>,
-    pub setup_status: RwLock<Option<Result<SetupStatus, RpcError>>>,
-    pub setup_result: RwLock<Option<(Arc<String>, SetupResult)>>,
-}
-
-impl AsRef<Jwk> for SetupContextSeed {
-    fn as_ref(&self) -> &Jwk {
-        &*CURRENT_SECRET
-    }
+    pub rpc_continuations: RpcContinuations,
 }
 
 #[derive(Clone)]
@@ -69,12 +82,12 @@ impl SetupContext {
                 )
             })?,
             disable_encryption: config.disable_encryption.unwrap_or(false),
+            progress: FullProgressTracker::new(),
+            task: OnceCell::new(),
+            result: OnceCell::new(),
             shutdown,
             datadir,
-            selected_v2_drive: RwLock::new(None),
-            cached_product_key: RwLock::new(None),
-            setup_status: RwLock::new(None),
-            setup_result: RwLock::new(None),
+            rpc_continuations: RpcContinuations::new(),
         })))
     }
     #[instrument(skip_all)]
@@ -97,6 +110,104 @@ impl SetupContext {
             .with_kind(crate::ErrorKind::Database)?;
         Ok(secret_store)
     }
+
+    pub fn run_setup<F, Fut>(&self, f: F) -> Result<(), Error>
+    where
+        F: FnOnce() -> Fut + Send + 'static,
+        Fut: Future<Output = Result<(SetupResult, RpcContext), Error>> + Send,
+    {
+        let local_ctx = self.clone();
+        self.task
+            .set(
+                tokio::spawn(async move {
+                    local_ctx
+                        .result
+                        .get_or_init(|| async {
+                            match f().await {
+                                Ok(res) => {
+                                    tracing::info!("Setup complete!");
+                                    Ok(res)
+                                }
+                                Err(e) => {
+                                    tracing::error!("Setup failed: {e}");
+                                    tracing::debug!("{e:?}");
+                                    Err(e)
+                                }
+                            }
+                        })
+                        .await;
+                    local_ctx.progress.complete();
+                })
+                .into(),
+            )
+            .map_err(|_| {
+                if self.result.initialized() {
+                    Error::new(eyre!("Setup already complete"), ErrorKind::InvalidRequest)
+                } else {
+                    Error::new(
+                        eyre!("Setup already in progress"),
+                        ErrorKind::InvalidRequest,
+                    )
+                }
+            })?;
+        Ok(())
+    }
+
+    pub async fn progress(&self) -> SetupProgress {
+        use axum::extract::ws;
+
+        let guid = Guid::new();
+        let progress_tracker = self.progress.clone();
+        let progress = progress_tracker.snapshot();
+        self.rpc_continuations
+            .add(
+                guid.clone(),
+                RpcContinuation::ws(
+                    |mut ws| async move {
+                        if let Err(e) = async {
+                            let mut stream =
+                                progress_tracker.stream(Some(Duration::from_millis(100)));
+                            while let Some(progress) = stream.next().await {
+                                ws.send(ws::Message::Text(
+                                    serde_json::to_string(&progress)
+                                        .with_kind(ErrorKind::Serialization)?,
+                                ))
+                                .await
+                                .with_kind(ErrorKind::Network)?;
+                                if progress.overall.is_complete() {
+                                    break;
+                                }
+                            }
+
+                            ws.normal_close("complete").await?;
+
+                            Ok::<_, Error>(())
+                        }
+                        .await
+                        {
+                            tracing::error!("Error in setup progress websocket: {e}");
+                            tracing::debug!("{e:?}");
+                        }
+                    },
+                    Duration::from_secs(30),
+                ),
+            )
+            .await;
+
+        SetupProgress { progress, guid }
+    }
+}
+
+impl AsRef<Jwk> for SetupContext {
+    fn as_ref(&self) -> &Jwk {
+        &*CURRENT_SECRET
+    }
+}
+
+impl AsRef<RpcContinuations> for SetupContext {
+    fn as_ref(&self) -> &RpcContinuations {
+        &self.rpc_continuations
+    }
 }
 
 impl Context for SetupContext {}