<commit_before>//! Creation of ar archives like for the lib and staticlib crate type

use std::collections::BTreeMap;
use std::convert::TryFrom;
use std::fs::File;
use std::io::{self, Read, Seek};
use std::path::{Path, PathBuf};

use rustc_codegen_ssa::back::archive::{find_library, ArchiveBuilder};
use rustc_codegen_ssa::METADATA_FILENAME;
use rustc_session::Session;

use object::read::archive::ArchiveFile;
use object::{Object, ObjectSymbol, ReadCache, SymbolKind};

#[derive(Debug)]
enum ArchiveEntry {
    FromArchive { archive_index: usize, file_range: (u64, u64) },
    File(PathBuf),
}

pub(crate) struct ArArchiveBuilder<'a> {
    sess: &'a Session,
    dst: PathBuf,
    lib_search_paths: Vec<PathBuf>,
    use_gnu_style_archive: bool,
    no_builtin_ranlib: bool,

    src_archives: Vec<File>,
    // Don't use `HashMap` here, as the order is important. `rust.metadata.bin` must always be at
    // the end of an archive for linkers to not get confused.
    entries: Vec<(String, ArchiveEntry)>,
}

impl<'a> ArchiveBuilder<'a> for ArArchiveBuilder<'a> {
    fn new(sess: &'a Session, output: &Path, input: Option<&Path>) -> Self {
        use rustc_codegen_ssa::back::link::archive_search_paths;

        let (src_archives, entries) = if let Some(input) = input {
            let read_cache = ReadCache::new(File::open(input).unwrap());
            let archive = ArchiveFile::parse(&read_cache).unwrap();
            let mut entries = Vec::new();

            for entry in archive.members() {
                let entry = entry.unwrap();
                entries.push((
                    String::from_utf8(entry.name().to_vec()).unwrap(),
                    ArchiveEntry::FromArchive { archive_index: 0, file_range: entry.file_range() },
                ));
            }

            (vec![read_cache.into_inner()], entries)
        } else {
            (vec![], Vec::new())
        };

        ArArchiveBuilder {
            sess,
            dst: output.to_path_buf(),
            lib_search_paths: archive_search_paths(sess),
            use_gnu_style_archive: sess.target.archive_format == "gnu",
            // FIXME fix builtin ranlib on macOS
            no_builtin_ranlib: sess.target.is_like_osx,

            src_archives,
            entries,
        }
    }

    fn src_files(&mut self) -> Vec<String> {
        self.entries.iter().map(|(name, _)| name.clone()).collect()
    }

    fn remove_file(&mut self, name: &str) {
        let index = self
            .entries
            .iter()
            .position(|(entry_name, _)| entry_name == name)
            .expect("Tried to remove file not existing in src archive");
        self.entries.remove(index);
    }

    fn add_file(&mut self, file: &Path) {
        self.entries.push((
            file.file_name().unwrap().to_str().unwrap().to_string(),
            ArchiveEntry::File(file.to_owned()),
        ));
    }

    fn add_native_library(&mut self, name: rustc_span::symbol::Symbol, verbatim: bool) {
        let location = find_library(name, verbatim, &self.lib_search_paths, self.sess);
        self.add_archive(location.clone(), |_| false).unwrap_or_else(|e| {
            panic!("failed to add native library {}: {}", location.to_string_lossy(), e);
        });
    }

    fn add_rlib(
        &mut self,
        rlib: &Path,
        name: &str,
        lto: bool,
        skip_objects: bool,
    ) -> io::Result<()> {
        let obj_start = name.to_owned();

        self.add_archive(rlib.to_owned(), move |fname: &str| {
            // Ignore metadata files, no matter the name.
            if fname == METADATA_FILENAME {
                return true;
            }

            // Don't include Rust objects if LTO is enabled
            if lto && fname.starts_with(&obj_start) && fname.ends_with(".o") {
                return true;
            }

            // Otherwise if this is *not* a rust object and we're skipping
            // objects then skip this file
            if skip_objects && (!fname.starts_with(&obj_start) || !fname.ends_with(".o")) {
                return true;
            }

            // ok, don't skip this
            false
        })
    }

    fn update_symbols(&mut self) {}

    fn build(mut self) {
        enum BuilderKind {
            Bsd(ar::Builder<File>),
            Gnu(ar::GnuBuilder<File>),
        }

        let sess = self.sess;

        let mut symbol_table = BTreeMap::new();

        let mut entries = Vec::new();

        for (entry_name, entry) in self.entries {
            // FIXME only read the symbol table of the object files to avoid having to keep all
            // object files in memory at once, or read them twice.
            let data = match entry {
                ArchiveEntry::FromArchive { archive_index, file_range } => {
                    // FIXME read symbols from symtab
                    let src_read_cache = &mut self.src_archives[archive_index];

                    src_read_cache.seek(io::SeekFrom::Start(file_range.0)).unwrap();
                    let mut data = std::vec::from_elem(0, usize::try_from(file_range.1).unwrap());
                    src_read_cache.read_exact(&mut data).unwrap();

                    data
                }
                ArchiveEntry::File(file) => std::fs::read(file).unwrap_or_else(|err| {
                    sess.fatal(&format!(
                        "error while reading object file during archive building: {}",
                        err
                    ));
                }),
            };

            if !self.no_builtin_ranlib {
                match object::File::parse(&*data) {
                    Ok(object) => {
                        symbol_table.insert(
                            entry_name.as_bytes().to_vec(),
                            object
                                .symbols()
                                .filter_map(|symbol| {
                                    if symbol.is_undefined()
                                        || symbol.is_local()
                                        || symbol.kind() != SymbolKind::Data
                                            && symbol.kind() != SymbolKind::Text
                                            && symbol.kind() != SymbolKind::Tls
                                    {
                                        None
                                    } else {
                                        symbol.name().map(|name| name.as_bytes().to_vec()).ok()
                                    }
                                })
                                .collect::<Vec<_>>(),
                        );
                    }
                    Err(err) => {
                        let err = err.to_string();
                        if err == "Unknown file magic" {
                            // Not an object file; skip it.
                        } else {
                            sess.fatal(&format!(
                                "error parsing `{}` during archive creation: {}",
                                entry_name, err
                            ));
                        }
                    }
                }
            }

            entries.push((entry_name, data));
        }

        let mut builder = if self.use_gnu_style_archive {
            BuilderKind::Gnu(
                ar::GnuBuilder::new(
                    File::create(&self.dst).unwrap_or_else(|err| {
                        sess.fatal(&format!(
                            "error opening destination during archive building: {}",
                            err
                        ));
                    }),
                    entries.iter().map(|(name, _)| name.as_bytes().to_vec()).collect(),
                    ar::GnuSymbolTableFormat::Size32,
                    symbol_table,
                )
                .unwrap(),
            )
        } else {
            BuilderKind::Bsd(
                ar::Builder::new(
                    File::create(&self.dst).unwrap_or_else(|err| {
                        sess.fatal(&format!(
                            "error opening destination during archive building: {}",
                            err
                        ));
                    }),
                    symbol_table,
                )
                .unwrap(),
            )
        };

        // Add all files
        for (entry_name, data) in entries.into_iter() {
            let header = ar::Header::new(entry_name.into_bytes(), data.len() as u64);
            match builder {
                BuilderKind::Bsd(ref mut builder) => builder.append(&header, &mut &*data).unwrap(),
                BuilderKind::Gnu(ref mut builder) => builder.append(&header, &mut &*data).unwrap(),
            }
        }

        // Finalize archive
        std::mem::drop(builder);

        if self.no_builtin_ranlib {
            let ranlib = crate::toolchain::get_toolchain_binary(self.sess, "ranlib");

            // Run ranlib to be able to link the archive
            let status = std::process::Command::new(ranlib)
                .arg(self.dst)
                .status()
                .expect("Couldn't run ranlib");

            if !status.success() {
                self.sess.fatal(&format!("Ranlib exited with code {:?}", status.code()));
            }
        }
    }

    fn inject_dll_import_lib(
        &mut self,
        _lib_name: &str,
        _dll_imports: &[rustc_middle::middle::cstore::DllImport],
        _tmpdir: &rustc_data_structures::temp_dir::MaybeTempDir,
    ) {
        bug!("injecting dll imports is not supported");
    }
}

impl<'a> ArArchiveBuilder<'a> {
    fn add_archive<F>(&mut self, archive_path: PathBuf, mut skip: F) -> io::Result<()>
    where
        F: FnMut(&str) -> bool + 'static,
    {
        let read_cache = ReadCache::new(std::fs::File::open(&archive_path)?);
        let archive = ArchiveFile::parse(&read_cache).unwrap();
        let archive_index = self.src_archives.len();

        for entry in archive.members() {
            let entry = entry.map_err(|err| io::Error::new(io::ErrorKind::InvalidData, err))?;
            let file_name = String::from_utf8(entry.name().to_vec())
                .map_err(|err| io::Error::new(io::ErrorKind::InvalidData, err))?;
            if !skip(&file_name) {
                self.entries.push((
                    file_name,
                    ArchiveEntry::FromArchive { archive_index, file_range: entry.file_range() },
                ));
            }
        }

        self.src_archives.push(read_cache.into_inner());
        Ok(())
    }
}
<commit_msg>Avoid converting filenames into strings where possible<commit_after>//! Creation of ar archives like for the lib and staticlib crate type

use std::collections::BTreeMap;
use std::convert::TryFrom;
use std::fs::File;
use std::io::{self, Read, Seek};
use std::path::{Path, PathBuf};

use rustc_codegen_ssa::back::archive::{find_library, ArchiveBuilder};
use rustc_codegen_ssa::METADATA_FILENAME;
use rustc_session::Session;

use object::read::archive::ArchiveFile;
use object::{Object, ObjectSymbol, ReadCache, SymbolKind};

#[derive(Debug)]
enum ArchiveEntry {
    FromArchive { archive_index: usize, file_range: (u64, u64) },
    File(PathBuf),
}

pub(crate) struct ArArchiveBuilder<'a> {
    sess: &'a Session,
    dst: PathBuf,
    lib_search_paths: Vec<PathBuf>,
    use_gnu_style_archive: bool,
    no_builtin_ranlib: bool,

    src_archives: Vec<File>,
    // Don't use `HashMap` here, as the order is important. `rust.metadata.bin` must always be at
    // the end of an archive for linkers to not get confused.
    entries: Vec<(Vec<u8>, ArchiveEntry)>,
}

