Nekomata/site_icons
Nekomata/site_icons
Archived
2
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Use default rust formatting

Browse source
This commit is contained in:
Markus Kohlhase 2023-09-15 15:25:39 +02:00
parent ed163b9a42
commit a1c38dbc9a
20 changed files with 1104 additions and 1102 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
.rustfmt.toml
View file

@ -1,3 +0,0 @@
indent_style = "Block"
reorder_imports = true
tab_spaces = 2

50
src/bin/site-icons.rs
View file

@ -6,37 +6,37 @@ use std::error::Error;
#[derive(Parser)] #[derive(Parser)]
struct Opts { struct Opts {
url: String, url: String,
#[clap(long)] #[clap(long)]
fast: bool, fast: bool,
#[clap(long)] #[clap(long)]
json: bool, json: bool,
#[clap(long)] #[clap(long)]
/// Print out errors that occurred for skipped items /// Print out errors that occurred for skipped items
debug: bool, debug: bool,
} }
#[tokio::main] #[tokio::main]
async fn main() -> Result<(), Box<dyn Error>> { async fn main() -> Result<(), Box<dyn Error>> {
let mut icons = SiteIcons::new(); let mut icons = SiteIcons::new();
let opts: Opts = Opts::parse(); let opts: Opts = Opts::parse();
if opts.debug { if opts.debug {
let mut builder = Builder::new(); let mut builder = Builder::new();
builder.filter_level(LevelFilter::Info); builder.filter_level(LevelFilter::Info);
builder.init(); builder.init();
}
let entries = icons.load_website(opts.url, opts.fast).await?;
if opts.json {
println!("{}", serde_json::to_string_pretty(&entries)?)
} else {
for icon in entries {
println!("{} {} {}", icon.url, icon.kind, icon.info);
} }
}
Ok(()) let entries = icons.load_website(opts.url, opts.fast).await?;
if opts.json {
println!("{}", serde_json::to_string_pretty(&entries)?)
} else {
for icon in entries {
println!("{} {} {}", icon.url, icon.kind, icon.info);
}
}
Ok(())
} }

184
src/html_parser/head.rs
View file

@ -8,9 +8,9 @@ use futures::Stream;
use futures::StreamExt; use futures::StreamExt;
use lol_html::{element, errors::RewritingError, HtmlRewriter, Settings}; use lol_html::{element, errors::RewritingError, HtmlRewriter, Settings};
use std::{ use std::{
cell::RefCell, cell::RefCell,
error::Error, error::Error,
fmt::{self, Display}, fmt::{self, Display},
}; };
use url::Url; use url::Url;
@ -18,104 +18,106 @@ use url::Url;
struct EndOfHead {} struct EndOfHead {}
impl Display for EndOfHead { impl Display for EndOfHead {
fn fmt(&self, _: &mut fmt::Formatter<'_>) -> fmt::Result { fn fmt(&self, _: &mut fmt::Formatter<'_>) -> fmt::Result {
Ok(()) Ok(())
} }
} }
impl Error for EndOfHead {} impl Error for EndOfHead {}
pub async fn parse_head( pub async fn parse_head(
url: &Url, url: &Url,
mut body: impl Stream<Item = Result<Vec<u8>, String>> + Unpin, mut body: impl Stream<Item = Result<Vec<u8>, String>> + Unpin,
) -> Result<Vec<Icon>, Box<dyn Error>> { ) -> Result<Vec<Icon>, Box<dyn Error>> {
let mut icons = Vec::new(); let mut icons = Vec::new();
let new_icons = RefCell::new(Vec::new()); let new_icons = RefCell::new(Vec::new());
{ {
let mut rewriter = HtmlRewriter::new( let mut rewriter = HtmlRewriter::new(
Settings { Settings {
element_content_handlers: vec![ element_content_handlers: vec![
element!("head", |head| { element!("head", |head| {
head.on_end_tag(|_| Err(Box::new(EndOfHead {})))?; head.on_end_tag(|_| Err(Box::new(EndOfHead {})))?;
Ok(()) Ok(())
}), }),
element!("link[rel~='manifest']", |manifest| { element!("link[rel~='manifest']", |manifest| {
if let Some(href) = manifest if let Some(href) = manifest
.get_attribute("href") .get_attribute("href")
.and_then(|href| url.join(&href).ok()) .and_then(|href| url.join(&href).ok())
{ {
new_icons.borrow_mut().push( new_icons.borrow_mut().push(
async { SiteIcons::load_manifest(href).await.unwrap_or(Vec::new()) } async {
.boxed_local() SiteIcons::load_manifest(href).await.unwrap_or(Vec::new())
.shared(), }
) .boxed_local()
.shared(),
)
}
Ok(())
}),
element!(
join_with!(
",",
"link[rel~='icon']",
"link[rel~='apple-touch-icon']",
"link[rel~='apple-touch-icon-precomposed']"
),
|link| {
let rel = link.get_attribute("rel").unwrap();
if let Some(href) = link
.get_attribute("href")
.and_then(|href| url.join(&href).ok())
{
let kind = if rel.contains("apple-touch-icon") {
IconKind::AppIcon
} else {
IconKind::SiteFavicon
};
let sizes = link.get_attribute("sizes");
new_icons.borrow_mut().push(
async {
Icon::load(href, kind, sizes)
.await
.map(|icon| vec![icon])
.unwrap_or(Vec::new())
}
.boxed_local()
.shared(),
)
};
Ok(())
}
),
],
..Settings::default()
},
|_: &[u8]| {},
);
while let Some(data) = poll_in_background(body.next(), join_all(icons.clone())).await {
let result = rewriter.write(&data?);
icons.extend(new_icons.borrow_mut().drain(..));
match result {
Err(RewritingError::ContentHandlerError(result)) => {
match result.downcast::<EndOfHead>() {
Ok(_) => break,
Err(err) => return Err(err),
};
}
result => result?,
} }
Ok(())
}),
element!(
join_with!(
",",
"link[rel~='icon']",
"link[rel~='apple-touch-icon']",
"link[rel~='apple-touch-icon-precomposed']"
),
|link| {
let rel = link.get_attribute("rel").unwrap();
if let Some(href) = link
.get_attribute("href")
.and_then(|href| url.join(&href).ok())
{
let kind = if rel.contains("apple-touch-icon") {
IconKind::AppIcon
} else {
IconKind::SiteFavicon
};
let sizes = link.get_attribute("sizes");
new_icons.borrow_mut().push(
async {
Icon::load(href, kind, sizes)
.await
.map(|icon| vec![icon])
.unwrap_or(Vec::new())
}
.boxed_local()
.shared(),
)
};
Ok(())
}
),
],
..Settings::default()
},
|_: &[u8]| {},
);
while let Some(data) = poll_in_background(body.next(), join_all(icons.clone())).await {
let result = rewriter.write(&data?);
icons.extend(new_icons.borrow_mut().drain(..));
match result {
Err(RewritingError::ContentHandlerError(result)) => {
match result.downcast::<EndOfHead>() {
Ok(_) => break,
Err(err) => return Err(err),
};
} }
result => result?,
}
} }
}
let icons = join_all(icons).await.into_iter().flatten().collect(); let icons = join_all(icons).await.into_iter().flatten().collect();
Ok(icons) Ok(icons)
} }

271
src/html_parser/site_logo.rs
View file

