From 3a795d3b3624a174239077ed8219830016579117 Mon Sep 17 00:00:00 2001 From: Lorenzo Delgado Date: Wed, 31 Jan 2024 16:36:20 +0100 Subject: [PATCH] refactor(graph-gateway): decouple query settings from auth context (#569) --- graph-gateway/src/client_query.rs | 8 ++++---- graph-gateway/src/client_query/auth.rs | 15 ++++++++++++++- graph-gateway/src/client_query/auth/common.rs | 8 ++++++++ graph-gateway/src/client_query/auth/context.rs | 15 --------------- graph-gateway/src/client_query/auth/studio.rs | 8 ++++++++ .../src/client_query/auth/subscriptions.rs | 6 ++++++ 6 files changed, 40 insertions(+), 20 deletions(-) diff --git a/graph-gateway/src/client_query.rs b/graph-gateway/src/client_query.rs index cfe1dfe9..d2adc2ce 100644 --- a/graph-gateway/src/client_query.rs +++ b/graph-gateway/src/client_query.rs @@ -274,12 +274,12 @@ async fn handle_client_query_inner( .grt_per_usd .ok_or_else(|| Error::Internal(anyhow!("missing exchange rate")))?; let mut budget = GRT(ctx.budgeter.query_fees_target.0 * grt_per_usd.0); - let user_settings = ctx.auth_handler.query_settings(&auth); - if let Some(user_budget) = user_settings.budget { + let query_settings = auth.query_settings(); + if let Some(user_budget) = query_settings.budget { // Security: Consumers can and will set their budget to unreasonably high values. // This `.min` prevents the budget from being set far beyond what it would be - // automatically. The reason this is important is because sometimes queries are - // subsidized and we would be at-risk to allow arbitrarily high values. + // automatically. The reason this is important is that sometimes queries are + // subsidized, and we would be at-risk to allow arbitrarily high values. let max_budget = GRT(budget.0 * UDecimal18::from(10)); budget = GRT(user_budget.0 * grt_per_usd.0).min(max_budget); } diff --git a/graph-gateway/src/client_query/auth.rs b/graph-gateway/src/client_query/auth.rs index b9ea5682..9292fb87 100644 --- a/graph-gateway/src/client_query/auth.rs +++ b/graph-gateway/src/client_query/auth.rs @@ -5,7 +5,8 @@ use thegraph::types::{DeploymentId, SubgraphId}; use crate::subgraph_studio::APIKey; -pub use self::context::{AuthContext, UserSettings}; +pub use self::common::QuerySettings; +pub use self::context::AuthContext; mod common; mod context; @@ -52,4 +53,16 @@ impl AuthToken { } } } + + /// Get the user settings associated with the auth token. + /// + /// See [`QuerySettings`] for more information. + pub fn query_settings(&self) -> QuerySettings { + match self { + AuthToken::StudioApiKey(api_key) => studio::get_query_settings(api_key), + AuthToken::SubscriptionsAuthToken(token_claims) => { + subscriptions::get_query_settings(token_claims) + } + } + } } diff --git a/graph-gateway/src/client_query/auth/common.rs b/graph-gateway/src/client_query/auth/common.rs index 466b6f0b..47173bd5 100644 --- a/graph-gateway/src/client_query/auth/common.rs +++ b/graph-gateway/src/client_query/auth/common.rs @@ -2,6 +2,8 @@ use std::sync::Arc; use thegraph::types::{DeploymentId, SubgraphId}; +use gateway_common::types::USD; + use crate::topology::Deployment; /// Check if the given deployments are authorized by the given authorized deployments. @@ -65,6 +67,12 @@ pub fn is_domain_authorized(authorized: &[&str], origin: &str) -> bool { .any(|pattern| match_domain(pattern, origin)) } +/// User query settings associated with an auth token. +#[derive(Debug)] +pub struct QuerySettings { + pub budget: Option, +} + #[cfg(test)] mod tests { use super::is_domain_authorized; diff --git a/graph-gateway/src/client_query/auth/context.rs b/graph-gateway/src/client_query/auth/context.rs index 95b37a40..7071238d 100644 --- a/graph-gateway/src/client_query/auth/context.rs +++ b/graph-gateway/src/client_query/auth/context.rs @@ -8,8 +8,6 @@ use axum::extract::FromRef; use eventuals::{Eventual, EventualExt, Ptr}; use tokio::sync::RwLock; -use gateway_common::types::USD; - use crate::subgraph_studio::APIKey; use crate::subscriptions::Subscription; use crate::topology::Deployment; @@ -52,11 +50,6 @@ impl FromRef for subscriptions::AuthContext { } } -#[derive(Debug)] -pub struct UserSettings { - pub budget: Option, -} - impl AuthContext { pub fn create( api_keys: Eventual>>>, @@ -133,12 +126,4 @@ impl AuthContext { } } } - - pub fn query_settings(&self, token: &AuthToken) -> UserSettings { - let budget = match token { - AuthToken::StudioApiKey(api_key) => api_key.max_budget, - _ => None, - }; - UserSettings { budget } - } } diff --git a/graph-gateway/src/client_query/auth/studio.rs b/graph-gateway/src/client_query/auth/studio.rs index f6d4a25d..f99b1b73 100644 --- a/graph-gateway/src/client_query/auth/studio.rs +++ b/graph-gateway/src/client_query/auth/studio.rs @@ -9,6 +9,7 @@ use crate::subgraph_studio::{APIKey, QueryStatus}; use crate::topology::Deployment; use super::common; +use super::common::QuerySettings; /// Errors that may occur when parsing a Studio API key. #[derive(Debug, thiserror::Error)] @@ -111,6 +112,13 @@ pub fn is_domain_authorized(api_key: &Arc, domain: &str) -> bool { common::is_domain_authorized(allowed_domains, domain) } +/// Get the user settings associated with the API key. +pub fn get_query_settings(api_key: &Arc) -> QuerySettings { + QuerySettings { + budget: api_key.max_budget, + } +} + pub async fn check_token( auth: &AuthContext, api_key: &Arc, diff --git a/graph-gateway/src/client_query/auth/subscriptions.rs b/graph-gateway/src/client_query/auth/subscriptions.rs index f76e26e1..691eaf2f 100644 --- a/graph-gateway/src/client_query/auth/subscriptions.rs +++ b/graph-gateway/src/client_query/auth/subscriptions.rs @@ -12,6 +12,7 @@ use crate::subscriptions::Subscription; use crate::topology::Deployment; use super::common; +use super::common::QuerySettings; /// App state (a.k.a [Context](crate::client_query::Context)) sub-state. pub struct AuthContext { @@ -91,6 +92,11 @@ pub fn is_domain_authorized(auth_token: &AuthTokenClaims, domain: &str) -> bool common::is_domain_authorized(&allowed_domains, domain) } +/// Get the user settings associated with the auth token. +pub fn get_query_settings(_auth: &AuthTokenClaims) -> QuerySettings { + QuerySettings { budget: None } +} + pub async fn check_token( auth: &AuthContext, auth_token: &AuthTokenClaims,