impl<'a> ArchiveBuilder<'a> for ArArchiveBuilder<'a> {
    fn new(sess: &'a Session, output: &Path, input: Option<&Path>) -> Self {
        use rustc_codegen_ssa::back::link::archive_search_paths;

        let (src_archives, entries) = if let Some(input) = input {
            let read_cache = ReadCache::new(File::open(input).unwrap());
            let archive = ArchiveFile::parse(&read_cache).unwrap();
            let mut entries = Vec::new();

            for entry in archive.members() {
                let entry = entry.unwrap();
                entries.push((
                    entry.name().to_vec(),
                    ArchiveEntry::FromArchive { archive_index: 0, file_range: entry.file_range() },
                ));
            }

            (vec![read_cache.into_inner()], entries)
        } else {
            (vec![], Vec::new())
        };

        ArArchiveBuilder {
            sess,
            dst: output.to_path_buf(),
            lib_search_paths: archive_search_paths(sess),
            use_gnu_style_archive: sess.target.archive_format == "gnu",
            // FIXME fix builtin ranlib on macOS
            no_builtin_ranlib: sess.target.is_like_osx,

            src_archives,
            entries,
        }
    }

    fn src_files(&mut self) -> Vec<String> {
        self.entries.iter().map(|(name, _)| String::from_utf8(name.clone()).unwrap()).collect()
    }

    fn remove_file(&mut self, name: &str) {
        let index = self
            .entries
            .iter()
            .position(|(entry_name, _)| entry_name == name.as_bytes())
            .expect("Tried to remove file not existing in src archive");
        self.entries.remove(index);
    }

    fn add_file(&mut self, file: &Path) {
        self.entries.push((
            file.file_name().unwrap().to_str().unwrap().to_string().into_bytes(),
            ArchiveEntry::File(file.to_owned()),
        ));
    }

    fn add_native_library(&mut self, name: rustc_span::symbol::Symbol, verbatim: bool) {
        let location = find_library(name, verbatim, &self.lib_search_paths, self.sess);
        self.add_archive(location.clone(), |_| false).unwrap_or_else(|e| {
            panic!("failed to add native library {}: {}", location.to_string_lossy(), e);
        });
    }

    fn add_rlib(
        &mut self,
        rlib: &Path,
        name: &str,
        lto: bool,
        skip_objects: bool,
    ) -> io::Result<()> {
        let obj_start = name.to_owned();

        self.add_archive(rlib.to_owned(), move |fname: &str| {
            // Ignore metadata files, no matter the name.
            if fname == METADATA_FILENAME {
                return true;
            }

            // Don't include Rust objects if LTO is enabled
            if lto && fname.starts_with(&obj_start) && fname.ends_with(".o") {
                return true;
            }

            // Otherwise if this is *not* a rust object and we're skipping
            // objects then skip this file
            if skip_objects && (!fname.starts_with(&obj_start) || !fname.ends_with(".o")) {
                return true;
            }

            // ok, don't skip this
            false
        })
    }

    fn update_symbols(&mut self) {}

    fn build(mut self) {
        enum BuilderKind {
            Bsd(ar::Builder<File>),
            Gnu(ar::GnuBuilder<File>),
        }

        let sess = self.sess;

        let mut symbol_table = BTreeMap::new();

        let mut entries = Vec::new();

        for (entry_name, entry) in self.entries {
            // FIXME only read the symbol table of the object files to avoid having to keep all
            // object files in memory at once, or read them twice.
            let data = match entry {
                ArchiveEntry::FromArchive { archive_index, file_range } => {
                    // FIXME read symbols from symtab
                    let src_read_cache = &mut self.src_archives[archive_index];

                    src_read_cache.seek(io::SeekFrom::Start(file_range.0)).unwrap();
                    let mut data = std::vec::from_elem(0, usize::try_from(file_range.1).unwrap());
                    src_read_cache.read_exact(&mut data).unwrap();

                    data
                }
                ArchiveEntry::File(file) => std::fs::read(file).unwrap_or_else(|err| {
                    sess.fatal(&format!(
                        "error while reading object file during archive building: {}",
                        err
                    ));
                }),
            };

            if !self.no_builtin_ranlib {
                match object::File::parse(&*data) {
                    Ok(object) => {
                        symbol_table.insert(
                            entry_name.to_vec(),
                            object
                                .symbols()
                                .filter_map(|symbol| {
                                    if symbol.is_undefined()
                                        || symbol.is_local()
                                        || symbol.kind() != SymbolKind::Data
                                            && symbol.kind() != SymbolKind::Text
                                            && symbol.kind() != SymbolKind::Tls
                                    {
                                        None
                                    } else {
                                        symbol.name().map(|name| name.as_bytes().to_vec()).ok()
                                    }
                                })
                                .collect::<Vec<_>>(),
                        );
                    }
                    Err(err) => {
                        let err = err.to_string();
                        if err == "Unknown file magic" {
                            // Not an object file; skip it.
                        } else {
                            sess.fatal(&format!(
                                "error parsing `{}` during archive creation: {}",
                                String::from_utf8_lossy(&entry_name),
                                err
                            ));
                        }
                    }
                }
            }

            entries.push((entry_name, data));
        }

        let mut builder = if self.use_gnu_style_archive {
            BuilderKind::Gnu(
                ar::GnuBuilder::new(
                    File::create(&self.dst).unwrap_or_else(|err| {
                        sess.fatal(&format!(
                            "error opening destination during archive building: {}",
                            err
                        ));
                    }),
                    entries.iter().map(|(name, _)| name.clone()).collect(),
                    ar::GnuSymbolTableFormat::Size32,
                    symbol_table,
                )
                .unwrap(),
            )
        } else {
            BuilderKind::Bsd(
                ar::Builder::new(
                    File::create(&self.dst).unwrap_or_else(|err| {
                        sess.fatal(&format!(
                            "error opening destination during archive building: {}",
                            err
                        ));
                    }),
                    symbol_table,
                )
                .unwrap(),
            )
        };

        // Add all files
        for (entry_name, data) in entries.into_iter() {
            let header = ar::Header::new(entry_name, data.len() as u64);
            match builder {
                BuilderKind::Bsd(ref mut builder) => builder.append(&header, &mut &*data).unwrap(),
                BuilderKind::Gnu(ref mut builder) => builder.append(&header, &mut &*data).unwrap(),
            }
        }

        // Finalize archive
        std::mem::drop(builder);

        if self.no_builtin_ranlib {
            let ranlib = crate::toolchain::get_toolchain_binary(self.sess, "ranlib");

            // Run ranlib to be able to link the archive
            let status = std::process::Command::new(ranlib)
                .arg(self.dst)
                .status()
                .expect("Couldn't run ranlib");

            if !status.success() {
                self.sess.fatal(&format!("Ranlib exited with code {:?}", status.code()));
            }
        }
    }

    fn inject_dll_import_lib(
        &mut self,
        _lib_name: &str,
        _dll_imports: &[rustc_middle::middle::cstore::DllImport],
        _tmpdir: &rustc_data_structures::temp_dir::MaybeTempDir,
    ) {
        bug!("injecting dll imports is not supported");
    }
}

impl<'a> ArArchiveBuilder<'a> {
    fn add_archive<F>(&mut self, archive_path: PathBuf, mut skip: F) -> io::Result<()>
    where
        F: FnMut(&str) -> bool + 'static,
    {
        let read_cache = ReadCache::new(std::fs::File::open(&archive_path)?);
        let archive = ArchiveFile::parse(&read_cache).unwrap();
        let archive_index = self.src_archives.len();

        for entry in archive.members() {
            let entry = entry.map_err(|err| io::Error::new(io::ErrorKind::InvalidData, err))?;
            let file_name = String::from_utf8(entry.name().to_vec())
                .map_err(|err| io::Error::new(io::ErrorKind::InvalidData, err))?;
            if !skip(&file_name) {
                self.entries.push((
                    file_name.into_bytes(),
                    ArchiveEntry::FromArchive { archive_index, file_range: entry.file_range() },
                ));
            }
        }

        self.src_archives.push(read_cache.into_inner());
        Ok(())
    }
}
<|endoftext|>


<commit_before>use std::{
    io::{self, Error, ErrorKind},
    mem,
    net::SocketAddr,
    os::unix::io::AsRawFd,
};

use async_trait::async_trait;
use socket2::{Domain, Protocol, SockAddr, Socket, Type};
use tokio::net::{TcpListener, TcpStream};

use crate::{
    config::RedirType,
    relay::{
        redir::{TcpListenerRedirExt, TcpStreamRedirExt},
        sys::sockaddr_to_std,
    },
};

#[async_trait]
impl TcpListenerRedirExt for TcpListener {
    async fn bind_redir(ty: RedirType, addr: &SocketAddr) -> io::Result<TcpListener> {
        match ty {
            RedirType::Redirect => {
                // REDIRECT rule doesn't need to set IP_TRANSPARENT
                TcpListener::bind(addr).await
            }
            RedirType::TProxy => {
                // TPROXY rule requires IP_TRANSPARENT
                create_redir_listener(addr)
            }
            _ => Err(Error::new(
                ErrorKind::InvalidInput,
                "not supported tcp transparent proxy type",
            )),
        }
    }
}

impl TcpStreamRedirExt for TcpStream {
    fn destination_addr(&self, ty: RedirType) -> io::Result<SocketAddr> {
        match ty {
            RedirType::Redirect => get_original_destination_addr(self),
            RedirType::TProxy => {
                // For TPROXY, uses getsockname() to retrieve original destination address
                self.local_addr()
            }
            _ => unreachable!("not supported tcp transparent proxy type"),
        }
    }
}

fn get_original_destination_addr(s: &TcpStream) -> io::Result<SocketAddr> {
    let fd = s.as_raw_fd();

    unsafe {
        let mut target_addr: libc::sockaddr_storage = mem::zeroed();
        let mut target_addr_len = mem::size_of_val(&target_addr) as libc::socklen_t;

        match s.local_addr()? {
            SocketAddr::V4(..) => {
                let ret = libc::getsockopt(
                    fd,
                    libc::SOL_IP,
                    libc::SO_ORIGINAL_DST,
                    &mut target_addr as *mut _ as *mut _,
                    &mut target_addr_len,
                );
                if ret != 0 {
                    let err = Error::last_os_error();
                    return Err(err);
                }
            }
            SocketAddr::V6(..) => {
                let ret = libc::getsockopt(
                    fd,
                    libc::SOL_IPV6,
                    libc::IP6T_SO_ORIGINAL_DST,
                    &mut target_addr as *mut _ as *mut _,
                    &mut target_addr_len,
                );

                if ret != 0 {
                    let err = Error::last_os_error();
                    return Err(err);
                }
            }
        }

        // Convert sockaddr_storage to SocketAddr
        sockaddr_to_std(&target_addr)
    }
}

