Stabilize (#24)

* refactor: clean-up

* refactor: extra logs plus try-catch

* refactor: do not use bare `except`

* refactor: clean-up redundant fields

* chore: pass code checks

* chore: code format

* refactor: code clean-up

* fix: refactoring stuff

* refactor: remove un-used file

* chore: code clean-up

* chore: code clean-up

* chore: code-format fix

* refactor: remove nostr.client wrapper

* refactor: code clean-up

* chore: code format

* refactor: remove `RelayList` class

* refactor: extract smaller methods with try-catch

* fix: better exception handling

* fix: remove redundant filters

* fix: simplify event

* chore: code format

* fix: code check

* fix: code check

* fix: simplify `REQ`

* fix: more clean-ups

* refactor: use simpler method

* refactor: re-order and rename

* fix: stop logic

* fix: subscription close before disconnect

* chore: play commit
This commit is contained in:
Vlad Stan 2023-11-01 17:46:42 +02:00 committed by GitHub
parent ab185bd2c4
commit 16ae9d15a1
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
20 changed files with 522 additions and 717 deletions

178
router.py
View file

@ -1,42 +1,61 @@
import asyncio
import json
from typing import List, Union
from typing import Dict, List
from fastapi import WebSocketDisconnect
from fastapi import WebSocket, WebSocketDisconnect
from loguru import logger
from lnbits.helpers import urlsafe_short_hash
from . import nostr
from .models import Event, Filter
from .nostr.filter import Filter as NostrFilter
from .nostr.filter import Filters as NostrFilters
from .nostr.message_pool import EndOfStoredEventsMessage, NoticeMessage
from . import nostr_client
from .nostr.message_pool import EndOfStoredEventsMessage, EventMessage, NoticeMessage
class NostrRouter:
received_subscription_events: dict[str, list[Event]] = {}
received_subscription_events: dict[str, List[EventMessage]] = {}
received_subscription_notices: list[NoticeMessage] = []
received_subscription_eosenotices: dict[str, EndOfStoredEventsMessage] = {}
def __init__(self, websocket):
self.subscriptions: List[str] = []
def __init__(self, websocket: WebSocket):
self.connected: bool = True
self.websocket = websocket
self.websocket: WebSocket = websocket
self.tasks: List[asyncio.Task] = []
self.original_subscription_ids = {}
self.original_subscription_ids: Dict[str, str] = {}
async def client_to_nostr(self):
"""Receives requests / data from the client and forwards it to relays. If the
request was a subscription/filter, registers it with the nostr client lib.
Remembers the subscription id so we can send back responses from the relay to this
client in `nostr_to_client`"""
while True:
@property
def subscriptions(self) -> List[str]:
return list(self.original_subscription_ids.keys())
def start(self):
self.connected = True
self.tasks.append(asyncio.create_task(self._client_to_nostr()))
self.tasks.append(asyncio.create_task(self._nostr_to_client()))
async def stop(self):
nostr_client.relay_manager.close_subscriptions(self.subscriptions)
self.connected = False
for t in self.tasks:
try:
t.cancel()
except Exception as _:
pass
try:
await self.websocket.close()
except Exception as _:
pass
async def _client_to_nostr(self):
"""
Receives requests / data from the client and forwards it to relays.
"""
while self.connected:
try:
json_str = await self.websocket.receive_text()
except WebSocketDisconnect:
self.connected = False
except WebSocketDisconnect as e:
logger.debug(e)
await self.stop()
break
try:
@ -44,15 +63,9 @@ class NostrRouter:
except Exception as e:
logger.debug(f"Failed to handle client message: '{str(e)}'.")
async def nostr_to_client(self):
"""Sends responses from relays back to the client. Polls the subscriptions of this client
stored in `my_subscriptions`. Then gets all responses for this subscription id from `received_subscription_events` which
is filled in tasks.py. Takes one response after the other and relays it back to the client. Reconstructs
the reponse manually because the nostr client lib we're using can't do it. Reconstructs the original subscription id
that we had previously rewritten in order to avoid collisions when multiple clients use the same id.
"""
while True and self.connected:
async def _nostr_to_client(self):
"""Sends responses from relays back to the client."""
while self.connected:
try:
await self._handle_subscriptions()
self._handle_notices()
@ -61,24 +74,6 @@ class NostrRouter:
await asyncio.sleep(0.1)
async def start(self):
self.tasks.append(asyncio.create_task(self.client_to_nostr()))
self.tasks.append(asyncio.create_task(self.nostr_to_client()))
async def stop(self):
for t in self.tasks:
try:
t.cancel()
except:
pass
for s in self.subscriptions:
try:
nostr.client.relay_manager.close_subscription(s)
except:
pass
self.connected = False
async def _handle_subscriptions(self):
for s in self.subscriptions:
if s in NostrRouter.received_subscription_events:
@ -86,8 +81,6 @@ class NostrRouter:
if s in NostrRouter.received_subscription_eosenotices:
await self._handle_received_subscription_eosenotices(s)
async def _handle_received_subscription_eosenotices(self, s):
try:
if s not in self.original_subscription_ids:
@ -95,7 +88,7 @@ class NostrRouter:
s_original = self.original_subscription_ids[s]
event_to_forward = ["EOSE", s_original]
del NostrRouter.received_subscription_eosenotices[s]
await self.websocket.send_text(json.dumps(event_to_forward))
except Exception as e:
logger.debug(e)
@ -104,97 +97,62 @@ class NostrRouter:
try:
if s not in NostrRouter.received_subscription_events:
return
while len(NostrRouter.received_subscription_events[s]):
my_event = NostrRouter.received_subscription_events[s].pop(0)
# event.to_message() does not include the subscription ID, we have to add it manually
event_json = {
"id": my_event.id,
"pubkey": my_event.public_key,
"created_at": my_event.created_at,
"kind": my_event.kind,
"tags": my_event.tags,
"content": my_event.content,
"sig": my_event.signature,
}
event_message = NostrRouter.received_subscription_events[s].pop(0)
event_json = event_message.event
# this reconstructs the original response from the relay
# reconstruct original subscription id
s_original = self.original_subscription_ids[s]
event_to_forward = ["EVENT", s_original, event_json]
await self.websocket.send_text(json.dumps(event_to_forward))
event_to_forward = f"""["EVENT", "{s_original}", {event_json}]"""
await self.websocket.send_text(event_to_forward)
except Exception as e:
logger.debug(e)
logger.debug(e) # there are 2900 errors here
def _handle_notices(self):
while len(NostrRouter.received_subscription_notices):
my_event = NostrRouter.received_subscription_notices.pop(0)
# note: we don't send it to the user because we don't know who should receive it
logger.info(f"Relay ('{my_event.url}') notice: '{my_event.content}']")
nostr.client.relay_manager.handle_notice(my_event)
def _marshall_nostr_filters(self, data: Union[dict, list]):
filters = data if isinstance(data, list) else [data]
filters = [Filter.parse_obj(f) for f in filters]
filter_list: list[NostrFilter] = []
for filter in filters:
filter_list.append(
NostrFilter(
event_ids=filter.ids, # type: ignore
kinds=filter.kinds, # type: ignore
authors=filter.authors, # type: ignore
since=filter.since, # type: ignore
until=filter.until, # type: ignore
event_refs=filter.e, # type: ignore
pubkey_refs=filter.p, # type: ignore
limit=filter.limit, # type: ignore
)
)
return NostrFilters(filter_list)
logger.info(f"[Relay '{my_event.url}'] Notice: '{my_event.content}']")
# Note: we don't send it to the user because
# we don't know who should receive it
nostr_client.relay_manager.handle_notice(my_event)
async def _handle_client_to_nostr(self, json_str):
"""Parses a (string) request from a client. If it is a subscription (REQ) or a CLOSE, it will
register the subscription in the nostr client library that we're using so we can
receive the callbacks on it later. Will rewrite the subscription id since we expect
multiple clients to use the router and want to avoid subscription id collisions
"""
json_data = json.loads(json_str)
assert len(json_data)
assert len(json_data), "Bad JSON array"
if json_data[0] == "REQ":
self._handle_client_req(json_data)
return
if json_data[0] == "CLOSE":
self._handle_client_close(json_data[1])
return
if json_data[0] == "EVENT":
nostr.client.relay_manager.publish_message(json_str)
nostr_client.relay_manager.publish_message(json_str)
return
def _handle_client_req(self, json_data):
subscription_id = json_data[1]
subscription_id_rewritten = urlsafe_short_hash()
self.original_subscription_ids[subscription_id_rewritten] = subscription_id
fltr = json_data[2:]
filters = self._marshall_nostr_filters(fltr)
filters = json_data[2:]
nostr.client.relay_manager.add_subscription(
subscription_id_rewritten, filters
)
request_rewritten = json.dumps([json_data[0], subscription_id_rewritten] + fltr)
self.subscriptions.append(subscription_id_rewritten)
nostr.client.relay_manager.publish_message(request_rewritten)
nostr_client.relay_manager.add_subscription(subscription_id_rewritten, filters)
def _handle_client_close(self, subscription_id):
subscription_id_rewritten = next((k for k, v in self.original_subscription_ids.items() if v == subscription_id), None)
subscription_id_rewritten = next(
(
k
for k, v in self.original_subscription_ids.items()
if v == subscription_id
),
None,
)
if subscription_id_rewritten:
self.original_subscription_ids.pop(subscription_id_rewritten)
nostr.client.relay_manager.close_subscription(subscription_id_rewritten)
nostr_client.relay_manager.close_subscription(subscription_id_rewritten)
else:
logger.debug(f"Failed to unsubscribe from '{subscription_id}.'")