From 84d56f3320ae5f70597571d488b3852dee9ddb85 Mon Sep 17 00:00:00 2001
From: Tamo <tamo@meilisearch.com>
Date: Wed, 2 Aug 2023 13:57:30 +0200
Subject: [PATCH] send the creation of api-key to zookeeper

---
 Cargo.lock                        |  1 +
 meilisearch-auth/Cargo.toml       |  1 +
 meilisearch-auth/src/error.rs     |  3 +++
 meilisearch-auth/src/lib.rs       | 18 ++++++++++++++++--
 meilisearch/src/routes/api_key.rs | 10 +++-------
 5 files changed, 24 insertions(+), 9 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 41a6fab29..46dab1af1 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -2689,6 +2689,7 @@ dependencies = [
  "sha2",
  "thiserror",
  "time",
+ "tokio",
  "uuid 1.3.3",
  "zookeeper-client",
 ]
diff --git a/meilisearch-auth/Cargo.toml b/meilisearch-auth/Cargo.toml
index bb9439622..fd60cf3ee 100644
--- a/meilisearch-auth/Cargo.toml
+++ b/meilisearch-auth/Cargo.toml
@@ -23,5 +23,6 @@ serde_json = { version = "1.0.95", features = ["preserve_order"] }
 sha2 = "0.10.6"
 thiserror = "1.0.40"
 time = { version = "0.3.20", features = ["serde-well-known", "formatting", "parsing", "macros"] }
+tokio = { version = "1.27.0", features = ["full"] }
 uuid = { version = "1.3.1", features = ["serde", "v4"] }
 zookeeper-client = "0.4.0"
diff --git a/meilisearch-auth/src/error.rs b/meilisearch-auth/src/error.rs
index 72cd96470..eb4cd9b48 100644
--- a/meilisearch-auth/src/error.rs
+++ b/meilisearch-auth/src/error.rs
@@ -2,6 +2,7 @@ use std::error::Error;
 
 use meilisearch_types::error::{Code, ErrorCode};
 use meilisearch_types::internal_error;
+use zookeeper_client as zk;
 
 pub type Result<T> = std::result::Result<T, AuthControllerError>;
 
@@ -19,6 +20,8 @@ internal_error!(
     AuthControllerError: meilisearch_types::milli::heed::Error,
     std::io::Error,
     serde_json::Error,
+    tokio::task::JoinError,
+    zk::Error,
     std::str::Utf8Error
 );
 
diff --git a/meilisearch-auth/src/lib.rs b/meilisearch-auth/src/lib.rs
index 861e47fa2..ce2d0c8a9 100644
--- a/meilisearch-auth/src/lib.rs
+++ b/meilisearch-auth/src/lib.rs
@@ -56,10 +56,24 @@ impl AuthController {
         self.store.used_size()
     }
 
-    pub fn create_key(&self, create_key: CreateApiKey) -> Result<Key> {
+    pub async fn create_key(&self, create_key: CreateApiKey) -> Result<Key> {
         match self.store.get_api_key(create_key.uid)? {
             Some(_) => Err(AuthControllerError::ApiKeyAlreadyExists(create_key.uid.to_string())),
-            None => self.store.put_api_key(create_key.to_key()),
+            None => {
+                let store = self.store.clone();
+                let key =
+                    tokio::task::spawn_blocking(move || store.put_api_key(create_key.to_key()))
+                        .await??;
+                if let Some(ref zk) = self.zk {
+                    zk.create(
+                        &format!("/auth/{}", key.uid),
+                        &serde_json::to_vec_pretty(&key)?,
+                        &zk::CreateOptions::new(zk::CreateMode::Persistent, zk::Acl::anyone_all()),
+                    )
+                    .await?;
+                }
+                Ok(key)
+            }
         }
     }
 
diff --git a/meilisearch/src/routes/api_key.rs b/meilisearch/src/routes/api_key.rs
index 597d04486..9b2193b01 100644
--- a/meilisearch/src/routes/api_key.rs
+++ b/meilisearch/src/routes/api_key.rs
@@ -41,14 +41,10 @@ pub async fn create_api_key(
     _req: HttpRequest,
 ) -> Result<HttpResponse, ResponseError> {
     let v = body.into_inner();
-    let res = tokio::task::spawn_blocking(move || -> Result<_, AuthControllerError> {
-        let key = auth_controller.create_key(v)?;
-        Ok(KeyView::from_key(key, &auth_controller))
-    })
-    .await
-    .map_err(|e| ResponseError::from_msg(e.to_string(), Code::Internal))??;
+    let key = auth_controller.create_key(v).await?;
+    let key = KeyView::from_key(key, &auth_controller);
 
-    Ok(HttpResponse::Created().json(res))
+    Ok(HttpResponse::Created().json(key))
 }
 
 #[derive(Deserr, Debug, Clone, Copy)]