initial server

src/expire_queue.rs

@@ -0,0 +1,161 @@
+use err_derive::Error;
+use priority_queue::PriorityQueue;
+use std::cmp::Ordering;
+use std::convert::TryInto;
+use std::fmt::Debug;
+use std::sync::{Arc, Mutex};
+use uuid::{Builder, Uuid, Variant, Version};
+
+#[derive(Debug, Clone, Hash, PartialEq, Eq, Copy)]
+pub struct UploadId(Uuid);
+
+#[derive(Debug, Error)]
+pub enum InvalidUploadIdError {
+    #[error(display = "Invalid string")]
+    InvalidString(#[error(source)] std::str::Utf8Error),
+    #[error(display = "Invalid upload id, uuid")]
+    InvalidUUID(#[error(source)] uuid::Error),
+}
+
+impl UploadId {
+    pub fn new(id: &str) -> Result<Self, InvalidUploadIdError> {
+        let uuid = Uuid::parse_str(id)?;
+        Ok(UploadId(uuid))
+    }
+
+    pub fn generate(expire: u64) -> Self {
+        use rand::RngCore;
+
+        let mut rng = rand::thread_rng();
+        let mut bytes = [0; 16];
+
+        rng.fill_bytes(&mut bytes);
+
+        let mut uuid_bytes = *Builder::from_bytes(bytes)
+            .set_variant(Variant::RFC4122)
+            .set_version(Version::Random)
+            .build()
+            .as_bytes();
+
+        // store the expire time in the top 7 bytes of the uuid
+        // since the uuid stores metadata in bytes 6 and 8, we only
+        // have 7 consecutive bytes of "free" space, so we discard 1 bytes
+        // from the expire time
+        // we xor the expire with the rng, to make the id look "nicer"
+        let expire_masked = expire ^ dbg!(u64::from_le_bytes(uuid_bytes[0..8].try_into().unwrap()));
+        uuid_bytes[9..].copy_from_slice(&expire_masked.to_le_bytes()[0..7]);
+
+        UploadId(Uuid::from_bytes(uuid_bytes))
+    }
+
+    pub fn get_expire(&self) -> u64 {
+        let uuid_bytes = *self.0.as_bytes();
+        let mut mask_bytes: [u8; 8] = uuid_bytes[0..8].try_into().unwrap();
+        mask_bytes[7] = 0;
+
+        let mut bytes = [0; 8];
+        bytes[0..7].copy_from_slice(&uuid_bytes[9..]);
+        let expire_masked = u64::from_ne_bytes(bytes);
+
+        expire_masked ^ u64::from_le_bytes(mask_bytes)
+    }
+
+    pub fn is_expired(&self, time: u64) -> bool {
+        return time >= self.get_expire();
+    }
+
+    pub fn as_string(&self) -> String {
+        format!("{}", self.0.to_simple())
+    }
+}
+
+#[test]
+fn test_generate_upload_id() {
+    let id = UploadId::generate(12345);
+    assert_eq!(12345, id.get_expire());
+}
+
+#[derive(Clone, Debug, Hash, PartialEq, Eq, Copy)]
+pub struct Expiration(u64);
+
+impl PartialOrd for Expiration {
+    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
+        // reverse the order so sort nearest expiration time first
+        other.0.partial_cmp(&self.0)
+    }
+}
+
+impl Ord for Expiration {
+    fn cmp(&self, other: &Self) -> Ordering {
+        // reverse the order so sort nearest expiration time first
+        other.0.cmp(&self.0)
+    }
+}
+
+impl From<u64> for Expiration {
+    fn from(from: u64) -> Self {
+        Expiration(from)
+    }
+}
+
+#[derive(Debug, Clone)]
+pub struct ExpireQueue {
+    queue: Arc<Mutex<PriorityQueue<UploadId, Expiration>>>,
+}
+
+impl ExpireQueue {
+    pub fn new() -> Self {
+        ExpireQueue {
+            queue: Arc::default(),
+        }
+    }
+
+    pub fn push(&self, key: UploadId) {
+        self.queue
+            .lock()
+            .unwrap()
+            .push(key, key.get_expire().into());
+    }
+
+    pub fn get_expired(&self, since: u64) -> Vec<UploadId> {
+        let expire = Expiration::from(since);
+        let mut queue = self.queue.lock().unwrap();
+
+        let mut expired = Vec::new();
+
+        while queue.peek().map(|(_, exp)| *exp >= expire).unwrap_or(false) {
+            expired.push(queue.pop().map(|(id, _)| id).unwrap())
+        }
+
+        expired
+    }
+
+    pub fn len(&self) -> usize {
+        self.queue.lock().unwrap().len()
+    }
+}
+
+#[test]
+fn test_queue() {
+    let queue = ExpireQueue::new();
+
+    let id1 = UploadId::generate(10);
+    let id2 = UploadId::generate(15);
+
+    queue.push(id1);
+    queue.push(id2);
+
+    assert_eq!(vec![id1], queue.get_expired(12));
+
+    assert_eq!(vec![id2], queue.get_expired(20));
+
+    let id3 = UploadId::generate(10);
+    let id4 = UploadId::generate(15);
+    let id5 = UploadId::generate(20);
+
+    queue.push(id3);
+    queue.push(id4);
+    queue.push(id5);
+
+    assert_eq!(vec![id3, id4, id5], queue.get_expired(20));
+}

src/main.rs

@@ -1,3 +1,130 @@
+#![feature(proc_macro_hygiene, decl_macro)]
+
+use crate::expire_queue::{ExpireQueue, InvalidUploadIdError, UploadId};
+use crate::token::{UploadToken, ValidTokens};
+use dotenv::dotenv;
+use rocket::http::RawStr;
+use rocket::request::FromParam;
+use rocket::response::NamedFile;
+use rocket::*;
+use std::collections::HashMap;
+use std::env;
+use std::env::current_dir;
+use std::fs::{create_dir, read_dir, remove_dir_all};
+use std::io;
+use std::path::PathBuf;
+use std::thread::{sleep, spawn};
+use std::time::{Duration, SystemTime, UNIX_EPOCH};
+
+mod expire_queue;
+mod token;
+
+impl<'r> FromParam<'r> for UploadId {
+    type Error = InvalidUploadIdError;
+
+    fn from_param(param: &'r RawStr) -> Result<Self, Self::Error> {
+        let param = param.url_decode()?;
+
+        UploadId::new(&param)
+    }
+}
+
+#[get("/")]
+fn home() -> &'static str {
+    "Home"
+}
+
+fn now() -> u64 {
+    let start = SystemTime::now();
+    start
+        .duration_since(UNIX_EPOCH)
+        .expect("Time went backwards")
+        .as_secs()
+}
+
+const THOUSAND_YEARS: u64 = 1000 * 356 * 24 * 60 * 60;
+
+#[put("/upload?<expire>&<name>", data = "<data>")]
+fn upload(
+    data: Data,
+    expire: Option<u64>,
+    name: String,
+    _token: UploadToken,
+    basedir: State<PathBuf>,
+    expire_queue: State<ExpireQueue>,
+) -> io::Result<String> {
+    let id = UploadId::generate(now() + expire.unwrap_or(THOUSAND_YEARS));
+    expire_queue.push(id);
+
+    let mut path = basedir.join(id.as_string());
+    create_dir(&path)?;
+    path.push(name);
+
+    data.stream_to_file(path)?;
+    Ok(id.as_string())
+}
+
+#[get("/<id>/<name..>")]
+fn download(id: UploadId, name: PathBuf, basedir: State<PathBuf>) -> Option<NamedFile> {
+    if id.is_expired(now()) {
+        None
+    } else {
+        let path = basedir.join(id.as_string()).join(name);
+        NamedFile::open(path).ok()
+    }
+}
+
 fn main() {
-    println!("Hello, world!");
+    dotenv().ok();
+
+    let mut env: HashMap<_, _> = env::vars().collect();
+    let expire_queue = ExpireQueue::new();
+
+    let tokens: ValidTokens = env
+        .remove("TOKENS")
+        .unwrap_or_default()
+        .split(',')
+        .collect();
+
+    let basedir = env
+        .remove("BASEDIR")
+        .map(PathBuf::from)
+        .unwrap_or_else(|| current_dir().unwrap_or_default().join("data"));
+
+    let expire_basedir = basedir.clone();
+    let expire_queue_clone = expire_queue.clone();
+
+    spawn(move || {
+        for dir_entry in read_dir(&expire_basedir).expect("Failed to list base directory") {
+            if let Ok(entry) = dir_entry {
+                if let Some(upload_id) = entry
+                    .file_name()
+                    .to_str()
+                    .and_then(|s| UploadId::new(s).ok())
+                {
+                    expire_queue_clone.push(upload_id);
+                }
+            }
+        }
+
+        println!(
+            "Loaded {} existing uploads for expiry",
+            expire_queue_clone.len()
+        );
+
+        loop {
+            for expired in expire_queue_clone.get_expired(now()) {
+                let _ = remove_dir_all(expire_basedir.join(expired.as_string()));
+            }
+
+            sleep(Duration::from_secs(5 * 60));
+        }
+    });
+
+    rocket::ignite()
+        .manage(tokens)
+        .manage(basedir.clone())
+        .manage(expire_queue)
+        .mount("/", routes![home, upload, download])
+        .launch();
 }

