<commit_before>#![deny(warnings)]

extern crate tar;
extern crate tree_magic;
extern crate lzma;
extern crate libflate;

use std::{env, process};
use std::io::{stdin, stdout, stderr, copy, Error, ErrorKind, Result, Read, Write, BufReader};
use std::fs::{self, File};
use std::os::unix::fs::OpenOptionsExt;
use std::path::Path;

use tar::{Archive, Builder, EntryType};
use lzma::LzmaReader;
use libflate::gzip::Decoder as GzipDecoder;

fn create_inner<T: Write>(input: &str, ar: &mut Builder<T>) -> Result<()> {
    if try!(fs::metadata(input)).is_dir() {
        for entry_result in try!(fs::read_dir(input)) {
            let entry = try!(entry_result);
            if try!(fs::metadata(entry.path())).is_dir() {
                try!(create_inner(entry.path().to_str().unwrap(), ar));
            } else {
                println!("{}", entry.path().display());
                try!(ar.append_path(entry.path()));
            }
        }
    } else {
        println!("{}", input);
        try!(ar.append_path(input));
    }

    Ok(())
}

fn create(input: &str, tar: &str) -> Result<()> {
    if tar == "-" {
        create_inner(input, &mut Builder::new(stdout()))
    } else {
        create_inner(input, &mut Builder::new(try!(File::create(tar))))
    }
}

fn list_inner<T: Read>(ar: &mut Archive<T>) -> Result<()> {
    for entry_result in try!(ar.entries()) {
        let entry = try!(entry_result);
        let path = try!(entry.path());
        println!("{}", path.display());
    }

    Ok(())
}

fn list(tar: &str) -> Result<()> {
    if tar == "-" {
        list_inner(&mut Archive::new(stdin()))
    } else {
        list_inner(&mut Archive::new(try!(File::open(tar))))
    }
}

fn extract_inner<T: Read>(ar: &mut Archive<T>, verbose: bool) -> Result<()> {
    for entry_result in try!(ar.entries()) {
        let mut entry = try!(entry_result);
        match entry.header().entry_type() {
            EntryType::Regular => {
                let mut file = {
                    let path = try!(entry.path());
                    if let Some(parent) = path.parent() {
                        try!(fs::create_dir_all(parent));
                    }
                    try!(
                        fs::OpenOptions::new()
                            .read(true)
                            .write(true)
                            .truncate(true)
                            .create(true)
                            .mode(entry.header().mode().unwrap_or(644))
                            .open(path)
                    )
                };
                try!(copy(&mut entry, &mut file));
            },
            EntryType::Directory => {
                try!(fs::create_dir_all(try!(entry.path())));
            },
            other => {
                panic!("Unsupported entry type {:?}", other);
            }
        }

        if verbose {
            println!("{}", entry.path()?.display());
        }
    }

    Ok(())
}

fn extract(tar: &str, verbose: bool) -> Result<()> {
    if tar == "-" {
        extract_inner(&mut Archive::new(stdin()), verbose)
    } else {
        let mime = tree_magic::from_filepath(Path::new(&tar));
        let file = BufReader::new(File::open(tar)?);
        if mime == "application/x-xz" {
            extract_inner(&mut Archive::new(LzmaReader::new_decompressor(file)
                                            .map_err(|e| Error::new(ErrorKind::Other, e))?), verbose)
        } else if mime == "application/gzip" {
            extract_inner(&mut Archive::new(GzipDecoder::new(file)
                                            .map_err(|e| Error::new(ErrorKind::Other, e))?), verbose)
        } else {
            extract_inner(&mut Archive::new(file), verbose)
        }
    }
}

fn main() {
    let mut args = env::args().skip(1);
    if let Some(op) = args.next() {
        match op.as_str() {
            "c" => if let Some(input) = args.next() {
                if let Err(err) = create(&input, "-") {
                    write!(stderr(), "tar: create: failed: {}\n", err).unwrap();
                    process::exit(1);
                }
            } else {
                write!(stderr(), "tar: create: no input specified: {}\n", op).unwrap();
                process::exit(1);
            },
            "cf" => if let Some(tar) = args.next() {
                if let Some(input) = args.next() {
                    if let Err(err) = create(&input, &tar) {
                        write!(stderr(), "tar: create: failed: {}\n", err).unwrap();
                        process::exit(1);
                    }
                } else {
                    write!(stderr(), "tar: create: no input specified: {}\n", op).unwrap();
                    process::exit(1);
                }
            } else {
                write!(stderr(), "tar: create: no tarfile specified: {}\n", op).unwrap();
                process::exit(1);
            },
            "t" | "tf" => {
                let tar = args.next().unwrap_or("-".to_string());
                if let Err(err) = list(&tar) {
                    write!(stderr(), "tar: list: failed: {}\n", err).unwrap();
                    process::exit(1);
                }
            },
            "x" | "xf" | "xvf" => {
                let tar = args.next().unwrap_or("-".to_string());
                let verbose = op.contains('v');
                if let Err(err) = extract(&tar, verbose) {
                    write!(stderr(), "tar: extract: failed: {}\n", err).unwrap();
                    process::exit(1);
                }
            },
            _ => {
                write!(stderr(), "tar: {}: unknown operation\n", op).unwrap();
                write!(stderr(), "tar: need to specify c[f] (create), t[f] (list), or x[f] (extract)\n").unwrap();
                process::exit(1);
            }
        }
    } else {
        write!(stderr(), "tar: no operation\n").unwrap();
        write!(stderr(), "tar: need to specify cf (create), tf (list), or xf (extract)\n").unwrap();
        process::exit(1);
    }
}
<commit_msg>tar: support -C, --directory<commit_after>#![deny(warnings)]

extern crate tar;
extern crate tree_magic;
extern crate lzma;
extern crate libflate;

use std::{env, process};
use std::io::{stdin, stdout, stderr, copy, Error, ErrorKind, Result, Read, Write, BufReader};
use std::fs::{self, File};
use std::os::unix::fs::OpenOptionsExt;
use std::path::Path;

use tar::{Archive, Builder, EntryType};
use lzma::LzmaReader;
use libflate::gzip::Decoder as GzipDecoder;

fn create_inner<T: Write>(input: &str, ar: &mut Builder<T>) -> Result<()> {
    if try!(fs::metadata(input)).is_dir() {
        for entry_result in try!(fs::read_dir(input)) {
            let entry = try!(entry_result);
            if try!(fs::metadata(entry.path())).is_dir() {
                try!(create_inner(entry.path().to_str().unwrap(), ar));
            } else {
                println!("{}", entry.path().display());
                try!(ar.append_path(entry.path()));
            }
        }
    } else {
        println!("{}", input);
        try!(ar.append_path(input));
    }

    Ok(())
}

fn create(input: &str, tar: &str) -> Result<()> {
    if tar == "-" {
        create_inner(input, &mut Builder::new(stdout()))
    } else {
        create_inner(input, &mut Builder::new(try!(File::create(tar))))
    }
}

fn list_inner<T: Read>(ar: &mut Archive<T>) -> Result<()> {
    for entry_result in try!(ar.entries()) {
        let entry = try!(entry_result);
        let path = try!(entry.path());
        println!("{}", path.display());
    }

    Ok(())
}

fn list(tar: &str) -> Result<()> {
    if tar == "-" {
        list_inner(&mut Archive::new(stdin()))
    } else {
        list_inner(&mut Archive::new(try!(File::open(tar))))
    }
}

fn extract_inner<T: Read>(ar: &mut Archive<T>, verbose: bool) -> Result<()> {
    for entry_result in try!(ar.entries()) {
        let mut entry = try!(entry_result);
        match entry.header().entry_type() {
            EntryType::Regular => {
                let mut file = {
                    let path = try!(entry.path());
                    if let Some(parent) = path.parent() {
                        try!(fs::create_dir_all(parent));
                    }
                    try!(
                        fs::OpenOptions::new()
                            .read(true)
                            .write(true)
                            .truncate(true)
                            .create(true)
                            .mode(entry.header().mode().unwrap_or(644))
                            .open(path)
                    )
                };
                try!(copy(&mut entry, &mut file));
            },
            EntryType::Directory => {
                try!(fs::create_dir_all(try!(entry.path())));
            },
            other => {
                panic!("Unsupported entry type {:?}", other);
            }
        }

        if verbose {
            println!("{}", entry.path()?.display());
        }
    }

    Ok(())
}

fn extract(tar: &str, verbose: bool) -> Result<()> {
    if tar == "-" {
        extract_inner(&mut Archive::new(stdin()), verbose)
    } else {
        let mime = tree_magic::from_filepath(Path::new(&tar));
        let file = BufReader::new(File::open(tar)?);
        if mime == "application/x-xz" {
            extract_inner(&mut Archive::new(LzmaReader::new_decompressor(file)
                                            .map_err(|e| Error::new(ErrorKind::Other, e))?), verbose)
        } else if mime == "application/gzip" {
            extract_inner(&mut Archive::new(GzipDecoder::new(file)
                                            .map_err(|e| Error::new(ErrorKind::Other, e))?), verbose)
        } else {
            extract_inner(&mut Archive::new(file), verbose)
        }
    }
}

fn main() {
    let mut args = env::args().skip(1);
    if let Some(op) = args.next() {
        match op.as_str() {
            "c" => if let Some(input) = args.next() {
                if let Err(err) = create(&input, "-") {
                    write!(stderr(), "tar: create: failed: {}\n", err).unwrap();
                    process::exit(1);
                }
            } else {
                write!(stderr(), "tar: create: no input specified: {}\n", op).unwrap();
                process::exit(1);
            },
            "cf" => if let Some(tar) = args.next() {
                if let Some(input) = args.next() {
                    if let Err(err) = create(&input, &tar) {
                        write!(stderr(), "tar: create: failed: {}\n", err).unwrap();
                        process::exit(1);
                    }
                } else {
                    write!(stderr(), "tar: create: no input specified: {}\n", op).unwrap();
                    process::exit(1);
                }
            } else {
                write!(stderr(), "tar: create: no tarfile specified: {}\n", op).unwrap();
                process::exit(1);
            },
            "t" | "tf" => {
                let tar = args.next().unwrap_or("-".to_string());
                if let Err(err) = list(&tar) {
                    write!(stderr(), "tar: list: failed: {}\n", err).unwrap();
                    process::exit(1);
                }
            },
            "x" | "xf" | "xvf" => {
                let mut tar = None;
                while let Some(arg) = args.next() {
                    if arg == "-C" {
                        env::set_current_dir(args.next().expect("-C requires path")).unwrap();
                    } else if arg.starts_with("--directory=") {
                        env::set_current_dir(&arg[12..]).unwrap();
                    } else if tar.is_none() {
                        tar = Some(arg);
                    }
                }
                let tar = tar.unwrap_or("-".to_string());

                let verbose = op.contains('v');
                if let Err(err) = extract(&tar, verbose) {
                    write!(stderr(), "tar: extract: failed: {}\n", err).unwrap();
                    process::exit(1);
                }
            },
            _ => {
                write!(stderr(), "tar: {}: unknown operation\n", op).unwrap();
                write!(stderr(), "tar: need to specify c[f] (create), t[f] (list), or x[f] (extract)\n").unwrap();
                process::exit(1);
            }
        }
    } else {
        write!(stderr(), "tar: no operation\n").unwrap();
        write!(stderr(), "tar: need to specify cf (create), tf (list), or xf (extract)\n").unwrap();
        process::exit(1);
    }
}
<|endoftext|>


