Rewrite to rust

Co-authored-by: Kurbanov Bulat <kurbanovbul@gmail.com>

src/config.rs

@@ -0,0 +1,35 @@
+use once_cell::sync::Lazy;
+
+
+fn get_env(env: &'static str) -> String {
+    std::env::var(env).unwrap_or_else(|_| panic!("Cannot get the {} env variable", env))
+}
+
+pub struct Config {
+    pub api_key: String,
+
+    pub telegram_api_url: String,
+    pub telegram_chat_id: i64,
+    pub telegram_temp_chat_id: i64,
+    pub bot_tokens: Vec<String>,
+
+    pub sentry_dsn: String,
+}
+
+
+impl Config {
+    pub fn load() -> Config {
+        Config {
+            api_key: get_env("API_KEY"),
+
+            telegram_api_url: get_env("API_URL"),
+            telegram_chat_id: get_env("TELEGRAM_CHAT_ID").parse().unwrap(),
+            telegram_temp_chat_id: get_env("TELEGRAM_TEMP_CHAT_ID").parse().unwrap(),
+            
+            bot_tokens: serde_json::from_str(&get_env("BOT_TOKENS")).unwrap(),
+            sentry_dsn: get_env("SENTRY_DSN"),
+        }
+    }
+}
+
+pub static CONFIG: Lazy<Config> = Lazy::new(Config::load);

src/core/bot.rs

@@ -0,0 +1,32 @@
+use std::sync::{
+    atomic::{AtomicUsize, Ordering},
+    Arc,
+};
+
+use once_cell::sync::Lazy;
+use teloxide::Bot;
+
+use crate::config::{self, CONFIG};
+
+pub struct RoundRobinBot {
+    bot_tokens: Arc<Vec<String>>,
+    current_index: AtomicUsize,
+}
+
+impl RoundRobinBot {
+    pub fn new(bot_tokens: Vec<String>) -> Self {
+        RoundRobinBot {
+            bot_tokens: Arc::new(bot_tokens),
+            current_index: AtomicUsize::new(0),
+        }
+    }
+
+    pub fn get_bot(&self) -> Bot {
+        let index = self.current_index.fetch_add(1, Ordering::Relaxed) % self.bot_tokens.len();
+        Bot::new(self.bot_tokens[index].clone())
+            .set_api_url(reqwest::Url::parse(CONFIG.telegram_api_url.as_str()).unwrap())
+    }
+}
+
+pub static ROUND_ROBIN_BOT: Lazy<RoundRobinBot> =
+    Lazy::new(|| RoundRobinBot::new(config::CONFIG.bot_tokens.clone()));

src/core/file_utils.rs

@@ -0,0 +1,104 @@
+use std::pin::Pin;
+
+use axum::body::Bytes;
+use futures::TryStreamExt;
+use serde::Serialize;
+use teloxide::{
+    net::Download,
+    requests::Requester,
+    types::{ChatId, InputFile, MessageId},
+    Bot,
+};
+use tokio::io::AsyncRead;
+use tokio_util::compat::FuturesAsyncReadCompatExt;
+
+use crate::config::CONFIG;
+use super::bot::ROUND_ROBIN_BOT;
+
+
+#[derive(Serialize)]
+pub struct UploadedFile {
+    pub backend: String,
+    pub data: MessageInfo,
+}
+
+#[derive(Serialize)]
+pub struct MessageInfo {
+    pub chat_id: i64,
+    pub message_id: i32,
+}
+
+
+pub async fn upload_file(
+    file: Bytes,
+    filename: String,
+    caption: Option<String>,
+) -> Result<UploadedFile, String> {
+    let bot = ROUND_ROBIN_BOT.get_bot();
+    let document = InputFile::memory(file).file_name(filename);
+
+    let mut request = bot.send_document(ChatId(CONFIG.telegram_chat_id), document);
+    request.caption = caption;
+
+    let result = request.await;
+
+    match result {
+        Ok(message) => Ok(UploadedFile {
+            backend: "bot".to_string(),
+            data: MessageInfo {
+                chat_id: message.chat.id.0,
+                message_id: message.id.0,
+            },
+        }),
+        Err(err) => Err(err.to_string()),
+    }
+}
+
+pub async fn download_file(chat_id: i64, message_id: i32) -> Option<BotDownloader> {
+    let bot = ROUND_ROBIN_BOT.get_bot();
+
+    let result = bot
+        .forward_message(
+            ChatId(CONFIG.telegram_temp_chat_id),
+            ChatId(chat_id),
+            MessageId(message_id),
+        )
+        .await;
+
+    match result {
+        Ok(message) => {
+            if message.document() == None {
+                return Option::None;
+            }
+
+            let file_id = message.document().unwrap().file.id.clone();
+            let path = bot.get_file(file_id.clone()).await.unwrap().path;
+
+            return Some(BotDownloader::new(bot, path));
+        }
+        Err(_) => None,
+    }
+}
+
+
+pub struct BotDownloader {
+    bot: Bot,
+    file_path: String,
+}
+
+impl BotDownloader {
+    pub fn new(bot: Bot, file_path: String) -> Self {
+        Self { bot, file_path }
+    }
+
+    pub fn get_async_read(self) -> Pin<Box<dyn AsyncRead + Send>> {
+        let stream = self.bot.download_file_stream(&self.file_path);
+
+        Box::pin(
+            stream
+                .map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e))
+                .into_async_read()
+                .compat()
+        )
+    }
+}

