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

@@ -46,12 +46,14 @@ This file is part of the iText (R) project.
 import com.itextpdf.io.LogMessageConstant;
 import com.itextpdf.io.util.MessageFormatUtil;
 import com.itextpdf.kernel.PdfException;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
 
 import java.io.Serializable;
 import java.util.ArrayList;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 /**
  * Algorithm for construction {@link PdfPages} tree
@@ -62,9 +64,9 @@ class PdfPagesTree implements Serializable {
 
     private final int leafSize = 10;
 
-    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;
@@ -78,9 +80,9 @@ class PdfPagesTree implements Serializable {
      */
     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)
@@ -425,10 +427,9 @@ private void loadPage(int pageNum) {
         } 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);
@@ -486,4 +487,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> implements Serializable {
+        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

@@ -56,14 +56,11 @@ This file is part of the iText (R) project.
 import com.itextpdf.kernel.pdf.xobject.PdfFormXObject;
 import com.itextpdf.kernel.pdf.xobject.PdfImageXObject;
 import com.itextpdf.kernel.utils.CompareTool;
+import com.itextpdf.test.AssertUtil;
 import com.itextpdf.test.ExtendedITextTest;
 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;
@@ -73,6 +70,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 {
@@ -86,6 +87,25 @@ public static void setup() {
         createDestinationFolder(destinationFolder);
     }
 
+    @Test
+    public void hugeNumberOfPagesWithOnePageTest() throws IOException {
+         PdfDocument pdfDoc = new PdfDocument(new PdfReader(sourceFolder + "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(sourceFolder + "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

@@ -0,0 +1,87 @@
+/*
+    This file is part of the iText (R) project.
+    Copyright (c) 1998-2022 iText Group NV
+    Authors: iText Software.
+
+    This program is free software; you can redistribute it and/or modify
+    it under the terms of the GNU Affero General Public License version 3
+    as published by the Free Software Foundation with the addition of the
+    following permission added to Section 15 as permitted in Section 7(a):
+    FOR ANY PART OF THE COVERED WORK IN WHICH THE COPYRIGHT IS OWNED BY
+    ITEXT GROUP. ITEXT GROUP DISCLAIMS THE WARRANTY OF NON INFRINGEMENT
+    OF THIRD PARTY RIGHTS
+
+    This program is distributed in the hope that it will be useful, but
+    WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
+    or FITNESS FOR A PARTICULAR PURPOSE.
+    See the GNU Affero General Public License for more details.
+    You should have received a copy of the GNU Affero General Public License
+    along with this program; if not, see http://www.gnu.org/licenses or write to
+    the Free Software Foundation, Inc., 51 Franklin Street, Fifth Floor,
+    Boston, MA, 02110-1301 USA, or download the license from the following URL:
+    http://itextpdf.com/terms-of-use/
+
+    The interactive user interfaces in modified source and object code versions
+    of this program must display Appropriate Legal Notices, as required under
+    Section 5 of the GNU Affero General Public License.
+
+    In accordance with Section 7(b) of the GNU Affero General Public License,
+    a covered work must retain the producer line in every PDF that is created
+    or manipulated using iText.
+
+    You can be released from the requirements of the license by purchasing
+    a commercial license. Buying such a license is mandatory as soon as you
+    develop commercial activities involving the iText software without
+    disclosing the source code of your own applications.
+    These activities include: offering paid services to customers as an ASP,
+    serving PDFs on the fly in a web application, shipping iText with a closed
+    source product.
+
+    For more information, please contact iText Software Corp. at this
+    address: [email protected]
+ */
+package com.itextpdf.kernel.pdf;
+
+import com.itextpdf.kernel.pdf.PdfPagesTree.NullUnlimitedList;
+import com.itextpdf.test.ExtendedITextTest;
+import com.itextpdf.test.annotations.type.UnitTest;
+
+import org.junit.Assert;
+import org.junit.Test;
+import org.junit.experimental.categories.Category;
+
+@Category(UnitTest.class)
+public class PdfPagesTreeTest extends ExtendedITextTest {
+    @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());
+    }
+}