<commit_before>// Copyright 2015 Nathan Sizemore <nathanrsizemore@gmail.com>
//
// 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/.
//
// This Source Code Form is "Incompatible
// With Secondary Licenses", as defined by
// the Mozilla Public License, v. 2.0.

use std::thread;
use std::thread::JoinHandle;
use std::net::TcpStream;
use std::ops::{DerefMut, Deref};
use std::os::unix::io::RawFd;
use std::collections::LinkedList;
use std::process::Command;
use std::sync::{Arc, Mutex};
use std::sync::mpsc::{
    channel,
    Sender,
    Receiver,
    TryRecvError
};

use super::types::*;
use super::epoll;
use super::epoll::util::*;
use super::epoll::EpollEvent;
use super::num_cpus;
use super::socket::Socket;
use super::simple_stream::SimpleStream;

pub struct EventLoop {
    /// Handle to event loop process
    prox: JoinHandle<()>,
    /// Sender<T> to proc
    prox_tx: Sender<TcpStream>
}

impl EventLoop {

    /// Returns a new EventLoop
    pub fn new(server_tx: Sender<EventTuple>) -> EventLoop {
        let (tx, rx): (Sender<TcpStream>, Receiver<TcpStream>) = channel();
        let prox = thread::Builder::new()
            .name("EventLoop".to_string())
            .spawn(move || {
                EventLoop::start(rx, server_tx);
            }).unwrap();

        EventLoop {
            prox: prox,
            prox_tx: tx
        }
    }

    /// Returns a clone of this EventLoop's Sender<TcpStream> channel
    pub fn sender(&self) -> Sender<TcpStream> {
        self.prox_tx.clone()
    }

    /// Main event loop
    fn start(rx: Receiver<TcpStream>, uspace_tx: Sender<EventTuple>) {
        // Master socket list
        let mut sockets = Arc::new(Mutex::new(LinkedList::<Socket>::new()));

        // Master id list
        // TODO - Figure out if this is needed to ensure collisions aren't made?
        let mut ids = Arc::new(Mutex::new(LinkedList::<u32>::new()));

        // Our error handler for this thread and it's children
        let (err_tx, err_rx): (Sender<()>, Receiver<()>) = channel();

        // Create an epoll instance
        let result = epoll::create1(0);
        if result.is_err() {
            let err = result.unwrap_err();
            panic!("Unable to create epoll instance: {}", err)
        }
        let epoll_instance = result.unwrap();

        // Start the epoll thread
        let c_sockets = sockets.clone();
        let c_epoll_instance = epoll_instance.clone();
        let epoll_prox = thread::Builder::new()
            .name("EpollLoop".to_string())
            .spawn(move || {
                EventLoop::epoll_loop(uspace_tx, c_sockets, c_epoll_instance, err_rx);
            }).unwrap();

        for new_stream in rx.iter() {
            match SimpleStream::new(new_stream) {
                Ok(s_stream) => {
                    // Add to master list
                    let mut socket = Socket::new(s_stream);

                    // TODO - Replace with recoverable version once into_inner() is stable
                    // Unfortunately, this is the only stable way to use mutexes at the moment
                    // Hopefully recovering from poisoning will be in 1.2
                    // let mut s_guard = match sockets.lock() {
                    //     Ok(guard) => guard,
                    //     Err(poisoned) => poisoned.into_inner()
                    // };
                    let mut s_guard = sockets.lock().unwrap();
                    let s_list = s_guard.deref_mut();
                    let s_fd = socket.raw_fd();
                    s_list.push_back(socket);

                    // Add to epoll
                    let event = Box::new(EpollEvent {
                        data: s_fd as u64,
                        events: (event_type::EPOLLIN | event_type::EPOLLET)
                        });
                    match epoll::ctl(epoll_instance, ctl_op::ADD, s_fd, event) {
                        Ok(()) => {},
                        Err(e) => println!("Epoll CtrlError: {}", e)
                    };
                }
                Err(e) => {
                    // TODO - write error to log file
                    println!("Error: {}", e);
                }
            };
        }

        // If we get here, shit is real bad. It means we have lost our
        // channel to the outside world.
        // Kill off event loop, so unwinding can begin
        err_tx.send(());
    }

    ///
    fn epoll_loop(uspace_tx: Sender<EventTuple>,
                  sockets: SocketList,
                  epoll_instance: RawFd,
                  err_rx: Receiver<()>) {

        // This is the maximum number of events we want to be notified of at one time
        let mut events = Vec::<EpollEvent>::with_capacity(100);

        loop {
            // Wait for epoll events
            match epoll::wait(epoll_instance, &mut events[..], -1) {
                Ok(num_events) => {
                    for x in 0..num_events {
                        let tx_clone = uspace_tx.clone();
                        let s_clone = sockets.clone();
                        EventLoop::handle_epoll_event(tx_clone, &events[x as usize], s_clone);
                    }
                }
                Err(e) => {
                    panic!("Error on epoll::wait(): {}", e)
                }
            }

            // If we receive anything on this end or lose communication, we need to halt,
            // because we have no where to send up events, our "user space" channel is
            // gone
            match err_rx.try_recv() {
                Ok(_) => {
                    println!("Error, terminating epoll_loop");
                    break;
                }
                Err(e) => {
                    match e {
                        TryRecvError::Empty => {}
                        _ => {
                            println!("Error, terminating epoll_loo");
                            break;
                        }
                    }
                }
            }
        }
    }

    /// Processes a read on the socket that is ready for reading
    ///
    /// An epoll event is passed to this function, with the fd that is ready to be read
    /// A traversal through the list of sockets is then performed in order to locate the
    /// SimpleStream instance that is wrapping that file descriptor.
    /// After a read is performed on that socket, it's data will be passed back to the
    /// "user-space" section of the Server to process the payload however it needs to.
    ///
    /// TODO - Create a faster way to gain a safe reference to the SimpleStream wrapping
    /// the fd instead of an O(n) lookup for every event. Currently lookup time is
    /// O(n) * MAX_EPOLL_EVENTS, where MAX_EPOLL_EVENTS is user defined in this function's
    /// caller.
    fn handle_epoll_event(uspace_tx: Sender<EventTuple>,
                          event: &EpollEvent,
                          sockets: SocketList) {

        // The fd we are traversing the list for
        let fd = event.data;

        // We need a handle before the list gets locked and used
        let s_list_clone = sockets.clone();

        // TODO - Replace with recoverable version once into_inner() is stable
        // Unfortunately, this is the only stable way to use mutexes at the moment
        // Hopefully recovering from poisoning will be in 1.2
        // let mut s_guard = match sockets.lock() {
        //     Ok(guard) => guard,
        //     Err(poisoned) => poisoned.into_inner()
        // };
        let mut s_guard = sockets.lock().unwrap();
        let s_list = s_guard.deref_mut();
        for socket in s_list.iter_mut() {
            if socket.raw_fd() as u64 == fd {
                // We've found the socket for the event passed from epoll
                // Lets attempt to read from it
                match socket.read() {
                    Ok(()) => {
                        for msg in socket.buffer().iter() {
                            let list_handle = s_list_clone.clone();
                            let socket_clone = socket.clone();
                            let msg_clone = msg.clone();
                            uspace_tx.send((list_handle, socket_clone, msg_clone));
                        }
                    }
                    Err(e) => {
                        // TODO - Figure out how to handle
                    }
                }
            }
        }
    }
}
<commit_msg>Debugging statements<commit_after>// Copyright 2015 Nathan Sizemore <nathanrsizemore@gmail.com>
//
// 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/.
//
// This Source Code Form is "Incompatible
// With Secondary Licenses", as defined by
// the Mozilla Public License, v. 2.0.

use std::thread;
use std::thread::JoinHandle;
use std::net::TcpStream;
use std::ops::{DerefMut, Deref};
use std::os::unix::io::RawFd;
use std::collections::LinkedList;
use std::process::Command;
use std::sync::{Arc, Mutex};
use std::sync::mpsc::{
    channel,
    Sender,
    Receiver,
    TryRecvError
};

use super::types::*;
use super::epoll;
use super::epoll::util::*;
use super::epoll::EpollEvent;
use super::num_cpus;
use super::socket::Socket;
use super::simple_stream::SimpleStream;

pub struct EventLoop {
    /// Handle to event loop process
    prox: JoinHandle<()>,
    /// Sender<T> to proc
    prox_tx: Sender<TcpStream>
}

impl EventLoop {

    /// Returns a new EventLoop
    pub fn new(server_tx: Sender<EventTuple>) -> EventLoop {
        let (tx, rx): (Sender<TcpStream>, Receiver<TcpStream>) = channel();
        let prox = thread::Builder::new()
            .name("EventLoop".to_string())
            .spawn(move || {
                EventLoop::start(rx, server_tx);
            }).unwrap();

        EventLoop {
            prox: prox,
            prox_tx: tx
        }
    }

    /// Returns a clone of this EventLoop's Sender<TcpStream> channel
    pub fn sender(&self) -> Sender<TcpStream> {
        self.prox_tx.clone()
    }

    /// Main event loop
    fn start(rx: Receiver<TcpStream>, uspace_tx: Sender<EventTuple>) {
        // Master socket list
        let mut sockets = Arc::new(Mutex::new(LinkedList::<Socket>::new()));

        // Master id list
        // TODO - Figure out if this is needed to ensure collisions aren't made?
        let mut ids = Arc::new(Mutex::new(LinkedList::<u32>::new()));

        // Our error handler for this thread and it's children
        let (err_tx, err_rx): (Sender<()>, Receiver<()>) = channel();

        // Create an epoll instance
        let result = epoll::create1(0);
        if result.is_err() {
            let err = result.unwrap_err();
            panic!("Unable to create epoll instance: {}", err)
        }
        let epoll_instance = result.unwrap();

        println!("efd: {}", epoll_instance);

        // Start the epoll thread
        let c_sockets = sockets.clone();
        let c_epoll_instance = epoll_instance.clone();
        let epoll_prox = thread::Builder::new()
            .name("EpollLoop".to_string())
            .spawn(move || {
                EventLoop::epoll_loop(uspace_tx, c_sockets, c_epoll_instance, err_rx);
            }).unwrap();

        for new_stream in rx.iter() {
            match SimpleStream::new(new_stream) {
                Ok(s_stream) => {
                    // Add to master list
                    let mut socket = Socket::new(s_stream);

                    // TODO - Replace with recoverable version once into_inner() is stable
                    // Unfortunately, this is the only stable way to use mutexes at the moment
                    // Hopefully recovering from poisoning will be in 1.2
                    // let mut s_guard = match sockets.lock() {
                    //     Ok(guard) => guard,
                    //     Err(poisoned) => poisoned.into_inner()
                    // };
                    let mut s_guard = sockets.lock().unwrap();
                    let s_list = s_guard.deref_mut();
                    let s_fd = socket.raw_fd();
                    s_list.push_back(socket);

                    // Add to epoll
                    let event = Box::new(EpollEvent {
                        data: s_fd as u64,
                        events: (event_type::EPOLLIN | event_type::EPOLLET)
                        });
                    match epoll::ctl(epoll_instance, ctl_op::ADD, s_fd, event) {
                        Ok(()) => {},
                        Err(e) => println!("Epoll CtrlError: {}", e)
                    };
                }
                Err(e) => {
                    // TODO - write error to log file
                    println!("Error: {}", e);
                }
            };
        }

        // If we get here, shit is real bad. It means we have lost our
        // channel to the outside world.
        // Kill off event loop, so unwinding can begin
        err_tx.send(());
    }

