From d080b5909e7889f7353f4582102299dc18996d8f Mon Sep 17 00:00:00 2001
From: Jan De Kock <jan@twofold.be>
Date: Sun, 16 Jun 2024 23:51:43 +0200
Subject: [PATCH] fix: Invalid mode fixes

---
 src/main/java/be/twofold/tinybcdec/BC6Decoder.java   |  6 +++++-
 src/main/java/be/twofold/tinybcdec/BC7Decoder.java   | 12 +++++++-----
 .../java/be/twofold/tinybcdec/BC6DecoderTest.java    | 11 +++++++++++
 .../java/be/twofold/tinybcdec/BC7DecoderTest.java    |  7 +++++++
 4 files changed, 30 insertions(+), 6 deletions(-)

diff --git a/src/main/java/be/twofold/tinybcdec/BC6Decoder.java b/src/main/java/be/twofold/tinybcdec/BC6Decoder.java
index e581303..690e895 100644
--- a/src/main/java/be/twofold/tinybcdec/BC6Decoder.java
+++ b/src/main/java/be/twofold/tinybcdec/BC6Decoder.java
@@ -31,7 +31,11 @@ final class BC6Decoder extends BlockDecoder {
     @SuppressWarnings("PointlessArithmeticExpression")
     public void decodeBlock(byte[] src, int srcPos, byte[] dst, int dstPos, int bytesPerLine) {
         Bits bits = Bits.from(src, srcPos);
-        Mode mode = MODES.get(mode(bits));
+        int modeIndex = mode(bits);
+        if (modeIndex >= MODES.size()) {
+            return;
+        }
+        Mode mode = MODES.get(modeIndex);
 
         int[] colors = new int[12];
         for (short op : mode.ops) {
diff --git a/src/main/java/be/twofold/tinybcdec/BC7Decoder.java b/src/main/java/be/twofold/tinybcdec/BC7Decoder.java
index ed77a18..2d623b8 100644
--- a/src/main/java/be/twofold/tinybcdec/BC7Decoder.java
+++ b/src/main/java/be/twofold/tinybcdec/BC7Decoder.java
@@ -84,7 +84,11 @@ final class BC7Decoder extends BlockDecoder {
     @Override
     public void decodeBlock(byte[] src, int srcPos, byte[] dst, int dstPos, int bytesPerLine) {
         Bits bits = Bits.from(src, srcPos);
-        Mode mode = MODES.get(mode(bits));
+        int modeIndex = mode(bits);
+        if (modeIndex >= MODES.size()) {
+            return;
+        }
+        Mode mode = MODES.get(modeIndex);
         int partition = bits.get(mode.pb);
         int rotation = bits.get(mode.rb);
         boolean selection = bits.get(mode.isb) != 0;
@@ -235,12 +239,10 @@ private int unpack(int i, int n) {
 
     private int mode(Bits bits) {
         int mode = 0;
-        while (true) {
-            if (bits.get1() != 0) {
-                return mode;
-            }
+        while (mode < 8 && bits.get1() == 0) {
             mode++;
         }
+        return mode;
     }
 
     private static final class Mode {
diff --git a/src/test/java/be/twofold/tinybcdec/BC6DecoderTest.java b/src/test/java/be/twofold/tinybcdec/BC6DecoderTest.java
index 51aba51..20a4d9f 100644
--- a/src/test/java/be/twofold/tinybcdec/BC6DecoderTest.java
+++ b/src/test/java/be/twofold/tinybcdec/BC6DecoderTest.java
@@ -30,4 +30,15 @@ void testBC6H_SF16() throws IOException {
         assertThat(actual).isEqualTo(expected);
     }
 
+    @Test
+    void testBC6InvalidBlock() {
+        byte[] src = new byte[16];
+        byte[] invalidModes = {0b10011, 0b10111, 0b11011, 0b11111};
+        for (byte invalidMode : invalidModes) {
+            src[0] = invalidMode;
+            BlockDecoder decoder = BlockDecoder.create(BlockFormat.BC6Unsigned, PixelOrder.RGB);
+            var actual = decoder.decode(4, 4, src, 0);
+            assertThat(actual).isEqualTo(new byte[16 * 3 * 2]);
+        }
+    }
 }
diff --git a/src/test/java/be/twofold/tinybcdec/BC7DecoderTest.java b/src/test/java/be/twofold/tinybcdec/BC7DecoderTest.java
index ac8dbcc..96496ea 100644
--- a/src/test/java/be/twofold/tinybcdec/BC7DecoderTest.java
+++ b/src/test/java/be/twofold/tinybcdec/BC7DecoderTest.java
@@ -20,4 +20,11 @@ void testBC7() throws IOException {
         assertThat(actual).isEqualTo(expected);
     }
 
+    @Test
+    void testBC7InvalidBlock() {
+        byte[] src = new byte[16];
+        byte[] actual = decoder.decode(4, 4, src, 0);
+        assertThat(actual).isEqualTo(new byte[16 * 4]);
+    }
+
 }