<commit_before>use std::env::*;
use std::env;
use std::fs;
use std::fs::File;
use std::path::Path;
use std::error::Error;
use std::io::prelude::*;
use std::io::BufReader;
use std::ascii::AsciiExt;
use hyper::Client;
use hyper::header::Connection;

extern crate rustc_serialize;
use self::rustc_serialize::json::*;

extern crate chrono;
use self::chrono::*;

#[derive(Clone, Debug)]
pub struct AWSCredentials {
    key: String,
    secret: String,
    token: Option<String>,
    expires_at: DateTime<UTC>
}

impl AWSCredentials {
    pub fn new(key:&str, secret:&str, token:Option<String>, expires_at:DateTime<UTC>) -> AWSCredentials {
        AWSCredentials {
            key: key.to_string(),
            secret: secret.to_string(),
            token: token,
            expires_at: expires_at,
        }
    }

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

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

    pub fn get_expires_at(&self) -> &DateTime<UTC> {
        &self.expires_at
    }

    pub fn get_token(&self) -> &Option<String> {
        &self.token
    }

    fn credentials_are_expired(&self) -> bool {
        //println!("Seeing if creds of {:?} are expired compared to {:?}", self.expires_at, UTC::now() + Duration::seconds(20));
        // This is a rough hack to hopefully avoid someone requesting creds then sitting on them
        // before issuing the request:
        if self.expires_at < UTC::now() + Duration::seconds(20) {
            return true;
        }
        return false;
    }
}

pub trait AWSCredentialsProvider {
    fn new() -> Self;
	fn get_credentials(&mut self) -> Result<&AWSCredentials, &str>;
}

pub struct EnvironmentCredentialsProvider {
    credentials: AWSCredentials
}

impl AWSCredentialsProvider for EnvironmentCredentialsProvider {
    fn new() -> EnvironmentCredentialsProvider {
        // expired by default
        return EnvironmentCredentialsProvider {credentials: AWSCredentials{ key: "".to_string(),
            secret: "".to_string(), token: None, expires_at: UTC::now() - Duration::seconds(600) } };
    }

	fn get_credentials(&mut self) -> Result<&AWSCredentials, &str> {
        if self.credentials.credentials_are_expired() {
            let env_key = match var("AWS_ACCESS_KEY_ID") {
                Ok(val) => val,
                Err(_) => return Err("No AWS_ACCESS_KEY_ID in environment")
            };
            let env_secret = match var("AWS_SECRET_ACCESS_KEY") {
                Ok(val) => val,
                Err(_) => return Err("No AWS_SECRET_ACCESS_KEY in environment")
            };

            if env_key.is_empty() || env_secret.is_empty() {
                return Err("Couldn't find either AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY or both in environment.");
            }

            self.credentials = AWSCredentials{key: env_key, secret: env_secret,
                 token: None, expires_at: UTC::now() + Duration::seconds(600)};
        }

		Ok(&self.credentials)
	}
}

pub struct FileCredentialsProvider {
    credentials: AWSCredentials
}

impl AWSCredentialsProvider for FileCredentialsProvider {
    fn new() -> FileCredentialsProvider {
        return FileCredentialsProvider {credentials: AWSCredentials{ key: "".to_string(),
            secret: "".to_string(), token: None, expires_at: UTC::now() - Duration::seconds(600) } };
    }

    fn get_credentials(&mut self) -> Result<&AWSCredentials, &str> {
        if self.credentials.credentials_are_expired() {
            // Default credentials file location:
            // ~/.aws/credentials (Linux/Mac)
            // %USERPROFILE%\.aws\credentials  (Windows)
            let mut profile_location;
            match env::home_dir() {
                Some(ref p) => profile_location = p.display().to_string() + "/.aws/credentials",
                None => return Err("Couldn't get your home dir.")
            }
            match get_credentials_from_file(profile_location) {
                Ok(creds) => {
                    self.credentials = AWSCredentials{ key: creds.get_aws_access_key_id().to_string(),
                        secret: creds.get_aws_secret_key().to_string(), token: None,
                        expires_at: UTC::now() + Duration::seconds(600) };
                }
                Err(_) => return Err("Couldn't get creds from file.")
            }
       }
       Ok(&self.credentials)
   }
}

// Finds and uses the first "aws_access_key_id" and "aws_secret_access_key" in the file.
fn get_credentials_from_file<'a>(file_with_path: String) -> Result<AWSCredentials, &'a str> {
    //println!("Looking for credentials file at {}", file_with_path);
    let path = Path::new(&file_with_path);
    let display = path.display();

    match fs::metadata(&path) {
        Err(_) => return Err("Couldn't stat credentials file."),
        Ok(metadata) => {
            if !metadata.is_file() {
                return Err("Couldn't open file.")
            }
        }
    };

    let file = match File::open(&path) {
        Err(why) => panic!("couldn't open {}: {}", display, Error::description(&why)),
        Ok(opened_file) => opened_file,
    };

    let mut access_key = String::new();
    let mut secret_key = String::new();
    let file_lines = BufReader::new(&file);
    for line in file_lines.lines() {
        let unwrapped_line : String = line.unwrap();
        if unwrapped_line.starts_with('#') {
            continue;
        }

        let lower_case_line = unwrapped_line.to_ascii_lowercase().to_string();

        if lower_case_line.contains("aws_access_key_id") {
            if access_key.is_empty() {
                let v: Vec<&str> = unwrapped_line.split("=").collect();
                if v.len() == 0 {
                    access_key = "".to_string();
                } else {
                    access_key = v[1].trim_matches(' ').to_string();
                }
            }
        } else if lower_case_line.contains("aws_secret_access_key") {
            if secret_key.is_empty() {
                let v: Vec<&str> = unwrapped_line.split("=").collect();
                if v.len() == 0 {
                    secret_key = "".to_string();
                } else {
                    secret_key = v[1].trim_matches(' ').to_string();
                }
            }
        }
    }

    if access_key.is_empty() || secret_key.is_empty() {
        return Err("Couldn't find either aws_access_key_id, aws_secret_access_key or both in profile file.");
    }

    return Ok(AWSCredentials{ key: access_key.to_string(), secret: secret_key.to_string(),
         token:None, expires_at: UTC::now() + Duration::seconds(600) });
}

pub struct IAMRoleCredentialsProvider {
    credentials: AWSCredentials
}

impl AWSCredentialsProvider for IAMRoleCredentialsProvider {
    fn new() -> IAMRoleCredentialsProvider {
        return IAMRoleCredentialsProvider {credentials: AWSCredentials{ key: "".to_string(),
        secret: "".to_string(), token: None, expires_at: UTC::now() - Duration::seconds(600) } };
    }

    fn get_credentials(&mut self) -> Result<&AWSCredentials, &str> {
        if self.credentials.credentials_are_expired() {
            // TODO: backoff and retry on failure.

            //println!("Calling IAM metadata");
            // for "real" use: http://169.254.169.254/latest/meta-data/iam/security-credentials/
            let mut address : String = "http://169.254.169.254/latest/meta-data/iam/security-credentials".to_string();
            let client = Client::new();
            let mut response;
            match client.get(&address)
                .header(Connection::close()).send() {
                    Err(_) => return Err("Couldn't connect to metadata service"), // add Why?
                    Ok(received_response) => response = received_response
                };

            let mut body = String::new();
            match response.read_to_string(&mut body) {
                Err(_) => return Err("Didn't get a parsable response body from metadata service"),
                Ok(_) => (),
            };

            address.push_str("/");
            address.push_str(&body);
            body = String::new();
            match client.get(&address)
                .header(Connection::close()).send() {
                    Err(_) => return Err("Didn't get a parseable response body from instance role details"),
                    Ok(received_response) => response = received_response
                };

            match response.read_to_string(&mut body) {
                Err(_) => return Err("Had issues with reading iam role response: {}"),
                Ok(_) => (),
            };

            let json_object : Json;
            match Json::from_str(&body) {
                Err(_) => return Err("Couldn't parse metadata response body."),
                Ok(val) => json_object = val
            };

            let mut access_key;
            match json_object.find("AccessKeyId") {
                None => return Err("Couldn't find AccessKeyId in response."),
                Some(val) => access_key = val.to_string().replace("\"", "")
            };

            let mut secret_key;
            match json_object.find("SecretAccessKey") {
                None => return Err("Couldn't find SecretAccessKey in response."),
                Some(val) => secret_key = val.to_string().replace("\"", "")
            };

            let mut expiration;
            match json_object.find("Expiration") {
                None => return Err("Couldn't find Expiration in response."),
                Some(val) => expiration = val.to_string().replace("\"", "")
            };

            let mut expiration_time;
            match expiration.parse::<DateTime<UTC>>() {
                Err(why) => panic!("Kabloey on parse: {}", why),
                Ok(val) => expiration_time = val
            };

            let mut token_from_response;
            match json_object.find("Token") {
                None => return Err("Couldn't find Token in response."),
                Some(val) => token_from_response = val.to_string().replace("\"", "")
            };

            self.credentials = AWSCredentials{ key: access_key.to_string(),
                secret: secret_key.to_string(),  token: Some(token_from_response.to_string()), expires_at: expiration_time };
        }

		Ok(&self.credentials)
	}
}

#[derive(Debug, Clone)]
pub struct DefaultAWSCredentialsProviderChain {
    credentials: AWSCredentials
}

// Chain the providers:
impl DefaultAWSCredentialsProviderChain {
    pub fn new() -> DefaultAWSCredentialsProviderChain {
        return DefaultAWSCredentialsProviderChain {credentials: AWSCredentials{ key: "".to_string(),
            secret: "".to_string(), token: None, expires_at: UTC::now() - Duration::seconds(600) } };
    }

    pub fn get_credentials(&mut self) -> &AWSCredentials {
        if self.credentials.credentials_are_expired() {
            // fetch creds in order: env, file, IAM
            let mut env_provider = EnvironmentCredentialsProvider::new();
            match env_provider.get_credentials() {
                Ok(creds) => {
                    //println!("Found creds in env");
                    self.credentials = AWSCredentials{ key: creds.get_aws_access_key_id().to_string(),
                        secret: creds.get_aws_secret_key().to_string(), token: None,
                        expires_at: UTC::now() + Duration::seconds(600) };
                    return &self.credentials;
                }
                Err(_) => () //println!("Whiffed on env")
            }

            let mut file_provider = FileCredentialsProvider::new();
            match file_provider.get_credentials() {
                Ok(creds) => {
                    //println!("Found creds in file");
                    self.credentials = AWSCredentials{ key: creds.get_aws_access_key_id().to_string(),
                        secret: creds.get_aws_secret_key().to_string(), token: None,
                        expires_at: UTC::now() + Duration::seconds(600) };
                    return &self.credentials;
                }
                Err(_) => () //println!("Whiffed on file")
            }

            let mut iam_provider = IAMRoleCredentialsProvider::new();
            match iam_provider.get_credentials() {
                Ok(creds) => {
                    //println!("Found creds via iam");
                    self.credentials = AWSCredentials{ key: creds.get_aws_access_key_id().to_string(),
                        secret: creds.get_aws_secret_key().to_string(), token: creds.get_token().clone(),
                        expires_at: UTC::now() + Duration::seconds(600) };
                    return &self.credentials;
                }
                Err(_) => panic!("Couldn't find AWS credentials in environment, default credential file location or IAM role.")
            }
        }
        return &self.credentials;
    }
}
<commit_msg>use Option<AWSCredentials> for EnvironmentCredentialsProvider<commit_after>use std::env::*;
use std::env;
use std::fs;
use std::fs::File;
use std::path::Path;
use std::error::Error;
use std::io::prelude::*;
use std::io::BufReader;
use std::ascii::AsciiExt;
use hyper::Client;
use hyper::header::Connection;