<commit_before>#![crate_name = "nohup"]

/*
 * This file is part of the uutils coreutils package.
 *
 * (c) 2014 Vsevolod Velichko <torkvemada@sorokdva.net>
 *
 * For the full copyright and license information, please view the LICENSE
 * file that was distributed with this source code.
 */

#![feature(macro_rules)]
extern crate getopts;
extern crate libc;
extern crate native;

use getopts::{optflag, getopts, usage};
use std::io::stdio::{stdin_raw, stdout_raw, stderr_raw};
use std::rt::rtio::{Open, Read, Append, Write};
use libc::funcs::posix88::unistd::{dup2, execvp};
use libc::consts::os::posix88::SIGHUP;
use libc::funcs::posix01::signal::signal;
use libc::consts::os::posix01::SIG_IGN;

#[path = "../common/util.rs"] mod util;
#[path = "../common/c_types.rs"] mod c_types;

static NAME: &'static str = "nohup";
static VERSION: &'static str = "1.0.0";

#[cfg(target_os = "macos")]
extern {
    fn _vprocmgr_detach_from_console(flags: u32) -> *const libc::c_int;
}

#[cfg(target_os = "macos")]
fn rewind_stdout<T: std::rt::rtio::RtioFileStream>(s: &mut T) {
    match s.seek(0, std::rt::rtio::SeekEnd) {
        Ok(_) => {}
        Err(f) => crash!(1, "{}", f.detail.unwrap())
    }
}

#[cfg(any(target_os = "linux", target_os = "freebsd"))]
unsafe fn _vprocmgr_detach_from_console(_: u32) -> *const libc::c_int { std::ptr::null() }

#[cfg(any(target_os = "linux", target_os = "freebsd"))]
fn rewind_stdout<T: std::rt::rtio::RtioFileStream>(_: &mut T) {}

pub fn uumain(args: Vec<String>) -> int {
    let program = &args[0];

    let options = [
        optflag("h", "help", "Show help and exit"),
        optflag("V", "version", "Show version and exit"),
    ];

    let opts = match getopts(args.tail(), options) {
        Ok(m) => m,
        Err(f) => {
            show_error!("{}", f);
            show_usage(program.as_slice(), options);
            return 1
        }
    };

    if opts.opt_present("V") { version(); return 0 }
    if opts.opt_present("h") { show_usage(program.as_slice(), options); return 0 }

    if opts.free.len() == 0 {
        show_error!("Missing operand: COMMAND");
        println!("Try `{:s} --help` for more information.", program.as_slice());
        return 1
    }
    replace_fds();

    unsafe { signal(SIGHUP, SIG_IGN) };

    if unsafe { _vprocmgr_detach_from_console(0) } != std::ptr::null() { crash!(2, "Cannot detach from console")};

    unsafe {
        // we ignore the memory leak here because it doesn't matter anymore
        let executable = opts.free[0].as_slice().to_c_str().unwrap();
        let mut args: Vec<*const i8> = opts.free.iter().map(|x| x.to_c_str().unwrap()).collect();
        args.push(std::ptr::null());
        execvp(executable as *const libc::c_char, args.as_ptr() as *mut *const libc::c_char) as int
    }
}

fn replace_fds() {
    let replace_stdin = stdin_raw().isatty();
    let replace_stdout = stdout_raw().isatty();
    let replace_stderr = stderr_raw().isatty();

    if replace_stdin {
        let devnull = "/dev/null".to_c_str();
        let new_stdin = match native::io::file::open(&devnull, Open, Read) {
            Ok(t) => t,
            Err(_) => {
                let e = std::io::IoError::last_error();
                crash!(2, "Cannot replace STDIN: {}", e)
            }
        };
        if unsafe { dup2(new_stdin.fd(), 0) } != 0 {
            crash!(2, "Cannot replace STDIN: {}", std::io::IoError::last_error())
        }
    }

    if replace_stdout {
        let mut new_stdout = find_stdout();

        rewind_stdout(&mut new_stdout);

        if unsafe { dup2(new_stdout.fd(), 1) } != 1 {
            crash!(2, "Cannot replace STDOUT: {}", std::io::IoError::last_error())
        }
    }

    if replace_stderr {
        if unsafe { dup2(1, 2) } != 2 {
            crash!(2, "Cannot replace STDERR: {}", std::io::IoError::last_error())
        }
    }
}

fn find_stdout() -> native::io::file::FileDesc {
    let localout = "nohup.out".to_c_str();
    match native::io::file::open(&localout, Append, Write) {
        Ok(t) => {
            show_warning!("Output is redirected to: nohup.out");
            t
        },
        Err(_) => {
            let e = std::io::IoError::last_error();
            let home = match std::os::getenv("HOME") {
                None => crash!(2, "Cannot replace STDOUT: {}", e),
                Some(h) => h
            };
            let mut homeoutpath = Path::new(home);
            homeoutpath.push("nohup.out");
            let homeout = homeoutpath.to_c_str();
            match native::io::file::open(&homeout, Append, Write) {
                Ok(t) => {
                    show_warning!("Output is redirected to: {}", homeoutpath.display());
                    t
                },
                Err(_) => {
                    let e = std::io::IoError::last_error();
                    crash!(2, "Cannot replace STDOUT: {}", e)
                }
            }
        }
    }
}

fn version() {
    println!("{} v{}", NAME, VERSION)
}

fn show_usage(program: &str, options: &[getopts::OptGroup]) {
    version();
    println!("Usage:");
    println!("  {:s} COMMAND [ARG]…", program);
    println!("  {:s} OPTION", program);
    println!("");
    print!("{:s}", usage(
            "Run COMMAND ignoring hangup signals.\n\
            If standard input is terminal, it'll be replaced with /dev/null.\n\
            If standard output is terminal, it'll be appended to nohup.out instead, \
            or $HOME/nohup.out, if nohup.out open failed.\n\
            If standard error is terminal, it'll be redirected to stdout.", options)
    );
}
<commit_msg>nohup: fix for latest Rust using terrible hacks<commit_after>#![crate_name = "nohup"]

/*
 * This file is part of the uutils coreutils package.
 *
 * (c) 2014 Vsevolod Velichko <torkvemada@sorokdva.net>
 *
 * For the full copyright and license information, please view the LICENSE
 * file that was distributed with this source code.
 */

#![feature(macro_rules)]
extern crate getopts;
extern crate libc;
extern crate native;

use getopts::{optflag, getopts, usage};
use std::io::stdio::{stdin_raw, stdout_raw, stderr_raw};
use std::io::{File, Open, Read, Append, Write};
use libc::funcs::posix88::unistd::{dup2, execvp};
use libc::consts::os::posix88::SIGHUP;
use libc::funcs::posix01::signal::signal;
use libc::consts::os::posix01::SIG_IGN;

#[path = "../common/util.rs"] mod util;
#[path = "../common/c_types.rs"] mod c_types;

static NAME: &'static str = "nohup";
static VERSION: &'static str = "1.0.0";

#[cfg(target_os = "macos")]
extern {
    fn _vprocmgr_detach_from_console(flags: u32) -> *const libc::c_int;
}

// BEGIN CODE TO DELETE AFTER https://github.com/rust-lang/rust/issues/18897 is fixed
struct HackyFile {
    pub fd: FileDesc,
    path: Path,
    last_nread: int
}

struct FileDesc {
    fd: libc::c_int,
    close_on_drop: bool
}

trait AsFileDesc {
    fn as_fd(&self) -> FileDesc;
}

impl AsFileDesc for File {
    fn as_fd(&self) -> FileDesc {
        let hack: HackyFile = unsafe { std::mem::transmute_copy(self) };
        hack.fd
    }
}
// END CODE TO DELETE

#[cfg(target_os = "macos")]
fn rewind_stdout(s: &mut FileDesc) {
    match s.seek(0, io::SeekEnd) {
        Ok(_) => {}
        Err(f) => crash!(1, "{}", f.detail.unwrap())
    }
}

#[cfg(any(target_os = "linux", target_os = "freebsd"))]
unsafe fn _vprocmgr_detach_from_console(_: u32) -> *const libc::c_int { std::ptr::null() }

#[cfg(any(target_os = "linux", target_os = "freebsd"))]
fn rewind_stdout(_: &mut FileDesc) {}

pub fn uumain(args: Vec<String>) -> int {
    let program = &args[0];

    let options = [
        optflag("h", "help", "Show help and exit"),
        optflag("V", "version", "Show version and exit"),
    ];

    let opts = match getopts(args.tail(), options) {
        Ok(m) => m,
        Err(f) => {
            show_error!("{}", f);
            show_usage(program.as_slice(), options);
            return 1
        }
    };

    if opts.opt_present("V") { version(); return 0 }
    if opts.opt_present("h") { show_usage(program.as_slice(), options); return 0 }

    if opts.free.len() == 0 {
        show_error!("Missing operand: COMMAND");
        println!("Try `{:s} --help` for more information.", program.as_slice());
        return 1
    }
    replace_fds();

    unsafe { signal(SIGHUP, SIG_IGN) };

    if unsafe { _vprocmgr_detach_from_console(0) } != std::ptr::null() { crash!(2, "Cannot detach from console")};

    unsafe {
        // we ignore the memory leak here because it doesn't matter anymore
        let executable = opts.free[0].as_slice().to_c_str().unwrap();
        let mut args: Vec<*const i8> = opts.free.iter().map(|x| x.to_c_str().unwrap()).collect();
        args.push(std::ptr::null());
        execvp(executable as *const libc::c_char, args.as_ptr() as *mut *const libc::c_char) as int
    }
}

