【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操作

功能介绍

  1. HTTP 代理:

    • 监听本地端口(默认8100),接收 HTTP 请求
    • 访问目标图片(路径以/image-public/开头)将被代理,转发到配置的目标服务器
  2. 图片访问统计:

    • 精确统计目标图片的访问次数
  3. 请求日志:

    • 详细记录每个请求的方法、路径和头部信息
    • 输出响应状态码和图片访问计数
  4. 错误处理:

    • 对于错误图片的请求,返回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

参考

posted @ 2024-09-28 01:18  VinciYan  阅读(21)  评论(0编辑  收藏  举报