作者:令川 | 发布时间:2024-09-03

Rust 实现一个请求代理服务器

前言

本文介绍一下如何使用 Rust 实现一个请求代理服务,重点是 axumreqwest 之间的 RequestResponse 的相互转换。

axumreqwest 分别是 Rust 生态中最为主流的 Web Server 框架、HTTP 请求库,二者都支持异步编程模型,且都依赖 tokio 异步运行时。因此使用它们共同实现代理服务器是个很不错的选择。

axum 的请求传给 reqwest

先来看下 axum 的请求中的 headersbody 如何传递给 reqwest

async fn handler(
    headers: axum::http::HeaderMap,
    // body: axum::body::Bytes,
    body: String,
) -> axum::response::Response {
    let mut req_headers = headers.clone();
    req_headers.remove(http::header::HOST);
    let client = reqwest::Client::new();
    let res = client
        .post("<https://lccl.cc>")
        .headers(req_headers)
        .body(body)
        .send()
        .await;

        // ...
}

对于 headers 的传递,这里特别注意要去除其中的 HOST ,因为传过来的请求的主机地址,一定是当前代理服务器的地址,而非被代理的服务器的地址。

如果发送的请求的 URL 和 headers 中的 HOST 不匹配,会导致请求发送失败。因此这里一定要先移除它。

由于 HOST 是 HTTP1.1 规定必须要有的,因此 reqwest 内部会在检测到外部没有传该值的时候,根据 URL 自动确定一个 HOST

headers.remove() 要求对象是可变的,因此用 let mut req_headers = headers.clone(); 复制了一份可变的 headers

在函数参数中提取 headers 时就声明可变性也是可以的,比如下面这种写法:

async fn handler(mut headers: axum::http::HeaderMap, ...) -> axum::response::Response {
    headers.remove(http::header::HOST);
    reqwest::Client::new()
        .post("<https://lccl.cc>")
        .headers(headers.clone())

		// ...
}

对于 body 的传递,通过 body: Stringbody: axum::body::Bytes 提取请求体都是可以的,并且都可以直接传递给 reqwest::Client()::new().body(body) ,还是比较方便的。

使用 body: axum::body::Body 提取请求体再传给 reqwest 构造的请求就比较麻烦了,涉及到数据类型的转换,这里就不介绍了,大家没有特别需求使用上面两种方式提取请求体即可。

reqwest 的响应传给 axum

reqwest 发送请求获取到响应后,根据它的状态码、响应头、响应体构造 axum 能处理的 Resonse 然后发送回客户端,代理请求的基本功能就算成形了。

状态码的传递比较简单,使用 res.status() 获取后直接传递即可。

headersbody 的传递方式就多种多样了,下面介绍几种方式。

为了减少重复代码,我们约定下面的 res 都是 reqwestsend() 请求发送成功之后返回的 Response 对象:

let res: reqwest::Response = reqwest::Client::new()
		.post("<https://lccl.cc>")
		// ...
		.send()
		.await
		.expect("请求发送失败");

先介绍个最为简单的方式——以字节流的方式传递响应体:

async fn handler(…) -> impl IntoResponse {
		// ...
		(
        res.status(),
        res.headers().clone(),
        axum::body::Body::from_stream(res.bytes_stream()),
    )
        .into_response()
}

这种方式之所以最为简单,是因为需要处理的内容较少。

主要是当 res.headers 中有 Transfer-Encoding: chunked 时,响应体是被一块一块的传输的,不是一次性发送完整的响应体。

并且这对键值对的设定是服务器决定的,和请求中的 headers 没有必然的联系。

这就会带来一个问题,如果我们把带有 Transfer-Encoding: chunkedheaders 传递给 axum 的响应,而 axum 的响应体又不是字节流的数据形式,那这个响应就会构建失败!Rust 中也没有报错!

使用 axum::body::Body::from_stream() 返回响应体时,会自动在 headers 中添加 Transfer-Encoding: chunked

因此使用这种方式传递响应体时,可以确保:

被代理的服务器的响应中,无论 headers 中有没有 Transfer-Encoding: chunked 、无论响应体是不是数据流的形式,都可以通过 res.bytes_stream() 拿到数据流形式的数据,所以这种传递方式最为稳妥、简单。

