wijjit_ssh.limits

Resource limits: how many sessions, from whom, for how long.

Without this module a Wijjit SSH server is unbounded in every direction that matters. Any peer can open sessions until the process runs out of memory, a forgotten ssh window holds a session slot forever, and a shutdown has no way to find the sessions it needs to drain.

Design

Everything here is pure bookkeeping and policy: no sockets, no asyncssh imports, and the clock is injectable. Sessions are reached only through the ManagedSession protocol. That is what lets the real assertions in test_limits.py run as fast unit tests with a fake clock, leaving the over-SSH tests to prove only that the wiring is connected - the same split that wijjit_ssh.input’s decoder tests already use.

Two chokepoints, not one

SPEC.md §8 lists “per-IP concurrency + connect rate limit” as one bullet, but they cannot share a hook, and the difference is load-bearing:

Per-IP limits and the rate limit are pre-authentication, checked when the TCP connection arrives. The entire point is to not spend a key exchange on an abusive peer, so waiting for auth would defeat them.
``max_sessions`` is inherently post-authentication. A session only exists once a channel is opened, which requires a successful userauth.

So per-IP counts connections while the global cap counts sessions. The per-IP session count is bounded transitively, since every session lives inside a connection.

There is no locking anywhere in this module. It is correct only because asyncio is single-threaded and none of these methods await: each runs to completion before another callback can observe the state. SessionRegistry.try_admit() is one call rather than a check followed by a register for exactly this reason - it makes the atomicity structural rather than a comment that a later refactor can invalidate.

class wijjit_ssh.limits.IdleTimer(*, idle_timeout, session_timeout, on_expire)[source]

Bases: object

Closes a session that has gone quiet, or that has simply run too long.

Two independent deadlines, because they answer different questions:

idle_timeout reclaims the forgotten ssh window - reset by every byte the client sends (poke()).
session_timeout caps total duration regardless of activity. It will interrupt someone who is actively working, which is why it is off by default and why it is a separate deadline rather than a bound on the idle one.

Owns real timers, which is why it lives outside SessionRegistry - keeping the registry free of them is what makes the registry testable without a loop.

Parameters:

idle_timeout (float or None) – Seconds of silence before expiry, or None to disable.
session_timeout (float or None) – Seconds since start() before expiry, or None to disable.
on_expire (callable) – (reason: str) -> None, called with "idle_timeout" or "session_timeout". Called at most once.

Examples

>>> timer = IdleTimer(
...     idle_timeout=600.0,
...     session_timeout=None,
...     on_expire=lambda reason: session.request_close(reason),
... )
>>> timer.start()
>>> timer.poke()      # on each byte from the client
>>> timer.cancel()    # on teardown

start()[source]

Arm both deadlines. Call once, when the session begins.

Return type:: None

poke()[source]

Reset the idle deadline. Call on every byte received from the client.

Cheap by construction: this runs on every keystroke, so it does no work beyond cancelling and rescheduling one timer handle. The absolute deadline is deliberately untouched.

Return type:: None

cancel()[source]

Disarm both deadlines. Idempotent; safe after expiry.

Return type:: None

class wijjit_ssh.limits.ManagedSession(*args, **kwargs)[source]

Bases: Protocol

What SessionRegistry needs of a session.

Structural, so the registry never imports the server and the tests never need a socket.

Variables:

session_id (str) – Correlation id; see new_session_id().
peer_ip (str) – Client address, for per-IP accounting.
username (str) – Authenticated username.
started_at (float) – Monotonic timestamp of admission.

request_close(reason, message=None)[source]

Ask the session to shut down cleanly. Must be idempotent.

Parameters:

reason (str)
message (str | None)

Return type:

None

abort()[source]

Force the session down now, having declined to exit cleanly.

Return type:: None

class wijjit_ssh.limits.Rejection(reason, message)[source]

Bases: object

Why a connection or session was refused, in both registers.

Variables:

reason (str) – Stable slug for logs and metrics, e.g. "server_full".
message (str) – Human-readable text for the client. Worth writing carefully: it is the only thing a locked-out user sees, and “try again later” versus “you have too many sessions open” is the difference between a support ticket and a self-service fix.

Parameters:

reason (str)
message (str)

class wijjit_ssh.limits.SessionRegistry(*, max_sessions=100, max_per_ip=10, connect_rate=0.0, connect_burst=20, clock=time.monotonic)[source]

Bases: object

Tracks live sessions and enforces the bounds around them.

See the module docstring for why connections and sessions are counted at different chokepoints, and why nothing here locks.

Parameters:

