m1guelpf · neelr · Jul 5, 2023
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -25,6 +25,7 @@ aide = { version = "0.11.0", features = ["axum"] }
 tokio = { version = "1.29.1", features = ["full"] }
 axum-jsonschema = { version = "0.6.0", features = ["aide"] }
 tracing-subscriber = { version = "0.3.17", features = ["env-filter"] }
+tower-http = { version = "0.4.1", features = ["validate-request", "auth"] }
 
 [build-dependencies]
 chrono = "0.4.26"
diff --git a/src/main.rs b/src/main.rs
@@ -14,6 +14,11 @@ mod similarity;
 
 #[tokio::main]
 async fn main() -> Result<()> {
+	let api_key : &str = option_env!("API_KEY").unwrap_or("");
+	if api_key == "" {
+		println!("Running in open mode, no API_KEY set");
+	}
+
 	tracing_subscriber::registry()
 		.with(tracing_subscriber::fmt::layer().with_filter(
 			EnvFilter::try_from_default_env().unwrap_or_else(|_| "tinyvector=info".into()),

diff --git a/src/routes/collection.rs b/src/routes/collection.rs
@@ -4,6 +4,7 @@ use aide::axum::{
 };
 use axum::{extract::Path, http::StatusCode, Extension};
 use axum_jsonschema::Json;
+use tower_http::validate_request::ValidateRequestHeaderLayer;
 use schemars::JsonSchema;
 use std::time::Instant;
 
@@ -13,16 +14,26 @@ use crate::{
 	similarity::Distance,
 };
 
+const API_KEY: Option<&str> = option_env!("API_KEY");
+
 pub fn handler() -> ApiRouter {
-	ApiRouter::new().nest(
+	let api_key_layer = ValidateRequestHeaderLayer::bearer(API_KEY.unwrap_or(""));
+
+	let route = ApiRouter::new().nest(
 		"/collections",
 		ApiRouter::new()
 			.api_route("/:collection_name", put(create_collection))
 			.api_route("/:collection_name", post(query_collection))
 			.api_route("/:collection_name", get(get_collection_info))
 			.api_route("/:collection_name", delete(delete_collection))
-			.api_route("/:collection_name/insert", post(insert_into_collection)),
-	)
+			.api_route("/:collection_name/insert", post(insert_into_collection))
+	);
+
+	if API_KEY.is_some() {
+		route.layer(api_key_layer)
+	} else {
+		route
+	}
 }
 
 /// Create a new collection