moparisthebest/rust-bucket
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

More refactoring

Browse Source
This commit is contained in:
Travis Burtrum 2018-08-16 23:33:37 -04:00
parent d99445ab11
commit 2b62a7e1da
1 changed files with 58 additions and 40 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

98
src/main.rs
View File

@ -32,7 +32,7 @@ use mpu::MultipartUpload;
const HOST: &'static str = "http://localhost:8000";
const UPLOAD_MAX_SIZE: u64 = 8 * 1024 * 1024;
type Result<T> = ::std::result::Result<T, Error>;
#[derive(Debug)]
pub enum Error {
@ -67,7 +67,7 @@ impl From<toml::de::Error> for Error {
}
}
#[derive(Deserialize, Serialize, Default, Debug)]
#[derive(Deserialize, Serialize, Debug)]
struct PasteInfo<'a> {
// these never change
key: Option<PasteID<'a>>, // key to update/delete paste with
@ -79,8 +79,21 @@ struct PasteInfo<'a> {
num_views: u32, // Only incremented if delete_after_num_views is Some, otherwise 0
}
impl<'a> Default for PasteInfo<'a> {
fn default() -> Self {
PasteInfo {
key: None,
delete_after: Some(SystemTime::now() + Duration::from_secs(2592000)), // default to 30 days
delete_after_num_views:None,
delete_if_not_viewed_in_last_seconds: None,
last_viewed: None,
num_views: 0,
}
}
}
impl<'a> PasteInfo<'a> {
fn read<P: AsRef<Path>>(path: P) -> ::std::result::Result<PasteInfo<'static>,Error> {
fn read<P: AsRef<Path>>(path: P) -> Result<PasteInfo<'static>> {
let mut f = File::open(path)?;
let mut input = String::new();
f.read_to_string(&mut input)?;
@ -88,7 +101,7 @@ impl<'a> PasteInfo<'a> {
Ok(paste_info)
}
fn write<P: AsRef<Path>>(&self, path: P) -> ::std::result::Result<(),Error> {
fn write<P: AsRef<Path>>(&self, path: P) -> Result<()> {
let toml = toml::to_string(&self)?;
fs::write(path, toml)?;
Ok(())
@ -101,7 +114,7 @@ impl<'a> PasteInfo<'a> {
|| self.delete_if_not_viewed_in_last_seconds.map(|n| (SystemTime::now() - n) > self.last_viewed.unwrap()).unwrap_or(false)
}
fn mark_viewed_and_write<P: AsRef<Path>>(&mut self, path: P) -> ::std::result::Result<(),Error> {
fn mark_viewed_and_write<P: AsRef<Path>>(&mut self, path: P) -> Result<()> {
let mut must_write = false;
if self.delete_after_num_views.is_some() {
must_write = true;
@ -178,27 +191,48 @@ impl Backend {
(format!("upload/{id}/file", id = id), format!("upload/{id}/info", id = id), format!("upload/{id}", id = id))
}
fn upload(&self, paste: Data, _key: Option<PasteID>) -> io::Result<String> {
fn upload(&self, paste: Data, _key: Option<PasteID>) -> Result<String> {
let (filename, info_filename, url) = self.new_paste();
PasteInfo::default().write(info_filename)?;
paste.stream_to_file(Path::new(&filename))?;
Ok(url)
}
fn upload_multipart(&self, paste: MultipartUpload) -> io::Result<String> {
fn upload_multipart(&self, paste: MultipartUpload) -> Result<String> {
let (filename, info_filename, url) = self.new_paste();
PasteInfo::default().write(info_filename)?;
paste.stream_to_file(Path::new(&filename))?;
Ok(url)
}
fn upload_string(&self, paste: &PasteForm) -> ::std::result::Result<String,Error> {
fn upload_string(&self, paste: &PasteForm) -> Result<String> {
let (filename, info_filename, url) = self.new_paste();
PasteInfo::from(paste).write(info_filename)?;
fs::write(filename, &paste.content)?;
Ok(url)
}
fn get(&self, id: PasteID) -> ::std::result::Result<content::Plain<File>,Error> {
fn upload_tcp_stream(&self, mut stream: TcpStream) -> Result<()>
{
let (filename, info_filename, url) = self.new_paste();
PasteInfo::default().write(info_filename)?;
let mut paste_file = File::create(&filename)?;
let timeout = Some(Duration::new(5, 0)); // todo: make this config store in struct
stream.set_read_timeout(timeout)?;
stream.set_write_timeout(timeout)?;
stream.write(&url.into_bytes())?;
stream.flush()?;
let upload_max_size: u64 = 8 * 1024 * 1024; // todo: make this config store in struct
copy(&mut stream, &mut paste_file, upload_max_size)?;
Ok(())
}
fn get(&self, id: PasteID) -> Result<content::Plain<File>> {
let (filename, info_filename, foldername) = self.file_paths(id);
let mut paste_info = PasteInfo::read(&info_filename)?;
// first check if we should delete this
@ -247,43 +281,43 @@ impl<'a> From<&'a PasteForm> for PasteInfo<'a> {
// todo: change /w to /, shouldn't conflict because of format, but it does currently
#[post("/w", format = "application/x-www-form-urlencoded", data = "<paste>")]
fn web_post(backend: &Backend, paste: LenientForm<PasteForm>) -> ::std::result::Result<String,Error> {
fn web_post(backend: &Backend, paste: LenientForm<PasteForm>) -> Result<String> {
backend.upload_string(paste.get())
}
// todo: change /w to /, shouldn't conflict because of format, but it does currently
#[post("/m", format = "multipart/form-data", data = "<paste>")]
fn mpu_post(backend: &Backend, paste: MultipartUpload) -> io::Result<String> {
fn mpu_post(backend: &Backend, paste: MultipartUpload) -> Result<String> {
backend.upload_multipart(paste)
}
#[put("/", data = "<paste>")]
fn upload_put(backend: &Backend, paste: Data) -> io::Result<String> {
fn upload_put(backend: &Backend, paste: Data) -> Result<String> {
backend.upload(paste, None)
}
#[post("/", data = "<paste>")]
fn upload_post(backend: &Backend, paste: Data) -> io::Result<String> {
fn upload_post(backend: &Backend, paste: Data) -> Result<String> {
backend.upload(paste, None)
}
#[patch("/", data = "<paste>")]
fn upload_patch(backend: &Backend, paste: Data) -> io::Result<String> {
fn upload_patch(backend: &Backend, paste: Data) -> Result<String> {
backend.upload(paste, None)
}
#[put("/<key>", data = "<paste>")]
fn upload_put_key(backend: &Backend, paste: Data, key: PasteID) -> io::Result<String> {
fn upload_put_key(backend: &Backend, paste: Data, key: PasteID) -> Result<String> {
backend.upload(paste, Some(key))
}
#[post("/<key>", data = "<paste>")]
fn upload_post_key(backend: &Backend, paste: Data, key: PasteID) -> io::Result<String> {
fn upload_post_key(backend: &Backend, paste: Data, key: PasteID) -> Result<String> {
backend.upload(paste, Some(key))
}
#[patch("/<key>", data = "<paste>")]
fn upload_patch_key(backend: &Backend, paste: Data, key: PasteID) -> io::Result<String> {
fn upload_patch_key(backend: &Backend, paste: Data, key: PasteID) -> Result<String> {
backend.upload(paste, Some(key))
}
@ -299,8 +333,9 @@ fn delete(id: PasteID, _key: PasteID) -> Option<content::Plain<File>> {
}
#[get("/")]
fn index() -> io::Result<NamedFile> {
NamedFile::open("static/index.html")
fn index() -> Result<NamedFile> {
let index = NamedFile::open("static/index.html")?;
Ok(index)
}
#[get("/static/<file..>")]
@ -320,7 +355,7 @@ fn rocket() -> rocket::Rocket {
}
// adapted from io::copy
fn copy<R: ?Sized, W: ?Sized>(reader: &mut R, writer: &mut W) -> io::Result<u64>
fn copy<R: ?Sized, W: ?Sized>(reader: &mut R, writer: &mut W, upload_max_size: u64) -> io::Result<u64>
where R: Read, W: Write
{
let mut buf : [u8; 8192] = [0; 8192];
@ -335,40 +370,23 @@ fn copy<R: ?Sized, W: ?Sized>(reader: &mut R, writer: &mut W) -> io::Result<u64>
};
writer.write_all(&buf[..len])?;
written += len as u64;
if written > UPLOAD_MAX_SIZE {
if written > upload_max_size {
return Err(std::io::Error::from(std::io::ErrorKind::InvalidData))
}
}
}
fn consume_paste(backend: &Backend, mut stream: TcpStream, timeout: Option<Duration>) -> io::Result<()>
{
let (filename, _info_filename, url) = backend.new_paste();
let mut paste_file = File::create(&filename)?;
stream.set_read_timeout(timeout)?;
stream.set_write_timeout(timeout)?;
stream.write(&url.into_bytes())?;
stream.flush()?;
copy(&mut stream, &mut paste_file)?;
Ok(())
}
fn run_tcp(){
// Bind the server's socket
thread::spawn(|| {
let backend = &Backend::PlainFile;
let timeout = Some(Duration::new(5, 0));
let listener = TcpListener::bind("127.0.0.1:12345").unwrap();
loop {
match listener.accept() {
Ok((mut stream, _addr)) => {
thread::spawn(move || {
consume_paste(backend, stream, timeout).is_ok(); // again we don't care about this error
backend.upload_tcp_stream(stream).is_ok(); // again we don't care about this error
});
},
Err(_e) => {