<commit_before>use std::env;
use std::path::{Path, PathBuf};
use std::default::Default;
use std::collections::{BTreeMap, BTreeSet};

use container::mount::{bind_mount};
use container::util::clean_dir;
use config::Config;
use config::Container;
use config::Settings;
use config::builders::PipSettings;
use super::capsule;
use super::packages;
use super::timer;
use path_util::{PathExt, ToRelative};
use file_util::create_dir;
use process_util::PROXY_ENV_VARS;

pub struct Context<'a> {
    pub config: &'a Config,
    pub container_name: String,
    pub container_config: &'a Container,
    pub ensure_dirs: BTreeSet<PathBuf>,
    pub empty_dirs: BTreeSet<PathBuf>,
    pub remove_dirs: BTreeSet<PathBuf>,
    pub cache_dirs: BTreeMap<PathBuf, String>,
    pub environ: BTreeMap<String, String>,

    pub settings: Settings,
    pub pip_settings: PipSettings,
    pub capsule: capsule::State,
    pub packages: BTreeSet<String>,
    pub build_deps: BTreeSet<String>,
    pub featured_packages: BTreeSet<packages::Package>,
    pub timelog: timer::TimeLog,
}

impl<'a> Context<'a> {
    pub fn new<'x>(cfg: &'x Config, name: String,
        container: &'x Container, settings: Settings)
        -> Context<'x>
    {
        let mut env: BTreeMap<String, String> = vec!(
            ("TERM".to_string(), "dumb".to_string()),
            ("HOME".to_string(), "/tmp".to_string()),
            ("PATH".to_string(),
             "/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin"
             .to_string()),
            ).into_iter().collect();
        if settings.proxy_env_vars {
            for k in &PROXY_ENV_VARS {
                if let Ok(v) = env::var(k) {
                    env.insert(k.to_string(), v);
                }
            }
        }
        return Context {
            config: cfg,
            container_name: name,
            container_config: container,
            ensure_dirs: vec!(
                PathBuf::from("proc"),
                PathBuf::from("sys"),
                PathBuf::from("dev"),
                PathBuf::from("work"),
                PathBuf::from("tmp"),
                ).into_iter().collect(),
            empty_dirs: vec!(
                PathBuf::from("tmp"),
                PathBuf::from("var/tmp"),
                ).into_iter().collect(),
            remove_dirs: vec!(
                ).into_iter().collect(),
            cache_dirs: vec!(
                ).into_iter().collect(),
            environ: env,
            settings: settings,
            pip_settings: Default::default(),
            capsule: Default::default(),
            packages: BTreeSet::new(),
            build_deps: BTreeSet::new(),
            featured_packages: BTreeSet::new(),
            timelog: timer::TimeLog::start(
                    &Path::new("/vagga/container/timings.log"))
                .map_err(|e| format!("Can't write timelog: {}", e))
                .unwrap(),
        };
    }

    pub fn add_cache_dir(&mut self, path: &Path, name: String)
        -> Result<(), String>
    {
        assert!(path.is_absolute());
        let path = path.rel();
        if self.cache_dirs.insert(path.to_path_buf(), name.clone()).is_none() {
            let cache_dir = Path::new("/vagga/cache").join(&name);
            if !cache_dir.exists() {
                try_msg!(create_dir(&cache_dir, false),
                     "Error creating cache dir: {err}");
            }
            let path = Path::new("/vagga/root").join(path);
            try_msg!(create_dir(&path, true),
                 "Error creating cache dir: {err}");
            try!(clean_dir(&path, false));
            try!(bind_mount(&cache_dir, &path));
        }
        return Ok(());
    }

    pub fn add_remove_dir(&mut self, path: &Path) {
        assert!(path.is_absolute());
        self.remove_dirs.insert(path.rel().to_path_buf());
    }

    pub fn add_empty_dir(&mut self, path: &Path) {
        assert!(path.is_absolute());
        self.empty_dirs.insert(path.rel().to_path_buf());
    }

    pub fn add_ensure_dir(&mut self, path: &Path) {
        assert!(path.is_absolute());
        self.ensure_dirs.insert(path.rel().to_path_buf());
    }

}
<commit_msg>Ensure that /run is always created in container<commit_after>use std::env;
use std::path::{Path, PathBuf};
use std::default::Default;
use std::collections::{BTreeMap, BTreeSet};

use container::mount::{bind_mount};
use container::util::clean_dir;
use config::Config;
use config::Container;
use config::Settings;
use config::builders::PipSettings;
use super::capsule;
use super::packages;
use super::timer;
use path_util::{PathExt, ToRelative};
use file_util::create_dir;
use process_util::PROXY_ENV_VARS;

pub struct Context<'a> {
    pub config: &'a Config,
    pub container_name: String,
    pub container_config: &'a Container,
    pub ensure_dirs: BTreeSet<PathBuf>,
    pub empty_dirs: BTreeSet<PathBuf>,
    pub remove_dirs: BTreeSet<PathBuf>,
    pub cache_dirs: BTreeMap<PathBuf, String>,
    pub environ: BTreeMap<String, String>,

    pub settings: Settings,
    pub pip_settings: PipSettings,
    pub capsule: capsule::State,
    pub packages: BTreeSet<String>,
    pub build_deps: BTreeSet<String>,
    pub featured_packages: BTreeSet<packages::Package>,
    pub timelog: timer::TimeLog,
}

impl<'a> Context<'a> {
    pub fn new<'x>(cfg: &'x Config, name: String,
        container: &'x Container, settings: Settings)
        -> Context<'x>
    {
        let mut env: BTreeMap<String, String> = vec!(
            ("TERM".to_string(), "dumb".to_string()),
            ("HOME".to_string(), "/tmp".to_string()),
            ("PATH".to_string(),
             "/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin"
             .to_string()),
            ).into_iter().collect();
        if settings.proxy_env_vars {
            for k in &PROXY_ENV_VARS {
                if let Ok(v) = env::var(k) {
                    env.insert(k.to_string(), v);
                }
            }
        }
        return Context {
            config: cfg,
            container_name: name,
            container_config: container,
            ensure_dirs: vec!(
                PathBuf::from("proc"),
                PathBuf::from("sys"),
                PathBuf::from("dev"),
                PathBuf::from("work"),
                PathBuf::from("tmp"),
                PathBuf::from("run"),
                ).into_iter().collect(),
            empty_dirs: vec!(
                PathBuf::from("tmp"),
                PathBuf::from("var/tmp"),
                ).into_iter().collect(),
            remove_dirs: vec!(
                ).into_iter().collect(),
            cache_dirs: vec!(
                ).into_iter().collect(),
            environ: env,
            settings: settings,
            pip_settings: Default::default(),
            capsule: Default::default(),
            packages: BTreeSet::new(),
            build_deps: BTreeSet::new(),
            featured_packages: BTreeSet::new(),
            timelog: timer::TimeLog::start(
                    &Path::new("/vagga/container/timings.log"))
                .map_err(|e| format!("Can't write timelog: {}", e))
                .unwrap(),
        };
    }

    pub fn add_cache_dir(&mut self, path: &Path, name: String)
        -> Result<(), String>
    {
        assert!(path.is_absolute());
        let path = path.rel();
        if self.cache_dirs.insert(path.to_path_buf(), name.clone()).is_none() {
            let cache_dir = Path::new("/vagga/cache").join(&name);
            if !cache_dir.exists() {
                try_msg!(create_dir(&cache_dir, false),
                     "Error creating cache dir: {err}");
            }
            let path = Path::new("/vagga/root").join(path);
            try_msg!(create_dir(&path, true),
                 "Error creating cache dir: {err}");
            try!(clean_dir(&path, false));
            try!(bind_mount(&cache_dir, &path));
        }
        return Ok(());
    }

    pub fn add_remove_dir(&mut self, path: &Path) {
        assert!(path.is_absolute());
        self.remove_dirs.insert(path.rel().to_path_buf());
    }

    pub fn add_empty_dir(&mut self, path: &Path) {
        assert!(path.is_absolute());
        self.empty_dirs.insert(path.rel().to_path_buf());
    }

    pub fn add_ensure_dir(&mut self, path: &Path) {
        assert!(path.is_absolute());
        self.ensure_dirs.insert(path.rel().to_path_buf());
    }

}
<|endoftext|>


<commit_before>// encoding.rs

// import items in baseNN.rs
import Base16         = baseNN::Base16;
import Base32         = baseNN::Base32;
import Base64         = baseNN::Base64;

// import constants in baseNN/base{16,32,64}.rs
import BASE16         = baseNN::base16::BASE16;
import BASE32         = baseNN::base32::BASE32_STD;
import BASE64         = baseNN::base64::BASE64_STD;
import BASE32_HEX     = baseNN::base32::BASE32_HEX;
import BASE64_URLSAFE = baseNN::base64::BASE64_URL;

// import structs in baseNN/base{16,32,64}.rs
import Base16Writer   = baseNN::base16::Base16Writer;
import Base32Writer   = baseNN::base32::Base32Writer;
import Base64Writer   = baseNN::base64::Base64Writer;
import Base32Reader   = baseNN::base32::Base32Reader;
import Base64Reader   = baseNN::base64::Base64Reader;

//===-------------------------------------------------------------------===//
//                              b a s e N N
//===-------------------------------------------------------------------===//

// XXX this line is required to resolve traits in baseNN.
export baseNN;

// export all BaseNN traits
export Base16;
export Base32;
export Base64;

// export constants
export BASE16;
export BASE32;
export BASE64;
export BASE32_HEX;
export BASE64_URLSAFE;

// export structs
export Base16Writer;
export Base32Writer;
export Base64Writer;
export Base32Reader;
export Base64Reader;

<commit_msg>encoding: BASE64_URLSAFE -> BASE64_URL<commit_after>// encoding.rs

// import items in baseNN.rs
import Base16         = baseNN::Base16;
import Base32         = baseNN::Base32;
import Base64         = baseNN::Base64;

// import constants in baseNN/base{16,32,64}.rs
import BASE16         = baseNN::base16::BASE16;
import BASE32         = baseNN::base32::BASE32_STD;
import BASE64         = baseNN::base64::BASE64_STD;
import BASE32_HEX     = baseNN::base32::BASE32_HEX;
import BASE64_URL     = baseNN::base64::BASE64_URL;

// import structs in baseNN/base{16,32,64}.rs
import Base16Writer   = baseNN::base16::Base16Writer;
import Base32Writer   = baseNN::base32::Base32Writer;
import Base64Writer   = baseNN::base64::Base64Writer;
import Base32Reader   = baseNN::base32::Base32Reader;
import Base64Reader   = baseNN::base64::Base64Reader;

//===-------------------------------------------------------------------===//
//                              b a s e N N
//===-------------------------------------------------------------------===//

// XXX this line is required to resolve traits in baseNN.
export baseNN;

// export all BaseNN traits
export Base16;
export Base32;
export Base64;

// export constants
export BASE16;
export BASE32;
export BASE64;
export BASE32_HEX;
export BASE64_URL;

// export structs
export Base16Writer;
export Base32Writer;
export Base64Writer;
export Base32Reader;
export Base64Reader;

<|endoftext|>


<commit_before>use name::*;
use art::*;
use std::hash;
use std::hash::Hash;
use std::result::Result;
use std::collections::HashMap;

//
// Cmd C ::= Skip
//        | C ; C
//        | While E C
//        | x := E
//
//     C ::= .. | Amb C C
//
// Exp E ::=
//        | x
//        | N
//        | E + E
//

/// Nominal, artful commands for the while language
#[derive(Show,Hash,PartialEq,Eq)]
pub enum Cmd<'x> {
    Skip,
    Seq(CmdB<'x>, CmdB<'x>),
    While(ExpB<'x>, CmdB<'x>),
    Assign(Var<'x>, ExpB<'x>),
    // ----
    Name(Name, CmdB<'x>),
    Art(Art<'x,CmdB<'x>>),
}
pub type CmdB<'x> = Box<Cmd<'x>>;

