#![allow(dead_code)] #[macro_use] extern crate log; use chain::{ Chain, }; use warp::{ Filter, Buf, reply::Response, }; use hyper::Body; use std::{ sync::Arc, fmt, error, net::{ SocketAddr, IpAddr, }, }; use tokio::{ sync::{ RwLock, mpsc, Notify, }, stream::{Stream,StreamExt,}, }; use serde::{ Serialize, Deserialize }; use futures::{ future::{ FutureExt, BoxFuture, join_all, }, }; use lazy_static::lazy_static; use cfg_if::cfg_if; macro_rules! if_debug { ($($tt:tt)*) => { cfg_if::cfg_if!{ if #[cfg(debug_assertions)] { $($tt)* } } } } macro_rules! status { ($code:expr) => { ::warp::http::status::StatusCode::from_u16($code).unwrap() }; } mod ext; use ext::*; mod util; mod range; mod sanitise; mod bytes; mod chunking; #[cfg(feature="api")] mod api; #[cfg(target_family="unix")] mod signals; mod config; mod msg; mod state; use state::State; mod save; mod ipfilt; mod forwarded_list; use forwarded_list::XForwardedFor; mod handle; mod feed; mod gen; mod sentance; const DEFAULT_LOG_LEVEL: &str = "warn"; fn init_log() { let level = match std::env::var_os("RUST_LOG") { None => { std::env::set_var("RUST_LOG", DEFAULT_LOG_LEVEL); std::borrow::Cow::Borrowed(std::ffi::OsStr::new(DEFAULT_LOG_LEVEL)) }, Some(w) => std::borrow::Cow::Owned(w), }; pretty_env_logger::init(); trace!("Initialising `genmarkov` ({}) v{} with log level {:?}.\n\tMade by {} with <3.\n\tLicensed with GPL v3 or later", std::env::args().next().unwrap(), env!("CARGO_PKG_VERSION"), level, env!("CARGO_PKG_AUTHORS")); } #[tokio::main] async fn main() { init_log(); let (config, ccache) = match config::load().await { Some(v) => { let cache = match v.try_gen_cache() { Ok(c) => c, Err(e) => { error!("Invalid config, cannot continue"); error!("{}", e); debug!("{:?}", e); return; }, }; (v, cache) }, _ => { let cfg = config::Config::default(); #[cfg(debug_assertions)] { if let Err(err) = cfg.save(config::DEFAULT_FILE_LOCATION).await { error!("Failed to create default config file: {}", err); } } let cache= cfg.try_gen_cache().unwrap(); (cfg, cache) }, }; debug!("Using config {:?}", config); trace!("With config cached: {:?}", ccache); let (chain_handle, chain) = handle::spawn(match save::load(&config.file).await { Ok(chain) => { info!("Loaded chain from {:?}", config.file); chain }, Err(e) => { warn!("Failed to load chain, creating new"); trace!("Error: {}", e); Chain::new() }, }, ccache.handler_settings.clone()); { let mut tasks = Vec::>::new(); tasks.push(chain_handle.map(|res| res.expect("Chain handle panicked")).boxed()); let (state, chain) = { let state = State::new(config, ccache, chain); let state2 = state.clone(); let saver = tokio::spawn(save::host(Box::new(state.clone()))); let chain = warp::any().map(move || state.clone()); tasks.push(saver.map(|res| res.expect("Saver panicked")).boxed()); (state2, chain) }; let client_ip = if state.config().trust_x_forwarded_for { warp::header("x-forwarded-for") .map(|ip: XForwardedFor| ip) .and_then(|x: XForwardedFor| async move { x.into_first().ok_or_else(|| warp::reject::not_found()) }) .or(warp::filters::addr::remote() .and_then(|x: Option| async move { x.map(|x| x.ip()).ok_or_else(|| warp::reject::not_found()) })) .unify().boxed() } else { warp::filters::addr::remote().and_then(|x: Option| async move {x.map(|x| x.ip()).ok_or_else(|| warp::reject::not_found())}).boxed() }; let ipfilter = warp::any() .and(chain) .and(client_ip) .and_then(|state: State, host: IpAddr| { async move { state.config().mask.check(&host) .map(|ci| { trace!("Accepting from rule {:?}", ci); host }) .map(move |host| (state, host)) .map_err(warp::reject::custom) } }).untuple_one(); let push = warp::put() .and(warp::path("put")) .and(ipfilter.clone()) .and(warp::body::content_length_limit(state.config().max_content_length)) .and(warp::body::stream()) .and_then(|state: State, host: IpAddr, buf| { async move { feed::full(&host, state, buf).await .map(|_| warp::reply::with_status(warp::reply(), status!(201))) .map_err(|_| warp::reject::not_found()) //(warp::reject::custom) //TODO: Recover rejection filter down below for custom error return } }) .recover(ipfilt::recover) .with(warp::log("markov::put")); cfg_if!{ if #[cfg(feature="api")] { let api = { let single = { let msz = state.config().max_gen_size; warp::post() .and(ipfilter.clone()) .and(warp::path("single")) .and(warp::path::param() .map(move |sz: usize| { if sz == 0 || (2..=msz).contains(&sz) { Some(sz) } else { None } }) .or(warp::any().map(|| None)) .unify()) .and(warp::body::content_length_limit(state.config().max_content_length)) .and(warp::body::aggregate()) .map(|_, x, y, z| (x,y,z)).untuple_one() .and_then(api::single) .with(warp::log("markov::api::single")) }; warp::path("api") .and(single) .recover(ipfilt::recover) .recover(api::error::rejection) }; } } let read = warp::get() .and(ipfilter.clone()) .and(warp::path::param().map(|opt: usize| Some(opt)) .or(warp::path::end().map(|| Option::::None)).unify()) .and_then(|state: State, host: IpAddr, num: Option| { async move { let (tx, rx) = mpsc::channel(state.config().max_gen_size); tokio::spawn(gen::body(state, num, tx)); Ok::<_, std::convert::Infallible>(Response::new(Body::wrap_stream(rx.filter_map(move |mut x| { if x.trim_in_place().len() != 0 { info!("{} <- {:?}", host, x); x.push('\n'); Some(Ok::<_, std::convert::Infallible>(x)) } else { None } })))) } }) .recover(ipfilt::recover) .with(warp::log("markov::read")); let sentance = warp::get() .and(warp::path("sentance")) //TODO: sanitise::Sentance::new_iter the body line .and(ipfilter.clone()) .and(warp::path::param().map(|opt: usize| Some(opt)) .or(warp::path::end().map(|| Option::::None)).unify()) .and_then(|state: State, host: IpAddr, num: Option| { async move { let (tx, rx) = mpsc::channel(state.config().max_gen_size); tokio::spawn(sentance::body(state, num, tx)); Ok::<_, std::convert::Infallible>(Response::new(Body::wrap_stream(rx.filter_map(move |mut x| { if x.trim_in_place().len() != 0 { info!("{} (sentance) <- {:?}", host, x); x.push(' '); Some(Ok::<_, std::convert::Infallible>(x)) } else { None } })))) } }) .recover(ipfilt::recover) .with(warp::log("markov::read::sentance")); let read = warp::path("get").and(read.or(sentance)); #[cfg(feature="api")] let read = read.or(api); #[cfg(target_family="unix")] tasks.push(tokio::spawn(signals::handle(state.clone())).map(|res| res.expect("Signal handler panicked")).boxed()); require_impl!(Send: async { let (server, init) = { let s2 = AssertNotSend::new(state.clone()); //temp clone the Arcs here for shutdown if server fails to bind, assert they cannot remain cloned across an await boundary. match bind::try_serve(warp::serve(push .or(read)), state.config().bindpoint.clone(), async move { tokio::signal::ctrl_c().await.unwrap(); state.shutdown(); }) { Ok((addr, server)) => { info!("Server bound on {:?}", addr); (server, s2.into_inner().into_initialiser()) }, Err(err) => { error!("Failed to bind server: {}", err); s2.into_inner().shutdown(); return; }, } }; tokio::join![ server, async move { cfg_if! { if #[cfg(feature="instant-init")] { trace!("Setting init"); } else { trace!("Setting init in 2 seconds for good measure."); tokio::time::delay_for(tokio::time::Duration::from_secs(2)).await; } } init.set().expect("Failed to initialise saver") }, ]; }).await; // Cleanup async move { trace!("Cleanup"); debug!("Waiting on {} tasks now", tasks.len()); join_all(tasks).await; } }.await; info!("Shut down gracefully") } mod bind;