<commit_before>extern crate libc;
extern crate errno;
extern crate rand;
extern crate dylib;

#[macro_use]
extern crate lazy_static;

pub use libc::{c_char, c_int, c_ulong, c_void, off_t, size_t, mode_t, ssize_t};

#[macro_use]
mod errors;
use errors::{OpenFunc, ReadFunc, WriteFunc, SeekFunc, CloseFunc, MmapFunc, Dup2Func, Dup3Func,
             IoctlFunc, BindFunc, StatFunc, FstatFunc, SocketFunc, ConnectFunc, SendRecvFunc,
             ERR_FDS, DELAY_FDS};
use self::errors::matches_addr;
use errors::{remove_fd_if_present, add_fd_if_old_present};

// These functions are designed to conform to their
//  libc counterparts, but may instead inject errors
//  depending on conditions defined in various environment
//  variables.

#[no_mangle]
pub extern "C" fn open64(filename_: *const c_char, flags: c_int, mode: mode_t) -> c_int {
    do_open!(filename_, flags, mode)
}

#[no_mangle]
pub extern "C" fn open(filename_: *const c_char, flags: c_int, mode: mode_t) -> c_int {
    do_open!(filename_, flags, mode)
}

#[no_mangle]
pub extern "C" fn creat(filename_: *const c_char, mode: mode_t) -> c_int {
    const FLAGS: c_int = libc::O_CREAT | libc::O_WRONLY | libc::O_TRUNC;

    // TODO: manpage says this is equivalent to creat() but
    //       should we just call creat() instead?
    do_open!(filename_, FLAGS, mode)
}

const SSIZE_ERR: ssize_t = -1isize;

#[no_mangle]
pub extern "C" fn read(fd: c_int, buf: *mut c_void, nbytes: c_int) -> ssize_t {
    lazy_static! {
        static ref READ_FUNC: ReadFunc = get_libc_func!(ReadFunc, "read");
    }

    injectFaults!(fd, "read", SSIZE_ERR);

    READ_FUNC(fd, buf, nbytes)
}

#[no_mangle]
pub extern "C" fn lseek(fd: c_int, offset: off_t, whence: c_int) -> off_t {
    lazy_static! {
        static ref SEEK_FUNC: SeekFunc = get_libc_func!(SeekFunc, "lseek");
    }

    injectFaults!(fd, "lseek", -1 as i64);

    SEEK_FUNC(fd, offset, whence)
}

#[no_mangle]
pub extern "C" fn write(fd: c_int, buf: *mut c_void, nbytes: c_int) -> ssize_t {
    lazy_static! {
        static ref WRITE_FUNC: WriteFunc = get_libc_func!(WriteFunc, "write");
    }

    injectFaults!(fd, "write", SSIZE_ERR);

    WRITE_FUNC(fd, buf, nbytes)
}

#[no_mangle]
pub extern "C" fn close(fd: c_int) -> c_int {
    lazy_static! {
        static ref CLOSE_FUNC: CloseFunc = get_libc_func!(CloseFunc, "close");
    }

    remove_fd_if_present(fd);

    CLOSE_FUNC(fd)
}

// For now we don't intercept this call for error injection,
//   only for fd tracking.
#[no_mangle]
pub extern "C" fn dup2(oldfd: c_int, newfd: c_int) -> c_int {
    lazy_static! {
        static ref DUP2_FUNC: Dup2Func = get_libc_func!(Dup2Func, "dup2");
    }

    add_fd_if_old_present(oldfd, newfd);

    DUP2_FUNC(oldfd, newfd)
}

// For now we don't intercept this call for error injection,
//   only for fd tracking.
#[no_mangle]
pub extern "C" fn dup3(oldfd: c_int, newfd: c_int, flags: c_int) -> c_int {
    lazy_static! {
        static ref DUP3_FUNC: Dup3Func = get_libc_func!(Dup3Func, "dup3");
    }

    add_fd_if_old_present(oldfd, newfd);

    DUP3_FUNC(oldfd, newfd, flags)
}

#[no_mangle]
pub extern "C" fn ioctl(fd: c_int, req: c_ulong, argp: *mut c_char) -> c_int {
    lazy_static! {
        static ref IOCTL_FUNC: IoctlFunc = get_libc_func!(IoctlFunc, "ioctl");
    }

    injectFaults!(fd, "ioctl", -1 as c_int);

    IOCTL_FUNC(fd, req, argp)
}

use libc::sockaddr;
use self::errors::socklen_t;
#[no_mangle]
pub extern "C" fn connect(sockfd: c_int, addr: *const sockaddr, addrlen: socklen_t) -> c_int {
    lazy_static! {
        static ref CONNECT_FUNC: ConnectFunc = get_libc_func!(ConnectFunc, "connect");
    }

    if unsafe { matches_addr(addr, "LIBFAULTINJ_ERROR_PATH") } {
        ERR_FDS.write().unwrap().insert(sockfd);
    }

    if unsafe { matches_addr(addr, "LIBFAULTINJ_DELAY_PATH") } {
        DELAY_FDS.write().unwrap().insert(sockfd);
    }

    CONNECT_FUNC(sockfd, addr, addrlen)
}

#[no_mangle]
pub extern "C" fn send(sockfd: c_int, buf: *mut c_void, len: size_t, flags: c_int) -> ssize_t {
    lazy_static! {
        static ref SEND_FUNC: SendRecvFunc = get_libc_func!(SendRecvFunc, "send");
    }

    injectFaults!(sockfd, "send", -1);

    SEND_FUNC(sockfd, buf, len, flags)
}

#[no_mangle]
pub extern "C" fn recv(sockfd: c_int, buf: *mut c_void, len: size_t, flags: c_int) -> ssize_t {
    lazy_static! {
        static ref RECV_FUNC: SendRecvFunc = get_libc_func!(SendRecvFunc, "recv");
    }

    injectFaults!(sockfd, "recv", -1);

    RECV_FUNC(sockfd, buf, len, flags)
}

#[no_mangle]
pub extern "C" fn bind(sockfd: c_int, addr: *const sockaddr, addrlen: u8) -> c_int {
    lazy_static! {
        static ref BIND_FUNC: BindFunc = get_libc_func!(BindFunc, "bind");
    }

    if unsafe { matches_addr(addr, "LIBFAULTINJ_ERROR_PATH") } {
        ERR_FDS.write().unwrap().insert(sockfd);
    }

    if unsafe { matches_addr(addr, "LIBFAULTINJ_DELAY_PATH") } {
        DELAY_FDS.write().unwrap().insert(sockfd);
    }

    BIND_FUNC(sockfd, addr, addrlen)
}

#[no_mangle]
pub extern "C" fn fstat(fd: c_int, buf: *const libc::stat) -> c_int {
    lazy_static! {
        static ref FSTAT_FUNC: FstatFunc = get_libc_func!(FstatFunc, "fstat"); // fixme
    }

    println!("intercepted fstat for {:?}", fd);
    injectFaults!(fd, "fstat", -1);

    FSTAT_FUNC(fd, buf)
}

// DISABLED -- these calls are disabled until problems caused can be addressed
#[no_mangle]
#[allow(private_no_mangle_fns)]
#[allow(dead_code)]
#[allow(unused_variables)]
extern "C" fn stat(pathname: *const c_char, buf: *mut libc::stat) -> c_int {
    lazy_static! {
        static ref STAT_FUNC: StatFunc = get_libc_func!(StatFunc, "stat"); // fixme
    }
    // TO BE IMPLEMENTED
    STAT_FUNC(pathname, buf) - 1
}

#[no_mangle]
#[allow(private_no_mangle_fns)]
#[allow(dead_code)]
#[allow(unused_variables)]
extern "C" fn socket(domain: c_int, type_: c_int, protocol: c_int) -> c_int {
    lazy_static! {
        static ref SOCKET_FUNC: SocketFunc = get_libc_func!(SocketFunc, "socket");
    }

    //  injection strategy TBD

    SOCKET_FUNC(domain, type_, protocol)
}

#[no_mangle]
#[allow(private_no_mangle_fns)]
#[allow(dead_code)]
#[allow(unused_variables)]
// pub
extern "C" fn lseek64(fd: c_int, offset: off_t, whence: c_int) -> off_t {
    // TODO -- create a macro to abstract seek/seek64?  define a off64_t type?
    -1 as off_t
}

// mmap() interception is disabled for now.  deadlocks on
//   malloc_init_hard()->mmap()->DynamicLibrary::open()->malloc_init_hard(), at least
//   on systems w/jemalloc.
#[no_mangle]
#[allow(private_no_mangle_fns)]
#[allow(dead_code)]
// pub
extern "C" fn mmap(addr: *mut c_void,
                   length_: size_t,
                   prot: c_int,
                   flags: c_int,
                   fd: c_int,
                   offset: off_t)
                   -> *mut c_void {
    lazy_static! {
        static ref MMAP_FUNC: MmapFunc = get_libc_func!(MmapFunc, "mmap");
    }

    injectFaults!(fd, "mmap", libc::MAP_FAILED);

    MMAP_FUNC(addr, length_, prot, flags, fd, offset)
}
<commit_msg>Update fault.rs<commit_after>extern crate libc;
extern crate errno;
extern crate rand;
extern crate dylib;

#[macro_use]
extern crate lazy_static;

pub use libc::{c_char, c_int, c_ulong, c_void, off_t, size_t, mode_t, ssize_t};

#[macro_use]
mod errors;
use errors::{OpenFunc, ReadFunc, WriteFunc, SeekFunc, CloseFunc, MmapFunc, Dup2Func, Dup3Func,
             IoctlFunc, BindFunc, StatFunc, FstatFunc, SocketFunc, ConnectFunc, SendRecvFunc,
             ERR_FDS, DELAY_FDS};
use self::errors::matches_addr;
use errors::{remove_fd_if_present, add_fd_if_old_present};

// These functions are designed to conform to their
//  libc counterparts, but may instead inject errors
//  depending on conditions defined in various environment
//  variables.

#[no_mangle]
pub extern "C" fn open64(filename_: *const c_char, flags: c_int, mode: mode_t) -> c_int {
    do_open!(filename_, flags, mode)
}

#[no_mangle]
pub extern "C" fn open(filename_: *const c_char, flags: c_int, mode: mode_t) -> c_int {
    do_open!(filename_, flags, mode)
}

#[no_mangle]
pub extern "C" fn creat(filename_: *const c_char, mode: mode_t) -> c_int {
    const FLAGS: c_int = libc::O_CREAT | libc::O_WRONLY | libc::O_TRUNC;

    // TODO: manpage says this is equivalent to creat() but
    //       should we just call creat() instead?
    do_open!(filename_, FLAGS, mode)
}

const SSIZE_ERR: ssize_t = -1isize;

#[no_mangle]
pub extern "C" fn read(fd: c_int, buf: *mut c_void, nbytes: c_int) -> ssize_t {
    lazy_static! {
        static ref READ_FUNC: ReadFunc = get_libc_func!(ReadFunc, "read");
    }

    injectFaults!(fd, "read", SSIZE_ERR);

    READ_FUNC(fd, buf, nbytes)
}

#[no_mangle]
pub extern "C" fn lseek(fd: c_int, offset: off_t, whence: c_int) -> off_t {
    lazy_static! {
        static ref SEEK_FUNC: SeekFunc = get_libc_func!(SeekFunc, "lseek");
    }

    injectFaults!(fd, "lseek", -1 as i64);

    SEEK_FUNC(fd, offset, whence)
}

#[no_mangle]
pub extern "C" fn write(fd: c_int, buf: *mut c_void, nbytes: c_int) -> ssize_t {
    lazy_static! {
        static ref WRITE_FUNC: WriteFunc = get_libc_func!(WriteFunc, "write");
    }

    injectFaults!(fd, "write", SSIZE_ERR);

    WRITE_FUNC(fd, buf, nbytes)
}

#[no_mangle]
pub extern "C" fn close(fd: c_int) -> c_int {
    lazy_static! {
        static ref CLOSE_FUNC: CloseFunc = get_libc_func!(CloseFunc, "close");
    }

    remove_fd_if_present(fd);

    CLOSE_FUNC(fd)
}

// For now we don't intercept this call for error injection,
//   only for fd tracking.
#[no_mangle]
pub extern "C" fn dup2(oldfd: c_int, newfd: c_int) -> c_int {
    lazy_static! {
        static ref DUP2_FUNC: Dup2Func = get_libc_func!(Dup2Func, "dup2");
    }

    add_fd_if_old_present(oldfd, newfd);

    DUP2_FUNC(oldfd, newfd)
}

// For now we don't intercept this call for error injection,
//   only for fd tracking.
#[no_mangle]
pub extern "C" fn dup3(oldfd: c_int, newfd: c_int, flags: c_int) -> c_int {
    lazy_static! {
        static ref DUP3_FUNC: Dup3Func = get_libc_func!(Dup3Func, "dup3");
    }

    add_fd_if_old_present(oldfd, newfd);

    DUP3_FUNC(oldfd, newfd, flags)
}

