Fix unsolicited AirGap responses on idle timeout

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;
+    }
+  }
 }

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) {