relay.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106

//! 別サーバへの中継。受けたリクエストをそのまま転送し、返ってきたものをそのまま返す。
//! 転送先は設定の relay_origin（環境変数 RELAY_ORIGIN）で決める。
//!
//! 行きも帰りも流したまま渡す。全部をメモリに読んでから送ることはしない。

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


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


/// 中継に使うクライアント。
/// 待ち時間の上限は「つなぐまで」と「次の一片が届くまで」に掛ける。
/// 全体に掛けると、大きな応答を流している途中で切れてしまう。
///
/// リダイレクトは追わない。中継先が返した 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 path = parts
        .uri
        .path_and_query()
        .map(|value| value.as_str())
        .unwrap_or("/");
    let url = format!("{}{}", state.conf.relay_origin, path);

    let mut upstream = state
        .client
        .request(parts.method, &url)
        .body(reqwest::Body::wrap_stream(body.into_data_stream()));
    for (name, value) in &parts.headers {
        if !SKIP_HEADER.contains(&name.as_str()) {
            upstream = upstream.header(name, value);
        }
    }

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

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

    if !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()) {
            built = built.header(name, value);
        }
    }

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