<commit_before>
extern crate std;
extern crate core;
extern crate num;
extern crate uuid;

use std::str::SendStr;

use super::def::*;
use super::serialize::*;
use super::reader::*;
use std::collections::TreeMap;

pub static CQL_VERSION_STRINGS:  [&'static str, .. 3] = ["3.0.0", "3.0.0", "3.0.0"];
pub static CQL_MAX_SUPPORTED_VERSION:u8 = 0x03;

macro_rules! serialize_and_check_io_error(
    ($writer: expr, $obj: ident, $version: expr, $msg: expr) => {
        match $obj.serialize($writer, $version) {
            Err(e) => return Err(RCError { kind: SerializeError, desc: format!("{} -> {}", $msg, e.desc).into_maybe_owned()}),
            _ => ()
        }
    }
)

type PrepsStore = TreeMap<String, Box<CqlPreparedStat>>;

pub struct CqlClient {
    socket: std::io::net::tcp::TcpStream,
    pub version: u8,
    prepared: PrepsStore
}

impl CqlClient {

    fn new(socket: std::io::net::tcp::TcpStream, version: u8) -> CqlClient {
        CqlClient {socket: socket, version: version, prepared: TreeMap::new()}
    }

    fn build_auth<'a>(&self, creds: &'a Vec<SendStr>, stream: i8) -> CqlRequest<'a> {
        return CqlRequest {
            version: self.version,
            flags: 0x00,
            stream: stream,
            opcode: OpcodeOptions,
            body: RequestCred(creds),
        };
    }

    fn build_options(&self) -> CqlRequest {
        return CqlRequest {
            version: self.version,
            flags: 0x00,
            stream: 0x01,
            opcode: OpcodeOptions,
            body: RequestOptions,
        };
    }

    pub fn exec_query(&mut self, query_str: &str, con: Consistency::Consistency) -> RCResult<CqlResponse> {
        let q = CqlRequest {
            version: self.version,
            flags: 0x00,
            stream: 0x01,
            opcode: OpcodeQuery,
            body: RequestQuery(query_str, con, 0)};

        let ref mut socket = self.socket;
        serialize_and_check_io_error!(socket, q, self.version, "Error serializing query");
        let res = read_and_check_io_error!(socket, read_cql_response, "Error reading query");
        Ok(res)
    }

    pub fn exec_prepared(&mut self, ps_id: &str, params: &[CqlValue], con: Consistency::Consistency) -> RCResult<CqlResponse> {
        let preps = match self.prepared.find(&String::from_str(ps_id)) {
            Some(ps) => ps,
            None => return Err(RCError::new(format!("Unknown prepared statement <{}>", ps_id), GenericError))
        };

        let q = CqlRequest {
            version: self.version,
            flags: 0x00,
            stream: 0x01,
            opcode: OpcodeExecute,
            body: RequestExec(*preps, params, con, 0x01),
        };

        serialize_and_check_io_error!(&mut self.socket, q, self.version, "Error serializing prepared statement execution");

        /* Code to debug prepared statements. Write to file the serialization of the request

        let path = Path::new("prepared_data.bin");
        let display = path.display();
        let mut file = match std::io::File::create(&path) {
            Err(why) => fail!("couldn't create {}: {}", display, why.desc),
            Ok(file) => file,
        };

        serialize_and_check_io_error!(&mut file, q, self.version, "Error serializing query");

        */

        Ok(read_and_check_io_error!(self.socket, read_cql_response, "Error reading prepared statement execution result"))
    }

    pub fn prepared_statement(&mut self, query_str: &str, query_id: &str) -> RCResult<()> {
        let q = CqlRequest {
            version: self.version,
            flags: 0x00,
            stream: 0x01,
            opcode: OpcodePrepare,
            body: RequestPrepare(query_str),
        };

        serialize_and_check_io_error!(&mut self.socket, q, self.version, "Error serializing prepared statement");

        let res = read_and_check_io_error!(&mut self.socket, read_cql_response, "Error reading query");
        match res.body {
            ResultPrepared(preps) => {
                self.prepared.insert(String::from_str(query_id), preps);
                Ok(())
            },
            _ => Err(RCError::new("Response does not contain prepared statement", ReadError))
        }
    }
}

fn build_startup(version: u8) -> CqlRequest {
    let body = CqlStringMap {
        pairs:vec![CqlPair{key: "CQL_VERSION", value: CQL_VERSION_STRINGS[(version-1) as uint]}],
    };
    return CqlRequest {
        version: version,
        flags: 0x00,
        stream: 0x01,
        opcode: OpcodeStartup,
        body: RequestStartup(body),
    };
}

fn send_startup(socket: &mut std::io::TcpStream, version: u8, creds: Option<&Vec<SendStr>>) -> RCResult<()> {
        let msg_startup = build_startup(version);
        serialize_and_check_error!(socket, msg_startup, version, "Error serializing startup message");

        let response = read_and_check_io_error!(socket, read_cql_response, "Error reding response");
        match response.body {
            ResponseReady =>  Ok(()),
            ResponseAuth(_) => {
                match creds {
                    Some(cred) => {
                        let msg_auth = CqlRequest {
                            version: version,
                            flags: 0x00,
                            stream: 0x01,
                            opcode: OpcodeOptions,
                            body: RequestCred(cred),
                        };
                        serialize_and_check_error!(socket, msg_auth, version, "Error serializing request (auth)");

                        let response = read_and_check_io_error!(socket, read_cql_response, "Error reding authenticaton response");
                        match response.body {
                            ResponseReady => Ok(()),
                            ResponseError(_, ref msg) => Err(RCError::new(format!("Error in authentication: {}", msg), ReadError)),
                            _ => Err(RCError::new("Server returned unknown message", ReadError))
                        }
                    },
                    None => Err(RCError::new("Credential should be provided for authentication", ReadError))
                }
            },
            ResponseError(_, ref msg) => Err(RCError::new(format!("Error connecting: {}", msg), ReadError)),
            _ => Err(RCError::new("Wrong response to startup", ReadError))
        }
}

pub fn connect(ip: &'static str, port: u16, creds:Option<&Vec<SendStr>>) -> RCResult<CqlClient> {

    let mut version = CQL_MAX_SUPPORTED_VERSION;

    while version >= 0x01 {
        let res = std::io::TcpStream::connect(ip.as_slice(), port);
        if res.is_err() {
            return Err(RCError::new(format!("Failed to connect to server at {}:{}", ip.as_slice(), port), ConnectionError));
        }

        let mut socket = res.unwrap();

        match send_startup(& mut socket, version, creds) {
            Ok(_) => return Ok(CqlClient::new(socket, version)),
            Err(e) => println!("Error connecting with protocol version v{}: {}", version, e.desc)
        }
        version -= 1;
    }
    Err(RCError::new("Unable to find suitable protocol version (v1, v2, v3)", ReadError))
}
<commit_msg>Use version for read_cql_response<commit_after>
extern crate std;
extern crate core;
extern crate num;
extern crate uuid;

use std::str::SendStr;

use super::def::*;
use super::serialize::*;
use super::reader::*;
use std::collections::TreeMap;

pub static CQL_VERSION_STRINGS:  [&'static str, .. 3] = ["3.0.0", "3.0.0", "3.0.0"];
pub static CQL_MAX_SUPPORTED_VERSION:u8 = 0x03;

macro_rules! serialize_and_check_io_error(
    ($writer: expr, $obj: ident, $version: expr, $msg: expr) => {
        match $obj.serialize($writer, $version) {
            Err(e) => return Err(RCError { kind: SerializeError, desc: format!("{} -> {}", $msg, e.desc).into_maybe_owned()}),
            _ => ()
        }
    }
)

type PrepsStore = TreeMap<String, Box<CqlPreparedStat>>;

pub struct CqlClient {
    socket: std::io::net::tcp::TcpStream,
    pub version: u8,
    prepared: PrepsStore
}

impl CqlClient {

    fn new(socket: std::io::net::tcp::TcpStream, version: u8) -> CqlClient {
        CqlClient {socket: socket, version: version, prepared: TreeMap::new()}
    }

    fn build_auth<'a>(&self, creds: &'a Vec<SendStr>, stream: i8) -> CqlRequest<'a> {
        return CqlRequest {
            version: self.version,
            flags: 0x00,
            stream: stream,
            opcode: OpcodeOptions,
            body: RequestCred(creds),
        };
    }

    fn build_options(&self) -> CqlRequest {
        return CqlRequest {
            version: self.version,
            flags: 0x00,
            stream: 0x01,
            opcode: OpcodeOptions,
            body: RequestOptions,
        };
    }

    pub fn exec_query(&mut self, query_str: &str, con: Consistency::Consistency) -> RCResult<CqlResponse> {
        let q = CqlRequest {
            version: self.version,
            flags: 0x00,
            stream: 0x01,
            opcode: OpcodeQuery,
            body: RequestQuery(query_str, con, 0)};

        let ref mut socket = self.socket;
        serialize_and_check_io_error!(socket, q, self.version, "Error serializing query");
        let res = read_and_check_io_error!(socket, read_cql_response, self.version, "Error reading query");
        Ok(res)
    }

    pub fn exec_prepared(&mut self, ps_id: &str, params: &[CqlValue], con: Consistency::Consistency) -> RCResult<CqlResponse> {
        let preps = match self.prepared.find(&String::from_str(ps_id)) {
            Some(ps) => ps,
            None => return Err(RCError::new(format!("Unknown prepared statement <{}>", ps_id), GenericError))
        };

        let q = CqlRequest {
            version: self.version,
            flags: 0x00,
            stream: 0x01,
            opcode: OpcodeExecute,
            body: RequestExec(*preps, params, con, 0x01),
        };

        serialize_and_check_io_error!(&mut self.socket, q, self.version, "Error serializing prepared statement execution");

        /* Code to debug prepared statements. Write to file the serialization of the request

        let path = Path::new("prepared_data.bin");
        let display = path.display();
        let mut file = match std::io::File::create(&path) {
            Err(why) => fail!("couldn't create {}: {}", display, why.desc),
            Ok(file) => file,
        };

        serialize_and_check_io_error!(&mut file, q, self.version, "Error serializing query");

        */

        Ok(read_and_check_io_error!(self.socket, read_cql_response, self.version, "Error reading prepared statement execution result"))
    }

    pub fn prepared_statement(&mut self, query_str: &str, query_id: &str) -> RCResult<()> {
        let q = CqlRequest {
            version: self.version,
            flags: 0x00,
            stream: 0x01,
            opcode: OpcodePrepare,
            body: RequestPrepare(query_str),
        };

        serialize_and_check_io_error!(&mut self.socket, q, self.version, "Error serializing prepared statement");

        let res = read_and_check_io_error!(&mut self.socket, read_cql_response, self.version, "Error reading query");
        match res.body {
            ResultPrepared(preps) => {
                self.prepared.insert(String::from_str(query_id), preps);
                Ok(())
            },
            _ => Err(RCError::new("Response does not contain prepared statement", ReadError))
        }
    }
}

fn build_startup(version: u8) -> CqlRequest {
    let body = CqlStringMap {
        pairs:vec![CqlPair{key: "CQL_VERSION", value: CQL_VERSION_STRINGS[(version-1) as uint]}],
    };
    return CqlRequest {
        version: version,
        flags: 0x00,
        stream: 0x01,
        opcode: OpcodeStartup,
        body: RequestStartup(body),
    };
}

