<commit_before>extern crate regex;

extern crate rusoto_core;
extern crate rusoto_sts;
extern crate rusoto_dynamodb;

pub mod SessionHandler {
    use std::fs::{File, OpenOptions};
    use std::io::BufReader;
    use std::io;
    use std::io::prelude::*;
    use std::io::{stdin, stdout};
    use std::env;
    use rusoto_core::{ProfileProvider, Region, default_tls_client};
    use rusoto_sts::{Sts, StsClient, GetSessionTokenRequest, GetSessionTokenResponse, GetSessionTokenError, Credentials};
    use std::collections::HashMap;

    struct AWSConfig {
        profiles: HashMap<String, AWSProfile>,
    }

    #[derive(Debug)]
    struct AWSProfile {
        source_profile: Option<String>,
        region: Option<String>,
        output: Option<String>,
        role_arn: Option<String>,
        mfa_serial: Option<String>,
        ot_session_name: Option<String>,
        ot_expiration: Option<String>,
        ot_source_profile: Option<String>,
        ot_role_arn: Option<String>,
    }

    impl AWSConfig {
        fn save(&self) -> Result<(), io::Error> {
            println!("Saving config to disk");
            let mut bluenine_config_path = env::home_dir().unwrap().display().to_string();
            bluenine_config_path.push_str("/.aws/bluenine_config");
            let mut file = OpenOptions::new()
                        .write(true)
                        .append(false)
                        .open(bluenine_config_path)
                        .unwrap();

            let mut prof = String::new();
            for (name, profile) in &self.profiles {
                prof.push_str(&format!("[profile {}]\n", name));

                let source_profile = &profile.source_profile;
                match source_profile {
                    &Some(ref source_profile) => {
                        prof.push_str(&format!("source_profile = {}\n", source_profile));
                    },
                    &None => {},
                };

                let region = &profile.region;
                match region {
                    &Some(ref region) => {
                        prof.push_str(&format!("region = {}\n", region));
                    },
                    &None => {},
                };

                let output = &profile.output;
                match output {
                    &Some(ref output) => {
                        prof.push_str(&format!("output = {}\n", output));
                    },
                    &None => {},
                };

                let role_arn = &profile.role_arn;
                match role_arn {
                    &Some(ref role_arn) => {
                        prof.push_str(&format!("role_arn = {}\n", role_arn));
                    },
                    &None => {},
                };

                let mfa_serial = &profile.mfa_serial;
                match mfa_serial {
                    &Some(ref mfa_serial) => {
                        prof.push_str(&format!("mfa_serial = {}\n", mfa_serial));
                    },
                    &None => {},
                };

                let ot_session_name = &profile.ot_session_name;
                match ot_session_name {
                    &Some(ref ot_session_name) => {
                        prof.push_str(&format!("ot_session_name = {}\n", ot_session_name));
                    },
                    &None => {},
                };

                let ot_expiration = &profile.ot_expiration;
                match ot_expiration {
                    &Some(ref ot_expiration) => {
                        prof.push_str(&format!("ot_expiration = {}\n", ot_expiration));
                    },
                    &None => {},
                };

                let ot_expiration = &profile.ot_expiration;
                match ot_expiration {
                    &Some(ref ot_expiration) => {
                        prof.push_str(&format!("ot_expiration = {}\n", ot_expiration));
                    },
                    &None => {},
                };

                let ot_source_profile = &profile.ot_source_profile;
                match ot_source_profile {
                    &Some(ref ot_source_profile) => {
                        prof.push_str(&format!("ot_source_profile = {}\n", ot_source_profile));
                    },
                    &None => {},
                };

                let ot_role_arn = &profile.ot_role_arn;
                match ot_role_arn {
                    &Some(ref ot_role_arn) => {
                        prof.push_str(&format!("ot_role_arn = {}\n", ot_role_arn));
                    },
                    &None => {},
                };

                prof.push_str("\n");
            }

            try!(file.write_all(prof.as_bytes()));
            Ok(())
        }

        fn get_profile(&self, name: &str) -> &AWSProfile {
            if (&self.profiles).contains_key(name) {
                &self.profiles[name]
            }
            else {
                panic!("Could not find profile {}", name);
            }
        }

        // fn remove_profile(&self, name: &str) {
        //     &self.profiles.remove(name);
        // }

        fn contains_profile(&self, name: &str) -> bool {
            if (&self.profiles).contains_key(name) {
                true
            }
            else {
                false
            }
        }
    }

    pub fn create(profile_name: &str) {
        println!("Loading config file");
        let aws_config = load_config();

        let mut session_name = profile_name.clone().to_string();
        session_name.push_str("-session");
        if aws_config.contains_profile(&session_name) {
            println!("Session {} for profile {} already exists", session_name, profile_name);
        }

        println!("Creating session for profile \"{}\"...", profile_name);
        let aws_profile = aws_config.get_profile(profile_name);

        let mut provider = ProfileProvider::new().unwrap();
        provider.set_profile(profile_name);
        // let region = &aws_profile.region;
        // match region {
        //     &Some(ref region) => region,
        //     &None => panic!("You must specify a region for profile {}", profile_name)
        // };
        // println!("Region: {:?}", region.as_ref().unwrap());
        let client = StsClient::new(default_tls_client().unwrap(), provider, Region::EuWest1);

        let mfa_serial = &aws_profile.mfa_serial;
        if mfa_serial.is_some() {
            print!("Enter AWS MFA code for profile [{}]: ", profile_name);
            stdout().flush();
            let mut token_code = String::new();
            stdin().read_line(&mut token_code)
                .ok()
                .expect("Couldn't read line");
            token_code.pop(); // Remove newline

            let request = GetSessionTokenRequest {
                duration_seconds: None,
                serial_number: Some(mfa_serial.as_ref().unwrap().to_owned()),
                token_code: Some(token_code),
            };
            let response = client.get_session_token(&request);
            match response {
                Ok(response) => {
                    match save_profile(profile_name, &aws_profile) {
                        Ok(_) => println!("Saved profile to config file."),
                        Err(err) => println!("Error saving profile config to file: {:?}", err)
                    };
                    let credentials = response.credentials.unwrap();
                    match save_credentials(profile_name, credentials) {
                        Ok(_) => println!("Saved to credentials file."),
                        Err(err) => println!("Error saving credentials to file: {:?}", err)
                    };
                },
                Err(err) => panic!("Failed to get session token for profile {}: {:?}", profile_name, err),
            };
        }
    }
    pub fn show(profile_name: &str) {
        println!("Showing config for profile {}...", profile_name);
    }

    pub fn refresh() {
        println!("Refreshing sessions...");
    }

    pub fn clean(profile_name: &str) {
        println!("Loading config file");
        let mut aws_config = load_config();
        let mut session_name = profile_name.clone().to_string();
        session_name.push_str("-session");

        if aws_config.contains_profile(&session_name) {
            println!("Cleaning session {}", session_name);
            aws_config.profiles.remove(&session_name);
            aws_config.save();
        }
        else {
            println!("No profiles to clean");
        }
    }

    fn load_config() -> AWSConfig {
        let aws_config_file = read_aws_config_file();
        let profiles = split_config_file(aws_config_file);

        // Create an instance of the AWSConfig struct
        let mut aws_config = AWSConfig {
            profiles: HashMap::new()
        };

        // Iterate over profile chunks
        for profile in profiles {
            // Create a vector of string that holds each line
            let split = profile.split("\n");
            let lines: Vec<String> = split.map(|s| s.to_string()).collect();

            // Get the profile name from the first line
            let profile_line = lines[0].to_owned();

            // Lets ignore the default for now
            if profile_line == "[default]" {
                continue;
            }

            let split2 = profile_line.split(" ");
            let mut words: Vec<String> = split2.map(|s| s.to_string()).collect();
            let mut profile_name = String::new();
            if words.len() > 1 {
                words[1].pop();
                profile_name = words[1].trim_right().to_owned();
            }

            let mut aws_profile = AWSProfile{
                source_profile: None,
                region: None,
                output: None,
                role_arn: None,
                mfa_serial: None,
                ot_expiration: None,
                ot_session_name: None,
                ot_source_profile: None,
                ot_role_arn: None,
            };

            for i in 1..lines.len() {
                let split = lines[i].split(" = ");
                let config: Vec<String> = split.map(|s| s.to_string()).collect();
                let key = &config[0];
                let value = config[1].to_owned();
                match key.as_ref() {
                    "mfa_serial" => aws_profile.mfa_serial = Some(value),
                    "output" => aws_profile.output = Some(value),
                    "region" => aws_profile.region = Some(value),
                    "role_arn" => aws_profile.role_arn = Some(value),
                    "source_profile" => aws_profile.source_profile = Some(value),
                    "ot_session_name" => aws_profile.ot_session_name = Some(value),
                    "ot_expiration" => aws_profile.ot_expiration = Some(value),
                    "ot_source_profile" => aws_profile.ot_source_profile = Some(value),
                    "ot_role_arn" => aws_profile.ot_role_arn = Some(value),
                    _ => (),
                }
            }

            if profile_name != "" {
                aws_config.profiles.insert(profile_name.clone(), aws_profile);
            }
        }

        aws_config
    }

    fn read_aws_config_file() -> String {
        let path = match env::home_dir() {
            Some(path) => path,
            None => panic!("Could not retrieve user's home directory."),
        };
        let config_file_path = format!("{}/.aws/config", path.display());

        let f = File::open(config_file_path).expect("Could not find AWS config file.");
        let mut buf_reader = BufReader::new(f);
        let mut contents = String::new();
        buf_reader.read_to_string(&mut contents).expect("Found config file but could not read it.");

        contents
    }

    fn split_config_file(aws_config: String) -> Vec<String> {
        let split = aws_config.split("\n\n");
        let mut profiles: Vec<String> = split.map(|s| s.to_string()).collect();;

        profiles.pop(); // Remove last element as it is always empty
        profiles
    }

    fn save_credentials(profile_name: &str, credentials: Credentials) -> Result<(), io::Error> {
        let mut aws_credentials_path = env::home_dir().unwrap().display().to_string();
        aws_credentials_path.push_str("/.aws/credentials");
        let mut file = OpenOptions::new()
                       .write(true)
                       .append(true)
                       .open(aws_credentials_path)
                       .unwrap();

        let mut creds = String::new();
        creds.push_str(&format!("[{}-session]\n", profile_name));
        creds.push_str(&format!("aws_access_key_id = {}\n", credentials.access_key_id));
        creds.push_str(&format!("expiration = {}\n", credentials.expiration));
        creds.push_str(&format!("aws_secret_access_key = {}\n", credentials.secret_access_key));
        creds.push_str(&format!("aws_session_token = {}\n\n", credentials.session_token));
        try!(file.write_all(creds.as_bytes()));
        Ok(())
    }

    fn save_profile(profile_name: &str, aws_profile: &AWSProfile) -> Result<(), io::Error> {
        let mut aws_config_path = env::home_dir().unwrap().display().to_string();
        aws_config_path.push_str("/.aws/config");
        let mut file = OpenOptions::new()
                       .write(true)
                       .append(true)
                       .open(aws_config_path)
                       .unwrap();

        let region = &aws_profile.region;
        match region {
            &Some(ref region) => {
                let mut prof = String::new();
                prof.push_str(&format!("[profile {}-session]\n", profile_name));
                prof.push_str(&format!("region = {}\n\n", region));
                try!(file.write_all(prof.as_bytes()));
                Ok(())
            },
            &None => {
                let region = "us-east-1";
                let mut prof = String::new();
                prof.push_str(&format!("[profile {}-session]\n", profile_name));
                prof.push_str(&format!("region = {}\n\n", region));
                try!(file.write_all(prof.as_bytes()));
                Ok(())
            },
        }
    }
}
<commit_msg>Add remove credentials function<commit_after>extern crate regex;

