<commit_before>use std::collections::BTreeMap;

use anyhow::Result;
use log::debug;

use crate::{
    analysis::dis,
    aspace::AddressSpace,
    module::{Module, Permissions},
    VA,
};
use lancelot_flirt::*;

/// make a best guess as to the reference target, found at `ref_offset` from
/// `va`.
fn get_ref(module: &Module, decoder: &zydis::Decoder, va: VA, ref_offset: u64) -> Option<VA> {
    for i in (1..=4u64).rev() {
        let candidate_insn_va = va + ref_offset - i;
        let mut insn_buf = [0u8; 16];

        if module.address_space.read_into(candidate_insn_va, &mut insn_buf).is_ok() {
            if let Ok(Some(insn)) = decoder.decode(&insn_buf) {
                let explicit_operands = insn
                    .operands
                    .iter()
                    .filter(|op| matches!(op.visibility, zydis::OperandVisibility::EXPLICIT));
                for (j, op) in explicit_operands.take(2).enumerate() {
                    match op.ty {
                        zydis::OperandType::MEMORY => {
                            if (op.mem.base == zydis::Register::NONE || op.mem.base == zydis::Register::RIP)
                                && op.mem.index == zydis::Register::NONE
                                && op.mem.scale == 0
                                && op.mem.disp.has_displacement
                                && insn.raw.disp_offset == i as u8
                            {
                                if let Ok(target) = insn.calc_absolute_address(candidate_insn_va, op) {
                                    if module.probe_va(target, Permissions::RX) {
                                        return Some(target);
                                    }
                                }
                            }
                            continue;
                        }
                        zydis::OperandType::IMMEDIATE => {
                            if insn.raw.imm[j].offset == i as u8 {
                                if let Ok(target) = insn.calc_absolute_address(candidate_insn_va, op) {
                                    if module.probe_va(target, Permissions::RX) {
                                        return Some(target);
                                    }
                                }
                            }
                            continue;
                        }
                        zydis::OperandType::POINTER => continue,
                        zydis::OperandType::REGISTER => continue,
                        zydis::OperandType::UNUSED => continue,
                    }
                }
            }
        }
    }

    None
}

pub fn match_flirt(module: &Module, sigs: &FlirtSignatureSet, va: VA) -> Result<Vec<FlirtSignature>> {
    fn match_flirt_inner(
        module: &Module,
        sigs: &FlirtSignatureSet,
        decoder: &zydis::Decoder,
        va: VA,
        cache: &mut BTreeMap<VA, Vec<FlirtSignature>>,
    ) -> Result<Vec<FlirtSignature>> {
        let sec = module
            .sections
            .iter()
            .find(|sec| sec.virtual_range.start <= va && va < sec.virtual_range.end)
            .unwrap();

        let size = sec.virtual_range.end - va;
        let buf = module.address_space.read_bytes(va, size as usize)?;

        debug!("flirt: matching: {:#x}", va);

        Ok(sigs
            .r#match(&buf)
            .iter()
            .filter(|sig| {
                let mut does_match_references = true;

                debug!("flirt: {:#x}: candidate: {:?}", va, sig);

                'names: for name in sig.names.iter() {
                    if let Symbol::Reference(Name {
                        offset,
                        name: wanted_name,
                    }) = name
                    {
                        // i dont know what this means.
                        assert!(*offset >= 0, "negative offset");

                        if let Some(target) = get_ref(module, &decoder, va, *offset as u64) {
                            // TODO: special case "."

                            // can't use entry because of mutable cache used to create cache entry.
                            #[allow(clippy::map_entry)]
                            if !cache.contains_key(&target) {
                                let target_sigs = match_flirt_inner(module, sigs, decoder, target, cache)
                                    .unwrap_or_else(|_| Default::default());
                                cache.insert(target, target_sigs);
                            }

                            let target_sigs = cache.get(&target).unwrap();

                            let mut does_name_match = false;
                            'sigs: for target_sig in target_sigs.iter() {
                                debug!("flirt: {:#x}: found reference: {:?} @ {:#x}", va, target_sig, offset);
                                for name in target_sig.names.iter() {
                                    match name {
                                        Symbol::Reference(_) => continue,
                                        Symbol::Local(Name {
                                            name: target_name,
                                            offset,
                                        }) => {
                                            if *offset == 0 && target_name == wanted_name {
                                                does_name_match = true;
                                                break 'sigs;
                                            }
                                        }
                                        Symbol::Public(Name {
                                            name: target_name,
                                            offset,
                                        }) => {
                                            if *offset == 0 && target_name == wanted_name {
                                                does_name_match = true;
                                                break 'sigs;
                                            }
                                        }
                                    }
                                }
                            }

                            if !does_name_match {
                                does_match_references = false;
                                break 'names;
                            }
                        } else {
                            does_match_references = false;
                            break;
                        }
                    }
                }

                does_match_references
            })
            .cloned()
            .cloned()
            .collect::<Vec<_>>())
    }

    let decoder = dis::get_disassembler(module)?;
    let mut cache = Default::default();
    match_flirt_inner(module, sigs, &decoder, va, &mut cache)
}
<commit_msg>flirt: add some doc<commit_after>use std::collections::BTreeMap;

use anyhow::Result;
use log::debug;

use crate::{
    analysis::dis,
    aspace::AddressSpace,
    module::{Module, Permissions},
    VA,
};
use lancelot_flirt::*;

/// make a best guess as to the reference target, found at `ref_offset` from
/// `va`.
fn get_ref(module: &Module, decoder: &zydis::Decoder, va: VA, ref_offset: u64) -> Option<VA> {
    for i in (1..=4u64).rev() {
        let candidate_insn_va = va + ref_offset - i;
        let mut insn_buf = [0u8; 16];

        if module.address_space.read_into(candidate_insn_va, &mut insn_buf).is_ok() {
            if let Ok(Some(insn)) = decoder.decode(&insn_buf) {
                let explicit_operands = insn
                    .operands
                    .iter()
                    .filter(|op| matches!(op.visibility, zydis::OperandVisibility::EXPLICIT));
                for (j, op) in explicit_operands.take(2).enumerate() {
                    match op.ty {
                        zydis::OperandType::MEMORY => {
                            if (op.mem.base == zydis::Register::NONE || op.mem.base == zydis::Register::RIP)
                                && op.mem.index == zydis::Register::NONE
                                && op.mem.scale == 0
                                && op.mem.disp.has_displacement
                                && insn.raw.disp_offset == i as u8
                            {
                                if let Ok(target) = insn.calc_absolute_address(candidate_insn_va, op) {
                                    if module.probe_va(target, Permissions::RX) {
                                        return Some(target);
                                    }
                                }
                            }
                            continue;
                        }
                        zydis::OperandType::IMMEDIATE => {
                            if insn.raw.imm[j].offset == i as u8 {
                                if let Ok(target) = insn.calc_absolute_address(candidate_insn_va, op) {
                                    if module.probe_va(target, Permissions::RX) {
                                        return Some(target);
                                    }
                                }
                            }
                            continue;
                        }
                        zydis::OperandType::POINTER => continue,
                        zydis::OperandType::REGISTER => continue,
                        zydis::OperandType::UNUSED => continue,
                    }
                }
            }
        }
    }

    None
}

pub fn match_flirt(module: &Module, sigs: &FlirtSignatureSet, va: VA) -> Result<Vec<FlirtSignature>> {
    fn match_flirt_inner(
        module: &Module,
        sigs: &FlirtSignatureSet,
        decoder: &zydis::Decoder,
        va: VA,
        cache: &mut BTreeMap<VA, Vec<FlirtSignature>>,
    ) -> Result<Vec<FlirtSignature>> {
        let sec = module
            .sections
            .iter()
            .find(|sec| sec.virtual_range.start <= va && va < sec.virtual_range.end)
            .unwrap();

        let size = sec.virtual_range.end - va;
        let buf = module.address_space.read_bytes(va, size as usize)?;

        debug!("flirt: matching: {:#x}", va);

        Ok(sigs
            .r#match(&buf)
            .iter()
            .filter(|sig| {
                let mut does_match_references = true;

                debug!("flirt: {:#x}: candidate: {:?}", va, sig);

                'names: for name in sig.names.iter() {
                    if let Symbol::Reference(Name {
                        offset,
                        name: wanted_name,
                    }) = name
                    {
                        // i dont know what this means.
                        assert!(*offset >= 0, "negative offset");

                        if let Some(target) = get_ref(module, &decoder, va, *offset as u64) {
                            if wanted_name == "." {
                                // special case: name "." matches anything
                                // not exactly sure, might instead match special data `ctype`.
                                // see: https://github.com/williballenthin/lancelot/issues/112#issuecomment-802379966

                                // TODO: get_ref only returns executable references.
                                continue;
                            }

                            // can't use entry API because of mutable cache used to create cache entry.
                            #[allow(clippy::map_entry)]
                            if !cache.contains_key(&target) {
                                let target_sigs = match_flirt_inner(module, sigs, decoder, target, cache)
                                    .unwrap_or_else(|_| Default::default());
                                cache.insert(target, target_sigs);
                            }

                            let target_sigs = cache.get(&target).unwrap();

                            let mut does_name_match = false;
                            'sigs: for target_sig in target_sigs.iter() {
                                debug!("flirt: {:#x}: found reference: {:?} @ {:#x}", va, target_sig, offset);
                                for name in target_sig.names.iter() {
                                    match name {
                                        Symbol::Reference(_) => continue,
                                        Symbol::Local(Name {
                                            name: target_name,
                                            offset,
                                        }) => {
                                            if *offset == 0 && target_name == wanted_name {
                                                does_name_match = true;
                                                break 'sigs;
                                            }
                                        }
                                        Symbol::Public(Name {
                                            name: target_name,
                                            offset,
                                        }) => {
                                            if *offset == 0 && target_name == wanted_name {
                                                does_name_match = true;
                                                break 'sigs;
                                            }
                                        }
                                    }
                                }
                            }

                            if !does_name_match {
                                does_match_references = false;
                                break 'names;
                            }
                        } else {
                            does_match_references = false;
                            break;
                        }
                    }
                }

                does_match_references
            })
            .cloned()
            .cloned()
            .collect::<Vec<_>>())
    }

    let decoder = dis::get_disassembler(module)?;
    let mut cache = Default::default();
    match_flirt_inner(module, sigs, &decoder, va, &mut cache)
}
<|endoftext|>


<commit_before>use block_cipher_operation::*;
use cbc::*;
use ciphertext::*;
use key::*;

pub struct MessageOperator {
	message: Vec<u8>,
	block_operator: String,
	IV: Vec<u8>,
	block_size: usize,
	key: Vec<u8>,
}

impl MessageOperator {
	pub fn encrypt(&self) -> Vec<u8> {
		let mut result_list: Vec<Vec<u8>> = Vec::new();
		let mut result: Vec<u8> = Vec::new();

		let mut last_iv = self.IV.to_owned();

		for chunk in self.message.chunks(self.block_size) {
			let mut cipher_block = self.make_operation(chunk, last_iv);

			let encrypted = cipher_block.encrypt();
			last_iv = encrypted.to_owned();

			result.extend(&encrypted[..]);
		}
		result
	}

	pub fn decrypt(&self) -> Vec<u8> {
		let mut result_list: Vec<Vec<u8>> = Vec::new();
		let mut result: Vec<u8> = Vec::new();

		let mut last_iv = self.IV.to_owned();

		for chunk in self.message.chunks(self.block_size).rev() {
			println!("{:?}", chunk);
			/*
			let mut cipher_block = self.make_operation(chunk, last_iv);

			let encrypted = cipher_block.encrypt();
			last_iv = encrypted.to_owned();

			result.extend(&encrypted[..]);
			*/
		}
		result
	}

