Merge pull request #207 from influxdata/crepererum/no-ipc-compression

feat: disallow compressed IPC data

Author: crepererum

Cargo.lock

@@ -8,6 +8,7 @@ license.workspace = true
 arrow.workspace = true
 
 [dev-dependencies]
+arrow = { workspace = true, features = ["ipc_compression"] }
 insta.workspace = true
 
 [lints]

arrow2bytes/Cargo.toml

@@ -0,0 +1,115 @@
+//! Scan IPC data for compressed data.
+//!
+//! This is a workaround until <https://github.com/apache/arrow-rs/issues/8917> is implemented.
+
+use std::io::{Cursor, Read, Seek};
+
+use arrow::{
+    error::ArrowError,
+    ipc::{BodyCompression, MessageHeader, root_as_message},
+};
+
+/// Detect and fail if there's compressed data.
+pub(crate) fn detect_compressed_data(bytes: &[u8]) -> Result<(), ArrowError> {
+    let mut reader = Cursor::new(bytes);
+
+    loop {
+        let Some(meta_len) = read_meta_len(&mut reader)? else {
+            break;
+        };
+        let mut meta = vec![0; meta_len];
+        reader.read_exact(&mut meta)?;
+        let msg = root_as_message(&meta).map_err(|err| {
+            ArrowError::ParseError(format!("Unable to get root as message: {err:?}"))
+        })?;
+
+        match msg.header_type() {
+            MessageHeader::Schema => {
+                // never compressed
+            }
+            MessageHeader::DictionaryBatch => {
+                if let Some(batch) = msg.header_as_dictionary_batch()
+                    && let Some(batch) = batch.data()
+                    && let Some(compression) = batch.compression()
+                {
+                    return Err(compression_err("dictionary batch", compression));
+                }
+            }
+            MessageHeader::RecordBatch => {
+                if let Some(batch) = msg.header_as_record_batch()
+                    && let Some(compression) = batch.compression()
+                {
+                    return Err(compression_err("record batch", compression));
+                }
+            }
+            x => {
+                return Err(ArrowError::ParseError(format!(
+                    "Unsupported message header type in IPC stream: '{x:?}'"
+                )));
+            }
+        }
+
+        let body_len = msg.bodyLength();
+        if body_len < 0 {
+            return Err(ArrowError::ParseError(format!(
+                "Invalid body length: {body_len}"
+            )));
+        }
+        reader.seek_relative(body_len)?;
+    }
+
+    Ok(())
+}
+
+/// Read the metadata length for the next message from the underlying stream.
+///
+/// # Returns
+/// - `Ok(None)` if the reader signals the end of stream with EOF on
+///   the first read
+/// - `Err(_)` if the reader returns an error other than EOF on the first
+///   read, or if the metadata length is less than 0.
+/// - `Ok(Some(_))` with the length otherwise.
+fn read_meta_len(reader: &mut Cursor<&[u8]>) -> Result<Option<usize>, ArrowError> {
+    const CONTINUATION_MARKER: [u8; 4] = [0xff; 4];
+    let mut meta_len: [u8; 4] = [0; 4];
+    match reader.read_exact(&mut meta_len) {
+        Ok(_) => {}
+        Err(e) => {
+            return if e.kind() == std::io::ErrorKind::UnexpectedEof {
+                // Handle EOF without the "0xFFFFFFFF 0x00000000"
+                // valid according to:
+                // https://arrow.apache.org/docs/format/Columnar.html#ipc-streaming-format
+                Ok(None)
+            } else {
+                Err(ArrowError::from(e))
+            };
+        }
+    }
+
+    let meta_len = {
+        // If a continuation marker is encountered, skip over it and read
+        // the size from the next four bytes.
+        if meta_len == CONTINUATION_MARKER {
+            reader.read_exact(&mut meta_len)?;
+        }
+
+        i32::from_le_bytes(meta_len)
+    };
+
+    if meta_len == 0 {
+        return Ok(None);
+    }
+
+    let meta_len = usize::try_from(meta_len)
+        .map_err(|_| ArrowError::ParseError(format!("Invalid metadata length: {meta_len}")))?;
+
+    Ok(Some(meta_len))
+}
+
+/// Generate error for encountered compression.
+fn compression_err(what: &'static str, compression: BodyCompression<'_>) -> ArrowError {
+    ArrowError::IpcError(format!(
+        "IPC {what} is compressed using {}, but compressed data MUST NOT cross the security boundary. If you want to handle compressed data, please decompress it within the guest.",
+        compression.codec().variant_name().unwrap_or("<unknown>")
+    ))
+}

arrow2bytes/src/compression_check.rs

