Allow ComputerCraft to reconnect to the same session using a token

computercraft/sigils.lua

@@ -73,6 +73,8 @@ local function init ()
   local wsContext = {
     wsUrl = config.server or DEFAULT_SERVER_URL,
     ws = nil,
+    reconnectToken = nil,
+    sessionId = string,
   }
 
   parallel.waitForAll(

computercraft/sigils/websocket.lua

@@ -24,27 +24,39 @@ local MESSAGE_TYPES = {
   GroupDel = true,
 }
 
----Request a session from the editor session server once
----@param ws Websocket ComputerCraft Websocket handle
+---Request a session from the editor session server once, or reconnect if
+---there's a reconnect token
+---@param wsContext table WebSocket context
 ---@return table response ConfirmationResponse as a Lua table
----@return string sessionId Session ID requested
-local function requestSessionOnce (ws)
-  local sessionId = Utils.randomString(4)
-  local req = {
-    type = 'SessionCreate',
-    reqId = Utils.randomString(20),
-    sessionId = sessionId,
-  }
-  ws.send(textutils.serializeJSON(req))
-
-  local res = ws.receive(5)
-  return textutils.unserializeJSON(res), sessionId
+local function requestSessionOnce (wsContext)
+  local req
+
+  if wsContext.reconnectToken then
+    req = {
+      type = 'SessionRejoin',
+      reqId = Utils.randomString(20),
+      ccReconnectToken = wsContext.reconnectToken,
+      sessionId = wsContext.sessionId,
+    }
+  else
+    wsContext.sessionId = Utils.randomString(4)
+    req = {
+      type = 'SessionCreate',
+      reqId = Utils.randomString(20),
+      sessionId = wsContext.sessionId,
+    }
+  end
+
+  wsContext.ws.send(textutils.serializeJSON(req))
+
+  local res = wsContext.ws.receive(5)
+  return textutils.unserializeJSON(res)
 end
 
 ---Connect to the editor session server and request a session, retrying if needed
 ---@param wsContext table Shared WebSocket context
 ---@param maxAttempts number Max attempts to connect and get a session
----@return boolean ok True if session was acquired, false otherwise
+---@return string ccReconnectToken A reconnect token for resuming the session if it breaks
 local function connectAndRequestSession (wsContext, maxAttempts)
   local attempts = 1
 
@@ -66,7 +78,7 @@ local function connectAndRequestSession (wsContext, maxAttempts)
   end
   wsContext.ws = ws
 
-  local res, sessionId = requestSessionOnce(ws)
+  local res = requestSessionOnce(wsContext)
   while res == nil or not res.ok do
     if attempts > maxAttempts then
       LOGGER:error(
@@ -76,17 +88,18 @@ local function connectAndRequestSession (wsContext, maxAttempts)
       return false
     end
     print('Trying to create session. Attempt', attempts)
-    res, sessionId = requestSessionOnce(ws)
+    os.sleep(3)
+    res = requestSessionOnce(wsContext)
     attempts = attempts + 1
   end
 
   print()
   print('Connection to editor server successful!')
   print('Press E again to end the session.')
   print('**')
-  print('** Insert code', sessionId, 'into web editor to edit pipes.')
+  print('** Insert code', wsContext.sessionId, 'into web editor to edit pipes.')
   print('**')
-  return true
+  return res.ccReconnectToken
 end
 
 ---Queue an OS event for a given message. The event name will always be in the
@@ -120,20 +133,23 @@ local function doWebSocket (wsContext)
   print('Press E to create a factory editing session.')
   while true do
     if state == 'WAIT-FOR-USER' then
+      wsContext.reconnectToken = nil
       local event, char = os.pullEvent('char')
       if char == 'e' then
         state = 'START-CONNECT'
       end
     elseif state == 'START-CONNECT' then
-      local established = connectAndRequestSession(wsContext, 5)
-      if established then
+      local reconnectToken = connectAndRequestSession(wsContext, 5)
+      if reconnectToken then
         state = 'CONNECTED'
+        wsContext.reconnectToken = reconnectToken
       else
         print()
         print(
           'Press E to try to create a factory editing session again ' ..
           'or press Q to stop all pipes and quit.'
         )
+        wsContext.reconnectToken = nil
         wsContext.ws = nil
         state = 'WAIT-FOR-USER'
       end
@@ -142,10 +158,20 @@ local function doWebSocket (wsContext)
         local ok, res, isBinary = pcall(function () return wsContext.ws.receive() end)
         if not ok then
           print()
-          print('Lost connection to editor session server.')
-          print('Press E to try to create a factory editing session again.')
-          wsContext.ws = nil
-          state = 'WAIT-FOR-USER'
+          print('Lost connection to editor session server. Attempting to reconnect.')
+          local reconnectToken = connectAndRequestSession(wsContext, 5)
+
+          if reconnectToken then
+            wsContext.reconnectToken = reconnectToken
+          else
+            print()
+            print('Lost connection to editor session server.')
+            print('Press E to try to create a factory editing session again.')
+            wsContext.reconnectToken = nil
+            wsContext.ws = nil
+            state = 'WAIT-FOR-USER'
+          end
+
         elseif res ~= nil and not isBinary then
           queueEventFromMessage(textutils.unserializeJSON(res))
         end
@@ -158,6 +184,7 @@ local function doWebSocket (wsContext)
             'Editor session closed. Press E to create a new editing session ' ..
             'or press Q to stop all pipes and quit.'
           )
+          wsContext.reconnectToken = nil
           wsContext.ws.close()
           wsContext.ws = nil
           state = 'WAIT-FOR-USER'

server/package.json

@@ -12,6 +12,7 @@
   "author": "",
   "license": "GPL-3.0-or-later",
   "dependencies": {
+    "uuid": "^9.0.1",
     "ws": "^8.16.0"
   },
   "devDependencies": {

server/src/index.ts

@@ -1,6 +1,7 @@
 import { WebSocket, WebSocketServer } from "ws";
-import { FailResponse, IdleTimeout, MessageType, Request, SessionCreateReq, SessionJoinReq, SuccessResponse } from "./types/messages";
+import { ConfirmationResponse, FailResponse, IdleTimeout, Message, MessageType, Request, SessionCreateReq, SessionCreateRes, SessionJoinReq, SessionRejoinReq, SuccessResponse } from "./types/messages";
 import { Session, SessionId } from "./types/session";
+import { v4 as uuidv4 } from "uuid";
 
 type Role = ('CC' | 'editor');
 
@@ -39,10 +40,15 @@ wss.on("connection", function connection(ws) {
             sessionId = joinSession(request as SessionJoinReq, ws);
             if (sessionId) role = 'editor';
             break;
+
+          case "SessionRejoin":
+            sessionId = rejoinSession(message as SessionRejoinReq, ws);
+            if (sessionId) role = 'CC';
+            break;
 
           default:
             const destination = role === 'CC' ? 'editor' : 'CC';
-            if ((role === 'CC' && !session.editor) || (role === 'editor' && !session.computerCraft)) {
+            if (role === 'CC' && !session.editor) {
               const res: FailResponse = {
                 type: "ConfirmationResponse",
                 respondingTo: message.type,
@@ -52,24 +58,25 @@ wss.on("connection", function connection(ws) {
                 message: `Tried sending a message to ${destination}, but it doesn't exist on this session.`
               };
               ws.send(JSON.stringify(res));
+            } else if (role === 'editor' && !session.computerCraft) {
+              queueRequestForCCForLater(message as Request, session);
             } else {
               relayMessage(json, sessionId, destination);
             }
             break;
         }
       } else if (message.type === "CcUpdatedFactory") {
-        const destination = role === 'CC' ? 'editor' : 'CC';
-        if ((role === 'CC' && !session.editor) || (role === 'editor' && !session.computerCraft)) {
+        if (!session.editor) {
           const res: FailResponse = {
             type: "ConfirmationResponse",
             respondingTo: message.type,
             ok: false,
             error: 'PeerNotConnected',
-            message: `Tried sending a message to ${destination}, but it doesn't exist on this session.`
+            message: `Tried sending a message to the editor, but it isn't connected on this session.`
           };
           ws.send(JSON.stringify(res));
         } else {
-          relayMessage(json, sessionId, destination);
+          relayMessage(json, sessionId, 'editor');
         }
       }
     } catch (error) {
@@ -88,8 +95,7 @@ wss.on("connection", function connection(ws) {
 
   ws.on("close", (data) => {
     if (role === "CC" && sessionId && sessions[sessionId]) {
-      sessions[sessionId].editor?.close();
-      delete sessions[sessionId];
+      sessions[sessionId].computerCraft = undefined;
     } else if (role === "editor" && sessionId && sessions[sessionId]) {
       sessions[sessionId].editor = undefined;
     }
@@ -117,6 +123,8 @@ function resetIdleTimer(session: Session) {
       session.editor.send(JSON.stringify(timeoutMsg));
       session.editor.close();
     }
+
+    delete sessions[session.id];
   }, 10 * 60 * 1000)
 }
 
@@ -189,13 +197,64 @@ function joinSession({ reqId, sessionId }: SessionJoinReq, editor: WebSocket): S
   return sessionId;
 }
 
+function rejoinSession({ reqId, sessionId, ccReconnectToken }: SessionRejoinReq, computerCraft: WebSocket) {
+  if (!sessions[sessionId]) {
+    const res: FailResponse = {
+      type: "ConfirmationResponse",
+      respondingTo: "SessionRejoin",
+      ok: false,
+      error: 'SessionIdNotExist',
+      message: "Cannot connect to an expired session ID",
+      reqId: reqId,
+    };
+    computerCraft.send(JSON.stringify(res));
+    return;
+  }
+
+  if (sessions[sessionId].ccReconnectToken !== ccReconnectToken) {
+    const res: FailResponse = {
+      type: "ConfirmationResponse",
+      respondingTo: "SessionRejoin",
+      ok: false,
+      error: 'BadReconnectToken',
+      message: "Reconnect token incorrect",
+      reqId: reqId,
+    };
+    computerCraft.send(JSON.stringify(res));
+    return;
+  }
+
+  const session = sessions[sessionId];
+  session.computerCraft = computerCraft;
+
+  const res: SessionCreateRes = {
+    type: "ConfirmationResponse",
+    respondingTo: "SessionRejoin",
+    ok: true,
+    reqId: reqId,
+    ccReconnectToken: ccReconnectToken,
+  };
+  computerCraft.send(JSON.stringify(res));
+
+  if (session.staleOutboxTimerId) {
+    clearTimeout(session.staleOutboxTimerId);
+    session.staleOutboxTimerId = undefined;
+  }
+
+  while (session.editorOutbox.length > 0) {
+    computerCraft.send(JSON.stringify(session.editorOutbox.pop()));
+  }
+
+  return sessionId as SessionId;
+}
+
 /**
  * Create a session and add the ComputerCraft computer to it
  * @param param0 Session create request
  * @param computerCraft Websocket of the CC computer
  * @returns Session ID on success, undefined on failure
  */
-function createSession({ reqId, sessionId }: SessionCreateReq, computerCraft: WebSocket): SessionId {
+function createSession({ reqId, sessionId }: SessionCreateReq, computerCraft: WebSocket) {
   if (sessionId in sessions) {
     const res: FailResponse = {
       type: "ConfirmationResponse",
@@ -209,11 +268,52 @@ function createSession({ reqId, sessionId }: SessionCreateReq, computerCraft: We
     return;
   }
 
+  const ccReconnectToken = uuidv4();
+
   sessions[sessionId] = {
     id: sessionId,
-    computerCraft: computerCraft
+    computerCraft: computerCraft,
+    ccReconnectToken: ccReconnectToken,
+    editorOutbox: [],
   };
 
-  sendGenericSuccess("SessionCreate", reqId, computerCraft);
-  return sessionId;
+  const res: SessionCreateRes = {
+    type: "ConfirmationResponse",
+    respondingTo: "SessionCreate",
+    ok: true,
+    reqId: reqId,
+    ccReconnectToken: ccReconnectToken,
+  };
+  computerCraft.send(JSON.stringify(res));
+
+  return sessionId as SessionId;
+}
+
+/**
+ * Queue a request to be sent to CC for after CC reconnects to the session.
+ * If there's not a timer already, this starts a timer which clears after
+ * CC reconnects, otherwise it sends sends failure ConfirmationResponses back
+ * to the editor and closes the editor's websocket
+ * @param message Message to queue for later.
+ */
+function queueRequestForCCForLater(request: Request, session: Session) {
+  session.editorOutbox.push(request);
+
+  if (!session.staleOutboxTimerId) {
+    session.staleOutboxTimerId = setTimeout(() => {
+      const failResponse: FailResponse = {
+        type: "ConfirmationResponse",
+        respondingTo: request.type,
+        reqId: request.reqId,
+        ok: false,
+        error: 'PeerNotConnected',
+        message: 'Tried sending a message to ComputerCraft, but it did not connect within 10 seconds.'
+      };
+      session.editor.send(JSON.stringify(failResponse));
+
+      session.editorOutbox = [];
+
+      session.editor.close();
+    }, 10 * 1000);
+  }
 }

server/src/types/errors.ts

@@ -3,5 +3,6 @@ export type ErrorType = (
   'PeerNotConnected' | // CC or editor peer tried to send a message to the other peer when the other peer is not connected
   'SessionIdNotExist' | // Editor tried connecting to a nonexistent session ID
   'SessionHasEditor' | // Editor tried connecting to a session ID that already has an editor
-  'SessionIdTaken' // CC tried requesting a new session using an ID that's already taken
+  'SessionIdTaken' | // CC tried requesting a new session using an ID that's already taken
+  'BadReconnectToken' // Reconnect token supplied is incorrect
 );