extern crate rustc_serialize;
use self::rustc_serialize::json::*;

extern crate chrono;
use self::chrono::*;

#[derive(Clone, Debug)]
pub struct AWSCredentials {
    key: String,
    secret: String,
    token: Option<String>,
    expires_at: DateTime<UTC>
}

impl AWSCredentials {
    pub fn new(key:&str, secret:&str, token:Option<String>, expires_at:DateTime<UTC>) -> AWSCredentials {
        AWSCredentials {
            key: key.to_string(),
            secret: secret.to_string(),
            token: token,
            expires_at: expires_at,
        }
    }

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

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

    pub fn get_expires_at(&self) -> &DateTime<UTC> {
        &self.expires_at
    }

    pub fn get_token(&self) -> &Option<String> {
        &self.token
    }

    fn credentials_are_expired(&self) -> bool {
        //println!("Seeing if creds of {:?} are expired compared to {:?}", self.expires_at, UTC::now() + Duration::seconds(20));
        // This is a rough hack to hopefully avoid someone requesting creds then sitting on them
        // before issuing the request:
        if self.expires_at < UTC::now() + Duration::seconds(20) {
            return true;
        }
        return false;
    }
}

pub trait AWSCredentialsProvider {
    fn new() -> Self;
	fn get_credentials(&mut self) -> Result<&AWSCredentials, &str>;
}

pub struct EnvironmentCredentialsProvider {
    credentials: Option<AWSCredentials>
}

impl AWSCredentialsProvider for EnvironmentCredentialsProvider {
    fn new() -> EnvironmentCredentialsProvider {
        // expired by default
        return EnvironmentCredentialsProvider { credentials: None }
    }

	fn get_credentials(&mut self) -> Result<&AWSCredentials, &str> {
        if self.credentials.is_none() || self.credentials.as_ref().unwrap().credentials_are_expired() {
           self.credentials = Some(try!(get_credentials_from_environment()));
        }
        Ok(self.credentials.as_ref().unwrap())
	}
}

fn get_credentials_from_environment<'a>() -> Result<AWSCredentials, &'a str> {
    let env_key = match var("AWS_ACCESS_KEY_ID") {
        Ok(val) => val,
        Err(_) => return Err("No AWS_ACCESS_KEY_ID in environment")
    };
    let env_secret = match var("AWS_SECRET_ACCESS_KEY") {
        Ok(val) => val,
        Err(_) => return Err("No AWS_SECRET_ACCESS_KEY in environment")
    };

    if env_key.is_empty() || env_secret.is_empty() {
        return Err("Couldn't find either AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY or both in environment.");
    }

    Ok(AWSCredentials{key: env_key, secret: env_secret,
       token: None, expires_at: UTC::now() + Duration::seconds(600)})
}

pub struct FileCredentialsProvider {
    credentials: AWSCredentials
}

impl AWSCredentialsProvider for FileCredentialsProvider {
    fn new() -> FileCredentialsProvider {
        return FileCredentialsProvider {credentials: AWSCredentials{ key: "".to_string(),
            secret: "".to_string(), token: None, expires_at: UTC::now() - Duration::seconds(600) } };
    }

    fn get_credentials(&mut self) -> Result<&AWSCredentials, &str> {
        if self.credentials.credentials_are_expired() {
            // Default credentials file location:
            // ~/.aws/credentials (Linux/Mac)
            // %USERPROFILE%\.aws\credentials  (Windows)
            let mut profile_location;
            match env::home_dir() {
                Some(ref p) => profile_location = p.display().to_string() + "/.aws/credentials",
                None => return Err("Couldn't get your home dir.")
            }
            match get_credentials_from_file(profile_location) {
                Ok(creds) => {
                    self.credentials = AWSCredentials{ key: creds.get_aws_access_key_id().to_string(),
                        secret: creds.get_aws_secret_key().to_string(), token: None,
                        expires_at: UTC::now() + Duration::seconds(600) };
                }
                Err(_) => return Err("Couldn't get creds from file.")
            }
       }
       Ok(&self.credentials)
   }
}

// Finds and uses the first "aws_access_key_id" and "aws_secret_access_key" in the file.
fn get_credentials_from_file<'a>(file_with_path: String) -> Result<AWSCredentials, &'a str> {
    //println!("Looking for credentials file at {}", file_with_path);
    let path = Path::new(&file_with_path);
    let display = path.display();

    match fs::metadata(&path) {
        Err(_) => return Err("Couldn't stat credentials file."),
        Ok(metadata) => {
            if !metadata.is_file() {
                return Err("Couldn't open file.")
            }
        }
    };

    let file = match File::open(&path) {
        Err(why) => panic!("couldn't open {}: {}", display, Error::description(&why)),
        Ok(opened_file) => opened_file,
    };

    let mut access_key = String::new();
    let mut secret_key = String::new();
    let file_lines = BufReader::new(&file);
    for line in file_lines.lines() {
        let unwrapped_line : String = line.unwrap();
        if unwrapped_line.starts_with('#') {
            continue;
        }

        let lower_case_line = unwrapped_line.to_ascii_lowercase().to_string();

        if lower_case_line.contains("aws_access_key_id") {
            if access_key.is_empty() {
                let v: Vec<&str> = unwrapped_line.split("=").collect();
                if v.len() == 0 {
                    access_key = "".to_string();
                } else {
                    access_key = v[1].trim_matches(' ').to_string();
                }
            }
        } else if lower_case_line.contains("aws_secret_access_key") {
            if secret_key.is_empty() {
                let v: Vec<&str> = unwrapped_line.split("=").collect();
                if v.len() == 0 {
                    secret_key = "".to_string();
                } else {
                    secret_key = v[1].trim_matches(' ').to_string();
                }
            }
        }
    }

    if access_key.is_empty() || secret_key.is_empty() {
        return Err("Couldn't find either aws_access_key_id, aws_secret_access_key or both in profile file.");
    }

    return Ok(AWSCredentials{ key: access_key.to_string(), secret: secret_key.to_string(),
         token:None, expires_at: UTC::now() + Duration::seconds(600) });
}

pub struct IAMRoleCredentialsProvider {
    credentials: AWSCredentials
}

impl AWSCredentialsProvider for IAMRoleCredentialsProvider {
    fn new() -> IAMRoleCredentialsProvider {
        return IAMRoleCredentialsProvider {credentials: AWSCredentials{ key: "".to_string(),
        secret: "".to_string(), token: None, expires_at: UTC::now() - Duration::seconds(600) } };
    }

    fn get_credentials(&mut self) -> Result<&AWSCredentials, &str> {
        if self.credentials.credentials_are_expired() {
            // TODO: backoff and retry on failure.

            //println!("Calling IAM metadata");
            // for "real" use: http://169.254.169.254/latest/meta-data/iam/security-credentials/
            let mut address : String = "http://169.254.169.254/latest/meta-data/iam/security-credentials".to_string();
            let client = Client::new();
            let mut response;
            match client.get(&address)
                .header(Connection::close()).send() {
                    Err(_) => return Err("Couldn't connect to metadata service"), // add Why?
                    Ok(received_response) => response = received_response
                };

            let mut body = String::new();
            match response.read_to_string(&mut body) {
                Err(_) => return Err("Didn't get a parsable response body from metadata service"),
                Ok(_) => (),
            };

            address.push_str("/");
            address.push_str(&body);
            body = String::new();
            match client.get(&address)
                .header(Connection::close()).send() {
                    Err(_) => return Err("Didn't get a parseable response body from instance role details"),
                    Ok(received_response) => response = received_response
                };

            match response.read_to_string(&mut body) {
                Err(_) => return Err("Had issues with reading iam role response: {}"),
                Ok(_) => (),
            };

            let json_object : Json;
            match Json::from_str(&body) {
                Err(_) => return Err("Couldn't parse metadata response body."),
                Ok(val) => json_object = val
            };

            let mut access_key;
            match json_object.find("AccessKeyId") {
                None => return Err("Couldn't find AccessKeyId in response."),
                Some(val) => access_key = val.to_string().replace("\"", "")
            };

            let mut secret_key;
            match json_object.find("SecretAccessKey") {
                None => return Err("Couldn't find SecretAccessKey in response."),
                Some(val) => secret_key = val.to_string().replace("\"", "")
            };

            let mut expiration;
            match json_object.find("Expiration") {
                None => return Err("Couldn't find Expiration in response."),
                Some(val) => expiration = val.to_string().replace("\"", "")
            };

            let mut expiration_time;
            match expiration.parse::<DateTime<UTC>>() {
                Err(why) => panic!("Kabloey on parse: {}", why),
                Ok(val) => expiration_time = val
            };

            let mut token_from_response;
            match json_object.find("Token") {
                None => return Err("Couldn't find Token in response."),
                Some(val) => token_from_response = val.to_string().replace("\"", "")
            };

            self.credentials = AWSCredentials{ key: access_key.to_string(),
                secret: secret_key.to_string(),  token: Some(token_from_response.to_string()), expires_at: expiration_time };
        }

		Ok(&self.credentials)
	}
}

#[derive(Debug, Clone)]
pub struct DefaultAWSCredentialsProviderChain {
    credentials: AWSCredentials
}

// Chain the providers:
impl DefaultAWSCredentialsProviderChain {
    pub fn new() -> DefaultAWSCredentialsProviderChain {
        return DefaultAWSCredentialsProviderChain {credentials: AWSCredentials{ key: "".to_string(),
            secret: "".to_string(), token: None, expires_at: UTC::now() - Duration::seconds(600) } };
    }

    pub fn get_credentials(&mut self) -> &AWSCredentials {
        if self.credentials.credentials_are_expired() {
            // fetch creds in order: env, file, IAM
            let mut env_provider = EnvironmentCredentialsProvider::new();
            match env_provider.get_credentials() {
                Ok(creds) => {
                    //println!("Found creds in env");
                    self.credentials = AWSCredentials{ key: creds.get_aws_access_key_id().to_string(),
                        secret: creds.get_aws_secret_key().to_string(), token: None,
                        expires_at: UTC::now() + Duration::seconds(600) };
                    return &self.credentials;
                }
                Err(_) => () //println!("Whiffed on env")
            }

            let mut file_provider = FileCredentialsProvider::new();
            match file_provider.get_credentials() {
                Ok(creds) => {
                    //println!("Found creds in file");
                    self.credentials = AWSCredentials{ key: creds.get_aws_access_key_id().to_string(),
                        secret: creds.get_aws_secret_key().to_string(), token: None,
                        expires_at: UTC::now() + Duration::seconds(600) };
                    return &self.credentials;
                }
                Err(_) => () //println!("Whiffed on file")
            }

            let mut iam_provider = IAMRoleCredentialsProvider::new();
            match iam_provider.get_credentials() {
                Ok(creds) => {
                    //println!("Found creds via iam");
                    self.credentials = AWSCredentials{ key: creds.get_aws_access_key_id().to_string(),
                        secret: creds.get_aws_secret_key().to_string(), token: creds.get_token().clone(),
                        expires_at: UTC::now() + Duration::seconds(600) };
                    return &self.credentials;
                }
                Err(_) => panic!("Couldn't find AWS credentials in environment, default credential file location or IAM role.")
            }
        }
        return &self.credentials;
    }
}

