RUST中编写简单TCP客户端和服务器代码

在本教程中，我们将学习如何使用 Rust 仅使用标准库编写一个简单的 netcat 客户端和服务器。 netcat 客户端就像网络领域的瑞士军刀。它类似于 PuTTY 和 telnet。您可以使用它连接到服务器并发送和接收数据。我们将创建一个既可以充当客户端又可以充当服务器的应用程序。

功能特点

  1. 我们的客户端将允许用户键入消息并将其发送到任何 TCP 套接字服务器，并在无限循环中显示服务器的响应。
  2. 我们的服务器将侦听来自客户端的传入 TCP 连接，显示来自客户端的消息，并将响应发送回客户端。

您可以在此处找到本教程的完整源代码 

依赖：

# Command line argument parsing.
clap = { version = "4.4.13", features = ["derive"] }

# Pretty logging.
femme = { version = "2.2.1" }
log = { version = "0.4.20" }

# Colorization and ANSI escape sequence codes.
r3bl_tui = { version = "0.5.1" }
r3bl_ansi_color = { version = "0.6.9" }


创建客户端
从最简单的客户端开始。

  • 我们将使用 std::net::TcpStream 创建一个 TCP 套接字客户端。
  • 为了建立 TCP 连接，我们需要一个 IP 地址和端口。 

fn main() {
    println!("Welcome to rtelnet");

    let cli_arg = CLIArg::parse();
    let address = cli_arg.address;
    let port = cli_arg.port;
    let socket_address = format!("{}:{}", address, port);

    if !cli_arg.log_disable {
        femme::start()
    }

    match match cli_arg.subcommand {
        CLISubcommand::Server => start_server(socket_address),
        CLISubcommand::Client => start_client(socket_address),
    } {
        Ok(_) => {
            println!("Program exited successfully");
        }
        Err(error) => {
            println!("Program exited with an error: {}", error);
        }
    }
}

执行客户端逻辑的函数如下所示。

fn start_client(socket_address: String) -> IOResult<()> {
    log::info!("Start client connection");
    let tcp_stream = TcpStream::connect(socket_address)?;
    let (mut reader, mut writer) = (BufReader::new(&tcp_stream), BufWriter::new(&tcp_stream));

    // Client loop.
    loop {
        // Read user input.
        let outgoing = {
            let mut it = String::new();
            let _ = stdin().read_line(&mut it)?;
            it.as_bytes().to_vec()
        };

        // Tx user input to writer.
        let _ = writer.write(&outgoing)?;
        writer.flush()?;

        // Rx response from reader.
        let incoming = {
            let mut it = vec![];
            let _ = reader.read_until(b'\n', &mut it);
            it
        };

        let display_msg = String::from_utf8_lossy(&incoming);
        let display_msg = display_msg.trim();

        let reset = SgrCode::Reset.to_string();
        let display_msg = format!("{}{}", display_msg, reset);
        println!("{}", display_msg);

        // Print debug.
        log::info!(
            "-> Tx: '{}', size: {} bytes{}",
            String::from_utf8_lossy(&outgoing).trim(),
            outgoing.len(),
            reset,
        );
        log::info!(
            "<- Rx: '{}', size: {} bytes{}",
            String::from_utf8_lossy(&incoming).trim(),
            incoming.len(),
            reset,
        );
    }
}

  • 我们为从 TcpStream::connect() 获得的 TcpStream 创建 BufReader 和 BufWriter。这是因为我们希望以块为单位读写数据，而不是一个字节一个字节地读写，这不仅是为了提高性能，也是为了简化逻辑。通过这两个结构体，我们可以非常方便地读写以新行（\n）分隔的数据块。
  • 有一个客户端循环会一直运行下去。这是因为我们想让客户端永远运行下去，这样用户就可以输入信息并将其发送到服务器，然后收到服务器的响应。
  • 如何退出客户端的无限循环？只有当用户按下 Ctrl+C 时，客户端才会退出。Rust 的默认行为是在这种情况下退出进程。这将中断 TCP 连接，导致服务器也退出。
  • 当我们从用户输入读取数据时，也会使用流，但不是 TcpStream，而是 stdin() 流。它的行为与 TcpStream 流非常相似。我们可以用新行（\n）分隔成块来读取数据。一旦用户键入一条信息并按下回车键，这条信息，例如"hi "和新行都会存储在 it 变量中，例如"hi/n"。然后，我们将字符串转换为字节数组，例如：[104, 105, 10]，再将其转换为 Vec<u8>。然后将其发送到服务器。出于 IO 性能考虑，我们必须调用 flush()，因为 BufWriter 会对数据进行缓冲，直到我们调用 flush() 才会将其发送到服务器。它将数据排成队列，分块发送，而不是一次发送一个字节。
  • 从服务器读取响应类似于我们已经看到的从 stdin() 读取响应。主线程会阻塞，直到有数据可以从服务器读取。或者 TCP 连接以任何方式出错（超时或以各种方式关闭）。如果出现错误，则该函数返回错误，主线程退出。请注意，start_client() 函数本身返回一个 IOResult，它只是 pub 类型 IOResult<T> = std::io::Result<T>; 的类型别名。错误处理非常简单。如果出现错误，我们会打印出来并退出程序。

