<commit_before>//! A Single-Producer, Multiple-Consumer queue.

use std::sync::mpsc::{channel, Receiver, RecvError, Sender, SendError};
use std::sync::{Arc, Mutex, MutexGuard};
use std::fmt;
use std::any::Any;
use std::error::Error;

#[derive(PartialEq, Eq, Clone, Copy, Debug)]
/// Error from broadcast module.
pub enum BroadcastError<T> {
    /// Send error
    SendError(T),
    /// Receive error
    RecvError,
}

impl<T: fmt::Display> fmt::Display for BroadcastError<T> {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            BroadcastError::SendError(ref t) =>
                write!(fmt, "could not send data on channel: {}", t),
            BroadcastError::RecvError =>
                write!(fmt, "could not receive data on channel"),
        }
    }
}

impl<T: Send + fmt::Display + fmt::Debug + Any> Error for BroadcastError<T> {
    fn description(&self) -> &str {
        match *self {
            BroadcastError::SendError(_) => "could not send data on channel",
            BroadcastError::RecvError => "could not receive data on channel",
        }
    }

    fn cause(&self) -> Option<&Error> {
        None
    }
}

impl<T> From<SendError<T>> for BroadcastError<T> {
    fn from(err: SendError<T>) -> BroadcastError<T> {
        let SendError(data) = err;
        BroadcastError::SendError(data)
    }
}

impl<T> From<RecvError> for BroadcastError<T> {
    fn from(_err: RecvError) -> BroadcastError<T> {
        BroadcastError::RecvError
    }
}

/// Struct that sends message on a broadcast pattern.
pub struct Broadcast<T> {
    inner: Arc<Inner<T>>,
}

impl<T> Broadcast<T> {
    /// Create a new Broadcast struct.
    pub fn new() -> Broadcast<T> {
        let inner = Arc::new(Inner { senders: Mutex::new(Vec::new()) });
        Broadcast { inner: inner }
    }

    /// Create a Consumer that listens to messages from the Broadcaster.
    pub fn consume(&self) -> Consumer<T> {
        let (b, c) = channel();
        self.inner.add_sender(b);
        Consumer { inner: self.inner.clone(), receiver: c }
    }
}

impl<T: Clone> Broadcast<T> {
    /// Send a message on the broadcast.
    pub fn send(&self, data: T) -> Result<(), BroadcastError<T>> {
        let guard = self.inner.read_senders();
        for s in guard.iter() {
            try!(s.send(data.clone()));
        }

        Ok(())
    }
}

struct Inner<T> {
    senders: Mutex<Vec<Sender<T>>>,
}

impl<T> Inner<T> {
    fn read_senders<'a>(&'a self) -> MutexGuard<'a, Vec<Sender<T>>> {
        self.senders.lock().unwrap()
    }

    fn add_sender(&self, sender: Sender<T>) {
        let mut vec = self.senders.lock().unwrap();
        vec.push(sender);
    }
}

/// Struct that receives messages from Broadcast.
pub struct Consumer<T> {
    inner: Arc<Inner<T>>,
    receiver: Receiver<T>,
}

impl<T> Consumer<T> {
    /// Receive a message from the Broadcast.
    ///
    /// This function will block.
    pub fn recv(&self) -> Result<T, BroadcastError<T>> {
        let data = try!(self.receiver.recv());
        Ok(data)
    }
}

impl<T> Clone for Consumer<T> {
    fn clone(&self) -> Self {
        let (s, r) = channel();
        self.inner.add_sender(s);
        Consumer {
            inner: self.inner.clone(),
            receiver: r,
        }
    }
}

/// Create a (Broadcast<T>, Consumer<T>) pair.
pub fn broadcast_channel<T: Clone>() -> (Broadcast<T>, Consumer<T>) {
    let broadcast = Broadcast::new();
    let consumer = broadcast.consume();
    (broadcast, consumer)
}

#[cfg(test)]
mod test {
    use broadcast::broadcast_channel;
    use super::Inner;

    use std::sync::{Arc, Mutex};
    use std::sync::mpsc::{channel};
    use std::thread::spawn;

    #[test]
    fn inner_iterator() {
        let (s1, r1) = channel();
        let (s2, r2) = channel();
        let inner = Arc::new(Inner { senders: Mutex::new(vec!(s1, s2)) });
        let guard = inner.read_senders();
        for s in guard.iter() {
            assert!(s.send(10u8).is_ok());
        }

        assert_eq!(r1.recv().unwrap(), 10u8);
        assert_eq!(r2.recv().unwrap(), 10u8);
    }

    #[test]
    fn sends_to_multiple_consumers() {
        let (p, c) = broadcast_channel();
        let c2 = c.clone();

        let res = p.send(9u8);
        assert!(res.is_ok());

        let res = c.recv();
        assert!(res.is_ok());
        assert_eq!(res.unwrap(), 9u8);

        let res = c2.recv();
        assert!(res.is_ok());
        assert_eq!(res.unwrap(), 9u8)
    }

    #[test]
    fn test_send_threads() {
        let (p, c1) = broadcast_channel();
        let c2 = c1.clone();

        let (s1, r1) = channel();
        let (s2, r2) = channel();

        let _thread = spawn(move || {
            assert_eq!(c1.recv().unwrap(), 9u8);
            s1.send(10u8).unwrap();
        });

        let _thread = spawn(move || {
            assert_eq!(c2.recv().unwrap(), 9u8);
            s2.send(10u8).unwrap();
        });

        assert!(p.send(9u8).is_ok());
        assert!(r1.recv().is_ok());
        assert!(r2.recv().is_ok());
    }
}
<commit_msg>Clone sender to ensure no allocations are made during send() call<commit_after>//! A Single-Producer, Multiple-Consumer queue.

use std::sync::mpsc::{channel, Receiver, RecvError, Sender, SendError};
use std::sync::{Arc, Mutex, MutexGuard};
use std::fmt;
use std::any::Any;
use std::error::Error;

#[derive(PartialEq, Eq, Clone, Copy, Debug)]
/// Error from broadcast module.
pub enum BroadcastError<T> {
    /// Send error
    SendError(T),
    /// Receive error
    RecvError,
}

impl<T: fmt::Display> fmt::Display for BroadcastError<T> {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            BroadcastError::SendError(ref t) =>
                write!(fmt, "could not send data on channel: {}", t),
            BroadcastError::RecvError =>
                write!(fmt, "could not receive data on channel"),
        }
    }
}

impl<T: Send + fmt::Display + fmt::Debug + Any> Error for BroadcastError<T> {
    fn description(&self) -> &str {
        match *self {
            BroadcastError::SendError(_) => "could not send data on channel",
            BroadcastError::RecvError => "could not receive data on channel",
        }
    }

    fn cause(&self) -> Option<&Error> {
        None
    }
}

impl<T> From<SendError<T>> for BroadcastError<T> {
    fn from(err: SendError<T>) -> BroadcastError<T> {
        let SendError(data) = err;
        BroadcastError::SendError(data)
    }
}

impl<T> From<RecvError> for BroadcastError<T> {
    fn from(_err: RecvError) -> BroadcastError<T> {
        BroadcastError::RecvError
    }
}

/// Struct that sends message on a broadcast pattern.
pub struct Broadcast<T> {
    inner: Arc<Inner<T>>,
}

impl<T> Broadcast<T> {
    /// Create a new Broadcast struct.
    pub fn new() -> Broadcast<T> {
        let inner = Arc::new(Inner { senders: Mutex::new(Vec::new()) });
        Broadcast { inner: inner }
    }

    /// Create a Consumer that listens to messages from the Broadcaster.
    pub fn consume(&self) -> Consumer<T> {
        let (b, c) = channel();
        self.inner.add_sender(b);
        Consumer { inner: self.inner.clone(), receiver: c }
    }
}

impl<T: Clone> Broadcast<T> {
    /// Send a message on the broadcast.
    pub fn send(&self, data: T) -> Result<(), BroadcastError<T>> {
        let guard = self.inner.read_senders();
        for s in guard.iter() {
            try!(s.send(data.clone()));
        }

        Ok(())
    }
}

struct Inner<T> {
    senders: Mutex<Vec<Sender<T>>>,
}

impl<T> Inner<T> {
    fn read_senders<'a>(&'a self) -> MutexGuard<'a, Vec<Sender<T>>> {
        self.senders.lock().unwrap()
    }

    fn add_sender(&self, sender: Sender<T>) {
        let mut vec = self.senders.lock().unwrap();
        // Clone the sender to ensure we receive a shared::Packet that
        // doese not perform allocations on sending
        vec.push(sender.clone());
    }
}

/// Struct that receives messages from Broadcast.
pub struct Consumer<T> {
    inner: Arc<Inner<T>>,
    receiver: Receiver<T>,
}

impl<T> Consumer<T> {
    /// Receive a message from the Broadcast.
    ///
    /// This function will block.
    pub fn recv(&self) -> Result<T, BroadcastError<T>> {
        let data = try!(self.receiver.recv());
        Ok(data)
    }
}

impl<T> Clone for Consumer<T> {
    fn clone(&self) -> Self {
        let (s, r) = channel();
        self.inner.add_sender(s);
        Consumer {
            inner: self.inner.clone(),
            receiver: r,
        }
    }
}

/// Create a (Broadcast<T>, Consumer<T>) pair.
pub fn broadcast_channel<T: Clone>() -> (Broadcast<T>, Consumer<T>) {
    let broadcast = Broadcast::new();
    let consumer = broadcast.consume();
    (broadcast, consumer)
}

#[cfg(test)]
mod test {
    use broadcast::broadcast_channel;
    use super::Inner;

    use std::sync::{Arc, Mutex};
    use std::sync::mpsc::{channel};
    use std::thread::spawn;

    #[test]
    fn inner_iterator() {
        let (s1, r1) = channel();
        let (s2, r2) = channel();
        let inner = Arc::new(Inner { senders: Mutex::new(vec!(s1, s2)) });
        let guard = inner.read_senders();
        for s in guard.iter() {
            assert!(s.send(10u8).is_ok());
        }

        assert_eq!(r1.recv().unwrap(), 10u8);
        assert_eq!(r2.recv().unwrap(), 10u8);
    }

    #[test]
    fn sends_to_multiple_consumers() {
        let (p, c) = broadcast_channel();
        let c2 = c.clone();

        let res = p.send(9u8);
        assert!(res.is_ok());

        let res = c.recv();
        assert!(res.is_ok());
        assert_eq!(res.unwrap(), 9u8);

        let res = c2.recv();
        assert!(res.is_ok());
        assert_eq!(res.unwrap(), 9u8)
    }

    #[test]
    fn test_send_threads() {
        let (p, c1) = broadcast_channel();
        let c2 = c1.clone();

        let (s1, r1) = channel();
        let (s2, r2) = channel();

        let _thread = spawn(move || {
            assert_eq!(c1.recv().unwrap(), 9u8);
            s1.send(10u8).unwrap();
        });

        let _thread = spawn(move || {
            assert_eq!(c2.recv().unwrap(), 9u8);
            s2.send(10u8).unwrap();
        });

        assert!(p.send(9u8).is_ok());
        assert!(r1.recv().is_ok());
        assert!(r2.recv().is_ok());
    }
}
<|endoftext|>


<commit_before>#[doc(hidden)];

//! Utilities for handling character classes.

use std::char;
use std::cmp::{min, max};

/// A range of codepoints.
pub type Range = (char, char);

