feature: make api_key reqs faster by using blake3

server/Cargo.toml

@@ -121,6 +121,7 @@ openssl = "0.10.64"
 utoipa-swagger-ui = { version = "6.0.0", features = ["actix-web"] }
 bb8-redis = "0.15.0"
 signal-hook = "0.3.17"
+blake3 = "1.5.1"
 
 
 [build-dependencies]

server/migrations/2024-04-11-052101_make-api-key-faster/down.sql

@@ -0,0 +1,3 @@
+-- This file should undo anything in `up.sql`
+ALTER TABLE user_api_key ALTER COLUMN api_key_hash SET NOT NULL;
+ALTER TABLE user_api_key DROP COLUMN blake3_hash;

server/migrations/2024-04-11-052101_make-api-key-faster/up.sql

@@ -0,0 +1,3 @@
+-- Your SQL goes here
+ALTER TABLE user_api_key ADD COLUMN blake3_hash TEXT;
+ALTER TABLE user_api_key ALTER COLUMN api_key_hash DROP NOT NULL;

server/src/af_middleware/auth_middleware.rs

@@ -43,12 +43,17 @@ where
             let tx_ctx =
                 sentry::TransactionContext::new("middleware", "get dataset, org, and/or user");
             let transaction = sentry::start_transaction(tx_ctx);
-            let org_id_span = transaction.start_child("orgid", "Getting organization id");
-
             let pool = req.app_data::<web::Data<Pool>>().unwrap().to_owned();
 
+            let get_user_span = transaction.start_child("get_user", "Getting user");
+
             let (http_req, pl) = req.parts_mut();
             let user = get_user(http_req, pl).await;
+            if let Some(ref user) = user {
+                req.extensions_mut().insert(user.clone());
+            };
+
+            get_user_span.finish();
 
             let org_id = match req.headers().get("TR-Dataset") {
                 Some(dataset_header) => {
@@ -62,48 +67,46 @@ where
                             ServiceError::BadRequest("Dataset must be valid UUID".to_string())
                         })?;
 
+                    let get_dataset_and_org_span = transaction
+                        .start_child("get_dataset_and_org", "Getting dataset and organization");
+
                     let dataset_org_plan_sub = get_dataset_and_organization_from_dataset_id_query(
                         dataset_id,
                         pool.clone(),
                     )
                     .await?;
 
+                    get_dataset_and_org_span.finish();
+
                     req.extensions_mut().insert(dataset_org_plan_sub.clone());
 
                     dataset_org_plan_sub.organization.organization.id
                 }
-                None => match req.headers().get("TR-Organization") {
-                    Some(org_header) => org_header
-                        .to_str()
-                        .map_err(|_| {
-                            Into::<Error>::into(ServiceError::BadRequest(
-                                "Could not convert Organization to str".to_string(),
-                            ))
-                        })?
-                        .parse::<uuid::Uuid>()
-                        .map_err(|_| {
-                            Into::<Error>::into(ServiceError::BadRequest(
-                                "Could not convert Organization to UUID".to_string(),
-                            ))
-                        })?,
-                    None => {
-                        let (http_req, pl) = req.parts_mut();
-                        let user = get_user(http_req, pl).await;
-                        if let Some(user) = user {
-                            req.extensions_mut().insert(user.clone());
-                        }
+                None => {
+                    if let Some(org_header) = req.headers().get("TR-Organization") {
+                        org_header
+                            .to_str()
+                            .map_err(|_| {
+                                Into::<Error>::into(ServiceError::BadRequest(
+                                    "Could not convert Organization to str".to_string(),
+                                ))
+                            })?
+                            .parse::<uuid::Uuid>()
+                            .map_err(|_| {
+                                Into::<Error>::into(ServiceError::BadRequest(
+                                    "Could not convert Organization to UUID".to_string(),
+                                ))
+                            })?
+                    } else {
                         let res = srv.call(req).await?;
-
-                        org_id_span.finish();
-                        transaction.finish();
-
                         return Ok(res);
                     }
-                },
+                }
             };
 
             if let Some(user) = user {
-                req.extensions_mut().insert(user.clone());
+                let find_user_org_span =
+                    transaction.start_child("find_user_org_role", "Finding user org role");
 
                 let user_org = user
                     .user_orgs
@@ -122,13 +125,14 @@ where
                 }?;
 
                 req.extensions_mut().insert(role);