fn create_redir_listener(addr: &SocketAddr) -> io::Result<TcpListener> {
    let domain = match *addr {
        SocketAddr::V4(..) => Domain::ipv4(),
        SocketAddr::V6(..) => Domain::ipv6(),
    };

    let socket = Socket::new(domain, Type::stream(), Some(Protocol::tcp()))?;

    // For Linux 2.4+ TPROXY
    // Sockets have to set IP_TRANSPARENT for retrieving original destination by getsockname()
    unsafe {
        let fd = socket.as_raw_fd();

        let enable: libc::c_int = 1;
        let ret = match *addr {
            SocketAddr::V4(..) => libc::setsockopt(
                fd,
                libc::SOL_IP,
                libc::IP_TRANSPARENT,
                &enable as *const _ as *const _,
                mem::size_of_val(&enable) as libc::socklen_t,
            ),
            SocketAddr::V6(..) => libc::setsockopt(
                fd,
                libc::SOL_IPV6,
                libc::IP_TRANSPARENT,
                &enable as *const _ as *const _,
                mem::size_of_val(&enable) as libc::socklen_t,
            ),
        };

        if ret != 0 {
            return Err(Error::last_os_error());
        }
    }

    // tokio requires non-blocked socket, and allow reuse addr
    socket.set_nonblocking(true)?;
    socket.set_reuse_address(true)?;

    let addr = SockAddr::from(*addr);

    // bind, listen as original
    socket.bind(&addr)?;
    socket.listen(1024)?; // backlogs = 1024 as mio's default

    TcpListener::from_std(socket.into_tcp_listener())
}
<commit_msg>Fixed IPv6 TProxy listener option: IPV6_TRANSPARENT<commit_after>use std::{
    io::{self, Error, ErrorKind},
    mem,
    net::SocketAddr,
    os::unix::io::AsRawFd,
};

use async_trait::async_trait;
use socket2::{Domain, Protocol, SockAddr, Socket, Type};
use tokio::net::{TcpListener, TcpStream};

use crate::{
    config::RedirType,
    relay::{
        redir::{TcpListenerRedirExt, TcpStreamRedirExt},
        sys::sockaddr_to_std,
    },
};

#[async_trait]
impl TcpListenerRedirExt for TcpListener {
    async fn bind_redir(ty: RedirType, addr: &SocketAddr) -> io::Result<TcpListener> {
        match ty {
            RedirType::Redirect => {
                // REDIRECT rule doesn't need to set IP_TRANSPARENT
                TcpListener::bind(addr).await
            }
            RedirType::TProxy => {
                // TPROXY rule requires IP_TRANSPARENT
                create_redir_listener(addr)
            }
            _ => Err(Error::new(
                ErrorKind::InvalidInput,
                "not supported tcp transparent proxy type",
            )),
        }
    }
}

impl TcpStreamRedirExt for TcpStream {
    fn destination_addr(&self, ty: RedirType) -> io::Result<SocketAddr> {
        match ty {
            RedirType::Redirect => get_original_destination_addr(self),
            RedirType::TProxy => {
                // For TPROXY, uses getsockname() to retrieve original destination address
                self.local_addr()
            }
            _ => unreachable!("not supported tcp transparent proxy type"),
        }
    }
}

fn get_original_destination_addr(s: &TcpStream) -> io::Result<SocketAddr> {
    let fd = s.as_raw_fd();

    unsafe {
        let mut target_addr: libc::sockaddr_storage = mem::zeroed();
        let mut target_addr_len = mem::size_of_val(&target_addr) as libc::socklen_t;

        match s.local_addr()? {
            SocketAddr::V4(..) => {
                let ret = libc::getsockopt(
                    fd,
                    libc::SOL_IP,
                    libc::SO_ORIGINAL_DST,
                    &mut target_addr as *mut _ as *mut _,
                    &mut target_addr_len,
                );
                if ret != 0 {
                    let err = Error::last_os_error();
                    return Err(err);
                }
            }
            SocketAddr::V6(..) => {
                let ret = libc::getsockopt(
                    fd,
                    libc::SOL_IPV6,
                    libc::IP6T_SO_ORIGINAL_DST,
                    &mut target_addr as *mut _ as *mut _,
                    &mut target_addr_len,
                );

                if ret != 0 {
                    let err = Error::last_os_error();
                    return Err(err);
                }
            }
        }

        // Convert sockaddr_storage to SocketAddr
        sockaddr_to_std(&target_addr)
    }
}

fn create_redir_listener(addr: &SocketAddr) -> io::Result<TcpListener> {
    let domain = match *addr {
        SocketAddr::V4(..) => Domain::ipv4(),
        SocketAddr::V6(..) => Domain::ipv6(),
    };

    let socket = Socket::new(domain, Type::stream(), Some(Protocol::tcp()))?;

    // For Linux 2.4+ TPROXY
    // Sockets have to set IP_TRANSPARENT for retrieving original destination by getsockname()
    unsafe {
        let fd = socket.as_raw_fd();

        let enable: libc::c_int = 1;
        let ret = match *addr {
            SocketAddr::V4(..) => libc::setsockopt(
                fd,
                libc::IPPROTO_IP,
                libc::IP_TRANSPARENT,
                &enable as *const _ as *const _,
                mem::size_of_val(&enable) as libc::socklen_t,
            ),
            SocketAddr::V6(..) => libc::setsockopt(
                fd,
                libc::IPPROTO_IPV6,
                libc::IPV6_TRANSPARENT,
                &enable as *const _ as *const _,
                mem::size_of_val(&enable) as libc::socklen_t,
            ),
        };

        if ret != 0 {
            return Err(Error::last_os_error());
        }
    }

    // tokio requires non-blocked socket, and allow reuse addr
    socket.set_nonblocking(true)?;
    socket.set_reuse_address(true)?;

    let addr = SockAddr::from(*addr);

    // bind, listen as original
    socket.bind(&addr)?;
    socket.listen(1024)?; // backlogs = 1024 as mio's default

    TcpListener::from_std(socket.into_tcp_listener())
}
<|endoftext|>


<commit_before>use std::env;

use cargo::core::Workspace;
use cargo::ops::{self, MessageFormat, Packages};
use cargo::util::{CliResult, Config};
use cargo::util::important_paths::{find_root_manifest_for_wd};

#[derive(Deserialize)]
pub struct Options {
    flag_target: Option<String>,
    flag_features: Vec<String>,
    flag_all_features: bool,
    flag_jobs: Option<u32>,
    flag_manifest_path: Option<String>,
    flag_no_default_features: bool,
    flag_no_deps: bool,
    flag_open: bool,
    flag_release: bool,
    flag_verbose: u32,
    flag_quiet: Option<bool>,
    flag_color: Option<String>,
    flag_message_format: MessageFormat,
    flag_package: Vec<String>,
    flag_lib: bool,
    flag_bin: Vec<String>,
    flag_bins: bool,
    flag_frozen: bool,
    flag_locked: bool,
    flag_all: bool,
    #[serde(rename = "flag_Z")]
    flag_z: Vec<String>,
}

pub const USAGE: &'static str = "
Build a package's documentation

Usage:
    cargo doc [options]

Options:
    -h, --help                   Print this message
    --open                       Opens the docs in a browser after the operation
    -p SPEC, --package SPEC ...  Package to document
    --all                        Document all packages in the workspace
    --no-deps                    Don't build documentation for dependencies
    -j N, --jobs N               Number of parallel jobs, defaults to # of CPUs
    --lib                        Document only this package's library
    --bin NAME                   Document only the specified binary
    --bins                       Document all binaries
    --release                    Build artifacts in release mode, with optimizations
    --features FEATURES          Space-separated list of features to also build
    --all-features               Build all available features
    --no-default-features        Do not build the `default` feature
    --target TRIPLE              Build for the target triple
    --manifest-path PATH         Path to the manifest to document
    -v, --verbose ...            Use verbose output (-vv very verbose/build.rs output)
    -q, --quiet                  No output printed to stdout
    --color WHEN                 Coloring: auto, always, never
    --message-format FMT         Error format: human, json [default: human]
    --frozen                     Require Cargo.lock and cache are up to date
    --locked                     Require Cargo.lock is up to date
    -Z FLAG ...                  Unstable (nightly-only) flags to Cargo

By default the documentation for the local package and all dependencies is
built. The output is all placed in `target/doc` in rustdoc's usual format.

All packages in the workspace are documented if the `--all` flag is supplied. The
`--all` flag is automatically assumed for a virtual manifest.
Note that `--exclude` has to be specified in conjunction with the `--all` flag.

If the --package argument is given, then SPEC is a package id specification
which indicates which package should be documented. If it is not given, then the
current package is documented. For more information on SPEC and its format, see
the `cargo help pkgid` command.
";

pub fn execute(options: Options, config: &mut Config) -> CliResult {
    debug!("executing; cmd=cargo-check; args={:?}",
           env::args().collect::<Vec<_>>());

    config.configure(options.flag_verbose,
                     options.flag_quiet,
                     &options.flag_color,
                     options.flag_frozen,
                     options.flag_locked,
                     &options.flag_z)?;

    let root = find_root_manifest_for_wd(options.flag_manifest_path, config.cwd())?;
    let ws = Workspace::new(&root, config)?;

    let spec = if options.flag_all || (ws.is_virtual() && options.flag_package.is_empty()) {
        Packages::All
    } else {
        Packages::Packages(&options.flag_package)
    };

    let empty = Vec::new();
    let doc_opts = ops::DocOptions {
        open_result: options.flag_open,
        compile_opts: ops::CompileOptions {
            config: config,
            jobs: options.flag_jobs,
            target: options.flag_target.as_ref().map(|t| &t[..]),
            features: &options.flag_features,
            all_features: options.flag_all_features,
            no_default_features: options.flag_no_default_features,
            spec: spec,
            filter: ops::CompileFilter::new(options.flag_lib,
                                            &options.flag_bin, options.flag_bins,
                                            &empty, false,
                                            &empty, false,
                                            &empty, false,
                                            false),
            message_format: options.flag_message_format,
            release: options.flag_release,
            mode: ops::CompileMode::Doc {
                deps: !options.flag_no_deps,
            },
            target_rustc_args: None,
            target_rustdoc_args: None,
        },
    };

    ops::doc(&ws, &doc_opts)?;
    Ok(())
}
<commit_msg>Support --exclude option for `cargo doc`<commit_after>use std::env;

use cargo::core::Workspace;
use cargo::ops::{self, MessageFormat, Packages};
use cargo::util::{CliResult, Config};
use cargo::util::important_paths::{find_root_manifest_for_wd};

#[derive(Deserialize)]
pub struct Options {
    flag_target: Option<String>,
    flag_features: Vec<String>,
    flag_all_features: bool,
    flag_jobs: Option<u32>,
    flag_manifest_path: Option<String>,
    flag_no_default_features: bool,
    flag_no_deps: bool,
    flag_open: bool,
    flag_release: bool,
    flag_verbose: u32,
    flag_quiet: Option<bool>,
    flag_color: Option<String>,
    flag_message_format: MessageFormat,
    flag_package: Vec<String>,
    flag_lib: bool,
    flag_bin: Vec<String>,
    flag_bins: bool,
    flag_frozen: bool,
    flag_locked: bool,
    flag_all: bool,
    flag_exclude: Vec<String>,
    #[serde(rename = "flag_Z")]
    flag_z: Vec<String>,
}