/// Expressions for the while language
#[derive(Show,Hash,PartialEq,Eq)]
pub enum Exp<'x> {
    Var(Var<'x>),
    Plus(ExpB<'x>,ExpB<'x>),
    Num(int),
}
pub type ExpB<'x> = Box<Exp<'x>>;

/// Variables for the while language
pub type Var<'x> = String;

pub trait Heap<L:Hash+Eq, V> : Sized {
    fn empty () -> Self ;
    fn update (Self, L, V) -> Self ;
    fn select (Self, L) -> (Self, Option<V>) ;
}

pub struct HashMapHeap { hm : HashMap<String, int> }

impl Heap<String, int> for HashMapHeap {
    fn empty () -> Self {
        HashMapHeap { hm : HashMap::new() }
    }
    fn update (hm:Self, l:String, v:int) -> Self {
        panic!("")
    }
    fn select (hm:Self, l:String) -> (Self, Option<int>) {
        panic!("")
    }
}

pub enum Error {
    VarNotDef(String),
}

pub fn eval_exp<'x>
    (heap:HashMapHeap, exp:Exp<'x>)
     -> (HashMapHeap, Result<int, Error>)
{
    match exp {
        Exp::Num(n) => (heap, Ok(n)),
        Exp::Var(v) => {
            let (heap,q) = Heap::select(heap, v.clone()) ;
            match q {
                None => (heap, Err(Error::VarNotDef(v))),
                Some(val) => (heap, Ok(val))
            }},

        Exp::Plus(l, r) => {
            let (heap, l) = eval_exp(heap, *l) ;
            let (heap, r) = eval_exp(heap, *r) ;
            match (l,r) {
                (Ok(l), Ok(r)) => (heap, Ok(l+r)),
                (Err(l), _) => (heap, Err(l)),
                (_, Err(r)) => (heap, Err(r))
            }},
    }
}

#[test]
pub fn doit () {
    println!("while lang!");
}
<commit_msg>finished spec for challenge problem language: 'while' (equiv 'imp')<commit_after>use name::*;
use art::*;
use std::hash;
use std::hash::Hash;
use std::result::Result;
use std::collections::HashMap;

//
// Cmd C ::= Skip
//        | C ; C
//        | While E C
//        | x := E
//
//     C ::= .. | Amb C C
//
// Exp E ::=
//        | x
//        | N
//        | E + E
//        | E < E
//

/// Nominal, artful commands for the while language
#[derive(Show,Hash,PartialEq,Eq,Clone)]
pub enum Cmd<'x> {
    Skip,
    //Continue,
    Assign(Var<'x>, ExpB<'x>),
    Seq(CmdB<'x>, CmdB<'x>),
    // WHILE-loops:
    While(ExpB<'x>, CmdB<'x>),
    If(ExpB<'x>, CmdB<'x>, CmdB<'x>),
    Again(ExpB<'x>, CmdB<'x>),
    Break,
    // ----
    Name(Name, CmdB<'x>),
    Art(Art<'x,CmdB<'x>>),
}
pub type CmdB<'x> = Box<Cmd<'x>>;

pub enum CmdCxt<'x> {
    Empty,
    Seq(CmdCxtB<'x>, CmdB<'x>),
    // ----
    Name(Name, CmdCxtB<'x>),
    Art(Art<'x,CmdCxtB<'x>>),
}
pub type CmdCxtB<'x> = Box<CmdCxt<'x>>;

/// Expressions for the while language
#[derive(Show,Hash,PartialEq,Eq,Clone)]
pub enum Exp<'x> {
    Var(Var<'x>),
    Plus(ExpB<'x>,ExpB<'x>),
    Less(ExpB<'x>,ExpB<'x>),
    Num(int),
}
pub type ExpB<'x> = Box<Exp<'x>>;

/// Variables for the while language
pub type Var<'x> = String;

pub trait Heap<L:Hash+Eq, V> {
    fn empty () -> Self ;
    fn update (Self, L, V) -> Self ;
    fn select (Self, L) -> (Self, Option<V>) ;
}

pub struct HashMapHeap { hm : HashMap<String, int> }

impl Heap<String, int> for HashMapHeap {
    fn empty () -> Self {
        HashMapHeap { hm : HashMap::new() }
    }
    fn update (hm:Self, l:String, v:int) -> Self {
        let mut hm = hm.hm.clone () ; // TODO: Avoid this clone.
        hm.insert(l,v) ;
        HashMapHeap{ hm : hm }
    }
    fn select (hm:Self, l:String) -> (Self, Option<int>) {
        let hmq = hm.hm.clone () ; // TODO: Avoid this clone.
        let x = hmq.get(&l) ;
        (hm, match x { None => None, Some(x) => Some(*x) })
    }
}

pub enum Error {
    VarNotDef(String),
}

pub fn eval_exp<'x>
    (heap:HashMapHeap, exp:Exp<'x>)
     -> (HashMapHeap, Result<int, Error>)
{
    match exp {
        Exp::Num(n) => (heap, Ok(n)),
        Exp::Var(v) => {
            let (heap,q) = Heap::select(heap, v.clone()) ;
            match q {
                None => (heap, Err(Error::VarNotDef(v))),
                Some(val) => (heap, Ok(val))
            }},

        Exp::Plus(l, r) => {
            let (heap, l) = eval_exp(heap, *l) ;
            let (heap, r) = eval_exp(heap, *r) ;
            match (l,r) {
                (Ok(l), Ok(r)) => (heap, Ok(l+r)),
                (Err(l), _) => (heap, Err(l)),
                (_, Err(r)) => (heap, Err(r))
            }},

        Exp::Less(l, r) => {
            let (heap, l) = eval_exp(heap, *l) ;
            let (heap, r) = eval_exp(heap, *r) ;
            match (l,r) {
                (Ok(l), Ok(r)) => (heap, Ok(if l < r { 1 } else { 0 })),
                (Err(l), _) => (heap, Err(l)),
                (_, Err(r)) => (heap, Err(r))
            }},

    }
}

fn do_break<'x> (cxt:CmdCxt<'x>) -> CmdCxt<'x> {
    match cxt {
        CmdCxt::Empty => CmdCxt::Empty, // Degenerate basecase.
        CmdCxt::Seq(cxt, cmd) =>
            match *cmd {
                Cmd::Again(_,_) => *cxt, // Normal basecase.
                // - - - - - - -
                // Keep searching for basecase...
                Cmd::Art(art) => do_break (CmdCxt::Seq(cxt, force(art))),
                // Cmd::Name(nm, cmd) => { do_break(CmdCxt::Seq(cxt, cmd)) },
                // - - - - - - -
                _ => do_break(*cxt),
            },
        // - - - - - - -
        CmdCxt::Name(_, cxt) => do_break(*cxt),
        CmdCxt::Art(art) => do_break(*force(art)),
    }
}

pub fn step_cmd<'x>
    (heap:HashMapHeap, cxt:CmdCxt<'x>, cmd:Cmd<'x>)
     -> (HashMapHeap, Option<(CmdCxt<'x>, Cmd<'x>)>)
{
    match (cxt, cmd) {
        (CmdCxt::Empty, Cmd::Skip) => (heap, None),
        (CmdCxt::Empty, Cmd::Break) => (heap, None),

        (CmdCxt::Seq(cxt, cmd), Cmd::Skip) => (heap, Some((*cxt, *cmd))),
        (cxt, Cmd::Break) => {
            let cxt = do_break(cxt) ;
            (heap, Some((cxt, Cmd::Skip)))
        }

        (cxt, Cmd::Seq(cmd1, cmd2)) => (heap, Some( (CmdCxt::Seq(box cxt, cmd2), *cmd1) )),

        // While becomes Again, which repeats again and again..
        // This case only fires when first entering the loop;
        // TODO: use it to extend the "loop path"
        (cxt, Cmd::While(exp, cmd)) => (heap, Some( (cxt, Cmd::Again(exp, cmd)) )),

        (cxt, Cmd::Again(exp, cmd)) =>
            (heap, Some( (CmdCxt::Seq(box cxt, box Cmd::Again(exp.clone(), cmd.clone())),
                          Cmd::If(exp, cmd, box Cmd::Break)
                          ) )),

        (cxt, Cmd::Assign(var, exp)) => {
            let (heap, r) = eval_exp (heap, *exp) ;
            match r {
                Err(error) => (heap, None),
                Ok(val) => {
                    let heap = Heap::update(heap, var, val) ;
                    (heap, Some(( cxt, Cmd::Skip )))
                }
            }
        },

        (cxt, Cmd::If(exp, cmd1, cmd2)) => {
            let (heap, r) = eval_exp (heap, *exp) ;
            match r {
                Err(error) => (heap, None),
                Ok(0) => { (heap, Some((cxt, *cmd2))) },
                Ok(_) => { (heap, Some((cxt, *cmd1))) },
            }
        },

        // - - - - - - - - -
        (cxt, Cmd::Art(art))         => { step_cmd(heap, cxt, *force(art)) },
        (cxt, Cmd::Name(nm, cmd))    => { step_cmd(heap, cxt, *cmd) }
        (CmdCxt::Art(art), cmd)      => { step_cmd(heap, *force(art), cmd) }
        (CmdCxt::Name(nm, cxt), cmd) => { step_cmd(heap, *cxt, cmd) }

    }
}

#[test]
pub fn doit () {
    println!("while lang!");
}
<|endoftext|>


<commit_before>use crate::thread::Time;
use crate::*;
use rustc_target::abi::{Align, Size};
use std::time::{Instant, SystemTime};

/// Implementation of the SYS_futex syscall.
pub fn futex<'tcx>(
    this: &mut MiriEvalContext<'_, 'tcx>,
    args: &[OpTy<'tcx, Tag>],
    dest: &PlaceTy<'tcx, Tag>,
) -> InterpResult<'tcx> {
    // The amount of arguments used depends on the type of futex operation.
    // The full futex syscall takes six arguments (excluding the syscall
    // number), which is also the maximum amount of arguments a linux syscall
    // can take on most architectures.
    // However, not all futex operations use all six arguments. The unused ones
    // may or may not be left out from the `syscall()` call.
    // Therefore we don't use `check_arg_count` here, but only check for the
    // number of arguments to fall within a range.
    if args.len() < 4 {
        throw_ub_format!(
            "incorrect number of arguments for `futex` syscall: got {}, expected at least 4",
            args.len()
        );
    }

    // The first three arguments (after the syscall number itself) are the same to all futex operations:
    //     (int *addr, int op, int val).
    // We checked above that these definitely exist.
    let addr = this.read_immediate(&args[1])?;
    let op = this.read_scalar(&args[2])?.to_i32()?;
    let val = this.read_scalar(&args[3])?.to_i32()?;

    let thread = this.get_active_thread();
    let addr_scalar = addr.to_scalar()?;

    let futex_private = this.eval_libc_i32("FUTEX_PRIVATE_FLAG")?;
    let futex_wait = this.eval_libc_i32("FUTEX_WAIT")?;
    let futex_wait_bitset = this.eval_libc_i32("FUTEX_WAIT_BITSET")?;
    let futex_wake = this.eval_libc_i32("FUTEX_WAKE")?;
    let futex_wake_bitset = this.eval_libc_i32("FUTEX_WAKE_BITSET")?;
    let futex_realtime = this.eval_libc_i32("FUTEX_CLOCK_REALTIME")?;

    // FUTEX_PRIVATE enables an optimization that stops it from working across processes.
    // Miri doesn't support that anyway, so we ignore that flag.
    match op & !futex_private {
        // FUTEX_WAIT: (int *addr, int op = FUTEX_WAIT, int val, const timespec *timeout)
        // Blocks the thread if *addr still equals val. Wakes up when FUTEX_WAKE is called on the same address,
        // or *timeout expires. `timeout == null` for an infinite timeout.
        //
        // FUTEX_WAIT_BITSET: (int *addr, int op = FUTEX_WAIT_BITSET, int val, const timespec *timeout, int *_ignored, unsigned int bitset)
        // This is identical to FUTEX_WAIT, except:
        //  - The timeout is absolute rather than relative.
        //  - You can specify the bitset to selecting what WAKE operations to respond to.
        op if op & !futex_realtime == futex_wait || op & !futex_realtime == futex_wait_bitset => {
            let wait_bitset = op & !futex_realtime == futex_wait_bitset;

            let bitset;

            if wait_bitset {
                if args.len() != 7 {
                    throw_ub_format!(
                        "incorrect number of arguments for `futex` syscall with `op=FUTEX_WAIT_BITSET`: got {}, expected 7",
                        args.len()
                    );
                }
                bitset = this.read_scalar(&args[6])?.to_u32()?;
            } else {
                if args.len() < 5 {
                    throw_ub_format!(
                        "incorrect number of arguments for `futex` syscall with `op=FUTEX_WAIT`: got {}, expected at least 5",
                        args.len()
                    );
                }
                bitset = u32::MAX;
            }

            if bitset == 0 {
                let einval = this.eval_libc("EINVAL")?;
                this.set_last_error(einval)?;
                this.write_scalar(Scalar::from_machine_isize(-1, this), dest)?;
                return Ok(());
            }

            // `deref_operand` but not actually dereferencing the ptr yet (it might be NULL!).
            let timeout = this.ref_to_mplace(&this.read_immediate(&args[4])?)?;
            let timeout_time = if this.ptr_is_null(timeout.ptr)? {
                None
            } else {
                this.check_no_isolation(
                    "`futex` syscall with `op=FUTEX_WAIT` and non-null timeout",
                )?;
                let duration = match this.read_timespec(&timeout)? {
                    Some(duration) => duration,
                    None => {
                        let einval = this.eval_libc("EINVAL")?;
                        this.set_last_error(einval)?;
                        this.write_scalar(Scalar::from_machine_isize(-1, this), dest)?;
                        return Ok(());
                    }
                };
                Some(if wait_bitset {
                    // FUTEX_WAIT_BITSET uses an absolute timestamp.
                    if op & futex_realtime != 0 {
                        Time::RealTime(SystemTime::UNIX_EPOCH.checked_add(duration).unwrap())
                    } else {
                        Time::Monotonic(this.machine.time_anchor.checked_add(duration).unwrap())
                    }
                } else {
                    // FUTEX_WAIT uses a relative timestamp.
                    if op & futex_realtime != 0 {
                        Time::RealTime(SystemTime::now().checked_add(duration).unwrap())
                    } else {
                        Time::Monotonic(Instant::now().checked_add(duration).unwrap())
                    }
                })
            };
            // Check the pointer for alignment and validity.
            // The API requires `addr` to be a 4-byte aligned pointer, and will
            // use the 4 bytes at the given address as an (atomic) i32.
            this.check_ptr_access_align(
                this.scalar_to_ptr(addr_scalar),
                Size::from_bytes(4),
                Align::from_bytes(4).unwrap(),
                CheckInAllocMsg::MemoryAccessTest,
            )?;
            // Read an `i32` through the pointer, regardless of any wrapper types.
            // It's not uncommon for `addr` to be passed as another type than `*mut i32`, such as `*const AtomicI32`.
            // FIXME: this fails if `addr` is not a pointer type.
            // The atomic ordering for futex(https://man7.org/linux/man-pages/man2/futex.2.html):
            //  "The load of the value of the futex word is an
            //   atomic memory access (i.e., using atomic machine instructions
            //   of the respective architecture).  This load, the comparison
            //   with the expected value, and starting to sleep are performed
            //   atomically and totally ordered with respect to other futex
            //   operations on the same futex word."
            // SeqCst is total order over all operations.
            // FIXME: check if this should be changed when weak memory orders are added.
            let futex_val = this
                .read_scalar_at_offset_atomic(
                    &addr.into(),
                    0,
                    this.machine.layouts.i32,
                    AtomicReadOp::SeqCst,
                )?
                .to_i32()?;
            if val == futex_val {
                // The value still matches, so we block the trait make it wait for FUTEX_WAKE.
                this.block_thread(thread);
                this.futex_wait(addr_scalar.to_machine_usize(this)?, thread, bitset);
                // Succesfully waking up from FUTEX_WAIT always returns zero.
                this.write_scalar(Scalar::from_machine_isize(0, this), dest)?;
                // Register a timeout callback if a timeout was specified.
                // This callback will override the return value when the timeout triggers.
                let dest = *dest;
                if let Some(timeout_time) = timeout_time {
                    this.register_timeout_callback(
                        thread,
                        timeout_time,
                        Box::new(move |this| {
                            this.unblock_thread(thread);
                            this.futex_remove_waiter(addr_scalar.to_machine_usize(this)?, thread);
                            let etimedout = this.eval_libc("ETIMEDOUT")?;
                            this.set_last_error(etimedout)?;
                            this.write_scalar(Scalar::from_machine_isize(-1, this), &dest)?;
                            Ok(())
                        }),
                    );
                }
            } else {
                // The futex value doesn't match the expected value, so we return failure
                // right away without sleeping: -1 and errno set to EAGAIN.
                let eagain = this.eval_libc("EAGAIN")?;
                this.set_last_error(eagain)?;
                this.write_scalar(Scalar::from_machine_isize(-1, this), dest)?;
            }
        }
        // FUTEX_WAKE: (int *addr, int op = FUTEX_WAKE, int val)
        // Wakes at most `val` threads waiting on the futex at `addr`.
        // Returns the amount of threads woken up.
        // Does not access the futex value at *addr.
        // FUTEX_WAKE_BITSET: (int *addr, int op = FUTEX_WAKE, int val, const timespect *_unused, int *_unused, unsigned int bitset)
        // Same as FUTEX_WAKE, but allows you to specify a bitset to select which threads to wake up.
        op if op == futex_wake || op == futex_wake_bitset => {
            let bitset;
            if op == futex_wake_bitset {
                if args.len() != 7 {
                    throw_ub_format!(
                        "incorrect number of arguments for `futex` syscall with `op=FUTEX_WAKE_BITSET`: got {}, expected 7",
                        args.len()
                    );
                }
                bitset = this.read_scalar(&args[6])?.to_u32()?;
            } else {
                bitset = u32::MAX;
            }
            if bitset == 0 {
                let einval = this.eval_libc("EINVAL")?;
                this.set_last_error(einval)?;
                this.write_scalar(Scalar::from_machine_isize(-1, this), dest)?;
                return Ok(());
            }
            let mut n = 0;
            for _ in 0..val {
                if let Some(thread) = this.futex_wake(addr_scalar.to_machine_usize(this)?, bitset) {
                    this.unblock_thread(thread);
                    this.unregister_timeout_callback_if_exists(thread);
                    n += 1;
                } else {
                    break;
                }
            }
            this.write_scalar(Scalar::from_machine_isize(n, this), dest)?;
        }
        op => throw_unsup_format!("Miri does not support `futex` syscall with op={}", op),
    }

    Ok(())
}
<commit_msg>Use `let = if;` instead of `let; if`.<commit_after>use crate::thread::Time;
use crate::*;
use rustc_target::abi::{Align, Size};
use std::time::{Instant, SystemTime};

