<commit_before>#![deny(warnings)]

extern crate time;
extern crate curl;
extern crate rustc_serialize;
extern crate openssl;

use std::io::prelude::*;

use curl::http;
use curl::http::body::ToBody;
use openssl::crypto::{hmac, hash};
use rustc_serialize::base64::{ToBase64, STANDARD};

pub struct Bucket {
    name: String,
    region: Option<String>,
    access_key: String,
    secret_key: String,
    proto: String,
}

impl Bucket {
    pub fn new(name: String,
               region: Option<String>,
               access_key: String,
               secret_key: String,
               proto: &str) -> Bucket {
        Bucket {
            name: name,
            region: region,
            access_key: access_key,
            secret_key: secret_key,
            proto: proto.to_string(),
        }
    }

    pub fn put<'a, 'b, T: ToBody<'b>>(&self, handle: &'a mut http::Handle,
                                      path: &str, content: T,
                                      content_type: &str)
                                      -> http::Request<'a, 'b> {
        let path = if path.starts_with("/") {&path[1..]} else {path};
        let host = self.host();
        let date = time::now().rfc822z().to_string();
        let auth = self.auth("PUT", &date, path, "", content_type);
        let url = format!("{}://{}/{}", self.proto, host, path);
        handle.put(&url[..], content)
              .header("Host", &host)
              .header("Date", &date)
              .header("Authorization", &auth)
              .content_type(content_type)
    }

    pub fn delete<'a, 'b>(&self, handle: &'a mut http::Handle, path: &str)
                          -> http::Request<'a, 'b> {
        let path = if path.starts_with("/") {&path[1..]} else {path};
        let host = self.host();
        let date = time::now().rfc822z().to_string();
        let auth = self.auth("DELETE", &date, path, "", "");
        let url = format!("{}://{}/{}", self.proto, host, path);
        handle.delete(&url[..])
              .header("Host", &host)
              .header("Date", &date)
              .header("Authorization", &auth)
    }

    pub fn host(&self) -> String {
        format!("{}.s3{}.amazonaws.com", self.name,
                match self.region {
                    Some(ref r) => format!("-{}", r),
                    None => String::new(),
                })
    }

    fn auth(&self, verb: &str, date: &str, path: &str,
            md5: &str, content_type: &str) -> String {
        let string = format!("{verb}\n{md5}\n{ty}\n{date}\n{headers}{resource}",
                             verb = verb,
                             md5 = md5,
                             ty = content_type,
                             date = date,
                             headers = "",
                             resource = format!("/{}/{}", self.name, path));
        let signature = {
            let mut hmac = hmac::HMAC::new(hash::Type::SHA1, self.secret_key.as_bytes());
            let _ = hmac.write_all(string.as_bytes());
            hmac.finish().to_base64(STANDARD)
        };
        format!("AWS {}:{}", self.access_key, signature)
    }
}
<commit_msg>Update s3 library for openssl and curl changes<commit_after>#![deny(warnings)]

extern crate time;
extern crate curl;
extern crate rustc_serialize;
extern crate openssl;

use std::io::prelude::*;

use curl::easy::{Easy, Transfer, List, ReadError};
use openssl::hash::MessageDigest;
use openssl::sign::Signer;
use openssl::pkey::PKey;
use rustc_serialize::base64::{ToBase64, STANDARD};

pub struct Bucket {
    name: String,
    region: Option<String>,
    access_key: String,
    secret_key: String,
    proto: String,
}

impl Bucket {
    pub fn new(name: String,
               region: Option<String>,
               access_key: String,
               secret_key: String,
               proto: &str) -> Bucket {
        Bucket {
            name: name,
            region: region,
            access_key: access_key,
            secret_key: secret_key,
            proto: proto.to_string(),
        }
    }

    pub fn put<'a, 'b>(&self,
                       easy: &'a mut Easy,
                       path: &str,
                       content: &'b mut Read,
                       content_type: &str)
                       -> Transfer<'a, 'b> {
        let path = if path.starts_with("/") {&path[1..]} else {path};
        let host = self.host();
        let date = time::now().rfc822z().to_string();
        let auth = self.auth("PUT", &date, path, "", content_type);
        let url = format!("{}://{}/{}", self.proto, host, path);

        let mut headers = List::new();
        headers.append(&format!("Host: {}", host)).unwrap();
        headers.append(&format!("Date: {}", date)).unwrap();
        headers.append(&format!("Authorization: {}", auth)).unwrap();
        headers.append(&format!("Content-Type: {}", content_type)).unwrap();

        easy.put(true).unwrap();
        easy.url(&url).unwrap();
        easy.http_headers(headers).unwrap();

        let mut transfer = easy.transfer();
        transfer.read_function(move |data| {
            content.read(data).map_err(|_| ReadError::Abort)
        }).unwrap();

        return transfer
    }

    pub fn delete<'a, 'b>(&self,
                          easy: &'a mut Easy,
                          path: &str)
                          -> Transfer<'a, 'b> {
        let path = if path.starts_with("/") {&path[1..]} else {path};
        let host = self.host();
        let date = time::now().rfc822z().to_string();
        let auth = self.auth("DELETE", &date, path, "", "");
        let url = format!("{}://{}/{}", self.proto, host, path);

        let mut headers = List::new();
        headers.append(&format!("Host: {}", host)).unwrap();
        headers.append(&format!("Date: {}", date)).unwrap();
        headers.append(&format!("Authorization: {}", auth)).unwrap();

        easy.custom_request("DELETE").unwrap();
        easy.url(&url).unwrap();
        easy.http_headers(headers).unwrap();

        return easy.transfer()
    }

    pub fn host(&self) -> String {
        format!("{}.s3{}.amazonaws.com", self.name,
                match self.region {
                    Some(ref r) => format!("-{}", r),
                    None => String::new(),
                })
    }

    fn auth(&self, verb: &str, date: &str, path: &str,
            md5: &str, content_type: &str) -> String {
        let string = format!("{verb}\n{md5}\n{ty}\n{date}\n{headers}{resource}",
                             verb = verb,
                             md5 = md5,
                             ty = content_type,
                             date = date,
                             headers = "",
                             resource = format!("/{}/{}", self.name, path));
        let signature = {
            let key = PKey::hmac(self.secret_key.as_bytes()).unwrap();
            let mut signer = Signer::new(MessageDigest::sha1(), &key).unwrap();
            signer.update(string.as_bytes()).unwrap();
            signer.finish().unwrap().to_base64(STANDARD)
        };
        format!("AWS {}:{}", self.access_key, signature)
    }
}
<|endoftext|>


<commit_before>use std::collections::{BTreeMap, HashMap};
use std::env;
use std::fs::File;
use std::process;

use applescript::{run as run_applescript};

use plist::Plist;

use tera::{Context, Tera};

use urls;

macro_rules! error(
  ($($arg:tt)*) => { { return Err(format!($($arg)*)) } }
);

/// Returns true/false if Safari is running.
pub fn is_safari_running() -> bool {
  let cmd_result = process::Command::new("ps")
    .arg("-eaf")
    .output()
    .expect("Unable to test if Safari is running.");
  for line in String::from_utf8(cmd_result.stdout).unwrap().lines() {
    if line.contains("Safari.app/Contents/MacOS/Safari") {
      return true
    }
  }
  false
}

/// Return a URL from a Safari window.
///
/// Given a (window, tab) pair, this function looks up the URL of the tab.
/// Note that it doesn't do any error handling, so will throw an execution
/// error if it fails.
///
/// * `window` - Window index.  1 is frontmost.  If None, assumes the
///              frontmost window.
/// * `tab` - Tab index.  1 is leftmost.  If None, assumes the frontmost tab.
///
pub fn get_url(window: Option<u32>, tab: Option<u32>) -> Result<String, String> {
  // If a tab isn't specified, assume the user wants the frontmost tab.
  let command = match window {
    Some(w_idx) => {
      match tab {
        Some(t_idx) => format!("tell application \"Safari\" to get URL of tab {} of window {}", t_idx, w_idx),
        None => format!("tell application \"Safari\" to get URL of document {}", w_idx),
      }
    },
    None => "tell application \"Safari\" to get URL of document 1".to_string(),
  };
  let output = run_applescript(&command);

  if output.status.success() {
    Ok(urls::tidy_url(output.stdout.trim()))
  } else {
    if output.stderr.contains("Invalid index") {
      error!("Invalid index: no such window or tab.")
    } else {
      error!("Unexpected error from osascript: {:?}", output.stderr)
    }
  }
}

/// Return a list of URLs from Safari.
///
/// This returns a list of URLs, one for every tab in Safari.  Iteration
/// order depends on AppleScript, which I don't think is guaranteed to be
/// stable (in particular, I think it depends on which window is frontmost).
///
pub fn get_all_urls() -> Vec<String> {
  let windows = get_window_tab_count_pairs();
  let mut urls = vec![];
  for window in windows {
    for tab in 1..window.tab_count {
      match get_url(Some(window.window_index), Some(tab)) {
        Ok(u) => urls.push(u),
        Err(_) => {},
      }
    }
  }
  urls
}

/// Convert URL patterns into AppleScript conditions.
///
/// These conditions can be used in an `if` statement in AppleScript to
/// decide whether a tab should be closed.  Three patterns are supported,
/// a limited regex syntax:
///
///     example.com             matches anywhere in the URL
///     ^http://examples.com    matches at the start of the URL
///     example.com/$           matches at the end of the URL
///
fn parse_conditions(url_patterns: Vec<&str>) -> Vec<String> {
  url_patterns.iter().map(|p|
    if p.starts_with("^") {
      format!("starts with \"{}\"", p.replace("^", ""))
    } else if p.ends_with("$") {
      format!("ends with \"{}\"", p.replace("$", ""))
    } else {
      format!("contains \"{}\"", p)
    }
  ).collect()
}

/// Tests for parse_conditions().
#[cfg(test)]
mod tests {
  use safari::parse_conditions;

  #[test]
  fn test_parse_conditions() {
    let patterns = vec!["github.com", "^facebook.com", "twitter.com$"];
    let expected = vec!["contains \"github.com\"", "starts with \"facebook.com\"", "ends with \"twitter.com\""];
    let actual = parse_conditions(patterns);
    assert_eq!(actual, expected);
  }
}

/// Close tabs in Safari that match URL patterns.
///
/// Takes a list of URL patterns, and tries to close any matching Safari tabs.
/// Bugginess in AppleScript means this isn't always perfect, but we can
/// have a go.  In general it will fail to close tabs, rather than close
/// the wrong tabs.
///
pub fn close_tabs(url_patterns: Vec<&str>) {
  let conditions = parse_conditions(url_patterns);

  let clean_tabs_template = include_str!("scripts/clean-tabs.scpt");
  let mut context = Context::new();
  context.add("conditions", &conditions);
  let script = Tera::one_off(&clean_tabs_template, context, false).unwrap();

  // Run it twice to get around weird AppleScript bugs.
  run_applescript(&script);
  run_applescript(&script);
}

