Reject empty fragments

Closes gh-895

Signed-off-by: Rossen Stoyanchev <[email protected]>

Author: rstoyanchev

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

@@ -23,7 +23,14 @@
 import io.netty.util.ReferenceCountUtil;
 import io.netty.util.collection.IntObjectHashMap;
 import io.netty.util.collection.IntObjectMap;
-import io.rsocket.frame.*;
+import io.rsocket.frame.FragmentationCodec;
+import io.rsocket.frame.FrameHeaderCodec;
+import io.rsocket.frame.FrameType;
+import io.rsocket.frame.PayloadFrameCodec;
+import io.rsocket.frame.RequestChannelFrameCodec;
+import io.rsocket.frame.RequestFireAndForgetFrameCodec;
+import io.rsocket.frame.RequestResponseFrameCodec;
+import io.rsocket.frame.RequestStreamFrameCodec;
 import java.util.concurrent.atomic.AtomicBoolean;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -221,31 +228,34 @@ void handleFollowsFlag(ByteBuf frame, int streamId, FrameType frameType) {
       putHeader(streamId, header);
     }
 
+    ByteBuf metadata = null;
     if (FrameHeaderCodec.hasMetadata(frame)) {
-      CompositeByteBuf metadata = getMetadata(streamId);
       switch (frameType) {
         case REQUEST_FNF:
-          metadata.addComponents(true, RequestFireAndForgetFrameCodec.metadata(frame).retain());
+          metadata = RequestFireAndForgetFrameCodec.metadata(frame);
           break;
         case REQUEST_STREAM:
-          metadata.addComponents(true, RequestStreamFrameCodec.metadata(frame).retain());
+          metadata = RequestStreamFrameCodec.metadata(frame);
           break;
         case REQUEST_RESPONSE:
-          metadata.addComponents(true, RequestResponseFrameCodec.metadata(frame).retain());
+          metadata = RequestResponseFrameCodec.metadata(frame);
           break;
         case REQUEST_CHANNEL:
-          metadata.addComponents(true, RequestChannelFrameCodec.metadata(frame).retain());
+          metadata = RequestChannelFrameCodec.metadata(frame);
           break;
           // Payload and synthetic types
         case PAYLOAD:
         case NEXT:
         case NEXT_COMPLETE:
         case COMPLETE:
-          metadata.addComponents(true, PayloadFrameCodec.metadata(frame).retain());
+          metadata = PayloadFrameCodec.metadata(frame);
           break;
         default:
           throw new IllegalStateException("unsupported fragment type");
       }
+      if (metadata != null) {
+        getMetadata(streamId).addComponents(true, metadata.retain());
+      }
     }
 
     ByteBuf data;
@@ -276,6 +286,10 @@ void handleFollowsFlag(ByteBuf frame, int streamId, FrameType frameType) {
 
     getData(streamId).addComponents(true, data);
     frame.release();
+
+    if ((metadata != null && metadata.readableBytes() == 0) && data.readableBytes() == 0) {
+      throw new IllegalStateException("Empty frame.");
+    }
   }
 
   void reassembleFrame(ByteBuf frame, SynchronousSink<ByteBuf> sink) {

rsocket-core/src/test/java/io/rsocket/fragmentation/FrameReassemblerTest.java

@@ -438,7 +438,7 @@ public void cancelBeforeAssembling() {
     Assert.assertFalse(reassembler.data.containsKey(1));
   }
 
-  @ParameterizedTest(name = "throws error if reassembling payload size exist {0}")
+  @ParameterizedTest(name = "throws error if reassembling payload size exceeds {0}")
   @ValueSource(ints = {64, 1024, 2048, 4096})
   public void errorTooBigPayload(int maxFrameLength) {
     List<ByteBuf> byteBufs =
@@ -478,6 +478,24 @@ public void errorTooBigPayload(int maxFrameLength) {
         .isExactlyInstanceOf(IllegalStateException.class);
   }
 
+  @DisplayName("throws error on empty fragment")
+  @Test
+  public void errorEmptyFrame() {
+    List<ByteBuf> byteBufs =
+        Arrays.asList(
+            RequestResponseFrameCodec.encode(
+                allocator, 1, true, Unpooled.wrappedBuffer(metadata), Unpooled.EMPTY_BUFFER),
+            PayloadFrameCodec.encode(
+                allocator, 1, true, false, true, Unpooled.EMPTY_BUFFER, Unpooled.EMPTY_BUFFER));
+
+    FrameReassembler reassembler = new FrameReassembler(allocator, Integer.MAX_VALUE);
+
+    Assertions.assertThatThrownBy(
+            Flux.fromIterable(byteBufs).handle(reassembler::reassembleFrame)::blockLast)
+        .hasMessage("Empty frame.")
+        .isExactlyInstanceOf(IllegalStateException.class);
+  }
+
   @DisplayName("dispose should clean up maps")
   @Test
   public void dispose() {

rsocket-core/src/test/java/io/rsocket/fragmentation/ReassembleDuplexConnectionTest.java

@@ -277,7 +277,7 @@ void reassembleNonFragmentableFrame() {
         .verifyComplete();
   }
 
-  @ParameterizedTest(name = "throws error if reassembling payload size exist {0}")
+  @ParameterizedTest(name = "throws error if reassembling payload size exceeds {0}")
   @ValueSource(ints = {64, 1024, 2048, 4096})
   public void errorTooBigPayload(int maxFrameLength) {
     List<ByteBuf> byteBufs =
@@ -331,4 +331,37 @@ public void errorTooBigPayload(int maxFrameLength) {
 
     allocator.assertHasNoLeaks();
   }
+
+  @DisplayName("throws error on empty fragment")
+  @Test
+  public void errorEmptyFrame() {
+    List<ByteBuf> byteBufs =
+        Arrays.asList(
+            RequestResponseFrameCodec.encode(
+                allocator, 1, true, Unpooled.wrappedBuffer(metadata), Unpooled.EMPTY_BUFFER),
+            PayloadFrameCodec.encode(
+                allocator, 1, true, false, true, Unpooled.EMPTY_BUFFER, Unpooled.EMPTY_BUFFER));
+
+    MonoProcessor<Void> onClose = MonoProcessor.create();
+
+    when(delegate.receive())
+        .thenReturn(
+            Flux.fromIterable(byteBufs)
+                .doOnDiscard(ReferenceCounted.class, ReferenceCountUtil::release));
+    when(delegate.onClose()).thenReturn(onClose);
+    when(delegate.alloc()).thenReturn(allocator);
+
+    new ReassemblyDuplexConnection(delegate, Integer.MAX_VALUE)
+        .receive()
+        .doFinally(__ -> onClose.onComplete())
+        .as(StepVerifier::create)
+        .expectErrorSatisfies(
+            t ->
+                Assertions.assertThat(t)
+                    .hasMessage("Empty frame.")
+                    .isExactlyInstanceOf(IllegalStateException.class))
+        .verify(Duration.ofSeconds(1));
+
+    allocator.assertHasNoLeaks();
+  }
 }