add checks for corrupt jumptables (#83)

* Fix typo in README

* add checks for corrupt jumptables

* fix ci?

* ci

* ci

* include file path in error message

---------

Co-authored-by: Daniel <[email protected]>

Author: normanrz

.github/workflows/python-module.yml

@@ -12,7 +12,7 @@ jobs:
         cargo test
 
   build_lin:
-    runs-on: ubuntu-20.04
+    runs-on: ubuntu-22.04
     steps:
     - uses: actions/checkout@v1
     - name: Build testing docker container
@@ -66,7 +66,7 @@ jobs:
            bash -c "\$PYBIN/python setup.py sdist bdist_wheel -p manylinux1_x86_64 && \$PYBIN/twine upload dist/*.whl dist/*.tar.gz"
 
   build_mac:
-    runs-on: macos-11
+    runs-on: macos-12
     strategy:
       max-parallel: 4
       matrix:
@@ -89,7 +89,7 @@ jobs:
         cd python
         export PATH=$PATH:$HOME/.cargo/bin
         pip install twine wheel
-        python setup.py bdist_wheel -p $(python -c "import distutils.util; print(distutils.util.get_platform().replace('x86_64', 'universal2'))")
+        python setup.py bdist_wheel -p macosx-10.9-universal2
         pip install dist/*.whl
     - name: Test
       run: |
@@ -110,7 +110,10 @@ jobs:
         twine upload dist/*.whl
 
   build_win:
-    runs-on: windows-2019
+    runs-on: windows-2022
+    defaults:
+      run:
+        shell: bash
     strategy:
       max-parallel: 4
       matrix:
@@ -123,14 +126,12 @@ jobs:
         python-version: ${{ matrix.python-version }}
         architecture: 'x64' # (x64 or x86)
     - name: Build
-      shell: bash
       run: |
         cd python
         pip install twine wheel
         python setup.py bdist_wheel -p $(python -c "import distutils.util; print(distutils.util.get_platform())")
         pip install dist/*.whl
     - name: Test
-      shell: bash
       run: |
         cd python
         python -c "import wkw"
@@ -142,7 +143,6 @@ jobs:
         python -c "import wkw"
         pytest tests -k "not big_read"
     - name: Publish
-      shell: bash
       if: startsWith(github.event.ref, 'refs/tags') && matrix.python-version == '3.8'
       env:
         TWINE_USERNAME: ${{ secrets.PYPI_USERNAME }}

README.md

@@ -98,7 +98,7 @@ If the file header indicates that the blocks were compressed using LZ4 (by
 having a blockType value of either 0x02 or 0x03), the file header is immediately
 followed by the jump table.
 
-The jump table is an array of N unsigned 64-bit intergers, where N is the number
+The jump table is an array of N unsigned 64-bit integers, where N is the number
 of blocks in the file. The n-th entry of the jump table contains the absolute
 address (relative to the beginning of the file) of the first byte after the data
 of block n.

python/tests/test_wkw.py

@@ -1,3 +1,4 @@
+from pathlib import Path
 import wkw
 import numpy as np
 import shutil
@@ -50,14 +51,12 @@ def test_non_negative_offsets():
 
 def test_empty_read():
     with wkw.Dataset.create("tests/tmp", wkw.Header(np.uint8)) as dataset:
-
         data = dataset.read((1, 1, 1), (0, 0, 0))
         assert data.shape == (1, 0, 0, 0)
 
 
 def test_readwrite():
     with wkw.Dataset.create("tests/tmp", wkw.Header(np.uint8)) as dataset:
-
         header_size = path.getsize(path.join("tests/tmp", "header.wkw"))
         test_data = generate_test_data(dataset.header.voxel_type)
 
@@ -96,7 +95,6 @@ def test_readwrite_live_compression_should_enforce_full_file_write():
         with wkw.Dataset.create(
             "tests/tmp", wkw.Header(np.uint8, block_type=BLOCK_TYPE_LZ4)
         ) as dataset:
-
             test_data = generate_test_data(dataset.header.voxel_type)
             dataset.write(POSITION, test_data)
 
@@ -145,7 +143,6 @@ def test_readwrite_live_compression_should_truncate():
 
 def test_compress():
     with wkw.Dataset.create("tests/tmp", wkw.Header(np.uint8)) as dataset:
-
         test_data = generate_test_data(dataset.header.voxel_type)
         dataset.write(POSITION, test_data)
 
@@ -269,7 +266,6 @@ def write_and_test_in_given_order(wkw_path, order):
 
 
 def test_multiple_writes_and_reads():
-
     mem_buffer = np.zeros((200, 200, 200), dtype=np.uint8, order="F")
     with wkw.Dataset.create("tests/tmp", wkw.Header(np.uint8)) as dataset:
         for i in range(10):
@@ -289,7 +285,6 @@ def test_multiple_writes_and_reads():
 
 
 def test_multi_channel_column_major_order():
-
     with wkw.Dataset.create(
         "tests/tmp", wkw.Header(np.uint8, num_channels=3)
     ) as dataset:
@@ -340,6 +335,31 @@ def test_invalid_dataset():
     print(excinfo.value)
 
 
+def test_dataset_with_invalid_jumptable():
+    with wkw.Dataset.create(
+        "tests/tmp",
+        wkw.Header(np.uint8, block_type=wkw.Header.BLOCK_TYPE_LZ4HC, file_len=4),
+    ) as dataset:
+        test_data = generate_test_data(
+            dataset.header.voxel_type, (32 * 4, 32 * 4, 32 * 4)
+        )
+        dataset.write((0, 0, 0), test_data)
+
+    # reset jumptable with zeros
+    data_path = Path("tests/tmp") / "z0" / "y0" / "x0.wkw"
+    wkw_file_bytes = bytearray(data_path.read_bytes())
+    zeros = b"\x00" * (4**3 * 8)
+    wkw_file_bytes[16 : (16 + (4**3) * 8)] = zeros
+    data_path.write_bytes(wkw_file_bytes)
+
+    with pytest.raises(wkw.wkw.WKWException) as excinfo:
+        with wkw.Dataset.open("tests/tmp") as dataset:
+            dataset.read((0, 0, 0), (32 * 4,) * 3)
+
+    print(excinfo)
+    assert "Corrupt jump table" in str(excinfo)
+
+
 def generate_test_data(dtype, size=SIZE, order="C"):
     return np.array(
         np.random.uniform(np.iinfo(dtype).min, np.iinfo(dtype).max, size).astype(dtype),

rust/src/dataset.rs

@@ -55,10 +55,12 @@ impl Dataset {
         }
 
         // create header file
-        let mut file = fs::File::create(&header_path).or(Err(format!(
-            "Could not create header file {:?}",
-            &header_path
-        )))?;
+        let mut file = fs::File::create(&header_path).or_else(|err| {
+            Err(format!(
+                "Could not create header file {:?}: {}",
+                &header_path, err
+            ))
+        })?;
 
         header.write(&mut file)
     }
@@ -104,7 +106,15 @@ impl Dataset {
 
                     // try to open file
                     if let Ok(mut file) = File::open(&cur_path) {
-                        file.read_mat(cur_src_pos, mat, cur_dst_pos)?;
+                        match file.read_mat(cur_src_pos, mat, cur_dst_pos) {
+                            Ok(_) => {}
+                            Err(err) => {
+                                return Err(format!(
+                                    "Error while reading from file {:?}: {}",
+                                    &cur_path, err
+                                ));
+                            }
+                        }
                     }
                 }
             }
@@ -176,8 +186,24 @@ impl Dataset {
                     let cur_src_pos = cur_box.min() - dst_pos;
                     let cur_dst_pos = cur_box.min() - cur_file_box.min();
 
-                    let mut file = File::open_or_create(&cur_path, &self.header)?;
-                    file.write_mat(cur_dst_pos, mat, cur_src_pos)?;
+                    let mut file = match File::open_or_create(&cur_path, &self.header) {
+                        Ok(file) => file,
+                        Err(err) => {
+                            return Err(format!(
+                                "Error while open file {:?} for writing: {}",
+                                &cur_path, err
+                            ));
+                        }
+                    };
+                    match file.write_mat(cur_dst_pos, mat, cur_src_pos) {
+                        Ok(_) => {}
+                        Err(err) => {
+                            return Err(format!(
+                                "Error while writing to file {:?}: {}",
+                                &cur_path, err
+                            ));
+                        }
+                    }
                 }
             }
         }

rust/src/file.rs

@@ -357,7 +357,15 @@ impl File {
         let block_size_raw = self.header.block_size();
 
         let buf_lz4_orig = &mut *self.disk_block_buf.as_mut().unwrap();
-        let buf_lz4 = &mut buf_lz4_orig[..block_size_lz4];
+        let buf_lz4 = match buf_lz4_orig.get_mut(..block_size_lz4) {
+            Some(buf_lz4) => buf_lz4,
+            None => {
+                return Err(format!(
+                    "Unexpected compressed block length {}",
+                    block_size_lz4
+                ))
+            }
+        };
 
         // read compressed block
         self.file

rust/src/header.rs

@@ -208,12 +208,18 @@ impl Header {
                 let jump_table = &*self.jump_table.as_ref().unwrap();
 
                 if block_idx == 0 {
-                    let block_size = jump_table[0] - self.data_offset;
-                    Ok(block_size as usize)
+                    let block_size = jump_table[0].checked_sub(self.data_offset);
+                    match block_size {
+                        Some(block_size) => Ok(block_size as usize),
+                        None => Err(String::from("Corrupt jump table")),
+                    }
                 } else if block_idx < self.file_vol() {
                     let block_idx = block_idx as usize;
-                    let block_size = jump_table[block_idx] - jump_table[block_idx - 1];
-                    Ok(block_size as usize)
+                    let block_size = jump_table[block_idx].checked_sub(jump_table[block_idx - 1]);
+                    match block_size {
+                        Some(block_size) => Ok(block_size as usize),
+                        None => Err(String::from("Corrupt jump table")),
+                    }
                 } else {
                     Err(String::from("Block index out of bounds"))
                 }