[Fix] Validate Pipe payload lengths before allocation

iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/sink/PipeReceiverTest.java

@@ -19,6 +19,7 @@
 
 package org.apache.iotdb.db.pipe.sink;
 
+import org.apache.iotdb.commons.conf.CommonConfig;
 import org.apache.iotdb.commons.conf.CommonDescriptor;
 import org.apache.iotdb.commons.pipe.sink.payload.thrift.request.IoTDBSinkRequestVersion;
 import org.apache.iotdb.commons.pipe.sink.payload.thrift.request.PipeRequestType;
@@ -38,6 +39,63 @@
 import java.util.Collections;
 
 public class PipeReceiverTest {
+
+  @Test
+  public void testMalformedPreHandshakePayloadsRejected() throws IOException {
+    final CommonConfig config = CommonDescriptor.getInstance().getConfig();
+    final long minimumReceiverMemory = config.getPipeMinimumReceiverMemory();
+    config.setPipeMinimumReceiverMemory(0);
+    try {
+      final IoTDBDataNodeReceiver receiver = new IoTDBDataNodeReceiver();
+      assertMalformedRequestRejected(
+          receiver,
+          PipeRequestType.HANDSHAKE_DATANODE_V1,
+          ByteBuffer.allocate(5).putInt(Integer.MAX_VALUE).put((byte) 1).flip());
+      assertMalformedRequestRejected(
+          receiver,
+          PipeRequestType.HANDSHAKE_DATANODE_V2,
+          ByteBuffer.allocate(9).putInt(1).putInt(Integer.MAX_VALUE).put((byte) 1).flip());
+      assertMalformedRequestRejected(
+          receiver,
+          PipeRequestType.TRANSFER_SLICE,
+          ByteBuffer.allocate(26)
+              .putInt(0)
+              .putShort(PipeRequestType.HANDSHAKE_DATANODE_V1.getType())
+              .putInt(0)
+              .putInt(Integer.MAX_VALUE)
+              .putInt(0)
+              .putInt(0)
+              .putInt(1)
+              .flip());
+      assertMalformedRequestRejected(
+          receiver,
+          PipeRequestType.TRANSFER_PIPE_RECEIVER_RUNTIME_INFO_CLEANUP,
+          ByteBuffer.allocate(12).putInt(Integer.MAX_VALUE).putLong(1).flip());
+    } finally {
+      config.setPipeMinimumReceiverMemory(minimumReceiverMemory);
+    }
+  }
+
+  private void assertMalformedRequestRejected(
+      final IoTDBDataNodeReceiver receiver, final PipeRequestType type, final ByteBuffer body)
+      throws IOException {
+    final TPipeTransferReq req = new TPipeTransferReq();
+    req.setVersion(IoTDBSinkRequestVersion.VERSION_1.getVersion());
+    req.setType(type.getType());
+    req.setBody(body);
+    final TPipeTransferReq compressedReq =
+        PipeTransferCompressedReq.toTPipeTransferReq(req, Collections.emptyList());
+
+    Assert.assertEquals(
+        TSStatusCode.PIPE_ERROR.getStatusCode(), receiver.receive(req).getStatus().getCode());
+    Assert.assertEquals(
+        TSStatusCode.PIPE_ERROR.getStatusCode(),
+        receiver.receive(compressedReq).getStatus().getCode());
+    Assert.assertEquals(
+        TSStatusCode.NOT_LOGIN.getStatusCode(),
+        receiver.receive(buildEmptyRawTabletTransferReq()).getStatus().getCode());
+  }
+
   @Test
   public void testUnauthenticatedPipeTransferRejected() {
     final IoTDBDataNodeReceiver receiver = new IoTDBDataNodeReceiver();

iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/payload/thrift/request/PipeTransferHandshakeV1Req.java

@@ -58,7 +58,7 @@ public final PipeTransferHandshakeV1Req convertToTPipeTransferReq(String timesta
 
   protected final PipeTransferHandshakeV1Req translateFromTPipeTransferReq(
       TPipeTransferReq transferReq) {
-    timestampPrecision = ReadWriteIOUtils.readString(transferReq.body.duplicate());
+    timestampPrecision = PipeTransferPayloadReader.readString(transferReq.body.duplicate());
 
     version = transferReq.version;
     type = transferReq.type;

iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/payload/thrift/request/PipeTransferHandshakeV2Req.java

@@ -26,6 +26,7 @@
 
 import java.io.DataOutputStream;
 import java.io.IOException;
+import java.nio.BufferUnderflowException;
 import java.nio.ByteBuffer;
 import java.util.HashMap;
 import java.util.Map;
@@ -66,9 +67,13 @@ protected final PipeTransferHandshakeV2Req translateFromTPipeTransferReq(
     Map<String, String> params = new HashMap<>();
     final ByteBuffer bodyBuffer = transferReq.body.duplicate();
     final int size = ReadWriteIOUtils.readInt(bodyBuffer);
+    // Each parameter needs at least two string length prefixes, even for null or empty strings.
+    if (size < 0 || size > bodyBuffer.remaining() / (2 * Integer.BYTES)) {
+      throw new BufferUnderflowException();
+    }
     for (int i = 0; i < size; ++i) {
-      final String key = ReadWriteIOUtils.readString(bodyBuffer);
-      final String value = ReadWriteIOUtils.readString(bodyBuffer);
+      final String key = PipeTransferPayloadReader.readString(bodyBuffer);
+      final String value = PipeTransferPayloadReader.readString(bodyBuffer);
       params.put(key, value);
     }
     this.params = params;

iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/payload/thrift/request/PipeTransferPayloadReader.java

@@ -0,0 +1,48 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.iotdb.commons.pipe.sink.payload.thrift.request;
+
+import org.apache.tsfile.utils.ReadWriteIOUtils;
+
+import java.nio.BufferUnderflowException;
+import java.nio.ByteBuffer;
+
+final class PipeTransferPayloadReader {
+
+  private PipeTransferPayloadReader() {}
+
+  static String readString(final ByteBuffer buffer) {
+    checkLength(buffer);
+    return ReadWriteIOUtils.readString(buffer);
+  }
+
+  static byte[] readBinary(final ByteBuffer buffer) {
+    checkLength(buffer);
+    return ReadWriteIOUtils.readBinary(buffer).getValues();
+  }
+
+  private static void checkLength(final ByteBuffer buffer) {
+    // Check the buffer limit before ReadWriteIOUtils allocates from the length prefix.
+    if (buffer.remaining() < Integer.BYTES
+        || buffer.getInt(buffer.position()) > buffer.remaining() - Integer.BYTES) {
+      throw new BufferUnderflowException();
+    }
+  }
+}

iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/payload/thrift/request/PipeTransferPipeReceiverRuntimeInfoCleanupReq.java

@@ -74,7 +74,7 @@ public static PipeTransferPipeReceiverRuntimeInfoCleanupReq fromTPipeTransferReq
     req.body = transferReq.body;
 
     final ByteBuffer bodyBuffer = transferReq.body.duplicate();
-    req.pipeName = ReadWriteIOUtils.readString(bodyBuffer);
+    req.pipeName = PipeTransferPayloadReader.readString(bodyBuffer);
     req.pipeCreationTime = ReadWriteIOUtils.readLong(bodyBuffer);
     return req;
   }

iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/payload/thrift/request/PipeTransferSliceReq.java

@@ -121,7 +121,7 @@ public static PipeTransferSliceReq fromTPipeTransferReq(final TPipeTransferReq t
     sliceReq.originReqType = ReadWriteIOUtils.readShort(transferReq.body);
     sliceReq.originBodySize = ReadWriteIOUtils.readInt(transferReq.body);
 
-    sliceReq.sliceBody = ReadWriteIOUtils.readBinary(transferReq.body).getValues();
+    sliceReq.sliceBody = PipeTransferPayloadReader.readBinary(transferReq.body);
 
     sliceReq.sliceIndex = ReadWriteIOUtils.readInt(transferReq.body);
     sliceReq.sliceCount = ReadWriteIOUtils.readInt(transferReq.body);

iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/pipe/sink/payload/thrift/request/PipeTransferRequestValidationTest.java

@@ -0,0 +1,221 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.iotdb.commons.pipe.sink.payload.thrift.request;
+
+import org.apache.iotdb.service.rpc.thrift.TPipeTransferReq;
+
+import org.junit.Assert;
+import org.junit.Test;
+
+import java.io.IOException;
+import java.nio.BufferUnderflowException;
+import java.nio.ByteBuffer;
+import java.util.HashMap;
+import java.util.Map;
+
+public class PipeTransferRequestValidationTest {
+
+  @Test(expected = BufferUnderflowException.class)
+  public void testHandshakeV1RejectsOversizedString() {
+    new HandshakeV1Req()
+        .translateFromTPipeTransferReq(
+            request(ByteBuffer.allocate(5).putInt(Integer.MAX_VALUE).put((byte) 1).flip()));
+  }
+
+  @Test(expected = BufferUnderflowException.class)
+  public void testHandshakeV2RejectsOversizedKey() {
+    new HandshakeV2Req()
+        .translateFromTPipeTransferReq(
+            request(
+                ByteBuffer.allocate(9).putInt(1).putInt(Integer.MAX_VALUE).put((byte) 1).flip()));
+  }
+
+  @Test(expected = BufferUnderflowException.class)
+  public void testHandshakeV2RejectsOversizedKeyWithValuePrefix() {
+    new HandshakeV2Req()
+        .translateFromTPipeTransferReq(
+            request(ByteBuffer.allocate(12).putInt(1).putInt(Integer.MAX_VALUE).putInt(0).flip()));
+  }
+
+  @Test(expected = BufferUnderflowException.class)
+  public void testHandshakeV2RejectsOversizedValue() {
+    new HandshakeV2Req()
+        .translateFromTPipeTransferReq(
+            request(
+                ByteBuffer.allocate(13)
+                    .putInt(1)
+                    .putInt(0)
+                    .putInt(Integer.MAX_VALUE)
+                    .put((byte) 1)
+                    .flip()));
+  }
+
+  @Test(expected = BufferUnderflowException.class)
+  public void testHandshakeRejectsStringBeyondBufferLimit() {
+    final ByteBuffer body = ByteBuffer.allocate(32);
+    body.putInt(7).putInt(5).put((byte) 'a').flip();
+    body.position(Integer.BYTES);
+    new HandshakeV1Req().translateFromTPipeTransferReq(request(body.asReadOnlyBuffer()));
+  }
+
+  @Test(expected = BufferUnderflowException.class)
+  public void testHandshakeRejectsTruncatedStringLength() {
+    new HandshakeV1Req().translateFromTPipeTransferReq(request(ByteBuffer.allocate(3)));
+  }
+
+  @Test(expected = BufferUnderflowException.class)
+  public void testHandshakeV2RejectsNegativeParameterCount() {
+    new HandshakeV2Req()
+        .translateFromTPipeTransferReq(request(ByteBuffer.allocate(4).putInt(-1).flip()));
+  }
+
+  @Test(expected = BufferUnderflowException.class)
+  public void testHandshakeV2RejectsParameterCountBeyondBody() {
+    new HandshakeV2Req()
+        .translateFromTPipeTransferReq(
+            request(ByteBuffer.allocate(12).putInt(Integer.MAX_VALUE).putInt(0).putInt(0).flip()));
+  }
+
+  @Test
+  public void testHandshakeV1RoundTripPreservesBody() throws IOException {
+    final HandshakeV1Req original = new HandshakeV1Req();
+    original.convertToTPipeTransferReq("ms");
+    final ByteBuffer body = original.body.duplicate();
+    final PipeTransferHandshakeV1Req decoded =
+        new HandshakeV1Req().translateFromTPipeTransferReq(original);
+    Assert.assertEquals("ms", decoded.getTimestampPrecision());
+    Assert.assertEquals(body, original.body);
+    Assert.assertEquals(original.version, decoded.version);
+    Assert.assertEquals(original.type, decoded.type);
+  }
+
+  @Test
+  public void testHandshakeV2RoundTripPreservesNullEmptyAndUnicode() throws IOException {
+    final Map<String, String> params = new HashMap<>();
+    params.put("", "");
+    params.put(null, null);
+    params.put("username", "用户");
+    final HandshakeV2Req original = new HandshakeV2Req();
+    original.convertToTPipeTransferReq(params);
+    final ByteBuffer body = ByteBuffer.allocateDirect(original.body.remaining() + Integer.BYTES);
+    body.putInt(7).put(original.body.duplicate()).flip();
+    body.position(Integer.BYTES);
+    original.body = body.asReadOnlyBuffer();
+
+    final PipeTransferHandshakeV2Req decoded =
+        new HandshakeV2Req().translateFromTPipeTransferReq(original);
+
+    Assert.assertEquals(params, decoded.getParams());
+    Assert.assertEquals(Integer.BYTES, original.body.position());
+    Assert.assertEquals(original.version, decoded.version);
+    Assert.assertEquals(original.type, decoded.type);
+    Assert.assertSame(original.body, decoded.body);
+  }
+
+  @Test
+  public void testHandshakeV2EmptyParameters() throws IOException {
+    final HandshakeV2Req original = new HandshakeV2Req();
+    original.convertToTPipeTransferReq(new HashMap<>());
+    Assert.assertTrue(
+        new HandshakeV2Req().translateFromTPipeTransferReq(original).getParams().isEmpty());
+  }
+
+  @Test(expected = BufferUnderflowException.class)
+  public void testSliceRejectsOversizedBinary() {
+    PipeTransferSliceReq.fromTPipeTransferReq(request(sliceBody(Integer.MAX_VALUE)));
+  }
+
+  @Test(expected = BufferUnderflowException.class)
+  public void testSliceRejectsBinaryBeyondRemainingBody() {
+    PipeTransferSliceReq.fromTPipeTransferReq(request(sliceBody(13)));
+  }
+
+  @Test
+  public void testSliceRoundTrip() throws IOException {
+    final PipeTransferSliceReq original =
+        PipeTransferSliceReq.toTPipeTransferReq(
+            7,
+            PipeRequestType.HANDSHAKE_DATANODE_V1.getType(),
+            0,
+            1,
+            ByteBuffer.wrap(new byte[] {1, 2, 3, 4}),
+            0,
+            4);
+    final PipeTransferSliceReq decoded = PipeTransferSliceReq.fromTPipeTransferReq(original);
+    Assert.assertEquals(original.getOrderId(), decoded.getOrderId());
+    Assert.assertEquals(original.getOriginReqType(), decoded.getOriginReqType());
+    Assert.assertEquals(original.getOriginBodySize(), decoded.getOriginBodySize());
+    Assert.assertEquals(original.getSliceIndex(), decoded.getSliceIndex());
+    Assert.assertEquals(original.getSliceCount(), decoded.getSliceCount());
+    Assert.assertEquals(original.version, decoded.version);
+    Assert.assertEquals(original.type, decoded.type);
+    Assert.assertArrayEquals(new byte[] {1, 2, 3, 4}, decoded.getSliceBody());
+  }
+
+  @Test(expected = BufferUnderflowException.class)
+  public void testCleanupRejectsOversizedPipeName() {
+    PipeTransferPipeReceiverRuntimeInfoCleanupReq.fromTPipeTransferReq(
+        request(ByteBuffer.allocate(12).putInt(Integer.MAX_VALUE).putLong(1).flip()));
+  }
+
+  @Test
+  public void testCleanupRoundTripPreservesBody() throws IOException {
+    final PipeTransferPipeReceiverRuntimeInfoCleanupReq original =
+        PipeTransferPipeReceiverRuntimeInfoCleanupReq.toTPipeTransferReq("pipe", 1);
+    final ByteBuffer body = original.body.duplicate();
+    final PipeTransferPipeReceiverRuntimeInfoCleanupReq decoded =
+        PipeTransferPipeReceiverRuntimeInfoCleanupReq.fromTPipeTransferReq(original);
+    Assert.assertEquals(original, decoded);
+    Assert.assertEquals(body, original.body);
+  }
+
+  private static ByteBuffer sliceBody(final int length) {
+    return ByteBuffer.allocate(26)
+        .putInt(0)
+        .putShort(PipeRequestType.HANDSHAKE_DATANODE_V1.getType())
+        .putInt(4)
+        .putInt(length)
+        .putInt(0)
+        .putInt(0)
+        .putInt(1)
+        .flip();
+  }
+
+  private static TPipeTransferReq request(final ByteBuffer body) {
+    final TPipeTransferReq req = new TPipeTransferReq();
+    req.version = IoTDBSinkRequestVersion.VERSION_1.getVersion();
+    req.body = body;
+    return req;
+  }
+
+  private static class HandshakeV1Req extends PipeTransferHandshakeV1Req {
+    @Override
+    protected PipeRequestType getPlanType() {
+      return PipeRequestType.HANDSHAKE_DATANODE_V1;
+    }
+  }
+
+  private static class HandshakeV2Req extends PipeTransferHandshakeV2Req {
+    @Override
+    protected PipeRequestType getPlanType() {
+      return PipeRequestType.HANDSHAKE_DATANODE_V2;
+    }
+  }
+}