#[no_mangle]
pub extern "C" fn ioctl(fd: c_int, req: c_ulong, argp: *mut c_char) -> c_int {
    lazy_static! {
        static ref IOCTL_FUNC: IoctlFunc = get_libc_func!(IoctlFunc, "ioctl");
    }

    injectFaults!(fd, "ioctl", -1 as c_int);

    IOCTL_FUNC(fd, req, argp)
}

use libc::sockaddr;
use self::errors::socklen_t;
#[no_mangle]
pub extern "C" fn connect(sockfd: c_int, addr: *const sockaddr, addrlen: socklen_t) -> c_int {
    lazy_static! {
        static ref CONNECT_FUNC: ConnectFunc = get_libc_func!(ConnectFunc, "connect");
    }

    if unsafe { matches_addr(addr, "LIBFAULTINJ_ERROR_PATH") } {
        ERR_FDS.write().unwrap().insert(sockfd);
    }

    if unsafe { matches_addr(addr, "LIBFAULTINJ_DELAY_PATH") } {
        DELAY_FDS.write().unwrap().insert(sockfd);
    }

    CONNECT_FUNC(sockfd, addr, addrlen)
}

#[no_mangle]
pub extern "C" fn send(sockfd: c_int, buf: *mut c_void, len: size_t, flags: c_int) -> ssize_t {
    lazy_static! {
        static ref SEND_FUNC: SendRecvFunc = get_libc_func!(SendRecvFunc, "send");
    }

    injectFaults!(sockfd, "send", -1);

    SEND_FUNC(sockfd, buf, len, flags)
}

#[no_mangle]
pub extern "C" fn recv(sockfd: c_int, buf: *mut c_void, len: size_t, flags: c_int) -> ssize_t {
    lazy_static! {
        static ref RECV_FUNC: SendRecvFunc = get_libc_func!(SendRecvFunc, "recv");
    }

    injectFaults!(sockfd, "recv", -1);

    RECV_FUNC(sockfd, buf, len, flags)
}

#[no_mangle]
pub extern "C" fn bind(sockfd: c_int, addr: *const sockaddr, addrlen: u8) -> c_int {
    lazy_static! {
        static ref BIND_FUNC: BindFunc = get_libc_func!(BindFunc, "bind");
    }

    if unsafe { matches_addr(addr, "LIBFAULTINJ_ERROR_PATH") } {
        ERR_FDS.write().unwrap().insert(sockfd);
    }

    if unsafe { matches_addr(addr, "LIBFAULTINJ_DELAY_PATH") } {
        DELAY_FDS.write().unwrap().insert(sockfd);
    }

    BIND_FUNC(sockfd, addr, addrlen)
}

#[no_mangle]
pub extern "C" fn fstat(fd: c_int, buf: *const libc::stat) -> c_int {
    lazy_static! {
        static ref FSTAT_FUNC: FstatFunc = get_libc_func!(FstatFunc, "fstat"); // fixme
    }

    injectFaults!(fd, "fstat", -1);

    FSTAT_FUNC(fd, buf)
}

// DISABLED -- these calls are disabled until problems caused can be addressed
#[no_mangle]
#[allow(private_no_mangle_fns)]
#[allow(dead_code)]
#[allow(unused_variables)]
extern "C" fn stat(pathname: *const c_char, buf: *mut libc::stat) -> c_int {
    lazy_static! {
        static ref STAT_FUNC: StatFunc = get_libc_func!(StatFunc, "stat"); // fixme
    }
    // TO BE IMPLEMENTED
    STAT_FUNC(pathname, buf) - 1
}

#[no_mangle]
#[allow(private_no_mangle_fns)]
#[allow(dead_code)]
#[allow(unused_variables)]
extern "C" fn socket(domain: c_int, type_: c_int, protocol: c_int) -> c_int {
    lazy_static! {
        static ref SOCKET_FUNC: SocketFunc = get_libc_func!(SocketFunc, "socket");
    }

    //  injection strategy TBD

    SOCKET_FUNC(domain, type_, protocol)
}

#[no_mangle]
#[allow(private_no_mangle_fns)]
#[allow(dead_code)]
#[allow(unused_variables)]
// pub
extern "C" fn lseek64(fd: c_int, offset: off_t, whence: c_int) -> off_t {
    // TODO -- create a macro to abstract seek/seek64?  define a off64_t type?
    -1 as off_t
}

// mmap() interception is disabled for now.  deadlocks on
//   malloc_init_hard()->mmap()->DynamicLibrary::open()->malloc_init_hard(), at least
//   on systems w/jemalloc.
#[no_mangle]
#[allow(private_no_mangle_fns)]
#[allow(dead_code)]
// pub
extern "C" fn mmap(addr: *mut c_void,
                   length_: size_t,
                   prot: c_int,
                   flags: c_int,
                   fd: c_int,
                   offset: off_t)
                   -> *mut c_void {
    lazy_static! {
        static ref MMAP_FUNC: MmapFunc = get_libc_func!(MmapFunc, "mmap");
    }

    injectFaults!(fd, "mmap", libc::MAP_FAILED);

    MMAP_FUNC(addr, length_, prot, flags, fd, offset)
}
<|endoftext|>


<commit_before>// Copyright (C) 2020 Jason Ish
//
// Permission is hereby granted, free of charge, to any person obtaining
// a copy of this software and associated documentation files (the
// "Software"), to deal in the Software without restriction, including
// without limitation the rights to use, copy, modify, merge, publish,
// distribute, sublicense, and/or sell copies of the Software, and to
// permit persons to whom the Software is furnished to do so, subject to
// the following conditions:
//
// The above copyright notice and this permission notice shall be
// included in all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
// EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
// MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
// NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE
// LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
// OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
// WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

use crate::eve::eve::EveJson;
use crate::prelude::*;
use maxminddb::{geoip2, Reader};
use std::sync::Mutex;
use std::time::{Duration, UNIX_EPOCH};

const DAYS_28: i64 = 86400 * 28;
const UPDATE_CHECK_TIMEOUT: u64 = 60;

struct Inner {
    reader: Reader<Vec<u8>>,
    last_modified: u64,
    last_update_check: std::time::Instant,
}

pub struct GeoIP {
    filename: String,
    inner: Mutex<Inner>,
}

impl GeoIP {
    pub fn open(filename: Option<String>) -> Result<GeoIP, Box<dyn std::error::Error>> {
        let (filename, reader) = if let Some(filename) = &filename {
            let reader = Reader::open_readfile(filename)?;
            (filename.clone(), reader)
        } else if let Some(filename) = find_database() {
            let reader = Reader::open_readfile(&filename)?;
            (filename, reader)
        } else {
            return Err("No database file found".into());
        };

        // Warn if database older than 4 weeks.
        let now = chrono::offset::Utc::now();
        let dt = chrono::NaiveDateTime::from_timestamp(reader.metadata.build_epoch as i64, 0);
        if (reader.metadata.build_epoch as i64) < now.timestamp() - DAYS_28 {
            warn!("GeoIP database older than 4 weeks: {}", dt);
        }
        info!("Loaded GeoIP database: {}: {}", filename, dt);

        let last_modified = match Self::get_last_modified(&filename) {
            Ok(last_modified) => last_modified,
            Err(err) => {
                error!(
                    "Failed to get last modified time for {}, file watch will not be enabled: {}",
                    filename, err
                );
                0
            }
        };

        let inner = Inner {
            reader: reader,
            last_modified: last_modified,
            last_update_check: std::time::Instant::now(),
        };

        let geoip = GeoIP {
            filename: filename,
            inner: Mutex::new(inner),
        };
        return Ok(geoip);
    }

    pub fn get_last_modified(filename: &str) -> Result<u64, Box<dyn std::error::Error>> {
        let last_modified = std::fs::metadata(filename)?
            .modified()?
            .duration_since(UNIX_EPOCH)?
            .as_secs();
        Ok(last_modified)
    }

    fn check_for_update(&self, inner: &mut Inner) -> bool {
        if inner.last_update_check.elapsed() < Duration::from_secs(UPDATE_CHECK_TIMEOUT) {
            return false;
        }
        let last_modified = match Self::get_last_modified(&self.filename) {
            Ok(last_modified) => last_modified,
            Err(err) => {
                warn!(
                    "Failed to get modification time for GeoIP database file {}: {}: ",
                    self.filename, err
                );
                return false;
            }
        };
        let mut updated = false;
        if last_modified <= inner.last_modified {
            debug!("GeoIP database file has not been updated");
        } else {
            debug!("GeoIP database file on disk has been updated");
            match Reader::open_readfile(&self.filename) {
                Err(err) => {
                    error!("Failed to open new GeoIP database file: {}", err);
                }
                Ok(new_reader) => {
                    inner.reader = new_reader;
                    inner.last_modified = last_modified;
                    updated = true;
                }
            }
        }
        inner.last_update_check = std::time::Instant::now();
        updated
    }

    pub fn lookup_city_from_str(
        &self,
        addr: &str,
    ) -> Result<geoip2::City, Box<dyn std::error::Error>> {
        let mut inner = self.inner.lock().unwrap();
        if self.check_for_update(&mut inner) {
            let build_time =
                chrono::NaiveDateTime::from_timestamp(inner.reader.metadata.build_epoch as i64, 0);
            info!("GeoIP database has been updated to {}", build_time);
        }
        let ip: std::net::IpAddr = std::str::FromStr::from_str(addr)?;
        let city = inner.reader.lookup(ip)?;
        Ok(city)
    }

    pub fn add_geoip_to_eve(&self, eve: &mut EveJson) {
        if let EveJson::String(addr) = &eve["dest_ip"] {
            if let Ok(city) = self.lookup_city_from_str(addr) {
                eve["geoip_destination"] = self.as_json(city);
            }
        }
        if let EveJson::String(addr) = &eve["src_ip"] {
            if let Ok(city) = self.lookup_city_from_str(addr) {
                eve["geoip_source"] = self.as_json(city);
            }
        }
    }

    fn as_json(&self, city: geoip2::City) -> serde_json::Value {
        let mut obj = serde_json::json!({});
        if let Some(city) = city.city {
            if let Some(names) = city.names {
                if let Some(name) = names.get("en") {
                    obj["city_name"] = name.to_string().into();
                }
            }
        }
        if let Some(country) = city.country {
            if let Some(names) = country.names {
                if let Some(name) = names.get("en") {
                    obj["country_name"] = name.to_string().into();
                }
            }
            if let Some(iso_code) = country.iso_code {
                obj["country_iso_code"] = iso_code.into();
            }
        }
        if let Some(subdivisions) = city.subdivisions {
            if let Some(subdivision) = &subdivisions.first() {
                if let Some(names) = &subdivision.names {
                    if let Some(name) = names.get("en") {
                        obj["region_name"] = name.to_string().into();
                    }
                }
                if let Some(iso_code) = &subdivision.iso_code {
                    obj["region_iso_code"] = iso_code.to_string().into();
                }
            }
        }
        if let Some(location) = city.location {
            let mut locobj = serde_json::json!({});
            let mut include = false;
            if let Some(lat) = location.latitude {
                locobj["lat"] = lat.into();
                include = true;
            }
            if let Some(lon) = location.longitude {
                locobj["lon"] = lon.into();
                include = true;
            }
            if include {
                obj["location"] = locobj;
            }
        }
        if let Some(continent) = city.continent {
            if let Some(names) = continent.names {
                if let Some(name) = names.get("en") {
                    obj["continent_name"] = name.to_string().into();
                }
            }
        }
        return obj;
    }
}

lazy_static! {
    static ref PATHS: Vec<&'static str> = {
        vec![
            "/etc/evebox/GeoLite2-City.mmdb",
            "/usr/local/share/GeoIP/GeoLite2-City.mmdb",
            "/usr/share/GeoIP/GeoLite2-City.mmdb",
        ]
    };
}

fn find_database() -> Option<String> {
    for filename in PATHS.iter() {
        if maxminddb::Reader::open_readfile(filename).is_ok() {
            debug!("Found geoip database file {}", filename);
            return Some(filename.to_string());
        }
    }
    None
}

#[cfg(test)]
mod test {

    #[test]
    fn lookup_example() {
        let db = super::GeoIP::open(None).unwrap();
        let _city = db.lookup_city_from_str("128.101.101.101").unwrap();
    }
}
<commit_msg>geoip/test: don't fail if db not found<commit_after>// Copyright (C) 2020 Jason Ish
//
// Permission is hereby granted, free of charge, to any person obtaining
// a copy of this software and associated documentation files (the
// "Software"), to deal in the Software without restriction, including
// without limitation the rights to use, copy, modify, merge, publish,
// distribute, sublicense, and/or sell copies of the Software, and to
// permit persons to whom the Software is furnished to do so, subject to
// the following conditions:
//
// The above copyright notice and this permission notice shall be
// included in all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
// EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
// MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
// NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE
// LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
// OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
// WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

