<commit_before>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this file,
 * You can obtain one at http://mozilla.org/MPL/2.0/. */

#[doc = "Rust wrappers around the raw JS apis"];

use std::libc::types::os::arch::c95::{size_t, c_uint};
use std::libc::{c_char, uintptr_t};
use std::cmp;
use std::rc;
use jsapi::*;
use jsval::{JSVal, NullValue};
use default_stacksize;
use default_heapsize;
use JSOPTION_VAROBJFIX;
use JSOPTION_METHODJIT;
use JSOPTION_TYPE_INFERENCE;
use ERR;
use std::ptr;
use std::ptr::null;
use result;
use result_obj;
use std::str::raw::from_c_str;
use std::cast;
use green::task::GreenTask;

// ___________________________________________________________________________
// friendly Rustic API to runtimes

pub type rt = rc::Rc<rt_rsrc>;

pub struct rt_rsrc {
    ptr : *JSRuntime,
}

impl Drop for rt_rsrc {
    fn drop(&mut self) {
        unsafe {
            JS_Finish(self.ptr);
        }
    }
}

pub fn new_runtime(p: *JSRuntime) -> rt {
    return rc::Rc::new(rt_rsrc {
        ptr: p
    })
}

pub trait RtUtils {
    fn cx(&self) -> rc::Rc<Cx>;
}

impl RtUtils for rc::Rc<rt_rsrc> {
    fn cx(&self) -> rc::Rc<Cx> {
        unsafe {
            new_context(JS_NewContext(self.borrow().ptr,
                                      default_stacksize as size_t), self.clone())
        }
    }
}

// FIXME: Is this safe once we have more than one stack segment?
extern fn gc_callback(rt: *JSRuntime, _status: JSGCStatus) {
    use std::rt::local::Local;
    use std::rt::task::Task;
    unsafe {
        let mut task = Local::borrow(None::<Task>);
        let green_task: ~GreenTask = task.get().maybe_take_runtime().unwrap();
        {
            let c = green_task.coroutine.get_ref();
            let start = c.current_stack_segment.start() as uintptr_t;
            let end = c.current_stack_segment.end() as uintptr_t;
            JS_SetNativeStackBounds(rt, cmp::min(start, end), cmp::max(start, end));
        }
        task.get().put_runtime(green_task);
    }
}

pub fn rt() -> rt {
    unsafe {
        let runtime = JS_Init(default_heapsize);
        JS_SetGCCallback(runtime, gc_callback);
        return new_runtime(runtime);
    }
}

// ___________________________________________________________________________
// contexts

pub struct Cx {
    ptr: *JSContext,
    rt: rt,
}

#[unsafe_destructor]
impl Drop for Cx {
    fn drop(&mut self) {
        unsafe {
            JS_DestroyContext(self.ptr);
        }
    }
}

pub fn new_context(ptr: *JSContext, rt: rt) -> rc::Rc<Cx> {
    return rc::Rc::new(Cx {
        ptr: ptr,
        rt: rt,
    })
}

pub trait CxUtils {
    fn rooted_obj(&self, obj: *JSObject) -> jsobj;
    fn new_compartment(&self, globcls: *JSClass) -> Result<rc::Rc<Compartment>,()>;
    fn new_compartment_with_global(&self, global: *JSObject) -> Result<rc::Rc<Compartment>,()>;
}

impl CxUtils for rc::Rc<Cx> {
    fn rooted_obj(&self, obj: *JSObject) -> jsobj {
        let cxptr = self.borrow().ptr;
        let jsobj = rc::Rc::new(jsobj_rsrc {cx: self.clone(), cxptr: cxptr, ptr: obj});
        unsafe {
            JS_AddObjectRoot(cxptr, &jsobj.borrow().ptr);
        }
        jsobj
    }

    fn new_compartment(&self, globcls: *JSClass) -> Result<rc::Rc<Compartment>,()> {
        unsafe {
            let ptr = self.borrow().ptr;
            let globobj = JS_NewGlobalObject(ptr, globcls, null());
            result(JS_InitStandardClasses(ptr, globobj)).and_then(|_ok| {
                Ok(rc::Rc::new(Compartment {
                    cx: self.clone(),
                    global_obj: self.rooted_obj(globobj),
                }))
            })
        }
    }

    fn new_compartment_with_global(&self, global: *JSObject) -> Result<rc::Rc<Compartment>,()> {
        Ok(rc::Rc::new(Compartment {
            cx: self.clone(),
            global_obj: self.rooted_obj(global),
        }))
    }
}

impl Cx {
    pub fn set_default_options_and_version(&self) {
        self.set_options(JSOPTION_VAROBJFIX | JSOPTION_METHODJIT |
                         JSOPTION_TYPE_INFERENCE);
        self.set_version(JSVERSION_LATEST);
    }

    pub fn set_options(&self, v: c_uint) {
        unsafe {
            JS_SetOptions(self.ptr, v);
        }
    }

    pub fn set_version(&self, v: i32) {
        unsafe {
            JS_SetVersion(self.ptr, v);
        }
    }

    pub fn set_logging_error_reporter(&self) {
        unsafe {
            JS_SetErrorReporter(self.ptr, reportError);
        }
    }

    pub fn set_error_reporter(&self, reportfn: extern "C" fn(*JSContext, *c_char, *JSErrorReport)) {
        unsafe {
            JS_SetErrorReporter(self.ptr, reportfn);
        }
    }

    pub fn evaluate_script(&self, glob: jsobj, script: ~str, filename: ~str, line_num: uint)
                    -> Result<(),()> {
        let script_utf16 = script.to_utf16();
        filename.to_c_str().with_ref(|filename_cstr| {
            let rval: JSVal = NullValue();
            debug!("Evaluating script from {:s} with content {}", filename, script);
            unsafe {
                if ERR == JS_EvaluateUCScript(self.ptr, glob.borrow().ptr,
                                              script_utf16.as_ptr(), script_utf16.len() as c_uint,
                                              filename_cstr, line_num as c_uint,
                                              &rval) {
                    debug!("...err!");
                    Err(())
                } else {
                    // we could return the script result but then we'd have
                    // to root it and so forth and, really, who cares?
                    debug!("...ok!");
                    Ok(())
                }
            }
        })
    }

    pub unsafe fn get_cx_private(&self) -> *() {
        cast::transmute(JS_GetContextPrivate(self.ptr))
    }

    pub unsafe fn set_cx_private(&self, data: *()) {
        JS_SetContextPrivate(self.ptr, cast::transmute(data));
    }

    pub unsafe fn get_obj_private(&self, obj: *JSObject) -> *() {
        cast::transmute(JS_GetPrivate(obj))
    }

    pub unsafe fn set_obj_private(&self, obj: *JSObject, data: *()) {
        JS_SetPrivate(obj, cast::transmute(data));
    }
}

pub extern fn reportError(_cx: *JSContext, msg: *c_char, report: *JSErrorReport) {
    unsafe {
        let fnptr = (*report).filename;
        let fname = if fnptr.is_not_null() {from_c_str(fnptr)} else {~"none"};
        let lineno = (*report).lineno;
        let msg = from_c_str(msg);
        error!("Error at {:s}:{}: {:s}\n", fname, lineno, msg);
    }
}

// ___________________________________________________________________________
// compartment

pub struct Compartment {
    cx: rc::Rc<Cx>,
    global_obj: jsobj,
}

impl Compartment {
    pub fn define_functions(&self, specvec: &'static [JSFunctionSpec]) -> Result<(),()> {
        unsafe {
            result(JS_DefineFunctions(self.cx.borrow().ptr,
                                      self.global_obj.borrow().ptr,
                                      specvec.as_ptr()))
        }
    }
    pub fn define_properties(&self, specvec: &'static [JSPropertySpec]) -> Result<(),()> {
        unsafe {
            result(JS_DefineProperties(self.cx.borrow().ptr,
                                       self.global_obj.borrow().ptr,
                                       specvec.as_ptr()))
        }
    }
    pub fn define_property(&self,
                           name: &'static str,
                           value: JSVal,
                           getter: JSPropertyOp, setter: JSStrictPropertyOp,
                           attrs: c_uint)
        -> Result<(),()> {
        unsafe {
            name.to_c_str().with_ref(|name| {
                result(JS_DefineProperty(self.cx.borrow().ptr,
                                         self.global_obj.borrow().ptr,
                                         name,
                                         value,
                                         Some(getter),
                                         Some(setter),
                                         attrs))
            })
        }
    }
    pub fn new_object(&self, classptr: *JSClass, proto: *JSObject, parent: *JSObject)
               -> Result<jsobj, ()> {
        unsafe {
            let obj = self.cx.rooted_obj(JS_NewObject(self.cx.borrow().ptr, classptr, proto, parent));
            result_obj(obj)
        }
    }
}

// ___________________________________________________________________________
// objects

pub type jsobj = rc::Rc<jsobj_rsrc>;

pub struct jsobj_rsrc {
    cx: rc::Rc<Cx>,
    cxptr: *JSContext,
    ptr: *JSObject,
}

#[unsafe_destructor]
impl Drop for jsobj_rsrc {
    fn drop(&mut self) {
        unsafe {
            JS_RemoveObjectRoot(self.cxptr, &self.ptr);
        }
    }
}

impl jsobj_rsrc {
    pub fn new_object(&self, cx: rc::Rc<Cx>, cxptr: *JSContext, ptr: *JSObject) -> jsobj {
        return rc::Rc::new(jsobj_rsrc {
            cx: cx,
            cxptr: cxptr,
            ptr: ptr
        })
    }
}

// ___________________________________________________________________________
// random utilities

pub trait to_jsstr {
    fn to_jsstr(self, cx: rc::Rc<Cx>) -> *JSString;
}

impl to_jsstr for ~str {
    fn to_jsstr(self, cx: rc::Rc<Cx>) -> *JSString {
        unsafe {
            let cbuf = cast::transmute(self.as_ptr());
            JS_NewStringCopyN(cx.borrow().ptr, cbuf, self.len() as size_t)
        }
    }
}

#[cfg(test)]
pub mod test {
    use super::rt;
    use super::{CxUtils, RtUtils};
    use super::super::global;
    use super::super::jsapi::{JS_GC, JS_GetRuntime};

    #[test]
    pub fn dummy() {
        let rt = rt();
        let cx = rt.cx();
        cx.borrow().set_default_options_and_version();
        cx.borrow().set_logging_error_reporter();
        cx.new_compartment(&global::BASIC_GLOBAL).and_then(|comp| {
            unsafe { JS_GC(JS_GetRuntime(cx.borrow().ptr)); }

            comp.borrow().define_functions(global::DEBUG_FNS);

            let s = ~"debug(22);";
            cx.borrow().evaluate_script(comp.borrow().global_obj.clone(), s, ~"test", 1u)
        });
    }

}
<commit_msg>Avoid scanning stack guard page, and remove unnecessary global object rooting.<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this file,
 * You can obtain one at http://mozilla.org/MPL/2.0/. */

#[doc = "Rust wrappers around the raw JS apis"];

use std::libc::types::os::arch::c95::{size_t, c_uint};
use std::libc::{c_char, uintptr_t};
use std::cmp;
use std::rc;
use std::rt::Runtime;
use jsapi::*;
use jsval::{JSVal, NullValue};
use default_stacksize;
use default_heapsize;
use JSOPTION_VAROBJFIX;
use JSOPTION_METHODJIT;
use JSOPTION_TYPE_INFERENCE;
use ERR;
use std::ptr;
use std::ptr::null;
use result;
use result_obj;
use std::str::raw::from_c_str;
use std::cast;
use green::task::GreenTask;

// ___________________________________________________________________________
// friendly Rustic API to runtimes

pub type rt = rc::Rc<rt_rsrc>;

pub struct rt_rsrc {
    ptr : *JSRuntime,
}

impl Drop for rt_rsrc {
    fn drop(&mut self) {
        unsafe {
            JS_Finish(self.ptr);
        }
    }
}

pub fn new_runtime(p: *JSRuntime) -> rt {
    return rc::Rc::new(rt_rsrc {
        ptr: p
    })
}

pub trait RtUtils {
    fn cx(&self) -> rc::Rc<Cx>;
}

impl RtUtils for rc::Rc<rt_rsrc> {
    fn cx(&self) -> rc::Rc<Cx> {
        unsafe {
            new_context(JS_NewContext(self.borrow().ptr,
                                      default_stacksize as size_t), self.clone())
        }
    }
}