    ///
    fn epoll_loop(uspace_tx: Sender<EventTuple>,
                  sockets: SocketList,
                  epoll_instance: RawFd,
                  err_rx: Receiver<()>) {

        // This is the maximum number of events we want to be notified of at one time
        let mut events = Vec::<EpollEvent>::with_capacity(100);

        loop {
            // Wait for epoll events
            match epoll::wait(epoll_instance, &mut events[..], -1) {
                Ok(num_events) => {
                    for x in 0..num_events {
                        let tx_clone = uspace_tx.clone();
                        let s_clone = sockets.clone();
                        EventLoop::handle_epoll_event(tx_clone, &events[x as usize], s_clone);
                    }
                }
                Err(e) => {
                    panic!("Error on epoll::wait(): {}", e)
                }
            }

            // If we receive anything on this end or lose communication, we need to halt,
            // because we have no where to send up events, our "user space" channel is
            // gone
            match err_rx.try_recv() {
                Ok(_) => {
                    println!("Error, terminating epoll_loop");
                    break;
                }
                Err(e) => {
                    match e {
                        TryRecvError::Empty => {}
                        _ => {
                            println!("Error, terminating epoll_loo");
                            break;
                        }
                    }
                }
            }
        }
    }

    /// Processes a read on the socket that is ready for reading
    ///
    /// An epoll event is passed to this function, with the fd that is ready to be read
    /// A traversal through the list of sockets is then performed in order to locate the
    /// SimpleStream instance that is wrapping that file descriptor.
    /// After a read is performed on that socket, it's data will be passed back to the
    /// "user-space" section of the Server to process the payload however it needs to.
    ///
    /// TODO - Create a faster way to gain a safe reference to the SimpleStream wrapping
    /// the fd instead of an O(n) lookup for every event. Currently lookup time is
    /// O(n) * MAX_EPOLL_EVENTS, where MAX_EPOLL_EVENTS is user defined in this function's
    /// caller.
    fn handle_epoll_event(uspace_tx: Sender<EventTuple>,
                          event: &EpollEvent,
                          sockets: SocketList) {

        // The fd we are traversing the list for
        let fd = event.data;

        // We need a handle before the list gets locked and used
        let s_list_clone = sockets.clone();

        // TODO - Replace with recoverable version once into_inner() is stable
        // Unfortunately, this is the only stable way to use mutexes at the moment
        // Hopefully recovering from poisoning will be in 1.2
        // let mut s_guard = match sockets.lock() {
        //     Ok(guard) => guard,
        //     Err(poisoned) => poisoned.into_inner()
        // };
        let mut s_guard = sockets.lock().unwrap();
        let s_list = s_guard.deref_mut();
        for socket in s_list.iter_mut() {
            if socket.raw_fd() as u64 == fd {
                // We've found the socket for the event passed from epoll
                // Lets attempt to read from it
                match socket.read() {
                    Ok(()) => {
                        for msg in socket.buffer().iter() {
                            let list_handle = s_list_clone.clone();
                            let socket_clone = socket.clone();
                            let msg_clone = msg.clone();
                            uspace_tx.send((list_handle, socket_clone, msg_clone));
                        }
                    }
                    Err(e) => {
                        // TODO - Figure out how to handle
                    }
                }
            }
        }
    }
}
<|endoftext|>


<commit_before>use scroll::{self, Pread};
use error::{self, Error};

#[repr(C)]
#[derive(Debug, PartialEq, Clone, Default)]
pub struct SectionTable {
    pub name: [u8; 8],
    pub real_name: Option<String>,
    pub virtual_size: u32,
    pub virtual_address: u32,
    pub size_of_raw_data: u32,
    pub pointer_to_raw_data: u32,
    pub pointer_to_relocations: u32,
    pub pointer_to_linenumbers: u32,
    pub number_of_relocations: u16,
    pub number_of_linenumbers: u16,
    pub characteristics: u32,
}

pub const SIZEOF_SECTION_TABLE: usize = 8 * 5;

// Based on https://github.com/llvm-mirror/llvm/blob/af7b1832a03ab6486c42a40d21695b2c03b2d8a3/lib/Object/COFFObjectFile.cpp#L70
// Decodes a string table entry in base 64 (//AAAAAA). Expects string without
// prefixed slashes.
fn base64_decode_string_entry(s: &str) -> Result<usize, ()> {
    assert!(s.len() <= 6, "String too long, possible overflow.");

    let mut val = 0;
    for c in s.bytes() {
        let v = match c {
            b'A'..=b'Z' => c - b'A' + 00, // 00..=25
            b'a'..=b'z' => c - b'a' + 26, // 26..=51
            b'0'..=b'9' => c - b'0' + 52, // 52..=61
            b'+'        => 62,            // 62
            b'/'        => 63,            // 63
            _           => return Err(())
        };
        val = val * 64 + v as usize;
    }
    Ok(val)
}

impl SectionTable {
    pub fn parse(bytes: &[u8], offset: &mut usize, string_table_offset: usize) -> error::Result<Self> {
        let mut table = SectionTable::default();
        let mut name = [0u8; 8];
        for i in 0..8 {
            name[i] = bytes.gread_with(offset, scroll::LE)?;
        }

        table.name = name;
        table.virtual_size = bytes.gread_with(offset, scroll::LE)?;
        table.virtual_address = bytes.gread_with(offset, scroll::LE)?;
        table.size_of_raw_data = bytes.gread_with(offset, scroll::LE)?;
        table.pointer_to_raw_data = bytes.gread_with(offset, scroll::LE)?;
        table.pointer_to_relocations = bytes.gread_with(offset, scroll::LE)?;
        table.pointer_to_linenumbers = bytes.gread_with(offset, scroll::LE)?;
        table.number_of_relocations = bytes.gread_with(offset, scroll::LE)?;
        table.number_of_linenumbers = bytes.gread_with(offset, scroll::LE)?;
        table.characteristics = bytes.gread_with(offset, scroll::LE)?;

        // Based on https://github.com/llvm-mirror/llvm/blob/af7b1832a03ab6486c42a40d21695b2c03b2d8a3/lib/Object/COFFObjectFile.cpp#L1054
        if name[0] == b'/' {
            let idx: usize = if name[1] == b'/' {
                let b64idx = name.pread::<&str>(2)?;
                base64_decode_string_entry(b64idx).map_err(|_|
                    Error::Malformed(format!("Invalid indirect section name //{}: base64 decoding failed", b64idx)))?
            } else {
                let name = name.pread::<&str>(1)?;
                name.parse().map_err(|err|
                    Error::Malformed(format!("Invalid indirect section name /{}: {}", name, err)))?
            };
            table.real_name = Some(bytes.pread::<&str>(string_table_offset + idx)?.to_string());
        }
        Ok(table)
    }

    pub fn name(&self) -> error::Result<&str> {
        match self.real_name.as_ref() {
            Some(s) => Ok(s),
            None => Ok(self.name.pread(0)?)
        }
    }
}

/// The section should not be padded to the next boundary. This flag is obsolete and is replaced
/// by `IMAGE_SCN_ALIGN_1BYTES`. This is valid only for object files.
pub const IMAGE_SCN_TYPE_NO_PAD: u32 = 0x00000008;
/// The section contains executable code.
pub const IMAGE_SCN_CNT_CODE: u32 = 0x00000020;
/// The section contains initialized data.
pub const IMAGE_SCN_CNT_INITIALIZED_DATA: u32 = 0x00000040;
///  The section contains uninitialized data.
pub const IMAGE_SCN_CNT_UNINITIALIZED_DATA: u32 = 0x00000080;
pub const IMAGE_SCN_LNK_OTHER: u32 = 0x00000100;
/// The section contains comments or other information. The .drectve section has this type.
/// This is valid for object files only.
pub const IMAGE_SCN_LNK_INFO: u32 = 0x00000200;
/// The section will not become part of the image. This is valid only for object files.
pub const IMAGE_SCN_LNK_REMOVE: u32 = 0x00000800;
/// The section contains COMDAT data. This is valid only for object files.
pub const IMAGE_SCN_LNK_COMDAT: u32 = 0x00001000;
/// The section contains data referenced through the global pointer (GP).
pub const IMAGE_SCN_GPREL: u32 = 0x00008000;
pub const IMAGE_SCN_MEM_PURGEABLE: u32 = 0x00020000;
pub const IMAGE_SCN_MEM_16BIT: u32 = 0x00020000;
pub const IMAGE_SCN_MEM_LOCKED: u32 = 0x00040000;
pub const IMAGE_SCN_MEM_PRELOAD: u32 = 0x00080000;

pub const IMAGE_SCN_ALIGN_1BYTES: u32 = 0x00100000;
pub const IMAGE_SCN_ALIGN_2BYTES: u32 = 0x00200000;
pub const IMAGE_SCN_ALIGN_4BYTES: u32 = 0x00300000;
pub const IMAGE_SCN_ALIGN_8BYTES: u32 = 0x00400000;
pub const IMAGE_SCN_ALIGN_16BYTES: u32 = 0x00500000;
pub const IMAGE_SCN_ALIGN_32BYTES: u32 = 0x00600000;
pub const IMAGE_SCN_ALIGN_64BYTES: u32 = 0x00700000;
pub const IMAGE_SCN_ALIGN_128BYTES: u32 = 0x00800000;
pub const IMAGE_SCN_ALIGN_256BYTES: u32 = 0x00900000;
pub const IMAGE_SCN_ALIGN_512BYTES: u32 = 0x00A00000;
pub const IMAGE_SCN_ALIGN_1024BYTES: u32 = 0x00B00000;
pub const IMAGE_SCN_ALIGN_2048BYTES: u32 = 0x00C00000;
pub const IMAGE_SCN_ALIGN_4096BYTES: u32 = 0x00D00000;
pub const IMAGE_SCN_ALIGN_8192BYTES: u32 = 0x00E00000;
pub const IMAGE_SCN_ALIGN_MASK: u32 = 0x00F00000;

