Merge pull request #26 from CoLearn-Dev/instant-server

- instant server; wait_user_init; protocol_attach; switch_to_generated_user; bugfix

Author: stneng

.github/workflows/check.yml

@@ -1,4 +1,8 @@
-on: [push]
+on:
+  push:
+    branches:
+      - main
+  pull_request:
 
 name: check
 

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "colink"
-version = "0.2.0"
+version = "0.2.1"
 edition = "2021"
 description = "CoLink Rust SDK"
 license = "MIT"
@@ -13,7 +13,6 @@ repository = "https://github.com/CoLearn-Dev/colink-sdk-rust-dev"
 async-trait = "0.1"
 base64 = "0.13.0"
 chrono = "0.4"
-ctrlc = { version = "3.2", features = ["termination"] }
 futures-lite = "1.12"
 lapin = "2.1"
 prost = "0.10"
@@ -34,9 +33,10 @@ prost-build = "0.10"
 tonic-build = "0.7"
 
 [features]
-default = ["extensions", "remote_storage", "variable_transfer", "registry", "policy_module"]
+default = ["extensions", "remote_storage", "variable_transfer", "registry", "policy_module", "instant_server"]
 extensions = []
 remote_storage = ["extensions"]
 variable_transfer = ["extensions", "remote_storage"]
 registry = []
 policy_module = []
+instant_server = []

README.md

@@ -9,7 +9,7 @@ CoLink SDK helps both application adnd protocol developers access the functional
 Add this to your Cargo.toml:
 ```toml
 [dependencies]
-colink = "0.2.0"
+colink = "0.2.1"
 ```
 
 ## Getting Started

examples/protocol_greetings_with_instant_server.rs

@@ -0,0 +1,68 @@
+#![allow(unused_variables)]
+use colink::{extensions::instant_server::InstantServer, CoLink, Participant, ProtocolEntry};
+
+struct Initiator;
+#[colink::async_trait]
+impl ProtocolEntry for Initiator {
+    async fn start(
+        &self,
+        cl: CoLink,
+        param: Vec<u8>,
+        participants: Vec<Participant>,
+    ) -> Result<(), Box<dyn std::error::Error + Send + Sync + 'static>> {
+        println!("initiator");
+        Ok(())
+    }
+}
+
+struct Receiver;
+#[colink::async_trait]
+impl ProtocolEntry for Receiver {
+    async fn start(
+        &self,
+        cl: CoLink,
+        param: Vec<u8>,
+        participants: Vec<Participant>,
+    ) -> Result<(), Box<dyn std::error::Error + Send + Sync + 'static>> {
+        println!("{}", String::from_utf8_lossy(&param));
+        cl.create_entry(&format!("tasks:{}:output", cl.get_task_id()?), &param)
+            .await?;
+        Ok(())
+    }
+}
+
+#[tokio::main]
+async fn main() -> Result<(), Box<dyn std::error::Error + Send + Sync + 'static>> {
+    let is0 = InstantServer::new();
+    let is1 = InstantServer::new();
+    let cl0 = is0.get_colink().switch_to_generated_user().await?;
+    let cl1 = is0.get_colink().switch_to_generated_user().await?;
+    colink::protocol_attach!(
+        cl0,
+        ("greetings:initiator", Initiator),
+        ("greetings:receiver", Receiver)
+    );
+    colink::protocol_attach!(
+        cl1,
+        ("greetings:initiator", Initiator),
+        ("greetings:receiver", Receiver)
+    );
+    let participants = vec![
+        Participant {
+            user_id: cl0.get_user_id()?,
+            role: "initiator".to_string(),
+        },
+        Participant {
+            user_id: cl1.get_user_id()?,
+            role: "receiver".to_string(),
+        },
+    ];
+    let task_id = cl0
+        .run_task("greetings", "test".as_bytes(), &participants, true)
+        .await?;
+    let res = cl1
+        .read_or_wait(&format!("tasks:{}:output", task_id))
+        .await?;
+    println!("{}", String::from_utf8_lossy(&res));
+    Ok(())
+}

examples/user_wait_task.rs

