|
@@ -0,0 +1,187 @@
|
|
|
+use crate::{Connection, Error};
|
|
|
+use nostr_rs_storage::RocksDb;
|
|
|
+use nostr_rs_types::{
|
|
|
+ relayer,
|
|
|
+ types::{Event, Filter, SubscriptionId},
|
|
|
+ Request, Response,
|
|
|
+};
|
|
|
+use parking_lot::{RwLock, RwLockReadGuard};
|
|
|
+use std::{collections::HashMap, ops::Deref};
|
|
|
+use tokio::{net::TcpStream, sync::mpsc::Sender};
|
|
|
+
|
|
|
+#[derive(Clone, Debug, Default, Eq, PartialEq, Hash)]
|
|
|
+pub struct SubscriptionType {
|
|
|
+ public_key: Option<Vec<u8>>,
|
|
|
+ id: Option<Vec<u8>>,
|
|
|
+ kind: Option<u32>,
|
|
|
+}
|
|
|
+
|
|
|
+type Subscriptions = HashMap<(u128, u128), (SubscriptionId, Sender<Response>)>;
|
|
|
+
|
|
|
+pub struct Relayer {
|
|
|
+ storage: RocksDb,
|
|
|
+ subscriptions: RwLock<HashMap<SubscriptionType, RwLock<Subscriptions>>>,
|
|
|
+ clients: RwLock<HashMap<u128, Connection>>,
|
|
|
+}
|
|
|
+
|
|
|
+impl Relayer {
|
|
|
+ pub fn new(storage: RocksDb) -> Self {
|
|
|
+ Self {
|
|
|
+ storage,
|
|
|
+ subscriptions: RwLock::new(HashMap::new()),
|
|
|
+ clients: RwLock::new(HashMap::new()),
|
|
|
+ }
|
|
|
+ }
|
|
|
+
|
|
|
+ pub async fn add_connection(&self, stream: TcpStream) -> Result<(), Error> {
|
|
|
+ let client = Connection::new(stream).await?;
|
|
|
+ let mut clients = self.clients.write();
|
|
|
+ clients.insert(client.conn_id, client);
|
|
|
+
|
|
|
+ Ok(())
|
|
|
+ }
|
|
|
+
|
|
|
+ pub async fn recv(&self, conn_id: u128, request: Request) -> Result<(), Error> {
|
|
|
+ let connections = self.clients.read();
|
|
|
+
|
|
|
+ let connection = connections
|
|
|
+ .get(&conn_id)
|
|
|
+ .ok_or(Error::UnknownConnection(conn_id))?;
|
|
|
+
|
|
|
+ match request {
|
|
|
+ Request::Event(event) => {
|
|
|
+ self.store_and_broadcast(event.deref());
|
|
|
+ }
|
|
|
+ Request::Request(request) => {
|
|
|
+ for filter in request.filters.into_iter() {
|
|
|
+ // Create subscription
|
|
|
+ let (conn_id, sub_id, receiver) =
|
|
|
+ connection.create_connection(request.subscription_id.deref().to_owned())?;
|
|
|
+ let mut subscriptions = self.subscriptions.write();
|
|
|
+ Self::get_indexes_from_filter(&filter)
|
|
|
+ .into_iter()
|
|
|
+ .for_each(|index| {
|
|
|
+ subscriptions
|
|
|
+ .entry(index)
|
|
|
+ .or_insert_with(|| RwLock::new(HashMap::new()))
|
|
|
+ .write()
|
|
|
+ .insert(
|
|
|
+ (conn_id, sub_id),
|
|
|
+ (request.subscription_id.clone(), receiver.clone()),
|
|
|
+ );
|
|
|
+ });
|
|
|
+
|
|
|
+ // Sent all events that match the filter that are stored in our database
|
|
|
+ self.storage
|
|
|
+ .get_by_filter(filter)?
|
|
|
+ .into_iter()
|
|
|
+ .for_each(|event| {
|
|
|
+ let _ = connection.send(
|
|
|
+ relayer::Event {
|
|
|
+ subscription_id: request.subscription_id.clone(),
|
|
|
+ event,
|
|
|
+ }
|
|
|
+ .into(),
|
|
|
+ );
|
|
|
+ });
|
|
|
+ }
|
|
|
+ }
|
|
|
+ Request::Close(_close) => {}
|
|
|
+ };
|
|
|
+ Ok(())
|
|
|
+ }
|
|
|
+
|
|
|
+ #[inline]
|
|
|
+ fn get_indexes_from_filter(filter: &Filter) -> Vec<SubscriptionType> {
|
|
|
+ let public_keys = if filter.references_to_public_key.is_empty() {
|
|
|
+ vec![None]
|
|
|
+ } else {
|
|
|
+ filter
|
|
|
+ .references_to_public_key
|
|
|
+ .iter()
|
|
|
+ .map(|public_key| Some((*public_key).to_vec()))
|
|
|
+ .collect()
|
|
|
+ };
|
|
|
+ let id = if filter.references_to_event.is_empty() {
|
|
|
+ vec![None]
|
|
|
+ } else {
|
|
|
+ filter
|
|
|
+ .references_to_event
|
|
|
+ .iter()
|
|
|
+ .map(|id| Some((*id).to_vec()))
|
|
|
+ .collect()
|
|
|
+ };
|
|
|
+ let kind = if filter.kinds.is_empty() {
|
|
|
+ vec![None]
|
|
|
+ } else {
|
|
|
+ filter
|
|
|
+ .kinds
|
|
|
+ .iter()
|
|
|
+ .map(|kind| Some((*kind).into()))
|
|
|
+ .collect()
|
|
|
+ };
|
|
|
+ let mut subs = vec![];
|
|
|
+
|
|
|
+ for public_key in public_keys.iter() {
|
|
|
+ for id in id.iter() {
|
|
|
+ for kind in kind.iter() {
|
|
|
+ subs.push(SubscriptionType {
|
|
|
+ public_key: public_key.clone(),
|
|
|
+ id: id.clone(),
|
|
|
+ kind: *kind,
|
|
|
+ });
|
|
|
+ }
|
|
|
+ }
|
|
|
+ }
|
|
|
+
|
|
|
+ subs
|
|
|
+ }
|
|
|
+
|
|
|
+ #[inline]
|
|
|
+ fn get_possible_listeners_from_event(event: &Event) -> Vec<SubscriptionType> {
|
|
|
+ let kind = event.kind().into();
|
|
|
+ let public_keys = [None, Some(event.author().as_ref().to_vec())];
|
|
|
+ let id = [None, Some(event.id.as_ref().to_vec())];
|
|
|
+ let kind = [None, Some(kind)];
|
|
|
+ let mut subs = vec![];
|
|
|
+
|
|
|
+ for public_key in public_keys.iter() {
|
|
|
+ for id in id.iter() {
|
|
|
+ for kind in kind.iter() {
|
|
|
+ subs.push(SubscriptionType {
|
|
|
+ public_key: public_key.clone(),
|
|
|
+ id: id.clone(),
|
|
|
+ kind: *kind,
|
|
|
+ });
|
|
|
+ }
|
|
|
+ }
|
|
|
+ }
|
|
|
+
|
|
|
+ subs
|
|
|
+ }
|
|
|
+
|
|
|
+ #[inline]
|
|
|
+ fn broadcast_to_subscribers(subscriptions: RwLockReadGuard<Subscriptions>, event: &Event) {
|
|
|
+ for (_, receiver) in subscriptions.iter() {
|
|
|
+ let _ = receiver.1.try_send(
|
|
|
+ relayer::Event {
|
|
|
+ subscription_id: receiver.0.clone(),
|
|
|
+ event: event.clone(),
|
|
|
+ }
|
|
|
+ .into(),
|
|
|
+ );
|
|
|
+ }
|
|
|
+ }
|
|
|
+
|
|
|
+ #[inline]
|
|
|
+ pub fn store_and_broadcast(&self, event: &Event) {
|
|
|
+ let _ = self.storage.store(event);
|
|
|
+ let subscriptions = self.subscriptions.read();
|
|
|
+
|
|
|
+ for subscription_type in Self::get_possible_listeners_from_event(event) {
|
|
|
+ if let Some(subscribers) = subscriptions.get(&subscription_type) {
|
|
|
+ Self::broadcast_to_subscribers(subscribers.read(), event);
|
|
|
+ }
|
|
|
+ }
|
|
|
+ }
|
|
|
+}
|