extern crate rusoto_core;
extern crate rusoto_sts;
extern crate rusoto_dynamodb;

pub mod SessionHandler {
    use std::fs::{File, OpenOptions};
    use std::io::BufReader;
    use std::io;
    use std::io::prelude::*;
    use std::io::{stdin, stdout};
    use std::env;
    use rusoto_core::{ProfileProvider, Region, default_tls_client};
    use rusoto_sts::{Sts, StsClient, GetSessionTokenRequest, GetSessionTokenResponse, GetSessionTokenError, Credentials};
    use std::collections::HashMap;

    struct AWSConfig {
        profiles: HashMap<String, AWSProfile>,
    }

    #[derive(Debug)]
    struct AWSProfile {
        source_profile: Option<String>,
        region: Option<String>,
        output: Option<String>,
        role_arn: Option<String>,
        mfa_serial: Option<String>,
        ot_session_name: Option<String>,
        ot_expiration: Option<String>,
        ot_source_profile: Option<String>,
        ot_role_arn: Option<String>,
    }

    impl AWSConfig {
        fn save(&self) -> Result<(), io::Error> {
            println!("Saving config to disk");
            let mut bluenine_config_path = env::home_dir().unwrap().display().to_string();
            bluenine_config_path.push_str("/.aws/bluenine_config");
            let mut file = OpenOptions::new()
                        .write(true)
                        .truncate(true)
                        .open(bluenine_config_path)
                        .unwrap();

            let mut prof = String::new();
            for (name, profile) in &self.profiles {
                prof.push_str(&format!("[profile {}]\n", name));

                let source_profile = &profile.source_profile;
                match source_profile {
                    &Some(ref source_profile) => {
                        prof.push_str(&format!("source_profile = {}\n", source_profile));
                    },
                    &None => {},
                };

                let region = &profile.region;
                match region {
                    &Some(ref region) => {
                        prof.push_str(&format!("region = {}\n", region));
                    },
                    &None => {},
                };

                let output = &profile.output;
                match output {
                    &Some(ref output) => {
                        prof.push_str(&format!("output = {}\n", output));
                    },
                    &None => {},
                };

                let role_arn = &profile.role_arn;
                match role_arn {
                    &Some(ref role_arn) => {
                        prof.push_str(&format!("role_arn = {}\n", role_arn));
                    },
                    &None => {},
                };

                let mfa_serial = &profile.mfa_serial;
                match mfa_serial {
                    &Some(ref mfa_serial) => {
                        prof.push_str(&format!("mfa_serial = {}\n", mfa_serial));
                    },
                    &None => {},
                };

                let ot_session_name = &profile.ot_session_name;
                match ot_session_name {
                    &Some(ref ot_session_name) => {
                        prof.push_str(&format!("ot_session_name = {}\n", ot_session_name));
                    },
                    &None => {},
                };

                let ot_expiration = &profile.ot_expiration;
                match ot_expiration {
                    &Some(ref ot_expiration) => {
                        prof.push_str(&format!("ot_expiration = {}\n", ot_expiration));
                    },
                    &None => {},
                };

                let ot_expiration = &profile.ot_expiration;
                match ot_expiration {
                    &Some(ref ot_expiration) => {
                        prof.push_str(&format!("ot_expiration = {}\n", ot_expiration));
                    },
                    &None => {},
                };

                let ot_source_profile = &profile.ot_source_profile;
                match ot_source_profile {
                    &Some(ref ot_source_profile) => {
                        prof.push_str(&format!("ot_source_profile = {}\n", ot_source_profile));
                    },
                    &None => {},
                };

                let ot_role_arn = &profile.ot_role_arn;
                match ot_role_arn {
                    &Some(ref ot_role_arn) => {
                        prof.push_str(&format!("ot_role_arn = {}\n", ot_role_arn));
                    },
                    &None => {},
                };

                prof.push_str("\n");
            }

            try!(file.write_all(prof.as_bytes()));
            Ok(())
        }

        fn get_profile(&self, name: &str) -> &AWSProfile {
            if (&self.profiles).contains_key(name) {
                &self.profiles[name]
            }
            else {
                panic!("Could not find profile {}", name);
            }
        }

        fn contains_profile(&self, name: &str) -> bool {
            if (&self.profiles).contains_key(name) {
                true
            }
            else {
                false
            }
        }
    }

    pub fn create(profile_name: &str) {
        println!("Loading config file");
        let aws_config = load_config();

        let mut session_name = profile_name.clone().to_string();
        session_name.push_str("-session");
        if aws_config.contains_profile(&session_name) {
            println!("Session {} for profile {} already exists", session_name, profile_name);
        }

        println!("Creating session for profile \"{}\"...", profile_name);
        let aws_profile = aws_config.get_profile(profile_name);

        let mut provider = ProfileProvider::new().unwrap();
        provider.set_profile(profile_name);
        // let region = &aws_profile.region;
        // match region {
        //     &Some(ref region) => region,
        //     &None => panic!("You must specify a region for profile {}", profile_name)
        // };
        // println!("Region: {:?}", region.as_ref().unwrap());
        let client = StsClient::new(default_tls_client().unwrap(), provider, Region::EuWest1);

        let mfa_serial = &aws_profile.mfa_serial;
        if mfa_serial.is_some() {
            print!("Enter AWS MFA code for profile [{}]: ", profile_name);
            stdout().flush();
            let mut token_code = String::new();
            stdin().read_line(&mut token_code)
                .ok()
                .expect("Couldn't read line");
            token_code.pop(); // Remove newline

            let request = GetSessionTokenRequest {
                duration_seconds: None,
                serial_number: Some(mfa_serial.as_ref().unwrap().to_owned()),
                token_code: Some(token_code),
            };
            let response = client.get_session_token(&request);
            match response {
                Ok(response) => {
                    match save_profile(profile_name, &aws_profile) {
                        Ok(_) => println!("Saved profile to config file."),
                        Err(err) => println!("Error saving profile config to file: {:?}", err)
                    };
                    let credentials = response.credentials.unwrap();
                    match save_credentials(profile_name, credentials) {
                        Ok(_) => println!("Saved to credentials file."),
                        Err(err) => println!("Error saving credentials to file: {:?}", err)
                    };
                },
                Err(err) => panic!("Failed to get session token for profile {}: {:?}", profile_name, err),
            };
        }
    }
    pub fn show(profile_name: &str) {
        println!("Showing config for profile {}...", profile_name);
    }

    pub fn refresh() {
        println!("Refreshing sessions...");
    }

    pub fn clean(profile_name: &str) {
        println!("Loading config file");
        let mut aws_config = load_config();
        let mut session_name = profile_name.clone().to_string();
        session_name.push_str("-session");

        if aws_config.contains_profile(&session_name) {
            println!("Cleaning session {}", session_name);
            aws_config.profiles.remove(&session_name);
            aws_config.save();
            // Remove credentials also
            remove_credentials(&session_name);
        }
        else {
            println!("No profiles to clean");
        }
    }

    fn load_config() -> AWSConfig {
        let aws_config_file = read_aws_config_file();
        let profiles = split_config_file(aws_config_file);

        // Create an instance of the AWSConfig struct
        let mut aws_config = AWSConfig {
            profiles: HashMap::new()
        };

        // Iterate over profile chunks
        for profile in profiles {
            // Create a vector of string that holds each line
            let split = profile.split("\n");
            let lines: Vec<String> = split.map(|s| s.to_string()).collect();

            // Get the profile name from the first line
            let profile_line = lines[0].to_owned();

            // Lets ignore the default for now
            if profile_line == "[default]" {
                continue;
            }

            let split2 = profile_line.split(" ");
            let mut words: Vec<String> = split2.map(|s| s.to_string()).collect();
            let mut profile_name = String::new();
            if words.len() > 1 {
                words[1].pop();
                profile_name = words[1].trim_right().to_owned();
            }

            let mut aws_profile = AWSProfile{
                source_profile: None,
                region: None,
                output: None,
                role_arn: None,
                mfa_serial: None,
                ot_expiration: None,
                ot_session_name: None,
                ot_source_profile: None,
                ot_role_arn: None,
            };

            for i in 1..lines.len() {
                let split = lines[i].split(" = ");
                let config: Vec<String> = split.map(|s| s.to_string()).collect();
                let key = &config[0];
                let value = config[1].to_owned();
                match key.as_ref() {
                    "mfa_serial" => aws_profile.mfa_serial = Some(value),
                    "output" => aws_profile.output = Some(value),
                    "region" => aws_profile.region = Some(value),
                    "role_arn" => aws_profile.role_arn = Some(value),
                    "source_profile" => aws_profile.source_profile = Some(value),
                    "ot_session_name" => aws_profile.ot_session_name = Some(value),
                    "ot_expiration" => aws_profile.ot_expiration = Some(value),
                    "ot_source_profile" => aws_profile.ot_source_profile = Some(value),
                    "ot_role_arn" => aws_profile.ot_role_arn = Some(value),
                    _ => (),
                }
            }

            if profile_name != "" {
                aws_config.profiles.insert(profile_name.clone(), aws_profile);
            }
        }

        aws_config
    }

    fn read_aws_config_file() -> String {
        let path = match env::home_dir() {
            Some(path) => path,
            None => panic!("Could not retrieve user's home directory."),
        };
        let config_file_path = format!("{}/.aws/bluenine_config", path.display());

        let f = File::open(config_file_path).expect("Could not find AWS config file.");
        let mut buf_reader = BufReader::new(f);
        let mut contents = String::new();
        buf_reader.read_to_string(&mut contents).expect("Found config file but could not read it.");

        contents
    }

    fn read_aws_credentials_file() -> String {
        let path = match env::home_dir() {
            Some(path) => path,
            None => panic!("Could not retrieve user's home directory."),
        };
        let config_file_path = format!("{}/.aws/bluenine_credentials", path.display());

        let f = File::open(config_file_path).expect("Could not find AWS credentials file.");
        let mut buf_reader = BufReader::new(f);
        let mut contents = String::new();
        buf_reader.read_to_string(&mut contents).expect("Found credentials file but could not read it.");

        contents
    }

    fn split_config_file(aws_config: String) -> Vec<String> {
        let split = aws_config.split("\n\n");
        let mut profiles: Vec<String> = split.map(|s| s.to_string()).collect();;

        profiles.pop(); // Remove last element as it is always empty
        profiles
    }

    fn save_credentials(profile_name: &str, credentials: Credentials) -> Result<(), io::Error> {
        let mut aws_credentials_path = env::home_dir().unwrap().display().to_string();
        aws_credentials_path.push_str("/.aws/bluenine_credentials");
        let mut file = OpenOptions::new()
                       .write(true)
                       .append(true)
                       .open(aws_credentials_path)
                       .unwrap();

        let mut creds = String::new();
        creds.push_str(&format!("[{}-session]\n", profile_name));
        creds.push_str(&format!("aws_access_key_id = {}\n", credentials.access_key_id));
        creds.push_str(&format!("expiration = {}\n", credentials.expiration));
        creds.push_str(&format!("aws_secret_access_key = {}\n", credentials.secret_access_key));
        creds.push_str(&format!("aws_session_token = {}\n\n", credentials.session_token));
        try!(file.write_all(creds.as_bytes()));
        Ok(())
    }

