Nano and timeout error fix #3

backend/server.js

@@ -2,20 +2,18 @@ const WebSocket = require('ws');
 const ssh2 = require('ssh2');
 const http = require('http');
 
-// Create an HTTP server to serve WebSocket connections
+// Create an HTTP server
 const server = http.createServer((req, res) => {
   res.writeHead(200, { 'Content-Type': 'text/plain' });
   res.end('WebSocket server is running\n');
 });
 
-// Create a WebSocket server attached to the HTTP server
+// Create a WebSocket server
 const wss = new WebSocket.Server({ server });
 
-// WebSocket connection handling
 wss.on('connection', (ws) => {
   console.log('WebSocket connection established');
-
-  let conn = null; // Declare SSH client outside to manage lifecycle
+  let conn = null;
 
   // Ping-Pong for WebSocket Keep-Alives
   const interval = setInterval(() => {
@@ -32,96 +30,84 @@ wss.on('connection', (ws) => {
 
   ws.on('message', (message) => {
     try {
-      const data = JSON.parse(message); // Try parsing the incoming message as JSON
+      const data = JSON.parse(message);
 
-      // Check if message contains SSH connection details
       if (data.host && data.username && data.password) {
         if (conn) {
-          conn.end(); // Close any previous connection before starting a new one
+          conn.end();
         }
 
-        conn = new ssh2.Client(); // Create a new SSH connection instance
+        conn = new ssh2.Client();
 
-        // When the SSH connection is ready
-        conn.on('ready', () => {
-          console.log('SSH Connection established');
-
-          // Start an interactive shell session
-          conn.shell((err, stream) => {
-            if (err) {
-              console.log(`SSH Error: ${err}`);
-              ws.send(`Error: ${err}`);
-              return;
-            }
-
-            // Handle data from SSH session
-            stream.on('data', (data) => {
-              console.log(`SSH Output: ${data}`);
-              ws.send(data.toString()); // Send the SSH output back to WebSocket client
-            });
-
-            // Handle stream close event
-            stream.on('close', () => {
-              console.log('SSH stream closed');
-              conn.end();
-            });
-
-            // When the WebSocket client sends a message (from terminal input), forward it to the SSH stream
-            ws.on('message', (message) => {
-              if (typeof message === 'string') {
-                try {
-                  const resizeEvent = JSON.parse(message);
-                  if (resizeEvent.type === 'resize') {
-                    stream.setWindow(
-                      resizeEvent.rows,
-                      resizeEvent.cols,
-                      resizeEvent.height,
-                      resizeEvent.width
-                    );
-                  }
-                } catch {
-                  stream.write(message);
-                }
+        conn
+          .on('ready', () => {
+            console.log('SSH Connection established');
+            conn.shell((err, stream) => {
+              if (err) {
+                ws.send(`Error: ${err.message}`);
+                return;
               }
+
+              stream.on('data', (data) => {
+                ws.send(data.toString());
+              });
+
+              stream.on('close', () => {
+                console.log('SSH Stream closed');
+                conn.end();
+              });
+
+              // Forward user input and resize events to the SSH stream
+              ws.on('message', (message) => {
+                if (typeof message === 'string') {
+                  try {
+                    const resizeEvent = JSON.parse(message);
+                    if (resizeEvent.type === 'resize') {
+                      stream.setWindow(
+                        resizeEvent.rows,
+                        resizeEvent.cols,
+                        resizeEvent.height,
+                        resizeEvent.width
+                      );
+                    }
+                  } catch {
+                    stream.write(message);
+                  }
+                }
+              });
             });
+          })
+          .on('error', (err) => {
+            console.log('SSH Error:', err.message);
+            ws.send(`SSH Error: ${err.message}`);
+          })
+          .on('close', () => {
+            console.log('SSH Connection closed');
+          })
+          .connect({
+            host: data.host,
+            port: 22,
+            username: data.username,
+            password: data.password,
+            keepaliveInterval: 20000, // Send SSH keepalive every 20 seconds
+            keepaliveCountMax: 5, // Allow 5 missed keepalives before disconnecting
           });
-        }).on('error', (err) => {
-          console.log('SSH Connection Error: ', err);
-          ws.send(`SSH Error: ${err}`);
-        }).connect({
-          host: data.host,       // Host provided from the client
-          port: 22,              // Default SSH port
-          username: data.username,  // Username provided from the client
-          password: data.password,  // Password provided from the client
-          keepaliveInterval: 15000, // Send keep-alive packets every 15 seconds
-          keepaliveCountMax: 3,    // Retry keep-alive 3 times before disconnecting
-        });
       }
     } catch (error) {
-      // If message is not valid JSON (i.e., terminal input), treat it as raw text and send it to SSH
-      console.log('Received non-JSON message, sending to SSH session:', message);
-      if (conn) {
-        const stream = conn._stream; // Access the SSH stream directly
-        if (stream && stream.writable) {
-          stream.write(message); // Write raw input message to SSH stream
-        }
-      } else {
-        console.error('SSH connection is not established yet.');
-      }
+      console.log('Non-JSON message received:', message);
     }
   });
 
-  // Handle WebSocket close event
   ws.on('close', () => {
     console.log('WebSocket closed');
     clearInterval(interval);
     if (conn) {
-      conn.end(); // Close SSH connection when WebSocket client disconnects
+      conn.end();
     }
   });
 });
 
-// Start the WebSocket server on port 8081
+// Start the server
 server.listen(8081, () => {
-  console.log('WebSocket server is listening on ws://localhost:8081');
+  console.log('WebSocket server is running on ws://localhost:8081');
 });