Feat/implement api key management

Cargo.toml

@@ -61,6 +61,7 @@ itertools = "0.14"
 
 # Webserver
 axum = { version = "0.8.1", features = ["macros"] }
+axum-extra = { version = "0.10.0", features = ["typed-header"] }
 tower-http = { version = "0.6.2", features = ["trace", "cors"] }
 http = "1.2.0"
 axum-openapi3 = { version = "0.2", features = ["derive"] }
@@ -97,6 +98,7 @@ memmap = "0.7.0"
 cuid2 = "0.1.4"
 
 mobc = "0.8"
+redact = { version = "0.1", features = ["serde"] }
 
 [build-dependencies]
 cc = "1"

api-test.hurl

@@ -10,8 +10,11 @@ HTTP 200
 collection_id: jsonpath "$[0]"
 
 POST {{base_url}}/v0/collections
+Authorization: Bearer my-master-api-key
 {
-    "id": "{{collection_id}}"
+    "id": "{{collection_id}}",
+    "read_api_key": "my-read-api-key",
+    "write_api_key": "my-write-api-key"
 }
 HTTP 201
 
@@ -21,6 +24,7 @@ HTTP 200
 jsonpath "$.document_count" == 0
 
 PATCH {{base_url}}/v0/collections/{{collection_id}}/documents
+Authorization: Bearer my-write-api-key
 [
     {
         "id": "1",
@@ -42,7 +46,7 @@ HTTP 200
 [Asserts]
 jsonpath "$.document_count" == 2
 
-POST {{base_url}}/v0/collections/{{collection_id}}/search
+POST {{base_url}}/v0/collections/{{collection_id}}/search?api-key=my-read-api-key
 [Options]
 retry: 10
 retry-interval: 500ms

config.yaml

@@ -10,6 +10,9 @@ log:
 
 writer_side:
     output: in-memory
+
+    # Replace the following value with your own API key
+    master_api_key: foobar
     config:
         data_dir: ./.data/writer
         # The maximum number of embeddings that can be stored in the queue

embedding-api-test.hurl

@@ -10,16 +10,20 @@ HTTP 200
 collection_id: jsonpath "$[0]"
 
 POST {{base_url}}/v0/collections
+Authorization: Bearer my-master-api-key
 {
     "id": "{{collection_id}}",
     "embeddings": {
         "model_name": "BGESmall",
         "document_fields": ["content"]
-    }
+    },
+    "read_api_key": "my-read-api-key",
+    "write_api_key": "my-write-api-key"
 }
 HTTP 201
 
 PATCH {{base_url}}/v0/collections/{{collection_id}}/documents
+Authorization: Bearer my-write-api-key
 [
     {
         "id": "1",
@@ -37,7 +41,8 @@ PATCH {{base_url}}/v0/collections/{{collection_id}}/documents
 HTTP 200
 
 
-POST {{base_url}}/v0/collections/{{collection_id}}/search
+POST {{base_url}}/v0/collections/{{collection_id}}/search?api-key=my-read-api-key
+Authorization: Bearer my-read-api-key
 [Options]
 retry: 10
 retry-interval: 500ms

src/collection_manager/dto.rs

@@ -1,7 +1,8 @@
 use std::collections::HashMap;
 
-use axum_openapi3::utoipa::ToSchema;
 use axum_openapi3::utoipa::{self, IntoParams};
+use axum_openapi3::utoipa::{PartialSchema, ToSchema};
+use redact::Secret;
 use serde::{de, Deserialize, Serialize};
 
 use crate::ai::OramaModel;
@@ -78,10 +79,39 @@ pub struct CreateCollectionEmbeddings {
     pub document_fields: Vec<String>,
 }
 
-#[derive(Debug, Serialize, Deserialize, ToSchema)]
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct ApiKey(pub Secret<String>);
+
+impl<'de> Deserialize<'de> for ApiKey {
+    fn deserialize<D>(deserializer: D) -> Result<ApiKey, D::Error>
+    where
+        D: de::Deserializer<'de>,
+    {
+        let s = String::deserialize(deserializer)?;
+
+        if s.is_empty() {
+            return Err(serde::de::Error::custom("API key cannot be empty"));
+        }
+
+        Ok(ApiKey(Secret::new(s)))
+    }
+}
+
+impl PartialSchema for ApiKey {
+    fn schema() -> utoipa::openapi::RefOr<utoipa::openapi::schema::Schema> {
+        String::schema()
+    }
+}
+impl ToSchema for ApiKey {}
+
+#[derive(Debug, Deserialize, ToSchema)]
 pub struct CreateCollection {
     pub id: CollectionId,
     pub description: Option<String>,
+
+    pub read_api_key: ApiKey,
+    pub write_api_key: ApiKey,
+
     #[schema(inline)]
     pub language: Option<LanguageDTO>,
     #[serde(default)]
@@ -98,6 +128,7 @@ impl TryFrom<serde_json::Value> for CreateCollection {
     }
 }
 
+#[derive(Debug, Serialize, Deserialize, ToSchema)]
 pub struct DeleteDocuments {
     pub document_ids: Vec<String>,
 }
@@ -362,7 +393,7 @@ pub struct DeleteHookParams {
     pub name: HookName,
 }
 
-#[derive(Deserialize, Clone, Serialize, ToSchema, IntoParams)]
+#[derive(Deserialize, Clone, Serialize, ToSchema)]
 pub struct ExecuteActionPayload {
     pub name: String, // we're not using an enum here since users will be able to define their own actions
     pub context: String,
@@ -421,7 +452,9 @@ mod test {
                     "model_name": "gte-small",
                     "document_fields": ["text"],
                 }
-            }
+            },
+            "read_api_key": "foo",
+            "write_api_key": "bar",
         })
         .try_into()
         .unwrap();

