dprint_core/communication/
reader_writer.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
use std::io::ErrorKind;
use std::io::Read;
use std::io::Result;
use std::io::Write;

const SUCCESS_BYTES: &[u8; 4] = &[255, 255, 255, 255];

pub struct MessageReader<TRead: Read + Unpin> {
  reader: TRead,
}

impl<TRead: Read + Unpin> MessageReader<TRead> {
  pub fn new(reader: TRead) -> Self {
    Self { reader }
  }

  /// Reads a u32 value.
  pub fn read_u32(&mut self) -> Result<u32> {
    let mut int_buf: [u8; 4] = [0; 4];
    self.reader.read_exact(&mut int_buf)?;
    Ok(u32::from_be_bytes(int_buf))
  }

  /// Reads a u32 value followed by a buffer.
  pub fn read_sized_bytes(&mut self) -> Result<Vec<u8>> {
    let size = self.read_u32()? as usize;
    self.read_bytes(size)
  }

  #[allow(clippy::read_zero_byte_vec)]
  pub fn read_bytes(&mut self, size: usize) -> Result<Vec<u8>> {
    let mut buf = Vec::new();
    if size > 0 {
      buf.try_reserve_exact(size)?;
      unsafe {
        buf.set_len(size);
      }
      self.reader.read_exact(&mut buf)?;
    }
    Ok(buf)
  }

  pub fn read_success_bytes(&mut self) -> Result<()> {
    let read_bytes = self.inner_read_success_bytes()?;
    if &read_bytes != SUCCESS_BYTES {
      let message = format!(
        "Catastrophic error reading from process. Did not receive the success bytes at end of message. Found: {:?}",
        read_bytes
      );
      Result::Err(std::io::Error::new(ErrorKind::InvalidData, message))
    } else {
      Ok(())
    }
  }

  fn inner_read_success_bytes(&mut self) -> Result<[u8; 4]> {
    let mut read_buf: [u8; 4] = [0; 4];
    self.reader.read_exact(&mut read_buf)?;
    Ok(read_buf)
  }
}

pub struct MessageWriter<TWrite: Write + Unpin> {
  writer: TWrite,
}

impl<TWrite: Write + Unpin> MessageWriter<TWrite> {
  pub fn new(writer: TWrite) -> Self {
    Self { writer }
  }

  pub fn send_u32(&mut self, value: u32) -> Result<()> {
    self.writer.write_all(&value.to_be_bytes())?;
    Ok(())
  }

  pub fn send_sized_bytes(&mut self, bytes: &[u8]) -> Result<()> {
    self.send_u32(bytes.len() as u32)?;
    if !bytes.is_empty() {
      self.writer.write_all(bytes)?;
    }
    Ok(())
  }

  pub fn send_success_bytes(&mut self) -> Result<()> {
    self.writer.write_all(SUCCESS_BYTES)?;
    self.writer.flush()?;
    Ok(())
  }

  pub fn flush(&mut self) -> Result<()> {
    self.writer.flush()
  }
}