Skip to main content

dprint_core/communication/
reader_writer.rs

1use std::io::ErrorKind;
2use std::io::Read;
3use std::io::Result;
4use std::io::Write;
5
6const SUCCESS_BYTES: &[u8; 4] = &[255, 255, 255, 255];
7
8pub struct MessageReader<TRead: Read + Unpin> {
9  reader: TRead,
10}
11
12impl<TRead: Read + Unpin> MessageReader<TRead> {
13  pub fn new(reader: TRead) -> Self {
14    Self { reader }
15  }
16
17  /// Reads a u32 value.
18  pub fn read_u32(&mut self) -> Result<u32> {
19    let mut int_buf: [u8; 4] = [0; 4];
20    self.reader.read_exact(&mut int_buf)?;
21    Ok(u32::from_be_bytes(int_buf))
22  }
23
24  /// Reads a u32 value followed by a buffer.
25  pub fn read_sized_bytes(&mut self) -> Result<Vec<u8>> {
26    let size = self.read_u32()? as usize;
27    self.read_bytes(size)
28  }
29
30  #[allow(clippy::read_zero_byte_vec)]
31  pub fn read_bytes(&mut self, size: usize) -> Result<Vec<u8>> {
32    let mut buf = Vec::new();
33    if size > 0 {
34      buf.try_reserve_exact(size)?;
35      unsafe {
36        buf.set_len(size);
37      }
38      self.reader.read_exact(&mut buf)?;
39    }
40    Ok(buf)
41  }
42
43  pub fn read_success_bytes(&mut self) -> Result<()> {
44    let read_bytes = self.inner_read_success_bytes()?;
45    if &read_bytes != SUCCESS_BYTES {
46      let message = format!(
47        "Catastrophic error reading from process. Did not receive the success bytes at end of message. Found: {:?}",
48        read_bytes
49      );
50      Result::Err(std::io::Error::new(ErrorKind::InvalidData, message))
51    } else {
52      Ok(())
53    }
54  }
55
56  fn inner_read_success_bytes(&mut self) -> Result<[u8; 4]> {
57    let mut read_buf: [u8; 4] = [0; 4];
58    self.reader.read_exact(&mut read_buf)?;
59    Ok(read_buf)
60  }
61}
62
63pub struct MessageWriter<TWrite: Write + Unpin> {
64  writer: TWrite,
65}
66
67impl<TWrite: Write + Unpin> MessageWriter<TWrite> {
68  pub fn new(writer: TWrite) -> Self {
69    Self { writer }
70  }
71
72  pub fn send_u32(&mut self, value: u32) -> Result<()> {
73    self.writer.write_all(&value.to_be_bytes())?;
74    Ok(())
75  }
76
77  pub fn send_sized_bytes(&mut self, bytes: &[u8]) -> Result<()> {
78    self.send_u32(bytes.len() as u32)?;
79    if !bytes.is_empty() {
80      self.writer.write_all(bytes)?;
81    }
82    Ok(())
83  }
84
85  pub fn send_success_bytes(&mut self) -> Result<()> {
86    self.writer.write_all(SUCCESS_BYTES)?;
87    self.writer.flush()?;
88    Ok(())
89  }
90
91  pub fn flush(&mut self) -> Result<()> {
92    self.writer.flush()
93  }
94}