跳转到内容
作者：令川 | 发布时间：2024-09-03

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

前言

本文介绍一下如何使用 Rust 实现一个请求代理服务，重点是 axum 和 reqwest 之间的 Request和 Response 的相互转换。

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

axum 的请求传给 reqwest

先来看下 axum 的请求中的 headers 和 body 如何传递给 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: String 或 body: axum::body::Bytes 提取请求体都是可以的，并且都可以直接传递给 reqwest::Client()::new().body(body) ，还是比较方便的。

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

reqwest 的响应传给 axum

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

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

而 headers 和 body 的传递方式就多种多样了，下面介绍几种方式。

为了减少重复代码，我们约定下面的 res 都是 reqwest 在 send() 请求发送成功之后返回的 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: chunked 的 headers 传递给 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:Bytes 或 bdoy: String 提取之后直接传递。

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

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

Over～

目录