transbeam/src/main.rs

242 lines
6.7 KiB
Rust
Raw Normal View History

mod download;
mod store;
2022-04-26 23:54:29 -04:00
mod upload;
mod zip;
use std::{fmt::Debug, fs::File, path::PathBuf, str::FromStr};
2022-04-26 23:54:29 -04:00
use actix_files::NamedFile;
2022-04-26 23:54:29 -04:00
use actix_web::{
get, http::StatusCode, middleware::Logger, post, web, App, HttpRequest, HttpResponse,
HttpServer, Responder,
2022-04-26 23:54:29 -04:00
};
use actix_web_actors::ws;
use bytesize::ByteSize;
use log::{error, warn};
use serde::{Deserialize, Serialize};
use store::FileStore;
use tokio::sync::RwLock;
2022-04-26 23:54:29 -04:00
2022-04-28 05:18:35 -04:00
const APP_NAME: &str = "transbeam";
struct AppState {
file_store: RwLock<FileStore>,
config: Config,
}
struct Config {
max_upload_size: u64,
max_lifetime: u16,
upload_password: String,
storage_dir: PathBuf,
static_dir: PathBuf,
reverse_proxy: bool,
mnemonic_codes: bool,
}
pub fn get_ip_addr(req: &HttpRequest, reverse_proxy: bool) -> String {
let conn_info = req.connection_info();
if reverse_proxy {
conn_info.realip_remote_addr()
} else {
conn_info.peer_addr()
}
.unwrap()
.to_owned()
}
pub fn log_auth_failure(ip_addr: &str) {
warn!("Incorrect authentication attempt from {}", ip_addr);
}
2022-04-26 23:54:29 -04:00
#[derive(Deserialize)]
struct DownloadRequest {
code: String,
}
#[get("/download")]
2022-04-28 05:18:35 -04:00
async fn handle_download(
req: HttpRequest,
download: web::Query<DownloadRequest>,
data: web::Data<AppState>,
2022-04-28 05:18:35 -04:00
) -> actix_web::Result<HttpResponse> {
let code = &download.code;
if !store::is_valid_storage_code(code) {
return download_not_found(req, data);
}
let info = data.file_store.read().await.lookup_file(code);
if let Some(info) = info {
let storage_path = data.config.storage_dir.join(code);
let file = File::open(&storage_path)?;
Ok(download::DownloadingFile {
file,
storage_path,
2022-04-28 05:18:35 -04:00
info,
}
.into_response(&req))
} else {
download_not_found(req, data)
}
}
fn download_not_found(
req: HttpRequest,
data: web::Data<AppState>,
) -> actix_web::Result<HttpResponse> {
let ip_addr = get_ip_addr(&req, data.config.reverse_proxy);
log_auth_failure(&ip_addr);
Ok(NamedFile::open(data.config.static_dir.join("404.html"))?
.set_status_code(StatusCode::NOT_FOUND)
.into_response(&req))
}
2022-04-26 23:54:29 -04:00
#[get("/upload")]
async fn handle_upload(
req: HttpRequest,
stream: web::Payload,
data: web::Data<AppState>,
) -> impl Responder {
if data.file_store.read().await.full() {
return Ok(HttpResponse::BadRequest().finish());
}
let ip_addr = get_ip_addr(&req, data.config.reverse_proxy);
ws::start(upload::Uploader::new(data, ip_addr), &req, stream)
}
#[derive(Deserialize)]
struct UploadPasswordCheck {
password: String,
}
#[post("/upload/check_password")]
async fn check_upload_password(
req: HttpRequest,
body: web::Json<UploadPasswordCheck>,
data: web::Data<AppState>,
) -> impl Responder {
let ip_addr = get_ip_addr(&req, data.config.reverse_proxy);
if body.password != data.config.upload_password {
log_auth_failure(&ip_addr);
HttpResponse::Forbidden().finish()
} else {
HttpResponse::NoContent().finish()
}
}
#[derive(Serialize)]
struct UploadLimits {
open: bool,
max_size: u64,
max_lifetime: u16,
}
#[get("/upload/limits.json")]
async fn upload_limits(data: web::Data<AppState>) -> impl Responder {
let file_store = data.file_store.read().await;
let open = !file_store.full();
let available_size = file_store.available_size();
let max_size = std::cmp::min(available_size, data.config.max_upload_size);
web::Json(UploadLimits {
open,
max_size,
max_lifetime: data.config.max_lifetime,
})
}
fn env_or<T: FromStr>(var: &str, default: T) -> T
where
<T as FromStr>::Err: Debug,
{
std::env::var(var)
.map(|val| {
val.parse::<T>()
.unwrap_or_else(|_| panic!("Invalid value {} for variable {}", val, var))
})
.unwrap_or(default)
}
fn env_or_else<T: FromStr, F: FnOnce() -> T>(var: &str, default: F) -> T
where
<T as FromStr>::Err: Debug,
{
std::env::var(var)
.map(|val| {
val.parse::<T>()
.unwrap_or_else(|_| panic!("Invalid value {} for variable {}", val, var))
})
.ok()
.unwrap_or_else(default)
2022-04-26 23:54:29 -04:00
}
#[actix_web::main]
async fn main() -> std::io::Result<()> {
dotenv::dotenv().ok();
2022-04-26 23:54:29 -04:00
env_logger::init();
let static_dir: PathBuf = env_or_else("TRANSBEAM_STATIC_DIR", || PathBuf::from("static"));
let storage_dir: PathBuf = env_or_else("TRANSBEAM_STORAGE_DIR", || PathBuf::from("storage"));
let port: u16 = env_or("TRANSBEAM_PORT", 8080);
let mnemonic_codes: bool = env_or("TRANSBEAM_MNEMONIC_CODES", true);
let reverse_proxy: bool = env_or("TRANSBEAM_REVERSE_PROXY", false);
let max_lifetime: u16 = env_or("TRANSBEAM_MAX_LIFETIME", 30);
let max_upload_size: u64 =
env_or::<ByteSize>("TRANSBEAM_MAX_UPLOAD_SIZE", ByteSize(16 * bytesize::GB)).as_u64();
let max_storage_size: u64 =
env_or::<ByteSize>("TRANSBEAM_MAX_STORAGE_SIZE", ByteSize(64 * bytesize::GB)).as_u64();
let upload_password: String =
std::env::var("TRANSBEAM_UPLOAD_PASSWORD").expect("TRANSBEAM_UPLOAD_PASSWORD must be set!");
2022-04-26 23:54:29 -04:00
let data = web::Data::new(AppState {
file_store: RwLock::new(FileStore::load(storage_dir.clone(), max_storage_size).await?),
config: Config {
max_upload_size,
max_lifetime,
upload_password,
storage_dir,
static_dir: static_dir.clone(),
reverse_proxy,
mnemonic_codes,
},
});
start_reaper(data.clone());
2022-04-27 00:53:32 -04:00
2022-04-26 23:54:29 -04:00
HttpServer::new(move || {
App::new()
.app_data(data.clone())
.wrap(if data.config.reverse_proxy {
Logger::new(r#"%{r}a "%r" %s %b "%{Referer}i" "%{User-Agent}i" %T"#)
} else {
Logger::default()
})
.service(handle_download)
.service(handle_upload)
.service(check_upload_password)
.service(upload_limits)
2022-04-27 00:53:32 -04:00
.service(actix_files::Files::new("/", static_dir.clone()).index_file("index.html"))
2022-04-26 23:54:29 -04:00
})
.bind((
if reverse_proxy {
"127.0.0.1"
} else {
"0.0.0.0"
},
port,
))?
2022-04-26 23:54:29 -04:00
.run()
.await?;
Ok(())
}
fn start_reaper(data: web::Data<AppState>) {
std::thread::spawn(move || {
actix_web::rt::System::new().block_on(async {
loop {
actix_web::rt::time::sleep(core::time::Duration::from_secs(86400)).await;
if let Err(e) = data.file_store.write().await.remove_expired_files().await {
error!("Error reaping expired files: {}", e);
}
}
});
});
}