From 3144e61c0ff37193181969a04737c330c44b0e71 Mon Sep 17 00:00:00 2001 From: Caideyipi <87789683+Caideyipi@users.noreply.github.com> Date: Tue, 16 Jun 2026 15:44:06 +0800 Subject: [PATCH 1/7] Implement UDP transport for air-gap sink --- .../protocol/airgap/IoTDBAirGapReceiver.java | 131 +++++++--- .../airgap/IoTDBAirGapReceiverAgent.java | 106 +++++++- .../thrift/IoTDBDataNodeReceiverAgent.java | 18 +- .../airgap/IoTDBAirGapReceiverTest.java | 37 +++ .../airgap/IoTDBDataRegionAirGapSinkTest.java | 107 +++++++- .../config/constant/PipeSinkConstant.java | 20 ++ .../pipe/sink/protocol/IoTDBAirGapSink.java | 240 +++++++++++++++++- 7 files changed, 618 insertions(+), 41 deletions(-) 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 825ff4c5eac2c..ff22ed8d9e0f7 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,14 @@ import org.slf4j.LoggerFactory; import java.io.BufferedInputStream; +import java.io.ByteArrayInputStream; import java.io.IOException; import java.io.InputStream; import java.io.OutputStream; +import java.net.DatagramPacket; +import java.net.DatagramSocket; import java.net.Socket; +import java.net.SocketAddress; import java.nio.ByteBuffer; import java.util.Arrays; import java.util.zip.CRC32; @@ -61,6 +65,9 @@ public class IoTDBAirGapReceiver extends WrappedRunnable { private final IoTDBDataNodeReceiverAgent agent; private boolean isELanguagePayload; + private OutputStream currentOutputStream; + private DatagramSocket currentDatagramSocket; + private SocketAddress currentRemoteSocketAddress; public IoTDBAirGapReceiver(final Socket socket, final long receiverId) { this.socket = socket; @@ -101,33 +108,11 @@ private void receive() throws IOException { try { final byte[] data = readData(inputStream); + currentOutputStream = socket.getOutputStream(); - // If check sum failed, it indicates that the length we read may not be correct. - // Namely, there may be remaining bytes in the socket stream, which will fail any subsequent - // attempts to read from that. - // We directly close the socket here. - if (!checkSum(data)) { - LOGGER.warn( - DataNodePipeMessages.PIPE_AIR_GAP_RECEIVER_CLOSED_BECAUSE_OF, receiverId, socket); - try { - fail(); - } finally { - socket.close(); - } - return; + if (!receive(data, null)) { + socket.close(); } - - // Removed the used checksum - final ByteBuffer byteBuffer = ByteBuffer.wrap(data, LONG_LEN, data.length - LONG_LEN); - - // Pseudo request, to reuse logic in IoTDBThriftReceiverAgent - final AirGapPseudoTPipeTransferRequest req = - (AirGapPseudoTPipeTransferRequest) - new AirGapPseudoTPipeTransferRequest() - .setVersion(ReadWriteIOUtils.readByte(byteBuffer)) - .setType(ReadWriteIOUtils.readShort(byteBuffer)) - .setBody(byteBuffer.slice()); - handleReq(req, System.currentTimeMillis()); } catch (final PipeConnectionException e) { LOGGER.info( DataNodePipeMessages.PIPE_AIR_GAP_RECEIVER_SOCKET_CLOSED_WHEN, @@ -145,9 +130,82 @@ private void receive() throws IOException { } } + boolean receive(final byte[] data, final String receiverKey) throws IOException { + try { + // If check sum failed, it indicates that the length we read may not be correct. + // Namely, there may be remaining bytes in the socket stream, which will fail any subsequent + // attempts to read from that. + if (!checkSum(data)) { + LOGGER.warn( + DataNodePipeMessages.PIPE_AIR_GAP_RECEIVER_CLOSED_BECAUSE_OF, receiverId, socket); + fail(); + return false; + } + + handleReq(toTPipeTransferReq(data), System.currentTimeMillis(), receiverKey); + return true; + } catch (final IOException e) { + throw e; + } catch (final Exception e) { + if (currentDatagramSocket != null) { + fail(); + } + throw e; + } finally { + currentOutputStream = null; + currentDatagramSocket = null; + currentRemoteSocketAddress = null; + } + } + + void receiveUdp( + final DatagramSocket datagramSocket, + final DatagramPacket packet, + final String receiverKey, + final byte[] buffer) + throws IOException { + isELanguagePayload = false; + currentDatagramSocket = datagramSocket; + currentRemoteSocketAddress = packet.getSocketAddress(); + boolean requestHandlingStarted = false; + try { + final byte[] data = readData(new ByteArrayInputStream(buffer, 0, packet.getLength())); + requestHandlingStarted = true; + receive(data, receiverKey); + } catch (final Exception e) { + if (!requestHandlingStarted) { + fail(); + } + throw e; + } finally { + currentOutputStream = null; + currentDatagramSocket = null; + currentRemoteSocketAddress = null; + } + } + + private AirGapPseudoTPipeTransferRequest toTPipeTransferReq(final byte[] data) { + // Removed the used checksum + final ByteBuffer byteBuffer = ByteBuffer.wrap(data, LONG_LEN, data.length - LONG_LEN); + + // Pseudo request, to reuse logic in IoTDBThriftReceiverAgent + return (AirGapPseudoTPipeTransferRequest) + new AirGapPseudoTPipeTransferRequest() + .setVersion(ReadWriteIOUtils.readByte(byteBuffer)) + .setType(ReadWriteIOUtils.readShort(byteBuffer)) + .setBody(byteBuffer.slice()); + } + private void handleReq(final AirGapPseudoTPipeTransferRequest req, final long startTime) throws IOException { - final TPipeTransferResp resp = agent.receive(req); + handleReq(req, startTime, null); + } + + private void handleReq( + final AirGapPseudoTPipeTransferRequest req, final long startTime, final String receiverKey) + throws IOException { + final TPipeTransferResp resp = + receiverKey == null ? agent.receive(req) : agent.receive(receiverKey, req); final TSStatus status = resp.getStatus(); if (status.getCode() == TSStatusCode.SUCCESS_STATUS.getStatusCode()) { @@ -170,7 +228,7 @@ private void handleReq(final AirGapPseudoTPipeTransferRequest req, final long st LOGGER.info(DataNodePipeMessages.TEMPORARY_UNAVAILABLE_EXCEPTION_ENCOUNTERED_AT_AIR_GAP); if (System.currentTimeMillis() - startTime < PipeConfig.getInstance().getPipeAirGapRetryMaxMs()) { - handleReq(req, startTime); + handleReq(req, startTime, receiverKey); } else { LOGGER.warn( DataNodePipeMessages.PIPE_AIR_GAP_RECEIVER_TEMPORARY_UNAVAILABLE_RETRY, receiverId); @@ -187,14 +245,23 @@ private void handleReq(final AirGapPseudoTPipeTransferRequest req, final long st } private void ok() throws IOException { - final OutputStream outputStream = socket.getOutputStream(); - outputStream.write(AirGapOneByteResponse.OK); - outputStream.flush(); + respond(AirGapOneByteResponse.OK); } private void fail() throws IOException { - final OutputStream outputStream = socket.getOutputStream(); - outputStream.write(AirGapOneByteResponse.FAIL); + respond(AirGapOneByteResponse.FAIL); + } + + private void respond(final byte[] response) throws IOException { + if (currentDatagramSocket != null && currentRemoteSocketAddress != null) { + currentDatagramSocket.send( + new DatagramPacket(response, response.length, currentRemoteSocketAddress)); + return; + } + + final OutputStream outputStream = + currentOutputStream != null ? currentOutputStream : socket.getOutputStream(); + outputStream.write(response); outputStream.flush(); } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/receiver/protocol/airgap/IoTDBAirGapReceiverAgent.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/receiver/protocol/airgap/IoTDBAirGapReceiverAgent.java index c14ce2054b7e9..0e8c642ce4716 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/receiver/protocol/airgap/IoTDBAirGapReceiverAgent.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/receiver/protocol/airgap/IoTDBAirGapReceiverAgent.java @@ -26,14 +26,22 @@ import org.apache.iotdb.commons.service.IService; import org.apache.iotdb.commons.service.ServiceType; import org.apache.iotdb.db.i18n.DataNodePipeMessages; +import org.apache.iotdb.db.pipe.agent.PipeDataNodeAgent; +import org.apache.iotdb.db.protocol.session.ClientSession; +import org.apache.iotdb.db.protocol.session.SessionManager; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import java.io.IOException; +import java.net.DatagramPacket; +import java.net.DatagramSocket; +import java.net.InetAddress; import java.net.ServerSocket; import java.net.Socket; +import java.util.Map; import java.util.Objects; +import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ExecutorService; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicLong; @@ -41,15 +49,21 @@ public class IoTDBAirGapReceiverAgent implements IService { private static final Logger LOGGER = LoggerFactory.getLogger(IoTDBAirGapReceiverAgent.class); + private static final int UDP_PACKET_MAX_SIZE_IN_BYTES = 65_507; private final ExecutorService listenExecutor = IoTDBThreadPoolFactory.newSingleThreadExecutor( ThreadName.PIPE_RECEIVER_AIR_GAP_AGENT.getName()); + private final ExecutorService udpListenExecutor = + IoTDBThreadPoolFactory.newSingleThreadExecutor( + ThreadName.PIPE_RECEIVER_AIR_GAP_AGENT.getName() + "-UDP"); private final AtomicBoolean allowSubmitListen = new AtomicBoolean(false); private ServerSocket serverSocket; + private DatagramSocket datagramSocket; private final AtomicLong receiverId = new AtomicLong(0); + private final Map udpClientSessions = new ConcurrentHashMap<>(); public void listen() { try { @@ -61,7 +75,9 @@ public void listen() { ThreadName.PIPE_AIR_GAP_RECEIVER.getName() + "-" + airGapReceiverId); airGapReceiverThread.start(); } catch (final IOException e) { - LOGGER.warn(DataNodePipeMessages.UNHANDLED_EXCEPTION_DURING_PIPE_AIR_GAP_RECEIVER, e); + if (allowSubmitListen.get()) { + LOGGER.warn(DataNodePipeMessages.UNHANDLED_EXCEPTION_DURING_PIPE_AIR_GAP_RECEIVER, e); + } } if (allowSubmitListen.get()) { @@ -69,32 +85,97 @@ public void listen() { } } + public void listenUdp() { + while (allowSubmitListen.get() && !datagramSocket.isClosed()) { + final byte[] buffer = new byte[UDP_PACKET_MAX_SIZE_IN_BYTES]; + final DatagramPacket packet = new DatagramPacket(buffer, buffer.length); + try { + datagramSocket.receive(packet); + + final long airGapReceiverId = receiverId.incrementAndGet(); + final IoTDBAirGapReceiver receiver = + new IoTDBAirGapReceiver(new Socket(), airGapReceiverId); + final String receiverKey = packet.getSocketAddress().toString(); + final boolean registeredSession = registerUdpSessionIfNecessary(packet); + try { + receiver.receiveUdp(datagramSocket, packet, receiverKey, buffer); + } finally { + if (registeredSession) { + SessionManager.getInstance().removeCurrSession(); + } + } + } catch (final IOException e) { + if (allowSubmitListen.get()) { + LOGGER.warn(DataNodePipeMessages.UNHANDLED_EXCEPTION_DURING_PIPE_AIR_GAP_RECEIVER, e); + } + } catch (final Exception e) { + LOGGER.warn(DataNodePipeMessages.UNHANDLED_EXCEPTION_DURING_PIPE_AIR_GAP_RECEIVER, e); + } + } + } + + private boolean registerUdpSessionIfNecessary(final DatagramPacket packet) { + final String receiverKey = packet.getSocketAddress().toString(); + final ClientSession session = + udpClientSessions.computeIfAbsent( + receiverKey, + key -> + new ClientSession(new DatagramClientSocket(packet.getAddress(), packet.getPort()))); + return SessionManager.getInstance().registerSession(session); + } + @Override public void start() throws StartupException { try { serverSocket = new ServerSocket(PipeConfig.getInstance().getPipeAirGapReceiverPort()); + datagramSocket = new DatagramSocket(PipeConfig.getInstance().getPipeAirGapReceiverPort()); } catch (final IOException e) { + if (Objects.nonNull(serverSocket)) { + try { + serverSocket.close(); + } catch (final IOException closeException) { + e.addSuppressed(closeException); + } + } throw new StartupException(e); } allowSubmitListen.set(true); listenExecutor.submit(this::listen); + udpListenExecutor.submit(this::listenUdp); LOGGER.info(DataNodePipeMessages.IOTDBAIRGAPRECEIVERAGENT_STARTED, serverSocket); } @Override public void stop() { + allowSubmitListen.set(false); + try { if (Objects.nonNull(serverSocket)) { serverSocket.close(); } + if (Objects.nonNull(datagramSocket)) { + datagramSocket.close(); + } } catch (final IOException e) { LOGGER.warn(DataNodePipeMessages.FAILED_TO_CLOSE_IOTDBAIRGAPRECEIVERAGENT_S_SERVER_SOCKET, e); } - allowSubmitListen.set(false); + udpClientSessions.forEach( + (key, session) -> { + final boolean registeredSession = SessionManager.getInstance().registerSession(session); + try { + PipeDataNodeAgent.receiver().thrift().handleClientExit(key); + } finally { + if (registeredSession) { + SessionManager.getInstance().removeCurrSession(); + } + } + }); + udpClientSessions.clear(); listenExecutor.shutdown(); + udpListenExecutor.shutdown(); LOGGER.info(DataNodePipeMessages.IOTDBAIRGAPRECEIVERAGENT_STOPPED, serverSocket); } @@ -103,4 +184,25 @@ public void stop() { public ServiceType getID() { return ServiceType.AIR_GAP_SERVICE; } + + private static class DatagramClientSocket extends Socket { + + private final InetAddress address; + private final int port; + + private DatagramClientSocket(final InetAddress address, final int port) { + this.address = address; + this.port = port; + } + + @Override + public InetAddress getInetAddress() { + return address; + } + + @Override + public int getPort() { + return port; + } + } } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/receiver/protocol/thrift/IoTDBDataNodeReceiverAgent.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/receiver/protocol/thrift/IoTDBDataNodeReceiverAgent.java index 8a995f7656d93..d761bb711fc90 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/receiver/protocol/thrift/IoTDBDataNodeReceiverAgent.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/receiver/protocol/thrift/IoTDBDataNodeReceiverAgent.java @@ -24,9 +24,13 @@ import org.apache.iotdb.commons.pipe.sink.payload.thrift.request.IoTDBSinkRequestVersion; import org.apache.iotdb.db.pipe.processor.twostage.exchange.receiver.TwoStageAggregateReceiver; +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; + public class IoTDBDataNodeReceiverAgent extends IoTDBReceiverAgent { private final ThreadLocal receiverThreadLocal = new ThreadLocal<>(); + private final Map receiverMap = new ConcurrentHashMap<>(); @Override protected void initConstructors() { @@ -38,16 +42,24 @@ protected void initConstructors() { @Override protected IoTDBReceiver getReceiverWithSpecifiedClient(final String ignore) { - return receiverThreadLocal.get(); + return ignore == null ? receiverThreadLocal.get() : receiverMap.get(ignore); } @Override protected void setReceiverWithSpecifiedClient(final String ignore, final IoTDBReceiver receiver) { - receiverThreadLocal.set(receiver); + if (ignore == null) { + receiverThreadLocal.set(receiver); + } else { + receiverMap.put(ignore, receiver); + } } @Override protected void removeReceiverWithSpecifiedClient(final String ignore) { - receiverThreadLocal.remove(); + if (ignore == null) { + receiverThreadLocal.remove(); + } else { + receiverMap.remove(ignore); + } } } 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..9bf59f0a50895 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 @@ -33,11 +33,13 @@ import org.apache.iotdb.service.rpc.thrift.TPipeTransferResp; import org.apache.tsfile.utils.BytesUtils; +import org.apache.tsfile.utils.ReadWriteIOUtils; import org.junit.Assert; import org.junit.Test; import java.io.ByteArrayInputStream; import java.io.ByteArrayOutputStream; +import java.io.DataOutputStream; import java.io.IOException; import java.io.InputStream; import java.io.OutputStream; @@ -45,6 +47,8 @@ import java.lang.reflect.Method; import java.net.Socket; import java.nio.ByteBuffer; +import java.util.Arrays; +import java.util.zip.CRC32; public class IoTDBAirGapReceiverTest { @@ -121,6 +125,19 @@ public void testTemporaryUnavailableRetryTimeoutReturnsFail() throws Exception { } } + @Test + public void testReceiveAirGapPayloadWithSpecifiedReceiverKey() throws Exception { + final RecordingSocket socket = new RecordingSocket(); + final IoTDBAirGapReceiver receiver = new IoTDBAirGapReceiver(socket, 4L); + final StubIoTDBDataNodeReceiverAgent stubAgent = new StubIoTDBDataNodeReceiverAgent(); + stubAgent.setStubReceiver( + "udp-client", new StubReceiver(new TSStatus(TSStatusCode.SUCCESS_STATUS.getStatusCode()))); + setField(receiver, "agent", stubAgent); + + Assert.assertTrue(receiver.receive(toAirGapData(toRawRequestBytes()), "udp-client")); + Assert.assertArrayEquals(AirGapOneByteResponse.OK, socket.getWrittenBytes()); + } + private static void setField(final Object target, final String fieldName, final Object value) throws Exception { final Field field = IoTDBAirGapReceiver.class.getDeclaredField(fieldName); @@ -147,6 +164,10 @@ private static class StubIoTDBDataNodeReceiverAgent extends IoTDBDataNodeReceive void setStubReceiver(final IoTDBReceiver receiver) { setReceiverWithSpecifiedClient(null, receiver); } + + void setStubReceiver(final String key, final IoTDBReceiver receiver) { + setReceiverWithSpecifiedClient(key, receiver); + } } private static class StubReceiver implements IoTDBReceiver { @@ -172,4 +193,20 @@ public IoTDBSinkRequestVersion getVersion() { return IoTDBSinkRequestVersion.VERSION_1; } } + + private static byte[] toRawRequestBytes() throws IOException { + try (final ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream(); + final DataOutputStream outputStream = new DataOutputStream(byteArrayOutputStream)) { + ReadWriteIOUtils.write(IoTDBSinkRequestVersion.VERSION_1.getVersion(), outputStream); + ReadWriteIOUtils.write((short) 0, outputStream); + return byteArrayOutputStream.toByteArray(); + } + } + + private static byte[] toAirGapData(final byte[] rawRequestBytes) { + final CRC32 crc32 = new CRC32(); + crc32.update(rawRequestBytes, 0, rawRequestBytes.length); + return BytesUtils.concatByteArrayList( + Arrays.asList(BytesUtils.longToBytes(crc32.getValue()), rawRequestBytes)); + } } diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/sink/protocol/airgap/IoTDBDataRegionAirGapSinkTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/sink/protocol/airgap/IoTDBDataRegionAirGapSinkTest.java index d49ef82ce0550..137d8b40dc875 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/sink/protocol/airgap/IoTDBDataRegionAirGapSinkTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/sink/protocol/airgap/IoTDBDataRegionAirGapSinkTest.java @@ -23,6 +23,7 @@ import org.apache.iotdb.commons.pipe.config.constant.PipeSinkConstant; import org.apache.iotdb.commons.pipe.config.plugin.configuraion.PipeTaskRuntimeConfiguration; import org.apache.iotdb.commons.pipe.config.plugin.env.PipeTaskSinkRuntimeEnvironment; +import org.apache.iotdb.commons.pipe.sink.payload.airgap.AirGapOneByteResponse; import org.apache.iotdb.commons.pipe.sink.payload.thrift.request.PipeRequestType; import org.apache.iotdb.db.pipe.event.common.heartbeat.PipeHeartbeatEvent; import org.apache.iotdb.db.pipe.event.common.tablet.PipeRawTabletInsertionEvent; @@ -39,12 +40,18 @@ import org.junit.Assert; import org.junit.Test; +import java.net.DatagramPacket; +import java.net.DatagramSocket; +import java.net.SocketTimeoutException; import java.nio.ByteBuffer; import java.util.ArrayList; import java.util.Arrays; import java.util.HashMap; import java.util.List; import java.util.Map; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicReference; +import java.util.zip.CRC32; public class IoTDBDataRegionAirGapSinkTest { @@ -104,7 +111,79 @@ public void testTransferTsFileBatchOverAirGap() throws Exception { } } + @Test + public void testUdpAirGapSlicing() throws Exception { + try (final DatagramSocket receiverSocket = new DatagramSocket(0); + final UdpTestingIoTDBDataRegionAirGapSink sink = + new UdpTestingIoTDBDataRegionAirGapSink()) { + final Map attributes = buildParameterAttributes(false); + attributes.put(PipeSinkConstant.CONNECTOR_AIR_GAP_TRANSPORT_KEY, "udp"); + attributes.put(PipeSinkConstant.CONNECTOR_AIR_GAP_UDP_PACKET_SIZE_KEY, "1024"); + + final PipeParameters parameters = new PipeParameters(attributes); + sink.validate(new PipeParameterValidator(parameters)); + sink.customize( + parameters, + new PipeTaskRuntimeConfiguration(new PipeTaskSinkRuntimeEnvironment("pipe", 1L, 1))); + + final List receivedRequests = new ArrayList<>(); + final AtomicBoolean senderFinished = new AtomicBoolean(false); + final AtomicReference receiverException = new AtomicReference<>(); + final Thread receiverThread = + new Thread( + () -> { + try { + receiverSocket.setSoTimeout(100); + while (!senderFinished.get()) { + final byte[] buffer = new byte[2048]; + final DatagramPacket packet = new DatagramPacket(buffer, buffer.length); + try { + receiverSocket.receive(packet); + receivedRequests.add(decodeAirGapDatagram(buffer, packet.getLength())); + receiverSocket.send( + new DatagramPacket( + AirGapOneByteResponse.OK, + AirGapOneByteResponse.OK.length, + packet.getAddress(), + packet.getPort())); + } catch (final SocketTimeoutException ignored) { + // Poll until the sender has finished all datagrams. + } + } + } catch (final Throwable ignored) { + receiverException.set(ignored); + } + }); + receiverThread.start(); + + final byte[] request = new byte[2500]; + request[0] = 1; + final byte[] type = org.apache.tsfile.utils.BytesUtils.shortToBytes((short) 11); + request[1] = type[0]; + request[2] = type[1]; + + try { + Assert.assertTrue(sink.sendUdp(receiverSocket.getLocalPort(), request)); + } finally { + senderFinished.set(true); + } + + receiverThread.join(5000); + Assert.assertFalse(receiverThread.isAlive()); + Assert.assertNull(receiverException.get()); + Assert.assertTrue(receivedRequests.size() > 1); + for (final TPipeTransferReq receivedRequest : receivedRequests) { + Assert.assertEquals(PipeRequestType.TRANSFER_SLICE.getType(), receivedRequest.type); + Assert.assertTrue(receivedRequest.body.remaining() < 1024); + } + } + } + private PipeParameters buildParameters(final boolean useTsFileBatch) { + return new PipeParameters(buildParameterAttributes(useTsFileBatch)); + } + + private Map buildParameterAttributes(final boolean useTsFileBatch) { final Map attributes = new HashMap<>(); attributes.put( PipeSinkConstant.CONNECTOR_KEY, @@ -115,7 +194,7 @@ private PipeParameters buildParameters(final boolean useTsFileBatch) { if (useTsFileBatch) { attributes.put(PipeSinkConstant.CONNECTOR_FORMAT_KEY, "tsfile"); } - return new PipeParameters(attributes); + return attributes; } private PipeRawTabletInsertionEvent createPipeRawTabletInsertionEvent( @@ -149,6 +228,22 @@ private static TPipeTransferReq toTPipeTransferReq(final byte[] requestBytes) { return req; } + private static TPipeTransferReq decodeAirGapDatagram(final byte[] datagram, final int length) { + final ByteBuffer buffer = ByteBuffer.wrap(datagram, 0, length); + final int payloadLength = ReadWriteIOUtils.readInt(buffer); + Assert.assertEquals(payloadLength, ReadWriteIOUtils.readInt(buffer)); + Assert.assertEquals(length - 2 * Integer.BYTES, payloadLength); + + final long expectedChecksum = ReadWriteIOUtils.readLong(buffer); + final byte[] payload = new byte[payloadLength - Long.BYTES]; + buffer.get(payload); + + final CRC32 crc32 = new CRC32(); + crc32.update(payload, 0, payload.length); + Assert.assertEquals(expectedChecksum, crc32.getValue()); + return toTPipeTransferReq(payload); + } + private static class RecordingIoTDBDataRegionAirGapSink extends IoTDBDataRegionAirGapSink { private final List sentRequests = new ArrayList<>(); @@ -180,4 +275,14 @@ public synchronized void setSoTimeout(final int timeout) { } } } + + private static class UdpTestingIoTDBDataRegionAirGapSink extends IoTDBDataRegionAirGapSink { + + private boolean sendUdp(final int port, final byte[] request) throws Exception { + try (final AirGapSocket socket = new AirGapSocket("127.0.0.1", port)) { + socket.connectUdp(1000); + return sendBytes(socket, request); + } + } + } } diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/config/constant/PipeSinkConstant.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/config/constant/PipeSinkConstant.java index 058b17e2f4f4c..98fedbbeb3535 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/config/constant/PipeSinkConstant.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/config/constant/PipeSinkConstant.java @@ -152,6 +152,26 @@ public class PipeSinkConstant { "sink.air-gap.handshake-timeout-ms"; public static final int CONNECTOR_AIR_GAP_HANDSHAKE_TIMEOUT_MS_DEFAULT_VALUE = 5000; + public static final String CONNECTOR_AIR_GAP_TRANSPORT_KEY = "connector.air-gap.transport"; + public static final String SINK_AIR_GAP_TRANSPORT_KEY = "sink.air-gap.transport"; + public static final String CONNECTOR_AIR_GAP_TRANSPORT_TCP_VALUE = "tcp"; + public static final String CONNECTOR_AIR_GAP_TRANSPORT_UDP_VALUE = "udp"; + public static final String CONNECTOR_AIR_GAP_TRANSPORT_DEFAULT_VALUE = + CONNECTOR_AIR_GAP_TRANSPORT_TCP_VALUE; + public static final Set CONNECTOR_AIR_GAP_TRANSPORT_SET = + Collections.unmodifiableSet( + new HashSet<>( + Arrays.asList( + CONNECTOR_AIR_GAP_TRANSPORT_TCP_VALUE, CONNECTOR_AIR_GAP_TRANSPORT_UDP_VALUE))); + + public static final String CONNECTOR_AIR_GAP_UDP_PACKET_SIZE_KEY = + "connector.air-gap.udp.packet-size-bytes"; + public static final String SINK_AIR_GAP_UDP_PACKET_SIZE_KEY = + "sink.air-gap.udp.packet-size-bytes"; + public static final int CONNECTOR_AIR_GAP_UDP_PACKET_SIZE_DEFAULT_VALUE = 60 * 1024; + public static final int CONNECTOR_AIR_GAP_UDP_PACKET_SIZE_MIN_VALUE = 1024; + public static final int CONNECTOR_AIR_GAP_UDP_PACKET_SIZE_MAX_VALUE = 65_507; + public static final String CONNECTOR_IOTDB_SYNC_CONNECTOR_VERSION_KEY = "connector.version"; public static final String SINK_IOTDB_SYNC_CONNECTOR_VERSION_KEY = "sink.version"; public static final String CONNECTOR_IOTDB_SYNC_CONNECTOR_VERSION_DEFAULT_VALUE = "1.1"; diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/protocol/IoTDBAirGapSink.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/protocol/IoTDBAirGapSink.java index 6d74102ff1940..6099109b78179 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/protocol/IoTDBAirGapSink.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/protocol/IoTDBAirGapSink.java @@ -25,25 +25,35 @@ import org.apache.iotdb.commons.pipe.config.PipeConfig; import org.apache.iotdb.commons.pipe.sink.payload.airgap.AirGapELanguageConstant; import org.apache.iotdb.commons.pipe.sink.payload.airgap.AirGapOneByteResponse; +import org.apache.iotdb.commons.pipe.sink.payload.thrift.common.PipeTransferSliceReqBuilder; import org.apache.iotdb.pipe.api.annotation.TableModel; import org.apache.iotdb.pipe.api.annotation.TreeModel; import org.apache.iotdb.pipe.api.customizer.configuration.PipeConnectorRuntimeConfiguration; +import org.apache.iotdb.pipe.api.customizer.parameter.PipeParameterValidator; import org.apache.iotdb.pipe.api.customizer.parameter.PipeParameters; import org.apache.iotdb.pipe.api.exception.PipeConnectionException; import org.apache.iotdb.pipe.api.exception.PipeException; import org.apache.iotdb.rpc.TSStatusCode; +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.slf4j.Logger; import org.slf4j.LoggerFactory; import java.io.BufferedOutputStream; +import java.io.DataOutputStream; import java.io.File; import java.io.IOException; import java.io.RandomAccessFile; +import java.net.DatagramPacket; +import java.net.DatagramSocket; +import java.net.InetAddress; import java.net.InetSocketAddress; import java.net.Socket; import java.net.SocketException; +import java.nio.ByteBuffer; import java.util.ArrayList; import java.util.Arrays; import java.util.HashMap; @@ -55,20 +65,36 @@ import static org.apache.iotdb.commons.pipe.config.constant.PipeSinkConstant.CONNECTOR_AIR_GAP_E_LANGUAGE_ENABLE_KEY; import static org.apache.iotdb.commons.pipe.config.constant.PipeSinkConstant.CONNECTOR_AIR_GAP_HANDSHAKE_TIMEOUT_MS_DEFAULT_VALUE; import static org.apache.iotdb.commons.pipe.config.constant.PipeSinkConstant.CONNECTOR_AIR_GAP_HANDSHAKE_TIMEOUT_MS_KEY; +import static org.apache.iotdb.commons.pipe.config.constant.PipeSinkConstant.CONNECTOR_AIR_GAP_TRANSPORT_DEFAULT_VALUE; +import static org.apache.iotdb.commons.pipe.config.constant.PipeSinkConstant.CONNECTOR_AIR_GAP_TRANSPORT_KEY; +import static org.apache.iotdb.commons.pipe.config.constant.PipeSinkConstant.CONNECTOR_AIR_GAP_TRANSPORT_SET; +import static org.apache.iotdb.commons.pipe.config.constant.PipeSinkConstant.CONNECTOR_AIR_GAP_TRANSPORT_UDP_VALUE; +import static org.apache.iotdb.commons.pipe.config.constant.PipeSinkConstant.CONNECTOR_AIR_GAP_UDP_PACKET_SIZE_DEFAULT_VALUE; +import static org.apache.iotdb.commons.pipe.config.constant.PipeSinkConstant.CONNECTOR_AIR_GAP_UDP_PACKET_SIZE_KEY; +import static org.apache.iotdb.commons.pipe.config.constant.PipeSinkConstant.CONNECTOR_AIR_GAP_UDP_PACKET_SIZE_MAX_VALUE; +import static org.apache.iotdb.commons.pipe.config.constant.PipeSinkConstant.CONNECTOR_AIR_GAP_UDP_PACKET_SIZE_MIN_VALUE; import static org.apache.iotdb.commons.pipe.config.constant.PipeSinkConstant.CONNECTOR_LOAD_BALANCE_PRIORITY_STRATEGY; import static org.apache.iotdb.commons.pipe.config.constant.PipeSinkConstant.CONNECTOR_LOAD_BALANCE_RANDOM_STRATEGY; import static org.apache.iotdb.commons.pipe.config.constant.PipeSinkConstant.CONNECTOR_LOAD_BALANCE_ROUND_ROBIN_STRATEGY; import static org.apache.iotdb.commons.pipe.config.constant.PipeSinkConstant.SINK_AIR_GAP_E_LANGUAGE_ENABLE_KEY; import static org.apache.iotdb.commons.pipe.config.constant.PipeSinkConstant.SINK_AIR_GAP_HANDSHAKE_TIMEOUT_MS_KEY; +import static org.apache.iotdb.commons.pipe.config.constant.PipeSinkConstant.SINK_AIR_GAP_TRANSPORT_KEY; +import static org.apache.iotdb.commons.pipe.config.constant.PipeSinkConstant.SINK_AIR_GAP_UDP_PACKET_SIZE_KEY; import static org.apache.iotdb.commons.utils.BasicStructureSerDeUtil.LONG_LEN; @TreeModel @TableModel public abstract class IoTDBAirGapSink extends IoTDBSink { + private static final int UDP_ENVELOPE_SIZE = 2 * Integer.BYTES + Long.BYTES; + private static final int UDP_SLICE_REQUEST_SERIALIZATION_RESERVED_SIZE = 64; + protected static class AirGapSocket extends Socket { private final TEndPoint endPoint; + private DatagramSocket datagramSocket; + private InetAddress datagramAddress; + private boolean isUdp; public AirGapSocket(final String ip, final int port) { this.endPoint = new TEndPoint(ip, port); @@ -78,9 +104,60 @@ public TEndPoint getEndPoint() { return endPoint; } + public void connectUdp(final int timeoutMs) throws IOException { + datagramAddress = InetAddress.getByName(endPoint.getIp()); + datagramSocket = new DatagramSocket(); + datagramSocket.connect(datagramAddress, endPoint.getPort()); + datagramSocket.setSoTimeout(timeoutMs); + isUdp = true; + } + + public InetAddress getDatagramAddress() { + return datagramAddress; + } + + public DatagramSocket getDatagramSocket() { + return datagramSocket; + } + + public boolean isUdp() { + return isUdp; + } + + @Override + public boolean isConnected() { + return isUdp + ? datagramSocket != null && datagramSocket.isConnected() && !datagramSocket.isClosed() + : super.isConnected(); + } + + @Override + public synchronized void setSoTimeout(final int timeout) throws SocketException { + if (isUdp && datagramSocket != null) { + datagramSocket.setSoTimeout(timeout); + } else { + super.setSoTimeout(timeout); + } + } + + @Override + public synchronized void close() throws IOException { + if (datagramSocket != null) { + datagramSocket.close(); + } + super.close(); + } + @Override public String toString() { - return "AirGapSocket{" + "endPoint=" + endPoint + "} (" + super.toString() + ")"; + return "AirGapSocket{" + + "endPoint=" + + endPoint + + ", transport=" + + (isUdp ? "udp" : "tcp") + + "} (" + + super.toString() + + ")"; } } @@ -98,12 +175,49 @@ public String toString() { private int handshakeTimeoutMs; private boolean eLanguageEnable; + private boolean useUdpTransport; + private int udpPacketSizeInBytes; // The air gap connector does not use clientManager thus we put handshake type here protected boolean supportModsIfIsDataNodeReceiver = true; private final Map failLogTimes = new HashMap<>(); + @Override + public void validate(final PipeParameterValidator validator) throws Exception { + super.validate(validator); + + final PipeParameters parameters = validator.getParameters(); + final String airGapTransport = + parameters + .getStringOrDefault( + Arrays.asList(CONNECTOR_AIR_GAP_TRANSPORT_KEY, SINK_AIR_GAP_TRANSPORT_KEY), + CONNECTOR_AIR_GAP_TRANSPORT_DEFAULT_VALUE) + .trim() + .toLowerCase(); + validator.validate( + arg -> CONNECTOR_AIR_GAP_TRANSPORT_SET.contains(airGapTransport), + String.format( + "Air gap transport should be one of %s, but got %s.", + CONNECTOR_AIR_GAP_TRANSPORT_SET, airGapTransport), + airGapTransport); + + final int packetSize = + parameters.getIntOrDefault( + Arrays.asList(CONNECTOR_AIR_GAP_UDP_PACKET_SIZE_KEY, SINK_AIR_GAP_UDP_PACKET_SIZE_KEY), + CONNECTOR_AIR_GAP_UDP_PACKET_SIZE_DEFAULT_VALUE); + validator.validate( + arg -> + packetSize >= CONNECTOR_AIR_GAP_UDP_PACKET_SIZE_MIN_VALUE + && packetSize <= CONNECTOR_AIR_GAP_UDP_PACKET_SIZE_MAX_VALUE, + String.format( + "UDP packet size should be in the range [%d, %d], but got %d.", + CONNECTOR_AIR_GAP_UDP_PACKET_SIZE_MIN_VALUE, + CONNECTOR_AIR_GAP_UDP_PACKET_SIZE_MAX_VALUE, + packetSize), + packetSize); + } + @Override public void customize( final PipeParameters parameters, final PipeConnectorRuntimeConfiguration configuration) @@ -145,6 +259,23 @@ public void customize( CONNECTOR_AIR_GAP_E_LANGUAGE_ENABLE_KEY, SINK_AIR_GAP_E_LANGUAGE_ENABLE_KEY), CONNECTOR_AIR_GAP_E_LANGUAGE_ENABLE_DEFAULT_VALUE); LOGGER.info(PipeMessages.AIR_GAP_CUSTOMIZED_E_LANGUAGE, eLanguageEnable); + + useUdpTransport = + CONNECTOR_AIR_GAP_TRANSPORT_UDP_VALUE.equals( + parameters + .getStringOrDefault( + Arrays.asList(CONNECTOR_AIR_GAP_TRANSPORT_KEY, SINK_AIR_GAP_TRANSPORT_KEY), + CONNECTOR_AIR_GAP_TRANSPORT_DEFAULT_VALUE) + .trim() + .toLowerCase()); + udpPacketSizeInBytes = + parameters.getIntOrDefault( + Arrays.asList(CONNECTOR_AIR_GAP_UDP_PACKET_SIZE_KEY, SINK_AIR_GAP_UDP_PACKET_SIZE_KEY), + CONNECTOR_AIR_GAP_UDP_PACKET_SIZE_DEFAULT_VALUE); + LOGGER.info( + "Air gap transport is {}, udp packet size is {} bytes.", + useUdpTransport ? "udp" : "tcp", + udpPacketSizeInBytes); } @Override @@ -179,8 +310,12 @@ public void handshake() throws Exception { final AirGapSocket socket = new AirGapSocket(ip, port); try { - socket.connect(new InetSocketAddress(ip, port), handshakeTimeoutMs); - socket.setKeepAlive(true); + if (useUdpTransport) { + socket.connectUdp(handshakeTimeoutMs); + } else { + socket.connect(new InetSocketAddress(ip, port), handshakeTimeoutMs); + socket.setKeepAlive(true); + } sockets.set(i, socket); LOGGER.info(PipeMessages.CONNECTED_TO_TARGET_SERVER, ip, port); failLogTimes.remove(nodeUrls.get(i)); @@ -318,6 +453,10 @@ protected boolean sendBytes(final AirGapSocket socket, byte[] bytes) throws IOEx String.format("Socket %s is closed, will try to handshake", socket)); } + if (socket.isUdp()) { + return sendBytesByUdp(socket, bytes); + } + final BufferedOutputStream outputStream = new BufferedOutputStream(socket.getOutputStream()); bytes = enrichWithLengthAndChecksum(bytes); outputStream.write(eLanguageEnable ? enrichWithELanguage(bytes) : bytes); @@ -328,6 +467,101 @@ protected boolean sendBytes(final AirGapSocket socket, byte[] bytes) throws IOEx return size > 0 && Arrays.equals(AirGapOneByteResponse.OK, response); } + private boolean sendBytesByUdp(final AirGapSocket socket, final byte[] bytes) throws IOException { + for (final byte[] requestBytes : sliceIfNeededForUdp(bytes)) { + if (!sendOneDatagram(socket, requestBytes)) { + return false; + } + } + return true; + } + + private boolean sendOneDatagram(final AirGapSocket socket, final byte[] requestBytes) + throws IOException { + final byte[] datagramBytes = + eLanguageEnable + ? enrichWithELanguage(enrichWithLengthAndChecksum(requestBytes)) + : enrichWithLengthAndChecksum(requestBytes); + if (datagramBytes.length > udpPacketSizeInBytes) { + throw new IOException( + String.format( + "Air gap UDP datagram size %d exceeds configured packet size %d.", + datagramBytes.length, udpPacketSizeInBytes)); + } + + final DatagramSocket datagramSocket = socket.getDatagramSocket(); + datagramSocket.send( + new DatagramPacket( + datagramBytes, + datagramBytes.length, + socket.getDatagramAddress(), + socket.getEndPoint().getPort())); + + final byte[] response = new byte[1]; + final DatagramPacket responsePacket = new DatagramPacket(response, response.length); + datagramSocket.receive(responsePacket); + return responsePacket.getLength() > 0 && Arrays.equals(AirGapOneByteResponse.OK, response); + } + + private List sliceIfNeededForUdp(final byte[] requestBytes) throws IOException { + final int rawPayloadSizeLimit = + udpPacketSizeInBytes + - UDP_ENVELOPE_SIZE + - (eLanguageEnable + ? AirGapELanguageConstant.E_LANGUAGE_PREFIX.length + + AirGapELanguageConstant.E_LANGUAGE_SUFFIX.length + : 0); + if (requestBytes.length <= rawPayloadSizeLimit) { + return Arrays.asList(requestBytes); + } + + final int sliceBodySizeLimit = + rawPayloadSizeLimit - UDP_SLICE_REQUEST_SERIALIZATION_RESERVED_SIZE; + if (sliceBodySizeLimit <= 0) { + throw new IOException( + String.format( + "Air gap UDP packet size %d is too small to transfer sliced requests.", + udpPacketSizeInBytes)); + } + + final TPipeTransferReq request = toTPipeTransferReq(requestBytes); + final int sliceOrderId = PipeTransferSliceReqBuilder.nextSliceOrderId(); + final int sliceCount = PipeTransferSliceReqBuilder.getSliceCount(request, sliceBodySizeLimit); + + final List slicedRequestBytes = new ArrayList<>(sliceCount); + for (int i = 0; i < sliceCount; i++) { + slicedRequestBytes.add( + toTPipeTransferBytes( + PipeTransferSliceReqBuilder.buildSliceReq( + request, sliceOrderId, i, sliceCount, sliceBodySizeLimit))); + } + return slicedRequestBytes; + } + + private TPipeTransferReq toTPipeTransferReq(final byte[] requestBytes) { + final ByteBuffer byteBuffer = ByteBuffer.wrap(requestBytes); + final TPipeTransferReq request = new TPipeTransferReq(); + request.version = ReadWriteIOUtils.readByte(byteBuffer); + request.type = ReadWriteIOUtils.readShort(byteBuffer); + request.body = byteBuffer.slice(); + return request; + } + + private byte[] toTPipeTransferBytes(final TPipeTransferReq request) throws IOException { + try (final PublicBAOS byteArrayOutputStream = new PublicBAOS(); + final DataOutputStream outputStream = new DataOutputStream(byteArrayOutputStream)) { + ReadWriteIOUtils.write(request.version, outputStream); + ReadWriteIOUtils.write(request.type, outputStream); + + final ByteBuffer bodyBuffer = request.body.duplicate(); + final byte[] body = new byte[bodyBuffer.remaining()]; + bodyBuffer.get(body); + outputStream.write(body); + + return Arrays.copyOf(byteArrayOutputStream.getBuf(), byteArrayOutputStream.size()); + } + } + protected boolean send(final AirGapSocket socket, final byte[] bytes) throws IOException { return send(null, 0, socket, bytes); } From 50aa7a86af2d93bbeb469ca75461509a9cbd6d78 Mon Sep 17 00:00:00 2001 From: Caideyipi <87789683+Caideyipi@users.noreply.github.com> Date: Tue, 16 Jun 2026 16:11:51 +0800 Subject: [PATCH 2/7] Update IoTDBDataRegionAirGapSinkTest.java --- .../sink/protocol/airgap/IoTDBDataRegionAirGapSinkTest.java | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/sink/protocol/airgap/IoTDBDataRegionAirGapSinkTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/sink/protocol/airgap/IoTDBDataRegionAirGapSinkTest.java index 137d8b40dc875..ed2568db08de6 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/sink/protocol/airgap/IoTDBDataRegionAirGapSinkTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/sink/protocol/airgap/IoTDBDataRegionAirGapSinkTest.java @@ -42,6 +42,7 @@ import java.net.DatagramPacket; import java.net.DatagramSocket; +import java.net.Socket; import java.net.SocketTimeoutException; import java.nio.ByteBuffer; import java.util.ArrayList; @@ -279,7 +280,8 @@ public synchronized void setSoTimeout(final int timeout) { private static class UdpTestingIoTDBDataRegionAirGapSink extends IoTDBDataRegionAirGapSink { private boolean sendUdp(final int port, final byte[] request) throws Exception { - try (final AirGapSocket socket = new AirGapSocket("127.0.0.1", port)) { + final AirGapSocket socket = new AirGapSocket("127.0.0.1", port); + try (final Socket ignored = socket) { socket.connectUdp(1000); return sendBytes(socket, request); } From 692e20607d79eaf180860d6931f0b423fce088a4 Mon Sep 17 00:00:00 2001 From: Caideyipi <87789683+Caideyipi@users.noreply.github.com> Date: Wed, 15 Jul 2026 15:31:14 +0800 Subject: [PATCH 3/7] Reuse thread pool --- .../protocol/airgap/IoTDBAirGapReceiverAgent.java | 11 ++++------- 1 file changed, 4 insertions(+), 7 deletions(-) diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/receiver/protocol/airgap/IoTDBAirGapReceiverAgent.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/receiver/protocol/airgap/IoTDBAirGapReceiverAgent.java index 0e8c642ce4716..35898dc2e0fe0 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/receiver/protocol/airgap/IoTDBAirGapReceiverAgent.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/receiver/protocol/airgap/IoTDBAirGapReceiverAgent.java @@ -50,13 +50,11 @@ public class IoTDBAirGapReceiverAgent implements IService { private static final Logger LOGGER = LoggerFactory.getLogger(IoTDBAirGapReceiverAgent.class); private static final int UDP_PACKET_MAX_SIZE_IN_BYTES = 65_507; + private static final int LISTENER_THREAD_COUNT = 2; private final ExecutorService listenExecutor = - IoTDBThreadPoolFactory.newSingleThreadExecutor( - ThreadName.PIPE_RECEIVER_AIR_GAP_AGENT.getName()); - private final ExecutorService udpListenExecutor = - IoTDBThreadPoolFactory.newSingleThreadExecutor( - ThreadName.PIPE_RECEIVER_AIR_GAP_AGENT.getName() + "-UDP"); + IoTDBThreadPoolFactory.newFixedThreadPool( + LISTENER_THREAD_COUNT, ThreadName.PIPE_RECEIVER_AIR_GAP_AGENT.getName()); private final AtomicBoolean allowSubmitListen = new AtomicBoolean(false); private ServerSocket serverSocket; @@ -142,7 +140,7 @@ public void start() throws StartupException { allowSubmitListen.set(true); listenExecutor.submit(this::listen); - udpListenExecutor.submit(this::listenUdp); + listenExecutor.submit(this::listenUdp); LOGGER.info(DataNodePipeMessages.IOTDBAIRGAPRECEIVERAGENT_STARTED, serverSocket); } @@ -175,7 +173,6 @@ public void stop() { }); udpClientSessions.clear(); listenExecutor.shutdown(); - udpListenExecutor.shutdown(); LOGGER.info(DataNodePipeMessages.IOTDBAIRGAPRECEIVERAGENT_STOPPED, serverSocket); } From 13cbf37dabd8955559cfcf03f21eff048624ae94 Mon Sep 17 00:00:00 2001 From: Caideyipi <87789683+Caideyipi@users.noreply.github.com> Date: Tue, 18 Aug 2026 09:45:05 +0800 Subject: [PATCH 4/7] Fix air-gap receiver Sonar issues --- .../receiver/protocol/airgap/IoTDBAirGapReceiver.java | 10 +++------- 1 file changed, 3 insertions(+), 7 deletions(-) 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 58928ccc0ed72..677112ce4709c 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 @@ -168,8 +168,9 @@ void receiveUdp( currentDatagramSocket = datagramSocket; currentRemoteSocketAddress = packet.getSocketAddress(); boolean requestHandlingStarted = false; - try { - final byte[] data = readData(new ByteArrayInputStream(buffer, 0, packet.getLength())); + try (final ByteArrayInputStream inputStream = + new ByteArrayInputStream(buffer, 0, packet.getLength())) { + final byte[] data = readData(inputStream); requestHandlingStarted = true; receive(data, receiverKey); } catch (final Exception e) { @@ -196,11 +197,6 @@ private AirGapPseudoTPipeTransferRequest toTPipeTransferReq(final byte[] data) { .setBody(byteBuffer.slice()); } - private void handleReq(final AirGapPseudoTPipeTransferRequest req, final long startTime) - throws IOException { - handleReq(req, startTime, null); - } - private void handleReq( final AirGapPseudoTPipeTransferRequest req, final long startTime, final String receiverKey) throws IOException { From a03062d41b39da7c8d042c873650720a6d9a2532 Mon Sep 17 00:00:00 2001 From: Caideyipi <87789683+Caideyipi@users.noreply.github.com> Date: Tue, 18 Aug 2026 15:18:36 +0800 Subject: [PATCH 5/7] Fix air gap receiver retry test --- .../receiver/protocol/airgap/IoTDBAirGapReceiverTest.java | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) 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 9bf59f0a50895..d5f89900dc520 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 @@ -114,9 +114,9 @@ public void testTemporaryUnavailableRetryTimeoutReturnsFail() throws Exception { final Method handleReq = IoTDBAirGapReceiver.class.getDeclaredMethod( - "handleReq", AirGapPseudoTPipeTransferRequest.class, long.class); + "handleReq", AirGapPseudoTPipeTransferRequest.class, long.class, String.class); handleReq.setAccessible(true); - handleReq.invoke(receiver, req, System.currentTimeMillis() - 10_000L); + handleReq.invoke(receiver, req, System.currentTimeMillis() - 10_000L, null); Assert.assertArrayEquals(AirGapOneByteResponse.FAIL, socket.getWrittenBytes()); } finally { From d94128240ed8ad36b00f2f87e5743a8a5841413d Mon Sep 17 00:00:00 2001 From: Caideyipi <87789683+Caideyipi@users.noreply.github.com> Date: Mon, 19 Oct 2026 10:53:12 +0800 Subject: [PATCH 6/7] Fix air-gap UDP ACK correlation and session cleanup --- .../iotdb/db/i18n/DataNodePipeMessages.java | 4 + .../iotdb/db/i18n/DataNodePipeMessages.java | 4 + .../protocol/airgap/IoTDBAirGapReceiver.java | 46 ++++++- .../airgap/IoTDBAirGapReceiverAgent.java | 116 +++++++++++++---- .../airgap/IoTDBAirGapReceiverAgentTest.java | 120 ++++++++++++++++++ .../airgap/IoTDBAirGapReceiverTest.java | 49 +++++++ .../airgap/IoTDBDataRegionAirGapSinkTest.java | 76 ++++++++++- .../iotdb/commons/i18n/PipeMessages.java | 8 ++ .../iotdb/commons/i18n/PipeMessages.java | 8 ++ .../pipe/sink/protocol/IoTDBAirGapSink.java | 59 ++++++++- 10 files changed, 449 insertions(+), 41 deletions(-) create mode 100644 iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/receiver/protocol/airgap/IoTDBAirGapReceiverAgentTest.java diff --git a/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodePipeMessages.java b/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodePipeMessages.java index 9d400f9d40b1c..9f8d9d1c17907 100644 --- a/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodePipeMessages.java +++ b/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodePipeMessages.java @@ -1367,6 +1367,10 @@ public final class DataNodePipeMessages { "Pipe air gap receiver {} started. Socket: {}"; public static final String PIPE_AIR_GAP_RECEIVER_TEMPORARY_UNAVAILABLE_RETRY = "Pipe air gap receiver {}: Temporary unavailable retry timed out, returning FAIL to sender."; + public static final String EXCEPTION_AIR_GAP_UDP_PACKET_IS_MISSING_ITS_REQUEST_ID_C4228CC8 = + "Air-gap UDP packet is missing its request ID."; + public static final String EXCEPTION_AIR_GAP_UDP_PACKET_CONTAINS_TRAILING_BYTES_2D34D418 = + "Air-gap UDP packet contains trailing bytes."; public static final String PIPE_DATA_TRANSPORT_ERROR = "Pipe data transport error, {}"; public static final String PIPE_INSERTING_ROW_CASTING_TYPE_FROM = "Pipe: Inserting row. Casting type from {} to {}."; diff --git a/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodePipeMessages.java b/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodePipeMessages.java index 4891920e079f0..cb783015b8a3b 100644 --- a/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodePipeMessages.java +++ b/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodePipeMessages.java @@ -1282,6 +1282,10 @@ public final class DataNodePipeMessages { "Pipe air gap receiver {}：监听数据时 socket {} 关闭。原因：{}"; public static final String PIPE_AIR_GAP_RECEIVER_STARTED_SOCKET = "Pipe air gap receiver {} 已启动。Socket：{}"; + public static final String EXCEPTION_AIR_GAP_UDP_PACKET_IS_MISSING_ITS_REQUEST_ID_C4228CC8 = + "Air-gap UDP 数据包缺少 request ID。"; + public static final String EXCEPTION_AIR_GAP_UDP_PACKET_CONTAINS_TRAILING_BYTES_2D34D418 = + "Air-gap UDP 数据包包含多余字节。"; public static final String PIPE_AIR_GAP_RECEIVER_TEMPORARY_UNAVAILABLE_RETRY = "Pipe air gap receiver {}：临时不可用重试超时，向 sender 返回 FAIL。"; public static final String PIPE_DATA_TRANSPORT_ERROR = "Pipe 数据传输错误，{}"; 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 36ac850dbd1ff..e3413b009b2ad 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 @@ -70,6 +70,7 @@ public class IoTDBAirGapReceiver extends WrappedRunnable { private OutputStream currentOutputStream; private DatagramSocket currentDatagramSocket; private SocketAddress currentRemoteSocketAddress; + private Long currentDatagramRequestId; public IoTDBAirGapReceiver(final Socket socket, final long receiverId) { this.socket = socket; @@ -147,6 +148,12 @@ private void receive() throws IOException { } boolean receive(final byte[] data, final String receiverKey) throws IOException { + return receive(data, receiverKey, null); + } + + private boolean receive( + final byte[] data, final String receiverKey, final Runnable onValidRequest) + throws IOException { try { // If check sum failed, it indicates that the length we read may not be correct. // Namely, there may be remaining bytes in the socket stream, which will fail any subsequent @@ -158,6 +165,9 @@ boolean receive(final byte[] data, final String receiverKey) throws IOException return false; } + if (onValidRequest != null) { + onValidRequest.run(); + } handleReq(toTPipeTransferReq(data), System.currentTimeMillis(), receiverKey); return true; } catch (final IOException e) { @@ -171,6 +181,7 @@ boolean receive(final byte[] data, final String receiverKey) throws IOException currentOutputStream = null; currentDatagramSocket = null; currentRemoteSocketAddress = null; + currentDatagramRequestId = null; } } @@ -178,17 +189,29 @@ void receiveUdp( final DatagramSocket datagramSocket, final DatagramPacket packet, final String receiverKey, - final byte[] buffer) + final byte[] buffer, + final Runnable onValidRequest) throws IOException { isELanguagePayload = false; currentDatagramSocket = datagramSocket; currentRemoteSocketAddress = packet.getSocketAddress(); boolean requestHandlingStarted = false; - try (final ByteArrayInputStream inputStream = - new ByteArrayInputStream(buffer, 0, packet.getLength())) { - final byte[] data = readData(inputStream); - requestHandlingStarted = true; - receive(data, receiverKey); + try { + if (packet.getLength() < Long.BYTES) { + throw new IOException( + DataNodePipeMessages.EXCEPTION_AIR_GAP_UDP_PACKET_IS_MISSING_ITS_REQUEST_ID_C4228CC8); + } + currentDatagramRequestId = ByteBuffer.wrap(buffer, 0, Long.BYTES).getLong(); + try (final ByteArrayInputStream inputStream = + new ByteArrayInputStream(buffer, Long.BYTES, packet.getLength() - Long.BYTES)) { + final byte[] data = readData(inputStream); + if (inputStream.available() != 0) { + throw new IOException( + DataNodePipeMessages.EXCEPTION_AIR_GAP_UDP_PACKET_CONTAINS_TRAILING_BYTES_2D34D418); + } + requestHandlingStarted = true; + receive(data, receiverKey, onValidRequest); + } } catch (final Exception e) { if (!requestHandlingStarted) { fail(); @@ -198,6 +221,7 @@ void receiveUdp( currentOutputStream = null; currentDatagramSocket = null; currentRemoteSocketAddress = null; + currentDatagramRequestId = null; } } @@ -272,8 +296,16 @@ private void fail() throws IOException { private void respond(final byte[] response) throws IOException { if (currentDatagramSocket != null && currentRemoteSocketAddress != null) { + if (currentDatagramRequestId == null) { + return; + } + final byte[] responseBytes = + ByteBuffer.allocate(Long.BYTES + response.length) + .putLong(currentDatagramRequestId) + .put(response) + .array(); currentDatagramSocket.send( - new DatagramPacket(response, response.length, currentRemoteSocketAddress)); + new DatagramPacket(responseBytes, responseBytes.length, currentRemoteSocketAddress)); return; } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/receiver/protocol/airgap/IoTDBAirGapReceiverAgent.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/receiver/protocol/airgap/IoTDBAirGapReceiverAgent.java index 35898dc2e0fe0..87683081d9be5 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/receiver/protocol/airgap/IoTDBAirGapReceiverAgent.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/receiver/protocol/airgap/IoTDBAirGapReceiverAgent.java @@ -39,10 +39,12 @@ import java.net.InetAddress; import java.net.ServerSocket; import java.net.Socket; +import java.net.SocketTimeoutException; import java.util.Map; import java.util.Objects; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ExecutorService; +import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicLong; @@ -50,6 +52,8 @@ public class IoTDBAirGapReceiverAgent implements IService { private static final Logger LOGGER = LoggerFactory.getLogger(IoTDBAirGapReceiverAgent.class); private static final int UDP_PACKET_MAX_SIZE_IN_BYTES = 65_507; + private static final int UDP_SESSION_CLEANUP_INTERVAL_MS = 60_000; + private static final int UDP_MAX_CLIENT_SESSIONS = 1_024; private static final int LISTENER_THREAD_COUNT = 2; private final ExecutorService listenExecutor = @@ -61,7 +65,7 @@ public class IoTDBAirGapReceiverAgent implements IService { private DatagramSocket datagramSocket; private final AtomicLong receiverId = new AtomicLong(0); - private final Map udpClientSessions = new ConcurrentHashMap<>(); + final Map udpClientSessions = new ConcurrentHashMap<>(); public void listen() { try { @@ -85,6 +89,7 @@ public void listen() { public void listenUdp() { while (allowSubmitListen.get() && !datagramSocket.isClosed()) { + removeExpiredUdpClientSessions(System.nanoTime()); final byte[] buffer = new byte[UDP_PACKET_MAX_SIZE_IN_BYTES]; final DatagramPacket packet = new DatagramPacket(buffer, buffer.length); try { @@ -94,14 +99,21 @@ public void listenUdp() { final IoTDBAirGapReceiver receiver = new IoTDBAirGapReceiver(new Socket(), airGapReceiverId); final String receiverKey = packet.getSocketAddress().toString(); - final boolean registeredSession = registerUdpSessionIfNecessary(packet); + final boolean[] registeredSession = {false}; try { - receiver.receiveUdp(datagramSocket, packet, receiverKey, buffer); + receiver.receiveUdp( + datagramSocket, + packet, + receiverKey, + buffer, + () -> registeredSession[0] = registerUdpSessionIfNecessary(packet)); } finally { - if (registeredSession) { + if (registeredSession[0]) { SessionManager.getInstance().removeCurrSession(); } } + } catch (final SocketTimeoutException e) { + removeExpiredUdpClientSessions(System.nanoTime()); } catch (final IOException e) { if (allowSubmitListen.get()) { LOGGER.warn(DataNodePipeMessages.UNHANDLED_EXCEPTION_DURING_PIPE_AIR_GAP_RECEIVER, e); @@ -112,14 +124,71 @@ public void listenUdp() { } } - private boolean registerUdpSessionIfNecessary(final DatagramPacket packet) { + boolean registerUdpSessionIfNecessary(final DatagramPacket packet) { final String receiverKey = packet.getSocketAddress().toString(); - final ClientSession session = - udpClientSessions.computeIfAbsent( - receiverKey, - key -> - new ClientSession(new DatagramClientSocket(packet.getAddress(), packet.getPort()))); - return SessionManager.getInstance().registerSession(session); + final long now = System.nanoTime(); + removeExpiredUdpClientSessions(now); + + UdpClientSession clientSession = udpClientSessions.get(receiverKey); + if (clientSession == null) { + ensureUdpClientSessionCapacity(); + clientSession = + new UdpClientSession( + new ClientSession(new DatagramClientSocket(packet.getAddress(), packet.getPort())), + now); + final UdpClientSession existing = udpClientSessions.putIfAbsent(receiverKey, clientSession); + if (existing != null) { + clientSession = existing; + } + } + clientSession.lastAccessTimeNanos = now; + return SessionManager.getInstance().registerSession(clientSession.session); + } + + private void ensureUdpClientSessionCapacity() { + if (udpClientSessions.size() < UDP_MAX_CLIENT_SESSIONS) { + return; + } + + String leastRecentlyUsedKey = null; + UdpClientSession leastRecentlyUsedSession = null; + for (final Map.Entry entry : udpClientSessions.entrySet()) { + if (leastRecentlyUsedSession == null + || entry.getValue().lastAccessTimeNanos < leastRecentlyUsedSession.lastAccessTimeNanos) { + leastRecentlyUsedKey = entry.getKey(); + leastRecentlyUsedSession = entry.getValue(); + } + } + if (leastRecentlyUsedKey != null) { + removeUdpClientSession(leastRecentlyUsedKey, leastRecentlyUsedSession); + } + } + + void removeExpiredUdpClientSessions(final long nowNanos) { + final long idleTimeoutNanos = + TimeUnit.MILLISECONDS.toNanos( + Math.max(1, PipeConfig.getInstance().getPipeSinkTransferTimeoutMs())); + udpClientSessions.forEach( + (key, session) -> { + if (nowNanos - session.lastAccessTimeNanos >= idleTimeoutNanos) { + removeUdpClientSession(key, session); + } + }); + } + + private void removeUdpClientSession(final String key, final UdpClientSession session) { + if (!udpClientSessions.remove(key, session)) { + return; + } + + final boolean registeredSession = SessionManager.getInstance().registerSession(session.session); + try { + PipeDataNodeAgent.receiver().thrift().handleClientExit(key); + } finally { + if (registeredSession) { + SessionManager.getInstance().removeCurrSession(); + } + } } @Override @@ -127,6 +196,7 @@ public void start() throws StartupException { try { serverSocket = new ServerSocket(PipeConfig.getInstance().getPipeAirGapReceiverPort()); datagramSocket = new DatagramSocket(PipeConfig.getInstance().getPipeAirGapReceiverPort()); + datagramSocket.setSoTimeout(UDP_SESSION_CLEANUP_INTERVAL_MS); } catch (final IOException e) { if (Objects.nonNull(serverSocket)) { try { @@ -160,18 +230,7 @@ public void stop() { LOGGER.warn(DataNodePipeMessages.FAILED_TO_CLOSE_IOTDBAIRGAPRECEIVERAGENT_S_SERVER_SOCKET, e); } - udpClientSessions.forEach( - (key, session) -> { - final boolean registeredSession = SessionManager.getInstance().registerSession(session); - try { - PipeDataNodeAgent.receiver().thrift().handleClientExit(key); - } finally { - if (registeredSession) { - SessionManager.getInstance().removeCurrSession(); - } - } - }); - udpClientSessions.clear(); + udpClientSessions.forEach(this::removeUdpClientSession); listenExecutor.shutdown(); LOGGER.info(DataNodePipeMessages.IOTDBAIRGAPRECEIVERAGENT_STOPPED, serverSocket); @@ -202,4 +261,15 @@ public int getPort() { return port; } } + + static class UdpClientSession { + + private final ClientSession session; + volatile long lastAccessTimeNanos; + + UdpClientSession(final ClientSession session, final long lastAccessTimeNanos) { + this.session = session; + this.lastAccessTimeNanos = lastAccessTimeNanos; + } + } } diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/receiver/protocol/airgap/IoTDBAirGapReceiverAgentTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/receiver/protocol/airgap/IoTDBAirGapReceiverAgentTest.java new file mode 100644 index 0000000000000..ffbebb5a66851 --- /dev/null +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/receiver/protocol/airgap/IoTDBAirGapReceiverAgentTest.java @@ -0,0 +1,120 @@ +/* + * 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.db.pipe.receiver.protocol.airgap; + +import org.apache.iotdb.common.rpc.thrift.TSStatus; +import org.apache.iotdb.commons.pipe.config.PipeConfig; +import org.apache.iotdb.commons.pipe.receiver.IoTDBReceiver; +import org.apache.iotdb.commons.pipe.sink.payload.thrift.request.IoTDBSinkRequestVersion; +import org.apache.iotdb.db.pipe.agent.PipeDataNodeAgent; +import org.apache.iotdb.db.pipe.receiver.protocol.thrift.IoTDBDataNodeReceiverAgent; +import org.apache.iotdb.db.protocol.session.ClientSession; +import org.apache.iotdb.rpc.TSStatusCode; +import org.apache.iotdb.service.rpc.thrift.TPipeTransferReq; +import org.apache.iotdb.service.rpc.thrift.TPipeTransferResp; + +import org.junit.Assert; +import org.junit.Test; + +import java.io.IOException; +import java.lang.reflect.Field; +import java.net.DatagramPacket; +import java.net.DatagramSocket; +import java.net.InetAddress; +import java.net.Socket; +import java.util.Map; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; + +public class IoTDBAirGapReceiverAgentTest { + + @Test + public void testExpiredUdpClientSessionIsRemoved() throws Exception { + final IoTDBAirGapReceiverAgent agent = new IoTDBAirGapReceiverAgent(); + final String sessionKey = "expired-client"; + final IoTDBAirGapReceiverAgent.UdpClientSession session = + new IoTDBAirGapReceiverAgent.UdpClientSession(new ClientSession(new Socket()), 0L); + agent.udpClientSessions.put(sessionKey, session); + + final AtomicBoolean receiverExited = new AtomicBoolean(false); + final IoTDBDataNodeReceiverAgent receiverAgent = PipeDataNodeAgent.receiver().thrift(); + final Field receiverMapField = IoTDBDataNodeReceiverAgent.class.getDeclaredField("receiverMap"); + receiverMapField.setAccessible(true); + @SuppressWarnings("unchecked") + final Map receiverMap = + (Map) receiverMapField.get(receiverAgent); + receiverMap.put( + sessionKey, + new IoTDBReceiver() { + @Override + public TPipeTransferResp receive(final TPipeTransferReq req) { + return new TPipeTransferResp(new TSStatus(TSStatusCode.SUCCESS_STATUS.getStatusCode())); + } + + @Override + public void handleExit() { + receiverExited.set(true); + } + + @Override + public IoTDBSinkRequestVersion getVersion() { + return IoTDBSinkRequestVersion.VERSION_1; + } + }); + + try { + final long timeoutNanos = + TimeUnit.MILLISECONDS.toNanos(PipeConfig.getInstance().getPipeSinkTransferTimeoutMs()); + agent.removeExpiredUdpClientSessions(timeoutNanos + 1); + + Assert.assertTrue(agent.udpClientSessions.isEmpty()); + Assert.assertTrue(receiverExited.get()); + } finally { + receiverMap.remove(sessionKey); + } + } + + @Test + public void testMalformedUdpPacketDoesNotCreateClientSession() throws Exception { + try (final DatagramSocket datagramSocket = new DatagramSocket()) { + final IoTDBAirGapReceiver receiver = new IoTDBAirGapReceiver(new Socket(), 1L); + final byte[] malformedPacket = new byte[] {1}; + final DatagramPacket packet = + new DatagramPacket( + malformedPacket, + malformedPacket.length, + InetAddress.getLoopbackAddress(), + datagramSocket.getLocalPort()); + final AtomicBoolean sessionRegistrationAttempted = new AtomicBoolean(false); + + Assert.assertThrows( + IOException.class, + () -> + receiver.receiveUdp( + datagramSocket, + packet, + "invalid-client", + malformedPacket, + () -> sessionRegistrationAttempted.set(true))); + + Assert.assertFalse(sessionRegistrationAttempted.get()); + } + } +} 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 60c6126943639..57cacb3879906 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 @@ -50,6 +50,8 @@ import java.io.OutputStream; import java.lang.reflect.Field; import java.lang.reflect.Method; +import java.net.DatagramPacket; +import java.net.DatagramSocket; import java.net.InetAddress; import java.net.Socket; import java.net.SocketException; @@ -164,6 +166,42 @@ public void testReceiveAirGapPayloadWithSpecifiedReceiverKey() throws Exception Assert.assertArrayEquals(AirGapOneByteResponse.OK, socket.getWrittenBytes()); } + @Test + public void testUdpResponseContainsRequestId() throws Exception { + try (final DatagramSocket receiverSocket = new DatagramSocket(0); + final DatagramSocket senderSocket = new DatagramSocket()) { + final long requestId = 73L; + final byte[] request = toUdpDatagram(requestId, toRawRequestBytes()); + senderSocket.send( + new DatagramPacket( + request, + request.length, + InetAddress.getLoopbackAddress(), + receiverSocket.getLocalPort())); + + final byte[] buffer = new byte[65_507]; + final DatagramPacket packet = new DatagramPacket(buffer, buffer.length); + receiverSocket.receive(packet); + + final RecordingSocket socket = new RecordingSocket(); + final IoTDBAirGapReceiver receiver = new IoTDBAirGapReceiver(socket, 5L); + final StubIoTDBDataNodeReceiverAgent stubAgent = new StubIoTDBDataNodeReceiverAgent(); + stubAgent.setStubReceiver( + "udp-client", + new StubReceiver(new TSStatus(TSStatusCode.SUCCESS_STATUS.getStatusCode()))); + setField(receiver, "agent", stubAgent); + receiver.receiveUdp(receiverSocket, packet, "udp-client", buffer, () -> {}); + + final byte[] response = new byte[Long.BYTES + 2]; + final DatagramPacket responsePacket = new DatagramPacket(response, response.length); + senderSocket.setSoTimeout(1000); + senderSocket.receive(responsePacket); + Assert.assertEquals(Long.BYTES + 1, responsePacket.getLength()); + Assert.assertEquals(requestId, ByteBuffer.wrap(response).getLong()); + Assert.assertEquals(AirGapOneByteResponse.OK[0], response[Long.BYTES]); + } + } + @Test public void testTemporaryUnavailableRetryUsesFreshRequestBody() throws Exception { assertTemporaryUnavailableRetryUsesFreshRequestBody(null); @@ -462,6 +500,17 @@ private static byte[] toAirGapData(final byte[] rawRequestBytes) { Arrays.asList(BytesUtils.longToBytes(crc32.getValue()), rawRequestBytes)); } + private static byte[] toUdpDatagram(final long requestId, final byte[] rawRequestBytes) { + final byte[] data = toAirGapData(rawRequestBytes); + final ByteBuffer datagram = + ByteBuffer.allocate(Long.BYTES + 2 * Integer.BYTES + data.length) + .putLong(requestId) + .putInt(data.length) + .putInt(data.length) + .put(data); + return datagram.array(); + } + private static class RuntimeCleanupReceiver implements IoTDBReceiver { private final String sessionKey; diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/sink/protocol/airgap/IoTDBDataRegionAirGapSinkTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/sink/protocol/airgap/IoTDBDataRegionAirGapSinkTest.java index e81de28d4c30d..1b858fb3b5b7b 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/sink/protocol/airgap/IoTDBDataRegionAirGapSinkTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/sink/protocol/airgap/IoTDBDataRegionAirGapSinkTest.java @@ -177,10 +177,14 @@ public void testUdpAirGapSlicing() throws Exception { try { receiverSocket.receive(packet); receivedRequests.add(decodeAirGapDatagram(buffer, packet.getLength())); + final ByteBuffer responseBuffer = + ByteBuffer.allocate(Long.BYTES + AirGapOneByteResponse.OK.length) + .putLong(ByteBuffer.wrap(buffer, 0, Long.BYTES).getLong()) + .put(AirGapOneByteResponse.OK); receiverSocket.send( new DatagramPacket( - AirGapOneByteResponse.OK, - AirGapOneByteResponse.OK.length, + responseBuffer.array(), + responseBuffer.array().length, packet.getAddress(), packet.getPort())); } catch (final SocketTimeoutException ignored) { @@ -216,6 +220,58 @@ public void testUdpAirGapSlicing() throws Exception { } } + @Test + public void testUdpAirGapIgnoresStaleAcknowledgement() throws Exception { + try (final DatagramSocket receiverSocket = new DatagramSocket(0); + final UdpTestingIoTDBDataRegionAirGapSink sink = + new UdpTestingIoTDBDataRegionAirGapSink()) { + final PipeParameters parameters = new PipeParameters(buildParameterAttributes(false)); + parameters.getAttribute().put(PipeSinkConstant.CONNECTOR_AIR_GAP_TRANSPORT_KEY, "udp"); + parameters.getAttribute().put(PipeSinkConstant.CONNECTOR_AIR_GAP_UDP_PACKET_SIZE_KEY, "1024"); + sink.validate(new PipeParameterValidator(parameters)); + sink.customize( + parameters, + new PipeTaskRuntimeConfiguration(new PipeTaskSinkRuntimeEnvironment("pipe", 1L, 1))); + + final AtomicBoolean secondPacketReceived = new AtomicBoolean(false); + final AtomicReference receiverException = new AtomicReference<>(); + final Thread receiverThread = + new Thread( + () -> { + try { + final byte[] buffer = new byte[2048]; + final DatagramPacket packet = new DatagramPacket(buffer, buffer.length); + receiverSocket.receive(packet); + final long requestId = ByteBuffer.wrap(buffer, 0, Long.BYTES).getLong(); + sendUdpAck(receiverSocket, packet, requestId); + receiverSocket.receive(packet); + sendUdpAck(receiverSocket, packet, requestId); + secondPacketReceived.set(true); + } catch (final Throwable t) { + receiverException.set(t); + } + }); + receiverThread.start(); + + try { + final byte[] request = new byte[2500]; + request[0] = 1; + request[1] = 0; + request[2] = 1; + Assert.assertThrows( + SocketTimeoutException.class, + () -> sink.sendUdp(receiverSocket.getLocalPort(), request)); + } finally { + receiverSocket.close(); + } + + receiverThread.join(5000); + Assert.assertFalse(receiverThread.isAlive()); + Assert.assertNull(receiverException.get()); + Assert.assertTrue(secondPacketReceived.get()); + } + } + private PipeParameters buildParameters(final boolean useTsFileBatch) { return new PipeParameters(buildParameterAttributes(useTsFileBatch)); } @@ -266,10 +322,10 @@ private static TPipeTransferReq toTPipeTransferReq(final byte[] requestBytes) { } private static TPipeTransferReq decodeAirGapDatagram(final byte[] datagram, final int length) { - final ByteBuffer buffer = ByteBuffer.wrap(datagram, 0, length); + final ByteBuffer buffer = ByteBuffer.wrap(datagram, Long.BYTES, length - Long.BYTES); final int payloadLength = ReadWriteIOUtils.readInt(buffer); Assert.assertEquals(payloadLength, ReadWriteIOUtils.readInt(buffer)); - Assert.assertEquals(length - 2 * Integer.BYTES, payloadLength); + Assert.assertEquals(length - Long.BYTES - 2 * Integer.BYTES, payloadLength); final long expectedChecksum = ReadWriteIOUtils.readLong(buffer); final byte[] payload = new byte[payloadLength - Long.BYTES]; @@ -281,6 +337,18 @@ private static TPipeTransferReq decodeAirGapDatagram(final byte[] datagram, fina return toTPipeTransferReq(payload); } + private static void sendUdpAck( + final DatagramSocket socket, final DatagramPacket requestPacket, final long requestId) + throws IOException { + final byte[] ack = + ByteBuffer.allocate(Long.BYTES + AirGapOneByteResponse.OK.length) + .putLong(requestId) + .put(AirGapOneByteResponse.OK) + .array(); + socket.send( + new DatagramPacket(ack, ack.length, requestPacket.getAddress(), requestPacket.getPort())); + } + private static class RecordingIoTDBDataRegionAirGapSink extends IoTDBDataRegionAirGapSink { private final List sentRequests = new ArrayList<>(); diff --git a/iotdb-core/node-commons/src/main/i18n/en/org/apache/iotdb/commons/i18n/PipeMessages.java b/iotdb-core/node-commons/src/main/i18n/en/org/apache/iotdb/commons/i18n/PipeMessages.java index 0aafc35842770..c09c42a188fc3 100644 --- a/iotdb-core/node-commons/src/main/i18n/en/org/apache/iotdb/commons/i18n/PipeMessages.java +++ b/iotdb-core/node-commons/src/main/i18n/en/org/apache/iotdb/commons/i18n/PipeMessages.java @@ -508,6 +508,14 @@ public final class PipeMessages { "IoTDBAirGapConnector is customized with handshakeTimeoutMs: {}."; public static final String AIR_GAP_CUSTOMIZED_E_LANGUAGE = "IoTDBAirGapConnector is customized with eLanguageEnable: {}."; + public static final String + EXCEPTION_AIR_GAP_UDP_DATAGRAM_SIZE_ARG_EXCEEDS_CONFIGURED_PACKET_SIZE_ARG_559D0654 = + "Air gap UDP datagram size %d exceeds configured packet size %d."; + public static final String EXCEPTION_TIMED_OUT_WAITING_FOR_THE_MATCHING_AIR_GAP_ACK_DD2F3706 = + "Timed out waiting for the matching air-gap ACK."; + public static final String + EXCEPTION_AIR_GAP_UDP_PACKET_SIZE_ARG_IS_TOO_SMALL_TO_TRANSFER_SLICED_REQUESTS_722B32AB = + "Air gap UDP packet size %d is too small to transfer sliced requests."; public static final String CONNECTED_TO_TARGET_SERVER = "Successfully connected to target server ip: {}, port: {}."; public static final String FAILED_TO_CLOSE_SOCKET = diff --git a/iotdb-core/node-commons/src/main/i18n/zh/org/apache/iotdb/commons/i18n/PipeMessages.java b/iotdb-core/node-commons/src/main/i18n/zh/org/apache/iotdb/commons/i18n/PipeMessages.java index 9746b71500d2d..b8fd1db6729a0 100644 --- a/iotdb-core/node-commons/src/main/i18n/zh/org/apache/iotdb/commons/i18n/PipeMessages.java +++ b/iotdb-core/node-commons/src/main/i18n/zh/org/apache/iotdb/commons/i18n/PipeMessages.java @@ -486,6 +486,14 @@ public final class PipeMessages { public static final String AIR_GAP_CUSTOMIZED_HANDSHAKE_TIMEOUT = "IoTDBAirGapConnector 自定义握手超时时间：{}。"; + public static final String + EXCEPTION_AIR_GAP_UDP_DATAGRAM_SIZE_ARG_EXCEEDS_CONFIGURED_PACKET_SIZE_ARG_559D0654 = + "Air gap UDP 数据报大小 %d 超过配置的包大小 %d。"; + public static final String EXCEPTION_TIMED_OUT_WAITING_FOR_THE_MATCHING_AIR_GAP_ACK_DD2F3706 = + "等待匹配的 air-gap ACK 超时。"; + public static final String + EXCEPTION_AIR_GAP_UDP_PACKET_SIZE_ARG_IS_TOO_SMALL_TO_TRANSFER_SLICED_REQUESTS_722B32AB = + "Air gap UDP 包大小 %d 太小，无法传输切片请求。"; public static final String AIR_GAP_CUSTOMIZED_E_LANGUAGE = "IoTDBAirGapConnector 自定义 eLanguageEnable：{}。"; public static final String CONNECTED_TO_TARGET_SERVER = diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/protocol/IoTDBAirGapSink.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/protocol/IoTDBAirGapSink.java index 4858c2cf94f8c..506b834ed77c6 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/protocol/IoTDBAirGapSink.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/protocol/IoTDBAirGapSink.java @@ -56,12 +56,15 @@ import java.net.InetSocketAddress; import java.net.Socket; import java.net.SocketException; +import java.net.SocketTimeoutException; import java.nio.ByteBuffer; import java.util.ArrayList; import java.util.Arrays; import java.util.HashMap; import java.util.List; import java.util.Map; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicLong; import java.util.zip.CRC32; import static org.apache.iotdb.commons.pipe.config.constant.PipeSinkConstant.CONNECTOR_AIR_GAP_E_LANGUAGE_ENABLE_DEFAULT_VALUE; @@ -90,6 +93,7 @@ public abstract class IoTDBAirGapSink extends IoTDBSink { private static final int UDP_ENVELOPE_SIZE = 2 * Integer.BYTES + Long.BYTES; + private static final int UDP_REQUEST_ID_SIZE = Long.BYTES; private static final int UDP_SLICE_REQUEST_SERIALIZATION_RESERVED_SIZE = 64; protected static class AirGapSocket extends Socket { @@ -185,6 +189,7 @@ public String toString() { protected boolean supportModsIfIsDataNodeReceiver = true; private final Map failLogTimes = new HashMap<>(); + private final AtomicLong udpRequestIdGenerator = new AtomicLong(); @Override public void validate(final PipeParameterValidator validator) throws Exception { @@ -484,15 +489,23 @@ private boolean sendBytesByUdp(final AirGapSocket socket, final byte[] bytes) th private boolean sendOneDatagram(final AirGapSocket socket, final byte[] requestBytes) throws IOException { - final byte[] datagramBytes = + final long requestId = udpRequestIdGenerator.incrementAndGet(); + final byte[] framedRequestBytes = eLanguageEnable ? enrichWithELanguage(enrichWithLengthAndChecksum(requestBytes)) : enrichWithLengthAndChecksum(requestBytes); + final byte[] datagramBytes = + ByteBuffer.allocate(UDP_REQUEST_ID_SIZE + framedRequestBytes.length) + .putLong(requestId) + .put(framedRequestBytes) + .array(); if (datagramBytes.length > udpPacketSizeInBytes) { throw new IOException( String.format( - "Air gap UDP datagram size %d exceeds configured packet size %d.", - datagramBytes.length, udpPacketSizeInBytes)); + PipeMessages + .EXCEPTION_AIR_GAP_UDP_DATAGRAM_SIZE_ARG_EXCEEDS_CONFIGURED_PACKET_SIZE_ARG_559D0654, + datagramBytes.length, + udpPacketSizeInBytes)); } final DatagramSocket datagramSocket = socket.getDatagramSocket(); @@ -503,16 +516,47 @@ private boolean sendOneDatagram(final AirGapSocket socket, final byte[] requestB socket.getDatagramAddress(), socket.getEndPoint().getPort())); - final byte[] response = new byte[1]; + final int originalTimeoutMs = datagramSocket.getSoTimeout(); + final long deadlineNanos = + originalTimeoutMs == 0 + ? Long.MAX_VALUE + : System.nanoTime() + TimeUnit.MILLISECONDS.toNanos(originalTimeoutMs); + final byte[] response = new byte[UDP_REQUEST_ID_SIZE + 1]; final DatagramPacket responsePacket = new DatagramPacket(response, response.length); - datagramSocket.receive(responsePacket); - return responsePacket.getLength() > 0 && Arrays.equals(AirGapOneByteResponse.OK, response); + try { + while (true) { + if (deadlineNanos != Long.MAX_VALUE) { + final long remainingNanos = deadlineNanos - System.nanoTime(); + if (remainingNanos <= 0) { + throw new SocketTimeoutException( + PipeMessages.EXCEPTION_TIMED_OUT_WAITING_FOR_THE_MATCHING_AIR_GAP_ACK_DD2F3706); + } + datagramSocket.setSoTimeout( + (int) Math.max(1, TimeUnit.NANOSECONDS.toMillis(remainingNanos))); + } + + responsePacket.setLength(response.length); + datagramSocket.receive(responsePacket); + if (responsePacket.getLength() != UDP_REQUEST_ID_SIZE + 1) { + continue; + } + + final ByteBuffer responseBuffer = ByteBuffer.wrap(response, 0, responsePacket.getLength()); + if (responseBuffer.getLong() != requestId) { + continue; + } + return Arrays.equals(AirGapOneByteResponse.OK, new byte[] {responseBuffer.get()}); + } + } finally { + datagramSocket.setSoTimeout(originalTimeoutMs); + } } private List sliceIfNeededForUdp(final byte[] requestBytes) throws IOException { final int rawPayloadSizeLimit = udpPacketSizeInBytes - UDP_ENVELOPE_SIZE + - UDP_REQUEST_ID_SIZE - (eLanguageEnable ? AirGapELanguageConstant.E_LANGUAGE_PREFIX.length + AirGapELanguageConstant.E_LANGUAGE_SUFFIX.length @@ -526,7 +570,8 @@ private List sliceIfNeededForUdp(final byte[] requestBytes) throws IOExc if (sliceBodySizeLimit <= 0) { throw new IOException( String.format( - "Air gap UDP packet size %d is too small to transfer sliced requests.", + PipeMessages + .EXCEPTION_AIR_GAP_UDP_PACKET_SIZE_ARG_IS_TOO_SMALL_TO_TRANSFER_SLICED_REQUESTS_722B32AB, udpPacketSizeInBytes)); } From 39dcb47d9c002d3e580429a1c99b75a7794d9d7b Mon Sep 17 00:00:00 2001 From: Caideyipi <87789683+Caideyipi@users.noreply.github.com> Date: Thu, 29 Oct 2026 21:16:12 +0800 Subject: [PATCH 7/7] Address air-gap UDP review comments --- .../thrift/IoTDBDataNodeReceiverAgent.java | 16 ++-- .../airgap/IoTDBDataRegionAirGapSinkTest.java | 61 +++++++++++++ .../pipe/sink/protocol/IoTDBAirGapSink.java | 91 ++++++++++--------- 3 files changed, 117 insertions(+), 51 deletions(-) diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/receiver/protocol/thrift/IoTDBDataNodeReceiverAgent.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/receiver/protocol/thrift/IoTDBDataNodeReceiverAgent.java index d761bb711fc90..57ed64ec520d1 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/receiver/protocol/thrift/IoTDBDataNodeReceiverAgent.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/receiver/protocol/thrift/IoTDBDataNodeReceiverAgent.java @@ -41,25 +41,25 @@ protected void initConstructors() { } @Override - protected IoTDBReceiver getReceiverWithSpecifiedClient(final String ignore) { - return ignore == null ? receiverThreadLocal.get() : receiverMap.get(ignore); + protected IoTDBReceiver getReceiverWithSpecifiedClient(final String key) { + return key == null ? receiverThreadLocal.get() : receiverMap.get(key); } @Override - protected void setReceiverWithSpecifiedClient(final String ignore, final IoTDBReceiver receiver) { - if (ignore == null) { + protected void setReceiverWithSpecifiedClient(final String key, final IoTDBReceiver receiver) { + if (key == null) { receiverThreadLocal.set(receiver); } else { - receiverMap.put(ignore, receiver); + receiverMap.put(key, receiver); } } @Override - protected void removeReceiverWithSpecifiedClient(final String ignore) { - if (ignore == null) { + protected void removeReceiverWithSpecifiedClient(final String key) { + if (key == null) { receiverThreadLocal.remove(); } else { - receiverMap.remove(ignore); + receiverMap.remove(key); } } } diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/sink/protocol/airgap/IoTDBDataRegionAirGapSinkTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/sink/protocol/airgap/IoTDBDataRegionAirGapSinkTest.java index 1b858fb3b5b7b..82000564cfef5 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/sink/protocol/airgap/IoTDBDataRegionAirGapSinkTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/sink/protocol/airgap/IoTDBDataRegionAirGapSinkTest.java @@ -164,6 +164,7 @@ public void testUdpAirGapSlicing() throws Exception { new PipeTaskRuntimeConfiguration(new PipeTaskSinkRuntimeEnvironment("pipe", 1L, 1))); final List receivedRequests = new ArrayList<>(); + final List receivedPacketSizes = new ArrayList<>(); final AtomicBoolean senderFinished = new AtomicBoolean(false); final AtomicReference receiverException = new AtomicReference<>(); final Thread receiverThread = @@ -176,6 +177,7 @@ public void testUdpAirGapSlicing() throws Exception { final DatagramPacket packet = new DatagramPacket(buffer, buffer.length); try { receiverSocket.receive(packet); + receivedPacketSizes.add(packet.getLength()); receivedRequests.add(decodeAirGapDatagram(buffer, packet.getLength())); final ByteBuffer responseBuffer = ByteBuffer.allocate(Long.BYTES + AirGapOneByteResponse.OK.length) @@ -213,6 +215,9 @@ public void testUdpAirGapSlicing() throws Exception { Assert.assertFalse(receiverThread.isAlive()); Assert.assertNull(receiverException.get()); Assert.assertTrue(receivedRequests.size() > 1); + for (final int packetSize : receivedPacketSizes) { + Assert.assertTrue(packetSize <= 1024); + } for (final TPipeTransferReq receivedRequest : receivedRequests) { Assert.assertEquals(PipeRequestType.TRANSFER_SLICE.getType(), receivedRequest.type); Assert.assertTrue(receivedRequest.body.remaining() < 1024); @@ -220,6 +225,62 @@ public void testUdpAirGapSlicing() throws Exception { } } + @Test + public void testUdpAirGapRejectsOversizedAcknowledgement() throws Exception { + try (final DatagramSocket receiverSocket = new DatagramSocket(0); + final UdpTestingIoTDBDataRegionAirGapSink sink = + new UdpTestingIoTDBDataRegionAirGapSink()) { + final PipeParameters parameters = new PipeParameters(buildParameterAttributes(false)); + parameters.getAttribute().put(PipeSinkConstant.CONNECTOR_AIR_GAP_TRANSPORT_KEY, "udp"); + sink.validate(new PipeParameterValidator(parameters)); + sink.customize( + parameters, + new PipeTaskRuntimeConfiguration(new PipeTaskSinkRuntimeEnvironment("pipe", 1L, 1))); + + final AtomicBoolean oversizedAcknowledgementSent = new AtomicBoolean(false); + final AtomicReference receiverException = new AtomicReference<>(); + final Thread receiverThread = + new Thread( + () -> { + try { + final byte[] buffer = new byte[2048]; + final DatagramPacket packet = new DatagramPacket(buffer, buffer.length); + receiverSocket.receive(packet); + final long requestId = ByteBuffer.wrap(buffer, 0, Long.BYTES).getLong(); + final byte[] oversizedAck = + ByteBuffer.allocate(Long.BYTES + AirGapOneByteResponse.OK.length + 1) + .putLong(requestId) + .put(AirGapOneByteResponse.OK) + .put((byte) 1) + .array(); + receiverSocket.send( + new DatagramPacket( + oversizedAck, + oversizedAck.length, + packet.getAddress(), + packet.getPort())); + oversizedAcknowledgementSent.set(true); + } catch (final Throwable t) { + receiverException.set(t); + } + }); + receiverThread.start(); + + try { + Assert.assertThrows( + SocketTimeoutException.class, + () -> sink.sendUdp(receiverSocket.getLocalPort(), new byte[] {1, 0, 1})); + } finally { + receiverSocket.close(); + } + + receiverThread.join(5000); + Assert.assertFalse(receiverThread.isAlive()); + Assert.assertNull(receiverException.get()); + Assert.assertTrue(oversizedAcknowledgementSent.get()); + } + } + @Test public void testUdpAirGapIgnoresStaleAcknowledgement() throws Exception { try (final DatagramSocket receiverSocket = new DatagramSocket(0); diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/protocol/IoTDBAirGapSink.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/protocol/IoTDBAirGapSink.java index 506b834ed77c6..bdffcfe473485 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/protocol/IoTDBAirGapSink.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/protocol/IoTDBAirGapSink.java @@ -39,14 +39,12 @@ 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.slf4j.Logger; import org.slf4j.LoggerFactory; import java.io.BufferedOutputStream; import java.io.DataInputStream; -import java.io.DataOutputStream; import java.io.File; import java.io.IOException; import java.io.RandomAccessFile; @@ -94,6 +92,7 @@ public abstract class IoTDBAirGapSink extends IoTDBSink { private static final int UDP_ENVELOPE_SIZE = 2 * Integer.BYTES + Long.BYTES; private static final int UDP_REQUEST_ID_SIZE = Long.BYTES; + private static final int UDP_ACK_SIZE = UDP_REQUEST_ID_SIZE + 1; private static final int UDP_SLICE_REQUEST_SERIALIZATION_RESERVED_SIZE = 64; protected static class AirGapSocket extends Socket { @@ -479,35 +478,56 @@ protected boolean sendBytes(final AirGapSocket socket, byte[] bytes) throws IOEx } private boolean sendBytesByUdp(final AirGapSocket socket, final byte[] bytes) throws IOException { - for (final byte[] requestBytes : sliceIfNeededForUdp(bytes)) { - if (!sendOneDatagram(socket, requestBytes)) { + for (final TPipeTransferReq request : sliceIfNeededForUdp(bytes)) { + if (!sendOneDatagram(socket, request)) { return false; } } return true; } - private boolean sendOneDatagram(final AirGapSocket socket, final byte[] requestBytes) + private boolean sendOneDatagram(final AirGapSocket socket, final TPipeTransferReq request) throws IOException { final long requestId = udpRequestIdGenerator.incrementAndGet(); - final byte[] framedRequestBytes = + final ByteBuffer body = request.body.duplicate(); + final int requestSize = Byte.BYTES + Short.BYTES + body.remaining(); + final int eLanguageSize = eLanguageEnable - ? enrichWithELanguage(enrichWithLengthAndChecksum(requestBytes)) - : enrichWithLengthAndChecksum(requestBytes); - final byte[] datagramBytes = - ByteBuffer.allocate(UDP_REQUEST_ID_SIZE + framedRequestBytes.length) - .putLong(requestId) - .put(framedRequestBytes) - .array(); - if (datagramBytes.length > udpPacketSizeInBytes) { + ? AirGapELanguageConstant.E_LANGUAGE_PREFIX.length + + AirGapELanguageConstant.E_LANGUAGE_SUFFIX.length + : 0; + final int datagramSize = UDP_REQUEST_ID_SIZE + UDP_ENVELOPE_SIZE + requestSize + eLanguageSize; + if (datagramSize > udpPacketSizeInBytes) { throw new IOException( String.format( PipeMessages .EXCEPTION_AIR_GAP_UDP_DATAGRAM_SIZE_ARG_EXCEEDS_CONFIGURED_PACKET_SIZE_ARG_559D0654, - datagramBytes.length, + datagramSize, udpPacketSizeInBytes)); } + final CRC32 crc32 = new CRC32(); + crc32.update(request.version); + crc32.update((request.type >>> Byte.SIZE) & 0xFF); + crc32.update(request.type & 0xFF); + crc32.update(body.duplicate()); + + final ByteBuffer datagramBuffer = ByteBuffer.allocate(datagramSize).putLong(requestId); + if (eLanguageEnable) { + datagramBuffer.put(AirGapELanguageConstant.E_LANGUAGE_PREFIX); + } + datagramBuffer + .putInt(requestSize + Long.BYTES) + .putInt(requestSize + Long.BYTES) + .putLong(crc32.getValue()) + .put(request.version) + .putShort(request.type) + .put(body); + if (eLanguageEnable) { + datagramBuffer.put(AirGapELanguageConstant.E_LANGUAGE_SUFFIX); + } + final byte[] datagramBytes = datagramBuffer.array(); + final DatagramSocket datagramSocket = socket.getDatagramSocket(); datagramSocket.send( new DatagramPacket( @@ -521,7 +541,8 @@ private boolean sendOneDatagram(final AirGapSocket socket, final byte[] requestB originalTimeoutMs == 0 ? Long.MAX_VALUE : System.nanoTime() + TimeUnit.MILLISECONDS.toNanos(originalTimeoutMs); - final byte[] response = new byte[UDP_REQUEST_ID_SIZE + 1]; + // One extra byte lets us detect datagrams that receive() would otherwise silently truncate. + final byte[] response = new byte[UDP_ACK_SIZE + 1]; final DatagramPacket responsePacket = new DatagramPacket(response, response.length); try { while (true) { @@ -537,7 +558,7 @@ private boolean sendOneDatagram(final AirGapSocket socket, final byte[] requestB responsePacket.setLength(response.length); datagramSocket.receive(responsePacket); - if (responsePacket.getLength() != UDP_REQUEST_ID_SIZE + 1) { + if (responsePacket.getLength() != UDP_ACK_SIZE) { continue; } @@ -545,14 +566,15 @@ private boolean sendOneDatagram(final AirGapSocket socket, final byte[] requestB if (responseBuffer.getLong() != requestId) { continue; } - return Arrays.equals(AirGapOneByteResponse.OK, new byte[] {responseBuffer.get()}); + return responseBuffer.get() == AirGapOneByteResponse.OK[0]; } } finally { datagramSocket.setSoTimeout(originalTimeoutMs); } } - private List sliceIfNeededForUdp(final byte[] requestBytes) throws IOException { + private List sliceIfNeededForUdp(final byte[] requestBytes) throws IOException { + final TPipeTransferReq request = toTPipeTransferReq(requestBytes); final int rawPayloadSizeLimit = udpPacketSizeInBytes - UDP_ENVELOPE_SIZE @@ -561,8 +583,8 @@ private List sliceIfNeededForUdp(final byte[] requestBytes) throws IOExc ? AirGapELanguageConstant.E_LANGUAGE_PREFIX.length + AirGapELanguageConstant.E_LANGUAGE_SUFFIX.length : 0); - if (requestBytes.length <= rawPayloadSizeLimit) { - return Arrays.asList(requestBytes); + if (Byte.BYTES + Short.BYTES + request.body.remaining() <= rawPayloadSizeLimit) { + return Arrays.asList(request); } final int sliceBodySizeLimit = @@ -575,18 +597,16 @@ private List sliceIfNeededForUdp(final byte[] requestBytes) throws IOExc udpPacketSizeInBytes)); } - final TPipeTransferReq request = toTPipeTransferReq(requestBytes); final int sliceOrderId = PipeTransferSliceReqBuilder.nextSliceOrderId(); final int sliceCount = PipeTransferSliceReqBuilder.getSliceCount(request, sliceBodySizeLimit); - final List slicedRequestBytes = new ArrayList<>(sliceCount); + final List slicedRequests = new ArrayList<>(sliceCount); for (int i = 0; i < sliceCount; i++) { - slicedRequestBytes.add( - toTPipeTransferBytes( - PipeTransferSliceReqBuilder.buildSliceReq( - request, sliceOrderId, i, sliceCount, sliceBodySizeLimit))); + slicedRequests.add( + PipeTransferSliceReqBuilder.buildSliceReq( + request, sliceOrderId, i, sliceCount, sliceBodySizeLimit)); } - return slicedRequestBytes; + return slicedRequests; } private TPipeTransferReq toTPipeTransferReq(final byte[] requestBytes) { @@ -598,21 +618,6 @@ private TPipeTransferReq toTPipeTransferReq(final byte[] requestBytes) { return request; } - private byte[] toTPipeTransferBytes(final TPipeTransferReq request) throws IOException { - try (final PublicBAOS byteArrayOutputStream = new PublicBAOS(); - final DataOutputStream outputStream = new DataOutputStream(byteArrayOutputStream)) { - ReadWriteIOUtils.write(request.version, outputStream); - ReadWriteIOUtils.write(request.type, outputStream); - - final ByteBuffer bodyBuffer = request.body.duplicate(); - final byte[] body = new byte[bodyBuffer.remaining()]; - bodyBuffer.get(body); - outputStream.write(body); - - return Arrays.copyOf(byteArrayOutputStream.getBuf(), byteArrayOutputStream.size()); - } - } - protected boolean send(final AirGapSocket socket, final byte[] bytes) throws IOException { return send(null, 0, socket, bytes); }