quicproquo/sdks/python/quicprochat/ffi.py

"""CFFI bindings to ``libquicprochat_ffi`` (the Rust C FFI layer).

This module loads the shared library and exposes a synchronous Python API
that mirrors the C functions in ``crates/quicprochat-ffi/src/lib.rs``.
"""

from __future__ import annotations

import json
import os
from pathlib import Path
from typing import Optional

import cffi

from quicprochat.types import (
    QpqError,
    AuthError,
    TimeoutError,
    ConnectionError,
)

# Status codes (must match crates/quicprochat-ffi/src/lib.rs).
QPQ_OK = 0
QPQ_ERROR = 1
QPQ_AUTH_FAILED = 2
QPQ_TIMEOUT = 3
QPQ_NOT_CONNECTED = 4

_CDEFS = """
    typedef struct QpqHandle QpqHandle;

    QpqHandle* qpq_connect(const char* server, const char* ca_cert, const char* server_name);
    int qpq_login(QpqHandle* handle, const char* username, const char* password);
    int qpq_send(QpqHandle* handle, const char* recipient, const uint8_t* message, size_t message_len);
    int qpq_receive(QpqHandle* handle, uint32_t timeout_ms, char** out_json);
    void qpq_disconnect(QpqHandle* handle);
    const char* qpq_last_error(const QpqHandle* handle);
    void qpq_free_string(char* ptr);
"""

_ffi = cffi.FFI()
_ffi.cdef(_CDEFS)
_lib: Optional[object] = None


def _load_lib() -> object:
    """Load the shared library, searching common paths."""
    global _lib
    if _lib is not None:
        return _lib

    search_paths = [
        # Explicit environment variable.
        os.environ.get("QPQ_LIB_PATH", ""),
        # Common cargo build output locations.
        str(Path(__file__).resolve().parents[3] / "target" / "release" / "libquicprochat_ffi.so"),
        str(Path(__file__).resolve().parents[3] / "target" / "debug" / "libquicprochat_ffi.so"),
        # macOS dylib.
        str(
            Path(__file__).resolve().parents[3]
            / "target"
            / "release"
            / "libquicprochat_ffi.dylib"
        ),
        str(
            Path(__file__).resolve().parents[3]
            / "target"
            / "debug"
            / "libquicprochat_ffi.dylib"
        ),
        # System library path.
        "libquicprochat_ffi.so",
    ]

    for path in search_paths:
        if not path:
            continue
        try:
            _lib = _ffi.dlopen(path)
            return _lib
        except OSError:
            continue

    raise OSError(
        "Could not find libquicprochat_ffi. Set QPQ_LIB_PATH or build with "
        "`cargo build --release -p quicprochat-ffi`."
    )


def _check_error(handle: object, code: int) -> None:
    """Raise an appropriate exception if the FFI call returned an error code."""
    if code == QPQ_OK:
        return

    lib = _load_lib()
    err_ptr = lib.qpq_last_error(handle)  # type: ignore[union-attr]
    msg = _ffi.string(err_ptr).decode("utf-8") if err_ptr != _ffi.NULL else "unknown error"

    if code == QPQ_AUTH_FAILED:
        raise AuthError(msg)
    if code == QPQ_TIMEOUT:
        raise TimeoutError(msg)
    if code == QPQ_NOT_CONNECTED:
        raise ConnectionError(msg)
    raise QpqError(msg)


class FfiTransport:
    """Synchronous transport wrapping ``libquicprochat_ffi``.

    Provides the same logical operations as ``QuicTransport`` but backed
    by the Rust client library through C FFI.

    Usage::

        ffi = FfiTransport.connect("127.0.0.1:5001", ca_cert="/path/to/ca.pem")
        ffi.login("alice", "password123")
        ffi.send("bob", b"hello")
        messages = ffi.receive(timeout_ms=5000)
        ffi.close()
    """

    def __init__(self, handle: object) -> None:
        self._handle = handle
        self._lib = _load_lib()

    @staticmethod
    def connect(
        addr: str,
        *,
        ca_cert_path: str = "",
        server_name: str = "",
    ) -> "FfiTransport":
        """Connect to a qpq server via the Rust FFI layer."""
        lib = _load_lib()

        server_c = _ffi.new("char[]", addr.encode("utf-8"))
        ca_c = _ffi.new("char[]", ca_cert_path.encode("utf-8"))

        if not server_name:
            host = addr.split(":")[0]
            server_name = host
        sn_c = _ffi.new("char[]", server_name.encode("utf-8"))

        handle = lib.qpq_connect(server_c, ca_c, sn_c)  # type: ignore[union-attr]
        if handle == _ffi.NULL:
            raise ConnectionError(f"qpq_connect failed for {addr}")

        return FfiTransport(handle)

    def login(self, username: str, password: str) -> None:
        """Authenticate with OPAQUE credentials."""
        u = _ffi.new("char[]", username.encode("utf-8"))
        p = _ffi.new("char[]", password.encode("utf-8"))
        code = self._lib.qpq_login(self._handle, u, p)  # type: ignore[union-attr]
        _check_error(self._handle, code)

    def send(self, recipient: str, message: bytes) -> None:
        """Send a message to a recipient (by username)."""
        r = _ffi.new("char[]", recipient.encode("utf-8"))
        m = _ffi.new("uint8_t[]", message)
        code = self._lib.qpq_send(self._handle, r, m, len(message))  # type: ignore[union-attr]
        _check_error(self._handle, code)

    def receive(self, timeout_ms: int = 5000) -> list[str]:
        """Receive pending messages, blocking up to *timeout_ms*.

        Returns a list of message strings (UTF-8).
        """
        out = _ffi.new("char**")
        code = self._lib.qpq_receive(self._handle, timeout_ms, out)  # type: ignore[union-attr]
        _check_error(self._handle, code)

        if out[0] == _ffi.NULL:
            return []

        json_str = _ffi.string(out[0]).decode("utf-8")
        self._lib.qpq_free_string(out[0])  # type: ignore[union-attr]
        return json.loads(json_str)  # type: ignore[no-any-return]

    def close(self) -> None:
        """Disconnect and free the handle."""
        if self._handle is not None:
            self._lib.qpq_disconnect(self._handle)  # type: ignore[union-attr]
            self._handle = None

    def __enter__(self) -> "FfiTransport":
        return self

    def __exit__(self, *args: object) -> None:
        self.close()