<commit_before>// Claxon -- A FLAC decoding library in Rust
// Copyright 2014 Ruud van Asseldonk
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// A copy of the License has been included in the root of the repository.

//! Claxon, a FLAC decoding library.
//!
//! Examples
//! ========
//!
//! The following example computes the root mean square (RMS) of a FLAC file.
//!
//! ```
//! # use claxon;
//! let mut reader = claxon::FlacReader::open("testsamples/pop.flac").unwrap();
//! let mut sqr_sum = 0.0;
//! let mut count = 0;
//! for sample in reader.samples() {
//!     let s = sample.unwrap() as f64;
//!     sqr_sum += s * s;
//!     count += 1;
//! }
//! println!("RMS is {}", (sqr_sum / count as f64).sqrt());
//! ```
//!
//! A simple way to decode a file to wav with Claxon and
//! [Hound](https://github.com/ruuda/hound):
//!
//! ```
//! # extern crate hound;
//! # extern crate claxon;
//! # use std::path::Path;
//! # fn decode_file(fname: &Path) {
//! let mut reader = claxon::FlacReader::open(fname).expect("failed to open FLAC stream");
//!
//! let spec = hound::WavSpec {
//!     channels: reader.streaminfo().channels as u16,
//!     sample_rate: reader.streaminfo().sample_rate,
//!     bits_per_sample: reader.streaminfo().bits_per_sample as u16,
//!     sample_format: hound::SampleFormat::Int,
//! };
//!
//! let fname_wav = fname.with_extension("wav");
//! let opt_wav_writer = hound::WavWriter::create(fname_wav, spec);
//! let mut wav_writer = opt_wav_writer.expect("failed to create wav file");
//!
//! for opt_sample in reader.samples() {
//!     let sample = opt_sample.expect("failed to decode FLAC stream");
//!     wav_writer.write_sample(sample).expect("failed to write wav file");
//! }
//! # }
//! ```
//!
//! Retrieving the artist metadata:
//!
//! ```
//! # use claxon;
//! let reader = claxon::FlacReader::open("testsamples/pop.flac").unwrap();
//! for artist in reader.get_tag("ARTIST") {
//!     println!("{}", artist);
//! }
//! ```
//!
//! For more examples, see the [examples](https://github.com/ruuda/claxon/tree/master/examples)
//! directory in the crate.

#![warn(missing_docs)]

use std::fs;
use std::io;
use std::mem;
use std::path;
use error::fmt_err;
use frame::FrameReader;
use input::{BufferedReader, ReadBytes};
use metadata::{MetadataBlock, MetadataBlockReader, StreamInfo, VorbisComment};

mod crc;
mod error;
pub mod frame;
pub mod input;
pub mod metadata;
pub mod subframe;

pub use error::{Error, Result};
pub use frame::Block;

/// A FLAC decoder that can decode the stream from the underlying reader.
///
/// TODO: Add an example.
pub struct FlacReader<R: io::Read> {
    streaminfo: StreamInfo,
    vorbis_comment: Option<VorbisComment>,
    #[allow(dead_code)] // TODO: Expose metadata nicely.
    metadata_blocks: Vec<MetadataBlock>,
    input: BufferedReader<R>,
}

/// Controls what metadata `FlacReader` reads when constructed.
///
/// The FLAC format contains a number of metadata blocks before the start of
/// audio data. Reading these is wasteful if the data is never used. The
/// `FlacReaderOptions` indicate which blocks to look for. As soon as all
/// desired blocks have been read, a `FlacReader` is returned without reading
/// remaining metadata blocks.
pub struct FlacReaderOptions {
    /// When true, read metadata blocks at least until a Vorbis comment block is found.
    ///
    /// When false, the `FlacReader` may be constructed without reading a
    /// Vorbis comment block, even if the stream contains one. Consequently,
    /// `FlacReader::tags()` and other tag-related methods may not return tag
    /// data.
    pub find_vorbis_comment: bool,

    /// When true, read metadata blocks at least until a seek table block is found.
    pub find_seek_table: bool,
}

/// An iterator that yields samples read from a `FlacReader`.
pub struct FlacSamples<R: ReadBytes> {
    frame_reader: FrameReader<R>,
    block: Block,
    sample: u32,
    channel: u32,

    /// If reading ever failed, this flag is set, so that the iterator knows not
    /// to return any new values.
    has_failed: bool,
}

// TODO: Add a `FlacIntoSamples`.

fn read_stream_header<R: ReadBytes>(input: &mut R) -> Result<()> {
    // A FLAC stream starts with a 32-bit header 'fLaC' (big endian).
    const FLAC_HEADER: u32 = 0x66_4c_61_43;

    // Some files start with ID3 tag data. The reference decoder supports this
    // for convenience. Claxon does not, but we can at least generate a helpful
    // error message if a file starts like this.
    const ID3_HEADER: u32 = 0x49_44_33_00;

    let header = try!(input.read_be_u32());
    if header != FLAC_HEADER {
        if (header & 0xff_ff_ff_00) == ID3_HEADER {
            fmt_err("stream starts with ID3 header rather than FLAC header")
        } else {
            fmt_err("invalid stream header")
        }
    } else {
        Ok(())
    }
}

impl<R: io::Read> FlacReader<R> {
    /// Attempts to create a reader that reads the FLAC format.
    ///
    /// The header and metadata blocks are read immediately. Audio frames
    /// will be read on demand.
    ///
    /// Claxon rejects files that claim to contain excessively large metadata
    /// blocks, to protect against denial of service attacks where a
    /// small damaged or malicous file could cause gigabytes of memory
    /// to be allocated. `Error::Unsupported` is returned in that case.
    pub fn new(reader: R) -> Result<FlacReader<R>> {
        let mut buf_reader = BufferedReader::new(reader);

        // A flac stream first of all starts with a stream header.
        try!(read_stream_header(&mut buf_reader));

        // Start a new scope, because the input reader must be available again
        // for the frame reader next.
        let (streaminfo, vorbis_comment, metadata_blocks) = {
            // Next are one or more metadata blocks. The flac specification
            // dictates that the streaminfo block is the first block. The metadata
            // block reader will yield at least one element, so the unwrap is safe.
            let mut metadata_iter = MetadataBlockReader::new(&mut buf_reader);
            let streaminfo_block = try!(metadata_iter.next().unwrap());
            let streaminfo = match streaminfo_block {
                MetadataBlock::StreamInfo(info) => info,
                _ => return fmt_err("streaminfo block missing"),
            };

            let mut vorbis_comment = None;

            // There might be more metadata blocks, read and store them.
            let mut metadata_blocks = Vec::new();
            for block_result in metadata_iter {
                match try!(block_result) {
                    MetadataBlock::VorbisComment(vc) => {
                        // The Vorbis comment block need not be present, but
                        // when it is, it must be unique.
                        if vorbis_comment.is_some() {
                            return fmt_err("encountered second Vorbis comment block")
                        } else {
                            vorbis_comment = Some(vc);
                            break;
                        }
                    }
                    MetadataBlock::StreamInfo(..) => {
                        return fmt_err("encountered second streaminfo block")
                    }
                    // Other blocks are currently not handled.
                    block => metadata_blocks.push(block),
                }
            }

            (streaminfo, vorbis_comment, metadata_blocks)
        };

        // The flac reader will contain the reader that will read frames.
        let flac_reader = FlacReader {
            streaminfo: streaminfo,
            vorbis_comment: vorbis_comment,
            metadata_blocks: metadata_blocks,
            input: buf_reader,
        };

        Ok(flac_reader)
    }

    /// Returns the streaminfo metadata.
    ///
    /// This contains information like the sample rate and number of channels.
    pub fn streaminfo(&self) -> StreamInfo {
        self.streaminfo
    }

    /// Returns the vendor string of the Vorbis comment block, if present.
    ///
    /// This string usually contains the name and version of the program that
    /// encoded the FLAC stream, such as `reference libFLAC 1.3.2 20170101`
    /// or `Lavf57.25.100`.
    pub fn vendor(&self) -> Option<&str> {
        self.vorbis_comment.as_ref().map(|vc| &vc.vendor[..])
    }

    /// Returns name-value pairs of Vorbis comments, such as `("ARTIST", "Queen")`.
    ///
    /// The name is supposed to be interpreted case-insensitively, and is
    /// guaranteed to consist of ASCII characters. Claxon does not normalize
    /// the casing of the name. Use `get_tag()` to do a case-insensitive lookup.
    ///
    /// Names need not be unique. For instance, multiple `ARTIST` comments might
    /// be present on a collaboration track.
    ///
    /// See https://www.xiph.org/vorbis/doc/v-comment.html for more details.
    pub fn tags<'a>(&'a self) -> metadata::Tags<'a> {
        match self.vorbis_comment.as_ref() {
            Some(vc) => metadata::Tags::new(&vc.comments[..]),
            None => metadata::Tags::new(&[]),
        }
    }

    /// Look up a Vorbis comment such as `ARTIST` in a case-insensitive way.
    ///
    /// Returns an iterator,  because tags may occur more than once. There could
    /// be multiple `ARTIST` tags on a collaboration track, for instance.
    ///
    /// Note that tag names are ASCII and never contain `'='`; trying to look up
    /// a non-ASCII tag will return no results. Furthermore, the Vorbis comment
    /// spec dictates that tag names should be handled case-insensitively, so
    /// this method performs a case-insensitive lookup.
    ///
    /// See also `tags()` for access to the raw tags.
    /// See https://www.xiph.org/vorbis/doc/v-comment.html for more details.
    pub fn get_tag<'a>(&'a self, tag_name: &'a str) -> metadata::GetTag<'a> {
        match self.vorbis_comment.as_ref() {
            Some(vc) => metadata::GetTag::new(&vc.comments[..], tag_name),
            None => metadata::GetTag::new(&[], tag_name),
        }
    }

    /// Returns an iterator that decodes a single frame on every iteration.
    /// TODO: It is not an iterator.
    ///
    /// This is a low-level primitive that gives you control over when decoding
    /// happens. The representation of the decoded audio is somewhat specific to
    /// the FLAC format. For a higher-level interface, see `samples()`.
    pub fn blocks<'r>(&'r mut self) -> FrameReader<&'r mut BufferedReader<R>> {
        FrameReader::new(&mut self.input)
    }

    /// Returns an iterator over all samples.
    ///
    /// The channel data is is interleaved. The iterator is streaming. That is,
    /// if you call this method once, read a few samples, and call this method
    /// again, the second iterator will not start again from the beginning of
    /// the file. It will continue somewhere after where the first iterator
    /// stopped, and it might skip some samples. (This is because FLAC divides
    /// a stream into blocks, which have to be decoded entirely. If you drop the
    /// iterator, you lose the unread samples in that block.)
    ///
    /// This is a user-friendly interface that trades performance for ease of
    /// use. If performance is an issue, consider using `blocks()` instead.
    ///
    /// This is a high-level interface to the decoder. The cost of retrieving
    /// the next sample can vary significantly, as sometimes a new block has to
    /// be decoded. Additionally, there is a cost to every iteration returning a
    /// `Result`. When a block has been decoded, iterating the samples in that
    /// block can never fail, but a match on every sample is required
    /// nonetheless. For more control over when decoding happens, and less error
    /// handling overhead, use `blocks()`.
    pub fn samples<'r>(&'r mut self) -> FlacSamples<&'r mut BufferedReader<R>> {
        FlacSamples {
            frame_reader: frame::FrameReader::new(&mut self.input),
            block: Block::empty(),
            sample: 0,
            channel: 0,
            has_failed: false,
        }
    }

    /// Destroys the FLAC reader and returns the underlying reader.
    ///
    /// Because the reader employs buffering internally, anything in the buffer
    /// will be lost.
    pub fn into_inner(self) -> R {
        self.input.into_inner()
    }
}

