<commit_before>use std::io;
use std::sync::{Arc, RwLock};
use std::time::Instant;

use console::Term;

/// Target for draw operations
///
/// This tells a progress bar or a multi progress object where to paint to.
/// The draw target is a stateful wrapper over a drawing destination and
/// internally optimizes how often the state is painted to the output
/// device.
#[derive(Debug)]
pub struct ProgressDrawTarget {
    kind: ProgressDrawTargetKind,
}

impl ProgressDrawTarget {
    /// Draw to a buffered stdout terminal at a max of 15 times a second.
    ///
    /// For more information see `ProgressDrawTarget::to_term`.
    pub fn stdout() -> ProgressDrawTarget {
        ProgressDrawTarget::term(Term::buffered_stdout(), 15)
    }

    /// Draw to a buffered stderr terminal at a max of 15 times a second.
    ///
    /// This is the default draw target for progress bars.  For more
    /// information see `ProgressDrawTarget::to_term`.
    pub fn stderr() -> ProgressDrawTarget {
        ProgressDrawTarget::term(Term::buffered_stderr(), 15)
    }

    /// Draw to a buffered stdout terminal at a max of `refresh_rate` times a second.
    ///
    /// For more information see `ProgressDrawTarget::to_term`.
    pub fn stdout_with_hz(refresh_rate: u64) -> ProgressDrawTarget {
        ProgressDrawTarget::term(Term::buffered_stdout(), refresh_rate)
    }

    /// Draw to a buffered stderr terminal at a max of `refresh_rate` times a second.
    ///
    /// For more information see `ProgressDrawTarget::to_term`.
    pub fn stderr_with_hz(refresh_rate: u64) -> ProgressDrawTarget {
        ProgressDrawTarget::term(Term::buffered_stderr(), refresh_rate)
    }

    /// Draw to a buffered stdout terminal without max framerate.
    ///
    /// This is useful when data is known to come in very slowly and
    /// not rendering some updates would be a problem (for instance
    /// when messages are used extensively).
    ///
    /// For more information see `ProgressDrawTarget::to_term`.
    pub fn stdout_nohz() -> ProgressDrawTarget {
        ProgressDrawTarget::term(Term::buffered_stdout(), None)
    }

    /// Draw to a buffered stderr terminal without max framerate.
    ///
    /// This is useful when data is known to come in very slowly and
    /// not rendering some updates would be a problem (for instance
    /// when messages are used extensively).
    ///
    /// For more information see `ProgressDrawTarget::to_term`.
    pub fn stderr_nohz() -> ProgressDrawTarget {
        ProgressDrawTarget::term(Term::buffered_stderr(), None)
    }

    pub(crate) fn new_remote(state: Arc<RwLock<MultiProgressState>>, idx: usize) -> Self {
        Self {
            kind: ProgressDrawTargetKind::Remote { state, idx },
        }
    }

    /// Draw to a terminal, optionally with a specific refresh rate.
    ///
    /// Progress bars are by default drawn to terminals however if the
    /// terminal is not user attended the entire progress bar will be
    /// hidden.  This is done so that piping to a file will not produce
    /// useless escape codes in that file.
    ///
    /// Will panic if refresh_rate is `Some(0)`. To disable rate limiting use `None` instead.
    #[allow(clippy::wrong_self_convention)]
    #[deprecated(since = "0.16.0", note = "Use `ProgressDrawTarget::term` instead")]
    pub fn to_term(term: Term, refresh_rate: impl Into<Option<u64>>) -> ProgressDrawTarget {
        ProgressDrawTarget::term(term, refresh_rate)
    }

    /// Draw to a terminal, optionally with a specific refresh rate.
    ///
    /// Progress bars are by default drawn to terminals however if the
    /// terminal is not user attended the entire progress bar will be
    /// hidden.  This is done so that piping to a file will not produce
    /// useless escape codes in that file.
    ///
    /// Will panic if refresh_rate is `Some(0)`. To disable rate limiting use `None` instead.
    pub fn term(term: Term, refresh_rate: impl Into<Option<u64>>) -> ProgressDrawTarget {
        ProgressDrawTarget {
            kind: ProgressDrawTargetKind::Term {
                term,
                last_line_count: 0,
                leaky_bucket: refresh_rate.into().map(|rate| LeakyBucket {
                    bucket: MAX_GROUP_SIZE,
                    leak_rate: rate as f64,
                    last_update: Instant::now(),
                }),
            },
        }
    }

    /// A hidden draw target.
    ///
    /// This forces a progress bar to be not rendered at all.
    pub fn hidden() -> ProgressDrawTarget {
        ProgressDrawTarget {
            kind: ProgressDrawTargetKind::Hidden,
        }
    }

    /// Returns true if the draw target is hidden.
    ///
    /// This is internally used in progress bars to figure out if overhead
    /// from drawing can be prevented.
    pub fn is_hidden(&self) -> bool {
        match self.kind {
            ProgressDrawTargetKind::Hidden => true,
            ProgressDrawTargetKind::Term { ref term, .. } => !term.is_term(),
            _ => false,
        }
    }

    /// Returns the current width of the draw target.
    pub(crate) fn width(&self) -> usize {
        match self.kind {
            ProgressDrawTargetKind::Term { ref term, .. } => term.size().1 as usize,
            ProgressDrawTargetKind::Remote { ref state, .. } => state.read().unwrap().width(),
            ProgressDrawTargetKind::Hidden => 0,
        }
    }

    /// Apply the given draw state (draws it).
    pub(crate) fn apply_draw_state(&mut self, draw_state: ProgressDrawState) -> io::Result<()> {
        let (term, last_line_count) = match self.kind {
            ProgressDrawTargetKind::Term {
                ref term,
                ref mut last_line_count,
                leaky_bucket: None,
            } => (term, last_line_count),
            ProgressDrawTargetKind::Term {
                ref term,
                ref mut last_line_count,
                leaky_bucket: Some(ref mut leaky_bucket),
            } => {
                if draw_state.force_draw || leaky_bucket.try_add_work() {
                    (term, last_line_count)
                } else {
                    // rate limited
                    return Ok(());
                }
            }
            ProgressDrawTargetKind::Remote { idx, ref state, .. } => {
                return state
                    .write()
                    .unwrap()
                    .draw(idx, draw_state)
                    .map_err(|e| io::Error::new(io::ErrorKind::Other, e));
            }
            // Hidden, finished, or no need to refresh yet
            _ => return Ok(()),
        };

        if !draw_state.lines.is_empty() && draw_state.move_cursor {
            term.move_cursor_up(*last_line_count)?;
        } else {
            term.clear_last_lines(*last_line_count)?;
        }

        let shift = match draw_state.alignment {
            MultiProgressAlignment::Bottom if draw_state.lines.len() < *last_line_count => {
                let shift = *last_line_count - draw_state.lines.len();
                for _ in 0..shift {
                    term.write_line("")?;
                }
                shift
            }
            _ => 0,
        };
        draw_state.draw_to_term(term)?;
        term.flush()?;
        *last_line_count = draw_state.lines.len() - draw_state.orphan_lines + shift;
        Ok(())
    }

    /// Properly disconnects from the draw target
    pub(crate) fn disconnect(&self) {
        match self.kind {
            ProgressDrawTargetKind::Term { .. } => {}
            ProgressDrawTargetKind::Remote { idx, ref state, .. } => {
                state
                    .write()
                    .unwrap()
                    .draw(
                        idx,
                        ProgressDrawState {
                            lines: vec![],
                            orphan_lines: 0,
                            force_draw: true,
                            move_cursor: false,
                            alignment: Default::default(),
                        },
                    )
                    .ok();
            }
            ProgressDrawTargetKind::Hidden => {}
        };
    }

    pub(crate) fn remote(&self) -> Option<(&Arc<RwLock<MultiProgressState>>, usize)> {
        match &self.kind {
            ProgressDrawTargetKind::Remote { state, idx } => Some((state, *idx)),
            _ => None,
        }
    }
}

#[derive(Debug)]
enum ProgressDrawTargetKind {
    Term {
        term: Term,
        last_line_count: usize,
        leaky_bucket: Option<LeakyBucket>,
    },
    Remote {
        state: Arc<RwLock<MultiProgressState>>,
        idx: usize,
    },
    Hidden,
}

#[derive(Debug)]
pub(crate) struct MultiProgressState {
    /// The collection of states corresponding to progress bars
    /// the state is None for bars that have not yet been drawn or have been removed
    pub(crate) draw_states: Vec<Option<ProgressDrawState>>,
    /// Set of removed bars, should have corresponding `None` elements in the `draw_states` vector
    pub(crate) free_set: Vec<usize>,
    /// Indices to the `draw_states` to maintain correct visual order
    pub(crate) ordering: Vec<usize>,
    /// Target for draw operation for MultiProgress
    pub(crate) draw_target: ProgressDrawTarget,
    /// Whether or not to just move cursor instead of clearing lines
    pub(crate) move_cursor: bool,
    /// Controls how the multi progress is aligned if some of its progress bars get removed, default is `Top`
    pub(crate) alignment: MultiProgressAlignment,
}

impl MultiProgressState {
    fn width(&self) -> usize {
        self.draw_target.width()
    }

    pub(crate) fn draw(&mut self, idx: usize, draw_state: ProgressDrawState) -> io::Result<()> {
        let force_draw = draw_state.force_draw;
        let mut orphan_lines = vec![];

        // Split orphan lines out of the draw state, if any
        let lines = if draw_state.orphan_lines > 0 {
            let split = draw_state.lines.split_at(draw_state.orphan_lines);
            orphan_lines.extend_from_slice(split.0);
            split.1.to_vec()
        } else {
            draw_state.lines
        };

        let draw_state = ProgressDrawState {
            lines,
            orphan_lines: 0,
            ..draw_state
        };

        self.draw_states[idx] = Some(draw_state);

        // the rest from here is only drawing, we can skip it.
        if self.draw_target.is_hidden() {
            return Ok(());
        }

        let mut lines = vec![];

        // Make orphaned lines appear at the top, so they can be properly
        // forgotten.
        let orphan_lines_count = orphan_lines.len();
        lines.append(&mut orphan_lines);

        for index in self.ordering.iter() {
            let draw_state = &self.draw_states[*index];
            if let Some(ref draw_state) = draw_state {
                lines.extend_from_slice(&draw_state.lines[..]);
            }
        }

        self.draw_target.apply_draw_state(ProgressDrawState {
            lines,
            orphan_lines: orphan_lines_count,
            force_draw: force_draw || orphan_lines_count > 0,
            move_cursor: self.move_cursor,
            alignment: self.alignment,
        })
    }

    pub(crate) fn len(&self) -> usize {
        self.draw_states.len() - self.free_set.len()
    }

    pub(crate) fn remove_idx(&mut self, idx: usize) {
        if self.free_set.contains(&idx) {
            return;
        }

        self.draw_states[idx].take();
        self.free_set.push(idx);
        self.ordering.retain(|&x| x != idx);

        assert!(
            self.len() == self.ordering.len(),
            "Draw state is inconsistent"
        );
    }
}

#[derive(Debug)]
pub(crate) struct LeakyBucket {
    leak_rate: f64,
    last_update: Instant,
    bucket: f64,
}

/// Rate limit but allow occasional bursts above desired rate
impl LeakyBucket {
    /// try to add some work to the bucket
    /// return false if the bucket is already full and the work should be skipped
    fn try_add_work(&mut self) -> bool {
        self.leak();
        if self.bucket < MAX_GROUP_SIZE {
            self.bucket += 1.0;
            true
        } else {
            false
        }
    }