/// Get the Bookmarks.plist dict for a given title
fn read_bookmarks_plist(title: &str) -> Result<Plist, String> {

  // All Safari data lives at ~/Library/Safari/Bookmarks.plist
  // TODO: There's probably a more idiomatic Rust-like way to write this.
  let mut plist_path = match env::home_dir() {
    Some(v) => v,
    None => error!("Unable to get home directory?"),
  };
  plist_path.push("Library/Safari/Bookmarks.plist");

  let file = match File::open(plist_path) {
    Ok(v) => v,
    Err(e) => error!("Unable to open ~/Library/Safari/Bookmarks.plist: {:?}", e),
  };

  let plist = match Plist::read(file) {
    Ok(v) => v,
    Err(e) => error!("Unable to read ~/Library/Safari/Bookmarks.plist: {:?}", e),
  };

  let data = match plist.as_dictionary() {
    Some(v) => v,
    None => error!("Unable to parse ~/Library/Safari/Bookmarks.plist as dictionary?"),
  };

  // The structure of Bookmarks.plist is as follows:
  //
  //     <dict>
  //       <key>Children</key>
  //       <array>
  //         <dict>
  //           <key>Title</key><string>History</string>
  //           ... dict data ...
  //         </dict>
  //         <dict>
  //           <key>Title</key><string>com.apple.ReadingList</string>
  //           ... dict data ...
  //         </dict>
  //         ... other array items ...
  //       </array>
  //     </dict>
  //
  let children = match data.get("Children") {
    Some(child_key) => match child_key.as_array() {
      Some(v) => v,
      None => error!("Top-level children key in ~/Library/Safari/Bookmarks.plist isn't an array?"),
    },
    None => error!("Unable to find top-level Children key in ~/Library/Safari/Bookmarks.plist"),
  };

  let mut matching_children = children.iter().filter(|d|
    match d.as_dictionary() {
      Some(dict) => match dict.get("Title") {
        Some(title_elem) => match title_elem.as_string() {
          Some(v) => (v == title),
          None => false,
        },
        None => false,
      },
      None => false,
    }
  );

  // Check we got one, and only one result.
  let result = match matching_children.next() {
    Some(v) => v,
    None => error!("Unable to find key {} in Bookmarks.plist", title),
  };

  match matching_children.next() {
    Some(_) => error!("Got more than one result for {} in Bookmarks.plist", title),
    None => {},
  };

  Ok(result.to_owned())
}

/// Return a list of URLs from Reading List.
///
/// Iteration order depends on the order in which they're stored in
/// Bookmarks.plist, which is usually (but not guaranteed to be) newest first.
///
pub fn get_reading_list_urls() -> Result<Vec<String>, String> {
  let plist = match read_bookmarks_plist("com.apple.ReadingList") {
    Ok(v) => v,
    Err(e) => return Err(e),
  };

  // TODO: All these unwrap() calls should probably be handled better
  let children = plist
    .as_dictionary().unwrap()
    .get("Children").unwrap()
    .as_array().unwrap();

  Ok(children
    .iter()
    .map(|child|
      child
        .as_dictionary().unwrap()
        .get("URLString").unwrap()
        .as_string().unwrap()
    )
    .map(|url| urls::tidy_url(url))
    .collect()
  )
}

/// Get the com.apple.Safari.plist preferences file
fn read_safari_plist() -> Result<BTreeMap<String, Plist>, String> {

  // All Safari data lives at ~/Library/Safari/Bookmarks.plist
  // TODO: There's probably a more idiomatic Rust-like way to write this.
  let mut plist_path = match env::home_dir() {
    Some(v) => v,
    None => error!("Unable to get home directory?"),
  };
  plist_path.push("Library/SyncedPreferences/com.apple.Safari.plist");

  let file = match File::open(plist_path) {
    Ok(v) => v,
    Err(e) => error!("Unable to open com.apple.Safari.plist: {:?}", e),
  };

  let plist = match Plist::read(file) {
    Ok(v) => v,
    Err(e) => error!("Unable to read com.apple.Safari.plist: {:?}", e),
  };

  let data = match plist.as_dictionary() {
    Some(v) => v,
    None => error!("Unable to parse com.apple.Safari.plist as dictionary?"),
  };

  // The structure of com.apple.Safari.plist is as follows:
  //
  //      <dict>
  //        <key>values</key>
  //        <dict>
  //          <dict>
  //            <key>[[ UUID ]]</key>
  //            <dict>
  //              ... dict data ...
  //              <key>value</key>
  //              <dict>
  //                ... dict data ...
  //                <key>DeviceName</key><string>[[ Device name ]]</string>
  //                <key>Tabs</key>
  //                <array>
  //                  <dict>
  //                    <key>Title</key><string>[[ Title ]]</string>
  //                    <key>URL</key><string>[[ URL ]]</string>
  //                  </dict>
  //                  ... array items for other tabs ...
  //                </array>
  //              </dict>
  //            </dict>
  //            ... dicts for other devices ...
  //          </dict>
  //        </dict>
  //      </dict>
  //
  match data.get("values") {
    Some(child_key) => match child_key.as_dictionary() {
      Some(v) => Ok(v.to_owned()),
      None => error!("Top-level values key in com.apple.Safari.plist isn't an dictionary?"),
    },
    None => error!("Unable to find top-level values key in com.apple.Safari.plist"),
  }
}

/// Return a list of devices in iCloud Tabs.
pub fn list_icloud_tabs_devices() -> Result<Vec<String>, String> {
  let plist = match read_safari_plist() {
    Ok(plist) => plist,
    Err(e) => return Err(e),
  };

  Ok(plist
    .values()
    .map(|value| value
      .as_dictionary().unwrap()
      .get("value").unwrap()
      .as_dictionary().unwrap()
      .get("DeviceName").unwrap()
      .as_string().unwrap()
      .to_owned()
    )
    .collect()
  )
}

/// Return a list of URLs from iCloud Tabs.
pub fn get_icloud_tabs_urls() -> Result<HashMap<String, Vec<String>>, String> {
  let plist = match read_safari_plist() {
    Ok(plist) => plist,
    Err(e) => return Err(e),
  };

  // Within the `values` dictionary, each device is as follows:
  //
  //    <key>[[ device UUID ]]</key>
  //    <dict>
  //      [[ device data ]]
  //      <dict>
  //        <key>Tabs</key>
  //        <key>DeviceName</key><string>[[ device name ]]</key>
  //        <key>Tabs</key>
  //        <array>
  //          <dict>
  //            <key>URL</key><string>[[ URL ]]</string>
  //            [[ other tab data ]]
  //          </dict>
  //          [[ other tabs ]]
  //
  let mut result: HashMap<String, Vec<String>> = HashMap::new();
  let devices_with_tabs = plist
    .values()
    .map(|value| value
      .as_dictionary().unwrap()
      .get("value").unwrap()
      .as_dictionary().unwrap()
    );
  for data in devices_with_tabs {
    let name: String = data
      .get("DeviceName").unwrap()
      .as_string().unwrap()
      .to_owned();

    // If a device is registered with iCloud but Safari isn't running or
    // there aren't any tabs open, there isn't a Tabs field.
    if !data.contains_key("Tabs") {
      continue;
    }

    let urls: Vec<String> = data
      .get("Tabs").unwrap()
      .as_array().unwrap()
      .iter()
      .map(|tab| tab
        .as_dictionary().unwrap()
        .get("URL").unwrap()
        .as_string().unwrap()
      )
      .map(|url| urls::tidy_url(url))
      .collect();
    result.insert(name, urls);
  }
  Ok(result)
}

struct SafariWindow {
  window_index: u32,
  tab_count: u32,
}

/// Get a list of (window, tab_count) pairs for the currently running
/// version of Safari.
///
/// There are lots of bugs in Safari's AppleScript handler, so knowing
/// that there are N windows does not imply you can look up the tabs for
/// each window 1, ..., N.  There might be gaps -- looking up a window
/// in the middle could crash the AppleScript handler.
fn get_window_tab_count_pairs() -> Vec<SafariWindow> {
  let mut pairs = vec![];
  let r = run_applescript(
    "tell application \"Safari\" to get count of windows"
  );
  let window_count = r.stdout.trim().parse::<u32>().unwrap();
  for window in 1..(window_count + 1) {
    let r = run_applescript(
      &format!("tell application \"Safari\" to get count of tabs of window {}",
               window)
    );
    if r.status.success() {
      pairs.push(SafariWindow {
        window_index: window,
        tab_count: r.stdout.trim().parse::<u32>().unwrap(),
      });
    }
  }
  pairs
}
<commit_msg>Split out getting a URL and a property on a Safari window<commit_after>use std::collections::{BTreeMap, HashMap};
use std::env;
use std::fs::File;
use std::process;

use applescript::{run as run_applescript};

use plist::Plist;

use tera::{Context, Tera};

use urls;

macro_rules! error(
  ($($arg:tt)*) => { { return Err(format!($($arg)*)) } }
);

/// Returns true/false if Safari is running.
pub fn is_safari_running() -> bool {
  let cmd_result = process::Command::new("ps")
    .arg("-eaf")
    .output()
    .expect("Unable to test if Safari is running.");
  for line in String::from_utf8(cmd_result.stdout).unwrap().lines() {
    if line.contains("Safari.app/Contents/MacOS/Safari") {
      return true
    }
  }
  false
}

/// Return a URL from a Safari window.
///
/// Given a (window, tab) pair, this function looks up the URL of the tab.
/// Note that it doesn't do any error handling, so will throw an execution
/// error if it fails.
///
/// * `window` - Window index.  1 is frontmost.  If None, assumes the
///              frontmost window.
/// * `tab` - Tab index.  1 is leftmost.  If None, assumes the frontmost tab.
///
pub fn get_url(window: Option<u32>, tab: Option<u32>) -> Result<String, String> {
  get_property(window, tab, "URL")
}

/// Look up a property on a Safari window.
///
/// Given a (window, tab) pair, this function looks up the property of that tab.
/// Note that it doesn't do any error handling, so will throw an execution
/// error if it fails.
///
/// * `window` - Window index.  1 is frontmost.  If None, assumes the
///              frontmost window.
/// * `tab` - Tab index.  1 is leftmost.  If None, assumes the frontmost tab.
/// * `property` - Name of the property, as defined in the OSA scripting
///       dictionary.  One of `URL`, `source` or `text`.
///
fn get_property(window: Option<u32>, tab: Option<u32>, property: &str) -> Result<String, String> {
  if (property != "URL") && (property != "source") && (property != "text") {
    error!("Unrecognised property: {:?}", property)
  }
  // If a tab isn't specified, assume the user wants the frontmost tab.
  let command = match window {
    Some(w_idx) => {
      match tab {
        Some(t_idx) => format!("tell application \"Safari\" to get {} of tab {} of window {}", property, t_idx, w_idx),
        None => format!("tell application \"Safari\" to get {} of document {}", property, w_idx),
      }
    },
    None => format!("tell application \"Safari\" to get {} of document 1", property),
  };
  let output = run_applescript(&command);

  if output.status.success() {
    Ok(urls::tidy_url(output.stdout.trim()))
  } else {
    if output.stderr.contains("Invalid index") {
      error!("Invalid index: no such window or tab.")
    } else {
      error!("Unexpected error from osascript: {:?}", output.stderr)
    }
  }
}

/// Tests for get_property().
#[cfg(test)]
mod tests_property {
  use safari::get_property;

  #[test]
  fn test_invalid_property_is_rejected() {
    let result = get_property(None, None, "fooble");
    assert!(result.is_err());
  }
}

/// Return a list of URLs from Safari.
///
/// This returns a list of URLs, one for every tab in Safari.  Iteration
/// order depends on AppleScript, which I don't think is guaranteed to be
/// stable (in particular, I think it depends on which window is frontmost).
///
pub fn get_all_urls() -> Vec<String> {
  let windows = get_window_tab_count_pairs();
  let mut urls = vec![];
  for window in windows {
    for tab in 1..window.tab_count {
      match get_url(Some(window.window_index), Some(tab)) {
        Ok(u) => urls.push(u),
        Err(_) => {},
      }
    }
  }
  urls
}

/// Convert URL patterns into AppleScript conditions.
///
/// These conditions can be used in an `if` statement in AppleScript to
/// decide whether a tab should be closed.  Three patterns are supported,
/// a limited regex syntax:
///
///     example.com             matches anywhere in the URL
///     ^http://examples.com    matches at the start of the URL
///     example.com/$           matches at the end of the URL
///
fn parse_conditions(url_patterns: Vec<&str>) -> Vec<String> {
  url_patterns.iter().map(|p|
    if p.starts_with("^") {
      format!("starts with \"{}\"", p.replace("^", ""))
    } else if p.ends_with("$") {
      format!("ends with \"{}\"", p.replace("$", ""))
    } else {
      format!("contains \"{}\"", p)
    }
  ).collect()
}

/// Tests for parse_conditions().
#[cfg(test)]
mod tests {
  use safari::parse_conditions;

