From 8ebedddfa1a380d81120230302720638cc30ae0b Mon Sep 17 00:00:00 2001 From: GatewayJ <835269233@qq.com> Date: Sun, 19 Jul 2026 23:34:20 +0800 Subject: [PATCH] fix(s3select): reject truncated object streams (#5027) * fix(s3select): reject truncated object streams * fix(s3select): validate raw stream before conversion --- Cargo.lock | 3 +- crates/s3select-api/Cargo.toml | 5 +- crates/s3select-api/src/object_store.rs | 224 +++++++++++------------- crates/s3select-api/src/storage_api.rs | 2 + 4 files changed, 109 insertions(+), 125 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 447701e42..3a1122751 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -9828,12 +9828,13 @@ dependencies = [ "http 1.4.2", "metrics", "parking_lot", - "pin-project-lite", "rustfs-common", "rustfs-ecstore", "rustfs-storage-api", + "rustfs-test-utils", "s3s", "serde_json", + "tempfile", "thiserror 2.0.18", "tokio", "tokio-util", diff --git a/crates/s3select-api/Cargo.toml b/crates/s3select-api/Cargo.toml index da3c6d047..d1c678575 100644 --- a/crates/s3select-api/Cargo.toml +++ b/crates/s3select-api/Cargo.toml @@ -37,7 +37,6 @@ rustfs-storage-api.workspace = true futures = { workspace = true } futures-core = { workspace = true } http.workspace = true -pin-project-lite.workspace = true s3s = { workspace = true, features = ["minio"] } serde_json = { workspace = true, features = ["raw_value"] } thiserror = { workspace = true } @@ -48,5 +47,9 @@ tracing.workspace = true transform-stream.workspace = true url.workspace = true +[dev-dependencies] +rustfs-test-utils.workspace = true +tempfile.workspace = true + [lib] doctest = false diff --git a/crates/s3select-api/src/object_store.rs b/crates/s3select-api/src/object_store.rs index 76d7e7f64..4e18c8931 100644 --- a/crates/s3select-api/src/object_store.rs +++ b/crates/s3select-api/src/object_store.rs @@ -28,7 +28,6 @@ use futures::pin_mut; use futures::{Stream, StreamExt, future::ready, stream}; use futures_core::stream::BoxStream; use http::{HeaderMap, HeaderValue, header::HeaderName}; -use pin_project_lite::pin_project; use rustfs_common::DEFAULT_DELIMITER; use s3s::S3Result; use s3s::dto::SelectObjectContentInput; @@ -39,12 +38,8 @@ use s3s::header::{ use s3s::s3_error; use std::collections::VecDeque; use std::ops::Range; -use std::pin::Pin; use std::sync::Arc; -use std::task::Poll; -use std::task::ready; use tokio::io::AsyncReadExt; -use tokio::io::{AsyncRead, ReadBuf}; use tokio_util::io::ReaderStream; use transform_stream::AsyncTryStream; @@ -516,23 +511,12 @@ impl ObjectStore for EcObjectStore { stream }; GetResultPayload::Stream(convert_field_delimiter_stream(stream, self.need_convert.then(|| self.delimiter.clone()))) - } else if self.need_convert { - let stream = bytes_stream( - ReaderStream::with_capacity( - ConvertStream::new(reader.stream, self.delimiter.clone()), - SELECT_DEFAULT_READ_BUFFER_SIZE, - ), - original_size as usize, - ) - .boxed(); - GetResultPayload::Stream(stream) } else { let stream = bytes_stream( ReaderStream::with_capacity(reader.stream, SELECT_DEFAULT_READ_BUFFER_SIZE), original_size as usize, - ) - .boxed(); - GetResultPayload::Stream(stream) + ); + GetResultPayload::Stream(convert_field_delimiter_stream(stream, self.need_convert.then(|| self.delimiter.clone()))) }; let meta = ObjectMeta { @@ -576,61 +560,6 @@ impl ObjectStore for EcObjectStore { } } -pin_project! { - struct ConvertStream { - inner: R, - converter: DelimiterConverter, - read_buf: Vec, - pending: Vec, - pending_pos: usize, - eof: bool, - } -} - -impl ConvertStream { - fn new(inner: R, delimiter: String) -> Self { - ConvertStream { - inner, - converter: DelimiterConverter::new(delimiter.into_bytes()), - read_buf: vec![0; SELECT_DEFAULT_READ_BUFFER_SIZE], - pending: Vec::new(), - pending_pos: 0, - eof: false, - } - } -} - -impl AsyncRead for ConvertStream { - #[tracing::instrument(level = "debug", skip_all)] - fn poll_read(self: Pin<&mut Self>, cx: &mut std::task::Context<'_>, buf: &mut ReadBuf<'_>) -> Poll> { - if buf.remaining() == 0 { - return Poll::Ready(Ok(())); - } - let this = self.project(); - loop { - if drain_pending(this.pending, this.pending_pos, buf) || *this.eof { - return Poll::Ready(Ok(())); - } - - let read_len = SELECT_DEFAULT_READ_BUFFER_SIZE.min(buf.remaining().max(1)); - let bytes_read = { - let mut read_buf = ReadBuf::new(&mut this.read_buf[..read_len]); - ready!(Pin::new(&mut *this.inner).poll_read(cx, &mut read_buf))?; - read_buf.filled().len() - }; - if bytes_read == 0 { - *this.eof = true; - *this.pending = this.converter.finish(); - *this.pending_pos = 0; - continue; - } - - *this.pending = this.converter.convert_chunk(&this.read_buf[..bytes_read]); - *this.pending_pos = 0; - } - } -} - struct DelimiterConverter { delimiter: Vec, carry: Vec, @@ -680,26 +609,6 @@ impl DelimiterConverter { } } -fn drain_pending(pending: &mut Vec, pending_pos: &mut usize, buf: &mut ReadBuf<'_>) -> bool { - if *pending_pos >= pending.len() { - pending.clear(); - *pending_pos = 0; - return false; - } - if buf.remaining() == 0 { - return false; - } - - let len = buf.remaining().min(pending.len() - *pending_pos); - buf.put_slice(&pending[*pending_pos..*pending_pos + len]); - *pending_pos += len; - if *pending_pos >= pending.len() { - pending.clear(); - *pending_pos = 0; - } - true -} - fn replace_symbol(delimiter: &[u8], slice: &[u8]) -> Vec { if delimiter.is_empty() { return slice.to_vec(); @@ -1069,6 +978,15 @@ where remaining -= bytes.len(); y.yield_ok(bytes).await; } + if remaining > 0 { + return Err(o_Error::Generic { + store: "EcObjectStore", + source: Box::new(std::io::Error::new( + std::io::ErrorKind::UnexpectedEof, + format!("object stream ended with {remaining} bytes remaining"), + )), + }); + } Ok(()) }) } @@ -1076,12 +994,15 @@ where #[cfg(test)] mod test { use super::{ - ConvertStream, SelectScanRange, bytes_stream, convert_field_delimiter_stream, extract_json_sub_path_from_expression, - find_delimiter, flatten_json_document_to_ndjson, http_range_spec_from_get_range, replace_symbol, scan_range_from_bounds, + SelectScanRange, bytes_stream, convert_field_delimiter_stream, extract_json_sub_path_from_expression, find_delimiter, + flatten_json_document_to_ndjson, http_range_spec_from_get_range, replace_symbol, scan_range_from_bounds, scan_range_read_start, scan_range_stream, select_read_headers, }; + use crate::storage_api::SelectPutObjReader; + use crate::storage_api::object_store::ObjectIO as _; use bytes::Bytes; use datafusion::object_store::{self, GetRange}; + use datafusion::object_store::{GetOptions, GetResultPayload, ObjectStore as _, path::Path}; use futures::{StreamExt, TryStreamExt, stream}; use s3s::dto::{ CSVInput, CSVOutput, ExpressionType, InputSerialization, OutputSerialization, SelectObjectContentInput, @@ -1095,8 +1016,6 @@ mod test { Arc, atomic::{AtomicUsize, Ordering}, }; - use tokio::io::AsyncReadExt; - use tokio_util::io::StreamReader; #[test] fn test_replace() { @@ -1104,32 +1023,6 @@ mod test { assert_eq!(result, b"dandan,is,best"); } - #[tokio::test] - async fn test_convert_stream_replaces_delimiter_across_chunks() { - let chunks = stream::iter(vec![ - Ok::<_, std::io::Error>(Bytes::from_static(b"a&")), - Ok::<_, std::io::Error>(Bytes::from_static(b"&b&&c")), - ]); - let reader = StreamReader::new(chunks); - let mut reader = ConvertStream::new(reader, "&&".to_string()); - let mut output = Vec::new(); - reader.read_to_end(&mut output).await.unwrap(); - assert_eq!(output, b"a,b,c"); - } - - #[tokio::test] - async fn test_convert_stream_replaces_delimiter_at_stream_end() { - let chunks = stream::iter(vec![ - Ok::<_, std::io::Error>(Bytes::from_static(b"a&")), - Ok::<_, std::io::Error>(Bytes::from_static(b"&")), - ]); - let reader = StreamReader::new(chunks); - let mut reader = ConvertStream::new(reader, "&&".to_string()); - let mut output = Vec::new(); - reader.read_to_end(&mut output).await.unwrap(); - assert_eq!(output, b"a,"); - } - #[tokio::test] async fn test_scan_range_stream_keeps_header_and_selected_record() { let chunks = stream::iter(vec![Ok::<_, std::io::Error>(Bytes::from_static(b"h1,h2\n1,a\n2,b\n3,c\n"))]); @@ -1301,6 +1194,69 @@ mod test { assert_eq!(output, b"a,"); } + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn test_get_opts_validates_raw_length_before_delimiter_conversion() { + let temp_root = tempfile::tempdir().expect("create s3select test temp root"); + let env = rustfs_test_utils::TestECStoreEnv::builder() + .base_dir(temp_root.path()) + .build() + .await; + let bucket = "s3select-multi-byte-delimiter"; + let object = "input.csv"; + let input_bytes = b"a&&1\n"; + env.make_bucket(bucket, false).await; + let mut reader = SelectPutObjReader::from_vec(input_bytes.to_vec()); + env.ecstore + .put_object(bucket, object, &mut reader, &Default::default()) + .await + .expect("put multi-byte-delimited test object"); + + let input = Arc::new(SelectObjectContentInput { + bucket: bucket.to_string(), + expected_bucket_owner: None, + key: object.to_string(), + sse_customer_algorithm: None, + sse_customer_key: None, + sse_customer_key_md5: None, + request: SelectObjectContentRequest { + expression: "SELECT * FROM s3object".to_string(), + expression_type: ExpressionType::from_static(ExpressionType::SQL), + input_serialization: InputSerialization { + csv: Some(CSVInput { + field_delimiter: Some("&&".to_string()), + ..Default::default() + }), + ..Default::default() + }, + output_serialization: OutputSerialization { + csv: Some(CSVOutput::default()), + ..Default::default() + }, + request_progress: None, + scan_range: None, + }, + }); + let store = super::EcObjectStore { + input, + need_convert: true, + delimiter: "&&".to_string(), + is_json_document: false, + json_sub_path: None, + store: env.ecstore, + }; + + let result = store + .get_opts(&Path::from(object), GetOptions::default()) + .await + .expect("read multi-byte-delimited test object"); + let GetResultPayload::Stream(stream) = result.payload else { + panic!("expected streaming object payload"); + }; + let chunks: Vec = stream.try_collect().await.expect("collect converted object stream"); + + assert_eq!(chunks.concat(), b"a,1\n"); + } + #[tokio::test] async fn test_bytes_stream_stops_at_content_length() { let poll_count = Arc::new(AtomicUsize::new(0)); @@ -1324,6 +1280,28 @@ mod test { assert_eq!(poll_count.load(Ordering::SeqCst), 1); } + #[tokio::test] + async fn test_bytes_stream_rejects_early_eof() { + let source = stream::iter(vec![Ok::<_, std::io::Error>(Bytes::from_static(b"ab"))]); + let output = bytes_stream(source, 4); + futures::pin_mut!(output); + + assert_eq!(output.next().await.expect("first stream item").expect("first chunk"), b"ab"[..]); + let err = output + .next() + .await + .expect("early EOF error") + .expect_err("short stream must fail"); + let object_store::Error::Generic { store, source } = err else { + panic!("expected generic object store error"); + }; + assert_eq!(store, "EcObjectStore"); + let source = source.downcast_ref::().expect("I/O error source"); + assert_eq!(source.kind(), std::io::ErrorKind::UnexpectedEof); + assert!(source.to_string().contains("2 bytes remaining")); + assert!(output.next().await.is_none()); + } + /// A JSON array is split into one NDJSON line per element. #[test] fn test_flatten_array_produces_one_line_per_element() { diff --git a/crates/s3select-api/src/storage_api.rs b/crates/s3select-api/src/storage_api.rs index b55b8b835..eba20036b 100644 --- a/crates/s3select-api/src/storage_api.rs +++ b/crates/s3select-api/src/storage_api.rs @@ -20,6 +20,8 @@ use rustfs_ecstore::api::error::{ is_err_object_not_found as select_is_err_object_not_found_from_backend, is_err_version_not_found as select_is_err_version_not_found_from_backend, }; +#[cfg(test)] +pub(crate) use rustfs_ecstore::api::object::PutObjReader as SelectPutObjReader; use rustfs_ecstore::api::runtime::object_store_handle as resolve_select_object_store_handle_from_backend; pub(crate) use rustfs_ecstore::api::set_disk::DEFAULT_READ_BUFFER_SIZE as SELECT_DEFAULT_READ_BUFFER_SIZE; pub(crate) use rustfs_ecstore::api::storage::ECStore as SelectStore;