<commit_before>use lookup;
use txt_records::{Spf, Word, Mechanism, Modifier};
use statistics::Statistics;

use ansi_term::Colour;
use chrono::{Local, Duration};
use chrono_humanize::HumanTime;
use error_chain::ChainedError;
use itertools::Itertools;
use std::cmp::Ordering;
use std::io::{self, Write};
use tabwriter::TabWriter;
use trust_dns::rr::{RData, Record};

pub trait OutputModule {
    fn output(&self, w: &mut Write) -> Result<()>;
    fn is_human(&self) -> bool;
}

pub struct DetailsOutput<'a> {
    responses: &'a [lookup::Result<lookup::Response>],
    human: bool,
}

impl<'a> DetailsOutput<'a> {
    pub fn new(responses: &'a [lookup::Result<lookup::Response>], human: bool) -> Self {
        DetailsOutput { responses, human }
    }
}

impl<'a> OutputModule for DetailsOutput<'a> {
    fn output(&self, w: &mut Write) -> Result<()> {
        for response in self.responses {
            match *response {
                Ok(ref r) => write_response(w, r, self.human).chain_err(|| ErrorKind::OutputError)?,
                Err(ref e) => print_error(w, e)?,
            }
        }

        Ok(())
    }

    fn is_human(&self) -> bool {
        self.human
    }
}

pub struct SummaryOutput<'a> {
    statistics: Statistics<'a>,
    human: bool,
}

impl<'a> SummaryOutput<'a> {
    pub fn new(responses: &'a [lookup::Result<lookup::Response>], human: bool) -> Self {
        let statistics = Statistics::from(responses);
        SummaryOutput { statistics, human }
    }
}

impl<'a> OutputModule for SummaryOutput<'a> {
    fn output(&self, w: &mut Write) -> Result<()> {
        writeln!(w, "Received {} (min {}, max {} records) answers from {} servers",
                 self.statistics.num_of_ok_samples,
                 self.statistics.min_num_of_records,
                 self.statistics.max_num_of_records,
                 self.statistics.num_of_samples,
        ).chain_err(|| ErrorKind::OutputError)?;
        let records: Vec<_> = self.statistics
            .record_counts
            .values()
            // TODO: Why do I need to specify a closure and not just a function?
            .sorted_by(|a, b| compare_records(a.0, b.0))
            .iter()
            .map(|rr| {
                let record = rr.0;
                let count = rr.1;
                format!("* {} ({})", fmt_record(record, self.human), count)
            })
            .collect();

        let mut tw = TabWriter::new(vec![]).padding(1);
        write!(&mut tw, "{}", records.join("\n")).chain_err(|| ErrorKind::OutputError)?;
        let out_str = String::from_utf8(tw.into_inner().unwrap()).unwrap();

        writeln!(w, "{}", out_str).chain_err(|| ErrorKind::OutputError)?;

        Ok(())
    }

    fn is_human(&self) -> bool {
        self.human
    }
}

fn write_response(f: &mut Write, r: &lookup::Response, human: bool) -> io::Result<()> {
    if r.answers.is_empty() {
        return writeln!(f, "DNS server {} has no records.", r.server);
    }
    let _ = write!(f, "DNS server {} responded with\n", r.server);
    let answers: Vec<String> = r.answers
        .iter()
        .sorted_by(|a, b| compare_records(a, b))
        .iter()
        .map(|answer|
            format!("* {} [expires {}]",
                    fmt_record(answer, human),
                    if human {
                        humanize_ttl(answer.ttl() as i64)
                    } else {
                        format!("in {} sec", answer.ttl())
                    }
            )
        )
        .collect();

    let mut tw = TabWriter::new(vec![]).padding(1);
    let _ = write!(&mut tw, "{}", answers.join("\n"));
    let out_str = String::from_utf8(tw.into_inner().unwrap()).unwrap();

    writeln!(f, "{}", out_str)
}

fn compare_records(a: &Record, b: &Record) -> Ordering {
    let a = record_type_to_ordinal(a);
    let b = record_type_to_ordinal(b);

    a.cmp(&b)
}

fn record_type_to_ordinal(r: &Record) -> u16 {
    match *r.rdata() {
        RData::SOA(_) => 1000,
        RData::NS(_) => 2000,
        RData::MX(ref mx) => 3000 + mx.preference(),
        RData::TXT(_) => 4000,
        RData::CNAME(_) => 5000,
        RData::A(_) => 6000,
        RData::AAAA(_) => 7000,
        RData::PTR(_) => 8000,
        _ => ::std::u16::MAX,
    }
}

fn humanize_ttl(ttl: i64) -> String {
    let dt = Local::now() + Duration::seconds(ttl);
    let ht = HumanTime::from(dt);

    format!("{}", ht)
}

fn fmt_record(r: &Record, human: bool) -> String {
    match *r.rdata() {
        RData::A(ip) => format!("IPv4:\t{}", ip),
        RData::AAAA(ip) => format!("IPv6:\t{}", ip),
        RData::CNAME(ref name) => {
            format!("CNAME:\t{}", Colour::Blue.paint(format!("{}", name)))
        }
        RData::MX(ref mx) => {
            format!(
                "MX:\t{} with preference {}",
                Colour::Yellow.paint(format!("{}", mx.exchange())),
                Colour::Yellow.paint(format!("{}", mx.preference()))
            )
        }
        RData::NS(ref name) => format!("NS:\t{}", Colour::Cyan.paint(format!("{}", name))),
        RData::SOA(ref soa) => {
            // TODO: time interval is in sec, use relative time
            format!(
                "SOA:\torigin NS {}, responsible party {}, serial {}, refresh {} sec, retry {} sec, expire {} sec, min {} sec",
                Colour::Green.paint(format!("{}", soa.mname())),
                Colour::Green.paint(format!("{}", soa.rname())),
                Colour::Green.paint(format!("{}", soa.serial())),
                Colour::Green.paint(format!("{}", soa.refresh())),
                Colour::Green.paint(format!("{}", soa.retry())),
                Colour::Green.paint(format!("{}", soa.expire())),
                Colour::Green.paint(format!("{}", soa.minimum()))
            )
        }
        RData::TXT(ref txt) => {
            format!("TXT:\t{}", Colour::Purple.paint(
                if human {
                    fmt_txt(txt.txt_data())
                } else {
                    txt.txt_data().join(" ")
                }
            ))
        }
        RData::PTR(ref ptr) => format!("PTR:\t{}", ptr.to_string()),
        ref x => {
            format!(
                " * unclassified answer: {}",
                Colour::Red.paint(format!("{:?}", x))
            )
        }
    }
}

fn fmt_txt(txts: &[String]) -> String {
    let fmts: Vec<_> = txts
        .iter()
        .map(|txt| {
            if let Ok(spf) = Spf::from_str(txt) {
                fmt_txt_spf(&spf)
            } else {
                format!("{}", txt)
            }
        })
        .collect();
    fmts.iter().join("\t* ")
}

fn fmt_txt_spf(spf: &Spf) -> String {
    let words: Vec<_> = spf.words
        .iter()
        .map(|w| {
            match w {
                &Word::Word(ref q, Mechanism::All) => format!("{:?} for all", q),
                &Word::Word(ref q, Mechanism::A) => format!("{:?} for A/AAAA record", q),
                &Word::Word(ref q, Mechanism::IPv4(range)) => format!("{:?} for IPv4 range {}", q, range),
                &Word::Word(ref q, Mechanism::IPv6(range)) => format!("{:?} for IPv6 range {}", q, range),
                &Word::Word(ref q, Mechanism::MX) => format!("{:?} for mail exchanges", q),
                &Word::Word(ref q, Mechanism::PTR) => format!("{:?} for reverse mapping", q),
                &Word::Word(ref q, Mechanism::Exists(domain)) => format!("{:?} for A/AAAA record according to {}", q, domain),
                &Word::Word(ref q, Mechanism::Include(domain)) => format!("{:?} for include from {}", q, domain),
                &Word::Modifier(Modifier::Redirect(query)) => format!("redirect to query {}", query),
                &Word::Modifier(Modifier::Exp(explanation)) => format!("explanation according to {}", explanation),
            }
        })
        .collect();
    format!("SPF version: {}\n\t* {}", spf.version, words.join("\n\t* "))
}

pub fn print_error<T: ChainedError>(w: &mut Write, err: &T) -> Result<()> {
    write!(w, "{} ", err).chain_err(|| ErrorKind::OutputError)?;
    for e in err.iter().skip(1) {
        write!(w, "because {}", e).chain_err(|| ErrorKind::OutputError)?;
    }
    writeln!(w).chain_err(|| ErrorKind::OutputError)?;

    // The backtrace is only available if run with `RUST_BACKTRACE=1`.
    if let Some(backtrace) = err.backtrace() {
        writeln!(w, "backtrace: {:?}", backtrace).chain_err(|| ErrorKind::OutputError)?;
    }

    Ok(())
}

error_chain! {
    errors {
        OutputError {
            description("Failed to write output")
            display("Failed to write output")
        }
    }
}
<commit_msg>Distinguish between ip and ip range for spf output<commit_after>use lookup;
use txt_records::{Spf, Word, Mechanism, Modifier};
use statistics::Statistics;

use ansi_term::Colour;
use chrono::{Local, Duration};
use chrono_humanize::HumanTime;
use error_chain::ChainedError;
use itertools::Itertools;
use std::cmp::Ordering;
use std::io::{self, Write};
use tabwriter::TabWriter;
use trust_dns::rr::{RData, Record};

pub trait OutputModule {
    fn output(&self, w: &mut Write) -> Result<()>;
    fn is_human(&self) -> bool;
}

pub struct DetailsOutput<'a> {
    responses: &'a [lookup::Result<lookup::Response>],
    human: bool,
}

impl<'a> DetailsOutput<'a> {
    pub fn new(responses: &'a [lookup::Result<lookup::Response>], human: bool) -> Self {
        DetailsOutput { responses, human }
    }
}

impl<'a> OutputModule for DetailsOutput<'a> {
    fn output(&self, w: &mut Write) -> Result<()> {
        for response in self.responses {
            match *response {
                Ok(ref r) => write_response(w, r, self.human).chain_err(|| ErrorKind::OutputError)?,
                Err(ref e) => print_error(w, e)?,
            }
        }

        Ok(())
    }

    fn is_human(&self) -> bool {
        self.human
    }
}

pub struct SummaryOutput<'a> {
    statistics: Statistics<'a>,
    human: bool,
}

impl<'a> SummaryOutput<'a> {
    pub fn new(responses: &'a [lookup::Result<lookup::Response>], human: bool) -> Self {
        let statistics = Statistics::from(responses);
        SummaryOutput { statistics, human }
    }
}

impl<'a> OutputModule for SummaryOutput<'a> {
    fn output(&self, w: &mut Write) -> Result<()> {
        writeln!(w, "Received {} (min {}, max {} records) answers from {} servers",
                 self.statistics.num_of_ok_samples,
                 self.statistics.min_num_of_records,
                 self.statistics.max_num_of_records,
                 self.statistics.num_of_samples,
        ).chain_err(|| ErrorKind::OutputError)?;
        let records: Vec<_> = self.statistics
            .record_counts
            .values()
            // TODO: Why do I need to specify a closure and not just a function?
            .sorted_by(|a, b| compare_records(a.0, b.0))
            .iter()
            .map(|rr| {
                let record = rr.0;
                let count = rr.1;
                format!("* {} ({})", fmt_record(record, self.human), count)
            })
            .collect();

        let mut tw = TabWriter::new(vec![]).padding(1);
        write!(&mut tw, "{}", records.join("\n")).chain_err(|| ErrorKind::OutputError)?;
        let out_str = String::from_utf8(tw.into_inner().unwrap()).unwrap();

        writeln!(w, "{}", out_str).chain_err(|| ErrorKind::OutputError)?;

        Ok(())
    }