    fn leak(&mut self) {
        let ticks = self.last_update.elapsed().as_secs_f64() * self.leak_rate;
        self.bucket -= ticks;
        if self.bucket < 0.0 {
            self.bucket = 0.0;
        }
        self.last_update = Instant::now();
    }
}

const MAX_GROUP_SIZE: f64 = 32.0;

/// The drawn state of an element.
#[derive(Clone, Debug)]
pub(crate) struct ProgressDrawState {
    /// The lines to print (can contain ANSI codes)
    pub lines: Vec<String>,
    /// The number of lines that shouldn't be reaped by the next tick.
    pub orphan_lines: usize,
    /// True if drawing should be forced.
    pub force_draw: bool,
    /// True if we should move the cursor up when possible instead of clearing lines.
    pub move_cursor: bool,
    /// Controls how the multi progress is aligned if some of its progress bars get removed, default is `Top`
    pub alignment: MultiProgressAlignment,
}

impl ProgressDrawState {
    pub(crate) fn new(lines: Vec<String>, force_draw: bool) -> Self {
        Self {
            lines,
            orphan_lines: 0,
            force_draw,
            move_cursor: false,
            alignment: Default::default(),
        }
    }

    pub fn draw_to_term(&self, term: &Term) -> io::Result<()> {
        for line in &self.lines {
            term.write_line(line)?;
        }
        Ok(())
    }
}

/// Vertical alignment of a multi progress.
///
/// The alignment controls how the multi progress is aligned if some of its progress bars get removed.
/// E.g. `Top` alignment (default), when _progress bar 2_ is removed:
/// ```ignore
/// [0/100] progress bar 1        [0/100] progress bar 1
/// [0/100] progress bar 2   =>   [0/100] progress bar 3
/// [0/100] progress bar 3
/// ```
///
/// `Bottom` alignment
/// ```ignore
/// [0/100] progress bar 1
/// [0/100] progress bar 2   =>   [0/100] progress bar 1
/// [0/100] progress bar 3        [0/100] progress bar 3
/// ```
#[derive(Debug, Copy, Clone)]
pub enum MultiProgressAlignment {
    Top,
    Bottom,
}

impl Default for MultiProgressAlignment {
    fn default() -> Self {
        Self::Top
    }
}
<commit_msg>Don't add an extra line after the progress bar<commit_after>use std::io;
use std::sync::{Arc, RwLock};
use std::time::Instant;

use console::Term;

/// Target for draw operations
///
/// This tells a progress bar or a multi progress object where to paint to.
/// The draw target is a stateful wrapper over a drawing destination and
/// internally optimizes how often the state is painted to the output
/// device.
#[derive(Debug)]
pub struct ProgressDrawTarget {
    kind: ProgressDrawTargetKind,
}

impl ProgressDrawTarget {
    /// Draw to a buffered stdout terminal at a max of 15 times a second.
    ///
    /// For more information see `ProgressDrawTarget::to_term`.
    pub fn stdout() -> ProgressDrawTarget {
        ProgressDrawTarget::term(Term::buffered_stdout(), 15)
    }

    /// Draw to a buffered stderr terminal at a max of 15 times a second.
    ///
    /// This is the default draw target for progress bars.  For more
    /// information see `ProgressDrawTarget::to_term`.
    pub fn stderr() -> ProgressDrawTarget {
        ProgressDrawTarget::term(Term::buffered_stderr(), 15)
    }

    /// Draw to a buffered stdout terminal at a max of `refresh_rate` times a second.
    ///
    /// For more information see `ProgressDrawTarget::to_term`.
    pub fn stdout_with_hz(refresh_rate: u64) -> ProgressDrawTarget {
        ProgressDrawTarget::term(Term::buffered_stdout(), refresh_rate)
    }

    /// Draw to a buffered stderr terminal at a max of `refresh_rate` times a second.
    ///
    /// For more information see `ProgressDrawTarget::to_term`.
    pub fn stderr_with_hz(refresh_rate: u64) -> ProgressDrawTarget {
        ProgressDrawTarget::term(Term::buffered_stderr(), refresh_rate)
    }

    /// Draw to a buffered stdout terminal without max framerate.
    ///
    /// This is useful when data is known to come in very slowly and
    /// not rendering some updates would be a problem (for instance
    /// when messages are used extensively).
    ///
    /// For more information see `ProgressDrawTarget::to_term`.
    pub fn stdout_nohz() -> ProgressDrawTarget {
        ProgressDrawTarget::term(Term::buffered_stdout(), None)
    }

    /// Draw to a buffered stderr terminal without max framerate.
    ///
    /// This is useful when data is known to come in very slowly and
    /// not rendering some updates would be a problem (for instance
    /// when messages are used extensively).
    ///
    /// For more information see `ProgressDrawTarget::to_term`.
    pub fn stderr_nohz() -> ProgressDrawTarget {
        ProgressDrawTarget::term(Term::buffered_stderr(), None)
    }

    pub(crate) fn new_remote(state: Arc<RwLock<MultiProgressState>>, idx: usize) -> Self {
        Self {
            kind: ProgressDrawTargetKind::Remote { state, idx },
        }
    }

    /// Draw to a terminal, optionally with a specific refresh rate.
    ///
    /// Progress bars are by default drawn to terminals however if the
    /// terminal is not user attended the entire progress bar will be
    /// hidden.  This is done so that piping to a file will not produce
    /// useless escape codes in that file.
    ///
    /// Will panic if refresh_rate is `Some(0)`. To disable rate limiting use `None` instead.
    #[allow(clippy::wrong_self_convention)]
    #[deprecated(since = "0.16.0", note = "Use `ProgressDrawTarget::term` instead")]
    pub fn to_term(term: Term, refresh_rate: impl Into<Option<u64>>) -> ProgressDrawTarget {
        ProgressDrawTarget::term(term, refresh_rate)
    }

    /// Draw to a terminal, optionally with a specific refresh rate.
    ///
    /// Progress bars are by default drawn to terminals however if the
    /// terminal is not user attended the entire progress bar will be
    /// hidden.  This is done so that piping to a file will not produce
    /// useless escape codes in that file.
    ///
    /// Will panic if refresh_rate is `Some(0)`. To disable rate limiting use `None` instead.
    pub fn term(term: Term, refresh_rate: impl Into<Option<u64>>) -> ProgressDrawTarget {
        ProgressDrawTarget {
            kind: ProgressDrawTargetKind::Term {
                term,
                last_line_count: 0,
                leaky_bucket: refresh_rate.into().map(|rate| LeakyBucket {
                    bucket: MAX_GROUP_SIZE,
                    leak_rate: rate as f64,
                    last_update: Instant::now(),
                }),
            },
        }
    }

    /// A hidden draw target.
    ///
    /// This forces a progress bar to be not rendered at all.
    pub fn hidden() -> ProgressDrawTarget {
        ProgressDrawTarget {
            kind: ProgressDrawTargetKind::Hidden,
        }
    }

    /// Returns true if the draw target is hidden.
    ///
    /// This is internally used in progress bars to figure out if overhead
    /// from drawing can be prevented.
    pub fn is_hidden(&self) -> bool {
        match self.kind {
            ProgressDrawTargetKind::Hidden => true,
            ProgressDrawTargetKind::Term { ref term, .. } => !term.is_term(),
            _ => false,
        }
    }

    /// Returns the current width of the draw target.
    pub(crate) fn width(&self) -> usize {
        match self.kind {
            ProgressDrawTargetKind::Term { ref term, .. } => term.size().1 as usize,
            ProgressDrawTargetKind::Remote { ref state, .. } => state.read().unwrap().width(),
            ProgressDrawTargetKind::Hidden => 0,
        }
    }

    /// Apply the given draw state (draws it).
    pub(crate) fn apply_draw_state(&mut self, draw_state: ProgressDrawState) -> io::Result<()> {
        let (term, last_line_count) = match self.kind {
            ProgressDrawTargetKind::Term {
                ref term,
                ref mut last_line_count,
                leaky_bucket: None,
            } => (term, last_line_count),
            ProgressDrawTargetKind::Term {
                ref term,
                ref mut last_line_count,
                leaky_bucket: Some(ref mut leaky_bucket),
            } => {
                if draw_state.force_draw || leaky_bucket.try_add_work() {
                    (term, last_line_count)
                } else {
                    // rate limited
                    return Ok(());
                }
            }
            ProgressDrawTargetKind::Remote { idx, ref state, .. } => {
                return state
                    .write()
                    .unwrap()
                    .draw(idx, draw_state)
                    .map_err(|e| io::Error::new(io::ErrorKind::Other, e));
            }
            // Hidden, finished, or no need to refresh yet
            _ => return Ok(()),
        };

        if !draw_state.lines.is_empty() && draw_state.move_cursor {
            term.move_cursor_up(*last_line_count)?;
        } else {
            clear_last_lines(term, *last_line_count)?;
        }

        let shift = match draw_state.alignment {
            MultiProgressAlignment::Bottom if draw_state.lines.len() < *last_line_count => {
                let shift = *last_line_count - draw_state.lines.len();
                for _ in 0..shift {
                    term.write_line("")?;
                }
                shift
            }
            _ => 0,
        };
        draw_state.draw_to_term(term)?;
        term.flush()?;
        *last_line_count = draw_state.lines.len() - draw_state.orphan_lines + shift;
        Ok(())
    }

    /// Properly disconnects from the draw target
    pub(crate) fn disconnect(&self) {
        match self.kind {
            ProgressDrawTargetKind::Term { .. } => {}
            ProgressDrawTargetKind::Remote { idx, ref state, .. } => {
                state
                    .write()
                    .unwrap()
                    .draw(
                        idx,
                        ProgressDrawState {
                            lines: vec![],
                            orphan_lines: 0,
                            force_draw: true,
                            move_cursor: false,
                            alignment: Default::default(),
                        },
                    )
                    .ok();
            }
            ProgressDrawTargetKind::Hidden => {}
        };
    }

    pub(crate) fn remote(&self) -> Option<(&Arc<RwLock<MultiProgressState>>, usize)> {
        match &self.kind {
            ProgressDrawTargetKind::Remote { state, idx } => Some((state, *idx)),
            _ => None,
        }
    }
}

#[derive(Debug)]
enum ProgressDrawTargetKind {
    Term {
        term: Term,
        last_line_count: usize,
        leaky_bucket: Option<LeakyBucket>,
    },
    Remote {
        state: Arc<RwLock<MultiProgressState>>,
        idx: usize,
    },
    Hidden,
}

#[derive(Debug)]
pub(crate) struct MultiProgressState {
    /// The collection of states corresponding to progress bars
    /// the state is None for bars that have not yet been drawn or have been removed
    pub(crate) draw_states: Vec<Option<ProgressDrawState>>,
    /// Set of removed bars, should have corresponding `None` elements in the `draw_states` vector
    pub(crate) free_set: Vec<usize>,
    /// Indices to the `draw_states` to maintain correct visual order
    pub(crate) ordering: Vec<usize>,
    /// Target for draw operation for MultiProgress
    pub(crate) draw_target: ProgressDrawTarget,
    /// Whether or not to just move cursor instead of clearing lines
    pub(crate) move_cursor: bool,
    /// Controls how the multi progress is aligned if some of its progress bars get removed, default is `Top`
    pub(crate) alignment: MultiProgressAlignment,
}

impl MultiProgressState {
    fn width(&self) -> usize {
        self.draw_target.width()
    }