impl FlacReader<fs::File> {
    /// Attempts to create a reader that reads from the specified file.
    ///
    /// This is a convenience constructor that opens a `File`, and constructs a
    /// `FlacReader` from it. There is no need to wrap the file in a
    /// `BufReader`, as the `FlacReader` employs buffering already.
    pub fn open<P: AsRef<path::Path>>(filename: P) -> Result<FlacReader<fs::File>> {
        let file = try!(fs::File::open(filename));
        FlacReader::new(file)
    }
}

impl<R: ReadBytes> Iterator for FlacSamples<R> {
    type Item = Result<i32>;

    fn next(&mut self) -> Option<Result<i32>> {
        // If the previous read failed, end iteration.
        if self.has_failed {
            return None;
        }

        // Iterate the samples channel interleaved, so first increment the
        // channel.
        self.channel += 1;

        // If that was the last channel, increment the sample number.
        if self.channel >= self.block.channels() {
            self.channel = 0;
            self.sample += 1;

            // If that was the last sample in the block, decode the next block.
            if self.sample >= self.block.duration() {
                self.sample = 0;

                // Replace the current block with an empty one so that we may
                // reuse the current buffer to decode again.
                let current_block = mem::replace(&mut self.block, Block::empty());

                match self.frame_reader.read_next_or_eof(current_block.into_buffer()) {
                    Ok(Some(next_block)) => {
                        self.block = next_block;
                    }
                    Ok(None) => {
                        // The stream ended with EOF.
                        // TODO: If a number of samples was specified in the
                        // streaminfo metadata block, verify that we did not
                        // read more or less samples.
                        return None;
                    }
                    Err(error) => {
                        self.has_failed = true;
                        return Some(Err(error));
                    }
                }
            }
        }

        Some(Ok(self.block.sample(self.channel, self.sample)))
    }
}
<commit_msg>Distinguish metadata-only reader in internal state<commit_after>// Claxon -- A FLAC decoding library in Rust
// Copyright 2014 Ruud van Asseldonk
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// A copy of the License has been included in the root of the repository.

//! Claxon, a FLAC decoding library.
//!
//! Examples
//! ========
//!
//! The following example computes the root mean square (RMS) of a FLAC file.
//!
//! ```
//! # use claxon;
//! let mut reader = claxon::FlacReader::open("testsamples/pop.flac").unwrap();
//! let mut sqr_sum = 0.0;
//! let mut count = 0;
//! for sample in reader.samples() {
//!     let s = sample.unwrap() as f64;
//!     sqr_sum += s * s;
//!     count += 1;
//! }
//! println!("RMS is {}", (sqr_sum / count as f64).sqrt());
//! ```
//!
//! A simple way to decode a file to wav with Claxon and
//! [Hound](https://github.com/ruuda/hound):
//!
//! ```
//! # extern crate hound;
//! # extern crate claxon;
//! # use std::path::Path;
//! # fn decode_file(fname: &Path) {
//! let mut reader = claxon::FlacReader::open(fname).expect("failed to open FLAC stream");
//!
//! let spec = hound::WavSpec {
//!     channels: reader.streaminfo().channels as u16,
//!     sample_rate: reader.streaminfo().sample_rate,
//!     bits_per_sample: reader.streaminfo().bits_per_sample as u16,
//!     sample_format: hound::SampleFormat::Int,
//! };
//!
//! let fname_wav = fname.with_extension("wav");
//! let opt_wav_writer = hound::WavWriter::create(fname_wav, spec);
//! let mut wav_writer = opt_wav_writer.expect("failed to create wav file");
//!
//! for opt_sample in reader.samples() {
//!     let sample = opt_sample.expect("failed to decode FLAC stream");
//!     wav_writer.write_sample(sample).expect("failed to write wav file");
//! }
//! # }
//! ```
//!
//! Retrieving the artist metadata:
//!
//! ```
//! # use claxon;
//! let reader = claxon::FlacReader::open("testsamples/pop.flac").unwrap();
//! for artist in reader.get_tag("ARTIST") {
//!     println!("{}", artist);
//! }
//! ```
//!
//! For more examples, see the [examples](https://github.com/ruuda/claxon/tree/master/examples)
//! directory in the crate.

#![warn(missing_docs)]

use std::fs;
use std::io;
use std::mem;
use std::path;
use error::fmt_err;
use frame::FrameReader;
use input::{BufferedReader, ReadBytes};
use metadata::{MetadataBlock, MetadataBlockReader, StreamInfo, VorbisComment};

mod crc;
mod error;
pub mod frame;
pub mod input;
pub mod metadata;
pub mod subframe;

pub use error::{Error, Result};
pub use frame::Block;

/// A FLAC decoder that can decode the stream from the underlying reader.
///
/// TODO: Add an example.
pub struct FlacReader<R: io::Read> {
    streaminfo: StreamInfo,
    vorbis_comment: Option<VorbisComment>,
    input: FlacReaderState<BufferedReader<R>>,
}

enum FlacReaderState<T> {
    /// When the reader is positioned at the beginning of a frame.
    Full(T),
    /// When the reader might not be positioned at the beginning of a frame.
    MetadataOnly(T),
}

/// Controls what metadata `FlacReader` reads when constructed.
///
/// The FLAC format contains a number of metadata blocks before the start of
/// audio data. Reading these is wasteful if the data is never used. The
/// `FlacReaderOptions` indicate which blocks to look for. As soon as all
/// desired blocks have been read, a `FlacReader` is returned without reading
/// remaining metadata blocks.
///
/// A few use cases:
///
/// * For a `FlacReader` that can read audio samples, set `read_until_samples`
///   to true.
/// * To read only the streaminfo, as quickly as possible, set both
///   `read_vorbis_comment` and `read_until_samples` to false.
///   The resulting reader cannot be used to read audio data.
/// * To read the streaminfo and tags, set `read_vorbis_comment` to true and
///   `read_until_samples` to false. The resulting reader cannot be used to
///   read audio data.
pub struct FlacReaderOptions {
    /// When true, read metadata blocks at least until a Vorbis comment block is found.
    ///
    /// When false, the `FlacReader` will be constructed without reading a
    /// Vorbis comment block, even if the stream contains one. Consequently,
    /// `FlacReader::tags()` and other tag-related methods will not return tag
    /// data.
    pub read_vorbis_comment: bool,

    /// When true, read all metadata blocks, up to the point where audio data starts.
    ///
    /// When false, the `FlacReader` is returned as soon as possible, but it
    /// will be unable to read audio samples.
    pub read_until_samples: bool,
}

/// An iterator that yields samples read from a `FlacReader`.
pub struct FlacSamples<R: ReadBytes> {
    frame_reader: FrameReader<R>,
    block: Block,
    sample: u32,
    channel: u32,

    /// If reading ever failed, this flag is set, so that the iterator knows not
    /// to return any new values.
    has_failed: bool,
}

// TODO: Add a `FlacIntoSamples`.

fn read_stream_header<R: ReadBytes>(input: &mut R) -> Result<()> {
    // A FLAC stream starts with a 32-bit header 'fLaC' (big endian).
    const FLAC_HEADER: u32 = 0x66_4c_61_43;

    // Some files start with ID3 tag data. The reference decoder supports this
    // for convenience. Claxon does not, but we can at least generate a helpful
    // error message if a file starts like this.
    const ID3_HEADER: u32 = 0x49_44_33_00;

    let header = try!(input.read_be_u32());
    if header != FLAC_HEADER {
        if (header & 0xff_ff_ff_00) == ID3_HEADER {
            fmt_err("stream starts with ID3 header rather than FLAC header")
        } else {
            fmt_err("invalid stream header")
        }
    } else {
        Ok(())
    }
}

impl<R: io::Read> FlacReader<R> {
    /// Attempts to create a reader that reads the FLAC format.
    ///
    /// The header and metadata blocks are read immediately. Audio frames
    /// will be read on demand.
    ///
    /// Claxon rejects files that claim to contain excessively large metadata
    /// blocks, to protect against denial of service attacks where a
    /// small damaged or malicous file could cause gigabytes of memory
    /// to be allocated. `Error::Unsupported` is returned in that case.
    pub fn new(reader: R) -> Result<FlacReader<R>> {
        let mut buf_reader = BufferedReader::new(reader);

        // A flac stream first of all starts with a stream header.
        try!(read_stream_header(&mut buf_reader));

        // Start a new scope, because the input reader must be available again
        // for the frame reader next.
        let (streaminfo, vorbis_comment) = {
            // Next are one or more metadata blocks. The flac specification
            // dictates that the streaminfo block is the first block. The metadata
            // block reader will yield at least one element, so the unwrap is safe.
            let mut metadata_iter = MetadataBlockReader::new(&mut buf_reader);
            let streaminfo_block = try!(metadata_iter.next().unwrap());
            let streaminfo = match streaminfo_block {
                MetadataBlock::StreamInfo(info) => info,
                _ => return fmt_err("streaminfo block missing"),
            };

            let mut vorbis_comment = None;

            // There might be more metadata blocks, read and store them.
            for block_result in metadata_iter {
                match try!(block_result) {
                    MetadataBlock::VorbisComment(vc) => {
                        // The Vorbis comment block need not be present, but
                        // when it is, it must be unique.
                        if vorbis_comment.is_some() {
                            return fmt_err("encountered second Vorbis comment block")
                        } else {
                            vorbis_comment = Some(vc);
                        }
                    }
                    MetadataBlock::StreamInfo(..) => {
                        return fmt_err("encountered second streaminfo block")
                    }
                    // Other blocks are currently not handled.
                    _block => {}
                }
            }

            (streaminfo, vorbis_comment)
        };

        // The flac reader will contain the reader that will read frames.
        let flac_reader = FlacReader {
            streaminfo: streaminfo,
            vorbis_comment: vorbis_comment,
            input: FlacReaderState::Full(buf_reader),
        };

        Ok(flac_reader)
    }

