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