    pub(crate) fn draw(&mut self, idx: usize, draw_state: ProgressDrawState) -> io::Result<()> {
        let force_draw = draw_state.force_draw;
        let mut orphan_lines = vec![];

        // Split orphan lines out of the draw state, if any
        let lines = if draw_state.orphan_lines > 0 {
            let split = draw_state.lines.split_at(draw_state.orphan_lines);
            orphan_lines.extend_from_slice(split.0);
            split.1.to_vec()
        } else {
            draw_state.lines
        };

        let draw_state = ProgressDrawState {
            lines,
            orphan_lines: 0,
            ..draw_state
        };

        self.draw_states[idx] = Some(draw_state);

        // the rest from here is only drawing, we can skip it.
        if self.draw_target.is_hidden() {
            return Ok(());
        }

        let mut lines = vec![];

        // Make orphaned lines appear at the top, so they can be properly
        // forgotten.
        let orphan_lines_count = orphan_lines.len();
        lines.append(&mut orphan_lines);

        for index in self.ordering.iter() {
            let draw_state = &self.draw_states[*index];
            if let Some(ref draw_state) = draw_state {
                lines.extend_from_slice(&draw_state.lines[..]);
            }
        }

        self.draw_target.apply_draw_state(ProgressDrawState {
            lines,
            orphan_lines: orphan_lines_count,
            force_draw: force_draw || orphan_lines_count > 0,
            move_cursor: self.move_cursor,
            alignment: self.alignment,
        })
    }

    pub(crate) fn len(&self) -> usize {
        self.draw_states.len() - self.free_set.len()
    }

    pub(crate) fn remove_idx(&mut self, idx: usize) {
        if self.free_set.contains(&idx) {
            return;
        }

        self.draw_states[idx].take();
        self.free_set.push(idx);
        self.ordering.retain(|&x| x != idx);

        assert!(
            self.len() == self.ordering.len(),
            "Draw state is inconsistent"
        );
    }
}

#[derive(Debug)]
pub(crate) struct LeakyBucket {
    leak_rate: f64,
    last_update: Instant,
    bucket: f64,
}

/// Rate limit but allow occasional bursts above desired rate
impl LeakyBucket {
    /// try to add some work to the bucket
    /// return false if the bucket is already full and the work should be skipped
    fn try_add_work(&mut self) -> bool {
        self.leak();
        if self.bucket < MAX_GROUP_SIZE {
            self.bucket += 1.0;
            true
        } else {
            false
        }
    }

    fn leak(&mut self) {
        let ticks = self.last_update.elapsed().as_secs_f64() * self.leak_rate;
        self.bucket -= ticks;
        if self.bucket < 0.0 {
            self.bucket = 0.0;
        }
        self.last_update = Instant::now();
    }
}

const MAX_GROUP_SIZE: f64 = 32.0;

/// The drawn state of an element.
#[derive(Clone, Debug)]
pub(crate) struct ProgressDrawState {
    /// The lines to print (can contain ANSI codes)
    pub lines: Vec<String>,
    /// The number of lines that shouldn't be reaped by the next tick.
    pub orphan_lines: usize,
    /// True if drawing should be forced.
    pub force_draw: bool,
    /// True if we should move the cursor up when possible instead of clearing lines.
    pub move_cursor: bool,
    /// Controls how the multi progress is aligned if some of its progress bars get removed, default is `Top`
    pub alignment: MultiProgressAlignment,
}

impl ProgressDrawState {
    pub(crate) fn new(lines: Vec<String>, force_draw: bool) -> Self {
        Self {
            lines,
            orphan_lines: 0,
            force_draw,
            move_cursor: false,
            alignment: Default::default(),
        }
    }

    pub fn draw_to_term(&self, term: &Term) -> io::Result<()> {
        let len = self.lines.len();
        for (idx, line) in self.lines.iter().enumerate() {
            if idx + 1 != len {
                term.write_line(line)?;
            } else {
                // Don't append a '\n' if this is the last line
                term.write_str(line)?;
            }
        }
        Ok(())
    }
}

/// Vertical alignment of a multi progress.
///
/// The alignment controls how the multi progress is aligned if some of its progress bars get removed.
/// E.g. `Top` alignment (default), when _progress bar 2_ is removed:
/// ```ignore
/// [0/100] progress bar 1        [0/100] progress bar 1
/// [0/100] progress bar 2   =>   [0/100] progress bar 3
/// [0/100] progress bar 3
/// ```
///
/// `Bottom` alignment
/// ```ignore
/// [0/100] progress bar 1
/// [0/100] progress bar 2   =>   [0/100] progress bar 1
/// [0/100] progress bar 3        [0/100] progress bar 3
/// ```
#[derive(Debug, Copy, Clone)]
pub enum MultiProgressAlignment {
    Top,
    Bottom,
}

impl Default for MultiProgressAlignment {
    fn default() -> Self {
        Self::Top
    }
}

/// Fork of console::clear_last_lines that assumes that the last line doesn't contain a '\n'
fn clear_last_lines(term: &Term, n: usize) -> io::Result<()> {
    term.move_cursor_up(n.saturating_sub(1))?;
    for i in 0..n {
        term.clear_line()?;
        if i + 1 != n {
            term.move_cursor_down(1)?;
        }
    }
    term.move_cursor_up(n.saturating_sub(1))?;
    Ok(())
}
<|endoftext|>


<commit_before>use failure::Error;
use hexplay::{self, HexViewBuilder, CODEPAGE_ASCII};
use metagoblin;
use metagoblin::Object;
use prettytable::Cell;
use prettytable::Row;

use format::*;
use Opt;

const SCALE: [char; 100] = [' '; 100];

fn minmaxscale(n: usize, min: usize, max: usize) -> usize {
    (SCALE.len() as f32 * (n as f32 - min as f32) / (max as f32 - min as f32)).floor() as usize
}

pub struct Meta<'a> {
    analysis: metagoblin::Analysis,
    bytes: &'a [u8],
    args: Opt,
}

impl<'a> Meta<'a> {
    pub fn new(object: &Object<'a>, bytes: &'a [u8], args: Opt) -> Self {
        let analysis = metagoblin::Analysis::new(object);
        Meta {
            analysis,
            bytes,
            args,
        }
    }
    pub fn print_ranges(&self) -> Result<(), Error> {
        let buffer = &self.bytes;
        let mut franges = self.analysis.franges.iter().collect::<Vec<_>>();
        franges.sort_by(|&(ref r1, _), &(ref r2, _)| r2.len().cmp(&r1.len()));
        let spaces = (0..100)
            .map(|e| if e % 2 == 0 { "0" } else { "1" })
            .collect::<String>();
        let mut table =
            new_table(row![b->"Name", b->"Tag", b->"Range", b->"Percent", b->"Size", b->spaces]);

        for &(ref range, ref data) in &franges {
            let size = range.len() - 1;
            if size == 0 {
                continue;
            }
            let scaled_min = minmaxscale(range.min as usize, 0, buffer.len());
            let scaled_max = minmaxscale(range.max as usize, 0, buffer.len());
            let scaled_size = minmaxscale(size as usize, 0, buffer.len());
            let mut chars = SCALE.clone();
            for i in scaled_min..scaled_max {
                chars[i] = '-';
            }
            chars[scaled_min] = '|';
            if scaled_max != 0 {
                chars[scaled_max - 1] = '|';
            }
            let chars = &chars.iter().cloned().collect::<String>();
            let name_cell = string_cell(&self.args, data.name().unwrap_or("None"));
            let range_cell = Cell::new(&format!("{:#x}..{:#x}", range.min, range.max));
            let size_cell = Cell::new(&format!("{}", size));
            table.add_row(Row::new(vec![
                name_cell,
                Cell::new(&format!("{:?}", data.tag)),
                range_cell,
                Cell::new(&format!("{}%", scaled_size)),
                size_cell,
                Cell::new(chars),
            ]));
        }

        table.print_tty(true);
        Ok(())
    }
    pub fn print_hex(&self) -> Result<(), Error> {
        let analysis = &self.analysis;
        let table = HexViewBuilder::new(&self.bytes)
            .row_width(16)
            .codepage(CODEPAGE_ASCII);

        let mut colors = Vec::new();
        use metagoblin::Tag::*;
        let franges = analysis.franges.iter().collect::<Vec<_>>();
        for window in franges.windows(2) {
            let (range, metadata) = window[0];
            let (_next, _) = window[1];

            // 0xde0..0xde8(8) -> Code - ".init_array" : RW
            // 0xde0..0x1000(544) -> Unknown - "PT_GNU_RELRO"
            // 0xde0..0x1030(592) -> Code - "PT_LOAD" : RW
            // 0xde8..0xdf0(8) -> Code - ".fini_array" : RW

            // 0xdf0..0xfd0(480) -> Meta - ".dynamic"
            // 0xfd0..0x1000(48) -> Code - ".got" : RW
            // 0x1000..0x1020(32) -> Code - ".got.plt" : RW
            // 0x1020..0x1030(16) -> Code - ".data" : RW
            // 0x1030..0x1038(8) -> Zero - ".bss" : RW
            // 0x1030..0x1041(17) -> Code - ".comment" :

            // we can rely on sortedness and this invariant:
            // if range.max <= next.min then this is the most specific tag for this memory range
            // _because_ any previous range, range2, must have this invariant:
            // range2.min <= range.min && range2.max >= range.max
            // e.g., it covers more area, or is exactly equal
            // nevertheless, its still unclear what to do in the case of overlapping ranges

            // if range.min < next.min {}
            // truncate
            //if range.max >= next.min {}

            //            print!(
            //                "{:#x}..{:#x}({}) -> ",
            //                range.min,
            //                range.max,
            //                range.len() - 1
            //            );
            //            print!(
            //                "{:?} - {:?}",
            //                metadata.tag,
            //                metadata.name().unwrap_or("None")
            //            );
            //            if let &Some(ref segment) = &metadata.memory {
            //                print!(" : {}", segment.permissions);
            //            }
            //            println!("");
            let color = match metadata.tag {
                Code => hexplay::color::red(),
                ASCII => hexplay::color::yellow_bold(),
                SymbolTable => hexplay::color::blue_bold(),
                Relocation => hexplay::color::green(),
                Zero => hexplay::color::white_bold(),
                Meta => hexplay::color::magenta(),
                Unknown => hexplay::color::black(),
                _ => hexplay::color::white(),
            };
            colors.push((color, range.min as usize..range.max as usize));
        }

        let table = table.add_colors(colors);

        let view = if self.args.color {
            table.force_color().finish()
        } else {
            table.finish()
        };

        view.print()?;
        println!();
        Ok(())
    }
}
<commit_msg>Skip zero sized ranges (#25)<commit_after>use failure::Error;
use hexplay::{self, HexViewBuilder, CODEPAGE_ASCII};
use metagoblin;
use metagoblin::{Object, Tag};
use prettytable::Cell;
use prettytable::Row;

use format::*;
use Opt;

const SCALE: [char; 100] = [' '; 100];

fn minmaxscale(n: usize, min: usize, max: usize) -> usize {
    (SCALE.len() as f32 * (n as f32 - min as f32) / (max as f32 - min as f32)).floor() as usize
}

pub struct Meta<'a> {
    analysis: metagoblin::Analysis,
    bytes: &'a [u8],
    args: Opt,
}