我们使用 reader.read_until(b'\n', &mut it);) 将数据从服务器读入传入变量。这是因为我们希望服务器发送给我们的数据是以新行（\n）结束的。因此，我们读取数据，直到遇到新行。这是一个阻塞调用，所以主线程会阻塞，直到有数据可以从服务器读取。请注意，\n 包含在传入变量中，就像在 stdin() 中一样。

  • 我们使用函数 String::from_utf8_lossy(&incoming); 将 incoming.Vec<u8> 转换为字符串：Vec<u8> 转换为字符串。我们在 String 上调用 .trim()，这样尾部的 \n 就会被去掉。
  • 请注意，trim() 返回的是一个 &str，因此如果要将其转换为字符串，必须通过这个表达式 format!("{}", String::from_utf8_lossy(&incoming).trim()) 函数运行。

这是一个教学示例，这种算法在某种程度上是为了演示如何在客户端和服务器之间来回发送字节，并让它们以某种方式解释这些字节。这种 "舞蹈 "的更正式版本称为 "协议"，如 HTTP、SMTP 等。

在循环的最后一步，从服务器读取输入数据后，我们将其打印到终端。由于服务器会向我们发送 ANSI 转义序列代码，对我们打印到终端的文本进行着色，因此我们希望在打印文本后重置颜色，以免污染我们的 stdout() 输出流。我们使用 SgrCode::Reset 代码来重置打印到终端的文本颜色。

创建服务器 
现在让我们创建服务器。我们将使用 std::net::TcpListener 创建一个 TCP 套接字服务器。为了建立 TCP 连接，我们需要一个 IP 地址和端口。

服务器代码与客户端代码非常相似。我们需要一个永远运行的服务器循环，首先读取（阻塞直到有可用数据），然后以新行（\n）为分界写入数据块。当读取器（又称输入 TCP 流）上没有数据可用时，就会出现 EOF，然后我们就会跳出这个循环并退出。当有数据输入时（以 \n 为分界），我们将对其进行处理，并向客户端发送响应。在处理这些数据时，我们会对其应用 lolcat 特效，因此客户端将获得一个非常丰富多彩的版本，就像他们发送到服务器的文本信息一样。

在实现服务器时，我们还会看到一件事，那就是必须产生多个线程来处理每个传入的客户端连接。客户端是单线程应用程序，而服务器则是多线程应用程序。客户端只需要处理单个 TCP 连接，而服务器则需要处理多个 TCP 连接，每个连接都来自运行货物运行客户端的不同客户端进程，并创建一个新的操作系统进程。幸运的是，Rust 从一开始就为无畏的并发性和并行性而生。

下面是我们服务器应用程序的主要功能：

pub fn start_server(socket_address: String) -> IOResult<()> {
    let tcp_listener = TcpListener::bind(socket_address)?;
    // Server connection accept loop.
    loop {
        log::info!("Waiting for a incoming connection...");
        let (tcp_stream, ..) = tcp_listener.accept()?; // This is a blocking call.

        // Spawn a new thread to handle this connection.
        thread::spawn(|| match handle_connection(tcp_stream) {
            Ok(_) => {
                log::info!("Successfully closed connection to client...");
            }
            Err(_) => {
                log::error!("Problem with client connection...");
            }
        });
    }
}

下面是关于服务器代码的一些注意事项：

