<commit_before>//! TODO Fill in

use std::{ptr, mem};
use xcb::{xkb, Connection};
use xcb::ffi::{self, xproto};
use nix::{self, libc};
use render;
use std::ffi::{CStr, CString};
use gdk_pixbuf::Pixbuf;
use glib::translate::ToGlibPtr;
use std::fmt::{self, Display, Formatter};
use std::process::{Command, Stdio};
use std::thread;
use std::default::Default;
use rlua::{self, Table, Lua, UserData, ToLua, Value, LightUserData, UserDataMethods, MetaMethod,
           AnyUserData};
use super::{XCB_CONNECTION_HANDLE, signal};
use super::xproperty::{XProperty, XPropertyType, PROPERTIES};

// TODO FIXME
// Often we are getting some raw pointers from the xcb replies
// we need to free them because the memory management for them is manual.

// TODO this isn't defined in the xcb crate, even though it should be.
// A patch should be open adding this to its generation scheme
extern "C" {
    fn xcb_xkb_get_names_value_list_unpack(buffer: *mut libc::c_void,
                                           nTypes: u8,
                                           indicators: u32,
                                           virtualMods: u16,
                                           groupNames: u8,
                                           nKeys: u8,
                                           nKeyAliases: u8,
                                           nRadioGroups: u8,
                                           which: u32,
                                           aux: *mut ffi::xkb::xcb_xkb_get_names_value_list_t) -> libc::c_int;
}

#[derive(Clone, Debug)]
pub struct AwesomeState {
    preferred_icon_size: u32
}

impl Default for AwesomeState {
    fn default() -> Self {
        AwesomeState {
            preferred_icon_size: 0
        }
    }
}

impl Display for AwesomeState {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        write!(f, "Awesome: {:p}", self)
    }
}

impl UserData for AwesomeState {
    fn add_methods(methods: &mut UserDataMethods<Self>) {
        fn index<'lua>(_: &'lua Lua,
                      (awesome, index): (AnyUserData<'lua>, Value<'lua>))
                      -> rlua::Result<rlua::Value<'lua>>
        {
            let table = awesome.get_user_value::<Table>()?;
            table.get::<_, Value>(index)
        };
        methods.add_meta_function(MetaMethod::Index, index);
    }
}

pub fn init(lua: &Lua) -> rlua::Result<()> {
    let awesome_table = lua.create_table()?;
    method_setup(lua, &awesome_table)?;
    property_setup(lua, &awesome_table)?;
    let globals = lua.globals();
    let awesome = lua.create_userdata(AwesomeState::default())?;
    awesome.set_user_value(awesome_table)?;
    globals.set("awesome", awesome)
}

fn method_setup<'lua>(lua: &'lua Lua, awesome_table: &Table<'lua>) -> rlua::Result<()> {
    // TODO Fill in rest
    awesome_table.set("connect_signal", lua.create_function(signal::global_connect_signal)?)?;
    awesome_table.set("disconnect_signal", lua.create_function(signal::global_disconnect_signal)?)?;
    awesome_table.set("emit_signal", lua.create_function(signal::global_emit_signal)?)?;
    awesome_table.set("xrdb_get_value", lua.create_function(xrdb_get_value)?)?;
    awesome_table.set("xkb_set_layout_group", lua.create_function(xkb_set_layout_group)?)?;
    awesome_table.set("xkb_get_layout_group", lua.create_function(xkb_get_layout_group)?)?;
    awesome_table.set("set_preferred_icon_size", lua.create_function(set_preferred_icon_size)?)?;
    awesome_table.set("register_xproperty", lua.create_function(register_xproperty)?)?;
    awesome_table.set("xkb_get_group_names", lua.create_function(xkb_get_group_names)?)?;
    awesome_table.set("set_xproperty", lua.create_function(set_xproperty)?)?;
    awesome_table.set("get_xproperty", lua.create_function(get_xproperty)?)?;
    awesome_table.set("systray", lua.create_function(systray)?)?;
    awesome_table.set("restart", lua.create_function(restart)?)?;
    awesome_table.set("load_image", lua.create_function(load_image)?)?;
    awesome_table.set("sync", lua.create_function(sync)?)?;
    awesome_table.set("exec", lua.create_function(exec)?)?;
    awesome_table.set("kill", lua.create_function(kill)?)?;
    awesome_table.set("quit", lua.create_function(quit)?)
}

fn property_setup<'lua>(lua: &'lua Lua, awesome_table: &Table<'lua>) -> rlua::Result<()> {
    // TODO Do properly
    awesome_table.set("version", "0".to_lua(lua)?)?;
    awesome_table.set("themes_path", "/usr/share/awesome/themes".to_lua(lua)?)?;
    awesome_table.set("conffile", "".to_lua(lua)?)
}

/// Registers a new X property
fn register_xproperty<'lua>(lua: &'lua Lua, (name_rust, v_type): (String, String)) -> rlua::Result<()> {
    let name = CString::new(name_rust.clone()).expect("XProperty was not CString");
    let arg_type = XPropertyType::from_string(v_type.clone())
        .ok_or(rlua::Error::RuntimeError(format!("{} not a valid xproperty", v_type)))?;
    unsafe {
        let raw_con = lua.named_registry_value::<LightUserData>(XCB_CONNECTION_HANDLE)?.0 as _;
        let atom_c = xproto::xcb_intern_atom_unchecked(raw_con,
                                                       false as u8,
                                                       name.to_bytes().len() as u16,
                                                       name.as_ptr());
        let atom_r = xproto::xcb_intern_atom_reply(raw_con, atom_c, ptr::null_mut());
        if atom_r.is_null() {
            return Ok(())
        }
        let new_property = XProperty::new(name_rust.clone(), arg_type, (*atom_r).atom);
        let mut properties = PROPERTIES.lock()
            .expect("Could not lock properties list");
        if let Some(found) = properties.iter().find(|&property| property == &new_property) {
            if found.type_ != new_property.type_ {
                return Err(rlua::Error::RuntimeError(
                    format!("property '{}' already registered with different type",
                            name_rust)));
            }
            return Ok(())
        }
        properties.push(new_property);
    }
    Ok(())
}

/// Get layout short names
fn xkb_get_group_names<'lua>(lua: &'lua Lua, _: ()) -> rlua::Result<Value<'lua>> {
    let xcb_con = lua.named_registry_value::<LightUserData>(XCB_CONNECTION_HANDLE)?.0;
    unsafe {
        let con = Connection::from_raw_conn(xcb_con as _);
        let raw_con = con.get_raw_conn();
        let id = xkb::ID_USE_CORE_KBD as _;
        // The structure here looks weird because we need to ensure Connection
        // isn't cleaned up even in the event of an error...
        let names_r = {
            let names_cookie = xkb::get_names_unchecked(&con, id, xkb::NAME_DETAIL_SYMBOLS);
            names_cookie.get_reply()
        };
        mem::forget(con);
        let names_r = match names_r {
            Ok(names_r) => {
                if names_r.ptr.is_null() {
                    warn!("Failed to get xkb symbols name");
                    return Ok(Value::Nil)
                }
                names_r
            },
            Err(err) => {
                warn!("Failed to get xkb symbols name {:?}", err);
                return Ok(Value::Nil)
            }
        };
        let buffer = ffi::xkb::xcb_xkb_get_names_value_list(names_r.ptr);
        if buffer.is_null() {
            warn!("Returned buffer was NULL");
            return Ok(Value::Nil)
        }
        let names_r_ptr = names_r.ptr;
        if names_r_ptr.is_null() {
            warn!("Name reply pointer was NULL");
            return Ok(Value::Nil)
        }
        let mut names_list: ffi::xkb::xcb_xkb_get_names_value_list_t = mem::uninitialized();
        xcb_xkb_get_names_value_list_unpack(buffer, (*names_r_ptr).nTypes, (*names_r_ptr).indicators, (*names_r_ptr).virtualMods,
                                            (*names_r_ptr).groupNames, (*names_r_ptr).nKeys, (*names_r_ptr).nKeyAliases,
                                            (*names_r_ptr).nRadioGroups, (*names_r_ptr).which, &mut names_list);
        let atom_name_c = ffi::xproto::xcb_get_atom_name_unchecked(raw_con, names_list.symbolsName);
        let atom_name_r = ffi::xproto::xcb_get_atom_name_reply(raw_con, atom_name_c, ptr::null_mut());
        if atom_name_r.is_null() {
            warn!("Failed to get atom symbols name");
            return Ok(Value::Nil)
        }
        let name_c = ffi::xproto::xcb_get_atom_name_name(atom_name_r);
        CStr::from_ptr(name_c).to_string_lossy().into_owned().to_lua(lua)
    }
}

/// Query & set information about the systray
fn systray<'lua>(_: &'lua Lua, _: ()) -> rlua::Result<(u32, Value)> {
    Ok((0, Value::Nil))
}

/// Restart Awesome by restarting the Lua thread
fn restart<'lua>(_: &'lua Lua, _: ()) -> rlua::Result<()> {
    use lua::{self, LuaQuery};
    if let Err(err) = lua::send(LuaQuery::Restart) {
        warn!("Could not restart Lua thread {:#?}", err);
    }
    Ok(())
}

/// Load an image from the given path
/// Returns either a cairo surface as light user data, nil and an error message
fn load_image<'lua>(lua: &'lua Lua, file_path: String) -> rlua::Result<Value<'lua>> {
    let pixbuf = Pixbuf::new_from_file(file_path.as_str())
        .map_err(|err| rlua::Error::RuntimeError(format!("{}", err)))?;
    let surface = render::load_surface_from_pixbuf(pixbuf);
    // UGH, I wanted to do to_glib_full, but that isn't defined apparently
    // So now I have to ignore the lifetime completely and just forget about the surface.
    let surface_ptr = surface.to_glib_none().0;
    ::std::mem::forget(surface);
    rlua::LightUserData(surface_ptr as _).to_lua(lua)
}

fn exec(_: &Lua, command: String) -> rlua::Result<()> {
    trace!("exec: \"{}\"", command);
    thread::Builder::new().name(command.clone()).spawn(|| {
        Command::new(command)
            .stdout(Stdio::null())
            .spawn()
            .expect("Could not spawn command")
    }).expect("Unable to spawn thread");
    Ok(())
}

/// Kills a PID with the given signal
///
/// Returns false if it could not send the signal to that process
fn kill(_: &Lua, (pid, sig): (libc::pid_t, libc::c_int)) -> rlua::Result<bool> {
    Ok(nix::sys::signal::kill(pid, sig).is_ok())
}

fn set_preferred_icon_size(lua: &Lua, val: u32) -> rlua::Result<()> {
    let mut awesome_state = lua.globals().get::<_, AwesomeState>("awesome")?;
    awesome_state.preferred_icon_size = val;
    lua.globals().set("awesome", awesome_state.to_lua(lua)?)

}

fn quit(_: &Lua, _: ()) -> rlua::Result<()> {
    ::rustwlc::terminate();
    Ok(())
}

/// No need to sync in Wayland
fn sync(_: &Lua, _: ()) -> rlua::Result<()> { Ok(()) }

fn set_xproperty(_: &Lua, _: Value) -> rlua::Result<()> {
    warn!("set_xproperty not supported");
    Ok(())
}

