<commit_before>extern crate reqwest;
extern crate serde;
extern crate serde_json;
extern crate serde_with;
extern crate url;

use reqwest::header::{Authorization, Bearer, Range};
use url::Url;

#[derive(Serialize, Deserialize, Debug)]
#[serde(rename_all = "camelCase")]
pub struct Bucket {
    id: String,
    name: String,
    location: String,
    self_link: String,
}

#[derive(Serialize, Deserialize, Debug)]
#[serde(rename_all = "camelCase")]
pub struct Object {
    id: String,
    name: String,
    bucket: String,
    self_link: String,
    #[serde(with = "serde_with::rust::display_fromstr")]
    size: u64,
}

#[derive(Serialize, Deserialize, Debug)]
#[serde(rename_all = "camelCase")]
pub struct ListObjectsResponse {
    next_page_token: Option<String>,
    items: Vec<Object>,
}

fn new_client() -> Result<reqwest::Client, reqwest::Error> {
    // NOTE(boulos): The service account needs both storage viewer (to see objects) and *project* viewer to see the Bucket.
    let creds = Bearer {
        token: "IMABADPERSON_HARDCODED_HERE_gcloud_auth_print_access_token_for_the_service_account".to_owned()
    };

    let mut headers = reqwest::header::Headers::new();
    headers.set(Authorization(creds));

    return reqwest::Client::builder()
        .default_headers(headers)
        .build();
}

fn get_bucket(url: Url) -> Result<Bucket, reqwest::Error> {
    println!("Looking to request: {:#?}", url);

    let client = new_client()?;

    let mut response = client.get(url)
        .send()
        .expect("Failed to send request");

    println!("{:#?}", response);

    let bucket = response.json::<Bucket>();
    println!("{:#?}", bucket);
    return bucket
}

fn get_object(url: Url) -> Result<Object, reqwest::Error> {
    println!("Looking to request: {:#?}", url);

    let client = new_client()?;

    let mut response = client.get(url)
        .send()
        .expect("Failed to send request");

    println!("{:#?}", response);

    let object = response.json::<Object>();
    println!("{:#?}", object);
    return object
}

pub fn get_bytes(obj: &Object, offset: u64, how_many: u64) -> Result<Vec<u8>, reqwest::Error> {
    println!("Asking for {} bytes at {} from the origin for {}", how_many, offset, obj.name);

    // Use the self_link from the object as the url, but add ?alt=media
    let mut object_url = Url::parse(&obj.self_link).unwrap();
    object_url.query_pairs_mut().append_pair("alt", "media");

    let client = new_client()?;

    let mut response = client.get(object_url)
        .header(Range::bytes(offset, offset + how_many))
        .send()
        .expect("Failed to send request");

    let mut buf: Vec<u8> = vec![];
    let written = response.copy_to(&mut buf)?;
    println!("Got back {} bytes", written);
    Ok(buf)
}

fn _do_one_list_object(bucket: &str, prefix: Option<&str>, delim: Option<&str>, token: Option<&str>) -> Result<ListObjectsResponse, reqwest::Error> {
    let mut list_url = Url::parse(bucket).unwrap();
    if let Some(prefix_str) = prefix {
        list_url.query_pairs_mut().append_pair("prefix", prefix_str);
    }

    if let Some(delim_str) = delim {
        list_url.query_pairs_mut().append_pair("delimiter", delim_str);
    }

    if let Some(token_str) = token {
        list_url.query_pairs_mut().append_pair("pageToken", token_str);
    }

    let client = new_client()?;

    let mut response = client.get(list_url)
        .send()
        .expect("Failed to send request");

    let list_response = response.json::<ListObjectsResponse>();
    return list_response;
}

fn list_objects(bucket: &str, prefix: Option<&str>, delim: Option<&str>) -> Result<Vec<Object>, reqwest::Error> {
    println!("Asking for a list from bucket '{}' with prefix '{:#?}' and delim = '{:#?}'", bucket, prefix, delim);

    let mut results: Vec<Object> = vec![];

    let mut page_token: String = String::from("");

    loop {
        let mut resp = match page_token.is_empty() {
            true => _do_one_list_object(bucket, prefix, delim, None)?,
            false => _do_one_list_object(bucket, prefix, delim, Some(&page_token))?,
        };

        results.append(&mut resp.items);

        match resp.next_page_token {
            Some(temp_token_str) => page_token = temp_token_str.clone(),
            None => break,
        }
    }

    Ok(results)
}

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

    #[test]
    fn get_landsat() {
        let landsat_url = Url::parse(
            "https://www.googleapis.com/storage/v1/b/gcp-public-data-landsat"
        ).unwrap();
        get_bucket(landsat_url);
    }

    #[test]
    fn get_private_bucket() {
        let bucket_url = Url::parse(
            "https://www.googleapis.com/storage/v1/b/boulos-vm-ml"
        ).unwrap();
        get_bucket(bucket_url);
    }

    #[test]
    fn get_private_object() {
        let object_url = Url::parse(
            "https://www.googleapis.com/storage/v1/b/boulos-hadoop/o/bdutil-staging%2fhadoop-m%2f20150202-172447-58j%2fbq-mapred-template.xml"
        ).unwrap();
        let object: Object = get_object(object_url).unwrap();
        println!("Object has {} bytes", object.size);

        let bytes: Vec<u8> = get_bytes(&object, 0, 769).unwrap();
        println!("Got back:\n {}", String::from_utf8(bytes).unwrap());

        let offset_bytes: Vec<u8> = get_bytes(&object, 6, 769).unwrap();
        println!("Got back:\n {}", String::from_utf8(offset_bytes).unwrap());

    }

    #[test]
    fn test_list_objects() {
        let object_url = "https://www.googleapis.com/storage/v1/b/boulos-hadoop/o";
        let prefix = "bdutil-staging";
        let delim = "%2F";

        let objects: Vec<Object> = list_objects(object_url, Some(prefix), Some(delim)).unwrap();
        println!("Got {} objects", objects.len());
        println!("Dump:\n\n{:#?}", objects);

        let all_objects: Vec<Object> = list_objects(object_url, None, None).unwrap();
        println!("Got {} objects", all_objects.len());
        println!("Dump:\n\n{:#?}", all_objects);

        let only_top_level: Vec<Object> = list_objects(object_url, None, Some(delim)).unwrap();
        println!("Got {} objects", only_top_level.len());
        println!("Dump:\n\n{:#?}", only_top_level);
    }

    #[test]
    fn test_list_paginated() {
        let object_url = "https://www.googleapis.com/storage/v1/b/gcp-public-data-landsat/o";
        let prefix = "LC08/PRE/044/034/";
        let delim = "%2F";

        let objects: Vec<Object> = list_objects(object_url, Some(prefix), Some(delim)).unwrap();
        println!("Got {} objects", objects.len());
    }
}
<commit_msg>Update http fetching for new hyper HTTP/2 support.<commit_after>extern crate reqwest;
extern crate serde;
extern crate serde_json;
extern crate serde_with;
extern crate url;

use url::Url;

#[derive(Serialize, Deserialize, Debug)]
#[serde(rename_all = "camelCase")]
pub struct Bucket {
    id: String,
    name: String,
    location: String,
    self_link: String,
}

#[derive(Serialize, Deserialize, Debug)]
#[serde(rename_all = "camelCase")]
pub struct Object {
    id: String,
    name: String,
    bucket: String,
    self_link: String,
    #[serde(with = "serde_with::rust::display_fromstr")]
    size: u64,
}

#[derive(Serialize, Deserialize, Debug)]
#[serde(rename_all = "camelCase")]
pub struct ListObjectsResponse {
    next_page_token: Option<String>,
    items: Vec<Object>,
}

// NOTE(boulos): The service account needs both storage viewer (to see objects) and *project* viewer to see the Bucket.
static GCLOUD_TOKEN: &'static str = "IMABADPERSON_HARDCODED_HERE_gcloud_auth_print_access_token_for_the_service_account";

fn new_client() -> Result<reqwest::Client, reqwest::Error> {
    // NOTE(boulos): Previously, I was passing some default headers
    // here. Now that bearer_auth is per request, this is less needed,
    // but we can change that later.
    return reqwest::Client::builder()
        .build();
}

fn get_bucket(url: Url) -> Result<Bucket, reqwest::Error> {
    println!("Looking to request: {:#?}", url);

    let client = new_client()?;

    let mut response = client.get(url)
        .bearer_auth(GCLOUD_TOKEN)
        .send()
        .expect("Failed to send request");

    println!("{:#?}", response);

    let bucket = response.json::<Bucket>();
    println!("{:#?}", bucket);
    return bucket
}

fn get_object(url: Url) -> Result<Object, reqwest::Error> {
    println!("Looking to request: {:#?}", url);

    let client = new_client()?;

    let mut response = client.get(url)
        .bearer_auth(GCLOUD_TOKEN)
        .send()
        .expect("Failed to send request");

    println!("{:#?}", response);

    let object = response.json::<Object>();
    println!("{:#?}", object);
    return object
}

pub fn get_bytes(obj: &Object, offset: u64, how_many: u64) -> Result<Vec<u8>, reqwest::Error> {
    println!("Asking for {} bytes at {} from the origin for {}", how_many, offset, obj.name);

    // Use the self_link from the object as the url, but add ?alt=media
    let mut object_url = Url::parse(&obj.self_link).unwrap();
    object_url.query_pairs_mut().append_pair("alt", "media");

    let client = new_client()?;

    let byte_range = format!("bytes={}-{}", offset, offset + how_many);

    let mut response = client.get(object_url)
        .header(reqwest::header::RANGE, byte_range)
        .bearer_auth(GCLOUD_TOKEN)
        .send()
        .expect("Failed to send request");

    let mut buf: Vec<u8> = vec![];
    let written = response.copy_to(&mut buf)?;
    println!("Got back {} bytes", written);
    Ok(buf)
}

fn _do_one_list_object(bucket: &str, prefix: Option<&str>, delim: Option<&str>, token: Option<&str>) -> Result<ListObjectsResponse, reqwest::Error> {
    let mut list_url = Url::parse(bucket).unwrap();
    if let Some(prefix_str) = prefix {
        list_url.query_pairs_mut().append_pair("prefix", prefix_str);
    }

    if let Some(delim_str) = delim {
        list_url.query_pairs_mut().append_pair("delimiter", delim_str);
    }

    if let Some(token_str) = token {
        list_url.query_pairs_mut().append_pair("pageToken", token_str);
    }

    let client = new_client()?;

    let mut response = client.get(list_url)
        .bearer_auth(GCLOUD_TOKEN)
        .send()
        .expect("Failed to send request");

    let list_response = response.json::<ListObjectsResponse>();
    return list_response;
}

