use std::{ collections::{BTreeSet, HashMap}, fs::File, io::{Read, Write}, num::NonZeroU32, ops::DerefMut, sync::atomic::Ordering }; use fast_image_resize as fir; use image::DynamicImage; use rustracing_jaeger::Span; use tiny_http::{Request, Response, ResponseBox, StatusCode}; use crate::{ db, db::DBConnection, dto, header, metrics, routes::{filters::UserInfo, fs::resize_dimensions, get_reply, AppError, ChannelReader} }; pub fn root(_: &Span, _: &mut Request, _: &mut DBConnection, info: UserInfo) -> Result { get_reply(&dto::responses::Root { statusCode: 200, rootId: info.0.root_id }) } pub fn node( span: &Span, _: &mut Request, db: &mut DBConnection, info: UserInfo, node: i32 ) -> Result { let guard_lock = DBConnection::get_lock(info.0.id); let _guard = guard_lock.read(span); let node = super::get_node_and_validate(span, &info.0, node, db).ok_or(AppError::BadRequest("Unknown node"))?; get_reply(&dto::responses::GetNode { statusCode: 200, id: node.id, name: node.name, isFile: node.is_file, preview: node.has_preview, parent: node.parent_id, size: node.size, children: (!node.is_file).then(|| { db.get_children(span, node.id) .iter() .map(|child| dto::responses::GetNodeEntry { id: child.id, name: child.name.clone(), isFile: child.is_file, preview: child.has_preview, parent: child.parent_id, size: child.size }) .collect() }) }) } pub fn path( span: &Span, _: &mut Request, db: &mut DBConnection, info: UserInfo, node: i32 ) -> Result { let guard_lock = DBConnection::get_lock(info.0.id); let _guard = guard_lock.read(span); let node = super::get_node_and_validate(span, &info.0, node, db).ok_or(AppError::BadRequest("Unknown node"))?; get_reply(&super::generate_path_dto(span, &node, db)) } pub fn create_node( span: &Span, _: &mut Request, db: &mut DBConnection, info: UserInfo, data: dto::requests::CreateNode, file: bool ) -> Result { let guard_lock = DBConnection::get_lock(info.0.id); let _guard = guard_lock.read(span); let node = super::create_node(span, data.name, &info.0, file, Some(data.parent), false, db); match node { Ok(v) => get_reply(&dto::responses::NewNode { statusCode: 200, id: v.id }), Err(v) => match v { super::CreateNodeResult::InvalidName => AppError::BadRequest("Invalid name").err(), super::CreateNodeResult::InvalidParent => AppError::BadRequest("Invalid parent").err(), super::CreateNodeResult::Exists(file, id) => get_reply(&dto::responses::NodeExists { statusCode: 200, id, exists: true, isFile: file }) } } } pub fn delete_node( span: &Span, _: &mut Request, db: &mut DBConnection, info: UserInfo, node: i32, pool: &db::DBPool ) -> Result { let guard_lock = DBConnection::get_lock(info.0.id); let inner_guard_lock = guard_lock.clone(); let _guard = guard_lock.read(span); let node = super::get_node_and_validate(span, &info.0, node, db).ok_or(AppError::BadRequest("Unknown node"))?; if node.parent_id.is_none() { return AppError::BadRequest("Can't delete root").err(); } let (tx, rx) = std::sync::mpsc::channel::(); let inner_pool = pool.clone(); tx.send("Waiting in queue\n".to_owned()).unwrap(); let inner_span = metrics::span("DELETE_POOL - queueing", span); super::DELETE_POOL.spawn(move || { let del_span = metrics::span("DELETE_POOL - deleting", &inner_span); drop(inner_span); let mut db = &mut DBConnection::from(inner_pool.get().unwrap()); let guard_lock = inner_guard_lock.clone(); let _guard = guard_lock.write(&del_span); super::delete_node(&del_span, &node, &tx, db.deref_mut()); }); Ok(Response::new( StatusCode::from(200), vec![header("content-type", "text/plain; charset=utf-8")], ChannelReader(rx), None, None ) .boxed()) } pub fn upload( span: &Span, req: &mut Request, db: &mut DBConnection, info: UserInfo, node: i32 ) -> Result { let guard_lock = DBConnection::get_lock(info.0.id); let _guard = guard_lock.read(span); let mut node = super::get_node_and_validate(span, &info.0, node, db).ok_or(AppError::BadRequest("Unknown node"))?; if !node.is_file { return AppError::BadRequest("Can't upload to a directory node").err(); } let mut file_size = 0_i64; let file_name = format!("./files/{}", node.id); let mut file_buf = Vec::::new(); { let _span = metrics::span("receive_file", span); let mut buf = vec![0_u8; 8 * 1024 * 1024]; let mut file = File::create(file_name.clone()).unwrap(); let reader = req.as_reader(); loop { let r = reader.read(&mut buf).unwrap(); if r == 0 { break; } if file_size < 20 * 1024 * 1024 { file_buf.write_all(&buf[..r]).unwrap(); } file.write_all(&buf[..r]).unwrap(); file_size += r as i64; } } metrics::DISK_USAGE .with_label_values(&[node.owner_id.to_string().as_str()]) .add(file_size - node.size.unwrap_or(0)); { let prev_span = metrics::span("generate_preview", span); node.has_preview = (|| { if file_size > file_buf.len() as i64 { return None; } let mime = mime_guess::from_path(std::path::Path::new(&node.name)).first()?.to_string(); let img = { let _span = metrics::span("generate_preview_load", &prev_span); image::load_from_memory_with_format( file_buf.as_slice(), image::ImageFormat::from_mime_type(mime)? ) .ok()? }; let img = { let _span = metrics::span("generate_preview_convert", &prev_span); let width = NonZeroU32::try_from(img.width()).unwrap(); let height = NonZeroU32::try_from(img.height()).unwrap(); match img { DynamicImage::ImageLuma8(v) => fir::Image::from_vec_u8(width, height, v.into_raw(), fir::PixelType::U8), DynamicImage::ImageLumaA8(v) => fir::Image::from_vec_u8(width, height, v.into_raw(), fir::PixelType::U8x2), DynamicImage::ImageRgb8(v) => fir::Image::from_vec_u8(width, height, v.into_raw(), fir::PixelType::U8x3), DynamicImage::ImageRgba8(v) => fir::Image::from_vec_u8(width, height, v.into_raw(), fir::PixelType::U8x4), DynamicImage::ImageLuma16(_) => fir::Image::from_vec_u8(width, height, img.to_luma8().into_raw(), fir::PixelType::U8), DynamicImage::ImageLumaA16(_) => fir::Image::from_vec_u8( width, height, img.to_luma_alpha8().into_raw(), fir::PixelType::U8x2 ), DynamicImage::ImageRgb16(_) => fir::Image::from_vec_u8( width, height, img.to_rgb8().into_raw(), fir::PixelType::U8x3 ), DynamicImage::ImageRgba16(_) => fir::Image::from_vec_u8( width, height, img.to_rgba8().into_raw(), fir::PixelType::U8x4 ), DynamicImage::ImageRgb32F(_) => fir::Image::from_vec_u8( width, height, img.to_rgb8().into_raw(), fir::PixelType::U8x3 ), DynamicImage::ImageRgba32F(_) => fir::Image::from_vec_u8( width, height, img.to_rgba8().into_raw(), fir::PixelType::U8x4 ), _ => fir::Image::from_vec_u8( width, height, img.to_rgba8().into_raw(), fir::PixelType::U8x4 ) } .expect("Failed to convert preview image") }; let img = { let _span = metrics::span("generate_preview_resize", &prev_span); let new_dim = resize_dimensions(img.width().get(), img.height().get()); let mut dst = fir::Image::new( NonZeroU32::try_from(new_dim.0).unwrap(), NonZeroU32::try_from(new_dim.1).unwrap(), img.pixel_type() ); fir::Resizer::new(fir::ResizeAlg::SuperSampling(fir::FilterType::Hamming, 2)) .resize(&img.view(), &mut dst.view_mut()) .expect("Failed to resize preview image"); dst }; let _span = metrics::span("generate_preview_save", &prev_span); let mut file = std::io::BufWriter::new( File::create(std::path::Path::new(&(file_name + "_preview.jpg"))) .expect("Failed to open preview image file") ); image::codecs::jpeg::JpegEncoder::new(&mut file) .encode( img.buffer(), img.width().get(), img.height().get(), match img.pixel_type() { fir::PixelType::U8 => image::ColorType::L8, fir::PixelType::U8x2 => image::ColorType::La8, fir::PixelType::U8x3 => image::ColorType::Rgb8, fir::PixelType::U8x4 => image::ColorType::Rgba8, _ => unreachable!() } ) .expect("Failed to save preview image"); Some(()) })() .is_some(); } node.size = Some(file_size); db.save_node(span, &node); get_reply(&dto::responses::Success { statusCode: 200 }) } pub fn create_zip( span: &Span, _: &mut Request, db: &mut DBConnection, info: UserInfo, data: dto::requests::CreateZip, pool: &db::DBPool ) -> Result { let guard_lock = DBConnection::get_lock(info.0.id); let inner_guard_lock = guard_lock.clone(); let _guard = guard_lock.read(span); let mut nodes: Vec = Vec::new(); for node in data.nodes.clone() { nodes.push(super::get_node_and_validate(span, &info.0, node, db).ok_or(AppError::BadRequest("Unknown node"))?); } let zip_nodes = BTreeSet::from_iter(data.nodes.iter().copied()); { let guard = super::ZIP_TO_PROGRESS.read(); if let Some(entry) = guard.get(&zip_nodes) { return get_reply(&dto::responses::CreateZipDone { statusCode: 200, done: entry.done.load(Ordering::Relaxed), progress: Some(entry.progress.load(Ordering::Relaxed)), total: Some(entry.total.load(Ordering::Relaxed)) }); } } let entry = { let mut guard = super::ZIP_TO_PROGRESS.write(); guard.insert( zip_nodes.clone(), std::sync::Arc::from(super::ZipProgressEntry { temp_id: super::NEXT_TEMP_ID.fetch_add(1, Ordering::Relaxed), done: std::sync::atomic::AtomicBool::new(false), progress: std::sync::atomic::AtomicU64::new(0), total: std::sync::atomic::AtomicU64::new(1), delete_after: std::sync::atomic::AtomicI64::new(0) }) ); guard.get(&zip_nodes).unwrap().clone() }; let inner_pool = pool.clone(); let inner_span = metrics::span("ZIP_POOL - queueing", span); super::ZIP_POOL.spawn(move || { let mut zip_span = metrics::span("ZIP_POOL - zipping", &inner_span); drop(inner_span); let db = &mut DBConnection::from(inner_pool.get().unwrap()); type NodeMap = HashMap; super::cleanup_temp_zips(&zip_span); let _guard = inner_guard_lock.read(&zip_span); fn get_path(node: &db::Inode, dirs: &NodeMap) -> String { let mut path = node.name.clone(); let mut _node = dirs.get(&node.parent_id.unwrap_or(-1)); while let Some(node) = _node { path.insert_str(0, &(node.name.clone() + "/")); _node = dirs.get(&node.parent_id.unwrap_or(-1)); } path } { let i_span = metrics::span("ZIP_POOL - calc total", &zip_span); nodes.iter().for_each(|node| { entry.total.fetch_add( super::get_total_size(&i_span, node.clone(), db), Ordering::Relaxed ); }); entry.total.fetch_sub(1, Ordering::Relaxed); } { let comp_span = metrics::span("ZIP_POOL - compressing total", &zip_span); let mut buf = vec![0_u8; 16 * 1024 * 1024]; let file = File::create(format!("./temp/{}", entry.temp_id)).expect("Failed to create temp file"); let mut zip = zip::ZipWriter::new(file); let zip_options = zip::write::FileOptions::default().large_file(true); let (files, dirs): (NodeMap, NodeMap) = { let _span = metrics::span("ZIP_POOL - gathering nodes", &comp_span); nodes .iter() .flat_map(|node| super::get_nodes_recursive(&comp_span, node.clone(), db)) .map(|node| (node.id, node)) .partition(|v| v.1.is_file) }; zip_span.set_tags(|| { vec![ rustracing::tag::Tag::new("files", files.len() as i64), rustracing::tag::Tag::new("dirs", dirs.len() as i64), ] }); { let _span = metrics::span("ZIP_POOL - dirs", &comp_span); dirs.values().for_each(|dir| { zip.add_directory(get_path(dir, &dirs), zip_options).expect("Failed to add dir to zip"); }); } { let _span = metrics::span("ZIP_POOL - files", &comp_span); files.values().for_each(|node| { zip.start_file(get_path(node, &dirs), zip_options).expect("Failed to start zip file"); let mut file = File::open(format!("./files/{}", node.id)).expect("Failed to open file for zip"); loop { let count = file.read(&mut buf).expect("Failed to read file for zip"); if count == 0 { break; } zip.write_all(&buf[..count]).expect("Failed to write zip"); entry.progress.fetch_add(count as u64, Ordering::Relaxed); } }); } zip.finish().expect("Failed to finish zip"); } entry.done.store(true, Ordering::Relaxed); entry.delete_after.store(chrono::Utc::now().timestamp() + 10 * 60, Ordering::Relaxed); }); get_reply(&dto::responses::CreateZipDone { statusCode: 200, done: false, progress: Some(0), total: Some(1) }) } pub fn download(span: &Span, req: &mut Request, db: &mut DBConnection) -> Result { let data: dto::requests::Download = serde_urlencoded::from_reader(req.as_reader()).map_err(|_| AppError::BadRequest("Invalid form data"))?; let info = crate::routes::filters::authorize_jwt(span, &data.jwtToken, db)?; let guard_lock = DBConnection::get_lock(info.0.id); let _guard = guard_lock.read(span); let node: db::Inode = super::get_node_and_validate(span, &info.0, data.id, db).ok_or(AppError::BadRequest("Unknown node"))?; if node.is_file { let file_name = format!("./files/{}", node.id); let resp = Response::from_file(File::open(std::path::Path::new(&file_name)).unwrap()) .with_header(header("content-type", "application/octet-stream")) .with_header(header( "content-disposition", &("attachment; filename=".to_owned() + &node.name) )) .boxed(); Ok(resp) } else { let nodes_key = BTreeSet::from([node.id]); let guard = super::ZIP_TO_PROGRESS.read(); let entry = guard.get(&nodes_key).ok_or(AppError::BadRequest("Unknown node"))?; if !entry.done.load(Ordering::Relaxed) { AppError::BadRequest("Unknown node").err() } else { let file_name = format!("./temp/{}", entry.temp_id); let resp = Response::from_file(File::open(std::path::Path::new(&file_name)).unwrap()) .with_header(header("content-type", "application/zip")) .with_header(header( "content-disposition", &("attachment; filename=".to_owned() + &node.name + ".zip") )) .boxed(); Ok(resp) } } } pub fn download_multi(span: &Span, req: &mut Request, db: &mut DBConnection) -> Result { let data: dto::requests::DownloadMulti = serde_urlencoded::from_reader(req.as_reader()).map_err(|_| AppError::BadRequest("Invalid form data"))?; let info = crate::routes::filters::authorize_jwt(span, &data.jwtToken, db)?; let guard_lock = DBConnection::get_lock(info.0.id); let _guard = guard_lock.read(span); let nodes: Vec = { let _span = metrics::span("parsing_nodes", span); data.id .split(',') .map(|v| { v.parse::().map_err(|_| AppError::BadRequest("Failed to parse")).map(|n| { super::get_node_and_validate(span, &info.0, n, db).ok_or(AppError::BadRequest("Unknown node")) }) }) .into_iter() .collect::, AppError>, AppError>>()?? }; let nodes_key = BTreeSet::from_iter(nodes.iter().map(|node| node.id)); let guard = super::ZIP_TO_PROGRESS.read(); let entry = guard.get(&nodes_key).ok_or(AppError::BadRequest("Unknown zip"))?; if !entry.done.load(Ordering::Relaxed) { AppError::BadRequest("Unfinished zip").err() } else { let file_name = format!("./temp/{}", entry.temp_id); let resp = Response::from_file(File::open(std::path::Path::new(&file_name)).unwrap()) .with_header(header("content-type", "application/zip")) .with_header(header( "content-disposition", "attachment; filename=files.zip" )) .boxed(); Ok(resp) } } pub fn download_preview( span: &Span, _: &mut Request, db: &mut DBConnection, info: UserInfo, node: i32 ) -> Result { let guard_lock = DBConnection::get_lock(info.0.id); let _guard = guard_lock.read(span); let node: db::Inode = super::get_node_and_validate(span, &info.0, node, db).ok_or(AppError::BadRequest("Unknown node"))?; if node.has_preview { let file = format!("./files/{}_preview.jpg", node.id); get_reply(&dto::responses::DownloadBase64 { statusCode: 200, data: "data:image/png;base64,".to_owned() + &base64::encode(std::fs::read(std::path::Path::new(&file)).unwrap()) }) } else { AppError::BadRequest("No preview").err() } } pub fn get_type( span: &Span, _: &mut Request, db: &mut DBConnection, info: UserInfo, node: i32 ) -> Result { let node: db::Inode = super::get_node_and_validate(span, &info.0, node, db).ok_or(AppError::BadRequest("Unknown node"))?; get_reply(&dto::responses::Type { statusCode: 200, _type: mime_guess::from_path(std::path::Path::new(&node.name)).first_or_octet_stream().to_string() }) } pub fn move_node( span: &Span, _: &mut Request, db: &mut DBConnection, info: UserInfo, data: dto::requests::MoveNode ) -> Result { let guard_lock = DBConnection::get_lock(info.0.id); let _guard = guard_lock.write(span); let target = super::get_node_and_validate(span, &info.0, data.target, db).ok_or(AppError::BadRequest("Invalid target"))?; if target.is_file { return AppError::BadRequest("Invalid target").err(); } let mut nodes = data .nodes .iter() .map(|v| super::get_node_and_validate(span, &info.0, *v, db)) .into_iter() .collect::>>() .ok_or(AppError::BadRequest("Invalid node"))?; for parent in super::get_node_path(span, target.clone(), db) { if nodes.contains(&parent) { return AppError::BadRequest("Can't move node into one of it's subfolders").err(); } } for child in db.get_children(span, target.id) { if nodes.iter().any(|n| n.name == child.name) { return AppError::BadRequest("Can't overwrite existing file").err(); } } nodes.iter_mut().for_each(|n| db.move_node(span, n, target.id)); get_reply(&dto::responses::Success { statusCode: 200 }) }