fix: harden terminal WebSocket edge cases
Deploy Dashboard (Dev) / deploy-dev (push) Successful in 2m44s

- Wrap tryRefreshSession in try-finally to ensure authRecoveryInFlight is always reset
- Fix variable scope issue: use const tabState instead of reassigning d
- Add error handling for process.stdin.write() to catch BrokenPipeError/OSError
- Verify WebSocket instance before closing in ping timeout callback
- Properly await conn.wait_closed() to prevent SSH connection leaks
- Gracefully cancel and wait for read_task with 2s timeout

These fixes prevent race conditions, resource leaks, and improve error recovery.
This commit is contained in:
Gan, Jimmy
2026-03-15 23:18:20 +08:00
parent d19707422d
commit b742ab8fef
2 changed files with 39 additions and 19 deletions
+15 -2
View File
@@ -190,18 +190,31 @@ async def ws_endpoint(websocket: WebSocket, host: str = Query("nas")):
rows = struct.unpack("!H", data[3:5])[0]
process.channel.change_terminal_size(cols, rows)
else:
process.stdin.write(data)
try:
process.stdin.write(data)
except (BrokenPipeError, OSError) as e:
logger.warning("Failed to write to stdin for %s: %s", host, e)
break
elif "text" in msg and msg["text"]:
if msg["text"] == "__ping__":
logger.debug("Received ping from %s, sending pong", host)
await websocket.send_text("__pong__")
continue
process.stdin.write(msg["text"].encode())
try:
process.stdin.write(msg["text"].encode())
except (BrokenPipeError, OSError) as e:
logger.warning("Failed to write to stdin for %s: %s", host, e)
break
except Exception as e:
logger.warning("WebSocket receive loop ended for %s: %s", host, e)
finally:
read_task.cancel()
try:
await asyncio.wait_for(read_task, timeout=2.0)
except (asyncio.TimeoutError, asyncio.CancelledError):
pass
process.close()
finally:
conn.close()
await conn.wait_closed()
await _release_session(session_id)
+24 -17
View File
@@ -300,7 +300,9 @@
ws.send("__ping__");
// Set timeout to detect missing pong
pongTimeout = setTimeout(() => {
if (ws.readyState === 1) {
// Check if still the same WebSocket instance and still open
const currentTab = tabData.get(tabId);
if (currentTab?.ws === ws && ws.readyState === 1) {
ws.close(1000, "keepalive ping timeout");
}
}, PONG_TIMEOUT);
@@ -343,8 +345,8 @@
};
ws.onclose = async (e) => {
const reason = e.reason || `Connection closed (code ${e.code})`;
d = tabData.get(tabId);
if (d) d.ws = null;
const tabState = tabData.get(tabId);
if (tabState) tabState.ws = null;
// Handle auth failures (code 1008)
if (e.code === 1008) {
@@ -359,18 +361,23 @@
// Try to refresh session
if (!closedManually && !authRecoveryInFlight) {
authRecoveryInFlight = true;
const refreshed = await tryRefreshSession();
authRecoveryInFlight = false;
if (refreshed && !closedManually) {
reconnecting = true;
updateTab(tabId, {
connected: false,
error: `Refreshing session (attempt ${consecutiveAuthFailures}/${MAX_AUTH_FAILURES})...`,
statusBanner: tab.persistent ? "Reconnecting to persistent session..." : "",
});
term.write(`\r\n\x1b[90m[Refreshing session...]\x1b[0m`);
void connect(false);
return;
try {
const refreshed = await tryRefreshSession();
if (refreshed && !closedManually) {
reconnecting = true;
updateTab(tabId, {
connected: false,
error: `Refreshing session (attempt ${consecutiveAuthFailures}/${MAX_AUTH_FAILURES})...`,
statusBanner: tab.persistent ? "Reconnecting to persistent session..." : "",
});
term.write(`\r\n\x1b[90m[Refreshing session...]\x1b[0m`);
void connect(false);
return;
}
} catch (err) {
console.error("Session refresh failed:", err);
} finally {
authRecoveryInFlight = false;
}
}
handleAuthExpired("Session expired — please log in again");
@@ -386,8 +393,8 @@
}
};
ws.onerror = () => {
d = tabData.get(tabId);
if (d) d.ws = null;
const tabState = tabData.get(tabId);
if (tabState) tabState.ws = null;
if (!closedManually) scheduleReconnect("Connection failed");
else updateTab(tabId, { connected: false, error: "Connection failed" });
};