	fn make_operation(&self, chunk: &[u8], IV: Vec<u8>) -> Box<BlockCipherOperation> {
		if self.block_operator == "CBC" {
			let padded_chunk = PaddedBytes::from_bytes(&chunk, self.block_size);

			return Box::new(CBC {
				input: padded_chunk.unwrap(),
				initialisation_vector: IV.to_owned(),
				key: RepeatingKey::new_bytes(&self.key)
			})
		}

		panic!("Unexpected BlockCipherOperation specified.");
	}
}

//-----------------------------------------------------------------------------

#[cfg(test)]
#[allow(non_snake_case)]
mod test {
	use super::*;

	use block_cipher_operation::*;
	use byte_conversion::*;
	use ciphertext::*;
	use key::*;

	#[test]
	fn MessageOperator_simple() {
		let block_size = 16;
		let message = "This is some short message to encrypt.";
		let block_operator = "CBC";
		let IV = "SomeIV";
		let key = "YELLOW SUBMARINE";

		let mo_encrypt = MessageOperator {
			message: readable_text_to_bytes(&message),
			block_operator: block_operator.to_string(),
			IV: RepeatingKey::new(&IV).of_length(block_size),
			block_size: 16,
			key: readable_text_to_bytes(&key),
		};

		let encrypted = mo_encrypt.encrypt();
		assert_eq!(16*3, encrypted.len());

		let mo_decrypt = MessageOperator {
			message: encrypted,
			block_operator: block_operator.to_string(),
			IV: RepeatingKey::new(&IV).of_length(block_size),
			block_size: 16,
			key: readable_text_to_bytes(&key),
		};

		assert_eq!(readable_text_to_bytes(&message), mo_decrypt.decrypt());
	}
}<commit_msg>Encrypt/Decrypt works for single blocks<commit_after>use block_cipher_operation::*;
use cbc::*;
use ciphertext::*;
use key::*;

pub struct MessageOperator {
	message: Vec<u8>,
	block_operator: String,
	IV: Vec<u8>,
	block_size: usize,
	key: Vec<u8>,
}

impl MessageOperator {
	pub fn encrypt(&self) -> Vec<u8> {
		let mut result: Vec<u8> = Vec::new();
		let mut last_iv = self.IV.to_owned();

		for chunk in self.message.chunks(self.block_size) {
			let mut cipher_block = self.make_operation(chunk, last_iv);

			let encrypted = cipher_block.encrypt();
			last_iv = encrypted.to_owned();

			result.extend(&encrypted[..]);
		}
		result
	}

	fn decrypt_raw(&self) -> Vec<u8> {
		let mut result: Vec<u8> = Vec::new();
		let mut last_iv = self.IV.to_owned();

		for chunk in self.message.chunks(self.block_size).rev() {
			let mut cipher_block = self.make_operation(chunk, last_iv);
			let decrypted = cipher_block.decrypt();

			last_iv = chunk.to_owned();

			result.extend(&decrypted[..]);
		}
		result
	}

	pub fn decrypt(&self) -> Vec<u8> {
		let decrypted = self.decrypt_raw();
		let last_char: u8 = *decrypted.last().unwrap();

		// Remove padding
		if last_char < self.block_size as u8 {
			let check_vec = vec![last_char; last_char as usize];
			if decrypted.ends_with(&check_vec[..]) { // Is there actually padding?
				let length = decrypted.len();
				let pos = (length as u8 - last_char) as usize;

				return decrypted.iter()
					.take(pos)
					.map(|x:&u8| *x)
					.collect();
			}
		}

		decrypted
	}

	fn make_operation(&self, chunk: &[u8], IV: Vec<u8>) -> Box<BlockCipherOperation> {
		if self.block_operator == "CBC" {
			let padded_chunk = PaddedBytes::from_bytes(&chunk, self.block_size);

			return Box::new(CBC {
				input: padded_chunk.unwrap(),
				initialisation_vector: IV.to_owned(),
				key: RepeatingKey::new_bytes(&self.key)
			})
		}

		panic!("Unexpected BlockCipherOperation specified.");
	}
}

//-----------------------------------------------------------------------------

#[cfg(test)]
#[allow(non_snake_case)]
mod test {
	use super::*;

	use block_cipher_operation::*;
	use byte_conversion::*;
	use ciphertext::*;
	use key::*;

	#[test]
	fn MessageOperator_simple() {
		let block_size = 16;
		let message = "12345";
		let block_operator = "CBC";
		let IV = "SomeIV";
		let key = "YELLOW SUBMARINE";

		let mo_encrypt = MessageOperator {
			message: readable_text_to_bytes(&message),
			block_operator: block_operator.to_string(),
			IV: RepeatingKey::new(&IV).of_length(block_size),
			block_size: block_size,
			key: readable_text_to_bytes(&key),
		};

		let encrypted = mo_encrypt.encrypt();
		assert_eq!(block_size, encrypted.len());

		let mo_decrypt = MessageOperator {
			message: encrypted,
			block_operator: block_operator.to_string(),
			IV: RepeatingKey::new(&IV).of_length(block_size),
			block_size: block_size,
			key: readable_text_to_bytes(&key),
		};

		assert_eq!(readable_text_to_bytes(&message), mo_decrypt.decrypt());
	}

	#[test]
	fn MessageOperator_messagelongerthanblocksize() {
		let block_size = 16;
		let message = "This is a longer message to decrypt";
		let block_operator = "CBC";
		let IV = "SomeIV";
		let key = "YELLOW SUBMARINE";

		let mo_encrypt = MessageOperator {
			message: readable_text_to_bytes(&message),
			block_operator: block_operator.to_string(),
			IV: RepeatingKey::new(&IV).of_length(block_size),
			block_size: block_size,
			key: readable_text_to_bytes(&key),
		};

		let encrypted = mo_encrypt.encrypt();
		assert_eq!(block_size*3, encrypted.len());

		let mo_decrypt = MessageOperator {
			message: encrypted,
			block_operator: block_operator.to_string(),
			IV: RepeatingKey::new(&IV).of_length(block_size),
			block_size: block_size,
			key: readable_text_to_bytes(&key),
		};

		assert_eq!(readable_text_to_bytes(&message), mo_decrypt.decrypt());
	}
}<|endoftext|>


<commit_before>//! A custom Lua callback mapping, where the default operation is done and
//! then a custom, user-defined Lua callback is ran.
use rustwlc::*;
use rustwlc::Size;
use ::lua::{self, LuaQuery, LuaSendError};
use ::layout::{try_lock_tree, Handle};
use ::convert::json::{size_to_json, point_to_json, geometry_to_json};

use super::default::Default;
use super::Mode;

pub struct CustomLua;

impl Mode for CustomLua {
    /// Triggered when an output is connected.
    /// Sends the id of the created output to Lua.
    fn output_created(&self, output: WlcOutput) -> bool {
        let result = Default.output_created(output);
        if let Some(id) = lookup_handle(output.into()) {
            send_to_lua(format!("way_cooler.on_output_created(\"{}\")", id));
        }
        result
    }

    /// Triggered when an output is disconnected.
    /// Sends the id of the destroyed output to Lua.
    fn output_destroyed(&self, output: WlcOutput) {
        // Get the output ID before it's destroyed
        let id = lookup_handle(output.into());
        Default.output_destroyed(output);
        if let Some(id) = id {
            send_to_lua(format!("way_cooler.on_output_destroyed(\"{}\")", id))
        }
    }

    /// Triggered when an output loses or gains focus.
    /// Sends the id of the output, and whether it gained or lost focus.
    fn output_focused(&self, output: WlcOutput, focused: bool) {
        Default.output_focused(output, focused);
        if let Some(id) = lookup_handle(output.into()) {
            send_to_lua(format!("way_cooler.on_output_focused(\"{}\", {})",
                                id, focused));
        }
    }

    /// Triggered when an output's resolution changes.
    /// Sends the id of the output, and the old and new size of the output.
    fn output_resolution(&self, output: WlcOutput,
                         old_size: Size, new_size: Size) {
        Default.output_resolution(output, old_size, new_size);
        if let Some(id) = lookup_handle(output.into()) {
            send_to_lua(format!(
                "way_cooler.on_output_resolution_change(\"{}\", {}, {})",
                id, size_to_lua(old_size), size_to_lua(new_size)));
        }
    }

    /// Triggered after the output is rendered.
    /// Sends the id of the output after it's rendered.
    fn output_render_post(&self, output: WlcOutput) {
        Default.output_render_post(output);
        if let Some(id) = lookup_handle(output.into()) {
            send_to_lua(format!(
                "way_cooler.on_output_render_post(\"{}\")", id));
        }
    }

    /// Triggered when a view moves outputs.
    /// Sends the id of the view, the id of the old output, and the id of the
    /// the new output.
    fn view_moved_to_output(&self,
                            view: WlcView,
                            o1: WlcOutput,
                            o2: WlcOutput) {
        Default.view_moved_to_output(view, o1, o2);
        let (view_id, o1_id, o2_id) = (lookup_handle(view.into()),
                                       lookup_handle(o1.into()),
                                       lookup_handle(o2.into()));
        match (view_id, o1_id, o2_id) {
            (Some(view_id), Some(o1_id), Some(o2_id)) => {
                send_to_lua(format!(
                    "way_cooler.on_view_moved_to_output(\"{}\", \"{}\", \"{}\")",
                    view_id, o1_id, o2_id))
            },
            _ => {}
        }
    }

    /// Triggered when a view is created.
    /// Sends the id of the new view.
    fn view_created(&self, view: WlcView) -> bool {
        let result = Default.view_created(view);
        if let Some(id) = lookup_handle(view.into()) {
            send_to_lua(format!("way_cooler.on_view_created(\"{}\")", id))
        }
        result
    }

    /// Triggered when a view is destroyed.
    /// Sends the id of the view after its destroyed.
    fn view_destroyed(&self, view: WlcView) {
        // Get id before it's destroyed.
        let id = lookup_handle(view.into());
        Default.view_destroyed(view);
        if let Some(id) = id {
            send_to_lua(format!("way_cooler.on_view_destroyed(\"{}\")", id))
        }
    }

    /// Triggered when a view gains or loses focus.
    /// Sends the id of the view after its focus changes. The focus is true if
    /// it gained focus, or false if it lost focus.
    fn view_focused(&self, view: WlcView, focused: bool) {
        Default.view_focused(view, focused);
        if let Some(id) = lookup_handle(view.into()) {
            send_to_lua(format!("way_cooler.on_view_focused(\"{}\", {})",
                                id, focused))
        }
    }

    /// Triggered when the properties of a view changes.
    /// Sends the id of the view that had its properties changed, and the
    /// changed properties as a bit mask.
    fn view_props_changed(&self, view: WlcView, prop: ViewPropertyType) {
        Default.view_props_changed(view, prop);
        if let Some(id) = lookup_handle(view.into()) {
            send_to_lua(format!("way_coooler.on_view_props_changed(\"{}\", {})",
                                id, prop.bits()))
        }
    }

    /// Triggered when a view has its state change.
    /// Sends the id of the view, the bit field for the state field,
    /// and a bool representing how it was toggled.
    fn view_request_state(&self,
                          view: WlcView,
                          state: ViewState,
                          toggle: bool) {
        Default.view_request_state(view, state, toggle);
        if let Some(id) = lookup_handle(view.into()) {
            send_to_lua(format!(
                "way_cooler.on_view_request_state(\"{}\", {}, {})",
                id, state.bits(), toggle))
        }
    }

    /// Triggered when a view requests a different geometry.
    /// Sends the id of the view, and the new geometry.
    fn view_request_geometry(&self, view: WlcView, geometry: Geometry) {
        Default.view_request_geometry(view, geometry);
        if let Some(id) = lookup_handle(view.into()) {
            send_to_lua(format!(
                "way_cooler.on_view_request_geometry(\"{}\", {})",
                id, geometry_to_lua(geometry)))
        }
    }

