Merge pull request #11 from chris-allan/fix-range-checks

Fix and expand range checks when performing tile retrieval

src/main/java/com/glencoesoftware/omero/zarr/ZarrPixelBuffer.java

@@ -28,7 +28,6 @@
 import java.util.Arrays;
 import java.util.List;
 import java.util.Map;
-import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ExecutionException;
 import java.util.stream.IntStream;
 
@@ -138,7 +137,7 @@ public ZarrPixelBuffer(Pixels pixels, Path root, Integer maxPlaneWidth,
                     int h = key.get(7);
                     int[] shape = new int[] { 1, 1, 1, h, w };
                     byte[] innerBuffer =
-                            new byte[length(shape) * getByteWidth()];
+                            new byte[(int) length(shape) * getByteWidth()];
                     setResolutionLevel(resolutionLevel);
                     return getTileDirect(z, c, t, x, y, w, h, innerBuffer);
                 });
@@ -181,25 +180,17 @@ public int getPixelsType() {
      * "https://numpy.org/doc/stable/reference/generated/numpy.shape.html">
      * numpy.shape</a> documentation
      */
-    private int length(int[] shape) {
-        return IntStream.of(shape).reduce(1, Math::multiplyExact);
+    private long length(int[] shape) {
+        return IntStream.of(shape)
+                .mapToLong(a -> (long) a)
+                .reduce(1, Math::multiplyExact);
     }
 
     private void read(byte[] buffer, int[] shape, int[] offset)
             throws IOException {
-        Integer sizeX = shape[4];
-        Integer sizeY = shape[3];
-        if((sizeX * sizeY) > (maxPlaneWidth*maxPlaneHeight)) {
-            throw new IllegalArgumentException(String.format(
-                    "Requested Region Size %d * %d > max plane size %d * %d", sizeX,
-                    sizeY, maxPlaneWidth, maxPlaneHeight));
-        }
-        if (sizeX < 0) {
-            throw new IllegalArgumentException("width < 0");
-        }
-        if (sizeY < 0) {
-            throw new IllegalArgumentException("height < 0");
-        }
+        // Check planar read size (sizeX and sizeY only)
+        checkReadSize(Arrays.copyOfRange(shape, 3, 5));
+
         try {
             ByteBuffer asByteBuffer = ByteBuffer.wrap(buffer);
             DataType dataType = array.getDataType();
@@ -326,31 +317,42 @@ public void close() throws IOException {
      * Implemented as specified by {@link PixelBuffer} I/F.
      * @see PixelBuffer#checkBounds(Integer, Integer, Integer, Integer, Integer)
      */
+    @Override
     public void checkBounds(Integer x, Integer y, Integer z, Integer c,
             Integer t)
             throws DimensionsOutOfBoundsException {
         if (x != null && (x > getSizeX() - 1 || x < 0)) {
             throw new DimensionsOutOfBoundsException("X '" + x
-                    + "' greater than sizeX '" + getSizeX() + "'.");
+                    + "' greater than sizeX '" + getSizeX() + "' or < '0'.");
         }
         if (y != null && (y > getSizeY() - 1 || y < 0)) {
             throw new DimensionsOutOfBoundsException("Y '" + y
-                    + "' greater than sizeY '" + getSizeY() + "'.");
+                    + "' greater than sizeY '" + getSizeY() + "' or < '0'.");
         }
 
         if (z != null && (z > getSizeZ() - 1 || z < 0)) {
             throw new DimensionsOutOfBoundsException("Z '" + z
-                    + "' greater than sizeZ '" + getSizeZ() + "'.");
+                    + "' greater than sizeZ '" + getSizeZ() + "' or < '0'.");
         }
 
         if (c != null && (c > getSizeC() - 1 || c < 0)) {
             throw new DimensionsOutOfBoundsException("C '" + c
-                    + "' greater than sizeC '" + getSizeC() + "'.");
+                    + "' greater than sizeC '" + getSizeC() + "' or < '0'.");
         }
 
         if (t != null && (t > getSizeT() - 1 || t < 0)) {
             throw new DimensionsOutOfBoundsException("T '" + t
-                    + "' greater than sizeT '" + getSizeT() + "'.");
+                    + "' greater than sizeT '" + getSizeT() + "' or < '0'.");
+        }
+    }
+
+    public void checkReadSize(int[] shape) {
+        long length = length(shape);
+        long maxLength = maxPlaneWidth * maxPlaneHeight;
+        if (length > maxLength) {
+            throw new IllegalArgumentException(String.format(
+                    "Requested shape %s > max plane size %d * %d",
+                    Arrays.toString(shape), maxPlaneWidth, maxPlaneHeight));
         }
     }
 