-            }
 
-            let res = srv.call(req).await?;
+                find_user_org_span.finish();
+            }
 
-            org_id_span.finish();
             transaction.finish();
 
+            let res = srv.call(req).await?;
+
             Ok(res)
         })
     }

server/src/data/models.rs

@@ -1920,33 +1920,36 @@ impl From<ApiKeyRole> for i32 {
     "created_at": "2021-01-01T00:00:00",
     "updated_at": "2021-01-01T00:00:00",
     "role": 1,
+    "blake3_hash": "hash",
 }))]
 #[diesel(table_name = user_api_key)]
 pub struct UserApiKey {
     pub id: uuid::Uuid,
     pub user_id: uuid::Uuid,
-    pub api_key_hash: String,
+    pub api_key_hash: Option<String>,
     pub name: String,
     pub created_at: chrono::NaiveDateTime,
     pub updated_at: chrono::NaiveDateTime,
     pub role: i32,
+    pub blake3_hash: Option<String>,
 }
 
 impl UserApiKey {
     pub fn from_details(
         user_id: uuid::Uuid,
-        api_key_hash: String,
+        blake3_hash: String,
         name: String,
         role: ApiKeyRole,
     ) -> Self {
         UserApiKey {
             id: uuid::Uuid::new_v4(),
             user_id,
-            api_key_hash,
+            api_key_hash: None,
             name,
             created_at: chrono::Utc::now().naive_local(),
             updated_at: chrono::Utc::now().naive_local(),
             role: role.into(),
+            blake3_hash: Some(blake3_hash),
         }
     }
 }

server/src/data/schema.rs

@@ -231,11 +231,12 @@ diesel::table! {
     user_api_key (id) {
         id -> Uuid,
         user_id -> Uuid,
-        api_key_hash -> Text,
+        api_key_hash -> Nullable<Text>,
         name -> Text,
         created_at -> Timestamp,
         updated_at -> Timestamp,
         role -> Int4,
+        blake3_hash -> Nullable<Text>,
     }
 }
 

server/src/operators/user_operator.rs

@@ -163,7 +163,7 @@ pub static SALT: Lazy<String> =
     Lazy::new(|| std::env::var("SALT").unwrap_or_else(|_| "supersecuresalt".to_string()));
 
 #[tracing::instrument]