fn send_startup(socket: &mut std::io::TcpStream, version: u8, creds: Option<&Vec<SendStr>>) -> RCResult<()> {
        let msg_startup = build_startup(version);
        serialize_and_check_error!(socket, msg_startup, version, "Error serializing startup message");

        let response = read_and_check_io_error!(socket, read_cql_response, version, "Error reding response");
        match response.body {
            ResponseReady =>  Ok(()),
            ResponseAuth(_) => {
                match creds {
                    Some(cred) => {
                        let msg_auth = CqlRequest {
                            version: version,
                            flags: 0x00,
                            stream: 0x01,
                            opcode: OpcodeOptions,
                            body: RequestCred(cred),
                        };
                        serialize_and_check_error!(socket, msg_auth, version, "Error serializing request (auth)");

                        let response = read_and_check_io_error!(socket, read_cql_response, version, "Error reding authenticaton response");
                        match response.body {
                            ResponseReady => Ok(()),
                            ResponseError(_, ref msg) => Err(RCError::new(format!("Error in authentication: {}", msg), ReadError)),
                            _ => Err(RCError::new("Server returned unknown message", ReadError))
                        }
                    },
                    None => Err(RCError::new("Credential should be provided for authentication", ReadError))
                }
            },
            ResponseError(_, ref msg) => Err(RCError::new(format!("Error connecting: {}", msg), ReadError)),
            _ => Err(RCError::new("Wrong response to startup", ReadError))
        }
}

pub fn connect(ip: &'static str, port: u16, creds:Option<&Vec<SendStr>>) -> RCResult<CqlClient> {

    let mut version = CQL_MAX_SUPPORTED_VERSION;

    while version >= 0x01 {
        let res = std::io::TcpStream::connect(ip.as_slice(), port);
        if res.is_err() {
            return Err(RCError::new(format!("Failed to connect to server at {}:{}", ip.as_slice(), port), ConnectionError));
        }

        let mut socket = res.unwrap();

        match send_startup(& mut socket, version, creds) {
            Ok(_) => return Ok(CqlClient::new(socket, version)),
            Err(e) => println!("Error connecting with protocol version v{}: {}", version, e.desc)
        }
        version -= 1;
    }
    Err(RCError::new("Unable to find suitable protocol version (v1, v2, v3)", ReadError))
}
<|endoftext|>


<commit_before>// Copyright (C) 2018 O.S. Systems Sofware LTDA
//
// SPDX-License-Identifier: MPL-2.0
//

use failure::Error;
use reqwest::Client;
use reqwest::header::{ContentType, Headers, UserAgent};

use std::time::Duration;

use firmware::Metadata;
use runtime_settings::RuntimeSettings;
use settings::Settings;

use update_package::UpdatePackage;

header! { (ApiContentType, "Api-Content-Type") => [String] }
header! { (ApiRetries, "Api-Retries") => [usize] }
header! { (AddExtraPoll, "Add-Extra-Poll") => [i64] }

pub struct Api<'a> {
    settings: &'a Settings,
    firmware: &'a Metadata,
    runtime_settings: &'a RuntimeSettings,
}

#[derive(Debug)]
pub enum ProbeResponse {
    NoUpdate,
    Update(UpdatePackage),
    ExtraPoll(i64),
}

impl<'a> Api<'a> {
    pub fn new(
        settings: &'a Settings,
        runtime_settings: &'a RuntimeSettings,
        firmware: &'a Metadata,
    ) -> Api<'a> {
        Api {
            settings,
            runtime_settings,
            firmware,
        }
    }

    fn client(&self) -> Result<Client, Error> {
        let mut headers = Headers::new();

        headers.set(UserAgent::new("updatehub/next"));
        headers.set(ContentType::json());
        headers.set(ApiContentType("application/vnd.updatehub-v1+json".into()));

        Ok(Client::builder()
            .timeout(Duration::from_secs(10))
            .default_headers(headers)
            .build()?)
    }

    pub fn probe(&self) -> Result<ProbeResponse, Error> {
        use reqwest::StatusCode;

        let mut response = self.client()?
            .post(&format!(
                "{}/upgrades",
                &self.settings.network.server_address
            ))
            .header(ApiRetries(self.runtime_settings.polling.retries))
            .json(&self.firmware)
            .send()?;

        match response.status() {
            StatusCode::NotFound => Ok(ProbeResponse::NoUpdate),
            StatusCode::Ok => {
                if let Some(extra_poll) = response.headers().get::<AddExtraPoll>() {
                    return Ok(ProbeResponse::ExtraPoll(extra_poll.0));
                }

                Ok(ProbeResponse::Update(
                    UpdatePackage::parse(&response.text()?)?,
                ))
            }
            _ => bail!("Invalid response. Status: {}", response.status()),
        }
    }

    pub fn download_object(&self, package_uid: &str, object: &str) -> Result<(), Error> {
        // FIXME: Discuss the need of packages inside the route
        let mut client = self.client()?.get(&format!(
            "{}/products/{}/packages/{}/objects/{}",
            &self.settings.network.server_address, &self.firmware.product_uid, package_uid, object
        ));

        let _ = client.send()?;

        Ok(())
    }
}

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

    pub enum FakeServer {
        NoUpdate,
        HasUpdate,
        ExtraPoll,
        ErrorOnce,
        InvalidHardware,
    }

    pub fn create_mock_server(server: FakeServer) -> Mock {
        use mockito::{mock, Matcher};
        use update_package::tests::get_update_json;

        fn fake_device_reply_body(identity: usize, hardware: &str) -> Matcher {
            Matcher::JSON(json!(
                        {
                            "product_uid": "229ffd7e08721d716163fc81a2dbaf6c90d449f0a3b009b6a2defe8a0b0d7381",
                            "version": "1.1",
                            "hardware": hardware,
                            "device_identity": {
                                "id1":[format!("value{}", identity)],
                                "id2":["value2"]
                            },
                            "device_attributes": {
                                "attr1":["attrvalue1"],
                                "attr2":["attrvalue2"]
                            }
                        }
            ))
        }

        match server {
            FakeServer::NoUpdate => mock("POST", "/upgrades")
                .match_header("Content-Type", "application/json")
                .match_header("Api-Content-Type", "application/vnd.updatehub-v1+json")
                .match_body(fake_device_reply_body(1, "board"))
                .with_status(404)
                .create(),

            FakeServer::HasUpdate => mock("POST", "/upgrades")
                .match_header("Content-Type", "application/json")
                .match_header("Api-Content-Type", "application/vnd.updatehub-v1+json")
                .match_body(fake_device_reply_body(2, "board"))
                .with_status(200)
                .with_body(&get_update_json().to_string())
                .create(),

            FakeServer::ExtraPoll => mock("POST", "/upgrades")
                .match_header("Content-Type", "application/json")
                .match_header("Api-Content-Type", "application/vnd.updatehub-v1+json")
                .match_body(fake_device_reply_body(3, "board"))
                .with_status(200)
                .with_header("Add-Extra-Poll", "10")
                .create(),

            FakeServer::ErrorOnce => mock("POST", "/upgrades")
                .match_header("Content-Type", "application/json")
                .match_header("Api-Content-Type", "application/vnd.updatehub-v1+json")
                .match_header("Api-Retries", "1")
                .match_body(fake_device_reply_body(1, "board"))
                .with_status(404)
                .create(),
            FakeServer::InvalidHardware => mock("POST", "/upgrades")
                .match_header("Content-Type", "application/json")
                .match_header("Api-Content-Type", "application/vnd.updatehub-v1+json")
                .match_body(fake_device_reply_body(4, "invalid"))
                .with_status(200)
                .with_body(&get_update_json().to_string())
                .create(),
        }
    }

    #[test]
    fn probe_requirements() {
        use firmware::tests::{create_fake_metadata, FakeDevice};

        let mock = create_mock_server(FakeServer::NoUpdate);
        let _ = Api::new(
            &Settings::default(),
            &RuntimeSettings::default(),
            &Metadata::new(&create_fake_metadata(FakeDevice::NoUpdate)).unwrap(),
        ).probe();
        mock.assert();
    }

    #[test]
    fn download_object() {
        use firmware::tests::{create_fake_metadata, FakeDevice};
        use mockito::mock;

        let metadata = Metadata::new(&create_fake_metadata(FakeDevice::NoUpdate)).unwrap();

        let m = mock(
            "GET",
            format!(
                "/products/{}/packages/{}/objects/{}",
                metadata.product_uid, "package_id", "object"
            ).as_str(),
        ).match_header("Content-Type", "application/json")
            .match_header("Api-Content-Type", "application/vnd.updatehub-v1+json")
            .with_status(200)
            .create();

        let _ = Api::new(&Settings::default(), &RuntimeSettings::default(), &metadata)
            .download_object("package_id", "object")
            .unwrap();

        m.assert();
    }
}
<commit_msg>client: Add support for object download resume<commit_after>// Copyright (C) 2018 O.S. Systems Sofware LTDA
//
// SPDX-License-Identifier: MPL-2.0
//

use failure::Error;
use reqwest::{Client, StatusCode};
use reqwest::header::{ByteRangeSpec, ContentType, Headers, Range, UserAgent};

use std::time::Duration;

use firmware::Metadata;
use runtime_settings::RuntimeSettings;
use settings::Settings;

use update_package::UpdatePackage;

header! { (ApiContentType, "Api-Content-Type") => [String] }
header! { (ApiRetries, "Api-Retries") => [usize] }
header! { (AddExtraPoll, "Add-Extra-Poll") => [i64] }

pub struct Api<'a> {
    settings: &'a Settings,
    firmware: &'a Metadata,
    runtime_settings: &'a RuntimeSettings,
}

#[derive(Debug)]
pub enum ProbeResponse {
    NoUpdate,
    Update(UpdatePackage),
    ExtraPoll(i64),
}

impl<'a> Api<'a> {
    pub fn new(
        settings: &'a Settings,
        runtime_settings: &'a RuntimeSettings,
        firmware: &'a Metadata,
    ) -> Api<'a> {
        Api {
            settings,
            runtime_settings,
            firmware,
        }
    }

    fn client(&self) -> Result<Client, Error> {
        let mut headers = Headers::new();

        headers.set(UserAgent::new("updatehub/next"));
        headers.set(ContentType::json());
        headers.set(ApiContentType("application/vnd.updatehub-v1+json".into()));

        Ok(Client::builder()
            .timeout(Duration::from_secs(10))
            .default_headers(headers)
            .build()?)
    }

    pub fn probe(&self) -> Result<ProbeResponse, Error> {
        let mut response = self.client()?
            .post(&format!(
                "{}/upgrades",
                &self.settings.network.server_address
            ))
            .header(ApiRetries(self.runtime_settings.polling.retries))
            .json(&self.firmware)
            .send()?;

        match response.status() {
            StatusCode::NotFound => Ok(ProbeResponse::NoUpdate),
            StatusCode::Ok => {
                if let Some(extra_poll) = response.headers().get::<AddExtraPoll>() {
                    return Ok(ProbeResponse::ExtraPoll(extra_poll.0));
                }

                Ok(ProbeResponse::Update(
                    UpdatePackage::parse(&response.text()?)?,
                ))
            }
            _ => bail!("Invalid response. Status: {}", response.status()),
        }
    }

    pub fn download_object(&self, package_uid: &str, object: &str) -> Result<(), Error> {
        use std::fs::{create_dir_all, OpenOptions};

        // FIXME: Discuss the need of packages inside the route
        let mut client = self.client()?.get(&format!(
            "{}/products/{}/packages/{}/objects/{}",
            &self.settings.network.server_address, &self.firmware.product_uid, package_uid, object
        ));

        let path = &self.settings.update.download_dir;
        if !&path.exists() {
            debug!("Creating directory to store the downloads.");
            create_dir_all(&path)?;
        }

        let file = path.join(object);
        if file.exists() {
            client.header(Range::Bytes(vec![
                ByteRangeSpec::AllFrom(file.metadata()?.len() - 1),
            ]));
        }

        let mut file = OpenOptions::new().create(true).append(true).open(&file)?;
        let mut response = client.send()?;
        if response.status().is_success() {
            response.copy_to(&mut file)?;
            return Ok(());
        }

        bail!("Couldn't download the object {}", object)
    }
}

