<commit_before>use crate::{
    arch::{gdt, interrupt::InterruptStack},
    context,
    ptrace,
    syscall,
    syscall::flag::{PTRACE_FLAG_IGNORE, PTRACE_STOP_PRE_SYSCALL, PTRACE_STOP_POST_SYSCALL},
};
use x86::msr;

pub unsafe fn init() {
    // IA32_STAR[31:0] are reserved.

    // The base selector of the two consecutive segments for kernel code and the immediately
    // suceeding stack (data).
    let syscall_cs_ss_base = (gdt::GDT_KERNEL_CODE as u16) << 3;
    // The base selector of the three consecutive segments (of which two are used) for user code
    // and user data. It points to a 32-bit code segment, which must be followed by a data segment
    // (stack), and a 64-bit code segment.
    let sysret_cs_ss_base = ((gdt::GDT_USER_CODE32_UNUSED as u16) << 3) | 3;
    let star_high = u32::from(syscall_cs_ss_base) | (u32::from(sysret_cs_ss_base) << 16);

    msr::wrmsr(msr::IA32_STAR, u64::from(star_high) << 32);
    msr::wrmsr(msr::IA32_LSTAR, syscall_instruction as u64);
    msr::wrmsr(msr::IA32_FMASK, 0x0300); // Clear trap flag and interrupt enable

    let efer = msr::rdmsr(msr::IA32_EFER);
    msr::wrmsr(msr::IA32_EFER, efer | 1);
}

macro_rules! with_interrupt_stack {
    (|$stack:ident| $code:block) => {{
        let allowed = ptrace::breakpoint_callback(PTRACE_STOP_PRE_SYSCALL, None)
            .and_then(|_| ptrace::next_breakpoint().map(|f| !f.contains(PTRACE_FLAG_IGNORE)));

        if allowed.unwrap_or(true) {
            // If the syscall is `clone`, the clone won't return here. Instead,
            // it'll return early and leave any undropped values. This is
            // actually GOOD, because any references are at that point UB
            // anyway, because they are based on the wrong stack.
            let $stack = &mut *$stack;
            (*$stack).scratch.rax = $code;
        }

        ptrace::breakpoint_callback(PTRACE_STOP_POST_SYSCALL, None);
    }}
}

#[no_mangle]
pub unsafe extern "C" fn __inner_syscall_instruction(stack: *mut InterruptStack) {
    let _guard = ptrace::set_process_regs(stack);
    with_interrupt_stack!(|stack| {
        // Set a restore point for clone
        let rbp;
        asm!("mov {}, rbp", out(reg) rbp);

        let scratch = &stack.scratch;
        syscall::syscall(scratch.rax, scratch.rdi, scratch.rsi, scratch.rdx, scratch.r10, scratch.r8, rbp, stack)
    });
}

macro_rules! fast_sysretq {
    () => { "
        pop rcx                 // Pop userspace return pointer
        sub rsp, 8              // Pop fake userspace CS
        pop r11                 // Pop rflags
        pop QWORD PTR gs:[0x70] // Pop userspace stack pointer
        mov rsp, gs:[0x70]      // Restore userspace stack pointer
        swapgs                  // Restore gs from TSS to user data
        ud2
        sysretq                 // Return into userspace; RCX=>RIP,R11=>RFLAGS
    " }
}

function!(syscall_instruction => {
    // Yes, this is magic. No, you don't need to understand
    "
        swapgs                    // Set gs segment to TSS
        mov gs:[0x70], rsp        // Save userspace stack pointer
        mov rsp, gs:[4]           // Load kernel stack pointer
        push QWORD PTR 5 * 8 + 3  // Push fake SS (resembling iret stack frame)
        push QWORD PTR gs:[0x70]  // Push userspace rsp
        push r11                  // Push rflags
        push QWORD PTR 6 * 8 + 3  // Push fake CS (resembling iret stack frame)
        push rcx                  // Push userspace return pointer
    ",

    // Push context registers
    "push rax\n",
    push_scratch!(),
    push_preserved!(),
    push_fs!(),

    // TODO: Map PTI
    // $crate::arch::x86_64::pti::map();

    // Call inner funtion
    "mov rdi, rsp\n",
    "call __inner_syscall_instruction\n",

    // TODO: Unmap PTI
    // $crate::arch::x86_64::pti::unmap();

    // Pop context registers
    pop_fs!(),
    pop_preserved!(),
    pop_scratch!(),

    // Return
    "
        // Test whether RCX is canonical. This is not strictly necessary, but could be fatal if
        // some kernel bug would allow RCX to be modified by user code.

        // Peek at the preserved RCX, to double-check that it is canonical.
        // Set ZF iff bit 47 (i.e. the bit that must be sign extended) is set.

        // TODO: Any hacky one-instruction canonicalness check?

        bt QWORD PTR [rsp], 47

        // Jump to the canonicalness check.
        jnz 2f

        // Otherwise, continue with the fast return.
        1:
        // Fast sysretq
    ",

    fast_sysretq!(),

    "
        // Jumps to aligned addresses are faster, and it is not always the case the userspace
        // doesn't call from higher half (though unlikely).
        .align 8, 0x90

        2:

        // Load userspace return RIP.
        mov rcx, [rsp]

        // Shift it to the right by 47, only getting the bits which must all be zero or one.
        shr rcx, 47
        // Since we have already checked that RCX[47] == 1, RCX[63:47] must all be one.
        cmp rcx, 0x1FFFF
        // If the address is canonical, go to the slower iretq.
        jne 3f
    ",

        fast_sysretq!(),

    "
        3:
        // Slow iretq

        xor rcx, rcx
        xor r11, r11
        swapgs
        iretq
    ",
});

interrupt_stack!(syscall, |stack| {
    with_interrupt_stack!(|stack| {
        {
            let contexts = context::contexts();
            let context = contexts.current();
            if let Some(current) = context {
                let current = current.read();
                println!("Warning: Context {} used deprecated `int 0x80` construct", *current.name.read());
            } else {
                println!("Warning: Unknown context used deprecated `int 0x80` construct");
            }
        }

        // Set a restore point for clone
        let rbp;
        asm!("mov {}, rbp", out(reg) rbp);

        let scratch = &stack.scratch;
        syscall::syscall(scratch.rax, stack.preserved.rbx, scratch.rcx, scratch.rdx, scratch.rsi, scratch.rdi, rbp, stack)
    })
});

function!(clone_ret => {
    // The address of this instruction is injected by `clone` in process.rs, on
    // top of the stack syscall->inner in this file, which is done using the rbp
    // register we save there.
    //
    // The top of our stack here is the address pointed to by rbp, which is:
    //
    // - the previous rbp
    // - the return location
    //
    // Our goal is to return from the parent function, inner, so we restore
    // rbp...
    "pop rbp\n",
    // ...and we return to the address at the top of the stack
    "ret\n",
});
<commit_msg>Simplify sysretq code.<commit_after>use crate::{
    arch::{gdt, interrupt::InterruptStack},
    context,
    ptrace,
    syscall,
    syscall::flag::{PTRACE_FLAG_IGNORE, PTRACE_STOP_PRE_SYSCALL, PTRACE_STOP_POST_SYSCALL},
};
use x86::msr;

pub unsafe fn init() {
    // IA32_STAR[31:0] are reserved.

    // The base selector of the two consecutive segments for kernel code and the immediately
    // suceeding stack (data).
    let syscall_cs_ss_base = (gdt::GDT_KERNEL_CODE as u16) << 3;
    // The base selector of the three consecutive segments (of which two are used) for user code
    // and user data. It points to a 32-bit code segment, which must be followed by a data segment
    // (stack), and a 64-bit code segment.
    let sysret_cs_ss_base = ((gdt::GDT_USER_CODE32_UNUSED as u16) << 3) | 3;
    let star_high = u32::from(syscall_cs_ss_base) | (u32::from(sysret_cs_ss_base) << 16);

    msr::wrmsr(msr::IA32_STAR, u64::from(star_high) << 32);
    msr::wrmsr(msr::IA32_LSTAR, syscall_instruction as u64);
    msr::wrmsr(msr::IA32_FMASK, 0x0300); // Clear trap flag and interrupt enable

    let efer = msr::rdmsr(msr::IA32_EFER);
    msr::wrmsr(msr::IA32_EFER, efer | 1);
}

macro_rules! with_interrupt_stack {
    (|$stack:ident| $code:block) => {{
        let allowed = ptrace::breakpoint_callback(PTRACE_STOP_PRE_SYSCALL, None)
            .and_then(|_| ptrace::next_breakpoint().map(|f| !f.contains(PTRACE_FLAG_IGNORE)));

        if allowed.unwrap_or(true) {
            // If the syscall is `clone`, the clone won't return here. Instead,
            // it'll return early and leave any undropped values. This is
            // actually GOOD, because any references are at that point UB
            // anyway, because they are based on the wrong stack.
            let $stack = &mut *$stack;
            (*$stack).scratch.rax = $code;
        }

        ptrace::breakpoint_callback(PTRACE_STOP_POST_SYSCALL, None);
    }}
}

#[no_mangle]
pub unsafe extern "C" fn __inner_syscall_instruction(stack: *mut InterruptStack) {
    let _guard = ptrace::set_process_regs(stack);
    with_interrupt_stack!(|stack| {
        // Set a restore point for clone
        let rbp;
        asm!("mov {}, rbp", out(reg) rbp);

        let scratch = &stack.scratch;
        syscall::syscall(scratch.rax, scratch.rdi, scratch.rsi, scratch.rdx, scratch.r10, scratch.r8, rbp, stack)
    });
}

function!(syscall_instruction => {
    // Yes, this is magic. No, you don't need to understand
    "
        swapgs                    // Set gs segment to TSS
        mov gs:[0x70], rsp        // Save userspace stack pointer
        mov rsp, gs:[4]           // Load kernel stack pointer
        push QWORD PTR 5 * 8 + 3  // Push fake SS (resembling iret stack frame)
        push QWORD PTR gs:[0x70]  // Push userspace rsp
        push r11                  // Push rflags
        push QWORD PTR 6 * 8 + 3  // Push fake CS (resembling iret stack frame)
        push rcx                  // Push userspace return pointer
    ",

    // Push context registers
    "push rax\n",
    push_scratch!(),
    push_preserved!(),
    push_fs!(),

    // TODO: Map PTI
    // $crate::arch::x86_64::pti::map();

    // Call inner funtion
    "mov rdi, rsp\n",
    "call __inner_syscall_instruction\n",

    // TODO: Unmap PTI
    // $crate::arch::x86_64::pti::unmap();

    // Pop context registers
    pop_fs!(),
    pop_preserved!(),
    pop_scratch!(),

    // Return
    //
    // We must test whether RCX is canonical. This is not strictly necessary, but could be
    // fatal if some kernel bug would allow RCX to be modified by user code.
    //
    // See https://xenproject.org/2012/06/13/the-intel-sysret-privilege-escalation/.
    //
    // This is not just theoretical; ptrace allows userspace to change rcx of target processes.
    "
        pop rcx                 // Pop userspace return pointer

        // Set ZF iff forbidden bit 47 (i.e. the bit that must be sign extended) is set.
        bt rcx, 47

        // If ZF was set, i.e. the address was invalid higher-half, so jump to the slower iretq and
        // handle the error without being able to execute attacker-controlled code!
        jmp 1f

        // Otherwise, continue with the fast sysretq.

        sub rsp, 8              // Pop fake userspace CS
        pop r11                 // Pop rflags
        pop QWORD PTR gs:[0x70] // Pop userspace stack pointer
        mov rsp, gs:[0x70]      // Restore userspace stack pointer
        swapgs                  // Restore gs from TSS to user data
        sysretq                 // Return into userspace; RCX=>RIP,R11=>RFLAGS

1:

        // Slow iretq
        push rcx
        xor rcx, rcx
        xor r11, r11
        swapgs
        iretq
    ",
});

interrupt_stack!(syscall, |stack| {
    with_interrupt_stack!(|stack| {
        {
            let contexts = context::contexts();
            let context = contexts.current();
            if let Some(current) = context {
                let current = current.read();
                println!("Warning: Context {} used deprecated `int 0x80` construct", *current.name.read());
            } else {
                println!("Warning: Unknown context used deprecated `int 0x80` construct");
            }
        }

        // Set a restore point for clone
        let rbp;
        asm!("mov {}, rbp", out(reg) rbp);

        let scratch = &stack.scratch;
        syscall::syscall(scratch.rax, stack.preserved.rbx, scratch.rcx, scratch.rdx, scratch.rsi, scratch.rdi, rbp, stack)
    })
});

function!(clone_ret => {
    // The address of this instruction is injected by `clone` in process.rs, on
    // top of the stack syscall->inner in this file, which is done using the rbp
    // register we save there.
    //
    // The top of our stack here is the address pointed to by rbp, which is:
    //
    // - the previous rbp
    // - the return location
    //
    // Our goal is to return from the parent function, inner, so we restore
    // rbp...
    "pop rbp\n",
    // ...and we return to the address at the top of the stack
    "ret\n",
});
<|endoftext|>


<commit_before>//
// FlightVars
// Copyright (c) 2015, 2016 Alvaro Polo
//
// 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 std::fmt;
use std::io;
use std::str::FromStr;

use byteorder::{BigEndian, ReadBytesExt};
use hex::FromHex;

#[derive(Debug, PartialEq)]
pub struct OffsetAddr(u16);

impl FromHex for OffsetAddr {
    type Error = io::Error;
    fn from_hex(s: &str) -> Result<Self, Self::Error> {
        let error = || io::Error::new(
            io::ErrorKind::InvalidInput,
            format!("invalid FSUIPC offset address in '{}'", s));
        let buf: Vec<u8> = try!(FromHex::from_hex(s).map_err(|_| error()));
        if buf.len() > 2 { Err(error()) }
        else {
            (&buf[..]).read_u16::<BigEndian>()
                .map(|v| OffsetAddr(v))
                .map_err(|_| error())
        }
    }
}

impl fmt::Display for OffsetAddr {
    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        write!(f, "{:x}", self.0)
    }
}

