<commit_before>// Copyright 2015 Nathan Sizemore <nathanrsizemore@gmail.com>
//
// This Source Code Form is subject to the
// terms of the Mozilla Public License, v.
// 2.0. If a copy of the MPL was not
// distributed with this file, You can
// obtain one at
// http://mozilla.org/MPL/2.0/.
//
// This Source Code Form is "Incompatible
// With Secondary Licenses", as defined by
// the Mozilla Public License, v. 2.0.

use std::{str, ptr};
use std::ffi::{CStr, CString};
use std::thread;
use std::net::TcpStream;
use std::sync::mpsc::{channel, Sender, Receiver};

use super::libc::{c_int, c_char};
use super::simple_stream::bstream::Bstream;

extern "C" {
    fn register_writer_tx(tx: *mut Sender<Vec<u8>>);
    fn register_stop_tx(tx: *mut Sender<()>);
}

/// Connects to the provided address, (eg "123.123.123.123:3000") and registers
/// the on data received handler
#[no_mangle]
pub extern "C" fn connect(address: *const c_char,
    handler: extern fn(*const c_char, c_int),
    on_connect_handler: extern fn(),
    on_disconnect_handler: extern fn()) -> c_int {

    let mut r_address;
    unsafe {
        r_address = CStr::from_ptr(address);
    }
    let s_address = r_address.to_bytes();
    let host_address = match str::from_utf8(s_address) {
        Ok(safe_str) => safe_str,
        Err(_) => {
            println!("Invalid host address");
            return -1 as c_int;
        }
    };

    // Create and register a way to kill this client
    let (k_tx, kill_rx): (Sender<()>, Receiver<()>) = channel();
    let kill_tx = k_tx.clone();
    let mut k_tx_ptr = Box::new(k_tx);
    unsafe {
        register_stop_tx(&mut *k_tx_ptr);
    }

    // Writer thread's channel
    let (w_tx, w_rx): (Sender<Vec<u8>>, Receiver<Vec<u8>>) = channel();
    let mut w_tx_ptr = Box::new(w_tx);
    unsafe {
        register_writer_tx(&mut *w_tx_ptr);
    }

    let result = TcpStream::connect(host_address);
    // if result.is_err() {
    //     println!("Error connecting to {} - {}", host_address, result.unwrap_err());
    //     return -1 as c_int;
    // }
    // println!("Connected");
    // on_connect_handler();

    // let stream = result.unwrap();
    // let client = Bstream::new(stream);
    //
    // let r_client = client.clone();
    // let r_kill_tx = kill_tx.clone();
    //
    // let w_client = client.clone();
    // let w_kill_tx = kill_tx.clone();
    //
    // // Start the reader thread
    // thread::Builder::new()
    //     .name("ReaderThread".to_string())
    //     .spawn(move||{
    //         reader_thread(r_client, handler, r_kill_tx)
    //     }).unwrap();
    //
    // // Start the writer thread
    // thread::Builder::new()
    //     .name("WriterThread".to_string())
    //     .spawn(move||{
    //         writer_thread(w_rx, w_client, w_kill_tx)
    //     }).unwrap();
    //
    // // Wait for the kill signal
    // match kill_rx.recv() {
    //     Ok(_) => { }
    //     Err(e) => {
    //         println!("Error on kill channel: {}", e);
    //         return -1 as c_int;
    //     }
    // };
    // on_disconnect_handler();

    // Exit out in standard C fashion
    0 as c_int
}

/// Writes the complete contents of buffer to the server
/// Returns -1 on error
#[no_mangle]
pub extern "C" fn send_to_writer(w_tx: *mut Sender<Vec<u8>>,
                                 buffer: *const c_char,
                                 count: c_int,
                                 k_tx: *mut Sender<()>) -> c_int {
    if count < 1 {
        println!("Error - count must be greater than zero");
        return -1 as c_int;
    }

    let num_elts = count as usize;
    let mut n_buffer = Vec::<u8>::with_capacity(num_elts);
    for x in 0..num_elts as isize {
        unsafe {
            n_buffer.push(ptr::read(buffer.offset(x)) as u8);
        }
    }

    unsafe {
        match (*w_tx).send(n_buffer) {
            Ok(_) => { }
            Err(e) => {
                println!("Error sending buffer: {}", e);
                let _ = (*k_tx).send(());
                return -1 as c_int;
            }
        };
    }

    0 as c_int
}

/// Forever listens to incoming data and when a complete message is received,
/// the passed callback is hit
fn reader_thread(client: Bstream,
                 event_handler: extern fn(*const c_char, c_int),
                 kill_tx: Sender<()>) {

    let mut reader = client.clone();
    loop {
        match reader.read() {
            Ok(buffer) => {
                // Launch the handler in a new thread
                thread::Builder::new()
                    .name("Reader-Worker".to_string())
                    .spawn(move||{
                        let slice = &buffer[..];
                        let c_buffer = CString::new(slice).unwrap();
                        event_handler(c_buffer.as_ptr(), buffer.len() as c_int);
                    }).unwrap();
            }
            Err(e) => {
                println!("Error: {}", e);
                break;
            }
        };
    }
    println!("Reader thread finished");
    let _ = kill_tx.send(());
}

/// Forever listens to Receiver<Vec<u8>> waiting on messages to come in
/// Once available, blocks until the entire message has been written
fn writer_thread(rx: Receiver<Vec<u8>>, client: Bstream, kill_tx: Sender<()>) {
    let mut writer = client.clone();
    loop {
        match rx.recv() {
            Ok(ref mut buffer) => {
                match writer.write(buffer) {
                    Ok(_) => { }
                    Err(e) => {
                        println!("Error: {}", e);
                        break;
                    }
                };
            }
            Err(e) => {
                println!("Error: {}", e);
                break;
            }
        };
    }
    println!("Writer thread finished");
    let _ = kill_tx.send(());
}
<commit_msg>Debugging<commit_after>// Copyright 2015 Nathan Sizemore <nathanrsizemore@gmail.com>
//
// This Source Code Form is subject to the
// terms of the Mozilla Public License, v.
// 2.0. If a copy of the MPL was not
// distributed with this file, You can
// obtain one at
// http://mozilla.org/MPL/2.0/.
//
// This Source Code Form is "Incompatible
// With Secondary Licenses", as defined by
// the Mozilla Public License, v. 2.0.

use std::{str, ptr};
use std::ffi::{CStr, CString};
use std::thread;
use std::net::TcpStream;
use std::sync::mpsc::{channel, Sender, Receiver};

use super::libc::{c_int, c_char};
use super::simple_stream::bstream::Bstream;

extern "C" {
    fn register_writer_tx(tx: *mut Sender<Vec<u8>>);
    fn register_stop_tx(tx: *mut Sender<()>);
}

/// Connects to the provided address, (eg "123.123.123.123:3000") and registers
/// the on data received handler
#[no_mangle]
pub extern "C" fn connect(address: *const c_char,
    handler: extern fn(*const c_char, c_int),
    on_connect_handler: extern fn(),
    on_disconnect_handler: extern fn()) -> c_int {

    let mut r_address;
    unsafe {
        r_address = CStr::from_ptr(address);
    }
    let s_address = r_address.to_bytes();
    let host_address = match str::from_utf8(s_address) {
        Ok(safe_str) => safe_str,
        Err(_) => {
            println!("Invalid host address");
            return -1 as c_int;
        }
    };

    // Create and register a way to kill this client
    let (k_tx, kill_rx): (Sender<()>, Receiver<()>) = channel();
    let kill_tx = k_tx.clone();
    let mut k_tx_ptr = Box::new(k_tx);
    unsafe {
        register_stop_tx(&mut *k_tx_ptr);
    }

    // Writer thread's channel
    let (w_tx, w_rx): (Sender<Vec<u8>>, Receiver<Vec<u8>>) = channel();
    let mut w_tx_ptr = Box::new(w_tx);
    unsafe {
        register_writer_tx(&mut *w_tx_ptr);
    }

    println!("host_address: {}", host_address);

    let result = TcpStream::connect(host_address);
    // if result.is_err() {
    //     println!("Error connecting to {} - {}", host_address, result.unwrap_err());
    //     return -1 as c_int;
    // }
    // println!("Connected");
    // on_connect_handler();

    // let stream = result.unwrap();
    // let client = Bstream::new(stream);
    //
    // let r_client = client.clone();
    // let r_kill_tx = kill_tx.clone();
    //
    // let w_client = client.clone();
    // let w_kill_tx = kill_tx.clone();
    //
    // // Start the reader thread
    // thread::Builder::new()
    //     .name("ReaderThread".to_string())
    //     .spawn(move||{
    //         reader_thread(r_client, handler, r_kill_tx)
    //     }).unwrap();
    //
    // // Start the writer thread
    // thread::Builder::new()
    //     .name("WriterThread".to_string())
    //     .spawn(move||{
    //         writer_thread(w_rx, w_client, w_kill_tx)
    //     }).unwrap();
    //
    // // Wait for the kill signal
    // match kill_rx.recv() {
    //     Ok(_) => { }
    //     Err(e) => {
    //         println!("Error on kill channel: {}", e);
    //         return -1 as c_int;
    //     }
    // };
    // on_disconnect_handler();

    // Exit out in standard C fashion
    0 as c_int
}

/// Writes the complete contents of buffer to the server
/// Returns -1 on error
#[no_mangle]
pub extern "C" fn send_to_writer(w_tx: *mut Sender<Vec<u8>>,
                                 buffer: *const c_char,
                                 count: c_int,
                                 k_tx: *mut Sender<()>) -> c_int {
    if count < 1 {
        println!("Error - count must be greater than zero");
        return -1 as c_int;
    }

    let num_elts = count as usize;
    let mut n_buffer = Vec::<u8>::with_capacity(num_elts);
    for x in 0..num_elts as isize {
        unsafe {
            n_buffer.push(ptr::read(buffer.offset(x)) as u8);
        }
    }

    unsafe {
        match (*w_tx).send(n_buffer) {
            Ok(_) => { }
            Err(e) => {
                println!("Error sending buffer: {}", e);
                let _ = (*k_tx).send(());
                return -1 as c_int;
            }
        };
    }

    0 as c_int
}

/// Forever listens to incoming data and when a complete message is received,
/// the passed callback is hit
fn reader_thread(client: Bstream,
                 event_handler: extern fn(*const c_char, c_int),
                 kill_tx: Sender<()>) {

    let mut reader = client.clone();
    loop {
        match reader.read() {
            Ok(buffer) => {
                // Launch the handler in a new thread
                thread::Builder::new()
                    .name("Reader-Worker".to_string())
                    .spawn(move||{
                        let slice = &buffer[..];
                        let c_buffer = CString::new(slice).unwrap();
                        event_handler(c_buffer.as_ptr(), buffer.len() as c_int);
                    }).unwrap();
            }
            Err(e) => {
                println!("Error: {}", e);
                break;
            }
        };
    }
    println!("Reader thread finished");
    let _ = kill_tx.send(());
}

/// Forever listens to Receiver<Vec<u8>> waiting on messages to come in
/// Once available, blocks until the entire message has been written
fn writer_thread(rx: Receiver<Vec<u8>>, client: Bstream, kill_tx: Sender<()>) {
    let mut writer = client.clone();
    loop {
        match rx.recv() {
            Ok(ref mut buffer) => {
                match writer.write(buffer) {
                    Ok(_) => { }
                    Err(e) => {
                        println!("Error: {}", e);
                        break;
                    }
                };
            }
            Err(e) => {
                println!("Error: {}", e);
                break;
            }
        };
    }
    println!("Writer thread finished");
    let _ = kill_tx.send(());
}
<|endoftext|>


<commit_before>use fragment::matching;
use ExclusionPattern;
use walkdir::{DirEntry, Error, WalkDir};
use std::path::PathBuf;
use IndexedPath;

pub struct Index {
    path: PathBuf,
    entries: Vec<IndexedPath>,
}

impl Index {
    pub fn new(path: PathBuf) -> Index {
        Index {
            path: path,
            entries: Vec::new()
        }
    }