@@ -22,6 +22,8 @@ use arrow::{
 #[cfg(test)]
 use insta as _;
 
+mod compression_check;
+
 /// Convert an [`Array`] to bytes.
 ///
 /// This is done by encoding writing this as a [`RecordBatch`] with a single [`Field`].
@@ -48,6 +50,8 @@ pub fn array2bytes(array: ArrayRef) -> Vec<u8> {
 ///
 /// See [`array2bytes`] for the reverse method and the format description.
 pub fn bytes2array(bytes: &[u8]) -> Result<ArrayRef, ArrowError> {
+    compression_check::detect_compressed_data(bytes)?;
+
     let cursor = Cursor::new(bytes);
     let mut reader = StreamReader::try_new(cursor, None)?;
     let Some(res) = reader.next() else {

arrow2bytes/src/lib.rs

@@ -8,20 +8,18 @@ use arrow::{
         ArrayRef, Int64Array, ListArray, RecordBatch, RecordBatchOptions, StringDictionaryBuilder,
     },
     datatypes::{DataType, Field, Int32Type, Schema},
-    ipc::writer::StreamWriter,
+    error::ArrowError,
+    ipc::{
+        CompressionType,
+        writer::{IpcWriteOptions, StreamWriter},
+    },
 };
 use datafusion_udf_wasm_arrow2bytes::{array2bytes, bytes2array};
 
 #[test]
 fn test_roundtrip() {
-    roundtrip(Arc::new(Int64Array::from_iter([Some(1), None, Some(3)])));
-
-    let mut builder = StringDictionaryBuilder::<Int32Type>::new();
-    builder.append("foo").unwrap();
-    builder.append_null();
-    builder.append("bar").unwrap();
-    builder.append("foo").unwrap();
-    roundtrip(Arc::new(builder.finish()));
+    roundtrip(int64_array());
+    roundtrip(string_dict_array());
 }
 
 #[test]
@@ -89,14 +87,8 @@ fn test_err_multiple_columns() {
         Field::new("a", DataType::Int64, true),
         Field::new("b", DataType::Int64, true),
     ]));
-    let batch = RecordBatch::try_new(
-        Arc::clone(&schema),
-        vec![
-            Arc::new(Int64Array::new_null(0)),
-            Arc::new(Int64Array::new_null(0)),
-        ],
-    )
-    .unwrap();
+    let batch =
+        RecordBatch::try_new(Arc::clone(&schema), vec![int64_array(), int64_array()]).unwrap();
     let mut writer =
         StreamWriter::try_new(Vec::new(), &schema).expect("writing to buffer never fails");
     writer.write(&batch).unwrap();
@@ -136,9 +128,66 @@ fn test_deeply_nested() {
     );
 }
 
+#[test]
+fn test_err_compression() {
+    insta::assert_snapshot!(
+        compression_err(int64_array(), CompressionType::LZ4_FRAME),
+        @"Ipc error: IPC record batch is compressed using LZ4_FRAME, but compressed data MUST NOT cross the security boundary. If you want to handle compressed data, please decompress it within the guest.",
+    );
+    insta::assert_snapshot!(
+        compression_err(int64_array(), CompressionType::ZSTD),
+        @"Ipc error: IPC record batch is compressed using ZSTD, but compressed data MUST NOT cross the security boundary. If you want to handle compressed data, please decompress it within the guest.",
+    );
+    insta::assert_snapshot!(
+        compression_err(string_dict_array(), CompressionType::LZ4_FRAME),
+        @"Ipc error: IPC dictionary batch is compressed using LZ4_FRAME, but compressed data MUST NOT cross the security boundary. If you want to handle compressed data, please decompress it within the guest.",
+    );
+    insta::assert_snapshot!(
+        compression_err(string_dict_array(), CompressionType::ZSTD),
+        @"Ipc error: IPC dictionary batch is compressed using ZSTD, but compressed data MUST NOT cross the security boundary. If you want to handle compressed data, please decompress it within the guest.",
+    );
+}
+
 #[track_caller]
 fn roundtrip(array: ArrayRef) {
     let bytes = array2bytes(Arc::clone(&array));
     let array2 = bytes2array(&bytes).unwrap();
     assert_eq!(&array, &array2);
 }
+
+/// Create a non-empty int64 array.
+fn int64_array() -> ArrayRef {
+    Arc::new(Int64Array::from_iter([Some(1), None, Some(3)]))
+}
+
+/// Create a non-empty dict-encoded string array.
+fn string_dict_array() -> ArrayRef {
+    let mut builder = StringDictionaryBuilder::<Int32Type>::new();
+    builder.append("foo").unwrap();
+    builder.append_null();
+    builder.append("bar").unwrap();
+    builder.append("foo").unwrap();
+    Arc::new(builder.finish())
+}
+
+#[track_caller]
+fn compression_err(array: ArrayRef, compression: CompressionType) -> ArrowError {
+    let schema = Arc::new(Schema::new(vec![Field::new(
+        "a",
+        array.data_type().clone(),
+        true,
+    )]));
+    let batch = RecordBatch::try_new(Arc::clone(&schema), vec![array]).unwrap();
+    let mut writer = StreamWriter::try_new_with_options(
+        Vec::new(),
+        &schema,
+        IpcWriteOptions::default()
+            .try_with_compression(Some(compression))
+            .unwrap(),
+    )
+    .expect("writing to buffer never fails");
+    writer.write(&batch).unwrap();
+    let bytes = writer.into_inner().unwrap();
+
+    bytes2array(&bytes).unwrap_err()
+}