重构了基于websocket的SSH终端连接方案

code-server-like/app/sockets/terminal_service.py

@@ -14,6 +14,7 @@ import termios
 import struct
 import fcntl
 import psutil
+
 TERM_INIT_CONFIG = {
     # instead of local server runnning this web terminal service
     # "domain" is the target that you want to access through local server (with this web terminal)
@@ -24,9 +25,108 @@ TERM_INIT_CONFIG = {
         'ssh': '/usr/bin/ssh'
     }
 }
+
 terminal_bp = Blueprint('terminal', __name__, url_prefix='/vsc-like')
 
 
+class TerminalSessionManager:
+    """终端会话管理器，用于统一管理所有终端会话"""
+    def __init__(self):
+        # 存储所有终端会话，key为session_id，value为会话配置
+        self.sessions = {}
+        # 线程锁，保证多线程安全
+        self.lock = threading.Lock()
+        # 最大终端数量限制
+        self.max_terminals = 100
+        # 终端空闲超时时间（秒）
+        self.idle_timeout = 3600  # 1小时
+    
+    def create_session(self):
+        """创建新的终端会话"""
+        with self.lock:
+            # 检查是否达到最大终端数量
+            if len(self.sessions) >= self.max_terminals:
+                return None, "Maximum number of terminals reached"
+            
+            # 生成唯一的会话ID和房间ID
+            session_id = str(uuid.uuid4())
+            room_id = f"terminal_{session_id}"
+            
+            # 创建会话配置
+            session_config = {
+                **TERM_INIT_CONFIG.copy(),
+                'session_id': session_id,
+                'room_id': room_id,
+                'created_at': time.time(),
+                'last_activity': time.time()
+            }
+            
+            # 保存会话
+            self.sessions[session_id] = session_config
+            
+            return session_config, None
+    
+    def get_session(self, session_id):
+        """获取会话配置"""
+        with self.lock:
+            return self.sessions.get(session_id)
+    
+    def update_session(self, session_id, updates):
+        """更新会话配置"""
+        with self.lock:
+            if session_id not in self.sessions:
+                return False
+            
+            # 更新会话配置
+            self.sessions[session_id].update(updates)
+            # 更新最后活动时间
+            self.sessions[session_id]['last_activity'] = time.time()
+            
+            return True
+    
+    def delete_session(self, session_id):
+        """删除会话"""
+        with self.lock:
+            if session_id in self.sessions:
+                del self.sessions[session_id]
+                return True
+            return False
+    
+    def get_session_by_room_id(self, room_id):
+        """通过房间ID获取会话"""
+        with self.lock:
+            for session_config in self.sessions.values():
+                if session_config.get('room_id') == room_id:
+                    return session_config
+            return None
+    
+    def cleanup_inactive_sessions(self):
+        """清理不活动的会话"""
+        with self.lock:
+            current_time = time.time()
+            expired_sessions = []
+            
+            # 找出超时的会话
+            for session_id, session_config in self.sessions.items():
+                if current_time - session_config['last_activity'] > self.idle_timeout:
+                    expired_sessions.append(session_id)
+            
+            # 删除超时的会话
+            for session_id in expired_sessions:
+                del self.sessions[session_id]
+            
+            return expired_sessions
+    
+    def get_active_session_count(self):
+        """获取当前活动会话数量"""
+        with self.lock:
+            return len(self.sessions)
+
+
+# 创建全局终端会话管理器实例
+terminal_manager = TerminalSessionManager()
+
+
 
 def set_winsize(fd, row, col, xpix=0, ypix=0):
     try:
@@ -123,28 +223,26 @@ def read_and_forward_pty_output(fd=None, pid=None, room_id=None, namespace=None)
                 pass
 
 class VSCLikeNameSpace(Namespace):
-    def create_terminal(self):
-        """创建新的终端并更新session配置"""
-        terminal_config = session.get('terminal_config', {})
+    def __init__(self, namespace=None):
+        super().__init__(namespace)
+        # 存储客户端socket_id到session_id的映射
+        self.socket_session_map = {}
+        # 线程锁，保证多线程安全
+        self.lock = threading.Lock()
+    
+    def create_terminal(self, session_id):
+        """创建新的终端并更新会话配置"""
+        # 获取会话配置
+        terminal_config = terminal_manager.get_session(session_id)
+        if not terminal_config:
+            current_app.logger.error(f"Session not found: {session_id}")
+            return None
         