我们正在使用 IOResult，就像客户端代码一样。我们经常调用 ? 操作符，它是在结果上进行匹配并在出现错误时提前返回的简写。

这只是最基本的错误处理，对于这个教学示例来说已经足够了。

请注意，即使在这个教学示例中，我们也没有使用 unwrap() 方法，因为一旦出现错误，就会引起panic恐慌。

我们总是使用 ? 操作符，如果出现错误，它会提前返回。

养成在测试之外使用 unwrap() 的习惯并不是一个好主意。这些习惯一旦养成就很难改掉。

你甚至可以在项目的顶层模块中添加 #![warn(clippy::unwrap_in_result)]，如果在测试之外使用 unwrap()，编译器就会发出警告。下面是一个例子example。

服务器要做的第一件事就是在给定的地址上预留一个端口。这叫做绑定，我们使用 TcpListener::bind(socket_address)?; 来完成。这还不能启动服务器。它只是在给定地址上保留一个端口，假设它是可用的。如果其他进程已经绑定了该端口，则会返回错误信息。

有了 TcpListener 实例后，我们就可以调用 accept() 开始监听传入的连接。这是一个阻塞调用，因此主线程会阻塞，直到有连接进入。一旦有传入连接，我们就会得到一个 TcpStream 实例，可以用它向客户端读写数据。这是一个阻塞调用。这意味着主线程在等待连接进入时，将无法做其他事情，比如处理其他传入连接。

这就是为什么我们要使用 thread::spawn() 来创建一个新线程，并由它来处理传入的连接。我们为每个传入连接生成一个新线程。这不是一个可扩展的解决方案，但对于这个教学示例来说已经足够好了。

现在，让我们来看看生成线程调用的 handle_connection() 函数。该函数用于处理从客户端传入的连接。它与客户端代码一起定义了我们的 "协议"。我们没有使用任何正式的协议，如 HTTP 或 SMTP。我们只是在客户端和服务器之间来回发送字节，并以某种方式解释它们，这就是我们的非正式协议。这段代码与客户端代码非常相似，包括循环、BufReader 和 BufWriter 结构。甚至还在寻找 EOF 以跳出循环。只不过，我们不会阻塞 stdin() 的输入。

fn handle_connection(tcp_stream: TcpStream) -> IOResult<()> {
    log::info!("Start handle connection");

    let reader = &mut BufReader::new(&tcp_stream);
    let write = &mut BufWriter::new(&tcp_stream);

    // Process client connection loop.
    loop {
        let mut incoming: Vec<u8> = vec![];

        // Read from reader.
        let num_bytes_read = reader.read_until(b'\n', &mut incoming)?;

        // Check for EOF. The stream is closed.
        if num_bytes_read == 0 {
            break;
        }

        // Process.
        let outgoing = process(&incoming);

        // Write to writer.
        write.write(&outgoing)?;
        let _ = write.flush()?;

        // Print debug.
        log::info!("-> Rx(bytes) : {:?}", &incoming);
        log::info!(
            "-> Rx(string): '{}', size: {} bytes",
            String::from_utf8_lossy(&incoming).trim(),
            incoming.len(),
        );
        log::info!(
            "<- Tx(string): '{}', size: {} bytes",
            String::from_utf8_lossy(&outgoing).trim(),
            outgoing.len()
        );
    }

    log::info!("End handle connection - connection closed");

    Ok(())
}


最后，让我们来看看 process() 函数，它将输入的字节转换为输出的字节。在这里，我们为应用程序添加了一些乐趣、色彩和亮点。我们用萝莉猫特效为传入字节着色，然后将其发送回客户端。

use r3bl_tui::ColorWheel;

fn process(incoming: &Vec<u8>) -> Vec<u8> {
    // 将输入字符串转换为字符串，并删除所有尾部空白（包括换行符）。
    let incoming = String::from_utf8_lossy(incoming);
    let incoming = incoming.trim();

    // Prepare outgoing payload.
    let outgoing = incoming.to_string();

    // Colorize it w/ a gradient.
    let outgoing = ColorWheel::lolcat_into_string(&outgoing);

    // 生成输出响应。在输出末尾添加换行符（以便客户端处理）。
    let outgoing = format!("{}\n", outgoing);

    // Return outgoing payload.
    outgoing.as_bytes().to_vec()
}