fn get_xproperty(_: &Lua, _: Value) -> rlua::Result<()> {
    warn!("get_xproperty not supported");
    Ok(())
}

fn xkb_set_layout_group(_: &Lua, _group: i32) -> rlua::Result<()> {
    warn!("xkb_set_layout_group not supported; Wait until wlroots");
    Ok(())
}

fn xkb_get_layout_group<'lua>(lua: &'lua Lua, _: ()) -> rlua::Result<Value<'lua>> {
    use xcb::ffi::xkb;
    unsafe {
        let raw_con = lua.named_registry_value::<LightUserData>(XCB_CONNECTION_HANDLE)?.0 as _;
        let state_c = xkb::xcb_xkb_get_state_unchecked(raw_con, xkb::XCB_XKB_ID_USE_CORE_KBD as _);
        let state_r = xkb::xcb_xkb_get_state_reply(raw_con, state_c, ptr::null_mut());
        if state_r.is_null() {
            warn!("State reply was NULL");
            return Ok(Value::Nil);
        }
        (*state_r).group.to_lua(lua)
    }
}

fn xrdb_get_value(_lua: &Lua, (_resource_class, _resource_name): (String, String))
                  -> rlua::Result<Value> {
    warn!("xrdb_get_value not supported");
    Ok(Value::Nil)
}
<commit_msg>awesome.set_preferred_icon_size: Modify state in place<commit_after>//! TODO Fill in

use std::{ptr, mem};
use xcb::{xkb, Connection};
use xcb::ffi::{self, xproto};
use nix::{self, libc};
use render;
use std::ffi::{CStr, CString};
use gdk_pixbuf::Pixbuf;
use glib::translate::ToGlibPtr;
use std::fmt::{self, Display, Formatter};
use std::process::{Command, Stdio};
use std::thread;
use std::default::Default;
use rlua::{self, Table, Lua, UserData, ToLua, Value, LightUserData, UserDataMethods, MetaMethod,
           AnyUserData};
use super::{XCB_CONNECTION_HANDLE, signal};
use super::xproperty::{XProperty, XPropertyType, PROPERTIES};

// TODO FIXME
// Often we are getting some raw pointers from the xcb replies
// we need to free them because the memory management for them is manual.

// TODO this isn't defined in the xcb crate, even though it should be.
// A patch should be open adding this to its generation scheme
extern "C" {
    fn xcb_xkb_get_names_value_list_unpack(buffer: *mut libc::c_void,
                                           nTypes: u8,
                                           indicators: u32,
                                           virtualMods: u16,
                                           groupNames: u8,
                                           nKeys: u8,
                                           nKeyAliases: u8,
                                           nRadioGroups: u8,
                                           which: u32,
                                           aux: *mut ffi::xkb::xcb_xkb_get_names_value_list_t) -> libc::c_int;
}

#[derive(Clone, Debug)]
pub struct AwesomeState {
    preferred_icon_size: u32
}

impl Default for AwesomeState {
    fn default() -> Self {
        AwesomeState {
            preferred_icon_size: 0
        }
    }
}

impl Display for AwesomeState {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        write!(f, "Awesome: {:p}", self)
    }
}

impl UserData for AwesomeState {
    fn add_methods(methods: &mut UserDataMethods<Self>) {
        fn index<'lua>(_: &'lua Lua,
                      (awesome, index): (AnyUserData<'lua>, Value<'lua>))
                      -> rlua::Result<rlua::Value<'lua>>
        {
            let table = awesome.get_user_value::<Table>()?;
            table.get::<_, Value>(index)
        };
        methods.add_meta_function(MetaMethod::Index, index);
    }
}

pub fn init(lua: &Lua) -> rlua::Result<()> {
    let awesome_table = lua.create_table()?;
    method_setup(lua, &awesome_table)?;
    property_setup(lua, &awesome_table)?;
    let globals = lua.globals();
    let awesome = lua.create_userdata(AwesomeState::default())?;
    awesome.set_user_value(awesome_table)?;
    globals.set("awesome", awesome)
}

fn method_setup<'lua>(lua: &'lua Lua, awesome_table: &Table<'lua>) -> rlua::Result<()> {
    // TODO Fill in rest
    awesome_table.set("connect_signal", lua.create_function(signal::global_connect_signal)?)?;
    awesome_table.set("disconnect_signal", lua.create_function(signal::global_disconnect_signal)?)?;
    awesome_table.set("emit_signal", lua.create_function(signal::global_emit_signal)?)?;
    awesome_table.set("xrdb_get_value", lua.create_function(xrdb_get_value)?)?;
    awesome_table.set("xkb_set_layout_group", lua.create_function(xkb_set_layout_group)?)?;
    awesome_table.set("xkb_get_layout_group", lua.create_function(xkb_get_layout_group)?)?;
    awesome_table.set("set_preferred_icon_size", lua.create_function(set_preferred_icon_size)?)?;
    awesome_table.set("register_xproperty", lua.create_function(register_xproperty)?)?;
    awesome_table.set("xkb_get_group_names", lua.create_function(xkb_get_group_names)?)?;
    awesome_table.set("set_xproperty", lua.create_function(set_xproperty)?)?;
    awesome_table.set("get_xproperty", lua.create_function(get_xproperty)?)?;
    awesome_table.set("systray", lua.create_function(systray)?)?;
    awesome_table.set("restart", lua.create_function(restart)?)?;
    awesome_table.set("load_image", lua.create_function(load_image)?)?;
    awesome_table.set("sync", lua.create_function(sync)?)?;
    awesome_table.set("exec", lua.create_function(exec)?)?;
    awesome_table.set("kill", lua.create_function(kill)?)?;
    awesome_table.set("quit", lua.create_function(quit)?)
}

fn property_setup<'lua>(lua: &'lua Lua, awesome_table: &Table<'lua>) -> rlua::Result<()> {
    // TODO Do properly
    awesome_table.set("version", "0".to_lua(lua)?)?;
    awesome_table.set("themes_path", "/usr/share/awesome/themes".to_lua(lua)?)?;
    awesome_table.set("conffile", "".to_lua(lua)?)
}

/// Registers a new X property
fn register_xproperty<'lua>(lua: &'lua Lua, (name_rust, v_type): (String, String)) -> rlua::Result<()> {
    let name = CString::new(name_rust.clone()).expect("XProperty was not CString");
    let arg_type = XPropertyType::from_string(v_type.clone())
        .ok_or(rlua::Error::RuntimeError(format!("{} not a valid xproperty", v_type)))?;
    unsafe {
        let raw_con = lua.named_registry_value::<LightUserData>(XCB_CONNECTION_HANDLE)?.0 as _;
        let atom_c = xproto::xcb_intern_atom_unchecked(raw_con,
                                                       false as u8,
                                                       name.to_bytes().len() as u16,
                                                       name.as_ptr());
        let atom_r = xproto::xcb_intern_atom_reply(raw_con, atom_c, ptr::null_mut());
        if atom_r.is_null() {
            return Ok(())
        }
        let new_property = XProperty::new(name_rust.clone(), arg_type, (*atom_r).atom);
        let mut properties = PROPERTIES.lock()
            .expect("Could not lock properties list");
        if let Some(found) = properties.iter().find(|&property| property == &new_property) {
            if found.type_ != new_property.type_ {
                return Err(rlua::Error::RuntimeError(
                    format!("property '{}' already registered with different type",
                            name_rust)));
            }
            return Ok(())
        }
        properties.push(new_property);
    }
    Ok(())
}

/// Get layout short names
fn xkb_get_group_names<'lua>(lua: &'lua Lua, _: ()) -> rlua::Result<Value<'lua>> {
    let xcb_con = lua.named_registry_value::<LightUserData>(XCB_CONNECTION_HANDLE)?.0;
    unsafe {
        let con = Connection::from_raw_conn(xcb_con as _);
        let raw_con = con.get_raw_conn();
        let id = xkb::ID_USE_CORE_KBD as _;
        // The structure here looks weird because we need to ensure Connection
        // isn't cleaned up even in the event of an error...
        let names_r = {
            let names_cookie = xkb::get_names_unchecked(&con, id, xkb::NAME_DETAIL_SYMBOLS);
            names_cookie.get_reply()
        };
        mem::forget(con);
        let names_r = match names_r {
            Ok(names_r) => {
                if names_r.ptr.is_null() {
                    warn!("Failed to get xkb symbols name");
                    return Ok(Value::Nil)
                }
                names_r
            },
            Err(err) => {
                warn!("Failed to get xkb symbols name {:?}", err);
                return Ok(Value::Nil)
            }
        };
        let buffer = ffi::xkb::xcb_xkb_get_names_value_list(names_r.ptr);
        if buffer.is_null() {
            warn!("Returned buffer was NULL");
            return Ok(Value::Nil)
        }
        let names_r_ptr = names_r.ptr;
        if names_r_ptr.is_null() {
            warn!("Name reply pointer was NULL");
            return Ok(Value::Nil)
        }
        let mut names_list: ffi::xkb::xcb_xkb_get_names_value_list_t = mem::uninitialized();
        xcb_xkb_get_names_value_list_unpack(buffer, (*names_r_ptr).nTypes, (*names_r_ptr).indicators, (*names_r_ptr).virtualMods,
                                            (*names_r_ptr).groupNames, (*names_r_ptr).nKeys, (*names_r_ptr).nKeyAliases,
                                            (*names_r_ptr).nRadioGroups, (*names_r_ptr).which, &mut names_list);
        let atom_name_c = ffi::xproto::xcb_get_atom_name_unchecked(raw_con, names_list.symbolsName);
        let atom_name_r = ffi::xproto::xcb_get_atom_name_reply(raw_con, atom_name_c, ptr::null_mut());
        if atom_name_r.is_null() {
            warn!("Failed to get atom symbols name");
            return Ok(Value::Nil)
        }
        let name_c = ffi::xproto::xcb_get_atom_name_name(atom_name_r);
        CStr::from_ptr(name_c).to_string_lossy().into_owned().to_lua(lua)
    }
}

/// Query & set information about the systray
fn systray<'lua>(_: &'lua Lua, _: ()) -> rlua::Result<(u32, Value)> {
    Ok((0, Value::Nil))
}

/// Restart Awesome by restarting the Lua thread
fn restart<'lua>(_: &'lua Lua, _: ()) -> rlua::Result<()> {
    use lua::{self, LuaQuery};
    if let Err(err) = lua::send(LuaQuery::Restart) {
        warn!("Could not restart Lua thread {:#?}", err);
    }
    Ok(())
}

/// Load an image from the given path
/// Returns either a cairo surface as light user data, nil and an error message
fn load_image<'lua>(lua: &'lua Lua, file_path: String) -> rlua::Result<Value<'lua>> {
    let pixbuf = Pixbuf::new_from_file(file_path.as_str())
        .map_err(|err| rlua::Error::RuntimeError(format!("{}", err)))?;
    let surface = render::load_surface_from_pixbuf(pixbuf);
    // UGH, I wanted to do to_glib_full, but that isn't defined apparently
    // So now I have to ignore the lifetime completely and just forget about the surface.
    let surface_ptr = surface.to_glib_none().0;
    ::std::mem::forget(surface);
    rlua::LightUserData(surface_ptr as _).to_lua(lua)
}