// FIXME: Is this safe once we have more than one stack segment?
extern fn gc_callback(rt: *JSRuntime, _status: JSGCStatus) {
    use std::rt::local::Local;
    use std::rt::task::Task;
    unsafe {
        let mut task = Local::borrow(None::<Task>);
        let green_task: ~GreenTask = task.get().maybe_take_runtime().unwrap();
        let (start, end) = green_task.stack_bounds();
        // libgreen adds a guard page to the stack which causes SpiderMonkey to
        // to choke, so we need to skip over that.
        let start = start + ::std::os::page_size();
        JS_SetNativeStackBounds(rt, cmp::min(start, end), cmp::max(start, end));
        task.get().put_runtime(green_task);
    }
}

pub fn rt() -> rt {
    unsafe {
        let runtime = JS_Init(default_heapsize);
        JS_SetGCCallback(runtime, gc_callback);
        return new_runtime(runtime);
    }
}

// ___________________________________________________________________________
// contexts

pub struct Cx {
    ptr: *JSContext,
    rt: rt,
}

#[unsafe_destructor]
impl Drop for Cx {
    fn drop(&mut self) {
        unsafe {
            JS_DestroyContext(self.ptr);
        }
    }
}

pub fn new_context(ptr: *JSContext, rt: rt) -> rc::Rc<Cx> {
    return rc::Rc::new(Cx {
        ptr: ptr,
        rt: rt,
    })
}

pub trait CxUtils {
    fn rooted_obj(&self, obj: *JSObject) -> jsobj;
    fn new_compartment(&self, globcls: *JSClass) -> Result<rc::Rc<Compartment>,()>;
    fn new_compartment_with_global(&self, global: *JSObject) -> Result<rc::Rc<Compartment>,()>;
}

impl CxUtils for rc::Rc<Cx> {
    fn rooted_obj(&self, obj: *JSObject) -> jsobj {
        let cxptr = self.borrow().ptr;
        let jsobj = rc::Rc::new(jsobj_rsrc {cx: self.clone(), cxptr: cxptr, ptr: obj});
        unsafe {
            JS_AddObjectRoot(cxptr, &jsobj.borrow().ptr);
        }
        jsobj
    }

    fn new_compartment(&self, globcls: *JSClass) -> Result<rc::Rc<Compartment>,()> {
        unsafe {
            let ptr = self.borrow().ptr;
            let globobj = JS_NewGlobalObject(ptr, globcls, null());
            result(JS_InitStandardClasses(ptr, globobj)).and_then(|_ok| {
                Ok(rc::Rc::new(Compartment {
                    cx: self.clone(),
                    global_obj: globobj,
                }))
            })
        }
    }

    fn new_compartment_with_global(&self, global: *JSObject) -> Result<rc::Rc<Compartment>,()> {
        Ok(rc::Rc::new(Compartment {
            cx: self.clone(),
            global_obj: global,
        }))
    }
}

impl Cx {
    pub fn set_default_options_and_version(&self) {
        self.set_options(JSOPTION_VAROBJFIX | JSOPTION_METHODJIT |
                         JSOPTION_TYPE_INFERENCE);
        self.set_version(JSVERSION_LATEST);
    }

    pub fn set_options(&self, v: c_uint) {
        unsafe {
            JS_SetOptions(self.ptr, v);
        }
    }

    pub fn set_version(&self, v: i32) {
        unsafe {
            JS_SetVersion(self.ptr, v);
        }
    }

    pub fn set_logging_error_reporter(&self) {
        unsafe {
            JS_SetErrorReporter(self.ptr, reportError);
        }
    }

    pub fn set_error_reporter(&self, reportfn: extern "C" fn(*JSContext, *c_char, *JSErrorReport)) {
        unsafe {
            JS_SetErrorReporter(self.ptr, reportfn);
        }
    }

    pub fn evaluate_script(&self, glob: *JSObject, script: ~str, filename: ~str, line_num: uint)
                    -> Result<(),()> {
        let script_utf16 = script.to_utf16();
        filename.to_c_str().with_ref(|filename_cstr| {
            let rval: JSVal = NullValue();
            debug!("Evaluating script from {:s} with content {}", filename, script);
            unsafe {
                if ERR == JS_EvaluateUCScript(self.ptr, glob,
                                              script_utf16.as_ptr(), script_utf16.len() as c_uint,
                                              filename_cstr, line_num as c_uint,
                                              &rval) {
                    debug!("...err!");
                    Err(())
                } else {
                    // we could return the script result but then we'd have
                    // to root it and so forth and, really, who cares?
                    debug!("...ok!");
                    Ok(())
                }
            }
        })
    }

    pub unsafe fn get_cx_private(&self) -> *() {
        cast::transmute(JS_GetContextPrivate(self.ptr))
    }

    pub unsafe fn set_cx_private(&self, data: *()) {
        JS_SetContextPrivate(self.ptr, cast::transmute(data));
    }

    pub unsafe fn get_obj_private(&self, obj: *JSObject) -> *() {
        cast::transmute(JS_GetPrivate(obj))
    }

    pub unsafe fn set_obj_private(&self, obj: *JSObject, data: *()) {
        JS_SetPrivate(obj, cast::transmute(data));
    }
}

pub extern fn reportError(_cx: *JSContext, msg: *c_char, report: *JSErrorReport) {
    unsafe {
        let fnptr = (*report).filename;
        let fname = if fnptr.is_not_null() {from_c_str(fnptr)} else {~"none"};
        let lineno = (*report).lineno;
        let msg = from_c_str(msg);
        error!("Error at {:s}:{}: {:s}\n", fname, lineno, msg);
    }
}

// ___________________________________________________________________________
// compartment

pub struct Compartment {
    cx: rc::Rc<Cx>,
    global_obj: *JSObject,
}

impl Compartment {
    pub fn define_functions(&self, specvec: &'static [JSFunctionSpec]) -> Result<(),()> {
        unsafe {
            result(JS_DefineFunctions(self.cx.borrow().ptr,
                                      self.global_obj,
                                      specvec.as_ptr()))
        }
    }
    pub fn define_properties(&self, specvec: &'static [JSPropertySpec]) -> Result<(),()> {
        unsafe {
            result(JS_DefineProperties(self.cx.borrow().ptr,
                                       self.global_obj,
                                       specvec.as_ptr()))
        }
    }
    pub fn define_property(&self,
                           name: &'static str,
                           value: JSVal,
                           getter: JSPropertyOp, setter: JSStrictPropertyOp,
                           attrs: c_uint)
        -> Result<(),()> {
        unsafe {
            name.to_c_str().with_ref(|name| {
                result(JS_DefineProperty(self.cx.borrow().ptr,
                                         self.global_obj,
                                         name,
                                         value,
                                         Some(getter),
                                         Some(setter),
                                         attrs))
            })
        }
    }
    pub fn new_object(&self, classptr: *JSClass, proto: *JSObject, parent: *JSObject)
               -> Result<jsobj, ()> {
        unsafe {
            let obj = self.cx.rooted_obj(JS_NewObject(self.cx.borrow().ptr, classptr, proto, parent));
            result_obj(obj)
        }
    }
}

// ___________________________________________________________________________
// objects

pub type jsobj = rc::Rc<jsobj_rsrc>;

pub struct jsobj_rsrc {
    cx: rc::Rc<Cx>,
    cxptr: *JSContext,
    ptr: *JSObject,
}

#[unsafe_destructor]
impl Drop for jsobj_rsrc {
    fn drop(&mut self) {
        unsafe {
            JS_RemoveObjectRoot(self.cxptr, &self.ptr);
        }
    }
}

impl jsobj_rsrc {
    pub fn new_object(&self, cx: rc::Rc<Cx>, cxptr: *JSContext, ptr: *JSObject) -> jsobj {
        return rc::Rc::new(jsobj_rsrc {
            cx: cx,
            cxptr: cxptr,
            ptr: ptr
        })
    }
}

// ___________________________________________________________________________
// random utilities

pub trait to_jsstr {
    fn to_jsstr(self, cx: rc::Rc<Cx>) -> *JSString;
}

impl to_jsstr for ~str {
    fn to_jsstr(self, cx: rc::Rc<Cx>) -> *JSString {
        unsafe {
            let cbuf = cast::transmute(self.as_ptr());
            JS_NewStringCopyN(cx.borrow().ptr, cbuf, self.len() as size_t)
        }
    }
}

#[cfg(test)]
pub mod test {
    use super::rt;
    use super::{CxUtils, RtUtils};
    use super::super::global;
    use super::super::jsapi::{JS_GC, JS_GetRuntime};

    #[test]
    pub fn dummy() {
        let rt = rt();
        let cx = rt.cx();
        cx.borrow().set_default_options_and_version();
        cx.borrow().set_logging_error_reporter();
        cx.new_compartment(&global::BASIC_GLOBAL).and_then(|comp| {
            unsafe { JS_GC(JS_GetRuntime(cx.borrow().ptr)); }

            comp.borrow().define_functions(global::DEBUG_FNS);

            let s = ~"debug(22);";
            cx.borrow().evaluate_script(comp.borrow().global_obj, s, ~"test", 1u)
        });
    }

}
<|endoftext|>


<commit_before>// Helpers for handling cast expressions, used in both
// typeck and codegen.

use crate::ty::{self, Ty};

use rustc_ast as ast;
use rustc_macros::HashStable;

/// Types that are represented as ints.
#[derive(Copy, Clone, Debug, PartialEq, Eq)]
pub enum IntTy {
    U(ast::UintTy),
    I,
    CEnum,
    Bool,
    Char,
}

// Valid types for the result of a non-coercion cast
#[derive(Copy, Clone, Debug, PartialEq, Eq)]
pub enum CastTy<'tcx> {
    /// Various types that are represented as ints and handled mostly
    /// in the same way, merged for easier matching.
    Int(IntTy),
    /// Floating-Point types
    Float,
    /// Function Pointers
    FnPtr,
    /// Raw pointers
    Ptr(ty::TypeAndMut<'tcx>),
}

/// Cast Kind. See RFC 401 (or librustc_typeck/check/cast.rs)
#[derive(Copy, Clone, Debug, TyEncodable, TyDecodable, HashStable)]
pub enum CastKind {
    CoercionCast,
    PtrPtrCast,
    PtrAddrCast,
    AddrPtrCast,
    NumericCast,
    EnumCast,
    PrimIntCast,
    U8CharCast,
    ArrayPtrCast,
    FnPtrPtrCast,
    FnPtrAddrCast,
}