/// Implementation of the SYS_futex syscall.
pub fn futex<'tcx>(
    this: &mut MiriEvalContext<'_, 'tcx>,
    args: &[OpTy<'tcx, Tag>],
    dest: &PlaceTy<'tcx, Tag>,
) -> InterpResult<'tcx> {
    // The amount of arguments used depends on the type of futex operation.
    // The full futex syscall takes six arguments (excluding the syscall
    // number), which is also the maximum amount of arguments a linux syscall
    // can take on most architectures.
    // However, not all futex operations use all six arguments. The unused ones
    // may or may not be left out from the `syscall()` call.
    // Therefore we don't use `check_arg_count` here, but only check for the
    // number of arguments to fall within a range.
    if args.len() < 4 {
        throw_ub_format!(
            "incorrect number of arguments for `futex` syscall: got {}, expected at least 4",
            args.len()
        );
    }

    // The first three arguments (after the syscall number itself) are the same to all futex operations:
    //     (int *addr, int op, int val).
    // We checked above that these definitely exist.
    let addr = this.read_immediate(&args[1])?;
    let op = this.read_scalar(&args[2])?.to_i32()?;
    let val = this.read_scalar(&args[3])?.to_i32()?;

    let thread = this.get_active_thread();
    let addr_scalar = addr.to_scalar()?;

    let futex_private = this.eval_libc_i32("FUTEX_PRIVATE_FLAG")?;
    let futex_wait = this.eval_libc_i32("FUTEX_WAIT")?;
    let futex_wait_bitset = this.eval_libc_i32("FUTEX_WAIT_BITSET")?;
    let futex_wake = this.eval_libc_i32("FUTEX_WAKE")?;
    let futex_wake_bitset = this.eval_libc_i32("FUTEX_WAKE_BITSET")?;
    let futex_realtime = this.eval_libc_i32("FUTEX_CLOCK_REALTIME")?;

    // FUTEX_PRIVATE enables an optimization that stops it from working across processes.
    // Miri doesn't support that anyway, so we ignore that flag.
    match op & !futex_private {
        // FUTEX_WAIT: (int *addr, int op = FUTEX_WAIT, int val, const timespec *timeout)
        // Blocks the thread if *addr still equals val. Wakes up when FUTEX_WAKE is called on the same address,
        // or *timeout expires. `timeout == null` for an infinite timeout.
        //
        // FUTEX_WAIT_BITSET: (int *addr, int op = FUTEX_WAIT_BITSET, int val, const timespec *timeout, int *_ignored, unsigned int bitset)
        // This is identical to FUTEX_WAIT, except:
        //  - The timeout is absolute rather than relative.
        //  - You can specify the bitset to selecting what WAKE operations to respond to.
        op if op & !futex_realtime == futex_wait || op & !futex_realtime == futex_wait_bitset => {
            let wait_bitset = op & !futex_realtime == futex_wait_bitset;

            let bitset = if wait_bitset {
                if args.len() != 7 {
                    throw_ub_format!(
                        "incorrect number of arguments for `futex` syscall with `op=FUTEX_WAIT_BITSET`: got {}, expected 7",
                        args.len()
                    );
                }
                this.read_scalar(&args[6])?.to_u32()?
            } else {
                if args.len() < 5 {
                    throw_ub_format!(
                        "incorrect number of arguments for `futex` syscall with `op=FUTEX_WAIT`: got {}, expected at least 5",
                        args.len()
                    );
                }
                u32::MAX
            };

            if bitset == 0 {
                let einval = this.eval_libc("EINVAL")?;
                this.set_last_error(einval)?;
                this.write_scalar(Scalar::from_machine_isize(-1, this), dest)?;
                return Ok(());
            }

            // `deref_operand` but not actually dereferencing the ptr yet (it might be NULL!).
            let timeout = this.ref_to_mplace(&this.read_immediate(&args[4])?)?;
            let timeout_time = if this.ptr_is_null(timeout.ptr)? {
                None
            } else {
                this.check_no_isolation(
                    "`futex` syscall with `op=FUTEX_WAIT` and non-null timeout",
                )?;
                let duration = match this.read_timespec(&timeout)? {
                    Some(duration) => duration,
                    None => {
                        let einval = this.eval_libc("EINVAL")?;
                        this.set_last_error(einval)?;
                        this.write_scalar(Scalar::from_machine_isize(-1, this), dest)?;
                        return Ok(());
                    }
                };
                Some(if wait_bitset {
                    // FUTEX_WAIT_BITSET uses an absolute timestamp.
                    if op & futex_realtime != 0 {
                        Time::RealTime(SystemTime::UNIX_EPOCH.checked_add(duration).unwrap())
                    } else {
                        Time::Monotonic(this.machine.time_anchor.checked_add(duration).unwrap())
                    }
                } else {
                    // FUTEX_WAIT uses a relative timestamp.
                    if op & futex_realtime != 0 {
                        Time::RealTime(SystemTime::now().checked_add(duration).unwrap())
                    } else {
                        Time::Monotonic(Instant::now().checked_add(duration).unwrap())
                    }
                })
            };
            // Check the pointer for alignment and validity.
            // The API requires `addr` to be a 4-byte aligned pointer, and will
            // use the 4 bytes at the given address as an (atomic) i32.
            this.check_ptr_access_align(
                this.scalar_to_ptr(addr_scalar),
                Size::from_bytes(4),
                Align::from_bytes(4).unwrap(),
                CheckInAllocMsg::MemoryAccessTest,
            )?;
            // Read an `i32` through the pointer, regardless of any wrapper types.
            // It's not uncommon for `addr` to be passed as another type than `*mut i32`, such as `*const AtomicI32`.
            // FIXME: this fails if `addr` is not a pointer type.
            // The atomic ordering for futex(https://man7.org/linux/man-pages/man2/futex.2.html):
            //  "The load of the value of the futex word is an
            //   atomic memory access (i.e., using atomic machine instructions
            //   of the respective architecture).  This load, the comparison
            //   with the expected value, and starting to sleep are performed
            //   atomically and totally ordered with respect to other futex
            //   operations on the same futex word."
            // SeqCst is total order over all operations.
            // FIXME: check if this should be changed when weak memory orders are added.
            let futex_val = this
                .read_scalar_at_offset_atomic(
                    &addr.into(),
                    0,
                    this.machine.layouts.i32,
                    AtomicReadOp::SeqCst,
                )?
                .to_i32()?;
            if val == futex_val {
                // The value still matches, so we block the trait make it wait for FUTEX_WAKE.
                this.block_thread(thread);
                this.futex_wait(addr_scalar.to_machine_usize(this)?, thread, bitset);
                // Succesfully waking up from FUTEX_WAIT always returns zero.
                this.write_scalar(Scalar::from_machine_isize(0, this), dest)?;
                // Register a timeout callback if a timeout was specified.
                // This callback will override the return value when the timeout triggers.
                let dest = *dest;
                if let Some(timeout_time) = timeout_time {
                    this.register_timeout_callback(
                        thread,
                        timeout_time,
                        Box::new(move |this| {
                            this.unblock_thread(thread);
                            this.futex_remove_waiter(addr_scalar.to_machine_usize(this)?, thread);
                            let etimedout = this.eval_libc("ETIMEDOUT")?;
                            this.set_last_error(etimedout)?;
                            this.write_scalar(Scalar::from_machine_isize(-1, this), &dest)?;
                            Ok(())
                        }),
                    );
                }
            } else {
                // The futex value doesn't match the expected value, so we return failure
                // right away without sleeping: -1 and errno set to EAGAIN.
                let eagain = this.eval_libc("EAGAIN")?;
                this.set_last_error(eagain)?;
                this.write_scalar(Scalar::from_machine_isize(-1, this), dest)?;
            }
        }
        // FUTEX_WAKE: (int *addr, int op = FUTEX_WAKE, int val)
        // Wakes at most `val` threads waiting on the futex at `addr`.
        // Returns the amount of threads woken up.
        // Does not access the futex value at *addr.
        // FUTEX_WAKE_BITSET: (int *addr, int op = FUTEX_WAKE, int val, const timespect *_unused, int *_unused, unsigned int bitset)
        // Same as FUTEX_WAKE, but allows you to specify a bitset to select which threads to wake up.
        op if op == futex_wake || op == futex_wake_bitset => {
            let bitset = if op == futex_wake_bitset {
                if args.len() != 7 {
                    throw_ub_format!(
                        "incorrect number of arguments for `futex` syscall with `op=FUTEX_WAKE_BITSET`: got {}, expected 7",
                        args.len()
                    );
                }
                this.read_scalar(&args[6])?.to_u32()?
            } else {
                u32::MAX
            };
            if bitset == 0 {
                let einval = this.eval_libc("EINVAL")?;
                this.set_last_error(einval)?;
                this.write_scalar(Scalar::from_machine_isize(-1, this), dest)?;
                return Ok(());
            }
            let mut n = 0;
            for _ in 0..val {
                if let Some(thread) = this.futex_wake(addr_scalar.to_machine_usize(this)?, bitset) {
                    this.unblock_thread(thread);
                    this.unregister_timeout_callback_if_exists(thread);
                    n += 1;
                } else {
                    break;
                }
            }
            this.write_scalar(Scalar::from_machine_isize(n, this), dest)?;
        }
        op => throw_unsup_format!("Miri does not support `futex` syscall with op={}", op),
    }

    Ok(())
}
<|endoftext|>


