Skip to content

Commit

Permalink
refactor(graph-gateway): decouple query settings from auth context (#569
Browse files Browse the repository at this point in the history
)
  • Loading branch information
LNSD authored Jan 31, 2024
1 parent 729f238 commit 3a795d3
Show file tree
Hide file tree
Showing 6 changed files with 40 additions and 20 deletions.
8 changes: 4 additions & 4 deletions graph-gateway/src/client_query.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}
Expand Down
15 changes: 14 additions & 1 deletion graph-gateway/src/client_query/auth.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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)
}
}
}
}
8 changes: 8 additions & 0 deletions graph-gateway/src/client_query/auth/common.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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<USD>,
}

#[cfg(test)]
mod tests {
use super::is_domain_authorized;
Expand Down
15 changes: 0 additions & 15 deletions graph-gateway/src/client_query/auth/context.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -52,11 +50,6 @@ impl FromRef<AuthContext> for subscriptions::AuthContext {
}
}

#[derive(Debug)]
pub struct UserSettings {
pub budget: Option<USD>,
}

impl AuthContext {
pub fn create(
api_keys: Eventual<Ptr<HashMap<String, Arc<APIKey>>>>,
Expand Down Expand Up @@ -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 }
}
}
8 changes: 8 additions & 0 deletions graph-gateway/src/client_query/auth/studio.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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)]
Expand Down Expand Up @@ -111,6 +112,13 @@ pub fn is_domain_authorized(api_key: &Arc<APIKey>, 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<APIKey>) -> QuerySettings {
QuerySettings {
budget: api_key.max_budget,
}
}

pub async fn check_token(
auth: &AuthContext,
api_key: &Arc<APIKey>,
Expand Down
6 changes: 6 additions & 0 deletions graph-gateway/src/client_query/auth/subscriptions.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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,
Expand Down

0 comments on commit 3a795d3

Please sign in to comment.