mirror of
https://codeberg.org/ashley/poke.git
synced 2024-11-26 22:58:16 +01:00
123 lines
3.5 KiB
Rust
123 lines
3.5 KiB
Rust
use std::time::Duration;
|
|
|
|
use actix_web::web::Bytes;
|
|
use encoding_rs::{Encoding, UTF_8_INIT};
|
|
use mime::Mime;
|
|
use reqwest::{
|
|
header::{self, CONTENT_TYPE},
|
|
Client, Response,
|
|
};
|
|
use scraper::Html;
|
|
use std::io::Write;
|
|
use tempfile::NamedTempFile;
|
|
|
|
use super::{result::Error, variables::MAX_BYTES};
|
|
|
|
lazy_static! {
|
|
static ref CLIENT: Client = reqwest::Client::builder()
|
|
.user_agent("Mozilla/5.0 (compatible; January/1.0; +https://github.com/revoltchat/january)")
|
|
.timeout(Duration::from_secs(15))
|
|
.connect_timeout(Duration::from_secs(5))
|
|
.build()
|
|
.expect("reqwest Client");
|
|
}
|
|
|
|
pub async fn fetch(url: &str) -> Result<(Response, Mime), Error> {
|
|
let resp = CLIENT
|
|
.get(url)
|
|
.send()
|
|
.await
|
|
.map_err(|_| Error::ReqwestFailed)?;
|
|
|
|
if !resp.status().is_success() {
|
|
return Err(Error::RequestFailed);
|
|
}
|
|
|
|
let content_type = resp
|
|
.headers()
|
|
.get(CONTENT_TYPE)
|
|
.ok_or(Error::MissingContentType)?
|
|
.to_str()
|
|
.map_err(|_| Error::ConversionFailed)?;
|
|
|
|
let mime: mime::Mime = content_type
|
|
.parse()
|
|
.map_err(|_| Error::FailedToParseContentType)?;
|
|
|
|
Ok((resp, mime))
|
|
}
|
|
|
|
pub async fn get_bytes(resp: &mut Response) -> Result<Bytes, Error> {
|
|
let content_length = resp.content_length().unwrap_or(0) as usize;
|
|
if content_length > *MAX_BYTES {
|
|
return Err(Error::ExceedsMaxBytes);
|
|
}
|
|
let mut bytes = Vec::with_capacity(content_length);
|
|
while let Some(chunk) = resp
|
|
.chunk()
|
|
.await
|
|
.map_err(|_| Error::FailedToConsumeBytes)?
|
|
{
|
|
if bytes.len() + chunk.len() > *MAX_BYTES {
|
|
return Err(Error::ExceedsMaxBytes);
|
|
}
|
|
bytes.extend(chunk)
|
|
}
|
|
Ok(Bytes::from(bytes))
|
|
}
|
|
|
|
pub async fn consume_fragment(mut resp: Response) -> Result<Html, Error> {
|
|
let bytes = get_bytes(&mut resp).await?;
|
|
|
|
let content_type = resp
|
|
.headers()
|
|
.get(header::CONTENT_TYPE)
|
|
.and_then(|value| value.to_str().ok())
|
|
.and_then(|value| value.parse::<Mime>().ok());
|
|
let encoding_name = content_type
|
|
.as_ref()
|
|
.and_then(|mime| mime.get_param("charset").map(|charset| charset.as_str()))
|
|
.unwrap_or("utf-8");
|
|
let encoding = Encoding::for_label(encoding_name.as_bytes()).unwrap_or(&UTF_8_INIT);
|
|
|
|
let (text, _, _) = encoding.decode(&bytes);
|
|
Ok(Html::parse_document(&text))
|
|
}
|
|
|
|
pub fn determine_video_size(path: &std::path::Path) -> Result<(isize, isize), Error> {
|
|
let data = ffprobe::ffprobe(path).map_err(|_| Error::ProbeError)?;
|
|
|
|
// Take the first valid stream.
|
|
for stream in data.streams {
|
|
if let (Some(w), Some(h)) = (stream.width, stream.height) {
|
|
if let (Ok(w), Ok(h)) = (w.try_into(), h.try_into()) {
|
|
return Ok((w, h));
|
|
}
|
|
}
|
|
}
|
|
|
|
Err(Error::ProbeError)
|
|
}
|
|
|
|
pub async fn consume_size(mut resp: Response, mime: Mime) -> Result<(isize, isize), Error> {
|
|
let bytes = get_bytes(&mut resp).await?;
|
|
|
|
match mime.type_() {
|
|
mime::IMAGE => {
|
|
if let Ok(size) = imagesize::blob_size(&bytes) {
|
|
Ok((size.width as isize, size.height as isize))
|
|
} else {
|
|
Err(Error::CouldNotDetermineImageSize)
|
|
}
|
|
}
|
|
mime::VIDEO => {
|
|
let mut tmp = NamedTempFile::new().map_err(|_| Error::CouldNotDetermineVideoSize)?;
|
|
|
|
tmp.write_all(&bytes)
|
|
.map_err(|_| Error::CouldNotDetermineVideoSize)?;
|
|
|
|
determine_video_size(tmp.path())
|
|
}
|
|
_ => unreachable!(),
|
|
}
|
|
}
|