/// The section contains extended relocations.
pub const IMAGE_SCN_LNK_NRELOC_OVFL: u32 = 0x01000000;
/// The section can be discarded as needed.
pub const IMAGE_SCN_MEM_DISCARDABLE: u32 = 0x02000000;
/// The section cannot be cached.
pub const IMAGE_SCN_MEM_NOT_CACHED: u32 = 0x04000000;
/// The section is not pageable.
pub const IMAGE_SCN_MEM_NOT_PAGED: u32 = 0x08000000;
/// The section can be shared in memory.
pub const IMAGE_SCN_MEM_SHARED: u32 = 0x10000000;
/// The section can be executed as code.
pub const IMAGE_SCN_MEM_EXECUTE: u32 = 0x20000000;
/// The section can be read.
pub const IMAGE_SCN_MEM_READ: u32 = 0x40000000;
/// The section can be written to.
pub const IMAGE_SCN_MEM_WRITE: u32 = 0x80000000;
<commit_msg>avoid using newer rust features<commit_after>use scroll::{self, Pread};
use error::{self, Error};

#[repr(C)]
#[derive(Debug, PartialEq, Clone, Default)]
pub struct SectionTable {
    pub name: [u8; 8],
    pub real_name: Option<String>,
    pub virtual_size: u32,
    pub virtual_address: u32,
    pub size_of_raw_data: u32,
    pub pointer_to_raw_data: u32,
    pub pointer_to_relocations: u32,
    pub pointer_to_linenumbers: u32,
    pub number_of_relocations: u16,
    pub number_of_linenumbers: u16,
    pub characteristics: u32,
}

pub const SIZEOF_SECTION_TABLE: usize = 8 * 5;

// Based on https://github.com/llvm-mirror/llvm/blob/af7b1832a03ab6486c42a40d21695b2c03b2d8a3/lib/Object/COFFObjectFile.cpp#L70
// Decodes a string table entry in base 64 (//AAAAAA). Expects string without
// prefixed slashes.
fn base64_decode_string_entry(s: &str) -> Result<usize, ()> {
    assert!(s.len() <= 6, "String too long, possible overflow.");

    let mut val = 0;
    for c in s.bytes() {
        let v = if b'A' <= c && c <= b'Z' {
            c - b'A' + 00 // 00..=25
        } else if b'a' <= c && c <= b'z' {
            c - b'a' + 26 // 26..=51
        } else if b'0' <= c && c <= b'9' {
            c - b'0' + 52 // 52..=61
        } else if c == b'+' {
            62            // 62
        } else if c == b'/' {
            63            // 63
        } else {
            return Err(())
        };
        val = val * 64 + v as usize;
    }
    Ok(val)
}

impl SectionTable {
    pub fn parse(bytes: &[u8], offset: &mut usize, string_table_offset: usize) -> error::Result<Self> {
        let mut table = SectionTable::default();
        let mut name = [0u8; 8];
        for i in 0..8 {
            name[i] = bytes.gread_with(offset, scroll::LE)?;
        }

        table.name = name;
        table.virtual_size = bytes.gread_with(offset, scroll::LE)?;
        table.virtual_address = bytes.gread_with(offset, scroll::LE)?;
        table.size_of_raw_data = bytes.gread_with(offset, scroll::LE)?;
        table.pointer_to_raw_data = bytes.gread_with(offset, scroll::LE)?;
        table.pointer_to_relocations = bytes.gread_with(offset, scroll::LE)?;
        table.pointer_to_linenumbers = bytes.gread_with(offset, scroll::LE)?;
        table.number_of_relocations = bytes.gread_with(offset, scroll::LE)?;
        table.number_of_linenumbers = bytes.gread_with(offset, scroll::LE)?;
        table.characteristics = bytes.gread_with(offset, scroll::LE)?;

        // Based on https://github.com/llvm-mirror/llvm/blob/af7b1832a03ab6486c42a40d21695b2c03b2d8a3/lib/Object/COFFObjectFile.cpp#L1054
        if name[0] == b'/' {
            let idx: usize = if name[1] == b'/' {
                let b64idx = name.pread::<&str>(2)?;
                base64_decode_string_entry(b64idx).map_err(|_|
                    Error::Malformed(format!("Invalid indirect section name //{}: base64 decoding failed", b64idx)))?
            } else {
                let name = name.pread::<&str>(1)?;
                name.parse().map_err(|err|
                    Error::Malformed(format!("Invalid indirect section name /{}: {}", name, err)))?
            };
            table.real_name = Some(bytes.pread::<&str>(string_table_offset + idx)?.to_string());
        }
        Ok(table)
    }

    pub fn name(&self) -> error::Result<&str> {
        match self.real_name.as_ref() {
            Some(s) => Ok(s),
            None => Ok(self.name.pread(0)?)
        }
    }
}

/// The section should not be padded to the next boundary. This flag is obsolete and is replaced
/// by `IMAGE_SCN_ALIGN_1BYTES`. This is valid only for object files.
pub const IMAGE_SCN_TYPE_NO_PAD: u32 = 0x00000008;
/// The section contains executable code.
pub const IMAGE_SCN_CNT_CODE: u32 = 0x00000020;
/// The section contains initialized data.
pub const IMAGE_SCN_CNT_INITIALIZED_DATA: u32 = 0x00000040;
///  The section contains uninitialized data.
pub const IMAGE_SCN_CNT_UNINITIALIZED_DATA: u32 = 0x00000080;
pub const IMAGE_SCN_LNK_OTHER: u32 = 0x00000100;
/// The section contains comments or other information. The .drectve section has this type.
/// This is valid for object files only.
pub const IMAGE_SCN_LNK_INFO: u32 = 0x00000200;
/// The section will not become part of the image. This is valid only for object files.
pub const IMAGE_SCN_LNK_REMOVE: u32 = 0x00000800;
/// The section contains COMDAT data. This is valid only for object files.
pub const IMAGE_SCN_LNK_COMDAT: u32 = 0x00001000;
/// The section contains data referenced through the global pointer (GP).
pub const IMAGE_SCN_GPREL: u32 = 0x00008000;
pub const IMAGE_SCN_MEM_PURGEABLE: u32 = 0x00020000;
pub const IMAGE_SCN_MEM_16BIT: u32 = 0x00020000;
pub const IMAGE_SCN_MEM_LOCKED: u32 = 0x00040000;
pub const IMAGE_SCN_MEM_PRELOAD: u32 = 0x00080000;

pub const IMAGE_SCN_ALIGN_1BYTES: u32 = 0x00100000;
pub const IMAGE_SCN_ALIGN_2BYTES: u32 = 0x00200000;
pub const IMAGE_SCN_ALIGN_4BYTES: u32 = 0x00300000;
pub const IMAGE_SCN_ALIGN_8BYTES: u32 = 0x00400000;
pub const IMAGE_SCN_ALIGN_16BYTES: u32 = 0x00500000;
pub const IMAGE_SCN_ALIGN_32BYTES: u32 = 0x00600000;
pub const IMAGE_SCN_ALIGN_64BYTES: u32 = 0x00700000;
pub const IMAGE_SCN_ALIGN_128BYTES: u32 = 0x00800000;
pub const IMAGE_SCN_ALIGN_256BYTES: u32 = 0x00900000;
pub const IMAGE_SCN_ALIGN_512BYTES: u32 = 0x00A00000;
pub const IMAGE_SCN_ALIGN_1024BYTES: u32 = 0x00B00000;
pub const IMAGE_SCN_ALIGN_2048BYTES: u32 = 0x00C00000;
pub const IMAGE_SCN_ALIGN_4096BYTES: u32 = 0x00D00000;
pub const IMAGE_SCN_ALIGN_8192BYTES: u32 = 0x00E00000;
pub const IMAGE_SCN_ALIGN_MASK: u32 = 0x00F00000;

/// The section contains extended relocations.
pub const IMAGE_SCN_LNK_NRELOC_OVFL: u32 = 0x01000000;
/// The section can be discarded as needed.
pub const IMAGE_SCN_MEM_DISCARDABLE: u32 = 0x02000000;
/// The section cannot be cached.
pub const IMAGE_SCN_MEM_NOT_CACHED: u32 = 0x04000000;
/// The section is not pageable.
pub const IMAGE_SCN_MEM_NOT_PAGED: u32 = 0x08000000;
/// The section can be shared in memory.
pub const IMAGE_SCN_MEM_SHARED: u32 = 0x10000000;
/// The section can be executed as code.
pub const IMAGE_SCN_MEM_EXECUTE: u32 = 0x20000000;
/// The section can be read.
pub const IMAGE_SCN_MEM_READ: u32 = 0x40000000;
/// The section can be written to.
pub const IMAGE_SCN_MEM_WRITE: u32 = 0x80000000;
<|endoftext|>


<commit_before>use crate::fmt;

/// A struct containing information about the location of a panic.
///
/// This structure is created by [`PanicInfo::location()`].
///
/// # Examples
///
/// ```should_panic
/// use std::panic;
///
/// panic::set_hook(Box::new(|panic_info| {
///     if let Some(location) = panic_info.location() {
///         println!("panic occurred in file '{}' at line {}", location.file(), location.line());
///     } else {
///         println!("panic occurred but can't get location information...");
///     }
/// }));
///
/// panic!("Normal panic");
/// ```
///
/// # Comparisons
///
/// Comparisons for equality and ordering are made in file, line, then column priority.
/// Files are compared as strings, not `Path`, which could be unexpected.
/// See [`Location::file`]'s documentation for more discussion.
#[lang = "panic_location"]
#[derive(Copy, Clone, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
#[stable(feature = "panic_hooks", since = "1.10.0")]
pub struct Location<'a> {
    file: &'a str,
    line: u32,
    col: u32,
}

impl<'a> Location<'a> {
    /// Returns the source location of the caller of this function. If that function's caller is
    /// annotated then its call location will be returned, and so on up the stack to the first call
    /// within a non-tracked function body.
    ///
    /// # Examples
    ///
    /// ```
    /// use std::panic::Location;
    ///
    /// /// Returns the [`Location`] at which it is called.
    /// #[track_caller]
    /// fn get_caller_location() -> &'static Location<'static> {
    ///     Location::caller()
    /// }
    ///
    /// /// Returns a [`Location`] from within this function's definition.
    /// fn get_just_one_location() -> &'static Location<'static> {
    ///     get_caller_location()
    /// }
    ///
    /// let fixed_location = get_just_one_location();
    /// assert_eq!(fixed_location.file(), file!());
    /// assert_eq!(fixed_location.line(), 14);
    /// assert_eq!(fixed_location.column(), 5);
    ///
    /// // running the same untracked function in a different location gives us the same result
    /// let second_fixed_location = get_just_one_location();
    /// assert_eq!(fixed_location.file(), second_fixed_location.file());
    /// assert_eq!(fixed_location.line(), second_fixed_location.line());
    /// assert_eq!(fixed_location.column(), second_fixed_location.column());
    ///
    /// let this_location = get_caller_location();
    /// assert_eq!(this_location.file(), file!());
    /// assert_eq!(this_location.line(), 28);
    /// assert_eq!(this_location.column(), 21);
    ///
    /// // running the tracked function in a different location produces a different value
    /// let another_location = get_caller_location();
    /// assert_eq!(this_location.file(), another_location.file());
    /// assert_ne!(this_location.line(), another_location.line());
    /// assert_ne!(this_location.column(), another_location.column());
    /// ```
    #[stable(feature = "track_caller", since = "1.46.0")]
    #[rustc_const_unstable(feature = "const_caller_location", issue = "76156")]
    #[track_caller]
    pub const fn caller() -> &'static Location<'static> {
        crate::intrinsics::caller_location()
    }
}