fn exec(_: &Lua, command: String) -> rlua::Result<()> {
    trace!("exec: \"{}\"", command);
    thread::Builder::new().name(command.clone()).spawn(|| {
        Command::new(command)
            .stdout(Stdio::null())
            .spawn()
            .expect("Could not spawn command")
    }).expect("Unable to spawn thread");
    Ok(())
}

/// Kills a PID with the given signal
///
/// Returns false if it could not send the signal to that process
fn kill(_: &Lua, (pid, sig): (libc::pid_t, libc::c_int)) -> rlua::Result<bool> {
    Ok(nix::sys::signal::kill(pid, sig).is_ok())
}

fn set_preferred_icon_size(lua: &Lua, val: u32) -> rlua::Result<()> {
    let awesome_state = lua.globals().get::<_, AnyUserData>("awesome")?;
    let mut awesome_state = awesome_state.borrow_mut::<AwesomeState>()?;
    awesome_state.preferred_icon_size = val;
    Ok(())
}

fn quit(_: &Lua, _: ()) -> rlua::Result<()> {
    ::rustwlc::terminate();
    Ok(())
}

/// No need to sync in Wayland
fn sync(_: &Lua, _: ()) -> rlua::Result<()> { Ok(()) }

fn set_xproperty(_: &Lua, _: Value) -> rlua::Result<()> {
    warn!("set_xproperty not supported");
    Ok(())
}

fn get_xproperty(_: &Lua, _: Value) -> rlua::Result<()> {
    warn!("get_xproperty not supported");
    Ok(())
}

fn xkb_set_layout_group(_: &Lua, _group: i32) -> rlua::Result<()> {
    warn!("xkb_set_layout_group not supported; Wait until wlroots");
    Ok(())
}

fn xkb_get_layout_group<'lua>(lua: &'lua Lua, _: ()) -> rlua::Result<Value<'lua>> {
    use xcb::ffi::xkb;
    unsafe {
        let raw_con = lua.named_registry_value::<LightUserData>(XCB_CONNECTION_HANDLE)?.0 as _;
        let state_c = xkb::xcb_xkb_get_state_unchecked(raw_con, xkb::XCB_XKB_ID_USE_CORE_KBD as _);
        let state_r = xkb::xcb_xkb_get_state_reply(raw_con, state_c, ptr::null_mut());
        if state_r.is_null() {
            warn!("State reply was NULL");
            return Ok(Value::Nil);
        }
        (*state_r).group.to_lua(lua)
    }
}

fn xrdb_get_value(_lua: &Lua, (_resource_class, _resource_name): (String, String))
                  -> rlua::Result<Value> {
    warn!("xrdb_get_value not supported");
    Ok(Value::Nil)
}
<|endoftext|>


<commit_before>extern crate termion;

extern crate chan_signal;

use ::backend;
use ::event::{Event, Key};
use self::termion::color as tcolor;
use self::termion::event::Key as TKey;
use self::termion::input::TermRead;
use self::termion::raw::IntoRawMode;
use self::termion::style as tstyle;
use std::cell::Cell;
use std::collections::BTreeMap;
use std::fmt;
use std::io::Write;
use std::thread;
use std::time;
use chan;

use ::theme;

pub struct Concrete {
    terminal: termion::raw::RawTerminal<::std::io::Stdout>,
    current_style: Cell<theme::ColorStyle>,
    colors: BTreeMap<i16, (Box<tcolor::Color>, Box<tcolor::Color>)>,

    input: chan::Receiver<Event>,
    resize: chan::Receiver<chan_signal::Signal>,
    timeout: Option<u32>,
}

trait Effectable {
    fn on(&self);
    fn off(&self);
}

struct ColorRef<'a>(&'a tcolor::Color);

impl<'a> tcolor::Color for ColorRef<'a> {
    #[inline]
    fn write_fg(&self, f: &mut fmt::Formatter) -> fmt::Result {
        self.0.write_fg(f)
    }

    #[inline]
    fn write_bg(&self, f: &mut fmt::Formatter) -> fmt::Result {
        self.0.write_bg(f)
    }
}

impl Effectable for theme::Effect {
    fn on(&self) {
        match *self {
            theme::Effect::Simple => (),
            theme::Effect::Reverse => print!("{}", tstyle::Invert),
        }
    }

    fn off(&self) {
        match *self {
            theme::Effect::Simple => (),
            theme::Effect::Reverse => print!("{}", tstyle::NoInvert),
        }
    }
}

fn apply_colors(fg: &tcolor::Color, bg: &tcolor::Color) {
    print!("{}{}", tcolor::Fg(ColorRef(fg)), tcolor::Bg(ColorRef(bg)));
}

impl Concrete {
    fn apply_colorstyle(&self, color_style: theme::ColorStyle) {
        let (ref fg, ref bg) = self.colors[&color_style.id()];
        apply_colors(&**fg, &**bg);
    }
}

impl backend::Backend for Concrete {
    fn init() -> Self {
        print!("{}", termion::cursor::Hide);

        let resize = chan_signal::notify(&[chan_signal::Signal::WINCH]);

        let terminal = ::std::io::stdout().into_raw_mode().unwrap();
        let (sender, receiver) = chan::async();

        // TODO: use signal_chan crate

        thread::spawn(move || {
            for key in ::std::io::stdin().keys() {
                if let Ok(key) = key {
                    sender.send(map_key(key))
                }
            }
        });

        let backend = Concrete {
            terminal: terminal,
            current_style: Cell::new(theme::ColorStyle::Background),
            colors: BTreeMap::new(),
            input: receiver,
            resize: resize,
            timeout: None,
        };

        backend
    }

    fn finish(&mut self) {
        print!("{}{}", termion::cursor::Show, termion::cursor::Goto(1, 1));
        self.clear();
    }

    fn init_color_style(&mut self, style: theme::ColorStyle,
                        foreground: &theme::Color, background: &theme::Color) {
        // Step 1: convert foreground and background into proper termion Color
        self.colors.insert(style.id(),
                           (colour_to_termion_colour(foreground),
                            colour_to_termion_colour(background)));
    }

    fn with_color<F: FnOnce()>(&self, color: theme::ColorStyle, f: F) {
        // TODO: actually use colors
        // TODO: careful! need to remember the previous state
        //                and apply it back

        let current_style = self.current_style.get();

        self.apply_colorstyle(color);

        self.current_style.set(color);
        f();
        self.current_style.set(current_style);

        self.apply_colorstyle(current_style);
    }

    fn with_effect<F: FnOnce()>(&self, effect: theme::Effect, f: F) {
        // TODO: actually use effects
        // TODO: careful! need to remember the previous state
        //                and apply it back
        effect.on();
        f();
        effect.off();
    }

    fn has_colors(&self) -> bool {
        // TODO: color support detection?
        true
    }

    fn screen_size(&self) -> (usize, usize) {
        let (x, y) = termion::terminal_size().unwrap_or((1, 1));
        (x as usize, y as usize)
    }

    fn clear(&self) {
        self.apply_colorstyle(theme::ColorStyle::Background);
        print!("{}", termion::clear::All);
    }

    fn refresh(&mut self) {
        self.terminal.flush().unwrap();
    }

    fn print_at(&self, (x, y): (usize, usize), text: &str) {
        // TODO: terminals are 1-based. Should we add 1 here?
        print!("{}{}", termion::cursor::Goto(x as u16, y as u16), text);
    }

    fn set_refresh_rate(&mut self, fps: u32) {
        // TODO: handle async refresh, when no input is entered.
        // Could be done with a timeout on the event polling,
        // if it was supportedd.
        self.timeout = Some(1000 / fps as u32);
    }

    fn poll_event(&self) -> Event {
        // TODO: select! on the input and SIGWINCH signal channel.
        // TODO: also handle timeout... recv_timeout?
        let input = &self.input;
        let resize = &self.resize;

        if let Some(timeout) = self.timeout {
            let timeout = chan::after_ms(timeout);
            chan_select!{
                timeout.recv() => return Event::Refresh,
                resize.recv() => return Event::WindowResize,
                input.recv() -> input => return input.unwrap(),
            }
        } else {
            chan_select!{
                resize.recv() => return Event::WindowResize,
                input.recv() -> input => return input.unwrap(),
            }
        }
    }
}

fn map_key(key: TKey) -> Event {
    match key {
        TKey::Esc => Event::Key(Key::Esc),
        TKey::Backspace => Event::Key(Key::Backspace),
        TKey::Left => Event::Key(Key::Left),
        TKey::Right => Event::Key(Key::Right),
        TKey::Up => Event::Key(Key::Up),
        TKey::Down => Event::Key(Key::Down),
        TKey::Home => Event::Key(Key::Home),
        TKey::End => Event::Key(Key::End),
        TKey::PageUp => Event::Key(Key::PageUp),
        TKey::PageDown => Event::Key(Key::PageDown),
        TKey::Delete => Event::Key(Key::Del),
        TKey::Insert => Event::Key(Key::Ins),
        TKey::F(i) if i < 12 => Event::Key(Key::from_f(i)),
        TKey::F(j) => Event::Unknown(-(j as i32)),
        TKey::Char('\n') => Event::Key(Key::Enter),
        TKey::Char('\t') => Event::Key(Key::Tab),
        TKey::Char(c) => Event::Char(c),
        TKey::Ctrl('c') => Event::Exit,
        TKey::Ctrl(c) => Event::CtrlChar(c),
        TKey::Alt(c) => Event::AltChar(c),
        _ => Event::Unknown(-1),
    }

}

fn colour_to_termion_colour(clr: &theme::Color) -> Box<tcolor::Color> {
    match *clr {
        theme::Color::Dark(theme::BaseColor::Black) => Box::new(tcolor::Black),
        theme::Color::Dark(theme::BaseColor::Red) => Box::new(tcolor::Red),
        theme::Color::Dark(theme::BaseColor::Green) => Box::new(tcolor::Green),
        theme::Color::Dark(theme::BaseColor::Yellow) => {
            Box::new(tcolor::Yellow)
        }
        theme::Color::Dark(theme::BaseColor::Blue) => Box::new(tcolor::Blue),
        theme::Color::Dark(theme::BaseColor::Magenta) => {
            Box::new(tcolor::Magenta)
        }
        theme::Color::Dark(theme::BaseColor::Cyan) => Box::new(tcolor::Cyan),
        theme::Color::Dark(theme::BaseColor::White) => Box::new(tcolor::White),

        theme::Color::Light(theme::BaseColor::Black) => {
            Box::new(tcolor::LightBlack)
        }
        theme::Color::Light(theme::BaseColor::Red) => {
            Box::new(tcolor::LightRed)
        }
        theme::Color::Light(theme::BaseColor::Green) => {
            Box::new(tcolor::LightGreen)
        }
        theme::Color::Light(theme::BaseColor::Yellow) => {
            Box::new(tcolor::LightYellow)
        }
        theme::Color::Light(theme::BaseColor::Blue) => {
            Box::new(tcolor::LightBlue)
        }
        theme::Color::Light(theme::BaseColor::Magenta) => {
            Box::new(tcolor::LightMagenta)
        }
        theme::Color::Light(theme::BaseColor::Cyan) => {
            Box::new(tcolor::LightCyan)
        }
        theme::Color::Light(theme::BaseColor::White) => {
            Box::new(tcolor::LightWhite)
        }

        theme::Color::Rgb(r, g, b) => Box::new(tcolor::Rgb(r, g, b)),
        theme::Color::RgbLowRes(r, g, b) => {
            Box::new(tcolor::AnsiValue::rgb(r, g, b))
        }
    }
}
<commit_msg>Fixed issue with resetting background color for Termion backend  - Background color was being set back to the theme's background color rather than the terminal's background color when the backend was `finish()`ed.  This resets the background color to the default instead.<commit_after>extern crate termion;

