<commit_before>//! Traits, implementations, and macros related to logging timely events.

extern crate time;

use std::cell::RefCell;
use std::io::Write;
use std::fs::File;

use ::Data;

use timely_communication::Allocate;
use ::progress::timestamp::RootTimestamp;
use ::progress::nested::product::Product;

use dataflow::scopes::root::Root;
use dataflow::operators::capture::{EventWriter, Event, EventPusher};

use abomonation::Abomonation;

static mut precise_time_ns_delta: Option<i64> = None;

/// Returns the value of an high resolution performance counter, in nanoseconds, rebased to be
/// roughly comparable to an unix timestamp.
/// Useful for comparing and merging logs from different machines (precision is limited by the
/// precision of the wall clock base; clock skew effects should be taken into consideration).
#[inline(always)]
fn get_precise_time_ns() -> u64 {
    (time::precise_time_ns() as i64 - unsafe { precise_time_ns_delta.unwrap() }) as u64
}

/// Logs `record` in `logger` if logging is enabled.
pub fn log<T: Logger>(logger: &'static ::std::thread::LocalKey<T>, record: T::Record) {
    if cfg!(feature = "logging") {
        logger.with(|x| x.log(record));
    }
}

/// Logging methods
pub trait Logger {
    /// The type of loggable record.
    type Record;
    /// Adds `record` to the log.
    fn log(&self, record: Self::Record);
    /// Called with some frequency; behavior unspecified.
    fn flush(&self);
}

/// Logs events to an underlying writer.
pub struct EventStreamLogger<T: Data, S: Write> {
    buffer: RefCell<Vec<(T,u64)>>,
    stream: RefCell<Option<EventWriter<Product<RootTimestamp, u64>, (T,u64), S>>>,
    last_time: RefCell<u64>,
}

impl<T: Data, S: Write> Logger for EventStreamLogger<T, S> {
    type Record = T;
    #[inline]
    fn log(&self, record: T) {
        self.buffer.borrow_mut().push((record, get_precise_time_ns()));
    }
    fn flush(&self) {
        // TODO : sends progress update even if nothing happened.
        // TODO : consider a minimum interval from prior update to
        // TODO : cut down on the amount on logging, at the expense
        // TODO : of freshness on the side of the reader.
        if let Some(ref mut writer) = *self.stream.borrow_mut() {
            let time = get_precise_time_ns();
            if self.buffer.borrow().len() > 0 {
                writer.push(Event::Messages(RootTimestamp::new(*self.last_time.borrow()), self.buffer.borrow().clone()));
            }
            writer.push(Event::Progress(vec![(RootTimestamp::new(*self.last_time.borrow()),-1), (RootTimestamp::new(time), 1)]));
            *self.last_time.borrow_mut() = time;
            self.buffer.borrow_mut().clear();
        }
        else {
            panic!("logging file not initialized!!!!");
        }
    }
}

impl<T: Data, S: Write> EventStreamLogger<T, S> {
    fn new() -> EventStreamLogger<T, S> {
        EventStreamLogger {
            buffer: RefCell::new(Vec::new()),
            stream: RefCell::new(None),
            last_time: RefCell::new(0),
        }
    }
    fn set(&self, stream: S) {
        let mut stream = EventWriter::new(stream);
        stream.push(Event::Progress(vec![(RootTimestamp::new(0), 1)]));
        *self.stream.borrow_mut() = Some(stream);
    }
}

impl<T: Data, S: Write> Drop for EventStreamLogger<T, S> {
    fn drop(&mut self) {
        if let Some(ref mut writer) = *self.stream.borrow_mut() {
            if self.buffer.borrow().len() > 0 {
                writer.push(Event::Messages(RootTimestamp::new(*self.last_time.borrow()), self.buffer.borrow().clone()));
            }
            writer.push(Event::Progress(vec![(RootTimestamp::new(*self.last_time.borrow()),-1)]));
        }
        else {
            panic!("logging file not initialized!!!!");
        }
    }
}

/// Initializes logging; called as part of `Root` initialization.
pub fn initialize<A: Allocate>(root: &mut Root<A>) {

    OPERATES.with(|x| x.set(File::create(format!("logs/operates-{}.abom", root.index())).unwrap()));
    CHANNELS.with(|x| x.set(File::create(format!("logs/channels-{}.abom", root.index())).unwrap()));
    MESSAGES.with(|x| x.set(File::create(format!("logs/messages-{}.abom", root.index())).unwrap()));
    PROGRESS.with(|x| x.set(File::create(format!("logs/progress-{}.abom", root.index())).unwrap()));
    SCHEDULE.with(|x| x.set(File::create(format!("logs/schedule-{}.abom", root.index())).unwrap()));
    GUARDED_MESSAGE.with(|x| x.set(File::create(format!("logs/guarded_message-{}.abom", root.index())).unwrap()));
    GUARDED_PROGRESS.with(|x| x.set(File::create(format!("logs/guarded_progress-{}.abom", root.index())).unwrap()));

    unsafe {
        precise_time_ns_delta = Some({
            let wall_time = time::get_time();
            let wall_time_ns = wall_time.nsec as i64 + wall_time.sec * 1000000000;
            time::precise_time_ns() as i64 - wall_time_ns
        });
    }

}

/// Flushes logs; called by `Root::step`.
pub fn flush_logs() {
    OPERATES.with(|x| x.flush());
    CHANNELS.with(|x| x.flush());
    PROGRESS.with(|x| x.flush());
    MESSAGES.with(|x| x.flush());
    SCHEDULE.with(|x| x.flush());
    GUARDED_MESSAGE.with(|x| x.flush());
    GUARDED_PROGRESS.with(|x| x.flush());
}

thread_local!{
    /// Logs operator creation.
    pub static OPERATES: EventStreamLogger<OperatesEvent, File> = EventStreamLogger::new();
    /// Logs channel creation.
    pub static CHANNELS: EventStreamLogger<ChannelsEvent, File> = EventStreamLogger::new();
    /// Logs progress transmission.
    pub static PROGRESS: EventStreamLogger<ProgressEvent, File> = EventStreamLogger::new();
    /// Logs message transmission.
    pub static MESSAGES: EventStreamLogger<MessagesEvent, File> = EventStreamLogger::new();
    /// Logs operator scheduling.
    pub static SCHEDULE: EventStreamLogger<ScheduleEvent, File> = EventStreamLogger::new();
    /// Logs delivery of message to an operator input.
    pub static GUARDED_MESSAGE: EventStreamLogger<bool, File> = EventStreamLogger::new();
    /// Logs delivery of notification to an operator.
    pub static GUARDED_PROGRESS: EventStreamLogger<bool, File> = EventStreamLogger::new();
}

#[derive(Debug, Clone)]
/// The creation of an `Operate` implementor.
pub struct OperatesEvent {
    /// Worker-unique identifier for the operator.
    pub id: usize,
    /// Sequence of nested scope identifiers indicating the path from the root to this instance.
    pub addr: Vec<usize>,
    /// A helpful name.
    pub name: String,
}

unsafe_abomonate!(OperatesEvent : addr, name);

#[derive(Debug, Clone)]
/// The creation of a channel between operators.
pub struct ChannelsEvent {
    /// Worker-unique identifier for the channel
    pub id: usize,
    /// Sequence of nested scope identifiers indicating the path from the root to this instance.
    pub scope_addr: Vec<usize>,
    /// Source descriptor, indicating operator index and output port.
    pub source: (usize, usize),
    /// Target descriptor, indicating operator index and input port.
    pub target: (usize, usize),
}

unsafe_abomonate!(ChannelsEvent : id, scope_addr, source, target);

#[derive(Debug, Clone)]
/// Send or receive of progress information.
pub struct ProgressEvent {
    /// `true` if the event is a send, and `false` if it is a receive.
    pub is_send: bool,
    /// Sequence of nested scope identifiers indicating the path from the root to this instance.
    pub addr: Vec<usize>,
    /// List of message updates, containing Target descriptor, timestamp as string, and delta.
    pub messages: Vec<(usize, usize, String, i64)>,
    /// List of capability updates, containing Source descriptor, timestamp as string, and delta.
    pub internal: Vec<(usize, usize, String, i64)>,
}

unsafe_abomonate!(ProgressEvent : is_send, addr, messages, internal);

#[derive(Debug, Clone)]
/// Message send or receive event
pub struct MessagesEvent {
    /// `true` if send event, `false` if receive event.
    pub is_send: bool,
    /// Channel identifier
    pub channel: usize,
    /// Source worker index.
    pub source: usize,
    /// Target worker index.
    pub target: usize,
    /// Message sequence number.
    pub seq_no: usize,
    /// Number of typed records in the message.
    pub length: usize,
}

unsafe_abomonate!(MessagesEvent);

/// Records the starting and stopping of an operator.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum StartStop {
    /// Operator starts.
    Start,
    /// Operator stops; did it have any activity?
    Stop {
        /// Did the operator perform non-trivial work.
        activity: bool
    },
}

impl Abomonation for StartStop { }

#[test]
fn start_stop_abomonation_roundtrip() {
    use abomonation::{encode, decode};
    fn check(data: StartStop) -> () {
        let mut bytes = Vec::new();
        unsafe {
            encode(&data, &mut bytes);
        }
        if let Some((result, remaining)) = unsafe { decode::<StartStop>(&mut bytes) } {
            assert!(result == &data);
            assert!(remaining.len() == 0);
        }
    }
    for data in vec![
        StartStop::Stop { activity: true },
        StartStop::Stop { activity: false },
        StartStop::Start,
    ] {
        check(data);
    }
}

#[derive(Debug, Clone)]
/// Operator start or stop.
pub struct ScheduleEvent {
    /// Worker-unique identifier for the operator, linkable to the identifiers in `OperatesEvent`.
    pub id: usize,
    /// `Start` if the operator is starting, `Stop` if it is stopping.
    /// activiy is true if it looks like some useful work was performed during this call (data was
    /// read or written, notifications were requested / delivered)
    pub start_stop: StartStop,
}

unsafe_abomonate!(ScheduleEvent : id, start_stop);
<commit_msg>Satisfy lint regarding global variable names<commit_after>//! Traits, implementations, and macros related to logging timely events.

extern crate time;

use std::cell::RefCell;
use std::io::Write;
use std::fs::File;

use ::Data;

use timely_communication::Allocate;
use ::progress::timestamp::RootTimestamp;
use ::progress::nested::product::Product;

use dataflow::scopes::root::Root;
use dataflow::operators::capture::{EventWriter, Event, EventPusher};

use abomonation::Abomonation;

static mut PRECISE_TIME_NS_DELTA: Option<i64> = None;

/// Returns the value of an high resolution performance counter, in nanoseconds, rebased to be
/// roughly comparable to an unix timestamp.
/// Useful for comparing and merging logs from different machines (precision is limited by the
/// precision of the wall clock base; clock skew effects should be taken into consideration).
#[inline(always)]
fn get_precise_time_ns() -> u64 {
    (time::precise_time_ns() as i64 - unsafe { PRECISE_TIME_NS_DELTA.unwrap() }) as u64
}