impl<'a> Location<'a> {
    #![unstable(
        feature = "panic_internals",
        reason = "internal details of the implementation of the `panic!` and related macros",
        issue = "none"
    )]
    #[doc(hidden)]
    pub const fn internal_constructor(file: &'a str, line: u32, col: u32) -> Self {
        Location { file, line, col }
    }

    /// Returns the name of the source file from which the panic originated.
    ///
    /// # `&str`, not `&Path`
    ///
    /// The returned name refers to a source path on the compiling system, but it isn't valid to
    /// represent this directly as a `&Path`. The compiled code may run on a different system with
    /// a different `Path` implementation than the system providing the contents and this library
    /// does not currently have a different "host path" type.
    ///
    /// The most surprising behavior occurs when "the same" file is reachable via multiple paths in
    /// the module system (usually using the `#[path = "..."]` attribute or similar), which can
    /// cause what appears to be identical code to return differing values from this function.
    ///
    /// # Cross-compilation
    ///
    /// This value is not suitable for passing to `Path::new` or similar constructors when the host
    /// platform and target platform differ.
    ///
    /// # Examples
    ///
    /// ```should_panic
    /// use std::panic;
    ///
    /// panic::set_hook(Box::new(|panic_info| {
    ///     if let Some(location) = panic_info.location() {
    ///         println!("panic occurred in file '{}'", location.file());
    ///     } else {
    ///         println!("panic occurred but can't get location information...");
    ///     }
    /// }));
    ///
    /// panic!("Normal panic");
    /// ```
    #[stable(feature = "panic_hooks", since = "1.10.0")]
    pub fn file(&self) -> &str {
        self.file
    }

    /// Returns the line number from which the panic originated.
    ///
    /// # Examples
    ///
    /// ```should_panic
    /// use std::panic;
    ///
    /// panic::set_hook(Box::new(|panic_info| {
    ///     if let Some(location) = panic_info.location() {
    ///         println!("panic occurred at line {}", location.line());
    ///     } else {
    ///         println!("panic occurred but can't get location information...");
    ///     }
    /// }));
    ///
    /// panic!("Normal panic");
    /// ```
    #[stable(feature = "panic_hooks", since = "1.10.0")]
    pub fn line(&self) -> u32 {
        self.line
    }

    /// Returns the column from which the panic originated.
    ///
    /// # Examples
    ///
    /// ```should_panic
    /// use std::panic;
    ///
    /// panic::set_hook(Box::new(|panic_info| {
    ///     if let Some(location) = panic_info.location() {
    ///         println!("panic occurred at column {}", location.column());
    ///     } else {
    ///         println!("panic occurred but can't get location information...");
    ///     }
    /// }));
    ///
    /// panic!("Normal panic");
    /// ```
    #[stable(feature = "panic_col", since = "1.25.0")]
    pub fn column(&self) -> u32 {
        self.col
    }
}

#[stable(feature = "panic_hook_display", since = "1.26.0")]
impl fmt::Display for Location<'_> {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(formatter, "{}:{}:{}", self.file, self.line, self.col)
    }
}
<commit_msg>Fix separation of public vs internal parts of Location<commit_after>use crate::fmt;

/// A struct containing information about the location of a panic.
///
/// This structure is created by [`PanicInfo::location()`].
///
/// # Examples
///
/// ```should_panic
/// use std::panic;
///
/// panic::set_hook(Box::new(|panic_info| {
///     if let Some(location) = panic_info.location() {
///         println!("panic occurred in file '{}' at line {}", location.file(), location.line());
///     } else {
///         println!("panic occurred but can't get location information...");
///     }
/// }));
///
/// panic!("Normal panic");
/// ```
///
/// # Comparisons
///
/// Comparisons for equality and ordering are made in file, line, then column priority.
/// Files are compared as strings, not `Path`, which could be unexpected.
/// See [`Location::file`]'s documentation for more discussion.
#[lang = "panic_location"]
#[derive(Copy, Clone, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
#[stable(feature = "panic_hooks", since = "1.10.0")]
pub struct Location<'a> {
    file: &'a str,
    line: u32,
    col: u32,
}

impl<'a> Location<'a> {
    /// Returns the source location of the caller of this function. If that function's caller is
    /// annotated then its call location will be returned, and so on up the stack to the first call
    /// within a non-tracked function body.
    ///
    /// # Examples
    ///
    /// ```
    /// use std::panic::Location;
    ///
    /// /// Returns the [`Location`] at which it is called.
    /// #[track_caller]
    /// fn get_caller_location() -> &'static Location<'static> {
    ///     Location::caller()
    /// }
    ///
    /// /// Returns a [`Location`] from within this function's definition.
    /// fn get_just_one_location() -> &'static Location<'static> {
    ///     get_caller_location()
    /// }
    ///
    /// let fixed_location = get_just_one_location();
    /// assert_eq!(fixed_location.file(), file!());
    /// assert_eq!(fixed_location.line(), 14);
    /// assert_eq!(fixed_location.column(), 5);
    ///
    /// // running the same untracked function in a different location gives us the same result
    /// let second_fixed_location = get_just_one_location();
    /// assert_eq!(fixed_location.file(), second_fixed_location.file());
    /// assert_eq!(fixed_location.line(), second_fixed_location.line());
    /// assert_eq!(fixed_location.column(), second_fixed_location.column());
    ///
    /// let this_location = get_caller_location();
    /// assert_eq!(this_location.file(), file!());
    /// assert_eq!(this_location.line(), 28);
    /// assert_eq!(this_location.column(), 21);
    ///
    /// // running the tracked function in a different location produces a different value
    /// let another_location = get_caller_location();
    /// assert_eq!(this_location.file(), another_location.file());
    /// assert_ne!(this_location.line(), another_location.line());
    /// assert_ne!(this_location.column(), another_location.column());
    /// ```
    #[stable(feature = "track_caller", since = "1.46.0")]
    #[rustc_const_unstable(feature = "const_caller_location", issue = "76156")]
    #[track_caller]
    pub const fn caller() -> &'static Location<'static> {
        crate::intrinsics::caller_location()
    }

    /// Returns the name of the source file from which the panic originated.
    ///
    /// # `&str`, not `&Path`
    ///
    /// The returned name refers to a source path on the compiling system, but it isn't valid to
    /// represent this directly as a `&Path`. The compiled code may run on a different system with
    /// a different `Path` implementation than the system providing the contents and this library
    /// does not currently have a different "host path" type.
    ///
    /// The most surprising behavior occurs when "the same" file is reachable via multiple paths in
    /// the module system (usually using the `#[path = "..."]` attribute or similar), which can
    /// cause what appears to be identical code to return differing values from this function.
    ///
    /// # Cross-compilation
    ///
    /// This value is not suitable for passing to `Path::new` or similar constructors when the host
    /// platform and target platform differ.
    ///
    /// # Examples
    ///
    /// ```should_panic
    /// use std::panic;
    ///
    /// panic::set_hook(Box::new(|panic_info| {
    ///     if let Some(location) = panic_info.location() {
    ///         println!("panic occurred in file '{}'", location.file());
    ///     } else {
    ///         println!("panic occurred but can't get location information...");
    ///     }
    /// }));
    ///
    /// panic!("Normal panic");
    /// ```
    #[stable(feature = "panic_hooks", since = "1.10.0")]
    pub fn file(&self) -> &str {
        self.file
    }

    /// Returns the line number from which the panic originated.
    ///
    /// # Examples
    ///
    /// ```should_panic
    /// use std::panic;
    ///
    /// panic::set_hook(Box::new(|panic_info| {
    ///     if let Some(location) = panic_info.location() {
    ///         println!("panic occurred at line {}", location.line());
    ///     } else {
    ///         println!("panic occurred but can't get location information...");
    ///     }
    /// }));
    ///
    /// panic!("Normal panic");
    /// ```
    #[stable(feature = "panic_hooks", since = "1.10.0")]
    pub fn line(&self) -> u32 {
        self.line
    }

    /// Returns the column from which the panic originated.
    ///
    /// # Examples
    ///
    /// ```should_panic
    /// use std::panic;
    ///
    /// panic::set_hook(Box::new(|panic_info| {
    ///     if let Some(location) = panic_info.location() {
    ///         println!("panic occurred at column {}", location.column());
    ///     } else {
    ///         println!("panic occurred but can't get location information...");
    ///     }
    /// }));
    ///
    /// panic!("Normal panic");
    /// ```
    #[stable(feature = "panic_col", since = "1.25.0")]
    pub fn column(&self) -> u32 {
        self.col
    }
}

#[unstable(
    feature = "panic_internals",
    reason = "internal details of the implementation of the `panic!` and related macros",
    issue = "none"
)]
impl<'a> Location<'a> {
    #[doc(hidden)]
    pub const fn internal_constructor(file: &'a str, line: u32, col: u32) -> Self {
        Location { file, line, col }
    }
}

#[stable(feature = "panic_hook_display", since = "1.26.0")]
impl fmt::Display for Location<'_> {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(formatter, "{}:{}:{}", self.file, self.line, self.col)
    }
}
<|endoftext|>


<commit_before>use std::io::{self, BufWriter, Read, Write};
use std::net::{ToSocketAddrs, SocketAddr};
use std::time::Duration;
use std::result;
use bytes::{BufMut, BytesMut};
#[cfg(unix)]
use std::os::unix::net::UnixStream;
#[cfg(unix)]
use std::os::unix::io::{AsRawFd, RawFd, FromRawFd, IntoRawFd};
#[cfg(windows)]
use std::os::windows::io::{AsRawSocket, RawSocket};
use postgres_protocol::message::frontend;
use postgres_protocol::message::backend;
use socket2::{Socket, SockAddr, Domain, Type};

use {Result, TlsMode};
use error;
use tls::TlsStream;
use params::{ConnectParams, Host};

const INITIAL_CAPACITY: usize = 8 * 1024;

pub struct MessageStream {
    stream: BufWriter<Box<TlsStream>>,
    in_buf: BytesMut,
    out_buf: Vec<u8>,
}

impl MessageStream {
    pub fn new(stream: Box<TlsStream>) -> MessageStream {
        MessageStream {
            stream: BufWriter::new(stream),
            in_buf: BytesMut::with_capacity(INITIAL_CAPACITY),
            out_buf: vec![],
        }
    }

    pub fn get_ref(&self) -> &Box<TlsStream> {
        self.stream.get_ref()
    }