#[derive(Debug, PartialEq)]
pub enum OffsetLength { Ub, Sb, Uw, Sw, Ud, Sd }

impl fmt::Display for OffsetLength {
    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        match *self {
            OffsetLength::Ub => write!(f, "UB"),
            OffsetLength::Sb => write!(f, "SB"),
            OffsetLength::Uw => write!(f, "UW"),
            OffsetLength::Sw => write!(f, "SW"),
            OffsetLength::Ud => write!(f, "UD"),
            OffsetLength::Sd => write!(f, "SD"),
        }
    }
}

impl FromStr for OffsetLength {
    type Err = io::Error;
    fn from_str(s: &str) -> Result<OffsetLength, io::Error> {
        match s {
            "UB" => Ok(OffsetLength::Ub),
            "SB" => Ok(OffsetLength::Sb),
            "UW" => Ok(OffsetLength::Uw),
            "SW" => Ok(OffsetLength::Sw),
            "UD" => Ok(OffsetLength::Ud),
            "SD" => Ok(OffsetLength::Sd),
            _ => Err(io::Error::new(
                io::ErrorKind::InvalidInput,
                format!("invalid FSUIPC offset length in '{}'", s))),
        }
    }
}

#[derive(Debug, PartialEq)]
pub enum Message {
    Begin { version: u16, client_id: String },
    WriteLvar { lvar: String, value: i64 },
    WriteOffset { address: OffsetAddr, length: OffsetLength, value: i64 }
}

impl Message {

    pub fn begin(version: u16, client_id: &str) -> Message {
        Message::Begin { version: version, client_id: client_id.to_string() }
    }

    pub fn write_lvar(lvar: &str, value: i64) -> Message {
        Message::WriteLvar { lvar: lvar.to_string(), value: value }
    }

    pub fn write_offset(addr: OffsetAddr, len: OffsetLength, value: i64) -> Message {
        Message::WriteOffset { address: addr, length: len, value: value }
    }

    pub fn decode<R: io::BufRead>(r: &mut R) -> io::Result<Message> {
        let mut line = String::new();
        try!(r.read_line(&mut line));
        let mut split = line.split_whitespace();
        let cmd = try!(split.next().ok_or(Message::input_error(&line)));
        match &cmd.to_uppercase()[..] {
            "BEGIN" => Message::decode_begin(&line, &mut split),
            "WRITE_LVAR" => Message::decode_write_lvar(&line, &mut split),
            "WRITE_OFFSET" => Message::decode_write_offset(&line, &mut split),
            _ => Err(Message::input_error(&line)),
        }
    }

    pub fn encode<W: io::Write>(&self, w: &mut W) -> io::Result<()> {
        match self {
            &Message::Begin { version, ref client_id } =>
                write!(w, "BEGIN {} {}\n", version, client_id),
            &Message::WriteLvar { ref lvar, value} =>
                write!(w, "WRITE_LVAR {} {}\n", lvar, value),
            &Message::WriteOffset { ref address, ref length, value } =>
                write!(w, "WRITE_OFFSET {}:{} {}\n", address, length, value),
        }
    }

    fn decode_begin<'a, I: Iterator<Item=&'a str>>(
            line: &str, args: &mut I) -> io::Result<Message> {
        let error = || Message::input_error(line);
        let version = try!(args.next().ok_or_else(&error));
        let client_id = try!(args.next().ok_or_else(&error));
        let version = try!(version.parse().map_err(|_| error()));
        Ok(Message::begin(version, &client_id))
    }

    fn decode_write_lvar<'a, I: Iterator<Item=&'a str>>(
            line: &str, args: &mut I) -> io::Result<Message> {
        let error = || Message::input_error(line);
        let lvar = try!(args.next().ok_or_else(&error));
        let value = try!(args.next().ok_or_else(&error));
        let value = try!(value.parse().map_err(|_| error()));
        Ok(Message::write_lvar(&lvar, value))
    }