/// Logs `record` in `logger` if logging is enabled.
pub fn log<T: Logger>(logger: &'static ::std::thread::LocalKey<T>, record: T::Record) {
    if cfg!(feature = "logging") {
        logger.with(|x| x.log(record));
    }
}

/// Logging methods
pub trait Logger {
    /// The type of loggable record.
    type Record;
    /// Adds `record` to the log.
    fn log(&self, record: Self::Record);
    /// Called with some frequency; behavior unspecified.
    fn flush(&self);
}

/// Logs events to an underlying writer.
pub struct EventStreamLogger<T: Data, S: Write> {
    buffer: RefCell<Vec<(T,u64)>>,
    stream: RefCell<Option<EventWriter<Product<RootTimestamp, u64>, (T,u64), S>>>,
    last_time: RefCell<u64>,
}

impl<T: Data, S: Write> Logger for EventStreamLogger<T, S> {
    type Record = T;
    #[inline]
    fn log(&self, record: T) {
        self.buffer.borrow_mut().push((record, get_precise_time_ns()));
    }
    fn flush(&self) {
        // TODO : sends progress update even if nothing happened.
        // TODO : consider a minimum interval from prior update to
        // TODO : cut down on the amount on logging, at the expense
        // TODO : of freshness on the side of the reader.
        if let Some(ref mut writer) = *self.stream.borrow_mut() {
            let time = get_precise_time_ns();
            if self.buffer.borrow().len() > 0 {
                writer.push(Event::Messages(RootTimestamp::new(*self.last_time.borrow()), self.buffer.borrow().clone()));
            }
            writer.push(Event::Progress(vec![(RootTimestamp::new(*self.last_time.borrow()),-1), (RootTimestamp::new(time), 1)]));
            *self.last_time.borrow_mut() = time;
            self.buffer.borrow_mut().clear();
        }
        else {
            panic!("logging file not initialized!!!!");
        }
    }
}

impl<T: Data, S: Write> EventStreamLogger<T, S> {
    fn new() -> EventStreamLogger<T, S> {
        EventStreamLogger {
            buffer: RefCell::new(Vec::new()),
            stream: RefCell::new(None),
            last_time: RefCell::new(0),
        }
    }
    fn set(&self, stream: S) {
        let mut stream = EventWriter::new(stream);
        stream.push(Event::Progress(vec![(RootTimestamp::new(0), 1)]));
        *self.stream.borrow_mut() = Some(stream);
    }
}

impl<T: Data, S: Write> Drop for EventStreamLogger<T, S> {
    fn drop(&mut self) {
        if let Some(ref mut writer) = *self.stream.borrow_mut() {
            if self.buffer.borrow().len() > 0 {
                writer.push(Event::Messages(RootTimestamp::new(*self.last_time.borrow()), self.buffer.borrow().clone()));
            }
            writer.push(Event::Progress(vec![(RootTimestamp::new(*self.last_time.borrow()),-1)]));
        }
        else {
            panic!("logging file not initialized!!!!");
        }
    }
}

/// Initializes logging; called as part of `Root` initialization.
pub fn initialize<A: Allocate>(root: &mut Root<A>) {

    OPERATES.with(|x| x.set(File::create(format!("logs/operates-{}.abom", root.index())).unwrap()));
    CHANNELS.with(|x| x.set(File::create(format!("logs/channels-{}.abom", root.index())).unwrap()));
    MESSAGES.with(|x| x.set(File::create(format!("logs/messages-{}.abom", root.index())).unwrap()));
    PROGRESS.with(|x| x.set(File::create(format!("logs/progress-{}.abom", root.index())).unwrap()));
    SCHEDULE.with(|x| x.set(File::create(format!("logs/schedule-{}.abom", root.index())).unwrap()));
    GUARDED_MESSAGE.with(|x| x.set(File::create(format!("logs/guarded_message-{}.abom", root.index())).unwrap()));
    GUARDED_PROGRESS.with(|x| x.set(File::create(format!("logs/guarded_progress-{}.abom", root.index())).unwrap()));

    unsafe {
        PRECISE_TIME_NS_DELTA = Some({
            let wall_time = time::get_time();
            let wall_time_ns = wall_time.nsec as i64 + wall_time.sec * 1000000000;
            time::precise_time_ns() as i64 - wall_time_ns
        });
    }

}

/// Flushes logs; called by `Root::step`.
pub fn flush_logs() {
    OPERATES.with(|x| x.flush());
    CHANNELS.with(|x| x.flush());
    PROGRESS.with(|x| x.flush());
    MESSAGES.with(|x| x.flush());
    SCHEDULE.with(|x| x.flush());
    GUARDED_MESSAGE.with(|x| x.flush());
    GUARDED_PROGRESS.with(|x| x.flush());
}

thread_local!{
    /// Logs operator creation.
    pub static OPERATES: EventStreamLogger<OperatesEvent, File> = EventStreamLogger::new();
    /// Logs channel creation.
    pub static CHANNELS: EventStreamLogger<ChannelsEvent, File> = EventStreamLogger::new();
    /// Logs progress transmission.
    pub static PROGRESS: EventStreamLogger<ProgressEvent, File> = EventStreamLogger::new();
    /// Logs message transmission.
    pub static MESSAGES: EventStreamLogger<MessagesEvent, File> = EventStreamLogger::new();
    /// Logs operator scheduling.
    pub static SCHEDULE: EventStreamLogger<ScheduleEvent, File> = EventStreamLogger::new();
    /// Logs delivery of message to an operator input.
    pub static GUARDED_MESSAGE: EventStreamLogger<bool, File> = EventStreamLogger::new();
    /// Logs delivery of notification to an operator.
    pub static GUARDED_PROGRESS: EventStreamLogger<bool, File> = EventStreamLogger::new();
}

#[derive(Debug, Clone)]
/// The creation of an `Operate` implementor.
pub struct OperatesEvent {
    /// Worker-unique identifier for the operator.
    pub id: usize,
    /// Sequence of nested scope identifiers indicating the path from the root to this instance.
    pub addr: Vec<usize>,
    /// A helpful name.
    pub name: String,
}

unsafe_abomonate!(OperatesEvent : addr, name);

#[derive(Debug, Clone)]
/// The creation of a channel between operators.
pub struct ChannelsEvent {
    /// Worker-unique identifier for the channel
    pub id: usize,
    /// Sequence of nested scope identifiers indicating the path from the root to this instance.
    pub scope_addr: Vec<usize>,
    /// Source descriptor, indicating operator index and output port.
    pub source: (usize, usize),
    /// Target descriptor, indicating operator index and input port.
    pub target: (usize, usize),
}

unsafe_abomonate!(ChannelsEvent : id, scope_addr, source, target);

#[derive(Debug, Clone)]
/// Send or receive of progress information.
pub struct ProgressEvent {
    /// `true` if the event is a send, and `false` if it is a receive.
    pub is_send: bool,
    /// Sequence of nested scope identifiers indicating the path from the root to this instance.
    pub addr: Vec<usize>,
    /// List of message updates, containing Target descriptor, timestamp as string, and delta.
    pub messages: Vec<(usize, usize, String, i64)>,
    /// List of capability updates, containing Source descriptor, timestamp as string, and delta.
    pub internal: Vec<(usize, usize, String, i64)>,
}

unsafe_abomonate!(ProgressEvent : is_send, addr, messages, internal);

#[derive(Debug, Clone)]
/// Message send or receive event
pub struct MessagesEvent {
    /// `true` if send event, `false` if receive event.
    pub is_send: bool,
    /// Channel identifier
    pub channel: usize,
    /// Source worker index.
    pub source: usize,
    /// Target worker index.
    pub target: usize,
    /// Message sequence number.
    pub seq_no: usize,
    /// Number of typed records in the message.
    pub length: usize,
}

unsafe_abomonate!(MessagesEvent);

/// Records the starting and stopping of an operator.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum StartStop {
    /// Operator starts.
    Start,
    /// Operator stops; did it have any activity?
    Stop {
        /// Did the operator perform non-trivial work.
        activity: bool
    },
}

impl Abomonation for StartStop { }

#[test]
fn start_stop_abomonation_roundtrip() {
    use abomonation::{encode, decode};
    fn check(data: StartStop) -> () {
        let mut bytes = Vec::new();
        unsafe {
            encode(&data, &mut bytes);
        }
        if let Some((result, remaining)) = unsafe { decode::<StartStop>(&mut bytes) } {
            assert!(result == &data);
            assert!(remaining.len() == 0);
        }
    }
    for data in vec![
        StartStop::Stop { activity: true },
        StartStop::Stop { activity: false },
        StartStop::Start,
    ] {
        check(data);
    }
}

#[derive(Debug, Clone)]
/// Operator start or stop.
pub struct ScheduleEvent {
    /// Worker-unique identifier for the operator, linkable to the identifiers in `OperatesEvent`.
    pub id: usize,
    /// `Start` if the operator is starting, `Stop` if it is stopping.
    /// activiy is true if it looks like some useful work was performed during this call (data was
    /// read or written, notifications were requested / delivered)
    pub start_stop: StartStop,
}

unsafe_abomonate!(ScheduleEvent : id, start_stop);
<|endoftext|>


<commit_before>/// Custom event handlers
use crate::{
    Cmd, EditMode, InputMode, InputState, KeyCode, KeyEvent, Modifiers, Refresher, RepeatCount,
};

use radix_trie::TrieKey;
use smallvec::{smallvec, SmallVec};

/// Input event
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub enum Event {
    /// Wildcard.
    /// Useful if you want to filter out some keys.
    Any,
    /// Key sequence
    // TODO Validate 2 ?
    KeySeq(SmallVec<[KeyEvent; 2]>),
    /// TODO Mouse event
    Mouse(),
}

impl Event {
    /// See [`KeyEvent::normalize`]
    pub(crate) fn normalize(mut self) -> Self {
        if let Event::KeySeq(ref mut keys) = self {
            for key in keys.iter_mut() {
                *key = KeyEvent::normalize(*key);
            }
        }
        self
    }

    /// Return `i`th key event
    pub fn get(&self, i: usize) -> Option<&KeyEvent> {
        if let Event::KeySeq(ref ks) = self {
            ks.get(i)
        } else {
            None
        }
    }
}

impl From<KeyEvent> for Event {
    fn from(k: KeyEvent) -> Event {
        Event::KeySeq(smallvec![k])
    }
}

