30

I have a struct with a BufStream<T> where T: Read+Write. The BufStream can be a TcpStream and I'd like to read n bytes from it. Not a fixed amount of bytes in a predefined buffer, but I have a string/stream which indicates the number of bytes to read next.

Is there a nice way to do that?

Shepmaster
  • 388,571
  • 95
  • 1,107
  • 1,366
buster
  • 1,071
  • 1
  • 10
  • 15

2 Answers2

34

Since Rust 1.6, Read::read_exact can be used to do this. If bytes_to_read is the number of bytes you need to read, possibly determined at runtime, and reader is the stream to read from:

let mut buf = vec![0u8; bytes_to_read];
reader.read_exact(&mut buf)?;

The part that wasn't clear to me from the read_exact documentation was that the target buffer can be a dynamically-allocated Vec.

Thanks to the Rust Gitter community for pointing me to this solution.

Shepmaster
  • 388,571
  • 95
  • 1,107
  • 1,366
rspeer
  • 3,539
  • 2
  • 25
  • 25
  • It's also not clear what happens if EOF is returned before bytes_to_read. – Andy Hayden Sep 05 '18 at 03:01
  • 3
    That's specified in the documentation: "If this function encounters an "end of file" before completely filling the buffer, it returns an error of the kind `ErrorKind::UnexpectedEof`." – rspeer Sep 06 '18 at 18:24
  • True, so I guess you can't actually do this. In the sense that it's not defined what happens to the buffer contents/impossible to recover. – Andy Hayden Sep 06 '18 at 18:48
  • How can you use a Vec as the target, when I attempt to do this with a `Vec::with_capacity(bytes_to_read)` the vec "length" is zero and nothing gets read – NathanFrasier Oct 11 '19 at 13:20
  • 3
    @Nanos Hopefully you've already figured this out elsewhere, but, the answer is you need to fill the `Vec` with some value first, (or unsafely set the length I guess). The `vec![0u8; bytes_to_read]` used above fills the `Vec` with `bytes_to_read` zeroes. – Ryan1729 Nov 03 '19 at 13:42
  • Same problem. `Vec::with_capacity` won't allocate until you append to it, so if you try to `read_exact` into it, it'll read 0 bytes. You can either call `vec.resize(bytes_to_read, 0);` or use the `vec!` macro - the macro is faster. – rjh May 31 '23 at 11:18
24

It sounds like you want Read::take and Read::read_to_end.

This will allow you to read data into a &mut Vec<u8>, which is useful when you want to reuse an existing buffer or don't have an appropriately sized slice already. This allows you to avoid initializing the data with dummy values before overwriting them with the newly-read information:

use std::{
    io::{prelude::*, BufReader},
    str,
};

fn read_n<R>(reader: R, bytes_to_read: u64) -> Vec<u8>
where
    R: Read,
{
    let mut buf = vec![];
    let mut chunk = reader.take(bytes_to_read);
    // Do appropriate error handling for your situation
    // Maybe it's OK if you didn't read enough bytes?
    let n = chunk.read_to_end(&mut buf).expect("Didn't read enough");
    assert_eq!(bytes_to_read as usize, n);
    buf
}

fn main() {
    let input_data = b"hello world";
    let mut reader = BufReader::new(&input_data[..]);

    let first = read_n(&mut reader, 5);
    let _ = read_n(&mut reader, 1);
    let second = read_n(&mut reader, 5);

    println!(
        "{:?}, {:?}",
        str::from_utf8(&first),
        str::from_utf8(&second)
    );
}

If you are worried that Read::take consumes the reader by reference, note that take comes from Read and Read is implemented for any mutable reference to a type that implements Read. You can also use Read::by_ref to create this mutable reference.

See also:

Shepmaster
  • 388,571
  • 95
  • 1,107
  • 1,366
  • 1
    Thanks.. doesn't take() consume the reader? My usecase is that i have a tcpstream (part of a struct) from which i read lines. Sometimes in between i don't need to read a line but a specific amount of bytes. Afterwards i'll need to read lines again... – buster May 23 '15 at 14:47
  • 2
    It does, but "the reader" can be a `&mut R` i.e a temporary mutable reference to another reader. Take uses a by-value API so that all use cases are supported. This model appears in some different corners of rust. – bluss May 23 '15 at 14:49
  • @buster as bluss points out, `&mut R` also implements the `Read` trait. My example uses that to good effect - that's why I pass in `&mut reader` in the main method of the example. – Shepmaster May 23 '15 at 14:50
  • @buster @Shepmaster how did you guys figure out that `Read::take` also accepts a `&mut reader`? Wondering how I would be able to figure that out myself next time. :) – Kaeros Sep 05 '19 at 20:00
  • 2
    @Kaeros [`take`](https://doc.rust-lang.org/std/io/trait.Read.html#method.take) comes from [`Read`](https://doc.rust-lang.org/std/io/trait.Read.html) and [`Read` is implemented for any mutable reference to a type that implements `Read`](https://doc.rust-lang.org/1.37.0/std/io/trait.Read.html#impl-Read-12). See also [Whats the idiomatic way reference BufReader/BufWriter when passing between functions?](https://stackoverflow.com/q/39464237/155423); – Shepmaster Sep 09 '19 at 18:11
  • Also `Read` has [`by_ref`](https://doc.rust-lang.org/std/io/trait.Read.html#method.by_ref) which specifically creates a sub-reader by reference. So even without knowing that you can just create a mutable reference and that's also `Read`, you can always call `reader.by_ref()` and pass *that* in. – Masklinn Apr 07 '21 at 09:44
  • The `buf` should be created with `.with_capacity(bytes_to_read)`, to avoid memory allocations. – Francesco Pasa Jul 24 '23 at 20:05