    /// Triggered when a view requests to be moved, but not resized.
    /// Sends the id of the view, and the new destination point.
    fn view_request_move(&self, view: WlcView, dest: Point) {
        Default.view_request_move(view, dest);
        if let Some(id) = lookup_handle(view.into()) {
            send_to_lua(format!("way_cooler.on_view_request_move(\"{}\", {})",
                id, point_to_lua(dest)))
        }
    }

    /// Triggered when a view requests to be resized.
    /// Sends the id of the view, the bit field of the resize edge,
    /// and the point where the mouse resized.
    fn view_request_resize(&self,
                           view: WlcView,
                           edge: ResizeEdge,
                           point: Point) {
        Default.view_request_resize(view, edge, point);
        if let Some(id) = lookup_handle(view.into()) {
            send_to_lua(format!(
                "way_cooler.on_view_request_resize(\"{}\", {}, {})",
                id, edge.bits(), point_to_lua(point)))
        }
    }

    /// Triggered just before a view is rendered.
    /// Used to render the borders.
    /// Sends the id of the view that will be rendered.
    fn view_pre_render(&self, view: WlcView) {
        Default.view_pre_render(view);
        if let Some(id) = lookup_handle(view.into()) {
            send_to_lua(format!("way_cooler.on_view_pre_render(\"{}\")", id))
        }
    }

    /// Triggered when a keyboard key is clicked.
    /// Sends the id of the view that it was sent to,
    /// the key that was sent, any modifiers that were sent,
    /// and the state the key was in (up or down).
    fn on_keyboard_key(&self,
                       view: WlcView,
                       time: u32,
                       mods: KeyboardModifiers,
                       key: u32,
                       state: KeyState) -> bool {
        let result = Default.on_keyboard_key(view, time, mods, key, state);
        if let Some(id) = lookup_handle(view.into()) {
            send_to_lua(format!(
                "way_cooler.on_keyboard_key(\"{id}\", {key}, {mods}, {state})",
                id = id,
                key = key,
                mods = mods.mods.bits(),
                state = state as i32))
        }
        result
    }

    /// Triggered when a user presses a button on the mouse.
    /// Sends the UUID of the view clicked, any modifiers that were held
    /// down, where absolutely on the screen it was clicked, and what button
    /// was used to Lua.
    fn on_pointer_button(&self,
                         view: WlcView,
                         time: u32,
                         mods: KeyboardModifiers,
                         button: u32,
                         state: ButtonState,
                         point: Point) -> bool {
        let result = Default.on_pointer_button(view, time, mods, button, state, point);
        if let Some(id) = lookup_handle(view.into()) {
            send_to_lua(format!("way_cooler.on_pointer_button(\
                                 \"{view}\", {mods}, {pos}, {button})",
                                view = id,
                                mods = mods.mods.bits(),
                                pos = point_to_lua(point),
                                button = button))
        }
        result
    }

    /// Triggered when the user scrolls using a mouse.
    /// Sends the UUID of the view clicked, any modifiers that were held down,
    /// the axis of the scroll wheel, and both heights.
    fn on_pointer_scroll(&self,
                         view: WlcView,
                         time: u32,
                         mods: KeyboardModifiers,
                         axis: ScrollAxis,
                         heights: [f64; 2]) -> bool {
        let result = Default.on_pointer_scroll(view, time, mods, axis, heights);
        if let Some(id) = lookup_handle(view.into()) {
            send_to_lua(format!(
                "way_cooler.on_pointer_scroll(\"{}\", {}, {}, {}, {})",
                id, mods.mods.bits(), axis as i32, heights[0], heights[1]))
        }
        result
    }

    /// Triggerd when a user moves the mouse.
    /// Sends the UUID of the view where it was moved over,
    /// and the absolute point on the screen where it was.
    fn on_pointer_motion(&self,
                         view: WlcView,
                         time: u32,
                         point: Point) -> bool {
        let result = Default.on_pointer_motion(view, time, point);
        if let Some(id) = lookup_handle(view.into()) {
            send_to_lua(format!(
                "way_cooler.on_pointer_motion(\"{}\", {})",
                id, point_to_lua(point)))
        }
        result
    }
}

/// Looks up the handle in the tree
fn lookup_handle(handle: Handle) -> Option<String> {
    if let Ok(tree) = try_lock_tree() {
        if let Ok(id) = tree.lookup_handle(handle.into()) {
            return Some(id.hyphenated().to_string())
        }
    }
    None
}

fn send_to_lua<Q: Into<String>>(msg: Q) {
    let msg = msg.into();
    match lua::send(LuaQuery::Execute(msg.clone())) {
        Ok(result) => {
            error!("Got back {:?}", result.recv())
        },
        Err(LuaSendError::Sender(err)) =>
            error!("Error while executing {:?}: {:?}", msg, err),
        Err(LuaSendError::ThreadUninitialized) =>
            error!("Thread was not initilazed yet, could not execute {:?}",
                   msg),
        Err(LuaSendError::ThreadClosed) => {
            error!("Thread closed, could not execute {:?}", msg)
        }
    }
}

fn size_to_lua(size: Size) -> String {
    format!("{}", size_to_json(size)).replace(":", "=").replace("\"", "")
}

fn point_to_lua(point: Point) -> String {
    format!("{}", point_to_json(point)).replace(":", "=").replace("\"", "")
}

fn geometry_to_lua(geometry: Geometry) -> String {
    format!("{}", geometry_to_json(geometry))
        .replace(":", "=").replace("\"", "")
}
<commit_msg>Fixed props_changed, removed error logging<commit_after>//! A custom Lua callback mapping, where the default operation is done and
//! then a custom, user-defined Lua callback is ran.
use rustwlc::*;
use rustwlc::Size;
use ::lua::{self, LuaQuery, LuaSendError};
use ::layout::{try_lock_tree, Handle};
use ::convert::json::{size_to_json, point_to_json, geometry_to_json};

use super::default::Default;
use super::Mode;

pub struct CustomLua;

impl Mode for CustomLua {
    /// Triggered when an output is connected.
    /// Sends the id of the created output to Lua.
    fn output_created(&self, output: WlcOutput) -> bool {
        let result = Default.output_created(output);
        if let Some(id) = lookup_handle(output.into()) {
            send_to_lua(format!("way_cooler.on_output_created(\"{}\")", id));
        }
        result
    }

    /// Triggered when an output is disconnected.
    /// Sends the id of the destroyed output to Lua.
    fn output_destroyed(&self, output: WlcOutput) {
        // Get the output ID before it's destroyed
        let id = lookup_handle(output.into());
        Default.output_destroyed(output);
        if let Some(id) = id {
            send_to_lua(format!("way_cooler.on_output_destroyed(\"{}\")", id))
        }
    }

    /// Triggered when an output loses or gains focus.
    /// Sends the id of the output, and whether it gained or lost focus.
    fn output_focused(&self, output: WlcOutput, focused: bool) {
        Default.output_focused(output, focused);
        if let Some(id) = lookup_handle(output.into()) {
            send_to_lua(format!("way_cooler.on_output_focused(\"{}\", {})",
                                id, focused));
        }
    }

    /// Triggered when an output's resolution changes.
    /// Sends the id of the output, and the old and new size of the output.
    fn output_resolution(&self, output: WlcOutput,
                         old_size: Size, new_size: Size) {
        Default.output_resolution(output, old_size, new_size);
        if let Some(id) = lookup_handle(output.into()) {
            send_to_lua(format!(
                "way_cooler.on_output_resolution_change(\"{}\", {}, {})",
                id, size_to_lua(old_size), size_to_lua(new_size)));
        }
    }

    /// Triggered after the output is rendered.
    /// Sends the id of the output after it's rendered.
    fn output_render_post(&self, output: WlcOutput) {
        Default.output_render_post(output);
        if let Some(id) = lookup_handle(output.into()) {
            send_to_lua(format!(
                "way_cooler.on_output_render_post(\"{}\")", id));
        }
    }

    /// Triggered when a view moves outputs.
    /// Sends the id of the view, the id of the old output, and the id of the
    /// the new output.
    fn view_moved_to_output(&self,
                            view: WlcView,
                            o1: WlcOutput,
                            o2: WlcOutput) {
        Default.view_moved_to_output(view, o1, o2);
        let (view_id, o1_id, o2_id) = (lookup_handle(view.into()),
                                       lookup_handle(o1.into()),
                                       lookup_handle(o2.into()));
        match (view_id, o1_id, o2_id) {
            (Some(view_id), Some(o1_id), Some(o2_id)) => {
                send_to_lua(format!(
                    "way_cooler.on_view_moved_to_output(\"{}\", \"{}\", \"{}\")",
                    view_id, o1_id, o2_id))
            },
            _ => {}
        }
    }

    /// Triggered when a view is created.
    /// Sends the id of the new view.
    fn view_created(&self, view: WlcView) -> bool {
        let result = Default.view_created(view);
        if let Some(id) = lookup_handle(view.into()) {
            send_to_lua(format!("way_cooler.on_view_created(\"{}\")", id))
        }
        result
    }

    /// Triggered when a view is destroyed.
    /// Sends the id of the view after its destroyed.
    fn view_destroyed(&self, view: WlcView) {
        // Get id before it's destroyed.
        let id = lookup_handle(view.into());
        Default.view_destroyed(view);
        if let Some(id) = id {
            send_to_lua(format!("way_cooler.on_view_destroyed(\"{}\")", id))
        }
    }

    /// Triggered when a view gains or loses focus.
    /// Sends the id of the view after its focus changes. The focus is true if
    /// it gained focus, or false if it lost focus.
    fn view_focused(&self, view: WlcView, focused: bool) {
        Default.view_focused(view, focused);
        if let Some(id) = lookup_handle(view.into()) {
            send_to_lua(format!("way_cooler.on_view_focused(\"{}\", {})",
                                id, focused))
        }
    }

    /// Triggered when the properties of a view changes.
    /// Sends the id of the view that had its properties changed, and the
    /// changed properties as a bit mask.
    fn view_props_changed(&self, view: WlcView, prop: ViewPropertyType) {
        Default.view_props_changed(view, prop);
        if let Some(id) = lookup_handle(view.into()) {
            send_to_lua(format!("way_cooler.on_view_props_changed(\"{}\", {})",
                                id, prop.bits()))
        }
    }

    /// Triggered when a view has its state change.
    /// Sends the id of the view, the bit field for the state field,
    /// and a bool representing how it was toggled.
    fn view_request_state(&self,
                          view: WlcView,
                          state: ViewState,
                          toggle: bool) {
        Default.view_request_state(view, state, toggle);
        if let Some(id) = lookup_handle(view.into()) {
            send_to_lua(format!(
                "way_cooler.on_view_request_state(\"{}\", {}, {})",
                id, state.bits(), toggle))
        }
    }

    /// Triggered when a view requests a different geometry.
    /// Sends the id of the view, and the new geometry.
    fn view_request_geometry(&self, view: WlcView, geometry: Geometry) {
        Default.view_request_geometry(view, geometry);
        if let Some(id) = lookup_handle(view.into()) {
            send_to_lua(format!(
                "way_cooler.on_view_request_geometry(\"{}\", {})",
                id, geometry_to_lua(geometry)))
        }
    }

    /// Triggered when a view requests to be moved, but not resized.
    /// Sends the id of the view, and the new destination point.
    fn view_request_move(&self, view: WlcView, dest: Point) {
        Default.view_request_move(view, dest);
        if let Some(id) = lookup_handle(view.into()) {
            send_to_lua(format!("way_cooler.on_view_request_move(\"{}\", {})",
                id, point_to_lua(dest)))
        }
    }

