diff --git a/packages/syft-rds/src/syft_rds/client.py b/packages/syft-rds/src/syft_rds/client.py index cb112eb5011..325b6bbe46f 100644 --- a/packages/syft-rds/src/syft_rds/client.py +++ b/packages/syft-rds/src/syft_rds/client.py @@ -229,6 +229,17 @@ def pair_with_mock_drive_service_connection( ) return cls._build_rds_pair_from_managers(ds_mgr, do_mgr) + @classmethod + def pair_with_mock_drive_and_mock_webrtc_connection( + cls, **kwargs: Any + ) -> tuple["SyftRDSClient", "SyftRDSClient"]: + """(ds, do) pair like ``pair_with_mock_drive_service_connection``, with + messages over WebRTC joined by a fake channel.""" + ds_mgr, do_mgr = SyftboxManager.pair_with_mock_drive_and_mock_webrtc_connection( + collection_specs=DATASET_COLLECTION_SPECS, **kwargs + ) + return cls._build_rds_pair_from_managers(ds_mgr, do_mgr) + @classmethod def _pair_with_google_drive_testing_connection( cls, **kwargs: Any diff --git a/syft/sync/connections/drive/drive_p2p.py b/syft/sync/connections/drive/drive_p2p.py index 50c3418bc99..825d5e2d4ac 100644 --- a/syft/sync/connections/drive/drive_p2p.py +++ b/syft/sync/connections/drive/drive_p2p.py @@ -82,6 +82,8 @@ class DriveP2PConnection(P2PConnection): methods below. """ + keeps_sent_messages = True + def __init__(self, cached_client: DriveCachedClient): self._cached_client = cached_client diff --git a/syft/sync/connections/protocols.py b/syft/sync/connections/protocols.py index 2df88176a38..6b744386c54 100644 --- a/syft/sync/connections/protocols.py +++ b/syft/sync/connections/protocols.py @@ -55,6 +55,10 @@ class P2PConnection(RouterConnection, ABC): is sent back to the peer. """ + # Whether sent messages stay where the peer reads them, as in Drive's + # outbox, so the owner compacts them. A live link drops them on receipt. + keeps_sent_messages: bool = False + @abstractmethod def send( self, peer: str, msg_id: str, data: bytes, kind: MessageKind diff --git a/syft/sync/connections/webrtc/__init__.py b/syft/sync/connections/webrtc/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/syft/sync/connections/webrtc/channel.py b/syft/sync/connections/webrtc/channel.py new file mode 100644 index 00000000000..d90aa327e0d --- /dev/null +++ b/syft/sync/connections/webrtc/channel.py @@ -0,0 +1,13 @@ +"""The pipe under the WebRTC core: one per peer, carrying frames as bytes.""" + +from abc import ABC, abstractmethod + + +class Channel(ABC): + """One ordered, reliable pipe to one peer. An aiortc data channel in + production, a ``FakeChannel`` in tests. It hands each frame that arrives + to the core's ``on_frame``.""" + + @abstractmethod + def send(self, frame: bytes) -> None: + """Send ``frame`` to the peer.""" diff --git a/syft/sync/connections/webrtc/fake_channel.py b/syft/sync/connections/webrtc/fake_channel.py new file mode 100644 index 00000000000..dd43669b81b --- /dev/null +++ b/syft/sync/connections/webrtc/fake_channel.py @@ -0,0 +1,26 @@ +"""An in-memory channel between two WebRTC cores, for tests.""" + +from syft.sync.connections.webrtc.channel import Channel +from syft.sync.connections.webrtc.webrtc_p2p import WebRTCP2PConnection + + +class FakeChannel(Channel): + """One end of an in-memory pipe, in place of an aiortc data channel. A + frame reaches the other core inside ``send``: no threads, no network.""" + + def __init__(self, sender: str, receiver: WebRTCP2PConnection): + self.sender = sender + self.receiver = receiver + + def send(self, frame: bytes) -> None: + self.receiver.on_frame(self.sender, frame) + + @staticmethod + def connect(a: WebRTCP2PConnection, b: WebRTCP2PConnection) -> None: + """Join ``a`` and ``b`` and open the channel. Both ends are attached + before it opens, as with a real channel, so the first receipts have + a way back.""" + a.attach(b.email, FakeChannel(sender=a.email, receiver=b)) + b.attach(a.email, FakeChannel(sender=b.email, receiver=a)) + a.on_open(b.email) + b.on_open(a.email) diff --git a/syft/sync/connections/webrtc/message_queue.py b/syft/sync/connections/webrtc/message_queue.py new file mode 100644 index 00000000000..9ac1d6bb6d7 --- /dev/null +++ b/syft/sync/connections/webrtc/message_queue.py @@ -0,0 +1,91 @@ +"""A queue of messages kept on disk, one file per message, so it survives a +restart and adding or removing one message never rewrites the others.""" + +from collections.abc import Collection +from pathlib import Path +from typing import NamedTuple + +from pydantic import BaseModel, ConfigDict + +from syft.sync.connections.types import MessageKind +from syft.sync.sync.caches.persisted_dict import locked_file + +# Bytes travel as base64 inside a JSON frame. +BYTES_AS_BASE64 = ConfigDict(ser_json_bytes="base64", val_json_bytes="base64") + + +class QueuedMessage(BaseModel): + """One message, as the sender's router encoded it.""" + + model_config = ConfigDict(frozen=True, **BYTES_AS_BASE64) + + msg_id: str + kind: MessageKind + data: bytes + + +class _Entry(NamedTuple): + """One message file, named ``..``.""" + + position: int + kind: MessageKind + msg_id: str + path: Path + + +class MessageQueue: + """Messages in one folder, oldest first. Every change holds the folder lock.""" + + def __init__(self, folder: Path): + self.folder = folder + + def messages( + self, kind: MessageKind | None = None, max_messages: int | None = None + ) -> list[QueuedMessage]: + """The messages, of ``kind`` when given, oldest first, at most + ``max_messages``. Only those files are read.""" + with locked_file(self.folder, exclusive=False): + entries = [e for e in self._entries() if kind is None or e.kind == kind] + return [ + QueuedMessage(msg_id=e.msg_id, kind=e.kind, data=e.path.read_bytes()) + for e in entries[:max_messages] + ] + + def add(self, message: QueuedMessage) -> None: + """Append ``message``, unless one with its ``msg_id`` is already here.""" + _require_file_name(message.msg_id) + with locked_file(self.folder): + entries = self._entries() + if any(e.msg_id == message.msg_id for e in entries): + return + position = entries[-1].position + 1 if entries else 0 + name = f"{position:012d}.{message.kind.value}.{message.msg_id}" + # Written under a hidden name first, so a reader never sees half. + tmp = self.folder / f".{name}.tmp" + tmp.write_bytes(message.data) + tmp.replace(self.folder / name) + + def remove(self, msg_ids: Collection[str]) -> list[str]: + """Remove the messages with these ids; return the ids that were here.""" + with locked_file(self.folder): + removed = [e for e in self._entries() if e.msg_id in msg_ids] + for entry in removed: + entry.path.unlink() + return [e.msg_id for e in removed] + + def _entries(self) -> list[_Entry]: + self.folder.mkdir(parents=True, exist_ok=True) + entries = [] + for path in self.folder.iterdir(): + if path.name.startswith("."): + continue + position, kind, msg_id = path.name.split(".", 2) + entries.append(_Entry(int(position), MessageKind(kind), msg_id, path)) + return sorted(entries, key=lambda e: e.position) + + +def _require_file_name(msg_id: str) -> None: + """``msg_id`` comes from the peer and names a file, so it may not leave the + queue folder or hide from it.""" + if msg_id != Path(msg_id).name or msg_id.startswith("."): + raise ValueError(f"msg_id {msg_id!r} is not a plain file name") diff --git a/syft/sync/connections/webrtc/webrtc_p2p.py b/syft/sync/connections/webrtc/webrtc_p2p.py new file mode 100644 index 00000000000..0d9aa1d17f1 --- /dev/null +++ b/syft/sync/connections/webrtc/webrtc_p2p.py @@ -0,0 +1,131 @@ +"""Messages between peers over WebRTC: the synchronous core. + +The core keeps the queues on disk, sends receipts and resends what the peer has +not confirmed. The ``Channel`` under it only carries frames: an aiortc data +channel in production, a ``FakeChannel`` in tests. +""" + +from collections import defaultdict +from collections.abc import Container +from pathlib import Path + +from pydantic import BaseModel, ConfigDict + +from syft.sync.connections.protocols import P2PConnection +from syft.sync.connections.types import MessageKind, RawMessage, SendStatus +from syft.sync.connections.webrtc.channel import Channel +from syft.sync.connections.webrtc.message_queue import ( + BYTES_AS_BASE64, + MessageQueue, + QueuedMessage, +) +from syft.sync.sync.constants import CACHE_DIR + +WEBRTC_QUEUES_DIR = "webrtc" + + +class Frame(BaseModel): + """What travels over a channel: a message, or the receipt for one.""" + + model_config = ConfigDict(frozen=True, **BYTES_AS_BASE64) + + message: QueuedMessage | None = None + # The msg_id of a message the peer stored. + receipt: str | None = None + + +class WebRTCP2PConnection(P2PConnection): + """Messages over one channel per peer, with a queue per peer and direction. + + ``send`` stores the message in the peer's outgoing queue before sending it, + so a restart loses nothing, and returns ``QUEUED``. The peer stores it in + its incoming queue and sends a receipt back. The message then leaves the + outgoing queue and ``delivered()`` reports it. + """ + + def __init__(self, email: str, syftbox_folder: Path): + self.email = email + self._folder = Path(syftbox_folder) / CACHE_DIR / WEBRTC_QUEUES_DIR + # peer email -> its channel + self._channels: dict[str, Channel] = {} + # Peers whose channel is open and got every older message again. + self._open: set[str] = set() + self._delivered: list[str] = [] + + def attach(self, peer: str, channel: Channel) -> None: + """Use ``channel`` for ``peer`` from now on. Receipts go over it right + away; messages once it opens.""" + self._channels[peer] = channel + + def on_open(self, peer: str) -> None: + """``peer``'s channel opened: send it every message it has not + confirmed, oldest first. New messages follow them, so none overtakes + an older one.""" + for message in self._outgoing(peer).messages(): + self._send_frame(peer, Frame(message=message)) + self._open.add(peer) + + def send( + self, peer: str, msg_id: str, data: bytes, kind: MessageKind + ) -> SendStatus: + message = QueuedMessage(msg_id=msg_id, kind=kind, data=data) + # Stored first, and not under the folder lock while it travels: a + # FakeChannel brings the receipt back inside this call. + self._outgoing(peer).add(message) + if peer in self._open: + self._send_frame(peer, Frame(message=message)) + return SendStatus.QUEUED + + def on_frame(self, peer: str, data: bytes) -> None: + """Called by ``peer``'s channel with each frame that arrives.""" + frame = Frame.model_validate_json(data) + if frame.message is not None: + self._incoming(peer).add(frame.message) + self._send_frame(peer, Frame(receipt=frame.message.msg_id)) + elif frame.receipt is not None: + self._delivered += self._outgoing(peer).remove({frame.receipt}) + + def delivered(self) -> list[str]: + delivered, self._delivered = self._delivered, [] + return delivered + + def receive( + self, + peer: str, + kind: MessageKind, + after: str | None = None, + max_messages: int | None = None, + already_received: Container[str] = frozenset(), + ) -> list[RawMessage]: + """The messages of ``kind`` from ``peer`` in the incoming queue, in the + order they arrived. ``after`` and ``already_received`` are not needed: + the queue holds only messages not processed yet.""" + return [ + RawMessage(sender=peer, msg_id=m.msg_id, kind=m.kind, data=m.data) + for m in self._incoming(peer).messages(kind, max_messages) + ] + + def mark_processed(self, message: RawMessage) -> None: + self.mark_all_processed([message]) + + def mark_all_processed(self, messages: list[RawMessage]) -> None: + """Remove ``messages`` from the incoming queues, one lock per sender.""" + # sender email -> msg_ids processed + by_sender: dict[str, set[str]] = defaultdict(set) + for message in messages: + by_sender[message.sender].add(message.msg_id) + for sender, msg_ids in by_sender.items(): + self._incoming(sender).remove(msg_ids) + + def _send_frame(self, peer: str, frame: Frame) -> None: + """Send ``frame`` if ``peer`` has a channel. A message without one + waits in the outgoing queue until ``on_open``.""" + channel = self._channels.get(peer) + if channel is not None: + channel.send(frame.model_dump_json().encode()) + + def _outgoing(self, peer: str) -> MessageQueue: + return MessageQueue(self._folder / "outgoing" / peer) + + def _incoming(self, peer: str) -> MessageQueue: + return MessageQueue(self._folder / "incoming" / peer) diff --git a/syft/sync/syftbox_manager.py b/syft/sync/syftbox_manager.py index a0322bbc11b..46daa2db6f1 100644 --- a/syft/sync/syftbox_manager.py +++ b/syft/sync/syftbox_manager.py @@ -63,6 +63,8 @@ from syft.sync.connections.config import ConnectionRouterConfig from syft.sync.platforms.gdrive_files_platform import GdriveFilesPlatform from syft.sync.connections.protocols import Connections +from syft.sync.connections.webrtc.fake_channel import FakeChannel +from syft.sync.connections.webrtc.webrtc_p2p import WebRTCP2PConnection from syft.sync.version.peer_manager import ( PeerManager, PeerManagerConfig, @@ -804,70 +806,180 @@ def pair_with_mock_drive_service_connection( Returns: Tuple of (ds_manager, do_manager) """ + do_config, ds_config = SyftboxManager._mock_pair_configs( + email1=email1, + email2=email2, + base_path1=base_path1, + base_path2=base_path2, + use_in_memory_cache=use_in_memory_cache, + persist_owner_state=persist_owner_state, + check_versions=check_versions, + encryption=encryption, + collection_specs=collection_specs, + ds_encryption_policy=ds_encryption_policy, + do_encryption_policy=do_encryption_policy, + ) + do_connections, ds_connections = SyftboxManager._mock_drive_connections( + do_config, ds_config + ) + return cls._test_pair( + do_config, do_connections, ds_config, ds_connections, add_peers + ) + + @classmethod + def pair_with_mock_drive_and_mock_webrtc_connection( + cls, + email1: str | None = None, + email2: str | None = None, + base_path1: str | None = None, + base_path2: str | None = None, + sync_automatically: bool = False, + add_peers: bool = True, + use_in_memory_cache: bool = True, + persist_owner_state: bool = True, + check_versions: bool = False, + encryption: bool = False, + collection_specs: list["CollectionSyncSpec"] | None = None, + ds_encryption_policy: EncryptionPolicy | None = None, + do_encryption_policy: EncryptionPolicy | None = None, + ): + """Like ``pair_with_mock_drive_service_connection``, with the same + arguments, but messages go over the real WebRTC core on both sides, + joined by a ``FakeChannel``. The mock Drive still does peering, datasets + and storage.""" + do_config, ds_config = SyftboxManager._mock_pair_configs( + email1=email1, + email2=email2, + base_path1=base_path1, + base_path2=base_path2, + use_in_memory_cache=use_in_memory_cache, + persist_owner_state=persist_owner_state, + check_versions=check_versions, + encryption=encryption, + collection_specs=collection_specs, + ds_encryption_policy=ds_encryption_policy, + do_encryption_policy=do_encryption_policy, + ) + do_connections, ds_connections = SyftboxManager._mock_drive_connections( + do_config, ds_config + ) + do_connections, ds_connections = SyftboxManager._with_fake_webrtc_p2p( + do_config, ds_config, do_connections, ds_connections + ) + return cls._test_pair( + do_config, do_connections, ds_config, ds_connections, add_peers + ) + + @staticmethod + def _mock_pair_configs( + email1: str | None = None, + email2: str | None = None, + base_path1: str | None = None, + base_path2: str | None = None, + use_in_memory_cache: bool = True, + persist_owner_state: bool = True, + check_versions: bool = False, + encryption: bool = False, + collection_specs: list["CollectionSyncSpec"] | None = None, + ds_encryption_policy: EncryptionPolicy | None = None, + do_encryption_policy: EncryptionPolicy | None = None, + ) -> tuple[SyftboxManagerConfig, SyftboxManagerConfig]: + """The configs of a DO and a DS for a test pair.""" default_policy = EncryptionPolicy.from_bool(encryption) - # Create configs using the existing base config generator + shared = dict( + use_in_memory_cache=use_in_memory_cache, + persist_owner_state=persist_owner_state, + check_versions=check_versions, + collection_specs=collection_specs, + ) do_config = SyftboxManagerConfig._base_config_for_testing( email=email1, syftbox_folder=base_path1, has_ds_role=False, has_do_role=True, - use_in_memory_cache=use_in_memory_cache, - persist_owner_state=persist_owner_state, - check_versions=check_versions, - collection_specs=collection_specs, encryption_policy=do_encryption_policy or default_policy, + **shared, ) - ds_config = SyftboxManagerConfig._base_config_for_testing( email=email2, syftbox_folder=base_path2, has_ds_role=True, has_do_role=False, - use_in_memory_cache=use_in_memory_cache, - persist_owner_state=persist_owner_state, - check_versions=check_versions, - collection_specs=collection_specs, encryption_policy=ds_encryption_policy or default_policy, + **shared, ) + return do_config, ds_config + @staticmethod + def _mock_drive_connections( + do_config: SyftboxManagerConfig, ds_config: SyftboxManagerConfig + ) -> tuple[Connections, Connections]: + """The connections of both managers over one mock Drive.""" do_cached_client, ds_cached_client = mock_drive_service.pair_with_mock_service( do_config.email, ds_config.email ) - do_manager = cls.from_config( - do_config, SyftboxManager._drive_connections(do_cached_client) + return ( + SyftboxManager._drive_connections(do_cached_client), + SyftboxManager._drive_connections(ds_cached_client), ) - ds_manager = cls.from_config( - ds_config, SyftboxManager._drive_connections(ds_cached_client) + + @staticmethod + def _with_fake_webrtc_p2p( + do_config: SyftboxManagerConfig, + ds_config: SyftboxManagerConfig, + do_connections: Connections, + ds_connections: Connections, + ) -> tuple[Connections, Connections]: + """Both connections with their p2p replaced by a WebRTC core, the two + cores joined by a ``FakeChannel``.""" + do_p2p = WebRTCP2PConnection(do_config.email, do_config.syftbox_folder) + ds_p2p = WebRTCP2PConnection(ds_config.email, ds_config.syftbox_folder) + FakeChannel.connect(do_p2p, ds_p2p) + return ( + do_connections.model_copy(update={"p2p_con": do_p2p}), + ds_connections.model_copy(update={"p2p_con": ds_p2p}), ) - # Set up callbacks for DS -> DO communication + @classmethod + def _test_pair( + cls, + do_config: SyftboxManagerConfig, + do_connections: Connections, + ds_config: SyftboxManagerConfig, + ds_connections: Connections, + add_peers: bool, + ) -> tuple["SyftboxManager", "SyftboxManager"]: + """(ds_manager, do_manager) built on these connections, with the + callbacks a test pair needs and both version files published.""" + do_manager = cls.from_config(do_config, do_connections) + ds_manager = cls.from_config(ds_config, ds_connections) + # DS -> DO: a written file is proposed to the DO ds_manager.file_writer.add_callback( - "write_file", - ds_manager.datasite_watcher_syncer.on_file_change, + "write_file", ds_manager.datasite_watcher_syncer.on_file_change ) - - # Set up callback for DO job handling + # The DO handles job files as they are written do_manager.datasite_owner_syncer.event_cache.add_callback( "on_event_local_write", do_manager.job_file_change_handler._handle_file_change, ) - - # Write version files ds_manager.peer_manager.write_own_peer_info() do_manager.peer_manager.write_own_peer_info() - if add_peers: - # DS creates peer request - ds_manager.add_peer(do_manager.email, sync=False) - # DO approves the peer request - do_manager.load_peers() - do_manager.approve_peer_request(ds_manager.email) - # DS refreshes peer state so it sees DO's accepted status and - # DO's advertised VersionInfo (incl. syft_client_install_source). - ds_manager.load_peers() - + SyftboxManager._approve_test_peers(ds_manager, do_manager) return ds_manager, do_manager + @staticmethod + def _approve_test_peers( + ds_manager: "SyftboxManager", do_manager: "SyftboxManager" + ) -> None: + """The DS asks the DO to peer, and the DO approves.""" + ds_manager.add_peer(do_manager.email, sync=False) + do_manager.load_peers() + do_manager.approve_peer_request(ds_manager.email) + # The DS sees the DO's accepted state and its advertised VersionInfo + # (incl. syft_client_install_source). + ds_manager.load_peers() + def add_peer( self, peer_email: str, diff --git a/syft/sync/sync/datasite_owner_syncer.py b/syft/sync/sync/datasite_owner_syncer.py index e33f6b869f9..3ffad11d4e7 100644 --- a/syft/sync/sync/datasite_owner_syncer.py +++ b/syft/sync/sync/datasite_owner_syncer.py @@ -979,9 +979,12 @@ def compact_outbox_if_needed( and deletes the originals. Returns the number of source messages compacted (0 if below threshold, - if no outbox folder, or if peer encryption is on — DO can't read - peer-encrypted bytes back to merge). + if no outbox folder, if the p2p connection keeps no sent messages, or + if peer encryption is on — DO can't read peer-encrypted bytes back to + merge). """ + if not self.connection_router.connections.p2p_con.keeps_sent_messages: + return 0 peer_store = self.connection_router.peer_store if peer_store.pair_encrypts(peer_store.get_cached_peer(recipient_email)): logger.debug( diff --git a/tests/unit/webrtc/__init__.py b/tests/unit/webrtc/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/webrtc/test_webrtc_end_to_end.py b/tests/unit/webrtc/test_webrtc_end_to_end.py new file mode 100644 index 00000000000..1c979843236 --- /dev/null +++ b/tests/unit/webrtc/test_webrtc_end_to_end.py @@ -0,0 +1,72 @@ +"""Managers that send their messages over WebRTC joined by a FakeChannel.""" + +from pathlib import Path + +from syft_rds import SyftRDSClient + +from syft.sync.connections.drive.drive_connections import drive_cached_client_of +from syft.sync.syftbox_manager import SyftboxManager +from tests.unit.test_sync_manager import path_for_job +from tests.unit.utils import grant_job_inbox_access + + +def drive_message_files(do_manager: SyftboxManager, ds_email: str) -> list[dict]: + """The files in the Drive inbox, its archive and the outbox between the two + peers. Over Drive a proposal lands in the inbox and is then archived.""" + client = drive_cached_client_of(do_manager._connection_router.connections) + folder_ids = [ + client.get_own_datasite_inbox_id(ds_email), + client.get_archive_folder_id_from_drive(ds_email), + client.get_own_datasite_outbox_id(ds_email), + ] + return [ + f + for folder_id in folder_ids + if folder_id is not None + for f in client.get_file_metadatas_from_folder(folder_id) + ] + + +def test_file_round_trip_goes_over_webrtc(): + ds_manager, do_manager = ( + SyftboxManager.pair_with_mock_drive_and_mock_webrtc_connection( + use_in_memory_cache=False + ) + ) + grant_job_inbox_access(do_manager, ds_manager.email) + job_path = path_for_job(do_manager.email, ds_manager.email) + + ds_manager._send_file_change(job_path, "Hello, world!") + do_manager.sync() + + assert (do_manager.syftbox_folder / job_path).read_text() == "Hello, world!" + + datasite_do = do_manager.syftbox_folder / do_manager.email + do_manager.datasite_owner_syncer.perm_context.open("results/").grant_read_access( + ds_manager.email + ) + (datasite_do / "results" / "result.txt").write_text("I am a result") + do_manager.sync() + ds_manager.sync() + + result_on_ds = ds_manager.syftbox_folder / do_manager.email / "results/result.txt" + assert result_on_ds.read_text() == "I am a result" + assert drive_message_files(do_manager, ds_manager.email) == [] + + +def test_job_round_trip_goes_over_webrtc(): + ds, do = SyftRDSClient.pair_with_mock_drive_and_mock_webrtc_connection( + use_in_memory_cache=False + ) + script = "mkdir -p code/outputs && echo 42 > code/outputs/result.txt" + + ds.submit_bash_job(do.email, script, "webrtc.job") + do.sync() + do.job_client.jobs[0].approve() + do.process_approved_jobs(stream_output=False, share_outputs_with_submitter=True) + do.sync() + ds.sync() + + [output_path] = ds.job_client.jobs[-1].output_paths + assert Path(output_path).read_text().strip() == "42" + assert drive_message_files(do.sync_engine, ds.email) == [] diff --git a/tests/unit/webrtc/test_webrtc_p2p.py b/tests/unit/webrtc/test_webrtc_p2p.py new file mode 100644 index 00000000000..727a9a8c04c --- /dev/null +++ b/tests/unit/webrtc/test_webrtc_p2p.py @@ -0,0 +1,109 @@ +"""The WebRTC core over a FakeChannel, without managers.""" + +import pytest + +from syft.sync.connections.types import MessageKind, SendStatus +from syft.sync.connections.webrtc.fake_channel import FakeChannel +from syft.sync.connections.webrtc.webrtc_p2p import WebRTCP2PConnection + +DO = "do@example.com" +DS = "ds@example.com" +PROPOSAL = MessageKind.PROPOSED_CHANGES +EVENTS = MessageKind.FILE_EVENTS + + +@pytest.fixture +def cores(tmp_path) -> tuple[WebRTCP2PConnection, WebRTCP2PConnection]: + do = WebRTCP2PConnection(DO, tmp_path / "do") + ds = WebRTCP2PConnection(DS, tmp_path / "ds") + FakeChannel.connect(do, ds) + return do, ds + + +def msg_ids(messages) -> list[str]: + return [m.msg_id for m in messages] + + +def test_sent_message_arrives_and_is_delivered_once(cores): + do, ds = cores + + assert ds.send(DO, "m1", b"job", PROPOSAL) is SendStatus.QUEUED + + [message] = do.receive(DS, PROPOSAL) + assert (message.sender, message.msg_id, message.data) == (DS, "m1", b"job") + assert ds.delivered() == ["m1"] + assert ds.delivered() == [] + + +def test_receive_returns_one_kind_oldest_first_up_to_max(cores): + do, ds = cores + do.send(DS, "e1", b"", EVENTS) + do.send(DS, "p1", b"", PROPOSAL) + do.send(DS, "e2", b"", EVENTS) + do.send(DS, "e3", b"", EVENTS) + + assert msg_ids(ds.receive(DO, EVENTS, max_messages=2)) == ["e1", "e2"] + assert msg_ids(ds.receive(DO, PROPOSAL)) == ["p1"] + + +def test_processed_message_is_not_received_again(cores): + do, ds = cores + ds.send(DO, "m1", b"job", PROPOSAL) + + do.mark_all_processed(do.receive(DS, PROPOSAL)) + + assert do.receive(DS, PROPOSAL) == [] + + +def test_message_sent_twice_is_received_once(cores): + do, ds = cores + + ds.send(DO, "m1", b"job", PROPOSAL) + ds.send(DO, "m1", b"job", PROPOSAL) + + assert msg_ids(do.receive(DS, PROPOSAL)) == ["m1"] + + +def test_messages_wait_on_disk_until_peer_connects(tmp_path): + do = WebRTCP2PConnection(DO, tmp_path / "do") + ds = WebRTCP2PConnection(DS, tmp_path / "ds") + ds.send(DO, "m1", b"first", PROPOSAL) + ds.send(DO, "m2", b"second", PROPOSAL) + assert do.receive(DS, PROPOSAL) == [] + + # A new core on the same folder, as after a restart, still has them. + restarted_ds = WebRTCP2PConnection(DS, tmp_path / "ds") + FakeChannel.connect(do, restarted_ds) + + assert msg_ids(do.receive(DS, PROPOSAL)) == ["m1", "m2"] + assert restarted_ds.delivered() == ["m1", "m2"] + + +def test_message_sent_before_channel_opens_follows_older_ones(tmp_path): + do = WebRTCP2PConnection(DO, tmp_path / "do") + ds = WebRTCP2PConnection(DS, tmp_path / "ds") + ds.send(DO, "m1", b"", PROPOSAL) + do.attach(DS, FakeChannel(sender=DO, receiver=ds)) + ds.attach(DO, FakeChannel(sender=DS, receiver=do)) + + ds.send(DO, "m2", b"", PROPOSAL) + assert do.receive(DS, PROPOSAL) == [] + + ds.on_open(DO) + assert msg_ids(do.receive(DS, PROPOSAL)) == ["m1", "m2"] + + +def test_msg_id_that_is_not_plain_file_name_is_refused(cores): + _, ds = cores + + with pytest.raises(ValueError): + ds.send(DO, "../escape", b"", PROPOSAL) + + +def test_restarted_receiver_keeps_unprocessed_messages(cores, tmp_path): + _, ds = cores + ds.send(DO, "m1", b"job", PROPOSAL) + + restarted_do = WebRTCP2PConnection(DO, tmp_path / "do") + + assert msg_ids(restarted_do.receive(DS, PROPOSAL)) == ["m1"]