path: root/packages/meshbay-node/tests/test_webrtc_transport.py

"""
Integration test: WebRTC DataChannel transport for browser clients.

Phase 9 milestone 9.1 — spike: validate aiortc WebRTC DataChannel works
for MNP protocol exchange (handshake, index_sync, file_request, file_chunk).

Uses local loopback (no STUN/ICE needed for localhost).
"""

import asyncio
import base64
import os
import struct
import time

import blake3
import jwt
import msgpack
import pytest
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey
from aiortc import RTCPeerConnection, RTCSessionDescription

from meshbay_common import MNP_VERSION
from meshbay_common.crypto import (
    generate_gek,
    pk_to_b64,
    chunk_key as derive_chunk_key,
    decrypt_chunk,
    verify_chunk_signature,
)
from meshbay_common.protocol import MNP
from meshbay_node.indexer import DirectoryIndexer
from meshbay_node.transport.webrtc_server import WebRTCTransport


@pytest.fixture
def sk_node():
    return Ed25519PrivateKey.generate()


@pytest.fixture
def sk_hub():
    return Ed25519PrivateKey.generate()


@pytest.fixture
def gek():
    return generate_gek()


@pytest.fixture
def shared_dir(tmp_path):
    d = tmp_path / "shared"
    d.mkdir()
    (d / "test.bin").write_bytes(os.urandom(2048))
    (d / "hello.txt").write_bytes(b"hello webrtc " * 50)
    return d


def _hub_pk_pem(sk_hub):
    return sk_hub.public_key().public_bytes(
        serialization.Encoding.PEM, serialization.PublicFormat.SubjectPublicKeyInfo)


def _make_jwt(sk_hub, groups=None):
    sk_pem = sk_hub.private_bytes(
        serialization.Encoding.PEM,
        serialization.PrivateFormat.PKCS8,
        serialization.NoEncryption(),
    )
    now = int(time.time())
    return jwt.encode({
        "iss": "test-hub", "sub": "user-001",
        "pk_user": "test", "hub_id": "test-hub",
        "jti": "test-jti-webrtc", "iat": now, "exp": now + 3600,
        "groups": groups or [],
    }, sk_pem, algorithm="EdDSA")


def _pack(obj: dict) -> bytes:
    data = msgpack.packb(obj, use_bin_type=True)
    return struct.pack(">I", len(data)) + data


def _unpack(raw: bytes) -> dict:
    length = struct.unpack(">I", raw[:4])[0]
    return msgpack.unpackb(raw[4:4 + length], raw=False)


@pytest.mark.asyncio
async def test_webrtc_datachannel_handshake(sk_node, sk_hub, gek, shared_dir):
    """WebRTC DataChannel: browser sends MNP handshake, node responds with handshake_ack."""
    hub_pk_pem = _hub_pk_pem(sk_hub)
    indexer = DirectoryIndexer(root=shared_dir, group_id="g", sk_node=sk_node, gek=gek)
    await indexer.initial_scan()

    transport = WebRTCTransport(
        sk_node=sk_node, hub_pk_pem=hub_pk_pem, gek=gek,
        shared_root=shared_dir, index=indexer.index,
        stun_servers=[],
    )

    browser_pc = RTCPeerConnection()
    received = asyncio.Queue()

    channel = browser_pc.createDataChannel("mnp")

    @channel.on("message")
    def on_msg(message):
        if isinstance(message, str):
            message = message.encode()
        received.put_nowait(_unpack(message))

    offer = await browser_pc.createOffer()
    await browser_pc.setLocalDescription(offer)

    answer_sdp, ice_candidates = await transport.handle_offer(
        browser_pc.localDescription.sdp, "peer-001")

    answer = RTCSessionDescription(sdp=answer_sdp, type="answer")
    await browser_pc.setRemoteDescription(answer)

    await asyncio.sleep(0.5)

    token = _make_jwt(sk_hub)
    channel.send(_pack({
        "type": MNP.HANDSHAKE,
        "v": MNP_VERSION,
        "token": token,
    }))

    msg = await asyncio.wait_for(received.get(), timeout=5.0)
    assert msg["type"] == MNP.HANDSHAKE_ACK
    assert msg["v"] == MNP_VERSION
    assert "node_pk" in msg

    await browser_pc.close()
    await transport.close_all()