/// A character class is a non-empty collection of ranges.
#[deriving(Clone)]
pub enum CharClass {
    priv CCOwned(~[Range]),
    priv CCStatic(&'static [Range])
}

impl CharClass {
    /// Create a class from a vector of ranges.
    pub fn new(mut r: ~[Range]) -> CharClass {
        if r.iter().any(|&(lo, hi)| lo > hi) {
            fail!("invalid range");
        }

        r.sort_by(|a, b| {
            // Sort by start value ascending, end value descending
            match a.lo().cmp(&b.lo()) {
                Equal => b.hi().cmp(&a.hi()),
                o => o
            }
        });

        if r.len() > 0 {
            // Coalesce overlapping ranges in place
            let mut cursor = 0;
            for i in range(1, r.len()) {
                if r[cursor].hi() == char::MAX {
                    break
                } else if r[i].lo() <= next_char(r[cursor].hi()) {
                    // Merge the two ranges
                    r[cursor] = (
                        min(r[cursor].lo(), r[i].lo()),
                        max(r[cursor].hi(), r[i].hi()));
                } else {
                    // Create a new range
                    cursor += 1;
                    r[cursor] = r[i];
                }
            }
            r.truncate(cursor + 1); r.shrink_to_fit();
            CCOwned(r)
        } else {
            fail!("char class cannot be empty")
        }
    }

    /// Create a class matching a single code point.
    pub fn from_char(c: char) -> CharClass {
        CharClass::new(~[(c, c)])
    }

    /// Create a class from a range of code points.
    pub fn from_range(lo: char, hi: char) -> CharClass {
        CharClass::new(~[(lo, hi)])
    }

    /// Combine several classes into one that subsumes them all.
    pub fn combine(classes: &[CharClass]) -> CharClass {
        let mut ranges: ~[Range] = ~[];
        for cc in classes.iter() {
            for r in cc.ranges().iter() {
                ranges.push(*r);
            }
        }
        CharClass::new(ranges)
    }

    ///
    /// Get the list of ranges contained in the character class.
    ///
    /// The returned vector is always sorted, with no overlapping
    /// ranges.
    ///
    pub fn ranges<'a>(&'a self) -> &'a [Range] {
        match *self {
            CCOwned(ref r) => {
                let r_borrow: &'a [Range] = *r;
                r_borrow
            },
            CCStatic(r) => {
                let r_borrow: &'a [Range] = r;
                r_borrow
            }
        }
    }

    /// If the class matches a single code point, return it; otherwise
    /// return `None`.
    pub fn to_char(&self) -> Option<char> {
        match self.ranges() {
            [(lo, hi)] if lo == hi => Some(lo),
            _ => None
        }
    }

    /// Determine if the character is contained in the class.
    pub fn includes(&self, c: char) -> bool {
        self.ranges().bsearch(|&(lo, hi)| {
            if c < lo {
                Greater
            } else if hi < c {
                Less
            } else {
                Equal
            }
        }).is_some()
    }

    /// Return the negation of the character class.
    pub fn negate(&self) -> CharClass {
        let ranges = self.ranges();
        let mut result: ~[Range] = ~[];

        let first_lo = ranges[0].lo();
        if first_lo != '\0' {
            result.push(('\0', prev_char(first_lo)));
        }

        let mut last_hi = ranges[0].hi();
        for &(lo, hi) in ranges.slice_from(1).iter() {
            result.push((next_char(last_hi), prev_char(lo)));
            last_hi = hi;
        }

        if last_hi != char::MAX {
            result.push((next_char(last_hi), char::MAX));
        }

        CCOwned(result)
    }
}

priv trait RangeUtils {
    fn lo(&self) -> char;
    fn hi(&self) -> char;
}

impl RangeUtils for (char, char) {
    #[inline]
    fn lo(&self) -> char {
        let &(lo, _) = self;
        lo
    }

    #[inline]
    fn hi(&self) -> char {
        let &(_, hi) = self;
        hi
    }
}

#[inline]
fn prev_char(c: char) -> char {
    char::from_u32(c as u32 - 1).unwrap()
}

#[inline]
fn next_char(c: char) -> char {
    char::from_u32(c as u32 + 1).unwrap()
}

#[cfg(test)]
mod test {
    use std::char;
    use super::CharClass;

    #[test]
    #[should_fail]
    fn invalid_range() {
        let _ = CharClass::new(~[('b', 'a')]);
    }

    #[test]
    #[should_fail]
    fn empty_class() {
        let _ = CharClass::new(~[]);
    }

    #[test]
    fn class_sorted() {
        let c = CharClass::new(~[('y', 'z'), ('a', 'b')]);
        assert_eq!(c.ranges(), [('a', 'b'), ('y', 'z')]);
    }

    #[test]
    fn class_optimize() {
        let c1 = CharClass::new(~[('a', 'b'), ('c', 'd')]);  // touching
        let c2 = CharClass::new(~[('a', 'c'), ('b', 'd')]);  // overlapping
        let c3 = CharClass::new(~[('a', 'd'), ('b', 'c')]);  // contained
        assert_eq!(c1.ranges(), [('a', 'd')]);
        assert_eq!(c2.ranges(), [('a', 'd')]);
        assert_eq!(c3.ranges(), [('a', 'd')]);
    }

    #[test]
    fn class_includes() {
        // This will fail if the bsearch is incorrect
        let c = CharClass::new(~[('a', 'b'), ('i', 'j'), ('y', 'z')]);
        assert!(c.includes('b'));
        assert!(c.includes('y'));
    }

    #[test]
    fn negate_simple() {
        let c = CharClass::new(~[('T', 's'), ('☻', '♪')]).negate();
        assert_eq!(c.ranges(), [('\0', 'S'), ('t', '☺'), ('♫', char::MAX)]);
    }

    #[test]
    fn negate_edge_case() {
        let c = CharClass::new(~[('\0', char::MAX)]).negate();
        assert_eq!(c.ranges(), []);
    }

    #[test]
    fn issue_1() {
        // See <https://github.com/lfairy/rose/issues/1>
        let c = CharClass::new(~[('c', 'c'), ('d', 'd'), ('z', 'z')]);
        assert_eq!(c.ranges(), [('c', 'd'), ('z', 'z')]);
    }
}

pub mod ascii;
<commit_msg>Make tests compile again<commit_after>#[doc(hidden)];

//! Utilities for handling character classes.

use std::char;
use std::cmp::{min, max};

/// A range of codepoints.
pub type Range = (char, char);

/// A character class is a non-empty collection of ranges.
#[deriving(Clone)]
pub enum CharClass {
    priv CCOwned(~[Range]),
    priv CCStatic(&'static [Range])
}

impl CharClass {
    /// Create a class from a vector of ranges.
    pub fn new(mut r: ~[Range]) -> CharClass {
        if r.iter().any(|&(lo, hi)| lo > hi) {
            fail!("invalid range");
        }

        r.sort_by(|a, b| {
            // Sort by start value ascending, end value descending
            match a.lo().cmp(&b.lo()) {
                Equal => b.hi().cmp(&a.hi()),
                o => o
            }
        });

        if r.len() > 0 {
            // Coalesce overlapping ranges in place
            let mut cursor = 0;
            for i in range(1, r.len()) {
                if r[cursor].hi() == char::MAX {
                    break
                } else if r[i].lo() <= next_char(r[cursor].hi()) {
                    // Merge the two ranges
                    r[cursor] = (
                        min(r[cursor].lo(), r[i].lo()),
                        max(r[cursor].hi(), r[i].hi()));
                } else {
                    // Create a new range
                    cursor += 1;
                    r[cursor] = r[i];
                }
            }
            r.truncate(cursor + 1); r.shrink_to_fit();
            CCOwned(r)
        } else {
            fail!("char class cannot be empty")
        }
    }

    /// Create a class matching a single code point.
    pub fn from_char(c: char) -> CharClass {
        CharClass::new(~[(c, c)])
    }

    /// Create a class from a range of code points.
    pub fn from_range(lo: char, hi: char) -> CharClass {
        CharClass::new(~[(lo, hi)])
    }

    /// Combine several classes into one that subsumes them all.
    pub fn combine(classes: &[CharClass]) -> CharClass {
        let mut ranges: ~[Range] = ~[];
        for cc in classes.iter() {
            for r in cc.ranges().iter() {
                ranges.push(*r);
            }
        }
        CharClass::new(ranges)
    }

    ///
    /// Get the list of ranges contained in the character class.
    ///
    /// The returned vector is always sorted, with no overlapping
    /// ranges.
    ///
    pub fn ranges<'a>(&'a self) -> &'a [Range] {
        match *self {
            CCOwned(ref r) => {
                let r_borrow: &'a [Range] = *r;
                r_borrow
            },
            CCStatic(r) => {
                let r_borrow: &'a [Range] = r;
                r_borrow
            }
        }
    }

    /// If the class matches a single code point, return it; otherwise
    /// return `None`.
    pub fn to_char(&self) -> Option<char> {
        match self.ranges() {
            [(lo, hi)] if lo == hi => Some(lo),
            _ => None
        }
    }

    /// Determine if the character is contained in the class.
    pub fn includes(&self, c: char) -> bool {
        self.ranges().bsearch(|&(lo, hi)| {
            if c < lo {
                Greater
            } else if hi < c {
                Less
            } else {
                Equal
            }
        }).is_some()
    }

    /// Return the negation of the character class.
    pub fn negate(&self) -> CharClass {
        let ranges = self.ranges();
        let mut result: ~[Range] = ~[];

        let first_lo = ranges[0].lo();
        if first_lo != '\0' {
            result.push(('\0', prev_char(first_lo)));
        }

        let mut last_hi = ranges[0].hi();
        for &(lo, hi) in ranges.slice_from(1).iter() {
            result.push((next_char(last_hi), prev_char(lo)));
            last_hi = hi;
        }

        if last_hi != char::MAX {
            result.push((next_char(last_hi), char::MAX));
        }

        CCOwned(result)
    }
}

priv trait RangeUtils {
    fn lo(&self) -> char;
    fn hi(&self) -> char;
}

impl RangeUtils for (char, char) {
    #[inline]
    fn lo(&self) -> char {
        let &(lo, _) = self;
        lo
    }

    #[inline]
    fn hi(&self) -> char {
        let &(_, hi) = self;
        hi
    }
}

#[inline]
fn prev_char(c: char) -> char {
    char::from_u32(c as u32 - 1).unwrap()
}

#[inline]
fn next_char(c: char) -> char {
    char::from_u32(c as u32 + 1).unwrap()
}

#[cfg(test)]
mod test {
    use std::char;
    use super::CharClass;

    #[test]
    #[should_fail]
    fn invalid_range() {
        let _ = CharClass::new(~[('b', 'a')]);
    }

    #[test]
    #[should_fail]
    fn empty_class() {
        let _ = CharClass::new(~[]);
    }

    #[test]
    fn class_sorted() {
        let c = CharClass::new(~[('y', 'z'), ('a', 'b')]);
        assert_eq!(c.ranges(), &[('a', 'b'), ('y', 'z')]);
    }

    #[test]
    fn class_optimize() {
        let c1 = CharClass::new(~[('a', 'b'), ('c', 'd')]);  // touching
        let c2 = CharClass::new(~[('a', 'c'), ('b', 'd')]);  // overlapping
        let c3 = CharClass::new(~[('a', 'd'), ('b', 'c')]);  // contained
        assert_eq!(c1.ranges(), &[('a', 'd')]);
        assert_eq!(c2.ranges(), &[('a', 'd')]);
        assert_eq!(c3.ranges(), &[('a', 'd')]);
    }

