You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
341 lines
8.4 KiB
341 lines
8.4 KiB
#![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::<BoxFuture<'static, ()>>::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<SocketAddr>| 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<SocketAddr>| 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::<usize>::None)).unify())
|
|
.and_then(|state: State, host: IpAddr, num: Option<usize>| {
|
|
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::<usize>::None)).unify())
|
|
.and_then(|state: State, host: IpAddr, num: Option<usize>| {
|
|
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;
|