  #[test]
  fn test_parse_conditions() {
    let patterns = vec!["github.com", "^facebook.com", "twitter.com$"];
    let expected = vec!["contains \"github.com\"", "starts with \"facebook.com\"", "ends with \"twitter.com\""];
    let actual = parse_conditions(patterns);
    assert_eq!(actual, expected);
  }
}

/// Close tabs in Safari that match URL patterns.
///
/// Takes a list of URL patterns, and tries to close any matching Safari tabs.
/// Bugginess in AppleScript means this isn't always perfect, but we can
/// have a go.  In general it will fail to close tabs, rather than close
/// the wrong tabs.
///
pub fn close_tabs(url_patterns: Vec<&str>) {
  let conditions = parse_conditions(url_patterns);

  let clean_tabs_template = include_str!("scripts/clean-tabs.scpt");
  let mut context = Context::new();
  context.add("conditions", &conditions);
  let script = Tera::one_off(&clean_tabs_template, context, false).unwrap();

  // Run it twice to get around weird AppleScript bugs.
  run_applescript(&script);
  run_applescript(&script);
}

/// Get the Bookmarks.plist dict for a given title
fn read_bookmarks_plist(title: &str) -> Result<Plist, String> {

  // All Safari data lives at ~/Library/Safari/Bookmarks.plist
  // TODO: There's probably a more idiomatic Rust-like way to write this.
  let mut plist_path = match env::home_dir() {
    Some(v) => v,
    None => error!("Unable to get home directory?"),
  };
  plist_path.push("Library/Safari/Bookmarks.plist");

  let file = match File::open(plist_path) {
    Ok(v) => v,
    Err(e) => error!("Unable to open ~/Library/Safari/Bookmarks.plist: {:?}", e),
  };

  let plist = match Plist::read(file) {
    Ok(v) => v,
    Err(e) => error!("Unable to read ~/Library/Safari/Bookmarks.plist: {:?}", e),
  };

  let data = match plist.as_dictionary() {
    Some(v) => v,
    None => error!("Unable to parse ~/Library/Safari/Bookmarks.plist as dictionary?"),
  };

  // The structure of Bookmarks.plist is as follows:
  //
  //     <dict>
  //       <key>Children</key>
  //       <array>
  //         <dict>
  //           <key>Title</key><string>History</string>
  //           ... dict data ...
  //         </dict>
  //         <dict>
  //           <key>Title</key><string>com.apple.ReadingList</string>
  //           ... dict data ...
  //         </dict>
  //         ... other array items ...
  //       </array>
  //     </dict>
  //
  let children = match data.get("Children") {
    Some(child_key) => match child_key.as_array() {
      Some(v) => v,
      None => error!("Top-level children key in ~/Library/Safari/Bookmarks.plist isn't an array?"),
    },
    None => error!("Unable to find top-level Children key in ~/Library/Safari/Bookmarks.plist"),
  };

  let mut matching_children = children.iter().filter(|d|
    match d.as_dictionary() {
      Some(dict) => match dict.get("Title") {
        Some(title_elem) => match title_elem.as_string() {
          Some(v) => (v == title),
          None => false,
        },
        None => false,
      },
      None => false,
    }
  );

  // Check we got one, and only one result.
  let result = match matching_children.next() {
    Some(v) => v,
    None => error!("Unable to find key {} in Bookmarks.plist", title),
  };

  match matching_children.next() {
    Some(_) => error!("Got more than one result for {} in Bookmarks.plist", title),
    None => {},
  };

  Ok(result.to_owned())
}

/// Return a list of URLs from Reading List.
///
/// Iteration order depends on the order in which they're stored in
/// Bookmarks.plist, which is usually (but not guaranteed to be) newest first.
///
pub fn get_reading_list_urls() -> Result<Vec<String>, String> {
  let plist = match read_bookmarks_plist("com.apple.ReadingList") {
    Ok(v) => v,
    Err(e) => return Err(e),
  };

  // TODO: All these unwrap() calls should probably be handled better
  let children = plist
    .as_dictionary().unwrap()
    .get("Children").unwrap()
    .as_array().unwrap();

  Ok(children
    .iter()
    .map(|child|
      child
        .as_dictionary().unwrap()
        .get("URLString").unwrap()
        .as_string().unwrap()
    )
    .map(|url| urls::tidy_url(url))
    .collect()
  )
}

/// Get the com.apple.Safari.plist preferences file
fn read_safari_plist() -> Result<BTreeMap<String, Plist>, String> {

  // All Safari data lives at ~/Library/Safari/Bookmarks.plist
  // TODO: There's probably a more idiomatic Rust-like way to write this.
  let mut plist_path = match env::home_dir() {
    Some(v) => v,
    None => error!("Unable to get home directory?"),
  };
  plist_path.push("Library/SyncedPreferences/com.apple.Safari.plist");

  let file = match File::open(plist_path) {
    Ok(v) => v,
    Err(e) => error!("Unable to open com.apple.Safari.plist: {:?}", e),
  };

  let plist = match Plist::read(file) {
    Ok(v) => v,
    Err(e) => error!("Unable to read com.apple.Safari.plist: {:?}", e),
  };

  let data = match plist.as_dictionary() {
    Some(v) => v,
    None => error!("Unable to parse com.apple.Safari.plist as dictionary?"),
  };

  // The structure of com.apple.Safari.plist is as follows:
  //
  //      <dict>
  //        <key>values</key>
  //        <dict>
  //          <dict>
  //            <key>[[ UUID ]]</key>
  //            <dict>
  //              ... dict data ...
  //              <key>value</key>
  //              <dict>
  //                ... dict data ...
  //                <key>DeviceName</key><string>[[ Device name ]]</string>
  //                <key>Tabs</key>
  //                <array>
  //                  <dict>
  //                    <key>Title</key><string>[[ Title ]]</string>
  //                    <key>URL</key><string>[[ URL ]]</string>
  //                  </dict>
  //                  ... array items for other tabs ...
  //                </array>
  //              </dict>
  //            </dict>
  //            ... dicts for other devices ...
  //          </dict>
  //        </dict>
  //      </dict>
  //
  match data.get("values") {
    Some(child_key) => match child_key.as_dictionary() {
      Some(v) => Ok(v.to_owned()),
      None => error!("Top-level values key in com.apple.Safari.plist isn't an dictionary?"),
    },
    None => error!("Unable to find top-level values key in com.apple.Safari.plist"),
  }
}

/// Return a list of devices in iCloud Tabs.
pub fn list_icloud_tabs_devices() -> Result<Vec<String>, String> {
  let plist = match read_safari_plist() {
    Ok(plist) => plist,
    Err(e) => return Err(e),
  };

  Ok(plist
    .values()
    .map(|value| value
      .as_dictionary().unwrap()
      .get("value").unwrap()
      .as_dictionary().unwrap()
      .get("DeviceName").unwrap()
      .as_string().unwrap()
      .to_owned()
    )
    .collect()
  )
}

/// Return a list of URLs from iCloud Tabs.
pub fn get_icloud_tabs_urls() -> Result<HashMap<String, Vec<String>>, String> {
  let plist = match read_safari_plist() {
    Ok(plist) => plist,
    Err(e) => return Err(e),
  };

  // Within the `values` dictionary, each device is as follows:
  //
  //    <key>[[ device UUID ]]</key>
  //    <dict>
  //      [[ device data ]]
  //      <dict>
  //        <key>Tabs</key>
  //        <key>DeviceName</key><string>[[ device name ]]</key>
  //        <key>Tabs</key>
  //        <array>
  //          <dict>
  //            <key>URL</key><string>[[ URL ]]</string>
  //            [[ other tab data ]]
  //          </dict>
  //          [[ other tabs ]]
  //
  let mut result: HashMap<String, Vec<String>> = HashMap::new();
  let devices_with_tabs = plist
    .values()
    .map(|value| value
      .as_dictionary().unwrap()
      .get("value").unwrap()
      .as_dictionary().unwrap()
    );
  for data in devices_with_tabs {
    let name: String = data
      .get("DeviceName").unwrap()
      .as_string().unwrap()
      .to_owned();

    // If a device is registered with iCloud but Safari isn't running or
    // there aren't any tabs open, there isn't a Tabs field.
    if !data.contains_key("Tabs") {
      continue;
    }

    let urls: Vec<String> = data
      .get("Tabs").unwrap()
      .as_array().unwrap()
      .iter()
      .map(|tab| tab
        .as_dictionary().unwrap()
        .get("URL").unwrap()
        .as_string().unwrap()
      )
      .map(|url| urls::tidy_url(url))
      .collect();
    result.insert(name, urls);
  }
  Ok(result)
}

struct SafariWindow {
  window_index: u32,
  tab_count: u32,
}

/// Get a list of (window, tab_count) pairs for the currently running
/// version of Safari.
///
/// There are lots of bugs in Safari's AppleScript handler, so knowing
/// that there are N windows does not imply you can look up the tabs for
/// each window 1, ..., N.  There might be gaps -- looking up a window
/// in the middle could crash the AppleScript handler.
fn get_window_tab_count_pairs() -> Vec<SafariWindow> {
  let mut pairs = vec![];
  let r = run_applescript(
    "tell application \"Safari\" to get count of windows"
  );
  let window_count = r.stdout.trim().parse::<u32>().unwrap();
  for window in 1..(window_count + 1) {
    let r = run_applescript(
      &format!("tell application \"Safari\" to get count of tabs of window {}",
               window)
    );
    if r.status.success() {
      pairs.push(SafariWindow {
        window_index: window,
        tab_count: r.stdout.trim().parse::<u32>().unwrap(),
      });
    }
  }
  pairs
}
<|endoftext|>


<commit_before>use std::collections::{BTreeMap, HashMap, HashSet};
use std::collections::hash_map::Entry::*;
use std::hash::{Hash, Hasher};
use std::sync::Arc;
use std::iter;
use std::str::CharIndices;
use std::ops;
use itertools::{GroupBy, Itertools};
use coalesce::Coalesce;

/// A Document contains an id and content.
/// Hashing and equality are based only on the id field.
#[derive(Clone, Eq, Debug, RustcEncodable, RustcDecodable)]
pub struct Document {
    id: String,
    content: String,
}

impl Document {
    /// Construct a new Document from an id and content.
    /// Both two arguments can be anything that can be turned into a String.
    pub fn new<S, T>(id: S, content: T) -> Document
        where S: Into<String>,
              T: Into<String>
    {
        Document { id: id.into(), content: content.into(), }
    }

    pub fn id(&self) -> &str {
        &self.id
    }

    pub fn content(&self) -> &str {
        &self.content
    }
}

impl Hash for Document {
    // Documents are unique only upon their id
    fn hash<H>(&self, state: &mut H) where H: Hasher {
        self.id.hash(state);
    }
}

impl PartialEq<Document> for Document {
    // Documents are unique only upon their id
    fn eq(&self, other: &Document) -> bool {
        self.id == other.id
    }
}

/// A SearchResult is the representation of a Document returned for a specific set of search
/// terms. It is unique upon the document and the vec of highlight indices. It also contains a
/// score and a String of the document's content post-highlighting.
#[derive(Clone, Debug, RustcEncodable)]
pub struct SearchResult {
    doc: Arc<Document>,
    highlights: Vec<(usize, usize)>,
    score: f32,
}

impl PartialEq for SearchResult {
    // SearchResult is unique upon its document and highlights
    fn eq(&self, other: &SearchResult) -> bool {
        self.doc == other.doc && self.highlights == other.highlights
    }
}

impl Eq for SearchResult {}

impl Hash for SearchResult {
    // SearchResult is unique upon its document and highlights
    fn hash<H>(&self, state: &mut H) where H: Hasher {
        self.doc.hash(state);
        self.highlights.hash(state);
    }
}