use crate::eve::eve::EveJson;
use crate::prelude::*;
use maxminddb::{geoip2, Reader};
use std::sync::Mutex;
use std::time::{Duration, UNIX_EPOCH};

const DAYS_28: i64 = 86400 * 28;
const UPDATE_CHECK_TIMEOUT: u64 = 60;

struct Inner {
    reader: Reader<Vec<u8>>,
    last_modified: u64,
    last_update_check: std::time::Instant,
}

pub struct GeoIP {
    filename: String,
    inner: Mutex<Inner>,
}

impl GeoIP {
    pub fn open(filename: Option<String>) -> Result<GeoIP, Box<dyn std::error::Error>> {
        let (filename, reader) = if let Some(filename) = &filename {
            let reader = Reader::open_readfile(filename)?;
            (filename.clone(), reader)
        } else if let Some(filename) = find_database() {
            let reader = Reader::open_readfile(&filename)?;
            (filename, reader)
        } else {
            return Err("No database file found".into());
        };

        // Warn if database older than 4 weeks.
        let now = chrono::offset::Utc::now();
        let dt = chrono::NaiveDateTime::from_timestamp(reader.metadata.build_epoch as i64, 0);
        if (reader.metadata.build_epoch as i64) < now.timestamp() - DAYS_28 {
            warn!("GeoIP database older than 4 weeks: {}", dt);
        }
        info!("Loaded GeoIP database: {}: {}", filename, dt);

        let last_modified = match Self::get_last_modified(&filename) {
            Ok(last_modified) => last_modified,
            Err(err) => {
                error!(
                    "Failed to get last modified time for {}, file watch will not be enabled: {}",
                    filename, err
                );
                0
            }
        };

        let inner = Inner {
            reader: reader,
            last_modified: last_modified,
            last_update_check: std::time::Instant::now(),
        };

        let geoip = GeoIP {
            filename: filename,
            inner: Mutex::new(inner),
        };
        return Ok(geoip);
    }

    pub fn get_last_modified(filename: &str) -> Result<u64, Box<dyn std::error::Error>> {
        let last_modified = std::fs::metadata(filename)?
            .modified()?
            .duration_since(UNIX_EPOCH)?
            .as_secs();
        Ok(last_modified)
    }

    fn check_for_update(&self, inner: &mut Inner) -> bool {
        if inner.last_update_check.elapsed() < Duration::from_secs(UPDATE_CHECK_TIMEOUT) {
            return false;
        }
        let last_modified = match Self::get_last_modified(&self.filename) {
            Ok(last_modified) => last_modified,
            Err(err) => {
                warn!(
                    "Failed to get modification time for GeoIP database file {}: {}: ",
                    self.filename, err
                );
                return false;
            }
        };
        let mut updated = false;
        if last_modified <= inner.last_modified {
            debug!("GeoIP database file has not been updated");
        } else {
            debug!("GeoIP database file on disk has been updated");
            match Reader::open_readfile(&self.filename) {
                Err(err) => {
                    error!("Failed to open new GeoIP database file: {}", err);
                }
                Ok(new_reader) => {
                    inner.reader = new_reader;
                    inner.last_modified = last_modified;
                    updated = true;
                }
            }
        }
        inner.last_update_check = std::time::Instant::now();
        updated
    }

    pub fn lookup_city_from_str(
        &self,
        addr: &str,
    ) -> Result<geoip2::City, Box<dyn std::error::Error>> {
        let mut inner = self.inner.lock().unwrap();
        if self.check_for_update(&mut inner) {
            let build_time =
                chrono::NaiveDateTime::from_timestamp(inner.reader.metadata.build_epoch as i64, 0);
            info!("GeoIP database has been updated to {}", build_time);
        }
        let ip: std::net::IpAddr = std::str::FromStr::from_str(addr)?;
        let city = inner.reader.lookup(ip)?;
        Ok(city)
    }

    pub fn add_geoip_to_eve(&self, eve: &mut EveJson) {
        if let EveJson::String(addr) = &eve["dest_ip"] {
            if let Ok(city) = self.lookup_city_from_str(addr) {
                eve["geoip_destination"] = self.as_json(city);
            }
        }
        if let EveJson::String(addr) = &eve["src_ip"] {
            if let Ok(city) = self.lookup_city_from_str(addr) {
                eve["geoip_source"] = self.as_json(city);
            }
        }
    }

    fn as_json(&self, city: geoip2::City) -> serde_json::Value {
        let mut obj = serde_json::json!({});
        if let Some(city) = city.city {
            if let Some(names) = city.names {
                if let Some(name) = names.get("en") {
                    obj["city_name"] = name.to_string().into();
                }
            }
        }
        if let Some(country) = city.country {
            if let Some(names) = country.names {
                if let Some(name) = names.get("en") {
                    obj["country_name"] = name.to_string().into();
                }
            }
            if let Some(iso_code) = country.iso_code {
                obj["country_iso_code"] = iso_code.into();
            }
        }
        if let Some(subdivisions) = city.subdivisions {
            if let Some(subdivision) = &subdivisions.first() {
                if let Some(names) = &subdivision.names {
                    if let Some(name) = names.get("en") {
                        obj["region_name"] = name.to_string().into();
                    }
                }
                if let Some(iso_code) = &subdivision.iso_code {
                    obj["region_iso_code"] = iso_code.to_string().into();
                }
            }
        }
        if let Some(location) = city.location {
            let mut locobj = serde_json::json!({});
            let mut include = false;
            if let Some(lat) = location.latitude {
                locobj["lat"] = lat.into();
                include = true;
            }
            if let Some(lon) = location.longitude {
                locobj["lon"] = lon.into();
                include = true;
            }
            if include {
                obj["location"] = locobj;
            }
        }
        if let Some(continent) = city.continent {
            if let Some(names) = continent.names {
                if let Some(name) = names.get("en") {
                    obj["continent_name"] = name.to_string().into();
                }
            }
        }
        return obj;
    }
}

lazy_static! {
    static ref PATHS: Vec<&'static str> = {
        vec![
            "/etc/evebox/GeoLite2-City.mmdb",
            "/usr/local/share/GeoIP/GeoLite2-City.mmdb",
            "/usr/share/GeoIP/GeoLite2-City.mmdb",
        ]
    };
}

fn find_database() -> Option<String> {
    for filename in PATHS.iter() {
        if maxminddb::Reader::open_readfile(filename).is_ok() {
            debug!("Found geoip database file {}", filename);
            return Some(filename.to_string());
        }
    }
    None
}

#[cfg(test)]
mod test {

    #[test]
    fn lookup_example() {
        if let Ok(db) = super::GeoIP::open(None) {
            let _city = db.lookup_city_from_str("128.101.101.101").unwrap();
        }
    }
}
<|endoftext|>


<commit_before>extern crate rand;
use rand::Rng;
use std::cmp::{Ord, Ordering};
use std::fmt;

// comands in config-file start with
pub const CONF_LINE_IDENTIFIER__CONTROL: char = ':';

// the default can be changed by one of the following switches
const DEFAULT_REWARDING_SCHEME: RewardingScheme = RewardingScheme::UnhideWhenGuessedChar;

// Keyword to switch rewarding scheme
// :traditional-rewarding
const UNHIDE_WHEN_LOST_LIVE_IDENTIFIER: &str = "traditional-rewarding";

// Keyword to switch rewarding scheme
// :success-rewarding
const UNHIDE_WHEN_GUESSED_CHAR_IDENTIFIER: &str = "success-rewarding";

// images in config file start with
pub const CONF_LINE_IDENTIFIER__IMAGE: char = '|';

const BIG_IMAGE: usize = 100; // sort algorithm <-> random algorithm

// first char of image lines must be '|'
const DEFAULT_IMAGES: & [&str] = &[
    "
|    ,,,,,
|   (o   o)
|    /. .\\
|   (_____)
|     : :
|    ##O##'
|  ,,,: :,,,
| _)\\ : : /(____
|{  \\     /  ___}
| \\/)     ((/
|  (_______)
|    :   :
|    :   :
|   / \\ / \\
|   \"\"\" \"\"\"
",
    "
|    |\\_|X|_/|
|   /         \\
| =(  O     O  )=
|  -\\    o    /-
|   / .-----. \\
| /_ | o   o |_ \\
|(U  |       |  U)
|   _|_     _|_
|  (   )---(   )
",
    "
|        _.---._    /\\\\
|     ./'       \"--`\\//
|   ./              o \\
|  /./\\  )______   \\__ \\
| ./  / /\\ \\   | \\ \\  \\ \\
|    / /  \\ \\  | |\\ \\  \\7
|     \"     \"    \"  \"        VK
",
    "
|       ,.
|      (_|,.
|     ,' /, )_______   _
|  __j o``-'        `.'-)'
| (\")                 \\'
|  `-j                |
|    `-._(           /
|       |_\\  |--^.  /
|      /_]'|_| /_)_/
|         /_]'  /_]'
# Author: hjw
",
    "
|        _
|       [ ]
|      (   )
|       |>|
|    __/===\\__
|   //| o=o |\\\\
| <]  | o=o |  [>
|     \\=====/
|    / / | \\ \\
|   <_________>
",
    "
|                          (_)(_)
|                          /     \\
|                         /       |
|                        /   \\  * |
|          ________     /    /\\__/
|  _      /        \\   /    /
| / \\    /  ____    \\_/    /
|//\\ \\  /  /    \\         /
|V  \\ \\/  /      \\       /
|    \\___/        \\_____/
",
    "
|         .-.
|        (. .)__,')
|        / V      )
|  ()    \\  (   \\/
|<)-`\\()  `._`._ \\
|  <).>=====<<==`'====
|   C-'`(>
# Author: hjw
",
    "
| >(. )
|  |  (     /)
|  |   \\___/ )
|  (   ----- )  >@)_//   >@)_//  >@)_//  >@)_//
|   \\_______/    (__)     (__)    (__)    (__)
|~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~
",
    "