    fn decode_write_offset<'a, I: Iterator<Item=&'a str>>(
            line: &str, args: &mut I) -> io::Result<Message> {
        let error = || Message::input_error(line);
        let mut offset = try!(args.next().ok_or_else(&error)).split(":");
        let addr = try!(offset.next().ok_or_else(&error));
        let addr: OffsetAddr = try!(OffsetAddr::from_hex(&addr).map_err(|_| error()));
        let len = try!(offset.next().ok_or_else(&error));
        let len = try!(len.parse().map_err(|_| error()));
        let value = try!(args.next().ok_or_else(&error));
        let value = try!(value.parse().map_err(|_| error()));
        Ok(Message::write_offset(addr, len, value))
    }

    fn input_error(line: &str) -> io::Error {
        io::Error::new(
            io::ErrorKind::InvalidInput,
            format!("invalid oacsp message in '{}'", line))
    }
}

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

    #[test]
    fn should_encode_begin_msg() {
        let mut buf = vec![];
        let msg = Message::begin(1, "arduino");
        msg.encode(&mut buf).unwrap();
        assert_eq!(buf, b"BEGIN 1 arduino\n")
    }

    #[test]
    fn should_decode_begin_msg() {
        let mut buf = &b"BEGIN 1 arduino\n"[..];
        let msg = Message::decode(&mut buf).unwrap();
        assert_eq!(msg, Message::begin(1, "arduino"));
    }

    #[test]
    fn should_encode_write_lvar_msg() {
        let mut buf = vec![];
        let msg = Message::write_lvar("foobar", 42);
        msg.encode(&mut buf).unwrap();
        assert_eq!(buf, b"WRITE_LVAR foobar 42\n")
    }

    #[test]
    fn should_decode_write_lvar_msg() {
        let mut buf = &b"WRITE_LVAR foobar 42\n"[..];
        let msg = Message::decode(&mut buf).unwrap();
        assert_eq!(msg, Message::write_lvar("foobar", 42));
    }

    #[test]
    fn should_encode_write_offset_msg() {
        let mut buf = vec![];
        let msg = Message::write_offset(OffsetAddr(0x1234), OffsetLength::Uw, 42);
        msg.encode(&mut buf).unwrap();
        assert_eq!(buf, b"WRITE_OFFSET 1234:UW 42\n")
    }

    #[test]
    fn should_decode_write_offset_msg() {
        let mut buf = &b"WRITE_OFFSET 1234:UW 42\n"[..];
        let msg = Message::decode(&mut buf).unwrap();
        assert_eq!(msg, Message::write_offset(OffsetAddr(0x1234), OffsetLength::Uw, 42));
    }
}
<commit_msg>Add `ObserveLvar` message handling in `oacsp` module<commit_after>//
// FlightVars
// Copyright (c) 2015, 2016 Alvaro Polo
//
// 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 std::fmt;
use std::io;
use std::str::FromStr;

use byteorder::{BigEndian, ReadBytesExt};
use hex::FromHex;

#[derive(Debug, PartialEq)]
pub struct OffsetAddr(u16);

impl FromHex for OffsetAddr {
    type Error = io::Error;
    fn from_hex(s: &str) -> Result<Self, Self::Error> {
        let error = || io::Error::new(
            io::ErrorKind::InvalidInput,
            format!("invalid FSUIPC offset address in '{}'", s));
        let buf: Vec<u8> = try!(FromHex::from_hex(s).map_err(|_| error()));
        if buf.len() > 2 { Err(error()) }
        else {
            (&buf[..]).read_u16::<BigEndian>()
                .map(|v| OffsetAddr(v))
                .map_err(|_| error())
        }
    }
}

impl fmt::Display for OffsetAddr {
    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        write!(f, "{:x}", self.0)
    }
}

#[derive(Debug, PartialEq)]
pub enum OffsetLength { Ub, Sb, Uw, Sw, Ud, Sd }

impl fmt::Display for OffsetLength {
    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        match *self {
            OffsetLength::Ub => write!(f, "UB"),
            OffsetLength::Sb => write!(f, "SB"),
            OffsetLength::Uw => write!(f, "UW"),
            OffsetLength::Sw => write!(f, "SW"),
            OffsetLength::Ud => write!(f, "UD"),
            OffsetLength::Sd => write!(f, "SD"),
        }
    }
}

impl FromStr for OffsetLength {
    type Err = io::Error;
    fn from_str(s: &str) -> Result<OffsetLength, io::Error> {
        match s {
            "UB" => Ok(OffsetLength::Ub),
            "SB" => Ok(OffsetLength::Sb),
            "UW" => Ok(OffsetLength::Uw),
            "SW" => Ok(OffsetLength::Sw),
            "UD" => Ok(OffsetLength::Ud),
            "SD" => Ok(OffsetLength::Sd),
            _ => Err(io::Error::new(
                io::ErrorKind::InvalidInput,
                format!("invalid FSUIPC offset length in '{}'", s))),
        }
    }
}

#[derive(Debug, PartialEq)]
pub enum Message {
    Begin { version: u16, client_id: String },
    WriteLvar { lvar: String, value: i64 },
    WriteOffset { address: OffsetAddr, length: OffsetLength, value: i64 },
    ObserveLvar { lvar: String }
}

impl Message {

    pub fn begin(version: u16, client_id: &str) -> Message {
        Message::Begin { version: version, client_id: client_id.to_string() }
    }

    pub fn write_lvar(lvar: &str, value: i64) -> Message {
        Message::WriteLvar { lvar: lvar.to_string(), value: value }
    }

    pub fn write_offset(addr: OffsetAddr, len: OffsetLength, value: i64) -> Message {
        Message::WriteOffset { address: addr, length: len, value: value }
    }

    pub fn obs_lvar(lvar: &str) -> Message {
        Message::ObserveLvar { lvar: lvar.to_string() }
    }

    pub fn decode<R: io::BufRead>(r: &mut R) -> io::Result<Message> {
        let mut line = String::new();
        try!(r.read_line(&mut line));
        let mut split = line.split_whitespace();
        let cmd = try!(split.next().ok_or(Message::input_error(&line)));
        match &cmd.to_uppercase()[..] {
            "BEGIN" => Message::decode_begin(&line, &mut split),
            "WRITE_LVAR" => Message::decode_write_lvar(&line, &mut split),
            "WRITE_OFFSET" => Message::decode_write_offset(&line, &mut split),
            "OBS_LVAR" => Message::decode_obs_lvar(&line, &mut split),
            _ => Err(Message::input_error(&line)),
        }
    }

    pub fn encode<W: io::Write>(&self, w: &mut W) -> io::Result<()> {
        match self {
            &Message::Begin { version, ref client_id } =>
                write!(w, "BEGIN {} {}\n", version, client_id),
            &Message::WriteLvar { ref lvar, value} =>
                write!(w, "WRITE_LVAR {} {}\n", lvar, value),
            &Message::WriteOffset { ref address, ref length, value } =>
                write!(w, "WRITE_OFFSET {}:{} {}\n", address, length, value),
            &Message::ObserveLvar { ref lvar } =>
                write!(w, "OBS_LVAR {}\n", lvar),
        }
    }

    fn decode_begin<'a, I: Iterator<Item=&'a str>>(
            line: &str, args: &mut I) -> io::Result<Message> {
        let error = || Message::input_error(line);
        let version = try!(args.next().ok_or_else(&error));
        let client_id = try!(args.next().ok_or_else(&error));
        let version = try!(version.parse().map_err(|_| error()));
        Ok(Message::begin(version, &client_id))
    }

    fn decode_write_lvar<'a, I: Iterator<Item=&'a str>>(
            line: &str, args: &mut I) -> io::Result<Message> {
        let error = || Message::input_error(line);
        let lvar = try!(args.next().ok_or_else(&error));
        let value = try!(args.next().ok_or_else(&error));
        let value = try!(value.parse().map_err(|_| error()));
        Ok(Message::write_lvar(&lvar, value))
    }

    fn decode_write_offset<'a, I: Iterator<Item=&'a str>>(
            line: &str, args: &mut I) -> io::Result<Message> {
        let error = || Message::input_error(line);
        let mut offset = try!(args.next().ok_or_else(&error)).split(":");
        let addr = try!(offset.next().ok_or_else(&error));
        let addr: OffsetAddr = try!(OffsetAddr::from_hex(&addr).map_err(|_| error()));
        let len = try!(offset.next().ok_or_else(&error));
        let len = try!(len.parse().map_err(|_| error()));
        let value = try!(args.next().ok_or_else(&error));
        let value = try!(value.parse().map_err(|_| error()));
        Ok(Message::write_offset(addr, len, value))
    }

    fn decode_obs_lvar<'a, I: Iterator<Item=&'a str>>(
            line: &str, args: &mut I) -> io::Result<Message> {
        let lvar = try!(args.next().ok_or(Message::input_error(line)));
        Ok(Message::obs_lvar(&lvar))
    }

