refactor: moved query auth function to utils

Author: parmesant

src/correlation/correlation_utils.rs

@@ -18,9 +18,8 @@
 
 use std::collections::HashSet;
 
-use actix_web::http::header::ContentType;
+use actix_web::{http::header::ContentType, Error};
 use chrono::Utc;
-use correlation_utils::user_auth_for_query;
 use datafusion::error::DataFusionError;
 use http::StatusCode;
 use itertools::Itertools;
@@ -31,12 +30,15 @@ use tokio::sync::RwLock;
 use tracing::{error, trace, warn};
 
 use crate::{
-    handlers::http::rbac::RBACError, option::CONFIG, query::QUERY_SESSION, rbac::map::SessionKey,
-    storage::ObjectStorageError, users::filters::FilterQuery, utils::get_hash,
+    handlers::http::rbac::RBACError,
+    option::CONFIG,
+    query::QUERY_SESSION,
+    rbac::{map::SessionKey, Users},
+    storage::ObjectStorageError,
+    users::filters::FilterQuery,
+    utils::{get_hash, user_auth_for_query},
 };
 
-pub mod correlation_utils;
-
 pub static CORRELATIONS: Lazy<Correlation> = Lazy::new(Correlation::default);
 
 #[derive(Debug, Default)]
@@ -77,11 +79,15 @@ impl Correlation {
         let correlations = self.0.read().await.iter().cloned().collect_vec();
 
         let mut user_correlations = vec![];
+        let permissions = Users.get_permissions(session_key);
+
         for c in correlations {
-            if user_auth_for_query(session_key, &c.table_configs)
-                .await
-                .is_ok()
-            {
+            let tables = &c
+                .table_configs
+                .iter()
+                .map(|t| t.table_name.clone())
+                .collect_vec();
+            if user_auth_for_query(&permissions, tables).is_ok() {
                 user_correlations.push(c);
             }
         }
@@ -220,7 +226,14 @@ impl CorrelationRequest {
         }
 
         // check if user has access to table
-        user_auth_for_query(session_key, &self.table_configs).await?;
+        let permissions = Users.get_permissions(session_key);
+        let tables = &self
+            .table_configs
+            .iter()
+            .map(|t| t.table_name.clone())
+            .collect_vec();
+
+        user_auth_for_query(&permissions, tables)?;
 
         // to validate table config, we need to check whether the mentioned fields
         // are present in the table or not
@@ -271,6 +284,8 @@ pub enum CorrelationError {
     Unauthorized,
     #[error("DataFusion Error: {0}")]
     DataFusion(#[from] DataFusionError),
+    #[error("{0}")]
+    ActixError(#[from] Error),
 }
 
 impl actix_web::ResponseError for CorrelationError {
@@ -283,6 +298,7 @@ impl actix_web::ResponseError for CorrelationError {
             Self::AnyhowError(_) => StatusCode::INTERNAL_SERVER_ERROR,
             Self::Unauthorized => StatusCode::BAD_REQUEST,
             Self::DataFusion(_) => StatusCode::INTERNAL_SERVER_ERROR,
+            Self::ActixError(_) => StatusCode::BAD_REQUEST,
         }
     }
 

src/correlation/mod.rs

@@ -34,9 +34,7 @@ use tonic::transport::{Identity, Server, ServerTlsConfig};
 use tonic_web::GrpcWebLayer;
 
 use crate::handlers::http::cluster::get_ingestor_info;
-use crate::handlers::http::query::{
-    authorize_and_set_filter_tags, into_query, update_schema_when_distributed,
-};
+use crate::handlers::http::query::{into_query, update_schema_when_distributed};
 use crate::handlers::livetail::cross_origin_config;
 use crate::metrics::QUERY_EXECUTE_TIME;
 use crate::option::CONFIG;
@@ -46,6 +44,7 @@ use crate::utils::arrow::flight::{
     send_to_ingester,
 };
 use crate::utils::time::TimeRange;
+use crate::utils::user_auth_for_query;
 use arrow_flight::{
     flight_service_server::FlightService, Action, ActionType, Criteria, Empty, FlightData,
     FlightDescriptor, FlightInfo, HandshakeRequest, HandshakeResponse, PutResult, SchemaAsIpc,
@@ -162,7 +161,7 @@ impl FlightService for AirServiceImpl {
             .map_err(|err| Status::internal(err.to_string()))?;
 
         // map payload to query
-        let mut query = into_query(&ticket, &session_state, time_range)
+        let query = into_query(&ticket, &session_state, time_range)
             .await
             .map_err(|_| Status::internal("Failed to parse query"))?;
 
@@ -202,7 +201,7 @@ impl FlightService for AirServiceImpl {
             rbac::Response::Authorized => (),
             rbac::Response::UnAuthorized => {
                 return Err(Status::permission_denied(
-                    "user is not authenticated to access this resource",
+                    "user is not authorized to access this resource",
                 ))
             }
             rbac::Response::ReloadRequired => {
@@ -212,7 +211,7 @@ impl FlightService for AirServiceImpl {
 
         let permissions = Users.get_permissions(&key);
 
-        authorize_and_set_filter_tags(&mut query, permissions, &streams).map_err(|_| {
+        user_auth_for_query(&permissions, &streams).map_err(|_| {
             Status::permission_denied("User Does not have permission to access this")
         })?;
         let time = Instant::now();

src/handlers/airplane.rs

@@ -18,17 +18,17 @@
 
 use actix_web::{web, HttpRequest, Responder};
 use bytes::Bytes;
+use itertools::Itertools;
 use relative_path::RelativePathBuf;
 
+use crate::rbac::Users;
+use crate::utils::user_auth_for_query;
 use crate::{
     option::CONFIG, storage::CORRELATIONS_ROOT_DIRECTORY,
     utils::actix::extract_session_key_from_req,
 };
 
-use crate::correlation::{
-    correlation_utils::user_auth_for_query, CorrelationConfig, CorrelationError,
-    CorrelationRequest, CORRELATIONS,
-};
+use crate::correlation::{CorrelationConfig, CorrelationError, CorrelationRequest, CORRELATIONS};
 
 pub async fn list(req: HttpRequest) -> Result<impl Responder, CorrelationError> {
     let session_key = extract_session_key_from_req(&req)
@@ -52,14 +52,17 @@ pub async fn get(req: HttpRequest) -> Result<impl Responder, CorrelationError> {
 
     let correlation = CORRELATIONS.get_correlation_by_id(correlation_id).await?;
 
-    if user_auth_for_query(&session_key, &correlation.table_configs)
-        .await
-        .is_ok()
-    {
-        Ok(web::Json(correlation))
-    } else {
-        Err(CorrelationError::Unauthorized)
-    }
+    let permissions = Users.get_permissions(&session_key);
+
+    let tables = &correlation
+        .table_configs
+        .iter()
+        .map(|t| t.table_name.clone())
+        .collect_vec();
+
+    user_auth_for_query(&permissions, tables)?;
+
+    Ok(web::Json(correlation))
 }
 
 pub async fn post(req: HttpRequest, body: Bytes) -> Result<impl Responder, CorrelationError> {
@@ -93,7 +96,14 @@ pub async fn modify(req: HttpRequest, body: Bytes) -> Result<impl Responder, Cor
 
     // validate whether user has access to this correlation object or not
     let correlation = CORRELATIONS.get_correlation_by_id(correlation_id).await?;
-    user_auth_for_query(&session_key, &correlation.table_configs).await?;
+    let permissions = Users.get_permissions(&session_key);
+    let tables = &correlation
+        .table_configs
+        .iter()
+        .map(|t| t.table_name.clone())
+        .collect_vec();
+
+    user_auth_for_query(&permissions, tables)?;
 
     let correlation_request: CorrelationRequest = serde_json::from_slice(&body)?;
     correlation_request.validate(&session_key).await?;
@@ -122,7 +132,14 @@ pub async fn delete(req: HttpRequest) -> Result<impl Responder, CorrelationError
     let correlation = CORRELATIONS.get_correlation_by_id(correlation_id).await?;
 
     // validate user's query auth
-    user_auth_for_query(&session_key, &correlation.table_configs).await?;
+    let permissions = Users.get_permissions(&session_key);
+    let tables = &correlation
+        .table_configs
+        .iter()
+        .map(|t| t.table_name.clone())
+        .collect_vec();
+
+    user_auth_for_query(&permissions, tables)?;
 
     // Delete from disk
     let store = CONFIG.storage().get_object_store();

src/handlers/http/correlation.rs

@@ -41,13 +41,13 @@ use crate::option::{Mode, CONFIG};
 use crate::query::error::ExecuteError;
 use crate::query::Query as LogicalQuery;
 use crate::query::{TableScanVisitor, QUERY_SESSION};
-use crate::rbac::role::{Action, Permission};
 use crate::rbac::Users;
 use crate::response::QueryResponse;
 use crate::storage::object_storage::commit_schema_to_storage;
 use crate::storage::ObjectStorageError;
 use crate::utils::actix::extract_session_key_from_req;
 use crate::utils::time::{TimeParseError, TimeRange};
+use crate::utils::user_auth_for_query;
 
 use super::modal::utils::logstream_utils::create_stream_and_schema_from_storage;
 
@@ -91,7 +91,7 @@ pub async fn query(req: HttpRequest, query_request: Query) -> Result<impl Respon
 
     let tables = visitor.into_inner();
     update_schema_when_distributed(&tables).await?;
-    let mut query: LogicalQuery = into_query(&query_request, &session_state, time_range).await?;
+    let query: LogicalQuery = into_query(&query_request, &session_state, time_range).await?;
 
     let creds = extract_session_key_from_req(&req)?;
     let permissions = Users.get_permissions(&creds);
@@ -100,7 +100,7 @@ pub async fn query(req: HttpRequest, query_request: Query) -> Result<impl Respon
         .first_table_name()
         .ok_or_else(|| QueryError::MalformedQuery("No table name found in query"))?;
 
-    authorize_and_set_filter_tags(&mut query, permissions, &tables)?;
+    user_auth_for_query(&permissions, &tables)?;
 
     let time = Instant::now();
     let (records, fields) = query.execute(table_name.clone()).await?;
@@ -153,49 +153,6 @@ pub async fn create_streams_for_querier() {
     }
 }
 
-pub fn authorize_and_set_filter_tags(
-    query: &mut LogicalQuery,
-    permissions: Vec<Permission>,
-    tables: &[String],
-) -> Result<(), QueryError> {
-    for table_name in tables.iter() {
-        // check authorization of this query if it references physical table;
-        let mut authorized = false;
-
-        let mut tags = Vec::new();
-
-        // in permission check if user can run query on the stream.
-        // also while iterating add any filter tags for this stream
-        for permission in &permissions {
-            match permission {
-                Permission::Stream(Action::All, _) => {
-                    authorized = true;
-                    break;
-                }
-                Permission::StreamWithTag(Action::Query, ref stream, tag)
-                    if stream == table_name || stream == "*" =>
-                {
-                    authorized = true;
-                    if let Some(tag) = tag {
-                        tags.push(tag.clone())
-                    }
-                }
-                _ => (),
-            }
-        }
-
-        if !authorized {
-            return Err(QueryError::Unauthorized);
-        }
-
-        if !tags.is_empty() {
-            query.filter_tag = Some(tags)
-        }
-    }
-
-    Ok(())
-}
-
 impl FromRequest for Query {
     type Error = actix_web::Error;
     type Future = Pin<Box<dyn Future<Output = Result<Self, Self::Error>>>>;