    fn remove_credentials(profile_name: &str) -> Result<(), io::Error> {
        println!("Saving credentials to disk");
        let aws_credentials_file = read_aws_credentials_file();
        let credentials = split_config_file(aws_credentials_file);
        let mut aws_credentials_path = env::home_dir().unwrap().display().to_string();
        aws_credentials_path.push_str("/.aws/bluenine_credentials");
        let mut file = OpenOptions::new()
                .write(true)
                .truncate(true)
                .open(aws_credentials_path)
                .unwrap();

        let mut creds = String::new();
        for credential in credentials {
            let split = credential.split("\n");
            let lines: Vec<String> = split.map(|s| s.to_string()).collect();

            // Get the profile name from the first line
            let profile_line = lines[0].to_owned();
            if profile_line == format!("[{}]", profile_name) {
                continue;
            }
            creds.push_str(&credential);
            creds.push_str("\n\n");
        }
        try!(file.write_all(creds.as_bytes()));
        Ok(())
    }

    fn save_profile(profile_name: &str, aws_profile: &AWSProfile) -> Result<(), io::Error> {
        let mut aws_config_path = env::home_dir().unwrap().display().to_string();
        aws_config_path.push_str("/.aws/bluenine_config");
        let mut file = OpenOptions::new()
                       .write(true)
                       .append(true)
                       .open(aws_config_path)
                       .unwrap();

        let region = &aws_profile.region;
        match region {
            &Some(ref region) => {
                let mut prof = String::new();
                prof.push_str(&format!("[profile {}-session]\n", profile_name));
                prof.push_str(&format!("region = {}\n\n", region));
                try!(file.write_all(prof.as_bytes()));
                Ok(())
            },
            &None => {
                let region = "us-east-1";
                let mut prof = String::new();
                prof.push_str(&format!("[profile {}-session]\n", profile_name));
                prof.push_str(&format!("region = {}\n\n", region));
                try!(file.write_all(prof.as_bytes()));
                Ok(())
            },
        }
    }
}
<|endoftext|>


<commit_before>extern crate librsync_sys as raw;

mod job;

use job::{Job, JobDriver};

use std::error;
use std::fmt::{self, Display, Formatter};
use std::io::{self, Read};
use std::ops::Deref;
use std::ptr;

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum SignatureType {
    MD4,
    Blake2,
}

#[derive(Debug)]
pub enum Error {
    Blocked,
    Io(io::Error),
    Syntax,
    Mem,
    InputEnded,
    BadMagic,
    Unimplemented,
    Corrupt,
    Internal,
    Param,
    Unknown(i32),
}

pub type Result<T> = std::result::Result<T, Error>;

pub struct Signature<R: Read> {
    driver: JobDriver<R>,
}

pub struct Delta<R: Read> {
    driver: JobDriver<R>,
    _sumset: Sumset,
}

struct Sumset(*mut raw::rs_signature_t);

impl<R: Read> Signature<R> {
    pub fn new(old: R,
               new_block_len: usize,
               strong_len: usize,
               sig_magic: SignatureType)
               -> Result<Self> {
        let job = unsafe { raw::rs_sig_begin(new_block_len, strong_len, sig_magic.as_raw()) };
        if job.is_null() {
            return Err(Error::BadMagic);
        }
        Ok(Signature { driver: JobDriver::new(old, Job(job)), })
    }

    pub fn into_inner(self) -> R {
        self.driver.into_inner()
    }
}

impl<R: Read> Read for Signature<R> {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        self.driver.read(buf)
    }
}

impl<R: Read> Delta<R> {
    pub fn new<R2: Read>(input: R, signature: R2) -> Result<Self> {
        // load the signature
        let sumset = unsafe {
            let mut sumset = ptr::null_mut();
            let job = raw::rs_loadsig_begin(&mut sumset);
            assert!(!job.is_null());
            let mut job = JobDriver::new(signature, Job(job));
            try!(job.consume_input());
            let sumset = Sumset(sumset);
            let res = raw::rs_build_hash_table(*sumset);
            if res != raw::RS_DONE {
                return Err(Error::from(res));
            }
            sumset
        };
        let job = unsafe { raw::rs_delta_begin(*sumset) };
        assert!(!job.is_null());
        Ok(Delta {
            driver: JobDriver::new(input, Job(job)),
            _sumset: sumset,
        })
    }
}

impl<R: Read> Read for Delta<R> {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        self.driver.read(buf)
    }
}

impl error::Error for Error {
    fn description(&self) -> &str {
        match *self {
            Error::Blocked => "blocked waiting for more data",
            Error::Io(ref err) => err.description(),
            Error::Syntax => "syntax error",
            Error::Mem => "out of memory",
            Error::InputEnded => "unexpected end of input file",
            Error::BadMagic => "bad magic number given",
            Error::Unimplemented => "unimplemented feature",
            Error::Corrupt => "unbelievable value in stream",
            Error::Internal => "internal error",
            Error::Param => "bad parameter",
            Error::Unknown(_) => "unknown error from librsync",
        }
    }
}

impl Display for Error {
    fn fmt(&self, fmt: &mut Formatter) -> fmt::Result {
        match *self {
            Error::Io(ref e) => write!(fmt, "{}", e),
            Error::Unknown(n) => write!(fmt, "unknown error {} from native library", n),
            _ => write!(fmt, "{}", std::error::Error::description(self)),
        }
    }
}

impl From<io::Error> for Error {
    fn from(err: io::Error) -> Error {
        Error::Io(err)
    }
}

impl From<raw::rs_result> for Error {
    fn from(err: raw::rs_result) -> Error {
        match err {
            raw::RS_BLOCKED => Error::Blocked,
            raw::RS_IO_ERROR => Error::Io(other_io_err("Unknown IO error from librsync")),
            raw::RS_SYNTAX_ERROR => Error::Syntax,
            raw::RS_MEM_ERROR => Error::Mem,
            raw::RS_INPUT_ENDED => Error::InputEnded,
            raw::RS_BAD_MAGIC => Error::BadMagic,
            raw::RS_UNIMPLEMENTED => Error::Unimplemented,
            raw::RS_CORRUPT => Error::Corrupt,
            raw::RS_INTERNAL_ERROR => Error::Internal,
            raw::RS_PARAM_ERROR => Error::Param,
            n => Error::Unknown(n),
        }
    }
}

impl SignatureType {
    pub fn from_raw(raw: raw::rs_magic_number) -> Option<Self> {
        match raw {
            raw::RS_MD4_SIG_MAGIC => Some(SignatureType::MD4),
            raw::RS_BLAKE2_SIG_MAGIC => Some(SignatureType::Blake2),
            _ => None,
        }
    }

    fn as_raw(&self) -> raw::rs_magic_number {
        match *self {
            SignatureType::MD4 => raw::RS_MD4_SIG_MAGIC,
            SignatureType::Blake2 => raw::RS_BLAKE2_SIG_MAGIC,
        }
    }
}

impl Drop for Sumset {
    fn drop(&mut self) {
        unsafe {
            if !self.0.is_null() {
                raw::rs_free_sumset(self.0);
            }
        }
    }
}

impl Deref for Sumset {
    type Target = *mut raw::rs_signature_t;
    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

fn other_io_err<T: AsRef<str>>(msg: T) -> io::Error {
    io::Error::new(io::ErrorKind::Other, msg.as_ref())
}

#[cfg(test)]
mod test {
    use super::*;
    use std::io::{Cursor, Read};

    #[test]
    fn signature() {
        let data = "this is a string to be tested";
        let cursor = Cursor::new(data);
        let mut sig = Signature::new(cursor, 10, 5, SignatureType::MD4).unwrap();
        let mut signature = Vec::new();
        let read = sig.read_to_end(&mut signature).unwrap();
        assert_eq!(read, signature.len());
        let expected = vec![0x72, 0x73, 0x01, 0x36, 0x00, 0x00, 0x00, 0x0a, 0x00, 0x00, 0x00,
                            0x05, 0x1b, 0x21, 0x04, 0x8b, 0xad, 0x3c, 0xbd, 0x19, 0x09, 0x1d,
                            0x1b, 0x04, 0xf0, 0x9d, 0x1f, 0x64, 0x31, 0xde, 0x15, 0xf4, 0x04,
                            0x87, 0x60, 0x96, 0x19, 0x50, 0x39];
        assert_eq!(signature, expected);
    }
}
<commit_msg>Added test on delta<commit_after>extern crate librsync_sys as raw;

mod job;

use job::{Job, JobDriver};

use std::error;
use std::fmt::{self, Display, Formatter};
use std::io::{self, Read};
use std::ops::Deref;
use std::ptr;

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum SignatureType {
    MD4,
    Blake2,
}

#[derive(Debug)]
pub enum Error {
    Blocked,
    Io(io::Error),
    Syntax,
    Mem,
    InputEnded,
    BadMagic,
    Unimplemented,
    Corrupt,
    Internal,
    Param,
    Unknown(i32),
}

pub type Result<T> = std::result::Result<T, Error>;

pub struct Signature<R: Read> {
    driver: JobDriver<R>,
}

pub struct Delta<R: Read> {
    driver: JobDriver<R>,
    _sumset: Sumset,
}

struct Sumset(*mut raw::rs_signature_t);

impl<R: Read> Signature<R> {
    pub fn new(old: R,
               new_block_len: usize,
               strong_len: usize,
               sig_magic: SignatureType)
               -> Result<Self> {
        let job = unsafe { raw::rs_sig_begin(new_block_len, strong_len, sig_magic.as_raw()) };
        if job.is_null() {
            return Err(Error::BadMagic);
        }
        Ok(Signature { driver: JobDriver::new(old, Job(job)), })
    }

    pub fn into_inner(self) -> R {
        self.driver.into_inner()
    }
}

impl<R: Read> Read for Signature<R> {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        self.driver.read(buf)
    }
}

impl<R: Read> Delta<R> {
    pub fn new<R2: Read>(input: R, signature: R2) -> Result<Self> {
        // load the signature
        let sumset = unsafe {
            let mut sumset = ptr::null_mut();
            let job = raw::rs_loadsig_begin(&mut sumset);
            assert!(!job.is_null());
            let mut job = JobDriver::new(signature, Job(job));
            try!(job.consume_input());
            let sumset = Sumset(sumset);
            let res = raw::rs_build_hash_table(*sumset);
            if res != raw::RS_DONE {
                return Err(Error::from(res));
            }
            sumset
        };
        let job = unsafe { raw::rs_delta_begin(*sumset) };
        assert!(!job.is_null());
        Ok(Delta {
            driver: JobDriver::new(input, Job(job)),
            _sumset: sumset,
        })
    }
}

impl<R: Read> Read for Delta<R> {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        self.driver.read(buf)
    }
}

impl error::Error for Error {
    fn description(&self) -> &str {
        match *self {
            Error::Blocked => "blocked waiting for more data",
            Error::Io(ref err) => err.description(),
            Error::Syntax => "syntax error",
            Error::Mem => "out of memory",
            Error::InputEnded => "unexpected end of input file",
            Error::BadMagic => "bad magic number given",
            Error::Unimplemented => "unimplemented feature",
            Error::Corrupt => "unbelievable value in stream",
            Error::Internal => "internal error",
            Error::Param => "bad parameter",
            Error::Unknown(_) => "unknown error from librsync",
        }
    }
}

impl Display for Error {
    fn fmt(&self, fmt: &mut Formatter) -> fmt::Result {
        match *self {
            Error::Io(ref e) => write!(fmt, "{}", e),
            Error::Unknown(n) => write!(fmt, "unknown error {} from native library", n),
            _ => write!(fmt, "{}", std::error::Error::description(self)),
        }
    }
}

