diff options
Diffstat (limited to 'src/utils')
| -rw-r--r-- | src/utils/path.rs | 2 | ||||
| -rw-r--r-- | src/utils/request.rs | 108 |
2 files changed, 72 insertions, 38 deletions
diff --git a/src/utils/path.rs b/src/utils/path.rs index 4e66337..83e51c7 100644 --- a/src/utils/path.rs +++ b/src/utils/path.rs @@ -42,7 +42,7 @@ pub async fn expand_glob_paths<T: AsRef<str>>(paths: &[T]) -> Result<Vec<String> Ok(new_paths) } -pub fn path_extension(path: &str) -> Option<String> { +pub fn get_patch_extension(path: &str) -> Option<String> { Path::new(&path) .extension() .map(|v| v.to_string_lossy().to_lowercase()) diff --git a/src/utils/request.rs b/src/utils/request.rs index 8733b75..039bf75 100644 --- a/src/utils/request.rs +++ b/src/utils/request.rs @@ -8,6 +8,7 @@ use tokio::io::AsyncWriteExt; pub const URL_LOADER: &str = "url"; pub const RECURSIVE_URL_LOADER: &str = "recursive_url"; +pub const MEDIA_URL_EXTENSION: &str = "media_url"; pub const DEFAULT_EXTENSION: &str = "txt"; lazy_static! { @@ -19,7 +20,11 @@ lazy_static! { }; } -pub async fn fetch(loaders: &HashMap<String, String>, path: &str) -> Result<(String, String)> { +pub async fn fetch( + loaders: &HashMap<String, String>, + path: &str, + allow_media: bool, +) -> Result<(String, String)> { if let Some(loader_command) = loaders.get(URL_LOADER) { let contents = run_loader_command(path, URL_LOADER, loader_command)?; return Ok((contents, DEFAULT_EXTENSION.into())); @@ -29,6 +34,9 @@ pub async fn fetch(loaders: &HashMap<String, String>, path: &str) -> Result<(Str Err(ref err) => bail!("{err}"), }; let mut res = client.get(path).send().await?; + if !res.status().is_success() { + bail!("Invalid status: {}", res.status()); + } let content_type = res .headers() .get(CONTENT_TYPE) @@ -37,45 +45,71 @@ pub async fn fetch(loaders: &HashMap<String, String>, path: &str) -> Result<(Str Some((mime, _)) => mime.trim(), None => v, }) - .unwrap_or_default(); - let extension = match content_type { - "application/pdf" => "pdf", - "application/vnd.openxmlformats-officedocument.wordprocessingml.document" => "docx", - "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet" => "xlsx", - "application/vnd.openxmlformats-officedocument.presentationml.presentation" => "pptx", - "application/vnd.oasis.opendocument.text" => "odt", - "application/vnd.oasis.opendocument.spreadsheet" => "ods", - "application/vnd.oasis.opendocument.presentation" => "odp", - "application/rtf" => "rtf", - "text/html" => "html", - _ => path + .map(|v| v.to_string()) + .unwrap_or_else(|| { + format!( + "_/{}", + get_patch_extension(path).unwrap_or_else(|| DEFAULT_EXTENSION.into()) + ) + }); + let mut is_media = false; + let extension = match content_type.as_str() { + "application/pdf" => "pdf".into(), + "application/vnd.openxmlformats-officedocument.wordprocessingml.document" => "docx".into(), + "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet" => "xlsx".into(), + "application/vnd.openxmlformats-officedocument.presentationml.presentation" => { + "pptx".into() + } + "application/vnd.oasis.opendocument.text" => "odt".into(), + "application/vnd.oasis.opendocument.spreadsheet" => "ods".into(), + "application/vnd.oasis.opendocument.presentation" => "odp".into(), + "application/rtf" => "rtf".into(), + "text/javascript" => "js".into(), + "text/html" => "html".into(), + _ => content_type .rsplit_once('/') - .and_then(|(_, pair)| pair.rsplit_once('.').map(|(_, ext)| ext)) - .unwrap_or(DEFAULT_EXTENSION), + .map(|(first, last)| { + if ["image", "video", "audio"].contains(&first) { + is_media = true; + MEDIA_URL_EXTENSION.into() + } else { + last.to_lowercase() + } + }) + .unwrap_or_else(|| DEFAULT_EXTENSION.into()), }; - let extension = extension.to_lowercase(); - let result = match loaders.get(&extension) { - Some(loader_command) => { - let save_path = temp_file("-download-", &format!(".{extension}")) - .display() - .to_string(); - let mut save_file = tokio::fs::File::create(&save_path).await?; - let mut size = 0; - while let Some(chunk) = res.chunk().await? { - size += chunk.len(); - save_file.write_all(&chunk).await?; - } - let contents = if size == 0 { - println!("{}", warning_text(&format!("No content at '{path}'"))); - String::new() - } else { - run_loader_command(&save_path, &extension, loader_command)? - }; - (contents, DEFAULT_EXTENSION.into()) + let result = if is_media { + if !allow_media { + bail!("Unexpected media type") } - None => { - let contents = res.text().await?; - (contents, extension) + let image_bytes = res.bytes().await?; + let image_base64 = base64_encode(&image_bytes); + let contents = format!("data:{};base64,{}", content_type, image_base64); + (contents, extension) + } else { + match loaders.get(&extension) { + Some(loader_command) => { + let save_path = temp_file("-download-", &format!(".{extension}")) + .display() + .to_string(); + let mut save_file = tokio::fs::File::create(&save_path).await?; + let mut size = 0; + while let Some(chunk) = res.chunk().await? { + size += chunk.len(); + save_file.write_all(&chunk).await?; + } + let contents = if size == 0 { + println!("{}", warning_text(&format!("No content at '{path}'"))); + String::new() + } else { + run_loader_command(&save_path, &extension, loader_command)? + }; + (contents, DEFAULT_EXTENSION.into()) + } + None => { + let contents = res.text().await?; + (contents, extension) + } } }; Ok(result) |