feat(services/january): website embed generation

crates/services/january/Cargo.toml

@@ -6,15 +6,21 @@ edition = "2021"
 [dependencies]
 # Utility
 mime = "0.3.17"
+regex = "1.11.0"
 tempfile = "3.13.0"
 lazy_static = "1.5.0"
 moka = { version = "0.12.8", features = ["future"] }
 
+# Web scraping
+scraper = "0.20.0"
+encoding_rs = "0.8.34"
+
 # Serialisation
 serde = { version = "1.0", features = ["derive"] }
 serde_json = "1.0.68"
 
 # Async runtime
+async-recursion = "1.1.1"
 tokio = { version = "1.0", features = ["full"] }
 
 # Web requests

crates/services/january/src/api.rs

@@ -1,15 +1,9 @@
 use axum::{extract::Query, response::IntoResponse, routing::get, Json, Router};
 use reqwest::header;
-use revolt_models::v0::Embed;
 use revolt_result::Result;
 use serde::{Deserialize, Serialize};
 use utoipa::ToSchema;
 
-use axum_extra::{
-    headers::{authorization::Bearer, Authorization},
-    TypedHeader,
-};
-
 use crate::requests::Request;
 
 pub static CACHE_CONTROL: &str = "public, max-age=600, immutable";
@@ -92,6 +86,6 @@ async fn proxy(Query(UrlQuery { url }): Query<UrlQuery>) -> Result<impl IntoResp
 async fn embed(
     Query(UrlQuery { url }): Query<UrlQuery>,
     // TypedHeader(Authorization(_bearer)): TypedHeader<Authorization<Bearer>>,
-) -> Result<Json<Embed>> {
-    Request::generate_embed(&url).await.map(Json)
+) -> Result<impl IntoResponse> {
+    Request::generate_embed(url).await.map(Json)
 }

crates/services/january/src/main.rs

@@ -11,6 +11,7 @@ use utoipa_scalar::{Scalar, Servable as ScalarServable};
 
 mod api;
 pub mod requests;
