diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/receiver/protocol/airgap/IoTDBAirGapReceiver.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/receiver/protocol/airgap/IoTDBAirGapReceiver.java index c337f1c60ef9d..9784efee5b24e 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/receiver/protocol/airgap/IoTDBAirGapReceiver.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/receiver/protocol/airgap/IoTDBAirGapReceiver.java @@ -40,10 +40,12 @@ import org.slf4j.LoggerFactory; import java.io.BufferedInputStream; +import java.io.FilterInputStream; import java.io.IOException; import java.io.InputStream; import java.io.OutputStream; import java.net.Socket; +import java.net.SocketTimeoutException; import java.nio.ByteBuffer; import java.util.Arrays; import java.util.zip.CRC32; @@ -97,7 +99,9 @@ public void runMayThrow() throws Throwable { } private void receive() throws IOException { - final InputStream inputStream = new BufferedInputStream(socket.getInputStream()); + final ReadProgressInputStream readProgressInputStream = + new ReadProgressInputStream(socket.getInputStream()); + final InputStream inputStream = new BufferedInputStream(readProgressInputStream); try { final byte[] data = readData(inputStream); @@ -128,6 +132,18 @@ private void receive() throws IOException { .setType(ReadWriteIOUtils.readShort(byteBuffer)) .setBody(byteBuffer.slice()); handleReq(req, System.currentTimeMillis()); + } catch (final SocketTimeoutException e) { + // It is normal for an air gap sender to remain idle. Only close the connection when the + // timeout occurs after a request has started, because the stream can no longer be decoded + // reliably in that case. Do not send FAIL without receiving a complete request. + if (readProgressInputStream.hasReadAnyByte()) { + LOGGER.warn( + DataNodePipeMessages.PIPE_AIR_GAP_RECEIVER_EXCEPTION_DURING_HANDLING, + receiverId, + socket, + e); + socket.close(); + } } catch (final PipeConnectionException e) { LOGGER.info( DataNodePipeMessages.PIPE_AIR_GAP_RECEIVER_SOCKET_CLOSED_WHEN, @@ -326,4 +342,35 @@ private void skipTillEnough(final InputStream inputStream, final long length) currentSkippedBytes += skippedBytes; } } + + private static class ReadProgressInputStream extends FilterInputStream { + + private boolean hasReadAnyByte; + + private ReadProgressInputStream(final InputStream inputStream) { + super(inputStream); + } + + @Override + public int read() throws IOException { + final int result = super.read(); + if (result >= 0) { + hasReadAnyByte = true; + } + return result; + } + + @Override + public int read(final byte[] buffer, final int offset, final int length) throws IOException { + final int result = super.read(buffer, offset, length); + if (result > 0) { + hasReadAnyByte = true; + } + return result; + } + + private boolean hasReadAnyByte() { + return hasReadAnyByte; + } + } } diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/receiver/protocol/airgap/IoTDBAirGapReceiverTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/receiver/protocol/airgap/IoTDBAirGapReceiverTest.java index e23db1f1ca8cd..1687adeafcd7c 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/receiver/protocol/airgap/IoTDBAirGapReceiverTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/receiver/protocol/airgap/IoTDBAirGapReceiverTest.java @@ -44,6 +44,7 @@ import java.lang.reflect.Field; import java.lang.reflect.Method; import java.net.Socket; +import java.net.SocketTimeoutException; import java.nio.ByteBuffer; public class IoTDBAirGapReceiverTest { @@ -121,6 +122,31 @@ public void testTemporaryUnavailableRetryTimeoutReturnsFail() throws Exception { } } + @Test + public void testIdleReadTimeoutDoesNotRespondOrCloseSocket() throws Exception { + final RecordingSocket socket = new RecordingSocket(new TimeoutInputStream(new byte[0])); + invokeReceive(new IoTDBAirGapReceiver(socket, 4L)); + + Assert.assertArrayEquals(new byte[0], socket.getWrittenBytes()); + Assert.assertFalse(socket.isClosed()); + } + + @Test + public void testPartialRequestReadTimeoutClosesSocketWithoutResponse() throws Exception { + final RecordingSocket socket = + new RecordingSocket(new TimeoutInputStream(new byte[] {(byte) 0xFF})); + invokeReceive(new IoTDBAirGapReceiver(socket, 5L)); + + Assert.assertArrayEquals(new byte[0], socket.getWrittenBytes()); + Assert.assertTrue(socket.isClosed()); + } + + private static void invokeReceive(final IoTDBAirGapReceiver receiver) throws Exception { + final Method receive = IoTDBAirGapReceiver.class.getDeclaredMethod("receive"); + receive.setAccessible(true); + receive.invoke(receiver); + } + private static void setField(final Object target, final String fieldName, final Object value) throws Exception { final Field field = IoTDBAirGapReceiver.class.getDeclaredField(fieldName); @@ -130,8 +156,22 @@ private static void setField(final Object target, final String fieldName, final private static class RecordingSocket extends Socket { + private final InputStream inputStream; private final ByteArrayOutputStream outputStream = new ByteArrayOutputStream(); + private RecordingSocket() { + this(new ByteArrayInputStream(new byte[0])); + } + + private RecordingSocket(final InputStream inputStream) { + this.inputStream = inputStream; + } + + @Override + public InputStream getInputStream() { + return inputStream; + } + @Override public OutputStream getOutputStream() { return outputStream; @@ -142,6 +182,35 @@ byte[] getWrittenBytes() { } } + private static class TimeoutInputStream extends InputStream { + + private final byte[] bytesBeforeTimeout; + private int position; + + private TimeoutInputStream(final byte[] bytesBeforeTimeout) { + this.bytesBeforeTimeout = bytesBeforeTimeout; + } + + @Override + public int read() throws IOException { + if (position >= bytesBeforeTimeout.length) { + throw new SocketTimeoutException("Test timeout"); + } + return bytesBeforeTimeout[position++] & 0xFF; + } + + @Override + public int read(final byte[] buffer, final int offset, final int length) throws IOException { + if (position >= bytesBeforeTimeout.length) { + throw new SocketTimeoutException("Test timeout"); + } + final int bytesToRead = Math.min(length, bytesBeforeTimeout.length - position); + System.arraycopy(bytesBeforeTimeout, position, buffer, offset, bytesToRead); + position += bytesToRead; + return bytesToRead; + } + } + private static class StubIoTDBDataNodeReceiverAgent extends IoTDBDataNodeReceiverAgent { void setStubReceiver(final IoTDBReceiver receiver) {