17. 从零用Rust编写正反向代理, Rust中一些功能的实现

精选原创

问蒙服务框架 2024-01-09 17:38:55 博主文章分类：wmproxy ©著作权

文章标签 数据 Rust 引用计数 文章分类 架构后端开发 yyds干货盘点

©著作权归作者所有：来自51CTO博客作者问蒙服务框架的原创作品，请联系作者获取转载授权，否则将追究法律责任

wmproxy

wmproxy是由Rust编写，已实现http/https代理，socks5代理，反向代理，静态文件服务器，内网穿透，配置热更新等，后续将实现websocket代理等，同时会将实现过程分享出来，感兴趣的可以一起造个轮子法

项目地址

gite: https://gitee.com/tickbh/wmproxy

github: https://github.com/tickbh/wmproxy

日志功能

为了更容易理解程序中发生的情况，我们可能想要添加一些日志语句。通常在编写应用程序时这很容易。「在某种程度上，日志记录与使用 println! 相同，只是你可以指定消息的重要性」。在rust中定义的日志级别有5种分别为error、warn、info、debug 和 trace 定义日志的级别是表示只关系这级别的日志及更高级别的日志:

定义log，则包含所有的级别定义warn，则只会显示error或者warn的消息

要向应用程序添加日志记录，你需要两样东西：

log crate，rust官方指定的日志级别库
一个实际将日志输出写到有用位置的适配器

当下我们选用的是流行的根据环境变量指定的适配器env_logger，它会根据环境变量中配置的值，日志等级，或者只开启指定的库等功能，或者不同的库分配不同的等级等。

在Linux或者MacOs上开启功能

env RUST_LOG=debug cargo run

在Windows PowerShell上开启功能

$env:RUST_LOG="debug"
cargo run

在Windows CMD上开启功能

set RUST_LOG="debug"
cargo run

如果我们指定库等级可以设置

RUST_LOG="info,wenmeng=warn,webparse=warn"

这样就可以减少第三方库打日志给程序带来的干扰

需要在Cargo.toml中引用

[dependencies]
log = "0.4.20"
env_logger = "0.10.0"

以下是示意代码

use log::{info, warn};
fn main() {
    env_logger::init();
    info!("欢迎使用软件wmproxy");
    warn!("现在已经成功启动");
}

用println!将会直接输出到stdout，当日志数据多的时候，无法进行关闭，做为第三方库，就不能干扰引用库的正常看日志，所以这只能调试的时候使用，或者少量的关键地方使用。

多个TcpListener的Accept

因为当前支持多个端口绑定，或者配置没有配置，存在None的情况，我们需要同时在一个线程中await所有的TcpListener。在这里我们先用的是tokio::select!对多个TcpListener同时进行await。如果此时我们没有绑定proxy的绑定地址，此时listener为None，但我们需要进行判断才知道他是否为None，如果我们用以下写法：

use tokio::net::TcpListener;
use std::io;

#[tokio::main]
async fn main() -> io::Result<()> {
    let mut listener: Option<TcpListener> = None;
    tokio::select! {
        // 加了if条件判断是否有值
        Ok((conn, addr)) = listener.as_mut().unwrap().accept(), if listener.is_some() => {
            println!("accept addr = {:?}", addr);
        }
    }
    Ok(())
}

此时我们试运行，依然报以下错误：

thread 'main' panicked at 'called `Option::unwrap()` on a `None` value', examples/udp.rs:9:46

也就是即使加了if条件我们也正确的执行我们的操作，因为tokio::select的每个分支必须返回Fut，此时如果为None，就不能返回Fut违反了该函数的定义，那么我们做以下封装：

async fn tcp_listen_work(listen: &Option<TcpListener>) -> Option<(TcpStream, SocketAddr)> {
    if listen.is_some() {
        match listen.as_ref().unwrap().accept().await {
            Ok((tcp, addr)) => Some((tcp, addr)),
            Err(_e) => None,
        }
    } else {
        // 如果为None的时候，就永远返回Poll::Pending
        let pend = std::future::pending();
        let () = pend.await;
        None
    }
}

如果为None的话，将其返回Poll::Pending，则该分支await的时候永远不会等到结果。那么最终的的代码示意如下：

#[tokio::main]
async fn main() -> io::Result<()> {

    let listener: Option<TcpListener> = TcpListener::bind("127.0.0.1:8090").await.ok();
    tokio::select! {
        Some((conn, addr)) = tcp_listen_work(&listener) => {
            println!("accept addr = {:?}", addr);
        }
    }
    Ok(())
}

另一种在反向代理的时候因为server的数量是不定的，所以监听的TcpListener也是不定的，此时我们用Vec<TcpListener>来做表示，那么此时，我们如何通过tokio::select来一次性await所有的accept呢？此时我们借助futures库中的select_all来监听，但是select_all又不允许空的Vec，因为他要返回一个Fut，空的无法返回一个Fut，所以此时我们也要对其进行封装：