<commit_before>use std::rt::io::net::ip::{SocketAddr, IpAddr, Ipv4Addr};
use knob::Settings;

pub trait SocketSettings {
  fn socket(&self) -> SocketAddr;
  fn port(&self) -> u16;
  fn ip(&self) -> IpAddr;
}

impl SocketSettings for Settings {
  fn socket(&self) -> SocketAddr {
    do self.fetch_with("addr") |addr| {
      match addr {
        Some(socket_addr) => { socket_addr },
        None => {
          let port: u16 = self.port();
          let ip: IpAddr = self.ip();
          SocketAddr { ip: ip, port: port }
        }
      }
    }
  }

  fn port(&self) -> u16 {
    self.fetch("port").unwrap_or(8080)
  }

  fn ip(&self) -> IpAddr {
    self.fetch("ip").unwrap_or(Ipv4Addr(127,0,0,1))
  }
}

#[cfg(test)]
mod tests {
  use super::*;
  #[test]
  fn test_insert() {
    let mut settings = Settings::new();
    settings.set("port", "12345");
    settings.set("ip", "127.0.0.1");
    settings.socket();
  }
}<commit_msg>Fix tests<commit_after>use std::rt::io::net::ip::{SocketAddr, IpAddr, Ipv4Addr};
use knob::Settings;

pub trait SocketSettings {
  fn socket(&self) -> SocketAddr;
  fn port(&self) -> u16;
  fn ip(&self) -> IpAddr;
}

impl SocketSettings for Settings {
  fn socket(&self) -> SocketAddr {
    do self.fetch_with("addr") |addr| {
      match addr {
        Some(socket_addr) => { socket_addr },
        None => {
          let port: u16 = self.port();
          let ip: IpAddr = self.ip();
          SocketAddr { ip: ip, port: port }
        }
      }
    }
  }

  fn port(&self) -> u16 {
    self.fetch("port").unwrap_or(8080)
  }

  fn ip(&self) -> IpAddr {
    self.fetch("ip").unwrap_or(Ipv4Addr(127,0,0,1))
  }
}

#[cfg(test)]
mod tests {
  use super::*;
  use knob::Settings;

  #[test]
  fn test_insert() {
    let mut settings = Settings::new();
    settings.set("port", "12345");
    settings.set("ip", "127.0.0.1");
    let socket = settings.socket();
    assert_eq!(socket.to_str(), ~"127.0.0.1:12345")
  }
}<|endoftext|>


<commit_before>// TODO: data xrefs

use super::*;
use zydis;
use byteorder::{ByteOrder, LittleEndian};
use std::mem;

fn analyze_operand_xrefs(
    // used to determine valid addresses
    layout: &ModuleLayout,
    // used to detect RIP-relative addressing
    pc: zydis::enums::register::Register,
    // used to compute relative addresses
    rva: Rva,
    // used to compute relative addresses
    insn: &zydis::ffi::DecodedInstruction,
    // the operand to inspect
    op: &zydis::ffi::DecodedOperand,
) -> Result<Option<Rva>, Error> {
    match op.ty {
        zydis::enums::OperandType::Unused => Err(Error::NotImplemented("xref from unused register")),
        zydis::enums::OperandType::Register => {
            // like: CALL rbx
            // TODO: for now, don't index unresolved indirect branches
            Ok(None)
        }
        zydis::enums::OperandType::Memory => {
            // like: .text:0000000180001041 FF 15 D1 78 07 00      call    cs:__imp_RtlVirtualUnwind_0
            //           0x0000000000001041:                       call    [0x0000000000079980]
            if pc == op.mem.base && op.mem.disp.has_displacement && op.mem.scale == 0 {
                // RIP-relative
                // this is the default encoding on x64.
                // tools like IDA automatically compute and display the target.
                // CALL [RIP + 0x401000]
                if let zydis::enums::register::Register::NONE = op.mem.index {
                    let target =
                        (rva as i64
                        // TODO: cast from rva (u64) to i64 is lossy.
                        + op.mem.disp.displacement
                        + i64::from(insn.length)) as Rva;

                    if layout.is_rva_valid(target) {
                        Ok(Some(target))
                    } else {
                        // TODO: record this anomaly somewhere.
                        warn!("problem: invalid xref target: memory not in sections");
                        Ok(None)
                    }
                } else {
                    // unsupported
                    // like: CALL [RIP + 4*RCX + 0x401000] ??
                    println!("CALL [RIP + 4*RCX + 0x401000] ??");
                    Err(Error::NotImplemented("xref from RIP-relative, non-zero index memory"))
                }
            } else if op.mem.base == zydis::enums::register::Register::NONE {
                // like: CALL [0x401000] ??
                println!("TODO: other OperandType::Memory branch");
                Err(Error::NotImplemented("xref from non-RIP-relative memory"))
            } else {
                // like: CALL [rbx]
                // like: CALL [rbx + 0x10]
                // TODO: for now, don't index unresolved indirect branches
                Ok(None)
            }
        }
        zydis::enums::OperandType::Pointer => {
            println!("TODO: operand: pointer");
            Err(Error::NotImplemented("xref from pointer"))
        }
        zydis::enums::OperandType::Immediate => {
            if !op.imm.is_relative {
                println!("TODO: absolute immediate operand");
                Err(Error::NotImplemented("xref from absolute immediate"))
            } else {
                let imm = if op.imm.is_signed {
                    u64_i64(op.imm.value)
                } else {
                    op.imm.value as i64
                };

                // TODO: cast from rva (u64) to i64 is lossy.
                let dst = (rva as i64 + imm + i64::from(insn.length)) as Rva;
                if layout.is_rva_valid(dst) {
                    Ok(Some(dst))
                } else {
                    // TODO: record this anomaly somewhere.
                    debug!("problem: invalid xref target: relative immediate not in sections");
                    Ok(None)
                }
            }
        }
    }
}