|           __
|           /(`o
|     ,-,  //  \\\\
|    (,,,) ||   V
|   (,,,,)\\//
|   (,,,/w)-'
|   \\,,/w)
|   `V/uu
|     / |
|     | |
|     o o
|     \\ |
|\\,/  ,\\|,.  \\,/
",
    "
|o
| \\_/\\o
|( Oo)                    \\|/
|(_=-)  .===O-  ~~Z~A~P~~ -O-
|/   \\_/U'                /|\\
|||  |_/
|\\\\  |
|{K ||
| | PP
| | ||
| (__\\\\
# Author: ac
",
    "
|      ______
|     /     /\\
|    /     /  \\
|   /_____/----\\_    (
|  \"     \"          ).
| _ ___          o (:') o
|(@))_))        o ~/~~\\~ o
|                o  o  o
",
    "
|                             _______     |\\
|                            |License|    | \\
|  _____                     | ~~*~~ |    |  \\
| |     |  (((        .--.   |_______|    |
| |DrJRO| ~OvO~ __   (////)               |
| |     | ( _ )|==|   \\__/                |
| |o    |  \\_/ |_(|  /    \\   _______     |
| |     | //|\\\\   \\\\//|  |\\\\  |__o__|     |
| |   __|//\\_/\\\\ __\\/ |__|//  |__o__|     |
| |  |==\"\"//=\\\\\"\"====|||||)   |__o__|     |
|_|__||_|_||_||_____||||||____|__o__|_____|
|    ||  (_) (_)    ||||||                \\
|    []             [(_)(_)
",
];

#[derive(PartialOrd, Eq, PartialEq, Debug, Copy, Clone)] //omitting Ord
pub struct ImageChar {
    pub point: (u8, u8),
    pub code: char,
}

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

// Ord enables us to v.sort()
impl Ord for ImageChar {
    fn cmp(&self, other: &Self) -> Ordering {
        fn weight(ic: &ImageChar) -> isize {
            let &ImageChar {
                point: (x, y),
                ..
            } = ic;
            // points near the lower left corner are light
            x as isize - y as isize
        }
        weight(&self).cmp(&weight(&other))
    }
}

#[derive(Debug, PartialEq)]
pub enum RewardingScheme {
    UnhideWhenLostLife,
    UnhideWhenGuessedChar,
}

#[derive(Debug)]
pub struct Image {
    pub ichars: Vec<ImageChar>,
    pub offset: (usize, usize),
    pub dimension: (u8, u8),
    pub visible_points: usize,
    pub rewarding_scheme: RewardingScheme,
}

impl fmt::Display for Image {
    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        let mut s = String::new();
        for ic in self.ichars.iter().take(self.visible_points) {
            let &ImageChar {
                point: (x, y),
                code,
            } = ic;
            s = s
                + "\x1b["
                + &(y as usize + 1 + self.offset.1).to_string()
                + ";"
                + &(x as usize + 1 + self.offset.0).to_string()
                + "f"
                + &code.to_string();
        }
        // after printing the image s, bring the cursor below
        write!(
            f,
            "{}\x1b[{};0f",
            s,
            (self.dimension.1 as usize + 1 + self.offset.1).to_string()
        )
    }
}

impl Image {
    pub fn new(string: &str, offset: (usize, usize)) -> Self {
        let mut v: Vec<ImageChar> = Vec::new();

        let mut rewarding_scheme: RewardingScheme = DEFAULT_REWARDING_SCHEME;
        for (y, line) in string
            // split in lines
            .lines()
            // interpret identifier line
            .filter_map(|l| {
                if l.starts_with(CONF_LINE_IDENTIFIER__CONTROL) {
                    if l[1..].trim().contains(UNHIDE_WHEN_LOST_LIVE_IDENTIFIER) {
                        rewarding_scheme = RewardingScheme::UnhideWhenLostLife;
                    }
                    if l[1..].trim().contains(UNHIDE_WHEN_GUESSED_CHAR_IDENTIFIER) {
                        rewarding_scheme = RewardingScheme::UnhideWhenGuessedChar;
                    }
                    None
                } else {
                    Some(l)
                }
            })
            // consider only lines starting with '|'
            .filter(|&l| l.starts_with(CONF_LINE_IDENTIFIER__IMAGE))
            .enumerate()
        //.inspect(|&(n,l)| println!("line {:?}: {:?} ", n,l))
        {
            let mut ii: Vec<_> = line
                .char_indices()
                // skip first char '|'
                .skip(1)
                // consider only chars != ' '
                .filter(|&(_, c)| c != ' ')
                // save in ImageChar object
                .map(|(x, c)| ImageChar {
                    point: (x as u8, y as u8),
                    code: c,
                })
                .collect();
            v.append(&mut ii);
        }

        // find dimensions
        let mut x_max = 0;
        let mut y_max = 0;
        for i in v.iter() {
            let &ImageChar {
                point: (x, y),
                ..
            } = i;
            if x > x_max {
                x_max = x
            };
            if y > y_max {
                y_max = y
            };
        }

        // order points
        let v_len = v.len();
        if v_len <= BIG_IMAGE {
            v.sort(); // Sort algorithm, see "impl Ord for ImageChar"
        } else {
            rand::thread_rng().shuffle(&mut v); // points appear randomly.
        }

        if v.len() == 0 {
            Self::new(rand::thread_rng().choose(&DEFAULT_IMAGES).unwrap(), offset)
        } else {
            Self {
                ichars: v,
                offset,
                dimension: (x_max,y_max),
                visible_points: v_len,
                rewarding_scheme,
            }
        }
    }

    pub fn disclose(&mut self, lives_frac: (usize, usize), guessed_chars_frac: (usize, usize)) {
        let l = self.ichars.len();

        let as_points = |(n, d)| (3 * l * (d - n) as usize / d as usize + 1 * l) / 4;

        self.visible_points = match self.rewarding_scheme {
            RewardingScheme::UnhideWhenGuessedChar => as_points(guessed_chars_frac),
            RewardingScheme::UnhideWhenLostLife => as_points(lives_frac),
        };
    }
}
<commit_msg>item name starts with its containing module's name<commit_after>extern crate rand;
use rand::Rng;
use std::cmp::{Ord, Ordering};
use std::fmt;

// comands in config-file start with
pub const CONF_LINE_IDENTIFIER__CONTROL: char = ':';

// the default can be changed by one of the following switches
const DEFAULT_REWARDING_SCHEME: RewardingScheme = RewardingScheme::UnhideWhenGuessedChar;

// Keyword to switch rewarding scheme
// :traditional-rewarding
const UNHIDE_WHEN_LOST_LIVE_IDENTIFIER: &str = "traditional-rewarding";

// Keyword to switch rewarding scheme
// :success-rewarding
const UNHIDE_WHEN_GUESSED_CHAR_IDENTIFIER: &str = "success-rewarding";

// images in config file start with
pub const CONF_LINE_IDENTIFIER__IMAGE: char = '|';

const BIG_IMAGE: usize = 100; // sort algorithm <-> random algorithm

// first char of image lines must be '|'
const DEFAULT_IMAGES: & [&str] = &[
    "
|    ,,,,,
|   (o   o)
|    /. .\\
|   (_____)
|     : :
|    ##O##'
|  ,,,: :,,,
| _)\\ : : /(____
|{  \\     /  ___}
| \\/)     ((/
|  (_______)
|    :   :
|    :   :
|   / \\ / \\
|   \"\"\" \"\"\"
",
    "
|    |\\_|X|_/|
|   /         \\
| =(  O     O  )=
|  -\\    o    /-
|   / .-----. \\
| /_ | o   o |_ \\
|(U  |       |  U)
|   _|_     _|_
|  (   )---(   )
",
    "
|        _.---._    /\\\\
|     ./'       \"--`\\//
|   ./              o \\
|  /./\\  )______   \\__ \\
| ./  / /\\ \\   | \\ \\  \\ \\
|    / /  \\ \\  | |\\ \\  \\7
|     \"     \"    \"  \"        VK
",
    "
|       ,.
|      (_|,.
|     ,' /, )_______   _
|  __j o``-'        `.'-)'
| (\")                 \\'
|  `-j                |
|    `-._(           /
|       |_\\  |--^.  /
|      /_]'|_| /_)_/
|         /_]'  /_]'
# Author: hjw
",
    "
|        _
|       [ ]
|      (   )
|       |>|
|    __/===\\__
|   //| o=o |\\\\
| <]  | o=o |  [>
|     \\=====/
|    / / | \\ \\
|   <_________>
",
    "
|                          (_)(_)
|                          /     \\
|                         /       |
|                        /   \\  * |
|          ________     /    /\\__/
|  _      /        \\   /    /
| / \\    /  ____    \\_/    /
|//\\ \\  /  /    \\         /
|V  \\ \\/  /      \\       /
|    \\___/        \\_____/
",
    "
|         .-.
|        (. .)__,')
|        / V      )
|  ()    \\  (   \\/
|<)-`\\()  `._`._ \\
|  <).>=====<<==`'====
|   C-'`(>
# Author: hjw
",
    "
| >(. )
|  |  (     /)
|  |   \\___/ )
|  (   ----- )  >@)_//   >@)_//  >@)_//  >@)_//
|   \\_______/    (__)     (__)    (__)    (__)
|~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~ ~
",
    "
|           __
|           /(`o
|     ,-,  //  \\\\
|    (,,,) ||   V
|   (,,,,)\\//
|   (,,,/w)-'
|   \\,,/w)
|   `V/uu
|     / |
|     | |
|     o o
|     \\ |
|\\,/  ,\\|,.  \\,/
",
    "
|o
| \\_/\\o
|( Oo)                    \\|/
|(_=-)  .===O-  ~~Z~A~P~~ -O-
|/   \\_/U'                /|\\
|||  |_/
|\\\\  |
|{K ||
| | PP
| | ||
| (__\\\\
# Author: ac
",
    "
|      ______
|     /     /\\
|    /     /  \\
|   /_____/----\\_    (
|  \"     \"          ).
| _ ___          o (:') o
|(@))_))        o ~/~~\\~ o
|                o  o  o
",
    "
|                             _______     |\\
|                            |License|    | \\
|  _____                     | ~~*~~ |    |  \\
| |     |  (((        .--.   |_______|    |
| |DrJRO| ~OvO~ __   (////)               |
| |     | ( _ )|==|   \\__/                |
| |o    |  \\_/ |_(|  /    \\   _______     |
| |     | //|\\\\   \\\\//|  |\\\\  |__o__|     |
| |   __|//\\_/\\\\ __\\/ |__|//  |__o__|     |
| |  |==\"\"//=\\\\\"\"====|||||)   |__o__|     |
|_|__||_|_||_||_____||||||____|__o__|_____|
|    ||  (_) (_)    ||||||                \\
|    []             [(_)(_)
",
];

#[derive(PartialOrd, Eq, PartialEq, Debug, Copy, Clone)] //omitting Ord
pub struct ImChar {
    pub point: (u8, u8),
    pub code: char,
}

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

// Ord enables us to v.sort()
impl Ord for ImChar {
    fn cmp(&self, other: &Self) -> Ordering {
        fn weight(ic: &ImChar) -> isize {
            let &ImChar {
                point: (x, y),
                ..
            } = ic;
            // points near the lower left corner are light
            x as isize - y as isize
        }
        weight(&self).cmp(&weight(&other))
    }
}

#[derive(Debug, PartialEq)]
pub enum RewardingScheme {
    UnhideWhenLostLife,
    UnhideWhenGuessedChar,
}

#[derive(Debug)]
pub struct Image {
    pub ichars: Vec<ImChar>,
    pub offset: (usize, usize),
    pub dimension: (u8, u8),
    pub visible_points: usize,
    pub rewarding_scheme: RewardingScheme,
}

impl fmt::Display for Image {
    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        let mut s = String::new();
        for ic in self.ichars.iter().take(self.visible_points) {
            let &ImChar {
                point: (x, y),
                code,
            } = ic;
            s = s
                + "\x1b["
                + &(y as usize + 1 + self.offset.1).to_string()
                + ";"
                + &(x as usize + 1 + self.offset.0).to_string()
                + "f"
                + &code.to_string();
        }
        // after printing the image s, bring the cursor below
        write!(
            f,
            "{}\x1b[{};0f",
            s,
            (self.dimension.1 as usize + 1 + self.offset.1).to_string()
        )
    }
}

impl Image {
    pub fn new(string: &str, offset: (usize, usize)) -> Self {
        let mut v: Vec<ImChar> = Vec::new();

        let mut rewarding_scheme: RewardingScheme = DEFAULT_REWARDING_SCHEME;
        for (y, line) in string
            // split in lines
            .lines()
            // interpret identifier line
            .filter_map(|l| {
                if l.starts_with(CONF_LINE_IDENTIFIER__CONTROL) {
                    if l[1..].trim().contains(UNHIDE_WHEN_LOST_LIVE_IDENTIFIER) {
                        rewarding_scheme = RewardingScheme::UnhideWhenLostLife;
                    }
                    if l[1..].trim().contains(UNHIDE_WHEN_GUESSED_CHAR_IDENTIFIER) {
                        rewarding_scheme = RewardingScheme::UnhideWhenGuessedChar;
                    }
                    None
                } else {
                    Some(l)
                }
            })
            // consider only lines starting with '|'
            .filter(|&l| l.starts_with(CONF_LINE_IDENTIFIER__IMAGE))
            .enumerate()
        //.inspect(|&(n,l)| println!("line {:?}: {:?} ", n,l))
        {
            let mut ii: Vec<_> = line
                .char_indices()
                // skip first char '|'
                .skip(1)
                // consider only chars != ' '
                .filter(|&(_, c)| c != ' ')
                // save in ImageChar object
                .map(|(x, c)| ImChar {
                    point: (x as u8, y as u8),
                    code: c,
                })
                .collect();
            v.append(&mut ii);
        }

        // find dimensions
        let mut x_max = 0;
        let mut y_max = 0;
        for i in v.iter() {
            let &ImChar {
                point: (x, y),
                ..
            } = i;
            if x > x_max {
                x_max = x
            };
            if y > y_max {
                y_max = y
            };
        }

        // order points
        let v_len = v.len();
        if v_len <= BIG_IMAGE {
            v.sort(); // Sort algorithm, see "impl Ord for ImageChar"
        } else {
            rand::thread_rng().shuffle(&mut v); // points appear randomly.
        }

        if v.len() == 0 {
            Self::new(rand::thread_rng().choose(&DEFAULT_IMAGES).unwrap(), offset)
        } else {
            Self {
                ichars: v,
                offset,
                dimension: (x_max,y_max),
                visible_points: v_len,
                rewarding_scheme,
            }
        }
    }

    pub fn disclose(&mut self, lives_frac: (usize, usize), guessed_chars_frac: (usize, usize)) {
        let l = self.ichars.len();

        let as_points = |(n, d)| (3 * l * (d - n) as usize / d as usize + 1 * l) / 4;

        self.visible_points = match self.rewarding_scheme {
            RewardingScheme::UnhideWhenGuessedChar => as_points(guessed_chars_frac),
            RewardingScheme::UnhideWhenLostLife => as_points(lives_frac),
        };
    }
}
<|endoftext|>


<commit_before>// FIXME: Should be bindgenned

use libc::{c_void, c_char, size_t};

#[nolink]
extern mod cg {
    pub fn CGDataProviderCreateWithFilename(filename: *c_char) -> CGDataProviderRef;
    pub fn CGDataProviderCreateWithData(info: *c_void,
                                        data: *c_void,
                                        size: size_t,
                                        releaseData: CGDataProviderReleaseDataCallback
                                        ) -> CGDataProviderRef;
    pub fn CGDataProviderRelease(provider: CGDataProviderRef);
    pub fn CGFontCreateWithDataProvider(provider: CGDataProviderRef) -> CGFontRef;
    pub fn CGFontRelease(font: CGFontRef);
}

pub type CGFontRef = *c_void;
pub type CGDataProviderRef = *c_void;
pub type CGDataProviderReleaseDataCallback = *c_void;
<commit_msg>Update for language changes<commit_after>// FIXME: Should be bindgenned

use libc::{c_void, c_char, size_t};

#[nolink]
pub extern mod cg {
    pub fn CGDataProviderCreateWithFilename(filename: *c_char) -> CGDataProviderRef;
    pub fn CGDataProviderCreateWithData(info: *c_void,
                                        data: *c_void,
                                        size: size_t,
                                        releaseData: CGDataProviderReleaseDataCallback
                                        ) -> CGDataProviderRef;
    pub fn CGDataProviderRelease(provider: CGDataProviderRef);
    pub fn CGFontCreateWithDataProvider(provider: CGDataProviderRef) -> CGFontRef;
    pub fn CGFontRelease(font: CGFontRef);
}

pub type CGFontRef = *c_void;
pub type CGDataProviderRef = *c_void;
pub type CGDataProviderReleaseDataCallback = *c_void;
<|endoftext|>


<commit_before>extern crate libc;
extern crate num;

use array::Array;
use dim4::Dim4;
use defines::Aftype;
use self::libc::{uint8_t, c_int, c_uint, c_double};
use self::num::Complex;

use std::vec::Vec;

type MutAfArray = *mut self::libc::c_longlong;
type MutDouble  = *mut self::libc::c_double;
type MutUint    = *mut self::libc::c_uint;
type AfArray    = self::libc::c_longlong;
type DimT       = self::libc::c_longlong;
type Intl       = self::libc::c_longlong;
type Uintl      = self::libc::c_ulonglong;

#[allow(dead_code)]
extern {
    fn af_constant(out: MutAfArray, val: c_double,
                   ndims: c_uint, dims: *const DimT, afdtype: c_int) -> c_int;

    fn af_constant_complex(out: MutAfArray, real: c_double, imag: c_double,
                           ndims: c_uint, dims: *const DimT, afdtype: c_int) -> c_int;

    fn af_constant_long(out: MutAfArray, val: Intl,
                        ndims: c_uint, dims: *const DimT) -> c_int;

    fn af_constant_ulong(out: MutAfArray, val: Uintl,
                         ndims: c_uint, dims: *const DimT) -> c_int;

    fn af_range(out: MutAfArray, ndims: c_uint, dims: *const DimT,
                seq_dims: c_int, afdtype: uint8_t) -> c_int;

    fn af_iota(out: MutAfArray, ndims: c_uint, dims: *const DimT,
               t_ndims: c_uint, tdims: *const DimT, afdtype: uint8_t) -> c_int;

    fn af_randu(out: MutAfArray, ndims: c_uint, dims: *const DimT, afdtype: uint8_t) -> c_int;
    fn af_randn(out: MutAfArray, ndims: c_uint, dims: *const DimT, afdtype: uint8_t) -> c_int;

    fn af_set_seed(seed: Uintl);
    fn af_get_seed(seed: *mut Uintl);

    fn af_identity(out: MutAfArray, ndims: c_uint, dims: *const DimT, afdtype: uint8_t) -> c_int;
    fn af_diag_create(out: MutAfArray, arr: AfArray, num: c_int) -> c_int;
    fn af_diag_extract(out: MutAfArray, arr: AfArray, num: c_int) -> c_int;
    fn af_join(out: MutAfArray, dim: c_int, first: AfArray, second: AfArray) -> c_int;
    fn af_join_many(out: MutAfArray, dim: c_int, n_arrays: c_uint, inpts: *const AfArray) -> c_int;

    fn af_tile(out: MutAfArray, arr: AfArray, x: c_uint, y: c_uint, z: c_uint, w: c_uint) -> c_int;
    fn af_reorder(o: MutAfArray, a: AfArray, x: c_uint, y: c_uint, z: c_uint, w: c_uint) -> c_int;
    fn af_shift(o: MutAfArray, a: AfArray, x: c_uint, y: c_uint, z: c_uint, w: c_uint) -> c_int;
    fn af_moddims(out: MutAfArray, arr: AfArray, ndims: c_uint, dims: *const DimT) -> c_int;

    fn af_flat(out: MutAfArray, arr: AfArray) -> c_int;
    fn af_flip(out: MutAfArray, arr: AfArray, dim: c_uint) -> c_int;
    fn af_lower(out: MutAfArray, arr: AfArray, is_unit_diag: c_int) -> c_int;
    fn af_upper(out: MutAfArray, arr: AfArray, is_unit_diag: c_int) -> c_int;
}

pub trait ConstGenerator {
    fn generate(&self, dims: Dim4) -> Array;
}

#[allow(unused_mut)]
impl ConstGenerator for i64 {
    fn generate(&self, dims: Dim4) -> Array {
        unsafe {
            let mut temp: i64 = 0;
            af_constant_long(&mut temp as MutAfArray, *self as Intl,
                             dims.ndims() as c_uint, dims.get().as_ptr() as *const DimT);
            Array::from(temp)
        }
    }
}

#[allow(unused_mut)]
impl ConstGenerator for u64 {
    fn generate(&self, dims: Dim4) -> Array {
        unsafe {
            let mut temp: i64 = 0;
            af_constant_ulong(&mut temp as MutAfArray, *self as Uintl,
                              dims.ndims() as c_uint, dims.get().as_ptr() as *const DimT);
            Array::from(temp)
        }
    }
}

#[allow(unused_mut)]
impl ConstGenerator for Complex<f32> {
    fn generate(&self, dims: Dim4) -> Array {
        unsafe {
            let mut temp: i64 = 0;
            af_constant_complex(&mut temp as MutAfArray,
                                (*self).re as c_double, (*self).im as c_double,
                                dims.ndims() as c_uint, dims.get().as_ptr() as *const DimT, 1);
            Array::from(temp)
        }
    }
}

#[allow(unused_mut)]
impl ConstGenerator for Complex<f64> {
    fn generate(&self, dims: Dim4) -> Array {
        unsafe {
            let mut temp: i64 = 0;
            af_constant_complex(&mut temp as MutAfArray,
                                (*self).re as c_double, (*self).im as c_double,
                                dims.ndims() as c_uint, dims.get().as_ptr() as *const DimT, 3);
            Array::from(temp)
        }
    }
}

macro_rules! cnst {
    ($rust_type:ty, $ffi_type:expr) => (
        #[allow(unused_mut)]
        impl ConstGenerator for $rust_type {
            fn generate(&self, dims: Dim4) -> Array {
                unsafe {
                    let mut temp: i64 = 0;
                    af_constant(&mut temp as MutAfArray, *self as u64 as c_double,
                                dims.ndims() as c_uint, dims.get().as_ptr() as *const DimT,
                                $ffi_type);
                    Array::from(temp)
                }
            }
        }
    )
}

cnst!(f32 , 0);
cnst!(f64 , 2);
cnst!(bool, 4);
cnst!(i32 , 5);
cnst!(u32 , 6);
cnst!(u8  , 7);

pub fn constant<T : ConstGenerator>(cnst: T, dims: Dim4) -> Array {
    cnst.generate(dims)
}

#[allow(unused_mut)]
pub fn range(dims: Dim4, seq_dim: i32, aftype: Aftype) -> Array {
    unsafe {
        let mut temp: i64 = 0;
        af_range(&mut temp as MutAfArray,
                 dims.ndims() as c_uint, dims.get().as_ptr() as *const DimT,
                 seq_dim as c_int, aftype as uint8_t);
        Array::from(temp)
    }
}

#[allow(unused_mut)]
pub fn iota(dims: Dim4, tdims: Dim4, aftype: Aftype) -> Array {
    unsafe {
        let mut temp: i64 = 0;
        af_iota(&mut temp as MutAfArray,
                dims.ndims() as c_uint, dims.get().as_ptr() as *const DimT,
                tdims.ndims() as c_uint, tdims.get().as_ptr() as *const DimT,
                aftype as uint8_t);
        Array::from(temp)
    }
}

pub fn set_seed(seed: u64) {
    unsafe { af_set_seed(seed as Uintl); }
}

#[allow(unused_mut)]
pub fn get_seed() -> u64 {
    unsafe {
        let mut temp: u64 = 0;
        af_get_seed(&mut temp as *mut Uintl);
        temp
    }
}

macro_rules! data_gen_def {
    ($fn_name:ident, $ffi_name: ident) => (
        #[allow(unused_mut)]
        pub fn $fn_name(dims: Dim4, aftype: Aftype) -> Array {
            unsafe {
                let mut temp: i64 = 0;
                $ffi_name(&mut temp as MutAfArray,
                          dims.ndims() as c_uint, dims.get().as_ptr() as *const DimT,
                          aftype as uint8_t);
                Array::from(temp)
            }
        }
    )
}

data_gen_def!(randu, af_randu);
data_gen_def!(randn, af_randn);
data_gen_def!(identity, af_identity);

#[allow(unused_mut)]
pub fn diag_create(input: &Array, dim: i32) -> Array {
    unsafe {
        let mut temp: i64 = 0;
        af_diag_create(&mut temp as MutAfArray, input.get() as AfArray, dim as c_int);
        Array::from(temp)
    }
}

#[allow(unused_mut)]
pub fn diag_extract(input: &Array, dim: i32) -> Array {
    unsafe {
        let mut temp: i64 = 0;
        af_diag_extract(&mut temp as MutAfArray, input.get() as AfArray, dim as c_int);
        Array::from(temp)
    }
}

#[allow(unused_mut)]
pub fn join(dim: i32, first: &Array, second: &Array) -> Array {
    unsafe {
        let mut temp: i64 = 0;
        af_join(&mut temp as MutAfArray, dim as c_int,
                first.get() as AfArray, second.get() as AfArray);
        Array::from(temp)
    }
}

#[allow(unused_mut)]
pub fn join_many(dim: i32, inputs: Vec<&Array>) -> Array {
    unsafe {
        let mut temp: i64 = 0;
        af_join_many(&mut temp as MutAfArray, dim as c_int,
                     inputs.len() as c_uint, inputs.as_ptr() as *const AfArray);
        Array::from(temp)
    }
}

// Using macro to implement join many wrapper
#[macro_export]
macro_rules! join_many {
    [$dim: expr; $($x:ident),+] => {
        {
            let mut temp_vec = Vec::new();
            $(
                temp_vec.push($x);
             )*
            let mut temp: i64 = 0;
            unsafe {
                let mut temp: i64 = 0;
                af_join_many(&mut temp as MutAfArray, $dim as c_int,
                             temp_vec.len() as c_uint, temp_vec.as_ptr() as *const AfArray);
                Array::from(temp)
            }
        }
    };
}

macro_rules! data_func_def {
    ($fn_name:ident, $ffi_name: ident) => (
        #[allow(unused_mut)]
        pub fn $fn_name(input: &Array, dims: Dim4) -> Array {
            unsafe {
                let mut temp: i64 = 0;
                $ffi_name(&mut temp as MutAfArray, input.get() as AfArray,
                          dims[0] as c_uint, dims[1] as c_uint,
                          dims[2] as c_uint, dims[3] as c_uint);
                Array::from(temp)
            }
        }
    )
}

data_func_def!(tile, af_tile);
data_func_def!(reorder, af_reorder);
data_func_def!(shift, af_shift);

#[allow(unused_mut)]
pub fn moddims(input: &Array, dims: Dim4) -> Array {
    unsafe {
        let mut temp: i64 = 0;
        af_moddims(&mut temp as MutAfArray, input.get() as AfArray,
                   dims.ndims() as c_uint, dims.get().as_ptr() as *const DimT);
        Array::from(temp)
    }
}

#[allow(unused_mut)]
pub fn flat(input: &Array) -> Array {
    unsafe {
        let mut temp: i64 = 0;
        af_flat(&mut temp as MutAfArray, input.get() as AfArray);
        Array::from(temp)
    }
}

#[allow(unused_mut)]
pub fn flip(input: &Array, dim: u32) -> Array {
    unsafe {
        let mut temp: i64 = 0;
        af_flip(&mut temp as MutAfArray, input.get() as AfArray, dim as c_uint);
        Array::from(temp)
    }
}

#[allow(unused_mut)]
pub fn lower(input: &Array, is_unit_diag: bool) -> Array {
    unsafe {
        let mut temp: i64 = 0;
        af_lower(&mut temp as MutAfArray, input.get() as AfArray, is_unit_diag as c_int);
        Array::from(temp)
    }
}

#[allow(unused_mut)]
pub fn upper(input: &Array, is_unit_diag: bool) -> Array {
    unsafe {
        let mut temp: i64 = 0;
        af_upper(&mut temp as MutAfArray, input.get() as AfArray, is_unit_diag as c_int);
        Array::from(temp)
    }
}
<commit_msg>specialize bool<commit_after>extern crate libc;
extern crate num;

use array::Array;
use dim4::Dim4;
use defines::Aftype;
use self::libc::{uint8_t, c_int, c_uint, c_double};
use self::num::Complex;

use std::vec::Vec;

type MutAfArray = *mut self::libc::c_longlong;
type MutDouble  = *mut self::libc::c_double;
type MutUint    = *mut self::libc::c_uint;
type AfArray    = self::libc::c_longlong;
type DimT       = self::libc::c_longlong;
type Intl       = self::libc::c_longlong;
type Uintl      = self::libc::c_ulonglong;

#[allow(dead_code)]
extern {
    fn af_constant(out: MutAfArray, val: c_double,
                   ndims: c_uint, dims: *const DimT, afdtype: c_int) -> c_int;

    fn af_constant_complex(out: MutAfArray, real: c_double, imag: c_double,
                           ndims: c_uint, dims: *const DimT, afdtype: c_int) -> c_int;

    fn af_constant_long(out: MutAfArray, val: Intl,
                        ndims: c_uint, dims: *const DimT) -> c_int;

    fn af_constant_ulong(out: MutAfArray, val: Uintl,
                         ndims: c_uint, dims: *const DimT) -> c_int;

    fn af_range(out: MutAfArray, ndims: c_uint, dims: *const DimT,
                seq_dims: c_int, afdtype: uint8_t) -> c_int;

    fn af_iota(out: MutAfArray, ndims: c_uint, dims: *const DimT,
               t_ndims: c_uint, tdims: *const DimT, afdtype: uint8_t) -> c_int;

    fn af_randu(out: MutAfArray, ndims: c_uint, dims: *const DimT, afdtype: uint8_t) -> c_int;
    fn af_randn(out: MutAfArray, ndims: c_uint, dims: *const DimT, afdtype: uint8_t) -> c_int;

    fn af_set_seed(seed: Uintl);
    fn af_get_seed(seed: *mut Uintl);

    fn af_identity(out: MutAfArray, ndims: c_uint, dims: *const DimT, afdtype: uint8_t) -> c_int;
    fn af_diag_create(out: MutAfArray, arr: AfArray, num: c_int) -> c_int;
    fn af_diag_extract(out: MutAfArray, arr: AfArray, num: c_int) -> c_int;
    fn af_join(out: MutAfArray, dim: c_int, first: AfArray, second: AfArray) -> c_int;
    fn af_join_many(out: MutAfArray, dim: c_int, n_arrays: c_uint, inpts: *const AfArray) -> c_int;

    fn af_tile(out: MutAfArray, arr: AfArray, x: c_uint, y: c_uint, z: c_uint, w: c_uint) -> c_int;
    fn af_reorder(o: MutAfArray, a: AfArray, x: c_uint, y: c_uint, z: c_uint, w: c_uint) -> c_int;
    fn af_shift(o: MutAfArray, a: AfArray, x: c_uint, y: c_uint, z: c_uint, w: c_uint) -> c_int;
    fn af_moddims(out: MutAfArray, arr: AfArray, ndims: c_uint, dims: *const DimT) -> c_int;

    fn af_flat(out: MutAfArray, arr: AfArray) -> c_int;
    fn af_flip(out: MutAfArray, arr: AfArray, dim: c_uint) -> c_int;
    fn af_lower(out: MutAfArray, arr: AfArray, is_unit_diag: c_int) -> c_int;
    fn af_upper(out: MutAfArray, arr: AfArray, is_unit_diag: c_int) -> c_int;
}

pub trait ConstGenerator {
    fn generate(&self, dims: Dim4) -> Array;
}

#[allow(unused_mut)]
impl ConstGenerator for i64 {
    fn generate(&self, dims: Dim4) -> Array {
        unsafe {
            let mut temp: i64 = 0;
            af_constant_long(&mut temp as MutAfArray, *self as Intl,
                             dims.ndims() as c_uint, dims.get().as_ptr() as *const DimT);
            Array::from(temp)
        }
    }
}

#[allow(unused_mut)]
impl ConstGenerator for u64 {
    fn generate(&self, dims: Dim4) -> Array {
        unsafe {
            let mut temp: i64 = 0;
            af_constant_ulong(&mut temp as MutAfArray, *self as Uintl,
                              dims.ndims() as c_uint, dims.get().as_ptr() as *const DimT);
            Array::from(temp)
        }
    }
}

#[allow(unused_mut)]
impl ConstGenerator for Complex<f32> {
    fn generate(&self, dims: Dim4) -> Array {
        unsafe {
            let mut temp: i64 = 0;
            af_constant_complex(&mut temp as MutAfArray,
                                (*self).re as c_double, (*self).im as c_double,
                                dims.ndims() as c_uint, dims.get().as_ptr() as *const DimT, 1);
            Array::from(temp)
        }
    }
}

#[allow(unused_mut)]
impl ConstGenerator for Complex<f64> {
    fn generate(&self, dims: Dim4) -> Array {
        unsafe {
            let mut temp: i64 = 0;
            af_constant_complex(&mut temp as MutAfArray,
                                (*self).re as c_double, (*self).im as c_double,
                                dims.ndims() as c_uint, dims.get().as_ptr() as *const DimT, 3);
            Array::from(temp)
        }
    }
}

#[allow(unused_mut)]
impl ConstGenerator for bool {
    fn generate(&self, dims: Dim4) -> Array {
        unsafe {
            let mut temp: i64 = 0;
            af_constant(&mut temp as MutAfArray, *self as c_int as c_double,
                              dims.ndims() as c_uint, dims.get().as_ptr() as *const DimT, 4);
            Array::from(temp)
        }
    }
}

macro_rules! cnst {
    ($rust_type:ty, $ffi_type:expr) => (
        #[allow(unused_mut)]
        impl ConstGenerator for $rust_type {
            fn generate(&self, dims: Dim4) -> Array {
                unsafe {
                    let mut temp: i64 = 0;
                    af_constant(&mut temp as MutAfArray, *self as c_double,
                                dims.ndims() as c_uint, dims.get().as_ptr() as *const DimT,
                                $ffi_type);
                    Array::from(temp)
                }
            }
        }
    )
}

cnst!(f32 , 0);
cnst!(f64 , 2);
cnst!(i32 , 5);
cnst!(u32 , 6);
cnst!(u8  , 7);

pub fn constant<T : ConstGenerator>(cnst: T, dims: Dim4) -> Array {
    cnst.generate(dims)
}

#[allow(unused_mut)]
pub fn range(dims: Dim4, seq_dim: i32, aftype: Aftype) -> Array {
    unsafe {
        let mut temp: i64 = 0;
        af_range(&mut temp as MutAfArray,
                 dims.ndims() as c_uint, dims.get().as_ptr() as *const DimT,
                 seq_dim as c_int, aftype as uint8_t);
        Array::from(temp)
    }
}

#[allow(unused_mut)]
pub fn iota(dims: Dim4, tdims: Dim4, aftype: Aftype) -> Array {
    unsafe {
        let mut temp: i64 = 0;
        af_iota(&mut temp as MutAfArray,
                dims.ndims() as c_uint, dims.get().as_ptr() as *const DimT,
                tdims.ndims() as c_uint, tdims.get().as_ptr() as *const DimT,
                aftype as uint8_t);
        Array::from(temp)
    }
}

pub fn set_seed(seed: u64) {
    unsafe { af_set_seed(seed as Uintl); }
}

#[allow(unused_mut)]
pub fn get_seed() -> u64 {
    unsafe {
        let mut temp: u64 = 0;
        af_get_seed(&mut temp as *mut Uintl);
        temp
    }
}

macro_rules! data_gen_def {
    ($fn_name:ident, $ffi_name: ident) => (
        #[allow(unused_mut)]
        pub fn $fn_name(dims: Dim4, aftype: Aftype) -> Array {
            unsafe {
                let mut temp: i64 = 0;
                $ffi_name(&mut temp as MutAfArray,
                          dims.ndims() as c_uint, dims.get().as_ptr() as *const DimT,
                          aftype as uint8_t);
                Array::from(temp)
            }
        }
    )
}

data_gen_def!(randu, af_randu);
data_gen_def!(randn, af_randn);
data_gen_def!(identity, af_identity);

#[allow(unused_mut)]
pub fn diag_create(input: &Array, dim: i32) -> Array {
    unsafe {
        let mut temp: i64 = 0;
        af_diag_create(&mut temp as MutAfArray, input.get() as AfArray, dim as c_int);
        Array::from(temp)
    }
}

#[allow(unused_mut)]
pub fn diag_extract(input: &Array, dim: i32) -> Array {
    unsafe {
        let mut temp: i64 = 0;
        af_diag_extract(&mut temp as MutAfArray, input.get() as AfArray, dim as c_int);
        Array::from(temp)
    }
}

#[allow(unused_mut)]
pub fn join(dim: i32, first: &Array, second: &Array) -> Array {
    unsafe {
        let mut temp: i64 = 0;
        af_join(&mut temp as MutAfArray, dim as c_int,
                first.get() as AfArray, second.get() as AfArray);
        Array::from(temp)
    }
}

#[allow(unused_mut)]
pub fn join_many(dim: i32, inputs: Vec<&Array>) -> Array {
    unsafe {
        let mut temp: i64 = 0;
        af_join_many(&mut temp as MutAfArray, dim as c_int,
                     inputs.len() as c_uint, inputs.as_ptr() as *const AfArray);
        Array::from(temp)
    }
}

// Using macro to implement join many wrapper
#[macro_export]
macro_rules! join_many {
    [$dim: expr; $($x:ident),+] => {
        {
            let mut temp_vec = Vec::new();
            $(
                temp_vec.push($x);
             )*
            let mut temp: i64 = 0;
            unsafe {
                let mut temp: i64 = 0;
                af_join_many(&mut temp as MutAfArray, $dim as c_int,
                             temp_vec.len() as c_uint, temp_vec.as_ptr() as *const AfArray);
                Array::from(temp)
            }
        }
    };
}

macro_rules! data_func_def {
    ($fn_name:ident, $ffi_name: ident) => (
        #[allow(unused_mut)]
        pub fn $fn_name(input: &Array, dims: Dim4) -> Array {
            unsafe {
                let mut temp: i64 = 0;
                $ffi_name(&mut temp as MutAfArray, input.get() as AfArray,
                          dims[0] as c_uint, dims[1] as c_uint,
                          dims[2] as c_uint, dims[3] as c_uint);
                Array::from(temp)
            }
        }
    )
}

data_func_def!(tile, af_tile);
data_func_def!(reorder, af_reorder);
data_func_def!(shift, af_shift);

#[allow(unused_mut)]
pub fn moddims(input: &Array, dims: Dim4) -> Array {
    unsafe {
        let mut temp: i64 = 0;
        af_moddims(&mut temp as MutAfArray, input.get() as AfArray,
                   dims.ndims() as c_uint, dims.get().as_ptr() as *const DimT);
        Array::from(temp)
    }
}

#[allow(unused_mut)]
pub fn flat(input: &Array) -> Array {
    unsafe {
        let mut temp: i64 = 0;
        af_flat(&mut temp as MutAfArray, input.get() as AfArray);
        Array::from(temp)
    }
}

#[allow(unused_mut)]
pub fn flip(input: &Array, dim: u32) -> Array {
    unsafe {
        let mut temp: i64 = 0;
        af_flip(&mut temp as MutAfArray, input.get() as AfArray, dim as c_uint);
        Array::from(temp)
    }
}

#[allow(unused_mut)]
pub fn lower(input: &Array, is_unit_diag: bool) -> Array {
    unsafe {
        let mut temp: i64 = 0;
        af_lower(&mut temp as MutAfArray, input.get() as AfArray, is_unit_diag as c_int);
        Array::from(temp)
    }
}

#[allow(unused_mut)]
pub fn upper(input: &Array, is_unit_diag: bool) -> Array {
    unsafe {
        let mut temp: i64 = 0;
        af_upper(&mut temp as MutAfArray, input.get() as AfArray, is_unit_diag as c_int);
        Array::from(temp)
    }
}
<|endoftext|>


<commit_before>use syntax::util::small_vector::SmallVector;

use syntax::{abi, ast, ast_util, codemap, parse};
use syntax::ptr::P;
use syntax::ext::base;
use syntax::ext::build::AstBuilder;

static BEFORE_EACH: &'static str = "before_each";
static AFTER_EACH:  &'static str = "after_each";
static BEFORE:      &'static str = "before";
static AFTER:       &'static str = "after";
static IT:          &'static str = "it";

/// Defines the state of a `describe!` macro as it is parsing.
struct DescribeState {
    name: Option<String>,
    before: Option<P<ast::Block>>,
    after: Option<P<ast::Block>>,
    before_each: Option<P<ast::Block>>,
    after_each: Option<P<ast::Block>>,
    tests: Vec<Test>
}

/// A test as a description and associated block.
struct Test {
    description: String,
    block: P<ast::Block>
}

/// Defines the overarching `describe!` syntax extension.
///
/// All other macros in stainless are actually "fake" in the sense
/// that they are detected and expanded inside of the implementation
/// of `describe!`.
pub fn describe(cx: &mut base::ExtCtxt, _: codemap::Span, tokens: &[ast::TokenTree]) -> Box<base::MacResult + 'static> {
    // Parse a full DescribeState from the input, emitting errors if used incorrectly.
    let state = parse_describe(DescribeState {
        name: None, before: None, after: None,
        before_each: None, after_each: None, tests: vec![]
    }, parse::tts_to_parser(cx.parse_sess(), tokens.to_vec(), cx.cfg()));

    let name = cx.ident_of(state.name.clone().unwrap().replace(" ", "_").as_slice());

    // Create tests from a full DescribeState
    let tests = create_tests(state, cx);
    let super_glob = cx.view_use_glob(codemap::DUMMY_SP, ast::Inherited, vec![cx.ident_of("super")]);

    let test_mod = cx.item_mod(codemap::DUMMY_SP, codemap::DUMMY_SP, name, vec![], vec![super_glob], tests);
    box MacItems { items: vec![test_mod] }
}

fn create_tests(state: DescribeState, cx: &mut base::ExtCtxt) -> Vec<P<ast::Item>> {
    // FIXME(reem): Implement before and after.
    let (_before, _after) = (state.before, state.after);

    let (before_each, after_each) = (state.before_each, state.after_each);
    let tests = state.tests;

    let test_attribute = cx.attribute(codemap::DUMMY_SP,
                                      cx.meta_word(codemap::DUMMY_SP, parse::token::InternedString::new("test")));

    tests.into_iter().map(|Test { description, block }| {
        let test_body = match (&before_each, &after_each) {
            (&None, &None) => block,

            (&Some(ref before), &None) => {
                P(ast::Block {
                    view_items: before.view_items + block.view_items,
                    stmts: before.stmts + block.stmts,
                    ..block.deref().clone()
                })
            },

            (&None, &Some(ref after)) => {
                P(ast::Block {
                    view_items: block.view_items + after.view_items,
                    stmts: block.stmts + after.stmts,
                    ..block.deref().clone()
                })
            },

            (&Some(ref before), &Some(ref after)) => {
                P(ast::Block {
                    view_items: before.view_items + block.view_items + after.view_items,
                    stmts: before.stmts + block.stmts + after.stmts,
                    ..block.deref().clone()
                })
            }
        };

        P(ast::Item {
            ident: cx.ident_of(description.replace(" ", "_").as_slice()),
            attrs: vec![test_attribute.clone()],
            id: ast::DUMMY_NODE_ID,
            node: ast::ItemFn(
                cx.fn_decl(vec![], cx.ty_nil()),
                ast::NormalFn,
                abi::Rust,
                ast_util::empty_generics(),
                test_body
            ),
            vis: ast::Inherited,
            span: codemap::DUMMY_SP
        })
    }).collect()
}

fn parse_describe(mut state: DescribeState, mut parser: parse::parser::Parser) -> DescribeState {
    // First parse the name of this describe block:
    let (name, _) = parser.parse_str();
    state.name = Some(name.get().to_string());

    // Move past the opening {
    if parse::token::LBRACE != parser.bump_and_get() {
        parser.fatal("Expected { after the description of a describe! block.");
    }

    // Now parse all tests and subsections:
    while parser.token != parse::token::RBRACE {
        // Get the name of this block, must be either:
        //     - before_each
        //     - after_each
        //     - before
        //     - after
        //     - it
        //
        // Any other top-level idents are not allowed.
        let block_name = parser.parse_ident();

        match block_name.as_str() {
            BEFORE_EACH => {
                if state.before_each.is_some() { parser.fatal("Only one `before_each` block is allowed per `describe!` block.") }
                state.before_each = Some(parser.parse_block());
            },

            AFTER_EACH => {
                if state.after_each.is_some() { parser.fatal("Only one `after_each` block is allowed per `describe!` block.") }
                state.after_each = Some(parser.parse_block());
            },

            BEFORE => {
                if state.before.is_some() { parser.fatal("Only one `before` block is allowed per `describe!` block.") }
                state.before = Some(parser.parse_block());
            },

            AFTER => {
                if state.after.is_some() { parser.fatal("Only one `after` block is allowed per `describe!` block.") }
                state.after = Some(parser.parse_block());
            },

            IT => {
                // Description of this `it` block.
                let (description, _) = parser.parse_str();

                state.tests.push(Test {
                    // Get as a String
                    description: description.get().to_string(),

                    // The associated block
                    block: parser.parse_block()
                });
            },

            banned => {
                // Illegal block name.
                let span = parser.span;
                parser.span_fatal(span, format!("Expected one of: `{}`, but found: `{}`",
                    format!("{}, {}, {}, {}, {}", BEFORE_EACH, AFTER_EACH, BEFORE, AFTER, IT).as_slice(),
                    banned).as_slice());
            }
        }
    }

    state
}

struct MacItems { items: Vec<P<ast::Item>> }

impl base::MacResult for MacItems {
    fn make_items(self: Box<MacItems>) -> Option<SmallVector<P<ast::Item>>> {
        Some(SmallVector::many(self.items.clone()))
    }
}

<commit_msg>(cleanup) Use builtin MacItems instead of homegrown version.<commit_after>use syntax::{abi, ast, ast_util, codemap, parse};
use syntax::ptr::P;
use syntax::ext::base;
use syntax::ext::build::AstBuilder;

static BEFORE_EACH: &'static str = "before_each";
static AFTER_EACH:  &'static str = "after_each";
static BEFORE:      &'static str = "before";
static AFTER:       &'static str = "after";
static IT:          &'static str = "it";

/// Defines the state of a `describe!` macro as it is parsing.
struct DescribeState {
    name: Option<String>,
    before: Option<P<ast::Block>>,
    after: Option<P<ast::Block>>,
    before_each: Option<P<ast::Block>>,
    after_each: Option<P<ast::Block>>,
    tests: Vec<Test>
}

/// A test as a description and associated block.
struct Test {
    description: String,
    block: P<ast::Block>
}

/// Defines the overarching `describe!` syntax extension.
///
/// All other macros in stainless are actually "fake" in the sense
/// that they are detected and expanded inside of the implementation
/// of `describe!`.
pub fn describe(cx: &mut base::ExtCtxt, _: codemap::Span, tokens: &[ast::TokenTree]) -> Box<base::MacResult + 'static> {
    // Parse a full DescribeState from the input, emitting errors if used incorrectly.
    let state = parse_describe(DescribeState {
        name: None, before: None, after: None,
        before_each: None, after_each: None, tests: vec![]
    }, parse::tts_to_parser(cx.parse_sess(), tokens.to_vec(), cx.cfg()));

    let name = cx.ident_of(state.name.clone().unwrap().replace(" ", "_").as_slice());

    // Create tests from a full DescribeState
    let tests = create_tests(state, cx);

    // Get a glob import of all items in scope to the module that `describe!` is called from.
    let super_glob = cx.view_use_glob(codemap::DUMMY_SP, ast::Inherited, vec![cx.ident_of("super")]);

    // Generate the new module.
    let test_mod = cx.item_mod(codemap::DUMMY_SP, codemap::DUMMY_SP, name, vec![], vec![super_glob], tests);

    // Export the new module.
    base::MacItems::new(vec![test_mod].into_iter())
}

fn create_tests(state: DescribeState, cx: &mut base::ExtCtxt) -> Vec<P<ast::Item>> {
    // FIXME(reem): Implement before and after.
    let (_before, _after) = (state.before, state.after);

    let (before_each, after_each) = (state.before_each, state.after_each);
    let tests = state.tests;

    let test_attribute = cx.attribute(codemap::DUMMY_SP,
                                      cx.meta_word(codemap::DUMMY_SP, parse::token::InternedString::new("test")));

    tests.into_iter().map(|Test { description, block }| {
        let test_body = match (&before_each, &after_each) {
            (&None, &None) => block,

            (&Some(ref before), &None) => {
                P(ast::Block {
                    view_items: before.view_items + block.view_items,
                    stmts: before.stmts + block.stmts,
                    ..block.deref().clone()
                })
            },

            (&None, &Some(ref after)) => {
                P(ast::Block {
                    view_items: block.view_items + after.view_items,
                    stmts: block.stmts + after.stmts,
                    ..block.deref().clone()
                })
            },

            (&Some(ref before), &Some(ref after)) => {
                P(ast::Block {
                    view_items: before.view_items + block.view_items + after.view_items,
                    stmts: before.stmts + block.stmts + after.stmts,
                    ..block.deref().clone()
                })
            }
        };

        P(ast::Item {
            ident: cx.ident_of(description.replace(" ", "_").as_slice()),
            attrs: vec![test_attribute.clone()],
            id: ast::DUMMY_NODE_ID,
            node: ast::ItemFn(
                cx.fn_decl(vec![], cx.ty_nil()),
                ast::NormalFn,
                abi::Rust,
                ast_util::empty_generics(),
                test_body
            ),
            vis: ast::Inherited,
            span: codemap::DUMMY_SP
        })
    }).collect()
}

fn parse_describe(mut state: DescribeState, mut parser: parse::parser::Parser) -> DescribeState {
    // First parse the name of this describe block:
    let (name, _) = parser.parse_str();
    state.name = Some(name.get().to_string());

    // Move past the opening {
    if parse::token::LBRACE != parser.bump_and_get() {
        parser.fatal("Expected { after the description of a describe! block.");
    }

    // Now parse all tests and subsections:
    while parser.token != parse::token::RBRACE {
        // Get the name of this block, must be either:
        //     - before_each
        //     - after_each
        //     - before
        //     - after
        //     - it
        //
        // Any other top-level idents are not allowed.
        let block_name = parser.parse_ident();

        match block_name.as_str() {
            BEFORE_EACH => {
                if state.before_each.is_some() { parser.fatal("Only one `before_each` block is allowed per `describe!` block.") }
                state.before_each = Some(parser.parse_block());
            },

            AFTER_EACH => {
                if state.after_each.is_some() { parser.fatal("Only one `after_each` block is allowed per `describe!` block.") }
                state.after_each = Some(parser.parse_block());
            },

            BEFORE => {
                if state.before.is_some() { parser.fatal("Only one `before` block is allowed per `describe!` block.") }
                state.before = Some(parser.parse_block());
            },

            AFTER => {
                if state.after.is_some() { parser.fatal("Only one `after` block is allowed per `describe!` block.") }
                state.after = Some(parser.parse_block());
            },

            IT => {
                // Description of this `it` block.
                let (description, _) = parser.parse_str();

                state.tests.push(Test {
                    // Get as a String
                    description: description.get().to_string(),

                    // The associated block
                    block: parser.parse_block()
                });
            },

            banned => {
                // Illegal block name.
                let span = parser.span;
                parser.span_fatal(span, format!("Expected one of: `{}`, but found: `{}`",
                    format!("{}, {}, {}, {}, {}", BEFORE_EACH, AFTER_EACH, BEFORE, AFTER, IT).as_slice(),
                    banned).as_slice());
            }
        }
    }

    state
}

<|endoftext|>


<commit_before>use std::collections::HashMap;
use std::fs::{self, File};
use std::path::{Path, PathBuf};

use mustache;

use error::{Error, Result};
use parser;

pub type Header = HashMap<String, String>;

#[derive(PartialEq, Debug)]
pub struct Document {
    data: Header,
    content: String,
}

impl Document {
    pub fn new(header: Header, content: &str) -> Document {
        Document { data: header, content: content.to_string() }
    }

    pub fn as_html(&self) -> Result<String> {
        let template = mustache::compile_str(&self.content);

        // Write the template to memory, then retrieve it as a string.
        let mut buf = Vec::<u8>::new();
        try!(template.render(&mut buf, &self.data));

        String::from_utf8(buf).map_err(Error::from)
    }

    pub fn render_to_file(&self, file_path: &Path, templates: &HashMap<String, String>) -> Result<()> {
        let template_path = try!(self.template());
        let template = try!(templates.get(&template_path.to_string())
            .ok_or(Error::TemplateNotFound));

        let parent_path = file_path.parent().ok_or(Error::missing_parent_path(&file_path));
        try!(fs::create_dir_all(&try!(parent_path)));

        let mut file = try!(File::create(file_path));
        let mut data = HashMap::new();

        data.insert("content", try!(self.as_html()));

        for (key, value) in self.data.iter() {
            data.insert(&key, value.clone());
        }

        let template = mustache::compile_str(&template);

        try!(template.render(&mut file, &data));

        info!("Created {}", file_path.display());
        Ok(())
    }

    fn template(&self) -> Result<&String> {
        self.data.get(&"template".to_string())
            .ok_or(Error::MissingTemplateField)
    }
}

pub fn load_documents_from_disk<F>(documents_path: &Path, mut criteria: F) -> Result<HashMap<PathBuf, Document>>
    where F : FnMut(&Path) -> bool
{
    use util;

    let mut documents = HashMap::new();

    try!(util::walk_dir(
        documents_path,
        &mut |path| {
            if !criteria(&path) || !path.is_file() {
                return Ok(());
            }

            // Read the document from the disk.
            let content = {
                use std::io::Read;

                let mut buf = String::new();
                let mut file = try!(File::open(&path));
                try!(file.read_to_string(&mut buf));
                buf
            };
            let document = try!(parser::parse_document(&content));
            let relative_path = try!(path.strip_prefix(documents_path));

            documents.insert(relative_path.with_extension("").to_path_buf(), document);

            Ok(())
        },
    ));

    Ok(documents)
}
<commit_msg>chore: replace all usages of `to_string`<commit_after>use std::collections::HashMap;
use std::fs::{self, File};
use std::path::{Path, PathBuf};

use mustache;

use error::{Error, Result};
use parser;

pub type Header = HashMap<String, String>;

#[derive(PartialEq, Debug)]
pub struct Document {
    data: Header,
    content: String,
}

impl Document {
    pub fn new(header: Header, content: &str) -> Document {
        Document { data: header, content: content.to_owned() }
    }

    pub fn as_html(&self) -> Result<String> {
        let template = mustache::compile_str(&self.content);

        // Write the template to memory, then retrieve it as a string.
        let mut buf = Vec::<u8>::new();
        try!(template.render(&mut buf, &self.data));

        String::from_utf8(buf).map_err(Error::from)
    }

    pub fn render_to_file(&self, file_path: &Path, templates: &HashMap<String, String>) -> Result<()> {
        let template_path = try!(self.template());
        let template = try!(templates.get(&template_path.to_owned())
            .ok_or(Error::TemplateNotFound));

        let parent_path = file_path.parent().ok_or(Error::missing_parent_path(&file_path));
        try!(fs::create_dir_all(&try!(parent_path)));

        let mut file = try!(File::create(file_path));
        let mut data = HashMap::new();

        data.insert("content", try!(self.as_html()));

        for (key, value) in self.data.iter() {
            data.insert(&key, value.clone());
        }

        let template = mustache::compile_str(&template);

        try!(template.render(&mut file, &data));

        info!("Created {}", file_path.display());
        Ok(())
    }

    fn template(&self) -> Result<&String> {
        self.data.get(&"template".to_owned())
            .ok_or(Error::MissingTemplateField)
    }
}

pub fn load_documents_from_disk<F>(documents_path: &Path, mut criteria: F) -> Result<HashMap<PathBuf, Document>>
    where F : FnMut(&Path) -> bool
{
    use util;

    let mut documents = HashMap::new();

    try!(util::walk_dir(
        documents_path,
        &mut |path| {
            if !criteria(&path) || !path.is_file() {
                return Ok(());
            }

            // Read the document from the disk.
            let content = {
                use std::io::Read;

                let mut buf = String::new();
                let mut file = try!(File::open(&path));
                try!(file.read_to_string(&mut buf));
                buf
            };
            let document = try!(parser::parse_document(&content));
            let relative_path = try!(path.strip_prefix(documents_path));

            documents.insert(relative_path.with_extension("").to_path_buf(), document);

            Ok(())
        },
    ));

    Ok(documents)
}
<|endoftext|>


<commit_before>use std::collections::BTreeMap;
use xref::{Xref, XrefEntry};
use super::{Object, ObjectId, Dictionary};
use object_stream::ObjectStream;
use byref::ByRef;

/// PDF document.
#[derive(Debug)]
pub struct Document {
	/// The version of the PDF specification to which the file conforms.
	pub version: String,

	/// The trailer gives the location of the cross-reference table and of certain special objects.
	pub trailer: Dictionary,

	/// The cross-reference table contains locations of the indirect objects.
	pub reference_table: Xref,

	/// The objects that make up the document contained in the file.
	pub objects: BTreeMap<ObjectId, Object>,

	/// The object streams which contains compressed objects.
	pub streams: BTreeMap<u32, ObjectStream>,

	/// Current maximum object id within the document.
	pub max_id: u32,
}

impl Document {

	/// Create new PDF document.
	pub fn new() -> Document {
		Document {
			version: "1.4".to_string(),
			trailer: Dictionary::new(),
			reference_table: Xref::new(0),
			objects: BTreeMap::new(),
			streams: BTreeMap::new(),
			max_id: 0,
		}
	}

	/// Create new PDF document.
	pub fn with_version<S: Into<String>>(version: S) -> Document {
		let mut document = Self::new();
		document.version = version.into();
		document
	}

	/// Get object by object id, will recursively dereference a referenced object.
	pub fn get_object(&self, id: ObjectId) -> Option<&Object> {
		if let Some(entry) = self.reference_table.get(id.0) {
			match *entry {
				XrefEntry::Normal { .. } => {
					if let Some(object) = self.objects.get(&id) {
						if let Some(id) = object.as_reference() {
							return self.get_object(id);
						} else {
							return Some(object);
						}
					}
				}
				XrefEntry::Compressed { container, index } => {
					if let Some(stream) = self.streams.get(&container) {
						if let Some(&(_id, ref object)) = stream.get_object(index as usize) {
							return Some(object);
						}
					}
				}
				_ => {},
			}
		}
		return None;
	}

	/// Traverse objects from trailer recursively, return all referenced object IDs.
	pub fn traverse_objects<A: Fn(&mut Object) -> ()>(&mut self, action: A) -> Vec<ObjectId> {
		fn traverse_array<A: Fn(&mut Object) -> ()>(array: &mut Vec<Object>, action: &A, refs: &mut Vec<ObjectId>) {
			for item in array.iter_mut() {
				traverse_object(item, action, refs);
			}
		}
		fn traverse_dictionary<A: Fn(&mut Object) -> ()>(dict: &mut Dictionary, action: &A, refs: &mut Vec<ObjectId>) {
			for (_, v) in dict.iter_mut() {
				traverse_object(v, action, refs);
			}
		}
		fn traverse_object<A: Fn(&mut Object) -> ()>(object: &mut Object, action: &A, refs: &mut Vec<ObjectId>) {
			action(object);
			match *object {
				Object::Array(ref mut array) => traverse_array(array, action, refs),
				Object::Dictionary(ref mut dict) => traverse_dictionary(dict, action, refs),
				Object::Stream(ref mut stream) => traverse_dictionary(&mut stream.dict, action, refs),
				Object::Reference(id) => {
					if !refs.contains(&id) {
						refs.push(id);
					}
				},
				_ => {}
			}
		}
		let mut refs = vec![];
		traverse_dictionary(&mut self.trailer, &action, &mut refs);
		let mut index = 0;
		while index < refs.len() {
			if let Some(object) = self.objects.get_mut(&refs[index]) {
				traverse_object(object, &action, &mut refs);
			}
			index += 1;
		}
		refs
	}

	/// Get catalog dictionary.
	pub fn catalog(&self) -> Option<&Dictionary> {
		self.trailer.get("Root").get_dict_by_ref(self)
	}

	/// Get page numbers and corresponding object ids.
	pub fn get_pages(&self) -> BTreeMap<u32, ObjectId> {
		fn collect_pages(doc: &Document, page_tree_id: ObjectId, page_number: &mut u32, pages: &mut BTreeMap<u32, ObjectId>) {
			if let Some(kids) = doc.get_object(page_tree_id).and_then(|obj|obj.as_dict()).and_then(|page_tree|page_tree.get("Kids")).and_then(|obj|obj.as_array()) {
				for kid in kids {
					if let Some(kid_id) = kid.as_reference() {
						if let Some(type_name) = doc.get_object(kid_id).and_then(|obj|obj.as_dict()).and_then(|dict|dict.type_name()) {
							match type_name {
								"Page" => {
									pages.insert(*page_number, kid_id);
									*page_number += 1;
								}
								"Pages" => {
									collect_pages(doc, kid_id, page_number, pages);
								}
								_ => {}
							}
						}
					}
				}
			}
		}

		let mut pages = BTreeMap::new();
		let mut page_number = 1;
		if let Some(page_tree_id) = self.catalog().and_then(|cat|cat.get("Pages")).and_then(|pages|pages.as_reference()) {
			collect_pages(self, page_tree_id, &mut page_number, &mut pages);
		}
		pages
	}
}
<commit_msg>Fix get_object for created document<commit_after>use std::collections::BTreeMap;
use xref::{Xref, XrefEntry};
use super::{Object, ObjectId, Dictionary};
use object_stream::ObjectStream;
use byref::ByRef;

/// PDF document.
#[derive(Debug)]
pub struct Document {
	/// The version of the PDF specification to which the file conforms.
	pub version: String,

	/// The trailer gives the location of the cross-reference table and of certain special objects.
	pub trailer: Dictionary,

	/// The cross-reference table contains locations of the indirect objects.
	pub reference_table: Xref,

	/// The objects that make up the document contained in the file.
	pub objects: BTreeMap<ObjectId, Object>,

	/// The object streams which contains compressed objects.
	pub streams: BTreeMap<u32, ObjectStream>,

	/// Current maximum object id within the document.
	pub max_id: u32,
}

impl Document {

	/// Create new PDF document.
	pub fn new() -> Document {
		Document {
			version: "1.4".to_string(),
			trailer: Dictionary::new(),
			reference_table: Xref::new(0),
			objects: BTreeMap::new(),
			streams: BTreeMap::new(),
			max_id: 0,
		}
	}

	/// Create new PDF document.
	pub fn with_version<S: Into<String>>(version: S) -> Document {
		let mut document = Self::new();
		document.version = version.into();
		document
	}

	/// Get object by object id, will recursively dereference a referenced object.
	pub fn get_object(&self, id: ObjectId) -> Option<&Object> {
		if let Some(object) = self.objects.get(&id) {
			return Some(object);
		}
		if let Some(entry) = self.reference_table.get(id.0) {
			match *entry {
				XrefEntry::Normal { .. } => {
					if let Some(object) = self.objects.get(&id) {
						if let Some(id) = object.as_reference() {
							return self.get_object(id);
						} else {
							return Some(object);
						}
					}
				}
				XrefEntry::Compressed { container, index } => {
					if let Some(stream) = self.streams.get(&container) {
						if let Some(&(_id, ref object)) = stream.get_object(index as usize) {
							return Some(object);
						}
					}
				}
				_ => {},
			}
		}
		return None;
	}

	/// Traverse objects from trailer recursively, return all referenced object IDs.
	pub fn traverse_objects<A: Fn(&mut Object) -> ()>(&mut self, action: A) -> Vec<ObjectId> {
		fn traverse_array<A: Fn(&mut Object) -> ()>(array: &mut Vec<Object>, action: &A, refs: &mut Vec<ObjectId>) {
			for item in array.iter_mut() {
				traverse_object(item, action, refs);
			}
		}
		fn traverse_dictionary<A: Fn(&mut Object) -> ()>(dict: &mut Dictionary, action: &A, refs: &mut Vec<ObjectId>) {
			for (_, v) in dict.iter_mut() {
				traverse_object(v, action, refs);
			}
		}
		fn traverse_object<A: Fn(&mut Object) -> ()>(object: &mut Object, action: &A, refs: &mut Vec<ObjectId>) {
			action(object);
			match *object {
				Object::Array(ref mut array) => traverse_array(array, action, refs),
				Object::Dictionary(ref mut dict) => traverse_dictionary(dict, action, refs),
				Object::Stream(ref mut stream) => traverse_dictionary(&mut stream.dict, action, refs),
				Object::Reference(id) => {
					if !refs.contains(&id) {
						refs.push(id);
					}
				},
				_ => {}
			}
		}
		let mut refs = vec![];
		traverse_dictionary(&mut self.trailer, &action, &mut refs);
		let mut index = 0;
		while index < refs.len() {
			if let Some(object) = self.objects.get_mut(&refs[index]) {
				traverse_object(object, &action, &mut refs);
			}
			index += 1;
		}
		refs
	}

	/// Get catalog dictionary.
	pub fn catalog(&self) -> Option<&Dictionary> {
		self.trailer.get("Root").get_dict_by_ref(self)
	}

	/// Get page numbers and corresponding object ids.
	pub fn get_pages(&self) -> BTreeMap<u32, ObjectId> {
		fn collect_pages(doc: &Document, page_tree_id: ObjectId, page_number: &mut u32, pages: &mut BTreeMap<u32, ObjectId>) {
			if let Some(kids) = doc.get_object(page_tree_id).and_then(|obj|obj.as_dict()).and_then(|page_tree|page_tree.get("Kids")).and_then(|obj|obj.as_array()) {
				for kid in kids {
					if let Some(kid_id) = kid.as_reference() {
						if let Some(type_name) = doc.get_object(kid_id).and_then(|obj|obj.as_dict()).and_then(|dict|dict.type_name()) {
							match type_name {
								"Page" => {
									pages.insert(*page_number, kid_id);
									*page_number += 1;
								}
								"Pages" => {
									collect_pages(doc, kid_id, page_number, pages);
								}
								_ => {}
							}
						}
					}
				}
			}
		}

		let mut pages = BTreeMap::new();
		let mut page_number = 1;
		if let Some(page_tree_id) = self.catalog().and_then(|cat|cat.get("Pages")).and_then(|pages|pages.as_reference()) {
			collect_pages(self, page_tree_id, &mut page_number, &mut pages);
		}
		pages
	}
}
<|endoftext|>