impl<'tcx> CastTy<'tcx> {
    /// Returns `Some` for integral/pointer casts.
    /// casts like unsizing casts will return `None`
    pub fn from_ty(t: Ty<'tcx>) -> Option<CastTy<'tcx>> {
        match *t.kind() {
            ty::Bool => Some(CastTy::Int(IntTy::Bool)),
            ty::Char => Some(CastTy::Int(IntTy::Char)),
            ty::Int(_) => Some(CastTy::Int(IntTy::I)),
            ty::Infer(ty::InferTy::IntVar(_)) => Some(CastTy::Int(IntTy::I)),
            ty::Infer(ty::InferTy::FloatVar(_)) => Some(CastTy::Float),
            ty::Uint(u) => Some(CastTy::Int(IntTy::U(u))),
            ty::Float(_) => Some(CastTy::Float),
            ty::Adt(d, _) if d.is_enum() && d.is_payloadfree() => Some(CastTy::Int(IntTy::CEnum)),
            ty::RawPtr(mt) => Some(CastTy::Ptr(mt)),
            ty::FnPtr(..) => Some(CastTy::FnPtr),
            _ => None,
        }
    }
}
<commit_msg>Rollup merge of #81131 - pierwill:edit-rustc_middle-cast, r=varkor<commit_after>// Helpers for handling cast expressions, used in both
// typeck and codegen.

use crate::ty::{self, Ty};

use rustc_ast as ast;
use rustc_macros::HashStable;

/// Types that are represented as ints.
#[derive(Copy, Clone, Debug, PartialEq, Eq)]
pub enum IntTy {
    U(ast::UintTy),
    I,
    CEnum,
    Bool,
    Char,
}

// Valid types for the result of a non-coercion cast
#[derive(Copy, Clone, Debug, PartialEq, Eq)]
pub enum CastTy<'tcx> {
    /// Various types that are represented as ints and handled mostly
    /// in the same way, merged for easier matching.
    Int(IntTy),
    /// Floating-point types.
    Float,
    /// Function pointers.
    FnPtr,
    /// Raw pointers.
    Ptr(ty::TypeAndMut<'tcx>),
}

/// Cast Kind. See [RFC 401](https://rust-lang.github.io/rfcs/0401-coercions.html)
/// (or librustc_typeck/check/cast.rs).
#[derive(Copy, Clone, Debug, TyEncodable, TyDecodable, HashStable)]
pub enum CastKind {
    CoercionCast,
    PtrPtrCast,
    PtrAddrCast,
    AddrPtrCast,
    NumericCast,
    EnumCast,
    PrimIntCast,
    U8CharCast,
    ArrayPtrCast,
    FnPtrPtrCast,
    FnPtrAddrCast,
}

impl<'tcx> CastTy<'tcx> {
    /// Returns `Some` for integral/pointer casts.
    /// Casts like unsizing casts will return `None`.
    pub fn from_ty(t: Ty<'tcx>) -> Option<CastTy<'tcx>> {
        match *t.kind() {
            ty::Bool => Some(CastTy::Int(IntTy::Bool)),
            ty::Char => Some(CastTy::Int(IntTy::Char)),
            ty::Int(_) => Some(CastTy::Int(IntTy::I)),
            ty::Infer(ty::InferTy::IntVar(_)) => Some(CastTy::Int(IntTy::I)),
            ty::Infer(ty::InferTy::FloatVar(_)) => Some(CastTy::Float),
            ty::Uint(u) => Some(CastTy::Int(IntTy::U(u))),
            ty::Float(_) => Some(CastTy::Float),
            ty::Adt(d, _) if d.is_enum() && d.is_payloadfree() => Some(CastTy::Int(IntTy::CEnum)),
            ty::RawPtr(mt) => Some(CastTy::Ptr(mt)),
            ty::FnPtr(..) => Some(CastTy::FnPtr),
            _ => None,
        }
    }
}
<|endoftext|>


<commit_before>use std;
use std::collections::HashMap;
use std::fs::File;
use std::io::Read;
use std::path::Path;

use rand;
use rand::Rng;

use aes::{pad, padding_valid, Aes128, MODE};
use aes::BLOCK_SIZE;

use serialize::from_base64_file;

use unstable_features::all_bytes;
use unstable_features::MoveFrom;

use xor::XOR;

use helper::ceil_div;

use errors::*;

use prefix_suffix_oracles::{DeterministicOracle, Oracle};
use prefix_suffix_oracles::{Oracle11, Oracle12, Oracle13, Oracle14, Oracle16};

fn matasano2_9() -> Result<()> {
    compare(
        [
            89,
            69,
            76,
            76,
            79,
            87,
            32,
            83,
            85,
            66,
            77,
            65,
            82,
            73,
            78,
            69,
            4,
            4,
            4,
            4,
        ].as_ref(),
        &pad(b"YELLOW SUBMARINE".as_ref(), 20)?,
    )
}

#[test]
fn aes_128_cbc() {
    let iv = [0; BLOCK_SIZE];
    let key = b"YELLOW SUBMARINE";
    let input = b"ABCDEFGHIJKLMNOP";
    assert_eq!(
        input.as_ref(),
        &input
            .encrypt(key, Some(&iv), MODE::CBC)
            .unwrap()
            .decrypt(key, Some(&iv), MODE::CBC)
            .unwrap()[..]
    );
}

fn matasano2_10() -> Result<()> {
    let key = b"YELLOW SUBMARINE";
    let input = from_base64_file(Path::new("data/10.txt"))?;
    let cleartext = input.decrypt(key, Some(&[0; BLOCK_SIZE]), MODE::CBC)?;

    //Read reference cleartext
    let path = Path::new("data/10.ref.txt");
    let mut file = File::open(&path)?;
    let mut cleartext_ref = String::new();
    file.read_to_string(&mut cleartext_ref)?;

    compare(cleartext_ref.as_bytes(), &cleartext)
}

fn uses_ecb(oracle: &mut Oracle11) -> Result<bool> {
    // Assumes that oracle prepends at most one block of jibber
    // TODO: Can we relax this condition?
    let input = vec![0; 3 * BLOCK_SIZE];
    let ciphertext = oracle.encrypt(&input)?;
    let blocks: Vec<&[u8]> = ciphertext.chunks(BLOCK_SIZE).skip(1).take(2).collect();
    Ok(blocks[0] == blocks[1])
}

fn prefix_plus_suffix_length<T: Oracle>(oracle: &T) -> Result<usize> {
    let initial = oracle.encrypt(&[])?.len();
    let input = [0; BLOCK_SIZE];
    //Would profit from range_inclusive
    if let Some(index) = (1..BLOCK_SIZE + 1).find(|&i| {
        if let Ok(ciphertext) = oracle.encrypt(&input[BLOCK_SIZE - i..]) {
            initial != ciphertext.len()
        } else {
            false
        }
    }) {
        Ok(initial - index)
    } else {
        bail!(
            "length of oracle output did not change, something is wrong with the provided oracle"
        );
    }
}

/* For an oracle prepending prefix and appending suffix to its input, this function returns
 * prefix.len()/BLOCK_SIZE, that is the number of blocks fully occupied by the prefix. */
fn prefix_blocks_count<T: DeterministicOracle>(oracle: &T) -> Result<usize> {
    if let Some(result) = oracle
        .encrypt(&[0])?
        .chunks(BLOCK_SIZE)
        .zip(oracle.encrypt(&[1])?.chunks(BLOCK_SIZE))
        .position(|(x, y)| x != y)
    {
        Ok(result)
    } else {
        bail!("no differing blocks found, something is wrong with the provided oracle");
    }
}

pub fn prefix_length<T: DeterministicOracle>(oracle: &T) -> Result<usize> {
    let n = prefix_blocks_count(oracle)?;
    let helper = |k: u8| -> Result<usize> {
        let u = vec![k; BLOCK_SIZE];

        let mut prev_u = oracle.encrypt(&u)?;

        for i in 0..BLOCK_SIZE {
            let cur_u = oracle.encrypt(&u[i + 1..])?;
            if prev_u.chunks(BLOCK_SIZE).nth(n) != cur_u.chunks(BLOCK_SIZE).nth(n) {
                return Ok(i);
            }
            prev_u = cur_u;
        }
        Ok(BLOCK_SIZE)
    };

    Ok(n * BLOCK_SIZE + std::cmp::min(helper(0)?, helper(1)?))
}

fn suffix_length<T: DeterministicOracle>(oracle: &T) -> Result<usize> {
    Ok(prefix_plus_suffix_length(oracle)? - prefix_length(oracle)?)
}

#[test]
fn test_length_functions() {
    let key = random_block();
    let mut prefix = Vec::new();
    let mut suffix = Vec::new();
    for _ in 0..64 {
        for _ in 0..64 {
            {
                let oracle = CommonPrefixSuffixOracle {
                    key: key.clone(),
                    prefix: prefix.clone(),
                    suffix: suffix.clone(),
                    mode: MODE::CTR,
                };
                println!("{}", prefix.len());
                println!("{}", prefix_length(&oracle).unwrap());
                assert!(prefix.len() == prefix_length(&oracle).unwrap());
                //assert!(suffix.len() == suffix_length(&oracle).unwrap());
            }
            suffix.push(1);
        }
        suffix.clear();
        prefix.push(0);
    }
}

fn matasano2_11() -> Result<()> {
    let mut oracle = Oracle11::new()?;
    let uses_ecb = uses_ecb(&mut oracle)?;
    oracle.verify_solution(uses_ecb)
}

fn decrypt_suffix<T: DeterministicOracle>(oracle: &T) -> Result<Vec<u8>> {
    // The following input is chosen in such a way that the cleartext in oracle looks as follows:
    //
    //            input start      input end
    //                ↓                ↓
    // <-- prefix --> 0 ... 0 || 0 ... 0 suffix[0] || suffix[1] ...
    //                ↑          ↑
    //            prefix_len  prefix_blocks*BLOCK_SIZE
    //
    // The resulting ciphertext is compared to oracle([input, u]). The u yielding a match is
    // equal to suffix[0].

    let prefix_len = prefix_length(oracle)?;
    let (prefix_blocks, prefix_padding) = ceil_div(prefix_len, BLOCK_SIZE);
    let suffix_len = suffix_length(oracle)?;

    let mut suffix = Vec::with_capacity(suffix_len);

    let mut input = vec![0; prefix_padding + BLOCK_SIZE - 1];
    let reference_ciphertexts = (0..BLOCK_SIZE)
        .map(|left_shift| oracle.encrypt(&input[left_shift..]))
        .collect::<Result<Vec<Vec<u8>>>>()?;

    for i in 0..suffix_len {
        let block = prefix_blocks + i / BLOCK_SIZE;
        //let block_range = block*BLOCK_SIZE .. (block + 1)*BLOCK_SIZE;
        let left_shift = i % BLOCK_SIZE;
        for u in all_bytes() {
            input.push(u);
            if reference_ciphertexts[left_shift][block * BLOCK_SIZE..(block + 1) * BLOCK_SIZE]
                == oracle.encrypt(&input[left_shift..])?
                    [block * BLOCK_SIZE..(block + 1) * BLOCK_SIZE]
            {
                suffix.push(u);
                break;
            }
            input.pop();
        }
    }
    Ok(suffix)
}

fn matasano2_12() -> Result<()> {
    let oracle = Oracle12::new()?;
    oracle.verify_suffix(&decrypt_suffix(&oracle)?)
}

pub fn decode_profile(u: &[u8], sep: u8) -> HashMap<&[u8], &[u8]> {
    let mut p = HashMap::new();
    for pair in u.split(|&x| x == sep) {
        let mut components = pair.split(|&x| x == b'=');
        p.insert(components.next().unwrap(), components.next().unwrap_or(&[]));
    }
    p
}

/* The following function works under the single assumption that the target value "user" (to be
   replaced by "admin") is stored at the very end of the profile. */
fn matasano2_13() -> Result<()> {
    let oracle = Oracle13::new()?;

    let prefix_len = prefix_length(&oracle)?;
    let (prefix_blocks, prefix_padding) = ceil_div(prefix_len, BLOCK_SIZE);
    let target_cleartext = b"admin".pad();
    let mut input = vec![0; prefix_padding];
    input.extend_from_slice(&target_cleartext);
    let target_last_block = oracle
        .encrypt(&input)?
        .split_off(prefix_blocks * BLOCK_SIZE);

    let (blocks, padding) = ceil_div(prefix_plus_suffix_length(&oracle)?, BLOCK_SIZE);
    let mut ciphertext = oracle.encrypt(&vec![0; padding + "user".len()])?;
    compare((blocks + 1) * BLOCK_SIZE, ciphertext.len())?;

    ciphertext[blocks * BLOCK_SIZE..].move_from2(target_last_block, 0, BLOCK_SIZE);

    oracle.verify_solution(&ciphertext)
}

fn matasano2_14() -> Result<()> {
    let oracle = Oracle14::new()?;
    oracle.verify_suffix(&decrypt_suffix(&oracle)?)
}

pub fn random_block() -> Vec<u8> {
    let mut rng = rand::thread_rng();
    rng.gen_iter().take(BLOCK_SIZE).collect()
}

fn matasano2_15() -> Result<()> {
    compare(true, b"ICE ICE BABY\x04\x04\x04\x04".padding_valid())?;
    compare(false, b"ICE ICE BABY\x05\x05\x05\x05".padding_valid())?;
    compare(false, b"ICE ICE BABY\x01\x02\x03\x04".padding_valid())?;
    compare(false, b"ICE ICE BABY\x03\x03\x03".padding_valid())?;
    compare(
        true,
        padding_valid(
            b"ICE ICE BABY\x0C\x0C\x0C\x0C\x0C\x0C\x0C\x0C\x0C\x0C\x0C\x0C",
            12,
        ).unwrap(),
    )?;
    Ok(())
}

fn matasano2_16() -> Result<()> {
    let oracle = Oracle16::new()?;

    let (blocks, padding) = ceil_div(prefix_plus_suffix_length(&oracle)?, BLOCK_SIZE);
    let mut ciphertext = oracle.encrypt(&vec![0; padding])?;
    compare((blocks + 1) * BLOCK_SIZE, ciphertext.len())?;

    let target_last_block = b";admin=true".pad();
    let current_last_block = vec![BLOCK_SIZE as u8; BLOCK_SIZE];
    let attack_bitflip = target_last_block.xor(&current_last_block);

    // Flip the next to last block
    ciphertext[(blocks - 1) * BLOCK_SIZE..blocks * BLOCK_SIZE].xor_inplace(&attack_bitflip);

    oracle.verify_solution(&ciphertext)
}

pub fn run() {
    println!("Set 2");
    run_exercise(matasano2_9, 9);
    run_exercise(matasano2_10, 10);
    run_exercise(matasano2_11, 11);
    run_exercise(matasano2_12, 12);
    run_exercise(matasano2_13, 13);
    run_exercise(matasano2_14, 14);
    run_exercise(matasano2_15, 15);
    run_exercise(matasano2_16, 16);
}
<commit_msg>Add some comments<commit_after>use std;
use std::collections::HashMap;
use std::fs::File;
use std::io::Read;
use std::path::Path;

use rand;
use rand::Rng;

use aes::{pad, padding_valid, Aes128, MODE};
use aes::BLOCK_SIZE;

use serialize::from_base64_file;

use unstable_features::all_bytes;
use unstable_features::MoveFrom;

use xor::XOR;

use helper::ceil_div;

use errors::*;

use prefix_suffix_oracles::{DeterministicOracle, Oracle};
use prefix_suffix_oracles::{Oracle11, Oracle12, Oracle13, Oracle14, Oracle16};

fn matasano2_9() -> Result<()> {
    compare(
        [
            89,
            69,
            76,
            76,
            79,
            87,
            32,
            83,
            85,
            66,
            77,
            65,
            82,
            73,
            78,
            69,
            4,
            4,
            4,
            4,
        ].as_ref(),
        &pad(b"YELLOW SUBMARINE".as_ref(), 20)?,
    )
}

#[test]
fn aes_128_cbc() {
    let iv = [0; BLOCK_SIZE];
    let key = b"YELLOW SUBMARINE";
    let input = b"ABCDEFGHIJKLMNOP";
    assert_eq!(
        input.as_ref(),
        &input
            .encrypt(key, Some(&iv), MODE::CBC)
            .unwrap()
            .decrypt(key, Some(&iv), MODE::CBC)
            .unwrap()[..]
    );
}

fn matasano2_10() -> Result<()> {
    let key = b"YELLOW SUBMARINE";
    let input = from_base64_file(Path::new("data/10.txt"))?;
    let cleartext = input.decrypt(key, Some(&[0; BLOCK_SIZE]), MODE::CBC)?;

    //Read reference cleartext
    let path = Path::new("data/10.ref.txt");
    let mut file = File::open(&path)?;
    let mut cleartext_ref = String::new();
    file.read_to_string(&mut cleartext_ref)?;

    compare(cleartext_ref.as_bytes(), &cleartext)
}

fn uses_ecb(oracle: &mut Oracle11) -> Result<bool> {
    // Assumes that oracle prepends at most one block of jibber
    // TODO: Can we relax this condition?
    let input = vec![0; 3 * BLOCK_SIZE];
    let ciphertext = oracle.encrypt(&input)?;
    let blocks: Vec<&[u8]> = ciphertext.chunks(BLOCK_SIZE).skip(1).take(2).collect();
    Ok(blocks[0] == blocks[1])
}

fn prefix_plus_suffix_length<T: Oracle>(oracle: &T) -> Result<usize> {
    let initial = oracle.encrypt(&[])?.len();
    let input = [0; BLOCK_SIZE];
    //Would profit from range_inclusive
    if let Some(index) = (1..BLOCK_SIZE + 1).find(|&i| {
        if let Ok(ciphertext) = oracle.encrypt(&input[BLOCK_SIZE - i..]) {
            initial != ciphertext.len()
        } else {
            false
        }
    }) {
        Ok(initial - index)
    } else {
        bail!(
            "length of oracle output did not change, something is wrong with the provided oracle"
        );
    }
}

/* For an oracle prepending prefix and appending suffix to its input, this function returns
 * prefix.len()/BLOCK_SIZE, that is the number of blocks fully occupied by the prefix.
 *
 * To determine this number, we pass two different cleartexts to the oracle and count the number
 * of identical blocks at the start of the corresponding ciphertexts. */
fn prefix_blocks_count<T: DeterministicOracle>(oracle: &T) -> Result<usize> {
    if let Some(result) = oracle
        .encrypt(&[0])?
        .chunks(BLOCK_SIZE)
        .zip(oracle.encrypt(&[1])?.chunks(BLOCK_SIZE))
        .position(|(x, y)| x != y)
    {
        Ok(result)
    } else {
        bail!("no differing blocks found, something is wrong with the provided oracle");
    }
}

// We look at the first block C not fully occupied by the prefix and fill it with a
// constant block B, say consisting of 0's.
// This part of the cleartext in oracle looks as follows:
//
//                               <------ B ------->
// prefix[?] prefix[?] prefix[?] 0 0 ... 0 || 0 0 0 suffix[0] suffix[1]
// <---------- C ------------------------>
//
// We then successively reduce the length of B until the ciphertext of C
// changes. This happens as soon as the cleartext in oracle looks as follows:
//
//                               <-- B -->
// prefix[?] prefix[?] prefix[?] 0 0 ... 0 suffix[0] || suffix[1]
// <---------- C ---------------------------------->
//
// This gives us the length of the prefix in C.
//
// We need to do this with two different constants because suffix[0] might accidentally
// coincide with the constant we have chosen.

pub fn prefix_length<T: DeterministicOracle>(oracle: &T) -> Result<usize> {
    let n = prefix_blocks_count(oracle)?;
    let helper = |k: u8| -> Result<usize> {
        let constant_block = vec![k; BLOCK_SIZE];

        let mut prev = oracle.encrypt(&constant_block)?;

        for i in 0..BLOCK_SIZE {
            let cur = oracle.encrypt(&constant_block[i + 1..])?;
            if prev.chunks(BLOCK_SIZE).nth(n) != cur.chunks(BLOCK_SIZE).nth(n) {
                return Ok(i);
            }
            prev = cur;
        }
        Ok(BLOCK_SIZE)
    };

    Ok(n * BLOCK_SIZE + std::cmp::min(helper(0)?, helper(1)?))
}

fn suffix_length<T: DeterministicOracle>(oracle: &T) -> Result<usize> {
    Ok(prefix_plus_suffix_length(oracle)? - prefix_length(oracle)?)
}

#[test]
fn test_length_functions() {
    let key = random_block();
    let mut prefix = Vec::new();
    let mut suffix = Vec::new();
    for _ in 0..64 {
        for _ in 0..64 {
            {
                let oracle = CommonPrefixSuffixOracle {
                    key: key.clone(),
                    prefix: prefix.clone(),
                    suffix: suffix.clone(),
                    mode: MODE::CTR,
                };
                println!("{}", prefix.len());
                println!("{}", prefix_length(&oracle).unwrap());
                assert!(prefix.len() == prefix_length(&oracle).unwrap());
                //assert!(suffix.len() == suffix_length(&oracle).unwrap());
            }
            suffix.push(1);
        }
        suffix.clear();
        prefix.push(0);
    }
}

fn matasano2_11() -> Result<()> {
    let mut oracle = Oracle11::new()?;
    let uses_ecb = uses_ecb(&mut oracle)?;
    oracle.verify_solution(uses_ecb)
}

fn decrypt_suffix<T: DeterministicOracle>(oracle: &T) -> Result<Vec<u8>> {
    // The following input is chosen in such a way that the cleartext in oracle looks as follows:
    //
    //            input start      input end
    //                ↓                ↓
    // <-- prefix --> 0 ... 0 || 0 ... 0 suffix[0] || suffix[1] ...
    //                ↑          ↑
    //            prefix_len  prefix_blocks*BLOCK_SIZE
    //
    // The resulting ciphertext is compared to oracle([input, u]). The u yielding a match is
    // equal to suffix[0].

    let prefix_len = prefix_length(oracle)?;
    let (prefix_blocks, prefix_padding) = ceil_div(prefix_len, BLOCK_SIZE);
    let suffix_len = suffix_length(oracle)?;

    let mut suffix = Vec::with_capacity(suffix_len);

    let mut input = vec![0; prefix_padding + BLOCK_SIZE - 1];
    let reference_ciphertexts = (0..BLOCK_SIZE)
        .map(|left_shift| oracle.encrypt(&input[left_shift..]))
        .collect::<Result<Vec<Vec<u8>>>>()?;

    for i in 0..suffix_len {
        let block = prefix_blocks + i / BLOCK_SIZE;
        //let block_range = block*BLOCK_SIZE .. (block + 1)*BLOCK_SIZE;
        let left_shift = i % BLOCK_SIZE;
        for u in all_bytes() {
            input.push(u);
            if reference_ciphertexts[left_shift][block * BLOCK_SIZE..(block + 1) * BLOCK_SIZE]
                == oracle.encrypt(&input[left_shift..])?
                    [block * BLOCK_SIZE..(block + 1) * BLOCK_SIZE]
            {
                suffix.push(u);
                break;
            }
            input.pop();
        }
    }
    Ok(suffix)
}

fn matasano2_12() -> Result<()> {
    let oracle = Oracle12::new()?;
    oracle.verify_suffix(&decrypt_suffix(&oracle)?)
}

pub fn decode_profile(u: &[u8], sep: u8) -> HashMap<&[u8], &[u8]> {
    let mut p = HashMap::new();
    for pair in u.split(|&x| x == sep) {
        let mut components = pair.split(|&x| x == b'=');
        p.insert(components.next().unwrap(), components.next().unwrap_or(&[]));
    }
    p
}

/* The following function works under the single assumption that the target value "user" (to be
   replaced by "admin") is stored at the very end of the profile. */
fn matasano2_13() -> Result<()> {
    let oracle = Oracle13::new()?;

    let prefix_len = prefix_length(&oracle)?;
    let (prefix_blocks, prefix_padding) = ceil_div(prefix_len, BLOCK_SIZE);
    let target_cleartext = b"admin".pad();
    let mut input = vec![0; prefix_padding];
    input.extend_from_slice(&target_cleartext);
    let target_last_block = oracle
        .encrypt(&input)?
        .split_off(prefix_blocks * BLOCK_SIZE);

    let (blocks, padding) = ceil_div(prefix_plus_suffix_length(&oracle)?, BLOCK_SIZE);
    let mut ciphertext = oracle.encrypt(&vec![0; padding + "user".len()])?;
    compare((blocks + 1) * BLOCK_SIZE, ciphertext.len())?;

    ciphertext[blocks * BLOCK_SIZE..].move_from2(target_last_block, 0, BLOCK_SIZE);

    oracle.verify_solution(&ciphertext)
}

fn matasano2_14() -> Result<()> {
    let oracle = Oracle14::new()?;
    oracle.verify_suffix(&decrypt_suffix(&oracle)?)
}

pub fn random_block() -> Vec<u8> {
    let mut rng = rand::thread_rng();
    rng.gen_iter().take(BLOCK_SIZE).collect()
}

fn matasano2_15() -> Result<()> {
    compare(true, b"ICE ICE BABY\x04\x04\x04\x04".padding_valid())?;
    compare(false, b"ICE ICE BABY\x05\x05\x05\x05".padding_valid())?;
    compare(false, b"ICE ICE BABY\x01\x02\x03\x04".padding_valid())?;
    compare(false, b"ICE ICE BABY\x03\x03\x03".padding_valid())?;
    compare(
        true,
        padding_valid(
            b"ICE ICE BABY\x0C\x0C\x0C\x0C\x0C\x0C\x0C\x0C\x0C\x0C\x0C\x0C",
            12,
        ).unwrap(),
    )?;
    Ok(())
}

fn matasano2_16() -> Result<()> {
    let oracle = Oracle16::new()?;

    let (blocks, padding) = ceil_div(prefix_plus_suffix_length(&oracle)?, BLOCK_SIZE);
    let mut ciphertext = oracle.encrypt(&vec![0; padding])?;
    compare((blocks + 1) * BLOCK_SIZE, ciphertext.len())?;

    let target_last_block = b";admin=true".pad();
    let current_last_block = vec![BLOCK_SIZE as u8; BLOCK_SIZE];
    let attack_bitflip = target_last_block.xor(&current_last_block);

    // Flip the next to last block
    ciphertext[(blocks - 1) * BLOCK_SIZE..blocks * BLOCK_SIZE].xor_inplace(&attack_bitflip);

    oracle.verify_solution(&ciphertext)
}

pub fn run() {
    println!("Set 2");
    run_exercise(matasano2_9, 9);
    run_exercise(matasano2_10, 10);
    run_exercise(matasano2_11, 11);
    run_exercise(matasano2_12, 12);
    run_exercise(matasano2_13, 13);
    run_exercise(matasano2_14, 14);
    run_exercise(matasano2_15, 15);
    run_exercise(matasano2_16, 16);
}
<|endoftext|>


<commit_before>use std::{env,
          ffi::OsString,
          fs as stdfs,
          path::PathBuf};

use crate::{common::ui::UI,
            hcore::{crypto::CACHE_KEY_PATH_ENV_VAR,
                    env as henv,
                    fs}};

use crate::{config,
            error::Result};

pub const ARTIFACT_PATH_ENVVAR: &str = "ARTIFACT_PATH";

const AUTH_TOKEN_ENVVAR: &str = "HAB_AUTH_TOKEN";
const BLDR_URL_ENVVAR: &str = "HAB_BLDR_URL";
const CTL_SECRET_ENVVAR: &str= "HAB_CTL_SECRET";
const ORIGIN_ENVVAR: &str = "HAB_ORIGIN";
const STUDIO_CMD: &str = "hab-studio";
const STUDIO_CMD_ENVVAR: &str = "HAB_STUDIO_BINARY";
const STUDIO_PACKAGE_IDENT: &str = "core/hab-studio";

fn set_env_var_from_config(env_var: &str, config_val: Option<String>) {
    if henv::var(env_var).is_err() {
        if let Some(val) = config_val {
            debug!("Setting {}={} via config file", env_var, val);
            env::set_var(env_var, val);
        }
    }
}

pub fn start(ui: &mut UI, args: &[OsString]) -> Result<()> {
    let config = config::load()?;

    set_env_var_from_config(AUTH_TOKEN_ENVVAR, config.auth_token);
    set_env_var_from_config(BLDR_URL_ENVVAR, config.bldr_url);
    set_env_var_from_config(CTL_SECRET_ENVVAR, config.ctl_secret);
    set_env_var_from_config(ORIGIN_ENVVAR, config.origin);

    if henv::var(CACHE_KEY_PATH_ENV_VAR).is_err() {
        let path = fs::cache_key_path(None::<&str>);
        debug!("Setting {}={}", CACHE_KEY_PATH_ENV_VAR, path.display());
        env::set_var(CACHE_KEY_PATH_ENV_VAR, &path);
    };

    let artifact_path = match henv::var(ARTIFACT_PATH_ENVVAR) {
        Ok(p) => PathBuf::from(p),
        Err(_) => {
            let path = fs::cache_artifact_path(None::<&str>);
            debug!("Setting {}={}", ARTIFACT_PATH_ENVVAR, path.display());
            env::set_var(ARTIFACT_PATH_ENVVAR, &path);
            path
        }
    };
    if !artifact_path.is_dir() {
        debug!("Creating artifact_path at: {}", artifact_path.display());
        stdfs::create_dir_all(&artifact_path)?;
    }

    inner::start(ui, args)
}

#[cfg(target_os = "linux")]
mod inner {
    use crate::{command::studio::docker,
                common::ui::{UIWriter,
                             UI},
                error::{Error,
                        Result},
                exec,
                hcore::{crypto::init,
                        env as henv,
                        fs::{am_i_root,
                             find_command},
                        os::process,
                        package::{PackageIdent,
                                  PackageInstall},
                        users::linux as group},
                VERSION};
    use std::{env,
              ffi::OsString,
              path::PathBuf,
              str::FromStr};