pub fn analyze_insn_xrefs(
    layout: &ModuleLayout,
    pc: zydis::enums::register::Register,
    rva: Rva,
    insn: &zydis::ffi::DecodedInstruction,
) -> Result<Vec<Xref>, Error> {
    match insn.mnemonic {
        // see InstructionCategory
        // syscall, sysexit, sysret
        // vmcall, vmmcall
        zydis::enums::mnemonic::Mnemonic::CALL => {
            // a CALL always has an operand, so assume this is ok.
            let op = Workspace::get_first_operand(insn).unwrap();

            let fallthrough = Xref {
                src: rva,
                dst: rva + u64::from(insn.length),
                typ: XrefType::Fallthrough,
            };

            match analyze_operand_xrefs(layout, pc, rva, insn, op)? {
                // TODO: fallthrough is not guaranteed if the function is noret
                Some(dst) => Ok(vec![Xref { src: rva, dst: dst, typ: XrefType::Call }, fallthrough]),
                None => Ok(vec![fallthrough]),
            }
        }
        zydis::enums::mnemonic::Mnemonic::RET
        | zydis::enums::mnemonic::Mnemonic::IRET
        | zydis::enums::mnemonic::Mnemonic::IRETD
        | zydis::enums::mnemonic::Mnemonic::IRETQ => Ok(vec![]),
        zydis::enums::mnemonic::Mnemonic::JMP => {
            // a JMP always has an operand, so assume this is ok.
            let op = Workspace::get_first_operand(insn).unwrap();

            match analyze_operand_xrefs(layout, pc, rva, insn, op)? {
                Some(dst) => Ok(vec![Xref { src: rva, dst: dst, typ: XrefType::UnconditionalJump }]),
                None => Ok(vec![]),
            }
        }
        zydis::enums::mnemonic::Mnemonic::JB
        | zydis::enums::mnemonic::Mnemonic::JBE
        | zydis::enums::mnemonic::Mnemonic::JCXZ
        | zydis::enums::mnemonic::Mnemonic::JECXZ
        | zydis::enums::mnemonic::Mnemonic::JKNZD
        | zydis::enums::mnemonic::Mnemonic::JKZD
        | zydis::enums::mnemonic::Mnemonic::JL
        | zydis::enums::mnemonic::Mnemonic::JLE
        | zydis::enums::mnemonic::Mnemonic::JNB
        | zydis::enums::mnemonic::Mnemonic::JNBE
        | zydis::enums::mnemonic::Mnemonic::JNL
        | zydis::enums::mnemonic::Mnemonic::JNLE
        | zydis::enums::mnemonic::Mnemonic::JNO
        | zydis::enums::mnemonic::Mnemonic::JNP
        | zydis::enums::mnemonic::Mnemonic::JNS
        | zydis::enums::mnemonic::Mnemonic::JNZ
        | zydis::enums::mnemonic::Mnemonic::JO
        | zydis::enums::mnemonic::Mnemonic::JP
        | zydis::enums::mnemonic::Mnemonic::JRCXZ
        | zydis::enums::mnemonic::Mnemonic::JS
        | zydis::enums::mnemonic::Mnemonic::JZ => {
            // a J* always has an operand, so assume this is ok.
            let op = Workspace::get_first_operand(insn).unwrap();

            let fallthrough = Xref {
                src: rva,
                dst: rva + u64::from(insn.length),
                typ: XrefType::Fallthrough,
            };

            match analyze_operand_xrefs(layout, pc, rva, insn, op)? {
                Some(dst) => Ok(vec![
                    Xref { src: rva, dst: dst, typ: XrefType::ConditionalJump, },
                    fallthrough,
                ]),
                None => Ok(vec![fallthrough]),
            }
        }
        zydis::enums::mnemonic::Mnemonic::CMOVB
        | zydis::enums::mnemonic::Mnemonic::CMOVBE
        | zydis::enums::mnemonic::Mnemonic::CMOVL
        | zydis::enums::mnemonic::Mnemonic::CMOVLE
        | zydis::enums::mnemonic::Mnemonic::CMOVNB
        | zydis::enums::mnemonic::Mnemonic::CMOVNBE
        | zydis::enums::mnemonic::Mnemonic::CMOVNL
        | zydis::enums::mnemonic::Mnemonic::CMOVNLE
        | zydis::enums::mnemonic::Mnemonic::CMOVNO
        | zydis::enums::mnemonic::Mnemonic::CMOVNP
        | zydis::enums::mnemonic::Mnemonic::CMOVNS
        | zydis::enums::mnemonic::Mnemonic::CMOVNZ
        | zydis::enums::mnemonic::Mnemonic::CMOVO
        | zydis::enums::mnemonic::Mnemonic::CMOVP
        | zydis::enums::mnemonic::Mnemonic::CMOVS
        | zydis::enums::mnemonic::Mnemonic::CMOVZ => Ok(vec![
            Xref {
                src: rva,
                dst: rva + u64::from(insn.length),
                typ: XrefType::Fallthrough
            },
            Xref {
                src: rva,
                dst: rva + u64::from(insn.length),
                typ: XrefType::ConditionalMove
            },
        ]),
        _ => Ok(vec![Xref {
            src: rva,
            dst: rva + u64::from(insn.length),
            typ: XrefType::Fallthrough,
        }]),
    }
}

pub fn find_insns(ws: &Workspace, predicate: fn(&Instruction) -> bool) -> Result<Vec<Rva>, Error> {
    let mut ret: Vec<Rva> = vec![];

        // empirically, doing this in parallel (par_extend/par_bridge) is not worth it.
        // the predicates are likely too simple,
        // so the majority of time is spent synchronizing.
        ret.extend(
            ws
            .iter_insns()
            .filter(|insn| predicate(insn))
            .map(|insn| match insn {
                Instruction::Valid {loc, ..} => loc.addr,
                Instruction::Invalid {loc, ..} => loc.addr,
            })
        );

    Ok(ret)
}

pub fn find_roots(ws: &Workspace) -> Result<Vec<Rva>, Error> {
    find_insns(ws, |insn| match insn {
        Instruction::Valid{loc, ..} => loc.xrefs.to.is_empty(),
        _ => false,
    })
}

pub fn find_call_targets(ws: &Workspace) -> Result<Vec<Rva>, Error> {
    find_insns(ws, |insn| match insn {
        Instruction::Valid{loc, ..} => loc.xrefs.to.iter().any(|xref| match xref.typ {
            XrefType::Call => true,
            _ => false,
        }),
        _ => false,
    })
}

pub fn find_branch_targets(ws: &Workspace) -> Result<Vec<Rva>, Error> {
    find_insns(ws, |insn| match insn {
        Instruction::Valid{loc, ..} => loc.xrefs.to.iter().any(|xref| match xref.typ {
            XrefType::UnconditionalJump => true,
            XrefType::ConditionalJump => true,
            _ => false,
        }),
        _ => false,
    })
}

fn get_section_by_name<'a>(ws: &'a Workspace, name: &[u8]) -> Result<Option<&'a Section>, Error> {
    match ws.get_obj()? {
        Object::PE(pe) => {
            match pe.sections.iter()
                .enumerate()
                .map(|(index, section)| (index, section.name))
                .find(|(_, sec_name)| sec_name == name) {
                    Some((index, _)) => Ok(Some(&ws.sections[index])),
                    None => Ok(None),
            }
        },
        _ => Err(Error::NotImplemented("analyzer for non-PE module"))
    }
}

// pub fn find_runtime_functions
//  on pe64, entries in .pdata are struct RUNTIME_FUNCTION
//
//  00000000 FunctionStart   dd ?                    ; offset rva
//  00000004 FunctionEnd     dd ?                    ; offset rva pastend
//  00000008 UnwindInfo      dd ?                    ; offset rva
//
// ref: https://stackoverflow.com/a/9794688/87207
//
// for k32, the first three entries should be:
//  - 1010
//  - 1060
//  - 1140
pub fn find_runtime_functions(ws: &Workspace) -> Result<Vec<Rva>, Error> {
    match ws.get_obj()? {
        Object::PE(_) => {
            let mut ret: Vec<Rva> = vec![];

            if let Some(sec) = get_section_by_name(ws, b".pdata\x00\x00")? {
                let layout = ws.get_section_layout();
                for runtime_function in sec.buf.chunks_exact(mem::size_of::<u32>() * 3) {
                    let rva = LittleEndian::read_u32(&runtime_function) as Rva;
                    if rva == 0 {
                        // must be end of (valid) .pdata section
                        break;
                    } else if layout.is_rva_valid(rva) {
                        ret.push(rva);
                    } else {
                        warn!("malformed .pdata section: entry points outside of image");
                        return Ok(vec![]);
                    }
                }
            }

            Ok(ret)
        },
        _ => Err(Error::NotImplemented("analyzer for non-PE module"))
    }
}

// pub fn find_fixups

// pub fn find_ptrs
//  relies on the image being loaded at the base address
//  also relies on loc vs insn to be figured out.

// for k32, there should be 1630
pub fn find_entrypoints(ws: &Workspace) -> Result<Vec<Rva>, Error> {
    match ws.get_obj()? {
        Object::PE(pe) => {
            let mut ret: Vec<Rva> = vec![];

            let optional_header = pe.header.optional_header.expect("optional header is not optional");
            ret.push(optional_header.standard_fields.address_of_entry_point);

            ret.extend(
                pe.exports.iter().map(|export| export.rva as Rva)
            );

            Ok(ret)
        },
        _ => Err(Error::NotImplemented("analyzer for non-PE module"))
    }
}

// pub fn find_sigs

// pub fn find_impossible_paths
// like:
//   MOV
//   MOV
//   INVALID
<commit_msg>analysis: better handle missing optional header<commit_after>// TODO: data xrefs

use super::*;
use zydis;
use byteorder::{ByteOrder, LittleEndian};
use std::mem;

fn analyze_operand_xrefs(
    // used to determine valid addresses
    layout: &ModuleLayout,
    // used to detect RIP-relative addressing
    pc: zydis::enums::register::Register,
    // used to compute relative addresses
    rva: Rva,
    // used to compute relative addresses
    insn: &zydis::ffi::DecodedInstruction,
    // the operand to inspect
    op: &zydis::ffi::DecodedOperand,
) -> Result<Option<Rva>, Error> {
    match op.ty {
        zydis::enums::OperandType::Unused => Err(Error::NotImplemented("xref from unused register")),
        zydis::enums::OperandType::Register => {
            // like: CALL rbx
            // TODO: for now, don't index unresolved indirect branches
            Ok(None)
        }
        zydis::enums::OperandType::Memory => {
            // like: .text:0000000180001041 FF 15 D1 78 07 00      call    cs:__imp_RtlVirtualUnwind_0
            //           0x0000000000001041:                       call    [0x0000000000079980]
            if pc == op.mem.base && op.mem.disp.has_displacement && op.mem.scale == 0 {
                // RIP-relative
                // this is the default encoding on x64.
                // tools like IDA automatically compute and display the target.
                // CALL [RIP + 0x401000]
                if let zydis::enums::register::Register::NONE = op.mem.index {
                    let target =
                        (rva as i64
                        // TODO: cast from rva (u64) to i64 is lossy.
                        + op.mem.disp.displacement
                        + i64::from(insn.length)) as Rva;

                    if layout.is_rva_valid(target) {
                        Ok(Some(target))
                    } else {
                        // TODO: record this anomaly somewhere.
                        warn!("problem: invalid xref target: memory not in sections");
                        Ok(None)
                    }
                } else {
                    // unsupported
                    // like: CALL [RIP + 4*RCX + 0x401000] ??
                    println!("CALL [RIP + 4*RCX + 0x401000] ??");
                    Err(Error::NotImplemented("xref from RIP-relative, non-zero index memory"))
                }
            } else if op.mem.base == zydis::enums::register::Register::NONE {
                // like: CALL [0x401000] ??
                println!("TODO: other OperandType::Memory branch");
                Err(Error::NotImplemented("xref from non-RIP-relative memory"))
            } else {
                // like: CALL [rbx]
                // like: CALL [rbx + 0x10]
                // TODO: for now, don't index unresolved indirect branches
                Ok(None)
            }
        }
        zydis::enums::OperandType::Pointer => {
            println!("TODO: operand: pointer");
            Err(Error::NotImplemented("xref from pointer"))
        }
        zydis::enums::OperandType::Immediate => {
            if !op.imm.is_relative {
                println!("TODO: absolute immediate operand");
                Err(Error::NotImplemented("xref from absolute immediate"))
            } else {
                let imm = if op.imm.is_signed {
                    u64_i64(op.imm.value)
                } else {
                    op.imm.value as i64
                };

                // TODO: cast from rva (u64) to i64 is lossy.
                let dst = (rva as i64 + imm + i64::from(insn.length)) as Rva;
                if layout.is_rva_valid(dst) {
                    Ok(Some(dst))
                } else {
                    // TODO: record this anomaly somewhere.
                    debug!("problem: invalid xref target: relative immediate not in sections");
                    Ok(None)
                }
            }
        }
    }
}

