package unit; import java.util.*; /** * Unit test for libjpeg-turbo-0001: CWE-407 O(P*C) linear color dedup in rdcolmap.c * * add_map_entry() performs a linear scan of the colormap (up to 256 entries) * for every pixel in a PPM colormap file. For large images that exhaust the * 256-color palette early, subsequent pixels each cost O(256) comparisons. * Total: O(W*H * C) comparisons. * * Fix: replace the linear scan with an O(1) open-addressing hash set, reset * once per _read_color_map() call (the single external entry point). * * JPEG_MAX_DIMENSION = 65500 -> W*H up to 4.3 billion pixels. * At 256 colors saturated: (4.3B - 256) * 256 = 1.1 trillion comparisons. * Speedup: ~256x for images that saturate the palette quickly. */ public class LibjpegTurbo0001Test { // ------------------------------------------------------------------------- // Unpatched: O(C) linear scan per pixel (models rdcolmap.c add_map_entry) // ------------------------------------------------------------------------- static class UnpatchedColorMap { private final int[] R = new int[256]; private final int[] G = new int[256]; private final int[] B = new int[256]; private int ncolors = 0; int comparisons = 0; /** Returns true if the color was new and added. */ boolean addMapEntry(int r, int g, int b) { for (int i = 0; i < ncolors; i++) { comparisons++; if (R[i] == r && G[i] == g && B[i] == b) return false; // already in map } if (ncolors >= 256) return false; // overflow guard R[ncolors] = r; G[ncolors] = g; B[ncolors] = b; ncolors++; return true; } } // ------------------------------------------------------------------------- // Patched: O(1) open-addressing hash set per pixel // ------------------------------------------------------------------------- static class PatchedColorMap { private static final int HASH_BITS = 9; // 512 slots private static final int HASH_SIZE = 1 << HASH_BITS; private static final int HASH_MASK = HASH_SIZE - 1; private final int[] slots = new int[HASH_SIZE]; // 0 == empty sentinel private final int[] R = new int[256]; private final int[] G = new int[256]; private final int[] B = new int[256]; private int ncolors = 0; int probes = 0; void reset() { Arrays.fill(slots, 0); ncolors = 0; probes = 0; } /** Returns true if the color was new and added. */ boolean addMapEntry(int r, int g, int b) { // Sentinel bit 24 ensures packed != 0 for any valid color. int key = (1 << 24) | (r << 16) | (g << 8) | b; // Fibonacci hashing for uniform distribution over HASH_BITS bits. int slot = (int)(((key & 0xFFFFFF) * 2654435761L) >>> (32 - HASH_BITS)) & HASH_MASK; while (slots[slot] != 0) { probes++; if (slots[slot] == key) return false; // already in map slot = (slot + 1) & HASH_MASK; } slots[slot] = key; if (ncolors >= 256) return false; R[ncolors] = r; G[ncolors] = g; B[ncolors] = b; ncolors++; return true; } } // ------------------------------------------------------------------------- // Helpers // ------------------------------------------------------------------------- /** * Simulate read_ppm_map: insert numPixels pixels where the first numUnique * are distinct colors and the rest repeat color (0,0,0). * Returns comparison count (unpatched) or probe count (patched). */ /** * Simulate read_ppm_map with P pixels randomly drawn from a palette of C colors. * Colors are pre-built as (0,i,0) for i=0..C-1 and placed in a random pixel order. * Each pixel lookup does an O(C) linear scan in the unpatched version. * * Expected unpatched comparisons per pixel: ~C/2 on average (found at middle) for * duplicate pixels, and 0..C-1 for the initial insertions. * For P pixels and palette size C: total ≈ P * C / 2. * * @param pixelColors pre-built array of P pixel colors (each an int = (r<<16)|(g<<8)|b) */ static long simulateUnpatched(int[] pixelColors) { UnpatchedColorMap cm = new UnpatchedColorMap(); for (int packed : pixelColors) { int r = (packed >> 16) & 0xFF; int g = (packed >> 8) & 0xFF; int b = packed & 0xFF; cm.addMapEntry(r, g, b); } return cm.comparisons; } static long simulatePatched(int[] pixelColors) { PatchedColorMap cm = new PatchedColorMap(); cm.reset(); for (int packed : pixelColors) { int r = (packed >> 16) & 0xFF; int g = (packed >> 8) & 0xFF; int b = packed & 0xFF; cm.addMapEntry(r, g, b); } return cm.probes; } /** Build a pixel array of numPixels pixels drawn from C distinct colors, * ordered so that the palette fills up first (best for showing O(P*C) growth). */ static int[] buildPixels(int numPixels, int numColors) { int[] pixels = new int[numPixels]; // First numColors pixels: each a distinct color (0, i, 0) for i=0..numColors-1 for (int i = 0; i < numColors && i < numPixels; i++) { pixels[i] = i & 0xFF; // packed as (0, 0, i) — b channel } // Remaining pixels: cycle through all colors (so each color appears P/C times) for (int i = numColors; i < numPixels; i++) { pixels[i] = (i % numColors) & 0xFF; } return pixels; } // ------------------------------------------------------------------------- // Tests // ------------------------------------------------------------------------- static void testCorrectness_smallImage() { UnpatchedColorMap u = new UnpatchedColorMap(); assert u.addMapEntry(255, 0, 0) : "red should be new"; assert u.addMapEntry(0, 255, 0) : "green should be new"; assert u.addMapEntry(0, 0, 255) : "blue should be new"; assert !u.addMapEntry(255, 0, 0) : "red should be duplicate"; assert !u.addMapEntry(0, 255, 0) : "green should be duplicate"; assert u.ncolors == 3 : "expected 3 colors"; PatchedColorMap p = new PatchedColorMap(); p.reset(); assert p.addMapEntry(255, 0, 0) : "red should be new"; assert p.addMapEntry(0, 255, 0) : "green should be new"; assert p.addMapEntry(0, 0, 255) : "blue should be new"; assert !p.addMapEntry(255, 0, 0) : "red should be duplicate"; assert !p.addMapEntry(0, 255, 0) : "green should be duplicate"; assert p.ncolors == 3 : "expected 3 colors"; System.out.println("testCorrectness_smallImage PASS"); } static void testCorrectness_duplicateColors() { UnpatchedColorMap u = new UnpatchedColorMap(); for (int i = 0; i < 100; i++) u.addMapEntry(42, 84, 168); assert u.ncolors == 1 : "100 identical pixels should yield 1 color"; PatchedColorMap p = new PatchedColorMap(); p.reset(); for (int i = 0; i < 100; i++) p.addMapEntry(42, 84, 168); assert p.ncolors == 1 : "100 identical pixels should yield 1 color"; System.out.println("testCorrectness_duplicateColors PASS"); } static void testCorrectness_paletteFull() { UnpatchedColorMap u = new UnpatchedColorMap(); for (int i = 0; i < 256; i++) { assert u.addMapEntry(i, 0, 0) : "color " + i + " should be new"; } assert !u.addMapEntry(255, 255, 0) : "257th color should be rejected (overflow)"; assert u.ncolors == 256 : "palette should be exactly 256"; PatchedColorMap p = new PatchedColorMap(); p.reset(); for (int i = 0; i < 256; i++) { assert p.addMapEntry(i, 0, 0) : "color " + i + " should be new"; } assert !p.addMapEntry(255, 255, 0) : "257th color should be rejected"; assert p.ncolors == 256 : "palette should be exactly 256"; System.out.println("testCorrectness_paletteFull PASS"); } static void testSpeedup_mediumImage() { // 10000 pixels, 256 distinct palette colors, cycling so each appears ~39 times. // Unpatched: each duplicate scans on average C/2 = 128 entries. // Total: ~256 insertions (avg 128 cmp each) + 9744 duplicates * 128 = ~1.27M cmp int pixels = 10_000; int colors = 256; int[] pix = buildPixels(pixels, colors); long unpatchedCmp = simulateUnpatched(pix); long patchedProbes = simulatePatched(pix); double ratio = (double) unpatchedCmp / Math.max(1, patchedProbes); System.out.printf("testSpeedup_medium: unpatched=%d cmp, patched=%d probes, ratio=%.1fx%n", unpatchedCmp, patchedProbes, ratio); assert unpatchedCmp > 500_000 : "unpatched should have >500K comparisons, got " + unpatchedCmp; assert patchedProbes < unpatchedCmp / 10 : "patched should be << unpatched"; System.out.println("testSpeedup_mediumImage PASS"); } static void testSpeedup_largeImage() { // 100000 pixels, 256 distinct palette colors. // Unpatched: ~100000 * 128 = ~12.8M comparisons vs near-zero probes for patched. int pixels = 100_000; int colors = 256; int[] pix = buildPixels(pixels, colors); long unpatchedCmp = simulateUnpatched(pix); long patchedProbes = simulatePatched(pix); double ratio = (double) unpatchedCmp / Math.max(1, patchedProbes); System.out.printf("testSpeedup_large: unpatched=%d cmp, patched=%d probes, ratio=%.1fx%n", unpatchedCmp, patchedProbes, ratio); assert ratio > 100.0 : "speedup ratio should be > 100x for large images, got " + ratio; System.out.println("testSpeedup_largeImage PASS"); } static void testHashNoCollision_allGrayscale() { PatchedColorMap p = new PatchedColorMap(); p.reset(); int added = 0; for (int i = 0; i < 256; i++) { if (p.addMapEntry(i, i, i)) added++; } assert added == 256 : "all 256 grayscale colors should be distinct, got " + added; // Second pass: all should now be duplicates int dups = 0; for (int i = 0; i < 256; i++) { if (!p.addMapEntry(i, i, i)) dups++; } assert dups == 256 : "all 256 should be detected as duplicates, got " + dups; System.out.println("testHashNoCollision_allGrayscale PASS"); } // ------------------------------------------------------------------------- // Main // ------------------------------------------------------------------------- public static void main(String[] args) { System.out.println("=== libjpeg-turbo-0001: rdcolmap add_map_entry O(P*C) -> O(P*1) ==="); testCorrectness_smallImage(); testCorrectness_duplicateColors(); testCorrectness_paletteFull(); testSpeedup_mediumImage(); testSpeedup_largeImage(); testHashNoCollision_allGrayscale(); System.out.println("ALL PASS"); } }