refactor: add attestation middleware (#506)

Author: gusinacio

crates/service/src/error.rs

@@ -1,7 +1,6 @@
 // Copyright 2023-, Edge & Node, GraphOps, and Semiotic Labs.
 // SPDX-License-Identifier: Apache-2.0
 
-use alloy::primitives::Address;
 use anyhow::Error;
 use axum::{
     response::{IntoResponse, Response},
@@ -27,12 +26,6 @@ pub enum IndexerServiceError {
 
     #[error("Issues with provided receipt: {0}")]
     ReceiptError(#[from] tap_core::Error),
-    #[error("No attestation signer found for allocation `{0}`")]
-    NoSignerForAllocation(Address),
-    #[error("Error while processing the request: {0}")]
-    ProcessingError(SubgraphServiceError),
-    #[error("Failed to sign attestation")]
-    FailedToSignAttestation,
 
     #[error("There was an error while accessing escrow account: {0}")]
     EscrowAccount(#[from] EscrowAccountsError),
@@ -48,12 +41,10 @@ impl IntoResponse for IndexerServiceError {
         }
 
         let status = match self {
-            NoSignerForAllocation(_) | FailedToSignAttestation => StatusCode::INTERNAL_SERVER_ERROR,
-
-            ReceiptError(_) | EscrowAccount(_) | ProcessingError(_) => StatusCode::BAD_REQUEST,
+            ReceiptError(_) | EscrowAccount(_) => StatusCode::BAD_REQUEST,
             ReceiptNotFound => StatusCode::PAYMENT_REQUIRED,
             DeploymentIdNotFound => StatusCode::INTERNAL_SERVER_ERROR,
-            AxumError(_) => StatusCode::BAD_REQUEST,
+            AxumError(_) => StatusCode::INTERNAL_SERVER_ERROR,
             SerializationError(_) => StatusCode::BAD_REQUEST,
         };
         tracing::error!(%self, "An IndexerServiceError occoured.");

crates/service/src/middleware.rs

@@ -1,19 +1,23 @@
 // Copyright 2023-, Edge & Node, GraphOps, and Semiotic Labs.
 // SPDX-License-Identifier: Apache-2.0
 
+mod allocation;
+mod attestation;
+mod attestation_signer;
 pub mod auth;
-mod inject_allocation;
-mod inject_context;
-mod inject_deployment;
-mod inject_labels;
-mod inject_receipt;
-mod inject_sender;
+mod deployment;
+mod labels;
 mod prometheus_metrics;
+mod sender;
+mod tap_context;
+mod tap_receipt;
 
-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, SenderState};
+pub use allocation::{allocation_middleware, Allocation, AllocationState};
+pub use attestation::{attestation_middleware, AttestationInput};
+pub use attestation_signer::{signer_middleware, AttestationState};
+pub use deployment::deployment_middleware;
+pub use labels::labels_middleware;
 pub use prometheus_metrics::PrometheusMetricsMiddlewareLayer;
+pub use sender::{sender_middleware, SenderState};
+pub use tap_context::context_middleware;
+pub use tap_receipt::receipt_middleware;

crates/service/src/middleware/inject_allocation.rs renamed to crates/service/src/middleware/allocation.rs

@@ -58,7 +58,7 @@ pub async fn allocation_middleware(
 
 #[cfg(test)]
 mod tests {
-    use crate::middleware::inject_allocation::Allocation;
+    use crate::middleware::allocation::Allocation;
 
     use super::{allocation_middleware, AllocationState};
 

crates/service/src/middleware/attestation.rs

@@ -0,0 +1,208 @@
+// Copyright 2023-, Edge & Node, GraphOps, and Semiotic Labs.
+// SPDX-License-Identifier: Apache-2.0
+
+use std::string::FromUtf8Error;
+
+use axum::{
+    body::to_bytes,
+    extract::Request,
+    middleware::Next,
+    response::{IntoResponse, Response},
+};
+use reqwest::StatusCode;
+use serde::Serialize;
+use thegraph_core::Attestation;
+
+use indexer_attestation::AttestationSigner;
+
+#[derive(Clone)]
+pub enum AttestationInput {
+    Attestable { req: String },
+    NotAttestable,
+}
+
+#[derive(Debug, Serialize)]
+#[cfg_attr(test, derive(serde::Deserialize))]
+pub struct IndexerResponsePayload {
+    #[serde(rename = "graphQLResponse")]
+    graphql_response: String,
+    attestation: Option<Attestation>,
+}
+
+/// Check if the query is attestable and generates attestation
+///
+/// Executes query -> return subgraph response: (string, attestable (bool))
+/// if attestable && allocation id:
+///     - look for signer
+///     - create attestation
+///     - return response with attestation
+/// else:
+///     - return with no attestation
+///
+/// Requires AttestationSigner
+pub async fn attestation_middleware(
+    request: Request,
+    next: Next,
+) -> Result<Response, AttestationError> {
+    let signer = request
+        .extensions()
+        .get::<AttestationSigner>()
+        .cloned()
+        .ok_or(AttestationError::CouldNotFindSigner)?;
+
+    let (parts, graphql_response) = next.run(request).await.into_parts();
+    let attestation_response = parts.extensions.get::<AttestationInput>();
+    let bytes = to_bytes(graphql_response, usize::MAX).await?;
+    let res = String::from_utf8(bytes.into())?;
+
+    let attestation = match attestation_response {
+        Some(AttestationInput::Attestable { req }) => Some(signer.create_attestation(req, &res)),
+        _ => None,
+    };
+
+    let response = serde_json::to_string(&IndexerResponsePayload {
+        graphql_response: res,
+        attestation,
+    })?;
+
+    Ok(Response::new(response.into()))
+}
+
+#[derive(thiserror::Error, Debug)]
+pub enum AttestationError {
+    #[error("Could not find signer for allocation")]
+    CouldNotFindSigner,
+
+    #[error("There was an AxumError: {0}")]
+    AxumError(#[from] axum::Error),
+
+    #[error("There was an error converting the response to UTF-8 string: {0}")]
+    FromUtf8Error(#[from] FromUtf8Error),
+
+    #[error("there was an error while serializing the response: {0}")]
+    SerializationError(#[from] serde_json::Error),
+}
+
+impl IntoResponse for AttestationError {
+    fn into_response(self) -> Response {
+        match self {
+            AttestationError::CouldNotFindSigner
+            | AttestationError::AxumError(_)
+            | AttestationError::FromUtf8Error(_)
+            | AttestationError::SerializationError(_) => StatusCode::INTERNAL_SERVER_ERROR,
+        }
+        .into_response()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use alloy::primitives::Address;
+    use axum::{
+        body::{to_bytes, Body},
+        http::{Request, Response},
+        middleware::from_fn,
+        routing::get,
+        Router,
+    };
+    use indexer_allocation::Allocation;
+    use indexer_attestation::AttestationSigner;
+    use reqwest::StatusCode;
+    use test_assets::{INDEXER_ALLOCATIONS, INDEXER_MNEMONIC};
+    use tower::ServiceExt;
+
+    use crate::middleware::{
+        attestation::IndexerResponsePayload, attestation_middleware, AttestationInput,
+    };
+
+    const REQUEST: &str = "request";
+    const RESPONSE: &str = "response";
+
+    fn allocation_signer() -> (Allocation, AttestationSigner) {
+        let allocation = INDEXER_ALLOCATIONS
+            .values()
+            .collect::<Vec<_>>()
+            .pop()
+            .unwrap()
+            .clone();
+        let signer =
+            AttestationSigner::new(&INDEXER_MNEMONIC.to_string(), &allocation, 1, Address::ZERO)
+                .unwrap();
+        (allocation, signer)
+    }
+
+    async fn payload_from_response(res: Response<Body>) -> IndexerResponsePayload {
+        let bytes = to_bytes(res.into_body(), usize::MAX).await.unwrap();
+
+        serde_json::from_slice(&bytes).unwrap()
+    }
+
+    async fn send_request(app: Router, signer: Option<AttestationSigner>) -> Response<Body> {
+        let mut request = Request::builder().uri("/");
+
+        if let Some(signer) = signer {
+            request = request.extension(signer);
+        }
+
+        app.oneshot(request.body(Body::empty()).unwrap())
+            .await
+            .unwrap()
+    }
+
+    #[tokio::test]
+    async fn test_create_attestation() {
+        let (allocation, signer) = allocation_signer();
+        let middleware = from_fn(attestation_middleware);
+
+        let handle = move |_: Request<Body>| async move {
+            let mut res = Response::new(RESPONSE.to_string());
+            res.extensions_mut().insert(AttestationInput::Attestable {
+                req: REQUEST.to_string(),
+            });
+            res
+        };
+
+        let app = Router::new().route("/", get(handle)).layer(middleware);
+
+        // with signer
+        let res = send_request(app, Some(signer.clone())).await;
+        assert_eq!(res.status(), StatusCode::OK);
+
+        let response = payload_from_response(res).await;
+        assert_eq!(response.graphql_response, RESPONSE.to_string());
+
+        let attestation = response.attestation.unwrap();
+        assert!(signer
+            .verify(&attestation, REQUEST, RESPONSE, &allocation.id)
+            .is_ok());
+    }
+
+    #[tokio::test]
+    async fn test_non_assignable() {
+        let (_, signer) = allocation_signer();
+        let handle = move |_: Request<Body>| async move { Response::new(RESPONSE.to_string()) };
+
+        let middleware = from_fn(attestation_middleware);
+        let app = Router::new().route("/", get(handle)).layer(middleware);
+
+        let res = send_request(app, Some(signer.clone())).await;
+        assert_eq!(res.status(), StatusCode::OK);
+
+        let response = payload_from_response(res).await;
+        assert_eq!(response.graphql_response, RESPONSE.to_string());
+        assert!(response.attestation.is_none());
+    }
+
+    #[tokio::test]
+    async fn test_no_signer() {
+        let handle = move |_: Request<Body>| async move {
+            Response::new(RESPONSE.to_string());
+        };
+
+        let middleware = from_fn(attestation_middleware);
+        let app = Router::new().route("/", get(handle)).layer(middleware);
+
+        let res = send_request(app, None).await;
+        assert_eq!(res.status(), StatusCode::INTERNAL_SERVER_ERROR);
+    }
+}