add api-key auth in http and grpc servers

Signed-off-by: Arshdeep54 <balarsh535@gmail.com>

Author: Arshdeep54

.env.example

@@ -5,8 +5,10 @@ HTTP_PORT=3000
 # gRPC Server
 GRPC_HOST=127.0.0.1
 GRPC_PORT=50051
-# required
-GRPC_ROOT_PASSWORD=your-secure-password
+
+# API keys used by both the HTTP and gRPC servers (required)
+# See keys.example.json for the file format
+VORTEXDB_KEYS_FILE=./keys.json
 
 # Database Configuration
 # Storage: inmemory, rocksdb

.gitignore

@@ -7,6 +7,7 @@
 .TODO
 /databases
 .env
+/keys.json
 __pycache__/
 build/
 .venv/

README.md

@@ -67,12 +67,23 @@ cp .env.example .env
 ```
 
 The following env vars are important, and are required to be set by the user:  
-`GRPC_ROOT_PASSWORD`  
+`VORTEXDB_KEYS_FILE`  
 `DIMENSION`  
 `DATA_PATH`  
 
 **NOTE**: `DATA_PATH` is the directory within the container where persistent data is stored
 
+**NOTE**: `VORTEXDB_KEYS_FILE` points to a JSON file of API keys shared by both the HTTP and gRPC servers, e.g.:
+```json
+{
+  "keys": [
+    { "name": "admin", "role": "readwrite", "key": "some-random-secret" },
+    { "name": "search-service", "role": "readonly", "key": "another-random-secret" }
+  ]
+}
+```
+Requests authenticate with an `api-key: <key>` header over HTTP, or an `authorization: Bearer <key>` header over gRPC. `readonly` keys can read/search but not insert, batch-insert, or delete; `readwrite` keys can do all of the above. See `keys.example.json`.
+
 
 Setting of the following env vars is optional, as they fallback to safe defaults, but recommended:  
 | .env Var | Function | Safe Default |

client/python/USAGE.md

@@ -18,7 +18,7 @@ pip install -e .
 The client communicates with VortexDB over gRPC and requires:
 
 - gRPC endpoint (host:port)
-- API key (maps to `GRPC_ROOT_PASSWORD` on the server)
+- API key (must match a `key` entry in the server's `VORTEXDB_KEYS_FILE`)
 
 These can be provided either:
 

crates/defs/src/auth.rs

@@ -0,0 +1,62 @@
+use serde::{Deserialize, Serialize};
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
+#[serde(rename_all = "lowercase")]
+pub enum ApiKeyRole {
+    ReadOnly,
+    ReadWrite,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct ApiKeyEntry {
+    pub name: String,
+    pub role: ApiKeyRole,
+    pub key: String,
+}
+
+#[derive(Debug, Clone, Default)]
+pub struct ApiKeyStore {
+    keys: Vec<ApiKeyEntry>,
+}
+
+fn constant_time_eq(a: &[u8], b: &[u8]) -> bool {
+    if a.len() != b.len() {
+        return false;
+    }
+    let mut diff = 0u8;
+    for (x, y) in a.iter().zip(b.iter()) {
+        diff |= x ^ y;
+    }
+    diff == 0
+}
+
+impl ApiKeyStore {
+    pub fn new(keys: Vec<ApiKeyEntry>) -> Self {
+        Self { keys }
+    }
+
+    pub fn find(&self, presented: &str) -> Option<&ApiKeyEntry> {
+        let mut matched = None;
+        for entry in &self.keys {
+            if constant_time_eq(entry.key.as_bytes(), presented.as_bytes()) {
+                matched = Some(entry);
+            }
+        }
+        matched
+    }
+
+    pub fn is_empty(&self) -> bool {
+        self.keys.is_empty()
+    }
+
+    pub fn duplicate_key(&self) -> Option<&str> {
+        for (i, entry) in self.keys.iter().enumerate() {
+            for other in &self.keys[i + 1..] {
+                if constant_time_eq(entry.key.as_bytes(), other.key.as_bytes()) {
+                    return Some(&entry.key);
+                }
+            }
+        }
+        None
+    }
+}

crates/defs/src/lib.rs

@@ -1,7 +1,9 @@
+pub mod auth;
 pub mod error;
 pub mod types;
 
 // Without re-exports, users would need to write defs::types::SomeType instead of just defs::SomeType. Re-exports simplify the API by flattening the module hierarchy. The * means "everything public" from that module.
+pub use auth::*;
 pub use error::*;
 use std::path::{Path, PathBuf};
 pub use types::*;

crates/grpc/src/interceptors.rs

@@ -1,36 +1,56 @@
+use std::sync::Arc;
+
+use defs::{ApiKeyEntry, ApiKeyRole, ApiKeyStore};
 use tonic::{Status, service::Interceptor};
 use tracing::{Level, event};
 
 use crate::constants::AUTHORIZATION_HEADER_KEY;
 
 #[derive(Clone)]
 pub struct AuthInterceptor {
-    root_password: String,
+    keys: Arc<ApiKeyStore>,
+}
+
+fn extract_bearer_token<T>(req: &tonic::Request<T>) -> Option<&str> {
+    req.metadata()
+        .get(AUTHORIZATION_HEADER_KEY)?
+        .to_str()
+        .ok()?
+        .strip_prefix("Bearer ")
+        .filter(|token| !token.is_empty())
+}
+
+pub fn require_write_role<T>(req: &tonic::Request<T>) -> Result<(), Status> {
+    match req.extensions().get::<ApiKeyEntry>() {
+        Some(entry) if entry.role == ApiKeyRole::ReadWrite => Ok(()),
+        Some(_) => Err(Status::permission_denied(
+            "This api key does not have write access",
+        )),
+        None => Err(Status::unauthenticated("Invalid credentials")),
+    }
 }
 
 impl Interceptor for AuthInterceptor {
-    fn call(&mut self, req: tonic::Request<()>) -> Result<tonic::Request<()>, Status> {
-        let auth_token = match req.metadata().get(AUTHORIZATION_HEADER_KEY) {
-            Some(t) => t,
-            None => return Err(Status::unauthenticated("Invalid credentials")),
-        };
-        if auth_token
-            .to_str()
-            .unwrap_or_default()
-            .strip_prefix("Bearer ")
-            .unwrap_or_default()
-            == self.root_password
-        {
-            Ok(req)
-        } else {
-            event!(Level::WARN, "Unauthorized Request");
-            Err(Status::unauthenticated("Invalid credentials"))
+    fn call(&mut self, mut req: tonic::Request<()>) -> Result<tonic::Request<()>, Status> {
+        let matched = extract_bearer_token(&req)
+            .and_then(|token| self.keys.find(token))
+            .cloned();
+
+        match matched {
+            Some(entry) => {
+                req.extensions_mut().insert(entry);
+                Ok(req)
+            }
+            None => {
+                event!(Level::WARN, "Unauthorized Request");
+                Err(Status::unauthenticated("Invalid credentials"))
+            }
         }
     }
 }
 
 impl AuthInterceptor {
-    pub fn new(root_password: String) -> AuthInterceptor {
-        AuthInterceptor { root_password }
+    pub fn new(keys: Arc<ApiKeyStore>) -> AuthInterceptor {
+        AuthInterceptor { keys }
     }
 }

crates/grpc/src/lib.rs

@@ -15,17 +15,13 @@ use utils::ServerEndpoint;
 pub async fn run_grpc_server(
     db: Arc<VectorDb>,
     addr: SocketAddr,
-    root_password: String,
+    keys: Arc<defs::ApiKeyStore>,
     logging: bool,
 ) -> Result<(), BoxError> {
     let vector_db_service = VectorDBService::new(db, logging);
-    run_server(
-        vector_db_service,
-        ServerEndpoint::Address(addr),
-        root_password,
-    )
-    .await
-    .map_err(|e| -> BoxError { e.to_string().into() })
+    run_server(vector_db_service, ServerEndpoint::Address(addr), keys)
+        .await
+        .map_err(|e| -> BoxError { e.to_string().into() })
 }
 
 #[cfg(test)]

crates/grpc/src/service.rs

@@ -38,6 +38,7 @@ impl VectorDb for VectorDBService {
         request: Request<InsertVectorRequest>,
     ) -> Result<Response<PointId>, Status> {
         log_rpc("insert_vector", self.logging);
+        interceptors::require_write_role(&request)?;
 
         let inner_request = request.into_inner();
 
@@ -157,6 +158,7 @@ impl VectorDb for VectorDBService {
 
     async fn delete_point(&self, request: Request<PointId>) -> Result<Response<()>, Status> {
         log_rpc("delete_point", self.logging);
+        interceptors::require_write_role(&request)?;
 
         let point_id = request.into_inner().id.unwrap_or_default().value;
 
@@ -179,6 +181,7 @@ impl VectorDb for VectorDBService {
         &self,
         request: tonic::Request<InsertVectorsBatchRequest>,
     ) -> Result<tonic::Response<InsertVectorsBatchResponse>, tonic::Status> {
+        interceptors::require_write_role(&request)?;
         let req = request.into_inner();
         let mut ids = Vec::with_capacity(req.vectors.len());
 
@@ -254,11 +257,11 @@ impl VectorDb for VectorDBService {
 pub async fn run_server(
     vector_db_service: VectorDBService,
     endpoint: ServerEndpoint,
-    root_password: String,
+    keys: Arc<defs::ApiKeyStore>,
 ) -> Result<(), Box<dyn std::error::Error>> {
     event!(Level::INFO, "Starting gRPC server at: {:?}", endpoint);
 
-    let auth_interceptor = interceptors::AuthInterceptor::new(root_password);
+    let auth_interceptor = interceptors::AuthInterceptor::new(keys);
 
     let router = Server::builder()
         .layer(InterceptorLayer::new(auth_interceptor))

crates/grpc/src/tests.rs

@@ -6,7 +6,7 @@ use crate::service::vectordb::{
 use crate::service::{VectorDBService, run_server};
 use crate::utils::ServerEndpoint;
 use api::DbConfig;
-use defs::Similarity;
+use defs::{ApiKeyEntry, ApiKeyRole, ApiKeyStore, Similarity};
 use index::{IndexType, hnsw::HnswConfig, kd_tree::KDTreeConfig};
 use std::net::SocketAddr;
 use std::sync::Arc;
@@ -17,6 +17,7 @@ use tonic::transport::Channel;
 // Inspired from https://github.com/hyperium/tonic/discussions/924#discussioncomment-9854088
 
 const TEST_AUTH_BEARER_TOKEN: &str = "123";
+const TEST_READONLY_BEARER_TOKEN: &str = "456";
 
 fn append_test_auth_header<T>(request: &mut tonic::Request<T>, token: &str) {
     let auth_value = format!("Bearer {}", token);
@@ -41,19 +42,28 @@ async fn start_test_server() -> Result<(SocketAddr, TempDir), Box<dyn std::error
 
     let vector_db_api = api::init_api(db_config)?;
 
+    let keys = Arc::new(ApiKeyStore::new(vec![
+        ApiKeyEntry {
+            name: "test".to_string(),
+            role: ApiKeyRole::ReadWrite,
+            key: TEST_AUTH_BEARER_TOKEN.to_string(),
+        },
+        ApiKeyEntry {
+            name: "test-readonly".to_string(),
+            role: ApiKeyRole::ReadOnly,
+            key: TEST_READONLY_BEARER_TOKEN.to_string(),
+        },
+    ]));
+
     let vector_db_service = VectorDBService::new(Arc::new(vector_db_api), false);
 
     let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await?;
     let listener_addr = listener.local_addr()?;
 
     tokio::spawn(async move {
-        let _ = run_server(
-            vector_db_service,
-            ServerEndpoint::Listener(listener),
-            TEST_AUTH_BEARER_TOKEN.to_string(),
-        )
-        .await
-        .inspect_err(|err| panic!("Could not start test server : {:?}", err));
+        let _ = run_server(vector_db_service, ServerEndpoint::Listener(listener), keys)
+            .await
+            .inspect_err(|err| panic!("Could not start test server : {:?}", err));
     });
 
     Ok((listener_addr, temp_dir))
@@ -254,3 +264,41 @@ async fn test_unauthorized_rpc() {
     // request must fail
     assert!(resp.is_err());
 }
+
+#[tokio::test]
+async fn test_readonly_key_cannot_write() {
+    let (server_addr, _temp_dir) = start_test_server().await.unwrap();
+    let mut client = create_test_client(server_addr).await.unwrap();
+
+    let mut request = tonic::Request::new(InsertVectorRequest {
+        vector: Some(DenseVector {
+            values: vec![1.0, 2.0, 3.0],
+        }),
+        payload: Some(Payload::default()),
+    });
+    append_test_auth_header(&mut request, TEST_READONLY_BEARER_TOKEN);
+
+    let resp = client.insert_vector(request).await;
+
+    assert_eq!(resp.unwrap_err().code(), tonic::Code::PermissionDenied);
+}
+
+#[tokio::test]
+async fn test_readonly_key_can_search() {
+    let (server_addr, _temp_dir) = start_test_server().await.unwrap();
+    let mut client = create_test_client(server_addr).await.unwrap();
+
+    let mut request = tonic::Request::new(SearchRequest {
+        query_vector: Some(DenseVector {
+            values: vec![1.0, 2.0, 3.0],
+        }),
+        similarity: 0,
+        limit: 1,
+        ef: 0,
+    });
+    append_test_auth_header(&mut request, TEST_READONLY_BEARER_TOKEN);
+
+    let resp = client.search_points(request).await;
+
+    assert!(resp.is_ok());
+}