graphprotocol · gusinacio · Nov 22, 2024 · Nov 21, 2024 · Nov 21, 2024 · Nov 21, 2024
diff --git a/.sqlx/query-6c05fc541bf0bb2af20fbe62747456055d5ebda5cb136d9d015f101ebbfe495f.json b/.sqlx/query-6c05fc541bf0bb2af20fbe62747456055d5ebda5cb136d9d015f101ebbfe495f.json
diff --git a/crates/service/src/error.rs b/crates/service/src/error.rs
@@ -15,18 +15,22 @@ use thiserror::Error;
 
 #[derive(Debug, Error)]
 pub enum IndexerServiceError {
+    #[error("No Tap receipt was found in the request")]
+    ReceiptNotFound,
+    #[error("Could not find deployment id")]
+    DeploymentIdNotFound,
+    #[error(transparent)]
+    AxumError(#[from] axum::Error),
+
+    #[error(transparent)]
+    SerializationError(#[from] serde_json::Error),
+
     #[error("Issues with provided receipt: {0}")]
     ReceiptError(#[from] tap_core::Error),
     #[error("No attestation signer found for allocation `{0}`")]
     NoSignerForAllocation(Address),
-    #[error("Invalid request body: {0}")]
-    InvalidRequest(anyhow::Error),
     #[error("Error while processing the request: {0}")]
     ProcessingError(SubgraphServiceError),
-    #[error("No valid receipt or free query auth token provided")]
-    Unauthorized,
-    #[error("Invalid free query auth token")]
-    InvalidFreeQueryAuthToken,
     #[error("Failed to sign attestation")]
     FailedToSignAttestation,
 
@@ -44,15 +48,13 @@ impl IntoResponse for IndexerServiceError {
         }
 
         let status = match self {
-            Unauthorized => StatusCode::UNAUTHORIZED,
-
             NoSignerForAllocation(_) | FailedToSignAttestation => StatusCode::INTERNAL_SERVER_ERROR,
 
-            ReceiptError(_)
-            | InvalidRequest(_)
-            | InvalidFreeQueryAuthToken
-            | EscrowAccount(_)
-            | ProcessingError(_) => StatusCode::BAD_REQUEST,
+            ReceiptError(_) | EscrowAccount(_) | ProcessingError(_) => StatusCode::BAD_REQUEST,
+            ReceiptNotFound => StatusCode::PAYMENT_REQUIRED,
+            DeploymentIdNotFound => StatusCode::INTERNAL_SERVER_ERROR,
+            AxumError(_) => StatusCode::BAD_REQUEST,
+            SerializationError(_) => StatusCode::BAD_REQUEST,
         };
         tracing::error!(%self, "An IndexerServiceError occoured.");
         (

diff --git a/crates/service/src/middleware.rs b/crates/service/src/middleware.rs
@@ -1,16 +1,19 @@
 // Copyright 2023-, Edge & Node, GraphOps, and Semiotic Labs.
 // SPDX-License-Identifier: Apache-2.0
 
+pub mod auth;
 mod inject_allocation;
+mod inject_context;
 mod inject_deployment;
 mod inject_labels;
 mod inject_receipt;
 mod inject_sender;
 mod prometheus_metrics;
 
 pub use inject_allocation::{allocation_middleware, Allocation, AllocationState};
+pub use inject_context::context_middleware;
 pub use inject_deployment::deployment_middleware;
 pub use inject_labels::labels_middleware;
 pub use inject_receipt::receipt_middleware;
-pub use inject_sender::{sender_middleware, Sender, SenderState};
+pub use inject_sender::{sender_middleware, SenderState};
 pub use prometheus_metrics::PrometheusMetricsMiddlewareLayer;
diff --git a/crates/service/src/middleware/auth.rs b/crates/service/src/middleware/auth.rs
@@ -0,0 +1,106 @@
+// Copyright 2023-, Edge & Node, GraphOps, and Semiotic Labs.
+// SPDX-License-Identifier: Apache-2.0
+
+mod bearer;
+mod or;
+mod tap;
+
+pub use bearer::Bearer;
+pub use or::OrExt;
+pub use tap::tap_receipt_authorize;
+
+#[cfg(test)]
+mod tests {
+    use std::time::Duration;
+
+    use alloy::primitives::{address, Address};
+    use axum::body::Body;
+    use axum::http::{Request, Response};
+    use reqwest::{header, StatusCode};
+    use sqlx::PgPool;
+    use tap_core::{manager::Manager, receipt::checks::CheckList};
+    use tower::{Service, ServiceBuilder, ServiceExt};
+    use tower_http::auth::AsyncRequireAuthorizationLayer;
+
+    use crate::middleware::auth::{self, Bearer, OrExt};
+    use crate::tap::IndexerTapContext;
+    use test_assets::{create_signed_receipt, TAP_EIP712_DOMAIN};
+
+    const ALLOCATION_ID: Address = address!("deadbeefcafebabedeadbeefcafebabedeadbeef");
+
+    async fn handle(_: Request<Body>) -> anyhow::Result<Response<Body>> {
+        Ok(Response::new(Body::default()))
+    }
+
+    #[sqlx::test(migrations = "../../migrations")]
+    async fn test_middleware_composition(pgpool: PgPool) {
+        let token = "test".to_string();
+        let context = IndexerTapContext::new(pgpool.clone(), TAP_EIP712_DOMAIN.clone()).await;
+        let tap_manager = Box::leak(Box::new(Manager::new(
+            TAP_EIP712_DOMAIN.clone(),
+            context,
+            CheckList::empty(),
+        )));
+        let metric = Box::leak(Box::new(
+            prometheus::register_counter_vec!(
+                "merge_checks_test",
+                "Failed queries to handler",
+                &["deployment"]
+            )
+            .unwrap(),
+        ));
+        let free_query = Bearer::new(&token);
+        let tap_auth = auth::tap_receipt_authorize(tap_manager, metric);
+        let authorize_requests = free_query.or(tap_auth);
+
+        let authorization_middleware = AsyncRequireAuthorizationLayer::new(authorize_requests);
+
+        let mut service = ServiceBuilder::new()
+            .layer(authorization_middleware)
+            .service_fn(handle);
+
+        let handle = service.ready().await.unwrap();
+
+        // should allow queries that contains the free token
+        // if the token does not match, return payment required
+        let mut req = Request::new(Default::default());
+        req.headers_mut().insert(
+            header::AUTHORIZATION,
+            format!("Bearer {token}").parse().unwrap(),
+        );
+        let res = handle.call(req).await.unwrap();
+        assert_eq!(res.status(), StatusCode::OK);
+
+        // if the token exists but is wrong, try the receipt
+        let mut req = Request::new(Default::default());
+        req.headers_mut()
+            .insert(header::AUTHORIZATION, "Bearer wrongtoken".parse().unwrap());
+        let res = handle.call(req).await.unwrap();
+        // we return the error from tap
+        assert_eq!(res.status(), StatusCode::PAYMENT_REQUIRED);
+
+        let receipt = create_signed_receipt(ALLOCATION_ID, 1, 1, 1).await;
+
+        // check with receipt
+        let mut req = Request::new(Default::default());
+        req.extensions_mut().insert(receipt);
+        let res = handle.call(req).await.unwrap();
+        assert_eq!(res.status(), StatusCode::OK);
+
+        // todo make this sleep better
+        tokio::time::sleep(Duration::from_millis(100)).await;
+
+        // verify receipts
+        let result = sqlx::query!("SELECT * FROM scalar_tap_receipts")
+            .fetch_all(&pgpool)
+            .await
+            .unwrap();
+        assert_eq!(result.len(), 1);
+
+        // if it has neither, should return unauthorized
+        // check no headers
+        let req = Request::new(Default::default());
+        let res = handle.call(req).await.unwrap();
+        assert_eq!(res.status(), StatusCode::PAYMENT_REQUIRED);
+    }
+}
diff --git a/crates/service/src/middleware/auth/bearer.rs b/crates/service/src/middleware/auth/bearer.rs
@@ -0,0 +1,67 @@
+// Copyright 2023-, Edge & Node, GraphOps, and Semiotic Labs.
+// SPDX-License-Identifier: Apache-2.0
+
+//! Bearer struct from tower-http but exposing the `new()` function
+//! to allow creation
+//!
+//! This code is from *tower-http*
+
+use std::{fmt, marker::PhantomData};
+
+use axum::http::{HeaderValue, Request, Response};
+use reqwest::{header, StatusCode};
+use tower_http::validate_request::ValidateRequest;
+
+pub struct Bearer<ResBody> {
+    header_value: HeaderValue,
+    _ty: PhantomData<fn() -> ResBody>,
+}
+
+impl<ResBody> Bearer<ResBody> {
+    pub fn new(token: &str) -> Self
+    where
+        ResBody: Default,
+    {
+        Self {
+            header_value: format!("Bearer {}", token)
+                .parse()
+                .expect("token is not a valid header value"),
+            _ty: PhantomData,
+        }
+    }
+}
+
+impl<ResBody> Clone for Bearer<ResBody> {
+    fn clone(&self) -> Self {
+        Self {
+            header_value: self.header_value.clone(),
+            _ty: PhantomData,
+        }
+    }
+}
+
+impl<ResBody> fmt::Debug for Bearer<ResBody> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_struct("Bearer")
+            .field("header_value", &self.header_value)
+            .finish()
+    }
+}
+
+impl<B, ResBody> ValidateRequest<B> for Bearer<ResBody>
+where
+    ResBody: Default,
+{
+    type ResponseBody = ResBody;
+
+    fn validate(&mut self, request: &mut Request<B>) -> Result<(), Response<Self::ResponseBody>> {
+        match request.headers().get(header::AUTHORIZATION) {
+            Some(actual) if actual == self.header_value => Ok(()),
+            _ => {
+                let mut res = Response::new(ResBody::default());
+                *res.status_mut() = StatusCode::UNAUTHORIZED;
+                Err(res)
+            }
+        }
+    }
+}