    const SUDO_CMD: &str = "sudo";

    pub fn start(ui: &mut UI, args: &[OsString]) -> Result<()> {
        rerun_with_sudo_if_needed(ui, &args)?;
        if is_docker_studio(&args) {
            docker::start_docker_studio(ui, args)
        } else {
            let command = match henv::var(super::STUDIO_CMD_ENVVAR) {
                Ok(command) => PathBuf::from(command),
                Err(_) => {
                    init();
                    let version: Vec<&str> = VERSION.split('/').collect();
                    let ident = PackageIdent::from_str(&format!("{}/{}",
                                                                super::STUDIO_PACKAGE_IDENT,
                                                                version[0]))?;
                    // This is a duplicate of the code in `hab pkg exec` and
                    // should be refactored as part of or after:
                    // https://github.com/habitat-sh/habitat/issues/6633
                    // https://github.com/habitat-sh/habitat/issues/6634
                    let pkg_install = PackageInstall::load(&ident, None)?;
                    let cmd_env = pkg_install.environment_for_command()?;
                    for (key, value) in cmd_env.into_iter() {
                        debug!("Setting: {}='{}'", key, value);
                        env::set_var(key, value);
                    }

                    let mut display_args = super::STUDIO_CMD.to_string();
                    for arg in args {
                        display_args.push(' ');
                        display_args.push_str(arg.to_string_lossy().as_ref());
                    }
                    debug!("Running: {}", display_args);

                    exec::command_from_min_pkg(ui, super::STUDIO_CMD, &ident)?
                }
            };

            if let Some(cmd) = find_command(command.to_string_lossy().as_ref()) {
                process::become_command(cmd, args)?;
                Ok(())
            } else {
                Err(Error::ExecCommandNotFound(command))
            }
        }
    }

