From 1369f7ac75ef994df0907a0e33dc3b80e6d851c4 Mon Sep 17 00:00:00 2001 From: Colin Constable Date: Tue, 7 Jul 2026 14:11:32 -0700 Subject: [PATCH] feat: asyncio client proof of concept (at_client.aio) RFC/proof-of-concept for an asyncio-native client, as a parallel package so the synchronous client is untouched: AsyncAtClient with PKAM auth, encrypted notify, and a monitor exposed as an async iterator with a cancellable heartbeat task, bounded reads, automatic reconnect and a seedable resume position. Only the connection layer is new (~330 lines); crypto, key handling and verb builders are reused from the existing package (they do no I/O). The structure makes the thread-lifecycle and wedged-read failure modes of the sync monitor unrepresentable: the heartbeat is a cancellable Task, reads are bounded by asyncio.wait_for, and cancelling the consumer tears everything down. Network-free tests in test/aio_client_test.py (monitor line parsing, decrypt round-trip with ivNonce and legacy zero IV, response framing). Verified live against the ephemeral environment interoperating with the released 0.2.70 sync client: sync->async decrypt, async->sync notify with first-contact shared-key creation, resume-from-epoch with no backlog replay, and clean cancellation. Runnable example in examples/aio_demo.py. --- at_client/aio/__init__.py | 4 + at_client/aio/atclient.py | 216 ++++++++++++++++++++++++++++++++++++ at_client/aio/connection.py | 112 +++++++++++++++++++ examples/aio_demo.py | 46 ++++++++ test/aio_client_test.py | 74 ++++++++++++ 5 files changed, 452 insertions(+) create mode 100644 at_client/aio/__init__.py create mode 100644 at_client/aio/atclient.py create mode 100644 at_client/aio/connection.py create mode 100644 examples/aio_demo.py create mode 100644 test/aio_client_test.py diff --git a/at_client/aio/__init__.py b/at_client/aio/__init__.py new file mode 100644 index 0000000..ab5e39c --- /dev/null +++ b/at_client/aio/__init__.py @@ -0,0 +1,4 @@ +from .atclient import AsyncAtClient, AtNotification +from .connection import AsyncAtConnection + +__all__ = ["AsyncAtClient", "AtNotification", "AsyncAtConnection"] diff --git a/at_client/aio/atclient.py b/at_client/aio/atclient.py new file mode 100644 index 0000000..d573fb2 --- /dev/null +++ b/at_client/aio/atclient.py @@ -0,0 +1,216 @@ +import asyncio +import base64 +import json +import ssl +import uuid +from dataclasses import dataclass, field + +from ..common.atsign import AtSign +from ..common.keys import SharedKey +from ..connections.address import Address +from ..exception.atexception import AtKeyNotFoundException, AtUnauthenticatedException +from ..util.encryptionutil import EncryptionUtil +from ..util.keysutil import KeysUtil +from ..util.verbbuilder import FromVerbBuilder, NotifyVerbBuilder, OperationEnum, PKAMVerbBuilder +from .connection import AsyncAtConnection, find_secondary + +LEGACY_IV = b"\x00" * 16 + + +@dataclass +class AtNotification: + """One decrypted notification from the monitor stream.""" + id: str + key: str + from_atsign: str + to_atsign: str + operation: str + epoch_millis: int + value: str | None + decrypted: bool + raw: dict = field(repr=False) + + +class AsyncAtClient: + """asyncio atProtocol client: PKAM auth, encrypted notify, and a monitor + exposed as an async iterator. + + Proof of concept. Reuses the synchronous SDK's crypto, key handling and verb + builders (all I/O-free); only the connection layer is new. Create with:: + + client = await AsyncAtClient.create("@alice") + async for notification in client.monitor(regex="my_namespace"): + ... + """ + + def __init__(self, atsign: AtSign, keys: dict, secondary_address: Address, + connection: AsyncAtConnection, verbose: bool = False): + self.atsign = atsign + self.keys = keys + self.secondary_address = secondary_address + self.connection = connection + self.verbose = verbose + self._shared_key_cache: dict = {} + + @classmethod + async def create(cls, atsign: str, root_address: Address = Address("root.atsign.org", 64), + secondary_address: Address = None, context: ssl.SSLContext = None, + verbose: bool = False) -> "AsyncAtClient": + at = AtSign(atsign) + keys = KeysUtil.load_keys(at) + if secondary_address is None: + secondary_address = await find_secondary(at, root_address, context, verbose) + connection = await AsyncAtConnection.connect(secondary_address, context, verbose) + client = cls(at, keys, secondary_address, connection, verbose) + await client._authenticate(connection) + return client + + async def _authenticate(self, connection: AsyncAtConnection): + challenge = (await connection.execute_command( + FromVerbBuilder().set_shared_by(self.atsign).build())).get_raw_data_response() + signature = EncryptionUtil.sign_sha256_rsa(challenge, self.keys[KeysUtil.pkam_private_key_name]) + response = await connection.execute_command(PKAMVerbBuilder().set_digest(signature).build()) + if response.get_raw_data_response() != "success": + raise AtUnauthenticatedException(f"PKAM failed: {response}") + + # ---------------------------------------------------------------- notify + async def notify(self, to: str, key_name: str, value: str, + namespace: str = None, ttl_ms: int = 60000) -> str: + """Encrypt value to `to` and send it as an update notification.""" + at_key = SharedKey(key_name, self.atsign, AtSign(to)) + if namespace: + at_key.set_namespace(namespace) + at_key.set_time_to_live(ttl_ms) + iv = EncryptionUtil.generate_iv_nonce() + at_key.metadata.iv_nonce = iv + shared_key = await self._encryption_key_shared_by_me(at_key) + encrypted = EncryptionUtil.aes_encrypt_from_base64(value, shared_key, iv) + command = NotifyVerbBuilder().with_at_key( + at_key, encrypted, OperationEnum.UPDATE, session_id=str(uuid.uuid4())).build() + return (await self.connection.execute_command(command)).get_raw_data_response() + + async def _encryption_key_shared_by_me(self, at_key: SharedKey) -> str: + """Fetch (or create, on first contact) the AES key we share with the recipient.""" + to_lookup = "shared_key." + at_key.shared_with.without_prefix + self.atsign.to_string() + response = await self.connection.execute_command("llookup:" + to_lookup, raise_exception=False) + if not response.is_error(): + return EncryptionUtil.rsa_decrypt_from_base64( + response.get_raw_data_response(), self.keys[KeysUtil.encryption_private_key_name]) + if not isinstance(response.get_exception(), AtKeyNotFoundException): + raise response.get_exception() + their_public = (await self.connection.execute_command( + "plookup:publickey" + at_key.shared_with.to_string())).get_raw_data_response() + aes_key = EncryptionUtil.generate_aes_key_base64() + for_us = EncryptionUtil.rsa_encrypt_to_base64(aes_key, self.keys[KeysUtil.encryption_public_key_name]) + for_them = EncryptionUtil.rsa_encrypt_to_base64(aes_key, their_public) + await self.connection.execute_command("update:" + to_lookup + " " + for_us) + await self.connection.execute_command( + "update:ttr:86400000:" + at_key.shared_with.to_string() + ":shared_key" + + self.atsign.to_string() + " " + for_them) + return aes_key + + async def _encryption_key_shared_by_other(self, sender: str) -> str: + """Fetch (and cache) the AES key `sender` shared with us, for decrypting.""" + if sender in self._shared_key_cache: + return self._shared_key_cache[sender] + response = await self.connection.execute_command("lookup:shared_key" + sender) + key = EncryptionUtil.rsa_decrypt_from_base64( + response.get_raw_data_response(), self.keys[KeysUtil.encryption_private_key_name]) + self._shared_key_cache[sender] = key + return key + + # --------------------------------------------------------------- monitor + async def monitor(self, regex: str = ".*", last_received_time: int = 0, + heartbeat_interval_s: float = 30.0, reconnect_delay_s: float = 3.0): + """Yield AtNotifications forever; reconnects and resumes automatically. + + Cancelling the consuming task stops the heartbeat and closes the + connection. `last_received_time` seeds the first monitor command; after + that the stream resumes from the newest notification it has processed. + """ + last = last_received_time + while True: + connection = None + heartbeat = None + try: + connection = await AsyncAtConnection.connect( + self.secondary_address, verbose=self.verbose) + await self._authenticate(connection) + await connection.write(f"monitor:{last} {regex}\n") + heartbeat = asyncio.create_task(self._heartbeat(connection, heartbeat_interval_s)) + async for notification in self._stream(connection, 3 * heartbeat_interval_s): + if notification.epoch_millis > last: + last = notification.epoch_millis + if notification.id == "-1": + continue # server stats notification + if notification.operation == "update": + await self._decrypt(notification) + yield notification + except (ConnectionError, asyncio.TimeoutError, asyncio.IncompleteReadError, OSError) as e: + if self.verbose: + print(f"[monitor {self.atsign}] {type(e).__name__}: {e}; reconnecting") + finally: + if heartbeat is not None: + heartbeat.cancel() + if connection is not None: + await connection.close() + await asyncio.sleep(reconnect_delay_s) + + async def _stream(self, connection: AsyncAtConnection, read_timeout_s: float): + """Yield parsed notifications from one monitor connection until EOF. + + The read timeout bounds a silently dead socket: heartbeat acks arrive + every interval, so a read that outlives three of them means the + connection is gone even if the OS never reports it. + """ + while True: + line = await asyncio.wait_for(connection.readline(), timeout=read_timeout_s) + if line == b"": + return # EOF: server closed the connection + notification = self._parse_monitor_line(line) + if notification is not None: + yield notification + + async def _heartbeat(self, connection: AsyncAtConnection, interval_s: float): + """Keep the monitor connection alive; its 'data:ok' acks are read (and + discarded) by the monitor loop, which also times out if they stop.""" + while True: + await asyncio.sleep(interval_s) + await connection.write("noop:0\n") + + def _parse_monitor_line(self, line: bytes) -> AtNotification | None: + index = line.find(b"notification:") + if index < 0: + return None # heartbeat ack, prompt remnant, or blank + data = json.loads(line[index + len(b"notification:"):].decode()) + return AtNotification( + id=str(data.get("id")), + key=str(data.get("key")), + from_atsign=str(data.get("from")), + to_atsign=str(data.get("to")), + operation=str(data.get("operation")), + epoch_millis=int(data.get("epochMillis", 0)), + value=None, + decrypted=False, + raw=data, + ) + + async def _decrypt(self, notification: AtNotification): + """Decrypt in place; on failure the notification is yielded undecrypted.""" + encrypted = notification.raw.get("value") + if encrypted is None: + return + metadata = notification.raw.get("metadata") or {} + iv_b64 = metadata.get("ivNonce") + iv = base64.b64decode(iv_b64) if iv_b64 else LEGACY_IV + try: + shared_key = await self._encryption_key_shared_by_other(notification.from_atsign) + notification.value = EncryptionUtil.aes_decrypt_from_base64( + encrypted.encode(), shared_key, iv) + notification.decrypted = True + except Exception as e: + if self.verbose: + print(f"[monitor {self.atsign}] could not decrypt {notification.key}: {e}") + + async def close(self): + await self.connection.close() diff --git a/at_client/aio/connection.py b/at_client/aio/connection.py new file mode 100644 index 0000000..6435b8c --- /dev/null +++ b/at_client/aio/connection.py @@ -0,0 +1,112 @@ +import asyncio +import ssl + +from ..connections.address import Address +from ..connections.response import Response +from ..exception.atexception import AtException + + +class AsyncAtConnection: + """An asyncio connection to an atServer, speaking the same wire protocol as + the synchronous AtConnection: commands are newline-terminated, responses end + at the server's '@' prompt or a newline.""" + + def __init__(self, address: Address, reader: asyncio.StreamReader, + writer: asyncio.StreamWriter, verbose: bool = False): + self.address = address + self._reader = reader + self._writer = writer + self._verbose = verbose + self._command_lock = asyncio.Lock() + + @classmethod + async def connect(cls, address: Address, context: ssl.SSLContext = None, + verbose: bool = False) -> "AsyncAtConnection": + context = context or ssl.create_default_context() + reader, writer = await asyncio.open_connection(address.host, address.port, ssl=context) + connection = cls(address, reader, writer, verbose) + await connection._read_frame() # consume the server's initial '@' prompt + return connection + + async def write(self, data: str): + self._writer.write(data.encode()) + await self._writer.drain() + + async def readline(self) -> bytes: + """Read one newline-terminated line (used by the monitor stream).""" + return await self._reader.readline() + + async def _read_frame(self) -> str: + """Accumulate until the '@' prompt or a newline, as AtConnection.read does.""" + response = b"" + while True: + chunk = await self._reader.read(1024) + if chunk == b"": + raise ConnectionError(f"connection to {self.address} closed by peer") + response += chunk + if chunk == b"@" or b"\n" in chunk: + return response.decode() + + async def execute_command(self, command: str, raise_exception: bool = True) -> Response: + """Send one command and parse its response (same rules as the sync client).""" + if not command.endswith("\n"): + command += "\n" + async with self._command_lock: + await self.write(command) + if self._verbose: + print(f"\tSENT: {command.strip()!r}") + raw = await self._read_frame() + if self._verbose: + print(f"\tRCVD: {raw!r}") + response = self._parse(raw) + if response.is_error() and raise_exception: + raise response.get_exception() + return response + + @staticmethod + def _parse(raw: str) -> Response: + if raw.endswith("@"): + raw = raw[:-1] + raw = raw.strip() + data_index = raw.find("data:") + error_index = raw.find("error:") + if data_index > -1: + return Response().set_raw_data_response(raw[data_index + len("data:"):].split("\n")[0]) + if error_index > -1: + return Response().set_raw_error_response(raw[error_index + len("error:"):]) + raise AtException(f"Invalid response from server: {raw}") + + async def close(self): + try: + self._writer.close() + await self._writer.wait_closed() + except Exception: + pass + + +async def find_secondary(atsign, root_address: Address, + context: ssl.SSLContext = None, verbose: bool = False) -> Address: + """Look up an atSign's secondary address on the root server. + + The root speaks a bare protocol: send the atSign without its '@' prefix, + receive 'host:port' (or 'null' when the atSign is not found). + """ + context = context or ssl.create_default_context() + reader, writer = await asyncio.open_connection(root_address.host, root_address.port, ssl=context) + try: + await reader.readuntil(b"@") # initial prompt + writer.write((atsign.without_prefix + "\n").encode()) + await writer.drain() + raw = (await reader.readuntil(b"@")).decode() + finally: + writer.close() + try: + await writer.wait_closed() + except Exception: + pass + response = raw.strip().rstrip("@").strip() + if verbose: + print(f"\troot lookup {atsign}: {response!r}") + if response == "null": + raise AtException(f"Root lookup returned null for {atsign}") + return Address.from_string(response) diff --git a/examples/aio_demo.py b/examples/aio_demo.py new file mode 100644 index 0000000..27c0c9a --- /dev/null +++ b/examples/aio_demo.py @@ -0,0 +1,46 @@ +#!/usr/bin/env python3 +"""asyncio client demo: subscribe to notifications as an async stream, and send one. + +Usage (keys in $HOME/.atsign/keys): + python examples/aio_demo.py @alice # subscribe + python examples/aio_demo.py @alice --to @bob --send hi # notify @bob, then subscribe + +Optional: --root host:port (default root.atsign.org:64), --regex +""" +import argparse +import asyncio + +from at_client.aio import AsyncAtClient +from at_client.connections.address import Address + + +async def main(): + parser = argparse.ArgumentParser() + parser.add_argument("atsign") + parser.add_argument("--root", default="root.atsign.org:64") + parser.add_argument("--regex", default=".*") + parser.add_argument("--to", help="atSign to notify") + parser.add_argument("--send", help="value to notify --to with") + args = parser.parse_args() + + client = await AsyncAtClient.create(args.atsign, root_address=Address.from_string(args.root)) + print(f"authenticated as {args.atsign}") + + if args.to and args.send: + notification_id = await client.notify(args.to, "demo", args.send, namespace="aiodemo") + print(f"notified {args.to}: {notification_id}") + + print(f"listening (regex={args.regex!r}) — Ctrl-C to stop") + try: + async for notification in client.monitor(regex=args.regex): + print(f"{notification.from_atsign} -> {notification.key}: " + f"{notification.value if notification.decrypted else ''}") + finally: + await client.close() + + +if __name__ == "__main__": + try: + asyncio.run(main()) + except KeyboardInterrupt: + pass diff --git a/test/aio_client_test.py b/test/aio_client_test.py new file mode 100644 index 0000000..2f07719 --- /dev/null +++ b/test/aio_client_test.py @@ -0,0 +1,74 @@ +import base64 +import json +import unittest + +from at_client.aio.atclient import AsyncAtClient, LEGACY_IV +from at_client.aio.connection import AsyncAtConnection +from at_client.util import EncryptionUtil + + +def _client_with_cached_shared_key(sender: str, aes_key: str) -> AsyncAtClient: + client = AsyncAtClient.__new__(AsyncAtClient) # bypass the network-connecting create() + client.verbose = False + client._shared_key_cache = {sender: aes_key} + return client + + +def _notification_line(value: str, aes_key: str, iv: bytes = None) -> bytes: + metadata = {} + if iv is not None: + metadata["ivNonce"] = base64.b64encode(iv).decode() + encrypted = EncryptionUtil.aes_encrypt_from_base64(value, aes_key, iv if iv is not None else LEGACY_IV) + data = {"id": "abc", "from": "@bob", "to": "@alice", "key": "@alice:demo.test@bob", + "operation": "update", "epochMillis": 1720000000000, "value": encrypted, + "isEncrypted": True, "metadata": metadata} + return b"notification: " + json.dumps(data).encode() + b"\n" + + +class AioClientTest(unittest.IsolatedAsyncioTestCase): + """Network-free tests for the asyncio client's parsing and decryption.""" + + def test_parse_monitor_line(self): + client = _client_with_cached_shared_key("@bob", EncryptionUtil.generate_aes_key_base64()) + line = _notification_line("hi", EncryptionUtil.generate_aes_key_base64()) + notification = client._parse_monitor_line(line) + self.assertEqual(notification.from_atsign, "@bob") + self.assertEqual(notification.operation, "update") + self.assertEqual(notification.epoch_millis, 1720000000000) + + def test_parse_monitor_line_tolerates_prompt_remnant(self): + client = _client_with_cached_shared_key("@bob", EncryptionUtil.generate_aes_key_base64()) + line = b"@alice@" + _notification_line("hi", EncryptionUtil.generate_aes_key_base64()) + self.assertIsNotNone(client._parse_monitor_line(line)) + + def test_parse_monitor_line_ignores_heartbeat_ack(self): + client = _client_with_cached_shared_key("@bob", EncryptionUtil.generate_aes_key_base64()) + self.assertIsNone(client._parse_monitor_line(b"data:ok\n")) + + async def test_decrypt_with_iv_nonce(self): + aes_key = EncryptionUtil.generate_aes_key_base64() + iv = EncryptionUtil.generate_iv_nonce() + client = _client_with_cached_shared_key("@bob", aes_key) + notification = client._parse_monitor_line(_notification_line("secret value", aes_key, iv)) + await client._decrypt(notification) + self.assertTrue(notification.decrypted) + self.assertEqual(notification.value, "secret value") + + async def test_decrypt_legacy_zero_iv(self): + aes_key = EncryptionUtil.generate_aes_key_base64() + client = _client_with_cached_shared_key("@bob", aes_key) + notification = client._parse_monitor_line(_notification_line("legacy value", aes_key)) + await client._decrypt(notification) + self.assertTrue(notification.decrypted) + self.assertEqual(notification.value, "legacy value") + + def test_response_parsing_matches_sync_rules(self): + response = AsyncAtConnection._parse("data:hello\n@alice@") + self.assertEqual(response.get_raw_data_response(), "hello") + response = AsyncAtConnection._parse("error:AT0015-key not found : nope@") + self.assertTrue(response.is_error()) + self.assertEqual(response.get_error_code(), "AT0015") + + +if __name__ == "__main__": + unittest.main()