fn list_objects(bucket: &str, prefix: Option<&str>, delim: Option<&str>) -> Result<Vec<Object>, reqwest::Error> {
    println!("Asking for a list from bucket '{}' with prefix '{:#?}' and delim = '{:#?}'", bucket, prefix, delim);

    let mut results: Vec<Object> = vec![];

    let mut page_token: String = String::from("");

    loop {
        let mut resp = match page_token.is_empty() {
            true => _do_one_list_object(bucket, prefix, delim, None)?,
            false => _do_one_list_object(bucket, prefix, delim, Some(&page_token))?,
        };

        results.append(&mut resp.items);

        match resp.next_page_token {
            Some(temp_token_str) => page_token = temp_token_str.clone(),
            None => break,
        }
    }

    Ok(results)
}

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

    #[test]
    fn get_landsat() {
        let landsat_url = Url::parse(
            "https://www.googleapis.com/storage/v1/b/gcp-public-data-landsat"
        ).unwrap();
        get_bucket(landsat_url);
    }

    #[test]
    fn get_private_bucket() {
        let bucket_url = Url::parse(
            "https://www.googleapis.com/storage/v1/b/boulos-vm-ml"
        ).unwrap();
        get_bucket(bucket_url);
    }

    #[test]
    fn get_private_object() {
        let object_url = Url::parse(
            "https://www.googleapis.com/storage/v1/b/boulos-hadoop/o/bdutil-staging%2fhadoop-m%2f20150202-172447-58j%2fbq-mapred-template.xml"
        ).unwrap();
        let object: Object = get_object(object_url).unwrap();
        println!("Object has {} bytes", object.size);

        let bytes: Vec<u8> = get_bytes(&object, 0, 769).unwrap();
        println!("Got back:\n {}", String::from_utf8(bytes).unwrap());

        let offset_bytes: Vec<u8> = get_bytes(&object, 6, 769).unwrap();
        println!("Got back:\n {}", String::from_utf8(offset_bytes).unwrap());

    }

    #[test]
    fn test_list_objects() {
        let object_url = "https://www.googleapis.com/storage/v1/b/boulos-hadoop/o";
        let prefix = "bdutil-staging";
        let delim = "%2F";

        let objects: Vec<Object> = list_objects(object_url, Some(prefix), Some(delim)).unwrap();
        println!("Got {} objects", objects.len());
        println!("Dump:\n\n{:#?}", objects);

        let all_objects: Vec<Object> = list_objects(object_url, None, None).unwrap();
        println!("Got {} objects", all_objects.len());
        println!("Dump:\n\n{:#?}", all_objects);

        let only_top_level: Vec<Object> = list_objects(object_url, None, Some(delim)).unwrap();
        println!("Got {} objects", only_top_level.len());
        println!("Dump:\n\n{:#?}", only_top_level);
    }

    #[test]
    fn test_list_paginated() {
        let object_url = "https://www.googleapis.com/storage/v1/b/gcp-public-data-landsat/o";
        let prefix = "LC08/PRE/044/034/";
        let delim = "%2F";

        let objects: Vec<Object> = list_objects(object_url, Some(prefix), Some(delim)).unwrap();
        println!("Got {} objects", objects.len());
    }
}
<|endoftext|>


<commit_before>use std::env::*;
use std::ffi::{OsStr, OsString};
use std::path::PathBuf;

use rand::distributions::Alphanumeric;
use rand::{thread_rng, Rng};

fn make_rand_name() -> OsString {
    let rng = thread_rng();
    let n = format!("TEST{}", rng.sample_iter(&Alphanumeric).take(10).collect::<String>());
    let n = OsString::from(n);
    assert!(var_os(&n).is_none());
    n
}

fn eq(a: Option<OsString>, b: Option<&str>) {
    assert_eq!(a.as_ref().map(|s| &**s), b.map(OsStr::new).map(|s| &*s));
}

#[test]
fn test_set_var() {
    let n = make_rand_name();
    set_var(&n, "VALUE");
    eq(var_os(&n), Some("VALUE"));
}

#[test]
fn test_remove_var() {
    let n = make_rand_name();
    set_var(&n, "VALUE");
    remove_var(&n);
    eq(var_os(&n), None);
}

#[test]
fn test_set_var_overwrite() {
    let n = make_rand_name();
    set_var(&n, "1");
    set_var(&n, "2");
    eq(var_os(&n), Some("2"));
    set_var(&n, "");
    eq(var_os(&n), Some(""));
}

#[test]
#[cfg_attr(target_os = "emscripten", ignore)]
fn test_var_big() {
    let mut s = "".to_string();
    let mut i = 0;
    while i < 100 {
        s.push_str("aaaaaaaaaa");
        i += 1;
    }
    let n = make_rand_name();
    set_var(&n, &s);
    eq(var_os(&n), Some(&s));
}

#[test]
#[cfg_attr(target_os = "emscripten", ignore)]
fn test_env_set_get_huge() {
    let n = make_rand_name();
    let s = "x".repeat(10000);
    set_var(&n, &s);
    eq(var_os(&n), Some(&s));
    remove_var(&n);
    eq(var_os(&n), None);
}

#[test]
fn test_env_set_var() {
    let n = make_rand_name();

    let mut e = vars_os();
    set_var(&n, "VALUE");
    assert!(!e.any(|(k, v)| { &*k == &*n && &*v == "VALUE" }));

    assert!(vars_os().any(|(k, v)| { &*k == &*n && &*v == "VALUE" }));
}

#[test]
#[cfg_attr(any(target_os = "emscripten", target_env = "sgx"), ignore)]
#[allow(deprecated)]
fn env_home_dir() {
    fn var_to_os_string(var: Result<String, VarError>) -> Option<OsString> {
        match var {
            Ok(var) => Some(OsString::from(var)),
            Err(VarError::NotUnicode(var)) => Some(var),
            _ => None,
        }
    }

    cfg_if::cfg_if! {
        if #[cfg(unix)] {
            let oldhome = var_to_os_string(var("HOME"));

            set_var("HOME", "/home/MountainView");
            assert_eq!(home_dir(), Some(PathBuf::from("/home/MountainView")));

            remove_var("HOME");
            if cfg!(target_os = "android") {
                assert!(home_dir().is_none());
            } else {
                // When HOME is not set, some platforms return `None`,
                // but others return `Some` with a default.
                // Just check that it is not "/home/MountainView".
                assert_ne!(home_dir(), Some(PathBuf::from("/home/MountainView")));
            }

            if let Some(oldhome) = oldhome { set_var("HOME", oldhome); }
        } else if #[cfg(windows)] {
            let oldhome = var_to_os_string(var("HOME"));
            let olduserprofile = var_to_os_string(var("USERPROFILE"));

            remove_var("HOME");
            remove_var("USERPROFILE");

            assert!(home_dir().is_some());

            set_var("HOME", "/home/MountainView");
            assert_eq!(home_dir(), Some(PathBuf::from("/home/MountainView")));

            remove_var("HOME");

            set_var("USERPROFILE", "/home/MountainView");
            assert_eq!(home_dir(), Some(PathBuf::from("/home/MountainView")));

            set_var("HOME", "/home/MountainView");
            set_var("USERPROFILE", "/home/PaloAlto");
            assert_eq!(home_dir(), Some(PathBuf::from("/home/MountainView")));

            remove_var("HOME");
            remove_var("USERPROFILE");

            if let Some(oldhome) = oldhome { set_var("HOME", oldhome); }
            if let Some(olduserprofile) = olduserprofile { set_var("USERPROFILE", olduserprofile); }
        }
    }
}
<commit_msg>Fix SGX CI<commit_after>use std::env::*;
use std::ffi::{OsStr, OsString};

use rand::distributions::Alphanumeric;
use rand::{thread_rng, Rng};

fn make_rand_name() -> OsString {
    let rng = thread_rng();
    let n = format!("TEST{}", rng.sample_iter(&Alphanumeric).take(10).collect::<String>());
    let n = OsString::from(n);
    assert!(var_os(&n).is_none());
    n
}

fn eq(a: Option<OsString>, b: Option<&str>) {
    assert_eq!(a.as_ref().map(|s| &**s), b.map(OsStr::new).map(|s| &*s));
}

#[test]
fn test_set_var() {
    let n = make_rand_name();
    set_var(&n, "VALUE");
    eq(var_os(&n), Some("VALUE"));
}

#[test]
fn test_remove_var() {
    let n = make_rand_name();
    set_var(&n, "VALUE");
    remove_var(&n);
    eq(var_os(&n), None);
}

#[test]
fn test_set_var_overwrite() {
    let n = make_rand_name();
    set_var(&n, "1");
    set_var(&n, "2");
    eq(var_os(&n), Some("2"));
    set_var(&n, "");
    eq(var_os(&n), Some(""));
}

#[test]
#[cfg_attr(target_os = "emscripten", ignore)]
fn test_var_big() {
    let mut s = "".to_string();
    let mut i = 0;
    while i < 100 {
        s.push_str("aaaaaaaaaa");
        i += 1;
    }
    let n = make_rand_name();
    set_var(&n, &s);
    eq(var_os(&n), Some(&s));
}

#[test]
#[cfg_attr(target_os = "emscripten", ignore)]
fn test_env_set_get_huge() {
    let n = make_rand_name();
    let s = "x".repeat(10000);
    set_var(&n, &s);
    eq(var_os(&n), Some(&s));
    remove_var(&n);
    eq(var_os(&n), None);
}

#[test]
fn test_env_set_var() {
    let n = make_rand_name();

    let mut e = vars_os();
    set_var(&n, "VALUE");
    assert!(!e.any(|(k, v)| { &*k == &*n && &*v == "VALUE" }));

    assert!(vars_os().any(|(k, v)| { &*k == &*n && &*v == "VALUE" }));
}

#[test]
#[cfg_attr(any(target_os = "emscripten", target_env = "sgx"), ignore)]
#[allow(deprecated)]
fn env_home_dir() {
    fn var_to_os_string(var: Result<String, VarError>) -> Option<OsString> {
        match var {
            Ok(var) => Some(OsString::from(var)),
            Err(VarError::NotUnicode(var)) => Some(var),
            _ => None,
        }
    }

    cfg_if::cfg_if! {
        if #[cfg(unix)] {
            use std::path::PathBuf;

            let oldhome = var_to_os_string(var("HOME"));

            set_var("HOME", "/home/MountainView");
            assert_eq!(home_dir(), Some(PathBuf::from("/home/MountainView")));

            remove_var("HOME");
            if cfg!(target_os = "android") {
                assert!(home_dir().is_none());
            } else {
                // When HOME is not set, some platforms return `None`,
                // but others return `Some` with a default.
                // Just check that it is not "/home/MountainView".
                assert_ne!(home_dir(), Some(PathBuf::from("/home/MountainView")));
            }

            if let Some(oldhome) = oldhome { set_var("HOME", oldhome); }
        } else if #[cfg(windows)] {
            use std::path::PathBuf;

            let oldhome = var_to_os_string(var("HOME"));
            let olduserprofile = var_to_os_string(var("USERPROFILE"));

            remove_var("HOME");
            remove_var("USERPROFILE");

            assert!(home_dir().is_some());

            set_var("HOME", "/home/MountainView");
            assert_eq!(home_dir(), Some(PathBuf::from("/home/MountainView")));

            remove_var("HOME");

            set_var("USERPROFILE", "/home/MountainView");
            assert_eq!(home_dir(), Some(PathBuf::from("/home/MountainView")));

            set_var("HOME", "/home/MountainView");
            set_var("USERPROFILE", "/home/PaloAlto");
            assert_eq!(home_dir(), Some(PathBuf::from("/home/MountainView")));

            remove_var("HOME");
            remove_var("USERPROFILE");

            if let Some(oldhome) = oldhome { set_var("HOME", oldhome); }
            if let Some(olduserprofile) = olduserprofile { set_var("USERPROFILE", olduserprofile); }
        }
    }
}
<|endoftext|>


<commit_before>use std::io::{Result, Error, ErrorKind};
use std::net::{ToSocketAddrs, SocketAddr, UdpSocket};
use std::time::Duration;
use url::{UrlParser, SchemeType};
use num;
use rand::{thread_rng, random, Rng};
use packet::{Packet, PacketType, OptionType};

const DEFAULT_RECEIVE_TIMEOUT: u64 = 5;  // 5s

pub struct CoAPClient {
    socket: UdpSocket,
    peer_addr: SocketAddr,
}

impl CoAPClient {
	/// Create a CoAP client with the peer address.
	pub fn new<A: ToSocketAddrs>(addr: A) -> Result<CoAPClient> {
		addr.to_socket_addrs().and_then(|mut iter| {
			match iter.next() {
				Some(SocketAddr::V4(a)) => {
					UdpSocket::bind("0.0.0.0:0").and_then(|s| {
						s.set_read_timeout(Some(Duration::new(DEFAULT_RECEIVE_TIMEOUT, 0))).and_then(|_| {
							Ok(CoAPClient {
								socket: s,
								peer_addr: SocketAddr::V4(a),
							})
						})
					})
				},
				Some(SocketAddr::V6(a)) => {
					UdpSocket::bind(":::0").and_then(|s| {
						s.set_read_timeout(Some(Duration::new(DEFAULT_RECEIVE_TIMEOUT, 0))).and_then(|_| {
							Ok(CoAPClient {
								socket: s,
								peer_addr: SocketAddr::V6(a),
							})
						})
					})
				},
				None => Err(Error::new(ErrorKind::Other, "no address"))
			}
		})
	}

