From 600afef0d22249581d7c1b955696079c83af8216 Mon Sep 17 00:00:00 2001 From: Caideyipi <87789683+Caideyipi@users.noreply.github.com> Date: Tue, 7 Jul 2026 10:31:11 +0800 Subject: [PATCH] Add pipe receiver memory protection (#18090) * Add pipe receiver memory protection * Add pipe receiver memory protection tests * Address receiver memory review comments (cherry picked from commit 71696f333af3f71e2bea946754436b90247b0d4e) --- .../thrift/IoTDBDataNodeReceiver.java | 117 ++++++++++++-- .../pipe/receiver/IoTDBFileReceiver.java | 39 ++++- .../common/PipeTransferSliceReqHandler.java | 37 +++++ .../request/PipeTransferCompressedReq.java | 21 +++ .../pipe/receiver/IoTDBFileReceiverTest.java | 89 +++++++++++ .../PipeTransferSliceReqBuilderTest.java | 53 +++++++ .../PipeTransferCompressedReqTest.java | 149 ++++++++++++++++++ 7 files changed, 489 insertions(+), 16 deletions(-) create mode 100644 iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/pipe/sink/payload/thrift/request/PipeTransferCompressedReqTest.java diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/receiver/protocol/thrift/IoTDBDataNodeReceiver.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/receiver/protocol/thrift/IoTDBDataNodeReceiver.java index 3cf0b434fea4c..21cab0a66fbe9 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/receiver/protocol/thrift/IoTDBDataNodeReceiver.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/receiver/protocol/thrift/IoTDBDataNodeReceiver.java @@ -152,6 +152,7 @@ public class IoTDBDataNodeReceiver extends IoTDBFileReceiver { private static final SessionManager SESSION_MANAGER = SessionManager.getInstance(); private PipeMemoryBlock allocatedMemoryBlock; + private final List allocatedSliceMemoryBlocks = new ArrayList<>(); static { try { @@ -173,7 +174,7 @@ public synchronized TPipeTransferResp receive(final TPipeTransferReq req) { if (PipeRequestType.isValidatedRequestType(rawRequestType)) { final PipeRequestType requestType = PipeRequestType.valueOf(rawRequestType); if (requestType != PipeRequestType.TRANSFER_SLICE) { - sliceReqHandler.clear(); + clearSliceReqHandler(); } switch (requestType) { case HANDSHAKE_DATANODE_V1: @@ -356,8 +357,18 @@ public synchronized TPipeTransferResp receive(final TPipeTransferReq req) { } case TRANSFER_COMPRESSED: { + long requestedMemorySizeInBytes = 0; try { - return receive(PipeTransferCompressedReq.fromTPipeTransferReq(req)); + requestedMemorySizeInBytes = + PipeTransferCompressedReq.getMaxAdditionalDecompressedLengthInBytes(req); + try (final PipeMemoryBlock ignored = + tryAllocateReceiverMemory(requestedMemorySizeInBytes)) { + return receive(PipeTransferCompressedReq.fromTPipeTransferReq(req)); + } + } catch (final PipeRuntimeOutOfMemoryCriticalException e) { + return new TPipeTransferResp( + getReceiverTemporaryUnavailableStatus( + "decompressing pipe transfer request", requestedMemorySizeInBytes, e)); } finally { PipeDataNodeReceiverMetrics.getInstance() .recordTransferCompressedTimer(System.nanoTime() - startTime); @@ -669,22 +680,97 @@ private String getConfigReceiverId() { } private TPipeTransferResp handleTransferSlice(final PipeTransferSliceReq pipeTransferSliceReq) { - final boolean isInorder = sliceReqHandler.receiveSlice(pipeTransferSliceReq); - if (!isInorder) { + final long sliceBodySizeInBytes = getSliceBodySizeInBytes(pipeTransferSliceReq); + long requestedMemorySizeInBytes = sliceBodySizeInBytes; + String memoryAction = "buffering sliced pipe transfer request"; + PipeMemoryBlock sliceMemoryBlock = null; + try { + sliceMemoryBlock = tryAllocateReceiverMemory(sliceBodySizeInBytes); + + final boolean isInorder = sliceReqHandler.receiveSlice(pipeTransferSliceReq); + if (!isInorder) { + closeMemoryBlock(sliceMemoryBlock); + clearSliceReqHandler(); + return new TPipeTransferResp( + RpcUtils.getStatus( + TSStatusCode.PIPE_TRANSFER_SLICE_OUT_OF_ORDER, + "Slice request is out of order, please check the request sequence.")); + } + + allocatedSliceMemoryBlocks.add(sliceMemoryBlock); + sliceMemoryBlock = null; + + if (pipeTransferSliceReq.getSliceIndex() + 1 < pipeTransferSliceReq.getSliceCount()) { + return new TPipeTransferResp( + RpcUtils.getStatus( + TSStatusCode.SUCCESS_STATUS, + "Slice received, waiting for more slices to complete the request.")); + } + + memoryAction = "assembling sliced pipe transfer request"; + requestedMemorySizeInBytes = pipeTransferSliceReq.getOriginBodySize(); + try (final PipeMemoryBlock ignored = tryAllocateReceiverMemory(requestedMemorySizeInBytes)) { + final Optional req = sliceReqHandler.makeReqIfComplete(); + if (!req.isPresent()) { + return new TPipeTransferResp( + RpcUtils.getStatus( + TSStatusCode.SUCCESS_STATUS, + "Slice received, waiting for more slices to complete the request.")); + } + clearSliceReqHandler(); + return receive(req.get()); + } + } catch (final PipeRuntimeOutOfMemoryCriticalException e) { + closeMemoryBlock(sliceMemoryBlock); + clearSliceReqHandler(); return new TPipeTransferResp( - RpcUtils.getStatus( - TSStatusCode.PIPE_TRANSFER_SLICE_OUT_OF_ORDER, - "Slice request is out of order, please check the request sequence.")); + getReceiverTemporaryUnavailableStatus(memoryAction, requestedMemorySizeInBytes, e)); + } catch (final RuntimeException e) { + closeMemoryBlock(sliceMemoryBlock); + clearSliceReqHandler(); + throw e; } - final Optional req = sliceReqHandler.makeReqIfComplete(); - if (!req.isPresent()) { - return new TPipeTransferResp( - RpcUtils.getStatus( - TSStatusCode.SUCCESS_STATUS, - "Slice received, waiting for more slices to complete the request.")); + } + + private long getSliceBodySizeInBytes(final PipeTransferSliceReq pipeTransferSliceReq) { + return pipeTransferSliceReq.getSliceBody() == null + ? 0 + : pipeTransferSliceReq.getSliceBody().length; + } + + private void clearSliceReqHandler() { + sliceReqHandler.clear(); + allocatedSliceMemoryBlocks.forEach(this::closeMemoryBlock); + allocatedSliceMemoryBlocks.clear(); + } + + private void closeMemoryBlock(final PipeMemoryBlock memoryBlock) { + if (Objects.nonNull(memoryBlock)) { + memoryBlock.close(); } - // sliceReqHandler will be cleared in the receive(req) method - return receive(req.get()); + } + + private PipeMemoryBlock tryAllocateReceiverMemory(final long requestedMemorySizeInBytes) + throws PipeRuntimeOutOfMemoryCriticalException { + return PipeDataNodeResourceManager.memory() + .forceAllocate(Math.max(requestedMemorySizeInBytes, 0)); + } + + @Override + protected TSStatus getReceiverTemporaryUnavailableStatus( + final String action, + final long requestedMemorySizeInBytes, + final PipeRuntimeOutOfMemoryCriticalException e) { + return new TSStatus(TSStatusCode.PIPE_RECEIVER_TEMPORARY_UNAVAILABLE_EXCEPTION.getStatusCode()) + .setMessage( + String.format( + "Temporarily out of memory when %s. Requested memory: %d bytes, used memory: %d " + + "bytes, free memory: %d bytes, total non-floating memory: %d bytes", + action, + requestedMemorySizeInBytes, + PipeDataNodeResourceManager.memory().getUsedMemorySizeInBytes(), + PipeDataNodeResourceManager.memory().getFreeMemorySizeInBytes(), + PipeDataNodeResourceManager.memory().getTotalNonFloatingMemorySizeInBytes())); } /** @@ -910,6 +996,7 @@ protected TSStatus login() { @Override public synchronized void handleExit() { + clearSliceReqHandler(); if (Objects.nonNull(configReceiverId.get())) { try { ClusterConfigTaskExecutor.getInstance().handlePipeConfigClientExit(configReceiverId.get()); diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/receiver/IoTDBFileReceiver.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/receiver/IoTDBFileReceiver.java index 7346ae2bccd48..431035502d329 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/receiver/IoTDBFileReceiver.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/receiver/IoTDBFileReceiver.java @@ -22,6 +22,7 @@ import org.apache.iotdb.common.rpc.thrift.TSStatus; import org.apache.iotdb.commons.conf.CommonDescriptor; import org.apache.iotdb.commons.exception.IllegalPathException; +import org.apache.iotdb.commons.exception.pipe.PipeRuntimeOutOfMemoryCriticalException; import org.apache.iotdb.commons.pipe.config.PipeConfig; import org.apache.iotdb.commons.pipe.config.constant.PipeSinkConstant; import org.apache.iotdb.commons.pipe.resource.log.PipeLogger; @@ -375,7 +376,7 @@ protected final TPipeTransferResp handleTransferFilePiece( final PipeTransferFilePieceReq req, final boolean isRequestThroughAirGap, final boolean isSingleFile) { - try { + try (final AutoCloseable ignored = tryAllocateMemoryForFilePiece(req)) { updateWritingFileIfNeeded(req.getFileName(), isSingleFile); // If the request is through air gap, the sender will resend the file piece from the beginning @@ -410,6 +411,22 @@ protected final TPipeTransferResp handleTransferFilePiece( writingFileWriter.write(req.getFilePiece()); return PipeTransferFilePieceResp.toTPipeTransferResp( RpcUtils.SUCCESS_STATUS, writingFileWriter.length()); + } catch (final PipeRuntimeOutOfMemoryCriticalException e) { + final TSStatus status = + getReceiverTemporaryUnavailableStatus( + "receiving pipe file piece", getFilePieceSizeInBytes(req), e); + PipeLogger.log( + LOGGER::warn, + e, + "Receiver id = %s: Failed to write file piece from req %s.", + receiverId.get(), + req); + try { + return PipeTransferFilePieceResp.toTPipeTransferResp( + status, PipeTransferFilePieceResp.ERROR_END_OFFSET); + } catch (Exception ex) { + return PipeTransferFilePieceResp.toTPipeTransferResp(status); + } } catch (final Exception e) { PipeLogger.log( LOGGER::warn, @@ -430,6 +447,26 @@ protected final TPipeTransferResp handleTransferFilePiece( } } + protected AutoCloseable tryAllocateMemoryForFilePiece(final PipeTransferFilePieceReq req) + throws PipeRuntimeOutOfMemoryCriticalException { + return () -> {}; + } + + protected TSStatus getReceiverTemporaryUnavailableStatus( + final String action, + final long requestedMemorySizeInBytes, + final PipeRuntimeOutOfMemoryCriticalException e) { + return new TSStatus(TSStatusCode.PIPE_RECEIVER_TEMPORARY_UNAVAILABLE_EXCEPTION.getStatusCode()) + .setMessage( + String.format( + "Temporarily out of memory when %s. Requested memory: %d bytes. Root cause: %s", + action, requestedMemorySizeInBytes, e.getMessage())); + } + + private static long getFilePieceSizeInBytes(final PipeTransferFilePieceReq req) { + return req.getFilePiece() == null ? 0 : req.getFilePiece().length; + } + protected final void updateWritingFileIfNeeded(final String fileName, final boolean isSingleFile) throws IOException { if (isFileExistedAndNameCorrect(fileName)) { diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/payload/thrift/common/PipeTransferSliceReqHandler.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/payload/thrift/common/PipeTransferSliceReqHandler.java index 09edb4bbe4a12..ead9ea277d623 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/payload/thrift/common/PipeTransferSliceReqHandler.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/payload/thrift/common/PipeTransferSliceReqHandler.java @@ -42,8 +42,14 @@ public class PipeTransferSliceReqHandler { private int sliceCount = -1; private final List sliceBodies = new ArrayList<>(); + private long receivedBodySize = 0; public boolean receiveSlice(final PipeTransferSliceReq req) { + if (!isValidSliceReq(req)) { + clear(); + return false; + } + if (orderId == -1 || originReqType == -1 || originBodySize == -1 @@ -58,6 +64,7 @@ public boolean receiveSlice(final PipeTransferSliceReq req) { originReqType = req.getOriginReqType(); originBodySize = req.getOriginBodySize(); sliceCount = req.getSliceCount(); + receivedBodySize = 0; } else { LOGGER.warn( "Invalid state: orderId={}, originReqType={}, originBodySize={}, sliceCount={}, sliceBodies.size={}", @@ -104,10 +111,39 @@ public boolean receiveSlice(final PipeTransferSliceReq req) { return false; } + if (receivedBodySize + req.getSliceBody().length > originBodySize) { + LOGGER.warn( + "Received slice body size {} exceeds origin body size {}", + receivedBodySize + req.getSliceBody().length, + originBodySize); + clear(); + return false; + } + sliceBodies.add(req.getSliceBody()); + receivedBodySize += req.getSliceBody().length; + + if (sliceBodies.size() == sliceCount && receivedBodySize != originBodySize) { + LOGGER.warn( + "Received slice body size {} is not equal to origin body size {}", + receivedBodySize, + originBodySize); + clear(); + return false; + } + return true; } + private boolean isValidSliceReq(final PipeTransferSliceReq req) { + return req.getOriginBodySize() >= 0 + && req.getSliceCount() > 0 + && req.getSliceIndex() >= 0 + && req.getSliceIndex() < req.getSliceCount() + && req.getSliceBody() != null + && req.getSliceBody().length <= req.getOriginBodySize(); + } + public Optional makeReqIfComplete() { if (sliceBodies.size() != sliceCount) { return Optional.empty(); @@ -131,5 +167,6 @@ public void clear() { originBodySize = -1; sliceCount = -1; sliceBodies.clear(); + receivedBodySize = 0; } } diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/payload/thrift/request/PipeTransferCompressedReq.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/payload/thrift/request/PipeTransferCompressedReq.java index 73bac4c05b401..7aa639140beca 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/payload/thrift/request/PipeTransferCompressedReq.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/payload/thrift/request/PipeTransferCompressedReq.java @@ -112,6 +112,27 @@ public static TPipeTransferReq fromTPipeTransferReq(final TPipeTransferReq trans return decompressedReq; } + /** Get the largest intermediate decompressed body size without consuming the request body. */ + public static int getMaxDecompressedLengthInBytes(final TPipeTransferReq transferReq) { + final ByteBuffer compressedBuffer = transferReq.body.duplicate(); + + int maxDecompressedLength = compressedBuffer.remaining(); + final int compressorsSize = ReadWriteIOUtils.readByte(compressedBuffer); + for (int i = 0; i < compressorsSize; ++i) { + ReadWriteIOUtils.readByte(compressedBuffer); + final int decompressedLength = ReadWriteIOUtils.readInt(compressedBuffer); + checkDecompressedLength(decompressedLength); + maxDecompressedLength = Math.max(maxDecompressedLength, decompressedLength); + } + return maxDecompressedLength; + } + + /** Get the largest additional decompressed body size beyond the current transfer frame. */ + public static int getMaxAdditionalDecompressedLengthInBytes(final TPipeTransferReq transferReq) { + final int transferFrameBodySize = transferReq.body.duplicate().remaining(); + return Math.max(0, getMaxDecompressedLengthInBytes(transferReq) - transferFrameBodySize); + } + /** This method is used to prevent decompression bomb attacks. */ private static void checkDecompressedLength(final int decompressedLength) throws IllegalArgumentException { diff --git a/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/pipe/receiver/IoTDBFileReceiverTest.java b/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/pipe/receiver/IoTDBFileReceiverTest.java index c692cb1513a4f..6601489e69f32 100644 --- a/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/pipe/receiver/IoTDBFileReceiverTest.java +++ b/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/pipe/receiver/IoTDBFileReceiverTest.java @@ -22,10 +22,13 @@ import org.apache.iotdb.common.rpc.thrift.TSStatus; import org.apache.iotdb.commons.conf.CommonDescriptor; import org.apache.iotdb.commons.exception.IllegalPathException; +import org.apache.iotdb.commons.exception.pipe.PipeRuntimeOutOfMemoryCriticalException; import org.apache.iotdb.commons.pipe.sink.payload.thrift.request.PipeRequestType; +import org.apache.iotdb.commons.pipe.sink.payload.thrift.request.PipeTransferFilePieceReq; import org.apache.iotdb.commons.pipe.sink.payload.thrift.request.PipeTransferFileSealReqV1; import org.apache.iotdb.commons.pipe.sink.payload.thrift.request.PipeTransferFileSealReqV2; import org.apache.iotdb.commons.pipe.sink.payload.thrift.request.PipeTransferHandshakeV1Req; +import org.apache.iotdb.commons.pipe.sink.payload.thrift.response.PipeTransferFilePieceResp; import org.apache.iotdb.rpc.TSStatusCode; import org.apache.iotdb.service.rpc.thrift.TPipeTransferReq; import org.apache.iotdb.service.rpc.thrift.TPipeTransferResp; @@ -138,10 +141,55 @@ public void testSealFileV1FailureDeletesTransferredFile() throws Exception { } } + @Test + public void testFilePieceMemoryAllocationFailureReturnsTemporaryUnavailable() throws Exception { + final Path baseDir = Files.createTempDirectory("iotdb-file-receiver-test"); + final DummyFileReceiver receiver = new DummyFileReceiver(baseDir.toFile()); + try { + receiver.setFailFilePieceMemoryAllocation(true); + + final TPipeTransferResp response = + receiver.writeFilePiece("normal.tsfile", 0, new byte[] {1, 2, 3}); + final PipeTransferFilePieceResp filePieceResp = + PipeTransferFilePieceResp.fromTPipeTransferResp(response); + + Assert.assertEquals( + TSStatusCode.PIPE_RECEIVER_TEMPORARY_UNAVAILABLE_EXCEPTION.getStatusCode(), + response.getStatus().getCode()); + Assert.assertTrue(response.getStatus().getMessage().contains("no memory for file piece")); + Assert.assertEquals( + PipeTransferFilePieceResp.ERROR_END_OFFSET, filePieceResp.getEndWritingOffset()); + Assert.assertFalse(receiver.getWritingFileInBaseDir("normal.tsfile").exists()); + } finally { + receiver.handleExit(); + } + } + + @Test + public void testFilePieceMemoryAllocationIsClosedAfterWrite() throws Exception { + final Path baseDir = Files.createTempDirectory("iotdb-file-receiver-test"); + final DummyFileReceiver receiver = new DummyFileReceiver(baseDir.toFile()); + try { + final TPipeTransferResp response = + receiver.writeFilePiece("normal.tsfile", 0, new byte[] {1, 2, 3}); + final PipeTransferFilePieceResp filePieceResp = + PipeTransferFilePieceResp.fromTPipeTransferResp(response); + + Assert.assertEquals( + TSStatusCode.SUCCESS_STATUS.getStatusCode(), response.getStatus().getCode()); + Assert.assertEquals(3, filePieceResp.getEndWritingOffset()); + Assert.assertEquals(1, receiver.getFilePieceMemoryCloseCount()); + } finally { + receiver.handleExit(); + } + } + private static class DummyFileReceiver extends IoTDBFileReceiver { private final File receiverFileBaseDir; private TSStatus loadFileV1Status = new TSStatus(TSStatusCode.SUCCESS_STATUS.getStatusCode()); + private boolean failFilePieceMemoryAllocation = false; + private int filePieceMemoryCloseCount = 0; DummyFileReceiver(final File baseDir) { receiverFileBaseDir = baseDir; @@ -166,6 +214,23 @@ void setLoadFileV1Status(final TSStatus status) { loadFileV1Status = status; } + void setFailFilePieceMemoryAllocation(final boolean failFilePieceMemoryAllocation) { + this.failFilePieceMemoryAllocation = failFilePieceMemoryAllocation; + } + + int getFilePieceMemoryCloseCount() { + return filePieceMemoryCloseCount; + } + + TPipeTransferResp writeFilePiece( + final String fileName, final long startWritingOffset, final byte[] filePiece) + throws IOException { + return handleTransferFilePiece( + DummyFilePieceReq.toTPipeTransferReq(fileName, startWritingOffset, filePiece), + false, + true); + } + TPipeTransferResp sealFileV1(final String fileName, final long fileLength) throws IOException { return handleTransferFileSealV1(DummyFileSealReqV1.toTPipeTransferReq(fileName, fileLength)); } @@ -223,6 +288,14 @@ protected TSStatus login() { return new TSStatus(200); } + @Override + protected AutoCloseable tryAllocateMemoryForFilePiece(final PipeTransferFilePieceReq req) { + if (failFilePieceMemoryAllocation) { + throw new PipeRuntimeOutOfMemoryCriticalException("no memory for file piece"); + } + return () -> filePieceMemoryCloseCount++; + } + @Override protected TSStatus loadFileV1( final PipeTransferFileSealReqV1 req, final String fileAbsolutePath) { @@ -247,6 +320,22 @@ public TPipeTransferResp receive(TPipeTransferReq req) { } } + private static class DummyFilePieceReq extends PipeTransferFilePieceReq { + + static DummyFilePieceReq toTPipeTransferReq( + final String fileName, final long startWritingOffset, final byte[] filePiece) + throws IOException { + return (DummyFilePieceReq) + new DummyFilePieceReq() + .convertToTPipeTransferReq(fileName, startWritingOffset, filePiece); + } + + @Override + protected PipeRequestType getPlanType() { + return PipeRequestType.TRANSFER_TS_FILE_PIECE; + } + } + private static class DummyHandshakeReq extends PipeTransferHandshakeV1Req { static DummyHandshakeReq toTPipeTransferReq(final String timestampPrecision) diff --git a/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/pipe/sink/payload/thrift/common/PipeTransferSliceReqBuilderTest.java b/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/pipe/sink/payload/thrift/common/PipeTransferSliceReqBuilderTest.java index 290ce3979807d..857a62ac75cfb 100644 --- a/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/pipe/sink/payload/thrift/common/PipeTransferSliceReqBuilderTest.java +++ b/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/pipe/sink/payload/thrift/common/PipeTransferSliceReqBuilderTest.java @@ -30,7 +30,9 @@ import org.junit.Before; import org.junit.Test; +import java.io.IOException; import java.nio.ByteBuffer; +import java.util.Optional; public class PipeTransferSliceReqBuilderTest { @@ -91,6 +93,57 @@ public void testShouldSliceOnlyForVersion1RequestsAboveThreshold() { createReq(IoTDBSinkRequestVersion.VERSION_1.getVersion(), 4), bodySizeLimit)); } + @Test + public void testSliceReqHandlerRejectsOversizedSlices() throws IOException { + final TPipeTransferReq req = createReq(IoTDBSinkRequestVersion.VERSION_1.getVersion(), 6); + final PipeTransferSliceReq firstSlice = + PipeTransferSliceReq.toTPipeTransferReq(7, req.getType(), 0, 2, req.body.duplicate(), 0, 4); + final PipeTransferSliceReq oversizedSecondSlice = + PipeTransferSliceReq.toTPipeTransferReq(7, req.getType(), 1, 2, req.body.duplicate(), 2, 6); + + final PipeTransferSliceReqHandler handler = new PipeTransferSliceReqHandler(); + Assert.assertTrue(handler.receiveSlice(firstSlice)); + Assert.assertFalse(handler.receiveSlice(oversizedSecondSlice)); + Assert.assertFalse(handler.makeReqIfComplete().isPresent()); + } + + @Test + public void testSliceReqHandlerAssemblesCompleteRequest() throws IOException { + final TPipeTransferReq req = createReq(IoTDBSinkRequestVersion.VERSION_1.getVersion(), 6); + final PipeTransferSliceReq firstSlice = + PipeTransferSliceReq.toTPipeTransferReq(7, req.getType(), 0, 2, req.body.duplicate(), 0, 4); + final PipeTransferSliceReq secondSlice = + PipeTransferSliceReq.toTPipeTransferReq(7, req.getType(), 1, 2, req.body.duplicate(), 4, 6); + + final PipeTransferSliceReqHandler handler = new PipeTransferSliceReqHandler(); + Assert.assertTrue(handler.receiveSlice(firstSlice)); + Assert.assertFalse(handler.makeReqIfComplete().isPresent()); + Assert.assertTrue(handler.receiveSlice(secondSlice)); + + final Optional completedReq = handler.makeReqIfComplete(); + Assert.assertTrue(completedReq.isPresent()); + + final TPipeTransferReq assembledReq = completedReq.get(); + Assert.assertEquals(IoTDBSinkRequestVersion.VERSION_1.getVersion(), assembledReq.getVersion()); + Assert.assertEquals(req.getType(), assembledReq.getType()); + + final byte[] assembledBody = new byte[assembledReq.body.remaining()]; + assembledReq.body.get(assembledBody); + Assert.assertArrayEquals(new byte[] {0, 1, 2, 3, 4, 5}, assembledBody); + } + + @Test + public void testSliceReqHandlerRejectsInvalidMetadata() throws IOException { + final TPipeTransferReq req = createReq(IoTDBSinkRequestVersion.VERSION_1.getVersion(), 2); + final PipeTransferSliceReq invalidSlice = + PipeTransferSliceReq.toTPipeTransferReq(7, req.getType(), 0, 0, req.body.duplicate(), 0, 1); + + final PipeTransferSliceReqHandler handler = new PipeTransferSliceReqHandler(); + + Assert.assertFalse(handler.receiveSlice(invalidSlice)); + Assert.assertFalse(handler.makeReqIfComplete().isPresent()); + } + private static TPipeTransferReq createReq(final byte version, final int bodySize) { final byte[] body = new byte[bodySize]; for (int i = 0; i < body.length; ++i) { diff --git a/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/pipe/sink/payload/thrift/request/PipeTransferCompressedReqTest.java b/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/pipe/sink/payload/thrift/request/PipeTransferCompressedReqTest.java new file mode 100644 index 0000000000000..a00530b8c6800 --- /dev/null +++ b/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/pipe/sink/payload/thrift/request/PipeTransferCompressedReqTest.java @@ -0,0 +1,149 @@ +/* + * 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.commons.pipe.sink.compressor.PipeCompressor; +import org.apache.iotdb.commons.pipe.sink.compressor.PipeCompressorFactory; +import org.apache.iotdb.service.rpc.thrift.TPipeTransferReq; + +import org.apache.tsfile.utils.BytesUtils; +import org.apache.tsfile.utils.PublicBAOS; +import org.apache.tsfile.utils.ReadWriteIOUtils; +import org.junit.Assert; +import org.junit.Test; + +import java.io.DataOutputStream; +import java.io.IOException; +import java.nio.ByteBuffer; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; + +public class PipeTransferCompressedReqTest { + + @Test + public void testPipeTransferCompressedReq() throws IOException { + final TPipeTransferReq originalReq = new TPipeTransferReq(); + originalReq.version = IoTDBSinkRequestVersion.VERSION_1.getVersion(); + originalReq.type = PipeRequestType.TRANSFER_TABLET_BINARY.getType(); + originalReq.body = ByteBuffer.wrap(new byte[] {1, 2, 3, 4}); + + final TPipeTransferReq compressedReq = + PipeTransferCompressedReq.toTPipeTransferReq( + originalReq, + Collections.singletonList( + PipeCompressorFactory.getCompressor( + PipeCompressor.PipeCompressionType.GZIP.getIndex()))); + final int compressedBodySize = compressedReq.body.remaining(); + final int maxDecompressedLength = + Math.max(compressedBodySize, originalReq.body.remaining() + 3); + Assert.assertEquals( + maxDecompressedLength, + PipeTransferCompressedReq.getMaxDecompressedLengthInBytes(compressedReq)); + Assert.assertEquals( + maxDecompressedLength - compressedBodySize, + PipeTransferCompressedReq.getMaxAdditionalDecompressedLengthInBytes(compressedReq)); + Assert.assertEquals(0, compressedReq.body.position()); + + final TPipeTransferReq decompressedReq = + PipeTransferCompressedReq.fromTPipeTransferReq(compressedReq); + + Assert.assertEquals(IoTDBSinkRequestVersion.VERSION_1.getVersion(), compressedReq.version); + Assert.assertEquals(PipeRequestType.TRANSFER_COMPRESSED.getType(), compressedReq.type); + Assert.assertEquals(originalReq.version, decompressedReq.version); + Assert.assertEquals(originalReq.type, decompressedReq.type); + Assert.assertArrayEquals(originalReq.getBody(), decompressedReq.getBody()); + } + + @Test + public void testAdditionalDecompressedLengthExcludesTransferFrameBody() throws IOException { + final TPipeTransferReq originalReq = new TPipeTransferReq(); + originalReq.version = IoTDBSinkRequestVersion.VERSION_1.getVersion(); + originalReq.type = PipeRequestType.TRANSFER_TABLET_BINARY.getType(); + final byte[] highlyCompressibleBody = new byte[16 * 1024]; + Arrays.fill(highlyCompressibleBody, (byte) 1); + originalReq.body = ByteBuffer.wrap(highlyCompressibleBody); + + final TPipeTransferReq compressedReq = + PipeTransferCompressedReq.toTPipeTransferReq( + originalReq, + Collections.singletonList( + PipeCompressorFactory.getCompressor( + PipeCompressor.PipeCompressionType.GZIP.getIndex()))); + final int compressedBodySize = compressedReq.body.remaining(); + final int maxDecompressedLength = + PipeTransferCompressedReq.getMaxDecompressedLengthInBytes(compressedReq); + + Assert.assertTrue(maxDecompressedLength > compressedBodySize); + Assert.assertEquals( + maxDecompressedLength - compressedBodySize, + PipeTransferCompressedReq.getMaxAdditionalDecompressedLengthInBytes(compressedReq)); + Assert.assertEquals(0, compressedReq.body.position()); + } + + @Test + public void testPipeTransferCompressedReqFromLegacyV13Body() throws IOException { + final TPipeTransferReq originalReq = new TPipeTransferReq(); + originalReq.version = IoTDBSinkRequestVersion.VERSION_1.getVersion(); + originalReq.type = PipeRequestType.TRANSFER_TABLET_BINARY.getType(); + originalReq.body = ByteBuffer.wrap(new byte[] {1, 2, 3, 4}); + + final TPipeTransferReq compressedReq = new TPipeTransferReq(); + compressedReq.version = IoTDBSinkRequestVersion.VERSION_1.getVersion(); + compressedReq.type = PipeRequestType.TRANSFER_COMPRESSED.getType(); + compressedReq.body = + serializeLegacyCompressedBody( + originalReq, + Collections.singletonList( + PipeCompressorFactory.getCompressor( + PipeCompressor.PipeCompressionType.GZIP.getIndex()))); + + final TPipeTransferReq decompressedReq = + PipeTransferCompressedReq.fromTPipeTransferReq(compressedReq); + + Assert.assertEquals(originalReq.version, decompressedReq.version); + Assert.assertEquals(originalReq.type, decompressedReq.type); + Assert.assertArrayEquals(originalReq.getBody(), decompressedReq.getBody()); + } + + private static ByteBuffer serializeLegacyCompressedBody( + final TPipeTransferReq originalReq, final List compressors) + throws IOException { + try (final PublicBAOS byteArrayOutputStream = new PublicBAOS(); + final DataOutputStream outputStream = new DataOutputStream(byteArrayOutputStream)) { + byte[] body = + BytesUtils.concatByteArrayList( + Arrays.asList( + new byte[] {originalReq.version}, + BytesUtils.shortToBytes(originalReq.type), + originalReq.getBody())); + + ReadWriteIOUtils.write((byte) compressors.size(), outputStream); + for (final PipeCompressor compressor : compressors) { + ReadWriteIOUtils.write(compressor.serialize(), outputStream); + ReadWriteIOUtils.write(body.length, outputStream); + body = compressor.compress(body); + } + outputStream.write(body); + + return ByteBuffer.wrap(byteArrayOutputStream.getBuf(), 0, byteArrayOutputStream.size()); + } + } +}