    #[test]
    fn class_includes() {
        // This will fail if the bsearch is incorrect
        let c = CharClass::new(~[('a', 'b'), ('i', 'j'), ('y', 'z')]);
        assert!(c.includes('b'));
        assert!(c.includes('y'));
    }

    #[test]
    fn negate_simple() {
        let c = CharClass::new(~[('T', 's'), ('☻', '♪')]).negate();
        assert_eq!(c.ranges(), &[('\0', 'S'), ('t', '☺'), ('♫', char::MAX)]);
    }

    #[test]
    fn negate_edge_case() {
        let c = CharClass::new(~[('\0', char::MAX)]).negate();
        assert_eq!(c.ranges(), &[]);
    }

    #[test]
    fn issue_1() {
        // See <https://github.com/lfairy/rose/issues/1>
        let c = CharClass::new(~[('c', 'c'), ('d', 'd'), ('z', 'z')]);
        assert_eq!(c.ranges(), &[('c', 'd'), ('z', 'z')]);
    }
}

pub mod ascii;
<|endoftext|>


<commit_before>#![crate_name = "chroot"]
#![feature(collections, rustc_private, path_ext)]

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

extern crate getopts;
extern crate libc;

use getopts::{optflag, optopt, getopts, usage};
use c_types::{get_pw_from_args, get_group};
use libc::funcs::posix88::unistd::{execvp, setgid, setuid};
use std::ffi::{CStr, CString};
use std::fs::PathExt;
use std::io::Write;
use std::iter::FromIterator;
use std::path::Path;

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

extern {
    fn chroot(path: *const libc::c_char) -> libc::c_int;
}

#[cfg(any(target_os = "macos", target_os = "freebsd"))]
extern {
    fn setgroups(size: libc::c_int, list: *const libc::gid_t) -> libc::c_int;
}

#[cfg(target_os = "linux")]
extern {
    fn setgroups(size: libc::size_t, list: *const libc::gid_t) -> libc::c_int;
}

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

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

    let options = [
        optopt("u", "user", "User (ID or name) to switch before running the program", "USER"),
        optopt("g", "group", "Group (ID or name) to switch to", "GROUP"),
        optopt("G", "groups", "Comma-separated list of groups to switch to", "GROUP1,GROUP2…"),
        optopt("", "userspec", "Colon-separated user and group to switch to. \
                                Same as -u USER -g GROUP. \
                                Userspec has higher preference than -u and/or -g", "USER:GROUP"),
        optflag("h", "help", "Show help"),
        optflag("V", "version", "Show program's version")
    ];

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

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

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

    let default_shell: &'static str = "/bin/sh";
    let default_option: &'static str = "-i";
    let user_shell = std::env::var("SHELL");

    let newroot = Path::new(&opts.free[0][..]);
    if !newroot.is_dir() {
        crash!(1, "cannot change root directory to `{}`: no such directory", newroot.display());
    }

    let command: Vec<&str> = match opts.free.len() {
        1 => {
            let shell: &str = match user_shell {
                Err(_) => default_shell,
                Ok(ref s) => &s[..],
            };
            vec!(shell, default_option)
        }
        _ => opts.free[1..opts.free.len()].iter().map(|x| &x[..]).collect()
    };

    set_context(&newroot, &opts);

    unsafe {
        let executable = CString::new(command[0]).unwrap().as_bytes_with_nul().as_ptr();
        let mut command_parts: Vec<*const u8> = command.iter().map(|x| CString::new(x.as_bytes()).unwrap().as_bytes_with_nul().as_ptr()).collect();
        command_parts.push(std::ptr::null());
        execvp(executable as *const libc::c_char, command_parts.as_ptr() as *mut *const libc::c_char) as i32
    }
}

fn set_context(root: &Path, options: &getopts::Matches) {
    let userspec_str = options.opt_str("userspec");
    let user_str = options.opt_str("user").unwrap_or_default();
    let group_str = options.opt_str("group").unwrap_or_default();
    let groups_str = options.opt_str("groups").unwrap_or_default();
    let userspec = match userspec_str {
        Some(ref u) => {
            let s: Vec<&str> = u.split(':').collect();
            if s.len() != 2 {
                crash!(1, "invalid userspec: `{}`", u)
            };
            s
        }
        None => Vec::new()
    };
    let user = if userspec.is_empty() { &user_str[..] } else { &userspec[0][..] };
    let group = if userspec.is_empty() { &group_str[..] } else { &userspec[1][..] };

    enter_chroot(root);

    set_groups_from_str(&groups_str[..]);
    set_main_group(&group[..]);
    set_user(&user[..]);
}

fn enter_chroot(root: &Path) {
    let root_str = root.display();
    std::env::set_current_dir(root).unwrap();
    let err = unsafe {
        chroot(CString::new(".".as_bytes()).unwrap().as_bytes_with_nul().as_ptr() as *const libc::c_char)
    };
    if err != 0 {
        crash!(1, "cannot chroot to {}: {}", root_str, strerror(err))
    };
}

fn set_main_group(group: &str) {
    if !group.is_empty() {
        let group_id = match get_group(group) {
            None => crash!(1, "no such group: {}", group),
            Some(g) => g.gr_gid
        };
        let err = unsafe { setgid(group_id) };
        if err != 0 {
            crash!(1, "cannot set gid to {}: {}", group_id, strerror(err))
        }
    }
}

#[cfg(any(target_os = "macos", target_os = "freebsd"))]
fn set_groups(groups: Vec<libc::gid_t>) -> libc::c_int {
    unsafe {
        setgroups(groups.len() as libc::c_int,
                  groups.as_ptr())
    }
}

#[cfg(target_os = "linux")]
fn set_groups(groups: Vec<libc::gid_t>) -> libc::c_int {
    unsafe {
        setgroups(groups.len() as libc::size_t,
                  groups.as_ptr())
    }
}

fn set_groups_from_str(groups: &str) {
    if !groups.is_empty() {
        let groups_vec: Vec<libc::gid_t> = FromIterator::from_iter(
            groups.split(',').map(
                |x| match get_group(x) {
                    None => crash!(1, "no such group: {}", x),
                    Some(g) => g.gr_gid
                })
            );
        let err = set_groups(groups_vec);
        if err != 0 {
            crash!(1, "cannot set groups: {}", strerror(err))
        }
    }
}

fn set_user(user: &str) {
    if !user.is_empty() {
        let user_id = get_pw_from_args(&vec!(String::from_str(user))).unwrap().pw_uid;
        let err = unsafe { setuid(user_id as libc::uid_t) };
        if err != 0 {
            crash!(1, "cannot set user to {}: {}", user, strerror(err))
        }
    }
}

fn strerror(errno: i32) -> String {
    unsafe {
        let err = libc::funcs::c95::string::strerror(errno) as *const libc::c_char;
        let bytes= CStr::from_ptr(err).to_bytes();
        String::from_utf8_lossy(bytes).to_string()
    }
}

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

fn help_menu(program: &str, options: &[getopts::OptGroup]) {
    version();
    println!("Usage:");
    println!("  {} [OPTION]… NEWROOT [COMMAND [ARG]…]", program);
    println!("");
    print!("{}", usage(
            "Run COMMAND with root directory set to NEWROOT.\n\
             If COMMAND is not specified, it defaults to '${SHELL} -i'. \
             If ${SHELL} is not set, /bin/sh is used.", options))
}
<commit_msg>working chroot implementation<commit_after>#![crate_name = "chroot"]
#![feature(collections, rustc_private, path_ext)]

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

extern crate getopts;
extern crate libc;

use getopts::{optflag, optopt, getopts, usage};
use c_types::{get_pw_from_args, get_group};
use libc::funcs::posix88::unistd::{setgid, setuid};
use std::ffi::{CString};
use std::fs::PathExt;
use std::io::{Error, Write};
use std::iter::FromIterator;
use std::path::Path;
use std::process::Command;

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

extern {
    fn chroot(path: *const libc::c_char) -> libc::c_int;
}

#[cfg(any(target_os = "macos", target_os = "freebsd"))]
extern {
    fn setgroups(size: libc::c_int, list: *const libc::gid_t) -> libc::c_int;
}

#[cfg(target_os = "linux")]
extern {
    fn setgroups(size: libc::size_t, list: *const libc::gid_t) -> libc::c_int;
}

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

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

    let options = [
        optopt("u", "user", "User (ID or name) to switch before running the program", "USER"),
        optopt("g", "group", "Group (ID or name) to switch to", "GROUP"),
        optopt("G", "groups", "Comma-separated list of groups to switch to", "GROUP1,GROUP2…"),
        optopt("", "userspec", "Colon-separated user and group to switch to. \
                                Same as -u USER -g GROUP. \
                                Userspec has higher preference than -u and/or -g", "USER:GROUP"),
        optflag("h", "help", "Show help"),
        optflag("V", "version", "Show program's version")
    ];

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

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

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

    let default_shell: &'static str = "/bin/sh";
    let default_option: &'static str = "-i";
    let user_shell = std::env::var("SHELL");

    let newroot = Path::new(&opts.free[0][..]);
    if !newroot.is_dir() {
        crash!(1, "cannot change root directory to `{}`: no such directory", newroot.display());
    }

    let command: Vec<&str> = match opts.free.len() {
        1 => {
            let shell: &str = match user_shell {
                Err(_) => default_shell,
                Ok(ref s) => s.as_ref(),
            };
            vec!(shell, default_option)
        },
        _ => opts.free[1..].iter().map(|x| &x[..]).collect()
    };

    set_context(&newroot, &opts);

    let pstatus = Command::new(command[0])
        .args(&command[1..])
        .status()
        .unwrap_or_else(|e| crash!(1, "Cannot exec: {}", e));

    if pstatus.success() {
        0
    } else {
        match pstatus.code() {
            Some(i) => i,
            None => -1,
        }
    }
}

fn set_context(root: &Path, options: &getopts::Matches) {
    let userspec_str = options.opt_str("userspec");
    let user_str = options.opt_str("user").unwrap_or_default();
    let group_str = options.opt_str("group").unwrap_or_default();
    let groups_str = options.opt_str("groups").unwrap_or_default();
    let userspec = match userspec_str {
        Some(ref u) => {
            let s: Vec<&str> = u.split(':').collect();
            if s.len() != 2 {
                crash!(1, "invalid userspec: `{}`", u)
            };
            s
        }
        None => Vec::new()
    };
    let user = if userspec.is_empty() { &user_str[..] } else { &userspec[0][..] };
    let group = if userspec.is_empty() { &group_str[..] } else { &userspec[1][..] };

    enter_chroot(root);

    set_groups_from_str(&groups_str[..]);
    set_main_group(&group[..]);
    set_user(&user[..]);
}

fn enter_chroot(root: &Path) {
    let root_str = root.display();
    std::env::set_current_dir(root).unwrap();
    let err = unsafe {
        chroot(CString::new(".".as_bytes()).unwrap().as_bytes_with_nul().as_ptr() as *const libc::c_char)
    };
    if err != 0 {
        crash!(1, "cannot chroot to {}: {}", root_str, Error::last_os_error())
    };
}

fn set_main_group(group: &str) {
    if !group.is_empty() {
        let group_id = match get_group(group) {
            None => crash!(1, "no such group: {}", group),
            Some(g) => g.gr_gid
        };
        let err = unsafe { setgid(group_id) };
        if err != 0 {
            crash!(1, "cannot set gid to {}: {}", group_id, Error::last_os_error())
        }
    }
}

