plezy/scripts/codegen/generate_relay_protocol.py

#!/usr/bin/env python3
"""Generate Dart and Go relay protocol constants from relay_protocol.json."""

from __future__ import annotations

import json
from pathlib import Path

ROOT = Path(__file__).resolve().parents[2]
SPEC_PATH = ROOT / "relay_protocol.json"
DART_PATH = ROOT / "lib/watch_together/services/relay_protocol.g.dart"
GO_PATH = ROOT / "server/relay_protocol_gen.go"

SUPPORTED_ID_PATTERN = r"^[A-Za-z0-9_-]+$"


def camel_to_pascal(value: str) -> str:
    return value[:1].upper() + value[1:]


def validated_id_pattern(spec: dict) -> str:
    try:
        pattern = spec["idPattern"]
    except KeyError:
        raise ValueError("idPattern is required") from None
    if not isinstance(pattern, str):
        raise ValueError("idPattern must be a string")
    if pattern != SUPPORTED_ID_PATTERN:
        raise ValueError(
            f"unsupported idPattern {pattern!r}; expected {SUPPORTED_ID_PATTERN!r}"
        )
    return pattern


def validated_reconnect_token_bytes(spec: dict) -> int:
    try:
        size = spec["reconnectTokenBytes"]
    except KeyError:
        raise ValueError("reconnectTokenBytes is required") from None
    if isinstance(size, bool) or not isinstance(size, int) or size <= 0:
        raise ValueError("reconnectTokenBytes must be a positive integer")
    return size


def reconnect_token_length(token_bytes: int) -> int:
    """Length of the unpadded base64url encoding of ``token_bytes`` bytes."""
    return (token_bytes * 4 + 2) // 3


def dart_source(spec: dict) -> str:
    id_pattern = validated_id_pattern(spec)
    token_bytes = validated_reconnect_token_bytes(spec)
    token_length = reconnect_token_length(token_bytes)
    lines = [
        "// Generated by scripts/codegen/generate_relay_protocol.py. Do not edit.",
        "",
        "abstract final class RelayProtocol {",
    ]
    lines.extend(
        [
            f"  static const int protocolVersion = {spec['protocolVersion']};",
            f"  static const int legacyProtocolVersion = {spec['legacyProtocolVersion']};",
            "",
        ]
    )
    for group in ("clientMessageTypes", "serverMessageTypes"):
        for name, value in spec[group].items():
            lines.append(f"  static const String {name} = {value!r};")
    for name, value in spec["errorCodes"].items():
        lines.append(f"  static const String {name}Code = {value!r};")
    lines.append("")
    for name, value in spec["limits"].items():
        lines.append(f"  static const int {name} = {value};")
    lines.extend(
        [
            "",
            f"  static const int reconnectTokenBytes = {token_bytes};",
            f"  static const int reconnectTokenLength = {token_length};",
            "",
            f"  static final RegExp _idPattern = RegExp(r{id_pattern!r});",
            f"  static final RegExp _reconnectTokenPattern = RegExp(r'^[A-Za-z0-9_-]{{{token_length}}}$');",
            "",
            "  static bool isValidSessionId(String value) =>",
            "      value.isNotEmpty && value.length <= maxSessionIdLength && _idPattern.hasMatch(value);",
            "",
            "  static bool isValidPeerId(String value) =>",
            "      value.isNotEmpty && value.length <= maxPeerIdLength && _idPattern.hasMatch(value);",
            "",
            "  static bool isValidReconnectToken(String value) => _reconnectTokenPattern.hasMatch(value);",
            "}",
            "",
        ]
    )
    return "\n".join(lines)


def go_source(spec: dict) -> str:
    validated_id_pattern(spec)
    token_bytes = validated_reconnect_token_bytes(spec)
    lines = [
        "// Code generated by scripts/codegen/generate_relay_protocol.py. DO NOT EDIT.",
        "",
        "package main",
        "",
        "const (",
    ]
    lines.extend(
        [
            f"\trelayProtocolVersion       = {spec['protocolVersion']}",
            f"\tlegacyRelayProtocolVersion = {spec['legacyProtocolVersion']}",
            "",
        ]
    )
    protocol_constants = []
    for group in ("clientMessageTypes", "serverMessageTypes"):
        protocol_constants.extend(
            (f"relayType{camel_to_pascal(name)}", f'"{value}"')
            for name, value in spec[group].items()
        )
    protocol_constants.extend(
        (f"relayError{camel_to_pascal(name)}", f'"{value}"')
        for name, value in spec["errorCodes"].items()
    )
    protocol_name_width = max(len(name) for name, _ in protocol_constants)
    lines.extend(
        f"\t{name:<{protocol_name_width}} = {value}"
        for name, value in protocol_constants
    )
    lines.append("")
    go_limit_names = {
        "maxRoomSize": "maxRoomSize",
        "maxMessageSize": "maxMessageSize",
        "maxSessionIdLength": "maxSessionIDLength",
        "maxPeerIdLength": "maxPeerIDLength",
    }
    limit_constants = [
        (go_limit_names[name], str(value)) for name, value in spec["limits"].items()
    ]
    limit_constants.append(("reconnectTokenSize", str(token_bytes)))
    limit_name_width = max(len(name) for name, _ in limit_constants)
    lines.extend(
        f"\t{name:<{limit_name_width}} = {value}"
        for name, value in limit_constants
    )
    lines.extend(
        [
            ")",
            "",
            "func supportedRelayProtocolVersion(version int) bool {",
            "\treturn version == legacyRelayProtocolVersion || version == relayProtocolVersion",
            "}",
            "",
            "func validRelayID(value string, maxLength int) bool {",
            "\tif len(value) == 0 || len(value) > maxLength {",
            "\t\treturn false",
            "\t}",
            "\tfor _, ch := range value {",
            "\t\tif (ch < 'a' || ch > 'z') && (ch < 'A' || ch > 'Z') &&",
            "\t\t\t(ch < '0' || ch > '9') && ch != '_' && ch != '-' {",
            "\t\t\treturn false",
            "\t\t}",
            "\t}",
            "\treturn true",
            "}",
            "",
        ]
    )
    return "\n".join(lines)


def main() -> None:
    spec = json.loads(SPEC_PATH.read_text(encoding="utf-8"))
    dart_output = dart_source(spec)
    go_output = go_source(spec)
    DART_PATH.write_text(dart_output, encoding="utf-8", newline="\n")
    GO_PATH.write_text(go_output, encoding="utf-8")


if __name__ == "__main__":
    main()