pub fn analyze_insn_xrefs(
    layout: &ModuleLayout,
    pc: zydis::enums::register::Register,
    rva: Rva,
    insn: &zydis::ffi::DecodedInstruction,
) -> Result<Vec<Xref>, Error> {
    match insn.mnemonic {
        // see InstructionCategory
        // syscall, sysexit, sysret
        // vmcall, vmmcall
        zydis::enums::mnemonic::Mnemonic::CALL => {
            // a CALL always has an operand, so assume this is ok.
            let op = Workspace::get_first_operand(insn).unwrap();

            let fallthrough = Xref {
                src: rva,
                dst: rva + u64::from(insn.length),
                typ: XrefType::Fallthrough,
            };

            match analyze_operand_xrefs(layout, pc, rva, insn, op)? {
                // TODO: fallthrough is not guaranteed if the function is noret
                Some(dst) => Ok(vec![Xref { src: rva, dst: dst, typ: XrefType::Call }, fallthrough]),
                None => Ok(vec![fallthrough]),
            }
        }
        zydis::enums::mnemonic::Mnemonic::RET
        | zydis::enums::mnemonic::Mnemonic::IRET
        | zydis::enums::mnemonic::Mnemonic::IRETD
        | zydis::enums::mnemonic::Mnemonic::IRETQ => Ok(vec![]),
        zydis::enums::mnemonic::Mnemonic::JMP => {
            // a JMP always has an operand, so assume this is ok.
            let op = Workspace::get_first_operand(insn).unwrap();

            match analyze_operand_xrefs(layout, pc, rva, insn, op)? {
                Some(dst) => Ok(vec![Xref { src: rva, dst: dst, typ: XrefType::UnconditionalJump }]),
                None => Ok(vec![]),
            }
        }
        zydis::enums::mnemonic::Mnemonic::JB
        | zydis::enums::mnemonic::Mnemonic::JBE
        | zydis::enums::mnemonic::Mnemonic::JCXZ
        | zydis::enums::mnemonic::Mnemonic::JECXZ
        | zydis::enums::mnemonic::Mnemonic::JKNZD
        | zydis::enums::mnemonic::Mnemonic::JKZD
        | zydis::enums::mnemonic::Mnemonic::JL
        | zydis::enums::mnemonic::Mnemonic::JLE
        | zydis::enums::mnemonic::Mnemonic::JNB
        | zydis::enums::mnemonic::Mnemonic::JNBE
        | zydis::enums::mnemonic::Mnemonic::JNL
        | zydis::enums::mnemonic::Mnemonic::JNLE
        | zydis::enums::mnemonic::Mnemonic::JNO
        | zydis::enums::mnemonic::Mnemonic::JNP
        | zydis::enums::mnemonic::Mnemonic::JNS
        | zydis::enums::mnemonic::Mnemonic::JNZ
        | zydis::enums::mnemonic::Mnemonic::JO
        | zydis::enums::mnemonic::Mnemonic::JP
        | zydis::enums::mnemonic::Mnemonic::JRCXZ
        | zydis::enums::mnemonic::Mnemonic::JS
        | zydis::enums::mnemonic::Mnemonic::JZ => {
            // a J* always has an operand, so assume this is ok.
            let op = Workspace::get_first_operand(insn).unwrap();

            let fallthrough = Xref {
                src: rva,
                dst: rva + u64::from(insn.length),
                typ: XrefType::Fallthrough,
            };

            match analyze_operand_xrefs(layout, pc, rva, insn, op)? {
                Some(dst) => Ok(vec![
                    Xref { src: rva, dst: dst, typ: XrefType::ConditionalJump, },
                    fallthrough,
                ]),
                None => Ok(vec![fallthrough]),
            }
        }
        zydis::enums::mnemonic::Mnemonic::CMOVB
        | zydis::enums::mnemonic::Mnemonic::CMOVBE
        | zydis::enums::mnemonic::Mnemonic::CMOVL
        | zydis::enums::mnemonic::Mnemonic::CMOVLE
        | zydis::enums::mnemonic::Mnemonic::CMOVNB
        | zydis::enums::mnemonic::Mnemonic::CMOVNBE
        | zydis::enums::mnemonic::Mnemonic::CMOVNL
        | zydis::enums::mnemonic::Mnemonic::CMOVNLE
        | zydis::enums::mnemonic::Mnemonic::CMOVNO
        | zydis::enums::mnemonic::Mnemonic::CMOVNP
        | zydis::enums::mnemonic::Mnemonic::CMOVNS
        | zydis::enums::mnemonic::Mnemonic::CMOVNZ
        | zydis::enums::mnemonic::Mnemonic::CMOVO
        | zydis::enums::mnemonic::Mnemonic::CMOVP
        | zydis::enums::mnemonic::Mnemonic::CMOVS
        | zydis::enums::mnemonic::Mnemonic::CMOVZ => Ok(vec![
            Xref {
                src: rva,
                dst: rva + u64::from(insn.length),
                typ: XrefType::Fallthrough
            },
            Xref {
                src: rva,
                dst: rva + u64::from(insn.length),
                typ: XrefType::ConditionalMove
            },
        ]),
        _ => Ok(vec![Xref {
            src: rva,
            dst: rva + u64::from(insn.length),
            typ: XrefType::Fallthrough,
        }]),
    }
}

pub fn find_insns(ws: &Workspace, predicate: fn(&Instruction) -> bool) -> Result<Vec<Rva>, Error> {
    let mut ret: Vec<Rva> = vec![];

        // empirically, doing this in parallel (par_extend/par_bridge) is not worth it.
        // the predicates are likely too simple,
        // so the majority of time is spent synchronizing.
        ret.extend(
            ws
            .iter_insns()
            .filter(|insn| predicate(insn))
            .map(|insn| match insn {
                Instruction::Valid {loc, ..} => loc.addr,
                Instruction::Invalid {loc, ..} => loc.addr,
            })
        );

    Ok(ret)
}

pub fn find_roots(ws: &Workspace) -> Result<Vec<Rva>, Error> {
    find_insns(ws, |insn| match insn {
        Instruction::Valid{loc, ..} => loc.xrefs.to.is_empty(),
        _ => false,
    })
}

pub fn find_call_targets(ws: &Workspace) -> Result<Vec<Rva>, Error> {
    find_insns(ws, |insn| match insn {
        Instruction::Valid{loc, ..} => loc.xrefs.to.iter().any(|xref| match xref.typ {
            XrefType::Call => true,
            _ => false,
        }),
        _ => false,
    })
}

pub fn find_branch_targets(ws: &Workspace) -> Result<Vec<Rva>, Error> {
    find_insns(ws, |insn| match insn {
        Instruction::Valid{loc, ..} => loc.xrefs.to.iter().any(|xref| match xref.typ {
            XrefType::UnconditionalJump => true,
            XrefType::ConditionalJump => true,
            _ => false,
        }),
        _ => false,
    })
}

fn get_section_by_name<'a>(ws: &'a Workspace, name: &[u8]) -> Result<Option<&'a Section>, Error> {
    match ws.get_obj()? {
        Object::PE(pe) => {
            match pe.sections.iter()
                .enumerate()
                .map(|(index, section)| (index, section.name))
                .find(|(_, sec_name)| sec_name == name) {
                    Some((index, _)) => Ok(Some(&ws.sections[index])),
                    None => Ok(None),
            }
        },
        _ => Err(Error::NotImplemented("analyzer for non-PE module"))
    }
}

// pub fn find_runtime_functions
//  on pe64, entries in .pdata are struct RUNTIME_FUNCTION
//
//  00000000 FunctionStart   dd ?                    ; offset rva
//  00000004 FunctionEnd     dd ?                    ; offset rva pastend
//  00000008 UnwindInfo      dd ?                    ; offset rva
//
// ref: https://stackoverflow.com/a/9794688/87207
//
// for k32, the first three entries should be:
//  - 1010
//  - 1060
//  - 1140
pub fn find_runtime_functions(ws: &Workspace) -> Result<Vec<Rva>, Error> {
    match ws.get_obj()? {
        Object::PE(_) => {
            let mut ret: Vec<Rva> = vec![];

            if let Some(sec) = get_section_by_name(ws, b".pdata\x00\x00")? {
                let layout = ws.get_section_layout();
                for runtime_function in sec.buf.chunks_exact(mem::size_of::<u32>() * 3) {
                    let rva = LittleEndian::read_u32(&runtime_function) as Rva;
                    if rva == 0 {
                        // must be end of (valid) .pdata section
                        break;
                    } else if layout.is_rva_valid(rva) {
                        ret.push(rva);
                    } else {
                        warn!("malformed .pdata section: entry points outside of image");
                        return Ok(vec![]);
                    }
                }
            }

            Ok(ret)
        },
        _ => Err(Error::NotImplemented("analyzer for non-PE module"))
    }
}

// pub fn find_fixups

// pub fn find_ptrs
//  relies on the image being loaded at the base address
//  also relies on loc vs insn to be figured out.

/// ```
/// use lancelot::*;
/// use lancelot::rsrc::*;
/// use lancelot::analysis::*;
/// let ws = get_workspace(Rsrc::TINY);
/// assert_eq!(0, find_entrypoints(&ws).unwrap().len());
///
/// //let ws = get_workspace(Rsrc::K32);
/// //assert_eq!(1630, find_entrypoints(&ws).unwrap().len());
/// ```
pub fn find_entrypoints(ws: &Workspace) -> Result<Vec<Rva>, Error> {
    match ws.get_obj()? {
        Object::PE(pe) => {
            let mut ret: Vec<Rva> = vec![];
            if let Some(opt) = pe.header.optional_header {
                ret.push(opt.standard_fields.address_of_entry_point);
            ret.extend(
                pe.exports.iter().map(|export| export.rva as Rva)
            );
            } else {
                info!("no optional header")
            }

            Ok(ret)
        },
        _ => Err(Error::NotImplemented("analyzer for non-PE module"))
    }
}

// pub fn find_sigs

// pub fn find_impossible_paths
// like:
//   MOV
//   MOV
//   INVALID
<|endoftext|>


<commit_before>use byteorder::{BigEndian, ReadBytesExt, WriteBytesExt};

use rustc_serialize::Encodable;
use rustc_serialize::Encoder;
use rustc_serialize::json;

use openssl;
use openssl::ssl;
use openssl::ssl::SslStream;
use openssl::ssl::error::SslError;

use std::ops::{Range, Index};
use std::net::TcpStream;
use std::io::{Cursor, Read, Write};
use std::path::Path;
use std::vec::Vec;
use std::collections::HashMap;

use num::pow;
use rand::{self, Rng};
use time;

#[derive(Debug)]
pub struct Payload<'a> {
    pub aps: PayloadAPS<'a>,
    pub info: Option<HashMap<&'a str, &'a str>>
}

#[derive(Debug)]
pub struct PayloadAPS<'a> {
    pub alert: PayloadAPSAlert<'a>,
    pub badge: Option<i32>,
    pub sound: Option<&'a str>,
    pub content_available: Option<i32>
}

#[derive(Debug)]
pub enum PayloadAPSAlert<'a> {
    Plain(&'a str),
    Localized(&'a str, Vec<&'a str>)
}

impl<'a> Encodable for Payload<'a> {
    fn encode<S: Encoder>(&self, encoder: &mut S) -> Result<(), S::Error> {
	match *self {
	    Payload{ref aps, ref info} => {
		if let Some(ref map) = *info {
	    	    encoder.emit_struct("Payload", 1 + map.len(), |encoder| {
			try!(encoder.emit_struct_field( "aps", 0usize, |encoder| aps.encode(encoder)));
			let mut index = 1usize;
			for (key, val) in map.iter() {
			    try!(encoder.emit_struct_field(key, index, |encoder| val.encode(encoder)));
			    index = index + 1;
			}
			Ok(())
	    	    })
		}
		else {
	    	    encoder.emit_struct("Payload", 1, |encoder| {
			try!(encoder.emit_struct_field( "aps", 0usize, |encoder| aps.encode(encoder)));
			Ok(())
	    	    })
		}
	    }
	}
    }
}