#[cfg(any(target_os = "macos", target_os = "freebsd"))]
fn set_groups(groups: Vec<libc::gid_t>) -> libc::c_int {
    unsafe {
        setgroups(groups.len() as libc::c_int,
                  groups.as_ptr())
    }
}

#[cfg(target_os = "linux")]
fn set_groups(groups: Vec<libc::gid_t>) -> libc::c_int {
    unsafe {
        setgroups(groups.len() as libc::size_t,
                  groups.as_ptr())
    }
}

fn set_groups_from_str(groups: &str) {
    if !groups.is_empty() {
        let groups_vec: Vec<libc::gid_t> = FromIterator::from_iter(
            groups.split(',').map(
                |x| match get_group(x) {
                    None => crash!(1, "no such group: {}", x),
                    Some(g) => g.gr_gid
                })
            );
        let err = set_groups(groups_vec);
        if err != 0 {
            crash!(1, "cannot set groups: {}", Error::last_os_error())
        }
    }
}

fn set_user(user: &str) {
    if !user.is_empty() {
        let user_id = get_pw_from_args(&vec!(String::from_str(user))).unwrap().pw_uid;
        let err = unsafe { setuid(user_id as libc::uid_t) };
        if err != 0 {
            crash!(1, "cannot set user to {}: {}", user, Error::last_os_error())
        }
    }
}

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

fn help_menu(program: &str, options: &[getopts::OptGroup]) {
    version();
    println!("Usage:");
    println!("  {} [OPTION]… NEWROOT [COMMAND [ARG]…]", program);
    println!("");
    print!("{}", usage(
            "Run COMMAND with root directory set to NEWROOT.\n\
             If COMMAND is not specified, it defaults to '${SHELL} -i'. \
             If ${SHELL} is not set, /bin/sh is used.", options))
}
<|endoftext|>


<commit_before>#![crate_name = "ansi_term"]
#![crate_type = "rlib"]
#![crate_type = "dylib"]
#![feature(plugin)]
#![plugin(regex_macros)]

//! This is a library for controlling colours and formatting, such as
//! red bold text or blue underlined text, on ANSI terminals.
//!
//! ```rust
//! extern crate ansi_term;
//! use ansi_term::Colour::{Black, Red, Green, Yellow, Blue, Purple, Cyan, Fixed};
//! use ansi_term::Style::Plain;
//! ```
//!
//! Simple Colours
//! --------------
//!
//! You can format strings by calling the `paint` method on a Colour
//! or a Style object, passing in the string you want to format. For
//! example, to get some red text, call the `paint` method on `Red`:
//!
//! ```rust
//! println!("This is in red: {}!", Red.paint("a red string"));
//! ```
//!
//! The `paint` method returns an `ANSIString` object, which will get
//! automatically converted to the correct sequence of escape codes when
//! used in a `println!` or `format!` macro, or anything else that
//! supports using the `Show` trait. This means that if you just want a
//! string of the escape codes without anything else, you can still use
//! the `to_string` method:
//!
//! ```rust
//! let red_string: String = Red.paint("another red string").to_string();
//! ```
//!
//! Bold, Underline, and Background
//! -------------------------------
//!
//! To do anything more complex than just foreground colours, you need
//! to use Style objects. Calling the `bold` or `underline` method on
//! a Colour returns a Style that has the appropriate property set on
//! it:
//!
//! ```rust
//! println!("Demonstrating {} and {}!",
//!          Blue.bold().paint("blue bold"),
//!          Yellow.underline().paint("yellow underline"));
//! ```
//!
//! These methods chain, so you can call them on existing Style
//! objects to set more than one particular properly, like so:
//!
//! ```rust
//! Blue.underline().bold().paint("Blue underline bold!")
//! ```
//!
//! You can set the background colour of a Style by using the `on`
//! method:
//!
//! ```rust
//! Blue.on(Yellow).paint("Blue on yellow!")
//! ```
//!
//! Finally, you can turn a Colour into a Style with the `normal`
//! method, though it'll produce the exact same string if you just use
//! the Colour. It's only useful if you're writing a method that can
//! return either normal or bold (or underline) styles, and need to
//! return a Style object from it.
//!
//! ```rust
//! Red.normal().paint("yet another red string")
//! ```
//!
//! Extended Colours
//! ----------------
//!
//! You can access the extended range of 256 colours by using the
//! Fixed constructor, which takes an argument of the colour number to
//! use. This can be used wherever you would use a Colour:
//!
//! ```rust
//! Fixed(134).paint("A sort of light purple.")
//! ```
//!
//! This even works for background colours:
//!
//! ```rust
//! Fixed(221).on(Fixed(124)).paint("Mustard in the ketchup.")
//! ```
//!
//! No Formatting
//! -------------
//!
//! Finally, for the sake of completeness, the Plain style provides
//! neither colours nor formatting.
//!
//! ```rust
//! Plain.paint("No colours here.")
//! ```

extern crate regex;

use Colour::{Black, Red, Green, Yellow, Blue, Purple, Cyan, White, Fixed};
use Style::{Plain, Foreground, Styled};
use std::fmt;

/// An ANSI String is a string coupled with the Style to display it
/// in a terminal.
///
/// Although not technically a string itself, it can be turned into
/// one with the `to_string` method.
#[derive(Copy)]
pub struct ANSIString<'a> {
    string: &'a str,
    style: Style,
}

impl<'a> ANSIString<'a> {
    /// Creates a new ANSI String with the given contents and style.
    pub fn new(contents: &'a str, style: Style) -> ANSIString<'a> {
        ANSIString { string: contents, style: style }
    }
}

impl<'a> fmt::Display for ANSIString<'a> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self.style {
            Plain => write!(f, "{}", self.string),
            Foreground(colour) => {
                try!(f.write_str("\x1B["));
                try!(colour.write_foreground_code(f));
                write!(f, "m{}\x1B[0m", self.string)
            },
            Styled { foreground, background, bold, underline } => {
                try!(f.write_str("\x1B["));

                if bold {
                    try!(f.write_str("1;"));
                }

                if underline {
                    try!(f.write_str("4;"));
                }

                match background {
                    Some(c) => {
                        try!(c.write_background_code(f));
                        try!(f.write_str(";"));
                    },
                    None => {},
                }

                try!(foreground.write_foreground_code(f));

                write!(f, "m{}\x1B[0m", self.string)
            }
        }
    }
}

/// A colour is one specific type of ANSI escape code, and can refer
/// to either the foreground or background colour.
///
/// These use the standard numeric sequences.
/// See http://invisible-island.net/xterm/ctlseqs/ctlseqs.html
#[derive(Copy)]
pub enum Colour {
    Black, Red, Green, Yellow, Blue, Purple, Cyan, White, Fixed(u8),
}

// I'm not beyond calling Colour Colour, rather than Color, but I did
// purposefully name this crate 'ansi-term' so people wouldn't get
// confused when they tried to install it.
//
// Only *after* they'd installed it.

impl Colour {
    fn write_foreground_code(self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            Black  => f.write_str("30"),
            Red    => f.write_str("31"),
            Green  => f.write_str("32"),
            Yellow => f.write_str("33"),
            Blue   => f.write_str("34"),
            Purple => f.write_str("35"),
            Cyan   => f.write_str("36"),
            White  => f.write_str("37"),
            Fixed(num) => write!(f, "38;5;{}", num),
        }
    }

    fn write_background_code(self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            Black  => f.write_str("40"),
            Red    => f.write_str("41"),
            Green  => f.write_str("42"),
            Yellow => f.write_str("43"),
            Blue   => f.write_str("44"),
            Purple => f.write_str("45"),
            Cyan   => f.write_str("46"),
            White  => f.write_str("47"),
            Fixed(num) => write!(f, "48;5;{}", num),
        }
    }

    /// Return a Style with the foreground colour set to this colour.
    pub fn normal(self) -> Style {
        Styled { foreground: self, background: None, bold: false, underline: false }
    }

    /// Paints the given text with this colour, returning an ANSI string.
    /// This is a short-cut so you don't have to use Blue.normal() just
    /// to get blue text.
    pub fn paint(self, input: &str) -> ANSIString {
        ANSIString::new(input, Foreground(self))
    }

    /// Returns a Style with the underline property set.
    pub fn underline(self) -> Style {
        Styled { foreground: self, background: None, bold: false, underline: true }
    }

    /// Returns a Style with the bold property set.
    pub fn bold(self) -> Style {
        Styled { foreground: self, background: None, bold: true, underline: false }
    }

    /// Returns a Style with the background colour property set.
    pub fn on(self, background: Colour) -> Style {
        Styled { foreground: self, background: Some(background), bold: false, underline: false }
    }
}

/// A style is a collection of properties that can format a string
/// using ANSI escape codes.
#[derive(Copy)]
pub enum Style {

    /// The Plain style provides no formatting.
    Plain,

    /// The Foreground style provides just a foreground colour.
    Foreground(Colour),

    /// The Styled style is a catch-all for anything more complicated
    /// than that. It's technically possible for there to be other
    /// cases, such as "bold foreground", but probably isn't worth it.
    Styled { foreground: Colour, background: Option<Colour>, bold: bool, underline: bool, },
}

impl Style {
    /// Paints the given text with this colour, returning an ANSI string.
    pub fn paint(self, input: &str) -> ANSIString {
        ANSIString::new(input, self)
    }

    /// Returns a Style with the bold property set.
    pub fn bold(self) -> Style {
        match self {
            Plain => Styled { foreground: White, background: None, bold: true, underline: false },
            Foreground(c) => Styled { foreground: c, background: None, bold: true, underline: false },
            Styled { foreground, background, bold: _, underline } => Styled { foreground: foreground, background: background, bold: true, underline: underline },
        }
    }

    /// Returns a Style with the underline property set.
    pub fn underline(self) -> Style {
        match self {
            Plain => Styled { foreground: White, background: None, bold: false, underline: true },
            Foreground(c) => Styled { foreground: c, background: None, bold: false, underline: true },
            Styled { foreground, background, bold, underline: _ } => Styled { foreground: foreground, background: background, bold: bold, underline: true },
        }
    }

    /// Returns a Style with the background colour property set.
    pub fn on(self, background: Colour) -> Style {
        match self {
            Plain => Styled { foreground: White,background: Some(background), bold: false, underline: false },
            Foreground(c) => Styled { foreground: c, background: Some(background), bold: false, underline: false },
            Styled { foreground, background: _, bold, underline } => Styled { foreground: foreground, background: Some(background), bold: bold, underline: underline },
        }
    }
}

#[cfg(test)]
mod tests {
    use super::Style::Plain;
    use super::Colour::{Black, Red, Green, Yellow, Blue, Purple, Cyan, White, Fixed};

    macro_rules! test {
        ($name: ident: $style: expr; $input: expr => $result: expr) => {
            #[test]
            fn $name() {
                assert_eq!($style.paint($input).to_string(), $result.to_string())
            }
        };
    }

