// // Copyright (c) 2016 KAMADA Ken'ichi. // All rights reserved. // // Redistribution and use in source and binary forms, with or without // modification, are permitted provided that the following conditions // are met: // 1. Redistributions of source code must retain the above copyright // notice, this list of conditions and the following disclaimer. // 2. Redistributions in binary form must reproduce the above copyright // notice, this list of conditions and the following disclaimer in the // documentation and/or other materials provided with the distribution. // // THIS SOFTWARE IS PROVIDED BY THE AUTHOR AND CONTRIBUTORS ``AS IS'' AND // ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE // IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE // ARE DISCLAIMED. IN NO EVENT SHALL THE AUTHOR OR CONTRIBUTORS BE LIABLE // FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL // DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS // OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) // HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT // LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY // OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF // SUCH DAMAGE. // use std::io; use std::io::Read as _; use crate::error::Error; const ASCII_0: u8 = 0x30; const ASCII_9: u8 = 0x39; pub fn read8(reader: &mut R) -> Result where R: io::Read { let mut buf = [0u8; 1]; reader.read_exact(&mut buf).and(Ok(buf[0])) } pub fn read16(reader: &mut R) -> Result where R: io::Read { let mut buf = [0u8; 2]; reader.read_exact(&mut buf)?; Ok(u16::from_be_bytes(buf)) } pub fn read64(reader: &mut R) -> Result where R: io::Read { let mut buf = [0u8; 8]; reader.read_exact(&mut buf)?; Ok(u64::from_be_bytes(buf)) } pub trait BufReadExt { fn discard_exact(&mut self, len: usize) -> io::Result<()>; fn is_eof(&mut self) -> io::Result; } impl BufReadExt for T where T: io::BufRead { fn discard_exact(&mut self, mut len: usize) -> io::Result<()> { while len > 0 { let consume_len = match self.fill_buf() { Ok(buf) if buf.is_empty() => return Err(io::Error::new( io::ErrorKind::UnexpectedEof, "unexpected EOF")), Ok(buf) => buf.len().min(len), Err(e) if e.kind() == io::ErrorKind::Interrupted => continue, Err(e) => return Err(e), }; self.consume(consume_len); len -= consume_len; } Ok(()) } fn is_eof(&mut self) -> io::Result { loop { match self.fill_buf() { Ok(buf) => return Ok(buf.is_empty()), Err(e) if e.kind() == io::ErrorKind::Interrupted => continue, Err(e) => return Err(e), } } } } pub trait ReadExt { fn read_exact_len(&mut self, buf: &mut Vec, len: usize) -> io::Result<()>; } impl ReadExt for T where T: io::Read { fn read_exact_len(&mut self, buf: &mut Vec, len: usize) -> io::Result<()> { // Using `vec![0; len]` and `read_exact` is more efficient but // less robust against broken files; a small file can easily // trigger OOM by a huge length value without actual data. // When the fallible allocation feature is stabilized, // we could revisit this. if self.take(len as u64).read_to_end(buf)? != len { return Err(io::Error::new( io::ErrorKind::UnexpectedEof, "unexpected EOF")); } Ok(()) } } // This function must not be called with more than 4 bytes. pub fn atou16(bytes: &[u8]) -> Result { debug_assert!(bytes.len() <= 4, "atou16 accepts up to 4 bytes"); if bytes.len() == 0 { return Err(Error::InvalidFormat("Not a number")); } let mut n = 0; for &c in bytes { if c < ASCII_0 || ASCII_9 < c { return Err(Error::InvalidFormat("Not a number")); } n = n * 10 + (c - ASCII_0) as u16; } Ok(n) } pub fn ctou32(c: u8) -> Result { if c < ASCII_0 || ASCII_9 < c { return Err(Error::InvalidFormat("Not a number")); } Ok((c - ASCII_0) as u32) } #[cfg(test)] mod tests { use std::io::ErrorKind; use std::io::Read; use super::*; #[test] fn discard_exact() { let mut buf = b"abc".as_ref(); buf.discard_exact(1).unwrap(); assert_eq!(buf, b"bc"); buf.discard_exact(2).unwrap(); assert_eq!(buf, b""); buf.discard_exact(1).unwrap_err(); } #[test] fn read8_len() { let data = []; assert_err_kind!(read8(&mut &data[..]), ErrorKind::UnexpectedEof); let data = [0x01]; assert_ok!(read8(&mut &data[..]), 0x01); let data = [0x01, 0x02]; let mut reader = &data[..]; let mut buf = Vec::new(); assert_ok!(read8(&mut reader), 0x01); assert_ok!(reader.read_to_end(&mut buf), 1); assert_eq!(buf, [0x02]); } #[test] fn read16_len() { let data = []; assert_err_kind!(read16(&mut &data[..]), ErrorKind::UnexpectedEof); let data = [0x01]; assert_err_kind!(read16(&mut &data[..]), ErrorKind::UnexpectedEof); let data = [0x01, 0x02]; assert_ok!(read16(&mut &data[..]), 0x0102); let data = [0x01, 0x02, 0x03]; let mut reader = &data[..]; let mut buf = Vec::new(); assert_ok!(read16(&mut reader), 0x0102); assert_ok!(reader.read_to_end(&mut buf), 1); assert_eq!(buf, [0x03]); } #[test] fn atou16_misc() { assert_ok!(atou16(b"0"), 0); assert_ok!(atou16(b"0010"), 10); assert_ok!(atou16(b"9999"), 9999); assert_err_pat!(atou16(b""), Error::InvalidFormat(_)); assert_err_pat!(atou16(b"/"), Error::InvalidFormat(_)); assert_err_pat!(atou16(b":"), Error::InvalidFormat(_)); assert_err_pat!(atou16(b"-1"), Error::InvalidFormat(_)); } }