impl SearchResult {
    fn new(doc: Arc<Document>, mut highlights: Vec<(usize, usize)>) -> SearchResult {
        SearchResult {
            score: highlights.iter()
                .map(|&(begin, end)| end - begin)
                .sum::<usize>() as f32 / (doc.content.len() as f32).sqrt(),
            doc: doc,
            highlights: { highlights.sort(); highlights },
        }
    }

    /// Returns the document
    pub fn doc(&self) -> &Arc<Document> {
        &self.doc
    }

    /// Returns the highlighted indices.
    ///
    /// Each `(usize, usize)` indicates the start and end of a term in the document's content
    /// that should be highlighted.
    pub fn highlights(&self) -> &Vec<(usize, usize)> {
        &self.highlights
    }

    /// Returns the search result's score.
    ///
    /// Score is computed by the product of the summed length of the matching terms and the inverse
    /// square root of the length of the document. Taking the square root of the document's length
    /// helps to combat bias toward short content.
    pub fn score(&self) -> f32 {
        self.score
    }

    pub fn highlight(&self, before: &str, after: &str) -> String {
        let mut begin_idx = 0;
        let mut parts = String::new();
        for &(begin, end) in &self.highlights {
            parts.push_str(&self.doc.content[begin_idx..begin]);
            parts.push_str(before);
            parts.push_str(&self.doc.content[begin..end]);
            parts.push_str(after);
            begin_idx = end;
        }
        parts.push_str(&self.doc.content[begin_idx..]);
        parts
    }
}

/// A basic implementation of an `Index`, the inverted index is a data structure that maps
/// from words to sets of Documents.
#[derive(Debug)]
pub struct InvertedIndex {
    index: BTreeMap<String, BTreeMap<String, Vec<(usize, usize)>>>,
    docs: BTreeMap<String, Arc<Document>>,
}

impl InvertedIndex {
    pub fn new() -> InvertedIndex {
        InvertedIndex {
            index: BTreeMap::new(),
            docs: BTreeMap::new(),
        }
    }

    /// A basic implementation of index, splits the document's content into whitespace-separated
    /// words, and inserts each word-document pair into the map.
    pub fn index(&mut self, doc: Document) {
        let doc = Arc::new(doc);
        let analyzed  = analyze_doc(doc.content());
        let previous_version = self.docs.insert(doc.id.clone(), doc.clone());
        if let Some(previous_version) = previous_version {
            let previous_analyzed = analyze_doc(previous_version.content());
            for (ngram, _) in previous_analyzed {
                let is_empty = {
                    let docs_for_ngram = self.index.get_mut(&ngram).unwrap();
                    docs_for_ngram.remove(&doc.id);
                    docs_for_ngram.is_empty()
                };
                if is_empty {
                    self.index.remove(&ngram);
                }
            }
        }

        for (ngram, highlighted) in analyzed {
            let mut highlights = self.index.entry(ngram)
                .or_insert_with(BTreeMap::new)
                .entry(doc.id.clone())
                .or_insert_with(Vec::new);
            let coalesce_idx = highlights.binary_search(&highlighted).err().unwrap();
            highlights.coalesce(coalesce_idx, highlighted);
        }
    }

    /// A basic search implementation that splits the query's content into whitespace-separated
    /// words, looks up the set of Documents for each word, and then concatenates the sets.
    pub fn search(&self, query: &str) -> Vec<SearchResult> {
        let unique_terms: HashSet<_> = query.split_whitespace().map(str::to_lowercase).collect();
        let map: HashMap<String, Vec<(usize, usize)>> = unique_terms.into_iter()
            .flat_map(|word| self.index.get(&word))
            .flat_map(BTreeMap::iter)
            .fold(HashMap::new(), |mut map: HashMap<String, Vec<(usize, usize)>>, (doc_id, highlights): (&String, &Vec<(usize, usize)>)| {
                match map.entry(doc_id.clone()) {
                    Vacant(entry) => { entry.insert(highlights.clone()); },
                    Occupied(mut entry) => {
                        let entry: &mut Vec<(usize, usize)> = entry.get_mut();
                        for highlight in highlights {
                            let coalesce_idx = entry.binary_search(highlight).err().unwrap();
                            entry.coalesce(coalesce_idx, highlight.clone());
                        }
                    }
                }
                map
            });
        let mut results: Vec<_> = map.into_iter()
            .map(|(doc_id, index_map)| SearchResult::new(self.docs[&doc_id].clone(), index_map.into_iter().collect()))
            .collect();
        results.sort_by(|result1, result2| result2.score.partial_cmp(&result1.score).unwrap());
        results
    }
}

fn analyze_doc(doc: &str)
-> iter::FlatMap<
    iter::Filter<
        GroupBy<bool, CharIndices, fn(&(usize, char)) -> bool>,
        fn(&(bool, Vec<(usize, char)>)) -> bool>,
    iter::Map<ops::Range<usize>, Ngrams>,
    fn((bool, Vec<(usize, char)>)) -> iter::Map<ops::Range<usize>, Ngrams>>
{
    doc.char_indices()
        .group_by(is_whitespace as fn(&(usize, char)) -> bool)
        .filter(not_whitespace as fn(&(bool, Vec<(usize, char)>)) -> bool)
        .flat_map(ngrams)
}

fn ngrams((_, chars): (bool, Vec<(usize, char)>)) -> iter::Map<ops::Range<usize>, Ngrams> {
    (1..chars.len() + 1).map(Ngrams::new(chars))
}

fn not_whitespace(&(is_whitespace, _): &(bool, Vec<(usize, char)>)) -> bool {
    !is_whitespace
}

fn is_whitespace(&(_, c): &(usize, char)) -> bool {
    c.is_whitespace()
}

struct Ngrams {
    chars: Vec<(usize, char)>
}

impl Ngrams {
    fn new(chars: Vec<(usize, char)>) -> Ngrams {
        Ngrams { chars: chars }
    }
}

impl Fn<(usize,)> for Ngrams {
    extern "rust-call" fn call(&self, (to,): (usize,)) -> (String, (usize, usize)) {
        let word = self.chars[..to].iter().flat_map(|&(_, c)| c.to_lowercase()).collect();
        let start = self.chars[0].0;
        let (last_idx, last_char) = self.chars[to - 1];
        let finish = last_idx + last_char.len_utf8();
        (word, (start, finish))
    }
}

impl FnMut<(usize,)> for Ngrams {
    extern "rust-call" fn call_mut(&mut self, to: (usize,)) -> (String, (usize, usize)) {
        self.call(to)
    }
}

impl FnOnce<(usize,)> for Ngrams {
    type Output = (String, (usize, usize));
    extern "rust-call" fn call_once(self, to: (usize,)) -> (String, (usize, usize)) {
        self.call(to)
    }
}

#[cfg(test)]
mod bench {
    use super::*;
    use std::sync::Arc;
    use std::collections::HashSet;

    #[test]
    fn test_search1() {
        let mut index = InvertedIndex::new();
        let doc1 = Document::new("1", "learn to program in rust today");
        index.index(doc1.clone());
        let doc2 = Document::new("2", "what did you today do");
        index.index(doc2.clone());
        let search_results = index.search("to");
        let expected = [
            SearchResult::new(Arc::new(doc1), vec![(6, 8), (25, 27)]),
            SearchResult::new(Arc::new(doc2), vec![(13, 15)])
        ];
        assert_eq!(search_results, expected.iter().cloned().collect::<Vec<_>>());
        assert_eq!("learn to program in rust \
                   today",
                   expected[0].highlight("", ""));
    }

    #[test]
    fn test_ngrams() {
        let mut index = InvertedIndex::new();
        let doc1 = Document::new("1", "learn to program in rust today");
        let doc2 = Document::new("2", "what did you today do");
        index.index(doc1.clone());
        index.index(doc2.clone());
        let search_results = index.search("to");
        let expected = [
            SearchResult::new(Arc::new(doc1), vec![(6, 8), (25, 27)]),
            SearchResult::new(Arc::new(doc2), vec![(13, 15)]),
        ];
        assert_eq!(search_results, expected.iter().cloned().collect::<Vec<_>>());
        assert_eq!("learn to program in rust \
                   today",
                   expected[0].highlight("", ""));

    }

    #[test]
    fn test_highlight1() {
        let mut index = InvertedIndex::new();
        let doc1 = Document::new("2", "what to do today");
        let doc2 = Document::new("3", "hey today");
        index.index(doc1.clone());
        index.index(doc2.clone());
        let search_results = index.search("to");
        let expected = [
            SearchResult::new(Arc::new(doc1), vec![(5, 7), (11, 13)]),
            SearchResult::new(Arc::new(doc2), vec![(4, 6)]),
        ];
        assert_eq!(search_results.into_iter().collect::<HashSet<_>>(),
                   expected.iter().cloned().collect());
        assert_eq!("what to do today",
                   expected[0].highlight("", ""));
    }

    #[test]
    fn test_highlight2() {
        let mut index = InvertedIndex::new();
        let doc1 = Document::new("2", "Won\u{2019}t this split the ecosystem? Will everyone use?");
        index.index(doc1.clone());
        let expected =  "Won\u{2019}t this split the *e*cosystem? Will *e*veryone use?";
        let search_results = index.search("e");
        assert_eq!(1, search_results.len());
        assert_eq!(search_results[0].highlight("*", "*"), expected);
    }

    #[test]
    fn test_unicode() {
        let mut index = InvertedIndex::new();
        let doc = Document::new("0", "嗨, 您好");
        index.index(doc.clone());
        let to_search = "您";
        let search_results = index.search(to_search);
        let &SearchResult { ref doc, ref highlights, .. } = search_results.iter().next().unwrap();
        let (begin, end) = highlights[0];
        assert_eq!(&doc.content()[begin..end], to_search);
    }

    #[test]
    fn test_update_doc() {
        let mut index = InvertedIndex::new();
        let doc = Document::new("0", "abc åäö");
        index.index(doc);
        let doc = Document::new("0", "different");
        index.index(doc);
        let search_results = index.search("å");
        assert!(search_results.is_empty());
        assert_eq!(index.docs.len(), 1);
    }

    #[test]
    fn test_ranking() {
        let mut index = InvertedIndex::new();
        let doc = Document::new("0", "beat");
        index.index(doc.clone());
        let doc2 = Document::new("1", "beast");
        index.index(doc2.clone());
        let search_results = index.search("be");
        assert_eq!(index.docs.len(), 2);
        assert_eq!(search_results.into_iter().map(|r| r.doc).collect::<HashSet<_>>(),
                [Arc::new(doc), Arc::new(doc2)].iter().cloned().collect());
    }

    #[test]
    fn test_duplicate_term() {
        let mut index = InvertedIndex::new();
        let doc = Document::new("0", "beat");
        index.index(doc.clone());
        let search_results = index.search("be be");
        assert_eq!(search_results.len(), 1);
    }

    #[test]
    fn test_duplicate_term2() {
        let mut index = InvertedIndex::new();
        let doc = Document::new("0", "beat");
        index.index(doc.clone());
        let search_results = index.search("be b");
        assert_eq!(search_results.len(), 1);
        assert_eq!(search_results[0], SearchResult::new(Arc::new(doc), vec![(0, 2)]));
    }

    #[test]
    fn test_lowercase_search() {
        let mut index = InvertedIndex::new();
        let doc = Document::new("0", "BeAt");
        index.index(doc.clone());
        let search_results = index.search("bE");
        assert_eq!(search_results.len(), 1);
        assert_eq!(search_results[0], SearchResult::new(Arc::new(doc), vec![(0, 2)]));
    }

    #[test]
    fn test_lowercase_index() {
        let mut index = InvertedIndex::new();
        let doc = Document::new("0", "BeAt");
        index.index(doc.clone());
        let search_results = index.search("be");
        assert_eq!(search_results.len(), 1);
        assert_eq!(search_results[0], SearchResult::new(Arc::new(doc), vec![(0, 2)]));
    }
}
<commit_msg>Remove unnecessary type annotation<commit_after>use std::collections::{BTreeMap, HashMap, HashSet};
use std::collections::hash_map::Entry::*;
use std::hash::{Hash, Hasher};
use std::sync::Arc;
use std::iter;
use std::str::CharIndices;
use std::ops;
use itertools::{GroupBy, Itertools};
use coalesce::Coalesce;

