|
| 1 | +// Copyright 2023-, Edge & Node, GraphOps, and Semiotic Labs. |
| 2 | +// SPDX-License-Identifier: Apache-2.0 |
| 3 | + |
| 4 | +use alloy::primitives::Address; |
| 5 | +use axum::{ |
| 6 | + extract::{Request, State}, |
| 7 | + middleware::Next, |
| 8 | + response::Response, |
| 9 | +}; |
| 10 | +use indexer_attestation::AttestationSigner; |
| 11 | +use std::collections::HashMap; |
| 12 | +use tokio::sync::watch; |
| 13 | + |
| 14 | +use super::Allocation; |
| 15 | + |
| 16 | +#[derive(Clone)] |
| 17 | +pub struct AttestationState { |
| 18 | + pub attestation_signers: watch::Receiver<HashMap<Address, AttestationSigner>>, |
| 19 | +} |
| 20 | + |
| 21 | +/// Injects the attestation signer to be used in the attestation |
| 22 | +/// |
| 23 | +/// Needs Allocation Extension |
| 24 | +pub async fn signer_middleware( |
| 25 | + State(state): State<AttestationState>, |
| 26 | + mut request: Request, |
| 27 | + next: Next, |
| 28 | +) -> Response { |
| 29 | + if let Some(Allocation(allocation_id)) = request.extensions().get::<Allocation>() { |
| 30 | + if let Some(signer) = state.attestation_signers.borrow().get(allocation_id) { |
| 31 | + request.extensions_mut().insert(signer.clone()); |
| 32 | + } |
| 33 | + } |
| 34 | + |
| 35 | + next.run(request).await |
| 36 | +} |
| 37 | + |
| 38 | +#[cfg(test)] |
| 39 | +mod tests { |
| 40 | + use crate::middleware::{allocation::Allocation, signer_middleware, AttestationState}; |
| 41 | + |
| 42 | + use axum::{body::Body, http::Request, middleware::from_fn_with_state, routing::get, Router}; |
| 43 | + use indexer_attestation::AttestationSigner; |
| 44 | + use indexer_monitor::attestation_signers; |
| 45 | + use reqwest::StatusCode; |
| 46 | + use test_assets::{DISPUTE_MANAGER_ADDRESS, INDEXER_ALLOCATIONS, INDEXER_MNEMONIC}; |
| 47 | + use tokio::sync::{mpsc::channel, watch}; |
| 48 | + use tower::Service; |
| 49 | + |
| 50 | + #[tokio::test] |
| 51 | + async fn test_attestation_signer_middleware() { |
| 52 | + let allocations = (*INDEXER_ALLOCATIONS).clone(); |
| 53 | + |
| 54 | + let allocation = **allocations.keys().collect::<Vec<_>>().first().unwrap(); |
| 55 | + |
| 56 | + let (_, allocations_rx) = watch::channel(allocations); |
| 57 | + let (_, dispute_manager_rx) = watch::channel(*DISPUTE_MANAGER_ADDRESS); |
| 58 | + let attestation_signers = attestation_signers( |
| 59 | + allocations_rx, |
| 60 | + INDEXER_MNEMONIC.clone(), |
| 61 | + 1, |
| 62 | + dispute_manager_rx, |
| 63 | + ); |
| 64 | + |
| 65 | + let expected_signer = attestation_signers |
| 66 | + .borrow() |
| 67 | + .get(&allocation) |
| 68 | + .unwrap() |
| 69 | + .clone(); |
| 70 | + |
| 71 | + let state = AttestationState { |
| 72 | + attestation_signers, |
| 73 | + }; |
| 74 | + |
| 75 | + let middleware = from_fn_with_state(state, signer_middleware); |
| 76 | + |
| 77 | + let (tx, mut rx) = channel(1); |
| 78 | + |
| 79 | + let handle = move |request: Request<Body>| async move { |
| 80 | + tx.send(request).await.unwrap(); |
| 81 | + Body::empty() |
| 82 | + }; |
| 83 | + |
| 84 | + let mut app = Router::new().route("/", get(handle)).layer(middleware); |
| 85 | + |
| 86 | + // with allocation |
| 87 | + let res = app |
| 88 | + .call( |
| 89 | + Request::builder() |
| 90 | + .uri("/") |
| 91 | + .extension(Allocation(allocation)) |
| 92 | + .body(Body::empty()) |
| 93 | + .unwrap(), |
| 94 | + ) |
| 95 | + .await |
| 96 | + .unwrap(); |
| 97 | + assert_eq!(res.status(), StatusCode::OK); |
| 98 | + |
| 99 | + let req = rx.recv().await.unwrap(); |
| 100 | + let signer = req.extensions().get::<AttestationSigner>().unwrap(); |
| 101 | + assert_eq!(*signer, expected_signer); |
| 102 | + |
| 103 | + // without allocation |
| 104 | + let res = app |
| 105 | + .call(Request::builder().uri("/").body(Body::empty()).unwrap()) |
| 106 | + .await |
| 107 | + .unwrap(); |
| 108 | + assert_eq!(res.status(), StatusCode::OK); |
| 109 | + |
| 110 | + let req = rx.recv().await.unwrap(); |
| 111 | + assert!(req.extensions().get::<AttestationSigner>().is_none()); |
| 112 | + } |
| 113 | +} |
0 commit comments