<|endoftext|>


<commit_before>use crypt32;
use std::io;
use std::mem;
use std::ptr;
use winapi;

use cert_context::CertContext;
use Inner;

pub struct CtlContext(winapi::PCCTL_CONTEXT);

unsafe impl Send for CtlContext {}
unsafe impl Sync for CtlContext {}

impl Drop for CtlContext {
	fn drop(&mut self) {
		unsafe {
			crypt32::CertFreeCTLContext(self.0);
		}
	}
}

impl Inner<winapi::PCCTL_CONTEXT> for CtlContext {
	unsafe fn from_inner(t: winapi::PCCTL_CONTEXT) -> CtlContext {
		CtlContext(t)
	}

	fn as_inner(&self) -> winapi::PCCTL_CONTEXT {
		self.0
	}

	fn get_mut(&mut self) -> &mut winapi::PCCTL_CONTEXT {
		&mut self.0
	}
}

impl CtlContext {
	pub fn builder() -> Builder {
		Builder {
			certificates: vec![],
			usages: vec![],
		}
	}
}

pub struct Builder {
	certificates: Vec<CertContext>,
	usages: Vec<Vec<u8>>,
}

impl Builder {
	pub fn certificate(&mut self, cert: CertContext) -> &mut Builder {
		self.certificates.push(cert);
		self
	}

	pub fn usage(&mut self, usage: &str) -> &mut Builder {
		let mut usage = usage.as_bytes().to_owned();
		usage.push(0);
		self.usages.push(usage);
		self
	}

	pub fn build(&self) -> io::Result<Vec<u8>> {
		unsafe {
			let encoding = winapi::X509_ASN_ENCODING | winapi::PKCS_7_ASN_ENCODING;

			let mut usages = self.usages.iter().map(|u| u.as_ptr()).collect::<Vec<_>>();
			let mut entry_data = vec![];
			let mut entries = vec![];
			for certificate in &self.certificates {
				let data = try!(cert_entry(certificate));
				entries.push(*(data.as_ptr() as *const winapi::CTL_ENTRY));
				entry_data.push(data);
			}

			let mut ctl_info: winapi::CTL_INFO = mem::zeroed();
			ctl_info.dwVersion = winapi::CTL_V1;
			ctl_info.SubjectUsage.cUsageIdentifier = usages.len() as winapi::DWORD;
			ctl_info.SubjectUsage.rgpszUsageIdentifier = usages.as_mut_ptr() as *mut winapi::LPSTR;
			let mut algorithm = winapi::szOID_OIWSEC_sha1.as_bytes().to_owned();
			algorithm.push(0);
			ctl_info.SubjectAlgorithm.pszObjId = algorithm.as_ptr() as winapi::LPSTR;
			ctl_info.cCTLEntry = entries.len() as winapi::DWORD;
			ctl_info.rgCTLEntry = entries.as_mut_ptr();

			let mut sign_info: winapi::CMSG_SIGNED_ENCODE_INFO = mem::zeroed();
			sign_info.cbSize = mem::size_of_val(&sign_info) as winapi::DWORD;
			let mut encoded_certs = self.certificates
										.iter()
										.map(|c| {
											winapi::CERT_BLOB {
												cbData: (*c.as_inner()).cbCertEncoded,
												pbData: (*c.as_inner()).pbCertEncoded,
											}
										})
										.collect::<Vec<_>>();
			sign_info.rgCertEncoded = encoded_certs.as_mut_ptr();
			sign_info.cCertEncoded = encoded_certs.len() as winapi::DWORD;

			let flags = winapi::CMSG_ENCODE_SORTED_CTL_FLAG |
				winapi::CMSG_ENCODE_HASHED_SUBJECT_IDENTIFIER_FLAG;

			let mut size = 0;

			let res = crypt32::CryptMsgEncodeAndSignCTL(
				encoding,
				&mut ctl_info,
				&mut sign_info,
				flags,
				ptr::null_mut(),
				&mut size);
			if res == winapi::FALSE {
				return Err(io::Error::last_os_error());
			}

			let mut encoded = vec![0; size as usize];

			let res = crypt32::CryptMsgEncodeAndSignCTL(
				encoding,
				&mut ctl_info,
				&mut sign_info,
				flags,
				encoded.as_mut_ptr() as *mut winapi::BYTE,
				&mut size);
			if res == winapi::FALSE {
				return Err(io::Error::last_os_error());
			}

			Ok(encoded)
		}
	}
}

fn cert_entry(cert: &CertContext) -> io::Result<Vec<u8>> {
	unsafe {
		let mut size = 0;

		let res = crypt32::CertCreateCTLEntryFromCertificateContextProperties(
			cert.as_inner(),
			0,
			ptr::null_mut(),
			winapi::CTL_ENTRY_FROM_PROP_CHAIN_FLAG,
			ptr::null_mut(),
			ptr::null_mut(),
			&mut size);
		if res == winapi::FALSE {
			return Err(io::Error::last_os_error());
		}

		let mut entry = vec![0; size as usize];
		let res = crypt32::CertCreateCTLEntryFromCertificateContextProperties(
			cert.as_inner(),
			0,
			ptr::null_mut(),
			winapi::CTL_ENTRY_FROM_PROP_CHAIN_FLAG,
			ptr::null_mut(),
			entry.as_mut_ptr() as winapi::PCTL_ENTRY,
			&mut size);
		if res == winapi::FALSE {
			Err(io::Error::last_os_error())
		} else {
			Ok(entry)
		}
	}
}

#[cfg(test)]
mod test {
	use std::io::Read;
	use std::fs::File;
	use cert_store::CertStore;
	use ctl_context::CtlContext;

	#[test]
	fn create_ctl() {
        let mut file = File::open(concat!(env!("CARGO_MANIFEST_DIR"),
                                          "/test/self-signed.badssl.com.cer"))
                           .unwrap();
        let mut cert = vec![];
        file.read_to_end(&mut cert).unwrap();

        let mut store = CertStore::memory().unwrap();
        let cert = store.add_encoded_certificate(&cert).unwrap();

        CtlContext::builder()
        	.certificate(cert)
            .usage("1.3.6.1.4.1.311.2.2.2")
            .build()
            .unwrap();
	}
}<commit_msg>Rename builder method<commit_after>use crypt32;
use std::io;
use std::mem;
use std::ptr;
use winapi;

use cert_context::CertContext;
use Inner;

pub struct CtlContext(winapi::PCCTL_CONTEXT);

unsafe impl Send for CtlContext {}
unsafe impl Sync for CtlContext {}

impl Drop for CtlContext {
	fn drop(&mut self) {
		unsafe {
			crypt32::CertFreeCTLContext(self.0);
		}
	}
}

impl Inner<winapi::PCCTL_CONTEXT> for CtlContext {
	unsafe fn from_inner(t: winapi::PCCTL_CONTEXT) -> CtlContext {
		CtlContext(t)
	}

	fn as_inner(&self) -> winapi::PCCTL_CONTEXT {
		self.0
	}

	fn get_mut(&mut self) -> &mut winapi::PCCTL_CONTEXT {
		&mut self.0
	}
}

impl CtlContext {
	pub fn builder() -> Builder {
		Builder {
			certificates: vec![],
			usages: vec![],
		}
	}
}

pub struct Builder {
	certificates: Vec<CertContext>,
	usages: Vec<Vec<u8>>,
}

impl Builder {
	pub fn certificate(&mut self, cert: CertContext) -> &mut Builder {
		self.certificates.push(cert);
		self
	}

	pub fn usage(&mut self, usage: &str) -> &mut Builder {
		let mut usage = usage.as_bytes().to_owned();
		usage.push(0);
		self.usages.push(usage);
		self
	}

	pub fn encode_and_sign(&self) -> io::Result<Vec<u8>> {
		unsafe {
			let encoding = winapi::X509_ASN_ENCODING | winapi::PKCS_7_ASN_ENCODING;

			let mut usages = self.usages.iter().map(|u| u.as_ptr()).collect::<Vec<_>>();
			let mut entry_data = vec![];
			let mut entries = vec![];
			for certificate in &self.certificates {
				let data = try!(cert_entry(certificate));
				entries.push(*(data.as_ptr() as *const winapi::CTL_ENTRY));
				entry_data.push(data);
			}

			let mut ctl_info: winapi::CTL_INFO = mem::zeroed();
			ctl_info.dwVersion = winapi::CTL_V1;
			ctl_info.SubjectUsage.cUsageIdentifier = usages.len() as winapi::DWORD;
			ctl_info.SubjectUsage.rgpszUsageIdentifier = usages.as_mut_ptr() as *mut winapi::LPSTR;
			let mut algorithm = winapi::szOID_OIWSEC_sha1.as_bytes().to_owned();
			algorithm.push(0);
			ctl_info.SubjectAlgorithm.pszObjId = algorithm.as_ptr() as winapi::LPSTR;
			ctl_info.cCTLEntry = entries.len() as winapi::DWORD;
			ctl_info.rgCTLEntry = entries.as_mut_ptr();

			let mut sign_info: winapi::CMSG_SIGNED_ENCODE_INFO = mem::zeroed();
			sign_info.cbSize = mem::size_of_val(&sign_info) as winapi::DWORD;
			let mut encoded_certs = self.certificates
										.iter()
										.map(|c| {
											winapi::CERT_BLOB {
												cbData: (*c.as_inner()).cbCertEncoded,
												pbData: (*c.as_inner()).pbCertEncoded,
											}
										})
										.collect::<Vec<_>>();
			sign_info.rgCertEncoded = encoded_certs.as_mut_ptr();
			sign_info.cCertEncoded = encoded_certs.len() as winapi::DWORD;

			let flags = winapi::CMSG_ENCODE_SORTED_CTL_FLAG |
				winapi::CMSG_ENCODE_HASHED_SUBJECT_IDENTIFIER_FLAG;

			let mut size = 0;

			let res = crypt32::CryptMsgEncodeAndSignCTL(
				encoding,
				&mut ctl_info,
				&mut sign_info,
				flags,
				ptr::null_mut(),
				&mut size);
			if res == winapi::FALSE {
				return Err(io::Error::last_os_error());
			}

			let mut encoded = vec![0; size as usize];

			let res = crypt32::CryptMsgEncodeAndSignCTL(
				encoding,
				&mut ctl_info,
				&mut sign_info,
				flags,
				encoded.as_mut_ptr() as *mut winapi::BYTE,
				&mut size);
			if res == winapi::FALSE {
				return Err(io::Error::last_os_error());
			}

			Ok(encoded)
		}
	}
}

