Delay console token consumption

Author: msinhore

host-agent/firecracker-agent.py

@@ -207,7 +207,7 @@ async def console_websocket(token: str, websocket: WebSocket):
     if not CONSOLE_TOKEN_MANAGER:
         await websocket.close(code=1008)
         return
-    info = CONSOLE_TOKEN_MANAGER.consume(token)
+    info = CONSOLE_TOKEN_MANAGER.peek(token)
     if not info:
         await websocket.close(code=1008)
         return
@@ -222,6 +222,7 @@ async def console_websocket(token: str, websocket: WebSocket):
         return
 
     await websocket.accept()
+    info = CONSOLE_TOKEN_MANAGER.consume(token)
     close_event = asyncio.Event()
 
     async def ws_to_tcp():

host-agent/utils/console_tokens.py

@@ -20,26 +20,27 @@ def issue(self, vm_name: str, port: int, password: Optional[str], session_name:
             "port": int(port),
             "password": password,
             "session": session_name,
-            "expires_at": time.time() + self.ttl,
+            "issued_at": time.time(),
         }
         with self._lock:
             self._tokens[token] = entry
         return token
 
-    def consume(self, token: str) -> Optional[Dict[str, object]]:
-        if not token:
-            return None
+    def peek(self, token: str) -> Optional[Dict[str, object]]:
         with self._lock:
-            entry = self._tokens.pop(token, None)
+            entry = self._tokens.get(token)
         if not entry:
             return None
-        if entry["expires_at"] < time.time():
+        if entry.get("issued_at", 0) + self.ttl < time.time():
+            self.consume(token)
             return None
         return entry
 
-    def purge_expired(self) -> None:
-        cutoff = time.time()
+    def consume(self, token: str) -> Optional[Dict[str, object]]:
         with self._lock:
-            expired = [tok for tok, info in self._tokens.items() if info["expires_at"] < cutoff]
-            for tok in expired:
-                self._tokens.pop(tok, None)
+            entry = self._tokens.pop(token, None)
+        if not entry:
+            return None
+        if entry.get("issued_at", 0) + self.ttl < time.time():
+            return None
+        return entry