@ -1,8 +1,8 @@
use crate::{utils::encode_svg, Icon, IconKind}; use crate::{utils::encode_svg, Icon, IconKind};
use futures::{Stream, StreamExt}; use futures::{Stream, StreamExt};
use html5ever::{ use html5ever::{
driver, driver,
tendril::{Tendril, TendrilSink}, tendril::{Tendril, TendrilSink},
}; };
use scraper::{ElementRef, Html}; use scraper::{ElementRef, Html};
use std::error::Error; use std::error::Error;
@ -11,147 +11,146 @@ use tldextract::TldOption;
use url::Url; use url::Url;
pub async fn parse_site_logo( pub async fn parse_site_logo(
url: &Url, url: &Url,
mut body: impl Stream<Item = Result<Vec<u8>, String>> + Unpin, mut body: impl Stream<Item = Result<Vec<u8>, String>> + Unpin,
is_blacklisted: impl Fn(&Url) -> bool, is_blacklisted: impl Fn(&Url) -> bool,
) -> Result<Icon, Box<dyn Error>> { ) -> Result<Icon, Box<dyn Error>> {
let mut parser = driver::parse_document(Html::new_document(), Default::default()); let mut parser = driver::parse_document(Html::new_document(), Default::default());
while let Some(data) = body.next().await { while let Some(data) = body.next().await {
if let Ok(data) = Tendril::try_from_byte_slice(&data?) { if let Ok(data) = Tendril::try_from_byte_slice(&data?) {
parser.process(data) parser.process(data)
}
}
let document = parser.finish();
let mut logos: Vec<_> = document
.select(selector!(
"a[href='/'] img, a[href='/'] svg",
"header img, header svg",
"img[src*=logo]",
"img[alt*=logo], svg[alt*=logo]",
"*[class*=logo] img, *[class*=logo] svg",
"*[id*=logo] img, *[id*=logo] svg",
"img[class*=logo], svg[class*=logo]",
"img[id*=logo], svg[id*=logo]",
))
.enumerate()
.filter_map(|(i, elem_ref)| {
let elem = elem_ref.value();
let ancestors = elem_ref
.ancestors()
.map(ElementRef::wrap)
.flatten()
.map(|elem_ref| elem_ref.value())
.collect::<Vec<_>>();
let skip_classnames = regex!("menu|search");
let should_skip = ancestors.iter().any(|ancestor| {
ancestor
.attr("class")
.map(|attr| skip_classnames.is_match(&attr.to_lowercase()))
.or_else(|| {
ancestor
.attr("id")
.map(|attr| skip_classnames.is_match(&attr.to_lowercase()))
})
.unwrap_or(false)
});
if should_skip {
return None;
}
let mut weight = 0;
// if in the header
if ancestors.iter().any(|element| element.name() == "header") {
weight += 2;
}
if i == 0 {
weight += 1;
}
let mentions = |attr_name, is_match: Box<dyn Fn(&str) -> bool>| {
ancestors.iter().chain(iter::once(&elem)).any(|ancestor| {
ancestor
.attr(attr_name)
.map(|attr| is_match(&attr.to_lowercase()))
.unwrap_or(false)
})
};
if mentions("href", Box::new(|attr| attr == "/")) {
weight += 5;
};
let mentions_logo = |attr_name| {
mentions(
attr_name,
Box::new(|attr| regex!("logo([^s]|$)").is_match(attr)),
)
};
if mentions_logo("class") || mentions_logo("id") {
weight += 3;
}
if mentions_logo("alt") {
weight += 2;
}
if mentions_logo("src") {
weight += 1;
}
if let Some(site_name) = url
.domain()
.and_then(|domain| TldOption::default().build().extract(domain).unwrap().domain)
{
// if the alt contains the site_name then highest priority
if site_name
.to_lowercase()
.split('-')
.any(|segment| mentions("alt", Box::new(move |attr| attr.contains(segment))))
{
weight += 10;
} }
} }
let href = if elem.name() == "svg" { let document = parser.finish();
Some(Url::parse(&encode_svg(&elem_ref.html())).unwrap())
} else {
elem.attr("src").and_then(|href| url.join(&href).ok())
};
if let Some(href) = &href { let mut logos: Vec<_> =
if is_blacklisted(href) { document
return None; .select(selector!(
"a[href='/'] img, a[href='/'] svg",
"header img, header svg",
"img[src*=logo]",
"img[alt*=logo], svg[alt*=logo]",
"*[class*=logo] img, *[class*=logo] svg",
"*[id*=logo] img, *[id*=logo] svg",
"img[class*=logo], svg[class*=logo]",
"img[id*=logo], svg[id*=logo]",
))
.enumerate()
.filter_map(|(i, elem_ref)| {
let elem = elem_ref.value();
let ancestors = elem_ref
.ancestors()
.map(ElementRef::wrap)
.flatten()
.map(|elem_ref| elem_ref.value())
.collect::<Vec<_>>();
let skip_classnames = regex!("menu|search");
let should_skip = ancestors.iter().any(|ancestor| {
ancestor
.attr("class")
.map(|attr| skip_classnames.is_match(&attr.to_lowercase()))
.or_else(|| {
ancestor
.attr("id")
.map(|attr| skip_classnames.is_match(&attr.to_lowercase()))
})
.unwrap_or(false)
});
if should_skip {
return None;
}
let mut weight = 0;
// if in the header
if ancestors.iter().any(|element| element.name() == "header") {
weight += 2;
}
if i == 0 {
weight += 1;
}
let mentions = |attr_name, is_match: Box<dyn Fn(&str) -> bool>| {
ancestors.iter().chain(iter::once(&elem)).any(|ancestor| {
ancestor
.attr(attr_name)
.map(|attr| is_match(&attr.to_lowercase()))
.unwrap_or(false)
})
};
if mentions("href", Box::new(|attr| attr == "/")) {
weight += 5;
};
let mentions_logo = |attr_name| {
mentions(
attr_name,
Box::new(|attr| regex!("logo([^s]|$)").is_match(attr)),
)
};
if mentions_logo("class") || mentions_logo("id") {
weight += 3;
}
if mentions_logo("alt") {
weight += 2;
}
if mentions_logo("src") {
weight += 1;
}
if let Some(site_name) = url
.domain()
.and_then(|domain| TldOption::default().build().extract(domain).unwrap().domain)
{
// if the alt contains the site_name then highest priority
if site_name.to_lowercase().split('-').any(|segment| {
mentions("alt", Box::new(move |attr| attr.contains(segment)))
}) {
weight += 10;
}
}
let href = if elem.name() == "svg" {
Some(Url::parse(&encode_svg(&elem_ref.html())).unwrap())
} else {
elem.attr("src").and_then(|href| url.join(&href).ok())
};
if let Some(href) = &href {
if is_blacklisted(href) {
return None;
}
}
href.map(|href| (href, elem_ref, weight))
})
.collect();
logos.sort_by(|(_, _, a_weight), (_, _, b_weight)| b_weight.cmp(a_weight));
// prefer <img> over svg
let mut prev_weight = None;
for (href, elem_ref, weight) in &logos {
if let Some(prev_weight) = prev_weight {
if weight != prev_weight {
break;
}
} }
} prev_weight = Some(weight);
href.map(|href| (href, elem_ref, weight)) if elem_ref.value().name() == "img" {
}) return Icon::load(href.clone(), IconKind::SiteLogo, None).await;
.collect(); }
logos.sort_by(|(_, _, a_weight), (_, _, b_weight)| b_weight.cmp(a_weight));
// prefer <img> over svg
let mut prev_weight = None;
for (href, elem_ref, weight) in &logos {
if let Some(prev_weight) = prev_weight {
if weight != prev_weight {
break;
}
} }
prev_weight = Some(weight);
if elem_ref.value().name() == "img" { match logos.into_iter().next() {
return Icon::load(href.clone(), IconKind::SiteLogo, None).await; Some((href, _, _)) => Icon::load(href.clone(), IconKind::SiteLogo, None).await,
None => Err("No site logo found".into()),
} }
}
match logos.into_iter().next() {
Some((href, _, _)) => Icon::load(href.clone(), IconKind::SiteLogo, None).await,
None => Err("No site logo found".into()),
}
} }

425
src/icon/icon_info.rs
View file

