relay.rs

//! 別サーバへの中継。受けたリクエストを転送し、返ってきたものを返す。
//! 転送先は設定の relay_origin（環境変数 RELAY_ORIGIN）で決める。
//!
//! 行きも帰りも流したまま渡す。全部をメモリに読んでから送ることはしない。
//!
//! mon のセッションのクッキーは、行きでは中継先へ渡さず、帰りでは中継先が書こうとしても通さない。
//! 中継先に id を知られると、そのセッションを使われる。書かれると、相手のセッションが差し替わる。

use axum::Router;
use axum::body::Body;
use axum::extract::{Request, State};
use axum::http::{HeaderMap, HeaderValue, header};
use axum::response::Response;
use axum::routing::any;
use crate::conf::Conf;
use crate::error::AppError;
use crate::ip;
use crate::session::Store;
use crate::state::AppState;


/// 転送しないヘッダ。接続そのものに関わるもの
const SKIP_HEADER: &[&str] = &[
    "connection",
    "host",
    "keep-alive",
    "proxy-authenticate",
    "proxy-authorization",
    "te",
    "trailer",
    "transfer-encoding",
    "upgrade",
];

/// 前段の情報を運ぶヘッダのうち、x-forwarded- で始まらないもの。
/// どれも相手が自分で書けるので、mon が最前段なら中継先へ渡さない
const FORWARDED_HEADER: &[&str] = &["forwarded", "x-real-ip"];


/// 中継に使うクライアント。
/// 待ち時間の上限は「つなぐまで」と「次の一片が届くまで」に掛ける。
/// 全体に掛けると、大きな応答を流している途中で切れてしまう。
///
/// リダイレクトは追わない。中継先が返した 302 は、そのまま相手へ渡すものである。
/// 追いかけると相手には最終の応答だけが届き、元の 302 と行き先が消える。
pub fn client(conf: &Conf) -> anyhow::Result<reqwest::Client> {
    let client = reqwest::Client::builder()
        .connect_timeout(conf.relay_timeout)
        .read_timeout(conf.relay_timeout)
        .redirect(reqwest::redirect::Policy::none())
        .build()?;
    Ok(client)
}


/// 中継のルート。`/` と、その下のすべてを受ける
pub fn router() -> Router<AppState> {
    Router::new()
        .route("/", any(relay))
        .route("/{*path}", any(relay))
}


async fn relay(State(state): State<AppState>, request: Request) -> Result<Response, AppError> {
    let (parts, body) = request.into_parts();
    let conf = &state.conf;
    let path = parts
        .uri
        .path_and_query()
        .map(|value| value.as_str())
        .unwrap_or("/");
    let url = format!("{}{}", conf.relay_origin, path);

    // CLIENT_IP_HEADER が空なら、mon は自分が最前段だとみなす（ip.rs と同じ考え方）。
    // そのとき forwarded 系は相手が書いたものなので、中継先へは渡さず、mon が書き直す
    let front = conf.client_ip_header.is_empty();
    let cookie = session_cookie(&state);

    let mut upstream = state
        .client
        .request(parts.method, &url)
        .body(reqwest::Body::wrap_stream(body.into_data_stream()));
    for (name, value) in &parts.headers {
        let text = name.as_str();
        let forwarded = text.starts_with("x-forwarded-") || FORWARDED_HEADER.contains(&text);
        let session = cookie.is_some() && name == header::COOKIE;

        if SKIP_HEADER.contains(&text) || (front && forwarded) || session {
            continue;
        }

        upstream = upstream.header(name, value);
    }

    // Cookie は、mon のセッションのものだけを外して渡す
    if let Some(cookie) = cookie
        && let Some(line) = cookie_without(&parts.headers, cookie)
    {
        upstream = upstream.header(header::COOKIE, line);
    }

    // 前段がいるなら、forwarded 系は前段が付けたものである。
    // 書き換えると中継先が数える位置がずれるので、ないときだけ足す
    if (front || !parts.headers.contains_key("x-forwarded-for"))
        && let Some(address) = ip::resolve(&parts.headers, &parts.extensions, conf)
    {
        upstream = upstream.header("x-forwarded-for", address.to_string());
    }

    // 中継先が canonical や OGP を組み立てるには、公開しているホストと scheme が要る。
    // host は転送しないので、これがないと中継先は自分の名前しか知れない
    if front || !parts.headers.contains_key("x-forwarded-proto") {
        upstream = upstream.header("x-forwarded-proto", parts.uri.scheme_str().unwrap_or("http"));
    }

    if (front || !parts.headers.contains_key("x-forwarded-host"))
        && let Some(host) = parts.headers.get(header::HOST)
    {
        upstream = upstream.header("x-forwarded-host", host);
    }

    let answer = upstream.send().await?;
    let mut built = Response::builder().status(answer.status());
    for (name, value) in answer.headers() {
        if SKIP_HEADER.contains(&name.as_str()) {
            continue;
        }

        // 中継先が mon のセッションのクッキーを書くと、相手のセッションが差し替わる
        if name == header::SET_COOKIE
            && let Some(cookie) = cookie
            && sets_cookie(value, cookie)
        {
            tracing::warn!("中継先が {cookie} のクッキーを書こうとしたので外しました");
            continue;
        }

        built = built.header(name, value);
    }

    built
        .body(Body::from_stream(answer.bytes_stream()))
        .map_err(|error| AppError::Internal(error.into()))
}


