path: root/src/decoder.rs

use std::cmp;
use std::io::{self, Read};

use encoding_rs::{Decoder, Encoding, UTF_8};

/// A BOM is at least 2 bytes and at most 3 bytes.
///
/// If fewer than 2 bytes are available to be read at the beginning of a
/// reader, then a BOM is `None`.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
struct Bom {
    bytes: [u8; 3],
    len: usize,
}

impl Bom {
    fn as_slice(&self) -> &[u8] {
        &self.bytes[0..self.len]
    }

    fn decoder(&self) -> Option<Decoder> {
        let bom = self.as_slice();
        if bom.len() < 3 {
            return None;
        }
        if let Some((enc, _)) = Encoding::for_bom(bom) {
            if enc != UTF_8 {
                return Some(enc.new_decoder_with_bom_removal());
            }
        }
        None
    }
}

/// `BomPeeker` wraps `R` and satisfies the `io::Read` interface while also
/// providing a peek at the BOM if one exists. Peeking at the BOM does not
/// advance the reader.
struct BomPeeker<R> {
    rdr: R,
    bom: Option<Bom>,
    nread: usize,
}

impl<R: io::Read> BomPeeker<R> {
    /// Create a new BomPeeker.
    ///
    /// The first three bytes can be read using the `peek_bom` method, but
    /// will not advance the reader.
    fn new(rdr: R) -> BomPeeker<R> {
        BomPeeker { rdr: rdr, bom: None, nread: 0 }
    }

    /// Peek at the first three bytes of the underlying reader.
    ///
    /// This does not advance the reader provided by `BomPeeker`.
    ///
    /// If the underlying reader does not have at least two bytes available,
    /// then `None` is returned.
    fn peek_bom(&mut self) -> io::Result<Bom> {
        if let Some(bom) = self.bom {
            return Ok(bom);
        }
        self.bom = Some(Bom { bytes: [0; 3], len: 0 });
        let mut buf = [0u8; 3];
        let bom_len = read_full(&mut self.rdr, &mut buf)?;
        self.bom = Some(Bom { bytes: buf, len: bom_len });
        Ok(self.bom.unwrap())
    }
}

impl<R: io::Read> io::Read for BomPeeker<R> {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        if self.nread < 3 {
            let bom = self.peek_bom()?;
            let bom = bom.as_slice();
            if self.nread < bom.len() {
                let rest = &bom[self.nread..];
                let len = cmp::min(buf.len(), rest.len());
                buf[..len].copy_from_slice(&rest[..len]);
                self.nread += len;
                return Ok(len);
            }
        }
        let nread = self.rdr.read(buf)?;
        self.nread += nread;
        Ok(nread)
    }
}

/// Like `io::Read::read_exact`, except it never returns `UnexpectedEof` and
/// instead returns the number of bytes read if EOF is seen before filling
/// `buf`.
fn read_full<R: io::Read>(
    mut rdr: R,
    mut buf: &mut [u8],
) -> io::Result<usize> {
    let mut nread = 0;
    while !buf.is_empty() {
        match rdr.read(buf) {
            Ok(0) => break,
            Ok(n) => {
                nread += n;
                let tmp = buf;
                buf = &mut tmp[n..];
            }
            Err(ref e) if e.kind() == io::ErrorKind::Interrupted => {}
            Err(e) => return Err(e),
        }
    }
    Ok(nread)
}

