provides reassembly payload size control

Signed-off-by: Oleh Dokuka <[email protected]>

Author: OlegDokuka

rsocket-core/src/main/java/io/rsocket/core/RSocketConnector.java

@@ -72,6 +72,7 @@
  * }</pre>
  */
 public class RSocketConnector {
+
   private static final String CLIENT_TAG = "client";
 
   private static final BiConsumer<RSocket, Invalidatable> INVALIDATE_FUNCTION =
@@ -91,6 +92,7 @@ public class RSocketConnector {
   private Supplier<Leases<?>> leasesSupplier;
 
   private int mtu = 0;
+  private int maxReassemblySize = Integer.MAX_VALUE;
   private PayloadDecoder payloadDecoder = PayloadDecoder.DEFAULT;
 
   private RSocketConnector() {}
@@ -395,6 +397,23 @@ public RSocketConnector lease(Supplier<Leases<? extends LeaseStats>> supplier) {
     return this;
   }
 
+  /**
+   * When this is set, frames reassembler control maximum payload size which can be reassembled.
+   *
+   * <p>By default this is not set in which case maximum reassembled payloads size is not
+   * controlled.
+   *
+   * @param maxReassemblySize the threshold size for reassembly, must be no less than 16,777,215
+   * @return the same instance for method chaining
+   * @see <a
+   *     href="https://github.com/rsocket/rsocket/blob/master/Protocol.md#fragmentation-and-reassembly">Fragmentation
+   *     and Reassembly</a>
+   */
+  public RSocketConnector reassemble(int maxReassemblySize) {
+    this.maxReassemblySize = ReassemblyDuplexConnection.assertMaxReassemblySize(maxReassemblySize);
+    return this;
+  }
+
   /**
    * When this is set, frames larger than the given maximum transmission unit (mtu) size value are
    * broken down into fragments to fit that size.
@@ -477,8 +496,9 @@ public Mono<RSocket> connect(Supplier<ClientTransport> transportSupplier) {
             .map(
                 connection ->
                     mtu > 0
-                        ? new FragmentationDuplexConnection(connection, mtu, "client")
-                        : new ReassemblyDuplexConnection(connection));
+                        ? new FragmentationDuplexConnection(
+                            connection, mtu, maxReassemblySize, "client")
+                        : new ReassemblyDuplexConnection(connection, maxReassemblySize));
     return connectionMono
         .flatMap(
             connection -> {

rsocket-core/src/main/java/io/rsocket/core/RSocketServer.java

@@ -66,6 +66,7 @@ public final class RSocketServer {
   private Supplier<Leases<?>> leasesSupplier = null;
 
   private int mtu = 0;
+  private int maxReassemblySize = Integer.MAX_VALUE;
   private PayloadDecoder payloadDecoder = PayloadDecoder.DEFAULT;
 
   private RSocketServer() {}
@@ -198,6 +199,23 @@ public RSocketServer lease(Supplier<Leases<?>> supplier) {
     return this;
   }
 
+  /**
+   * When this is set, frames reassembler control maximum payload size which can be reassembled.
+   *
+   * <p>By default this is not set in which case maximum reassembled payloads size is not
+   * controlled.
+   *
+   * @param maxReassemblySize the threshold size for reassembly, must be no less than 16,777,215
+   * @return the same instance for method chaining
+   * @see <a
+   *     href="https://github.com/rsocket/rsocket/blob/master/Protocol.md#fragmentation-and-reassembly">Fragmentation
+   *     and Reassembly</a>
+   */
+  public RSocketServer reassemble(int maxReassemblySize) {
+    this.maxReassemblySize = ReassemblyDuplexConnection.assertMaxReassemblySize(maxReassemblySize);
+    return this;
+  }
+
   /**
    * When this is set, frames larger than the given maximum transmission unit (mtu) size value are
    * fragmented.
@@ -301,8 +319,8 @@ public Mono<Void> apply(DuplexConnection connection) {
   private Mono<Void> acceptor(ServerSetup serverSetup, DuplexConnection connection) {
     connection =
         mtu > 0
-            ? new FragmentationDuplexConnection(connection, mtu, "server")
-            : new ReassemblyDuplexConnection(connection);
+            ? new FragmentationDuplexConnection(connection, mtu, maxReassemblySize, "server")
+            : new ReassemblyDuplexConnection(connection, maxReassemblySize);
 
     ClientServerInputMultiplexer multiplexer =
         new ClientServerInputMultiplexer(connection, interceptors, false);

rsocket-core/src/main/java/io/rsocket/fragmentation/FragmentationDuplexConnection.java

@@ -39,12 +39,14 @@
  */
 public final class FragmentationDuplexConnection extends ReassemblyDuplexConnection
     implements DuplexConnection {
+
   public static final int MIN_MTU_SIZE = 64;
+
   private static final Logger logger = LoggerFactory.getLogger(FragmentationDuplexConnection.class);
-  private final DuplexConnection delegate;
-  private final int mtu;
-  private final FrameReassembler frameReassembler;
-  private final String type;
+
+  final DuplexConnection delegate;
+  final int mtu;
+  final String type;
 
   /**
    * Class constructor.
@@ -53,16 +55,14 @@ public final class FragmentationDuplexConnection extends ReassemblyDuplexConnect
    * @param mtu the fragment size, greater than {@link #MIN_MTU_SIZE}
    * @param type a label to use for logging purposes
    */
-  public FragmentationDuplexConnection(DuplexConnection delegate, int mtu, String type) {
-    super(delegate);
+  public FragmentationDuplexConnection(
+      DuplexConnection delegate, int mtu, int maxReassemblySize, String type) {
+    super(delegate, maxReassemblySize);
 
     Objects.requireNonNull(delegate, "delegate must not be null");
     this.delegate = delegate;
     this.mtu = assertMtu(mtu);
-    this.frameReassembler = new FrameReassembler(delegate.alloc());
     this.type = type;
-
-    delegate.onClose().doFinally(s -> frameReassembler.dispose()).subscribe();
   }
 
   private boolean shouldFragment(FrameType frameType, int readableBytes) {

rsocket-core/src/main/java/io/rsocket/fragmentation/FrameReassembler.java

@@ -48,10 +48,12 @@ final class FrameReassembler extends AtomicBoolean implements Disposable {
   final IntObjectMap<CompositeByteBuf> metadata;
   final IntObjectMap<CompositeByteBuf> data;
 
-  private final ByteBufAllocator allocator;
+  final ByteBufAllocator allocator;
+  final int maxReassemblySize;
 
-  public FrameReassembler(ByteBufAllocator allocator) {
+  public FrameReassembler(ByteBufAllocator allocator, int maxReassemblySize) {
     this.allocator = allocator;
+    this.maxReassemblySize = maxReassemblySize;
     this.headers = new IntObjectHashMap<>();
     this.metadata = new IntObjectHashMap<>();
     this.data = new IntObjectHashMap<>();
@@ -100,6 +102,16 @@ synchronized CompositeByteBuf getMetadata(int streamId) {
     return byteBuf;
   }
 
+  synchronized int getMetadataSize(int streamId) {
+    CompositeByteBuf byteBuf = metadata.get(streamId);
+
+    if (byteBuf == null) {
+      return 0;
+    }
+
+    return byteBuf.readableBytes();
+  }
+
   synchronized CompositeByteBuf getData(int streamId) {
     CompositeByteBuf byteBuf = data.get(streamId);
 
@@ -111,6 +123,16 @@ synchronized CompositeByteBuf getData(int streamId) {
     return byteBuf;
   }
 
+  synchronized int getDataSize(int streamId) {
+    CompositeByteBuf byteBuf = data.get(streamId);
+
+    if (byteBuf == null) {
+      return 0;
+    }
+
+    return byteBuf.readableBytes();
+  }
+
   @Nullable
   synchronized ByteBuf removeHeader(int streamId) {
     return headers.remove(streamId);
@@ -151,6 +173,19 @@ void cancelAssemble(int streamId) {
   void handleNoFollowsFlag(ByteBuf frame, SynchronousSink<ByteBuf> sink, int streamId) {
     ByteBuf header = removeHeader(streamId);
     if (header != null) {
+
+      int maxReassemblySize = this.maxReassemblySize;
+      if (maxReassemblySize != Integer.MAX_VALUE) {
+        int currentPayloadSize = getMetadataSize(streamId) + getDataSize(streamId);
+        if (currentPayloadSize + frame.readableBytes() - FrameHeaderCodec.size()
+            > maxReassemblySize) {
+          header.release();
+          frame.release();
+          cancelAssemble(streamId);
+          throw new IllegalStateException("Reassembled payload went out of allowed size");
+        }
+      }
+
       if (FrameHeaderCodec.hasMetadata(header)) {
         ByteBuf assembledFrame = assembleFrameWithMetadata(frame, streamId, header);
         sink.next(assembledFrame);
@@ -177,6 +212,17 @@ void handleFollowsFlag(ByteBuf frame, int streamId, FrameType frameType) {
       putHeader(streamId, header);
     }
 
+    int maxReassemblySize = this.maxReassemblySize;
+    if (maxReassemblySize != Integer.MAX_VALUE) {
+      int currentPayloadSize = getMetadataSize(streamId) + getDataSize(streamId);
+      if (currentPayloadSize + frame.readableBytes() - FrameHeaderCodec.size()
+          > maxReassemblySize) {
+        cancelAssemble(streamId);
+        frame.release();
+        throw new IllegalStateException("Reassembled payload went out of allowed size");
+      }
+    }
+
     if (FrameHeaderCodec.hasMetadata(frame)) {
       CompositeByteBuf metadata = getMetadata(streamId);
       switch (frameType) {
@@ -226,6 +272,7 @@ void handleFollowsFlag(ByteBuf frame, int streamId, FrameType frameType) {
         data = PayloadFrameCodec.data(frame).retain();
         break;
       default:
+        frame.release();
         throw new IllegalStateException("unsupported fragment type");
     }
 

rsocket-core/src/main/java/io/rsocket/fragmentation/ReassemblyDuplexConnection.java

@@ -19,6 +19,7 @@
 import io.netty.buffer.ByteBuf;
 import io.netty.buffer.ByteBufAllocator;
 import io.rsocket.DuplexConnection;
+import io.rsocket.frame.FrameLengthCodec;
 import java.util.Objects;
 import org.reactivestreams.Publisher;
 import reactor.core.publisher.Flux;
@@ -36,14 +37,26 @@ public class ReassemblyDuplexConnection implements DuplexConnection {
   private final FrameReassembler frameReassembler;
 
   /** Constructor with the underlying delegate to receive frames from. */
-  public ReassemblyDuplexConnection(DuplexConnection delegate) {
+  public ReassemblyDuplexConnection(DuplexConnection delegate, int maxReassemblySize) {
     Objects.requireNonNull(delegate, "delegate must not be null");
     this.delegate = delegate;
-    this.frameReassembler = new FrameReassembler(delegate.alloc());
+    this.frameReassembler = new FrameReassembler(delegate.alloc(), maxReassemblySize);
 
     delegate.onClose().doFinally(s -> frameReassembler.dispose()).subscribe();
   }
 
+  public static int assertMaxReassemblySize(int maxReassemblySize) {
+    if (maxReassemblySize < FrameLengthCodec.FRAME_LENGTH_MASK) {
+      String msg =
+          String.format(
+              "The smallest allowed maxReassemblySize size is %d bytes, provided: %d",
+              FrameLengthCodec.FRAME_LENGTH_MASK, maxReassemblySize);
+      throw new IllegalArgumentException(msg);
+    } else {
+      return maxReassemblySize;
+    }
+  }
+
   @Override
   public Mono<Void> send(Publisher<ByteBuf> frames) {
     return delegate.send(frames);