Skip to content

fix(console): accept Luna SQL chunks #293

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Zxc4wonder merged 1 commit into from
Sep 20, 2026
+472 −21
Merged

fix(console): accept Luna SQL chunks #293

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
85 changes: 64 additions & 21 deletions ...ramework/src/main/java/org/jumpserver/chen/framework/console/SQLChunkTransferManager.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,14 +16,17 @@ 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;
private static final Duration DEFAULT_TIMEOUT = Duration.ofSeconds(30);
private static final ScheduledThreadPoolExecutor TIMEOUT_EXECUTOR = createTimeoutExecutor();

private final Object lock = new Object();
private final Map<String, ChunkTransfer> transfers = new HashMap<>();
private final Map<TransferKey, ChunkTransfer> transfers = new HashMap<>();
private final ScheduledExecutorService scheduler;
private final Duration timeout;
private boolean closed;
Expand All @@ -42,7 +45,7 @@ final class SQLChunkTransferManager implements AutoCloseable {

Optional<String> receiveChunk(Object rawData) {
Map<?, ?> data = requireMap(rawData);
String requestId = requireRequestId(data);
TransferKey transferKey = requireTransferKey(data);
int total;
int index;
String chunk;
Expand All @@ -54,20 +57,24 @@ Optional<String> 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();
}
if (transfer.total != total) {
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) {
Expand All @@ -80,34 +87,41 @@ Optional<String> receiveChunk(Object rawData) {

transfer.chunks[index] = chunk;
transfer.receivedChunks += 1;
return assembleIfReady(transfer);
return assembleIfReady(transferKey, transfer);
}
}

Optional<String> 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);
}
}

Expand All @@ -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;
}
Expand All @@ -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<String> assembleIfReady(ChunkTransfer transfer) {
private Optional<String> assembleIfReady(TransferKey transferKey, ChunkTransfer transfer) {
if (!transfer.completeReceived || transfer.receivedChunks != transfer.total) {
return Optional.empty();
}
Expand All @@ -169,14 +183,25 @@ private Optional<String> 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);
}
}
}
Expand All @@ -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);
Expand All @@ -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) {
Expand Down Expand Up @@ -286,4 +323,10 @@ private ChunkTransfer(int total) {
this.chunks = new String[total];
}
}

private record TransferKey(String requestId) {
private boolean isLegacy() {
return requestId == null;
}
}
}
Loading
Loading