    fn input_error(line: &str) -> io::Error {
        io::Error::new(
            io::ErrorKind::InvalidInput,
            format!("invalid oacsp message in '{}'", line))
    }
}

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

    #[test]
    fn should_encode_begin_msg() {
        let mut buf = vec![];
        let msg = Message::begin(1, "arduino");
        msg.encode(&mut buf).unwrap();
        assert_eq!(buf, b"BEGIN 1 arduino\n")
    }

    #[test]
    fn should_decode_begin_msg() {
        let mut buf = &b"BEGIN 1 arduino\n"[..];
        let msg = Message::decode(&mut buf).unwrap();
        assert_eq!(msg, Message::begin(1, "arduino"));
    }

    #[test]
    fn should_encode_write_lvar_msg() {
        let mut buf = vec![];
        let msg = Message::write_lvar("foobar", 42);
        msg.encode(&mut buf).unwrap();
        assert_eq!(buf, b"WRITE_LVAR foobar 42\n")
    }

    #[test]
    fn should_decode_write_lvar_msg() {
        let mut buf = &b"WRITE_LVAR foobar 42\n"[..];
        let msg = Message::decode(&mut buf).unwrap();
        assert_eq!(msg, Message::write_lvar("foobar", 42));
    }

    #[test]
    fn should_encode_write_offset_msg() {
        let mut buf = vec![];
        let msg = Message::write_offset(OffsetAddr(0x1234), OffsetLength::Uw, 42);
        msg.encode(&mut buf).unwrap();
        assert_eq!(buf, b"WRITE_OFFSET 1234:UW 42\n")
    }

    #[test]
    fn should_decode_write_offset_msg() {
        let mut buf = &b"WRITE_OFFSET 1234:UW 42\n"[..];
        let msg = Message::decode(&mut buf).unwrap();
        assert_eq!(msg, Message::write_offset(OffsetAddr(0x1234), OffsetLength::Uw, 42));
    }

    #[test]
    fn should_encode_obs_lvar_msg() {
        let mut buf = vec![];
        let msg = Message::obs_lvar("foobar");
        msg.encode(&mut buf).unwrap();
        assert_eq!(buf, b"OBS_LVAR foobar\n")
    }

    #[test]
    fn should_decode_obs_lvar_msg() {
        let mut buf = &b"OBS_LVAR foobar\n"[..];
        let msg = Message::decode(&mut buf).unwrap();
        assert_eq!(msg, Message::obs_lvar("foobar"));
    }
}
<|endoftext|>


<commit_before>#![allow(unsafe_code)]

use std::{
    fmt::{self, Debug},
    ops::Deref,
    sync::atomic::Ordering::{Acquire, Release},
};

use crossbeam_epoch::{unprotected, Atomic, Guard, Owned, Shared};

use crate::debug_delay;

/// A node in the lock-free `Stack`.
#[derive(Debug)]
pub struct Node<T: Send + 'static> {
    pub(crate) inner: T,
    pub(crate) next: Atomic<Node<T>>,
}

impl<T: Send + 'static> Drop for Node<T> {
    fn drop(&mut self) {
        unsafe {
            let mut cursor = self.next.load(Acquire, unprotected());

            while !cursor.is_null() {
                // we carefully unset the next pointer here to avoid
                // a stack overflow when freeing long lists.
                let node = cursor.into_owned();
                cursor = node.next.swap(Shared::null(), Acquire, unprotected());
                drop(node);
            }
        }
    }
}

/// A simple lock-free stack, with the ability to atomically
/// append or entirely swap-out entries.
pub struct Stack<T: Send + 'static> {
    head: Atomic<Node<T>>,
}

impl<T: Send + 'static> Default for Stack<T> {
    fn default() -> Self {
        Self { head: Atomic::null() }
    }
}

impl<T: Send + 'static> Drop for Stack<T> {
    fn drop(&mut self) {
        unsafe {
            let curr = self.head.load(Acquire, unprotected());
            if !curr.as_raw().is_null() {
                drop(curr.into_owned());
            }
        }
    }
}

impl<T> Debug for Stack<T>
where
    T: Clone + Debug + Send + 'static + Sync,
{
    fn fmt(
        &self,
        formatter: &mut fmt::Formatter<'_>,
    ) -> Result<(), fmt::Error> {
        let guard = crossbeam_epoch::pin();
        let head = self.head(&guard);
        let iter = StackIter::from_ptr(head, &guard);

        formatter.write_str("Stack [")?;
        let mut written = false;
        for node in iter {
            if written {
                formatter.write_str(", ")?;
            }
            formatter.write_str(&*format!("({:?}) ", &node))?;
            node.fmt(formatter)?;
            written = true;
        }
        formatter.write_str("]")?;
        Ok(())
    }
}

impl<T: Send + 'static> Deref for Node<T> {
    type Target = T;
    fn deref(&self) -> &T {
        &self.inner
    }
}

impl<T: Send + Sync + 'static> Stack<T> {
    /// Add an item to the stack, spinning until successful.
    pub fn push(&self, inner: T, guard: &Guard) {
        debug_delay();
        let node = Owned::new(Node { inner, next: Atomic::null() });

        unsafe {
            let node = node.into_shared(guard);

            loop {
                let head = self.head(guard);
                node.deref().next.store(head, Release);
                if self.head.compare_and_set(head, node, Release, guard).is_ok()
                {
                    return;
                }
            }
        }
    }

    /// Clears the stack and returns all items
    pub fn take_iter<'a>(
        &self,
        guard: &'a Guard,
    ) -> impl Iterator<Item = &'a T> {
        debug_delay();
        let node = self.head.swap(Shared::null(), Release, guard);

        let iter = StackIter { inner: node, guard };

        if !node.is_null() {
            unsafe {
                guard.defer_destroy(node);
            }
        }

        iter
    }

    /// Pop the next item off the stack. Returns None if nothing is there.
    #[cfg(any(test, feature = "event_log"))]
    pub(crate) fn pop(&self, guard: &Guard) -> Option<T> {
        use std::ptr;
        use std::sync::atomic::Ordering::SeqCst;
        debug_delay();
        let mut head = self.head(guard);
        loop {
            match unsafe { head.as_ref() } {
                Some(h) => {
                    let next = h.next.load(Acquire, guard);
                    match self.head.compare_and_set(head, next, Release, guard)
                    {
                        Ok(_) => unsafe {
                            // we unset the next pointer before destruction
                            // to avoid double-frees.
                            h.next.store(Shared::default(), SeqCst);
                            guard.defer_destroy(head);
                            return Some(ptr::read(&h.inner));
                        },
                        Err(h) => head = h.current,
                    }
                }
                None => return None,
            }
        }
    }

    /// Returns the current head pointer of the stack, which can
    /// later be used as the key for cas and cap operations.
    pub fn head<'g>(&self, guard: &'g Guard) -> Shared<'g, Node<T>> {
        self.head.load(Acquire, guard)
    }
}

/// An iterator over nodes in a lock-free stack.
pub struct StackIter<'a, T>
where
    T: Send + 'static + Sync,
{
    inner: Shared<'a, Node<T>>,
    guard: &'a Guard,
}

impl<'a, T> StackIter<'a, T>
where
    T: 'a + Send + 'static + Sync,
{
    /// Creates a `StackIter` from a pointer to one.
    pub fn from_ptr<'b>(
        ptr: Shared<'b, Node<T>>,
        guard: &'b Guard,
    ) -> StackIter<'b, T> {
        StackIter { inner: ptr, guard }
    }
}

impl<'a, T> Iterator for StackIter<'a, T>
where
    T: Send + 'static + Sync,
{
    type Item = &'a T;

    fn next(&mut self) -> Option<Self::Item> {
        debug_delay();
        if self.inner.is_null() {
            None
        } else {
            unsafe {
                let ret = &self.inner.deref().inner;
                self.inner = self.inner.deref().next.load(Acquire, self.guard);
                Some(ret)
            }
        }
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        let mut size = 0;
        let mut cursor = self.inner;

        while !cursor.is_null() {
            unsafe {
                cursor = cursor.deref().next.load(Acquire, self.guard);
            }
            size += 1;
        }

        (size, Some(size))
    }
}

#[test]
fn basic_functionality() {
    use crossbeam_epoch::pin;
    use crossbeam_utils::CachePadded;
    use std::sync::Arc;
    use std::thread;

    let guard = pin();
    let ll = Arc::new(Stack::default());
    assert_eq!(ll.pop(&guard), None);
    ll.push(CachePadded::new(1), &guard);
    let ll2 = Arc::clone(&ll);
    let t = thread::spawn(move || {
        let guard = pin();
        ll2.push(CachePadded::new(2), &guard);
        ll2.push(CachePadded::new(3), &guard);
        ll2.push(CachePadded::new(4), &guard);
        guard.flush();
    });
    t.join().unwrap();
    ll.push(CachePadded::new(5), &guard);
    assert_eq!(ll.pop(&guard), Some(CachePadded::new(5)));
    assert_eq!(ll.pop(&guard), Some(CachePadded::new(4)));
    let ll3 = Arc::clone(&ll);
    let t = thread::spawn(move || {
        let guard = pin();
        assert_eq!(ll3.pop(&guard), Some(CachePadded::new(3)));
        assert_eq!(ll3.pop(&guard), Some(CachePadded::new(2)));
        guard.flush();
    });
    t.join().unwrap();
    assert_eq!(ll.pop(&guard), Some(CachePadded::new(1)));
    let ll4 = Arc::clone(&ll);
    let t = thread::spawn(move || {
        let guard = pin();
        assert_eq!(ll4.pop(&guard), None);
        guard.flush();
    });
    t.join().unwrap();
    drop(ll);
    guard.flush();
    drop(guard);
}
<commit_msg>Rename StackIter to Iter<commit_after>#![allow(unsafe_code)]

