key type to accomodate query via api key

Author: nikhilsinhaparseable

src/apikeys.rs

@@ -39,24 +39,42 @@ pub struct ApiKeyStore {
     pub keys: RwLock<HashMap<String, HashMap<Ulid, ApiKey>>>,
 }
 
+/// Type of API key, determining how it can be used.
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
+#[serde(rename_all = "lowercase")]
+pub enum KeyType {
+    /// Used as a substitute for basic auth on ingestion endpoints
+    Ingestion,
+    /// Used as a substitute for basic auth on query endpoints (global query access)
+    Query,
+}
+
 #[derive(Debug, Clone, Serialize, Deserialize)]
 #[serde(rename_all = "camelCase")]
 pub struct ApiKey {
     pub key_id: Ulid,
     pub api_key: String,
     pub key_name: String,
+    #[serde(default = "default_key_type")]
+    pub key_type: KeyType,
     pub created_by: String,
     pub created_at: DateTime<Utc>,
     pub modified_at: DateTime<Utc>,
     #[serde(default)]
     pub tenant: Option<String>,
 }
 
+fn default_key_type() -> KeyType {
+    KeyType::Ingestion
+}
+
 /// Request body for creating a new API key
 #[derive(Debug, Deserialize)]
 #[serde(rename_all = "camelCase")]
 pub struct CreateApiKeyRequest {
     pub key_name: String,
+    #[serde(default = "default_key_type")]
+    pub key_type: KeyType,
 }
 
 /// Response for list keys (api_key masked to last 4 chars)