@pytest.mark.asyncio
async def test_webrtc_datachannel_file_transfer(sk_node, sk_hub, gek, shared_dir):
    """WebRTC DataChannel: full file transfer — handshake, index, fetch chunk, decrypt."""
    hub_pk_pem = _hub_pk_pem(sk_hub)
    indexer = DirectoryIndexer(root=shared_dir, group_id="g", sk_node=sk_node, gek=gek)
    await indexer.initial_scan()

    transport = WebRTCTransport(
        sk_node=sk_node, hub_pk_pem=hub_pk_pem, gek=gek,
        shared_root=shared_dir, index=indexer.index,
        stun_servers=[],
    )

    browser_pc = RTCPeerConnection()
    received = asyncio.Queue()

    channel = browser_pc.createDataChannel("mnp")

    @channel.on("open")
    def on_open():
        token = _make_jwt(sk_hub)
        channel.send(_pack({
            "type": MNP.HANDSHAKE,
            "v": MNP_VERSION,
            "token": token,
        }))

    buf = bytearray()

    @channel.on("message")
    def on_msg(message):
        if isinstance(message, str):
            message = message.encode()
        buf.extend(message)
        while len(buf) >= 4:
            length = struct.unpack(">I", buf[:4])[0]
            if len(buf) < 4 + length:
                break
            msg_bytes = bytes(buf[4:4 + length])
            del buf[:4 + length]
            received.put_nowait(msgpack.unpackb(msg_bytes, raw=False))

    offer = await browser_pc.createOffer()
    await browser_pc.setLocalDescription(offer)

    answer_sdp, _ = await transport.handle_offer(
        browser_pc.localDescription.sdp, "peer-002")
    await browser_pc.setRemoteDescription(
        RTCSessionDescription(sdp=answer_sdp, type="answer"))

    # 1) Handshake ack
    ack = await asyncio.wait_for(received.get(), timeout=5.0)
    assert ack["type"] == MNP.HANDSHAKE_ACK

    # 2) Request index
    channel.send(_pack({"type": MNP.INDEX_SYNC, "v": MNP_VERSION}))
    idx_msg = await asyncio.wait_for(received.get(), timeout=5.0)
    assert idx_msg["type"] == MNP.INDEX_SYNC
    assert "index_b64" in idx_msg

    # 3) Request file chunk
    entry = next(e for e in indexer.index.entries if e.name == "test.bin")
    channel.send(_pack({
        "type": MNP.FILE_REQUEST,
        "v": MNP_VERSION,
        "file_id": entry.id,
        "chunk_index": 0,
    }))

    chunk_msg = await asyncio.wait_for(received.get(), timeout=5.0)
    assert chunk_msg["type"] == MNP.FILE_CHUNK

    # 4) Verify and decrypt
    ct = base64.b64decode(chunk_msg["ct_b64"])
    nonce = base64.b64decode(chunk_msg["nonce_b64"])
    ct_hash = base64.b64decode(chunk_msg["ct_hash_b64"])
    pt_hash = base64.b64decode(chunk_msg["pt_hash_b64"])
    sig = base64.b64decode(chunk_msg["sig_b64"])
    file_hash = base64.b64decode(chunk_msg["file_hash_b64"])

    pk_node = sk_node.public_key()
    verify_chunk_signature(pk_node, 0, nonce, ct_hash, sig)
    assert blake3.blake3(ct).digest() == ct_hash

    ckey = derive_chunk_key(gek, file_hash, 0)
    plaintext = decrypt_chunk(ckey, nonce, ct)
    assert blake3.blake3(plaintext).digest() == pt_hash

    original = (shared_dir / "test.bin").read_bytes()
    assert plaintext == original

    await browser_pc.close()
    await transport.close_all()


@pytest.mark.asyncio
async def test_webrtc_invalid_jwt_rejected(sk_node, sk_hub, gek, shared_dir):
    """WebRTC DataChannel: invalid JWT is rejected with error."""
    hub_pk_pem = _hub_pk_pem(sk_hub)
    indexer = DirectoryIndexer(root=shared_dir, group_id="g", sk_node=sk_node, gek=gek)
    await indexer.initial_scan()

    transport = WebRTCTransport(
        sk_node=sk_node, hub_pk_pem=hub_pk_pem, gek=gek,
        shared_root=shared_dir, index=indexer.index,
        stun_servers=[],
    )

    browser_pc = RTCPeerConnection()
    received = asyncio.Queue()

    channel = browser_pc.createDataChannel("mnp")

    @channel.on("message")
    def on_msg(message):
        if isinstance(message, str):
            message = message.encode()
        received.put_nowait(_unpack(message))

    offer = await browser_pc.createOffer()
    await browser_pc.setLocalDescription(offer)

    answer_sdp, _ = await transport.handle_offer(
        browser_pc.localDescription.sdp, "peer-003")
    await browser_pc.setRemoteDescription(
        RTCSessionDescription(sdp=answer_sdp, type="answer"))

    await asyncio.sleep(0.5)

    channel.send(_pack({
        "type": MNP.HANDSHAKE,
        "v": MNP_VERSION,
        "token": "invalid.jwt.token",
    }))

    msg = await asyncio.wait_for(received.get(), timeout=5.0)
    assert msg["type"] == "error"
    assert "JWT" in msg["detail"] or "Invalid" in msg["detail"]

    await browser_pc.close()
    await transport.close_all()


@pytest.mark.asyncio
async def test_webrtc_request_before_handshake_rejected(sk_node, sk_hub, gek, shared_dir):
    """WebRTC DataChannel: request without handshake is rejected."""
    hub_pk_pem = _hub_pk_pem(sk_hub)
    indexer = DirectoryIndexer(root=shared_dir, group_id="g", sk_node=sk_node, gek=gek)
    await indexer.initial_scan()

    transport = WebRTCTransport(
        sk_node=sk_node, hub_pk_pem=hub_pk_pem, gek=gek,
        shared_root=shared_dir, index=indexer.index,
        stun_servers=[],
    )

    browser_pc = RTCPeerConnection()
    received = asyncio.Queue()

    channel = browser_pc.createDataChannel("mnp")

    @channel.on("message")
    def on_msg(message):
        if isinstance(message, str):
            message = message.encode()
        received.put_nowait(_unpack(message))

    offer = await browser_pc.createOffer()
    await browser_pc.setLocalDescription(offer)

    answer_sdp, _ = await transport.handle_offer(
        browser_pc.localDescription.sdp, "peer-004")
    await browser_pc.setRemoteDescription(
        RTCSessionDescription(sdp=answer_sdp, type="answer"))

    await asyncio.sleep(0.5)

    channel.send(_pack({"type": MNP.INDEX_SYNC, "v": MNP_VERSION}))

    msg = await asyncio.wait_for(received.get(), timeout=5.0)
    assert msg["type"] == "error"
    assert "Handshake required" in msg["detail"]

    await browser_pc.close()
    await transport.close_all()