extern crate chan_signal;

use ::backend;
use ::event::{Event, Key};
use self::termion::color as tcolor;
use self::termion::event::Key as TKey;
use self::termion::input::TermRead;
use self::termion::raw::IntoRawMode;
use self::termion::style as tstyle;
use std::cell::Cell;
use std::collections::BTreeMap;
use std::fmt;
use std::io::Write;
use std::thread;
use std::time;
use chan;

use ::theme;

pub struct Concrete {
    terminal: termion::raw::RawTerminal<::std::io::Stdout>,
    current_style: Cell<theme::ColorStyle>,
    colors: BTreeMap<i16, (Box<tcolor::Color>, Box<tcolor::Color>)>,

    input: chan::Receiver<Event>,
    resize: chan::Receiver<chan_signal::Signal>,
    timeout: Option<u32>,
}

trait Effectable {
    fn on(&self);
    fn off(&self);
}

struct ColorRef<'a>(&'a tcolor::Color);

impl<'a> tcolor::Color for ColorRef<'a> {
    #[inline]
    fn write_fg(&self, f: &mut fmt::Formatter) -> fmt::Result {
        self.0.write_fg(f)
    }

    #[inline]
    fn write_bg(&self, f: &mut fmt::Formatter) -> fmt::Result {
        self.0.write_bg(f)
    }
}

impl Effectable for theme::Effect {
    fn on(&self) {
        match *self {
            theme::Effect::Simple => (),
            theme::Effect::Reverse => print!("{}", tstyle::Invert),
        }
    }

    fn off(&self) {
        match *self {
            theme::Effect::Simple => (),
            theme::Effect::Reverse => print!("{}", tstyle::NoInvert),
        }
    }
}

fn apply_colors(fg: &tcolor::Color, bg: &tcolor::Color) {
    print!("{}{}", tcolor::Fg(ColorRef(fg)), tcolor::Bg(ColorRef(bg)));
}

impl Concrete {
    fn apply_colorstyle(&self, color_style: theme::ColorStyle) {
        let (ref fg, ref bg) = self.colors[&color_style.id()];
        apply_colors(&**fg, &**bg);
    }
}

impl backend::Backend for Concrete {
    fn init() -> Self {
        print!("{}", termion::cursor::Hide);

        let resize = chan_signal::notify(&[chan_signal::Signal::WINCH]);

        let terminal = ::std::io::stdout().into_raw_mode().unwrap();
        let (sender, receiver) = chan::async();

        // TODO: use signal_chan crate

        thread::spawn(move || {
            for key in ::std::io::stdin().keys() {
                if let Ok(key) = key {
                    sender.send(map_key(key))
                }
            }
        });

        let backend = Concrete {
            terminal: terminal,
            current_style: Cell::new(theme::ColorStyle::Background),
            colors: BTreeMap::new(),
            input: receiver,
            resize: resize,
            timeout: None,
        };

        backend
    }

    fn finish(&mut self) {
        print!("{}{}", termion::cursor::Show, termion::cursor::Goto(1, 1));
        print!("{}[49m{}[39m{}", 27 as char, 27 as char, termion::clear::All);
    }

    fn init_color_style(&mut self, style: theme::ColorStyle,
                        foreground: &theme::Color, background: &theme::Color) {
        // Step 1: convert foreground and background into proper termion Color
        self.colors.insert(style.id(),
                           (colour_to_termion_colour(foreground),
                            colour_to_termion_colour(background)));
    }

    fn with_color<F: FnOnce()>(&self, color: theme::ColorStyle, f: F) {
        // TODO: actually use colors
        // TODO: careful! need to remember the previous state
        //                and apply it back

        let current_style = self.current_style.get();

        self.apply_colorstyle(color);

        self.current_style.set(color);
        f();
        self.current_style.set(current_style);

        self.apply_colorstyle(current_style);
    }

    fn with_effect<F: FnOnce()>(&self, effect: theme::Effect, f: F) {
        // TODO: actually use effects
        // TODO: careful! need to remember the previous state
        //                and apply it back
        effect.on();
        f();
        effect.off();
    }

    fn has_colors(&self) -> bool {
        // TODO: color support detection?
        true
    }

    fn screen_size(&self) -> (usize, usize) {
        let (x, y) = termion::terminal_size().unwrap_or((1, 1));
        (x as usize, y as usize)
    }

    fn clear(&self) {
        self.apply_colorstyle(theme::ColorStyle::Background);
        print!("{}", termion::clear::All);
    }

    fn refresh(&mut self) {
        self.terminal.flush().unwrap();
    }

    fn print_at(&self, (x, y): (usize, usize), text: &str) {
        // TODO: terminals are 1-based. Should we add 1 here?
        print!("{}{}", termion::cursor::Goto(x as u16, y as u16), text);
    }

    fn set_refresh_rate(&mut self, fps: u32) {
        // TODO: handle async refresh, when no input is entered.
        // Could be done with a timeout on the event polling,
        // if it was supportedd.
        self.timeout = Some(1000 / fps as u32);
    }

    fn poll_event(&self) -> Event {
        // TODO: select! on the input and SIGWINCH signal channel.
        // TODO: also handle timeout... recv_timeout?
        let input = &self.input;
        let resize = &self.resize;

        if let Some(timeout) = self.timeout {
            let timeout = chan::after_ms(timeout);
            chan_select!{
                timeout.recv() => return Event::Refresh,
                resize.recv() => return Event::WindowResize,
                input.recv() -> input => return input.unwrap(),
            }
        } else {
            chan_select!{
                resize.recv() => return Event::WindowResize,
                input.recv() -> input => return input.unwrap(),
            }
        }
    }
}

fn map_key(key: TKey) -> Event {
    match key {
        TKey::Esc => Event::Key(Key::Esc),
        TKey::Backspace => Event::Key(Key::Backspace),
        TKey::Left => Event::Key(Key::Left),
        TKey::Right => Event::Key(Key::Right),
        TKey::Up => Event::Key(Key::Up),
        TKey::Down => Event::Key(Key::Down),
        TKey::Home => Event::Key(Key::Home),
        TKey::End => Event::Key(Key::End),
        TKey::PageUp => Event::Key(Key::PageUp),
        TKey::PageDown => Event::Key(Key::PageDown),
        TKey::Delete => Event::Key(Key::Del),
        TKey::Insert => Event::Key(Key::Ins),
        TKey::F(i) if i < 12 => Event::Key(Key::from_f(i)),
        TKey::F(j) => Event::Unknown(-(j as i32)),
        TKey::Char('\n') => Event::Key(Key::Enter),
        TKey::Char('\t') => Event::Key(Key::Tab),
        TKey::Char(c) => Event::Char(c),
        TKey::Ctrl('c') => Event::Exit,
        TKey::Ctrl(c) => Event::CtrlChar(c),
        TKey::Alt(c) => Event::AltChar(c),
        _ => Event::Unknown(-1),
    }

}

fn colour_to_termion_colour(clr: &theme::Color) -> Box<tcolor::Color> {
    match *clr {
        theme::Color::Dark(theme::BaseColor::Black) => Box::new(tcolor::Black),
        theme::Color::Dark(theme::BaseColor::Red) => Box::new(tcolor::Red),
        theme::Color::Dark(theme::BaseColor::Green) => Box::new(tcolor::Green),
        theme::Color::Dark(theme::BaseColor::Yellow) => {
            Box::new(tcolor::Yellow)
        }
        theme::Color::Dark(theme::BaseColor::Blue) => Box::new(tcolor::Blue),
        theme::Color::Dark(theme::BaseColor::Magenta) => {
            Box::new(tcolor::Magenta)
        }
        theme::Color::Dark(theme::BaseColor::Cyan) => Box::new(tcolor::Cyan),
        theme::Color::Dark(theme::BaseColor::White) => Box::new(tcolor::White),

        theme::Color::Light(theme::BaseColor::Black) => {
            Box::new(tcolor::LightBlack)
        }
        theme::Color::Light(theme::BaseColor::Red) => {
            Box::new(tcolor::LightRed)
        }
        theme::Color::Light(theme::BaseColor::Green) => {
            Box::new(tcolor::LightGreen)
        }
        theme::Color::Light(theme::BaseColor::Yellow) => {
            Box::new(tcolor::LightYellow)
        }
        theme::Color::Light(theme::BaseColor::Blue) => {
            Box::new(tcolor::LightBlue)
        }
        theme::Color::Light(theme::BaseColor::Magenta) => {
            Box::new(tcolor::LightMagenta)
        }
        theme::Color::Light(theme::BaseColor::Cyan) => {
            Box::new(tcolor::LightCyan)
        }
        theme::Color::Light(theme::BaseColor::White) => {
            Box::new(tcolor::LightWhite)
        }

        theme::Color::Rgb(r, g, b) => Box::new(tcolor::Rgb(r, g, b)),
        theme::Color::RgbLowRes(r, g, b) => {
            Box::new(tcolor::AnsiValue::rgb(r, g, b))
        }
    }
}
<|endoftext|>


<commit_before>// Copyright 2018 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Implementation of compiling the compiler and standard library, in "check" mode.

use compile::{run_cargo, std_cargo, test_cargo, rustc_cargo, rustc_cargo_env, add_to_sysroot};
use builder::{RunConfig, Builder, ShouldRun, Step};
use tool::{prepare_tool_cargo, SourceType};
use {Compiler, Mode};
use cache::{INTERNER, Interned};
use std::path::PathBuf;

#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
pub struct Std {
    pub target: Interned<String>,
}

impl Step for Std {
    type Output = ();
    const DEFAULT: bool = true;

    fn should_run(run: ShouldRun) -> ShouldRun {
        run.all_krates("std")
    }

    fn make_run(run: RunConfig) {
        run.builder.ensure(Std {
            target: run.target,
        });
    }

    fn run(self, builder: &Builder) {
        let target = self.target;
        let compiler = builder.compiler(0, builder.config.build);

        let mut cargo = builder.cargo(compiler, Mode::Std, target, "check");
        std_cargo(builder, &compiler, target, &mut cargo);

        let _folder = builder.fold_output(|| format!("stage{}-std", compiler.stage));
        builder.info(&format!("Checking std artifacts ({} -> {})", &compiler.host, target));
        run_cargo(builder,
                  &mut cargo,
                  vec![],
                  &libstd_stamp(builder, compiler, target),
                  true);

        let libdir = builder.sysroot_libdir(compiler, target);
        add_to_sysroot(&builder, &libdir, &libstd_stamp(builder, compiler, target));
    }
}

#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
pub struct Rustc {
    pub target: Interned<String>,
}

impl Step for Rustc {
    type Output = ();
    const ONLY_HOSTS: bool = true;
    const DEFAULT: bool = true;