    fn is_human(&self) -> bool {
        self.human
    }
}

fn write_response(f: &mut Write, r: &lookup::Response, human: bool) -> io::Result<()> {
    if r.answers.is_empty() {
        return writeln!(f, "DNS server {} has no records.", r.server);
    }
    let _ = write!(f, "DNS server {} responded with\n", r.server);
    let answers: Vec<String> = r.answers
        .iter()
        .sorted_by(|a, b| compare_records(a, b))
        .iter()
        .map(|answer|
            format!("* {} [expires {}]",
                    fmt_record(answer, human),
                    if human {
                        humanize_ttl(answer.ttl() as i64)
                    } else {
                        format!("in {} sec", answer.ttl())
                    }
            )
        )
        .collect();

    let mut tw = TabWriter::new(vec![]).padding(1);
    let _ = write!(&mut tw, "{}", answers.join("\n"));
    let out_str = String::from_utf8(tw.into_inner().unwrap()).unwrap();

    writeln!(f, "{}", out_str)
}

fn compare_records(a: &Record, b: &Record) -> Ordering {
    let a = record_type_to_ordinal(a);
    let b = record_type_to_ordinal(b);

    a.cmp(&b)
}

fn record_type_to_ordinal(r: &Record) -> u16 {
    match *r.rdata() {
        RData::SOA(_) => 1000,
        RData::NS(_) => 2000,
        RData::MX(ref mx) => 3000 + mx.preference(),
        RData::TXT(_) => 4000,
        RData::CNAME(_) => 5000,
        RData::A(_) => 6000,
        RData::AAAA(_) => 7000,
        RData::PTR(_) => 8000,
        _ => ::std::u16::MAX,
    }
}

fn humanize_ttl(ttl: i64) -> String {
    let dt = Local::now() + Duration::seconds(ttl);
    let ht = HumanTime::from(dt);

    format!("{}", ht)
}

fn fmt_record(r: &Record, human: bool) -> String {
    match *r.rdata() {
        RData::A(ip) => format!("IPv4:\t{}", ip),
        RData::AAAA(ip) => format!("IPv6:\t{}", ip),
        RData::CNAME(ref name) => {
            format!("CNAME:\t{}", Colour::Blue.paint(format!("{}", name)))
        }
        RData::MX(ref mx) => {
            format!(
                "MX:\t{} with preference {}",
                Colour::Yellow.paint(format!("{}", mx.exchange())),
                Colour::Yellow.paint(format!("{}", mx.preference()))
            )
        }
        RData::NS(ref name) => format!("NS:\t{}", Colour::Cyan.paint(format!("{}", name))),
        RData::SOA(ref soa) => {
            // TODO: time interval is in sec, use relative time
            format!(
                "SOA:\torigin NS {}, responsible party {}, serial {}, refresh {} sec, retry {} sec, expire {} sec, min {} sec",
                Colour::Green.paint(format!("{}", soa.mname())),
                Colour::Green.paint(format!("{}", soa.rname())),
                Colour::Green.paint(format!("{}", soa.serial())),
                Colour::Green.paint(format!("{}", soa.refresh())),
                Colour::Green.paint(format!("{}", soa.retry())),
                Colour::Green.paint(format!("{}", soa.expire())),
                Colour::Green.paint(format!("{}", soa.minimum()))
            )
        }
        RData::TXT(ref txt) => {
            format!("TXT:\t{}", Colour::Purple.paint(
                if human {
                    fmt_txt(txt.txt_data())
                } else {
                    txt.txt_data().join(" ")
                }
            ))
        }
        RData::PTR(ref ptr) => format!("PTR:\t{}", ptr.to_string()),
        ref x => {
            format!(
                " * unclassified answer: {}",
                Colour::Red.paint(format!("{:?}", x))
            )
        }
    }
}

fn fmt_txt(txts: &[String]) -> String {
    let fmts: Vec<_> = txts
        .iter()
        .map(|txt| {
            if let Ok(spf) = Spf::from_str(txt) {
                fmt_txt_spf(&spf)
            } else {
                format!("{}", txt)
            }
        })
        .collect();
    fmts.iter().join("\t* ")
}

fn fmt_txt_spf(spf: &Spf) -> String {
    let words: Vec<_> = spf.words
        .iter()
        .map(|w| {
            match w {
                &Word::Word(ref q, Mechanism::All) => format!("{:?} for all", q),
                &Word::Word(ref q, Mechanism::A) => format!("{:?} for A/AAAA record", q),
                &Word::Word(ref q, Mechanism::IPv4(range)) if range.contains("/") => format!("{:?} for IPv4 range {}", q, range),
                &Word::Word(ref q, Mechanism::IPv4(range)) => format!("{:?} for IPv4 {}", q, range),
                &Word::Word(ref q, Mechanism::IPv6(range)) if range.contains("/") => format!("{:?} for IPv6 range {}", q, range),
                &Word::Word(ref q, Mechanism::IPv6(range)) => format!("{:?} for IPv6 {}", q, range),
                &Word::Word(ref q, Mechanism::MX) => format!("{:?} for mail exchanges", q),
                &Word::Word(ref q, Mechanism::PTR) => format!("{:?} for reverse mapping", q),
                &Word::Word(ref q, Mechanism::Exists(domain)) => format!("{:?} for A/AAAA record according to {}", q, domain),
                &Word::Word(ref q, Mechanism::Include(domain)) => format!("{:?} for include from {}", q, domain),
                &Word::Modifier(Modifier::Redirect(query)) => format!("redirect to query {}", query),
                &Word::Modifier(Modifier::Exp(explanation)) => format!("explanation according to {}", explanation),
            }
        })
        .collect();
    format!("SPF version: {}\n\t* {}", spf.version, words.join("\n\t* "))
}

pub fn print_error<T: ChainedError>(w: &mut Write, err: &T) -> Result<()> {
    write!(w, "{} ", err).chain_err(|| ErrorKind::OutputError)?;
    for e in err.iter().skip(1) {
        write!(w, "because {}", e).chain_err(|| ErrorKind::OutputError)?;
    }
    writeln!(w).chain_err(|| ErrorKind::OutputError)?;

    // The backtrace is only available if run with `RUST_BACKTRACE=1`.
    if let Some(backtrace) = err.backtrace() {
        writeln!(w, "backtrace: {:?}", backtrace).chain_err(|| ErrorKind::OutputError)?;
    }

    Ok(())
}

error_chain! {
    errors {
        OutputError {
            description("Failed to write output")
            display("Failed to write output")
        }
    }
}
<|endoftext|>


<commit_before>//! Decoding and Encoding of PNG Images
//!
//! PNG (Portable Network Graphics) is an image format that supports lossless compression.
//!
//! # Related Links
//! * <http://www.w3.org/TR/PNG/> - The PNG Specification
//!

use std::convert::TryFrom;
use std::io::{self, Read, Write};

use crate::color::{ColorType, ExtendedColorType};
use crate::error::{
    DecodingError, ImageError, ImageResult, LimitError, LimitErrorKind, ParameterError, ParameterErrorKind
};
use crate::image::{ImageDecoder, ImageEncoder, ImageFormat};

/// PNG Reader
///
/// This reader will try to read the png one row at a time,
/// however for interlaced png files this is not possible and
/// these are therefore read at once.
pub struct PNGReader<R: Read> {
    reader: png::Reader<R>,
    buffer: Vec<u8>,
    index: usize,
}

impl<R: Read> PNGReader<R> {
    fn new(mut reader: png::Reader<R>) -> ImageResult<PNGReader<R>> {
        let len = reader.output_buffer_size();
        // Since interlaced images do not come in
        // scanline order it is almost impossible to
        // read them in a streaming fashion, however
        // this shouldn't be a too big of a problem
        // as most interlaced images should fit in memory.
        let buffer = if reader.info().interlaced {
            let mut buffer = vec![0; len];
            reader.next_frame(&mut buffer).map_err(ImageError::from_png)?;
            buffer
        } else {
            Vec::new()
        };

        Ok(PNGReader {
            reader,
            buffer,
            index: 0,
        })
    }
}

impl<R: Read> Read for PNGReader<R> {
    fn read(&mut self, mut buf: &mut [u8]) -> io::Result<usize> {
        // io::Write::write for slice cannot fail
        let readed = buf.write(&self.buffer[self.index..]).unwrap();

        let mut bytes = readed;
        self.index += readed;

        while self.index >= self.buffer.len() {
            match self.reader.next_row()? {
                Some(row) => {
                    // Faster to copy directly to external buffer
                    let readed  = buf.write(row).unwrap();
                    bytes += readed;

                    self.buffer = (&row[readed..]).to_owned();
                    self.index = 0;
                }
                None => return Ok(bytes)
            }
        }

        Ok(bytes)
    }

    fn read_to_end(&mut self, buf: &mut Vec<u8>) -> io::Result<usize> {
        let mut bytes = self.buffer.len();
        buf.extend_from_slice(&self.buffer);
        self.buffer = Vec::new();
        self.index = 0;

        while let Some(row) = self.reader.next_row()? {
            buf.extend_from_slice(row);
            bytes += row.len();
        }

        Ok(bytes)
    }
}

/// PNG decoder
pub struct PngDecoder<R: Read> {
    color_type: ColorType,
    reader: png::Reader<R>,
}