impl<'a> Meta<'a> {
    pub fn new(object: &Object<'a>, bytes: &'a [u8], args: Opt) -> Self {
        let analysis = metagoblin::Analysis::new(object);
        Meta {
            analysis,
            bytes,
            args,
        }
    }
    pub fn print_ranges(&self) -> Result<(), Error> {
        let buffer = &self.bytes;
        let mut franges = self.analysis.franges.iter().collect::<Vec<_>>();
        franges.sort_by(|&(ref r1, _), &(ref r2, _)| r2.len().cmp(&r1.len()));
        let spaces = (0..100)
            .map(|e| if e % 2 == 0 { "0" } else { "1" })
            .collect::<String>();
        let mut table =
            new_table(row![b->"Name", b->"Tag", b->"Range", b->"Percent", b->"Size", b->spaces]);

        for &(ref range, ref data) in &franges {
            if let Tag::Zero = data.tag {
                continue;
            }
            let size = range.len() - 1;
            if size == 0 {
                continue;
            }
            let scaled_min = minmaxscale(range.min as usize, 0, buffer.len());
            let scaled_max = minmaxscale(range.max as usize, 0, buffer.len());
            let scaled_size = minmaxscale(size as usize, 0, buffer.len());
            let mut chars = SCALE.clone();
            for i in scaled_min..scaled_max {
                chars[i] = '-';
            }
            chars[scaled_min] = '|';
            if scaled_max != 0 {
                chars[scaled_max - 1] = '|';
            }
            let chars = &chars.iter().cloned().collect::<String>();
            let name_cell = string_cell(&self.args, data.name().unwrap_or("None"));
            let range_cell = Cell::new(&format!("{:#x}..{:#x}", range.min, range.max));
            let size_cell = Cell::new(&format!("{}", size));
            table.add_row(Row::new(vec![
                name_cell,
                Cell::new(&format!("{:?}", data.tag)),
                range_cell,
                Cell::new(&format!("{}%", scaled_size)),
                size_cell,
                Cell::new(chars),
            ]));
        }

        table.print_tty(true);
        Ok(())
    }
    pub fn print_hex(&self) -> Result<(), Error> {
        let analysis = &self.analysis;
        let table = HexViewBuilder::new(&self.bytes)
            .row_width(16)
            .codepage(CODEPAGE_ASCII);

        let mut colors = Vec::new();
        use metagoblin::Tag::*;
        let franges = analysis.franges.iter().collect::<Vec<_>>();
        for window in franges.windows(2) {
            let (range, metadata) = window[0];
            let (_next, _) = window[1];

            // 0xde0..0xde8(8) -> Code - ".init_array" : RW
            // 0xde0..0x1000(544) -> Unknown - "PT_GNU_RELRO"
            // 0xde0..0x1030(592) -> Code - "PT_LOAD" : RW
            // 0xde8..0xdf0(8) -> Code - ".fini_array" : RW

            // 0xdf0..0xfd0(480) -> Meta - ".dynamic"
            // 0xfd0..0x1000(48) -> Code - ".got" : RW
            // 0x1000..0x1020(32) -> Code - ".got.plt" : RW
            // 0x1020..0x1030(16) -> Code - ".data" : RW
            // 0x1030..0x1038(8) -> Zero - ".bss" : RW
            // 0x1030..0x1041(17) -> Code - ".comment" :

            // we can rely on sortedness and this invariant:
            // if range.max <= next.min then this is the most specific tag for this memory range
            // _because_ any previous range, range2, must have this invariant:
            // range2.min <= range.min && range2.max >= range.max
            // e.g., it covers more area, or is exactly equal
            // nevertheless, its still unclear what to do in the case of overlapping ranges

            // if range.min < next.min {}
            // truncate
            //if range.max >= next.min {}

            //            print!(
            //                "{:#x}..{:#x}({}) -> ",
            //                range.min,
            //                range.max,
            //                range.len() - 1
            //            );
            //            print!(
            //                "{:?} - {:?}",
            //                metadata.tag,
            //                metadata.name().unwrap_or("None")
            //            );
            //            if let &Some(ref segment) = &metadata.memory {
            //                print!(" : {}", segment.permissions);
            //            }
            //            println!("");
            let color = match metadata.tag {
                Code => hexplay::color::red(),
                ASCII => hexplay::color::yellow_bold(),
                SymbolTable => hexplay::color::blue_bold(),
                Relocation => hexplay::color::green(),
                Zero => hexplay::color::white_bold(),
                Meta => hexplay::color::magenta(),
                Unknown => hexplay::color::black(),
                _ => hexplay::color::white(),
            };
            colors.push((color, range.min as usize..range.max as usize));
        }

        let table = table.add_colors(colors);

        let view = if self.args.color {
            table.force_color().finish()
        } else {
            table.finish()
        };

        view.print()?;
        println!();
        Ok(())
    }
}
<|endoftext|>


<commit_before>pub const MAX_CHANNELS: usize = 8;