@ -6,269 +6,272 @@ use mime::MediaType;
use reqwest::{header::*, Url}; use reqwest::{header::*, Url};
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Serialize};
use std::{ use std::{
cmp::Ordering, cmp::Ordering,
convert::TryFrom, convert::TryFrom,
error::Error, error::Error,
fmt::{self, Display}, fmt::{self, Display},
io, io,
}; };
enum IconKind { enum IconKind {
SVG, SVG,
PNG, PNG,
JPEG, JPEG,
ICO, ICO,
GIF, GIF,
} }
#[derive(Clone, Debug, Serialize, Deserialize, Eq, PartialEq)] #[derive(Clone, Debug, Serialize, Deserialize, Eq, PartialEq)]
#[serde(tag = "type")] #[serde(tag = "type")]
#[serde(rename_all = "lowercase")] #[serde(rename_all = "lowercase")]
pub enum IconInfo { pub enum IconInfo {
PNG { size: IconSize }, PNG { size: IconSize },
JPEG { size: IconSize }, JPEG { size: IconSize },
ICO { sizes: IconSizes }, ICO { sizes: IconSizes },
GIF { size: IconSize }, GIF { size: IconSize },
SVG { size: Option<IconSize> }, SVG { size: Option<IconSize> },
} }
impl IconInfo { impl IconInfo {
async fn decode<R: AsyncRead + Unpin>( async fn decode<R: AsyncRead + Unpin>(
reader: &mut R, reader: &mut R,
kind: Option<IconKind>, kind: Option<IconKind>,
) -> Result<IconInfo, Box<dyn Error>> { ) -> Result<IconInfo, Box<dyn Error>> {
let mut header = [0; 2]; let mut header = [0; 2];
reader.read_exact(&mut header).await?; reader.read_exact(&mut header).await?;
match (kind, &header) { match (kind, &header) {
(Some(IconKind::SVG), bytes) => { (Some(IconKind::SVG), bytes) => {
let size = get_svg_size(bytes, reader).await?; let size = get_svg_size(bytes, reader).await?;
Ok(IconInfo::SVG { size }) Ok(IconInfo::SVG { size })
} }
(_, &[0x60, byte_two]) => { (_, &[0x60, byte_two]) => {
let size = get_svg_size(&[0x60, byte_two], reader).await?; let size = get_svg_size(&[0x60, byte_two], reader).await?;
Ok(IconInfo::SVG { size }) Ok(IconInfo::SVG { size })
} }
(Some(IconKind::PNG), _) | (_, b"\x89P") => { (Some(IconKind::PNG), _) | (_, b"\x89P") => {
let size = get_png_size(reader).await?; let size = get_png_size(reader).await?;
Ok(IconInfo::PNG { size }) Ok(IconInfo::PNG { size })
} }
(Some(IconKind::ICO), _) | (_, &[0x00, 0x00]) => { (Some(IconKind::ICO), _) | (_, &[0x00, 0x00]) => {
let sizes = get_ico_sizes(reader).await?; let sizes = get_ico_sizes(reader).await?;
Ok(IconInfo::ICO { sizes }) Ok(IconInfo::ICO { sizes })
} }
(Some(IconKind::JPEG), _) | (_, &[0xFF, 0xD8]) => { (Some(IconKind::JPEG), _) | (_, &[0xFF, 0xD8]) => {
let size = get_jpeg_size(reader).await?; let size = get_jpeg_size(reader).await?;
Ok(IconInfo::JPEG { size }) Ok(IconInfo::JPEG { size })
} }
(Some(IconKind::GIF), _) | (_, b"GI") => { (Some(IconKind::GIF), _) | (_, b"GI") => {
let size = get_gif_size(reader).await?; let size = get_gif_size(reader).await?;
Ok(IconInfo::GIF { size }) Ok(IconInfo::GIF { size })
} }
_ => Err(format!("unknown icon type ({:?})", header).into()), _ => Err(format!("unknown icon type ({:?})", header).into()),
}
} }
}
pub async fn load( pub async fn load(
url: Url, url: Url,
headers: HeaderMap, headers: HeaderMap,
sizes: Option<String>, sizes: Option<String>,
) -> Result<IconInfo, Box<dyn Error>> { ) -> Result<IconInfo, Box<dyn Error>> {
let sizes = sizes.as_ref().and_then(|s| IconSizes::try_from(s).ok()); let sizes = sizes.as_ref().and_then(|s| IconSizes::try_from(s).ok());
let (mime, mut body): (_, Box<dyn AsyncRead + Unpin>) = match url.scheme() { let (mime, mut body): (_, Box<dyn AsyncRead + Unpin>) = match url.scheme() {
"data" => { "data" => {
let url = url.to_string(); let url = url.to_string();
let url = DataUrl::process(&url).map_err(|_| "failed to parse data uri")?; let url = DataUrl::process(&url).map_err(|_| "failed to parse data uri")?;
let mime = url.mime_type().to_string().parse::<MediaType>()?; let mime = url.mime_type().to_string().parse::<MediaType>()?;
let body = Cursor::new( let body = Cursor::new(
url url.decode_to_vec()
.decode_to_vec() .map_err(|_| "failed to decode data uri body")?
.map_err(|_| "failed to decode data uri body")? .0,
.0, );
);
(mime, Box::new(body)) (mime, Box::new(body))
} }
_ => { _ => {
let res = CLIENT let res = CLIENT
.get(url) .get(url)
.headers(headers) .headers(headers)
.send() .send()
.await? .await?
.error_for_status()?; .error_for_status()?;
if !res.status().is_success() { if !res.status().is_success() {
return Err("failed to fetch".into()); return Err("failed to fetch".into());
};
let mime = res
.headers()
.get(CONTENT_TYPE)
.ok_or("no content type")?
.to_str()?
.parse::<MediaType>()?;
let body = res
.bytes_stream()
.map(|result| {
result.map_err(|error| {
io::Error::new(io::ErrorKind::Other, error.to_string())
})
})
.into_async_read();
(mime, Box::new(body))
}
}; };
let mime = res let kind = match (mime.type_(), mime.subtype()) {
.headers() (mime::IMAGE, mime::PNG) => {
.get(CONTENT_TYPE) if let Some(sizes) = sizes {
.ok_or("no content type")? return Ok(IconInfo::PNG {
.to_str()? size: *sizes.largest(),
.parse::<MediaType>()?; });
}
Some(IconKind::PNG)
}
let body = res (mime::IMAGE, mime::JPEG) => {
.bytes_stream() if let Some(sizes) = sizes {
.map(|result| { return Ok(IconInfo::JPEG {
result.map_err(|error| io::Error::new(io::ErrorKind::Other, error.to_string())) size: *sizes.largest(),
}) });
.into_async_read(); }
Some(IconKind::JPEG)
}
(mime, Box::new(body)) (mime::IMAGE, "x-icon") | (mime::IMAGE, "vnd.microsoft.icon") => {
} if let Some(sizes) = sizes {
}; return Ok(IconInfo::ICO { sizes });
}
let kind = match (mime.type_(), mime.subtype()) { Some(IconKind::ICO)
(mime::IMAGE, mime::PNG) => { }
if let Some(sizes) = sizes {
return Ok(IconInfo::PNG {
size: *sizes.largest(),
});
}
Some(IconKind::PNG)
}
(mime::IMAGE, mime::JPEG) => { (mime::IMAGE, mime::GIF) => {
if let Some(sizes) = sizes { if let Some(sizes) = sizes {
return Ok(IconInfo::JPEG { return Ok(IconInfo::GIF {
size: *sizes.largest(), size: *sizes.largest(),
}); });
} }
Some(IconKind::JPEG)
}
(mime::IMAGE, "x-icon") | (mime::IMAGE, "vnd.microsoft.icon") => { Some(IconKind::GIF)
if let Some(sizes) = sizes { }
return Ok(IconInfo::ICO { sizes });
}
Some(IconKind::ICO) (mime::IMAGE, mime::SVG) | (mime::TEXT, mime::PLAIN) => {
} if let Some(sizes) = sizes {
return Ok(IconInfo::SVG {
size: Some(*sizes.largest()),
});
}
(mime::IMAGE, mime::GIF) => { Some(IconKind::SVG)
if let Some(sizes) = sizes { }
return Ok(IconInfo::GIF {
size: *sizes.largest(),
});
}
Some(IconKind::GIF) _ => None,
} };
(mime::IMAGE, mime::SVG) | (mime::TEXT, mime::PLAIN) => { IconInfo::decode(&mut body, kind).await
if let Some(sizes) = sizes {
return Ok(IconInfo::SVG {
size: Some(*sizes.largest()),
});
}
Some(IconKind::SVG)
}
_ => None,
};
IconInfo::decode(&mut body, kind).await
}
pub fn size(&self) -> Option<&IconSize> {
match self {
IconInfo::ICO { sizes } => Some(sizes.largest()),
IconInfo::PNG { size } | IconInfo::JPEG { size } | IconInfo::GIF { size } => Some(size),
IconInfo::SVG { size } => size.as_ref(),
} }
}
pub fn sizes(&self) -> Option<IconSizes> { pub fn size(&self) -> Option<&IconSize> {
match self { match self {
IconInfo::ICO { sizes } => Some((*sizes).clone()), IconInfo::ICO { sizes } => Some(sizes.largest()),
IconInfo::PNG { size } | IconInfo::JPEG { size } | IconInfo::GIF { size } => { IconInfo::PNG { size } | IconInfo::JPEG { size } | IconInfo::GIF { size } => Some(size),
Some((*size).into()) IconInfo::SVG { size } => size.as_ref(),
} }
IconInfo::SVG { size } => size.map(|size| size.into()),
} }
}
pub fn mime_type(&self) -> &'static str { pub fn sizes(&self) -> Option<IconSizes> {
match self { match self {
IconInfo::PNG { .. } => "image/png", IconInfo::ICO { sizes } => Some((*sizes).clone()),
IconInfo::JPEG { .. } => "image/jpeg", IconInfo::PNG { size } | IconInfo::JPEG { size } | IconInfo::GIF { size } => {
IconInfo::ICO { .. } => "image/x-icon", Some((*size).into())
IconInfo::GIF { .. } => "image/gif", }
IconInfo::SVG { .. } => "image/svg+xml", IconInfo::SVG { size } => size.map(|size| size.into()),
}
}
pub fn mime_type(&self) -> &'static str {
match self {
IconInfo::PNG { .. } => "image/png",
IconInfo::JPEG { .. } => "image/jpeg",
IconInfo::ICO { .. } => "image/x-icon",
IconInfo::GIF { .. } => "image/gif",
IconInfo::SVG { .. } => "image/svg+xml",
}
} }
}
} }
impl Display for IconInfo { impl Display for IconInfo {
fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> { fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
match self { match self {
IconInfo::PNG { size } => write!(f, "png {}", size), IconInfo::PNG { size } => write!(f, "png {}", size),
IconInfo::JPEG { size } => write!(f, "jpeg {}", size), IconInfo::JPEG { size } => write!(f, "jpeg {}", size),
IconInfo::GIF { size } => write!(f, "gif {}", size), IconInfo::GIF { size } => write!(f, "gif {}", size),
IconInfo::ICO { sizes } => write!(f, "ico {}", sizes), IconInfo::ICO { sizes } => write!(f, "ico {}", sizes),
IconInfo::SVG { size } => { IconInfo::SVG { size } => {
write!( write!(
f, f,
"svg{}", "svg{}",
if let Some(size) = size { if let Some(size) = size {
format!(" {}", size) format!(" {}", size)
} else { } else {
"".to_string() "".to_string()
} }
) )
} }
}
} }
}
} }
impl Ord for IconInfo { impl Ord for IconInfo {
fn cmp(&self, other: &Self) -> Ordering { fn cmp(&self, other: &Self) -> Ordering {
match (self, other) { match (self, other) {
(IconInfo::SVG { size }, IconInfo::SVG { size: other_size }) => match (size, other_size) { (IconInfo::SVG { size }, IconInfo::SVG { size: other_size }) => {
(Some(_), None) => Ordering::Less, match (size, other_size) {
(None, Some(_)) => Ordering::Greater, (Some(_), None) => Ordering::Less,
(Some(size), Some(other_size)) => size.cmp(other_size), (None, Some(_)) => Ordering::Greater,
(None, None) => Ordering::Equal, (Some(size), Some(other_size)) => size.cmp(other_size),
}, (None, None) => Ordering::Equal,
(IconInfo::SVG { .. }, _) => Ordering::Less, }
(_, IconInfo::SVG { .. }) => Ordering::Greater, }
(IconInfo::SVG { .. }, _) => Ordering::Less,
(_, IconInfo::SVG { .. }) => Ordering::Greater,
_ => { _ => {
let size = self.size().unwrap(); let size = self.size().unwrap();
let other_size = other.size().unwrap(); let other_size = other.size().unwrap();
size.cmp(other_size).then_with(|| match (self, other) { size.cmp(other_size).then_with(|| match (self, other) {
(IconInfo::PNG { .. }, IconInfo::PNG { .. }) => Ordering::Equal, (IconInfo::PNG { .. }, IconInfo::PNG { .. }) => Ordering::Equal,
(IconInfo::PNG { .. }, _) => Ordering::Less, (IconInfo::PNG { .. }, _) => Ordering::Less,
(_, IconInfo::PNG { .. }) => Ordering::Greater, (_, IconInfo::PNG { .. }) => Ordering::Greater,
(IconInfo::GIF { .. }, IconInfo::GIF { .. }) => Ordering::Equal, (IconInfo::GIF { .. }, IconInfo::GIF { .. }) => Ordering::Equal,
(IconInfo::GIF { .. }, _) => Ordering::Less, (IconInfo::GIF { .. }, _) => Ordering::Less,
(_, IconInfo::GIF { .. }) => Ordering::Greater, (_, IconInfo::GIF { .. }) => Ordering::Greater,
(IconInfo::JPEG { .. }, IconInfo::JPEG { .. }) => Ordering::Equal, (IconInfo::JPEG { .. }, IconInfo::JPEG { .. }) => Ordering::Equal,
(IconInfo::JPEG { .. }, _) => Ordering::Less, (IconInfo::JPEG { .. }, _) => Ordering::Less,
(_, IconInfo::JPEG { .. }) => Ordering::Greater, (_, IconInfo::JPEG { .. }) => Ordering::Greater,
(IconInfo::ICO { .. }, IconInfo::ICO { .. }) => Ordering::Equal, (IconInfo::ICO { .. }, IconInfo::ICO { .. }) => Ordering::Equal,
(IconInfo::ICO { .. }, _) => Ordering::Less, (IconInfo::ICO { .. }, _) => Ordering::Less,
(_, IconInfo::ICO { .. }) => Ordering::Greater, (_, IconInfo::ICO { .. }) => Ordering::Greater,
_ => unreachable!(), _ => unreachable!(),
}) })
} }
}
} }
}
} }
impl PartialOrd for IconInfo { impl PartialOrd for IconInfo {
fn partial_cmp(&self, other: &Self) -> Option<Ordering> { fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
Some(self.cmp(other)) Some(self.cmp(other))
} }
} }

