Merge pull request #3 from sigp/gzip-and-ssz-relay

realbigsean · web-flow · commit 76c5c47f8ae7 · 2024-10-04T18:31:57.000-07:00
Gzip and ssz relay
diff --git a/Cargo.toml b/Cargo.toml
@@ -17,6 +17,7 @@ eth2 = { git = "https://github.com/realbigsean/lighthouse.git", rev = "10ce60633
 ethereum_serde_utils = "0.5.2"
 ethereum_ssz = "0.5.4"
 ethereum_ssz_derive = "0.5.4"
+flate2 = "1.0.34"
 futures = "0.3.30"
 http = "1"
 reqwest = { version = "0.12.5", features = ["json"] }
diff --git a/relay-api-types/src/lib.rs b/relay-api-types/src/lib.rs
@@ -264,3 +264,56 @@ pub type GetValidatorsResponse = Vec<ValidatorsResponse>;
 pub type GetDeliveredPayloadsResponse = Vec<BidTraceV2>;
 pub type GetReceivedBidsResponse = Vec<BidTraceV2WithTimestamp>;
 pub type GetValidatorRegistrationResponse = SignedValidatorRegistrationData;
+
+// Headers
+#[derive(Default)]
+pub enum ContentType {
+    #[default]
+    Json,
+    Ssz,
+}
+
+impl std::fmt::Display for ContentType {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            ContentType::Json => write!(f, "application/json"),
+            ContentType::Ssz => write!(f, "application/octet-stream"),
+        }
+    }
+}
+
+impl From<String> for ContentType {
+    fn from(value: String) -> Self {
+        match value.as_str() {
+            "application/json" => ContentType::Json,
+            "application/octet-stream" => ContentType::Ssz,
+            _ => panic!("unknown content type: {}", value),
+        }
+    }
+}
+
+#[derive(Default)]
+pub enum ContentEncoding {
+    Gzip,
+    #[default]
+    None,
+}
+
+impl std::fmt::Display for ContentEncoding {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            ContentEncoding::Gzip => write!(f, "gzip"),
+            ContentEncoding::None => write!(f, ""),
+        }
+    }
+}
+
+impl From<String> for ContentEncoding {
+    fn from(value: String) -> Self {
+        match value.as_ref() {
+            "gzip" => ContentEncoding::Gzip,
+            "" => ContentEncoding::None,
+            _ => panic!("unknown content encoding: {}", value),
+        }
+    }
+}
diff --git a/relay-client/src/lib.rs b/relay-client/src/lib.rs
@@ -1,4 +1,8 @@
 use futures::{Stream, StreamExt};
+use http::header::InvalidHeaderValue;
+use http::header::CONTENT_ENCODING;
+use http::header::CONTENT_TYPE;
+use http::HeaderValue;
 pub use relay_api_types::*;
 use reqwest::Client;
 use reqwest::Url;
@@ -13,6 +17,7 @@ pub enum Error {
     StatusCode(http::StatusCode),
     InvalidUrl(Url),
     WebSocket(tokio_tungstenite::tungstenite::Error),
+    InvalidHeader(InvalidHeaderValue),
 }
 
 impl From<reqwest::Error> for Error {
@@ -27,6 +32,12 @@ impl From<tokio_tungstenite::tungstenite::Error> for Error {
     }
 }
 