impl From<io::Error> for Error {
    fn from(err: io::Error) -> Error {
        Error::Io(err)
    }
}

impl From<raw::rs_result> for Error {
    fn from(err: raw::rs_result) -> Error {
        match err {
            raw::RS_BLOCKED => Error::Blocked,
            raw::RS_IO_ERROR => Error::Io(other_io_err("Unknown IO error from librsync")),
            raw::RS_SYNTAX_ERROR => Error::Syntax,
            raw::RS_MEM_ERROR => Error::Mem,
            raw::RS_INPUT_ENDED => Error::InputEnded,
            raw::RS_BAD_MAGIC => Error::BadMagic,
            raw::RS_UNIMPLEMENTED => Error::Unimplemented,
            raw::RS_CORRUPT => Error::Corrupt,
            raw::RS_INTERNAL_ERROR => Error::Internal,
            raw::RS_PARAM_ERROR => Error::Param,
            n => Error::Unknown(n),
        }
    }
}

impl SignatureType {
    pub fn from_raw(raw: raw::rs_magic_number) -> Option<Self> {
        match raw {
            raw::RS_MD4_SIG_MAGIC => Some(SignatureType::MD4),
            raw::RS_BLAKE2_SIG_MAGIC => Some(SignatureType::Blake2),
            _ => None,
        }
    }

    fn as_raw(&self) -> raw::rs_magic_number {
        match *self {
            SignatureType::MD4 => raw::RS_MD4_SIG_MAGIC,
            SignatureType::Blake2 => raw::RS_BLAKE2_SIG_MAGIC,
        }
    }
}

impl Drop for Sumset {
    fn drop(&mut self) {
        unsafe {
            if !self.0.is_null() {
                raw::rs_free_sumset(self.0);
            }
        }
    }
}

impl Deref for Sumset {
    type Target = *mut raw::rs_signature_t;
    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

fn other_io_err<T: AsRef<str>>(msg: T) -> io::Error {
    io::Error::new(io::ErrorKind::Other, msg.as_ref())
}

#[cfg(test)]
mod test {
    use super::*;
    use std::io::{Cursor, Read};

    const DATA: &'static str = "this is a string to be tested";
    const DATA2: &'static str = "this is another string to be tested";

    fn data_signature() -> Vec<u8> {
        vec![0x72, 0x73, 0x01, 0x36, 0x00, 0x00, 0x00, 0x0a, 0x00, 0x00, 0x00, 0x05, 0x1b, 0x21,
             0x04, 0x8b, 0xad, 0x3c, 0xbd, 0x19, 0x09, 0x1d, 0x1b, 0x04, 0xf0, 0x9d, 0x1f, 0x64,
             0x31, 0xde, 0x15, 0xf4, 0x04, 0x87, 0x60, 0x96, 0x19, 0x50, 0x39]
    }

    #[test]
    fn signature() {
        let cursor = Cursor::new(DATA);
        let mut sig = Signature::new(cursor, 10, 5, SignatureType::MD4).unwrap();
        let mut signature = Vec::new();
        let read = sig.read_to_end(&mut signature).unwrap();
        assert_eq!(read, signature.len());
        assert_eq!(signature, data_signature());
    }

    #[test]
    fn delta() {
        let sig = data_signature();
        let sig = Cursor::new(sig);
        let input = Cursor::new(DATA2);
        let mut job = Delta::new(input, sig).unwrap();
        let mut delta = Vec::new();
        let read = job.read_to_end(&mut delta).unwrap();
    }
}
<|endoftext|>


<commit_before>#![feature(vec_push_all)]

pub struct JsonEncoder {
    buffer: Vec<u8>
}

impl JsonEncoder {
    pub fn new() -> JsonEncoder {
        JsonEncoder{buffer: Vec::new()}
    }

    #[inline(always)]
    pub fn encode_raw(&mut self, raw: &[u8]) {
        self.buffer.push_all(raw);
    }

    pub fn into_vec(self) -> Vec<u8> {
        self.buffer
    }

    #[inline(always)]
    pub fn encode_raw_str(&mut self, raw_str: &str) {
        self.buffer.push_all(raw_str.as_bytes());
    }

    fn escape_bytes(&mut self, bytes: &[u8]) {
        let mut start = 0;

        for (i, byte) in bytes.iter().enumerate() {
            let escaped = match *byte {
                b'"' => b"\\\"",
                b'\\' => b"\\\\",
                b'\x08' => b"\\b",
                b'\x0c' => b"\\f",
                b'\n' => b"\\n",
                b'\r' => b"\\r",
                b'\t' => b"\\t",
                _ => {
                    continue;
                }
            };

            if start < i {
                self.encode_raw(&bytes[start..i]);
            }

            self.encode_raw(escaped);

            start = i + 1;
        }

        if start != bytes.len() {
            self.encode_raw(&bytes[start..]);
        }
    }

    pub fn encode_u64_decimal(&mut self, value: u64) {
        self.buffer.push(b'"');
        self._encode_u64_decimal(value);
        self.buffer.push(b'"');
    }

    // encodes as decimal string
    #[inline(always)]
    fn _encode_u64_decimal(&mut self, value: u64) {
        const CHARS: &'static [u8] = b"0123456789";
        const MAX_DIGITS: usize = 20;

        if value == 0 {
            self.buffer.push(b'0');
            return;
        }

        let mut digits: [u8; MAX_DIGITS] = [b'0'; MAX_DIGITS];
        let mut i = MAX_DIGITS;
        let mut n = value;
        while n > 0 {
            i -= 1;
            digits[i] = CHARS[(n % 10) as usize];
            n = n / 10;
        }

        self.encode_raw(&digits[i..]);
    }

    #[inline(always)]
    pub fn encode_str(&mut self, s: &str) {
        self.buffer.push(b'"');
        self.escape_bytes(s.as_bytes());
        self.buffer.push(b'"');
    }

    #[inline(always)]
    pub fn encode_i32(&mut self, value: i32) {
        if value == 0 {
            self.buffer.push(b'0');
        } else if value > 0 {
            self.encode_u31(value as u32);
        } else {
            self.buffer.push(b'-');
            self.encode_u31((-value) as u32);
        }
    }

    // encodes a 31-bit unsigned integer != 0
    fn encode_u31(&mut self, value: u32) {
        const CHARS: &'static [u8] = b"0123456789";
        const MAX_DIGITS: usize = 10;
        debug_assert!(value != 0);

        let mut digits: [u8; MAX_DIGITS] = [b'0'; MAX_DIGITS];
        let mut i = MAX_DIGITS;
        let mut n = value;
        while n > 0 {
            i -= 1;
            digits[i] = CHARS[(n % 10) as usize];
            n = n / 10;
        }

        self.encode_raw(&digits[i..]);
    }

    pub fn encode_obj<F>(&mut self, f: F) where F: Fn(&mut JsonObjectEncoder) {
        self.buffer.push(b'{');
        {
            f(&mut JsonObjectEncoder {js: self, needs_sep: false});
        }
        self.buffer.push(b'}');
    }

    pub fn encode_array<F>(&mut self, f: F) where F: Fn(&mut JsonArrayEncoder) {
        self.buffer.push(b'[');
        {
            f(&mut JsonArrayEncoder {js: self, needs_sep: false});
        }
        self.buffer.push(b']');
    }

}

pub struct JsonObjectEncoder<'a> {
    js: &'a mut JsonEncoder,
    needs_sep: bool,
}

impl<'a> JsonObjectEncoder<'a> {
    // XXX: name MAY NOT include escapable characters
    #[inline(always)]
    pub fn encode_field<F:Fn(&mut JsonEncoder)>(&mut self, name: &str, f: F) {
        if self.needs_sep {
            self.js.buffer.push(b',');
        } else {
            self.needs_sep = true;
        }
        self.js.buffer.push(b'"');
        self.js.encode_raw_str(name);
        self.js.buffer.push(b'"');
        self.js.buffer.push(b':');

        f(self.js);
    }
}

pub struct JsonArrayEncoder<'a> {
    js: &'a mut JsonEncoder,
    needs_sep: bool,
}

impl<'a> JsonArrayEncoder<'a> {
    #[inline(always)]
    pub fn encode_elm<F:Fn(&mut JsonEncoder)>(&mut self, f: F) {
        if self.needs_sep {
            self.js.buffer.push(b',');
        } else {
            self.needs_sep = true;
        }
        f(self.js);
    }
}

#[test]
fn test_json_obj_encoder() {
    use std::str;

    let mut js = JsonEncoder::new();
    js.encode_obj(|_|{});
    assert_eq!(b"{}", &js.into_vec()[..]);

    let mut js = JsonEncoder::new();
    js.encode_obj(|jso| {
        jso.encode_field("total", |js| js.encode_i32(31));
    });
    assert_eq!(b"{\"total\":31}", &js.into_vec()[..]);

    let mut js = JsonEncoder::new();
    js.encode_obj(|jso| {
        jso.encode_field("total", |js| js.encode_i32(31));
        jso.encode_field("next", |js| js.encode_str("abc"));
    });
    assert_eq!(b"{\"total\":31,\"next\":\"abc\"}", &js.into_vec()[..]);

    let mut js = JsonEncoder::new();
    js.encode_obj(|jso| {
        jso.encode_field("total", |js| js.encode_i32(31));
        jso.encode_field("next", |js| js.encode_str("abc"));
        jso.encode_field("tags", |js| {
                js.encode_array(|jsa| {
                        jsa.encode_elm(|js| js.encode_i32(1));
                        jsa.encode_elm(|js| js.encode_i32(2));
                });
        });
    });
    assert_eq!("{\"total\":31,\"next\":\"abc\",\"tags\":[1,2]}", str::from_utf8(&js.into_vec()[..]).unwrap());
}
<commit_msg>Add with_buf() and minor name changes<commit_after>#![feature(vec_push_all)]

pub struct JsonEncoder {
    buffer: Vec<u8>
}

impl JsonEncoder {
    pub fn new() -> JsonEncoder {
        JsonEncoder{buffer: Vec::new()}
    }

    #[inline(always)]
    pub fn encode_raw(&mut self, raw: &[u8]) {
        self.buffer.push_all(raw);
    }

    pub fn into_vec(self) -> Vec<u8> {
        self.buffer
    }

    pub fn with_buf<F>(&mut self, f: F) where F: Fn(&mut Vec<u8>) {
        f(&mut self.buffer);
    }

    #[inline(always)]
    pub fn encode_raw_str(&mut self, raw_str: &str) {
        self.buffer.push_all(raw_str.as_bytes());
    }

    fn escape_bytes(&mut self, bytes: &[u8]) {
        let mut start = 0;

        for (i, byte) in bytes.iter().enumerate() {
            let escaped = match *byte {
                b'"' => b"\\\"",
                b'\\' => b"\\\\",
                b'\x08' => b"\\b",
                b'\x0c' => b"\\f",
                b'\n' => b"\\n",
                b'\r' => b"\\r",
                b'\t' => b"\\t",
                _ => {
                    continue;
                }
            };

            if start < i {
                self.encode_raw(&bytes[start..i]);
            }

            self.encode_raw(escaped);

            start = i + 1;
        }

        if start != bytes.len() {
            self.encode_raw(&bytes[start..]);
        }
    }

    pub fn encode_decimal_str(&mut self, value: u64) {
        self.buffer.push(b'"');
        self.encode_u64_decimal(value);
        self.buffer.push(b'"');
    }