22
src/icon/icon_size/gif.rs
View file

@ -2,23 +2,23 @@ use super::IconSize;
use byteorder::{LittleEndian, ReadBytesExt}; use byteorder::{LittleEndian, ReadBytesExt};
use futures::prelude::*; use futures::prelude::*;
use std::{ use std::{
error::Error, error::Error,
io::{Cursor, Seek, SeekFrom}, io::{Cursor, Seek, SeekFrom},
}; };
pub async fn get_gif_size<R: AsyncRead + Unpin>( pub async fn get_gif_size<R: AsyncRead + Unpin>(
reader: &mut R, reader: &mut R,
) -> Result<IconSize, Box<dyn Error>> { ) -> Result<IconSize, Box<dyn Error>> {
let mut header = [0; 8]; let mut header = [0; 8];
reader.read_exact(&mut header).await?; reader.read_exact(&mut header).await?;
let header = &mut Cursor::new(header); let header = &mut Cursor::new(header);
assert_slice_eq!(header, 0, b"F8", "bad header"); assert_slice_eq!(header, 0, b"F8", "bad header");
header.seek(SeekFrom::Start(4))?; header.seek(SeekFrom::Start(4))?;
let width = header.read_u16::<LittleEndian>()? as u32; let width = header.read_u16::<LittleEndian>()? as u32;
let height = header.read_u16::<LittleEndian>()? as u32; let height = header.read_u16::<LittleEndian>()? as u32;
Ok(IconSize::new(width, height)) Ok(IconSize::new(width, height))
} }

86
src/icon/icon_size/ico.rs
View file

@ -2,58 +2,58 @@ use super::{png::get_png_size, IconSize, IconSizes};
use byteorder::{LittleEndian, ReadBytesExt as _}; use byteorder::{LittleEndian, ReadBytesExt as _};
use futures::prelude::*; use futures::prelude::*;
use std::{ use std::{
convert::TryInto, convert::TryInto,
error::Error, error::Error,
io::{Cursor, Seek, SeekFrom}, io::{Cursor, Seek, SeekFrom},
}; };
const ICO_TYPE: u16 = 1; const ICO_TYPE: u16 = 1;
const INDEX_SIZE: u16 = 16; const INDEX_SIZE: u16 = 16;
pub async fn get_ico_sizes<R: AsyncRead + Unpin>( pub async fn get_ico_sizes<R: AsyncRead + Unpin>(
reader: &mut R, reader: &mut R,
) -> Result<IconSizes, Box<dyn Error>> { ) -> Result<IconSizes, Box<dyn Error>> {
let mut offset = 4; let mut offset = 4;
let mut header = [0; 4]; let mut header = [0; 4];
reader.read_exact(&mut header).await?; reader.read_exact(&mut header).await?;
let mut header = Cursor::new(header); let mut header = Cursor::new(header);
let icon_type = header.read_u16::<LittleEndian>()?; let icon_type = header.read_u16::<LittleEndian>()?;
if icon_type != ICO_TYPE { if icon_type != ICO_TYPE {
return Err("bad header".into()); return Err("bad header".into());
}
let icon_count = header.read_u16::<LittleEndian>()?;
let mut data = vec![0; (icon_count * INDEX_SIZE) as usize];
reader.read_exact(&mut data).await?;
offset += data.len();
let mut data = Cursor::new(data);
let mut sizes = Vec::new();
for i in 0..icon_count {
data.seek(SeekFrom::Start((INDEX_SIZE * i) as _))?;
let width = data.read_u8()?;
let height = data.read_u8()?;
if width == 0 && height == 0 {
data.seek(SeekFrom::Current(10))?;
let image_offset = data.read_u32::<LittleEndian>()?;
let mut data = vec![0; image_offset as usize - offset];
reader.read_exact(&mut data).await?;
offset += data.len();
let size = get_png_size(reader).await;
if let Ok(size) = size {
sizes.push(size);
}
} else {
sizes.push(IconSize::new(width as _, height as _))
} }
}
Ok(sizes.try_into()?) let icon_count = header.read_u16::<LittleEndian>()?;
let mut data = vec![0; (icon_count * INDEX_SIZE) as usize];
reader.read_exact(&mut data).await?;
offset += data.len();
let mut data = Cursor::new(data);
let mut sizes = Vec::new();
for i in 0..icon_count {
data.seek(SeekFrom::Start((INDEX_SIZE * i) as _))?;
let width = data.read_u8()?;
let height = data.read_u8()?;
if width == 0 && height == 0 {
data.seek(SeekFrom::Current(10))?;
let image_offset = data.read_u32::<LittleEndian>()?;
let mut data = vec![0; image_offset as usize - offset];
reader.read_exact(&mut data).await?;
offset += data.len();
let size = get_png_size(reader).await;
if let Ok(size) = size {
sizes.push(size);
}
} else {
sizes.push(IconSize::new(width as _, height as _))
}
}
Ok(sizes.try_into()?)
} }

120
src/icon/icon_size/icon_sizes.rs
View file

@ -3,11 +3,11 @@ use itertools::Itertools;
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Serialize};
use serde_json::Value; use serde_json::Value;
use std::{ use std::{
cmp::Ordering, cmp::Ordering,
convert::{TryFrom, TryInto}, convert::{TryFrom, TryInto},
error::Error, error::Error,
fmt::{self, Display}, fmt::{self, Display},
ops::Deref, ops::Deref,
}; };
use vec1::{vec1, Vec1}; use vec1::{vec1, Vec1};
@ -16,99 +16,99 @@ use vec1::{vec1, Vec1};
pub struct IconSizes(Vec1<IconSize>); pub struct IconSizes(Vec1<IconSize>);
impl Display for IconSizes { impl Display for IconSizes {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
f.write_str(&self.0.iter().join(" ")) f.write_str(&self.0.iter().join(" "))
} }
} }
impl IconSizes { impl IconSizes {
pub fn add_size(&mut self, size: IconSize) { pub fn add_size(&mut self, size: IconSize) {
match self.0.binary_search(&size) { match self.0.binary_search(&size) {
Ok(_) => {} Ok(_) => {}
Err(pos) => self.0.insert(pos, size), Err(pos) => self.0.insert(pos, size),
}
} }
}
pub fn largest(&self) -> &IconSize { pub fn largest(&self) -> &IconSize {
self.0.first() self.0.first()
} }
} }
impl TryFrom<&str> for IconSizes { impl TryFrom<&str> for IconSizes {
type Error = Box<dyn Error>; type Error = Box<dyn Error>;
fn try_from(sizes_str: &str) -> Result<Self, Self::Error> { fn try_from(sizes_str: &str) -> Result<Self, Self::Error> {
let size_strs = sizes_str.split(" "); let size_strs = sizes_str.split(" ");
let mut sizes = Vec::new(); let mut sizes = Vec::new();
for size in size_strs { for size in size_strs {
if let Ok(size) = serde_json::from_value(Value::String(size.to_string())) { if let Ok(size) = serde_json::from_value(Value::String(size.to_string())) {
sizes.push(size); sizes.push(size);
} }
}
Ok(sizes.try_into()?)
} }
Ok(sizes.try_into()?)
}
} }
impl TryFrom<&String> for IconSizes { impl TryFrom<&String> for IconSizes {
type Error = Box<dyn Error>; type Error = Box<dyn Error>;
fn try_from(sizes_str: &String) -> Result<Self, Self::Error> { fn try_from(sizes_str: &String) -> Result<Self, Self::Error> {
IconSizes::try_from(sizes_str.as_str()) IconSizes::try_from(sizes_str.as_str())
} }
} }
impl TryFrom<String> for IconSizes { impl TryFrom<String> for IconSizes {
type Error = Box<dyn Error>; type Error = Box<dyn Error>;
fn try_from(sizes_str: String) -> Result<Self, Self::Error> { fn try_from(sizes_str: String) -> Result<Self, Self::Error> {
IconSizes::try_from(sizes_str.as_str()) IconSizes::try_from(sizes_str.as_str())
} }
} }
impl Deref for IconSizes { impl Deref for IconSizes {
type Target = Vec1<IconSize>; type Target = Vec1<IconSize>;
fn deref(&self) -> &Vec1<IconSize> { fn deref(&self) -> &Vec1<IconSize> {
&self.0 &self.0
} }
} }
impl IntoIterator for IconSizes { impl IntoIterator for IconSizes {
type Item = IconSize; type Item = IconSize;
type IntoIter = std::vec::IntoIter<Self::Item>; type IntoIter = std::vec::IntoIter<Self::Item>;
fn into_iter(self) -> Self::IntoIter { fn into_iter(self) -> Self::IntoIter {
self.0.into_iter() self.0.into_iter()
} }
} }
impl Ord for IconSizes { impl Ord for IconSizes {
fn cmp(&self, other: &Self) -> Ordering { fn cmp(&self, other: &Self) -> Ordering {
self.largest().cmp(&other.largest()) self.largest().cmp(&other.largest())
} }
} }
impl PartialOrd for IconSizes { impl PartialOrd for IconSizes {
fn partial_cmp(&self, other: &Self) -> Option<Ordering> { fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
Some(self.cmp(other)) Some(self.cmp(other))
} }
} }
impl TryFrom<Vec<IconSize>> for IconSizes { impl TryFrom<Vec<IconSize>> for IconSizes {
type Error = String; type Error = String;
fn try_from(mut vec: Vec<IconSize>) -> Result<Self, Self::Error> { fn try_from(mut vec: Vec<IconSize>) -> Result<Self, Self::Error> {
vec.sort(); vec.sort();
Ok(IconSizes( Ok(IconSizes(
vec.try_into().map_err(|_| "must contain a size")?, vec.try_into().map_err(|_| "must contain a size")?,
)) ))
} }
} }
impl From<IconSize> for IconSizes { impl From<IconSize> for IconSizes {
fn from(size: IconSize) -> Self { fn from(size: IconSize) -> Self {
IconSizes(vec1![size]) IconSizes(vec1![size])
} }
} }

