nwcprovider/nwcp.py
2026-08-26 13:06:39 +02:00

689 lines
24 KiB
Python

import asyncio
import hashlib
import json
import random
import time
from collections.abc import Awaitable, Callable
from typing import Any, Union
from coincurve import PublicKeyXOnly
from lnbits.helpers import encrypt_internal_message
from lnbits.settings import settings
from loguru import logger
from pynostr.key import PrivateKey
from websockets.legacy.client import connect
class RateLimit:
backoff: int = 0
last_attempt_time: int = 0
class MainSubscription:
def __init__(self):
self.requests_sub_id: str | None = None
self.responses_sub_id: str | None = None
self.requests_eose = False
self.responses_eose = False
self.events: dict[str, dict] = {}
self.responses: list[str] = []
self.seen_requests: dict[str, int] = {}
def get_stale(self) -> list[dict]:
"""
Get all the pending events that do not have a response yet.
"""
pending_events = []
for [event_id, event] in self.events.items():
if event_id not in self.responses:
pending_events.append(event)
return pending_events
def register_response(self, event_id: str):
"""
Register a response for a request event (not stale anymore)
"""
if event_id not in self.responses:
self.responses.append(event_id)
def gc(self, expire: int | None = None):
"""
Garbage collection, remove all the events that have a response older
than expire seconds (defaults to 1 hour if 0 or None)
and all seen requests that are expired
"""
expire = expire or 1 * 60 * 60
now = int(time.time())
deleted_ids = []
for event_id, event in list(self.events.items()):
if event_id in self.responses:
if now - event["created_at"] > expire:
del self.events[event_id]
deleted_ids.append(event_id)
self.responses = [
event_id for event_id in self.responses if event_id not in deleted_ids
]
if len(deleted_ids) > 0:
logger.debug("Garbage collected " + str(len(deleted_ids)) + " events")
# Clean seen requests
for event_id, expiry in list(self.seen_requests.items()):
if expiry < now:
del self.seen_requests[event_id]
class Config:
arbitrary_types_allowed = True
class NWCServiceProvider:
def __init__(
self,
private_key_hex: str | None = None,
relay: str | None = None,
handle_missed_events: int = 0,
):
if not relay: # Connect to nostrclient
relay = "nostrclient"
if relay == "nostrclient":
relay = f"ws://localhost:{settings.port}/nostrclient/api/v1/relay"
elif relay == "nostrclient:private":
relay_endpoint = encrypt_internal_message("relay")
relay = (
f"ws://localhost:{settings.port}/nostrclient/api/v1/{relay_endpoint}"
)
self.relay = relay
if not private_key_hex: # Create random key
self.private_key = PrivateKey()
self.private_key_hex = self.private_key.hex()
else:
self.private_key = PrivateKey.from_hex(private_key_hex)
self.private_key_hex = private_key_hex
self.public_key = self.private_key.public_key
if not self.public_key:
raise Exception("Invalid public key")
self.public_key_hex = self.public_key.hex()
# List of supported methods
self.supported_methods: list[str] = []
# Keep track of the number of subscriptions (used for unique subid)
self.subscriptions_count: int = 0
# Request listeners, listen to specific methods
self.request_listeners: dict[
str,
Callable[
[NWCServiceProvider, str, dict],
Awaitable[list[tuple[dict | None, dict | None, list]]],
],
] = {}
# Reconnect task (if the connection is lost)
self.reconnect_task = None
# Garbage collection loop
self.gc_task = None
# Periodic info event resend loop
self.info_event_task = None
# Requests are handled independently from the relay receive loop.
self.request_tasks: set[asyncio.Task[list[dict]]] = set()
# Subscription
self.sub: MainSubscription | None = None
self.rate_limit: dict[str, RateLimit] = {}
# websocket connection
self.ws = None
# if True the websocket is connected
self.connected = False
# if True the instance is shutting down
self.shutdown = False
# process missed events that are not older than
# handle_missed_events seconds (0 to disable)
# (handles reboots)
self.handle_missed_events = handle_missed_events
self.event_max_age = self.handle_missed_events or 5 * 60
logger.info(
"NWC Service is ready. relay: "
+ str(self.relay)
+ " pubkey: "
+ self.public_key_hex
)
async def _gc_loop(self):
while not self._is_shutting_down():
if self.sub:
self.sub.gc(self.handle_missed_events)
await asyncio.sleep(60)
def get_supported_methods(self):
"""
Returns the list of supported methods by this service provider.
"""
return self.supported_methods
def add_request_listener(
self,
method: str,
listener: Callable[
["NWCServiceProvider", str, dict],
Awaitable[list[tuple[dict | None, dict | None, list]]],
],
):
"""
Adds a request listener for a specific method.
Args:
method (str): The method name.
listener (Callable[
["NWCServiceProvider", str, Dict], List[Tuple[Dict, Dict]]
]): The listener function
"""
if method not in self.supported_methods:
self.supported_methods.append(method)
self.request_listeners[method] = listener
async def start(self):
"""
Starts the NWC service provider.
"""
self.reconnect_task = asyncio.create_task(self._connect_to_relay())
self.gc_task = asyncio.create_task(self._gc_loop())
self.info_event_task = asyncio.create_task(self._info_event_loop())
def _json_dumps(self, data: Union[dict, list]) -> str:
"""
Converts a Python dictionary to a JSON string with compact encoding.
Args:
data (Dict): The dictionary to be converted.
Returns:
str: The compact JSON string.
"""
if isinstance(data, dict):
data = {k: v for k, v in data.items() if v is not None}
return json.dumps(data, separators=(",", ":"), ensure_ascii=False)
def _is_shutting_down(self) -> bool:
"""
Returns True if the instance is shutting down.
"""
return self.shutdown or not settings.lnbits_running
async def _send(self, data: list[Any]):
"""
Sends data to the relay.
Args:
data (Dict): The data to be sent.
"""
if not self.ws:
raise Exception("Websocket connection is not established")
if self._is_shutting_down():
logger.warning("Trying to send data while shutting down")
return
await self._wait_for_connection() # ensure the connection is established
tx = self._json_dumps(data)
await self.ws.send(tx)
def _get_new_subid(self) -> str:
"""
Generates a unique subscription id.
Returns:
str: The generated 64 characters long subscription id (eg. lnbits0abc...)
"""
subid = "lnbitsnwcs" + str(self.subscriptions_count)
self.subscriptions_count += 1
max_length = 64
chars = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"
n = max_length - len(subid)
if n > 0:
for _ in range(n):
subid += chars[random.randint(0, len(chars) - 1)]
return subid
async def _wait_for_connection(self):
"""
Waits until the connection is established.
"""
while not self.connected:
if self._is_shutting_down():
raise Exception("Connection is closing")
logger.debug("Waiting for connection...")
await asyncio.sleep(1)
async def _ratelimit(self, unit: str, max_sleep_time: int = 120) -> None:
limit: RateLimit | None = self.rate_limit.get(unit)
if not limit:
self.rate_limit[unit] = limit = RateLimit()
if time.time() - limit.last_attempt_time > max_sleep_time:
# reset backoff if action lasted more than max_sleep_time
limit.backoff = 0
else:
# increase backoff
limit.backoff = (
min(limit.backoff * 2, max_sleep_time) if limit.backoff > 0 else 1
)
logger.debug("Sleeping for " + str(limit.backoff) + " seconds before " + unit)
await asyncio.sleep(limit.backoff)
limit.last_attempt_time = int(time.time())
def _create_subscription(self) -> MainSubscription:
sub = MainSubscription()
self.sub = sub
return sub
async def _subscribe(self):
"""
[Re]Subscribe to receive nip 47 requests and responses from the relay
"""
sub = self._create_subscription()
# Create requests subscription
req_filter = {
"kinds": [23194],
"#p": [self.public_key_hex],
# Since the last handle_missed_events seconds (handles reboots)
"since": int(time.time()) - self.handle_missed_events,
}
sub.requests_sub_id = self._get_new_subid()
# Create responses subscription (needed to track previosly responded requests)
res_filter = {
"kinds": [23195],
"authors": [self.public_key_hex],
"since": int(time.time()) - self.handle_missed_events,
}
sub.responses_sub_id = self._get_new_subid()
# Subscribe
await self._send(["REQ", sub.requests_sub_id, req_filter])
await self._send(["REQ", sub.responses_sub_id, res_filter])
async def _on_connection(self, _):
"""
On connection callback, announce the service provider
methods and subscribe to nip67 events.
"""
# Send info event
await self._send_info_event()
# Resubscribe to nwc events
await self._subscribe()
async def _send_info_event(self):
"""
Build and publish the NWC service info event (kind 13194).
"""
event = {
"kind": 13194,
"content": " ".join(self.supported_methods),
"created_at": int(time.time()),
"tags": [["p", self.public_key_hex]],
}
self._sign_event(event)
await self._send(["EVENT", event])
async def _info_event_loop(self):
"""
Periodically resend the service info event (kind 13194) so that the
provider can recover if the relay silently dropped the event without
closing the WebSocket connection.
"""
while not self._is_shutting_down():
await asyncio.sleep(60)
if self.connected and not self._is_shutting_down():
try:
await self._send_info_event()
except Exception as e:
logger.warning("Error resending info event: " + str(e))
async def _handle_request(self, event: dict) -> list[dict]:
"""
Handle a nwc request
"""
if not self.sub:
raise Exception("Subscription is not established")
sub = self.sub
expire = sub.seen_requests.get(event["id"])
if expire or event["created_at"] < int(time.time() - self.event_max_age):
raise Exception("Event is too old or already handled")
expiration = self._extract_expiration_from_tags(event["tags"])
if expiration <= 0:
expiration = int(time.time() + self.event_max_age)
sub.seen_requests[event["id"]] = expiration
nwc_pubkey = event["pubkey"]
content = event["content"]
# Decrypt the content
content = self.private_key.decrypt_message(content, nwc_pubkey)
# Deserialize content
content = json.loads(content)
# Handle request
method = content["method"]
listener = self.request_listeners.get(method, None)
outs: list[dict[str, Any]] = []
if not listener:
outs.append(
{
"error": {
"code": "NOT_IMPLEMENTED",
"message": "Method "
+ method
+ " is not implemented by this service provider",
}
}
)
else:
try:
results = await listener(self, nwc_pubkey, content)
for result in results:
r = result[0]
e = result[1]
t = result[2] if len(result) > 2 else []
out = {"result": r, "error": e, "tags": t}
outs.append(out)
except Exception as e:
outs.append({"error": {"code": "INTERNAL", "message": str(e)}})
sent_events = []
for out in outs:
# Finalize output
content = {}
content["result_type"] = method
if "result" in out:
content["result"] = out["result"]
if "error" in out:
content["error"] = out["error"]
raw_tags = out.get("tags")
tags = list(raw_tags) if isinstance(raw_tags, list) else []
# Prepare response event
res: dict = {
"kind": 23195,
"created_at": int(time.time()),
"tags": tags,
"content": self._json_dumps(content),
}
# Reference request
res["tags"].append(["e", event["id"]])
# Reference user
res["tags"].append(["p", nwc_pubkey])
# Finalize response event
res["content"] = self.private_key.encrypt_message(
res["content"], nwc_pubkey
)
self._sign_event(res)
# Register response for this request, so we knows it is not stale
if self.sub:
self.sub.register_response(event["id"])
# Send response event
await self._send(["EVENT", res])
# Track sent events
sent_events.append(res)
return sent_events
def _log_request_task_exception(self, task: asyncio.Future[list[dict]]) -> None:
if task.cancelled():
return
exception = task.exception()
if exception:
logger.error("Error handling request: " + str(exception))
def _dispatch_request(self, event: dict) -> None:
task = asyncio.create_task(self._handle_request(event))
self.request_tasks.add(task)
task.add_done_callback(self.request_tasks.discard)
task.add_done_callback(self._log_request_task_exception)
def _extract_expiration_from_tags(self, tags: list) -> int:
expiration = -1
for tag in tags:
try:
if tag[0] == "expiration" and len(tag) > 1:
expiration = int(tag[1])
break
except Exception:
pass
return expiration
async def _on_event_message(self, msg):
if not self.sub:
return
sub_id = msg[1]
event = msg[2]
# Ensure the event is valid (do not trust relays)
if not self._verify_event(event):
raise Exception("Invalid event signature")
tags = event["tags"]
expiration = self._extract_expiration_from_tags(tags)
# Handle event expiration if the relay doesn't support nip 40
if expiration > 0 and expiration < int(time.time()):
logger.debug("Event expired")
return
if event["kind"] == 23194 and sub_id == self.sub.requests_sub_id:
# Ensure the request is for this service provider
valid_p = any(
tag[0] == "p" and tag[1] == self.public_key_hex for tag in tags
)
if not valid_p:
raise Exception("Unexpected request from another service")
# Track request
self.sub.events[event["id"]] = event
# if eose was received for both subscriptions, we handle the request
# in realtime if not, we do nothing since the request may be
# already handled or stale, all stale requests will be handled
# later when eose is received
if self.sub.requests_eose and self.sub.responses_eose:
self._dispatch_request(event)
elif event["kind"] == 23195 and sub_id == self.sub.responses_sub_id:
# Ensure the response is from this service provider
if event["pubkey"] != self.public_key_hex:
raise Exception("Unexpected response from another service")
# Register as response for each e tag (request event id)
# Note: usually we expect only one "e" tag, but we are handling
# multiple "e" tags just in case
etag = next((tag[1] for tag in tags if tag[0] == "e"), None)
if etag:
self.sub.register_response(etag)
async def _on_eose_message(self, msg):
if not self.sub:
return
sub_id = msg[1]
# Track EOSE
if sub_id == self.sub.requests_sub_id:
self.sub.requests_eose = True
elif sub_id == self.sub.responses_sub_id:
self.sub.responses_eose = True
# When both EOSE are receives, handle all the stale requests
# Note: All the requests that were received prior to the
# service connection and do not have a response yet,
# are considered stale, we will process them now
if self.sub.requests_eose and self.sub.responses_eose:
stales = self.sub.get_stale()
for stale in stales:
self._dispatch_request(stale)
async def _on_closed_message(self, msg):
if not self.sub:
return
# Subscription was closed remotely.
sub_id = msg[1]
info = msg[2] or "" if len(msg) > 2 else ""
# Resubscribe if one of the main subscriptions was closed
if sub_id == self.sub.requests_sub_id or sub_id == self.sub.responses_sub_id:
logger.warning(
"Subscription "
+ sub_id
+ " was closed remotely: "
+ info
+ " ... resubscribing..."
)
await self._ratelimit("subscribing")
await self._subscribe()
async def _on_message(self, _, message: str):
"""
Handle incoming messages from the relay.
"""
try:
msg = json.loads(message)
if msg[0] == "EVENT": # Event message
await self._on_event_message(msg)
elif msg[0] == "EOSE":
await self._on_eose_message(msg)
elif msg[0] == "CLOSED":
await self._on_closed_message(msg)
elif msg[0] == "NOTICE":
# A message from the relay, mostly useless, but we log it anyway
logger.info("Notice from relay " + self.relay + ": " + str(msg[1]))
elif msg[0] == "OK":
pass
else:
raise Exception("Unknown message type " + str(msg[0]))
except Exception as e:
logger.error("Error parsing event: " + str(e))
async def _connect_to_relay(self):
"""
Initiate websocket connection to the relay.
"""
await asyncio.sleep(1)
logger.debug("Connecting to NWC relay " + self.relay)
while (
not self._is_shutting_down()
): # Reconnect until the instance is shutting down
logger.debug("Creating new connection...")
try:
async with connect(self.relay) as ws:
self.ws = ws
self.connected = True
await self._on_connection(ws)
while (
not self._is_shutting_down()
): # receive messages until the instance is shutting down
try:
reply = await ws.recv()
if isinstance(reply, bytes):
reply = reply.decode("utf-8")
await self._on_message(ws, reply)
except Exception as e:
logger.debug("Error receiving message: " + str(e))
break
logger.debug("Connection to NWC relay closed")
except Exception as e:
logger.error("Error connecting to NWC relay: " + str(e))
await asyncio.sleep(5)
# the connection was closed, so we set the connected flag to False
# this will make the methods calling _wait_for_connection() to wait
# until the connection is re-established
self.connected = False
if not self._is_shutting_down():
# Wait some time before reconnecting
logger.debug("Reconnecting to NWC relay...")
await self._ratelimit("connecting")
def _verify_event(self, event: dict) -> bool:
"""
Verify the event signature
Args:
event (Dict): The event to verify.
Returns:
bool: True if the event signature is valid, False otherwise.
"""
signature_data = self._json_dumps(
[
0,
event["pubkey"],
event["created_at"],
event["kind"],
event["tags"],
event["content"],
]
)
event_id = hashlib.sha256(signature_data.encode()).hexdigest()
if event_id != event["id"]: # Invalid event id
return False
pubkey_hex = event["pubkey"]
pubkey = PublicKeyXOnly(bytes.fromhex(pubkey_hex))
if not pubkey.verify(bytes.fromhex(event["sig"]), bytes.fromhex(event_id)):
return False
return True
def _sign_event(self, event: dict) -> dict:
"""
Signs the event (in place)
Args:
event (Dict): The event to be signed.
Returns:
Dict: The input event with the signature added.
"""
signature_data = self._json_dumps(
[
0,
self.public_key_hex,
event["created_at"],
event["kind"],
event["tags"],
event["content"],
]
)
event_id = hashlib.sha256(signature_data.encode()).hexdigest()
event["id"] = event_id
event["pubkey"] = self.public_key_hex
signature = self.private_key.sign(bytes.fromhex(event_id))
event["sig"] = signature.hex() # type: ignore
return event
async def cleanup(self):
logger.debug("Closing NWC Service Provider connection")
self.shutdown = True # Mark for shutdown
# close tasks
try:
if self.reconnect_task:
self.reconnect_task.cancel()
except Exception as e:
logger.warning("Error closing reconnection task: " + str(e))
try:
if self.gc_task:
self.gc_task.cancel()
except Exception as e:
logger.warning("Error closing gc loop: " + str(e))
try:
if self.info_event_task:
self.info_event_task.cancel()
except Exception as e:
logger.warning("Error closing info event loop: " + str(e))
request_tasks = list(self.request_tasks)
for task in request_tasks:
task.cancel()
if request_tasks:
await asyncio.gather(*request_tasks, return_exceptions=True)
self.request_tasks.clear()
# close the websocket
try:
if self.ws:
await self.ws.close()
except Exception as e:
logger.warning("Error closing websocket connection: " + str(e))
class Config:
arbitrary_types_allowed = True