diff --git a/grpc/Cargo.toml b/grpc/Cargo.toml index 01f1eb75c..748524dba 100644 --- a/grpc/Cargo.toml +++ b/grpc/Cargo.toml @@ -48,10 +48,14 @@ tls-rustls = [ "dep:rustls-platform-verifier", "dep:rustls-webpki", ] +gzip = ["dep:flate2"] [dependencies] base64 = "0.22" bytes = "1.10.1" +flate2 = { version = "1.0", optional = true, default-features = false, features = [ + "zlib-rs", +] } futures = { version = "0.3", default-features = false, optional = true } hickory-resolver = { version = "0.26.1", optional = true } http = "1.1.0" diff --git a/grpc/src/codec.rs b/grpc/src/codec.rs new file mode 100644 index 000000000..f2098dbff --- /dev/null +++ b/grpc/src/codec.rs @@ -0,0 +1,26 @@ +/* + * + * Copyright 2026 gRPC authors. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to + * deal in the Software without restriction, including without limitation the + * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or + * sell copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING + * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS + * IN THE SOFTWARE. + * + */ + +pub(crate) mod compression; +pub(crate) mod message; diff --git a/grpc/src/codec/compression.rs b/grpc/src/codec/compression.rs new file mode 100644 index 000000000..224b237d3 --- /dev/null +++ b/grpc/src/codec/compression.rs @@ -0,0 +1,61 @@ +/* + * + * Copyright 2026 gRPC authors. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to + * deal in the Software without restriction, including without limitation the + * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or + * sell copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING + * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS + * IN THE SOFTWARE. + * + */ + +use bytes::Buf; +use bytes::BufMut; + +#[cfg(feature = "gzip")] +mod gzip; + +pub(crate) mod registry; + +/// A trait for compressing outgoing gRPC payloads. +pub trait Compressor: Send + Sync + 'static { + /// The canonical gRPC content coding name (e.g., "gzip"). + fn name(&self) -> &str; + + /// Compress data from `source` into `destination`. + /// + /// # Errors + /// + /// Returns an `io::Error` if compression fails. Implementations should gracefully + /// handle constrained `destination` buffers by returning an error rather than panicking + /// (e.g., by verifying `destination.remaining_mut()` is sufficient before writing). + fn compress(&self, source: &mut dyn Buf, destination: &mut dyn BufMut) -> Result<(), String>; +} + +/// A trait for decompressing incoming gRPC payloads. +pub trait Decompressor: Send + Sync + 'static { + /// The canonical gRPC content coding name (e.g., "gzip"). + fn name(&self) -> &str; + + /// Decompress data from `source` into `destination`. + /// + /// # Errors + /// + /// Returns an `io::Error` if decompression fails. Implementations should gracefully + /// handle constrained `destination` buffers by returning an error rather than panicking + /// (e.g., by verifying `destination.remaining_mut()` is sufficient before writing). + fn decompress(&self, source: &mut dyn Buf, destination: &mut dyn BufMut) -> Result<(), String>; +} diff --git a/grpc/src/codec/compression/gzip.rs b/grpc/src/codec/compression/gzip.rs new file mode 100644 index 000000000..d715eb849 --- /dev/null +++ b/grpc/src/codec/compression/gzip.rs @@ -0,0 +1,414 @@ +/* + * + * Copyright 2026 gRPC authors. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to + * deal in the Software without restriction, including without limitation the + * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or + * sell copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING + * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS + * IN THE SOFTWARE. + * + */ + +use bytes::Buf; +use bytes::BufMut; +use flate2::Compress; +use flate2::Compression as FlateCompression; +use flate2::Decompress; +use flate2::FlushCompress; +use flate2::FlushDecompress; +use flate2::Status; + +use crate::codec::compression::Compressor; +use crate::codec::compression::Decompressor; + +/// The base-2 logarithm of the gzip sliding window size. 15 is the maximum +/// (and standard) value; `flate2` adds the gzip framing on top of it. +const GZIP_WINDOW_BITS: u8 = 15; + +/// A gzip compression implementation. +#[derive(Debug, Clone, Copy)] +pub struct Gzip { + level: FlateCompression, +} + +impl Gzip { + /// Creates a new gzip compression implementation. + pub fn new() -> Self { + Self { + level: FlateCompression::new(6), + } + } +} + +impl Default for Gzip { + fn default() -> Self { + Self::new() + } +} + +impl Compressor for Gzip { + fn name(&self) -> &str { + "gzip" + } + + /// Compresses `source` into `destination`, writing the gzip output directly + /// into the destination's spare capacity (no intermediate buffer). + /// + /// # Errors + /// + /// Returns an error if compression fails or if `destination` is a + /// fixed-capacity sink too small to hold the output. On error, the contents + /// of `destination` are unspecified and should be discarded by the caller. + fn compress(&self, source: &mut dyn Buf, destination: &mut dyn BufMut) -> Result<(), String> { + let mut compressor = Compress::new_gzip(self.level, GZIP_WINDOW_BITS); + loop { + let input = source.chunk(); + // Once every source byte lives in the current chunk, ask the codec + // to emit the final gzip trailer via `Finish`. + let flush = if source.remaining() == input.len() { + FlushCompress::Finish + } else { + FlushCompress::None + }; + + // Writable spare capacity of the destination. `chunk_mut()` grows + // growable sinks (`Vec`/`BytesMut`); a fixed sink may yield an empty + // region once full. + // SAFETY: `as_uninit_slice_mut` exposes the destination's + // uninitialized spare capacity. We only ever write into it below (via + // the codec) and never read it before writing. + let output = unsafe { destination.chunk_mut().as_uninit_slice_mut() }; + if output.is_empty() { + return Err("gzip: compression destination buffer is full".to_string()); + } + + let before_in = compressor.total_in(); + let before_out = compressor.total_out(); + let status = compressor + .compress_uninit(input, output, flush) + .map_err(|e| e.to_string())?; + let consumed = (compressor.total_in() - before_in) as usize; + let produced = (compressor.total_out() - before_out) as usize; + + source.advance(consumed); + // SAFETY: the codec initialized exactly `produced` bytes at the start + // of the region obtained above, and `produced <= output.len()` per + // flate2's contract. + unsafe { destination.advance_mut(produced) }; + + match status { + Status::StreamEnd => return Ok(()), + Status::Ok | Status::BufError => { + if consumed == 0 && produced == 0 { + return Err("gzip: compression stalled with no progress".to_string()); + } + } + } + } + } +} + +impl Decompressor for Gzip { + fn name(&self) -> &str { + "gzip" + } + + /// Decompresses `source` into `destination`, writing the inflated output + /// directly into the destination's spare capacity (no intermediate buffer). + /// + /// `source` must contain exactly one gzip member. A gRPC message frame + /// carries a single compressed member, so any bytes remaining after the end + /// of the gzip stream — trailing garbage or a concatenated additional + /// member — are treated as malformed and rejected. + /// + /// # Errors + /// + /// Returns an error if the stream is corrupt or truncated, if there is + /// trailing data after the gzip member, or if `destination` is a + /// fixed-capacity sink too small to hold the output. On error, the contents + /// of `destination` are unspecified and should be discarded by the caller. + fn decompress(&self, source: &mut dyn Buf, destination: &mut dyn BufMut) -> Result<(), String> { + let mut decompressor = Decompress::new_gzip(GZIP_WINDOW_BITS); + loop { + let input = source.chunk(); + + // Writable spare capacity of the destination. `chunk_mut()` grows + // growable sinks (`Vec`/`BytesMut`); a fixed sink may yield an empty + // region once full. + // SAFETY: `as_uninit_slice_mut` exposes the destination's + // uninitialized spare capacity. We only ever write into it below (via + // the codec) and never read it before writing. + let output = unsafe { destination.chunk_mut().as_uninit_slice_mut() }; + if output.is_empty() { + return Err("gzip: decompression destination buffer is full".to_string()); + } + + let before_in = decompressor.total_in(); + let before_out = decompressor.total_out(); + let status = decompressor + .decompress_uninit(input, output, FlushDecompress::None) + .map_err(|e| e.to_string())?; + let consumed = (decompressor.total_in() - before_in) as usize; + let produced = (decompressor.total_out() - before_out) as usize; + + source.advance(consumed); + // SAFETY: the codec initialized exactly `produced` bytes at the start + // of the region obtained above, and `produced <= output.len()` per + // flate2's contract. + unsafe { destination.advance_mut(produced) }; + + match status { + Status::StreamEnd => { + // Reject any bytes after the end of the gzip stream (see the + // single-member contract in the method docs). + if source.has_remaining() { + return Err("gzip: trailing data after end of gzip stream".to_string()); + } + return Ok(()); + } + Status::Ok | Status::BufError => { + // No forward progress means the input is truncated (needs + // more bytes) with nothing left to feed it. + if consumed == 0 && produced == 0 { + return Err("gzip: truncated or corrupt gzip stream".to_string()); + } + } + } + } + } +} + +#[cfg(test)] +mod tests { + use bytes::Bytes; + + use super::*; + + #[test] + fn gzip_compress_decompress() { + let compressor = Gzip::new(); + let data = Bytes::from_static(b"hello world"); + let mut compressed = Vec::new(); + compressor + .compress(&mut data.clone(), &mut compressed) + .unwrap(); + + assert_ne!(compressed.as_slice(), data); + let mut decompressed = Vec::new(); + compressor + .decompress(&mut compressed.as_slice(), &mut decompressed) + .unwrap(); + assert_eq!(data, decompressed.as_slice()); + } + + /// Compresses then decompresses `data` through growable sinks that start + /// empty, forcing repeated `chunk_mut()` growth, and asserts the round trip + /// reproduces the original bytes. + fn assert_round_trip(data: &[u8]) { + let gzip = Gzip::new(); + + let mut compressed = Vec::new(); + gzip.compress(&mut Bytes::copy_from_slice(data), &mut compressed) + .unwrap(); + + let mut decompressed = Vec::new(); + gzip.decompress(&mut compressed.as_slice(), &mut decompressed) + .unwrap(); + + assert_eq!(data, decompressed.as_slice()); + } + + #[test] + fn round_trip_empty() { + assert_round_trip(b""); + } + + #[test] + fn round_trip_small() { + assert_round_trip(b"hello world"); + } + + #[test] + fn round_trip_large_compressible() { + // > 32KB of highly compressible data to force multiple grow cycles. + let data = vec![b'a'; 100 * 1024]; + assert_round_trip(&data); + } + + #[test] + fn compress_into_too_small_buffer_errors() { + let gzip = Gzip::new(); + let data = vec![b'x'; 8 * 1024]; + + // A fixed, non-growable sink far too small for the gzip output. + let mut tiny = [0u8; 4]; + let mut sink: &mut [u8] = &mut tiny; + let result = gzip.compress(&mut Bytes::copy_from_slice(&data), &mut sink); + + // Must return Err (not panic, not overrun). + assert!(result.is_err()); + } + + #[test] + fn decompress_into_too_small_buffer_errors() { + let gzip = Gzip::new(); + let data = vec![b'y'; 8 * 1024]; + let mut compressed = Vec::new(); + gzip.compress(&mut Bytes::copy_from_slice(&data), &mut compressed) + .unwrap(); + + // Decompress into a fixed sink too small to hold the 8KB output. + let mut tiny = [0u8; 16]; + let mut sink: &mut [u8] = &mut tiny; + let result = gzip.decompress(&mut compressed.as_slice(), &mut sink); + + assert!(result.is_err()); + } + + #[test] + fn compress_does_not_write_past_fixed_region() { + // Guard-byte canary: the destination is a sub-slice of a larger backing + // array; the trailing sentinel bytes must remain untouched even when the + // destination fills up and compression fails. + let gzip = Gzip::new(); + let data = vec![b'z'; 8 * 1024]; + + const DEST_LEN: usize = 8; + const SENTINEL: u8 = 0xAB; + let mut backing = [SENTINEL; DEST_LEN + 16]; + { + let (dest, _guard) = backing.split_at_mut(DEST_LEN); + let mut sink: &mut [u8] = dest; + // Expected to error since the destination is far too small. + let _ = gzip.compress(&mut Bytes::copy_from_slice(&data), &mut sink); + } + // The guard region after the destination must be pristine. + assert!(backing[DEST_LEN..].iter().all(|&b| b == SENTINEL)); + } + + #[test] + fn decompress_split_across_header() { + // The gzip header is 10 bytes; split the source inside it (at byte 4) so + // the decoder must resume header parsing across a chunk boundary. + let data = b"streaming across the gzip header boundary"; + let mut compressed = Vec::new(); + Gzip::new() + .compress(&mut Bytes::copy_from_slice(data), &mut compressed) + .unwrap(); + assert!(compressed.len() > 4); + + let mut source = (&compressed[..4]).chain(&compressed[4..]); + let mut decompressed = Vec::new(); + Gzip::new() + .decompress(&mut source, &mut decompressed) + .unwrap(); + assert_eq!(decompressed.as_slice(), data); + } + + #[test] + fn decompress_split_across_trailer() { + // The 8-byte CRC32/ISIZE trailer sits at the end; split within it so the + // decoder must consume the trailer across a chunk boundary. + let data = b"streaming across the gzip trailer boundary"; + let mut compressed = Vec::new(); + Gzip::new() + .compress(&mut Bytes::copy_from_slice(data), &mut compressed) + .unwrap(); + let split = compressed.len() - 4; + + let mut source = (&compressed[..split]).chain(&compressed[split..]); + let mut decompressed = Vec::new(); + Gzip::new() + .decompress(&mut source, &mut decompressed) + .unwrap(); + assert_eq!(decompressed.as_slice(), data); + } + + #[test] + fn decompress_truncated_stream_errors() { + // Dropping the trailing bytes of a valid stream leaves the decoder + // wanting more input it never receives -> truncation error. + let mut compressed = Vec::new(); + Gzip::new() + .compress( + &mut Bytes::from_static(b"a truncated gzip stream must be rejected"), + &mut compressed, + ) + .unwrap(); + let truncated = &compressed[..compressed.len() - 4]; + + let mut decompressed = Vec::new(); + let result = Gzip::new().decompress(&mut &truncated[..], &mut decompressed); + assert!(result.is_err()); + } + + #[test] + fn decompress_corrupt_trailer_errors() { + // Flipping a byte in the CRC32/ISIZE trailer must fail integrity checks. + let mut compressed = Vec::new(); + Gzip::new() + .compress( + &mut Bytes::from_static(b"corrupt trailer must be detected"), + &mut compressed, + ) + .unwrap(); + let last = compressed.len() - 1; + compressed[last] ^= 0xFF; + + let mut decompressed = Vec::new(); + let result = Gzip::new().decompress(&mut compressed.as_slice(), &mut decompressed); + assert!(result.is_err()); + } + + #[test] + fn decompress_trailing_data_errors() { + // Matches gRPC C-core: bytes after a complete gzip member are trailing + // data and must be rejected (a gRPC frame carries exactly one member). + let data = b"member followed by trailing garbage"; + let mut compressed = Vec::new(); + Gzip::new() + .compress(&mut Bytes::copy_from_slice(data), &mut compressed) + .unwrap(); + compressed.extend_from_slice(b"trailing garbage not part of the stream"); + + let mut decompressed = Vec::new(); + let result = Gzip::new().decompress(&mut compressed.as_slice(), &mut decompressed); + assert!(result.is_err()); + } + + #[test] + fn decompress_multi_member_errors() { + // Matches gRPC C-core: a concatenated second gzip member is trailing + // data after the first member's end, so decoding must error rather than + // silently decode only the first member. + let mut compressed = Vec::new(); + Gzip::new() + .compress(&mut Bytes::from_static(b"first member"), &mut compressed) + .unwrap(); + let mut second_member = Vec::new(); + Gzip::new() + .compress( + &mut Bytes::from_static(b"second member"), + &mut second_member, + ) + .unwrap(); + compressed.extend_from_slice(&second_member); + + let mut decompressed = Vec::new(); + let result = Gzip::new().decompress(&mut compressed.as_slice(), &mut decompressed); + assert!(result.is_err()); + } +} diff --git a/grpc/src/codec/compression/registry.rs b/grpc/src/codec/compression/registry.rs new file mode 100644 index 000000000..cc3357737 --- /dev/null +++ b/grpc/src/codec/compression/registry.rs @@ -0,0 +1,253 @@ +/* + * + * Copyright 2026 gRPC authors. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to + * deal in the Software without restriction, including without limitation the + * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or + * sell copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING + * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS + * IN THE SOFTWARE. + * + */ + +use std::collections::HashMap; +use std::sync::Arc; +use std::sync::LazyLock; + +use crate::codec::compression::Compressor; +use crate::codec::compression::Decompressor; + +const IDENTITY_ENCODING: &str = "identity"; + +/// The immutable data backing a [`CompressionRegistry`]. +struct RegistryInner { + compressors: HashMap>, + decompressors: HashMap>, + accept_encodings: Arc<[String]>, +} + +/// Computes the sorted `grpc-accept-encoding` list from the registered +/// decompressors, always placing `identity` last. +fn compute_accept_encodings( + decompressors: &HashMap>, +) -> Arc<[String]> { + let mut encodings: Vec = decompressors + .keys() + .filter(|k| k.as_str() != IDENTITY_ENCODING) + .cloned() + .collect(); + encodings.sort_unstable(); + encodings.push(IDENTITY_ENCODING.to_owned()); + encodings.into() +} + +/// A builder for assembling a [`CompressionRegistry`]. +#[derive(Default)] +pub struct CompressionRegistryBuilder { + compressors: HashMap>, + decompressors: HashMap>, +} + +impl CompressionRegistryBuilder { + /// Creates a builder with no codecs registered. + pub fn new() -> Self { + Self::default() + } + + /// Creates a builder pre-populated with the built-in codecs (`gzip`). + pub fn with_defaults() -> Self { + let mut builder = Self::default(); + builder.register_builtin_defaults(); + builder + } + + /// Registers a compressor, keyed by its [`Compressor::name`]. + /// + /// An existing registration for the same encoding name is overwritten. + pub fn register_compressor(mut self, compressor: Arc) -> Self { + self.compressors + .insert(compressor.name().to_owned(), compressor); + self + } + + /// Registers a decompressor, keyed by its [`Decompressor::name`]. + /// + /// An existing registration for the same encoding name is overwritten. + pub fn register_decompressor(mut self, decompressor: Arc) -> Self { + self.decompressors + .insert(decompressor.name().to_owned(), decompressor); + self + } + + /// Builds a [`CompressionRegistry`] from the registered codecs. + pub fn build(self) -> CompressionRegistry { + let accept_encodings = compute_accept_encodings(&self.decompressors); + CompressionRegistry { + inner: Arc::new(RegistryInner { + compressors: self.compressors, + decompressors: self.decompressors, + accept_encodings, + }), + } + } + + /// Registers the built-in codecs that ship out of the box. + fn register_builtin_defaults(&mut self) { + // Register built-in compressors. + #[cfg(feature = "gzip")] + self.compressors.insert( + "gzip".to_owned(), + Arc::new(crate::codec::compression::gzip::Gzip::default()), + ); + + // Register built-in decompressors. + #[cfg(feature = "gzip")] + self.decompressors.insert( + "gzip".to_owned(), + Arc::new(crate::codec::compression::gzip::Gzip::default()), + ); + } +} + +/// A set of compression codecs, looked up by their encoding name. +/// +/// Build one with [`CompressionRegistryBuilder`], or use the default set of +/// built-in codecs via [`CompressionRegistry::global`]. Cheap to clone. +#[derive(Clone)] +pub struct CompressionRegistry { + inner: Arc, +} + +/// The default registry containing the built-in codecs. +static GLOBAL_COMPRESSION_REGISTRY: LazyLock = + LazyLock::new(|| CompressionRegistryBuilder::with_defaults().build()); + +impl CompressionRegistry { + /// Returns the default registry containing the built-in codecs. + pub fn global() -> Self { + GLOBAL_COMPRESSION_REGISTRY.clone() + } + + /// Returns the compressor registered for the given encoding name, or `None` + /// if none is registered. + pub fn get_compressor(&self, name: &str) -> Option> { + self.inner.compressors.get(name).cloned() + } + + /// Returns the decompressor registered for the given encoding name, or + /// `None` if none is registered. + pub fn get_decompressor(&self, name: &str) -> Option> { + self.inner.decompressors.get(name).cloned() + } + + /// Returns the encoding names to advertise in the `grpc-accept-encoding` + /// header. + pub fn accept_encodings(&self) -> &[String] { + &self.inner.accept_encodings + } +} + +#[cfg(test)] +mod tests { + use bytes::Buf; + use bytes::BufMut; + + use super::*; + + #[derive(Debug, Clone, Copy)] + struct MockCompression; + + impl Compressor for MockCompression { + fn name(&self) -> &str { + "mock" + } + + fn compress( + &self, + _source: &mut dyn Buf, + _destination: &mut dyn BufMut, + ) -> Result<(), String> { + Ok(()) + } + } + + impl Decompressor for MockCompression { + fn name(&self) -> &str { + "mock" + } + + fn decompress( + &self, + _source: &mut dyn Buf, + _destination: &mut dyn BufMut, + ) -> Result<(), String> { + Ok(()) + } + } + + #[test] + fn test_default_compressors_populated() { + let registry = CompressionRegistry::global(); + + // Verify gzip is present by default. + #[cfg(feature = "gzip")] + { + assert!(registry.get_compressor("gzip").is_some()); + assert!(registry.get_decompressor("gzip").is_some()); + } + } + + #[test] + fn accept_encoding() { + let registry = CompressionRegistry::global(); + let encodings = registry.accept_encodings(); + assert!(encodings.iter().any(|e| e == "identity")); + } + + #[test] + fn test_builder_registration_and_overwrite() { + // A freshly built registry does not contain the mock codec. + let registry = CompressionRegistryBuilder::new().build(); + assert!(registry.get_compressor("mock").is_none()); + + // Registering the mock makes it available. + let registry = CompressionRegistryBuilder::new() + .register_compressor(Arc::new(MockCompression)) + .build(); + assert!(registry.get_compressor("mock").is_some()); + + // Registering again with the same name overwrites correctly. + let registry = CompressionRegistryBuilder::new() + .register_compressor(Arc::new(MockCompression)) + .register_compressor(Arc::new(MockCompression)) + .build(); + assert!(registry.get_compressor("mock").is_some()); + } + + #[test] + fn test_accept_encodings_header_update() { + // A registry without the mock decompressor must not advertise it. + let registry = CompressionRegistryBuilder::new().build(); + assert!(!registry.accept_encodings().iter().any(|e| e == "mock")); + + // Registering a mock decompressor advertises it, and `identity` is + // always present. + let registry = CompressionRegistryBuilder::new() + .register_decompressor(Arc::new(MockCompression)) + .build(); + assert!(registry.accept_encodings().iter().any(|e| e == "mock")); + assert!(registry.accept_encodings().iter().any(|e| e == "identity")); + } +} diff --git a/grpc/src/codec/message.rs b/grpc/src/codec/message.rs new file mode 100644 index 000000000..a1e54dce0 --- /dev/null +++ b/grpc/src/codec/message.rs @@ -0,0 +1,131 @@ +use std::any::TypeId; +use std::collections::VecDeque; + +use bytes::Buf; +use bytes::Bytes; + +use crate::core::MessageType; +use crate::core::RecvMessage; +use crate::core::SendMessage; + +/// An immutable value-type struct representing an incoming raw gRPC message. +pub(crate) struct IncomingRawMessage { + buf: Box, + compressed: bool, +} + +impl IncomingRawMessage { + /// Constructs a new `IncomingRawMessage` initialized with a cheap empty buffer. + pub(crate) fn new() -> Self { + Self { + buf: Box::new(Bytes::new()), + compressed: false, + } + } + + /// Destructures the message by value into its raw payload buffer and compression flag. + pub(crate) fn into_parts(self) -> (Box, bool) { + (self.buf, self.compressed) + } + + /// Safely sets the per-message compression flag. + pub(crate) fn set_compressed(&mut self, compressed: bool) { + self.compressed = compressed; + } +} + +impl Default for IncomingRawMessage { + fn default() -> Self { + Self::new() + } +} + +impl MessageType for IncomingRawMessage { + type Target<'a> = IncomingRawMessage; +} + +impl RecvMessage for IncomingRawMessage { + fn decode(&mut self, data: &mut dyn Buf) -> Result<(), String> { + // Directly updates the immutable value-type container's inner buffer + self.buf = Box::new(data.copy_to_bytes(data.remaining())); + Ok(()) + } + + unsafe fn _ptr_for(&mut self, id: TypeId) -> Option<*mut ()> { + if id == TypeId::of::() { + Some(self as *mut IncomingRawMessage as *mut ()) + } else { + None + } + } +} + +/// A custom `Buf` implementation that streams sequentially through a deque of `Bytes` chunks. +struct ChunkedBuf { + chunks: VecDeque, +} + +impl Buf for ChunkedBuf { + fn remaining(&self) -> usize { + self.chunks.iter().map(|b| b.len()).sum() + } + + fn chunk(&self) -> &[u8] { + self.chunks.front().map(|b| b.chunk()).unwrap_or(&[]) + } + + fn advance(&mut self, mut cnt: usize) { + while cnt > 0 { + if let Some(front) = self.chunks.front_mut() { + let len = front.len(); + if cnt >= len { + cnt -= len; + self.chunks.pop_front(); + } else { + front.advance(cnt); + break; + } + } else { + break; + } + } + } +} + +/// A raw outgoing message usable for configuring SendOptions cleanly. +/// Stores data as a hybrid enum to allow zero-copy outbound serialization +/// without allocating a `VecDeque` for standard contiguous messages. +pub(crate) enum RawMessage { + Contiguous(Bytes), + Chunks(VecDeque), +} + +impl RawMessage { + pub(crate) fn from_buf(mut buf: impl Buf) -> Self { + let remaining = buf.remaining(); + if buf.chunk().len() == remaining { + RawMessage::Contiguous(buf.copy_to_bytes(remaining)) + } else { + let mut chunks = VecDeque::new(); + while buf.has_remaining() { + let chunk_len = buf.chunk().len(); + chunks.push_back(buf.copy_to_bytes(chunk_len)); + } + RawMessage::Chunks(chunks) + } + } +} + +impl SendMessage for RawMessage { + fn encode(&self) -> Result, String> { + match self { + RawMessage::Contiguous(bytes) => Ok(Box::new(bytes.clone())), + RawMessage::Chunks(chunks) => { + // `Bytes` clones are cheap $O(1)$ reference bumps, preserving idempotency safely. + Ok(Box::new(ChunkedBuf { + chunks: chunks.clone(), + })) + } + } + } +} diff --git a/grpc/src/lib.rs b/grpc/src/lib.rs index eeae2df3c..ba8642aad 100644 --- a/grpc/src/lib.rs +++ b/grpc/src/lib.rs @@ -51,6 +51,7 @@ pub mod attributes; pub mod client; +pub(crate) mod codec; pub mod core; pub mod credentials; pub mod metadata; diff --git a/grpc/src/server/interceptor.rs b/grpc/src/server/interceptor.rs index a65521e2e..adedc223a 100644 --- a/grpc/src/server/interceptor.rs +++ b/grpc/src/server/interceptor.rs @@ -29,6 +29,8 @@ use crate::server::Handle; use crate::server::RecvStream; use crate::server::SendStream; +pub mod compression; + /// A trait which allows intercepting an incoming RPC call to a [`Handle`] implementation. #[trait_variant::make(Send)] pub trait Intercept: Sync + 'static { diff --git a/grpc/src/server/interceptor/compression.rs b/grpc/src/server/interceptor/compression.rs new file mode 100644 index 000000000..c1d98f532 --- /dev/null +++ b/grpc/src/server/interceptor/compression.rs @@ -0,0 +1,1514 @@ +use std::sync::Arc; + +use bytes::BufMut; +use bytes::BytesMut; + +use crate::StatusCodeError; +use crate::StatusError; +use crate::client::CallOptions; +use crate::codec::compression::Compressor; +use crate::codec::compression::Decompressor; +use crate::codec::compression::registry::CompressionRegistry; +use crate::codec::message::IncomingRawMessage; +use crate::codec::message::RawMessage; +use crate::core::RecvMessage; +use crate::core::RequestHeaders; +use crate::core::ResponseHeaders; +use crate::core::Trailers; +use crate::metadata::MetadataMap; +use crate::server::Handle; +use crate::server::RecvStream; +use crate::server::ResponseStreamItem; +use crate::server::SendOptions; +use crate::server::SendStream; +use crate::server::interceptor::Intercept; + +const DEFAULT_DECOMPRESSION_LIMIT: usize = 4 * 1024 * 1024; +const INITIAL_COMPRESSION_BUFFER_CAPACITY: usize = 8192; + +const GRPC_ENCODING_HEADER: &str = "grpc-encoding"; +const GRPC_ACCEPT_ENCODING_HEADER: &str = "grpc-accept-encoding"; +const IDENTITY_ENCODING: &str = "identity"; + +/// A gRPC server interceptor that manages automatic payload compression and +/// decompression based on client headers and server registry capabilities. +/// +/// # Examples +/// +/// ```rust,ignore +/// use std::sync::Arc; +/// use tonic_server_grpc::codec::compression::global_codec_registry; +/// use tonic_server_grpc::server::interceptor::compression::ServerCompressionInterceptor; +/// +/// let resolver = Arc::new(global_codec_registry()); +/// let interceptor = ServerCompressionInterceptor::new(resolver) +/// .with_decompression_limit(8 * 1024 * 1024); +/// ``` +#[derive(Clone)] +pub struct ServerCompressionInterceptor { + registry: CompressionRegistry, + decompression_limit: usize, + default_send_compressor: Option, +} + +impl ServerCompressionInterceptor { + /// Creates a new compression interceptor using the provided registry. + pub fn new(registry: CompressionRegistry) -> Self { + Self { + registry, + decompression_limit: DEFAULT_DECOMPRESSION_LIMIT, + default_send_compressor: None, + } + } + + /// Configures a custom byte ceiling for decompression bomb mitigation. + pub fn with_decompression_limit(mut self, limit: usize) -> Self { + self.decompression_limit = limit; + self + } + + /// Sets a global default compressor to use for responses if the application handler does not specify one. + pub fn with_default_send_compressor(mut self, encoding: &str) -> Self { + self.default_send_compressor = Some(encoding.to_string()); + self + } +} + +impl Default for ServerCompressionInterceptor { + fn default() -> Self { + Self::new(CompressionRegistry::global()) + } +} + +impl Intercept for ServerCompressionInterceptor { + async fn intercept( + &self, + headers: RequestHeaders, + options: CallOptions, + tx: &mut impl SendStream, + rx: impl RecvStream + 'static, + next: &impl Handle, + ) -> Trailers { + let decompressor = match resolve_decompressor(&self.registry, headers.metadata()) { + Ok(d) => d, + Err(err) => { + let mut trailers = Trailers::new(Err(err.status)); + if let Some(accept_str) = &err.accept_encodings + && let Ok(val) = accept_str.parse() + { + trailers + .metadata_mut() + .insert(GRPC_ACCEPT_ENCODING_HEADER, val); + } + return trailers; + } + }; + + let accepted_encodings = headers + .metadata() + .get_all(GRPC_ACCEPT_ENCODING_HEADER) + .iter() + .map(|v| v.to_str()) + .flat_map(|v| v.split(',')) + .map(str::trim) + .map(String::from) + .collect::>(); + + let request_encoding = headers + .metadata() + .get(GRPC_ENCODING_HEADER) + .map(|v| v.to_str()) + .map(String::from); + + let fallback_encoding = self.default_send_compressor.clone().or(request_encoding); + + let pending = PendingNegotiation { + registry: self.registry.clone(), + accepted_encodings, + fallback_encoding, + }; + + let mut wrapped_tx = CompressedSendStream::new(tx, pending); + + let active_rx = decompressor.map(|codec| ActiveDecompressor { + codec, + buf: BytesMut::with_capacity(INITIAL_COMPRESSION_BUFFER_CAPACITY), + }); + + let wrapped_rx = DecompressedRecvStream { + inner: rx, + decompression_limit: self.decompression_limit, + active: active_rx, + }; + + next.handle(headers, options, &mut wrapped_tx, wrapped_rx) + .await + } +} + +/// State wrapper for an active stream compressor. +/// +/// Holds the compressor implementation and a buffer used to incrementally +/// compress outbound gRPC messages. The buffer is retained to avoid reallocation +/// between messages. +struct ActiveCompressor { + codec: Arc, + buf: BytesMut, +} + +impl ActiveCompressor { + fn new(codec: Arc) -> Self { + Self { + codec, + buf: BytesMut::with_capacity(INITIAL_COMPRESSION_BUFFER_CAPACITY), + } + } +} + +struct PendingNegotiation { + registry: CompressionRegistry, + accepted_encodings: Vec, + fallback_encoding: Option, +} + +impl PendingNegotiation { + /// Resolves the final state for the compressor based on outbound headers. + fn resolve(&self, headers: &mut ResponseHeaders) -> Result { + let Some((enc, should_inject_encoding_header)) = self.negotiate_encoding(headers) else { + return Ok(SendCompressorState::Disabled); + }; + + match self.registry.get_compressor(&enc) { + Some(codec) => { + if should_inject_encoding_header && let Ok(val) = codec.name().parse() { + headers.metadata_mut().insert(GRPC_ENCODING_HEADER, val); + } + Ok(SendCompressorState::Active(ActiveCompressor::new(codec))) + } + None => Ok(SendCompressorState::Disabled), + } + } + + /// Determines which encoding to use and whether it needs to be injected. + fn negotiate_encoding(&self, headers: &mut ResponseHeaders) -> Option<(String, bool)> { + // 1. Check if handler provided a valid override + if let Some(enc) = Self::get_handler_encoding(headers) { + if self.accepted_encodings.contains(&enc) { + return Some((enc, false)); + } + // Lenient conflict resolution: strip the invalid header + // and send uncompressed response. + headers.metadata_mut().remove(GRPC_ENCODING_HEADER); + return None; + } + + // 2. Check fallback (global default or symmetric) + if let Some(enc) = &self.fallback_encoding + && self.accepted_encodings.contains(enc) + { + return Some((enc.clone(), true)); + } + None + } + + fn get_handler_encoding(headers: &ResponseHeaders) -> Option { + headers + .metadata() + .get(GRPC_ENCODING_HEADER) + .map(|v| v.to_str()) + .filter(|&enc| enc != IDENTITY_ENCODING) + .map(String::from) + } +} + +enum SendCompressorState { + Pending(PendingNegotiation), + Active(ActiveCompressor), + Disabled, +} + +/// Transparent stream adapter that intercepts outbound messages and applies compression. +/// +/// If a compressor is negotiated and active, each message is compressed before being sent +/// to the underlying transport. +struct CompressedSendStream<'a, S: SendStream> { + inner: &'a mut S, + state: SendCompressorState, +} + +impl<'a, S: SendStream> CompressedSendStream<'a, S> { + fn new(inner: &'a mut S, pending: PendingNegotiation) -> Self { + Self { + inner, + state: SendCompressorState::Pending(pending), + } + } +} + +impl<'a, S: SendStream> SendStream for CompressedSendStream<'a, S> { + async fn send<'b>( + &mut self, + item: ResponseStreamItem<'b>, + options: SendOptions, + ) -> Result<(), ()> { + match item { + ResponseStreamItem::Headers(mut headers) => { + if let SendCompressorState::Pending(pending) = &self.state { + self.state = pending.resolve(&mut headers)?; + } else { + // gRPC strictly allows Initial Metadata (Headers) to be sent only once. + // If the state is no longer Pending, Headers were already processed. + return Err(()); + } + + self.inner + .send(ResponseStreamItem::Headers(headers), options) + .await + } + ResponseStreamItem::Message(msg) => { + let active = match &mut self.state { + SendCompressorState::Active(active) => active, + SendCompressorState::Disabled => { + let mut options = options; + options.disable_compression = true; + return self + .inner + .send(ResponseStreamItem::Message(msg), options) + .await; + } + SendCompressorState::Pending(_) => { + // gRPC strictly requires headers to precede messages. + // If the handler attempts to send a message before headers, abort the stream. + return Err(()); + } + }; + + if options.disable_compression { + // disable_compression is already true — message is uncompressed. + return self + .inner + .send(ResponseStreamItem::Message(msg), options) + .await; + } + + let mut buf = msg.encode().map_err(|_| ())?; + // TODO: Implement capacity shrinking to avoid memory leaks on long-lived streams. + // If capacity is excessive (e.g. > 8MB), replace with a new `BytesMut` instead. + active.buf.clear(); + active + .codec + .compress(&mut *buf, &mut active.buf) + .map_err(|_| ())?; + + let raw_msg = RawMessage::from_buf(active.buf.split().freeze()); + // Signal downstream that this message is compressed. + // TODO(sauravz): disable_compression defaults to false, so messages + // that bypass the compression interceptor entirely will also appear + // as "compressed" to the framing layer. Consider adding an explicit + // is_compressed field to SendOptions once a breaking change is feasible. + let mut options = options; + options.disable_compression = false; + self.inner + .send(ResponseStreamItem::Message(&raw_msg), options) + .await + } + } + } +} + +/// State wrapper for an active stream decompressor. +/// +/// Holds the decompressor implementation and a buffer used to incrementally +/// decompress inbound gRPC messages. The buffer is retained to avoid reallocation +/// between messages. +struct ActiveDecompressor { + codec: Arc, + buf: BytesMut, +} + +/// Transparent stream adapter that intercepts inbound messages and applies decompression. +/// +/// If an `ActiveDecompressor` is present, each message is decompressed before being yielded +/// to the application handler. +struct DecompressedRecvStream { + inner: R, + decompression_limit: usize, + active: Option, +} + +impl RecvStream for DecompressedRecvStream { + /// Fetches the next incoming gRPC message, destructuring the raw buffer directly. + async fn next(&mut self, msg: &mut dyn RecvMessage) -> Option> { + if let Some(active) = &mut self.active { + let mut raw_msg = IncomingRawMessage::new(); + let res = self.inner.next(&mut raw_msg).await?; + if res.is_err() { + return Some(Err(())); + } + + let (mut source_buf, is_compressed) = raw_msg.into_parts(); + + if is_compressed { + // TODO: Implement capacity shrinking to avoid memory leaks on long-lived streams. + // If capacity is excessive (e.g. > 8MB), replace with a new `BytesMut` instead. + active.buf.clear(); + let mut limited_dst = (&mut active.buf).limit(self.decompression_limit); + if active + .codec + .decompress(&mut *source_buf, &mut limited_dst) + .is_err() + { + return Some(Err(())); + } + let mut payload = active.buf.split().freeze(); + if msg.decode(&mut payload).is_err() { + return Some(Err(())); + } + } else if msg.decode(&mut *source_buf).is_err() { + return Some(Err(())); + } + Some(Ok(())) + } else { + self.inner.next(msg).await + } + } +} + +/// Bundled error context returned by pure codec resolvers, holding both the gRPC +/// status and any optional pushback trailer context (like supported encodings). +#[derive(Debug)] +struct ResolverError { + /// The primary gRPC status error (e.g., Unimplemented or Internal). + status: StatusError, + /// An optional comma-separated string of supported encodings, to be attached + /// as the `grpc-accept-encoding` trailing header upon Unimplemented errors. + accept_encodings: Option, +} + +/// Inspects incoming headers to determine if the client encoded the request. +/// +/// If `grpc-encoding` is present and is not `identity`, this function queries the registry +/// for an appropriate decompressor. If the encoding is unsupported, an error is returned. +fn resolve_decompressor( + registry: &CompressionRegistry, + metadata: &MetadataMap, +) -> Result>, ResolverError> { + let recv_encoding = metadata + .get(GRPC_ENCODING_HEADER) + .map(|v| v.to_str()) + .filter(|&enc| enc != IDENTITY_ENCODING); + + if let Some(encoding) = recv_encoding { + match registry.get_decompressor(encoding) { + Some(decompressor) => Ok(Some(decompressor)), + None => { + let status = StatusError::new( + StatusCodeError::Unimplemented, + format!("Compression encoding {} not supported", encoding), + ); + let accept_encodings = Some(registry.accept_encodings().join(",")); + Err(ResolverError { + status, + accept_encodings, + }) + } + } + } else { + Ok(None) + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use bytes::Buf; + use bytes::Bytes; + use tokio::sync::Mutex; + + use super::*; + use crate::codec::compression::registry::CompressionRegistryBuilder; + use crate::core::ResponseHeaders; + use crate::core::SendMessage; + use crate::core::Trailers; + use crate::server::interceptor::HandleExt; + + /// A fake compressor/decompressor used for testing. + /// It "compresses" by prepending `[compressed]` to the payload, + /// and "decompresses" by verifying and stripping that prefix. + #[derive(Debug, Clone, Copy)] + struct MockCodec; + + impl Compressor for MockCodec { + fn name(&self) -> &str { + "mock" + } + fn compress( + &self, + src: &mut dyn Buf, + dst: &mut dyn bytes::buf::BufMut, + ) -> Result<(), String> { + dst.put_slice(b"[compressed]"); + dst.put_slice(&src.copy_to_bytes(src.remaining())); + Ok(()) + } + } + + impl Decompressor for MockCodec { + fn name(&self) -> &str { + "mock" + } + fn decompress( + &self, + src: &mut dyn Buf, + dst: &mut dyn bytes::buf::BufMut, + ) -> Result<(), String> { + let bytes = src.copy_to_bytes(src.remaining()); + if bytes.starts_with(b"[compressed]") { + let payload = &bytes[12..]; + if dst.remaining_mut() < payload.len() { + return Err("limit reached".to_string()); + } + dst.put_slice(payload); + Ok(()) + } else { + Err("not compressed".to_string()) + } + } + } + + /// Builds a compression registry that only supports the "mock" codec. + fn mock_registry() -> CompressionRegistry { + CompressionRegistryBuilder::new() + .register_compressor(Arc::new(MockCodec)) + .register_decompressor(Arc::new(MockCodec)) + .build() + } + + /// A fake network send stream that intercepts and stores outgoing messages + /// and headers so that tests can assert what was sent back to the client. + struct MockSendStream { + messages: Arc>>, + headers: Arc>>, + /// Captures the `disable_compression` value for each message sent. + msg_disable_compression: Arc>>, + } + impl SendStream for MockSendStream { + async fn send<'a>( + &mut self, + item: ResponseStreamItem<'a>, + opts: SendOptions, + ) -> Result<(), ()> { + match item { + ResponseStreamItem::Headers(h) => { + *self.headers.lock().await = Some(h); + } + ResponseStreamItem::Message(msg) => { + let mut buf = msg.encode().unwrap(); + self.messages + .lock() + .await + .push(buf.copy_to_bytes(buf.remaining())); + self.msg_disable_compression + .lock() + .await + .push(opts.disable_compression); + } + } + Ok(()) + } + } + + /// A fake network receive stream that yields hardcoded byte arrays + /// to simulate incoming client messages. + struct MockRecvStream { + items: Vec>, + } + impl RecvStream for MockRecvStream { + async fn next(&mut self, msg: &mut dyn RecvMessage) -> Option> { + if self.items.is_empty() { + return None; + } + let item = self.items.remove(0); + match item { + Ok(bytes) => { + let is_compressed = + bytes.starts_with(b"[compressed]") || bytes.as_ref() == b"bad payload"; + if let Some(raw_msg) = msg.downcast_mut::() { + raw_msg.set_compressed(is_compressed); + } + let mut buf = bytes; + if msg.decode(&mut buf).is_err() { + Some(Err(())) + } else { + Some(Ok(())) + } + } + Err(()) => Some(Err(())), + } + } + } + + /// A simple mock gRPC service handler that echoes back the string "echo". + struct MockHandler; + impl Handle for MockHandler { + async fn handle( + &self, + _headers: RequestHeaders, + _options: CallOptions, + tx: &mut impl SendStream, + mut rx: impl RecvStream + 'static, + ) -> Trailers { + let _ = tx + .send( + ResponseStreamItem::Headers(ResponseHeaders::new()), + SendOptions::default(), + ) + .await; + struct StringMsg(String); + impl RecvMessage for StringMsg { + fn decode(&mut self, data: &mut dyn Buf) -> Result<(), String> { + let b = data.copy_to_bytes(data.remaining()); + self.0 = String::from_utf8(b.to_vec()).unwrap(); + Ok(()) + } + } + impl SendMessage for StringMsg { + fn encode(&self) -> Result, String> { + Ok(Box::new(Bytes::from(self.0.clone()))) + } + } + + while let Some(Ok(())) = rx.next(&mut StringMsg(String::new())).await { + let _ = tx + .send( + ResponseStreamItem::Message(&StringMsg("echo".into())), + SendOptions::default(), + ) + .await; + } + Trailers::new(Ok(())) + } + } + + #[tokio::test] + async fn test_unknown_incoming_encoding() { + let registry = mock_registry(); + let mut tx = MockSendStream { + messages: Arc::new(Mutex::new(Vec::new())), + headers: Arc::new(Mutex::new(None)), + msg_disable_compression: Arc::new(Mutex::new(Vec::new())), + }; + let interceptor = ServerCompressionInterceptor::new(registry); + let chain = MockHandler.with_interceptor(interceptor); + + let mut headers = RequestHeaders::new(); + headers + .metadata_mut() + .insert(GRPC_ENCODING_HEADER, "unknown".parse().unwrap()); + + let trailers = chain + .handle( + headers, + CallOptions::default(), + &mut tx, + MockRecvStream { items: vec![] }, + ) + .await; + assert_eq!( + trailers.status().as_ref().unwrap_err().code(), + StatusCodeError::Unimplemented + ); + assert_eq!( + trailers + .metadata() + .get(GRPC_ACCEPT_ENCODING_HEADER) + .expect("Expected grpc-accept-encoding trailer to be present") + .to_str(), + "mock,identity" + ); + } + + #[tokio::test] + async fn test_identity_incoming_encoding() { + let registry = mock_registry(); + let messages = Arc::new(Mutex::new(Vec::new())); + let mut tx = MockSendStream { + messages: messages.clone(), + headers: Arc::new(Mutex::new(None)), + msg_disable_compression: Arc::new(Mutex::new(Vec::new())), + }; + let interceptor = ServerCompressionInterceptor::new(registry); + let chain = MockHandler.with_interceptor(interceptor); + + let mut headers = RequestHeaders::new(); + headers + .metadata_mut() + .insert("grpc-encoding", "identity".parse().unwrap()); + let rx = MockRecvStream { + items: vec![Ok(Bytes::from_static(b"hello"))], + }; + + let trailers = chain + .handle(headers, CallOptions::default(), &mut tx, rx) + .await; + assert!(trailers.status().is_ok()); + assert_eq!(messages.lock().await.len(), 1); + } + + #[tokio::test] + async fn test_supported_incoming_encoding() { + let registry = mock_registry(); + let messages = Arc::new(Mutex::new(Vec::new())); + let mut tx = MockSendStream { + messages: messages.clone(), + headers: Arc::new(Mutex::new(None)), + msg_disable_compression: Arc::new(Mutex::new(Vec::new())), + }; + let interceptor = ServerCompressionInterceptor::new(registry); + let chain = MockHandler.with_interceptor(interceptor); + + let mut headers = RequestHeaders::new(); + headers + .metadata_mut() + .insert("grpc-encoding", "mock".parse().unwrap()); + let rx = MockRecvStream { + items: vec![Ok(Bytes::from_static(b"[compressed]hello"))], + }; + + let trailers = chain + .handle(headers, CallOptions::default(), &mut tx, rx) + .await; + assert!(trailers.status().is_ok()); + assert_eq!(messages.lock().await.len(), 1); + } + + #[tokio::test] + async fn test_decompression_failure() { + struct FailingDecompHandler; + impl Handle for FailingDecompHandler { + async fn handle( + &self, + _h: RequestHeaders, + _o: CallOptions, + _tx: &mut impl SendStream, + mut rx: impl RecvStream + 'static, + ) -> Trailers { + struct StringMsg; + impl RecvMessage for StringMsg { + fn decode(&mut self, _data: &mut dyn Buf) -> Result<(), String> { + Ok(()) + } + } + let res = rx.next(&mut StringMsg).await; + assert!(matches!(res, Some(Err(())))); + Trailers::new(Err(StatusError::new( + crate::status::StatusCodeError::Internal, + "decompression failed", + ))) + } + } + + let registry = mock_registry(); + let mut tx = MockSendStream { + messages: Arc::new(Mutex::new(Vec::new())), + headers: Arc::new(Mutex::new(None)), + msg_disable_compression: Arc::new(Mutex::new(Vec::new())), + }; + let interceptor = ServerCompressionInterceptor::new(registry); + let chain = FailingDecompHandler.with_interceptor(interceptor); + + let mut headers = RequestHeaders::new(); + headers + .metadata_mut() + .insert("grpc-encoding", "mock".parse().unwrap()); + let rx = MockRecvStream { + items: vec![Ok(Bytes::from_static(b"bad payload"))], + }; + + let trailers = chain + .handle(headers, CallOptions::default(), &mut tx, rx) + .await; + assert_eq!( + trailers.status().as_ref().unwrap_err().code(), + crate::status::StatusCodeError::Internal + ); + } + + #[tokio::test] + async fn test_supported_outgoing_encoding() { + let registry = mock_registry(); + let messages = Arc::new(Mutex::new(Vec::new())); + let resp_headers = Arc::new(Mutex::new(None)); + let mut tx = MockSendStream { + messages: messages.clone(), + headers: resp_headers.clone(), + msg_disable_compression: Arc::new(Mutex::new(Vec::new())), + }; + let interceptor = ServerCompressionInterceptor::new(registry); + let chain = MockHandler.with_interceptor(interceptor); + + let mut headers = RequestHeaders::new(); + headers + .metadata_mut() + .insert("grpc-encoding", "mock".parse().unwrap()); + headers + .metadata_mut() + .insert("grpc-accept-encoding", "mock".parse().unwrap()); + let rx = MockRecvStream { + items: vec![Ok(Bytes::from_static(b"[compressed]hello"))], + }; + + let trailers = chain + .handle(headers, CallOptions::default(), &mut tx, rx) + .await; + assert!(trailers.status().is_ok()); + + let h = resp_headers.lock().await.take().unwrap(); + assert_eq!(h.metadata().get("grpc-encoding").unwrap().to_str(), "mock"); + + let msgs = messages.lock().await; + assert!(msgs[0].starts_with(b"[compressed]")); + } + + #[tokio::test] + async fn test_disable_compression_option() { + struct DisableHandler; + impl Handle for DisableHandler { + async fn handle( + &self, + _h: RequestHeaders, + _o: CallOptions, + tx: &mut impl SendStream, + _rx: impl RecvStream + 'static, + ) -> Trailers { + let _ = tx + .send( + ResponseStreamItem::Headers(ResponseHeaders::new()), + SendOptions::default(), + ) + .await; + + struct StringMsg; + impl SendMessage for StringMsg { + fn encode(&self) -> Result, String> { + Ok(Box::new(Bytes::from_static(b"echo"))) + } + } + let opts = SendOptions { + disable_compression: true, + ..Default::default() + }; + let _ = tx.send(ResponseStreamItem::Message(&StringMsg), opts).await; + Trailers::new(Ok(())) + } + } + + let registry = mock_registry(); + let messages = Arc::new(Mutex::new(Vec::new())); + let mut tx = MockSendStream { + messages: messages.clone(), + headers: Arc::new(Mutex::new(None)), + msg_disable_compression: Arc::new(Mutex::new(Vec::new())), + }; + let interceptor = ServerCompressionInterceptor::new(registry); + let chain = DisableHandler.with_interceptor(interceptor); + + let mut headers = RequestHeaders::new(); + headers + .metadata_mut() + .insert("grpc-encoding", "mock".parse().unwrap()); + headers + .metadata_mut() + .insert("grpc-accept-encoding", "mock".parse().unwrap()); + + let _ = chain + .handle( + headers, + CallOptions::default(), + &mut tx, + MockRecvStream { items: vec![] }, + ) + .await; + + let msgs = messages.lock().await; + assert_eq!(msgs[0], Bytes::from_static(b"echo")); + } + + /// Verifies that when compression is active and applied, the compression + /// interceptor sets `disable_compression = false` to signal downstream + /// that the message payload is compressed. + #[tokio::test] + async fn test_compressed_message_signals_disable_compression_false() { + let registry = mock_registry(); + let messages = Arc::new(Mutex::new(Vec::new())); + let msg_disable_compression = Arc::new(Mutex::new(Vec::new())); + let mut tx = MockSendStream { + messages: messages.clone(), + headers: Arc::new(Mutex::new(None)), + msg_disable_compression: msg_disable_compression.clone(), + }; + + let interceptor = ServerCompressionInterceptor::new(registry); + let chain = MockHandler.with_interceptor(interceptor); + + let mut headers = RequestHeaders::new(); + headers + .metadata_mut() + .insert("grpc-encoding", "mock".parse().unwrap()); + headers + .metadata_mut() + .insert("grpc-accept-encoding", "mock".parse().unwrap()); + let rx = MockRecvStream { + items: vec![Ok(Bytes::from_static(b"[compressed]hello"))], + }; + + let trailers = chain + .handle(headers, CallOptions::default(), &mut tx, rx) + .await; + assert!(trailers.status().is_ok()); + + let msgs = messages.lock().await; + assert_eq!(msgs.len(), 1); + assert!(msgs[0].starts_with(b"[compressed]")); + + let flags = msg_disable_compression.lock().await; + assert_eq!(flags.len(), 1); + assert!( + !flags[0], + "disable_compression should be false for a compressed message" + ); + } + + /// Verifies that when compression is disabled (no grpc-accept-encoding), + /// the compression interceptor sets `disable_compression = true` to signal + /// downstream that the message payload is NOT compressed. + #[tokio::test] + async fn test_disabled_state_signals_disable_compression_true() { + let registry = mock_registry(); + let messages = Arc::new(Mutex::new(Vec::new())); + let msg_disable_compression = Arc::new(Mutex::new(Vec::new())); + let mut tx = MockSendStream { + messages: messages.clone(), + headers: Arc::new(Mutex::new(None)), + msg_disable_compression: msg_disable_compression.clone(), + }; + + let interceptor = ServerCompressionInterceptor::new(registry); + let chain = MockHandler.with_interceptor(interceptor); + + // No grpc-accept-encoding → compression disabled + let headers = RequestHeaders::new(); + let rx = MockRecvStream { + items: vec![Ok(Bytes::from_static(b"hello"))], + }; + + let trailers = chain + .handle(headers, CallOptions::default(), &mut tx, rx) + .await; + assert!(trailers.status().is_ok()); + + let msgs = messages.lock().await; + assert_eq!(msgs.len(), 1); + assert_eq!(msgs[0], Bytes::from_static(b"echo")); + + let flags = msg_disable_compression.lock().await; + assert_eq!(flags.len(), 1); + assert!( + flags[0], + "disable_compression should be true when compression is disabled" + ); + } + + /// Verifies that when compression is negotiated but disabled per-message + /// via SendOptions, the compression interceptor preserves + /// `disable_compression = true` to signal that this particular message + /// is NOT compressed. + #[tokio::test] + async fn test_per_message_disable_signals_disable_compression_true() { + struct PerMsgDisableHandler; + impl Handle for PerMsgDisableHandler { + async fn handle( + &self, + _h: RequestHeaders, + _o: CallOptions, + tx: &mut impl SendStream, + _rx: impl RecvStream + 'static, + ) -> Trailers { + let _ = tx + .send( + ResponseStreamItem::Headers(ResponseHeaders::new()), + SendOptions::default(), + ) + .await; + + struct EchoMsg; + impl SendMessage for EchoMsg { + fn encode(&self) -> Result, String> { + Ok(Box::new(Bytes::from_static(b"echo"))) + } + } + let opts = SendOptions { + disable_compression: true, + ..Default::default() + }; + let _ = tx.send(ResponseStreamItem::Message(&EchoMsg), opts).await; + Trailers::new(Ok(())) + } + } + + let registry = mock_registry(); + let messages = Arc::new(Mutex::new(Vec::new())); + let msg_disable_compression = Arc::new(Mutex::new(Vec::new())); + let mut tx = MockSendStream { + messages: messages.clone(), + headers: Arc::new(Mutex::new(None)), + msg_disable_compression: msg_disable_compression.clone(), + }; + + let interceptor = ServerCompressionInterceptor::new(registry); + let chain = PerMsgDisableHandler.with_interceptor(interceptor); + + let mut headers = RequestHeaders::new(); + headers + .metadata_mut() + .insert("grpc-encoding", "mock".parse().unwrap()); + headers + .metadata_mut() + .insert("grpc-accept-encoding", "mock".parse().unwrap()); + + let trailers = chain + .handle( + headers, + CallOptions::default(), + &mut tx, + MockRecvStream { items: vec![] }, + ) + .await; + assert!(trailers.status().is_ok()); + + let msgs = messages.lock().await; + assert_eq!(msgs.len(), 1); + assert_eq!(msgs[0], Bytes::from_static(b"echo")); + + let flags = msg_disable_compression.lock().await; + assert_eq!(flags.len(), 1); + assert!( + flags[0], + "disable_compression should be true when per-message compression is disabled" + ); + } + + #[tokio::test] + async fn test_multi_value_accept_encoding() { + let registry = mock_registry(); + let messages = Arc::new(Mutex::new(Vec::new())); + let resp_headers = Arc::new(Mutex::new(None)); + let mut tx = MockSendStream { + messages: messages.clone(), + headers: resp_headers.clone(), + msg_disable_compression: Arc::new(Mutex::new(Vec::new())), + }; + let interceptor = ServerCompressionInterceptor::new(registry); + let chain = MockHandler.with_interceptor(interceptor); + + let mut headers = RequestHeaders::new(); + headers + .metadata_mut() + .insert("grpc-encoding", "mock".parse().unwrap()); + headers.metadata_mut().insert( + "grpc-accept-encoding", + "gzip, mock, identity".parse().unwrap(), + ); + let rx = MockRecvStream { + items: vec![Ok(Bytes::from_static(b"[compressed]hello"))], + }; + + let trailers = chain + .handle(headers, CallOptions::default(), &mut tx, rx) + .await; + assert!(trailers.status().is_ok()); + + let h = resp_headers.lock().await.take().unwrap(); + assert_eq!(h.metadata().get("grpc-encoding").unwrap().to_str(), "mock"); + + let msgs = messages.lock().await; + assert!(msgs[0].starts_with(b"[compressed]")); + } + + #[tokio::test] + async fn test_missing_accept_encoding() { + let registry = mock_registry(); + let messages = Arc::new(Mutex::new(Vec::new())); + let resp_headers = Arc::new(Mutex::new(None)); + let mut tx = MockSendStream { + messages: messages.clone(), + headers: resp_headers.clone(), + msg_disable_compression: Arc::new(Mutex::new(Vec::new())), + }; + let interceptor = ServerCompressionInterceptor::new(registry); + let chain = MockHandler.with_interceptor(interceptor); + + let headers = RequestHeaders::new(); // No accept-encoding header + let rx = MockRecvStream { + items: vec![Ok(Bytes::from_static(b"hello"))], + }; + + let trailers = chain + .handle(headers, CallOptions::default(), &mut tx, rx) + .await; + assert!(trailers.status().is_ok()); + + let h = resp_headers.lock().await.take().unwrap(); + assert!(h.metadata().get("grpc-encoding").is_none()); + + let msgs = messages.lock().await; + assert_eq!(msgs[0], Bytes::from_static(b"echo")); + } + + #[tokio::test] + async fn test_asymmetric_compression_global_default() { + let registry = mock_registry(); + let messages = Arc::new(Mutex::new(Vec::new())); + let resp_headers = Arc::new(Mutex::new(None)); + let mut tx = MockSendStream { + messages: messages.clone(), + headers: resp_headers.clone(), + msg_disable_compression: Arc::new(Mutex::new(Vec::new())), + }; + let interceptor = + ServerCompressionInterceptor::new(registry).with_default_send_compressor("mock"); + let chain = MockHandler.with_interceptor(interceptor); + + let mut headers = RequestHeaders::new(); + // Request is uncompressed (no grpc-encoding), but client accepts "mock" + headers + .metadata_mut() + .insert("grpc-accept-encoding", "mock".parse().unwrap()); + + let rx = MockRecvStream { + items: vec![Ok(Bytes::from_static(b"hello"))], + }; + + let trailers = chain + .handle(headers, CallOptions::default(), &mut tx, rx) + .await; + assert!(trailers.status().is_ok()); + + let h = resp_headers.lock().await.take().unwrap(); + assert_eq!(h.metadata().get("grpc-encoding").unwrap().to_str(), "mock"); + + let msgs = messages.lock().await; + assert!(msgs[0].starts_with(b"[compressed]")); + } + + #[tokio::test] + async fn test_asymmetric_compression_handler_override() { + struct OverrideHandler; + impl Handle for OverrideHandler { + async fn handle( + &self, + _h: RequestHeaders, + _o: CallOptions, + tx: &mut impl SendStream, + mut rx: impl RecvStream + 'static, + ) -> Trailers { + let mut headers = ResponseHeaders::new(); + headers + .metadata_mut() + .insert("grpc-encoding", "mock".parse().unwrap()); + let _ = tx + .send(ResponseStreamItem::Headers(headers), SendOptions::default()) + .await; + + struct StringMsg; + impl SendMessage for StringMsg { + fn encode(&self) -> Result, String> { + Ok(Box::new(Bytes::from_static(b"echo"))) + } + } + impl RecvMessage for StringMsg { + fn decode(&mut self, _data: &mut dyn Buf) -> Result<(), String> { + Ok(()) + } + } + let _ = rx.next(&mut StringMsg).await; + let _ = tx + .send( + ResponseStreamItem::Message(&StringMsg), + SendOptions::default(), + ) + .await; + Trailers::new(Ok(())) + } + } + + let registry = mock_registry(); + let messages = Arc::new(Mutex::new(Vec::new())); + let resp_headers = Arc::new(Mutex::new(None)); + let mut tx = MockSendStream { + messages: messages.clone(), + headers: resp_headers.clone(), + msg_disable_compression: Arc::new(Mutex::new(Vec::new())), + }; + let interceptor = ServerCompressionInterceptor::new(registry); + let chain = OverrideHandler.with_interceptor(interceptor); + + let mut headers = RequestHeaders::new(); + // Request is uncompressed, but client accepts "mock" + headers + .metadata_mut() + .insert("grpc-accept-encoding", "mock".parse().unwrap()); + let rx = MockRecvStream { + items: vec![Ok(Bytes::from_static(b"hello"))], + }; + + let trailers = chain + .handle(headers, CallOptions::default(), &mut tx, rx) + .await; + assert!(trailers.status().is_ok()); + + let h = resp_headers.lock().await.take().unwrap(); + assert_eq!(h.metadata().get("grpc-encoding").unwrap().to_str(), "mock"); + + let msgs = messages.lock().await; + assert!(msgs[0].starts_with(b"[compressed]")); + } + + #[tokio::test] + async fn test_asymmetric_compression_invalid_handler_override() { + struct OverrideHandler; + impl Handle for OverrideHandler { + async fn handle( + &self, + _h: RequestHeaders, + _o: CallOptions, + tx: &mut impl SendStream, + mut rx: impl RecvStream + 'static, + ) -> Trailers { + let mut headers = ResponseHeaders::new(); + // Handler tries to force "mock", but client won't accept it. + headers + .metadata_mut() + .insert("grpc-encoding", "mock".parse().unwrap()); + let _ = tx + .send(ResponseStreamItem::Headers(headers), SendOptions::default()) + .await; + + struct StringMsg; + impl SendMessage for StringMsg { + fn encode(&self) -> Result, String> { + Ok(Box::new(Bytes::from_static(b"echo"))) + } + } + impl RecvMessage for StringMsg { + fn decode(&mut self, _data: &mut dyn Buf) -> Result<(), String> { + Ok(()) + } + } + let _ = rx.next(&mut StringMsg).await; + let _ = tx + .send( + ResponseStreamItem::Message(&StringMsg), + SendOptions::default(), + ) + .await; + Trailers::new(Ok(())) + } + } + + let registry = mock_registry(); + let messages = Arc::new(Mutex::new(Vec::new())); + let resp_headers = Arc::new(Mutex::new(None)); + let mut tx = MockSendStream { + messages: messages.clone(), + headers: resp_headers.clone(), + msg_disable_compression: Arc::new(Mutex::new(Vec::new())), + }; + let interceptor = ServerCompressionInterceptor::new(registry); + let chain = OverrideHandler.with_interceptor(interceptor); + + let mut headers = RequestHeaders::new(); + // Client ONLY accepts gzip + headers + .metadata_mut() + .insert("grpc-accept-encoding", "gzip".parse().unwrap()); + let rx = MockRecvStream { + items: vec![Ok(Bytes::from_static(b"hello"))], + }; + + let trailers = chain + .handle(headers, CallOptions::default(), &mut tx, rx) + .await; + assert!(trailers.status().is_ok()); + + // The interceptor should have stripped the invalid "mock" header! + let h = resp_headers.lock().await.take().unwrap(); + assert!(h.metadata().get("grpc-encoding").is_none()); + + // And the message should NOT be compressed. + let msgs = messages.lock().await; + assert_eq!(msgs[0], Bytes::from_static(b"echo")); + } + + #[tokio::test] + async fn test_underlying_stream_error_propagation() { + struct ErrorPropHandler; + impl Handle for ErrorPropHandler { + async fn handle( + &self, + _h: RequestHeaders, + _o: CallOptions, + _tx: &mut impl SendStream, + mut rx: impl RecvStream + 'static, + ) -> Trailers { + struct StringMsg; + impl RecvMessage for StringMsg { + fn decode(&mut self, _data: &mut dyn Buf) -> Result<(), String> { + Ok(()) + } + } + let res = rx.next(&mut StringMsg).await; + assert!(matches!(res, Some(Err(())))); + Trailers::new(Err(StatusError::new( + crate::status::StatusCodeError::Internal, + "propagated", + ))) + } + } + + let registry = mock_registry(); + let mut tx = MockSendStream { + messages: Arc::new(Mutex::new(Vec::new())), + headers: Arc::new(Mutex::new(None)), + msg_disable_compression: Arc::new(Mutex::new(Vec::new())), + }; + let interceptor = ServerCompressionInterceptor::new(registry); + let chain = ErrorPropHandler.with_interceptor(interceptor); + + let mut headers = RequestHeaders::new(); + headers + .metadata_mut() + .insert("grpc-encoding", "mock".parse().unwrap()); + let rx = MockRecvStream { + items: vec![Err(())], + }; + + let trailers = chain + .handle(headers, CallOptions::default(), &mut tx, rx) + .await; + assert_eq!( + trailers.status().as_ref().unwrap_err().code(), + crate::status::StatusCodeError::Internal + ); + } + + #[tokio::test] + async fn test_compression_encoding_failure() { + struct FailingEncodeHandler; + impl Handle for FailingEncodeHandler { + async fn handle( + &self, + _h: RequestHeaders, + _o: CallOptions, + tx: &mut impl SendStream, + _rx: impl RecvStream + 'static, + ) -> Trailers { + // Send headers first to transition state machine to Active + let _ = tx + .send( + ResponseStreamItem::Headers(ResponseHeaders::new()), + SendOptions::default(), + ) + .await; + + struct BadMsg; + impl SendMessage for BadMsg { + fn encode(&self) -> Result, String> { + Err("encode failed".into()) + } + } + let res = tx + .send(ResponseStreamItem::Message(&BadMsg), SendOptions::default()) + .await; + assert!(res.is_err()); + Trailers::new(Ok(())) + } + } + + let registry = mock_registry(); + let mut tx = MockSendStream { + messages: Arc::new(Mutex::new(Vec::new())), + headers: Arc::new(Mutex::new(None)), + msg_disable_compression: Arc::new(Mutex::new(Vec::new())), + }; + let interceptor = ServerCompressionInterceptor::new(registry); + let chain = FailingEncodeHandler.with_interceptor(interceptor); + + let mut headers = RequestHeaders::new(); + headers + .metadata_mut() + .insert("grpc-encoding", "mock".parse().unwrap()); + headers + .metadata_mut() + .insert("grpc-accept-encoding", "mock".parse().unwrap()); + + let _ = chain + .handle( + headers, + CallOptions::default(), + &mut tx, + MockRecvStream { items: vec![] }, + ) + .await; + } + + #[tokio::test] + async fn test_post_decompression_decoding_failure() { + struct FailingDecodeHandler; + impl Handle for FailingDecodeHandler { + async fn handle( + &self, + _h: RequestHeaders, + _o: CallOptions, + _tx: &mut impl SendStream, + mut rx: impl RecvStream + 'static, + ) -> Trailers { + struct BadMsg; + impl RecvMessage for BadMsg { + fn decode(&mut self, _data: &mut dyn Buf) -> Result<(), String> { + Err("decode failed".into()) + } + } + let res = rx.next(&mut BadMsg).await; + assert!(matches!(res, Some(Err(())))); + Trailers::new(Ok(())) + } + } + + let registry = mock_registry(); + let mut tx = MockSendStream { + messages: Arc::new(Mutex::new(Vec::new())), + headers: Arc::new(Mutex::new(None)), + msg_disable_compression: Arc::new(Mutex::new(Vec::new())), + }; + let interceptor = ServerCompressionInterceptor::new(registry); + let chain = FailingDecodeHandler.with_interceptor(interceptor); + + let mut headers = RequestHeaders::new(); + headers + .metadata_mut() + .insert("grpc-encoding", "mock".parse().unwrap()); + let rx = MockRecvStream { + items: vec![Ok(Bytes::from_static(b"[compressed]valid bytes"))], + }; + + let _ = chain + .handle(headers, CallOptions::default(), &mut tx, rx) + .await; + } + + #[tokio::test] + async fn test_decompression_limit_exceeded() { + struct LimitHandler; + impl Handle for LimitHandler { + async fn handle( + &self, + _h: RequestHeaders, + _o: CallOptions, + _tx: &mut impl SendStream, + mut rx: impl RecvStream + 'static, + ) -> Trailers { + struct StringMsg; + impl RecvMessage for StringMsg { + fn decode(&mut self, _data: &mut dyn Buf) -> Result<(), String> { + Ok(()) + } + } + let res = rx.next(&mut StringMsg).await; + assert!(matches!(res, Some(Err(())))); + Trailers::new(Err(StatusError::new( + crate::status::StatusCodeError::Internal, + "limit exceeded", + ))) + } + } + + let registry = mock_registry(); + let mut tx = MockSendStream { + messages: Arc::new(Mutex::new(Vec::new())), + headers: Arc::new(Mutex::new(None)), + msg_disable_compression: Arc::new(Mutex::new(Vec::new())), + }; + let interceptor = ServerCompressionInterceptor::new(registry).with_decompression_limit(3); + let chain = LimitHandler.with_interceptor(interceptor); + + let mut headers = RequestHeaders::new(); + headers + .metadata_mut() + .insert("grpc-encoding", "mock".parse().unwrap()); + let rx = MockRecvStream { + items: vec![Ok(Bytes::from_static( + b"[compressed]long payload exceeding limit", + ))], + }; + + let trailers = chain + .handle(headers, CallOptions::default(), &mut tx, rx) + .await; + assert_eq!( + trailers.status().as_ref().unwrap_err().code(), + crate::status::StatusCodeError::Internal + ); + } + + #[tokio::test] + async fn test_get_compressor_none() { + // A registry that can decompress "mock" but has no matching compressor, + // so the response send path is disabled (uncompressed) while the request + // is still decompressed. + let registry = CompressionRegistryBuilder::new() + .register_decompressor(Arc::new(MockCodec)) + .build(); + let messages = Arc::new(Mutex::new(Vec::new())); + let headers = Arc::new(Mutex::new(None)); + let mut tx = MockSendStream { + messages: messages.clone(), + headers: headers.clone(), + msg_disable_compression: Arc::new(Mutex::new(Vec::new())), + }; + + let interceptor = ServerCompressionInterceptor::new(registry); + let chain = MockHandler.with_interceptor(interceptor); + + let mut req_headers = RequestHeaders::new(); + req_headers + .metadata_mut() + .insert(GRPC_ENCODING_HEADER, "mock".parse().unwrap()); + req_headers + .metadata_mut() + .insert(GRPC_ACCEPT_ENCODING_HEADER, "mock".parse().unwrap()); + + let rx = MockRecvStream { + items: vec![Ok(Bytes::from_static(b"[compressed]hello"))], + }; + + let trailers = chain + .handle(req_headers, CallOptions::default(), &mut tx, rx) + .await; + + assert!(trailers.status().is_ok()); + + let h = headers.lock().await.take().unwrap(); + assert!(h.metadata().get(GRPC_ENCODING_HEADER).is_none()); + + let msgs = messages.lock().await; + assert_eq!(msgs.len(), 1); + assert_eq!(msgs[0], Bytes::from_static(b"echo")); + } + + #[test] + fn test_default_interceptor() { + let _interceptor = ServerCompressionInterceptor::default(); + } +}