-pub fn hash_password(password: &str) -> Result<String, ServiceError> {
+pub fn hash_argon2_api_key(password: &str) -> Result<String, ServiceError> {
     let config = Config {
         secret: SECRET_KEY.as_bytes(),
         ..Config::original()
@@ -173,6 +173,11 @@ pub fn hash_password(password: &str) -> Result<String, ServiceError> {
     })
 }
 
+#[tracing::instrument]
+pub fn hash_api_key(password: &str) -> String {
+    blake3::hash(password.as_bytes()).to_string()
+}
+
 #[tracing::instrument(skip(pool))]
 pub async fn set_user_api_key_query(
     user_id: uuid::Uuid,
@@ -181,7 +186,7 @@ pub async fn set_user_api_key_query(
     pool: web::Data<Pool>,
 ) -> Result<String, ServiceError> {
     let raw_api_key = generate_api_key();
-    let hashed_api_key = hash_password(&raw_api_key)?;
+    let hashed_api_key = hash_api_key(&raw_api_key);
 
     let mut conn = pool.get().await.unwrap();
 
@@ -206,7 +211,7 @@ pub async fn get_user_from_api_key_query(
     use crate::data::schema::user_organizations::dsl as user_organizations_columns;
     use crate::data::schema::users::dsl as users_columns;
 
-    let api_key_hash = hash_password(api_key)?;
+    let api_key_hash = hash_api_key(api_key);
 
     let mut conn = pool.get().await.unwrap();
 
@@ -218,7 +223,7 @@ pub async fn get_user_from_api_key_query(
                     .on(organization_columns::id.eq(user_organizations_columns::organization_id)),
             )
             .inner_join(user_api_key_columns::user_api_key)
-            .filter(user_api_key_columns::api_key_hash.eq(api_key_hash))
+            .filter(user_api_key_columns::blake3_hash.eq(api_key_hash.clone()))
             .select((
                 User::as_select(),
                 UserOrganization::as_select(),
@@ -257,7 +262,67 @@ pub async fn get_user_from_api_key_query(
                 .collect::<Vec<Organization>>();
             Ok(SlimUser::from_details(user, user_orgs, orgs))
         }
-        None => Err(ServiceError::BadRequest("User not found".to_string())),
+        None => {
+            let argon2_hash = hash_argon2_api_key(api_key)?;
+
+            let user_orgs_orgs: Vec<(User, UserOrganization, Organization, UserApiKey)> =
+                users_columns::users
+                    .inner_join(user_organizations_columns::user_organizations)
+                    .inner_join(organization_columns::organizations.on(
+                        organization_columns::id.eq(user_organizations_columns::organization_id),
+                    ))
+                    .inner_join(user_api_key_columns::user_api_key)
+                    .filter(user_api_key_columns::api_key_hash.eq(argon2_hash.clone()))
+                    .select((
+                        User::as_select(),
+                        UserOrganization::as_select(),
+                        Organization::as_select(),
+                        UserApiKey::as_select(),
+                    ))
+                    .load::<(User, UserOrganization, Organization, UserApiKey)>(&mut conn)
+                    .await
+                    .map_err(|_| ServiceError::BadRequest("API Key Incorrect".to_string()))?;
+
+            match user_orgs_orgs.get(0) {
+                Some(first_user_org) => {
+                    let user = first_user_org.0.clone();
+                    let mut user_orgs = user_orgs_orgs
+                        .iter()
+                        .map(|user_org| user_org.1.clone())
+                        .collect::<Vec<UserOrganization>>();
+
+                    user_orgs.iter_mut().for_each(|user_org| {
+                        if user_orgs_orgs
+                            .iter()
+                            .find(|user_org_org| user_org_org.1.id == user_org.id)
+                            .unwrap()
+                            .3
+                            .role
+                            == 0
+                        {
+                            user_org.role = 0;
+                        }
+                    });
+
+                    let orgs = user_orgs_orgs
+                        .iter()
+                        .map(|user_org_org| user_org_org.2.clone())
+                        .collect::<Vec<Organization>>();
+
+                    diesel::update(
+                        user_api_key_columns::user_api_key
+                            .filter(user_api_key_columns::api_key_hash.eq(argon2_hash)),
+                    )
+                    .set(user_api_key_columns::blake3_hash.eq(api_key_hash))
+                    .execute(&mut conn)
+                    .await
+                    .map_err(|_| ServiceError::BadRequest("Error updating api key".to_string()))?;
+
+                    Ok(SlimUser::from_details(user, user_orgs, orgs))
+                }
+                None => Err(ServiceError::BadRequest("API Key Incorrect".to_string())),
+            }
+        }
     }
 }
 
@@ -432,7 +497,7 @@ pub async fn create_default_user(api_key: &str, pool: web::Data<Pool>) -> Result
     use crate::data::schema::user_organizations::dsl as user_organizations_columns;
     use crate::data::schema::users::dsl as users_columns;
 
-    let api_key_hash = hash_password(api_key)?;
+    let api_key_hash = hash_api_key(api_key);
 
     let mut conn = pool.get_ref().get().await.unwrap();