hottub/src/util/flaresolverr.rs

use std::{collections::HashMap, env, sync::Arc};

use serde_json::Value;
use serde_json::json;
use tokio::sync::Mutex;
use wreq::{Client, Proxy};
use wreq_util::Emulation;

#[derive(serde::Serialize, serde::Deserialize, Debug, Clone)]
pub struct FlareSolverrRequest {
    pub cmd: String,
    pub url: String,
    pub maxTimeout: u32,
}

#[derive(serde::Serialize, serde::Deserialize, Debug)]
pub struct FlaresolverrCookie {
    pub name: String,
    pub value: String,
    pub domain: String,
    pub path: String,
    pub expires: f64,
    pub size: u64,
    pub httpOnly: bool,
    pub secure: bool,
    pub session: bool,
    pub sameSite: Option<String>,
    pub priority: String,
    pub sameParty: bool,
    pub sourceScheme: String,
    pub sourcePort: u32,
    pub partitionKey: Option<String>,
}

#[derive(serde::Serialize, serde::Deserialize, Debug)]
pub struct FlareSolverrSolution {
    pub url: String,
    pub status: u32,
    pub response: String,
    pub headers: HashMap<String, String>,
    pub cookies: Vec<FlaresolverrCookie>,
    pub userAgent: String,
}

#[derive(serde::Serialize, serde::Deserialize, Debug)]
pub struct FlareSolverrResponse {
    pub status: String,
    pub message: String,
    pub solution: FlareSolverrSolution,
    pub startTimestamp: u64,
    pub endTimestamp: u64,
    pub version: String,
}

#[derive(Clone)]
pub struct Flaresolverr {
    url: String,
    proxy: bool,
}

#[derive(Debug, Default)]
struct SessionState {
    ready_session: Option<String>,
}

fn global_session_state() -> &'static Arc<Mutex<SessionState>> {
    static STATE: std::sync::OnceLock<Arc<Mutex<SessionState>>> = std::sync::OnceLock::new();
    STATE.get_or_init(|| Arc::new(Mutex::new(SessionState::default())))
}

impl Flaresolverr {
    pub fn new(url: String) -> Self {
        Self { url, proxy: false }
    }

    pub fn set_proxy(&mut self, proxy: bool) {
        self.proxy = proxy;
    }

    async fn post_payload(
        &self,
        payload: Value,
    ) -> Result<Value, Box<dyn std::error::Error + Send + Sync>> {
        let client = Client::builder().emulation(Emulation::Firefox136).build()?;

        let mut req = client
            .post(&self.url)
            .header("Content-Type", "application/json")
            .json(&payload);

        if self.proxy {
            if let Ok(proxy_url) = env::var("BURP_URL") {
                match Proxy::all(&proxy_url) {
                    Ok(proxy) => {
                        req = req.proxy(proxy);
                    }
                    Err(e) => {
                        eprintln!("Invalid proxy URL '{}': {}", proxy_url, e);
                    }
                }
            }
        }

        let response = req.send().await?;
        let body = response.json::<Value>().await?;
        if body
            .get("status")
            .and_then(Value::as_str)
            .is_some_and(|status| status.eq_ignore_ascii_case("error"))
        {
            let message = body
                .get("message")
                .and_then(Value::as_str)
                .unwrap_or("FlareSolverr returned status=error");
            return Err(message.to_string().into());
        }
        Ok(body)
    }

    async fn create_session(&self) -> Result<String, Box<dyn std::error::Error + Send + Sync>> {
        let body = self
            .post_payload(json!({ "cmd": "sessions.create" }))
            .await?;
        let session = body
            .get("session")
            .and_then(Value::as_str)
            .ok_or("sessions.create response missing `session`")?;
        Ok(session.to_string())
    }

    async fn destroy_session(
        &self,
        session: &str,
    ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
        let _ = self
            .post_payload(json!({
                "cmd": "sessions.destroy",
                "session": session,
            }))
            .await?;
        Ok(())
    }

    async fn solve_with_session(
        &self,
        request: FlareSolverrRequest,
        session: &str,
    ) -> Result<FlareSolverrResponse, Box<dyn std::error::Error + Send + Sync>> {
        let body = self
            .post_payload(json!({
                "cmd": request.cmd,
                "url": request.url,
                "maxTimeout": request.maxTimeout,
                "session": session,
            }))
            .await?;
        let typed = serde_json::from_value::<FlareSolverrResponse>(body)?;
        Ok(typed)
    }

    async fn ensure_ready_session_locked(
        &self,
        state: &mut SessionState,
    ) -> Result<String, Box<dyn std::error::Error + Send + Sync>> {
        if let Some(existing) = state.ready_session.clone() {
            return Ok(existing);
        }
        let created = self.create_session().await?;
        state.ready_session = Some(created.clone());
        Ok(created)
    }

    pub async fn solve(
        &self,
        request: FlareSolverrRequest,
    ) -> Result<FlareSolverrResponse, Box<dyn std::error::Error + Send + Sync>> {
        // Keep one ready session globally and rotate it per solve:
        // - solve with current ready session
        // - create replacement session in parallel
        // - destroy old session
        // - keep replacement as new ready session
        let session_state = global_session_state().clone();
        let mut state = session_state.lock().await;
        let active_session = self.ensure_ready_session_locked(&mut state).await?;

        let replacement_creator = {
            let solver = self.clone();
            tokio::spawn(async move { solver.create_session().await })
        };

        let solve_result = self.solve_with_session(request, &active_session).await;

        let replacement_session = match replacement_creator.await {
            Ok(Ok(session)) => session,
            Ok(Err(error)) => {
                eprintln!(
                    "FlareSolverr replacement session creation failed, retrying inline: {}",
                    error
                );
                self.create_session().await?
            }
            Err(join_error) => {
                eprintln!(
                    "FlareSolverr replacement task join failed, retrying inline: {}",
                    join_error
                );
                self.create_session().await?
            }
        };

        if let Err(error) = self.destroy_session(&active_session).await {
            eprintln!(
                "FlareSolverr session cleanup failed for session '{}': {}",
                active_session, error
            );
        }

        state.ready_session = Some(replacement_session);
        solve_result
    }
}