    pub fn write_message<F, E>(&mut self, f: F) -> result::Result<(), E>
    where
        F: FnOnce(&mut Vec<u8>) -> result::Result<(), E>,
        E: From<io::Error>,
    {
        self.out_buf.clear();
        f(&mut self.out_buf)?;
        self.stream.write_all(&self.out_buf).map_err(From::from)
    }

    pub fn read_message(&mut self) -> io::Result<backend::Message> {
        loop {
            match backend::Message::parse(&mut self.in_buf) {
                Ok(Some(message)) => return Ok(message),
                Ok(None) => self.read_in()?,
                Err(e) => return Err(e),
            }
        }
    }

    fn read_in(&mut self) -> io::Result<()> {
        self.in_buf.reserve(1);
        match self.stream.get_mut().read(
            unsafe { self.in_buf.bytes_mut() },
        ) {
            Ok(0) => Err(io::Error::new(
                io::ErrorKind::UnexpectedEof,
                "unexpected EOF",
            )),
            Ok(n) => {
                unsafe { self.in_buf.advance_mut(n) };
                Ok(())
            }
            Err(e) => Err(e),
        }
    }

    pub fn read_message_timeout(
        &mut self,
        timeout: Duration,
    ) -> io::Result<Option<backend::Message>> {
        if self.in_buf.is_empty() {
            self.set_read_timeout(Some(timeout))?;
            let r = self.read_in();
            self.set_read_timeout(None)?;

            match r {
                Ok(()) => {}
                Err(ref e)
                    if e.kind() == io::ErrorKind::WouldBlock ||
                           e.kind() == io::ErrorKind::TimedOut => return Ok(None),
                Err(e) => return Err(e),
            }
        }

        self.read_message().map(Some)
    }

    pub fn read_message_nonblocking(&mut self) -> io::Result<Option<backend::Message>> {
        if self.in_buf.is_empty() {
            self.set_nonblocking(true)?;
            let r = self.read_in();
            self.set_nonblocking(false)?;

            match r {
                Ok(()) => {}
                Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => return Ok(None),
                Err(e) => return Err(e),
            }
        }

        self.read_message().map(Some)
    }

    pub fn flush(&mut self) -> io::Result<()> {
        self.stream.flush()
    }

    fn set_read_timeout(&self, timeout: Option<Duration>) -> io::Result<()> {
        // self.stream.get_ref().get_ref().0.set_read_timeout(timeout)
        let s = unsafe {
            ::std::net::TcpStream::from_raw_fd(self.stream.get_ref().get_ref().0.as_raw_fd())
        };
        let r = s.set_read_timeout(timeout);
        ::std::mem::forget(s);
        r
    }

    fn set_nonblocking(&self, nonblock: bool) -> io::Result<()> {
        self.stream.get_ref().get_ref().0.set_nonblocking(nonblock)
    }
}

/// A connection to the Postgres server.
///
/// It implements `Read`, `Write` and `TlsStream`, as well as `AsRawFd` on
/// Unix platforms and `AsRawSocket` on Windows platforms.
#[derive(Debug)]
pub struct Stream(Socket);

impl Read for Stream {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        self.0.read(buf)
    }
}

impl Write for Stream {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.0.write(buf)
    }

    fn flush(&mut self) -> io::Result<()> {
        self.0.flush()
    }
}

impl TlsStream for Stream {
    fn get_ref(&self) -> &Stream {
        self
    }

    fn get_mut(&mut self) -> &mut Stream {
        self
    }
}

#[cfg(unix)]
impl AsRawFd for Stream {
    fn as_raw_fd(&self) -> RawFd {
        self.0.as_raw_fd()
    }
}

#[cfg(windows)]
impl AsRawSocket for Stream {
    fn as_raw_socket(&self) -> RawSocket {
        self.0.as_raw_socket()
    }
}

fn open_socket(params: &ConnectParams) -> Result<Socket> {
    let port = params.port();
    match *params.host() {
        Host::Tcp(ref host) => {
            let mut error = None;
            for addr in (&**host, port).to_socket_addrs()? {
                let domain = match addr {
                    SocketAddr::V4(_) => Domain::ipv4(),
                    SocketAddr::V6(_) => Domain::ipv6(),
                };
                let socket = Socket::new(domain, Type::stream(), None)?;
                let addr = SockAddr::from(addr);
                let r = match params.connect_timeout() {
                    Some(timeout) => socket.connect_timeout(&addr, timeout),
                    None => socket.connect(&addr),
                };
                match r {
                    Ok(()) => return Ok(socket),
                    Err(e) => error = Some(e),
                }
            }

            Err(
                error
                    .unwrap_or_else(|| {
                        io::Error::new(
                            io::ErrorKind::InvalidInput,
                            "could not resolve any addresses",
                        )
                    })
                    .into(),
            )
        }
        #[cfg(unix)]
        Host::Unix(ref path) => {
            let path = path.join(&format!(".s.PGSQL.{}", port));
            Ok(UnixStream::connect(&path).map(|s| unsafe {
                Socket::from_raw_fd(s.into_raw_fd())
            })?)
        }
        #[cfg(not(unix))]
        Host::Unix(..) => {
            Err(
                io::Error::new(
                    io::ErrorKind::InvalidInput,
                    "unix sockets are not supported on this system",
                ).into(),
            )
        }
    }
}

pub fn initialize_stream(params: &ConnectParams, tls: TlsMode) -> Result<Box<TlsStream>> {
    let mut socket = Stream(open_socket(params)?);

    let (tls_required, handshaker) = match tls {
        TlsMode::None => return Ok(Box::new(socket)),
        TlsMode::Prefer(handshaker) => (false, handshaker),
        TlsMode::Require(handshaker) => (true, handshaker),
    };

    let mut buf = vec![];
    frontend::ssl_request(&mut buf);
    socket.write_all(&buf)?;
    socket.flush()?;

    let mut b = [0; 1];
    socket.read_exact(&mut b)?;
    if b[0] == b'N' {
        if tls_required {
            return Err(error::tls("the server does not support TLS".into()));
        } else {
            return Ok(Box::new(socket));
        }
    }

    let host = match *params.host() {
        Host::Tcp(ref host) => host,
        // Postgres doesn't support TLS over unix sockets
        Host::Unix(_) => return Err(::bad_response().into()),
    };

    handshaker.tls_handshake(host, socket).map_err(error::tls)
}
<commit_msg>Fix read timeout handling<commit_after>use std::io::{self, BufWriter, Read, Write};
use std::net::{ToSocketAddrs, SocketAddr};
use std::time::Duration;
use std::result;
use bytes::{BufMut, BytesMut};
#[cfg(unix)]
use std::os::unix::net::UnixStream;
#[cfg(unix)]
use std::os::unix::io::{AsRawFd, RawFd, FromRawFd, IntoRawFd};
#[cfg(windows)]
use std::os::windows::io::{AsRawSocket, RawSocket};
use postgres_protocol::message::frontend;
use postgres_protocol::message::backend;
use socket2::{Socket, SockAddr, Domain, Type};

use {Result, TlsMode};
use error;
use tls::TlsStream;
use params::{ConnectParams, Host};

const INITIAL_CAPACITY: usize = 8 * 1024;

pub struct MessageStream {
    stream: BufWriter<Box<TlsStream>>,
    in_buf: BytesMut,
    out_buf: Vec<u8>,
}

impl MessageStream {
    pub fn new(stream: Box<TlsStream>) -> MessageStream {
        MessageStream {
            stream: BufWriter::new(stream),
            in_buf: BytesMut::with_capacity(INITIAL_CAPACITY),
            out_buf: vec![],
        }
    }

    pub fn get_ref(&self) -> &Box<TlsStream> {
        self.stream.get_ref()
    }

    pub fn write_message<F, E>(&mut self, f: F) -> result::Result<(), E>
    where
        F: FnOnce(&mut Vec<u8>) -> result::Result<(), E>,
        E: From<io::Error>,
    {
        self.out_buf.clear();
        f(&mut self.out_buf)?;
        self.stream.write_all(&self.out_buf).map_err(From::from)
    }

    pub fn read_message(&mut self) -> io::Result<backend::Message> {
        loop {
            match backend::Message::parse(&mut self.in_buf) {
                Ok(Some(message)) => return Ok(message),
                Ok(None) => self.read_in()?,
                Err(e) => return Err(e),
            }
        }
    }

    fn read_in(&mut self) -> io::Result<()> {
        self.in_buf.reserve(1);
        match self.stream.get_mut().read(
            unsafe { self.in_buf.bytes_mut() },
        ) {
            Ok(0) => Err(io::Error::new(
                io::ErrorKind::UnexpectedEof,
                "unexpected EOF",
            )),
            Ok(n) => {
                unsafe { self.in_buf.advance_mut(n) };
                Ok(())
            }
            Err(e) => Err(e),
        }
    }

    pub fn read_message_timeout(
        &mut self,
        timeout: Duration,
    ) -> io::Result<Option<backend::Message>> {
        if self.in_buf.is_empty() {
            self.set_read_timeout(Some(timeout))?;
            let r = self.read_in();
            self.set_read_timeout(None)?;

            match r {
                Ok(()) => {}
                Err(ref e)
                    if e.kind() == io::ErrorKind::WouldBlock ||
                           e.kind() == io::ErrorKind::TimedOut => return Ok(None),
                Err(e) => return Err(e),
            }
        }

        self.read_message().map(Some)
    }

    pub fn read_message_nonblocking(&mut self) -> io::Result<Option<backend::Message>> {
        if self.in_buf.is_empty() {
            self.set_nonblocking(true)?;
            let r = self.read_in();
            self.set_nonblocking(false)?;

            match r {
                Ok(()) => {}
                Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => return Ok(None),
                Err(e) => return Err(e),
            }
        }

        self.read_message().map(Some)
    }

    pub fn flush(&mut self) -> io::Result<()> {
        self.stream.flush()
    }

    fn set_read_timeout(&self, timeout: Option<Duration>) -> io::Result<()> {
        self.stream.get_ref().get_ref().0.set_read_timeout(timeout)
    }

    fn set_nonblocking(&self, nonblock: bool) -> io::Result<()> {
        self.stream.get_ref().get_ref().0.set_nonblocking(nonblock)
    }
}

/// A connection to the Postgres server.
///
/// It implements `Read`, `Write` and `TlsStream`, as well as `AsRawFd` on
/// Unix platforms and `AsRawSocket` on Windows platforms.
#[derive(Debug)]
pub struct Stream(Socket);

impl Read for Stream {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        self.0.read(buf)
    }
}

impl Write for Stream {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.0.write(buf)
    }

    fn flush(&mut self) -> io::Result<()> {
        self.0.flush()
    }
}

impl TlsStream for Stream {
    fn get_ref(&self) -> &Stream {
        self
    }

    fn get_mut(&mut self) -> &mut Stream {
        self
    }
}