use std::{
    fmt::{self, Debug},
    ops::Deref,
    sync::atomic::Ordering::{Acquire, Release},
};

use crossbeam_epoch::{unprotected, Atomic, Guard, Owned, Shared};

use crate::debug_delay;

/// A node in the lock-free `Stack`.
#[derive(Debug)]
pub struct Node<T: Send + 'static> {
    pub(crate) inner: T,
    pub(crate) next: Atomic<Node<T>>,
}

impl<T: Send + 'static> Drop for Node<T> {
    fn drop(&mut self) {
        unsafe {
            let mut cursor = self.next.load(Acquire, unprotected());

            while !cursor.is_null() {
                // we carefully unset the next pointer here to avoid
                // a stack overflow when freeing long lists.
                let node = cursor.into_owned();
                cursor = node.next.swap(Shared::null(), Acquire, unprotected());
                drop(node);
            }
        }
    }
}

/// A simple lock-free stack, with the ability to atomically
/// append or entirely swap-out entries.
pub struct Stack<T: Send + 'static> {
    head: Atomic<Node<T>>,
}

impl<T: Send + 'static> Default for Stack<T> {
    fn default() -> Self {
        Self { head: Atomic::null() }
    }
}

impl<T: Send + 'static> Drop for Stack<T> {
    fn drop(&mut self) {
        unsafe {
            let curr = self.head.load(Acquire, unprotected());
            if !curr.as_raw().is_null() {
                drop(curr.into_owned());
            }
        }
    }
}

impl<T> Debug for Stack<T>
where
    T: Clone + Debug + Send + 'static + Sync,
{
    fn fmt(
        &self,
        formatter: &mut fmt::Formatter<'_>,
    ) -> Result<(), fmt::Error> {
        let guard = crossbeam_epoch::pin();
        let head = self.head(&guard);
        let iter = Iter::from_ptr(head, &guard);

        formatter.write_str("Stack [")?;
        let mut written = false;
        for node in iter {
            if written {
                formatter.write_str(", ")?;
            }
            formatter.write_str(&*format!("({:?}) ", &node))?;
            node.fmt(formatter)?;
            written = true;
        }
        formatter.write_str("]")?;
        Ok(())
    }
}

impl<T: Send + 'static> Deref for Node<T> {
    type Target = T;
    fn deref(&self) -> &T {
        &self.inner
    }
}

impl<T: Send + Sync + 'static> Stack<T> {
    /// Add an item to the stack, spinning until successful.
    pub fn push(&self, inner: T, guard: &Guard) {
        debug_delay();
        let node = Owned::new(Node { inner, next: Atomic::null() });

        unsafe {
            let node = node.into_shared(guard);

            loop {
                let head = self.head(guard);
                node.deref().next.store(head, Release);
                if self.head.compare_and_set(head, node, Release, guard).is_ok()
                {
                    return;
                }
            }
        }
    }

    /// Clears the stack and returns all items
    pub fn take_iter<'a>(
        &self,
        guard: &'a Guard,
    ) -> impl Iterator<Item = &'a T> {
        debug_delay();
        let node = self.head.swap(Shared::null(), Release, guard);

        let iter = Iter { inner: node, guard };

        if !node.is_null() {
            unsafe {
                guard.defer_destroy(node);
            }
        }

        iter
    }

    /// Pop the next item off the stack. Returns None if nothing is there.
    #[cfg(any(test, feature = "event_log"))]
    pub(crate) fn pop(&self, guard: &Guard) -> Option<T> {
        use std::ptr;
        use std::sync::atomic::Ordering::SeqCst;
        debug_delay();
        let mut head = self.head(guard);
        loop {
            match unsafe { head.as_ref() } {
                Some(h) => {
                    let next = h.next.load(Acquire, guard);
                    match self.head.compare_and_set(head, next, Release, guard)
                    {
                        Ok(_) => unsafe {
                            // we unset the next pointer before destruction
                            // to avoid double-frees.
                            h.next.store(Shared::default(), SeqCst);
                            guard.defer_destroy(head);
                            return Some(ptr::read(&h.inner));
                        },
                        Err(h) => head = h.current,
                    }
                }
                None => return None,
            }
        }
    }

    /// Returns the current head pointer of the stack, which can
    /// later be used as the key for cas and cap operations.
    pub fn head<'g>(&self, guard: &'g Guard) -> Shared<'g, Node<T>> {
        self.head.load(Acquire, guard)
    }
}

/// An iterator over nodes in a lock-free stack.
pub struct Iter<'a, T>
where
    T: Send + 'static + Sync,
{
    inner: Shared<'a, Node<T>>,
    guard: &'a Guard,
}

impl<'a, T> Iter<'a, T>
where
    T: 'a + Send + 'static + Sync,
{
    /// Creates a `Iter` from a pointer to one.
    pub fn from_ptr<'b>(
        ptr: Shared<'b, Node<T>>,
        guard: &'b Guard,
    ) -> Iter<'b, T> {
        Iter { inner: ptr, guard }
    }
}

impl<'a, T> Iterator for Iter<'a, T>
where
    T: Send + 'static + Sync,
{
    type Item = &'a T;

    fn next(&mut self) -> Option<Self::Item> {
        debug_delay();
        if self.inner.is_null() {
            None
        } else {
            unsafe {
                let ret = &self.inner.deref().inner;
                self.inner = self.inner.deref().next.load(Acquire, self.guard);
                Some(ret)
            }
        }
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        let mut size = 0;
        let mut cursor = self.inner;

        while !cursor.is_null() {
            unsafe {
                cursor = cursor.deref().next.load(Acquire, self.guard);
            }
            size += 1;
        }

        (size, Some(size))
    }
}

#[test]
fn basic_functionality() {
    use crossbeam_epoch::pin;
    use crossbeam_utils::CachePadded;
    use std::sync::Arc;
    use std::thread;

    let guard = pin();
    let ll = Arc::new(Stack::default());
    assert_eq!(ll.pop(&guard), None);
    ll.push(CachePadded::new(1), &guard);
    let ll2 = Arc::clone(&ll);
    let t = thread::spawn(move || {
        let guard = pin();
        ll2.push(CachePadded::new(2), &guard);
        ll2.push(CachePadded::new(3), &guard);
        ll2.push(CachePadded::new(4), &guard);
        guard.flush();
    });
    t.join().unwrap();
    ll.push(CachePadded::new(5), &guard);
    assert_eq!(ll.pop(&guard), Some(CachePadded::new(5)));
    assert_eq!(ll.pop(&guard), Some(CachePadded::new(4)));
    let ll3 = Arc::clone(&ll);
    let t = thread::spawn(move || {
        let guard = pin();
        assert_eq!(ll3.pop(&guard), Some(CachePadded::new(3)));
        assert_eq!(ll3.pop(&guard), Some(CachePadded::new(2)));
        guard.flush();
    });
    t.join().unwrap();
    assert_eq!(ll.pop(&guard), Some(CachePadded::new(1)));
    let ll4 = Arc::clone(&ll);
    let t = thread::spawn(move || {
        let guard = pin();
        assert_eq!(ll4.pop(&guard), None);
        guard.flush();
    });
    t.join().unwrap();
    drop(ll);
    guard.flush();
    drop(guard);
}
<|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 {
    name: String
}

impl CustomLua {
    fn new(name: String) -> Self {
        CustomLua { name: name }
    }
}

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("\"", "")
}
<commit_msg>Revert "CustomLua has a Name"<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>use std::collections::HashMap;
use std::io::timer::sleep;
use std::time::duration::Duration;
use std::comm::Disconnected;

use std::fmt::Show;

use std::sync::{Arc, Mutex};

use std::task::TaskBuilder;

use super::replication::Committable;

pub trait Intercommunication < T: Committable + Send + Show > {
    fn new() -> Self;
    fn register(&mut self, host: String) -> Endpoint < T >;
    fn receive(&mut self) -> Option < Package < T > >;
    fn send(&mut self, recipient: String, package: Package < T >);
    fn is_debug(&self) -> bool;
}

pub struct DefaultIntercommunication < T: Committable + Send > {
    receiver: Receiver < Package < T > >,
    sender: Sender < Package < T > >,
    senders: HashMap < String, Sender < Package < T > > >,

    pub is_debug: bool,
}

pub struct Endpoint < T: Committable + Send > {
    host: String,
    tx: Sender < Package < T > >,
    rx: Receiver < Package < T > >,
}

#[deriving(Clone, Show)]
pub struct AppendLog < T: Committable + Send > {
    pub committed_offset: uint,
    pub node_list: Vec < String >,
    pub enqueue: Option < AppendLogEntry < T > >,
}

#[deriving(Clone, Show)]
pub struct AppendLogEntry < T: Committable + Send > {
    pub offset: uint,
    pub entry: T,
}

