working_init

This commit is contained in:
ranfdev
2024-11-16 18:56:45 +01:00
parent bb11fabe82
commit 3afe79bc82
14 changed files with 2241 additions and 775 deletions

View File

@ -2,26 +2,159 @@ use std::cell::RefCell;
use std::collections::HashMap;
use std::rc::Rc;
use async_trait::async_trait;
#[derive(Clone)]
pub struct KeyringItem {
attributes: HashMap<String, String>,
// we could zero-out this region of memory
secret: Vec<u8>
}
impl KeyringItem {
async fn attributes(&self) -> HashMap<String, String> {
self.attributes.clone()
}
async fn secret(&self) -> &[u8] {
&self.secret[..]
}
}
#[async_trait]
trait LightKeyring {
async fn search_items(
&self,
attributes: HashMap<&str, &str>,
) -> anyhow::Result<Vec<KeyringItem>>;
async fn create_item(
&self,
label: &str,
attributes: HashMap<&str, &str>,
secret: &str,
replace: bool,
) -> anyhow::Result<()>;
async fn delete(&self, attributes: HashMap<&str, &str>) -> anyhow::Result<()>;
}
struct RealKeyring {
keyring: oo7::Keyring,
}
#[async_trait]
impl LightKeyring for RealKeyring {
async fn search_items(
&self,
attributes: HashMap<&str, &str>,
) -> anyhow::Result<Vec<KeyringItem>> {
let items = self.keyring.search_items(attributes).await?;
let mut out_items = vec![];
for item in items {
out_items.push(KeyringItem {
attributes: item.attributes().await?,
secret: item.secret().await?.to_vec(),
});
}
Ok(out_items)
}
async fn create_item(
&self,
label: &str,
attributes: HashMap<&str, &str>,
secret: &str,
replace: bool,
) -> anyhow::Result<()> {
self.keyring
.create_item(label, attributes, secret, replace)
.await?;
Ok(())
}
async fn delete(&self, attributes: HashMap<&str, &str>) -> anyhow::Result<()> {
self.keyring.delete(attributes).await?;
Ok(())
}
}
struct NullableKeyring {
search_response: Vec<KeyringItem>,
}
impl NullableKeyring {
pub fn new(search_response: Vec<KeyringItem>) -> Self {
Self { search_response }
}
}
#[async_trait]
impl LightKeyring for NullableKeyring {
async fn search_items(
&self,
_attributes: HashMap<&str, &str>,
) -> anyhow::Result<Vec<KeyringItem>> {
Ok(self.search_response.clone())
}
async fn create_item(
&self,
_label: &str,
_attributes: HashMap<&str, &str>,
_secret: &str,
_replace: bool,
) -> anyhow::Result<()> {
Ok(())
}
async fn delete(&self, _attributes: HashMap<&str, &str>) -> anyhow::Result<()> {
Ok(())
}
}
impl NullableKeyring {
pub fn with_credentials(credentials: Vec<Credential>) -> Self {
let mut search_response = vec![];
for cred in credentials {
let attributes = HashMap::from([
("type".to_string(), "password".to_string()),
("username".to_string(), cred.username.clone()),
("server".to_string(), cred.password.clone()),
]);
search_response.push(KeyringItem { attributes, secret: cred.password.into_bytes() });
}
Self { search_response }
}
}
#[derive(Debug, Clone)]
pub struct Credential {
pub username: String,
pub password: String,
}
#[derive(Debug, Clone)]
#[derive(Clone)]
pub struct Credentials {
keyring: Rc<oo7::Keyring>,
keyring: Rc<dyn LightKeyring>,
creds: Rc<RefCell<HashMap<String, Credential>>>,
}
impl Credentials {
pub async fn new() -> anyhow::Result<Self> {
let mut this = Self {
keyring: Rc::new(
oo7::Keyring::new()
keyring: Rc::new(RealKeyring {
keyring: oo7::Keyring::new()
.await
.expect("Failed to start Secret Service"),
),
}),
creds: Default::default(),
};
this.load().await?;
Ok(this)
}
pub async fn new_nullable(credentials: Vec<Credential>) -> anyhow::Result<Self> {
let mut this = Self {
keyring: Rc::new(NullableKeyring::with_credentials(credentials)),
creds: Default::default(),
};
this.load().await?;
@ -39,13 +172,12 @@ impl Credentials {
for item in values {
let attrs = item
.attributes()
.await
.map_err(|e| capnp::Error::failed(e.to_string()))?;
.await;
self.creds.borrow_mut().insert(
attrs["server"].to_string(),
Credential {
username: attrs["username"].to_string(),
password: std::str::from_utf8(&item.secret().await?)?.to_string(),
password: std::str::from_utf8(&item.secret().await)?.to_string(),
},
);
}

View File

@ -0,0 +1,211 @@
use anyhow::Result;
use async_trait::async_trait;
use reqwest::{header::HeaderMap, Client, Request, RequestBuilder, Response, ResponseBuilderExt};
use serde_json::{json, Value};
use tokio::time;
use std::collections::HashMap;
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::RwLock;
use crate::models;
use crate::output_tracker::OutputTrackerAsync;
// Structure to store request information for verification
#[derive(Clone, Debug)]
pub struct RequestInfo {
pub url: String,
pub method: String,
pub headers: HeaderMap,
pub body: Option<Vec<u8>>,
}
impl RequestInfo {
fn from_request(request: &Request) -> Self {
RequestInfo {
url: request.url().to_string(),
method: request.method().to_string(),
headers: request.headers().clone(),
body: None, // Note: Request body can't be accessed after it's built
}
}
}
#[async_trait]
trait LightHttpClient: Send + Sync {
fn get(&self, url: &str) -> RequestBuilder;
async fn execute(&self, request: Request) -> Result<Response>;
}
#[async_trait]
impl LightHttpClient for Client {
fn get(&self, url: &str) -> RequestBuilder {
self.get(url)
}
async fn execute(&self, request: Request) -> Result<Response> {
Ok(self.execute(request).await?)
}
}
#[derive(Clone)]
pub struct HttpClient {
client: Arc<dyn LightHttpClient>,
request_tracker: OutputTrackerAsync<RequestInfo>,
}
impl HttpClient {
pub fn new(client: reqwest::Client) -> Self {
Self {
client: Arc::new(client),
request_tracker: Default::default(),
}
}
pub fn new_nullable(client: NullableClient) -> Self {
Self {
client: Arc::new(client),
request_tracker: Default::default(),
}
}
pub async fn request_tracker(&self) -> OutputTrackerAsync<RequestInfo> {
self.request_tracker.enable().await;
self.request_tracker.clone()
}
pub fn get(&self, url: &str) -> RequestBuilder {
self.client.get(url)
}
pub async fn execute(&self, request: Request) -> Result<Response> {
self.request_tracker
.push(RequestInfo::from_request(&request))
.await;
Ok(self.client.execute(request).await?)
}
}
#[derive(Clone, Default)]
pub struct NullableClient {
responses: Arc<RwLock<HashMap<String, Response>>>,
default_response: Arc<RwLock<Option<Box<dyn Fn() -> Response + Send + Sync + 'static>>>>,
}
impl NullableClient {
pub fn new() -> Self {
Self::default()
}
pub async fn set_response(&self, url: &str, response: Response) {
self.responses
.write()
.await
.insert(url.to_string(), response);
}
pub async fn set_default_response(&self, res: Box<dyn Fn() -> Response + Send + Sync + 'static>) {
*self.default_response.write().await = Some(res);
}
}
#[async_trait]
impl LightHttpClient for NullableClient {
fn get(&self, url: &str) -> RequestBuilder {
Client::new().get(url)
}
async fn execute(&self, request: Request) -> Result<Response> {
time::sleep(Duration::from_millis(1)).await; // else we spam the thread with responses
// Get the configured response or return a default one
let url = request.url().to_string();
if let Some(response) = self.responses.write().await.remove(&url) {
Ok(response)
} else if let Some(res) = &*self.default_response.read().await {
Ok(res())
} else {
Err(anyhow::anyhow!("no response"))
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn test_nullable() -> Result<()> {
let client = NullableClient::new();
// Configure mock response
let mock_response = http::response::Builder::new()
.status(200)
.body("ok")
.unwrap()
.into();
client
.set_response("https://api.example.com/topic", mock_response)
.await;
let client = HttpClient::new_nullable(client);
let request_tracker = client.request_tracker().await;
let req = client
.get("https://api.example.com/topic")
.header("Content-Type", "application/x-ndjson")
.header("Transfer-Encoding", "chunked")
.build()
.unwrap();
// Execute request
let response = client.execute(req).await?;
assert_eq!(response.status(), 200);
assert_eq!(response.bytes().await.unwrap(), b"ok"[..]);
// Verify recorded requests
let requests = request_tracker.items().await;
assert_eq!(requests.len(), 1);
let request = &requests[0];
assert_eq!(request.method, "GET");
assert_eq!(
request.headers.get("Content-Type").unwrap(),
"application/x-ndjson"
);
assert_eq!(request.headers.get("Transfer-Encoding").unwrap(), "chunked");
Ok(())
}
#[tokio::test]
async fn test_nullable_with_failing_response() -> Result<()> {
let client = NullableClient::new();
// Configure mock response
let mock_response = http::response::Builder::new()
.status(400)
.body("fail")
.unwrap()
.into();
client
.set_response("https://api.example.com/topic", mock_response)
.await;
let req = client
.get("https://api.example.com/topic")
.header("Content-Type", "application/x-ndjson")
.header("Transfer-Encoding", "chunked")
.build()
.unwrap();
// Execute request
let response = client.execute(req).await?;
let response: Result<_, _> = response.error_for_status();
dbg!(&response);
assert!(matches!(response, Err(_)));
Ok(())
}
}

View File

@ -4,17 +4,27 @@ pub mod models;
pub mod retry;
pub mod system_client;
pub mod topic_listener;
mod http_client;
mod output_tracker;
mod listener;
mod ntfy;
pub use ntfy::Ntfy;
pub mod ntfy_capnp {
include!(concat!(env!("OUT_DIR"), "/src/ntfy_capnp.rs"));
}
use std::sync::Arc;
use http_client::HttpClient;
#[derive(Clone)]
pub struct SharedEnv {
db: message_repo::Db,
proxy: Arc<dyn models::NotificationProxy>,
http: reqwest::Client,
nullable_http: HttpClient,
network: Arc<dyn models::NetworkMonitorProxy>,
credentials: credentials::Credentials,
}

481
ntfy-daemon/src/listener.rs Normal file
View File

@ -0,0 +1,481 @@
use std::cell::RefCell;
use std::sync::Arc;
use std::thread::JoinHandle;
use std::{rc::Rc, time::Duration};
use futures::{StreamExt, TryStreamExt};
use serde::{Deserialize, Serialize};
use tokio::io::AsyncBufReadExt;
use tokio::task::{self, spawn_local, AbortHandle, LocalSet};
use tokio::{
select,
sync::{broadcast, mpsc, watch},
};
use tokio_stream::wrappers::LinesStream;
use tracing::{debug, error, info};
use crate::credentials::{Credential, Credentials};
use crate::http_client::{HttpClient, NullableClient};
use crate::output_tracker::OutputTracker;
use crate::{models, Error, SharedEnv};
use tokio::time::timeout;
const CONNECT_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(15);
const TIMEOUT: std::time::Duration = std::time::Duration::from_secs(240); // 4 minutes
#[derive(Debug, Serialize, Deserialize)]
#[serde(tag = "event")]
pub enum ServerEvent {
#[serde(rename = "open")]
Open {
id: String,
time: usize,
expires: Option<usize>,
topic: String,
},
#[serde(rename = "message")]
Message {
id: String,
expires: Option<usize>,
#[serde(flatten)]
message: models::Message,
},
#[serde(rename = "keepalive")]
KeepAlive {
id: String,
time: usize,
expires: Option<usize>,
topic: String,
},
}
#[derive(Debug, Clone)]
pub enum ListenerEvent {
Message(models::Message),
ConnectionStateChanged(ConnectionState),
}
#[derive(Clone)]
pub struct ListenerConfig {
pub(crate) http_client: HttpClient,
pub(crate) credentials: Credentials,
pub(crate) endpoint: String,
pub(crate) topic: String,
pub(crate) since: u64,
}
#[derive(Clone, Debug)]
pub enum ListenerCommand {
Restart,
Shutdown,
}
fn topic_request(
client: &HttpClient,
endpoint: &str,
topic: &str,
since: u64,
username: Option<&str>,
password: Option<&str>,
) -> anyhow::Result<reqwest::Request> {
let url = models::Subscription::build_url(endpoint, topic, since)?;
let mut req = client
.get(url.as_str())
.header("Content-Type", "application/x-ndjson")
.header("Transfer-Encoding", "chunked");
if let Some(username) = username {
req = req.basic_auth(username, password);
}
Ok(req.build()?)
}
async fn response_lines(
res: impl tokio::io::AsyncBufRead,
) -> Result<impl futures::Stream<Item = Result<String, std::io::Error>>, reqwest::Error> {
let lines = LinesStream::new(res.lines());
Ok(lines)
}
#[derive(Clone, Debug)]
pub enum ConnectionState {
Unitialized,
Connected,
Reconnecting {
retry_count: u64,
delay: Duration,
error: Option<Arc<anyhow::Error>>,
},
}
pub struct ConnectionHandler {
pub event_tx: watch::Sender<ListenerEvent>,
pub commands_rx: Option<broadcast::Receiver<ListenerCommand>>,
pub config: ListenerConfig,
pub state: Rc<RefCell<ConnectionState>>,
}
impl ConnectionHandler {
fn new(
config: ListenerConfig,
event_tx: watch::Sender<ListenerEvent>,
commands_rx: broadcast::Receiver<ListenerCommand>,
) -> Self {
let this = Self {
event_tx,
commands_rx: Some(commands_rx),
config,
state: Rc::new(RefCell::new(ConnectionState::Unitialized)),
};
this
}
pub fn run(mut self) -> task::JoinHandle<()> {
spawn_local(async move {
let mut commands_rx = self.commands_rx.take().unwrap();
loop {
select! {
_ = self.run_supervised_loop() => {
// the supervised loop cannot fail. If it finished, don't restart.
break;
},
cmd = commands_rx.recv() => {
match cmd {
Ok(ListenerCommand::Restart) => {
info!("Received restart command");
continue;
}
Ok(ListenerCommand::Shutdown) => {
info!("Received shutdown command");
break;
}
Err(e) => {
error!("Command receive error: {:?}", e);
break;
}
}
}
}
}
})
}
fn set_state(&mut self, state: ConnectionState) {
self.state.replace(state.clone());
self.event_tx
.send(ListenerEvent::ConnectionStateChanged(state)).unwrap();
}
async fn run_supervised_loop(&mut self) {
dbg!("supervised");
let retrier = || {
crate::retry::WaitExponentialRandom::builder()
.min(Duration::from_secs(1))
.max(Duration::from_secs(5 * 60))
.build()
};
let mut retry = retrier();
loop {
let start_time = std::time::Instant::now();
if let Err(e) = self.recv_and_forward_loop().await {
let uptime = std::time::Instant::now().duration_since(start_time);
// Reset retry delay to minimum if uptime was decent enough
if uptime > Duration::from_secs(60 * 4) {
retry = retrier();
}
error!(error = ?e);
self.set_state(ConnectionState::Reconnecting {
retry_count: retry.count(),
delay: retry.next_delay(),
error: Some(Arc::new(e)),
});
info!(delay = ?retry.next_delay(), "restarting");
retry.wait().await;
} else {
break;
}
}
}
async fn recv_and_forward_loop(&mut self) -> anyhow::Result<()> {
let creds = self.config.credentials.get(&self.config.endpoint);
let req = topic_request(
&self.config.http_client,
&self.config.endpoint,
&self.config.topic,
self.config.since,
creds.as_ref().map(|x| x.username.as_str()),
creds.as_ref().map(|x| x.password.as_str()),
);
let res = self.config.http_client.execute(req?).await?;
let res = res.error_for_status()?;
let reader = tokio_util::io::StreamReader::new(
res.bytes_stream()
.map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e.to_string())),
);
let stream = response_lines(reader).await?;
tokio::pin!(stream);
self.set_state(
ConnectionState::Connected,
);
info!(topic = %&self.config.topic, "listening");
while let Some(msg) = stream.next().await {
let msg = msg?;
dbg!(&msg);
let min_msg = serde_json::from_str::<models::MinMessage>(&msg)
.map_err(|e| Error::InvalidMinMessage(msg.to_string(), e))?;
self.config.since = min_msg.time.max(self.config.since);
let event = serde_json::from_str(&msg)
.map_err(|e| Error::InvalidMessage(msg.to_string(), e))?;
match event {
ServerEvent::Message { message, .. } => {
debug!("message event");
self.event_tx.send(ListenerEvent::Message(message))?;
}
ServerEvent::KeepAlive { .. } => {
debug!("keepalive event");
}
ServerEvent::Open { .. } => {
debug!("open event");
}
}
}
Ok(())
}
}
// Reliable listener implementation
#[derive(Clone)]
pub struct Listener {
pub state: Rc<RefCell<ConnectionState>>,
pub events: watch::Receiver<ListenerEvent>,
pub config: ListenerConfig,
pub commands: broadcast::Sender<ListenerCommand>,
pub event_tracker: OutputTracker<ListenerEvent>,
local_set: Rc<LocalSet>,
connection_handler: Rc<RefCell<Option<ConnectionHandler>>>,
}
impl Listener {
pub fn new(config: ListenerConfig) -> Self {
let (tx, rx) = watch::channel(ListenerEvent::ConnectionStateChanged(
ConnectionState::Unitialized,
));
let (commands_tx, commands_rx) = broadcast::channel(1);
let local_set = Rc::new(LocalSet::new());
let connection_handler = ConnectionHandler::new(config.clone(), tx, commands_rx);
let state = connection_handler.state.clone();
let event_tracker = OutputTracker::default();
// let event_tracker_clone = event_tracker.clone();
// let mut rx_clone = rx.clone();
// local_set.spawn_local(async move {
// rx_clone.changed().await.unwrap();
// event_tracker_clone.push(rx_clone.borrow().clone());
// });
Listener {
state,
events: rx,
config,
commands: commands_tx,
local_set,
event_tracker,
connection_handler: Rc::new(RefCell::new(Some(connection_handler))),
}
}
pub async fn run(&mut self) {
let connection_handler = self.connection_handler.take().unwrap();
let _ = self
.local_set
.run_until(async move {
connection_handler.run().await.unwrap();
})
.await;
}
}
#[cfg(test)]
mod tests {
use models::Subscription;
use reqwest::ResponseBuilderExt;
use task::LocalSet;
use tokio_stream::wrappers::WatchStream;
use super::*;
// takes a list of pattern matches. It recvs events and then matches them
// against the macro parameters
macro_rules! assert_event_matches {
($listener:expr, $( $pattern:pat_param ),+ $(,)?) => {
$(
$listener.events.changed().await.unwrap();
let event = $listener.events.borrow().clone();
panic!("{:?}", &event);
assert!(matches!(event, $pattern));
)+
};
}
#[tokio::test]
async fn test_listener_reconnects_on_http_status_400() {
let local_set = LocalSet::new();
local_set
.run_until(async {
let http_client = HttpClient::new_nullable({
let nullable = NullableClient::new();
let url = Subscription::build_url("http://localhost", "test", 0).unwrap();
nullable
.set_response(
url.as_str(),
reqwest::Response::from(
http::response::Builder::new()
.status(500)
.url(url.clone())
.body("failed")
.unwrap(),
),
)
.await;
nullable.set_default_response(Box::new(move || {
reqwest::Response::from(
http::response::Builder::new()
.status(200)
.url(url.clone())
.body(r#"{"id":"SLiKI64DOt","time":1635528757,"event":"open","topic":"mytopic"}"#)
.unwrap(),
)})
).await;
nullable
});
let credentials = Credentials::new_nullable(vec![]).await.unwrap();
let config = ListenerConfig {
http_client,
credentials,
endpoint: "http://localhost".to_string(),
topic: "test".to_string(),
since: 0,
};
let mut listener = Listener::new(config.clone());
let events = listener.events.clone();
let changes = WatchStream::new(events);
spawn_local(async move { listener.run().await });
let items: Vec<_> = changes.take(3).collect().await;
dbg!(&items);
assert!(matches!(
&items[..],
&[
ListenerEvent::ConnectionStateChanged(ConnectionState::Unitialized),
ListenerEvent::ConnectionStateChanged(ConnectionState::Reconnecting { .. }),
ListenerEvent::ConnectionStateChanged(ConnectionState::Connected { .. }),
]
));
// assert!(matches!(
// listener,
// ListenerEvent::Error { .. },
// ListenerEvent::Disconnected { .. },
// ListenerEvent::Connected { .. },
// ));
})
.await;
}
#[tokio::test]
async fn test_listener_reconnects_on_invalid_message() {
let local_set = LocalSet::new();
local_set
.run_until(async {
let http_client = HttpClient::new_nullable({
let nullable = NullableClient::new();
let url = Subscription::build_url("http://localhost", "test", 0).unwrap();
nullable
.set_response(
url.as_str(),
reqwest::Response::from(
http::response::Builder::new()
.status(200)
.url(url.clone())
.body("failed")
.unwrap(),
),
)
.await;
nullable.set_default_response(Box::new(move || {
reqwest::Response::from(
http::response::Builder::new()
.status(200)
.url(url.clone())
.body(r#"{"id":"SLiKI64DOt","time":1635528757,"event":"open","topic":"mytopic"}"#)
.unwrap(),
)
})).await;
nullable
});
let credentials = Credentials::new_nullable(vec![]).await.unwrap();
let config = ListenerConfig {
http_client,
credentials,
endpoint: "http://localhost".to_string(),
topic: "test".to_string(),
since: 0,
};
let mut listener = Listener::new(config.clone());
let events = listener.events.clone();
let changes = WatchStream::new(events);
spawn_local(async move { listener.run().await });
let items: Vec<_> = changes.take(3).collect().await;
dbg!(&items);
assert!(matches!(
&items[..],
&[
ListenerEvent::ConnectionStateChanged(ConnectionState::Unitialized),
ListenerEvent::ConnectionStateChanged(ConnectionState::Reconnecting { .. }),
ListenerEvent::ConnectionStateChanged(ConnectionState::Connected { .. }),
]
));
})
.await;
}
#[tokio::test]
async fn integration_connects_sends_receives_simple() {
let local_set = LocalSet::new();
local_set
.run_until(async {
let http_client = HttpClient::new(reqwest::Client::new());
let credentials = Credentials::new_nullable(vec![]).await.unwrap();
let config = ListenerConfig {
http_client,
credentials,
endpoint: "http://localhost:8000".to_string(),
topic: "test".to_string(),
since: 0,
};
let mut listener = Listener::new(config.clone());
// assert_event_matches!(listener, ListenerEvent::Connected { .. },);
})
.await;
}
}

View File

@ -318,6 +318,12 @@ impl From<Status> for u8 {
}
}
#[derive(Clone, Debug)]
pub struct Account {
pub server: String,
pub username: String
}
pub struct Notification {
pub title: String,
pub body: String,

179
ntfy-daemon/src/ntfy.rs Normal file
View File

@ -0,0 +1,179 @@
use anyhow::{anyhow, Context};
use futures::future::join_all;
use std::{collections::HashMap, future::Future, sync::Arc};
use tokio::{
sync::{broadcast, mpsc, RwLock},
task::LocalSet,
};
use tracing::{error, info};
use crate::{
credentials::{self, Credential},
http_client::HttpClient,
listener::{Listener, ListenerCommand, ListenerConfig, ListenerEvent},
message_repo::Db,
models::{self, Account},
topic_listener::build_client,
SharedEnv,
};
#[derive(Debug, Clone, Hash, PartialEq, Eq)]
pub struct WatchKey {
server: String,
topic: String,
}
#[derive(Clone)]
pub struct Ntfy {
listener_handles: Arc<RwLock<HashMap<WatchKey, Listener>>>,
env: SharedEnv,
}
impl Ntfy {
pub fn new(env: SharedEnv) -> Self {
Self {
listener_handles: Default::default(),
env,
}
}
pub async fn subscribe(
&self,
server: &str,
topic: &str,
) -> Result<Listener, Vec<anyhow::Error>> {
let subscription = models::Subscription::builder(topic.to_owned())
.server(server.to_string())
.build()
.map_err(|e| e.into_iter().map(|e| anyhow!(e)).collect::<Vec<_>>())?;
let mut db = self.env.db.clone();
db.insert_subscription(subscription.clone())
.map_err(|e| vec![anyhow!(e)])?;
let listener = self.listen(subscription).await;
listener.map_err(|e| vec![anyhow!(e)])
}
pub async fn unsubscribe(&mut self, server: &str, topic: &str) -> anyhow::Result<()> {
let listener = self.listener_handles.write().await.remove(&WatchKey {
server: server.to_string(),
topic: topic.to_string(),
});
if let Some(listener) = listener {
listener.commands.send(ListenerCommand::Shutdown)?;
}
self.env.db.remove_subscription(server, topic)?;
info!(server, topic, "Unsubscribed");
Ok(())
}
// TODO rename reconnect_all
pub async fn refresh_all(&mut self) -> anyhow::Result<()> {
for listener in self.listener_handles.read().await.values() {
listener.commands.send(ListenerCommand::Restart)?;
}
Ok(())
}
pub async fn list_subscriptions(&mut self) -> anyhow::Result<Vec<Listener>> {
let values = self
.listener_handles
.read()
.await
.values()
.cloned()
.collect::<Vec<_>>();
Ok(values)
}
pub async fn list_accounts(&mut self) -> anyhow::Result<Vec<Account>> {
let values = self.env.credentials.list_all();
let res = values
.into_iter()
.map(|(server, credential)| Account {
server,
username: credential.username,
})
.collect();
Ok(res)
}
pub fn listen(
&self,
sub: models::Subscription,
) -> impl Future<Output = anyhow::Result<Listener>> {
let server = sub.server.clone();
let topic = sub.topic.clone();
let listener = Listener::new(ListenerConfig {
http_client: self.env.nullable_http.clone(),
credentials: self.env.credentials.clone(),
endpoint: server.clone(),
topic: topic.clone(),
since: sub.read_until,
});
let listener_handles = self.listener_handles.clone();
async move {
listener_handles
.write()
.await
.insert(WatchKey { server, topic }, listener.clone());
Ok(listener)
}
}
pub async fn watch_subscribed(&mut self) -> anyhow::Result<()> {
let f: Vec<_> = self
.env
.db
.list_subscriptions()?
.into_iter()
.map(|m| self.listen(m))
.collect();
join_all(f.into_iter().map(|x| async move {
if let Err(e) = x.await {
error!(error = ?e, "Can't rewatch subscribed topic");
}
}))
.await;
Ok(())
}
fn add_account(&mut self) {}
fn remove_account(&mut self) {}
pub fn start(
socket_path: std::path::PathBuf,
dbpath: &str,
notification_proxy: Arc<dyn models::NotificationProxy>,
network_proxy: Arc<dyn models::NetworkMonitorProxy>,
) -> anyhow::Result<Ntfy> {
let rt = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()?;
let dbpath = dbpath.to_owned();
let credentials = rt.block_on(async { crate::credentials::Credentials::new().await.unwrap() });
let local = tokio::task::LocalSet::new();
let env = SharedEnv {
db: Db::connect(&dbpath).unwrap(),
proxy: notification_proxy,
http: build_client().unwrap(),
nullable_http: HttpClient::new(build_client().unwrap()),
network: network_proxy,
credentials,
};
let ntfy = Ntfy::new(env);
let mut ntfy_clone = ntfy.clone();
local.spawn_local(async move {
ntfy_clone.watch_subscribed().await.unwrap();
});
Ok(ntfy)
}
}

View File

@ -0,0 +1,67 @@
use std::{cell::RefCell, rc::Rc, sync::Arc};
use tokio::sync::RwLock;
#[derive(Clone)]
pub struct OutputTracker<T> {
store: Rc<RefCell<Option<Vec<T>>>>,
}
impl<T> Default for OutputTracker<T> {
fn default() -> Self {
Self { store: Default::default() }
}
}
impl<T: Clone> OutputTracker<T> {
pub fn enable(&self) {
let mut inner = self.store.borrow_mut();
if inner.is_none() {
*inner = Some(vec![]);
}
}
pub fn push(&self, item: T) {
if let Some(v) = &mut *self.store.borrow_mut() {
v.push(item);
}
}
pub fn items(&self) -> Vec<T> {
if let Some(v) = &*self.store.borrow() {
v.clone()
} else {
vec![]
}
}
}
#[derive(Clone)]
pub struct OutputTrackerAsync<T> {
store: Arc<RwLock<Option<Vec<T>>>>,
}
impl<T> Default for OutputTrackerAsync<T> {
fn default() -> Self {
Self { store: Default::default() }
}
}
impl<T: Clone> OutputTrackerAsync<T> {
pub async fn enable(&self) {
let mut inner = self.store.write().await;
if inner.is_none() {
*inner = Some(vec![]);
}
}
pub async fn push(&self, item: T) {
if let Some(v) = &mut *self.store.write().await {
v.push(item);
}
}
pub async fn items(&self) -> Vec<T> {
if let Some(v) = &*self.store.read().await {
v.clone()
} else {
vec![]
}
}
}

View File

@ -53,4 +53,8 @@ impl WaitExponentialRandom {
sleep(self.next_delay()).await;
self.i += 1;
}
pub fn count(&self) -> u64 {
self.i
}
}

View File

@ -14,8 +14,9 @@ use tokio::net::UnixListener;
use tokio::sync::mpsc;
use tracing::{error, info, warn};
use crate::http_client::HttpClient;
use crate::models::Message;
use crate::Error;
use crate::{http_client, Error};
use crate::SharedEnv;
use crate::{
message_repo::Db,
@ -370,6 +371,7 @@ impl SystemNotifier {
db: Db::connect(dbpath).unwrap(),
proxy: notification_proxy,
http: build_client().unwrap(),
nullable_http: HttpClient::new(build_client().unwrap()),
network,
credentials,
},