fn replace_fds() {
    let replace_stdin = stdin_raw().isatty();
    let replace_stdout = stdout_raw().isatty();
    let replace_stderr = stderr_raw().isatty();

    if replace_stdin {
        let new_stdin = match File::open_mode(&Path::new("/dev/null"), Open, Read) {
            Ok(t) => t,
            Err(e) => {
                crash!(2, "Cannot replace STDIN: {}", e)
            }
        };
        if unsafe { dup2(new_stdin.as_fd().fd, 0) } != 0 {
            crash!(2, "Cannot replace STDIN: {}", std::io::IoError::last_error())
        }
    }

    if replace_stdout {
        let new_stdout = find_stdout();
        let mut fd = new_stdout.as_fd();

        rewind_stdout(&mut fd);

        if unsafe { dup2(fd.fd, 1) } != 1 {
            crash!(2, "Cannot replace STDOUT: {}", std::io::IoError::last_error())
        }
    }

    if replace_stderr {
        if unsafe { dup2(1, 2) } != 2 {
            crash!(2, "Cannot replace STDERR: {}", std::io::IoError::last_error())
        }
    }
}

fn find_stdout() -> File {
    match File::open_mode(&Path::new("nohup.out"), Append, Write) {
        Ok(t) => {
            show_warning!("Output is redirected to: nohup.out");
            t
        },
        Err(e) => {
            let home = match std::os::getenv("HOME") {
                None => crash!(2, "Cannot replace STDOUT: {}", e),
                Some(h) => h
            };
            let mut homeout = Path::new(home);
            homeout.push("nohup.out");
            match File::open_mode(&homeout, Append, Write) {
                Ok(t) => {
                    show_warning!("Output is redirected to: {}", homeout.display());
                    t
                },
                Err(e) => {
                    crash!(2, "Cannot replace STDOUT: {}", e)
                }
            }
        }
    }
}

fn version() {
    println!("{} v{}", NAME, VERSION)
}

fn show_usage(program: &str, options: &[getopts::OptGroup]) {
    version();
    println!("Usage:");
    println!("  {:s} COMMAND [ARG]…", program);
    println!("  {:s} OPTION", program);
    println!("");
    print!("{:s}", usage(
            "Run COMMAND ignoring hangup signals.\n\
            If standard input is terminal, it'll be replaced with /dev/null.\n\
            If standard output is terminal, it'll be appended to nohup.out instead, \
            or $HOME/nohup.out, if nohup.out open failed.\n\
            If standard error is terminal, it'll be redirected to stdout.", options)
    );
}
<|endoftext|>


<commit_before>mod byte;
mod ring;
mod sink;
mod slice;
mod source;

pub use self::byte::{ByteBuf, MutByteBuf, ROByteBuf};
pub use self::ring::RingBuf;
pub use self::slice::{SliceBuf, MutSliceBuf};

use {BufError, RopeBuf};
use std::{cmp, fmt, io, ptr};

/// A trait for values that provide sequential read access to bytes.
pub trait Buf {

    /// Returns the number of bytes that can be accessed from the Buf
    fn remaining(&self) -> usize;

    /// Returns a slice starting at the current Buf position and of length
    /// between 0 and `Buf::remaining()`.
    fn bytes<'a>(&'a self) -> &'a [u8];

    /// Advance the internal cursor of the Buf
    fn advance(&mut self, cnt: usize);

    /// Returns true if there are any more bytes to consume
    fn has_remaining(&self) -> bool {
        self.remaining() > 0
    }

    /// Read bytes from the `Buf` into the given slice and advance the cursor by
    /// the number of bytes read.
    ///
    /// If there are fewer bytes remaining than is needed to satisfy the
    /// request (aka `dst.len()` > self.remaining()`), then
    /// `Err(BufError::Overflow)` is returned.
    ///
    /// ```
    /// use bytes::{SliceBuf, Buf};
    ///
    /// let mut buf = SliceBuf::wrap(b"hello world");
    /// let mut dst = [0; 5];
    ///
    /// buf.read_slice(&mut dst);
    /// assert_eq!(b"hello", &dst);
    /// assert_eq!(6, buf.remaining());
    /// ```
    fn read_slice(&mut self, dst: &mut [u8]) -> usize {
        let mut off = 0;
        let len = cmp::min(dst.len(), self.remaining());

        while off < len {
            let mut cnt;

            unsafe {
                let src = self.bytes();
                cnt = cmp::min(src.len(), len - off);

                ptr::copy_nonoverlapping(
                    src.as_ptr(), dst[off..].as_mut_ptr(), cnt);

                off += src.len();
            }

            self.advance(cnt);
        }

        len
    }

    /// Read a single byte from the `Buf`
    fn read_byte(&mut self) -> Option<u8> {
        let mut dst = [0];

        if self.read_slice(&mut dst) == 0 {
            return None;
        }

        Some(dst[0])
    }
}

/// An extension trait providing extra functions applicable to all `Buf` values.
pub trait BufExt {

    /// Read bytes from this Buf into the given sink and advance the cursor by
    /// the number of bytes read.
    fn read<S: Sink>(&mut self, dst: S) -> Result<usize, S::Error>;
}

/// A trait for values that provide sequential write access to bytes.
pub trait MutBuf : Sized {

    /// Returns the number of bytes that can be accessed from the Buf
    fn remaining(&self) -> usize;

    /// Advance the internal cursor of the Buf
    fn advance(&mut self, cnt: usize);

    /// Returns true if there are any more bytes to consume
    fn has_remaining(&self) -> bool {
        self.remaining() > 0
    }

    /// Returns a mutable slice starting at the current Buf position and of
    /// length between 0 and `Buf::remaining()`.
    fn mut_bytes<'a>(&'a mut self) -> &'a mut [u8];

    /// Read bytes from this Buf into the given slice and advance the cursor by
    /// the number of bytes read.
    ///
    /// If there are fewer bytes remaining than is needed to satisfy the
    /// request (aka `dst.len()` > self.remaining()`), then
    /// `Err(BufError::Overflow)` is returned.
    ///
    /// ```
    /// use bytes::{MutSliceBuf, Buf, MutBuf};
    ///
    /// let mut dst = [0; 6];
    ///
    /// {
    ///     let mut buf = MutSliceBuf::wrap(&mut dst);
    ///     buf.write_slice(b"hello");
    ///
    ///     assert_eq!(1, buf.remaining());
    /// }
    ///
    /// assert_eq!(b"hello\0", &dst);
    /// ```
    fn write_slice(&mut self, src: &[u8]) -> usize {
        let mut off = 0;
        let len = cmp::min(src.len(), self.remaining());

        while off < len {
            let mut cnt;

            unsafe {
                let dst = self.mut_bytes();
                cnt = cmp::min(dst.len(), len - off);

                ptr::copy_nonoverlapping(
                    src[off..].as_ptr(),
                    dst.as_mut_ptr(),
                    cnt);

                off += cnt;
            }

            self.advance(cnt);
        }

        len
    }

    /// Write a single byte to the `MuBuf`
    fn write_byte(&mut self, byte: u8) -> bool {
        let src = [byte];

        if self.write_slice(&src) == 0 {
            return false;
        }

        true
    }
}

/// An extension trait providing extra functions applicable to all `MutBuf` values.
pub trait MutBufExt {

    /// Write bytes from the given source into the current `MutBuf` and advance
    /// the cursor by the number of bytes written.
    fn write<S: Source>(&mut self, src: S) -> Result<usize, S::Error>;
}

/*
 *
 * ===== *Ext impls =====
 *
 */

impl<B: Buf> BufExt for B {
    fn read<S: Sink>(&mut self, dst: S) -> Result<usize, S::Error> {
        dst.sink(self)
    }
}

impl<B: MutBuf> MutBufExt for B {
    fn write<S: Source>(&mut self, src: S) -> Result<usize, S::Error> {
        src.fill(self)
    }
}

/*
 *
 * ===== Sink / Source =====
 *
 */

/// A value that reads bytes from a Buf into itself
pub trait Sink {
    type Error;

    fn sink<B: Buf>(self, buf: &mut B) -> Result<usize, Self::Error>;
}

/// A value that writes bytes from itself into a `MutBuf`.
pub trait Source {
    type Error;

    fn fill<B: MutBuf>(self, buf: &mut B) -> Result<usize, Self::Error>;
}

impl<'a> Sink for &'a mut [u8] {
    type Error = BufError;

    fn sink<B: Buf>(self, buf: &mut B) -> Result<usize, BufError> {
        Ok(buf.read_slice(self))
    }
}

impl<'a> Sink for &'a mut Vec<u8> {
    type Error = BufError;

    fn sink<B: Buf>(self, buf: &mut B) -> Result<usize, BufError> {
        use std::slice;

        self.clear();

        let rem = buf.remaining();
        let cap = self.capacity();

        // Ensure that the vec is big enough
        if rem > self.capacity() {
            self.reserve(rem - cap);
        }

        unsafe {
            {
                let dst = &mut self[..];
                let cnt = buf.read_slice(slice::from_raw_parts_mut(dst.as_mut_ptr(), rem));

                debug_assert!(cnt == rem);
            }

            self.set_len(rem);
        }

        Ok(rem)
    }
}

impl<'a> Source for &'a [u8] {
    type Error = BufError;

    fn fill<B: MutBuf>(self, buf: &mut B) -> Result<usize, BufError> {
        Ok(buf.write_slice(self))
    }
}

impl<'a> Source for &'a Vec<u8> {
    type Error = BufError;

    fn fill<B: MutBuf>(self, buf: &mut B) -> Result<usize, BufError> {
        Ok(buf.write_slice(self.as_ref()))
    }
}

impl<'a, R: io::Read+'a> Source for &'a mut R {
    type Error = io::Error;

    fn fill<B: MutBuf>(self, buf: &mut B) -> Result<usize, io::Error> {
        let mut cnt = 0;

        while buf.has_remaining() {
            let i = try!(self.read(buf.mut_bytes()));

            if i == 0 {
                break;
            }

            buf.advance(i);
            cnt += i;
        }

        Ok(cnt)
    }
}

/*
 *
 * ===== Buf impls =====
 *
 */

impl Buf for Box<Buf+'static> {
    fn remaining(&self) -> usize {
        (**self).remaining()
    }

    fn bytes(&self) -> &[u8] {
        (**self).bytes()
    }

    fn advance(&mut self, cnt: usize) {
        (**self).advance(cnt);
    }

    fn read_slice(&mut self, dst: &mut [u8]) -> usize {
        (**self).read_slice(dst)
    }
}

impl fmt::Debug for Box<Buf+'static> {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        write!(fmt, "Box<Buf> {{ remaining: {} }}", self.remaining())
    }
}

/*
 *
 * ===== Read impls =====
 *
 */