impl < T: Committable + Send + Show > Intercommunication < T > for DefaultIntercommunication < T > {
    fn new() -> DefaultIntercommunication < T > {
        let (tx, rx) = channel();

        DefaultIntercommunication {
            senders: HashMap::new(),
            sender: tx,
            receiver: rx,
            is_debug: false,
        }
    }

    fn register(&mut self, host: String) -> Endpoint < T > {
        let (tx, rx) = channel();

        self.senders.insert(host.clone(), tx);

        Endpoint {
            host: host,
            rx: rx,
            tx: self.sender.clone(),
        }
    }

    fn receive(&mut self) -> Option < Package < T > > {
        match self.receiver.try_recv() {
            Ok(package) => Some(package),
            _ => None,
        }
    }

    fn send(&mut self, recipient: String, package: Package < T >) {
        match self.senders.find(&recipient) {
            Some(tx) => {
                // be more careful at sending
                match tx.send_opt(package) {
                    Err(_) => (),
                    _ => (),
                }
            },
            None => (),
        }
    }

    fn is_debug(&self) -> bool {
        self.is_debug
    }
}

impl < T: Committable + Send > Endpoint < T > {
    pub fn send (&self, host: String, package: PackageDetails < T >) {
        self.tx.send(Pack(self.host.clone(), host, package));
    }

    pub fn listen_block_with_timeout(&self) -> Option < Package < T > > {
        for _ in range(0, 10u) {
            match self.listen() {
                Some(x) => return Some(x),
                _ => ()
            }

            sleep(Duration::milliseconds(2));
        }

        None
    }

    pub fn listen(&self) -> Option < Package < T > > {
        match self.rx.try_recv() {
            Ok(result) => Some(result),
            _ => None,
        }
    }
}

#[deriving(Show)]
pub enum PackageDetails < T: Committable + Send > {
    Ack,

    LeaderQuery,

    // LeaderQueryResponse(leader_host)
    LeaderQueryResponse(Option < String >),

    // AppendQuery(log)
    AppendQuery(AppendLog < T >),

    // Persisted(entry_offset)
    Persisted(uint),

    // RequestVote(term
    RequestVote(uint),

    // Vote(term)
    Vote(uint),
}

#[deriving(Show)]
pub enum Package < T: Committable + Send > {
    // Pack(from, to, package)
    Pack(String, String, PackageDetails < T >),
}

pub fn start < T: Committable + Send + Clone + Show, I: Intercommunication < T > + Send >(intercommunication: I) -> Sender < int > {
    let mutex = Arc::new(Mutex::new(intercommunication));
    let (exit_tx, exit_rx) = channel();

    TaskBuilder::new().named("intercommunication").spawn(proc() {
        loop {
            let mut intercommunication = mutex.lock();

            match intercommunication.receive() {
                Some(Pack(from, to, package)) => {
                    if intercommunication.is_debug() {
                        println!("sent package {} to host {} from {}", package, to, from);
                    }

                    intercommunication.send(to.clone(), Pack(from, to, package));
                },
                None => ()
            }

            match exit_rx.try_recv() {
                Ok(_) | Err(Disconnected) => break,
                _ => (),
            }

            intercommunication.cond.signal();

            sleep(Duration::milliseconds(2));
        }
    });

    exit_tx
}
<commit_msg>Intercommunication: make Endpoint fields public<commit_after>use std::collections::HashMap;
use std::io::timer::sleep;
use std::time::duration::Duration;
use std::comm::Disconnected;

use std::fmt::Show;

use std::sync::{Arc, Mutex};

use std::task::TaskBuilder;

use super::replication::Committable;

pub trait Intercommunication < T: Committable + Send + Show > {
    fn new() -> Self;
    fn register(&mut self, host: String) -> Endpoint < T >;
    fn receive(&mut self) -> Option < Package < T > >;
    fn send(&mut self, recipient: String, package: Package < T >);
    fn is_debug(&self) -> bool;
}

pub struct DefaultIntercommunication < T: Committable + Send > {
    receiver: Receiver < Package < T > >,
    sender: Sender < Package < T > >,
    senders: HashMap < String, Sender < Package < T > > >,

    pub is_debug: bool,
}

pub struct Endpoint < T: Committable + Send > {
    pub host: String,
    pub tx: Sender < Package < T > >,
    pub rx: Receiver < Package < T > >,
}

#[deriving(Clone, Show)]
pub struct AppendLog < T: Committable + Send > {
    pub committed_offset: uint,
    pub node_list: Vec < String >,
    pub enqueue: Option < AppendLogEntry < T > >,
}

#[deriving(Clone, Show)]
pub struct AppendLogEntry < T: Committable + Send > {
    pub offset: uint,
    pub entry: T,
}

impl < T: Committable + Send + Show > Intercommunication < T > for DefaultIntercommunication < T > {
    fn new() -> DefaultIntercommunication < T > {
        let (tx, rx) = channel();

        DefaultIntercommunication {
            senders: HashMap::new(),
            sender: tx,
            receiver: rx,
            is_debug: false,
        }
    }

    fn register(&mut self, host: String) -> Endpoint < T > {
        let (tx, rx) = channel();

        self.senders.insert(host.clone(), tx);

        Endpoint {
            host: host,
            rx: rx,
            tx: self.sender.clone(),
        }
    }

    fn receive(&mut self) -> Option < Package < T > > {
        match self.receiver.try_recv() {
            Ok(package) => Some(package),
            _ => None,
        }
    }

    fn send(&mut self, recipient: String, package: Package < T >) {
        match self.senders.find(&recipient) {
            Some(tx) => {
                // be more careful at sending
                match tx.send_opt(package) {
                    Err(_) => (),
                    _ => (),
                }
            },
            None => (),
        }
    }

    fn is_debug(&self) -> bool {
        self.is_debug
    }
}

impl < T: Committable + Send > Endpoint < T > {
    pub fn send (&self, host: String, package: PackageDetails < T >) {
        self.tx.send(Pack(self.host.clone(), host, package));
    }

    pub fn listen_block_with_timeout(&self) -> Option < Package < T > > {
        for _ in range(0, 10u) {
            match self.listen() {
                Some(x) => return Some(x),
                _ => ()
            }

            sleep(Duration::milliseconds(2));
        }

        None
    }

    pub fn listen(&self) -> Option < Package < T > > {
        match self.rx.try_recv() {
            Ok(result) => Some(result),
            _ => None,
        }
    }
}

#[deriving(Show)]
pub enum PackageDetails < T: Committable + Send > {
    Ack,

    LeaderQuery,

    // LeaderQueryResponse(leader_host)
    LeaderQueryResponse(Option < String >),

    // AppendQuery(log)
    AppendQuery(AppendLog < T >),

    // Persisted(entry_offset)
    Persisted(uint),

    // RequestVote(term
    RequestVote(uint),

    // Vote(term)
    Vote(uint),
}

#[deriving(Show)]
pub enum Package < T: Committable + Send > {
    // Pack(from, to, package)
    Pack(String, String, PackageDetails < T >),
}

pub fn start < T: Committable + Send + Clone + Show, I: Intercommunication < T > + Send >(intercommunication: I) -> Sender < int > {
    let mutex = Arc::new(Mutex::new(intercommunication));
    let (exit_tx, exit_rx) = channel();

    TaskBuilder::new().named("intercommunication").spawn(proc() {
        loop {
            let mut intercommunication = mutex.lock();

            match intercommunication.receive() {
                Some(Pack(from, to, package)) => {
                    if intercommunication.is_debug() {
                        println!("sent package {} to host {} from {}", package, to, from);
                    }

                    intercommunication.send(to.clone(), Pack(from, to, package));
                },
                None => ()
            }

            match exit_rx.try_recv() {
                Ok(_) | Err(Disconnected) => break,
                _ => (),
            }

            intercommunication.cond.signal();

            sleep(Duration::milliseconds(2));
        }
    });

    exit_tx
}
<|endoftext|>


<commit_before>use super::error::{Error, Result};
use super::util;
use ansi_term::Colour;
use difference::{Changeset, Difference};
use lazy_static::lazy_static;
use regex::bytes::{Regex as ByteRegex, RegexBuilder as ByteRegexBuilder};
use regex::Regex;
use std::char;
use std::fs;
use std::io::{self, Write};
use std::path::Path;
use std::str;
use tree_sitter::{Language, LogType, Parser, Query};

lazy_static! {
    static ref HEADER_REGEX: ByteRegex = ByteRegexBuilder::new(r"^===+\r?\n([^=]*)\r?\n===+\r?\n")
        .multi_line(true)
        .build()
        .unwrap();
    static ref DIVIDER_REGEX: ByteRegex = ByteRegexBuilder::new(r"^---+\r?\n")
        .multi_line(true)
        .build()
        .unwrap();
    static ref WHITESPACE_REGEX: Regex = Regex::new(r"\s+").unwrap();
    static ref SEXP_FIELD_REGEX: Regex = Regex::new(r" \w+: \(").unwrap();
}

#[derive(Debug, PartialEq, Eq)]
pub enum TestEntry {
    Group {
        name: String,
        children: Vec<TestEntry>,
    },
    Example {
        name: String,
        input: Vec<u8>,
        output: String,
        has_fields: bool,
    },
}

