amagi 0.1.1

Rust SDK, CLI, and Web API service skeleton for multi-platform social web adapters.
Documentation
use axum::{
    Json,
    http::{HeaderMap, StatusCode},
};

use super::super::state::AppState;
use super::{CatalogErrorResponse, FetchErrorResponse};
use crate::catalog::Platform;
use crate::error::AppError;
use crate::platforms::bilibili::BilibiliFetcher;
use crate::platforms::douyin::DouyinFetcher;
use crate::platforms::kuaishou::KuaishouFetcher;
use crate::platforms::twitter::TwitterFetcher;
use crate::platforms::xiaohongshu::XiaohongshuFetcher;

pub(super) type CatalogResult<T> = Result<Json<T>, (StatusCode, Json<CatalogErrorResponse>)>;
pub(super) type FetchResult<T> = Result<Json<T>, (StatusCode, Json<FetchErrorResponse>)>;

const GENERIC_COOKIE_HEADER: &str = "x-amagi-cookie";
const BILIBILI_COOKIE_HEADER: &str = "x-amagi-bilibili-cookie";
const DOUYIN_COOKIE_HEADER: &str = "x-amagi-douyin-cookie";
const KUAISHOU_COOKIE_HEADER: &str = "x-amagi-kuaishou-cookie";
const TWITTER_COOKIE_HEADER: &str = "x-amagi-twitter-cookie";
const XIAOHONGSHU_COOKIE_HEADER: &str = "x-amagi-xiaohongshu-cookie";

enum RequestCookieOverride {
    UseConfigured,
    Override(String),
    Clear,
}

pub(super) fn fetch_error_response(error: AppError) -> (StatusCode, Json<FetchErrorResponse>) {
    let status = match &error {
        AppError::Io(_) | AppError::InvalidRequestConfig(_) => StatusCode::INTERNAL_SERVER_ERROR,
        AppError::Json(_) | AppError::Http(_) | AppError::UpstreamResponse { .. } => {
            StatusCode::BAD_GATEWAY
        }
    };

    (
        status,
        Json(FetchErrorResponse {
            error: "fetch_failed",
            detail: error.to_string(),
        }),
    )
}

pub(super) fn bilibili_fetcher(state: &AppState, headers: &HeaderMap) -> BilibiliFetcher {
    BilibiliFetcher::new(platform_client_for_request(
        state,
        Platform::Bilibili,
        headers,
    ))
}

pub(super) fn douyin_fetcher(state: &AppState, headers: &HeaderMap) -> DouyinFetcher {
    DouyinFetcher::new(platform_client_for_request(
        state,
        Platform::Douyin,
        headers,
    ))
}

pub(super) fn kuaishou_fetcher(state: &AppState, headers: &HeaderMap) -> KuaishouFetcher {
    KuaishouFetcher::new(platform_client_for_request(
        state,
        Platform::Kuaishou,
        headers,
    ))
}

pub(super) fn xiaohongshu_fetcher(state: &AppState, headers: &HeaderMap) -> XiaohongshuFetcher {
    XiaohongshuFetcher::new(platform_client_for_request(
        state,
        Platform::Xiaohongshu,
        headers,
    ))
}

pub(super) fn twitter_fetcher(state: &AppState, headers: &HeaderMap) -> TwitterFetcher {
    TwitterFetcher::new(platform_client_for_request(
        state,
        Platform::Twitter,
        headers,
    ))
}

fn platform_client_for_request(
    state: &AppState,
    platform: Platform,
    headers: &HeaderMap,
) -> crate::client::PlatformClient {
    let mut client = state.client.platform(platform);

    match request_cookie_override(headers, platform) {
        RequestCookieOverride::UseConfigured => client,
        RequestCookieOverride::Override(cookie) => {
            strip_cookie_headers(&mut client.request.headers);
            client.cookie = Some(cookie);
            client
        }
        RequestCookieOverride::Clear => {
            strip_cookie_headers(&mut client.request.headers);
            client.cookie = None;
            client
        }
    }
}

fn request_cookie_override(headers: &HeaderMap, platform: Platform) -> RequestCookieOverride {
    for header_name in [platform_cookie_header(platform), GENERIC_COOKIE_HEADER] {
        if let Some(value) = headers.get(header_name) {
            let Ok(value) = value.to_str() else {
                continue;
            };

            let trimmed = value.trim();
            return if trimmed.is_empty() {
                RequestCookieOverride::Clear
            } else {
                RequestCookieOverride::Override(trimmed.to_owned())
            };
        }
    }

    RequestCookieOverride::UseConfigured
}

fn platform_cookie_header(platform: Platform) -> &'static str {
    match platform {
        Platform::Bilibili => BILIBILI_COOKIE_HEADER,
        Platform::Douyin => DOUYIN_COOKIE_HEADER,
        Platform::Kuaishou => KUAISHOU_COOKIE_HEADER,
        Platform::Twitter => TWITTER_COOKIE_HEADER,
        Platform::Xiaohongshu => XIAOHONGSHU_COOKIE_HEADER,
    }
}

fn strip_cookie_headers(headers: &mut std::collections::BTreeMap<String, String>) {
    headers.remove("Cookie");
    headers.remove("cookie");
}

#[cfg(test)]
mod tests {
    use axum::http::{HeaderMap, HeaderValue};

    use super::{
        GENERIC_COOKIE_HEADER, RequestCookieOverride, TWITTER_COOKIE_HEADER,
        request_cookie_override,
    };
    use crate::catalog::Platform;

    #[test]
    fn prefers_platform_cookie_header_over_generic_header() {
        let mut headers = HeaderMap::new();
        headers.insert(
            GENERIC_COOKIE_HEADER,
            HeaderValue::from_static("generic=value"),
        );
        headers.insert(
            TWITTER_COOKIE_HEADER,
            HeaderValue::from_static("auth_token=platform"),
        );

        let override_value = request_cookie_override(&headers, Platform::Twitter);

        assert!(matches!(
            override_value,
            RequestCookieOverride::Override(value) if value == "auth_token=platform"
        ));
    }

    #[test]
    fn uses_generic_cookie_header_when_platform_header_is_missing() {
        let mut headers = HeaderMap::new();
        headers.insert(
            GENERIC_COOKIE_HEADER,
            HeaderValue::from_static("auth_token=generic"),
        );

        let override_value = request_cookie_override(&headers, Platform::Twitter);

        assert!(matches!(
            override_value,
            RequestCookieOverride::Override(value) if value == "auth_token=generic"
        ));
    }

    #[test]
    fn blank_cookie_header_clears_configured_cookie() {
        let mut headers = HeaderMap::new();
        headers.insert(TWITTER_COOKIE_HEADER, HeaderValue::from_static("   "));

        let override_value = request_cookie_override(&headers, Platform::Twitter);

        assert!(matches!(override_value, RequestCookieOverride::Clear));
    }
}