asonix/pict-rs
1
1
Fork 1
Code Issues 2 Pull requests 3 Actions Projects Releases 133 Wiki Activity

Implement delete tokens

Browse source
This commit is contained in:
asonix 2020-06-06 19:29:15 -05:00
parent 7298e500cb
commit d60862a0d0
3 changed files with 199 additions and 10 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

16
src/error.rs
View file

@ -20,6 +20,9 @@ pub enum UploadError {
#[error("Error processing image, {0}")]
Image(#[from] image::error::ImageError),
#[error("Error interacting with filesystem, {0}")]
Io(#[from] std::io::Error),
#[error("Panic in blocking operation")]
Canceled,
@ -34,6 +37,18 @@ pub enum UploadError {
#[error("Alias directed to missing file")]
MissingFile,
#[error("Provided token did not match expected token")]
InvalidToken,
}
impl From<sled::transaction::TransactionError<UploadError>> for UploadError {
fn from(e: sled::transaction::TransactionError<UploadError>) -> Self {
match e {
sled::transaction::TransactionError::Abort(t) => t,
sled::transaction::TransactionError::Storage(e) => e.into(),
}
}
}
impl From<actix_form_data::Error> for UploadError {
@ -61,6 +76,7 @@ impl ResponseError for UploadError {
StatusCode::BAD_REQUEST
}
UploadError::MissingAlias => StatusCode::NOT_FOUND,
UploadError::InvalidToken => StatusCode::FORBIDDEN,
_ => StatusCode::INTERNAL_SERVER_ERROR,
}
}

26
src/main.rs
View file

@ -76,7 +76,10 @@ fn from_ext(ext: std::ffi::OsString) -> mime::Mime {
}
/// Handle responding to succesful uploads
async fn upload(value: Value) -> Result<HttpResponse, UploadError> {
async fn upload(
value: Value,
manager: web::Data<UploadManager>,
) -> Result<HttpResponse, UploadError> {
let images = value
.map()
.and_then(|mut m| m.remove("images"))
@ -92,13 +95,28 @@ async fn upload(value: Value) -> Result<HttpResponse, UploadError> {
.and_then(|s| s.to_str())
{
info!("Uploaded {} as {:?}", image.filename, saved_as);
files.push(serde_json::json!({ "file": saved_as }));
let delete_token = manager.delete_token(saved_as.to_owned()).await?;
files.push(serde_json::json!({
"file": saved_as,
"delete_token": delete_token,
}));
}
}
Ok(HttpResponse::Created().json(serde_json::json!({ "msg": "ok", "files": files })))
}
async fn delete(
manager: web::Data<UploadManager>,
path_entries: web::Path<(String, String)>,
) -> Result<HttpResponse, UploadError> {
let (alias, token) = path_entries.into_inner();
manager.delete(token, alias).await?;
Ok(HttpResponse::NoContent().finish())
}
/// Serve original files
async fn serve(
manager: web::Data<UploadManager>,
@ -248,6 +266,10 @@ async fn main() -> Result<(), anyhow::Error> {
.route(web::post().to(upload)),
)
.service(web::resource("/{filename}").route(web::get().to(serve)))
.service(
web::resource("/delete/{delete_token}/{filename}")
.route(web::delete().to(delete)),
)
.service(
web::resource("/{size}/{filename}").route(web::get().to(serve_resized)),
),

167
src/upload_manager.rs
View file

@ -1,6 +1,7 @@
use crate::{error::UploadError, safe_save_file, to_ext, ACCEPTED_MIMES};
use actix_web::web;
use futures::stream::{Stream, StreamExt};
use log::warn;
use sha2::Digest;
use std::{path::PathBuf, pin::Pin, sync::Arc};
@ -12,8 +13,8 @@ pub struct UploadManager {
struct UploadManagerInner {
hasher: sha2::Sha256,
image_dir: PathBuf,
db: sled::Db,
alias_tree: sled::Tree,
db: sled::Db,
}
type UploadStream = Pin<Box<dyn Stream<Item = Result<bytes::Bytes, actix_form_data::Error>>>>;
@ -43,10 +44,7 @@ impl UploadManager {
let mut sled_dir = root_dir.clone();
sled_dir.push("db");
// sled automatically creates it's own directories
//
// This is technically a blocking operation but it's fine because it happens before we
// start handling requests
let db = sled::open(sled_dir)?;
let db = web::block(move || sled::open(sled_dir)).await?;
root_dir.push("files");
@ -63,6 +61,108 @@ impl UploadManager {
})
}
pub(crate) async fn delete(&self, alias: String, token: String) -> Result<(), UploadError> {
use sled::Transactional;
let db = self.inner.db.clone();
let alias_tree = self.inner.alias_tree.clone();
let alias2 = alias.clone();
let hash = web::block(move || {
[&*db, &alias_tree].transaction(|v| {
let db = &v[0];
let alias_tree = &v[1];
// -- GET TOKEN --
let existing_token = alias_tree
.remove(delete_key(&alias2).as_bytes())?
.ok_or(trans_err(UploadError::MissingAlias))?;
// Bail if invalid token
if existing_token != token {
warn!("Invalid delete token");
return Err(trans_err(UploadError::InvalidToken));
}
// -- GET ID FOR HASH TREE CLEANUP --
let id = alias_tree
.remove(alias_id_key(&alias2).as_bytes())?
.ok_or(trans_err(UploadError::MissingAlias))?;
let id = String::from_utf8(id.to_vec()).map_err(|e| trans_err(e.into()))?;
// -- GET HASH FOR HASH TREE CLEANUP --
let hash = alias_tree
.remove(alias2.as_bytes())?
.ok_or(trans_err(UploadError::MissingAlias))?;
// -- REMOVE HASH TREE ELEMENT --
db.remove(alias_key(&hash, &id))?;
Ok(hash)
})
})
.await?;
// -- CHECK IF ANY OTHER ALIASES EXIST
let db = self.inner.db.clone();
let (start, end) = alias_key_bounds(&hash);
let any_aliases = web::block(move || {
Ok(db.range(start..end).next().is_some()) as Result<bool, UploadError>
})
.await?;
// Bail if there are existing aliases
if any_aliases {
return Ok(());
}
// -- DELETE HASH ENTRY --
let db = self.inner.db.clone();
let real_filename = web::block(move || {
let real_filename = db.remove(&hash)?.ok_or(UploadError::MissingFile)?;
Ok(real_filename) as Result<sled::IVec, UploadError>
})
.await?;
let real_filename = String::from_utf8(real_filename.to_vec())?;
let image_dir = self.image_dir();
web::block(move || blocking_delete_all_by_filename(image_dir, &real_filename)).await?;
Ok(())
}
/// Generate a delete token for an alias
pub(crate) async fn delete_token(&self, alias: String) -> Result<String, UploadError> {
use rand::distributions::{Alphanumeric, Distribution};
let rng = rand::thread_rng();
let s: String = Alphanumeric.sample_iter(rng).take(10).collect();
let delete_token = s.clone();
let alias_tree = self.inner.alias_tree.clone();
let key = delete_key(&alias);
let res = web::block(move || {
alias_tree.compare_and_swap(
key.as_bytes(),
None as Option<sled::IVec>,
Some(s.as_bytes()),
)
})
.await?;
if let Err(sled::CompareAndSwapError {
current: Some(ivec),
..
}) = res
{
let s = String::from_utf8(ivec.to_vec())?;
return Ok(s);
}
Ok(delete_token)
}
/// Upload the file, discarding bytes if it's already present, or saving if it's new
pub(crate) async fn upload(
&self,
@ -206,9 +306,7 @@ impl UploadManager {
let db = self.inner.db.clone();
let id = web::block(move || db.generate_id()).await?.to_string();
let mut key = hash.to_vec();
key.extend(id.as_bytes());
let key = alias_key(hash, &id);
let db = self.inner.db.clone();
let alias2 = alias.clone();
let res = web::block(move || {
@ -217,6 +315,10 @@ impl UploadManager {
.await?;
if res.is_ok() {
let alias_tree = self.inner.alias_tree.clone();
let key = alias_id_key(&alias);
web::block(move || alias_tree.insert(key.as_bytes(), id.as_bytes())).await?;
break;
}
}
@ -255,6 +357,55 @@ impl UploadManager {
}
}
fn blocking_delete_all_by_filename(mut dir: PathBuf, filename: &str) -> Result<(), UploadError> {
for res in std::fs::read_dir(dir.clone())? {
let entry = res?;
if entry.path().is_dir() {
blocking_delete_all_by_filename(entry.path(), filename)?;
}
}
dir.push(filename);
if dir.is_file() {
std::fs::remove_file(dir)?;
}
Ok(())
}
fn trans_err(e: UploadError) -> sled::transaction::ConflictableTransactionError<UploadError> {
sled::transaction::ConflictableTransactionError::Abort(e)
}
fn file_name(name: String, content_type: mime::Mime) -> String {
format!("{}{}", name, to_ext(content_type))
}
fn alias_key(hash: &[u8], id: &str) -> Vec<u8> {
let mut key = hash.to_vec();
// add a separator to the key between the hash and the ID
key.extend(&[0]);
key.extend(id.as_bytes());
key
}
fn alias_key_bounds(hash: &[u8]) -> (Vec<u8>, Vec<u8>) {
let mut start = hash.to_vec();
start.extend(&[0]);
let mut end = hash.to_vec();
end.extend(&[1]);
(start, end)
}
fn alias_id_key(alias: &str) -> String {
format!("{}/id", alias)
}
fn delete_key(alias: &str) -> String {
format!("{}/delete", alias)
}