/// A Document contains an id and content.
/// Hashing and equality are based only on the id field.
#[derive(Clone, Eq, Debug, RustcEncodable, RustcDecodable)]
pub struct Document {
    id: String,
    content: String,
}

impl Document {
    /// Construct a new Document from an id and content.
    /// Both two arguments can be anything that can be turned into a String.
    pub fn new<S, T>(id: S, content: T) -> Document
        where S: Into<String>,
              T: Into<String>
    {
        Document { id: id.into(), content: content.into(), }
    }

    pub fn id(&self) -> &str {
        &self.id
    }

    pub fn content(&self) -> &str {
        &self.content
    }
}

impl Hash for Document {
    // Documents are unique only upon their id
    fn hash<H>(&self, state: &mut H) where H: Hasher {
        self.id.hash(state);
    }
}

impl PartialEq<Document> for Document {
    // Documents are unique only upon their id
    fn eq(&self, other: &Document) -> bool {
        self.id == other.id
    }
}

/// A SearchResult is the representation of a Document returned for a specific set of search
/// terms. It is unique upon the document and the vec of highlight indices. It also contains a
/// score and a String of the document's content post-highlighting.
#[derive(Clone, Debug, RustcEncodable)]
pub struct SearchResult {
    doc: Arc<Document>,
    highlights: Vec<(usize, usize)>,
    score: f32,
}

impl PartialEq for SearchResult {
    // SearchResult is unique upon its document and highlights
    fn eq(&self, other: &SearchResult) -> bool {
        self.doc == other.doc && self.highlights == other.highlights
    }
}

impl Eq for SearchResult {}

impl Hash for SearchResult {
    // SearchResult is unique upon its document and highlights
    fn hash<H>(&self, state: &mut H) where H: Hasher {
        self.doc.hash(state);
        self.highlights.hash(state);
    }
}

impl SearchResult {
    fn new(doc: Arc<Document>, mut highlights: Vec<(usize, usize)>) -> SearchResult {
        SearchResult {
            score: highlights.iter()
                .map(|&(begin, end)| end - begin)
                .sum::<usize>() as f32 / (doc.content.len() as f32).sqrt(),
            doc: doc,
            highlights: { highlights.sort(); highlights },
        }
    }

    /// Returns the document
    pub fn doc(&self) -> &Arc<Document> {
        &self.doc
    }

    /// Returns the highlighted indices.
    ///
    /// Each `(usize, usize)` indicates the start and end of a term in the document's content
    /// that should be highlighted.
    pub fn highlights(&self) -> &Vec<(usize, usize)> {
        &self.highlights
    }

    /// Returns the search result's score.
    ///
    /// Score is computed by the product of the summed length of the matching terms and the inverse
    /// square root of the length of the document. Taking the square root of the document's length
    /// helps to combat bias toward short content.
    pub fn score(&self) -> f32 {
        self.score
    }

    pub fn highlight(&self, before: &str, after: &str) -> String {
        let mut begin_idx = 0;
        let mut parts = String::new();
        for &(begin, end) in &self.highlights {
            parts.push_str(&self.doc.content[begin_idx..begin]);
            parts.push_str(before);
            parts.push_str(&self.doc.content[begin..end]);
            parts.push_str(after);
            begin_idx = end;
        }
        parts.push_str(&self.doc.content[begin_idx..]);
        parts
    }
}

/// A basic implementation of an `Index`, the inverted index is a data structure that maps
/// from words to sets of Documents.
#[derive(Debug)]
pub struct InvertedIndex {
    index: BTreeMap<String, BTreeMap<String, Vec<(usize, usize)>>>,
    docs: BTreeMap<String, Arc<Document>>,
}

impl InvertedIndex {
    pub fn new() -> InvertedIndex {
        InvertedIndex {
            index: BTreeMap::new(),
            docs: BTreeMap::new(),
        }
    }

    /// A basic implementation of index, splits the document's content into whitespace-separated
    /// words, and inserts each word-document pair into the map.
    pub fn index(&mut self, doc: Document) {
        let doc = Arc::new(doc);
        let analyzed  = analyze_doc(doc.content());
        let previous_version = self.docs.insert(doc.id.clone(), doc.clone());
        if let Some(previous_version) = previous_version {
            let previous_analyzed = analyze_doc(previous_version.content());
            for (ngram, _) in previous_analyzed {
                let is_empty = {
                    let docs_for_ngram = self.index.get_mut(&ngram).unwrap();
                    docs_for_ngram.remove(&doc.id);
                    docs_for_ngram.is_empty()
                };
                if is_empty {
                    self.index.remove(&ngram);
                }
            }
        }

        for (ngram, highlighted) in analyzed {
            let mut highlights = self.index.entry(ngram)
                .or_insert_with(BTreeMap::new)
                .entry(doc.id.clone())
                .or_insert_with(Vec::new);
            let coalesce_idx = highlights.binary_search(&highlighted).err().unwrap();
            highlights.coalesce(coalesce_idx, highlighted);
        }
    }

    /// A basic search implementation that splits the query's content into whitespace-separated
    /// words, looks up the set of Documents for each word, and then concatenates the sets.
    pub fn search(&self, query: &str) -> Vec<SearchResult> {
        let unique_terms: HashSet<_> = query.split_whitespace().map(str::to_lowercase).collect();
        let map: HashMap<String, Vec<(usize, usize)>> = unique_terms.into_iter()
            .flat_map(|word| self.index.get(&word))
            .flat_map(BTreeMap::iter)
            .fold(HashMap::new(), |mut map, (doc_id, highlights)| {
                match map.entry(doc_id.clone()) {
                    Vacant(entry) => { entry.insert(highlights.clone()); },
                    Occupied(mut entry) => {
                        let entry: &mut Vec<(usize, usize)> = entry.get_mut();
                        for highlight in highlights {
                            let coalesce_idx = entry.binary_search(highlight).err().unwrap();
                            entry.coalesce(coalesce_idx, highlight.clone());
                        }
                    }
                }
                map
            });
        let mut results: Vec<_> = map.into_iter()
            .map(|(doc_id, index_map)| SearchResult::new(self.docs[&doc_id].clone(), index_map.into_iter().collect()))
            .collect();
        results.sort_by(|result1, result2| result2.score.partial_cmp(&result1.score).unwrap());
        results
    }
}

fn analyze_doc(doc: &str)
-> iter::FlatMap<
    iter::Filter<
        GroupBy<bool, CharIndices, fn(&(usize, char)) -> bool>,
        fn(&(bool, Vec<(usize, char)>)) -> bool>,
    iter::Map<ops::Range<usize>, Ngrams>,
    fn((bool, Vec<(usize, char)>)) -> iter::Map<ops::Range<usize>, Ngrams>>
{
    doc.char_indices()
        .group_by(is_whitespace as fn(&(usize, char)) -> bool)
        .filter(not_whitespace as fn(&(bool, Vec<(usize, char)>)) -> bool)
        .flat_map(ngrams)
}

fn ngrams((_, chars): (bool, Vec<(usize, char)>)) -> iter::Map<ops::Range<usize>, Ngrams> {
    (1..chars.len() + 1).map(Ngrams::new(chars))
}

fn not_whitespace(&(is_whitespace, _): &(bool, Vec<(usize, char)>)) -> bool {
    !is_whitespace
}

fn is_whitespace(&(_, c): &(usize, char)) -> bool {
    c.is_whitespace()
}

struct Ngrams {
    chars: Vec<(usize, char)>
}

impl Ngrams {
    fn new(chars: Vec<(usize, char)>) -> Ngrams {
        Ngrams { chars: chars }
    }
}

impl Fn<(usize,)> for Ngrams {
    extern "rust-call" fn call(&self, (to,): (usize,)) -> (String, (usize, usize)) {
        let word = self.chars[..to].iter().flat_map(|&(_, c)| c.to_lowercase()).collect();
        let start = self.chars[0].0;
        let (last_idx, last_char) = self.chars[to - 1];
        let finish = last_idx + last_char.len_utf8();
        (word, (start, finish))
    }
}

impl FnMut<(usize,)> for Ngrams {
    extern "rust-call" fn call_mut(&mut self, to: (usize,)) -> (String, (usize, usize)) {
        self.call(to)
    }
}

impl FnOnce<(usize,)> for Ngrams {
    type Output = (String, (usize, usize));
    extern "rust-call" fn call_once(self, to: (usize,)) -> (String, (usize, usize)) {
        self.call(to)
    }
}

#[cfg(test)]
mod bench {
    use super::*;
    use std::sync::Arc;
    use std::collections::HashSet;

    #[test]
    fn test_search1() {
        let mut index = InvertedIndex::new();
        let doc1 = Document::new("1", "learn to program in rust today");
        index.index(doc1.clone());
        let doc2 = Document::new("2", "what did you today do");
        index.index(doc2.clone());
        let search_results = index.search("to");
        let expected = [
            SearchResult::new(Arc::new(doc1), vec![(6, 8), (25, 27)]),
            SearchResult::new(Arc::new(doc2), vec![(13, 15)])
        ];
        assert_eq!(search_results, expected.iter().cloned().collect::<Vec<_>>());
        assert_eq!("learn to program in rust \
                   today",
                   expected[0].highlight("", ""));
    }

    #[test]
    fn test_ngrams() {
        let mut index = InvertedIndex::new();
        let doc1 = Document::new("1", "learn to program in rust today");
        let doc2 = Document::new("2", "what did you today do");
        index.index(doc1.clone());
        index.index(doc2.clone());
        let search_results = index.search("to");
        let expected = [
            SearchResult::new(Arc::new(doc1), vec![(6, 8), (25, 27)]),
            SearchResult::new(Arc::new(doc2), vec![(13, 15)]),
        ];
        assert_eq!(search_results, expected.iter().cloned().collect::<Vec<_>>());
        assert_eq!("learn to program in rust \
                   today",
                   expected[0].highlight("", ""));

    }

    #[test]
    fn test_highlight1() {
        let mut index = InvertedIndex::new();
        let doc1 = Document::new("2", "what to do today");
        let doc2 = Document::new("3", "hey today");
        index.index(doc1.clone());
        index.index(doc2.clone());
        let search_results = index.search("to");
        let expected = [
            SearchResult::new(Arc::new(doc1), vec![(5, 7), (11, 13)]),
            SearchResult::new(Arc::new(doc2), vec![(4, 6)]),
        ];
        assert_eq!(search_results.into_iter().collect::<HashSet<_>>(),
                   expected.iter().cloned().collect());
        assert_eq!("what to do today",
                   expected[0].highlight("", ""));
    }

    #[test]
    fn test_highlight2() {
        let mut index = InvertedIndex::new();
        let doc1 = Document::new("2", "Won\u{2019}t this split the ecosystem? Will everyone use?");
        index.index(doc1.clone());
        let expected =  "Won\u{2019}t this split the *e*cosystem? Will *e*veryone use?";
        let search_results = index.search("e");
        assert_eq!(1, search_results.len());
        assert_eq!(search_results[0].highlight("*", "*"), expected);
    }

    #[test]
    fn test_unicode() {
        let mut index = InvertedIndex::new();
        let doc = Document::new("0", "嗨, 您好");
        index.index(doc.clone());
        let to_search = "您";
        let search_results = index.search(to_search);
        let &SearchResult { ref doc, ref highlights, .. } = search_results.iter().next().unwrap();
        let (begin, end) = highlights[0];
        assert_eq!(&doc.content()[begin..end], to_search);
    }

    #[test]
    fn test_update_doc() {
        let mut index = InvertedIndex::new();
        let doc = Document::new("0", "abc åäö");
        index.index(doc);
        let doc = Document::new("0", "different");
        index.index(doc);
        let search_results = index.search("å");
        assert!(search_results.is_empty());
        assert_eq!(index.docs.len(), 1);
    }