pub const USAGE: &'static str = "
Build a package's documentation

Usage:
    cargo doc [options]

Options:
    -h, --help                   Print this message
    --open                       Opens the docs in a browser after the operation
    -p SPEC, --package SPEC ...  Package to document
    --all                        Document all packages in the workspace
    --exclude SPEC ...           Exclude packages from the build
    --no-deps                    Don't build documentation for dependencies
    -j N, --jobs N               Number of parallel jobs, defaults to # of CPUs
    --lib                        Document only this package's library
    --bin NAME                   Document only the specified binary
    --bins                       Document all binaries
    --release                    Build artifacts in release mode, with optimizations
    --features FEATURES          Space-separated list of features to also build
    --all-features               Build all available features
    --no-default-features        Do not build the `default` feature
    --target TRIPLE              Build for the target triple
    --manifest-path PATH         Path to the manifest to document
    -v, --verbose ...            Use verbose output (-vv very verbose/build.rs output)
    -q, --quiet                  No output printed to stdout
    --color WHEN                 Coloring: auto, always, never
    --message-format FMT         Error format: human, json [default: human]
    --frozen                     Require Cargo.lock and cache are up to date
    --locked                     Require Cargo.lock is up to date
    -Z FLAG ...                  Unstable (nightly-only) flags to Cargo

By default the documentation for the local package and all dependencies is
built. The output is all placed in `target/doc` in rustdoc's usual format.

All packages in the workspace are documented if the `--all` flag is supplied. The
`--all` flag is automatically assumed for a virtual manifest.
Note that `--exclude` has to be specified in conjunction with the `--all` flag.

If the --package argument is given, then SPEC is a package id specification
which indicates which package should be documented. If it is not given, then the
current package is documented. For more information on SPEC and its format, see
the `cargo help pkgid` command.
";

pub fn execute(options: Options, config: &mut Config) -> CliResult {
    debug!("executing; cmd=cargo-check; args={:?}",
           env::args().collect::<Vec<_>>());

    config.configure(options.flag_verbose,
                     options.flag_quiet,
                     &options.flag_color,
                     options.flag_frozen,
                     options.flag_locked,
                     &options.flag_z)?;

    let root = find_root_manifest_for_wd(options.flag_manifest_path, config.cwd())?;
    let ws = Workspace::new(&root, config)?;

    let spec = Packages::from_flags(options.flag_all,
                                    &options.flag_exclude,
                                    &options.flag_package)?;

    let empty = Vec::new();
    let doc_opts = ops::DocOptions {
        open_result: options.flag_open,
        compile_opts: ops::CompileOptions {
            config: config,
            jobs: options.flag_jobs,
            target: options.flag_target.as_ref().map(|t| &t[..]),
            features: &options.flag_features,
            all_features: options.flag_all_features,
            no_default_features: options.flag_no_default_features,
            spec: spec,
            filter: ops::CompileFilter::new(options.flag_lib,
                                            &options.flag_bin, options.flag_bins,
                                            &empty, false,
                                            &empty, false,
                                            &empty, false,
                                            false),
            message_format: options.flag_message_format,
            release: options.flag_release,
            mode: ops::CompileMode::Doc {
                deps: !options.flag_no_deps,
            },
            target_rustc_args: None,
            target_rustdoc_args: None,
        },
    };

    ops::doc(&ws, &doc_opts)?;
    Ok(())
}
<|endoftext|>


<commit_before>// Copyright 2016 Mozilla
//
// 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.

///! This module defines some core types that support find expressions: sources,
///! variables, expressions, etc.
///! These are produced as 'fuel' by the query parser, consumed by the query
///! translator and executor.
///!
///! Many of these types are defined as simple structs that are little more than
///! a richer type alias: a variable, for example, is really just a fancy kind
///! of string.
///!
///! At some point in the future, we might consider reducing copying and memory
///! usage by recasting all of these string-holding structs and enums in terms
///! of string references, with those references being slices of some parsed
///! input query string, and valid for the lifetime of that string.
///!
///! For now, for the sake of simplicity, all of these strings are heap-allocated.
///!
///! Furthermore, we might cut out some of the chaff here: each time a 'tagged'
///! type is used within an enum, we have an opportunity to simplify and use the
///! inner type directly in conjunction with matching on the enum. Before diving
///! deeply into this it's worth recognizing that this loss of 'sovereignty' is
///! a tradeoff against well-typed function signatures and other such boundaries.

extern crate edn;

use edn::{BigInt, OrderedFloat, NamespacedKeyword, PlainSymbol};

pub type SrcVarName = String;          // Do not include the required syntactic '$'.

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Variable(pub PlainSymbol);

pub trait FromValue<T> {
    fn from_value(v: &edn::Value) -> Option<T>;
}

/// If the provided EDN value is a PlainSymbol beginning with '?', return
/// it wrapped in a Variable. If not, return None.
impl FromValue<Variable> for Variable {
    fn from_value(v: &edn::Value) -> Option<Variable> {
        if let edn::Value::PlainSymbol(ref s) = *v {
            Variable::from_symbol(s)
        } else {
            None
        }
    }
}

impl Variable {
    pub fn from_symbol(sym: &PlainSymbol) -> Option<Variable> {
        if sym.is_var_symbol() {
            Some(Variable(sym.clone()))
        } else {
            None
        }
    }
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub enum SrcVar {
    DefaultSrc,
    NamedSrc(SrcVarName),
}

impl FromValue<SrcVar> for SrcVar {
    fn from_value(v: &edn::Value) -> Option<SrcVar> {
        if let edn::Value::PlainSymbol(ref s) = *v {
            SrcVar::from_symbol(s)
        } else {
            None
        }
    }
}

impl SrcVar {
    pub fn from_symbol(sym: &PlainSymbol) -> Option<SrcVar> {
        if sym.is_src_symbol() {
            Some(SrcVar::NamedSrc(sym.plain_name().to_string()))
        } else {
            None
        }
    }
}

/// These are the scalar values representable in EDN.
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum NonIntegerConstant {
    Boolean(bool),
    BigInteger(BigInt),
    Float(OrderedFloat<f64>),
    Text(String),
}

pub enum FnArg {
    Variable(Variable),
    SrcVar(SrcVar),
    EntidOrInteger(i64),
    Ident(NamespacedKeyword),
    Constant(NonIntegerConstant),
}

/// e, a, tx can't be values -- no strings, no floats -- and so
/// they can only be variables, entity IDs, ident keywords, or
/// placeholders.
/// This encoding allows us to represent integers that aren't
/// entity IDs. That'll get filtered out in the context of the
/// database.
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum PatternNonValuePlace {
    Placeholder,
    Variable(Variable),
    Entid(i64),                       // Will always be +ve. See #190.
    Ident(NamespacedKeyword),
}

impl PatternNonValuePlace {
    // I think we'll want move variants, so let's leave these here for now.
    #[allow(dead_code)]
    fn into_pattern_value_place(self) -> PatternValuePlace {
        match self {
            PatternNonValuePlace::Placeholder => PatternValuePlace::Placeholder,
            PatternNonValuePlace::Variable(x) => PatternValuePlace::Variable(x),
            PatternNonValuePlace::Entid(x)    => PatternValuePlace::EntidOrInteger(x),
            PatternNonValuePlace::Ident(x)    => PatternValuePlace::IdentOrKeyword(x),
        }
    }

    fn to_pattern_value_place(&self) -> PatternValuePlace {
        match *self {
            PatternNonValuePlace::Placeholder     => PatternValuePlace::Placeholder,
            PatternNonValuePlace::Variable(ref x) => PatternValuePlace::Variable(x.clone()),
            PatternNonValuePlace::Entid(x)        => PatternValuePlace::EntidOrInteger(x),
            PatternNonValuePlace::Ident(ref x)    => PatternValuePlace::IdentOrKeyword(x.clone()),
        }
    }
}

impl FromValue<PatternNonValuePlace> for PatternNonValuePlace {
    fn from_value(v: &edn::Value) -> Option<PatternNonValuePlace> {
        match v {
            &edn::Value::Integer(x) => if x >= 0 {
                Some(PatternNonValuePlace::Entid(x))
            } else {
                None
            },
            &edn::Value::PlainSymbol(ref x) => if x.0.as_str() == "_" {
                Some(PatternNonValuePlace::Placeholder)
            } else {
                if let Some(v) = Variable::from_symbol(x) {
                    Some(PatternNonValuePlace::Variable(v))
                } else {
                    None
                }
            },
            &edn::Value::NamespacedKeyword(ref x) =>
                Some(PatternNonValuePlace::Ident(x.clone())),
            _ => None,
        }
    }
}

/// The `v` part of a pattern can be much broader: it can represent
/// integers that aren't entity IDs (particularly negative integers),
/// strings, and all the rest. We group those under `Constant`.
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum PatternValuePlace {
    Placeholder,
    Variable(Variable),
    EntidOrInteger(i64),
    IdentOrKeyword(NamespacedKeyword),
    Constant(NonIntegerConstant),
}

impl FromValue<PatternValuePlace> for PatternValuePlace {
    fn from_value(v: &edn::Value) -> Option<PatternValuePlace> {
        match v {
            &edn::Value::Integer(x) =>
                Some(PatternValuePlace::EntidOrInteger(x)),
            &edn::Value::PlainSymbol(ref x) if x.0.as_str() == "_" =>
                Some(PatternValuePlace::Placeholder),
            &edn::Value::PlainSymbol(ref x) if x.is_var_symbol() =>
                Some(PatternValuePlace::Variable(Variable(x.clone()))),
            &edn::Value::NamespacedKeyword(ref x) =>
                Some(PatternValuePlace::IdentOrKeyword(x.clone())),
            &edn::Value::Boolean(x) =>
                Some(PatternValuePlace::Constant(NonIntegerConstant::Boolean(x))),
            &edn::Value::Float(x) =>
                Some(PatternValuePlace::Constant(NonIntegerConstant::Float(x))),
            &edn::Value::BigInteger(ref x) =>
                Some(PatternValuePlace::Constant(NonIntegerConstant::BigInteger(x.clone()))),
            &edn::Value::Text(ref x) =>
                Some(PatternValuePlace::Constant(NonIntegerConstant::Text(x.clone()))),
            _ => None,
        }
    }
}

impl PatternValuePlace {
    // I think we'll want move variants, so let's leave these here for now.
    #[allow(dead_code)]
    fn into_pattern_non_value_place(self) -> Option<PatternNonValuePlace> {
        match self {
            PatternValuePlace::Placeholder       => Some(PatternNonValuePlace::Placeholder),
            PatternValuePlace::Variable(x)       => Some(PatternNonValuePlace::Variable(x)),
            PatternValuePlace::EntidOrInteger(x) => if x >= 0 {
                Some(PatternNonValuePlace::Entid(x))
            } else {
                None
            },
            PatternValuePlace::IdentOrKeyword(x) => Some(PatternNonValuePlace::Ident(x)),
            PatternValuePlace::Constant(_)       => None,
        }
    }