/// A reader that transcodes to UTF-8. The source encoding is determined by
/// inspecting the BOM from the stream read from `R`, if one exists. If a
/// UTF-16 BOM exists, then the source stream is transcoded to UTF-8 with
/// invalid UTF-16 sequences translated to the Unicode replacement character.
/// In all other cases, the underlying reader is passed through unchanged.
///
/// `R` is the type of the underlying reader and `B` is the type of an internal
/// buffer used to store the results of transcoding.
///
/// Note that not all methods on `io::Read` work with this implementation.
/// For example, the `bytes` adapter method attempts to read a single byte at
/// a time, but this implementation requires a buffer of size at least `4`. If
/// a buffer of size less than 4 is given, then an error is returned.
pub struct DecodeReader<R, B> {
    /// The underlying reader, wrapped in a peeker for reading a BOM if one
    /// exists.
    rdr: BomPeeker<R>,
    /// The internal buffer to store transcoded bytes before they are read by
    /// callers.
    buf: B,
    /// The current position in `buf`. Subsequent reads start here.
    pos: usize,
    /// The number of transcoded bytes in `buf`. Subsequent reads end here.
    buflen: usize,
    /// Whether this is the first read or not (in which we inspect the BOM).
    first: bool,
    /// Whether a "last" read has occurred. After this point, EOF will always
    /// be returned.
    last: bool,
    /// The underlying text decoder derived from the BOM, if one exists.
    decoder: Option<Decoder>,
}

impl<R: io::Read, B: AsMut<[u8]>> DecodeReader<R, B> {
    /// Create a new transcoder that converts a source stream to valid UTF-8.
    ///
    /// If an encoding is specified, then it is used to transcode `rdr` to
    /// UTF-8. Otherwise, if no encoding is specified, and if a UTF-16 BOM is
    /// found, then the corresponding UTF-16 encoding is used to transcode
    /// `rdr` to UTF-8. In all other cases, `rdr` is assumed to be at least
    /// ASCII-compatible and passed through untouched.
    ///
    /// Errors in the encoding of `rdr` are handled with the Unicode
    /// replacement character. If no encoding of `rdr` is specified, then
    /// errors are not handled.
    pub fn new(
        rdr: R,
        buf: B,
        enc: Option<&'static Encoding>,
    ) -> DecodeReader<R, B> {
        DecodeReader {
            rdr: BomPeeker::new(rdr),
            buf: buf,
            buflen: 0,
            pos: 0,
            first: enc.is_none(),
            last: false,
            decoder: enc.map(|enc| enc.new_decoder_with_bom_removal()),
        }
    }

    /// Fill the internal buffer from the underlying reader.
    ///
    /// If there are unread bytes in the internal buffer, then we move them
    /// to the beginning of the internal buffer and fill the remainder.
    ///
    /// If the internal buffer is too small to read additional bytes, then an
    /// error is returned.
    #[inline(always)] // massive perf benefit (???)
    fn fill(&mut self) -> io::Result<()> {
        if self.pos < self.buflen {
            if self.buflen >= self.buf.as_mut().len() {
                return Err(io::Error::new(
                    io::ErrorKind::Other,
                    "DecodeReader: internal buffer exhausted"));
            }
            let newlen = self.buflen - self.pos;
            let mut tmp = Vec::with_capacity(newlen);
            tmp.extend_from_slice(&self.buf.as_mut()[self.pos..self.buflen]);
            self.buf.as_mut()[..newlen].copy_from_slice(&tmp);
            self.buflen = newlen;
        } else {
            self.buflen = 0;
        }
        self.pos = 0;
        self.buflen +=
            self.rdr.read(&mut self.buf.as_mut()[self.buflen..])?;
        Ok(())
    }

    /// Transcode the inner stream to UTF-8 in `buf`. This assumes that there
    /// is a decoder capable of transcoding the inner stream to UTF-8. This
    /// returns the number of bytes written to `buf`.
    ///
    /// When this function returns, exactly one of the following things will
    /// be true:
    ///
    /// 1. A non-zero number of bytes were written to `buf`.
    /// 2. The underlying reader reached EOF.
    /// 3. An error is returned: the internal buffer ran out of room.
    /// 4. An I/O error occurred.
    ///
    /// Note that `buf` must have at least 4 bytes of space.
    fn transcode(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        assert!(buf.len() >= 4);
        if self.last {
            return Ok(0);
        }
        if self.pos >= self.buflen {
            self.fill()?;
        }
        let mut nwrite = 0;
        loop {
            let (_, nin, nout, _) =
                self.decoder.as_mut().unwrap().decode_to_utf8(
                    &self.buf.as_mut()[self.pos..self.buflen], buf, false);
            self.pos += nin;
            nwri