    #[test]
    fn test_ranking() {
        let mut index = InvertedIndex::new();
        let doc = Document::new("0", "beat");
        index.index(doc.clone());
        let doc2 = Document::new("1", "beast");
        index.index(doc2.clone());
        let search_results = index.search("be");
        assert_eq!(index.docs.len(), 2);
        assert_eq!(search_results.into_iter().map(|r| r.doc).collect::<HashSet<_>>(),
                [Arc::new(doc), Arc::new(doc2)].iter().cloned().collect());
    }

    #[test]
    fn test_duplicate_term() {
        let mut index = InvertedIndex::new();
        let doc = Document::new("0", "beat");
        index.index(doc.clone());
        let search_results = index.search("be be");
        assert_eq!(search_results.len(), 1);
    }

    #[test]
    fn test_duplicate_term2() {
        let mut index = InvertedIndex::new();
        let doc = Document::new("0", "beat");
        index.index(doc.clone());
        let search_results = index.search("be b");
        assert_eq!(search_results.len(), 1);
        assert_eq!(search_results[0], SearchResult::new(Arc::new(doc), vec![(0, 2)]));
    }

    #[test]
    fn test_lowercase_search() {
        let mut index = InvertedIndex::new();
        let doc = Document::new("0", "BeAt");
        index.index(doc.clone());
        let search_results = index.search("bE");
        assert_eq!(search_results.len(), 1);
        assert_eq!(search_results[0], SearchResult::new(Arc::new(doc), vec![(0, 2)]));
    }

    #[test]
    fn test_lowercase_index() {
        let mut index = InvertedIndex::new();
        let doc = Document::new("0", "BeAt");
        index.index(doc.clone());
        let search_results = index.search("be");
        assert_eq!(search_results.len(), 1);
        assert_eq!(search_results[0], SearchResult::new(Arc::new(doc), vec![(0, 2)]));
    }
}
<|endoftext|>


<commit_before>use std::ops::{Not, BitAnd, BitAndAssign, BitOr, BitOrAssign, BitXor, BitXorAssign, Shl, ShlAssign, Shr, ShrAssign};
use wrapping::{WrappingAdd, WrappingSub};

/// A type designating data that will only be used in a constant time manner
pub struct Secret<T: Copy> {
    inner: T
}

impl<T: Copy> Copy for Secret<T> { }
impl<T: Copy> Clone for Secret<T> {
    fn clone(&self) -> Secret<T> {
        *self
    }
}

impl<T: Copy> Secret<T> {
    pub fn new(val: T) -> Secret<T> {
        Secret { inner: val }
    }

    /// Expose the data inside this wrapper. TODO: warning about use and greppability
    pub fn expose(self) -> T {
        self.inner
    }
}

macro_rules! pod_impl {
    { $tr:ident, $method:ident, $t:ty } => {
        impl $tr<$t> for Secret<$t> {
            type Output = Secret<<$t as $tr<$t>>::Output>;
            fn $method(self, rhs: $t) -> Self::Output {
                Secret::new($tr::$method(self.expose(), rhs))
            }
        }

        impl $tr<Secret<$t>> for $t {
            type Output = Secret<<$t as $tr<$t>>::Output>;
            fn $method(self, rhs: Secret<$t>) -> Self::Output {
                Secret::new($tr::$method(self, rhs.expose()))
            }
        }

        impl $tr<Secret<$t>> for Secret<$t> {
            type Output = Secret<<$t as $tr<$t>>::Output>;
            fn $method(self, rhs: Secret<$t>) -> Self::Output {
                Secret::new($tr::$method(self.expose(), rhs.expose()))
            }
        }
    };

    { $tr:ident, $method:ident, $trassign:ident, $methodassign: ident, $t:ty } => {
        pod_impl! { $tr, $method, $t }

        impl $trassign<Secret<$t>> for Secret<$t> {
            fn $methodassign(&mut self, rhs: Secret<$t>) {
                $trassign::$methodassign(&mut self.inner, rhs.expose());
            }
        }

        impl $trassign<$t> for Secret<$t> {
            fn $methodassign(&mut self, rhs: $t) {
                $trassign::$methodassign(&mut self.inner, rhs);
            }
        }
    };

    { shift, $tr:ident, $method:ident, $trassign:ident, $methodassign: ident, $t:ty } => {
        impl $tr<usize> for Secret<$t> {
            type Output = Secret<<$t as $tr<$t>>::Output>;
            fn $method(self, rhs: usize) -> Self::Output {
                Secret::new($tr::$method(self.expose(), rhs))
            }
        }

        impl $trassign<usize> for Secret<$t> {
            fn $methodassign(&mut self, rhs: usize) {
                $trassign::$methodassign(&mut self.inner, rhs);
            }
        }
    }
}

macro_rules! pod_impls {
    { $t:ty } => {
        impl Not for Secret<$t> {
            type Output = Secret<<$t as Not>::Output>;
            fn not(self) -> Self::Output {
                Secret::new(!self.expose())
            }
        }

        pod_impl! { BitAnd, bitand, BitAndAssign, bitand_assign, $t }
        pod_impl! { BitOr, bitor, BitOrAssign, bitor_assign, $t }
        pod_impl! { BitXor, bitxor, BitXorAssign, bitxor_assign, $t }
        pod_impl! { shift, Shl, shl, ShlAssign, shl_assign, $t }
        pod_impl! { shift, Shr, shr, ShrAssign, shr_assign, $t }
        pod_impl! { WrappingAdd, wrapping_add, $t }
        pod_impl! { WrappingSub, wrapping_sub, $t }
    };
}

pod_impls! { u8 }
pod_impls! { u16 }
pod_impls! { u32 }
pod_impls! { u64 }
pod_impls! { usize }
pod_impls! { i8 }
pod_impls! { i16 }
pod_impls! { i32 }
pod_impls! { i64 }
pod_impls! { isize }
<commit_msg>From impls for Secret<commit_after>use std::convert::From;
use std::ops::{Not, BitAnd, BitAndAssign, BitOr, BitOrAssign, BitXor, BitXorAssign, Shl, ShlAssign, Shr, ShrAssign};
use wrapping::{WrappingAdd, WrappingSub};

/// A type designating data that will only be used in a constant time manner
pub struct Secret<T: Copy> {
    inner: T
}

impl<T: Copy> Copy for Secret<T> { }
impl<T: Copy> Clone for Secret<T> {
    fn clone(&self) -> Secret<T> {
        *self
    }
}

impl<T: Copy> Secret<T> {
    pub fn new(val: T) -> Secret<T> {
        Secret { inner: val }
    }

    /// Expose the data inside this wrapper. TODO: warning about use and greppability
    pub fn expose(self) -> T {
        self.inner
    }
}

macro_rules! pod_impl {
    { $tr:ident, $method:ident, $t:ty } => {
        impl $tr<$t> for Secret<$t> {
            type Output = Secret<<$t as $tr<$t>>::Output>;
            fn $method(self, rhs: $t) -> Self::Output {
                Secret::new($tr::$method(self.expose(), rhs))
            }
        }

        impl $tr<Secret<$t>> for $t {
            type Output = Secret<<$t as $tr<$t>>::Output>;
            fn $method(self, rhs: Secret<$t>) -> Self::Output {
                Secret::new($tr::$method(self, rhs.expose()))
            }
        }

        impl $tr<Secret<$t>> for Secret<$t> {
            type Output = Secret<<$t as $tr<$t>>::Output>;
            fn $method(self, rhs: Secret<$t>) -> Self::Output {
                Secret::new($tr::$method(self.expose(), rhs.expose()))
            }
        }
    };

    { $tr:ident, $method:ident, $trassign:ident, $methodassign: ident, $t:ty } => {
        pod_impl! { $tr, $method, $t }

        impl $trassign<Secret<$t>> for Secret<$t> {
            fn $methodassign(&mut self, rhs: Secret<$t>) {
                $trassign::$methodassign(&mut self.inner, rhs.expose());
            }
        }

        impl $trassign<$t> for Secret<$t> {
            fn $methodassign(&mut self, rhs: $t) {
                $trassign::$methodassign(&mut self.inner, rhs);
            }
        }
    };

    { shift, $tr:ident, $method:ident, $trassign:ident, $methodassign: ident, $t:ty } => {
        impl $tr<usize> for Secret<$t> {
            type Output = Secret<<$t as $tr<$t>>::Output>;
            fn $method(self, rhs: usize) -> Self::Output {
                Secret::new($tr::$method(self.expose(), rhs))
            }
        }

        impl $trassign<usize> for Secret<$t> {
            fn $methodassign(&mut self, rhs: usize) {
                $trassign::$methodassign(&mut self.inner, rhs);
            }
        }
    }
}

macro_rules! pod_impls {
    { $t:ty } => {
        impl Not for Secret<$t> {
            type Output = Secret<<$t as Not>::Output>;
            fn not(self) -> Self::Output {
                Secret::new(!self.expose())
            }
        }

        pod_impl! { BitAnd, bitand, BitAndAssign, bitand_assign, $t }
        pod_impl! { BitOr, bitor, BitOrAssign, bitor_assign, $t }
        pod_impl! { BitXor, bitxor, BitXorAssign, bitxor_assign, $t }
        pod_impl! { shift, Shl, shl, ShlAssign, shl_assign, $t }
        pod_impl! { shift, Shr, shr, ShrAssign, shr_assign, $t }
        pod_impl! { WrappingAdd, wrapping_add, $t }
        pod_impl! { WrappingSub, wrapping_sub, $t }
    };

    { $t:ty, $($rest:ty),* } => {
        pod_impls! { $t }
        pod_impls! { $($rest),* }

        $(
        impl From<$rest> for Secret<$t> {
            fn from(val: $rest) -> Secret<$t> {
                Secret::new(From::from(val))
            }
        }

        impl From<Secret<$rest>> for Secret<$t> {
            fn from(val: Secret<$rest>) -> Secret<$t> {
                Secret::new(From::from(val.expose()))
            }
        }
        )*
    };
}

pod_impls! { u64, u32, u16, u8 }
pod_impls! { i64, i32, i16, i8 }
<|endoftext|>


<commit_before>
use std::io::{TcpListener};
use std::io::{Listener, Acceptor};
use std::io::{IoResult};
use std::io::net;
use std::io;
use std::collections::hashmap::{HashMap};

use channel::{ChannelEvent};
use client::{SharedClient, Client, ClientId};
use msg::{MessageHandler};
use cmd;

/// Forwards the message to a channel
pub struct ChannelProxy {
    name: String,
    tx: Sender<ChannelEvent>,
    server_tx: Sender<Event>
}

impl ChannelProxy {
    pub fn new(name: String,
           tx: Sender<ChannelEvent>,
           server_tx: Sender<Event>) -> ChannelProxy {
        ChannelProxy {
            name: name,
            tx: tx,
            server_tx: server_tx
        }
    }
    pub fn send(&self, event: ChannelEvent) {
        match self.tx.send_opt(event) {
            Ok(_) => {},
            Err(_) => {
                let _ = self.server_tx.send_opt(ChannelLost(self.name.clone()));
            }
        }
    }
}

pub struct Server {
    host: String,
    ip: String,
    port: u16,
    tx: Option<Sender<Event>>,
    // TODO put unregisterd clients in a staging Map
    clients: HashMap<ClientId, SharedClient>,
    pub registered: HashMap<String, SharedClient>,
    pub channels: HashMap<String, ChannelProxy>
}

/// Enumeration of the events the server can receive
pub enum Event {
    /// Message received from a client
    MessageReceived(ClientId, Box<MessageHandler + Send>),
    /// Connection to a client established
    ClientConnected(Client),
    /// The task of Channel(name) failed
    ChannelLost(String),
}

/// Convenience function to run the server
pub fn run_server(host: &str) -> IoResult<Server> {
    let server = try!(Server::new(host));
    server.serve_forever()
}

