vendor encoders from async-compression

This integrates the encoders from
https://github.com/Nemo157/async-compression so we can use them other a
stream that flushes data regularly

Author: Geal

Cargo.lock

@@ -286,6 +286,7 @@ dependencies = [
  "axum",
  "backtrace",
  "base64 0.20.0",
+ "brotli",
  "buildstructor 0.5.2",
  "bytes",
  "ci_info",
@@ -401,6 +402,8 @@ dependencies = [
  "wiremock",
  "wsl",
  "yaml-rust",
+ "zstd",
+ "zstd-safe",
 ]
 
 [[package]]
@@ -7073,3 +7076,33 @@ dependencies = [
  "quote",
  "syn 2.0.13",
 ]
+
+[[package]]
+name = "zstd"
+version = "0.12.3+zstd.1.5.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "76eea132fb024e0e13fd9c2f5d5d595d8a967aa72382ac2f9d39fcc95afd0806"
+dependencies = [
+ "zstd-safe",
+]
+
+[[package]]
+name = "zstd-safe"
+version = "6.0.5+zstd.1.5.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d56d9e60b4b1758206c238a10165fbcae3ca37b01744e394c463463f6529d23b"
+dependencies = [
+ "libc",
+ "zstd-sys",
+]
+
+[[package]]
+name = "zstd-sys"
+version = "2.0.8+zstd.1.5.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5556e6ee25d32df2586c098bbfa278803692a20d0ab9565e049480d52707ec8c"
+dependencies = [
+ "cc",
+ "libc",
+ "pkg-config",
+]

apollo-router/Cargo.toml

@@ -203,6 +203,10 @@ yaml-rust = "0.4.5"
 wsl = "0.1.0"
 tokio-rustls = "0.23.4"
 http-serde = "1.1.2"
+memchr = "2.5.0"
+brotli = "3.3.4"
+zstd = "0.12.3"
+zstd-safe = "6.0.5"
 
 [target.'cfg(macos)'.dependencies]
 uname = "0.1.1"

apollo-router/src/axum_factory/compression/codec/brotli/encoder.rs

@@ -0,0 +1,110 @@
+// All code from this module is extracted from https://github.com/Nemo157/async-compression and is under MIT or Apache-2 licence
+// it will be removed when we find a long lasting solution to https://github.com/Nemo157/async-compression/issues/154
+use crate::axum_factory::compression::{codec::Encode, util::PartialBuffer};
+use std::{
+    fmt,
+    io::{Error, ErrorKind, Result},
+};
+
+use brotli::enc::{
+    backward_references::BrotliEncoderParams,
+    encode::{
+        BrotliEncoderCompressStream, BrotliEncoderCreateInstance, BrotliEncoderHasMoreOutput,
+        BrotliEncoderIsFinished, BrotliEncoderOperation, BrotliEncoderStateStruct,
+    },
+    StandardAlloc,
+};
+
+pub(crate) struct BrotliEncoder {
+    state: BrotliEncoderStateStruct<StandardAlloc>,
+}
+
+impl BrotliEncoder {
+    pub(crate) fn new(params: BrotliEncoderParams) -> Self {
+        let mut state = BrotliEncoderCreateInstance(StandardAlloc::default());
+        state.params = params;
+        Self { state }
+    }
+
+    fn encode(
+        &mut self,
+        input: &mut PartialBuffer<impl AsRef<[u8]>>,
+        output: &mut PartialBuffer<impl AsRef<[u8]> + AsMut<[u8]>>,
+        op: BrotliEncoderOperation,
+    ) -> Result<()> {
+        let in_buf = input.unwritten();
+        let mut out_buf = output.unwritten_mut();
+
+        let mut input_len = 0;
+        let mut output_len = 0;
+
+        if BrotliEncoderCompressStream(
+            &mut self.state,
+            op,
+            &mut in_buf.len(),
+            in_buf,
+            &mut input_len,
+            &mut out_buf.len(),
+            out_buf,
+            &mut output_len,
+            &mut None,
+            &mut |_, _, _, _| (),
+        ) <= 0
+        {
+            return Err(Error::new(ErrorKind::Other, "brotli error"));
+        }
+
+        input.advance(input_len);
+        output.advance(output_len);
+
+        Ok(())
+    }
+}
+
+impl Encode for BrotliEncoder {
+    fn encode(
+        &mut self,
+        input: &mut PartialBuffer<impl AsRef<[u8]>>,
+        output: &mut PartialBuffer<impl AsRef<[u8]> + AsMut<[u8]>>,
+    ) -> Result<()> {
+        self.encode(
+            input,
+            output,
+            BrotliEncoderOperation::BROTLI_OPERATION_PROCESS,
+        )
+    }
+
+    fn flush(
+        &mut self,
+        output: &mut PartialBuffer<impl AsRef<[u8]> + AsMut<[u8]>>,
+    ) -> Result<bool> {
+        self.encode(
+            &mut PartialBuffer::new(&[][..]),
+            output,
+            BrotliEncoderOperation::BROTLI_OPERATION_FLUSH,
+        )?;
+
+        Ok(BrotliEncoderHasMoreOutput(&self.state) == 0)
+    }
+
+    fn finish(
+        &mut self,
+        output: &mut PartialBuffer<impl AsRef<[u8]> + AsMut<[u8]>>,
+    ) -> Result<bool> {
+        self.encode(
+            &mut PartialBuffer::new(&[][..]),
+            output,
+            BrotliEncoderOperation::BROTLI_OPERATION_FINISH,
+        )?;
+
+        Ok(BrotliEncoderIsFinished(&self.state) == 1)
+    }
+}
+
+impl fmt::Debug for BrotliEncoder {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_struct("BrotliEncoder")
+            .field("compress", &"<no debug>")
+            .finish()
+    }
+}

