Allow PdfStreams as a referenced object for MCR

DEVSIX-8283

Author: Alexander Pliushchou

kernel/src/main/java/com/itextpdf/kernel/exceptions/KernelExceptionMessageConstant.java

@@ -251,6 +251,8 @@ public final class KernelExceptionMessageConstant {
             "Invalid PDF Type 3 Function object, \"Functions\" array should exist and can't be empty.";
     public static final String INVALID_RANGE_ARRAY = "Invalid range array.";
     public static final String INVALID_OFFSET_FOR_THIS_OBJECT = "Invalid offset for object {0}.";
+    public static final String INVALID_OBJECT_REFERENCE_TYPE = "Object reference has unsupported type, " +
+            "supported types are dictionaries and streams";
     public static final String INVALID_XREF_STREAM = "Invalid xref stream.";
     public static final String INVALID_XREF_TABLE = "Invalid xref table.";
     public static final String IO_EXCEPTION = "I/O exception.";

kernel/src/main/java/com/itextpdf/kernel/pdf/tagging/ParentTreeHandler.java

@@ -197,8 +197,13 @@ private void registerMcr(PdfMcr mcr, boolean registeringOnInit) {
             }
             pageMcrs.putXObjectMcr(stmIndRef, mcr);
         } else if (mcr instanceof PdfObjRef) {
-            PdfDictionary obj = ((PdfDictionary) mcr.getPdfObject()).getAsDictionary(PdfName.Obj);
-            if (obj == null || obj.isFlushed()) {
+            PdfObject mcrObj = ((PdfDictionary) mcr.getPdfObject()).get(PdfName.Obj);
+            if (!(mcrObj instanceof PdfDictionary)) {
+                throw new PdfException(KernelExceptionMessageConstant.INVALID_OBJECT_REFERENCE_TYPE);
+            }
+
+            PdfDictionary obj = (PdfDictionary) mcrObj;
+            if (obj.isFlushed()) {
                 throw new PdfException(
                         KernelExceptionMessageConstant.WHEN_ADDING_OBJECT_REFERENCE_TO_THE_TAG_TREE_IT_MUST_BE_CONNECTED_TO_NOT_FLUSHED_OBJECT);
             }
@@ -247,15 +252,6 @@ public void unregisterMcr(PdfMcr mcrToUnregister) {
                 }
                 structTreeRoot.setModified();
             } else if (mcrToUnregister instanceof PdfObjRef) {
-                PdfDictionary obj = ((PdfDictionary) mcrToUnregister.getPdfObject()).getAsDictionary(PdfName.Obj);
-                if (obj != null && !obj.isFlushed()) {
-                    PdfNumber n = obj.getAsNumber(PdfName.StructParent);
-                    if (n != null) {
-                        pageMcrs.getObjRefs().remove(n.intValue());
-                        structTreeRoot.setModified();
-                        return;
-                    }
-                }
                 for (Map.Entry<Integer, PdfMcr> entry : pageMcrs.getObjRefs().entrySet()) {
                     if (entry.getValue().getPdfObject() == mcrToUnregister.getPdfObject()) {
                         pageMcrs.getObjRefs().remove(entry.getKey());

kernel/src/main/java/com/itextpdf/kernel/pdf/tagging/PdfObjRef.java

@@ -25,6 +25,7 @@ This file is part of the iText (R) project.
 import com.itextpdf.kernel.pdf.PdfDictionary;
 import com.itextpdf.kernel.pdf.PdfName;
 import com.itextpdf.kernel.pdf.PdfNumber;
+import com.itextpdf.kernel.pdf.PdfObject;
 import com.itextpdf.kernel.pdf.annot.PdfAnnotation;
 
 public class PdfObjRef extends PdfMcr {
@@ -55,6 +56,11 @@ public PdfDictionary getPageObject() {
     }
 
     public PdfDictionary getReferencedObject() {
-        return ((PdfDictionary) getPdfObject()).getAsDictionary(PdfName.Obj);
+        PdfObject obj = ((PdfDictionary) getPdfObject()).get(PdfName.Obj);
+        if(obj instanceof PdfDictionary){
+            return (PdfDictionary) obj;
+        } else {
+            return null;
+        }
     }
 }

kernel/src/main/java/com/itextpdf/kernel/pdf/tagging/StructureTreeCopier.java

@@ -354,13 +354,14 @@ private static PdfDictionary copyObject(PdfDictionary source, PdfDictionary dest
         } else {
             copied = source.copyTo(copyingParams.getToDocument(), ignoreKeysForCopy, true);
 
-            PdfDictionary obj = source.getAsDictionary(PdfName.Obj);
-            if (obj != null) {
+            PdfObject obj = source.get(PdfName.Obj);
+            if (obj instanceof PdfDictionary) {
+                PdfDictionary objDic = (PdfDictionary) obj;
                 // Link annotations could be not added to the toDocument, so we need to identify this case.
                 // When obj.copyTo is called, and annotation was already copied, we would get this already created copy.
                 // If it was already copied and added, /P key would be set. Otherwise /P won't be set.
-                obj = obj.copyTo(copyingParams.getToDocument(), Arrays.asList(PdfName.P), false);
-                copied.put(PdfName.Obj, obj);
+                objDic = objDic.copyTo(copyingParams.getToDocument(), Arrays.asList(PdfName.P), false);
+                copied.put(PdfName.Obj, objDic);
             }
 
             PdfDictionary nsDict = source.getAsDictionary(PdfName.NS);
@@ -456,7 +457,7 @@ private static PdfObject copyObjectKid(PdfObject kid, PdfDictionary copiedParent
                     PdfMcr mcr;
                     if (copiedKid.containsKey(PdfName.Obj)) {
                         mcr = new PdfObjRef(copiedKid, new PdfStructElem(copiedParent));
-                        PdfDictionary contentItemObject = copiedKid.getAsDictionary(PdfName.Obj);
+                        PdfDictionary contentItemObject = (PdfDictionary) copiedKid.get(PdfName.Obj);
                         if (PdfName.Link.equals(contentItemObject.getAsName(PdfName.Subtype))
                                 && !contentItemObject.containsKey(PdfName.P)) {
                             // Some link annotations may be not copied, because their destination page is not copied.

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

@@ -358,6 +358,44 @@ public void test06() throws IOException {
         assertTrue(checkParentTree(outFile, cmpFile));
     }
 
+    @Test
+    public void objRefAsStreamTest() throws IOException, InterruptedException {
+        String pdf = sourceFolder + "objRefAsStream.pdf";
+        String outPdf = destinationFolder + "objRefAsStream.pdf";
+        String cmpPdf = sourceFolder + "cmp_objRefAsStream.pdf";
+
+        PdfDocument taggedPdf = new PdfDocument(new PdfReader(pdf), CompareTool.createTestPdfWriter(outPdf));
+        taggedPdf.close();
+
+        Assert.assertNull(new CompareTool().compareByContent(outPdf, cmpPdf, destinationFolder, "diff"));
+    }
+
+    @Test
+    @LogMessages(messages = {
+            @LogMessage(messageTemplate = IoLogMessageConstant.TAG_STRUCTURE_INIT_FAILED)
+    })
+    public void objRefAsInvalidType() throws IOException {
+        String pdf = sourceFolder + "objRefAsInvalidType.pdf";
+        PdfDocument doc = new PdfDocument(new PdfReader(pdf));
+        Assert.assertNull(doc.getStructTreeRoot());
+    }
+
+    @Test
+    public void unregisterObjRefAsStreamTest() throws IOException, InterruptedException {
+        String pdf = sourceFolder + "objRefAsStream.pdf";
+        String outPdf = destinationFolder + "objRefAsStreamUnregisterMcr.pdf";
+        String cmpPdf = sourceFolder + "cmp_objRefAsStreamUnregisterMcr.pdf";
+
+        PdfDocument taggedPdf = new PdfDocument(new PdfReader(pdf), CompareTool.createTestPdfWriter(outPdf));
+
+        PdfStructElem elem = (PdfStructElem) taggedPdf.getStructTreeRoot().getKids().get(0).getKids().get(0);
+        elem.removeKid(0);
+        taggedPdf.close();
+
+        Assert.assertNull(new CompareTool().compareByContent(outPdf, cmpPdf, destinationFolder, "diff"));
+    }
+
+
     @Test
     @LogMessages(messages = {
             @LogMessage(messageTemplate = KernelLogMessageConstant.STRUCT_PARENT_INDEX_MISSED_AND_RECREATED, count = 4)

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

@@ -383,6 +383,23 @@ public void copyPagesLinkAnnotationTest() throws IOException, InterruptedExcepti
         Assert.assertNull(new CompareTool().compareByContent(outFileName, cmpFileName, destinationFolder));
     }
 
+    @Test
+    public void objRefAsStreamCopyTest() throws IOException, InterruptedException {
+        String pdf = sourceFolder + "objRefAsStream.pdf";
+        String outPdf = destinationFolder + "objRefAsStreamCopy.pdf";
+        String cmpPdf = sourceFolder + "cmp_objRefAsStreamCopy.pdf";
+
+        PdfDocument pdfFile = new PdfDocument(new PdfReader(pdf));
+        PdfDocument copiedFile = new PdfDocument(CompareTool.createTestPdfWriter(outPdf));
+        copiedFile.setTagged();
+        pdfFile.copyPagesTo(1, 1, copiedFile);
+
+        pdfFile.close();
+        copiedFile.close();
+
+        Assert.assertNull(new CompareTool().compareByContent(outPdf, cmpPdf, destinationFolder, "diff"));
+    }
+
     @Test
     public void copyDocWithFullDDictionary() throws IOException, InterruptedException {
         String outFileName = destinationFolder + "copyDocWithDDictionary.pdf";

kernel/src/test/java/com/itextpdf/kernel/pdf/tagging/PdfObjRefUnitTest.java

@@ -0,0 +1,36 @@
+package com.itextpdf.kernel.pdf.tagging;
+
+import com.itextpdf.kernel.pdf.*;
+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 PdfObjRefUnitTest extends ExtendedITextTest {
+
+    @Test
+    public void refObjAsStreamTest() {
+       PdfDictionary ref = new PdfStream();
+       ref.put(PdfName.Name, new PdfString("reference"));
+
+       PdfDictionary obj = new PdfDictionary();
+       obj.put(PdfName.Obj, ref);
+
+       PdfObjRef objRef = new PdfObjRef(obj, new PdfStructElem(new PdfDictionary()));
+
+       Assert.assertTrue(objRef.getReferencedObject() instanceof PdfStream);
+       Assert.assertTrue(objRef.getReferencedObject().containsKey(PdfName.Name));
+    }
+
+    @Test
+    public void refObjAsInvalidTypeTest() {
+        PdfDictionary obj = new PdfDictionary();
+        obj.put(PdfName.Obj, new PdfString("incorrect type"));
+
+        PdfObjRef objRef = new PdfObjRef(obj, new PdfStructElem(new PdfDictionary()));
+
+        Assert.assertNull(objRef.getReferencedObject());
+    }
+}