	// Execute a request with the coap url and a specific timeout. Default timeout is 5s.
	pub fn request_with_timeout(url: &str, timeout: Option<Duration>) -> Result<Packet> {
		let mut url_parser = UrlParser::new();
		url_parser.scheme_type_mapper(Self::coap_scheme_type_mapper);

		match url_parser.parse(url) {
			Ok(url_params) => {
				let mut packet = Packet::new();
				packet.header.set_version(1);
				packet.header.set_type(PacketType::Confirmable);
				packet.header.set_code("0.01");

				let message_id = thread_rng().gen_range(0, num::pow(2u32, 16)) as u16;
				packet.header.set_message_id(message_id);

				let mut token: Vec<u8> = vec!(1, 1, 1, 1);
				for x in token.iter_mut() {
				    *x = random()
				}
				packet.set_token(token.clone());

				let domain = match url_params.domain() {
					Some(d) => d,
					None => return Err(Error::new(ErrorKind::InvalidInput, "domain error"))
				};
				let port = url_params.port_or_default().unwrap();

				if let Some(path) = url_params.path() {
					for p in path.iter() {
						packet.add_option(OptionType::UriPath, p.clone().into_bytes().to_vec());
					}
				};

				let client = try!(Self::new((domain, port)));
				try!(client.send(&packet));

				try!(client.set_receive_timeout(timeout));
				match client.receive() {
				 	Ok(receive_packet) => {
				 		if receive_packet.header.get_message_id() == message_id
				 			&& *receive_packet.get_token() == token {
				 				return Ok(receive_packet)
				 			} else {
				 				return Err(Error::new(ErrorKind::Other, "receive invalid data"))
				 			}
				 	},
				 	Err(e) => Err(e)
				}
			},
			Err(_) => Err(Error::new(ErrorKind::InvalidInput, "url error"))
		}
	}

	/// Execute a request with the coap url.
	pub fn request(url: &str) -> Result<Packet> {
		Self::request_with_timeout(url, Some(Duration::new(DEFAULT_RECEIVE_TIMEOUT, 0)))
	}

	/// Response the client with the specifc payload.
	pub fn reply(&self, request_packet: &Packet, payload: Vec<u8>) -> Result<()> {
		let mut packet = Packet::new();

		packet.header.set_version(1);
		let response_type = match request_packet.header.get_type() {
			PacketType::Confirmable => PacketType::Acknowledgement,
			PacketType::NonConfirmable => PacketType::NonConfirmable,
			_ => return Err(Error::new(ErrorKind::InvalidInput, "request type error"))
		};
		packet.header.set_type(response_type);
		packet.header.set_code("2.05");
		packet.header.set_message_id(request_packet.header.get_message_id());
		packet.set_token(request_packet.get_token().clone());
		packet.payload = payload;
		self.send(&packet)
	}

	/// Execute a request.
	pub fn send(&self, packet: &Packet) -> Result<()> {
		match packet.to_bytes() {
			Ok(bytes) => {
				let size = try!(self.socket.send_to(&bytes[..], self.peer_addr));
				if size == bytes.len() {
					Ok(())
				} else {
					Err(Error::new(ErrorKind::Other, "send length error"))
				}
			},
			Err(_) => Err(Error::new(ErrorKind::InvalidInput, "packet error"))
		}
	}

	/// Receive a response.
	pub fn receive(&self) -> Result<Packet> {
		let mut buf = [0; 1500];

		let (nread, _src) = try!(self.socket.recv_from(&mut buf));
		match Packet::from_bytes(&buf[..nread]) {
			Ok(packet) => Ok(packet),
			Err(_) => Err(Error::new(ErrorKind::InvalidInput, "packet error"))
		}
	}

	/// Set the receive timeout.
	pub fn set_receive_timeout(&self, dur: Option<Duration>) -> Result<()> {
	    self.socket.set_read_timeout(dur)
	}

	fn coap_scheme_type_mapper(scheme: &str) -> SchemeType {
		match scheme {
			"coap" => SchemeType::Relative(5683),
			_ => SchemeType::NonRelative,
		}
	}
}

#[cfg(test)]
mod test {
	use super::*;
	use std::time::Duration;
	use std::io::ErrorKind;
	use packet::{Packet, PacketType};
	use server::CoAPServer;

	#[test]
	fn test_request_error_url() {
		assert!(CoAPClient::request("http://127.0.0.1").is_err());
		assert!(CoAPClient::request("coap://127.0.0.").is_err());
		assert!(CoAPClient::request("127.0.0.1").is_err());
	}

	#[test]
	fn test_reply_error() {
		let client = CoAPClient::new("127.0.0.1:5683").unwrap();
		let mut packet = Packet::new();
		packet.header.set_type(PacketType::Acknowledgement);
		assert!(client.reply(&packet, b"Test".to_vec()).is_err());
	}

	fn request_handler(_: Packet, _: CoAPClient) {
	}

	#[test]
	fn test_request_timeout() {
		let mut server = CoAPServer::new("127.0.0.1:5684").unwrap();
		server.handle(request_handler).unwrap();

		let error = CoAPClient::request_with_timeout("coap://127.0.0.1:5684/Rust", Some(Duration::new(1, 0))).unwrap_err();
		assert_eq!(error.kind(), ErrorKind::WouldBlock);
	}
}<commit_msg>missing /<commit_after>use std::io::{Result, Error, ErrorKind};
use std::net::{ToSocketAddrs, SocketAddr, UdpSocket};
use std::time::Duration;
use url::{UrlParser, SchemeType};
use num;
use rand::{thread_rng, random, Rng};
use packet::{Packet, PacketType, OptionType};

const DEFAULT_RECEIVE_TIMEOUT: u64 = 5;  // 5s

pub struct CoAPClient {
    socket: UdpSocket,
    peer_addr: SocketAddr,
}

impl CoAPClient {
	/// Create a CoAP client with the peer address.
	pub fn new<A: ToSocketAddrs>(addr: A) -> Result<CoAPClient> {
		addr.to_socket_addrs().and_then(|mut iter| {
			match iter.next() {
				Some(SocketAddr::V4(a)) => {
					UdpSocket::bind("0.0.0.0:0").and_then(|s| {
						s.set_read_timeout(Some(Duration::new(DEFAULT_RECEIVE_TIMEOUT, 0))).and_then(|_| {
							Ok(CoAPClient {
								socket: s,
								peer_addr: SocketAddr::V4(a),
							})
						})
					})
				},
				Some(SocketAddr::V6(a)) => {
					UdpSocket::bind(":::0").and_then(|s| {
						s.set_read_timeout(Some(Duration::new(DEFAULT_RECEIVE_TIMEOUT, 0))).and_then(|_| {
							Ok(CoAPClient {
								socket: s,
								peer_addr: SocketAddr::V6(a),
							})
						})
					})
				},
				None => Err(Error::new(ErrorKind::Other, "no address"))
			}
		})
	}

	/// Execute a request with the coap url and a specific timeout. Default timeout is 5s.
	pub fn request_with_timeout(url: &str, timeout: Option<Duration>) -> Result<Packet> {
		let mut url_parser = UrlParser::new();
		url_parser.scheme_type_mapper(Self::coap_scheme_type_mapper);

		match url_parser.parse(url) {
			Ok(url_params) => {
				let mut packet = Packet::new();
				packet.header.set_version(1);
				packet.header.set_type(PacketType::Confirmable);
				packet.header.set_code("0.01");

				let message_id = thread_rng().gen_range(0, num::pow(2u32, 16)) as u16;
				packet.header.set_message_id(message_id);

				let mut token: Vec<u8> = vec!(1, 1, 1, 1);
				for x in token.iter_mut() {
				    *x = random()
				}
				packet.set_token(token.clone());

				let domain = match url_params.domain() {
					Some(d) => d,
					None => return Err(Error::new(ErrorKind::InvalidInput, "domain error"))
				};
				let port = url_params.port_or_default().unwrap();

				if let Some(path) = url_params.path() {
					for p in path.iter() {
						packet.add_option(OptionType::UriPath, p.clone().into_bytes().to_vec());
					}
				};

				let client = try!(Self::new((domain, port)));
				try!(client.send(&packet));

				try!(client.set_receive_timeout(timeout));
				match client.receive() {
				 	Ok(receive_packet) => {
				 		if receive_packet.header.get_message_id() == message_id
				 			&& *receive_packet.get_token() == token {
				 				return Ok(receive_packet)
				 			} else {
				 				return Err(Error::new(ErrorKind::Other, "receive invalid data"))
				 			}
				 	},
				 	Err(e) => Err(e)
				}
			},
			Err(_) => Err(Error::new(ErrorKind::InvalidInput, "url error"))
		}
	}

	/// Execute a request with the coap url.
	pub fn request(url: &str) -> Result<Packet> {
		Self::request_with_timeout(url, Some(Duration::new(DEFAULT_RECEIVE_TIMEOUT, 0)))
	}

	/// Response the client with the specifc payload.
	pub fn reply(&self, request_packet: &Packet, payload: Vec<u8>) -> Result<()> {
		let mut packet = Packet::new();

		packet.header.set_version(1);
		let response_type = match request_packet.header.get_type() {
			PacketType::Confirmable => PacketType::Acknowledgement,
			PacketType::NonConfirmable => PacketType::NonConfirmable,
			_ => return Err(Error::new(ErrorKind::InvalidInput, "request type error"))
		};
		packet.header.set_type(response_type);
		packet.header.set_code("2.05");
		packet.header.set_message_id(request_packet.header.get_message_id());
		packet.set_token(request_packet.get_token().clone());
		packet.payload = payload;
		self.send(&packet)
	}

	/// Execute a request.
	pub fn send(&self, packet: &Packet) -> Result<()> {
		match packet.to_bytes() {
			Ok(bytes) => {
				let size = try!(self.socket.send_to(&bytes[..], self.peer_addr));
				if size == bytes.len() {
					Ok(())
				} else {
					Err(Error::new(ErrorKind::Other, "send length error"))
				}
			},
			Err(_) => Err(Error::new(ErrorKind::InvalidInput, "packet error"))
		}
	}

	/// Receive a response.
	pub fn receive(&self) -> Result<Packet> {
		let mut buf = [0; 1500];

		let (nread, _src) = try!(self.socket.recv_from(&mut buf));
		match Packet::from_bytes(&buf[..nread]) {
			Ok(packet) => Ok(packet),
			Err(_) => Err(Error::new(ErrorKind::InvalidInput, "packet error"))
		}
	}

	/// Set the receive timeout.
	pub fn set_receive_timeout(&self, dur: Option<Duration>) -> Result<()> {
	    self.socket.set_read_timeout(dur)
	}

	fn coap_scheme_type_mapper(scheme: &str) -> SchemeType {
		match scheme {
			"coap" => SchemeType::Relative(5683),
			_ => SchemeType::NonRelative,
		}
	}
}

#[cfg(test)]
mod test {
	use super::*;
	use std::time::Duration;
	use std::io::ErrorKind;
	use packet::{Packet, PacketType};
	use server::CoAPServer;

	#[test]
	fn test_request_error_url() {
		assert!(CoAPClient::request("http://127.0.0.1").is_err());
		assert!(CoAPClient::request("coap://127.0.0.").is_err());
		assert!(CoAPClient::request("127.0.0.1").is_err());
	}

	#[test]
	fn test_reply_error() {
		let client = CoAPClient::new("127.0.0.1:5683").unwrap();
		let mut packet = Packet::new();
		packet.header.set_type(PacketType::Acknowledgement);
		assert!(client.reply(&packet, b"Test".to_vec()).is_err());
	}

	fn request_handler(_: Packet, _: CoAPClient) {
	}

	#[test]
	fn test_request_timeout() {
		let mut server = CoAPServer::new("127.0.0.1:5684").unwrap();
		server.handle(request_handler).unwrap();

		let error = CoAPClient::request_with_timeout("coap://127.0.0.1:5684/Rust", Some(Duration::new(1, 0))).unwrap_err();
		assert_eq!(error.kind(), ErrorKind::WouldBlock);
	}
}<|endoftext|>