fn cert_entry(cert: &CertContext) -> io::Result<Vec<u8>> {
	unsafe {
		let mut size = 0;

		let res = crypt32::CertCreateCTLEntryFromCertificateContextProperties(
			cert.as_inner(),
			0,
			ptr::null_mut(),
			winapi::CTL_ENTRY_FROM_PROP_CHAIN_FLAG,
			ptr::null_mut(),
			ptr::null_mut(),
			&mut size);
		if res == winapi::FALSE {
			return Err(io::Error::last_os_error());
		}

		let mut entry = vec![0; size as usize];
		let res = crypt32::CertCreateCTLEntryFromCertificateContextProperties(
			cert.as_inner(),
			0,
			ptr::null_mut(),
			winapi::CTL_ENTRY_FROM_PROP_CHAIN_FLAG,
			ptr::null_mut(),
			entry.as_mut_ptr() as winapi::PCTL_ENTRY,
			&mut size);
		if res == winapi::FALSE {
			Err(io::Error::last_os_error())
		} else {
			Ok(entry)
		}
	}
}

#[cfg(test)]
mod test {
	use std::io::Read;
	use std::fs::File;
	use cert_store::CertStore;
	use ctl_context::CtlContext;

	#[test]
	fn create_ctl() {
        let mut file = File::open(concat!(env!("CARGO_MANIFEST_DIR"),
                                          "/test/self-signed.badssl.com.cer"))
                           .unwrap();
        let mut cert = vec![];
        file.read_to_end(&mut cert).unwrap();

        let mut store = CertStore::memory().unwrap();
        let cert = store.add_encoded_certificate(&cert).unwrap();

        CtlContext::builder()
        	.certificate(cert)
            .usage("1.3.6.1.4.1.311.2.2.2")
            .encode_and_sign()
            .unwrap();
	}
}<|endoftext|>


<commit_before>#[link(
    name = "depviz",
    vers = "0.1.0",
    uuid = "aecc116d-c1bd-4847-9435-f020d565e181"
)];

#[author = "KokaKiwi <kokakiwi@kokakiwi.net>"];
#[license = "MIT"];

#[crate_type = "bin"];

#[feature(managed_boxes)];

extern mod extra;
extern mod syntax;

use std::os;
use std::path::Path;
use extra::getopts::groups;
use depviz::Node;

mod depviz;

static BRIEF_USAGE: &'static str = "Rust crate dependencies visualization.";

fn print_usage(program: &str, opts: &[groups::OptGroup])
{
    println!("Usage: {:s} [OPTIONS] FILE", program);
    println!("");
    print(groups::usage(BRIEF_USAGE, opts));
}

fn main_args(args: &[~str]) -> int
{
    let program = args[0].as_slice();

    let opts = ~[
        groups::optflag("h", "help", "Show this help and exit."),
    ];

    let matches = match groups::getopts(args, opts) {
        Ok(m) => m,
        Err(f) => fail!(f.to_err_msg()),
    };

    if matches.opts_present([~"h", ~"help"])
    {
        print_usage(program, opts);
        return 0;
    }

    if matches.free.len() <= 1
    {
        print_usage(program, opts);
        return 1;
    }

    let filename = matches.free[1];
    let path = Path::new(filename);
    let name = match path.filestem_str() {
        Some(s) => s.to_owned(),
        None => fail!(),
    };

    let root = depviz::construct::construct_crate(name, path);
    print_tree(root, ~"");

    return 0;
}

fn print_tree(node: &Node, indent: ~str)
{
    print(indent);
    println!("{}: {}", node.name, node.path.display());

    for child in node.children.iter()
    {
        print_tree(*child, indent + "    ")
    }
}

fn main()
{
    os::set_exit_status(main_args(os::args()));
}
<commit_msg>Oh str, come on. <3<commit_after>#[link(
    name = "depviz",
    vers = "0.1.0",
    uuid = "aecc116d-c1bd-4847-9435-f020d565e181"
)];

#[author = "KokaKiwi <kokakiwi@kokakiwi.net>"];
#[license = "MIT"];

#[crate_type = "bin"];

#[feature(managed_boxes)];

extern mod extra;
extern mod syntax;

use std::os;
use std::path::Path;
use extra::getopts::groups;
use depviz::Node;

mod depviz;

static BRIEF_USAGE: &'static str = "Rust crate dependencies visualization.";

fn print_usage(program: &str, opts: &[groups::OptGroup])
{
    println!("Usage: {:s} [OPTIONS] FILE", program);
    println!("");
    print(groups::usage(BRIEF_USAGE, opts));
}

fn main_args(args: &[~str]) -> int
{
    let program = args[0].as_slice();

    let opts = ~[
        groups::optflag("h", "help", "Show this help and exit."),
    ];

    let matches = match groups::getopts(args, opts) {
        Ok(m) => m,
        Err(f) => fail!(f.to_err_msg()),
    };

    if matches.opts_present([~"h", ~"help"])
    {
        print_usage(program, opts);
        return 0;
    }

    if matches.free.len() <= 1
    {
        print_usage(program, opts);
        return 1;
    }

    let filename = matches.free[1];
    let path = Path::new(filename);
    let name = match path.filestem_str() {
        Some(s) => s.to_owned(),
        None => fail!(),
    };

    let root = depviz::construct::construct_crate(name, path);
    print_tree(root, "");

    return 0;
}

fn print_tree(node: &Node, indent: &str)
{
    print(indent);
    println!("{}: {}", node.name, node.path.display());

    for child in node.children.iter()
    {
        print_tree(*child, indent + "    ")
    }
}

fn main()
{
    os::set_exit_status(main_args(os::args()));
}
<|endoftext|>


<commit_before>use std::cell::Cell;
use std::collections::HashMap;
use std::iter::repeat;

use tree::{Tree, Visibility, Visitor};

pub struct Config {
    pub colored: bool,
    pub show_conditional: bool,
    pub show_external: bool,
    pub show_types: bool,
}

pub struct DotPrinter<'a> {
    config: Config,
    tree: &'a Tree,
}

impl<'a> Visitor for DotPrinter<'a> {
    fn visit(&self, tree: &Tree, _: &[(usize, usize)], parents: &[&str]) {
        if !self.config.show_conditional {
            if let Tree::Module { ref condition, .. } = *tree {
                if condition.is_some() {
                    return;
                }
            }
        }

        let mut parent_name = "".to_string();

        for parent in parents.iter().skip(1) {
            parent_name += "::";
            parent_name += parent;
        }
        let name = parent_name.clone() + "::" + tree.name();
        print!("\t\"{}\" [label=\"{}\"", name, tree.name());

        if self.config.colored {
            let kind = match *tree {
                Tree::Crate { .. } => "color=green".to_string(),
                Tree::Module { ref visibility, .. } => match *visibility {
                    Visibility::Public => "color=green".to_string(),
                    Visibility::Private => "color=gold".to_string(),
                },
                Tree::Orphan { .. } => "color=red".to_string(),
            };
            print!(",{}", kind);
        }

        if let Tree::Module { ref condition, .. } = *tree {
            if condition.is_some() {
                print!(",style=dotted");
            }
        }

        println!("];");

        if !parents.is_empty() {
            if parent_name == "" {
                println!("\t\"::{}\" -> \"{}\" [weight=100,len=1];", parents[0], name);
            } else {
                println!("\t\"{}\" -> \"{}\" [weight=100,len=1];", parent_name, name);
            }
        }

        if let Tree::Module { ref uses, .. } = *tree {
            let mut used_modules = HashMap::new();
            for (visibility, use_name) in uses {
                let use_name = fix_supers(&use_name, &name);
                let use_name = fix_selfs(&use_name, &name);
                let use_name = make_absolute(&use_name);

                let use_data = UseModuleFinder::new(&use_name);
                self.tree.accept(&mut vec![], &mut vec![], &use_data);

                if let Some(module_name) = use_data.module.take() {
                    let (strip_str, module_name) = if module_name != "::" {
                        (module_name.clone() + "::", module_name)
                    } else {
                        (module_name, String::from("::") + parents[0])
                    };

                    let display_name = if use_name.to_string() + "::" == strip_str {
                        "self"
                    } else {
                        use_name.split(strip_str.as_str()).last().unwrap()
                    };
                    // Insert into used_modules hashmap
                    used_modules
                        .entry(module_name)
                        .or_insert_with(|| vec![])
                        .push((*visibility, display_name.to_string()));
                } else if self.config.show_external {
                    let attrs = if self.config.colored {
                        match visibility {
                            Visibility::Public => "color=green".to_string(),
                            Visibility::Private => "color=gold".to_string(),
                        }
                    } else {
                        "".to_string()
                    };
                    println!("\t\"{}\" -> \"{}\" [{}];", name, &use_name[2..], attrs);
                }
            }

            for (key, val) in used_modules {
                let types = if self.config.show_types {
                    val.iter()
                        .map(|(_, name)| name.as_ref())
                        .collect::<Vec<&str>>()
                        .join("\\n")
                } else {
                    "".to_string()
                };

                let attrs = if self.config.colored {
                    if val.iter().any(|(vis, _)| &Visibility::Public == vis) {
                        ",color=green"
                    } else {
                        ",color=gold"
                    }
                } else {
                    ""
                };

                println!(
                    "\t\"{}\" -> \"{}\" [penwidth={},label=\"{}\"{}];",
                    name,
                    key,
                    val.len(),
                    types,
                    attrs,
                );
            }
        }
    }
}

/// This method resolves relative paths that start with `super`
fn fix_supers(use_name: &str, name: &str) -> String {
    let mut use_name_iter = use_name.split("::");
    let mut opt_super = use_name_iter.next().unwrap();
    if opt_super == "super" {
        let name_segments = name.split("::").collect::<Vec<&str>>();
        let mut name_iter = name_segments.iter().rev();

        while opt_super == "super" {
            name_iter.next();

            if let Some(segment) = use_name_iter.next() {
                opt_super = segment;
            } else {
                opt_super = "";
            }
        }

        name_iter
            .rev()
            .map(|s| *s)
            .chain([opt_super].iter().cloned())
            .chain(use_name_iter)
            .collect::<Vec<&str>>()
            .join("::")
    } else {
        use_name.to_string()
    }
}

/// This method resolves relative paths that start with `self`
fn fix_selfs(use_name: &str, name: &str) -> String {
    let mut use_name_iter = use_name.split("::");
    let opt_self = use_name_iter.next().unwrap();
    if opt_self == "self" {
        [name]
            .iter()
            .map(|s| *s)
            .chain(use_name_iter)
            .collect::<Vec<&str>>()
            .join("::")
    } else {
        use_name.to_string()
    }
}

fn make_absolute(use_name: &str) -> String {
    if use_name.starts_with("::") {
        use_name.to_string()
    } else {
        "::".to_string() + use_name
    }
}

impl<'a> DotPrinter<'a> {
    pub fn new(config: Config, tree: &'a Tree) -> Self {
        DotPrinter { config, tree }
    }
}

/// Bit of a hack to determine if this use string is pointing to some internal type
/// or to some external type. And if it is internal we also find out the module it belongs to
struct UseModuleFinder<'a> {
    pub name: &'a str,
    pub module: Cell<Option<String>>,
}

