keyrelay/src/control_server.h

/**
 * keyrelay -- the control server: HTTP for the page, WebSocket for the keystrokes.
 *
 * Shape of it:
 *   - one controller at a time. A second WebSocket upgrade is refused outright rather
 *     than queued or swapped silently, because two controllers typing into one machine
 *     is exactly the ambiguity a keyboard must not have. The page says so (E3).
 *     Read-only HTTP is *not* 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).
 *   - 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.
 *   - 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.
 *
 * Bytes in are decoded by kr::ws::Reader (host-tested); commands by kr::proto::parse
 * (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 "config.h"
#include "proto.h"
#include "ws_codec.h"

namespace kr {

// What the board knows about itself, as shown to the controller. The JSON field
// names in StatusJson below are the contract the page in web_page.h reads.
struct StatusSnapshot {
  const char *version = "0";
  uint32_t uptimeS = 0;
  bool usbHost = false;      // TinyUSB is attached to a target
  bool armed = false;
  bool netUp = false;
  const char *net = "down";  // "sta" | "ap" | "down"
  const char *ssid = "";
  const char *ip = "";
  uint32_t keysSent = 0;
  uint32_t chordsSent = 0;
  uint32_t dropped = 0;
  uint32_t typedChars = 0;   // characters sent by pasted text (a count, never the text)
  uint16_t typing = 0;       // characters still queued in the paste in progress
  bool rollover = false;
  // The mouse: whether it can report at all, which buttons the board believes are
  // down, and how many reports have gone out. `mouse` is false when the interface is
  // switched off on the board or when the target has forced boot protocol -- in both
  // cases the keyboard is untouched and M= frames are refused, not ignored.
  bool mouse = false;
  uint8_t mouseButtons = 0;
  uint32_t mouseFrames = 0;
  uint8_t ledState = 0;      // the target's caps/num/scroll LEDs, per its HID reports
  bool defaultToken = false;
  bool apForced = false;
};

// 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,
};

/**
 * Renders the status frame. Split out of the class so it is callable from the serial
 * console too, and so the field names stay in one place.
 */
size_t statusJson(const StatusSnapshot &s, char *out, size_t cap);

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);

  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; }

 private:
  enum class Phase : uint8_t { None, Http, WebSocket };

  void acceptNew(uint32_t nowMs);
  void serveHttp(uint32_t nowMs);
  void serveReadonly(WiFiClient &c, const ws::Request &req);
  void serveViewer(uint32_t nowMs);
  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);

  WiFiServer _server;
  WiFiClient _client;
  Phase _phase = Phase::None;

  char _http[kHttpHeaderCap];
  size_t _httpLen = 0;

  // A second socket, read-only, for when the controller slot is taken.
  WiFiClient _viewer;
  char _viewerHttp[kHttpHeaderCap];
  size_t _viewerLen = 0;
  uint32_t _viewerAtMs = 0;

  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;

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

}  // namespace kr