@@ -66,18 +84,25 @@ pub struct ApiKeyListEntry {
     pub key_id: Ulid,
     pub api_key: String,
     pub key_name: String,
+    pub key_type: KeyType,
     pub created_by: String,
     pub created_at: DateTime<Utc>,
     pub modified_at: DateTime<Utc>,
 }
 
 impl ApiKey {
-    pub fn new(key_name: String, created_by: String, tenant: Option<String>) -> Self {
+    pub fn new(
+        key_name: String,
+        key_type: KeyType,
+        created_by: String,
+        tenant: Option<String>,
+    ) -> Self {
         let now = Utc::now();
         Self {
             key_id: Ulid::new(),
             api_key: uuid::Uuid::new_v4().to_string(),
             key_name,
+            key_type,
             created_by,
             created_at: now,
             modified_at: now,
@@ -96,6 +121,7 @@ impl ApiKey {
             key_id: self.key_id,
             api_key: masked,
             key_name: self.key_name.clone(),
+            key_type: self.key_type,
             created_by: self.created_by.clone(),
             created_at: self.created_at,
             modified_at: self.modified_at,
@@ -228,24 +254,24 @@ impl ApiKeyStore {
             .ok_or_else(|| ApiKeyError::KeyNotFound(key_id.to_string()))
     }
 
-    /// Validate an API key for ingestion. Returns true if the key is valid.
+    /// Validate an API key against a required key type. Returns true if the
+    /// key is valid AND its type matches the required type.
     /// For multi-tenant: checks the key belongs to the specified tenant.
     /// For single-tenant: checks the key exists globally.
-    pub async fn validate_key(&self, api_key_value: &str, tenant_id: &Option<String>) -> bool {
+    pub async fn validate_key(
+        &self,
+        api_key_value: &str,
+        tenant_id: &Option<String>,
+        required_type: KeyType,
+    ) -> bool {
         let map = self.keys.read().await;
-        if let Some(tenant_id) = tenant_id {
-            // Multi-tenant: check keys for the specific tenant
-            if let Some(tenant_keys) = map.get(tenant_id) {
-                return tenant_keys.values().any(|k| k.api_key == api_key_value);
-            }
-            false
-        } else {
-            // Single-tenant: check keys under DEFAULT_TENANT
-            if let Some(tenant_keys) = map.get(DEFAULT_TENANT) {
-                return tenant_keys.values().any(|k| k.api_key == api_key_value);
-            }
-            false
+        let tenant = tenant_id.as_deref().unwrap_or(DEFAULT_TENANT);
+        if let Some(tenant_keys) = map.get(tenant) {
+            return tenant_keys
+                .values()
+                .any(|k| k.api_key == api_key_value && k.key_type == required_type);
         }
+        false
     }
 
     /// Insert an API key directly into memory (used for sync from prism)

src/handlers/http/middleware.rs

@@ -156,29 +156,69 @@ where
         let mut header_error = None;
         let user_and_tenant_id = get_user_and_tenant(&self.action, &mut req, &mut header_error);
 
-        // If X-API-KEY header is present for ingestion, short-circuit normal auth
-        if let Some(api_key) = extract_api_key(&req, &self.action) {
+        // If X-API-KEY header is present and the action supports API key auth,
+        // short-circuit the normal auth flow.
+        if let Some(api_key) = extract_api_key(&req)
+            && let Some(required_type) = api_key_type_for_action(&self.action)
+        {
+            struct SessionCleanupGuard(Option<Ulid>);
+            impl Drop for SessionCleanupGuard {
+                fn drop(&mut self) {
+                    if let Some(sid) = self.0 {
+                        mut_sessions().remove_session(&SessionKey::SessionId(sid));
+                    }
+                }
+            }
+
             let suspension = check_suspension(req.request(), self.action);
             let tenant_id = req
                 .headers()
                 .get(TENANT_ID)
                 .and_then(|v| v.to_str().ok())
                 .map(String::from);
+
+            // For Query keys, set up a short-lived session with global reader
+            // permissions so the query handler's per-stream auth passes.
+            let query_session_id = if required_type == crate::apikeys::KeyType::Query {
+                let session_id = Ulid::new();
+                let session_key = SessionKey::SessionId(session_id);
+                mut_sessions().track_new(
+                    format!("api-key:{session_id}"),
+                    session_key.clone(),
+                    Utc::now() + TimeDelta::minutes(5),
+                    query_api_key_permissions(),
+                    &tenant_id,
+                );
+                req.extensions_mut().insert(session_key);
+                Some(session_id)
+            } else {
+                None
+            };
+
             let fut = self.service.call(req);
 
             return Box::pin(async move {
-                if let Some(err) = header_error {
-                    return Err(err);
-                }
-                if let rbac::Response::Suspended(msg) = suspension {
-                    return Err(ErrorBadRequest(msg));
-                }
+                let _guard = SessionCleanupGuard(query_session_id);
+                let result: Result<ServiceResponse<B>, Error> = async {
+                    if let Some(err) = header_error {
+                        return Err(err);
+                    }
+                    if let rbac::Response::Suspended(msg) = suspension {
+                        return Err(ErrorBadRequest(msg));
+                    }
 
-                use crate::apikeys::API_KEYS;
-                if API_KEYS.validate_key(&api_key, &tenant_id).await {
-                    return fut.await;
+                    use crate::apikeys::API_KEYS;
+                    if API_KEYS
+                        .validate_key(&api_key, &tenant_id, required_type)
+                        .await
+                    {
+                        return fut.await;
+                    }
+                    Err(ErrorUnauthorized("Invalid API key"))
                 }
-                Err(ErrorUnauthorized("Invalid API key"))
+                .await;
+
+                result
             });
         }
 
@@ -263,18 +303,35 @@ fn extract_kinesis_headers(req: &mut ServiceRequest) {
     }
 }
 
-/// Extract X-API-KEY header value if present and action is Ingest.
-fn extract_api_key(req: &ServiceRequest, action: &Action) -> Option<String> {
-    if action.eq(&Action::Ingest) {
-        req.headers()
-            .get("x-api-key")
-            .and_then(|v| v.to_str().ok())
-            .map(String::from)
-    } else {
-        None
+/// Extract X-API-KEY header value if present (independent of action).
+fn extract_api_key(req: &ServiceRequest) -> Option<String> {
+    req.headers()
+        .get("x-api-key")
+        .and_then(|v| v.to_str().ok())
+        .map(String::from)
+}
+
+/// Map an Action to the KeyType required for API key auth on that action.
+/// Returns None if the action doesn't support API key auth.
+fn api_key_type_for_action(action: &Action) -> Option<crate::apikeys::KeyType> {
+    use crate::apikeys::KeyType;
+    match action {
+        Action::Ingest => Some(KeyType::Ingestion),
+        Action::Query => Some(KeyType::Query),
+        _ => None,
     }
 }
 
+/// Build the set of permissions granted to a Query API key session.
+/// Equivalent to the Reader privilege with no resource restriction
+/// (global query access across all streams in the tenant).
+fn query_api_key_permissions() -> Vec<crate::rbac::role::Permission> {
+    crate::rbac::role::RoleBuilder::from(&crate::rbac::role::model::DefaultPrivilege::Reader {
+        resource: None,
+    })
+    .build()
+}
+
 #[inline]
 fn get_user_and_tenant(
     action: &Action,