#[cfg(unix)]
impl AsRawFd for Stream {
    fn as_raw_fd(&self) -> RawFd {
        self.0.as_raw_fd()
    }
}

#[cfg(windows)]
impl AsRawSocket for Stream {
    fn as_raw_socket(&self) -> RawSocket {
        self.0.as_raw_socket()
    }
}

fn open_socket(params: &ConnectParams) -> Result<Socket> {
    let port = params.port();
    match *params.host() {
        Host::Tcp(ref host) => {
            let mut error = None;
            for addr in (&**host, port).to_socket_addrs()? {
                let domain = match addr {
                    SocketAddr::V4(_) => Domain::ipv4(),
                    SocketAddr::V6(_) => Domain::ipv6(),
                };
                let socket = Socket::new(domain, Type::stream(), None)?;
                let addr = SockAddr::from(addr);
                let r = match params.connect_timeout() {
                    Some(timeout) => socket.connect_timeout(&addr, timeout),
                    None => socket.connect(&addr),
                };
                match r {
                    Ok(()) => return Ok(socket),
                    Err(e) => error = Some(e),
                }
            }

            Err(
                error
                    .unwrap_or_else(|| {
                        io::Error::new(
                            io::ErrorKind::InvalidInput,
                            "could not resolve any addresses",
                        )
                    })
                    .into(),
            )
        }
        #[cfg(unix)]
        Host::Unix(ref path) => {
            let path = path.join(&format!(".s.PGSQL.{}", port));
            Ok(UnixStream::connect(&path).map(|s| unsafe {
                Socket::from_raw_fd(s.into_raw_fd())
            })?)
        }
        #[cfg(not(unix))]
        Host::Unix(..) => {
            Err(
                io::Error::new(
                    io::ErrorKind::InvalidInput,
                    "unix sockets are not supported on this system",
                ).into(),
            )
        }
    }
}

pub fn initialize_stream(params: &ConnectParams, tls: TlsMode) -> Result<Box<TlsStream>> {
    let mut socket = Stream(open_socket(params)?);

    let (tls_required, handshaker) = match tls {
        TlsMode::None => return Ok(Box::new(socket)),
        TlsMode::Prefer(handshaker) => (false, handshaker),
        TlsMode::Require(handshaker) => (true, handshaker),
    };

    let mut buf = vec![];
    frontend::ssl_request(&mut buf);
    socket.write_all(&buf)?;
    socket.flush()?;

    let mut b = [0; 1];
    socket.read_exact(&mut b)?;
    if b[0] == b'N' {
        if tls_required {
            return Err(error::tls("the server does not support TLS".into()));
        } else {
            return Ok(Box::new(socket));
        }
    }

    let host = match *params.host() {
        Host::Tcp(ref host) => host,
        // Postgres doesn't support TLS over unix sockets
        Host::Unix(_) => return Err(::bad_response().into()),
    };

    handshaker.tls_handshake(host, socket).map_err(error::tls)
}
<|endoftext|>


<commit_before>//! Serializes and deserializes Ardite Schema Definitions from different data
//! formats such as JSON and YAML.

use std::io::prelude::*;
use std::io::BufReader;
use std::fs::File;
use std::path::PathBuf;

use regex::Regex;
use serde_json;
use serde_yaml;

use error::{Error, ErrorCode};
use schema::{Definition, Type, Schema};
use value::Value;

mod types;

use schema::serde::types::*;

/// Gets an Ardite Schema Definition from a file. Aims to support mainly the
/// JSON and YAML formats.
pub fn from_file(path: PathBuf) -> Result<Definition, Error> {
  let extension = path.extension().map_or("", |s| s.to_str().unwrap());
  let mut file = try!(File::open(&path));
  match extension {
    "json" => {
      let reader = BufReader::new(file);
      let definition: SerdeDefinition = try!(serde_json::from_reader(reader));
      Ok(try!(serde_definition_into_definition(definition)))
    },
    "yml" => {
      let mut string = String::new();
      try!(file.read_to_string(&mut string));
      let definition: SerdeDefinition = try!(serde_yaml::from_str(&string));
      Ok(try!(serde_definition_into_definition(definition)))
    },
    _ => Err(Error::new(
      ErrorCode::NotAcceptable,
      format!("File extension '{}' cannot be deserialized in '{}'.", extension, path.display()),
      Some("Use a recognizable file extension like '.json' or '.yml'.".to_owned())
    ))
  }
}

/// Transforms the intermediary type into the useful type.
fn serde_definition_into_definition(serde_definition: SerdeDefinition) -> Result<Definition, Error> {
  let mut definition = Definition::new();

  for (key, schema) in serde_definition.types.into_iter() {
    let mut type_ = Type::new(key);
    type_.set_boxed_schema(try!(serde_schema_into_schema(schema)));
    definition.add_type(type_);
  }

  Ok(definition)
}

/// Transforms the intermediary type into the useful type.
fn serde_schema_into_schema(serde_schema: SerdeSchema) -> Result<Box<Schema + 'static>, Error> {
  match serde_schema.type_ {
    Some(type_) => match type_.as_str() {
      "null" => Ok(Box::new(Schema::null())),
      "boolean" => Ok(Box::new(Schema::boolean())),
      type_ @ "number" | type_ @ "integer" => {
        let mut schema = Schema::number();
        if type_ == "integer" { schema.set_multiple_of(1.0); }
        else if let Some(multiple_of) = serde_schema.multiple_of { schema.set_multiple_of(multiple_of); }
        if let Some(minimum) = serde_schema.minimum { schema.set_minimum(minimum); }
        if let Some(maximum) = serde_schema.maximum { schema.set_maximum(maximum); }
        if serde_schema.exclusive_minimum.unwrap_or(false) { schema.enable_exclusive_minimum(); }
        if serde_schema.exclusive_maximum.unwrap_or(false) { schema.enable_exclusive_maximum(); }
        Ok(Box::new(schema))
      },
      "string" => {
        let mut schema = Schema::string();
        if let Some(min_length) = serde_schema.min_length { schema.set_min_length(min_length); }
        if let Some(max_length) = serde_schema.max_length { schema.set_max_length(max_length); }
        if let Some(pattern) = serde_schema.pattern.and_then(|p| Regex::new(&p).ok()) { schema.set_pattern(pattern); }
        Ok(Box::new(schema))
      },
      "array" => {
        let mut schema = Schema::array();
        if let Some(items) = serde_schema.items { schema.set_boxed_items(try!(serde_schema_into_schema(*items))); }
        Ok(Box::new(schema))
      },
      "object" => {
        let mut schema = Schema::object();
        schema.set_required(serde_schema.required.unwrap_or_default());
        if serde_schema.additional_properties.unwrap_or(false) { schema.enable_additional_properties(); }
        for (key, value) in serde_schema.properties.unwrap_or_default() {
          schema.add_boxed_property(key, try!(serde_schema_into_schema(value)));
        }
        Ok(Box::new(schema))
      },
      _ => Err(Error::invalid(
        format!("Invalid type '{}'.", type_),
        format!("Use a permitted type like 'string' and not '{}'.", type_)
      ))
    },
    None => {
      if let Some(enum_) = serde_schema.enum_ {
        Ok(Box::new(Schema::enum_(enum_.into_iter().map(Value::String).collect())))
      } else {
        Err(Error::invalid("No schema type specified.", "Set a `type` property or an `enum` property."))
      }
    }
  }
}

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

  use std::path::PathBuf;

  use schema::definition::create_basic;

  #[test]
  fn test_basic_json() {
    assert_eq!(from_file(PathBuf::from("tests/fixtures/definitions/basic.json")).unwrap(), create_basic());
  }

  #[test]
  fn test_basic_yaml() {
    assert_eq!(from_file(PathBuf::from("tests/fixtures/definitions/basic.yml")).unwrap(), create_basic());
  }
}
<commit_msg>Use serde_yaml::from_reader<commit_after>//! Serializes and deserializes Ardite Schema Definitions from different data
//! formats such as JSON and YAML.

use std::io::BufReader;
use std::fs::File;
use std::path::PathBuf;

use regex::Regex;
use serde_json;
use serde_yaml;

use error::{Error, ErrorCode};
use schema::{Definition, Type, Schema};
use value::Value;

mod types;

use schema::serde::types::*;

/// Gets an Ardite Schema Definition from a file. Aims to support mainly the
/// JSON and YAML formats.
pub fn from_file(path: PathBuf) -> Result<Definition, Error> {
  let extension = path.extension().map_or("", |s| s.to_str().unwrap());
  let file = try!(File::open(&path));
  match extension {
    "json" => {
      let reader = BufReader::new(file);
      let definition: SerdeDefinition = try!(serde_json::from_reader(reader));
      Ok(try!(serde_definition_into_definition(definition)))
    },
    "yml" => {
      let reader = BufReader::new(file);
      let definition: SerdeDefinition = try!(serde_yaml::from_reader(reader));
      Ok(try!(serde_definition_into_definition(definition)))
    },
    _ => Err(Error::new(
      ErrorCode::NotAcceptable,
      format!("File extension '{}' cannot be deserialized in '{}'.", extension, path.display()),
      Some("Use a recognizable file extension like '.json' or '.yml'.".to_owned())
    ))
  }
}

/// Transforms the intermediary type into the useful type.
fn serde_definition_into_definition(serde_definition: SerdeDefinition) -> Result<Definition, Error> {
  let mut definition = Definition::new();

  for (key, schema) in serde_definition.types.into_iter() {
    let mut type_ = Type::new(key);
    type_.set_boxed_schema(try!(serde_schema_into_schema(schema)));
    definition.add_type(type_);
  }

  Ok(definition)
}