<commit_before>use std::net::{TcpStream, ToSocketAddrs};
use openssl::ssl::{SslContext, SslStream};
use std::io::{Error, ErrorKind, Read, Result, Write};
use regex::Regex;

static TAG_PREFIX: &'static str = "a";
const INITIAL_TAG: i32 = 1;

/// Stream to interface with the IMAP server. This interface is only for the command stream.
pub struct IMAPStream<T> {
	stream: T,
	tag: u32
}

pub struct IMAPMailbox {
	pub flags: String,
	pub exists: u32,
	pub recent: u32,
	pub unseen: Option<u32>,
	pub permanent_flags: Option<String>,
	pub uid_next: Option<u32>,
	pub uid_validity: Option<u32>
}

impl IMAPStream<TcpStream> {
	///
	pub fn connect<A: ToSocketAddrs>(addr: A) -> Result<IMAPStream<TcpStream>> {
		match TcpStream::connect(addr) {
			Ok(stream) => {
				let mut socket = IMAPStream { stream: stream, tag: INITIAL_TAG};

				try!(socket.read_greeting());
				Ok(socket)
			},
			Err(e) => Err(e)
		}
	}
}

impl IMAPStream<SslStream<TcpStream>> {
	///
	pub fn secure_connect<A: ToSocketAddrs>(addr: A, ssl_context: SslContext) -> Result<IMAPStream<SslStream<TcpStream>>> {
		match TcpStream::connect(addr) {
			Ok(stream) => {
				let mut socket = IMAPStream { stream: SslStream::connect(&ssl_context, stream).unwrap(), tag: INITIAL_TAG};

				try!(socket.read_greeting());
				Ok(socket)
			},
			Err(e) => Err(e)
		}
	}
}

impl<T: Read+Write> IMAPStream<T> {

	/// Log in to the IMAP server.
	pub fn login(&mut self, username: & str, password: & str) -> Result<()> {
		self.run_command_and_check_ok(&format!("LOGIN {} {}", username, password).to_string())
	}

	/// Selects a mailbox
	pub fn select(&mut self, mailbox_name: &str) -> Result<IMAPMailbox> {
		match self.run_command(&format!("SELECT {}", mailbox_name).to_string()) {
			Ok(lines) => self.parse_select_or_examine(lines),
			Err(e) => Err(e)
		}
	}

	fn parse_select_or_examine(&mut self, lines: Vec<String>) -> Result<IMAPMailbox> {
		let exists_regex = match Regex::new(r"^\* (\d+) EXISTS\r\n") {
    		Ok(re) => re,
    		Err(err) => panic!("{}", err),
		};

		let recent_regex = match Regex::new(r"^\* (\d+) RECENT\r\n") {
    		Ok(re) => re,
    		Err(err) => panic!("{}", err),
		};

		let flags_regex = match Regex::new(r"^\* FLAGS (.+)\r\n") {
    		Ok(re) => re,
    		Err(err) => panic!("{}", err),
		};

		let unseen_regex = match Regex::new(r"^OK \[UNSEEN (\d+)\](.*)\r\n") {
    		Ok(re) => re,
    		Err(err) => panic!("{}", err),
		};

		let uid_validity_regex = match Regex::new(r"^OK \[UIDVALIDITY (\d+)\](.*)\r\n") {
    		Ok(re) => re,
    		Err(err) => panic!("{}", err),
		};

		let uid_next_regex = match Regex::new(r"^OK \[UIDNEXT (\d+)\](.*)\r\n") {
    		Ok(re) => re,
    		Err(err) => panic!("{}", err),
		};

		let permanent_flags_regex = match Regex::new(r"^OK \[PERMANENTFLAGS (.+)\]\r\n") {
    		Ok(re) => re,
    		Err(err) => panic!("{}", err),
		};

		//Check Ok
		match self.parse_response_ok(lines.clone()) {
			Ok(_) => (),
			Err(e) => return Err(e)
		};

		let mut mailbox = IMAPMailbox{
			flags: "".to_string(),
			exists: 0,
			recent: 0,
			unseen: None,
			permanent_flags: None,
			uid_next: None,
			uid_validity: None
		};

		for line in lines.iter() {
			if exists_regex.is_match(line) {
				let cap = exists_regex.captures(line).unwrap();
				mailbox.exists = cap.at(1).unwrap().parse::<u32>().unwrap();
			} else if recent_regex.is_match(line) {
				let cap = recent_regex.captures(line).unwrap();
				mailbox.recent = cap.at(1).unwrap().parse::<u32>().unwrap();
			} else if flags_regex.is_match(line) {
				let cap = flags_regex.captures(line).unwrap();
				mailbox.flags = cap.at(1).unwrap().to_string();
			} else if unseen_regex.is_match(line) {
				let cap = unseen_regex.captures(line).unwrap();
				mailbox.unseen = Some(cap.at(1).unwrap().parse::<u32>().unwrap());
			} else if uid_validity_regex.is_match(line) {
				let cap = uid_validity_regex.captures(line).unwrap();
				mailbox.uid_validity = Some(cap.at(1).unwrap().parse::<u32>().unwrap());
			} else if uid_next_regex.is_match(line) {
				let cap = uid_next_regex.captures(line).unwrap();
				mailbox.uid_next = Some(cap.at(1).unwrap().parse::<u32>().unwrap());
			} else if permanent_flags_regex.is_match(line) {
				let cap = permanent_flags_regex.captures(line).unwrap();
				mailbox.permanent_flags = Some(cap.at(1).unwrap().to_string());
			}
		}

		return Ok(mailbox);
	}

	/// Examine is identical to Select, but the selected mailbox is identified as read-only
	pub fn examine(&mut self, mailbox_name: &str) -> Result<IMAPMailbox> {
		match self.run_command(&format!("EXAMINE {}", mailbox_name).to_string()) {
			Ok(lines) => self.parse_select_or_examine(lines),
			Err(e) => Err(e)
		}
	}

	/// Fetch retreives data associated with a message in the mailbox.
	pub fn fetch(&mut self, sequence_set: &str, query: &str) -> Result<Vec<String>> {
		self.run_command(&format!("FETCH {} {}", sequence_set, query).to_string())
	}

	/// Noop always succeeds, and it does nothing.
	pub fn noop(&mut self) -> Result<()> {
		self.run_command_and_check_ok("NOOP")
	}

	/// Logout informs the server that the client is done with the connection.
	pub fn logout(&mut self) -> Result<()> {
		self.run_command_and_check_ok("LOGOUT")
	}

	/// Create creates a mailbox with the given name.
	pub fn create(&mut self, mailbox_name: &str) -> Result<()> {
		self.run_command_and_check_ok(&format!("CREATE {}", mailbox_name).to_string())
	}

	/// Delete permanently removes the mailbox with the given name.
	pub fn delete(&mut self, mailbox_name: &str) -> Result<()> {
		self.run_command_and_check_ok(&format!("DELETE {}", mailbox_name).to_string())
	}

	/// Rename changes the name of a mailbox.
	pub fn rename(&mut self, current_mailbox_name: &str, new_mailbox_name: &str) -> Result<()> {
		self.run_command_and_check_ok(&format!("RENAME {} {}", current_mailbox_name, new_mailbox_name).to_string())
	}

	/// Subscribe adds the specified mailbox name to the server's set of "active" or "subscribed"
	/// mailboxes as returned by the LSUB command.
	pub fn subscribe(&mut self, mailbox: &str) -> Result<()> {
		self.run_command_and_check_ok(&format!("SUBSCRIBE {}", mailbox).to_string())
	}

	/// Unsubscribe removes the specified mailbox name from the server's set of "active" or "subscribed"
	/// mailboxes as returned by the LSUB command.
	pub fn unsubscribe(&mut self, mailbox: &str) -> Result<()> {
		self.run_command_and_check_ok(&format!("UNSUBSCRIBE {}", mailbox).to_string())
	}

	/// Capability requests a listing of capabilities that the server supports.
	pub fn capability(&mut self) -> Result<Vec<String>> {
		match self.run_command(&format!("CAPABILITY").to_string()) {
			Ok(lines) => self.parse_capability(lines),
			Err(e) => Err(e)
		}
	}

	fn parse_capability(&mut self, lines: Vec<String>) -> Result<Vec<String>> {
		let capability_regex = match Regex::new(r"^\* CAPABILITY (.*)\r\n") {
    		Ok(re) => re,
    		Err(err) => panic!("{}", err),
		};

		//Check Ok
		match self.parse_response_ok(lines.clone()) {
			Ok(_) => (),
			Err(e) => return Err(e)
		};

		for line in lines.iter() {
			if capability_regex.is_match(line) {
				let cap = capability_regex.captures(line).unwrap();
				let capabilities_str = cap.at(1).unwrap();
				return Ok(capabilities_str.split(' ').map(|x| x.to_string()).collect());
			}
		}

		Err(Error::new(ErrorKind::Other, "Error parsing capabilities response"))
	}

	/// Expunge permanently removes all messages that have the \Deleted flag set from the currently
	/// selected mailbox.
	pub fn expunge(&mut self) -> Result<()> {
		self.run_command_and_check_ok("CHECK")
	}

	/// Check requests a checkpoint of the currently selected mailbox.
	pub fn check(&mut self) -> Result<()> {
		self.run_command_and_check_ok("CHECK")
	}

	/// Close permanently removes all messages that have the \Deleted flag set from the currently
	/// selected mailbox, and returns to the authenticated state from the selected state.
	pub fn close(&mut self) -> Result<()> {
		self.run_command_and_check_ok("CLOSE")
	}

	/// Copy copies the specified message to the end of the specified destination mailbox.
	pub fn copy(&mut self, sequence_set: &str, mailbox_name: &str) -> Result<()> {
		self.run_command_and_check_ok(&format!("COPY {} {}", sequence_set, mailbox_name).to_string())
	}

	pub fn run_command_and_check_ok(&mut self, command: &str) -> Result<()> {
		match self.run_command(command) {
			Ok(lines) => self.parse_response_ok(lines),
			Err(e) => Err(e)
		}
	}

	pub fn run_command(&mut self, untagged_command: &str) -> Result<Vec<String>> {
		let command = self.create_command(untagged_command.to_string());

		match self.stream.write_fmt(format_args!("{}", &*command)) {
			Ok(_) => (),
			Err(_) => return Err(Error::new(ErrorKind::Other, "Failed to write")),
		};

		let ret = match self.read_response() {
			Ok(lines) => Ok(lines),
			Err(_) => Err(Error::new(ErrorKind::Other, "Failed to read")),
		};

		self.tag += 1;

		return ret;
	}

	fn parse_response_ok(&mut self, lines: Vec<String>) -> Result<()> {
		let ok_regex = match Regex::new(r"^([a-zA-Z0-9]+) ([a-zA-Z0-9]+)(.*)") {
    		Ok(re) => re,
    		Err(err) => panic!("{}", err),
		};
		let last_line = lines.last().unwrap();

		for cap in ok_regex.captures_iter(last_line) {
			let response_type = cap.at(2).unwrap_or("");
			if response_type == "OK" {
				return Ok(());
			}
		}

		return Err(Error::new(ErrorKind::Other, format!("Invalid Response: {}", last_line).to_string()));
	}

	fn read_response(&mut self) -> Result<Vec<String>> {
		//Carriage return
		let cr = 0x0d;
		//Line Feed
		let lf = 0x0a;
		let mut found_tag_line = false;
		let start_str = format!("{}{} ", TAG_PREFIX, self.tag);
		let mut lines: Vec<String> = Vec::new();

		while !found_tag_line {
			let mut line_buffer: Vec<u8> = Vec::new();
			while line_buffer.len() < 2 || (line_buffer[line_buffer.len()-1] != lf && line_buffer[line_buffer.len()-2] != cr) {
					let byte_buffer: &mut [u8] = &mut [0];
					match self.stream.read(byte_buffer) {
						Ok(_) => {},
						Err(_) => return Err(Error::new(ErrorKind::Other, "Failed to read the response")),
					}
					line_buffer.push(byte_buffer[0]);
			}

			let line = String::from_utf8(line_buffer).unwrap();

			lines.push(line.clone());

			if (&*line).starts_with(&*start_str) {
				found_tag_line = true;
			}
		}

		Ok(lines)
	}