async fn multi_tcp_listen_work(listens: &mut Vec<TcpListener>) -> (io::Result<(TcpStream, SocketAddr)>, usize) {
    if !listens.is_empty() {
        let (conn, index, _) = select_all(listens.iter_mut()
                .map(|listener| listener.accept().boxed())).await;
        (conn, index)
    } else {
        let pend = std::future::pending();
        let () = pend.await;
        unreachable!()
    }
}

此时监听从8091-8099，我们的最终代码：

#[tokio::main]
async fn main() -> io::Result<()> {
    let listener: Option<TcpListener> = TcpListener::bind("127.0.0.1:8090").await.ok();
    let mut listeners = vec![];
    for i in 8091..8099 {
        listeners.push(TcpListener::bind(format!("127.0.0.1:{}", i)).await?);
    }
    tokio::select! {
        Some((conn, addr)) = tcp_listen_work(&listener) => {
            println!("accept addr = {:?}", addr);
        }
        (result, index) = multi_tcp_listen_work(&mut listeners) => {
            println!("index receiver = {:?}", index)
        }
    }
    Ok(())
}

如果此时我们用

telnet 127.0.0.1 8098

那么我们就可以看到输出：

index receiver = 7

表示代码已正确的执行。

Rust中数据在多个线程中的共享

Rust中每个对象的所有权都仅只能有一个对象拥有，那么我们数据在在多个地方共享的时候可以怎么办呢？在单线程中，我们可以用use std::rc::Rc;

Rc的特点

单线程的引用计数
不可变引用
非线程安全，即仅能在单线程中使用 Rc引用计数中还有一个弱引用称为Weak，弱引用表示持有对象的一个指针，但是不添加引用计数，也不会影响数据删除，不保证一定能取得到数据。因为其不能修改数据，所以也常用RefCell做配合，来做引用计数的修改。以下是一个父类子类用弱引用计数实现的方案：

use std::rc::Rc;
use std::rc::Weak;
use std::cell::RefCell;

/// 父类拥有者
struct Owner {
    name: String,
    gadgets: RefCell<Vec<Weak<Gadget>>>,
}

/// 子类对象
struct Gadget {
    id: i32,
    owner: Rc<Owner>,
}

fn main() {
    let gadget_owner: Rc<Owner> = Rc::new(
        Owner {
            name: "wmproxy".to_string(),
            gadgets: RefCell::new(vec![]),
        }
    );
    
    // 生成两个小工具
    let gadget1 = Rc::new(
        Gadget {
            id: 1,
            owner: Rc::clone(&gadget_owner),
        }
    );
    let gadget2 = Rc::new(
        Gadget {
            id: 2,
            owner: Rc::clone(&gadget_owner),
        }
    );

    {
        let mut gadgets = gadget_owner.gadgets.borrow_mut();
        gadgets.push(Rc::downgrade(&gadget1));
        gadgets.push(Rc::downgrade(&gadget2));
    }

    for gadget_weak in gadget_owner.gadgets.borrow().iter() {
        let gadget = gadget_weak.upgrade().unwrap();
        println!("小工具 {} 的拥有者：{}", gadget.id, gadget.owner.name);
    }
}

因为其并未实现Send函数，所以无法在多线程种传递。在多线程中，我们需要用Arc，但是在Arc获取可变对象的时候有限制，必须他是唯一引用的时候才能修改。

use std::sync::Arc;
fn main() {
    let mut x = Arc::new(3);
    *Arc::get_mut(&mut x).unwrap() = 4;
    assert_eq!(*x, 4);
    
    let _y = Arc::clone(&x);
    assert!(Arc::get_mut(&mut x).is_none());
}

所以我们在多线程中的引用需要修改的时候，通常会用Atomic或者Mutex来做数据的写入的唯一性。

#![allow(unused)]
fn main() {
    use std::sync::{Arc, Mutex};
    use std::thread;
    use std::sync::mpsc::channel;
    
    const N: usize = 10;
    
    let data = Arc::new(Mutex::new(0));
    
    let (tx, rx) = channel();
    for _ in 0..N {
        let (data, tx) = (Arc::clone(&data), tx.clone());
        thread::spawn(move || {
            // 共享数据data，保证在线程中只会同时有一个对象拥有修改权限，也相当于拥有所有权，10个线程，每个线程+1，最终结果必须等于10
            let mut data = data.lock().unwrap();
            *data += 1;
            if *data == N {
                tx.send(()).unwrap();
            }
        });
    }
    rx.recv().unwrap();
    assert!(*data.lock().unwrap() == 10);
}