#[cfg(test)]
pub mod tests {
    use super::*;
    use firmware::tests::{create_fake_metadata, FakeDevice};
    use mockito::{mock, Mock};

    pub enum FakeServer {
        NoUpdate,
        HasUpdate,
        ExtraPoll,
        ErrorOnce,
        InvalidHardware,
    }

    pub fn create_mock_server(server: FakeServer) -> Mock {
        use mockito::Matcher;
        use update_package::tests::get_update_json;

        fn fake_device_reply_body(identity: usize, hardware: &str) -> Matcher {
            Matcher::JSON(json!(
                        {
                            "product_uid": "229ffd7e08721d716163fc81a2dbaf6c90d449f0a3b009b6a2defe8a0b0d7381",
                            "version": "1.1",
                            "hardware": hardware,
                            "device_identity": {
                                "id1":[format!("value{}", identity)],
                                "id2":["value2"]
                            },
                            "device_attributes": {
                                "attr1":["attrvalue1"],
                                "attr2":["attrvalue2"]
                            }
                        }
            ))
        }

        match server {
            FakeServer::NoUpdate => mock("POST", "/upgrades")
                .match_header("Content-Type", "application/json")
                .match_header("Api-Content-Type", "application/vnd.updatehub-v1+json")
                .match_body(fake_device_reply_body(1, "board"))
                .with_status(404)
                .create(),

            FakeServer::HasUpdate => mock("POST", "/upgrades")
                .match_header("Content-Type", "application/json")
                .match_header("Api-Content-Type", "application/vnd.updatehub-v1+json")
                .match_body(fake_device_reply_body(2, "board"))
                .with_status(200)
                .with_body(&get_update_json().to_string())
                .create(),

            FakeServer::ExtraPoll => mock("POST", "/upgrades")
                .match_header("Content-Type", "application/json")
                .match_header("Api-Content-Type", "application/vnd.updatehub-v1+json")
                .match_body(fake_device_reply_body(3, "board"))
                .with_status(200)
                .with_header("Add-Extra-Poll", "10")
                .create(),

            FakeServer::ErrorOnce => mock("POST", "/upgrades")
                .match_header("Content-Type", "application/json")
                .match_header("Api-Content-Type", "application/vnd.updatehub-v1+json")
                .match_header("Api-Retries", "1")
                .match_body(fake_device_reply_body(1, "board"))
                .with_status(404)
                .create(),
            FakeServer::InvalidHardware => mock("POST", "/upgrades")
                .match_header("Content-Type", "application/json")
                .match_header("Api-Content-Type", "application/vnd.updatehub-v1+json")
                .match_body(fake_device_reply_body(4, "invalid"))
                .with_status(200)
                .with_body(&get_update_json().to_string())
                .create(),
        }
    }

    #[test]
    fn probe_requirements() {
        use firmware::tests::{create_fake_metadata, FakeDevice};

        let mock = create_mock_server(FakeServer::NoUpdate);
        let _ = Api::new(
            &Settings::default(),
            &RuntimeSettings::default(),
            &Metadata::new(&create_fake_metadata(FakeDevice::NoUpdate)).unwrap(),
        ).probe();
        mock.assert();
    }

    #[test]
    fn download_object() {
        let metadata = Metadata::new(&create_fake_metadata(FakeDevice::NoUpdate)).unwrap();
        use mktemp::Temp;
        use std::fs::File;
        use std::io::Read;

        let m1 = mock(
            "GET",
            format!(
                "/products/{}/packages/{}/objects/{}",
                metadata.product_uid, "package_id", "object"
            ).as_str(),
        ).match_header("Content-Type", "application/json")
            .match_header("Api-Content-Type", "application/vnd.updatehub-v1+json")
            .with_status(200)
            .with_body("1234")
            .create();

        let m2 = mock(
            "GET",
            format!(
                "/products/{}/packages/{}/objects/{}",
                metadata.product_uid, "package_id", "object"
            ).as_str(),
        ).match_header("Content-Type", "application/json")
            .match_header("Api-Content-Type", "application/vnd.updatehub-v1+json")
            .match_header("Range", "bytes=3-")
            .with_status(200)
            .with_body("567890")
            .create();

        let mut settings = Settings::default();
        settings.update.download_dir = Temp::new_dir().unwrap().to_path_buf();

        // Download the object.
        let _ = Api::new(&settings, &RuntimeSettings::default(), &metadata)
            .download_object("package_id", "object")
            .expect("Failed to download the object.");

        // Verify it has been downloaded successfully.
        let mut downloaded = String::new();
        let _ = File::open(&settings.update.download_dir.join("object"))
            .expect("Failed to open destination object.")
            .read_to_string(&mut downloaded);

        m1.assert();

        assert_eq!(downloaded, "1234".to_string());

        // Download the remaining bytes of the object.
        let _ = Api::new(&settings, &RuntimeSettings::default(), &metadata)
            .download_object("package_id", "object")
            .expect("Failed to download the object.");

        // Verify it has been downloaded successfully.
        let mut downloaded = String::new();
        let _ = File::open(&settings.update.download_dir.join("object"))
            .expect("Failed to open destination object.")
            .read_to_string(&mut downloaded);

        m2.assert();

        assert_eq!(downloaded, "1234567890".to_string());
    }
}
<|endoftext|>


<commit_before>use crate::{Applicability, Handler, Level, StashKey};
use crate::{Diagnostic, DiagnosticId, DiagnosticStyledString};

use log::debug;
use rustc_span::{MultiSpan, Span};
use std::fmt::{self, Debug};
use std::ops::{Deref, DerefMut};
use std::thread::panicking;

/// Used for emitting structured error messages and other diagnostic information.
///
/// If there is some state in a downstream crate you would like to
/// access in the methods of `DiagnosticBuilder` here, consider
/// extending `HandlerFlags`, accessed via `self.handler.flags`.
#[must_use]
#[derive(Clone)]
pub struct DiagnosticBuilder<'a>(Box<DiagnosticBuilderInner<'a>>);

/// This is a large type, and often used as a return value, especially within
/// the frequently-used `PResult` type. In theory, return value optimization
/// (RVO) should avoid unnecessary copying. In practice, it does not (at the
/// time of writing). The split between `DiagnosticBuilder` and
/// `DiagnosticBuilderInner` exists to avoid many `memcpy` calls.
#[must_use]
#[derive(Clone)]
struct DiagnosticBuilderInner<'a> {
    handler: &'a Handler,
    diagnostic: Diagnostic,
    allow_suggestions: bool,
}