    /// Finds all files inside and beneath the index path
    /// and adds them to the index entries vector.
    pub fn populate(&mut self, exclusions: Option<Vec<ExclusionPattern>>, case_sensitive: bool) {
        // The entries listed by read_dir include the root index path; we want
        // relative paths, so we get this length so that we can strip it.
        let prefix_length = match self.path.to_str() {
            Some(path) => path.len() + 1,
            None => return,
        };

        // Start indexing at the specified path.
        let filtered_entries = WalkDir::new(&self.path).into_iter().filter_entry(|entry| {
            if let Some(ref exclusions) = exclusions {
                !exclusions.iter().any(|exclusion| {
                    exclusion.matches(entry.path().to_string_lossy().as_ref())
                })
            } else {
                true
            }
        });

        for entry in filtered_entries {
            relative_entry_path(entry, prefix_length).map(|entry_path| {
                self.entries.push(
                    IndexedPath::new(&entry_path, case_sensitive)
                );
            });
        }
    }

    pub fn find(&self, term: &str, limit: usize) -> Vec<&PathBuf> {
        matching::find(term, &self.entries, limit)
            .into_iter()
            .map(|result| &result.path)
            .collect()
    }
}

/// Transforms a DirEntry object into an optional relative path string,
/// returning None if any errors occur or if the entry is not a file.
fn relative_entry_path(entry: Result<DirEntry, Error>, prefix_length: usize) -> Option<String> {
    entry.ok().and_then(|e| {
        // Limit path look-ups to files.
        e.path().metadata().ok().and_then(|metadata| {
            if metadata.is_file() {
                // Map the absolute path to a relative one.
                e.path().to_str().map(|path| path[prefix_length..].to_string())
            } else {
                None
            }
        })
    })
}

#[cfg(test)]
mod tests {
    extern crate fragment;

    use super::{Index, IndexedPath, ExclusionPattern};
    use std::path::PathBuf;

    #[test]
    fn populate_adds_all_files_to_entries() {
        let path = PathBuf::from("tests/sample");
        let mut index = Index::new(path);
        let expected_entries = vec![IndexedPath::new("directory/Capitalized_file", true),
                                    IndexedPath::new("directory/nested_file", true),
                                    IndexedPath::new("root_file", true)];
        index.populate(None, true);
        index.entries.sort();

        assert_eq!(index.entries, expected_entries);
    }

    #[test]
    fn populate_respects_exclusions() {
        let path = PathBuf::from("tests/sample");
        let mut index = Index::new(path);
        let expected_entries = vec![IndexedPath::new("root_file", true)];
        index.populate(Some(vec![ExclusionPattern::new("**/directory").unwrap()]), true);

        assert_eq!(index.entries, expected_entries);
    }

    #[test]
    fn populate_lowercases_entries_when_case_sensitive_is_false() {
        let path = PathBuf::from("tests/sample");
        let mut index = Index::new(path);
        let expected_entries = vec![IndexedPath::new("directory/capitalized_file", false),
                                    IndexedPath::new("directory/nested_file", false),
                                    IndexedPath::new("root_file", false)];
        index.populate(None, false);
        index.entries.sort();

        assert_eq!(index.entries, expected_entries);
    }

    #[test]
    fn populate_lowercases_entries_when_case_sensitive_is_true() {
        let path = PathBuf::from("tests/sample");
        let mut index = Index::new(path);
        let expected_entries = vec![IndexedPath::new("directory/Capitalized_file", true),
                                    IndexedPath::new("directory/nested_file", true),
                                    IndexedPath::new("root_file", true)];
        index.populate(None, true);
        index.entries.sort();

        assert_eq!(index.entries, expected_entries);
    }

    #[test]
    fn find_defers_to_matching_module() {
        let path = PathBuf::from("tests/sample");
        let mut index = Index::new(path);
        index.populate(None, true);
        let term = "root";
        let limit = 5;

        // Get a string version of the results (PathBuf doesn't implement the display trait).
        let results: Vec<String> = index.find(term, limit)
                                        .iter()
                                        .map(|r| r.to_string_lossy().into_owned())
                                        .collect();

        assert_eq!(results, vec!["root_file".to_string()]);
    }

    #[test]
    fn find_always_uses_case_sensitive_queries() {
        let path = PathBuf::from("tests/sample");
        let mut index = Index::new(path);
        index.populate(None, true);
        let term = "Root";
        let limit = 5;

        // Get a string version of the results (PathBuf doesn't implement the display trait).
        let results: Vec<String> = index.find(term, limit)
                                        .iter()
                                        .map(|r| r.to_string_lossy().into_owned())
                                        .collect();

        assert!(results.is_empty());
    }
}
<commit_msg>Remove superfluous/outdated tests<commit_after>use fragment::matching;
use ExclusionPattern;
use walkdir::{DirEntry, Error, WalkDir};
use std::path::PathBuf;
use IndexedPath;

pub struct Index {
    path: PathBuf,
    entries: Vec<IndexedPath>,
}

impl Index {
    pub fn new(path: PathBuf) -> Index {
        Index {
            path: path,
            entries: Vec::new()
        }
    }

    /// Finds all files inside and beneath the index path
    /// and adds them to the index entries vector.
    pub fn populate(&mut self, exclusions: Option<Vec<ExclusionPattern>>, case_sensitive: bool) {
        // The entries listed by read_dir include the root index path; we want
        // relative paths, so we get this length so that we can strip it.
        let prefix_length = match self.path.to_str() {
            Some(path) => path.len() + 1,
            None => return,
        };

        // Start indexing at the specified path.
        let filtered_entries = WalkDir::new(&self.path).into_iter().filter_entry(|entry| {
            if let Some(ref exclusions) = exclusions {
                !exclusions.iter().any(|exclusion| {
                    exclusion.matches(entry.path().to_string_lossy().as_ref())
                })
            } else {
                true
            }
        });

        for entry in filtered_entries {
            relative_entry_path(entry, prefix_length).map(|entry_path| {
                self.entries.push(
                    IndexedPath::new(&entry_path, case_sensitive)
                );
            });
        }
    }

    pub fn find(&self, term: &str, limit: usize) -> Vec<&PathBuf> {
        matching::find(term, &self.entries, limit)
            .into_iter()
            .map(|result| &result.path)
            .collect()
    }
}

/// Transforms a DirEntry object into an optional relative path string,
/// returning None if any errors occur or if the entry is not a file.
fn relative_entry_path(entry: Result<DirEntry, Error>, prefix_length: usize) -> Option<String> {
    entry.ok().and_then(|e| {
        // Limit path look-ups to files.
        e.path().metadata().ok().and_then(|metadata| {
            if metadata.is_file() {
                // Map the absolute path to a relative one.
                e.path().to_str().map(|path| path[prefix_length..].to_string())
            } else {
                None
            }
        })
    })
}

#[cfg(test)]
mod tests {
    use super::{Index, IndexedPath, ExclusionPattern};
    use std::path::PathBuf;

    #[test]
    fn populate_respects_exclusions() {
        let path = PathBuf::from("tests/sample");
        let mut index = Index::new(path);
        let expected_entries = vec![IndexedPath::new("root_file", true)];
        index.populate(Some(vec![ExclusionPattern::new("**/directory").unwrap()]), true);

        assert_eq!(index.entries, expected_entries);
    }

    #[test]
    fn populate_lowercases_entries_when_case_sensitive_is_false() {
        let path = PathBuf::from("tests/sample");
        let mut index = Index::new(path);
        let expected_entries = vec![IndexedPath::new("directory/capitalized_file", false),
                                    IndexedPath::new("directory/nested_file", false),
                                    IndexedPath::new("root_file", false)];
        index.populate(None, false);
        index.entries.sort();

        assert_eq!(index.entries, expected_entries);
    }

    #[test]
    fn populate_lowercases_entries_when_case_sensitive_is_true() {
        let path = PathBuf::from("tests/sample");
        let mut index = Index::new(path);
        let expected_entries = vec![IndexedPath::new("directory/Capitalized_file", true),
                                    IndexedPath::new("directory/nested_file", true),
                                    IndexedPath::new("root_file", true)];
        index.populate(None, true);
        index.entries.sort();

        assert_eq!(index.entries, expected_entries);
    }

    #[test]
    fn find_defers_to_matching_module() {
        let path = PathBuf::from("tests/sample");
        let mut index = Index::new(path);
        index.populate(None, true);
        let term = "root";
        let limit = 5;

        // Get a string version of the results (PathBuf doesn't implement the display trait).
        let results: Vec<String> = index.find(term, limit)
                                        .iter()
                                        .map(|r| r.to_string_lossy().into_owned())
                                        .collect();

        assert_eq!(results, vec!["root_file".to_string()]);
    }
}
<|endoftext|>


<commit_before>//! Create new `Streams` connected to external inputs.

use std::rc::Rc;
use std::cell::RefCell;
use std::default::Default;

use progress::frontier::Antichain;
use progress::{Operate, Timestamp};
use progress::nested::subgraph::Source;
use progress::count_map::CountMap;
use progress::timestamp::RootTimestamp;
use progress::nested::product::Product;

use timely_communication::Allocate;
use {Data, Push};
use dataflow::channels::Content;
use dataflow::channels::pushers::{Tee, Counter};

use dataflow::{Stream, Scope};
use dataflow::scopes::{Child, Root};

// TODO : This is an exogenous input, but it would be nice to wrap a Subgraph in something
// TODO : more like a harness, with direct access to its inputs.

// NOTE : This only takes a &self, not a &mut self, which works but is a bit weird.
// NOTE : Experiments with &mut indicate that the borrow of 'a lives for too long.
// NOTE : Might be able to fix with another lifetime parameter, say 'c: 'a.

/// Create a new `Stream` and `Handle` through which to supply input.
pub trait Input<'a, A: Allocate, T: Timestamp+Ord> {
    /// Create a new `Stream` and `Handle` through which to supply input.
    ///
    /// The `new_input` method returns a pair `(Handle, Stream)` where the `Stream` can be used
    /// immediately for timely dataflow construction, and the `Handle` is later used to introduce
    /// data into the timely dataflow computation.
    ///
    /// The `Handle` also provides a means to indicate
    /// to timely dataflow that the input has advanced beyond certain timestamps, allowing timely
    /// to issue progress notifications.
    ///
    /// #Examples
    /// ```
    /// use timely::*;
    /// use timely::dataflow::Scope;
    /// use timely::dataflow::operators::{Input, Inspect};
    ///
    /// // construct and execute a timely dataflow
    /// timely::execute(Configuration::Thread, |root| {
    ///
    ///     // add an input and base computation off of it
    ///     let mut input = root.scoped(|scope| {
    ///         let (input, stream) = scope.new_input();
    ///         stream.inspect(|x| println!("hello {:?}", x));
    ///         input
    ///     });
    ///
    ///     // introduce input, advance computation
    ///     for round in 0..10 {
    ///         input.send(round);
    ///         input.advance_to(round + 1);
    ///         root.step();
    ///     }
    /// });
    /// ```
    fn new_input<D: Data>(&mut self) -> (Handle<T, D>, Stream<Child<'a, Root<A>, T>, D>);
}

impl<'a, A: Allocate, T: Timestamp+Ord> Input<'a, A, T> for Child<'a, Root<A>, T> {
    fn new_input<D: Data>(&mut self) -> (Handle<T, D>, Stream<Child<'a, Root<A>, T>, D>) {

        let (output, registrar) = Tee::<Product<RootTimestamp, T>, D>::new();
        let produced = Rc::new(RefCell::new(CountMap::new()));
        let helper = Handle::new(Counter::new(output, produced.clone()));
        let copies = self.peers();

        let index = self.add_operator(Operator {
            progress: helper.progress.clone(),
            messages: produced.clone(),
            copies:   copies,
        });

        (helper, Stream::new(Source { index: index, port: 0 }, registrar, self.clone()))
    }
}

struct Operator<T:Timestamp+Ord> {
    progress:   Rc<RefCell<CountMap<Product<RootTimestamp, T>>>>,           // times closed since last asked
    messages:   Rc<RefCell<CountMap<Product<RootTimestamp, T>>>>,           // messages sent since last asked
    copies:     usize,
}

impl<T:Timestamp+Ord> Operate<Product<RootTimestamp, T>> for Operator<T> {
    fn name(&self) -> String { "Input".to_owned() }
    fn inputs(&self) -> usize { 0 }
    fn outputs(&self) -> usize { 1 }

    fn get_internal_summary(&mut self) -> (Vec<Vec<Antichain<<Product<RootTimestamp, T> as Timestamp>::Summary>>>,
                                           Vec<CountMap<Product<RootTimestamp, T>>>) {
        let mut map = CountMap::new();
        map.update(&Default::default(), self.copies as i64);
        (Vec::new(), vec![map])
    }