    // encodes as decimal string
    #[inline(always)]
    fn encode_u64_decimal(&mut self, value: u64) {
        const CHARS: &'static [u8] = b"0123456789";
        const MAX_DIGITS: usize = 20;

        if value == 0 {
            self.buffer.push(b'0');
            return;
        }

        let mut digits: [u8; MAX_DIGITS] = [b'0'; MAX_DIGITS];
        let mut i = MAX_DIGITS;
        let mut n = value;
        while n > 0 {
            i -= 1;
            digits[i] = CHARS[(n % 10) as usize];
            n = n / 10;
        }

        self.encode_raw(&digits[i..]);
    }

    #[inline(always)]
    pub fn encode_str(&mut self, s: &str) {
        self.buffer.push(b'"');
        self.escape_bytes(s.as_bytes());
        self.buffer.push(b'"');
    }

    #[inline(always)]
    pub fn encode_i32(&mut self, value: i32) {
        if value == 0 {
            self.buffer.push(b'0');
        } else if value > 0 {
            self.encode_u31(value as u32);
        } else {
            self.buffer.push(b'-');
            self.encode_u31((-value) as u32);
        }
    }

    // encodes a 31-bit unsigned integer != 0
    fn encode_u31(&mut self, value: u32) {
        const CHARS: &'static [u8] = b"0123456789";
        const MAX_DIGITS: usize = 10;
        debug_assert!(value != 0);

        let mut digits: [u8; MAX_DIGITS] = [b'0'; MAX_DIGITS];
        let mut i = MAX_DIGITS;
        let mut n = value;
        while n > 0 {
            i -= 1;
            digits[i] = CHARS[(n % 10) as usize];
            n = n / 10;
        }

        self.encode_raw(&digits[i..]);
    }

    pub fn encode_obj<F>(&mut self, f: F) where F: Fn(&mut JsonObjectEncoder) {
        self.buffer.push(b'{');
        {
            f(&mut JsonObjectEncoder {js: self, needs_sep: false});
        }
        self.buffer.push(b'}');
    }

    pub fn encode_array<F>(&mut self, f: F) where F: Fn(&mut JsonArrayEncoder) {
        self.buffer.push(b'[');
        {
            f(&mut JsonArrayEncoder {js: self, needs_sep: false});
        }
        self.buffer.push(b']');
    }

}

pub struct JsonObjectEncoder<'a> {
    js: &'a mut JsonEncoder,
    needs_sep: bool,
}

impl<'a> JsonObjectEncoder<'a> {
    // XXX: name MAY NOT include escapable characters
    #[inline(always)]
    pub fn encode_field<F:Fn(&mut JsonEncoder)>(&mut self, name: &str, f: F) {
        if self.needs_sep {
            self.js.buffer.push(b',');
        } else {
            self.needs_sep = true;
        }
        self.js.buffer.push(b'"');
        self.js.encode_raw_str(name);
        self.js.buffer.push(b'"');
        self.js.buffer.push(b':');

        f(self.js);
    }
}

pub struct JsonArrayEncoder<'a> {
    js: &'a mut JsonEncoder,
    needs_sep: bool,
}

impl<'a> JsonArrayEncoder<'a> {
    #[inline(always)]
    pub fn encode_elm<F:Fn(&mut JsonEncoder)>(&mut self, f: F) {
        if self.needs_sep {
            self.js.buffer.push(b',');
        } else {
            self.needs_sep = true;
        }
        f(self.js);
    }
}

#[test]
fn test_json_obj_encoder() {
    use std::str;

    let mut js = JsonEncoder::new();
    js.encode_obj(|_|{});
    assert_eq!(b"{}", &js.into_vec()[..]);

    let mut js = JsonEncoder::new();
    js.encode_obj(|jso| {
        jso.encode_field("total", |js| js.encode_i32(31));
    });
    assert_eq!(b"{\"total\":31}", &js.into_vec()[..]);

    let mut js = JsonEncoder::new();
    js.encode_obj(|jso| {
        jso.encode_field("total", |js| js.encode_i32(31));
        jso.encode_field("next", |js| js.encode_str("abc"));
    });
    assert_eq!(b"{\"total\":31,\"next\":\"abc\"}", &js.into_vec()[..]);

    let mut js = JsonEncoder::new();
    js.encode_obj(|jso| {
        jso.encode_field("total", |js| js.encode_i32(31));
        jso.encode_field("next", |js| js.encode_str("abc"));
        jso.encode_field("tags", |js| {
                js.encode_array(|jsa| {
                        jsa.encode_elm(|js| js.encode_i32(1));
                        jsa.encode_elm(|js| js.encode_i32(2));
                });
        });
    });
    assert_eq!("{\"total\":31,\"next\":\"abc\",\"tags\":[1,2]}", str::from_utf8(&js.into_vec()[..]).unwrap());
}
<|endoftext|>


<commit_before>//! # Mime
//!
//! Mime is now Media Type, technically, but `Mime` is more immediately
//! understandable, so the main type here is `Mime`.
//!
//! ## What is Mime?
//!
//! Example mime string: `text/plain;charset=utf-8`
//!
//! ```rust
//! # #[macro_use] extern crate mime;
//! # fn main() {
//! let plain_text: mime::Mime = "text/plain;charset=utf-8".parse().unwrap();
//! assert_eq!(plain_text, mime!(Text/Plain; Charset=Utf8));
//! # }
//! ```

#![doc(html_root_url = "https://hyperium.github.io/mime.rs")]
#![cfg_attr(test, deny(warnings))]
#![cfg_attr(all(feature = "nightly", test), feature(test))]

#[macro_use]
extern crate log;

#[cfg(feature = "nightly")]
#[cfg(test)]
extern crate test;

use std::ascii::AsciiExt;
use std::fmt;
use std::iter::Enumerate;
use std::str::{FromStr, Chars};

macro_rules! inspect(
    ($s:expr, $t:expr) => ({
        let t = $t;
        trace!("inspect {}: {:?}", $s, t);
        t
    })
);

/// Mime, or Media Type. Encapsulates common registers types.
///
/// Consider that a traditional mime type contains a "top level type",
/// a "sub level type", and 0-N "parameters". And they're all strings.
/// Strings everywhere. Strings mean typos. Rust has type safety. We should
/// use types!
///
/// So, Mime bundles together this data into types so the compiler can catch
/// your typos.
///
/// This improves things so you use match without Strings:
///
/// ```rust
/// use mime::{Mime, TopLevel, SubLevel};
///
/// let mime: Mime = "application/json".parse().unwrap();
///
/// match mime {
///     Mime(TopLevel::Application, SubLevel::Json, _) => println!("matched json!"),
///     _ => ()
/// }
/// ```
#[derive(Clone, Debug)]
pub struct Mime<T: AsRef<[Param]> = Vec<Param>>(pub TopLevel, pub SubLevel, pub T);

impl<LHS: AsRef<[Param]>, RHS: AsRef<[Param]>> PartialEq<Mime<RHS>> for Mime<LHS> {
    fn eq(&self, other: &Mime<RHS>) -> bool {
        self.0 == other.0 && self.1 == other.1 && self.2.as_ref() == other.2.as_ref()
    }
}

/// Easily create a Mime without having to import so many enums.
///
/// # Example
///
/// ```
/// # #[macro_use] extern crate mime;
///
/// # fn main() {
/// let json = mime!(Application/Json);
/// let plain = mime!(Text/Plain; Charset=Utf8);
/// let text = mime!(Text/Html; Charset=("bar"), ("baz")=("quux"));
/// let img = mime!(Image/_);
/// # }
/// ```
#[macro_export]
macro_rules! mime {
    ($top:tt / $sub:tt) => (
        mime!($top / $sub;)
    );

    ($top:tt / $sub:tt ; $($attr:tt = $val:tt),*) => (
        $crate::Mime(
            __mime__ident_or_ext!(TopLevel::$top),
            __mime__ident_or_ext!(SubLevel::$sub),
            [ $((__mime__ident_or_ext!(Attr::$attr), __mime__ident_or_ext!(Value::$val))),* ]
        )
    );
}

#[doc(hidden)]
#[macro_export]
macro_rules! __mime__ident_or_ext {
    ($enoom:ident::_) => (
        $crate::$enoom::Star
    );
    ($enoom:ident::($inner:expr)) => (
        $crate::$enoom::Ext($inner.to_string())
    );
    ($enoom:ident::$var:ident) => (
        $crate::$enoom::$var
    )
}

macro_rules! enoom {
    (pub enum $en:ident; $ext:ident; $($ty:ident, $text:expr;)*) => (

        #[derive(Clone, Debug)]
        pub enum $en {
            $($ty),*,
            $ext(String)
        }

        impl PartialEq for $en {
            fn eq(&self, other: &$en) -> bool {
                match (self, other) {
                    $( (&$en::$ty, &$en::$ty) => true ),*,
                    (&$en::$ext(ref a), &$en::$ext(ref b)) => a == b,
                    _ => self.to_string() == other.to_string()
                }
            }
        }

        impl fmt::Display for $en {
            fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
                fmt.write_str(match *self {
                    $($en::$ty => $text),*,
                    $en::$ext(ref s) => s
                })
            }
        }

        impl FromStr for $en {
            type Err = ();
            fn from_str(s: &str) -> Result<$en, ()> {
                Ok(match s {
                    $(_s if _s == $text => $en::$ty),*,
                    s => $en::$ext(inspect!(stringify!($ext), s).to_string())
                })
            }
        }
    )
}

enoom! {
    pub enum TopLevel;
    Ext;
    Star, "*";
    Text, "text";
    Image, "image";
    Audio, "audio";
    Video, "video";
    Application, "application";
    Multipart, "multipart";
    Message, "message";
    Model, "model";
}

enoom! {
    pub enum SubLevel;
    Ext;
    Star, "*";

    // common text/*
    Plain, "plain";
    Html, "html";
    Xml, "xml";
    Javascript, "javascript";
    Css, "css";

    // common application/*
    Json, "json";
    WwwFormUrlEncoded, "x-www-form-urlencoded";

    // multipart/*
    FormData, "form-data";

    // common image/*
    Png, "png";
    Gif, "gif";
    Bmp, "bmp";
    Jpeg, "jpeg";
}

enoom! {
    pub enum Attr;
    Ext;
    Charset, "charset";
    Boundary, "boundary";
    Q, "q";
}

enoom! {
    pub enum Value;
    Ext;
    Utf8, "utf-8";
}

pub type Param = (Attr, Value);

impl<T: AsRef<[Param]>> fmt::Display for Mime<T> {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        let Mime(ref top, ref sub, ref params) = *self;
        try!(write!(fmt, "{}/{}", top, sub));
        fmt_params(params.as_ref(), fmt)
    }
}

