diff --git a/src/api/icons.rs b/src/api/icons.rs index 51cb504..8ae3c83 100644 --- a/src/api/icons.rs +++ b/src/api/icons.rs @@ -8,6 +8,8 @@ use rocket::Route; use reqwest::{header::HeaderMap, Client, Response}; +use rocket::http::{Cookie}; + use regex::Regex; use soup::prelude::*; @@ -147,7 +149,7 @@ struct IconList { /// favicon_location1 = get_icon_url("github.com"); /// favicon_location2 = get_icon_url("gitlab.com"); /// ``` -fn get_icon_url(domain: &str) -> Result { +fn get_icon_url(domain: &str) -> Result<(String, String), Error> { // Default URL with secure and insecure schemes let ssldomain = format!("https://{}", domain); let httpdomain = format!("http://{}", domain); @@ -155,10 +157,20 @@ fn get_icon_url(domain: &str) -> Result { // Create the iconlist let mut iconlist: Vec = Vec::new(); + // Create the cookie_str to fill it all the cookies from the response + // These cookies can be used to request/download the favicon image. + // Some sites have extra security in place with for example XSRF Tokens. + let mut cookie_str = String::new(); + let resp = get_page(&ssldomain).or_else(|_| get_page(&httpdomain)); if let Ok(content) = resp { // Extract the URL from the respose in case redirects occured (like @ gitlab.com) let url = content.url().clone(); + let raw_cookies = content.headers().get_all("set-cookie"); + raw_cookies.iter().for_each(|raw_cookie| { + let cookie = Cookie::parse(raw_cookie.to_str().unwrap_or_default()).unwrap(); + cookie_str.push_str(&format!("{}={}; ", cookie.name(), cookie.value())); + }); // Add the default favicon.ico to the list with the domain the content responded from. iconlist.push(IconList { priority: 35, href: url.join("/favicon.ico").unwrap().into_string() }); @@ -188,11 +200,16 @@ fn get_icon_url(domain: &str) -> Result { iconlist.sort_by_key(|x| x.priority); // There always is an icon in the list, so no need to check if it exists, and just return the first one - Ok(iconlist.remove(0).href) + Ok((iconlist.remove(0).href, cookie_str)) } fn get_page(url: &str) -> Result { - CLIENT.get(url).send()?.error_for_status().map_err(Into::into) + //CLIENT.get(url).send()?.error_for_status().map_err(Into::into) + get_page_with_cookies(url, "") +} + +fn get_page_with_cookies(url: &str, cookie_str: &str) -> Result { + CLIENT.get(url).header("cookie", cookie_str).send()?.error_for_status().map_err(Into::into) } /// Returns a Integer with the priority of the type of the icon which to prefer. @@ -244,10 +261,10 @@ fn get_icon_priority(href: &str, sizes: &str) -> u8 { } fn download_icon(domain: &str) -> Result, Error> { - let url = get_icon_url(&domain)?; + let (url, cookie_str) = get_icon_url(&domain)?; info!("Downloading icon for {} via {}...", domain, url); - let mut res = get_page(&url)?; + let mut res = get_page_with_cookies(&url, &cookie_str)?; let mut buffer = Vec::new(); res.copy_to(&mut buffer)?;