    fn is_docker_studio(args: &[OsString]) -> bool {
        if cfg!(not(target_os = "linux")) {
            return false;
        }

        for arg in args.iter() {
            let str_arg = arg.to_string_lossy();
            if str_arg == "-D" {
                return true;
            }
        }

        false
    }

    fn has_docker_group() -> bool {
        let current_user = group::get_current_username().unwrap();
        let docker_members = group::get_members_by_groupname("docker");
        docker_members.map_or(false, |d| d.contains(&current_user))
    }

    fn rerun_with_sudo_if_needed(ui: &mut UI, args: &[OsString]) -> Result<()> {
        // If I have root permissions or if I am executing a docker studio
        // and have the appropriate group - early return, we are done.
        if am_i_root() || (is_docker_studio(args) && has_docker_group()) {
            return Ok(());
        }

        // Otherwise we will try to re-run this program using `sudo`
        match find_command(SUDO_CMD) {
            Some(sudo_prog) => {
                let mut args: Vec<OsString> = vec!["-p".into(),
                                                   "[sudo hab-studio] password for %u: ".into(),
                                                   "-E".into(),];
                args.append(&mut env::args_os().collect());
                process::become_command(sudo_prog, &args)?;
                Ok(())
            }
            None => {
                ui.warn(format!("Could not find the `{}' command, is it in your PATH?",
                                SUDO_CMD))?;
                ui.warn("Running Habitat Studio requires root or administrator privileges. \
                         Please retry this command as a super user or use a privilege-granting \
                         facility such as sudo.")?;
                ui.br()?;
                Err(Error::RootRequired)
            }
        }
    }
}

#[cfg(not(target_os = "linux"))]
mod inner {
    use crate::{command::studio::docker,
                common::ui::UI,
                error::{Error,
                        Result},
                exec,
                hcore::{crypto::init,
                        env as henv,
                        fs::find_command,
                        os::process,
                        package::PackageIdent},
                VERSION};
    use std::{ffi::OsString,
              path::PathBuf,
              str::FromStr};

    pub fn start(_ui: &mut UI, args: &[OsString]) -> Result<()> {
        if is_windows_studio(&args) {
            start_windows_studio(_ui, args)
        } else {
            docker::start_docker_studio(_ui, args)
        }
    }

    pub fn start_windows_studio(ui: &mut UI, args: &[OsString]) -> Result<()> {
        let command = match henv::var(super::STUDIO_CMD_ENVVAR) {
            Ok(command) => PathBuf::from(command),
            Err(_) => {
                init();
                let version: Vec<&str> = VERSION.split('/').collect();
                let ident = PackageIdent::from_str(&format!("{}/{}",
                                                            super::STUDIO_PACKAGE_IDENT,
                                                            version[0]))?;
                exec::command_from_min_pkg(ui, super::STUDIO_CMD, &ident)?
            }
        };

        if let Some(cmd) = find_command(command.to_string_lossy().as_ref()) {
            process::become_command(cmd, args)?;
        } else {
            return Err(Error::ExecCommandNotFound(command));
        }
        Ok(())
    }

    fn is_windows_studio(args: &[OsString]) -> bool {
        if cfg!(not(target_os = "windows")) {
            return false;
        }

        for arg in args.iter() {
            let str_arg = arg.to_string_lossy();
            if str_arg == "-D" {
                return false;
            }
        }

        // -w/--windows is deprecated and should be removed in a post 0.64.0 release
        for arg in args.iter() {
            let str_arg = arg.to_string_lossy().to_lowercase();
            if str_arg == "--windows" || str_arg == "-w" {
                return true;
            }
        }

        true
    }
}
<commit_msg>rustfmt cleanup: `cargo +nightly-2019-03-04 fmt<commit_after>use std::{env,
          ffi::OsString,
          fs as stdfs,
          path::PathBuf};

use crate::{common::ui::UI,
            hcore::{crypto::CACHE_KEY_PATH_ENV_VAR,
                    env as henv,
                    fs}};

use crate::{config,
            error::Result};

pub const ARTIFACT_PATH_ENVVAR: &str = "ARTIFACT_PATH";

const AUTH_TOKEN_ENVVAR: &str = "HAB_AUTH_TOKEN";
const BLDR_URL_ENVVAR: &str = "HAB_BLDR_URL";
const CTL_SECRET_ENVVAR: &str = "HAB_CTL_SECRET";
const ORIGIN_ENVVAR: &str = "HAB_ORIGIN";
const STUDIO_CMD: &str = "hab-studio";
const STUDIO_CMD_ENVVAR: &str = "HAB_STUDIO_BINARY";
const STUDIO_PACKAGE_IDENT: &str = "core/hab-studio";

fn set_env_var_from_config(env_var: &str, config_val: Option<String>) {
    if henv::var(env_var).is_err() {
        if let Some(val) = config_val {
            debug!("Setting {}={} via config file", env_var, val);
            env::set_var(env_var, val);
        }
    }
}

pub fn start(ui: &mut UI, args: &[OsString]) -> Result<()> {
    let config = config::load()?;

    set_env_var_from_config(AUTH_TOKEN_ENVVAR, config.auth_token);
    set_env_var_from_config(BLDR_URL_ENVVAR, config.bldr_url);
    set_env_var_from_config(CTL_SECRET_ENVVAR, config.ctl_secret);
    set_env_var_from_config(ORIGIN_ENVVAR, config.origin);

    if henv::var(CACHE_KEY_PATH_ENV_VAR).is_err() {
        let path = fs::cache_key_path(None::<&str>);
        debug!("Setting {}={}", CACHE_KEY_PATH_ENV_VAR, path.display());
        env::set_var(CACHE_KEY_PATH_ENV_VAR, &path);
    };

    let artifact_path = match henv::var(ARTIFACT_PATH_ENVVAR) {
        Ok(p) => PathBuf::from(p),
        Err(_) => {
            let path = fs::cache_artifact_path(None::<&str>);
            debug!("Setting {}={}", ARTIFACT_PATH_ENVVAR, path.display());
            env::set_var(ARTIFACT_PATH_ENVVAR, &path);
            path
        }
    };
    if !artifact_path.is_dir() {
        debug!("Creating artifact_path at: {}", artifact_path.display());
        stdfs::create_dir_all(&artifact_path)?;
    }

    inner::start(ui, args)
}

#[cfg(target_os = "linux")]
mod inner {
    use crate::{command::studio::docker,
                common::ui::{UIWriter,
                             UI},
                error::{Error,
                        Result},
                exec,
                hcore::{crypto::init,
                        env as henv,
                        fs::{am_i_root,
                             find_command},
                        os::process,
                        package::{PackageIdent,
                                  PackageInstall},
                        users::linux as group},
                VERSION};
    use std::{env,
              ffi::OsString,
              path::PathBuf,
              str::FromStr};

    const SUDO_CMD: &str = "sudo";

    pub fn start(ui: &mut UI, args: &[OsString]) -> Result<()> {
        rerun_with_sudo_if_needed(ui, &args)?;
        if is_docker_studio(&args) {
            docker::start_docker_studio(ui, args)
        } else {
            let command = match henv::var(super::STUDIO_CMD_ENVVAR) {
                Ok(command) => PathBuf::from(command),
                Err(_) => {
                    init();
                    let version: Vec<&str> = VERSION.split('/').collect();
                    let ident = PackageIdent::from_str(&format!("{}/{}",
                                                                super::STUDIO_PACKAGE_IDENT,
                                                                version[0]))?;
                    // This is a duplicate of the code in `hab pkg exec` and
                    // should be refactored as part of or after:
                    // https://github.com/habitat-sh/habitat/issues/6633
                    // https://github.com/habitat-sh/habitat/issues/6634
                    let pkg_install = PackageInstall::load(&ident, None)?;
                    let cmd_env = pkg_install.environment_for_command()?;
                    for (key, value) in cmd_env.into_iter() {
                        debug!("Setting: {}='{}'", key, value);
                        env::set_var(key, value);
                    }

                    let mut display_args = super::STUDIO_CMD.to_string();
                    for arg in args {
                        display_args.push(' ');
                        display_args.push_str(arg.to_string_lossy().as_ref());
                    }
                    debug!("Running: {}", display_args);

                    exec::command_from_min_pkg(ui, super::STUDIO_CMD, &ident)?
                }
            };

            if let Some(cmd) = find_command(command.to_string_lossy().as_ref()) {
                process::become_command(cmd, args)?;
                Ok(())
            } else {
                Err(Error::ExecCommandNotFound(command))
            }
        }
    }

    fn is_docker_studio(args: &[OsString]) -> bool {
        if cfg!(not(target_os = "linux")) {
            return false;
        }

        for arg in args.iter() {
            let str_arg = arg.to_string_lossy();
            if str_arg == "-D" {
                return true;
            }
        }

        false
    }

    fn has_docker_group() -> bool {
        let current_user = group::get_current_username().unwrap();
        let docker_members = group::get_members_by_groupname("docker");
        docker_members.map_or(false, |d| d.contains(&current_user))
    }

    fn rerun_with_sudo_if_needed(ui: &mut UI, args: &[OsString]) -> Result<()> {
        // If I have root permissions or if I am executing a docker studio
        // and have the appropriate group - early return, we are done.
        if am_i_root() || (is_docker_studio(args) && has_docker_group()) {
            return Ok(());
        }

        // Otherwise we will try to re-run this program using `sudo`
        match find_command(SUDO_CMD) {
            Some(sudo_prog) => {
                let mut args: Vec<OsString> = vec!["-p".into(),
                                                   "[sudo hab-studio] password for %u: ".into(),
                                                   "-E".into(),];
                args.append(&mut env::args_os().collect());
                process::become_command(sudo_prog, &args)?;
                Ok(())
            }
            None => {
                ui.warn(format!("Could not find the `{}' command, is it in your PATH?",
                                SUDO_CMD))?;
                ui.warn("Running Habitat Studio requires root or administrator privileges. \
                         Please retry this command as a super user or use a privilege-granting \
                         facility such as sudo.")?;
                ui.br()?;
                Err(Error::RootRequired)
            }
        }
    }
}