    fn pull_internal_progress(&mut self,_messages_consumed: &mut [CountMap<Product<RootTimestamp, T>>],
                                         frontier_progress: &mut [CountMap<Product<RootTimestamp, T>>],
                                         messages_produced: &mut [CountMap<Product<RootTimestamp, T>>]) -> bool
    {
        self.messages.borrow_mut().drain_into(&mut messages_produced[0]);
        self.progress.borrow_mut().drain_into(&mut frontier_progress[0]);
        false
    }

    fn notify_me(&self) -> bool { false }
}

/// A handle to an input `Stream`, used to introduce data to a timely dataflow computation.
pub struct Handle<T: Timestamp+Ord, D: Data> {
    // frontier: Rc<RefCell<MutableAntichain<Product<RootTimestamp, T>>>>,   // times available for sending
    progress: Rc<RefCell<CountMap<Product<RootTimestamp, T>>>>,           // times closed since last asked
    pusher: Counter<Product<RootTimestamp, T>, D, Tee<Product<RootTimestamp, T>, D>>,
    buffer: Vec<D>,
    now_at: Product<RootTimestamp, T>,
}

// an input helper's state is either uninitialized, with now_at == None, or at some specific time.
// if now_at == None it has a hold on Default::default(), else it has a hold on the specific time.
// if now_at == None the pusher has not been opened, else it is open with the specific time.

impl<T:Timestamp+Ord, D: Data> Handle<T, D> {
    fn new(pusher: Counter<Product<RootTimestamp, T>, D, Tee<Product<RootTimestamp, T>, D>>) -> Handle<T, D> {
        Handle {
            // frontier: Rc::new(RefCell::new(MutableAntichain::new_bottom(Default::default()))),
            progress: Rc::new(RefCell::new(CountMap::new())),
            pusher: pusher,
            buffer: Vec::with_capacity(Content::<D>::default_length()),
            now_at: Default::default(),
        }
    }

    // flushes any data we are sitting on. may need to initialize self.now_at if no one has yet.
    fn flush(&mut self) {
        Content::push_at(&mut self.buffer, self.now_at, &mut self.pusher);
    }

    // closes the current epoch, flushing if needed, shutting if needed, and updating the frontier.
    fn close_epoch(&mut self) {
        if !self.buffer.is_empty() { self.flush(); }
        self.pusher.done();
        self.progress.borrow_mut().update(&self.now_at, -1);
    }

    #[inline(always)]
    /// Sends one record into the corresponding timely dataflow `Stream`, at the current epoch.
    pub fn send(&mut self, data: D) {
        // assert!(self.buffer.capacity() == Content::<D>::default_length());
        self.buffer.push(data);
        if self.buffer.len() == self.buffer.capacity() {
            self.flush();
        }
    }

    /// Advances the current epoch to `next`.
    ///
    /// This method allows timely dataflow to issue progress notifications as it can now determine
    /// that this input can no longer produce data at earlier timestamps.
    pub fn advance_to(&mut self, next: T) {
        assert!(next > self.now_at.inner);
        self.close_epoch();
        self.now_at = RootTimestamp::new(next);
        self.progress.borrow_mut().update(&self.now_at, 1);
    }

    /// Closes the input.
    ///
    /// This method allows timely dataflow to issue all progress notifications blocked by this input
    /// and to begin to shut down operators, as this input can no longer produce data.
    pub fn close(self) { }

    /// Reports the current epoch.
    pub fn epoch(&self) -> &T {
        &self.now_at.inner
    }

    /// Reports the current timestamp.
    pub fn time(&self) -> &Product<RootTimestamp, T> {
        &self.now_at
    }
}

impl<T:Timestamp+Ord, D: Data> Drop for Handle<T, D> {
    fn drop(&mut self) {
        self.close_epoch();
    }
}
<commit_msg>added send_batch<commit_after>//! Create new `Streams` connected to external inputs.

use std::rc::Rc;
use std::cell::RefCell;
use std::default::Default;

use progress::frontier::Antichain;
use progress::{Operate, Timestamp};
use progress::nested::subgraph::Source;
use progress::count_map::CountMap;
use progress::timestamp::RootTimestamp;
use progress::nested::product::Product;

use timely_communication::Allocate;
use {Data, Push};
use dataflow::channels::Content;
use dataflow::channels::pushers::{Tee, Counter};

use dataflow::{Stream, Scope};
use dataflow::scopes::{Child, Root};

// TODO : This is an exogenous input, but it would be nice to wrap a Subgraph in something
// TODO : more like a harness, with direct access to its inputs.

// NOTE : This only takes a &self, not a &mut self, which works but is a bit weird.
// NOTE : Experiments with &mut indicate that the borrow of 'a lives for too long.
// NOTE : Might be able to fix with another lifetime parameter, say 'c: 'a.

/// Create a new `Stream` and `Handle` through which to supply input.
pub trait Input<'a, A: Allocate, T: Timestamp+Ord> {
    /// Create a new `Stream` and `Handle` through which to supply input.
    ///
    /// The `new_input` method returns a pair `(Handle, Stream)` where the `Stream` can be used
    /// immediately for timely dataflow construction, and the `Handle` is later used to introduce
    /// data into the timely dataflow computation.
    ///
    /// The `Handle` also provides a means to indicate
    /// to timely dataflow that the input has advanced beyond certain timestamps, allowing timely
    /// to issue progress notifications.
    ///
    /// #Examples
    /// ```
    /// use timely::*;
    /// use timely::dataflow::Scope;
    /// use timely::dataflow::operators::{Input, Inspect};
    ///
    /// // construct and execute a timely dataflow
    /// timely::execute(Configuration::Thread, |root| {
    ///
    ///     // add an input and base computation off of it
    ///     let mut input = root.scoped(|scope| {
    ///         let (input, stream) = scope.new_input();
    ///         stream.inspect(|x| println!("hello {:?}", x));
    ///         input
    ///     });
    ///
    ///     // introduce input, advance computation
    ///     for round in 0..10 {
    ///         input.send(round);
    ///         input.advance_to(round + 1);
    ///         root.step();
    ///     }
    /// });
    /// ```
    fn new_input<D: Data>(&mut self) -> (Handle<T, D>, Stream<Child<'a, Root<A>, T>, D>);
}

impl<'a, A: Allocate, T: Timestamp+Ord> Input<'a, A, T> for Child<'a, Root<A>, T> {
    fn new_input<D: Data>(&mut self) -> (Handle<T, D>, Stream<Child<'a, Root<A>, T>, D>) {

        let (output, registrar) = Tee::<Product<RootTimestamp, T>, D>::new();
        let produced = Rc::new(RefCell::new(CountMap::new()));
        let helper = Handle::new(Counter::new(output, produced.clone()));
        let copies = self.peers();

        let index = self.add_operator(Operator {
            progress: helper.progress.clone(),
            messages: produced.clone(),
            copies:   copies,
        });

        (helper, Stream::new(Source { index: index, port: 0 }, registrar, self.clone()))
    }
}

struct Operator<T:Timestamp+Ord> {
    progress:   Rc<RefCell<CountMap<Product<RootTimestamp, T>>>>,           // times closed since last asked
    messages:   Rc<RefCell<CountMap<Product<RootTimestamp, T>>>>,           // messages sent since last asked
    copies:     usize,
}

impl<T:Timestamp+Ord> Operate<Product<RootTimestamp, T>> for Operator<T> {
    fn name(&self) -> String { "Input".to_owned() }
    fn inputs(&self) -> usize { 0 }
    fn outputs(&self) -> usize { 1 }

    fn get_internal_summary(&mut self) -> (Vec<Vec<Antichain<<Product<RootTimestamp, T> as Timestamp>::Summary>>>,
                                           Vec<CountMap<Product<RootTimestamp, T>>>) {
        let mut map = CountMap::new();
        map.update(&Default::default(), self.copies as i64);
        (Vec::new(), vec![map])
    }

    fn pull_internal_progress(&mut self,_messages_consumed: &mut [CountMap<Product<RootTimestamp, T>>],
                                         frontier_progress: &mut [CountMap<Product<RootTimestamp, T>>],
                                         messages_produced: &mut [CountMap<Product<RootTimestamp, T>>]) -> bool
    {
        self.messages.borrow_mut().drain_into(&mut messages_produced[0]);
        self.progress.borrow_mut().drain_into(&mut frontier_progress[0]);
        false
    }

    fn notify_me(&self) -> bool { false }
}

/// A handle to an input `Stream`, used to introduce data to a timely dataflow computation.
pub struct Handle<T: Timestamp+Ord, D: Data> {
    // frontier: Rc<RefCell<MutableAntichain<Product<RootTimestamp, T>>>>,   // times available for sending
    progress: Rc<RefCell<CountMap<Product<RootTimestamp, T>>>>,           // times closed since last asked
    pusher: Counter<Product<RootTimestamp, T>, D, Tee<Product<RootTimestamp, T>, D>>,
    buffer: Vec<D>,
    now_at: Product<RootTimestamp, T>,
}

// an input helper's state is either uninitialized, with now_at == None, or at some specific time.
// if now_at == None it has a hold on Default::default(), else it has a hold on the specific time.
// if now_at == None the pusher has not been opened, else it is open with the specific time.

impl<T:Timestamp+Ord, D: Data> Handle<T, D> {
    fn new(pusher: Counter<Product<RootTimestamp, T>, D, Tee<Product<RootTimestamp, T>, D>>) -> Handle<T, D> {
        Handle {
            // frontier: Rc::new(RefCell::new(MutableAntichain::new_bottom(Default::default()))),
            progress: Rc::new(RefCell::new(CountMap::new())),
            pusher: pusher,
            buffer: Vec::with_capacity(Content::<D>::default_length()),
            now_at: Default::default(),
        }
    }

    // flushes any data we are sitting on. may need to initialize self.now_at if no one has yet.
    fn flush(&mut self) {
        Content::push_at(&mut self.buffer, self.now_at, &mut self.pusher);
    }

    // closes the current epoch, flushing if needed, shutting if needed, and updating the frontier.
    fn close_epoch(&mut self) {
        if !self.buffer.is_empty() { self.flush(); }
        self.pusher.done();
        self.progress.borrow_mut().update(&self.now_at, -1);
    }

    #[inline(always)]
    /// Sends one record into the corresponding timely dataflow `Stream`, at the current epoch.
    pub fn send(&mut self, data: D) {
        // assert!(self.buffer.capacity() == Content::<D>::default_length());
        self.buffer.push(data);
        if self.buffer.len() == self.buffer.capacity() {
            self.flush();
        }
    }

    /// Sends a batch of records into the corresponding timely dataflow `Stream`, at the current epoch.
    ///
    /// This method flushes single elements previously sent with `send`, to keep the insertion order.
    pub fn send_batch(&mut self, buffer: &mut Vec<D>) {
        self.flush();
        Content::push_at(buffer, self.now_at, &mut self.pusher);
    }

    /// Advances the current epoch to `next`.
    ///
    /// This method allows timely dataflow to issue progress notifications as it can now determine
    /// that this input can no longer produce data at earlier timestamps.
    pub fn advance_to(&mut self, next: T) {
        assert!(next > self.now_at.inner);
        self.close_epoch();
        self.now_at = RootTimestamp::new(next);
        self.progress.borrow_mut().update(&self.now_at, 1);
    }

    /// Closes the input.
    ///
    /// This method allows timely dataflow to issue all progress notifications blocked by this input
    /// and to begin to shut down operators, as this input can no longer produce data.
    pub fn close(self) { }

    /// Reports the current epoch.
    pub fn epoch(&self) -> &T {
        &self.now_at.inner
    }

    /// Reports the current timestamp.
    pub fn time(&self) -> &Product<RootTimestamp, T> {
        &self.now_at
    }
}

impl<T:Timestamp+Ord, D: Data> Drop for Handle<T, D> {
    fn drop(&mut self) {
        self.close_epoch();
    }
}
<|endoftext|>


<commit_before>extern crate ring;
extern crate udt;

use std::net::{UdpSocket, SocketAddr, IpAddr};
use std::str;
use std::fmt;
use udt::{UdtSocket, UdtError, UdtOpts, SocketType, SocketFamily};
use ring::aead;
use ring::rand;

// TODO config
const UDT_BUF_SIZE: i32 = 4096000;
pub const MAX_MESSAGE_SIZE: usize = 1024000;

pub mod crypto {
    use ring::aead;
    use ring::aead::{SealingKey, OpeningKey, Algorithm};
    use ring::rand::{SystemRandom, SecureRandom};
    static ALGORITHM: &'static Algorithm = &aead::AES_256_GCM;

