diff --git a/backend/framework/src/main/java/org/jumpserver/chen/framework/console/SQLChunkTransferManager.java b/backend/framework/src/main/java/org/jumpserver/chen/framework/console/SQLChunkTransferManager.java index 35b4a571..fc2f2928 100644 --- a/backend/framework/src/main/java/org/jumpserver/chen/framework/console/SQLChunkTransferManager.java +++ b/backend/framework/src/main/java/org/jumpserver/chen/framework/console/SQLChunkTransferManager.java @@ -16,6 +16,9 @@ final class SQLChunkTransferManager implements AutoCloseable { static final int MAX_CHUNK_SIZE = 4096; static final int MAX_CHUNKS = 1024; + // Luna's current chunk protocol has no requestId. A manager belongs to one QueryConsole, + // whose packets are processed serially, so one reserved key safely represents that transfer. + private static final TransferKey LEGACY_TRANSFER_KEY = new TransferKey(null); private static final int MAX_ACTIVE_TRANSFERS = 2; private static final int MAX_TRACKED_TRANSFERS = 64; private static final int MAX_REQUEST_ID_LENGTH = 128; @@ -23,7 +26,7 @@ final class SQLChunkTransferManager implements AutoCloseable { private static final ScheduledThreadPoolExecutor TIMEOUT_EXECUTOR = createTimeoutExecutor(); private final Object lock = new Object(); - private final Map transfers = new HashMap<>(); + private final Map transfers = new HashMap<>(); private final ScheduledExecutorService scheduler; private final Duration timeout; private boolean closed; @@ -42,7 +45,7 @@ final class SQLChunkTransferManager implements AutoCloseable { Optional receiveChunk(Object rawData) { Map data = requireMap(rawData); - String requestId = requireRequestId(data); + TransferKey transferKey = requireTransferKey(data); int total; int index; String chunk; @@ -54,13 +57,13 @@ Optional receiveChunk(Object rawData) { } chunk = requireChunk(data); } catch (IllegalArgumentException e) { - reject(requestId); + reject(transferKey); throw e; } synchronized (lock) { ensureOpen(); - ChunkTransfer transfer = getOrCreate(requestId, total); + ChunkTransfer transfer = getOrCreate(transferKey, total); if (transfer.terminal) { return Optional.empty(); } @@ -68,6 +71,10 @@ Optional receiveChunk(Object rawData) { rejectLocked(transfer); throw new IllegalArgumentException("total does not match the existing transfer"); } + if (transferKey.isLegacy() && index != transfer.receivedChunks) { + rejectLocked(transfer); + throw new IllegalArgumentException("legacy chunks must arrive in order"); + } String existing = transfer.chunks[index]; if (existing != null) { @@ -80,34 +87,41 @@ Optional receiveChunk(Object rawData) { transfer.chunks[index] = chunk; transfer.receivedChunks += 1; - return assembleIfReady(transfer); + return assembleIfReady(transferKey, transfer); } } Optional receiveComplete(Object rawData) { Map data = requireMap(rawData); - String requestId = requireRequestId(data); + TransferKey transferKey = requireTransferKey(data); int total; try { total = requireTotal(data); } catch (IllegalArgumentException e) { - reject(requestId); + rejectAndForgetLegacy(transferKey); throw e; } synchronized (lock) { ensureOpen(); - ChunkTransfer transfer = getOrCreate(requestId, total); + ChunkTransfer transfer = getOrCreate(transferKey, total); if (transfer.terminal) { + forgetLegacyTransferLocked(transferKey, transfer); return Optional.empty(); } if (transfer.total != total) { rejectLocked(transfer); + forgetLegacyTransferLocked(transferKey, transfer); throw new IllegalArgumentException("total does not match the existing transfer"); } + if (transferKey.isLegacy() && transfer.receivedChunks != transfer.total) { + rejectLocked(transfer); + forgetLegacyTransferLocked(transferKey, transfer); + throw new IllegalArgumentException("legacy transfer completed before all chunks arrived"); + } transfer.completeReceived = true; - return assembleIfReady(transfer); + return assembleIfReady(transferKey, transfer); } } @@ -128,8 +142,8 @@ public void close() { } } - private ChunkTransfer getOrCreate(String requestId, int total) { - ChunkTransfer existing = transfers.get(requestId); + private ChunkTransfer getOrCreate(TransferKey transferKey, int total) { + ChunkTransfer existing = transfers.get(transferKey); if (existing != null) { return existing; } @@ -144,21 +158,21 @@ private ChunkTransfer getOrCreate(String requestId, int total) { } ChunkTransfer transfer = new ChunkTransfer(total); - transfers.put(requestId, transfer); + transfers.put(transferKey, transfer); try { transfer.timeoutFuture = scheduler.schedule( - () -> expire(requestId, transfer), + () -> expire(transferKey, transfer), timeout.toMillis(), TimeUnit.MILLISECONDS ); } catch (RuntimeException e) { - transfers.remove(requestId, transfer); + transfers.remove(transferKey, transfer); throw e; } return transfer; } - private Optional assembleIfReady(ChunkTransfer transfer) { + private Optional assembleIfReady(TransferKey transferKey, ChunkTransfer transfer) { if (!transfer.completeReceived || transfer.receivedChunks != transfer.total) { return Optional.empty(); } @@ -169,14 +183,25 @@ private Optional assembleIfReady(ChunkTransfer transfer) { } String assembled = sql.toString(); rejectLocked(transfer); + forgetLegacyTransferLocked(transferKey, transfer); return Optional.of(assembled); } - private void reject(String requestId) { + private void reject(TransferKey transferKey) { + synchronized (lock) { + ChunkTransfer transfer = transfers.get(transferKey); + if (transfer != null) { + rejectLocked(transfer); + } + } + } + + private void rejectAndForgetLegacy(TransferKey transferKey) { synchronized (lock) { - ChunkTransfer transfer = transfers.get(requestId); + ChunkTransfer transfer = transfers.get(transferKey); if (transfer != null) { rejectLocked(transfer); + forgetLegacyTransferLocked(transferKey, transfer); } } } @@ -191,14 +216,23 @@ private void rejectLocked(ChunkTransfer transfer) { Arrays.fill(transfer.chunks, null); } - private void expire(String requestId, ChunkTransfer expectedTransfer) { + private void expire(TransferKey transferKey, ChunkTransfer expectedTransfer) { synchronized (lock) { - if (transfers.remove(requestId, expectedTransfer)) { + if (transfers.remove(transferKey, expectedTransfer)) { rejectLocked(expectedTransfer); } } } + private void forgetLegacyTransferLocked(TransferKey transferKey, ChunkTransfer transfer) { + if (!transferKey.isLegacy() || !transfers.remove(transferKey, transfer)) { + return; + } + if (transfer.timeoutFuture != null) { + transfer.timeoutFuture.cancel(false); + } + } + private void clearTransfersLocked() { for (ChunkTransfer transfer : transfers.values()) { rejectLocked(transfer); @@ -222,14 +256,17 @@ private void ensureOpen() { return map; } - private static String requireRequestId(Map data) { + private static TransferKey requireTransferKey(Map data) { Object value = data.get("requestId"); + if (value == null) { + return LEGACY_TRANSFER_KEY; + } if (!(value instanceof String requestId) || requestId.isBlank() || requestId.length() > MAX_REQUEST_ID_LENGTH) { throw new IllegalArgumentException("requestId is invalid"); } - return requestId; + return new TransferKey(requestId); } private static int requireTotal(Map data) { @@ -286,4 +323,10 @@ private ChunkTransfer(int total) { this.chunks = new String[total]; } } + + private record TransferKey(String requestId) { + private boolean isLegacy() { + return requestId == null; + } + } } diff --git a/backend/framework/src/test/java/org/jumpserver/chen/framework/console/SQLChunkTransferManagerTest.java b/backend/framework/src/test/java/org/jumpserver/chen/framework/console/SQLChunkTransferManagerTest.java new file mode 100644 index 00000000..a20417ac --- /dev/null +++ b/backend/framework/src/test/java/org/jumpserver/chen/framework/console/SQLChunkTransferManagerTest.java @@ -0,0 +1,408 @@ +package org.jumpserver.chen.framework.console; + +import org.junit.jupiter.api.Test; + +import java.time.Duration; +import java.util.HashMap; +import java.util.Map; +import java.util.Optional; +import java.util.concurrent.FutureTask; +import java.util.concurrent.ScheduledFuture; +import java.util.concurrent.ScheduledThreadPoolExecutor; +import java.util.concurrent.TimeUnit; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +class SQLChunkTransferManagerTest { + + @Test + void assemblesCurrentLunaPayloadWithoutRequestId() { + String sql = "x".repeat(10_722); + + try (SQLChunkTransferManager manager = new SQLChunkTransferManager()) { + Optional assembled = sendLegacyChunks(manager, sql); + + assertTrue(assembled.isPresent()); + assertEquals(sql, assembled.orElseThrow()); + } + } + + @Test + void acceptsConsecutiveLegacyTransfers() { + try (SQLChunkTransferManager manager = new SQLChunkTransferManager()) { + assertEquals("first", sendLegacyChunks(manager, "first").orElseThrow()); + assertEquals("second", sendLegacyChunks(manager, "second").orElseThrow()); + } + } + + @Test + void rejectsLegacyCompletionBeforeLastChunkAndRecovers() { + try (SQLChunkTransferManager manager = new SQLChunkTransferManager()) { + assertFalse(manager.receiveChunk(legacyChunk("first-", 0, 2)).isPresent()); + assertThrows( + IllegalArgumentException.class, + () -> manager.receiveComplete(legacyComplete(2)) + ); + + assertEquals("recovered", sendLegacyChunks(manager, "recovered").orElseThrow()); + } + } + + @Test + void rejectsOutOfOrderLegacyChunksWithoutExecuting() { + try (SQLChunkTransferManager manager = new SQLChunkTransferManager()) { + assertThrows( + IllegalArgumentException.class, + () -> manager.receiveChunk(legacyChunk("second", 1, 2)) + ); + assertFalse(manager.receiveComplete(legacyComplete(2)).isPresent()); + + assertEquals("ordered", sendLegacyChunks(manager, "ordered").orElseThrow()); + } + } + + @Test + void clearsRejectedLegacyTransferAtCompletionBoundary() { + try (SQLChunkTransferManager manager = new SQLChunkTransferManager()) { + assertFalse(manager.receiveChunk(legacyChunk("first", 0, 2)).isPresent()); + assertThrows( + IllegalArgumentException.class, + () -> manager.receiveChunk(legacyChunk("different", 0, 2)) + ); + assertFalse(manager.receiveComplete(legacyComplete(2)).isPresent()); + + assertEquals("recovered", sendLegacyChunks(manager, "recovered").orElseThrow()); + } + } + + @Test + void keepsExplicitRequestIdTransfersIsolated() { + try (SQLChunkTransferManager manager = new SQLChunkTransferManager()) { + assertFalse(manager.receiveChunk(chunk("request-a", "a1", 0, 2)).isPresent()); + assertFalse(manager.receiveChunk(chunk("request-b", "b1", 0, 2)).isPresent()); + assertFalse(manager.receiveComplete(complete("request-a", 2)).isPresent()); + assertFalse(manager.receiveComplete(complete("request-b", 2)).isPresent()); + + assertEquals("a1a2", manager.receiveChunk(chunk("request-a", "a2", 1, 2)).orElseThrow()); + assertEquals("b1b2", manager.receiveChunk(chunk("request-b", "b2", 1, 2)).orElseThrow()); + } + } + + @Test + void rejectsInvalidProvidedRequestId() { + try (SQLChunkTransferManager manager = new SQLChunkTransferManager()) { + Map blankRequestId = legacyChunk("sql", 0, 1); + blankRequestId.put("requestId", " "); + Map numericRequestId = legacyChunk("sql", 0, 1); + numericRequestId.put("requestId", 123); + + assertThrows(IllegalArgumentException.class, () -> manager.receiveChunk(blankRequestId)); + assertThrows(IllegalArgumentException.class, () -> manager.receiveChunk(numericRequestId)); + } + } + + @Test + void acceptsChunkAtExactSizeLimit() { + String sql = "x".repeat(SQLChunkTransferManager.MAX_CHUNK_SIZE); + + try (SQLChunkTransferManager manager = new SQLChunkTransferManager()) { + assertFalse(manager.receiveChunk(legacyChunk(sql, 0, 1)).isPresent()); + assertEquals(sql, manager.receiveComplete(legacyComplete(1)).orElseThrow()); + } + } + + @Test + void rejectsChunkAboveSizeLimitAndRecovers() { + String oversized = "x".repeat(SQLChunkTransferManager.MAX_CHUNK_SIZE + 1); + + try (SQLChunkTransferManager manager = new SQLChunkTransferManager()) { + assertThrows( + IllegalArgumentException.class, + () -> manager.receiveChunk(legacyChunk(oversized, 0, 1)) + ); + assertThrows( + IllegalArgumentException.class, + () -> manager.receiveComplete(legacyComplete(1)) + ); + + assertEquals("recovered", sendLegacyChunks(manager, "recovered").orElseThrow()); + } + } + + @Test + void acceptsMaximumChunkCountAndRejectsValuesOutsideBounds() { + try (SQLChunkTransferManager manager = new SQLChunkTransferManager()) { + for (int index = 0; index < SQLChunkTransferManager.MAX_CHUNKS; index++) { + assertFalse(manager.receiveChunk( + legacyChunk("x", index, SQLChunkTransferManager.MAX_CHUNKS) + ).isPresent()); + } + assertEquals( + "x".repeat(SQLChunkTransferManager.MAX_CHUNKS), + manager.receiveComplete(legacyComplete(SQLChunkTransferManager.MAX_CHUNKS)).orElseThrow() + ); + + assertThrows(IllegalArgumentException.class, () -> manager.receiveChunk(legacyChunk("sql", 0, 0))); + assertThrows(IllegalArgumentException.class, () -> manager.receiveChunk(legacyChunk("sql", 0, -1))); + assertThrows( + IllegalArgumentException.class, + () -> manager.receiveChunk(legacyChunk("sql", 0, SQLChunkTransferManager.MAX_CHUNKS + 1)) + ); + } + } + + @Test + void rejectsIndexesOutsideBounds() { + try (SQLChunkTransferManager manager = new SQLChunkTransferManager()) { + assertThrows(IllegalArgumentException.class, () -> manager.receiveChunk(legacyChunk("sql", -1, 1))); + assertThrows(IllegalArgumentException.class, () -> manager.receiveChunk(legacyChunk("sql", 1, 1))); + } + } + + @Test + void rejectsNonIntegralOrNonFiniteNumericFields() { + try (SQLChunkTransferManager manager = new SQLChunkTransferManager()) { + assertThrows( + IllegalArgumentException.class, + () -> manager.receiveChunk(chunkWithNumbers("sql", 0.5, 1.0)) + ); + assertThrows( + IllegalArgumentException.class, + () -> manager.receiveChunk(chunkWithNumbers("sql", Double.NaN, 1.0)) + ); + assertThrows( + IllegalArgumentException.class, + () -> manager.receiveChunk(chunkWithNumbers("sql", 0.0, Double.POSITIVE_INFINITY)) + ); + assertThrows( + IllegalArgumentException.class, + () -> manager.receiveChunk(Map.of("chunk", "sql", "index", "0", "total", 1.0)) + ); + } + } + + @Test + void rejectsMissingOrInvalidChunkPayload() { + try (SQLChunkTransferManager manager = new SQLChunkTransferManager()) { + assertThrows( + IllegalArgumentException.class, + () -> manager.receiveChunk(Map.of("index", 0.0, "total", 1.0)) + ); + assertThrows( + IllegalArgumentException.class, + () -> manager.receiveChunk(Map.of("chunk", 42, "index", 0.0, "total", 1.0)) + ); + assertThrows(IllegalArgumentException.class, () -> manager.receiveChunk(null)); + assertThrows(IllegalArgumentException.class, () -> manager.receiveChunk("not-a-map")); + } + } + + @Test + void enforcesRequestIdLengthBoundary() { + String maximumRequestId = "r".repeat(128); + String oversizedRequestId = "r".repeat(129); + + try (SQLChunkTransferManager manager = new SQLChunkTransferManager()) { + assertFalse(manager.receiveChunk(chunk(maximumRequestId, "sql", 0, 1)).isPresent()); + assertEquals("sql", manager.receiveComplete(complete(maximumRequestId, 1)).orElseThrow()); + assertThrows( + IllegalArgumentException.class, + () -> manager.receiveChunk(chunk(oversizedRequestId, "sql", 0, 1)) + ); + } + } + + @Test + void explicitRequestIdSupportsOutOfOrderChunksAndEarlyComplete() { + try (SQLChunkTransferManager manager = new SQLChunkTransferManager()) { + assertFalse(manager.receiveChunk(chunk("request", "third", 2, 3)).isPresent()); + assertFalse(manager.receiveComplete(complete("request", 3)).isPresent()); + assertFalse(manager.receiveChunk(chunk("request", "first-", 0, 3)).isPresent()); + + assertEquals( + "first-second-third", + manager.receiveChunk(chunk("request", "second-", 1, 3)).orElseThrow() + ); + } + } + + @Test + void explicitRequestIdIgnoresMatchingDuplicateAndRejectsConflictingDuplicate() { + try (SQLChunkTransferManager manager = new SQLChunkTransferManager()) { + assertFalse(manager.receiveChunk(chunk("request", "first", 0, 2)).isPresent()); + assertFalse(manager.receiveChunk(chunk("request", "first", 0, 2)).isPresent()); + assertThrows( + IllegalArgumentException.class, + () -> manager.receiveChunk(chunk("request", "different", 0, 2)) + ); + assertFalse(manager.receiveComplete(complete("request", 2)).isPresent()); + } + } + + @Test + void rejectsLegacyTotalMismatchAndRecoversAtCompletionBoundary() { + try (SQLChunkTransferManager manager = new SQLChunkTransferManager()) { + assertFalse(manager.receiveChunk(legacyChunk("first", 0, 2)).isPresent()); + assertThrows( + IllegalArgumentException.class, + () -> manager.receiveChunk(legacyChunk("second", 1, 3)) + ); + assertFalse(manager.receiveComplete(legacyComplete(2)).isPresent()); + + assertEquals("recovered", sendLegacyChunks(manager, "recovered").orElseThrow()); + } + } + + @Test + void rejectsLegacyCompletionTotalMismatchAndRecoversImmediately() { + try (SQLChunkTransferManager manager = new SQLChunkTransferManager()) { + assertFalse(manager.receiveChunk(legacyChunk("first", 0, 2)).isPresent()); + assertThrows( + IllegalArgumentException.class, + () -> manager.receiveComplete(legacyComplete(3)) + ); + + assertEquals("recovered", sendLegacyChunks(manager, "recovered").orElseThrow()); + } + } + + @Test + void cancelAllDiscardsAnActiveLegacyTransfer() { + try (SQLChunkTransferManager manager = new SQLChunkTransferManager()) { + assertFalse(manager.receiveChunk(legacyChunk("stale", 0, 2)).isPresent()); + manager.cancelAll(); + + assertEquals("fresh", sendLegacyChunks(manager, "fresh").orElseThrow()); + } + } + + @Test + void closeRejectsFurtherPackets() { + SQLChunkTransferManager manager = new SQLChunkTransferManager(); + manager.close(); + + assertThrows( + IllegalStateException.class, + () -> manager.receiveChunk(legacyChunk("sql", 0, 1)) + ); + } + + @Test + void timeoutDiscardsAnIncompleteLegacyTransfer() { + ManualScheduler scheduler = new ManualScheduler(); + try (SQLChunkTransferManager manager = new SQLChunkTransferManager(scheduler, Duration.ofSeconds(30))) { + assertFalse(manager.receiveChunk(legacyChunk("stale", 0, 2)).isPresent()); + scheduler.runScheduled(); + + assertEquals("fresh", sendLegacyChunks(manager, "fresh").orElseThrow()); + } finally { + scheduler.shutdownNow(); + } + } + + @Test + void limitsConcurrentActiveTransfers() { + try (SQLChunkTransferManager manager = new SQLChunkTransferManager()) { + assertFalse(manager.receiveChunk(chunk("request-a", "a", 0, 2)).isPresent()); + assertFalse(manager.receiveChunk(chunk("request-b", "b", 0, 2)).isPresent()); + + assertThrows( + IllegalArgumentException.class, + () -> manager.receiveChunk(chunk("request-c", "c", 0, 2)) + ); + } + } + + @Test + void preservesSurrogatePairSplitAtChunkBoundary() { + String sql = "x".repeat(SQLChunkTransferManager.MAX_CHUNK_SIZE - 1) + "😀" + "tail"; + + try (SQLChunkTransferManager manager = new SQLChunkTransferManager()) { + assertEquals(sql, sendLegacyChunks(manager, sql).orElseThrow()); + } + } + + private static Optional sendLegacyChunks(SQLChunkTransferManager manager, String sql) { + int chunkSize = SQLChunkTransferManager.MAX_CHUNK_SIZE; + int total = Math.max(1, (sql.length() + chunkSize - 1) / chunkSize); + for (int index = 0; index < total; index++) { + String chunk = sql.substring(index * chunkSize, Math.min(sql.length(), (index + 1) * chunkSize)); + assertFalse(manager.receiveChunk(legacyChunk(chunk, index, total)).isPresent()); + } + return manager.receiveComplete(legacyComplete(total)); + } + + private static Map legacyChunk(String chunk, int index, int total) { + Map data = new HashMap<>(); + data.put("chunk", chunk); + data.put("index", (double) index); + data.put("total", (double) total); + return data; + } + + private static Map legacyComplete(int total) { + Map data = new HashMap<>(); + data.put("total", (double) total); + return data; + } + + private static Map chunkWithNumbers(String chunk, Object index, Object total) { + Map data = new HashMap<>(); + data.put("chunk", chunk); + data.put("index", index); + data.put("total", total); + return data; + } + + private static Map chunk(String requestId, String chunk, int index, int total) { + Map data = legacyChunk(chunk, index, total); + data.put("requestId", requestId); + return data; + } + + private static Map complete(String requestId, int total) { + Map data = legacyComplete(total); + data.put("requestId", requestId); + return data; + } + + private static final class ManualScheduler extends ScheduledThreadPoolExecutor { + private ManualScheduledFuture scheduled; + + private ManualScheduler() { + super(1); + } + + @Override + public ScheduledFuture schedule(Runnable command, long delay, TimeUnit unit) { + this.scheduled = new ManualScheduledFuture(command); + return this.scheduled; + } + + private void runScheduled() { + if (this.scheduled == null) { + throw new IllegalStateException("no task scheduled"); + } + this.scheduled.run(); + } + } + + private static final class ManualScheduledFuture extends FutureTask implements ScheduledFuture { + private ManualScheduledFuture(Runnable runnable) { + super(runnable, null); + } + + @Override + public long getDelay(TimeUnit unit) { + return 0; + } + + @Override + public int compareTo(java.util.concurrent.Delayed ignored) { + return 0; + } + } +}