/// Irc server
impl Server {
    /// Creates a new IRC server instance.
    pub fn new(host: &str) -> IoResult<Server> {
        let addresses = try!(net::get_host_addresses(host));
        debug!("{}", addresses)
        let ip = match addresses.as_slice().get(0) {
            Some(ip) => ip,
            None => return Err(io::IoError {
                kind: io::OtherIoError,
                desc: "cannnot resolve ip",
                detail: None
            })
        };
        Ok(Server {
            host: host.to_string(),
            ip: format!("{}", ip),
            port: 6667,
            tx: None,
            clients: HashMap::new(),
            registered: HashMap::new(),
            channels: HashMap::new()
        })
    }

    /// Starts the main loop and listens on the specified host and port.
    pub fn serve_forever(mut self) -> IoResult<Server> {
        // todo change this to a more general event dispatching loop
        let message_rx = try!(self.start_listening());
        for event in message_rx.iter() {
            match event {
                MessageReceived(client_id, handler) => {
                    let client = match self.clients.find(&client_id) {
                        Some(client) => Some(client.clone()),
                        None => None
                    };
                    match client {
                        Some(client) => handler.invoke(&mut self, client),
                        None => error!(
                            "Client {} not found when sending message.",
                            client_id
                        )
                    }
                },
                ClientConnected(client) => {
                    let client = client.as_shared();
                    self.clients.insert(client.borrow().id(), client);
                }
                ChannelLost(name) => {
                    self.channels.remove(&name);
                }
            }
        }
        Ok(self)
    }

    fn start_listening(&mut self) -> IoResult<Receiver<(Event)>>  {
        let listener = TcpListener::bind(self.ip.as_slice(), self.port);
        info!("started listening on {}:{} ({})", self.ip, self.port, self.host);
        let acceptor = try!(listener.listen());
        let (tx, rx) = channel();
        self.tx = Some(tx.clone());
        let host = self.host.clone();
        spawn(proc() {
            let mut a = acceptor; // https://github.com/rust-lang/rust/issues/11958
            for maybe_stream in a.incoming() {
                match maybe_stream {
                    Err(err) => { error!("{}", err) }
                    Ok(stream) => {
                        match Client::listen(host.clone(), stream, tx.clone()) {
                            Ok(()) => {
                            },
                            Err(err) => {
                                error!("{}", err)
                            }
                        }
                    }
                }
            }
        });
        Ok(rx)
    }

    /// Getter for hostname
    pub fn host(&self) -> &str {
        self.host.as_slice()
    }

    /// Getter for hostname
    pub fn tx(&self) -> Option<Sender<Event>> {
        self.tx.clone()
    }

    pub fn remove_client(&mut self, client: &SharedClient) {
        let client = client.borrow();
        self.registered.remove(&client.nickname);
        self.clients.remove(&client.id());
    }

    /// Sends a welcome message to a newly registered client
    pub fn send_welcome_msg(&self, client: &SharedClient) {
        client.borrow().send_response(cmd::RPL_WELCOME, None, None)
    }
}<commit_msg>listen only on ipv4<commit_after>
use std::io::{TcpListener};
use std::io::{Listener, Acceptor};
use std::io::{IoResult};
use std::io::net;
use std::io;
use std::collections::hashmap::{HashMap};

use channel::{ChannelEvent};
use client::{SharedClient, Client, ClientId};
use msg::{MessageHandler};
use cmd;

/// Forwards the message to a channel
pub struct ChannelProxy {
    name: String,
    tx: Sender<ChannelEvent>,
    server_tx: Sender<Event>
}

impl ChannelProxy {
    pub fn new(name: String,
           tx: Sender<ChannelEvent>,
           server_tx: Sender<Event>) -> ChannelProxy {
        ChannelProxy {
            name: name,
            tx: tx,
            server_tx: server_tx
        }
    }
    pub fn send(&self, event: ChannelEvent) {
        match self.tx.send_opt(event) {
            Ok(_) => {},
            Err(_) => {
                let _ = self.server_tx.send_opt(ChannelLost(self.name.clone()));
            }
        }
    }
}

pub struct Server {
    host: String,
    ip: String,
    port: u16,
    tx: Option<Sender<Event>>,
    // TODO put unregisterd clients in a staging Map
    clients: HashMap<ClientId, SharedClient>,
    pub registered: HashMap<String, SharedClient>,
    pub channels: HashMap<String, ChannelProxy>
}

/// Enumeration of the events the server can receive
pub enum Event {
    /// Message received from a client
    MessageReceived(ClientId, Box<MessageHandler + Send>),
    /// Connection to a client established
    ClientConnected(Client),
    /// The task of Channel(name) failed
    ChannelLost(String),
}

/// Convenience function to run the server
pub fn run_server(host: &str) -> IoResult<Server> {
    let server = try!(Server::new(host));
    server.serve_forever()
}

/// Irc server
impl Server {
    /// Creates a new IRC server instance.
    pub fn new(host: &str) -> IoResult<Server> {
        let addresses = try!(net::get_host_addresses(host));
        debug!("{}", addresses)
        // Listen only on ipv4 for now…
        let ip = match addresses.iter().filter(
            |&v| match v { &net::ip::Ipv4Addr(_, _, _, _) => true, _ => false }
        ).nth(0) {
            Some(ip) => ip,
            None => return Err(io::IoError {
                kind: io::OtherIoError,
                desc: "cannnot resolve ip",
                detail: None
            })
        };
        Ok(Server {
            host: host.to_string(),
            ip: format!("{}", ip),
            port: 6667,
            tx: None,
            clients: HashMap::new(),
            registered: HashMap::new(),
            channels: HashMap::new()
        })
    }

    /// Starts the main loop and listens on the specified host and port.
    pub fn serve_forever(mut self) -> IoResult<Server> {
        // todo change this to a more general event dispatching loop
        let message_rx = try!(self.start_listening());
        for event in message_rx.iter() {
            match event {
                MessageReceived(client_id, handler) => {
                    let client = match self.clients.find(&client_id) {
                        Some(client) => Some(client.clone()),
                        None => None
                    };
                    match client {
                        Some(client) => handler.invoke(&mut self, client),
                        None => error!(
                            "Client {} not found when sending message.",
                            client_id
                        )
                    }
                },
                ClientConnected(client) => {
                    let client = client.as_shared();
                    self.clients.insert(client.borrow().id(), client);
                }
                ChannelLost(name) => {
                    self.channels.remove(&name);
                }
            }
        }
        Ok(self)
    }

    fn start_listening(&mut self) -> IoResult<Receiver<(Event)>>  {
        let listener = TcpListener::bind(self.ip.as_slice(), self.port);
        info!("started listening on {}:{} ({})", self.ip, self.port, self.host);
        let acceptor = try!(listener.listen());
        let (tx, rx) = channel();
        self.tx = Some(tx.clone());
        let host = self.host.clone();
        spawn(proc() {
            let mut a = acceptor; // https://github.com/rust-lang/rust/issues/11958
            for maybe_stream in a.incoming() {
                match maybe_stream {
                    Err(err) => { error!("{}", err) }
                    Ok(stream) => {
                        match Client::listen(host.clone(), stream, tx.clone()) {
                            Ok(()) => {
                            },
                            Err(err) => {
                                error!("{}", err)
                            }
                        }
                    }
                }
            }
        });
        Ok(rx)
    }

    /// Getter for hostname
    pub fn host(&self) -> &str {
        self.host.as_slice()
    }

    /// Getter for hostname
    pub fn tx(&self) -> Option<Sender<Event>> {
        self.tx.clone()
    }

    pub fn remove_client(&mut self, client: &SharedClient) {
        let client = client.borrow();
        self.registered.remove(&client.nickname);
        self.clients.remove(&client.id());
    }

    /// Sends a welcome message to a newly registered client
    pub fn send_welcome_msg(&self, client: &SharedClient) {
        client.borrow().send_response(cmd::RPL_WELCOME, None, None)
    }
}<|endoftext|>


<commit_before>use std::env;
use std::io::ErrorKind;
use std::process;
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::sync::Arc;
use std::time::Duration;
use std::net::SocketAddr;
use hex;
use time;

use byteorder::{LittleEndian, WriteBytesExt};

use mio::net::UdpSocket;
use mio::{Events, Poll, PollOpt, Ready, Token};
use mio_extras::timer::Timer;

use config;
use config::ServerConfig;
use kms;
use key::{LongTermKey, OnlineKey};
use merkle::MerkleTree;
use {Error, RtMessage, Tag};
use {MIN_REQUEST_LENGTH, VERSION};

macro_rules! check_ctrlc {
    ($keep_running:expr) => {
        if !$keep_running.load(Ordering::Acquire) {
            warn!("Ctrl-C caught, exiting...");
            return true;
        }
    }
}

const MESSAGE: Token = Token(0);
const STATUS: Token = Token(1);

fn make_response(srep: &RtMessage, cert_bytes: &[u8], path: &[u8], idx: u32) -> RtMessage {
    let mut index = [0; 4];
    (&mut index as &mut [u8])
        .write_u32::<LittleEndian>(idx)
        .unwrap();

    let sig_bytes = srep.get_field(Tag::SIG).unwrap();
    let srep_bytes = srep.get_field(Tag::SREP).unwrap();

    let mut response = RtMessage::new(5);
    response.add_field(Tag::SIG, sig_bytes).unwrap();
    response.add_field(Tag::PATH, path).unwrap();
    response.add_field(Tag::SREP, srep_bytes).unwrap();
    response.add_field(Tag::CERT, cert_bytes).unwrap();
    response.add_field(Tag::INDX, &index).unwrap();

    response
}

// extract the client's nonce from its request
fn nonce_from_request(buf: &[u8], num_bytes: usize) -> Result<&[u8], Error> {
    if num_bytes < MIN_REQUEST_LENGTH as usize {
        return Err(Error::RequestTooShort);
    }

    let tag_count = &buf[..4];
    let expected_nonc = &buf[8..12];
    let expected_pad = &buf[12..16];

    let tag_count_is_2 = tag_count == [0x02, 0x00, 0x00, 0x00];
    let tag1_is_nonc = expected_nonc == Tag::NONC.wire_value();
    let tag2_is_pad = expected_pad == Tag::PAD.wire_value();

    if tag_count_is_2 && tag1_is_nonc && tag2_is_pad {
        Ok(&buf[0x10..0x50])
    } else {
        Err(Error::InvalidRequest)
    }
}

pub struct Server {
    config: Box<ServerConfig>,
    online_key: OnlineKey,
    cert_bytes: Vec<u8>,

    response_counter: AtomicUsize,
    num_bad_requests: u64,

    socket: UdpSocket,
    keep_running: Arc<AtomicBool>,
    poll_duration: Option<Duration>,
    timer: Timer<()>,
    poll: Poll,
    events: Events,
    merkle: MerkleTree,
    requests: Vec<(Vec<u8>, SocketAddr)>,
    buf: [u8; 65_536],

    public_key: String,
}

impl Server {
    pub fn new(config: Box<ServerConfig>) -> Server {

        let mut online_key = OnlineKey::new();
        let public_key: String;

        let cert_bytes = {
            let seed = kms::load_seed(&config).unwrap();
            println!("Seed: {:?}", seed);
            let mut long_term_key = LongTermKey::new(&seed);
            public_key = hex::encode(long_term_key.public_key());

            long_term_key.make_cert(&online_key).encode().unwrap()
        };

        let response_counter = AtomicUsize::new(0);
        let keep_running = Arc::new(AtomicBool::new(true));

        let sock_addr = config.socket_addr().expect("");
        let socket = UdpSocket::bind(&sock_addr).expect("failed to bind to socket");
        let poll_duration = Some(Duration::from_millis(100));

        let mut timer: Timer<()> = Timer::default();
        timer.set_timeout(config.status_interval(), ());

        let poll = Poll::new().unwrap();
        poll.register(&socket, MESSAGE, Ready::readable(), PollOpt::edge())
            .unwrap();
        poll.register(&timer, STATUS, Ready::readable(), PollOpt::edge())
            .unwrap();

        let mut merkle = MerkleTree::new();
        let mut requests = Vec::with_capacity(config.batch_size() as usize);

        Server {
            config,
            online_key,
            cert_bytes,

            response_counter,
            num_bad_requests: 0,
            socket,
            keep_running,
            poll_duration,
            timer,
            poll,
            events: Events::with_capacity(32),
            merkle,
            requests,
            buf: [0u8; 65_536],

            public_key
        }

    }