	fn read_greeting(&mut self) -> Result<()> {
		//Carriage return
		let cr = 0x0d;
		//Line Feed
		let lf = 0x0a;

		let mut line_buffer: Vec<u8> = Vec::new();
		while line_buffer.len() < 2 || (line_buffer[line_buffer.len()-1] != lf && line_buffer[line_buffer.len()-2] != cr) {
				let byte_buffer: &mut [u8] = &mut [0];
				match self.stream.read(byte_buffer) {
					Ok(_) => {},
					Err(_) => return Err(Error::new(ErrorKind::Other, "Failed to read the response")),
				}
				line_buffer.push(byte_buffer[0]);
		}

		Ok(())
	}

	fn create_command(&mut self, command: String) -> String {
		let command = format!("{}{} {}\r\n", TAG_PREFIX, self.tag, command);
		return command;
	}
}
<commit_msg>Fix initial tag tyype and remove regex creation checks<commit_after>use std::net::{TcpStream, ToSocketAddrs};
use openssl::ssl::{SslContext, SslStream};
use std::io::{Error, ErrorKind, Read, Result, Write};
use regex::Regex;

static TAG_PREFIX: &'static str = "a";
const INITIAL_TAG: u32 = 1;

/// Stream to interface with the IMAP server. This interface is only for the command stream.
pub struct IMAPStream<T> {
	stream: T,
	tag: u32
}

pub struct IMAPMailbox {
	pub flags: String,
	pub exists: u32,
	pub recent: u32,
	pub unseen: Option<u32>,
	pub permanent_flags: Option<String>,
	pub uid_next: Option<u32>,
	pub uid_validity: Option<u32>
}

impl IMAPStream<TcpStream> {
	///
	pub fn connect<A: ToSocketAddrs>(addr: A) -> Result<IMAPStream<TcpStream>> {
		match TcpStream::connect(addr) {
			Ok(stream) => {
				let mut socket = IMAPStream { stream: stream, tag: INITIAL_TAG};

				try!(socket.read_greeting());
				Ok(socket)
			},
			Err(e) => Err(e)
		}
	}
}

impl IMAPStream<SslStream<TcpStream>> {
	///
	pub fn secure_connect<A: ToSocketAddrs>(addr: A, ssl_context: SslContext) -> Result<IMAPStream<SslStream<TcpStream>>> {
		match TcpStream::connect(addr) {
			Ok(stream) => {
				let mut socket = IMAPStream { stream: SslStream::connect(&ssl_context, stream).unwrap(), tag: INITIAL_TAG};

				try!(socket.read_greeting());
				Ok(socket)
			},
			Err(e) => Err(e)
		}
	}
}

impl<T: Read+Write> IMAPStream<T> {

	/// Log in to the IMAP server.
	pub fn login(&mut self, username: & str, password: & str) -> Result<()> {
		self.run_command_and_check_ok(&format!("LOGIN {} {}", username, password).to_string())
	}

	/// Selects a mailbox
	pub fn select(&mut self, mailbox_name: &str) -> Result<IMAPMailbox> {
		match self.run_command(&format!("SELECT {}", mailbox_name).to_string()) {
			Ok(lines) => self.parse_select_or_examine(lines),
			Err(e) => Err(e)
		}
	}

	fn parse_select_or_examine(&mut self, lines: Vec<String>) -> Result<IMAPMailbox> {
		let exists_regex = Regex::new(r"^\* (\d+) EXISTS\r\n").unwrap();

		let recent_regex = Regex::new(r"^\* (\d+) RECENT\r\n").unwrap();

		let flags_regex = Regex::new(r"^\* FLAGS (.+)\r\n").unwrap();

		let unseen_regex = Regex::new(r"^OK \[UNSEEN (\d+)\](.*)\r\n").unwrap();

		let uid_validity_regex = Regex::new(r"^OK \[UIDVALIDITY (\d+)\](.*)\r\n").unwrap();

		let uid_next_regex = Regex::new(r"^OK \[UIDNEXT (\d+)\](.*)\r\n").unwrap();

		let permanent_flags_regex = Regex::new(r"^OK \[PERMANENTFLAGS (.+)\]\r\n").unwrap();

		//Check Ok
		match self.parse_response_ok(lines.clone()) {
			Ok(_) => (),
			Err(e) => return Err(e)
		};

		let mut mailbox = IMAPMailbox{
			flags: "".to_string(),
			exists: 0,
			recent: 0,
			unseen: None,
			permanent_flags: None,
			uid_next: None,
			uid_validity: None
		};

		for line in lines.iter() {
			if exists_regex.is_match(line) {
				let cap = exists_regex.captures(line).unwrap();
				mailbox.exists = cap.at(1).unwrap().parse::<u32>().unwrap();
			} else if recent_regex.is_match(line) {
				let cap = recent_regex.captures(line).unwrap();
				mailbox.recent = cap.at(1).unwrap().parse::<u32>().unwrap();
			} else if flags_regex.is_match(line) {
				let cap = flags_regex.captures(line).unwrap();
				mailbox.flags = cap.at(1).unwrap().to_string();
			} else if unseen_regex.is_match(line) {
				let cap = unseen_regex.captures(line).unwrap();
				mailbox.unseen = Some(cap.at(1).unwrap().parse::<u32>().unwrap());
			} else if uid_validity_regex.is_match(line) {
				let cap = uid_validity_regex.captures(line).unwrap();
				mailbox.uid_validity = Some(cap.at(1).unwrap().parse::<u32>().unwrap());
			} else if uid_next_regex.is_match(line) {
				let cap = uid_next_regex.captures(line).unwrap();
				mailbox.uid_next = Some(cap.at(1).unwrap().parse::<u32>().unwrap());
			} else if permanent_flags_regex.is_match(line) {
				let cap = permanent_flags_regex.captures(line).unwrap();
				mailbox.permanent_flags = Some(cap.at(1).unwrap().to_string());
			}
		}

		return Ok(mailbox);
	}

	/// Examine is identical to Select, but the selected mailbox is identified as read-only
	pub fn examine(&mut self, mailbox_name: &str) -> Result<IMAPMailbox> {
		match self.run_command(&format!("EXAMINE {}", mailbox_name).to_string()) {
			Ok(lines) => self.parse_select_or_examine(lines),
			Err(e) => Err(e)
		}
	}

	/// Fetch retreives data associated with a message in the mailbox.
	pub fn fetch(&mut self, sequence_set: &str, query: &str) -> Result<Vec<String>> {
		self.run_command(&format!("FETCH {} {}", sequence_set, query).to_string())
	}

	/// Noop always succeeds, and it does nothing.
	pub fn noop(&mut self) -> Result<()> {
		self.run_command_and_check_ok("NOOP")
	}

	/// Logout informs the server that the client is done with the connection.
	pub fn logout(&mut self) -> Result<()> {
		self.run_command_and_check_ok("LOGOUT")
	}

	/// Create creates a mailbox with the given name.
	pub fn create(&mut self, mailbox_name: &str) -> Result<()> {
		self.run_command_and_check_ok(&format!("CREATE {}", mailbox_name).to_string())
	}

	/// Delete permanently removes the mailbox with the given name.
	pub fn delete(&mut self, mailbox_name: &str) -> Result<()> {
		self.run_command_and_check_ok(&format!("DELETE {}", mailbox_name).to_string())
	}

	/// Rename changes the name of a mailbox.
	pub fn rename(&mut self, current_mailbox_name: &str, new_mailbox_name: &str) -> Result<()> {
		self.run_command_and_check_ok(&format!("RENAME {} {}", current_mailbox_name, new_mailbox_name).to_string())
	}

	/// Subscribe adds the specified mailbox name to the server's set of "active" or "subscribed"
	/// mailboxes as returned by the LSUB command.
	pub fn subscribe(&mut self, mailbox: &str) -> Result<()> {
		self.run_command_and_check_ok(&format!("SUBSCRIBE {}", mailbox).to_string())
	}

	/// Unsubscribe removes the specified mailbox name from the server's set of "active" or "subscribed"
	/// mailboxes as returned by the LSUB command.
	pub fn unsubscribe(&mut self, mailbox: &str) -> Result<()> {
		self.run_command_and_check_ok(&format!("UNSUBSCRIBE {}", mailbox).to_string())
	}

	/// Capability requests a listing of capabilities that the server supports.
	pub fn capability(&mut self) -> Result<Vec<String>> {
		match self.run_command(&format!("CAPABILITY").to_string()) {
			Ok(lines) => self.parse_capability(lines),
			Err(e) => Err(e)
		}
	}

	fn parse_capability(&mut self, lines: Vec<String>) -> Result<Vec<String>> {
		let capability_regex = match Regex::new(r"^\* CAPABILITY (.*)\r\n") {
    		Ok(re) => re,
    		Err(err) => panic!("{}", err),
		};

		//Check Ok
		match self.parse_response_ok(lines.clone()) {
			Ok(_) => (),
			Err(e) => return Err(e)
		};

		for line in lines.iter() {
			if capability_regex.is_match(line) {
				let cap = capability_regex.captures(line).unwrap();
				let capabilities_str = cap.at(1).unwrap();
				return Ok(capabilities_str.split(' ').map(|x| x.to_string()).collect());
			}
		}

		Err(Error::new(ErrorKind::Other, "Error parsing capabilities response"))
	}

	/// Expunge permanently removes all messages that have the \Deleted flag set from the currently
	/// selected mailbox.
	pub fn expunge(&mut self) -> Result<()> {
		self.run_command_and_check_ok("CHECK")
	}

	/// Check requests a checkpoint of the currently selected mailbox.
	pub fn check(&mut self) -> Result<()> {
		self.run_command_and_check_ok("CHECK")
	}

	/// Close permanently removes all messages that have the \Deleted flag set from the currently
	/// selected mailbox, and returns to the authenticated state from the selected state.
	pub fn close(&mut self) -> Result<()> {
		self.run_command_and_check_ok("CLOSE")
	}

	/// Copy copies the specified message to the end of the specified destination mailbox.
	pub fn copy(&mut self, sequence_set: &str, mailbox_name: &str) -> Result<()> {
		self.run_command_and_check_ok(&format!("COPY {} {}", sequence_set, mailbox_name).to_string())
	}

	pub fn run_command_and_check_ok(&mut self, command: &str) -> Result<()> {
		match self.run_command(command) {
			Ok(lines) => self.parse_response_ok(lines),
			Err(e) => Err(e)
		}
	}

	pub fn run_command(&mut self, untagged_command: &str) -> Result<Vec<String>> {
		let command = self.create_command(untagged_command.to_string());

		match self.stream.write_fmt(format_args!("{}", &*command)) {
			Ok(_) => (),
			Err(_) => return Err(Error::new(ErrorKind::Other, "Failed to write")),
		};

		let ret = match self.read_response() {
			Ok(lines) => Ok(lines),
			Err(_) => Err(Error::new(ErrorKind::Other, "Failed to read")),
		};

		self.tag += 1;

		return ret;
	}

	fn parse_response_ok(&mut self, lines: Vec<String>) -> Result<()> {
		let ok_regex = match Regex::new(r"^([a-zA-Z0-9]+) ([a-zA-Z0-9]+)(.*)") {
    		Ok(re) => re,
    		Err(err) => panic!("{}", err),
		};
		let last_line = lines.last().unwrap();

		for cap in ok_regex.captures_iter(last_line) {
			let response_type = cap.at(2).unwrap_or("");
			if response_type == "OK" {
				return Ok(());
			}
		}

		return Err(Error::new(ErrorKind::Other, format!("Invalid Response: {}", last_line).to_string()));
	}