macro_rules! impl_read {
    ($ty:ty) => {
        impl io::Read for $ty {
            fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
                if !self.has_remaining() {
                    return Ok(0);
                }

                Ok(self.read_slice(buf))
            }
        }
    }
}

impl_read!(ByteBuf);
impl_read!(ROByteBuf);
impl_read!(RopeBuf);
impl_read!(Box<Buf+'static>);

macro_rules! impl_write {
    ($ty:ty) => {
        impl io::Write for $ty {
            fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
                if !self.has_remaining() {
                    return Ok(0);
                }

                Ok(self.write_slice(buf))
            }

            fn flush(&mut self) -> io::Result<()> {
                Ok(())
            }
        }
    }
}

impl_write!(MutByteBuf);
<commit_msg>Fix Buf and MutBuf documentation<commit_after>mod byte;
mod ring;
mod sink;
mod slice;
mod source;

pub use self::byte::{ByteBuf, MutByteBuf, ROByteBuf};
pub use self::ring::RingBuf;
pub use self::slice::{SliceBuf, MutSliceBuf};

use {BufError, RopeBuf};
use std::{cmp, fmt, io, ptr};

/// A trait for values that provide sequential read access to bytes.
pub trait Buf {

    /// Returns the number of bytes that can be accessed from the Buf
    fn remaining(&self) -> usize;

    /// Returns a slice starting at the current Buf position and of length
    /// between 0 and `Buf::remaining()`.
    fn bytes<'a>(&'a self) -> &'a [u8];

    /// Advance the internal cursor of the Buf
    fn advance(&mut self, cnt: usize);

    /// Returns true if there are any more bytes to consume
    fn has_remaining(&self) -> bool {
        self.remaining() > 0
    }

    /// Read bytes from the `Buf` into the given slice and advance the cursor by
    /// the number of bytes read.
    /// Returns the number of bytes read.
    ///
    /// ```
    /// use bytes::{SliceBuf, Buf};
    ///
    /// let mut buf = SliceBuf::wrap(b"hello world");
    /// let mut dst = [0; 5];
    ///
    /// buf.read_slice(&mut dst);
    /// assert_eq!(b"hello", &dst);
    /// assert_eq!(6, buf.remaining());
    /// ```
    fn read_slice(&mut self, dst: &mut [u8]) -> usize {
        let mut off = 0;
        let len = cmp::min(dst.len(), self.remaining());

        while off < len {
            let mut cnt;

            unsafe {
                let src = self.bytes();
                cnt = cmp::min(src.len(), len - off);

                ptr::copy_nonoverlapping(
                    src.as_ptr(), dst[off..].as_mut_ptr(), cnt);

                off += src.len();
            }

            self.advance(cnt);
        }

        len
    }

    /// Read a single byte from the `Buf`
    fn read_byte(&mut self) -> Option<u8> {
        let mut dst = [0];

        if self.read_slice(&mut dst) == 0 {
            return None;
        }

        Some(dst[0])
    }
}

/// An extension trait providing extra functions applicable to all `Buf` values.
pub trait BufExt {

    /// Read bytes from this Buf into the given sink and advance the cursor by
    /// the number of bytes read.
    fn read<S: Sink>(&mut self, dst: S) -> Result<usize, S::Error>;
}

/// A trait for values that provide sequential write access to bytes.
pub trait MutBuf : Sized {

    /// Returns the number of bytes that can be written to the MutBuf
    fn remaining(&self) -> usize;

    /// Advance the internal cursor of the MutBuf
    fn advance(&mut self, cnt: usize);

    /// Returns true iff there is any more space for bytes to be written
    fn has_remaining(&self) -> bool {
        self.remaining() > 0
    }

    /// Returns a mutable slice starting at the current MutBuf position and of
    /// length between 0 and `MutBuf::remaining()`.
    fn mut_bytes<'a>(&'a mut self) -> &'a mut [u8];

    /// Write bytes from the given slice into the `MutBuf` and advance the
    /// cursor by the number of bytes written.
    /// Returns the number of bytes written.
    ///
    /// ```
    /// use bytes::{MutSliceBuf, Buf, MutBuf};
    ///
    /// let mut dst = [0; 6];
    ///
    /// {
    ///     let mut buf = MutSliceBuf::wrap(&mut dst);
    ///     buf.write_slice(b"hello");
    ///
    ///     assert_eq!(1, buf.remaining());
    /// }
    ///
    /// assert_eq!(b"hello\0", &dst);
    /// ```
    fn write_slice(&mut self, src: &[u8]) -> usize {
        let mut off = 0;
        let len = cmp::min(src.len(), self.remaining());

        while off < len {
            let mut cnt;

            unsafe {
                let dst = self.mut_bytes();
                cnt = cmp::min(dst.len(), len - off);

                ptr::copy_nonoverlapping(
                    src[off..].as_ptr(),
                    dst.as_mut_ptr(),
                    cnt);

                off += cnt;
            }

            self.advance(cnt);
        }

        len
    }

    /// Write a single byte to the `MuBuf`
    fn write_byte(&mut self, byte: u8) -> bool {
        let src = [byte];

        if self.write_slice(&src) == 0 {
            return false;
        }

        true
    }
}

/// An extension trait providing extra functions applicable to all `MutBuf` values.
pub trait MutBufExt {

    /// Write bytes from the given source into the current `MutBuf` and advance
    /// the cursor by the number of bytes written.
    fn write<S: Source>(&mut self, src: S) -> Result<usize, S::Error>;
}

/*
 *
 * ===== *Ext impls =====
 *
 */

impl<B: Buf> BufExt for B {
    fn read<S: Sink>(&mut self, dst: S) -> Result<usize, S::Error> {
        dst.sink(self)
    }
}

impl<B: MutBuf> MutBufExt for B {
    fn write<S: Source>(&mut self, src: S) -> Result<usize, S::Error> {
        src.fill(self)
    }
}

/*
 *
 * ===== Sink / Source =====
 *
 */

/// A value that reads bytes from a Buf into itself
pub trait Sink {
    type Error;

    fn sink<B: Buf>(self, buf: &mut B) -> Result<usize, Self::Error>;
}

/// A value that writes bytes from itself into a `MutBuf`.
pub trait Source {
    type Error;

    fn fill<B: MutBuf>(self, buf: &mut B) -> Result<usize, Self::Error>;
}

impl<'a> Sink for &'a mut [u8] {
    type Error = BufError;

    fn sink<B: Buf>(self, buf: &mut B) -> Result<usize, BufError> {
        Ok(buf.read_slice(self))
    }
}

impl<'a> Sink for &'a mut Vec<u8> {
    type Error = BufError;

    fn sink<B: Buf>(self, buf: &mut B) -> Result<usize, BufError> {
        use std::slice;

        self.clear();

        let rem = buf.remaining();
        let cap = self.capacity();

        // Ensure that the vec is big enough
        if rem > self.capacity() {
            self.reserve(rem - cap);
        }

        unsafe {
            {
                let dst = &mut self[..];
                let cnt = buf.read_slice(slice::from_raw_parts_mut(dst.as_mut_ptr(), rem));

                debug_assert!(cnt == rem);
            }

            self.set_len(rem);
        }

        Ok(rem)
    }
}

impl<'a> Source for &'a [u8] {
    type Error = BufError;

    fn fill<B: MutBuf>(self, buf: &mut B) -> Result<usize, BufError> {
        Ok(buf.write_slice(self))
    }
}

impl<'a> Source for &'a Vec<u8> {
    type Error = BufError;

    fn fill<B: MutBuf>(self, buf: &mut B) -> Result<usize, BufError> {
        Ok(buf.write_slice(self.as_ref()))
    }
}

impl<'a, R: io::Read+'a> Source for &'a mut R {
    type Error = io::Error;

    fn fill<B: MutBuf>(self, buf: &mut B) -> Result<usize, io::Error> {
        let mut cnt = 0;

        while buf.has_remaining() {
            let i = try!(self.read(buf.mut_bytes()));

            if i == 0 {
                break;
            }

            buf.advance(i);
            cnt += i;
        }

        Ok(cnt)
    }
}

/*
 *
 * ===== Buf impls =====
 *
 */

impl Buf for Box<Buf+'static> {
    fn remaining(&self) -> usize {
        (**self).remaining()
    }

    fn bytes(&self) -> &[u8] {
        (**self).bytes()
    }

    fn advance(&mut self, cnt: usize) {
        (**self).advance(cnt);
    }

    fn read_slice(&mut self, dst: &mut [u8]) -> usize {
        (**self).read_slice(dst)
    }
}

impl fmt::Debug for Box<Buf+'static> {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        write!(fmt, "Box<Buf> {{ remaining: {} }}", self.remaining())
    }
}

/*
 *
 * ===== Read impls =====
 *
 */

macro_rules! impl_read {
    ($ty:ty) => {
        impl io::Read for $ty {
            fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
                if !self.has_remaining() {
                    return Ok(0);
                }

                Ok(self.read_slice(buf))
            }
        }
    }
}

impl_read!(ByteBuf);
impl_read!(ROByteBuf);
impl_read!(RopeBuf);
impl_read!(Box<Buf+'static>);

macro_rules! impl_write {
    ($ty:ty) => {
        impl io::Write for $ty {
            fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
                if !self.has_remaining() {
                    return Ok(0);
                }

                Ok(self.write_slice(buf))
            }

            fn flush(&mut self) -> io::Result<()> {
                Ok(())
            }
        }
    }
}

impl_write!(MutByteBuf);
<|endoftext|>


<commit_before>#![cfg_attr(not(any(unix, windows, target_os = "redox")), allow(unused_imports))]

use std::collections::HashMap;
use std::convert::AsRef;
use std::env;
use std::fs;
use std::io::Read;
use std::path::{Path, PathBuf};

use url::Url;

use crate::errors::Result;
use crate::types::{RawToken, SourceMap, Token};

/// Helper for sourcemap generation
///
/// This helper exists because generating and modifying `SourceMap`
/// objects is generally not very comfortable.  As a general aid this
/// type can help.
pub struct SourceMapBuilder {
    file: Option<String>,
    name_map: HashMap<String, u32>,
    names: Vec<String>,
    tokens: Vec<RawToken>,
    source_map: HashMap<String, u32>,
    sources: Vec<String>,
    source_contents: Vec<Option<String>>,
}

#[cfg(any(unix, windows, target_os = "redox"))]
fn resolve_local_reference(base: &Url, reference: &str) -> Option<PathBuf> {
    let url = match base.join(reference) {
        Ok(url) => {
            if url.scheme() != "file" {
                return None;
            }
            url
        }
        Err(_) => {
            return None;
        }
    };

    url.to_file_path().ok()
}