    /// Triggered when a view requests to be resized.
    /// Sends the id of the view, the bit field of the resize edge,
    /// and the point where the mouse resized.
    fn view_request_resize(&self,
                           view: WlcView,
                           edge: ResizeEdge,
                           point: Point) {
        Default.view_request_resize(view, edge, point);
        if let Some(id) = lookup_handle(view.into()) {
            send_to_lua(format!(
                "way_cooler.on_view_request_resize(\"{}\", {}, {})",
                id, edge.bits(), point_to_lua(point)))
        }
    }

    /// Triggered just before a view is rendered.
    /// Used to render the borders.
    /// Sends the id of the view that will be rendered.
    fn view_pre_render(&self, view: WlcView) {
        Default.view_pre_render(view);
        if let Some(id) = lookup_handle(view.into()) {
            send_to_lua(format!("way_cooler.on_view_pre_render(\"{}\")", id))
        }
    }

    /// Triggered when a keyboard key is clicked.
    /// Sends the id of the view that it was sent to,
    /// the key that was sent, any modifiers that were sent,
    /// and the state the key was in (up or down).
    fn on_keyboard_key(&self,
                       view: WlcView,
                       time: u32,
                       mods: KeyboardModifiers,
                       key: u32,
                       state: KeyState) -> bool {
        let result = Default.on_keyboard_key(view, time, mods, key, state);
        if let Some(id) = lookup_handle(view.into()) {
            send_to_lua(format!(
                "way_cooler.on_keyboard_key(\"{id}\", {key}, {mods}, {state})",
                id = id,
                key = key,
                mods = mods.mods.bits(),
                state = state as i32))
        }
        result
    }

    /// Triggered when a user presses a button on the mouse.
    /// Sends the UUID of the view clicked, any modifiers that were held
    /// down, where absolutely on the screen it was clicked, and what button
    /// was used to Lua.
    fn on_pointer_button(&self,
                         view: WlcView,
                         time: u32,
                         mods: KeyboardModifiers,
                         button: u32,
                         state: ButtonState,
                         point: Point) -> bool {
        let result = Default.on_pointer_button(view, time, mods, button, state, point);
        if let Some(id) = lookup_handle(view.into()) {
            send_to_lua(format!("way_cooler.on_pointer_button(\
                                 \"{view}\", {mods}, {pos}, {button})",
                                view = id,
                                mods = mods.mods.bits(),
                                pos = point_to_lua(point),
                                button = button))
        }
        result
    }

    /// Triggered when the user scrolls using a mouse.
    /// Sends the UUID of the view clicked, any modifiers that were held down,
    /// the axis of the scroll wheel, and both heights.
    fn on_pointer_scroll(&self,
                         view: WlcView,
                         time: u32,
                         mods: KeyboardModifiers,
                         axis: ScrollAxis,
                         heights: [f64; 2]) -> bool {
        let result = Default.on_pointer_scroll(view, time, mods, axis, heights);
        if let Some(id) = lookup_handle(view.into()) {
            send_to_lua(format!(
                "way_cooler.on_pointer_scroll(\"{}\", {}, {}, {}, {})",
                id, mods.mods.bits(), axis as i32, heights[0], heights[1]))
        }
        result
    }

    /// Triggerd when a user moves the mouse.
    /// Sends the UUID of the view where it was moved over,
    /// and the absolute point on the screen where it was.
    fn on_pointer_motion(&self,
                         view: WlcView,
                         time: u32,
                         point: Point) -> bool {
        let result = Default.on_pointer_motion(view, time, point);
        if let Some(id) = lookup_handle(view.into()) {
            send_to_lua(format!(
                "way_cooler.on_pointer_motion(\"{}\", {})",
                id, point_to_lua(point)))
        }
        result
    }
}

/// Looks up the handle in the tree
fn lookup_handle(handle: Handle) -> Option<String> {
    if let Ok(tree) = try_lock_tree() {
        if let Ok(id) = tree.lookup_handle(handle.into()) {
            return Some(id.hyphenated().to_string())
        }
    }
    None
}

fn send_to_lua<Q: Into<String>>(msg: Q) {
    let msg = msg.into();
    match lua::send(LuaQuery::Execute(msg.clone())) {
        Ok(result) => {},
        Err(LuaSendError::Sender(err)) =>
            error!("Error while executing {:?}: {:?}", msg, err),
        Err(LuaSendError::ThreadUninitialized) =>
            error!("Thread was not initilazed yet, could not execute {:?}",
                   msg),
        Err(LuaSendError::ThreadClosed) => {
            error!("Thread closed, could not execute {:?}", msg)
        }
    }
}

fn size_to_lua(size: Size) -> String {
    format!("{}", size_to_json(size)).replace(":", "=").replace("\"", "")
}

fn point_to_lua(point: Point) -> String {
    format!("{}", point_to_json(point)).replace(":", "=").replace("\"", "")
}

fn geometry_to_lua(geometry: Geometry) -> String {
    format!("{}", geometry_to_json(geometry))
        .replace(":", "=").replace("\"", "")
}
<|endoftext|>


<commit_before>extern crate pink_spider;

use pink_spider::model::Playlist;

pub fn main() {
    println!("Start crawling playlists...");

    let mut playlists = Playlist::find_actives();
    println!("{} playlists", playlists.len());
    for playlist in playlists.iter_mut() {
        println!("Fetch latest tracks of [{}] {}: {}",
                 playlist.provider,
                 playlist.id,
                 playlist.title);
        let tracks = match playlist.fetch_tracks() {
            Ok(tracks) => tracks,
            Err(e)     => {
                println!("{}", e);
                vec![]
            }
        };
        println!("{} tracks are added", tracks.len());
    }

    println!("Complete crawling playlists...");
}
<commit_msg>Better logging for playlist_crawler<commit_after>extern crate chrono;
extern crate pink_spider;

use std::time::Instant;
use chrono::Duration;
use pink_spider::model::Playlist;

pub fn main() {
    println!("Start crawling playlists...");
    let now = Instant::now();
    let mut playlists = Playlist::find_actives();
    println!("{} playlists", playlists.len());
    let total = playlists.len();
    let mut index = 0;
    for playlist in playlists.iter_mut() {
        println!("[{}/{}] Fetch latest tracks of [{}] {}: {}",
                 index,
                 total,
                 playlist.provider,
                 playlist.id,
                 playlist.title);
        let tracks = match playlist.fetch_tracks() {
            Ok(tracks) => tracks,
            Err(e)     => {
                println!("{}", e);
                vec![]
            }
        };
        println!("{} tracks are added", tracks.len());
        index += 1;
    }

    println!("Complete crawling playlists... total {} ms",
             Duration::from_std(now.elapsed()).unwrap().num_milliseconds());
}
<|endoftext|>


<commit_before>#![crate_type = "bin"]
#![feature(collections)]

extern crate gtk;
extern crate glib;
extern crate sysinfo;

use gtk::{WindowTrait, ContainerTrait, WidgetTrait};
use glib::{Type, timeout};
use sysinfo::*;
use gtk::signal::Inhibit;
use gtk::signal::{WidgetSignals, TreeViewSignals};
use std::collections::VecMap;
use std::str::FromStr;
use std::rc::Rc;
use std::cell::{RefCell, RefMut};
use std::ops::DerefMut;

fn append_column(title: &str, v: &mut Vec<gtk::TreeViewColumn>) {
    let l = v.len();
    let renderer = gtk::CellRendererText::new().unwrap();

    v.push(gtk::TreeViewColumn::new().unwrap());
    let tmp = v.get_mut(l).unwrap();

    tmp.set_title(title);
    tmp.pack_start(&renderer, true);
    tmp.add_attribute(&renderer, "text", l as i32);
}

fn create_and_fill_model(list_store: &mut gtk::ListStore, pid: i64, name: &str, cpu: f32, memory: u64) {
    if name.len() < 1 {
        return;
    }
    let mut top_level = gtk::TreeIter::new();

    list_store.append(&mut top_level);
    list_store.set_string(&top_level, 0, &format!("{}", pid));
    list_store.set_string(&top_level, 1, name);
    list_store.set_string(&top_level, 2, &format!("{}", cpu));
    list_store.set_string(&top_level, 3, &format!("{}", memory));
}

fn update_window(w: &mut (&mut gtk::ListStore, &mut sysinfo::System)) -> i32 {
    let mut model = w.0.get_model().unwrap();
    let mut iter = gtk::TreeIter::new();

    w.1.refresh();
    let mut entries : VecMap<Processus> = (*w.1.get_processus_list()).clone();
    let mut nb = model.iter_n_children(None);
    let mut i = 0;

    while i < nb {
        if model.iter_nth_child(&mut iter, None, i) {
            let pid : i64 = i64::from_str(&model.get_value(&iter, 0).get_string().unwrap()).unwrap();
            let mut to_delete = false;

            match entries.get(&(pid as usize)) {
                Some(p) => {
                    //w.0.set_string(&iter, 1, &p.name);
                    w.0.set_string(&iter, 2, &format!("{}", p.cpu_usage));
                    w.0.set_string(&iter, 3, &format!("{}", p.memory));
                    to_delete = true;
                }
                None => {
                    w.0.remove(&iter);
                }
            }
            if to_delete {
                entries.remove(&(pid as usize));
                nb = model.iter_n_children(None);
                i += 1;
            }
        } else {
            i += 1;
        }
    }
    for (_, pro) in entries {
        create_and_fill_model(&mut w.0, pro.pid, &pro.name, pro.cpu_usage, pro.memory);
    }
    1
}

fn main() {
    gtk::init();

    let mut window = gtk::Window::new(gtk::WindowType::TopLevel).unwrap();
    let mut sys = sysinfo::System::new();

    window.set_title("Processus viewer");
    window.set_window_position(gtk::WindowPosition::Center);

    window.connect_delete_event(|_, _| {
        gtk::main_quit();
        Inhibit(true)
    });

    let mut left_tree = gtk::TreeView::new().unwrap();
    let mut scroll = gtk::ScrolledWindow::new(None, None).unwrap();
    let mut current_pid : Rc<RefCell<Option<String>>> = Rc::new(RefCell::new(None));
    let mut current_pid1 = current_pid.clone();

    scroll.set_min_content_height(800);
    scroll.set_min_content_width(600);

    let mut columns : Vec<gtk::TreeViewColumn> = Vec::new();

    append_column("pid", &mut columns);
    append_column("process name", &mut columns);
    append_column("cpu usage", &mut columns);
    append_column("memory usage", &mut columns);

    for i in columns {
        left_tree.append_column(&i);
    }

    let mut list_store = gtk::ListStore::new(&[Type::String, Type::String, Type::String, Type::String]).unwrap();
    sys.refresh();
    for (_, pro) in sys.get_processus_list() {
        create_and_fill_model(&mut list_store, pro.pid, &pro.name, pro.cpu_usage, pro.memory);
    }

    left_tree.set_model(&list_store.get_model().unwrap());
    left_tree.set_headers_visible(true);
    scroll.add(&left_tree);
    let mut vertical_layout = gtk::Box::new(gtk::Orientation::Vertical, 0).unwrap();
    let mut kill_button = gtk::Button::new_with_label("End task").unwrap();

    left_tree.connect_cursor_changed(move |tree_view| {
        let mut path = gtk::TreePath::new().unwrap();
        tree_view.get_cursor(Some(&mut path), None);
        let mut iter = gtk::TreeIter::new();
        let model = tree_view.get_model().unwrap();

        if model.get_iter(&mut iter, &path) {
            let mut tmp = current_pid1.borrow_mut();
            *tmp = Some(model.get_value(&iter, 0).get_string().unwrap());
        }
        println!("current pid : {:?}", current_pid);
    });
    kill_button.set_sensitive(false);
    timeout::add(1500, update_window, &mut (&mut list_store, &mut sys));

    vertical_layout.add(&scroll);
    vertical_layout.add(&kill_button);
    window.add(&vertical_layout);
    window.show_all();
    gtk::main();
}
<commit_msg>Add signal handling<commit_after>#![crate_type = "bin"]
#![feature(collections)]