#[cfg(not(target_os = "linux"))]
mod inner {
    use crate::{command::studio::docker,
                common::ui::UI,
                error::{Error,
                        Result},
                exec,
                hcore::{crypto::init,
                        env as henv,
                        fs::find_command,
                        os::process,
                        package::PackageIdent},
                VERSION};
    use std::{ffi::OsString,
              path::PathBuf,
              str::FromStr};

    pub fn start(_ui: &mut UI, args: &[OsString]) -> Result<()> {
        if is_windows_studio(&args) {
            start_windows_studio(_ui, args)
        } else {
            docker::start_docker_studio(_ui, args)
        }
    }

    pub fn start_windows_studio(ui: &mut UI, args: &[OsString]) -> Result<()> {
        let command = match henv::var(super::STUDIO_CMD_ENVVAR) {
            Ok(command) => PathBuf::from(command),
            Err(_) => {
                init();
                let version: Vec<&str> = VERSION.split('/').collect();
                let ident = PackageIdent::from_str(&format!("{}/{}",
                                                            super::STUDIO_PACKAGE_IDENT,
                                                            version[0]))?;
                exec::command_from_min_pkg(ui, super::STUDIO_CMD, &ident)?
            }
        };

        if let Some(cmd) = find_command(command.to_string_lossy().as_ref()) {
            process::become_command(cmd, args)?;
        } else {
            return Err(Error::ExecCommandNotFound(command));
        }
        Ok(())
    }

    fn is_windows_studio(args: &[OsString]) -> bool {
        if cfg!(not(target_os = "windows")) {
            return false;
        }

        for arg in args.iter() {
            let str_arg = arg.to_string_lossy();
            if str_arg == "-D" {
                return false;
            }
        }

        // -w/--windows is deprecated and should be removed in a post 0.64.0 release
        for arg in args.iter() {
            let str_arg = arg.to_string_lossy().to_lowercase();
            if str_arg == "--windows" || str_arg == "-w" {
                return true;
            }
        }

        true
    }
}
<|endoftext|>


<commit_before>use std::{env,
          ffi::OsString,
          fs as stdfs,
          path::PathBuf};

use crate::{common::ui::UI,
            hcore::{crypto::CACHE_KEY_PATH_ENV_VAR,
                    env as henv,
                    fs}};

use crate::{config,
            error::Result,
            BLDR_URL_ENVVAR,
            CTL_SECRET_ENVVAR,
            ORIGIN_ENVVAR};

use habitat_core::AUTH_TOKEN_ENVVAR;

pub const ARTIFACT_PATH_ENVVAR: &str = "ARTIFACT_PATH";

const STUDIO_CMD: &str = "hab-studio";
const STUDIO_CMD_ENVVAR: &str = "HAB_STUDIO_BINARY";
const STUDIO_PACKAGE_IDENT: &str = "core/hab-studio";

fn set_env_var_from_config(env_var: &str, config_val: Option<String>, sensitive: bool) {
    if henv::var(env_var).is_err() {
        if let Some(val) = config_val {
            if sensitive {
                debug!("Setting {}=REDACTED (sensitive) via config file", env_var)
            } else {
                debug!("Setting {}={} via config file", env_var, val)
            }
            env::set_var(env_var, val);
        }
    }
}

pub fn start(ui: &mut UI, args: &[OsString]) -> Result<()> {
    let config = config::load()?;

    set_env_var_from_config(AUTH_TOKEN_ENVVAR, config.auth_token, true);
    set_env_var_from_config(BLDR_URL_ENVVAR, config.bldr_url, false);
    set_env_var_from_config(CTL_SECRET_ENVVAR, config.ctl_secret, true);
    set_env_var_from_config(ORIGIN_ENVVAR, config.origin, false);

    if henv::var(CACHE_KEY_PATH_ENV_VAR).is_err() {
        let path = fs::cache_key_path(None::<&str>);
        debug!("Setting {}={}", CACHE_KEY_PATH_ENV_VAR, path.display());
        env::set_var(CACHE_KEY_PATH_ENV_VAR, &path);
    };

    let artifact_path = match henv::var(ARTIFACT_PATH_ENVVAR) {
        Ok(p) => PathBuf::from(p),
        Err(_) => {
            let path = fs::cache_artifact_path(None::<&str>);
            debug!("Setting {}={}", ARTIFACT_PATH_ENVVAR, path.display());
            env::set_var(ARTIFACT_PATH_ENVVAR, &path);
            path
        }
    };
    if !artifact_path.is_dir() {
        debug!("Creating artifact_path at: {}", artifact_path.display());
        stdfs::create_dir_all(&artifact_path)?;
    }

    inner::start(ui, args)
}

#[cfg(target_os = "linux")]
mod inner {
    use crate::{command::studio::docker,
                common::ui::{UIWriter,
                             UI},
                error::{Error,
                        Result},
                exec,
                hcore::{crypto::init,
                        env as henv,
                        fs::{am_i_root,
                             find_command},
                        os::process,
                        package::{PackageIdent,
                                  PackageInstall},
                        users::linux as group},
                VERSION};
    use std::{env,
              ffi::OsString,
              path::PathBuf,
              str::FromStr};

    const SUDO_CMD: &str = "sudo";

    pub fn start(ui: &mut UI, args: &[OsString]) -> Result<()> {
        rerun_with_sudo_if_needed(ui, &args)?;
        if is_docker_studio(&args) {
            docker::start_docker_studio(ui, args)
        } else {
            let command = match henv::var(super::STUDIO_CMD_ENVVAR) {
                Ok(command) => PathBuf::from(command),
                Err(_) => {
                    init();
                    let version: Vec<&str> = VERSION.split('/').collect();
                    let ident = PackageIdent::from_str(&format!("{}/{}",
                                                                super::STUDIO_PACKAGE_IDENT,
                                                                version[0]))?;
                    // This is a duplicate of the code in `hab pkg exec` and
                    // should be refactored as part of or after:
                    // https://github.com/habitat-sh/habitat/issues/6633
                    // https://github.com/habitat-sh/habitat/issues/6634
                    let pkg_install = PackageInstall::load(&ident, None)?;
                    let cmd_env = pkg_install.environment_for_command()?;
                    for (key, value) in cmd_env.into_iter() {
                        debug!("Setting: {}='{}'", key, value);
                        env::set_var(key, value);
                    }

                    let mut display_args = super::STUDIO_CMD.to_string();
                    for arg in args {
                        display_args.push(' ');
                        display_args.push_str(arg.to_string_lossy().as_ref());
                    }
                    debug!("Running: {}", display_args);

                    exec::command_from_min_pkg(ui, super::STUDIO_CMD, &ident)?
                }
            };

            if let Some(cmd) = find_command(command.to_string_lossy().as_ref()) {
                process::become_command(cmd, args)?;
                Ok(())
            } else {
                Err(Error::ExecCommandNotFound(command))
            }
        }
    }

    fn is_docker_studio(args: &[OsString]) -> bool {
        if cfg!(not(target_os = "linux")) {
            return false;
        }

        for arg in args.iter() {
            let str_arg = arg.to_string_lossy();
            if str_arg == "-D" {
                return true;
            }
        }

        false
    }

    fn has_docker_group() -> bool {
        let current_user = group::get_current_username().unwrap();
        let docker_members = group::get_members_by_groupname("docker");
        docker_members.map_or(false, |d| d.contains(&current_user))
    }

    fn rerun_with_sudo_if_needed(ui: &mut UI, args: &[OsString]) -> Result<()> {
        // If I have root permissions or if I am executing a docker studio
        // and have the appropriate group - early return, we are done.
        if am_i_root() || (is_docker_studio(args) && has_docker_group()) {
            return Ok(());
        }

        // Otherwise we will try to re-run this program using `sudo`
        match find_command(SUDO_CMD) {
            Some(sudo_prog) => {
                let mut args: Vec<OsString> = vec!["-p".into(),
                                                   "[sudo hab-studio] password for %u: ".into(),
                                                   "-E".into(),];
                args.append(&mut env::args_os().collect());
                process::become_command(sudo_prog, &args)?;
                Ok(())
            }
            None => {
                ui.warn(format!("Could not find the `{}' command, is it in your PATH?",
                                SUDO_CMD))?;
                ui.warn("Running Habitat Studio requires root or administrator privileges. \
                         Please retry this command as a super user or use a privilege-granting \
                         facility such as sudo.")?;
                ui.br()?;
                Err(Error::RootRequired)
            }
        }
    }
}

#[cfg(not(target_os = "linux"))]
mod inner {
    use crate::{command::studio::docker,
                common::ui::UI,
                error::{Error,
                        Result},
                exec,
                hcore::{crypto::init,
                        env as henv,
                        fs::find_command,
                        os::process,
                        package::PackageIdent},
                VERSION};
    use std::{ffi::OsString,
              path::PathBuf,
              str::FromStr};

    pub fn start(_ui: &mut UI, args: &[OsString]) -> Result<()> {
        if is_windows_studio(&args) {
            start_windows_studio(_ui, args)
        } else {
            docker::start_docker_studio(_ui, args)
        }
    }

    pub fn start_windows_studio(ui: &mut UI, args: &[OsString]) -> Result<()> {
        let command = match henv::var(super::STUDIO_CMD_ENVVAR) {
            Ok(command) => PathBuf::from(command),
            Err(_) => {
                init();
                let version: Vec<&str> = VERSION.split('/').collect();
                let ident = PackageIdent::from_str(&format!("{}/{}",
                                                            super::STUDIO_PACKAGE_IDENT,
                                                            version[0]))?;
                exec::command_from_min_pkg(ui, super::STUDIO_CMD, &ident)?
            }
        };

        if let Some(cmd) = find_command(command.to_string_lossy().as_ref()) {
            process::become_command(cmd, args)?;
        } else {
            return Err(Error::ExecCommandNotFound(command));
        }
        Ok(())
    }