    fn should_run(run: ShouldRun) -> ShouldRun {
        run.all_krates("rustc-main")
    }

    fn make_run(run: RunConfig) {
        run.builder.ensure(Rustc {
            target: run.target,
        });
    }

    /// Build the compiler.
    ///
    /// This will build the compiler for a particular stage of the build using
    /// the `compiler` targeting the `target` architecture. The artifacts
    /// created will also be linked into the sysroot directory.
    fn run(self, builder: &Builder) {
        let compiler = builder.compiler(0, builder.config.build);
        let target = self.target;

        let mut cargo = builder.cargo(compiler, Mode::Rustc, target, "check");
        rustc_cargo(builder, &mut cargo);

        let _folder = builder.fold_output(|| format!("stage{}-rustc", compiler.stage));
        builder.info(&format!("Checking compiler artifacts ({} -> {})", &compiler.host, target));
        run_cargo(builder,
                  &mut cargo,
                  vec![],
                  &librustc_stamp(builder, compiler, target),
                  true);

        let libdir = builder.sysroot_libdir(compiler, target);
        add_to_sysroot(&builder, &libdir, &librustc_stamp(builder, compiler, target));
    }
}

#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
pub struct CodegenBackend {
    pub target: Interned<String>,
    pub backend: Interned<String>,
}

impl Step for CodegenBackend {
    type Output = ();
    const ONLY_HOSTS: bool = true;
    const DEFAULT: bool = true;

    fn should_run(run: ShouldRun) -> ShouldRun {
        run.all_krates("rustc_codegen_llvm")
    }

    fn make_run(run: RunConfig) {
        let backend = run.builder.config.rust_codegen_backends.get(0);
        let backend = backend.cloned().unwrap_or_else(|| {
            INTERNER.intern_str("llvm")
        });
        run.builder.ensure(CodegenBackend {
            target: run.target,
            backend,
        });
    }

    fn run(self, builder: &Builder) {
        let compiler = builder.compiler(0, builder.config.build);
        let target = self.target;
        let backend = self.backend;

        let out_dir = builder.cargo_out(compiler, Mode::Codegen, target);
        builder.clear_if_dirty(&out_dir, &librustc_stamp(builder, compiler, target));

        let mut cargo = builder.cargo(compiler, Mode::Codegen, target, "check");
        cargo.arg("--manifest-path").arg(builder.src.join("src/librustc_codegen_llvm/Cargo.toml"));
        rustc_cargo_env(builder, &mut cargo);

        // We won't build LLVM if it's not available, as it shouldn't affect `check`.

        let _folder = builder.fold_output(|| format!("stage{}-rustc_codegen_llvm", compiler.stage));
        run_cargo(builder,
                  &mut cargo,
                  vec![],
                  &codegen_backend_stamp(builder, compiler, target, backend),
                  true);
    }
}

#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
pub struct Test {
    pub target: Interned<String>,
}

impl Step for Test {
    type Output = ();
    const DEFAULT: bool = true;

    fn should_run(run: ShouldRun) -> ShouldRun {
        run.all_krates("test")
    }

    fn make_run(run: RunConfig) {
        run.builder.ensure(Test {
            target: run.target,
        });
    }

    fn run(self, builder: &Builder) {
        let compiler = builder.compiler(0, builder.config.build);
        let target = self.target;

        let mut cargo = builder.cargo(compiler, Mode::Test, target, "check");
        test_cargo(builder, &compiler, target, &mut cargo);

        let _folder = builder.fold_output(|| format!("stage{}-test", compiler.stage));
        builder.info(&format!("Checking test artifacts ({} -> {})", &compiler.host, target));
        run_cargo(builder,
                  &mut cargo,
                  vec![],
                  &libtest_stamp(builder, compiler, target),
                  true);

        let libdir = builder.sysroot_libdir(compiler, target);
        add_to_sysroot(builder, &libdir, &libtest_stamp(builder, compiler, target));
    }
}

#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
pub struct Rustdoc {
    pub target: Interned<String>,
}

impl Step for Rustdoc {
    type Output = ();
    const ONLY_HOSTS: bool = true;
    const DEFAULT: bool = true;

    fn should_run(run: ShouldRun) -> ShouldRun {
        run.path("src/tools/rustdoc")
    }

    fn make_run(run: RunConfig) {
        run.builder.ensure(Rustdoc {
            target: run.target,
        });
    }

    fn run(self, builder: &Builder) {
        let compiler = builder.compiler(0, builder.config.build);
        let target = self.target;

        let stage_out = builder.stage_out(compiler, Mode::ToolRustc);
        builder.clear_if_dirty(&stage_out, &libstd_stamp(builder, compiler, target));
        builder.clear_if_dirty(&stage_out, &libtest_stamp(builder, compiler, target));
        builder.clear_if_dirty(&stage_out, &librustc_stamp(builder, compiler, target));

        let mut cargo = prepare_tool_cargo(builder,
                                           compiler,
                                           Mode::ToolRustc,
                                           target,
                                           "check",
                                           "src/tools/rustdoc",
                                           SourceType::InTree);

        let _folder = builder.fold_output(|| format!("stage{}-rustdoc", compiler.stage));
        println!("Checking rustdoc artifacts ({} -> {})", &compiler.host, target);
        run_cargo(builder,
                  &mut cargo,
                  vec![],
                  &rustdoc_stamp(builder, compiler, target),
                  true);

        let libdir = builder.sysroot_libdir(compiler, target);
        add_to_sysroot(&builder, &libdir, &rustdoc_stamp(builder, compiler, target));
        builder.cargo(compiler, Mode::ToolRustc, target, "clean");
    }
}

/// Cargo's output path for the standard library in a given stage, compiled
/// by a particular compiler for the specified target.
pub fn libstd_stamp(builder: &Builder, compiler: Compiler, target: Interned<String>) -> PathBuf {
    builder.cargo_out(compiler, Mode::Std, target).join(".libstd-check.stamp")
}

/// Cargo's output path for libtest in a given stage, compiled by a particular
/// compiler for the specified target.
pub fn libtest_stamp(builder: &Builder, compiler: Compiler, target: Interned<String>) -> PathBuf {
    builder.cargo_out(compiler, Mode::Test, target).join(".libtest-check.stamp")
}

/// Cargo's output path for librustc in a given stage, compiled by a particular
/// compiler for the specified target.
pub fn librustc_stamp(builder: &Builder, compiler: Compiler, target: Interned<String>) -> PathBuf {
    builder.cargo_out(compiler, Mode::Rustc, target).join(".librustc-check.stamp")
}

/// Cargo's output path for librustc_codegen_llvm in a given stage, compiled by a particular
/// compiler for the specified target and backend.
fn codegen_backend_stamp(builder: &Builder,
                         compiler: Compiler,
                         target: Interned<String>,
                         backend: Interned<String>) -> PathBuf {
    builder.cargo_out(compiler, Mode::Codegen, target)
         .join(format!(".librustc_codegen_llvm-{}-check.stamp", backend))
}

/// Cargo's output path for rustdoc in a given stage, compiled by a particular
/// compiler for the specified target.
pub fn rustdoc_stamp(builder: &Builder, compiler: Compiler, target: Interned<String>) -> PathBuf {
    builder.cargo_out(compiler, Mode::ToolRustc, target)
        .join(".rustdoc-check.stamp")
}
<commit_msg>Add check build dependencies<commit_after>// Copyright 2018 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Implementation of compiling the compiler and standard library, in "check" mode.

use compile::{run_cargo, std_cargo, test_cargo, rustc_cargo, rustc_cargo_env, add_to_sysroot};
use builder::{RunConfig, Builder, ShouldRun, Step};
use tool::{prepare_tool_cargo, SourceType};
use {Compiler, Mode};
use cache::{INTERNER, Interned};
use std::path::PathBuf;

#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
pub struct Std {
    pub target: Interned<String>,
}

impl Step for Std {
    type Output = ();
    const DEFAULT: bool = true;

    fn should_run(run: ShouldRun) -> ShouldRun {
        run.all_krates("std")
    }

    fn make_run(run: RunConfig) {
        run.builder.ensure(Std {
            target: run.target,
        });
    }

    fn run(self, builder: &Builder) {
        let target = self.target;
        let compiler = builder.compiler(0, builder.config.build);

        let mut cargo = builder.cargo(compiler, Mode::Std, target, "check");
        std_cargo(builder, &compiler, target, &mut cargo);

        let _folder = builder.fold_output(|| format!("stage{}-std", compiler.stage));
        builder.info(&format!("Checking std artifacts ({} -> {})", &compiler.host, target));
        run_cargo(builder,
                  &mut cargo,
                  vec![],
                  &libstd_stamp(builder, compiler, target),
                  true);

        let libdir = builder.sysroot_libdir(compiler, target);
        add_to_sysroot(&builder, &libdir, &libstd_stamp(builder, compiler, target));
    }
}

#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
pub struct Rustc {
    pub target: Interned<String>,
}

impl Step for Rustc {
    type Output = ();
    const ONLY_HOSTS: bool = true;
    const DEFAULT: bool = true;

    fn should_run(run: ShouldRun) -> ShouldRun {
        run.all_krates("rustc-main")
    }

    fn make_run(run: RunConfig) {
        run.builder.ensure(Rustc {
            target: run.target,
        });
    }

    /// Build the compiler.
    ///
    /// This will build the compiler for a particular stage of the build using
    /// the `compiler` targeting the `target` architecture. The artifacts
    /// created will also be linked into the sysroot directory.
    fn run(self, builder: &Builder) {
        let compiler = builder.compiler(0, builder.config.build);
        let target = self.target;

        builder.ensure(Test { target });

        let mut cargo = builder.cargo(compiler, Mode::Rustc, target, "check");
        rustc_cargo(builder, &mut cargo);

        let _folder = builder.fold_output(|| format!("stage{}-rustc", compiler.stage));
        builder.info(&format!("Checking compiler artifacts ({} -> {})", &compiler.host, target));
        run_cargo(builder,
                  &mut cargo,
                  vec![],
                  &librustc_stamp(builder, compiler, target),
                  true);

        let libdir = builder.sysroot_libdir(compiler, target);
        add_to_sysroot(&builder, &libdir, &librustc_stamp(builder, compiler, target));
    }
}

#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
pub struct CodegenBackend {
    pub target: Interned<String>,
    pub backend: Interned<String>,
}

impl Step for CodegenBackend {
    type Output = ();
    const ONLY_HOSTS: bool = true;
    const DEFAULT: bool = true;

    fn should_run(run: ShouldRun) -> ShouldRun {
        run.all_krates("rustc_codegen_llvm")
    }

    fn make_run(run: RunConfig) {
        let backend = run.builder.config.rust_codegen_backends.get(0);
        let backend = backend.cloned().unwrap_or_else(|| {
            INTERNER.intern_str("llvm")
        });
        run.builder.ensure(CodegenBackend {
            target: run.target,
            backend,
        });
    }