	fn read_response(&mut self) -> Result<Vec<String>> {
		//Carriage return
		let cr = 0x0d;
		//Line Feed
		let lf = 0x0a;
		let mut found_tag_line = false;
		let start_str = format!("{}{} ", TAG_PREFIX, self.tag);
		let mut lines: Vec<String> = Vec::new();

		while !found_tag_line {
			let mut line_buffer: Vec<u8> = Vec::new();
			while line_buffer.len() < 2 || (line_buffer[line_buffer.len()-1] != lf && line_buffer[line_buffer.len()-2] != cr) {
					let byte_buffer: &mut [u8] = &mut [0];
					match self.stream.read(byte_buffer) {
						Ok(_) => {},
						Err(_) => return Err(Error::new(ErrorKind::Other, "Failed to read the response")),
					}
					line_buffer.push(byte_buffer[0]);
			}

			let line = String::from_utf8(line_buffer).unwrap();

			lines.push(line.clone());

			if (&*line).starts_with(&*start_str) {
				found_tag_line = true;
			}
		}

		Ok(lines)
	}

	fn read_greeting(&mut self) -> Result<()> {
		//Carriage return
		let cr = 0x0d;
		//Line Feed
		let lf = 0x0a;

		let mut line_buffer: Vec<u8> = Vec::new();
		while line_buffer.len() < 2 || (line_buffer[line_buffer.len()-1] != lf && line_buffer[line_buffer.len()-2] != cr) {
				let byte_buffer: &mut [u8] = &mut [0];
				match self.stream.read(byte_buffer) {
					Ok(_) => {},
					Err(_) => return Err(Error::new(ErrorKind::Other, "Failed to read the response")),
				}
				line_buffer.push(byte_buffer[0]);
		}

		Ok(())
	}

	fn create_command(&mut self, command: String) -> String {
		let command = format!("{}{} {}\r\n", TAG_PREFIX, self.tag, command);
		return command;
	}
}
<|endoftext|>


<commit_before>use std::iter::Iterator;
use std::ptr;
use std::thread;

use mongo_c_driver_wrapper::bindings;
use bson::{Bson,Document,oid};

use super::BsoncError;
use super::bsonc;
use super::client::Client;
use super::database::Database;
use super::flags::QueryFlag;
use super::collection::{Collection,TailOptions};
use super::CommandAndFindOptions;

use super::Result;

pub enum CreatedBy<'a> {
    Client(&'a Client<'a>),
    Database(&'a Database<'a>),
    Collection(&'a Collection<'a>)
}

pub struct Cursor<'a> {
    _created_by:       CreatedBy<'a>,
    inner:             *mut bindings::mongoc_cursor_t,
    tailing:           bool,
    tail_wait_time_ms: u32,
    // Become owner of bsonc because the cursor needs it
    // to be allocated for it's entire lifetime
    _fields:           Option<bsonc::Bsonc>
}

impl<'a> Cursor<'a> {
    pub fn new(
        created_by: CreatedBy<'a>,
        inner:      *mut bindings::mongoc_cursor_t,
        fields:     Option<bsonc::Bsonc>
    ) -> Cursor<'a> {
        assert!(!inner.is_null());
        Cursor {
            _created_by:       created_by,
            inner:             inner,
            tailing:           false,
            tail_wait_time_ms: 0,
            _fields:           fields
        }
    }

    fn is_alive(&self) -> bool {
        assert!(!self.inner.is_null());
        unsafe {
            bindings::mongoc_cursor_is_alive(self.inner) == 1
        }
    }

    fn more(&self) -> bool {
        assert!(!self.inner.is_null());
        unsafe {
            bindings::mongoc_cursor_more(self.inner) == 1
        }
    }

    fn error(&self) -> BsoncError {
        assert!(!self.inner.is_null());
        let mut error = BsoncError::empty();
        unsafe {
            bindings::mongoc_cursor_error(
                self.inner,
                error.mut_inner()
            )
        };
        error
    }
}

impl<'a> Iterator for Cursor<'a> {
    type Item = Result<Document>;

    fn next(&mut self) -> Option<Self::Item> {
        assert!(!self.inner.is_null());

        loop {
            if !self.more() {
                return None
            }

            // The C driver writes the document to memory and sets an
            // already existing pointer to it.
            let mut bson_ptr: *const bindings::bson_t = ptr::null();
            let success = unsafe {
                bindings::mongoc_cursor_next(
                    self.inner,
                    &mut bson_ptr
                )
            };

            // Fetch error that might have occurred while getting
            // the next cursor.
            let error = self.error();

            if success == 0 {
                if error.is_empty() {
                    if self.tailing && self.is_alive() {
                        // Since there was no error, this is a tailing cursor
                        // and the cursor is alive we'll wait before trying again.
                        thread::sleep_ms(self.tail_wait_time_ms);
                        continue;
                    } else {
                        // No result, no error and cursor not tailing so we must
                        // be at the end.
                        return None
                    }
                } else {
                    // There was an error
                    return Some(Err(error.into()))
                }
            }
            assert!(!bson_ptr.is_null());

            // Parse and return bson document.
            let bsonc = bsonc::Bsonc::from_ptr(bson_ptr);
            match bsonc.as_document() {
                Ok(document) => return Some(Ok(document)),
                Err(error)   => return Some(Err(error.into()))
            }
        }
    }
}

impl<'a> Drop for Cursor<'a> {
    fn drop(&mut self) {
        assert!(!self.inner.is_null());
        unsafe {
            bindings::mongoc_cursor_destroy(self.inner);
        }
    }
}

/// Cursor that will reconnect and resume tailing a collection
/// at the right point if the connection fails.
pub struct TailingCursor<'a> {
    collection:   &'a Collection<'a>,
    query:        Document,
    find_options: CommandAndFindOptions,
    tail_options: TailOptions,
    cursor:       Option<Cursor<'a>>,
    last_seen_id: Option<oid::ObjectId>,
    retry_count:  u32
}

impl<'a> TailingCursor<'a> {
    pub fn new(
        collection:   &'a Collection<'a>,
        query:        Document,
        find_options: CommandAndFindOptions,
        tail_options: TailOptions
    ) -> TailingCursor<'a> {
        // Add flags to make query tailable
        let mut find_options = find_options;
        find_options.query_flags.add(QueryFlag::TailableCursor);
        find_options.query_flags.add(QueryFlag::AwaitData);

        TailingCursor {
            collection:   collection,
            query:        query,
            find_options: find_options,
            tail_options: tail_options,
            cursor:       None,
            last_seen_id: None,
            retry_count:  0
        }
    }
}

impl<'a> Iterator for TailingCursor<'a> {
    type Item = Result<Document>;

    fn next(&mut self) -> Option<Self::Item> {
        loop {
            // Start a scope so we're free to set the cursor to None at the end.
            {
                if self.cursor.is_none() {
                    // Add the last seen id to the query if it's present.
                    match self.last_seen_id.take() {
                        Some(id) => {
                            self.query.insert("_id".to_string(), Bson::Document(doc! { "$gt" => id }));
                        },
                        None => ()
                    };

                    // Set the cursor
                    self.cursor = match self.collection.find(&self.query, Some(&self.find_options)) {
                        Ok(mut c)  => {
                            c.tailing           = true;
                            c.tail_wait_time_ms = self.tail_options.wait_time_ms;
                            Some(c)
                        },
                        Err(e) => return Some(Err(e.into()))
                    };
                }

                let cursor = match self.cursor {
                    Some(ref mut c) => c,
                    None => panic!("It should be impossible to not have a cursor here")
                };

                match cursor.next() {
                    Some(next_result) => {
                        match next_result {
                            Ok(next) => {
                                // This was successfull, so reset retry count and return result.
                                self.retry_count = 0;
                                return Some(Ok(next))
                            },
                            Err(e) => {
                                // Retry if we haven't exceeded the maximum number of retries.
                                if self.retry_count >= self.tail_options.max_retries {
                                    return Some(Err(e.into()))
                                }
                            }
                        }
                    },
                    None => ()
                };
            }

            // We made it to the end, so we weren't able to get the next item from
            // the cursor. We need to reconnect in the next iteration of the loop.
            self.retry_count += 1;
            self.cursor      = None;
        }
    }
}
<commit_msg>Use insert_bson in cursor<commit_after>use std::iter::Iterator;
use std::ptr;
use std::thread;

use mongo_c_driver_wrapper::bindings;
use bson::{Bson,Document,oid};

use super::BsoncError;
use super::bsonc;
use super::client::Client;
use super::database::Database;
use super::flags::QueryFlag;
use super::collection::{Collection,TailOptions};
use super::CommandAndFindOptions;

use super::Result;

pub enum CreatedBy<'a> {
    Client(&'a Client<'a>),
    Database(&'a Database<'a>),
    Collection(&'a Collection<'a>)
}

pub struct Cursor<'a> {
    _created_by:       CreatedBy<'a>,
    inner:             *mut bindings::mongoc_cursor_t,
    tailing:           bool,
    tail_wait_time_ms: u32,
    // Become owner of bsonc because the cursor needs it
    // to be allocated for it's entire lifetime
    _fields:           Option<bsonc::Bsonc>
}

impl<'a> Cursor<'a> {
    pub fn new(
        created_by: CreatedBy<'a>,
        inner:      *mut bindings::mongoc_cursor_t,
        fields:     Option<bsonc::Bsonc>
    ) -> Cursor<'a> {
        assert!(!inner.is_null());
        Cursor {
            _created_by:       created_by,
            inner:             inner,
            tailing:           false,
            tail_wait_time_ms: 0,
            _fields:           fields
        }
    }

    fn is_alive(&self) -> bool {
        assert!(!self.inner.is_null());
        unsafe {
            bindings::mongoc_cursor_is_alive(self.inner) == 1
        }
    }

    fn more(&self) -> bool {
        assert!(!self.inner.is_null());
        unsafe {
            bindings::mongoc_cursor_more(self.inner) == 1
        }
    }

    fn error(&self) -> BsoncError {
        assert!(!self.inner.is_null());
        let mut error = BsoncError::empty();
        unsafe {
            bindings::mongoc_cursor_error(
                self.inner,
                error.mut_inner()
            )
        };
        error
    }
}

impl<'a> Iterator for Cursor<'a> {
    type Item = Result<Document>;

    fn next(&mut self) -> Option<Self::Item> {
        assert!(!self.inner.is_null());

        loop {
            if !self.more() {
                return None
            }

            // The C driver writes the document to memory and sets an
            // already existing pointer to it.
            let mut bson_ptr: *const bindings::bson_t = ptr::null();
            let success = unsafe {
                bindings::mongoc_cursor_next(
                    self.inner,
                    &mut bson_ptr
                )
            };

            // Fetch error that might have occurred while getting
            // the next cursor.
            let error = self.error();

            if success == 0 {
                if error.is_empty() {
                    if self.tailing && self.is_alive() {
                        // Since there was no error, this is a tailing cursor
                        // and the cursor is alive we'll wait before trying again.
                        thread::sleep_ms(self.tail_wait_time_ms);
                        continue;
                    } else {
                        // No result, no error and cursor not tailing so we must
                        // be at the end.
                        return None
                    }
                } else {
                    // There was an error
                    return Some(Err(error.into()))
                }
            }
            assert!(!bson_ptr.is_null());

            // Parse and return bson document.
            let bsonc = bsonc::Bsonc::from_ptr(bson_ptr);
            match bsonc.as_document() {
                Ok(document) => return Some(Ok(document)),
                Err(error)   => return Some(Err(error.into()))
            }
        }
    }
}

impl<'a> Drop for Cursor<'a> {
    fn drop(&mut self) {
        assert!(!self.inner.is_null());
        unsafe {
            bindings::mongoc_cursor_destroy(self.inner);
        }
    }
}

/// Cursor that will reconnect and resume tailing a collection
/// at the right point if the connection fails.
pub struct TailingCursor<'a> {
    collection:   &'a Collection<'a>,
    query:        Document,
    find_options: CommandAndFindOptions,
    tail_options: TailOptions,
    cursor:       Option<Cursor<'a>>,
    last_seen_id: Option<oid::ObjectId>,
    retry_count:  u32
}

