Merge pull request #24 from CoLearn-Dev/registry

- registry support

Author: stneng

Cargo.lock

@@ -1,6 +1,6 @@
 [package]
 name = "colink-server"
-version = "0.1.7"
+version = "0.1.8"
 edition = "2021"
 
 [dependencies]
@@ -28,4 +28,5 @@ url = "2.2"
 uuid = { version = "0.8", features = ["v4"] }
 
 [build-dependencies]
+prost-build = "0.10"
 tonic-build = "0.7"

Cargo.toml

@@ -1,4 +1,6 @@
 fn main() {
     tonic_build::compile_protos("proto/colink.proto")
         .unwrap_or_else(|e| panic!("Failed to compile protos {:?}", e));
+    prost_build::compile_protos(&["proto/colink_registry.proto"], &["proto/"])
+        .unwrap_or_else(|e| panic!("Failed to compile protos {:?}", e));
 }

build.rs

@@ -1,11 +1,16 @@
 use super::utils::*;
 use crate::colink_proto::*;
+pub use colink_registry_proto::UserRecord;
 use openssl::sha::sha256;
 use prost::Message;
 use secp256k1::{ecdsa::Signature, PublicKey, Secp256k1};
 use tonic::{Request, Response, Status};
 use uuid::Uuid;
 