impl Default for TestEntry {
    fn default() -> Self {
        TestEntry::Group {
            name: String::new(),
            children: Vec::new(),
        }
    }
}

pub fn run_tests_at_path(
    language: Language,
    path: &Path,
    debug: bool,
    debug_graph: bool,
    filter: Option<&str>,
) -> Result<()> {
    let test_entry = parse_tests(path)?;
    let mut _log_session = None;
    let mut parser = Parser::new();
    parser.set_language(language).map_err(|e| e.to_string())?;

    if debug_graph {
        _log_session = Some(util::log_graphs(&mut parser, "log.html")?);
    } else if debug {
        parser.set_logger(Some(Box::new(|log_type, message| {
            if log_type == LogType::Lex {
                io::stderr().write(b"  ").unwrap();
            }
            write!(&mut io::stderr(), "{}\n", message).unwrap();
        })));
    }

    let mut failures = Vec::new();
    if let TestEntry::Group { children, .. } = test_entry {
        for child in children {
            run_tests(&mut parser, child, filter, 0, &mut failures)?;
        }
    }

    if failures.len() > 0 {
        println!("");

        if failures.len() == 1 {
            println!("1 failure:")
        } else {
            println!("{} failures:", failures.len())
        }

        print_diff_key();
        for (i, (name, actual, expected)) in failures.iter().enumerate() {
            println!("\n  {}. {}:", i + 1, name);
            print_diff(actual, expected);
        }
        Error::err(String::new())
    } else {
        Ok(())
    }
}

pub fn check_queries_at_path(language: Language, path: &Path) -> Result<()> {
    for entry in fs::read_dir(path)? {
        let entry = entry?;
        let hidden = entry.file_name().to_str().unwrap_or("").starts_with(".");
        if !hidden {
            let content = fs::read_to_string(entry.path()).map_err(Error::wrap(|| {
                format!("Error reading query file {:?}", entry.file_name())
            }))?;
            Query::new(language, &content).map_err(Error::wrap(|| {
                format!("Error in query file {:?}", entry.file_name())
            }))?;
        }
    }
    Ok(())
}

pub fn print_diff_key() {
    println!(
        "\n{} / {}",
        Colour::Green.paint("expected"),
        Colour::Red.paint("actual")
    );
}

pub fn print_diff(actual: &String, expected: &String) {
    let changeset = Changeset::new(actual, expected, " ");
    print!("    ");
    for diff in &changeset.diffs {
        match diff {
            Difference::Same(part) => {
                print!("{}{}", part, changeset.split);
            }
            Difference::Add(part) => {
                print!("{}{}", Colour::Green.paint(part), changeset.split);
            }
            Difference::Rem(part) => {
                print!("{}{}", Colour::Red.paint(part), changeset.split);
            }
        }
    }
    println!("");
}

fn run_tests(
    parser: &mut Parser,
    test_entry: TestEntry,
    filter: Option<&str>,
    mut indent_level: i32,
    failures: &mut Vec<(String, String, String)>,
) -> Result<()> {
    match test_entry {
        TestEntry::Example {
            name,
            input,
            output,
            has_fields,
        } => {
            if let Some(filter) = filter {
                if !name.contains(filter) {
                    return Ok(());
                }
            }
            let tree = parser.parse(&input, None).unwrap();
            let mut actual = tree.root_node().to_sexp();
            if !has_fields {
                actual = strip_sexp_fields(actual);
            }
            for _ in 0..indent_level {
                print!("  ");
            }
            if actual == output {
                println!("✓ {}", Colour::Green.paint(&name));
            } else {
                println!("✗ {}", Colour::Red.paint(&name));
                failures.push((name, actual, output));
            }
        }
        TestEntry::Group { name, children } => {
            for _ in 0..indent_level {
                print!("  ");
            }
            println!("{}:", name);
            indent_level += 1;
            for child in children {
                run_tests(parser, child, filter, indent_level, failures)?;
            }
        }
    }
    Ok(())
}

pub fn parse_tests(path: &Path) -> io::Result<TestEntry> {
    let name = path
        .file_stem()
        .and_then(|s| s.to_str())
        .unwrap_or("")
        .to_string();
    if path.is_dir() {
        let mut children = Vec::new();
        for entry in fs::read_dir(path)? {
            let entry = entry?;
            let hidden = entry.file_name().to_str().unwrap_or("").starts_with(".");
            if !hidden {
                children.push(parse_tests(&entry.path())?);
            }
        }
        Ok(TestEntry::Group { name, children })
    } else {
        let content = fs::read_to_string(path)?;
        Ok(parse_test_content(name, content))
    }
}

pub fn strip_sexp_fields(sexp: String) -> String {
    SEXP_FIELD_REGEX.replace_all(&sexp, " (").to_string()
}

fn parse_test_content(name: String, content: String) -> TestEntry {
    let mut children = Vec::new();
    let bytes = content.as_bytes();
    let mut prev_name = String::new();
    let mut prev_header_end = 0;

    // Identify all of the test descriptions using the `======` headers.
    for (header_start, header_end) in HEADER_REGEX
        .find_iter(&bytes)
        .map(|m| (m.start(), m.end()))
        .chain(Some((bytes.len(), bytes.len())))
    {
        // Find the longest line of dashes following each test description.
        // That is the divider between input and expected output.
        if prev_header_end > 0 {
            let divider_match = DIVIDER_REGEX
                .find_iter(&bytes[prev_header_end..header_start])
                .map(|m| (prev_header_end + m.start(), prev_header_end + m.end()))
                .max_by_key(|(start, end)| end - start);
            if let Some((divider_start, divider_end)) = divider_match {
                if let Ok(output) = str::from_utf8(&bytes[divider_end..header_start]) {
                    let mut input = bytes[prev_header_end..divider_start].to_vec();

                    // Remove trailing newline from the input.
                    input.pop();
                    if input.last() == Some(&b'\r') {
                        input.pop();
                    }

                    // Normalize the whitespace in the expected output.
                    let output = WHITESPACE_REGEX.replace_all(output.trim(), " ").to_string();
                    let output = output.replace(" )", ")");

                    // Identify if the expected output has fields indicated. If not, then
                    // fields will not be checked.
                    let has_fields = SEXP_FIELD_REGEX.is_match(&output);

                    children.push(TestEntry::Example {
                        name: prev_name,
                        input,
                        output,
                        has_fields,
                    });
                }
            }
        }
        prev_name = String::from_utf8_lossy(&bytes[header_start..header_end])
            .trim_matches(|c| char::is_whitespace(c) || c == '=')
            .to_string();
        prev_header_end = header_end;
    }
    TestEntry::Group { name, children }
}

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

    #[test]
    fn test_parse_test_content() {
        let entry = parse_test_content(
            "the-filename".to_string(),
            r#"
===============
The first test
===============

a b c

---

(a
    (b c))

================
The second test
================
d
---
(d)
        "#
            .trim()
            .to_string(),
        );

        assert_eq!(
            entry,
            TestEntry::Group {
                name: "the-filename".to_string(),
                children: vec![
                    TestEntry::Example {
                        name: "The first test".to_string(),
                        input: "\na b c\n".as_bytes().to_vec(),
                        output: "(a (b c))".to_string(),
                        has_fields: false,
                    },
                    TestEntry::Example {
                        name: "The second test".to_string(),
                        input: "d".as_bytes().to_vec(),
                        output: "(d)".to_string(),
                        has_fields: false,
                    },
                ]
            }
        );
    }

    #[test]
    fn test_parse_test_content_with_dashes_in_source_code() {
        let entry = parse_test_content(
            "the-filename".to_string(),
            r#"
==================
Code with dashes
==================
abc
---
defg
----
hijkl
-------

(a (b))

=========================
Code ending with dashes
=========================
abc
-----------
-------------------

(c (d))
        "#
            .trim()
            .to_string(),
        );

        assert_eq!(
            entry,
            TestEntry::Group {
                name: "the-filename".to_string(),
                children: vec![
                    TestEntry::Example {
                        name: "Code with dashes".to_string(),
                        input: "abc\n---\ndefg\n----\nhijkl".as_bytes().to_vec(),
                        output: "(a (b))".to_string(),
                        has_fields: false,
                    },
                    TestEntry::Example {
                        name: "Code ending with dashes".to_string(),
                        input: "abc\n-----------".as_bytes().to_vec(),
                        output: "(c (d))".to_string(),
                        has_fields: false,
                    },
                ]
            }
        );
    }
}
<commit_msg>Fix error in test command when queries folder does not exist<commit_after>use super::error::{Error, Result};
use super::util;
use ansi_term::Colour;
use difference::{Changeset, Difference};
use lazy_static::lazy_static;
use regex::bytes::{Regex as ByteRegex, RegexBuilder as ByteRegexBuilder};
use regex::Regex;
use std::char;
use std::fs;
use std::io::{self, Write};
use std::path::Path;
use std::str;
use tree_sitter::{Language, LogType, Parser, Query};