    /// Returns the streaminfo metadata.
    ///
    /// This contains information like the sample rate and number of channels.
    pub fn streaminfo(&self) -> StreamInfo {
        self.streaminfo
    }

    /// Returns the vendor string of the Vorbis comment block, if present.
    ///
    /// This string usually contains the name and version of the program that
    /// encoded the FLAC stream, such as `reference libFLAC 1.3.2 20170101`
    /// or `Lavf57.25.100`.
    pub fn vendor(&self) -> Option<&str> {
        self.vorbis_comment.as_ref().map(|vc| &vc.vendor[..])
    }

    /// Returns name-value pairs of Vorbis comments, such as `("ARTIST", "Queen")`.
    ///
    /// The name is supposed to be interpreted case-insensitively, and is
    /// guaranteed to consist of ASCII characters. Claxon does not normalize
    /// the casing of the name. Use `get_tag()` to do a case-insensitive lookup.
    ///
    /// Names need not be unique. For instance, multiple `ARTIST` comments might
    /// be present on a collaboration track.
    ///
    /// See https://www.xiph.org/vorbis/doc/v-comment.html for more details.
    pub fn tags<'a>(&'a self) -> metadata::Tags<'a> {
        match self.vorbis_comment.as_ref() {
            Some(vc) => metadata::Tags::new(&vc.comments[..]),
            None => metadata::Tags::new(&[]),
        }
    }

    /// Look up a Vorbis comment such as `ARTIST` in a case-insensitive way.
    ///
    /// Returns an iterator,  because tags may occur more than once. There could
    /// be multiple `ARTIST` tags on a collaboration track, for instance.
    ///
    /// Note that tag names are ASCII and never contain `'='`; trying to look up
    /// a non-ASCII tag will return no results. Furthermore, the Vorbis comment
    /// spec dictates that tag names should be handled case-insensitively, so
    /// this method performs a case-insensitive lookup.
    ///
    /// See also `tags()` for access to the raw tags.
    /// See https://www.xiph.org/vorbis/doc/v-comment.html for more details.
    pub fn get_tag<'a>(&'a self, tag_name: &'a str) -> metadata::GetTag<'a> {
        match self.vorbis_comment.as_ref() {
            Some(vc) => metadata::GetTag::new(&vc.comments[..], tag_name),
            None => metadata::GetTag::new(&[], tag_name),
        }
    }

    /// Returns an iterator that decodes a single frame on every iteration.
    /// TODO: It is not an iterator.
    ///
    /// This is a low-level primitive that gives you control over when decoding
    /// happens. The representation of the decoded audio is somewhat specific to
    /// the FLAC format. For a higher-level interface, see `samples()`.
    pub fn blocks<'r>(&'r mut self) -> FrameReader<&'r mut BufferedReader<R>> {
        match self.input {
            FlacReaderState::Full(ref mut inp) => FrameReader::new(inp),
            FlacReaderState::MetadataOnly(..) =>
                panic!("FlacReaderOptions::read_until_samples must be set \
                       to be able to use FlacReader::blocks()"),
        }
    }

    /// Returns an iterator over all samples.
    ///
    /// The channel data is is interleaved. The iterator is streaming. That is,
    /// if you call this method once, read a few samples, and call this method
    /// again, the second iterator will not start again from the beginning of
    /// the file. It will continue somewhere after where the first iterator
    /// stopped, and it might skip some samples. (This is because FLAC divides
    /// a stream into blocks, which have to be decoded entirely. If you drop the
    /// iterator, you lose the unread samples in that block.)
    ///
    /// This is a user-friendly interface that trades performance for ease of
    /// use. If performance is an issue, consider using `blocks()` instead.
    ///
    /// This is a high-level interface to the decoder. The cost of retrieving
    /// the next sample can vary significantly, as sometimes a new block has to
    /// be decoded. Additionally, there is a cost to every iteration returning a
    /// `Result`. When a block has been decoded, iterating the samples in that
    /// block can never fail, but a match on every sample is required
    /// nonetheless. For more control over when decoding happens, and less error
    /// handling overhead, use `blocks()`.
    pub fn samples<'r>(&'r mut self) -> FlacSamples<&'r mut BufferedReader<R>> {
        match self.input {
            FlacReaderState::Full(ref mut inp) => {
                FlacSamples {
                    frame_reader: frame::FrameReader::new(inp),
                    block: Block::empty(),
                    sample: 0,
                    channel: 0,
                    has_failed: false,
                }
            }
            FlacReaderState::MetadataOnly(..) => {
                panic!("FlacReaderOptions::read_until_samples must be set \
                       to be able to use FlacReader::samples()")
            }
        }
    }

    /// Destroys the FLAC reader and returns the underlying reader.
    ///
    /// Because the reader employs buffering internally, anything in the buffer
    /// will be lost.
    pub fn into_inner(self) -> R {
        match self.input {
            FlacReaderState::Full(inp) => inp.into_inner(),
            FlacReaderState::MetadataOnly(inp) => inp.into_inner(),
        }
    }
}

impl FlacReader<fs::File> {
    /// Attempts to create a reader that reads from the specified file.
    ///
    /// This is a convenience constructor that opens a `File`, and constructs a
    /// `FlacReader` from it. There is no need to wrap the file in a
    /// `BufReader`, as the `FlacReader` employs buffering already.
    pub fn open<P: AsRef<path::Path>>(filename: P) -> Result<FlacReader<fs::File>> {
        let file = try!(fs::File::open(filename));
        FlacReader::new(file)
    }
}

impl<R: ReadBytes> Iterator for FlacSamples<R> {
    type Item = Result<i32>;

    fn next(&mut self) -> Option<Result<i32>> {
        // If the previous read failed, end iteration.
        if self.has_failed {
            return None;
        }

        // Iterate the samples channel interleaved, so first increment the
        // channel.
        self.channel += 1;

        // If that was the last channel, increment the sample number.
        if self.channel >= self.block.channels() {
            self.channel = 0;
            self.sample += 1;

            // If that was the last sample in the block, decode the next block.
            if self.sample >= self.block.duration() {
                self.sample = 0;

                // Replace the current block with an empty one so that we may
                // reuse the current buffer to decode again.
                let current_block = mem::replace(&mut self.block, Block::empty());

                match self.frame_reader.read_next_or_eof(current_block.into_buffer()) {
                    Ok(Some(next_block)) => {
                        self.block = next_block;
                    }
                    Ok(None) => {
                        // The stream ended with EOF.
                        // TODO: If a number of samples was specified in the
                        // streaminfo metadata block, verify that we did not
                        // read more or less samples.
                        return None;
                    }
                    Err(error) => {
                        self.has_failed = true;
                        return Some(Err(error));
                    }
                }
            }
        }

        Some(Ok(self.block.sample(self.channel, self.sample)))
    }
}
<|endoftext|>


<commit_before>/*!
 * A basic library for finding primes, using the Sieve of Eratosthenes.
 *
 * To use, see `PrimeSet`, a class which handles the Sieve and has multiple methods for iterating
 * over primes.
*/

#[warn(non_camel_case_types)]
#[warn(non_snake_case)]
#[warn(unused_qualifications)]
#[warn(non_upper_case_globals)]
#[warn(missing_docs)]

#[cfg(test)]
extern crate test;

use std::ops::Index;
use std::slice;
use std::iter;
use std::num::Float;

#[cfg(test)]
use test::Bencher;

fn sqrt_floor(n : uint) -> uint {
    (n as f64).sqrt().floor() as uint
}

/// A prime generator, using the Sieve of Eratosthenes.
pub struct PrimeSet {
    lst : Vec<uint>
}

/// An iterator over generated primes. Created by PrimeSet::iter or
/// PrimeSet::generator
pub struct PrimeSetIter<'a> {
    p : &'a mut PrimeSet,
    n : uint,
    expand : bool
}

impl PrimeSet {
    /// A new prime penerator, primed with 2 and 3
    pub fn new() -> PrimeSet {
        PrimeSet{lst:vec!(2,3)}
    }

    /// Finds one more prime, and adds it to the list
    pub fn expand(&mut self) {
        let mut l = self.lst[self.lst.len()-1] + 2;
        let mut sql = sqrt_floor(l);
        let mut remainder = 0;
        loop {
            for &n in self.lst.iter() {
                remainder = l % n;
                if remainder == 0 || n > sql {
                    break;
                }
            };

            if remainder != 0 {
                self.lst.push(l);
                break;
            };

            l += 2;
            sql = sqrt_floor(l);
        }
    }

    /// Number of primes found so far
    pub fn len(&self) -> uint {
        self.lst.len()
    }

    /// Return all primes found so far as a slice
    pub fn list<'a>(&'a self) -> &'a [uint] {
        self.lst.as_slice()
    }

    /// Iterator over all primes not yet found
    pub fn generator<'a>(&'a mut self) -> PrimeSetIter<'a> {
        let myn = self.len();
        PrimeSetIter{p:self, n:myn, expand:true}
    }

    /// Iterator over all primes, starting with 2. If you don't care about the "state" of the
    /// PrimeSet, this is what you want!
    pub fn iter<'a>(&'a mut self) -> PrimeSetIter<'a> {
        PrimeSetIter{p:self, n:0, expand:true}
    }

    //~ pub fn iter_once(&'self mut self) -> PrimeSetIter<'self> {
        //~ PrimeSetIter{p:self, n:0, expand:false}
    //~ }

    /// Iterator over just the primes found so far
    pub fn iter_vec<'a>(&'a self) -> slice::Items<'a, uint> {
        self.lst.iter()
    }

    /// Find the next largest prime from a number
    /// Returns (idx, prime)
    /// Note that if n is prime, then the output will be (idx, n)
    pub fn find(&mut self, n: uint) -> (uint, uint) {
        while n > *(self.lst.last().unwrap_or(&0)){
            self.expand();
        }
        self.find_vec(n).unwrap()
    }

    /// Check if a number is prime
    pub fn is_prime(&mut self, n: uint) -> bool {
        if n <= 1 {return false;}
        if n == 2 {return true;} // otherwise we get 2 % 2 == 0!
        for m in self.iter() {
            if n % m == 0 {return false;};
            if m*m > n {return true;};
        }
        panic!("This iterator should not be empty.");
    }

    /// Find the next largest prime from a number, if it is within the already-found list
    /// Returns (idx, prime)
    /// Note that if n is prime, then the output will be (idx, n)
    pub fn find_vec(&self, n: uint) -> Option<(uint, uint)> {
        if n > *(self.lst.last().unwrap_or(&0)){ return None;}

        let mut base : uint = 0;
        let mut lim : uint = self.len();

        while lim != 0 {
            let ix = base + (lim >> 1);
            match self.lst[ix].cmp(&n) {
                Equal => return Some((ix, self.lst[ix])),
                Less => {
                    base = ix + 1;
                    lim -= 1;
                }
                Greater => ()
            }
            lim >>= 1;
        }
        return Some((base, self.lst[base]));
    }

    /// Get the nth prime, even if we haven't yet found it
    pub fn get(&mut self, index : &uint) -> &uint {
		for _ in range(0, (*index as int) + 1 - (self.lst.len() as int)){
			self.expand();
		}
        self.lst.index(index)
	}

	/// Get the prime factors of a number, starting from 2, including repeats
	pub fn prime_factors(&mut self, n: uint) -> Vec<uint> {
		if n == 1 {return Vec::new();}
		let mut curn = n;
		let mut m = ((curn as f64).sqrt()).ceil() as uint;
		let mut lst: Vec<uint> = Vec::new();
		for p in self.iter() {
			while curn % p == 0 {
				lst.push(p);
				curn /= p;
				if curn == 1 {return lst;}
				m = ((curn as f64).sqrt()).ceil() as uint;
			}

			if p > m {
				lst.push(p);
				return lst;
			}
		}
		panic!("This should be unreachable.");
	}
}