    pub fn get_keep_running(&self) -> Arc<AtomicBool> {
        return self.keep_running.clone()
    }

    pub fn process_events(&mut self) -> bool {
        self.poll.poll(&mut self.events, self.poll_duration).expect("poll failed");

        for event in self.events.iter() {
            match event.token() {
                MESSAGE => {
                    let mut done = false;

                    'process_batch: loop {
                        check_ctrlc!(self.keep_running);

                        self.merkle.reset();
                        self.requests.clear();

                        let resp_start = self.response_counter.load(Ordering::SeqCst);

                        for i in 0..self.config.batch_size() {
                            match self.socket.recv_from(&mut self.buf) {
                                Ok((num_bytes, src_addr)) => {
                                    if let Ok(nonce) = nonce_from_request(&self.buf, num_bytes) {
                                        self.requests.push((Vec::from(nonce), src_addr));
                                        self.merkle.push_leaf(nonce);
                                    } else {
                                        self.num_bad_requests += 1;
                                        info!(
                                            "Invalid request ({} bytes) from {} (#{} in batch, resp #{})",
                                            num_bytes, src_addr, i, resp_start + i as usize
                                        );
                                    }
                                }
                                Err(e) => match e.kind() {
                                    ErrorKind::WouldBlock => {
                                        done = true;
                                        break;
                                    }
                                    _ => {
                                        error!(
                                            "Error receiving from socket: {:?}: {:?}",
                                            e.kind(),
                                            e
                                        );
                                        break;
                                    }
                                },
                            };
                        }

                        if self.requests.is_empty() {
                            break 'process_batch;
                        }

                        let merkle_root = self.merkle.compute_root();
                        let srep = self.online_key.make_srep(time::get_time(), &merkle_root);

                        for (i, &(ref nonce, ref src_addr)) in self.requests.iter().enumerate() {
                            let paths = self.merkle.get_paths(i);

                            let resp = make_response(&srep, &self.cert_bytes, &paths, i as u32);
                            let resp_bytes = resp.encode().unwrap();

                            let bytes_sent = self.socket
                                .send_to(&resp_bytes, &src_addr)
                                .expect("send_to failed");
                            let num_responses = self.response_counter.fetch_add(1, Ordering::SeqCst);

                            info!(
                                "Responded {} bytes to {} for '{}..' (#{} in batch, resp #{})",
                                bytes_sent,
                                src_addr,
                                hex::encode(&nonce[0..4]),
                                i,
                                num_responses
                            );
                        }
                        if done {
                            break 'process_batch;
                        }
                    }
                }

                STATUS => {
                    info!(
                        "responses {}, invalid requests {}",
                        self.response_counter.load(Ordering::SeqCst),
                        self.num_bad_requests
                    );

                    self.timer.set_timeout(self.config.status_interval(), ());
                }

                _ => unreachable!(),
            }
        }
        false
    }

    pub fn send_to_self(&mut self, data: &[u8]) {
        self.socket.send_to(data, &self.socket.local_addr().unwrap());
    }

    pub fn get_public_key(&self) -> &str {
        return &self.public_key
    }

    pub fn get_online_key(&self) -> &OnlineKey {
        return &self.online_key
    }

    pub fn get_config(&self) -> &Box<ServerConfig> {
        return &self.config
    }
}
<commit_msg>More changes to server<commit_after>use std::env;
use std::io::ErrorKind;
use std::process;
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::sync::Arc;
use std::time::Duration;
use std::net::SocketAddr;
use hex;
use time;

use byteorder::{LittleEndian, WriteBytesExt};

use mio::net::UdpSocket;
use mio::{Events, Poll, PollOpt, Ready, Token};
use mio_extras::timer::Timer;

use config;
use config::ServerConfig;
use kms;
use key::{LongTermKey, OnlineKey};
use merkle::MerkleTree;
use {Error, RtMessage, Tag};
use {MIN_REQUEST_LENGTH, VERSION};

macro_rules! check_ctrlc {
    ($keep_running:expr) => {
        if !$keep_running.load(Ordering::Acquire) {
            warn!("Ctrl-C caught, exiting...");
            return true;
        }
    }
}

const MESSAGE: Token = Token(0);
const STATUS: Token = Token(1);

fn make_response(srep: &RtMessage, cert_bytes: &[u8], path: &[u8], idx: u32) -> RtMessage {
    let mut index = [0; 4];
    (&mut index as &mut [u8])
        .write_u32::<LittleEndian>(idx)
        .unwrap();

    let sig_bytes = srep.get_field(Tag::SIG).unwrap();
    let srep_bytes = srep.get_field(Tag::SREP).unwrap();

    let mut response = RtMessage::new(5);
    response.add_field(Tag::SIG, sig_bytes).unwrap();
    response.add_field(Tag::PATH, path).unwrap();
    response.add_field(Tag::SREP, srep_bytes).unwrap();
    response.add_field(Tag::CERT, cert_bytes).unwrap();
    response.add_field(Tag::INDX, &index).unwrap();

    response
}

// extract the client's nonce from its request
fn nonce_from_request(buf: &[u8], num_bytes: usize) -> Result<&[u8], Error> {
    if num_bytes < MIN_REQUEST_LENGTH as usize {
        return Err(Error::RequestTooShort);
    }

    let tag_count = &buf[..4];
    let expected_nonc = &buf[8..12];
    let expected_pad = &buf[12..16];

    let tag_count_is_2 = tag_count == [0x02, 0x00, 0x00, 0x00];
    let tag1_is_nonc = expected_nonc == Tag::NONC.wire_value();
    let tag2_is_pad = expected_pad == Tag::PAD.wire_value();

    if tag_count_is_2 && tag1_is_nonc && tag2_is_pad {
        Ok(&buf[0x10..0x50])
    } else {
        Err(Error::InvalidRequest)
    }
}

pub struct Server {
    config: Box<ServerConfig>,
    online_key: OnlineKey,
    cert_bytes: Vec<u8>,

    response_counter: AtomicUsize,
    num_bad_requests: u64,

    socket: UdpSocket,
    fake_client_socket: UdpSocket,
    keep_running: Arc<AtomicBool>,
    poll_duration: Option<Duration>,
    timer: Timer<()>,
    poll: Poll,
    events: Events,
    merkle: MerkleTree,
    requests: Vec<(Vec<u8>, SocketAddr)>,
    buf: [u8; 65_536],

    public_key: String,
}

impl Server {
    pub fn new(config: Box<ServerConfig>) -> Server {

        let mut online_key = OnlineKey::new();
        let public_key: String;

        let cert_bytes = {
            let seed = kms::load_seed(&config).unwrap();
            println!("Seed: {:?}", seed);
            let mut long_term_key = LongTermKey::new(&seed);
            public_key = hex::encode(long_term_key.public_key());

            long_term_key.make_cert(&online_key).encode().unwrap()
        };

        let response_counter = AtomicUsize::new(0);
        let keep_running = Arc::new(AtomicBool::new(true));

        let sock_addr = config.socket_addr().expect("");
        let socket = UdpSocket::bind(&sock_addr).expect("failed to bind to socket");
        let fake_client_socket = UdpSocket::bind(&"127.0.0.1:0".parse().unwrap()).unwrap();

        let poll_duration = Some(Duration::from_millis(100));

        let mut timer: Timer<()> = Timer::default();
        //timer.set_timeout(config.status_interval(), ());

        let poll = Poll::new().unwrap();
        poll.register(&socket, MESSAGE, Ready::readable(), PollOpt::edge())
            .unwrap();
        //poll.register(&timer, STATUS, Ready::readable(), PollOpt::edge())
        //    .unwrap();

        let mut merkle = MerkleTree::new();
        let mut requests = Vec::with_capacity(config.batch_size() as usize);

        Server {
            config,
            online_key,
            cert_bytes,

            response_counter,
            num_bad_requests: 0,
            socket,
            fake_client_socket,
            keep_running,
            poll_duration,
            timer,
            poll,
            events: Events::with_capacity(32),
            merkle,
            requests,
            buf: [0u8; 65_536],

            public_key
        }

    }

    pub fn get_keep_running(&self) -> Arc<AtomicBool> {
        return self.keep_running.clone()
    }

    pub fn process_events(&mut self) -> bool {
        self.poll.poll(&mut self.events, self.poll_duration).expect("poll failed");

        for event in self.events.iter() {
            match event.token() {
                MESSAGE => {
                    let mut done = false;

                    'process_batch: loop {
                        check_ctrlc!(self.keep_running);

                        let resp_start = self.response_counter.load(Ordering::SeqCst);

                        for i in 0..self.config.batch_size() {
                            match self.socket.recv_from(&mut self.buf) {
                                Ok((num_bytes, src_addr)) => {
                                    info!("Read bytes: {}", num_bytes);
                                    match nonce_from_request(&self.buf, num_bytes) {
                                        Ok(nonce) => {
                                            self.requests.push((Vec::from(nonce), src_addr));
                                            self.merkle.push_leaf(nonce);
                                        },
                                        Err(e) => {
                                           self.num_bad_requests += 1;

                                             info!(
                                                "Invalid request: '{:?}' ({} bytes) from {} (#{} in batch, resp #{})",
                                                e, num_bytes, src_addr, i, resp_start + i as usize
                                            );
                                        }
                                    }
                                },
                                Err(e) => match e.kind() {
                                    ErrorKind::WouldBlock => {
                                        done = true;
                                        break;
                                    }
                                    _ => {
                                        error!(
                                            "Error receiving from socket: {:?}: {:?}",
                                            e.kind(),
                                            e
                                        );
                                        break;
                                    }
                                },
                            };
                        }

                        if self.requests.is_empty() {
                            break 'process_batch;
                        }

                        let merkle_root = self.merkle.compute_root();
                        let srep = self.online_key.make_srep(time::get_time(), &merkle_root);

                        for (i, &(ref nonce, ref src_addr)) in self.requests.iter().enumerate() {
                            let paths = self.merkle.get_paths(i);

                            let resp = make_response(&srep, &self.cert_bytes, &paths, i as u32);
                            let resp_bytes = resp.encode().unwrap();

                            let bytes_sent = self.socket
                                .send_to(&resp_bytes, &src_addr)
                                .expect("send_to failed");
                            let num_responses = self.response_counter.fetch_add(1, Ordering::SeqCst);

                            info!(
                                "Responded {} bytes to {} for '{}..' (#{} in batch, resp #{})",
                                bytes_sent,
                                src_addr,
                                hex::encode(&nonce[0..4]),
                                i,
                                num_responses
                            );
                        }

                        self.merkle.reset();
                        self.requests.clear();

                        if done {
                            break 'process_batch;
                        }
                    }
                }

                STATUS => {
                    info!(
                        "responses {}, invalid requests {}",
                        self.response_counter.load(Ordering::SeqCst),
                        self.num_bad_requests
                    );

                    self.timer.set_timeout(self.config.status_interval(), ());
                }

                _ => unreachable!(),
            }
        }
        false
    }

    pub fn send_to_self(&mut self, data: &[u8]) {
        self.response_counter.store(0,  Ordering::SeqCst);;
        self.num_bad_requests = 0;
        let res = self.fake_client_socket.send_to(data, &self.socket.local_addr().unwrap());
        info!("Sent to self: {:?}", res);
    }

    pub fn get_public_key(&self) -> &str {
        return &self.public_key
    }

    pub fn get_online_key(&self) -> &OnlineKey {
        return &self.online_key
    }

    pub fn get_config(&self) -> &Box<ServerConfig> {
        return &self.config
    }
}
<|endoftext|>