最近我使用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
參考
- https://github.com/VinciYan/proxy_counter.git
- Getting Started | hyper