impl Index<uint, uint> for PrimeSet {
    fn index(&self, index: &uint) -> &uint {
        self.lst.index(index)
    }
}

impl<'a> Iterator<uint> for PrimeSetIter<'a> {
    fn next(&mut self) -> Option<uint> {
        while self.n >= self.p.len(){
            match self.expand {
                true => self.p.expand(),
                false => return None
            }
        }
        self.n += 1;

        let m = self.p.lst[self.n-1];

        Some(m)
    }
}

/// Find the first factor (other than 1) of a number
fn firstfac(x: uint) -> uint {
    let m = ((x as f64).sqrt()).ceil() as uint;
    if x % 2 == 0 { return 2; };
    for n in iter::range_step(3, m + 1, 2) {
        if x % n == 0 { return n; };
    }
    return x;
}

/// Find all prime factors of a number
pub fn factors(x: uint) -> Vec<uint> {
    let mut lst: Vec<uint> = Vec::new();
    let mut curn = x;
    loop  {
        let m = firstfac(curn);
        lst.push(m);
        if m == curn { break  } else { curn /= m };
    }
    return lst
}

/// Find all unique prime factors of a number
pub fn factors_uniq(x: uint) -> Vec<uint> {
    let mut lst: Vec<uint> = Vec::new();
    let mut curn = x;
    loop  {
        let m = firstfac(curn);
        lst.push(m);
        if curn == m { break ; }
        while curn % m == 0 { curn /= m; }
        if curn == 1 { break ; }
    }
    return lst
}

/// Test whether a number is prime. Checks every odd number up to sqrt(n).
pub fn is_prime(n : uint) -> bool {
    if n <= 1 {return false;}
    firstfac(n) == n
}

#[test]
fn test_iter(){
    let mut pset = PrimeSet::new();
    let first_few = [2u,3,5,7,11,13,17,19,23];
    for (m, &n) in pset.iter().zip(first_few.iter()) {
        assert_eq!(m, n);
        break;
    }
}

#[test]
fn test_primes(){
    let mut pset = PrimeSet::new();

    // note: some are repeated, because the pset list grows as it goes

    assert!(!pset.is_prime(1));
    assert!(!is_prime(1));
    assert!(pset.is_prime(2));
    assert!(is_prime(2));
    assert!(pset.is_prime(13));
    assert!(is_prime(13));
    assert!(!pset.is_prime(45));
    assert!(!is_prime(45));
    assert!(!pset.is_prime(13*13));
    assert!(!is_prime(13*13));
    assert!(pset.is_prime(13));
    assert!(pset.is_prime(7));
    assert!(is_prime(7));
    assert!(!pset.is_prime(9));
    assert!(!is_prime(9));
    assert!(pset.is_prime(5));
    assert!(is_prime(5));
}

#[test]
fn test_factors(){
	let mut pset = PrimeSet::new();

	assert_eq!(pset.prime_factors(1), vec!());
	assert_eq!(pset.prime_factors(2), vec!(2));
	assert_eq!(pset.prime_factors(3), vec!(3));
	assert_eq!(pset.prime_factors(4), vec!(2,2));
	assert_eq!(pset.prime_factors(5), vec!(5));
	assert_eq!(pset.prime_factors(6), vec!(2,3));
	assert_eq!(pset.prime_factors(9), vec!(3,3));

	pset = PrimeSet::new();
	assert_eq!(pset.prime_factors(12), vec!(2,2,3));
	assert_eq!(pset.prime_factors(100), vec!(2,2,5,5));
	assert_eq!(pset.prime_factors(121), vec!(11, 11));
}

#[bench]
fn bench_primes(b : &mut Bencher){
	b.iter(|| {
		let mut pset = PrimeSet::new();
		 let (_, _) = pset.find(1_000_000);
		 //~ let (idx, n) = pset.find(10000000);
		 //~ println!("Prime {}: {}", idx, n);
		 })
}
<commit_msg>added documentation<commit_after>/*!
A basic library for finding primes, using the Sieve of Eratosthenes. This library provides methods
for generating primes, testing whether a number is prime, and factorizing numbers. Most methods
generate primes lazily, so only enough primes will be generated for the given test, and primes are
cached for later use.

To use, see `PrimeSet`, a class which handles the Sieve and has multiple methods for iterating
over primes.

This also provides a few functions unconnected to `PrimeSet`, which will be faster for the first case,
but slower in the long term as they do not use any caching of primes.
*/

#[warn(non_camel_case_types)]
#[warn(non_snake_case)]
#[warn(unused_qualifications)]
#[warn(non_upper_case_globals)]
#[warn(missing_docs)]

#[cfg(test)]
extern crate test;

use std::ops::Index;
use std::slice;
use std::iter;
use std::num::Float;

#[cfg(test)]
use test::Bencher;

fn sqrt_floor(n : uint) -> uint {
    (n as f64).sqrt().floor() as uint
}

/** A prime generator, using the Sieve of Eratosthenes.

Create with `let mut pset = PrimeSet::new()`, and then use `pset.iter()` to iterate over all primes.
**/
pub struct PrimeSet {
    lst : Vec<uint>
}

/// An iterator over generated primes. Created by PrimeSet::iter or
/// PrimeSet::generator
pub struct PrimeSetIter<'a> {
    p : &'a mut PrimeSet,
    n : uint,
    expand : bool
}

impl PrimeSet {
    /// A new prime generator, primed with 2 and 3
    pub fn new() -> PrimeSet {
        PrimeSet{lst:vec!(2,3)}
    }

    /// Finds one more prime, and adds it to the list
    pub fn expand(&mut self) {
        let mut l = self.lst[self.lst.len()-1] + 2;
        let mut sql = sqrt_floor(l);
        let mut remainder = 0;
        loop {
            for &n in self.lst.iter() {
                remainder = l % n;
                if remainder == 0 || n > sql {
                    break;
                }
            };

            if remainder != 0 {
                self.lst.push(l);
                break;
            };

            l += 2;
            sql = sqrt_floor(l);
        }
    }

    /// Number of primes found so far
    pub fn len(&self) -> uint {
        self.lst.len()
    }

    /// Return all primes found so far as a slice
    pub fn list<'a>(&'a self) -> &'a [uint] {
        self.lst.as_slice()
    }

    /// Iterator over all primes not yet found
    pub fn generator<'a>(&'a mut self) -> PrimeSetIter<'a> {
        let myn = self.len();
        PrimeSetIter{p:self, n:myn, expand:true}
    }

    /// Iterator over all primes, starting with 2. If you don't care about the "state" of the
    /// PrimeSet, this is what you want!
    pub fn iter<'a>(&'a mut self) -> PrimeSetIter<'a> {
        PrimeSetIter{p:self, n:0, expand:true}
    }

    //~ pub fn iter_once(&'self mut self) -> PrimeSetIter<'self> {
        //~ PrimeSetIter{p:self, n:0, expand:false}
    //~ }

    /// Iterator over just the primes found so far
    pub fn iter_vec<'a>(&'a self) -> slice::Items<'a, uint> {
        self.lst.iter()
    }

    /// Find the next largest prime from a number
    /// Returns (idx, prime)
    /// Note that if n is prime, then the output will be (idx, n)
    pub fn find(&mut self, n: uint) -> (uint, uint) {
        while n > *(self.lst.last().unwrap_or(&0)){
            self.expand();
        }
        self.find_vec(n).unwrap()
    }

    /// Check if a number is prime
    /// Note that this only requires primes up to n.sqrt() to be generated, and will generate
    /// them as necessary on its own.
    pub fn is_prime(&mut self, n: uint) -> bool {
        if n <= 1 {return false;}
        if n == 2 {return true;} // otherwise we get 2 % 2 == 0!
        for m in self.iter() {
            if n % m == 0 {return false;};
            if m*m > n {return true;};
        }
        panic!("This iterator should not be empty.");
    }

    /// Find the next largest prime from a number, if it is within the already-found list
    /// Returns (idx, prime)
    /// Note that if n is prime, then the output will be (idx, n)
    pub fn find_vec(&self, n: uint) -> Option<(uint, uint)> {
        if n > *(self.lst.last().unwrap_or(&0)){ return None;}

        let mut base : uint = 0;
        let mut lim : uint = self.len();

        while lim != 0 {
            let ix = base + (lim >> 1);
            match self.lst[ix].cmp(&n) {
                Equal => return Some((ix, self.lst[ix])),
                Less => {
                    base = ix + 1;
                    lim -= 1;
                }
                Greater => ()
            }
            lim >>= 1;
        }
        return Some((base, self.lst[base]));
    }

    /// Get the nth prime, even if we haven't yet found it
    pub fn get(&mut self, index : &uint) -> &uint {
		for _ in range(0, (*index as int) + 1 - (self.lst.len() as int)){
			self.expand();
		}
        self.lst.index(index)
	}

	/// Get the prime factors of a number, starting from 2, including repeats
	pub fn prime_factors(&mut self, n: uint) -> Vec<uint> {
		if n == 1 {return Vec::new();}
		let mut curn = n;
		let mut m = ((curn as f64).sqrt()).ceil() as uint;
		let mut lst: Vec<uint> = Vec::new();
		for p in self.iter() {
			while curn % p == 0 {
				lst.push(p);
				curn /= p;
				if curn == 1 {return lst;}
				m = ((curn as f64).sqrt()).ceil() as uint;
			}

			if p > m {
				lst.push(p);
				return lst;
			}
		}
		panic!("This should be unreachable.");
	}
}

