Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 8 additions & 2 deletions apps/control-plane/src/termflow_control_plane/api/terminal.py
Original file line number Diff line number Diff line change
Expand Up @@ -109,7 +109,10 @@ async def _send_terminal_events(
) -> None:
while True:
event = await terminal.next_event()
if await persisted_authentication_epoch(repositories) != auth_epoch:
router_epoch = getattr(terminal_router, "authentication_epoch", None)
if (
isinstance(router_epoch, int) and router_epoch != auth_epoch
) or await persisted_authentication_epoch(repositories) != auth_epoch:
await terminal_router.request_close(terminal, "client_closed")
with suppress(RuntimeError):
await websocket.close(code=4401, reason="Authentication epoch changed")
Expand Down Expand Up @@ -219,7 +222,10 @@ async def _receive_terminal_input(
incoming = await websocket.receive()
if incoming["type"] == "websocket.disconnect":
return
if await persisted_authentication_epoch(repositories) != auth_epoch:
router_epoch = getattr(terminal_router, "authentication_epoch", None)
if (
isinstance(router_epoch, int) and router_epoch != auth_epoch
) or await persisted_authentication_epoch(repositories) != auth_epoch:
terminal.terminate("client_closed")
with suppress(RuntimeError):
await websocket.close(code=4401, reason="Authentication epoch changed")
Expand Down
17 changes: 15 additions & 2 deletions apps/control-plane/src/termflow_control_plane/auth/epoch.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,11 +8,24 @@

from termflow_control_plane.persistence.repositories import RepositoryBundle

_EPOCH_CACHE_TTL_SECONDS = 1.0

async def persisted_authentication_epoch(repositories: RepositoryBundle) -> int:

async def persisted_authentication_epoch(
repositories: RepositoryBundle,
*,
ttl_seconds: float = _EPOCH_CACHE_TTL_SECONDS,
) -> int:
"""Finish the short DB read before propagating transport cancellation."""

lookup = asyncio.create_task(repositories.auth_state.get())
auth_state = repositories.auth_state
cached = getattr(auth_state, "cached_epoch", None)
if callable(cached):
epoch = cached(ttl_seconds=ttl_seconds)
if isinstance(epoch, int):
return epoch

lookup = asyncio.create_task(auth_state.get())
try:
state = await asyncio.shield(lookup)
except asyncio.CancelledError:
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -50,6 +50,7 @@ def __init__(
self._local_close: LocalTerminalClose | None = None
self._stream_id: UUID | None = None
self._last_seq = 0
self._resuming = False
self.close_requested = False
self.remote_closed = False
self.input_bytes = 0
Expand Down Expand Up @@ -110,10 +111,15 @@ def observe_opened(self, opened: TerminalOpenedPayload) -> None:
if self._stream_id != opened.stream_id:
self._stream_id = opened.stream_id
self._last_seq = 0
self._resuming = False

def output_state(self, output: TerminalOutputPayload) -> str:
if self._stream_id != output.stream_id:
return "gap"
if self._resuming:
if output.seq > self._last_seq + 1:
return "duplicate"
self._resuming = False
if output.seq <= self._last_seq:
return "duplicate"
return "next" if output.seq == self._last_seq + 1 else "gap"
Expand All @@ -134,6 +140,7 @@ def prepare_resume(self, stream_id: UUID, after_seq: int) -> bool:
self._queued_bytes = 0
self._available.clear()
self._last_seq = after_seq
self._resuming = True
return True


Expand All @@ -147,6 +154,10 @@ def __init__(self, *, queue_max_messages: int, queue_max_bytes: int) -> None:
self._lock = asyncio.Lock()
self._auth_epoch = 1

@property
def authentication_epoch(self) -> int:
return self._auth_epoch

async def register(
self,
instance_id: UUID,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
import json
import math
import secrets
import time
from collections.abc import Mapping
from dataclasses import dataclass, field
from datetime import UTC, datetime, timedelta
Expand Down Expand Up @@ -556,12 +557,24 @@ async def count_since(self, since: datetime) -> int:
class AuthStateRepository:
def __init__(self, sessions: async_sessionmaker[AsyncSession]) -> None:
self._sessions = sessions
self._cached_epoch: int | None = None
self._cached_epoch_at: float = 0.0

def cached_epoch(self, *, ttl_seconds: float = 1.0) -> int | None:
if (
self._cached_epoch is not None
and (time.monotonic() - self._cached_epoch_at) < ttl_seconds
):
return self._cached_epoch
return None

async def get(self) -> AuthenticationState:
async with self._sessions() as session:
state = await session.get(AuthenticationState, 1)
if state is None:
raise RuntimeError("authentication state singleton is missing")
self._cached_epoch = int(state.epoch)
self._cached_epoch_at = time.monotonic()
return state

async def configure_totp(
Expand Down Expand Up @@ -742,6 +755,8 @@ async def rotate_credentials(
)
)
await session.commit()
self._cached_epoch = int(epoch)
self._cached_epoch_at = time.monotonic()
return int(epoch)

async def reset_and_increment_epoch(
Expand Down Expand Up @@ -795,6 +810,8 @@ async def reset_and_increment_epoch(
)
)
await session.commit()
self._cached_epoch = int(epoch)
self._cached_epoch_at = time.monotonic()
return int(epoch)


Expand Down
Loading
Loading