const BASE: u32 = 0x0010ffff + 1;
const BASE_CONTROL: u32 = 0x02000000;
const BASE_META: u32 = 0x04000000;
const BASE_SHIFT: u32 = 0x01000000;
const ESCAPE: u32 = 27;
const PAGE_UP: u32 = BASE + 1;
const PAGE_DOWN: u32 = PAGE_UP + 1;
const DOWN: u32 = PAGE_DOWN + 1;
const UP: u32 = DOWN + 1;
const LEFT: u32 = UP + 1;
const RIGHT: u32 = LEFT + 1;
const HOME: u32 = RIGHT + 1;
const END: u32 = HOME + 1;
const DELETE: u32 = END + 1;
const INSERT: u32 = DELETE + 1;
//const F1: u32 = INSERT + 1;
const MOUSE: u32 = /*F24 + 1*/ INSERT + 25;
const PASTE_START: u32 = MOUSE + 1;
const PASTE_FINISH: u32 = PASTE_START + 1;
const ANY: u32 = PASTE_FINISH + 1;

impl KeyEvent {
    fn encode(&self) -> u32 {
        let mut u = match self.0 {
            KeyCode::UnknownEscSeq => 0,
            KeyCode::Backspace => u32::from('H') | BASE_CONTROL,
            KeyCode::BackTab => u32::from('I') | BASE_CONTROL | BASE_SHIFT,
            KeyCode::BracketedPasteStart => PASTE_START,
            KeyCode::BracketedPasteEnd => PASTE_FINISH,
            KeyCode::Char(c) => u32::from(c),
            KeyCode::Delete => DELETE,
            KeyCode::Down => DOWN,
            KeyCode::End => END,
            KeyCode::Enter => u32::from('M') | BASE_CONTROL,
            KeyCode::F(i) => INSERT + i as u32,
            KeyCode::Esc => ESCAPE,
            KeyCode::Home => HOME,
            KeyCode::Insert => INSERT,
            KeyCode::Left => LEFT,
            KeyCode::Null => 0,
            KeyCode::PageDown => PAGE_DOWN,
            KeyCode::PageUp => PAGE_UP,
            KeyCode::Right => RIGHT,
            KeyCode::Tab => u32::from('I') | BASE_CONTROL,
            KeyCode::Up => UP,
        };
        if self.1.contains(Modifiers::CTRL) {
            u |= BASE_CONTROL;
        }
        if self.1.contains(Modifiers::ALT) {
            u |= BASE_META;
        }
        if self.1.contains(Modifiers::SHIFT) {
            u |= BASE_SHIFT;
        }
        u
    }
}

impl TrieKey for Event {
    fn encode_bytes(&self) -> Vec<u8> {
        match self {
            Event::Any => ANY.to_be_bytes().to_vec(),
            Event::KeySeq(keys) => {
                let mut dst = Vec::with_capacity(keys.len() * 4);
                for key in keys {
                    dst.extend_from_slice(&key.encode().to_be_bytes());
                }
                dst
            }
            Event::Mouse() => MOUSE.to_be_bytes().to_vec(),
        }
    }
}

/// Event handler
pub enum EventHandler {
    /// unconditional command
    Simple(Cmd),
    /// handler behaviour depends on input state
    Conditional(Box<dyn ConditionalEventHandler>),
    /* invoke multiple actions
     * TODO Macro(), */
}

impl From<Cmd> for EventHandler {
    fn from(c: Cmd) -> EventHandler {
        EventHandler::Simple(c)
    }
}

/// Give access to user input.
pub struct EventContext<'r> {
    mode: EditMode,
    input_mode: InputMode,
    wrt: &'r dyn Refresher,
}

impl<'r> EventContext<'r> {
    pub(crate) fn new(is: &InputState, wrt: &'r dyn Refresher) -> Self {
        EventContext {
            mode: is.mode,
            input_mode: is.input_mode,
            wrt,
        }
    }

    /// emacs or vi mode
    pub fn mode(&self) -> EditMode {
        self.mode
    }

    /// vi input mode
    pub fn input_mode(&self) -> InputMode {
        self.input_mode
    }

    /// Returns `true` if there is a hint displayed.
    pub fn has_hint(&self) -> bool {
        self.wrt.has_hint()
    }

    /// currently edited line
    pub fn line(&self) -> &str {
        self.wrt.line()
    }

    /// Current cursor position (byte position)
    pub fn pos(&self) -> usize {
        self.wrt.pos()
    }
}

/// May behave differently depending on:
///  * edit mode (emacs vs vi)
///  * vi input mode (insert vs replace vs command modes)
///  * empty line
///  * cursor position
///  * repeat count
///  * original key pressed (when same command is bound to different key)
///  * hint
///  * ...
pub trait ConditionalEventHandler: Send + Sync {
    /// Takes the current input state and
    /// returns the command to be performed or `None` to perform the default
    /// one.
    fn handle(
        &self,
        evt: &Event,
        n: RepeatCount,
        positive: bool,
        ctx: &EventContext,
    ) -> Option<Cmd>;
}

#[cfg(test)]
mod test {
    use super::{Event, EventHandler};
    use crate::{Cmd, KeyEvent};
    use radix_trie::Trie;
    use smallvec::smallvec;

    #[test]
    fn encode() {
        let mut trie = Trie::new();
        let evt = Event::KeySeq(smallvec![KeyEvent::ctrl('X'), KeyEvent::ctrl('E')]);
        trie.insert(evt.clone(), EventHandler::from(Cmd::Noop));
        let prefix = Event::from(KeyEvent::ctrl('X'));
        let subtrie = trie.get_raw_descendant(&prefix);
        assert!(subtrie.is_some());
        let subtrie = subtrie.unwrap();
        let sub_result = subtrie.get(&evt);
        assert!(sub_result.is_ok());
        assert!(sub_result.unwrap().is_some());
        let prefix = Event::from(KeyEvent::ctrl('O'));
        let subtrie = trie.get_raw_descendant(&prefix);
        assert!(subtrie.is_none())
    }
}
<commit_msg>Fix KeyEvent::encode<commit_after>/// Custom event handlers
use crate::{
    Cmd, EditMode, InputMode, InputState, KeyCode, KeyEvent, Modifiers, Refresher, RepeatCount,
};

use radix_trie::TrieKey;
use smallvec::{smallvec, SmallVec};

/// Input event
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub enum Event {
    /// Wildcard.
    /// Useful if you want to filter out some keys.
    Any,
    /// Key sequence
    // TODO Validate 2 ?
    KeySeq(SmallVec<[KeyEvent; 2]>),
    /// TODO Mouse event
    Mouse(),
}

impl Event {
    /// See [`KeyEvent::normalize`]
    pub(crate) fn normalize(mut self) -> Self {
        if let Event::KeySeq(ref mut keys) = self {
            for key in keys.iter_mut() {
                *key = KeyEvent::normalize(*key);
            }
        }
        self
    }

    /// Return `i`th key event
    pub fn get(&self, i: usize) -> Option<&KeyEvent> {
        if let Event::KeySeq(ref ks) = self {
            ks.get(i)
        } else {
            None
        }
    }
}

impl From<KeyEvent> for Event {
    fn from(k: KeyEvent) -> Event {
        Event::KeySeq(smallvec![k])
    }
}

const BASE: u32 = 0x0010ffff + 1;
const BASE_CONTROL: u32 = 0x02000000;
const BASE_META: u32 = 0x04000000;
const BASE_SHIFT: u32 = 0x01000000;
const ESCAPE: u32 = 27;
const PAGE_UP: u32 = BASE + 1;
const PAGE_DOWN: u32 = PAGE_UP + 1;
const DOWN: u32 = PAGE_DOWN + 1;
const UP: u32 = DOWN + 1;
const LEFT: u32 = UP + 1;
const RIGHT: u32 = LEFT + 1;
const HOME: u32 = RIGHT + 1;
const END: u32 = HOME + 1;
const DELETE: u32 = END + 1;
const INSERT: u32 = DELETE + 1;
//const F1: u32 = INSERT + 1;
const MOUSE: u32 = /*F24 + 1*/ INSERT + 25;
const PASTE_START: u32 = MOUSE + 1;
const PASTE_FINISH: u32 = PASTE_START + 1;
const ANY: u32 = PASTE_FINISH + 1;

impl KeyEvent {
    fn encode(&self) -> u32 {
        let mut u = match self.0 {
            KeyCode::UnknownEscSeq => 0,
            KeyCode::Backspace => u32::from('\x7f'),
            KeyCode::BackTab => u32::from('\t') | BASE_SHIFT,
            KeyCode::BracketedPasteStart => PASTE_START,
            KeyCode::BracketedPasteEnd => PASTE_FINISH,
            KeyCode::Char(c) => u32::from(c),
            KeyCode::Delete => DELETE,
            KeyCode::Down => DOWN,
            KeyCode::End => END,
            KeyCode::Enter => u32::from('\r'),
            KeyCode::F(i) => INSERT + i as u32,
            KeyCode::Esc => ESCAPE,
            KeyCode::Home => HOME,
            KeyCode::Insert => INSERT,
            KeyCode::Left => LEFT,
            KeyCode::Null => 0,
            KeyCode::PageDown => PAGE_DOWN,
            KeyCode::PageUp => PAGE_UP,
            KeyCode::Right => RIGHT,
            KeyCode::Tab => u32::from('\t'),
            KeyCode::Up => UP,
        };
        if self.1.contains(Modifiers::CTRL) {
            u |= BASE_CONTROL;
        }
        if self.1.contains(Modifiers::ALT) {
            u |= BASE_META;
        }
        if self.1.contains(Modifiers::SHIFT) {
            u |= BASE_SHIFT;
        }
        u
    }
}

impl TrieKey for Event {
    fn encode_bytes(&self) -> Vec<u8> {
        match self {
            Event::Any => ANY.to_be_bytes().to_vec(),
            Event::KeySeq(keys) => {
                let mut dst = Vec::with_capacity(keys.len() * 4);
                for key in keys {
                    dst.extend_from_slice(&key.encode().to_be_bytes());
                }
                dst
            }
            Event::Mouse() => MOUSE.to_be_bytes().to_vec(),
        }
    }
}

/// Event handler
pub enum EventHandler {
    /// unconditional command
    Simple(Cmd),
    /// handler behaviour depends on input state
    Conditional(Box<dyn ConditionalEventHandler>),
    /* invoke multiple actions
     * TODO Macro(), */
}

impl From<Cmd> for EventHandler {
    fn from(c: Cmd) -> EventHandler {
        EventHandler::Simple(c)
    }
}

/// Give access to user input.
pub struct EventContext<'r> {
    mode: EditMode,
    input_mode: InputMode,
    wrt: &'r dyn Refresher,
}

impl<'r> EventContext<'r> {
    pub(crate) fn new(is: &InputState, wrt: &'r dyn Refresher) -> Self {
        EventContext {
            mode: is.mode,
            input_mode: is.input_mode,
            wrt,
        }
    }

    /// emacs or vi mode
    pub fn mode(&self) -> EditMode {
        self.mode
    }

    /// vi input mode
    pub fn input_mode(&self) -> InputMode {
        self.input_mode
    }

    /// Returns `true` if there is a hint displayed.
    pub fn has_hint(&self) -> bool {
        self.wrt.has_hint()
    }