apollo-router/src/axum_factory/compression/codec/brotli/mod.rs

@@ -0,0 +1,5 @@
+// All code from this module is extracted from https://github.com/Nemo157/async-compression and is under MIT or Apache-2 licence
+// it will be removed when we find a long lasting solution to https://github.com/Nemo157/async-compression/issues/154
+mod encoder;
+
+pub(crate) use self::encoder::BrotliEncoder;

apollo-router/src/axum_factory/compression/codec/deflate/encoder.rs

@@ -0,0 +1,46 @@
+// All code from this module is extracted from https://github.com/Nemo157/async-compression and is under MIT or Apache-2 licence
+// it will be removed when we find a long lasting solution to https://github.com/Nemo157/async-compression/issues/154
+use crate::axum_factory::compression::{
+    codec::{Encode, FlateEncoder},
+    util::PartialBuffer,
+};
+use std::io::Result;
+
+use flate2::Compression;
+
+#[derive(Debug)]
+pub(crate) struct DeflateEncoder {
+    inner: FlateEncoder,
+}
+
+impl DeflateEncoder {
+    pub(crate) fn new(level: Compression) -> Self {
+        Self {
+            inner: FlateEncoder::new(level, false),
+        }
+    }
+}
+
+impl Encode for DeflateEncoder {
+    fn encode(
+        &mut self,
+        input: &mut PartialBuffer<impl AsRef<[u8]>>,
+        output: &mut PartialBuffer<impl AsRef<[u8]> + AsMut<[u8]>>,
+    ) -> Result<()> {
+        self.inner.encode(input, output)
+    }
+
+    fn flush(
+        &mut self,
+        output: &mut PartialBuffer<impl AsRef<[u8]> + AsMut<[u8]>>,
+    ) -> Result<bool> {
+        self.inner.flush(output)
+    }
+
+    fn finish(
+        &mut self,
+        output: &mut PartialBuffer<impl AsRef<[u8]> + AsMut<[u8]>>,
+    ) -> Result<bool> {
+        self.inner.finish(output)
+    }
+}

apollo-router/src/axum_factory/compression/codec/deflate/mod.rs

@@ -0,0 +1,5 @@
+// All code from this module is extracted from https://github.com/Nemo157/async-compression and is under MIT or Apache-2 licence
+// it will be removed when we find a long lasting solution to https://github.com/Nemo157/async-compression/issues/154
+mod encoder;
+
+pub(crate) use self::encoder::DeflateEncoder;

