130 lines
3.4 KiB
Rust
130 lines
3.4 KiB
Rust
extern crate multipart;
|
|
extern crate tiny_http;
|
|
|
|
use blake2::{Blake2s256, Digest};
|
|
use lmdb::LmdbResultExt;
|
|
use lmdb_zero as lmdb;
|
|
use multipart::server::Multipart;
|
|
use std::{
|
|
io::{self, Read},
|
|
sync::Arc,
|
|
thread,
|
|
};
|
|
use tiny_http::{Request, Response};
|
|
|
|
fn main() {
|
|
let env = Arc::new(unsafe {
|
|
lmdb::EnvBuilder::new()
|
|
.unwrap()
|
|
.open("./db", lmdb::open::Flags::empty(), 0o600)
|
|
.unwrap()
|
|
});
|
|
|
|
let db = Arc::new(
|
|
lmdb::Database::open(env.clone(), None, &lmdb::DatabaseOptions::defaults()).unwrap(),
|
|
);
|
|
|
|
let server =
|
|
Arc::new(tiny_http::Server::http("localhost:8000").expect("Could not bind localhost:8000"));
|
|
|
|
let mut guards = Vec::with_capacity(4);
|
|
|
|
for _ in 0..4 {
|
|
let server = server.clone();
|
|
let db = db.clone();
|
|
let env = env.clone();
|
|
|
|
let guard = thread::spawn(move || loop {
|
|
let request = server.recv().unwrap();
|
|
|
|
match request.url() {
|
|
"/upload" => {
|
|
let mut txn = lmdb::WriteTransaction::new(env.clone()).unwrap();
|
|
|
|
process_upload(request, &mut txn, db.clone()).unwrap();
|
|
|
|
txn.commit().unwrap();
|
|
}
|
|
s if s.starts_with("/get/") => {
|
|
let mut txn = lmdb::ReadTransaction::new(env.clone()).unwrap();
|
|
process_get(request, &mut txn, db.clone()).unwrap();
|
|
}
|
|
_ => todo!(),
|
|
}
|
|
});
|
|
|
|
guards.push(guard);
|
|
}
|
|
|
|
for guard in guards {
|
|
guard.join().unwrap();
|
|
}
|
|
}
|
|
|
|
fn process_upload(
|
|
mut request: Request,
|
|
txn: &mut lmdb::WriteTransaction<'_>,
|
|
db: Arc<lmdb::Database<'_>>,
|
|
) -> io::Result<()> {
|
|
if let Ok(form) = Multipart::from_request(&mut request) {
|
|
if let Some(mut entry) = form.into_entry().into_result()? {
|
|
let mut data: Vec<u8> = Vec::with_capacity(20000);
|
|
entry.data.read_to_end(&mut data)?;
|
|
|
|
let data_hash = Blake2s256::digest(&data);
|
|
|
|
let mut accessor = txn.access();
|
|
if accessor
|
|
.get::<[u8], [u8]>(&db, data_hash.as_slice())
|
|
.to_opt()
|
|
.unwrap()
|
|
.is_some()
|
|
{
|
|
println!("found dupe!");
|
|
request.respond(Response::from_string(base64_url::encode(&data_hash)))?;
|
|
return Ok(());
|
|
}
|
|
|
|
accessor
|
|
.put(&db, data_hash.as_slice(), &data, lmdb::put::Flags::empty())
|
|
.unwrap();
|
|
|
|
request.respond(Response::from_string(base64_url::encode(&data_hash)))?;
|
|
}
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
fn process_get(
|
|
request: Request,
|
|
txn: &mut lmdb::ReadTransaction<'_>,
|
|
db: Arc<lmdb::Database<'_>>,
|
|
) -> io::Result<()> {
|
|
let get_id = request
|
|
.url()
|
|
.strip_prefix("/get/")
|
|
.and_then(|s| {
|
|
let mut out: [u8; 32] = [0; 32];
|
|
base64_url::decode_to_slice(s, &mut out).ok()?;
|
|
Some(out)
|
|
})
|
|
.unwrap();
|
|
|
|
let access = txn.access();
|
|
|
|
if let Some(data) = access.get::<[u8], [u8]>(&db, &get_id).to_opt().unwrap() {
|
|
request.respond(Response::new(
|
|
200.into(),
|
|
vec![],
|
|
data,
|
|
Some(data.len()),
|
|
None,
|
|
))?;
|
|
} else {
|
|
request.respond(Response::empty(404))?;
|
|
}
|
|
|
|
Ok(())
|
|
}
|