90
src/icon/icon_size/jpeg.rs
View file

@ -5,58 +5,58 @@ use futures::{AsyncRead, AsyncReadExt as _};
use super::IconSize; use super::IconSize;
async fn read_u16_be<R: AsyncRead + Unpin>(reader: &mut R) -> Result<u16, Box<dyn Error>> { async fn read_u16_be<R: AsyncRead + Unpin>(reader: &mut R) -> Result<u16, Box<dyn Error>> {
let mut buf = [0u8; 2]; let mut buf = [0u8; 2];
reader.read_exact(&mut buf).await?; reader.read_exact(&mut buf).await?;
Ok(u16::from_be_bytes(buf)) Ok(u16::from_be_bytes(buf))
} }
pub async fn get_jpeg_size<R: AsyncRead + Unpin>( pub async fn get_jpeg_size<R: AsyncRead + Unpin>(
reader: &mut R, reader: &mut R,
) -> Result<IconSize, Box<dyn Error>> { ) -> Result<IconSize, Box<dyn Error>> {
let mut marker = [0; 2]; let mut marker = [0; 2];
let mut depth = 0i32; let mut depth = 0i32;
loop { loop {
// Read current marker (FF XX) // Read current marker (FF XX)
reader.read_exact(&mut marker).await?; reader.read_exact(&mut marker).await?;
if marker[0] != 0xFF { if marker[0] != 0xFF {
// Did not read a marker. Assume image is corrupt. // Did not read a marker. Assume image is corrupt.
return Err("invalid jpeg".into()); return Err("invalid jpeg".into());
}
let page = marker[1];
// Check for valid SOFn markers. C4, C8, and CC aren't dimension markers.
if (page >= 0xC0 && page <= 0xC3)
|| (page >= 0xC5 && page <= 0xC7)
|| (page >= 0xC9 && page <= 0xCB)
|| (page >= 0xCD && page <= 0xCF)
{
// Only get outside image size
if depth == 0 {
// Correct marker, go forward 3 bytes so we're at height offset
reader.read_exact(&mut [0; 3]).await?;
break;
}
} else if page == 0xD8 {
depth += 1;
} else if page == 0xD9 {
depth -= 1;
if depth < 0 {
return Err("invalid jpeg".into());
}
}
// Read the marker length and skip over it entirely
let page_size = read_u16_be(reader).await? as i64;
reader
.read_exact(&mut vec![0; (page_size - 2) as usize])
.await?;
} }
let page = marker[1]; let height = read_u16_be(reader).await?;
let width = read_u16_be(reader).await?;
// Check for valid SOFn markers. C4, C8, and CC aren't dimension markers. Ok(IconSize::new(width as _, height as _))
if (page >= 0xC0 && page <= 0xC3)
|| (page >= 0xC5 && page <= 0xC7)
|| (page >= 0xC9 && page <= 0xCB)
|| (page >= 0xCD && page <= 0xCF)
{
// Only get outside image size
if depth == 0 {
// Correct marker, go forward 3 bytes so we're at height offset
reader.read_exact(&mut [0; 3]).await?;
break;
}
} else if page == 0xD8 {
depth += 1;
} else if page == 0xD9 {
depth -= 1;
if depth < 0 {
return Err("invalid jpeg".into());
}
}
// Read the marker length and skip over it entirely
let page_size = read_u16_be(reader).await? as i64;
reader
.read_exact(&mut vec![0; (page_size - 2) as usize])
.await?;
}
let height = read_u16_be(reader).await?;
let width = read_u16_be(reader).await?;
Ok(IconSize::new(width as _, height as _))
} }

104
src/icon/icon_size/mod.rs
View file

@ -14,87 +14,87 @@ pub use svg::*;
use serde::{de, Deserialize, Deserializer, Serialize, Serializer}; use serde::{de, Deserialize, Deserializer, Serialize, Serializer};
use std::{ use std::{
cmp::{self, Ordering}, cmp::{self, Ordering},
error::Error, error::Error,
fmt::{self, Display}, fmt::{self, Display},
io::{Read, Seek, SeekFrom}, io::{Read, Seek, SeekFrom},
}; };
#[serde_as] #[serde_as]
#[derive(Clone, Copy, Debug, PartialEq, Eq)] #[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct IconSize { pub struct IconSize {
pub width: u32, pub width: u32,
pub height: u32, pub height: u32,
} }
impl Display for IconSize { impl Display for IconSize {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
write!(f, "{}x{}", self.width, self.height) write!(f, "{}x{}", self.width, self.height)
} }
} }
impl IconSize { impl IconSize {
pub fn new(width: u32, height: u32) -> Self { pub fn new(width: u32, height: u32) -> Self {
Self { width, height } Self { width, height }
} }
pub fn max_rect(&self) -> u32 { pub fn max_rect(&self) -> u32 {
cmp::max(self.width, self.height) cmp::max(self.width, self.height)
} }
} }
impl Ord for IconSize { impl Ord for IconSize {
fn cmp(&self, other: &Self) -> Ordering { fn cmp(&self, other: &Self) -> Ordering {
other.max_rect().cmp(&self.max_rect()) other.max_rect().cmp(&self.max_rect())
} }
} }
impl PartialOrd for IconSize { impl PartialOrd for IconSize {
fn partial_cmp(&self, other: &Self) -> Option<Ordering> { fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
Some(self.cmp(other)) Some(self.cmp(other))
} }
} }
impl Serialize for IconSize { impl Serialize for IconSize {
fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error> fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
where where
S: Serializer, S: Serializer,
{ {
serializer.collect_str(self) serializer.collect_str(self)
} }
} }
impl<'de> Deserialize<'de> for IconSize { impl<'de> Deserialize<'de> for IconSize {
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error> fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
where where
D: Deserializer<'de>, D: Deserializer<'de>,
{ {
let value: String = Deserialize::deserialize(deserializer)?; let value: String = Deserialize::deserialize(deserializer)?;
let mut split = value.split("x"); let mut split = value.split("x");
let width = split let width = split
.next() .next()
.ok_or(de::Error::custom("expected width"))? .ok_or(de::Error::custom("expected width"))?
.parse() .parse()
.map_err(de::Error::custom)?; .map_err(de::Error::custom)?;
let height = split let height = split
.next() .next()
.ok_or(de::Error::custom("expected height"))? .ok_or(de::Error::custom("expected height"))?
.parse() .parse()
.map_err(de::Error::custom)?; .map_err(de::Error::custom)?;
Ok(IconSize::new(width, height)) Ok(IconSize::new(width, height))
} }
} }
fn slice_eq<T: Read + Seek + Unpin>( fn slice_eq<T: Read + Seek + Unpin>(
cur: &mut T, cur: &mut T,
offset: u64, offset: u64,
slice: &[u8], slice: &[u8],
) -> Result<bool, Box<dyn Error>> { ) -> Result<bool, Box<dyn Error>> {
cur.seek(SeekFrom::Start(offset))?; cur.seek(SeekFrom::Start(offset))?;
let mut buffer = vec![0; slice.len()]; let mut buffer = vec![0; slice.len()];
cur.read_exact(&mut buffer)?; cur.read_exact(&mut buffer)?;
Ok(buffer == slice) Ok(buffer == slice)
} }