    fn to_pattern_non_value_place(&self) -> Option<PatternNonValuePlace> {
        match *self {
            PatternValuePlace::Placeholder           => Some(PatternNonValuePlace::Placeholder),
            PatternValuePlace::Variable(ref x)       => Some(PatternNonValuePlace::Variable(x.clone())),
            PatternValuePlace::EntidOrInteger(x)     => if x >= 0 {
                Some(PatternNonValuePlace::Entid(x))
            } else {
                None
            },
            PatternValuePlace::IdentOrKeyword(ref x) => Some(PatternNonValuePlace::Ident(x.clone())),
            PatternValuePlace::Constant(_)           => None,
        }
    }
}

/*
pub enum PullPattern {
    Constant(Constant),
    Variable(Variable),
}

pub struct Pull {
    pub src: SrcVar,
    pub var: Variable,
    pub pattern: PullPattern,      // Constant, variable, or plain variable.
}
*/

/*
pub struct Aggregate {
    pub fn_name: String,
    pub args: Vec<FnArg>,
}
*/

#[derive(Clone,Debug,Eq,PartialEq)]
pub enum Element {
    Variable(Variable),
    // Aggregate(Aggregate),   // TODO
    // Pull(Pull),             // TODO
}

/// A definition of the first part of a find query: the
/// `[:find ?foo ?bar…]` bit.
///
/// There are four different kinds of find specs, allowing you to query for
/// a single value, a collection of values from different entities, a single
/// tuple (relation), or a collection of tuples.
///
/// Examples:
///
/// ```rust
/// # extern crate edn;
/// # extern crate mentat_query;
/// # use edn::PlainSymbol;
/// # use mentat_query::{Element, FindSpec, Variable};
///
/// # fn main() {
///
///   let elements = vec![
///     Element::Variable(Variable(PlainSymbol("?foo".to_string()))),
///     Element::Variable(Variable(PlainSymbol("?bar".to_string()))),
///   ];
///   let rel = FindSpec::FindRel(elements);
///
///   if let FindSpec::FindRel(elements) = rel {
///     assert_eq!(2, elements.len());
///   }
///
/// # }
/// ```
///
#[derive(Clone,Debug,Eq,PartialEq)]
pub enum FindSpec {
    /// Returns an array of arrays.
    FindRel(Vec<Element>),

    /// Returns an array of scalars, usually homogeneous.
    /// This is equivalent to mapping over the results of a `FindRel`,
    /// returning the first value of each.
    FindColl(Element),

    /// Returns a single tuple: a heterogeneous array of scalars. Equivalent to
    /// taking the first result from a `FindRel`.
    FindTuple(Vec<Element>),

    /// Returns a single scalar value. Equivalent to taking the first result
    /// from a `FindColl`.
    FindScalar(Element),
}

/// Returns true if the provided `FindSpec` returns at most one result.
pub fn is_unit_limited(spec: &FindSpec) -> bool {
    match spec {
        &FindSpec::FindScalar(..) => true,
        &FindSpec::FindTuple(..)  => true,
        &FindSpec::FindRel(..)    => false,
        &FindSpec::FindColl(..)   => false,
    }
}

/// Returns true if the provided `FindSpec` cares about distinct results.
///
/// I use the words "cares about" because find is generally defined in terms of producing distinct
/// results at the Datalog level.
///
/// Two of the find specs (scalar and tuple) produce only a single result. Those don't need to be
/// run with `SELECT DISTINCT`, because we're only consuming a single result. Those queries will be
/// run with `LIMIT 1`.
///
/// Additionally, some projections cannot produce duplicate results: `[:find (max ?x) …]`, for
/// example.
///
/// This function gives us the hook to add that logic when we're ready.
///
/// Beyond this, `DISTINCT` is not always needed. For example, in some kinds of accumulation or
/// sampling projections we might not need to do it at the SQL level because we're consuming into
/// a dupe-eliminating data structure like a Set, or we know that a particular query cannot produce
/// duplicate results.
pub fn requires_distinct(spec: &FindSpec) -> bool {
    return !is_unit_limited(spec);
}

// Note that the "implicit blank" rule applies.
// A pattern with a reversed attribute — :foo/_bar — is reversed
// at the point of parsing. These `Pattern` instances only represent
// one direction.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct Pattern {
    pub source: Option<SrcVar>,
    pub entity: PatternNonValuePlace,
    pub attribute: PatternNonValuePlace,
    pub value: PatternValuePlace,
    pub tx: PatternNonValuePlace,
}

impl Pattern {
    pub fn new(src: Option<SrcVar>,
               e: PatternNonValuePlace,
               a: PatternNonValuePlace,
               v: PatternValuePlace,
               tx: PatternNonValuePlace) -> Option<Pattern> {
        let aa = a.clone();       // Too tired of fighting borrow scope for now.
        if let PatternNonValuePlace::Ident(ref k) = aa {
            if k.is_backward() {
                // e and v have different types; we must convert them.
                // Not every parseable value is suitable for the entity field!
                // As such, this is a failable constructor.
                let e_v = e.to_pattern_value_place();
                if let Some(v_e) = v.to_pattern_non_value_place() {
                    return Some(Pattern {
                        source: src,
                        entity: v_e,
                        attribute: PatternNonValuePlace::Ident(k.to_reversed()),
                        value: e_v,
                        tx: tx,
                    });
                } else {
                    return None;
                }
            }
        }
        Some(Pattern {
            source: src,
            entity: e,
            attribute: a,
            value: v,
            tx: tx,
        })
    }
}

#[allow(dead_code)]
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum WhereClause {
    Not,
    NotJoin,
    Or,
    OrJoin,
    Pred,
    WhereFn,
    RuleExpr,
    Pattern(Pattern),
}

#[allow(dead_code)]
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct FindQuery {
    pub find_spec: FindSpec,
    pub default_source: SrcVar,
    pub with: Vec<Variable>,
    pub in_vars: Vec<Variable>,
    pub in_sources: Vec<SrcVar>,
    pub where_clauses: Vec<WhereClause>,
    // TODO: in_rules;
}
<commit_msg>Re-export NamespacedKeyword and PlainSymbol out of query. r=jsantell<commit_after>// Copyright 2016 Mozilla
//
// 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.

///! This module defines some core types that support find expressions: sources,
///! variables, expressions, etc.
///! These are produced as 'fuel' by the query parser, consumed by the query
///! translator and executor.
///!
///! Many of these types are defined as simple structs that are little more than
///! a richer type alias: a variable, for example, is really just a fancy kind
///! of string.
///!
///! At some point in the future, we might consider reducing copying and memory
///! usage by recasting all of these string-holding structs and enums in terms
///! of string references, with those references being slices of some parsed
///! input query string, and valid for the lifetime of that string.
///!
///! For now, for the sake of simplicity, all of these strings are heap-allocated.
///!
///! Furthermore, we might cut out some of the chaff here: each time a 'tagged'
///! type is used within an enum, we have an opportunity to simplify and use the
///! inner type directly in conjunction with matching on the enum. Before diving
///! deeply into this it's worth recognizing that this loss of 'sovereignty' is
///! a tradeoff against well-typed function signatures and other such boundaries.

extern crate edn;

use edn::{BigInt, OrderedFloat};
pub use edn::{NamespacedKeyword, PlainSymbol};

pub type SrcVarName = String;          // Do not include the required syntactic '$'.

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Variable(pub PlainSymbol);

pub trait FromValue<T> {
    fn from_value(v: &edn::Value) -> Option<T>;
}

/// If the provided EDN value is a PlainSymbol beginning with '?', return
/// it wrapped in a Variable. If not, return None.
impl FromValue<Variable> for Variable {
    fn from_value(v: &edn::Value) -> Option<Variable> {
        if let edn::Value::PlainSymbol(ref s) = *v {
            Variable::from_symbol(s)
        } else {
            None
        }
    }
}

impl Variable {
    pub fn from_symbol(sym: &PlainSymbol) -> Option<Variable> {
        if sym.is_var_symbol() {
            Some(Variable(sym.clone()))
        } else {
            None
        }
    }
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub enum SrcVar {
    DefaultSrc,
    NamedSrc(SrcVarName),
}

impl FromValue<SrcVar> for SrcVar {
    fn from_value(v: &edn::Value) -> Option<SrcVar> {
        if let edn::Value::PlainSymbol(ref s) = *v {
            SrcVar::from_symbol(s)
        } else {
            None
        }
    }
}

impl SrcVar {
    pub fn from_symbol(sym: &PlainSymbol) -> Option<SrcVar> {
        if sym.is_src_symbol() {
            Some(SrcVar::NamedSrc(sym.plain_name().to_string()))
        } else {
            None
        }
    }
}

/// These are the scalar values representable in EDN.
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum NonIntegerConstant {
    Boolean(bool),
    BigInteger(BigInt),
    Float(OrderedFloat<f64>),
    Text(String),
}

pub enum FnArg {
    Variable(Variable),
    SrcVar(SrcVar),
    EntidOrInteger(i64),
    Ident(NamespacedKeyword),
    Constant(NonIntegerConstant),
}

/// e, a, tx can't be values -- no strings, no floats -- and so
/// they can only be variables, entity IDs, ident keywords, or
/// placeholders.
/// This encoding allows us to represent integers that aren't
/// entity IDs. That'll get filtered out in the context of the
/// database.
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum PatternNonValuePlace {
    Placeholder,
    Variable(Variable),
    Entid(i64),                       // Will always be +ve. See #190.
    Ident(NamespacedKeyword),
}

impl PatternNonValuePlace {
    // I think we'll want move variants, so let's leave these here for now.
    #[allow(dead_code)]
    fn into_pattern_value_place(self) -> PatternValuePlace {
        match self {
            PatternNonValuePlace::Placeholder => PatternValuePlace::Placeholder,
            PatternNonValuePlace::Variable(x) => PatternValuePlace::Variable(x),
            PatternNonValuePlace::Entid(x)    => PatternValuePlace::EntidOrInteger(x),
            PatternNonValuePlace::Ident(x)    => PatternValuePlace::IdentOrKeyword(x),
        }
    }