/// In general, the `DiagnosticBuilder` uses deref to allow access to
/// the fields and methods of the embedded `diagnostic` in a
/// transparent way. *However,* many of the methods are intended to
/// be used in a chained way, and hence ought to return `self`. In
/// that case, we can't just naively forward to the method on the
/// `diagnostic`, because the return type would be a `&Diagnostic`
/// instead of a `&DiagnosticBuilder<'a>`. This `forward!` macro makes
/// it easy to declare such methods on the builder.
macro_rules! forward {
    // Forward pattern for &self -> &Self
    (
        $(#[$attrs:meta])*
        pub fn $n:ident(&self, $($name:ident: $ty:ty),* $(,)?) -> &Self
    ) => {
        $(#[$attrs])*
        pub fn $n(&self, $($name: $ty),*) -> &Self {
            self.diagnostic.$n($($name),*);
            self
        }
    };

    // Forward pattern for &mut self -> &mut Self
    (
        $(#[$attrs:meta])*
        pub fn $n:ident(&mut self, $($name:ident: $ty:ty),* $(,)?) -> &mut Self
    ) => {
        $(#[$attrs])*
        pub fn $n(&mut self, $($name: $ty),*) -> &mut Self {
            self.0.diagnostic.$n($($name),*);
            self
        }
    };

    // Forward pattern for &mut self -> &mut Self, with S: Into<MultiSpan>
    // type parameter. No obvious way to make this more generic.
    (
        $(#[$attrs:meta])*
        pub fn $n:ident<S: Into<MultiSpan>>(
            &mut self,
            $($name:ident: $ty:ty),*
            $(,)?
        ) -> &mut Self
    ) => {
        $(#[$attrs])*
        pub fn $n<S: Into<MultiSpan>>(&mut self, $($name: $ty),*) -> &mut Self {
            self.0.diagnostic.$n($($name),*);
            self
        }
    };
}

impl<'a> Deref for DiagnosticBuilder<'a> {
    type Target = Diagnostic;

    fn deref(&self) -> &Diagnostic {
        &self.0.diagnostic
    }
}

impl<'a> DerefMut for DiagnosticBuilder<'a> {
    fn deref_mut(&mut self) -> &mut Diagnostic {
        &mut self.0.diagnostic
    }
}

impl<'a> DiagnosticBuilder<'a> {
    /// Emit the diagnostic.
    pub fn emit(&mut self) {
        self.0.handler.emit_diagnostic(&self);
        self.cancel();
    }

    /// Emit the diagnostic unless `delay` is true,
    /// in which case the emission will be delayed as a bug.
    ///
    /// See `emit` and `delay_as_bug` for details.
    pub fn emit_unless(&mut self, delay: bool) {
        if delay {
            self.delay_as_bug();
        } else {
            self.emit();
        }
    }

    /// Stashes diagnostic for possible later improvement in a different,
    /// later stage of the compiler. The diagnostic can be accessed with
    /// the provided `span` and `key` through `.steal_diagnostic` on `Handler`.
    ///
    /// As with `buffer`, this is unless the handler has disabled such buffering.
    pub fn stash(self, span: Span, key: StashKey) {
        if let Some((diag, handler)) = self.into_diagnostic() {
            handler.stash_diagnostic(span, key, diag);
        }
    }

    /// Converts the builder to a `Diagnostic` for later emission,
    /// unless handler has disabled such buffering.
    pub fn into_diagnostic(mut self) -> Option<(Diagnostic, &'a Handler)> {
        if self.0.handler.flags.dont_buffer_diagnostics
            || self.0.handler.flags.treat_err_as_bug.is_some()
        {
            self.emit();
            return None;
        }

        let handler = self.0.handler;

        // We need to use `ptr::read` because `DiagnosticBuilder` implements `Drop`.
        let diagnostic;
        unsafe {
            diagnostic = std::ptr::read(&self.0.diagnostic);
            std::mem::forget(self);
        };
        // Logging here is useful to help track down where in logs an error was
        // actually emitted.
        debug!("buffer: diagnostic={:?}", diagnostic);

        Some((diagnostic, handler))
    }

    /// Buffers the diagnostic for later emission,
    /// unless handler has disabled such buffering.
    pub fn buffer(self, buffered_diagnostics: &mut Vec<Diagnostic>) {
        buffered_diagnostics.extend(self.into_diagnostic().map(|(diag, _)| diag));
    }

    /// Convenience function for internal use, clients should use one of the
    /// span_* methods instead.
    pub fn sub<S: Into<MultiSpan>>(
        &mut self,
        level: Level,
        message: &str,
        span: Option<S>,
    ) -> &mut Self {
        let span = span.map(|s| s.into()).unwrap_or_else(|| MultiSpan::new());
        self.0.diagnostic.sub(level, message, span, None);
        self
    }

    /// Delay emission of this diagnostic as a bug.
    ///
    /// This can be useful in contexts where an error indicates a bug but
    /// typically this only happens when other compilation errors have already
    /// happened. In those cases this can be used to defer emission of this
    /// diagnostic as a bug in the compiler only if no other errors have been
    /// emitted.
    ///
    /// In the meantime, though, callsites are required to deal with the "bug"
    /// locally in whichever way makes the most sense.
    pub fn delay_as_bug(&mut self) {
        self.level = Level::Bug;
        self.0.handler.delay_as_bug(self.0.diagnostic.clone());
        self.cancel();
    }

    /// Adds a span/label to be included in the resulting snippet.
    /// This is pushed onto the `MultiSpan` that was created when the
    /// diagnostic was first built. If you don't call this function at
    /// all, and you just supplied a `Span` to create the diagnostic,
    /// then the snippet will just include that `Span`, which is
    /// called the primary span.
    pub fn span_label(&mut self, span: Span, label: impl Into<String>) -> &mut Self {
        self.0.diagnostic.span_label(span, label);
        self
    }

    /// Labels all the given spans with the provided label.
    /// See `span_label` for more information.
    pub fn span_labels(
        &mut self,
        spans: impl IntoIterator<Item = Span>,
        label: impl AsRef<str>,
    ) -> &mut Self {
        let label = label.as_ref();
        for span in spans {
            self.0.diagnostic.span_label(span, label);
        }
        self
    }

    forward!(pub fn note_expected_found(
        &mut self,
        expected_label: &dyn fmt::Display,
        expected: DiagnosticStyledString,
        found_label: &dyn fmt::Display,
        found: DiagnosticStyledString,
    ) -> &mut Self);

    forward!(pub fn note_expected_found_extra(
        &mut self,
        expected_label: &dyn fmt::Display,
        expected: DiagnosticStyledString,
        found_label: &dyn fmt::Display,
        found: DiagnosticStyledString,
        expected_extra: &dyn fmt::Display,
        found_extra: &dyn fmt::Display,
    ) -> &mut Self);

    forward!(pub fn note_unsuccessfull_coercion(
        &mut self,
        expected: DiagnosticStyledString,
        found: DiagnosticStyledString,
    ) -> &mut Self);

    forward!(pub fn note(&mut self, msg: &str) -> &mut Self);
    forward!(pub fn span_note<S: Into<MultiSpan>>(
        &mut self,
        sp: S,
        msg: &str,
    ) -> &mut Self);
    forward!(pub fn warn(&mut self, msg: &str) -> &mut Self);
    forward!(pub fn span_warn<S: Into<MultiSpan>>(&mut self, sp: S, msg: &str) -> &mut Self);
    forward!(pub fn help(&mut self, msg: &str) -> &mut Self);
    forward!(pub fn span_help<S: Into<MultiSpan>>(
        &mut self,
        sp: S,
        msg: &str,
    ) -> &mut Self);

    pub fn multipart_suggestion(
        &mut self,
        msg: &str,
        suggestion: Vec<(Span, String)>,
        applicability: Applicability,
    ) -> &mut Self {
        if !self.0.allow_suggestions {
            return self;
        }
        self.0.diagnostic.multipart_suggestion(msg, suggestion, applicability);
        self
    }

    pub fn tool_only_multipart_suggestion(
        &mut self,
        msg: &str,
        suggestion: Vec<(Span, String)>,
        applicability: Applicability,
    ) -> &mut Self {
        if !self.0.allow_suggestions {
            return self;
        }
        self.0.diagnostic.tool_only_multipart_suggestion(msg, suggestion, applicability);
        self
    }

    pub fn span_suggestion(
        &mut self,
        sp: Span,
        msg: &str,
        suggestion: String,
        applicability: Applicability,
    ) -> &mut Self {
        if !self.0.allow_suggestions {
            return self;
        }
        self.0.diagnostic.span_suggestion(sp, msg, suggestion, applicability);
        self
    }

    pub fn span_suggestions(
        &mut self,
        sp: Span,
        msg: &str,
        suggestions: impl Iterator<Item = String>,
        applicability: Applicability,
    ) -> &mut Self {
        if !self.0.allow_suggestions {
            return self;
        }
        self.0.diagnostic.span_suggestions(sp, msg, suggestions, applicability);
        self
    }

    pub fn span_suggestion_short(
        &mut self,
        sp: Span,
        msg: &str,
        suggestion: String,
        applicability: Applicability,
    ) -> &mut Self {
        if !self.0.allow_suggestions {
            return self;
        }
        self.0.diagnostic.span_suggestion_short(sp, msg, suggestion, applicability);
        self
    }

    pub fn span_suggestion_hidden(
        &mut self,
        sp: Span,
        msg: &str,
        suggestion: String,
        applicability: Applicability,
    ) -> &mut Self {
        if !self.0.allow_suggestions {
            return self;
        }
        self.0.diagnostic.span_suggestion_hidden(sp, msg, suggestion, applicability);
        self
    }

    pub fn tool_only_span_suggestion(
        &mut self,
        sp: Span,
        msg: &str,
        suggestion: String,
        applicability: Applicability,
    ) -> &mut Self {
        if !self.0.allow_suggestions {
            return self;
        }
        self.0.diagnostic.tool_only_span_suggestion(sp, msg, suggestion, applicability);
        self
    }

    forward!(pub fn set_span<S: Into<MultiSpan>>(&mut self, sp: S) -> &mut Self);
    forward!(pub fn code(&mut self, s: DiagnosticId) -> &mut Self);

    pub fn allow_suggestions(&mut self, allow: bool) -> &mut Self {
        self.0.allow_suggestions = allow;
        self
    }

    /// Convenience function for internal use, clients should use one of the
    /// struct_* methods on Handler.
    crate fn new(handler: &'a Handler, level: Level, message: &str) -> DiagnosticBuilder<'a> {
        DiagnosticBuilder::new_with_code(handler, level, None, message)
    }

    /// Convenience function for internal use, clients should use one of the
    /// struct_* methods on Handler.
    crate fn new_with_code(
        handler: &'a Handler,
        level: Level,
        code: Option<DiagnosticId>,
        message: &str,
    ) -> DiagnosticBuilder<'a> {
        let diagnostic = Diagnostic::new_with_code(level, code, message);
        DiagnosticBuilder::new_diagnostic(handler, diagnostic)
    }

    /// Creates a new `DiagnosticBuilder` with an already constructed
    /// diagnostic.
    crate fn new_diagnostic(handler: &'a Handler, diagnostic: Diagnostic) -> DiagnosticBuilder<'a> {
        debug!("Created new diagnostic");
        DiagnosticBuilder(Box::new(DiagnosticBuilderInner {
            handler,
            diagnostic,
            allow_suggestions: true,
        }))
    }
}

impl<'a> Debug for DiagnosticBuilder<'a> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        self.0.diagnostic.fmt(f)
    }
}

/// Destructor bomb - a `DiagnosticBuilder` must be either emitted or canceled
/// or we emit a bug.
impl<'a> Drop for DiagnosticBuilder<'a> {
    fn drop(&mut self) {
        if !panicking() && !self.cancelled() {
            let mut db = DiagnosticBuilder::new(
                self.0.handler,
                Level::Bug,
                "the following error was constructed but not emitted",
            );
            db.emit();
            self.emit();
            panic!();
        }
    }
}

#[macro_export]
macro_rules! struct_span_err {
    ($session:expr, $span:expr, $code:ident, $($message:tt)*) => ({
        $session.struct_span_err_with_code(
            $span,
            &format!($($message)*),
            $crate::error_code!($code),
        )
    })
}

#[macro_export]
macro_rules! error_code {
    ($code:ident) => {{ $crate::DiagnosticId::Error(stringify!($code).to_owned()) }};
}
<commit_msg>Fix a leak in `DiagnosticBuilder::into_diagnostic`.<commit_after>use crate::{Applicability, Handler, Level, StashKey};
use crate::{Diagnostic, DiagnosticId, DiagnosticStyledString};

use log::debug;
use rustc_span::{MultiSpan, Span};
use std::fmt::{self, Debug};
use std::ops::{Deref, DerefMut};
use std::thread::panicking;

/// Used for emitting structured error messages and other diagnostic information.
///
/// If there is some state in a downstream crate you would like to
/// access in the methods of `DiagnosticBuilder` here, consider
/// extending `HandlerFlags`, accessed via `self.handler.flags`.
#[must_use]
#[derive(Clone)]
pub struct DiagnosticBuilder<'a>(Box<DiagnosticBuilderInner<'a>>);

/// This is a large type, and often used as a return value, especially within
/// the frequently-used `PResult` type. In theory, return value optimization
/// (RVO) should avoid unnecessary copying. In practice, it does not (at the
/// time of writing). The split between `DiagnosticBuilder` and
/// `DiagnosticBuilderInner` exists to avoid many `memcpy` calls.
#[must_use]
#[derive(Clone)]
struct DiagnosticBuilderInner<'a> {
    handler: &'a Handler,
    diagnostic: Diagnostic,
    allow_suggestions: bool,
}

/// In general, the `DiagnosticBuilder` uses deref to allow access to
/// the fields and methods of the embedded `diagnostic` in a
/// transparent way. *However,* many of the methods are intended to
/// be used in a chained way, and hence ought to return `self`. In
/// that case, we can't just naively forward to the method on the
/// `diagnostic`, because the return type would be a `&Diagnostic`
/// instead of a `&DiagnosticBuilder<'a>`. This `forward!` macro makes
/// it easy to declare such methods on the builder.
macro_rules! forward {
    // Forward pattern for &self -> &Self
    (
        $(#[$attrs:meta])*
        pub fn $n:ident(&self, $($name:ident: $ty:ty),* $(,)?) -> &Self
    ) => {
        $(#[$attrs])*
        pub fn $n(&self, $($name: $ty),*) -> &Self {
            self.diagnostic.$n($($name),*);
            self
        }
    };

    // Forward pattern for &mut self -> &mut Self
    (
        $(#[$attrs:meta])*
        pub fn $n:ident(&mut self, $($name:ident: $ty:ty),* $(,)?) -> &mut Self
    ) => {
        $(#[$attrs])*
        pub fn $n(&mut self, $($name: $ty),*) -> &mut Self {
            self.0.diagnostic.$n($($name),*);
            self
        }
    };

    // Forward pattern for &mut self -> &mut Self, with S: Into<MultiSpan>
    // type parameter. No obvious way to make this more generic.
    (
        $(#[$attrs:meta])*
        pub fn $n:ident<S: Into<MultiSpan>>(
            &mut self,
            $($name:ident: $ty:ty),*
            $(,)?
        ) -> &mut Self
    ) => {
        $(#[$attrs])*
        pub fn $n<S: Into<MultiSpan>>(&mut self, $($name: $ty),*) -> &mut Self {
            self.0.diagnostic.$n($($name),*);
            self
        }
    };
}

impl<'a> Deref for DiagnosticBuilder<'a> {
    type Target = Diagnostic;

    fn deref(&self) -> &Diagnostic {
        &self.0.diagnostic
    }
}

impl<'a> DerefMut for DiagnosticBuilder<'a> {
    fn deref_mut(&mut self) -> &mut Diagnostic {
        &mut self.0.diagnostic
    }
}

impl<'a> DiagnosticBuilder<'a> {
    /// Emit the diagnostic.
    pub fn emit(&mut self) {
        self.0.handler.emit_diagnostic(&self);
        self.cancel();
    }

    /// Emit the diagnostic unless `delay` is true,
    /// in which case the emission will be delayed as a bug.
    ///
    /// See `emit` and `delay_as_bug` for details.
    pub fn emit_unless(&mut self, delay: bool) {
        if delay {
            self.delay_as_bug();
        } else {
            self.emit();
        }
    }

    /// Stashes diagnostic for possible later improvement in a different,
    /// later stage of the compiler. The diagnostic can be accessed with
    /// the provided `span` and `key` through `.steal_diagnostic` on `Handler`.
    ///
    /// As with `buffer`, this is unless the handler has disabled such buffering.
    pub fn stash(self, span: Span, key: StashKey) {
        if let Some((diag, handler)) = self.into_diagnostic() {
            handler.stash_diagnostic(span, key, diag);
        }
    }

    /// Converts the builder to a `Diagnostic` for later emission,
    /// unless handler has disabled such buffering.
    pub fn into_diagnostic(mut self) -> Option<(Diagnostic, &'a Handler)> {
        if self.0.handler.flags.dont_buffer_diagnostics
            || self.0.handler.flags.treat_err_as_bug.is_some()
        {
            self.emit();
            return None;
        }

        let handler = self.0.handler;

        // We must use `Level::Cancelled` for `dummy` to avoid an ICE about an
        // unused diagnostic.
        let dummy = Diagnostic::new(Level::Cancelled, "");
        let diagnostic = std::mem::replace(&mut self.0.diagnostic, dummy);

        // Logging here is useful to help track down where in logs an error was
        // actually emitted.
        debug!("buffer: diagnostic={:?}", diagnostic);

        Some((diagnostic, handler))
    }

    /// Buffers the diagnostic for later emission,
    /// unless handler has disabled such buffering.
    pub fn buffer(self, buffered_diagnostics: &mut Vec<Diagnostic>) {
        buffered_diagnostics.extend(self.into_diagnostic().map(|(diag, _)| diag));
    }

    /// Convenience function for internal use, clients should use one of the
    /// span_* methods instead.
    pub fn sub<S: Into<MultiSpan>>(
        &mut self,
        level: Level,
        message: &str,
        span: Option<S>,
    ) -> &mut Self {
        let span = span.map(|s| s.into()).unwrap_or_else(|| MultiSpan::new());
        self.0.diagnostic.sub(level, message, span, None);
        self
    }

    /// Delay emission of this diagnostic as a bug.
    ///
    /// This can be useful in contexts where an error indicates a bug but
    /// typically this only happens when other compilation errors have already
    /// happened. In those cases this can be used to defer emission of this
    /// diagnostic as a bug in the compiler only if no other errors have been
    /// emitted.
    ///
    /// In the meantime, though, callsites are required to deal with the "bug"
    /// locally in whichever way makes the most sense.
    pub fn delay_as_bug(&mut self) {
        self.level = Level::Bug;
        self.0.handler.delay_as_bug(self.0.diagnostic.clone());
        self.cancel();
    }

    /// Adds a span/label to be included in the resulting snippet.
    /// This is pushed onto the `MultiSpan` that was created when the
    /// diagnostic was first built. If you don't call this function at
    /// all, and you just supplied a `Span` to create the diagnostic,
    /// then the snippet will just include that `Span`, which is
    /// called the primary span.
    pub fn span_label(&mut self, span: Span, label: impl Into<String>) -> &mut Self {
        self.0.diagnostic.span_label(span, label);
        self
    }

    /// Labels all the given spans with the provided label.
    /// See `span_label` for more information.
    pub fn span_labels(
        &mut self,
        spans: impl IntoIterator<Item = Span>,
        label: impl AsRef<str>,
    ) -> &mut Self {
        let label = label.as_ref();
        for span in spans {
            self.0.diagnostic.span_label(span, label);
        }
        self
    }

    forward!(pub fn note_expected_found(
        &mut self,
        expected_label: &dyn fmt::Display,
        expected: DiagnosticStyledString,
        found_label: &dyn fmt::Display,
        found: DiagnosticStyledString,
    ) -> &mut Self);

    forward!(pub fn note_expected_found_extra(
        &mut self,
        expected_label: &dyn fmt::Display,
        expected: DiagnosticStyledString,
        found_label: &dyn fmt::Display,
        found: DiagnosticStyledString,
        expected_extra: &dyn fmt::Display,
        found_extra: &dyn fmt::Display,
    ) -> &mut Self);

    forward!(pub fn note_unsuccessfull_coercion(
        &mut self,
        expected: DiagnosticStyledString,
        found: DiagnosticStyledString,
    ) -> &mut Self);

    forward!(pub fn note(&mut self, msg: &str) -> &mut Self);
    forward!(pub fn span_note<S: Into<MultiSpan>>(
        &mut self,
        sp: S,
        msg: &str,
    ) -> &mut Self);
    forward!(pub fn warn(&mut self, msg: &str) -> &mut Self);
    forward!(pub fn span_warn<S: Into<MultiSpan>>(&mut self, sp: S, msg: &str) -> &mut Self);
    forward!(pub fn help(&mut self, msg: &str) -> &mut Self);
    forward!(pub fn span_help<S: Into<MultiSpan>>(
        &mut self,
        sp: S,
        msg: &str,
    ) -> &mut Self);

    pub fn multipart_suggestion(
        &mut self,
        msg: &str,
        suggestion: Vec<(Span, String)>,
        applicability: Applicability,
    ) -> &mut Self {
        if !self.0.allow_suggestions {
            return self;
        }
        self.0.diagnostic.multipart_suggestion(msg, suggestion, applicability);
        self
    }

    pub fn tool_only_multipart_suggestion(
        &mut self,
        msg: &str,
        suggestion: Vec<(Span, String)>,
        applicability: Applicability,
    ) -> &mut Self {
        if !self.0.allow_suggestions {
            return self;
        }
        self.0.diagnostic.tool_only_multipart_suggestion(msg, suggestion, applicability);
        self
    }

    pub fn span_suggestion(
        &mut self,
        sp: Span,
        msg: &str,
        suggestion: String,
        applicability: Applicability,
    ) -> &mut Self {
        if !self.0.allow_suggestions {
            return self;
        }
        self.0.diagnostic.span_suggestion(sp, msg, suggestion, applicability);
        self
    }

    pub fn span_suggestions(
        &mut self,
        sp: Span,
        msg: &str,
        suggestions: impl Iterator<Item = String>,
        applicability: Applicability,
    ) -> &mut Self {
        if !self.0.allow_suggestions {
            return self;
        }
        self.0.diagnostic.span_suggestions(sp, msg, suggestions, applicability);
        self
    }

    pub fn span_suggestion_short(
        &mut self,
        sp: Span,
        msg: &str,
        suggestion: String,
        applicability: Applicability,
    ) -> &mut Self {
        if !self.0.allow_suggestions {
            return self;
        }
        self.0.diagnostic.span_suggestion_short(sp, msg, suggestion, applicability);
        self
    }

    pub fn span_suggestion_hidden(
        &mut self,
        sp: Span,
        msg: &str,
        suggestion: String,
        applicability: Applicability,
    ) -> &mut Self {
        if !self.0.allow_suggestions {
            return self;
        }
        self.0.diagnostic.span_suggestion_hidden(sp, msg, suggestion, applicability);
        self
    }

    pub fn tool_only_span_suggestion(
        &mut self,
        sp: Span,
        msg: &str,
        suggestion: String,
        applicability: Applicability,
    ) -> &mut Self {
        if !self.0.allow_suggestions {
            return self;
        }
        self.0.diagnostic.tool_only_span_suggestion(sp, msg, suggestion, applicability);
        self
    }

    forward!(pub fn set_span<S: Into<MultiSpan>>(&mut self, sp: S) -> &mut Self);
    forward!(pub fn code(&mut self, s: DiagnosticId) -> &mut Self);

    pub fn allow_suggestions(&mut self, allow: bool) -> &mut Self {
        self.0.allow_suggestions = allow;
        self
    }

    /// Convenience function for internal use, clients should use one of the
    /// struct_* methods on Handler.
    crate fn new(handler: &'a Handler, level: Level, message: &str) -> DiagnosticBuilder<'a> {
        DiagnosticBuilder::new_with_code(handler, level, None, message)
    }

    /// Convenience function for internal use, clients should use one of the
    /// struct_* methods on Handler.
    crate fn new_with_code(
        handler: &'a Handler,
        level: Level,
        code: Option<DiagnosticId>,
        message: &str,
    ) -> DiagnosticBuilder<'a> {
        let diagnostic = Diagnostic::new_with_code(level, code, message);
        DiagnosticBuilder::new_diagnostic(handler, diagnostic)
    }

    /// Creates a new `DiagnosticBuilder` with an already constructed
    /// diagnostic.
    crate fn new_diagnostic(handler: &'a Handler, diagnostic: Diagnostic) -> DiagnosticBuilder<'a> {
        debug!("Created new diagnostic");
        DiagnosticBuilder(Box::new(DiagnosticBuilderInner {
            handler,
            diagnostic,
            allow_suggestions: true,
        }))
    }
}

impl<'a> Debug for DiagnosticBuilder<'a> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        self.0.diagnostic.fmt(f)
    }
}

/// Destructor bomb - a `DiagnosticBuilder` must be either emitted or canceled
/// or we emit a bug.
impl<'a> Drop for DiagnosticBuilder<'a> {
    fn drop(&mut self) {
        if !panicking() && !self.cancelled() {
            let mut db = DiagnosticBuilder::new(
                self.0.handler,
                Level::Bug,
                "the following error was constructed but not emitted",
            );
            db.emit();
            self.emit();
            panic!();
        }
    }
}

#[macro_export]
macro_rules! struct_span_err {
    ($session:expr, $span:expr, $code:ident, $($message:tt)*) => ({
        $session.struct_span_err_with_code(
            $span,
            &format!($($message)*),
            $crate::error_code!($code),
        )
    })
}

#[macro_export]
macro_rules! error_code {
    ($code:ident) => {{ $crate::DiagnosticId::Error(stringify!($code).to_owned()) }};
}
<|endoftext|>


<commit_before>use filter;
use filter_tree;
use walkdir::WalkDir;
use walkdir::DirEntry;
use parser;
use nom::IResult;
use nom::Needed;
use std::fs;
use pretty_bytes::converter::convert as pretty_bytes_convert;
use std::time;
use chrono::{Local, TimeZone};
use mime_guess;
use std::process::Command;
use liquid;
use liquid::{Renderable, Context, Value};

fn format_mimetype(t: mime_guess::Mime, machine_readable: bool) -> String {
    format!("{}", t)
}

fn format_filetype(ft: fs::FileType, machine_readable: bool) -> String {
    if ft.is_file() {
        String::from("file")
    } else if ft.is_dir() {
        String::from("dir")
    } else if ft.is_symlink() {
        String::from("slink")
    } else {
        String::from("unkown")
    }
}

fn format_filesize(size: u64, machine_readable: bool) -> String {
    if machine_readable {
        return format!("{}", size)
    } else {
        pretty_bytes_convert(size as f64)
    }
}

fn format_systime(t: time::SystemTime, machine_readable: bool) -> String {
    let duration = t.duration_since(time::UNIX_EPOCH).unwrap();
    if machine_readable {
        format!("{}", duration.as_secs())
    } else {
       format!("{}", Local.timestamp(duration.as_secs() as i64, 0).format("%F %T"))
    }
}

pub struct Query {
    attributes: Vec<filter::Attribute>,
    directories: Vec<String>,
    filters: filter_tree::FilterTree,
    command: Option<String>,
    machine_mode: bool,
}

impl Query {
    pub fn new(attributes: Option<Vec<filter::Attribute>>, directories: Option<Vec<String>>, filters: Option<filter_tree::FilterTree>, command: Option<String>) -> Query{
        let attr = attributes.unwrap_or(vec![filter::Attribute::Name]);
        let dirs = directories.unwrap_or(vec![String::from(".")]);
        let filters = filters.unwrap_or(filter_tree::FilterTree::new(None, None, None));
        Query{attributes: attr, directories: dirs, filters: filters, command: command, machine_mode: false}
    }

    pub fn parse(inp: &str) -> Query {
        match parser::query(inp.as_bytes()){
            IResult::Done(leftovers, q)     => {
                if leftovers.len() > 1 {
                    panic!("Could not parse from here -> {}\nDid you make a typo?\n\n", String::from_utf8_lossy(leftovers).into_owned());
                };
                q
            },
            IResult::Error(e)               => panic!("Syntax error {}", e),
            IResult::Incomplete(n)          => match n {
                    Needed::Unknown         => panic!("Need more input, but I haven't got a clou how much!"),
                    Needed::Size(n)         => panic!("Need {}bytes more input!", n),
                },
        }
    }

    fn print_attributes(&self, entry: &DirEntry) {
        let mut print_string = String::from("");
        for attribute in &self.attributes {
            let attr_str = match *attribute {
                filter::Attribute::Name	    => format!("{}", entry.path().display()),
                filter::Attribute::Basename	=> format!("{}", entry.file_name().to_str().unwrap()),
                filter::Attribute::Size	    => format_filesize(entry.metadata().unwrap().len(), self.machine_mode),
                filter::Attribute::Mtime	=> format_systime(entry.metadata().unwrap().modified().unwrap(), self.machine_mode),
                filter::Attribute::Ctime	=> {
                    match entry.metadata().unwrap().created() {
                        Ok(t)   => format_systime(t, self.machine_mode),
                        Err(_)  => String::from("N/A"),
                    }},
                filter::Attribute::Atime    => {
                    match entry.metadata().unwrap().accessed() {
                        Ok(t)   => format_systime(t, self.machine_mode),
                        Err(_)  => String::from("N/A"),
                    }},
                filter::Attribute::Filetype	=> format_filetype(entry.metadata().unwrap().file_type(), self.machine_mode),
                filter::Attribute::Mimetype	=> {
                    let filepath = format!("{}", entry.path().display());
                    format_mimetype(mime_guess::guess_mime_type(filepath), self.machine_mode)
                },
                filter::Attribute::Inode	=> format!("{}", entry.ino()),
            };
            if !print_string.is_empty() {
                print_string.push(',');
            }
            print_string.push_str(attr_str.as_str());
        }
        println!("{}", print_string);
    }

    fn setup_context(&self, dir_entry: &DirEntry) -> Context {
        // TODO does this need speed improvement? All attributes will be needed rarely...
        let mut context = Context::new();
        context.set_val("name", Value::Str(String::from(dir_entry.path().to_str().unwrap())));
        context.set_val("basename", Value::Str(String::from(dir_entry.file_name().to_str().unwrap())));
        context.set_val("size", Value::Str(format_filesize(dir_entry.metadata().unwrap().len(), self.machine_mode)));
        context.set_val("mtime", Value::Str(format_systime(dir_entry.metadata().unwrap().modified().unwrap(), self.machine_mode)));
        let ctime = match dir_entry.metadata().unwrap().created() {
                Ok(t)   => format_systime(t, self.machine_mode),
                Err(_)  => String::from("N/A"),
        };
        context.set_val("ctime", Value::Str(ctime));

        let atime = match dir_entry.metadata().unwrap().accessed() {
                Ok(t)   => format_systime(t, self.machine_mode),
                Err(_)  => String::from("N/A"),
            };
        context.set_val("atime", Value::Str(atime));

        context.set_val("filetype", Value::Str(format_filetype(dir_entry.metadata().unwrap().file_type(), self.machine_mode)));
        context.set_val("mimetype", Value::Str(format_mimetype(mime_guess::guess_mime_type(dir_entry.path().to_str().unwrap()), self.machine_mode)));
        context.set_val("inode", Value::Str(format!("{}", dir_entry.ino())));
        return context;
    }

    fn run_command(&self, dir_entry: &DirEntry) {
        match self.command {
            None    => {},
            Some(ref c) => {
                let template = liquid::parse(c, Default::default()).unwrap();
                let mut context = self.setup_context(dir_entry);
                let output = template.render(&mut context);
                match output {
                    Ok(res) => Command::new("sh").arg("-c").arg(res.unwrap()).spawn().expect("Failed to start command."),
                    Err(e)    => panic!("Command template error: {}", e),
                };
            },
        }
    }

    pub fn execute(&mut self, max_depth: usize, machine_mode: bool) {
        if machine_mode {
            self.machine_mode = true
        }

        for dir in &self.directories {
            let dir_iter = WalkDir::new(dir).max_depth(max_depth);

            'files: for entry in dir_iter {
                let entry = match entry{
                    Ok(e)  => e,
                    Err(e) => {
                        println!("Error: {}", e);
                        continue 'files;
                    }
                };
                if self.filters.test(&entry) != true {
                        continue 'files;
                }
                self.print_attributes(&entry);
                self.run_command(&entry);
            }
        }
    }
}
<commit_msg>Fixed bug: default attribute for printing is name<commit_after>use filter;
use filter_tree;
use walkdir::WalkDir;
use walkdir::DirEntry;
use parser;
use nom::IResult;
use nom::Needed;
use std::fs;
use pretty_bytes::converter::convert as pretty_bytes_convert;
use std::time;
use chrono::{Local, TimeZone};
use mime_guess;
use std::process::Command;
use liquid;
use liquid::{Renderable, Context, Value};

fn format_mimetype(t: mime_guess::Mime, machine_readable: bool) -> String {
    format!("{}", t)
}

fn format_filetype(ft: fs::FileType, machine_readable: bool) -> String {
    if ft.is_file() {
        String::from("file")
    } else if ft.is_dir() {
        String::from("dir")
    } else if ft.is_symlink() {
        String::from("slink")
    } else {
        String::from("unkown")
    }
}

fn format_filesize(size: u64, machine_readable: bool) -> String {
    if machine_readable {
        return format!("{}", size)
    } else {
        pretty_bytes_convert(size as f64)
    }
}

fn format_systime(t: time::SystemTime, machine_readable: bool) -> String {
    let duration = t.duration_since(time::UNIX_EPOCH).unwrap();
    if machine_readable {
        format!("{}", duration.as_secs())
    } else {
       format!("{}", Local.timestamp(duration.as_secs() as i64, 0).format("%F %T"))
    }
}

pub struct Query {
    attributes: Vec<filter::Attribute>,
    directories: Vec<String>,
    filters: filter_tree::FilterTree,
    command: Option<String>,
    machine_mode: bool,
}

impl Query {
    pub fn new(attributes: Option<Vec<filter::Attribute>>, directories: Option<Vec<String>>, filters: Option<filter_tree::FilterTree>, command: Option<String>) -> Query{
        let mut attr = attributes.unwrap_or(vec![filter::Attribute::Name]);
        if attr.len() < 1 {
            attr.push(filter::Attribute::Name);
        }
        let dirs = directories.unwrap_or(vec![String::from(".")]);
        let filters = filters.unwrap_or(filter_tree::FilterTree::new(None, None, None));
        Query{attributes: attr, directories: dirs, filters: filters, command: command, machine_mode: false}
    }

    pub fn parse(inp: &str) -> Query {
        match parser::query(inp.as_bytes()){
            IResult::Done(leftovers, q)     => {
                if leftovers.len() > 1 {
                    panic!("Could not parse from here -> {}\nDid you make a typo?\n\n", String::from_utf8_lossy(leftovers).into_owned());
                };
                q
            },
            IResult::Error(e)               => panic!("Syntax error {}", e),
            IResult::Incomplete(n)          => match n {
                    Needed::Unknown         => panic!("Need more input, but I haven't got a clou how much!"),
                    Needed::Size(n)         => panic!("Need {}bytes more input!", n),
                },
        }
    }

    fn print_attributes(&self, entry: &DirEntry) {
        let mut print_string = String::from("");
        for attribute in &self.attributes {
            let attr_str = match *attribute {
                filter::Attribute::Name	    => format!("{}", entry.path().display()),
                filter::Attribute::Basename	=> format!("{}", entry.file_name().to_str().unwrap()),
                filter::Attribute::Size	    => format_filesize(entry.metadata().unwrap().len(), self.machine_mode),
                filter::Attribute::Mtime	=> format_systime(entry.metadata().unwrap().modified().unwrap(), self.machine_mode),
                filter::Attribute::Ctime	=> {
                    match entry.metadata().unwrap().created() {
                        Ok(t)   => format_systime(t, self.machine_mode),
                        Err(_)  => String::from("N/A"),
                    }},
                filter::Attribute::Atime    => {
                    match entry.metadata().unwrap().accessed() {
                        Ok(t)   => format_systime(t, self.machine_mode),
                        Err(_)  => String::from("N/A"),
                    }},
                filter::Attribute::Filetype	=> format_filetype(entry.metadata().unwrap().file_type(), self.machine_mode),
                filter::Attribute::Mimetype	=> {
                    let filepath = format!("{}", entry.path().display());
                    format_mimetype(mime_guess::guess_mime_type(filepath), self.machine_mode)
                },
                filter::Attribute::Inode	=> format!("{}", entry.ino()),
            };
            if !print_string.is_empty() {
                print_string.push(',');
            }
            print_string.push_str(attr_str.as_str());
        }
        println!("{}", print_string);
    }

    fn setup_context(&self, dir_entry: &DirEntry) -> Context {
        // TODO does this need speed improvement? All attributes will be needed rarely...
        let mut context = Context::new();
        context.set_val("name", Value::Str(String::from(dir_entry.path().to_str().unwrap())));
        context.set_val("basename", Value::Str(String::from(dir_entry.file_name().to_str().unwrap())));
        context.set_val("size", Value::Str(format_filesize(dir_entry.metadata().unwrap().len(), self.machine_mode)));
        context.set_val("mtime", Value::Str(format_systime(dir_entry.metadata().unwrap().modified().unwrap(), self.machine_mode)));
        let ctime = match dir_entry.metadata().unwrap().created() {
                Ok(t)   => format_systime(t, self.machine_mode),
                Err(_)  => String::from("N/A"),
        };
        context.set_val("ctime", Value::Str(ctime));

        let atime = match dir_entry.metadata().unwrap().accessed() {
                Ok(t)   => format_systime(t, self.machine_mode),
                Err(_)  => String::from("N/A"),
            };
        context.set_val("atime", Value::Str(atime));

        context.set_val("filetype", Value::Str(format_filetype(dir_entry.metadata().unwrap().file_type(), self.machine_mode)));
        context.set_val("mimetype", Value::Str(format_mimetype(mime_guess::guess_mime_type(dir_entry.path().to_str().unwrap()), self.machine_mode)));
        context.set_val("inode", Value::Str(format!("{}", dir_entry.ino())));
        return context;
    }

    fn run_command(&self, dir_entry: &DirEntry) {
        match self.command {
            None    => {},
            Some(ref c) => {
                let template = liquid::parse(c, Default::default()).unwrap();
                let mut context = self.setup_context(dir_entry);
                let output = template.render(&mut context);
                match output {
                    Ok(res) => Command::new("sh").arg("-c").arg(res.unwrap()).spawn().expect("Failed to start command."),
                    Err(e)    => panic!("Command template error: {}", e),
                };
            },
        }
    }

    pub fn execute(&mut self, max_depth: usize, machine_mode: bool) {
        if machine_mode {
            self.machine_mode = true
        }

        for dir in &self.directories {
            let dir_iter = WalkDir::new(dir).max_depth(max_depth);

            'files: for entry in dir_iter {
                let entry = match entry{
                    Ok(e)  => e,
                    Err(e) => {
                        println!("Error: {}", e);
                        continue 'files;
                    }
                };
                if self.filters.test(&entry) != true {
                        continue 'files;
                }
                self.print_attributes(&entry);
                self.run_command(&entry);
            }
        }
    }
}
<|endoftext|>


<commit_before>use key::Key;
use pattern::{self, Pattern};

pub struct Query {
    patterns: Vec<Pattern>,
    string: String, // TODO: make Cow
}

impl Query {
    fn new(string: String) -> Self {
        Query {
            patterns: pattern::patterns_from_str(string.as_ref()),
            string: string,
        }
    }

    pub fn test(&self, haystack: &[u8]) -> bool {
        self.patterns.iter().all(|p| p.test(haystack))
    }
}

impl AsRef<str> for Query {
    fn as_ref(&self) -> &str {
        self.string.as_ref()
    }
}

pub struct QueryEditor {
    cursor_position: usize,
    string: String,
}

impl QueryEditor {
    pub fn new<S: Into<String>>(string: S) -> QueryEditor {
        QueryEditor {
            cursor_position: 0,
            string: string.into(),
        }
    }

    pub fn put_key(&mut self, key: Key) {
        use key::Key::*;
        match key {
            CtrlA => {
                self.cursor_position = 0;
            }
            CtrlB => {
                if self.cursor_position > 0 {
                    self.cursor_position -= 1;
                }
            }
            CtrlD => {
                if self.cursor_position < self.string.len() {
                    let cursor = self.cursor_position;
                    self.string.remove(cursor);
                }
            }
            CtrlE => {
                self.cursor_position = self.string.len();
            }
            CtrlF => {
                if self.cursor_position < self.string.len() {
                    self.cursor_position += 1;
                }
            }
            CtrlH | Del => {
                if self.cursor_position > 0 {
                    let cursor = self.cursor_position;
                    self.string.remove(cursor - 1);
                    self.cursor_position -= 1;
                }
            }
            CtrlK => {
                let cursor = self.cursor_position;
                self.string.truncate(cursor);
            }
            CtrlW => {
                let cursor = self.cursor_position;
                let word_end = self.string[0..cursor].rfind(|ch| ch != ' ').unwrap_or(0);
                let word_start = self.string[0..word_end].rfind(' ').and_then(|i| Some(i + 1)).unwrap_or(0);
                self.string.truncate(word_start);
                self.cursor_position = word_start;
            }
            Char(ch) => {
                let cursor = self.cursor_position;
                self.string.insert(cursor, ch);
                self.cursor_position += 1;
            }
            _  => { }
        }
    }

    pub fn cursor_position(&self) -> usize {
        self.cursor_position as usize
    }

    pub fn query(&self) -> Query {
        Query::new(self.string.clone())
    }
}

impl AsRef<str> for QueryEditor {
    fn as_ref(&self) -> &str {
        self.string.as_ref()
    }
}
<commit_msg>Move cursor at the end of the initial query<commit_after>use unicode_width::UnicodeWidthStr;

use key::Key;
use pattern::{self, Pattern};

pub struct Query {
    patterns: Vec<Pattern>,
    string: String, // TODO: make Cow
}

impl Query {
    fn new(string: String) -> Self {
        Query {
            patterns: pattern::patterns_from_str(string.as_ref()),
            string: string,
        }
    }

    pub fn test(&self, haystack: &[u8]) -> bool {
        self.patterns.iter().all(|p| p.test(haystack))
    }
}

impl AsRef<str> for Query {
    fn as_ref(&self) -> &str {
        self.string.as_ref()
    }
}

pub struct QueryEditor {
    cursor_position: usize,
    string: String,
}

impl QueryEditor {
    pub fn new<S: Into<String>>(string: S) -> QueryEditor {
        let string = string.into();
        QueryEditor {
            cursor_position: UnicodeWidthStr::width(AsRef::<str>::as_ref(&string)),
            string: string.into(),
        }
    }

    pub fn put_key(&mut self, key: Key) {
        use key::Key::*;
        match key {
            CtrlA => {
                self.cursor_position = 0;
            }
            CtrlB => {
                if self.cursor_position > 0 {
                    self.cursor_position -= 1;
                }
            }
            CtrlD => {
                if self.cursor_position < self.string.len() {
                    let cursor = self.cursor_position;
                    self.string.remove(cursor);
                }
            }
            CtrlE => {
                self.cursor_position = self.string.len();
            }
            CtrlF => {
                if self.cursor_position < self.string.len() {
                    self.cursor_position += 1;
                }
            }
            CtrlH | Del => {
                if self.cursor_position > 0 {
                    let cursor = self.cursor_position;
                    self.string.remove(cursor - 1);
                    self.cursor_position -= 1;
                }
            }
            CtrlK => {
                let cursor = self.cursor_position;
                self.string.truncate(cursor);
            }
            CtrlW => {
                let cursor = self.cursor_position;
                let word_end = self.string[0..cursor].rfind(|ch| ch != ' ').unwrap_or(0);
                let word_start = self.string[0..word_end].rfind(' ').and_then(|i| Some(i + 1)).unwrap_or(0);
                self.string.truncate(word_start);
                self.cursor_position = word_start;
            }
            Char(ch) => {
                let cursor = self.cursor_position;
                self.string.insert(cursor, ch);
                self.cursor_position += 1;
            }
            _  => { }
        }
    }

    pub fn cursor_position(&self) -> usize {
        self.cursor_position as usize
    }

    pub fn query(&self) -> Query {
        Query::new(self.string.clone())
    }
}

impl AsRef<str> for QueryEditor {
    fn as_ref(&self) -> &str {
        self.string.as_ref()
    }
}
<|endoftext|>


<commit_before>use rustc_serialize::json::Json;

use super::Document;

#[derive(Debug)]
pub enum QuerySyntaxError {
    ExpectedObject,
    ExpectedString,
    ExpectedArray,
    UnknownQueryType(String),
    NoQuery,
    FilteredNoFilter,
    FilteredNoQuery,
    MissingQueryString,
    MultiMatchMissingFields,
}

#[derive(Debug)]
pub enum Filter {
    Term(String, Json),
    Prefix(String, String),
    And(Vec<Filter>),
    Or(Vec<Filter>),
    Not(Box<Filter>),
}

impl Filter {
    pub fn matches(&self, doc: &Document) -> bool {
        match *self {
            Filter::Term(ref field, ref value) => {
                let obj = doc.data.as_object().unwrap();

                if let Some(field_value) = obj.get(field) {
                    return field_value == value;
                }

                false
            }
            Filter::Prefix(ref field, ref value) => {
                let obj = doc.data.as_object().unwrap();

                if let Some(field_value) = obj.get(field) {
                    if let Json::String(ref field_value) = *field_value {
                        return field_value.starts_with(value);
                    }
                }

                false
            }
            Filter::And(ref filters) => {
                for filter in filters.iter() {
                    if !filter.matches(doc) {
                        return false;
                    }
                }

                true
            }
            Filter::Or(ref filters) => {
                for filter in filters.iter() {
                    if filter.matches(doc) {
                        return true;
                    }
                }

                false
            }
            Filter::Not(ref filter) => !filter.matches(doc),
        }
    }
}

pub fn parse_filter(json: &Json) -> Filter {
    let filter_json = json.as_object().unwrap();
    let first_key = filter_json.keys().nth(0).unwrap();

    if first_key == "term" {
        let filter_json = filter_json.get("term").unwrap().as_object().unwrap();
        let first_key = filter_json.keys().nth(0).unwrap();

        Filter::Term(first_key.clone(), filter_json.get(first_key).unwrap().clone())
    } else if first_key == "prefix" {
        let filter_json = filter_json.get("prefix").unwrap().as_object().unwrap();
        let first_key = filter_json.keys().nth(0).unwrap();
        let value = filter_json.get(first_key).unwrap().as_string().unwrap();

        Filter::Prefix(first_key.clone(), value.to_owned())
    } else if first_key == "and" {
        Filter::And(filter_json.get("and").unwrap()
                               .as_array().unwrap()
                               .iter().map(|f| parse_filter(f))
                               .collect::<Vec<_>>(),)
    } else if first_key == "or" {
        Filter::Or(filter_json.get("or").unwrap()
                               .as_array().unwrap()
                               .iter().map(|f| parse_filter(f))
                               .collect::<Vec<_>>(),)
    } else if first_key == "not" {
        Filter::Not(Box::new(parse_filter(filter_json.get("not").unwrap())))
    } else {
        Filter::Term("not".to_owned(), Json::String("implemented".to_owned()))
    }
}

#[derive(Debug)]
pub enum Query {
    Match{field: String, query: String},
    MultiMatch{fields: Vec<String>, query: String},
    Filtered{query: Box<Query>, filter: Box<Filter>},
}

impl Query {
    pub fn matches(&self, doc: &Document) -> bool {
        match *self {
            Query::Match{ref field, ref query} => {
                let obj = doc.data.as_object().unwrap();

                if let Some(field_value) = obj.get(field) {
                    let mut field_value = field_value.as_string().unwrap().to_lowercase();
                    let mut query = query.to_lowercase();

                    return field_value.contains(&query);
                }

                false
            }
            Query::MultiMatch{ref fields, ref query} => {
                // TODO
                false
            }
            Query::Filtered{ref query, ref filter} => {
                if filter.matches(doc) {
                    query.matches(doc)
                } else {
                    false
                }
            }
        }
    }
}

pub fn parse_match_query(json: &Json) -> Result<Query, QuerySyntaxError> {
    let json_object = try!(json.as_object().ok_or(QuerySyntaxError::ExpectedObject));
    let first_key = try!(json_object.keys().nth(0).ok_or(QuerySyntaxError::NoQuery));

    Ok(Query::Match {
        field: first_key.clone(),
        query: json_object.get(first_key).unwrap().as_string().unwrap().to_owned(),
    })
}

pub fn parse_multi_match_query(json: &Json) -> Result<Query, QuerySyntaxError> {
    let json_object = try!(json.as_object().ok_or(QuerySyntaxError::ExpectedObject));

    // Convert "fields" into a Vec<String>
    let fields_json = try!(json_object.get("fields").ok_or(QuerySyntaxError::MultiMatchMissingFields));
    let fields = try!(fields_json.as_array().ok_or(QuerySyntaxError::ExpectedArray))
                      .iter().map(|s| s.as_string().unwrap().to_owned())
                      .collect::<Vec<_>>();

    let query_json = try!(json_object.get("query").ok_or(QuerySyntaxError::MissingQueryString));
    let query = try!(query_json.as_string().ok_or(QuerySyntaxError::ExpectedString)).to_owned();

    Ok(Query::MultiMatch {
        fields: fields,
        query: query,
    })
}

pub fn parse_filtered_query(json: &Json) -> Result<Query, QuerySyntaxError> {
    let json_object = try!(json.as_object().ok_or(QuerySyntaxError::ExpectedObject));

    let filter_json = try!(json_object.get("filter").ok_or(QuerySyntaxError::FilteredNoFilter));
    let filter = parse_filter(filter_json);

    let query_json = try!(json_object.get("query").ok_or(QuerySyntaxError::FilteredNoQuery));
    let query = try!(parse_query(query_json));

    Ok(Query::Filtered {
        filter: Box::new(filter),
        query: Box::new(query),
    })
}

pub fn parse_query(json: &Json) -> Result<Query, QuerySyntaxError> {
    let json_object = try!(json.as_object().ok_or(QuerySyntaxError::ExpectedObject));
    let first_key = try!(json_object.keys().nth(0).ok_or(QuerySyntaxError::NoQuery));

    if first_key == "match" {
        let inner_query = json_object.get("match").unwrap();
        Ok(try!(parse_match_query(inner_query)))
    } else if first_key == "multi_match" {
        let inner_query = json_object.get("multi_match").unwrap();
        Ok(try!(parse_multi_match_query(inner_query)))
    } else if first_key == "filtered" {
        let inner_query = json_object.get("filtered").unwrap();
        Ok(try!(parse_filtered_query(inner_query)))
    } else {
        Err(QuerySyntaxError::UnknownQueryType(first_key.clone()))
    }
}
<commit_msg>Added some tests for parse_query<commit_after>use rustc_serialize::json::Json;

use super::Document;

#[derive(Debug, PartialEq)]
pub enum QuerySyntaxError {
    ExpectedObject,
    ExpectedString,
    ExpectedArray,
    UnknownQueryType(String),
    NoQuery,
    FilteredNoFilter,
    FilteredNoQuery,
    MissingQueryString,
    MultiMatchMissingFields,
}

#[derive(Debug, PartialEq)]
pub enum Filter {
    Term(String, Json),
    Prefix(String, String),
    And(Vec<Filter>),
    Or(Vec<Filter>),
    Not(Box<Filter>),
}

impl Filter {
    pub fn matches(&self, doc: &Document) -> bool {
        match *self {
            Filter::Term(ref field, ref value) => {
                let obj = doc.data.as_object().unwrap();

                if let Some(field_value) = obj.get(field) {
                    return field_value == value;
                }

                false
            }
            Filter::Prefix(ref field, ref value) => {
                let obj = doc.data.as_object().unwrap();

                if let Some(field_value) = obj.get(field) {
                    if let Json::String(ref field_value) = *field_value {
                        return field_value.starts_with(value);
                    }
                }

                false
            }
            Filter::And(ref filters) => {
                for filter in filters.iter() {
                    if !filter.matches(doc) {
                        return false;
                    }
                }

                true
            }
            Filter::Or(ref filters) => {
                for filter in filters.iter() {
                    if filter.matches(doc) {
                        return true;
                    }
                }

                false
            }
            Filter::Not(ref filter) => !filter.matches(doc),
        }
    }
}

pub fn parse_filter(json: &Json) -> Filter {
    let filter_json = json.as_object().unwrap();
    let first_key = filter_json.keys().nth(0).unwrap();

    if first_key == "term" {
        let filter_json = filter_json.get("term").unwrap().as_object().unwrap();
        let first_key = filter_json.keys().nth(0).unwrap();

        Filter::Term(first_key.clone(), filter_json.get(first_key).unwrap().clone())
    } else if first_key == "prefix" {
        let filter_json = filter_json.get("prefix").unwrap().as_object().unwrap();
        let first_key = filter_json.keys().nth(0).unwrap();
        let value = filter_json.get(first_key).unwrap().as_string().unwrap();

        Filter::Prefix(first_key.clone(), value.to_owned())
    } else if first_key == "and" {
        Filter::And(filter_json.get("and").unwrap()
                               .as_array().unwrap()
                               .iter().map(|f| parse_filter(f))
                               .collect::<Vec<_>>(),)
    } else if first_key == "or" {
        Filter::Or(filter_json.get("or").unwrap()
                               .as_array().unwrap()
                               .iter().map(|f| parse_filter(f))
                               .collect::<Vec<_>>(),)
    } else if first_key == "not" {
        Filter::Not(Box::new(parse_filter(filter_json.get("not").unwrap())))
    } else {
        Filter::Term("not".to_owned(), Json::String("implemented".to_owned()))
    }
}

#[derive(Debug, PartialEq)]
pub enum Query {
    Match{field: String, query: String},
    MultiMatch{fields: Vec<String>, query: String},
    Filtered{query: Box<Query>, filter: Box<Filter>},
}

impl Query {
    pub fn matches(&self, doc: &Document) -> bool {
        match *self {
            Query::Match{ref field, ref query} => {
                let obj = doc.data.as_object().unwrap();

                if let Some(field_value) = obj.get(field) {
                    let mut field_value = field_value.as_string().unwrap().to_lowercase();
                    let mut query = query.to_lowercase();

                    return field_value.contains(&query);
                }

                false
            }
            Query::MultiMatch{ref fields, ref query} => {
                // TODO
                false
            }
            Query::Filtered{ref query, ref filter} => {
                if filter.matches(doc) {
                    query.matches(doc)
                } else {
                    false
                }
            }
        }
    }
}

pub fn parse_match_query(json: &Json) -> Result<Query, QuerySyntaxError> {
    let json_object = try!(json.as_object().ok_or(QuerySyntaxError::ExpectedObject));
    let first_key = try!(json_object.keys().nth(0).ok_or(QuerySyntaxError::NoQuery));

    Ok(Query::Match {
        field: first_key.clone(),
        query: json_object.get(first_key).unwrap().as_string().unwrap().to_owned(),
    })
}

pub fn parse_multi_match_query(json: &Json) -> Result<Query, QuerySyntaxError> {
    let json_object = try!(json.as_object().ok_or(QuerySyntaxError::ExpectedObject));

    // Convert "fields" into a Vec<String>
    let fields_json = try!(json_object.get("fields").ok_or(QuerySyntaxError::MultiMatchMissingFields));
    let fields = try!(fields_json.as_array().ok_or(QuerySyntaxError::ExpectedArray))
                      .iter().map(|s| s.as_string().unwrap().to_owned())
                      .collect::<Vec<_>>();

    let query_json = try!(json_object.get("query").ok_or(QuerySyntaxError::MissingQueryString));
    let query = try!(query_json.as_string().ok_or(QuerySyntaxError::ExpectedString)).to_owned();

    Ok(Query::MultiMatch {
        fields: fields,
        query: query,
    })
}

pub fn parse_filtered_query(json: &Json) -> Result<Query, QuerySyntaxError> {
    let json_object = try!(json.as_object().ok_or(QuerySyntaxError::ExpectedObject));

    let filter_json = try!(json_object.get("filter").ok_or(QuerySyntaxError::FilteredNoFilter));
    let filter = parse_filter(filter_json);

    let query_json = try!(json_object.get("query").ok_or(QuerySyntaxError::FilteredNoQuery));
    let query = try!(parse_query(query_json));

    Ok(Query::Filtered {
        filter: Box::new(filter),
        query: Box::new(query),
    })
}

pub fn parse_query(json: &Json) -> Result<Query, QuerySyntaxError> {
    let json_object = try!(json.as_object().ok_or(QuerySyntaxError::ExpectedObject));
    let first_key = try!(json_object.keys().nth(0).ok_or(QuerySyntaxError::NoQuery));

    if first_key == "match" {
        let inner_query = json_object.get("match").unwrap();
        Ok(try!(parse_match_query(inner_query)))
    } else if first_key == "multi_match" {
        let inner_query = json_object.get("multi_match").unwrap();
        Ok(try!(parse_multi_match_query(inner_query)))
    } else if first_key == "filtered" {
        let inner_query = json_object.get("filtered").unwrap();
        Ok(try!(parse_filtered_query(inner_query)))
    } else {
        Err(QuerySyntaxError::UnknownQueryType(first_key.clone()))
    }
}

#[cfg(test)]
mod tests {
    use rustc_serialize::json::Json;
    use super::{Query, Filter, QuerySyntaxError, parse_query};

    #[test]
    fn test_match_query() {
        let query = parse_query(&Json::from_str("
            {
                \"match\": {
                    \"title\": \"Hello world!\"
                }
            }
        ").unwrap());

        assert_eq!(query, Ok(Query::Match{
            field: "title".to_owned(),
            query: "Hello world!".to_owned(),
        }))
    }

    #[test]
    fn test_match_query_without_field() {
        let query = parse_query(&Json::from_str("
            {
                \"match\": {}
            }
        ").unwrap());

        assert_eq!(query, Err(QuerySyntaxError::NoQuery))
    }

    #[test]
    fn test_multi_match_query() {
        let query = parse_query(&Json::from_str("
            {
                \"multi_match\": {
                    \"fields\": [\"title\", \"body\"],
                    \"query\": \"Hello world!\"
                }
            }
        ").unwrap());

        assert_eq!(query, Ok(Query::MultiMatch{
            fields: vec!["title".to_owned(), "body".to_owned()],
            query: "Hello world!".to_owned(),
        }))
    }

    #[test]
    fn test_multi_match_query_without_fields() {
        let query = parse_query(&Json::from_str("
            {
                \"multi_match\": {
                    \"query\": \"Hello world!\"
                }
            }
        ").unwrap());

        assert_eq!(query, Err(QuerySyntaxError::MultiMatchMissingFields))
    }

    #[test]
    fn test_multi_match_query_without_query() {
        let query = parse_query(&Json::from_str("
            {
                \"multi_match\": {
                    \"fields\": [\"title\", \"body\"]
                }
            }
        ").unwrap());

        assert_eq!(query, Err(QuerySyntaxError::MissingQueryString))
    }

    #[test]
    fn test_filtered_query() {
        let query = parse_query(&Json::from_str("
            {
                \"filtered\": {
                    \"query\": {
                        \"match\": {
                            \"title\": \"Hello world!\"
                        }
                    },
                    \"filter\": {
                        \"term\": {
                            \"date\": \"2016-01-25\"
                        }
                    }
                }
            }
        ").unwrap());

        assert_eq!(query, Ok(Query::Filtered{
            query: Box::new(Query::Match{
                field: "title".to_owned(),
                query: "Hello world!".to_owned(),
            }),
            filter: Box::new(Filter::Term(
                "date".to_owned(),
                Json::from_str("\"2016-01-25\"").unwrap(),
            ))
        }))
    }

    #[test]
    fn test_filtered_query_without_query() {
        let query = parse_query(&Json::from_str("
            {
                \"filtered\": {
                    \"filter\": {
                        \"term\": {
                            \"date\": \"2016-01-25\"
                        }
                    }
                }
            }
        ").unwrap());

        assert_eq!(query, Err(QuerySyntaxError::FilteredNoQuery))
    }

    #[test]
    fn test_filtered_query_without_filter() {
        let query = parse_query(&Json::from_str("
            {
                \"filtered\": {
                    \"query\": {
                        \"match\": {
                            \"title\": \"Hello world!\"
                        }
                    }
                }
            }
        ").unwrap());

        assert_eq!(query, Err(QuerySyntaxError::FilteredNoFilter))
    }
}
<|endoftext|>


<commit_before>use std::default::Default;

/// A struct specifying the runtime configuration of a pool.
///
/// `Config` implements `Default`, which provides a set of reasonable default
/// values.
#[deriving(Encodable, Decodable)]
pub struct Config {
    /// The number of connections managed by the pool.
    ///
    /// Must be positive.
    ///
    /// Defaults to 10.
    pub pool_size: uint,
    /// The number of tasks that the pool will use for asynchronous operations
    /// such as connection creation and health checks.
    ///
    /// Must be positive.
    ///
    /// Defaults to 3.
    pub helper_tasks: uint,
    /// If true, the health of a connection will be verified before it is
    /// checked out of the pool.
    ///
    /// Defaults to false.
    pub test_on_check_out: bool,
}

impl Default for Config {
    fn default() -> Config {
        Config {
            pool_size: 10,
            helper_tasks: 3,
            test_on_check_out: false,
        }
    }
}

impl Config {
    /// Determines if the configuration is valid
    pub fn validate(&self) -> Result<(), &'static str> {
        if self.pool_size == 0 {
            return Err("pool_size must be positive");
        }

        if self.helper_tasks == 0 {
            return Err("helper_tasks must be positive");
        }

        Ok(())
    }
}

<commit_msg>Clarify docs<commit_after>use std::default::Default;

/// A struct specifying the runtime configuration of a pool.
///
/// `Config` implements `Default`, which provides a set of reasonable default
/// values.
#[deriving(Encodable, Decodable)]
pub struct Config {
    /// The number of connections managed by the pool.
    ///
    /// Must be positive.
    ///
    /// Defaults to 10.
    pub pool_size: uint,
    /// The number of tasks that the pool will use for asynchronous operations
    /// such as connection creation and health checks.
    ///
    /// Must be positive.
    ///
    /// Defaults to 3.
    pub helper_tasks: uint,
    /// If true, the health of a connection will be verified via a call to
    /// `PoolManager::is_valid` before it is checked out of the pool.
    ///
    /// Defaults to false.
    pub test_on_check_out: bool,
}

impl Default for Config {
    fn default() -> Config {
        Config {
            pool_size: 10,
            helper_tasks: 3,
            test_on_check_out: false,
        }
    }
}

impl Config {
    /// Determines if the configuration is valid
    pub fn validate(&self) -> Result<(), &'static str> {
        if self.pool_size == 0 {
            return Err("pool_size must be positive");
        }

        if self.helper_tasks == 0 {
            return Err("helper_tasks must be positive");
        }

        Ok(())
    }
}

<|endoftext|>