    test!(plain:                 Plain;                             "text/plain" => "text/plain");
    test!(red:                   Red;                               "hi" => "\x1B[31mhi\x1B[0m");
    test!(black:                 Black.normal();                    "hi" => "\x1B[30mhi\x1B[0m");
    test!(yellow_bold:           Yellow.bold();                     "hi" => "\x1B[1;33mhi\x1B[0m");
    test!(yellow_bold_2:         Yellow.normal().bold();            "hi" => "\x1B[1;33mhi\x1B[0m");
    test!(blue_underline:        Blue.underline();                  "hi" => "\x1B[4;34mhi\x1B[0m");
    test!(green_bold_ul:         Green.bold().underline();          "hi" => "\x1B[1;4;32mhi\x1B[0m");
    test!(green_bold_ul_2:       Green.underline().bold();          "hi" => "\x1B[1;4;32mhi\x1B[0m");
    test!(purple_on_white:       Purple.on(White);                  "hi" => "\x1B[47;35mhi\x1B[0m");
    test!(purple_on_white_2:     Purple.normal().on(White);         "hi" => "\x1B[47;35mhi\x1B[0m");
    test!(cyan_bold_on_white:    Cyan.bold().on(White);             "hi" => "\x1B[1;47;36mhi\x1B[0m");
    test!(cyan_ul_on_white:      Cyan.underline().on(White);        "hi" => "\x1B[4;47;36mhi\x1B[0m");
    test!(cyan_bold_ul_on_white: Cyan.bold().underline().on(White); "hi" => "\x1B[1;4;47;36mhi\x1B[0m");
    test!(cyan_ul_bold_on_white: Cyan.underline().bold().on(White); "hi" => "\x1B[1;4;47;36mhi\x1B[0m");
    test!(fixed:                 Fixed(100);                        "hi" => "\x1B[38;5;100mhi\x1B[0m");
    test!(fixed_on_purple:       Fixed(100).on(Purple);             "hi" => "\x1B[45;38;5;100mhi\x1B[0m");
    test!(fixed_on_fixed:        Fixed(100).on(Fixed(200));         "hi" => "\x1B[48;5;200;38;5;100mhi\x1B[0m");
}
<commit_msg>Allow bold and underline with no foreground<commit_after>#![crate_name = "ansi_term"]
#![crate_type = "rlib"]
#![crate_type = "dylib"]
#![feature(plugin)]
#![plugin(regex_macros)]

//! This is a library for controlling colours and formatting, such as
//! red bold text or blue underlined text, on ANSI terminals.
//!
//! ```rust
//! extern crate ansi_term;
//! use ansi_term::Colour::{Black, Red, Green, Yellow, Blue, Purple, Cyan, Fixed};
//! use ansi_term::Style::Plain;
//! ```
//!
//! Simple Colours
//! --------------
//!
//! You can format strings by calling the `paint` method on a Colour
//! or a Style object, passing in the string you want to format. For
//! example, to get some red text, call the `paint` method on `Red`:
//!
//! ```rust
//! println!("This is in red: {}!", Red.paint("a red string"));
//! ```
//!
//! The `paint` method returns an `ANSIString` object, which will get
//! automatically converted to the correct sequence of escape codes when
//! used in a `println!` or `format!` macro, or anything else that
//! supports using the `Show` trait. This means that if you just want a
//! string of the escape codes without anything else, you can still use
//! the `to_string` method:
//!
//! ```rust
//! let red_string: String = Red.paint("another red string").to_string();
//! ```
//!
//! Bold, Underline, and Background
//! -------------------------------
//!
//! To do anything more complex than just foreground colours, you need
//! to use Style objects. Calling the `bold` or `underline` method on
//! a Colour returns a Style that has the appropriate property set on
//! it:
//!
//! ```rust
//! println!("Demonstrating {} and {}!",
//!          Blue.bold().paint("blue bold"),
//!          Yellow.underline().paint("yellow underline"));
//! ```
//!
//! These methods chain, so you can call them on existing Style
//! objects to set more than one particular properly, like so:
//!
//! ```rust
//! Blue.underline().bold().paint("Blue underline bold!")
//! ```
//!
//! You can set the background colour of a Style by using the `on`
//! method:
//!
//! ```rust
//! Blue.on(Yellow).paint("Blue on yellow!")
//! ```
//!
//! Finally, you can turn a Colour into a Style with the `normal`
//! method, though it'll produce the exact same string if you just use
//! the Colour. It's only useful if you're writing a method that can
//! return either normal or bold (or underline) styles, and need to
//! return a Style object from it.
//!
//! ```rust
//! Red.normal().paint("yet another red string")
//! ```
//!
//! Extended Colours
//! ----------------
//!
//! You can access the extended range of 256 colours by using the
//! Fixed constructor, which takes an argument of the colour number to
//! use. This can be used wherever you would use a Colour:
//!
//! ```rust
//! Fixed(134).paint("A sort of light purple.")
//! ```
//!
//! This even works for background colours:
//!
//! ```rust
//! Fixed(221).on(Fixed(124)).paint("Mustard in the ketchup.")
//! ```
//!
//! No Formatting
//! -------------
//!
//! Finally, for the sake of completeness, the Plain style provides
//! neither colours nor formatting.
//!
//! ```rust
//! Plain.paint("No colours here.")
//! ```

extern crate regex;

use Colour::{Black, Red, Green, Yellow, Blue, Purple, Cyan, White, Fixed};
use Style::{Plain, Foreground, Styled};
use std::fmt;

/// An ANSI String is a string coupled with the Style to display it
/// in a terminal.
///
/// Although not technically a string itself, it can be turned into
/// one with the `to_string` method.
#[derive(Copy)]
pub struct ANSIString<'a> {
    string: &'a str,
    style: Style,
}

impl<'a> ANSIString<'a> {
    /// Creates a new ANSI String with the given contents and style.
    pub fn new(contents: &'a str, style: Style) -> ANSIString<'a> {
        ANSIString { string: contents, style: style }
    }
}

impl<'a> fmt::Display for ANSIString<'a> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self.style {
            Plain => write!(f, "{}", self.string),
            Foreground(colour) => {
                try!(f.write_str("\x1B["));
                try!(colour.write_foreground_code(f));
                write!(f, "m{}\x1B[0m", self.string)
            },
            Styled { foreground, background, bold, underline } => {
                let mut semicolon = false;
                try!(f.write_str("\x1B["));

                if bold {
                    try!(f.write_str("1"));
                    semicolon = true;
                }

                if underline {
                    if semicolon { try!(f.write_str(";")); }
                    try!(f.write_str("4"));
                    semicolon = true;
                }

                match background {
                    Some(c) => {
                        if semicolon { try!(f.write_str(";")); }
                        try!(c.write_background_code(f));
                        semicolon = true;
                    },
                    None => {},
                }

                if let Some(fg) = foreground {
                    if semicolon { try!(f.write_str(";")); }
                    try!(fg.write_foreground_code(f));
                }

                write!(f, "m{}\x1B[0m", self.string)
            }
        }
    }
}

/// A colour is one specific type of ANSI escape code, and can refer
/// to either the foreground or background colour.
///
/// These use the standard numeric sequences.
/// See http://invisible-island.net/xterm/ctlseqs/ctlseqs.html
#[derive(Copy)]
pub enum Colour {
    Black, Red, Green, Yellow, Blue, Purple, Cyan, White, Fixed(u8),
}

// I'm not beyond calling Colour Colour, rather than Color, but I did
// purposefully name this crate 'ansi-term' so people wouldn't get
// confused when they tried to install it.
//
// Only *after* they'd installed it.

impl Colour {
    fn write_foreground_code(self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            Black  => f.write_str("30"),
            Red    => f.write_str("31"),
            Green  => f.write_str("32"),
            Yellow => f.write_str("33"),
            Blue   => f.write_str("34"),
            Purple => f.write_str("35"),
            Cyan   => f.write_str("36"),
            White  => f.write_str("37"),
            Fixed(num) => write!(f, "38;5;{}", num),
        }
    }

    fn write_background_code(self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            Black  => f.write_str("40"),
            Red    => f.write_str("41"),
            Green  => f.write_str("42"),
            Yellow => f.write_str("43"),
            Blue   => f.write_str("44"),
            Purple => f.write_str("45"),
            Cyan   => f.write_str("46"),
            White  => f.write_str("47"),
            Fixed(num) => write!(f, "48;5;{}", num),
        }
    }

    /// Return a Style with the foreground colour set to this colour.
    pub fn normal(self) -> Style {
        Styled { foreground: Some(self), background: None, bold: false, underline: false }
    }

    /// Paints the given text with this colour, returning an ANSI string.
    /// This is a short-cut so you don't have to use Blue.normal() just
    /// to get blue text.
    pub fn paint(self, input: &str) -> ANSIString {
        ANSIString::new(input, Foreground(self))
    }

    /// Returns a Style with the underline property set.
    pub fn underline(self) -> Style {
        Styled { foreground: Some(self), background: None, bold: false, underline: true }
    }

    /// Returns a Style with the bold property set.
    pub fn bold(self) -> Style {
        Styled { foreground: Some(self), background: None, bold: true, underline: false }
    }

    /// Returns a Style with the background colour property set.
    pub fn on(self, background: Colour) -> Style {
        Styled { foreground: Some(self), background: Some(background), bold: false, underline: false }
    }
}

/// A style is a collection of properties that can format a string
/// using ANSI escape codes.
#[derive(Copy)]
pub enum Style {

    /// The Plain style provides no formatting.
    Plain,

    /// The Foreground style provides just a foreground colour.
    Foreground(Colour),

    /// The Styled style is a catch-all for anything more complicated
    /// than that. It's technically possible for there to be other
    /// cases, such as "bold foreground", but probably isn't worth it.
    Styled { foreground: Option<Colour>, background: Option<Colour>, bold: bool, underline: bool },
}

impl Style {
    /// Paints the given text with this colour, returning an ANSI string.
    pub fn paint(self, input: &str) -> ANSIString {
        ANSIString::new(input, self)
    }

    /// Returns a Style with the bold property set.
    pub fn bold(self) -> Style {
        match self {
            Plain => Styled { foreground: None, background: None, bold: true, underline: false },
            Foreground(c) => Styled { foreground: Some(c), background: None, bold: true, underline: false },
            Styled { foreground, background, bold: _, underline } => Styled { foreground: foreground, background: background, bold: true, underline: underline },
        }
    }

    /// Returns a Style with the underline property set.
    pub fn underline(self) -> Style {
        match self {
            Plain => Styled { foreground: None, background: None, bold: false, underline: true },
            Foreground(c) => Styled { foreground: Some(c), background: None, bold: false, underline: true },
            Styled { foreground, background, bold, underline: _ } => Styled { foreground: foreground, background: background, bold: bold, underline: true },
        }
    }

    /// Returns a Style with the background colour property set.
    pub fn on(self, background: Colour) -> Style {
        match self {
            Plain => Styled { foreground: None, background: Some(background), bold: false, underline: false },
            Foreground(c) => Styled { foreground: Some(c), background: Some(background), bold: false, underline: false },
            Styled { foreground, background: _, bold, underline } => Styled { foreground: foreground, background: Some(background), bold: bold, underline: underline },
        }
    }
}

#[cfg(test)]
mod tests {
    use super::Style::Plain;
    use super::Colour::{Black, Red, Green, Yellow, Blue, Purple, Cyan, White, Fixed};

    macro_rules! test {
        ($name: ident: $style: expr; $input: expr => $result: expr) => {
            #[test]
            fn $name() {
                assert_eq!($style.paint($input).to_string(), $result.to_string())
            }
        };
    }