impl<R: Read> PngDecoder<R> {
    /// Creates a new decoder that decodes from the stream ```r```
    pub fn new(r: R) -> ImageResult<PngDecoder<R>> {
        let limits = png::Limits {
            bytes: usize::max_value(),
        };
        let mut decoder = png::Decoder::new_with_limits(r, limits);
        // By default the PNG decoder will scale 16 bpc to 8 bpc, so custom
        // transformations must be set. EXPAND preserves the default behavior
        // expanding bpc < 8 to 8 bpc.
        decoder.set_transformations(png::Transformations::EXPAND);
        let (_, mut reader) = decoder.read_info().map_err(ImageError::from_png)?;
        let (color_type, bits) = reader.output_color_type();
        let color_type = match (color_type, bits) {
            (png::ColorType::Grayscale, png::BitDepth::Eight) => ColorType::L8,
            (png::ColorType::Grayscale, png::BitDepth::Sixteen) => ColorType::L16,
            (png::ColorType::GrayscaleAlpha, png::BitDepth::Eight) => ColorType::La8,
            (png::ColorType::GrayscaleAlpha, png::BitDepth::Sixteen) => ColorType::La16,
            (png::ColorType::RGB, png::BitDepth::Eight) => ColorType::Rgb8,
            (png::ColorType::RGB, png::BitDepth::Sixteen) => ColorType::Rgb16,
            (png::ColorType::RGBA, png::BitDepth::Eight) => ColorType::Rgba8,
            (png::ColorType::RGBA, png::BitDepth::Sixteen) => ColorType::Rgba16,

            (png::ColorType::Grayscale, png::BitDepth::One) =>
                return Err(ImageError::UnsupportedColor(ExtendedColorType::L1)),
            (png::ColorType::GrayscaleAlpha, png::BitDepth::One) =>
                return Err(ImageError::UnsupportedColor(ExtendedColorType::La1)),
            (png::ColorType::RGB, png::BitDepth::One) =>
                return Err(ImageError::UnsupportedColor(ExtendedColorType::Rgb1)),
            (png::ColorType::RGBA, png::BitDepth::One) =>
                return Err(ImageError::UnsupportedColor(ExtendedColorType::Rgba1)),

            (png::ColorType::Grayscale, png::BitDepth::Two) =>
                return Err(ImageError::UnsupportedColor(ExtendedColorType::L2)),
            (png::ColorType::GrayscaleAlpha, png::BitDepth::Two) =>
                return Err(ImageError::UnsupportedColor(ExtendedColorType::La2)),
            (png::ColorType::RGB, png::BitDepth::Two) =>
                return Err(ImageError::UnsupportedColor(ExtendedColorType::Rgb2)),
            (png::ColorType::RGBA, png::BitDepth::Two) =>
                return Err(ImageError::UnsupportedColor(ExtendedColorType::Rgba2)),

            (png::ColorType::Grayscale, png::BitDepth::Four) =>
                return Err(ImageError::UnsupportedColor(ExtendedColorType::L4)),
            (png::ColorType::GrayscaleAlpha, png::BitDepth::Four) =>
                return Err(ImageError::UnsupportedColor(ExtendedColorType::La4)),
            (png::ColorType::RGB, png::BitDepth::Four) =>
                return Err(ImageError::UnsupportedColor(ExtendedColorType::Rgb4)),
            (png::ColorType::RGBA, png::BitDepth::Four) =>
                return Err(ImageError::UnsupportedColor(ExtendedColorType::Rgba4)),

            (png::ColorType::Indexed, bits) =>
                return Err(ImageError::UnsupportedColor(ExtendedColorType::Unknown(bits as u8))),
        };

        Ok(PngDecoder { color_type, reader })
    }
}

impl<'a, R: 'a + Read> ImageDecoder<'a> for PngDecoder<R> {
    type Reader = PNGReader<R>;

    fn dimensions(&self) -> (u32, u32) {
        self.reader.info().size()
    }

    fn color_type(&self) -> ColorType {
        self.color_type
    }

    fn into_reader(self) -> ImageResult<Self::Reader> {
        PNGReader::new(self.reader)
    }

    fn read_image(mut self, buf: &mut [u8]) -> ImageResult<()> {
        use byteorder::{BigEndian, ByteOrder, NativeEndian};

        assert_eq!(u64::try_from(buf.len()), Ok(self.total_bytes()));
        self.reader.next_frame(buf).map_err(ImageError::from_png)?;
        // PNG images are big endian. For 16 bit per channel and larger types,
        // the buffer may need to be reordered to native endianness per the
        // contract of `read_image`.
        // TODO: assumes equal channel bit depth.
        let bpc = self.color_type().bytes_per_pixel() / self.color_type().channel_count();
        match bpc {
            1 => (),  // No reodering necessary for u8
            2 => buf.chunks_mut(2).for_each(|c| {
                let v = BigEndian::read_u16(c);
                NativeEndian::write_u16(c, v)
            }),
            _ => unreachable!(),
        }
        Ok(())
    }

    fn scanline_bytes(&self) -> u64 {
        let width = self.reader.info().width;
        self.reader.output_line_size(width) as u64
    }
}

/// PNG encoder
pub struct PNGEncoder<W: Write> {
    w: W,
}

impl<W: Write> PNGEncoder<W> {
    /// Create a new encoder that writes its output to ```w```
    pub fn new(w: W) -> PNGEncoder<W> {
        PNGEncoder { w }
    }

    /// Encodes the image ```data```
    /// that has dimensions ```width``` and ```height```
    /// and ```ColorType``` ```c```
    pub fn encode(self, data: &[u8], width: u32, height: u32, color: ColorType) -> ImageResult<()> {
        let (ct, bits) = match color {
            ColorType::L8 => (png::ColorType::Grayscale, png::BitDepth::Eight),
            ColorType::L16 => (png::ColorType::Grayscale,png::BitDepth::Sixteen),
            ColorType::La8 => (png::ColorType::GrayscaleAlpha, png::BitDepth::Eight),
            ColorType::La16 => (png::ColorType::GrayscaleAlpha,png::BitDepth::Sixteen),
            ColorType::Rgb8 => (png::ColorType::RGB, png::BitDepth::Eight),
            ColorType::Rgb16 => (png::ColorType::RGB,png::BitDepth::Sixteen),
            ColorType::Rgba8 => (png::ColorType::RGBA, png::BitDepth::Eight),
            ColorType::Rgba16 => (png::ColorType::RGBA,png::BitDepth::Sixteen),
            _ => return Err(ImageError::UnsupportedColor(color.into())),
        };

        let mut encoder = png::Encoder::new(self.w, width, height);
        encoder.set_color(ct);
        encoder.set_depth(bits);
        let mut writer = encoder.write_header().map_err(|e| ImageError::IoError(e.into()))?;
        writer.write_image_data(data).map_err(|e| ImageError::IoError(e.into()))
    }
}

impl<W: Write> ImageEncoder for PNGEncoder<W> {
    fn write_image(
        self,
        buf: &[u8],
        width: u32,
        height: u32,
        color_type: ColorType,
    ) -> ImageResult<()> {
        use byteorder::{BigEndian, ByteOrder, NativeEndian};

        // PNG images are big endian. For 16 bit per channel and larger types,
        // the buffer may need to be reordered to big endian per the
        // contract of `write_image`.
        // TODO: assumes equal channel bit depth.
        let bpc = color_type.bytes_per_pixel() / color_type.channel_count();
        match bpc {
            1 => self.encode(buf, width, height, color_type),  // No reodering necessary for u8
            2 => {
                // Because the buffer is immutable and the PNG encoder does not
                // yet take Write/Read traits, create a temporary buffer for
                // big endian reordering.
                let mut reordered = vec![0; buf.len()];
                buf.chunks(2)
                    .zip(reordered.chunks_mut(2))
                    .for_each(|(b, r)| BigEndian::write_u16(r, NativeEndian::read_u16(b)));
                self.encode(&reordered, width, height, color_type)
            },
            _ => unreachable!(),
        }
    }
}

