From 6d5f00152127c2004f4101bfe3e0258ecab0f0f1 Mon Sep 17 00:00:00 2001
From: Lorenzo Delgado <lorenzo@edgeandnode.com>
Date: Tue, 23 Jan 2024 15:46:21 +0100
Subject: [PATCH] refactor(graph-gateway): use require_auth client query
 middleware (#550)

---
 graph-gateway/src/client_query.rs | 24 +++++++-----------------
 graph-gateway/src/main.rs         | 15 +++++++++------
 2 files changed, 16 insertions(+), 23 deletions(-)

diff --git a/graph-gateway/src/client_query.rs b/graph-gateway/src/client_query.rs
index da40a67c..03e67769 100644
--- a/graph-gateway/src/client_query.rs
+++ b/graph-gateway/src/client_query.rs
@@ -4,13 +4,14 @@ use std::time::{Duration, Instant};
 
 use alloy_primitives::{Address, BlockNumber, U256};
 use alloy_sol_types::Eip712Domain;
-use anyhow::{anyhow, Context as _};
+use anyhow::anyhow;
 use axum::extract::OriginalUri;
 use axum::http::Uri;
 use axum::{
     body::Bytes,
     extract::{Path, State},
     http::{header, HeaderMap, Response, StatusCode},
+    Extension,
 };
 use cost_model::{Context as AgoraContext, CostModel};
 use eventuals::Ptr;
@@ -71,6 +72,7 @@ pub struct QueryBody {
 
 pub async fn handle_query(
     State(ctx): State<Context>,
+    Extension(auth): Extension<AuthToken>,
     OriginalUri(original_uri): OriginalUri,
     Path(params): Path<BTreeMap<String, String>>,
     headers: HeaderMap,
@@ -81,16 +83,6 @@ pub async fn handle_query(
     let start_time = Instant::now();
     let timestamp = unix_timestamp();
 
-    let bearer_token = headers
-        .get(header::AUTHORIZATION)
-        .and_then(|h| h.to_str().ok())
-        .map(|header| header.trim_start_matches("Bearer").trim())
-        .unwrap_or("");
-    let auth = ctx
-        .auth_handler
-        .parse_bearer_token(bearer_token)
-        .context("Invalid auth");
-
     let resolved_deployments = resolve_subgraph_deployments(&ctx.network, &params).await;
 
     // This is very useful for investigating gateway logs in production.
@@ -131,22 +123,20 @@ pub async fn handle_query(
         }
     }
 
-    tracing::debug!(%bearer_token);
-
+    // TODO(LNSD): Move origin header parsing to an extractor.
     let domain = headers
         .get(header::ORIGIN)
         .and_then(|v| v.to_str().ok())
         .and_then(|v| Some(v.parse::<Url>().ok()?.host_str()?.to_string()))
         .unwrap_or("".to_string());
 
-    let result = match (auth, resolved_deployments) {
-        (Ok(auth), Ok((deployments, _))) => {
+    let result = match resolved_deployments {
+        Ok((deployments, _)) => {
             handle_client_query_inner(&ctx, deployments, payload, auth, domain)
                 .instrument(span.clone())
                 .await
         }
-        (Err(auth_err), _) => Err(Error::Auth(auth_err)),
-        (_, Err(subgraph_resolution_err)) => Err(subgraph_resolution_err),
+        Err(subgraph_resolution_err) => Err(subgraph_resolution_err),
     };
 
     let deployment: Option<String> = result
diff --git a/graph-gateway/src/main.rs b/graph-gateway/src/main.rs
index b49234c8..a91c1afb 100644
--- a/graph-gateway/src/main.rs
+++ b/graph-gateway/src/main.rs
@@ -49,6 +49,7 @@ use graph_gateway::client_query::context::Context;
 use graph_gateway::client_query::legacy_auth_adapter::legacy_auth_adapter;
 use graph_gateway::client_query::query_id::SetQueryIdLayer;
 use graph_gateway::client_query::query_tracing::QueryTracingLayer;
+use graph_gateway::client_query::require_auth::RequireAuthorizationLayer;
 use graph_gateway::config::{Config, ExchangeRateProvider};
 use graph_gateway::indexer_client::IndexerClient;
 use graph_gateway::indexers::indexing;
@@ -357,7 +358,7 @@ async fn main() {
         .with_state(client_query_ctx)
         .layer(
             // ServiceBuilder works by composing all layers into one such that they run top to
-            // bottom, and then the response would bubble back up through the layers in reverse.
+            // bottom, and then the response would bubble back up through the layers in reverse
             tower::ServiceBuilder::new()
                 .layer(
                     CorsLayer::new()
@@ -365,17 +366,19 @@ async fn main() {
                         .allow_headers(cors::Any)
                         .allow_methods([http::Method::OPTIONS, http::Method::POST]),
                 )
-                // Set up the query tracing span.
+                // Set up the query tracing span
                 .layer(QueryTracingLayer::new(config.graph_env_id.clone()))
-                // Set the query ID on the request.
+                // Set the query ID on the request
                 .layer(SetQueryIdLayer::new(gateway_id))
-                // Handle legacy in-path auth, and convert it into a header.
-                .layer(middleware::from_fn(legacy_auth_adapter)),
+                // Handle legacy in-path auth, and convert it into a header
+                .layer(middleware::from_fn(legacy_auth_adapter))
+                // Require the query to be authorized
+                .layer(RequireAuthorizationLayer::new(auth_handler.clone())),
         );
 
     let router = Router::new()
         .route("/", routing::get(|| async { "Ready to roll!" }))
-        // This path is required by NGINX ingress controller.
+        // This path is required by NGINX ingress controller
         .route("/ready", routing::get(|| async { "Ready" }))
         .route(
             "/collect-receipts",