feat(executor): handle auto decompression of subgraph responses

Author: ardatan

Cargo.lock

@@ -45,11 +45,11 @@ reqwest-retry = { workspace = true }
 reqwest-middleware = { workspace = true }
 vrl = { workspace = true }
 serde_json = { workspace = true }
+tokio-util = { workspace = true }
 
 mimalloc = { version = "0.1.48", features = ["v3"] }
 moka = { version = "0.12.10", features = ["future"] }
 ulid = "1.2.1"
-tokio-util = "0.7.16"
 cookie = "0.18.1"
 regex-automata = "0.4.10"
 arc-swap = "1.7.1"

bin/router/Cargo.toml

@@ -37,6 +37,7 @@ strum = { version = "0.27.2", features = ["derive"] }
 ntex-http = "0.1.15"
 ordered-float = "4.2.0"
 hyper-tls = { version = "0.6.0", features = ["vendored"] }
+async-compression = { version = "0.4.3", features = ["all"]}
 hyper-util = { version = "0.1.16", features = [
   "client",
   "client-legacy",

lib/executor/Cargo.toml

@@ -0,0 +1,107 @@
+use async_compression::futures::bufread::{
+    BrotliDecoder, DeflateDecoder, GzipDecoder, ZstdDecoder,
+};
+use bytes::Bytes;
+use futures::{future::BoxFuture, AsyncReadExt};
+
+use crate::executors::error::SubgraphExecutorError;
+
+pub enum CompressionType {
+    Gzip,
+    Deflate,
+    Brotli,
+    Zstd,
+    Multiple(Vec<CompressionType>),
+    Identity,
+}
+
+
+impl CompressionType {
+    pub fn header_value(&self) -> String {
+        match self {
+            CompressionType::Gzip => "gzip".to_string(),
+            CompressionType::Deflate => "deflate".to_string(),
+            CompressionType::Brotli => "br".to_string(),
+            CompressionType::Zstd => "zstd".to_string(),
+            CompressionType::Identity => "identity".to_string(),
+            CompressionType::Multiple(types) => types
+                .iter()
+                .map(|t| t.header_value().to_string())
+                .collect::<Vec<String>>()
+                .join(", ")
+                .to_string(),
+        }
+    }
+    pub fn from_encoding_header(
+        encoding: &str,
+    ) -> Result<CompressionType, SubgraphExecutorError> {
+        let encodings: Vec<&str> = encoding.split(',').map(|s| s.trim()).collect();
+        if encodings.len() > 1 {
+            let types = encodings
+                .iter()
+                .map(|&e| CompressionType::from_encoding_header(e))
+                .collect::<Result<Vec<CompressionType>, SubgraphExecutorError>>()?;
+            Ok(CompressionType::Multiple(types))
+        } else {
+            match encodings[0].to_lowercase().as_str() {
+                "gzip" => Ok(CompressionType::Gzip),
+                "deflate" => Ok(CompressionType::Deflate),
+                "br" => Ok(CompressionType::Brotli),
+                "zstd" => Ok(CompressionType::Zstd),
+                "identity" => Ok(CompressionType::Identity),
+                "none" => Ok(CompressionType::Identity),
+                _ => Err(SubgraphExecutorError::UnknownEncoding(encoding.to_string())),
+            }
+        }
+    }
+    pub fn decompress<'a>(&'a self, body: Bytes) -> BoxFuture<'a, Result<Bytes, SubgraphExecutorError>> {
+        Box::pin(async move {
+            match self {
+                CompressionType::Gzip => {
+                    let mut decoder = GzipDecoder::new(body.as_ref());
+                    let mut buf = Vec::new();
+                    let _ = decoder
+                        .read_to_end(&mut buf)
+                        .await
+                        .map_err(|e| SubgraphExecutorError::DecompressionFailed(self.header_value(), e.to_string()));
+                    Ok(Bytes::from(buf))
+                }
+                CompressionType::Deflate => {
+                    let mut decoder = DeflateDecoder::new(body.as_ref());
+                    let mut buf = Vec::new();
+                    let _ = decoder
+                        .read_to_end(&mut buf)
+                        .await
+                        .map_err(|e| SubgraphExecutorError::DecompressionFailed(self.header_value(), e.to_string()));
+                    Ok(Bytes::from(buf))
+                }
+                CompressionType::Brotli => {
+                    let mut decoder = BrotliDecoder::new(body.as_ref());
+                    let mut buf = Vec::new();
+                    let _ = decoder
+                        .read_to_end(&mut buf)
+                        .await
+                        .map_err(|e| SubgraphExecutorError::DecompressionFailed(self.header_value(), e.to_string()));
+                    Ok(Bytes::from(buf))
+                }
+                CompressionType::Zstd => {
+                    let mut decoder = ZstdDecoder::new(body.as_ref());
+                    let mut buf = Vec::new();
+                    let _ = decoder
+                        .read_to_end(&mut buf)
+                        .await
+                        .map_err(|e| SubgraphExecutorError::DecompressionFailed(self.header_value(), e.to_string()));
+                    Ok(Bytes::from(buf))
+                }
+                CompressionType::Multiple(types) => {
+                    let mut decompressed_body = body;
+                    for ctype in types {
+                        decompressed_body = ctype.decompress(decompressed_body).await?;
+                    }
+                    Ok(decompressed_body)
+                }
+                CompressionType::Identity => Ok(body),
+            }
+        })
+    }
+}

lib/executor/src/executors/compression.rs

@@ -20,6 +20,10 @@ pub enum SubgraphExecutorError {
     RequestFailure(String, String),
     #[error("Failed to serialize variable \"{0}\": {1}")]
     VariablesSerializationFailure(String, String),
+    #[error("An unknown encoding \"{0}\" was specified in the 'Content-Encoding' header.")]
+    UnknownEncoding(String),
+    #[error("Decompression failed for encoding \"{0}\": {1}")]
+    DecompressionFailed(String, String),
 }
 
 impl From<SubgraphExecutorError> for GraphQLError {
@@ -75,7 +79,9 @@ impl SubgraphExecutorError {
             SubgraphExecutorError::RequestFailure(_, _) => "SUBGRAPH_REQUEST_FAILURE",
             SubgraphExecutorError::VariablesSerializationFailure(_, _) => {
                 "SUBGRAPH_VARIABLES_SERIALIZATION_FAILURE"
-            }
+            },
+            SubgraphExecutorError::UnknownEncoding(_) => "SUBGRAPH_UNKNOWN_ENCODING",
+            SubgraphExecutorError::DecompressionFailed(_, _) => "SUBGRAPH_DECOMPRESSION_FAILED",
         }
     }
 }