impl Index<uint, uint> for PrimeSet {
    fn index(&self, index: &uint) -> &uint {
        self.lst.index(index)
    }
}

impl<'a> Iterator<uint> for PrimeSetIter<'a> {
    fn next(&mut self) -> Option<uint> {
        while self.n >= self.p.len(){
            match self.expand {
                true => self.p.expand(),
                false => return None
            }
        }
        self.n += 1;

        let m = self.p.lst[self.n-1];

        Some(m)
    }
}

/// Find the first factor (other than 1) of a number
fn firstfac(x: uint) -> uint {
    let m = ((x as f64).sqrt()).ceil() as uint;
    if x % 2 == 0 { return 2; };
    for n in iter::range_step(3, m + 1, 2) {
        if x % n == 0 { return n; };
    }
    return x;
}

/// Find all prime factors of a number
/// Does not use a PrimeSet, but simply counts upwards
pub fn factors(x: uint) -> Vec<uint> {
    if x <= 1 {return vec!()};
    let mut lst: Vec<uint> = Vec::new();
    let mut curn = x;
    loop  {
        let m = firstfac(curn);
        lst.push(m);
        if m == curn { break  } else { curn /= m };
    }
    return lst
}

/// Find all unique prime factors of a number
pub fn factors_uniq(x: uint) -> Vec<uint> {
    let mut lst: Vec<uint> = Vec::new();
    let mut curn = x;
    loop  {
        let m = firstfac(curn);
        lst.push(m);
        if curn == m { break ; }
        while curn % m == 0 { curn /= m; }
        if curn == 1 { break ; }
    }
    return lst
}

/// Test whether a number is prime. Checks every odd number up to sqrt(n).
pub fn is_prime(n : uint) -> bool {
    if n <= 1 {return false;}
    firstfac(n) == n
}

#[test]
fn test_iter(){
    let mut pset = PrimeSet::new();
    let first_few = [2u,3,5,7,11,13,17,19,23];
    for (m, &n) in pset.iter().zip(first_few.iter()) {
        assert_eq!(m, n);
    }
}

#[test]
fn test_primes(){
    let mut pset = PrimeSet::new();

    // note: some are repeated, because the pset list grows as it goes

    assert!(!pset.is_prime(1));
    assert!(!is_prime(1));
    assert!(pset.is_prime(2));
    assert!(is_prime(2));
    assert!(pset.is_prime(13));
    assert!(is_prime(13));
    assert!(!pset.is_prime(45));
    assert!(!is_prime(45));
    assert!(!pset.is_prime(13*13));
    assert!(!is_prime(13*13));
    assert!(pset.is_prime(13));
    assert!(pset.is_prime(7));
    assert!(is_prime(7));
    assert!(!pset.is_prime(9));
    assert!(!is_prime(9));
    assert!(pset.is_prime(5));
    assert!(is_prime(5));
}

#[test]
fn test_factors(){
	let mut pset = PrimeSet::new();

	let ns = [  (1, vec!()),
                (2, vec!(2)),
                (3, vec!(3)),
                (4, vec!(2,2)),
                (5, vec!(5)),
                (6, vec!(2,3)),
	            (9, vec!(3,3)),
	            (12, vec!(2,2,3)),
	            (121, vec!(11,11)),
	            (144, vec!(2,2,2,2,3,3)),
	            (10_000_000, vec!(2,2,2,2,2,2,2,5,5,5,5,5,5,5)),
                (100, vec!(2,2,5,5)),
                (121, vec!(11, 11)),
                ];

    for &(n, ref v) in ns.iter(){
        assert_eq!(pset.prime_factors(n), *v);
        assert_eq!(factors(n), *v);
    }

	pset = PrimeSet::new();
	assert_eq!(pset.prime_factors(12), vec!(2,2,3));

}

#[bench]
fn bench_primes(b : &mut Bencher){
	b.iter(|| {
		let mut pset = PrimeSet::new();
		 let (_, _) = pset.find(1_000_000);
		 //~ let (idx, n) = pset.find(1_000_000);
		 //~ println!("Prime {}: {}", idx, n);
		 })
}
<|endoftext|>


<commit_before>#![crate_name = "ssdb"]
//[comment = "SSDB client"]
//#![crate_type = "dylib"]
//#![crate_type = "rlib"]
extern crate bufstream;

mod ssdb{
    use std::net::{TcpStream};
    use std::io::{BufRead,Read,Write,Result,Error,ErrorKind};
    use std::string::String;
    use std::convert::AsRef;
    use bufstream::BufStream;

    pub enum SSDBResult {
        Nil,
        Int(String,i64),
        Data(String,Vec<u8>)
        //List(&[Result]),
        //Error(&str),
        //Status(&str)
    }

    pub struct Client<S:Read+Write>{
        ip:String,
        port:u16,
        stream:Option<BufStream<S>>
    }

    fn invalid_input(desc: &'static str, detail: &str) -> Error {
        Error::new(ErrorKind::InvalidInput, desc)
        // {kind: InvalidInput, desc: desc, detail: Some(String::from_str(detail))}
    }

    fn read_byte<S: Read+Write>(stream:&mut BufStream<S>) ->Result<u8>{
        let buf: &mut[u8;1]=&mut [0];
        let nread = try!(stream.read(buf));
        if nread < 1 {
            Err(invalid_input("Read failed!","Could not read enough bytes"))
        } else {
            Ok(buf[0])
        }
    }

    fn read_item<S: Read+Write>(stream: &mut BufStream<S>) -> Result<String> {
        let l = &mut String::new();
        stream.read_line(l);
        let len=match l.trim().as_ref(){
            ""=>{return Ok("".to_string());0},
            _=>{l.trim().parse::<usize>().unwrap()}
        };
        let mut s = vec![];
        s.reserve(len);
        for _ in 0..len{
            s.push(try!(read_byte(stream)));
        }
        read_byte(stream);

        Ok(String::from_utf8(s).unwrap())
    }

    pub fn read_status<S: Read+Write>(stream: &mut BufStream<S>)->Result<()> {
        let resp=read_item(stream).unwrap();
        match resp.as_ref() {
            "ok"=>Ok(()),
            _=>Err(invalid_input("Operation failed!",resp.as_ref()))
        }
        //if resp.as_ref()!="ok" {
        //    return Err(invalid_input("Operation failed!",resp.as_ref()));
        //}
    }

    impl Client<TcpStream>{
        pub fn send(&mut self, data:Vec<&str>)->Result<Vec<String>>{
            let mut v=vec![];
            match self.stream{
                Some(ref mut stream) => {
                    for item in data.iter(){
                        writeln!(stream,"{}",item.len());
                        writeln!(stream,"{}",item);
                    }
                    writeln!(stream,"");
                    //println!("data:{}",String::from_utf8(bufwriter.unwrap()).unwrap().as_slice());
                    //stream.write(bufwriter.unwrap().as_slice());
                    try!(stream.flush());
                    //try!(read_status(stream));//read_item(stream).unwrap();
                    loop{
                        let s=try!(read_item(stream));//try!(stream.read_line());
                        if s=="" { break; }
                        v.push(s);
                    }
                },
                _ =>{}
            }
            Ok(v)
        }

        pub fn send_req(&mut self, cmd:&str, param:Vec<&str>)->Vec<String>{
            match self.stream{
                Some(ref mut stream) =>{
                    writeln!(stream, "{}", cmd.len());
                    writeln!(stream, "{}", cmd);
                },
                None => {}
            }
            self.send(param).unwrap()
        }

        pub fn parse(&mut self, cmd:&str, resp:Vec<String>)->SSDBResult{
            match cmd {
                "getbit" | "setbit" | "countbit" | "strlen"
                    | "set" | "setx" | "setnx" | "zset" | "hset"
                    | "qpush" | "qpush_front" | "qpush_back"
                    | "del" | "zdel" | "hdel" | "hsize" | "zsize" | "qsize"
                    | "hclear" | "zclear" | "qclear"
                    | "multi_set" | "multi_del" | "multi_hset" | "multi_hdel" | "multi_zset" | "multi_zdel"
                    | "incr" | "decr" | "zincr" | "zdecr" | "hincr" | "hdecr"
                    | "zget" | "zrank" | "zrrank" | "zcount" | "zsum" | "zremrangebyrank" | "zremrangebyscore"
                    =>SSDBResult::Int(resp[0].clone(), resp[1].parse::<i64>().unwrap()),
                "get"=>SSDBResult::Data(resp[0].clone(), resp[1].clone().into_bytes()),
                _ =>SSDBResult::Nil
            }
        }

        pub fn get(&mut self, key:&str)->SSDBResult{
            let resp=self.send_req("get",vec![key]);
            self.parse("get",resp)
        }

        pub fn set(&mut self, key:&str, value:&str)->SSDBResult{
            let resp=self.send_req("set",vec![key, value]);
            self.parse("set",resp)
        }

        pub fn new(ip:&str,port:i32)->Client<TcpStream>{
            Client{ip:ip.to_string(),
                port:port as u16,
                stream:None
            }
        }

        pub fn connect(&mut self) {
            if !self.stream.is_none(){
                panic!("can not call connect() twice!");
            }

            let stream=TcpStream::connect((self.ip.as_ref(),self.port)).unwrap();
            self.stream=Some(BufStream::new(stream));
            println!("connected!");
        }

        pub fn close(&mut self){
            match self.stream {
                Some(ref mut stream) =>{ drop(stream)},
                None =>{}
            }
        }
    }
}

#[test]
fn it_works() {

    //Require ssdb daemon running

    let mut foo = ssdb::Client::new("127.0.0.1",8888);
    foo.connect();

    foo.set("测试","foo");
    let resp=match foo.get("测试"){ssdb::SSDBResult::Data(_,resp)=>resp,_=>vec![]};
    //println!("{:?}",String::from_utf8(resp).unwrap());

    assert!(String::from_utf8(resp).unwrap()=="foo");

    foo.close();
}
<commit_msg>fix compiling issue under rust 1.10<commit_after>#![crate_name = "ssdb"]
//[comment = "SSDB client"]
//#![crate_type = "dylib"]
//#![crate_type = "rlib"]
extern crate bufstream;

mod ssdb{
    use std::net::{TcpStream};
    use std::io::{BufRead,Read,Write,Result,Error,ErrorKind};
    use std::string::String;
    use std::convert::AsRef;
    use bufstream::BufStream;

    pub enum SSDBResult {
        Nil,
        Int(String,i64),
        Data(String,Vec<u8>)
    }

    pub struct Client<S:Read+Write>{
        ip:String,
        port:u16,
        stream:Option<BufStream<S>>
    }

    fn invalid_input(desc: &'static str, detail: &str) -> Error {
        Error::new(ErrorKind::InvalidInput, desc)
    }