    /// currently edited line
    pub fn line(&self) -> &str {
        self.wrt.line()
    }

    /// Current cursor position (byte position)
    pub fn pos(&self) -> usize {
        self.wrt.pos()
    }
}

/// May behave differently depending on:
///  * edit mode (emacs vs vi)
///  * vi input mode (insert vs replace vs command modes)
///  * empty line
///  * cursor position
///  * repeat count
///  * original key pressed (when same command is bound to different key)
///  * hint
///  * ...
pub trait ConditionalEventHandler: Send + Sync {
    /// Takes the current input state and
    /// returns the command to be performed or `None` to perform the default
    /// one.
    fn handle(
        &self,
        evt: &Event,
        n: RepeatCount,
        positive: bool,
        ctx: &EventContext,
    ) -> Option<Cmd>;
}

#[cfg(test)]
mod test {
    use super::{Event, EventHandler};
    use crate::{Cmd, KeyCode, KeyEvent, Modifiers};
    use radix_trie::Trie;
    use smallvec::smallvec;

    #[test]
    fn encode() {
        let mut trie = Trie::new();
        let evt = Event::KeySeq(smallvec![KeyEvent::ctrl('X'), KeyEvent::ctrl('E')]);
        trie.insert(evt.clone(), EventHandler::from(Cmd::Noop));
        let prefix = Event::from(KeyEvent::ctrl('X'));
        let subtrie = trie.get_raw_descendant(&prefix);
        assert!(subtrie.is_some());
        let subtrie = subtrie.unwrap();
        let sub_result = subtrie.get(&evt);
        assert!(sub_result.is_ok());
        assert!(sub_result.unwrap().is_some());
        let prefix = Event::from(KeyEvent::ctrl('O'));
        let subtrie = trie.get_raw_descendant(&prefix);
        assert!(subtrie.is_none())
    }

    #[test]
    fn no_collision() {
        use {Event as E, EventHandler as H, KeyCode as C, KeyEvent as K, Modifiers as M};
        let mut trie = Trie::new();
        trie.insert(E::from(K(C::Backspace, M::NONE)), H::from(Cmd::Noop));
        trie.insert(E::from(K(C::Enter, M::NONE)), H::from(Cmd::Noop));
        trie.insert(E::from(K(C::Tab, M::NONE)), H::from(Cmd::Noop));
        trie.insert(E::from(K(C::Backspace, M::CTRL)), H::from(Cmd::Noop));
        trie.insert(E::from(K(C::Enter, M::CTRL)), H::from(Cmd::Noop));
        trie.insert(E::from(K(C::Tab, M::CTRL)), H::from(Cmd::Noop));
    }
}
<|endoftext|>


<commit_before>use {Action, InputMap};
use input::Button;
use piston_window::Size;
use std::convert::Into;
use std::default::Default;

/// Convenience object for constructing an InputMap
pub struct InputMapBuilder<A: Action> {
    input_remappings: Vec<(Button, A)>,
    x_axis_motion_inverted: bool,
    y_axis_motion_inverted: bool,
    x_axis_scroll_inverted: bool,
    y_axis_scroll_inverted: bool,
    viewport_size: Size
}

impl<A: Action> InputMapBuilder<A> {
    pub fn new() -> Self { Default::default() }

    pub fn x_scroll_inverted(mut self, invert: bool) -> Self {
        self.x_axis_scroll_inverted = invert;
        self
    }

    pub fn get_x_scroll_inverted(&self) -> &bool {
        &self.x_axis_scroll_inverted
    }

    pub fn y_scroll_inverted(mut self, invert: bool) -> Self {
        self.y_axis_scroll_inverted = invert;
        self
    }

    pub fn get_y_scroll_inverted(&self) -> &bool {
        &self.y_axis_scroll_inverted
    }

    pub fn x_motion_inverted(mut self, invert: bool) -> Self {
        self.x_axis_motion_inverted = invert;
        self
    }
    pub fn get_x_motion_inverted(&self) -> &bool {
        &self.x_axis_motion_inverted
    }
    pub fn y_motion_inverted(mut self, invert: bool) -> Self {
        self.y_axis_motion_inverted = invert;
        self
    }
    pub fn get_y_motion_inverted(&self) -> &bool {
        &self.y_axis_motion_inverted
    }

    pub fn viewport_size(mut self, sz: Size) -> Self {
        self.viewport_size = sz;
        self
    }

    pub fn get_viewport_size(&self) -> &Size {
        &self.viewport_size
    }

    pub fn with_action_mapping(mut self, but: Button, act: A) -> Self {
        self.input_remappings.push((but, act));
        self
    }
}

impl<A: Action> Default for InputMapBuilder<A> {
    fn default() -> Self {
        InputMapBuilder {
            input_remappings: vec![],
            x_axis_motion_inverted: false,
            y_axis_motion_inverted: false,
            x_axis_scroll_inverted: false,
            y_axis_scroll_inverted: false,
            viewport_size: (800, 600).into()
        }
    }
}

impl<A: Action> Into<InputMap<A>> for InputMapBuilder<A> {
    fn into(self) -> InputMap<A> {
        let mut imap = InputMap::new(self.viewport_size);

        imap.mouse_translator.x_axis_motion_inverted = self.x_axis_motion_inverted;
        imap.mouse_translator.y_axis_motion_inverted = self.y_axis_motion_inverted;
        imap.mouse_translator.x_axis_scroll_inverted = self.x_axis_scroll_inverted;
        imap.mouse_translator.y_axis_scroll_inverted = self.y_axis_scroll_inverted;

        //TODO: set key remappings
        imap.keymap.btn_map.reserve(self.input_remappings.len());

        imap
    }
}

<commit_msg>Rename imap var to input_map for InputMap::into, and add a build function which calls into()<commit_after>use {Action, InputMap};
use input::Button;
use piston_window::Size;
use std::convert::Into;
use std::default::Default;

/// Convenience object for constructing an InputMap
pub struct InputMapBuilder<A: Action> {
    input_remappings: Vec<(Button, A)>,
    x_axis_motion_inverted: bool,
    y_axis_motion_inverted: bool,
    x_axis_scroll_inverted: bool,
    y_axis_scroll_inverted: bool,
    viewport_size: Size
}

impl<A: Action> InputMapBuilder<A> {
    pub fn new() -> Self { Default::default() }

    pub fn x_scroll_inverted(mut self, invert: bool) -> Self {
        self.x_axis_scroll_inverted = invert;
        self
    }

    pub fn get_x_scroll_inverted(&self) -> &bool {
        &self.x_axis_scroll_inverted
    }

    pub fn y_scroll_inverted(mut self, invert: bool) -> Self {
        self.y_axis_scroll_inverted = invert;
        self
    }

    pub fn get_y_scroll_inverted(&self) -> &bool {
        &self.y_axis_scroll_inverted
    }

    pub fn x_motion_inverted(mut self, invert: bool) -> Self {
        self.x_axis_motion_inverted = invert;
        self
    }
    pub fn get_x_motion_inverted(&self) -> &bool {
        &self.x_axis_motion_inverted
    }
    pub fn y_motion_inverted(mut self, invert: bool) -> Self {
        self.y_axis_motion_inverted = invert;
        self
    }
    pub fn get_y_motion_inverted(&self) -> &bool {
        &self.y_axis_motion_inverted
    }

    pub fn viewport_size(mut self, sz: Size) -> Self {
        self.viewport_size = sz;
        self
    }

    pub fn get_viewport_size(&self) -> &Size {
        &self.viewport_size
    }

    pub fn with_action_mapping(mut self, but: Button, act: A) -> Self {
        self.input_remappings.push((but, act));
        self
    }

    pub fn build(self) -> InputMap<A> {self.into()}
}

impl<A: Action> Default for InputMapBuilder<A> {
    fn default() -> Self {
        InputMapBuilder {
            input_remappings: vec![],
            x_axis_motion_inverted: false,
            y_axis_motion_inverted: false,
            x_axis_scroll_inverted: false,
            y_axis_scroll_inverted: false,
            viewport_size: (800, 600).into()
        }
    }
}

impl<A: Action> Into<InputMap<A>> for InputMapBuilder<A> {
    fn into(self) -> InputMap<A> {
        let mut input_map = InputMap::new(self.viewport_size);

        input_map.mouse_translator.x_axis_motion_inverted = self.x_axis_motion_inverted;
        input_map.mouse_translator.y_axis_motion_inverted = self.y_axis_motion_inverted;
        input_map.mouse_translator.x_axis_scroll_inverted = self.x_axis_scroll_inverted;
        input_map.mouse_translator.y_axis_scroll_inverted = self.y_axis_scroll_inverted;

        //TODO: set key remappings
        input_map.keymap.btn_map.reserve(self.input_remappings.len());

        input_map
    }
}

<|endoftext|>


<commit_before>// Copyright 2013-2015, The Gtk-rs Project Developers.
// See the COPYRIGHT file at the top-level directory of this distribution.
// Licensed under the MIT license, see the LICENSE file or <http://opensource.org/licenses/MIT>

#![cfg_attr(not(feature = "v3_10"), allow(unused_imports))]

use libc::{c_char, ssize_t};

use glib::Object;
use glib::object::{Downcast, IsA};
use glib::translate::*;
use ffi;

glib_wrapper! {
    pub struct Builder(Object<ffi::GtkBuilder>);

    match fn {
        get_type => || ffi::gtk_builder_get_type(),
    }
}

impl Builder {
    pub fn new() -> Builder {
        assert_initialized_main_thread!();
        unsafe { from_glib_full(ffi::gtk_builder_new()) }
    }

    #[cfg(feature = "v3_10")]
    pub fn new_from_file(file_name: &str) -> Builder {
        assert_initialized_main_thread!();
        unsafe { from_glib_full(ffi::gtk_builder_new_from_file(file_name.to_glib_none().0)) }
    }

    #[cfg(feature = "v3_10")]
    pub fn new_from_resource(resource_path: &str) -> Builder {
        assert_initialized_main_thread!();
        unsafe {
            from_glib_full(ffi::gtk_builder_new_from_resource(resource_path.to_glib_none().0))
        }
    }

    #[cfg(feature = "v3_10")]
    pub fn new_from_string(string: &str) -> Builder {
        assert_initialized_main_thread!();
        unsafe {
            // Don't need a null-terminated string here
            from_glib_full(
                ffi::gtk_builder_new_from_string(string.as_ptr() as *const c_char,
                    string.len() as ssize_t))
        }
    }

    pub fn get_object<T: IsA<Object>>(&self, name: &str) -> Option<T> {
        unsafe {
            Option::<Object>::from_glib_none(
                ffi::gtk_builder_get_object(self.to_glib_none().0, name.to_glib_none().0))
                .and_then(|obj| obj.downcast().ok())
        }
    }
}
<commit_msg>Add methods to Builder object<commit_after>// Copyright 2013-2015, The Gtk-rs Project Developers.
// See the COPYRIGHT file at the top-level directory of this distribution.
// Licensed under the MIT license, see the LICENSE file or <http://opensource.org/licenses/MIT>

