Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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

Large diffs are not rendered by default.

Original file line number Diff line number Diff line change
Expand Up @@ -37,7 +37,7 @@ public ConsoleContext resolve(String submittedNodeKey, String consoleType) {

private static boolean isAllowedNodeType(String consoleType, String nodeType) {
return switch (consoleType) {
case Connect.CONSOLE_TYPE_QUERY -> QUERY_NODE_TYPES.contains(nodeType);
case Connect.CONSOLE_TYPE_QUERY, Connect.CONSOLE_TYPE_CONSOLE -> QUERY_NODE_TYPES.contains(nodeType);
case Connect.CONSOLE_TYPE_DATA_VIEW -> DATA_VIEW_NODE_TYPES.contains(nodeType);
default -> false;
};
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@
@EqualsAndHashCode(callSuper = true)
@Data
public class DataView extends SQLResult {
private final String id;
private final String title;
private final StateManager<DataViewState> stateManager;
private LoadDataInterface loadDataInterface;
Expand All @@ -38,8 +39,13 @@ public class DataView extends SQLResult {
private Logger consoleLogger;

public DataView(String title, PacketIO packetIO, Logger logger) {
this(title, title, packetIO, logger);
}

public DataView(String id, String title, PacketIO packetIO, Logger logger) {
this.id = id;
this.title = title;
this.state = new DataViewState(title);
this.state = new DataViewState(this.id, title);
this.stateManager = new StateManager<>(this.state, packetIO);
this.consoleLogger = logger;
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -4,10 +4,16 @@

@Data
public class UpdateDataView {
private String id;
private String title;
private DataViewData data;

public UpdateDataView(String title, DataViewData data) {
this(title, title, data);
}

public UpdateDataView(String id, String title, DataViewData data) {
this.id = id;
this.title = title;
this.data = data;
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
public class Connect {

public static final String CONSOLE_TYPE_QUERY = "query";
public static final String CONSOLE_TYPE_CONSOLE = "console";
public static final String CONSOLE_TYPE_DATA_VIEW = "data_view";

private String nodeKey;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -8,14 +8,16 @@
@EqualsAndHashCode(callSuper = true)
@Data
public class DataViewState extends State {
private String id;
private int page;
private int limit;
private int total;
private boolean pinned;
private boolean paged;

public DataViewState(String title) {
public DataViewState(String id, String title) {
super(title);
this.id = id;
this.paged = true;
this.pinned = false;
this.total = 0;
Expand Down
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
package org.jumpserver.chen.framework.session;

import lombok.extern.slf4j.Slf4j;
import org.jumpserver.chen.framework.console.Console;

import java.util.Map;
import java.util.UUID;
Expand Down Expand Up @@ -31,6 +32,13 @@ public static void unregisterSession(String token) {
log.info("session {} unregistered, current session count {}", token, instance.getCurrentSessionCount());
}

public static boolean registerConsole(String token, String consoleId, Console console) {
return instance.store.computeIfPresent(token, (ignored, session) -> {
session.getConsoles().put(consoleId, console);
return session;
}) != null;
}

public int getCurrentSessionCount() {
return instance.store.size();
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -23,10 +23,12 @@
import java.sql.Connection;
import java.sql.SQLException;
import java.time.LocalDateTime;
import java.util.ArrayList;
import java.util.List;
import java.util.Locale;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicBoolean;

@Slf4j
public class BaseSession implements Session {
Expand All @@ -44,6 +46,7 @@ public class BaseSession implements Session {

@Getter
Map<String, Console> consoles = new ConcurrentHashMap<>();
private final AtomicBoolean closeStarted = new AtomicBoolean(false);

@Getter
@Setter
Expand Down Expand Up @@ -201,10 +204,22 @@ public boolean isActive() {

@Override
public void close() {
if (!this.beginClose()) {
return;
}
this.closeSessionResources();
}

protected final boolean beginClose() {
return this.closeStarted.compareAndSet(false, true);
}

protected final void closeSessionResources() {
if (this.getController() != null) {
this.getController().cancelAllDialogs();
}
SessionManager.unregisterSession(this.getWebToken());
this.closeConsoles();
this.getDatasource().close();
this.getPacketIO().close();
var path = this.getTempPath();
Expand All @@ -213,6 +228,24 @@ public void close() {
}
}

private void closeConsoles() {
var detached = new ArrayList<Map.Entry<String, Console>>();
while (!this.consoles.isEmpty()) {
for (var entry : this.consoles.entrySet()) {
if (this.consoles.remove(entry.getKey(), entry.getValue())) {
detached.add(Map.entry(entry.getKey(), entry.getValue()));
}
}
}
for (var entry : detached) {
try {
entry.getValue().close();
} catch (RuntimeException e) {
log.warn("close console failed, consoleId={}", entry.getKey(), e);
}
}
}

@Override
public void close(String message, Object... args) {
this.close();
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -267,6 +267,13 @@ private void startWaitIdleTime() {

@Override
public void close() {
if (!this.beginClose()) {
return;
}
this.closeJmsSessionResources();
}

private void closeJmsSessionResources() {
if (this.getController() != null) {
this.getController().cancelAllDialogs();
}
Expand All @@ -279,23 +286,27 @@ public void close() {
}

} finally {
super.close();
super.closeSessionResources();
}
}

public void close(String message, String reason, Object... args) {
if (!this.beginClose()) {
return;
}
SessionManager.setContext(this.getWebToken());
try {
this.getPacketIO().sendPacket("session_close", null);

this.getPacketIO().sendPacket("session_close", null);

var dialog = new Dialog(MessageUtils.get("SessionFinished"));
dialog.setBody(MessageUtils.get(message, args));
this.getController().showDialog(dialog);

this.recordLifecycle(ServiceOuterClass.SessionLifecycleLogRequest.EventType.AssetConnectFinished, reason);
this.closed = true;
var dialog = new Dialog(MessageUtils.get("SessionFinished"));
dialog.setBody(MessageUtils.get(message, args));
this.getController().showDialog(dialog);

this.close();
this.recordLifecycle(ServiceOuterClass.SessionLifecycleLogRequest.EventType.AssetConnectFinished, reason);
this.closed = true;
} finally {
this.closeJmsSessionResources();
}
}

private void finishedJmsSession() {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -142,7 +142,12 @@ private void onConnectPacket(WebSocketSession session, Packet packet) {
Console console = this.createConsole(connect.getType(), webSess.getDatasource(), session, context);
if (console != null) {
this.setDatabaseContext(console);
webSess.getConsoles().put(session.getId(), console);
var token = (String) session.getAttributes().get("token");
if (!SessionManager.registerConsole(token, session.getId(), console)) {
console.close();
this.closeWebSocket(session);
return;
}
console.onInit(connect);
log.info("User {} open a console ", webSess.getUsername());
}
Expand All @@ -152,6 +157,7 @@ protected Console createConsole(String type, org.jumpserver.chen.framework.datas
WebSocketSession session, ConsoleContext context) {
return switch (type) {
case Connect.CONSOLE_TYPE_QUERY -> new QueryConsole(datasource, session, context);
case Connect.CONSOLE_TYPE_CONSOLE -> new QueryConsole(datasource, session, context, true);
case Connect.CONSOLE_TYPE_DATA_VIEW -> new DataViewConsole(datasource, session, context);
default -> null;
};
Expand Down