    test!(plain:                 Plain;                             "text/plain" => "text/plain");
    test!(red:                   Red;                               "hi" => "\x1B[31mhi\x1B[0m");
    test!(black:                 Black.normal();                    "hi" => "\x1B[30mhi\x1B[0m");
    test!(yellow_bold:           Yellow.bold();                     "hi" => "\x1B[1;33mhi\x1B[0m");
    test!(yellow_bold_2:         Yellow.normal().bold();            "hi" => "\x1B[1;33mhi\x1B[0m");
    test!(blue_underline:        Blue.underline();                  "hi" => "\x1B[4;34mhi\x1B[0m");
    test!(green_bold_ul:         Green.bold().underline();          "hi" => "\x1B[1;4;32mhi\x1B[0m");
    test!(green_bold_ul_2:       Green.underline().bold();          "hi" => "\x1B[1;4;32mhi\x1B[0m");
    test!(purple_on_white:       Purple.on(White);                  "hi" => "\x1B[47;35mhi\x1B[0m");
    test!(purple_on_white_2:     Purple.normal().on(White);         "hi" => "\x1B[47;35mhi\x1B[0m");
    test!(cyan_bold_on_white:    Cyan.bold().on(White);             "hi" => "\x1B[1;47;36mhi\x1B[0m");
    test!(cyan_ul_on_white:      Cyan.underline().on(White);        "hi" => "\x1B[4;47;36mhi\x1B[0m");
    test!(cyan_bold_ul_on_white: Cyan.bold().underline().on(White); "hi" => "\x1B[1;4;47;36mhi\x1B[0m");
    test!(cyan_ul_bold_on_white: Cyan.underline().bold().on(White); "hi" => "\x1B[1;4;47;36mhi\x1B[0m");
    test!(fixed:                 Fixed(100);                        "hi" => "\x1B[38;5;100mhi\x1B[0m");
    test!(fixed_on_purple:       Fixed(100).on(Purple);             "hi" => "\x1B[45;38;5;100mhi\x1B[0m");
    test!(fixed_on_fixed:        Fixed(100).on(Fixed(200));         "hi" => "\x1B[48;5;200;38;5;100mhi\x1B[0m");
    test!(bold:                  Plain.bold();                      "hi" => "\x1B[1mhi\x1B[0m");
    test!(underline:             Plain.underline();                 "hi" => "\x1B[4mhi\x1B[0m");
    test!(bunderline:            Plain.bold().underline();          "hi" => "\x1B[1;4mhi\x1B[0m");
}
<|endoftext|>


<commit_before>use std::sync::{ Once, ONCE_INIT };
use std::intrinsics::abort;
use std::{ mem, ptr };
use rand::{ thread_rng, Rng };
use rand::os::OsRng;
use errno::{ Errno, set_errno };

const GARBAGE_VALUE: u8 = 0xd0;
const CANARY_SIZE: usize = 16;
static ALLOC_INIT: Once = ONCE_INIT;
static mut PAGE_SIZE: usize = 0;
static mut CANARY: [u8; CANARY_SIZE] = [0; CANARY_SIZE];

// -- alloc init --

unsafe fn alloc_init() {
    #[cfg(unix)] {
        let page_size = ::libc::sysconf(::libc::_SC_PAGESIZE);
        if page_size > 0 {
            PAGE_SIZE = page_size as usize;
        }
    }

    #[cfg(windows)] {
        let si = mem::uninitialized();
        ::kernel32::GetSystemInfo(si);
        PAGE_SIZE = ptr::read(si).dwPageSize as usize;
    }

    if PAGE_SIZE < CANARY_SIZE || PAGE_SIZE < mem::size_of::<usize>() {
        abort();
    }

    match OsRng::new() {
        Ok(mut rng) => rng.fill_bytes(&mut CANARY),
        Err(_) => thread_rng().fill_bytes(&mut CANARY)
    }
}

// -- aligned alloc/free --

#[cfg(unix)]
unsafe fn alloc_aligned<T>(size: usize) -> *mut T {
    let mut memptr = mem::uninitialized();
    match ::libc::posix_memalign(&mut memptr, PAGE_SIZE, size) {
        0 => memptr as *mut T,
        ::libc::EINVAL => panic!("EINVAL: invalid alignmen. {}", PAGE_SIZE),
        ::libc::ENOMEM => ptr::null_mut(),
        _ => unreachable!()
    }
}

#[cfg(windows)]
unsafe fn alloc_aligned<T>(size: usize) -> *mut T {
    ::kernel32::VirtualAlloc(
        ptr::null(),
        size,
        ::winapi::MEM_COMMIT | ::winapi::MEM_RESERVE,
        ::winapi::PAGE_READWRITE
    ) as *mut T
}

#[cfg(unix)]
unsafe fn free_aligned<T>(memptr: *mut T) {
    ::libc::free(memptr as *mut ::libc::c_void);
}

#[cfg(windows)]
unsafe fn free_aligned<T>(memptr: *mut T) {
    ::kernel32::VirtualFree(memptr as ::winapi::LPVOID, 0, ::winapi::MEM_RELEASE);
}

// -- malloc --

#[inline]
unsafe fn page_round(size: usize) -> usize {
    let page_mask = PAGE_SIZE - 1;
    (size + page_mask) & (!page_mask)
}

pub unsafe fn unprotected_ptr_from_user_ptr<T>(memptr: *const T) -> *mut T {
    let canary_ptr = memptr.offset(-(mem::size_of_val(&CANARY) as isize));
    let page_mask = PAGE_SIZE - 1;
    let unprotected_ptr_u = canary_ptr as usize & !page_mask;
    if unprotected_ptr_u <= PAGE_SIZE * 2 {
        abort();
    }
    unprotected_ptr_u as *mut T
}

unsafe fn _malloc<T>(size: usize) -> *mut T {
    ALLOC_INIT.call_once(|| alloc_init());

    if size >= ::std::usize::MAX - PAGE_SIZE * 4 {
        set_errno(Errno(::libc::ENOMEM));
        return ptr::null_mut();
    }
    if PAGE_SIZE <= mem::size_of_val(&CANARY) || PAGE_SIZE < mem::size_of::<usize>() {
        abort();
    }

    // aligned alloc ptr
    let size_with_canary = mem::size_of_val(&CANARY) + size;
    let unprotected_size = page_round(size_with_canary);
    let total_size = PAGE_SIZE + PAGE_SIZE + unprotected_size + PAGE_SIZE;
    let base_ptr = alloc_aligned(total_size);
    if base_ptr.is_null() {
        return base_ptr;
    }

    // canary offset
    let unprotected_ptr = base_ptr.offset(PAGE_SIZE as isize * 2);
    ::mprotect(base_ptr.offset(PAGE_SIZE as isize), PAGE_SIZE, ::Prot::NoAccess);
    ptr::copy(
        CANARY.as_ptr(),
        unprotected_ptr.offset(unprotected_size as isize) as *mut u8,
        mem::size_of_val(&CANARY)
    );

    // mprotect ptr
    ::mprotect(unprotected_ptr.offset(unprotected_size as isize), PAGE_SIZE, ::Prot::NoAccess);
    ::mlock(unprotected_ptr, unprotected_size);
    let canary_ptr = unprotected_ptr
        .offset(page_round(size_with_canary) as isize)
        .offset(-(size_with_canary as isize));
    let user_ptr = canary_ptr.offset(mem::size_of_val(&CANARY) as isize);
    ptr::copy(CANARY.as_ptr(), canary_ptr as *mut u8, mem::size_of_val(&CANARY));
    ptr::write(base_ptr as *mut usize, unprotected_size);
    ::mprotect(base_ptr, PAGE_SIZE, ::Prot::ReadOnly);

    debug_assert_eq!(unprotected_ptr_from_user_ptr(user_ptr), unprotected_ptr);

    user_ptr
}

pub unsafe fn malloc<T>(size: usize) -> *mut T {
    let memptr = _malloc(size);
    if !memptr.is_null() {
        ::memset(memptr, GARBAGE_VALUE as i32, size);
    }
    memptr
}

pub unsafe fn allocarray<T>(count: usize, size: usize) -> *mut T {
    if count > mem::size_of::<usize>() && size >= ::std::usize::MAX / count {
        set_errno(Errno(::libc::ENOMEM));
        return ptr::null_mut();
    }
    malloc(count * size)
}

pub unsafe fn free<T>(memptr: *mut T) {
    if memptr.is_null() { return () };

    let canary_ptr = memptr.offset(-(mem::size_of_val(&CANARY) as isize));
    let unprotected_ptr = unprotected_ptr_from_user_ptr(memptr);
    let base_ptr = unprotected_ptr.offset(-(PAGE_SIZE as isize * 2));
    let unprotected_size = ptr::read(base_ptr as *const usize);
    let total_size = PAGE_SIZE + PAGE_SIZE + unprotected_size + PAGE_SIZE;
    ::mprotect(base_ptr, total_size, ::Prot::ReadWrite);

    debug_assert!(::memcmp(canary_ptr as *const u8, CANARY.as_ptr(), mem::size_of_val(&CANARY)));
    debug_assert!(::memcmp(
        unprotected_ptr.offset(unprotected_size as isize) as *const u8,
        CANARY.as_ptr(),
        mem::size_of_val(&CANARY)
    ));

    ::munlock(unprotected_ptr, unprotected_size);
    free_aligned(base_ptr);
}

// -- unprotected mprotect --

pub unsafe fn unprotected_mprotect<T>(ptr: *mut T, prot: ::Prot) -> bool {
    let unprotected_ptr = unprotected_ptr_from_user_ptr(ptr);
    let base_ptr = unprotected_ptr.offset(-(PAGE_SIZE as isize * 2));
    let unprotected_size = ptr::read(base_ptr as *const usize);
    ::mprotect(unprotected_ptr, unprotected_size, prot)
}

#[cfg(all(unix, test))]
mod test {
    use std::mem;

    #[should_panic]
    #[test]
    fn mprotect_2_test() {
        use nix::sys::signal;

        super::ALLOC_INIT.call_once(|| unsafe { super::alloc_init() });

        extern fn sigsegv(_: i32) { panic!() }
        let sigaction = signal::SigAction::new(
            signal::SigHandler::Handler(sigsegv),
            signal::SA_SIGINFO,
            signal::SigSet::empty(),
        );
        unsafe { signal::sigaction(signal::SIGSEGV, &sigaction).ok() };

        let x: *mut u8 = unsafe { super::alloc_aligned(16 * mem::size_of::<u8>()) };
        unsafe { ::mprotect(x, 16 * mem::size_of::<u8>(), ::Prot::NoAccess) };

        unsafe { ::memzero(x, 16 * mem::size_of::<u8>()) }; // SIGSEGV!
    }
}
<commit_msg>add some doc<commit_after>use std::sync::{ Once, ONCE_INIT };
use std::intrinsics::abort;
use std::{ mem, ptr };
use rand::{ thread_rng, Rng };
use rand::os::OsRng;
use errno::{ Errno, set_errno };

const GARBAGE_VALUE: u8 = 0xd0;
const CANARY_SIZE: usize = 16;
static ALLOC_INIT: Once = ONCE_INIT;
static mut PAGE_SIZE: usize = 0;
static mut CANARY: [u8; CANARY_SIZE] = [0; CANARY_SIZE];

// -- alloc init --

unsafe fn alloc_init() {
    #[cfg(unix)] {
        let page_size = ::libc::sysconf(::libc::_SC_PAGESIZE);
        if page_size > 0 {
            PAGE_SIZE = page_size as usize;
        }
    }

    #[cfg(windows)] {
        let si = mem::uninitialized();
        ::kernel32::GetSystemInfo(si);
        PAGE_SIZE = ptr::read(si).dwPageSize as usize;
    }

    if PAGE_SIZE < CANARY_SIZE || PAGE_SIZE < mem::size_of::<usize>() {
        abort();
    }

    match OsRng::new() {
        Ok(mut rng) => rng.fill_bytes(&mut CANARY),
        Err(_) => thread_rng().fill_bytes(&mut CANARY)
    }
}

