Reject frames less than min MTU size (#913)

rstoyanchev · web-flow · commit 2c873d6bf886 · 2020-08-10T17:41:40.000+03:00
This is the proposed fix for #895 for the 1.1 codebase where the logic for re-assembling is in a different class than 1.0. Signed-off-by: Rossen Stoyanchev <rstoyanchev@vmware.com>
diff --git a/rsocket-core/src/main/java/io/rsocket/core/FireAndForgetResponderSubscriber.java b/rsocket-core/src/main/java/io/rsocket/core/FireAndForgetResponderSubscriber.java
@@ -68,7 +68,7 @@ private FireAndForgetResponderSubscriber() {
 
     this.frames =
         ReassemblyUtils.addFollowingFrame(
-            allocator.compositeBuffer(), firstFrame, maxInboundPayloadSize);
+            allocator.compositeBuffer(), firstFrame, true, maxInboundPayloadSize);
   }
 
   @Override
@@ -92,7 +92,8 @@ public void handleNext(ByteBuf followingFrame, boolean hasFollows, boolean isLas
     final CompositeByteBuf frames = this.frames;
 
     try {
-      ReassemblyUtils.addFollowingFrame(frames, followingFrame, this.maxInboundPayloadSize);
+      ReassemblyUtils.addFollowingFrame(
+          frames, followingFrame, hasFollows, this.maxInboundPayloadSize);
     } catch (IllegalStateException t) {
       this.requesterResponderSupport.remove(this.streamId, this);
 
diff --git a/rsocket-core/src/main/java/io/rsocket/core/ReassemblyUtils.java b/rsocket-core/src/main/java/io/rsocket/core/ReassemblyUtils.java
@@ -102,7 +102,7 @@ static <T extends RequesterFrameHandler> void handleNextSupport(
     if (frames == null) {
       frames =
           ReassemblyUtils.addFollowingFrame(
-              allocator.compositeBuffer(), frame, maxInboundPayloadSize);
+              allocator.compositeBuffer(), frame, hasFollows, maxInboundPayloadSize);
       instance.setFrames(frames);
 
       long previousState = markReassembling(updater, instance);
@@ -113,7 +113,8 @@ static <T extends RequesterFrameHandler> void handleNextSupport(
       }
     } else {
       try {
-        frames = ReassemblyUtils.addFollowingFrame(frames, frame, maxInboundPayloadSize);
+        frames =
+            ReassemblyUtils.addFollowingFrame(frames, frame, hasFollows, maxInboundPayloadSize);
       } catch (IllegalStateException t) {
         if (isTerminated(updater.get(instance))) {
           return;
@@ -160,7 +161,10 @@ static <T extends RequesterFrameHandler> void handleNextSupport(
   }
 
   static CompositeByteBuf addFollowingFrame(
-      CompositeByteBuf frames, ByteBuf followingFrame, int maxInboundPayloadSize) {
+      CompositeByteBuf frames,
+      ByteBuf followingFrame,
+      boolean hasFollows,
+      int maxInboundPayloadSize) {
     int readableBytes = frames.readableBytes();
     if (readableBytes == 0) {
       return frames.addComponent(true, followingFrame.retain());
@@ -169,6 +173,9 @@ static CompositeByteBuf addFollowingFrame(
             > maxInboundPayloadSize) {
       throw new IllegalStateException(
           String.format(ILLEGAL_REASSEMBLED_PAYLOAD_SIZE, maxInboundPayloadSize));
+    } else if (followingFrame.readableBytes() < MIN_MTU_SIZE - 3 && hasFollows) {
+      // FIXME: check MIN_MTU_SIZE only (currently fragments have size of 61)
+      throw new IllegalStateException("Fragment is too small.");
     }
 
     final boolean hasMetadata = FrameHeaderCodec.hasMetadata(followingFrame);
diff --git a/rsocket-core/src/main/java/io/rsocket/core/RequestChannelResponderSubscriber.java b/rsocket-core/src/main/java/io/rsocket/core/RequestChannelResponderSubscriber.java
@@ -104,7 +104,7 @@ public RequestChannelResponderSubscriber(
 
     this.frames =
         ReassemblyUtils.addFollowingFrame(
-            allocator.compositeBuffer(), firstFrame, maxInboundPayloadSize);
+            allocator.compositeBuffer(), firstFrame, true, maxInboundPayloadSize);
     STATE.lazySet(this, REASSEMBLING_FLAG);
   }
 
@@ -491,7 +491,7 @@ public void handleNext(ByteBuf frame, boolean hasFollows, boolean isLastPayload)
     if (frames == null) {
       frames =
           ReassemblyUtils.addFollowingFrame(
-              this.allocator.compositeBuffer(), frame, this.maxInboundPayloadSize);
+              this.allocator.compositeBuffer(), frame, hasFollows, this.maxInboundPayloadSize);
       this.frames = frames;
 
       long previousState = markReassembling(STATE, this);
@@ -502,7 +502,9 @@ public void handleNext(ByteBuf frame, boolean hasFollows, boolean isLastPayload)
       }
     } else {
       try {
-        frames = ReassemblyUtils.addFollowingFrame(frames, frame, this.maxInboundPayloadSize);
+        frames =
+            ReassemblyUtils.addFollowingFrame(
+                frames, frame, hasFollows, this.maxInboundPayloadSize);
       } catch (IllegalStateException e) {
         if (isTerminated(this.state)) {
           return;
diff --git a/rsocket-core/src/main/java/io/rsocket/core/RequestResponseResponderSubscriber.java b/rsocket-core/src/main/java/io/rsocket/core/RequestResponseResponderSubscriber.java
@@ -81,7 +81,7 @@ public RequestResponseResponderSubscriber(
     this.handler = handler;
     this.frames =
         ReassemblyUtils.addFollowingFrame(
-            allocator.compositeBuffer(), firstFrame, maxInboundPayloadSize);
+            allocator.compositeBuffer(), firstFrame, true, maxInboundPayloadSize);
   }
 
   public RequestResponseResponderSubscriber(
@@ -218,7 +218,7 @@ public void handleNext(ByteBuf frame, boolean hasFollows, boolean isLastPayload)
     }
 
     try {
-      ReassemblyUtils.addFollowingFrame(frames, frame, this.maxInboundPayloadSize);
+      ReassemblyUtils.addFollowingFrame(frames, frame, hasFollows, this.maxInboundPayloadSize);
     } catch (IllegalStateException t) {
       S.lazySet(this, Operators.cancelledSubscription());
 
diff --git a/rsocket-core/src/main/java/io/rsocket/core/RequestStreamResponderSubscriber.java b/rsocket-core/src/main/java/io/rsocket/core/RequestStreamResponderSubscriber.java
@@ -84,7 +84,7 @@ public RequestStreamResponderSubscriber(
     this.handler = handler;
     this.frames =
         ReassemblyUtils.addFollowingFrame(
-            allocator.compositeBuffer(), firstFrame, maxInboundPayloadSize);
+            allocator.compositeBuffer(), firstFrame, true, maxInboundPayloadSize);
   }
 
   public RequestStreamResponderSubscriber(
@@ -258,7 +258,8 @@ public void handleNext(ByteBuf followingFrame, boolean hasFollows, boolean isLas
     }
 
     try {
-      ReassemblyUtils.addFollowingFrame(frames, followingFrame, this.maxInboundPayloadSize);
+      ReassemblyUtils.addFollowingFrame(
+          frames, followingFrame, hasFollows, this.maxInboundPayloadSize);
     } catch (IllegalStateException t) {
       // if subscription is null, it means that streams has not yet reassembled all the fragments
       // and fragmentation of the first frame was cancelled before
diff --git a/rsocket-core/src/test/java/io/rsocket/core/RSocketRequesterTest.java b/rsocket-core/src/test/java/io/rsocket/core/RSocketRequesterTest.java
@@ -1332,6 +1332,36 @@ public void errorTooBigPayload(
     leaksTrackingByteBufAllocator.assertHasNoLeaks();
   }
 
+  @ParameterizedTest(name = "throws error if fragment before the last is < min MTU {0}")
+  @MethodSource("requestNInteractions")
+  public void errorFragmentTooSmall(
+      FrameType frameType,
+      BiFunction<ClientSocketRule, Payload, Publisher<Payload>> requestFunction) {
+    final int mtu = 32;
+    final LeaksTrackingByteBufAllocator leaksTrackingByteBufAllocator =
+        LeaksTrackingByteBufAllocator.instrument(ByteBufAllocator.DEFAULT);
+
+    final Payload requestPayload = genericPayload(leaksTrackingByteBufAllocator);
+    final Payload responsePayload = fixedSizePayload(leaksTrackingByteBufAllocator, 156);
+    List<ByteBuf> fragments = prepareFragments(leaksTrackingByteBufAllocator, mtu, responsePayload);
+    responsePayload.release();
+
+    StepVerifier.create(requestFunction.apply(rule, requestPayload))
+        .then(() -> rule.connection.addToReceivedBuffer(fragments.toArray(new ByteBuf[0])))
+        .expectErrorMessage("Fragment is too small.")
+        .verify();
+
+    FrameAssert.assertThat(rule.connection.getSent().poll()).typeOf(frameType).hasNoLeaks();
+
+    if (frameType == REQUEST_CHANNEL) {
+      FrameAssert.assertThat(rule.connection.getSent().poll()).typeOf(COMPLETE).hasNoLeaks();
+    }
+
+    FrameAssert.assertThat(rule.connection.getSent().poll()).typeOf(CANCEL).hasNoLeaks();
+
+    leaksTrackingByteBufAllocator.assertHasNoLeaks();
+  }
+
   public static class ClientSocketRule extends AbstractSocketRule<RSocketRequester> {
     @Override
     protected RSocketRequester newRSocket() {
diff --git a/rsocket-core/src/test/java/io/rsocket/core/RSocketResponderTest.java b/rsocket-core/src/test/java/io/rsocket/core/RSocketResponderTest.java
@@ -1002,6 +1002,55 @@ public Flux<Payload> requestChannel(Publisher<Payload> payloads) {
     rule.assertHasNoLeaks();
   }
 
+  @ParameterizedTest(name = "throws error if fragment before the last is < min MTU {0}")
+  @MethodSource("requestCases")
+  public void errorFragmentTooSmall(FrameType frameType) {
+    final int mtu = 32;
+    AtomicReference<Payload> receivedPayload = new AtomicReference<>();
+    rule.setAcceptingSocket(
+        new RSocket() {
+          @Override
+          public Mono<Void> fireAndForget(Payload payload) {
+            receivedPayload.set(payload);
+            return Mono.empty();
+          }
+
+          @Override
+          public Mono<Payload> requestResponse(Payload payload) {
+            receivedPayload.set(payload);
+            return Mono.just(genericPayload(rule.allocator));
+          }
+
+          @Override
+          public Flux<Payload> requestStream(Payload payload) {
+            receivedPayload.set(payload);
+            return Flux.just(genericPayload(rule.allocator));
+          }
+
+          @Override
+          public Flux<Payload> requestChannel(Publisher<Payload> payloads) {
+            Flux.from(payloads).subscribe(receivedPayload::set, null, null, s -> s.request(1));
+            return Flux.just(genericPayload(rule.allocator));
+          }
+        });
+    final Payload randomPayload = fixedSizePayload(rule.allocator, 156);
+    List<ByteBuf> fragments = prepareFragments(rule.allocator, mtu, randomPayload, frameType);
+    randomPayload.release();
+
+    rule.connection.addToReceivedBuffer(fragments.toArray(new ByteBuf[0]));
+
+    PayloadAssert.assertThat(receivedPayload.get()).isNull();
+
+    if (frameType != REQUEST_FNF) {
+      FrameAssert.assertThat(rule.connection.getSent().poll())
+          .typeOf(ERROR)
+          .hasData("Failed to reassemble payload. Cause: Fragment is too small.")
+          .hasNoLeaks();
+    }
+
+    rule.assertHasNoLeaks();
+  }
+
   @ParameterizedTest
   @MethodSource("requestCases")
   void receivingRequestOnStreamIdThaIsAlreadyInUseMUSTBeIgnored_ReassemblyCase(

Original file line number	Diff line number	Diff line change
`@@ -81,7 +81,7 @@ public RequestResponseResponderSubscriber(`
`81`	`81`	`this.handler = handler;`
`82`	`82`	`this.frames =`
`83`	`83`	`ReassemblyUtils.addFollowingFrame(`
`84`		`- allocator.compositeBuffer(), firstFrame, maxInboundPayloadSize);`
	`84`	`+ allocator.compositeBuffer(), firstFrame, true, maxInboundPayloadSize);`
`85`	`85`	`}`
`86`	`86`
`87`	`87`	`public RequestResponseResponderSubscriber(`
`@@ -218,7 +218,7 @@ public void handleNext(ByteBuf frame, boolean hasFollows, boolean isLastPayload)`
`218`	`218`	`}`
`219`	`219`
`220`	`220`	`try {`
`221`		`- ReassemblyUtils.addFollowingFrame(frames, frame, this.maxInboundPayloadSize);`
	`221`	`+ ReassemblyUtils.addFollowingFrame(frames, frame, hasFollows, this.maxInboundPayloadSize);`
`222`	`222`	`} catch (IllegalStateException t) {`
`223`	`223`	`S.lazySet(this, Operators.cancelledSubscription());`
`224`	`224`