impl SourceMapBuilder {
    /// Creates a new source map builder and sets the file.
    pub fn new(file: Option<&str>) -> SourceMapBuilder {
        SourceMapBuilder {
            file: file.map(str::to_owned),
            name_map: HashMap::new(),
            names: vec![],
            tokens: vec![],
            source_map: HashMap::new(),
            sources: vec![],
            source_contents: vec![],
        }
    }

    /// Sets the file for the sourcemap (optional)
    pub fn set_file(&mut self, value: Option<&str>) {
        self.file = value.map(str::to_owned);
    }

    /// Returns the currently set file.
    pub fn get_file(&self) -> Option<&str> {
        self.file.as_ref().map(|x| &x[..])
    }

    /// Registers a new source with the builder and returns the source ID.
    pub fn add_source(&mut self, src: &str) -> u32 {
        let count = self.sources.len() as u32;
        let id = *self.source_map.entry(src.into()).or_insert(count);
        if id == count {
            self.sources.push(src.into());
        }
        id
    }

    /// Changes the source name for an already set source.
    pub fn set_source(&mut self, src_id: u32, src: &str) {
        assert!(src_id != !0, "Cannot set sources for tombstone source id");
        self.sources[src_id as usize] = src.to_string();
    }

    /// Looks up a source name for an ID.
    pub fn get_source(&self, src_id: u32) -> Option<&str> {
        self.sources.get(src_id as usize).map(|x| &x[..])
    }

    /// Sets the source contents for an already existing source.
    pub fn set_source_contents(&mut self, src_id: u32, contents: Option<&str>) {
        assert!(src_id != !0, "Cannot set sources for tombstone source id");
        if self.sources.len() > self.source_contents.len() {
            self.source_contents.resize(self.sources.len(), None);
        }
        self.source_contents[src_id as usize] = contents.map(str::to_owned);
    }

    /// Returns the current source contents for a source.
    pub fn get_source_contents(&self, src_id: u32) -> Option<&str> {
        self.source_contents
            .get(src_id as usize)
            .and_then(|x| x.as_ref().map(|x| &x[..]))
    }

    /// Checks if a given source ID has source contents available.
    pub fn has_source_contents(&self, src_id: u32) -> bool {
        self.get_source_contents(src_id).is_some()
    }

    /// Loads source contents from locally accessible files if referenced
    /// accordingly.  Returns the number of loaded source contents
    #[cfg(any(unix, windows, target_os = "redox"))]
    pub fn load_local_source_contents(&mut self, base_path: Option<&Path>) -> Result<usize> {
        let mut abs_path = env::current_dir()?;
        if let Some(path) = base_path {
            abs_path.push(path);
        }
        let base_url = Url::from_directory_path(&abs_path).unwrap();

        let mut to_read = vec![];
        for (source, &src_id) in self.source_map.iter() {
            if self.has_source_contents(src_id) {
                continue;
            }
            if let Some(path) = resolve_local_reference(&base_url, &source) {
                to_read.push((src_id, path));
            }
        }

        let rv = to_read.len();
        for (src_id, path) in to_read {
            if let Ok(mut f) = fs::File::open(&path) {
                let mut contents = String::new();
                if f.read_to_string(&mut contents).is_ok() {
                    self.set_source_contents(src_id, Some(&contents));
                }
            }
        }

        Ok(rv)
    }

    /// Registers a name with the builder and returns the name ID.
    pub fn add_name(&mut self, name: &str) -> u32 {
        let count = self.names.len() as u32;
        let id = *self.name_map.entry(name.into()).or_insert(count);
        if id == count {
            self.names.push(name.into());
        }
        id
    }

    /// Adds a new mapping to the builder.
    pub fn add(
        &mut self,
        dst_line: u32,
        dst_col: u32,
        src_line: u32,
        src_col: u32,
        source: Option<&str>,
        name: Option<&str>,
    ) -> RawToken {
        let src_id = match source {
            Some(source) => self.add_source(source),
            None => !0,
        };
        let name_id = match name {
            Some(name) => self.add_name(name),
            None => !0,
        };
        let raw = RawToken {
            dst_line,
            dst_col,
            src_line,
            src_col,
            src_id,
            name_id,
        };
        self.tokens.push(raw);
        raw
    }

    /// Shortcut for adding a new mapping based of an already existing token,
    /// optionally removing the name.
    pub fn add_token(&mut self, token: &Token<'_>, with_name: bool) -> RawToken {
        let name = if with_name { token.get_name() } else { None };
        self.add(
            token.get_dst_line(),
            token.get_dst_col(),
            token.get_src_line(),
            token.get_src_col(),
            token.get_source(),
            name,
        )
    }

    /// Strips common prefixes from the sources in the builder
    pub fn strip_prefixes<S: AsRef<str>>(&mut self, prefixes: &[S]) {
        for source in self.sources.iter_mut() {
            for prefix in prefixes {
                let mut prefix = prefix.as_ref().to_string();
                if !prefix.ends_with('/') {
                    prefix.push('/');
                }
                if source.starts_with(&prefix) {
                    *source = source[prefix.len()..].to_string();
                    break;
                }
            }
        }
    }