pub struct Frame {
  pub header: Header,
  pub footer: Footer,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum ChannelAssignment {
  Independent = 0,
  LeftSide    = 1,
  RightSide   = 2,
  MiddleSide  = 3,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum NumberType {
  Frame(u32),
  Sample(u64),
}

#[derive(Debug, PartialEq, Eq)]
pub struct Header {
  pub block_size: u32,
  pub sample_rate: u32,
  pub channels: u8,
  pub channel_assignment: ChannelAssignment,
  pub bits_per_sample: usize,
  pub number: NumberType,
  pub crc: u8,
}

#[derive(Debug, PartialEq, Eq)]
pub struct Footer(pub u16);
<commit_msg>Add documentation about MAX_CHANNELS<commit_after>/// Maximum number of channels supported in the FLAC format.
pub const MAX_CHANNELS: usize = 8;

pub struct Frame {
  pub header: Header,
  pub footer: Footer,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum ChannelAssignment {
  Independent = 0,
  LeftSide    = 1,
  RightSide   = 2,
  MiddleSide  = 3,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum NumberType {
  Frame(u32),
  Sample(u64),
}

#[derive(Debug, PartialEq, Eq)]
pub struct Header {
  pub block_size: u32,
  pub sample_rate: u32,
  pub channels: u8,
  pub channel_assignment: ChannelAssignment,
  pub bits_per_sample: usize,
  pub number: NumberType,
  pub crc: u8,
}

#[derive(Debug, PartialEq, Eq)]
pub struct Footer(pub u16);
<|endoftext|>


<commit_before>use std::slice;
use std::io;

use image;
use image::ImageResult;
use image::ImageDecoder;
use color;

use super::lzw::LZWReader;

const IMAGEDESCRIPTOR: u8 = 0x2C;
const EXTENSION: u8 = 0x21;
const APPLICATION: u8 = 0xFF;
const GRAPHICCONTROL: u8 = 0xF9;
const COMMENT: u8 = 0xFE;
const TRAILER: u8 = 0x3B;

/// The Representation of a GIF decoder
pub struct GIFDecoder <R> {
    r: R,

    width: u16,
    height: u16,

    global_table: [(u8, u8, u8), ..256],
    local_table: Option<Vec<(u8, u8, u8)>>,

    delay: u16,
    image: Vec<u8>,

    global_backgroud_index: Option<u8>,
    local_transparent_index: Option<u8>,

    have_header: bool,
    decoded_rows: u32,
}

impl<R: Reader> GIFDecoder<R> {
    /// Create a new GIFDecoder from the Reader ```r```.
    /// This function takes ownership of the Reader.
    pub fn new(r: R) -> GIFDecoder<R> {
        GIFDecoder {
            r: r,

            width: 0,
            height: 0,

            global_table: [(0u8, 0u8, 0u8), ..256],
            local_table: None,

            delay: 0,
            image: Vec::new(),

            global_backgroud_index: None,
            local_transparent_index: None,

            have_header: false,
            decoded_rows: 0,
        }
    }

    ///Returns the display delay in 100th's of a second for the currently
    ///decoded image.
    pub fn delay(&mut self) -> ImageResult<u16> {

        let _ = try!(self.read_metadata());

        Ok(self.delay)
    }

    fn read_header(&mut self) -> ImageResult<()> {
        let signature = try!(self.r.read_exact(3));
        let version   = try!(self.r.read_exact(3));

        if signature.as_slice() != "GIF".as_bytes() {
            Err(image::ImageError::FormatError("GIF signature not found.".to_string()))
        } else if version.as_slice() != "87a".as_bytes() &&
                  version.as_slice() != "89a".as_bytes() {
            Err(image::ImageError::UnsupportedError(
                format!("GIF version {} is not supported.", version)
            ))
        } else {
            Ok(())
        }
    }

    fn read_block(&mut self) -> io::IoResult<Vec<u8>> {
        let size = try!(self.r.read_u8());
        Ok(try!(self.r.read_exact(size as uint)))
    }

    fn read_image_data(&mut self) -> ImageResult<Vec<u8>> {
        let minimum_code_size = try!(self.r.read_u8());

        if minimum_code_size > 8 {
            return Err(image::ImageError::FormatError(format!("Invalid code size {}.", minimum_code_size)))
        }

        let mut data = Vec::new();
        loop {
            let b = try!(self.read_block());

            if b.len() == 0 {
                break
            }

            data = data + b;
        }

        let m = io::MemReader::new(data);
        let mut lzw = LZWReader::new(m, minimum_code_size);
        let b = lzw.read_to_end().unwrap();

        Ok(b)
    }

    fn read_image_descriptor(&mut self) -> ImageResult<()> {
        let image_left   = try!(self.r.read_le_u16());
        let image_top    = try!(self.r.read_le_u16());
        let image_width  = try!(self.r.read_le_u16());
        let image_height = try!(self.r.read_le_u16());

        let fields = try!(self.r.read_u8());

        let local_table = fields & 80 != 0;
        let interlace   = fields & 40 != 0;
        let table_size  = fields & 7;

        if interlace {
            return Err(image::ImageError::UnsupportedError("Interlaced images are not supported.".to_string()))
        }

        if local_table {
            let n   = 1 << (table_size + 1) as uint;
            let buf = try!(self.r.read_exact(3 * n));
            let mut b = Vec::from_elem(n, (0u8, 0u8, 0u8));

            for (i, rgb) in buf.as_slice().chunks(3).enumerate() {
                b.as_mut_slice()[i] = (rgb[0], rgb[1], rgb[2]);
            }

            self.local_table = Some(b);
        }

        let indices = try!(self.read_image_data());

        {

            let trans_index = if self.local_transparent_index.is_some() {
                self.local_transparent_index
            } else {
                self.global_backgroud_index
            };

            let table = if self.local_table.is_some() {
                self.local_table.as_ref().unwrap().as_slice()
            } else {
                self.global_table.as_slice()
            };

            expand_image(
                table,
                indices.as_slice(),
                image_top as uint,
                image_left as uint,
                image_width as uint,
                image_height as uint,
                self.width as uint * 3,
                trans_index,
                self.image.as_mut_slice()
            );
        }

        self.local_table = None;
        self.local_transparent_index = None;

        Ok(())
    }

    fn read_extension(&mut self) -> ImageResult<()> {
        let identifier = try!(self.r.read_u8());

        match identifier {
            APPLICATION    => try!(self.read_application_extension()),
            GRAPHICCONTROL => try!(self.read_graphic_control_extension()),
            COMMENT 	   => try!(self.read_comment_extension()),
            _              => return Err(image::ImageError::UnsupportedError(
                                  format!("Identifier {} is not supported.", identifier))
                              )
        }

        Ok(())
    }

    fn read_comment_extension(&mut self) -> ImageResult<()> {
        loop {
            let b = try!(self.read_block());

            if b.len() == 0 {
                break
            }
        }

        Ok(())
    }

    fn read_graphic_control_extension(&mut self) -> ImageResult<()> {
        let size   = try!(self.r.read_u8());
        assert!(size == 4);

        let fields = try!(self.r.read_u8());
        self.delay = try!(self.r.read_le_u16());
        let trans  = try!(self.r.read_u8());

        if fields & 1 != 0 {
            self.local_transparent_index = Some(trans);
        }

        let _disposal = (fields & 0x1C) >> 2;
        let _term = try!(self.r.read_u8());

        Ok(())
    }

    fn read_application_extension(&mut self) -> ImageResult<()> {
        let size = try!(self.r.read_u8());
        let _ = try!(self.r.read_exact(size as uint));

        loop {
            let b = try!(self.read_block());

            if b.len() == 0 {
                break
            }
        }

        Ok(())
    }

    fn read_logical_screen_descriptor(&mut self) -> ImageResult<()> {
        self.width  = try!(self.r.read_le_u16());
        self.height = try!(self.r.read_le_u16());
        self.image  = Vec::from_elem(self.width as uint * self.height as uint * 3, 0u8);

        let fields = try!(self.r.read_u8());

        let global_table = fields & 0x80 != 0;

        let entries = if global_table {
            1 << ((fields & 7) + 1) as uint
        } else {
            0u
        };

        let b = try!(self.r.read_u8());
        if global_table {
            self.global_backgroud_index = Some(b);
        }

        let _aspect_ratio = try!(self.r.read_u8());

        let buf = try!(self.r.read_exact(3 * entries));

        for (i, rgb) in buf.as_slice().chunks(3).enumerate() {
            self.global_table.as_mut_slice()[i] = (rgb[0], rgb[1], rgb[2]);
        }

        Ok(())
    }

    fn read_metadata(&mut self) -> ImageResult<()> {
        if !self.have_header {
            let _ = try!(self.read_header());
            let _ = try!(self.read_logical_screen_descriptor());
            self.have_header = true;
        }

        Ok(())
    }
}

impl<R: Reader> ImageDecoder for GIFDecoder<R> {
    fn dimensions(&mut self) -> ImageResult<(u32, u32)> {
        let _ = try!(self.read_metadata());
        Ok((self.width as u32, self.height as u32))
    }

    fn colortype(&mut self) -> ImageResult<color::ColorType> {
        let _ = try!(self.read_metadata());
        Ok(color::ColorType::RGB(8))
    }

    fn row_len(&mut self) -> ImageResult<uint> {
        let _ = try!(self.read_metadata());
        Ok(3 * self.width as uint)
    }

    fn read_scanline(&mut self, buf: &mut [u8]) -> ImageResult<u32> {
        let _ = try!(self.read_metadata());
        if self.decoded_rows == self.height as u32 || self.image.len() == 0 {
            let _ = try!(self.read_image());
        }

        let rlen  = buf.len();
        let slice = self.image.slice(self.decoded_rows as uint * rlen,
        self.decoded_rows as uint * rlen + rlen);

        slice::bytes::copy_memory(buf, slice);
        self.decoded_rows += 1;

        Ok(self.decoded_rows)
    }

    fn read_image(&mut self) -> ImageResult<Vec<u8>> {
        let _ = try!(self.read_metadata());
        loop {
            let block = try!(self.r.read_u8());

            match block {
                EXTENSION => try!(self.read_extension()),
                IMAGEDESCRIPTOR => {
                    let _ = try!(self.read_image_descriptor());
                    return Ok(self.image.clone())
                }
                TRAILER => break,
                _       => return Err(image::ImageError::UnsupportedError(
                            format!("Block type {} is not supported.", block))
                        )
            }
        }

        Err(image::ImageError::ImageEnd)
    }
}

fn expand_image(palette: &[(u8, u8, u8)],
                indices: &[u8],
                y0: uint,
                x0: uint,
                width: uint,
                height: uint,
                stride: uint,
                trans_index: Option<u8>,
                image: &mut [u8]) {

    for y in range(0, height) {
        for x in range(0, width) {
            let index = indices[y * width + x];

            if trans_index == Some(index) {
                continue
            }

            let (r, g, b) = palette[index as uint];

            image[(y0 + y) * stride + x0 * 3 + x * 3 + 0] = r;
            image[(y0 + y) * stride + x0 * 3 + x * 3 + 1] = g;
            image[(y0 + y) * stride + x0 * 3 + x * 3 + 2] = b;
        }
    }
}
<commit_msg>Use space in comments in gif<commit_after>use std::slice;
use std::io;

use image;
use image::ImageResult;
use image::ImageDecoder;
use color;

use super::lzw::LZWReader;

const IMAGEDESCRIPTOR: u8 = 0x2C;
const EXTENSION: u8 = 0x21;
const APPLICATION: u8 = 0xFF;
const GRAPHICCONTROL: u8 = 0xF9;
const COMMENT: u8 = 0xFE;
const TRAILER: u8 = 0x3B;

/// The Representation of a GIF decoder
pub struct GIFDecoder <R> {
    r: R,

    width: u16,
    height: u16,

    global_table: [(u8, u8, u8), ..256],
    local_table: Option<Vec<(u8, u8, u8)>>,

    delay: u16,
    image: Vec<u8>,

    global_backgroud_index: Option<u8>,
    local_transparent_index: Option<u8>,

    have_header: bool,
    decoded_rows: u32,
}

impl<R: Reader> GIFDecoder<R> {
    /// Create a new GIFDecoder from the Reader ```r```.
    /// This function takes ownership of the Reader.
    pub fn new(r: R) -> GIFDecoder<R> {
        GIFDecoder {
            r: r,

            width: 0,
            height: 0,

            global_table: [(0u8, 0u8, 0u8), ..256],
            local_table: None,

            delay: 0,
            image: Vec::new(),

            global_backgroud_index: None,
            local_transparent_index: None,

            have_header: false,
            decoded_rows: 0,
        }
    }

    /// Returns the display delay in 100th's of a second for the currently
    /// decoded image.
    pub fn delay(&mut self) -> ImageResult<u16> {

        let _ = try!(self.read_metadata());

        Ok(self.delay)
    }

    fn read_header(&mut self) -> ImageResult<()> {
        let signature = try!(self.r.read_exact(3));
        let version   = try!(self.r.read_exact(3));

        if signature.as_slice() != "GIF".as_bytes() {
            Err(image::ImageError::FormatError("GIF signature not found.".to_string()))
        } else if version.as_slice() != "87a".as_bytes() &&
                  version.as_slice() != "89a".as_bytes() {
            Err(image::ImageError::UnsupportedError(
                format!("GIF version {} is not supported.", version)
            ))
        } else {
            Ok(())
        }
    }

    fn read_block(&mut self) -> io::IoResult<Vec<u8>> {
        let size = try!(self.r.read_u8());
        Ok(try!(self.r.read_exact(size as uint)))
    }

    fn read_image_data(&mut self) -> ImageResult<Vec<u8>> {
        let minimum_code_size = try!(self.r.read_u8());

        if minimum_code_size > 8 {
            return Err(image::ImageError::FormatError(format!("Invalid code size {}.", minimum_code_size)))
        }

        let mut data = Vec::new();
        loop {
            let b = try!(self.read_block());

            if b.len() == 0 {
                break
            }

            data = data + b;
        }

        let m = io::MemReader::new(data);
        let mut lzw = LZWReader::new(m, minimum_code_size);
        let b = lzw.read_to_end().unwrap();

        Ok(b)
    }

    fn read_image_descriptor(&mut self) -> ImageResult<()> {
        let image_left   = try!(self.r.read_le_u16());
        let image_top    = try!(self.r.read_le_u16());
        let image_width  = try!(self.r.read_le_u16());
        let image_height = try!(self.r.read_le_u16());

        let fields = try!(self.r.read_u8());

        let local_table = fields & 80 != 0;
        let interlace   = fields & 40 != 0;
        let table_size  = fields & 7;

        if interlace {
            return Err(image::ImageError::UnsupportedError("Interlaced images are not supported.".to_string()))
        }

        if local_table {
            let n   = 1 << (table_size + 1) as uint;
            let buf = try!(self.r.read_exact(3 * n));
            let mut b = Vec::from_elem(n, (0u8, 0u8, 0u8));

            for (i, rgb) in buf.as_slice().chunks(3).enumerate() {
                b.as_mut_slice()[i] = (rgb[0], rgb[1], rgb[2]);
            }

            self.local_table = Some(b);
        }

        let indices = try!(self.read_image_data());

        {

            let trans_index = if self.local_transparent_index.is_some() {
                self.local_transparent_index
            } else {
                self.global_backgroud_index
            };

            let table = if self.local_table.is_some() {
                self.local_table.as_ref().unwrap().as_slice()
            } else {
                self.global_table.as_slice()
            };

            expand_image(
                table,
                indices.as_slice(),
                image_top as uint,
                image_left as uint,
                image_width as uint,
                image_height as uint,
                self.width as uint * 3,
                trans_index,
                self.image.as_mut_slice()
            );
        }

        self.local_table = None;
        self.local_transparent_index = None;

        Ok(())
    }

    fn read_extension(&mut self) -> ImageResult<()> {
        let identifier = try!(self.r.read_u8());

        match identifier {
            APPLICATION    => try!(self.read_application_extension()),
            GRAPHICCONTROL => try!(self.read_graphic_control_extension()),
            COMMENT 	   => try!(self.read_comment_extension()),
            _              => return Err(image::ImageError::UnsupportedError(
                                  format!("Identifier {} is not supported.", identifier))
                              )
        }

        Ok(())
    }

    fn read_comment_extension(&mut self) -> ImageResult<()> {
        loop {
            let b = try!(self.read_block());

            if b.len() == 0 {
                break
            }
        }

        Ok(())
    }

    fn read_graphic_control_extension(&mut self) -> ImageResult<()> {
        let size   = try!(self.r.read_u8());
        assert!(size == 4);

        let fields = try!(self.r.read_u8());
        self.delay = try!(self.r.read_le_u16());
        let trans  = try!(self.r.read_u8());

        if fields & 1 != 0 {
            self.local_transparent_index = Some(trans);
        }

        let _disposal = (fields & 0x1C) >> 2;
        let _term = try!(self.r.read_u8());

        Ok(())
    }

    fn read_application_extension(&mut self) -> ImageResult<()> {
        let size = try!(self.r.read_u8());
        let _ = try!(self.r.read_exact(size as uint));

        loop {
            let b = try!(self.read_block());

            if b.len() == 0 {
                break
            }
        }

        Ok(())
    }

    fn read_logical_screen_descriptor(&mut self) -> ImageResult<()> {
        self.width  = try!(self.r.read_le_u16());
        self.height = try!(self.r.read_le_u16());
        self.image  = Vec::from_elem(self.width as uint * self.height as uint * 3, 0u8);

        let fields = try!(self.r.read_u8());

        let global_table = fields & 0x80 != 0;

        let entries = if global_table {
            1 << ((fields & 7) + 1) as uint
        } else {
            0u
        };

        let b = try!(self.r.read_u8());
        if global_table {
            self.global_backgroud_index = Some(b);
        }

        let _aspect_ratio = try!(self.r.read_u8());

        let buf = try!(self.r.read_exact(3 * entries));

        for (i, rgb) in buf.as_slice().chunks(3).enumerate() {
            self.global_table.as_mut_slice()[i] = (rgb[0], rgb[1], rgb[2]);
        }

        Ok(())
    }

    fn read_metadata(&mut self) -> ImageResult<()> {
        if !self.have_header {
            let _ = try!(self.read_header());
            let _ = try!(self.read_logical_screen_descriptor());
            self.have_header = true;
        }

        Ok(())
    }
}

impl<R: Reader> ImageDecoder for GIFDecoder<R> {
    fn dimensions(&mut self) -> ImageResult<(u32, u32)> {
        let _ = try!(self.read_metadata());
        Ok((self.width as u32, self.height as u32))
    }

    fn colortype(&mut self) -> ImageResult<color::ColorType> {
        let _ = try!(self.read_metadata());
        Ok(color::ColorType::RGB(8))
    }

    fn row_len(&mut self) -> ImageResult<uint> {
        let _ = try!(self.read_metadata());
        Ok(3 * self.width as uint)
    }

    fn read_scanline(&mut self, buf: &mut [u8]) -> ImageResult<u32> {
        let _ = try!(self.read_metadata());
        if self.decoded_rows == self.height as u32 || self.image.len() == 0 {
            let _ = try!(self.read_image());
        }

        let rlen  = buf.len();
        let slice = self.image.slice(self.decoded_rows as uint * rlen,
        self.decoded_rows as uint * rlen + rlen);

        slice::bytes::copy_memory(buf, slice);
        self.decoded_rows += 1;

        Ok(self.decoded_rows)
    }

    fn read_image(&mut self) -> ImageResult<Vec<u8>> {
        let _ = try!(self.read_metadata());
        loop {
            let block = try!(self.r.read_u8());

            match block {
                EXTENSION => try!(self.read_extension()),
                IMAGEDESCRIPTOR => {
                    let _ = try!(self.read_image_descriptor());
                    return Ok(self.image.clone())
                }
                TRAILER => break,
                _       => return Err(image::ImageError::UnsupportedError(
                            format!("Block type {} is not supported.", block))
                        )
            }
        }

        Err(image::ImageError::ImageEnd)
    }
}

fn expand_image(palette: &[(u8, u8, u8)],
                indices: &[u8],
                y0: uint,
                x0: uint,
                width: uint,
                height: uint,
                stride: uint,
                trans_index: Option<u8>,
                image: &mut [u8]) {

    for y in range(0, height) {
        for x in range(0, width) {
            let index = indices[y * width + x];

            if trans_index == Some(index) {
                continue
            }

            let (r, g, b) = palette[index as uint];

            image[(y0 + y) * stride + x0 * 3 + x * 3 + 0] = r;
            image[(y0 + y) * stride + x0 * 3 + x * 3 + 1] = g;
            image[(y0 + y) * stride + x0 * 3 + x * 3 + 2] = b;
        }
    }
}
<|endoftext|>


<commit_before>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at https://mozilla.org/MPL/2.0/. */

<%namespace name="helpers" file="/helpers.mako.rs" />

<% data.new_style_struct("InheritedUI", inherited=True, gecko_name="UI") %>

${helpers.predefined_type(
    "cursor",
    "Cursor",
    "computed::Cursor::auto()",
    engines="gecko servo-2013 servo-2020",
    initial_specified_value="specified::Cursor::auto()",
    animation_value_type="discrete",
    spec="https://drafts.csswg.org/css-ui/#cursor",
)}

// NB: `pointer-events: auto` (and use of `pointer-events` in anything that isn't SVG, in fact)
// is nonstandard, slated for CSS4-UI.
// TODO(pcwalton): SVG-only values.
${helpers.single_keyword(
    "pointer-events",
    "auto none",
    engines="gecko servo-2013 servo-2020",
    animation_value_type="discrete",
    extra_gecko_values="visiblepainted visiblefill visiblestroke visible painted fill stroke all",
    spec="https://www.w3.org/TR/SVG11/interact.html#PointerEventsProperty",
)}

${helpers.single_keyword(
    "-moz-user-input",
    "auto none",
    engines="gecko",
    gecko_ffi_name="mUserInput",
    gecko_enum_prefix="StyleUserInput",
    animation_value_type="discrete",
    spec="Nonstandard (https://developer.mozilla.org/en-US/docs/Web/CSS/-moz-user-input)",
)}

${helpers.single_keyword(
    "-moz-user-modify",
    "read-only read-write write-only",
    engines="gecko",
    gecko_ffi_name="mUserModify",
    gecko_enum_prefix="StyleUserModify",
    needs_conversion=True,
    animation_value_type="discrete",
    spec="Nonstandard (https://developer.mozilla.org/en-US/docs/Web/CSS/-moz-user-modify)",
)}

${helpers.single_keyword(
    "-moz-user-focus",
    "none ignore normal select-after select-before select-menu select-same select-all",
    engines="gecko",
    gecko_ffi_name="mUserFocus",
    gecko_enum_prefix="StyleUserFocus",
    animation_value_type="discrete",
    spec="Nonstandard (https://developer.mozilla.org/en-US/docs/Web/CSS/-moz-user-focus)",
)}

${helpers.predefined_type(
    "caret-color",
    "ColorOrAuto",
    "generics::color::ColorOrAuto::Auto",
    engines="gecko",
    spec="https://drafts.csswg.org/css-ui/#caret-color",
    animation_value_type="AnimatedCaretColor",
    boxed=True,
    ignored_when_colors_disabled=True,
)}

${helpers.predefined_type(
    "scrollbar-color",
    "ui::ScrollbarColor",
    "Default::default()",
    engines="gecko",
    spec="https://drafts.csswg.org/css-scrollbars-1/#scrollbar-color",
    gecko_pref="layout.css.scrollbar-color.enabled",
    # Surprisingly, yes the computed value of scrollbar-color has no effect on
    # Gecko scrollbar elements, since the value only matters on the scrollable
    # element itself.
    has_effect_on_gecko_scrollbars=False,
    animation_value_type="ScrollbarColor",
    boxed=True,
    ignored_when_colors_disabled=True,
    enabled_in="chrome",
)}
<commit_msg>style: convert NS_STYLE_POINTER_EVENTS_* to an enum class in nsStyleConsts.h<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at https://mozilla.org/MPL/2.0/. */

<%namespace name="helpers" file="/helpers.mako.rs" />

<% data.new_style_struct("InheritedUI", inherited=True, gecko_name="UI") %>

${helpers.predefined_type(
    "cursor",
    "Cursor",
    "computed::Cursor::auto()",
    engines="gecko servo-2013 servo-2020",
    initial_specified_value="specified::Cursor::auto()",
    animation_value_type="discrete",
    spec="https://drafts.csswg.org/css-ui/#cursor",
)}

// NB: `pointer-events: auto` (and use of `pointer-events` in anything that isn't SVG, in fact)
// is nonstandard, slated for CSS4-UI.
// TODO(pcwalton): SVG-only values.
${helpers.single_keyword(
    "pointer-events",
    "auto none",
    engines="gecko servo-2013 servo-2020",
    animation_value_type="discrete",
    extra_gecko_values="visiblepainted visiblefill visiblestroke visible painted fill stroke all",
    spec="https://www.w3.org/TR/SVG11/interact.html#PointerEventsProperty",
    gecko_enum_prefix="StylePointerEvents",
)}

${helpers.single_keyword(
    "-moz-user-input",
    "auto none",
    engines="gecko",
    gecko_ffi_name="mUserInput",
    gecko_enum_prefix="StyleUserInput",
    animation_value_type="discrete",
    spec="Nonstandard (https://developer.mozilla.org/en-US/docs/Web/CSS/-moz-user-input)",
)}

${helpers.single_keyword(
    "-moz-user-modify",
    "read-only read-write write-only",
    engines="gecko",
    gecko_ffi_name="mUserModify",
    gecko_enum_prefix="StyleUserModify",
    needs_conversion=True,
    animation_value_type="discrete",
    spec="Nonstandard (https://developer.mozilla.org/en-US/docs/Web/CSS/-moz-user-modify)",
)}

${helpers.single_keyword(
    "-moz-user-focus",
    "none ignore normal select-after select-before select-menu select-same select-all",
    engines="gecko",
    gecko_ffi_name="mUserFocus",
    gecko_enum_prefix="StyleUserFocus",
    animation_value_type="discrete",
    spec="Nonstandard (https://developer.mozilla.org/en-US/docs/Web/CSS/-moz-user-focus)",
)}

${helpers.predefined_type(
    "caret-color",
    "ColorOrAuto",
    "generics::color::ColorOrAuto::Auto",
    engines="gecko",
    spec="https://drafts.csswg.org/css-ui/#caret-color",
    animation_value_type="AnimatedCaretColor",
    boxed=True,
    ignored_when_colors_disabled=True,
)}

${helpers.predefined_type(
    "scrollbar-color",
    "ui::ScrollbarColor",
    "Default::default()",
    engines="gecko",
    spec="https://drafts.csswg.org/css-scrollbars-1/#scrollbar-color",
    gecko_pref="layout.css.scrollbar-color.enabled",
    # Surprisingly, yes the computed value of scrollbar-color has no effect on
    # Gecko scrollbar elements, since the value only matters on the scrollable
    # element itself.
    has_effect_on_gecko_scrollbars=False,
    animation_value_type="ScrollbarColor",
    boxed=True,
    ignored_when_colors_disabled=True,
    enabled_in="chrome",
)}
<|endoftext|>


<commit_before>extern crate openssl;
extern crate openssl_probe;

use self::openssl::error::ErrorStack;
use self::openssl::hash::MessageDigest;
use self::openssl::nid::Nid;
use self::openssl::pkcs12::{ParsedPkcs12, Pkcs12};
use self::openssl::pkey::PKey;
use self::openssl::ssl::{
    self, MidHandshakeSslStream, SslAcceptor, SslConnector, SslContextBuilder, SslMethod,
    SslVerifyMode,
};
use self::openssl::x509::{X509, X509VerifyResult};
use std::error;
use std::fmt;
use std::io;
use std::sync::{Once, ONCE_INIT};

use {Protocol, TlsAcceptorBuilder, TlsConnectorBuilder};
use self::openssl::pkey::Private;

#[cfg(have_min_max_version)]
fn supported_protocols(
    min: Option<Protocol>,
    max: Option<Protocol>,
    ctx: &mut SslContextBuilder,
) -> Result<(), ErrorStack> {
    use self::openssl::ssl::SslVersion;

    fn cvt(p: Protocol) -> SslVersion {
        match p {
            Protocol::Sslv3 => SslVersion::SSL3,
            Protocol::Tlsv10 => SslVersion::TLS1,
            Protocol::Tlsv11 => SslVersion::TLS1_1,
            Protocol::Tlsv12 => SslVersion::TLS1_2,
            Protocol::__NonExhaustive => unreachable!(),
        }
    }

    ctx.set_min_proto_version(min.map(cvt))?;
    ctx.set_max_proto_version(max.map(cvt))?;

    Ok(())
}

#[cfg(not(have_min_max_version))]
fn supported_protocols(
    min: Option<Protocol>,
    max: Option<Protocol>,
    ctx: &mut SslContextBuilder,
) -> Result<(), ErrorStack> {
    use self::openssl::ssl::SslOptions;

    let no_ssl_mask = SslOptions::NO_SSLV2
        | SslOptions::NO_SSLV3
        | SslOptions::NO_TLSV1
        | SslOptions::NO_TLSV1_1
        | SslOptions::NO_TLSV1_2;

    ctx.clear_options(no_ssl_mask);
    let mut options = SslOptions::empty();
    options |= match min {
        None => SslOptions::empty(),
        Some(Protocol::Sslv3) => SslOptions::NO_SSLV2,
        Some(Protocol::Tlsv10) => SslOptions::NO_SSLV2 | SslOptions::NO_SSLV3,
        Some(Protocol::Tlsv11) => {
            SslOptions::NO_SSLV2 | SslOptions::NO_SSLV3 | SslOptions::NO_TLSV1
        }
        Some(Protocol::Tlsv12) => {
            SslOptions::NO_SSLV2
                | SslOptions::NO_SSLV3
                | SslOptions::NO_TLSV1
                | SslOptions::NO_TLSV1_1
        }
        Some(Protocol::__NonExhaustive) => unreachable!(),
    };
    options |= match max {
        None | Some(Protocol::Tlsv12) => SslOptions::empty(),
        Some(Protocol::Tlsv11) => SslOptions::NO_TLSV1_2,
        Some(Protocol::Tlsv10) => SslOptions::NO_TLSV1_1 | SslOptions::NO_TLSV1_2,
        Some(Protocol::Sslv3) => {
            SslOptions::NO_TLSV1 | SslOptions::NO_TLSV1_1 | SslOptions::NO_TLSV1_2
        }
        Some(Protocol::__NonExhaustive) => unreachable!(),
    };

    ctx.set_options(options);

    Ok(())
}

fn init_trust() {
    static ONCE: Once = ONCE_INIT;
    ONCE.call_once(|| openssl_probe::init_ssl_cert_env_vars());
}

#[cfg(target_os = "android")]
fn load_android_root_certs(connector: &mut SslContextBuilder) -> Result<(), Error> {
    use std::fs;

    if let Ok(dir) = fs::read_dir("/system/etc/security/cacerts") {
        let certs = dir
            .filter_map(|r| r.ok())
            .filter_map(|e| fs::read(e.path()).ok())
            .filter_map(|b| X509::from_pem(&b).ok());
        for cert in certs {
            if let Err(err) = connector.cert_store_mut().add_cert(cert) {
                debug!("load_android_root_certs error: {:?}", err);
            }
        }
    }

    Ok(())
}

#[derive(Debug)]
pub enum Error {
    Normal(ErrorStack),
    Ssl(ssl::Error, X509VerifyResult),
}

impl error::Error for Error {
    fn description(&self) -> &str {
        match *self {
            Error::Normal(ref e) => error::Error::description(e),
            Error::Ssl(ref e, _) => error::Error::description(e),
        }
    }

    fn cause(&self) -> Option<&error::Error> {
        match *self {
            Error::Normal(ref e) => error::Error::cause(e),
            Error::Ssl(ref e, _) => error::Error::cause(e),
        }
    }
}

impl fmt::Display for Error {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            Error::Normal(ref e) => fmt::Display::fmt(e, fmt),
            Error::Ssl(ref e, X509VerifyResult::OK) => fmt::Display::fmt(e, fmt),
            Error::Ssl(ref e, v) => write!(fmt, "{} ({})", e, v),
        }
    }
}

impl From<ErrorStack> for Error {
    fn from(err: ErrorStack) -> Error {
        Error::Normal(err)
    }
}

#[derive(Clone)]
pub struct Identity {
    pkey: PKey<Private>,
    cert: X509,
    chain: Vec<X509>,
}

impl Identity {
    pub fn from_pkcs12(buf: &[u8], pass: &str) -> Result<Identity, Error> {
        let pkcs12 = Pkcs12::from_der(buf)?;
        let parsed = pkcs12.parse(pass)?;
        Ok(Identity {
            pkey: parsed.pkey,
            cert: parsed.cert,
            chain: parsed.chain.into_iter().flat_map(|x| x).collect(),
        })
    }
}

#[derive(Clone)]
pub struct Certificate(X509);

impl Certificate {
    pub fn from_der(buf: &[u8]) -> Result<Certificate, Error> {
        let cert = X509::from_der(buf)?;
        Ok(Certificate(cert))
    }

    pub fn from_pem(buf: &[u8]) -> Result<Certificate, Error> {
        let cert = X509::from_pem(buf)?;
        Ok(Certificate(cert))
    }

    pub fn to_der(&self) -> Result<Vec<u8>, Error> {
        let der = self.0.to_der()?;
        Ok(der)
    }
}

pub struct MidHandshakeTlsStream<S>(MidHandshakeSslStream<S>);

impl<S> fmt::Debug for MidHandshakeTlsStream<S>
where
    S: fmt::Debug,
{
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt::Debug::fmt(&self.0, fmt)
    }
}

impl<S> MidHandshakeTlsStream<S> {
    pub fn get_ref(&self) -> &S {
        self.0.get_ref()
    }

