TelegramTwitterMediaBot/crates/xmedia-bot/src/main.rs

use dotenv::dotenv;
use std::time::Duration;
use teloxide::dptree::endpoint;
use teloxide::prelude::*;
use teloxide::stop::StopToken;
use teloxide::types::{ChatId, InputFile, MessageId};
use teloxide::update_listeners::{self, UpdateListener, webhooks};
use tokio::sync::watch;
use x_media::site;

mod config;
mod ctx;
mod db;
mod handlers;
mod link_cache;
mod media_sender;
mod photo;
mod queue;
mod rate_limit;
mod send;
mod state;

use ctx::CONTEXT;
use handlers::{CONFIG, TASK_QUEUE};

/// Docker `stop` / `compose down` delivers SIGTERM, which teloxide's ctrlc
/// handler (SIGINT only) never sees — without this the process would die
/// before the graceful shutdown below (admin notice, queue drain). Stopping
/// the token unwinds the dispatcher exactly like Ctrl+C does.
#[cfg(unix)]
fn spawn_sigterm_handler(stop_token: StopToken) {
    tokio::spawn(async move {
        let mut sigterm = tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())
            .expect("failed to install SIGTERM handler");
        sigterm.recv().await;
        log::info!("SIGTERM received, stopping the dispatcher");
        stop_token.stop();
    });
}

#[cfg(not(unix))]
fn spawn_sigterm_handler(_stop_token: StopToken) {}

/// A leftover temp file must be at least this old before the startup sweep
/// touches it. Orphans come from a *previous* run; anything younger could
/// belong to a second instance sharing the temp directory (a misconfiguration,
/// but one that must not cost it its in-flight download).
const ORPHAN_TEMP_AGE: Duration = Duration::from_secs(3600);

/// Removes this project's own leftover temp entries (`x_media::TEMP_FILE_PREFIX`)
/// from `dir` once they are older than `older_than`. Returns how many were
/// removed. Entries that are not ours, or are too young, or cannot be dated,
/// are left alone: the OS temp directory is shared, and the marker prefix plus
/// the age gate are the only two things that make deleting here safe.
fn sweep_temp_dir(dir: &std::path::Path, older_than: Duration) -> usize {
    let Ok(entries) = std::fs::read_dir(dir) else {
        return 0;
    };
    let cutoff = std::time::SystemTime::now() - older_than;
    let mut removed = 0;
    for entry in entries.flatten() {
        let name = entry.file_name();
        if !name
            .to_string_lossy()
            .starts_with(x_media::TEMP_FILE_PREFIX)
        {
            continue;
        }
        let old_enough = entry
            .metadata()
            .and_then(|meta| meta.modified())
            .is_ok_and(|modified| modified < cutoff);
        if !old_enough {
            continue;
        }
        let path = entry.path();
        let result = if entry.file_type().is_ok_and(|kind| kind.is_dir()) {
            std::fs::remove_dir_all(&path)
        } else {
            std::fs::remove_file(&path)
        };
        match result {
            Ok(()) => removed += 1,
            // Not worth a warning per entry: a file another process removed
            // first (or one we may not delete) is not a problem here.
            Err(e) => log::debug!("could not remove orphaned temp entry {path:?}: {e}"),
        }
    }
    removed
}