18
src/icon/icon_size/png.rs
View file

@ -4,17 +4,17 @@ use futures::prelude::*;
use std::{error::Error, io::Cursor}; use std::{error::Error, io::Cursor};
pub async fn get_png_size<R: AsyncRead + Unpin>( pub async fn get_png_size<R: AsyncRead + Unpin>(
reader: &mut R, reader: &mut R,
) -> Result<IconSize, Box<dyn Error>> { ) -> Result<IconSize, Box<dyn Error>> {
let mut header = [0; 22]; let mut header = [0; 22];
reader.read_exact(&mut header).await?; reader.read_exact(&mut header).await?;
let header = &mut Cursor::new(header); let header = &mut Cursor::new(header);
assert_slice_eq!(header, 0, b"NG\r\n\x1a\n", "bad header"); assert_slice_eq!(header, 0, b"NG\r\n\x1a\n", "bad header");
assert_slice_eq!(header, 10, b"IHDR", "bad header"); assert_slice_eq!(header, 10, b"IHDR", "bad header");
let width = header.read_u32::<BigEndian>()?; let width = header.read_u32::<BigEndian>()?;
let height = header.read_u32::<BigEndian>()?; let height = header.read_u32::<BigEndian>()?;
Ok(IconSize::new(width, height)) Ok(IconSize::new(width, height))
} }

100
src/icon/icon_size/svg.rs
View file

@ -4,65 +4,67 @@ use lol_html::{element, HtmlRewriter, Settings};
use std::{cell::RefCell, error::Error}; use std::{cell::RefCell, error::Error};
fn parse_size<S: ToString>(size: S) -> Option<u32> { fn parse_size<S: ToString>(size: S) -> Option<u32> {
size size.to_string()
.to_string() .parse::<f64>()
.parse::<f64>() .ok()
.ok() .map(|size| size.round() as u32)
.map(|size| size.round() as u32)
} }
pub async fn get_svg_size<R: AsyncRead + Unpin>( pub async fn get_svg_size<R: AsyncRead + Unpin>(
first_bytes: &[u8; 2], first_bytes: &[u8; 2],
reader: &mut R, reader: &mut R,
) -> Result<Option<IconSize>, Box<dyn Error>> { ) -> Result<Option<IconSize>, Box<dyn Error>> {
let size = RefCell::new(None); let size = RefCell::new(None);
let mut rewriter = HtmlRewriter::new( let mut rewriter = HtmlRewriter::new(
Settings { Settings {
element_content_handlers: vec![ element_content_handlers: vec![
// Rewrite insecure hyperlinks // Rewrite insecure hyperlinks
element!("svg", |el| { element!("svg", |el| {
let viewbox = el.get_attribute("viewbox"); let viewbox = el.get_attribute("viewbox");
let width = el.get_attribute("width").and_then(parse_size); let width = el.get_attribute("width").and_then(parse_size);
let height = el.get_attribute("height").and_then(parse_size); let height = el.get_attribute("height").and_then(parse_size);
*size.borrow_mut() = Some(if let (Some(width), Some(height)) = (width, height) { *size.borrow_mut() =
Some(IconSize::new(width, height)) Some(if let (Some(width), Some(height)) = (width, height) {
} else if let Some(viewbox) = viewbox { Some(IconSize::new(width, height))
regex!(r"^-?\d+\s+-?\d+\s+(\d+\.?[\d]?)\s+(\d+\.?[\d]?)") } else if let Some(viewbox) = viewbox {
.captures(&viewbox) regex!(r"^-?\d+\s+-?\d+\s+(\d+\.?[\d]?)\s+(\d+\.?[\d]?)")
.map(|captures| { .captures(&viewbox)
let width = parse_size(captures.get(1).unwrap().as_str()).unwrap(); .map(|captures| {
let height = parse_size(captures.get(2).unwrap().as_str()).unwrap(); let width =
IconSize::new(width, height) parse_size(captures.get(1).unwrap().as_str()).unwrap();
}) let height =
} else { parse_size(captures.get(2).unwrap().as_str()).unwrap();
None IconSize::new(width, height)
}); })
} else {
None
});
Ok(()) Ok(())
}), }),
], ],
..Settings::default() ..Settings::default()
}, },
|_: &[u8]| {}, |_: &[u8]| {},
); );
rewriter.write(first_bytes)?; rewriter.write(first_bytes)?;
let mut buffer = [0; 100]; let mut buffer = [0; 100];
loop { loop {
let n = reader.read(&mut buffer).await?; let n = reader.read(&mut buffer).await?;
if n == 0 { if n == 0 {
return Err("invalid svg".into()); return Err("invalid svg".into());
}
rewriter.write(&buffer[..n])?;
if let Some(size) = *size.borrow() {
return Ok(size);
}
} }
rewriter.write(&buffer[..n])?;
if let Some(size) = *size.borrow() {
return Ok(size);
}
}
} }

157
src/icon/mod.rs
View file

@ -8,116 +8,115 @@ use itertools::Itertools;
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Serialize};
use serde_with::{DeserializeFromStr, SerializeDisplay}; use serde_with::{DeserializeFromStr, SerializeDisplay};
use std::{ use std::{
cmp::Ordering, cmp::Ordering,
collections::HashMap, collections::HashMap,
convert::TryInto, convert::TryInto,
error::Error, error::Error,
fmt::{self, Display}, fmt::{self, Display},
hash::{Hash, Hasher}, hash::{Hash, Hasher},
str::FromStr, str::FromStr,
}; };
use url::Url; use url::Url;
#[derive(Debug, Clone, PartialOrd, PartialEq, Ord, Eq, SerializeDisplay, DeserializeFromStr)] #[derive(Debug, Clone, PartialOrd, PartialEq, Ord, Eq, SerializeDisplay, DeserializeFromStr)]
pub enum IconKind { pub enum IconKind {
AppIcon, AppIcon,
SiteFavicon, SiteFavicon,
SiteLogo, SiteLogo,
} }
impl Display for IconKind { impl Display for IconKind {
fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> { fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
f.write_str(match self { f.write_str(match self {
IconKind::SiteLogo => "site_logo", IconKind::SiteLogo => "site_logo",
IconKind::AppIcon => "app_icon", IconKind::AppIcon => "app_icon",
IconKind::SiteFavicon => "site_favicon", IconKind::SiteFavicon => "site_favicon",
}) })
} }
} }
impl FromStr for IconKind { impl FromStr for IconKind {
type Err = String; type Err = String;
fn from_str(kind: &str) -> Result<Self, Self::Err> { fn from_str(kind: &str) -> Result<Self, Self::Err> {
match kind { match kind {
"site_logo" => Ok(IconKind::SiteLogo), "site_logo" => Ok(IconKind::SiteLogo),
"app_icon" => Ok(IconKind::AppIcon), "app_icon" => Ok(IconKind::AppIcon),
"site_favicon" => Ok(IconKind::SiteFavicon), "site_favicon" => Ok(IconKind::SiteFavicon),
_ => Err("unknown icon kind!".into()), _ => Err("unknown icon kind!".into()),
}
} }
}
} }
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct Icon { pub struct Icon {
pub url: Url, pub url: Url,
pub headers: HashMap<String, String>, pub headers: HashMap<String, String>,
pub kind: IconKind, pub kind: IconKind,
#[serde(flatten)] #[serde(flatten)]
pub info: IconInfo, pub info: IconInfo,
} }
impl Hash for Icon { impl Hash for Icon {
fn hash<H: Hasher>(&self, state: &mut H) { fn hash<H: Hasher>(&self, state: &mut H) {
( (
&self.url, &self.url,
self self.headers
.headers .iter()
.iter() .sorted_by_key(|(key, _)| *key)
.sorted_by_key(|(key, _)| *key) .collect::<Vec<_>>(),
.collect::<Vec<_>>(), )
) .hash(state);
.hash(state); }
}
} }
impl Icon { impl Icon {
pub fn new(url: Url, kind: IconKind, info: IconInfo) -> Self { pub fn new(url: Url, kind: IconKind, info: IconInfo) -> Self {
Icon::new_with_headers(url, HashMap::new(), kind, info) Icon::new_with_headers(url, HashMap::new(), kind, info)
}
pub fn new_with_headers(
url: Url,
headers: HashMap<String, String>,
kind: IconKind,
info: IconInfo,
) -> Self {
Self {
url,
headers,
kind,
info,
} }
}
pub async fn load( pub fn new_with_headers(
url: Url, url: Url,
kind: IconKind, headers: HashMap<String, String>,
sizes: Option<String>, kind: IconKind,
) -> Result<Self, Box<dyn Error>> { info: IconInfo,
Icon::load_with_headers(url, HashMap::new(), kind, sizes).await ) -> Self {
} Self {
url,
headers,
kind,
info,
}
}
pub async fn load_with_headers( pub async fn load(
url: Url, url: Url,
headers: HashMap<String, String>, kind: IconKind,
kind: IconKind, sizes: Option<String>,
sizes: Option<String>, ) -> Result<Self, Box<dyn Error>> {
) -> Result<Self, Box<dyn Error>> { Icon::load_with_headers(url, HashMap::new(), kind, sizes).await
let info = IconInfo::load(url.clone(), (&headers).try_into().unwrap(), sizes).await?; }
Ok(Icon::new_with_headers(url, headers, kind, info)) pub async fn load_with_headers(
} url: Url,
headers: HashMap<String, String>,
kind: IconKind,
sizes: Option<String>,
) -> Result<Self, Box<dyn Error>> {
let info = IconInfo::load(url.clone(), (&headers).try_into().unwrap(), sizes).await?;
Ok(Icon::new_with_headers(url, headers, kind, info))
}
} }
impl Ord for Icon { impl Ord for Icon {
fn cmp(&self, other: &Self) -> Ordering { fn cmp(&self, other: &Self) -> Ordering {
self.info.cmp(&other.info) self.info.cmp(&other.info)
} }
} }
impl PartialOrd for Icon { impl PartialOrd for Icon {
fn partial_cmp(&self, other: &Self) -> Option<Ordering> { fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
Some(self.cmp(other)) Some(self.cmp(other))
} }
} }

