refactor(graph-gateway): add query selector extractor and check (#561)

graph-gateway/src/client_query.rs

@@ -8,7 +8,7 @@ use anyhow::anyhow;
 use axum::extract::OriginalUri;
 use axum::{
     body::Bytes,
-    extract::{Path, State},
+    extract::State,
     http::{HeaderMap, Response, StatusCode},
     Extension,
 };
@@ -20,7 +20,7 @@ use prost::bytes::Buf;
 use rand::{rngs::SmallRng, SeedableRng as _};
 use serde::Deserialize;
 use serde_json::value::RawValue;
-use thegraph::types::{attestation, BlockPointer, DeploymentId, SubgraphId};
+use thegraph::types::{attestation, BlockPointer, DeploymentId};
 use tokio::sync::mpsc;
 use toolshed::buffer_queue::QueueWriter;
 use tracing::Instrument;
@@ -44,6 +44,7 @@ use indexer_selection::{
 };
 
 use crate::block_constraints::{block_constraints, make_query_deterministic};
+use crate::client_query::query_selector::QuerySelector;
 use crate::indexer_client::{check_block_error, IndexerClient, ResponsePayload};
 use crate::reports::{self, serialize_attestation, KafkaClient};
 use crate::topology::{Deployment, GraphNetwork, Subgraph};
@@ -61,6 +62,7 @@ mod graphql;
 mod l2_forwarding;
 pub mod legacy_auth_adapter;
 pub mod query_id;
+mod query_selector;
 pub mod query_tracing;
 pub mod require_auth;
 
@@ -74,28 +76,32 @@ pub async fn handle_query(
     State(ctx): State<Context>,
     Extension(auth): Extension<AuthToken>,
     OriginalUri(original_uri): OriginalUri,
-    Path(params): Path<BTreeMap<String, String>>,
+    selector: QuerySelector,
     headers: HeaderMap,
     payload: Bytes,
 ) -> Response<String> {
-    let span = tracing::span::Span::current();
-
     let start_time = Instant::now();
     let timestamp = unix_timestamp();
 
-    let resolved_deployments = resolve_subgraph_deployments(&ctx.network, &params).await;
+    // Check if the query selector is authorized by the auth token
+    match &selector {
+        QuerySelector::Subgraph(id) => {
+            if !auth.is_subgraph_authorized(id) {
+                return graphql::error_response(Error::Auth(anyhow!(
+                    "Subgraph not authorized by user"
+                )));
+            }
+        }
+        QuerySelector::Deployment(id) => {
+            if !auth.is_deployment_authorized(id) {
+                return graphql::error_response(Error::Auth(anyhow!(
+                    "Deployment not authorized by user"
+                )));
+            }
+        }
+    }
 
-    // This is very useful for investigating gateway logs in production
-    let selector = match &resolved_deployments {
-        Ok((_, Some(subgraph))) => subgraph.id.to_string(),
-        Ok((deployments, None)) => deployments
-            .iter()
-            .map(|d| d.id.to_string())
-            .collect::<Vec<_>>()
-            .join(","),
-        Err(_) => "".to_string(),
-    };
-    span.record("selector", tracing::field::display(selector));
+    let resolved_deployments = resolve_subgraph_deployments(&ctx.network, &selector).await;
 
     // We only resolve a subgraph when a subgraph ID is given as a URL param.
     let subgraph = resolved_deployments
@@ -126,7 +132,7 @@ pub async fn handle_query(
     let result = match resolved_deployments {
         Ok((deployments, _)) => {
             handle_client_query_inner(&ctx, deployments, payload, auth)
-                .instrument(span.clone())
+                .in_current_span()
                 .await
         }
         Err(subgraph_resolution_err) => Err(subgraph_resolution_err),
@@ -173,36 +179,27 @@ pub async fn handle_query(
 
 async fn resolve_subgraph_deployments(
     network: &GraphNetwork,
-    params: &BTreeMap<String, String>,
+    selector: &QuerySelector,
 ) -> Result<(Vec<Arc<Deployment>>, Option<Subgraph>), Error> {
-    if let Some(id) = params.get("subgraph_id") {
-        // Parse the subgraph ID
-        let id: SubgraphId = id
-            .parse()
-            .map_err(|_| Error::SubgraphNotFound(anyhow!("invalid subgraph ID: {id}")))?;
-
-        // Get the subgraph by ID
-        let subgraph = network
-            .subgraph_by_id(&id)
-            .ok_or_else(|| Error::SubgraphNotFound(anyhow!("{id}")))?;
-
-        // Get the subgraph's deployments (versions = deployments)
-        let versions = subgraph.deployments.clone();
-        Ok((versions, Some(subgraph)))
-    } else if let Some(id) = params.get("deployment_id") {
-        // Parse the deployment ID
-        let id: DeploymentId = id
-            .parse()
-            .map_err(|_| Error::SubgraphNotFound(anyhow!("invalid deployment ID: {id}")))?;
-
-        // Get the deployment by ID, no subgraph
-        let deployment = network
-            .deployment_by_id(&id)
-            .ok_or_else(|| Error::SubgraphNotFound(anyhow!("deployment not found: {id}")))?;
-
-        Ok((vec![deployment], None))
-    } else {
-        Err(Error::SubgraphNotFound(anyhow!("missing identifier")))
+    match selector {
+        QuerySelector::Subgraph(subgraph_id) => {
+            // Get the subgraph by ID
+            let subgraph = network
+                .subgraph_by_id(subgraph_id)
+                .ok_or_else(|| Error::SubgraphNotFound(anyhow!("{subgraph_id}")))?;
+
+            // Get the subgraph's deployments (versions = deployments)
+            let versions = subgraph.deployments.clone();
+            Ok((versions, Some(subgraph)))
+        }
+        QuerySelector::Deployment(deployment_id) => {
+            // Get the deployment by ID, no subgraph
+            let deployment = network.deployment_by_id(deployment_id).ok_or_else(|| {
+                Error::SubgraphNotFound(anyhow!("deployment not found: {deployment_id}"))
+            })?;
+
+            Ok((vec![deployment], None))
+        }
     }
 }
 

graph-gateway/src/client_query/auth.rs

@@ -1,6 +1,7 @@
 use std::sync::Arc;
 
 use thegraph::subscriptions::auth::AuthTokenClaims;
+use thegraph::types::{DeploymentId, SubgraphId};
 
 use crate::subgraph_studio::APIKey;
 
@@ -20,6 +21,28 @@ pub enum AuthToken {
 }
 
 impl AuthToken {
+    /// Check if the given subgraph is authorized for this auth token.
+    pub fn is_subgraph_authorized(&self, subgraph: &SubgraphId) -> bool {
+        match self {
+            AuthToken::StudioApiKey(api_key) => studio::is_subgraph_authorized(api_key, subgraph),
+            AuthToken::SubscriptionsAuthToken(claims) => {
+                subscriptions::is_subgraph_authorized(claims, subgraph)
+            }
+        }
+    }
+
+    /// Check if the given deployment is authorized for this auth token.
+    pub fn is_deployment_authorized(&self, deployment: &DeploymentId) -> bool {
+        match self {
+            AuthToken::StudioApiKey(api_key) => {
+                studio::is_deployment_authorized(api_key, deployment)
+            }
+            AuthToken::SubscriptionsAuthToken(claims) => {
+                subscriptions::is_deployment_authorized(claims, deployment)
+            }
+        }
+    }
+
     /// Check if the given origin domain is authorized for this auth token.    
     pub fn is_domain_authorized(&self, domain: &str) -> bool {
         match self {

graph-gateway/src/client_query/auth/common.rs

@@ -17,6 +17,11 @@ pub fn are_deployments_authorized(
             .any(|deployment| authorized.contains(&deployment.id))
 }
 
+/// Check if the given deployment is authorized by the given authorized deployments.
+pub fn is_deployment_authorized(authorized: &[DeploymentId], deployment: &DeploymentId) -> bool {
+    authorized.is_empty() || authorized.contains(deployment)
+}
+
 /// Check if any of the given deployments are authorized by the given authorized subgraphs.
 ///
 /// If the authorized subgraphs set is empty, all deployments are considered authorized.
@@ -33,6 +38,11 @@ pub fn are_subgraphs_authorized(
         })
 }
 
+/// Check if the given subgraph is authorized by the given authorized subgraphs.
+pub fn is_subgraph_authorized(authorized: &[SubgraphId], subgraph: &SubgraphId) -> bool {
+    authorized.is_empty() || authorized.contains(subgraph)
+}
+
 /// Check if the query origin domain is authorized.
 ///
 /// If the authorized domain starts with a `*`, it is considered a wildcard

graph-gateway/src/client_query/auth/studio.rs

@@ -3,6 +3,7 @@ use std::sync::Arc;
 
 use anyhow::bail;
 use eventuals::{Eventual, Ptr};
+use thegraph::types::{DeploymentId, SubgraphId};
 
 use crate::subgraph_studio::{APIKey, QueryStatus};
 use crate::topology::Deployment;
@@ -87,6 +88,18 @@ pub fn parse_bearer_token(auth: &AuthContext, token: &str) -> anyhow::Result<Arc
         .ok_or_else(|| anyhow::anyhow!("API key not found"))
 }
 
+/// Check if the given deployment is authorized by the given API key.
+pub fn is_deployment_authorized(api_key: &Arc<APIKey>, deployment: &DeploymentId) -> bool {
+    let allowed_deployments = &api_key.deployments;
+    common::is_deployment_authorized(allowed_deployments, deployment)
+}
+
+/// Check if the given subgraph is authorized by the given API key.
+pub fn is_subgraph_authorized(api_key: &Arc<APIKey>, subgraph: &SubgraphId) -> bool {
+    let allowed_subgraphs = &api_key.subgraphs;
+    common::is_subgraph_authorized(allowed_subgraphs, subgraph)
+}
+
 /// Check if the given domain is authorized by the given API key.
 pub fn is_domain_authorized(api_key: &Arc<APIKey>, domain: &str) -> bool {
     let allowed_domains = &api_key

graph-gateway/src/client_query/auth/subscriptions.rs

@@ -5,6 +5,7 @@ use std::sync::{atomic, Arc};
 use alloy_primitives::Address;
 use eventuals::{Eventual, Ptr};
 use thegraph::subscriptions::auth::{parse_auth_token, verify_auth_token_claims, AuthTokenClaims};
+use thegraph::types::{DeploymentId, SubgraphId};
 use tokio::sync::RwLock;
 
 use crate::subscriptions::Subscription;
@@ -66,6 +67,18 @@ pub fn parse_bearer_token(_auth: &AuthContext, token: &str) -> anyhow::Result<Au
     Ok(claims)
 }
 
+/// Check if the given deployment is authorized by the given API key.
+pub fn is_deployment_authorized(auth_token: &AuthTokenClaims, deployment: &DeploymentId) -> bool {
+    let allowed_deployments = &auth_token.allowed_deployments;
+    common::is_deployment_authorized(allowed_deployments, deployment)
+}
+
+/// Check if the given subgraph is authorized by the given API key.
+pub fn is_subgraph_authorized(auth_token: &AuthTokenClaims, subgraph: &SubgraphId) -> bool {
+    let allowed_subgraphs = &auth_token.allowed_subgraphs;
+    common::is_subgraph_authorized(allowed_subgraphs, subgraph)
+}
+
 /// Check if the given domain is authorized by the auth token claims.
 pub fn is_domain_authorized(auth_token: &AuthTokenClaims, domain: &str) -> bool {
     // Get domain allowlist