当前位置：首页 > news >正文

【Rust网络编程】开发一个图片代理和统计服务

news 来源：原创 2024/9/28 9:08:52

最近我使用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 {// 构建新的 URIlet 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 Foundlet 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