#[tokio::main]
async fn main() {
    dotenv().ok();
    // Without RUST_LOG nothing at all was logged (env_logger falls back to
    // `error`), so a deployment that forgot the variable looked like a bot
    // with no logs; and at `debug` the HTTP client's own lines (hyper_util,
    // reqwest) outnumbered the bot's by two to one. The timed builder adds
    // the timestamp the plain `init` omitted, so a line can be compared with
    // a user's report. An explicit RUST_LOG still wins outright — but a blank
    // one (`RUST_LOG=` in `.env`, which is not "unset") must not silence the
    // log the way its absence used to.
    let filter = std::env::var("RUST_LOG")
        .ok()
        .filter(|value| !value.trim().is_empty())
        .unwrap_or_else(|| "info,hyper_util=warn,reqwest=warn".to_string());
    pretty_env_logger::formatted_timed_builder()
        .parse_filters(&filter)
        .init();
    log::info!("Starting bot");

    // Temp media (downloaded files, ugoira/remux dirs) is cleaned up by
    // `TempDir`/`NamedTempFile` on drop — which a killed process never runs.
    // Without this sweep every hard restart left its downloads behind (up to
    // hundreds of MB each) and nothing could tell them apart from a live
    // process's files or from anything else in the OS temp dir. See
    // [`sweep_temp_dir`] for why the age gate makes that safe.
    let orphans = sweep_temp_dir(&std::env::temp_dir(), ORPHAN_TEMP_AGE);
    if orphans > 0 {
        log::info!("swept {orphans} orphaned temp file(s) from a previous run");
    }

    let bot = Bot::from_env();
    // Force the queue workers' shared Bot to initialize now so a missing
    // token fails at startup, not on the first queued task.
    let _ = &*send::BOT;

    // Register the command list with Telegram (client `/` menu).
    if let Err(e) = handlers::register_commands(&bot).await {
        log::warn!("failed to register commands: {e}");
    }

    // The effective tunables, so an operator can see what the process actually
    // resolved (a mistyped DATA_DIR or a forgotten TTL override is otherwise
    // invisible until it bites). The proxy URL is never printed — it may embed
    // credentials — and admin ids are chat identifiers, so they stay at debug.
    let quote_chars = match CONFIG.caption_quote_text_chars {
        0 => "off".to_string(),
        n => format!("{n} chars"),
    };
    log::info!(
        "config: {} admin(s), state {}, edit-message TTL {}s, link cache TTL {}s, caption quote {quote_chars}, proxy={}",
        CONFIG.admin_ids.len(),
        crate::handlers::db_path().display(),
        CONFIG.edit_message_ttl.as_secs(),
        CONFIG.link_cache_ttl.as_secs(),
        if std::env::var("TELOXIDE_PROXY").is_ok() {
            "yes"
        } else {
            "no"
        }
    );
    log::debug!("config: admin ids {:?}", CONFIG.admin_ids);

    // Startup repair, before any worker runs: a queued retry whose media was a
    // local file (ugoira MP4, bsky remux, a downloaded temp file) can never
    // succeed after a restart — the registry that kept those files alive is in
    // memory — so those rows are re-fetched from their post instead of
    // dead-lettering the user's link.
    let repaired = handlers::repair_lost_local_media(&CONTEXT).await;
    if repaired > 0 {
        log::info!("startup repair: re-fetched {repaired} queued task(s)");
    }

    // Queue worker: handles typed tasks, dead-letters failed sends to the
    // task's chat. Both closures use the shared context (the queue requires
    // 'static handlers, and the statics are process-wide anyway).
    TASK_QUEUE
        .start(
            |payload| send::handle_task(&CONTEXT, payload),
            |payload, message| send::dead_letter_notify(&CONTEXT, payload, message),
        )
        .await;
    log::info!("task queue worker started");

    // URL job workers: bounded channel + fixed pool for per-URL work.
    handlers::start_url_workers().await;
    log::info!("url workers started");

    // Site login validation (user request): a failed login notifies the
    // admin and the site disables itself for this process (pixiv).
    let failures = site::validate_all().await;
    if failures.is_empty() {
        log::info!("site logins validated");
    } else {
        for (site_id, message) in &failures {
            log::error!("{site_id} login failed: {message}");
            if let Some(admin) = CONFIG.admin_ids.first() {
                let _ = bot
                    .send_message(ChatId(*admin), format!("{site_id} login failed: {message}"))
                    .await;
            }
        }
    }

    // Background sweep: expires the edit prompts and prunes what has aged out.
    log::info!(
        "edit-expiry sweep: every {}s, ttl {}",
        SWEEP_INTERVAL.as_secs(),
        CONFIG.edit_message_ttl.as_secs()
    );
    let (stop_tx, stop_rx) = watch::channel(false);
    {
        let bot = bot.clone();
        tokio::spawn(async move {
            periodic_sweep(crate::ctx::AppContext::from_statics(&bot), stop_rx).await;
        });
    }

    let handler = dptree::entry()
        .branch(Update::filter_message().branch(endpoint(handlers::message_handler)))
        .branch(Update::filter_inline_query().branch(endpoint(handlers::inline_query_handler)))
        .branch(Update::filter_callback_query().branch(endpoint(handlers::callback_query_handler)));

    let mut dispatcher = Dispatcher::builder(bot.clone(), handler)
        .enable_ctrlc_handler()
        .build();

    if CONFIG.webhook_enabled {
        log::info!("running in webhook mode");
        let url = CONFIG.webhook_url.clone().expect("WEBHOOK_URL is not set");
        // `webhooks::axum` calls set_webhook itself (with the full options,
        // secret token included) — no explicit registration here.
        let listen = CONFIG.webhook_listen.expect("WEBHOOK_LISTEN is not set");
        let port = CONFIG.webhook_port.expect("WEBHOOK_PORT is not set");
        // No secret, no webhook: without one the axum listener accepts any
        // POST, and a forged update can impersonate anyone — admins included.
        let secret = CONFIG
            .webhook_secret_token
            .clone()
            .expect("WEBHOOK_SECRET_TOKEN is not set (required in webhook mode)");
        let mut options = webhooks::Options::new((listen, port).into(), url).secret_token(secret);
        if let Some(cert) = &CONFIG.webhook_cert {
            options = options.certificate(InputFile::file(cert));
        }

        let mut listener = webhooks::axum(bot.clone(), options)
            .await
            .expect("Failed to create webhook listener");
        let stop_token = listener.stop_token();
        spawn_sigterm_handler(stop_token);

        dispatcher
            .dispatch_with_listener(
                listener,
                LoggingErrorHandler::with_custom_text("Error from update listener"),
            )
            .await;
    } else {
        log::info!("running in polling mode");
        // Same listener `dispatch()` builds internally — using
        // `dispatch_with_listener` just exposes its stop token so SIGTERM can
        // unwind the dispatcher before the graceful shutdown below.
        let mut listener = update_listeners::polling_default(bot.clone()).await;
        let stop_token = listener.stop_token();
        spawn_sigterm_handler(stop_token);

        dispatcher
            .dispatch_with_listener(
                listener,
                LoggingErrorHandler::with_custom_text("Error from update listener"),
            )
            .await;
    }

    // Graceful stop (Ctrl+C / SIGTERM): stop the sweep, notify the admin,
    // drain the queue. Bounded: a worker mid-download (30 s timeout) or a
    // long ugoira encode must not hold the shutdown hostage forever.
    log::info!("Stopping bot");
    const SHUTDOWN_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30);
    let shutdown = async {
        let _ = stop_tx.send(true);
        handlers::stop_url_workers().await;
        if let Some(admin) = CONFIG.admin_ids.first() {
            let _ = bot.send_message(ChatId(*admin), "Shutting down...").await;
        }
        TASK_QUEUE.stop().await;
    };
    if tokio::time::timeout(SHUTDOWN_TIMEOUT, shutdown)
        .await
        .is_err()
    {
        log::warn!("graceful shutdown timed out after {SHUTDOWN_TIMEOUT:?}; exiting");
    } else {
        log::info!("Bot stopped");
    }
}