    pub struct Handler {
        _working_nonce_buf: [u8; 32],
        _working_seal_buf: [u8; super::MAX_MESSAGE_SIZE],
        rand: SystemRandom,
        opening_key: OpeningKey,
        sealing_key: SealingKey,
    }

    pub fn gen_key() -> Vec<u8> {
        let rand = SystemRandom::new();
        let mut keybytes = vec![0u8; ALGORITHM.key_len()];
        rand.fill(&mut keybytes);
        keybytes
    }

    impl Handler {
        pub fn new(key: &[u8]) -> Handler {
            Handler {
                _working_seal_buf: [0u8; super::MAX_MESSAGE_SIZE],
                _working_nonce_buf: [0u8; 32],
                rand: SystemRandom::new(),
                opening_key: aead::OpeningKey::new(ALGORITHM, key).unwrap(),
                sealing_key: aead::SealingKey::new(ALGORITHM, key).unwrap(),
            }
        }

        pub fn seal(&mut self, buf: &mut [u8], len: usize) -> Result<usize, ()> {
            let nonce_len = ALGORITHM.nonce_len();
            let max_suffix_len = ALGORITHM.max_overhead_len();

            assert!(nonce_len < u8::max_value() as usize,
                    "Uh, why is the nonce size this big?");

            assert!(len <= buf.len() - max_suffix_len,
                    "Buffer doesn't have enough suffix padding.");

            let mut nonce = &mut self._working_nonce_buf[..nonce_len];
            self.rand.fill(&mut nonce).unwrap();

            let mut sealed = &mut self._working_seal_buf[..len + max_suffix_len];
            sealed[0..len].copy_from_slice(&buf[..len]);
            match aead::seal_in_place(&self.sealing_key,
                                      &nonce,
                                      &mut sealed,
                                      max_suffix_len,
                                      &[]) {
                Ok(seal_len) => {
                    buf[..nonce_len].copy_from_slice(&nonce[..]);
                    buf[nonce_len..nonce_len+seal_len].copy_from_slice(&sealed[..seal_len]);
                    Ok(nonce_len + seal_len)
                }
                Err(e) => {
                    Err(())
                }
            }
        }

        pub fn open(&mut self, buf: &mut [u8]) -> Result<usize, String> {
            let nonce_len = ALGORITHM.nonce_len();

            if buf.len() < nonce_len {
                return Err("msg not long enough to contain nonce".into());
            } else if buf.len() > super::MAX_MESSAGE_SIZE {
                return Err("max message size exceeded".into());
            }

            let nonce = buf[..nonce_len].to_owned();

            aead::open_in_place(&self.opening_key, &nonce, nonce_len, buf, &[])
                .map_err(|_| String::from("decrypt failed"))
        }
    }

    // Tests for the crypto module
    #[cfg(test)]
    mod test {

        #[test]
        fn raw_roundtrip() {
            use ring::aead;
            use ring::aead::{SealingKey, OpeningKey, Algorithm};
            use ring::rand::{SystemRandom, SecureRandom};

            let rng = SystemRandom::new();
            let mut key_bytes = vec![0u8; super::ALGORITHM.key_len()];
            let mut nonce_bytes = vec![0u8; super::ALGORITHM.nonce_len()];
            rng.fill(&mut key_bytes).unwrap();
            rng.fill(&mut nonce_bytes).unwrap();
            let key = SealingKey::new(super::ALGORITHM, &key_bytes).unwrap();

            let data = [1u8; 1350];
            let out_suffix_capacity = super::ALGORITHM.max_overhead_len();
            let mut in_out = vec![1u8; data.len() + out_suffix_capacity];
            aead::seal_in_place(&key, &nonce_bytes,
                                &mut in_out, out_suffix_capacity,
                                &[]).unwrap();

            let opening_key = OpeningKey::new(super::ALGORITHM, &key_bytes).unwrap();
            let len = aead::open_in_place(&opening_key, &nonce_bytes,
                                          0, &mut in_out, &[]).unwrap();

            assert_eq!(len, 1350);
            assert_eq!(&in_out[..len], &data[..]);
        }

        #[test]
        fn roundtrip() {
            use ::rand;
            use ::rand::distributions::{IndependentSample, Range};
            // generate some data, seal it, and then make sure it unseals to the same thing
            let mut rng = rand::thread_rng();
            let between = Range::new(10, 10000);

            let key = super::gen_key();
            let mut handler = super::Handler::new(&key);
            let data_size: usize = between.ind_sample(&mut rng);
            let mut data = vec![0u8; super::super::MAX_MESSAGE_SIZE];
            for i in 0..data_size {
                data[i] = rand::random();
            }

            let orig = data[..data_size].to_owned();

            let cipher_len = handler.seal(&mut data, data_size).unwrap();
            let decrypted_len = handler.open(&mut data[..cipher_len]).unwrap();
            assert_eq!(decrypted_len, data_size);
            assert_eq!(orig, &data[..decrypted_len], "original and decrypted don't match!");
        }

        #[test]
        fn key_sanity() {
            use std::collections::HashSet;

            let mut set: HashSet<Vec<u8>> = HashSet::with_capacity(10000);

            for _ in 0..10000 {
                let key = super::gen_key();
                assert!(set.insert(key));
            }
        }
    }

    #[cfg(all(feature = "nightly", test))]
    mod bench {
        extern crate test;
        const DATA_SIZE: usize = 1300;

        #[bench]
        fn bench_raw_seal(b: &mut test::Bencher) {
            use ring::aead;
            use ring::aead::{SealingKey, OpeningKey, Algorithm};
            use ring::rand::{SystemRandom, SecureRandom};

            let rng = SystemRandom::new();
            let mut key_bytes = vec![0u8; super::ALGORITHM.key_len()];
            let mut nonce_bytes = vec![0u8; super::ALGORITHM.nonce_len()];
            rng.fill(&mut key_bytes).unwrap();
            let key = SealingKey::new(super::ALGORITHM, &key_bytes).unwrap();

            let data = [1u8; DATA_SIZE];
            let out_suffix_capacity = super::ALGORITHM.max_overhead_len();
            let mut in_out = vec![1u8; data.len() + out_suffix_capacity];

            b.bytes = DATA_SIZE as u64;
            b.iter(|| {
                rng.fill(&mut nonce_bytes).unwrap();
                aead::seal_in_place(&key, &nonce_bytes, &mut in_out,
                                    out_suffix_capacity, &[]).unwrap()
            })
        }

        #[bench]
        fn bench_raw_open(b: &mut test::Bencher) {
            use ring::aead;
            use ring::aead::{SealingKey, OpeningKey, Algorithm};
            use ring::rand::{SystemRandom, SecureRandom};

            let rng = SystemRandom::new();
            let mut key_bytes = vec![0u8; super::ALGORITHM.key_len()];
            let mut nonce_bytes = vec![0u8; super::ALGORITHM.nonce_len()];
            rng.fill(&mut key_bytes).unwrap();
            rng.fill(&mut nonce_bytes).unwrap();
            let key = SealingKey::new(super::ALGORITHM, &key_bytes).unwrap();
            let opening_key = OpeningKey::new(super::ALGORITHM, &key_bytes).unwrap();

            let data = [1u8; DATA_SIZE];
            let out_suffix_capacity = super::ALGORITHM.max_overhead_len();
            let mut in_out = vec![1u8; data.len() + out_suffix_capacity];

            b.bytes = DATA_SIZE as u64;

            let sealed_len = aead::seal_in_place(&key, &nonce_bytes, &mut in_out,
                                                 out_suffix_capacity, &[]).unwrap();
            b.iter(|| aead::open_in_place(&opening_key, &nonce_bytes,
                                               0, &mut in_out, &[]))
        }

        #[bench]
        fn bench_seal(b: &mut test::Bencher) {
            let key = super::gen_key();
            let mut handler = super::Handler::new(&key);
            let mut buf = vec![0u8; super::super::MAX_MESSAGE_SIZE];
            b.bytes = DATA_SIZE as u64;
            b.iter(|| handler.seal(&mut buf, DATA_SIZE))
        }

        #[bench]
        fn bench_open(b: &mut test::Bencher) {
            let key = super::gen_key();
            let mut handler = super::Handler::new(&key);
            let mut buf = vec![0u8; super::super::MAX_MESSAGE_SIZE];
            let sealed_len = handler.seal(&mut buf, DATA_SIZE).unwrap();
            b.bytes = DATA_SIZE as u64;
            b.iter(|| handler.open(&mut buf[..sealed_len]))
        }
    }
}

fn new_udt_socket() -> UdtSocket {
    udt::init();
    let sock = UdtSocket::new(SocketFamily::AFInet, SocketType::Datagram).unwrap();
    sock.setsockopt(UdtOpts::UDP_RCVBUF, UDT_BUF_SIZE).unwrap();
    sock.setsockopt(UdtOpts::UDP_SNDBUF, UDT_BUF_SIZE).unwrap();
    sock
}

fn send(sock: &UdtSocket, key: &aead::SealingKey, buf: &mut [u8], len: usize) -> Result<(), UdtError> {
    unimplemented!();
    // FIXME don't unwrap, create an Error struct that can handle everything
    // if let Ok(sealed_len) = crypto::seal(buf, len, key) {
    //     sock.sendmsg(&buf[..sealed_len]).map(|_| ())
    // } else {
    //     Err(UdtError {
    //         err_code: -1,
    //         err_msg: "encryption failure".into(),
    //     })
    // }
}

fn recv(sock: &UdtSocket, key: &aead::OpeningKey, buf: &mut [u8]) -> Result<usize, UdtError> {
    unimplemented!();
    // let size = try!(sock.recvmsg(buf));
    // crypto::open(&mut buf[..size], key).map_err(|_| {
    //     UdtError {
    //         err_code: -1,
    //         err_msg: String::from("decryption failure"),
    //     }
    // })
}

pub struct PortRange {
    start: u16,
    end: u16,
}

pub trait Transceiver {
    fn send(&self, buf: &[u8]) -> Result<(), UdtError>;
    fn recv(&self, buf: &mut [u8]) -> Result<usize, UdtError>;
    fn close(&self) -> Result<(), UdtError>;
}

pub struct Server {
    pub ip_addr: IpAddr,
    pub port: u16,
    crypto: crypto::Handler,
    sock: UdtSocket,
}

pub struct Client {
    addr: SocketAddr,
    sock: UdtSocket,
    crypto: crypto::Handler,
}

pub struct ServerConnection<'a> {
    crypto: &'a crypto::Handler,
    sock: UdtSocket,
}

impl Client {
    pub fn new(addr: SocketAddr, key: &[u8]) -> Client {
        let sock = new_udt_socket();
        Client {
            addr: addr,
            sock: sock,
            crypto: crypto::Handler::new(key),
        }
    }

    pub fn connect(&self) -> Result<(), UdtError> {
        self.sock.connect(self.addr)
    }
}

impl Transceiver for Client {
    fn send(&self, buf: &[u8]) -> Result<(), UdtError> {
        unimplemented!();
        // send(&self.sock, &self.sealing_key, buf)
    }

    fn recv(&self, buf: &mut [u8]) -> Result<usize, UdtError> {
        unimplemented!();
        // recv(&self.sock, &self.opening_key, buf)
    }

    fn close(&self) -> Result<(), UdtError> {
        self.sock.close()
    }
}

impl Server {
    pub fn get_open_port(range: &PortRange) -> Result<u16, ()> {
        for p in range.start..range.end {
            if let Ok(_) = UdpSocket::bind(&format!("0.0.0.0:{}", p)[..]) {
                return Ok(p);
            }
        }
        Err(())
    }

    pub fn new(ip_addr: IpAddr, port: u16, key: &[u8]) -> Server {
        let sock = new_udt_socket();
        sock.bind(SocketAddr::new(ip_addr, port)).unwrap();
        Server {
            sock: sock,
            ip_addr: ip_addr,
            port: port,
            crypto: crypto::Handler::new(key),
        }
    }

    pub fn listen(&self) -> Result<(), UdtError> {
        self.sock.listen(2)
    }

    pub fn accept(&self) -> Result<ServerConnection, UdtError> {
        self.sock.accept().map(|(sock, _)| {
            ServerConnection {
                crypto: &self.crypto,
                sock: sock,
            }
        })
    }
}

impl<'a> ServerConnection<'a> {
    pub fn getpeer(&self) -> Result<SocketAddr, UdtError> {
        self.sock.getpeername()
    }
}

impl<'a> Transceiver for ServerConnection<'a> {
    fn send(&self, buf: &[u8]) -> Result<(), UdtError> {
        unimplemented!();
        // send(&self.sock, self.key, buf)
    }