src/core/mod.rs

@@ -0,0 +1,3 @@
+pub mod bot;
+pub mod views;
+pub mod file_utils;

src/core/views.rs

@@ -0,0 +1,103 @@
+use axum::{
+    body::Bytes,
+    extract::{DefaultBodyLimit, Path},
+    http::{self, Request, StatusCode},
+    middleware::{self, Next},
+    response::{IntoResponse, Response},
+    routing::{get, post},
+    Router,
+};
+use axum_prometheus::PrometheusMetricLayer;
+use axum_typed_multipart::{TryFromMultipart, TypedMultipart};
+use tokio_util::io::ReaderStream;
+use tower_http::trace::{self, TraceLayer};
+use tracing::Level;
+
+use crate::config::CONFIG;
+
+use super::file_utils::{download_file, upload_file};
+
+const BODY_LIMIT: usize = 4 * (2 << 30); // bytes: 4GB
+
+
+async fn auth(req: Request<axum::body::Body>, next: Next) -> Result<Response, StatusCode> {
+    let auth_header = req
+        .headers()
+        .get(http::header::AUTHORIZATION)
+        .and_then(|header| header.to_str().ok());
+
+    let auth_header = if let Some(auth_header) = auth_header {
+        auth_header
+    } else {
+        return Err(StatusCode::UNAUTHORIZED);
+    };
+
+    if auth_header != CONFIG.api_key {
+        return Err(StatusCode::UNAUTHORIZED);
+    }
+
+    Ok(next.run(req).await)
+}
+
+
+pub async fn get_router() -> Router {
+    let (prometheus_layer, metric_handle) = PrometheusMetricLayer::pair();
+
+    let app_router = Router::new()
+        .route("/upload", post(upload))
+        .route("/download_by_message/:chat_id/:message_id", get(download))
+        .layer(DefaultBodyLimit::max(BODY_LIMIT))
+        .layer(middleware::from_fn(auth))
+        .layer(prometheus_layer);
+
+    let metric_router =
+        Router::new().route("/metrics", get(|| async move { metric_handle.render() }));
+
+    Router::new()
+        .nest("/api/v1/files", app_router)
+        .nest("/", metric_router)
+        .layer(
+            TraceLayer::new_for_http()
+                .make_span_with(trace::DefaultMakeSpan::new().level(Level::INFO))
+                .on_response(trace::DefaultOnResponse::new().level(Level::INFO)),
+        )
+}
+
+
+#[derive(TryFromMultipart)]
+pub struct UploadFileRequest {
+    #[form_data(limit = "unlimited")]
+    file: Bytes,
+    filename: String,
+    caption: Option<String>,
+}
+
+
+async fn upload(data: TypedMultipart<UploadFileRequest>) -> impl IntoResponse {
+    let result = match upload_file(
+        data.file.clone(),
+        data.filename.to_string(),
+        data.caption.clone(),
+    )
+    .await
+    {
+        Ok(file) => serde_json::to_string(&file),
+        Err(err) => Ok(err),
+    };
+
+    result.unwrap()
+}
+
+
+async fn download(Path(chat_id): Path<i64>, Path(message_id): Path<i32>) -> impl IntoResponse {
+    let downloader = download_file(chat_id, message_id).await; 
+
+    let data = match downloader {
+        Some(v) => v.get_async_read(),
+        None => return StatusCode::NOT_FOUND.into_response()
+    };
+
+    let reader = ReaderStream::new(data);
+    
+    axum::body::Body::from_stream(reader).into_response()
+}

src/main.rs

@@ -0,0 +1,36 @@
+mod config;
+mod core;
+
+use std::{net::SocketAddr, str::FromStr};
+use sentry::{integrations::debug_images::DebugImagesIntegration, types::Dsn, ClientOptions};
+
+use crate::core::views::get_router;
+
+
+#[tokio::main]
+async fn main() {
+    dotenv::dotenv().ok();
+
+    let options = ClientOptions {
+        dsn: Some(Dsn::from_str(&config::CONFIG.sentry_dsn).unwrap()),
+        default_integrations: false,
+        ..Default::default()
+    }
+    .add_integration(DebugImagesIntegration::new());
+
+    let _guard = sentry::init(options);
+
+    tracing_subscriber::fmt()
+        .with_target(false)
+        .compact()
+        .init();
+
+    let addr = SocketAddr::from(([0, 0, 0, 0], 8080));
+
+    let app = get_router().await;
+
+    println!("Start webserver...");
+    let listener = tokio::net::TcpListener::bind(&addr).await.unwrap();
+    axum::serve(listener, app).await.unwrap();
+    println!("Webserver shutdown...");
+}