feat:new console

Author: Zxc4wonder

backend/framework/src/main/java/org/jumpserver/chen/framework/console/QueryConsole.java

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

backend/framework/src/main/java/org/jumpserver/chen/framework/console/context/ConsoleContextResolver.java

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

backend/framework/src/main/java/org/jumpserver/chen/framework/console/dataview/DataView.java

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

backend/framework/src/main/java/org/jumpserver/chen/framework/console/dataview/UpdateDataView.java

@@ -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;

backend/framework/src/main/java/org/jumpserver/chen/framework/console/entity/request/Connect.java

@@ -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;

backend/framework/src/main/java/org/jumpserver/chen/framework/console/state/DataViewState.java

@@ -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;
@@ -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();
     }

backend/framework/src/main/java/org/jumpserver/chen/framework/session/SessionManager.java

@@ -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 {
@@ -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
@@ -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();
@@ -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();

backend/framework/src/main/java/org/jumpserver/chen/framework/session/impl/BaseSession.java

@@ -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();
         }
@@ -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() {

backend/framework/src/main/java/org/jumpserver/chen/framework/session/impl/JMSSession.java

@@ -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());
         }
@@ -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;
         };