keyrelay/src/control_server.h

/**
 * keyrelay -- the control server: HTTP for the page, WebSocket for the keystrokes,
 * over a plaintext socket or a TLS session -- into one controller slot.
 *
 * Shape of it:
 *   - one controller at a time. Two controllers typing into one machine is exactly the
 *     ambiguity a keyboard must not have, so the second socket never *drives* one. What
 *     it may now do (phase 03) is stand in the spare slot as a **pending claim**: it
 *     completes the same handshake, authenticates with the same passphrase against the
 *     same lockout counters, and may then send nothing but `K` and the deliberate `X`
 *     -- every arm, key state, chord, text, mouse and release frame is refused by the
 *     one whitelist in claim.h and the claim is dropped. `X` hands the bridge over
 *     *through the funnel*: the incumbent is told (`E12`), closeLink() runs (which
 *     disarms and releases every key and button -- bug 10's fix, un-narrowed), and only
 *     then is the claimant adopted as the controller, DISARMED. No path reachable from
 *     `X` can call arm(); the newcomer arms with an ordinary `A1` or not at all
 *     (handOverClaim). Before phase 03 the answer was `503`/`E3` outright; read-only
 *     HTTP is *not* and was never refused: what is exclusive is the keyboard, not the
 *     board. Refusing everything meant pressing reload on your own tab failed, and a
 *     live controller could not be inspected by anything else (serveViewer).
 *   - two listeners, ONE state machine, ONE transport view. :80 and (KR_TLS) :443 both
 *     feed the same accept decision -- tls::admitConnection -- the same serveHttp() /
 *     serveWebSocket() routing and the same serveReadonly(); the only thing a socket's
 *     origin changes is which Conn it got bound to, and the only place the difference is
 *     visible at all is the plaintext policy in serveReadonly() (:80 answers a redirect or
 *     426 instead of content, compiled out entirely by -D KR_TLS=0). A socket is HTTP,
 *     WebSocket or none regardless of port, so a TLS controller and a plaintext controller
 *     can never both hold the keyboard.
 *   - non-blocking throughout: tick() runs once per loop() and consumes only what is
 *     already buffered. The keystroke path must not wait on the network stack -- under
 *     TLS that includes the handshake (Conn::poll(), one library call per tick) and every
 *     response write (queued, pumped under kPageBytesPerTick, never waited on); in
 *     plaintext it includes the write that used to be able to sit in a one-second select
 *     ten times over (see Conn::writeRaw).
 *   - nothing is trusted before the passphrase. The socket opens, must present T<token>
 *     inside kAuthTimeoutMs, and is closed otherwise; repeated failures lock the
 *     evaluation out for kTokenLockMs. TLS is transport, never an arming input: a
 *     TLS session that authenticates but does not arm types nothing.
 *
 * Bytes in are decoded by kr::ws::Reader (host-tested); commands by kr::proto::parse
 * (host-tested); the transport and the plaintext-port tables live in conn.h (its pure
 * half host-tested). What is left here is routing, framing and lifecycle -- the part that
 * genuinely needs a socket to be worth testing.
 */
#pragma once

#include <WiFi.h>

#include "claim.h"
#include "config.h"
#include "conn.h"
#include "proto.h"
#include "status.h"
#include "tls_session.h"
#include "ws_codec.h"

namespace kr {

// StatusSnapshot and statusJson() live in status.h, which includes nothing but stdio and
// config.h. That is not tidiness: it is what lets tests/test_proto.cpp render the real
// frame and measure it against kStatusJsonCap instead of taking that on trust from a
// device. Include status.h rather than re-declaring anything here.
// Error codes, named. The wire form is `E<code>:<text>`; see proto.h for the list.
enum ErrorCode : uint8_t {
  kErrNotAuthed = 1,
  kErrLockedOut = 2,
  kErrBusy = 3,
  kErrNotArmed = 4,
  kErrNoTarget = 5,
  kErrBadMessage = 6,
  kErrQueueFull = 7,
  kErrTextFull = 8,
  kErrUntypeable = 9,
  kErrMouseDisabled = 10,
  kErrMouseBootProtocol = 11,
  // "you were DROPPED", sent on the incumbent's socket just before the handover closes
  // it -- distinct from kErrBusy ("you were refused"), because the controller that was
  // never let in and the controller that was deliberately displaced deserve different
  // sentences on their screens (proto.h carries the list).
  kErrTakenOver = 12,
};

class ControlServer {
 public:
  // ctx-first function pointers rather than std::function: no allocation, and the
  // wiring in setup() reads as a list of the callbacks the server may invoke.
  using OnMessage = void (*)(void *ctx, const proto::Message &msg);
  using OnLinkChange = void (*)(void *ctx, bool linked);
  using FillStatus = void (*)(void *ctx, StatusSnapshot &out);
  using GetToken = const char *(*)(void *ctx);