impl<'a> UseModuleFinder<'a> {
    pub fn new(name: &'a str) -> Self {
        UseModuleFinder {
            name,
            module: Cell::new(None),
        }
    }
}

impl<'a> Visitor for UseModuleFinder<'a> {
    fn visit(&self, tree: &Tree, _: &[(usize, usize)], parents: &[&str]) {
        let tree_path: Vec<&str> = parents
            .iter()
            .chain([tree.name()].iter())
            .cloned()
            .collect();
        let last_index: usize = if parents.is_empty() {
            0
        } else {
            parents.len() - 1
        };
        let module_name = "::".to_string() + &tree_path[1..].join("::");

        for (index, (segment, parent)) in self
            .name
            .split("::")
            .zip(tree_path.iter().chain(repeat(&"")))
            .skip(1) // Skip the crate name
            .enumerate()
        {
            if index == last_index {
                if &segment != parent {
                    if index != 0 {
                        let old_value = self.module.take();

                        if old_value.is_none()
                            || (old_value.is_some()
                                && old_value.as_ref().unwrap().split("::").count() < index)
                        {
                            self.module.set(Some(module_name.to_string()));
                        } else {
                            self.module.set(old_value);
                        }
                    }
                } else {
                    self.module.set(Some(module_name.to_string()));
                }
            } else if &segment != parent {
                return;
            }
        }
    }
}
<commit_msg>Use cloned (clippy)<commit_after>use std::cell::Cell;
use std::collections::HashMap;
use std::iter::repeat;

use tree::{Tree, Visibility, Visitor};

pub struct Config {
    pub colored: bool,
    pub show_conditional: bool,
    pub show_external: bool,
    pub show_types: bool,
}

pub struct DotPrinter<'a> {
    config: Config,
    tree: &'a Tree,
}

impl<'a> Visitor for DotPrinter<'a> {
    fn visit(&self, tree: &Tree, _: &[(usize, usize)], parents: &[&str]) {
        if !self.config.show_conditional {
            if let Tree::Module { ref condition, .. } = *tree {
                if condition.is_some() {
                    return;
                }
            }
        }

        let mut parent_name = "".to_string();

        for parent in parents.iter().skip(1) {
            parent_name += "::";
            parent_name += parent;
        }
        let name = parent_name.clone() + "::" + tree.name();
        print!("\t\"{}\" [label=\"{}\"", name, tree.name());

        if self.config.colored {
            let kind = match *tree {
                Tree::Crate { .. } => "color=green".to_string(),
                Tree::Module { ref visibility, .. } => match *visibility {
                    Visibility::Public => "color=green".to_string(),
                    Visibility::Private => "color=gold".to_string(),
                },
                Tree::Orphan { .. } => "color=red".to_string(),
            };
            print!(",{}", kind);
        }

        if let Tree::Module { ref condition, .. } = *tree {
            if condition.is_some() {
                print!(",style=dotted");
            }
        }

        println!("];");

        if !parents.is_empty() {
            if parent_name == "" {
                println!("\t\"::{}\" -> \"{}\" [weight=100,len=1];", parents[0], name);
            } else {
                println!("\t\"{}\" -> \"{}\" [weight=100,len=1];", parent_name, name);
            }
        }

        if let Tree::Module { ref uses, .. } = *tree {
            let mut used_modules = HashMap::new();
            for (visibility, use_name) in uses {
                let use_name = fix_supers(&use_name, &name);
                let use_name = fix_selfs(&use_name, &name);
                let use_name = make_absolute(&use_name);

                let use_data = UseModuleFinder::new(&use_name);
                self.tree.accept(&mut vec![], &mut vec![], &use_data);

                if let Some(module_name) = use_data.module.take() {
                    let (strip_str, module_name) = if module_name != "::" {
                        (module_name.clone() + "::", module_name)
                    } else {
                        (module_name, String::from("::") + parents[0])
                    };

                    let display_name = if use_name.to_string() + "::" == strip_str {
                        "self"
                    } else {
                        use_name.split(strip_str.as_str()).last().unwrap()
                    };
                    // Insert into used_modules hashmap
                    used_modules
                        .entry(module_name)
                        .or_insert_with(|| vec![])
                        .push((*visibility, display_name.to_string()));
                } else if self.config.show_external {
                    let attrs = if self.config.colored {
                        match visibility {
                            Visibility::Public => "color=green".to_string(),
                            Visibility::Private => "color=gold".to_string(),
                        }
                    } else {
                        "".to_string()
                    };
                    println!("\t\"{}\" -> \"{}\" [{}];", name, &use_name[2..], attrs);
                }
            }

            for (key, val) in used_modules {
                let types = if self.config.show_types {
                    val.iter()
                        .map(|(_, name)| name.as_ref())
                        .collect::<Vec<&str>>()
                        .join("\\n")
                } else {
                    "".to_string()
                };

                let attrs = if self.config.colored {
                    if val.iter().any(|(vis, _)| &Visibility::Public == vis) {
                        ",color=green"
                    } else {
                        ",color=gold"
                    }
                } else {
                    ""
                };

                println!(
                    "\t\"{}\" -> \"{}\" [penwidth={},label=\"{}\"{}];",
                    name,
                    key,
                    val.len(),
                    types,
                    attrs,
                );
            }
        }
    }
}

/// This method resolves relative paths that start with `super`
fn fix_supers(use_name: &str, name: &str) -> String {
    let mut use_name_iter = use_name.split("::");
    let mut opt_super = use_name_iter.next().unwrap();
    if opt_super == "super" {
        let name_segments = name.split("::").collect::<Vec<&str>>();
        let mut name_iter = name_segments.iter().rev();

        while opt_super == "super" {
            name_iter.next();

            if let Some(segment) = use_name_iter.next() {
                opt_super = segment;
            } else {
                opt_super = "";
            }
        }

        name_iter
            .rev()
            .cloned()
            .chain([opt_super].iter().cloned())
            .chain(use_name_iter)
            .collect::<Vec<&str>>()
            .join("::")
    } else {
        use_name.to_string()
    }
}

/// This method resolves relative paths that start with `self`
fn fix_selfs(use_name: &str, name: &str) -> String {
    let mut use_name_iter = use_name.split("::");
    let opt_self = use_name_iter.next().unwrap();
    if opt_self == "self" {
        [name]
            .iter()
            .cloned()
            .chain(use_name_iter)
            .collect::<Vec<&str>>()
            .join("::")
    } else {
        use_name.to_string()
    }
}

fn make_absolute(use_name: &str) -> String {
    if use_name.starts_with("::") {
        use_name.to_string()
    } else {
        "::".to_string() + use_name
    }
}

impl<'a> DotPrinter<'a> {
    pub fn new(config: Config, tree: &'a Tree) -> Self {
        DotPrinter { config, tree }
    }
}

/// Bit of a hack to determine if this use string is pointing to some internal type
/// or to some external type. And if it is internal we also find out the module it belongs to
struct UseModuleFinder<'a> {
    pub name: &'a str,
    pub module: Cell<Option<String>>,
}

impl<'a> UseModuleFinder<'a> {
    pub fn new(name: &'a str) -> Self {
        UseModuleFinder {
            name,
            module: Cell::new(None),
        }
    }
}

impl<'a> Visitor for UseModuleFinder<'a> {
    fn visit(&self, tree: &Tree, _: &[(usize, usize)], parents: &[&str]) {
        let tree_path: Vec<&str> = parents
            .iter()
            .chain([tree.name()].iter())
            .cloned()
            .collect();
        let last_index: usize = if parents.is_empty() {
            0
        } else {
            parents.len() - 1
        };
        let module_name = "::".to_string() + &tree_path[1..].join("::");

        for (index, (segment, parent)) in self
            .name
            .split("::")
            .zip(tree_path.iter().chain(repeat(&"")))
            .skip(1) // Skip the crate name
            .enumerate()
        {
            if index == last_index {
                if &segment != parent {
                    if index != 0 {
                        let old_value = self.module.take();

                        if old_value.is_none()
                            || (old_value.is_some()
                                && old_value.as_ref().unwrap().split("::").count() < index)
                        {
                            self.module.set(Some(module_name.to_string()));
                        } else {
                            self.module.set(old_value);
                        }
                    }
                } else {
                    self.module.set(Some(module_name.to_string()));
                }
            } else if &segment != parent {
                return;
            }
        }
    }
}
<|endoftext|>


<commit_before>pub const MAX_CHANNELS: usize = 8;

