2021-10-16 09:50:11 -07:00
|
|
|
#![warn(clippy::nursery, clippy::pedantic)]
|
|
|
|
|
2021-10-26 23:51:05 -07:00
|
|
|
use std::convert::Infallible;
|
2021-10-16 09:50:11 -07:00
|
|
|
use std::sync::Arc;
|
2021-10-24 18:07:48 -07:00
|
|
|
use std::time::Duration;
|
2021-10-16 09:50:11 -07:00
|
|
|
|
|
|
|
use anyhow::Result;
|
|
|
|
use axum::body::Bytes;
|
|
|
|
use axum::extract::{Extension, Path, TypedHeader};
|
|
|
|
use axum::handler::{get, post};
|
2021-10-19 02:18:33 -07:00
|
|
|
use axum::http::header::EXPIRES;
|
|
|
|
use axum::http::StatusCode;
|
2021-10-27 01:49:06 -07:00
|
|
|
use axum::response::Html;
|
2021-10-26 23:51:05 -07:00
|
|
|
use axum::{service, AddExtensionLayer, Router};
|
2021-10-21 18:35:54 -07:00
|
|
|
use chrono::Utc;
|
2021-10-19 02:18:33 -07:00
|
|
|
use headers::HeaderMap;
|
2021-10-27 01:49:06 -07:00
|
|
|
use omegaupload_common::{Expiration, API_ENDPOINT};
|
2021-10-19 02:18:33 -07:00
|
|
|
use rand::thread_rng;
|
|
|
|
use rand::Rng;
|
2021-10-24 18:07:48 -07:00
|
|
|
use rocksdb::{ColumnFamilyDescriptor, IteratorMode};
|
2021-10-19 02:18:33 -07:00
|
|
|
use rocksdb::{Options, DB};
|
2021-10-16 09:50:11 -07:00
|
|
|
use tokio::task;
|
2021-10-26 23:51:05 -07:00
|
|
|
use tower_http::services::ServeDir;
|
2021-10-23 10:10:55 -07:00
|
|
|
use tracing::{error, instrument, trace};
|
2021-10-21 18:35:54 -07:00
|
|
|
use tracing::{info, warn};
|
2021-10-16 09:50:11 -07:00
|
|
|
|
2021-10-19 02:18:33 -07:00
|
|
|
use crate::short_code::ShortCode;
|
2021-10-16 09:50:11 -07:00
|
|
|
|
|
|
|
mod short_code;
|
|
|
|
|
2021-10-24 18:07:48 -07:00
|
|
|
const BLOB_CF_NAME: &str = "blob";
|
|
|
|
const META_CF_NAME: &str = "meta";
|
|
|
|
|
2021-10-16 09:50:11 -07:00
|
|
|
#[tokio::main]
|
|
|
|
async fn main() -> Result<()> {
|
2021-10-24 18:07:48 -07:00
|
|
|
const PASTE_DB_PATH: &str = "database";
|
2021-10-16 09:50:11 -07:00
|
|
|
const SHORT_CODE_SIZE: usize = 12;
|
|
|
|
|
|
|
|
tracing_subscriber::fmt::init();
|
|
|
|
|
2021-10-24 18:07:48 -07:00
|
|
|
let mut db_options = Options::default();
|
|
|
|
db_options.create_if_missing(true);
|
|
|
|
db_options.create_missing_column_families(true);
|
|
|
|
db_options.set_compression_type(rocksdb::DBCompressionType::Zstd);
|
|
|
|
let db = Arc::new(DB::open_cf_descriptors(
|
|
|
|
&db_options,
|
|
|
|
PASTE_DB_PATH,
|
|
|
|
[
|
|
|
|
ColumnFamilyDescriptor::new(BLOB_CF_NAME, Options::default()),
|
|
|
|
ColumnFamilyDescriptor::new(META_CF_NAME, Options::default()),
|
|
|
|
],
|
|
|
|
)?);
|
2021-10-16 09:50:11 -07:00
|
|
|
|
2021-10-25 17:31:30 -07:00
|
|
|
set_up_expirations(&db);
|
2021-10-16 09:50:11 -07:00
|
|
|
|
2021-10-26 23:51:05 -07:00
|
|
|
let root_service = service::get(ServeDir::new("static"))
|
|
|
|
.handle_error(|_| Ok::<_, Infallible>(StatusCode::NOT_FOUND));
|
|
|
|
|
|
|
|
axum::Server::bind(&"0.0.0.0:8080".parse()?)
|
2021-10-16 09:50:11 -07:00
|
|
|
.serve(
|
|
|
|
Router::new()
|
|
|
|
.route("/", post(upload::<SHORT_CODE_SIZE>))
|
|
|
|
.route(
|
|
|
|
"/:code",
|
2021-10-27 01:49:06 -07:00
|
|
|
get(|| async { Html(include_str!("../../dist/index.html")) }),
|
2021-10-26 23:51:05 -07:00
|
|
|
)
|
|
|
|
.nest("/static", root_service)
|
|
|
|
.route(
|
2021-10-27 01:49:06 -07:00
|
|
|
&format!("{}{}", API_ENDPOINT.to_string(), "/:code"),
|
2021-10-16 09:50:11 -07:00
|
|
|
get(paste::<SHORT_CODE_SIZE>).delete(delete::<SHORT_CODE_SIZE>),
|
|
|
|
)
|
|
|
|
.layer(AddExtensionLayer::new(db))
|
|
|
|
.into_make_service(),
|
|
|
|
)
|
|
|
|
.await?;
|
|
|
|
|
|
|
|
// Must be called for correct shutdown
|
2021-10-24 18:07:48 -07:00
|
|
|
DB::destroy(&Options::default(), PASTE_DB_PATH)?;
|
2021-10-16 09:50:11 -07:00
|
|
|
Ok(())
|
|
|
|
}
|
|
|
|
|
2021-10-25 17:31:30 -07:00
|
|
|
fn set_up_expirations(db: &Arc<DB>) {
|
2021-10-21 18:35:54 -07:00
|
|
|
let mut corrupted = 0;
|
|
|
|
let mut expired = 0;
|
|
|
|
let mut pending = 0;
|
|
|
|
|
|
|
|
info!("Setting up cleanup timers, please wait...");
|
|
|
|
|
2021-10-24 18:07:48 -07:00
|
|
|
let meta_cf = db.cf_handle(META_CF_NAME).unwrap();
|
|
|
|
|
2021-10-25 17:31:30 -07:00
|
|
|
let db_ref = Arc::clone(db);
|
2021-10-24 18:07:48 -07:00
|
|
|
|
|
|
|
let delete_entry = move |key: &[u8]| {
|
|
|
|
let blob_cf = db_ref.cf_handle(BLOB_CF_NAME).unwrap();
|
|
|
|
let meta_cf = db_ref.cf_handle(META_CF_NAME).unwrap();
|
|
|
|
if let Err(e) = db_ref.delete_cf(blob_cf, &key) {
|
|
|
|
warn!("{}", e);
|
|
|
|
}
|
|
|
|
if let Err(e) = db_ref.delete_cf(meta_cf, &key) {
|
|
|
|
warn!("{}", e);
|
|
|
|
}
|
|
|
|
};
|
|
|
|
|
|
|
|
for (key, value) in db.iterator_cf(meta_cf, IteratorMode::Start) {
|
2021-10-25 17:31:30 -07:00
|
|
|
let expiration = if let Ok(value) = bincode::deserialize::<Expiration>(&value) {
|
2021-10-21 18:35:54 -07:00
|
|
|
value
|
|
|
|
} else {
|
|
|
|
corrupted += 1;
|
2021-10-24 18:07:48 -07:00
|
|
|
delete_entry(&key);
|
2021-10-21 18:35:54 -07:00
|
|
|
continue;
|
|
|
|
};
|
|
|
|
|
2021-10-25 17:31:30 -07:00
|
|
|
let expiration_time = match expiration {
|
2021-10-24 18:07:48 -07:00
|
|
|
Expiration::BurnAfterReading => {
|
|
|
|
panic!("Got burn after reading expiration time? Invariant violated");
|
2021-10-21 18:35:54 -07:00
|
|
|
}
|
2021-10-24 18:07:48 -07:00
|
|
|
Expiration::UnixTime(time) => time,
|
|
|
|
};
|
|
|
|
|
|
|
|
let sleep_duration = (expiration_time - Utc::now()).to_std().unwrap_or_default();
|
2021-10-25 17:31:30 -07:00
|
|
|
if sleep_duration == Duration::default() {
|
|
|
|
expired += 1;
|
|
|
|
delete_entry(&key);
|
|
|
|
} else {
|
2021-10-24 18:07:48 -07:00
|
|
|
pending += 1;
|
|
|
|
let delete_entry_ref = delete_entry.clone();
|
|
|
|
task::spawn_blocking(move || async move {
|
|
|
|
tokio::time::sleep(sleep_duration).await;
|
|
|
|
delete_entry_ref(&key);
|
|
|
|
});
|
2021-10-21 18:35:54 -07:00
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
if corrupted == 0 {
|
|
|
|
info!("No corrupted pastes found.");
|
|
|
|
} else {
|
|
|
|
warn!("Found {} corrupted pastes.", corrupted);
|
|
|
|
}
|
2021-10-25 17:31:30 -07:00
|
|
|
|
2021-10-21 18:35:54 -07:00
|
|
|
info!("Found {} expired pastes.", expired);
|
|
|
|
info!("Found {} active pastes.", pending);
|
|
|
|
info!("Cleanup timers have been initialized.");
|
|
|
|
}
|
|
|
|
|
2021-10-23 10:10:55 -07:00
|
|
|
#[instrument(skip(db, body), err)]
|
2021-10-16 09:50:11 -07:00
|
|
|
async fn upload<const N: usize>(
|
|
|
|
Extension(db): Extension<Arc<DB>>,
|
|
|
|
maybe_expires: Option<TypedHeader<Expiration>>,
|
|
|
|
body: Bytes,
|
|
|
|
) -> Result<Vec<u8>, StatusCode> {
|
|
|
|
if body.is_empty() {
|
|
|
|
return Err(StatusCode::BAD_REQUEST);
|
|
|
|
}
|
|
|
|
|
|
|
|
// 3GB max; this is a soft-limit of RocksDb
|
|
|
|
if body.len() >= 3_221_225_472 {
|
|
|
|
return Err(StatusCode::PAYLOAD_TOO_LARGE);
|
|
|
|
}
|
|
|
|
|
|
|
|
let mut new_key = None;
|
|
|
|
|
2021-10-23 10:10:55 -07:00
|
|
|
trace!("Generating short code...");
|
|
|
|
|
2021-10-16 09:50:11 -07:00
|
|
|
// Try finding a code; give up after 1000 attempts
|
|
|
|
// Statistics show that this is very unlikely to happen
|
2021-10-23 10:10:55 -07:00
|
|
|
for i in 0..1000 {
|
2021-10-19 02:18:33 -07:00
|
|
|
let code: ShortCode<N> = thread_rng().sample(short_code::Generator);
|
2021-10-16 09:50:11 -07:00
|
|
|
let db = Arc::clone(&db);
|
|
|
|
let key = code.as_bytes();
|
2021-10-24 18:07:48 -07:00
|
|
|
let query = task::spawn_blocking(move || {
|
|
|
|
db.key_may_exist_cf(db.cf_handle(META_CF_NAME).unwrap(), key)
|
|
|
|
})
|
|
|
|
.await;
|
2021-10-16 09:50:11 -07:00
|
|
|
if matches!(query, Ok(false)) {
|
|
|
|
new_key = Some(key);
|
2021-10-23 10:10:55 -07:00
|
|
|
trace!("Found new key after {} attempts.", i);
|
|
|
|
break;
|
2021-10-16 09:50:11 -07:00
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
let key = if let Some(key) = new_key {
|
|
|
|
key
|
|
|
|
} else {
|
2021-10-23 10:10:55 -07:00
|
|
|
error!("Failed to generate a valid short code!");
|
2021-10-16 09:50:11 -07:00
|
|
|
return Err(StatusCode::INTERNAL_SERVER_ERROR);
|
|
|
|
};
|
|
|
|
|
2021-10-23 10:10:55 -07:00
|
|
|
trace!("Serializing paste...");
|
2021-10-16 09:50:11 -07:00
|
|
|
|
2021-10-23 10:10:55 -07:00
|
|
|
trace!("Finished serializing paste.");
|
|
|
|
|
2021-10-21 18:35:54 -07:00
|
|
|
let db_ref = Arc::clone(&db);
|
2021-10-24 18:07:48 -07:00
|
|
|
match task::spawn_blocking(move || {
|
|
|
|
let blob_cf = db_ref.cf_handle(BLOB_CF_NAME).unwrap();
|
|
|
|
let meta_cf = db_ref.cf_handle(META_CF_NAME).unwrap();
|
|
|
|
let data = bincode::serialize(&body).expect("bincode to serialize");
|
|
|
|
db_ref.put_cf(blob_cf, key, data)?;
|
|
|
|
let expires = maybe_expires.map(|v| v.0).unwrap_or_default();
|
|
|
|
let meta = bincode::serialize(&expires).expect("bincode to serialize");
|
|
|
|
if db_ref.put_cf(meta_cf, key, meta).is_err() {
|
|
|
|
// try and roll back on metadata write failure
|
|
|
|
db_ref.delete_cf(blob_cf, key)?;
|
|
|
|
}
|
|
|
|
Result::<_, anyhow::Error>::Ok(())
|
|
|
|
})
|
|
|
|
.await
|
|
|
|
{
|
2021-10-21 18:35:54 -07:00
|
|
|
Ok(Ok(_)) => {
|
|
|
|
if let Some(expires) = maybe_expires {
|
2021-10-24 18:07:48 -07:00
|
|
|
if let Expiration::UnixTime(expiration_time) = expires.0 {
|
|
|
|
let sleep_duration =
|
|
|
|
(expiration_time - Utc::now()).to_std().unwrap_or_default();
|
|
|
|
|
|
|
|
task::spawn_blocking(move || async move {
|
|
|
|
tokio::time::sleep(sleep_duration).await;
|
|
|
|
let blob_cf = db.cf_handle(BLOB_CF_NAME).unwrap();
|
|
|
|
let meta_cf = db.cf_handle(META_CF_NAME).unwrap();
|
|
|
|
if let Err(e) = db.delete_cf(blob_cf, key) {
|
|
|
|
warn!("{}", e);
|
|
|
|
}
|
|
|
|
if let Err(e) = db.delete_cf(meta_cf, key) {
|
2021-10-21 18:35:54 -07:00
|
|
|
warn!("{}", e);
|
|
|
|
}
|
2021-10-24 18:07:48 -07:00
|
|
|
});
|
2021-10-21 18:35:54 -07:00
|
|
|
}
|
|
|
|
}
|
|
|
|
}
|
2021-10-16 09:50:11 -07:00
|
|
|
e => {
|
|
|
|
error!("Failed to insert paste into db: {:?}", e);
|
|
|
|
return Err(StatusCode::INTERNAL_SERVER_ERROR);
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
Ok(Vec::from(key))
|
|
|
|
}
|
|
|
|
|
|
|
|
#[instrument(skip(db), err)]
|
|
|
|
async fn paste<const N: usize>(
|
|
|
|
Extension(db): Extension<Arc<DB>>,
|
|
|
|
Path(url): Path<ShortCode<N>>,
|
2021-10-19 02:18:33 -07:00
|
|
|
) -> Result<(HeaderMap, Bytes), StatusCode> {
|
2021-10-16 09:50:11 -07:00
|
|
|
let key = url.as_bytes();
|
|
|
|
|
2021-10-24 18:07:48 -07:00
|
|
|
let metadata: Expiration = {
|
|
|
|
let meta_cf = db.cf_handle(META_CF_NAME).unwrap();
|
|
|
|
let query_result = db.get_cf(meta_cf, key).map_err(|e| {
|
2021-10-16 09:50:11 -07:00
|
|
|
error!("Failed to fetch initial query: {}", e);
|
|
|
|
StatusCode::INTERNAL_SERVER_ERROR
|
|
|
|
})?;
|
|
|
|
|
|
|
|
let data = match query_result {
|
|
|
|
Some(data) => data,
|
|
|
|
None => return Err(StatusCode::NOT_FOUND),
|
|
|
|
};
|
|
|
|
|
|
|
|
bincode::deserialize(&data).map_err(|_| {
|
|
|
|
error!("Failed to deserialize data?!");
|
|
|
|
StatusCode::INTERNAL_SERVER_ERROR
|
|
|
|
})?
|
|
|
|
};
|
|
|
|
|
2021-10-24 18:07:48 -07:00
|
|
|
// Check if paste has expired.
|
|
|
|
if let Expiration::UnixTime(expires) = metadata {
|
|
|
|
if expires < Utc::now() {
|
|
|
|
task::spawn_blocking(move || {
|
|
|
|
let blob_cf = db.cf_handle(BLOB_CF_NAME).unwrap();
|
|
|
|
let meta_cf = db.cf_handle(META_CF_NAME).unwrap();
|
|
|
|
if let Err(e) = db.delete_cf(blob_cf, &key) {
|
|
|
|
warn!("{}", e);
|
|
|
|
}
|
|
|
|
if let Err(e) = db.delete_cf(meta_cf, &key) {
|
|
|
|
warn!("{}", e);
|
|
|
|
}
|
|
|
|
})
|
2021-10-16 09:50:11 -07:00
|
|
|
.await
|
|
|
|
.map_err(|e| {
|
|
|
|
error!("Failed to join handle: {}", e);
|
|
|
|
StatusCode::INTERNAL_SERVER_ERROR
|
|
|
|
})?;
|
|
|
|
|
2021-10-24 18:07:48 -07:00
|
|
|
return Err(StatusCode::NOT_FOUND);
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
let paste: Bytes = {
|
|
|
|
// not sure if perf of get_pinned is better than spawn_blocking
|
|
|
|
let blob_cf = db.cf_handle(BLOB_CF_NAME).unwrap();
|
|
|
|
let query_result = db.get_pinned_cf(blob_cf, key).map_err(|e| {
|
|
|
|
error!("Failed to fetch initial query: {}", e);
|
2021-10-16 09:50:11 -07:00
|
|
|
StatusCode::INTERNAL_SERVER_ERROR
|
|
|
|
})?;
|
|
|
|
|
2021-10-24 18:07:48 -07:00
|
|
|
let data = match query_result {
|
|
|
|
Some(data) => data,
|
|
|
|
None => return Err(StatusCode::NOT_FOUND),
|
|
|
|
};
|
2021-10-16 09:50:11 -07:00
|
|
|
|
2021-10-24 18:07:48 -07:00
|
|
|
bincode::deserialize(&data).map_err(|_| {
|
|
|
|
error!("Failed to deserialize data?!");
|
|
|
|
StatusCode::INTERNAL_SERVER_ERROR
|
|
|
|
})?
|
|
|
|
};
|
|
|
|
|
|
|
|
// Check if we need to burn after read
|
|
|
|
if matches!(metadata, Expiration::BurnAfterReading) {
|
2021-10-16 09:50:11 -07:00
|
|
|
let join_handle = task::spawn_blocking(move || db.delete(key))
|
|
|
|
.await
|
|
|
|
.map_err(|e| {
|
|
|
|
error!("Failed to join handle: {}", e);
|
|
|
|
StatusCode::INTERNAL_SERVER_ERROR
|
|
|
|
})?;
|
|
|
|
|
|
|
|
join_handle.map_err(|e| {
|
|
|
|
error!("Failed to burn paste after read: {}", e);
|
|
|
|
StatusCode::INTERNAL_SERVER_ERROR
|
|
|
|
})?;
|
|
|
|
}
|
|
|
|
|
2021-10-19 02:18:33 -07:00
|
|
|
let mut map = HeaderMap::new();
|
2021-10-24 18:07:48 -07:00
|
|
|
map.insert(EXPIRES, metadata.into());
|
|
|
|
|
|
|
|
Ok((map, paste))
|
2021-10-16 09:50:11 -07:00
|
|
|
}
|
|
|
|
|
|
|
|
#[instrument(skip(db))]
|
|
|
|
async fn delete<const N: usize>(
|
|
|
|
Extension(db): Extension<Arc<DB>>,
|
|
|
|
Path(url): Path<ShortCode<N>>,
|
|
|
|
) -> StatusCode {
|
2021-10-24 18:07:48 -07:00
|
|
|
match task::spawn_blocking(move || {
|
|
|
|
let blob_cf = db.cf_handle(BLOB_CF_NAME).unwrap();
|
|
|
|
let meta_cf = db.cf_handle(META_CF_NAME).unwrap();
|
|
|
|
if let Err(e) = db.delete_cf(blob_cf, url.as_bytes()) {
|
|
|
|
warn!("{}", e);
|
|
|
|
return Err(());
|
|
|
|
}
|
|
|
|
|
|
|
|
if let Err(e) = db.delete_cf(meta_cf, url.as_bytes()) {
|
|
|
|
warn!("{}", e);
|
|
|
|
return Err(());
|
|
|
|
}
|
|
|
|
|
|
|
|
Ok(())
|
|
|
|
})
|
|
|
|
.await
|
|
|
|
{
|
|
|
|
Ok(_) => StatusCode::OK,
|
2021-10-16 09:50:11 -07:00
|
|
|
_ => StatusCode::INTERNAL_SERVER_ERROR,
|
|
|
|
}
|
|
|
|
}
|