  void begin(OnMessage onMessage, OnLinkChange onLinkChange, FillStatus fillStatus,
             GetToken getToken, void *ctx);
  void tick(uint32_t nowMs);

  // Push a status frame now. Called on the 1 Hz heartbeat and immediately after an
  // arm/disarm, so the page's one source of truth is the board, not its own guess.
  void sendStatus();
  void sendError(uint8_t code, const char *text);
  // The bare `OK` the protocol promises as "accepted": what `X` from the incumbent
  // controller gets (it already holds the bridge; there is nothing to do but say so).
  void sendOk();

  bool linked() const { return _phase == Phase::WebSocket && _authed; }
  /**
   * True from the first successful authentication until reboot. The silence timer needs
   * this as well as lastInboundMs(): an armed board with no controller at all (armed by
   * its own button, or by the serial console) has nothing to hear from, and must not be
   * disarmed for it -- but a board whose controller came and went must never stay armed
   * on the strength of a link that no longer exists.
   */
  bool controllerSeen() const { return _controllerSeen; }
  uint32_t lastInboundMs() const { return _lastInboundMs; }
  uint32_t rejected() const { return _rejected; }
  uint32_t authFails() const { return _authFails; }

  // Successful bridge handovers (phase 03): a count of events, never addresses, user
  // agents or anything else that would identify the other session (invariant 4 -- the
  // README's never-stored list is about exactly this temptation).
  uint32_t takeovers() const { return _takeovers; }

  // Frames the transport could not take (its queue was full) -- a periodic status frame
  // loses its slot, the next tick's drain resumes, and nothing is ever half-sent. Either
  // transport: the queue is the transport view's, so the count means one thing everywhere.
  uint32_t txDropped() const { return _txDropped; }

  // Connections refused because the slot table was full -- either listener, either
  // transport. A number, about a refusal; the refusal itself is visible to the peer
  // (TLS alert then close on :443, immediate close on :80).
  uint32_t slotRefused() const { return _slotRefused; }

#if KR_TLS
  // The identity the TLS listener serves. Must be set before begin(), and to nullptr
  // (or left unset) when there is no usable identity: begin() then does not open :443
  // at all, rather than open a listener that fails every handshake. The PemPair lives
  // in main.cpp next to the NVS code that loaded it; the server only reads it, at
  // session-open time -- copying 2 kB into every server object would be silly.
  void setIdentity(const tls::PemPair *identity) { _identity = identity; }

  // Live TLS sessions, both slots -- the heap-cost readout main.cpp logs whenever this
  // changes. A number, about numbers; nothing about content (invariant 4). Walked over
  // the STORAGE, not the slot pointers: a bridge handover swaps which object serves
  // which slot (handOverClaim), and the count means the same thing either way.
  uint32_t tlsLiveSessions() const {
    return (_tlsStore[0].alive() ? 1u : 0u) + (_tlsStore[1].alive() ? 1u : 0u);
  }
#endif

 private:
  // No separate "handshake in progress" phase: the handshake is what Conn::ready() and
  // Conn::poll() are for, and a slot in Phase::Http that is not ready yet is a slot still
  // handshaking. A phase for it would be a second place to remember the same fact, which
  // is how the plaintext and TLS paths used to drift.
  enum class Phase : uint8_t { None, Http, WebSocket };

  // What the slot does when its outbound queue finally drains: the response is finished,
  // so the link goes. Nothing else closes a link on a *write* boundary -- and a response
  // that is still draining must never be cut short, or the page arrives truncated and
  // parses as garbage.
  enum class AfterDrain : uint8_t { Nothing, CloseLink, StopViewer };

