update + favicon

src/favicon.rs

@@ -0,0 +1,30 @@
+use based::get_pg;
+
+
+pub async fn download_favicon(domain: &str) -> Option<Vec<u8>> {
+    let mut favicon_url = url::Url::parse(&format!("https://{}", domain)).ok()?;
+    favicon_url.set_path("/favicon.ico");
+
+    log::info!("Fetching favicon from: {}", favicon_url);
+
+    let response = reqwest::get(favicon_url).await.ok()?;
+
+    if !response.status().is_success() {
+        return None;
+    }
+
+    let favicon_data = response.bytes().await.ok()?.to_vec();
+
+    Some(favicon_data)
+}
+
+pub async fn download_favicons_for_sites(sites: Vec<String>) {
+    for site in sites {
+        if let Some(fav) = download_favicon(&site).await {
+            sqlx::query("INSERT INTO site_favicon VALUES ($1, $2)")
+            .bind(site)
+            .bind(fav)
+            .execute(get_pg!()).await.unwrap();
+        }
+    }
+}

src/main.rs

@@ -4,23 +4,27 @@ use rocket::routes;
 
 mod archive;
 mod pages;
+mod favicon;
 
 #[rocket::launch]
 async fn launch() -> _ {
     env_logger::init();
 
-    // let pg = get_pg!();
-    // sqlx::migrate!("./migrations").run(pg).await.unwrap();
+    let pg = get_pg!();
+    sqlx::migrate!("./migrations").run(pg).await.unwrap();
 
     let arc = WebsiteArchive::new("./websites");
 
+    favicon::download_favicons_for_sites(arc.domains()).await;
+
     rocket::build()
         .mount(
             "/",
             routes![
                 pages::index,
                 pages::render_website,
-                pages::domain_info_route
+                pages::domain_info_route,
+                pages::favicon_route
             ],
         )
         .manage(arc)

src/pages/mod.rs

@@ -1,11 +1,10 @@
 use std::path::PathBuf;
 
 use based::{
-    page::Shell,
-    request::{respond_html, RequestContext, StringResponse},
+    get_pg, page::Shell, request::{assets::DataResponse, respond_html, RawResponse, RequestContext, StringResponse}
 };
 use maud::{html, PreEscaped};
-use rocket::{get, State};
+use rocket::{get, Data, State};
 
 use crate::archive::{PathEntry, WebsiteArchive};
 
@@ -25,6 +24,19 @@ pub async fn render_page(content: PreEscaped<String>, ctx: RequestContext) -> St
     .await
 }
 
+#[get("/favicon/<domain>")]
+pub async fn favicon_route(domain: &str) -> Option<DataResponse> {
+    let fav: Option<(Vec<u8>,)> = sqlx::query_as("SELECT favicon FROM site_favicon WHERE domain = $1")
+    .bind(domain)
+    .fetch_optional(get_pg!()).await.unwrap();
+
+    if let Some(fav_data) = fav {
+        return Some(DataResponse::new(fav_data.0, "image/x-icon", Some(60 * 60 * 24 * 5)));
+    }
+
+    None
+}
+
 #[get("/")]
 pub async fn index(ctx: RequestContext, arc: &State<WebsiteArchive>) -> StringResponse {
     let websites = arc.domains();