impl<'a> TailingCursor<'a> {
    pub fn new(
        collection:   &'a Collection<'a>,
        query:        Document,
        find_options: CommandAndFindOptions,
        tail_options: TailOptions
    ) -> TailingCursor<'a> {
        // Add flags to make query tailable
        let mut find_options = find_options;
        find_options.query_flags.add(QueryFlag::TailableCursor);
        find_options.query_flags.add(QueryFlag::AwaitData);

        TailingCursor {
            collection:   collection,
            query:        query,
            find_options: find_options,
            tail_options: tail_options,
            cursor:       None,
            last_seen_id: None,
            retry_count:  0
        }
    }
}

impl<'a> Iterator for TailingCursor<'a> {
    type Item = Result<Document>;

    fn next(&mut self) -> Option<Self::Item> {
        loop {
            // Start a scope so we're free to set the cursor to None at the end.
            {
                if self.cursor.is_none() {
                    // Add the last seen id to the query if it's present.
                    match self.last_seen_id.take() {
                        Some(id) => {
                            self.query.insert_bson("_id".to_string(), Bson::Document(doc!{ "$gt" => id }));
                        },
                        None => ()
                    };

                    // Set the cursor
                    self.cursor = match self.collection.find(&self.query, Some(&self.find_options)) {
                        Ok(mut c)  => {
                            c.tailing           = true;
                            c.tail_wait_time_ms = self.tail_options.wait_time_ms;
                            Some(c)
                        },
                        Err(e) => return Some(Err(e.into()))
                    };
                }

                let cursor = match self.cursor {
                    Some(ref mut c) => c,
                    None => panic!("It should be impossible to not have a cursor here")
                };

                match cursor.next() {
                    Some(next_result) => {
                        match next_result {
                            Ok(next) => {
                                // This was successfull, so reset retry count and return result.
                                self.retry_count = 0;
                                return Some(Ok(next))
                            },
                            Err(e) => {
                                // Retry if we haven't exceeded the maximum number of retries.
                                if self.retry_count >= self.tail_options.max_retries {
                                    return Some(Err(e.into()))
                                }
                            }
                        }
                    },
                    None => ()
                };
            }

            // We made it to the end, so we weren't able to get the next item from
            // the cursor. We need to reconnect in the next iteration of the loop.
            self.retry_count += 1;
            self.cursor      = None;
        }
    }
}
<|endoftext|>


<commit_before>//! # Glulx memory mapping
//!
//! The Glulx machine memory consists of a ROM and RAM, which themselves
//! are separated into several different groups.
//!
//!
//! ## Program Provided Values
//!
//! The following values are defined by the given program:
//!
//! * RAMSTART -- Specifies the starting point of the RAM
//! * EXTSTART -- Specifies start of zeroed section of RAM
//! * ENDMEM -- Specifies the ending point of the RAM
//!
//!
//! ## ROM
//!
//! The ROM is defined by the following:
//!
//! * The ROM ranges from 0x00000000 to RAMSTART.
//! * The Header ranges from 0x00000000 to 0x00000024.
//! * The ROM is read-only.
//! * The ROM must be at least 0x100 bytes long.
//! * The ROM usually (but not always) contains all the executable code
//! and constants for the loaded program.
//!
//!
//! ## RAM
//!
//! The RAM is defined by the following:
//!
//! * The RAM ranges from RAMSTART to ENDMEM
//! * The RAM from EXTSTART to ENDMEM is initialized to all zeros.
//! * The RAM can be 0x0 bytes long.
//!
//! ## Misc
//!
//! * RAMSTART, EXTSTART, and ENDMEM must be aligned on 0x100 byte
//! boundries
//! * A Glulx gamefile only stores data from 0x0 to EXTSTART.

use byteorder::{BigEndian, ByteOrder};

/// Glulx magic number
const MAGIC_NUMBER: u32 = 0x476C756C;

/// Minimum supported glulx target
const MIN_VERSION: u32 = 0x0020000;

/// Maximum supported glulx target
const MAX_VERSION: u32 = 0x00301FF;

/// Struct representing a glulx memory object.
pub struct GlulxMemory {
    heap_mode: bool,
    memory: Vec<u8>,
}

impl GlulxMemory {

    /// This takes a rom, validates it to verify that it is a valid
    /// Glulx rom, and retuns either an error string or a `GlulxMemory`
    /// TODO: Replace error string with custom glulx error type.
    pub fn from_rom(rom: Vec<u8>) -> Result<GlulxMemory, String> {
        // Validate magic number.
        let magic_number = BigEndian::read_u32(&rom[0x0..0x4]);
        if magic_number != MAGIC_NUMBER {
            return Err("executable code starts with invalid magic number"
                .to_string());
        }

        // Validate glulx version.
        let glulx_version = BigEndian::read_u32(&rom[0x4..0x8]);
        if glulx_version < MIN_VERSION {
            return Err("executable code glulx version is less than 2.0.0"
                .to_string());
        } else if glulx_version > MAX_VERSION {
            return Err("executable code glulx version is greater than 3.1.x"
                .to_string());
        }

        // Validate ramstart value.
        let ramstart = BigEndian::read_u32(&rom[0x8..0xC]);
        if ramstart < 0x100 {
            return Err("ramstart is less than 0x100 bytes".to_string());
        } else if ramstart % 0x100 != 0 {
            return Err("ramstart is misaligned".to_string());
        }

        // Validate extstart value.
        let extstart = BigEndian::read_u32(&rom[0xC..0x10]);
        if extstart < ramstart {
            return Err("extstart is less than ramstart".to_string());
        } else if extstart % 0x100 != 0 {
            return Err("extstart is misaligned".to_string());
        }

        // Validate endmem value.
        let endmem = BigEndian::read_u32(&rom[0x10..0x14]);
        if endmem < extstart {
            return Err("endmem is less than extstart".to_string());
        } else if endmem % 0x100 != 0 {
            return Err("endmem is misaligned".to_string());
        }

        // Validate stack size.
        let stack_size = BigEndian::read_u32(&rom[0x14..0x18]);
        if stack_size % 0x100 != 0 {
            return Err("stack size is misaligned".to_string());
        }

        if endmem as usize >= rom.len() {
            return Err("rom size is not correct".to_string());
        }

        let checksum = BigEndian::read_u32(&rom[0x20..0x24]);
        let sum = {
            let mut sum = 0u32;
            for i in 0..rom.len()/4 {
                sum = sum.wrapping_add(BigEndian::read_u32(&rom[i*4..]));
            }
            sum.wrapping_sub(checksum)
        };
        if checksum != sum {
            return Err("rom checksum is not correct".to_string());
        }

        let mut rom = rom;
        let ext_size = (endmem - extstart) as usize;
        rom.reserve_exact(ext_size);
        rom.resize(endmem as usize, 0x0);

        Ok(GlulxMemory { heap_mode: false, memory: rom })
    }

    pub fn zero_range(&mut self, size: u32, ptr: u32){
        for i in ptr as usize..(ptr + size) as usize {
            self.memory[i] = 0;
        }
    }

    pub fn copy_range(&mut self, size: u32, from_ptr: u32, to_ptr: u32){
        let to_ptr = to_ptr as usize;
        let from_ptr = from_ptr as usize;

        for i in 0x0..size as usize {
            self.memory[to_ptr + i] = self.memory[from_ptr + i];
        }
    }

    // Header value functions.

    /// The glulx magic number, stored from `0x0..0x4` in the header.
    /// This number is unique to all glulx games. This value should be
    /// `0x476C756C`, which is equivalent to `b"Glul"`.
    fn magic_number(&self) -> u32 {
        BigEndian::read_u32(&self.memory[..0x4])
    }

    /// The glulx version number, stored from `0x4..0x8` in the header.
    /// First two bytes store the major version number. The next byte
    /// stores the minor version number. The final byte stores the patch
    /// version number.
    pub fn glulx_version(&self) -> u32 {
        BigEndian::read_u32(&self.memory[0x4..0x8])
    }

    /// The address indicating the start of the RAM, stored from
    /// `0x4..0x8` in the header. TODO: MOAR DATA
    fn ramstart(&self) -> u32 {
        BigEndian::read_u32(&self.memory[0x8..0xC])
    }

    fn extstart(&self) -> u32 {
        BigEndian::read_u32(&self.memory[0xC..0x10])
    }

    fn endmem(&self) -> u32 {
        BigEndian::read_u32(&self.memory[0x10..0x14])
    }

    pub fn stack_size(&self) -> u32 {
        BigEndian::read_u32(&self.memory[0x14..0x18])
    }

    pub fn start_func(&self) -> u32 {
        BigEndian::read_u32(&self.memory[0x18..0x1C])
    }

    fn decoding_tbl(&self) -> u32 {
        BigEndian::read_u32(&self.memory[0x1C..0x20])
    }

    /// The sum of the intial contents of memory, considered as an array
    /// of `u32`s. When calculated, the checksum value is considered to
    /// be `0`.
    fn checksum(&self) -> u32 {
        BigEndian::read_u32(&self.memory[0x20..0x24])
    }

    pub fn get_mem_size(&self) -> u32 {
        self.memory.len() as u32
    }

    /// Sets the memory size to the given value. This call only works if
    /// there is no active heap, the given value is a multiple of 0x100,
    /// and the value is greater than the ENDMEM value identified in the
    /// header.
    pub fn set_mem_size(&mut self, value: u32) -> u32 {
        if self.heap_mode
                && value % 0x100 == 0
                && value >= self.endmem() {
            self.memory.resize(value as usize, 0x0);
            0
        } else {
            1
        }
    }
}

pub trait Memory<T> {
    fn read(&self, ptr: u32) -> T;
    fn write(&mut self, ptr: u32, value: T);
    fn ram_read(&self, ptr: u32) -> T;
    fn ram_write(&mut self, ptr: u32, value: T);
}

impl Memory<u8> for GlulxMemory {
    fn read(&self, ptr: u32) -> u8 {
        self.memory[ptr as usize]
    }

    fn write(&mut self, ptr: u32, value: u8) {
        self.memory[ptr as usize] = value
    }

    fn ram_read(&self, ptr: u32) -> u8 {
        let ptr = (ptr + self.ramstart()) as usize;
        self.memory[ptr]
    }

    fn ram_write(&mut self, ptr: u32, value: u8) {
        let ptr = (ptr + self.ramstart()) as usize;
        self.memory[ptr] = value
    }
}

impl Memory<u16> for GlulxMemory {
    fn read(&self, ptr: u32) -> u16 {
        BigEndian::read_u16(&self.memory[ptr as usize..])
    }

    fn write(&mut self, ptr: u32, value: u16) {
        BigEndian::write_u16(&mut self.memory[ptr as usize..], value)
    }

    fn ram_read(&self, ptr: u32) -> u16 {
        let ptr = (ptr + self.ramstart()) as usize;
        BigEndian::read_u16(&self.memory[ptr..])
    }

    fn ram_write(&mut self, ptr: u32, value: u16) {
        let ptr = (ptr + self.ramstart()) as usize;
        BigEndian::write_u16(&mut self.memory[ptr..], value)
    }
}

impl Memory<u32> for GlulxMemory {
    fn read(&self, ptr: u32) -> u32 {
        BigEndian::read_u32(&self.memory[ptr as usize..])
    }

    fn write(&mut self, ptr: u32, value: u32) {
        BigEndian::write_u32(&mut self.memory[ptr as usize..], value)
    }

    fn ram_read(&self, ptr: u32) -> u32 {
        let ptr = (ptr + self.ramstart()) as usize;
        BigEndian::read_u32(&self.memory[ptr..])
    }

    fn ram_write(&mut self, ptr: u32, value: u32) {
        let ptr = (ptr + self.ramstart()) as usize;
        BigEndian::write_u32(&mut self.memory[ptr..], value)
    }
}

impl Memory<i32> for GlulxMemory {
    fn read(&self, ptr: u32) -> i32 {
        BigEndian::read_i32(&self.memory[ptr as usize..])
    }

    fn write(&mut self, ptr: u32, value: i32) {
        BigEndian::write_i32(&mut self.memory[ptr as usize..], value)
    }