    pub fn get_mut(&mut self) -> &mut S {
        self.0.get_mut()
    }
}

impl<S> MidHandshakeTlsStream<S>
where
    S: io::Read + io::Write,
{
    pub fn handshake(self) -> Result<TlsStream<S>, HandshakeError<S>> {
        match self.0.handshake() {
            Ok(s) => Ok(TlsStream(s)),
            Err(e) => Err(e.into()),
        }
    }
}

pub enum HandshakeError<S> {
    Failure(Error),
    WouldBlock(MidHandshakeTlsStream<S>),
}

impl<S> From<ssl::HandshakeError<S>> for HandshakeError<S> {
    fn from(e: ssl::HandshakeError<S>) -> HandshakeError<S> {
        match e {
            ssl::HandshakeError::SetupFailure(e) => HandshakeError::Failure(e.into()),
            ssl::HandshakeError::Failure(e) => {
                let v = e.ssl().verify_result();
                HandshakeError::Failure(Error::Ssl(e.into_error(), v))
            }
            ssl::HandshakeError::WouldBlock(s) => {
                HandshakeError::WouldBlock(MidHandshakeTlsStream(s))
            }
        }
    }
}

impl<S> From<ErrorStack> for HandshakeError<S> {
    fn from(e: ErrorStack) -> HandshakeError<S> {
        HandshakeError::Failure(e.into())
    }
}

#[derive(Clone)]
pub struct TlsConnector {
    connector: SslConnector,
    use_sni: bool,
    accept_invalid_hostnames: bool,
    accept_invalid_certs: bool,
}

impl TlsConnector {
    pub fn new(builder: &TlsConnectorBuilder) -> Result<TlsConnector, Error> {
        init_trust();

        let mut connector = SslConnector::builder(SslMethod::tls())?;
        if let Some(ref identity) = builder.identity {
            connector.set_certificate(&identity.0.cert)?;
            connector.set_private_key(&identity.0.pkey)?;
            for cert in identity.0.chain.iter().rev() {
                connector.add_extra_chain_cert(cert.to_owned())?;
            }
        }
        supported_protocols(builder.min_protocol, builder.max_protocol, &mut connector)?;

        for cert in &builder.root_certificates {
            if let Err(err) = connector.cert_store_mut().add_cert((cert.0).0.clone()) {
                debug!("add_cert error: {:?}", err);
            }
        }

        #[cfg(target_os = "android")]
        load_android_root_certs(&mut connector)?;

        Ok(TlsConnector {
            connector: connector.build(),
            use_sni: builder.use_sni,
            accept_invalid_hostnames: builder.accept_invalid_hostnames,
            accept_invalid_certs: builder.accept_invalid_certs,
        })
    }