reqwest::Response 提供的其他的获取响应体数据的方式,都需要通过 .await 拿到完整的数据,这时如果代理服务器的响应的 headers 中还有 Transfer-Encondig: chunked ,响应就会构建失败,例如使用 res.bytes().await 获取服务器的响应体:

async fn handler(…) -> impl IntoResponse {
		// ...
		
		(
        res.status(),
        res.headers().clone(),
        axum::body::Bytes::from(res.bytes().await.unwrap()),
    )
        .into_response()
}

这时就会响应失败。

当然了,如果 res.headers 中没有 Transfer-Encoding: chunked ,这种传递方式也没问题。

如果确实不想用 res.bytes_stream() ,比较保险的的做法是在 headers 中删除这个键值对:

async fn handler(…) -> impl IntoResponse {
		// ...
		
		let mut new_headers = res.headers().clone();
    new_headers.remove(http::header::TRANSFER_ENCODING);
		
		(
        res.status(),
        new_headers,
        // 下面这些方式都是可以的,可以看到它们都有 .await,即等到完整响应体返回完成
        axum::body::Bytes::from(res.bytes().await.unwrap())
        axum::body::Bytes::from(res.chunk().await.unwrap().unwrap())
        axum::body::Body::from(res.bytes().await.unwrap())
        axum::body::Body::from(res.chunk().await.unwrap().unwrap())
        axum::body::Body::from(res.text().await.unwrap())
    )
        .into_response()
}

再介绍下使用 axum::response::Response::builder() 的写法:

async fn handler(…) -> Response {
		// ...
		
    let mut builder = Response::builder().status(res.status());
    *builder.headers_mut().unwrap() = res.headers().clone();
    builder
        .body(axum::body::Body::from_stream(res.bytes_stream()))
        .unwrap()
}

实际案例:代理 Github 获取 Token 接口

下面写一个实际的请求代理函数:

/// 代理请求:POST <https://github.com/login/oauth/access_token>
async fn github_access_token_proxy(
    mut headers: axum::http::HeaderMap,
    body: axum::body::Bytes,
) -> impl IntoResponse {
    headers.remove(http::header::HOST);
    let res = reqwest::Client::new()
        .post("<https://github.com/login/oauth/access_token>")
        .headers(headers)
        .body(body)
        .send()
        .await;
    match res {
        Err(_) => (StatusCode::INTERNAL_SERVER_ERROR).into_response(),
        Ok(res) => (
            res.status(),
            res.headers().clone(),
            Body::from_stream(res.bytes_stream()),
        ).into_response()
    }
}

JavaScript 版本的代理函数

同样是上面的代理函数,看下 JavaScript 版本的代码:

async function(req: Request): Promise<Response> {
  return fetch("<https://github.com/login/oauth/access_token>", {
    method: "POST",
    headers: req.headers,
    body: req.body,
  });
}

😂

使用 fetch 时, headers 中的 HOST 是不允许自定义的,来源见:MDN:禁止修改的标头

底层会根据传入的 URL 自动设置一个合理的 HOST 值。

同时 fetch 的结果 Promise<Response> 是可以作为响应直接返回的,不存在 Transfer-Encoding 和响应体数据格式不一致所带来的问题,也就不需要额外处理。

总结

Rust 作为一门系统级的编程语言,需要考虑的事情确实偏底层一些。

作为一名前端开发,可能很少有机会去了解请求头中的 HOST 如果和 URL 不匹配会发生什么、响应头中 Transfer-Encoding: chunked 具体代表什么含义、作为一个代理服务器如何把这种响应透传出去……

说回 Rust。

如果想要把 axum 的请求的 headers 传递给 reqwest ,可以通过 headers: axum::http::HashMap 获取,然后记得先 remove(http::headers::HOST) 再传给 reqwest 构造的请求。

如果想要把 axum 的请求体传给 reqwest ,可以通过 body: axum::body:Bytesbdoy: String 提取之后直接传递。

如果想要把 reqwest::Response 转换为 axum 能接收的响应,最简单的方式:

(
		res.status(),
		res.headers().clone(),
		axum::body::Body::from_stream(res.bytes_strem())
).into_response()

Over~

目录 / Contents

空。

令川 · 记