    fn ram_read(&self, ptr: u32) -> i32 {
        let ptr = (ptr + self.ramstart()) as usize;
        BigEndian::read_i32(&self.memory[ptr..])
    }

    fn ram_write(&mut self, ptr: u32, value: i32) {
        let ptr = (ptr + self.ramstart()) as usize;
        BigEndian::write_i32(&mut self.memory[ptr..], value)
    }
}

impl Memory<f32> for GlulxMemory {
    fn read(&self, ptr: u32) -> f32 {
        BigEndian::read_f32(&self.memory[ptr as usize..])
    }

    fn write(&mut self, ptr: u32, value: f32) {
        BigEndian::write_f32(&mut self.memory[ptr as usize..], value)
    }

    fn ram_read(&self, ptr: u32) -> f32 {
        let ptr = (ptr + self.ramstart()) as usize;
        BigEndian::read_f32(&self.memory[ptr..])
    }

    fn ram_write(&mut self, ptr: u32, value: f32) {
        let ptr = (ptr + self.ramstart()) as usize;
        BigEndian::write_f32(&mut self.memory[ptr..], value)
    }
}
<commit_msg>fix the rom len changed just made<commit_after>//! # Glulx memory mapping
//!
//! The Glulx machine memory consists of a ROM and RAM, which themselves
//! are separated into several different groups.
//!
//!
//! ## Program Provided Values
//!
//! The following values are defined by the given program:
//!
//! * RAMSTART -- Specifies the starting point of the RAM
//! * EXTSTART -- Specifies start of zeroed section of RAM
//! * ENDMEM -- Specifies the ending point of the RAM
//!
//!
//! ## ROM
//!
//! The ROM is defined by the following:
//!
//! * The ROM ranges from 0x00000000 to RAMSTART.
//! * The Header ranges from 0x00000000 to 0x00000024.
//! * The ROM is read-only.
//! * The ROM must be at least 0x100 bytes long.
//! * The ROM usually (but not always) contains all the executable code
//! and constants for the loaded program.
//!
//!
//! ## RAM
//!
//! The RAM is defined by the following:
//!
//! * The RAM ranges from RAMSTART to ENDMEM
//! * The RAM from EXTSTART to ENDMEM is initialized to all zeros.
//! * The RAM can be 0x0 bytes long.
//!
//! ## Misc
//!
//! * RAMSTART, EXTSTART, and ENDMEM must be aligned on 0x100 byte
//! boundries
//! * A Glulx gamefile only stores data from 0x0 to EXTSTART.

use byteorder::{BigEndian, ByteOrder};

/// Glulx magic number
const MAGIC_NUMBER: u32 = 0x476C756C;

/// Minimum supported glulx target
const MIN_VERSION: u32 = 0x0020000;

/// Maximum supported glulx target
const MAX_VERSION: u32 = 0x00301FF;

/// Struct representing a glulx memory object.
pub struct GlulxMemory {
    heap_mode: bool,
    memory: Vec<u8>,
}

impl GlulxMemory {

    /// This takes a rom, validates it to verify that it is a valid
    /// Glulx rom, and retuns either an error string or a `GlulxMemory`
    /// TODO: Replace error string with custom glulx error type.
    pub fn from_rom(rom: Vec<u8>) -> Result<GlulxMemory, String> {
        // Validate magic number.
        let magic_number = BigEndian::read_u32(&rom[0x0..0x4]);
        if magic_number != MAGIC_NUMBER {
            return Err("executable code starts with invalid magic number"
                .to_string());
        }

        // Validate glulx version.
        let glulx_version = BigEndian::read_u32(&rom[0x4..0x8]);
        if glulx_version < MIN_VERSION {
            return Err("executable code glulx version is less than 2.0.0"
                .to_string());
        } else if glulx_version > MAX_VERSION {
            return Err("executable code glulx version is greater than 3.1.x"
                .to_string());
        }

        // Validate ramstart value.
        let ramstart = BigEndian::read_u32(&rom[0x8..0xC]);
        if ramstart < 0x100 {
            return Err("ramstart is less than 0x100 bytes".to_string());
        } else if ramstart % 0x100 != 0 {
            return Err("ramstart is misaligned".to_string());
        }

        // Validate extstart value.
        let extstart = BigEndian::read_u32(&rom[0xC..0x10]);
        if extstart < ramstart {
            return Err("extstart is less than ramstart".to_string());
        } else if extstart % 0x100 != 0 {
            return Err("extstart is misaligned".to_string());
        }

        // Validate endmem value.
        let endmem = BigEndian::read_u32(&rom[0x10..0x14]);
        if endmem < extstart {
            return Err("endmem is less than extstart".to_string());
        } else if endmem % 0x100 != 0 {
            return Err("endmem is misaligned".to_string());
        }

        // Validate stack size.
        let stack_size = BigEndian::read_u32(&rom[0x14..0x18]);
        if stack_size % 0x100 != 0 {
            return Err("stack size is misaligned".to_string());
        }

        if (endmem as usize) < rom.len() {
            return Err("rom size is not correct".to_string());
        }

        let checksum = BigEndian::read_u32(&rom[0x20..0x24]);
        let sum = {
            let mut sum = 0u32;
            for i in 0..rom.len()/4 {
                sum = sum.wrapping_add(BigEndian::read_u32(&rom[i*4..]));
            }
            sum.wrapping_sub(checksum)
        };
        if checksum != sum {
            return Err("rom checksum is not correct".to_string());
        }

        let mut rom = rom;
        let ext_size = (endmem - extstart) as usize;
        rom.reserve_exact(ext_size);
        rom.resize(endmem as usize, 0x0);

        Ok(GlulxMemory { heap_mode: false, memory: rom })
    }

    pub fn zero_range(&mut self, size: u32, ptr: u32){
        for i in ptr as usize..(ptr + size) as usize {
            self.memory[i] = 0;
        }
    }

    pub fn copy_range(&mut self, size: u32, from_ptr: u32, to_ptr: u32){
        let to_ptr = to_ptr as usize;
        let from_ptr = from_ptr as usize;

        for i in 0x0..size as usize {
            self.memory[to_ptr + i] = self.memory[from_ptr + i];
        }
    }

    // Header value functions.

    /// The glulx magic number, stored from `0x0..0x4` in the header.
    /// This number is unique to all glulx games. This value should be
    /// `0x476C756C`, which is equivalent to `b"Glul"`.
    fn magic_number(&self) -> u32 {
        BigEndian::read_u32(&self.memory[..0x4])
    }

    /// The glulx version number, stored from `0x4..0x8` in the header.
    /// First two bytes store the major version number. The next byte
    /// stores the minor version number. The final byte stores the patch
    /// version number.
    pub fn glulx_version(&self) -> u32 {
        BigEndian::read_u32(&self.memory[0x4..0x8])
    }

    /// The address indicating the start of the RAM, stored from
    /// `0x4..0x8` in the header. TODO: MOAR DATA
    fn ramstart(&self) -> u32 {
        BigEndian::read_u32(&self.memory[0x8..0xC])
    }

    fn extstart(&self) -> u32 {
        BigEndian::read_u32(&self.memory[0xC..0x10])
    }

    fn endmem(&self) -> u32 {
        BigEndian::read_u32(&self.memory[0x10..0x14])
    }

    pub fn stack_size(&self) -> u32 {
        BigEndian::read_u32(&self.memory[0x14..0x18])
    }

    pub fn start_func(&self) -> u32 {
        BigEndian::read_u32(&self.memory[0x18..0x1C])
    }

    fn decoding_tbl(&self) -> u32 {
        BigEndian::read_u32(&self.memory[0x1C..0x20])
    }

    /// The sum of the intial contents of memory, considered as an array
    /// of `u32`s. When calculated, the checksum value is considered to
    /// be `0`.
    fn checksum(&self) -> u32 {
        BigEndian::read_u32(&self.memory[0x20..0x24])
    }

    pub fn get_mem_size(&self) -> u32 {
        self.memory.len() as u32
    }

    /// Sets the memory size to the given value. This call only works if
    /// there is no active heap, the given value is a multiple of 0x100,
    /// and the value is greater than the ENDMEM value identified in the
    /// header.
    pub fn set_mem_size(&mut self, value: u32) -> u32 {
        if self.heap_mode
                && value % 0x100 == 0
                && value >= self.endmem() {
            self.memory.resize(value as usize, 0x0);
            0
        } else {
            1
        }
    }
}

pub trait Memory<T> {
    fn read(&self, ptr: u32) -> T;
    fn write(&mut self, ptr: u32, value: T);
    fn ram_read(&self, ptr: u32) -> T;
    fn ram_write(&mut self, ptr: u32, value: T);
}

impl Memory<u8> for GlulxMemory {
    fn read(&self, ptr: u32) -> u8 {
        self.memory[ptr as usize]
    }

    fn write(&mut self, ptr: u32, value: u8) {
        self.memory[ptr as usize] = value
    }

    fn ram_read(&self, ptr: u32) -> u8 {
        let ptr = (ptr + self.ramstart()) as usize;
        self.memory[ptr]
    }

    fn ram_write(&mut self, ptr: u32, value: u8) {
        let ptr = (ptr + self.ramstart()) as usize;
        self.memory[ptr] = value
    }
}

impl Memory<u16> for GlulxMemory {
    fn read(&self, ptr: u32) -> u16 {
        BigEndian::read_u16(&self.memory[ptr as usize..])
    }

    fn write(&mut self, ptr: u32, value: u16) {
        BigEndian::write_u16(&mut self.memory[ptr as usize..], value)
    }

    fn ram_read(&self, ptr: u32) -> u16 {
        let ptr = (ptr + self.ramstart()) as usize;
        BigEndian::read_u16(&self.memory[ptr..])
    }

    fn ram_write(&mut self, ptr: u32, value: u16) {
        let ptr = (ptr + self.ramstart()) as usize;
        BigEndian::write_u16(&mut self.memory[ptr..], value)
    }
}

impl Memory<u32> for GlulxMemory {
    fn read(&self, ptr: u32) -> u32 {
        BigEndian::read_u32(&self.memory[ptr as usize..])
    }

    fn write(&mut self, ptr: u32, value: u32) {
        BigEndian::write_u32(&mut self.memory[ptr as usize..], value)
    }

    fn ram_read(&self, ptr: u32) -> u32 {
        let ptr = (ptr + self.ramstart()) as usize;
        BigEndian::read_u32(&self.memory[ptr..])
    }

    fn ram_write(&mut self, ptr: u32, value: u32) {
        let ptr = (ptr + self.ramstart()) as usize;
        BigEndian::write_u32(&mut self.memory[ptr..], value)
    }
}

impl Memory<i32> for GlulxMemory {
    fn read(&self, ptr: u32) -> i32 {
        BigEndian::read_i32(&self.memory[ptr as usize..])
    }

    fn write(&mut self, ptr: u32, value: i32) {
        BigEndian::write_i32(&mut self.memory[ptr as usize..], value)
    }

    fn ram_read(&self, ptr: u32) -> i32 {
        let ptr = (ptr + self.ramstart()) as usize;
        BigEndian::read_i32(&self.memory[ptr..])
    }

    fn ram_write(&mut self, ptr: u32, value: i32) {
        let ptr = (ptr + self.ramstart()) as usize;
        BigEndian::write_i32(&mut self.memory[ptr..], value)
    }
}

impl Memory<f32> for GlulxMemory {
    fn read(&self, ptr: u32) -> f32 {
        BigEndian::read_f32(&self.memory[ptr as usize..])
    }

    fn write(&mut self, ptr: u32, value: f32) {
        BigEndian::write_f32(&mut self.memory[ptr as usize..], value)
    }

    fn ram_read(&self, ptr: u32) -> f32 {
        let ptr = (ptr + self.ramstart()) as usize;
        BigEndian::read_f32(&self.memory[ptr..])
    }

    fn ram_write(&mut self, ptr: u32, value: f32) {
        let ptr = (ptr + self.ramstart()) as usize;
        BigEndian::write_f32(&mut self.memory[ptr..], value)
    }
}
<|endoftext|>