mirror of
https://github.com/rustfs/rustfs.git
synced 2026-09-02 10:18:10 +00:00
GatewayJ
1571 lines
61 KiB
Rust
1571 lines
61 KiB
Rust
// Copyright 2024 RustFS Team
|
|
//
|
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
|
// you may not use this file except in compliance with the License.
|
|
// You may obtain a copy of the License at
|
|
//
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
//
|
|
// Unless required by applicable law or agreed to in writing, software
|
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
// See the License for the specific language governing permissions and
|
|
// limitations under the License.
|
|
|
|
use crate::{
|
|
MAX_ERROR_SOURCE_DEPTH, SelectError, SelectInputMetrics, metrics::SelectInputMetricsRecorder,
|
|
query::session::QueryExecutionGuard,
|
|
};
|
|
use async_compression::tokio::bufread::BzDecoder;
|
|
use bytes::{Buf as _, Bytes};
|
|
use datafusion::object_store::{Error as ObjectStoreError, Result as ObjectStoreResult};
|
|
use flate2::bufread::GzDecoder;
|
|
use futures::{StreamExt, stream};
|
|
use futures_core::stream::BoxStream;
|
|
use std::{
|
|
error::Error as StdError,
|
|
io::{self, BufRead as _, Read as _},
|
|
pin::Pin,
|
|
sync::Arc,
|
|
task::{Context, Poll},
|
|
};
|
|
use tokio::{
|
|
io::{AsyncRead, AsyncReadExt, BufReader, ReadBuf},
|
|
sync::{mpsc, oneshot},
|
|
};
|
|
use tokio_stream::wrappers::ReceiverStream;
|
|
use tokio_util::io::{ReaderStream, StreamReader};
|
|
|
|
pub(crate) const MAX_SELECT_RECORD_BYTES: usize = 1024 * 1024;
|
|
const MAX_SELECT_PROCESSED_BYTES: u64 = 5 * 1024 * 1024 * 1024 * 1024;
|
|
pub(crate) const SELECT_DECODE_CHUNK_BYTES: usize = 64 * 1024;
|
|
const DECOMPRESSION_CHANNEL_CAPACITY: usize = 2;
|
|
|
|
pub(crate) type SelectInputReader = Box<dyn AsyncRead + Unpin + Send>;
|
|
|
|
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
|
|
pub(crate) enum CompressionFormat {
|
|
Gzip,
|
|
Bzip2,
|
|
}
|
|
|
|
impl CompressionFormat {
|
|
fn name(self) -> &'static str {
|
|
match self {
|
|
Self::Gzip => "GZIP",
|
|
Self::Bzip2 => "BZIP2",
|
|
}
|
|
}
|
|
|
|
fn invalid_header_error(self) -> SelectError {
|
|
SelectError::InvalidCompressionFormatForObject {
|
|
compression: self.name(),
|
|
}
|
|
}
|
|
}
|
|
|
|
pub(crate) fn processed_bytes_limit() -> u64 {
|
|
// Processed throughput is independent of compression ratio and live
|
|
// memory; use the S3 Select object-size ceiling as the absolute bound.
|
|
MAX_SELECT_PROCESSED_BYTES
|
|
}
|
|
|
|
pub(crate) fn compressed_input_reader(
|
|
reader: SelectInputReader,
|
|
compressed_size: u64,
|
|
format: CompressionFormat,
|
|
input_metrics: Arc<SelectInputMetrics>,
|
|
max_processed_bytes: u64,
|
|
query_guard: Option<QueryExecutionGuard>,
|
|
) -> SelectInputReader {
|
|
// rustfs-zip does not expose Select's streaming metrics, member validation,
|
|
// typed errors, or cancellation contract, so the protocol adapter lives here.
|
|
let input_metrics = input_metrics.recorder();
|
|
let reader = ScannedReader::new(reader, compressed_size, input_metrics.clone());
|
|
let reader = CooperativeReader::new(reader);
|
|
let decoder = match format {
|
|
CompressionFormat::Gzip => blocking_gzip_reader(Box::new(reader), query_guard),
|
|
CompressionFormat::Bzip2 => blocking_bzip2_reader(Box::new(Bzip2HeaderValidatingReader::new(reader)), query_guard),
|
|
};
|
|
Box::new(CooperativeReader::new(ProcessedReader::new(decoder, input_metrics, max_processed_bytes)))
|
|
}
|
|
|
|
fn blocking_gzip_reader(reader: SelectInputReader, query_guard: Option<QueryExecutionGuard>) -> SelectInputReader {
|
|
// The bounded bridge keeps RFC 1952 decoding off Tokio workers while
|
|
// preserving member validation, backpressure, and reader cancellation.
|
|
let (compressed_tx, compressed_rx) = mpsc::channel(DECOMPRESSION_CHANNEL_CAPACITY);
|
|
let (decoded_tx, decoded_rx) = mpsc::channel(DECOMPRESSION_CHANNEL_CAPACITY);
|
|
|
|
let decoded_closed = decoded_tx.clone();
|
|
drop(tokio::spawn(async move {
|
|
let mut stream = ReaderStream::with_capacity(reader, SELECT_DECODE_CHUNK_BYTES);
|
|
loop {
|
|
let item = tokio::select! {
|
|
biased;
|
|
_ = decoded_closed.closed() => break,
|
|
item = stream.next() => item,
|
|
};
|
|
let Some(item) = item else {
|
|
break;
|
|
};
|
|
let sent = tokio::select! {
|
|
biased;
|
|
_ = decoded_closed.closed() => false,
|
|
result = compressed_tx.send(item) => result.is_ok(),
|
|
};
|
|
if !sent {
|
|
break;
|
|
}
|
|
}
|
|
}));
|
|
|
|
let blocking_output = decoded_tx.clone();
|
|
spawn_decoder_thread("s3select-gzip", decoded_tx, query_guard, move || {
|
|
decode_gzip(compressed_rx, blocking_output)
|
|
});
|
|
|
|
Box::new(StreamReader::new(ReceiverStream::new(decoded_rx)))
|
|
}
|
|
|
|
fn blocking_bzip2_reader(reader: SelectInputReader, query_guard: Option<QueryExecutionGuard>) -> SelectInputReader {
|
|
let (decoded_tx, decoded_rx) = mpsc::channel(DECOMPRESSION_CHANNEL_CAPACITY);
|
|
let runtime = tokio::runtime::Handle::current();
|
|
let blocking_output = decoded_tx.clone();
|
|
// A dedicated thread avoids occupying Tokio's blocking pool while the
|
|
// decoder waits for asynchronous object reads. Query admission bounds the
|
|
// number of concurrent Select decoder threads.
|
|
spawn_decoder_thread("s3select-bzip2", decoded_tx, query_guard, move || {
|
|
runtime.block_on(decode_bzip2(reader, blocking_output))
|
|
});
|
|
Box::new(StreamReader::new(ReceiverStream::new(decoded_rx)))
|
|
}
|
|
|
|
fn spawn_decoder_thread(
|
|
name: &'static str,
|
|
decoded: mpsc::Sender<io::Result<Bytes>>,
|
|
query_guard: Option<QueryExecutionGuard>,
|
|
task: impl FnOnce() -> io::Result<()> + Send + 'static,
|
|
) {
|
|
let (finished_tx, finished_rx) = oneshot::channel();
|
|
let spawn_result = std::thread::Builder::new().name(name.to_string()).spawn(move || {
|
|
let _query_guard = query_guard;
|
|
let _ = finished_tx.send(task());
|
|
});
|
|
if spawn_result.is_err() {
|
|
let _ = decoded.try_send(Err(io::Error::other(SelectError::InternalError)));
|
|
return;
|
|
}
|
|
|
|
drop(tokio::spawn(async move {
|
|
let result = finished_rx
|
|
.await
|
|
.unwrap_or_else(|_| Err(io::Error::other(SelectError::InternalError)));
|
|
if let Err(error) = result {
|
|
let _ = decoded.send(Err(error)).await;
|
|
}
|
|
}));
|
|
}
|
|
|
|
async fn decode_bzip2(reader: SelectInputReader, decoded: mpsc::Sender<io::Result<Bytes>>) -> io::Result<()> {
|
|
let reader = BufReader::with_capacity(SELECT_DECODE_CHUNK_BYTES, reader);
|
|
let mut decoder = BzDecoder::new(reader);
|
|
decoder.multiple_members(true);
|
|
let mut buffer = vec![0; SELECT_DECODE_CHUNK_BYTES];
|
|
loop {
|
|
let read = tokio::select! {
|
|
biased;
|
|
_ = decoded.closed() => return Ok(()),
|
|
result = decoder.read(&mut buffer) => result?,
|
|
};
|
|
if read == 0 {
|
|
return Ok(());
|
|
}
|
|
let bytes = Bytes::copy_from_slice(&buffer[..read]);
|
|
tokio::select! {
|
|
biased;
|
|
_ = decoded.closed() => return Ok(()),
|
|
result = decoded.send(Ok(bytes)) => {
|
|
if result.is_err() {
|
|
return Ok(());
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
fn decode_gzip(compressed: mpsc::Receiver<io::Result<Bytes>>, decoded: mpsc::Sender<io::Result<Bytes>>) -> io::Result<()> {
|
|
let mut reader = io::BufReader::with_capacity(SELECT_DECODE_CHUNK_BYTES, BlockingChannelReader::new(compressed));
|
|
let mut buffer = vec![0; SELECT_DECODE_CHUNK_BYTES];
|
|
let mut decoded_member = false;
|
|
loop {
|
|
if decoded.is_closed() {
|
|
return Ok(());
|
|
}
|
|
|
|
if reader.fill_buf()?.is_empty() {
|
|
return if decoded_member {
|
|
Ok(())
|
|
} else {
|
|
Err(io::Error::new(io::ErrorKind::UnexpectedEof, SelectError::TruncatedInput))
|
|
};
|
|
}
|
|
let header = read_gzip_header(&mut reader).map_err(|error| {
|
|
if decoded_member
|
|
&& find_error_source::<SelectError>(&error)
|
|
.is_some_and(|source| matches!(source, SelectError::InvalidCompressionFormatForObject { .. }))
|
|
{
|
|
io::Error::new(error.kind(), SelectError::TruncatedInput)
|
|
} else {
|
|
error
|
|
}
|
|
})?;
|
|
let mut decoder = GzDecoder::new(std::io::Read::chain(io::Cursor::new(header), reader));
|
|
loop {
|
|
let read = decoder.read(&mut buffer)?;
|
|
if read == 0 {
|
|
break;
|
|
}
|
|
if decoded.blocking_send(Ok(Bytes::copy_from_slice(&buffer[..read]))).is_err() {
|
|
return Ok(());
|
|
}
|
|
}
|
|
let (_, remaining) = decoder.into_inner().into_inner();
|
|
reader = remaining;
|
|
decoded_member = true;
|
|
}
|
|
}
|
|
|
|
fn read_gzip_header<R: io::BufRead>(reader: &mut R) -> io::Result<[u8; 10]> {
|
|
let mut fixed = [0; 10];
|
|
read_gzip_exact(reader, &mut fixed)?;
|
|
if fixed[..3] != [0x1f, 0x8b, 0x08] || fixed[3] & 0xe0 != 0 {
|
|
return Err(invalid_gzip_header_error());
|
|
}
|
|
let flags = fixed[3];
|
|
let mut header_crc = (flags & GZIP_FLAG_HEADER_CRC != 0).then(|| crc_fast::Digest::new(crc_fast::CrcAlgorithm::Crc32IsoHdlc));
|
|
update_gzip_header_crc(&mut header_crc, &fixed);
|
|
|
|
if flags & GZIP_FLAG_EXTRA != 0 {
|
|
let mut length = [0; 2];
|
|
read_gzip_exact(reader, &mut length)?;
|
|
update_gzip_header_crc(&mut header_crc, &length);
|
|
let extra_len = usize::from(u16::from_le_bytes(length));
|
|
read_gzip_header_bytes(reader, extra_len, &mut header_crc)?;
|
|
}
|
|
if flags & GZIP_FLAG_NAME != 0 {
|
|
read_gzip_text_field(reader, &mut header_crc)?;
|
|
}
|
|
if flags & GZIP_FLAG_COMMENT != 0 {
|
|
read_gzip_text_field(reader, &mut header_crc)?;
|
|
}
|
|
if let Some(digest) = header_crc {
|
|
let expected =
|
|
u16::try_from(digest.finalize() & u64::from(u16::MAX)).map_err(|_| io::Error::other(SelectError::InternalError))?;
|
|
let mut actual = [0; 2];
|
|
read_gzip_exact(reader, &mut actual)?;
|
|
if u16::from_le_bytes(actual) != expected {
|
|
return Err(invalid_gzip_header_error());
|
|
}
|
|
}
|
|
|
|
fixed[3] &= GZIP_FLAG_TEXT;
|
|
Ok(fixed)
|
|
}
|
|
|
|
fn read_gzip_exact<R: io::Read>(reader: &mut R, bytes: &mut [u8]) -> io::Result<()> {
|
|
reader.read_exact(bytes).map_err(|error| {
|
|
if error.kind() == io::ErrorKind::UnexpectedEof && !error_chain_contains::<CompressedSourceReadError>(&error) {
|
|
io::Error::new(io::ErrorKind::UnexpectedEof, SelectError::TruncatedInput)
|
|
} else {
|
|
error
|
|
}
|
|
})
|
|
}
|
|
|
|
fn read_gzip_header_bytes<R: io::BufRead>(
|
|
reader: &mut R,
|
|
mut remaining: usize,
|
|
header_crc: &mut Option<crc_fast::Digest>,
|
|
) -> io::Result<()> {
|
|
while remaining > 0 {
|
|
let available = reader.fill_buf()?;
|
|
if available.is_empty() {
|
|
return Err(io::Error::new(io::ErrorKind::UnexpectedEof, SelectError::TruncatedInput));
|
|
}
|
|
let consumed = remaining.min(available.len());
|
|
update_gzip_header_crc(header_crc, &available[..consumed]);
|
|
reader.consume(consumed);
|
|
remaining -= consumed;
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
fn read_gzip_text_field<R: io::BufRead>(reader: &mut R, header_crc: &mut Option<crc_fast::Digest>) -> io::Result<()> {
|
|
loop {
|
|
let available = reader.fill_buf()?;
|
|
if available.is_empty() {
|
|
return Err(io::Error::new(io::ErrorKind::UnexpectedEof, SelectError::TruncatedInput));
|
|
}
|
|
let terminator = available.iter().position(|byte| *byte == 0);
|
|
let consumed = terminator.map_or(available.len(), |position| position + 1);
|
|
update_gzip_header_crc(header_crc, &available[..consumed]);
|
|
reader.consume(consumed);
|
|
if terminator.is_some() {
|
|
return Ok(());
|
|
}
|
|
}
|
|
}
|
|
|
|
fn update_gzip_header_crc(header_crc: &mut Option<crc_fast::Digest>, bytes: &[u8]) {
|
|
if let Some(digest) = header_crc {
|
|
digest.update(bytes);
|
|
}
|
|
}
|
|
|
|
fn invalid_gzip_header_error() -> io::Error {
|
|
io::Error::new(io::ErrorKind::InvalidData, CompressionFormat::Gzip.invalid_header_error())
|
|
}
|
|
|
|
struct BlockingChannelReader {
|
|
receiver: mpsc::Receiver<io::Result<Bytes>>,
|
|
current: Bytes,
|
|
}
|
|
|
|
impl BlockingChannelReader {
|
|
fn new(receiver: mpsc::Receiver<io::Result<Bytes>>) -> Self {
|
|
Self {
|
|
receiver,
|
|
current: Bytes::new(),
|
|
}
|
|
}
|
|
}
|
|
|
|
impl io::Read for BlockingChannelReader {
|
|
fn read(&mut self, buffer: &mut [u8]) -> io::Result<usize> {
|
|
while self.current.is_empty() {
|
|
match self.receiver.blocking_recv() {
|
|
Some(Ok(bytes)) => self.current = bytes,
|
|
Some(Err(error)) => return Err(error),
|
|
None => return Ok(0),
|
|
}
|
|
}
|
|
let read = buffer.len().min(self.current.len());
|
|
buffer[..read].copy_from_slice(&self.current[..read]);
|
|
self.current.advance(read);
|
|
Ok(read)
|
|
}
|
|
}
|
|
|
|
pub(crate) fn compressed_input_stream(
|
|
reader: SelectInputReader,
|
|
compressed_size: u64,
|
|
format: CompressionFormat,
|
|
input_metrics: Arc<SelectInputMetrics>,
|
|
record_delimiter: Vec<u8>,
|
|
max_processed_bytes: u64,
|
|
query_guard: Option<QueryExecutionGuard>,
|
|
) -> ObjectStoreResult<BoxStream<'static, ObjectStoreResult<Bytes>>> {
|
|
let record_size = RecordSizeTracker::new(record_delimiter).map_err(select_object_store_error)?;
|
|
let reader = compressed_input_reader(reader, compressed_size, format, input_metrics, max_processed_bytes, query_guard);
|
|
let stream = ReaderStream::with_capacity(reader, SELECT_DECODE_CHUNK_BYTES);
|
|
Ok(stream::try_unfold((stream, record_size), |(mut stream, mut record_size)| async move {
|
|
match stream.next().await {
|
|
Some(Ok(bytes)) => {
|
|
record_size.observe(&bytes).map_err(select_object_store_error)?;
|
|
Ok(Some((bytes, (stream, record_size))))
|
|
}
|
|
Some(Err(error)) => Err(input_io_error(error)),
|
|
None => {
|
|
record_size.finish().map_err(select_object_store_error)?;
|
|
Ok(None)
|
|
}
|
|
}
|
|
})
|
|
.boxed())
|
|
}
|
|
|
|
pub(crate) fn input_io_error(source: io::Error) -> ObjectStoreError {
|
|
let source: Box<dyn StdError + Send + Sync> = match find_error_source::<SelectError>(&source) {
|
|
Some(error) => Box::new(error.clone()),
|
|
None => Box::new(source),
|
|
};
|
|
ObjectStoreError::Generic {
|
|
store: "EcObjectStore",
|
|
source,
|
|
}
|
|
}
|
|
|
|
fn select_object_store_error(source: SelectError) -> ObjectStoreError {
|
|
ObjectStoreError::Generic {
|
|
store: "EcObjectStore",
|
|
source: Box::new(source),
|
|
}
|
|
}
|
|
|
|
#[derive(Debug, thiserror::Error)]
|
|
#[error("compressed object source read failed")]
|
|
struct CompressedSourceReadError {
|
|
#[source]
|
|
source: io::Error,
|
|
}
|
|
|
|
struct ScannedReader<R> {
|
|
inner: tokio::io::Take<R>,
|
|
input_metrics: SelectInputMetricsRecorder,
|
|
}
|
|
|
|
impl<R: AsyncRead + Unpin> ScannedReader<R> {
|
|
fn new(reader: R, compressed_size: u64, input_metrics: SelectInputMetricsRecorder) -> Self {
|
|
Self {
|
|
inner: reader.take(compressed_size),
|
|
input_metrics,
|
|
}
|
|
}
|
|
}
|
|
|
|
impl<R: AsyncRead + Unpin> AsyncRead for ScannedReader<R> {
|
|
fn poll_read(mut self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll<io::Result<()>> {
|
|
let before = buf.filled().len();
|
|
match Pin::new(&mut self.inner).poll_read(cx, buf) {
|
|
Poll::Pending => Poll::Pending,
|
|
Poll::Ready(Err(source)) => Poll::Ready(Err(source_read_error(source))),
|
|
Poll::Ready(Ok(())) => {
|
|
let read = buf.filled().len() - before;
|
|
if read == 0 && self.inner.limit() > 0 {
|
|
let source = io::Error::new(
|
|
io::ErrorKind::UnexpectedEof,
|
|
format!("compressed object stream ended with {} bytes remaining", self.inner.limit()),
|
|
);
|
|
return Poll::Ready(Err(source_read_error(source)));
|
|
}
|
|
self.input_metrics.record_scanned(read);
|
|
Poll::Ready(Ok(()))
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
struct CooperativeReader<R> {
|
|
inner: R,
|
|
bytes_since_yield: usize,
|
|
yield_pending: bool,
|
|
}
|
|
|
|
impl<R> CooperativeReader<R> {
|
|
fn new(inner: R) -> Self {
|
|
Self {
|
|
inner,
|
|
bytes_since_yield: 0,
|
|
yield_pending: false,
|
|
}
|
|
}
|
|
}
|
|
|
|
impl<R: AsyncRead + Unpin> AsyncRead for CooperativeReader<R> {
|
|
fn poll_read(mut self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll<io::Result<()>> {
|
|
if self.yield_pending {
|
|
self.yield_pending = false;
|
|
self.bytes_since_yield = 0;
|
|
cx.waker().wake_by_ref();
|
|
return Poll::Pending;
|
|
}
|
|
if buf.remaining() == 0 {
|
|
return Poll::Ready(Ok(()));
|
|
}
|
|
|
|
let remaining_budget = SELECT_DECODE_CHUNK_BYTES - self.bytes_since_yield;
|
|
let read_limit = remaining_budget.min(buf.remaining());
|
|
let read = {
|
|
let unfilled = buf.initialize_unfilled_to(read_limit);
|
|
let mut limited = ReadBuf::new(unfilled);
|
|
match Pin::new(&mut self.inner).poll_read(cx, &mut limited) {
|
|
Poll::Pending => return Poll::Pending,
|
|
Poll::Ready(Err(error)) => return Poll::Ready(Err(error)),
|
|
Poll::Ready(Ok(())) => limited.filled().len(),
|
|
}
|
|
};
|
|
buf.advance(read);
|
|
self.bytes_since_yield += read;
|
|
self.yield_pending = read > 0 && self.bytes_since_yield == SELECT_DECODE_CHUNK_BYTES;
|
|
Poll::Ready(Ok(()))
|
|
}
|
|
}
|
|
|
|
fn source_read_error(source: io::Error) -> io::Error {
|
|
let kind = source.kind();
|
|
io::Error::new(kind, CompressedSourceReadError { source })
|
|
}
|
|
|
|
struct Bzip2HeaderValidatingReader<R> {
|
|
inner: R,
|
|
position: usize,
|
|
pending_error: Option<SelectError>,
|
|
}
|
|
|
|
impl<R> Bzip2HeaderValidatingReader<R> {
|
|
fn new(inner: R) -> Self {
|
|
Self {
|
|
inner,
|
|
position: 0,
|
|
pending_error: None,
|
|
}
|
|
}
|
|
}
|
|
|
|
impl<R: AsyncRead + Unpin> AsyncRead for Bzip2HeaderValidatingReader<R> {
|
|
fn poll_read(mut self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll<io::Result<()>> {
|
|
if let Some(error) = self.pending_error.take() {
|
|
return Poll::Ready(Err(io::Error::new(io::ErrorKind::InvalidData, error)));
|
|
}
|
|
if buf.remaining() == 0 {
|
|
return Poll::Ready(Ok(()));
|
|
}
|
|
if self.position == 4 {
|
|
return Pin::new(&mut self.inner).poll_read(cx, buf);
|
|
}
|
|
|
|
let before = buf.filled().len();
|
|
match Pin::new(&mut self.inner).poll_read(cx, buf) {
|
|
Poll::Pending => Poll::Pending,
|
|
Poll::Ready(Err(error)) => Poll::Ready(Err(error)),
|
|
Poll::Ready(Ok(())) => {
|
|
let after = buf.filled().len();
|
|
if after == before {
|
|
return Poll::Ready(Err(io::Error::new(io::ErrorKind::UnexpectedEof, SelectError::TruncatedInput)));
|
|
}
|
|
if let Err(error) = validate_bzip2_header(&mut self.position, &buf.filled()[before..after]) {
|
|
buf.set_filled(before + error.offset);
|
|
if error.offset == 0 {
|
|
return Poll::Ready(Err(io::Error::new(io::ErrorKind::InvalidData, error.source)));
|
|
}
|
|
self.pending_error = Some(error.source);
|
|
}
|
|
Poll::Ready(Ok(()))
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
struct HeaderValidationError {
|
|
offset: usize,
|
|
source: SelectError,
|
|
}
|
|
|
|
fn validate_bzip2_header(position: &mut usize, bytes: &[u8]) -> Result<(), HeaderValidationError> {
|
|
for (offset, byte) in bytes.iter().copied().enumerate() {
|
|
let valid = match *position {
|
|
0 => byte == b'B',
|
|
1 => byte == b'Z',
|
|
2 => byte == b'h',
|
|
3 => matches!(byte, b'1'..=b'9'),
|
|
_ => break,
|
|
};
|
|
if !valid {
|
|
return Err(HeaderValidationError {
|
|
offset,
|
|
source: CompressionFormat::Bzip2.invalid_header_error(),
|
|
});
|
|
}
|
|
*position += 1;
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
const GZIP_FLAG_HEADER_CRC: u8 = 0x02;
|
|
const GZIP_FLAG_EXTRA: u8 = 0x04;
|
|
const GZIP_FLAG_NAME: u8 = 0x08;
|
|
const GZIP_FLAG_COMMENT: u8 = 0x10;
|
|
const GZIP_FLAG_TEXT: u8 = 0x01;
|
|
|
|
struct ProcessedReader<R> {
|
|
inner: R,
|
|
input_metrics: SelectInputMetricsRecorder,
|
|
processed_bytes: u64,
|
|
max_processed_bytes: u64,
|
|
}
|
|
|
|
impl<R> ProcessedReader<R> {
|
|
fn new(inner: R, input_metrics: SelectInputMetricsRecorder, max_processed_bytes: u64) -> Self {
|
|
Self {
|
|
inner,
|
|
input_metrics,
|
|
processed_bytes: 0,
|
|
max_processed_bytes,
|
|
}
|
|
}
|
|
}
|
|
|
|
impl<R: AsyncRead + Unpin> AsyncRead for ProcessedReader<R> {
|
|
fn poll_read(mut self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll<io::Result<()>> {
|
|
let before = buf.filled().len();
|
|
match Pin::new(&mut self.inner).poll_read(cx, buf) {
|
|
Poll::Pending => Poll::Pending,
|
|
Poll::Ready(Err(error)) => Poll::Ready(Err(classify_decoder_error(error))),
|
|
Poll::Ready(Ok(())) => {
|
|
let read = buf.filled().len() - before;
|
|
let read = u64::try_from(read).unwrap_or(u64::MAX);
|
|
let Some(processed_bytes) = self.processed_bytes.checked_add(read) else {
|
|
return Poll::Ready(Err(processed_bytes_limit_error()));
|
|
};
|
|
if processed_bytes > self.max_processed_bytes {
|
|
return Poll::Ready(Err(processed_bytes_limit_error()));
|
|
}
|
|
self.processed_bytes = processed_bytes;
|
|
self.input_metrics.record_processed(buf.filled().len() - before);
|
|
Poll::Ready(Ok(()))
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
fn processed_bytes_limit_error() -> io::Error {
|
|
io::Error::new(io::ErrorKind::OutOfMemory, SelectError::ResourceExhausted)
|
|
}
|
|
|
|
fn classify_decoder_error(error: io::Error) -> io::Error {
|
|
if error_chain_contains::<CompressedSourceReadError>(&error) || error_chain_contains::<SelectError>(&error) {
|
|
return error;
|
|
}
|
|
|
|
let select_error = if error.kind() == io::ErrorKind::OutOfMemory {
|
|
SelectError::ResourceExhausted
|
|
} else {
|
|
SelectError::TruncatedInput
|
|
};
|
|
io::Error::new(error.kind(), select_error)
|
|
}
|
|
|
|
fn error_chain_contains<T: StdError + 'static>(error: &(dyn StdError + 'static)) -> bool {
|
|
find_error_source::<T>(error).is_some()
|
|
}
|
|
|
|
fn find_error_source<'a, T: StdError + 'static>(error: &'a (dyn StdError + 'static)) -> Option<&'a T> {
|
|
let mut current = Some(error);
|
|
for _ in 0..MAX_ERROR_SOURCE_DEPTH {
|
|
let Some(error) = current else {
|
|
break;
|
|
};
|
|
if let Some(error) = error.downcast_ref::<T>() {
|
|
return Some(error);
|
|
}
|
|
current = error
|
|
.downcast_ref::<io::Error>()
|
|
.and_then(|error| error.get_ref())
|
|
.map(|source| source as &(dyn StdError + 'static))
|
|
.or_else(|| error.source());
|
|
}
|
|
None
|
|
}
|
|
|
|
struct RecordSizeTracker {
|
|
delimiter: Vec<u8>,
|
|
prefix: Vec<usize>,
|
|
record_bytes: usize,
|
|
matched: usize,
|
|
}
|
|
|
|
impl RecordSizeTracker {
|
|
fn new(delimiter: Vec<u8>) -> Result<Self, SelectError> {
|
|
if delimiter.is_empty() {
|
|
return Err(SelectError::InvalidDataSource);
|
|
}
|
|
|
|
let mut prefix = vec![0; delimiter.len()];
|
|
let mut matched = 0;
|
|
for index in 1..delimiter.len() {
|
|
while matched > 0 && delimiter[index] != delimiter[matched] {
|
|
matched = prefix[matched - 1];
|
|
}
|
|
if delimiter[index] == delimiter[matched] {
|
|
matched += 1;
|
|
}
|
|
prefix[index] = matched;
|
|
}
|
|
|
|
Ok(Self {
|
|
delimiter,
|
|
prefix,
|
|
record_bytes: 0,
|
|
matched: 0,
|
|
})
|
|
}
|
|
|
|
fn observe(&mut self, bytes: &[u8]) -> Result<(), SelectError> {
|
|
if self.delimiter.len() == 1 {
|
|
return self.observe_single_byte_delimiter(bytes);
|
|
}
|
|
|
|
for &byte in bytes {
|
|
self.record_bytes = self.record_bytes.checked_add(1).ok_or(SelectError::OverMaxRecordSize)?;
|
|
while self.matched > 0 && byte != self.delimiter[self.matched] {
|
|
self.matched = self.prefix[self.matched - 1];
|
|
}
|
|
if byte == self.delimiter[self.matched] {
|
|
self.matched += 1;
|
|
}
|
|
if self.matched == self.delimiter.len() {
|
|
let payload_bytes = self.record_bytes - self.delimiter.len();
|
|
if payload_bytes > MAX_SELECT_RECORD_BYTES {
|
|
return Err(SelectError::OverMaxRecordSize);
|
|
}
|
|
self.record_bytes = 0;
|
|
self.matched = 0;
|
|
} else if self.record_bytes - self.matched > MAX_SELECT_RECORD_BYTES {
|
|
return Err(SelectError::OverMaxRecordSize);
|
|
}
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
fn observe_single_byte_delimiter(&mut self, mut bytes: &[u8]) -> Result<(), SelectError> {
|
|
let delimiter = self.delimiter[0];
|
|
while let Some(index) = bytes.iter().position(|byte| *byte == delimiter) {
|
|
let payload_bytes = self.record_bytes.checked_add(index).ok_or(SelectError::OverMaxRecordSize)?;
|
|
if payload_bytes > MAX_SELECT_RECORD_BYTES {
|
|
return Err(SelectError::OverMaxRecordSize);
|
|
}
|
|
self.record_bytes = 0;
|
|
bytes = &bytes[index + 1..];
|
|
}
|
|
self.record_bytes = self
|
|
.record_bytes
|
|
.checked_add(bytes.len())
|
|
.ok_or(SelectError::OverMaxRecordSize)?;
|
|
if self.record_bytes > MAX_SELECT_RECORD_BYTES {
|
|
return Err(SelectError::OverMaxRecordSize);
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
fn finish(&self) -> Result<(), SelectError> {
|
|
if self.record_bytes > MAX_SELECT_RECORD_BYTES {
|
|
Err(SelectError::OverMaxRecordSize)
|
|
} else {
|
|
Ok(())
|
|
}
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
pub(crate) async fn encode_compressed_fixture(format: CompressionFormat, input: &[u8]) -> Vec<u8> {
|
|
use tokio::io::AsyncWriteExt as _;
|
|
|
|
let cursor = std::io::Cursor::new(Vec::new());
|
|
match format {
|
|
CompressionFormat::Gzip => {
|
|
let mut encoder = async_compression::tokio::write::GzipEncoder::new(cursor);
|
|
encoder.write_all(input).await.expect("gzip fixture should encode");
|
|
encoder.shutdown().await.expect("gzip fixture should finish");
|
|
encoder.into_inner().into_inner()
|
|
}
|
|
CompressionFormat::Bzip2 => {
|
|
let mut encoder = async_compression::tokio::write::BzEncoder::new(cursor);
|
|
encoder.write_all(input).await.expect("bzip2 fixture should encode");
|
|
encoder.shutdown().await.expect("bzip2 fixture should finish");
|
|
encoder.into_inner().into_inner()
|
|
}
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use crate::SelectInputMetricsSnapshot;
|
|
use futures::TryStreamExt;
|
|
use std::{io::Cursor, sync::Mutex as StdMutex, thread::ThreadId};
|
|
use tokio::io::{AsyncWriteExt, DuplexStream};
|
|
|
|
struct ThreadRecordingReader {
|
|
inner: Cursor<Vec<u8>>,
|
|
thread_id: Arc<StdMutex<Option<ThreadId>>>,
|
|
}
|
|
|
|
impl AsyncRead for ThreadRecordingReader {
|
|
fn poll_read(self: Pin<&mut Self>, cx: &mut Context<'_>, buffer: &mut ReadBuf<'_>) -> Poll<io::Result<()>> {
|
|
let this = self.get_mut();
|
|
let mut thread_id = this.thread_id.lock().expect("thread recorder mutex should not be poisoned");
|
|
thread_id.get_or_insert_with(|| std::thread::current().id());
|
|
drop(thread_id);
|
|
Pin::new(&mut this.inner).poll_read(cx, buffer)
|
|
}
|
|
}
|
|
|
|
struct ErrorAfterReader {
|
|
inner: Cursor<Vec<u8>>,
|
|
end: u64,
|
|
failed: bool,
|
|
}
|
|
|
|
impl AsyncRead for ErrorAfterReader {
|
|
fn poll_read(mut self: Pin<&mut Self>, cx: &mut Context<'_>, buffer: &mut ReadBuf<'_>) -> Poll<io::Result<()>> {
|
|
if self.inner.position() < self.end {
|
|
return Pin::new(&mut self.inner).poll_read(cx, buffer);
|
|
}
|
|
if !self.failed {
|
|
self.failed = true;
|
|
return Poll::Ready(Err(io::Error::new(io::ErrorKind::ConnectionReset, "injected source failure")));
|
|
}
|
|
Poll::Ready(Ok(()))
|
|
}
|
|
}
|
|
|
|
async fn decode(format: CompressionFormat, compressed: Vec<u8>) -> (ObjectStoreResult<Vec<u8>>, SelectInputMetricsSnapshot) {
|
|
let compressed_len = u64::try_from(compressed.len()).expect("fixture length should fit in u64");
|
|
let metrics = Arc::new(SelectInputMetrics::default());
|
|
let stream = compressed_input_stream(
|
|
Box::new(Cursor::new(compressed)),
|
|
compressed_len,
|
|
format,
|
|
Arc::clone(&metrics),
|
|
b"\n".to_vec(),
|
|
u64::MAX,
|
|
None,
|
|
)
|
|
.expect("record delimiter should be valid");
|
|
let result = stream.try_collect::<Vec<_>>().await.map(|chunks| chunks.concat());
|
|
(result, metrics.snapshot())
|
|
}
|
|
|
|
fn select_error(error: &ObjectStoreError) -> Option<SelectError> {
|
|
find_error_source::<SelectError>(error).cloned()
|
|
}
|
|
|
|
#[test]
|
|
fn protocol_limits_match_the_s3_select_contract() {
|
|
assert_eq!(MAX_SELECT_RECORD_BYTES, 1_048_576);
|
|
assert_eq!(MAX_SELECT_PROCESSED_BYTES, 5_497_558_138_880);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn gzip_and_bzip2_preserve_bytes_and_metric_boundaries() {
|
|
const INPUT: &[u8] = b"name,age\nAlice,30\n";
|
|
for format in [CompressionFormat::Gzip, CompressionFormat::Bzip2] {
|
|
let compressed = encode_compressed_fixture(format, INPUT).await;
|
|
let compressed_len = u64::try_from(compressed.len()).expect("fixture length should fit in u64");
|
|
let (decoded, metrics) = decode(format, compressed).await;
|
|
|
|
assert_eq!(decoded.expect("valid compressed input should decode"), INPUT);
|
|
assert_eq!(metrics.bytes_scanned, compressed_len);
|
|
assert_eq!(
|
|
metrics.bytes_processed,
|
|
u64::try_from(INPUT.len()).expect("fixture length should fit in u64")
|
|
);
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn source_read_errors_are_not_reclassified_as_truncated_input() {
|
|
for format in [CompressionFormat::Gzip, CompressionFormat::Bzip2] {
|
|
let compressed = encode_compressed_fixture(format, b"name\nAlice\n").await;
|
|
let compressed_len = u64::try_from(compressed.len()).expect("fixture length should fit in u64");
|
|
let expected_size = u64::try_from(compressed.len() + 1).expect("fixture length should fit in u64");
|
|
let stream = compressed_input_stream(
|
|
Box::new(ErrorAfterReader {
|
|
inner: Cursor::new(compressed),
|
|
end: compressed_len,
|
|
failed: false,
|
|
}),
|
|
expected_size,
|
|
format,
|
|
Arc::new(SelectInputMetrics::default()),
|
|
b"\n".to_vec(),
|
|
u64::MAX,
|
|
None,
|
|
)
|
|
.expect("record delimiter should be valid");
|
|
|
|
let error = stream
|
|
.try_collect::<Vec<_>>()
|
|
.await
|
|
.expect_err("a storage read failure must terminate decoding");
|
|
|
|
assert_eq!(select_error(&error), None, "{format:?} must not report TruncatedInput: {error:?}");
|
|
assert!(
|
|
find_error_source::<CompressedSourceReadError>(&error).is_some(),
|
|
"{format:?} must preserve the source error class"
|
|
);
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn source_read_error_inside_gzip_header_is_not_truncated_input() {
|
|
let compressed = encode_compressed_fixture(CompressionFormat::Gzip, b"name\nAlice\n").await;
|
|
let partial_header = compressed[..5].to_vec();
|
|
let stream = compressed_input_stream(
|
|
Box::new(Cursor::new(partial_header)),
|
|
10,
|
|
CompressionFormat::Gzip,
|
|
Arc::new(SelectInputMetrics::default()),
|
|
b"\n".to_vec(),
|
|
u64::MAX,
|
|
None,
|
|
)
|
|
.expect("record delimiter should be valid");
|
|
|
|
let error = stream
|
|
.try_collect::<Vec<_>>()
|
|
.await
|
|
.expect_err("a source failure inside the GZIP header must terminate decoding");
|
|
|
|
assert_eq!(select_error(&error), None, "source failure must not report TruncatedInput: {error:?}");
|
|
assert!(
|
|
find_error_source::<CompressedSourceReadError>(&error).is_some(),
|
|
"the source error class must survive GZIP header validation"
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn concatenated_members_decode_in_order() {
|
|
for format in [CompressionFormat::Gzip, CompressionFormat::Bzip2] {
|
|
let mut compressed = encode_compressed_fixture(format, b"a\n").await;
|
|
compressed.extend_from_slice(&encode_compressed_fixture(format, b"b\n").await);
|
|
let (decoded, _) = decode(format, compressed).await;
|
|
assert_eq!(decoded.expect("valid concatenated members should decode"), b"a\nb\n");
|
|
}
|
|
|
|
let mut compressed = encode_compressed_fixture(CompressionFormat::Gzip, b"a\n").await;
|
|
let (second, _) = gzip_with_optional_header_crc(b"b\n").await;
|
|
compressed.extend_from_slice(&second);
|
|
let (decoded, _) = decode(CompressionFormat::Gzip, compressed).await;
|
|
assert_eq!(decoded.expect("optional headers must work in later GZIP members"), b"a\nb\n");
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn invalid_initial_headers_are_typed_as_invalid_compression() {
|
|
for format in [CompressionFormat::Gzip, CompressionFormat::Bzip2] {
|
|
let (result, _) = decode(format, b"not compressed\n".to_vec()).await;
|
|
let error = result.expect_err("invalid compression header must fail");
|
|
let compression = format.name();
|
|
assert_eq!(select_error(&error), Some(SelectError::InvalidCompressionFormatForObject { compression }));
|
|
assert_eq!(
|
|
select_error(&error).expect("typed compression error").to_string(),
|
|
format!("{compression} is not applicable to the queried object. Please correct the request and try again.")
|
|
);
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn fixed_header_variants_are_validated_before_decoding() {
|
|
let encoded = encode_compressed_fixture(CompressionFormat::Gzip, b"name\nAlice\n").await;
|
|
for (position, invalid) in [(2, 0), (3, 0x20)] {
|
|
let mut mutated = encoded.clone();
|
|
mutated[position] = invalid;
|
|
let (result, _) = decode(CompressionFormat::Gzip, mutated).await;
|
|
assert_eq!(
|
|
select_error(&result.expect_err("invalid GZIP fixed header must fail")),
|
|
Some(CompressionFormat::Gzip.invalid_header_error())
|
|
);
|
|
}
|
|
|
|
let (result, _) = decode(CompressionFormat::Bzip2, b"BZh0".to_vec()).await;
|
|
assert_eq!(
|
|
select_error(&result.expect_err("invalid BZIP2 block size must fail")),
|
|
Some(CompressionFormat::Bzip2.invalid_header_error())
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn incomplete_valid_headers_are_typed_as_truncated_input() {
|
|
for (format, prefixes) in [
|
|
(CompressionFormat::Gzip, vec![b"".as_slice(), b"\x1f", b"\x1f\x8b", b"\x1f\x8b\x08"]),
|
|
(CompressionFormat::Bzip2, vec![b"".as_slice(), b"B", b"BZ", b"BZh"]),
|
|
] {
|
|
for prefix in prefixes {
|
|
let (result, _) = decode(format, prefix.to_vec()).await;
|
|
assert_eq!(
|
|
select_error(&result.expect_err("incomplete compression header must fail")),
|
|
Some(SelectError::TruncatedInput)
|
|
);
|
|
}
|
|
}
|
|
}
|
|
|
|
async fn gzip_with_optional_header_crc(input: &[u8]) -> (Vec<u8>, usize) {
|
|
const EXTRA: &[u8] = b"s3-select";
|
|
const FILE_NAME: &[u8] = b"select.csv";
|
|
const COMMENT: &[u8] = b"fixture";
|
|
|
|
let encoded = encode_compressed_fixture(CompressionFormat::Gzip, input).await;
|
|
let mut header = encoded[..10].to_vec();
|
|
header[3] = GZIP_FLAG_EXTRA | GZIP_FLAG_NAME | GZIP_FLAG_COMMENT | GZIP_FLAG_HEADER_CRC;
|
|
header.extend_from_slice(
|
|
&u16::try_from(EXTRA.len())
|
|
.expect("GZIP extra fixture should fit its length field")
|
|
.to_le_bytes(),
|
|
);
|
|
header.extend_from_slice(EXTRA);
|
|
header.extend_from_slice(FILE_NAME);
|
|
header.push(0);
|
|
header.extend_from_slice(COMMENT);
|
|
header.push(0);
|
|
let mut digest = crc_fast::Digest::new(crc_fast::CrcAlgorithm::Crc32IsoHdlc);
|
|
digest.update(&header);
|
|
let crc = digest.finalize();
|
|
let crc = u16::try_from(crc & u64::from(u16::MAX)).expect("masked header CRC should fit in u16");
|
|
let crc_offset = header.len();
|
|
header.extend_from_slice(&crc.to_le_bytes());
|
|
header.extend_from_slice(&encoded[10..]);
|
|
(header, crc_offset)
|
|
}
|
|
|
|
async fn gzip_with_text_header(input: &[u8], flag: u8, field_bytes: usize) -> Vec<u8> {
|
|
let encoded = encode_compressed_fixture(CompressionFormat::Gzip, input).await;
|
|
let mut member = encoded[..10].to_vec();
|
|
member[3] = flag;
|
|
member.extend(std::iter::repeat_n(b'x', field_bytes));
|
|
member.push(0);
|
|
member.extend_from_slice(&encoded[10..]);
|
|
member
|
|
}
|
|
|
|
async fn gzip_with_text_header_crc(input: &[u8], flag: u8, field_bytes: usize) -> (Vec<u8>, usize) {
|
|
let encoded = encode_compressed_fixture(CompressionFormat::Gzip, input).await;
|
|
let mut member = encoded[..10].to_vec();
|
|
member[3] = flag | GZIP_FLAG_HEADER_CRC;
|
|
member.extend(std::iter::repeat_n(b'x', field_bytes));
|
|
member.push(0);
|
|
let mut digest = crc_fast::Digest::new(crc_fast::CrcAlgorithm::Crc32IsoHdlc);
|
|
digest.update(&member);
|
|
let crc = u16::try_from(digest.finalize() & u64::from(u16::MAX)).expect("masked header CRC should fit in u16");
|
|
let crc_offset = member.len();
|
|
member.extend_from_slice(&crc.to_le_bytes());
|
|
member.extend_from_slice(&encoded[10..]);
|
|
(member, crc_offset)
|
|
}
|
|
|
|
async fn gzip_with_extra_header_crc(input: &[u8], extra_bytes: usize) -> (Vec<u8>, usize) {
|
|
let encoded = encode_compressed_fixture(CompressionFormat::Gzip, input).await;
|
|
let mut member = encoded[..10].to_vec();
|
|
member[3] = GZIP_FLAG_EXTRA | GZIP_FLAG_HEADER_CRC;
|
|
member.extend_from_slice(
|
|
&u16::try_from(extra_bytes)
|
|
.expect("GZIP extra fixture should fit its length field")
|
|
.to_le_bytes(),
|
|
);
|
|
member.extend(std::iter::repeat_n(b'x', extra_bytes));
|
|
let mut digest = crc_fast::Digest::new(crc_fast::CrcAlgorithm::Crc32IsoHdlc);
|
|
digest.update(&member);
|
|
let crc = u16::try_from(digest.finalize() & u64::from(u16::MAX)).expect("masked header CRC should fit in u16");
|
|
let crc_offset = member.len();
|
|
member.extend_from_slice(&crc.to_le_bytes());
|
|
member.extend_from_slice(&encoded[10..]);
|
|
(member, crc_offset)
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn gzip_optional_header_crc_is_validated_without_error_strings() {
|
|
const INPUT: &[u8] = b"name\nAlice\n";
|
|
|
|
let (encoded, crc_offset) = gzip_with_optional_header_crc(INPUT).await;
|
|
let (decoded, _) = decode(CompressionFormat::Gzip, encoded.clone()).await;
|
|
assert_eq!(decoded.expect("valid optional GZIP header should decode"), INPUT);
|
|
|
|
let mut corrupt = encoded;
|
|
corrupt[crc_offset] ^= 0xff;
|
|
let (result, _) = decode(CompressionFormat::Gzip, corrupt).await;
|
|
assert_eq!(
|
|
select_error(&result.expect_err("invalid GZIP header CRC must fail")),
|
|
Some(SelectError::InvalidCompressionFormatForObject {
|
|
compression: CompressionFormat::Gzip.name(),
|
|
})
|
|
);
|
|
|
|
let mut truncated = encode_compressed_fixture(CompressionFormat::Gzip, INPUT).await;
|
|
truncated.truncate(10);
|
|
truncated[3] = GZIP_FLAG_NAME;
|
|
truncated.extend_from_slice(b"unterminated-name");
|
|
let (result, _) = decode(CompressionFormat::Gzip, truncated).await;
|
|
assert_eq!(
|
|
select_error(&result.expect_err("incomplete optional GZIP header must fail")),
|
|
Some(SelectError::TruncatedInput)
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn gzip_optional_header_can_cross_source_chunks() {
|
|
const INPUT: &[u8] = b"name\nAlice\n";
|
|
const LARGE_GZIP_TEXT_FIELD_BYTES: usize = SELECT_DECODE_CHUNK_BYTES * 3 + 17;
|
|
|
|
let encoded = encode_compressed_fixture(CompressionFormat::Gzip, INPUT).await;
|
|
let mut with_name = encoded[..10].to_vec();
|
|
with_name[3] = GZIP_FLAG_NAME;
|
|
with_name.extend(std::iter::repeat_n(b'x', LARGE_GZIP_TEXT_FIELD_BYTES));
|
|
with_name.push(0);
|
|
with_name.extend_from_slice(&encoded[10..]);
|
|
let expected_scanned = u64::try_from(with_name.len()).expect("fixture length should fit in u64");
|
|
let (mut writer, reader) = tokio::io::duplex(257);
|
|
let writer = tokio::spawn(async move {
|
|
writer
|
|
.write_all(&with_name)
|
|
.await
|
|
.expect("fixture source should accept bytes");
|
|
writer.shutdown().await.expect("fixture source should close");
|
|
});
|
|
let metrics = Arc::new(SelectInputMetrics::default());
|
|
let stream = compressed_input_stream(
|
|
Box::new(reader),
|
|
expected_scanned,
|
|
CompressionFormat::Gzip,
|
|
Arc::clone(&metrics),
|
|
b"\n".to_vec(),
|
|
u64::MAX,
|
|
None,
|
|
)
|
|
.expect("record delimiter should be valid");
|
|
let decoded = stream
|
|
.try_collect::<Vec<_>>()
|
|
.await
|
|
.map(|chunks| chunks.concat())
|
|
.expect("chunked optional GZIP names should decode");
|
|
writer.await.expect("fixture writer should complete");
|
|
|
|
assert_eq!(decoded, INPUT);
|
|
assert_eq!(metrics.snapshot().bytes_scanned, expected_scanned);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn valid_512_byte_text_fields_decode_in_every_gzip_member() {
|
|
for flag in [GZIP_FLAG_NAME, GZIP_FLAG_COMMENT] {
|
|
let mut encoded = encode_compressed_fixture(CompressionFormat::Gzip, b"a\n").await;
|
|
encoded.extend_from_slice(&gzip_with_text_header(b"b\n", flag, 512).await);
|
|
let (decoded, _) = decode(CompressionFormat::Gzip, encoded).await;
|
|
assert_eq!(decoded.expect("RFC 1952 does not limit zero-terminated text fields"), b"a\nb\n");
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn long_text_field_header_crc_accumulates_across_source_chunks() {
|
|
const FIELD_BYTES: usize = SELECT_DECODE_CHUNK_BYTES * 3 + 17;
|
|
|
|
for flag in [GZIP_FLAG_NAME, GZIP_FLAG_COMMENT] {
|
|
let first = encode_compressed_fixture(CompressionFormat::Gzip, b"a\n").await;
|
|
let (second, crc_offset) = gzip_with_text_header_crc(b"b\n", flag, FIELD_BYTES).await;
|
|
|
|
let mut valid = first.clone();
|
|
valid.extend_from_slice(&second);
|
|
let (decoded, _) = decode(CompressionFormat::Gzip, valid).await;
|
|
assert_eq!(decoded.expect("multi-chunk GZIP header CRC should validate"), b"a\nb\n");
|
|
|
|
let mut corrupt_second = second;
|
|
corrupt_second[crc_offset] ^= 0xff;
|
|
let mut corrupt = first;
|
|
corrupt.extend_from_slice(&corrupt_second);
|
|
let (result, _) = decode(CompressionFormat::Gzip, corrupt).await;
|
|
assert_eq!(
|
|
select_error(&result.expect_err("corrupt multi-chunk GZIP header CRC must fail")),
|
|
Some(SelectError::TruncatedInput)
|
|
);
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn maximum_extra_field_header_crc_accumulates_across_source_chunks() {
|
|
const INPUT: &[u8] = b"name\nAlice\n";
|
|
let (valid, crc_offset) = gzip_with_extra_header_crc(INPUT, usize::from(u16::MAX)).await;
|
|
|
|
let (decoded, _) = decode(CompressionFormat::Gzip, valid.clone()).await;
|
|
assert_eq!(decoded.expect("maximum GZIP extra field should decode"), INPUT);
|
|
|
|
let mut corrupt = valid;
|
|
corrupt[crc_offset] ^= 0xff;
|
|
let (result, _) = decode(CompressionFormat::Gzip, corrupt).await;
|
|
assert_eq!(
|
|
select_error(&result.expect_err("corrupt multi-chunk GZIP extra-field CRC must fail")),
|
|
Some(CompressionFormat::Gzip.invalid_header_error())
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn processed_byte_limit_is_independent_of_live_memory_budget() {
|
|
assert_eq!(processed_bytes_limit(), MAX_SELECT_PROCESSED_BYTES);
|
|
assert!(processed_bytes_limit() > 80 * 1024 * 1024);
|
|
}
|
|
|
|
#[test]
|
|
fn cooperative_reader_yields_after_bounded_ready_input() {
|
|
let source = Cursor::new(vec![0_u8; SELECT_DECODE_CHUNK_BYTES + 1]);
|
|
let mut reader = CooperativeReader::new(source);
|
|
let mut output = vec![0_u8; SELECT_DECODE_CHUNK_BYTES + 1];
|
|
let mut read_buf = ReadBuf::new(&mut output);
|
|
let waker = futures::task::noop_waker_ref();
|
|
let mut context = Context::from_waker(waker);
|
|
|
|
assert!(Pin::new(&mut reader).poll_read(&mut context, &mut read_buf).is_ready());
|
|
assert_eq!(read_buf.filled().len(), SELECT_DECODE_CHUNK_BYTES);
|
|
assert!(Pin::new(&mut reader).poll_read(&mut context, &mut read_buf).is_pending());
|
|
assert!(Pin::new(&mut reader).poll_read(&mut context, &mut read_buf).is_ready());
|
|
assert_eq!(read_buf.filled().len(), SELECT_DECODE_CHUNK_BYTES + 1);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn bzip2_decoded_output_yields_at_the_cooperative_quantum() {
|
|
let input = vec![b'x'; SELECT_DECODE_CHUNK_BYTES * 2];
|
|
let compressed = encode_compressed_fixture(CompressionFormat::Bzip2, &input).await;
|
|
let compressed_len = u64::try_from(compressed.len()).expect("fixture length should fit in u64");
|
|
let mut reader = compressed_input_reader(
|
|
Box::new(Cursor::new(compressed)),
|
|
compressed_len,
|
|
CompressionFormat::Bzip2,
|
|
Arc::new(SelectInputMetrics::default()),
|
|
processed_bytes_limit(),
|
|
None,
|
|
);
|
|
let mut output = vec![0; SELECT_DECODE_CHUNK_BYTES * 2];
|
|
let mut read_buf = ReadBuf::new(&mut output);
|
|
while read_buf.filled().len() < SELECT_DECODE_CHUNK_BYTES {
|
|
futures::future::poll_fn(|cx| Pin::new(&mut reader).poll_read(cx, &mut read_buf))
|
|
.await
|
|
.expect("valid BZIP2 input should decode");
|
|
}
|
|
assert_eq!(read_buf.filled().len(), SELECT_DECODE_CHUNK_BYTES);
|
|
|
|
let waker = futures::task::noop_waker_ref();
|
|
let mut context = Context::from_waker(waker);
|
|
assert!(Pin::new(&mut reader).poll_read(&mut context, &mut read_buf).is_pending());
|
|
}
|
|
|
|
#[tokio::test(flavor = "current_thread")]
|
|
async fn bzip2_decoder_polls_source_off_runtime_thread() {
|
|
let input = b"name\nAlice\n";
|
|
let compressed = encode_compressed_fixture(CompressionFormat::Bzip2, input).await;
|
|
let compressed_len = u64::try_from(compressed.len()).expect("fixture length should fit in u64");
|
|
let runtime_thread = std::thread::current().id();
|
|
let source_thread = Arc::new(StdMutex::new(None));
|
|
let source = ThreadRecordingReader {
|
|
inner: Cursor::new(compressed),
|
|
thread_id: Arc::clone(&source_thread),
|
|
};
|
|
let mut reader = compressed_input_reader(
|
|
Box::new(source),
|
|
compressed_len,
|
|
CompressionFormat::Bzip2,
|
|
Arc::new(SelectInputMetrics::default()),
|
|
processed_bytes_limit(),
|
|
None,
|
|
);
|
|
let mut decoded = Vec::new();
|
|
|
|
reader
|
|
.read_to_end(&mut decoded)
|
|
.await
|
|
.expect("valid BZIP2 input should decode off the runtime thread");
|
|
|
|
assert_eq!(decoded, input);
|
|
assert_ne!(
|
|
source_thread
|
|
.lock()
|
|
.expect("thread recorder mutex should not be poisoned")
|
|
.expect("compressed source should be polled"),
|
|
runtime_thread,
|
|
"BZIP2 decoding must not poll codec work on a Tokio runtime worker"
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn decoder_thread_holds_query_admission_until_exit() {
|
|
let admission = Arc::new(tokio::sync::Semaphore::new(1));
|
|
let permit = Arc::new(
|
|
Arc::clone(&admission)
|
|
.try_acquire_owned()
|
|
.expect("query admission should be available"),
|
|
);
|
|
let (started_tx, started_rx) = oneshot::channel();
|
|
let (release_tx, release_rx) = std::sync::mpsc::channel();
|
|
let (decoded_tx, _decoded_rx) = mpsc::channel(1);
|
|
|
|
spawn_decoder_thread("s3select-guard-test", decoded_tx, Some(permit), move || {
|
|
let _ = started_tx.send(());
|
|
release_rx.recv().map_err(io::Error::other)
|
|
});
|
|
|
|
started_rx.await.expect("decoder thread should start");
|
|
assert!(Arc::clone(&admission).try_acquire_owned().is_err());
|
|
release_tx.send(()).expect("decoder thread should be releasable");
|
|
let recovered = tokio::time::timeout(std::time::Duration::from_secs(1), Arc::clone(&admission).acquire_owned())
|
|
.await
|
|
.expect("decoder thread should release admission promptly")
|
|
.expect("query admission should remain open");
|
|
drop(recovered);
|
|
}
|
|
|
|
#[test]
|
|
fn bzip2_decoder_does_not_wait_for_tokio_blocking_pool() {
|
|
let runtime = tokio::runtime::Builder::new_multi_thread()
|
|
.worker_threads(1)
|
|
.max_blocking_threads(1)
|
|
.enable_all()
|
|
.build()
|
|
.expect("test runtime should build");
|
|
|
|
runtime.block_on(async {
|
|
let (blocker_started_tx, blocker_started_rx) = tokio::sync::oneshot::channel();
|
|
let (release_tx, release_rx) = std::sync::mpsc::channel();
|
|
let blocker = tokio::task::spawn_blocking(move || {
|
|
let _ = blocker_started_tx.send(());
|
|
let _ = release_rx.recv();
|
|
});
|
|
blocker_started_rx.await.expect("blocking worker should be occupied");
|
|
|
|
let input = b"name\nAlice\n";
|
|
let compressed = encode_compressed_fixture(CompressionFormat::Bzip2, input).await;
|
|
let decode_result =
|
|
tokio::time::timeout(std::time::Duration::from_secs(2), decode(CompressionFormat::Bzip2, compressed)).await;
|
|
|
|
let _ = release_tx.send(());
|
|
blocker.await.expect("blocking worker should finish");
|
|
|
|
let (decoded, _) = decode_result.expect("BZIP2 decoder must not queue behind Tokio blocking work");
|
|
assert_eq!(decoded.expect("valid BZIP2 input should decode"), input);
|
|
});
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn processed_reader_streams_past_the_default_memory_budget() {
|
|
let expected = 64 * 1024 * 1024 + 1;
|
|
let metrics = Arc::new(SelectInputMetrics::default());
|
|
let source = tokio::io::repeat(b'x').take(expected);
|
|
let mut reader = ProcessedReader::new(source, metrics.recorder(), processed_bytes_limit());
|
|
let mut buffer = vec![0; SELECT_DECODE_CHUNK_BYTES];
|
|
let mut total = 0_u64;
|
|
|
|
loop {
|
|
let read = reader
|
|
.read(&mut buffer)
|
|
.await
|
|
.expect("streaming input within the expansion budget should pass");
|
|
if read == 0 {
|
|
break;
|
|
}
|
|
total += u64::try_from(read).expect("read buffer length fits in u64");
|
|
}
|
|
|
|
assert_eq!(total, expected);
|
|
assert_eq!(metrics.snapshot().bytes_processed, expected);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn truncated_and_corrupt_gzip_are_typed_as_truncated_input() {
|
|
let encoded = encode_compressed_fixture(CompressionFormat::Gzip, b"name\nAlice\n").await;
|
|
|
|
let mut truncated = encoded.clone();
|
|
truncated.truncate(truncated.len() - 1);
|
|
let (result, _) = decode(CompressionFormat::Gzip, truncated).await;
|
|
assert_eq!(
|
|
select_error(&result.expect_err("truncated gzip trailer must fail")),
|
|
Some(SelectError::TruncatedInput)
|
|
);
|
|
|
|
let mut corrupt = encoded;
|
|
let checksum_index = corrupt.len() - 8;
|
|
corrupt[checksum_index] ^= 0xff;
|
|
let (result, _) = decode(CompressionFormat::Gzip, corrupt).await;
|
|
assert_eq!(
|
|
select_error(&result.expect_err("corrupt gzip checksum must fail")),
|
|
Some(SelectError::TruncatedInput)
|
|
);
|
|
|
|
let mut corrupt_size = encode_compressed_fixture(CompressionFormat::Gzip, b"name\nAlice\n").await;
|
|
let size_index = corrupt_size.len() - 1;
|
|
corrupt_size[size_index] ^= 0xff;
|
|
let (result, _) = decode(CompressionFormat::Gzip, corrupt_size).await;
|
|
assert_eq!(
|
|
select_error(&result.expect_err("corrupt gzip uncompressed size must fail")),
|
|
Some(SelectError::TruncatedInput)
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn truncated_bzip2_is_typed_as_truncated_input() {
|
|
let mut encoded = encode_compressed_fixture(CompressionFormat::Bzip2, b"name\nAlice\n").await;
|
|
encoded.truncate(encoded.len() - 1);
|
|
let (result, _) = decode(CompressionFormat::Bzip2, encoded).await;
|
|
assert_eq!(
|
|
select_error(&result.expect_err("truncated bzip2 trailer must fail")),
|
|
Some(SelectError::TruncatedInput)
|
|
);
|
|
|
|
let mut corrupt = encode_compressed_fixture(CompressionFormat::Bzip2, b"name\nAlice\n").await;
|
|
let checksum_index = corrupt.len() - 2;
|
|
corrupt[checksum_index] ^= 0xff;
|
|
let (result, _) = decode(CompressionFormat::Bzip2, corrupt).await;
|
|
assert_eq!(
|
|
select_error(&result.expect_err("corrupt bzip2 checksum must fail")),
|
|
Some(SelectError::TruncatedInput)
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn trailing_non_member_bytes_fail_closed() {
|
|
for format in [CompressionFormat::Gzip, CompressionFormat::Bzip2] {
|
|
let mut encoded = encode_compressed_fixture(format, b"name\nAlice\n").await;
|
|
encoded.extend_from_slice(b"trailing garbage");
|
|
let (result, _) = decode(format, encoded).await;
|
|
assert_eq!(
|
|
select_error(&result.expect_err("trailing bytes must not be ignored")),
|
|
Some(SelectError::TruncatedInput)
|
|
);
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn oversized_compressed_record_fails_before_unbounded_buffering() {
|
|
let mut input = vec![b'x'; MAX_SELECT_RECORD_BYTES + 1];
|
|
input.push(b'\n');
|
|
let compressed = encode_compressed_fixture(CompressionFormat::Gzip, &input).await;
|
|
let (result, _) = decode(CompressionFormat::Gzip, compressed).await;
|
|
assert_eq!(
|
|
select_error(&result.expect_err("oversized compressed record must fail")),
|
|
Some(SelectError::OverMaxRecordSize)
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn one_megabyte_record_is_accepted_with_or_without_delimiter() {
|
|
for terminated in [false, true] {
|
|
let mut input = vec![b'x'; MAX_SELECT_RECORD_BYTES];
|
|
if terminated {
|
|
input.push(b'\n');
|
|
}
|
|
let compressed = encode_compressed_fixture(CompressionFormat::Gzip, &input).await;
|
|
let (decoded, _) = decode(CompressionFormat::Gzip, compressed).await;
|
|
assert_eq!(decoded.expect("record at the protocol limit should decode"), input);
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn processed_byte_limit_rejects_many_small_records() {
|
|
let input = b"{}\n".repeat(1024);
|
|
let compressed = encode_compressed_fixture(CompressionFormat::Gzip, &input).await;
|
|
let compressed_len = u64::try_from(compressed.len()).expect("fixture length should fit in u64");
|
|
let exact_stream = compressed_input_stream(
|
|
Box::new(Cursor::new(compressed.clone())),
|
|
compressed_len,
|
|
CompressionFormat::Gzip,
|
|
Arc::new(SelectInputMetrics::default()),
|
|
b"\n".to_vec(),
|
|
u64::try_from(input.len()).expect("fixture length should fit in u64"),
|
|
None,
|
|
)
|
|
.expect("record delimiter should be valid");
|
|
assert_eq!(
|
|
exact_stream
|
|
.try_collect::<Vec<_>>()
|
|
.await
|
|
.expect("decoded bytes at the processed limit should pass")
|
|
.concat(),
|
|
input
|
|
);
|
|
|
|
let stream = compressed_input_stream(
|
|
Box::new(Cursor::new(compressed)),
|
|
compressed_len,
|
|
CompressionFormat::Gzip,
|
|
Arc::new(SelectInputMetrics::default()),
|
|
b"\n".to_vec(),
|
|
u64::try_from(input.len() - 1).expect("fixture length should fit in u64"),
|
|
None,
|
|
)
|
|
.expect("record delimiter should be valid");
|
|
|
|
let error = stream
|
|
.try_collect::<Vec<_>>()
|
|
.await
|
|
.expect_err("decoded bytes over the decompression budget must fail");
|
|
assert_eq!(select_error(&error), Some(SelectError::ResourceExhausted));
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn oversized_final_record_with_partial_delimiter_fails_at_eof() {
|
|
let mut input = vec![b'x'; MAX_SELECT_RECORD_BYTES + 1];
|
|
input.push(b'\r');
|
|
let compressed = encode_compressed_fixture(CompressionFormat::Gzip, &input).await;
|
|
let compressed_len = u64::try_from(compressed.len()).expect("fixture length should fit in u64");
|
|
let metrics = Arc::new(SelectInputMetrics::default());
|
|
let stream = compressed_input_stream(
|
|
Box::new(Cursor::new(compressed)),
|
|
compressed_len,
|
|
CompressionFormat::Gzip,
|
|
metrics,
|
|
b"\r\n".to_vec(),
|
|
u64::MAX,
|
|
None,
|
|
)
|
|
.expect("record delimiter should be valid");
|
|
|
|
let error = stream
|
|
.try_collect::<Vec<_>>()
|
|
.await
|
|
.expect_err("oversized unterminated record must fail");
|
|
assert_eq!(select_error(&error), Some(SelectError::OverMaxRecordSize));
|
|
}
|
|
|
|
struct DropObservedReader {
|
|
inner: DuplexStream,
|
|
dropped: Arc<std::sync::atomic::AtomicBool>,
|
|
}
|
|
|
|
impl AsyncRead for DropObservedReader {
|
|
fn poll_read(mut self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll<io::Result<()>> {
|
|
Pin::new(&mut self.inner).poll_read(cx, buf)
|
|
}
|
|
}
|
|
|
|
impl Drop for DropObservedReader {
|
|
fn drop(&mut self) {
|
|
self.dropped.store(true, std::sync::atomic::Ordering::Release);
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn dropping_decoder_stream_releases_source_without_background_work() {
|
|
let input = b"a\n".repeat(MAX_SELECT_RECORD_BYTES);
|
|
for format in [CompressionFormat::Gzip, CompressionFormat::Bzip2] {
|
|
let compressed = encode_compressed_fixture(format, &input).await;
|
|
let compressed_len = u64::try_from(compressed.len() + 1).expect("fixture length should fit in u64");
|
|
let (source, mut peer) = tokio::io::duplex(compressed.len());
|
|
peer.write_all(&compressed).await.expect("write compressed fixture");
|
|
let dropped = Arc::new(std::sync::atomic::AtomicBool::new(false));
|
|
let admission = Arc::new(tokio::sync::Semaphore::new(1));
|
|
let permit = Arc::new(
|
|
Arc::clone(&admission)
|
|
.try_acquire_owned()
|
|
.expect("query admission should be available"),
|
|
);
|
|
let mut stream = compressed_input_stream(
|
|
Box::new(DropObservedReader {
|
|
inner: source,
|
|
dropped: Arc::clone(&dropped),
|
|
}),
|
|
compressed_len,
|
|
format,
|
|
Arc::new(SelectInputMetrics::default()),
|
|
b"\n".to_vec(),
|
|
u64::MAX,
|
|
Some(permit),
|
|
)
|
|
.expect("record delimiter should be valid");
|
|
let first = tokio::time::timeout(std::time::Duration::from_secs(1), stream.next())
|
|
.await
|
|
.expect("decoder should produce a chunk before source EOF")
|
|
.expect("decoder stream should produce a chunk")
|
|
.expect("valid partial decode should succeed");
|
|
assert!(!first.is_empty());
|
|
|
|
drop(stream);
|
|
|
|
tokio::time::timeout(std::time::Duration::from_secs(1), async {
|
|
while !dropped.load(std::sync::atomic::Ordering::Acquire) {
|
|
tokio::task::yield_now().await;
|
|
}
|
|
})
|
|
.await
|
|
.expect("dropping decoded output must cancel the blocked source read");
|
|
let recovered = tokio::time::timeout(std::time::Duration::from_secs(1), Arc::clone(&admission).acquire_owned())
|
|
.await
|
|
.expect("decoder exit should release query admission")
|
|
.expect("query admission should remain open");
|
|
drop(recovered);
|
|
}
|
|
}
|
|
}
|