1use std::io;
28use std::io::Read as _;
29
30use crate::error::Error;
31
32const ASCII_0: u8 = 0x30;
33const ASCII_9: u8 = 0x39;
34
35pub fn read8<R>(reader: &mut R) -> Result<u8, io::Error> where R: io::Read {
36 let mut buf = [0u8; 1];
37 reader.read_exact(&mut buf).and(Ok(buf[0]))
38}
39
40pub fn read16<R>(reader: &mut R) -> Result<u16, io::Error> where R: io::Read {
41 let mut buf = [0u8; 2];
42 reader.read_exact(&mut buf)?;
43 Ok(u16::from_be_bytes(buf))
44}
45
46pub fn read64<R>(reader: &mut R) -> Result<u64, io::Error> where R: io::Read {
47 let mut buf = [0u8; 8];
48 reader.read_exact(&mut buf)?;
49 Ok(u64::from_be_bytes(buf))
50}
51
52pub trait BufReadExt {
53 fn discard_exact(&mut self, len: usize) -> io::Result<()>;
54 fn is_eof(&mut self) -> io::Result<bool>;
55}
56
57impl<T> BufReadExt for T where T: io::BufRead {
58 fn discard_exact(&mut self, mut len: usize) -> io::Result<()> {
59 while len > 0 {
60 let consume_len = match self.fill_buf() {
61 Ok(buf) if buf.is_empty() =>
62 return Err(io::Error::new(
63 io::ErrorKind::UnexpectedEof, "unexpected EOF")),
64 Ok(buf) => buf.len().min(len),
65 Err(e) if e.kind() == io::ErrorKind::Interrupted => continue,
66 Err(e) => return Err(e),
67 };
68 self.consume(consume_len);
69 len -= consume_len;
70 }
71 Ok(())
72 }
73
74 fn is_eof(&mut self) -> io::Result<bool> {
75 loop {
76 match self.fill_buf() {
77 Ok(buf) => return Ok(buf.is_empty()),
78 Err(e) if e.kind() == io::ErrorKind::Interrupted => continue,
79 Err(e) => return Err(e),
80 }
81 }
82 }
83}
84
85pub trait ReadExt {
86 fn read_exact_len(&mut self, buf: &mut Vec<u8>, len: usize)
87 -> io::Result<()>;
88}
89
90impl<T> ReadExt for T where T: io::Read {
91 fn read_exact_len(&mut self, buf: &mut Vec<u8>, len: usize)
92 -> io::Result<()> {
93 if self.take(len as u64).read_to_end(buf)? != len {
99 return Err(io::Error::new(
100 io::ErrorKind::UnexpectedEof, "unexpected EOF"));
101 }
102 Ok(())
103 }
104}
105
106pub fn atou16(bytes: &[u8]) -> Result<u16, Error> {
108 debug_assert!(bytes.len() <= 4);
109 if bytes.len() == 0 {
110 return Err(Error::InvalidFormat("Not a number"));
111 }
112 let mut n = 0;
113 for &c in bytes {
114 if c < ASCII_0 || ASCII_9 < c {
115 return Err(Error::InvalidFormat("Not a number"));
116 }
117 n = n * 10 + (c - ASCII_0) as u16;
118 }
119 Ok(n)
120}
121
122pub fn ctou32(c: u8) -> Result<u32, Error> {
123 if c < ASCII_0 || ASCII_9 < c {
124 return Err(Error::InvalidFormat("Not a number"));
125 }
126 Ok((c - ASCII_0) as u32)
127}
128
129#[cfg(test)]
130mod tests {
131 use std::io::ErrorKind;
132 use std::io::Read;
133 use super::*;
134
135 #[test]
136 fn discard_exact() {
137 let mut buf = b"abc".as_ref();
138 buf.discard_exact(1).unwrap();
139 assert_eq!(buf, b"bc");
140 buf.discard_exact(2).unwrap();
141 assert_eq!(buf, b"");
142 buf.discard_exact(1).unwrap_err();
143 }
144
145 #[test]
146 fn read8_len() {
147 let data = [];
148 assert_err_kind!(read8(&mut &data[..]), ErrorKind::UnexpectedEof);
149 let data = [0x01];
150 assert_ok!(read8(&mut &data[..]), 0x01);
151 let data = [0x01, 0x02];
152 let mut reader = &data[..];
153 let mut buf = Vec::new();
154 assert_ok!(read8(&mut reader), 0x01);
155 assert_ok!(reader.read_to_end(&mut buf), 1);
156 assert_eq!(buf, [0x02]);
157 }
158
159 #[test]
160 fn read16_len() {
161 let data = [];
162 assert_err_kind!(read16(&mut &data[..]), ErrorKind::UnexpectedEof);
163 let data = [0x01];
164 assert_err_kind!(read16(&mut &data[..]), ErrorKind::UnexpectedEof);
165 let data = [0x01, 0x02];
166 assert_ok!(read16(&mut &data[..]), 0x0102);
167 let data = [0x01, 0x02, 0x03];
168 let mut reader = &data[..];
169 let mut buf = Vec::new();
170 assert_ok!(read16(&mut reader), 0x0102);
171 assert_ok!(reader.read_to_end(&mut buf), 1);
172 assert_eq!(buf, [0x03]);
173 }
174
175 #[test]
176 fn atou16_misc() {
177 assert_ok!(atou16(b"0"), 0);
178 assert_ok!(atou16(b"0010"), 10);
179 assert_ok!(atou16(b"9999"), 9999);
180 assert_err_pat!(atou16(b""), Error::InvalidFormat(_));
181 assert_err_pat!(atou16(b"/"), Error::InvalidFormat(_));
182 assert_err_pat!(atou16(b":"), Error::InvalidFormat(_));
183 assert_err_pat!(atou16(b"-1"), Error::InvalidFormat(_));
184 }
185}