#![cfg_attr(not(feature = "v3_10"), allow(unused_imports))]

use libc::{c_char, ssize_t};

use glib::Object;
use glib::object::{Downcast, IsA};
use glib::translate::*;
use ffi;

glib_wrapper! {
    pub struct Builder(Object<ffi::GtkBuilder>);

    match fn {
        get_type => || ffi::gtk_builder_get_type(),
    }
}

impl Builder {
    pub fn new() -> Builder {
        assert_initialized_main_thread!();
        unsafe { from_glib_full(ffi::gtk_builder_new()) }
    }

    #[cfg(feature = "v3_10")]
    pub fn new_from_file(file_name: &str) -> Builder {
        assert_initialized_main_thread!();
        unsafe { from_glib_full(ffi::gtk_builder_new_from_file(file_name.to_glib_none().0)) }
    }

    #[cfg(feature = "v3_10")]
    pub fn new_from_resource(resource_path: &str) -> Builder {
        assert_initialized_main_thread!();
        unsafe {
            from_glib_full(ffi::gtk_builder_new_from_resource(resource_path.to_glib_none().0))
        }
    }

    #[cfg(feature = "v3_10")]
    pub fn new_from_string(string: &str) -> Builder {
        assert_initialized_main_thread!();
        unsafe {
            // Don't need a null-terminated string here
            from_glib_full(
                ffi::gtk_builder_new_from_string(string.as_ptr() as *const c_char,
                    string.len() as ssize_t))
        }
    }

    pub fn get_object<T: IsA<Object>>(&self, name: &str) -> Option<T> {
        unsafe {
            Option::<Object>::from_glib_none(
                ffi::gtk_builder_get_object(self.to_glib_none().0, name.to_glib_none().0))
                .and_then(|obj| obj.downcast().ok())
        }
    }

    pub fn add_from_file(&self, file_name: &str) -> bool {
        unsafe {
            let mut ptr = ::std::ptr::null_mut();
            ffi::gtk_builder_add_from_file(self.to_glib_none().0,
                                           file_name.to_glib_none().0,
                                           &mut ptr) > 0
        }
    }

    pub fn add_from_resource(&self, resource_path: &str) -> bool {
        unsafe {
            let mut ptr = ::std::ptr::null_mut();
            ffi::gtk_builder_add_from_resource(self.to_glib_none().0,
                                               resource_path.to_glib_none().0,
                                               &mut ptr) > 0
        }
    }

    pub fn add_from_string(&self, string: &str) -> bool {
        unsafe {
            let mut ptr = ::std::ptr::null_mut();
            ffi::gtk_builder_add_from_string(self.to_glib_none().0,
                                             string.as_ptr() as *const c_char,
                                             string.len(), &mut ptr) > 0
        }
    }
}
<|endoftext|>


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

use {Opcode, ResponseCode, Header, QueryType, QueryClass};

/// Allows to build a DNS packet
///
/// Both query and answer packets may be built with this interface, although,
/// much of functionality is not implemented yet.
pub struct Builder {
    buf: Vec<u8>,
}

impl Builder {
    /// Creates a new query
    ///
    /// Initially all sections are empty. You're expected to fill
    /// the questions section with `add_question`
    pub fn new_query(id: u16, recursion: bool) -> Builder {
        let mut buf = Vec::with_capacity(512);
        let head = Header {
            id: id,
            query: true,
            opcode: Opcode::StandardQuery,
            authoritative: false,
            truncated: false,
            recursion_desired: recursion,
            recursion_available: false,
            response_code: ResponseCode::NoError,
            questions: 0,
            answers: 0,
            nameservers: 0,
            additional: 0,
        };
        buf.extend([0u8; 12].iter());
        head.write(&mut buf[..12]);
        Builder { buf: buf }
    }
    /// Adds a question to the packet
    ///
    /// # Panics
    ///
    /// * Answers, nameservers or additional section has already been written
    /// * There are already 65535 questions in the buffer.
    /// * When name is invalid
    pub fn add_question(&mut self, qname: &str,
        qtype: QueryType, qclass: QueryClass)
        -> &mut Builder
    {
        println!("BUF {:?}", self.buf);
        if &self.buf[6..12] != b"\x00\x00\x00\x00\x00\x00" {
            panic!("Too late to add a question");
        }
        self.write_name(qname);
        self.buf.write_u16::<BigEndian>(qtype as u16).unwrap();
        self.buf.write_u16::<BigEndian>(qclass as u16).unwrap();
        let oldq = BigEndian::read_u16(&self.buf[4..6]);
        if oldq == 65535 {
            panic!("Too many questions");
        }
        BigEndian::write_u16(&mut self.buf[4..6], oldq+1);
        self
    }
    fn write_name(&mut self, name: &str) {
        for part in name.split('.') {
            assert!(part.len() < 63);
            let ln = part.len() as u8;
            self.buf.push(ln);
            self.buf.extend(part.as_bytes());
        }
        self.buf.push(0);
    }
    /// Returns the final packet
    ///
    /// When packet is not truncated method returns `Ok(packet)`. If
    /// packet is truncated the method returns `Err(packet)`. In both
    /// cases the packet is fully valid.
    ///
    /// In the server implementation you may use
    /// `x.build().unwrap_or_else(|x| x)`.
    ///
    /// In the client implementation it's probably unwise to send truncated
    /// packet, as it doesn't make sense. Even panicking may be more
    /// appropriate.
    // TODO(tailhook) does the truncation make sense for TCP, and how
    // to treat it for EDNS0?
    pub fn build(mut self) -> Result<Vec<u8>,Vec<u8>> {
        // TODO(tailhook) optimize labels
        if self.buf.len() > 512 {
            Header::set_truncated(&mut self.buf[..12]);
            Err(self.buf)
        } else {
            Ok(self.buf)
        }
    }
}

#[cfg(test)]
mod test {
    use QueryType as QT;
    use QueryClass as QC;
    use super::Builder;

    #[test]
    fn build_query() {
        let mut bld = Builder::new_query(1573, true);
        bld.add_question("example.com", QT::A, QC::IN);
        let result = b"\x06%\x01\x00\x00\x01\x00\x00\x00\x00\x00\x00\
                      \x07example\x03com\x00\x00\x01\x00\x01";
        assert_eq!(&bld.build().unwrap()[..], &result[..]);
    }
}
<commit_msg>Remove debugging print<commit_after>use byteorder::{ByteOrder, BigEndian, WriteBytesExt};

use {Opcode, ResponseCode, Header, QueryType, QueryClass};

/// Allows to build a DNS packet
///
/// Both query and answer packets may be built with this interface, although,
/// much of functionality is not implemented yet.
pub struct Builder {
    buf: Vec<u8>,
}

impl Builder {
    /// Creates a new query
    ///
    /// Initially all sections are empty. You're expected to fill
    /// the questions section with `add_question`
    pub fn new_query(id: u16, recursion: bool) -> Builder {
        let mut buf = Vec::with_capacity(512);
        let head = Header {
            id: id,
            query: true,
            opcode: Opcode::StandardQuery,
            authoritative: false,
            truncated: false,
            recursion_desired: recursion,
            recursion_available: false,
            response_code: ResponseCode::NoError,
            questions: 0,
            answers: 0,
            nameservers: 0,
            additional: 0,
        };
        buf.extend([0u8; 12].iter());
        head.write(&mut buf[..12]);
        Builder { buf: buf }
    }
    /// Adds a question to the packet
    ///
    /// # Panics
    ///
    /// * Answers, nameservers or additional section has already been written
    /// * There are already 65535 questions in the buffer.
    /// * When name is invalid
    pub fn add_question(&mut self, qname: &str,
        qtype: QueryType, qclass: QueryClass)
        -> &mut Builder
    {
        if &self.buf[6..12] != b"\x00\x00\x00\x00\x00\x00" {
            panic!("Too late to add a question");
        }
        self.write_name(qname);
        self.buf.write_u16::<BigEndian>(qtype as u16).unwrap();
        self.buf.write_u16::<BigEndian>(qclass as u16).unwrap();
        let oldq = BigEndian::read_u16(&self.buf[4..6]);
        if oldq == 65535 {
            panic!("Too many questions");
        }
        BigEndian::write_u16(&mut self.buf[4..6], oldq+1);
        self
    }
    fn write_name(&mut self, name: &str) {
        for part in name.split('.') {
            assert!(part.len() < 63);
            let ln = part.len() as u8;
            self.buf.push(ln);
            self.buf.extend(part.as_bytes());
        }
        self.buf.push(0);
    }
    /// Returns the final packet
    ///
    /// When packet is not truncated method returns `Ok(packet)`. If
    /// packet is truncated the method returns `Err(packet)`. In both
    /// cases the packet is fully valid.
    ///
    /// In the server implementation you may use
    /// `x.build().unwrap_or_else(|x| x)`.
    ///
    /// In the client implementation it's probably unwise to send truncated
    /// packet, as it doesn't make sense. Even panicking may be more
    /// appropriate.
    // TODO(tailhook) does the truncation make sense for TCP, and how
    // to treat it for EDNS0?
    pub fn build(mut self) -> Result<Vec<u8>,Vec<u8>> {
        // TODO(tailhook) optimize labels
        if self.buf.len() > 512 {
            Header::set_truncated(&mut self.buf[..12]);
            Err(self.buf)
        } else {
            Ok(self.buf)
        }
    }
}

#[cfg(test)]
mod test {
    use QueryType as QT;
    use QueryClass as QC;
    use super::Builder;

    #[test]
    fn build_query() {
        let mut bld = Builder::new_query(1573, true);
        bld.add_question("example.com", QT::A, QC::IN);
        let result = b"\x06%\x01\x00\x00\x01\x00\x00\x00\x00\x00\x00\
                      \x07example\x03com\x00\x00\x01\x00\x01";
        assert_eq!(&bld.build().unwrap()[..], &result[..]);
    }
}
<|endoftext|>


<commit_before>/// `start` and `end` are 0-based from LSB (least significant bit).
/// `start` is inclusive and `end` is exclusive.
pub fn get_bits(bits: u32, start: usize, end: usize) -> u32 {
    // !0 is equal to 0xffffffff
    let mask = !(!0 << (end - start + 1));
    (bits >> start) & mask
}

pub fn set_bit(bits: u32, bit_index: usize) -> u32 {
    let mask = 1 << bit_index;
    bits | mask
}

pub fn clear_bit(bits: u32, bit_index: usize) -> u32 {
    let mask = !(1 << bit_index);
    bits & mask
}