    fn recv(&self, buf: &mut [u8]) -> Result<usize, UdtError> {
        unimplemented!();
        // recv(&self.sock, self.key, buf)
    }

    fn close(&self) -> Result<(), UdtError> {
        self.sock.close()
    }
}

impl<'a> PortRange {
    fn new(start: u16, end: u16) -> Result<PortRange, &'a str> {
        if start > end {
            Err("range end must be greater than or equal to start")
        } else {
            Ok(PortRange {
                start: start,
                end: end,
            })
        }
    }

    pub fn from(s: &str) -> Result<PortRange, &'a str> {
        let sections: Vec<&str> = s.split('-').collect();
        if sections.len() != 2 {
            return Err("Range must be specified in the form of \"<start>-<end>\"");
        }
        let (start, end) = (sections[0].parse::<u16>(), sections[1].parse::<u16>());
        if start.is_err() || end.is_err() {
            return Err("improperly formatted port range");
        }
        PortRange::new(start.unwrap(), end.unwrap())
    }
}

impl fmt::Display for PortRange {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}-{}", self.start, self.end)
    }
}
<commit_msg>allocation-less open<commit_after>extern crate ring;
extern crate udt;

use std::net::{UdpSocket, SocketAddr, IpAddr};
use std::str;
use std::fmt;
use udt::{UdtSocket, UdtError, UdtOpts, SocketType, SocketFamily};
use ring::aead;
use ring::rand;

// TODO config
const UDT_BUF_SIZE: i32 = 4096000;
pub const MAX_MESSAGE_SIZE: usize = 1024000;

pub mod crypto {
    use ring::aead;
    use ring::aead::{SealingKey, OpeningKey, Algorithm};
    use ring::rand::{SystemRandom, SecureRandom};
    static ALGORITHM: &'static Algorithm = &aead::AES_256_GCM;

    pub struct Handler {
        _working_nonce_buf: [u8; 32],
        _working_seal_buf: [u8; super::MAX_MESSAGE_SIZE],
        rand: SystemRandom,
        opening_key: OpeningKey,
        sealing_key: SealingKey,
    }

    pub fn gen_key() -> Vec<u8> {
        let rand = SystemRandom::new();
        let mut keybytes = vec![0u8; ALGORITHM.key_len()];
        rand.fill(&mut keybytes);
        keybytes
    }

    impl Handler {
        pub fn new(key: &[u8]) -> Handler {
            Handler {
                _working_seal_buf: [0u8; super::MAX_MESSAGE_SIZE],
                _working_nonce_buf: [0u8; 32],
                rand: SystemRandom::new(),
                opening_key: aead::OpeningKey::new(ALGORITHM, key).unwrap(),
                sealing_key: aead::SealingKey::new(ALGORITHM, key).unwrap(),
            }
        }

        pub fn seal(&mut self, buf: &mut [u8], len: usize) -> Result<usize, ()> {
            let nonce_len = ALGORITHM.nonce_len();
            let max_suffix_len = ALGORITHM.max_overhead_len();

            assert!(nonce_len < u8::max_value() as usize,
                    "Uh, why is the nonce size this big?");

            assert!(len <= buf.len() - max_suffix_len,
                    "Buffer doesn't have enough suffix padding.");

            let mut nonce = &mut self._working_nonce_buf[..nonce_len];
            self.rand.fill(&mut nonce).unwrap();

            let mut sealed = &mut self._working_seal_buf[..len + max_suffix_len];
            sealed[0..len].copy_from_slice(&buf[..len]);
            match aead::seal_in_place(&self.sealing_key,
                                      &nonce,
                                      &mut sealed,
                                      max_suffix_len,
                                      &[]) {
                Ok(seal_len) => {
                    buf[..nonce_len].copy_from_slice(&nonce[..]);
                    buf[nonce_len..nonce_len+seal_len].copy_from_slice(&sealed[..seal_len]);
                    Ok(nonce_len + seal_len)
                }
                Err(e) => {
                    Err(())
                }
            }
        }

        pub fn open(&mut self, buf: &mut [u8]) -> Result<usize, String> {
            let nonce_len = ALGORITHM.nonce_len();

            if buf.len() < nonce_len {
                return Err("msg not long enough to contain nonce".into());
            } else if buf.len() > super::MAX_MESSAGE_SIZE {
                return Err("max message size exceeded".into());
            }

            let nonce = &mut self._working_nonce_buf[..nonce_len];
            nonce.copy_from_slice(&buf[..nonce_len]);

            aead::open_in_place(&self.opening_key, &nonce, nonce_len, buf, &[])
                .map_err(|_| String::from("decrypt failed"))
        }
    }

    // Tests for the crypto module
    #[cfg(test)]
    mod test {

        #[test]
        fn raw_roundtrip() {
            use ring::aead;
            use ring::aead::{SealingKey, OpeningKey, Algorithm};
            use ring::rand::{SystemRandom, SecureRandom};

            let rng = SystemRandom::new();
            let mut key_bytes = vec![0u8; super::ALGORITHM.key_len()];
            let mut nonce_bytes = vec![0u8; super::ALGORITHM.nonce_len()];
            rng.fill(&mut key_bytes).unwrap();
            rng.fill(&mut nonce_bytes).unwrap();
            let key = SealingKey::new(super::ALGORITHM, &key_bytes).unwrap();

            let data = [1u8; 1350];
            let out_suffix_capacity = super::ALGORITHM.max_overhead_len();
            let mut in_out = vec![1u8; data.len() + out_suffix_capacity];
            aead::seal_in_place(&key, &nonce_bytes,
                                &mut in_out, out_suffix_capacity,
                                &[]).unwrap();

            let opening_key = OpeningKey::new(super::ALGORITHM, &key_bytes).unwrap();
            let len = aead::open_in_place(&opening_key, &nonce_bytes,
                                          0, &mut in_out, &[]).unwrap();

            assert_eq!(len, 1350);
            assert_eq!(&in_out[..len], &data[..]);
        }

        #[test]
        fn roundtrip() {
            use ::rand;
            use ::rand::distributions::{IndependentSample, Range};
            // generate some data, seal it, and then make sure it unseals to the same thing
            let mut rng = rand::thread_rng();
            let between = Range::new(10, 10000);

            let key = super::gen_key();
            let mut handler = super::Handler::new(&key);
            let data_size: usize = between.ind_sample(&mut rng);
            let mut data = vec![0u8; super::super::MAX_MESSAGE_SIZE];
            for i in 0..data_size {
                data[i] = rand::random();
            }

            let orig = data[..data_size].to_owned();

            let cipher_len = handler.seal(&mut data, data_size).unwrap();
            let decrypted_len = handler.open(&mut data[..cipher_len]).unwrap();
            assert_eq!(decrypted_len, data_size);
            assert_eq!(orig, &data[..decrypted_len], "original and decrypted don't match!");
        }

        #[test]
        fn key_sanity() {
            use std::collections::HashSet;

            let mut set: HashSet<Vec<u8>> = HashSet::with_capacity(10000);

            for _ in 0..10000 {
                let key = super::gen_key();
                assert!(set.insert(key));
            }
        }
    }

    #[cfg(all(feature = "nightly", test))]
    mod bench {
        extern crate test;
        const DATA_SIZE: usize = 1300;

        #[bench]
        fn bench_raw_seal(b: &mut test::Bencher) {
            use ring::aead;
            use ring::aead::{SealingKey, OpeningKey, Algorithm};
            use ring::rand::{SystemRandom, SecureRandom};

            let rng = SystemRandom::new();
            let mut key_bytes = vec![0u8; super::ALGORITHM.key_len()];
            let mut nonce_bytes = vec![0u8; super::ALGORITHM.nonce_len()];
            rng.fill(&mut key_bytes).unwrap();
            let key = SealingKey::new(super::ALGORITHM, &key_bytes).unwrap();

            let data = [1u8; DATA_SIZE];
            let out_suffix_capacity = super::ALGORITHM.max_overhead_len();
            let mut in_out = vec![1u8; data.len() + out_suffix_capacity];

            b.bytes = DATA_SIZE as u64;
            b.iter(|| {
                rng.fill(&mut nonce_bytes).unwrap();
                aead::seal_in_place(&key, &nonce_bytes, &mut in_out,
                                    out_suffix_capacity, &[]).unwrap()
            })
        }

        #[bench]
        fn bench_raw_open(b: &mut test::Bencher) {
            use ring::aead;
            use ring::aead::{SealingKey, OpeningKey, Algorithm};
            use ring::rand::{SystemRandom, SecureRandom};

            let rng = SystemRandom::new();
            let mut key_bytes = vec![0u8; super::ALGORITHM.key_len()];
            let mut nonce_bytes = vec![0u8; super::ALGORITHM.nonce_len()];
            rng.fill(&mut key_bytes).unwrap();
            rng.fill(&mut nonce_bytes).unwrap();
            let key = SealingKey::new(super::ALGORITHM, &key_bytes).unwrap();
            let opening_key = OpeningKey::new(super::ALGORITHM, &key_bytes).unwrap();

            let data = [1u8; DATA_SIZE];
            let out_suffix_capacity = super::ALGORITHM.max_overhead_len();
            let mut in_out = vec![1u8; data.len() + out_suffix_capacity];

            b.bytes = DATA_SIZE as u64;

            let sealed_len = aead::seal_in_place(&key, &nonce_bytes, &mut in_out,
                                                 out_suffix_capacity, &[]).unwrap();
            b.iter(|| aead::open_in_place(&opening_key, &nonce_bytes,
                                               0, &mut in_out, &[]))
        }

        #[bench]
        fn bench_seal(b: &mut test::Bencher) {
            let key = super::gen_key();
            let mut handler = super::Handler::new(&key);
            let mut buf = vec![0u8; super::super::MAX_MESSAGE_SIZE];
            b.bytes = DATA_SIZE as u64;
            b.iter(|| handler.seal(&mut buf, DATA_SIZE))
        }

        #[bench]
        fn bench_open(b: &mut test::Bencher) {
            let key = super::gen_key();
            let mut handler = super::Handler::new(&key);
            let mut buf = vec![0u8; super::super::MAX_MESSAGE_SIZE];
            let sealed_len = handler.seal(&mut buf, DATA_SIZE).unwrap();
            b.bytes = DATA_SIZE as u64;
            b.iter(|| handler.open(&mut buf[..sealed_len]))
        }
    }
}

fn new_udt_socket() -> UdtSocket {
    udt::init();
    let sock = UdtSocket::new(SocketFamily::AFInet, SocketType::Datagram).unwrap();
    sock.setsockopt(UdtOpts::UDP_RCVBUF, UDT_BUF_SIZE).unwrap();
    sock.setsockopt(UdtOpts::UDP_SNDBUF, UDT_BUF_SIZE).unwrap();
    sock
}

fn send(sock: &UdtSocket, key: &aead::SealingKey, buf: &mut [u8], len: usize) -> Result<(), UdtError> {
    unimplemented!();
    // FIXME don't unwrap, create an Error struct that can handle everything
    // if let Ok(sealed_len) = crypto::seal(buf, len, key) {
    //     sock.sendmsg(&buf[..sealed_len]).map(|_| ())
    // } else {
    //     Err(UdtError {
    //         err_code: -1,
    //         err_msg: "encryption failure".into(),
    //     })
    // }
}

fn recv(sock: &UdtSocket, key: &aead::OpeningKey, buf: &mut [u8]) -> Result<usize, UdtError> {
    unimplemented!();
    // let size = try!(sock.recvmsg(buf));
    // crypto::open(&mut buf[..size], key).map_err(|_| {
    //     UdtError {
    //         err_code: -1,
    //         err_msg: String::from("decryption failure"),
    //     }
    // })
}

pub struct PortRange {
    start: u16,
    end: u16,
}

pub trait Transceiver {
    fn send(&self, buf: &[u8]) -> Result<(), UdtError>;
    fn recv(&self, buf: &mut [u8]) -> Result<usize, UdtError>;
    fn close(&self) -> Result<(), UdtError>;
}

pub struct Server {
    pub ip_addr: IpAddr,
    pub port: u16,
    crypto: crypto::Handler,
    sock: UdtSocket,
}

pub struct Client {
    addr: SocketAddr,
    sock: UdtSocket,
    crypto: crypto::Handler,
}

pub struct ServerConnection<'a> {
    crypto: &'a crypto::Handler,
    sock: UdtSocket,
}

