path: root/openpgp/src/parse/partial_body.rs

use std;
use std::cmp;
use std::io;
use std::io::{Error,ErrorKind};

use buffered_reader::{buffered_reader_generic_read_impl, BufferedReader};
use super::BodyLength;
use parse::body_length_new_format;

/// A `BufferedReader` that transparently handles OpenPGP's chunking
/// scheme.  This implicitly implements a limitor.
pub struct BufferedReaderPartialBodyFilter<T: BufferedReader<C>, C> {
    // The underlying reader.
    reader: T,

    // The amount of unread data in the current partial body chunk.
    // That is, if `buffer` contains 10 bytes and
    // `partial_body_length` is 20, then there are 30 bytes of
    // unprocessed (unconsumed) data in the current chunk.
    partial_body_length: u32,
    // Whether this is the last partial body chuck.
    last: bool,

    // Sometimes we have to double buffer.  This happens if the caller
    // requests X bytes and that chunk straddles a partial body length
    // boundary.
    buffer: Option<Box<[u8]>>,
    // The position within the buffer.
    cursor: usize,

    // The user-defined cookie.
    cookie: C,
}

impl<T: BufferedReader<C>, C> std::fmt::Debug
        for BufferedReaderPartialBodyFilter<T, C> {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        f.debug_struct("BufferedReaderPartialBodyFilter")
            .field("reader", &self.reader)
            .field("partial_body_length", &self.partial_body_length)
            .field("last", &self.last)
            .field("buffer (bytes left)",
                   &if let Some(ref buffer) = self.buffer {
                       Some(buffer.len())
                   } else {
                       None
                   })
            .finish()
    }
}

impl<T: BufferedReader<C>, C> BufferedReaderPartialBodyFilter<T, C> {
    /// Create a new BufferedReaderPartialBodyFilter object.
    /// `partial_body_length` is the amount of data in the initial
    /// partial body chunk.
    pub fn with_cookie(reader: T, partial_body_length: u32, cookie: C) -> Self {
        BufferedReaderPartialBodyFilter {
            reader: reader,
            partial_body_length: partial_body_length,
            last: false,
            buffer: None,
            cursor: 0,
            cookie: cookie,
        }
    }

    // Make sure that the local buffer contains `amount` bytes.
    fn do_fill_buffer (&mut self, amount: usize) -> Result<(), std::io::Error> {
        // eprintln!("BufferedReaderPartialBodyFilter::do_fill_buffer(\
        //            amount: {}) (partial body length: {}, last: {})",
        //           amount, self.partial_body_length, self.last);

        // We want to avoid double buffering as much as possible.
        // Thus, we only buffer as much as needed.
        let mut buffer = vec![0; amount];
        let mut amount_buffered = 0;

        if let Some(ref old_buffer) = self.buffer {
            // The amount of data that is left in the old buffer.
            let amount_left = old_buffer.len() - self.cursor;

            // This function should only be called if we actually need
            // to read something.
            assert!(amount > amount_left);

            amount_buffered = amount_left;

            // Copy the data that is still in buffer.
            buffer[..amount_buffered]
                .copy_from_slice(&old_buffer[self.cursor..]);

        }

        let mut err = None;

        loop {
            let to_read = cmp::min(
                // Data in current chunk.
                self.partial_body_length as usize,
                // Space left in the buffer.
                buffer.len() - amount_buffered);
            //println!("Trying to buffer {} bytes (partial body length: {}; space: {})",
            //         to_read, self.partial_body_length,
            //         buffer.len() - amount_buffered);
            if to_read > 0 {
                let result = self.reader.read(
                    &mut buffer[amount_buffered..amount_buffered + to_read]);
                match result {
                    Ok(did_read) => {
                        //println!("Buffered {} bytes", did_read);
                        amount_buffered += did_read;
                        self.partial_body_length -= did_read as u32;

                        if did_read < to_read {
                            // Short read => EOF.  We're done.
                            // (Although the underlying message is
                            // probably corrupt.)
                            break;
                        }
                    },
                    Err(e) => {
                        //println!("Err reading: {:?}", e);
                        err = Some(e);
                        break;
                    },
                }
            }

            if amount_buffered == amount || self.last {
                // We're read enough or we've read everything.
                break;
            }

            // Read the next partial body length header.
            assert_eq!(self.partial_body_length, 0);

            match body_length_new_format(&mut self.reader) {
                Ok(BodyLength::Full(len)) => {
                    //println!("Last chunk: {} bytes", len);
                    self.last = true;
                    self.partial_body_length = len;
                },
                Ok(BodyLength::Partial(len)) => {
                    //println!("Next chunk: {} bytes", len);
                    self.partial_body_length = len;
                },
                Ok(BodyLength::Indeterminate) => {
                    // A new format packet can't return Indeterminate.
                    unreachable!();
                },
                Err(e) => {
                    //println!("Err reading next chuck: {:?}", e);
                    err = Some(e);
                    break;
                }
            }
        }

        buffer.truncate(amount_buffered);
        buffer.shrink_to_fit();

        // We're done.
        self.buffer = Some(buffer.into_boxed_slice());
        self.cursor = 0;

        if let Some(err) = err {
            return Err(err)
        } else {
            return Ok(());
        }
    }

    fn data_helper(&mut self, amount: usize, hard: bool, and_consume: bool)
                   -> Result<&[u8], std::io::Error> {
        let mut need_fill = false;

        //println!("BufferedReaderPartialBodyFilter::data_helper({})", amount);

        if let Some(ref buffer) = self.buffer {
            // We have some data buffered locally.

            //println!("  Reading from buffer");

            let amount_buffered = buffer.len() - self.cursor;
            if amount > amount_buffered {
                // The requested amount exceeds what is in the buffer.
                // Read more.

                // We can't call self.do_fill_buffer here, because self
                // is borrowed.  Set a flag and do it after the borrow
                // ends.
                need_fill = true;
            }
        } else {
            // We don't have any data buffered.

            assert_eq!(self.cursor, 0);

            if amount <= self.partial_body_length as usize
                || /* Short read.  */ self.last {
                // The amount of data that the caller requested does
                // not exceed the amount of data in the current chunk.
                // As such, there is no need to double buffer.

                //println!("  Reading from inner reader");

                let result = if hard && and_consume {
                    self.