Merge pull request #7 from ahdark-services/feat/stickers-export-handl…

…er/webm-convert

Feat: webm convertion for stickers export handler

docker/rust-basic/Dockerfile

@@ -15,9 +15,10 @@ WORKDIR /app
 
 ARG COMPONENT
 
-RUN apt update && \
-    apt install -y openssl libssl-dev ca-certificates && \
-    rm -rf /var/lib/apt/lists/*
+RUN apt update
+RUN apt install -y openssl libssl-dev ca-certificates
+RUN if [ "${COMPONENT}" = "stickers-export-handler" ]; then apt install -y ffmpeg; fi
+RUN rm -rf /var/lib/apt/lists/*
 
 COPY --from=builder /usr/src/pegasus/target/release/${COMPONENT} /app/entry
 

rust-components/stickers-export-handler/Cargo.toml

@@ -20,3 +20,4 @@ teloxide = { workspace = true, features = ["macros"] }
 serde = { workspace = true, features = ["derive"] }
 image = { version = "0.25", features = [] }
 anyhow = "1.0"
+async-tempfile = "0.5"

rust-components/stickers-export-handler/src/convert.rs

@@ -0,0 +1,91 @@
+use std::io::Cursor;
+
+use image::codecs::{png, webp};
+use image::ImageFormat;
+
+pub(crate) fn convert_webp_to_png(input_buffer: Vec<u8>) -> anyhow::Result<Vec<u8>> {
+    let input_image = image::load_from_memory(&input_buffer)
+        .map_err(|err| anyhow::anyhow!("Failed to load image from memory: {}", err))?;
+
+    let mut output_buffer = Vec::new();
+    let mut output_cursor = Cursor::new(&mut output_buffer);
+    input_image
+        .write_to(&mut output_cursor, ImageFormat::Png)
+        .map_err(|err| anyhow::anyhow!("Failed to write image to buffer: {}", err))?;
+
+    Ok(output_buffer)
+}
+
+pub(crate) async fn convert_webm_to_gif(input_buffer: Vec<u8>) -> anyhow::Result<Vec<u8>> {
+    let mut input_file = async_tempfile::TempFile::new_with_name("input.webm")
+        .await
+        .map_err(|err| anyhow::anyhow!("Failed to create input file: {}", err))?;
+
+    let palette_file = async_tempfile::TempFile::new_with_name("palette.png")
+        .await
+        .map_err(|err| anyhow::anyhow!("Failed to create palette file: {}", err))?;
+
+    let mut output_file = async_tempfile::TempFile::new_with_name("output.gif")
+        .await
+        .map_err(|err| anyhow::anyhow!("Failed to create output file: {}", err))?;
+
+    log::debug!(
+        "input_file: {}, palette_file: {}, output_file: {}",
+        input_file.file_path().to_str().unwrap(),
+        palette_file.file_path().to_str().unwrap(),
+        output_file.file_path().to_str().unwrap()
+    );
+
+    tokio::io::AsyncWriteExt::write_all(&mut input_file, &input_buffer)
+        .await
+        .map_err(|err| anyhow::anyhow!("Failed to write input file: {}", err))?;
+
+    // create a palette from the input file
+    let out = tokio::process::Command::new("ffmpeg")
+        .arg("-i")
+        .arg(input_file.file_path())
+        .arg("-vf")
+        .arg("palettegen")
+        .arg(palette_file.file_path())
+        .arg("-y")
+        .output()
+        .await
+        .map_err(|err| anyhow::anyhow!("Failed to generate palette: {}", err))?;
+    if !out.status.success() {
+        return Err(anyhow::anyhow!(
+            "Failed to generate palette: {}",
+            String::from_utf8_lossy(&out.stderr)
+        ));
+    }
+
+    // convert the input file to a gif using the palette
+    let out = tokio::process::Command::new("ffmpeg")
+        .arg("-i")
+        .arg(input_file.file_path())
+        .arg("-i")
+        .arg(palette_file.file_path())
+        .arg("-filter_complex")
+        .arg("paletteuse")
+        .arg(output_file.file_path())
+        .arg("-y")
+        .output()
+        .await
+        .map_err(|err| anyhow::anyhow!("Failed to convert webm to gif: {}", err))?;
+    if !out.status.success() {
+        return Err(anyhow::anyhow!(
+            "Failed to convert webm to gif: {}",
+            String::from_utf8_lossy(&out.stderr)
+        ));
+    }
+
+    let mut output_buffer = Vec::new();
+    tokio::io::AsyncReadExt::read_to_end(&mut output_file, &mut output_buffer)
+        .await
+        .map_err(|err| anyhow::anyhow!("Failed to read output file: {}", err))?;
+
+    if output_buffer.is_empty() {
+        return Err(anyhow::anyhow!("Empty output buffer"));
+    }
+
+    Ok(output_buffer)
+}

rust-components/stickers-export-handler/src/handlers.rs

@@ -1,12 +1,11 @@
-use std::io::Cursor;
-
-use image::ImageError;
-use image::ImageFormat::Png;
+use image::guess_format;
 use teloxide::net::Download;
 use teloxide::prelude::*;
 use teloxide::types::{InputFile, MediaKind, MessageKind};
 use teloxide::utils::command::BotCommands;
 
+use crate::convert::{convert_webm_to_gif, convert_webp_to_png};
+
 #[derive(BotCommands, Clone)]
 #[command(rename_rule = "lowercase")]
 pub(crate) enum Command {
@@ -36,6 +35,12 @@ pub(crate) async fn export_sticker_handler(bot: Bot, message: Message) -> anyhow
         return Ok(());
     };
 
+    let pending_message = bot
+        .send_message(message.chat.id, "Processing...")
+        .reply_to_message_id(message.id)
+        .send()
+        .await?;
+
     let media_sticker = if let MediaKind::Sticker(media_sticker) = &reply_to_message.media_kind {
         media_sticker
     } else {
@@ -44,26 +49,53 @@ pub(crate) async fn export_sticker_handler(bot: Bot, message: Message) -> anyhow
     };
 
     // convert sticker to png
-
     let file = bot.get_file(&media_sticker.sticker.file.id).send().await?;
     let mut buffer = Vec::new();
-    bot.download_file(&file.path, &mut buffer).await?;
+    log::debug!(
+        "Downloading sticker: {}({}), chat id: {}",
+        media_sticker.sticker.file.id,
+        file.path,
+        message.chat.id,
+    );
 
-    let img = image::load_from_memory(&buffer)?;
-    let output_data = tokio::task::spawn_blocking(move || -> Result<Vec<u8>, ImageError> {
-        let mut bytes: Vec<u8> = Vec::new();
-        let mut cursor = Cursor::new(&mut bytes);
-        img.write_to(&mut cursor, Png)?;
-        Ok(bytes)
-    })
-    .await??;
+    match bot.download_file(&file.path, &mut buffer).await {
+        Ok(_) => {}
+        Err(err) => {
+            send_error_message!(
+                bot,
+                message,
+                &format!("Failed to download sticker: {}", err)
+            );
+            return Err(err.into());
+        }
+    };
 
-    // send png
+    let input_file = match guess_format(&buffer) {
+        Ok(_) => match convert_webp_to_png(buffer) {
+            Ok(buf) => InputFile::memory(buf).file_name("sticker.png"),
+            Err(err) => {
+                send_error_message!(bot, message, &format!("Failed to convert sticker: {}", err));
+                return Err(err.into());
+            }
+        },
+        Err(_) => match convert_webm_to_gif(buffer).await {
+            Ok(buf) => InputFile::memory(buf).file_name("sticker.gif"),
+            Err(err) => {
+                send_error_message!(bot, message, &format!("Failed to convert sticker: {}", err));
+                return Err(err.into());
+            }
+        },
+    };
 
-    bot.send_photo(message.chat.id, InputFile::memory(output_data))
+    // send png
+    bot.send_document(message.chat.id, input_file)
         .reply_to_message_id(message.id)
         .send()
         .await?;
 
+    bot.delete_message(pending_message.chat.id, pending_message.id)
+        .send()
+        .await?;
+
     Ok(())
 }

rust-components/stickers-export-handler/src/main.rs

@@ -9,6 +9,7 @@ use pegasus_common::{observability, settings};
 
 use crate::run::run;
 
+mod convert;
 mod handlers;
 mod run;