2020-04-13 22:03:06 +02:00
|
|
|
use flodgatt::config;
|
2020-04-24 01:28:26 +02:00
|
|
|
use flodgatt::request::{Handler, Subscription};
|
|
|
|
use flodgatt::response::{RedisManager, SseStream, WsStream};
|
2020-04-22 20:38:22 +02:00
|
|
|
use flodgatt::Error;
|
2019-02-19 20:29:32 +01:00
|
|
|
|
2020-04-13 22:03:06 +02:00
|
|
|
use futures::{future::lazy, stream::Stream as _};
|
|
|
|
use std::fs;
|
|
|
|
use std::net::SocketAddr;
|
|
|
|
use std::os::unix::fs::PermissionsExt;
|
|
|
|
use std::time::Instant;
|
|
|
|
use tokio::net::UnixListener;
|
2020-04-24 01:28:26 +02:00
|
|
|
use tokio::sync::mpsc;
|
2020-04-13 22:03:06 +02:00
|
|
|
use tokio::timer::Interval;
|
|
|
|
use warp::ws::Ws2;
|
|
|
|
use warp::Filter;
|
2020-01-06 03:58:18 +01:00
|
|
|
|
2020-04-22 20:38:22 +02:00
|
|
|
fn main() -> Result<(), Error> {
|
2020-04-13 22:03:06 +02:00
|
|
|
config::merge_dotenv()?;
|
|
|
|
pretty_env_logger::try_init()?;
|
2020-04-15 02:37:49 +02:00
|
|
|
let (postgres_cfg, redis_cfg, cfg) = config::from_env(dotenv::vars().collect())?;
|
|
|
|
let poll_freq = *redis_cfg.polling_interval;
|
2019-10-09 02:35:26 +02:00
|
|
|
|
2020-04-15 02:37:49 +02:00
|
|
|
let request = Handler::new(&postgres_cfg, *cfg.whitelist_mode)?;
|
2020-04-24 01:28:26 +02:00
|
|
|
let shared_manager = RedisManager::try_from(&redis_cfg)?.into_arc();
|
2019-07-10 04:20:11 +02:00
|
|
|
|
2019-07-06 02:08:50 +02:00
|
|
|
// Server Sent Events
|
2020-04-13 22:03:06 +02:00
|
|
|
let sse_manager = shared_manager.clone();
|
|
|
|
let sse = request
|
|
|
|
.sse_subscription()
|
2019-09-05 03:48:29 +02:00
|
|
|
.and(warp::sse())
|
2020-04-13 22:03:06 +02:00
|
|
|
.map(move |subscription: Subscription, sse: warp::sse::Sse| {
|
|
|
|
log::info!("Incoming SSE request for {:?}", subscription.timeline);
|
2020-04-22 20:38:22 +02:00
|
|
|
let mut manager = sse_manager.lock().unwrap_or_else(RedisManager::recover);
|
2020-04-24 01:28:26 +02:00
|
|
|
let (event_tx, event_rx) = mpsc::unbounded_channel();
|
|
|
|
manager.subscribe(&subscription, event_tx);
|
|
|
|
let sse_stream = SseStream::new(subscription);
|
|
|
|
sse_stream.send_events(sse, event_rx)
|
2020-04-13 22:03:06 +02:00
|
|
|
})
|
2020-03-27 17:00:48 +01:00
|
|
|
.with(warp::reply::with::header("Connection", "keep-alive"));
|
2019-04-28 23:28:57 +02:00
|
|
|
|
2019-07-06 02:08:50 +02:00
|
|
|
// WebSocket
|
2020-04-13 22:03:06 +02:00
|
|
|
let ws_manager = shared_manager.clone();
|
|
|
|
let ws = request
|
|
|
|
.ws_subscription()
|
2019-09-09 19:06:24 +02:00
|
|
|
.and(warp::ws::ws2())
|
2020-03-27 17:00:48 +01:00
|
|
|
.map(move |subscription: Subscription, ws: Ws2| {
|
|
|
|
log::info!("Incoming websocket request for {:?}", subscription.timeline);
|
2020-04-22 20:38:22 +02:00
|
|
|
let mut manager = ws_manager.lock().unwrap_or_else(RedisManager::recover);
|
2020-04-24 01:28:26 +02:00
|
|
|
let (event_tx, event_rx) = mpsc::unbounded_channel();
|
|
|
|
manager.subscribe(&subscription, event_tx);
|
2020-04-13 22:03:06 +02:00
|
|
|
let token = subscription.access_token.clone().unwrap_or_default(); // token sent for security
|
2020-04-24 01:28:26 +02:00
|
|
|
let ws_stream = WsStream::new(subscription);
|
2019-09-11 06:13:45 +02:00
|
|
|
|
2020-04-24 01:28:26 +02:00
|
|
|
(
|
|
|
|
ws.on_upgrade(move |ws| ws_stream.send_to(ws, event_rx)),
|
|
|
|
token,
|
|
|
|
)
|
2020-03-27 17:00:48 +01:00
|
|
|
})
|
2019-07-04 16:57:15 +02:00
|
|
|
.map(|(reply, token)| warp::reply::with_header(reply, "sec-websocket-protocol", token));
|
2019-05-09 05:02:01 +02:00
|
|
|
|
2020-04-13 22:03:06 +02:00
|
|
|
#[cfg(feature = "stub_status")]
|
|
|
|
#[rustfmt::skip]
|
|
|
|
let status = {
|
|
|
|
let (r1, r3) = (shared_manager.clone(), shared_manager.clone());
|
|
|
|
request.health().map(|| "OK")
|
|
|
|
.or(request.status()
|
2020-04-22 20:38:22 +02:00
|
|
|
.map(move || r1.lock().unwrap_or_else(RedisManager::recover).count()))
|
2020-04-13 22:03:06 +02:00
|
|
|
.or(request.status_per_timeline()
|
2020-04-22 20:38:22 +02:00
|
|
|
.map(move || r3.lock().unwrap_or_else(RedisManager::recover).list()))
|
2020-04-13 22:03:06 +02:00
|
|
|
};
|
|
|
|
#[cfg(not(feature = "stub_status"))]
|
|
|
|
let status = request.health().map(|| "OK");
|
|
|
|
|
2019-10-04 00:02:23 +02:00
|
|
|
let cors = warp::cors()
|
|
|
|
.allow_any_origin()
|
|
|
|
.allow_methods(cfg.cors.allowed_methods)
|
|
|
|
.allow_headers(cfg.cors.allowed_headers);
|
|
|
|
|
2020-04-13 22:03:06 +02:00
|
|
|
let streaming_server = move || {
|
|
|
|
let manager = shared_manager.clone();
|
|
|
|
let stream = Interval::new(Instant::now(), poll_freq)
|
|
|
|
.map_err(|e| log::error!("{}", e))
|
2020-04-22 20:38:22 +02:00
|
|
|
.for_each(move |_| {
|
|
|
|
let mut manager = manager.lock().unwrap_or_else(RedisManager::recover);
|
|
|
|
manager.poll_broadcast().map_err(Error::log)
|
|
|
|
});
|
2020-04-17 23:07:10 +02:00
|
|
|
|
2020-04-13 22:03:06 +02:00
|
|
|
warp::spawn(lazy(move || stream));
|
|
|
|
warp::serve(ws.or(sse).with(cors).or(status).recover(Handler::err))
|
2020-04-05 16:54:42 +02:00
|
|
|
};
|
2020-01-07 23:27:46 +01:00
|
|
|
|
2020-01-08 15:51:25 +01:00
|
|
|
if let Some(socket) = &*cfg.unix_socket {
|
2020-03-20 01:54:23 +01:00
|
|
|
log::info!("Using Unix socket {}", socket);
|
2020-01-09 23:54:57 +01:00
|
|
|
fs::remove_file(socket).unwrap_or_default();
|
2020-04-15 02:37:49 +02:00
|
|
|
let incoming = UnixListener::bind(socket)?.incoming();
|
|
|
|
fs::set_permissions(socket, PermissionsExt::from_mode(0o666))?;
|
2020-04-13 22:03:06 +02:00
|
|
|
tokio::run(lazy(|| streaming_server().serve_incoming(incoming)));
|
2019-10-04 00:02:23 +02:00
|
|
|
} else {
|
2020-04-01 21:35:24 +02:00
|
|
|
let server_addr = SocketAddr::new(*cfg.address, *cfg.port);
|
2020-04-13 22:03:06 +02:00
|
|
|
tokio::run(lazy(move || streaming_server().bind(server_addr)));
|
|
|
|
}
|
2020-04-22 20:38:22 +02:00
|
|
|
Err(Error::Unrecoverable) // only get here if there's an unrecoverable error in poll_broadcast.
|
2019-02-11 09:45:14 +01:00
|
|
|
}
|