+pub mod website_embed;
 
 #[tokio::main]
 async fn main() -> Result<(), std::io::Error> {

crates/services/january/src/requests.rs

@@ -1,19 +1,23 @@
+use encoding_rs::{Encoding, UTF_8_INIT};
+use lazy_static::lazy_static;
+use mime::Mime;
+use regex::Regex;
+use reqwest::{
+    header::{self, CONTENT_TYPE},
+    redirect, Client, Response,
+};
+use revolt_config::report_internal_error;
+use revolt_files::{create_thumbnail, decode_image, image_size_vec, is_valid_image, video_size};
+use revolt_models::v0::{Embed, Image, ImageSize, Video};
+use revolt_result::{create_error, Result};
 use std::{
     io::{Cursor, Write},
     time::Duration,
 };
 
-use lazy_static::lazy_static;
-use mime::Mime;
-use reqwest::{header::CONTENT_TYPE, redirect, Client, Response};
-use revolt_config::report_internal_error;
-use revolt_files::{create_thumbnail, decode_image, image_size_vec, is_valid_image, video_size};
-use revolt_models::v0::{Embed, Image, ImageSize, Video};
-use revolt_result::{create_error, Result};
-
 lazy_static! {
+    /// Request client
     static ref CLIENT: Client = reqwest::Client::builder()
-        .user_agent("Mozilla/5.0 (compatible; January/2.0; +https://github.com/revoltchat/backend)")
         .timeout(Duration::from_secs(10)) // TODO config
         .connect_timeout(Duration::from_secs(5)) // TODO config
         .redirect(redirect::Policy::custom(|attempt| {
@@ -28,6 +32,9 @@ lazy_static! {
         .build()
         .expect("reqwest Client");
 
+    /// Spoof User Agent as Discord
+    static ref RE_USER_AGENT_SPOOFING_AS_DISCORD: Regex = Regex::new("^(?:(?:https?:)?//)?(?:(?:vx|fx)?twitter|(?:fixv|fixup)?x).com").expect("valid regex");
+
     /// Cache for proxy results
     static ref PROXY_CACHE: moka::future::Cache<String, Result<(String, Vec<u8>)>> = moka::future::Cache::builder()
         .max_capacity(10_000) // TODO config
@@ -35,7 +42,7 @@ lazy_static! {
         .build();
 
     /// Cache for embed results
-    static ref EMBED_CACHE: moka::future::Cache<String, Result<Embed>> = moka::future::Cache::builder()
+    static ref EMBED_CACHE: moka::future::Cache<String, Embed> = moka::future::Cache::builder()
         .max_capacity(1_000) // TODO config
         .time_to_live(Duration::from_secs(60)) // TODO config
         .build();
@@ -98,57 +105,126 @@ impl Request {
         }
     }
 
-    /// Generate embed for a given URL
-    pub async fn generate_embed(url: &str) -> Result<Embed> {
+    /// Fetch metadata for an image
+    pub async fn fetch_image_metadata(
+        url: &str,
+        request: Option<Request>,
+    ) -> Result<Option<Image>> {
         if let Some(hit) = EMBED_CACHE.get(url).await {
-            hit
-        } else {
-            let Request { response, mime } = Request::new(url).await?;
-
-            match (mime.type_(), mime.subtype()) {
-                (_, mime::HTML) => {
-                    // let mut metadata = Metadata::from(resp, url).await?;
-                    // metadata.resolve_external().await;
-
-                    // if metadata.is_none() {
-                    //     return Ok(Embed::None);
-                    // }
-
-                    // Ok(Embed::Website(metadata))
-                    todo!()
-                }
-                (mime::IMAGE, _) => {
-                    if let Some((width, height)) =
-                        image_size_vec(&report_internal_error!(response.bytes().await)?)
-                    {
-                        Ok(Embed::Image(Image {
-                            url: url.to_owned(),
-                            width,
-                            height,
-                            size: ImageSize::Large,
-                        }))
-                    } else {
-                        Ok(Embed::None)
-                    }
-                }
-                (mime::VIDEO, _) => {
-                    let mut file = report_internal_error!(tempfile::NamedTempFile::new())?;
-                    report_internal_error!(
-                        file.write_all(&report_internal_error!(response.bytes().await)?)
-                    )?;
-
-                    if let Some((width, height)) = video_size(&file) {
-                        Ok(Embed::Video(Video {
-                            url: url.to_owned(),
-                            width: width as usize,
-                            height: height as usize,
-                        }))
-                    } else {
-                        Ok(Embed::None)
-                    }
-                }
-                _ => Ok(Embed::None),
+            match hit {
+                Embed::Image(img) => Ok(Some(img)),
+                _ => Ok(None),
             }
+        } else {
+            let response = if let Some(Request { response, .. }) = request {
+                response
+            } else {
+                let Request { response, mime } = Request::new(url).await?;
+                if matches!(mime.type_(), mime::IMAGE) {
+                    response
+                } else {
+                    return Err(create_error!(LabelMe));
+                }
+            };
+
+            if let Some((width, height)) =
+                image_size_vec(&report_internal_error!(response.bytes().await)?)
+            {
+                Ok(Some(Image {
+                    url: url.to_owned(),
+                    width,
+                    height,
+                    size: ImageSize::Large,
+                }))
+            } else {
+                Ok(None)
+            }
+        }
+    }
+
+    /// Fetch metadata for an video
+    pub async fn fetch_video_metadata(
+        url: &str,
+        request: Option<Request>,
+    ) -> Result<Option<Video>> {
+        if let Some(hit) = EMBED_CACHE.get(url).await {
+            match hit {
+                Embed::Video(vid) => Ok(Some(vid)),
+                _ => Ok(None),
+            }
+        } else {
+            let response = if let Some(Request { response, .. }) = request {
+                response
+            } else {
+                let Request { response, mime } = Request::new(&url).await?;
+                if matches!(mime.type_(), mime::VIDEO) {
+                    response
+                } else {
+                    return Err(create_error!(LabelMe));
+                }
+            };
+
+            let mut file = report_internal_error!(tempfile::NamedTempFile::new())?;
+            report_internal_error!(
+                file.write_all(&report_internal_error!(response.bytes().await)?)
+            )?;
+
+            if let Some((width, height)) = video_size(&file) {
+                Ok(Some(Video {
+                    url: url.to_owned(),
+                    width: width as usize,
+                    height: height as usize,
+                }))
+            } else {
+                Ok(None)
+            }
+        }
+    }
+
+    /// Generate embed for a given URL
+    pub async fn generate_embed(url: String) -> Result<Embed> {
+        if let Some(hit) = EMBED_CACHE.get(&url).await {
+            Ok(hit)
+        } else {
+            let request = Request::new(&url).await?;
+            let embed = match (request.mime.type_(), request.mime.subtype()) {
+                (_, mime::HTML) => {
+                    let content_type = request
+                        .response
+                        .headers()
+                        .get(header::CONTENT_TYPE)
+                        .and_then(|value| value.to_str().ok())
+                        .and_then(|value| value.parse::<Mime>().ok());
+
+                    let encoding_name = content_type
+                        .as_ref()
+                        .and_then(|mime| mime.get_param("charset").map(|charset| charset.as_str()))
+                        .unwrap_or("utf-8");
+
+                    let encoding =
+                        Encoding::for_label(encoding_name.as_bytes()).unwrap_or(&UTF_8_INIT);
+
+                    let bytes = report_internal_error!(request.response.bytes().await)?;
+                    let (text, _, _) = encoding.decode(&bytes);
+
+                    crate::website_embed::create_website_embed(&url, &text)
+                        .await
+                        .map(Embed::Website)
+                        .unwrap_or_default()
+                }
+                (mime::IMAGE, _) => Request::fetch_image_metadata(&url, Some(request))
+                    .await
+                    .map(|res| res.map(Embed::Image).unwrap_or_default())
+                    .unwrap_or_default(),
+                (mime::VIDEO, _) => Request::fetch_video_metadata(&url, Some(request))
+                    .await
+                    .map(|res| res.map(Embed::Video).unwrap_or_default())
+                    .unwrap_or_default(),
+                _ => Embed::None,
+            };
+
+            EMBED_CACHE.insert(url.to_owned(), embed.clone()).await;
+            Ok(embed)
         }
     }
 
@@ -156,6 +232,14 @@ impl Request {
     pub async fn new(url: &str) -> Result<Request> {
         let response = CLIENT
             .get(url)
+            .header(
+                "User-Agent",
+                if RE_USER_AGENT_SPOOFING_AS_DISCORD.is_match(url) {
+                    "Mozilla/5.0 (compatible; Discordbot/2.0; +https://discordapp.com)"
+                } else {
+                    "Mozilla/5.0 (compatible; January/2.0; +https://github.com/revoltchat/backend)"
+                },
+            )
             .send()
             .await
             .map_err(|_| create_error!(ProxyError))?;

crates/services/january/src/website_embed.rs

@@ -0,0 +1,304 @@
+use std::collections::HashMap;
+
+use lazy_static::lazy_static;
+use regex::Regex;
+use revolt_models::v0::{
+    BandcampType, Image, ImageSize, LightspeedType, Special, TwitchType, Video, WebsiteMetadata,
+};
+use scraper::{Html, Selector};
+
+/// Create website metadata from URL and document
+pub async fn create_website_embed(original_url: &str, document: &str) -> Option<WebsiteMetadata> {
+    let document = Html::parse_document(document);
+
+    // create selectors
+    let meta_selector = Selector::parse("meta").ok()?;
+    let link_selector = Selector::parse("link").ok()?;
+
+    // extract meta tags
+    let mut meta = HashMap::new();
+    for el in document.select(&meta_selector) {
+        let node = el.value();
+
+        if let (Some(property), Some(content)) = (
+            node.attr("property").or_else(|| node.attr("name")),
+            node.attr("content"),
+        ) {
+            meta.insert(property.to_string(), content.to_string());
+        }
+    }
+
+    // extract rel links
+    let mut link = HashMap::new();
+    for el in document.select(&link_selector) {
+        let node = el.value();
+
+        if let (Some(property), Some(content)) = (node.attr("rel"), node.attr("href")) {
+            link.insert(property.to_string(), content.to_string());
+        }
+    }
+
+    // build metadata
+    let mut metadata = WebsiteMetadata {
+        title: meta
+            .remove("og:title")
+            .or_else(|| meta.remove("twitter:title"))
+            .or_else(|| meta.remove("title"))
+            .map(|s| s.trim().to_owned()),
+        description: meta
+            .remove("og:description")
+            .or_else(|| meta.remove("twitter:description"))
+            .or_else(|| meta.remove("description"))
+            .map(|s| s.trim().to_owned()),
+        image: meta
+            .remove("og:image")
+            .or_else(|| meta.remove("og:image:secure_url"))
+            .or_else(|| meta.remove("twitter:image"))
+            .or_else(|| meta.remove("twitter:image:src"))
+            .map(|s| s.trim().to_owned())
+            .map(|mut url| {
+                // If relative URL, prepend root URL. Also if root URL ends with a slash, remove it.
+                if let Some(ch) = url.chars().next() {
+                    if ch == '/' {
+                        url = format!("{}{}", &original_url.trim_end_matches('/'), url);
+                    }
+                }
+
+                let mut size = ImageSize::Preview;
+                if let Some(card) = meta.remove("twitter:card") {
+                    if &card == "summary_large_image" {
+                        size = ImageSize::Large;
+                    }
+                }
+
+                Image {
+                    url: url.to_owned(),
+                    width: meta
+                        .remove("og:image:width")
+                        .unwrap_or_default()
+                        .parse()
+                        .unwrap_or(0),
+                    height: meta
+                        .remove("og:image:height")
+                        .unwrap_or_default()
+                        .parse()
+                        .unwrap_or(0),
+                    size,
+                }
+            }),
+        video: meta
+            .remove("og:video")
+            .or_else(|| meta.remove("og:video:url"))
+            .or_else(|| meta.remove("og:video:secure_url"))
+            .map(|s| s.trim().to_owned())
+            .map(|mut url| {
+                // If relative URL, prepend root URL. Also if root URL ends with a slash, remove it.
+                if let Some(ch) = url.chars().next() {
+                    if ch == '/' {
+                        url = format!("{}{}", &original_url.trim_end_matches('/'), url);
+                    }
+                }
+
+                Video {
+                    url: url.to_owned(),
+                    width: meta
+                        .remove("og:video:width")
+                        .unwrap_or_default()
+                        .parse()
+                        .unwrap_or(0),
+                    height: meta
+                        .remove("og:video:height")
+                        .unwrap_or_default()
+                        .parse()
+                        .unwrap_or(0),
+                }
+            }),
+        icon_url: link
+            .remove("apple-touch-icon")
+            .or_else(|| link.remove("icon"))
+            .map(|s| s.trim().to_owned())
+            .map(|mut v| {
+                // If relative URL, prepend root URL.
+                if let Some(ch) = v.chars().next() {
+                    if ch == '/' {
+                        v = format!("{}{}", &original_url.trim_end_matches('/'), v);
+                    }
+                }
+
+                v
+            }),
+        colour: meta.remove("theme-color").map(|s| s.trim().to_owned()),
+        site_name: meta.remove("og:site_name").map(|s| s.trim().to_owned()),
+        url: meta
+            .remove("og:url")
+            .or_else(|| Some(original_url.to_owned())),
+        original_url: Some(original_url.to_owned()),
+        special: None,
+    };
+
+    // populate extra metadata for popular websites
+    populate_special(original_url.to_owned(), &mut metadata);
+
+    // TODO: these do not work because compiler is tripping:
+
+    // fetch video size if missing
+    // if let Some(Video { width, height, url }) = &metadata.video {
+    //     if width == &0 || height == &0 {
+    //         metadata.video = match crate::requests::Request::fetch_video_metadata(url, None).await {
+    //             Ok(Some(video)) => Some(video),
+    //             _ => None,
+    //         }
+    //     }
+    // }
+
+    // fetch image size if missing
+    // if let Some(Image {
+    //     width, height, url, ..
+    // }) = &metadata.image
+    // {
+    //     if width == &0 || height == &0 {
+    //         metadata.image = match crate::requests::Request::fetch_image_metadata(url, None).await {
+    //             Ok(Some(image)) => Some(image),
+    //             _ => None,
+    //         }
+    //     }
+    // }
+
+    // truncate data
+    metadata.truncate();
+
+    // if it's empty, don't return anything
+    if metadata.is_empty() {
+        None
+    } else {
+        Some(metadata)
+    }
+}
+
+pub fn populate_special(original_url: String, metadata: &mut WebsiteMetadata) {
+    lazy_static! {
+        static ref RE_YOUTUBE: Regex = Regex::new("^(?:(?:https?:)?//)?(?:(?:www|m)\\.)?(?:(?:youtube\\.com|youtu.be))(?:/(?:[\\w\\-]+\\?v=|embed/|v/)?)([\\w\\-]+)(?:\\S+)?$").unwrap();
+
+        static ref RE_LIGHTSPEED: Regex = Regex::new("^(?:https?://)?(?:[\\w]+\\.)?lightspeed\\.tv/([a-z0-9_]{4,25})").unwrap();
+
+        static ref RE_TWITCH: Regex = Regex::new("^(?:https?://)?(?:www\\.|go\\.)?twitch\\.tv/([a-z0-9_]+)($|\\?)").unwrap();
+        static ref RE_TWITCH_VOD: Regex = Regex::new("^(?:https?://)?(?:www\\.|go\\.)?twitch\\.tv/videos/([0-9]+)($|\\?)").unwrap();
+        static ref RE_TWITCH_CLIP: Regex = Regex::new("^(?:https?://)?(?:www\\.|go\\.)?twitch\\.tv/(?:[a-z0-9_]+)/clip/([A-z0-9_-]+)($|\\?)").unwrap();
+
+        static ref RE_SPOTIFY: Regex = Regex::new("^(?:https?://)?open.spotify.com/(track|user|artist|album|playlist)/([A-z0-9]+)").unwrap();
+        static ref RE_SOUNDCLOUD: Regex = Regex::new("^(?:https?://)?soundcloud.com/([a-zA-Z0-9-]+)/([A-z0-9-]+)").unwrap();
+        static ref RE_BANDCAMP: Regex = Regex::new("^(?:https?://)?(?:[A-z0-9_-]+).bandcamp.com/(track|album)/([A-z0-9_-]+)").unwrap();
+        static ref RE_APPLE_MUSIC: Regex = Regex::new("^(?:https?://)?music\\.apple\\.com/(?:[a-z]{2}/)?album/(?:[a-zA-Z0-9-]+)/(\\d+)(?:\\?i=(\\d+))?").unwrap();
+
+        static ref RE_STREAMABLE: Regex = Regex::new("^(?:https?://)?(?:www\\.)?streamable\\.com/([\\w\\d-]+)").unwrap();
+
+        static ref RE_GIF: Regex = Regex::new("^(?:https?://)?(www\\.)?(gifbox\\.me/view|yiffbox\\.me/view|tenor\\.com/view|giphy\\.com/gifs|gfycat\\.com|redgifs\\.com/watch)/[\\w\\d-]+").unwrap();
+    }
+
+    let url = metadata
+        .url
+        .as_ref()
+        .or(metadata.original_url.as_ref())
+        .unwrap_or(&original_url)
+        .as_str();
+
+    metadata.special = if let Some(captures) = RE_STREAMABLE.captures_iter(url).next() {
+        Some(Special::Streamable {
+            id: captures[1].to_string(),
+        })
+    } else if let Some(captures) = RE_YOUTUBE.captures_iter(url).next() {
+        lazy_static! {
+            static ref RE_TIMESTAMP: Regex = Regex::new("(?:\\?|&)(?:t|start)=([\\w]+)").unwrap();
+        }
+
+        if let Some(video) = &metadata.video {
+            if let Some(timestamp_captures) = RE_TIMESTAMP.captures_iter(&video.url).next() {
+                Some(Special::YouTube {
+                    id: captures[1].to_string(),
+                    timestamp: Some(timestamp_captures[1].to_string()),
+                })
+            } else {
+                Some(Special::YouTube {
+                    id: captures[1].to_string(),
+                    timestamp: None,
+                })
+            }
+        } else {
+            None
+        }
+    } else if let Some(captures) = RE_LIGHTSPEED.captures_iter(url).next() {
+        Some(Special::Lightspeed {
+            id: captures[1].to_string(),
+            content_type: LightspeedType::Channel,
+        })
+    } else if let Some(captures) = RE_TWITCH.captures_iter(url).next() {
+        Some(Special::Twitch {
+            id: captures[1].to_string(),
+            content_type: TwitchType::Channel,
+        })
+    } else if let Some(captures) = RE_TWITCH_VOD.captures_iter(url).next() {
+        Some(Special::Twitch {
+            id: captures[1].to_string(),
+            content_type: TwitchType::Video,
+        })
+    } else if let Some(captures) = RE_TWITCH_CLIP.captures_iter(url).next() {
+        Some(Special::Twitch {
+            id: captures[1].to_string(),
+            content_type: TwitchType::Clip,
+        })
+    } else if let Some(captures) = RE_SPOTIFY.captures_iter(url).next() {
+        Some(Special::Spotify {
+            content_type: captures[1].to_string(),
+            id: captures[2].to_string(),
+        })
+    } else if RE_SOUNDCLOUD.is_match(url) {
+        Some(Special::Soundcloud)
+    } else if RE_BANDCAMP.is_match(url) {
+        lazy_static! {
+            static ref RE_TRACK: Regex = Regex::new("track=(\\d+)").unwrap();
+            static ref RE_ALBUM: Regex = Regex::new("album=(\\d+)").unwrap();
+        }
+
+        if let Some(video) = &metadata.video {
+            if let Some(captures) = RE_TRACK.captures_iter(&video.url).next() {
+                Some(Special::Bandcamp {
+                    content_type: BandcampType::Track,
+                    id: captures[1].to_string(),
+                })
+            } else {
+                RE_ALBUM
+                    .captures_iter(&video.url)
+                    .next()
+                    .map(|captures| Special::Bandcamp {
+                        content_type: BandcampType::Album,
+                        id: captures[1].to_string(),
+                    })
+            }
+        } else {
+            None
+        }
+    } else if RE_GIF.is_match(url) {
+        Some(Special::GIF)
+    } else {
+        RE_APPLE_MUSIC
+            .captures_iter(url)
+            .next()
+            .map(|captures| Special::AppleMusic {
+                album_id: captures[1].to_string(),
+                track_id: captures.get(2).map(|m| m.as_str().to_string()),
+            })
+    };
+
+    // add colours for popular websites
+    if let Some(special) = &metadata.special {
+        match special {
+            Special::YouTube { .. } => metadata.colour = Some("#FF424F".to_string()),
+            Special::Twitch { .. } => metadata.colour = Some("#7B68EE".to_string()),
+            Special::Lightspeed { .. } => metadata.colour = Some("#7445D9".to_string()),
+            Special::Spotify { .. } => metadata.colour = Some("#1ABC9C".to_string()),
+            Special::Soundcloud { .. } => metadata.colour = Some("#FF7F50".to_string()),
+            Special::AppleMusic { .. } => metadata.colour = Some("#FA233B".to_string()),
+            _ => {}
+        }
+    }
+}