extern crate gtk;
extern crate glib;
extern crate sysinfo;

use gtk::{WindowTrait, ContainerTrait, WidgetTrait, ButtonSignals};
use glib::{Type, timeout};
use sysinfo::*;
use gtk::signal::Inhibit;
use gtk::signal::{WidgetSignals, TreeViewSignals};
use std::collections::VecMap;
use std::str::FromStr;
use std::rc::Rc;
use std::cell::{RefCell, RefMut};
use std::ops::DerefMut;

fn append_column(title: &str, v: &mut Vec<gtk::TreeViewColumn>) {
    let l = v.len();
    let renderer = gtk::CellRendererText::new().unwrap();

    v.push(gtk::TreeViewColumn::new().unwrap());
    let tmp = v.get_mut(l).unwrap();

    tmp.set_title(title);
    tmp.pack_start(&renderer, true);
    tmp.add_attribute(&renderer, "text", l as i32);
}

fn create_and_fill_model(list_store: &mut gtk::ListStore, pid: i64, name: &str, cpu: f32, memory: u64) {
    if name.len() < 1 {
        return;
    }
    let mut top_level = gtk::TreeIter::new();

    list_store.append(&mut top_level);
    list_store.set_string(&top_level, 0, &format!("{}", pid));
    list_store.set_string(&top_level, 1, name);
    list_store.set_string(&top_level, 2, &format!("{}", cpu));
    list_store.set_string(&top_level, 3, &format!("{}", memory));
}

fn update_window(w: &mut (&mut gtk::ListStore, Rc<RefCell<sysinfo::System>>)) -> i32 {
    let mut model = w.0.get_model().unwrap();
    let mut iter = gtk::TreeIter::new();

    (*w.1.borrow_mut()).refresh();
    let mut entries : VecMap<Processus> = ((*w.1.borrow()).get_processus_list()).clone();
    let mut nb = model.iter_n_children(None);
    let mut i = 0;

    while i < nb {
        if model.iter_nth_child(&mut iter, None, i) {
            let pid : i64 = i64::from_str(&model.get_value(&iter, 0).get_string().unwrap()).unwrap();
            let mut to_delete = false;

            match entries.get(&(pid as usize)) {
                Some(p) => {
                    //w.0.set_string(&iter, 1, &p.name);
                    w.0.set_string(&iter, 2, &format!("{}", p.cpu_usage));
                    w.0.set_string(&iter, 3, &format!("{}", p.memory));
                    to_delete = true;
                }
                None => {
                    w.0.remove(&iter);
                }
            }
            if to_delete {
                entries.remove(&(pid as usize));
                nb = model.iter_n_children(None);
                i += 1;
            }
        } else {
            i += 1;
        }
    }
    for (_, pro) in entries {
        create_and_fill_model(&mut w.0, pro.pid, &pro.name, pro.cpu_usage, pro.memory);
    }
    1
}

fn main() {
    gtk::init();

    let mut window = gtk::Window::new(gtk::WindowType::TopLevel).unwrap();
    let mut sys : Rc<RefCell<sysinfo::System>> = Rc::new(RefCell::new(sysinfo::System::new()));
    let mut sys1 = sys.clone();

    window.set_title("Processus viewer");
    window.set_window_position(gtk::WindowPosition::Center);

    window.connect_delete_event(|_, _| {
        gtk::main_quit();
        Inhibit(true)
    });

    let mut left_tree = gtk::TreeView::new().unwrap();
    let mut scroll = gtk::ScrolledWindow::new(None, None).unwrap();
    let mut current_pid : Rc<RefCell<Option<String>>> = Rc::new(RefCell::new(None));
    let mut current_pid1 = current_pid.clone();
    let mut current_pid2 = current_pid.clone();
    let mut kill_button : Rc<RefCell<gtk::Button>> = Rc::new(RefCell::new(gtk::Button::new_with_label("End task").unwrap()));
    let mut kill_button1 = kill_button.clone();

    scroll.set_min_content_height(800);
    scroll.set_min_content_width(600);

    let mut columns : Vec<gtk::TreeViewColumn> = Vec::new();

    append_column("pid", &mut columns);
    append_column("process name", &mut columns);
    append_column("cpu usage", &mut columns);
    append_column("memory usage", &mut columns);

    for i in columns {
        left_tree.append_column(&i);
    }

    let mut list_store = gtk::ListStore::new(&[Type::String, Type::String, Type::String, Type::String]).unwrap();
    (*sys.borrow_mut()).refresh();
    for (_, pro) in (*sys.borrow()).get_processus_list() {
        create_and_fill_model(&mut list_store, pro.pid, &pro.name, pro.cpu_usage, pro.memory);
    }

    left_tree.set_model(&list_store.get_model().unwrap());
    left_tree.set_headers_visible(true);
    scroll.add(&left_tree);
    let mut vertical_layout = gtk::Box::new(gtk::Orientation::Vertical, 0).unwrap();

    left_tree.connect_cursor_changed(move |tree_view| {
        match tree_view.get_selection() {
            Some(selection) => {
                let model = tree_view.get_model().unwrap();
                let mut iter = gtk::TreeIter::new();

                if selection.get_selected(&model, &mut iter) {
                    let pid = model.get_value(&iter, 0).get_string();
                    let mut tmp = current_pid1.borrow_mut();
                    *tmp = pid;
                }
                (*kill_button1.borrow_mut()).set_sensitive((*current_pid.borrow()).is_some());
            }
            None => {
                let mut tmp = current_pid1.borrow_mut();
                *tmp = None;
            }
        }
        /*let mut path = gtk::TreePath::new().unwrap();
        println!("{:?}", path.unwrap_pointer());
        tree_view.get_cursor(Some(&mut path), None);
        println!("{:?}", path.unwrap_pointer());
        let model = tree_view.get_model().unwrap();

        if model.get_iter(&mut iter, &path) {
            let mut tmp = current_pid1.borrow_mut();
            *tmp = Some(model.get_value(&iter, 0).get_string().unwrap());
        }
        println!("current pid : {:?}", current_pid);*/
    });
    (*kill_button.borrow_mut()).set_sensitive(false);
    (*kill_button.borrow_mut()).connect_clicked(move |_| {
        let tmp = (*current_pid2.borrow_mut()).is_some() ;

        if tmp {
            let s = (*current_pid2.borrow()).clone();
            match (*sys.borrow_mut()).get_processus(i64::from_str(&s.unwrap()).unwrap()) {
                Some(p) => {
                    p.kill(Signal::Kill);
                },
                None => {}
            };
        }
    });
    timeout::add(1500, update_window, &mut (&mut list_store, sys1.clone()));

    vertical_layout.add(&scroll);
    vertical_layout.add(&(*kill_button.borrow_mut()));
    window.add(&vertical_layout);
    window.show_all();
    gtk::main();
}
<|endoftext|>


<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/. */

use dom::bindings::utils::{Reflectable, Reflector, reflect_dom_object};
use dom::bindings::error::{Fallible};
use dom::bindings::codegen::BindingDeclarations::FormDataBinding;
use dom::bindings::js::{JS, JSRef, Temporary, OptionalUnrootable};
use dom::blob::Blob;
use dom::htmlformelement::HTMLFormElement;
use dom::window::Window;
use servo_util::str::DOMString;

use collections::hashmap::HashMap;

#[deriving(Encodable)]
pub enum FormDatum {
    StringData(DOMString),
    BlobData { blob: JS<Blob>, name: DOMString }
}

#[deriving(Encodable)]
pub struct FormData {
    pub data: HashMap<DOMString, FormDatum>,
    pub reflector_: Reflector,
    pub window: JS<Window>,
    pub form: Option<JS<HTMLFormElement>>
}

impl FormData {
    pub fn new_inherited(form: Option<JSRef<HTMLFormElement>>, window: &JSRef<Window>) -> FormData {
        FormData {
            data: HashMap::new(),
            reflector_: Reflector::new(),
            window: window.unrooted(),
            form: form.unrooted(),
        }
    }

    pub fn new(form: Option<JSRef<HTMLFormElement>>, window: &JSRef<Window>) -> Temporary<FormData> {
        reflect_dom_object(box FormData::new_inherited(form, window), window, FormDataBinding::Wrap)
    }

    pub fn Constructor(window: &JSRef<Window>, form: Option<JSRef<HTMLFormElement>>) -> Fallible<Temporary<FormData>> {
        Ok(FormData::new(form, window))
    }
}

pub trait FormDataMethods {
    fn Append(&mut self, name: DOMString, value: &JSRef<Blob>, filename: Option<DOMString>);
    fn Append_(&mut self, name: DOMString, value: DOMString);
}

impl<'a> FormDataMethods for JSRef<'a, FormData> {
    fn Append(&mut self, name: DOMString, value: &JSRef<Blob>, filename: Option<DOMString>) {
        let blob = BlobData {
            blob: value.unrooted(),
            name: filename.unwrap_or("default".to_owned())
        };
        self.data.insert(name.clone(), blob);
    }

    fn Append_(&mut self, name: DOMString, value: DOMString) {
        self.data.insert(name, StringData(value));
    }
}

impl Reflectable for FormData {
    fn reflector<'a>(&'a self) -> &'a Reflector {
        &self.reflector_
    }

    fn mut_reflector<'a>(&'a mut self) -> &'a mut Reflector {
        &mut self.reflector_
    }
}
<commit_msg>Use Cell/RefCell for interior mutability of FormData, BlobData.<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/. */

use dom::bindings::utils::{Reflectable, Reflector, reflect_dom_object};
use dom::bindings::error::{Fallible};
use dom::bindings::codegen::BindingDeclarations::FormDataBinding;
use dom::bindings::js::{JS, JSRef, Temporary, OptionalUnrootable};
use dom::blob::Blob;
use dom::htmlformelement::HTMLFormElement;
use dom::window::Window;
use servo_util::str::DOMString;
use collections::hashmap::HashMap;

#[deriving(Encodable)]
pub enum FormDatum {
    StringData(DOMString),
    BlobData { blob: JS<Blob>, name: DOMString }
}

#[deriving(Encodable)]
pub struct FormData {
    pub data: HashMap<DOMString, FormDatum>,
    pub reflector_: Reflector,
    pub window: JS<Window>,
    pub form: Option<JS<HTMLFormElement>>
}

impl FormData {
    pub fn new_inherited(form: Option<JSRef<HTMLFormElement>>, window: &JSRef<Window>) -> FormData {
        FormData {
            data: HashMap::new(),
            reflector_: Reflector::new(),
            window: window.unrooted(),
            form: form.unrooted(),
        }
    }

    pub fn new(form: Option<JSRef<HTMLFormElement>>, window: &JSRef<Window>) -> Temporary<FormData> {
        reflect_dom_object(box FormData::new_inherited(form, window), window, FormDataBinding::Wrap)
    }

    pub fn Constructor(window: &JSRef<Window>, form: Option<JSRef<HTMLFormElement>>) -> Fallible<Temporary<FormData>> {
        Ok(FormData::new(form, window))
    }
}

pub trait FormDataMethods {
    fn Append(&mut self, name: DOMString, value: &JSRef<Blob>, filename: Option<DOMString>);
    fn Append_(&mut self, name: DOMString, value: DOMString);
}

impl<'a> FormDataMethods for JSRef<'a, FormData> {
    fn Append(&mut self, name: DOMString, value: &JSRef<Blob>, filename: Option<DOMString>) {
        let blob = BlobData {
            blob: value.unrooted(),
            name: filename.unwrap_or("default".to_owned())
        };
        self.data.insert(name.clone(), blob);
    }

    fn Append_(&mut self, name: DOMString, value: DOMString) {
        self.data.insert(name, StringData(value));
    }
}

impl Reflectable for FormData {
    fn reflector<'a>(&'a self) -> &'a Reflector {
        &self.reflector_
    }

    fn mut_reflector<'a>(&'a mut self) -> &'a mut Reflector {
        &mut self.reflector_
    }
}
<|endoftext|>