-        # 如果没有房间ID，创建一个新的
-        if not terminal_config.get('room_id'):
-            session_id = str(uuid.uuid4())
-            room_id = f"terminal_{session_id}"
-            join_room(room_id)
-            terminal_config.update({
-                **TERM_INIT_CONFIG.copy(),
-                'session_id': session_id,
-                'room_id': room_id
-            })
-        else:
         room_id = terminal_config['room_id']
-        
         current_app.logger.debug(f"Creating new terminal for room: {room_id}")
         
         # create child process attached to a pty we can read from and write to
         (child_pid, fd) = pty.fork()
-        lesson_path = session.get('path')
         if child_pid == 0:
             # this is the child process fork.
             # anything printed here will show up in the pty, including the output
@@ -157,14 +255,14 @@ class VSCLikeNameSpace(Namespace):
                     print("Terminal type not specified, exit")
                     os._exit(1)
                 
-                path = TERM_INIT_CONFIG.get('client_path', {}).get(term_type, None)
+                path = terminal_config.get('client_path', {}).get(term_type, None)
                 if not path or not os.path.exists(path):
                     print(f"Can't locate {term_type} binary at {path}, exit")
                     os._exit(1)
                 
                 # 获取连接参数
                 username = terminal_config.get('username')
-                domain = terminal_config.get('domain', TERM_INIT_CONFIG['domain'])
+                domain = terminal_config.get('domain')
                 port = terminal_config.get('port')
                 
                 if not username or not domain:
@@ -188,11 +286,12 @@ class VSCLikeNameSpace(Namespace):
                 print(f"Error in child process: {e}")
                 os._exit(1)
         else:
-            # 更新会话配置，保存到session中
-            terminal_config['fd'] = fd
-            terminal_config['child_pid'] = child_pid
-            session['terminal_config'] = terminal_config
-            session.modified = True
+            # 更新会话配置，保存到终端管理器
+            updates = {
+                'fd': fd,
+                'child_pid': child_pid
+            }
+            terminal_manager.update_session(session_id, updates)
             
             # 设置初始窗口大小
             set_winsize(fd, 50, 50)
@@ -208,31 +307,41 @@ class VSCLikeNameSpace(Namespace):
     
     def on_connect(self):
         """new client connected"""
-        # 为每个客户端生成唯一的会话ID和房间ID
-        session_id = str(uuid.uuid4())
-        room_id = f"terminal_{session_id}"
+        # 创建新的终端会话
+        terminal_config, error = terminal_manager.create_session()
+        if not terminal_config:
+            current_app.logger.error(f"Failed to create terminal session: {error}")
+            return
+        
+        session_id = terminal_config['session_id']
+        room_id = terminal_config['room_id']
         
         # 加入房间
         join_room(room_id)
         
-        # 确保terminal_config存在于session中，每个客户端有独立配置
-        terminal_config = {
-            **TERM_INIT_CONFIG.copy(),
-            'session_id': session_id,
-            'room_id': room_id
-        }
-        session['terminal_config'] = terminal_config
-        session.modified = True
+        # 保存socket_id到session_id的映射
+        socket_id = request.sid
+        with self.lock:
+            self.socket_session_map[socket_id] = session_id
         
-        current_app.logger.debug(f"Client connected, session_id: {session_id}, room_id: {room_id}")
+        current_app.logger.debug(f"Client connected, session_id: {session_id}, room_id: {room_id}, socket_id: {socket_id}")
         
         # 创建终端
-        self.create_terminal()
+        self.create_terminal(session_id)
         current_app.logger.debug("background task running")
 
     def on_pty_input(self, data):
         """write to the child pty, which now is the ssh process from this machine to the 'domain' configured
     """
+        # 获取socket_id对应的session_id
+        socket_id = request.sid
+        with self.lock:
+            session_id = self.socket_session_map.get(socket_id)
+        
+        if not session_id:
+            current_app.logger.error(f"No session found for socket: {socket_id}")
+            return
+        
         # 输入验证
         if not isinstance(data, dict):
             current_app.logger.error("Invalid input format: expected dictionary")
@@ -249,34 +358,47 @@ class VSCLikeNameSpace(Namespace):
             current_app.logger.warning("Input truncated due to length")
         
         try:
-            terminal_config = session.get('terminal_config', {})
+            terminal_config = terminal_manager.get_session(session_id)
+            if not terminal_config:
+                current_app.logger.error(f"Session not found: {session_id}")
+                return
+            
             child_pid = terminal_config.get('child_pid')
             
             # 如果没有子进程，重新创建终端
             if not child_pid:
-                current_app.logger.debug("No child process found, creating new terminal")
-                terminal_config = self.create_terminal()
+                current_app.logger.debug(f"No child process found for session: {session_id}, creating new terminal")
+                terminal_config = self.create_terminal(session_id)
+                if not terminal_config:
+                    current_app.logger.error(f"Failed to create new terminal for session: {session_id}")
+                    return
                 child_pid = terminal_config.get('child_pid')
                 if not child_pid:
-                    current_app.logger.error("Failed to create new terminal")
+                    current_app.logger.error(f"Failed to get child_pid for session: {session_id}")
                     return
             
             # 检查子进程是否存在且运行中
             try:
                 child_process = psutil.Process(child_pid)
                 if child_process.status() not in ('running', 'sleeping'):
-                    current_app.logger.debug("Child process not running, creating new terminal")
-                    terminal_config = self.create_terminal()
+                    current_app.logger.debug(f"Child process not running for session: {session_id}, creating new terminal")
+                    terminal_config = self.create_terminal(session_id)
+                    if not terminal_config:
+                        current_app.logger.error(f"Failed to create new terminal for session: {session_id}")
+                        return
                     child_pid = terminal_config.get('child_pid')
                     if not child_pid:
-                        current_app.logger.error("Failed to create new terminal")
+                        current_app.logger.error(f"Failed to get child_pid for session: {session_id}")
                         return
             except psutil.NoSuchProcess:
-                current_app.logger.debug("Child process not found, creating new terminal")
-                terminal_config = self.create_terminal()
+                current_app.logger.debug(f"Child process not found for session: {session_id}, creating new terminal")
+                terminal_config = self.create_terminal(session_id)
+                if not terminal_config:
+                    current_app.logger.error(f"Failed to create new terminal for session: {session_id}")
+                    return
                 child_pid = terminal_config.get('child_pid')
                 if not child_pid:
-                    current_app.logger.error("Failed to create new terminal")
+                    current_app.logger.error(f"Failed to get child_pid for session: {session_id}")
                     return
             
             # 获取文件描述符并写入数据
@@ -286,18 +408,27 @@ class VSCLikeNameSpace(Namespace):
                     os.write(fd, input_data.encode())
                 except (OSError, IOError) as e:
                     # File descriptor closed or invalid, create new terminal
-                    current_app.logger.debug(f"Error writing to file descriptor: {e}, creating new terminal")
-                    terminal_config = self.create_terminal()
+                    current_app.logger.debug(f"Error writing to file descriptor: {e}, creating new terminal for session: {session_id}")
+                    self.create_terminal(session_id)
         except Exception as e:
             current_app.logger.error(f"Error in pty_input handling: {e}")
             # 发生异常时尝试重新创建终端
             try:
-                self.create_terminal()
+                self.create_terminal(session_id)
             except Exception as create_err:
                 current_app.logger.error(f"Failed to recreate terminal: {create_err}")
 
     
     def on_resize(self, data):
+        # 获取socket_id对应的session_id
+        socket_id = request.sid
+        with self.lock:
+            session_id = self.socket_session_map.get(socket_id)
+        
+        if not session_id:
+            current_app.logger.error(f"No session found for socket: {socket_id}")
+            return
+        
         # 输入验证
         if not isinstance(data, dict):
             current_app.logger.error("Invalid resize data format: expected dictionary")
@@ -317,34 +448,47 @@ class VSCLikeNameSpace(Namespace):
             return
         
         try:
-            terminal_config = session.get('terminal_config', {})
+            terminal_config = terminal_manager.get_session(session_id)
+            if not terminal_config:
+                current_app.logger.error(f"Session not found: {session_id}")
+                return
+            
             child_pid = terminal_config.get('child_pid')
             
             # 如果没有子进程，重新创建终端
             if not child_pid:
-                current_app.logger.debug("No child process found for resize, creating new terminal")
-                terminal_config = self.create_terminal()
+                current_app.logger.debug(f"No child process found for resize, creating new terminal for session: {session_id}")
+                terminal_config = self.create_terminal(session_id)
+                if not terminal_config:
+                    current_app.logger.error(f"Failed to create new terminal for resize, session: {session_id}")
+                    return
                 child_pid = terminal_config.get('child_pid')
                 if not child_pid:
-                    current_app.logger.error("Failed to create new terminal for resize")
+                    current_app.logger.error(f"Failed to get child_pid for resize, session: {session_id}")
                     return
             
             # 检查子进程是否存在且运行中
             try:
                 child_process = psutil.Process(child_pid)
                 if child_process.status() not in ('running', 'sleeping'):
-                    current_app.logger.debug("Child process not running for resize, creating new terminal")
-                    terminal_config = self.create_terminal()
+                    current_app.logger.debug(f"Child process not running for resize, creating new terminal for session: {session_id}")
+                    terminal_config = self.create_terminal(session_id)
+                    if not terminal_config:
+                        current_app.logger.error(f"Failed to create new terminal for resize, session: {session_id}")
+                        return
                     child_pid = terminal_config.get('child_pid')
                     if not child_pid:
-                        current_app.logger.error("Failed to create new terminal for resize")
+                        current_app.logger.error(f"Failed to get child_pid for resize, session: {session_id}")
                         return
             except psutil.NoSuchProcess:
-                current_app.logger.debug("Child process not found for resize, creating new terminal")
-                terminal_config = self.create_terminal()
+                current_app.logger.debug(f"Child process not found for resize, creating new terminal for session: {session_id}")
+                terminal_config = self.create_terminal(session_id)
+                if not terminal_config:
+                    current_app.logger.error(f"Failed to create new terminal for resize, session: {session_id}")
+                    return
                 child_pid = terminal_config.get('child_pid')
                 if not child_pid:
-                    current_app.logger.error("Failed to create new terminal for resize")
+                    current_app.logger.error(f"Failed to get child_pid for resize, session: {session_id}")
                     return
             
             # 获取文件描述符并调整窗口大小
@@ -355,24 +499,37 @@ class VSCLikeNameSpace(Namespace):
                     set_winsize(fd, rows, cols)
                 except (OSError, IOError) as e:
                     # 文件描述符无效，创建新终端
-                    current_app.logger.debug(f"Error resizing terminal: {e}, creating new terminal")
-                    terminal_config = self.create_terminal()
+                    current_app.logger.debug(f"Error resizing terminal: {e}, creating new terminal for session: {session_id}")
+                    self.create_terminal(session_id)
         except Exception as e:
             current_app.logger.error(f"Error in resize handling: {e}")
             # 发生异常时尝试重新创建终端
             try:
-                self.create_terminal()
+                self.create_terminal(session_id)
             except Exception as create_err:
                 current_app.logger.error(f"Failed to recreate terminal for resize: {create_err}")
 
     def on_disconnect(self):
-        terminal_config = session.get('terminal_config', {})
+        # 获取socket_id对应的session_id
+        socket_id = request.sid
+        with self.lock:
+            session_id = self.socket_session_map.pop(socket_id, None)
+        
+        if not session_id:
+            current_app.logger.error(f"No session found for socket: {socket_id}")
+            return
+        
+        # 获取会话配置
+        terminal_config = terminal_manager.get_session(session_id)
+        if not terminal_config:
+            current_app.logger.error(f"Session not found: {session_id}")
+            return
+        
         child_pid = terminal_config.get('child_pid')
         fd = terminal_config.get('fd')
         room_id = terminal_config.get('room_id')
-        session_id = terminal_config.get('session_id')
         
-        current_app.logger.debug(f"Client disconnecting, session_id: {session_id}, room_id: {room_id}")
+        current_app.logger.debug(f"Client disconnecting, session_id: {session_id}, room_id: {room_id}, socket_id: {socket_id}")
         
         # 退出房间
         if room_id:
@@ -415,9 +572,9 @@ class VSCLikeNameSpace(Namespace):
             except Exception as err:
                 current_app.logger.error(f'Error terminating process: {err}')
         
-        # Clear session config for this client
-        session.pop('terminal_config', None)
-        session.modified = True
+        # 从终端管理器中删除会话
+        terminal_manager.delete_session(session_id)
+        current_app.logger.debug(f"Session deleted: {session_id}")
         
         current_app.logger.debug(f"Client disconnected, session_id: {session_id}, room_id: {room_id}")
         current_app.logger.debug('Client disconnected')