    fn is_windows_studio(args: &[OsString]) -> bool {
        if cfg!(not(target_os = "windows")) {
            return false;
        }

        for arg in args.iter() {
            let str_arg = arg.to_string_lossy();
            if str_arg == "-D" {
                return false;
            }
        }

        // -w/--windows is deprecated and should be removed in a post 0.64.0 release
        for arg in args.iter() {
            let str_arg = arg.to_string_lossy().to_lowercase();
            if str_arg == "--windows" || str_arg == "-w" {
                return true;
            }
        }

        true
    }
}
<commit_msg>Add Sensitivity enum to ensure sensitive information is not printed<commit_after>use std::{env,
          ffi::OsString,
          fs as stdfs,
          path::PathBuf};

use crate::{common::ui::UI,
            hcore::{crypto::CACHE_KEY_PATH_ENV_VAR,
                    env as henv,
                    fs}};

use crate::{config,
            error::Result,
            BLDR_URL_ENVVAR,
            CTL_SECRET_ENVVAR,
            ORIGIN_ENVVAR};

use habitat_core::AUTH_TOKEN_ENVVAR;

pub const ARTIFACT_PATH_ENVVAR: &str = "ARTIFACT_PATH";

const STUDIO_CMD: &str = "hab-studio";
const STUDIO_CMD_ENVVAR: &str = "HAB_STUDIO_BINARY";
const STUDIO_PACKAGE_IDENT: &str = "core/hab-studio";

enum Sensitivity {
    PrintValue,
    NoPrintValue,
}

fn set_env_var_from_config(env_var: &str, config_val: Option<String>, sensitive: Sensitivity) {
    if henv::var(env_var).is_err() {
        if let Some(val) = config_val {
            match sensitive {
                Sensitivity::NoPrintValue =>  debug!("Setting {}=REDACTED (sensitive) via config file", env_var),
                Sensitivity::PrintValue => debug!("Setting {}={} via config file", env_var, val),
            }
            env::set_var(env_var, val);
        }
    }
}

pub fn start(ui: &mut UI, args: &[OsString]) -> Result<()> {
    let config = config::load()?;

    set_env_var_from_config(AUTH_TOKEN_ENVVAR, config.auth_token, Sensitivity::NoPrintValue);
    set_env_var_from_config(BLDR_URL_ENVVAR, config.bldr_url, Sensitivity::PrintValue);
    set_env_var_from_config(CTL_SECRET_ENVVAR, config.ctl_secret, Sensitivity::NoPrintValue);
    set_env_var_from_config(ORIGIN_ENVVAR, config.origin, Sensitivity::PrintValue);

    if henv::var(CACHE_KEY_PATH_ENV_VAR).is_err() {
        let path = fs::cache_key_path(None::<&str>);
        debug!("Setting {}={}", CACHE_KEY_PATH_ENV_VAR, path.display());
        env::set_var(CACHE_KEY_PATH_ENV_VAR, &path);
    };

    let artifact_path = match henv::var(ARTIFACT_PATH_ENVVAR) {
        Ok(p) => PathBuf::from(p),
        Err(_) => {
            let path = fs::cache_artifact_path(None::<&str>);
            debug!("Setting {}={}", ARTIFACT_PATH_ENVVAR, path.display());
            env::set_var(ARTIFACT_PATH_ENVVAR, &path);
            path
        }
    };
    if !artifact_path.is_dir() {
        debug!("Creating artifact_path at: {}", artifact_path.display());
        stdfs::create_dir_all(&artifact_path)?;
    }

    inner::start(ui, args)
}

#[cfg(target_os = "linux")]
mod inner {
    use crate::{command::studio::docker,
                common::ui::{UIWriter,
                             UI},
                error::{Error,
                        Result},
                exec,
                hcore::{crypto::init,
                        env as henv,
                        fs::{am_i_root,
                             find_command},
                        os::process,
                        package::{PackageIdent,
                                  PackageInstall},
                        users::linux as group},
                VERSION};
    use std::{env,
              ffi::OsString,
              path::PathBuf,
              str::FromStr};

    const SUDO_CMD: &str = "sudo";

    pub fn start(ui: &mut UI, args: &[OsString]) -> Result<()> {
        rerun_with_sudo_if_needed(ui, &args)?;
        if is_docker_studio(&args) {
            docker::start_docker_studio(ui, args)
        } else {
            let command = match henv::var(super::STUDIO_CMD_ENVVAR) {
                Ok(command) => PathBuf::from(command),
                Err(_) => {
                    init();
                    let version: Vec<&str> = VERSION.split('/').collect();
                    let ident = PackageIdent::from_str(&format!("{}/{}",
                                                                super::STUDIO_PACKAGE_IDENT,
                                                                version[0]))?;
                    // This is a duplicate of the code in `hab pkg exec` and
                    // should be refactored as part of or after:
                    // https://github.com/habitat-sh/habitat/issues/6633
                    // https://github.com/habitat-sh/habitat/issues/6634
                    let pkg_install = PackageInstall::load(&ident, None)?;
                    let cmd_env = pkg_install.environment_for_command()?;
                    for (key, value) in cmd_env.into_iter() {
                        debug!("Setting: {}='{}'", key, value);
                        env::set_var(key, value);
                    }

                    let mut display_args = super::STUDIO_CMD.to_string();
                    for arg in args {
                        display_args.push(' ');
                        display_args.push_str(arg.to_string_lossy().as_ref());
                    }
                    debug!("Running: {}", display_args);

                    exec::command_from_min_pkg(ui, super::STUDIO_CMD, &ident)?
                }
            };

            if let Some(cmd) = find_command(command.to_string_lossy().as_ref()) {
                process::become_command(cmd, args)?;
                Ok(())
            } else {
                Err(Error::ExecCommandNotFound(command))
            }
        }
    }

    fn is_docker_studio(args: &[OsString]) -> bool {
        if cfg!(not(target_os = "linux")) {
            return false;
        }

        for arg in args.iter() {
            let str_arg = arg.to_string_lossy();
            if str_arg == "-D" {
                return true;
            }
        }

        false
    }

    fn has_docker_group() -> bool {
        let current_user = group::get_current_username().unwrap();
        let docker_members = group::get_members_by_groupname("docker");
        docker_members.map_or(false, |d| d.contains(&current_user))
    }

    fn rerun_with_sudo_if_needed(ui: &mut UI, args: &[OsString]) -> Result<()> {
        // If I have root permissions or if I am executing a docker studio
        // and have the appropriate group - early return, we are done.
        if am_i_root() || (is_docker_studio(args) && has_docker_group()) {
            return Ok(());
        }

        // Otherwise we will try to re-run this program using `sudo`
        match find_command(SUDO_CMD) {
            Some(sudo_prog) => {
                let mut args: Vec<OsString> = vec!["-p".into(),
                                                   "[sudo hab-studio] password for %u: ".into(),
                                                   "-E".into(),];
                args.append(&mut env::args_os().collect());
                process::become_command(sudo_prog, &args)?;
                Ok(())
            }
            None => {
                ui.warn(format!("Could not find the `{}' command, is it in your PATH?",
                                SUDO_CMD))?;
                ui.warn("Running Habitat Studio requires root or administrator privileges. \
                         Please retry this command as a super user or use a privilege-granting \
                         facility such as sudo.")?;
                ui.br()?;
                Err(Error::RootRequired)
            }
        }
    }
}

#[cfg(not(target_os = "linux"))]
mod inner {
    use crate::{command::studio::docker,
                common::ui::UI,
                error::{Error,
                        Result},
                exec,
                hcore::{crypto::init,
                        env as henv,
                        fs::find_command,
                        os::process,
                        package::PackageIdent},
                VERSION};
    use std::{ffi::OsString,
              path::PathBuf,
              str::FromStr};

    pub fn start(_ui: &mut UI, args: &[OsString]) -> Result<()> {
        if is_windows_studio(&args) {
            start_windows_studio(_ui, args)
        } else {
            docker::start_docker_studio(_ui, args)
        }
    }

    pub fn start_windows_studio(ui: &mut UI, args: &[OsString]) -> Result<()> {
        let command = match henv::var(super::STUDIO_CMD_ENVVAR) {
            Ok(command) => PathBuf::from(command),
            Err(_) => {
                init();
                let version: Vec<&str> = VERSION.split('/').collect();
                let ident = PackageIdent::from_str(&format!("{}/{}",
                                                            super::STUDIO_PACKAGE_IDENT,
                                                            version[0]))?;
                exec::command_from_min_pkg(ui, super::STUDIO_CMD, &ident)?
            }
        };

        if let Some(cmd) = find_command(command.to_string_lossy().as_ref()) {
            process::become_command(cmd, args)?;
        } else {
            return Err(Error::ExecCommandNotFound(command));
        }
        Ok(())
    }

    fn is_windows_studio(args: &[OsString]) -> bool {
        if cfg!(not(target_os = "windows")) {
            return false;
        }

        for arg in args.iter() {
            let str_arg = arg.to_string_lossy();
            if str_arg == "-D" {
                return false;
            }
        }

        // -w/--windows is deprecated and should be removed in a post 0.64.0 release
        for arg in args.iter() {
            let str_arg = arg.to_string_lossy().to_lowercase();
            if str_arg == "--windows" || str_arg == "-w" {
                return true;
            }
        }

        true
    }
}
<|endoftext|>


<commit_before>use layout::{AnonymousBlock, BlockNode, InlineNode, LayoutBox, Rect};
use css::{Value, Color};

pub struct Canvas {
    pub pixels: Vec<Color>,
    pub width: usize,
    pub height: usize,
}

/// Paint a tree of LayoutBoxes to an array of pixels.
pub fn paint(layout_root: &LayoutBox, bounds: Rect) -> Canvas {
    let display_list = build_display_list(layout_root);
    let mut canvas = Canvas::new(bounds.width as usize, bounds.height as usize);
    for item in display_list {
        canvas.paint_item(&item);
    }
    return canvas;
}

#[derive(Debug)]
pub enum DisplayCommand {
    SolidColor(Color, Rect),
}

pub type DisplayList = Vec<DisplayCommand>;

pub fn build_display_list(layout_root: &LayoutBox) -> DisplayList {
    let mut list = Vec::new();
    render_layout_box(&mut list, layout_root);
    return list;
}

fn render_layout_box(list: &mut DisplayList, layout_box: &LayoutBox) {
    render_background(list, layout_box);
    render_borders(list, layout_box);
    for child in &layout_box.children {
        render_layout_box(list, child);
    }
}

fn render_background(list: &mut DisplayList, layout_box: &LayoutBox) {
    get_color(layout_box, "background").map(|color|
        list.push(DisplayCommand::SolidColor(color, layout_box.dimensions.border_box())));
}

fn render_borders(list: &mut DisplayList, layout_box: &LayoutBox) {
    let color = match get_color(layout_box, "border-color") {
        Some(color) => color,
        _ => return
    };

    let d = &layout_box.dimensions;
    let border_box = d.border_box();

    // Left border
    list.push(DisplayCommand::SolidColor(color, Rect {
        x: border_box.x,
        y: border_box.y,
        width: d.border.left,
        height: border_box.height,
    }));

    // Right border
    list.push(DisplayCommand::SolidColor(color, Rect {
        x: border_box.x + border_box.width - d.border.right,
        y: border_box.y,
        width: d.border.right,
        height: border_box.height,
    }));

    // Top border
    list.push(DisplayCommand::SolidColor(color, Rect {
        x: border_box.x,
        y: border_box.y,
        width: border_box.width,
        height: d.border.top,
    }));

    // Bottom border
    list.push(DisplayCommand::SolidColor(color, Rect {
        x: border_box.x,
        y: border_box.y + border_box.height - d.border.bottom,
        width: border_box.width,
        height: d.border.bottom,
    }));
}

/// Return the specified color for CSS property `name`, or None if no color was specified.
fn get_color(layout_box: &LayoutBox, name: &str) -> Option<Color> {
    match layout_box.box_type {
        BlockNode(style) | InlineNode(style) => match style.value(name) {
            Some(Value::ColorValue(color)) => Some(color),
            _ => None
        },
        AnonymousBlock => None
    }
}

impl Canvas {
    /// Create a blank canvas
    fn new(width: usize, height: usize) -> Canvas {
        let white = Color { r: 255, g: 255, b: 255, a: 255 };
        return Canvas {
            pixels: vec![white; width * height],
            width: width,
            height: height,
        }
    }