@@ -24,7 +24,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error + Send + Sync + 'static>
         role: "query_from_registries".to_string(),
     }];
     let task_id = cl
-        .run_task("registry", target_user.as_bytes(), &participants, false)
+        .run_task("registry", &payload, &participants, false)
         .await?;
     println!(
         "Task {} has been created, waiting for it to finish...",

src/application.rs

@@ -45,23 +45,19 @@ impl CoLink {
         }
     }
 
-    pub fn ca_certificate(self, ca_certificate: &str) -> Self {
+    pub fn ca_certificate(mut self, ca_certificate: &str) -> Self {
         let ca_certificate = std::fs::read(ca_certificate).unwrap();
         let ca_certificate = Certificate::from_pem(ca_certificate);
-        Self {
-            ca_certificate: Some(ca_certificate),
-            ..self
-        }
+        self.ca_certificate = Some(ca_certificate);
+        self
     }
 
-    pub fn identity(self, client_cert: &str, client_key: &str) -> Self {
+    pub fn identity(mut self, client_cert: &str, client_key: &str) -> Self {
         let client_cert = std::fs::read(client_cert).unwrap();
         let client_key = std::fs::read(client_key).unwrap();
         let identity = Identity::from_pem(client_cert, client_key);
-        Self {
-            identity: Some(identity),
-            ..self
-        }
+        self.identity = Some(identity);
+        self
     }
 
     async fn _grpc_connect(&self, address: &str) -> Result<CoLinkClient<Channel>, Error> {

src/extensions.rs

@@ -1,5 +1,7 @@
 #[cfg(feature = "extensions")]
 mod get_participant_id;
+#[cfg(feature = "instant_server")]
+pub mod instant_server;
 #[cfg(feature = "extensions")]
 mod lock;
 #[cfg(feature = "policy_module")]
@@ -10,7 +12,11 @@ mod read_or_wait;
 pub mod registry;
 #[cfg(feature = "remote_storage")]
 mod remote_storage;
+#[cfg(feature = "extensions")]
+mod switch_to_generated_user;
 #[cfg(feature = "variable_transfer")]
 mod variable_transfer;
 #[cfg(feature = "extensions")]
 mod wait_task;
+#[cfg(feature = "extensions")]
+mod wait_user_init;

src/extensions/instant_server.rs

@@ -0,0 +1,105 @@
+use crate::{utils::get_colink_home, CoLink};
+use rand::Rng;
+use std::{
+    path::Path,
+    process::{Child, Command, Stdio},
+};
+
+pub struct InstantServer {
+    id: String,
+    port: i32,
+    host_token: String,
+    process: Child,
+}
+
+impl Drop for InstantServer {
+    fn drop(&mut self) {
+        Command::new("pkill")
+            .arg("-9")
+            .arg("-P")
+            .arg(&self.process.id().to_string())
+            .stdout(Stdio::null())
+            .stderr(Stdio::null())
+            .status()
+            .unwrap();
+        self.process.kill().unwrap();
+        let colink_home = get_colink_home().unwrap();
+        let working_dir = Path::new(&colink_home)
+            .join("instant_servers")
+            .join(self.id.clone());
+        std::fs::remove_dir_all(&working_dir).unwrap();
+    }
+}
+
+impl Default for InstantServer {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl InstantServer {
+    pub fn new() -> Self {
+        let colink_home = get_colink_home().unwrap();
+        let program = Path::new(&colink_home).join("colink-server");
+        if std::fs::metadata(program.clone()).is_err() {
+            Command::new("bash")
+                .arg("-c")
+                .arg("bash -c \"$(curl -fsSL https://raw.githubusercontent.com/CoLearn-Dev/colinkctl/main/install_colink.sh)\"")
+                .status()
+                .unwrap();
+        }
+        let instant_server_id = uuid::Uuid::new_v4().to_string();
+        let mut port = rand::thread_rng().gen_range(10000..20000);
+        while std::net::TcpStream::connect(&format!("127.0.0.1:{}", port)).is_ok() {
+            port = rand::thread_rng().gen_range(10000..20000);
+        }
+        let working_dir = Path::new(&colink_home)
+            .join("instant_servers")
+            .join(instant_server_id.clone());
+        std::fs::create_dir_all(&working_dir).unwrap();
+        let child = Command::new(program)
+            .args([
+                "--address",
+                "0.0.0.0",
+                "--port",
+                &port.to_string(),
+                "--mq-amqp",
+                "amqp://guest:guest@localhost:5672",
+                "--mq-api",
+                "http://guest:guest@localhost:15672/api",
+                "--mq-prefix",
+                &format!("colink-instant-server-{}", port),
+                "--core-uri",
+                &format!("http://127.0.0.1:{}", port),
+                "--inter-core-reverse-mode",
+            ])
+            .env("COLINK_HOME", colink_home)
+            .current_dir(working_dir.clone())
+            .stdout(Stdio::null())
+            .stderr(Stdio::null())
+            .spawn()
+            .unwrap();
+        loop {
+            if std::fs::metadata(working_dir.join("host_token.txt")).is_ok()
+                && std::net::TcpStream::connect(&format!("127.0.0.1:{}", port)).is_ok()
+            {
+                break;
+            }
+            std::thread::sleep(core::time::Duration::from_millis(10));
+        }
+        let host_token: String =
+            String::from_utf8_lossy(&std::fs::read(working_dir.join("host_token.txt")).unwrap())
+                .parse()
+                .unwrap();
+        Self {
+            id: instant_server_id,
+            port,
+            host_token,
+            process: child,
+        }
+    }
+
+    pub fn get_colink(&self) -> CoLink {
+        CoLink::new(&format!("http://127.0.0.1:{}", self.port), &self.host_token)
+    }
+}

src/extensions/switch_to_generated_user.rs

@@ -0,0 +1,22 @@
+use crate::decode_jwt_without_validation;
+
+type Error = Box<dyn std::error::Error + Send + Sync + 'static>;
+
+impl crate::application::CoLink {
+    async fn generate_user_and_import(&self) -> Result<String, Error> {
+        let auth_content = decode_jwt_without_validation(&self.jwt)?;
+        let expiration_timestamp = auth_content.exp;
+        let (pk, sk) = crate::generate_user();
+        let (_, core_pub_key, _) = self.request_info().await?;
+        let (signature_timestamp, sig) =
+            crate::prepare_import_user_signature(&pk, &sk, &core_pub_key, expiration_timestamp);
+        self.import_user(&pk, signature_timestamp, expiration_timestamp, &sig)
+            .await
+    }
+
+    pub async fn switch_to_generated_user(&self) -> Result<Self, Error> {
+        let cl = Self::new(&self.core_addr, &self.generate_user_and_import().await?);
+        cl.wait_user_init().await?;
+        Ok(cl)
+    }
+}

src/extensions/wait_user_init.rs

@@ -0,0 +1,40 @@
+use crate::{colink_proto::*, utils::get_path_timestamp};
+use prost::Message;
+use tracing::debug;
+
+type Error = Box<dyn std::error::Error + Send + Sync + 'static>;
+
+impl crate::application::CoLink {
+    pub async fn wait_user_init(&self) -> Result<(), Error> {
+        let is_initialized_key = "_internal:_is_initialized";
+        let start_timestamp = match self
+            .read_entries(&[StorageEntry {
+                key_name: is_initialized_key.to_string(),
+                ..Default::default()
+            }])
+            .await
+        {
+            Ok(res) => {
+                if res[0].payload[0] == 1 {
+                    return Ok(());
+                }
+                get_path_timestamp(&res[0].key_path) + 1
+            }
+            Err(_) => 0,
+        };
+        let queue_name = self
+            .subscribe(is_initialized_key, Some(start_timestamp))
+            .await?;
+        let mut subscriber = self.new_subscriber(&queue_name).await?;
+        loop {
+            let data = subscriber.get_next().await?;
+            debug!("Received [{}]", String::from_utf8_lossy(&data));
+            let message: SubscriptionMessage = Message::decode(&*data).unwrap();
+            if message.change_type != "delete" && message.payload[0] == 1 {
+                break;
+            }
+        }
+        self.unsubscribe(&queue_name).await?;
+        Ok(())
+    }
+}