use std::str::FromStr; use std::path::PathBuf; use std::path::Path; use anyhow::anyhow; use anyhow::Error; use anyhow::Context; use anyhow::Result; use tonic::Request; use tonic::Response; use tonic::Status; use resiter::AndThen; use resiter::Map; use fss::VersionReply; use fss::VersionRequest; use fss::IndexFileRequest; use fss::IndexFileReply; use fss::SearchQueryRequest; use fss::SearchResponse; mod config; mod cli; mod ft; mod schema; mod fss { tonic::include_proto!("fss"); // The string specified here must match the proto package name } #[derive(getset::CopyGetters)] pub struct Server { index: tantivy::Index, #[getset(get_copy = "pub")] field_path: tantivy::schema::Field, #[getset(get_copy = "pub")] field_ft: tantivy::schema::Field, #[getset(get_copy = "pub")] field_size: tantivy::schema::Field, #[getset(get_copy = "pub")] field_created: tantivy::schema::Field, #[getset(get_copy = "pub")] field_modified: tantivy::schema::Field, #[getset(get_copy = "pub")] field_indexed: tantivy::schema::Field, #[getset(get_copy = "pub")] field_body: tantivy::schema::Field, } impl Server { fn write_file_to_index(&self, filepath: &Path) -> Result<()> { let ext = filepath .extension() .map(ToOwned::to_owned) .and_then(|osstr| osstr.to_str().map(|s| s.to_string())) .ok_or_else(|| anyhow!("Path {} is not UTF8", filepath.display()))?; let mut doc = tantivy::Document::default(); doc.add_text(self.field_path(), filepath.display().to_string()); doc.add_text(self.field_ft(), &ext); { use std::os::linux::fs::MetadataExt; let meta = filepath.metadata()?; doc.add_i64(self.field_created(), meta.st_atime()); doc.add_i64(self.field_modified(), meta.st_mtime()); if let Ok(now) = std::time::SystemTime::now().duration_since(std::time::SystemTime::UNIX_EPOCH).map(|d| d.as_secs()) { doc.add_u64(self.field_indexed(), now) } } let doc = self.parse_for_file(doc, filepath, &ext)?; let mut index_writer = self.index.writer(50_000_000)?; index_writer.add_document(doc); index_writer.commit()?; Ok(()) } fn parse_for_file(&self, doc: tantivy::Document, filepath: &Path, ext: &str) -> Result { use crate::ft::FileTypeParser; use crate::ft::TextFileParser; match ext { "txt" => Ok(TextFileParser {}), _ => Err(anyhow!("No parser available for {}", ext)) }? .parse(&self, filepath, ext, doc) } fn search(&self, query_str: &str) -> Result> { let reader = self.index .reader_builder() .reload_policy(tantivy::ReloadPolicy::OnCommit) .try_into()?; let searcher = reader.searcher(); let query_parser = tantivy::query::QueryParser::for_index(&self.index, vec![ self.field_path, self.field_ft, self.field_size, self.field_created, self.field_modified, self.field_indexed, self.field_body, ]); let query = query_parser.parse_query(query_str)?; let top_docs = searcher.search(&query, &tantivy::collector::TopDocs::with_limit(10))?; let results = top_docs .into_iter() .map(|(_score, adr)| searcher.doc(adr).map_err(Error::from)) .and_then_ok(|retr| { retr.get_all(self.field_path) .map(|value| { value.text().ok_or_else(|| anyhow!("Not a text value..")) }) .map_ok(PathBuf::from) .collect::>>() }) .collect::>>>()? .into_iter() .map(Vec::into_iter) .flatten() .collect(); Ok(results) } } #[tonic::async_trait] impl fss::fss_server::Fss for Server { async fn get_version(&self, _request: Request) -> Result, Status> { let reply = fss::VersionReply { version: version::version!().to_string(), }; Ok(Response::new(reply)) } async fn index_file(&self, request: Request) -> Result, Status> { let error = match self.write_file_to_index(request.get_ref().path.as_ref()) { Ok(()) => false, Err(e) => { log::error!("Error writing to index: {} -> {:?}", request.get_ref().path, e); true } }; let reply = fss::IndexFileReply { error }; Ok(Response::new(reply)) } async fn search_query(&self, request: Request) -> Result, Status> { let (success, pathes) = match self.search(&request.get_ref().query) { Ok(pathes) => (true, pathes), Err(e) => { log::error!("Error searching with '{}' -> {:?}", request.get_ref().query, e); (false, vec![]) } }; let reply = fss::SearchResponse { success, pathes: pathes.into_iter().map(|pb| pb.display().to_string()).collect(), }; Ok(Response::new(reply)) } } #[tokio::main] async fn main() -> Result<()> { let cli = crate::cli::server_app(); let matches = cli.get_matches(); let _ = env_logger::try_init()?; let mut config = ::config::Config::default(); { let xdg = xdg::BaseDirectories::with_prefix("fss")?; let xdg_config = xdg.find_config_file("config.toml") .ok_or_else(|| anyhow!("No configuration file found with XDG: {}", xdg.get_config_home().display()))?; log::debug!("Configuration file found with XDG: {}", xdg_config.display()); config.merge(::config::File::from(xdg_config).required(false)) .context("Failed to load config.toml from XDG configuration directory")?; } let config = config.try_into::()?; let server_addr = matches.value_of("server_addr").unwrap_or_else(|| config.server_addr()); let server_port = matches.value_of("server_port").map(usize::from_str).unwrap_or_else(|| Ok(*config.server_port()))?; let addr = format!("{}:{}", server_addr, server_port).parse()?; let server = { let index_path = tantivy::directory::MmapDirectory::open(config.database_path())?; let schema = crate::schema::schema(); let index = tantivy::Index::open_or_create(index_path, schema.clone())?; let field_path = schema.get_field("path").ok_or_else(|| anyhow!("BUG"))?; let field_ft = schema.get_field("ft").ok_or_else(|| anyhow!("BUG"))?; let field_size = schema.get_field("size").ok_or_else(|| anyhow!("BUG"))?; let field_created = schema.get_field("created").ok_or_else(|| anyhow!("BUG"))?; let field_modified = schema.get_field("modified").ok_or_else(|| anyhow!("BUG"))?; let field_indexed = schema.get_field("indexed").ok_or_else(|| anyhow!("BUG"))?; let field_body = schema.get_field("body").ok_or_else(|| anyhow!("BUG"))?; Server { index, field_path, field_ft, field_size, field_created, field_modified, field_indexed, field_body, } }; tonic::transport::Server::builder() .add_service(fss::fss_server::FssServer::new(server)) .serve(addr) .await .map_err(Error::from) .map(|_| ()) }