    fn paint_item(&mut self, item: &DisplayCommand) {
        match item {
            &DisplayCommand::SolidColor(color, rect) => {
                // Clip the rectangle to the canvas boundaries.
                let x0 = rect.x.clamp(0.0, self.width as f32) as usize;
                let y0 = rect.y.clamp(0.0, self.height as f32) as usize;
                let x1 = (rect.x + rect.width).clamp(0.0, self.width as f32) as usize;
                let y1 = (rect.y + rect.height).clamp(0.0, self.height as f32) as usize;

                for y in (y0 .. y1) {
                    for x in (x0 .. x1) {
                        // TODO: alpha compositing with existing pixel
                        self.pixels[y * self.width + x] = color;
                    }
                }
            }
        }
    }
}

trait Clamp {
    fn clamp(self, lower: Self, upper: Self) -> Self;
}
impl Clamp for f32 {
    fn clamp(self, lower: f32, upper: f32) -> f32 {
        self.max(lower).min(upper)
    }
}
<commit_msg>Fix unnecessary parentheses warning around `for`<commit_after>use layout::{AnonymousBlock, BlockNode, InlineNode, LayoutBox, Rect};
use css::{Value, Color};

pub struct Canvas {
    pub pixels: Vec<Color>,
    pub width: usize,
    pub height: usize,
}

/// Paint a tree of LayoutBoxes to an array of pixels.
pub fn paint(layout_root: &LayoutBox, bounds: Rect) -> Canvas {
    let display_list = build_display_list(layout_root);
    let mut canvas = Canvas::new(bounds.width as usize, bounds.height as usize);
    for item in display_list {
        canvas.paint_item(&item);
    }
    return canvas;
}

#[derive(Debug)]
pub enum DisplayCommand {
    SolidColor(Color, Rect),
}

pub type DisplayList = Vec<DisplayCommand>;

pub fn build_display_list(layout_root: &LayoutBox) -> DisplayList {
    let mut list = Vec::new();
    render_layout_box(&mut list, layout_root);
    return list;
}

fn render_layout_box(list: &mut DisplayList, layout_box: &LayoutBox) {
    render_background(list, layout_box);
    render_borders(list, layout_box);
    for child in &layout_box.children {
        render_layout_box(list, child);
    }
}

fn render_background(list: &mut DisplayList, layout_box: &LayoutBox) {
    get_color(layout_box, "background").map(|color|
        list.push(DisplayCommand::SolidColor(color, layout_box.dimensions.border_box())));
}

fn render_borders(list: &mut DisplayList, layout_box: &LayoutBox) {
    let color = match get_color(layout_box, "border-color") {
        Some(color) => color,
        _ => return
    };

    let d = &layout_box.dimensions;
    let border_box = d.border_box();

    // Left border
    list.push(DisplayCommand::SolidColor(color, Rect {
        x: border_box.x,
        y: border_box.y,
        width: d.border.left,
        height: border_box.height,
    }));

    // Right border
    list.push(DisplayCommand::SolidColor(color, Rect {
        x: border_box.x + border_box.width - d.border.right,
        y: border_box.y,
        width: d.border.right,
        height: border_box.height,
    }));

    // Top border
    list.push(DisplayCommand::SolidColor(color, Rect {
        x: border_box.x,
        y: border_box.y,
        width: border_box.width,
        height: d.border.top,
    }));

    // Bottom border
    list.push(DisplayCommand::SolidColor(color, Rect {
        x: border_box.x,
        y: border_box.y + border_box.height - d.border.bottom,
        width: border_box.width,
        height: d.border.bottom,
    }));
}

/// Return the specified color for CSS property `name`, or None if no color was specified.
fn get_color(layout_box: &LayoutBox, name: &str) -> Option<Color> {
    match layout_box.box_type {
        BlockNode(style) | InlineNode(style) => match style.value(name) {
            Some(Value::ColorValue(color)) => Some(color),
            _ => None
        },
        AnonymousBlock => None
    }
}

impl Canvas {
    /// Create a blank canvas
    fn new(width: usize, height: usize) -> Canvas {
        let white = Color { r: 255, g: 255, b: 255, a: 255 };
        return Canvas {
            pixels: vec![white; width * height],
            width: width,
            height: height,
        }
    }

    fn paint_item(&mut self, item: &DisplayCommand) {
        match item {
            &DisplayCommand::SolidColor(color, rect) => {
                // Clip the rectangle to the canvas boundaries.
                let x0 = rect.x.clamp(0.0, self.width as f32) as usize;
                let y0 = rect.y.clamp(0.0, self.height as f32) as usize;
                let x1 = (rect.x + rect.width).clamp(0.0, self.width as f32) as usize;
                let y1 = (rect.y + rect.height).clamp(0.0, self.height as f32) as usize;

                for y in y0 .. y1 {
                    for x in x0 .. x1 {
                        // TODO: alpha compositing with existing pixel
                        self.pixels[y * self.width + x] = color;
                    }
                }
            }
        }
    }
}

trait Clamp {
    fn clamp(self, lower: Self, upper: Self) -> Self;
}
impl Clamp for f32 {
    fn clamp(self, lower: f32, upper: f32) -> f32 {
        self.max(lower).min(upper)
    }
}
<|endoftext|>


<commit_before>//! Pixel manipulations.

use image::Pixel;
use conv::ValueInto;
use math::cast;
use definitions::Clamp;

/// Adds pixels with the given weights. Results are clamped to prevent arithmetical overflows.
pub fn weighted_sum<P: Pixel>(left: P, right: P, left_weight: f32, right_weight: f32) -> P
where
    P::Subpixel: ValueInto<f32> + Clamp<f32>,
{
    left.map2(&right, |p, q| {
        weighted_channel_sum(p, q, left_weight, right_weight)
    })
}

/// Equivalent to `weighted_sum(left, right, left_weight, 1 - left_weight)`.
pub fn interpolate<P: Pixel>(left: P, right: P, left_weight: f32) -> P
where
    P::Subpixel: ValueInto<f32> + Clamp<f32>,
{
    weighted_sum(left, right, left_weight, 1.0 - left_weight)
}

#[inline(always)]
fn weighted_channel_sum<C>(left: C, right: C, left_weight: f32, right_weight: f32) -> C
where
    C: ValueInto<f32> + Clamp<f32>,
{
    Clamp::clamp(cast(left) * left_weight + cast(right) * right_weight)
}

#[cfg(test)]
mod test {
    use super::*;
    use image::{Rgb, Luma};
    use test::{Bencher, black_box};

    #[test]
    fn test_weighted_channel_sum() {
        // Midpoint
        assert_eq!(weighted_channel_sum(10u8, 20u8, 0.5, 0.5), 15u8);
        // Mainly left
        assert_eq!(weighted_channel_sum(10u8, 20u8, 0.9, 0.1), 11u8);
        // Clamped
        assert_eq!(weighted_channel_sum(150u8, 150u8, 1.8, 0.8), 255u8);
    }

    #[test]
    fn test_weighted_sum() {
        let left = Rgb([10u8, 20u8, 30u8]);
        let right = Rgb([100u8, 80u8, 60u8]);
        let sum = weighted_sum(left, right, 0.7, 0.3);
        assert_eq!(sum, Rgb([37, 38, 39]));
    }

    #[bench]
    fn bench_weighted_sum_rgb(b: &mut Bencher) {
        b.iter(|| {
            let left = black_box(Rgb([10u8, 20u8, 33u8]));
            let right = black_box(Rgb([80u8, 70u8, 60u8]));
            let left_weight = black_box(0.3);
            let right_weight = black_box(0.7);
            black_box(weighted_sum(left, right, left_weight, right_weight));
        })
    }

    #[bench]
    fn bench_weighted_sum_gray(b: &mut Bencher) {
        b.iter(|| {
            let left = black_box(Luma([10u8]));
            let right = black_box(Luma([80u8]));
            let left_weight = black_box(0.3);
            let right_weight = black_box(0.7);
            black_box(weighted_sum(left, right, left_weight, right_weight));
        })
    }

    #[bench]
    fn bench_interpolate_rgb(b: &mut Bencher) {
        b.iter(|| {
            let left = black_box(Rgb([10u8, 20u8, 33u8]));
            let right = black_box(Rgb([80u8, 70u8, 60u8]));
            let left_weight = black_box(0.3);
            black_box(interpolate(left, right, left_weight));
        })
    }

    #[bench]
    fn bench_interpolate_gray(b: &mut Bencher) {
        b.iter(|| {
            let left = black_box(Luma([10u8]));
            let right = black_box(Luma([80u8]));
            let left_weight = black_box(0.3);
            black_box(interpolate(left, right, left_weight));
        })
    }
}
<commit_msg>Add doctests for pixelops module<commit_after>//! Pixel manipulations.

use image::Pixel;
use conv::ValueInto;
use math::cast;
use definitions::Clamp;

/// Adds pixels with the given weights. Results are clamped to prevent arithmetical overflows.
///
/// # Examples
/// ```
/// # extern crate image;
/// # extern crate imageproc;
/// # fn main() {
/// use image::Rgb;
/// use imageproc::pixelops::weighted_sum;
///
/// let left = Rgb([10u8, 20u8, 30u8]);
/// let right = Rgb([100u8, 80u8, 60u8]);
///
/// let sum = weighted_sum(left, right, 0.7, 0.3);
/// assert_eq!(sum, Rgb([37, 38, 39]));
/// # }
/// ```
pub fn weighted_sum<P: Pixel>(left: P, right: P, left_weight: f32, right_weight: f32) -> P
where
    P::Subpixel: ValueInto<f32> + Clamp<f32>,
{
    left.map2(&right, |p, q| {
        weighted_channel_sum(p, q, left_weight, right_weight)
    })
}

/// Equivalent to `weighted_sum(left, right, left_weight, 1 - left_weight)`.
///
/// # Examples
/// ```
/// # extern crate image;
/// # extern crate imageproc;
/// # fn main() {
/// use image::Rgb;
/// use imageproc::pixelops::interpolate;
///
/// let left = Rgb([10u8, 20u8, 30u8]);
/// let right = Rgb([100u8, 80u8, 60u8]);
///
/// let sum = interpolate(left, right, 0.7);
/// assert_eq!(sum, Rgb([37, 38, 39]));
/// # }
/// ```
pub fn interpolate<P: Pixel>(left: P, right: P, left_weight: f32) -> P
where
    P::Subpixel: ValueInto<f32> + Clamp<f32>,
{
    weighted_sum(left, right, left_weight, 1.0 - left_weight)
}

#[inline(always)]
fn weighted_channel_sum<C>(left: C, right: C, left_weight: f32, right_weight: f32) -> C
where
    C: ValueInto<f32> + Clamp<f32>,
{
    Clamp::clamp(cast(left) * left_weight + cast(right) * right_weight)
}

#[cfg(test)]
mod test {
    use super::*;
    use image::{Rgb, Luma};
    use test::{Bencher, black_box};

    #[test]
    fn test_weighted_channel_sum() {
        // Midpoint
        assert_eq!(weighted_channel_sum(10u8, 20u8, 0.5, 0.5), 15u8);
        // Mainly left
        assert_eq!(weighted_channel_sum(10u8, 20u8, 0.9, 0.1), 11u8);
        // Clamped
        assert_eq!(weighted_channel_sum(150u8, 150u8, 1.8, 0.8), 255u8);
    }

    #[bench]
    fn bench_weighted_sum_rgb(b: &mut Bencher) {
        b.iter(|| {
            let left = black_box(Rgb([10u8, 20u8, 33u8]));
            let right = black_box(Rgb([80u8, 70u8, 60u8]));
            let left_weight = black_box(0.3);
            let right_weight = black_box(0.7);
            black_box(weighted_sum(left, right, left_weight, right_weight));
        })
    }

    #[bench]
    fn bench_weighted_sum_gray(b: &mut Bencher) {
        b.iter(|| {
            let left = black_box(Luma([10u8]));
            let right = black_box(Luma([80u8]));
            let left_weight = black_box(0.3);
            let right_weight = black_box(0.7);
            black_box(weighted_sum(left, right, left_weight, right_weight));
        })
    }

    #[bench]
    fn bench_interpolate_rgb(b: &mut Bencher) {
        b.iter(|| {
            let left = black_box(Rgb([10u8, 20u8, 33u8]));
            let right = black_box(Rgb([80u8, 70u8, 60u8]));
            let left_weight = black_box(0.3);
            black_box(interpolate(left, right, left_weight));
        })
    }

    #[bench]
    fn bench_interpolate_gray(b: &mut Bencher) {
        b.iter(|| {
            let left = black_box(Luma([10u8]));
            let right = black_box(Luma([80u8]));
            let left_weight = black_box(0.3);
            black_box(interpolate(left, right, left_weight));
        })
    }
}
<|endoftext|>