vector/sources/util/grpc/decompression.rs
1use std::{
2 cmp,
3 future::Future,
4 io::{self, Write},
5 mem,
6 pin::Pin,
7 sync::LazyLock,
8 task::{Context, Poll},
9};
10
11use bytes::{Buf, BufMut, BytesMut};
12use flate2::write::GzDecoder;
13use futures_util::FutureExt;
14use http::{HeaderValue, Request, Response};
15use hyper::{
16 Body,
17 body::{HttpBody, Sender},
18};
19use tokio::{pin, select};
20use tonic::{Status, body::BoxBody, metadata::AsciiMetadataValue};
21use tower::{Layer, Service};
22use vector_lib::internal_event::{
23 ByteSize, BytesReceived, InternalEventHandle as _, Protocol, Registered,
24};
25
26use crate::internal_events::{GrpcError, GrpcInvalidCompressionSchemeError};
27use crate::sources::util::decompression::{
28 CappedDecoder, DecompressedSizeLimitExceeded, is_decompressed_size_limit_error,
29 max_decompressed_size_bytes, max_zlib_compressed_frame_size_bytes,
30};
31
32// Every gRPC message has a five byte header:
33// - a compressed flag (u8, 0/1 for compressed/decompressed)
34// - a length prefix, indicating the number of remaining bytes to read (u32)
35const GRPC_MESSAGE_HEADER_LEN: usize = mem::size_of::<u8>() + mem::size_of::<u32>();
36// Fixed container framing a valid frame adds on top of zlib's worst-case expansion. Added to the
37// compressed-frame pre-filter so a small cap does not reject a valid zstd frame (or a typical gzip
38// frame) whose decompressed size is within the cap.
39//
40// zstd's frame header is bounded at 22 bytes: 4 magic + 1 frame-header descriptor + 1 window
41// descriptor + 4 dictionary ID + 8 frame-content size + 4 content checksum (RFC 8878 §3.1.1), so
42// 22 fully covers zstd. gzip's mandatory framing is only 18 bytes (10 header + 8 trailer), but its
43// optional FNAME, FCOMMENT and FEXTRA fields are unbounded (RFC 1952 §2.3.1): a gzip frame carrying
44// more than this slack in those optional fields could still be rejected here. Encoders don't emit
45// them in practice, so 22 covers the realistic case; the prefilter is only a cheap wire-size guard
46// and the authoritative per-output cap is still enforced during decompression.
47const GRPC_COMPRESSED_FRAME_OVERHEAD_SLACK: usize = 22;
48const GRPC_ENCODING_HEADER: &str = "grpc-encoding";
49const GRPC_ACCEPT_ENCODING_HEADER: &str = "grpc-accept-encoding";
50
51// The encodings this layer advertises to clients via `grpc-accept-encoding`.
52// Each variant maps to a `CompressionScheme` (or `None` for `identity`) through
53// `to_scheme`, so adding a variant here forces the decompression match to be
54// updated and the advertised list cannot drift from the schemes actually handled.
55#[derive(Clone, Copy)]
56enum AdvertisedEncoding {
57 Gzip,
58 Zstd,
59 Identity,
60}
61
62impl AdvertisedEncoding {
63 const ALL: &'static [Self] = &[Self::Gzip, Self::Zstd, Self::Identity];
64
65 const fn as_str(self) -> &'static str {
66 match self {
67 Self::Gzip => "gzip",
68 Self::Zstd => "zstd",
69 Self::Identity => "identity",
70 }
71 }
72
73 fn parse(s: &str) -> Option<Self> {
74 Self::ALL.iter().copied().find(|e| e.as_str() == s)
75 }
76
77 // `identity` is the gRPC no-op encoding: the request body is already
78 // uncompressed, so there's nothing to decompress.
79 const fn to_scheme(self) -> Option<CompressionScheme> {
80 match self {
81 Self::Gzip => Some(CompressionScheme::Gzip),
82 Self::Zstd => Some(CompressionScheme::Zstd),
83 Self::Identity => None,
84 }
85 }
86}
87
88// Advertised to clients via `grpc-accept-encoding`. Derived from
89// `AdvertisedEncoding::ALL` so this layer is the single owner of gRPC compression
90// negotiation for all Vector gRPC sources and the header value cannot drift from
91// the set of schemes actually handled.
92static GRPC_ACCEPT_ENCODING_VALUE: LazyLock<String> = LazyLock::new(|| {
93 AdvertisedEncoding::ALL
94 .iter()
95 .map(|e| e.as_str())
96 .collect::<Vec<_>>()
97 .join(",")
98});
99
100enum CompressionScheme {
101 Gzip,
102 Zstd,
103}
104
105impl CompressionScheme {
106 fn from_encoding_header(req: &Request<Body>) -> Result<Option<Self>, Status> {
107 req.headers()
108 .get(GRPC_ENCODING_HEADER)
109 .map(|s| {
110 s.to_str().map(|s| s.to_string()).map_err(|_| {
111 Status::unimplemented(format!(
112 "`{GRPC_ENCODING_HEADER}` contains non-visible characters and is not a valid encoding"
113 ))
114 })
115 })
116 .transpose()
117 .and_then(|value| match value {
118 None => Ok(None),
119 Some(scheme) => match AdvertisedEncoding::parse(&scheme) {
120 Some(encoding) => Ok(encoding.to_scheme()),
121 None => Err(Status::unimplemented(format!(
122 "compression scheme `{scheme}` is not supported"
123 ))),
124 },
125 })
126 .map_err(|mut status| {
127 status.metadata_mut().insert(
128 GRPC_ACCEPT_ENCODING_HEADER,
129 AsciiMetadataValue::try_from(GRPC_ACCEPT_ENCODING_VALUE.as_str())
130 .expect("advertised encoding value must be valid ASCII"),
131 );
132 status
133 })
134 }
135}
136
137#[derive(Default)]
138enum State {
139 #[default]
140 WaitingForHeader,
141 Forward {
142 overall_len: usize,
143 },
144 Decompress {
145 remaining: usize,
146 },
147}
148
149/// Maps a decompressor `io::Error` to a gRPC [`Status`]: an oversized payload becomes
150/// `out_of_range` (a client fault, matching the existing >4GB handling) while anything else falls
151/// back to `internal` with `internal_msg`.
152fn decompressor_error_to_status(error: &io::Error, internal_msg: &'static str) -> Status {
153 if is_decompressed_size_limit_error(error) {
154 Status::out_of_range("decompressed message exceeds the maximum allowed size")
155 } else {
156 Status::internal(internal_msg)
157 }
158}
159
160/// A `Write` sink that appends into a `Vec` but refuses to grow past `max_len`, so a streaming
161/// decompressor errors out *during* decompression rather than first materializing an oversized
162/// output and only then having its size checked.
163struct LimitedWriter {
164 buf: Vec<u8>,
165 max_len: usize,
166}
167
168impl LimitedWriter {
169 const fn new(buf: Vec<u8>, max_len: usize) -> Self {
170 Self { buf, max_len }
171 }
172
173 fn into_inner(self) -> Vec<u8> {
174 self.buf
175 }
176}
177
178impl Write for LimitedWriter {
179 fn write(&mut self, data: &[u8]) -> io::Result<usize> {
180 if self.buf.len().saturating_add(data.len()) > self.max_len {
181 return Err(io::Error::other(DecompressedSizeLimitExceeded));
182 }
183 self.buf.extend_from_slice(data);
184 Ok(data.len())
185 }
186
187 fn flush(&mut self) -> io::Result<()> {
188 Ok(())
189 }
190}
191
192enum Decompressor {
193 Gzip {
194 decoder: Box<GzDecoder<LimitedWriter>>,
195 },
196 Zstd {
197 compressed: Vec<u8>,
198 output_buf: Vec<u8>,
199 limit: usize,
200 },
201}
202
203impl Decompressor {
204 fn new(scheme: &CompressionScheme) -> Result<Self, io::Error> {
205 // Create the backing buffer for the decompressor and set the compression flag to false (0)
206 // and pre-allocate the space for the length prefix, which we'll fill out once we've
207 // finalized the decompressor.
208 let buf = vec![0; GRPC_MESSAGE_HEADER_LEN];
209 // Cap the decompressed output so a compression bomb on this unauthenticated gRPC
210 // listener cannot drive unbounded allocation.
211 let limit = max_decompressed_size_bytes();
212 match scheme {
213 // The gzip output buffer already holds the 5-byte header, so the sink may grow to the
214 // header plus the decompressed cap; anything larger errors mid-decompression.
215 CompressionScheme::Gzip => Ok(Decompressor::Gzip {
216 decoder: Box::new(GzDecoder::new(LimitedWriter::new(
217 buf,
218 GRPC_MESSAGE_HEADER_LEN.saturating_add(limit),
219 ))),
220 }),
221 CompressionScheme::Zstd => Ok(Decompressor::Zstd {
222 compressed: Vec::new(),
223 output_buf: buf,
224 limit,
225 }),
226 }
227 }
228
229 fn write_all(&mut self, data: &[u8]) -> io::Result<()> {
230 match self {
231 // The `LimitedWriter` bounds the decompressed output, so a gzip bomb errors here rather
232 // than materializing an oversized buffer before the size is checked.
233 Decompressor::Gzip { decoder } => decoder.write_all(data),
234 Decompressor::Zstd { compressed, .. } => {
235 compressed.extend_from_slice(data);
236 Ok(())
237 }
238 }
239 }
240
241 fn finish(self) -> io::Result<Vec<u8>> {
242 match self {
243 Decompressor::Gzip { decoder } => (*decoder).finish().map(LimitedWriter::into_inner),
244 // Decode directly into output_buf to avoid a temporary intermediate Vec that
245 // decode_all would produce; peak memory is compressed + decompressed rather than
246 // compressed + 2 × decompressed. The output is bounded by `limit` via
247 // CappedDecoder, and zstd's internal window is bounded to match so a crafted
248 // frame cannot force a large window allocation.
249 Decompressor::Zstd {
250 compressed,
251 mut output_buf,
252 limit,
253 } => {
254 // The `CappedReader` errors out once the decompressed output would exceed the cap,
255 // so a zstd bomb surfaces as a size-limit error here rather than overflowing
256 // `output_buf`.
257 let mut reader =
258 CappedDecoder::zstd_with_limit(io::Cursor::new(&compressed), limit)?
259 .into_reader();
260 io::copy(&mut reader, &mut output_buf)?;
261 Ok(output_buf)
262 }
263 }
264 }
265}
266
267async fn drive_body_decompression(
268 mut source: Body,
269 mut destination: Sender,
270 scheme: Option<CompressionScheme>,
271) -> Result<usize, Status> {
272 let mut state = State::default();
273 let mut buf = BytesMut::new();
274 let mut decompressor: Option<Decompressor> = None;
275 let mut bytes_received = 0;
276
277 // Drain all message chunks from the body first.
278 while let Some(result) = source.data().await {
279 let chunk = result.map_err(|_| Status::internal("failed to read from underlying body"))?;
280 buf.put(chunk);
281
282 let maybe_message = loop {
283 match state {
284 State::WaitingForHeader => {
285 // If we don't have enough data yet to even read the gRPC message header, we can't do anything yet.
286 if buf.len() < GRPC_MESSAGE_HEADER_LEN {
287 break None;
288 }
289
290 // Extract the compressed flag and length prefix.
291 let (is_compressed, message_len) = {
292 let header = &buf[..GRPC_MESSAGE_HEADER_LEN];
293
294 let message_len_raw: u32 = header[1..]
295 .try_into()
296 .map(u32::from_be_bytes)
297 .expect("there must be four bytes remaining in the header slice");
298 let message_len = message_len_raw
299 .try_into()
300 .expect("Vector does not support 16-bit platforms");
301
302 (header[0] == 1, message_len)
303 };
304
305 // Now, if the message is not compressed, then put ourselves into forward mode, where we'll wait for
306 // the rest of the message to come in -- decoding isn't streaming so there's no benefit there --
307 // before we emit it.
308 //
309 // If the message _is_ compressed, we do roughly the same thing but we shove it into the
310 // decompressor incrementally because there's no good reason to make both the internal buffer and
311 // the decompressor buffer expand if we don't have to.
312 if is_compressed {
313 // Per the gRPC compression spec, the compressed flag requires a
314 // negotiated encoding. Reject frames that set it under identity
315 // (or with no `grpc-encoding` header) rather than silently
316 // falling back to gzip and masking client/server mismatches.
317 if scheme.is_none() {
318 return Err(Status::internal(
319 "received compressed frame but no compression scheme was negotiated",
320 ));
321 }
322
323 // Reject a compressed payload whose declared wire size could not
324 // legitimately decompress within the cap, before we buffer any of it. The
325 // bound (decompressed cap plus zlib's worst-case expansion, shared with the
326 // logstash source) keeps a peer from advertising a huge length and
327 // slow-streaming bytes to grow the decompressor's input buffer unbounded.
328 //
329 // The zlib expansion factor covers gzip and zstd too, but those formats add
330 // a few bytes of fixed container framing (gzip header/trailer, zstd frame
331 // header) on top. Add a small fixed slack so a tiny cap cannot falsely
332 // reject a valid gzip/zstd frame whose decompressed size is within the cap;
333 // the authoritative per-output cap is still enforced during decompression.
334 let compressed_frame_limit = max_zlib_compressed_frame_size_bytes()
335 .saturating_add(GRPC_COMPRESSED_FRAME_OVERHEAD_SLACK);
336 if message_len > compressed_frame_limit {
337 return Err(Status::out_of_range(
338 "compressed message length exceeds the maximum allowed size",
339 ));
340 }
341
342 // We skip the header in the buffer because it doesn't matter to the decompressor and we
343 // recreate it anyways.
344 buf.advance(GRPC_MESSAGE_HEADER_LEN);
345
346 state = State::Decompress {
347 remaining: message_len,
348 };
349 } else {
350 // Reject an identity (uncompressed) message larger than the cap before
351 // buffering it to `overall_len`, so a large declared length cannot drive
352 // unbounded buffering here ahead of tonic's own decode-size limit.
353 if message_len > max_decompressed_size_bytes() {
354 return Err(Status::out_of_range(
355 "message length exceeds the maximum allowed size",
356 ));
357 }
358
359 let overall_len = GRPC_MESSAGE_HEADER_LEN + message_len;
360 state = State::Forward { overall_len };
361 }
362 }
363 State::Forward { overall_len } => {
364 // All we're doing at this point is waiting until we have all the bytes for the current gRPC message
365 // before we emit them to the caller.
366 if buf.len() < overall_len {
367 break None;
368 }
369
370 // Now that we have all the bytes we need, slice them out of our internal buffer, reset our state,
371 // and hand the message back to the caller.
372 let message = buf.split_to(overall_len).freeze();
373 state = State::WaitingForHeader;
374
375 bytes_received += overall_len;
376
377 break Some(message);
378 }
379 State::Decompress { ref mut remaining } => {
380 if *remaining > 0 {
381 // We're waiting for `remaining` more bytes to feed to the decompressor before we finalize it and
382 // generate our new chunk of data. We might have data in our internal buffer, so try and drain that
383 // first before polling the underlying body for more.
384 let available = buf.len();
385 if available > 0 {
386 // Write the lesser of what the buffer has, or what is remaining for the current message, into
387 // the decompressor. This is _technically_ synchronous but there's really no way to do it
388 // asynchronously since we already have the data, and that's the only asynchronous part.
389 let to_take = cmp::min(available, *remaining);
390 let d = match &mut decompressor {
391 Some(d) => d,
392 slot @ None => {
393 let scheme = scheme.as_ref().expect(
394 "compressed frames without a negotiated scheme are rejected earlier",
395 );
396 slot.insert(Decompressor::new(scheme).map_err(|_| {
397 Status::internal("failed to initialize decompressor")
398 })?)
399 }
400 };
401 if let Err(error) = d.write_all(&buf[..to_take]) {
402 return Err(decompressor_error_to_status(
403 &error,
404 "failed to write to decompressor",
405 ));
406 }
407
408 *remaining -= to_take;
409 buf.advance(to_take);
410 } else {
411 break None;
412 }
413 } else {
414 // We don't need any more data, so consume the decompressor, finalize it by updating the length
415 // prefix, and then pass it back to the caller.
416 let result = decompressor
417 .take()
418 .expect("consumed decompressor when no decompressor was present")
419 .finish();
420
421 // Decompression can fail here either because the payload exceeded the size
422 // cap (an oversized-request client fault) or, for malformed input, during
423 // finalization; map the former to `out_of_range` and treat anything else as
424 // an internal error.
425 let mut buf = result.map_err(|error| {
426 decompressor_error_to_status(
427 &error,
428 "reached impossible error during decompressor finalization",
429 )
430 })?;
431 bytes_received += buf.len();
432
433 // Write the length of our decompressed message in the pre-allocated slot for the message's length prefix.
434 let message_len_actual = buf.len() - GRPC_MESSAGE_HEADER_LEN;
435 let message_len = u32::try_from(message_len_actual).map_err(|_| {
436 Status::out_of_range("messages greater than 4GB are not supported")
437 })?;
438
439 let message_len_bytes = message_len.to_be_bytes();
440 let message_len_slot = &mut buf[1..GRPC_MESSAGE_HEADER_LEN];
441 message_len_slot.copy_from_slice(&message_len_bytes[..]);
442
443 // Reset our state before returning the decompressed message.
444 state = State::WaitingForHeader;
445
446 break Some(buf.into());
447 }
448 }
449 }
450 };
451
452 if let Some(message) = maybe_message {
453 // We got a decompressed (or passthrough) message chunk, so just forward it to the destination.
454 if destination.send_data(message).await.is_err() {
455 return Err(Status::internal("destination body abnormally closed"));
456 }
457 }
458 }
459
460 // When we've exhausted all the message chunks, we try sending any trailers that came in on the underlying body.
461 let result = source.trailers().await;
462 let maybe_trailers =
463 result.map_err(|_| Status::internal("error reading trailers from underlying body"))?;
464 if let Some(trailers) = maybe_trailers
465 && destination.send_trailers(trailers).await.is_err()
466 {
467 return Err(Status::internal("destination body abnormally closed"));
468 }
469
470 Ok(bytes_received)
471}
472
473async fn drive_request<F, E>(
474 source: Body,
475 destination: Sender,
476 inner: F,
477 bytes_received: Registered<BytesReceived>,
478 scheme: Option<CompressionScheme>,
479) -> Result<Response<BoxBody>, E>
480where
481 F: Future<Output = Result<Response<BoxBody>, E>>,
482 E: std::fmt::Display,
483{
484 let body_decompression = drive_body_decompression(source, destination, scheme);
485
486 pin!(inner);
487 pin!(body_decompression);
488
489 let mut body_eof = false;
490 let mut body_bytes_received = 0;
491
492 let mut result = loop {
493 select! {
494 biased;
495
496 // Drive the inner future, as this will be consuming the message chunks we give it.
497 result = &mut inner => break result,
498
499 // Drive the core decompression loop, reading chunks from the underlying body, decompressing them if needed,
500 // and eventually handling trailers at the end, if they're present.
501 result = &mut body_decompression, if !body_eof => match result {
502 Err(e) => break Ok(e.to_http()),
503 Ok(bytes_received) => {
504 body_bytes_received = bytes_received;
505 body_eof = true;
506 },
507 }
508 }
509 };
510
511 // If the response indicates success, then emit the necessary metrics
512 // otherwise emit the error.
513 match &result {
514 Ok(res) if res.status().is_success() => {
515 bytes_received.emit(ByteSize(body_bytes_received));
516 }
517 Ok(res) => {
518 emit!(GrpcError {
519 error: format!("Received {}", res.status())
520 });
521 }
522 Err(error) => {
523 emit!(GrpcError { error: &error });
524 }
525 };
526
527 // Advertise the set of compression schemes this layer can accept to the client.
528 // Since this layer is the single owner of compression negotiation, individual
529 // services no longer call `.accept_compressed(..)` and therefore tonic would not
530 // set this header itself.
531 if let Ok(res) = result.as_mut() {
532 res.headers_mut().insert(
533 GRPC_ACCEPT_ENCODING_HEADER,
534 HeaderValue::from_str(&GRPC_ACCEPT_ENCODING_VALUE)
535 .expect("advertised encoding value must be valid ASCII"),
536 );
537 }
538
539 result
540}
541
542#[derive(Clone)]
543pub struct DecompressionAndMetrics<S> {
544 inner: S,
545 bytes_received: Registered<BytesReceived>,
546}
547
548impl<S> Service<Request<Body>> for DecompressionAndMetrics<S>
549where
550 S: Service<Request<Body>, Response = Response<BoxBody>> + Clone + Send + 'static,
551 S::Future: Send + 'static,
552 S::Error: std::fmt::Display,
553{
554 type Response = Response<BoxBody>;
555 type Error = S::Error;
556 type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;
557
558 fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
559 self.inner.poll_ready(cx)
560 }
561
562 fn call(&mut self, req: Request<Body>) -> Self::Future {
563 match CompressionScheme::from_encoding_header(&req) {
564 // There was a header for the encoding, but it was either invalid data or a scheme we don't support.
565 Err(status) => {
566 emit!(GrpcInvalidCompressionSchemeError { status: &status });
567 Box::pin(async move { Ok(status.to_http()) })
568 }
569
570 // The request either isn't using compression, or it has indicated compression may be used and we know we
571 // can support decompression based on the indicated compression scheme... so wrap the body to decompress, if
572 // need be, and then track the bytes that flowed through.
573 Ok(scheme) => {
574 let (destination, decompressed_body) = Body::channel();
575 let (mut req_parts, req_body) = req.into_parts();
576 // Since this layer owns compression negotiation and is about to hand the
577 // inner service a fully decompressed body (with the per-message compressed
578 // flag cleared), strip the `grpc-encoding` header so tonic's codegen treats
579 // the request as uncompressed and does not try to validate the encoding
580 // against any per-service `accept_compressed(..)` configuration.
581 if scheme.is_some() {
582 req_parts.headers.remove(GRPC_ENCODING_HEADER);
583 }
584 let mapped_req = Request::from_parts(req_parts, decompressed_body);
585
586 let inner = self.inner.call(mapped_req);
587
588 drive_request(
589 req_body,
590 destination,
591 inner,
592 self.bytes_received.clone(),
593 scheme,
594 )
595 .boxed()
596 }
597 }
598 }
599}
600
601/// A layer for decompressing Tonic request payloads and emitting telemetry for the payload sizes.
602///
603/// In some cases, we configure `tonic` to use compression on requests to save CPU and throughput when sending those
604/// large requests. In the case of Vector-to-Vector communication, this means the Vector v2 source may deal with
605/// compressed requests. The code already transparently handles decompression, but as part of our component
606/// specification, we have specific goals around what event representations we pay attention to.
607///
608/// In the case of tracking bytes sent/received, we always want to track the number of bytes received _after_
609/// decompression to faithfully represent the amount of data being processed by Vector. This poses a problem with the
610/// out-of-the-box `tonic` codegen as there is no hook whatsoever to inspect the raw request payload (after
611/// decompression, if it was compressed at all) prior to the payload being decoded as a Protocol Buffers payload.
612///
613/// This layer wraps the incoming body in our own body type, which allows us to do two things: decompress the payload
614/// before it enters the decoding phase, and emit metrics based on the decompressed payload.
615///
616/// Since we can see the decompressed bytes, and also know if the underlying service responded successfully -- i.e. the
617/// request was valid, and was processed -- we can now report the number of bytes (after decompression) that were
618/// received _and_ processed correctly.
619///
620/// The supported compression schemes are gzip and zstd.
621#[derive(Clone, Default)]
622pub struct DecompressionAndMetricsLayer;
623
624impl<S> Layer<S> for DecompressionAndMetricsLayer {
625 type Service = DecompressionAndMetrics<S>;
626
627 fn layer(&self, inner: S) -> Self::Service {
628 DecompressionAndMetrics {
629 inner,
630 bytes_received: register!(BytesReceived::from(Protocol::from("grpc"))),
631 }
632 }
633}