impl Client {
    pub fn new(addr: SocketAddr, key: &[u8]) -> Client {
        let sock = new_udt_socket();
        Client {
            addr: addr,
            sock: sock,
            crypto: crypto::Handler::new(key),
        }
    }

    pub fn connect(&self) -> Result<(), UdtError> {
        self.sock.connect(self.addr)
    }
}

impl Transceiver for Client {
    fn send(&self, buf: &[u8]) -> Result<(), UdtError> {
        unimplemented!();
        // send(&self.sock, &self.sealing_key, buf)
    }

    fn recv(&self, buf: &mut [u8]) -> Result<usize, UdtError> {
        unimplemented!();
        // recv(&self.sock, &self.opening_key, buf)
    }

    fn close(&self) -> Result<(), UdtError> {
        self.sock.close()
    }
}

impl Server {
    pub fn get_open_port(range: &PortRange) -> Result<u16, ()> {
        for p in range.start..range.end {
            if let Ok(_) = UdpSocket::bind(&format!("0.0.0.0:{}", p)[..]) {
                return Ok(p);
            }
        }
        Err(())
    }

    pub fn new(ip_addr: IpAddr, port: u16, key: &[u8]) -> Server {
        let sock = new_udt_socket();
        sock.bind(SocketAddr::new(ip_addr, port)).unwrap();
        Server {
            sock: sock,
            ip_addr: ip_addr,
            port: port,
            crypto: crypto::Handler::new(key),
        }
    }

    pub fn listen(&self) -> Result<(), UdtError> {
        self.sock.listen(2)
    }

    pub fn accept(&self) -> Result<ServerConnection, UdtError> {
        self.sock.accept().map(|(sock, _)| {
            ServerConnection {
                crypto: &self.crypto,
                sock: sock,
            }
        })
    }
}

impl<'a> ServerConnection<'a> {
    pub fn getpeer(&self) -> Result<SocketAddr, UdtError> {
        self.sock.getpeername()
    }
}

impl<'a> Transceiver for ServerConnection<'a> {
    fn send(&self, buf: &[u8]) -> Result<(), UdtError> {
        unimplemented!();
        // send(&self.sock, self.key, buf)
    }

    fn recv(&self, buf: &mut [u8]) -> Result<usize, UdtError> {
        unimplemented!();
        // recv(&self.sock, self.key, buf)
    }

    fn close(&self) -> Result<(), UdtError> {
        self.sock.close()
    }
}

impl<'a> PortRange {
    fn new(start: u16, end: u16) -> Result<PortRange, &'a str> {
        if start > end {
            Err("range end must be greater than or equal to start")
        } else {
            Ok(PortRange {
                start: start,
                end: end,
            })
        }
    }

    pub fn from(s: &str) -> Result<PortRange, &'a str> {
        let sections: Vec<&str> = s.split('-').collect();
        if sections.len() != 2 {
            return Err("Range must be specified in the form of \"<start>-<end>\"");
        }
        let (start, end) = (sections[0].parse::<u16>(), sections[1].parse::<u16>());
        if start.is_err() || end.is_err() {
            return Err("improperly formatted port range");
        }
        PortRange::new(start.unwrap(), end.unwrap())
    }
}

impl fmt::Display for PortRange {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}-{}", self.start, self.end)
    }
}
<|endoftext|>


<commit_before>use libc;
use std::collections::hash_map::Entry;
use std::collections::HashMap;
use std::ptr;
use std::rc::Rc;

use handle::{self, RsvgHandle};
use node::{Node, RsvgNode};
use util::utf8_cstr;

pub enum RsvgDefs {}

pub struct Defs {
    handle: *const RsvgHandle,
    nodes: HashMap<String, Rc<Node>>,
    externs: HashMap<String, *const RsvgHandle>,
}

impl Defs {
    fn new(handle: *const RsvgHandle) -> Defs {
        Defs {
            handle,
            nodes: Default::default(),
            externs: Default::default(),
        }
    }

    pub fn insert(&mut self, id: &str, node: &Rc<Node>) {
        self.nodes.entry(id.to_string()).or_insert(node.clone());
    }

    pub fn lookup(&mut self, name: &str) -> Option<&Rc<Node>> {
        match name.rfind('#') {
            None => None,
            Some(p) if p == 0 => self.nodes.get(&name[1..]),
            Some(p) => {
                let handle = self.get_extern_handle(&name[..p]);
                if handle.is_null() {
                    None
                } else {
                    handle::get_defs(handle).nodes.get(&name[(p + 1)..])
                }
            }
        }
    }

    fn get_extern_handle(&mut self, possibly_relative_uri: &str) -> *const RsvgHandle {
        handle::resolve_uri(self.handle, possibly_relative_uri).map_or(
            ptr::null(),
            |uri| match self.externs.entry(uri) {
                Entry::Occupied(e) => *(e.get()),
                Entry::Vacant(e) => {
                    let handle = handle::load_extern(self.handle, e.key());
                    if !handle.is_null() {
                        e.insert(handle);
                    }
                    handle
                }
            },
        )
    }
}

/// Represents a possibly non-canonical URI with an optional fragment identifier
///
/// Sometimes in SVG element references (e.g. the `href` in the `<feImage>` element) we
/// must decide between referencing an external file, or using a plain fragment identifier
/// like `href="#foo"` as a reference to an SVG element in the same file as the one being
/// processes.  This enum makes that distinction.
#[derive(Debug, PartialEq)]
pub enum Reference<'a> {
    PlainUri(&'a str),
    FragmentId(&'a str),
    UriWithFragmentId(&'a str, &'a str),
}

impl<'a> Reference<'a> {
    pub fn parse(s: &str) -> Result<Reference, ()> {
        let (uri, fragment) = match s.rfind('#') {
            None => (Some(s), None),
            Some(p) if p == 0 => (None, Some(&s[1..])),
            Some(p) => (Some(&s[..p]), Some(&s[(p + 1)..])),
        };

        match (uri, fragment) {
            (None, Some(f)) if f.len() == 0 => Err(()),
            (None, Some(f)) => Ok(Reference::FragmentId(f)),
            (Some(u), _) if u.len() == 0 => Err(()),
            (Some(u), None) => Ok(Reference::PlainUri(u)),
            (Some(_u), Some(f)) if f.len() == 0 => Err(()),
            (Some(u), Some(f)) => Ok(Reference::UriWithFragmentId(u, f)),
            (_, _) => Err(()),
        }
    }
}

#[no_mangle]
pub extern "C" fn rsvg_defs_new(handle: *const RsvgHandle) -> *mut RsvgDefs {
    Box::into_raw(Box::new(Defs::new(handle))) as *mut RsvgDefs
}

#[no_mangle]
pub extern "C" fn rsvg_defs_free(defs: *mut RsvgDefs) {
    assert!(!defs.is_null());

    unsafe {
        let defs = { &mut *(defs as *mut Defs) };
        Box::from_raw(defs);
    }
}

#[no_mangle]
pub extern "C" fn rsvg_defs_lookup(
    defs: *mut RsvgDefs,
    name: *const libc::c_char,
) -> *const RsvgNode {
    assert!(!defs.is_null());
    assert!(!name.is_null());

    let defs = unsafe { &mut *(defs as *mut Defs) };
    let name = unsafe { utf8_cstr(name) };

    match defs.lookup(name) {
        Some(n) => n as *const RsvgNode,
        None => ptr::null(),
    }
}

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

    #[test]
    fn reference_kinds() {
        assert_eq!(Reference::parse("uri"), Ok(Reference::PlainUri("uri")));
        assert_eq!(
            Reference::parse("#fragment"),
            Ok(Reference::FragmentId("fragment"))
        );
        assert_eq!(
            Reference::parse("uri#fragment"),
            Ok(Reference::UriWithFragmentId("uri", "fragment"))
        );
    }

    #[test]
    fn reference_errors() {
        assert!(Reference::parse("").is_err());
        assert!(Reference::parse("#").is_err());
        assert!(Reference::parse("uri#").is_err());
    }
}
<commit_msg>Defs::lookup() - Use Reference::parse() instead of doing the parsing here<commit_after>use libc;
use std::collections::hash_map::Entry;
use std::collections::HashMap;
use std::ptr;
use std::rc::Rc;

use handle::{self, RsvgHandle};
use node::{Node, RsvgNode};
use util::utf8_cstr;

pub enum RsvgDefs {}

pub struct Defs {
    handle: *const RsvgHandle,
    nodes: HashMap<String, Rc<Node>>,
    externs: HashMap<String, *const RsvgHandle>,
}

impl Defs {
    fn new(handle: *const RsvgHandle) -> Defs {
        Defs {
            handle,
            nodes: Default::default(),
            externs: Default::default(),
        }
    }

    pub fn insert(&mut self, id: &str, node: &Rc<Node>) {
        self.nodes.entry(id.to_string()).or_insert(node.clone());
    }

    /// Returns a node from an URI reference, or `None`
    ///
    /// This may return a node within the same RSVG handle, or a node in a secondary RSVG
    /// handle that is referenced by the current one.  If the element's id is not found,
    /// returns `None`.
    pub fn lookup(&mut self, name: &str) -> Option<&Rc<Node>> {
        if let Ok(reference) = Reference::parse(name) {
            match reference {
                Reference::PlainUri(_) => None,
                Reference::FragmentId(fragment) => self.nodes.get(fragment),
                Reference::UriWithFragmentId(uri, fragment) => {
                    let handle = self.get_extern_handle(uri);
                    if handle.is_null() {
                        None
                    } else {
                        handle::get_defs(handle).nodes.get(fragment)
                    }
                }
            }
        } else {
            None
        }
    }

    fn get_extern_handle(&mut self, possibly_relative_uri: &str) -> *const RsvgHandle {
        handle::resolve_uri(self.handle, possibly_relative_uri).map_or(
            ptr::null(),
            |uri| match self.externs.entry(uri) {
                Entry::Occupied(e) => *(e.get()),
                Entry::Vacant(e) => {
                    let handle = handle::load_extern(self.handle, e.key());
                    if !handle.is_null() {
                        e.insert(handle);
                    }
                    handle
                }
            },
        )
    }
}

/// Represents a possibly non-canonical URI with an optional fragment identifier
///
/// Sometimes in SVG element references (e.g. the `href` in the `<feImage>` element) we
/// must decide between referencing an external file, or using a plain fragment identifier
/// like `href="#foo"` as a reference to an SVG element in the same file as the one being
/// processes.  This enum makes that distinction.
#[derive(Debug, PartialEq)]
pub enum Reference<'a> {
    PlainUri(&'a str),
    FragmentId(&'a str),
    UriWithFragmentId(&'a str, &'a str),
}

impl<'a> Reference<'a> {
    pub fn parse(s: &str) -> Result<Reference, ()> {
        let (uri, fragment) = match s.rfind('#') {
            None => (Some(s), None),
            Some(p) if p == 0 => (None, Some(&s[1..])),
            Some(p) => (Some(&s[..p]), Some(&s[(p + 1)..])),
        };

        match (uri, fragment) {
            (None, Some(f)) if f.len() == 0 => Err(()),
            (None, Some(f)) => Ok(Reference::FragmentId(f)),
            (Some(u), _) if u.len() == 0 => Err(()),
            (Some(u), None) => Ok(Reference::PlainUri(u)),
            (Some(_u), Some(f)) if f.len() == 0 => Err(()),
            (Some(u), Some(f)) => Ok(Reference::UriWithFragmentId(u, f)),
            (_, _) => Err(()),
        }
    }
}

#[no_mangle]
pub extern "C" fn rsvg_defs_new(handle: *const RsvgHandle) -> *mut RsvgDefs {
    Box::into_raw(Box::new(Defs::new(handle))) as *mut RsvgDefs
}

#[no_mangle]
pub extern "C" fn rsvg_defs_free(defs: *mut RsvgDefs) {
    assert!(!defs.is_null());

    unsafe {
        let defs = { &mut *(defs as *mut Defs) };
        Box::from_raw(defs);
    }
}

#[no_mangle]
pub extern "C" fn rsvg_defs_lookup(
    defs: *mut RsvgDefs,
    name: *const libc::c_char,
) -> *const RsvgNode {
    assert!(!defs.is_null());
    assert!(!name.is_null());

    let defs = unsafe { &mut *(defs as *mut Defs) };
    let name = unsafe { utf8_cstr(name) };

    match defs.lookup(name) {
        Some(n) => n as *const RsvgNode,
        None => ptr::null(),
    }
}

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

    #[test]
    fn reference_kinds() {
        assert_eq!(Reference::parse("uri"), Ok(Reference::PlainUri("uri")));
        assert_eq!(
            Reference::parse("#fragment"),
            Ok(Reference::FragmentId("fragment"))
        );
        assert_eq!(
            Reference::parse("uri#fragment"),
            Ok(Reference::UriWithFragmentId("uri", "fragment"))
        );
    }

    #[test]
    fn reference_errors() {
        assert!(Reference::parse("").is_err());
        assert!(Reference::parse("#").is_err());
        assert!(Reference::parse("uri#").is_err());
    }
}
<|endoftext|>