    fn read_byte<S: Read+Write>(stream:&mut BufStream<S>) ->Result<u8>{
        let buf: &mut[u8;1]=&mut [0];
        let nread = try!(stream.read(buf));
        if nread < 1 {
            Err(invalid_input("Read failed!","Could not read enough bytes"))
        } else {
            Ok(buf[0])
        }
    }

    fn read_item<S: Read+Write>(stream: &mut BufStream<S>) -> Result<String> {
        let l = &mut String::new();
        stream.read_line(l);
        let len=match l.trim().as_ref(){
            ""=>{return Ok("".to_string())},
            _=>{l.trim().parse::<usize>().unwrap_or(0)}
        };
        let mut s = vec![];
        s.reserve(len);
        for _ in 0..len{
            s.push(try!(read_byte(stream)));
        }
        read_byte(stream);

        Ok(String::from_utf8(s).unwrap_or(String::new()))
    }

    pub fn read_status<S: Read+Write>(stream: &mut BufStream<S>)->Result<()> {
        match read_item(stream).unwrap_or("".to_owned()).as_ref(){
            "ok"=>Ok(()),
            resp@_=>Err(invalid_input("Operation failed!", resp.as_ref()))
        }
    }

    impl Client<TcpStream>{
        pub fn send(&mut self, data:Vec<&str>)->Result<Vec<String>>{
            let mut v=vec![];
            match self.stream{
                Some(ref mut stream) => {
                    for item in data.iter(){
                        writeln!(stream,"{}",item.len());
                        writeln!(stream,"{}",item);
                    }
                    writeln!(stream,"");
                    try!(stream.flush());
                    loop{
                        let s=try!(read_item(stream));
                        if s=="" { break; }
                        v.push(s);
                    }
                },
                _ =>{panic!("Connect not open!");}
            }
            Ok(v)
        }

        pub fn send_req(&mut self, cmd:&str, param:Vec<&str>)->Vec<String>{

            match self.stream{
                Some(ref mut stream) =>{
                    writeln!(stream, "{}", cmd.len());
                    writeln!(stream, "{}", cmd);
                },
                _ => {panic!("Connect not open!");}
            }
            self.send(param).unwrap_or(vec![])
        }

        pub fn parse(&mut self, cmd:&str, resp:Vec<String>)->SSDBResult{
            match cmd {
                "getbit" | "setbit" | "countbit" | "strlen"
                    | "set" | "setx" | "setnx" | "zset" | "hset"
                    | "qpush" | "qpush_front" | "qpush_back"
                    | "del" | "zdel" | "hdel" | "hsize" | "zsize" | "qsize"
                    | "hclear" | "zclear" | "qclear"
                    | "multi_set" | "multi_del" | "multi_hset" | "multi_hdel" | "multi_zset" | "multi_zdel"
                    | "incr" | "decr" | "zincr" | "zdecr" | "hincr" | "hdecr"
                    | "zget" | "zrank" | "zrrank" | "zcount" | "zsum" | "zremrangebyrank" | "zremrangebyscore"
                    =>SSDBResult::Int(resp[0].clone(), resp[1].parse::<i64>().unwrap_or(0)),
                "get"=>SSDBResult::Data(resp[0].clone(), resp[1].clone().into_bytes()),
                _ =>SSDBResult::Nil
            }
        }

        pub fn get(&mut self, key:&str)->SSDBResult{
            let resp=self.send_req("get",vec![key]);
            self.parse("get",resp)
        }

        pub fn set(&mut self, key:&str, value:&str)->SSDBResult{
            let resp=self.send_req("set",vec![key, value]);
            self.parse("set",resp)
        }

        pub fn new(ip:&str,port:i32)->Client<TcpStream>{
            Client{
                ip:ip.to_string(),
                port:port as u16,
                stream:None,
            }
        }

        pub fn connect(&mut self){
            if !self.stream.is_none(){
                panic!("can not call connect() twice!");
            }

            let stream=TcpStream::connect((self.ip.as_ref(),self.port)).unwrap();
            self.stream=Some(BufStream::new(stream));
            println!("connected!");
        }

        pub fn close(&mut self){
            match self.stream {
                Some(ref mut stream) =>{ stream.flush(); },
                _ =>{}
            }
            self.stream=None;
        }

    }
}

#[test]
fn it_works() {
    use std::thread;

    let mut handles=vec![];

    //Require ssdb daemon running
    for i in 0..3{
        let h=thread::spawn(move ||{
            let key="测试".to_owned()+i.to_string().as_ref();
            let mut foo=ssdb::Client::new("127.0.0.1",8888);
            foo.connect();
            foo.set(key.as_ref(),"foo");
            let resp=match foo.get(key.as_ref()){ssdb::SSDBResult::Data(_,resp)=>resp,_=>vec![]};
            assert!(String::from_utf8(resp).unwrap()=="foo");
            foo.close();
        });
        handles.push(h);
    }
    for h in handles{
        h.join();
    }
    //println!("{:?}",String::from_utf8(resp).unwrap());
    //foo.set("测试","foo");
}
<|endoftext|>


<commit_before>extern crate emailaddress;
extern crate iron;
extern crate openssl;
extern crate rand;
extern crate redis;
extern crate rustc_serialize;
extern crate serde_json;
extern crate time;
extern crate url;
extern crate urlencoded;

pub mod email;
pub mod oidc;

use emailaddress::EmailAddress;
use iron::headers::ContentType;
use iron::middleware::Handler;
use iron::modifiers;
use iron::prelude::*;
use iron::status;
use openssl::bn::BigNum;
use openssl::crypto::hash;
use openssl::crypto::pkey::PKey;
use serde_json::builder::ObjectBuilder;
use serde_json::de::from_reader;
use serde_json::value::Value;
use rand::{OsRng, Rng};
use redis::Client;
use rustc_serialize::base64::{self, ToBase64};
use std::collections::BTreeMap;
use std::fs::File;
use std::io::{BufReader, Write};
use std::iter::Iterator;
use time::now_utc;
use urlencoded::{UrlEncodedBody, UrlEncodedQuery};

/// Helper function for returning an Iron response with JSON data.
///
/// Serializes the argument value to JSON and returns a HTTP 200 response
/// code with the serialized JSON as the body.
fn json_response(obj: &Value) -> IronResult<Response> {
    let content = serde_json::to_string(&obj).unwrap();
    let mut rsp = Response::with((status::Ok, content));
    rsp.headers.set(ContentType::json());
    Ok(rsp)
}

/// Configuration data for a "famous" identity provider.
///
/// Used as part of the `AppConfig` struct.
#[derive(Clone)]
struct ProviderConfig {
    /// URL pointing to the OpenID configuration document.
    discovery: String,
    /// Client ID issued for this daemon instance by the provider.
    client_id: String,
    /// Secret issued for this daemon instance by the provider.
    secret: String,
    /// Issuer origin as used in identity tokens issued by the provider.
    /// Used to check that the issuer in the token matches expectations.
    issuer: String,
}

/// Configuration data for this daemon instance.
#[derive(Clone)]
pub struct AppConfig {
    /// Version of the daemon (used in the `WelcomeHandler`).
    version: String,
    /// Base URL where this daemon can be found. This is used to construct
    /// callback URLs.
    base_url: String,
    /// Private key used to sign identity tokens.
    priv_key: PKey,
    /// Redis database connector.
    store: Client,
    /// Duration in seconds for expiry of data stored in Redis.
    expire_keys: usize,
    /// Email sender (email address, then human-readable name).
    sender: (String, String),
    /// Duration in seconds for expiration of identity tokens.
    token_validity: i64,
    /// A map of "famous" identity providers. Each key is an email domain,
    /// the value holds the configuration required to act as an application
    /// doing OpenID Connect against the provider.
    providers: BTreeMap<String, ProviderConfig>,
}

/// Implementation with single method to read configuration from JSON.
impl AppConfig {
    pub fn from_json_file(file_name: &str) -> AppConfig {

        let config_file = File::open(file_name).unwrap();
        let config_value: Value = from_reader(BufReader::new(config_file)).unwrap();
        let config = config_value.as_object().unwrap();

        // The `private_key_file` key in the JSON object should hold a file
        // name pointing to a PEM-encoded private RSA key.
        let key_file_name = config["private_key_file"].as_string().unwrap();
        let priv_key_file = File::open(key_file_name).unwrap();
        let mut reader = BufReader::new(priv_key_file);
        let sender = config["sender"].as_object().unwrap();
        AppConfig {
            // Use the crate's version as defined in Cargo.toml.
            version: env!("CARGO_PKG_VERSION").to_string(),
            base_url: config["base_url"].as_string().unwrap().to_string(),
            priv_key: PKey::private_key_from_pem(&mut reader).unwrap(),
            store: Client::open(config["redis"].as_string().unwrap()).unwrap(),
            expire_keys: config["expire_keys"].as_u64().unwrap() as usize,
            sender: (
                sender["address"].as_string().unwrap().to_string(),
                sender["name"].as_string().unwrap().to_string(),
            ),
            token_validity: config["token_validity"].as_i64().unwrap(),
            providers: config["providers"].as_object().unwrap().iter()
                .map(|(host, params)| {
                    let pobj = params.as_object().unwrap();
                    (host.clone(), ProviderConfig {
                        discovery: pobj["discovery"].as_string().unwrap().to_string(),
                        client_id: pobj["client_id"].as_string().unwrap().to_string(),
                        secret: pobj["secret"].as_string().unwrap().to_string(),
                        issuer: pobj["issuer"].as_string().unwrap().to_string(),
                    })
                })
                .collect::<BTreeMap<String, ProviderConfig>>(),
        }

    }
}

/// Iron handler for the root path, returns human-friendly message.
///
/// This is not actually used in the protocol.
pub struct WelcomeHandler { pub app: AppConfig }
impl Handler for WelcomeHandler {
    fn handle(&self, _: &mut Request) -> IronResult<Response> {
        json_response(&ObjectBuilder::new()
            .insert("ladaemon", "Welcome")
            .insert("version", &self.app.version)
            .unwrap())
    }
}

/// Iron handler to return the OpenID Discovery document.
///
/// Most of this is hard-coded for now, although the URLs are constructed by
/// using the base URL as configured in the `base_url` configuration value.
pub struct OIDConfigHandler { pub app: AppConfig }
impl Handler for OIDConfigHandler {
    fn handle(&self, _: &mut Request) -> IronResult<Response> {
        json_response(&ObjectBuilder::new()
            .insert("issuer", &self.app.base_url)
            .insert("authorization_endpoint",
                    format!("{}/auth", self.app.base_url))
            .insert("jwks_uri", format!("{}/keys.json", self.app.base_url))
            .insert("scopes_supported", vec!["openid", "email"])
            .insert("claims_supported",
                    vec!["aud", "email", "email_verified", "exp", "iat", "iss", "sub"])
            .insert("response_types_supported", vec!["id_token"])
            .insert("response_modes_supported", vec!["form_post"])
            .insert("grant_types_supported", vec!["implicit"])
            .insert("subject_types_supported", vec!["public"])
            .insert("id_token_signing_alg_values_supported", vec!["RS256"])
            .unwrap())
    }
}