<commit_before>use std::error::Error;

use diesel::dsl::sql;
use diesel::*;

use super::data_structures::*;
use super::table_data::TableName;
use crate::print_schema::ColumnSorting;

table! {
    sqlite_master (name) {
        name -> VarChar,
    }
}

table! {
    pragma_table_info (cid) {
        cid ->Integer,
        name -> VarChar,
        type_name -> VarChar,
        notnull -> Bool,
        dflt_value -> Nullable<VarChar>,
        pk -> Bool,
        hidden -> Integer,
    }
}

table! {
    pragma_foreign_key_list {
        id -> Integer,
        seq -> Integer,
        _table -> VarChar,
        from -> VarChar,
        to -> VarChar,
        on_update -> VarChar,
        on_delete -> VarChar,
        _match -> VarChar,
    }
}

pub fn load_table_names(
    connection: &mut SqliteConnection,
    schema_name: Option<&str>,
) -> Result<Vec<TableName>, Box<dyn Error + Send + Sync + 'static>> {
    use self::sqlite_master::dsl::*;

    if schema_name.is_some() {
        return Err("sqlite cannot infer schema for databases other than the \
                    main database"
            .into());
    }

    Ok(sqlite_master
        .select(name)
        .filter(name.not_like("\\_\\_%").escape('\\'))
        .filter(name.not_like("sqlite%"))
        .filter(sql::<sql_types::Bool>("type='table'"))
        .order(name)
        .load::<String>(connection)?
        .into_iter()
        .map(TableName::from_name)
        .collect())
}

pub fn load_foreign_key_constraints(
    connection: &mut SqliteConnection,
    schema_name: Option<&str>,
) -> Result<Vec<ForeignKeyConstraint>, Box<dyn Error + Send + Sync + 'static>> {
    let tables = load_table_names(connection, schema_name)?;
    let rows = tables
        .into_iter()
        .map(|child_table| {
            let query = format!("PRAGMA FOREIGN_KEY_LIST('{}')", child_table.sql_name);
            Ok(sql::<pragma_foreign_key_list::SqlType>(&query)
                .load::<ForeignKeyListRow>(connection)?
                .into_iter()
                .map(|row| {
                    let parent_table = TableName::from_name(row.parent_table);
                    ForeignKeyConstraint {
                        child_table: child_table.clone(),
                        parent_table,
                        foreign_key: row.foreign_key.clone(),
                        foreign_key_rust_name: row.foreign_key,
                        primary_key: row.primary_key,
                    }
                })
                .collect())
        })
        .collect::<QueryResult<Vec<Vec<_>>>>()?;
    Ok(rows.into_iter().flatten().collect())
}

#[derive(PartialEq, Eq, PartialOrd, Ord)]
struct SqliteVersion {
    major: u32,
    minor: u32,
    patch: u32,
}

impl SqliteVersion {
    pub fn new(major: u32, minor: u32, patch: u32) -> SqliteVersion {
        SqliteVersion {
            major,
            minor,
            patch,
        }
    }
}

fn get_sqlite_version(
    conn: &mut SqliteConnection
) -> SqliteVersion {
    let query = "SELECT sqlite_version()";
    let result = sql::<sql_types::Text>(&query).load::<String>(conn).unwrap();
    let parts = result[0]
        .split('.')
        .map(|part| { part.parse().unwrap() })
        .collect::<Vec<u32>>();
    assert_eq!(parts.len(), 3);
    SqliteVersion::new(
        parts[0],
        parts[1],
        parts[2],
    )
}

pub fn get_table_data(
    conn: &mut SqliteConnection,
    table: &TableName,
    column_sorting: &ColumnSorting,
) -> QueryResult<Vec<ColumnInformation>> {
    let sqlite_version = get_sqlite_version(conn);
    let query = if sqlite_version >= SqliteVersion::new(3, 26, 0) {
        format!("PRAGMA TABLE_XINFO('{}')", &table.sql_name)
    } else {
        format!("PRAGMA TABLE_INFO('{}')", &table.sql_name)
    };
    let mut result = sql::<pragma_table_info::SqlType>(&query).load(conn)?;
    match column_sorting {
        ColumnSorting::OrdinalPosition => {}
        ColumnSorting::Name => {
            result.sort_by(|a: &ColumnInformation, b: &ColumnInformation| {
                a.column_name.partial_cmp(&b.column_name).unwrap()
            });
        }
    };
    Ok(result)
}

#[derive(Queryable)]
struct FullTableInfo {
    _cid: i32,
    name: String,
    _type_name: String,
    _not_null: bool,
    _dflt_value: Option<String>,
    primary_key: bool,
    _hidden: i32,
}

#[derive(Queryable)]
struct ForeignKeyListRow {
    _id: i32,
    _seq: i32,
    parent_table: String,
    foreign_key: String,
    primary_key: String,
    _on_update: String,
    _on_delete: String,
    _match: String,
}

pub fn get_primary_keys(
    conn: &mut SqliteConnection,
    table: &TableName,
) -> QueryResult<Vec<String>> {
    let sqlite_version = get_sqlite_version(conn);
    let query = if sqlite_version >= SqliteVersion::new(3, 26, 0) {
        format!("PRAGMA TABLE_XINFO('{}')", &table.sql_name)
    } else {
        format!("PRAGMA TABLE_INFO('{}')", &table.sql_name)
    };
    let results = sql::<pragma_table_info::SqlType>(&query).load::<FullTableInfo>(conn)?;
    Ok(results
        .into_iter()
        .filter_map(|i| if i.primary_key { Some(i.name) } else { None })
        .collect())
}

pub fn determine_column_type(
    attr: &ColumnInformation,
) -> Result<ColumnType, Box<dyn Error + Send + Sync + 'static>> {
    let mut type_name = attr.type_name.to_lowercase();
    if type_name == "generated always" {
        type_name.clear();
    }

    let path = if is_bool(&type_name) {
        String::from("Bool")
    } else if is_smallint(&type_name) {
        String::from("SmallInt")
    } else if is_bigint(&type_name) {
        String::from("BigInt")
    } else if type_name.contains("int") {
        String::from("Integer")
    } else if is_text(&type_name) {
        String::from("Text")
    } else if is_binary(&type_name) {
        String::from("Binary")
    } else if is_float(&type_name) {
        String::from("Float")
    } else if is_double(&type_name) {
        String::from("Double")
    } else if type_name == "datetime" || type_name == "timestamp" {
        String::from("Timestamp")
    } else if type_name == "date" {
        String::from("Date")
    } else if type_name == "time" {
        String::from("Time")
    } else {
        return Err(format!("Unsupported type: {}", type_name).into());
    };

    Ok(ColumnType {
        schema: None,
        rust_name: path.clone(),
        sql_name: path,
        is_array: false,
        is_nullable: attr.nullable,
        is_unsigned: false,
    })
}

fn is_text(type_name: &str) -> bool {
    type_name.contains("char") || type_name.contains("clob") || type_name.contains("text")
}

fn is_binary(type_name: &str) -> bool {
    type_name.contains("blob") || type_name.contains("binary") || type_name.is_empty()
}

fn is_bool(type_name: &str) -> bool {
    type_name == "boolean" || type_name.contains("tiny") && type_name.contains("int")
}

fn is_smallint(type_name: &str) -> bool {
    type_name == "int2" || type_name.contains("small") && type_name.contains("int")
}

fn is_bigint(type_name: &str) -> bool {
    type_name == "int8" || type_name.contains("big") && type_name.contains("int")
}

fn is_float(type_name: &str) -> bool {
    type_name.contains("float") || type_name.contains("real")
}

fn is_double(type_name: &str) -> bool {
    type_name.contains("double") || type_name.contains("num") || type_name.contains("dec")
}

#[test]
fn load_table_names_returns_nothing_when_no_tables_exist() {
    let mut conn = SqliteConnection::establish(":memory:").unwrap();
    assert_eq!(
        Vec::<TableName>::new(),
        load_table_names(&mut conn, None).unwrap()
    );
}

#[test]
fn load_table_names_includes_tables_that_exist() {
    let mut conn = SqliteConnection::establish(":memory:").unwrap();
    conn.execute("CREATE TABLE users (id INTEGER PRIMARY KEY AUTOINCREMENT)")
        .unwrap();
    let table_names = load_table_names(&mut conn, None).unwrap();
    assert!(table_names.contains(&TableName::from_name("users")));
}

#[test]
fn load_table_names_excludes_diesel_metadata_tables() {
    let mut conn = SqliteConnection::establish(":memory:").unwrap();
    conn.execute("CREATE TABLE __diesel_metadata (id INTEGER PRIMARY KEY AUTOINCREMENT)")
        .unwrap();
    let table_names = load_table_names(&mut conn, None).unwrap();
    assert!(!table_names.contains(&TableName::from_name("__diesel_metadata")));
}

#[test]
fn load_table_names_excludes_sqlite_metadata_tables() {
    let mut conn = SqliteConnection::establish(":memory:").unwrap();
    conn.execute("CREATE TABLE __diesel_metadata (id INTEGER PRIMARY KEY AUTOINCREMENT)")
        .unwrap();
    conn.execute("CREATE TABLE users (id INTEGER PRIMARY KEY AUTOINCREMENT)")
        .unwrap();
    let table_names = load_table_names(&mut conn, None);
    assert_eq!(vec![TableName::from_name("users")], table_names.unwrap());
}

#[test]
fn load_table_names_excludes_views() {
    let mut conn = SqliteConnection::establish(":memory:").unwrap();
    conn.execute("CREATE TABLE users (id INTEGER PRIMARY KEY AUTOINCREMENT)")
        .unwrap();
    conn.execute("CREATE VIEW answer AS SELECT 42").unwrap();
    let table_names = load_table_names(&mut conn, None);
    assert_eq!(vec![TableName::from_name("users")], table_names.unwrap());
}

#[test]
fn load_table_names_returns_error_when_given_schema_name() {
    let mut conn = SqliteConnection::establish(":memory:").unwrap();
    // We're testing the error message rather than using #[should_panic]
    // to ensure this is returning an error and not actually panicking.
    let table_names = load_table_names(&mut conn, Some("stuff"));
    match table_names {
        Ok(_) => panic!("Expected load_table_names to return an error"),
        Err(e) => assert!(e.to_string().starts_with(
            "sqlite cannot infer \
             schema for databases"
        )),
    }
}

#[test]
fn load_table_names_output_is_ordered() {
    let mut conn = SqliteConnection::establish(":memory:").unwrap();
    conn.execute("CREATE TABLE bbb (id INTEGER PRIMARY KEY AUTOINCREMENT)")
        .unwrap();
    conn.execute("CREATE TABLE aaa (id INTEGER PRIMARY KEY AUTOINCREMENT)")
        .unwrap();
    conn.execute("CREATE TABLE ccc (id INTEGER PRIMARY KEY AUTOINCREMENT)")
        .unwrap();

    let table_names = load_table_names(&mut conn, None)
        .unwrap()
        .iter()
        .map(|table| table.to_string())
        .collect::<Vec<_>>();
    assert_eq!(vec!["aaa", "bbb", "ccc"], table_names);
}