  // A response built but not yet written: the status line and headers, plus a body that
  // lives somewhere else -- a string literal in flash, the rendered status JSON, or the
  // 48 KB page in PROGMEM. Every pointer here is to storage that outlives the pumping of
  // the response (literals and kIndexHtml live forever; the JSON lives in this struct,
  // which the caller keeps until the response is queued), which is what makes "build once,
  // queue on either transport" safe.
  struct BuiltResponse {
    enum class Body : uint8_t { None, Text, Json };
    char head[kHttpResponseHeadCap];
    size_t headLen = 0;
    Body body = Body::None;
    const char *text = nullptr;   // Body::Text: a literal, or kIndexHtml in PROGMEM
    size_t textLen = 0;
    bool textInFlash = false;     // true for the page: streamed through a bounded chunk
    char json[kStatusJsonCap];    // Body::Json: rendered BEFORE buildResponse is called
    size_t jsonLen = 0;
  };

  // One response formatter for every route, every transport and both ports; the header
  // block itself is rendered by the pure kr::renderHttpResponseHead() in conn.h, which is
  // what makes its sizing host-testable.
  static void buildResponse(BuiltResponse &rsp, const char *status,
                            const char *contentType, BuiltResponse::Body body,
                            const char *text, size_t textLen,
                            const char *extra = nullptr, bool textInFlash = false);

  void acceptNew(uint32_t nowMs);
  void serveHttp(uint32_t nowMs);
  // The read-only routes, built rather than written: the routing answer must have exactly
  // one source (invariant 8) while the transport answer differs only in which Conn the
  // built response is queued on. `plaintext` is what selects the :80 policy (redirect /
  // 426) in a TLS build and is always true -- and inert -- in the KR_TLS=0 rollback.
  void serveReadonly(const ws::Request &req, BuiltResponse &rsp, bool plaintext);
  void serveViewer(uint32_t nowMs);
  // The pending claim's half of the spare slot: the WebSocket that lives there once
  // an upgrade turned Readonly into Claim (claim.h for the rules).
  void serveClaim(uint32_t nowMs);
  bool handleClaimFrame(uint32_t nowMs);
  // The bridge handover (`X`): tell the incumbent, run the release funnel, adopt the
  // claimant -- in that order and no other. Named handOver, never takeover, so a
  // reader cannot mistake it for an arming path: it is the opposite, the board lands
  // disarmed and nothing here may call arm().
  void handOverClaim(uint32_t nowMs);
  void dropClaim();
  void sendClaimText(const char *text, size_t len);
  void sendClaimError(uint8_t code, const char *text);
  void stopViewer();
  void serveWebSocket(uint32_t nowMs);
  void closeLink(const char *why);
  bool handleWsFrame(uint32_t nowMs);
  void sendPing();
  void writeTextFrame(const char *text, size_t len);
  size_t readSome(uint8_t *dst, size_t cap);

  // Queue a built response on a slot's transport. One response writer for both slots and
  // both transports -- the grep for a second HTTP response writer in src/ is part of this
  // task's completion criteria -- so `done` says what that slot does when the bytes have
  // all left, and the two callers differ only in that answer.
  bool queueResponse(Conn &conn, AfterDrain &after, AfterDrain done,
                     const BuiltResponse &rsp);
  void respondAndClose(const BuiltResponse &rsp);  // controller slot
  void respondViewer(const BuiltResponse &rsp);    // read-only slot
  // Drain both slots' queues under the per-tick budget, and run the AfterDrain decisions.
  void pumpSlots(uint32_t nowMs);
  bool spareBusy() const { return _viewOpen; }  // the read-only slot, either transport

#if KR_TLS
  void startTlsListener();
  void admitTls(uint32_t nowMs);                // accept at most one TLS socket per tick
#endif

  WiFiServer _server;
  // The controller slot: a Conn over one of these two, plus the slot's request buffer.
  // The WiFiClient and tls::Session objects are the owners; Conn is a view of one of them.
  WiFiClient _client;
  Conn _ctl;
  uint8_t _ctlTx[kConnTxCap];
  AfterDrain _ctlAfter = AfterDrain::Nothing;
  Phase _phase = Phase::None;

  char _http[kHttpHeaderCap];
  size_t _httpLen = 0;

  // A second socket, read-only, for when the controller slot is taken. Same shape as the
  // controller slot -- a Conn over whichever transport answered -- because the spare slot
  // is one slot and a second set of route handling would be a second answer to give.
  WiFiClient _viewer;
  Conn _view;
  uint8_t _viewTx[kConnViewTxCap];
  AfterDrain _viewAfter = AfterDrain::Nothing;
  bool _viewOpen = false;
  char _viewerHttp[kHttpHeaderCap];
  size_t _viewerLen = 0;
  uint32_t _viewerAtMs = 0;