/// Helper function to encode a `BigNum` as URL-safe base64-encoded bytes.
///
/// This is used for the public RSA key components returned by the
/// `KeysHandler`.
fn json_big_num(n: &BigNum) -> String {
    n.to_vec().to_base64(base64::URL_SAFE)
}

/// Helper function to build a session ID for a login attempt.
///
/// Put the email address, the client ID (RP origin) and some randomness into
/// a SHA256 hash, and encode it with URL-safe bas64 encoding. This is used
/// as the key in Redis, as well as the state for OAuth authentication.
fn session_id(email: &EmailAddress, client_id: &str) -> String {
    let mut rng = OsRng::new().unwrap();
    let mut bytes_iter = rng.gen_iter();
    let rand_bytes: Vec<u8> = (0..16).map(|_| bytes_iter.next().unwrap()).collect();

    let mut hasher = hash::Hasher::new(hash::Type::SHA256);
    hasher.write(email.to_string().as_bytes()).unwrap();
    hasher.write(client_id.as_bytes()).unwrap();
    hasher.write(&rand_bytes).unwrap();
    hasher.finish().to_base64(base64::URL_SAFE)
}

/// Iron handler for the JSON Web Key Set document.
///
/// Currently only supports a single RSA key (as configured), which is
/// published with the `"base"` key ID, scoped to signing usage. Relying
/// Parties will need to fetch this data to be able to verify identity tokens
/// issued by this daemon instance.
pub struct KeysHandler { pub app: AppConfig }
impl Handler for KeysHandler {
    fn handle(&self, _: &mut Request) -> IronResult<Response> {
        let rsa = self.app.priv_key.get_rsa();
        json_response(&ObjectBuilder::new()
            .insert_array("keys", |builder| {
                builder.push_object(|builder| {
                    builder.insert("kty", "RSA")
                        .insert("alg", "RS256")
                        .insert("use", "sig")
                        .insert("kid", "base")
                        .insert("n", json_big_num(&rsa.n().unwrap()))
                        .insert("e", json_big_num(&rsa.e().unwrap()))
                })
            })
            .unwrap())
    }
}

/// Iron handler for authentication requests from the RP.
///
/// Calls the `oidc::request()` function if the provided email address's
/// domain matches one of the configured famous providers. Otherwise, calls the
/// `email::request()` function to allow authentication through the email loop.
pub struct AuthHandler { pub app: AppConfig }
impl Handler for AuthHandler {
    fn handle(&self, req: &mut Request) -> IronResult<Response> {
        let params = req.get_ref::<UrlEncodedBody>().unwrap();
        let email_addr = EmailAddress::new(&params.get("login_hint").unwrap()[0]).unwrap();
        if self.app.providers.contains_key(&email_addr.domain) {

            // OIDC authentication. Using 302 Found for redirection here. Note
            // that, per RFC 7231, a user agent MAY change the request method
            // from POST to GET for the subsequent request.
            let auth_url = oidc::request(&self.app, params);
            Ok(Response::with((status::Found, modifiers::Redirect(auth_url))))

        } else {

            // Email loop authentication. For now, returns a JSON response;
            // empty if successful, otherwise contains an error.
            let obj = email::request(&self.app, params);
            json_response(&obj)

        }
    }
}

/// Helper method to create a JWT for a given email address and origin.
///
/// Builds the JSON payload and header, encoding with (URL-safe)
/// base64-encoding, then hashing and signing with the provided private key.
/// Returns the full JWT.
fn create_jwt(app: &AppConfig, email: &str, origin: &str) -> String {

    let now = now_utc().to_timespec().sec;
    let payload = serde_json::to_string(
        &ObjectBuilder::new()
            .insert("aud", origin)
            .insert("email", email)
            .insert("email_verified", email)
            .insert("exp", now + app.token_validity)
            .insert("iat", now)
            .insert("iss", &app.base_url)
            .insert("sub", email)
            .unwrap()
        ).unwrap();
    let header = serde_json::to_string(
        &ObjectBuilder::new()
            .insert("kid", "base")
            .insert("alg", "RS256")
            .unwrap()
        ).unwrap();

    let mut input = Vec::<u8>::new();
    input.extend(header.as_bytes().to_base64(base64::URL_SAFE).into_bytes());
    input.push(b'.');
    input.extend(payload.as_bytes().to_base64(base64::URL_SAFE).into_bytes());
    let sha256 = hash::hash(hash::Type::SHA256, &input);
    let sig = app.priv_key.sign(&sha256);
    input.push(b'.');
    input.extend(sig.to_base64(base64::URL_SAFE).into_bytes());
    String::from_utf8(input).unwrap()

}

/// HTML template used to have the user agent POST the identity token built
/// by the daemon instance to the RP's `redirect_uri`.
const FORWARD_TEMPLATE: &'static str = r#"<!DOCTYPE html>
<html>
  <head>
    <title>Let's Auth</title>
    <script>
      document.addEventListener('DOMContentLoaded', function() {
        document.getElementById('form').submit();
      });
    </script>
  </head>
  <body>
    <form id="form" action="{{ return_url }}" method="post">
      <input type="hidden" name="id_token" value="{{ jwt }}">
    </form>
  </body>
</html>"#;

/// Helper function for sending an identity token to the Relying Party.
///
/// Builds the JWT header and payload JSON data and signs it with the
/// configured private RSA key. Then uses `FORWARD_TEMPLATE` to embed the token
/// in a form that's POSTed to the RP's `redirect_uri` as soon as the page
/// is loaded.
fn send_jwt_response(jwt: &str, redirect: &str) -> IronResult<Response> {
    let html = FORWARD_TEMPLATE.replace("{{ return_url }}", redirect)
        .replace("{{ jwt }}", jwt);
    let mut rsp = Response::with((status::Ok, html));
    rsp.headers.set(ContentType::html());
    Ok(rsp)
}

/// Iron handler for one-time pad email loop confirmation.
///
/// Retrieves the session based session ID and the expected one-time pad. If
/// an error occurs, returns a JSON object with an error message. Otherwise,
/// send an identity token to the RP. TODO: the former is obviously wrong.
pub struct ConfirmHandler { pub app: AppConfig }
impl Handler for ConfirmHandler {
    fn handle(&self, req: &mut Request) -> IronResult<Response> {
        let params = req.get_ref::<UrlEncodedQuery>().unwrap();
        let session_id = &params.get("session").unwrap()[0];
        let code = &params.get("code").unwrap()[0];
        let res = email::verify(&self.app, session_id, code);
        if res.is_err() {
            json_response(&ObjectBuilder::new()
                          .insert("error", res.unwrap_err())
                          .unwrap())
        } else {
            let (jwt, redirect) = res.unwrap();
            send_jwt_response(&jwt, &redirect)
        }
    }
}

/// Iron handler for OAuth callbacks
///
/// After the user allows or denies the Authentication Request with the famous
/// identity provider, they will be redirected back to the callback handler.
/// Verify the callback data and return the resulting token or error.
pub struct CallbackHandler { pub app: AppConfig }
impl Handler for CallbackHandler {
    fn handle(&self, req: &mut Request) -> IronResult<Response> {
        let params = req.get_ref::<UrlEncodedQuery>().unwrap();
        let session = &params.get("state").unwrap()[0];
        let code = &params.get("code").unwrap()[0];
        let res = oidc::verify(&self.app, session, code);
        if res.is_err() {
            json_response(&ObjectBuilder::new()
                          .insert("error", res.unwrap_err())
                          .unwrap())
        } else {
            let (jwt, redirect) = res.unwrap();
            send_jwt_response(&jwt, &redirect)
        }
    }
}
<commit_msg>Unify handling of authentication results<commit_after>extern crate emailaddress;
extern crate iron;
extern crate openssl;
extern crate rand;
extern crate redis;
extern crate rustc_serialize;
extern crate serde_json;
extern crate time;
extern crate url;
extern crate urlencoded;

pub mod email;
pub mod oidc;

use emailaddress::EmailAddress;
use iron::headers::ContentType;
use iron::middleware::Handler;
use iron::modifiers;
use iron::prelude::*;
use iron::status;
use openssl::bn::BigNum;
use openssl::crypto::hash;
use openssl::crypto::pkey::PKey;
use serde_json::builder::ObjectBuilder;
use serde_json::de::from_reader;
use serde_json::value::Value;
use rand::{OsRng, Rng};
use redis::Client;
use rustc_serialize::base64::{self, ToBase64};
use std::collections::BTreeMap;
use std::fs::File;
use std::io::{BufReader, Write};
use std::iter::Iterator;
use time::now_utc;
use urlencoded::{UrlEncodedBody, UrlEncodedQuery};

/// Helper function for returning an Iron response with JSON data.
///
/// Serializes the argument value to JSON and returns a HTTP 200 response
/// code with the serialized JSON as the body.
fn json_response(obj: &Value) -> IronResult<Response> {
    let content = serde_json::to_string(&obj).unwrap();
    let mut rsp = Response::with((status::Ok, content));
    rsp.headers.set(ContentType::json());
    Ok(rsp)
}

/// Configuration data for a "famous" identity provider.
///
/// Used as part of the `AppConfig` struct.
#[derive(Clone)]
struct ProviderConfig {
    /// URL pointing to the OpenID configuration document.
    discovery: String,
    /// Client ID issued for this daemon instance by the provider.
    client_id: String,
    /// Secret issued for this daemon instance by the provider.
    secret: String,
    /// Issuer origin as used in identity tokens issued by the provider.
    /// Used to check that the issuer in the token matches expectations.
    issuer: String,
}

/// Configuration data for this daemon instance.
#[derive(Clone)]
pub struct AppConfig {
    /// Version of the daemon (used in the `WelcomeHandler`).
    version: String,
    /// Base URL where this daemon can be found. This is used to construct
    /// callback URLs.
    base_url: String,
    /// Private key used to sign identity tokens.
    priv_key: PKey,
    /// Redis database connector.
    store: Client,
    /// Duration in seconds for expiry of data stored in Redis.
    expire_keys: usize,
    /// Email sender (email address, then human-readable name).
    sender: (String, String),
    /// Duration in seconds for expiration of identity tokens.
    token_validity: i64,
    /// A map of "famous" identity providers. Each key is an email domain,
    /// the value holds the configuration required to act as an application
    /// doing OpenID Connect against the provider.
    providers: BTreeMap<String, ProviderConfig>,
}