@@ -454,6 +456,11 @@ public PixelData getTile(
         checkBounds(x, y, z, c, t);
         //Check check bottom-right of tile in bounds
         checkBounds(x + w - 1, y + h - 1, z, c, t);
+        //Check planar read size (sizeX and sizeY only), essential that this
+        //happens before the similar check in read().  Otherwise we will
+        //potentially allocate massive inner buffers in the tile cache
+        //asynchronous entry builder that will never be used.
+        checkReadSize(new int[] { w, h });
 
         List<List<Integer>> keys = new ArrayList<List<Integer>>();
         List<Integer> key = null;
@@ -476,10 +483,7 @@ public PixelData getTile(
             // of channels can be unpredictable.
             tileCache.synchronous().invalidateAll();
         }
-        CompletableFuture<Map<List<Integer>, byte[]>> future =
-                tileCache.getAll(keys);
-        Map<List<Integer>, byte[]> values = future.join();
-        return toPixelData(values.get(key));
+        return toPixelData(tileCache.getAll(keys).join().get(key));
     }
 
     @Override

src/test/java/com/glencoesoftware/omero/zarr/ZarrPixelBufferTest.java

@@ -41,12 +41,11 @@
 
 import com.bc.zarr.ZarrArray;
 import com.bc.zarr.ZarrGroup;
+import com.fasterxml.jackson.databind.ObjectMapper;
 import com.github.benmanes.caffeine.cache.Caffeine;
 
 import picocli.CommandLine;
 import com.glencoesoftware.bioformats2raw.Converter;
-import com.glencoesoftware.omero.zarr.ZarrPixelsService;
-import com.glencoesoftware.omero.zarr.ZarrPixelBuffer;
 
 import loci.formats.FormatTools;
 import loci.formats.in.FakeReader;
@@ -503,8 +502,39 @@ public void testGetTileLargerThanImage()
         }
     }
 
-    @Test
-    public void testTileExceedsMax() throws IOException, InvalidRangeException {
+    @Test(expected = IllegalArgumentException.class)
+    public void testTileIntegerOverflow()
+            throws IOException, InvalidRangeException {
+        int sizeT = 1;
+        int sizeC = 3;
+        int sizeZ = 1;
+        int sizeY = 1;
+        int sizeX = 1;
+        int resolutions = 1;
+        Pixels pixels = new Pixels(
+                null, null, sizeX, sizeY, sizeZ, sizeC, sizeT, "", null);
+        Path output = writeTestZarr(
+                sizeT, sizeC, sizeZ, sizeY, sizeX, "uint16",
+                resolutions);
+
+        // Hack the .zarray so we can appear as though we have more data than
+        // we actually have written above.
+        ObjectMapper mapper = new ObjectMapper();
+        HashMap<String, Object> zArray = mapper.readValue(
+                Files.readAllBytes(output.resolve("0/0/.zarray")),
+                HashMap.class);
+        List<Integer> shape = (List<Integer>) zArray.get("shape");
+        shape.set(3, 50000);
+        shape.set(4, 50000);
+        mapper.writeValue(output.resolve("0/0/.zarray").toFile(), zArray);
+        try (ZarrPixelBuffer zpbuf =
+                createPixelBuffer(pixels, output.resolve("0"), 32, 32)) {
+            zpbuf.getTile(0, 0, 0, 0, 0, 50000, 50000);
+        }
+    }
+
+    @Test(expected = IllegalArgumentException.class)
+    public void testTileExceedsMinMax() throws IOException {
         int sizeT = 1;
         int sizeC = 3;
         int sizeZ = 1;
@@ -519,8 +549,9 @@ public void testTileExceedsMax() throws IOException, InvalidRangeException {
 
         try (ZarrPixelBuffer zpbuf =
                 createPixelBuffer(pixels, output.resolve("0"), 32, 32)) {
-            PixelData pixelData = zpbuf.getTile(0, 0, 0, 0, 0, 32, 33);
-            Assert.assertNull(pixelData);
+            Assert.assertNull(zpbuf.getTile(0, 0, 0, 0, 0, 32, 33));
+            // Throws exception
+            zpbuf.getTile(0, 0, 0, -1, 0, 1, 1);
         }
     }