lazy_static! {
    static ref HEADER_REGEX: ByteRegex = ByteRegexBuilder::new(r"^===+\r?\n([^=]*)\r?\n===+\r?\n")
        .multi_line(true)
        .build()
        .unwrap();
    static ref DIVIDER_REGEX: ByteRegex = ByteRegexBuilder::new(r"^---+\r?\n")
        .multi_line(true)
        .build()
        .unwrap();
    static ref WHITESPACE_REGEX: Regex = Regex::new(r"\s+").unwrap();
    static ref SEXP_FIELD_REGEX: Regex = Regex::new(r" \w+: \(").unwrap();
}

#[derive(Debug, PartialEq, Eq)]
pub enum TestEntry {
    Group {
        name: String,
        children: Vec<TestEntry>,
    },
    Example {
        name: String,
        input: Vec<u8>,
        output: String,
        has_fields: bool,
    },
}

impl Default for TestEntry {
    fn default() -> Self {
        TestEntry::Group {
            name: String::new(),
            children: Vec::new(),
        }
    }
}

pub fn run_tests_at_path(
    language: Language,
    path: &Path,
    debug: bool,
    debug_graph: bool,
    filter: Option<&str>,
) -> Result<()> {
    let test_entry = parse_tests(path)?;
    let mut _log_session = None;
    let mut parser = Parser::new();
    parser.set_language(language).map_err(|e| e.to_string())?;

    if debug_graph {
        _log_session = Some(util::log_graphs(&mut parser, "log.html")?);
    } else if debug {
        parser.set_logger(Some(Box::new(|log_type, message| {
            if log_type == LogType::Lex {
                io::stderr().write(b"  ").unwrap();
            }
            write!(&mut io::stderr(), "{}\n", message).unwrap();
        })));
    }

    let mut failures = Vec::new();
    if let TestEntry::Group { children, .. } = test_entry {
        for child in children {
            run_tests(&mut parser, child, filter, 0, &mut failures)?;
        }
    }

    if failures.len() > 0 {
        println!("");

        if failures.len() == 1 {
            println!("1 failure:")
        } else {
            println!("{} failures:", failures.len())
        }

        print_diff_key();
        for (i, (name, actual, expected)) in failures.iter().enumerate() {
            println!("\n  {}. {}:", i + 1, name);
            print_diff(actual, expected);
        }
        Error::err(String::new())
    } else {
        Ok(())
    }
}

pub fn check_queries_at_path(language: Language, path: &Path) -> Result<()> {
    if path.exists() {
        for entry in fs::read_dir(path)? {
            let entry = entry?;
            let hidden = entry.file_name().to_str().unwrap_or("").starts_with(".");
            if !hidden {
                let content = fs::read_to_string(entry.path()).map_err(Error::wrap(|| {
                    format!("Error reading query file {:?}", entry.file_name())
                }))?;
                Query::new(language, &content).map_err(Error::wrap(|| {
                    format!("Error in query file {:?}", entry.file_name())
                }))?;
            }
        }
    }
    Ok(())
}

pub fn print_diff_key() {
    println!(
        "\n{} / {}",
        Colour::Green.paint("expected"),
        Colour::Red.paint("actual")
    );
}

pub fn print_diff(actual: &String, expected: &String) {
    let changeset = Changeset::new(actual, expected, " ");
    print!("    ");
    for diff in &changeset.diffs {
        match diff {
            Difference::Same(part) => {
                print!("{}{}", part, changeset.split);
            }
            Difference::Add(part) => {
                print!("{}{}", Colour::Green.paint(part), changeset.split);
            }
            Difference::Rem(part) => {
                print!("{}{}", Colour::Red.paint(part), changeset.split);
            }
        }
    }
    println!("");
}

fn run_tests(
    parser: &mut Parser,
    test_entry: TestEntry,
    filter: Option<&str>,
    mut indent_level: i32,
    failures: &mut Vec<(String, String, String)>,
) -> Result<()> {
    match test_entry {
        TestEntry::Example {
            name,
            input,
            output,
            has_fields,
        } => {
            if let Some(filter) = filter {
                if !name.contains(filter) {
                    return Ok(());
                }
            }
            let tree = parser.parse(&input, None).unwrap();
            let mut actual = tree.root_node().to_sexp();
            if !has_fields {
                actual = strip_sexp_fields(actual);
            }
            for _ in 0..indent_level {
                print!("  ");
            }
            if actual == output {
                println!("✓ {}", Colour::Green.paint(&name));
            } else {
                println!("✗ {}", Colour::Red.paint(&name));
                failures.push((name, actual, output));
            }
        }
        TestEntry::Group { name, children } => {
            for _ in 0..indent_level {
                print!("  ");
            }
            println!("{}:", name);
            indent_level += 1;
            for child in children {
                run_tests(parser, child, filter, indent_level, failures)?;
            }
        }
    }
    Ok(())
}

pub fn parse_tests(path: &Path) -> io::Result<TestEntry> {
    let name = path
        .file_stem()
        .and_then(|s| s.to_str())
        .unwrap_or("")
        .to_string();
    if path.is_dir() {
        let mut children = Vec::new();
        for entry in fs::read_dir(path)? {
            let entry = entry?;
            let hidden = entry.file_name().to_str().unwrap_or("").starts_with(".");
            if !hidden {
                children.push(parse_tests(&entry.path())?);
            }
        }
        Ok(TestEntry::Group { name, children })
    } else {
        let content = fs::read_to_string(path)?;
        Ok(parse_test_content(name, content))
    }
}

pub fn strip_sexp_fields(sexp: String) -> String {
    SEXP_FIELD_REGEX.replace_all(&sexp, " (").to_string()
}

fn parse_test_content(name: String, content: String) -> TestEntry {
    let mut children = Vec::new();
    let bytes = content.as_bytes();
    let mut prev_name = String::new();
    let mut prev_header_end = 0;

    // Identify all of the test descriptions using the `======` headers.
    for (header_start, header_end) in HEADER_REGEX
        .find_iter(&bytes)
        .map(|m| (m.start(), m.end()))
        .chain(Some((bytes.len(), bytes.len())))
    {
        // Find the longest line of dashes following each test description.
        // That is the divider between input and expected output.
        if prev_header_end > 0 {
            let divider_match = DIVIDER_REGEX
                .find_iter(&bytes[prev_header_end..header_start])
                .map(|m| (prev_header_end + m.start(), prev_header_end + m.end()))
                .max_by_key(|(start, end)| end - start);
            if let Some((divider_start, divider_end)) = divider_match {
                if let Ok(output) = str::from_utf8(&bytes[divider_end..header_start]) {
                    let mut input = bytes[prev_header_end..divider_start].to_vec();

                    // Remove trailing newline from the input.
                    input.pop();
                    if input.last() == Some(&b'\r') {
                        input.pop();
                    }

                    // Normalize the whitespace in the expected output.
                    let output = WHITESPACE_REGEX.replace_all(output.trim(), " ").to_string();
                    let output = output.replace(" )", ")");

                    // Identify if the expected output has fields indicated. If not, then
                    // fields will not be checked.
                    let has_fields = SEXP_FIELD_REGEX.is_match(&output);

                    children.push(TestEntry::Example {
                        name: prev_name,
                        input,
                        output,
                        has_fields,
                    });
                }
            }
        }
        prev_name = String::from_utf8_lossy(&bytes[header_start..header_end])
            .trim_matches(|c| char::is_whitespace(c) || c == '=')
            .to_string();
        prev_header_end = header_end;
    }
    TestEntry::Group { name, children }
}

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

    #[test]
    fn test_parse_test_content() {
        let entry = parse_test_content(
            "the-filename".to_string(),
            r#"
===============
The first test
===============

a b c

---

(a
    (b c))

================
The second test
================
d
---
(d)
        "#
            .trim()
            .to_string(),
        );

        assert_eq!(
            entry,
            TestEntry::Group {
                name: "the-filename".to_string(),
                children: vec![
                    TestEntry::Example {
                        name: "The first test".to_string(),
                        input: "\na b c\n".as_bytes().to_vec(),
                        output: "(a (b c))".to_string(),
                        has_fields: false,
                    },
                    TestEntry::Example {
                        name: "The second test".to_string(),
                        input: "d".as_bytes().to_vec(),
                        output: "(d)".to_string(),
                        has_fields: false,
                    },
                ]
            }
        );
    }

    #[test]
    fn test_parse_test_content_with_dashes_in_source_code() {
        let entry = parse_test_content(
            "the-filename".to_string(),
            r#"
==================
Code with dashes
==================
abc
---
defg
----
hijkl
-------

(a (b))

=========================
Code ending with dashes
=========================
abc
-----------
-------------------

(c (d))
        "#
            .trim()
            .to_string(),
        );

        assert_eq!(
            entry,
            TestEntry::Group {
                name: "the-filename".to_string(),
                children: vec![
                    TestEntry::Example {
                        name: "Code with dashes".to_string(),
                        input: "abc\n---\ndefg\n----\nhijkl".as_bytes().to_vec(),
                        output: "(a (b))".to_string(),
                        has_fields: false,
                    },
                    TestEntry::Example {
                        name: "Code ending with dashes".to_string(),
                        input: "abc\n-----------".as_bytes().to_vec(),
                        output: "(c (d))".to_string(),
                        has_fields: false,
                    },
                ]
            }
        );
    }
}
<|endoftext|>