+mod colink_registry_proto {
+    include!(concat!(env!("OUT_DIR"), "/colink_registry.rs"));
+}
+
 impl crate::server::MyService {
     pub async fn _create_task(&self, request: Request<Task>) -> Result<Response<Task>, Status> {
         Self::check_privilege_in(request.metadata(), &["user"])?;
@@ -28,24 +33,13 @@ impl crate::server::MyService {
         self._internal_storage_update(&user_id, &format!("tasks:{}", task_id), &payload)
             .await?;
         for i in 1..task.participants.len() {
-            let core_addr = self
-                ._internal_storage_read(
-                    &user_id,
-                    &format!("known_users:{}:core_addr", &task.participants[i].user_id),
-                )
+            let (core_addr, guest_jwt) = self
+                .query_user_record(&user_id, &task.participants[i].user_id)
                 .await?;
-            let core_addr = String::from_utf8(core_addr).unwrap();
             let mut client = match self._grpc_connect(&core_addr).await {
                 Ok(client) => client,
                 Err(e) => return Err(Status::internal(format!("{}", e))),
             };
-            let guest_jwt = self
-                ._internal_storage_read(
-                    &user_id,
-                    &format!("known_users:{}:guest_jwt", &task.participants[i].user_id),
-                )
-                .await?;
-            let guest_jwt = String::from_utf8(guest_jwt).unwrap();
             client
                 .inter_core_sync_task(generate_request(&guest_jwt, task.clone()))
                 .await?;
@@ -136,24 +130,13 @@ impl crate::server::MyService {
         drop(task_storage_mutex);
 
         if task.require_agreement && user_status != "ignored" {
-            let core_addr = self
-                ._internal_storage_read(
-                    &user_id,
-                    &format!("known_users:{}:core_addr", &task.participants[0].user_id),
-                )
+            let (core_addr, guest_jwt) = self
+                .query_user_record(&user_id, &task.participants[0].user_id)
                 .await?;
-            let core_addr = String::from_utf8(core_addr).unwrap();
             let mut client = match self._grpc_connect(&core_addr).await {
                 Ok(client) => client,
                 Err(e) => return Err(Status::internal(format!("{}", e))),
             };
-            let guest_jwt = self
-                ._internal_storage_read(
-                    &user_id,
-                    &format!("known_users:{}:guest_jwt", &task.participants[0].user_id),
-                )
-                .await?;
-            let guest_jwt = String::from_utf8(guest_jwt).unwrap();
             client
                 .inter_core_sync_task(generate_request(
                     &guest_jwt,
@@ -304,24 +287,13 @@ impl crate::server::MyService {
                 // The initiator should broadcast the status change.
                 if task.participants[0].user_id == user_id && task.participants.len() > 2 {
                     for i in 1..task.participants.len() {
-                        let core_addr = self
-                            ._internal_storage_read(
-                                &user_id,
-                                &format!("known_users:{}:core_addr", &task.participants[i].user_id),
-                            )
+                        let (core_addr, guest_jwt) = self
+                            .query_user_record(&user_id, &task.participants[i].user_id)
                             .await?;
-                        let core_addr = String::from_utf8(core_addr).unwrap();
                         let mut client = match self._grpc_connect(&core_addr).await {
                             Ok(client) => client,
                             Err(e) => return Err(Status::internal(format!("{}", e))),
                         };
-                        let guest_jwt = self
-                            ._internal_storage_read(
-                                &user_id,
-                                &format!("known_users:{}:guest_jwt", &task.participants[i].user_id),
-                            )
-                            .await?;
-                        let guest_jwt = String::from_utf8(guest_jwt).unwrap();
                         client
                             .inter_core_sync_task(generate_request(&guest_jwt, task.clone()))
                             .await?;
@@ -437,6 +409,89 @@ impl crate::server::MyService {
         Ok(())
     }
 
+    async fn query_user_record(
+        &self,
+        user_id: &str,
+        query_user_id: &str,
+    ) -> Result<(String, String), Status> {
+        let mut counter = 0;
+        while self
+            ._internal_storage_read(
+                user_id,
+                &format!("known_users:{}:core_addr", &query_user_id),
+            )
+            .await
+            .is_err()
+            || self
+                ._internal_storage_read(
+                    user_id,
+                    &format!("known_users:{}:guest_jwt", &query_user_id),
+                )
+                .await
+                .is_err()
+        {
+            if counter == 0 {
+                let user = UserRecord {
+                    user_id: query_user_id.to_string(),
+                    ..Default::default()
+                };
+                let mut payload = vec![];
+                user.encode(&mut payload).unwrap();
+                let mut local_task = Task {
+                    task_id: Uuid::new_v4().to_string(),
+                    protocol_name: "registry".to_string(),
+                    protocol_param: payload,
+                    participants: vec![Participant {
+                        user_id: user_id.to_string(),
+                        role: "query".to_string(),
+                    }],
+                    require_agreement: false,
+                    status: "started".to_string(),
+                    expiration_time: chrono::Utc::now().timestamp() + 86400,
+                    ..Default::default()
+                };
+                local_task
+                    .decisions
+                    .resize(local_task.participants.len(), Default::default());
+                local_task.decisions[0] = self
+                    .generate_decision(true, false, "", user_id, &local_task)
+                    .await?;
+                let mut payload = vec![];
+                local_task.encode(&mut payload).unwrap();
+                self._internal_storage_update(
+                    user_id,
+                    &format!("tasks:{}", local_task.task_id),
+                    &payload,
+                )
+                .await?;
+                let task_storage_mutex = self.task_storage_mutex.lock().await;
+                self.add_task_new_status(user_id, &local_task).await?;
+                drop(task_storage_mutex);
+            }
+            // We choose 1 second as the retry interval and retry 30 times.
+            tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
+            counter += 1;
+            if counter > 30 {
+                break;
+            }
+        }
+        let core_addr = self
+            ._internal_storage_read(
+                user_id,
+                &format!("known_users:{}:core_addr", &query_user_id),
+            )
+            .await?;
+        let core_addr = String::from_utf8(core_addr).unwrap();
+        let guest_jwt = self
+            ._internal_storage_read(
+                user_id,
+                &format!("known_users:{}:guest_jwt", &query_user_id),
+            )
+            .await?;
+        let guest_jwt = String::from_utf8(guest_jwt).unwrap();
+        Ok((core_addr, guest_jwt))
+    }
+
     /**
      * This function only checks the validity of the signature, user's decision will not be checked.
      */