Infinispan: circular reference on marshalling leads to dos

@@ -5,6 +5,7 @@
 import org.infinispan.protostream.ProtobufTagMarshaller;
 import org.infinispan.protostream.impl.BaseMarshallerDelegate;
 import org.infinispan.protostream.impl.ByteArrayOutputStreamEx;
+import org.infinispan.protostream.impl.Log;
 import org.infinispan.protostream.impl.TagWriterImpl;
 
 /**
@@ -16,6 +17,8 @@
 @SuppressWarnings("unused")
 public class GeneratedMarshallerBase {
 
+   private static final Log log = Log.LogFactory.getLog(GeneratedMarshallerBase.class);
+
    /**
     * Invoked by generated code.
     */
@@ -38,6 +41,11 @@ protected final <T> void writeMessage(BaseMarshallerDelegate<T> marshallerDelega
     * Invoked by generated code.
     */
    protected final <T> void writeNestedMessage(BaseMarshallerDelegate<T> marshallerDelegate, ProtobufTagMarshaller.WriteContext ctx, int fieldNumber, T message) throws IOException {
+      int maxNestedMessageDepth = ctx.getSerializationContext().getConfiguration().maxNestedMessageDepth();
+      if (ctx.depth() >= maxNestedMessageDepth) {
+         throw log.maxNestedMessageDepth(maxNestedMessageDepth, message.getClass());
+      }
+
       ByteArrayOutputStreamEx baos = new ByteArrayOutputStreamEx();
       TagWriterImpl nested = TagWriterImpl.newNestedInstance(ctx, baos);
       writeMessage(marshallerDelegate, nested, message);

@@ -0,0 +1,24 @@
+package org.infinispan.protostream.exception;
+
+public class ProtoStreamException extends RuntimeException {
+
+   public ProtoStreamException() {
+      super();
+   }
+
+   public ProtoStreamException(Throwable cause) {
+      super(cause);
+   }
+
+   public ProtoStreamException(String msg) {
+      super(msg);
+   }
+
+   public ProtoStreamException(String msg, Throwable cause) {
+      super(msg, cause);
+   }
+
+   public ProtoStreamException(String message, Throwable cause, boolean enableSuppression, boolean writableStackTrace) {
+      super(message, cause, enableSuppression, writableStackTrace);
+   }
+}

@@ -5,6 +5,7 @@
 import java.io.IOException;
 
 import org.infinispan.protostream.MalformedProtobufException;
+import org.infinispan.protostream.exception.ProtoStreamException;
 import org.jboss.logging.BasicLogger;
 import org.jboss.logging.Logger;
 import org.jboss.logging.annotations.Cause;
@@ -45,6 +46,10 @@ default MalformedProtobufException messageTruncated() {
    @Message(value = "Ran out of buffer space", id = 7)
    IOException outOfWriteBufferSpace(@Cause Throwable cause);
 
+   @Message(value = "The nested message depth appears to be larger than the configured limit of '%s'." +
+         "It is possible that the entity to marshall with type '%s' can have some circular dependencies.", id = 8)
+   ProtoStreamException maxNestedMessageDepth(int maxNestedMessageDepth, Class<?> entityType);
+
    class LogFactory {
       public static Log getLog(Class<?> clazz) {
          return Logger.getMessageLogger(Log.class, clazz.getName());

@@ -15,6 +15,8 @@
  */
 public interface Configuration {
 
+   int DEFAULT_MAX_NESTED_DEPTH = 100;
+
    /**
     * The name of the TypeId annotation. This optional annotation defines a unique positive integer type identifier for
     * each message or enum type. This can be used alternatively instead of the fully qualified type name during
@@ -35,6 +37,13 @@ public interface Configuration {
     */
    boolean logOutOfSequenceWrites();
 
+   /**
+    * The max nested message depth to apply to all {@link org.infinispan.protostream.annotations.impl.GeneratedMarshallerBase}s.
+    * This value is used as way to avoid recurring on circular dependencies without the need to maintain the list of already visited entities.
+    * Default to {@link #DEFAULT_MAX_NESTED_DEPTH}
+    */
+   int maxNestedMessageDepth();
+
    WrappingConfig wrappingConfig();
 
    interface WrappingConfig {
@@ -84,6 +93,8 @@ interface Builder {
 
       Builder setLogOutOfSequenceWrites(boolean logOutOfSequenceWrites);
 
+      Builder maxNestedMessageDepth(int maxNestedMessageDepth);
+
       /**
        * Should we log a warning every time we encounter an undefined documentation annotation? This is {@code true} by
        * default.

@@ -23,11 +23,15 @@ public final class ConfigurationImpl implements Configuration {
 
    private final AnnotationsConfigImpl annotationsConfig;
 
+   private final int maxNestedMessageDepth;
+
    private ConfigurationImpl(boolean logOutOfSequenceReads, boolean logOutOfSequenceWrites,
+                             int maxNestedMessageDepth,
                              WrappedMessageTypeIdMapper wrappedMessageTypeIdMapper,
                              Map<String, AnnotationConfigurationImpl> annotations, boolean logUndefinedAnnotations) {
       this.logOutOfSequenceReads = logOutOfSequenceReads;
       this.logOutOfSequenceWrites = logOutOfSequenceWrites;
+      this.maxNestedMessageDepth = maxNestedMessageDepth;
       this.wrappingConfig = new WrappingConfigImpl(wrappedMessageTypeIdMapper);
       this.annotationsConfig = new AnnotationsConfigImpl(annotations, logUndefinedAnnotations);
    }
@@ -42,6 +46,11 @@ public boolean logOutOfSequenceWrites() {
       return logOutOfSequenceWrites;
    }
 
+   @Override
+   public int maxNestedMessageDepth() {
+      return maxNestedMessageDepth;
+   }
+
    @Override
    public WrappingConfig wrappingConfig() {
       return wrappingConfig;
@@ -114,6 +123,8 @@ public static final class BuilderImpl implements Builder {
 
       private boolean logOutOfSequenceWrites = true;
 
+      private int maxNestedMessageDepth = Configuration.DEFAULT_MAX_NESTED_DEPTH;
+
       private WrappingConfigBuilderImpl wrappingConfigBuilder = null;
 
       private AnnotationsConfigBuilderImpl annotationsConfigBuilder = null;
@@ -180,6 +191,12 @@ public Builder setLogOutOfSequenceWrites(boolean logOutOfSequenceWrites) {
          return this;
       }
 
+      @Override
+      public Builder maxNestedMessageDepth(int maxNestedMessageDepth) {
+         this.maxNestedMessageDepth = maxNestedMessageDepth;
+         return this;
+      }
+
       @Override
       public WrappingConfigBuilderImpl wrappingConfig() {
          if (wrappingConfigBuilder == null) {
@@ -228,7 +245,7 @@ public Configuration build() {
 
          // TypeId is the only predefined annotation. If there are more than one then we know we have at least one user defined.
          boolean logUndefinedAnnotations = annotationsConfig().logUndefinedAnnotations == null ? annotations.size() > 1 : annotationsConfig().logUndefinedAnnotations;
-         return new ConfigurationImpl(logOutOfSequenceReads, logOutOfSequenceWrites,
+         return new ConfigurationImpl(logOutOfSequenceReads, logOutOfSequenceWrites, maxNestedMessageDepth,
                wrappingConfig().wrappedMessageTypeIdMapper,
                annotations, logUndefinedAnnotations);
       }

@@ -34,6 +34,8 @@ public final class TagWriterImpl implements TagWriter, ProtobufTagMarshaller.Wri
 
    private final TagWriterImpl parent;
 
+   private final int depth;
+
    // lazily initialized
    private Map<Object, Object> params = null;
 
@@ -43,12 +45,14 @@ public final class TagWriterImpl implements TagWriter, ProtobufTagMarshaller.Wri
 
    private TagWriterImpl(TagWriterImpl parent, Encoder encoder) {
       this.parent = parent;
+      this.depth = parent.depth + 1;
       this.serCtx = parent.serCtx;
       this.encoder = encoder;
    }
 
    private TagWriterImpl(SerializationContextImpl serCtx, Encoder encoder) {
       this.parent = null;
+      this.depth = 0;
       this.serCtx = serCtx;
       this.encoder = encoder;
    }
@@ -274,6 +278,11 @@ public TagWriter getWriter() {
       return this;
    }
 
+   @Override
+   public int depth() {
+      return depth;
+   }
+
    /**
     * @deprecated this will be removed in 5.0 together with {@link org.infinispan.protostream.MessageMarshaller}
     */

@@ -45,5 +45,7 @@ interface ReadContext extends OperationContext {
     */
    interface WriteContext extends OperationContext {
       TagWriter getWriter();
+
+      int depth();
    }
 }