From 00d1ff37ba51c1cb7625b5665233afe7c39a7151 Mon Sep 17 00:00:00 2001 From: root Date: Sun, 2 Nov 2025 00:11:16 +0000 Subject: [PATCH] first commit --- .gitignore | 3 + Cargo.toml | 18 +++++ src/index.rs | 188 +++++++++++++++++++++++++++++++++++++++++++++++++++ src/main.rs | 37 ++++++++++ src/serve.rs | 123 +++++++++++++++++++++++++++++++++ 5 files changed, 369 insertions(+) create mode 100644 .gitignore create mode 100644 Cargo.toml create mode 100644 src/index.rs create mode 100644 src/main.rs create mode 100644 src/serve.rs diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..c0c260e --- /dev/null +++ b/.gitignore @@ -0,0 +1,3 @@ +/target +.env +Cargo.lock diff --git a/Cargo.toml b/Cargo.toml new file mode 100644 index 0000000..f83858f --- /dev/null +++ b/Cargo.toml @@ -0,0 +1,18 @@ +[package] +name = "file-srv2" +version = "0.1.0" +edition = "2024" + +[dependencies] +axum = "0.8.6" +dotenv = "0.15.0" +indicatif = "0.18.2" +meilisearch-sdk = "0.30.0" +serde = { version = "1.0.228", features = ["derive"] } +serde_json = "1.0.145" +tokio = { version = "1.48.0", features = ["macros", "rt-multi-thread"] } +tower = { version = "0.5.2", features = ["util"] } +tower-http = { version = "0.6.6", features = ["fs", "trace"] } +tracing = "0.1.41" +tracing-subscriber = "0.3.20" +walkdir = "2.5.0" diff --git a/src/index.rs b/src/index.rs new file mode 100644 index 0000000..b4df3ef --- /dev/null +++ b/src/index.rs @@ -0,0 +1,188 @@ +use indicatif::{ProgressBar, ProgressStyle}; +use meilisearch_sdk::client::Client; +use meilisearch_sdk::settings::Settings; +use std::{collections::hash_map::ExtractIf, path::PathBuf, time::Duration}; +use tokio::{fs, time::sleep}; +use tracing::{error, info}; + +#[derive(serde::Serialize)] +pub struct FileEntry { + pub name: String, + pub ext: String, + pub path: String, + pub url: String, + pub size: i64, + pub preview: String, + pub id: i64, +} + +pub async fn index( + serve_path: &String, + url: &String, + meilisearch_key: &String, + meilisearch_url: &String, +) { + let pb = ProgressBar::new(0); + pb.set_style(ProgressStyle::default_bar() + .template("{msg}\n{spinner:.green} [{elapsed_precise}] [{wide_bar:.cyan/blue}] {pos}/{len} ({per_sec}, {eta})").unwrap() + .progress_chars("#>-")); + pb.set_message("Running index..."); + + // Main loop + loop { + info!("Running index"); + delete_index(meilisearch_key, meilisearch_url).await; + + let mut i = 0; + let mut idx = 0; + let mut failed_preview = 0; + let mut files: Vec = Vec::new(); + for res in walkdir::WalkDir::new(&serve_path) { + let entry = match res { + Ok(entry) => entry, + Err(e) => { + error!("Failed to get entry: {}", e); + continue; + } + }; + // Make sure the entry is a file + if entry.file_type().is_dir() { + continue; + } + + let ext = entry.path().extension(); + let ext = match ext { + None => "none".to_string(), + Some(ok) => ok.to_string_lossy().to_string(), + }; + + if ext != "zip" { + continue; + } + + let name = entry.file_name().to_string_lossy().to_string(); + let path = entry.path().to_string_lossy().to_string(); + let size = entry.metadata().unwrap().len() as i64; + let file_url = format!("{}{}", url, path[serve_path.len() + 1..].to_string()); + let parent_dir = path[..path.len() - ext.len() - 1].to_string(); + + let mut preview = "None".to_string(); + + for file in walkdir::WalkDir::new(&parent_dir) { + let f = match file { + Ok(f) => f, + Err(e) => { + error!("Failed to read file while looking for preview: {}", e); + failed_preview += 1; + continue; + } + }; + + let file_name = f.file_name().to_string_lossy().to_string(); + if file_name.contains("Preview") { + let preview_path = f.path().to_string_lossy().to_string(); + preview = format!( + "{}{}", + url, + preview_path[serve_path.len() + 1..].to_string() + ); + } + } + + let file_entry = FileEntry { + name, + ext, + path, + url: file_url, + size, + preview, + id: idx, + }; + + files.push(file_entry); + i += 1; + idx += 1; + pb.inc(1); + if i > 10000 { + send_index_chunk(files, meilisearch_key, meilisearch_url).await; + files = Vec::new(); + i = 0; + } + } + + // Push any removing entries + send_index_chunk(files, meilisearch_key, meilisearch_url).await; + pb.finish_and_clear(); + info!("Done! Total: {}, Failed previews: {}", idx, failed_preview); + sleep(Duration::from_secs(900)).await; + } +} + +async fn send_index_chunk( + files: Vec, + meilisearch_key: &String, + meilisearch_url: &String, +) { + let client = match Client::new(meilisearch_url, Some(meilisearch_key)) { + Ok(client) => client, + Err(e) => { + error!("Failed to create meilisearch client: \n\t{}", e); + return; + } + }; + + let searchable_attributes = ["name", "path", "ext"]; + let ranking_rules = ["words", "typo", "attribute", "exactness", "cost::asc"]; + + let settings = Settings::new() + .with_ranking_rules(ranking_rules) + .with_searchable_attributes(searchable_attributes); + + let result = client + .index("shared_files") + .set_settings(&settings) + .await + .unwrap() + .wait_for_completion(&client, None, None) + .await + .unwrap(); + + if result.is_failure() { + error!("Failed to set index: \n\t{:?}", result.unwrap_failure()); + return; + } + + let result = client + .index("shared_files") + .add_or_update(&files, Some("id")) + .await + .unwrap() + .wait_for_completion(&client, None, None) + .await + .unwrap(); + + if result.is_failure() { + error!("Failed to add files: \n\t{:?}", result.unwrap_failure()); + return; + } +} + +async fn delete_index(meilisearch_key: &String, meilisearch_url: &String) { + let client = match Client::new(meilisearch_url, Some(meilisearch_key)) { + Ok(ok) => ok, + Err(e) => { + error!("Failed to create meilisearch client: \n\t{}", e); + return; + } + }; + + let index = client.index("shared_files"); + + index + .delete() + .await + .unwrap() + .wait_for_completion(&client, None, None) + .await + .unwrap(); +} diff --git a/src/main.rs b/src/main.rs new file mode 100644 index 0000000..af8fb40 --- /dev/null +++ b/src/main.rs @@ -0,0 +1,37 @@ +mod index; +mod serve; + +use std::process::exit; +use tracing::{debug, error, info}; + +#[tokio::main] +async fn main() { + tracing_subscriber::fmt::init(); + debug!("Loading .env file"); + if let Err(e) = dotenv::dotenv() { + error!("Failed to get .env file: {}", e); + exit(1); + } + + let bind_addr = std::env::var("BIND_ADDR").unwrap_or_else(|_| "127.0.0.1:8080".to_string()); + let serve_path = std::env::var("SERVE_PATH").unwrap_or_else(|_| "./files/".to_string()); + let url = std::env::var("URL").unwrap_or_else(|_| "http://localhost:8080/files".to_string()); + let meilisearch_url = + std::env::var("MEILISEARCH_URL").unwrap_or_else(|_| "http://localhost:7700".to_string()); + let meilisearch_key = match std::env::var("MEILISEARCH_KEY") { + Ok(key) => key, + Err(e) => { + error!( + "Failed to get meilisearch key from environment. Error: \n\t{}\n\n Exiting...", + e + ); + exit(1); + } + }; + + info!("Starting services..."); + tokio::join!( + serve::serve(bind_addr, &serve_path, &meilisearch_key, &meilisearch_url), + index::index(&serve_path, &url, &meilisearch_key, &meilisearch_url), + ); +} diff --git a/src/serve.rs b/src/serve.rs new file mode 100644 index 0000000..903ef65 --- /dev/null +++ b/src/serve.rs @@ -0,0 +1,123 @@ +use axum::Router; +use axum::extract::{Query, State}; +use axum::routing::get; +use meilisearch_sdk::client::Client; +use meilisearch_sdk::search::SearchResult; +use serde::{Deserialize, Serialize}; +use std::net::SocketAddr; +use tower_http::services::ServeDir; +use tower_http::trace::TraceLayer; +use tracing::{error, info}; + +#[derive(Debug, Deserialize)] +pub struct SearchInput { + pub q: String, + pub p: Option, + pub s: Option, +} + +#[derive(Debug, Clone, Serialize)] +pub struct Root { + pub results: Vec>, + pub metadata: Metadata, +} + +#[derive(Serialize, Debug, Clone)] +pub struct Metadata { + pub page: usize, + pub total_pages: usize, + pub page_size: usize, +} + +#[derive(Clone)] +pub struct AppContext { + pub meilisearch_client: Client, +} + +#[derive(Serialize, Deserialize, Debug, Clone)] +pub struct FileEntryDisplay { + pub name: String, + pub ext: String, + pub path: String, + pub url: String, + pub size: i64, + pub preview: String, +} + +pub async fn serve( + addr: String, + path: &String, + meilisearch_key: &String, + meilisearch_url: &String, +) { + let addr = SocketAddr::from(addr.parse::().unwrap()); + let listener = tokio::net::TcpListener::bind(addr).await.unwrap(); + let client = match Client::new(meilisearch_url, Some(meilisearch_key)) { + Ok(client) => client, + Err(e) => { + error!( + "Failed to create meilisearch client for http context: {}", + e + ); + return; + } + }; + + let context = AppContext { + meilisearch_client: client, + }; + let app = Router::new() + .nest_service("/files", ServeDir::new(path)) + .route("/search", get(search)) + .with_state(context); + tracing::info!("listening on {}", listener.local_addr().unwrap()); + axum::serve(listener, app.layer(TraceLayer::new_for_http())) + .await + .unwrap(); +} + +pub async fn search(query: Query, state: State) -> Result { + info!( + "Got query for {} p: {:?} s: {:?}", + query.q, query.p, query.s + ); + let client = &state.meilisearch_client; + + let size = query.s; + let index = client.index("shared_files"); + let mut search = index.search(); + let request = search.with_query(query.q.as_str()); + + if let Some(page) = query.p { + request.page = Some(page); + } + if let Some(size) = query.s { + request.hits_per_page = Some(size); + } + + let res = match request.execute::().await { + Ok(ok) => ok, + Err(e) => { + return Err(format!("Failed with error: {}", e)); + } + }; + + let hits = res.hits; + let total_pages = res.total_pages; + + let metadata = Metadata { + page: res.page.unwrap_or(0), + total_pages: total_pages.unwrap_or(0), + page_size: size.unwrap_or(0), + }; + + let root = Root { + results: hits, + metadata, + }; + + match serde_json::to_string(&root) { + Ok(json_str) => Ok(json_str), + Err(e) => Err(format!("Failed to serialize: {}", e)), + } +}