    pub fn connect<S>(&self, domain: &str, stream: S) -> Result<TlsStream<S>, HandshakeError<S>>
    where
        S: io::Read + io::Write,
    {
        let mut ssl = self
            .connector
            .configure()?
            .use_server_name_indication(self.use_sni)
            .verify_hostname(!self.accept_invalid_hostnames);
        if self.accept_invalid_certs {
            ssl.set_verify(SslVerifyMode::NONE);
        }

        let s = ssl.connect(domain, stream)?;
        Ok(TlsStream(s))
    }
}

#[derive(Clone)]
pub struct TlsAcceptor(SslAcceptor);

impl TlsAcceptor {
    pub fn new(builder: &TlsAcceptorBuilder) -> Result<TlsAcceptor, Error> {
        let mut acceptor = SslAcceptor::mozilla_intermediate(SslMethod::tls())?;
        acceptor.set_private_key(&builder.identity.0.pkey)?;
        acceptor.set_certificate(&builder.identity.0.cert)?;
        for cert in builder.identity.0.chain.iter().rev() {
            acceptor.add_extra_chain_cert(cert.to_owned())?;
        }
        supported_protocols(builder.min_protocol, builder.max_protocol, &mut acceptor)?;

        Ok(TlsAcceptor(acceptor.build()))
    }

    pub fn accept<S>(&self, stream: S) -> Result<TlsStream<S>, HandshakeError<S>>
    where
        S: io::Read + io::Write,
    {
        let s = self.0.accept(stream)?;
        Ok(TlsStream(s))
    }
}

pub struct TlsStream<S>(ssl::SslStream<S>);

impl<S: fmt::Debug> fmt::Debug for TlsStream<S> {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt::Debug::fmt(&self.0, fmt)
    }
}

impl<S> TlsStream<S> {
    pub fn get_ref(&self) -> &S {
        self.0.get_ref()
    }

    pub fn get_mut(&mut self) -> &mut S {
        self.0.get_mut()
    }
}

impl<S: io::Read + io::Write> TlsStream<S> {
    pub fn buffered_read_size(&self) -> Result<usize, Error> {
        Ok(self.0.ssl().pending())
    }

    pub fn peer_certificate(&self) -> Result<Option<Certificate>, Error> {
        Ok(self.0.ssl().peer_certificate().map(Certificate))
    }

    pub fn tls_server_end_point(&self) -> Result<Option<Vec<u8>>, Error> {
        let cert = if self.0.ssl().is_server() {
            self.0.ssl().certificate().map(|x| x.to_owned())
        } else {
            self.0.ssl().peer_certificate()
        };

        let cert = match cert {
            Some(cert) => cert,
            None => return Ok(None),
        };

        let algo_nid = cert.signature_algorithm().object().nid();
        let signature_algorithms = match algo_nid.signature_algorithms() {
            Some(algs) => algs,
            None => return Ok(None),
        };

        let md = match signature_algorithms.digest {
            Nid::MD5 | Nid::SHA1 => MessageDigest::sha256(),
            nid => match MessageDigest::from_nid(nid) {
                Some(md) => md,
                None => return Ok(None),
            },
        };

        let digest = cert.digest(md)?;

        Ok(Some(digest.to_vec()))
    }

    pub fn shutdown(&mut self) -> io::Result<()> {
        match self.0.shutdown() {
            Ok(_) => Ok(()),
            Err(ref e) if e.code() == ssl::ErrorCode::ZERO_RETURN => Ok(()),
            Err(e) => Err(e
                .into_io_error()
                .unwrap_or_else(|e| io::Error::new(io::ErrorKind::Other, e))),
        }
    }
}

impl<S: io::Read + io::Write> io::Read for TlsStream<S> {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        self.0.read(buf)
    }
}

impl<S: io::Read + io::Write> io::Write for TlsStream<S> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.0.write(buf)
    }

    fn flush(&mut self) -> io::Result<()> {
        self.0.flush()
    }
}
<commit_msg>Remove unused import<commit_after>extern crate openssl;
extern crate openssl_probe;

use self::openssl::error::ErrorStack;
use self::openssl::hash::MessageDigest;
use self::openssl::nid::Nid;
use self::openssl::pkcs12::Pkcs12;
use self::openssl::pkey::PKey;
use self::openssl::ssl::{
    self, MidHandshakeSslStream, SslAcceptor, SslConnector, SslContextBuilder, SslMethod,
    SslVerifyMode,
};
use self::openssl::x509::{X509, X509VerifyResult};
use std::error;
use std::fmt;
use std::io;
use std::sync::{Once, ONCE_INIT};

use {Protocol, TlsAcceptorBuilder, TlsConnectorBuilder};
use self::openssl::pkey::Private;

#[cfg(have_min_max_version)]
fn supported_protocols(
    min: Option<Protocol>,
    max: Option<Protocol>,
    ctx: &mut SslContextBuilder,
) -> Result<(), ErrorStack> {
    use self::openssl::ssl::SslVersion;

    fn cvt(p: Protocol) -> SslVersion {
        match p {
            Protocol::Sslv3 => SslVersion::SSL3,
            Protocol::Tlsv10 => SslVersion::TLS1,
            Protocol::Tlsv11 => SslVersion::TLS1_1,
            Protocol::Tlsv12 => SslVersion::TLS1_2,
            Protocol::__NonExhaustive => unreachable!(),
        }
    }

    ctx.set_min_proto_version(min.map(cvt))?;
    ctx.set_max_proto_version(max.map(cvt))?;

    Ok(())
}