impl FromStr for Mime {
    type Err = ();
    fn from_str(raw: &str) -> Result<Mime, ()> {
        let ascii = raw.to_ascii_lowercase(); // lifetimes :(
        let len = ascii.len();
        let mut iter = ascii.chars().enumerate();
        let mut params = vec![];
        // toplevel
        let mut start;
        let mut top;
        loop {
            match inspect!("top iter", iter.next()) {
                Some((0, c)) if is_restricted_name_first_char(c) => (),
                Some((i, c)) if i > 0 && is_restricted_name_char(c) => (),
                Some((i, '/')) if i > 0 => match FromStr::from_str(&ascii[..i]) {
                    Ok(t) => {
                        top = t;
                        start = i + 1;
                        break;
                    }
                    Err(_) => return Err(())
                },
                _ => return Err(()) // EOF and no toplevel is no Mime
            };

        }

        // sublevel
        let mut sub;
        loop {
            match inspect!("sub iter", iter.next()) {
                Some((i, c)) if i == start && is_restricted_name_first_char(c) => (),
                Some((i, c)) if i > start && is_restricted_name_char(c) => (),
                Some((i, ';')) if i > start => match FromStr::from_str(&ascii[start..i]) {
                    Ok(s) => {
                        sub = s;
                        start = i + 1;
                        break;
                    }
                    Err(_) => return Err(())
                },
                None => match FromStr::from_str(&ascii[start..]) {
                    Ok(s) => return Ok(Mime(top, s, params)),
                    Err(_) => return Err(())
                },
                _ => return Err(())
            };
        }

        // params
        debug!("starting params, len={}", len);
        loop {
            match inspect!("param", param_from_str(raw, &ascii, &mut iter, start)) {
                Some((p, end)) => {
                    params.push(p);
                    start = end;
                    if start >= len {
                        break;
                    }
                }
                None => break
            }
        }

        Ok(Mime(top, sub, params))
    }
}

fn param_from_str(raw: &str, ascii: &str, iter: &mut Enumerate<Chars>, mut start: usize) -> Option<(Param, usize)> {
    let mut attr;
    debug!("param_from_str, start={}", start);
    loop {
        match inspect!("attr iter", iter.next()) {
            Some((i, ' ')) if i == start => start = i + 1,
            Some((i, c)) if i == start && is_restricted_name_first_char(c) => (),
            Some((i, c)) if i > start && is_restricted_name_char(c) => (),
            Some((i, '=')) if i > start => match FromStr::from_str(&ascii[start..i]) {
                Ok(a) => {
                    attr = inspect!("attr", a);
                    start = i + 1;
                    break;
                },
                Err(_) => return None
            },
            _ => return None
        }
    }

    let mut value;
    // values must be restrict-name-char or "anything goes"
    let mut is_quoted = false;

    {
        let substr = |a,b| { if attr==Attr::Charset { &ascii[a..b] } else { &raw[a..b] } };
        let endstr = |a| { if attr==Attr::Charset { &ascii[a..] } else { &raw[a..] } };
        loop {
            match inspect!("value iter", iter.next()) {
                Some((i, '"')) if i == start => {
                    debug!("quoted");
                    is_quoted = true;
                    start = i + 1;
                },
                Some((i, c)) if i == start && is_restricted_name_first_char(c) => (),
                Some((i, '"')) if i > start && is_quoted => match FromStr::from_str(substr(start,i)) {
                    Ok(v) => {
                        value = v;
                        start = i + 1;
                        break;
                    },
                    Err(_) => return None
                },
                Some((i, c)) if i > start && is_quoted || is_restricted_name_char(c) => (),
                Some((i, ';')) if i > start => match FromStr::from_str(substr(start,i)) {
                    Ok(v) => {
                        value = v;
                        start = i + 1;
                        break;
                    },
                    Err(_) => return None
                },
                None => match FromStr::from_str(endstr(start)) {
                    Ok(v) => {
                        value = v;
                        start = raw.len();
                        break;
                    },
                    Err(_) => return None
                },

                _ => return None
            }
        }
    }

    Some(((attr, value), start))
}

// From [RFC6838](http://tools.ietf.org/html/rfc6838#section-4.2):
//
// > All registered media types MUST be assigned top-level type and
// > subtype names.  The combination of these names serves to uniquely
// > identify the media type, and the subtype name facet (or the absence
// > of one) identifies the registration tree.  Both top-level type and
// > subtype names are case-insensitive.
// >
// > Type and subtype names MUST conform to the following ABNF:
// >
// >     type-name = restricted-name
// >     subtype-name = restricted-name
// >
// >     restricted-name = restricted-name-first *126restricted-name-chars
// >     restricted-name-first  = ALPHA / DIGIT
// >     restricted-name-chars  = ALPHA / DIGIT / "!" / "#" /
// >                              "$" / "&" / "-" / "^" / "_"
// >     restricted-name-chars =/ "." ; Characters before first dot always
// >                                  ; specify a facet name
// >     restricted-name-chars =/ "+" ; Characters after last plus always
// >                                  ; specify a structured syntax suffix
//
fn is_restricted_name_first_char(c: char) -> bool {
    match c {
        'a'...'z' |
        '0'...'9' => true,
        _ => false
    }
}

fn is_restricted_name_char(c: char) -> bool {
    if is_restricted_name_first_char(c) {
        true
    } else {
        match c {
            '!' |
            '#' |
            '$' |
            '&' |
            '-' |
            '^' |
            '.' |
            '+' |
            '_' => true,
            _ => false
        }
    }
}

#[inline]
fn fmt_params(params: &[Param], fmt: &mut fmt::Formatter) -> fmt::Result {
    for param in params.iter() {
        try!(fmt_param(param, fmt));
    }
    Ok(())
}

#[inline]
fn fmt_param(param: &Param, fmt: &mut fmt::Formatter) -> fmt::Result {
    let (ref attr, ref value) = *param;
    write!(fmt, "; {}={}", attr, value)
}

#[cfg(test)]
mod tests {
    use std::str::FromStr;
    #[cfg(feature = "nightly")]
    use test::Bencher;
    use super::Mime;

    #[test]
    fn test_mime_show() {
        let mime = mime!(Text/Plain);
        assert_eq!(mime.to_string(), "text/plain".to_string());
        let mime = mime!(Text/Plain; Charset=Utf8);
        assert_eq!(mime.to_string(), "text/plain; charset=utf-8".to_string());
    }

    #[test]
    fn test_mime_from_str() {
        assert_eq!(Mime::from_str("text/plain").unwrap(), mime!(Text/Plain));
        assert_eq!(Mime::from_str("TEXT/PLAIN").unwrap(), mime!(Text/Plain));
        assert_eq!(Mime::from_str("text/plain; charset=utf-8").unwrap(), mime!(Text/Plain; Charset=Utf8));
        assert_eq!(Mime::from_str("text/plain;charset=\"utf-8\"").unwrap(), mime!(Text/Plain; Charset=Utf8));
        assert_eq!(Mime::from_str("text/plain; charset=utf-8; foo=bar").unwrap(),
            mime!(Text/Plain; Charset=Utf8, ("foo")=("bar")));
    }

    #[test]
    fn test_case_sensitive_values() {
        assert_eq!(Mime::from_str("multipart/form-data; boundary=ABCDEFG").unwrap(),
                   mime!(Multipart/FormData; Boundary=("ABCDEFG")));
        assert_eq!(Mime::from_str("multipart/form-data; charset=BASE64; boundary=ABCDEFG").unwrap(),
                   mime!(Multipart/FormData; Charset=("base64"), Boundary=("ABCDEFG")));
    }

    #[cfg(feature = "nightly")]
    #[bench]
    fn bench_show(b: &mut Bencher) {
        let mime = mime!(Text/Plain; Charset=Utf8, ("foo")=("bar"));
        b.bytes = mime.to_string().as_bytes().len() as u64;
        b.iter(|| mime.to_string())
    }

    #[cfg(feature = "nightly")]
    #[bench]
    fn bench_from_str(b: &mut Bencher) {
        let s = "text/plain; charset=utf-8; foo=bar";
        b.bytes = s.as_bytes().len() as u64;
        b.iter(|| s.parse::<Mime>())
    }
}
<commit_msg>make mime! use vec! to not break compat<commit_after>//! # Mime
//!
//! Mime is now Media Type, technically, but `Mime` is more immediately
//! understandable, so the main type here is `Mime`.
//!
//! ## What is Mime?
//!
//! Example mime string: `text/plain;charset=utf-8`
//!
//! ```rust
//! # #[macro_use] extern crate mime;
//! # fn main() {
//! let plain_text: mime::Mime = "text/plain;charset=utf-8".parse().unwrap();
//! assert_eq!(plain_text, mime!(Text/Plain; Charset=Utf8));
//! # }
//! ```

#![doc(html_root_url = "https://hyperium.github.io/mime.rs")]
#![cfg_attr(test, deny(warnings))]
#![cfg_attr(all(feature = "nightly", test), feature(test))]

#[macro_use]
extern crate log;

#[cfg(feature = "nightly")]
#[cfg(test)]
extern crate test;

use std::ascii::AsciiExt;
use std::fmt;
use std::iter::Enumerate;
use std::str::{FromStr, Chars};

macro_rules! inspect(
    ($s:expr, $t:expr) => ({
        let t = $t;
        trace!("inspect {}: {:?}", $s, t);
        t
    })
);

/// Mime, or Media Type. Encapsulates common registers types.
///
/// Consider that a traditional mime type contains a "top level type",
/// a "sub level type", and 0-N "parameters". And they're all strings.
/// Strings everywhere. Strings mean typos. Rust has type safety. We should
/// use types!
///
/// So, Mime bundles together this data into types so the compiler can catch
/// your typos.
///
/// This improves things so you use match without Strings:
///
/// ```rust
/// use mime::{Mime, TopLevel, SubLevel};
///
/// let mime: Mime = "application/json".parse().unwrap();
///
/// match mime {
///     Mime(TopLevel::Application, SubLevel::Json, _) => println!("matched json!"),
///     _ => ()
/// }
/// ```
#[derive(Clone, Debug)]
pub struct Mime<T: AsRef<[Param]> = Vec<Param>>(pub TopLevel, pub SubLevel, pub T);

impl<LHS: AsRef<[Param]>, RHS: AsRef<[Param]>> PartialEq<Mime<RHS>> for Mime<LHS> {
    fn eq(&self, other: &Mime<RHS>) -> bool {
        self.0 == other.0 && self.1 == other.1 && self.2.as_ref() == other.2.as_ref()
    }
}

/// Easily create a Mime without having to import so many enums.
///
/// # Example
///
/// ```
/// # #[macro_use] extern crate mime;
///
/// # fn main() {
/// let json = mime!(Application/Json);
/// let plain = mime!(Text/Plain; Charset=Utf8);
/// let text = mime!(Text/Html; Charset=("bar"), ("baz")=("quux"));
/// let img = mime!(Image/_);
/// # }
/// ```
#[macro_export]
macro_rules! mime {
    ($top:tt / $sub:tt) => (
        mime!($top / $sub;)
    );

    ($top:tt / $sub:tt ; $($attr:tt = $val:tt),*) => (
        $crate::Mime(
            __mime__ident_or_ext!(TopLevel::$top),
            __mime__ident_or_ext!(SubLevel::$sub),
            vec![ $((__mime__ident_or_ext!(Attr::$attr), __mime__ident_or_ext!(Value::$val))),* ]
        )
    );
}

#[doc(hidden)]
#[macro_export]
macro_rules! __mime__ident_or_ext {
    ($enoom:ident::_) => (
        $crate::$enoom::Star
    );
    ($enoom:ident::($inner:expr)) => (
        $crate::$enoom::Ext($inner.to_string())
    );
    ($enoom:ident::$var:ident) => (
        $crate::$enoom::$var
    )
}

macro_rules! enoom {
    (pub enum $en:ident; $ext:ident; $($ty:ident, $text:expr;)*) => (

        #[derive(Clone, Debug)]
        pub enum $en {
            $($ty),*,
            $ext(String)
        }

        impl PartialEq for $en {
            fn eq(&self, other: &$en) -> bool {
                match (self, other) {
                    $( (&$en::$ty, &$en::$ty) => true ),*,
                    (&$en::$ext(ref a), &$en::$ext(ref b)) => a == b,
                    _ => self.to_string() == other.to_string()
                }
            }
        }

        impl fmt::Display for $en {
            fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
                fmt.write_str(match *self {
                    $($en::$ty => $text),*,
                    $en::$ext(ref s) => s
                })
            }
        }

        impl FromStr for $en {
            type Err = ();
            fn from_str(s: &str) -> Result<$en, ()> {
                Ok(match s {
                    $(_s if _s == $text => $en::$ty),*,
                    s => $en::$ext(inspect!(stringify!($ext), s).to_string())
                })
            }
        }
    )
}

