Read all data from TCP stream in Rust

One way to do this, is:

1. Set a read timeout on the TcpStream using set_read_timeout. This avoids a hanging client in case the server has stopped sending but left the connection open.
2. Read from the stream in a loop, using a BufReader, and aggregate the read bytes.
3. Break out of the loop in case the read timed out or the server closed the connection.

Here's some example code:

use std::{
    io::{self, BufReader, LineWriter, Read, Write},
    net::TcpStream,
    time::Duration,
};

fn main() -> io::Result<()> {
    println!("Start of program");

    // This is just to get data via TCP, we could of course parse the HTTP header for
    // the message length
    let server_host = "www.neverssl.com";
    let server_port = "80";

    let stream = TcpStream::connect(server_host.to_string() + ":" + server_port)?;

    let mut reader = BufReader::new(&stream);
    let mut writer = LineWriter::new(&stream);

    // Don't block indefinitely on `reader.read` when there's no data to read
    let max_read_time = Duration::from_millis(200);
    stream.set_read_timeout(Some(max_read_time))?;

    writer.write_all(format!("GET / HTTP/1.1\r\nhost: {server_host}\r\n\r\n").as_bytes())?;

    let mut total_bytes_read = vec![];
    let bytes_to_read_per_attempt = 1024;
    let mut read_attempt_nr = 0;
    loop {
        read_attempt_nr += 1;
        println!("Read attempt nr {read_attempt_nr}");
        let mut cur_buffer = vec![0; bytes_to_read_per_attempt];

        // If the reader has no data but the server hasn't closed the connection,
        // by default `reader.read()` would block until the server closes the connection.
        // Hence, we need `stream.set_read_timeout`.
        let nr_of_bytes_read = match reader.read(&mut cur_buffer) {
            Ok(nr_of_bytes_read) => nr_of_bytes_read,
            Err(err) => {
                if err.kind() == io::ErrorKind::WouldBlock || err.kind() == io::ErrorKind::TimedOut
                {
                    println!("Read attempt timed out");
                    break;
                } else {
                    return Err(err);
                }
            }
        };

        // Reading zero bytes indicates that the server closed the connection
        if nr_of_bytes_read == 0 {
            println!("Read zero bytes → Connection seems closed");
            break;
        }

        // Remove the excess null bytes at the end of cur_buffer
        cur_buffer.truncate(nr_of_bytes_read);
        total_bytes_read.append(&mut cur_buffer);

        println!("Read {nr_of_bytes_read} bytes in attempt nr {read_attempt_nr}");
    }
    let response = String::from_utf8_lossy(&total_bytes_read).to_string();
    let nr_of_bytes_received = total_bytes_read.len();

    println!("Server response : {response}");
    println!("{nr_of_bytes_received} bytes received");

    Ok(println!("End of program"))
}

This client assumes that if it hasn't received data from the server for 200 milliseconds (while the connection is still open), that the server's message is complete.

posted about 2 years ago

CC BY-SA 4.0

2y ago

Matthias Braun‭

177 reputation 6 3 24 12