#[cfg(not(have_min_max_version))]
fn supported_protocols(
    min: Option<Protocol>,
    max: Option<Protocol>,
    ctx: &mut SslContextBuilder,
) -> Result<(), ErrorStack> {
    use self::openssl::ssl::SslOptions;

    let no_ssl_mask = SslOptions::NO_SSLV2
        | SslOptions::NO_SSLV3
        | SslOptions::NO_TLSV1
        | SslOptions::NO_TLSV1_1
        | SslOptions::NO_TLSV1_2;

    ctx.clear_options(no_ssl_mask);
    let mut options = SslOptions::empty();
    options |= match min {
        None => SslOptions::empty(),
        Some(Protocol::Sslv3) => SslOptions::NO_SSLV2,
        Some(Protocol::Tlsv10) => SslOptions::NO_SSLV2 | SslOptions::NO_SSLV3,
        Some(Protocol::Tlsv11) => {
            SslOptions::NO_SSLV2 | SslOptions::NO_SSLV3 | SslOptions::NO_TLSV1
        }
        Some(Protocol::Tlsv12) => {
            SslOptions::NO_SSLV2
                | SslOptions::NO_SSLV3
                | SslOptions::NO_TLSV1
                | SslOptions::NO_TLSV1_1
        }
        Some(Protocol::__NonExhaustive) => unreachable!(),
    };
    options |= match max {
        None | Some(Protocol::Tlsv12) => SslOptions::empty(),
        Some(Protocol::Tlsv11) => SslOptions::NO_TLSV1_2,
        Some(Protocol::Tlsv10) => SslOptions::NO_TLSV1_1 | SslOptions::NO_TLSV1_2,
        Some(Protocol::Sslv3) => {
            SslOptions::NO_TLSV1 | SslOptions::NO_TLSV1_1 | SslOptions::NO_TLSV1_2
        }
        Some(Protocol::__NonExhaustive) => unreachable!(),
    };

    ctx.set_options(options);

    Ok(())
}

fn init_trust() {
    static ONCE: Once = ONCE_INIT;
    ONCE.call_once(|| openssl_probe::init_ssl_cert_env_vars());
}

#[cfg(target_os = "android")]
fn load_android_root_certs(connector: &mut SslContextBuilder) -> Result<(), Error> {
    use std::fs;

    if let Ok(dir) = fs::read_dir("/system/etc/security/cacerts") {
        let certs = dir
            .filter_map(|r| r.ok())
            .filter_map(|e| fs::read(e.path()).ok())
            .filter_map(|b| X509::from_pem(&b).ok());
        for cert in certs {
            if let Err(err) = connector.cert_store_mut().add_cert(cert) {
                debug!("load_android_root_certs error: {:?}", err);
            }
        }
    }

    Ok(())
}

#[derive(Debug)]
pub enum Error {
    Normal(ErrorStack),
    Ssl(ssl::Error, X509VerifyResult),
}

impl error::Error for Error {
    fn description(&self) -> &str {
        match *self {
            Error::Normal(ref e) => error::Error::description(e),
            Error::Ssl(ref e, _) => error::Error::description(e),
        }
    }

    fn cause(&self) -> Option<&error::Error> {
        match *self {
            Error::Normal(ref e) => error::Error::cause(e),
            Error::Ssl(ref e, _) => error::Error::cause(e),
        }
    }
}

impl fmt::Display for Error {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            Error::Normal(ref e) => fmt::Display::fmt(e, fmt),
            Error::Ssl(ref e, X509VerifyResult::OK) => fmt::Display::fmt(e, fmt),
            Error::Ssl(ref e, v) => write!(fmt, "{} ({})", e, v),
        }
    }
}

impl From<ErrorStack> for Error {
    fn from(err: ErrorStack) -> Error {
        Error::Normal(err)
    }
}

#[derive(Clone)]
pub struct Identity {
    pkey: PKey<Private>,
    cert: X509,
    chain: Vec<X509>,
}

impl Identity {
    pub fn from_pkcs12(buf: &[u8], pass: &str) -> Result<Identity, Error> {
        let pkcs12 = Pkcs12::from_der(buf)?;
        let parsed = pkcs12.parse(pass)?;
        Ok(Identity {
            pkey: parsed.pkey,
            cert: parsed.cert,
            chain: parsed.chain.into_iter().flat_map(|x| x).collect(),
        })
    }
}

#[derive(Clone)]
pub struct Certificate(X509);

impl Certificate {
    pub fn from_der(buf: &[u8]) -> Result<Certificate, Error> {
        let cert = X509::from_der(buf)?;
        Ok(Certificate(cert))
    }

    pub fn from_pem(buf: &[u8]) -> Result<Certificate, Error> {
        let cert = X509::from_pem(buf)?;
        Ok(Certificate(cert))
    }

    pub fn to_der(&self) -> Result<Vec<u8>, Error> {
        let der = self.0.to_der()?;
        Ok(der)
    }
}

pub struct MidHandshakeTlsStream<S>(MidHandshakeSslStream<S>);

impl<S> fmt::Debug for MidHandshakeTlsStream<S>
where
    S: fmt::Debug,
{
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt::Debug::fmt(&self.0, fmt)
    }
}

impl<S> MidHandshakeTlsStream<S> {
    pub fn get_ref(&self) -> &S {
        self.0.get_ref()
    }

    pub fn get_mut(&mut self) -> &mut S {
        self.0.get_mut()
    }
}

impl<S> MidHandshakeTlsStream<S>
where
    S: io::Read + io::Write,
{
    pub fn handshake(self) -> Result<TlsStream<S>, HandshakeError<S>> {
        match self.0.handshake() {
            Ok(s) => Ok(TlsStream(s)),
            Err(e) => Err(e.into()),
        }
    }
}

pub enum HandshakeError<S> {
    Failure(Error),
    WouldBlock(MidHandshakeTlsStream<S>),
}

impl<S> From<ssl::HandshakeError<S>> for HandshakeError<S> {
    fn from(e: ssl::HandshakeError<S>) -> HandshakeError<S> {
        match e {
            ssl::HandshakeError::SetupFailure(e) => HandshakeError::Failure(e.into()),
            ssl::HandshakeError::Failure(e) => {
                let v = e.ssl().verify_result();
                HandshakeError::Failure(Error::Ssl(e.into_error(), v))
            }
            ssl::HandshakeError::WouldBlock(s) => {
                HandshakeError::WouldBlock(MidHandshakeTlsStream(s))
            }
        }
    }
}

impl<S> From<ErrorStack> for HandshakeError<S> {
    fn from(e: ErrorStack) -> HandshakeError<S> {
        HandshakeError::Failure(e.into())
    }
}

#[derive(Clone)]
pub struct TlsConnector {
    connector: SslConnector,
    use_sni: bool,
    accept_invalid_hostnames: bool,
    accept_invalid_certs: bool,
}

impl TlsConnector {
    pub fn new(builder: &TlsConnectorBuilder) -> Result<TlsConnector, Error> {
        init_trust();

        let mut connector = SslConnector::builder(SslMethod::tls())?;
        if let Some(ref identity) = builder.identity {
            connector.set_certificate(&identity.0.cert)?;
            connector.set_private_key(&identity.0.pkey)?;
            for cert in identity.0.chain.iter().rev() {
                connector.add_extra_chain_cert(cert.to_owned())?;
            }
        }
        supported_protocols(builder.min_protocol, builder.max_protocol, &mut connector)?;

        for cert in &builder.root_certificates {
            if let Err(err) = connector.cert_store_mut().add_cert((cert.0).0.clone()) {
                debug!("add_cert error: {:?}", err);
            }
        }

        #[cfg(target_os = "android")]
        load_android_root_certs(&mut connector)?;

        Ok(TlsConnector {
            connector: connector.build(),
            use_sni: builder.use_sni,
            accept_invalid_hostnames: builder.accept_invalid_hostnames,
            accept_invalid_certs: builder.accept_invalid_certs,
        })
    }

    pub fn connect<S>(&self, domain: &str, stream: S) -> Result<TlsStream<S>, HandshakeError<S>>
    where
        S: io::Read + io::Write,
    {
        let mut ssl = self
            .connector
            .configure()?
            .use_server_name_indication(self.use_sni)
            .verify_hostname(!self.accept_invalid_hostnames);
        if self.accept_invalid_certs {
            ssl.set_verify(SslVerifyMode::NONE);
        }

        let s = ssl.connect(domain, stream)?;
        Ok(TlsStream(s))
    }
}

#[derive(Clone)]
pub struct TlsAcceptor(SslAcceptor);

impl TlsAcceptor {
    pub fn new(builder: &TlsAcceptorBuilder) -> Result<TlsAcceptor, Error> {
        let mut acceptor = SslAcceptor::mozilla_intermediate(SslMethod::tls())?;
        acceptor.set_private_key(&builder.identity.0.pkey)?;
        acceptor.set_certificate(&builder.identity.0.cert)?;
        for cert in builder.identity.0.chain.iter().rev() {
            acceptor.add_extra_chain_cert(cert.to_owned())?;
        }
        supported_protocols(builder.min_protocol, builder.max_protocol, &mut acceptor)?;

        Ok(TlsAcceptor(acceptor.build()))
    }

    pub fn accept<S>(&self, stream: S) -> Result<TlsStream<S>, HandshakeError<S>>
    where
        S: io::Read + io::Write,
    {
        let s = self.0.accept(stream)?;
        Ok(TlsStream(s))
    }
}

pub struct TlsStream<S>(ssl::SslStream<S>);

impl<S: fmt::Debug> fmt::Debug for TlsStream<S> {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt::Debug::fmt(&self.0, fmt)
    }
}

impl<S> TlsStream<S> {
    pub fn get_ref(&self) -> &S {
        self.0.get_ref()
    }

    pub fn get_mut(&mut self) -> &mut S {
        self.0.get_mut()
    }
}

impl<S: io::Read + io::Write> TlsStream<S> {
    pub fn buffered_read_size(&self) -> Result<usize, Error> {
        Ok(self.0.ssl().pending())
    }

    pub fn peer_certificate(&self) -> Result<Option<Certificate>, Error> {
        Ok(self.0.ssl().peer_certificate().map(Certificate))
    }

    pub fn tls_server_end_point(&self) -> Result<Option<Vec<u8>>, Error> {
        let cert = if self.0.ssl().is_server() {
            self.0.ssl().certificate().map(|x| x.to_owned())
        } else {
            self.0.ssl().peer_certificate()
        };

        let cert = match cert {
            Some(cert) => cert,
            None => return Ok(None),
        };

        let algo_nid = cert.signature_algorithm().object().nid();
        let signature_algorithms = match algo_nid.signature_algorithms() {
            Some(algs) => algs,
            None => return Ok(None),
        };

        let md = match signature_algorithms.digest {
            Nid::MD5 | Nid::SHA1 => MessageDigest::sha256(),
            nid => match MessageDigest::from_nid(nid) {
                Some(md) => md,
                None => return Ok(None),
            },
        };

        let digest = cert.digest(md)?;

        Ok(Some(digest.to_vec()))
    }

    pub fn shutdown(&mut self) -> io::Result<()> {
        match self.0.shutdown() {
            Ok(_) => Ok(()),
            Err(ref e) if e.code() == ssl::ErrorCode::ZERO_RETURN => Ok(()),
            Err(e) => Err(e
                .into_io_error()
                .unwrap_or_else(|e| io::Error::new(io::ErrorKind::Other, e))),
        }
    }
}

impl<S: io::Read + io::Write> io::Read for TlsStream<S> {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        self.0.read(buf)
    }
}

impl<S: io::Read + io::Write> io::Write for TlsStream<S> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.0.write(buf)
    }

    fn flush(&mut self) -> io::Result<()> {
        self.0.flush()
    }
}
<|endoftext|>