/// Transforms the intermediary type into the useful type.
fn serde_schema_into_schema(serde_schema: SerdeSchema) -> Result<Box<Schema + 'static>, Error> {
  match serde_schema.type_ {
    Some(type_) => match type_.as_str() {
      "null" => Ok(Box::new(Schema::null())),
      "boolean" => Ok(Box::new(Schema::boolean())),
      type_ @ "number" | type_ @ "integer" => {
        let mut schema = Schema::number();
        if type_ == "integer" { schema.set_multiple_of(1.0); }
        else if let Some(multiple_of) = serde_schema.multiple_of { schema.set_multiple_of(multiple_of); }
        if let Some(minimum) = serde_schema.minimum { schema.set_minimum(minimum); }
        if let Some(maximum) = serde_schema.maximum { schema.set_maximum(maximum); }
        if serde_schema.exclusive_minimum.unwrap_or(false) { schema.enable_exclusive_minimum(); }
        if serde_schema.exclusive_maximum.unwrap_or(false) { schema.enable_exclusive_maximum(); }
        Ok(Box::new(schema))
      },
      "string" => {
        let mut schema = Schema::string();
        if let Some(min_length) = serde_schema.min_length { schema.set_min_length(min_length); }
        if let Some(max_length) = serde_schema.max_length { schema.set_max_length(max_length); }
        if let Some(pattern) = serde_schema.pattern.and_then(|p| Regex::new(&p).ok()) { schema.set_pattern(pattern); }
        Ok(Box::new(schema))
      },
      "array" => {
        let mut schema = Schema::array();
        if let Some(items) = serde_schema.items { schema.set_boxed_items(try!(serde_schema_into_schema(*items))); }
        Ok(Box::new(schema))
      },
      "object" => {
        let mut schema = Schema::object();
        schema.set_required(serde_schema.required.unwrap_or_default());
        if serde_schema.additional_properties.unwrap_or(false) { schema.enable_additional_properties(); }
        for (key, value) in serde_schema.properties.unwrap_or_default() {
          schema.add_boxed_property(key, try!(serde_schema_into_schema(value)));
        }
        Ok(Box::new(schema))
      },
      _ => Err(Error::invalid(
        format!("Invalid type '{}'.", type_),
        format!("Use a permitted type like 'string' and not '{}'.", type_)
      ))
    },
    None => {
      if let Some(enum_) = serde_schema.enum_ {
        Ok(Box::new(Schema::enum_(enum_.into_iter().map(Value::String).collect())))
      } else {
        Err(Error::invalid("No schema type specified.", "Set a `type` property or an `enum` property."))
      }
    }
  }
}

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

  use std::path::PathBuf;

  use schema::definition::create_basic;

  #[test]
  fn test_basic_json() {
    assert_eq!(from_file(PathBuf::from("tests/fixtures/definitions/basic.json")).unwrap(), create_basic());
  }

  #[test]
  fn test_basic_yaml() {
    assert_eq!(from_file(PathBuf::from("tests/fixtures/definitions/basic.yml")).unwrap(), create_basic());
  }
}
<|endoftext|>


<commit_before>//
//   Copyright 2016, 2017 Andrew Hunter
//
//   Licensed under the Apache License, Version 2.0 (the "License");
//   you may not use this file except in compliance with the License.
//   You may obtain a copy of the License at
//
//       http://www.apache.org/licenses/LICENSE-2.0
//
//   Unless required by applicable law or agreed to in writing, software
//   distributed under the License is distributed on an "AS IS" BASIS,
//   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//   See the License for the specific language governing permissions and
//   limitations under the License.
//

//!
//! A DFA that matches transitions against symbol ranges.
//!

use std::mem::*;

use super::dfa_builder::*;
use super::pattern_matcher::*;
use super::symbol_range::*;
use super::state_machine::*;

///
/// DFA that decides on transitions based on non-overlapping, sorted lists of input symbols
///
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SymbolRangeDfa<InputSymbol: Ord, OutputSymbol> {
    //
    // Indexes of where each state starts in the transition table (it ends at the start of the next state)
    //
    states: Vec<usize>,

    //
    // The transitions making up this DFA
    //
    transitions: Vec<(SymbolRange<InputSymbol>, StateId)>,

    //
    // The accepting symbol for each state
    //
    accept: Vec<Option<OutputSymbol>>
}

///
/// DFA builder that creates RangeDfas
///
pub struct SymbolRangeDfaBuilder<InputSymbol: Ord, OutputSymbol> {
    states: Vec<usize>,
    transitions: Vec<(SymbolRange<InputSymbol>, StateId)>,
    accept: Vec<Option<OutputSymbol>>
}

impl<InputSymbol: Ord, OutputSymbol> SymbolRangeDfaBuilder<InputSymbol, OutputSymbol> {
    pub fn new() -> SymbolRangeDfaBuilder<InputSymbol, OutputSymbol> {
        SymbolRangeDfaBuilder { states: vec![], transitions: vec![], accept: vec![] }
    }
}

impl<InputSymbol: Ord, OutputSymbol> DfaBuilder<SymbolRange<InputSymbol>, OutputSymbol, SymbolRangeDfa<InputSymbol, OutputSymbol>> for SymbolRangeDfaBuilder<InputSymbol, OutputSymbol> {
    fn start_state(&mut self) {
        // Begin the next state
        self.states.push(self.transitions.len());
        self.accept.push(None);
    }

    fn transition(&mut self, symbol: SymbolRange<InputSymbol>, target_state: StateId) {
        self.transitions.push((symbol, target_state));
    }

    fn accept(&mut self, symbol: OutputSymbol) {
        self.accept.pop();
        self.accept.push(Some(symbol));
    }

    fn build(self) -> SymbolRangeDfa<InputSymbol, OutputSymbol> {
        // Turn into a RangeDfa
        let mut result = SymbolRangeDfa { states: self.states, transitions: self.transitions, accept: self.accept };

        // 'Cap' the last state so we don't need to special-case it later
        // ie, we can always find the index of the last symbol by looking at the next state and don't need to handle the final state differently
        result.states.push(result.transitions.len());

        result
    }
}

impl<InputSymbol: Ord+Clone, OutputSymbol> StateMachine<SymbolRange<InputSymbol>, OutputSymbol> for SymbolRangeDfa<InputSymbol, OutputSymbol> {
    ///
    /// Returns the number of states in this state machine
    ///
    /// Note that if state x exists then state x-1 is also expected to exist provided x > 0. This means that this returns the
    /// first unused state in this state machine.
    ///
    fn count_states(&self) -> StateId {
        (self.states.len()-1) as StateId
    }

    ///
    /// Returns the transitions for a particular symbol
    ///
    fn get_transitions_for_state(&self, state: StateId) -> Vec<(SymbolRange<InputSymbol>, StateId)> {
        let mut result = vec![];

        let start_index = self.states[state as usize];
        let end_index   = self.states[(state+1) as usize];

        for transit_index in start_index..end_index {
            let (ref range, target_state) = self.transitions[transit_index];

            result.push((range.clone(), target_state));
        }

        result
    }

    ///
    /// If a state is an accepting state, then this returns the output symbol that should be produced if this is the longest match
    ///
    fn output_symbol_for_state(&self, state: StateId) -> Option<&OutputSymbol> {
        self.accept[state as usize].as_ref()
    }
}

///
/// A state of a symbol range state machine
///
#[derive(Clone)]
pub struct SymbolRangeState<'a, InputSymbol: Ord+'a, OutputSymbol: 'a> {
    // The current state of the state machine
    state: StateId,

    // The number of symbols that have been processed so far
    count: usize,

    // If something other than none, the most recent accepting state
    accept: Option<(usize, &'a OutputSymbol)>,

    // The state machine this is running
    state_machine: &'a SymbolRangeDfa<InputSymbol, OutputSymbol>
}

impl<InputSymbol: Ord, OutputSymbol> SymbolRangeDfa<InputSymbol, OutputSymbol> {
    ///
    /// Returns a `MatchAction` for the initial state of the DFA
    ///
    pub fn start<'a>(&'a self) -> MatchAction<'a, OutputSymbol, SymbolRangeState<'a, InputSymbol, OutputSymbol>> {
        // TODO: if state 0 is accepting, then this will erroneously not move straight to the accepting state
        if let Some(ref outputsymbol) = self.accept[0] {
            More(SymbolRangeState { state: 0, count: 0, accept: Some((0, outputsymbol)), state_machine: self })
        } else {
            More(SymbolRangeState { state: 0, count: 0, accept: None, state_machine: self })
        }
    }

    ///
    /// Returns a description of this DFA
    ///
    pub fn description(&self) -> String {
        let state_size          = size_of::<usize>() * self.states.len();
        let transitions_size    = size_of::<(SymbolRange<InputSymbol>, StateId)>() * self.transitions.len();
        let accept_size         = size_of::<Option<OutputSymbol>>() * self.accept.len();
        let total_size          = state_size + transitions_size + accept_size;

        format!("SymbolRangeDfa: {} states, {} total transitions. {} bytes", self.states.len(), self.transitions.len(), total_size)
    }
}

impl<'a, InputSymbol: Ord+'a, OutputSymbol: 'a> MatchingState<'a, InputSymbol, OutputSymbol> for SymbolRangeState<'a, InputSymbol, OutputSymbol> {
    fn next(self, symbol: InputSymbol) -> MatchAction<'a, OutputSymbol, Self> {
        // The transition range is defined by the current state
        let start_transition    = self.state_machine.states[self.state as usize];
        let end_transition      = self.state_machine.states[self.state as usize+1];

        // See if there is an input symbol matching this transition
        // TODO: consider binary searching for states with large numbers of transitions? (Do these occur regularly in patterns that people use?)
        for transit in start_transition..end_transition {
            // Test this transition
            let (ref range, new_state) = self.state_machine.transitions[transit];

            if range.includes(&symbol) {
                // Found a transition to a new state: result will be `More(new state)`
                let new_count = self.count+1;

                // If the new state is an accepting state, then remember it in case we reach a rejecting state later
                let new_accept = if let Some(ref output) = self.state_machine.accept[new_state as usize] {
                    Some((new_count, output))
                } else {
                    self.accept
                };

                // Action is 'More'
                // TODO: might be an option to return Accept or Reject here if the new state has no transitions
                // (Possible performance advantage, but depends on the regex and input conditions)
                return More(SymbolRangeState { state: new_state, count: new_count, accept: new_accept, state_machine: self.state_machine });
            }
        }

        // No matches: finish the state machine
        self.finish()
    }

    fn finish(self) -> MatchAction<'a, OutputSymbol, Self> {
        if let Some(accept_state) = self.accept {
            // We found an accepting state earlier on, so return that
            let (length, symbol) = accept_state;
            Accept(length, symbol)
        } else {
            // No accepting state was found while this state machine was running
            Reject
        }
    }
}

#[cfg(test)]
mod test {
    use super::super::dfa_builder::*;
    use super::super::symbol_range::*;
    use super::super::pattern_matcher::*;
    use super::super::state_machine::*;
    use super::*;

    #[test]
    fn can_build_state_machine() {
        let mut builder = SymbolRangeDfaBuilder::new();

        // State 0: '0', move to state 1
        builder.start_state();
        builder.transition(SymbolRange::new(0, 0), 1);

        // State 1: accept, output symbol "Success"
        builder.start_state();
        builder.accept("Success");

        // Create the state machine
        let state_machine = builder.build();

        assert!(state_machine.count_states() == 2);
        assert!(state_machine.output_symbol_for_state(0) == None);
        assert!(state_machine.output_symbol_for_state(1) == Some(&"Success"));
        assert!(state_machine.get_transitions_for_state(0) == vec![(SymbolRange::new(0,0), 1)]);
    }