    fn to_pattern_value_place(&self) -> PatternValuePlace {
        match *self {
            PatternNonValuePlace::Placeholder     => PatternValuePlace::Placeholder,
            PatternNonValuePlace::Variable(ref x) => PatternValuePlace::Variable(x.clone()),
            PatternNonValuePlace::Entid(x)        => PatternValuePlace::EntidOrInteger(x),
            PatternNonValuePlace::Ident(ref x)    => PatternValuePlace::IdentOrKeyword(x.clone()),
        }
    }
}

impl FromValue<PatternNonValuePlace> for PatternNonValuePlace {
    fn from_value(v: &edn::Value) -> Option<PatternNonValuePlace> {
        match v {
            &edn::Value::Integer(x) => if x >= 0 {
                Some(PatternNonValuePlace::Entid(x))
            } else {
                None
            },
            &edn::Value::PlainSymbol(ref x) => if x.0.as_str() == "_" {
                Some(PatternNonValuePlace::Placeholder)
            } else {
                if let Some(v) = Variable::from_symbol(x) {
                    Some(PatternNonValuePlace::Variable(v))
                } else {
                    None
                }
            },
            &edn::Value::NamespacedKeyword(ref x) =>
                Some(PatternNonValuePlace::Ident(x.clone())),
            _ => None,
        }
    }
}

/// The `v` part of a pattern can be much broader: it can represent
/// integers that aren't entity IDs (particularly negative integers),
/// strings, and all the rest. We group those under `Constant`.
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum PatternValuePlace {
    Placeholder,
    Variable(Variable),
    EntidOrInteger(i64),
    IdentOrKeyword(NamespacedKeyword),
    Constant(NonIntegerConstant),
}

impl FromValue<PatternValuePlace> for PatternValuePlace {
    fn from_value(v: &edn::Value) -> Option<PatternValuePlace> {
        match v {
            &edn::Value::Integer(x) =>
                Some(PatternValuePlace::EntidOrInteger(x)),
            &edn::Value::PlainSymbol(ref x) if x.0.as_str() == "_" =>
                Some(PatternValuePlace::Placeholder),
            &edn::Value::PlainSymbol(ref x) if x.is_var_symbol() =>
                Some(PatternValuePlace::Variable(Variable(x.clone()))),
            &edn::Value::NamespacedKeyword(ref x) =>
                Some(PatternValuePlace::IdentOrKeyword(x.clone())),
            &edn::Value::Boolean(x) =>
                Some(PatternValuePlace::Constant(NonIntegerConstant::Boolean(x))),
            &edn::Value::Float(x) =>
                Some(PatternValuePlace::Constant(NonIntegerConstant::Float(x))),
            &edn::Value::BigInteger(ref x) =>
                Some(PatternValuePlace::Constant(NonIntegerConstant::BigInteger(x.clone()))),
            &edn::Value::Text(ref x) =>
                Some(PatternValuePlace::Constant(NonIntegerConstant::Text(x.clone()))),
            _ => None,
        }
    }
}

impl PatternValuePlace {
    // I think we'll want move variants, so let's leave these here for now.
    #[allow(dead_code)]
    fn into_pattern_non_value_place(self) -> Option<PatternNonValuePlace> {
        match self {
            PatternValuePlace::Placeholder       => Some(PatternNonValuePlace::Placeholder),
            PatternValuePlace::Variable(x)       => Some(PatternNonValuePlace::Variable(x)),
            PatternValuePlace::EntidOrInteger(x) => if x >= 0 {
                Some(PatternNonValuePlace::Entid(x))
            } else {
                None
            },
            PatternValuePlace::IdentOrKeyword(x) => Some(PatternNonValuePlace::Ident(x)),
            PatternValuePlace::Constant(_)       => None,
        }
    }

    fn to_pattern_non_value_place(&self) -> Option<PatternNonValuePlace> {
        match *self {
            PatternValuePlace::Placeholder           => Some(PatternNonValuePlace::Placeholder),
            PatternValuePlace::Variable(ref x)       => Some(PatternNonValuePlace::Variable(x.clone())),
            PatternValuePlace::EntidOrInteger(x)     => if x >= 0 {
                Some(PatternNonValuePlace::Entid(x))
            } else {
                None
            },
            PatternValuePlace::IdentOrKeyword(ref x) => Some(PatternNonValuePlace::Ident(x.clone())),
            PatternValuePlace::Constant(_)           => None,
        }
    }
}

/*
pub enum PullPattern {
    Constant(Constant),
    Variable(Variable),
}

pub struct Pull {
    pub src: SrcVar,
    pub var: Variable,
    pub pattern: PullPattern,      // Constant, variable, or plain variable.
}
*/

/*
pub struct Aggregate {
    pub fn_name: String,
    pub args: Vec<FnArg>,
}
*/

#[derive(Clone,Debug,Eq,PartialEq)]
pub enum Element {
    Variable(Variable),
    // Aggregate(Aggregate),   // TODO
    // Pull(Pull),             // TODO
}

/// A definition of the first part of a find query: the
/// `[:find ?foo ?bar…]` bit.
///
/// There are four different kinds of find specs, allowing you to query for
/// a single value, a collection of values from different entities, a single
/// tuple (relation), or a collection of tuples.
///
/// Examples:
///
/// ```rust
/// # extern crate edn;
/// # extern crate mentat_query;
/// # use edn::PlainSymbol;
/// # use mentat_query::{Element, FindSpec, Variable};
///
/// # fn main() {
///
///   let elements = vec![
///     Element::Variable(Variable(PlainSymbol("?foo".to_string()))),
///     Element::Variable(Variable(PlainSymbol("?bar".to_string()))),
///   ];
///   let rel = FindSpec::FindRel(elements);
///
///   if let FindSpec::FindRel(elements) = rel {
///     assert_eq!(2, elements.len());
///   }
///
/// # }
/// ```
///
#[derive(Clone,Debug,Eq,PartialEq)]
pub enum FindSpec {
    /// Returns an array of arrays.
    FindRel(Vec<Element>),

    /// Returns an array of scalars, usually homogeneous.
    /// This is equivalent to mapping over the results of a `FindRel`,
    /// returning the first value of each.
    FindColl(Element),

    /// Returns a single tuple: a heterogeneous array of scalars. Equivalent to
    /// taking the first result from a `FindRel`.
    FindTuple(Vec<Element>),

    /// Returns a single scalar value. Equivalent to taking the first result
    /// from a `FindColl`.
    FindScalar(Element),
}

/// Returns true if the provided `FindSpec` returns at most one result.
pub fn is_unit_limited(spec: &FindSpec) -> bool {
    match spec {
        &FindSpec::FindScalar(..) => true,
        &FindSpec::FindTuple(..)  => true,
        &FindSpec::FindRel(..)    => false,
        &FindSpec::FindColl(..)   => false,
    }
}

/// Returns true if the provided `FindSpec` cares about distinct results.
///
/// I use the words "cares about" because find is generally defined in terms of producing distinct
/// results at the Datalog level.
///
/// Two of the find specs (scalar and tuple) produce only a single result. Those don't need to be
/// run with `SELECT DISTINCT`, because we're only consuming a single result. Those queries will be
/// run with `LIMIT 1`.
///
/// Additionally, some projections cannot produce duplicate results: `[:find (max ?x) …]`, for
/// example.
///
/// This function gives us the hook to add that logic when we're ready.
///
/// Beyond this, `DISTINCT` is not always needed. For example, in some kinds of accumulation or
/// sampling projections we might not need to do it at the SQL level because we're consuming into
/// a dupe-eliminating data structure like a Set, or we know that a particular query cannot produce
/// duplicate results.
pub fn requires_distinct(spec: &FindSpec) -> bool {
    return !is_unit_limited(spec);
}

// Note that the "implicit blank" rule applies.
// A pattern with a reversed attribute — :foo/_bar — is reversed
// at the point of parsing. These `Pattern` instances only represent
// one direction.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct Pattern {
    pub source: Option<SrcVar>,
    pub entity: PatternNonValuePlace,
    pub attribute: PatternNonValuePlace,
    pub value: PatternValuePlace,
    pub tx: PatternNonValuePlace,
}

impl Pattern {
    pub fn new(src: Option<SrcVar>,
               e: PatternNonValuePlace,
               a: PatternNonValuePlace,
               v: PatternValuePlace,
               tx: PatternNonValuePlace) -> Option<Pattern> {
        let aa = a.clone();       // Too tired of fighting borrow scope for now.
        if let PatternNonValuePlace::Ident(ref k) = aa {
            if k.is_backward() {
                // e and v have different types; we must convert them.
                // Not every parseable value is suitable for the entity field!
                // As such, this is a failable constructor.
                let e_v = e.to_pattern_value_place();
                if let Some(v_e) = v.to_pattern_non_value_place() {
                    return Some(Pattern {
                        source: src,
                        entity: v_e,
                        attribute: PatternNonValuePlace::Ident(k.to_reversed()),
                        value: e_v,
                        tx: tx,
                    });
                } else {
                    return None;
                }
            }
        }
        Some(Pattern {
            source: src,
            entity: e,
            attribute: a,
            value: v,
            tx: tx,
        })
    }
}

#[allow(dead_code)]
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum WhereClause {
    Not,
    NotJoin,
    Or,
    OrJoin,
    Pred,
    WhereFn,
    RuleExpr,
    Pattern(Pattern),
}

#[allow(dead_code)]
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct FindQuery {
    pub find_spec: FindSpec,
    pub default_source: SrcVar,
    pub with: Vec<Variable>,
    pub in_vars: Vec<Variable>,
    pub in_sources: Vec<SrcVar>,
    pub where_clauses: Vec<WhereClause>,
    // TODO: in_rules;
}
<|endoftext|>


<commit_before>#[macro_use]
extern crate log;
extern crate env_logger;
extern crate httparse;
extern crate mio;
extern crate slab;
extern crate time;

use std::ascii::AsciiExt;
use std::env;
use std::io::{self, Read, Write};
use std::fmt;
use std::mem;
use std::net::SocketAddr;
use std::slice;
use std::str;

use mio::tcp::{TcpStream, TcpListener};
use mio::{Poll, Token, EventSet, PollOpt, Events};
use slab::Slab;

const LISTENER: mio::Token = mio::Token(0);

struct Server<'a> {
    count: u64,
    listener: &'a TcpListener,
    connections: Slab<Connection, mio::Token>,
}