+impl From<InvalidHeaderValue> for Error {
+    fn from(value: InvalidHeaderValue) -> Self {
+        Error::InvalidHeader(value)
+    }
+}
+
 #[derive(Clone)]
 pub struct RelayClient {
     client: Client,
@@ -45,6 +56,34 @@ impl RelayClient {
     where
         T: for<'de> Deserialize<'de>,
     {
+        self.build_response_with_headers(response, <_>::default(), <_>::default())
+            .await
+    }
+
+    async fn build_response_with_headers<T>(
+        &self,
+        mut response: reqwest::Response,
+        content_type: ContentType,
+        content_encoding: ContentEncoding,
+    ) -> Result<T, Error>
+    where
+        T: for<'de> Deserialize<'de>,
+    {
+        response.headers_mut().insert(
+            CONTENT_TYPE,
+            HeaderValue::from_str(content_type.to_string().as_str())?,
+        );
+
+        match content_encoding {
+            ContentEncoding::Gzip => {
+                response.headers_mut().insert(
+                    CONTENT_ENCODING,
+                    HeaderValue::from_str(content_encoding.to_string().as_str())?,
+                );
+            }
+            ContentEncoding::None => {}
+        }
+
         let status = response.status();
         let text = response.text().await;
 
@@ -62,6 +101,8 @@ impl RelayClient {
         &self,
         query_params: SubmitBlockQueryParams,
         body: SubmitBlockRequest<E>,
+        content_type: ContentType,
+        content_encoding: ContentEncoding,
     ) -> Result<(), Error>
     where
         E: EthSpec,
@@ -78,7 +119,8 @@ impl RelayClient {
             .send()
             .await?;
 
-        self.build_response(response).await
+        self.build_response_with_headers(response, content_type, content_encoding)
+            .await
     }
 
     pub async fn get_validators<E>(&self) -> Result<GetValidatorsResponse, Error>
@@ -149,6 +191,8 @@ impl RelayClient {
         &self,
         query_params: SubmitBlockQueryParams,
         body: SignedHeaderSubmission<E>,
+        content_type: ContentType,
+        content_encoding: ContentEncoding,
     ) -> Result<(), Error>
     where
         E: EthSpec,
@@ -165,13 +209,16 @@ impl RelayClient {
             .send()
             .await?;
 
-        self.build_response(response).await
+        self.build_response_with_headers(response, content_type, content_encoding)
+            .await
     }
 
     pub async fn submit_block_optimistic_v2<E>(
         &self,
         query_params: SubmitBlockQueryParams,
         body: SubmitBlockRequest<E>,
+        content_type: ContentType,
+        content_encoding: ContentEncoding,
     ) -> Result<(), Error>
     where
         E: EthSpec,
@@ -188,17 +235,24 @@ impl RelayClient {
             .send()
             .await?;
 
-        self.build_response(response).await
+        self.build_response_with_headers(response, content_type, content_encoding)
+            .await
     }
 
-    pub async fn submit_cancellation(&self, body: SignedCancellation) -> Result<(), Error> {
+    pub async fn submit_cancellation(
+        &self,
+        body: SignedCancellation,
+        content_type: ContentType,
+        content_encoding: ContentEncoding,
+    ) -> Result<(), Error> {
         let mut url = self.base_url.clone();
         url.path_segments_mut()
             .map_err(|_| Error::InvalidUrl(self.base_url.clone()))?
             .extend(&["relay", "v1", "builder", "cancel_bid"]);
         let response = self.client.post(url).json(&body).send().await?;
 
-        self.build_response(response).await
+        self.build_response_with_headers(response, content_type, content_encoding)
+            .await
     }
 
     pub async fn subscribe_top_bids(
diff --git a/relay-server/Cargo.toml b/relay-server/Cargo.toml
@@ -8,6 +8,7 @@ async-trait.workspace = true
 axum.workspace = true
 bytes.workspace = true
 ethereum_ssz.workspace = true
+flate2.workspace = true
 futures.workspace = true
 http.workspace = true
 relay-api-types = { path = "../relay-api-types" }
diff --git a/relay-server/src/server.rs b/relay-server/src/server.rs
@@ -7,17 +7,20 @@ use axum::{
     extract::{FromRequest, Query, Request, State},
     response::{IntoResponse, Response},
     routing::{get, post},
-    Json, RequestExt, Router,
+    Router,
 };
 use bytes::Bytes;
+use flate2::read::GzDecoder;
 use futures::{sink::SinkExt, stream::StreamExt};
+use http::header::CONTENT_ENCODING;
 use http::{header::CONTENT_TYPE, HeaderValue, StatusCode};
 use relay_api_types::{
-    ErrorResponse, GetDeliveredPayloadsQueryParams, GetReceivedBidsQueryParams,
-    GetValidatorRegistrationQueryParams, SignedCancellation, SignedHeaderSubmission,
-    SubmitBlockQueryParams, SubmitBlockRequest,
+    ContentEncoding, ContentType, ErrorResponse, GetDeliveredPayloadsQueryParams,
+    GetReceivedBidsQueryParams, GetValidatorRegistrationQueryParams, SignedCancellation,
+    SignedHeaderSubmission, SubmitBlockQueryParams, SubmitBlockRequest,
 };
 use serde::Serialize;
+use std::io::Read;
 use std::net::SocketAddr;
 use tracing::error;
 use types::eth_spec::EthSpec;
@@ -139,7 +142,7 @@ where
 async fn submit_block<I, A, E>(
     Query(query_params): Query<SubmitBlockQueryParams>,
     State(api_impl): State<I>,
-    JsonOrSsz(body): JsonOrSsz<SubmitBlockRequest<E>>,
+    JsonOrSszMaybeGzipped(body): JsonOrSszMaybeGzipped<SubmitBlockRequest<E>>,
 ) -> Result<Response<Body>, StatusCode>
 where
     E: EthSpec,
@@ -155,7 +158,7 @@ where
 async fn submit_block_optimistic_v2<I, A, E>(
     Query(query_params): Query<SubmitBlockQueryParams>,
     State(api_impl): State<I>,
-    JsonOrSsz(body): JsonOrSsz<SubmitBlockRequest<E>>,
+    JsonOrSszMaybeGzipped(body): JsonOrSszMaybeGzipped<SubmitBlockRequest<E>>,
 ) -> Result<Response<Body>, StatusCode>
 where
     E: EthSpec,
@@ -174,7 +177,7 @@ where
 async fn submit_header<I, A, E>(
     Query(query_params): Query<SubmitBlockQueryParams>,
     State(api_impl): State<I>,
-    JsonOrSsz(body): JsonOrSsz<SignedHeaderSubmission<E>>,
+    JsonOrSszMaybeGzipped(body): JsonOrSszMaybeGzipped<SignedHeaderSubmission<E>>,
 ) -> Result<Response<Body>, StatusCode>
 where
     E: EthSpec,
@@ -189,7 +192,7 @@ where
 #[tracing::instrument(skip_all)]
 async fn submit_cancellation<I, A, E>(
     State(api_impl): State<I>,
-    JsonOrSsz(body): JsonOrSsz<SignedCancellation>,
+    JsonOrSszMaybeGzipped(body): JsonOrSszMaybeGzipped<SignedCancellation>,
 ) -> Result<Response<Body>, StatusCode>
 where
     E: EthSpec,
@@ -357,28 +360,50 @@ where
 
 #[must_use]
 #[derive(Debug, Clone, Copy, Default)]
-struct JsonOrSsz<T>(T);
+struct JsonOrSszMaybeGzipped<T>(T);
 
 #[async_trait]
-impl<T, S> FromRequest<S> for JsonOrSsz<T>
+impl<T, S> FromRequest<S> for JsonOrSszMaybeGzipped<T>
 where
     T: serde::de::DeserializeOwned + ssz::Decode + 'static,
     S: Send + Sync,
 {
     type Rejection = Response;
 
     async fn from_request(req: Request, _state: &S) -> Result<Self, Self::Rejection> {
-        let content_type_header = req.headers().get(CONTENT_TYPE);
-        let content_type = content_type_header.and_then(|value| value.to_str().ok());
+        let headers = req.headers().clone();
+        let content_type = headers
+            .get(CONTENT_TYPE)
+            .and_then(|value| value.to_str().ok());
+        let content_encoding = headers
+            .get(CONTENT_ENCODING)
+            .and_then(|value| value.to_str().ok());
+
+        let bytes = Bytes::from_request(req, _state)
+            .await
+            .map_err(IntoResponse::into_response)?;
+
+        let decoded_bytes = if content_encoding == Some(&ContentEncoding::Gzip.to_string()) {
+            let mut decoder = GzDecoder::new(&bytes[..]);
+            let mut decoded = Vec::new();
+            decoder
+                .read_to_end(&mut decoded)
+                .map_err(|_| StatusCode::BAD_REQUEST.into_response())?;
+            decoded
+        } else {
+            bytes.to_vec()
+        };
 
         if let Some(content_type) = content_type {
-            if content_type.starts_with("application/json") {
-                let Json(payload) = req.extract().await.map_err(IntoResponse::into_response)?;
+            if content_type.starts_with(&ContentType::Json.to_string()) {
+                let payload: T = serde_json::from_slice(&decoded_bytes)
+                    .map_err(|_| StatusCode::BAD_REQUEST.into_response())?;
                 return Ok(Self(payload));
             }
 
-            if content_type.starts_with("application/octet-stream") {
-                let Ssz(payload) = req.extract().await.map_err(IntoResponse::into_response)?;
+            if content_type.starts_with(&ContentType::Ssz.to_string()) {
+                let payload = T::from_ssz_bytes(&decoded_bytes)
+                    .map_err(|_| StatusCode::BAD_REQUEST.into_response())?;
                 return Ok(Self(payload));
             }
         }
