24、Rust 中 TCP 和 UDP 编程及 TCP 反向代理实现

Rust 中 TCP 和 UDP 编程及 TCP 反向代理实现

1. 引言

IP 地址和套接字是使用 Rust 标准库进行网络编程的基础数据结构。接下来，我们将学习如何使用 Rust 编写基于 TCP 和 UDP 协议进行通信的程序。

2. 创建项目

首先，创建一个名为 tcpudp 的新项目，用于编写 TCP 和 UDP 的服务器与客户端代码：

cargo new tcpudp && cd tcpudp

3. UDP 通信编程

3.1 UDP 服务器编写

UDP 服务器通过 UdpSocket::bind 绑定到本地套接字，创建固定大小的缓冲区，并在循环中监听传入的数据流。如果接收到数据，会生成一个新线程将数据回显给发送者。以下是 udp-server.rs 的代码：

use std::str;
use std::thread;

fn main() {
    let socket = UdpSocket::bind("127.0.0.1:3000").expect(
        "Unable to bind to port");
    let mut buffer = [0; 1024];
    loop {
        let socket_new = socket.try_clone().expect(
            "Unable to clone socket");
        match socket_new.recv_from(&mut buffer) {
            Ok((num_bytes, src_addr)) => {
                thread::spawn(move || {
                    let send_buffer = &mut 
                        buffer[..num_bytes];
                    println!(
                        "Received from client:{}",
                        str::from_utf8(
                            send_buffer).unwrap()
                    );
                    let response_string =
                        format!("Received this: {}", 
                            String::from_utf8_lossy(
                            send_buffer));
                    socket_new
                        .send_to(&response_string
                            .as_bytes(), &src_addr)
                        .expect("error in sending datagram 
                            to remote socket");
                });
            }
            Err(err) => {
                println!("Error in receiving datagrams over 
                    UDP: {}", err);
            }
        }
    }
}

3.2 UDP 客户端编写

UDP 客户端先绑定到本地端口，允许操作系统选择临时 IP 地址和端口发送数据报，然后尝试连接到服务器所在的远程套接字。连接成功后，打印出对等方的套接字地址，并使用 send 方法向远程套接字发送消息。以下是 udp-client.rs 的代码：

use std::net::UdpSocket;
fn main() {
    // Create a local UDP socket
    let socket = UdpSocket::bind("0.0.0.0:0").expect(
        "Unable to bind to socket");
    // Connect the socket to a remote socket
    socket
        .connect("127.0.0.1:3000")
        .expect("Could not connect to UDP server");
    println!("socket peer addr is {:?}", 
        socket.peer_addr());
    // Send a datagram to the remote socket
    socket
        .send("Hello: sent using send() call".as_bytes())
        .expect("Unable to send bytes");
}

3.3 运行 UDP 服务器和客户端

运行 UDP 服务器：

cargo run --bin  udp-server

在另一个终端运行 UDP 客户端：

cargo run --bin  udp-client

此时，服务器将接收到客户端发送的消息。

4. TCP 通信编程

4.1 TCP 服务器编写

TCP 服务器使用 TcpListener::bind 创建一个监听套接字，通过 incoming 方法获取传入连接的迭代器。对于每个连接，读取数据并打印，然后将数据回显给客户端。以下是 tcp-server.rs 的代码：

use std::io::{Read, Write};
use std::net::TcpListener;
fn main() {
    let connection_listener = TcpListener::bind(
        "127.0.0.1:3000").unwrap();
    println!("Running on port 3000");
    for stream in connection_listener.incoming() {
        let mut stream = stream.unwrap();
        println!("Connection established");
        let mut buffer = [0; 100];
        stream.read(&mut buffer).unwrap();
        println!("Received from client: {}", 
            String::from_utf8_lossy(&buffer));
        stream.write(&mut buffer).unwrap();
    }
}

4.2 TCP 客户端编写

TCP 客户端使用 TcpStream::connect 连接到服务器所在的远程套接字，向服务器发送数据，然后读取服务器的响应。以下是 tcp-client.rs 的代码：

use std::io::{Read, Write};
use std::net::TcpStream;
use std::str;
fn main() {
    let mut stream = TcpStream::connect(
        "localhost:3000").unwrap();
    let msg_to_send = "Hello from TCP client";
    stream.write(msg_to_send.as_bytes()).unwrap();
    let mut buffer = [0; 200];
    stream.read(&mut buffer).unwrap();
    println!(
        "Got echo back from server:{:?}",
        str::from_utf8(&buffer)
            .unwrap()
            .trim_end_matches(char::from(0))
    );
}

4.3 运行 TCP 服务器和客户端

运行 TCP 服务器：

cargo run --bin  tcp-server

在另一个终端运行 TCP 客户端：

cargo run --bin  tcp-client

客户端发送的消息将被服务器接收并回显。

5. TCP 反向代理项目

5.1 代理服务器概述

代理服务器是在互联网上跨多个网络导航时使用的中间软件服务，分为正向代理和反向代理。正向代理为客户端向互联网发出请求提供中介服务，反向代理为服务器提供中介服务。我们将重点实现反向代理的核心功能，即把客户端的请求转发到后端源服务器，并将源服务器的响应路由回请求客户端。

5.2 创建项目

创建一个名为 tcpproxy 的新项目，用于编写源服务器和代理服务器的代码：

cargo new tcpproxy && cd tcpproxy

创建两个文件： tcpproxy/src/bin/origin.rs 和 tcpproxy/src/bin/proxy.rs 。

5.3 源服务器编写

源服务器是一个 TCP 服务器，理解有限的 HTTP 语义，它将完成以下任务：
- 接收传入的 HTTP 请求。
- 提取请求的第一行（HTTP 请求行）。
- 接受特定路由上的 GET HTTP 请求（例如， /order/status/1 ）。
- 返回订单状态。

以下是源服务器的代码实现：

// tcpproxy/src/bin/origin.rs
use std::io::{Read, Write};
use std::net::TcpListener;
use std::net::{IpAddr, Ipv4Addr, SocketAddr};
use std::str;
use std::str::FromStr;
use std::string::ParseError;

#[derive(Debug)]
struct RequestLine {
    method: Option<String>,
    path: Option<String>,
    protocol: Option<String>,
}

impl RequestLine {
    fn method(&self) -> String {
        if let Some(method) = &self.method {
            method.to_string()
        } else {
            String::from("")
        }
    }
    fn path(&self) -> String {
        if let Some(path) = &self.path {
            path.to_string()
        } else {
            String::from("")
        }
    }
    fn get_order_number(&self) -> String {
        let path = self.path();
        let path_tokens: Vec<String> = path.split("/").map(
            |s| s.parse().unwrap()).collect();
        path_tokens[path_tokens.len() - 1].clone()
    }
}

impl FromStr for RequestLine {
    type Err = ParseError;

    fn from_str(msg: &str) -> Result<Self, Self::Err> {
        let mut msg_tokens = msg.split_ascii_whitespace();
        let method = match msg_tokens.next() {
            Some(token) => Some(String::from(token)),
            None => None,
        };
        let path = match msg_tokens.next() {
            Some(token) => Some(String::from(token)),
            None => None,
        };
        let protocol = match msg_tokens.next() {
            Some(token) => Some(String::from(token)),
            None => None,
        };

        Ok(Self {
            method: method,
            path: path,
            protocol: protocol,
        })
    }
}

fn main() {
    // Start the origin server
    let port = 3000;
    let socket_addr = SocketAddr::new(IpAddr::V4(
        Ipv4Addr::new(127, 0, 0, 1)), port);
    let connection_listener = TcpListener::bind(
        socket_addr).unwrap();

    println!("Running on port: {}", port);

    for stream in connection_listener.incoming() {
        // Read the first line of incoming HTTP request 
        // and convert it into RequestLine struct
        let mut stream = stream.unwrap();
        let mut buffer = [0; 200];
        stream.read(&mut buffer).unwrap();
        let req_line = "";
        let string_request_line =
            if let Some(line) = str::from_utf8(
                &buffer).unwrap().lines().next() {
                line
            } else {
                println!("Invalid request line received");
                req_line
            }; 
        let req_line = RequestLine::from_str(
            string_request_line).unwrap();

        // Construct the HTTP response string 
        let html_response_string;
        let order_status;

        println!("len is {}", req_line.get_order_number()
            .len());

        if req_line.method() != "GET"
            || !req_line.path().starts_with(
               "/order/status")
            || req_line.get_order_number().len() == 0
        {
            if req_line.get_order_number().len() == 0 {
                order_status = format!("Please provide 
                    valid order number");
            } else {
                order_status = format!("Sorry,this page is 
                    not found");
            }

            html_response_string = format!(
                "HTTP/1.1 404 Not Found\nContent-Type: 
                    text/html\nContent-Length:{}\n\n{}",
                order_status.len(),
                order_status
            );
        } else {
            order_status = format!(
                "Order status for order number {} is: 
                    Shipped\n",
                req_line.get_order_number()
            );
            html_response_string = format!(
                "HTTP/1.1 200 OK\nContent-Type: 
                    text/html\nContent-Length:{}\n\n{}",
                order_status.len(),
                order_status
            );
        }

        stream.write(html_response_string.as_bytes()).unwrap();
    }
}

运行源服务器：

cargo run --bin origin

在浏览器中输入以下 URL 进行测试：
- localhost:3000/order/status/2 ，应显示 Order status for order number 2 is: Shipped 。
- localhost:3000/invalid/path ，应显示 Sorry, this page is not found 。
- localhost:3000/order/status/ ，应显示 Please provide valid order number 。

以下是 UDP 和 TCP 通信的流程对比表格：
| 通信协议 | 服务器操作步骤 | 客户端操作步骤 |
| ---- | ---- | ---- |
| UDP | 1. 绑定本地套接字
2. 创建缓冲区
3. 循环监听数据
4. 接收数据后生成新线程回显 | 1. 绑定本地端口
2. 连接远程套接字
3. 打印对等方地址
4. 发送消息 |
| TCP | 1. 绑定监听套接字
2. 获取连接迭代器
3. 读取数据
4. 回显数据 | 1. 连接远程套接字
2. 发送数据
3. 读取响应 |

以下是源服务器处理请求的 mermaid 流程图：

graph TD;
    A[接收连接] --> B[读取请求第一行];
    B --> C[转换为 RequestLine 结构体];
    C --> D{是否为 GET 请求};
    D -- 否 --> E[构造 404 响应];
    D -- 是 --> F{路径是否以 /order/status 开头};
    F -- 否 --> E;
    F -- 是 --> G{是否提供订单号};
    G -- 否 --> E;
    G -- 是 --> H[构造 200 响应];
    E --> I[发送响应];
    H --> I;

Rust 中 TCP 和 UDP 编程及 TCP 反向代理实现

5.4 反向代理服务器编写

反向代理服务器的代码将接收客户端请求，并将其转发到源服务器，然后将源服务器的响应返回给客户端。以下是反向代理服务器的代码实现：

// tcpproxy/src/bin/proxy.rs
use std::env;
use std::io::{Read, Write};
use std::net::{TcpListener, TcpStream};
use std::process::exit;
use std::thread;

fn handle_connection(proxy_stream: &mut TcpStream, 
    origin_stream: &mut TcpStream) {
    let mut in_buffer: Vec<u8> = vec![0; 200];
    let mut out_buffer: Vec<u8> = vec![0; 200];
    // Read incoming request to proxy_stream
    if let Err(err) = proxy_stream.read(&mut in_buffer) {
        println!("Error in reading from incoming proxy 
            stream: {}", err);
    } else {
        println!(
            "1: Incoming client request: {}",
            String::from_utf8_lossy(&in_buffer)
        );
    }
    // Write the byte stream to origin_stream
    let _ = origin_stream.write(&mut in_buffer).unwrap();
    println!("2: Forwarding request to origin server\n");
    // Read response from the backend server
    let _ = origin_stream.read(&mut out_buffer).unwrap();
    println!(
        "3: Received response from origin server: {}",
        String::from_utf8_lossy(&out_buffer)
    );
    // Write response back to the proxy client
    let _ = proxy_stream.write(&mut out_buffer).unwrap();
    println!("4: Forwarding response back to client");
}

fn main() {
    // Accept command-line parameters for proxy_stream and 
    // origin_stream
    let args: Vec<_> = env::args().collect();
    if args.len() < 3 {
        eprintln!("Please provide proxy-from and proxy-to 
            addresses");
        exit(2);
    }
    let proxy_server = &args[1];
    let origin_server = &args[2];
    // Start a socket server on proxy_stream
    let proxy_listener;
    if let Ok(proxy) = TcpListener::bind(proxy_server) {
        proxy_listener = proxy;
        let addr = proxy_listener.local_addr()
            .unwrap().ip();
        let port = proxy_listener.local_addr().unwrap()
            .port();
        if let Err(_err) = TcpStream::connect(
            origin_server) {
            println!("Please re-start the origin server");
            exit(1);
        }
        println!("Running on Addr:{}, Port:{}\n", addr, 
            port);
    } else {
        eprintln!("Unable to bind to specified proxy 
            port");
        exit(1);
    }
    // Listen for incoming connections from proxy_server 
    // and read byte stream
    let mut thread_handles = Vec::new();
    for proxy_stream in proxy_listener.incoming() {
        let mut proxy_stream = proxy_stream.expect("Error 
            in incoming TCP connection");
        // Establish a new TCP connection to origin_stream
        let mut origin_stream =
            TcpStream::connect(origin_server).expect(
                "Please re-start the origin server");
        let handle =
            thread::spawn(move || handle_connection(&mut  
                proxy_stream, &mut origin_stream));
        thread_handles.push(handle);
    }
    for handle in thread_handles {
        handle.join().expect("Unable to join child 
            thread");
    }
}

运行反向代理服务器时，需要提供两个命令行参数，分别对应反向代理服务器和源服务器的套接字地址。运行命令如下：

cargo run --bin proxy localhost:3001 localhost:3000

在浏览器中输入以下 URL 进行测试：
- localhost:3001/order/status/2 ，应显示 Order status for order number 2 is: Shipped 。
- localhost:3001/invalid/path ，应显示 Sorry, this page is not found 。
- localhost:3001/order/status/ ，应显示 Please provide valid order number 。

以下是反向代理服务器处理请求的步骤列表：
1. 接受命令行参数，分别为反向代理服务器和源服务器的套接字地址。
2. 启动反向代理服务器的套接字监听。
3. 尝试连接源服务器，如果失败则输出错误信息并退出。
4. 监听客户端的连接请求，对于每个连接，生成一个新线程处理。
5. 在新线程中，读取客户端请求，将其转发到源服务器。
6. 读取源服务器的响应，将其返回给客户端。

以下是反向代理服务器处理请求的 mermaid 流程图：

graph TD;
    A[启动监听] --> B[接受客户端连接];
    B --> C[连接源服务器];
    C --> D[读取客户端请求];
    D --> E[转发请求到源服务器];
    E --> F[读取源服务器响应];
    F --> G[返回响应给客户端];
    G --> B;

总结

通过上述内容，我们学习了如何使用 Rust 进行 UDP 和 TCP 通信编程，以及如何实现一个简单的 TCP 反向代理服务器。具体包括：
- 掌握了 UDP 服务器和客户端的编写方法，包括绑定套接字、发送和接收数据。
- 学会了 TCP 服务器和客户端的编写，包括监听连接、读取和写入数据。
- 理解了反向代理服务器的工作原理，并实现了一个基本的反向代理服务器，能够将客户端请求转发到源服务器，并将响应返回给客户端。

这些知识和代码示例可以帮助我们在实际项目中构建高效、稳定的网络应用程序。同时，我们可以根据实际需求对代码进行扩展和优化，例如添加更多的错误处理、实现负载均衡和缓存等功能。