From 5016e30cf298834f45af7abfd4811d2d70868e34 Mon Sep 17 00:00:00 2001
From: BlackDex <black.dex@gmail.com>
Date: Thu, 31 Jan 2019 15:49:58 +0100
Subject: [PATCH] Added cookies to the icon download request.

Some sites use XSRF Tokens, or other Tokens to verify a subseqense
response. The cookies which are sent during the page request are now
used when downloading the favicon.

A site which uses this is mijn.ing.nl.
---
 src/api/icons.rs | 27 ++++++++++++++++++++++-----
 1 file changed, 22 insertions(+), 5 deletions(-)

diff --git a/src/api/icons.rs b/src/api/icons.rs
index 51cb5041..8ae3c83f 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<String, Error> {
+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<String, Error> {
     // Create the iconlist
     let mut iconlist: Vec<IconList> = 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<String, Error> {
     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<Response, Error> {
-    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<Response, Error> {
+    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<Vec<u8>, 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)?;