diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/sink/PipeReceiverTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/sink/PipeReceiverTest.java index 7262920a630d8..84574095fb96f 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/sink/PipeReceiverTest.java +++ b/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(); diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/payload/thrift/request/PipeTransferHandshakeV1Req.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/payload/thrift/request/PipeTransferHandshakeV1Req.java index 4495ee1d0e330..c07eb570bac99 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/payload/thrift/request/PipeTransferHandshakeV1Req.java +++ b/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; diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/payload/thrift/request/PipeTransferHandshakeV2Req.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/payload/thrift/request/PipeTransferHandshakeV2Req.java index 472679494729b..5668abb12e63f 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/payload/thrift/request/PipeTransferHandshakeV2Req.java +++ b/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 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; diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/payload/thrift/request/PipeTransferPayloadReader.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/payload/thrift/request/PipeTransferPayloadReader.java new file mode 100644 index 0000000000000..e3475b04d26df --- /dev/null +++ b/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(); + } + } +} diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/payload/thrift/request/PipeTransferPipeReceiverRuntimeInfoCleanupReq.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/payload/thrift/request/PipeTransferPipeReceiverRuntimeInfoCleanupReq.java index 4c261b9ee31db..efa7319f4268e 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/payload/thrift/request/PipeTransferPipeReceiverRuntimeInfoCleanupReq.java +++ b/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; } diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/payload/thrift/request/PipeTransferSliceReq.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/payload/thrift/request/PipeTransferSliceReq.java index a7753d41bd0ce..9c8a41e115a10 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/payload/thrift/request/PipeTransferSliceReq.java +++ b/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); diff --git a/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/pipe/sink/payload/thrift/request/PipeTransferRequestValidationTest.java b/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/pipe/sink/payload/thrift/request/PipeTransferRequestValidationTest.java new file mode 100644 index 0000000000000..2f30c6eec0f07 --- /dev/null +++ b/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 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; + } + } +}