max_sessions (int, optional) – Concurrent sessions server-wide. Default 100.
max_per_ip (int, optional) – Concurrent connections from one IP. Default 10.
connect_rate (float, optional) – Sustained connections/second/IP; 0 disables. Default 0.
connect_burst (int, optional) – Bucket capacity for connect_rate. Default 20.
clock (callable, optional) – Monotonic clock, injectable for tests.

Examples

>>> registry = SessionRegistry(max_sessions=2)
>>> registry.check_connection("10.0.0.7") is None    # allowed
True
>>> registry.connection_opened("10.0.0.7")
>>> registry.active_connections
1

check_connection(peer_ip)[source]

Decide whether to accept a new TCP connection from peer_ip.

Called before authentication, so this is cheap on purpose: it must cost far less than the key exchange it is declining to perform.

Does not record the connection - call connection_opened() for that, and only if this returned None.

Parameters:: peer_ip (str) – Client address.
Returns:: None to accept.
Return type:: Rejection or None

connection_opened(peer_ip)[source]

Record an accepted connection.

Parameters:: peer_ip (str) – Client address.
Return type:: None

connection_closed(peer_ip)[source]

Release a connection previously passed to connection_opened().

Tolerates an unknown IP: this is called from a connection_lost callback, which must never raise, and an over-release would otherwise underflow the count and permanently loosen the limit for that peer.

Parameters:: peer_ip (str) – Client address.
Return type:: None

try_admit(session)[source]

Register session if there is room, atomically.

Deliberately one call rather than a check followed by a register: on a single-threaded loop a non-awaiting method is atomic by construction, and collapsing the two makes that structural instead of a comment a later refactor could invalidate.

Parameters:: session (ManagedSession) – The session asking to start.
Returns:: None if admitted.
Return type:: Rejection or None

release(session)[source]

Deregister a session. Idempotent.

Parameters:: session (ManagedSession) – The session that has ended.
Return type:: None

property active_sessions: int: How many sessions are live right now.

property active_connections: int: How many connections are live right now, across all peers.

connections_from(peer_ip)[source]

How many connections are live from one peer.

Parameters:: peer_ip (str) – Client address.
Return type:: int

sessions()[source]

A snapshot of the live sessions.

A copy, because callers iterate it while sessions close themselves and mutate the underlying dict.

Return type:: list[ManagedSession]

async drain(*, reason, message, grace)[source]

Ask every session to end, and wait up to grace for them to.

Clean exit matters here beyond tidiness: a session that ends properly runs the app’s teardown, which leaves the alternate screen buffer and restores the client’s terminal. A session that is aborted skips that and leaves a real person with a wedged terminal. So sessions are asked first, and only killed if they will not go.

Parameters:

reason (str) – Slug recorded for each session, e.g. "server_shutdown".
message (str or None) – Text shown to each client.
grace (float) – Seconds to wait before forcing. 0 forces immediately.

Returns:

How many sessions had to be aborted. 0 means everyone left cleanly.

Return type:

int

class wijjit_ssh.limits.TokenBucket(rate, burst, *, clock=time.monotonic)[source]

Bases: object

Classic token bucket: sustained rate per second, up to burst at once.

Chosen over a fixed window because a window lets a peer make burst connections at the end of one window and burst more at the start of the next - twice the intended rate, at the worst possible moment. A bucket refills continuously, so the sustained rate holds across any interval.

Refill is computed lazily from the clock on each consume() rather than on a timer, so an idle bucket costs nothing and there is no task to cancel.

Parameters:

rate (float) – Tokens added per second. 0 disables the bucket entirely - consume() always allows. This is the default posture: see connect_rate.
burst (float) – Maximum tokens held; the bucket starts full, so a fresh peer may make burst connections immediately.
clock (callable, optional) – Returns monotonic seconds. Injectable so tests need no sleeping.

Examples

>>> bucket = TokenBucket(rate=1.0, burst=2)
>>> bucket.consume(), bucket.consume()
(True, True)
>>> bucket.consume()          # burst exhausted, refill is 1/second
False

property enabled: bool: Whether this bucket limits anything at all.

property is_full: bool

Whether the bucket has refilled to capacity.

A full bucket is indistinguishable from a freshly constructed one, which is what makes it safe to forget - see SessionRegistry.connection_closed(). Read-only: unlike consume() this does not fold the refill into the stored state.

Returns:: True when nothing is currently being throttled. Always True for a disabled bucket.
Return type:: bool

consume(amount=1.0)[source]

Take amount tokens if available.

Parameters:: amount (float, optional) – Tokens to take. Default 1.
Returns:: True if taken (the caller may proceed); False if the bucket is dry. Always True when rate is 0.
Return type:: bool