diff --git a/dagshub/data_engine/model/datasource.py b/dagshub/data_engine/model/datasource.py index c3bfd862..ab8a80e8 100644 --- a/dagshub/data_engine/model/datasource.py +++ b/dagshub/data_engine/model/datasource.py @@ -155,6 +155,21 @@ def to_dict(self, ds: "Datasource") -> Dict[str, Any]: _metadata_contexts: Dict[Union[int, str], "MetadataContextManager"] = {} +def _reject_brotli_parquet(file_path: str) -> None: + # CVE-2026-111497 is an out-of-bounds write in the Brotli decoder. + # pyarrow wheels bundle that decoder. The Parquet footer is not Brotli, + # so this check reads the footer and does not decode column pages. + # Remove the check after a pyarrow wheel bundles Brotli 1.0.10, 1.1.1, or 1.2.1. + import pyarrow.parquet as pq + + metadata = pq.ParquetFile(file_path).metadata + for row_group_index in range(metadata.num_row_groups): + row_group = metadata.row_group(row_group_index) + for column_index in range(row_group.num_columns): + if row_group.column(column_index).compression == "BROTLI": + raise ValueError("Parquet files that use Brotli compression are not supported") + + class Datasource: def __init__( self, @@ -1630,6 +1645,7 @@ def _convert_file_to_df(file_path: str): if file_path.lower().endswith(".csv"): df = pandas.read_csv(file_path) elif file_path.lower().endswith(".parquet"): + _reject_brotli_parquet(file_path) df = pandas.read_parquet(file_path) elif file_path.lower().endswith(".zip"): df = pandas.read_csv(file_path, compression="zip") diff --git a/tests/data_engine/test_brotli_parquet.py b/tests/data_engine/test_brotli_parquet.py new file mode 100644 index 00000000..baa30b58 --- /dev/null +++ b/tests/data_engine/test_brotli_parquet.py @@ -0,0 +1,28 @@ +import pyarrow as pa +import pyarrow.parquet as pq +import pytest + +from dagshub.data_engine.model.datasource import Datasource, _reject_brotli_parquet + + +def _write_parquet(path, compression): + table = pa.table({"path": ["a.txt"], "value": [1]}) + pq.write_table(table, path, compression=compression) + + +def test_reject_brotli_parquet(tmp_path): + path = tmp_path / "meta.parquet" + _write_parquet(path, "brotli") + + with pytest.raises(ValueError, match="Brotli"): + _reject_brotli_parquet(str(path)) + + +def test_snappy_parquet_is_allowed(tmp_path): + path = tmp_path / "meta.parquet" + _write_parquet(path, "snappy") + + _reject_brotli_parquet(str(path)) + frame = Datasource._convert_file_to_df(str(path)) + + assert frame["path"].tolist() == ["a.txt"]