    fn run(self, builder: &Builder) {
        let compiler = builder.compiler(0, builder.config.build);
        let target = self.target;
        let backend = self.backend;

        builder.ensure(Rustc { target });

        let mut cargo = builder.cargo(compiler, Mode::Codegen, target, "check");
        cargo.arg("--manifest-path").arg(builder.src.join("src/librustc_codegen_llvm/Cargo.toml"));
        rustc_cargo_env(builder, &mut cargo);

        // We won't build LLVM if it's not available, as it shouldn't affect `check`.

        let _folder = builder.fold_output(|| format!("stage{}-rustc_codegen_llvm", compiler.stage));
        run_cargo(builder,
                  &mut cargo,
                  vec![],
                  &codegen_backend_stamp(builder, compiler, target, backend),
                  true);
    }
}

#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
pub struct Test {
    pub target: Interned<String>,
}

impl Step for Test {
    type Output = ();
    const DEFAULT: bool = true;

    fn should_run(run: ShouldRun) -> ShouldRun {
        run.all_krates("test")
    }

    fn make_run(run: RunConfig) {
        run.builder.ensure(Test {
            target: run.target,
        });
    }

    fn run(self, builder: &Builder) {
        let compiler = builder.compiler(0, builder.config.build);
        let target = self.target;

        builder.ensure(Std { target });

        let mut cargo = builder.cargo(compiler, Mode::Test, target, "check");
        test_cargo(builder, &compiler, target, &mut cargo);

        let _folder = builder.fold_output(|| format!("stage{}-test", compiler.stage));
        builder.info(&format!("Checking test artifacts ({} -> {})", &compiler.host, target));
        run_cargo(builder,
                  &mut cargo,
                  vec![],
                  &libtest_stamp(builder, compiler, target),
                  true);

        let libdir = builder.sysroot_libdir(compiler, target);
        add_to_sysroot(builder, &libdir, &libtest_stamp(builder, compiler, target));
    }
}

#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
pub struct Rustdoc {
    pub target: Interned<String>,
}

impl Step for Rustdoc {
    type Output = ();
    const ONLY_HOSTS: bool = true;
    const DEFAULT: bool = true;

    fn should_run(run: ShouldRun) -> ShouldRun {
        run.path("src/tools/rustdoc")
    }

    fn make_run(run: RunConfig) {
        run.builder.ensure(Rustdoc {
            target: run.target,
        });
    }

    fn run(self, builder: &Builder) {
        let compiler = builder.compiler(0, builder.config.build);
        let target = self.target;

        builder.ensure(Rustc { target });

        let mut cargo = prepare_tool_cargo(builder,
                                           compiler,
                                           Mode::ToolRustc,
                                           target,
                                           "check",
                                           "src/tools/rustdoc",
                                           SourceType::InTree);

        let _folder = builder.fold_output(|| format!("stage{}-rustdoc", compiler.stage));
        println!("Checking rustdoc artifacts ({} -> {})", &compiler.host, target);
        run_cargo(builder,
                  &mut cargo,
                  vec![],
                  &rustdoc_stamp(builder, compiler, target),
                  true);

        let libdir = builder.sysroot_libdir(compiler, target);
        add_to_sysroot(&builder, &libdir, &rustdoc_stamp(builder, compiler, target));
        builder.cargo(compiler, Mode::ToolRustc, target, "clean");
    }
}

/// Cargo's output path for the standard library in a given stage, compiled
/// by a particular compiler for the specified target.
pub fn libstd_stamp(builder: &Builder, compiler: Compiler, target: Interned<String>) -> PathBuf {
    builder.cargo_out(compiler, Mode::Std, target).join(".libstd-check.stamp")
}

/// Cargo's output path for libtest in a given stage, compiled by a particular
/// compiler for the specified target.
pub fn libtest_stamp(builder: &Builder, compiler: Compiler, target: Interned<String>) -> PathBuf {
    builder.cargo_out(compiler, Mode::Test, target).join(".libtest-check.stamp")
}

/// Cargo's output path for librustc in a given stage, compiled by a particular
/// compiler for the specified target.
pub fn librustc_stamp(builder: &Builder, compiler: Compiler, target: Interned<String>) -> PathBuf {
    builder.cargo_out(compiler, Mode::Rustc, target).join(".librustc-check.stamp")
}

/// Cargo's output path for librustc_codegen_llvm in a given stage, compiled by a particular
/// compiler for the specified target and backend.
fn codegen_backend_stamp(builder: &Builder,
                         compiler: Compiler,
                         target: Interned<String>,
                         backend: Interned<String>) -> PathBuf {
    builder.cargo_out(compiler, Mode::Codegen, target)
         .join(format!(".librustc_codegen_llvm-{}-check.stamp", backend))
}

/// Cargo's output path for rustdoc in a given stage, compiled by a particular
/// compiler for the specified target.
pub fn rustdoc_stamp(builder: &Builder, compiler: Compiler, target: Interned<String>) -> PathBuf {
    builder.cargo_out(compiler, Mode::ToolRustc, target)
        .join(".rustdoc-check.stamp")
}
<|endoftext|>


<commit_before>mod opt;

use std::{
    net::SocketAddr,
    str,
    sync::atomic::{AtomicBool, AtomicU64, Ordering::SeqCst},
};

use crate::{cmd::run::Response, err, r, Result};
use crossbeam_channel::{Receiver, Sender};
use crossbeam_skiplist::SkipMap;
use futures::prelude::*;
use romio::TcpStream;
use scram::client::{ScramClient, ServerFinal, ServerFirst};
use serde::{Deserialize, Serialize};

pub use self::opt::*;

const NULL_BYTE: u8 = b'\0';
const BUF_SIZE: usize = 512;
const PROTOCOL_VERSION: usize = 0;

#[derive(Debug)]
pub struct Connection {
    db: String,
    stream: TcpStream,
    broken: AtomicBool,
    multiplex: Option<Multiplex>,
}

enum Version {
    V1_0 = 0x34c2_bdc3,
}

impl r {
    /// Create a new connection to the database server
    ///
    /// You can change the default connection options using [Opts].
    ///
    /// **Example:** Open a connection using the default host and port,
    /// specifying the default database.
    ///
    /// ```rust
    /// # use reql::{r, cmd::connect::Opts};
    /// #
    /// r.connect(Opts::default().db("marvel"))
    /// # ;
    /// ```
    ///
    /// The connection is created asynchronously, so you will have to `await`
    /// the result to get an actual connection.
    ///
    /// **Example:** Open a new connection to the database.
    ///
    /// ```rust
    /// # use reql::{r, cmd::connect::Opts};
    /// #
    /// let opts = Opts::default()
    ///     .host([127, 0, 0, 1])
    ///     .port(28015)
    ///     .db("marvel");
    /// r.connect(opts)
    /// # ;
    /// ```
    ///
    /// **Example:** Open a new connection to the database, specifying a
    /// user/password combination for authentication.
    ///
    /// ```rust
    /// # use reql::{r, cmd::connect::Opts};
    /// #
    /// let opts = Opts::default()
    ///     .host([127, 0, 0, 1])
    ///     .port(28015)
    ///     .db("marvel")
    ///     .user("herofinder")
    ///     .password("metropolis");
    /// r.connect(opts)
    /// # ;
    /// ```
    ///
    /// [Opts]: cmd/connect/struct.Opts.html
    pub async fn connect<O>(self, opts: O) -> Result<Connection>
    where
        O: Into<Option<Opts>>,
    {
        let opt = opts.into().unwrap_or_default();
        let addr = SocketAddr::new(opt.host, opt.port);
        let stream = await!(TcpStream::connect(&addr))?;
        let multiplex = if opt.multiplex {
            // Start counting from 1 because we want to use 0 to detect when the
            // token wraps over. If we allow tokens to be reused, the client may
            // return data not meant for that particular connection which may be
            // a security risk. This effectively means that a connection may be
            // used by `run` only up to `usize::max_value()` times.
            let counter = AtomicU64::new(1);
            let controller = SkipMap::new();
            Some(Multiplex {
                counter,
                controller,
            })
        } else {
            None
        };
        let conn = Connection {
            stream,
            multiplex,
            db: String::new(),
            broken: AtomicBool::new(false),
        };
        let handshake = HandShake {
            conn,
            buf: [0u8; BUF_SIZE],
        };
        await!(handshake.greet(opt))
    }
}

struct HandShake {
    // this should be enough for the handshake messages
    buf: [u8; BUF_SIZE],
    conn: Connection,
}

impl HandShake {
    // Performs the actual handshake
    //
    // This method optimises message exchange as suggested in the RethinkDB
    // documentation by sending message 3 right after message 1, without waiting
    // for message 2 first.
    async fn greet(mut self, opt: Opts) -> Result<Connection> {
        // Send the version we support
        let version = (Version::V1_0 as u32).to_le_bytes();
        await!(self.conn.stream.write_all(&version))?; // message 1

        // Send client first message
        let scram = ScramClient::new(&opt.user, &opt.password, None)?;
        let (scram, msg) = client_first(scram)?;
        await!(self.conn.stream.write_all(&msg))?; // message 3

        // Receive supported versions
        await!(self.conn.stream.read(&mut self.buf))?; // message 2
        let (len, info) = self.read_buf(0);
        ServerInfo::validate(info)?;

        // Receive server first message
        let offset = len + 1;
        let resp = if offset < BUF_SIZE && self.buf[offset] != NULL_BYTE {
            self.read_buf(offset).1
        } else {
            await!(self.conn.stream.read(&mut self.buf))?; // message 4
            self.read_buf(0).1
        };
        let info = AuthResponse::from_slice(resp)?;
        let auth = match info.authentication {
            Some(auth) => auth,
            None => {
                let msg = String::from("server did not send authentication info");
                return Err(err::Driver::Other(msg))?;
            }
        };

        // Send client final message
        let (scram, msg) = client_final(scram, &auth)?;
        await!(self.conn.stream.write_all(&msg))?; // message 5

        // Receive server final message
        await!(self.conn.stream.read(&mut self.buf))?; // message 6
        server_final(scram, self.read_buf(0).1)?;

        self.conn.db = opt.db;
        Ok(self.conn)
    }

    fn read_buf(&self, offset: usize) -> (usize, &[u8]) {
        let len = (&self.buf[offset..])
            .iter()
            .take_while(|x| **x != NULL_BYTE)
            .count();
        let max = offset + len;
        (max, &self.buf[offset..max])
    }
}

// We are going to use &str for `server_version` because it is safe to do so.
// Unfortunately, the other fields that are using String, are doing so because
// because they can potentially contain an escaped double quote which is not
// supported by serde in &str.
#[derive(Serialize, Deserialize, Debug)]
struct ServerInfo<'a> {
    success: bool,
    min_protocol_version: usize,
    max_protocol_version: usize,
    server_version: &'a str,
}

impl<'a> ServerInfo<'a> {
    fn validate(resp: &[u8]) -> Result<()> {
        match serde_json::from_slice::<ServerInfo>(resp) {
            Ok(info) => {
                if !info.success {
                    return Err(err::Runtime::Internal(resp.to_owned()))?;
                }
                if PROTOCOL_VERSION < info.min_protocol_version
                    || info.max_protocol_version < PROTOCOL_VERSION
                {
                    let msg = format!(
                        "unsupported protocol version {version}, expected between {min} and {max}",
                        version = PROTOCOL_VERSION,
                        min = info.min_protocol_version,
                        max = info.max_protocol_version,
                    );
                    Err(err::Driver::Other(msg))?
                }
                Ok(())
            }
            Err(_) => {
                let msg = str::from_utf8(resp)?;
                Err(err::Driver::Other(msg.to_owned()))?
            }
        }
    }
}