#[test]
fn load_foreign_key_constraints_loads_foreign_keys() {
    let mut connection = SqliteConnection::establish(":memory:").unwrap();

    connection.execute("CREATE TABLE table_1 (id)").unwrap();
    connection
        .execute("CREATE TABLE table_2 (id, fk_one REFERENCES table_1(id))")
        .unwrap();
    connection
        .execute("CREATE TABLE table_3 (id, fk_two REFERENCES table_2(id))")
        .unwrap();

    let table_1 = TableName::from_name("table_1");
    let table_2 = TableName::from_name("table_2");
    let table_3 = TableName::from_name("table_3");
    let fk_one = ForeignKeyConstraint {
        child_table: table_2.clone(),
        parent_table: table_1.clone(),
        foreign_key: "fk_one".into(),
        foreign_key_rust_name: "fk_one".into(),
        primary_key: "id".into(),
    };
    let fk_two = ForeignKeyConstraint {
        child_table: table_3.clone(),
        parent_table: table_2.clone(),
        foreign_key: "fk_two".into(),
        foreign_key_rust_name: "fk_two".into(),
        primary_key: "id".into(),
    };
    let fks = load_foreign_key_constraints(&mut connection, None).unwrap();
    assert_eq!(vec![fk_one, fk_two], fks);
}
<commit_msg>Formatting<commit_after>use std::error::Error;

use diesel::dsl::sql;
use diesel::*;

use super::data_structures::*;
use super::table_data::TableName;
use crate::print_schema::ColumnSorting;

table! {
    sqlite_master (name) {
        name -> VarChar,
    }
}

table! {
    pragma_table_info (cid) {
        cid ->Integer,
        name -> VarChar,
        type_name -> VarChar,
        notnull -> Bool,
        dflt_value -> Nullable<VarChar>,
        pk -> Bool,
        hidden -> Integer,
    }
}

table! {
    pragma_foreign_key_list {
        id -> Integer,
        seq -> Integer,
        _table -> VarChar,
        from -> VarChar,
        to -> VarChar,
        on_update -> VarChar,
        on_delete -> VarChar,
        _match -> VarChar,
    }
}

pub fn load_table_names(
    connection: &mut SqliteConnection,
    schema_name: Option<&str>,
) -> Result<Vec<TableName>, Box<dyn Error + Send + Sync + 'static>> {
    use self::sqlite_master::dsl::*;

    if schema_name.is_some() {
        return Err("sqlite cannot infer schema for databases other than the \
                    main database"
            .into());
    }

    Ok(sqlite_master
        .select(name)
        .filter(name.not_like("\\_\\_%").escape('\\'))
        .filter(name.not_like("sqlite%"))
        .filter(sql::<sql_types::Bool>("type='table'"))
        .order(name)
        .load::<String>(connection)?
        .into_iter()
        .map(TableName::from_name)
        .collect())
}

pub fn load_foreign_key_constraints(
    connection: &mut SqliteConnection,
    schema_name: Option<&str>,
) -> Result<Vec<ForeignKeyConstraint>, Box<dyn Error + Send + Sync + 'static>> {
    let tables = load_table_names(connection, schema_name)?;
    let rows = tables
        .into_iter()
        .map(|child_table| {
            let query = format!("PRAGMA FOREIGN_KEY_LIST('{}')", child_table.sql_name);
            Ok(sql::<pragma_foreign_key_list::SqlType>(&query)
                .load::<ForeignKeyListRow>(connection)?
                .into_iter()
                .map(|row| {
                    let parent_table = TableName::from_name(row.parent_table);
                    ForeignKeyConstraint {
                        child_table: child_table.clone(),
                        parent_table,
                        foreign_key: row.foreign_key.clone(),
                        foreign_key_rust_name: row.foreign_key,
                        primary_key: row.primary_key,
                    }
                })
                .collect())
        })
        .collect::<QueryResult<Vec<Vec<_>>>>()?;
    Ok(rows.into_iter().flatten().collect())
}

#[derive(PartialEq, Eq, PartialOrd, Ord)]
struct SqliteVersion {
    major: u32,
    minor: u32,
    patch: u32,
}

impl SqliteVersion {
    pub fn new(major: u32, minor: u32, patch: u32) -> SqliteVersion {
        SqliteVersion {
            major,
            minor,
            patch,
        }
    }
}

fn get_sqlite_version(conn: &mut SqliteConnection) -> SqliteVersion {
    let query = "SELECT sqlite_version()";
    let result = sql::<sql_types::Text>(&query).load::<String>(conn).unwrap();
    let parts = result[0]
        .split('.')
        .map(|part| part.parse().unwrap())
        .collect::<Vec<u32>>();
    assert_eq!(parts.len(), 3);
    SqliteVersion::new(parts[0], parts[1], parts[2])
}

pub fn get_table_data(
    conn: &mut SqliteConnection,
    table: &TableName,
    column_sorting: &ColumnSorting,
) -> QueryResult<Vec<ColumnInformation>> {
    let sqlite_version = get_sqlite_version(conn);
    let query = if sqlite_version >= SqliteVersion::new(3, 26, 0) {
        format!("PRAGMA TABLE_XINFO('{}')", &table.sql_name)
    } else {
        format!("PRAGMA TABLE_INFO('{}')", &table.sql_name)
    };
    let mut result = sql::<pragma_table_info::SqlType>(&query).load(conn)?;
    match column_sorting {
        ColumnSorting::OrdinalPosition => {}
        ColumnSorting::Name => {
            result.sort_by(|a: &ColumnInformation, b: &ColumnInformation| {
                a.column_name.partial_cmp(&b.column_name).unwrap()
            });
        }
    };
    Ok(result)
}

#[derive(Queryable)]
struct FullTableInfo {
    _cid: i32,
    name: String,
    _type_name: String,
    _not_null: bool,
    _dflt_value: Option<String>,
    primary_key: bool,
    _hidden: i32,
}

#[derive(Queryable)]
struct ForeignKeyListRow {
    _id: i32,
    _seq: i32,
    parent_table: String,
    foreign_key: String,
    primary_key: String,
    _on_update: String,
    _on_delete: String,
    _match: String,
}

pub fn get_primary_keys(
    conn: &mut SqliteConnection,
    table: &TableName,
) -> QueryResult<Vec<String>> {
    let sqlite_version = get_sqlite_version(conn);
    let query = if sqlite_version >= SqliteVersion::new(3, 26, 0) {
        format!("PRAGMA TABLE_XINFO('{}')", &table.sql_name)
    } else {
        format!("PRAGMA TABLE_INFO('{}')", &table.sql_name)
    };
    let results = sql::<pragma_table_info::SqlType>(&query).load::<FullTableInfo>(conn)?;
    Ok(results
        .into_iter()
        .filter_map(|i| if i.primary_key { Some(i.name) } else { None })
        .collect())
}

pub fn determine_column_type(
    attr: &ColumnInformation,
) -> Result<ColumnType, Box<dyn Error + Send + Sync + 'static>> {
    let mut type_name = attr.type_name.to_lowercase();
    if type_name == "generated always" {
        type_name.clear();
    }

    let path = if is_bool(&type_name) {
        String::from("Bool")
    } else if is_smallint(&type_name) {
        String::from("SmallInt")
    } else if is_bigint(&type_name) {
        String::from("BigInt")
    } else if type_name.contains("int") {
        String::from("Integer")
    } else if is_text(&type_name) {
        String::from("Text")
    } else if is_binary(&type_name) {
        String::from("Binary")
    } else if is_float(&type_name) {
        String::from("Float")
    } else if is_double(&type_name) {
        String::from("Double")
    } else if type_name == "datetime" || type_name == "timestamp" {
        String::from("Timestamp")
    } else if type_name == "date" {
        String::from("Date")
    } else if type_name == "time" {
        String::from("Time")
    } else {
        return Err(format!("Unsupported type: {}", type_name).into());
    };

    Ok(ColumnType {
        schema: None,
        rust_name: path.clone(),
        sql_name: path,
        is_array: false,
        is_nullable: attr.nullable,
        is_unsigned: false,
    })
}

fn is_text(type_name: &str) -> bool {
    type_name.contains("char") || type_name.contains("clob") || type_name.contains("text")
}

fn is_binary(type_name: &str) -> bool {
    type_name.contains("blob") || type_name.contains("binary") || type_name.is_empty()
}

fn is_bool(type_name: &str) -> bool {
    type_name == "boolean" || type_name.contains("tiny") && type_name.contains("int")
}

fn is_smallint(type_name: &str) -> bool {
    type_name == "int2" || type_name.contains("small") && type_name.contains("int")
}

fn is_bigint(type_name: &str) -> bool {
    type_name == "int8" || type_name.contains("big") && type_name.contains("int")
}

fn is_float(type_name: &str) -> bool {
    type_name.contains("float") || type_name.contains("real")
}

fn is_double(type_name: &str) -> bool {
    type_name.contains("double") || type_name.contains("num") || type_name.contains("dec")
}

#[test]
fn load_table_names_returns_nothing_when_no_tables_exist() {
    let mut conn = SqliteConnection::establish(":memory:").unwrap();
    assert_eq!(
        Vec::<TableName>::new(),
        load_table_names(&mut conn, None).unwrap()
    );
}

#[test]
fn load_table_names_includes_tables_that_exist() {
    let mut conn = SqliteConnection::establish(":memory:").unwrap();
    conn.execute("CREATE TABLE users (id INTEGER PRIMARY KEY AUTOINCREMENT)")
        .unwrap();
    let table_names = load_table_names(&mut conn, None).unwrap();
    assert!(table_names.contains(&TableName::from_name("users")));
}

#[test]
fn load_table_names_excludes_diesel_metadata_tables() {
    let mut conn = SqliteConnection::establish(":memory:").unwrap();
    conn.execute("CREATE TABLE __diesel_metadata (id INTEGER PRIMARY KEY AUTOINCREMENT)")
        .unwrap();
    let table_names = load_table_names(&mut conn, None).unwrap();
    assert!(!table_names.contains(&TableName::from_name("__diesel_metadata")));
}

#[test]
fn load_table_names_excludes_sqlite_metadata_tables() {
    let mut conn = SqliteConnection::establish(":memory:").unwrap();
    conn.execute("CREATE TABLE __diesel_metadata (id INTEGER PRIMARY KEY AUTOINCREMENT)")
        .unwrap();
    conn.execute("CREATE TABLE users (id INTEGER PRIMARY KEY AUTOINCREMENT)")
        .unwrap();
    let table_names = load_table_names(&mut conn, None);
    assert_eq!(vec![TableName::from_name("users")], table_names.unwrap());
}

#[test]
fn load_table_names_excludes_views() {
    let mut conn = SqliteConnection::establish(":memory:").unwrap();
    conn.execute("CREATE TABLE users (id INTEGER PRIMARY KEY AUTOINCREMENT)")
        .unwrap();
    conn.execute("CREATE VIEW answer AS SELECT 42").unwrap();
    let table_names = load_table_names(&mut conn, None);
    assert_eq!(vec![TableName::from_name("users")], table_names.unwrap());
}

#[test]
fn load_table_names_returns_error_when_given_schema_name() {
    let mut conn = SqliteConnection::establish(":memory:").unwrap();
    // We're testing the error message rather than using #[should_panic]
    // to ensure this is returning an error and not actually panicking.
    let table_names = load_table_names(&mut conn, Some("stuff"));
    match table_names {
        Ok(_) => panic!("Expected load_table_names to return an error"),
        Err(e) => assert!(e.to_string().starts_with(
            "sqlite cannot infer \
             schema for databases"
        )),
    }
}

#[test]
fn load_table_names_output_is_ordered() {
    let mut conn = SqliteConnection::establish(":memory:").unwrap();
    conn.execute("CREATE TABLE bbb (id INTEGER PRIMARY KEY AUTOINCREMENT)")
        .unwrap();
    conn.execute("CREATE TABLE aaa (id INTEGER PRIMARY KEY AUTOINCREMENT)")
        .unwrap();
    conn.execute("CREATE TABLE ccc (id INTEGER PRIMARY KEY AUTOINCREMENT)")
        .unwrap();

    let table_names = load_table_names(&mut conn, None)
        .unwrap()
        .iter()
        .map(|table| table.to_string())
        .collect::<Vec<_>>();
    assert_eq!(vec!["aaa", "bbb", "ccc"], table_names);
}