358
src/icons.rs
View file

@ -11,198 +11,200 @@ use url::Url;
use vec1::Vec1; use vec1::Vec1;
pub struct SiteIcons { pub struct SiteIcons {
blacklist: Option<Box<dyn Fn(&Url) -> bool>>, blacklist: Option<Box<dyn Fn(&Url) -> bool>>,
} }
#[derive(Debug, Clone)] #[derive(Debug, Clone)]
enum LoadedKind { enum LoadedKind {
DefaultManifest(Option<Vec1<Icon>>), DefaultManifest(Option<Vec1<Icon>>),
HeadTags(Option<Vec1<Icon>>), HeadTags(Option<Vec1<Icon>>),
DefaultFavicon(Option<Icon>), DefaultFavicon(Option<Icon>),
SiteLogo(Option<Icon>), SiteLogo(Option<Icon>),
} }
impl SiteIcons { impl SiteIcons {
pub fn new() -> Self { pub fn new() -> Self {
SiteIcons { blacklist: None } SiteIcons { blacklist: None }
}
pub fn new_with_blacklist(blacklist: impl Fn(&Url) -> bool + 'static) -> Self {
SiteIcons {
blacklist: Some(Box::new(blacklist)),
} }
}
pub fn is_blacklisted(&self, url: &Url) -> bool { pub fn new_with_blacklist(blacklist: impl Fn(&Url) -> bool + 'static) -> Self {
if let Some(is_blacklisted) = &self.blacklist { SiteIcons {
is_blacklisted(url) blacklist: Some(Box::new(blacklist)),
} else {
false
}
}
pub async fn load_website<U: IntoUrl>(
&mut self,
url: U,
best_matches_only: bool,
) -> Result<Vec<Icon>, Box<dyn Error>> {
let url = url.into_url()?;
let manifest_urls = vec![
push_url(&url, "manifest.json"),
push_url(&url, "manifest.webmanifest"),
url.join("/manifest.json")?,
url.join("/manifest.webmanifest")?,
]
.into_iter()
.unique();
let favicon_urls = vec![
push_url(&url, "favicon.svg"),
url.join("/favicon.svg")?,
push_url(&url, "favicon.ico"),
url.join("/favicon.ico")?,
]
.into_iter()
.unique();
let html_response = async {
let res = CLIENT
.get(url.clone())
.header(ACCEPT, "text/html")
.send()
.await
.ok()?
.error_for_status()
.ok()?;
let url = res.url().clone();
if self.is_blacklisted(&url) {
None
} else {
let body = res.bytes_stream().map(|res| {
res
.map(|bytes| bytes.to_vec())
.map_err(|err| err.to_string())
});
let mut publisher = Publisher::new(128);
let subscriber = publisher.subscribe();
Some((
url,
async move { StreamPublisher::new(&mut publisher, body).await }.shared(),
subscriber,
))
}
}
.shared();
let mut futures = vec![
async {
let html_response = html_response.clone().await;
LoadedKind::HeadTags(match html_response {
Some((url, _, body)) => html_parser::parse_head(&url, body)
.await
.ok()
.and_then(|icons| icons.try_into().ok()),
None => None,
})
}
.boxed_local(),
async {
let html_response = html_response.clone().await;
LoadedKind::SiteLogo(match html_response {
Some((url, complete, body)) => {
let (icons, _) = join!(
html_parser::parse_site_logo(&url, body, |url| self.is_blacklisted(url)),
complete
);
icons.ok()
}
None => None,
})
}
.boxed_local(),
async {
let manifests = join_all(manifest_urls.map(|url| SiteIcons::load_manifest(url))).await;
LoadedKind::DefaultManifest(
manifests
.into_iter()
.find_map(|manifest| manifest.ok().and_then(|icons| icons.try_into().ok())),
)
}
.boxed_local(),
async {
let favicons =
join_all(favicon_urls.map(|url| Icon::load(url.clone(), IconKind::SiteFavicon, None)))
.await;
LoadedKind::DefaultFavicon(favicons.into_iter().find_map(|favicon| favicon.ok()))
}
.boxed_local(),
];
let mut icons: Vec<Icon> = Vec::new();
let mut found_best_match = false;
let mut previous_loads = Vec::new();
while !futures.is_empty() {
let (loaded, index, _) = select_all(&mut futures).await;
futures.remove(index);
match loaded.clone() {
LoadedKind::DefaultManifest(manifest_icons) => {
if let Some(manifest_icons) = manifest_icons {
icons.extend(manifest_icons);
found_best_match = true;
}
} }
LoadedKind::DefaultFavicon(favicon) => { }
if let Some(favicon) = favicon {
icons.push(favicon);
if previous_loads pub fn is_blacklisted(&self, url: &Url) -> bool {
.iter() if let Some(is_blacklisted) = &self.blacklist {
.any(|kind| matches!(kind, LoadedKind::HeadTags(_))) is_blacklisted(url)
{ } else {
found_best_match = true; false
}
}
pub async fn load_website<U: IntoUrl>(
&mut self,
url: U,
best_matches_only: bool,
) -> Result<Vec<Icon>, Box<dyn Error>> {
let url = url.into_url()?;
let manifest_urls = vec![
push_url(&url, "manifest.json"),
push_url(&url, "manifest.webmanifest"),
url.join("/manifest.json")?,
url.join("/manifest.webmanifest")?,
]
.into_iter()
.unique();
let favicon_urls = vec![
push_url(&url, "favicon.svg"),
url.join("/favicon.svg")?,
push_url(&url, "favicon.ico"),
url.join("/favicon.ico")?,
]
.into_iter()
.unique();
let html_response = async {
let res = CLIENT
.get(url.clone())
.header(ACCEPT, "text/html")
.send()
.await
.ok()?
.error_for_status()
.ok()?;
let url = res.url().clone();
if self.is_blacklisted(&url) {
None
} else {
let body = res.bytes_stream().map(|res| {
res.map(|bytes| bytes.to_vec())
.map_err(|err| err.to_string())
});
let mut publisher = Publisher::new(128);
let subscriber = publisher.subscribe();
Some((
url,
async move { StreamPublisher::new(&mut publisher, body).await }.shared(),
subscriber,
))
} }
}
} }
LoadedKind::HeadTags(head_icons) => { .shared();
if let Some(head_icons) = head_icons {
icons.extend(head_icons); let mut futures = vec![
found_best_match = true; async {
} else if previous_loads let html_response = html_response.clone().await;
.iter()
.any(|kind| matches!(kind, LoadedKind::DefaultFavicon(Some(_)))) LoadedKind::HeadTags(match html_response {
{ Some((url, _, body)) => html_parser::parse_head(&url, body)
found_best_match = true; .await
} .ok()
.and_then(|icons| icons.try_into().ok()),
None => None,
})
}
.boxed_local(),
async {
let html_response = html_response.clone().await;
LoadedKind::SiteLogo(match html_response {
Some((url, complete, body)) => {
let (icons, _) = join!(
html_parser::parse_site_logo(&url, body, |url| self
.is_blacklisted(url)),
complete
);
icons.ok()
}
None => None,
})
}
.boxed_local(),
async {
let manifests =
join_all(manifest_urls.map(|url| SiteIcons::load_manifest(url))).await;
LoadedKind::DefaultManifest(
manifests
.into_iter()
.find_map(|manifest| manifest.ok().and_then(|icons| icons.try_into().ok())),
)
}
.boxed_local(),
async {
let favicons = join_all(
favicon_urls.map(|url| Icon::load(url.clone(), IconKind::SiteFavicon, None)),
)
.await;
LoadedKind::DefaultFavicon(favicons.into_iter().find_map(|favicon| favicon.ok()))
}
.boxed_local(),
];
let mut icons: Vec<Icon> = Vec::new();
let mut found_best_match = false;
let mut previous_loads = Vec::new();
while !futures.is_empty() {
let (loaded, index, _) = select_all(&mut futures).await;
futures.remove(index);
match loaded.clone() {
LoadedKind::DefaultManifest(manifest_icons) => {
if let Some(manifest_icons) = manifest_icons {
icons.extend(manifest_icons);
found_best_match = true;
}
}
LoadedKind::DefaultFavicon(favicon) => {
if let Some(favicon) = favicon {
icons.push(favicon);
if previous_loads
.iter()
.any(|kind| matches!(kind, LoadedKind::HeadTags(_)))
{
found_best_match = true;
}
}
}
LoadedKind::HeadTags(head_icons) => {
if let Some(head_icons) = head_icons {
icons.extend(head_icons);
found_best_match = true;
} else if previous_loads
.iter()
.any(|kind| matches!(kind, LoadedKind::DefaultFavicon(Some(_))))
{
found_best_match = true;
}
}
LoadedKind::SiteLogo(logo) => {
if let Some(logo) = logo {
icons.push(logo);
}
}
}
previous_loads.push(loaded);
icons.sort();
icons = icons.into_iter().unique().collect();
if best_matches_only && found_best_match {
break;
}
} }
LoadedKind::SiteLogo(logo) => {
if let Some(logo) = logo {
icons.push(logo);
}
}
}
previous_loads.push(loaded); Ok(icons)
icons.sort();
icons = icons.into_iter().unique().collect();
if best_matches_only && found_best_match {
break;
}
} }
Ok(icons)
}
} }