/// How often [`periodic_sweep`] runs.
const SWEEP_INTERVAL: Duration = Duration::from_secs(300);

/// The background sweep: rewrites the expired edit prompts in place, prunes the
/// link cache, the idle rate-limit buckets and the idle inline-query entries,
/// and reports the queue only when it is not empty.
///
/// Takes the shared [`crate::ctx::AppContext`] — the collaborators as one
/// bundle, production assembling it from the statics and tests from tempdir
/// stores — so a test can drive a tick with a paused clock: a sleeping task
/// nothing drives is how the queue's own sweep kept a missing worker wake-up.
async fn periodic_sweep(ctx: crate::ctx::AppContext<'_>, mut stop: watch::Receiver<bool>) {
    loop {
        tokio::select! {
            _ = stop.changed() => break,
            _ = tokio::time::sleep(SWEEP_INTERVAL) => {}
        }
        let removed = ctx
            .chat_store
            .prune_expired(ctx.config.edit_message_ttl)
            .await;
        let pruned = ctx.link_cache.prune(ctx.config.link_cache_ttl).await;
        if pruned > 0 {
            log::info!("link cache: pruned {pruned} expired entr(ies)");
        }
        let idle_limiters = crate::rate_limit::prune_idle();
        if idle_limiters > 0 {
            log::debug!("rate limiter: dropped {idle_limiters} idle bucket(s)");
        }
        // Entries past Telegram's own inline cache window: a repeat is sent to
        // the bot again anyway, so keeping them would suppress a fetch the user
        // is waiting for (and the map grew one entry per user, forever).
        let idle_inline = handlers::prune_idle_states();
        if idle_inline > 0 {
            log::debug!("inline queries: dropped {idle_inline} idle entry(ies)");
        }
        // Only speaks up when the queue is not empty: a healthy bot has nothing
        // to report, and a periodic "0 pending" line is noise that hides the
        // lines that matter.
        if let Some((pending, oldest_run_after)) = ctx.task_queue.pending_backlog().await {
            let overdue = crate::db::now_f64() - oldest_run_after;
            if overdue >= 0.0 {
                log::info!("queue: {pending} pending task(s), oldest {overdue:.0}s overdue");
            } else {
                log::info!(
                    "queue: {pending} pending task(s), oldest retry in {:.0}s",
                    -overdue
                );
            }
        }
        for (chat_id, prompt_message_id) in removed {
            // Rewritten in place, not announced: the sweep is a background
            // timer, and a fresh message would wake the chat up to a full TTL
            // later about a prompt the user already walked away from. The edit
            // drops the buttons too. If the prompt was already deleted this
            // fails with a 400 "message to edit not found" — log and ignore.
            if let Err(e) = ctx
                .sender
                .edit_message_text(
                    ChatId(chat_id),
                    MessageId(prompt_message_id as i32),
                    send::EDIT_PROMPT_EXPIRED_TEXT.to_string(),
                )
                .await
            {
                log::info!("edit-expiry sweep: prompt message gone: {e}");
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn sweep_removes_only_our_old_temp_entries() {
        let dir = tempfile::tempdir().unwrap();
        let old = std::time::SystemTime::now() - Duration::from_secs(7200);
        let make = |name: &str, aged: bool| {
            let path = dir.path().join(name);
            std::fs::write(&path, b"x").unwrap();
            if aged {
                let file = std::fs::File::options().write(true).open(&path).unwrap();
                file.set_modified(old).unwrap();
            }
            path
        };
        let ours_old = make(&format!("{}photo-old.jpg", x_media::TEMP_FILE_PREFIX), true);
        let ours_fresh = make(
            &format!("{}photo-new.jpg", x_media::TEMP_FILE_PREFIX),
            false,
        );
        let theirs = make("someone-elses-file", true);

        assert_eq!(sweep_temp_dir(dir.path(), Duration::from_secs(3600)), 1);
        assert!(!ours_old.exists(), "an old leftover of ours is removed");
        assert!(ours_fresh.exists(), "a fresh file may belong to a live run");
        assert!(
            theirs.exists(),
            "files without our prefix are never touched"
        );

        // A caller with no age gate also reaches the directory branch (aging a
        // *directory* is not portable, so the gate is what the first half
        // above proves): the fresh dir and file go, the unrelated file stays.
        let leftover_dir = dir
            .path()
            .join(format!("{}ugoira", x_media::TEMP_FILE_PREFIX));
        std::fs::create_dir(&leftover_dir).unwrap();
        std::fs::write(leftover_dir.join("frame.png"), b"x").unwrap();
        assert_eq!(sweep_temp_dir(dir.path(), Duration::ZERO), 2);
        assert!(
            !leftover_dir.exists(),
            "leftover dirs go with their contents"
        );
        assert!(!ours_fresh.exists(), "no age gate: ours, however fresh");
        assert!(theirs.exists());
    }

    /// The sweep's tick: an expired prompt is rewritten in place (buttons
    /// dropped) while a live one is left alone. Driven through the loop's own
    /// timer on a paused clock — the loop is what a hand-called helper would
    /// leave untested, which is how the queue's sweep kept a missing wake-up.
    #[tokio::test(start_paused = true)]
    async fn the_sweep_expires_only_the_prompts_past_their_ttl() {
        use crate::ctx::test_support::{
            FORWARDED_ID, PROMPT_ID, TestStores, api_error, seed_prompt,
        };
        use crate::media_sender::test_support::MockSender;
        use crate::state::EditMessage;

        // The interval is pinned here because no assertion on the edits can see
        // it: a shorter interval produces the same single edit (the record is
        // gone after the first tick), and the paused clock can jump past the
        // boundary while a tick's DB work is in flight.
        assert_eq!(SWEEP_INTERVAL, Duration::from_secs(300));

        let config = crate::config::Config::load();
        let stores = TestStores::new();
        let sender = MockSender::scripted(vec![], || {
            api_error("Bad Request: message to edit not found")
        });
        let ctx = stores.ctx(&sender);
        // Chat 1 holds a prompt past its ttl; chat 2 a live one.
        let stale = crate::db::unix_now() - config.edit_message_ttl.as_secs() as i64 - 1;
        seed_prompt(&ctx, "", stale).await;
        stores
            .chat_store()
            .update(2, |data| {
                data.edit_message.insert(
                    PROMPT_ID,
                    EditMessage {
                        url: "https://x.com/u/status/1".into(),
                        chat_id: 2,
                        forward_message_ids: vec![FORWARDED_ID],
                        template: String::new(),
                        created_at: crate::db::unix_now(),
                    },
                );
            })
            .await;

        let (stop_tx, stop_rx) = watch::channel(false);
        let sweep = periodic_sweep(stores.ctx(&sender), stop_rx);
        tokio::pin!(sweep);

        // One second short of the interval: nothing has been touched. The
        // select is what polls the loop (a pinned future nobody awaits never
        // runs), and the paused clock makes this the loop's own timer.
        tokio::select! {
            _ = &mut sweep => unreachable!("the sweep only returns on stop"),
            _ = tokio::time::sleep(SWEEP_INTERVAL - Duration::from_secs(1)) => {}
        }
        assert!(
            sender.edited_texts().is_empty(),
            "the sweep ran before its interval"
        );

        // The second that crosses the interval: the tick fires.
        tokio::select! {
            _ = &mut sweep => unreachable!("the sweep only returns on stop"),
            _ = tokio::time::sleep(Duration::from_secs(2)) => {}
        }

        assert_eq!(
            sender.edited_texts(),
            vec![(1, PROMPT_ID, send::EDIT_PROMPT_EXPIRED_TEXT.to_string())],
            "exactly the expired prompt, rewritten in place"
        );
        assert!(
            !ctx.chat_store
                .get(1)
                .await
                .edit_message
                .contains_key(&PROMPT_ID),
            "the expired record is gone"
        );
        assert!(
            ctx.chat_store
                .get(2)
                .await
                .edit_message
                .contains_key(&PROMPT_ID),
            "a live prompt keeps its record and its buttons"
        );

        stop_tx.send(true).unwrap();
        sweep.await;
    }
}