【Rust网络编程】开发一个图片代理和统计服务
最近我使用Rust开发了一个代理服务。可以用于代理和统计图片资源的访问
例如:
http://127.0.0.1:8100/image-public/0a1e65f4-7ced-4ef0-ba7d-12ec4d14a0d4.png
->http://xxx.com:45004/image-public/0a1e65f4-7ced-4ef0-ba7d-12ec4d14a0d4.png
项目特点
- 高性能:使用Rust语言编写
- 异步处理:基于Tokio运行时,实现高并发的异步I/O操作
- 精确统计:准确记录目标图片的访问次数
技术栈
- Rust 编程语言
- Hyper:用于HTTP服务器和客户端的快速、安全框架。性能好,偏底层,应用广泛,知名的reqwest和axum等都使用了hyper,已成为Rust网络程序生态的重要基石之一
- Tokio:异步运行时,提供高效的I/O操作
功能介绍
-
HTTP 代理:
- 监听本地端口(默认8100),接收 HTTP 请求
- 访问目标图片(路径以/image-public/开头)将被代理,转发到配置的目标服务器
-
图片访问统计:
- 精确统计目标图片的访问次数
-
请求日志:
- 详细记录每个请求的方法、路径和头部信息
- 输出响应状态码和图片访问计数
-
错误处理:
- 对于错误图片的请求,返回404 Not Found响应
代码
核心代码如下:
#![deny(warnings)]
use std::net::SocketAddr;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::str::FromStr;
use bytes::Bytes;
use http_body_util::{combinators::BoxBody, BodyExt, Empty, Full};
use hyper::client::conn::http1::Builder;
use hyper::server::conn::http1;
use hyper::service::service_fn;
use hyper::upgrade::Upgraded;
use hyper::{Method, Request, Response, StatusCode, Uri};
use tokio::net::{TcpListener, TcpStream};
#[path = "../benches/support/mod.rs"]
mod support;
use support::TokioIo;
// 图片下载计数器
static IMAGE_DOWNLOAD_COUNT: AtomicUsize = AtomicUsize::new(0);
#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
let addr = SocketAddr::from(([127, 0, 0, 1], 8100));
let listener = TcpListener::bind(addr).await?;
println!("正在监听 http://{}", addr);
loop {
let (stream, _) = listener.accept().await?;
let io = TokioIo::new(stream);
tokio::task::spawn(async move {
if let Err(err) = http1::Builder::new()
.preserve_header_case(true)
.title_case_headers(true)
.serve_connection(io, service_fn(proxy))
.with_upgrades()
.await
{
println!("服务连接失败: {:?}", err);
}
});
}
}
async fn proxy(
req: Request<hyper::body::Incoming>,
) -> Result<Response<BoxBody<Bytes, hyper::Error>>, hyper::Error> {
println!("收到请求: 方法={:?}, 路径={}, 头部={:?}", req.method(), req.uri().path(), req.headers());
println!("请求: {:?}", req);
if Method::CONNECT == req.method() {
if let Some(addr) = host_addr(req.uri()) {
tokio::task::spawn(async move {
match hyper::upgrade::on(req).await {
Ok(upgraded) => {
if let Err(e) = tunnel(upgraded, addr).await {
eprintln!("服务器 IO 错误: {}", e);
};
}
Err(e) => eprintln!("升级错误: {}", e),
}
});
Ok(Response::new(empty()))
} else {
eprintln!("CONNECT 主机不是 socket 地址: {:?}", req.uri());
let mut resp = Response::new(full("CONNECT 必须连接到 socket 地址"));
*resp.status_mut() = StatusCode::BAD_REQUEST;
Ok(resp)
}
} else {
// 检查是否是目标图片下载请求
let is_target_image = req.uri().path().starts_with("/image-public/");
if is_target_image {
// 构建新的 URI
let new_uri = format!("http://xxx.com:45004{}", req.uri().path());
let new_uri = Uri::from_str(&new_uri).expect("无效的 URI");
// 保存原始路径
let original_path = req.uri().path().to_string();
// 创建新的请求
let (parts, body) = req.into_parts();
let mut new_req = Request::new(body);
*new_req.method_mut() = parts.method;
*new_req.uri_mut() = new_uri;
*new_req.version_mut() = parts.version;
*new_req.headers_mut() = parts.headers;
// 连接到实际的服务器
let stream = TcpStream::connect(("xxx.com", 45004)).await.unwrap();
let io = TokioIo::new(stream);
let (mut sender, conn) = Builder::new()
.preserve_header_case(true)
.title_case_headers(true)
.handshake(io)
.await?;
tokio::task::spawn(async move {
if let Err(err) = conn.await {
println!("连接失败: {:?}", err);
}
});
let resp = sender.send_request(new_req).await?;
// 如果是目标图片请求,增加计数器
if resp.status().is_success() || resp.status() == StatusCode::NOT_MODIFIED {
let count = IMAGE_DOWNLOAD_COUNT.fetch_add(1, Ordering::SeqCst);
println!("目标图片请求成功。状态码: {}. 总计数: {}", resp.status(), count + 1);
} else if resp.status() == StatusCode::NOT_FOUND {
println!("目标图片不存在。路径: {}", original_path);
let mut not_found_resp = Response::new(full("Image Not Found"));
*not_found_resp.status_mut() = StatusCode::NOT_FOUND;
return Ok(not_found_resp);
}
Ok(resp.map(|b| b.boxed()))
} else {
// 对于非目标图片请求,返回 404 Not Found
let mut resp = Response::new(full("Not Found"));
*resp.status_mut() = StatusCode::NOT_FOUND;
Ok(resp)
}
}
}
fn host_addr(uri: &http::Uri) -> Option<String> {
uri.authority().and_then(|auth| Some(auth.to_string()))
}
fn empty() -> BoxBody<Bytes, hyper::Error> {
Empty::<Bytes>::new()
.map_err(|never| match never {})
.boxed()
}
fn full<T: Into<Bytes>>(chunk: T) -> BoxBody<Bytes, hyper::Error> {
Full::new(chunk.into())
.map_err(|never| match never {})
.boxed()
}
async fn tunnel(upgraded: Upgraded, addr: String) -> std::io::Result<()> {
let mut server = TcpStream::connect(addr).await?;
let mut upgraded = TokioIo::new(upgraded);
let (from_client, from_server) =
tokio::io::copy_bidirectional(&mut upgraded, &mut server).await?;
println!(
"客户端写入 {} 字节并接收 {} 字节",
from_client, from_server
);
Ok(())
}
完整代码参考我的仓库:https://github.com/VinciYan/proxy_counter.git
运行效果
收到请求: 方法=GET, 路径=/image-public/0a1e65f4-7ced-4ef0-ba7d-12ec4d14a0d4.png, 头部={"content-type": "application/json", "user-agent": "PostmanRuntime/7.42.0", "accept": "*/*", "postman-token": "9fe5ee1a-ad8e-4e0d-8f65-e82090115795", "host": "127.0.0.1:8100", "accept-encoding": "gzip, deflate, br", "connection": "keep-alive", "content-length": "75"}
请求: Request { method: GET, uri: /image-public/0a1e65f4-7ced-4ef0-ba7d-12ec4d14a0d4.png, version: HTTP/1.1, headers: {"content-type": "application/json", "user-agent": "PostmanRun
time/7.42.0", "accept": "*/*", "postman-token": "9fe5ee1a-ad8e-4e0d-8f65-e82090115795", "host": "127.0.0.1:8100", "accept-encoding": "gzip, deflate, br", "connection": "keep-alive", "content-length": "75"}, body: Body(Streaming) }
目标图片请求成功。状态码: 200 OK. 总计数: 3