pub struct Frame {
  pub header: Header,
  pub footer: Footer,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum ChannelAssignment {
  Independent = 0,
  LeftSide    = 1,
  RightSide   = 2,
  MiddleSide  = 3,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum NumberType {
  Frame(u32),
  Sample(u64),
}

#[derive(Debug, PartialEq, Eq)]
pub struct Header {
  pub block_size: u32,
  pub sample_rate: u32,
  pub channel_assignment: ChannelAssignment,
  pub bits_per_sample: usize,
  pub number: NumberType,
  pub crc: u8,
}

#[derive(Debug, PartialEq, Eq)]
pub struct Footer(pub u16);
<commit_msg>Add channel member to Header<commit_after>pub const MAX_CHANNELS: usize = 8;

pub struct Frame {
  pub header: Header,
  pub footer: Footer,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum ChannelAssignment {
  Independent = 0,
  LeftSide    = 1,
  RightSide   = 2,
  MiddleSide  = 3,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum NumberType {
  Frame(u32),
  Sample(u64),
}

#[derive(Debug, PartialEq, Eq)]
pub struct Header {
  pub block_size: u32,
  pub sample_rate: u32,
  pub channels: u8,
  pub channel_assignment: ChannelAssignment,
  pub bits_per_sample: usize,
  pub number: NumberType,
  pub crc: u8,
}

#[derive(Debug, PartialEq, Eq)]
pub struct Footer(pub u16);
<|endoftext|>


<commit_before>/*
 Simplex module
 Reference implementation:
    http://webstaff.itn.liu.se/~stegu/simplexnoise/SimplexNoise.java
*/
use std::rand;
use std::rand::{Rng, SeedableRng};
use lcgrng::LCG;
use gen::NoiseGen;

pub struct Simplex {
    seed: u32,
    zoom: f64,
    // 512 to remove the need for bounding array indicies
    perm: [uint, ..256]
}

impl Simplex {
    pub fn new_rand(zoom: f64) -> Simplex {
        let mut rng = rand::task_rng(); // for getting a random seed
        let mut simp = Simplex { seed: rng.gen(), zoom: 1.0/zoom, perm: [0, ..256] };
        simp.init_perm();
        simp
    }

    pub fn from_seed(seed: u32, zoom: f64) -> Simplex {
        let mut simp = Simplex { seed: seed, zoom: 1.0/zoom, perm: [0, ..256] };
        simp.init_perm();
        simp
    }

    pub fn set_seed(&mut self, seed: u32) -> () {
        self.seed = seed;
        self.init_perm();
    }

    pub fn get_seed(&mut self) -> u32 {
        self.seed
    }

    fn init_perm(&mut self) -> () {
        let mut rng: LCG = SeedableRng::from_seed(self.seed);

        for i in range(0, 256) {
            self.perm[i] = i as uint;
        }
        rng.shuffle(self.perm);
    }

    fn dot3(g: &[int, ..3], x: f64, y: f64, z: f64) -> f64 {
        ((g[0] as f64)*x) + ((g[1] as f64)*y) + ((g[2] as f64)*z)
    }

    fn dot2(g: &[int, ..2], x: f64, y: f64) -> f64 {
        ((g[0] as f64)*x) + ((g[1] as f64)*y)
    }
}

impl NoiseGen for Simplex {
    fn get_value2d(&mut self, xx: f64, yy: f64) -> f64 {
        let x = xx * self.zoom;
        let y = yy * self.zoom;
        // sqrt(3) = 1.7320508075688772935274463415059
        static sqrt3: f64 = 1.7320508075688772935274463415059;
        static F2: f64 = 0.5*(sqrt3-1.0);
        static G2: f64 = (3.0-sqrt3)/6.0;

        static grad:[[int, ..2], ..12] = [[1,1],[-1,1],[1,-1],[-1,-1],
                                          [1,0],[-1,0],[1,0],[-1,0],
                                          [0,1],[0,-1],[0,1],[0,-1]];

        let mut n0: f64;
        let mut n1: f64;
        let mut n2: f64;

        let s: f64 = (x+y)*F2;
        let i = (x+s).floor() as int;
        let j = (y+s).floor() as int;

        let t: f64 = ((i+j) as f64)*G2;
        let x0: f64 = x-((i as f64)-t);
        let y0: f64 = y-((j as f64)-t);

        let mut i1: uint;
        let mut j1: uint;

        if x0 > y0 {
            i1 = 1;
            j1 = 0;
        } else {
            i1 = 0;
            j1 = 1;
        }

        let x1: f64 = x0-(i1 as f64)+G2;
        let y1: f64 = y0-(j1 as f64)+G2;
        let x2: f64 = x0-1.0+2.0*G2;
        let y2: f64 = y0-1.0+2.0*G2;

        let ii = (i & 0xff) as uint;
        let jj = (j & 0xff) as uint;

        let gi0 = (self.perm[(ii+   self.perm[(jj   )&255])&255] % 12) as uint;
        let gi1 = (self.perm[(ii+i1+self.perm[(jj+j1)&255])&255] % 12) as uint;
        let gi2 = (self.perm[(ii+1u+self.perm[(jj+1u)&255])&255] % 12) as uint;

        let t0: f64 = 0.5-(x0*x0)-(y0*y0);
        if t0 < 0.0 {
            n0 = 0.0;
        } else {
            n0 = t0*t0*t0*t0*Simplex::dot2(&grad[gi0], x0, y0);
        }

        let t1: f64 = 0.5-(x1*x1)-(y1*y1);
        if t1 < 0.0 {
            n1 = 0.0;
        } else {
            n1 = t1*t1*t1*t1*Simplex::dot2(&grad[gi1], x1, y1);
        }

        let t2: f64 = 0.5-(x2*x2)-(y2*y2);
        if t2 < 0.0 {
            n2 = 0.0;
        } else {
            n2 = t2*t2*t2*t2*Simplex::dot2(&grad[gi2], x2, y2);
        }

        70.0*(n0+n1+n2)
    }

    fn get_value3d(&mut self, xx: f64, yy: f64, zz: f64) -> f64 {
        let x = xx * self.zoom;
        let y = yy * self.zoom;
        let z = zz * self.zoom;

        static F3: f64 = 1.0/3.0;
        static G3: f64 = 1.0/6.0;

        static grad: [[int, ..3], ..12] = [[1,1,0],[-1,1,0],[1,-1,0],[-1,-1,0],
                                           [1,0,1],[-1,0,1],[1,0,-1],[-1,0,-1],
                                           [0,1,1],[0,-1,1],[0,1,-1],[0,-1,-1]];

        let mut n0: f64;
        let mut n1: f64;
        let mut n2: f64;
        let mut n3: f64;

        let s: f64 = (x+y+z)*F3;
        let i = (x+s).floor() as int;
        let j = (y+s).floor() as int;
        let k = (z+s).floor() as int;

        let t: f64 = ((i+j+k) as f64)*G3;
        let x0: f64 = x-((i as f64)-t);
        let y0: f64 = y-((j as f64)-t);
        let z0: f64 = z-((k as f64)-t);

        let mut i1: uint;
        let mut j1: uint;
        let mut k1: uint;
        let mut i2: uint;
        let mut j2: uint;
        let mut k2: uint;

        if x0 >= y0 {
            if y0 >= z0 {
                i1=1; j1=0; k1=0; i2=1; j2=1; k2=0;
            }else if x0 >= z0 {
                i1=1; j1=0; k1=0; i2=1; j2=0; k2=1;
            }else {
                i1=0; j1=0; k1=1; i2=1; j2=0; k2=1;
            }
        }else { // x0 < y0
            if y0 < z0 {
                i1=0; j1=0; k1=1; i2=0; j2=1; k2=1;
            }else if x0 < z0 {
                i1=0; j1=1; k1=0; i2=0; j2=1; k2=1;
            }else {
                i1=0; j1=1; k1=0; i2=1; j2=1; k2=0;
            }
        }

        let x1: f64 = x0 - (i1 as f64) + G3;
        let y1: f64 = y0 - (j1 as f64) + G3;
        let z1: f64 = z0 - (k1 as f64) + G3;
        let x2: f64 = x0 - (i2 as f64) + (2.0*G3);
        let y2: f64 = y0 - (j2 as f64) + (2.0*G3);
        let z2: f64 = z0 - (k2 as f64) + (2.0*G3);
        let x3: f64 = x0 - 1.0 + (3.0*G3);
        let y3: f64 = y0 - 1.0 + (3.0*G3);
        let z3: f64 = z0 - 1.0 + (3.0*G3);

        let ii = (i & 0xff) as uint;
        let jj = (j & 0xff) as uint;
        let kk = (k & 0xff) as uint;

        let gi0 = (self.perm[(ii+   self.perm[(jj+   self.perm[ kk        ])&255])&255]%12) as uint;
        let gi1 = (self.perm[(ii+i1+self.perm[(jj+j1+self.perm[(kk+k1)&255])&255])&255]%12) as uint;
        let gi2 = (self.perm[(ii+i2+self.perm[(jj+j2+self.perm[(kk+k2)&255])&255])&255]%12) as uint;
        let gi3 = (self.perm[(ii+1u+self.perm[(jj+1u+self.perm[(kk+1u)&255])&255])&255]%12) as uint;

        let t0 = 0.5-(x0*x0)-(y0*y0)-(z0*z0);
        if t0 < 0.0 {
            n0 = 0.0;
        } else {
            n0 = t0*t0*t0*t0*Simplex::dot3(&grad[gi0], x0, y0, z0);
        }

        let t1 = 0.5-(x1*x1)-(y1*y1)-(z1*z1);
        if t1 < 0.0 {
            n1 = 0.0;
        } else {
            n1 = t1*t1*t1*t1*Simplex::dot3(&grad[gi1], x1, y1, z1);
        }

        let t2 = 0.5-(x2*x2)-(y2*y2)-(z2*z2);
        if t2 < 0.0 {
            n2 = 0.0;
        } else {
            n2 = t2*t2*t2*t2*Simplex::dot3(&grad[gi2], x2, y2, z2);
        }

        let t3 = 0.5-(x3*x3)-(y3*y3)-(z3*z3);
        if t3 < 0.0 {
            n3 = 0.0;
        } else {
            n3 = t3*t3*t3*t3*Simplex::dot3(&grad[gi3], x3, y3, z3);
        }

        32.0*(n0+n1+n2+n3)
    }
}
<commit_msg>Now use 2D OpenSimplex noise<commit_after>/*
 Simplex module
 Reference implementation:
    http://webstaff.itn.liu.se/~stegu/simplexnoise/SimplexNoise.java
*/
use std::rand;
use std::rand::{Rng, SeedableRng};
use lcgrng::LCG;
use gen::NoiseGen;

pub struct Simplex {
    seed: u32,
    zoom: f64,
    // 512 to remove the need for bounding array indicies
    perm: [u8, ..256]
}

impl Simplex {
    pub fn new_rand(zoom: f64) -> Simplex {
        let mut rng = rand::task_rng(); // for getting a random seed
        let mut simp = Simplex { seed: rng.gen(), zoom: 1.0/zoom, perm: [0, ..256] };
        simp.init_perm();
        simp
    }

    pub fn from_seed(seed: u32, zoom: f64) -> Simplex {
        let mut simp = Simplex { seed: seed, zoom: 1.0/zoom, perm: [0, ..256] };
        simp.init_perm();
        simp
    }

    pub fn set_seed(&mut self, seed: u32) -> () {
        self.seed = seed;
        self.init_perm();
    }

    pub fn get_seed(&mut self) -> u32 {
        self.seed
    }

    fn init_perm(&mut self) -> () {
        let mut rng: LCG = SeedableRng::from_seed(self.seed);

        for i in range(0, 256) {
            self.perm[i] = i as u8;
        }
        rng.shuffle(self.perm);
    }

    fn extrapolate_2d(&mut self, xsb: int, ysb: int, dx: f64, dy: f64) -> f64 {
        static GRAD_2D: [i8, ..16] = [
             5,  2,    2,  5,
            -5,  2,   -2,  5,
             5, -2,    2, -5,
            -5, -2,   -2, -5,
        ];

        let xsb_idx = (xsb & 0xFF) as uint;
        let ysb_idx = (ysb & 0xFF) as uint;
        let idx = (self.perm[((self.perm[xsb_idx] as uint) + ysb_idx) & 0xFF] & 0x0E) as uint;

        ((GRAD_2D[idx] as f64) * dx) + ((GRAD_2D[idx+1u] as f64) * dy)
    }
}

impl NoiseGen for Simplex {
    // OpenSimplex implimentation: https://gist.github.com/KdotJPG/b1270127455a94ac5d19
    fn get_value2d(&mut self, x: f64, y: f64) -> f64 {
        static STRETCH_CONSTANT: f64 = -0.211324865405187;
        static SQUISH_CONSTANT: f64 = 0.366025403784439;
        static NORM_CONSTANT: f64 = 47.0;

        let x = x * self.zoom;
        let y = y * self.zoom;

        let stretch_offset = (x + y) * STRETCH_CONSTANT;
        let xs = x + stretch_offset;
        let ys = y + stretch_offset;

        // Floor to get grid coordinates of rhombus (stretched square) super cell origin
        let xsb = xs.floor() as int;
        let ysb = ys.floor() as int;

        // Skew out to get actual coordinates of rhombus origin
        let squish_offset = ((xsb + ysb) as f64) * SQUISH_CONSTANT;
        let xb = (xsb as f64) + squish_offset;
        let yb = (ysb as f64) + squish_offset;

        // Computer grid coordinates relative to rhombus origin
        let xins = xs - (xsb as f64);
        let yins = ys - (ysb as f64);

        // Sum those together to get a value that determines which region we're in
        let in_sum = xins + yins;

        // Position relative to origin point
        let dx0 = x - xb;
        let dy0 = y - yb;

        // Contribution (1,0)
        let dx1 = dx0 - 1.0 - SQUISH_CONSTANT;
        let dy1 = dy0 - 0.0 - SQUISH_CONSTANT;
        let attn1 = 2.0 - (dx1 * dx1) - (dy1 * dy1);

        let v1: f64 = if attn1 > 0.0 {
            let attn1_4 = attn1*attn1*attn1*attn1;
            attn1_4*self.extrapolate_2d(xsb+1, ysb+0, dx1, dy1)
        } else {
            0.0
        };

        // Contribution (0,1)
        let dx2 = dx0 - 0.0 - SQUISH_CONSTANT;
        let dy2 = dy0 - 1.0 - SQUISH_CONSTANT;
        let attn2 = 2.0 - (dx2 * dx2) - (dy2 * dy2);

        let v2: f64 = if attn2 > 0.0 {
            let attn2_4 = attn2*attn2*attn2*attn2;
            attn2_4*self.extrapolate_2d(xsb+0, ysb+1, dx2, dy2)
        } else {
            0.0
        };

        let (xsv_ext, ysv_ext, dx_ext, dy_ext) = if in_sum <= 1.0 { // We're inside the triangle (2-Simplex) at (0,0)
            let zins = 1.0 - in_sum;
            if zins > xins || zins > yins { // (0,0) is one of the closest two triangular vertices
                if xins > yins {
                    (xsb+1, ysb-1, dx0-1.0, dy0+1.0)
                } else {
                    (xsb-1, ysb+1, dx0+1.0, dy0-1.0)
                }
            }else { // (1,0) and (0,1) are the closest two vertices.
                (xsb+1, ysb+1, dx0-1.0-(2.0*SQUISH_CONSTANT), dy0-1.0-(2.0*SQUISH_CONSTANT))
            }
        }else { // We're inside the triangle (2-Simplex) at (1,1)
            let zins = 2.0 - in_sum;
            if zins < xins || zins < yins { // (0,0) is one of the closest two triangular vertices
                if xins > yins {
                    (xsb+2, ysb+0, dx0-2.0-(2.0*SQUISH_CONSTANT), dy0+0.0-(2.0*SQUISH_CONSTANT))
                }else {
                    (xsb+0, ysb+2, dx0-0.0-(2.0*SQUISH_CONSTANT), dy0-2.0-(2.0*SQUISH_CONSTANT))
                }
            }else { // (1,0) and (0,1) are the closest two vertices.
                (xsb, ysb, dx0, dy0)
            }
        };
        // In the original implementation this was done in the above,
        // but I think this is nicer
        let (xsb, ysb, dx0, dy0) = if in_sum <= 1.0 {
            (xsb, ysb, dx0, dy0)
        }else {
            (xsb + 1,
             ysb + 1,
             dx0 - 1.0 - (2.0 * SQUISH_CONSTANT),
             dy0 - 1.0 - (2.0 * SQUISH_CONSTANT))
        };

        // Contribution (0,0) or (1,1)
        let attn0 = 2.0 - (dx0 * dx0) - (dy0 * dy0);
        let v0: f64 = if attn0 > 0.0 {
            let attn0_4 = attn0*attn0*attn0*attn0;
            attn0_4*self.extrapolate_2d(xsb, ysb, dx0, dy0)
        } else {
            0.0
        };

        // Extra vertex
        let attn_ext = 2.0 - (dx_ext * dx_ext) - (dy_ext * dy_ext);
        let v_ext: f64 = if attn_ext > 0.0 {
            let attn_ext_4 = attn_ext*attn_ext*attn_ext*attn_ext;
            attn_ext_4*self.extrapolate_2d(xsv_ext, ysv_ext, dx_ext, dy_ext)
        } else {
            0.0
        };

        (v0 + v1 + v2 + v_ext) / NORM_CONSTANT
    }

    fn get_value3d(&mut self, x: f64, y: f64, z: f64) -> f64 {
        // TODO: All of this
        x + y + z
    }
}
<|endoftext|>


<commit_before>
use monster::incubation::FindAndTake;

mod youtube;
mod soundcloud;
mod twitch;

use std::vec::Vec;

macro_rules! regex(
    ($s:expr) => (regex::Regex::new($s).unwrap());
);

struct Module {
    checker: Box<Fn(i32) -> bool>,
    action: Box<Fn(&mut Registry, i32) -> i32>
}

struct Registry {
    modules: Vec<Module>
}

impl Registry {
	pub fn new() -> Registry{
        Registry {modules: Vec::new()}
    }

    fn register(&mut self, module: Module) {
        self.modules.push(module);
    }

    fn handle_some_data(&mut self, data: i32) {
	    let module = self.modules.find_and_take(|&(_, module)| (module.checker)(data));

	    (module.action)(self, data);

	    self.modules.push(module);
	}
}

pub fn init_handlers() -> Registry {
	let mut registry = Registry::new();

	registry
}<commit_msg>Fixes<commit_after>
use monster::incubation::FindAndTake;

mod youtube;
mod soundcloud;
mod twitch;

use std::vec::Vec;

macro_rules! regex(
    ($s:expr) => (regex::Regex::new($s).unwrap());
);

pub struct Module {
    checker: Box<Fn(i32) -> bool>,
    action: Box<Fn(&mut Registry, i32) -> i32>
}

pub struct Registry {
    modules: Vec<Module>
}

impl Registry {
	pub fn new() -> Registry{
        Registry {modules: Vec::new()}
    }

    fn register(&mut self, module: Module) {
        self.modules.push(module);
    }

    fn handle_some_data(&mut self, data: i32) {
	    let module = self.modules.find_and_take(|module| (module.checker)(data));

        if let Some(module) = module {
	       (module.action)(self, data);
	       self.modules.push(module);
        }
	}
}

pub fn init_handlers() -> Registry {
	let mut registry = Registry::new();

	registry
}<|endoftext|>


<commit_before>use byteorder::{ReadBytesExt, LittleEndian};
use std::io::{Read, Seek, SeekFrom};

use color::ColorType;
use image::{DecodingResult, ImageResult, ImageDecoder, ImageError};

use bmp::BMPDecoder;
use png::PNGDecoder;
use self::InnerDecoder::*;

// http://www.w3.org/TR/PNG-Structure.html
// The first eight bytes of a PNG file always contain the following (decimal) values:
const PNG_SIGNATURE: [u8; 8] = [137, 80, 78, 71, 13, 10, 26, 10];

/// An ico decoder
pub struct ICODecoder<R: Read> {
    selected_entry: DirEntry,
    inner_decoder: InnerDecoder<R>,
}

enum InnerDecoder<R: Read> {
    BMP(BMPDecoder<R>),
    PNG(PNGDecoder<R>)
}

#[derive(Clone, Copy, Default)]
struct DirEntry {
    width: u8,
    height: u8,
    color_count: u8,
    reserved: u8,

    num_color_planes: u16,
    bits_per_pixel: u16,

    image_length: u32,
    image_offset: u32,
}

impl<R: Read + Seek> ICODecoder<R> {
    /// Create a new decoder that decodes from the stream ```r```
    pub fn new(mut r: R) -> ImageResult<ICODecoder<R>> {
        let entries = try!(read_entries(&mut r));
        let entry = try!(best_entry(entries));
        let decoder = try!(entry.decoder(r));

        Ok(ICODecoder {
            selected_entry: entry,
            inner_decoder: decoder,
        })
    }
}

fn read_entries<R: Read>(r: &mut R) -> ImageResult<Vec<DirEntry>> {
    let _reserved = try!(r.read_u16::<LittleEndian>());
    let _type = try!(r.read_u16::<LittleEndian>());
    let count = try!(r.read_u16::<LittleEndian>());
    (0..count).map(|_| read_entry(r)).collect()
}

fn read_entry<R: Read>(r: &mut R) -> ImageResult<DirEntry> {
    let mut entry = DirEntry::default();

    entry.width = try!(r.read_u8());
    entry.height = try!(r.read_u8());
    entry.color_count = try!(r.read_u8());
    entry.reserved = try!(r.read_u8());

    entry.num_color_planes = try!(r.read_u16::<LittleEndian>());
    entry.bits_per_pixel = try!(r.read_u16::<LittleEndian>());

    entry.image_length = try!(r.read_u32::<LittleEndian>());
    entry.image_offset = try!(r.read_u32::<LittleEndian>());

    Ok(entry)
}

/// Find the entry with the highest (color depth, size).
fn best_entry(mut entries: Vec<DirEntry>) -> ImageResult<DirEntry> {
    let mut best = try!(entries.pop().ok_or(ImageError::ImageEnd));
    let mut best_score = (best.bits_per_pixel, best.real_width() as u32 * best.real_height() as u32);

    for entry in entries.into_iter() {
        let score = (entry.bits_per_pixel, entry.real_width() as u32 * entry.real_height() as u32);
        if score > best_score {
            best = entry;
            best_score = score;
        }
    }
    Ok(best)
}

impl DirEntry {
    fn real_width(&self) -> u16 {
        match self.width {
            0 => 256,
            w => w as u16
        }
    }

    fn real_height(&self) -> u16 {
        match self.height {
            0 => 256,
            h => h as u16
        }
    }

    fn seek_to_start<R: Read + Seek>(&self, r: &mut R) -> ImageResult<()> {
        try!(r.seek(SeekFrom::Start(self.image_offset as u64)));
        Ok(())
    }

    fn is_png<R: Read + Seek>(&self, r: &mut R) -> ImageResult<bool> {
        try!(self.seek_to_start(r));

        // Read the first 8 bytes to sniff the image.
        let mut signature = [0u8; 8];
        try!(r.read_exact(&mut signature));

        Ok(signature == PNG_SIGNATURE)
    }

    fn decoder<R: Read + Seek>(&self, mut r: R) -> ImageResult<InnerDecoder<R>> {
        let is_png = try!(self.is_png(&mut r));
        try!(self.seek_to_start(&mut r));

        if is_png {
            Ok(PNG(PNGDecoder::new(r)))
        } else {
            let mut decoder = BMPDecoder::new(r);
            try!(decoder.read_metadata_in_ico_format());
            Ok(BMP(decoder))
        }
    }
}

impl<R: Read + Seek> ImageDecoder for ICODecoder<R> {
    fn dimensions(&mut self) -> ImageResult<(u32, u32)> {
        match self.inner_decoder {
            BMP(ref mut decoder) => decoder.dimensions(),
            PNG(ref mut decoder) => decoder.dimensions()
        }
    }

    fn colortype(&mut self) -> ImageResult<ColorType> {
        match self.inner_decoder {
            BMP(ref mut decoder) => decoder.colortype(),
            PNG(ref mut decoder) => decoder.colortype()
        }
    }

    fn row_len(&mut self) -> ImageResult<usize> {
        match self.inner_decoder {
            BMP(ref mut decoder) => decoder.row_len(),
            PNG(ref mut decoder) => decoder.row_len()
        }
    }

    fn read_scanline(&mut self, buf: &mut [u8]) -> ImageResult<u32> {
        match self.inner_decoder {
            BMP(ref mut decoder) => decoder.read_scanline(buf),
            PNG(ref mut decoder) => decoder.read_scanline(buf)
        }
    }

    fn read_image(&mut self) -> ImageResult<DecodingResult> {
        match self.inner_decoder {
            PNG(ref mut decoder) => {
                decoder.read_image()
            }
            BMP(ref mut decoder) => {
                let (width, height) = try!(decoder.dimensions());

                // The ICO decoder needs an alpha chanel to apply the AND mask.
                if try!(decoder.colortype()) != ColorType::RGBA(8) {
                    return Err(ImageError::UnsupportedError("Unsupported color type".to_string()))
                }
                let mut pixel_data = match try!(decoder.read_image()) {
                    DecodingResult::U8(v) => v,
                    _ => unreachable!()
                };

                // If there's an AND mask following the image, read and apply it.
                let r = decoder.reader();
                let mask_start = try!(r.seek(SeekFrom::Current(0)));
                let mask_end = (self.selected_entry.image_offset + self.selected_entry.image_length) as u64;
                let mask_length = mask_end - mask_start;

                if mask_length > 0 {
                    // A mask row contains 1 bit per pixel, padded to 4 bytes.
                    let mask_row_bytes = ((width + 31) / 32) * 4;
                    let expected_length = (mask_row_bytes * height) as u64;
                    if mask_length < expected_length {
                        return Err(ImageError::ImageEnd)
                    }

                    for y in 0..height {
                        let mut x = 0;
                        for _ in 0..mask_row_bytes {
                            // Apply the bits of each byte until we reach the end of the row.
                            let mask_byte = try!(r.read_u8());
                            for bit in (0..8).rev() {
                                if x >= width { break }
                                if mask_byte & (1 << bit) != 0 {
                                    // Set alpha channel to transparent.
                                    pixel_data[((height - y - 1) * width + x) as usize * 4 + 3] = 0;
                                }
                                x += 1;
                            }
                        }
                    }
                }
                Ok(DecodingResult::U8(pixel_data))
            }
        }
    }
}
<commit_msg>Add more checks in the ICO decoder<commit_after>use byteorder::{ReadBytesExt, LittleEndian};
use std::io::{Read, Seek, SeekFrom};

use color::ColorType;
use image::{DecodingResult, ImageResult, ImageDecoder, ImageError};

use bmp::BMPDecoder;
use png::PNGDecoder;
use self::InnerDecoder::*;

// http://www.w3.org/TR/PNG-Structure.html
// The first eight bytes of a PNG file always contain the following (decimal) values:
const PNG_SIGNATURE: [u8; 8] = [137, 80, 78, 71, 13, 10, 26, 10];

/// An ico decoder
pub struct ICODecoder<R: Read> {
    selected_entry: DirEntry,
    inner_decoder: InnerDecoder<R>,
}

enum InnerDecoder<R: Read> {
    BMP(BMPDecoder<R>),
    PNG(PNGDecoder<R>)
}

#[derive(Clone, Copy, Default)]
struct DirEntry {
    width: u8,
    height: u8,
    color_count: u8,
    reserved: u8,

    num_color_planes: u16,
    bits_per_pixel: u16,

    image_length: u32,
    image_offset: u32,
}

impl<R: Read + Seek> ICODecoder<R> {
    /// Create a new decoder that decodes from the stream ```r```
    pub fn new(mut r: R) -> ImageResult<ICODecoder<R>> {
        let entries = try!(read_entries(&mut r));
        let entry = try!(best_entry(entries));
        let decoder = try!(entry.decoder(r));

        Ok(ICODecoder {
            selected_entry: entry,
            inner_decoder: decoder,
        })
    }
}

fn read_entries<R: Read>(r: &mut R) -> ImageResult<Vec<DirEntry>> {
    let _reserved = try!(r.read_u16::<LittleEndian>());
    let _type = try!(r.read_u16::<LittleEndian>());
    let count = try!(r.read_u16::<LittleEndian>());
    (0..count).map(|_| read_entry(r)).collect()
}

fn read_entry<R: Read>(r: &mut R) -> ImageResult<DirEntry> {
    let mut entry = DirEntry::default();

    entry.width = try!(r.read_u8());
    entry.height = try!(r.read_u8());
    entry.color_count = try!(r.read_u8());
    // Reserved value (not used)
    entry.reserved = try!(r.read_u8());

    // This may be either the number of color planes (0 or 1), or the horizontal coordinate
    // of the hotspot for CUR files.
    entry.num_color_planes = try!(r.read_u16::<LittleEndian>());
    if entry.num_color_planes > 256 {
        return Err(ImageError::FormatError(
            "ICO image entry has a too large color planes/hotspot value".to_string()
        ));
    }

    // This may be either the bit depth (may be 0 meaning unspecified),
    // or the vertical coordinate of the hotspot for CUR files.
    entry.bits_per_pixel = try!(r.read_u16::<LittleEndian>());
    if entry.bits_per_pixel > 256 {
        return Err(ImageError::FormatError(
            "ICO image entry has a too large bits per pixel/hotspot value".to_string()
        ));
    }

    entry.image_length = try!(r.read_u32::<LittleEndian>());
    entry.image_offset = try!(r.read_u32::<LittleEndian>());

    Ok(entry)
}

/// Find the entry with the highest (color depth, size).
fn best_entry(mut entries: Vec<DirEntry>) -> ImageResult<DirEntry> {
    let mut best = try!(entries.pop().ok_or(ImageError::ImageEnd));
    let mut best_score = (best.bits_per_pixel, best.real_width() as u32 * best.real_height() as u32);

    for entry in entries.into_iter() {
        let score = (entry.bits_per_pixel, entry.real_width() as u32 * entry.real_height() as u32);
        if score > best_score {
            best = entry;
            best_score = score;
        }
    }
    Ok(best)
}

impl DirEntry {
    fn real_width(&self) -> u16 {
        match self.width {
            0 => 256,
            w => w as u16
        }
    }

    fn real_height(&self) -> u16 {
        match self.height {
            0 => 256,
            h => h as u16
        }
    }

    fn matches_dimensions(&self, width: u32, height: u32) -> bool {
        u32::from(self.real_width()) == width &&
            u32::from(self.real_height()) == height
    }

    fn seek_to_start<R: Read + Seek>(&self, r: &mut R) -> ImageResult<()> {
        try!(r.seek(SeekFrom::Start(self.image_offset as u64)));
        Ok(())
    }

    fn is_png<R: Read + Seek>(&self, r: &mut R) -> ImageResult<bool> {
        try!(self.seek_to_start(r));

        // Read the first 8 bytes to sniff the image.
        let mut signature = [0u8; 8];
        try!(r.read_exact(&mut signature));

        Ok(signature == PNG_SIGNATURE)
    }

    fn decoder<R: Read + Seek>(&self, mut r: R) -> ImageResult<InnerDecoder<R>> {
        let is_png = try!(self.is_png(&mut r));
        try!(self.seek_to_start(&mut r));

        if is_png {
            Ok(PNG(PNGDecoder::new(r)))
        } else {
            let mut decoder = BMPDecoder::new(r);
            try!(decoder.read_metadata_in_ico_format());
            Ok(BMP(decoder))
        }
    }
}

impl<R: Read + Seek> ImageDecoder for ICODecoder<R> {
    fn dimensions(&mut self) -> ImageResult<(u32, u32)> {
        match self.inner_decoder {
            BMP(ref mut decoder) => decoder.dimensions(),
            PNG(ref mut decoder) => decoder.dimensions()
        }
    }

    fn colortype(&mut self) -> ImageResult<ColorType> {
        match self.inner_decoder {
            BMP(ref mut decoder) => decoder.colortype(),
            PNG(ref mut decoder) => decoder.colortype()
        }
    }

    fn row_len(&mut self) -> ImageResult<usize> {
        match self.inner_decoder {
            BMP(ref mut decoder) => decoder.row_len(),
            PNG(ref mut decoder) => decoder.row_len()
        }
    }

    fn read_scanline(&mut self, buf: &mut [u8]) -> ImageResult<u32> {
        match self.inner_decoder {
            BMP(ref mut decoder) => decoder.read_scanline(buf),
            PNG(ref mut decoder) => decoder.read_scanline(buf)
        }
    }

    fn read_image(&mut self) -> ImageResult<DecodingResult> {
        match self.inner_decoder {
            PNG(ref mut decoder) => {
                if self.selected_entry.image_length < PNG_SIGNATURE.len() as u32 {
                    return Err(ImageError::FormatError(
                        "Entry specified a length that is shorter than PNG header!".to_string()
                    ));
                }

                // Check if the image dimensions match the ones in the image data.
                let (width, height) = try!(decoder.dimensions());
                if !self.selected_entry.matches_dimensions(width, height) {
                    return Err(ImageError::FormatError(
                        "Entry and PNG dimensions do not match!".to_string())
                    );

                }

                // Embedded PNG images can only be of the 32BPP RGBA format.
                // https://blogs.msdn.microsoft.com/oldnewthing/20101022-00/?p=12473/
                let color_type = try!(decoder.colortype());
                if let ColorType::RGBA(8) = color_type {} else {
                    return Err(ImageError::FormatError(
                        "The PNG is not in RGBA format!".to_string()
                    ));
                }

                decoder.read_image()
            }
            BMP(ref mut decoder) => {
                let (width, height) = try!(decoder.dimensions());
                if !self.selected_entry.matches_dimensions(width, height) {
                    return Err(ImageError::FormatError(
                        "Entry({:?}) and BMP({:?}) dimensions do not match!".to_string()
                    ));
                }

                // The ICO decoder needs an alpha chanel to apply the AND mask.
                if try!(decoder.colortype()) != ColorType::RGBA(8) {
                    return Err(ImageError::UnsupportedError("Unsupported color type".to_string()))
                }

                let mut pixel_data = match try!(decoder.read_image()) {
                    DecodingResult::U8(v) => v,
                    _ => unreachable!()
                };

                // If there's an AND mask following the image, read and apply it.
                let r = decoder.reader();
                let mask_start = try!(r.seek(SeekFrom::Current(0)));
                let mask_end = (self.selected_entry.image_offset + self.selected_entry.image_length) as u64;
                let mask_length = mask_end - mask_start;

                if mask_length > 0 {
                    // A mask row contains 1 bit per pixel, padded to 4 bytes.
                    let mask_row_bytes = ((width + 31) / 32) * 4;
                    let expected_length = (mask_row_bytes * height) as u64;
                    if mask_length < expected_length {
                        return Err(ImageError::ImageEnd)
                    }

                    for y in 0..height {
                        let mut x = 0;
                        for _ in 0..mask_row_bytes {
                            // Apply the bits of each byte until we reach the end of the row.
                            let mask_byte = try!(r.read_u8());
                            for bit in (0..8).rev() {
                                if x >= width { break }
                                if mask_byte & (1 << bit) != 0 {
                                    // Set alpha channel to transparent.
                                    pixel_data[((height - y - 1) * width + x) as usize * 4 + 3] = 0;
                                }
                                x += 1;
                            }
                        }
                    }
                }
                Ok(DecodingResult::U8(pixel_data))
            }
        }
    }
}
<|endoftext|>