impl<'a> Encodable for PayloadAPS<'a> {
    fn encode<S: Encoder>(&self, encoder: &mut S) -> Result<(), S::Error> {
        match *self {
	    PayloadAPS{ref alert, ref badge, ref sound, ref content_available} => {
	        let mut count = 1;
	        if badge.is_some() { count = count + 1; }
	        if sound.is_some() { count = count + 1; }
	        if content_available.is_some() { count = count + 1; }

	        let mut index = 0usize;
    	        encoder.emit_struct("PayloadAPS", count, |encoder| {
	            try!(encoder.emit_struct_field( "alert", index, |encoder| alert.encode(encoder)));
	            index = index + 1;
	            if badge.is_some() {
		        try!(encoder.emit_struct_field( "badge", index, |encoder| badge.unwrap().encode(encoder)));
		        index = index + 1;
	            }
	            if sound.is_some() {
		        try!(encoder.emit_struct_field( "sound", index, |encoder| sound.unwrap().encode(encoder)));
		        index = index + 1;
	            }
	            if content_available.is_some() {
		        try!(encoder.emit_struct_field( "content-available", index, |encoder| content_available.unwrap().encode(encoder)));
		        index = index + 1;
	            }
	            Ok(())
    	        })
            }
        }
    }
}

impl<'a> Encodable for PayloadAPSAlert<'a> {
    fn encode<S: Encoder>(&self, encoder: &mut S) -> Result<(), S::Error> {
	match *self {
	    PayloadAPSAlert::Plain(ref str) => {
    	        encoder.emit_str(str)
	    },
	    PayloadAPSAlert::Localized(ref key, ref args) => {
    	        encoder.emit_struct("PayloadAPSAlert", 2, |encoder| {
		    try!(encoder.emit_struct_field( "loc-key", 0usize, |encoder| key.encode(encoder)));
		    try!(encoder.emit_struct_field( "loc-args", 1usize, |encoder| args.encode(encoder)));
		    Ok(())
		})
	    }
	}
    }
}

#[allow(dead_code)]
fn hex_to_int(hex: &str) -> u32 {
    let mut total = 0u32;
    let mut n = hex.to_string().len();

    for c in hex.chars() {
        n = n - 1;
	match c {
	    '0'...'9' => {
		total += pow(16, n) * ((c as u32) - ('0' as u32));
	    },
	    'a'...'f' => {
		total += pow(16, n) * ((c as u32) - ('a' as u32) + 10);
	    },
	    _ => {
	    }
	}
    }

    return total;
}

#[allow(dead_code)]
fn convert_to_binary(device_token: &str) -> Vec<u8> {
    let mut device_token_bytes: Vec<u8> = Vec::new();
    for i in 0..8 {
	let string = device_token.to_string();
	let range = Range{start:i*8, end:i*8+8};
	let sub_str = string.index(range);

	let sub_str_num = hex_to_int(sub_str);
	let mut sub_str_bytes = vec![];
	let _ = sub_str_bytes.write_u32::<BigEndian>(sub_str_num);

        for s in sub_str_bytes.iter() {
            device_token_bytes.push(*s);
        }
    }

    return device_token_bytes;
}

#[allow(dead_code)]
pub fn convert_to_token(binary: &[u8]) -> String {
    let mut token = "".to_string();
    for i in 0..8 {
	let range = Range{start:i*4, end:i*4+4};
	let sub_slice = binary.index(range);

	let mut rdr = Cursor::new(sub_slice.to_vec());
	let num = rdr.read_u32::<BigEndian>().unwrap();

	token = format!("{}{:x}", token, num);
    }
    return token;
}

#[allow(dead_code)]
pub fn convert_to_timestamp(binary: &[u8]) -> u32 {
    let mut rdr = Cursor::new(binary.to_vec());
    let num = rdr.read_u32::<BigEndian>().unwrap();

    return num;
}

pub struct APNS<'a> {
    pub sandbox: bool,
    pub certificate: &'a Path,
    pub private_key: &'a Path,
    pub ca_certificate: &'a Path,
}

impl<'a> APNS<'a> {
    pub fn new(sandbox: bool, cert_file: &'a Path, private_key_file: &'a Path, ca_file: &'a Path) -> APNS<'a> {
	APNS{sandbox: sandbox, certificate: cert_file, private_key: private_key_file, ca_certificate: ca_file}
    }

    #[allow(dead_code)]
    pub fn get_feedback(&self) -> Result<Vec<(u32, String)>, SslError> {
	let apns_feedback_production = "feedback.push.apple.com:2196";
	let apns_feedback_development = "feedback.sandbox.push.apple.com:2196";

	let apns_feedback_url = if self.sandbox { apns_feedback_development } else { apns_feedback_production };
	let mut stream = try!(get_ssl_stream(apns_feedback_url, self.certificate, self.private_key, self.ca_certificate));

	let mut tokens: Vec<(u32, String)> = Vec::new();
        let mut read_buffer = [0u8; 38];
        loop {
            match stream.read(&mut read_buffer) {
                Ok(size) => {
                    if size != 38 {
                        break;
                    }
                },
                Err(..) => {
                    /*				return Result::Err(SslError::StreamError(error));*/
                    break;
                }
            }

	    let time_range = Range{start:0, end:4};
	    let time_slice = read_buffer.index(time_range);
	    let time = convert_to_timestamp(time_slice);

	    let token_range = Range{start:6, end:38};
	    let token_slice = read_buffer.index(token_range);
            let token = convert_to_token(token_slice);

	    tokens.push((time, token));
        }

        return Result::Ok(tokens);
    }

    #[allow(dead_code)]
    pub fn send_payload(&self, payload: Payload, device_token: &str) {
        let notification_bytes = get_notification_bytes(payload, device_token);

        let apns_url_production = "gateway.push.apple.com:2195";
        let apns_url_development = "gateway.sandbox.push.apple.com:2195";

        let apns_url = if self.sandbox { apns_url_development } else { apns_url_production };

        let ssl_result = get_ssl_stream(apns_url, self.certificate, self.private_key, self.ca_certificate);
        match ssl_result {
            Ok(mut ssls) => {
                if let Err(error) = ssls.write_all(&notification_bytes) {
                    println!("ssl_stream write error {:?}", error);
                }
                else if ssls.pending() == 6 {
                    let mut read_buffer = [0u8; 6];
                    match ssls.read(&mut read_buffer) {
                        Ok(size) => {
                            for c in read_buffer.iter() {
                                print!("{}", c);
                            }
                            println!("ssl_stream read size {:?}", size);
                        }
                        Err(error) => {
                            println!("ssl_stream read error {:?}", error);
                        }
                    }
                }
                // let _ = ssls.flush();
            },
            Err(error) => {
                println!("failed to get_ssl_stream error {:?}", error);
            }
        };
    }
}

fn get_notification_bytes(payload: Payload, device_token: &str) -> Vec<u8> {
    let payload_str = match json::encode(&payload) {
        Ok(json_str) => { json_str.to_string() }
        Err(error) => {
            println!("json encode error {:?}", error);
            return vec![];
        }
    };

    let payload_bytes = payload_str.into_bytes();
    let device_token_bytes: Vec<u8> = convert_to_binary(device_token);

    let mut notification_buffer: Vec<u8> = vec![];
    let mut message_buffer: Vec<u8> = vec![];

    // Device token
    let mut device_token_length = vec![];
    let _ = device_token_length.write_u16::<BigEndian>(device_token_bytes.len() as u16);

    message_buffer.push(1u8);
    for s in device_token_length.iter() {
        message_buffer.push(*s);
    }
    for s in device_token_bytes.iter() {
        message_buffer.push(*s);
    }

    // Payload
    let mut payload_length = vec![];
    let _ = payload_length.write_u16::<BigEndian>(payload_bytes.len() as u16);

    message_buffer.push(2u8);
    for s in payload_length.iter() {
        message_buffer.push(*s);
    }
    for s in payload_bytes.iter() {
        message_buffer.push(*s);
    }

    // Notification identifier
    let payload_id = rand::thread_rng().gen();
    let mut payload_id_be = vec![];
    let _ = payload_id_be.write_u32::<BigEndian>(payload_id);

    let mut payload_id_length = vec![];
    let _ = payload_id_length.write_u16::<BigEndian>(payload_id_be.len() as u16);

    message_buffer.push(3u8);
    for s in payload_id_length.iter() {
        message_buffer.push(*s);
    }
    for s in payload_id_be.iter() {
        message_buffer.push(*s);
    }

    //	Expiration date
    let time = time::now().to_timespec().sec + 86400;	// expired after one day
    let mut exp_date_be = vec![];
    let _ = exp_date_be.write_u32::<BigEndian>(time as u32);

    let mut exp_date_length = vec![];
    let _ = exp_date_length.write_u16::<BigEndian>(exp_date_be.len() as u16);

    message_buffer.push(4u8);
    for s in exp_date_length.iter() {
        message_buffer.push(*s);
    }
    for s in exp_date_be.iter() {
        message_buffer.push(*s);
    }

    // Priority
    let mut priority_length = vec![];
    let _ = priority_length.write_u16::<BigEndian>(1u16);

    message_buffer.push(5u8);
    for s in priority_length.iter() {
        message_buffer.push(*s);
    }
    message_buffer.push(10u8);

    let mut message_buffer_length = vec![];
    let _ = message_buffer_length.write_u32::<BigEndian>(message_buffer.len() as u32);

    let command = 2u8;
    notification_buffer.push(command);
    for s in message_buffer_length.iter() {
        notification_buffer.push(*s);
    }
    for s in message_buffer.iter() {
        notification_buffer.push(*s);
    }

    return notification_buffer;
}

fn get_ssl_stream(url: &str, cert_file: &Path, private_key_file: &Path, ca_file: &Path) -> Result<SslStream<TcpStream>, SslError> {
    let mut context = try!(ssl::SslContext::new(ssl::SslMethod::Sslv23));

    if let Err(error) = context.set_CA_file(ca_file) {
	println!("set_CA_file error {:?}", error);
    }
    if let Err(error) = context.set_certificate_file(cert_file, openssl::x509::X509FileType::PEM) {
	println!("set_certificate_file error {:?}", error);
    }
    if let Err(error) = context.set_private_key_file(private_key_file, openssl::x509::X509FileType::PEM) {
	println!("set_private_key_file error {:?}", error);
    }

    let tcp_conn = match TcpStream::connect(url) {
	Ok(conn) => { conn },
	Err(error) => {
            println!("tcp_stream connect error {:?}", error);
	    return Result::Err(SslError::StreamError(error));
	}
    };

    return SslStream::new(&context, tcp_conn);
}

<commit_msg>more reliable when sending notification<commit_after>use byteorder::{BigEndian, ReadBytesExt, WriteBytesExt};

use rustc_serialize::Encodable;
use rustc_serialize::Encoder;
use rustc_serialize::json;

use openssl;
use openssl::ssl;
use openssl::ssl::SslStream;
use openssl::ssl::error::SslError;

use std::ops::{Range, Index};
use std::net::TcpStream;
use std::io::{Cursor, Read, Write};
use std::path::Path;
use std::vec::Vec;
use std::collections::HashMap;

use num::pow;
use rand::{self, Rng};
use time;

#[derive(Debug)]
pub struct Payload<'a> {
    pub aps: PayloadAPS<'a>,
    pub info: Option<HashMap<&'a str, &'a str>>
}