    /// Converts the builder into a sourcemap.
    pub fn into_sourcemap(self) -> SourceMap {
        let contents = if !self.source_contents.is_empty() {
            Some(self.source_contents)
        } else {
            None
        };
        SourceMap::new(self.file, self.tokens, self.names, self.sources, contents)
    }
}
<commit_msg>Add SourceMapBuilder.add_raw (#25)<commit_after>#![cfg_attr(not(any(unix, windows, target_os = "redox")), allow(unused_imports))]

use std::collections::HashMap;
use std::convert::AsRef;
use std::env;
use std::fs;
use std::io::Read;
use std::path::{Path, PathBuf};

use url::Url;

use crate::errors::Result;
use crate::types::{RawToken, SourceMap, Token};

/// Helper for sourcemap generation
///
/// This helper exists because generating and modifying `SourceMap`
/// objects is generally not very comfortable.  As a general aid this
/// type can help.
pub struct SourceMapBuilder {
    file: Option<String>,
    name_map: HashMap<String, u32>,
    names: Vec<String>,
    tokens: Vec<RawToken>,
    source_map: HashMap<String, u32>,
    sources: Vec<String>,
    source_contents: Vec<Option<String>>,
}

#[cfg(any(unix, windows, target_os = "redox"))]
fn resolve_local_reference(base: &Url, reference: &str) -> Option<PathBuf> {
    let url = match base.join(reference) {
        Ok(url) => {
            if url.scheme() != "file" {
                return None;
            }
            url
        }
        Err(_) => {
            return None;
        }
    };

    url.to_file_path().ok()
}

impl SourceMapBuilder {
    /// Creates a new source map builder and sets the file.
    pub fn new(file: Option<&str>) -> SourceMapBuilder {
        SourceMapBuilder {
            file: file.map(str::to_owned),
            name_map: HashMap::new(),
            names: vec![],
            tokens: vec![],
            source_map: HashMap::new(),
            sources: vec![],
            source_contents: vec![],
        }
    }

    /// Sets the file for the sourcemap (optional)
    pub fn set_file(&mut self, value: Option<&str>) {
        self.file = value.map(str::to_owned);
    }

    /// Returns the currently set file.
    pub fn get_file(&self) -> Option<&str> {
        self.file.as_ref().map(|x| &x[..])
    }

    /// Registers a new source with the builder and returns the source ID.
    pub fn add_source(&mut self, src: &str) -> u32 {
        let count = self.sources.len() as u32;
        let id = *self.source_map.entry(src.into()).or_insert(count);
        if id == count {
            self.sources.push(src.into());
        }
        id
    }

    /// Changes the source name for an already set source.
    pub fn set_source(&mut self, src_id: u32, src: &str) {
        assert!(src_id != !0, "Cannot set sources for tombstone source id");
        self.sources[src_id as usize] = src.to_string();
    }

    /// Looks up a source name for an ID.
    pub fn get_source(&self, src_id: u32) -> Option<&str> {
        self.sources.get(src_id as usize).map(|x| &x[..])
    }

    /// Sets the source contents for an already existing source.
    pub fn set_source_contents(&mut self, src_id: u32, contents: Option<&str>) {
        assert!(src_id != !0, "Cannot set sources for tombstone source id");
        if self.sources.len() > self.source_contents.len() {
            self.source_contents.resize(self.sources.len(), None);
        }
        self.source_contents[src_id as usize] = contents.map(str::to_owned);
    }

    /// Returns the current source contents for a source.
    pub fn get_source_contents(&self, src_id: u32) -> Option<&str> {
        self.source_contents
            .get(src_id as usize)
            .and_then(|x| x.as_ref().map(|x| &x[..]))
    }

    /// Checks if a given source ID has source contents available.
    pub fn has_source_contents(&self, src_id: u32) -> bool {
        self.get_source_contents(src_id).is_some()
    }

    /// Loads source contents from locally accessible files if referenced
    /// accordingly.  Returns the number of loaded source contents
    #[cfg(any(unix, windows, target_os = "redox"))]
    pub fn load_local_source_contents(&mut self, base_path: Option<&Path>) -> Result<usize> {
        let mut abs_path = env::current_dir()?;
        if let Some(path) = base_path {
            abs_path.push(path);
        }
        let base_url = Url::from_directory_path(&abs_path).unwrap();

        let mut to_read = vec![];
        for (source, &src_id) in self.source_map.iter() {
            if self.has_source_contents(src_id) {
                continue;
            }
            if let Some(path) = resolve_local_reference(&base_url, &source) {
                to_read.push((src_id, path));
            }
        }

        let rv = to_read.len();
        for (src_id, path) in to_read {
            if let Ok(mut f) = fs::File::open(&path) {
                let mut contents = String::new();
                if f.read_to_string(&mut contents).is_ok() {
                    self.set_source_contents(src_id, Some(&contents));
                }
            }
        }

        Ok(rv)
    }

    /// Registers a name with the builder and returns the name ID.
    pub fn add_name(&mut self, name: &str) -> u32 {
        let count = self.names.len() as u32;
        let id = *self.name_map.entry(name.into()).or_insert(count);
        if id == count {
            self.names.push(name.into());
        }
        id
    }

    /// Adds a new mapping to the builder.
    pub fn add(
        &mut self,
        dst_line: u32,
        dst_col: u32,
        src_line: u32,
        src_col: u32,
        source: Option<&str>,
        name: Option<&str>,
    ) -> RawToken {
        let src_id = match source {
            Some(source) => self.add_source(source),
            None => !0,
        };
        let name_id = match name {
            Some(name) => self.add_name(name),
            None => !0,
        };
        let raw = RawToken {
            dst_line,
            dst_col,
            src_line,
            src_col,
            src_id,
            name_id,
        };
        self.tokens.push(raw);
        raw
    }

    /// Adds a new mapping to the builder.
    pub fn add_raw(
        &mut self,
        dst_line: u32,
        dst_col: u32,
        src_line: u32,
        src_col: u32,
        source: Option<u32>,
        name: Option<u32>,
    ) -> RawToken {
        let src_id = match source {
            Some(source) => source,
            None => !0,
        };
        let name_id = match name {
            Some(name) => name,
            None => !0,
        };
        let raw = RawToken {
            dst_line,
            dst_col,
            src_line,
            src_col,
            src_id,
            name_id,
        };
        self.tokens.push(raw);
        raw
    }

    /// Shortcut for adding a new mapping based of an already existing token,
    /// optionally removing the name.
    pub fn add_token(&mut self, token: &Token<'_>, with_name: bool) -> RawToken {
        let name = if with_name { token.get_name() } else { None };
        self.add(
            token.get_dst_line(),
            token.get_dst_col(),
            token.get_src_line(),
            token.get_src_col(),
            token.get_source(),
            name,
        )
    }

    /// Strips common prefixes from the sources in the builder
    pub fn strip_prefixes<S: AsRef<str>>(&mut self, prefixes: &[S]) {
        for source in self.sources.iter_mut() {
            for prefix in prefixes {
                let mut prefix = prefix.as_ref().to_string();
                if !prefix.ends_with('/') {
                    prefix.push('/');
                }
                if source.starts_with(&prefix) {
                    *source = source[prefix.len()..].to_string();
                    break;
                }
            }
        }
    }

    /// Converts the builder into a sourcemap.
    pub fn into_sourcemap(self) -> SourceMap {
        let contents = if !self.source_contents.is_empty() {
            Some(self.source_contents)
        } else {
            None
        };
        SourceMap::new(self.file, self.tokens, self.names, self.sources, contents)
    }
}
<|endoftext|>


<commit_before>extern crate rustc_serialize;
extern crate chrono;
extern crate kuchiki;
extern crate regex;
extern crate url;

use std::io::Cursor;
use std::io::BufReader;

use kuchiki::traits::*;
use kuchiki::NodeRef;

use model::ListTopicItem;
use model::ShowItem;
use model::ShowReplyItem;
use model::UrlQueryItem;

use regex::Regex;
use url::Url;
use std::collections::HashMap;

pub struct Builder { }

impl Builder {
    pub fn new() -> Self {
        Builder {}
    }

    pub fn show_item(&mut self, document: &NodeRef,  url: &str) -> ShowItem {
        parse_show_item(&document, &url)
    }

    pub fn url_query_item(&mut self, url: &str) -> UrlQueryItem {
        parse_url_query_item(&url)
    }
}

fn parse_show_item(document: &NodeRef, url: &str) -> ShowItem {

    let url_query = parse_url_query_item(&url);

    let (title, reply_count) = {
        let repliers_tr = document.select(".repliers tr").unwrap().next().unwrap();
        let repliers_header = repliers_tr.as_node()
                                         .select(".repliers_header")
                                         .unwrap()
                                         .last()
                                         .unwrap();
        let divs = repliers_header.as_node().select("div").unwrap().collect::<Vec<_>>();

        let topic_data = divs.iter()
                             .enumerate()
                             .map(|(index, div)| {
                                 let s_trimmed = div.text_contents().trim().to_string();
                                 if index == 1 {
                                     let re = Regex::new(r"^(?P<count>\d+)個回應$").unwrap();
                                     let cap = re.captures(&s_trimmed).unwrap();
                                     // String::from(cap.name("count").unwrap_or("0"))
                                     cap.name("count").unwrap_or("0").to_string()
                                 } else {
                                     s_trimmed
                                 }
                             })
                             .collect::<Vec<_>>();

        if topic_data.len() < 2 {
            panic!("length of topic_data is invalid.")
        }

        (topic_data.get(0).unwrap().to_string(), // return as title
         topic_data.get(1).unwrap().to_string() /* return as reply_count */)
    };

    let (page, max_page) = {

        let page_select = document.select("select[name='page']").unwrap().last().unwrap();
        let page_str = page_select.as_node()
                                  .select("option[selected='selected']")
                                  .unwrap()
                                  .next()
                                  .unwrap();
        let max_page_str = page_select.as_node().select("option").unwrap().last().unwrap();

        let page = page_str.text_contents().trim().to_string().parse::<usize>().unwrap_or(0);
        let max_page = max_page_str.text_contents()
                                   .trim()
                                   .to_string()
                                   .parse::<usize>()
                                   .unwrap_or(0);

        (page, max_page)
    };

    let replies = parse_show_reply_items(&document);

    ShowItem {
        url_query: url_query,
        replies: replies,
        page: page,
        max_page: max_page,
        reply_count: String::from(reply_count),
        title: String::from(title),
    }
}

fn parse_show_reply_items(document: &NodeRef) -> Vec<ShowReplyItem> {

    let replies_data = document.select(".repliers tr[userid][username]")
                               .unwrap()
                               .collect::<Vec<_>>();

    replies_data.iter()
                .enumerate()
                .map(|(index, tr)| {

                    let tr_attrs = (&tr.attributes).borrow();
                    let userid = tr_attrs.get("userid").unwrap();
                    let username = tr_attrs.get("username").unwrap();

                    let content_elm = tr.as_node()
                                        .select(".repliers_right .ContentGrid")
                                        .unwrap()
                                        .next()
                                        .unwrap(); // first

                    let mut buff = Cursor::new(Vec::new());
                    let serialize_result = content_elm.as_node().serialize(&mut buff);
                    let vec = buff.into_inner();
                    let content = String::from_utf8(vec).unwrap();

                    let datatime = tr.as_node()
                                     .select(".repliers_right span")
                                     .unwrap()
                                     .last()
                                     .unwrap()
                                     .text_contents();

                    ShowReplyItem {
                        userid: String::from(userid),
                        username: String::from(username),
                        content: String::from(content),
                        published_at: String::from(datatime),
                    }
                })
                .collect::<Vec<_>>()
}

fn parse_url_query_item(url_str: &str) -> UrlQueryItem {
    let url = Url::parse(&url_str).unwrap();
    let query = url.query().unwrap();
    let re = Regex::new(r"(\\?|&)(?P<key>[^&=]+)=(?P<value>[^&]+)").unwrap();

    let (channel, message) = {

        let mut map = HashMap::new();

        for cap in re.captures_iter(query) {
            let key = cap.name("key").unwrap_or("").to_string();
            let value = cap.name("value").unwrap_or("").to_string();
            map.entry(key).or_insert(value);
        }

        if map.len() < 2 {
            panic!("length of map is invalid.")
        }

        (
            map.get("type").unwrap().to_string(),
            map.get("message").unwrap().to_string()
        )
    };

    UrlQueryItem {
        message: String::from(message)
    }

}
<commit_msg>[FIX] url can have 3 item (type, message, page)<commit_after>extern crate rustc_serialize;
extern crate chrono;
extern crate kuchiki;
extern crate regex;
extern crate url;

use std::io::Cursor;
use std::io::BufReader;

use kuchiki::traits::*;
use kuchiki::NodeRef;

use model::ListTopicItem;
use model::ShowItem;
use model::ShowReplyItem;
use model::UrlQueryItem;

use regex::Regex;
use url::Url;
use std::collections::HashMap;

pub struct Builder { }

impl Builder {
    pub fn new() -> Self {
        Builder {}
    }

    pub fn show_item(&mut self, document: &NodeRef,  url: &str) -> ShowItem {
        parse_show_item(&document, &url)
    }

    pub fn url_query_item(&mut self, url: &str) -> UrlQueryItem {
        parse_url_query_item(&url)
    }
}

fn parse_show_item(document: &NodeRef, url: &str) -> ShowItem {

    let url_query = parse_url_query_item(&url);

    let (title, reply_count) = {
        let repliers_tr = document.select(".repliers tr").unwrap().next().unwrap();
        let repliers_header = repliers_tr.as_node()
                                         .select(".repliers_header")
                                         .unwrap()
                                         .last()
                                         .unwrap();
        let divs = repliers_header.as_node().select("div").unwrap().collect::<Vec<_>>();

        let topic_data = divs.iter()
                             .enumerate()
                             .map(|(index, div)| {
                                 let s_trimmed = div.text_contents().trim().to_string();
                                 if index == 1 {
                                     let re = Regex::new(r"^(?P<count>\d+)個回應$").unwrap();
                                     let cap = re.captures(&s_trimmed).unwrap();
                                     // String::from(cap.name("count").unwrap_or("0"))
                                     cap.name("count").unwrap_or("0").to_string()
                                 } else {
                                     s_trimmed
                                 }
                             })
                             .collect::<Vec<_>>();

        if topic_data.len() < 2 {
            panic!("length of topic_data is invalid.")
        }

        (topic_data.get(0).unwrap().to_string(), // return as title
         topic_data.get(1).unwrap().to_string() /* return as reply_count */)
    };

    let (page, max_page) = {

        let page_select = document.select("select[name='page']").unwrap().last().unwrap();
        let page_str = page_select.as_node()
                                  .select("option[selected='selected']")
                                  .unwrap()
                                  .next()
                                  .unwrap();
        let max_page_str = page_select.as_node().select("option").unwrap().last().unwrap();

        let page = page_str.text_contents().trim().to_string().parse::<usize>().unwrap_or(0);
        let max_page = max_page_str.text_contents()
                                   .trim()
                                   .to_string()
                                   .parse::<usize>()
                                   .unwrap_or(0);

        (page, max_page)
    };

    let replies = parse_show_reply_items(&document);

    ShowItem {
        url_query: url_query,
        replies: replies,
        page: page,
        max_page: max_page,
        reply_count: String::from(reply_count),
        title: String::from(title),
    }
}

fn parse_show_reply_items(document: &NodeRef) -> Vec<ShowReplyItem> {

    let replies_data = document.select(".repliers tr[userid][username]")
                               .unwrap()
                               .collect::<Vec<_>>();

    replies_data.iter()
                .enumerate()
                .map(|(index, tr)| {

                    let tr_attrs = (&tr.attributes).borrow();
                    let userid = tr_attrs.get("userid").unwrap();
                    let username = tr_attrs.get("username").unwrap();

                    let content_elm = tr.as_node()
                                        .select(".repliers_right .ContentGrid")
                                        .unwrap()
                                        .next()
                                        .unwrap(); // first

                    let mut buff = Cursor::new(Vec::new());
                    let serialize_result = content_elm.as_node().serialize(&mut buff);
                    let vec = buff.into_inner();
                    let content = String::from_utf8(vec).unwrap();

                    let datatime = tr.as_node()
                                     .select(".repliers_right span")
                                     .unwrap()
                                     .last()
                                     .unwrap()
                                     .text_contents();

                    ShowReplyItem {
                        userid: String::from(userid),
                        username: String::from(username),
                        content: String::from(content),
                        published_at: String::from(datatime),
                    }
                })
                .collect::<Vec<_>>()
}

fn parse_url_query_item(url_str: &str) -> UrlQueryItem {
    let url = Url::parse(&url_str).unwrap();
    let query = url.query().unwrap();
    let re = Regex::new(r"(\\?|&)(?P<key>[^&=]+)=(?P<value>[^&]+)").unwrap();

    let (channel, message) = {

        let mut map = HashMap::new();

        for cap in re.captures_iter(query) {
            let key = cap.name("key").unwrap_or("").to_string();
            let value = cap.name("value").unwrap_or("").to_string();
            map.entry(key).or_insert(value);
        }

        // if map.len() < 2 {
        //     panic!("length of map is invalid.")
        // }

        (
            map.get("type").unwrap().to_string(),
            map.get("message").unwrap().to_string()
        )
    };

    UrlQueryItem {
        message: String::from(message)
    }

}
<|endoftext|>


<commit_before>use std::error::Error;
use std::collections::HashMap;
use emailaddress::EmailAddress;
use iron::Url;
use serde_json::de::from_reader;
use serde_json::value::Value;
use super::error::{BrokerError, BrokerResult};
use super::hyper::client::Client as HttpClient;
use super::hyper::header::ContentType as HyContentType;
use super::hyper::header::Headers;
use super::{Config, create_jwt};
use super::crypto::{session_id, verify_jws};
use super::store_cache::CacheKey;
use time::now_utc;
use url::percent_encoding::{utf8_percent_encode, QUERY_ENCODE_SET};

/// Macro used to extract a typed field from a JSON Value.
///
/// Will return from the caller with a BrokerError if the field is missing or its value is an
/// incompatible type. `descr` is used to format the error message.
///
/// ```
/// let foo = try_get_json_field!(value, "foo", as_str, "example document");
/// ```
macro_rules! try_get_json_field {
    ( $input:expr, $key:tt, $conv:ident, $descr:expr ) => {
        try!($input.find($key).and_then(|v| v.$conv()).ok_or_else(|| {
            BrokerError::Provider(format!("{} missing from {}", $key, $descr))
        }))
    }
}

/// Helper method to issue an OAuth authorization request.
///
/// When an authentication request comes in and matches one of the "famous"
/// identity providers configured in the `Config`, we redirect the client
/// to an Authentication Request URL, which we discover by reading the
/// provider's configured Discovery URL. We pass in the client ID we received
/// when pre-registering for the provider, as well as a callback URL which the
/// user will be redirected back to after confirming (or denying) the
/// Authentication Request. Included in the request is a nonce which we can
/// later use to definitively match the callback to this request.
pub fn request(app: &Config, email_addr: EmailAddress, client_id: &str, nonce: &str, redirect_uri: &Url)
               -> BrokerResult<Url> {

    let session = session_id(&email_addr, client_id);

    // Store the nonce and the RP's `redirect_uri` in Redis for use in the
    // callback handler.
    try!(app.store.store_session(&session, &[
        ("type", "oidc"),
        ("email", &email_addr.to_string()),
        ("client_id", client_id),
        ("nonce", nonce),
        ("redirect", &redirect_uri.to_string()),
    ]));

    let client = HttpClient::new();

    // Retrieve the provider's Discovery document and extract the
    // `authorization_endpoint` from it.
    let domain = &email_addr.domain;
    let provider = &app.providers[domain];
    let config_obj: Value = try!(
        app.store.cache.fetch_json_url(
            &app.store,
            CacheKey::Discovery { domain: &email_addr.domain },
            &client,
            &provider.discovery_url
        ).map_err(|e| {
            BrokerError::Provider(format!("could not fetch {}'s discovery document: {}",
                                          domain, e.description()))
        })
    );
    let descr = format!("{}'s discovery document", domain);
    let authz_base = try_get_json_field!(config_obj, "authorization_endpoint", as_str, descr);

    // Create the URL to redirect to, properly escaping all parameters.
    Url::parse(&vec![
        authz_base,
        "?",
        "client_id=",
        &utf8_percent_encode(&provider.client_id, QUERY_ENCODE_SET).to_string(),
        "&response_type=code",
        "&scope=",
        &utf8_percent_encode("openid email", QUERY_ENCODE_SET).to_string(),
        "&redirect_uri=",
        &utf8_percent_encode(&format!("{}/callback", &app.public_url),
                             QUERY_ENCODE_SET).to_string(),
        "&state=",
        &utf8_percent_encode(&session, QUERY_ENCODE_SET).to_string(),
        "&login_hint=",
        &utf8_percent_encode(&email_addr.to_string(), QUERY_ENCODE_SET).to_string(),
    ].join("")).map_err(|_| {
        BrokerError::Provider(format!("failed to build valid authorization URL from {}'s 'authorization_endpoint'", domain))
    })

}

pub fn canonicalize_google(email: String) -> String {
    let at = email.find('@').unwrap();
    let (user, domain) = email.split_at(at);
    let domain = &domain[1..];
    let user = &user.replace(".", ""); // Ignore dots

    // Trim plus addresses
    let user = match user.find('+') {
        Some(pos) => user.split_at(pos).0,
        None => user,
    };

    // Normalize googlemail.com to gmail.com
    let domain = if domain == "googlemail.com" {
        "gmail.com"
    } else {
        domain
    };
    user.to_string() + "@" + domain
}

pub fn canonicalized(email: &str) -> String {
    let normalized = email.to_lowercase();
    if normalized.ends_with("@gmail.com") || normalized.ends_with("@googlemail.com") {
        canonicalize_google(normalized)
    } else {
        normalized
    }
}

/// Helper method to verify OAuth authentication result.
///
/// Match the returned email address and nonce against our Redis data, then
/// extract the identity token returned by the provider and verify it. Return
/// an identity token for the RP if successful, or an error message otherwise.
pub fn verify(app: &Config, stored: &HashMap<String, String>, code: &str)
              -> BrokerResult<(String, String)> {

    let email_addr = EmailAddress::new(&stored["email"]).unwrap();
    let origin = &stored["client_id"];
    let nonce = &stored["nonce"];

    let client = HttpClient::new();

    // Request the provider's Discovery document to get the
    // `token_endpoint` and `jwks_uri` values from it. TODO: save these
    // when requesting the Discovery document in the `oauth_request()`
    // function, and/or cache them by provider host.
    let domain = &email_addr.domain;
    let provider = &app.providers[domain];
    let config_obj: Value = try!(
        app.store.cache.fetch_json_url(
            &app.store,
            CacheKey::Discovery { domain: domain },
            &client,
            &provider.discovery_url
        ).map_err(|e| {
            BrokerError::Provider(format!("could not fetch {}'s discovery document: {}",
                                          domain, e.description()))
        })
    );
    let descr = format!("{}'s discovery document", domain);
    let token_url = try_get_json_field!(config_obj, "token_endpoint", as_str, descr);
    let jwks_url = try_get_json_field!(config_obj, "jwks_uri", as_str, descr);

    // Create form data for the Token Request, where we exchange the code
    // received in this callback request for an identity token (while
    // proving our identity by passing the client secret value).
    let body: String = vec![
        "code=",
        &utf8_percent_encode(code, QUERY_ENCODE_SET).to_string(),
        "&client_id=",
        &utf8_percent_encode(&provider.client_id, QUERY_ENCODE_SET).to_string(),
        "&client_secret=",
        &utf8_percent_encode(&provider.secret, QUERY_ENCODE_SET).to_string(),
        "&redirect_uri=",
        &utf8_percent_encode(&format!("{}/callback", &app.public_url),
                             QUERY_ENCODE_SET).to_string(),
        "&grant_type=authorization_code",
    ].join("");

    // Send the Token Request and extract the `id_token` from the response.
    let token_obj: Value = {
        let mut headers = Headers::new();
        headers.set(HyContentType::form_url_encoded());
        try!(
            client.post(token_url).headers(headers).body(&body).send()
                .map_err(BrokerError::Http)
                .and_then(|rsp| from_reader(rsp).map_err(|_| {
                    BrokerError::Provider("failed to parse response as JSON".to_string())
                }))
                .map_err(|e| {
                    BrokerError::Provider(format!("{} token request failed: {}",
                                                  domain, e.description()))
                })
        )
    };
    let descr = format!("{}'s token response", domain);
    let id_token = try_get_json_field!(token_obj, "id_token", as_str, descr);

    // Grab the keys from the provider, then verify the signature.
    let jwt_payload = {
        let keys_obj: Value = try!(
            app.store.cache.fetch_json_url(
                &app.store,
                CacheKey::KeySet { domain: domain },
                &client,
                &jwks_url
            ).map_err(|e| {
                BrokerError::Provider(format!("could not fetch {}'s keys: {}",
                                              domain, e.description()))
            })
        );
        try!(
            verify_jws(id_token, &keys_obj).map_err(|_| {
                BrokerError::Provider(format!("could not verify the token received from {}", domain))
            })
        )
    };

    // Verify the claims contained in the token.
    let descr = format!("{}'s token payload", domain);
    let iss = try_get_json_field!(jwt_payload, "iss", as_str, descr);
    let aud = try_get_json_field!(jwt_payload, "aud", as_str, descr);
    let token_addr = try_get_json_field!(jwt_payload, "email", as_str, descr);
    let exp = try_get_json_field!(jwt_payload, "exp", as_i64, descr);
    let issuer_origin = vec!["https://", &provider.issuer_domain].join("");
    assert!(iss == provider.issuer_domain || iss == issuer_origin);
    assert!(aud == provider.client_id);
    assert!(canonicalized(token_addr) == canonicalized(&email_addr.to_string()));
    let now = now_utc().to_timespec().sec;
    assert!(now < exp);

    // If everything is okay, build a new identity token and send it
    // to the relying party.
    let id_token = create_jwt(app, &email_addr.to_string(), origin, nonce);
    let redirect = &stored["redirect"];
    Ok((id_token, redirect.to_string()))

}

#[cfg(test)]
mod tests {
    use super::canonicalized;
    #[test]
    fn test_canonicalized_google() {
        assert_eq!(canonicalized("example.foo+bar@example.com"),
                   "example.foo+bar@example.com");
        assert_eq!(canonicalized("example@gmail.com"),
                   "example@gmail.com");
        assert_eq!(canonicalized("example@googlemail.com"),
                   "example@gmail.com");
        assert_eq!(canonicalized("example.foo@gmail.com"),
                   "examplefoo@gmail.com");
        assert_eq!(canonicalized("example+bar@gmail.com"),
                   "example@gmail.com");
        assert_eq!(canonicalized("example.foo+bar@googlemail.com"),
                   "examplefoo@gmail.com");
    }
}
<commit_msg>Lowercase domains when looking up provider info<commit_after>use std::error::Error;
use std::collections::HashMap;
use emailaddress::EmailAddress;
use iron::Url;
use serde_json::de::from_reader;
use serde_json::value::Value;
use super::error::{BrokerError, BrokerResult};
use super::hyper::client::Client as HttpClient;
use super::hyper::header::ContentType as HyContentType;
use super::hyper::header::Headers;
use super::{Config, create_jwt};
use super::crypto::{session_id, verify_jws};
use super::store_cache::CacheKey;
use time::now_utc;
use url::percent_encoding::{utf8_percent_encode, QUERY_ENCODE_SET};

/// Macro used to extract a typed field from a JSON Value.
///
/// Will return from the caller with a BrokerError if the field is missing or its value is an
/// incompatible type. `descr` is used to format the error message.
///
/// ```
/// let foo = try_get_json_field!(value, "foo", as_str, "example document");
/// ```
macro_rules! try_get_json_field {
    ( $input:expr, $key:tt, $conv:ident, $descr:expr ) => {
        try!($input.find($key).and_then(|v| v.$conv()).ok_or_else(|| {
            BrokerError::Provider(format!("{} missing from {}", $key, $descr))
        }))
    }
}

/// Helper method to issue an OAuth authorization request.
///
/// When an authentication request comes in and matches one of the "famous"
/// identity providers configured in the `Config`, we redirect the client
/// to an Authentication Request URL, which we discover by reading the
/// provider's configured Discovery URL. We pass in the client ID we received
/// when pre-registering for the provider, as well as a callback URL which the
/// user will be redirected back to after confirming (or denying) the
/// Authentication Request. Included in the request is a nonce which we can
/// later use to definitively match the callback to this request.
pub fn request(app: &Config, email_addr: EmailAddress, client_id: &str, nonce: &str, redirect_uri: &Url)
               -> BrokerResult<Url> {

    let session = session_id(&email_addr, client_id);

    // Store the nonce and the RP's `redirect_uri` in Redis for use in the
    // callback handler.
    try!(app.store.store_session(&session, &[
        ("type", "oidc"),
        ("email", &email_addr.to_string()),
        ("client_id", client_id),
        ("nonce", nonce),
        ("redirect", &redirect_uri.to_string()),
    ]));

    let client = HttpClient::new();

    // Retrieve the provider's Discovery document and extract the
    // `authorization_endpoint` from it.
    let domain = &email_addr.domain.to_lowercase();
    let provider = &app.providers[domain];
    let config_obj: Value = try!(
        app.store.cache.fetch_json_url(
            &app.store,
            CacheKey::Discovery { domain: &domain },
            &client,
            &provider.discovery_url
        ).map_err(|e| {
            BrokerError::Provider(format!("could not fetch {}'s discovery document: {}",
                                          domain, e.description()))
        })
    );
    let descr = format!("{}'s discovery document", domain);
    let authz_base = try_get_json_field!(config_obj, "authorization_endpoint", as_str, descr);

    // Create the URL to redirect to, properly escaping all parameters.
    Url::parse(&vec![
        authz_base,
        "?",
        "client_id=",
        &utf8_percent_encode(&provider.client_id, QUERY_ENCODE_SET).to_string(),
        "&response_type=code",
        "&scope=",
        &utf8_percent_encode("openid email", QUERY_ENCODE_SET).to_string(),
        "&redirect_uri=",
        &utf8_percent_encode(&format!("{}/callback", &app.public_url),
                             QUERY_ENCODE_SET).to_string(),
        "&state=",
        &utf8_percent_encode(&session, QUERY_ENCODE_SET).to_string(),
        "&login_hint=",
        &utf8_percent_encode(&email_addr.to_string(), QUERY_ENCODE_SET).to_string(),
    ].join("")).map_err(|_| {
        BrokerError::Provider(format!("failed to build valid authorization URL from {}'s 'authorization_endpoint'", domain))
    })

}

pub fn canonicalize_google(email: String) -> String {
    let at = email.find('@').unwrap();
    let (user, domain) = email.split_at(at);
    let domain = &domain[1..];
    let user = &user.replace(".", ""); // Ignore dots

    // Trim plus addresses
    let user = match user.find('+') {
        Some(pos) => user.split_at(pos).0,
        None => user,
    };

    // Normalize googlemail.com to gmail.com
    let domain = if domain == "googlemail.com" {
        "gmail.com"
    } else {
        domain
    };
    user.to_string() + "@" + domain
}

pub fn canonicalized(email: &str) -> String {
    let normalized = email.to_lowercase();
    if normalized.ends_with("@gmail.com") || normalized.ends_with("@googlemail.com") {
        canonicalize_google(normalized)
    } else {
        normalized
    }
}

/// Helper method to verify OAuth authentication result.
///
/// Match the returned email address and nonce against our Redis data, then
/// extract the identity token returned by the provider and verify it. Return
/// an identity token for the RP if successful, or an error message otherwise.
pub fn verify(app: &Config, stored: &HashMap<String, String>, code: &str)
              -> BrokerResult<(String, String)> {

    let email_addr = EmailAddress::new(&stored["email"]).unwrap();
    let origin = &stored["client_id"];
    let nonce = &stored["nonce"];

    let client = HttpClient::new();

    // Request the provider's Discovery document to get the
    // `token_endpoint` and `jwks_uri` values from it. TODO: save these
    // when requesting the Discovery document in the `oauth_request()`
    // function, and/or cache them by provider host.
    let domain = &email_addr.domain.to_lowercase();
    let provider = &app.providers[domain];
    let config_obj: Value = try!(
        app.store.cache.fetch_json_url(
            &app.store,
            CacheKey::Discovery { domain: domain },
            &client,
            &provider.discovery_url
        ).map_err(|e| {
            BrokerError::Provider(format!("could not fetch {}'s discovery document: {}",
                                          domain, e.description()))
        })
    );
    let descr = format!("{}'s discovery document", domain);
    let token_url = try_get_json_field!(config_obj, "token_endpoint", as_str, descr);
    let jwks_url = try_get_json_field!(config_obj, "jwks_uri", as_str, descr);

    // Create form data for the Token Request, where we exchange the code
    // received in this callback request for an identity token (while
    // proving our identity by passing the client secret value).
    let body: String = vec![
        "code=",
        &utf8_percent_encode(code, QUERY_ENCODE_SET).to_string(),
        "&client_id=",
        &utf8_percent_encode(&provider.client_id, QUERY_ENCODE_SET).to_string(),
        "&client_secret=",
        &utf8_percent_encode(&provider.secret, QUERY_ENCODE_SET).to_string(),
        "&redirect_uri=",
        &utf8_percent_encode(&format!("{}/callback", &app.public_url),
                             QUERY_ENCODE_SET).to_string(),
        "&grant_type=authorization_code",
    ].join("");

    // Send the Token Request and extract the `id_token` from the response.
    let token_obj: Value = {
        let mut headers = Headers::new();
        headers.set(HyContentType::form_url_encoded());
        try!(
            client.post(token_url).headers(headers).body(&body).send()
                .map_err(BrokerError::Http)
                .and_then(|rsp| from_reader(rsp).map_err(|_| {
                    BrokerError::Provider("failed to parse response as JSON".to_string())
                }))
                .map_err(|e| {
                    BrokerError::Provider(format!("{} token request failed: {}",
                                                  domain, e.description()))
                })
        )
    };
    let descr = format!("{}'s token response", domain);
    let id_token = try_get_json_field!(token_obj, "id_token", as_str, descr);

    // Grab the keys from the provider, then verify the signature.
    let jwt_payload = {
        let keys_obj: Value = try!(
            app.store.cache.fetch_json_url(
                &app.store,
                CacheKey::KeySet { domain: domain },
                &client,
                &jwks_url
            ).map_err(|e| {
                BrokerError::Provider(format!("could not fetch {}'s keys: {}",
                                              domain, e.description()))
            })
        );
        try!(
            verify_jws(id_token, &keys_obj).map_err(|_| {
                BrokerError::Provider(format!("could not verify the token received from {}", domain))
            })
        )
    };

    // Verify the claims contained in the token.
    let descr = format!("{}'s token payload", domain);
    let iss = try_get_json_field!(jwt_payload, "iss", as_str, descr);
    let aud = try_get_json_field!(jwt_payload, "aud", as_str, descr);
    let token_addr = try_get_json_field!(jwt_payload, "email", as_str, descr);
    let exp = try_get_json_field!(jwt_payload, "exp", as_i64, descr);
    let issuer_origin = vec!["https://", &provider.issuer_domain].join("");
    assert!(iss == provider.issuer_domain || iss == issuer_origin);
    assert!(aud == provider.client_id);
    assert!(canonicalized(token_addr) == canonicalized(&email_addr.to_string()));
    let now = now_utc().to_timespec().sec;
    assert!(now < exp);

    // If everything is okay, build a new identity token and send it
    // to the relying party.
    let id_token = create_jwt(app, &email_addr.to_string(), origin, nonce);
    let redirect = &stored["redirect"];
    Ok((id_token, redirect.to_string()))

}

#[cfg(test)]
mod tests {
    use super::canonicalized;
    #[test]
    fn test_canonicalized_google() {
        assert_eq!(canonicalized("example.foo+bar@example.com"),
                   "example.foo+bar@example.com");
        assert_eq!(canonicalized("example@gmail.com"),
                   "example@gmail.com");
        assert_eq!(canonicalized("example@googlemail.com"),
                   "example@gmail.com");
        assert_eq!(canonicalized("example.foo@gmail.com"),
                   "examplefoo@gmail.com");
        assert_eq!(canonicalized("example+bar@gmail.com"),
                   "example@gmail.com");
        assert_eq!(canonicalized("example.foo+bar@googlemail.com"),
                   "examplefoo@gmail.com");
    }

    #[test]
    fn test_canonicalized_casing() {
        assert_eq!(canonicalized("EXAMPLE.FOO+BAR@EXAMPLE.COM"),
                   "example.foo+bar@example.com");
        assert_eq!(canonicalized("EXAMPLE@GOOGLEMAIL.COM"),
                   "example@gmail.com");
    }
}
<|endoftext|>