apollo-router/src/axum_factory/compression/codec/flate/encoder.rs

@@ -0,0 +1,104 @@
+// All code from this module is extracted from https://github.com/Nemo157/async-compression and is under MIT or Apache-2 licence
+// it will be removed when we find a long lasting solution to https://github.com/Nemo157/async-compression/issues/154
+use std::io::{Error, ErrorKind, Result};
+
+use flate2::{Compress, Compression, FlushCompress, Status};
+
+use crate::axum_factory::compression::{codec::Encode, util::PartialBuffer};
+
+#[derive(Debug)]
+pub(crate) struct FlateEncoder {
+    compress: Compress,
+    flushed: bool,
+}
+
+impl FlateEncoder {
+    pub(crate) fn new(level: Compression, zlib_header: bool) -> Self {
+        Self {
+            compress: Compress::new(level, zlib_header),
+            flushed: true,
+        }
+    }
+
+    fn encode(
+        &mut self,
+        input: &mut PartialBuffer<impl AsRef<[u8]>>,
+        output: &mut PartialBuffer<impl AsRef<[u8]> + AsMut<[u8]>>,
+        flush: FlushCompress,
+    ) -> Result<Status> {
+        let prior_in = self.compress.total_in();
+        let prior_out = self.compress.total_out();
+
+        let status = self
+            .compress
+            .compress(input.unwritten(), output.unwritten_mut(), flush)?;
+
+        input.advance((self.compress.total_in() - prior_in) as usize);
+        output.advance((self.compress.total_out() - prior_out) as usize);
+
+        Ok(status)
+    }
+}
+
+impl Encode for FlateEncoder {
+    fn encode(
+        &mut self,
+        input: &mut PartialBuffer<impl AsRef<[u8]>>,
+        output: &mut PartialBuffer<impl AsRef<[u8]> + AsMut<[u8]>>,
+    ) -> Result<()> {
+        self.flushed = false;
+        match self.encode(input, output, FlushCompress::None)? {
+            Status::Ok => Ok(()),
+            Status::StreamEnd => unreachable!(),
+            Status::BufError => Err(Error::new(ErrorKind::Other, "unexpected BufError")),
+        }
+    }
+
+    fn flush(
+        &mut self,
+        output: &mut PartialBuffer<impl AsRef<[u8]> + AsMut<[u8]>>,
+    ) -> Result<bool> {
+        // We need to keep track of whether we've already flushed otherwise we'll just keep writing
+        // out sync blocks continuously and probably never complete flushing.
+        if self.flushed {
+            return Ok(true);
+        }
+
+        self.encode(
+            &mut PartialBuffer::new(&[][..]),
+            output,
+            FlushCompress::Sync,
+        )?;
+
+        loop {
+            let old_len = output.written().len();
+            self.encode(
+                &mut PartialBuffer::new(&[][..]),
+                output,
+                FlushCompress::None,
+            )?;
+            if output.written().len() == old_len {
+                break;
+            }
+        }
+
+        self.flushed = true;
+        Ok(!output.unwritten().is_empty())
+    }
+
+    fn finish(
+        &mut self,
+        output: &mut PartialBuffer<impl AsRef<[u8]> + AsMut<[u8]>>,
+    ) -> Result<bool> {
+        self.flushed = false;
+        match self.encode(
+            &mut PartialBuffer::new(&[][..]),
+            output,
+            FlushCompress::Finish,
+        )? {
+            Status::Ok => Ok(false),
+            Status::StreamEnd => Ok(true),
+            Status::BufError => Err(Error::new(ErrorKind::Other, "unexpected BufError")),
+        }
+    }
+}

apollo-router/src/axum_factory/compression/codec/flate/mod.rs

@@ -0,0 +1,5 @@
+// All code from this module is extracted from https://github.com/Nemo157/async-compression and is under MIT or Apache-2 licence
+// it will be removed when we find a long lasting solution to https://github.com/Nemo157/async-compression/issues/154
+mod encoder;
+
+pub(crate) use self::encoder::FlateEncoder;