10
src/lib.rs
View file

@ -35,12 +35,12 @@ pub use icons::*;
use once_cell::sync::Lazy; use once_cell::sync::Lazy;
use reqwest::{ use reqwest::{
header::{HeaderMap, HeaderValue, USER_AGENT}, header::{HeaderMap, HeaderValue, USER_AGENT},
Client, Client,
}; };
static CLIENT: Lazy<Client> = Lazy::new(|| { static CLIENT: Lazy<Client> = Lazy::new(|| {
let mut headers = HeaderMap::new(); let mut headers = HeaderMap::new();
headers.insert(USER_AGENT, HeaderValue::from_str("Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/88.0.4324.104 Safari/537.36").unwrap()); headers.insert(USER_AGENT, HeaderValue::from_str("Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/88.0.4324.104 Safari/537.36").unwrap());
Client::builder().default_headers(headers).build().unwrap() Client::builder().default_headers(headers).build().unwrap()
}); });

52
src/manifest.rs
View file

@ -8,49 +8,47 @@ use url::Url;
#[derive(Debug, Deserialize)] #[derive(Debug, Deserialize)]
struct ManifestIcon { struct ManifestIcon {
src: String, src: String,
sizes: Option<String>, sizes: Option<String>,
} }
#[derive(Debug, Deserialize)] #[derive(Debug, Deserialize)]
struct Manifest { struct Manifest {
icons: Vec<ManifestIcon>, icons: Vec<ManifestIcon>,
} }
impl SiteIcons { impl SiteIcons {
pub async fn load_manifest<U: IntoUrl>(url: U) -> Result<Vec<Icon>, Box<dyn Error>> { pub async fn load_manifest<U: IntoUrl>(url: U) -> Result<Vec<Icon>, Box<dyn Error>> {
let url = url.into_url()?; let url = url.into_url()?;
Ok(load_manifest_cached(url).await?) Ok(load_manifest_cached(url).await?)
} }
} }
#[cached(sync_writes = true)] #[cached(sync_writes = true)]
async fn load_manifest_cached(url: Url) -> Result<Vec<Icon>, String> { async fn load_manifest_cached(url: Url) -> Result<Vec<Icon>, String> {
let url = &url; let url = &url;
let manifest: Manifest = CLIENT let manifest: Manifest = CLIENT
.get(url.clone()) .get(url.clone())
.send() .send()
.await .await
.map_err(|e| format!("{}: {:?}", url, e))? .map_err(|e| format!("{}: {:?}", url, e))?
.error_for_status() .error_for_status()
.map_err(|e| format!("{}: {:?}", url, e))? .map_err(|e| format!("{}: {:?}", url, e))?
.json() .json()
.await .await
.map_err(|e| format!("{}: {:?}", url, e))?; .map_err(|e| format!("{}: {:?}", url, e))?;
Ok( Ok(join_all(manifest.icons.into_iter().map(|icon| async move {
join_all(manifest.icons.into_iter().map(|icon| async move { if let Ok(src) = url.join(&icon.src) {
if let Ok(src) = url.join(&icon.src) { Icon::load(src, IconKind::AppIcon, icon.sizes).await.ok()
Icon::load(src, IconKind::AppIcon, icon.sizes).await.ok() } else {
} else { None
None }
}
})) }))
.await .await
.into_iter() .into_iter()
.filter_map(|icon| icon) .filter_map(|icon| icon)
.collect(), .collect())
)
} }

62
src/utils/background_poll.rs
View file

@ -1,43 +1,43 @@
use std::{ use std::{
pin::Pin, pin::Pin,
task::{Context, Poll}, task::{Context, Poll},
}; };
use futures::Future; use futures::Future;
pub async fn poll_in_background<F, B, FO, BO>(future: F, background_future: B) -> FO pub async fn poll_in_background<F, B, FO, BO>(future: F, background_future: B) -> FO
where where
F: Future<Output = FO> + Unpin,
B: Future<Output = BO> + Unpin,
{
struct BackgroundPoller<F, B> {
future: F,
background_future: B,
}
impl<F, B, FO, BO> Future for BackgroundPoller<F, B>
where
F: Future<Output = FO> + Unpin, F: Future<Output = FO> + Unpin,
B: Future<Output = BO> + Unpin, B: Future<Output = BO> + Unpin,
{ {
type Output = FO; struct BackgroundPoller<F, B> {
future: F,
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { background_future: B,
let this = self.get_mut();
let result = Pin::new(&mut this.future).poll(cx);
if result.is_pending() {
let _ = Pin::new(&mut this.background_future).poll(cx);
}
result
} }
}
BackgroundPoller { impl<F, B, FO, BO> Future for BackgroundPoller<F, B>
future, where
background_future, F: Future<Output = FO> + Unpin,
} B: Future<Output = BO> + Unpin,
.await {
type Output = FO;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let this = self.get_mut();
let result = Pin::new(&mut this.future).poll(cx);
if result.is_pending() {
let _ = Pin::new(&mut this.background_future).poll(cx);
}
result
}
}
BackgroundPoller {
future,
background_future,
}
.await
} }

8
src/utils/macros.rs
View file

@ -12,10 +12,10 @@ macro_rules! join_with {
} }
macro_rules! regex { macro_rules! regex {
($re:literal $(,)?) => {{ ($re:literal $(,)?) => {{
static RE: once_cell::sync::OnceCell<regex::Regex> = once_cell::sync::OnceCell::new(); static RE: once_cell::sync::OnceCell<regex::Regex> = once_cell::sync::OnceCell::new();
RE.get_or_init(|| regex::Regex::new($re).unwrap()) RE.get_or_init(|| regex::Regex::new($re).unwrap())
}}; }};
} }
macro_rules! assert_slice_eq { macro_rules! assert_slice_eq {

6
src/utils/mod.rs
View file

@ -10,7 +10,7 @@ pub use svg_encoder::*;
use url::Url; use url::Url;
pub fn push_url(url: &Url, segment: &str) -> Url { pub fn push_url(url: &Url, segment: &str) -> Url {
let mut url = url.clone(); let mut url = url.clone();
url.path_segments_mut().unwrap().push(segment); url.path_segments_mut().unwrap().push(segment);
url url
} }

80
src/utils/svg_encoder.rs
View file

@ -2,52 +2,52 @@ use percent_encoding::{utf8_percent_encode, AsciiSet, CONTROLS};
use std::borrow::Cow; use std::borrow::Cow;
const DATA_URI: &AsciiSet = &CONTROLS const DATA_URI: &AsciiSet = &CONTROLS
.add(b'\r') .add(b'\r')
.add(b'\n') .add(b'\n')
.add(b'%') .add(b'%')
.add(b'#') .add(b'#')
.add(b'(') .add(b'(')
.add(b')') .add(b')')
.add(b'<') .add(b'<')
.add(b'>') .add(b'>')
.add(b'?') .add(b'?')
.add(b'[') .add(b'[')
.add(b'\\') .add(b'\\')
.add(b']') .add(b']')
.add(b'^') .add(b'^')
.add(b'`') .add(b'`')
.add(b'{') .add(b'{')
.add(b'|') .add(b'|')
.add(b'}'); .add(b'}');
pub fn encode_svg(svg: &str) -> String { pub fn encode_svg(svg: &str) -> String {
// add namespace // add namespace
let encoded = if !svg.contains("http://www.w3.org/2000/svg") { let encoded = if !svg.contains("http://www.w3.org/2000/svg") {
regex!("<svg").replace(svg, "<svg xmlns='http://www.w3.org/2000/svg'") regex!("<svg").replace(svg, "<svg xmlns='http://www.w3.org/2000/svg'")
} else { } else {
svg.into() svg.into()
}; };
// use single quotes instead of double to avoid encoding. // use single quotes instead of double to avoid encoding.
let mut encoded = regex!("\"").replace_all(&encoded, "'"); let mut encoded = regex!("\"").replace_all(&encoded, "'");
// remove a fill=none attribute // remove a fill=none attribute
if let Some(captures) = regex!("^[^>]+fill='?(none)'?").captures(&encoded) { if let Some(captures) = regex!("^[^>]+fill='?(none)'?").captures(&encoded) {
let index = captures.get(1).unwrap(); let index = captures.get(1).unwrap();
let mut result = String::new(); let mut result = String::new();
for (i, c) in encoded.chars().enumerate() { for (i, c) in encoded.chars().enumerate() {
if i < index.start() || i >= index.end() { if i < index.start() || i >= index.end() {
result.push(c); result.push(c);
} }
}
encoded = Cow::from(result);
} }
encoded = Cow::from(result);
}
// remove whitespace // remove whitespace
let encoded = regex!(r">\s{1,}<").replace_all(&encoded, "><"); let encoded = regex!(r">\s{1,}<").replace_all(&encoded, "><");
let encoded = regex!(r"\s{2,}").replace_all(&encoded, " "); let encoded = regex!(r"\s{2,}").replace_all(&encoded, " ");
let encoded = utf8_percent_encode(&encoded, DATA_URI); let encoded = utf8_percent_encode(&encoded, DATA_URI);
format!("data:image/svg+xml,{}", encoded) format!("data:image/svg+xml,{}", encoded)
} }