    #[test]
    fn can_accept_single_symbol() {
        let mut builder = SymbolRangeDfaBuilder::new();

        // State 0: '0', move to state 1
        builder.start_state();
        builder.transition(SymbolRange::new(0, 0), 1);

        // State 1: accept, output symbol "Success"
        builder.start_state();
        builder.accept("Success");

        // Create the state machine
        let state_machine = builder.build();

        // Run the first state
        let mut action = state_machine.start();

        if let More(next_state) = action {
            action = next_state.next(0);
        }

        if let More(next_state) = action {
            action = next_state.next(0);

            // Should have reached an accepting state (read one character)
            if let Accept(count, symbol) = action {
                // One symbol accepted
                assert!(count == 1);

                // Output symbol correct
                assert!(symbol == &"Success");
            } else {
                // Should have accepted here (the second '0' is rejected)
                assert!(false);
            }
        } else {
            // State machine did not accept the character
            assert!(false);
        }
    }
}
<commit_msg>Add some code to compress overlapping symbol ranges into a single symbol range<commit_after>//
//   Copyright 2016, 2017 Andrew Hunter
//
//   Licensed under the Apache License, Version 2.0 (the "License");
//   you may not use this file except in compliance with the License.
//   You may obtain a copy of the License at
//
//       http://www.apache.org/licenses/LICENSE-2.0
//
//   Unless required by applicable law or agreed to in writing, software
//   distributed under the License is distributed on an "AS IS" BASIS,
//   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//   See the License for the specific language governing permissions and
//   limitations under the License.
//

//!
//! A DFA that matches transitions against symbol ranges.
//!

use std::mem::*;

use super::countable::*;
use super::dfa_builder::*;
use super::pattern_matcher::*;
use super::symbol_range::*;
use super::state_machine::*;

///
/// DFA that decides on transitions based on non-overlapping, sorted lists of input symbols
///
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SymbolRangeDfa<InputSymbol: Ord, OutputSymbol> {
    //
    // Indexes of where each state starts in the transition table (it ends at the start of the next state)
    //
    states: Vec<usize>,

    //
    // The transitions making up this DFA
    //
    transitions: Vec<(SymbolRange<InputSymbol>, StateId)>,

    //
    // The accepting symbol for each state
    //
    accept: Vec<Option<OutputSymbol>>
}

///
/// DFA builder that creates RangeDfas
///
pub struct SymbolRangeDfaBuilder<InputSymbol: Ord+Countable, OutputSymbol> {
    states: Vec<usize>,
    transitions: Vec<(SymbolRange<InputSymbol>, StateId)>,
    accept: Vec<Option<OutputSymbol>>
}

impl<InputSymbol: Ord+Countable, OutputSymbol> SymbolRangeDfaBuilder<InputSymbol, OutputSymbol> {
    pub fn new() -> SymbolRangeDfaBuilder<InputSymbol, OutputSymbol> {
        SymbolRangeDfaBuilder { states: vec![], transitions: vec![], accept: vec![] }
    }
}

impl<InputSymbol: Ord+Countable+Clone, OutputSymbol> DfaBuilder<SymbolRange<InputSymbol>, OutputSymbol, SymbolRangeDfa<InputSymbol, OutputSymbol>> for SymbolRangeDfaBuilder<InputSymbol, OutputSymbol> {
    fn start_state(&mut self) {
        // Join any adjoining transitions
        if let Some(start_index) = self.states.last() {
            let mut index = start_index+1;

            while index < self.transitions.len() {
                let (prev_symbols, prev_state) = self.transitions[index-1].clone();
                let (next_symbols, next_state) = self.transitions[index].clone();

                if prev_state == next_state && prev_symbols.highest.next() == next_symbols.lowest {
                    // The previous transition and the next transition can be merged
                    self.transitions[index-1] = (SymbolRange { lowest: prev_symbols.lowest, highest: next_symbols.highest }, prev_state);

                    self.transitions.remove(index);
                    index -= 1;
                }

                index += 1;
            }
        }

        // Begin the next state
        self.states.push(self.transitions.len());
        self.accept.push(None);
    }

    fn transition(&mut self, symbol: SymbolRange<InputSymbol>, target_state: StateId) {
        self.transitions.push((symbol, target_state));
    }

    fn accept(&mut self, symbol: OutputSymbol) {
        self.accept.pop();
        self.accept.push(Some(symbol));
    }

    fn build(self) -> SymbolRangeDfa<InputSymbol, OutputSymbol> {
        // Turn into a RangeDfa
        let mut result = SymbolRangeDfa { states: self.states, transitions: self.transitions, accept: self.accept };

        // 'Cap' the last state so we don't need to special-case it later
        // ie, we can always find the index of the last symbol by looking at the next state and don't need to handle the final state differently
        result.states.push(result.transitions.len());

        result
    }
}

impl<InputSymbol: Ord+Clone, OutputSymbol> StateMachine<SymbolRange<InputSymbol>, OutputSymbol> for SymbolRangeDfa<InputSymbol, OutputSymbol> {
    ///
    /// Returns the number of states in this state machine
    ///
    /// Note that if state x exists then state x-1 is also expected to exist provided x > 0. This means that this returns the
    /// first unused state in this state machine.
    ///
    fn count_states(&self) -> StateId {
        (self.states.len()-1) as StateId
    }

    ///
    /// Returns the transitions for a particular symbol
    ///
    fn get_transitions_for_state(&self, state: StateId) -> Vec<(SymbolRange<InputSymbol>, StateId)> {
        let mut result = vec![];

        let start_index = self.states[state as usize];
        let end_index   = self.states[(state+1) as usize];

        for transit_index in start_index..end_index {
            let (ref range, target_state) = self.transitions[transit_index];

            result.push((range.clone(), target_state));
        }

        result
    }

    ///
    /// If a state is an accepting state, then this returns the output symbol that should be produced if this is the longest match
    ///
    fn output_symbol_for_state(&self, state: StateId) -> Option<&OutputSymbol> {
        self.accept[state as usize].as_ref()
    }
}

///
/// A state of a symbol range state machine
///
#[derive(Clone)]
pub struct SymbolRangeState<'a, InputSymbol: Ord+'a, OutputSymbol: 'a> {
    // The current state of the state machine
    state: StateId,

    // The number of symbols that have been processed so far
    count: usize,

    // If something other than none, the most recent accepting state
    accept: Option<(usize, &'a OutputSymbol)>,

    // The state machine this is running
    state_machine: &'a SymbolRangeDfa<InputSymbol, OutputSymbol>
}

impl<InputSymbol: Ord, OutputSymbol> SymbolRangeDfa<InputSymbol, OutputSymbol> {
    ///
    /// Returns a `MatchAction` for the initial state of the DFA
    ///
    pub fn start<'a>(&'a self) -> MatchAction<'a, OutputSymbol, SymbolRangeState<'a, InputSymbol, OutputSymbol>> {
        // TODO: if state 0 is accepting, then this will erroneously not move straight to the accepting state
        if let Some(ref outputsymbol) = self.accept[0] {
            More(SymbolRangeState { state: 0, count: 0, accept: Some((0, outputsymbol)), state_machine: self })
        } else {
            More(SymbolRangeState { state: 0, count: 0, accept: None, state_machine: self })
        }
    }

    ///
    /// Returns a description of this DFA
    ///
    pub fn description(&self) -> String {
        let state_size          = size_of::<usize>() * self.states.len();
        let transitions_size    = size_of::<(SymbolRange<InputSymbol>, StateId)>() * self.transitions.len();
        let accept_size         = size_of::<Option<OutputSymbol>>() * self.accept.len();
        let total_size          = state_size + transitions_size + accept_size;

        format!("SymbolRangeDfa: {} states, {} total transitions. {} bytes", self.states.len(), self.transitions.len(), total_size)
    }
}

impl<'a, InputSymbol: Ord+'a, OutputSymbol: 'a> MatchingState<'a, InputSymbol, OutputSymbol> for SymbolRangeState<'a, InputSymbol, OutputSymbol> {
    fn next(self, symbol: InputSymbol) -> MatchAction<'a, OutputSymbol, Self> {
        // The transition range is defined by the current state
        let start_transition    = self.state_machine.states[self.state as usize];
        let end_transition      = self.state_machine.states[self.state as usize+1];

        // See if there is an input symbol matching this transition
        // TODO: consider binary searching for states with large numbers of transitions? (Do these occur regularly in patterns that people use?)
        for transit in start_transition..end_transition {
            // Test this transition
            let (ref range, new_state) = self.state_machine.transitions[transit];

            if range.includes(&symbol) {
                // Found a transition to a new state: result will be `More(new state)`
                let new_count = self.count+1;

                // If the new state is an accepting state, then remember it in case we reach a rejecting state later
                let new_accept = if let Some(ref output) = self.state_machine.accept[new_state as usize] {
                    Some((new_count, output))
                } else {
                    self.accept
                };

                // Action is 'More'
                // TODO: might be an option to return Accept or Reject here if the new state has no transitions
                // (Possible performance advantage, but depends on the regex and input conditions)
                return More(SymbolRangeState { state: new_state, count: new_count, accept: new_accept, state_machine: self.state_machine });
            }
        }

        // No matches: finish the state machine
        self.finish()
    }

    fn finish(self) -> MatchAction<'a, OutputSymbol, Self> {
        if let Some(accept_state) = self.accept {
            // We found an accepting state earlier on, so return that
            let (length, symbol) = accept_state;
            Accept(length, symbol)
        } else {
            // No accepting state was found while this state machine was running
            Reject
        }
    }
}

#[cfg(test)]
mod test {
    use super::super::dfa_builder::*;
    use super::super::symbol_range::*;
    use super::super::pattern_matcher::*;
    use super::super::state_machine::*;
    use super::*;

    #[test]
    fn can_build_state_machine() {
        let mut builder = SymbolRangeDfaBuilder::new();

        // State 0: '0', move to state 1
        builder.start_state();
        builder.transition(SymbolRange::new(0, 0), 1);

        // State 1: accept, output symbol "Success"
        builder.start_state();
        builder.accept("Success");

        // Create the state machine
        let state_machine = builder.build();

        assert!(state_machine.count_states() == 2);
        assert!(state_machine.output_symbol_for_state(0) == None);
        assert!(state_machine.output_symbol_for_state(1) == Some(&"Success"));
        assert!(state_machine.get_transitions_for_state(0) == vec![(SymbolRange::new(0,0), 1)]);
    }

    #[test]
    fn can_accept_single_symbol() {
        let mut builder = SymbolRangeDfaBuilder::new();

        // State 0: '0', move to state 1
        builder.start_state();
        builder.transition(SymbolRange::new(0, 0), 1);

        // State 1: accept, output symbol "Success"
        builder.start_state();
        builder.accept("Success");

        // Create the state machine
        let state_machine = builder.build();

        // Run the first state
        let mut action = state_machine.start();

        if let More(next_state) = action {
            action = next_state.next(0);
        }

        if let More(next_state) = action {
            action = next_state.next(0);

            // Should have reached an accepting state (read one character)
            if let Accept(count, symbol) = action {
                // One symbol accepted
                assert!(count == 1);

                // Output symbol correct
                assert!(symbol == &"Success");
            } else {
                // Should have accepted here (the second '0' is rejected)
                assert!(false);
            }
        } else {
            // State machine did not accept the character
            assert!(false);
        }
    }
}
<|endoftext|>