pub fn toggle_bit(bits: u32, bit_index: usize) -> u32 {
    let mask = 1 << bit_index;
    bits ^ mask
}

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

    #[test]
    fn test_get_bits() {
        assert_eq!(0xff, get_bits(0xff000000, 24, 32));
    }

    #[test]
    fn test_set_bit() {
        assert_eq!(0b1011, set_bit(0b1010, 0));
        assert_eq!(0b1011, set_bit(0b1011, 0));
        assert_eq!(0b1010, set_bit(0b0010, 3));
    }

    #[test]
    fn test_clear_bit() {
        assert_eq!(0b1010, clear_bit(0b1011, 0));
        assert_eq!(0b1010, clear_bit(0b1010, 0));
        assert_eq!(0b0010, clear_bit(0b1010, 3));
    }

    #[test]
    fn test_toggle_bit() {
        assert_eq!(0b1011, toggle_bit(0b1010, 0));
        assert_eq!(0b1010, toggle_bit(0b1011, 0));
        assert_eq!(0b1010, toggle_bit(0b0010, 3));
    }
}
<commit_msg>adjusted bit_ops to be inclusive, added more docs<commit_after>/// Returns the range of bits from index `start` to index `end`.
///
/// `start` and `end` are both inclusive and 0-based from LSB (least significant bit).
pub fn get_bits(bits: u32, start: usize, end: usize) -> u32 {
    // !0 is equal to 0xffffffff
    let range = (end + 1) - start;
    let mask = !(!0 << range);
    (bits >> start) & mask
}

/// Returns the value of the bit at `bit_index`.
///
/// `bit_index` is 0-based from LSB (least significant bit).
pub fn get_bit(bits: u32, bit_index: usize) -> u32 {
    get_bits(bits, bit_index, bit_index)
}

/// Sets the bit at `bit_index` to 1.
///
/// `bit_index` is 0-based from LSB (least significant bit).
pub fn set_bit(bits: u32, bit_index: usize) -> u32 {
    let mask = 1 << bit_index;
    bits | mask
}

/// Sets the bit at `bit_index` to 0.
///
/// `bit_index` is 0-based from LSB (least significant bit).
pub fn clear_bit(bits: u32, bit_index: usize) -> u32 {
    let mask = !(1 << bit_index);
    bits & mask
}

/// Flips the bit at `bit_index` to 0.
///
/// `bit_index` is 0-based from LSB (least significant bit).
pub fn toggle_bit(bits: u32, bit_index: usize) -> u32 {
    let mask = 1 << bit_index;
    bits ^ mask
}

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

    #[test]
    fn test_get_bits() {
        assert_eq!(0xff, get_bits(0xff000000, 24, 31));
        assert_eq!(0b_101, get_bits(0b110101, 0, 2));

        let bit_32 = 0b1_00111010_01101000110011111000000;
        assert_eq!(0b1, get_bits(bit_32, 31, 31));
        assert_eq!(0b1, get_bit(bit_32, 31));
        assert_eq!(0b00111010, get_bits(bit_32, 23, 30));
        assert_eq!(0b01101000110011111000000, get_bits(bit_32, 0, 22));

        let bit_32 = 0b0_01011110_11001011011100110101010;
        assert_eq!(0b0, get_bits(bit_32, 31, 31));
        assert_eq!(0b0, get_bit(bit_32, 31));
        assert_eq!(0b01011110, get_bits(bit_32, 23, 30));
        assert_eq!(0b11001011011100110101010, get_bits(bit_32, 0, 22));
    }

    #[test]
    fn test_set_bit() {
        assert_eq!(0b1011, set_bit(0b1010, 0));
        assert_eq!(0b1011, set_bit(0b1011, 0));
        assert_eq!(0b1010, set_bit(0b0010, 3));
    }

    #[test]
    fn test_clear_bit() {
        assert_eq!(0b1010, clear_bit(0b1011, 0));
        assert_eq!(0b1010, clear_bit(0b1010, 0));
        assert_eq!(0b0010, clear_bit(0b1010, 3));
    }

    #[test]
    fn test_toggle_bit() {
        assert_eq!(0b1011, toggle_bit(0b1010, 0));
        assert_eq!(0b1010, toggle_bit(0b1011, 0));
        assert_eq!(0b1010, toggle_bit(0b0010, 3));
    }
}
<|endoftext|>


<commit_before>// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use std::fmt;

use {Args, Format};

/// A message that has been localized and can be formatted in a
/// locale-aware manner.
///
/// While a `Message` can be created directly from [`Format`]
/// components, it is easiest to create it from [`icu::parse`].
///
/// A message can be formatted, returning a `String`:
///
/// ```
/// use message_format::*;
///
/// let m = icu::parse("{name} went to {place}.").unwrap();
/// assert_eq!(&m.format_message(&arg("name", "Jacob").arg("place", "the store")),
///            "Jacob went to the store.");
/// ```
///
/// It can also be written to a stream, but this is more cumbersome:
///
///
/// ```
/// use message_format::*;
///
/// let m = icu::parse("{name} went to {place}.").unwrap();
/// let mut output = String::new();
/// m.write_message(&mut output, &arg("name", "Jacob").arg("place", "the store")).unwrap();
/// assert_eq!(output, "Jacob went to the store.");
/// ```
///
/// [`Format`]: trait.Format.html
/// [`icu::parse`]: icu/fn.parse.html
pub struct Message {
    #[doc(hidden)]
    pub parts: Vec<Box<Format>>,
}

impl Message {
    /// Construct a message from constituent parts.
    pub fn new(parts: Vec<Box<Format>>) -> Self {
        Message { parts: parts }
    }

    /// Format a message, returning a string.
    pub fn format_message<'f>(&'f self, args: &'f Args<'f>) -> String {
        let mut output = String::new();
        let _ = self.write_message(&mut output, args);
        output
    }

    /// Write a message to a stream.
    pub fn write_message<'f>(&'f self, stream: &mut fmt::Write, args: &'f Args<'f>) -> fmt::Result {
        for part in &self.parts {
            try!(part.apply_format(stream, args));
        }
        Ok(())
    }
}
<commit_msg>Message: explain why `parts` is `pub`.<commit_after>// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use std::fmt;

use {Args, Format};

/// A message that has been localized and can be formatted in a
/// locale-aware manner.
///
/// While a `Message` can be created directly from [`Format`]
/// components, it is easiest to create it from [`icu::parse`].
///
/// A message can be formatted, returning a `String`:
///
/// ```
/// use message_format::*;
///
/// let m = icu::parse("{name} went to {place}.").unwrap();
/// assert_eq!(&m.format_message(&arg("name", "Jacob").arg("place", "the store")),
///            "Jacob went to the store.");
/// ```
///
/// It can also be written to a stream, but this is more cumbersome:
///
///
/// ```
/// use message_format::*;
///
/// let m = icu::parse("{name} went to {place}.").unwrap();
/// let mut output = String::new();
/// m.write_message(&mut output, &arg("name", "Jacob").arg("place", "the store")).unwrap();
/// assert_eq!(output, "Jacob went to the store.");
/// ```
///
/// [`Format`]: trait.Format.html
/// [`icu::parse`]: icu/fn.parse.html
pub struct Message {
    // This is pub due to icu::ast::PluralFormat. Once we address that, we
    // can make this private again.
    #[doc(hidden)]
    pub parts: Vec<Box<Format>>,
}

impl Message {
    /// Construct a message from constituent parts.
    pub fn new(parts: Vec<Box<Format>>) -> Self {
        Message { parts: parts }
    }

    /// Format a message, returning a string.
    pub fn format_message<'f>(&'f self, args: &'f Args<'f>) -> String {
        let mut output = String::new();
        let _ = self.write_message(&mut output, args);
        output
    }

    /// Write a message to a stream.
    pub fn write_message<'f>(&'f self, stream: &mut fmt::Write, args: &'f Args<'f>) -> fmt::Result {
        for part in &self.parts {
            try!(part.apply_format(stream, args));
        }
        Ok(())
    }
}
<|endoftext|>


<commit_before>use address::Address;
use messages::outgoing;
use mio;
use node;
use routing_table::{InsertOutcome, InsertionResult, RoutingTable};
use std::collections::HashMap;
use std::io::Cursor;
use std::net::{UdpSocket, ToSocketAddrs, SocketAddr};
use std::sync::mpsc;
use std::thread;
use stun;
use transaction::{TransactionId, TransactionIdGenerator};

#[derive(Debug)]
pub enum Event {
    ReceivedPacket(Address, Vec<u8>),
    Shutdown
}

pub enum OneshotTask {
    Incoming(Vec<u8>),
    StartBootstrap,
    SendPacket(Address, Vec<u8>),
    Shutdown
}

pub enum ScheduledTask {
    ContinueBootstrap,
    ContinueHealthCheck,
    ContinueRefresh
}

enum TableAction {
    Bootstrap(mio::Timeout),
    HealthCheck(mio::Timeout),
    RefreshBucket(mio::Timeout)
}

pub type TaskSender = mio::Sender<OneshotTask>;

enum Status {
    Bootstrapping,
    Idle
}

pub struct Network {
    host: SocketAddr,
    routing_table: RoutingTable,
    self_node: Box<node::Node + 'static>,
    transaction_ids: TransactionIdGenerator,
    status: Status,
    pending_actions: HashMap<TransactionId, TableAction>,
    event_listeners: Vec<mpsc::Sender<Event>>,
    shutdown_signals: Vec<mpsc::Sender<mpsc::Sender<()>>>
}

impl Network {
    pub fn new<T: ToSocketAddrs>(self_address: Address, host: T, routers: Vec<Box<node::Node>>) -> Network {
        let host = host.to_socket_addrs().unwrap().next().unwrap();
        let mapped_host = stun::get_mapped_address(host).unwrap();
        let self_node = node::UdpNode::new(self_address, mapped_host);
        let routing_table = RoutingTable::new(8, self_address, routers);

        Network {
            host: host,
            routing_table: routing_table,
            self_node: Box::new(self_node),
            transaction_ids: TransactionIdGenerator::new(),
            status: Status::Idle,
            pending_actions: HashMap::new(),
            event_listeners: vec![],
            shutdown_signals: vec![]
        }
    }

    pub fn run(mut self) -> TaskSender {
        let mut event_loop_config = mio::EventLoopConfig::new();
        event_loop_config.notify_capacity(16384);
        let mut event_loop = mio::EventLoop::configured(event_loop_config).unwrap();

        let shutdown_signal = create_incoming_udp_channel(self.host, event_loop.channel());
        self.shutdown_signals.push(shutdown_signal);

        event_loop.channel().send(OneshotTask::StartBootstrap).unwrap();
        info!("Running server at {:?}", self.self_node);
        let mut handler = Handler::new(self);
        let task_sender = event_loop.channel();
        thread::spawn(move|| event_loop.run(&mut handler).unwrap());
        task_sender
    }

    pub fn register_event_listener(&mut self, event_listener: mpsc::Sender<Event>) {
        self.event_listeners.push(event_listener);
    }

