add node depth limit

Author: spirit-at-canva

commonmark/src/main/java/org/commonmark/node/AbstractVisitor.java

@@ -7,6 +7,19 @@
  * call {@link #visitChildren}.
  */
 public abstract class AbstractVisitor implements Visitor {
+    private final int maxDepth;
+    private int currentDepth;
+
+    public AbstractVisitor() {
+        this(Integer.MAX_VALUE);
+    }
+
+    protected AbstractVisitor(int maxDepth) {
+        if (maxDepth < 0) {
+            throw new IllegalArgumentException("maxDepth must be >= 0");
+        }
+        this.maxDepth = maxDepth;
+    }
 
     @Override
     public void visit(BlockQuote blockQuote) {
@@ -129,12 +142,20 @@ public void visit(CustomNode customNode) {
      * @param parent the parent node whose children should be visited
      */
     protected void visitChildren(Node parent) {
+        if (currentDepth >= maxDepth) {
+            return;
+        }
         Node node = parent.getFirstChild();
         while (node != null) {
             // A subclass of this visitor might modify the node, resulting in getNext returning a different node or no
             // node after visiting it. So get the next node before visiting.
             Node next = node.getNext();
-            node.accept(this);
+            currentDepth++;
+            try {
+                node.accept(this);
+            } finally {
+                currentDepth--;
+            }
             node = next;
         }
     }

commonmark/src/test/java/org/commonmark/test/AbstractVisitorTest.java

@@ -3,9 +3,38 @@
 import org.commonmark.node.*;
 import org.junit.jupiter.api.Test;
 
+import java.util.ArrayList;
+import java.util.List;
+
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
 
 public class AbstractVisitorTest {
+    @Test
+    public void maxDepthMustBeZeroOrGreater() {
+        assertThatThrownBy(() -> new RecordingVisitor(-1))
+                .isInstanceOf(IllegalArgumentException.class);
+    }
+
+    @Test
+    public void maxDepthZeroVisitsOnlyRoot() {
+        var paragraph = paragraphTree();
+        var visitor = new RecordingVisitor(0);
+
+        paragraph.accept(visitor);
+
+        assertThat(visitor.visited).containsExactly("paragraph");
+    }
+
+    @Test
+    public void maxDepthOneVisitsDirectChildrenButNotGrandchildren() {
+        var paragraph = paragraphTree();
+        var visitor = new RecordingVisitor(1);
+
+        paragraph.accept(visitor);
+
+        assertThat(visitor.visited).containsExactly("paragraph", "emphasis", "text:tail");
+    }
 
     @Test
     public void replacingNodeInVisitorShouldNotDestroyVisitOrder() {
@@ -34,4 +63,39 @@ private static void assertCode(String expectedLiteral, Node node) {
         Code code = (Code) node;
         assertThat(code.getLiteral()).isEqualTo(expectedLiteral);
     }
+
+    private static Paragraph paragraphTree() {
+        var paragraph = new Paragraph();
+        var emphasis = new Emphasis();
+        emphasis.appendChild(new Text("nested"));
+        paragraph.appendChild(emphasis);
+        paragraph.appendChild(new Text("tail"));
+        return paragraph;
+    }
+
+    private static final class RecordingVisitor extends AbstractVisitor {
+        private final List<String> visited = new ArrayList<>();
+
+        private RecordingVisitor(int maxDepth) {
+            super(maxDepth);
+        }
+
+        @Override
+        public void visit(Paragraph paragraph) {
+            visited.add("paragraph");
+            super.visit(paragraph);
+        }
+
+        @Override
+        public void visit(Emphasis emphasis) {
+            visited.add("emphasis");
+            super.visit(emphasis);
+        }
+
+        @Override
+        public void visit(Text text) {
+            visited.add("text:" + text.getLiteral());
+            super.visit(text);
+        }
+    }
 }