impl<'a> Server<'a> {
    fn new(listener: &'a TcpListener) -> Server<'a> {
        Server {
            count: 0,
            listener: listener,
            connections: Slab::new_starting_at(mio::Token(1), 1024),
        }
    }

    fn ready(&mut self,
             poll: &Poll,
             token: Token,
             events: EventSet) {
        debug!("{:?} {:?}", token, events);
        if token == LISTENER {
            while let Ok(Some(socket)) = self.listener.accept() {
                debug!("accepted");
                self.count += 1;
                // socket.0.set_nodelay(true).unwrap();
                let token = self.connections
                    .insert_with(move |_| Connection::new(socket.0))
                    .unwrap();
                poll.register(&self.connections[token].socket,
                              token,
                              EventSet::readable() | EventSet::writable(),
                              PollOpt::edge()).unwrap();
                self.try_connection(token, EventSet::all());
            }
        } else {
            self.try_connection(token, events);
        }
    }

    fn try_connection(&mut self, token: Token, events: EventSet) {
        let res = self.connections[token].ready(events);

        if res.is_err() || self.connections[token].closed {
            if let Err(ref e) = res {
                info!("error: {:?}\non: {:?} {:?}", e,
                       token, self.connections[token].socket);
            }
            debug!("removing");
            self.connections.remove(token);
        }
    }
}

struct Connection {
    socket: TcpStream,
    input: Vec<u8>,
    output: Output,
    keepalive: bool,
    closed: bool,
    read_closed: bool,
    events: EventSet,
}

struct Output {
    buf: Vec<u8>,
}

impl Connection {
    fn new(socket: TcpStream) -> Connection {
        Connection {
            socket: socket,
            input: Vec::with_capacity(2048),
            output: Output {
                buf: Vec::with_capacity(2048),
            },
            keepalive: false,
            read_closed: false,
            closed: false,
            events: EventSet::none(),
        }
    }

    fn ready(&mut self, events: EventSet) -> io::Result<()> {
        self.events = self.events | events;
        while self.events.is_readable() && !self.read_closed {
            let before = self.input.len();
            let eof = try!(read(&mut self.socket,
                                &mut self.input,
                                &mut self.events));
            if eof {
                debug!("eof");
            }
            self.read_closed = eof;
            if self.input.len() == before {
                break
            }

            while self.input.len() > 0 {
                let (req, amt) = match try!(parse(&self.input)) {
                    Some(pair) => pair,
                    None => {
                        debug!("need more data for a request");
                        return Ok(())
                    }
                };
                let request = Request {
                    inner: req,
                    amt: amt,
                    data: mem::replace(&mut self.input, Vec::new()),
                };
                debug!("got a request");
                self.keepalive = request.version() >= 1;
                self.keepalive = self.keepalive || request.headers().any(|s| {
                    s.0.eq_ignore_ascii_case("connection") &&
                        s.1.eq_ignore_ascii_case(b"keep-alive")
                });

                let response = process(&request);
                response.to_bytes(&mut self.output.buf);

                self.input = request.into_input_buf();
                if !self.keepalive {
                    debug!("disabling keepalive");
                    self.read_closed = true;
                }
            }
        }

        if self.events.is_writable() && self.output.buf.len() > 0 {
            let done = try!(write(&mut self.socket,
                                  &mut self.output,
                                  &mut self.events));
            if done {
                debug!("wrote response");
                if !self.keepalive || self.read_closed {
                    self.closed = true;
                }
            }
        }

        if self.read_closed && self.output.buf.len() == 0 {
            self.closed = true;
        }
        Ok(())
    }
}

fn process(r: &Request) -> Response {
    assert!(r.path() == "/plaintext");
    let mut r = Response::new();
    r.header("Content-Type", "text/plain; charset=UTF-8")
     .body("Hello, World!");
    return r
}

type Slice = (usize, usize);

#[allow(dead_code)]
pub struct Request {
    inner: RawRequest,
    data: Vec<u8>,
    amt: usize,
}

struct RawRequest {
    method: Slice,
    path: Slice,
    version: u8,
    headers: Vec<(Slice, Slice)>,
}

pub struct Headers<'a> {
    iter: slice::Iter<'a, (Slice, Slice)>,
    req: &'a Request,
}

impl Request {
    pub fn method(&self) -> &str {
        str::from_utf8(self.get(&self.inner.method)).unwrap()
    }

    pub fn path(&self) -> &str {
        str::from_utf8(self.get(&self.inner.path)).unwrap()
    }

    pub fn version(&self) -> u8 {
        self.inner.version
    }

    pub fn headers(&self) -> Headers {
        Headers {
            iter: self.inner.headers.iter(),
            req: self,
        }
    }

    pub fn into_input_buf(mut self) -> Vec<u8> {
        self.data.drain(..self.amt);
        self.data
    }

    fn get(&self, s: &Slice) -> &[u8] {
        &self.data[s.0..s.1]
    }
}

impl<'a> Iterator for Headers<'a> {
    type Item = (&'a str, &'a [u8]);

    fn next(&mut self) -> Option<(&'a str, &'a [u8])> {
        self.iter.next().map(|&(ref k, ref v)| {
            (str::from_utf8(self.req.get(k)).unwrap(), self.req.get(v))
        })
    }
}

fn parse(buf: &[u8]) -> io::Result<Option<(RawRequest, usize)>> {
    let mut headers = [httparse::EMPTY_HEADER; 16];
    let mut r = httparse::Request::new(&mut headers);
    let status = try!(r.parse(&buf).map_err(|_| {
        io::Error::new(io::ErrorKind::Other, "failed to parse")
    }));
    return match status {
        httparse::Status::Complete(amt) => {
            debug!("ok {:?}", String::from_utf8_lossy(&buf[..amt]));
            Ok(Some((RawRequest {
                method: slice(buf, r.method.unwrap().as_bytes()),
                path: slice(buf, r.path.unwrap().as_bytes()),
                version: r.version.unwrap(),
                headers: r.headers.iter().map(|h| {
                    (slice(buf, h.name.as_bytes()), slice(buf, &h.value))
                }).collect(),
            }, amt)))
        }
        httparse::Status::Partial => Ok(None),
    };

    fn slice(buf: &[u8], inner: &[u8]) -> Slice {
        let start = inner.as_ptr() as usize - buf.as_ptr() as usize;
        assert!(start < buf.len());
        (start, start + inner.len())
    }
}

pub struct Response {
    headers: Vec<(String, String)>,
    response: String,
}

impl Response {
    pub fn new() -> Response {
        Response {
            headers: Vec::new(),
            response: String::new(),
        }
    }

    pub fn header(&mut self, name: &str, val: &str) -> &mut Response {
        self.headers.push((name.to_string(), val.to_string()));
        self
    }

    pub fn body(&mut self, s: &str) -> &mut Response {
        self.response = s.to_string();
        self
    }

    fn to_bytes(&self, into: &mut Vec<u8>) {
        use std::fmt::Write;

        write!(FastWrite(into), "\
            HTTP/1.1 200 OK\r\n\
            Server: Example\r\n\
            Date: {}\r\n\
            Content-Length: {}\r\n\
        ", time::now().rfc822(), self.response.len()).unwrap();
        for &(ref k, ref v) in &self.headers {
            extend(into, k.as_bytes());
            extend(into, b": ");
            extend(into, v.as_bytes());
            extend(into, b"\r\n");
        }
        extend(into, b"\r\n");
        extend(into, self.response.as_bytes());
    }
}

// TODO: impl fmt::Write for Vec<u8>
//
// Right now `write!` on `Vec<u8>` goes through io::Write and is not super
// speedy, so inline a less-crufty implementation here which doesn't go through
// io::Error.
struct FastWrite<'a>(&'a mut Vec<u8>);

impl<'a> fmt::Write for FastWrite<'a> {
    fn write_str(&mut self, s: &str) -> fmt::Result {
        extend(self.0, s.as_bytes());
        Ok(())
    }
}

// TODO: why does extend_from_slice not optimize?
fn extend(dst: &mut Vec<u8>, data: &[u8]) {
    use std::ptr;
    dst.reserve(data.len());
    let prev = dst.len();
    unsafe {
        ptr::copy_nonoverlapping(data.as_ptr(),
                                 dst.as_mut_ptr().offset(prev as isize),
                                 data.len());
        dst.set_len(prev + data.len());
    }
}

fn read(socket: &mut TcpStream,
        input: &mut Vec<u8>,
        events: &mut EventSet) -> io::Result<bool> {
    match socket.read(unsafe { slice_to_end(input) }) {
        Ok(0) => return Ok(true),
        Ok(n) => {
            let len = input.len();
            unsafe { input.set_len(len + n); }
            return Ok(false)
        }
        Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
            *events = *events & !EventSet::readable();
            return Ok(false)
        }
        Err(e) => return Err(e),
    }

    unsafe fn slice_to_end(v: &mut Vec<u8>) -> &mut [u8] {
        use std::slice;
        if v.capacity() == 0 {
            v.reserve(16);
        }
        if v.capacity() == v.len() {
            v.reserve(1);
        }
        slice::from_raw_parts_mut(v.as_mut_ptr().offset(v.len() as isize),
                                  v.capacity() - v.len())
    }
}

fn write(socket: &mut TcpStream,
         output: &mut Output,
         events: &mut EventSet) -> io::Result<bool> {
    assert!(output.buf.len() > 0);
    loop {
        match socket.write(&output.buf) {
            Ok(0) => {
                return Err(io::Error::new(io::ErrorKind::Other, "early eof2"))
            }
            Ok(n) => {
                output.buf.drain(..n);
                if output.buf.len() == 0 {
                    return Ok(true)
                }
            }
            Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
                *events = *events & !EventSet::writable();
                return Ok(false)
            }
            Err(e) => return Err(e),
        }
    }
}

fn main() {
    env_logger::init().unwrap();

    let poll = mio::Poll::new().unwrap();

    let addr = env::args().nth(1).unwrap_or("127.0.0.1:8080".to_string());
    let addr = addr.parse::<SocketAddr>().unwrap();
    let listener = TcpListener::bind(&addr).unwrap();
    poll.register(&listener,
                  LISTENER,
                  EventSet::readable(),
                  PollOpt::level()).unwrap();

    let mut events = Events::new();
    let mut server = Server::new(&listener);
    loop {
        poll.poll(&mut events, None).unwrap();

        for i in 0..events.len() {
            let event = events.get(i).unwrap();
            server.ready(&poll, event.token(), event.kind());
        }
    }
}
<commit_msg>Add catch_unwind to mio-singlethread<commit_after>#[macro_use]
extern crate log;
extern crate env_logger;
extern crate httparse;
extern crate mio;
extern crate slab;
extern crate time;

use std::ascii::AsciiExt;
use std::env;
use std::fmt;
use std::io::{self, Read, Write};
use std::mem;
use std::net::SocketAddr;
use std::panic;
use std::slice;
use std::str;

use mio::tcp::{TcpStream, TcpListener};
use mio::{Poll, Token, EventSet, PollOpt, Events};
use slab::Slab;

const LISTENER: mio::Token = mio::Token(0);