    fn handle_incoming(&mut self, data: Vec<u8>, event_loop: &mut mio::EventLoop<Handler>) {
        use messages::incoming::{Message, Query, Response, self};
        let mut data = Cursor::new(data);
        let message = incoming::parse_from_reader(&mut data).unwrap();

        match message {
            Message::Query(transaction_id, origin, query) => {
                let origin_address = origin.address();
                match query {
                    Query::FindNode(target) => {
                        let response: Vec<u8> = outgoing::create_find_node_response(
                                transaction_id,
                                &self.self_node,
                                self.routing_table.nearest_to(&target, false));
                        origin.send(response);

                        self.insert_node(origin).unwrap();
                        {
                            if let Some(origin) = self.routing_table.find_node(&origin_address) {
                                origin.received_query(transaction_id);
                            }
                        }
                    },
                    Query::Packet(payload) => {
                        for listener in &self.event_listeners {
                            listener.send(Event::ReceivedPacket(origin_address, payload.clone())).unwrap();
                        }
                        let response = outgoing::create_packet_response(
                            transaction_id, &self.self_node);
                        origin.send(response);
                    },
                    Query::Ping => {
                        let response = outgoing::create_ping_response(
                            transaction_id,
                            &self.self_node);
                        origin.send(response);

                        self.insert_node(origin).unwrap();
                        if let Some(origin) = self.routing_table.find_node(&origin_address) {
                            origin.received_query(transaction_id);
                        }
                    }
                }
            }
            Message::Response(transaction_id, origin, response) => {
                let origin_address = origin.address();
                match response {
                    Response::FindNode(mut nodes) => {
                        let mut encounted_new_node = false;
                        let mut tail = vec![origin];
                        for node in nodes.drain(..).chain(tail.drain(..)) {
                            match self.insert_node(node) {
                                Ok(InsertOutcome::Inserted) => {
                                    encounted_new_node = true;
                                }
                                Err(error) => { panic!(error) }
                                _ => { }
                            }
                        }

                        {
                            if let Some(mut origin) = self.routing_table.find_node(&origin_address) {
                                origin.received_response(transaction_id);
                            }
                        }

                        match self.pending_actions.remove(&transaction_id) {
                            Some(TableAction::Bootstrap(timeout)) => {
                                event_loop.clear_timeout(timeout);
                                if encounted_new_node {
                                    self.continue_bootstrap(event_loop);
                                } else {
                                    self.continue_health_check(event_loop);
                                    self.continue_refresh(event_loop)
                                }
                            }
                            Some(TableAction::HealthCheck(_)) => {
                            },
                            Some(TableAction::RefreshBucket(_)) => {
                            }
                            None => { }
                        }

                        match self.status {
                            Status::Bootstrapping => {
                                self.status = Status::Idle;
                            }
                            _ => { }
                        }
                    }
                    Response::Packet => {
                        if let Some(mut origin) = self.routing_table.find_node(&origin_address) {
                            origin.received_response(transaction_id);
                        }
                    }
                    Response::Ping => {
                        if let Some(mut origin) = self.routing_table.find_node(&origin_address) {
                            origin.received_response(transaction_id);
                        }
                    }
                }
            }
        }
    }

    fn start_bootstrap(&mut self, event_loop: &mut mio::EventLoop<Handler>) {
        self.status = Status::Bootstrapping;
        self.continue_bootstrap(event_loop);
    }

    fn continue_bootstrap(&mut self, event_loop: &mut mio::EventLoop<Handler>) {
        let address = &self.self_node.address();
        let transaction_id = self.find_node(address);
        let timeout = event_loop.timeout_ms(ScheduledTask::ContinueBootstrap, 1000).unwrap();
        self.pending_actions.insert(transaction_id, TableAction::Bootstrap(timeout));
    }

    fn continue_health_check(&mut self, event_loop: &mut mio::EventLoop<Handler>) {
        let transaction_id = self.health_check();
        let timeout = event_loop.timeout_ms(ScheduledTask::ContinueHealthCheck, 1000).unwrap();
        self.pending_actions.insert(transaction_id, TableAction::HealthCheck(timeout));
    }

    fn continue_refresh(&mut self, event_loop: &mut mio::EventLoop<Handler>) {
        let timeout = event_loop.timeout_ms(ScheduledTask::ContinueRefresh, 1000).unwrap();

        match self.refresh_bucket() {
            Some(transaction_id) => {
                self.pending_actions.insert(transaction_id, TableAction::RefreshBucket(timeout));
            }
            None => { }
        }
    }

    fn find_node(&mut self, address: &Address) -> TransactionId {
        let transaction_id = self.transaction_ids.generate();
        let query = outgoing::create_find_node_query(
            transaction_id,
            &self.self_node,
            address);
        for node in self.routing_table.nearest() {
            node.send(query.clone());
            node.sent_query(transaction_id);
        }
        transaction_id
    }

    fn health_check(&mut self) -> TransactionId {
        self.routing_table.remove_bad_nodes();
        let transaction_id = self.transaction_ids.generate();
        if let Some(node) = self.routing_table.questionable_nodes().get_mut(0) {
            let query = outgoing::create_ping_query(
                transaction_id, &self.self_node);
            node.send(query);
            node.sent_query(transaction_id);
        }
        transaction_id
    }

    fn refresh_bucket(&mut self) -> Option<TransactionId> {
        if let Some(address) = self.address_to_find_for_refresh() {
            Some(self.find_node(&address))
        } else {
            None
        }
    }

    fn address_to_find_for_refresh(&self) -> Option<Address> {
        match self.routing_table.bucket_needing_refresh() {
            Some(bucket) => {
                Some(bucket.random_address_in_space())
            }
            None => None
        }
    }

    fn insert_node(&mut self, node: Box<node::Node>) -> InsertionResult {
        self.routing_table.insert(node, &self.self_node, &mut self.transaction_ids)
    }

    fn send_packet(&mut self, recipient: Address, payload: Vec<u8>, _event_loop: &mut mio::EventLoop<Handler>) {
        for node in self.routing_table.nearest_to(&recipient, false) {
            let transaction_id = self.transaction_ids.generate();
            let query = outgoing::create_packet_query(
                transaction_id, &self.self_node, payload.clone());
            node.send(query.clone());
            node.sent_query(transaction_id);
        }
    }

    fn shutdown(&mut self, event_loop: &mut mio::EventLoop<Handler>) {
        let sentinels: Vec<mpsc::Receiver<()>> = self.shutdown_signals.iter().map(|signal| {
            let (sentinal_sender, sentinal_receiver) = mpsc::channel();
            signal.send(sentinal_sender).expect("couldn't send shutdown signal");
            sentinal_receiver
        }).collect();
        for sentinel in sentinels {
            sentinel.recv().unwrap();
        }
        event_loop.shutdown();
        for listener in &self.event_listeners {
            listener.send(Event::Shutdown).unwrap();
        }
    }
}

struct Handler {
    network: Network
}

impl Handler {
    fn new(network: Network) -> Handler {
        Handler {
            network: network
        }
    }
}

impl mio::Handler for Handler {
    type Timeout = ScheduledTask;
    type Message = OneshotTask;

    fn notify(&mut self, event_loop: &mut mio::EventLoop<Handler>, task: OneshotTask) {
        match task {
            OneshotTask::Incoming(data) => self.network.handle_incoming(data, event_loop),
            OneshotTask::StartBootstrap => self.network.start_bootstrap(event_loop),
            OneshotTask::SendPacket(recipient, payload) =>
                self.network.send_packet(recipient, payload, event_loop),
            OneshotTask::Shutdown => self.network.shutdown(event_loop)
        }
    }

    fn timeout(&mut self, event_loop: &mut mio::EventLoop<Handler>, timeout: ScheduledTask) {
        match timeout {
            ScheduledTask::ContinueBootstrap => self.network.continue_bootstrap(event_loop),
            ScheduledTask::ContinueHealthCheck => self.network.continue_health_check(event_loop),
            ScheduledTask::ContinueRefresh => self.network.continue_refresh(event_loop)
        }
    }
}

// TODO: I'd rather have some kind of select on two channels instead of this nested non-blocking
// recv. It only works because of the 10ms timeout on the socket.
fn create_incoming_udp_channel(host: SocketAddr, sender: TaskSender) -> mpsc::Sender<mpsc::Sender<()>> {
    use std::time::Duration;

    let (shutdown_sender, shutdown_receiver) = mpsc::channel::<mpsc::Sender<()>>();

    thread::spawn(move || {
        let host = ("0.0.0.0", host.port());
        let socket = UdpSocket::bind(host).unwrap();
        socket.set_read_timeout(Some(Duration::from_millis(10))).unwrap();

        loop {
            match shutdown_receiver.try_recv() {
                Err(mpsc::TryRecvError::Empty) => {
                    let mut buf = [0; 4096];
                    match socket.recv(&mut buf) {
                        Ok(size) => {
                            sender.send(OneshotTask::Incoming(buf[..size].iter().cloned().collect()))
                                .unwrap_or_else(|err| info!("Couldn't handling incoming: {:?}", err));
                        }
                        Err(e) => warn!("Error receiving from server: {}", e)
                    }
                }

                Err(mpsc::TryRecvError::Disconnected) => {
                    // The other end of the shutdown_receiver was dropped
                    break
                }

                Ok(sentinel) => {
                    sentinel.send(()).unwrap();
                    break;
                }
            }
        }

    });
    shutdown_sender
}
<commit_msg>Comment out log spam<commit_after>use address::Address;
use messages::outgoing;
use mio;
use node;
use routing_table::{InsertOutcome, InsertionResult, RoutingTable};
use std::collections::HashMap;
use std::io::Cursor;
use std::net::{UdpSocket, ToSocketAddrs, SocketAddr};
use std::sync::mpsc;
use std::thread;
use stun;
use transaction::{TransactionId, TransactionIdGenerator};

#[derive(Debug)]
pub enum Event {
    ReceivedPacket(Address, Vec<u8>),
    Shutdown
}

pub enum OneshotTask {
    Incoming(Vec<u8>),
    StartBootstrap,
    SendPacket(Address, Vec<u8>),
    Shutdown
}

pub enum ScheduledTask {
    ContinueBootstrap,
    ContinueHealthCheck,
    ContinueRefresh
}

enum TableAction {
    Bootstrap(mio::Timeout),
    HealthCheck(mio::Timeout),
    RefreshBucket(mio::Timeout)
}

pub type TaskSender = mio::Sender<OneshotTask>;

enum Status {
    Bootstrapping,
    Idle
}

pub struct Network {
    host: SocketAddr,
    routing_table: RoutingTable,
    self_node: Box<node::Node + 'static>,
    transaction_ids: TransactionIdGenerator,
    status: Status,
    pending_actions: HashMap<TransactionId, TableAction>,
    event_listeners: Vec<mpsc::Sender<Event>>,
    shutdown_signals: Vec<mpsc::Sender<mpsc::Sender<()>>>
}