  // ---- the pending claim (phase 03) ------------------------------------------
  // What the spare slot is for, and how far its request has got. The role type and
  // the frame whitelist live in claim.h (pure, host-tested) because the whitelist is
  // the whole safety argument of this slot and an argument nothing executes is one a
  // later edit can undo silently; `ViewerRole` is spelled here so the slot reads the
  // way the task that added it named it.
  //
  // What each field costs, in the static RAM this board has no PSRAM of:
  //   _viewerRole, _viewerPhase   1 byte each (uint8_t enums; they pack together)
  //   _claimAuthed                1 byte   -- the claim's OWN auth flag. The
  //                                          controller's `_authed` stays a separate
  //                                          fact on purpose: linked() is the
  //                                          controller link, and a claim must never
  //                                          make one look live
  //   _claimAtMs                  4 bytes  -- the claim's own deadline anchor; the
  //                                          controller-silence clock (_lastInboundMs)
  //                                          is not it and is never moved by a claim
  //   _claimMsg + _claimReader    64 bytes + ws::Reader's ~40 B of decode state (the
  //                                          payload IS _claimMsg; the reader owns no
  //                                          buffer of its own)
  // Total ~110 B. The claim buffer is a QUARTER of the controller's kWsMessageCap --
  // 64 against 256 -- because the widest frame a claim may ever send is `T<token>` at
  // kTokenMax + 1 = 33 bytes (config.h re-derives this); anything larger is refused by
  // the codec, not buffered, so this slot cannot be made to grow.
  using ViewerRole = claim::ViewerRole;
  ViewerRole _viewerRole = ViewerRole::Readonly;
  Phase _viewerPhase = Phase::None;
  bool _claimAuthed = false;
  uint32_t _claimAtMs = 0;
  uint8_t _claimMsg[kClaimMsgCap];
  ws::Reader _claimReader{_claimMsg, sizeof(_claimMsg)};
  // Task 02 adopts the claim's Conn on the bridge handover; until then the claim
  // travels on `_view`, the slot it authenticated on, and nothing else.

  uint8_t _wsBuf[kWsMessageCap];
  ws::Reader _reader{_wsBuf, sizeof(_wsBuf)};

  bool _authed = false;
  bool _controllerSeen = false;
  uint32_t _connectedAtMs = 0;
  uint32_t _lastInboundMs = 0;
  uint32_t _lastStatusMs = 0;
  uint32_t _lastPingMs = 0;
  uint32_t _lockUntilMs = 0;
  uint32_t _authFails = 0;
  uint32_t _rejected = 0;
  uint32_t _txDropped = 0;
  uint32_t _slotRefused = 0;
  uint32_t _takeovers = 0;   // completed handovers: counters only, never identities

#if KR_TLS
  // The encrypted listener: a raw LWIP socket rather than WiFiServer, because the
  // session must own the accepted fd -- NetworkClient closes its fd through a shared
  // handle the moment the last copy dies, which is the double-close tls::Session is
  // written to make impossible. bind/listen happen in startTlsListener().
  int _tlsListen = -1;
  const tls::PemPair *_identity = nullptr;

  // Two session objects (kTlsMaxSessions of RAM is what sized them), and two pointers
  // that say which one the controller slot uses. Normally _tlsCtl is store[0] and the
  // spare's is store[1]; a bridge handover SWAPS the pointers so the claimant's live
  // session is named by the controller slot without one crypto context having to move
  // inside another (mbedtls contexts own heap buffers and are not moveable state --
  // a pointer swap is a fact about slot naming, which is all the handover needs).
  tls::Session _tlsStore[kTlsMaxSessions];
  static_assert(kTlsMaxSessions == 2, "the slot pointers below name store[0] and store[1]");
  tls::Session *_tlsCtl = &_tlsStore[0];    // the controller slot, when _ctl.isTls()
  tls::Session *_tlsView = &_tlsStore[1];   // the spare slot,    when _view.isTls()
#endif

  OnMessage _onMessage = nullptr;
  OnLinkChange _onLinkChange = nullptr;
  FillStatus _fillStatus = nullptr;
  GetToken _getToken = nullptr;
  void *_ctx = nullptr;
};

}  // namespace kr