enoom! {
    pub enum TopLevel;
    Ext;
    Star, "*";
    Text, "text";
    Image, "image";
    Audio, "audio";
    Video, "video";
    Application, "application";
    Multipart, "multipart";
    Message, "message";
    Model, "model";
}

enoom! {
    pub enum SubLevel;
    Ext;
    Star, "*";

    // common text/*
    Plain, "plain";
    Html, "html";
    Xml, "xml";
    Javascript, "javascript";
    Css, "css";

    // common application/*
    Json, "json";
    WwwFormUrlEncoded, "x-www-form-urlencoded";

    // multipart/*
    FormData, "form-data";

    // common image/*
    Png, "png";
    Gif, "gif";
    Bmp, "bmp";
    Jpeg, "jpeg";
}

enoom! {
    pub enum Attr;
    Ext;
    Charset, "charset";
    Boundary, "boundary";
    Q, "q";
}

enoom! {
    pub enum Value;
    Ext;
    Utf8, "utf-8";
}

pub type Param = (Attr, Value);

impl<T: AsRef<[Param]>> fmt::Display for Mime<T> {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        let Mime(ref top, ref sub, ref params) = *self;
        try!(write!(fmt, "{}/{}", top, sub));
        fmt_params(params.as_ref(), fmt)
    }
}

impl FromStr for Mime {
    type Err = ();
    fn from_str(raw: &str) -> Result<Mime, ()> {
        let ascii = raw.to_ascii_lowercase(); // lifetimes :(
        let len = ascii.len();
        let mut iter = ascii.chars().enumerate();
        let mut params = vec![];
        // toplevel
        let mut start;
        let mut top;
        loop {
            match inspect!("top iter", iter.next()) {
                Some((0, c)) if is_restricted_name_first_char(c) => (),
                Some((i, c)) if i > 0 && is_restricted_name_char(c) => (),
                Some((i, '/')) if i > 0 => match FromStr::from_str(&ascii[..i]) {
                    Ok(t) => {
                        top = t;
                        start = i + 1;
                        break;
                    }
                    Err(_) => return Err(())
                },
                _ => return Err(()) // EOF and no toplevel is no Mime
            };

        }

        // sublevel
        let mut sub;
        loop {
            match inspect!("sub iter", iter.next()) {
                Some((i, c)) if i == start && is_restricted_name_first_char(c) => (),
                Some((i, c)) if i > start && is_restricted_name_char(c) => (),
                Some((i, ';')) if i > start => match FromStr::from_str(&ascii[start..i]) {
                    Ok(s) => {
                        sub = s;
                        start = i + 1;
                        break;
                    }
                    Err(_) => return Err(())
                },
                None => match FromStr::from_str(&ascii[start..]) {
                    Ok(s) => return Ok(Mime(top, s, params)),
                    Err(_) => return Err(())
                },
                _ => return Err(())
            };
        }

        // params
        debug!("starting params, len={}", len);
        loop {
            match inspect!("param", param_from_str(raw, &ascii, &mut iter, start)) {
                Some((p, end)) => {
                    params.push(p);
                    start = end;
                    if start >= len {
                        break;
                    }
                }
                None => break
            }
        }

        Ok(Mime(top, sub, params))
    }
}

fn param_from_str(raw: &str, ascii: &str, iter: &mut Enumerate<Chars>, mut start: usize) -> Option<(Param, usize)> {
    let mut attr;
    debug!("param_from_str, start={}", start);
    loop {
        match inspect!("attr iter", iter.next()) {
            Some((i, ' ')) if i == start => start = i + 1,
            Some((i, c)) if i == start && is_restricted_name_first_char(c) => (),
            Some((i, c)) if i > start && is_restricted_name_char(c) => (),
            Some((i, '=')) if i > start => match FromStr::from_str(&ascii[start..i]) {
                Ok(a) => {
                    attr = inspect!("attr", a);
                    start = i + 1;
                    break;
                },
                Err(_) => return None
            },
            _ => return None
        }
    }

    let mut value;
    // values must be restrict-name-char or "anything goes"
    let mut is_quoted = false;

    {
        let substr = |a,b| { if attr==Attr::Charset { &ascii[a..b] } else { &raw[a..b] } };
        let endstr = |a| { if attr==Attr::Charset { &ascii[a..] } else { &raw[a..] } };
        loop {
            match inspect!("value iter", iter.next()) {
                Some((i, '"')) if i == start => {
                    debug!("quoted");
                    is_quoted = true;
                    start = i + 1;
                },
                Some((i, c)) if i == start && is_restricted_name_first_char(c) => (),
                Some((i, '"')) if i > start && is_quoted => match FromStr::from_str(substr(start,i)) {
                    Ok(v) => {
                        value = v;
                        start = i + 1;
                        break;
                    },
                    Err(_) => return None
                },
                Some((i, c)) if i > start && is_quoted || is_restricted_name_char(c) => (),
                Some((i, ';')) if i > start => match FromStr::from_str(substr(start,i)) {
                    Ok(v) => {
                        value = v;
                        start = i + 1;
                        break;
                    },
                    Err(_) => return None
                },
                None => match FromStr::from_str(endstr(start)) {
                    Ok(v) => {
                        value = v;
                        start = raw.len();
                        break;
                    },
                    Err(_) => return None
                },

                _ => return None
            }
        }
    }

    Some(((attr, value), start))
}

// From [RFC6838](http://tools.ietf.org/html/rfc6838#section-4.2):
//
// > All registered media types MUST be assigned top-level type and
// > subtype names.  The combination of these names serves to uniquely
// > identify the media type, and the subtype name facet (or the absence
// > of one) identifies the registration tree.  Both top-level type and
// > subtype names are case-insensitive.
// >
// > Type and subtype names MUST conform to the following ABNF:
// >
// >     type-name = restricted-name
// >     subtype-name = restricted-name
// >
// >     restricted-name = restricted-name-first *126restricted-name-chars
// >     restricted-name-first  = ALPHA / DIGIT
// >     restricted-name-chars  = ALPHA / DIGIT / "!" / "#" /
// >                              "$" / "&" / "-" / "^" / "_"
// >     restricted-name-chars =/ "." ; Characters before first dot always
// >                                  ; specify a facet name
// >     restricted-name-chars =/ "+" ; Characters after last plus always
// >                                  ; specify a structured syntax suffix
//
fn is_restricted_name_first_char(c: char) -> bool {
    match c {
        'a'...'z' |
        '0'...'9' => true,
        _ => false
    }
}

fn is_restricted_name_char(c: char) -> bool {
    if is_restricted_name_first_char(c) {
        true
    } else {
        match c {
            '!' |
            '#' |
            '$' |
            '&' |
            '-' |
            '^' |
            '.' |
            '+' |
            '_' => true,
            _ => false
        }
    }
}

#[inline]
fn fmt_params(params: &[Param], fmt: &mut fmt::Formatter) -> fmt::Result {
    for param in params.iter() {
        try!(fmt_param(param, fmt));
    }
    Ok(())
}

#[inline]
fn fmt_param(param: &Param, fmt: &mut fmt::Formatter) -> fmt::Result {
    let (ref attr, ref value) = *param;
    write!(fmt, "; {}={}", attr, value)
}

#[cfg(test)]
mod tests {
    use std::str::FromStr;
    #[cfg(feature = "nightly")]
    use test::Bencher;
    use super::Mime;

    #[test]
    fn test_mime_show() {
        let mime = mime!(Text/Plain);
        assert_eq!(mime.to_string(), "text/plain".to_string());
        let mime = mime!(Text/Plain; Charset=Utf8);
        assert_eq!(mime.to_string(), "text/plain; charset=utf-8".to_string());
    }

    #[test]
    fn test_mime_from_str() {
        assert_eq!(Mime::from_str("text/plain").unwrap(), mime!(Text/Plain));
        assert_eq!(Mime::from_str("TEXT/PLAIN").unwrap(), mime!(Text/Plain));
        assert_eq!(Mime::from_str("text/plain; charset=utf-8").unwrap(), mime!(Text/Plain; Charset=Utf8));
        assert_eq!(Mime::from_str("text/plain;charset=\"utf-8\"").unwrap(), mime!(Text/Plain; Charset=Utf8));
        assert_eq!(Mime::from_str("text/plain; charset=utf-8; foo=bar").unwrap(),
            mime!(Text/Plain; Charset=Utf8, ("foo")=("bar")));
    }

    #[test]
    fn test_case_sensitive_values() {
        assert_eq!(Mime::from_str("multipart/form-data; boundary=ABCDEFG").unwrap(),
                   mime!(Multipart/FormData; Boundary=("ABCDEFG")));
        assert_eq!(Mime::from_str("multipart/form-data; charset=BASE64; boundary=ABCDEFG").unwrap(),
                   mime!(Multipart/FormData; Charset=("base64"), Boundary=("ABCDEFG")));
    }

    #[cfg(feature = "nightly")]
    #[bench]
    fn bench_show(b: &mut Bencher) {
        let mime = mime!(Text/Plain; Charset=Utf8, ("foo")=("bar"));
        b.bytes = mime.to_string().as_bytes().len() as u64;
        b.iter(|| mime.to_string())
    }

    #[cfg(feature = "nightly")]
    #[bench]
    fn bench_from_str(b: &mut Bencher) {
        let s = "text/plain; charset=utf-8; foo=bar";
        b.bytes = s.as_bytes().len() as u64;
        b.iter(|| s.parse::<Mime>())
    }
}
<|endoftext|>


<commit_before>extern crate libc;
use std::ffi::CString;

#[allow(dead_code,
        non_snake_case,
        non_camel_case_types,
        non_upper_case_globals)]
pub mod libpq;

// Implementation of initialization and callbacks.

pub unsafe extern fn init(cb: *mut libpq::OutputPluginCallbacks) {
    (*cb).startup_cb = Some(startup);
    (*cb).begin_cb = Some(begin);
    (*cb).change_cb = Some(change);
    (*cb).commit_cb = Some(commit);
    (*cb).shutdown_cb = Some(shutdown);
}
/*
pub type LogicalOutputPluginInit =
    ::std::option::Option<extern "C" fn(cb: *mut Struct_OutputPluginCallbacks)
                              -> ()>;
*/

extern fn startup(ctx: *mut libpq::Struct_LogicalDecodingContext,
                  options: *mut libpq::OutputPluginOptions,
                  is_init: libpq::_bool) {
    unsafe {
        (*options).output_type = libpq::OUTPUT_PLUGIN_TEXTUAL_OUTPUT;
    }
}
/*
pub type LogicalDecodeStartupCB =
    ::std::option::Option<extern "C" fn(ctx:
                                            *mut Struct_LogicalDecodingContext,
                                        options: *mut OutputPluginOptions,
                                        is_init: _bool) -> ()>;
 */