#[derive(Serialize, Deserialize, Debug)]
struct AuthRequest {
    protocol_version: usize,
    authentication_method: &'static str,
    authentication: String,
}

fn client_first<'a>(scram: ScramClient<'a>) -> Result<(ServerFirst<'a>, Vec<u8>)> {
    let (scram, client_first) = scram.client_first();
    let ar = AuthRequest {
        protocol_version: PROTOCOL_VERSION,
        authentication_method: "SCRAM-SHA-256",
        authentication: client_first,
    };
    let mut msg = serde_json::to_vec(&ar)?;
    msg.push(NULL_BYTE);
    Ok((scram, msg))
}

#[derive(Serialize, Deserialize, Debug)]
struct AuthConfirmation {
    authentication: String,
}

fn client_final<'a>(scram: ServerFirst<'a>, auth: &str) -> Result<(ServerFinal, Vec<u8>)> {
    let scram = scram.handle_server_first(auth)?;
    let (scram, client_final) = scram.client_final();
    let conf = AuthConfirmation {
        authentication: client_final,
    };
    let mut msg = serde_json::to_vec(&conf)?;
    msg.push(NULL_BYTE);
    Ok((scram, msg))
}

#[derive(Serialize, Deserialize, Debug)]
struct AuthResponse {
    success: bool,
    authentication: Option<String>,
    error_code: Option<usize>,
    error: Option<String>,
}

impl AuthResponse {
    fn from_slice(resp: &[u8]) -> Result<Self> {
        match serde_json::from_slice::<AuthResponse>(resp) {
            Ok(info) => {
                if !info.success {
                    // If error code is between 10 and 20, this is an auth error
                    if let Some(10...20) = info.error_code {
                        if let Some(msg) = info.error {
                            return Err(err::Driver::Auth(msg))?;
                        }
                    }
                    return Err(err::Runtime::Internal(resp.to_owned()))?;
                }
                Ok(info)
            }
            Err(_) => {
                let msg = str::from_utf8(resp)?;
                Err(err::Driver::Other(msg.to_owned()))?
            }
        }
    }
}

fn server_final(scram: ServerFinal, resp: &[u8]) -> Result<()> {
    let info = AuthResponse::from_slice(resp)?;
    if let Some(auth) = info.authentication {
        if let Err(err) = scram.handle_server_final(&auth) {
            return Err(err)?;
        }
    }
    Ok(())
}

impl Connection {
    /// Change the default database on this connection
    ///
    /// **Example:** Change the default database so that we don’t need to
    /// specify the database when referencing a table.
    ///
    /// ```rust
    /// # use reql::r;
    /// # use futures::executor::block_on;
    /// # let mut conn = block_on(r.connect(None)).unwrap();
    /// conn.use_db("marvel");
    /// r.table("heroes") // refers to r.db("marvel").table("heroes")
    /// # ;
    /// ```
    pub fn use_db(&mut self, name: &str) {
        self.db = name.to_owned();
    }

    #[doc(hidden)]
    pub fn broken(&self) -> bool {
        self.broken.load(SeqCst)
    }

    pub(super) fn stream(&self) -> &TcpStream {
        &self.stream
    }

    pub(super) fn mark_broken(&self) {
        self.broken.store(true, SeqCst);
    }

    pub(super) fn token(&self) -> Result<RequestId> {
        match self.multiplex {
            Some(Multiplex { ref counter, .. }) => {
                let id = counter.fetch_add(1, SeqCst);
                // the token has wrapped over
                if id == 0 {
                    self.mark_broken();
                    return Err(err::Driver::TokenOverflow)?;
                }
                Ok(id)
            }
            None => Ok(1),
        }
    }

    pub(super) fn controller(&self) -> Option<&Controller> {
        match self.multiplex {
            Some(Multiplex { ref controller, .. }) => Some(controller),
            None => None,
        }
    }
}

pub(super) type RequestId = u64;
pub(super) type Channel = (Sender<Response>, Receiver<Response>);
pub(super) type Controller = SkipMap<RequestId, Channel>;

#[derive(Debug)]
pub(super) struct Multiplex {
    counter: AtomicU64,
    controller: Controller,
}

#[cfg(test)]
mod tests {
    use crate::r;
    use futures::executor::block_on;

    #[test]
    fn driver_can_connect() -> crate::Result<()> {
        block_on(
            async {
                await!(r.connect(None))?;
                Ok(())
            },
        )
    }
}
<commit_msg>silence clippy on version comparison<commit_after>mod opt;

use std::{
    net::SocketAddr,
    str,
    sync::atomic::{AtomicBool, AtomicU64, Ordering::SeqCst},
};

use crate::{cmd::run::Response, err, r, Result};
use crossbeam_channel::{Receiver, Sender};
use crossbeam_skiplist::SkipMap;
use futures::prelude::*;
use romio::TcpStream;
use scram::client::{ScramClient, ServerFinal, ServerFirst};
use serde::{Deserialize, Serialize};

pub use self::opt::*;

const NULL_BYTE: u8 = b'\0';
const BUF_SIZE: usize = 512;
const PROTOCOL_VERSION: usize = 0;

#[derive(Debug)]
pub struct Connection {
    db: String,
    stream: TcpStream,
    broken: AtomicBool,
    multiplex: Option<Multiplex>,
}

enum Version {
    V1_0 = 0x34c2_bdc3,
}

impl r {
    /// Create a new connection to the database server
    ///
    /// You can change the default connection options using [Opts].
    ///
    /// **Example:** Open a connection using the default host and port,
    /// specifying the default database.
    ///
    /// ```rust
    /// # use reql::{r, cmd::connect::Opts};
    /// #
    /// r.connect(Opts::default().db("marvel"))
    /// # ;
    /// ```
    ///
    /// The connection is created asynchronously, so you will have to `await`
    /// the result to get an actual connection.
    ///
    /// **Example:** Open a new connection to the database.
    ///
    /// ```rust
    /// # use reql::{r, cmd::connect::Opts};
    /// #
    /// let opts = Opts::default()
    ///     .host([127, 0, 0, 1])
    ///     .port(28015)
    ///     .db("marvel");
    /// r.connect(opts)
    /// # ;
    /// ```
    ///
    /// **Example:** Open a new connection to the database, specifying a
    /// user/password combination for authentication.
    ///
    /// ```rust
    /// # use reql::{r, cmd::connect::Opts};
    /// #
    /// let opts = Opts::default()
    ///     .host([127, 0, 0, 1])
    ///     .port(28015)
    ///     .db("marvel")
    ///     .user("herofinder")
    ///     .password("metropolis");
    /// r.connect(opts)
    /// # ;
    /// ```
    ///
    /// [Opts]: cmd/connect/struct.Opts.html
    pub async fn connect<O>(self, opts: O) -> Result<Connection>
    where
        O: Into<Option<Opts>>,
    {
        let opt = opts.into().unwrap_or_default();
        let addr = SocketAddr::new(opt.host, opt.port);
        let stream = await!(TcpStream::connect(&addr))?;
        let multiplex = if opt.multiplex {
            // Start counting from 1 because we want to use 0 to detect when the
            // token wraps over. If we allow tokens to be reused, the client may
            // return data not meant for that particular connection which may be
            // a security risk. This effectively means that a connection may be
            // used by `run` only up to `usize::max_value()` times.
            let counter = AtomicU64::new(1);
            let controller = SkipMap::new();
            Some(Multiplex {
                counter,
                controller,
            })
        } else {
            None
        };
        let conn = Connection {
            stream,
            multiplex,
            db: String::new(),
            broken: AtomicBool::new(false),
        };
        let handshake = HandShake {
            conn,
            buf: [0u8; BUF_SIZE],
        };
        await!(handshake.greet(opt))
    }
}

struct HandShake {
    // this should be enough for the handshake messages
    buf: [u8; BUF_SIZE],
    conn: Connection,
}

impl HandShake {
    // Performs the actual handshake
    //
    // This method optimises message exchange as suggested in the RethinkDB
    // documentation by sending message 3 right after message 1, without waiting
    // for message 2 first.
    async fn greet(mut self, opt: Opts) -> Result<Connection> {
        // Send the version we support
        let version = (Version::V1_0 as u32).to_le_bytes();
        await!(self.conn.stream.write_all(&version))?; // message 1

        // Send client first message
        let scram = ScramClient::new(&opt.user, &opt.password, None)?;
        let (scram, msg) = client_first(scram)?;
        await!(self.conn.stream.write_all(&msg))?; // message 3

        // Receive supported versions
        await!(self.conn.stream.read(&mut self.buf))?; // message 2
        let (len, info) = self.read_buf(0);
        ServerInfo::validate(info)?;

        // Receive server first message
        let offset = len + 1;
        let resp = if offset < BUF_SIZE && self.buf[offset] != NULL_BYTE {
            self.read_buf(offset).1
        } else {
            await!(self.conn.stream.read(&mut self.buf))?; // message 4
            self.read_buf(0).1
        };
        let info = AuthResponse::from_slice(resp)?;
        let auth = match info.authentication {
            Some(auth) => auth,
            None => {
                let msg = String::from("server did not send authentication info");
                return Err(err::Driver::Other(msg))?;
            }
        };

        // Send client final message
        let (scram, msg) = client_final(scram, &auth)?;
        await!(self.conn.stream.write_all(&msg))?; // message 5

        // Receive server final message
        await!(self.conn.stream.read(&mut self.buf))?; // message 6
        server_final(scram, self.read_buf(0).1)?;

        self.conn.db = opt.db;
        Ok(self.conn)
    }

    fn read_buf(&self, offset: usize) -> (usize, &[u8]) {
        let len = (&self.buf[offset..])
            .iter()
            .take_while(|x| **x != NULL_BYTE)
            .count();
        let max = offset + len;
        (max, &self.buf[offset..max])
    }
}

// We are going to use &str for `server_version` because it is safe to do so.
// Unfortunately, the other fields that are using String, are doing so because
// because they can potentially contain an escaped double quote which is not
// supported by serde in &str.
#[derive(Serialize, Deserialize, Debug)]
struct ServerInfo<'a> {
    success: bool,
    min_protocol_version: usize,
    max_protocol_version: usize,
    server_version: &'a str,
}

impl<'a> ServerInfo<'a> {
    fn validate(resp: &[u8]) -> Result<()> {
        #[allow(clippy::absurd_extreme_comparisons)]
        match serde_json::from_slice::<ServerInfo>(resp) {
            Ok(info) => {
                if !info.success {
                    return Err(err::Runtime::Internal(resp.to_owned()))?;
                }
                if PROTOCOL_VERSION < info.min_protocol_version
                    || info.max_protocol_version < PROTOCOL_VERSION
                {
                    let msg = format!(
                        "unsupported protocol version {version}, expected between {min} and {max}",
                        version = PROTOCOL_VERSION,
                        min = info.min_protocol_version,
                        max = info.max_protocol_version,
                    );
                    Err(err::Driver::Other(msg))?
                }
                Ok(())
            }
            Err(_) => {
                let msg = str::from_utf8(resp)?;
                Err(err::Driver::Other(msg.to_owned()))?
            }
        }
    }
}

