[WIP] MTP migration

src/main.rs

@@ -9,23 +9,32 @@ use crate::transport::omikron_connection;
 use crate::util::file_util::get_directory;
 use crate::util::logger::PrintType;
 use crate::util::logger::startup;
-use base64::Engine as _;
 use dotenv::from_path;
-use mtp_crypto::{Keyring, PublicKeyBundle};
+use mtp::files::{load_keyring as load_keyring_file, save_keyring, save_public_key_bundle};
+use mtp_crypto::Keyring;
 use once_cell::sync::Lazy;
 use rustls::crypto::aws_lc_rs::default_provider;
 use std::env;
 use std::path::Path;
 
-static KEYRING_ENV: Lazy<String> = Lazy::new(|| env::var("KEYRING").unwrap());
-fn load_keyring() -> Keyring {
-    let bytes = base64::engine::general_purpose::STANDARD
-        .decode(&*KEYRING_ENV)
-        .expect("Invalid KEYRING env var: not valid base64");
-    Keyring::from_bytes(&bytes).expect("Invalid KEYRING env var: failed to deserialize")
+const KEYRING_PATH: &str = "./omega.mk";
+
+static KEYRING: Lazy<Keyring> = Lazy::new(|| {
+    load_keyring_file(KEYRING_PATH).unwrap_or_else(|_| {
+        let kr = Keyring::generate();
+        save_keyring(&kr, KEYRING_PATH).expect("Failed to save generated keyring");
+        save_public_key_bundle(&kr.public_key_bundle(), KEYRING_PATH)
+            .expect("Failed to save generated public key bundle");
+        eprintln!("Generated new keyring at {}", KEYRING_PATH);
+        kr
+    })
+});
+
+pub fn get_keyring() -> &'static Keyring {
+    &KEYRING
 }
-pub fn get_public_key_bundle() -> PublicKeyBundle {
-    load_keyring().public_key_bundle()
+pub fn load_keyring() -> Keyring {
+    Keyring::from_bytes(&KEYRING.to_bytes()).unwrap()
 }
 
 #[tokio::main]

src/server/api.rs

@@ -1,4 +1,4 @@
-use crate::get_public_key_bundle;
+use crate::load_keyring;
 use crate::sql::sql;
 use crate::sql::sql::{get_by_user_id, get_omikron_by_id};
 use crate::sql::user_online_tracker::get_iota_primary_omikron_connection;
@@ -192,7 +192,7 @@ pub async fn handle(path: &str, body_string: Option<String>) -> HttpResponse {
         ["api", "get", "public_key"] => {
             let mut res = JsonValue::new_object();
             res["status"] = "success".into();
-            let bundle = get_public_key_bundle();
+            let bundle = load_keyring().public_key_bundle();
             res["public_key"] = base64::engine::general_purpose::STANDARD
                 .encode(bundle.as_bytes())
                 .into();

src/transport/omikron_connection.rs

@@ -71,8 +71,6 @@ pub struct WaitingTask {
 pub struct OmikronConnection {
     id: u64,
     sender: Mutex<Option<Sender>>,
-    challenge: RwLock<String>,
-    pub_key: RwLock<Option<Vec<u8>>>,
     pub ping: RwLock<i64>,
     waiting_tasks: DashMap<u32, WaitingTask>,
     cleanup_handle: std::sync::Mutex<Option<tokio::task::JoinHandle<()>>>,
@@ -91,12 +89,10 @@ impl OmikronConnection {
     // Construction
     // -------------------------------------------------------------------------
 
-    pub fn new(sender: Sender) -> Arc<Self> {
+    pub fn new(sender: Sender, id: u64) -> Arc<Self> {
         let conn = Arc::new(Self {
-            id: rand::random(),
+            id,
             sender: Mutex::new(Some(sender)),
-            challenge: RwLock::new(String::new()),
-            pub_key: RwLock::new(None),
             ping: RwLock::new(-1),
             waiting_tasks: DashMap::new(),
             cleanup_handle: std::sync::Mutex::new(None),
@@ -163,11 +159,14 @@ impl OmikronConnection {
             return Ok(());
         }
 
-        // Handle ping regardless of auth state
+        // Handle ping regardless of message type
         if cv.is_type(CommunicationType::Ping) {
             return self.handle_ping(cv).await;
         }
-        return Ok(());
+
+        // Authentication is completed by the mtp host before this connection exists.
+        let omikron_id = self.id as i64;
+        self.clone().handle_authenticated(cv, omikron_id).await
     }
 
     async fn handle_authenticated(
@@ -1054,7 +1053,7 @@ pub async fn get_by_omikron_id(
         .ok()
         .map(|(bundle, _ip_address)| bundle)
 }
-pub async fn complete_register(pub_key: PublicKeyBundle, description: Option<String>) -> u64 {
+pub async fn complete_register(_pub_key: PublicKeyBundle, _description: Option<String>) -> u64 {
     0
 }
 
@@ -1097,7 +1096,8 @@ pub async fn start(port: u16) -> Result<(), Box<dyn std::error::Error>> {
 
     while let Ok(Some(mut connection)) = host.accept().await {
         tokio::spawn(async move {
-            let conn = OmikronConnection::new(connection.sender);
+            let conn = OmikronConnection::new(connection.sender, connection.client_id);
+            omikron_manager::add_omikron(conn.clone()).await;
             conn.handle(&mut connection.receiver).await;
         });
     }