// -- aligned alloc/free --

#[cfg(unix)]
unsafe fn alloc_aligned<T>(size: usize) -> *mut T {
    let mut memptr = mem::uninitialized();
    match ::libc::posix_memalign(&mut memptr, PAGE_SIZE, size) {
        0 => memptr as *mut T,
        ::libc::EINVAL => panic!("EINVAL: invalid alignmen. {}", PAGE_SIZE),
        ::libc::ENOMEM => ptr::null_mut(),
        _ => unreachable!()
    }
}

#[cfg(windows)]
unsafe fn alloc_aligned<T>(size: usize) -> *mut T {
    ::kernel32::VirtualAlloc(
        ptr::null(),
        size,
        ::winapi::MEM_COMMIT | ::winapi::MEM_RESERVE,
        ::winapi::PAGE_READWRITE
    ) as *mut T
}

#[cfg(unix)]
unsafe fn free_aligned<T>(memptr: *mut T) {
    ::libc::free(memptr as *mut ::libc::c_void);
}

#[cfg(windows)]
unsafe fn free_aligned<T>(memptr: *mut T) {
    ::kernel32::VirtualFree(memptr as ::winapi::LPVOID, 0, ::winapi::MEM_RELEASE);
}

// -- malloc --

#[inline]
unsafe fn page_round(size: usize) -> usize {
    let page_mask = PAGE_SIZE - 1;
    (size + page_mask) & (!page_mask)
}

unsafe fn unprotected_ptr_from_user_ptr<T>(memptr: *const T) -> *mut T {
    let canary_ptr = memptr.offset(-(mem::size_of_val(&CANARY) as isize));
    let page_mask = PAGE_SIZE - 1;
    let unprotected_ptr_u = canary_ptr as usize & !page_mask;
    if unprotected_ptr_u <= PAGE_SIZE * 2 {
        abort();
    }
    unprotected_ptr_u as *mut T
}

unsafe fn _malloc<T>(size: usize) -> *mut T {
    ALLOC_INIT.call_once(|| alloc_init());

    if size >= ::std::usize::MAX - PAGE_SIZE * 4 {
        set_errno(Errno(::libc::ENOMEM));
        return ptr::null_mut();
    }
    if PAGE_SIZE <= mem::size_of_val(&CANARY) || PAGE_SIZE < mem::size_of::<usize>() {
        abort();
    }

    // aligned alloc ptr
    let size_with_canary = mem::size_of_val(&CANARY) + size;
    let unprotected_size = page_round(size_with_canary);
    let total_size = PAGE_SIZE + PAGE_SIZE + unprotected_size + PAGE_SIZE;
    let base_ptr = alloc_aligned(total_size);
    if base_ptr.is_null() {
        return base_ptr;
    }

    // canary offset
    let unprotected_ptr = base_ptr.offset(PAGE_SIZE as isize * 2);
    ::mprotect(base_ptr.offset(PAGE_SIZE as isize), PAGE_SIZE, ::Prot::NoAccess);
    ptr::copy(
        CANARY.as_ptr(),
        unprotected_ptr.offset(unprotected_size as isize) as *mut u8,
        mem::size_of_val(&CANARY)
    );

    // mprotect ptr
    ::mprotect(unprotected_ptr.offset(unprotected_size as isize), PAGE_SIZE, ::Prot::NoAccess);
    ::mlock(unprotected_ptr, unprotected_size);
    let canary_ptr = unprotected_ptr
        .offset(page_round(size_with_canary) as isize)
        .offset(-(size_with_canary as isize));
    let user_ptr = canary_ptr.offset(mem::size_of_val(&CANARY) as isize);
    ptr::copy(CANARY.as_ptr(), canary_ptr as *mut u8, mem::size_of_val(&CANARY));
    ptr::write(base_ptr as *mut usize, unprotected_size);
    ::mprotect(base_ptr, PAGE_SIZE, ::Prot::ReadOnly);

    debug_assert_eq!(unprotected_ptr_from_user_ptr(user_ptr), unprotected_ptr);

    user_ptr
}

/// Secure malloc.
pub unsafe fn malloc<T>(size: usize) -> *mut T {
    let memptr = _malloc(size);
    if !memptr.is_null() {
        ::memset(memptr, GARBAGE_VALUE as i32, size);
    }
    memptr
}

/// Alloc array.
pub unsafe fn allocarray<T>(count: usize, size: usize) -> *mut T {
    if count > mem::size_of::<usize>() && size >= ::std::usize::MAX / count {
        set_errno(Errno(::libc::ENOMEM));
        return ptr::null_mut();
    }
    malloc(count * size)
}

/// Secure free.
pub unsafe fn free<T>(memptr: *mut T) {
    if memptr.is_null() { return () };

    // get unprotected ptr
    let canary_ptr = memptr.offset(-(mem::size_of_val(&CANARY) as isize));
    let unprotected_ptr = unprotected_ptr_from_user_ptr(memptr);
    let base_ptr = unprotected_ptr.offset(-(PAGE_SIZE as isize * 2));
    let unprotected_size = ptr::read(base_ptr as *const usize);
    let total_size = PAGE_SIZE + PAGE_SIZE + unprotected_size + PAGE_SIZE;
    ::mprotect(base_ptr, total_size, ::Prot::ReadWrite);

    // check
    debug_assert!(::memcmp(canary_ptr as *const u8, CANARY.as_ptr(), mem::size_of_val(&CANARY)));
    debug_assert!(::memcmp(
        unprotected_ptr.offset(unprotected_size as isize) as *const u8,
        CANARY.as_ptr(),
        mem::size_of_val(&CANARY)
    ));

    // free
    ::munlock(unprotected_ptr, unprotected_size);
    free_aligned(base_ptr);
}

// -- unprotected mprotect --

/// Secure mprotect.
pub unsafe fn unprotected_mprotect<T>(ptr: *mut T, prot: ::Prot) -> bool {
    let unprotected_ptr = unprotected_ptr_from_user_ptr(ptr);
    let base_ptr = unprotected_ptr.offset(-(PAGE_SIZE as isize * 2));
    let unprotected_size = ptr::read(base_ptr as *const usize);
    ::mprotect(unprotected_ptr, unprotected_size, prot)
}

#[cfg(all(unix, test))]
mod test {
    use std::mem;

    #[should_panic]
    #[test]
    fn mprotect_2_test() {
        use nix::sys::signal;

        super::ALLOC_INIT.call_once(|| unsafe { super::alloc_init() });

        extern fn sigsegv(_: i32) { panic!() }
        let sigaction = signal::SigAction::new(
            signal::SigHandler::Handler(sigsegv),
            signal::SA_SIGINFO,
            signal::SigSet::empty(),
        );
        unsafe { signal::sigaction(signal::SIGSEGV, &sigaction).ok() };

        let x: *mut u8 = unsafe { super::alloc_aligned(16 * mem::size_of::<u8>()) };
        unsafe { ::mprotect(x, 16 * mem::size_of::<u8>(), ::Prot::NoAccess) };

        unsafe { ::memzero(x, 16 * mem::size_of::<u8>()) }; // SIGSEGV!
    }
}
<|endoftext|>


<commit_before>//
// zhtta.rs
//
// Running on Rust 0.8
//
// Starting code for PS3
//
// Note: it would be very unwise to run this server on a machine that is
// on the Internet and contains any sensitive files!
//
// University of Virginia - cs4414 Fall 2013
// Weilin Xu and David Evans
// Version 0.3

extern mod extra;

use std::rt::io::*;
use std::rt::io::net::ip::SocketAddr;
use std::io::println;
use std::cell::Cell;
use std::{os, str, io};
use extra::arc;
use std::comm::*;
use extra::priority_queue;

static PORT:    int = 4414;
static IP: &'static str = "127.0.0.1";

struct sched_msg {
    stream: Option<std::rt::io::net::tcp::TcpStream>,
    filepath: ~std::path::PosixPath,
    topPriority: int,
    fileSize: u64
}

impl Ord for sched_msg {
    fn lt(&self, other: &sched_msg) -> bool {

        let mut retVal:bool = false;

        //If self has a lower geographic priority, then it is less than
        if (self.topPriority < other.topPriority) {
            retVal = true;
        }
        else if (self.topPriority == other.topPriority) {

            //If the two have equal priorities, then we will sort by file size
            //Smaller files get a higher priority
            if (self.fileSize > other.fileSize) {
                retVal = true;
            }
        }

        retVal
    }
}

