@ -1,19 +1,19 @@
use std ::{ use std ::{
collections ::HashMap , collections ::HashMap ,
fs ::{ create_dir_all , remove_file , symlink_metadata , File } ,
io ::prelude ::* ,
net ::{ IpAddr , ToSocketAddrs } , net ::{ IpAddr , ToSocketAddrs } ,
sync ::{ Arc , RwLock } , sync ::{ Arc , RwLock } ,
time ::{ Duration , SystemTime } , time ::{ Duration , SystemTime } ,
} ; } ;
use bytes ::{ Buf , Bytes , BytesMut } ;
use futures ::{ stream ::StreamExt , TryFutureExt } ;
use once_cell ::sync ::Lazy ; use once_cell ::sync ::Lazy ;
use regex ::Regex ; use regex ::Regex ;
use reqwest ::{ blocking::Client , blocking ::Response , header } ; use reqwest ::{ header, Client , Response } ;
use rocket ::{ use rocket ::{ http ::ContentType , response ::Redirect , Route } ;
http ::ContentType , use tokio ::{
response::{ Content , Redirect } , fs::{ create_dir_all , remove_file , symlink_metadata , File } ,
Route , io::{ AsyncReadExt , AsyncWriteExt } ,
} ; } ;
use crate ::{ use crate ::{
@ -104,27 +104,23 @@ fn icon_google(domain: String) -> Option<Redirect> {
} }
#[ get( " /<domain>/icon.png " ) ] #[ get( " /<domain>/icon.png " ) ]
fn icon_internal ( domain : String ) -> Cached < Content < Vec < u8 > > > { async fn icon_internal ( domain : String ) -> Cached < ( ContentType , Vec < u8 > ) > {
const FALLBACK_ICON : & [ u8 ] = include_bytes! ( "../static/images/fallback-icon.png" ) ; const FALLBACK_ICON : & [ u8 ] = include_bytes! ( "../static/images/fallback-icon.png" ) ;
if ! is_valid_domain ( & domain ) { if ! is_valid_domain ( & domain ) {
warn ! ( "Invalid domain: {}" , domain ) ; warn ! ( "Invalid domain: {}" , domain ) ;
return Cached ::ttl ( return Cached ::ttl (
Content ( ContentType ::new ( "image" , "png" ) , FALLBACK_ICON . to_vec ( ) ) , ( ContentType ::new ( "image" , "png" ) , FALLBACK_ICON . to_vec ( ) ) ,
CONFIG . icon_cache_negttl ( ) , CONFIG . icon_cache_negttl ( ) ,
true , true ,
) ; ) ;
} }
match get_icon ( & domain ) { match get_icon ( & domain ) . await {
Some ( ( icon , icon_type ) ) = > { Some ( ( icon , icon_type ) ) = > {
Cached ::ttl ( Content ( ContentType ::new ( "image" , icon_type ) , icon ) , CONFIG . icon_cache_ttl ( ) , true ) Cached ::ttl ( ( ContentType ::new ( "image" , icon_type ) , icon ) , CONFIG . icon_cache_ttl ( ) , true )
} }
_ = > Cached ::ttl ( _ = > Cached ::ttl ( ( ContentType ::new ( "image" , "png" ) , FALLBACK_ICON . to_vec ( ) ) , CONFIG . icon_cache_negttl ( ) , true ) ,
Content ( ContentType ::new ( "image" , "png" ) , FALLBACK_ICON . to_vec ( ) ) ,
CONFIG . icon_cache_negttl ( ) ,
true ,
) ,
} }
} }
@ -317,15 +313,15 @@ fn is_domain_blacklisted(domain: &str) -> bool {
is_blacklisted is_blacklisted
} }
fn get_icon ( domain : & str ) -> Option < ( Vec < u8 > , String ) > { async fn get_icon ( domain : & str ) -> Option < ( Vec < u8 > , String ) > {
let path = format! ( "{}/{}.png" , CONFIG . icon_cache_folder ( ) , domain ) ; let path = format! ( "{}/{}.png" , CONFIG . icon_cache_folder ( ) , domain ) ;
// Check for expiration of negatively cached copy
// Check for expiration of negatively cached copy
if icon_is_negcached ( & path ) { if icon_is_negcached ( & path ) . await {
return None ; return None ;
} }
if let Some ( icon ) = get_cached_icon ( & path ) { if let Some ( icon ) = get_cached_icon ( & path ) . await {
let icon_type = match get_icon_type ( & icon ) { let icon_type = match get_icon_type ( & icon ) {
Some ( x ) = > x , Some ( x ) = > x ,
_ = > "x-icon" , _ = > "x-icon" ,
@ -338,31 +334,31 @@ fn get_icon(domain: &str) -> Option<(Vec<u8>, String)> {
} }
// Get the icon, or None in case of error
// Get the icon, or None in case of error
match download_icon ( domain ) { match download_icon ( domain ) . await {
Ok ( ( icon , icon_type ) ) = > { Ok ( ( icon , icon_type ) ) = > {
save_icon ( & path , & icon ) save_icon ( & path , & icon ) .await ;
Some ( ( icon icon_type . unwrap_or ( "x-icon" ) . to_string ( ) ) ) Some ( ( icon .to_vec ( ) , icon_type . unwrap_or ( "x-icon" ) . to_string ( ) ) )
} }
Err ( e ) = > { Err ( e ) = > {
warn ! ( "Unable to download icon: {:?}" , e ) ; warn ! ( "Unable to download icon: {:?}" , e ) ;
let miss_indicator = path + ".miss" ; let miss_indicator = path + ".miss" ;
save_icon ( & miss_indicator , & [ ] ) save_icon ( & miss_indicator , & [ ] ) .await ;
None None
} }
} }
} }
fn get_cached_icon ( path : & str ) -> Option < Vec < u8 > > { async fn get_cached_icon ( path : & str ) -> Option < Vec < u8 > > {
// Check for expiration of successfully cached copy
// Check for expiration of successfully cached copy
if icon_is_expired ( path ) { if icon_is_expired ( path ) . await {
return None ; return None ;
} }
// Try to read the cached icon, and return it if it exists
// Try to read the cached icon, and return it if it exists
if let Ok ( mut f ) = File ::open ( path ) { if let Ok ( mut f ) = File ::open ( path ) . await {
let mut buffer = Vec ::new ( ) ; let mut buffer = Vec ::new ( ) ;
if f . read_to_end ( & mut buffer ) . is_ok ( ) { if f . read_to_end ( & mut buffer ) . await . is_ok ( ) {
return Some ( buffer ) ; return Some ( buffer ) ;
} }
} }
@ -370,22 +366,22 @@ fn get_cached_icon(path: &str) -> Option<Vec<u8>> {
None None
} }
fn file_is_expired ( path : & str , ttl : u64 ) -> Result < bool , Error > { async fn file_is_expired ( path : & str , ttl : u64 ) -> Result < bool , Error > {
let meta = symlink_metadata ( path ) ? ; let meta = symlink_metadata ( path ) . await ? ;
let modified = meta . modified ( ) ? ; let modified = meta . modified ( ) ? ;
let age = SystemTime ::now ( ) . duration_since ( modified ) ? ; let age = SystemTime ::now ( ) . duration_since ( modified ) ? ;
Ok ( ttl > 0 & & ttl < = age . as_secs ( ) ) Ok ( ttl > 0 & & ttl < = age . as_secs ( ) )
} }
fn icon_is_negcached ( path : & str ) -> bool { async fn icon_is_negcached ( path : & str ) -> bool {
let miss_indicator = path . to_owned ( ) + ".miss" ; let miss_indicator = path . to_owned ( ) + ".miss" ;
let expired = file_is_expired ( & miss_indicator , CONFIG . icon_cache_negttl ( ) ) let expired = file_is_expired ( & miss_indicator , CONFIG . icon_cache_negttl ( ) ) .await ;
match expired { match expired {
// No longer negatively cached, drop the marker
// No longer negatively cached, drop the marker
Ok ( true ) = > { Ok ( true ) = > {
if let Err ( e ) = remove_file ( & miss_indicator ) { if let Err ( e ) = remove_file ( & miss_indicator ) . await {
error ! ( "Could not remove negative cache indicator for icon {:?}: {:?}" , path , e ) ; error ! ( "Could not remove negative cache indicator for icon {:?}: {:?}" , path , e ) ;
} }
false false
@ -397,8 +393,8 @@ fn icon_is_negcached(path: &str) -> bool {
} }
} }
fn icon_is_expired ( path : & str ) -> bool { async fn icon_is_expired ( path : & str ) -> bool {
let expired = file_is_expired ( path , CONFIG . icon_cache_ttl ( ) ) let expired = file_is_expired ( path , CONFIG . icon_cache_ttl ( ) ) .await ;
expired . unwrap_or ( true ) expired . unwrap_or ( true )
} }
@ -521,13 +517,13 @@ struct IconUrlResult {
/// let icon_result = get_icon_url("github.com")?;
/// let icon_result = get_icon_url("github.com")?;
/// let icon_result = get_icon_url("vaultwarden.discourse.group")?;
/// let icon_result = get_icon_url("vaultwarden.discourse.group")?;
/// ```
/// ```
fn get_icon_url ( domain : & str ) -> Result < IconUrlResult , Error > { async fn get_icon_url ( domain : & str ) -> Result < IconUrlResult , Error > {
// Default URL with secure and insecure schemes
// Default URL with secure and insecure schemes
let ssldomain = format! ( "https://{}" , domain ) ; let ssldomain = format! ( "https://{}" , domain ) ;
let httpdomain = format! ( "http://{}" , domain ) ; let httpdomain = format! ( "http://{}" , domain ) ;
// First check the domain as given during the request for both HTTPS and HTTP.
// First check the domain as given during the request for both HTTPS and HTTP.
let resp = match get_page ( & ssldomain ) . or_else ( | _ | get_page ( & httpdomain ) ) { let resp = match get_page ( & ssldomain ) . or_else ( | _ | get_page ( & httpdomain ) ) . await {
Ok ( c ) = > Ok ( c ) , Ok ( c ) = > Ok ( c ) ,
Err ( e ) = > { Err ( e ) = > {
let mut sub_resp = Err ( e ) ; let mut sub_resp = Err ( e ) ;
@ -546,7 +542,7 @@ fn get_icon_url(domain: &str) -> Result<IconUrlResult, Error> {
let httpbase = format! ( "http://{}" , base_domain ) ; let httpbase = format! ( "http://{}" , base_domain ) ;
debug ! ( "[get_icon_url]: Trying without subdomains '{}'" , base_domain ) ; debug ! ( "[get_icon_url]: Trying without subdomains '{}'" , base_domain ) ;
sub_resp = get_page ( & sslbase ) . or_else ( | _ | get_page ( & httpbase ) ) sub_resp = get_page ( & sslbase ) . or_else ( | _ | get_page ( & httpbase ) ) .await ;
} }
// When the domain is not an IP, and has less then 2 dots, try to add www. infront of it.
// When the domain is not an IP, and has less then 2 dots, try to add www. infront of it.
@ -557,7 +553,7 @@ fn get_icon_url(domain: &str) -> Result<IconUrlResult, Error> {
let httpwww = format! ( "http://{}" , www_domain ) ; let httpwww = format! ( "http://{}" , www_domain ) ;
debug ! ( "[get_icon_url]: Trying with www. prefix '{}'" , www_domain ) ; debug ! ( "[get_icon_url]: Trying with www. prefix '{}'" , www_domain ) ;
sub_resp = get_page ( & sslwww ) . or_else ( | _ | get_page ( & httpwww ) ) sub_resp = get_page ( & sslwww ) . or_else ( | _ | get_page ( & httpwww ) ) .await ;
} }
} }
@ -581,7 +577,7 @@ fn get_icon_url(domain: &str) -> Result<IconUrlResult, Error> {
iconlist . push ( Icon ::new ( 35 , String ::from ( url . join ( "/favicon.ico" ) . unwrap ( ) ) ) ) ; iconlist . push ( Icon ::new ( 35 , String ::from ( url . join ( "/favicon.ico" ) . unwrap ( ) ) ) ) ;
// 384KB should be more than enough for the HTML, though as we only really need the HTML header.
// 384KB should be more than enough for the HTML, though as we only really need the HTML header.
let mut limited_reader = content. take ( 384 * 1024 ) ; let mut limited_reader = stream_to_bytes_limit( content , 384 * 1024 ) . await ? . reader ( ) ;
use html5ever ::tendril ::TendrilSink ; use html5ever ::tendril ::TendrilSink ;
let dom = html5ever ::parse_document ( markup5ever_rcdom ::RcDom ::default ( ) , Default ::default ( ) ) let dom = html5ever ::parse_document ( markup5ever_rcdom ::RcDom ::default ( ) , Default ::default ( ) )
@ -607,11 +603,11 @@ fn get_icon_url(domain: &str) -> Result<IconUrlResult, Error> {
} ) } )
} }
fn get_page ( url : & str ) -> Result < Response , Error > { async fn get_page ( url : & str ) -> Result < Response , Error > {
get_page_with_referer ( url , "" ) get_page_with_referer ( url , "" ) . await
} }
fn get_page_with_referer ( url : & str , referer : & str ) -> Result < Response , Error > { async fn get_page_with_referer ( url : & str , referer : & str ) -> Result < Response , Error > {
if is_domain_blacklisted ( url ::Url ::parse ( url ) . unwrap ( ) . host_str ( ) . unwrap_or_default ( ) ) { if is_domain_blacklisted ( url ::Url ::parse ( url ) . unwrap ( ) . host_str ( ) . unwrap_or_default ( ) ) {
warn ! ( "Favicon '{}' resolves to a blacklisted domain or IP!" , url ) ; warn ! ( "Favicon '{}' resolves to a blacklisted domain or IP!" , url ) ;
} }
@ -621,7 +617,7 @@ fn get_page_with_referer(url: &str, referer: &str) -> Result<Response, Error> {
client = client . header ( "Referer" , referer ) client = client . header ( "Referer" , referer )
} }
match client . send ( ) { match client . send ( ) . await {
Ok ( c ) = > c . error_for_status ( ) . map_err ( Into ::into ) , Ok ( c ) = > c . error_for_status ( ) . map_err ( Into ::into ) ,
Err ( e ) = > err_silent ! ( format! ( "{}" , e ) ) , Err ( e ) = > err_silent ! ( format! ( "{}" , e ) ) ,
} }
@ -706,14 +702,14 @@ fn parse_sizes(sizes: Option<&str>) -> (u16, u16) {
( width , height ) ( width , height )
} }
fn download_icon ( domain : & str ) -> Result < ( Vec < u8 > , Option < & str > ) , Error > { async fn download_icon ( domain : & str ) -> Result < ( Bytes , Option < & str > ) , Error > {
if is_domain_blacklisted ( domain ) { if is_domain_blacklisted ( domain ) {
err_silent ! ( "Domain is blacklisted" , domain ) err_silent ! ( "Domain is blacklisted" , domain )
} }
let icon_result = get_icon_url ( domain ) ? ; let icon_result = get_icon_url ( domain ) . await ? ;
let mut buffer = Vec new ( ) ; let mut buffer = Bytes new ( ) ;
let mut icon_type : Option < & str > = None ; let mut icon_type : Option < & str > = None ;
use data_url ::DataUrl ; use data_url ::DataUrl ;
@ -722,8 +718,12 @@ fn download_icon(domain: &str) -> Result<(Vec<u8>, Option<&str>), Error> {
if icon . href . starts_with ( "data:image" ) { if icon . href . starts_with ( "data:image" ) {
let datauri = DataUrl ::process ( & icon . href ) . unwrap ( ) ; let datauri = DataUrl ::process ( & icon . href ) . unwrap ( ) ;
// Check if we are able to decode the data uri
// Check if we are able to decode the data uri
match datauri . decode_to_vec ( ) { let mut body = BytesMut ::new ( ) ;
Ok ( ( body , _fragment ) ) = > { match datauri . decode ::< _ , ( ) > ( | bytes | {
body . extend_from_slice ( bytes ) ;
Ok ( ( ) )
} ) {
Ok ( _ ) = > {
// Also check if the size is atleast 67 bytes, which seems to be the smallest png i could create
// Also check if the size is atleast 67 bytes, which seems to be the smallest png i could create
if body . len ( ) > = 67 { if body . len ( ) > = 67 {
// Check if the icon type is allowed, else try an icon from the list.
// Check if the icon type is allowed, else try an icon from the list.
@ -733,16 +733,16 @@ fn download_icon(domain: &str) -> Result<(Vec<u8>, Option<&str>), Error> {
continue ; continue ;
} }
info ! ( "Extracted icon from data:image uri for {}" , domain ) ; info ! ( "Extracted icon from data:image uri for {}" , domain ) ;
buffer = body buffer = body .freeze ( ) ;
break ; break ;
} }
} }
_ = > debug ! ( "Extracted icon from data:image uri is invalid" ) , _ = > debug ! ( "Extracted icon from data:image uri is invalid" ) ,
} ; } ;
} else { } else {
match get_page_with_referer ( & icon . href , & icon_result . referer ) { match get_page_with_referer ( & icon . href , & icon_result . referer ) . await {
Ok ( mut res ) = > { Ok ( res ) = > {
res. copy_to ( & mut buffer ) ? ; buffer = stream_to_bytes_limit ( res , 512 * 1024 ) . await ? ; // 512 KB for each icon max
// Check if the icon type is allowed, else try an icon from the list.
// Check if the icon type is allowed, else try an icon from the list.
icon_type = get_icon_type ( & buffer ) ; icon_type = get_icon_type ( & buffer ) ;
if icon_type . is_none ( ) { if icon_type . is_none ( ) {
@ -765,13 +765,13 @@ fn download_icon(domain: &str) -> Result<(Vec<u8>, Option<&str>), Error> {
Ok ( ( buffer , icon_type ) ) Ok ( ( buffer , icon_type ) )
} }
fn save_icon ( path : & str , icon : & [ u8 ] ) { async fn save_icon ( path : & str , icon : & [ u8 ] ) {
match File ::create ( path ) { match File ::create ( path ) . await {
Ok ( mut f ) = > { Ok ( mut f ) = > {
f . write_all ( icon ) . expect ( "Error writing icon file" ) ; f . write_all ( icon ) . await . expect ( "Error writing icon file" ) ;
} }
Err ( ref e ) if e . kind ( ) = = std ::io ::ErrorKind ::NotFound = > { Err ( ref e ) if e . kind ( ) = = std ::io ::ErrorKind ::NotFound = > {
create_dir_all ( & CONFIG . icon_cache_folder ( ) ) . expect ( "Error creating icon cache folder" ) ; create_dir_all ( & CONFIG . icon_cache_folder ( ) ) . await . expect ( "Error creating icon cache folder" ) ;
} }
Err ( e ) = > { Err ( e ) = > {
warn ! ( "Unable to save icon: {:?}" , e ) ; warn ! ( "Unable to save icon: {:?}" , e ) ;
@ -820,8 +820,6 @@ impl reqwest::cookie::CookieStore for Jar {
} }
fn cookies ( & self , url : & url ::Url ) -> Option < header ::HeaderValue > { fn cookies ( & self , url : & url ::Url ) -> Option < header ::HeaderValue > {
use bytes ::Bytes ;
let cookie_store = self . 0. read ( ) . unwrap ( ) ; let cookie_store = self . 0. read ( ) . unwrap ( ) ;
let s = cookie_store let s = cookie_store
. get_request_values ( url ) . get_request_values ( url )
@ -836,3 +834,12 @@ impl reqwest::cookie::CookieStore for Jar {
header ::HeaderValue ::from_maybe_shared ( Bytes ::from ( s ) ) . ok ( ) header ::HeaderValue ::from_maybe_shared ( Bytes ::from ( s ) ) . ok ( )
} }
} }
async fn stream_to_bytes_limit ( res : Response , max_size : usize ) -> Result < Bytes , reqwest ::Error > {
let mut stream = res . bytes_stream ( ) . take ( max_size ) ;
let mut buf = BytesMut ::new ( ) ;
while let Some ( chunk ) = stream . next ( ) . await {
buf . extend ( chunk ? ) ;
}
Ok ( buf . freeze ( ) )
}
Daniel García
3 years ago