refactor: use attestation middleware

Signed-off-by: Gustavo Inacio <[email protected]>

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,8 @@ 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,8 +43,6 @@ impl IntoResponse for IndexerServiceError {
         }
 
         let status = match self {
-            NoSignerForAllocation(_) | FailedToSignAttestation => StatusCode::INTERNAL_SERVER_ERROR,
-
             ReceiptError(_) | EscrowAccount(_) | ProcessingError(_) => StatusCode::BAD_REQUEST,
             ReceiptNotFound => StatusCode::PAYMENT_REQUIRED,
             DeploymentIdNotFound => StatusCode::INTERNAL_SERVER_ERROR,

crates/service/src/middleware.rs

@@ -18,4 +18,6 @@ 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 inject_attestation_signer::{signer_middleware, AttestationState};
+pub use attestation::{attestation_middleware, AttestationInput};
 pub use prometheus_metrics::PrometheusMetricsMiddlewareLayer;

crates/service/src/routes/request_handler.rs

@@ -3,51 +3,28 @@
 
 use std::sync::Arc;
 
-use crate::{error::IndexerServiceError, middleware::Allocation};
+use crate::error::IndexerServiceError;
 use axum::{
     extract::{Path, State},
     response::IntoResponse,
-    Extension,
 };
-use reqwest::StatusCode;
 use thegraph_core::DeploymentId;
 use tracing::trace;
 
-use crate::service::{AttestationOutput, IndexerServiceResponse, IndexerServiceState};
+use crate::service::IndexerServiceState;
 
 pub async fn request_handler(
     Path(manifest_id): Path<DeploymentId>,
-    Extension(Allocation(allocation_id)): Extension<Allocation>,
     State(state): State<Arc<IndexerServiceState>>,
     req: String,
 ) -> Result<impl IntoResponse, IndexerServiceError> {
     trace!("Handling request for deployment `{manifest_id}`");
 
-    // Check if we have an attestation signer for the allocation the receipt was created for
-    let signer = state
-        .attestation_signers
-        .borrow()
-        .get(&allocation_id)
-        .cloned()
-        .ok_or_else(|| (IndexerServiceError::NoSignerForAllocation(allocation_id)))?;
-
     let response = state
         .service_impl
-        .process_request(manifest_id, &req)
+        .process_request(manifest_id, req)
         .await
         .map_err(IndexerServiceError::ProcessingError)?;
 
-    let res = response
-        .as_str()
-        .map_err(|_| IndexerServiceError::FailedToSignAttestation)?;
-
-    let attestation = AttestationOutput::Attestation(
-        response
-            .is_attestable()
-            .then(|| signer.create_attestation(&req, res)),
-    );
-
-    let response = response.finalize(attestation);
-
-    Ok((StatusCode::OK, response))
+    Ok(response)
 }

crates/service/src/service.rs

@@ -9,13 +9,12 @@ use anyhow::anyhow;
 use async_graphql::{EmptySubscription, Schema};
 use async_graphql_axum::GraphQL;
 use axum::{
+    http::{HeaderValue, Response},
     routing::{post, post_service},
-    Json, Router,
+    Router,
 };
 use indexer_config::{Config, DipsConfig};
-use reqwest::Url;
-use serde::{de::DeserializeOwned, Serialize};
-use serde_json::{json, Value};
+use reqwest::{header::CONTENT_TYPE, Url};
 use sqlx::PgPool;
 use thegraph_core::attestation::eip712_domain;
 use thegraph_core::DeploymentId;
@@ -26,6 +25,7 @@ use crate::{
         self,
         dips::{AgreementStore, InMemoryAgreementStore},
     },
+    middleware::AttestationInput,
     routes::dips::Price,
     service::indexer_service::{IndexerServiceOptions, IndexerServiceRelease},
 };
@@ -38,42 +38,6 @@ mod tap_receipt_header;
 pub use indexer_service::{AttestationOutput, IndexerServiceResponse, IndexerServiceState};
 pub use tap_receipt_header::TapReceipt;
 