<commit_before>use {
    BackEnd,
    Borrowed,
    Color,
    Draw,
    Field,
    Image,
    Matrix2d,
    Rectangle,
    Value,
    View,
};
use vecmath::{
    identity,
};
use triangulation::{
    rect_tri_list_xy_f32,
    rect_tri_list_rgba_f32,
    rect_tri_list_uv_f32,
};
use internal::{
    CanColor,
    CanRectangle,
    CanTransform,
    HasColor,
    HasRectangle,
    HasTransform,
};

/// An image rectangle context.
pub struct ImageRectangleColorContext<'a> {
    /// Base/original transformation.
    pub base: Field<'a, Matrix2d>,
    /// Current transformation.
    pub transform: Field<'a, Matrix2d>,
    /// Current rectangle.
    pub rect: Field<'a, Rectangle>,
    /// Current image.
    pub image: Field<'a, Image>,
    /// Current color.
    pub color: Field<'a, Color>,
}

impl<'a> HasTransform<'a, Matrix2d> for ImageRectangleColorContext<'a> {
    #[inline(always)]
    fn get_transform(&'a self) -> &'a Matrix2d {
        self.transform.get()
    }
}

impl<'a> CanTransform<'a, ImageRectangleColorContext<'a>, Matrix2d> for ImageRectangleColorContext<'a> {
    #[inline(always)]
    fn transform(&'a self, value: Matrix2d) -> ImageRectangleColorContext<'a> {
        ImageRectangleColorContext {
            base: Borrowed(self.base.get()),
            transform: Value(value),
            rect: Borrowed(self.rect.get()),
            image: Borrowed(self.image.get()),
            color: Borrowed(self.color.get()),
        }
    }
}

impl<'a> HasColor<'a, Color> for ImageRectangleColorContext<'a> {
    #[inline(always)]
    fn get_color(&'a self) -> &'a Color {
        self.color.get()
    }
}

impl<'a> CanColor<'a, ImageRectangleColorContext<'a>, Color> for ImageRectangleColorContext<'a> {
    #[inline(always)]
    fn color(&'a self, value: Color) -> ImageRectangleColorContext<'a> {
        ImageRectangleColorContext {
            base: Borrowed(self.base.get()),
            transform: Borrowed(self.transform.get()),
            color: Value(value),
            rect: Borrowed(self.rect.get()),
            image: Borrowed(self.image.get()),
        }
    }
}

impl<'a> HasRectangle<'a, Rectangle> for ImageRectangleColorContext<'a> {
    #[inline(always)]
    fn get_rectangle(&'a self) -> &'a Rectangle {
        self.rect.get()
    }
}

impl<'a> CanRectangle<'a, ImageRectangleColorContext<'a>, Rectangle> for ImageRectangleColorContext<'a> {
    #[inline(always)]
    fn rectangle(&'a self, rect: Rectangle) -> ImageRectangleColorContext<'a> {
        ImageRectangleColorContext {
            base: Borrowed(self.base.get()),
            transform: Borrowed(self.transform.get()),
            rect: Value(rect),
            image: Borrowed(self.image.get()),
            color: Borrowed(self.color.get()),
        }
    }
}

impl<'a> Draw<'a> for ImageRectangleColorContext<'a> {
    fn draw<B: BackEnd>(&'a self, back_end: &mut B) {
        if back_end.supports_single_texture()
        && back_end.supports_tri_list_xy_f32_rgba_f32_uv_f32() {
            let rect = self.rect.get();
            let color = self.color.get();
            let color: [f32, ..4] = [color[0], color[1], color[2], color[3]];
            let texture_id = self.image.get().texture_id;
            // Complete transparency does not need to be rendered.
            if color[3] == 0.0 { return; }
            // Turn on alpha blending if not completely opaque or if the texture has alpha channel.
            let needs_alpha = color[3] != 1.0 || back_end.has_texture_alpha(texture_id);
            if needs_alpha { back_end.enable_alpha_blend(); }
            back_end.enable_single_texture(texture_id);
            back_end.tri_list_xy_f32_rgba_f32_uv_f32(
                rect_tri_list_xy_f32(self.transform.get(), rect),
                rect_tri_list_rgba_f32(color),
                rect_tri_list_uv_f32(self.image.get())
            );
            back_end.disable_single_texture();
            if needs_alpha { back_end.disable_alpha_blend(); }
        } else {
            unimplemented!();
        }
    }
}

impl<'a> View<'a> for ImageRectangleColorContext<'a> {
    #[inline(always)]
    fn view(&'a self) -> ImageRectangleColorContext<'a> {
        ImageRectangleColorContext {
            base: Borrowed(self.base.get()),
            transform: Borrowed(self.base.get()),
            rect: Borrowed(self.rect.get()),
            image: Borrowed(self.image.get()),
            color: Borrowed(self.color.get()),
        }
    }

    #[inline(always)]
    fn reset(&'a self) -> ImageRectangleColorContext<'a> {
        ImageRectangleColorContext {
            base: Borrowed(self.base.get()),
            transform: Value(identity()),
            rect: Borrowed(self.rect.get()),
            image: Borrowed(self.image.get()),
            color: Borrowed(self.color.get()),
        }
    }

    #[inline(always)]
    fn store_view(&'a self) -> ImageRectangleColorContext<'a> {
        ImageRectangleColorContext {
            base: Borrowed(self.transform.get()),
            transform: Borrowed(self.transform.get()),
            rect: Borrowed(self.rect.get()),
            image: Borrowed(self.image.get()),
            color: Borrowed(self.color.get()),
        }
    }
}

<commit_msg>Implemented `Clear` for `ImageRectangleColorContext`<commit_after>use {
    BackEnd,
    Borrowed,
    Clear,
    Color,
    Draw,
    Field,
    Image,
    Matrix2d,
    Rectangle,
    Value,
    View,
};
use vecmath::{
    identity,
};
use triangulation::{
    rect_tri_list_xy_f32,
    rect_tri_list_rgba_f32,
    rect_tri_list_uv_f32,
};
use internal::{
    CanColor,
    CanRectangle,
    CanTransform,
    HasColor,
    HasRectangle,
    HasTransform,
};

/// An image rectangle context.
pub struct ImageRectangleColorContext<'a> {
    /// Base/original transformation.
    pub base: Field<'a, Matrix2d>,
    /// Current transformation.
    pub transform: Field<'a, Matrix2d>,
    /// Current rectangle.
    pub rect: Field<'a, Rectangle>,
    /// Current image.
    pub image: Field<'a, Image>,
    /// Current color.
    pub color: Field<'a, Color>,
}

impl<'a> HasTransform<'a, Matrix2d> for ImageRectangleColorContext<'a> {
    #[inline(always)]
    fn get_transform(&'a self) -> &'a Matrix2d {
        self.transform.get()
    }
}

impl<'a> CanTransform<'a, ImageRectangleColorContext<'a>, Matrix2d> for ImageRectangleColorContext<'a> {
    #[inline(always)]
    fn transform(&'a self, value: Matrix2d) -> ImageRectangleColorContext<'a> {
        ImageRectangleColorContext {
            base: Borrowed(self.base.get()),
            transform: Value(value),
            rect: Borrowed(self.rect.get()),
            image: Borrowed(self.image.get()),
            color: Borrowed(self.color.get()),
        }
    }
}

impl<'a> HasColor<'a, Color> for ImageRectangleColorContext<'a> {
    #[inline(always)]
    fn get_color(&'a self) -> &'a Color {
        self.color.get()
    }
}

impl<'a> CanColor<'a, ImageRectangleColorContext<'a>, Color> for ImageRectangleColorContext<'a> {
    #[inline(always)]
    fn color(&'a self, value: Color) -> ImageRectangleColorContext<'a> {
        ImageRectangleColorContext {
            base: Borrowed(self.base.get()),
            transform: Borrowed(self.transform.get()),
            color: Value(value),
            rect: Borrowed(self.rect.get()),
            image: Borrowed(self.image.get()),
        }
    }
}

impl<'a> HasRectangle<'a, Rectangle> for ImageRectangleColorContext<'a> {
    #[inline(always)]
    fn get_rectangle(&'a self) -> &'a Rectangle {
        self.rect.get()
    }
}

impl<'a> CanRectangle<'a, ImageRectangleColorContext<'a>, Rectangle> for ImageRectangleColorContext<'a> {
    #[inline(always)]
    fn rectangle(&'a self, rect: Rectangle) -> ImageRectangleColorContext<'a> {
        ImageRectangleColorContext {
            base: Borrowed(self.base.get()),
            transform: Borrowed(self.transform.get()),
            rect: Value(rect),
            image: Borrowed(self.image.get()),
            color: Borrowed(self.color.get()),
        }
    }
}

impl<'a> Draw<'a> for ImageRectangleColorContext<'a> {
    fn draw<B: BackEnd>(&'a self, back_end: &mut B) {
        if back_end.supports_single_texture()
        && back_end.supports_tri_list_xy_f32_rgba_f32_uv_f32() {
            let rect = self.rect.get();
            let color = self.color.get();
            let color: [f32, ..4] = [color[0], color[1], color[2], color[3]];
            let texture_id = self.image.get().texture_id;
            // Complete transparency does not need to be rendered.
            if color[3] == 0.0 { return; }
            // Turn on alpha blending if not completely opaque or if the texture has alpha channel.
            let needs_alpha = color[3] != 1.0 || back_end.has_texture_alpha(texture_id);
            if needs_alpha { back_end.enable_alpha_blend(); }
            back_end.enable_single_texture(texture_id);
            back_end.tri_list_xy_f32_rgba_f32_uv_f32(
                rect_tri_list_xy_f32(self.transform.get(), rect),
                rect_tri_list_rgba_f32(color),
                rect_tri_list_uv_f32(self.image.get())
            );
            back_end.disable_single_texture();
            if needs_alpha { back_end.disable_alpha_blend(); }
        } else {
            unimplemented!();
        }
    }
}

impl<'a> View<'a> for ImageRectangleColorContext<'a> {
    #[inline(always)]
    fn view(&'a self) -> ImageRectangleColorContext<'a> {
        ImageRectangleColorContext {
            base: Borrowed(self.base.get()),
            transform: Borrowed(self.base.get()),
            rect: Borrowed(self.rect.get()),
            image: Borrowed(self.image.get()),
            color: Borrowed(self.color.get()),
        }
    }

    #[inline(always)]
    fn reset(&'a self) -> ImageRectangleColorContext<'a> {
        ImageRectangleColorContext {
            base: Borrowed(self.base.get()),
            transform: Value(identity()),
            rect: Borrowed(self.rect.get()),
            image: Borrowed(self.image.get()),
            color: Borrowed(self.color.get()),
        }
    }

    #[inline(always)]
    fn store_view(&'a self) -> ImageRectangleColorContext<'a> {
        ImageRectangleColorContext {
            base: Borrowed(self.transform.get()),
            transform: Borrowed(self.transform.get()),
            rect: Borrowed(self.rect.get()),
            image: Borrowed(self.image.get()),
            color: Borrowed(self.color.get()),
        }
    }
}

impl<'a> Clear for ImageRectangleColorContext<'a> {
    fn clear<B: BackEnd>(&self, back_end: &mut B) {
        if back_end.supports_clear_rgba() {
            let color = self.color.get();
            back_end.clear_rgba(color[0], color[1], color[2], color[3]);
        }
    }
}

<|endoftext|>


<commit_before>//! Convenience functions for making requests over HTTP
//!
//! This module provides a few functions that make it easier to make
//! requests over HTTP in a consistent manner e.g. with the 'User-Agent` header
//! set and respecting cache control headers in responses. In addition to reducing
//! the number of network requests for the client, several APIs ask clients
//! to implement caching to reduce load on their servers.

use std::fs::create_dir_all;
use std::io::Write;
use std::{env, fs::File, io, path::Path};

use eyre::Result;
use http_cache_reqwest::{CACacheManager, Cache, CacheMode, HttpCache};
use once_cell::sync::Lazy;
use reqwest::header::{HeaderMap, HeaderName, HeaderValue};
use reqwest::Response;
use reqwest_middleware::{ClientBuilder, ClientWithMiddleware};
use serde::de::DeserializeOwned;
use serde::Serialize;
use tempfile::NamedTempFile;

// Re-exports for consumers of this crate
pub use reqwest;
pub use reqwest::header as headers;
pub use reqwest_middleware;
pub use serde;
pub use serde_json;
pub use tempfile;
pub use url;

pub static USER_AGENT: &str = concat!("stencila/", env!("CARGO_PKG_VERSION"),);

/// Get the directory of the HTTP cache
pub fn cache_dir() -> String {
    let user_cache_dir = dirs::cache_dir().unwrap_or_else(|| env::current_dir().unwrap());
    match env::consts::OS {
        "macos" | "windows" => user_cache_dir.join("Stencila").join("HTTP-Cache"),
        _ => user_cache_dir.join("stencila").join("http-cache"),
    }
    .to_string_lossy()
    .to_string()
}

pub static CLIENT: Lazy<ClientWithMiddleware> = Lazy::new(|| {
    let client = reqwest::Client::builder()
        .user_agent(USER_AGENT)
        .build()
        .expect("Should be able to build HTTP client");
    let caching_middleware = Cache(HttpCache {
        mode: CacheMode::Default,
        manager: CACacheManager { path: cache_dir() },
        options: None,
    });
    ClientBuilder::new(client).with(caching_middleware).build()
});

/// Convert an array of tuples to a `reqwest::HeaderMap`
fn header_map(headers: &[(HeaderName, String)]) -> Result<HeaderMap> {
    let mut header_map = HeaderMap::new();
    header_map.insert(headers::ACCEPT, HeaderValue::from_str("application/json")?);
    for (key, value) in headers {
        header_map.insert(key, value.parse()?);
    }
    Ok(header_map)
}

/// Send a GET request
pub async fn get<T: DeserializeOwned>(url: &str) -> Result<T> {
    get_with(url, &[]).await
}

/// Send a GET request with additional request headers
pub async fn get_with<T: DeserializeOwned>(
    url: &str,
    headers: &[(HeaderName, String)],
) -> Result<T> {
    let response = CLIENT
        .get(url)
        .headers(header_map(headers)?)
        .send()
        .await?
        .error_for_status()?;

    let json = response.json().await?;

    Ok(json)
}

/// Send a POST request
pub async fn post<B: Serialize, T: DeserializeOwned>(url: &str, body: B) -> Result<T> {
    post_with(url, body, &[]).await
}

/// Send a POST request with additional request headers
pub async fn post_with<B: Serialize, T: DeserializeOwned>(
    url: &str,
    body: B,
    headers: &[(HeaderName, String)],
) -> Result<T> {
    let response = post_response(url, body, headers).await?;
    let json = response.json().await?;
    Ok(json)
}

/// Send a POST request with additional request headers and return response
pub async fn post_response<B: Serialize>(
    url: &str,
    body: B,
    headers: &[(HeaderName, String)],
) -> Result<Response> {
    let response = CLIENT
        .post(url)
        .json(&body)
        .headers(header_map(headers)?)
        .send()
        .await?
        .error_for_status()?;
    Ok(response)
}

/// Send a DELETE request
pub async fn delete(url: &str) -> Result<()> {
    delete_with(url, &[]).await
}

/// Send a DELETE request with additional request headers
pub async fn delete_with(url: &str, headers: &[(HeaderName, String)]) -> Result<()> {
    CLIENT
        .delete(url)
        .headers(header_map(headers)?)
        .send()
        .await?
        .error_for_status()?;
    Ok(())
}

/// Download a file from to a path
pub async fn download(url: &str, path: &Path) -> Result<()> {
    download_with(url, path, &[]).await
}

/// Download a file from to a path with additional request headers
pub async fn download_with(url: &str, path: &Path, headers: &[(HeaderName, String)]) -> Result<()> {
    let response = CLIENT
        .get(url)
        .headers(header_map(headers)?)
        .send()
        .await?
        .error_for_status()?;

    if let Some(parent) = path.parent() {
        create_dir_all(parent)?
    }

    let bytes = response.bytes().await?;
    let mut file = File::create(&path)?;
    io::copy(&mut bytes.as_ref(), &mut file)?;
    file.flush()?;

    Ok(())
}

/// Download a file from to a path synchronously
pub fn download_sync(url: &str, path: &Path) -> Result<()> {
    download_with_sync(url, path, &[])
}

/// Download a file from to a path synchronously with additional request headers
pub fn download_with_sync(url: &str, path: &Path, headers: &[(HeaderName, String)]) -> Result<()> {
    let url = url.to_owned();
    let path = path.to_owned();
    let headers = headers.to_vec();
    let (sender, receiver) = std::sync::mpsc::channel();
    tokio::spawn(async move {
        let result = download_with(&url, &path, &headers).await;
        sender.send(result)
    });
    receiver.recv()?
}

/// Download a file from to a temporary file
pub async fn download_temp(url: &str, ext: Option<&str>) -> Result<NamedTempFile> {
    download_temp_with(url, ext, &[]).await
}

/// Download a file from to a temporary file with additional request headers
///
/// Returns a `NamedTempFile` which will remove the temporary file when it
/// is dropped. Be aware of that and the security implications of long-lived temp files:
/// https://docs.rs/tempfile/latest/tempfile/struct.NamedTempFile.html
pub async fn download_temp_with(
    url: &str,
    ext: Option<&str>,
    headers: &[(HeaderName, String)],
) -> Result<NamedTempFile> {
    let suffix = ext.map_or_else(String::new, |ext| [".", ext].concat());
    let temp = tempfile::Builder::new().suffix(&suffix).tempfile()?;
    download_with(url, temp.path(), headers).await?;
    Ok(temp)
}
<commit_msg>feat(HTTP utilities): Expose `get_response`<commit_after>//! Convenience functions for making requests over HTTP
//!
//! This module provides a few functions that make it easier to make
//! requests over HTTP in a consistent manner e.g. with the 'User-Agent` header
//! set and respecting cache control headers in responses. In addition to reducing
//! the number of network requests for the client, several APIs ask clients
//! to implement caching to reduce load on their servers.

use std::fs::create_dir_all;
use std::io::Write;
use std::{env, fs::File, io, path::Path};

use eyre::Result;
use http_cache_reqwest::{CACacheManager, Cache, CacheMode, HttpCache};
use once_cell::sync::Lazy;
use reqwest::header::{HeaderMap, HeaderName, HeaderValue};
use reqwest::Response;
use reqwest_middleware::{ClientBuilder, ClientWithMiddleware};
use serde::de::DeserializeOwned;
use serde::Serialize;
use tempfile::NamedTempFile;

// Re-exports for consumers of this crate
pub use reqwest;
pub use reqwest::header as headers;
pub use reqwest_middleware;
pub use serde;
pub use serde_json;
pub use tempfile;
pub use url;

pub static USER_AGENT: &str = concat!("stencila/", env!("CARGO_PKG_VERSION"),);

/// Get the directory of the HTTP cache
pub fn cache_dir() -> String {
    let user_cache_dir = dirs::cache_dir().unwrap_or_else(|| env::current_dir().unwrap());
    match env::consts::OS {
        "macos" | "windows" => user_cache_dir.join("Stencila").join("HTTP-Cache"),
        _ => user_cache_dir.join("stencila").join("http-cache"),
    }
    .to_string_lossy()
    .to_string()
}

pub static CLIENT: Lazy<ClientWithMiddleware> = Lazy::new(|| {
    let client = reqwest::Client::builder()
        .user_agent(USER_AGENT)
        .build()
        .expect("Should be able to build HTTP client");
    let caching_middleware = Cache(HttpCache {
        mode: CacheMode::Default,
        manager: CACacheManager { path: cache_dir() },
        options: None,
    });
    ClientBuilder::new(client).with(caching_middleware).build()
});

/// Convert an array of tuples to a `reqwest::HeaderMap`
fn header_map(headers: &[(HeaderName, String)]) -> Result<HeaderMap> {
    let mut header_map = HeaderMap::new();
    header_map.insert(headers::ACCEPT, HeaderValue::from_str("application/json")?);
    for (key, value) in headers {
        header_map.insert(key, value.parse()?);
    }
    Ok(header_map)
}

/// Send a GET request
pub async fn get<T: DeserializeOwned>(url: &str) -> Result<T> {
    get_with(url, &[]).await
}

/// Send a GET request with additional request headers
pub async fn get_with<T: DeserializeOwned>(
    url: &str,
    headers: &[(HeaderName, String)],
) -> Result<T> {
    let response = get_response(url, headers).await?;
    let json = response.error_for_status()?.json().await?;
    Ok(json)
}

/// Send a GET request with additional request headers and return response
///
/// Use this when you want to handle response errors or content yourself.
pub async fn get_response(url: &str, headers: &[(HeaderName, String)]) -> Result<Response> {
    let response = CLIENT.get(url).headers(header_map(headers)?).send().await?;
    Ok(response)
}

/// Send a POST request
pub async fn post<B: Serialize, T: DeserializeOwned>(url: &str, body: B) -> Result<T> {
    post_with(url, body, &[]).await
}

/// Send a POST request with additional request headers
pub async fn post_with<B: Serialize, T: DeserializeOwned>(
    url: &str,
    body: B,
    headers: &[(HeaderName, String)],
) -> Result<T> {
    let response = post_response(url, body, headers).await?;
    let json = response.error_for_status()?.json().await?;
    Ok(json)
}

/// Send a POST request with additional request headers and return response
pub async fn post_response<B: Serialize>(
    url: &str,
    body: B,
    headers: &[(HeaderName, String)],
) -> Result<Response> {
    let response = CLIENT
        .post(url)
        .json(&body)
        .headers(header_map(headers)?)
        .send()
        .await?;
    Ok(response)
}

/// Send a DELETE request
pub async fn delete(url: &str) -> Result<()> {
    delete_with(url, &[]).await
}

/// Send a DELETE request with additional request headers
pub async fn delete_with(url: &str, headers: &[(HeaderName, String)]) -> Result<()> {
    CLIENT
        .delete(url)
        .headers(header_map(headers)?)
        .send()
        .await?
        .error_for_status()?;
    Ok(())
}

/// Download a file from to a path
pub async fn download(url: &str, path: &Path) -> Result<()> {
    download_with(url, path, &[]).await
}

/// Download a file from to a path with additional request headers
pub async fn download_with(url: &str, path: &Path, headers: &[(HeaderName, String)]) -> Result<()> {
    let response = CLIENT
        .get(url)
        .headers(header_map(headers)?)
        .send()
        .await?
        .error_for_status()?;

    if let Some(parent) = path.parent() {
        create_dir_all(parent)?
    }

    let bytes = response.bytes().await?;
    let mut file = File::create(&path)?;
    io::copy(&mut bytes.as_ref(), &mut file)?;
    file.flush()?;

    Ok(())
}

/// Download a file from to a path synchronously
pub fn download_sync(url: &str, path: &Path) -> Result<()> {
    download_with_sync(url, path, &[])
}

/// Download a file from to a path synchronously with additional request headers
pub fn download_with_sync(url: &str, path: &Path, headers: &[(HeaderName, String)]) -> Result<()> {
    let url = url.to_owned();
    let path = path.to_owned();
    let headers = headers.to_vec();
    let (sender, receiver) = std::sync::mpsc::channel();
    tokio::spawn(async move {
        let result = download_with(&url, &path, &headers).await;
        sender.send(result)
    });
    receiver.recv()?
}

/// Download a file from to a temporary file
pub async fn download_temp(url: &str, ext: Option<&str>) -> Result<NamedTempFile> {
    download_temp_with(url, ext, &[]).await
}

/// Download a file from to a temporary file with additional request headers
///
/// Returns a `NamedTempFile` which will remove the temporary file when it
/// is dropped. Be aware of that and the security implications of long-lived temp files:
/// https://docs.rs/tempfile/latest/tempfile/struct.NamedTempFile.html
pub async fn download_temp_with(
    url: &str,
    ext: Option<&str>,
    headers: &[(HeaderName, String)],
) -> Result<NamedTempFile> {
    let suffix = ext.map_or_else(String::new, |ext| [".", ext].concat());
    let temp = tempfile::Builder::new().suffix(&suffix).tempfile()?;
    download_with(url, temp.path(), headers).await?;
    Ok(temp)
}
<|endoftext|>