/// mon のセッションのクッキー名。セッションを使っていなければ None である。
/// 使っていないなら、同じ名前のクッキーは中継先のものかもしれないので触らない
fn session_cookie(state: &AppState) -> Option<&str> {
    if matches!(state.session, Store::None) {
        return None;
    }

    Some(state.conf.session_cookie.as_str())
}


/// Cookie から `name` のものを外し、1 行につなぐ。何も残らなければ None を返す。
/// Cookie は複数行に分かれて届くことがあるので、全部を見る
fn cookie_without(headers: &HeaderMap, name: &str) -> Option<String> {
    let rest: Vec<&str> = headers
        .get_all(header::COOKIE)
        .iter()
        .filter_map(|value| value.to_str().ok())
        .flat_map(|line| line.split(';'))
        .map(str::trim)
        .filter(|pair| !pair.is_empty() && cookie_name(pair) != name)
        .collect();

    if rest.is_empty() {
        return None;
    }

    Some(rest.join("; "))
}


/// Set-Cookie が `name` のクッキーを書くものか
fn sets_cookie(value: &HeaderValue, name: &str) -> bool {
    value
        .to_str()
        .ok()
        .and_then(|line| line.split(';').next())
        .is_some_and(|pair| cookie_name(pair) == name)
}


/// `name=value` の name の部分
fn cookie_name(pair: &str) -> &str {
    pair.split_once('=').map_or(pair, |(name, _)| name).trim()
}


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

    fn cookie_headers(lines: &[&str]) -> HeaderMap {
        let mut headers = HeaderMap::new();
        for line in lines {
            headers.append(header::COOKIE, HeaderValue::from_str(line).unwrap());
        }

        headers
    }

    #[test]
    fn only_the_session_cookie_is_removed() {
        let headers = cookie_headers(&["sid=abc; theme=dark", "lang=ja"]);
        assert_eq!(cookie_without(&headers, "sid").unwrap(), "theme=dark; lang=ja");
    }

    #[test]
    fn nothing_is_left_when_only_the_session_cookie_is_sent() {
        let headers = cookie_headers(&["sid=abc"]);
        assert_eq!(cookie_without(&headers, "sid"), None);
    }

    /// 名前の前半が同じだけのクッキーは外さない
    #[test]
    fn a_similar_name_is_kept() {
        let headers = cookie_headers(&["sidebar=open; sid=abc"]);
        assert_eq!(cookie_without(&headers, "sid").unwrap(), "sidebar=open");
    }

    #[test]
    fn a_set_cookie_for_the_session_is_recognized() {
        assert!(sets_cookie(&HeaderValue::from_static("sid=planted; Path=/"), "sid"));
        assert!(!sets_cookie(&HeaderValue::from_static("theme=dark; Path=/"), "sid"));
        assert!(!sets_cookie(&HeaderValue::from_static("sidebar=open"), "sid"));
    }
}