From 924b7f2f8d55db14715d054e225b5711b93ec678 Mon Sep 17 00:00:00 2001 From: lablans Date: Tue, 3 Dec 2024 15:36:05 +0000 Subject: [PATCH] Refactor/rename auth headers --- src/blaze.rs | 50 ++++++++++++++++++++++++------------ src/config.rs | 57 ++++++++++++++++++++++++++++++++++++----- src/exporter.rs | 30 ++++++++++++++++------ src/intermediate_rep.rs | 26 ++++++++++++++----- src/main.rs | 3 ++- 5 files changed, 127 insertions(+), 39 deletions(-) diff --git a/src/blaze.rs b/src/blaze.rs index 18aa63c..0597af7 100644 --- a/src/blaze.rs +++ b/src/blaze.rs @@ -1,9 +1,13 @@ +use reqwest::header; +use reqwest::header::HeaderName; +use reqwest::header::HeaderValue; use reqwest::StatusCode; use serde::Deserialize; use serde::Serialize; use serde_json::Value; use tracing::{debug, warn, info}; +use crate::config::FocusBackend; use crate::errors::FocusError; use crate::util; use crate::util::get_json_field; @@ -21,28 +25,42 @@ pub struct AstQuery { pub payload: String, } -pub async fn check_availability() -> bool { +pub(crate) struct Blaze; - debug!("Checking Blaze availability..."); +impl FocusBackend for Blaze { + fn make_authheader(apikey: &str) -> Result<(HeaderName, HeaderValue), FocusError> { + let name = header::AUTHORIZATION; + let value = HeaderValue::from_str(apikey) + .map_err(|e| FocusError::ConfigurationError(format!("Invalid value \"{}\" in apikey for blaze backend: {}", apikey, e)))?; + Ok((name, value)) + } - let resp = match CONFIG.client - .get(format!("{}metadata", CONFIG.endpoint_url)) - .send() - .await - { - Ok(response) => response, - Err(e) => { - warn!("Error making Blaze request: {:?}", e); - return false; - } - }; + async fn check_availability() -> bool { + debug!("Checking Blaze availability..."); - if resp.status().is_success() { - return true; + let resp = match CONFIG.client + .get(format!("{}metadata", CONFIG.endpoint_url)) + .send() + .await + { + Ok(response) => response, + Err(e) => { + warn!("Error making Blaze request: {:?}", e); + return false; + } + }; + + if resp.status().is_success() { + return true; + } + false } - false + + // TODO: Refactor other functions into a focus backend following a common trait. } +// TODO: Respect blaze auth header + pub async fn post_library(library: String) -> Result<(), FocusError> { debug!("Creating a Library..."); diff --git a/src/config.rs b/src/config.rs index d45a408..e8bce4c 100644 --- a/src/config.rs +++ b/src/config.rs @@ -3,12 +3,12 @@ use std::path::PathBuf; use beam_lib::AppId; use clap::Parser; -use reqwest::{header::HeaderValue, Url}; +use reqwest::{header::{HeaderName, HeaderValue}, Url}; use once_cell::sync::Lazy; use reqwest::{Certificate, Client, Proxy}; use tracing::{debug, info, warn}; -use crate::errors::FocusError; +use crate::{blaze, errors::FocusError, exporter}; #[derive(clap::ValueEnum, Clone, PartialEq, Debug)] pub enum Obfuscate { @@ -65,7 +65,7 @@ struct CliArgs { #[clap(long, env, value_parser)] beam_app_id_long: String, - /// This applications beam API key + /// This application's beam API key #[clap(long, env, value_parser)] api_key: String, @@ -158,9 +158,18 @@ struct CliArgs { #[clap(long, env, value_parser)] provider_icon: Option, - /// Authorization header + //TODO: Support blaze auth +/* /// Blaze Backend: API key for authentication #[clap(long, env, value_parser)] - auth_header: Option, + backend_blaze_apikey: Option,*/ + + /// Exporter backend: API key for authentication + #[clap(long, env, value_parser)] + backend_exporter_apikey: Option, + + /// AST2SQL transformer backend: API key for authentication + #[clap(long, env, value_parser)] + backend_ast2sql_apikey: Option, /// Postgres connection string #[cfg(feature = "query-sql")] @@ -197,7 +206,9 @@ pub(crate) struct Config { pub client: Client, pub provider: Option, pub provider_icon: Option, - pub auth_header: Option, + pub backend_blaze_authheader: Option<(HeaderName, HeaderValue)>, + pub backend_exporter_authheader: Option<(HeaderName, HeaderValue)>, + pub backend_ast2sql_authheader: Option<(HeaderName, HeaderValue)>, #[cfg(feature = "query-sql")] pub postgres_connection_string: Option, #[cfg(feature = "query-sql")] @@ -219,6 +230,31 @@ impl Config { let client = prepare_reqwest_client(&tls_ca_certificates)?; dbg!(cli_args.endpoint_url.clone()); dbg!(cli_args.blaze_url.clone()); + + let backend_blaze_authheader = { + // if let Some(apikey) = cli_args.backend_blaze_apikey.as_ref() { + // Some(blaze::Blaze::make_authheader(apikey)?) + // } else { + None + // } + }; + + let backend_exporter_authheader = { + if let Some(apikey) = cli_args.backend_exporter_apikey.as_ref() { + Some(exporter::Exporter::make_authheader(apikey)?) + } else { + None + } + }; + + let backend_ast2sql_authheader = { + if let Some(apikey) = cli_args.backend_ast2sql_apikey.as_ref() { + Some(exporter::Exporter::make_authheader(apikey)?) + } else { + None + } + }; + let config = Config { beam_proxy_url: cli_args.beam_proxy_url, beam_app_id_long: AppId::new_unchecked(cli_args.beam_app_id_long), @@ -242,7 +278,9 @@ impl Config { queries_to_cache: cli_args.queries_to_cache, provider: cli_args.provider, provider_icon: cli_args.provider_icon, - auth_header: cli_args.auth_header, + backend_blaze_authheader, + backend_exporter_authheader, + backend_ast2sql_authheader, #[cfg(feature = "query-sql")] postgres_connection_string: cli_args.postgres_connection_string, #[cfg(feature = "query-sql")] @@ -313,3 +351,8 @@ pub fn prepare_reqwest_client(certs: &Vec) -> Result Result<(HeaderName, HeaderValue), FocusError>; + async fn check_availability() -> bool; +} \ No newline at end of file diff --git a/src/exporter.rs b/src/exporter.rs index c3eb932..8049ea0 100644 --- a/src/exporter.rs +++ b/src/exporter.rs @@ -1,11 +1,11 @@ -use reqwest::{header::{self, HeaderMap, HeaderValue}, StatusCode}; +use reqwest::{header::{self, HeaderMap, HeaderName, HeaderValue}, StatusCode}; use serde::Deserialize; use serde::Serialize; use serde_json::Value; use std::str; use tracing::{debug, warn}; -use crate::config::CONFIG; +use crate::config::{FocusBackend, CONFIG}; use crate::errors::FocusError; use crate::util; @@ -35,6 +35,24 @@ const EXECUTE: Params = Params { done: "executed", }; +pub(crate) struct Exporter; + +impl FocusBackend for Exporter { + fn make_authheader(apikey: &str) -> Result<(HeaderName, HeaderValue), FocusError> { + let name = HeaderName::from_static("x-api-key"); + let value = HeaderValue::from_str(apikey) + .map_err(|e| FocusError::ConfigurationError(format!("Invalid value \"{}\" in apikey for exporter backend: {}", apikey, e)))?; + Ok((name, value)) + } + + async fn check_availability() -> bool { + // TODO: Implement + true + } + + // TODO: Refactor other functions into a focus backend following a common trait. +} + pub async fn post_exporter_query(body: &String, task_type: TaskType) -> Result { let Some(exporter_url) = &CONFIG.exporter_url else { return Err(FocusError::MissingExporterEndpoint); @@ -42,12 +60,8 @@ pub async fn post_exporter_query(body: &String, task_type: TaskType) -> Result Result<(header::HeaderName, HeaderValue), FocusError> { + let name = header::AUTHORIZATION; + let value = HeaderValue::from_str(apikey) + .map_err(|e| FocusError::ConfigurationError(format!("Invalid value \"{}\" in apikey for ast2sql backend: {}", apikey, e)))?; + Ok((name, value)) + } + + async fn check_availability() -> bool { + // TODO: Implement + true + } +} + pub async fn post_ast(ast: ast::Ast) -> Result { debug!("Posting AST..."); @@ -26,12 +42,8 @@ pub async fn post_ast(ast: ast::Ast) -> Result { HeaderValue::from_static("application/json"), ); - if let Some(auth_header_value) = CONFIG.auth_header.clone() { - headers.insert( - header::AUTHORIZATION, - HeaderValue::from_str(auth_header_value.as_str()) - .map_err(FocusError::InvalidHeaderValue)?, - ); + if let Some(auth_header) = CONFIG.backend_ast2sql_authheader.clone() { + headers.insert(auth_header.0, auth_header.1); } let resp = CONFIG diff --git a/src/main.rs b/src/main.rs index 6aa85b4..8d83243 100644 --- a/src/main.rs +++ b/src/main.rs @@ -20,6 +20,7 @@ mod db; use base64::engine::general_purpose; use base64::{engine::general_purpose::STANDARD as BASE64, Engine as _}; use beam_lib::{TaskRequest, TaskResult}; +use config::FocusBackend; use futures_util::future::BoxFuture; use futures_util::FutureExt; use laplace_rs::ObfCache; @@ -159,7 +160,7 @@ async fn main_loop() -> ExitCode { }, }; let endpoint_service_available: fn() -> BoxFuture<'static, bool> = match CONFIG.endpoint_type { - EndpointType::Blaze => || blaze::check_availability().boxed(), + EndpointType::Blaze => || blaze::Blaze::check_availability().boxed(), EndpointType::Omop => || async { true }.boxed(), // TODO health check #[cfg(feature = "query-sql")] EndpointType::BlazeAndSql => || blaze::check_availability().boxed(),