#[derive(Debug)]
pub struct PayloadAPS<'a> {
    pub alert: PayloadAPSAlert<'a>,
    pub badge: Option<i32>,
    pub sound: Option<&'a str>,
    pub content_available: Option<i32>
}

#[derive(Debug)]
pub enum PayloadAPSAlert<'a> {
    Plain(&'a str),
    Localized(&'a str, Vec<&'a str>)
}

impl<'a> Encodable for Payload<'a> {
    fn encode<S: Encoder>(&self, encoder: &mut S) -> Result<(), S::Error> {
		match *self {
		    Payload{ref aps, ref info} => {
				if let Some(ref map) = *info {
					encoder.emit_struct("Payload", 1 + map.len(), |encoder| {
					try!(encoder.emit_struct_field( "aps", 0usize, |encoder| aps.encode(encoder)));
					let mut index = 1usize;
					for (key, val) in map.iter() {
					    try!(encoder.emit_struct_field(key, index, |encoder| val.encode(encoder)));
					    index = index + 1;
					}
					Ok(())
					})
				}
				else {
					encoder.emit_struct("Payload", 1, |encoder| {
					try!(encoder.emit_struct_field( "aps", 0usize, |encoder| aps.encode(encoder)));
					Ok(())
					})
				}
		    }
		}
    }
}

impl<'a> Encodable for PayloadAPS<'a> {
    fn encode<S: Encoder>(&self, encoder: &mut S) -> Result<(), S::Error> {
        match *self {
			PayloadAPS{ref alert, ref badge, ref sound, ref content_available} => {
				let mut count = 1;
		        if badge.is_some() { count = count + 1; }
		        if sound.is_some() { count = count + 1; }
		        if content_available.is_some() { count = count + 1; }

		        let mut index = 0usize;
				encoder.emit_struct("PayloadAPS", count, |encoder| {
					try!(encoder.emit_struct_field( "alert", index, |encoder| alert.encode(encoder)));
					index = index + 1;
					if badge.is_some() {
						try!(encoder.emit_struct_field( "badge", index, |encoder| badge.unwrap().encode(encoder)));
						index = index + 1;
		            }
		            if sound.is_some() {
						try!(encoder.emit_struct_field( "sound", index, |encoder| sound.unwrap().encode(encoder)));
						index = index + 1;
		            }
		            if content_available.is_some() {
						try!(encoder.emit_struct_field( "content-available", index, |encoder| content_available.unwrap().encode(encoder)));
						index = index + 1;
		            }
					Ok(())
				})
			}
		}
    }
}

impl<'a> Encodable for PayloadAPSAlert<'a> {
    fn encode<S: Encoder>(&self, encoder: &mut S) -> Result<(), S::Error> {
		match *self {
		    PayloadAPSAlert::Plain(ref str) => {
				encoder.emit_str(str)
		    },
		    PayloadAPSAlert::Localized(ref key, ref args) => {
				encoder.emit_struct("PayloadAPSAlert", 2, |encoder| {
					try!(encoder.emit_struct_field( "loc-key", 0usize, |encoder| key.encode(encoder)));
					try!(encoder.emit_struct_field( "loc-args", 1usize, |encoder| args.encode(encoder)));
					Ok(())
					})
			}
		}
	}
}

#[allow(dead_code)]
fn hex_to_int(hex: &str) -> u32 {
    let mut total = 0u32;
    let mut n = hex.to_string().len();

    for c in hex.chars() {
        n = n - 1;
		match c {
		    '0'...'9' => {
				total += pow(16, n) * ((c as u32) - ('0' as u32));
		    },
		    'a'...'f' => {
				total += pow(16, n) * ((c as u32) - ('a' as u32) + 10);
		    },
		    _ => {

		    }
		}
    }

    return total;
}

#[allow(dead_code)]
fn convert_to_binary(device_token: &str) -> Vec<u8> {
    let mut device_token_bytes: Vec<u8> = Vec::new();
    for i in 0..8 {
		let string = device_token.to_string();
		let range = Range{start:i*8, end:i*8+8};
		let sub_str = string.index(range);

		let sub_str_num = hex_to_int(sub_str);
		let mut sub_str_bytes = vec![];
		let _ = sub_str_bytes.write_u32::<BigEndian>(sub_str_num);

        for s in sub_str_bytes.iter() {
            device_token_bytes.push(*s);
        }
    }

    return device_token_bytes;
}

#[allow(dead_code)]
pub fn convert_to_token(binary: &[u8]) -> String {
    let mut token = "".to_string();
    for i in 0..8 {
		let range = Range{start:i*4, end:i*4+4};
		let sub_slice = binary.index(range);

		let mut rdr = Cursor::new(sub_slice.to_vec());
		let num = rdr.read_u32::<BigEndian>().unwrap();

		token = format!("{}{:x}", token, num);
    }
    return token;
}

#[allow(dead_code)]
pub fn convert_to_timestamp(binary: &[u8]) -> u32 {
    let mut rdr = Cursor::new(binary.to_vec());
    let num = rdr.read_u32::<BigEndian>().unwrap();

    return num;
}

pub struct APNS<'a> {
    pub sandbox: bool,
    pub certificate: &'a Path,
    pub private_key: &'a Path,
    pub ca_certificate: &'a Path,
}

impl<'a> APNS<'a> {
    pub fn new(sandbox: bool, cert_file: &'a Path, private_key_file: &'a Path, ca_file: &'a Path) -> APNS<'a> {
		APNS{sandbox: sandbox, certificate: cert_file, private_key: private_key_file, ca_certificate: ca_file}
    }

    #[allow(dead_code)]
    pub fn get_feedback(&self) -> Result<Vec<(u32, String)>, SslError> {
		let apns_feedback_production = "feedback.push.apple.com:2196";
		let apns_feedback_development = "feedback.sandbox.push.apple.com:2196";

		let apns_feedback_url = if self.sandbox { apns_feedback_development } else { apns_feedback_production };
		let mut stream = try!(get_ssl_stream(apns_feedback_url, self.certificate, self.private_key, self.ca_certificate));

		let mut tokens: Vec<(u32, String)> = Vec::new();
        let mut read_buffer = [0u8; 38];
        loop {
            match stream.read(&mut read_buffer) {
                Ok(size) => {
                    if size != 38 {
                        break;
                    }
                },
                Err(..) => {
                    /*				return Result::Err(SslError::StreamError(error));*/
                    break;
                }
            }
		    let time_range = Range{start:0, end:4};
		    let time_slice = read_buffer.index(time_range);
		    let time = convert_to_timestamp(time_slice);

		    let token_range = Range{start:6, end:38};
		    let token_slice = read_buffer.index(token_range);

            let token = convert_to_token(token_slice);
			tokens.push((time, token));
        }

        return Result::Ok(tokens);
    }

    #[allow(dead_code)]
    pub fn send_payload(&self, payload: Payload, device_token: &str) {
        let notification_bytes = get_notification_bytes(payload, device_token);

        let apns_url_production = "gateway.push.apple.com:2195";
        let apns_url_development = "gateway.sandbox.push.apple.com:2195";

        let apns_url = if self.sandbox { apns_url_development } else { apns_url_production };

        let ssl_result = get_ssl_stream(apns_url, self.certificate, self.private_key, self.ca_certificate);
        match ssl_result {
            Ok(mut ssls) => {
                if let Err(error) = ssls.write_all(&notification_bytes) {
                    println!("ssl_stream write error {:?}", error);
                }

				let _ = ssls.flush();

				// Read possible error code response
				if ssls.pending() == 6 {
                    let mut read_buffer = [0u8; 6];
                    match ssls.read(&mut read_buffer) {
                        Ok(size) => {
                            for c in read_buffer.iter() {
                                print!("{}", c);
                            }
                            println!("ssl_stream read size {:?}", size);
                        }
                        Err(error) => {
                            println!("ssl_stream read error {:?}", error);
                        }
                    }
                }
            },
            Err(error) => {
                println!("failed to get_ssl_stream error {:?}", error);
            }
        };
    }
}

fn get_notification_bytes(payload: Payload, device_token: &str) -> Vec<u8> {
    let payload_str = match json::encode(&payload) {
        Ok(json_str) => { json_str.to_string() }
        Err(error) => {
            println!("json encode error {:?}", error);
            return vec![];
        }
    };

    let payload_bytes = payload_str.into_bytes();
    let device_token_bytes: Vec<u8> = convert_to_binary(device_token);

    let mut notification_buffer: Vec<u8> = vec![];
    let mut message_buffer: Vec<u8> = vec![];

    // Device token
    let mut device_token_length = vec![];
    let _ = device_token_length.write_u16::<BigEndian>(device_token_bytes.len() as u16);

    message_buffer.push(1u8);
    for s in device_token_length.iter() {
        message_buffer.push(*s);
    }
    for s in device_token_bytes.iter() {
        message_buffer.push(*s);
    }

    // Payload
    let mut payload_length = vec![];
    let _ = payload_length.write_u16::<BigEndian>(payload_bytes.len() as u16);

    message_buffer.push(2u8);
    for s in payload_length.iter() {
        message_buffer.push(*s);
    }
    for s in payload_bytes.iter() {
        message_buffer.push(*s);
    }

    // Notification identifier
    let payload_id = rand::thread_rng().gen();
    let mut payload_id_be = vec![];
    let _ = payload_id_be.write_u32::<BigEndian>(payload_id);

    let mut payload_id_length = vec![];
    let _ = payload_id_length.write_u16::<BigEndian>(payload_id_be.len() as u16);

    message_buffer.push(3u8);
    for s in payload_id_length.iter() {
        message_buffer.push(*s);
    }
    for s in payload_id_be.iter() {
        message_buffer.push(*s);
    }

    //	Expiration date
    let time = time::now().to_timespec().sec + 86400;	// expired after one day
    let mut exp_date_be = vec![];
    let _ = exp_date_be.write_u32::<BigEndian>(time as u32);

    let mut exp_date_length = vec![];
    let _ = exp_date_length.write_u16::<BigEndian>(exp_date_be.len() as u16);

    message_buffer.push(4u8);
    for s in exp_date_length.iter() {
        message_buffer.push(*s);
    }
    for s in exp_date_be.iter() {
        message_buffer.push(*s);
    }

    // Priority
    let mut priority_length = vec![];
    let _ = priority_length.write_u16::<BigEndian>(1u16);

    message_buffer.push(5u8);
    for s in priority_length.iter() {
        message_buffer.push(*s);
    }
    message_buffer.push(10u8);

    let mut message_buffer_length = vec![];
    let _ = message_buffer_length.write_u32::<BigEndian>(message_buffer.len() as u32);

    let command = 2u8;
    notification_buffer.push(command);
    for s in message_buffer_length.iter() {
        notification_buffer.push(*s);
    }
    for s in message_buffer.iter() {
        notification_buffer.push(*s);
    }

    return notification_buffer;
}

fn get_ssl_stream(url: &str, cert_file: &Path, private_key_file: &Path, ca_file: &Path) -> Result<SslStream<TcpStream>, SslError> {
    let mut context = try!(ssl::SslContext::new(ssl::SslMethod::Sslv23));

    if let Err(error) = context.set_CA_file(ca_file) {
		println!("set_CA_file error {:?}", error);
    }
    if let Err(error) = context.set_certificate_file(cert_file, openssl::x509::X509FileType::PEM) {
		println!("set_certificate_file error {:?}", error);
    }
    if let Err(error) = context.set_private_key_file(private_key_file, openssl::x509::X509FileType::PEM) {
		println!("set_private_key_file error {:?}", error);
    }

    let tcp_conn = match TcpStream::connect(url) {
		Ok(conn) => { conn },
		Err(error) => {
			println!("tcp_stream connect error {:?}", error);
			return Result::Err(SslError::StreamError(error));
		}
	};

    return SslStream::new(&context, tcp_conn);
}
<|endoftext|>