/// Implementation with single method to read configuration from JSON.
impl AppConfig {
    pub fn from_json_file(file_name: &str) -> AppConfig {

        let config_file = File::open(file_name).unwrap();
        let config_value: Value = from_reader(BufReader::new(config_file)).unwrap();
        let config = config_value.as_object().unwrap();

        // The `private_key_file` key in the JSON object should hold a file
        // name pointing to a PEM-encoded private RSA key.
        let key_file_name = config["private_key_file"].as_string().unwrap();
        let priv_key_file = File::open(key_file_name).unwrap();
        let mut reader = BufReader::new(priv_key_file);
        let sender = config["sender"].as_object().unwrap();
        AppConfig {
            // Use the crate's version as defined in Cargo.toml.
            version: env!("CARGO_PKG_VERSION").to_string(),
            base_url: config["base_url"].as_string().unwrap().to_string(),
            priv_key: PKey::private_key_from_pem(&mut reader).unwrap(),
            store: Client::open(config["redis"].as_string().unwrap()).unwrap(),
            expire_keys: config["expire_keys"].as_u64().unwrap() as usize,
            sender: (
                sender["address"].as_string().unwrap().to_string(),
                sender["name"].as_string().unwrap().to_string(),
            ),
            token_validity: config["token_validity"].as_i64().unwrap(),
            providers: config["providers"].as_object().unwrap().iter()
                .map(|(host, params)| {
                    let pobj = params.as_object().unwrap();
                    (host.clone(), ProviderConfig {
                        discovery: pobj["discovery"].as_string().unwrap().to_string(),
                        client_id: pobj["client_id"].as_string().unwrap().to_string(),
                        secret: pobj["secret"].as_string().unwrap().to_string(),
                        issuer: pobj["issuer"].as_string().unwrap().to_string(),
                    })
                })
                .collect::<BTreeMap<String, ProviderConfig>>(),
        }

    }
}

/// Iron handler for the root path, returns human-friendly message.
///
/// This is not actually used in the protocol.
pub struct WelcomeHandler { pub app: AppConfig }
impl Handler for WelcomeHandler {
    fn handle(&self, _: &mut Request) -> IronResult<Response> {
        json_response(&ObjectBuilder::new()
            .insert("ladaemon", "Welcome")
            .insert("version", &self.app.version)
            .unwrap())
    }
}

/// Iron handler to return the OpenID Discovery document.
///
/// Most of this is hard-coded for now, although the URLs are constructed by
/// using the base URL as configured in the `base_url` configuration value.
pub struct OIDConfigHandler { pub app: AppConfig }
impl Handler for OIDConfigHandler {
    fn handle(&self, _: &mut Request) -> IronResult<Response> {
        json_response(&ObjectBuilder::new()
            .insert("issuer", &self.app.base_url)
            .insert("authorization_endpoint",
                    format!("{}/auth", self.app.base_url))
            .insert("jwks_uri", format!("{}/keys.json", self.app.base_url))
            .insert("scopes_supported", vec!["openid", "email"])
            .insert("claims_supported",
                    vec!["aud", "email", "email_verified", "exp", "iat", "iss", "sub"])
            .insert("response_types_supported", vec!["id_token"])
            .insert("response_modes_supported", vec!["form_post"])
            .insert("grant_types_supported", vec!["implicit"])
            .insert("subject_types_supported", vec!["public"])
            .insert("id_token_signing_alg_values_supported", vec!["RS256"])
            .unwrap())
    }
}

/// Helper function to encode a `BigNum` as URL-safe base64-encoded bytes.
///
/// This is used for the public RSA key components returned by the
/// `KeysHandler`.
fn json_big_num(n: &BigNum) -> String {
    n.to_vec().to_base64(base64::URL_SAFE)
}

/// Helper function to build a session ID for a login attempt.
///
/// Put the email address, the client ID (RP origin) and some randomness into
/// a SHA256 hash, and encode it with URL-safe bas64 encoding. This is used
/// as the key in Redis, as well as the state for OAuth authentication.
fn session_id(email: &EmailAddress, client_id: &str) -> String {
    let mut rng = OsRng::new().unwrap();
    let mut bytes_iter = rng.gen_iter();
    let rand_bytes: Vec<u8> = (0..16).map(|_| bytes_iter.next().unwrap()).collect();

    let mut hasher = hash::Hasher::new(hash::Type::SHA256);
    hasher.write(email.to_string().as_bytes()).unwrap();
    hasher.write(client_id.as_bytes()).unwrap();
    hasher.write(&rand_bytes).unwrap();
    hasher.finish().to_base64(base64::URL_SAFE)
}

/// Iron handler for the JSON Web Key Set document.
///
/// Currently only supports a single RSA key (as configured), which is
/// published with the `"base"` key ID, scoped to signing usage. Relying
/// Parties will need to fetch this data to be able to verify identity tokens
/// issued by this daemon instance.
pub struct KeysHandler { pub app: AppConfig }
impl Handler for KeysHandler {
    fn handle(&self, _: &mut Request) -> IronResult<Response> {
        let rsa = self.app.priv_key.get_rsa();
        json_response(&ObjectBuilder::new()
            .insert_array("keys", |builder| {
                builder.push_object(|builder| {
                    builder.insert("kty", "RSA")
                        .insert("alg", "RS256")
                        .insert("use", "sig")
                        .insert("kid", "base")
                        .insert("n", json_big_num(&rsa.n().unwrap()))
                        .insert("e", json_big_num(&rsa.e().unwrap()))
                })
            })
            .unwrap())
    }
}

/// Iron handler for authentication requests from the RP.
///
/// Calls the `oidc::request()` function if the provided email address's
/// domain matches one of the configured famous providers. Otherwise, calls the
/// `email::request()` function to allow authentication through the email loop.
pub struct AuthHandler { pub app: AppConfig }
impl Handler for AuthHandler {
    fn handle(&self, req: &mut Request) -> IronResult<Response> {
        let params = req.get_ref::<UrlEncodedBody>().unwrap();
        let email_addr = EmailAddress::new(&params.get("login_hint").unwrap()[0]).unwrap();
        if self.app.providers.contains_key(&email_addr.domain) {

            // OIDC authentication. Using 302 Found for redirection here. Note
            // that, per RFC 7231, a user agent MAY change the request method
            // from POST to GET for the subsequent request.
            let auth_url = oidc::request(&self.app, params);
            Ok(Response::with((status::Found, modifiers::Redirect(auth_url))))

        } else {

            // Email loop authentication. For now, returns a JSON response;
            // empty if successful, otherwise contains an error.
            let obj = email::request(&self.app, params);
            json_response(&obj)

        }
    }
}

/// Helper method to create a JWT for a given email address and origin.
///
/// Builds the JSON payload and header, encoding with (URL-safe)
/// base64-encoding, then hashing and signing with the provided private key.
/// Returns the full JWT.
fn create_jwt(app: &AppConfig, email: &str, origin: &str) -> String {

    let now = now_utc().to_timespec().sec;
    let payload = serde_json::to_string(
        &ObjectBuilder::new()
            .insert("aud", origin)
            .insert("email", email)
            .insert("email_verified", email)
            .insert("exp", now + app.token_validity)
            .insert("iat", now)
            .insert("iss", &app.base_url)
            .insert("sub", email)
            .unwrap()
        ).unwrap();
    let header = serde_json::to_string(
        &ObjectBuilder::new()
            .insert("kid", "base")
            .insert("alg", "RS256")
            .unwrap()
        ).unwrap();

    let mut input = Vec::<u8>::new();
    input.extend(header.as_bytes().to_base64(base64::URL_SAFE).into_bytes());
    input.push(b'.');
    input.extend(payload.as_bytes().to_base64(base64::URL_SAFE).into_bytes());
    let sha256 = hash::hash(hash::Type::SHA256, &input);
    let sig = app.priv_key.sign(&sha256);
    input.push(b'.');
    input.extend(sig.to_base64(base64::URL_SAFE).into_bytes());
    String::from_utf8(input).unwrap()

}

/// HTML template used to have the user agent POST the identity token built
/// by the daemon instance to the RP's `redirect_uri`.
const FORWARD_TEMPLATE: &'static str = r#"<!DOCTYPE html>
<html>
  <head>
    <title>Let's Auth</title>
    <script>
      document.addEventListener('DOMContentLoaded', function() {
        document.getElementById('form').submit();
      });
    </script>
  </head>
  <body>
    <form id="form" action="{{ return_url }}" method="post">
      <input type="hidden" name="id_token" value="{{ jwt }}">
    </form>
  </body>
</html>"#;

/// Helper function for returning result to the Relying Party.
///
/// Takes a `Result` from one of the verification functions and embeds it in
/// a form in the `FORWARD_TEMPLATE`, from where it's POSTED to the RP's
/// `redirect_ur` as soon as the page has loaded. Result can either be an error
/// message or a JWT asserting the user's email address identity.
/// TODO: return error to RP instead of in a simple HTTP response.
fn return_to_relier(result: Result<(String, String), &'static str>)
                    -> IronResult<Response> {

    if result.is_err() {
        return json_response(&ObjectBuilder::new()
                            .insert("error", result.unwrap_err())
                            .unwrap());
    }

    let (jwt, redirect) = result.unwrap();
    let html = FORWARD_TEMPLATE.replace("{{ return_url }}", &redirect)
        .replace("{{ jwt }}", &jwt);
    let mut rsp = Response::with((status::Ok, html));
    rsp.headers.set(ContentType::html());
    Ok(rsp)

}

/// Iron handler for one-time pad email loop confirmation.
///
/// Retrieves the session based session ID and the expected one-time pad.
/// Verify the code and return the resulting token or error to the RP.
pub struct ConfirmHandler { pub app: AppConfig }
impl Handler for ConfirmHandler {
    fn handle(&self, req: &mut Request) -> IronResult<Response> {
        let params = req.get_ref::<UrlEncodedQuery>().unwrap();
        let session_id = &params.get("session").unwrap()[0];
        let code = &params.get("code").unwrap()[0];
        return_to_relier(email::verify(&self.app, session_id, code))
    }
}

/// Iron handler for OAuth callbacks
///
/// After the user allows or denies the Authentication Request with the famous
/// identity provider, they will be redirected back to the callback handler.
/// Verify the callback data and return the resulting token or error.
pub struct CallbackHandler { pub app: AppConfig }
impl Handler for CallbackHandler {
    fn handle(&self, req: &mut Request) -> IronResult<Response> {
        let params = req.get_ref::<UrlEncodedQuery>().unwrap();
        let session = &params.get("state").unwrap()[0];
        let code = &params.get("code").unwrap()[0];
        return_to_relier(oidc::verify(&self.app, session, code))
    }
}
<|endoftext|>