extern fn begin(ctx: *mut libpq::Struct_LogicalDecodingContext,
                txn: *mut libpq::ReorderBufferTXN) {
    unsafe {
        let is_last = 1;                                  // True in C language
        let s = CString::new("{ \"BEGIN\": %u }").unwrap();
        libpq::OutputPluginPrepareWrite(ctx, is_last);
        libpq::appendStringInfo((*ctx).out, s.as_ptr(), (*txn).xid);
        libpq::OutputPluginWrite(ctx, is_last);
    }
}
/*
pub type LogicalDecodeBeginCB =
    ::std::option::Option<extern "C" fn(arg1:
                                            *mut Struct_LogicalDecodingContext,
                                        txn: *mut ReorderBufferTXN) -> ()>;
 */

extern fn change(ctx: *mut libpq::Struct_LogicalDecodingContext,
                 txn: *mut libpq::ReorderBufferTXN,
                 relation: libpq::Relation,
                 change: *mut libpq::ReorderBufferChange) {
    unsafe {
        let last = 1;                                     // True in C language
        libpq::OutputPluginPrepareWrite(ctx, last);
        append_change(relation, change, (*ctx).out);
        libpq::OutputPluginWrite(ctx, last);
    }
}
/*
pub type LogicalDecodeChangeCB =
    ::std::option::Option<extern "C" fn(arg1:
                                            *mut Struct_LogicalDecodingContext,
                                        txn: *mut ReorderBufferTXN,
                                        relation: Relation,
                                        change: *mut ReorderBufferChange)
                              -> ()>;
 */

extern fn commit(ctx: *mut libpq::Struct_LogicalDecodingContext,
                 txn: *mut libpq::ReorderBufferTXN,
                 lsn: libpq::XLogRecPtr) {
    unsafe {
        let last = 1;                                     // True in C language
        let s = CString::new("{ \"COMMIT\": %u }").unwrap();
        libpq::OutputPluginPrepareWrite(ctx, last);
        libpq::appendStringInfo((*ctx).out, s.as_ptr(), (*txn).xid);
        libpq::OutputPluginWrite(ctx, last);
    }
}
/*
pub type LogicalDecodeCommitCB =
    ::std::option::Option<extern "C" fn(arg1:
                                            *mut Struct_LogicalDecodingContext,
                                        txn: *mut ReorderBufferTXN,
                                        commit_lsn: XLogRecPtr) -> ()>;
 */

extern fn shutdown(ctx: *mut libpq::Struct_LogicalDecodingContext) {
  // Do nothing.
}
/*
pub type LogicalDecodeShutdownCB =
    ::std::option::Option<extern "C" fn(arg1:
                                            *mut Struct_LogicalDecodingContext)
                              -> ()>;

 */

unsafe fn append_change(relation: libpq::Relation,
                        change: *mut libpq::ReorderBufferChange,
                        out: libpq::StringInfo) {
    let tuple_desc = (*relation).rd_att;
    let tuples = (*change).data.tp();
    let tuple_new = (*tuples).newtuple;
    let tuple_old = (*tuples).oldtuple;
    let token = match (*change).action {
        libpq::REORDER_BUFFER_CHANGE_INSERT => "INSERT",
        libpq::REORDER_BUFFER_CHANGE_UPDATE => "UPDATE",
        libpq::REORDER_BUFFER_CHANGE_DELETE => "DELETE",
        _ => panic!("Unrecognized change action!")
    };
    append("{ ", out);
    append("\"", out);
    append(token, out);
    append("\": ", out);
    append_tuple_buf_as_json(tuple_new, tuple_desc, out);
    if !tuple_old.is_null() {
        append(", ", out);
        append(" \"@\": ", out);
        append_tuple_buf_as_json(tuple_old, tuple_desc, out);
    }
    append(" }\n", out);
}

unsafe fn append_tuple_buf_as_json(data: *mut libpq::ReorderBufferTupleBuf,
                                   desc: libpq::TupleDesc,
                                   out: libpq::StringInfo) {
    if !data.is_null() {
        let heap_tuple = &mut (*data).tuple;
        let datum = libpq::heap_copy_tuple_as_datum(heap_tuple, desc);
        let empty_oid: libpq::Oid = 0;
        let json = libpq::DirectFunctionCall1Coll(Some(row_to_json),
                                                  empty_oid,
                                                  datum);
        let json_output_function: libpq::Oid = 322;     // TODO: Dynamic lookup
        let text = libpq::OidOutputFunctionCall(json_output_function, json);
        libpq::appendStringInfoString(out, text);
    } else {
        append("{}", out);
    }
}

unsafe fn append<T: Into<Vec<u8>>>(t: T, out: libpq::StringInfo) {
    libpq::appendStringInfoString(out, CString::new(t).unwrap().as_ptr());
}

extern fn row_to_json(fcinfo: libpq::FunctionCallInfo) -> libpq::Datum {
    // We wrap the unsafe call to make it safe.
    unsafe {
        libpq::row_to_json(fcinfo)
    }
}

// Symbols Postgres needs to find.

#[allow(non_snake_case)]
#[no_mangle]
pub unsafe extern fn _PG_init() { }

#[allow(non_snake_case)]
#[no_mangle]
pub unsafe extern fn
    _PG_output_plugin_init(cb: *mut libpq::OutputPluginCallbacks) { init(cb); }
<commit_msg>Lowercasify everything<commit_after>extern crate libc;
use std::ffi::CString;

#[allow(dead_code,
        non_snake_case,
        non_camel_case_types,
        non_upper_case_globals)]
pub mod libpq;

// Implementation of initialization and callbacks.

pub unsafe extern fn init(cb: *mut libpq::OutputPluginCallbacks) {
    (*cb).startup_cb = Some(startup);
    (*cb).begin_cb = Some(begin);
    (*cb).change_cb = Some(change);
    (*cb).commit_cb = Some(commit);
    (*cb).shutdown_cb = Some(shutdown);
}
/*
pub type LogicalOutputPluginInit =
    ::std::option::Option<extern "C" fn(cb: *mut Struct_OutputPluginCallbacks)
                              -> ()>;
*/

extern fn startup(ctx: *mut libpq::Struct_LogicalDecodingContext,
                  options: *mut libpq::OutputPluginOptions,
                  is_init: libpq::_bool) {
    unsafe {
        (*options).output_type = libpq::OUTPUT_PLUGIN_TEXTUAL_OUTPUT;
    }
}
/*
pub type LogicalDecodeStartupCB =
    ::std::option::Option<extern "C" fn(ctx:
                                            *mut Struct_LogicalDecodingContext,
                                        options: *mut OutputPluginOptions,
                                        is_init: _bool) -> ()>;
 */

extern fn begin(ctx: *mut libpq::Struct_LogicalDecodingContext,
                txn: *mut libpq::ReorderBufferTXN) {
    unsafe {
        let is_last = 1;                                  // True in C language
        let s = CString::new("{ \"begin\": %u }").unwrap();
        libpq::OutputPluginPrepareWrite(ctx, is_last);
        libpq::appendStringInfo((*ctx).out, s.as_ptr(), (*txn).xid);
        libpq::OutputPluginWrite(ctx, is_last);
    }
}
/*
pub type LogicalDecodeBeginCB =
    ::std::option::Option<extern "C" fn(arg1:
                                            *mut Struct_LogicalDecodingContext,
                                        txn: *mut ReorderBufferTXN) -> ()>;
 */

extern fn change(ctx: *mut libpq::Struct_LogicalDecodingContext,
                 txn: *mut libpq::ReorderBufferTXN,
                 relation: libpq::Relation,
                 change: *mut libpq::ReorderBufferChange) {
    unsafe {
        let last = 1;                                     // True in C language
        libpq::OutputPluginPrepareWrite(ctx, last);
        append_change(relation, change, (*ctx).out);
        libpq::OutputPluginWrite(ctx, last);
    }
}
/*
pub type LogicalDecodeChangeCB =
    ::std::option::Option<extern "C" fn(arg1:
                                            *mut Struct_LogicalDecodingContext,
                                        txn: *mut ReorderBufferTXN,
                                        relation: Relation,
                                        change: *mut ReorderBufferChange)
                              -> ()>;
 */

extern fn commit(ctx: *mut libpq::Struct_LogicalDecodingContext,
                 txn: *mut libpq::ReorderBufferTXN,
                 lsn: libpq::XLogRecPtr) {
    unsafe {
        let last = 1;                                     // True in C language
        let s = CString::new("{ \"commit\": %u }").unwrap();
        libpq::OutputPluginPrepareWrite(ctx, last);
        libpq::appendStringInfo((*ctx).out, s.as_ptr(), (*txn).xid);
        libpq::OutputPluginWrite(ctx, last);
    }
}
/*
pub type LogicalDecodeCommitCB =
    ::std::option::Option<extern "C" fn(arg1:
                                            *mut Struct_LogicalDecodingContext,
                                        txn: *mut ReorderBufferTXN,
                                        commit_lsn: XLogRecPtr) -> ()>;
 */

extern fn shutdown(ctx: *mut libpq::Struct_LogicalDecodingContext) {
  // Do nothing.
}
/*
pub type LogicalDecodeShutdownCB =
    ::std::option::Option<extern "C" fn(arg1:
                                            *mut Struct_LogicalDecodingContext)
                              -> ()>;

 */

unsafe fn append_change(relation: libpq::Relation,
                        change: *mut libpq::ReorderBufferChange,
                        out: libpq::StringInfo) {
    let tuple_desc = (*relation).rd_att;
    let tuples = (*change).data.tp();
    let tuple_new = (*tuples).newtuple;
    let tuple_old = (*tuples).oldtuple;
    let token = match (*change).action {
        libpq::REORDER_BUFFER_CHANGE_INSERT => "insert",
        libpq::REORDER_BUFFER_CHANGE_UPDATE => "update",
        libpq::REORDER_BUFFER_CHANGE_DELETE => "delete",
        _ => panic!("Unrecognized change action!")
    };
    append("{ ", out);
    append("\"", out);
    append(token, out);
    append("\": ", out);
    append_tuple_buf_as_json(tuple_new, tuple_desc, out);
    if !tuple_old.is_null() {
        append(", \"@\": ", out);
        append_tuple_buf_as_json(tuple_old, tuple_desc, out);
    }
    append(" }", out);
}

unsafe fn append_tuple_buf_as_json(data: *mut libpq::ReorderBufferTupleBuf,
                                   desc: libpq::TupleDesc,
                                   out: libpq::StringInfo) {
    if !data.is_null() {
        let heap_tuple = &mut (*data).tuple;
        let datum = libpq::heap_copy_tuple_as_datum(heap_tuple, desc);
        let empty_oid: libpq::Oid = 0;
        let json = libpq::DirectFunctionCall1Coll(Some(row_to_json),
                                                  empty_oid,
                                                  datum);
        let json_output_function: libpq::Oid = 322;     // TODO: Dynamic lookup
        let text = libpq::OidOutputFunctionCall(json_output_function, json);
        libpq::appendStringInfoString(out, text);
    } else {
        append("{}", out);
    }
}

unsafe fn append<T: Into<Vec<u8>>>(t: T, out: libpq::StringInfo) {
    libpq::appendStringInfoString(out, CString::new(t).unwrap().as_ptr());
}

extern fn row_to_json(fcinfo: libpq::FunctionCallInfo) -> libpq::Datum {
    // We wrap the unsafe call to make it safe.
    unsafe {
        libpq::row_to_json(fcinfo)
    }
}

// Symbols Postgres needs to find.

#[allow(non_snake_case)]
#[no_mangle]
pub unsafe extern fn _PG_init() { }

#[allow(non_snake_case)]
#[no_mangle]
pub unsafe extern fn
    _PG_output_plugin_init(cb: *mut libpq::OutputPluginCallbacks) { init(cb); }
<|endoftext|>