impl Network {
    pub fn new<T: ToSocketAddrs>(self_address: Address, host: T, routers: Vec<Box<node::Node>>) -> Network {
        let host = host.to_socket_addrs().unwrap().next().unwrap();
        let mapped_host = stun::get_mapped_address(host).unwrap();
        let self_node = node::UdpNode::new(self_address, mapped_host);
        let routing_table = RoutingTable::new(8, self_address, routers);

        Network {
            host: host,
            routing_table: routing_table,
            self_node: Box::new(self_node),
            transaction_ids: TransactionIdGenerator::new(),
            status: Status::Idle,
            pending_actions: HashMap::new(),
            event_listeners: vec![],
            shutdown_signals: vec![]
        }
    }

    pub fn run(mut self) -> TaskSender {
        let mut event_loop_config = mio::EventLoopConfig::new();
        event_loop_config.notify_capacity(16384);
        let mut event_loop = mio::EventLoop::configured(event_loop_config).unwrap();

        let shutdown_signal = create_incoming_udp_channel(self.host, event_loop.channel());
        self.shutdown_signals.push(shutdown_signal);

        event_loop.channel().send(OneshotTask::StartBootstrap).unwrap();
        info!("Running server at {:?}", self.self_node);
        let mut handler = Handler::new(self);
        let task_sender = event_loop.channel();
        thread::spawn(move|| event_loop.run(&mut handler).unwrap());
        task_sender
    }

    pub fn register_event_listener(&mut self, event_listener: mpsc::Sender<Event>) {
        self.event_listeners.push(event_listener);
    }

    fn handle_incoming(&mut self, data: Vec<u8>, event_loop: &mut mio::EventLoop<Handler>) {
        use messages::incoming::{Message, Query, Response, self};
        let mut data = Cursor::new(data);
        let message = incoming::parse_from_reader(&mut data).unwrap();

        match message {
            Message::Query(transaction_id, origin, query) => {
                let origin_address = origin.address();
                match query {
                    Query::FindNode(target) => {
                        let response: Vec<u8> = outgoing::create_find_node_response(
                                transaction_id,
                                &self.self_node,
                                self.routing_table.nearest_to(&target, false));
                        origin.send(response);

                        self.insert_node(origin).unwrap();
                        {
                            if let Some(origin) = self.routing_table.find_node(&origin_address) {
                                origin.received_query(transaction_id);
                            }
                        }
                    },
                    Query::Packet(payload) => {
                        for listener in &self.event_listeners {
                            listener.send(Event::ReceivedPacket(origin_address, payload.clone())).unwrap();
                        }
                        let response = outgoing::create_packet_response(
                            transaction_id, &self.self_node);
                        origin.send(response);
                    },
                    Query::Ping => {
                        let response = outgoing::create_ping_response(
                            transaction_id,
                            &self.self_node);
                        origin.send(response);

                        self.insert_node(origin).unwrap();
                        if let Some(origin) = self.routing_table.find_node(&origin_address) {
                            origin.received_query(transaction_id);
                        }
                    }
                }
            }
            Message::Response(transaction_id, origin, response) => {
                let origin_address = origin.address();
                match response {
                    Response::FindNode(mut nodes) => {
                        let mut encounted_new_node = false;
                        let mut tail = vec![origin];
                        for node in nodes.drain(..).chain(tail.drain(..)) {
                            match self.insert_node(node) {
                                Ok(InsertOutcome::Inserted) => {
                                    encounted_new_node = true;
                                }
                                Err(error) => { panic!(error) }
                                _ => { }
                            }
                        }

                        {
                            if let Some(mut origin) = self.routing_table.find_node(&origin_address) {
                                origin.received_response(transaction_id);
                            }
                        }

                        match self.pending_actions.remove(&transaction_id) {
                            Some(TableAction::Bootstrap(timeout)) => {
                                event_loop.clear_timeout(timeout);
                                if encounted_new_node {
                                    self.continue_bootstrap(event_loop);
                                } else {
                                    self.continue_health_check(event_loop);
                                    self.continue_refresh(event_loop)
                                }
                            }
                            Some(TableAction::HealthCheck(_)) => {
                            },
                            Some(TableAction::RefreshBucket(_)) => {
                            }
                            None => { }
                        }

                        match self.status {
                            Status::Bootstrapping => {
                                self.status = Status::Idle;
                            }
                            _ => { }
                        }
                    }
                    Response::Packet => {
                        if let Some(mut origin) = self.routing_table.find_node(&origin_address) {
                            origin.received_response(transaction_id);
                        }
                    }
                    Response::Ping => {
                        if let Some(mut origin) = self.routing_table.find_node(&origin_address) {
                            origin.received_response(transaction_id);
                        }
                    }
                }
            }
        }
    }

    fn start_bootstrap(&mut self, event_loop: &mut mio::EventLoop<Handler>) {
        self.status = Status::Bootstrapping;
        self.continue_bootstrap(event_loop);
    }

    fn continue_bootstrap(&mut self, event_loop: &mut mio::EventLoop<Handler>) {
        let address = &self.self_node.address();
        let transaction_id = self.find_node(address);
        let timeout = event_loop.timeout_ms(ScheduledTask::ContinueBootstrap, 1000).unwrap();
        self.pending_actions.insert(transaction_id, TableAction::Bootstrap(timeout));
    }

    fn continue_health_check(&mut self, event_loop: &mut mio::EventLoop<Handler>) {
        let transaction_id = self.health_check();
        let timeout = event_loop.timeout_ms(ScheduledTask::ContinueHealthCheck, 1000).unwrap();
        self.pending_actions.insert(transaction_id, TableAction::HealthCheck(timeout));
    }

    fn continue_refresh(&mut self, event_loop: &mut mio::EventLoop<Handler>) {
        let timeout = event_loop.timeout_ms(ScheduledTask::ContinueRefresh, 1000).unwrap();

        match self.refresh_bucket() {
            Some(transaction_id) => {
                self.pending_actions.insert(transaction_id, TableAction::RefreshBucket(timeout));
            }
            None => { }
        }
    }

    fn find_node(&mut self, address: &Address) -> TransactionId {
        let transaction_id = self.transaction_ids.generate();
        let query = outgoing::create_find_node_query(
            transaction_id,
            &self.self_node,
            address);
        for node in self.routing_table.nearest() {
            node.send(query.clone());
            node.sent_query(transaction_id);
        }
        transaction_id
    }

    fn health_check(&mut self) -> TransactionId {
        self.routing_table.remove_bad_nodes();
        let transaction_id = self.transaction_ids.generate();
        if let Some(node) = self.routing_table.questionable_nodes().get_mut(0) {
            let query = outgoing::create_ping_query(
                transaction_id, &self.self_node);
            node.send(query);
            node.sent_query(transaction_id);
        }
        transaction_id
    }

    fn refresh_bucket(&mut self) -> Option<TransactionId> {
        if let Some(address) = self.address_to_find_for_refresh() {
            Some(self.find_node(&address))
        } else {
            None
        }
    }

    fn address_to_find_for_refresh(&self) -> Option<Address> {
        match self.routing_table.bucket_needing_refresh() {
            Some(bucket) => {
                Some(bucket.random_address_in_space())
            }
            None => None
        }
    }

    fn insert_node(&mut self, node: Box<node::Node>) -> InsertionResult {
        self.routing_table.insert(node, &self.self_node, &mut self.transaction_ids)
    }

    fn send_packet(&mut self, recipient: Address, payload: Vec<u8>, _event_loop: &mut mio::EventLoop<Handler>) {
        for node in self.routing_table.nearest_to(&recipient, false) {
            let transaction_id = self.transaction_ids.generate();
            let query = outgoing::create_packet_query(
                transaction_id, &self.self_node, payload.clone());
            node.send(query.clone());
            node.sent_query(transaction_id);
        }
    }

    fn shutdown(&mut self, event_loop: &mut mio::EventLoop<Handler>) {
        let sentinels: Vec<mpsc::Receiver<()>> = self.shutdown_signals.iter().map(|signal| {
            let (sentinal_sender, sentinal_receiver) = mpsc::channel();
            signal.send(sentinal_sender).expect("couldn't send shutdown signal");
            sentinal_receiver
        }).collect();
        for sentinel in sentinels {
            sentinel.recv().unwrap();
        }
        event_loop.shutdown();
        for listener in &self.event_listeners {
            listener.send(Event::Shutdown).unwrap();
        }
    }
}

struct Handler {
    network: Network
}

impl Handler {
    fn new(network: Network) -> Handler {
        Handler {
            network: network
        }
    }
}

impl mio::Handler for Handler {
    type Timeout = ScheduledTask;
    type Message = OneshotTask;

    fn notify(&mut self, event_loop: &mut mio::EventLoop<Handler>, task: OneshotTask) {
        match task {
            OneshotTask::Incoming(data) => self.network.handle_incoming(data, event_loop),
            OneshotTask::StartBootstrap => self.network.start_bootstrap(event_loop),
            OneshotTask::SendPacket(recipient, payload) =>
                self.network.send_packet(recipient, payload, event_loop),
            OneshotTask::Shutdown => self.network.shutdown(event_loop)
        }
    }

    fn timeout(&mut self, event_loop: &mut mio::EventLoop<Handler>, timeout: ScheduledTask) {
        match timeout {
            ScheduledTask::ContinueBootstrap => self.network.continue_bootstrap(event_loop),
            ScheduledTask::ContinueHealthCheck => self.network.continue_health_check(event_loop),
            ScheduledTask::ContinueRefresh => self.network.continue_refresh(event_loop)
        }
    }
}

// TODO: I'd rather have some kind of select on two channels instead of this nested non-blocking
// recv. It only works because of the 10ms timeout on the socket.
fn create_incoming_udp_channel(host: SocketAddr, sender: TaskSender) -> mpsc::Sender<mpsc::Sender<()>> {
    use std::time::Duration;

    let (shutdown_sender, shutdown_receiver) = mpsc::channel::<mpsc::Sender<()>>();

    thread::spawn(move || {
        let host = ("0.0.0.0", host.port());
        let socket = UdpSocket::bind(host).unwrap();
        socket.set_read_timeout(Some(Duration::from_millis(10))).unwrap();

        loop {
            match shutdown_receiver.try_recv() {
                Err(mpsc::TryRecvError::Empty) => {
                    let mut buf = [0; 4096];
                    match socket.recv(&mut buf) {
                        Ok(size) => {
                            sender.send(OneshotTask::Incoming(buf[..size].iter().cloned().collect()))
                                .unwrap_or_else(|err| info!("Couldn't handling incoming: {:?}", err));
                        }
                        Err(e) => {
                            //warn!("Error receiving from server: {}", e)
                        }
                    }
                }

                Err(mpsc::TryRecvError::Disconnected) => {
                    // The other end of the shutdown_receiver was dropped
                    break
                }

                Ok(sentinel) => {
                    sentinel.send(()).unwrap();
                    break;
                }
            }
        }

    });
    shutdown_sender
}
<|endoftext|>