Provide a NullUnlimitedList implementation to prevent OOM exception while opening document with enormous page count

DEVSIX-7793

Author: introfog

kernel/src/main/java/com/itextpdf/kernel/pdf/PdfPagesTree.java

@@ -22,19 +22,20 @@ This file is part of the iText (R) project.
  */
 package com.itextpdf.kernel.pdf;
 
-import com.itextpdf.io.logs.IoLogMessageConstant;
 import com.itextpdf.commons.utils.MessageFormatUtil;
-import com.itextpdf.kernel.exceptions.PdfException;
+import com.itextpdf.io.logs.IoLogMessageConstant;
 import com.itextpdf.kernel.exceptions.KernelExceptionMessageConstant;
+import com.itextpdf.kernel.exceptions.PdfException;
 
+import java.util.ArrayList;
+import java.util.HashMap;
 import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
 import java.util.Set;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-import java.util.ArrayList;
-import java.util.List;
-
 /**
  * Algorithm for construction {@link PdfPages} tree
  */
@@ -44,9 +45,9 @@ class PdfPagesTree {
 
     private final int leafSize = DEFAULT_LEAF_SIZE;
 
-    private List<PdfIndirectReference> pageRefs;
+    private NullUnlimitedList<PdfIndirectReference> pageRefs;
     private List<PdfPages> parents;
-    private List<PdfPage> pages;
+    private NullUnlimitedList<PdfPage> pages;
     private PdfDocument document;
     private boolean generated = false;
     private PdfPages root;
@@ -60,9 +61,9 @@ class PdfPagesTree {
      */
     public PdfPagesTree(PdfCatalog pdfCatalog) {
         this.document = pdfCatalog.getDocument();
-        this.pageRefs = new ArrayList<>();
+        this.pageRefs = new NullUnlimitedList<>();
         this.parents = new ArrayList<>();
-        this.pages = new ArrayList<>();
+        this.pages = new NullUnlimitedList<>();
         if (pdfCatalog.getPdfObject().containsKey(PdfName.Pages)) {
             PdfDictionary pages = pdfCatalog.getPdfObject().getAsDictionary(PdfName.Pages);
             if (pages == null) {
@@ -451,10 +452,9 @@ private void loadPage(int pageNum, Set<PdfIndirectReference> processedParents) {
         } else {
             int from = parent.getFrom();
 
-            // Possible exception in case kids.getSize() < parent.getCount().
-            // In any case parent.getCount() has higher priority.
             // NOTE optimization? when we already found needed index
-            for (int i = 0; i < parent.getCount(); i++) {
+            final int pageCount = Math.min(parent.getCount(), kids.size());
+            for (int i = 0; i < pageCount; i++) {
                 PdfObject kid = kids.get(i, false);
                 if (kid instanceof PdfIndirectReference) {
                     pageRefs.set(from + i, (PdfIndirectReference) kid);
@@ -512,4 +512,99 @@ private void correctPdfPagesFromProperty(int index, int correction) {
             }
         }
     }
+
+    /**
+     * The class represents a list which allows null elements, but doesn't allocate a memory for them, in the rest of
+     * cases it behaves like usual {@link ArrayList} and should have the same complexity (because keys are unique
+     * integers, so collisions are impossible). Class doesn't implement {@code List} interface because it provides
+     * only methods which are in use in {@link PdfPagesTree} class.
+     *
+     * @param <T> elements of the list
+     */
+    static final class NullUnlimitedList<T> {
+        private final Map<Integer, T> map = new HashMap<>();
+        private int size = 0;
+
+        // O(1)
+        public void add(T element) {
+            if (element == null) {
+                size++;
+                return;
+            }
+            map.put(size++, element);
+        }
+
+        // In worth scenario O(n^2) but it is mostly impossible because keys shouldn't have
+        // collisions at all (they are integers). So in average should be O(n).
+        public void add(int index, T element) {
+            if (index < 0 || index > size) {
+                return;
+            }
+            size++;
+            // Shifts the element currently at that position (if any) and any
+            // subsequent elements to the right (adds one to their indices).
+            T previous = map.get(index);
+            for (int i = index + 1; i < size; i++) {
+                T currentToAdd = previous;
+                previous = map.get(i);
+                this.set(i, currentToAdd);
+            }
+
+            this.set(index, element);
+        }
+
+        // average O(1), worth O(n) (mostly impossible in case when keys are integers)
+        public T get(int index) {
+            return map.get(index);
+        }
+
+        // average O(1), worth O(n) (mostly impossible in case when keys are integers)
+        public void set(int index, T element) {
+            if (element == null) {
+                map.remove(index);
+            } else {
+                map.put(index, element);
+            }
+        }
+
+        // O(n)
+        public int indexOf(T element) {
+            if (element == null) {
+                for (int i = 0; i < size; i++) {
+                    if (!map.containsKey(i)) {
+                        return i;
+                    }
+                }
+                return -1;
+            }
+            for (Map.Entry<Integer, T> entry : map.entrySet()) {
+                if (element.equals(entry.getValue())) {
+                    return entry.getKey();
+                }
+            }
+            return -1;
+        }
+
+        // In worth scenario O(n^2) but it is mostly impossible because keys shouldn't have
+        // collisions at all (they are integers). So in average should be O(n).
+        public void remove(int index) {
+            if (index < 0 || index >= size) {
+                return;
+            }
+            map.remove(index);
+            // Shifts any subsequent elements to the left (subtracts one from their indices).
+            T previous = map.get(size - 1);
+            for (int i = size - 2; i >= index; i--) {
+                T current = previous;
+                previous = map.get(i);
+                this.set(i, current);
+            }
+            map.remove(--size);
+        }
+
+        // O(1)
+        public int size() {
+            return size;
+        }
+    }
 }

kernel/src/test/java/com/itextpdf/kernel/pdf/PdfPagesTest.java

@@ -22,13 +22,13 @@ This file is part of the iText (R) project.
  */
 package com.itextpdf.kernel.pdf;
 
-import com.itextpdf.io.logs.IoLogMessageConstant;
+import com.itextpdf.commons.utils.MessageFormatUtil;
 import com.itextpdf.io.image.ImageDataFactory;
+import com.itextpdf.io.logs.IoLogMessageConstant;
 import com.itextpdf.io.source.RandomAccessSourceFactory;
-import com.itextpdf.commons.utils.MessageFormatUtil;
-import com.itextpdf.kernel.exceptions.PdfException;
 import com.itextpdf.kernel.colors.ColorConstants;
 import com.itextpdf.kernel.exceptions.KernelExceptionMessageConstant;
+import com.itextpdf.kernel.exceptions.PdfException;
 import com.itextpdf.kernel.geom.PageSize;
 import com.itextpdf.kernel.geom.Rectangle;
 import com.itextpdf.kernel.pdf.annot.PdfAnnotation;
@@ -42,10 +42,6 @@ This file is part of the iText (R) project.
 import com.itextpdf.test.annotations.LogMessage;
 import com.itextpdf.test.annotations.LogMessages;
 import com.itextpdf.test.annotations.type.IntegrationTest;
-import org.junit.Assert;
-import org.junit.BeforeClass;
-import org.junit.Test;
-import org.junit.experimental.categories.Category;
 
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
@@ -55,6 +51,10 @@ This file is part of the iText (R) project.
 import java.util.List;
 import java.util.Random;
 import java.util.Set;
+import org.junit.Assert;
+import org.junit.BeforeClass;
+import org.junit.Test;
+import org.junit.experimental.categories.Category;
 
 @Category(IntegrationTest.class)
 public class PdfPagesTest extends ExtendedITextTest {
@@ -67,6 +67,25 @@ public static void setup() {
         createDestinationFolder(DESTINATION_FOLDER);
     }
 
+    @Test
+    public void hugeNumberOfPagesWithOnePageTest() throws IOException {
+         PdfDocument pdfDoc = new PdfDocument(new PdfReader(SOURCE_FOLDER + "hugeNumberOfPagesWithOnePage.pdf"),
+                 new PdfWriter(new ByteArrayOutputStream()));
+         PdfPage page = new PdfPage(pdfDoc, pdfDoc.getDefaultPageSize());
+         AssertUtil.doesNotThrow(() -> pdfDoc.addPage(1, page));
+    }
+
+    @Test
+    public void countDontCorrespondToRealTest() throws IOException {
+        PdfDocument pdfDoc = new PdfDocument(new PdfReader(SOURCE_FOLDER + "countDontCorrespondToReal.pdf"),
+                new PdfWriter(new ByteArrayOutputStream()));
+        PdfPage page = new PdfPage(pdfDoc, pdfDoc.getDefaultPageSize());
+        AssertUtil.doesNotThrow(() -> pdfDoc.addPage(1, page));
+
+        // we don't expect that Count will be different from real number of pages
+        Assert.assertThrows(NullPointerException.class, () -> pdfDoc.close());
+    }
+
     @Test
     public void simplePagesTest() throws IOException {
         String filename = "simplePagesTest.pdf";

kernel/src/test/java/com/itextpdf/kernel/pdf/PdfPagesTreeTest.java

@@ -23,14 +23,11 @@ This file is part of the iText (R) project.
 package com.itextpdf.kernel.pdf;
 
 import com.itextpdf.io.source.ByteArrayOutputStream;
-import com.itextpdf.kernel.events.PdfDocumentEvent;
-import com.itextpdf.kernel.pdf.layer.PdfLayer;
-import com.itextpdf.kernel.utils.CompareTool;
+import com.itextpdf.kernel.pdf.PdfPagesTree.NullUnlimitedList;
 import com.itextpdf.test.AssertUtil;
 import com.itextpdf.test.ExtendedITextTest;
 import com.itextpdf.test.annotations.type.UnitTest;
 
-import java.io.IOException;
 import org.junit.Assert;
 import org.junit.Test;
 import org.junit.experimental.categories.Category;
@@ -42,4 +39,37 @@ public void generateTreeDocHasNoPagesTest() {
         PdfDocument pdfDoc = new PdfDocument(new PdfWriter(new ByteArrayOutputStream()));
         AssertUtil.doesNotThrow(() -> pdfDoc.close());
     }
+
+    @Test
+    public void nullUnlimitedListAddTest() {
+        NullUnlimitedList<String> list = new NullUnlimitedList<>();
+        list.add("hey");
+        list.add("bye");
+        Assert.assertEquals(2, list.size());
+        list.add(-1, "hello");
+        list.add(3, "goodbye");
+        Assert.assertEquals(2, list.size());
+    }
+
+    @Test
+    public void nullUnlimitedListIndexOfTest() {
+        NullUnlimitedList<String> list = new NullUnlimitedList<>();
+        list.add("hey");
+        list.add(null);
+        list.add("bye");
+        list.add(null);
+        Assert.assertEquals(4, list.size());
+        Assert.assertEquals(1, list.indexOf(null));
+    }
+
+    @Test
+    public void nullUnlimitedListRemoveTest() {
+        NullUnlimitedList<String> list = new NullUnlimitedList<>();
+        list.add("hey");
+        list.add("bye");
+        Assert.assertEquals(2, list.size());
+        list.remove(-1);
+        list.remove(2);
+        Assert.assertEquals(2, list.size());
+    }
 }