#[test]
fn load_foreign_key_constraints_loads_foreign_keys() {
    let mut connection = SqliteConnection::establish(":memory:").unwrap();

    connection.execute("CREATE TABLE table_1 (id)").unwrap();
    connection
        .execute("CREATE TABLE table_2 (id, fk_one REFERENCES table_1(id))")
        .unwrap();
    connection
        .execute("CREATE TABLE table_3 (id, fk_two REFERENCES table_2(id))")
        .unwrap();

    let table_1 = TableName::from_name("table_1");
    let table_2 = TableName::from_name("table_2");
    let table_3 = TableName::from_name("table_3");
    let fk_one = ForeignKeyConstraint {
        child_table: table_2.clone(),
        parent_table: table_1.clone(),
        foreign_key: "fk_one".into(),
        foreign_key_rust_name: "fk_one".into(),
        primary_key: "id".into(),
    };
    let fk_two = ForeignKeyConstraint {
        child_table: table_3.clone(),
        parent_table: table_2.clone(),
        foreign_key: "fk_two".into(),
        foreign_key_rust_name: "fk_two".into(),
        primary_key: "id".into(),
    };
    let fks = load_foreign_key_constraints(&mut connection, None).unwrap();
    assert_eq!(vec![fk_one, fk_two], fks);
}
<|endoftext|>


<commit_before>/*!
An actor type
*/

use pipes::{Port, Chan, SharedChan, stream};

/**
The client reference to an actor

Actors are only referred to by opaque handles parameterized
over the actor's message type, which can be considered the
actor's interface.
*/
struct ActorRef<M: Send> {
    chan: Chan<M>,
}

impl<M: Send> ActorRef<M> {
    fn send(msg: M) {
        self.chan.send(move msg);
    }
}

/// The local actor interface
trait Actor<M> {
    fn handle(msg: M) -> bool;
}

/// A helper function used by actor constructors
fn spawn<A: Actor<M>, M: Send>(f: ~fn() -> A) -> ActorRef<M> {
    let (chan, port) = stream();
    do task::spawn |move f, move port| {
        let actor = f();
        loop {
            let msg = port.recv();
            if !actor.handle(move msg) {
                break;
            }
        }
    }

    return ActorRef {
        chan: move chan
    }
}

struct SharedActorRef<M: Send> {
    chan: SharedChan<M>
}

impl<M: Send> SharedActorRef<M> {
    fn send(msg: M) {
        self.chan.send(move msg);
    }
}

fn SharedActorRef<M: Send>(actor: ActorRef<M>) -> SharedActorRef<M> {
    let chan = match move actor {
        ActorRef {
            chan: move chan
        } => {
            move chan
        }
    };

    SharedActorRef {
        chan: SharedChan(move chan)
    }
}

#[cfg(test)]
mod test {

    enum HelloMsg {
        Exit(Chan<()>)
    }

    struct HelloActor {
        name: ~str
    }

    impl HelloActor: Actor<HelloMsg> {
        fn handle(msg: HelloMsg) -> bool {
            match msg {
                Exit(chan) => {
                    chan.send(());
                    return false;
                }
            }
        }
    }

    fn HelloActor(name: ~str) -> ActorRef<HelloMsg> {
        do spawn |move name| {
            HelloActor {
                name: copy name
            }
        }
    }

    #[test]
    fn test_exit() {
        let actor = HelloActor(~"bob");
        let (chan, port) = stream();
        actor.send(Exit(move chan));
        port.recv();
    }
}<commit_msg>Add a SharedActor test<commit_after>/*!
An actor type
*/

use pipes::{Port, Chan, SharedChan, stream};

/**
The client reference to an actor

Actors are only referred to by opaque handles parameterized
over the actor's message type, which can be considered the
actor's interface.
*/
struct ActorRef<M: Send> {
    chan: Chan<M>,
}

impl<M: Send> ActorRef<M> {
    fn send(&self, msg: M) {
        self.chan.send(move msg);
    }
}

/// The local actor interface
trait Actor<M> {
    fn handle(&self, msg: M) -> bool;
}

/// A helper function used by actor constructors
fn spawn<A: Actor<M>, M: Send>(f: ~fn() -> A) -> ActorRef<M> {
    let (chan, port) = stream();
    do task::spawn |move f, move port| {
        let actor = f();
        loop {
            let msg = port.recv();
            if !actor.handle(move msg) {
                break;
            }
        }
    }

    return ActorRef {
        chan: move chan
    }
}

struct SharedActorRef<M: Send> {
    chan: SharedChan<M>
}

impl<M: Send> SharedActorRef<M> {
    fn send(&self, msg: M) {
        self.chan.send(move msg);
    }

    fn clone(&self) -> SharedActorRef<M> {
        SharedActorRef {
            chan: self.chan.clone()
        }
    }
}

fn SharedActorRef<M: Send>(actor: ActorRef<M>) -> SharedActorRef<M> {
    let chan = match move actor {
        ActorRef {
            chan: move chan
        } => {
            move chan
        }
    };

    SharedActorRef {
        chan: SharedChan(move chan)
    }
}

#[cfg(test)]
mod test {

    enum HelloMsg {
        GetName(Chan<~str>),
        Exit(Chan<()>)
    }

    struct HelloActor {
        name: ~str
    }

    impl HelloActor: Actor<HelloMsg> {
        fn handle(&self, msg: HelloMsg) -> bool {
            match msg {
                GetName(chan) => chan.send(copy self.name),
                Exit(chan) => {
                    chan.send(());
                    return false;
                }
            }

            return true;
        }
    }

    fn HelloActor(name: ~str) -> ActorRef<HelloMsg> {
        do spawn |move name| {
            HelloActor {
                name: copy name
            }
        }
    }

    #[test]
    fn test_exit() {
        let actor = HelloActor(~"bob");
        let (chan, port) = stream();
        actor.send(Exit(move chan));
        port.recv();
    }

    #[test]
    fn test_shared() {
        let actor = HelloActor(~"bob");
        let actor1 = SharedActorRef(move actor);
        let actor2 = actor1.clone();

        let (chan1, port1) = stream();
        actor1.send(GetName(move chan1));
        let (chan2, port2) = stream();
        actor2.send(GetName(move chan2));

        assert port1.recv() == ~"bob";
        assert port2.recv() == ~"bob";

        let (chan, port) = stream();
        actor1.send(Exit(move chan));
        port.recv();
    }

}<|endoftext|>


<commit_before>#[doc(hidden)]
pub mod pg;

mod limit_clause;
mod offset_clause;
mod order_clause;
mod select_statement;
mod where_clause;
pub mod update_statement;
mod delete_statement;
mod functions;

pub use self::functions::*;
#[doc(hidden)]
pub use self::select_statement::SelectStatement;
pub use self::update_statement::{IncompleteUpdateStatement, AsChangeset, Changeset, UpdateTarget};

use expression::Expression;
use std::error::Error;
use types::NativeSqlType;

pub type Binds = Vec<Option<Vec<u8>>>;
pub type BuildQueryResult = Result<(), Box<Error>>;

pub trait QueryBuilder {
    fn push_sql(&mut self, sql: &str);
    fn push_identifier(&mut self, identifier: &str) -> BuildQueryResult;
    fn push_bound_value(&mut self, tpe: &NativeSqlType, binds: Option<Vec<u8>>);
}

pub trait Query: QueryFragment {
    type SqlType: NativeSqlType;
}

impl<'a, T: Query> Query for &'a T where &'a T: QueryFragment {
    type SqlType = T::SqlType;
}

pub trait QueryFragment {
    fn to_sql(&self, out: &mut QueryBuilder) -> BuildQueryResult;
}

impl<T: Expression> QueryFragment for T {
    fn to_sql(&self, out: &mut QueryBuilder) -> BuildQueryResult {
        Expression::to_sql(self, out)
    }
}

pub trait AsQuery {
    type SqlType: NativeSqlType;
    type Query: Query<SqlType=Self::SqlType>;

    fn as_query(self) -> Self::Query;
}

impl<T: Query> AsQuery for T {
    type SqlType = <Self as Query>::SqlType;
    type Query = Self;

    fn as_query(self) -> Self::Query {
        self
    }
}
<commit_msg>Document the top level `query_builder` module<commit_after>//! Contains traits responsible for the actual construction of SQL statements
#[doc(hidden)]
pub mod pg;

mod limit_clause;
mod offset_clause;
mod order_clause;
mod select_statement;
mod where_clause;
#[doc(hidden)]
pub mod update_statement;
mod delete_statement;
mod functions;

pub use self::functions::*;
#[doc(hidden)]
pub use self::select_statement::SelectStatement;
#[doc(inline)]
pub use self::update_statement::{IncompleteUpdateStatement, AsChangeset, Changeset, UpdateTarget};

use expression::Expression;
use std::error::Error;
use types::NativeSqlType;

#[doc(hidden)]
pub type Binds = Vec<Option<Vec<u8>>>;
pub type BuildQueryResult = Result<(), Box<Error>>;

/// Apps should not need to concern themselves with this trait.
///
/// This is the trait used to actually construct a SQL query. You will take one
/// of these as an argument if you're implementing
/// [`Expression`](../expression/trait.Expression.html) or
/// [`QueryFragment`](trait.QueryFragment.html) manually.
pub trait QueryBuilder {
    fn push_sql(&mut self, sql: &str);
    fn push_identifier(&mut self, identifier: &str) -> BuildQueryResult;
    fn push_bound_value(&mut self, tpe: &NativeSqlType, binds: Option<Vec<u8>>);
}

/// A complete SQL query with a return type. This can be a select statement, or
/// a command such as `update` or `insert` with a `RETURNING` clause. Unlike
/// [`Expression`](../expression/trait.Expression.html), types implementing this
/// trait are guaranteed to be executable on their own.
pub trait Query: QueryFragment {
    type SqlType: NativeSqlType;
}

impl<'a, T: Query> Query for &'a T where &'a T: QueryFragment {
    type SqlType = T::SqlType;
}

/// An untyped fragment of SQL. This may be a complete SQL command (such as
/// an update statement without a `RETURNING` clause), or a subsection (such as
/// our internal types used to represent a `WHERE` clause). All methods on
/// [`Connection`](../struct.Connection.html) that execute a query require this
/// trait to be implemented.
pub trait QueryFragment {
    fn to_sql(&self, out: &mut QueryBuilder) -> BuildQueryResult;
}

impl<T: Expression> QueryFragment for T {
    fn to_sql(&self, out: &mut QueryBuilder) -> BuildQueryResult {
        Expression::to_sql(self, out)
    }
}

/// Types that can be converted into a complete, typed SQL query. This is used
/// internally to automatically add the right select clause when none is
/// specified, or to automatically add `RETURNING *` in certain contexts
pub trait AsQuery {
    type SqlType: NativeSqlType;
    type Query: Query<SqlType=Self::SqlType>;

    fn as_query(self) -> Self::Query;
}

impl<T: Query> AsQuery for T {
    type SqlType = <Self as Query>::SqlType;
    type Query = Self;

    fn as_query(self) -> Self::Query {
        self
    }
}
<|endoftext|>


<commit_before>/// Imports used by specific backends
pub use base_decode::BaseDecode;
pub use base_encode::BaseEncode;
pub use clap::{App, Arg, ArgMatches};
pub use compiler_options::CompilerOptions;
pub use converter::Converter;
pub use core::*;
pub use dynamic_converter::DynamicConverter;
pub use dynamic_decode::DynamicDecode;
pub use dynamic_encode::DynamicEncode;
pub use environment::Environment;
pub use for_context::ForContext;
pub use into_bytes::IntoBytes;
pub use naming::{CamelCase, FromNaming, Naming, SnakeCase};
pub use options::Options;
pub use package_processor::PackageProcessor;
pub use package_utils::PackageUtils;
pub use value_builder::ValueBuilder;
<commit_msg>backend: remove imports<commit_after><|endoftext|>