impl ImageError {
    fn from_png(err: png::DecodingError) -> ImageError {
        use png::DecodingError::*;
        match err {
            IoError(err) => ImageError::IoError(err),
            Format(message) => ImageError::Decoding(DecodingError::with_message(
                ImageFormat::Png.into(),
                message,
            )),
            LimitsExceeded => ImageError::Limits(LimitError::from_kind(
                LimitErrorKind::InsufficientMemory,
            )),
            // Other is used when the buffer to `Reader::next_frame` is too small.
            Other(message) => ImageError::Parameter(ParameterError::from_kind(
                ParameterErrorKind::Generic(message.into_owned())
            )),
            err @ InvalidSignature
            | err @ CrcMismatch { .. }
            | err @ CorruptFlateStream => {
                ImageError::Decoding(DecodingError::new(
                    ImageFormat::Png.into(),
                    err,
                ))
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use crate::image::ImageDecoder;
    use std::io::Read;
    use super::*;

    #[test]
    fn ensure_no_decoder_off_by_one() {
        let dec = PngDecoder::new(std::fs::File::open("tests/images/png/bugfixes/debug_triangle_corners_widescreen.png").unwrap())
            .expect("Unable to read PNG file (does it exist?)");

        assert_eq![(2000, 1000), dec.dimensions()];

        assert_eq![
            ColorType::Rgb8,
            dec.color_type(),
            "Image MUST have the Rgb8 format"
        ];

        let correct_bytes = dec
            .into_reader()
            .expect("Unable to read file")
            .bytes()
            .map(|x| x.expect("Unable to read byte"))
            .collect::<Vec<u8>>();

        assert_eq![6_000_000, correct_bytes.len()];
    }

    #[test]
    fn underlying_error() {
        use std::error::Error;

        let mut not_png = std::fs::read("tests/images/png/bugfixes/debug_triangle_corners_widescreen.png").unwrap();
        not_png[0] = 0;

        let error = PngDecoder::new(&not_png[..]).err().unwrap();
        let _ = error
            .source()
            .unwrap()
            .downcast_ref::<png::DecodingError>()
            .expect("Caused by a png error");
    }
}
<commit_msg>Replace with_message() with bubbling of png crate error in PNG decoder<commit_after>//! Decoding and Encoding of PNG Images
//!
//! PNG (Portable Network Graphics) is an image format that supports lossless compression.
//!
//! # Related Links
//! * <http://www.w3.org/TR/PNG/> - The PNG Specification
//!

use std::convert::TryFrom;
use std::io::{self, Read, Write};

use crate::color::{ColorType, ExtendedColorType};
use crate::error::{
    DecodingError, ImageError, ImageResult, LimitError, LimitErrorKind, ParameterError, ParameterErrorKind
};
use crate::image::{ImageDecoder, ImageEncoder, ImageFormat};

/// PNG Reader
///
/// This reader will try to read the png one row at a time,
/// however for interlaced png files this is not possible and
/// these are therefore read at once.
pub struct PNGReader<R: Read> {
    reader: png::Reader<R>,
    buffer: Vec<u8>,
    index: usize,
}

impl<R: Read> PNGReader<R> {
    fn new(mut reader: png::Reader<R>) -> ImageResult<PNGReader<R>> {
        let len = reader.output_buffer_size();
        // Since interlaced images do not come in
        // scanline order it is almost impossible to
        // read them in a streaming fashion, however
        // this shouldn't be a too big of a problem
        // as most interlaced images should fit in memory.
        let buffer = if reader.info().interlaced {
            let mut buffer = vec![0; len];
            reader.next_frame(&mut buffer).map_err(ImageError::from_png)?;
            buffer
        } else {
            Vec::new()
        };

        Ok(PNGReader {
            reader,
            buffer,
            index: 0,
        })
    }
}

impl<R: Read> Read for PNGReader<R> {
    fn read(&mut self, mut buf: &mut [u8]) -> io::Result<usize> {
        // io::Write::write for slice cannot fail
        let readed = buf.write(&self.buffer[self.index..]).unwrap();

        let mut bytes = readed;
        self.index += readed;

        while self.index >= self.buffer.len() {
            match self.reader.next_row()? {
                Some(row) => {
                    // Faster to copy directly to external buffer
                    let readed  = buf.write(row).unwrap();
                    bytes += readed;

                    self.buffer = (&row[readed..]).to_owned();
                    self.index = 0;
                }
                None => return Ok(bytes)
            }
        }

        Ok(bytes)
    }

    fn read_to_end(&mut self, buf: &mut Vec<u8>) -> io::Result<usize> {
        let mut bytes = self.buffer.len();
        buf.extend_from_slice(&self.buffer);
        self.buffer = Vec::new();
        self.index = 0;

        while let Some(row) = self.reader.next_row()? {
            buf.extend_from_slice(row);
            bytes += row.len();
        }

        Ok(bytes)
    }
}

/// PNG decoder
pub struct PngDecoder<R: Read> {
    color_type: ColorType,
    reader: png::Reader<R>,
}

impl<R: Read> PngDecoder<R> {
    /// Creates a new decoder that decodes from the stream ```r```
    pub fn new(r: R) -> ImageResult<PngDecoder<R>> {
        let limits = png::Limits {
            bytes: usize::max_value(),
        };
        let mut decoder = png::Decoder::new_with_limits(r, limits);
        // By default the PNG decoder will scale 16 bpc to 8 bpc, so custom
        // transformations must be set. EXPAND preserves the default behavior
        // expanding bpc < 8 to 8 bpc.
        decoder.set_transformations(png::Transformations::EXPAND);
        let (_, mut reader) = decoder.read_info().map_err(ImageError::from_png)?;
        let (color_type, bits) = reader.output_color_type();
        let color_type = match (color_type, bits) {
            (png::ColorType::Grayscale, png::BitDepth::Eight) => ColorType::L8,
            (png::ColorType::Grayscale, png::BitDepth::Sixteen) => ColorType::L16,
            (png::ColorType::GrayscaleAlpha, png::BitDepth::Eight) => ColorType::La8,
            (png::ColorType::GrayscaleAlpha, png::BitDepth::Sixteen) => ColorType::La16,
            (png::ColorType::RGB, png::BitDepth::Eight) => ColorType::Rgb8,
            (png::ColorType::RGB, png::BitDepth::Sixteen) => ColorType::Rgb16,
            (png::ColorType::RGBA, png::BitDepth::Eight) => ColorType::Rgba8,
            (png::ColorType::RGBA, png::BitDepth::Sixteen) => ColorType::Rgba16,

            (png::ColorType::Grayscale, png::BitDepth::One) =>
                return Err(ImageError::UnsupportedColor(ExtendedColorType::L1)),
            (png::ColorType::GrayscaleAlpha, png::BitDepth::One) =>
                return Err(ImageError::UnsupportedColor(ExtendedColorType::La1)),
            (png::ColorType::RGB, png::BitDepth::One) =>
                return Err(ImageError::UnsupportedColor(ExtendedColorType::Rgb1)),
            (png::ColorType::RGBA, png::BitDepth::One) =>
                return Err(ImageError::UnsupportedColor(ExtendedColorType::Rgba1)),

            (png::ColorType::Grayscale, png::BitDepth::Two) =>
                return Err(ImageError::UnsupportedColor(ExtendedColorType::L2)),
            (png::ColorType::GrayscaleAlpha, png::BitDepth::Two) =>
                return Err(ImageError::UnsupportedColor(ExtendedColorType::La2)),
            (png::ColorType::RGB, png::BitDepth::Two) =>
                return Err(ImageError::UnsupportedColor(ExtendedColorType::Rgb2)),
            (png::ColorType::RGBA, png::BitDepth::Two) =>
                return Err(ImageError::UnsupportedColor(ExtendedColorType::Rgba2)),

            (png::ColorType::Grayscale, png::BitDepth::Four) =>
                return Err(ImageError::UnsupportedColor(ExtendedColorType::L4)),
            (png::ColorType::GrayscaleAlpha, png::BitDepth::Four) =>
                return Err(ImageError::UnsupportedColor(ExtendedColorType::La4)),
            (png::ColorType::RGB, png::BitDepth::Four) =>
                return Err(ImageError::UnsupportedColor(ExtendedColorType::Rgb4)),
            (png::ColorType::RGBA, png::BitDepth::Four) =>
                return Err(ImageError::UnsupportedColor(ExtendedColorType::Rgba4)),

            (png::ColorType::Indexed, bits) =>
                return Err(ImageError::UnsupportedColor(ExtendedColorType::Unknown(bits as u8))),
        };

        Ok(PngDecoder { color_type, reader })
    }
}

impl<'a, R: 'a + Read> ImageDecoder<'a> for PngDecoder<R> {
    type Reader = PNGReader<R>;

    fn dimensions(&self) -> (u32, u32) {
        self.reader.info().size()
    }

    fn color_type(&self) -> ColorType {
        self.color_type
    }

    fn into_reader(self) -> ImageResult<Self::Reader> {
        PNGReader::new(self.reader)
    }

    fn read_image(mut self, buf: &mut [u8]) -> ImageResult<()> {
        use byteorder::{BigEndian, ByteOrder, NativeEndian};

        assert_eq!(u64::try_from(buf.len()), Ok(self.total_bytes()));
        self.reader.next_frame(buf).map_err(ImageError::from_png)?;
        // PNG images are big endian. For 16 bit per channel and larger types,
        // the buffer may need to be reordered to native endianness per the
        // contract of `read_image`.
        // TODO: assumes equal channel bit depth.
        let bpc = self.color_type().bytes_per_pixel() / self.color_type().channel_count();
        match bpc {
            1 => (),  // No reodering necessary for u8
            2 => buf.chunks_mut(2).for_each(|c| {
                let v = BigEndian::read_u16(c);
                NativeEndian::write_u16(c, v)
            }),
            _ => unreachable!(),
        }
        Ok(())
    }

    fn scanline_bytes(&self) -> u64 {
        let width = self.reader.info().width;
        self.reader.output_line_size(width) as u64
    }
}

/// PNG encoder
pub struct PNGEncoder<W: Write> {
    w: W,
}

impl<W: Write> PNGEncoder<W> {
    /// Create a new encoder that writes its output to ```w```
    pub fn new(w: W) -> PNGEncoder<W> {
        PNGEncoder { w }
    }

    /// Encodes the image ```data```
    /// that has dimensions ```width``` and ```height```
    /// and ```ColorType``` ```c```
    pub fn encode(self, data: &[u8], width: u32, height: u32, color: ColorType) -> ImageResult<()> {
        let (ct, bits) = match color {
            ColorType::L8 => (png::ColorType::Grayscale, png::BitDepth::Eight),
            ColorType::L16 => (png::ColorType::Grayscale,png::BitDepth::Sixteen),
            ColorType::La8 => (png::ColorType::GrayscaleAlpha, png::BitDepth::Eight),
            ColorType::La16 => (png::ColorType::GrayscaleAlpha,png::BitDepth::Sixteen),
            ColorType::Rgb8 => (png::ColorType::RGB, png::BitDepth::Eight),
            ColorType::Rgb16 => (png::ColorType::RGB,png::BitDepth::Sixteen),
            ColorType::Rgba8 => (png::ColorType::RGBA, png::BitDepth::Eight),
            ColorType::Rgba16 => (png::ColorType::RGBA,png::BitDepth::Sixteen),
            _ => return Err(ImageError::UnsupportedColor(color.into())),
        };

        let mut encoder = png::Encoder::new(self.w, width, height);
        encoder.set_color(ct);
        encoder.set_depth(bits);
        let mut writer = encoder.write_header().map_err(|e| ImageError::IoError(e.into()))?;
        writer.write_image_data(data).map_err(|e| ImageError::IoError(e.into()))
    }
}

impl<W: Write> ImageEncoder for PNGEncoder<W> {
    fn write_image(
        self,
        buf: &[u8],
        width: u32,
        height: u32,
        color_type: ColorType,
    ) -> ImageResult<()> {
        use byteorder::{BigEndian, ByteOrder, NativeEndian};

        // PNG images are big endian. For 16 bit per channel and larger types,
        // the buffer may need to be reordered to big endian per the
        // contract of `write_image`.
        // TODO: assumes equal channel bit depth.
        let bpc = color_type.bytes_per_pixel() / color_type.channel_count();
        match bpc {
            1 => self.encode(buf, width, height, color_type),  // No reodering necessary for u8
            2 => {
                // Because the buffer is immutable and the PNG encoder does not
                // yet take Write/Read traits, create a temporary buffer for
                // big endian reordering.
                let mut reordered = vec![0; buf.len()];
                buf.chunks(2)
                    .zip(reordered.chunks_mut(2))
                    .for_each(|(b, r)| BigEndian::write_u16(r, NativeEndian::read_u16(b)));
                self.encode(&reordered, width, height, color_type)
            },
            _ => unreachable!(),
        }
    }
}

impl ImageError {
    fn from_png(err: png::DecodingError) -> ImageError {
        use png::DecodingError::*;
        match err {
            IoError(err) => ImageError::IoError(err),
            err @ Format(_) => ImageError::Decoding(DecodingError::new(
                ImageFormat::Png.into(),
                err,
            )),
            LimitsExceeded => ImageError::Limits(LimitError::from_kind(
                LimitErrorKind::InsufficientMemory,
            )),
            // Other is used when the buffer to `Reader::next_frame` is too small.
            Other(message) => ImageError::Parameter(ParameterError::from_kind(
                ParameterErrorKind::Generic(message.into_owned())
            )),
            err @ InvalidSignature
            | err @ CrcMismatch { .. }
            | err @ CorruptFlateStream => {
                ImageError::Decoding(DecodingError::new(
                    ImageFormat::Png.into(),
                    err,
                ))
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use crate::image::ImageDecoder;
    use std::io::Read;
    use super::*;

    #[test]
    fn ensure_no_decoder_off_by_one() {
        let dec = PngDecoder::new(std::fs::File::open("tests/images/png/bugfixes/debug_triangle_corners_widescreen.png").unwrap())
            .expect("Unable to read PNG file (does it exist?)");

        assert_eq![(2000, 1000), dec.dimensions()];

        assert_eq![
            ColorType::Rgb8,
            dec.color_type(),
            "Image MUST have the Rgb8 format"
        ];

        let correct_bytes = dec
            .into_reader()
            .expect("Unable to read file")
            .bytes()
            .map(|x| x.expect("Unable to read byte"))
            .collect::<Vec<u8>>();

        assert_eq![6_000_000, correct_bytes.len()];
    }

    #[test]
    fn underlying_error() {
        use std::error::Error;

        let mut not_png = std::fs::read("tests/images/png/bugfixes/debug_triangle_corners_widescreen.png").unwrap();
        not_png[0] = 0;

        let error = PngDecoder::new(&not_png[..]).err().unwrap();
        let _ = error
            .source()
            .unwrap()
            .downcast_ref::<png::DecodingError>()
            .expect("Caused by a png error");
    }
}
<|endoftext|>


<commit_before>use std::iter;
use byteorder::{BigEndian, ByteOrder};

pub const K: [u64; 80] = [
    0x428a2f98d728ae22,
    0x7137449123ef65cd,
    0xb5c0fbcfec4d3b2f,
    0xe9b5dba58189dbbc,
    0x3956c25bf348b538,
    0x59f111f1b605d019,
    0x923f82a4af194f9b,
    0xab1c5ed5da6d8118,
    0xd807aa98a3030242,
    0x12835b0145706fbe,
    0x243185be4ee4b28c,
    0x550c7dc3d5ffb4e2,
    0x72be5d74f27b896f,
    0x80deb1fe3b1696b1,
    0x9bdc06a725c71235,
    0xc19bf174cf692694,
    0xe49b69c19ef14ad2,
    0xefbe4786384f25e3,
    0x0fc19dc68b8cd5b5,
    0x240ca1cc77ac9c65,
    0x2de92c6f592b0275,
    0x4a7484aa6ea6e483,
    0x5cb0a9dcbd41fbd4,
    0x76f988da831153b5,
    0x983e5152ee66dfab,
    0xa831c66d2db43210,
    0xb00327c898fb213f,
    0xbf597fc7beef0ee4,
    0xc6e00bf33da88fc2,
    0xd5a79147930aa725,
    0x06ca6351e003826f,
    0x142929670a0e6e70,
    0x27b70a8546d22ffc,
    0x2e1b21385c26c926,
    0x4d2c6dfc5ac42aed,
    0x53380d139d95b3df,
    0x650a73548baf63de,
    0x766a0abb3c77b2a8,
    0x81c2c92e47edaee6,
    0x92722c851482353b,
    0xa2bfe8a14cf10364,
    0xa81a664bbc423001,
    0xc24b8b70d0f89791,
    0xc76c51a30654be30,
    0xd192e819d6ef5218,
    0xd69906245565a910,
    0xf40e35855771202a,
    0x106aa07032bbd1b8,
    0x19a4c116b8d2d0c8,
    0x1e376c085141ab53,
    0x2748774cdf8eeb99,
    0x34b0bcb5e19b48a8,
    0x391c0cb3c5c95a63,
    0x4ed8aa4ae3418acb,
    0x5b9cca4f7763e373,
    0x682e6ff3d6b2b8a3,
    0x748f82ee5defb2fc,
    0x78a5636f43172f60,
    0x84c87814a1f0ab72,
    0x8cc702081a6439ec,
    0x90befffa23631e28,
    0xa4506cebde82bde9,
    0xbef9a3f7b2c67915,
    0xc67178f2e372532b,
    0xca273eceea26619c,
    0xd186b8c721c0c207,
    0xeada7dd6cde0eb1e,
    0xf57d4f7fee6ed178,
    0x06f067aa72176fba,
    0x0a637dc5a2c898a6,
    0x113f9804bef90dae,
    0x1b710b35131c471b,
    0x28db77f523047d84,
    0x32caab7b40c72493,
    0x3c9ebe0a15c9bebc,
    0x431d67c49c100d4c,
    0x4cc5d4becb3e42b6,
    0x597f299cfc657e2a,
    0x5fcb6fab3ad6faec,
    0x6c44198c4a475817,
];

pub struct SHA512 {}

impl SHA512 {
    pub fn digest(message: &[u8]) -> [u8; 64] {
        Self::get_digest(message)
    }
}

impl Digest for SHA512 {
    const OUTPUT_LEN: usize = 64;
    const INITIAL_STATE: [u64; 8] = [
        0x6a09e667f3bcc908,
        0xbb67ae8584caa73b,
        0x3c6ef372fe94f82b,
        0xa54ff53a5f1d36f1,
        0x510e527fade682d1,
        0x9b05688c2b3e6c1f,
        0x1f83d9abfb41bd6b,
        0x5be0cd19137e2179,
    ];
}

pub struct SHA384 {}

impl SHA384 {
    // TODO: change the return type and/or do something else to reduce duplication
    pub fn digest(message: &[u8]) -> [u8; Self::OUTPUT_LEN] {
        let mut digest = [0; Self::OUTPUT_LEN];
        digest.copy_from_slice(&Self::get_digest(message)[..Self::OUTPUT_LEN]);
        digest
    }
}

impl Digest for SHA384 {
    const OUTPUT_LEN: usize = 48;
    const INITIAL_STATE: [u64; 8] = [
        0xcbbb9d5dc1059ed8,
        0x629a292a367cd507,
        0x9159015a3070dd17,
        0x152fecd8f70e5939,
        0x67332667ffc00b31,
        0x8eb44a8768581511,
        0xdb0c2e0d64f98fa7,
        0x47b5481dbefa4fa4,
    ];
}

trait Digest {
    const OUTPUT_LEN: usize;
    const INITIAL_STATE: [u64; 8];

    fn get_digest(message: &[u8]) -> [u8; 64] {
        let mut sha = SHA(Self::INITIAL_STATE);
        sha.process(message);
        sha.digest(Self::OUTPUT_LEN)
    }
}

struct SHA([u64; 8]);

impl SHA {
    fn process(&mut self, message: &[u8]) {
        let mut message = message.to_vec();
        Self::pad(&mut message);
        let mut w = [0; 80];
        for chunk in message.chunks(128) {
            for (wt, long) in w.iter_mut().zip(chunk.chunks(8)) {
                *wt = BigEndian::read_u64(long);
            }
            for t in 16..80 {
                w[t] = Self::ssig1(w[t - 2])
                    .wrapping_add(w[t - 7])
                    .wrapping_add(Self::ssig0(w[t - 15]))
                    .wrapping_add(w[t - 16]);
            }
            let mut a = self.0[0];
            let mut b = self.0[1];
            let mut c = self.0[2];
            let mut d = self.0[3];
            let mut e = self.0[4];
            let mut f = self.0[5];
            let mut g = self.0[6];
            let mut h = self.0[7];
            for (&kt, &wt) in K.iter().zip(w.iter()) {
                let t1 = h.wrapping_add(Self::bsig1(e))
                    .wrapping_add(Self::ch(e, f, g))
                    .wrapping_add(kt)
                    .wrapping_add(wt);
                let t2 = Self::bsig0(a).wrapping_add(Self::maj(a, b, c));
                h = g;
                g = f;
                f = e;
                e = d.wrapping_add(t1);
                d = c;
                c = b;
                b = a;
                a = t1.wrapping_add(t2);
            }
            self.0[0] = self.0[0].wrapping_add(a);
            self.0[1] = self.0[1].wrapping_add(b);
            self.0[2] = self.0[2].wrapping_add(c);
            self.0[3] = self.0[3].wrapping_add(d);
            self.0[4] = self.0[4].wrapping_add(e);
            self.0[5] = self.0[5].wrapping_add(f);
            self.0[6] = self.0[6].wrapping_add(g);
            self.0[7] = self.0[7].wrapping_add(h);
        }
    }

    fn digest(self, len: usize) -> [u8; 64] {
        let mut digest = [0; 64];
        for (chunk, &long) in digest.chunks_mut(8).zip(&self.0).take(len / 8) {
            BigEndian::write_u64(chunk, long);
        }
        digest
    }

    fn ch(x: u64, y: u64, z: u64) -> u64 {
        (x & y) ^ (!x & z)
    }

    fn maj(x: u64, y: u64, z: u64) -> u64 {
        (x & y) ^ (x & z) ^ (y & z)
    }

    fn bsig0(x: u64) -> u64 {
        x.rotate_right(28) ^ x.rotate_right(34) ^ x.rotate_right(39)
    }

    fn bsig1(x: u64) -> u64 {
        x.rotate_right(14) ^ x.rotate_right(18) ^ x.rotate_right(41)
    }

    fn ssig0(x: u64) -> u64 {
        x.rotate_right(1) ^ x.rotate_right(8) ^ (x >> 7)
    }

    fn ssig1(x: u64) -> u64 {
        x.rotate_right(19) ^ x.rotate_right(61) ^ (x >> 6)
    }

    /// Only supports messages with at most 2^64 - 1 bits for now
    fn pad(bytes: &mut Vec<u8>) {
        let len = len(bytes);
        bytes.push(0x80);
        let padding = (128 + 112 - bytes.len() % 128) % 128;
        bytes.extend(iter::repeat(0).take(padding));
        bytes.extend_from_slice(&[0; 8]);
        bytes.extend_from_slice(&len);
    }
}

fn len(bytes: &[u8]) -> [u8; 8] {
    let mut len = [0; 8];
    BigEndian::write_u64(&mut len, 8 * bytes.len() as u64);
    len
}

#[cfg(test)]
mod tests {
    use sha::*;
    use test_helpers::*;

    #[test]
    fn test_pad() {
        let mut message = vec![0b01100001, 0b01100010, 0b01100011, 0b01100100, 0b01100101];
        let expected = h2b(
            "6162636465800000000000000000000000000000000000000000000000000000\
             0000000000000000000000000000000000000000000000000000000000000000\
             0000000000000000000000000000000000000000000000000000000000000000\
             0000000000000000000000000000000000000000000000000000000000000028",
        );
        SHA::pad(&mut message);
        assert_eq!(expected, message);
    }

    #[test]
    fn test_digest() {
        let mut expected = h2b(
            "cf83e1357eefb8bdf1542850d66d8007d620e4050b5715dc83f4a921d36ce9ce\
             47d0d13c5d85f2b0ff8318d2877eec2f63b931bd47417a81a538327af927da3e",
        );
        let actual = SHA512::digest(&[]);
        assert_eq!(expected, actual.to_vec());

        expected = h2b(
            "38b060a751ac96384cd9327eb1b1e36a21fdb71114be07434c0cc7bf63f6e1da\
             274edebfe76f65fbd51ad2f14898b95b",
        );
        let actual = SHA384::digest(&[]);
        assert_eq!(expected, actual.to_vec());
    }
}
<commit_msg>sha: make K private<commit_after>use std::iter;
use byteorder::{BigEndian, ByteOrder};

const K: [u64; 80] = [
    0x428a2f98d728ae22,
    0x7137449123ef65cd,
    0xb5c0fbcfec4d3b2f,
    0xe9b5dba58189dbbc,
    0x3956c25bf348b538,
    0x59f111f1b605d019,
    0x923f82a4af194f9b,
    0xab1c5ed5da6d8118,
    0xd807aa98a3030242,
    0x12835b0145706fbe,
    0x243185be4ee4b28c,
    0x550c7dc3d5ffb4e2,
    0x72be5d74f27b896f,
    0x80deb1fe3b1696b1,
    0x9bdc06a725c71235,
    0xc19bf174cf692694,
    0xe49b69c19ef14ad2,
    0xefbe4786384f25e3,
    0x0fc19dc68b8cd5b5,
    0x240ca1cc77ac9c65,
    0x2de92c6f592b0275,
    0x4a7484aa6ea6e483,
    0x5cb0a9dcbd41fbd4,
    0x76f988da831153b5,
    0x983e5152ee66dfab,
    0xa831c66d2db43210,
    0xb00327c898fb213f,
    0xbf597fc7beef0ee4,
    0xc6e00bf33da88fc2,
    0xd5a79147930aa725,
    0x06ca6351e003826f,
    0x142929670a0e6e70,
    0x27b70a8546d22ffc,
    0x2e1b21385c26c926,
    0x4d2c6dfc5ac42aed,
    0x53380d139d95b3df,
    0x650a73548baf63de,
    0x766a0abb3c77b2a8,
    0x81c2c92e47edaee6,
    0x92722c851482353b,
    0xa2bfe8a14cf10364,
    0xa81a664bbc423001,
    0xc24b8b70d0f89791,
    0xc76c51a30654be30,
    0xd192e819d6ef5218,
    0xd69906245565a910,
    0xf40e35855771202a,
    0x106aa07032bbd1b8,
    0x19a4c116b8d2d0c8,
    0x1e376c085141ab53,
    0x2748774cdf8eeb99,
    0x34b0bcb5e19b48a8,
    0x391c0cb3c5c95a63,
    0x4ed8aa4ae3418acb,
    0x5b9cca4f7763e373,
    0x682e6ff3d6b2b8a3,
    0x748f82ee5defb2fc,
    0x78a5636f43172f60,
    0x84c87814a1f0ab72,
    0x8cc702081a6439ec,
    0x90befffa23631e28,
    0xa4506cebde82bde9,
    0xbef9a3f7b2c67915,
    0xc67178f2e372532b,
    0xca273eceea26619c,
    0xd186b8c721c0c207,
    0xeada7dd6cde0eb1e,
    0xf57d4f7fee6ed178,
    0x06f067aa72176fba,
    0x0a637dc5a2c898a6,
    0x113f9804bef90dae,
    0x1b710b35131c471b,
    0x28db77f523047d84,
    0x32caab7b40c72493,
    0x3c9ebe0a15c9bebc,
    0x431d67c49c100d4c,
    0x4cc5d4becb3e42b6,
    0x597f299cfc657e2a,
    0x5fcb6fab3ad6faec,
    0x6c44198c4a475817,
];

pub struct SHA512 {}

impl SHA512 {
    pub fn digest(message: &[u8]) -> [u8; 64] {
        Self::get_digest(message)
    }
}

impl Digest for SHA512 {
    const OUTPUT_LEN: usize = 64;
    const INITIAL_STATE: [u64; 8] = [
        0x6a09e667f3bcc908,
        0xbb67ae8584caa73b,
        0x3c6ef372fe94f82b,
        0xa54ff53a5f1d36f1,
        0x510e527fade682d1,
        0x9b05688c2b3e6c1f,
        0x1f83d9abfb41bd6b,
        0x5be0cd19137e2179,
    ];
}

pub struct SHA384 {}

impl SHA384 {
    // TODO: change the return type and/or do something else to reduce duplication
    pub fn digest(message: &[u8]) -> [u8; Self::OUTPUT_LEN] {
        let mut digest = [0; Self::OUTPUT_LEN];
        digest.copy_from_slice(&Self::get_digest(message)[..Self::OUTPUT_LEN]);
        digest
    }
}

impl Digest for SHA384 {
    const OUTPUT_LEN: usize = 48;
    const INITIAL_STATE: [u64; 8] = [
        0xcbbb9d5dc1059ed8,
        0x629a292a367cd507,
        0x9159015a3070dd17,
        0x152fecd8f70e5939,
        0x67332667ffc00b31,
        0x8eb44a8768581511,
        0xdb0c2e0d64f98fa7,
        0x47b5481dbefa4fa4,
    ];
}

trait Digest {
    const OUTPUT_LEN: usize;
    const INITIAL_STATE: [u64; 8];

    fn get_digest(message: &[u8]) -> [u8; 64] {
        let mut sha = SHA(Self::INITIAL_STATE);
        sha.process(message);
        sha.digest(Self::OUTPUT_LEN)
    }
}

struct SHA([u64; 8]);

impl SHA {
    fn process(&mut self, message: &[u8]) {
        let mut message = message.to_vec();
        Self::pad(&mut message);
        let mut w = [0; 80];
        for chunk in message.chunks(128) {
            for (wt, long) in w.iter_mut().zip(chunk.chunks(8)) {
                *wt = BigEndian::read_u64(long);
            }
            for t in 16..80 {
                w[t] = Self::ssig1(w[t - 2])
                    .wrapping_add(w[t - 7])
                    .wrapping_add(Self::ssig0(w[t - 15]))
                    .wrapping_add(w[t - 16]);
            }
            let mut a = self.0[0];
            let mut b = self.0[1];
            let mut c = self.0[2];
            let mut d = self.0[3];
            let mut e = self.0[4];
            let mut f = self.0[5];
            let mut g = self.0[6];
            let mut h = self.0[7];
            for (&kt, &wt) in K.iter().zip(w.iter()) {
                let t1 = h.wrapping_add(Self::bsig1(e))
                    .wrapping_add(Self::ch(e, f, g))
                    .wrapping_add(kt)
                    .wrapping_add(wt);
                let t2 = Self::bsig0(a).wrapping_add(Self::maj(a, b, c));
                h = g;
                g = f;
                f = e;
                e = d.wrapping_add(t1);
                d = c;
                c = b;
                b = a;
                a = t1.wrapping_add(t2);
            }
            self.0[0] = self.0[0].wrapping_add(a);
            self.0[1] = self.0[1].wrapping_add(b);
            self.0[2] = self.0[2].wrapping_add(c);
            self.0[3] = self.0[3].wrapping_add(d);
            self.0[4] = self.0[4].wrapping_add(e);
            self.0[5] = self.0[5].wrapping_add(f);
            self.0[6] = self.0[6].wrapping_add(g);
            self.0[7] = self.0[7].wrapping_add(h);
        }
    }

    fn digest(self, len: usize) -> [u8; 64] {
        let mut digest = [0; 64];
        for (chunk, &long) in digest.chunks_mut(8).zip(&self.0).take(len / 8) {
            BigEndian::write_u64(chunk, long);
        }
        digest
    }

    fn ch(x: u64, y: u64, z: u64) -> u64 {
        (x & y) ^ (!x & z)
    }

    fn maj(x: u64, y: u64, z: u64) -> u64 {
        (x & y) ^ (x & z) ^ (y & z)
    }

    fn bsig0(x: u64) -> u64 {
        x.rotate_right(28) ^ x.rotate_right(34) ^ x.rotate_right(39)
    }

    fn bsig1(x: u64) -> u64 {
        x.rotate_right(14) ^ x.rotate_right(18) ^ x.rotate_right(41)
    }

    fn ssig0(x: u64) -> u64 {
        x.rotate_right(1) ^ x.rotate_right(8) ^ (x >> 7)
    }

    fn ssig1(x: u64) -> u64 {
        x.rotate_right(19) ^ x.rotate_right(61) ^ (x >> 6)
    }

    /// Only supports messages with at most 2^64 - 1 bits for now
    fn pad(bytes: &mut Vec<u8>) {
        let len = len(bytes);
        bytes.push(0x80);
        let padding = (128 + 112 - bytes.len() % 128) % 128;
        bytes.extend(iter::repeat(0).take(padding));
        bytes.extend_from_slice(&[0; 8]);
        bytes.extend_from_slice(&len);
    }
}

fn len(bytes: &[u8]) -> [u8; 8] {
    let mut len = [0; 8];
    BigEndian::write_u64(&mut len, 8 * bytes.len() as u64);
    len
}

#[cfg(test)]
mod tests {
    use sha::*;
    use test_helpers::*;

    #[test]
    fn test_pad() {
        let mut message = vec![0b01100001, 0b01100010, 0b01100011, 0b01100100, 0b01100101];
        let expected = h2b(
            "6162636465800000000000000000000000000000000000000000000000000000\
             0000000000000000000000000000000000000000000000000000000000000000\
             0000000000000000000000000000000000000000000000000000000000000000\
             0000000000000000000000000000000000000000000000000000000000000028",
        );
        SHA::pad(&mut message);
        assert_eq!(expected, message);
    }

    #[test]
    fn test_digest() {
        let mut expected = h2b(
            "cf83e1357eefb8bdf1542850d66d8007d620e4050b5715dc83f4a921d36ce9ce\
             47d0d13c5d85f2b0ff8318d2877eec2f63b931bd47417a81a538327af927da3e",
        );
        let actual = SHA512::digest(&[]);
        assert_eq!(expected, actual.to_vec());

        expected = h2b(
            "38b060a751ac96384cd9327eb1b1e36a21fdb71114be07434c0cc7bf63f6e1da\
             274edebfe76f65fbd51ad2f14898b95b",
        );
        let actual = SHA384::digest(&[]);
        assert_eq!(expected, actual.to_vec());
    }
}
<|endoftext|>


<commit_before>use core::prelude::*;
use core::atomic::{AtomicUsize, ATOMIC_USIZE_INIT, Ordering};
use interrupts::apic;
use collections::Vec;
use machine::{outb};

extern {
    fn smp_init_vector();
}

#[allow(non_upper_case_globals)]
pub static processor_count: AtomicUsize = ATOMIC_USIZE_INIT;

#[no_mangle]
pub static SMP_STACK_PTR: AtomicUsize = ATOMIC_USIZE_INIT;
#[no_mangle]
pub static SMP_CR3: AtomicUsize = ATOMIC_USIZE_INIT;

fn send_startup_interrupt(address: u32, id: u8) {
    outb(0x70, 0x0F);
    outb(0x71, 0x0A);
    unsafe {*(0x469 as *mut u32) = address; }
    let interrupt = 0x4500;
    apic::send_interrupt(id, interrupt);
    apic::wait_for_delivery();
    let interrupt = (address >> 12) | 0x4600;
    apic::send_interrupt(id, interrupt);
    apic::wait_for_delivery();
}

pub fn init(processors: &Vec<u8>) {
    let (self_id, is_bsp) = apic::whoami();
    if is_bsp {
        let startup_address: u32 = smp_init_vector as u32;
        for (number, id) in processors.iter().enumerate() {
            let stack_loc = 0x120_000 + (8192 * number);
            SMP_STACK_PTR.store(0x120_000 + (8192 * number), Ordering::SeqCst);
            if *id != self_id {
                unsafe {
                    // Trigger the page fault in advance
                    *(stack_loc as *mut u32) = 0xDEADBEEF;
                    *((stack_loc - 4) as *mut u32) = 0;
                    *((stack_loc - 0x1000 - 4) as *mut u32) = 0;
                }
                let old_processors_count = processor_count.load(Ordering::Relaxed);
                send_startup_interrupt(startup_address, *id);
                while processor_count.load(Ordering::SeqCst) <= old_processors_count { };
            }
        }
    }
}
<commit_msg>Use stack_loc everywhere<commit_after>use core::prelude::*;
use core::atomic::{AtomicUsize, ATOMIC_USIZE_INIT, Ordering};
use interrupts::apic;
use collections::Vec;
use machine::{outb};

extern {
    fn smp_init_vector();
}

#[allow(non_upper_case_globals)]
pub static processor_count: AtomicUsize = ATOMIC_USIZE_INIT;

#[no_mangle]
pub static SMP_STACK_PTR: AtomicUsize = ATOMIC_USIZE_INIT;
#[no_mangle]
pub static SMP_CR3: AtomicUsize = ATOMIC_USIZE_INIT;

fn send_startup_interrupt(address: u32, id: u8) {
    outb(0x70, 0x0F);
    outb(0x71, 0x0A);
    unsafe {*(0x469 as *mut u32) = address; }
    let interrupt = 0x4500;
    apic::send_interrupt(id, interrupt);
    apic::wait_for_delivery();
    let interrupt = (address >> 12) | 0x4600;
    apic::send_interrupt(id, interrupt);
    apic::wait_for_delivery();
}

pub fn init(processors: &Vec<u8>) {
    let (self_id, is_bsp) = apic::whoami();
    if is_bsp {
        let startup_address: u32 = smp_init_vector as u32;
        for (number, id) in processors.iter().enumerate() {
            let stack_loc = 0x120_000 + (8192 * number);
            SMP_STACK_PTR.store(stack_loc, Ordering::SeqCst);
            if *id != self_id {
                unsafe {
                    // Trigger the page fault in advance
                    *(stack_loc as *mut u32) = 0xDEADBEEF;
                    *((stack_loc - 4) as *mut u32) = 0;
                    *((stack_loc - 0x1000 - 4) as *mut u32) = 0;
                }
                let old_processors_count = processor_count.load(Ordering::Relaxed);
                send_startup_interrupt(startup_address, *id);
                while processor_count.load(Ordering::SeqCst) <= old_processors_count { };
            }
        }
    }
}
<|endoftext|>


<commit_before>//! **UNSTABLE:** Structured key-value pairs.
//!
//! This module is unstable and breaking changes may be made
//! at any time. See [the tracking issue](https://github.com/rust-lang-nursery/log/issues/328)
//! for more details.
//!
//! Add the `kv_unstable` feature to your `Cargo.toml` to enable
//! this module:
//!
//! ```toml
//! [dependencies.log]
//! features = ["kv_unstable"]
//! ```

mod error;
mod source;
mod key;
pub mod value;

pub use self::error::Error;
pub use self::source::{Source, Visitor};
pub use self::key::{Key, ToKey};
pub use self::value::{Value, ToValue};
<commit_msg>inline docs for kv::value<commit_after>//! **UNSTABLE:** Structured key-value pairs.
//!
//! This module is unstable and breaking changes may be made
//! at any time. See [the tracking issue](https://github.com/rust-lang-nursery/log/issues/328)
//! for more details.
//!
//! Add the `kv_unstable` feature to your `Cargo.toml` to enable
//! this module:
//!
//! ```toml
//! [dependencies.log]
//! features = ["kv_unstable"]
//! ```

mod error;
mod key;
mod source;

pub mod value;

pub use self::error::Error;
pub use self::key::{Key, ToKey};
pub use self::source::{Source, Visitor};

#[doc(inline)]
pub use self::value::{ToValue, Value};
<|endoftext|>


<commit_before>use mysql::conn::MyOpts;
use mysql::conn::pool;
use mysql::conn::pool::{MyPooledConn,MyPool};
use mysql::value::from_value;

use std::thread::sleep_ms;

use lib::DownloadError;
use std::error::Error;
use lib::downloader::DownloadDB;
use CONFIG;

///Move result value out, return with none on err & print
macro_rules! try_reoption { ($e:expr) => (match $e { Ok(x) => x, Err(e) => {warn!("{}",e);return None; },}) }
macro_rules! try_return { ($e:expr) => (match $e { Ok(x) => x, Err(e) => {warn!("{}",e);return; },}) }

macro_rules! try_option { ($e:expr) => (match $e { Some(x) => x, None => return None }) }

/// Connect to DBMS, retry on failure.
pub fn db_connect(opts: MyOpts, sleep_time: u32) -> MyPool {
    loop {
        match pool::MyPool::new(opts.clone()) {
            Ok(conn) => {return conn;},
            Err(err) => error!("Unable to establish a connection: {}",err),
        };
        sleep_ms(sleep_time);
    }
}

///Set the state of the current query, code dependent, see QueryCodes
pub fn set_query_state(pool: & pool::MyPool,qid: &i64 , state: &str, finished: bool){ // same here
    let mut conn = pool.get_conn().unwrap();
    let progress: i32 = if finished {
        100
    }else{
        0
    };
    let mut stmt = try_return!(conn.prepare("UPDATE querydetails SET status = ? , progress = ? WHERE qid = ?"));
    let result = stmt.execute((&state,&progress,qid)); // why is this var needed ?!
    match result {
        Ok(_) => (),
        Err(why) => error!("Error setting query state: {}",why),
    }
}

pub fn set_null_state(pool: & pool::MyPool, qid: &i64){
	let mut conn = pool.get_conn().unwrap();
	let mut stmt = try_return!(conn.prepare("UPDATE querydetails SET status = NULL WHERE qid = ?"));
	let result = stmt.execute((qid,));
	match result {
	    Ok(_) => (),
	    Err(why) => error!("Error setting query sate: {}", why),
	}
}

/// Update status code for query entrys
pub fn set_query_code(conn: & mut MyPooledConn, code: &i8, qid: &i64) -> Result<(), DownloadError> { // same here
    let mut stmt = conn.prepare("UPDATE querydetails SET code = ? WHERE qid = ?").unwrap();
    let result = stmt.execute((&code,&qid));
    match result {
        Ok(_) => Ok(()),
        Err(why) => Err(DownloadError::DBError(why.description().into())),
    }
}

/// Update progress steps for db entrys
pub fn update_steps(pool: & pool::MyPool, qid: &i64, step: i32, max_steps: i32, finished: bool){
    set_query_state(&pool,qid, &format!("{}|{}", step, max_steps), finished);
}

/// Add file to db including it's name & fid based on the qid
pub fn add_file_entry(pool: & pool::MyPool, fid: &i64, name: &str, real_name: &str){
    trace!("name: {}",name);
    let mut conn = pool.get_conn().unwrap();
    let mut stmt = conn.prepare("INSERT INTO files (fid,name,rname,valid) VALUES (?,?,?,?)").unwrap();
    let result = stmt.execute((fid,&real_name,&name,&1)); // why is this var needed ?!
    match result {
        Ok(_) => (),
        Err(why) => error!("Error adding file: {}",why),
    }
}

/// Add query status msg for error reporting
pub fn add_query_status(pool: & pool::MyPool, qid: &i64, status: &str){
    let mut conn = pool.get_conn().unwrap();
    let mut stmt = conn.prepare("INSERT INTO querystatus (qid,msg) VALUES (?,?)").unwrap();
    let result = stmt.execute((&qid,&status));
    match result {
        Ok(_) => (),
        Err(why) => error!("Error inserting querystatus: {}",why),
    }
}

/// Request an entry from the DB to handle
pub fn request_entry(pool: & pool::MyPool) -> Option<DownloadDB> {
    let mut conn = try_reoption!(pool.get_conn());
    let mut stmt = try_reoption!(conn.prepare("SELECT queries.qid,url,`type`,quality,zip,`from`,`to` FROM querydetails \
                    INNER JOIN queries \
                    ON querydetails.qid = queries.qid \
                    LEFT JOIN playlists ON queries.qid = playlists.qid \
                    WHERE querydetails.code = -1 \
                    ORDER BY queries.created \
                    LIMIT 1"));
    let mut result = try_reoption!(stmt.execute(()));
    let result = try_reoption!(try_option!(result.next())); // result.next().'Some'->value.'unwrap'

    trace!("Result: {:?}", result[0]);
    trace!("result str: {}", result[1].into_str());
    let from;
    let to;
    let compress;
    let playlist = from_value::<Option<i16>>(result[4].clone()).is_some();
    if playlist {
        compress = from_value::<Option<i16>>(result[4].clone()).unwrap() == 1;
        from = from_value::<Option<i32>>(result[5].clone()).unwrap();
        to = from_value::<Option<i32>>(result[6].clone()).unwrap();
    } else {
        from = 0;
        to = 0;
        compress = false;
    }
    let download_db = DownloadDB { url: from_value::<String>(result[1].clone()),
                                    quality: from_value::<i16>(result[3].clone()),
                                    qid: from_value::<i64>(result[0].clone()),
                                    folder: CONFIG.general.temp_dir.clone(),
                                    pool: pool.clone(),
                                    playlist: playlist,
                                    compress: compress,
                                    to: to,
                                    from: from,
                                    source_type: from_value::<i16>(result[2].clone()) };
    Some(download_db)
}

/// Set DBMS connection settings
pub fn mysql_options() -> MyOpts {
    MyOpts {
        tcp_addr: Some(CONFIG.db.ip.clone()),
        tcp_port: CONFIG.db.port,
        user: Some(CONFIG.db.user.clone()),
        pass: Some(CONFIG.db.password.clone()),
        db_name: Some(CONFIG.db.db.clone()),
        ..Default::default() // set others to default
    }
}<commit_msg>fixed name insertion error changed file-insert error to be fatal, causing the files to be deleted this is necessary as those files would be existing, but untracked since they're not listed in the DB<commit_after>use mysql::conn::MyOpts;
use mysql::conn::pool;
use mysql::conn::pool::{MyPooledConn,MyPool};
use mysql::value::from_value;

use std::thread::sleep_ms;

use lib::DownloadError;
use std::error::Error;
use lib::downloader::DownloadDB;
use CONFIG;

///Move result value out, return with none on err & print
macro_rules! try_reoption { ($e:expr) => (match $e { Ok(x) => x, Err(e) => {warn!("{}",e);return None; },}) }
macro_rules! try_return { ($e:expr) => (match $e { Ok(x) => x, Err(e) => {warn!("{}",e);return; },}) }

macro_rules! try_option { ($e:expr) => (match $e { Some(x) => x, None => return None }) }

/// Connect to DBMS, retry on failure.
pub fn db_connect(opts: MyOpts, sleep_time: u32) -> MyPool {
    loop {
        match pool::MyPool::new(opts.clone()) {
            Ok(conn) => {return conn;},
            Err(err) => error!("Unable to establish a connection: {}",err),
        };
        sleep_ms(sleep_time);
    }
}

///Set the state of the current query, code dependent, see QueryCodes
pub fn set_query_state(pool: & pool::MyPool,qid: &i64 , state: &str, finished: bool){ // same here
    let mut conn = pool.get_conn().unwrap();
    let progress: i32 = if finished {
        100
    }else{
        0
    };
    let mut stmt = try_return!(conn.prepare("UPDATE querydetails SET status = ? , progress = ? WHERE qid = ?"));
    let result = stmt.execute((&state,&progress,qid)); // why is this var needed ?!
    match result {
        Ok(_) => (),
        Err(why) => error!("Error setting query state: {}",why),
    }
}

pub fn set_null_state(pool: & pool::MyPool, qid: &i64){
	let mut conn = pool.get_conn().unwrap();
	let mut stmt = try_return!(conn.prepare("UPDATE querydetails SET status = NULL WHERE qid = ?"));
	let result = stmt.execute((qid,));
	match result {
	    Ok(_) => (),
	    Err(why) => error!("Error setting query sate: {}", why),
	}
}

/// Update status code for query entrys
pub fn set_query_code(conn: & mut MyPooledConn, code: &i8, qid: &i64) -> Result<(), DownloadError> { // same here
    let mut stmt = conn.prepare("UPDATE querydetails SET code = ? WHERE qid = ?").unwrap();
    let result = stmt.execute((&code,&qid));
    match result {
        Ok(_) => Ok(()),
        Err(why) => Err(DownloadError::DBError(why.description().into())),
    }
}

/// Update progress steps for db entrys
pub fn update_steps(pool: & pool::MyPool, qid: &i64, step: i32, max_steps: i32, finished: bool){
    set_query_state(&pool,qid, &format!("{}|{}", step, max_steps), finished);
}

/// Add file to db including it's name & fid based on the qid
pub fn add_file_entry(pool: & pool::MyPool, fid: &i64, name: &str, real_name: &str) -> Result<(), DownloadError> {
    trace!("name: {:?}",name);
    let mut conn = pool.get_conn().unwrap();
    let mut stmt = conn.prepare("INSERT INTO files (fid,rname,name,valid) VALUES (?,?,?,?)").unwrap();
    try!(stmt.execute((fid,&real_name,&name,&1))); // why is this var needed ?!
//    match result {
//        Ok(_) => (),
//        Err(why) => error!("Error adding file: {}",why),
//    }
	Ok(())
}

/// Add query status msg for error reporting
pub fn add_query_status(pool: & pool::MyPool, qid: &i64, status: &str){
    let mut conn = pool.get_conn().unwrap();
    let mut stmt = conn.prepare("INSERT INTO querystatus (qid,msg) VALUES (?,?)").unwrap();
    let result = stmt.execute((&qid,&status));
    match result {
        Ok(_) => (),
        Err(why) => error!("Error inserting querystatus: {}",why),
    }
}

/// Request an entry from the DB to handle
pub fn request_entry(pool: & pool::MyPool) -> Option<DownloadDB> {
    let mut conn = try_reoption!(pool.get_conn());
    let mut stmt = try_reoption!(conn.prepare("SELECT queries.qid,url,`type`,quality,zip,`from`,`to` FROM querydetails \
                    INNER JOIN queries \
                    ON querydetails.qid = queries.qid \
                    LEFT JOIN playlists ON queries.qid = playlists.qid \
                    WHERE querydetails.code = -1 \
                    ORDER BY queries.created \
                    LIMIT 1"));
    let mut result = try_reoption!(stmt.execute(()));
    let result = try_reoption!(try_option!(result.next())); // result.next().'Some'->value.'unwrap'

    trace!("Result: {:?}", result[0]);
    trace!("result str: {}", result[1].into_str());
    let from;
    let to;
    let compress;
    let playlist = from_value::<Option<i16>>(result[4].clone()).is_some();
    if playlist {
        compress = from_value::<Option<i16>>(result[4].clone()).unwrap() == 1;
        from = from_value::<Option<i32>>(result[5].clone()).unwrap();
        to = from_value::<Option<i32>>(result[6].clone()).unwrap();
    } else {
        from = 0;
        to = 0;
        compress = false;
    }
    let download_db = DownloadDB { url: from_value::<String>(result[1].clone()),
                                    quality: from_value::<i16>(result[3].clone()),
                                    qid: from_value::<i64>(result[0].clone()),
                                    folder: CONFIG.general.temp_dir.clone(),
                                    pool: pool.clone(),
                                    playlist: playlist,
                                    compress: compress,
                                    to: to,
                                    from: from,
                                    source_type: from_value::<i16>(result[2].clone()) };
    Some(download_db)
}

/// Set DBMS connection settings
pub fn mysql_options() -> MyOpts {
    MyOpts {
        tcp_addr: Some(CONFIG.db.ip.clone()),
        tcp_port: CONFIG.db.port,
        user: Some(CONFIG.db.user.clone()),
        pass: Some(CONFIG.db.password.clone()),
        db_name: Some(CONFIG.db.db.clone()),
        ..Default::default() // set others to default
    }
}<|endoftext|>


<commit_before>extern crate walkdir;

use std::env;
use std::process::Command;

use walkdir::{DirEntry, WalkDir};

#[cfg(all(not(feature = "rustc-dep-of-std"), not(feature = "with_submodule")))]
fn build() {
	#[cfg(windows)]
	let out_dir = env::temp_dir().to_str().unwrap().to_owned();
	#[cfg(not(windows))]
	let out_dir = env::var("OUT_DIR").unwrap();
	let profile = env::var("PROFILE").expect("PROFILE was not set");

	let _output = Command::new("cargo")
		.current_dir(out_dir.clone())
		.arg("download")
		.arg("--output")
		.arg(out_dir.clone() + "/rusty-hermit")
		.arg("--extract")
		.arg("rusty-hermit")
		.output()
		.expect("Unable to download rusty-hermit. Please install `cargo-download`.");

	let mut cmd = Command::new("cargo");
	cmd.current_dir(out_dir.clone() + "/rusty-hermit")
		.arg("build")
		.arg("-Z")
		.arg("build-std=core,alloc")
		.arg("--target")
		.arg("x86_64-unknown-hermit-kernel");

	if profile == "release" {
		cmd.arg("--release");
	}

	#[cfg(feature = "instrument")]
	cmd.env("RUSTFLAGS", "-Z instrument-mcount");
	// if instrument is not set, ensure that instrument is not in environment variables!
	#[cfg(not(feature = "instrument"))]
	cmd.env(
		"RUSTFLAGS",
		env::var("RUSTFLAGS")
			.unwrap_or_else(|_| "".into())
			.replace("-Z instrument-mcount", ""),
	);

	let output = cmd.output().expect("Unable to build kernel");
	let stdout = std::string::String::from_utf8(output.stdout);
	let stderr = std::string::String::from_utf8(output.stderr);
	println!("Build libhermit-rs output-status: {}", output.status);
	println!("Build libhermit-rs output-stdout: {}", stdout.unwrap());
	println!("Build libhermit-rs output-stderr: {}", stderr.unwrap());
	assert!(output.status.success());

	println!(
		"cargo:rustc-link-search=native={}/rusty-hermit/target/x86_64-unknown-hermit-kernel/{}",
		out_dir.clone(),
		profile
	);
	println!("cargo:rustc-link-lib=static=hermit");
}

#[cfg(all(not(feature = "rustc-dep-of-std"), feature = "with_submodule"))]
fn build() {
	let out_dir = env::var("OUT_DIR").unwrap();
	let target_dir = out_dir.clone() + "/target";
	let profile = env::var("PROFILE").expect("PROFILE was not set");

	let mut cmd = Command::new("cargo");
	cmd.current_dir("../libhermit-rs")
		.arg("build")
		.arg("-Z")
		.arg("build-std=core,alloc")
		.arg("--target")
		.arg("x86_64-unknown-hermit-kernel")
		.arg("--target-dir")
		.arg(target_dir);

	if profile == "release" {
		cmd.arg("--release");
	}

	#[cfg(feature = "instrument")]
	cmd.env("RUSTFLAGS", "-Z instrument-mcount");
	// if instrument is not set, ensure that instrument is not in environment variables!
	#[cfg(not(feature = "instrument"))]
	cmd.env(
		"RUSTFLAGS",
		env::var("RUSTFLAGS")
			.unwrap_or_else(|_| "".into())
			.replace("-Z instrument-mcount", ""),
	);

	let output = cmd.output().expect("Unable to build kernel");
	let stdout = std::string::String::from_utf8(output.stdout);
	let stderr = std::string::String::from_utf8(output.stderr);
	println!("Build libhermit-rs output-status: {}", output.status);
	println!("Build libhermit-rs output-stdout: {}", stdout.unwrap());
	println!("Build libhermit-rs output-stderr: {}", stderr.unwrap());
	assert!(output.status.success());

	println!(
		"cargo:rustc-link-search=native={}/target/x86_64-unknown-hermit-kernel/{}",
		out_dir, profile
	);
	println!("cargo:rustc-link-lib=static=hermit");
}

#[cfg(all(not(feature = "rustc-dep-of-std"), feature = "with_submodule"))]
fn configure_cargo_rerun_if_changed() {
	fn is_not_ignored(entry: &DirEntry) -> bool {
		// Ignore .git .vscode and target directories, but not .cargo or .github
		if entry.depth() == 1
			&& entry.path().is_dir()
			&& (entry.path().ends_with("target")
				|| entry.path().ends_with(".git")
				|| entry.path().ends_with(".vscode"))
		{
			return false;
		}
		true
	}

	WalkDir::new("../libhermit-rs")
		.into_iter()
		.filter_entry(|e| is_not_ignored(e))
		.filter_map(|v| v.ok())
		.filter_map(|v| v.path().canonicalize().ok())
		.for_each(|x| println!("cargo:rerun-if-changed={}", x.display()));
}

fn main() {
	#[cfg(all(not(feature = "rustc-dep-of-std"), feature = "with_submodule"))]
	configure_cargo_rerun_if_changed();
	#[cfg(not(feature = "rustc-dep-of-std"))]
	build();
}
<commit_msg>Build: Refactor to reduce duplicated Code (#32)<commit_after>extern crate walkdir;

use std::env;
use std::path::{Path, PathBuf};
use std::process::Command;

use walkdir::{DirEntry, WalkDir};

fn build_hermit(src_dir: &Path, target_dir_opt: Option<&Path>) {
	let profile = env::var("PROFILE").expect("PROFILE was not set");
	let mut cmd = Command::new("cargo");
	cmd.current_dir(src_dir)
		.arg("build")
		.arg("-Z")
		.arg("build-std=core,alloc")
		.arg("--target")
		.arg("x86_64-unknown-hermit-kernel");

	if let Some(target_dir) = target_dir_opt {
		cmd.arg("--target-dir").arg(target_dir);
	}
	let target_dir = match target_dir_opt {
		Some(target_dir) => target_dir.to_path_buf(),
		None => src_dir.join("target"),
	};

	if profile == "release" {
		cmd.arg("--release");
	}

	#[cfg(feature = "instrument")]
	cmd.env("RUSTFLAGS", "-Z instrument-mcount");
	// if instrument is not set, ensure that instrument is not in environment variables!
	#[cfg(not(feature = "instrument"))]
	cmd.env(
		"RUSTFLAGS",
		env::var("RUSTFLAGS")
			.unwrap_or_else(|_| "".into())
			.replace("-Z instrument-mcount", ""),
	);

	let output = cmd.output().expect("Unable to build kernel");
	let stdout = std::string::String::from_utf8(output.stdout);
	let stderr = std::string::String::from_utf8(output.stderr);

	println!("Build libhermit-rs output-status: {}", output.status);
	println!("Build libhermit-rs output-stdout: {}", stdout.unwrap());
	println!("Build libhermit-rs output-stderr: {}", stderr.unwrap());
	assert!(output.status.success());

	let lib_location = target_dir
		.join("x86_64-unknown-hermit-kernel")
		.join(&profile)
		.canonicalize()
		.unwrap(); // Must exist after building
	println!("cargo:rustc-link-search=native={}", lib_location.display());
	println!("cargo:rustc-link-lib=static=hermit");

	//HERMIT_LOG_LEVEL_FILTER sets the log level filter at compile time
	// Doesn't actually rebuild atm - see: https://github.com/rust-lang/cargo/issues/8306
	println!("cargo:rerun-if-env-changed=HERMIT_LOG_LEVEL_FILTER");
}

#[cfg(all(not(feature = "rustc-dep-of-std"), not(feature = "with_submodule")))]
fn build() {
	#[cfg(windows)]
	let out_dir = env::temp_dir();
	#[cfg(not(windows))]
	let out_dir = PathBuf::from(env::var("OUT_DIR").unwrap());
	let src_dir = out_dir.join("rusty-hermit");

	let _output = Command::new("cargo")
		.current_dir(out_dir)
		.arg("download")
		.arg("--output")
		.arg(src_dir.clone().into_os_string())
		.arg("--extract")
		.arg("rusty-hermit")
		.output()
		.expect("Unable to download rusty-hermit. Please install `cargo-download`.");

	build_hermit(src_dir.as_ref(), None);
}

#[cfg(all(not(feature = "rustc-dep-of-std"), feature = "with_submodule"))]
fn build() {
	let out_dir = PathBuf::from(env::var("OUT_DIR").unwrap());
	let target_dir = out_dir.clone().join("target");
	let src_dir = env::current_dir()
		.unwrap()
		.parent()
		.unwrap()
		.join("libhermit-rs");

	build_hermit(src_dir.as_ref(), Some(target_dir.as_ref()));
	configure_cargo_rerun_if_changed(src_dir.as_ref());
}

#[cfg(all(not(feature = "rustc-dep-of-std"), feature = "with_submodule"))]
fn configure_cargo_rerun_if_changed(src_dir: &Path) {
	fn is_not_ignored(entry: &DirEntry) -> bool {
		// Ignore .git .vscode and target directories, but not .cargo or .github
		if entry.depth() == 1
			&& entry.path().is_dir()
			&& (entry.path().ends_with("target")
				|| entry.path().ends_with(".git")
				|| entry.path().ends_with(".vscode"))
		{
			return false;
		}
		true
	}

	WalkDir::new(src_dir)
		.into_iter()
		.filter_entry(|e| is_not_ignored(e))
		.filter_map(|v| v.ok())
		.filter_map(|v| v.path().canonicalize().ok())
		.for_each(|x| println!("cargo:rerun-if-changed={}", x.display()));
}

fn main() {
	#[cfg(not(feature = "rustc-dep-of-std"))]
	build();
}
<|endoftext|>