struct Server<'a> {
    count: u64,
    listener: &'a TcpListener,
    connections: Slab<Connection, mio::Token>,
}

impl<'a> Server<'a> {
    fn new(listener: &'a TcpListener) -> Server<'a> {
        Server {
            count: 0,
            listener: listener,
            connections: Slab::new_starting_at(mio::Token(1), 1024),
        }
    }

    fn ready(&mut self,
             poll: &Poll,
             token: Token,
             events: EventSet) {
        debug!("{:?} {:?}", token, events);
        if token == LISTENER {
            while let Ok(Some(socket)) = self.listener.accept() {
                debug!("accepted");
                self.count += 1;
                // socket.0.set_nodelay(true).unwrap();
                let token = self.connections
                    .insert_with(move |_| Connection::new(socket.0))
                    .unwrap();
                poll.register(&self.connections[token].socket,
                              token,
                              EventSet::readable() | EventSet::writable(),
                              PollOpt::edge()).unwrap();
                self.try_connection(token, EventSet::all());
            }
        } else {
            self.try_connection(token, events);
        }
    }

    fn try_connection(&mut self, token: Token, events: EventSet) {
        // Simulate a `catch_unwind` that a real server would do anyway
        let res = panic::catch_unwind(panic::AssertUnwindSafe(|| {
            self.connections[token].ready(events)
        })).expect("oh no it panicked!");

        if res.is_err() || self.connections[token].closed {
            if let Err(ref e) = res {
                info!("error: {:?}\non: {:?} {:?}", e,
                       token, self.connections[token].socket);
            }
            debug!("removing");
            self.connections.remove(token);
        }
    }
}

struct Connection {
    socket: TcpStream,
    input: Vec<u8>,
    output: Output,
    keepalive: bool,
    closed: bool,
    read_closed: bool,
    events: EventSet,
}

struct Output {
    buf: Vec<u8>,
}

impl Connection {
    fn new(socket: TcpStream) -> Connection {
        Connection {
            socket: socket,
            input: Vec::with_capacity(2048),
            output: Output {
                buf: Vec::with_capacity(2048),
            },
            keepalive: false,
            read_closed: false,
            closed: false,
            events: EventSet::none(),
        }
    }

    fn ready(&mut self, events: EventSet) -> io::Result<()> {
        self.events = self.events | events;
        while self.events.is_readable() && !self.read_closed {
            let before = self.input.len();
            let eof = try!(read(&mut self.socket,
                                &mut self.input,
                                &mut self.events));
            if eof {
                debug!("eof");
            }
            self.read_closed = eof;
            if self.input.len() == before {
                break
            }

            while self.input.len() > 0 {
                let (req, amt) = match try!(parse(&self.input)) {
                    Some(pair) => pair,
                    None => {
                        debug!("need more data for a request");
                        return Ok(())
                    }
                };
                let request = Request {
                    inner: req,
                    amt: amt,
                    data: mem::replace(&mut self.input, Vec::new()),
                };
                debug!("got a request");
                self.keepalive = request.version() >= 1;
                self.keepalive = self.keepalive || request.headers().any(|s| {
                    s.0.eq_ignore_ascii_case("connection") &&
                        s.1.eq_ignore_ascii_case(b"keep-alive")
                });

                let response = process(&request);
                response.to_bytes(&mut self.output.buf);

                self.input = request.into_input_buf();
                if !self.keepalive {
                    debug!("disabling keepalive");
                    self.read_closed = true;
                }
            }
        }

        if self.events.is_writable() && self.output.buf.len() > 0 {
            let done = try!(write(&mut self.socket,
                                  &mut self.output,
                                  &mut self.events));
            if done {
                debug!("wrote response");
                if !self.keepalive || self.read_closed {
                    self.closed = true;
                }
            }
        }

        if self.read_closed && self.output.buf.len() == 0 {
            self.closed = true;
        }
        Ok(())
    }
}

fn process(r: &Request) -> Response {
    assert!(r.path() == "/plaintext");
    let mut r = Response::new();
    r.header("Content-Type", "text/plain; charset=UTF-8")
     .body("Hello, World!");
    return r
}

type Slice = (usize, usize);

#[allow(dead_code)]
pub struct Request {
    inner: RawRequest,
    data: Vec<u8>,
    amt: usize,
}

struct RawRequest {
    method: Slice,
    path: Slice,
    version: u8,
    headers: Vec<(Slice, Slice)>,
}

pub struct Headers<'a> {
    iter: slice::Iter<'a, (Slice, Slice)>,
    req: &'a Request,
}

impl Request {
    pub fn method(&self) -> &str {
        str::from_utf8(self.get(&self.inner.method)).unwrap()
    }

    pub fn path(&self) -> &str {
        str::from_utf8(self.get(&self.inner.path)).unwrap()
    }

    pub fn version(&self) -> u8 {
        self.inner.version
    }

    pub fn headers(&self) -> Headers {
        Headers {
            iter: self.inner.headers.iter(),
            req: self,
        }
    }

    pub fn into_input_buf(mut self) -> Vec<u8> {
        self.data.drain(..self.amt);
        self.data
    }

    fn get(&self, s: &Slice) -> &[u8] {
        &self.data[s.0..s.1]
    }
}

impl<'a> Iterator for Headers<'a> {
    type Item = (&'a str, &'a [u8]);

    fn next(&mut self) -> Option<(&'a str, &'a [u8])> {
        self.iter.next().map(|&(ref k, ref v)| {
            (str::from_utf8(self.req.get(k)).unwrap(), self.req.get(v))
        })
    }
}

fn parse(buf: &[u8]) -> io::Result<Option<(RawRequest, usize)>> {
    let mut headers = [httparse::EMPTY_HEADER; 16];
    let mut r = httparse::Request::new(&mut headers);
    let status = try!(r.parse(&buf).map_err(|_| {
        io::Error::new(io::ErrorKind::Other, "failed to parse")
    }));
    return match status {
        httparse::Status::Complete(amt) => {
            debug!("ok {:?}", String::from_utf8_lossy(&buf[..amt]));
            Ok(Some((RawRequest {
                method: slice(buf, r.method.unwrap().as_bytes()),
                path: slice(buf, r.path.unwrap().as_bytes()),
                version: r.version.unwrap(),
                headers: r.headers.iter().map(|h| {
                    (slice(buf, h.name.as_bytes()), slice(buf, &h.value))
                }).collect(),
            }, amt)))
        }
        httparse::Status::Partial => Ok(None),
    };

    fn slice(buf: &[u8], inner: &[u8]) -> Slice {
        let start = inner.as_ptr() as usize - buf.as_ptr() as usize;
        assert!(start < buf.len());
        (start, start + inner.len())
    }
}

pub struct Response {
    headers: Vec<(String, String)>,
    response: String,
}

impl Response {
    pub fn new() -> Response {
        Response {
            headers: Vec::new(),
            response: String::new(),
        }
    }

    pub fn header(&mut self, name: &str, val: &str) -> &mut Response {
        self.headers.push((name.to_string(), val.to_string()));
        self
    }

    pub fn body(&mut self, s: &str) -> &mut Response {
        self.response = s.to_string();
        self
    }

    fn to_bytes(&self, into: &mut Vec<u8>) {
        use std::fmt::Write;

        write!(FastWrite(into), "\
            HTTP/1.1 200 OK\r\n\
            Server: Example\r\n\
            Date: {}\r\n\
            Content-Length: {}\r\n\
        ", time::now().rfc822(), self.response.len()).unwrap();
        for &(ref k, ref v) in &self.headers {
            extend(into, k.as_bytes());
            extend(into, b": ");
            extend(into, v.as_bytes());
            extend(into, b"\r\n");
        }
        extend(into, b"\r\n");
        extend(into, self.response.as_bytes());
    }
}

// TODO: impl fmt::Write for Vec<u8>
//
// Right now `write!` on `Vec<u8>` goes through io::Write and is not super
// speedy, so inline a less-crufty implementation here which doesn't go through
// io::Error.
struct FastWrite<'a>(&'a mut Vec<u8>);

impl<'a> fmt::Write for FastWrite<'a> {
    fn write_str(&mut self, s: &str) -> fmt::Result {
        extend(self.0, s.as_bytes());
        Ok(())
    }
}

// TODO: why does extend_from_slice not optimize?
fn extend(dst: &mut Vec<u8>, data: &[u8]) {
    use std::ptr;
    dst.reserve(data.len());
    let prev = dst.len();
    unsafe {
        ptr::copy_nonoverlapping(data.as_ptr(),
                                 dst.as_mut_ptr().offset(prev as isize),
                                 data.len());
        dst.set_len(prev + data.len());
    }
}

fn read(socket: &mut TcpStream,
        input: &mut Vec<u8>,
        events: &mut EventSet) -> io::Result<bool> {
    match socket.read(unsafe { slice_to_end(input) }) {
        Ok(0) => return Ok(true),
        Ok(n) => {
            let len = input.len();
            unsafe { input.set_len(len + n); }
            return Ok(false)
        }
        Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
            *events = *events & !EventSet::readable();
            return Ok(false)
        }
        Err(e) => return Err(e),
    }

    unsafe fn slice_to_end(v: &mut Vec<u8>) -> &mut [u8] {
        use std::slice;
        if v.capacity() == 0 {
            v.reserve(16);
        }
        if v.capacity() == v.len() {
            v.reserve(1);
        }
        slice::from_raw_parts_mut(v.as_mut_ptr().offset(v.len() as isize),
                                  v.capacity() - v.len())
    }
}

fn write(socket: &mut TcpStream,
         output: &mut Output,
         events: &mut EventSet) -> io::Result<bool> {
    assert!(output.buf.len() > 0);
    loop {
        match socket.write(&output.buf) {
            Ok(0) => {
                return Err(io::Error::new(io::ErrorKind::Other, "early eof2"))
            }
            Ok(n) => {
                output.buf.drain(..n);
                if output.buf.len() == 0 {
                    return Ok(true)
                }
            }
            Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
                *events = *events & !EventSet::writable();
                return Ok(false)
            }
            Err(e) => return Err(e),
        }
    }
}

fn main() {
    env_logger::init().unwrap();

    let poll = mio::Poll::new().unwrap();

    let addr = env::args().nth(1).unwrap_or("127.0.0.1:8080".to_string());
    let addr = addr.parse::<SocketAddr>().unwrap();
    let listener = TcpListener::bind(&addr).unwrap();
    poll.register(&listener,
                  LISTENER,
                  EventSet::readable(),
                  PollOpt::level()).unwrap();

    let mut events = Events::new();
    let mut server = Server::new(&listener);
    loop {
        poll.poll(&mut events, None).unwrap();

        for i in 0..events.len() {
            let event = events.get(i).unwrap();
            server.ready(&poll, event.token(), event.kind());
        }
    }
}
<|endoftext|>