src/collection_manager/sides/read/collection.rs

@@ -12,6 +12,7 @@ use committed::CommittedCollection;
 use dashmap::DashMap;
 use dump::{CollectionInfo, CollectionInfoV1};
 use merge::{merge_bool_field, merge_number_field, merge_string_field, merge_vector_field};
+use redact::Secret;
 use serde::{Deserialize, Serialize};
 use tokio::{
     join,
@@ -28,8 +29,8 @@ use crate::{
     ai::{AIService, OramaModel},
     collection_manager::{
         dto::{
-            self, BM25Scorer, FacetDefinition, FacetResult, FieldId, Filter, Limit, NumberFilter,
-            Properties, SearchMode, SearchParams,
+            self, ApiKey, BM25Scorer, FacetDefinition, FacetResult, FieldId, Filter, Limit,
+            NumberFilter, Properties, SearchMode, SearchParams,
         },
         sides::{CollectionWriteOperation, Offset, OramaModelSerializable},
     },
@@ -48,6 +49,7 @@ use super::IndexesConfig;
 #[derive(Debug)]
 pub struct CollectionReader {
     id: CollectionId,
+    read_api_key: ApiKey,
     ai_service: Arc<AIService>,
     nlp_service: Arc<NLPService>,
 
@@ -70,12 +72,14 @@ pub struct CollectionReader {
 impl CollectionReader {
     pub fn try_new(
         id: CollectionId,
+        read_api_key: ApiKey,
         ai_service: Arc<AIService>,
         nlp_service: Arc<NLPService>,
         _: IndexesConfig,
     ) -> Result<Self> {
         Ok(Self {
             id,
+            read_api_key,
             ai_service,
             nlp_service,
             document_count: AtomicU64::new(0),
@@ -92,6 +96,13 @@ impl CollectionReader {
         })
     }
 
+    pub fn check_read_api_key(&self, api_key: ApiKey) -> Result<()> {
+        if api_key != self.read_api_key {
+            return Err(anyhow!("Invalid read api key"));
+        }
+        Ok(())
+    }
+
     #[inline]
     pub fn get_id(&self) -> CollectionId {
         self.id.clone()
@@ -139,6 +150,7 @@ impl CollectionReader {
             .context("Cannot read previous collection info")?;
 
         let dump::CollectionInfo::V1(collection_info) = collection_info;
+        self.read_api_key = ApiKey(Secret::new(collection_info.read_api_key));
 
         for (field_name, (field_id, field_type)) in collection_info.fields {
             let typed_field: TypedField = match field_type {
@@ -218,6 +230,7 @@ impl CollectionReader {
             debug!("No previous collection info found, creating a new one");
             CollectionInfoV1 {
                 fields: Default::default(),
+                read_api_key: self.read_api_key.0.expose_secret().clone(),
                 id: self.id.clone(),
                 used_models: Default::default(),
                 number_field_infos: Default::default(),
@@ -1121,6 +1134,7 @@ mod dump {
     #[derive(Debug, Serialize, Deserialize)]
     pub struct CollectionInfoV1 {
         pub id: CollectionId,
+        pub read_api_key: String,
         pub fields: Vec<(String, (FieldId, TypedField))>,
         pub used_models: Vec<(OramaModelSerializable, Vec<FieldId>)>,
         pub number_field_infos: Vec<(FieldId, committed::fields::NumberFieldInfo)>,