Repository: Hopding/pdf-lib
Pull request #441 (2020-05-09 20:36:34)

Fix page removal bug
Addresses #440


Patch:
diff --git a/src/core/errors.ts b/src/core/errors.ts
index 207d61c..146bd7e 100644
--- a/src/core/errors.ts
+++ b/src/core/errors.ts
@@ -85,6 +85,20 @@ export class InvalidPDFDateStringError extends Error {
   }
 }
 
+export class InvalidTargetIndexError extends Error {
+  constructor(targetIndex: number, Count: number) {
+    const msg = `Invalid targetIndex specified: targetIndex=${targetIndex} must be less than Count=${Count}`;
+    super(msg);
+  }
+}
+
+export class CorruptPageTreeError extends Error {
+  constructor(targetIndex: number, operation: string) {
+    const msg = `Failed to ${operation} at targetIndex=${targetIndex} due to corrupt page tree: It is likely that one or more 'Count' entries are invalid`;
+    super(msg);
+  }
+}
+
 /***** Parser Errors ******/
 
 export interface Position {
diff --git a/src/core/structures/PDFPageTree.ts b/src/core/structures/PDFPageTree.ts
index 601275f..a8368dc 100644
--- a/src/core/structures/PDFPageTree.ts
+++ b/src/core/structures/PDFPageTree.ts
@@ -5,6 +5,7 @@ import PDFNumber from 'src/core/objects/PDFNumber';
 import PDFRef from 'src/core/objects/PDFRef';
 import PDFContext from 'src/core/PDFContext';
 import PDFPageLeaf from 'src/core/structures/PDFPageLeaf';
+import { InvalidTargetIndexError, CorruptPageTreeError } from 'src/core/errors';
 
 export type TreeNode = PDFPageTree | PDFPageLeaf;
 
@@ -40,11 +41,7 @@ class PDFPageTree extends PDFDict {
 
   pushLeafNode(leafRef: PDFRef): void {
     const Kids = this.Kids();
-    Kids.push(leafRef);
-    this.ascend((node) => {
-      const Count = node.Count();
-      node.set(PDFName.of('Count'), PDFNumber.of(Count.asNumber() + 1));
-    });
+    this.insertLeafKid(Kids.size(), leafRef);
   }
 
   /**
@@ -58,85 +55,98 @@ class PDFPageTree extends PDFDict {
    */
   insertLeafNode(leafRef: PDFRef, targetIndex: number): PDFRef | undefined {
     const Kids = this.Kids();
-    const kidSize = Kids.size();
+    const Count = this.Count().asNumber();
+
+    if (targetIndex > Count) {
+      throw new InvalidTargetIndexError(targetIndex, Count);
+    }
 
-    let kidIdx = 0;
-    let pageIdx = 0;
-    while (pageIdx < targetIndex) {
-      if (kidIdx >= kidSize) {
-        throw new Error(`Index out of bounds: ${kidIdx}/${kidSize}`);
+    let leafsRemainingUntilTarget = targetIndex;
+    for (let idx = 0, len = Kids.size(); idx < len; idx++) {
+      if (leafsRemainingUntilTarget === 0) {
+        // Insert page and return
+        this.insertLeafKid(idx, leafRef);
+        return undefined;
       }
 
-      const kidRef = Kids.get(kidIdx++) as PDFRef;
-      const kid = this.context.lookup(kidRef) as TreeNode;
+      const kidRef = Kids.get(idx) as PDFRef;
+      const kid = this.context.lookup(kidRef);
 
       if (kid instanceof PDFPageTree) {
-        const kidCount = kid.Count().asNumber();
-        if (pageIdx + kidCount > targetIndex) {
-          return kid.insertLeafNode(leafRef, targetIndex - pageIdx) || kidRef;
+        if (kid.Count().asNumber() > leafsRemainingUntilTarget) {
+          // Dig in
+          return (
+            kid.insertLeafNode(leafRef, leafsRemainingUntilTarget) || kidRef
+          );
         } else {
-          pageIdx += kidCount;
+          // Move on
+          leafsRemainingUntilTarget -= kid.Count().asNumber();
         }
-      } else {
-        pageIdx += 1;
+      }
+
+      if (kid instanceof PDFPageLeaf) {
+        // Move on
+        leafsRemainingUntilTarget -= 1;
       }
     }
 
-    Kids.insert(kidIdx, leafRef);
-    this.ascend((node) => {
-      const Count = node.Count();
-      node.set(PDFName.of('Count'), PDFNumber.of(Count.asNumber() + 1));
-    });
+    if (leafsRemainingUntilTarget === 0) {
+      // Insert page at the end and return
+      this.insertLeafKid(Kids.size(), leafRef);
+      return undefined;
+    }
 
-    return undefined;
+    // Should never get here if `targetIndex` is valid
+    throw new CorruptPageTreeError(targetIndex, 'insertLeafNode');
   }
 
   /**
    * Removes the leaf node at the specified index (zero-based) from this page
    * tree. Also decrements the `Count` of each page tree in the hierarchy to
    * account for the removed page.
+   *
+   * If `prune` is true, then intermediate tree nodes will be removed from the
+   * tree if they contain 0 children after the leaf node is removed.
    */
-  removeLeafNode(targetIndex: number): void {
+  removeLeafNode(targetIndex: number, prune = true): void {
     const Kids = this.Kids();
-    const kidSize = Kids.size();
+    const Count = this.Count().asNumber();
 
-    let kidIdx = 0;
-    let pageIdx = 0;
-    while (pageIdx < targetIndex) {
-      if (kidIdx >= kidSize) {
-        throw new Error(`Index out of bounds: ${kidIdx}/${kidSize - 1} (a)`);
-      }
+    if (targetIndex >= Count) {
+      throw new InvalidTargetIndexError(targetIndex, Count);
+    }
 
-      const kidRef = Kids.get(kidIdx++) as PDFRef;
-      const kid = this.context.lookup(kidRef) as TreeNode;
+    let leafsRemainingUntilTarget = targetIndex;
+    for (let idx = 0, len = Kids.size(); idx < len; idx++) {
+      const kidRef = Kids.get(idx) as PDFRef;
+      const kid = this.context.lookup(kidRef);
 
       if (kid instanceof PDFPageTree) {
-        const kidCount = kid.Count().asNumber();
-        if (pageIdx + kidCount > targetIndex) {
-          kid.removeLeafNode(targetIndex - pageIdx);
+        if (kid.Count().asNumber() > leafsRemainingUntilTarget) {
+          // Dig in
+          kid.removeLeafNode(leafsRemainingUntilTarget, prune);
+          if (prune && kid.Kids().size() === 0) Kids.remove(idx);
           return;
         } else {
-          pageIdx += kidCount;
+          // Move on
+          leafsRemainingUntilTarget -= kid.Count().asNumber();
         }
-      } else {
-        pageIdx += 1;
       }
-    }
 
-    if (kidIdx >= kidSize) {
-      throw new Error(`Index out of bounds: ${kidIdx}/${kidSize - 1} (b)`);
+      if (kid instanceof PDFPageLeaf) {
+        if (leafsRemainingUntilTarget === 0) {
+          // Remove page and return
+          this.removeKid(idx);
+          return;
+        } else {
+          // Move on
+          leafsRemainingUntilTarget -= 1;
+        }
+      }
     }
 
-    const target = Kids.lookup(kidIdx);
-    if (target instanceof PDFPageTree) {
-      target.removeLeafNode(0);
-    } else {
-      Kids.remove(kidIdx);
-      this.ascend((node) => {
-        const Count = node.Count();
-        node.set(PDFName.of('Count'), PDFNumber.of(Count.asNumber() - 1));
-      });
-    }
+    // Should never get here if `targetIndex` is valid
+    throw new CorruptPageTreeError(targetIndex, 'removeLeafNode');
   }
 
   ascend(visitor: (node: PDFPageTree) => any): void {
@@ -155,6 +165,31 @@ class PDFPageTree extends PDFDict {
       visitor(kid, kidRef);
     }
   }
+
+  private insertLeafKid(kidIdx: number, leafRef: PDFRef): void {
+    const Kids = this.Kids();
+
+    this.ascend((node) => {
+      const newCount = node.Count().asNumber() + 1;
+      node.set(PDFName.of('Count'), PDFNumber.of(newCount));
+    });
+
+    Kids.insert(kidIdx, leafRef);
+  }
+
+  private removeKid(kidIdx: number): void {
+    const Kids = this.Kids();
+
+    const kid = Kids.lookup(kidIdx);
+    if (kid instanceof PDFPageLeaf) {
+      this.ascend((node) => {
+        const newCount = node.Count().asNumber() - 1;
+        node.set(PDFName.of('Count'), PDFNumber.of(newCount));
+      });
+    }
+
+    Kids.remove(kidIdx);
+  }
 }
 
 export default PDFPageTree;


Test patch:
diff --git a/apps/node/tests/test3.ts b/apps/node/tests/test3.ts
index 1ed0084..2af8939 100644
--- a/apps/node/tests/test3.ts
+++ b/apps/node/tests/test3.ts
@@ -66,8 +66,8 @@ export default async (assets: Assets) => {
     color: hotPink,
   });
   page2.drawLine({
-    start: { x: 30, y: 50 },
-    end: { x: 30 + lastPageTextWidth, y: 50 },
+    start: { x: 30, y: 205 },
+    end: { x: 30 + lastPageTextWidth, y: 205 },
     color: hotPink,
     thickness: 5,
   });
diff --git a/apps/rn/src/tests/test3.js b/apps/rn/src/tests/test3.js
index afc78df..4bf0d9d 100644
--- a/apps/rn/src/tests/test3.js
+++ b/apps/rn/src/tests/test3.js
@@ -62,8 +62,8 @@ export default async () => {
     color: hotPink,
   });
   page2.drawLine({
-    start: { x: 30, y: 50 },
-    end: { x: 30 + lastPageTextWidth, y: 50 },
+    start: { x: 30, y: 205 },
+    end: { x: 30 + lastPageTextWidth, y: 205 },
     color: hotPink,
     thickness: 5,
   });
diff --git a/apps/web/test3.html b/apps/web/test3.html
index e386027..9a2ebd4 100644
--- a/apps/web/test3.html
+++ b/apps/web/test3.html
@@ -116,8 +116,8 @@
         color: hotPink,
       });
       page2.drawLine({
-        start: { x: 30, y: 50 },
-        end: { x: 30 + lastPageTextWidth, y: 50 },
+        start: { x: 30, y: 205 },
+        end: { x: 30 + lastPageTextWidth, y: 205 },
         color: hotPink,
         thickness: 5,
       });
diff --git a/tests/core/structures/PDFPageTree.spec.ts b/tests/core/structures/PDFPageTree.spec.ts
index 7ce9e1f..0a70160 100644
--- a/tests/core/structures/PDFPageTree.spec.ts
+++ b/tests/core/structures/PDFPageTree.spec.ts
@@ -382,7 +382,7 @@ describe(`PDFPageTree`, () => {
       expect(lvl2Tree1.Kids().get(1)).toBe(lvl3Tree1Ref);
       expect(lvl4Tree1.Kids().get(0)).toBe(leafRef3);
 
-      lvl1Tree1.removeLeafNode(2);
+      lvl1Tree1.removeLeafNode(2, false);
 
       expect(lvl1Tree1.Count().asNumber()).toBe(3);
       expect(lvl2Tree1.Count().asNumber()).toBe(1);
@@ -433,10 +433,10 @@ describe(`PDFPageTree`, () => {
       expect(() => buildTree().removeLeafNode(2)).not.toThrow();
 
       expect(() => buildTree().removeLeafNode(3)).toThrow(
-        'Index out of bounds: 3/2 (b)',
+        'Invalid targetIndex specified: targetIndex=3 must be less than Count=3',
       );
       expect(() => buildTree().removeLeafNode(4)).toThrow(
-        'Index out of bounds: 3/2 (a)',
+        'Invalid targetIndex specified: targetIndex=4 must be less than Count=3',
       );
     });
 
@@ -483,11 +483,90 @@ describe(`PDFPageTree`, () => {
       expect(() => buildTree().removeLeafNode(7)).not.toThrow();
 
       expect(() => buildTree().removeLeafNode(8)).toThrow(
-        'Index out of bounds: 2/1 (b)',
+        'Invalid targetIndex specified: targetIndex=8 must be less than Count=8',
       );
       expect(() => buildTree().removeLeafNode(9)).toThrow(
-        'Index out of bounds: 2/1 (a)',
+        'Invalid targetIndex specified: targetIndex=9 must be less than Count=8',
+      );
+    });
+
+    it(`throws an error when the page tree has no kids`, () => {
+      const buildTree = () => {
+        const { pageTree } = pageUtils();
+        const [, lvl1Tree1] = pageTree();
+        return lvl1Tree1;
+      };
+
+      expect(() => buildTree().removeLeafNode(0)).toThrow(
+        'Invalid targetIndex specified: targetIndex=0 must be less than Count=0',
+      );
+      expect(() => buildTree().removeLeafNode(1)).toThrow(
+        'Invalid targetIndex specified: targetIndex=1 must be less than Count=0',
+      );
+    });
+
+    it(`throws an error when the page tree is invalid`, () => {
+      const buildTree = () => {
+        const { pageTree } = pageUtils();
+        const [, lvl1Tree1] = pageTree();
+        return lvl1Tree1;
+      };
+
+      const tree = buildTree();
+      tree.set(PDFName.of('Count'), PDFNumber.of(21));
+
+      expect(() => tree.removeLeafNode(0)).toThrow(
+        `Failed to removeLeafNode at targetIndex=0 due to corrupt page tree: It is likely that one or more 'Count' entries are invalid`,
+      );
+      expect(() => tree.removeLeafNode(1)).toThrow(
+        `Failed to removeLeafNode at targetIndex=1 due to corrupt page tree: It is likely that one or more 'Count' entries are invalid`,
       );
     });
+
+    it(`removes tree nodes that have 0 children after a leaf node is removed`, () => {
+      const { pageTree, pageLeaf } = pageUtils();
+
+      const [tree1Ref, tree1] = pageTree();
+      const [tree2Ref, tree2] = pageTree(tree1Ref);
+      const [tree3Ref, tree3] = pageTree(tree1Ref);
+      const [tree4Ref, tree4] = pageTree(tree1Ref);
+      const [tree5Ref, tree5] = pageTree(tree3Ref);
+
+      const [leafRef1] = pageLeaf(tree2Ref);
+      const [leafRef2] = pageLeaf(tree4Ref);
+      const [leafRef3] = pageLeaf(tree5Ref);
+      const [leafRef4] = pageLeaf(tree5Ref);
+
+      pushTreeNodes(tree1, tree2Ref, tree3Ref, tree4Ref);
+      pushTreeNodes(tree3, tree5Ref);
+
+      pushLeafNodes(tree2, leafRef1);
+      pushLeafNodes(tree4, leafRef2);
+      pushLeafNodes(tree5, leafRef3, leafRef4);
+
+      expect(tree1.Kids().size()).toBe(3);
+      expect(tree1.Kids().get(0)).toBe(tree2Ref);
+      expect(tree1.Kids().get(1)).toBe(tree3Ref);
+      expect(tree1.Kids().get(2)).toBe(tree4Ref);
+
+      expect(tree3.Kids().size()).toBe(1);
+      expect(tree3.Kids().get(0)).toBe(tree5Ref);
+
+      tree1.removeLeafNode(1);
+
+      expect(tree1.Kids().size()).toBe(3);
+      expect(tree1.Kids().get(0)).toBe(tree2Ref);
+      expect(tree1.Kids().get(1)).toBe(tree3Ref);
+      expect(tree1.Kids().get(2)).toBe(tree4Ref);
+
+      expect(tree3.Kids().size()).toBe(1);
+      expect(tree3.Kids().get(0)).toBe(tree5Ref);
+
+      tree1.removeLeafNode(1);
+
+      expect(tree1.Kids().size()).toBe(2);
+      expect(tree1.Kids().get(0)).toBe(tree2Ref);
+      expect(tree1.Kids().get(1)).toBe(tree4Ref);
+    });
   });
 });