-#[derive(Debug)]
-pub struct SubgraphServiceResponse {
-    inner: String,
-    attestable: bool,
-}
-
-impl SubgraphServiceResponse {
-    pub fn new(inner: String, attestable: bool) -> Self {
-        Self { inner, attestable }
-    }
-}
-
-impl IndexerServiceResponse for SubgraphServiceResponse {
-    type Data = Json<Value>;
-    type Error = SubgraphServiceError; // not used
-
-    fn is_attestable(&self) -> bool {
-        self.attestable
-    }
-
-    fn as_str(&self) -> Result<&str, Self::Error> {
-        Ok(self.inner.as_str())
-    }
-
-    fn finalize(self, attestation: AttestationOutput) -> Self::Data {
-        let (attestation_key, attestation_value) = match attestation {
-            AttestationOutput::Attestation(attestation) => ("attestation", json!(attestation)),
-            AttestationOutput::Attestable => ("attestable", json!(self.is_attestable())),
-        };
-        Json(json!({
-            "graphQLResponse": self.inner,
-            attestation_key: attestation_value,
-        }))
-    }
-}
-
 #[derive(Clone)]
 pub struct SubgraphServiceState {
     pub config: &'static Config,
@@ -95,11 +59,11 @@ impl SubgraphService {
 }
 
 impl SubgraphService {
-    pub async fn process_request<Request: DeserializeOwned + Send + std::fmt::Debug + Serialize>(
+    pub async fn process_request(
         &self,
         deployment: DeploymentId,
-        request: &Request,
-    ) -> Result<SubgraphServiceResponse, SubgraphServiceError> {
+        req: String,
+    ) -> Result<Response<String>, SubgraphServiceError> {
         let deployment_url = self
             .state
             .graph_node_query_base_url
@@ -110,7 +74,8 @@ impl SubgraphService {
             .state
             .graph_node_client
             .post(deployment_url)
-            .json(request)
+            .body(req.clone())
+            .header(CONTENT_TYPE, HeaderValue::from_static("application/json"))
             .send()
             .await
             .map_err(SubgraphServiceError::QueryForwardingError)?;
@@ -126,8 +91,16 @@ impl SubgraphService {
             .text()
             .await
             .map_err(SubgraphServiceError::QueryForwardingError)?;
+        let attestation_input = if attestable {
+            AttestationInput::Attestable { req }
+        } else {
+            AttestationInput::NotAttestable
+        };
+
+        let mut response = Response::new(body);
+        response.extensions_mut().insert(attestation_input);
 
-        Ok(SubgraphServiceResponse::new(body, attestable))
+        Ok(response)
     }
 }
 

crates/service/src/service/indexer_service.rs

@@ -11,7 +11,6 @@ use axum::{
     serve, Json, Router, ServiceExt,
 };
 use build_info::BuildInfo;
-use indexer_attestation::AttestationSigner;
 use indexer_monitor::{
     attestation_signers, deployment_to_allocation, dispute_manager, escrow_accounts,
     indexer_allocations, DeploymentDetails, SubgraphClient,
@@ -24,8 +23,8 @@ use std::{
     collections::HashMap, error::Error, net::SocketAddr, path::PathBuf, sync::Arc, time::Duration,
 };
 use tap_core::{manager::Manager, receipt::checks::CheckList, tap_eip712_domain};
-use thegraph_core::{Address, Attestation};
-use tokio::{net::TcpListener, signal, sync::watch::Receiver};
+use thegraph_core::Attestation;
+use tokio::{net::TcpListener, signal};
 use tower::ServiceBuilder;
 use tower_governor::{governor::GovernorConfigBuilder, GovernorLayer};
 use tower_http::{
@@ -40,10 +39,11 @@ use tracing::{error, info, info_span, warn};
 use crate::{
     metrics::{FAILED_RECEIPT, HANDLER_FAILURE, HANDLER_HISTOGRAM},
     middleware::{
-        allocation_middleware,
+        allocation_middleware, attestation_middleware,
         auth::{self, Bearer, OrExt},
         context_middleware, deployment_middleware, labels_middleware, receipt_middleware,
-        sender_middleware, AllocationState, PrometheusMetricsMiddlewareLayer, SenderState,
+        sender_middleware, signer_middleware, AllocationState, AttestationState,
+        PrometheusMetricsMiddlewareLayer, SenderState,
     },
     routes::{health, request_handler, static_subgraph_request_handler},
     tap::IndexerTapContext,
@@ -98,7 +98,6 @@ pub struct IndexerServiceOptions {
 
 pub struct IndexerServiceState {
     pub config: Config,
-    pub attestation_signers: Receiver<HashMap<Address, AttestationSigner>>,
     pub service_impl: SubgraphService,
 }
 
@@ -276,7 +275,6 @@ pub async fn run(options: IndexerServiceOptions) -> Result<(), anyhow::Error> {
 
     let state = Arc::new(IndexerServiceState {
         config: options.config.clone(),
-        attestation_signers,
         service_impl: options.service_impl,
     });
 
@@ -386,6 +384,9 @@ pub async fn run(options: IndexerServiceOptions) -> Result<(), anyhow::Error> {
         escrow_accounts,
         domain_separator,
     };
+    let attestation_state = AttestationState {
+        attestation_signers,
+    };
 
     let service_builder = ServiceBuilder::new()
         // inject deployment id
@@ -404,7 +405,11 @@ pub async fn run(options: IndexerServiceOptions) -> Result<(), anyhow::Error> {
             HANDLER_FAILURE.clone(),
         ))
         // tap context
-        .layer(from_fn(context_middleware));
+        .layer(from_fn(context_middleware))
+        // inject signer
+        .layer(from_fn_with_state(attestation_state, signer_middleware))
+        // create attestation
+        .layer(from_fn(attestation_middleware));
 
     request_handler_route = request_handler_route.layer(service_builder);