src/token.rs

@@ -0,0 +1,58 @@
+use err_derive::Error;
+use rocket::http::Status;
+use rocket::outcome::Outcome;
+use rocket::request::{self, FromRequest, Request};
+use rocket::State;
+use std::iter::FromIterator;
+
+#[derive(PartialOrd, PartialEq, Debug)]
+pub struct UploadToken(String);
+
+pub struct ValidTokens(Vec<UploadToken>);
+
+impl<'a> FromIterator<&'a str> for ValidTokens {
+    fn from_iter<T: IntoIterator<Item = &'a str>>(iter: T) -> Self {
+        let tokens: Vec<_> = iter
+            .into_iter()
+            .map(|token| token.to_string())
+            .map(UploadToken)
+            .collect();
+        ValidTokens(tokens)
+    }
+}
+
+impl ValidTokens {
+    pub fn contains(&self, token: &str) -> bool {
+        self.0.iter().any(|accepted| accepted.0 == token)
+    }
+}
+
+#[derive(Debug, Error)]
+pub enum UploadTokenError {
+    #[error(display = "Wrong number of upload tokens provided")]
+    BadCount,
+    #[error(display = "No upload token provided")]
+    Missing,
+    #[error(display = "Invalid upload token")]
+    Invalid,
+}
+
+impl<'a, 'r> FromRequest<'a, 'r> for UploadToken {
+    type Error = UploadTokenError;
+
+    fn from_request(request: &'a Request<'r>) -> request::Outcome<Self, Self::Error> {
+        let accepted_tokens = request
+            .guard::<State<ValidTokens>>()
+            .expect("No tokens configured");
+        let keys: Vec<_> = request.headers().get("x-upload-token").collect();
+
+        match keys.len() {
+            0 => Outcome::Failure((Status::Unauthorized, UploadTokenError::Missing)),
+            1 if accepted_tokens.contains(keys[0]) => {
+                Outcome::Success(UploadToken(keys[0].to_string()))
+            }
+            1 => Outcome::Failure((Status::Unauthorized, UploadTokenError::Invalid)),
+            _ => Outcome::Failure((Status::BadRequest, UploadTokenError::BadCount)),
+        }
+    }
+}