#[derive(Serialize, Deserialize, Debug)]
struct AuthRequest {
    protocol_version: usize,
    authentication_method: &'static str,
    authentication: String,
}

fn client_first<'a>(scram: ScramClient<'a>) -> Result<(ServerFirst<'a>, Vec<u8>)> {
    let (scram, client_first) = scram.client_first();
    let ar = AuthRequest {
        protocol_version: PROTOCOL_VERSION,
        authentication_method: "SCRAM-SHA-256",
        authentication: client_first,
    };
    let mut msg = serde_json::to_vec(&ar)?;
    msg.push(NULL_BYTE);
    Ok((scram, msg))
}

#[derive(Serialize, Deserialize, Debug)]
struct AuthConfirmation {
    authentication: String,
}

fn client_final<'a>(scram: ServerFirst<'a>, auth: &str) -> Result<(ServerFinal, Vec<u8>)> {
    let scram = scram.handle_server_first(auth)?;
    let (scram, client_final) = scram.client_final();
    let conf = AuthConfirmation {
        authentication: client_final,
    };
    let mut msg = serde_json::to_vec(&conf)?;
    msg.push(NULL_BYTE);
    Ok((scram, msg))
}

#[derive(Serialize, Deserialize, Debug)]
struct AuthResponse {
    success: bool,
    authentication: Option<String>,
    error_code: Option<usize>,
    error: Option<String>,
}

impl AuthResponse {
    fn from_slice(resp: &[u8]) -> Result<Self> {
        match serde_json::from_slice::<AuthResponse>(resp) {
            Ok(info) => {
                if !info.success {
                    // If error code is between 10 and 20, this is an auth error
                    if let Some(10...20) = info.error_code {
                        if let Some(msg) = info.error {
                            return Err(err::Driver::Auth(msg))?;
                        }
                    }
                    return Err(err::Runtime::Internal(resp.to_owned()))?;
                }
                Ok(info)
            }
            Err(_) => {
                let msg = str::from_utf8(resp)?;
                Err(err::Driver::Other(msg.to_owned()))?
            }
        }
    }
}

fn server_final(scram: ServerFinal, resp: &[u8]) -> Result<()> {
    let info = AuthResponse::from_slice(resp)?;
    if let Some(auth) = info.authentication {
        if let Err(err) = scram.handle_server_final(&auth) {
            return Err(err)?;
        }
    }
    Ok(())
}

impl Connection {
    /// Change the default database on this connection
    ///
    /// **Example:** Change the default database so that we don’t need to
    /// specify the database when referencing a table.
    ///
    /// ```rust
    /// # use reql::r;
    /// # use futures::executor::block_on;
    /// # let mut conn = block_on(r.connect(None)).unwrap();
    /// conn.use_db("marvel");
    /// r.table("heroes") // refers to r.db("marvel").table("heroes")
    /// # ;
    /// ```
    pub fn use_db(&mut self, name: &str) {
        self.db = name.to_owned();
    }

    #[doc(hidden)]
    pub fn broken(&self) -> bool {
        self.broken.load(SeqCst)
    }

    pub(super) fn stream(&self) -> &TcpStream {
        &self.stream
    }

    pub(super) fn mark_broken(&self) {
        self.broken.store(true, SeqCst);
    }

    pub(super) fn token(&self) -> Result<RequestId> {
        match self.multiplex {
            Some(Multiplex { ref counter, .. }) => {
                let id = counter.fetch_add(1, SeqCst);
                // the token has wrapped over
                if id == 0 {
                    self.mark_broken();
                    return Err(err::Driver::TokenOverflow)?;
                }
                Ok(id)
            }
            None => Ok(1),
        }
    }

    pub(super) fn controller(&self) -> Option<&Controller> {
        match self.multiplex {
            Some(Multiplex { ref controller, .. }) => Some(controller),
            None => None,
        }
    }
}

pub(super) type RequestId = u64;
pub(super) type Channel = (Sender<Response>, Receiver<Response>);
pub(super) type Controller = SkipMap<RequestId, Channel>;

#[derive(Debug)]
pub(super) struct Multiplex {
    counter: AtomicU64,
    controller: Controller,
}

#[cfg(test)]
mod tests {
    use crate::r;
    use futures::executor::block_on;

    #[test]
    fn driver_can_connect() -> crate::Result<()> {
        block_on(
            async {
                await!(r.connect(None))?;
                Ok(())
            },
        )
    }
}
<|endoftext|>


<commit_before>use std::collections::HashSet;
use std::fmt::Display;
use std::io::{Result, Write};

use analysis::general::StatusedTypeId;
use config::Config;
use git::repo_hash;
use version::VERSION;

pub fn start_comments<W: Write>(w: &mut W, conf: &Config) -> Result<()>{
    try!(writeln!(w, "// This file was generated by gir ({}) from gir-files ({})",
            VERSION, repo_hash(&conf.girs_dir).unwrap_or_else(|_| "???".into())));
    try!(writeln!(w, "// DO NOT EDIT"));

    Ok(())
}

pub fn uses<W: Write>(w: &mut W, used_types: &HashSet<String>) -> Result<()>{
    let v = vec![
        "",
        "use glib::translate::*;",
        "use glib::types;",
        "use ffi;",
        "",
        "use object::*;",
    ];
    try!(write_vec(w, &v));

    let mut used_types: Vec<String> = used_types.iter()
        .map(|s| s.clone()).collect();
    used_types.sort_by(|a, b| a.cmp(b));

    for name in used_types {
        try!(writeln!(w, "use {};", name));
    }

    Ok(())
}

pub fn objects_child_type<W: Write>(w: &mut W, type_name: &str, glib_name: &str) -> Result<()>{
    try!(writeln!(w, ""));
    try!(writeln!(w, "pub type {} = Object<ffi::{}>;", type_name, glib_name));

    Ok(())
}

pub fn impl_parents<W: Write>(w: &mut W, type_name: &str, parents: &Vec<StatusedTypeId>) -> Result<()>{
    try!(writeln!(w, ""));
    for stid in parents {
        //TODO: don't generate for parents without traits
        try!(writeln!(w, "unsafe impl Upcast<{}> for {} {{ }}", stid.name, type_name));
    }

    Ok(())
}

pub fn impl_interfaces<W: Write>(w: &mut W, type_name: &str, implements: &Vec<StatusedTypeId>) -> Result<()>{
    try!(writeln!(w, ""));
    for stid in implements {
        try!(writeln!(w, "unsafe impl Upcast<{}> for {} {{ }}", stid.name, type_name));
    }

    Ok(())
}

pub fn impl_static_type<W: Write>(w: &mut W, type_name: &str, glib_func_name: &str) -> Result<()>{
    try!(writeln!(w, ""));
    try!(writeln!(w, "impl types::StaticType for {} {{", type_name));
    try!(writeln!(w, "{}#[inline]", tabs(1)));
    try!(writeln!(w, "{}fn static_type() -> types::Type {{", tabs(1)));
    try!(writeln!(w, "{}unsafe {{ from_glib(ffi::{}()) }}", tabs(2), glib_func_name));
    try!(writeln!(w, "{}}}", tabs(1)));
    try!(writeln!(w, "}}"));

    Ok(())
}

//TODO: convert to macro with usage
//format!(indent!(5, "format:{}"), 6)
pub fn tabs(num: i32) -> String {
    (0..num).map(|_| "    ").collect::<String>()
}

//TODO: use in non sys codegen
pub fn fix_parameter_name(name: &str) -> String{
    let bad_names = vec![
        "self",
        "type",
        "box",
    ];
    if bad_names.contains(&name) {
        format!("{}_", name)
    } else {
        name.into()
    }
}

pub fn write_vec<W: Write, T: Display>(w: &mut W, v: &Vec<T>) -> Result<()> {
    for s in v {
        try!(writeln!(w, "{}", s));
    }
    Ok(())
}
<commit_msg>Add another reserved word<commit_after>use std::collections::HashSet;
use std::fmt::Display;
use std::io::{Result, Write};

use analysis::general::StatusedTypeId;
use config::Config;
use git::repo_hash;
use version::VERSION;

pub fn start_comments<W: Write>(w: &mut W, conf: &Config) -> Result<()>{
    try!(writeln!(w, "// This file was generated by gir ({}) from gir-files ({})",
            VERSION, repo_hash(&conf.girs_dir).unwrap_or_else(|_| "???".into())));
    try!(writeln!(w, "// DO NOT EDIT"));

    Ok(())
}

pub fn uses<W: Write>(w: &mut W, used_types: &HashSet<String>) -> Result<()>{
    let v = vec![
        "",
        "use glib::translate::*;",
        "use glib::types;",
        "use ffi;",
        "",
        "use object::*;",
    ];
    try!(write_vec(w, &v));

    let mut used_types: Vec<String> = used_types.iter()
        .map(|s| s.clone()).collect();
    used_types.sort_by(|a, b| a.cmp(b));

    for name in used_types {
        try!(writeln!(w, "use {};", name));
    }

    Ok(())
}

pub fn objects_child_type<W: Write>(w: &mut W, type_name: &str, glib_name: &str) -> Result<()>{
    try!(writeln!(w, ""));
    try!(writeln!(w, "pub type {} = Object<ffi::{}>;", type_name, glib_name));

    Ok(())
}

pub fn impl_parents<W: Write>(w: &mut W, type_name: &str, parents: &Vec<StatusedTypeId>) -> Result<()>{
    try!(writeln!(w, ""));
    for stid in parents {
        //TODO: don't generate for parents without traits
        try!(writeln!(w, "unsafe impl Upcast<{}> for {} {{ }}", stid.name, type_name));
    }

    Ok(())
}

pub fn impl_interfaces<W: Write>(w: &mut W, type_name: &str, implements: &Vec<StatusedTypeId>) -> Result<()>{
    try!(writeln!(w, ""));
    for stid in implements {
        try!(writeln!(w, "unsafe impl Upcast<{}> for {} {{ }}", stid.name, type_name));
    }

    Ok(())
}

pub fn impl_static_type<W: Write>(w: &mut W, type_name: &str, glib_func_name: &str) -> Result<()>{
    try!(writeln!(w, ""));
    try!(writeln!(w, "impl types::StaticType for {} {{", type_name));
    try!(writeln!(w, "{}#[inline]", tabs(1)));
    try!(writeln!(w, "{}fn static_type() -> types::Type {{", tabs(1)));
    try!(writeln!(w, "{}unsafe {{ from_glib(ffi::{}()) }}", tabs(2), glib_func_name));
    try!(writeln!(w, "{}}}", tabs(1)));
    try!(writeln!(w, "}}"));

    Ok(())
}

//TODO: convert to macro with usage
//format!(indent!(5, "format:{}"), 6)
pub fn tabs(num: i32) -> String {
    (0..num).map(|_| "    ").collect::<String>()
}

//TODO: use in non sys codegen
pub fn fix_parameter_name(name: &str) -> String{
    let bad_names = vec![
        "self",
        "type",
        "box",
        "where",
    ];
    if bad_names.contains(&name) {
        format!("{}_", name)
    } else {
        name.into()
    }
}

pub fn write_vec<W: Write, T: Display>(w: &mut W, v: &Vec<T>) -> Result<()> {
    for s in v {
        try!(writeln!(w, "{}", s));
    }
    Ok(())
}
<|endoftext|>