fn main() {
    //let req_vec: ~[sched_msg] = ~[];
    let req_vec: ~priority_queue::PriorityQueue<sched_msg> = ~priority_queue::PriorityQueue::new();
    let shared_req_vec = arc::RWArc::new(req_vec);
    let add_vec = shared_req_vec.clone();
    let take_vec = shared_req_vec.clone();

    let (port, chan) = stream();
    let chan = SharedChan::new(chan);

    let safe_count: ~[uint] = ~[0];
    let shared_visit = arc::RWArc::new(safe_count);

    // dequeue file requests, and send responses.
    // FIFO
    do spawn {
        let (sm_port, sm_chan) = stream();

        // a task for sending responses.
        do spawn {
            loop {
                let mut tf: sched_msg = sm_port.recv(); // wait for the dequeued request to handle
                match io::read_whole_file(tf.filepath) { // killed if file size is larger than memory size.
                    Ok(file_data) => {
                        println(fmt!("begin serving file [%?]", tf.filepath));
                        // A web server should always reply a HTTP header for any legal HTTP request.
                        //tf.stream.write("HTTP/1.1 200 OK\r\nContent-Type: application/octet-stream; charset=UTF-8\r\n\r\n".as_bytes());

                        /*
                        let fileType = match tf.filepath.filetype() {
                            Some(s) => s,
                            None => &""
                        };

                        let httpHeader: ~str = match fileType {
                            ".html" | ".htm" | ".php" => ~"HTTP/1.1 200 OK\r\nContent-Type: text/html; charset=UTF-8\r\n\r\n",
                            _ => ~"HTTP/1.1 200 OK\r\nContent-Type: application/octet-stream; charset=UTF-8\r\n\r\n"
                        };

                        tf.stream.write(httpHeader.as_bytes());

                        */

                        //tf.stream.flush();
                        tf.stream.write(file_data);
                        println(fmt!("finish file [%?]", tf.filepath));
                    }
                    Err(err) => {
                        println("ERROR IN SEND");
                        println(err);
                    }
                }
            }
        }

        loop {
            port.recv(); // wait for arrving notification
            do take_vec.write |vec| {
                if ((*vec).len() > 0) {
                    // LIFO didn't make sense in service scheduling, so we modify it as FIFO by using shift_opt() rather than pop().
                    //let tf_opt: Option<sched_msg> = (*vec).shift_opt();
                    //let tf = tf_opt.unwrap();

                    //Since we are using a priority queue, we will use pop()
                    let tf = (*vec).pop();
                    println(fmt!("shift from queue, size: %ud", (*vec).len()));
                    sm_chan.send(tf); // send the request to send-response-task to serve.
                }
            }
        }
    }

    let ip = match FromStr::from_str(IP) { Some(ip) => ip,
                                           None => { println(fmt!("Error: Invalid IP address <%s>", IP));
                                                     return;},
                                         };

    let socket = net::tcp::TcpListener::bind(SocketAddr {ip: ip, port: PORT as u16});

    println(fmt!("Listening on %s:%d ...", ip.to_str(), PORT));
    let mut acceptor = socket.listen().unwrap();

    loop {
        let stream = acceptor.accept();
        let stream = Cell::new(stream);

        // Start a new task to handle the each connection
        let child_chan = chan.clone();
        let child_add_vec = add_vec.clone();
        let child_arc = shared_visit.clone();
        do spawn {

            let mut actual_count:uint = 0;

            //We will use the cloned RWARC
            do child_arc.write |count_vec| {

                let prev_count = (*count_vec).pop();
                let new_count: uint = prev_count + 1;

                actual_count = new_count;

                (*count_vec).push(new_count);
            }

            let stream = stream.take();

            match stream {
                Some(s) => {

                    let mut stream = s;

                    let mut buf = [0, ..500];
                    stream.read(buf);
                    let request_str = str::from_utf8(buf);

                    let req_group : ~[&str]= request_str.splitn_iter(' ', 3).collect();
                    if req_group.len() > 2 {
                        let path = req_group[1];
                        println(fmt!("Request for path: \n%?", path));

                        let file_path = ~os::getcwd().push(path.replace("/../", ""));
                        if !os::path_exists(file_path) || os::path_is_dir(file_path) {
                            println(fmt!("Request received:\n%s", request_str));
                            let response: ~str = fmt!(
                                "HTTP/1.1 200 OK\r\nContent-Type: text/html; charset=UTF-8\r\n\r\n
                                 <doctype !html><title>Hello, Rust!</title>

Greetings, Krusty!

Visitor count: %u

                                 \r\n", actual_count);

                            stream.write(response.as_bytes());
                        }
                        else {
                            // Requests scheduling

                            let mut streamPriority: int = 0;

                            let ipStr: ~str = match (stream).peer_name() {
                                Some(pr) => pr.ip.to_str(),
                                None => ~"0.0.0.0"
                            };

                            let ipSplit: ~[~str] = ipStr.split_iter('.').filter(|&x| x != "")
                                 .map(|x| x.to_owned()).collect();

                            if ( (ipSplit[0] == ~"127" && ipSplit[1] == ~"0") || (ipSplit[0] == ~"128" && ipSplit[1] == ~"143")
                                || (ipSplit[0] == ~"137" && ipSplit[1] == ~"54") ) {

                                streamPriority = 1;
                            }

                            let fileType = match file_path.filetype() {
                                Some(s) => s,
                                None => &""
                            };

                            let httpHeader: ~str = match fileType {
                                ".html" | ".htm" | ".php" => ~"HTTP/1.1 200 OK\r\nContent-Type: text/html; charset=UTF-8\r\n\r\n",
                                _ => ~"HTTP/1.1 200 OK\r\nContent-Type: application/octet-stream; charset=UTF-8\r\n\r\n"
                            };

                            stream.write(httpHeader.as_bytes());

                            let fileInfo = match std::rt::io::file::stat(file_path) {
                                Some(s) => s,
                                None => fail!("Could not access file stats")
                            };

                            let msg: sched_msg = sched_msg{stream: Some(stream), filepath: file_path.clone(), topPriority: streamPriority, fileSize: fileInfo.size};
                            let (sm_port, sm_chan) = std::comm::stream();
                            sm_chan.send(msg);

                            do child_add_vec.write |vec| {
                                let msg = sm_port.recv();
                                (*vec).push(msg); // enqueue new request.
                                println("add to queue");
                            }
                            child_chan.send(""); //notify the new arriving request.
                            println(fmt!("get file request: %?", file_path));
                        }
                    }
                    println!("connection terminates")

                },
                None => ()
            }

        }
    }
}<commit_msg>finished checking 3<commit_after>//
// zhtta.rs
//
// Running on Rust 0.8
//
// Starting code for PS3
//
// Note: it would be very unwise to run this server on a machine that is
// on the Internet and contains any sensitive files!
//
// University of Virginia - cs4414 Fall 2013
// Weilin Xu and David Evans
// Version 0.3

extern mod extra;

use std::rt::io::*;
use std::rt::io::net::ip::SocketAddr;
use std::io::println;
use std::cell::Cell;
use std::{os, str, io};
use extra::arc;
use std::comm::*;
use extra::priority_queue;

static PORT:    int = 4414;
static IP: &'static str = "127.0.0.1";

struct sched_msg {
    stream: Option<std::rt::io::net::tcp::TcpStream>,
    filepath: ~std::path::PosixPath,
    topPriority: int,
    fileSize: u64
}

impl Ord for sched_msg {
    fn lt(&self, other: &sched_msg) -> bool {

        let mut retVal:bool = false;

        //If self has a lower geographic priority, then it is less than
        if (self.topPriority < other.topPriority) {
            retVal = true;
        }
        else if (self.topPriority == other.topPriority) {

            //If the two have equal priorities, then we will sort by file size
            //Smaller files get a higher priority
            if (self.fileSize > other.fileSize) {
                retVal = true;
            }
        }

        retVal
    }
}

fn main() {
    //let req_vec: ~[sched_msg] = ~[];
    let req_vec: ~priority_queue::PriorityQueue<sched_msg> = ~priority_queue::PriorityQueue::new();
    let shared_req_vec = arc::RWArc::new(req_vec);
    let add_vec = shared_req_vec.clone();
    let take_vec = shared_req_vec.clone();

    let (port, chan) = stream();
    let chan = SharedChan::new(chan);

    let safe_count: ~[uint] = ~[0];
    let shared_visit = arc::RWArc::new(safe_count);

    // dequeue file requests, and send responses.
    // FIFO
    do spawn {
        let (sm_port, sm_chan) = stream();

        // a task for sending responses.
        do spawn {
            loop {
                let mut tf: sched_msg = sm_port.recv(); // wait for the dequeued request to handle
                match io::read_whole_file(tf.filepath) { // killed if file size is larger than memory size.
                    Ok(file_data) => {
                        println(fmt!("begin serving file [%?]", tf.filepath));

                        /*
                        let fileType = match tf.filepath.filetype() {
                            Some(s) => s,
                            None => &""
                        };

                        let httpHeader: ~str = match fileType {
                            ".html" | ".htm" | ".php" => ~"HTTP/1.1 200 OK\r\nContent-Type: text/html; charset=UTF-8\r\n\r\n",
                            _ => ~"HTTP/1.1 200 OK\r\nContent-Type: application/octet-stream; charset=UTF-8\r\n\r\n"
                        };

                        tf.stream.write(httpHeader.as_bytes());
                        tf.stream.flush();
                        */

                        tf.stream.write(file_data);
                        println(fmt!("finish file [%?]", tf.filepath));
                    }
                    Err(err) => {
                        println("ERROR IN SEND");
                        println(err);
                    }
                }
            }
        }

        loop {
            port.recv(); // wait for arrving notification
            do take_vec.write |vec| {
                if ((*vec).len() > 0) {

                    //Since we are using a priority queue, we will use pop()
                    let tf = (*vec).pop();
                    println(fmt!("shift from queue, size: %ud", (*vec).len()));
                    sm_chan.send(tf); // send the request to send-response-task to serve.
                }
            }
        }
    }

    let ip = match FromStr::from_str(IP) { Some(ip) => ip,
                                           None => { println(fmt!("Error: Invalid IP address <%s>", IP));
                                                     return;},
                                         };

    let socket = net::tcp::TcpListener::bind(SocketAddr {ip: ip, port: PORT as u16});

    println(fmt!("Listening on %s:%d ...", ip.to_str(), PORT));
    let mut acceptor = socket.listen().unwrap();

    loop {
        let stream = acceptor.accept();
        let stream = Cell::new(stream);

        // Start a new task to handle the each connection
        let child_chan = chan.clone();
        let child_add_vec = add_vec.clone();
        let child_arc = shared_visit.clone();
        do spawn {

            let mut actual_count:uint = 0;

            //We will use the cloned RWARC
            do child_arc.write |count_vec| {

                let prev_count = (*count_vec).pop();
                let new_count: uint = prev_count + 1;

                actual_count = new_count;

                (*count_vec).push(new_count);
            }

            let stream = stream.take();

            match stream {
                Some(s) => {

                    let mut stream = s;

                    let mut buf = [0, ..500];
                    stream.read(buf);
                    let request_str = str::from_utf8(buf);

                    let req_group : ~[&str]= request_str.splitn_iter(' ', 3).collect();
                    if req_group.len() > 2 {
                        let path = req_group[1];
                        println(fmt!("Request for path: \n%?", path));

                        let file_path = ~os::getcwd().push(path.replace("/../", ""));
                        if !os::path_exists(file_path) || os::path_is_dir(file_path) {
                            println(fmt!("Request received:\n%s", request_str));
                            let response: ~str = fmt!(
                                "HTTP/1.1 200 OK\r\nContent-Type: text/html; charset=UTF-8\r\n\r\n
                                 <doctype !html><title>Hello, Rust!</title>

Greetings, Krusty!

Visitor count: %u

                                 \r\n", actual_count);

                            stream.write(response.as_bytes());
                        }
                        else {
                            // Requests scheduling

                            let mut streamPriority: int = 0;

                            //Retrieving the requesting IP address
                            let ipStr: ~str = match (stream).peer_name() {
                                Some(pr) => pr.ip.to_str(),
                                None => ~"0.0.0.0"
                            };

                            //Split the IP address by '.' so that we can compare
                            let ipSplit: ~[~str] = ipStr.split_iter('.').filter(|&x| x != "")
                                 .map(|x| x.to_owned()).collect();

                            //Assign priority based on geography or if localhost
                            if ( (ipSplit[0] == ~"127" && ipSplit[1] == ~"0") || (ipSplit[0] == ~"128" && ipSplit[1] == ~"143")
                                || (ipSplit[0] == ~"137" && ipSplit[1] == ~"54") ) {

                                streamPriority = 1;
                            }

                            //Get filetype in order to determine HTTP header type
                            let fileType = match file_path.filetype() {
                                Some(s) => s,
                                None => &""
                            };

                            //In order to optimize for the benchmark, we will send the HTTP header quickly before adding to the queue
                            let httpHeader: ~str = match fileType {
                                ".html" | ".htm" | ".php" => ~"HTTP/1.1 200 OK\r\nContent-Type: text/html; charset=UTF-8\r\n\r\n",
                                _ => ~"HTTP/1.1 200 OK\r\nContent-Type: application/octet-stream; charset=UTF-8\r\n\r\n"
                            };

                            stream.write(httpHeader.as_bytes());
                            stream.flush();

                            //Retrieve file info for additional latency fixes
                            let fileInfo = match std::rt::io::file::stat(file_path) {
                                Some(s) => s,
                                None => fail!("Could not access file stats")
                            };

                            let msg: sched_msg = sched_msg{stream: Some(stream), filepath: file_path.clone(), topPriority: streamPriority, fileSize: fileInfo.size};
                            let (sm_port, sm_chan) = std::comm::stream();
                            sm_chan.send(msg);

                            do child_add_vec.write |vec| {
                                let msg = sm_port.recv();
                                (*vec).push(msg); // enqueue new request.
                                println("add to queue");
                            }
                            child_chan.send(""); //notify the new arriving request.
                            println(fmt!("get file request: %?", file_path));
                        }
                    }
                    println!("connection terminates")

                },
                None => ()
            }

        }
    }
}<|endoftext|>